Geometry#

Distance and displacement functions for open boundary conditions (molecules), periodic boundary conditions (solids), and spherical geometry (quantum Hall).

Open boundary conditions#

jaqmc.geometry.obc.pair_displacements_within(positions)[source]#

Computes pairwise displacements and distances within one species.

Parameters:

positions (Array) – Particle positions with shape (n_particles, ndim).

Returns:

  • disp – Pairwise displacement vectors r_i - r_j with shape (n_particles, n_particles, ndim).

  • r – Pairwise distances with shape (n_particles, n_particles).

Return type:

tuple[Array, Array]

Raises:

ValueError – If positions does not have shape (n_particles, ndim).

Examples

>>> from jax import numpy as jnp
>>> pos = jnp.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
>>> disp, r = pair_displacements_within(pos)
>>> disp.shape
(3, 3, 2)
>>> r.shape
(3, 3)
jaqmc.geometry.obc.pair_displacements_between(positions_a, positions_b)[source]#

Computes pairwise displacements and distances across two species.

Parameters:
  • positions_a (Array) – Positions of the first species, shape (n_a, ndim).

  • positions_b (Array) – Positions of the second species, shape (n_b, ndim).

Returns:

  • disp – Pairwise displacement vectors r_a - r_b with shape (n_a, n_b, ndim).

  • r – Pairwise distances with shape (n_a, n_b).

Return type:

tuple[Array, Array]

Raises:

ValueError – If inputs are not two-dimensional or if their spatial dimensions (ndim) do not match.

Periodic boundary conditions#

class jaqmc.geometry.pbc.DistanceType(*values)[source]#

Periodic distance functions for solid-state systems.

Maps particle separations into smooth representations that respect periodicity. Both functions take an (over-complete) set of real-space lattice vectors \(\mathbf{a}_i\) and reciprocal vectors \(\mathbf{b}_i\) produced by get_symmetry_lat(), and compute fractional projections \(\omega_i = \mathbf{b}_i \cdot \mathbf{r}\) (wrapped to \([-\pi, \pi]\)).

nu[source]#

Polynomial distance. Defines two smooth, periodic-compatible polynomials:

\[f(\omega) = \lvert\omega\rvert \bigl(1 - \tfrac{1}{4}\lvert\omega/\pi\rvert^3\bigr)\]
\[g(\omega) = \omega \bigl(1 - \tfrac{3}{2}\lvert\omega/\pi\rvert + \tfrac{1}{2}\lvert\omega/\pi\rvert^2\bigr)\]

and computes the distance as:

\[d(\mathbf{r}) = \sqrt{ \sum_i \lVert\mathbf{a}_i\rVert^2 f(\omega_i)^2 + \sum_{i \neq j} (\mathbf{a}_i \cdot \mathbf{a}_j) \, g(\omega_i)\, g(\omega_j) }\]

Produces 3D relative coordinates \(\sum_i g(\omega_i)\,\mathbf{a}_i\). Works well for most systems.

tri[source]#

Trigonometric distance. Uses the metric tensor \(G_{ij} = \mathbf{a}_i \cdot \mathbf{a}_j\) and:

\[V_{ij} = \sin(\omega_i)\sin(\omega_j) + (1 - \cos(\omega_i))(1 - \cos(\omega_j))\]

to compute:

\[d(\mathbf{r}) = \sqrt{\sum_{ij} V_{ij}\, G_{ij}}\]

Produces 6D relative coordinates by concatenating \(\sum_i \sin(\omega_i)\,\mathbf{a}_i\) and \(\sum_i \cos(\omega_i)\,\mathbf{a}_i\). More expressive than nu at the cost of doubling the feature dimension.

class jaqmc.geometry.pbc.SymmetryType(*values)[source]#

Lattice symmetry types for periodic feature construction.

Expands the primitive reciprocal basis \(\mathbf{b}_\text{base}\) into an overcomplete set \(\mathbf{b} = M \mathbf{b}_\text{base}\) by applying integer linear combinations. The expanded vectors are used by the periodic distance functions (DistanceType) to construct symmetry-aware features. The primitive basis alone may miss symmetry-equivalent directions of the crystal.

minimal[source]#

Identity matrix (3 vectors). No expansion.

fcc[source]#

Adds \([1,1,1]\) combination (4 vectors).

bcc[source]#

Adds \([1,-1,0]\), \([1,0,-1]\), \([0,1,-1]\) combinations (6 vectors).

hexagonal[source]#

Adds \([-1,-1,0]\) combination (4 vectors).

jaqmc.geometry.pbc.build_distance_fn(lattice)[source]#

Computes minimal image distance between particles under PBC.

Parameters:

lattice (Array) – Lattice vectors with shape (ndim, ndim).

Returns:

A function dist_fn(pos_a, pos_b) -> (disp, r).

jaqmc.geometry.pbc.wrap_positions(positions, lattice)[source]#

Wraps positions into the primary unit cell.

Parameters:
  • positions (Array) – Particle positions with shape (…, ndim).

  • lattice (Array) – Lattice vectors with shape (ndim, ndim), where each row is a lattice vector.

Return type:

Array

Returns:

Wrapped positions with the same shape as positions.

jaqmc.geometry.pbc.get_symmetry_lat(lattice, sym_type=SymmetryType.minimal)[source]#

Expands reciprocal lattice vectors to include high-symmetry directions.

This function generates a specific set of reciprocal lattice vectors by applying integer linear transformations to the primitive reciprocal basis. This is necessary to capture symmetry-equivalent vectors that may not be aligned with the principal axes.

Parameters:
  • lattice (Array) – The primitive cell lattice.

  • sym_type (SymmetryType, default: <SymmetryType.minimal: 'minimal'>) – The type of symmetry to apply.

Returns:

A new Cell object with updated reciprocal_lattice.

jaqmc.geometry.pbc.get_distance_function(distance_type)[source]#

Returns the distance function corresponding to the given DistanceType.

Parameters:

distance_type (DistanceType) – Type of periodic distance (‘nu’ or ‘tri’).

Returns:

Function that computes distance.

Raises:

ValueError – If distance_type is unknown.

jaqmc.geometry.pbc.make_pbc_gaussian_proposal(lattice)[source]#

Creates a gaussian proposal that wraps positions to the primary cell.

Parameters:

lattice (Array) – Lattice vectors with shape (ndim, ndim).

Return type:

SamplingProposal

Returns:

A sampling proposal function that respects PBC.

jaqmc.geometry.pbc.scaled_f(w)[source]#

Function f used in polynomial distance (nu_distance).

\[f(w) = |w| (1 - |w/\pi|^3 / 4)\]
Parameters:

w (Array) – the fractional coordinates scaled to the range \([-\pi, \pi]\), \(w = \mathbf{r} \cdot \mathbf{b}\), where \(\mathbf{b}\) is the reciprocal space vectors.

Return type:

Array

Returns:

Scaled value of w.

jaqmc.geometry.pbc.scaled_g(w)[source]#

Function g used in polynomial distance (nu_distance).

\[g(w) = w (1 - 1.5 |w/\pi| + 0.5 |w/\pi|^2)\]
Parameters:

w (Array) – the fractional coordinates scaled to the range \([-\pi, \pi]\), \(w = \mathbf{r} \cdot \mathbf{b}\), where \(\mathbf{b}\) is the reciprocal space vectors.

Return type:

Array

Returns:

Scaled value of w.

Spherical geometry#

jaqmc.geometry.sphere.sphere_proposal(rngs, x, stddev)[source]#

Propose MCMC moves on the sphere, operating on PyTree leaves.

Applies a rotation-based spherical move to each leaf array in x. Each leaf should have shape (..., 2) where the last axis contains (theta, phi) spherical coordinates.

Parameters:
  • rngs (PRNGKey) – Random key.

  • x – Current configuration (array or PyTree of arrays).

  • stddev (float | Array) – Gaussian width of the angular move.

Returns:

New configuration with the same structure as x.