keeping external forces only for one rod for now
@@ -9,8 +9,8 @@ from ._linalg import _batch_matmul, _batch_matvec, _batch_cross
# the base class for rod external forcing
# also the no forcing class
class NoForces:
- def __init__(self, rod_list):
- self.rod_list = rod_list
+ def __init__(self, rod):
+ self.rod = rod
def apply_forces(self):
pass
@@ -21,24 +21,22 @@ class NoForces:
# apply gravity on the list of rod
class GravityForces(NoForces):
- def __init__(self, rod_list, gravity):
- self.rod_list = rod_list
+ def __init__(self, rod, gravity):
+ self.rod = rod
self.gravity = gravity
def apply_forces(self):
- for rod in self.rod_list:
- rod.external_forces += np.outer(self.gravity, rod.mass)
+ self.rod.external_forces += np.outer(self.gravity, self.rod.mass)
# puts constant forces on endpoints
# can be modified for temporal variation
class EndpointForces(NoForces):
- def __init__(self, rod_list, start_force, end_force):
- self.rod_list = rod_list
+ def __init__(self, rod, start_force, end_force):
+ self.rod = rod
self.start_force = start_force
self.end_force = end_force
def apply_forces(self):
- for rod in self.rod_list:
- rod.external_forces[..., 0] += self.start_force
- rod.external_forces[..., -1] += self.end_force
+ self.rod.external_forces[..., 0] += self.start_force
+ self.rod.external_forces[..., -1] += self.end_force
