Case Setup

Case

Simulation setup.

class jax_sph.case_setup.SimulationSetup(cfg)[source]

Parent class for all simulation setups.

initialize()[source]

Initialize and return everything needed for the numerical setup.

This function sets up all physical and numerical quantities including:

  • reference values and case specific parameters

  • integration time step

  • relaxation or simulation mode

  • equation of state

  • boundary conditions

  • state dictionary

  • function including external force and boundary conditions

Returns: A tuple containing the following elements:

  • cfg (DictConfig): configuration arguments

  • box_size (ndarray): size of the simulation box starting at 0.0

  • state (dict): dictionary containing all field values

  • g_ext_fn (Callable): external force function

  • bc_fn (Callable): boundary conditions function (e.g. velocity at walls)

  • nw_fn (Callable): jit-able wall normal funct. when moving walls, else None

  • eos (Callable): equation of state function

  • key (PRNGKey): random key for sampling

  • displacement_fn (Callable): displacement function for edge features

  • shift_fn (Callable): shift function for position updates

jax_sph.case_setup.set_relaxation(Case, cfg)[source]

Make a relaxation case from a SimulationSetup instance.

Create a child class of a particular SimulationSetup instance and overwrite:

  • _init_pos{2|3}D

  • _box_size{2|3}D

  • _tag{2|3}D

  • _init_velocity{2|3}D

  • _external_acceleration_fn

  • _boundary_conditions_fn

jax_sph.case_setup.load_case(case_root: str, case_py_file: str) → SimulationSetup[source]

Load Case class from Python file.

Parameters:

  • case_root (str) – Path to the case root directory, e.g. “cases/”.

  • case_py_file (str) – Name of the Python case file, e.g. “db.py”.

Partition

Neighbors search backends.

jax_sph.partition.get_particle_cells(idx, cl_capacity, N)[source]

Given a cell list idx of shape (nx, ny, nz, cell_capacity), we first enumerate each cell and then return a list of shape (N,) containing the number of the cell each particle belongs to.

jax_sph.partition.neighbor_list(displacement_or_metric: Callable[[Any, Any], Any] | Callable[[Any, Any], float], box_size: Any, r_cutoff: float, backend: str = 'jaxmd_vmap', dr_threshold: float = 0.0, capacity_multiplier: float = 1.25, disable_cell_list: bool = False, mask_self: bool = True, custom_mask_function: Callable[[Any], Any] | None = None, fractional_coordinates: bool = False, format: NeighborListFormat = NeighborListFormat.Sparse, num_particles_max: int | None = None, num_partitions: int = 1, pbc: Array | None = None) → Callable[[Any, NeighborList | None, int | None], NeighborList][source]

Neighbor lists wrapper. Its arguments are mainly based on the jax-md ones.

Parameters:

  • displacement – A function d(R_a, R_b) that computes the displacement between pairs of points.

  • box_size – Either a float specifying the size of the box or an array of shape [spatial_dim] specifying the box size in each spatial dimension.

  • r_cutoff – A scalar specifying the neighborhood radius.

  • dr_threshold – A scalar specifying the maximum distance particles can move before rebuilding the neighbor list.

  • backend – The backend to use. Can be one of: 1) jaxmd_vmap - the default jax-md neighbor list which vectorizes the computations. 2) jaxmd_scan - a modified jax-md neighbor list which serializes the search into num_partitions chunks to improve the memory efficiency. 3) matscipy - a jit-able implementation with the matscipy neighbor list backend, which runs on CPU and takes variable number of particles smaller or equal to num_particles.

  • capacity_multiplier – A floating point scalar specifying the fractional increase in maximum neighborhood occupancy we allocate compared with the maximum in the example positions.

  • disable_cell_list – An optional boolean. If set to True then the neighbor list is constructed using only distances. This can be useful for debugging but should generally be left as False.

  • mask_self – An optional boolean. Determines whether points can consider themselves to be their own neighbors.

  • custom_mask_function – An optional function. Takes the neighbor array and masks selected elements. Note: The input array to the function is (n_particles, m) where the index of particle 1 is in index in the first dimension of the array, the index of particle 2 is given by the value in the array

  • fractional_coordinates – An optional boolean. Specifies whether positions will be supplied in fractional coordinates in the unit cube, \([0, 1]^d\). If this is set to True then the box_size will be set to 1.0 and the cell size used in the cell list will be set to cutoff / box_size.

  • format – The format of the neighbor list; see the NeighborListFormat() enum for details about the different choices for formats. Defaults to Dense.

  • num_particles_max – only used with the matscipy backend. Based on the largest particles system in a dataset.

  • num_partitions – only used with the jaxmd_scan backend

  • pbc – only used with the matscipy backend. Defines the boundary conditions for each dimension individually. Can have shape (2,) or (3,).

  • **static_kwargs – kwargs that get threaded through the calculation of example positions.

Returns:

A NeighborListFns object that contains a method to allocate a new neighbor list and a method to update an existing neighbor list.