Skip to content

Geometry

jax_morph.geometry.free_space #

free_space()

Free-space (unbounded) metric: displacement = Ra - Rb, shift = R + dR.

Returns the same module-level singleton functions on every call. The state stores the space as static fields (displacement/shift), so sharing one identity keeps every free-space state the same PyTree structure - a jitted simulate/logp compiles once, not once per initial state.

Returns:

  • –

    A (displacement, shift) pair implementing an unbounded space.

jax_morph.geometry.periodic #

periodic(box)

Periodic orthorhombic box of side lengths box (scalar or (dim,)); minimum-image.

Memoised on the box value so periodic(box) returns the same function objects for equal boxes - equal-box states then share one static space identity and do not force a recompile (same discipline as free_space and the generated state-class cache).

Parameters:

  • box –

    Scalar side length or vector of per-axis side lengths.

Returns:

  • –

    A memoized (displacement, shift) pair implementing minimum-image wrapping.

jax_morph.geometry.pairwise_displacements #

pairwise_displacements(positions, displacement)

(N, N, dim) array with [i, j] = displacement(positions[i], positions[j]).

The fundamental pairwise geometry primitive; distances are safe_norm(disp, axis=-1) at the call site (some steps want only displacements, some only distances).

Parameters:

  • positions –

    Cell positions of shape (N, dim).

  • displacement –

    Binary displacement function supplied by the state's space.

Returns:

  • –

    Pairwise displacement vectors of shape (N, N, dim).

jax_morph.geometry.neighbor_sum #

neighbor_sum(pair_values, alive)

Reduce a per-pair tensor over sources: sum_j pair_values[i, j] over alive j != i.

pair_values has shape (N, N, *rest) and the result is (N, *rest). A dense reduction combinator that keeps custom pairwise layers short with uniform alive/self masking. It consumes a materialized dense N x N tensor, so a sparse/neighbor-list backend replaces both it AND its callers - a real future migration, not a drop-in.

Parameters:

  • pair_values –

    Per-pair values of shape (N, N, *rest).

  • alive –

    Boolean liveness mask of shape (N,).

Returns:

  • –

    Per-cell sums of shape (N, *rest), excluding self and dead-cell pairs.