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)