Skip to content

Commit

Permalink
refactor: add mapping function to memory_block_rigid_body to avoid co…
Browse files Browse the repository at this point in the history 
…de repetition
  • Loading branch information
@sy-cui
sy-cui committed Sep 7, 2023
1 parent 854898d commit c6aeb21
Showing 1 changed file with 81 additions and 55 deletions.
136 changes: 81 additions & 55 deletions elastica/memory_block/memory_block_rigid_body.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
__doc__ = """Create block-structure class for collection of rigid body systems."""
import numpy as np
from typing import Sequence, List
from typing import Sequence, Literal

from elastica.rigidbody import RigidBodyBase
from elastica.rigidbody.data_structures import _RigidRodSymplecticStepperMixin
Expand All @@ -15,15 +15,15 @@ def __init__(self, systems: Sequence, system_idx_list: Sequence[np.int64]):
self.system_idx_list = np.array(system_idx_list, dtype=np.int64)

# Allocate block structure using system collection.
self.allocate_block_variables_scalars(systems)
self.allocate_block_variables_vectors(systems)
self.allocate_block_variables_matrix(systems)
self.allocate_block_variables_for_symplectic_stepper(systems)
self._allocate_block_variables_scalars(systems)
self._allocate_block_variables_vectors(systems)
self._allocate_block_variables_matrix(systems)
self._allocate_block_variables_for_symplectic_stepper(systems)

# Initialize the mixin class for symplectic time-stepper.
_RigidRodSymplecticStepperMixin.__init__(self)

def allocate_block_variables_scalars(self, systems: Sequence):
def _allocate_block_variables_scalars(self, systems: Sequence):
"""
This function takes system collection and allocates the variables for
block-structure and references allocated variables back to the systems.
Expand Down Expand Up @@ -55,19 +55,14 @@ def allocate_block_variables_scalars(self, systems: Sequence):
(len(map_scalar_dofs_in_rigid_bodies), self.n_elems)
)

for k, v in map_scalar_dofs_in_rigid_bodies.items():
self.__dict__[k] = np.lib.stride_tricks.as_strided(
self.scalar_dofs_in_rigid_bodies[v], (self.n_elems,)
)

for k, v in map_scalar_dofs_in_rigid_bodies.items():
for system_idx, system in enumerate(systems):
self.__dict__[k][..., system_idx : system_idx + 1] = system.__dict__[k]
system.__dict__[k] = np.ndarray.view(
self.__dict__[k][..., system_idx : system_idx + 1]
)
self._map_system_properties_to_block_memory(
mapping_dict=map_scalar_dofs_in_rigid_bodies,
systems=systems,
block_memory=self.scalar_dofs_in_rigid_bodies,
value_type="scalar",
)

def allocate_block_variables_vectors(self, systems: Sequence):
def _allocate_block_variables_vectors(self, systems: Sequence):
"""
This function takes system collection and allocates the vector variables for
block-structure and references allocated vector variables back to the systems.
Expand Down Expand Up @@ -96,21 +91,14 @@ def allocate_block_variables_vectors(self, systems: Sequence):
(len(map_vector_dofs_in_rigid_bodies), 3 * self.n_elems)
)

for k, v in map_vector_dofs_in_rigid_bodies.items():
self.__dict__[k] = np.lib.stride_tricks.as_strided(
self.vector_dofs_in_rigid_bodies[v], (3, self.n_elems)
)

for k, v in map_vector_dofs_in_rigid_bodies.items():
for system_idx, system in enumerate(systems):
self.__dict__[k][..., system_idx : system_idx + 1] = system.__dict__[
k
].copy()
system.__dict__[k] = np.ndarray.view(
self.__dict__[k][..., system_idx : system_idx + 1]
)
self._map_system_properties_to_block_memory(
mapping_dict=map_vector_dofs_in_rigid_bodies,
systems=systems,
block_memory=self.vector_dofs_in_rigid_bodies,
value_type="vector",
)

def allocate_block_variables_matrix(self, systems: Sequence):
def _allocate_block_variables_matrix(self, systems: Sequence):
"""
This function takes system collection and allocates the matrix variables for
block-structure and references allocated matrix variables back to the systems.
Expand Down Expand Up @@ -139,21 +127,14 @@ def allocate_block_variables_matrix(self, systems: Sequence):
(len(map_matrix_dofs_in_rigid_bodies), 9 * self.n_elems)
)

for k, v in map_matrix_dofs_in_rigid_bodies.items():
self.__dict__[k] = np.lib.stride_tricks.as_strided(
self.matrix_dofs_in_rigid_bodies[v], (3, 3, self.n_elems)
)

for k, v in map_matrix_dofs_in_rigid_bodies.items():
for system_idx, system in enumerate(systems):
self.__dict__[k][..., system_idx : system_idx + 1] = system.__dict__[
k
].copy()
system.__dict__[k] = np.ndarray.view(
self.__dict__[k][..., system_idx : system_idx + 1]
)
self._map_system_properties_to_block_memory(
mapping_dict=map_matrix_dofs_in_rigid_bodies,
systems=systems,
block_memory=self.matrix_dofs_in_rigid_bodies,
value_type="tensor",
)

def allocate_block_variables_for_symplectic_stepper(self, systems: Sequence):
def _allocate_block_variables_for_symplectic_stepper(self, systems: Sequence):
"""
This function takes system collection and allocates the variables used by symplectic
stepper for block-structure and references allocated variables back to the systems.
Expand Down Expand Up @@ -181,11 +162,12 @@ def allocate_block_variables_for_symplectic_stepper(self, systems: Sequence):
"alpha_collection": 3,
}
self.rate_collection = np.zeros((len(map_rate_collection), 3 * self.n_elems))
for k, v in map_rate_collection.items():
self.__dict__[k] = np.lib.stride_tricks.as_strided(
self.rate_collection[v], (3, self.n_elems)
)

self._map_system_properties_to_block_memory(
mapping_dict=map_rate_collection,
systems=systems,
block_memory=self.rate_collection,
value_type="vector",
)
# For Dynamic state update of position Verlet create references
self.v_w_collection = np.lib.stride_tricks.as_strided(
self.rate_collection[0:2], (2, 3 * self.n_elems)
Expand All @@ -195,11 +177,55 @@ def allocate_block_variables_for_symplectic_stepper(self, systems: Sequence):
self.rate_collection[2:-1], (2, 3 * self.n_elems)
)

for k, v in map_rate_collection.items():
def _map_system_properties_to_block_memory(
self,
mapping_dict: dict,
systems: Sequence,
block_memory: np.ndarray,
value_type: Literal["scalar", "vector", "tensor"],
) -> None:
"""Map system (rigid bodies) properties to memory blocks.
Parameters
----------
mapping_dict: dict
Dictionary with attribute names as keys and block row index as values.
systems: Sequence
A sequence containing Cosserat rod objects to map from.
block_memory: ndarray
Memory block that, at the end of the method execution, contains all designated
attributes of all systems.
value_type: str
A string that indicates the shape of the attribute.
Options among "scalar", "vector", and "tensor".
"""
if value_type == "scalar":
view_shape = (self.n_elems,)

elif value_type == "vector":
view_shape = (3, self.n_elems)

elif value_type == "tensor":
view_shape = (3, 3, self.n_elems)

else:
raise ValueError(
"Incorrect value type. Must be one of scalar, vector, and tensor."
)

for k, v in mapping_dict.items():
self.__dict__[k] = np.lib.stride_tricks.as_strided(
block_memory[v],
shape=view_shape,
)

for system_idx, system in enumerate(systems):
self.__dict__[k][..., system_idx : system_idx + 1] = system.__dict__[
k
].copy()
self.__dict__[k][..., system_idx : system_idx + 1] = (
system.__dict__[k]
if value_type == "scalar"
else system.__dict__[k].copy()
)
system.__dict__[k] = np.ndarray.view(
self.__dict__[k][..., system_idx : system_idx + 1]
)

0 comments on commit c6aeb21

Please sign in to comment.