Merge branch 'update-0.3.3' into wip/354-dissipation

elastica/experimental/connection_contact_joint/parallel_connection.py

@@ -7,6 +7,8 @@
 from elastica._linalg import (
     _batch_norm,
     _batch_matvec,
+    _batch_cross,
+    _batch_matrix_transpose,
 )
 
 
@@ -149,13 +151,13 @@ def _apply_forces(
         rod_two_external_forces,
     ):
 
-        rod_one_to_rod_two_connection_vec = (
-            rod_one_director_collection[:, :, index_one].T
-            @ rod_one_direction_vec_in_material_frame
+        rod_one_to_rod_two_connection_vec = _batch_matvec(
+            _batch_matrix_transpose(rod_one_director_collection[:, :, index_one]),
+            rod_one_direction_vec_in_material_frame,
         )
-        rod_two_to_rod_one_connection_vec = (
-            rod_two_director_collection[:, :, index_two].T
-            @ rod_two_direction_vec_in_material_frame
+        rod_two_to_rod_one_connection_vec = _batch_matvec(
+            _batch_matrix_transpose(rod_two_director_collection[:, :, index_two]),
+            rod_two_direction_vec_in_material_frame,
         )
 
         # Compute element positions
@@ -212,33 +214,42 @@ def _apply_forces(
         # we apply a repulsive force.
         center_distance = rod_two_element_position - rod_one_element_position
         center_distance_unit_vec = center_distance / np.linalg.norm(center_distance)
-        penetration = np.linalg.norm(center_distance) - (
+        penetration_strain = np.linalg.norm(center_distance) - (
             rod_one_radius[index_one]
             + offset_rod_one
             + rod_two_radius[index_two]
             + offset_rod_two
         )
 
-        round(penetration, 12)
+        np.round_(penetration_strain, 12, penetration_strain)
         # Contact present only if rods penetrate to each other
-        if penetration < 0:
-            # Hertzian contact
-            contact_force = (
-                -k_repulsive * np.abs(penetration) ** (1.5) * center_distance_unit_vec
-            )
-        else:
-            contact_force = np.zeros(
-                3,
-            )
+        idx_penetrate = np.where(penetration_strain < 0)[0]
+        k_contact = np.zeros(index_one.shape[0])
+        k_contact_temp = -k_repulsive * np.abs(penetration_strain) ** (1.5)
+        k_contact[idx_penetrate] += k_contact_temp[idx_penetrate]
+        contact_force = k_contact * center_distance_unit_vec
 
         # Add contact forces
         total_force += contact_force
 
         # Re-distribute forces from elements to nodes.
-        rod_one_external_forces[..., index_one] += 0.5 * total_force
-        rod_one_external_forces[..., index_one + 1] += 0.5 * total_force
-        rod_two_external_forces[..., index_two] -= 0.5 * total_force
-        rod_two_external_forces[..., index_two + 1] -= 0.5 * total_force
+        block_size = index_one.shape[0]
+        for k in range(block_size):
+            rod_one_external_forces[0, index_one[k]] += 0.5 * total_force[0, k]
+            rod_one_external_forces[1, index_one[k]] += 0.5 * total_force[1, k]
+            rod_one_external_forces[2, index_one[k]] += 0.5 * total_force[2, k]
+
+            rod_one_external_forces[0, index_one[k] + 1] += 0.5 * total_force[0, k]
+            rod_one_external_forces[1, index_one[k] + 1] += 0.5 * total_force[1, k]
+            rod_one_external_forces[2, index_one[k] + 1] += 0.5 * total_force[2, k]
+
+            rod_two_external_forces[0, index_two[k]] -= 0.5 * total_force[0, k]
+            rod_two_external_forces[1, index_two[k]] -= 0.5 * total_force[1, k]
+            rod_two_external_forces[2, index_two[k]] -= 0.5 * total_force[2, k]
+
+            rod_two_external_forces[0, index_two[k] + 1] -= 0.5 * total_force[0, k]
+            rod_two_external_forces[1, index_two[k] + 1] -= 0.5 * total_force[1, k]
+            rod_two_external_forces[2, index_two[k] + 1] -= 0.5 * total_force[2, k]
 
         return (
             rod_one_rd2,
@@ -278,12 +289,31 @@ def _apply_torques(
         torque_on_rod_one = np.cross(rod_one_rd2, spring_force)
         torque_on_rod_two = np.cross(rod_two_rd2, -spring_force)
 
-        torque_on_rod_one_material_frame = (
-            rod_one_director_collection[:, :, index_one] @ torque_on_rod_one
+        torque_on_rod_one_material_frame = _batch_matvec(
+            rod_one_director_collection[:, :, index_one], torque_on_rod_one
         )
-        torque_on_rod_two_material_frame = (
-            rod_two_director_collection[:, :, index_two] @ torque_on_rod_two
+        torque_on_rod_two_material_frame = _batch_matvec(
+            rod_two_director_collection[:, :, index_two], torque_on_rod_two
         )
 
-        rod_one_external_torques[..., index_one] += torque_on_rod_one_material_frame
-        rod_two_external_torques[..., index_two] += torque_on_rod_two_material_frame
+        blocksize = index_one.shape[0]
+        for k in range(blocksize):
+            rod_one_external_torques[
+                0, index_one[k]
+            ] += torque_on_rod_one_material_frame[0, k]
+            rod_one_external_torques[
+                1, index_one[k]
+            ] += torque_on_rod_one_material_frame[1, k]
+            rod_one_external_torques[
+                2, index_one[k]
+            ] += torque_on_rod_one_material_frame[2, k]
+
+            rod_two_external_torques[
+                0, index_two[k]
+            ] += torque_on_rod_two_material_frame[0, k]
+            rod_two_external_torques[
+                1, index_two[k]
+            ] += torque_on_rod_two_material_frame[1, k]
+            rod_two_external_torques[
+                2, index_two[k]
+            ] += torque_on_rod_two_material_frame[2, k]