jaxdem.utils#

Utility functions used to set up simulations and analyze the output.

Functions

grid_state(*, n_per_axis, spacing[, radius, ...])

Create a state where particles sit on a rectangular lattice.

random_state(*, N, dim[, box_size, ...])

Generate N non-overlap-checked particles uniformly in an axis-aligned box.
jaxdem.utils.grid_state(*, n_per_axis: Sequence[int], spacing: Array | ndarray | bool | number | bool | int | float | complex, radius: float = 0.5, mass: float = 1.0, jitter: float = 0.0, key: Array | None = None) → State[source][source]#

Create a state where particles sit on a rectangular lattice.

Parameters:

  • n_per_axis (tuple[int]) – Number of spheres along each axis.

  • spacing (tuple[float] | float) – Centre-to-centre distance; scalar is broadcast to every axis.

  • radius (float) – Shared radius / mass for all particles.

  • mass (float) – Shared radius / mass for all particles.

  • jitter (float) – Add a uniform random offset in the range [-jitter, +jitter] for non-perfect grids (useful to break symmetry).

  • key (PRNG key) – Required when jitter > 0.

Return type:

State

jaxdem.utils.random_state(*, N: int, dim: int, box_size: Array | ndarray | bool | number | bool | int | float | complex | None = None, box_anchor: Array | ndarray | bool | number | bool | int | float | complex | None = None, radius_range: Array | ndarray | bool | number | bool | int | float | complex | None = None, mass_range: Array | ndarray | bool | number | bool | int | float | complex | None = None, vel_range: Array | ndarray | bool | number | bool | int | float | complex | None = None, seed: int = 0) → State[source][source]#

Generate N non-overlap-checked particles uniformly in an axis-aligned box.

Parameters:

  • N – Number of particles.

  • dim – Spatial dimension (2 or 3).

  • box_size – Edge lengths of the domain.

  • box_anchor – Coordinate of the lower box corner.

  • radius_range – min and max values that the radius can take.

  • mass_range – min and max values that the radius can take.

  • vel_range – min and max values that the velocity components can take.

  • seed – Integer for reproducibility.

Returns:

A fully-initialised State instance.

Return type:

State