jaxdem.utils.grid_state#

jaxdem.utils.grid_state(*, n_per_axis: Sequence[int], spacing: ArrayLike | float, radius: float = 1.0, mass: float = 1.0, jitter: float = 0.0, vel_range: ArrayLike | None = None, radius_range: ArrayLike | None = None, mass_range: ArrayLike | None = None, seed: int = 0, key: jax.Array | None = None) → State[source]#

Create a state where particles sit on a rectangular lattice.

Specify *_range arguments to sample random particle radii, masses and velocities. The function interprets each range as (min, max) bounds of a uniform distribution. Without a range, the function uses the corresponding radius or mass argument for all particles and samples the velocity components in [-1, 1].

Parameters:

n_per_axis (tuple[int]) – Number of spheres along each axis.
spacing (tuple[float] | float) – Center-to-center distance. A scalar value applies to every axis.
radius (float) – Shared radius / mass for all particles when the corresponding range is not provided.
mass (float) – Shared radius / mass for all particles when the corresponding range is not provided.
jitter (float) – Add a uniform random offset in the range [-jitter, +jitter] for non-perfect grids (useful to break symmetry).
vel_range (ArrayLike | None) – (min, max) values for the velocity components, radii and masses.
radius_range (ArrayLike | None) – (min, max) values for the velocity components, radii and masses.
mass_range (ArrayLike | None) – (min, max) values for the velocity components, radii and masses.
seed (int) – Integer seed used when key is not supplied.
key (PRNG key, optional) – Controls randomness. If None, the function creates a key from seed.

Return type:

State