def make_indexed_storage(backend, iterable, idx=None):
index = make_Index(backend).from_ndarray(np.array(iterable))
if idx is not None:
result = make_IndexedStorage(backend).indexed(idx, index)
else:
result = index
return result
def test_adaptive_sdm_gamma(backend, gamma, idx, n, cell_id, dt_left, dt, dt_max, is_first_in_pair, expected_dt_left, expected_n_substep):
# Arrange
_gamma = backend.Storage.from_ndarray(np.asarray(gamma))
_idx = make_Index(backend).from_ndarray(np.asarray(idx))
_n = make_IndexedStorage(backend).from_ndarray(_idx, np.asarray(n))
_cell_id = backend.Storage.from_ndarray(np.asarray(cell_id))
_dt_left = backend.Storage.from_ndarray(np.asarray(dt_left))
_is_first_in_pair = make_PairIndicator(backend)(len(n))
_is_first_in_pair.indicator[:] = np.asarray(is_first_in_pair)
_n_substep = backend.Storage.from_ndarray(np.zeros_like(dt_left, dtype=int))
_dt_min = backend.Storage.from_ndarray(np.zeros_like(dt_left))
dt_range = (np.nan, dt_max)
# Act
backend.adaptive_sdm_gamma(_gamma, _n, _cell_id, _dt_left, dt, dt_range, _is_first_in_pair, _n_substep, _dt_min)
# Assert
np.testing.assert_array_almost_equal(_dt_left.to_ndarray(), np.asarray(expected_dt_left))
expected_gamma = np.empty_like(np.asarray(gamma))
for i in range(len(idx)):
if is_first_in_pair[i]:
expected_gamma[i // 2] = (dt - np.asarray(expected_dt_left[cell_id[i]])) / dt * np.asarray(gamma)[
i // 2]
np.testing.assert_array_almost_equal(_gamma.to_ndarray(), expected_gamma)
np.testing.assert_array_equal(_n_substep, np.asarray(expected_n_substep))
def test_remove_zero_n_or_flagged(backend):
# Arrange
n_sd = 44
idx = make_Index(backend).identity_index(n_sd)
data = np.ones(n_sd).astype(np.int64)
data[0], data[n_sd // 2], data[-1] = 0, 0, 0
data = backend.Storage.from_ndarray(data)
data = make_IndexedStorage(backend).indexed(storage=data, idx=idx)
# Act
idx.remove_zero_n_or_flagged(data)
# Assert
assert len(idx) == n_sd - 3
assert (backend.Storage.to_ndarray(data)[idx.to_ndarray()[:len(idx)]] >
0).all()
def test_compute_gamma(backend):
# Arrange
n = 87
prob = np.linspace(0, 3, n, endpoint=True)
rand = np.linspace(0, 1, n, endpoint=False)
expected = lambda p, r: p // 1 + (r < p - p // 1)
n_sd = 2
for p in prob:
for r in rand:
# Act
prob_arr = backend.Storage.from_ndarray(
np.full((n_sd // 2, ), p))
rand_arr = backend.Storage.from_ndarray(
np.full((n_sd // 2, ), r))
idx = make_Index(backend).from_ndarray(np.arange(n_sd))
mult = make_IndexedStorage(backend).from_ndarray(
idx,
np.asarray([expected(p, r),
1]).astype(backend.Storage.INT))
_ = backend.Storage.from_ndarray(np.zeros(n_sd // 2))
cell_id = backend.Storage.from_ndarray(
np.zeros(n_sd, dtype=backend.Storage.INT))
indicator = make_PairIndicator(backend)(n_sd)
indicator.indicator[0] = 1
indicator.indicator[1] = 0
backend.compute_gamma(prob_arr,
rand_arr,
mult,
cell_id=cell_id,
is_first_in_pair=indicator,
collision_rate=_,
collision_rate_deficit=_)
# Assert
assert expected(p, r) == prob_arr.to_ndarray()[0]
def __init__(self, n_sd, backend):
assert isinstance(backend, object)
self.__n_sd = n_sd
self.backend = backend
self.formulae = backend.formulae
self.environment = None
self.particles: (Particles, None) = None
self.dynamics = {}
self.products = {}
self.observers = []
self.n_steps = 0
self.sorting_scheme = 'default'
self.condensation_solver = None
self.Index = make_Index(backend)
self.PairIndicator = make_PairIndicator(backend)
self.PairwiseStorage = make_PairwiseStorage(backend)
self.IndexedStorage = make_IndexedStorage(backend)
self.timers = {}
self.null = self.Storage.empty(0, dtype=float)
