jaxdem.utils#

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

jaxdem.utils.unit(v: Array) → Array[source][source]#

Normalize vectors along the last axis. v: (…, D) returns: (…, D), unit vectors; zeros map to zeros.

jaxdem.utils.signed_angle(v1: Array, v2: Array) → Array[source][source]#

Directional angle from v1 -> v2 around normal \(\hat{z}\) (right-hand rule), in \([-\pi, \pi)\).

jaxdem.utils.signed_angle_x(v1: Array) → Array[source][source]#

Directional angle from v1 -> \(\hat{x}\) around normal \(\hat{z}\), in \((-\pi, \pi]\).

jaxdem.utils.angle(v1: Array, v2: Array) → Array[source][source]#

angle from v1 -> v2 in \([0, \pi]\)

jaxdem.utils.angle_x(v1: Array) → Array[source][source]#

angle from v1 -> \(\hat{x}\) in \([0, \pi]\)

jaxdem.utils.grid_state(*, n_per_axis: Sequence[int], spacing: Array | ndarray | bool | number | bool | int | float | complex, radius: float = 1.0, mass: float = 1.0, jitter: float = 0.0, vel_range: 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, seed: int = 0, key: Array | None = None) → State[source][source]#

Create a state where particles sit on a rectangular lattice.

Random values can be sampled for particle radii, masses and velocities by specifying *_range arguments, which are interpreted as (min, max) bounds for a uniform distribution. When a range is not provided the corresponding radius or mass argument is used for all particles and the velocity components are sampled in [-1, 1].

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 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 a key will be created from seed.

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

jaxdem.utils.encode_callable(fn: Callable) → str[source][source]#

Return a dotted path like ‘jax._src.nn.functions.gelu’.

jaxdem.utils.decode_callable(path: str) → Callable[source][source]#

Import a callable from a dotted path string.

Modules

angles

Utility functions to compute angles between vectors.

gridState

Utility functions to initialize states with particles arranged in a grid.

linalg

Utility functions to help with linear algebra.

randomState

Utility functions to randomly initialize states.

serialization