Source code for elastica.external_forces

__doc__ = """ Numba implementation module for boundary condition implementations that apply
external forces to the system."""


import numpy as np
from elastica._linalg import _batch_matvec
from elastica.typing import SystemType, RodType
from elastica.utils import _bspline

from numba import njit
from elastica._linalg import _batch_product_i_k_to_ik


[docs]class NoForces: """ This is the base class for external forcing boundary conditions applied to rod-like objects. Notes ----- Every new external forcing class must be derived from NoForces class. """
[docs] def __init__(self): """ NoForces class does not need any input parameters. """ pass
[docs] def apply_forces(self, system: SystemType, time: np.float64 = 0.0): """Apply forces to a rod-like object. In NoForces class, this routine simply passes. Parameters ---------- system : SystemType Rod or rigid-body object time : float The time of simulation. """ pass
[docs] def apply_torques(self, system: SystemType, time: np.float64 = 0.0): """Apply torques to a rod-like object. In NoForces class, this routine simply passes. Parameters ---------- system : SystemType Rod or rigid-body object time : float The time of simulation. """ pass
[docs]class GravityForces(NoForces): """ This class applies a constant gravitational force to the entire rod. Attributes ---------- acc_gravity: numpy.ndarray 1D (dim) array containing data with 'float' type. Gravitational acceleration vector. """
[docs] def __init__(self, acc_gravity=np.array([0.0, -9.80665, 0.0])): """ Parameters ---------- acc_gravity: numpy.ndarray 1D (dim) array containing data with 'float' type. Gravitational acceleration vector. """ super(GravityForces, self).__init__() self.acc_gravity = acc_gravity
[docs] def apply_forces(self, system: SystemType, time=0.0): self.compute_gravity_forces( self.acc_gravity, system.mass, system.external_forces )
@staticmethod @njit(cache=True) def compute_gravity_forces(acc_gravity, mass, external_forces): """ This function add gravitational forces on the nodes. We are using njit decorated function to increase the speed. Parameters ---------- acc_gravity: numpy.ndarray 1D (dim) array containing data with 'float' type. Gravitational acceleration vector. mass: numpy.ndarray 1D (blocksize) array containing data with 'float' type. Mass on the nodes. external_forces: numpy.ndarray 2D (dim, blocksize) array containing data with 'float' type. External force vector. """ inplace_addition(external_forces, _batch_product_i_k_to_ik(acc_gravity, mass))
[docs]class EndpointForces(NoForces): """ This class applies constant forces on the endpoint nodes. Attributes ---------- start_force: numpy.ndarray 1D (dim) array containing data with 'float' type. Force applied to first node of the system. end_force: numpy.ndarray 1D (dim) array containing data with 'float' type. Force applied to last node of the system. ramp_up_time: float Applied forces are ramped up until ramp up time. """
[docs] def __init__(self, start_force, end_force, ramp_up_time): """ Parameters ---------- start_force: numpy.ndarray 1D (dim) array containing data with 'float' type. Force applied to first node of the system. end_force: numpy.ndarray 1D (dim) array containing data with 'float' type. Force applied to last node of the system. ramp_up_time: float Applied forces are ramped up until ramp up time. """ super(EndpointForces, self).__init__() self.start_force = start_force self.end_force = end_force assert ramp_up_time > 0.0 self.ramp_up_time = ramp_up_time
[docs] def apply_forces(self, system: SystemType, time=0.0): self.compute_end_point_forces( system.external_forces, self.start_force, self.end_force, time, self.ramp_up_time, )
@staticmethod @njit(cache=True) def compute_end_point_forces( external_forces, start_force, end_force, time, ramp_up_time ): """ Compute end point forces that are applied on the rod using numba njit decorator. Parameters ---------- external_forces: numpy.ndarray 2D (dim, blocksize) array containing data with 'float' type. External force vector. start_force: numpy.ndarray 1D (dim) array containing data with 'float' type. end_force: numpy.ndarray 1D (dim) array containing data with 'float' type. Force applied to last node of the system. time: float ramp_up_time: float Applied forces are ramped up until ramp up time. """ factor = min(1.0, time / ramp_up_time) external_forces[..., 0] += start_force * factor external_forces[..., -1] += end_force * factor
[docs]class UniformTorques(NoForces): """ This class applies a uniform torque to the entire rod. Attributes ---------- torque: numpy.ndarray 2D (dim, 1) array containing data with 'float' type. Total torque applied to a rod-like object. """
[docs] def __init__(self, torque, direction=np.array([0.0, 0.0, 0.0])): """ Parameters ---------- torque: float Torque magnitude applied to a rod-like object. direction: numpy.ndarray 1D (dim) array containing data with 'float' type. Direction in which torque applied. """ super(UniformTorques, self).__init__() self.torque = torque * direction
[docs] def apply_torques(self, system: SystemType, time: np.float64 = 0.0): n_elems = system.n_elems torque_on_one_element = ( _batch_product_i_k_to_ik(self.torque, np.ones((n_elems))) / n_elems ) system.external_torques += _batch_matvec( system.director_collection, torque_on_one_element )
[docs]class UniformForces(NoForces): """ This class applies a uniform force to the entire rod. Attributes ---------- force: numpy.ndarray 2D (dim, 1) array containing data with 'float' type. Total force applied to a rod-like object. """
[docs] def __init__(self, force, direction=np.array([0.0, 0.0, 0.0])): """ Parameters ---------- force: float Force magnitude applied to a rod-like object. direction: numpy.ndarray 1D (dim) array containing data with 'float' type. Direction in which force applied. """ super(UniformForces, self).__init__() self.force = (force * direction).reshape(3, 1)
[docs] def apply_forces(self, rod: RodType, time: np.float64 = 0.0): force_on_one_element = self.force / rod.n_elems rod.external_forces += force_on_one_element # Because mass of first and last node is half rod.external_forces[..., 0] -= 0.5 * force_on_one_element[:, 0] rod.external_forces[..., -1] -= 0.5 * force_on_one_element[:, 0]
[docs]class MuscleTorques(NoForces): """ This class applies muscle torques along the body. The applied muscle torques are treated as applied external forces. This class can apply muscle torques as a traveling wave with a beta spline or only as a traveling wave. For implementation details refer to Gazzola et. al. RSoS. (2018). Attributes ---------- direction: numpy.ndarray 2D (dim, 1) array containing data with 'float' type. Muscle torque direction. angular_frequency: float Angular frequency of traveling wave. wave_number: float Wave number of traveling wave. phase_shift: float Phase shift of traveling wave. ramp_up_time: float Applied muscle torques are ramped up until ramp up time. my_spline: numpy.ndarray 1D (blocksize) array containing data with 'float' type. Generated spline. """
[docs] def __init__( self, base_length, b_coeff, period, wave_number, phase_shift, direction, rest_lengths, ramp_up_time, with_spline=False, ): """ Parameters ---------- base_length: float Rest length of the rod-like object. b_coeff: nump.ndarray 1D array containing data with 'float' type. Beta coefficients for beta-spline. period: float Period of traveling wave. wave_number: float Wave number of traveling wave. phase_shift: float Phase shift of traveling wave. direction: numpy.ndarray 1D (dim) array containing data with 'float' type. Muscle torque direction. ramp_up_time: float Applied muscle torques are ramped up until ramp up time. with_spline: boolean Option to use beta-spline. """ super(MuscleTorques, self).__init__() self.direction = direction # Direction torque applied self.angular_frequency = 2.0 * np.pi / period self.wave_number = wave_number self.phase_shift = phase_shift assert ramp_up_time > 0.0 self.ramp_up_time = ramp_up_time # s is the position of nodes on the rod, we go from node=1 to node=nelem-1, because there is no # torques applied by first and last node on elements. Reason is that we cannot apply torque in an # infinitesimal segment at the beginning and end of rod, because there is no additional element # (at element=-1 or element=n_elem+1) to provide internal torques to cancel out an external # torque. This coupled with the requirement that the sum of all muscle torques has # to be zero results in this condition. self.s = np.cumsum(rest_lengths) self.s /= self.s[-1] if with_spline: assert b_coeff.size != 0, "Beta spline coefficient array (t_coeff) is empty" my_spline, ctr_pts, ctr_coeffs = _bspline(b_coeff) self.my_spline = my_spline(self.s) else: self.my_spline = np.full_like(self.s, fill_value=1.0)
[docs] def apply_torques(self, rod: RodType, time: np.float64 = 0.0): self.compute_muscle_torques( time, self.my_spline, self.s, self.angular_frequency, self.wave_number, self.phase_shift, self.ramp_up_time, self.direction, rod.director_collection, rod.external_torques, )
@staticmethod @njit(cache=True) def compute_muscle_torques( time, my_spline, s, angular_frequency, wave_number, phase_shift, ramp_up_time, direction, director_collection, external_torques, ): # Ramp up the muscle torque factor = min(1.0, time / ramp_up_time) # From the node 1 to node nelem-1 # Magnitude of the torque. Am = beta(s) * sin(2pi*t/T + 2pi*s/lambda + phi) # There is an inconsistency with paper and Elastica cpp implementation. In paper sign in # front of wave number is positive, in Elastica cpp it is negative. torque_mag = ( factor * my_spline * np.sin(angular_frequency * time - wave_number * s + phase_shift) ) # Head and tail of the snake is opposite compared to elastica cpp. We need to iterate torque_mag # from last to first element. torque = _batch_product_i_k_to_ik(direction, torque_mag[::-1]) inplace_addition( external_torques[..., 1:], _batch_matvec(director_collection, torque)[..., 1:], ) inplace_substraction( external_torques[..., :-1], _batch_matvec(director_collection[..., :-1], torque[..., 1:]), )
@njit(cache=True) def inplace_addition(external_force_or_torque, force_or_torque): """ This function does inplace addition. First argument `external_force_or_torque` is the system.external_forces or system.external_torques. Second argument force or torque vector to be added. Parameters ---------- external_force_or_torque: numpy.ndarray 2D (dim, blocksize) array containing data with 'float' type. force_or_torque: numpy.ndarray 2D (dim, blocksize) array containing data with 'float' type. """ blocksize = force_or_torque.shape[1] for i in range(3): for k in range(blocksize): external_force_or_torque[i, k] += force_or_torque[i, k] @njit(cache=True) def inplace_substraction(external_force_or_torque, force_or_torque): """ This function does inplace substraction. First argument `external_force_or_torque` is the system.external_forces or system.external_torques. Second argument force or torque vector to be substracted. Parameters ---------- external_force_or_torque: numpy.ndarray 2D (dim, blocksize) array containing data with 'float' type. force_or_torque: numpy.ndarray 2D (dim, blocksize) array containing data with 'float' type. """ blocksize = force_or_torque.shape[1] for i in range(3): for k in range(blocksize): external_force_or_torque[i, k] -= force_or_torque[i, k]
[docs]class EndpointForcesSinusoidal(NoForces): """ This class applies sinusoidally varying forces to the ends of a rod. Forces are applied in a plane, which is defined by the tangent_direction and normal_direction. Attributes ---------- start_force_mag: float Magnitude of the force that is applied to the start of the rod (node 0). end_force_mag: float Magnitude of the force that is applied to the end of the rod (node -1). ramp_up_time: float Applied forces are applied in the normal direction until time reaches ramp_up_time. normal_direction: np.ndarray An array (3,) contains type float. This is the normal direction of the rod. roll_direction: np.ndarray An array (3,) contains type float. This is the direction perpendicular to rod tangent, and rod normal. Notes ----- In order to see example how to use this class, see joint examples. """
[docs] def __init__( self, start_force_mag, end_force_mag, ramp_up_time=0.0, tangent_direction=np.array([0, 0, 1]), normal_direction=np.array([0, 1, 0]), ): """ Parameters ---------- start_force_mag: float Magnitude of the force that is applied to the start of the system (node 0). end_force_mag: float Magnitude of the force that is applied to the end of the system (node -1). ramp_up_time: float Applied muscle torques are ramped up until ramp up time. tangent_direction: np.ndarray An array (3,) contains type float. This is the tangent direction of the system, or normal of the plane that forces applied. normal_direction: np.ndarray An array (3,) contains type float. This is the normal direction of the system. """ super(EndpointForcesSinusoidal, self).__init__() # Start force self.start_force_mag = start_force_mag self.end_force_mag = end_force_mag # Applied force directions self.normal_direction = normal_direction self.roll_direction = np.cross(normal_direction, tangent_direction) assert ramp_up_time >= 0.0 self.ramp_up_time = ramp_up_time
[docs] def apply_forces(self, system: SystemType, time=0.0): if time < self.ramp_up_time: # When time smaller than ramp up time apply the force in normal direction # First pull the rod upward or downward direction some time. start_force = -2.0 * self.start_force_mag * self.normal_direction end_force = -2.0 * self.end_force_mag * self.normal_direction system.external_forces[..., 0] += start_force system.external_forces[..., -1] += end_force else: # When time is greater than ramp up time, forces are applied in normal # and roll direction or forces are in a plane perpendicular to the # direction. # First force applied to start of the rod roll_forces_start = ( self.start_force_mag * np.cos(0.5 * np.pi * (time - self.ramp_up_time)) * self.roll_direction ) normal_forces_start = ( self.start_force_mag * np.sin(0.5 * np.pi * (time - self.ramp_up_time)) * self.normal_direction ) start_force = roll_forces_start + normal_forces_start # Now force applied to end of the rod roll_forces_end = ( self.end_force_mag * np.cos(0.5 * np.pi * (time - self.ramp_up_time)) * self.roll_direction ) normal_forces_end = ( self.end_force_mag * np.sin(0.5 * np.pi * (time - self.ramp_up_time)) * self.normal_direction ) end_force = roll_forces_end + normal_forces_end # Update external forces system.external_forces[..., 0] += start_force system.external_forces[..., -1] += end_force