Source code for elastica._elastica_numpy._joint

__doc__ = """ Joint between rods module of Elastica Numpy implementation """

import numpy as np
from elastica.utils import Tolerance, MaxDimension
from elastica.joint import FreeJoint


[docs]class ExternalContact(FreeJoint): """ Assumes that the second entity is a rigid body for now, can be changed at a later time Most of the cylinder-cylinder contact SHOULD be implemented as given in this paper: http://larochelle.sdsmt.edu/publications/2005-2009/Collision%20Detection%20of%20Cylindrical%20Rigid%20Bodies%20Using%20Line%20Geometry.pdf but, it isn't (the elastica-cpp kernels are implented)! This is maybe to speed-up the kernel, but it's potentially dangerous as it does not deal with "end" conditions correctly. """
[docs] def __init__(self, k, nu): super().__init__(k, nu) # 0 is min, 1 is max self.aabb_rod = np.empty((MaxDimension.value(), 2)) self.aabb_cylinder = np.empty((MaxDimension.value(), 2))
# Should be a free function, can be jitted def __aabbs_not_intersecting(self, aabb_one, aabb_two): # FIXME : Not memory friendly, Not early exit first_max_second_min = aabb_one[..., 1] < aabb_two[..., 0] first_min_second_max = aabb_two[..., 1] < aabb_one[..., 0] # Returns true if aabbs not intersecting, else if it intersects returns false return np.any(np.logical_or(first_max_second_min, first_min_second_max)) def __find_min_dist(self, x1, e1, x2, e2): """Assumes x2, e2 is one elment for now Will definitely get speedup from numba """ def _batch_inner(first, second): return np.einsum("ij,ij->j", first, second) def _batch_inner_spec(first, second): return np.einsum("ij,i->j", first, second) # Compute distances for all cases def _difference_impl_for_multiple_st(ax1, ae1, at1, ax2, ae2, as2): return ( ax2.reshape(3, 1, -1) + np.einsum("i,jk->ijk", ae2, as2) - ax1.reshape(3, 1, -1) - np.einsum("ik,jk->ijk", ae1, at1) ) def _difference_impl(ax1, ae1, at1, ax2, ae2, as2): return ( ax2.reshape(3, -1) + np.einsum("i,j->ij", ae2, as2) - ax1.reshape(3, -1) - np.einsum("ik,k->ik", ae1, at1) ) e1e1 = _batch_inner(e1, e1) e1e2 = _batch_inner_spec(e1, e2) e2e2 = np.inner(e2, e2) x1e1 = _batch_inner(x1, e1) x1e2 = _batch_inner_spec(x1, e2) x2e1 = _batch_inner_spec(e1, x2) x2e2 = np.inner(x2, e2) # Parameteric representation of line # First line is x_1 + e_1 * t # Second line is x_2 + e_2 * s s = np.empty_like(e1e1) t = np.empty_like(e1e1) parallel_condition = np.abs(1.0 - e1e2 ** 2 / (e1e1 * e2e2)) parallel_idx = parallel_condition < 1e-5 not_parallel_idx = np.bitwise_not(parallel_idx) anything_not_parallel = np.any(not_parallel_idx) if np.any(parallel_idx): # Some are parallel, so do processing t[parallel_idx] = (x2e1[parallel_idx] - x1e1[parallel_idx]) / e1e1[ parallel_idx ] # Comes from taking dot of e1 with a normal t[parallel_idx] = np.clip( t[parallel_idx], 0.0, 1.0 ) # Note : the out version doesn't work s[parallel_idx] = ( x1e2[parallel_idx] + t[parallel_idx] * e1e2[parallel_idx] - x2e2 ) / e2e2 # Same as before s[parallel_idx] = np.clip(s[parallel_idx], 0.0, 1.0) if anything_not_parallel: # Using the Cauchy-Binet formula on eq(7) in docstring referenc s[not_parallel_idx] = ( e1e1[not_parallel_idx] * (x1e2[not_parallel_idx] - x2e2) + e1e2[not_parallel_idx] * (x2e1[not_parallel_idx] - x1e1[not_parallel_idx]) ) / (e1e1[not_parallel_idx] * e2e2 - (e1e2[not_parallel_idx]) ** 2) t[not_parallel_idx] = ( e1e2[not_parallel_idx] * s[not_parallel_idx] + x2e1[not_parallel_idx] - x1e1[not_parallel_idx] ) / e1e1[not_parallel_idx] if anything_not_parallel: # Done here rather than the other loop to avoid # creating copies by selection of selections such as s[not_parallel_idx][idx] # which creates copies and bor # Remnants for non-parallel indices # as parallel selections are always clipped idx1 = s < 0.0 idx2 = s > 1.0 idx3 = t < 0.0 idx4 = t > 1.0 idx = idx1 | idx2 | idx3 | idx4 if np.any(idx): local_e1e1 = e1e1[idx] local_e1e2 = e1e2[idx] local_x1e1 = x1e1[idx] local_x1e2 = x1e2[idx] local_x2e1 = x2e1[idx] potential_t = np.empty((4, local_e1e1.shape[0])) potential_s = np.empty_like(potential_t) potential_dist = np.empty_like(potential_t) # Fill in the possibilities potential_t[0, ...] = (local_x2e1 - local_x1e1) / local_e1e1 potential_s[0, ...] = 0.0 potential_t[1, ...] = ( local_x2e1 + local_e1e2 - local_x1e1 ) / local_e1e1 potential_s[1, ...] = 1.0 potential_t[2, ...] = 0.0 potential_s[2, ...] = (local_x1e2 - x2e2) / e2e2 potential_t[3, ...] = 1.0 potential_s[3, ...] = (local_x1e2 + local_e1e2 - x2e2) / e2e2 np.clip(potential_t, 0.0, 1.0, out=potential_t) np.clip(potential_s, 0.0, 1.0, out=potential_s) potential_difference = _difference_impl_for_multiple_st( x1[..., idx], e1[..., idx], potential_t, x2, e2, potential_s ) np.sqrt( np.einsum( "ijk,ijk->jk", potential_difference, potential_difference ), out=potential_dist, ) min_idx = np.expand_dims(np.argmin(potential_dist, axis=0), axis=0) s[idx] = np.take_along_axis(potential_s, min_idx, axis=0)[ 0 ] # [0] at the end reduces the dimension, you can also squeeze t[idx] = np.take_along_axis(potential_t, min_idx, axis=0)[ 0 ] # [0] at the end reduces the dimension, you can also squeeze return _difference_impl(x1, e1, t, x2, e2, s)
[docs] def apply_forces(self, rod_one, index_one, cylinder_two, index_two): del index_one, index_two # First, check for a global AABB bounding box, and see whether that # intersects # FIXME : Optimization : multiple passes over same array to do min/max max_possible_dimension = np.zeros((3,)) max_possible_dimension[...] = np.amax(rod_one.radius) + np.amax(rod_one.lengths) self.aabb_rod[..., 0] = ( np.amin(rod_one.position_collection, axis=1) - max_possible_dimension ) self.aabb_rod[..., 1] = ( np.amax(rod_one.position_collection, axis=1) + max_possible_dimension ) cylinder_dimensions_in_world_FOR = cylinder_two.director_collection[ ..., 0 ].T @ np.array( [[cylinder_two.radius], [cylinder_two.radius], [0.5 * cylinder_two.length]] ) np.amax( np.abs(cylinder_dimensions_in_world_FOR), axis=1, out=max_possible_dimension ) self.aabb_cylinder[..., 0] = ( cylinder_two.position_collection[..., 0] - max_possible_dimension ) self.aabb_cylinder[..., 1] = ( cylinder_two.position_collection[..., 0] + max_possible_dimension ) if self.__aabbs_not_intersecting(self.aabb_cylinder, self.aabb_rod): return x_rod = rod_one.position_collection[..., :-1] # Discount last node r_rod = rod_one.radius l_rod = rod_one.lengths # We need at start of the element x_cyl = ( cylinder_two.position_collection[..., 0] - 0.5 * cylinder_two.length * cylinder_two.director_collection[2, :, 0] ) r_cyl = cylinder_two.radius # scalar l_cyl = cylinder_two.length sum_r = r_rod + r_cyl # Element-wise bounding box, if outside then don't worry del_x = x_rod - np.expand_dims(x_cyl, axis=1) norm_del_x = np.sqrt(np.einsum("ij,ij->j", del_x, del_x)) idx = norm_del_x <= (sum_r + l_rod + l_cyl) # If everything is out, don't bother to process if not np.any(idx): # idx[0] = True return # Process only the selected idx elements x_selected = x_rod[..., idx] edge_selected = l_rod[idx] * rod_one.tangents[..., idx] edge_cylinder = l_cyl * cylinder_two.director_collection[2, :, 0] # find the shortest line segment between the two centerline # segments : differs from normal cylinder-cylinder intersection distance_vector_collection = self.__find_min_dist( x_selected, edge_selected, x_cyl, edge_cylinder ) distance_vector_lengths = np.sqrt( np.einsum( "ij,ij->j", distance_vector_collection, distance_vector_collection ) ) distance_vector_collection /= distance_vector_lengths gamma = sum_r[idx] - distance_vector_lengths interacting_idx = gamma > -1e-5 # This step is necessary to scatter the masked results back to the original array # With pure masking, information regarding where in the original array we had True/False is lost # For example if idx = a < 0.2 => idx = [True False True] # b = a[idx]. nidx = b < -0.2 => nidx = [True False]. # We now need selections of a such that we somehow do and AND with idx & nidx # But there's a shape mismatch. To overcome that find where in the idx vector we have Trues # widx = np.where(idx)[0] => widx = [0, 2] => idx[widx] = [True True] # Thus idx[widx] &= nidx gives now idx = [True False False] # Final [0] needed because where returns a one-tuple w_idx = np.where(idx)[0] # Get subset of interacting indices from original as explained above idx[w_idx] &= interacting_idx rod_elemental_forces = rod_one.external_forces + rod_one.internal_forces # No better way to do this? padded_idx = np.hstack((idx, False)) # Because at the end there's one more node r_padded_idx = np.roll(padded_idx, 1) rod_elemental_forces = 0.5 * ( rod_elemental_forces[..., padded_idx] + rod_elemental_forces[..., r_padded_idx] ) equilibrium_forces = -rod_elemental_forces + cylinder_two.external_forces pruned_distance_vector_collection = distance_vector_collection[ ..., interacting_idx ] normal_force = np.einsum( "ij,ij->j", equilibrium_forces, pruned_distance_vector_collection ) # Following line same as np.where(normal_force < 0.0, -normal_force, 0.0) np.abs(np.minimum(normal_force, 0.0), out=normal_force) # CHECK FOR GAMMA > 0.0? # Make it a float array of 1's and 0's # mask = (gamma[interacting_idx] > 0.0) * 1.0 mask = np.heaviside(gamma[interacting_idx], 0.0) contact_force = self.k * gamma[interacting_idx] interpenetration_velocity = ( 0.5 * ( rod_one.velocity_collection[..., padded_idx] + rod_one.velocity_collection[..., r_padded_idx] ) - cylinder_two.velocity_collection ) contact_damping_force = self.nu * np.einsum( "ij,ij->j", interpenetration_velocity, pruned_distance_vector_collection ) # magnitude* direction net_contact_force = ( normal_force + 0.5 * mask * (contact_damping_force + contact_force) ) * pruned_distance_vector_collection ## Equivalent statments # padded_idx_mask = padded_idx * 1.0 # r_padded_idx_mask = r_padded_idx * 1.0 padded_idx_mask = np.heaviside(padded_idx, 0.0) r_padded_idx_mask = np.heaviside(r_padded_idx, 0.0) padded_idx_mask[0] = 0.5 r_padded_idx_mask[-1] = 0.5 padded_idx_mask = padded_idx_mask[padded_idx] r_padded_idx_mask = r_padded_idx_mask[r_padded_idx] # Add it to the rods at the end of the day force_on_ith_element = net_contact_force * padded_idx_mask force_on_i_plus_oneth_element = net_contact_force * r_padded_idx_mask rod_one.external_forces[..., padded_idx] -= force_on_ith_element rod_one.external_forces[..., r_padded_idx] -= force_on_i_plus_oneth_element cylinder_two.external_forces[..., 0] += np.sum( force_on_ith_element + force_on_i_plus_oneth_element, axis=1 )