def overload_operator_kinematic_numba(
n_nodes,
prefac,
position_collection,
director_collection,
velocity_collection,
omega_collection,
):
"""overloaded += operator
The add for directors is customized to reflect Rodrigues' rotation
formula.
Parameters
----------
scaled_deriv_array : np.ndarray containing dt * (v, ω),
as retured from _DynamicState's `kinematic_rates` method
Returns
-------
self : _KinematicState instance with inplace modified data
Caveats
-------
Takes a np.ndarray and not a _KinematicState object (as one expects).
This is done for efficiency reasons, see _DynamicState's `kinematic_rates`
method
"""
# x += v*dt
for i in range(3):
for k in range(n_nodes):
position_collection[i, k] += prefac * velocity_collection[i, k]
rotation_matrix = _get_rotation_matrix(1.0, prefac * omega_collection)
director_collection[:] = _batch_matmul(rotation_matrix, director_collection)
return
def test_batch_matmul(blocksize):
input_first_matrix_collection = np.random.randn(3, 3, blocksize)
input_second_matrix_collection = np.random.randn(3, 3, blocksize)
test_matrix_collection = _batch_matmul(input_first_matrix_collection,
input_second_matrix_collection)
correct_matrix_collection = np.empty((3, 3, blocksize))
for i in range(blocksize):
correct_matrix_collection[..., i] = np.dot(
input_first_matrix_collection[..., i],
input_second_matrix_collection[..., i],
)
assert_allclose(test_matrix_collection, correct_matrix_collection)
def _rotate(director_collection, scale: float, axis_collection):
"""
Does alibi rotations
https://en.wikipedia.org/wiki/Rotation_matrix#Ambiguities
Parameters
----------
director_collection
scale
axis_collection
Returns
-------
# TODO Finish documentation
"""
# return _batch_matmul(
# director_collection, _get_rotation_matrix(scale, axis_collection)
# )
return _batch_matmul(_get_rotation_matrix(scale, axis_collection),
director_collection)