grid

Grid operations for crystalline systems.

This module provides functions for working with real and reciprocal space grids in crystalline systems. It includes utilities for:

• Generating G-vectors and R-vectors

• \(k\)-point sampling for Brillouin zone integration

• Frequency space operations and masks

• Grid transformations between real and reciprocal space

jrystal.grid.cubic_mask(grid_sizes: List | Array) → Bool[Array, 'x y z']

Create a cubic mask for frequency components in reciprocal space.

Generates a mask that enables only certain frequency components in a cubic region of reciprocal space. This is particularly useful when dealing with electron density (\(\rho\)), which contains more frequency components than the wavefunction (\(\psi\)) since \(\rho = |\psi|^2\).

The mask is constructed by considering the frequency mixing that occurs when squaring the wavefunction, ensuring all relevant frequency components are included.

Parameters:

grid_sizes (Union[List, jax.Array]) – Grid dimensions in reciprocal space.

Returns:

A boolean array of shape (*grid_sizes) where True indicates allowed frequency components in the cubic mask.

Return type:

Bool[Array, ‘x y z’]

jrystal.grid.estimate_max_cutoff_energy(cell_vectors: Float[Array, '3 3'], mask: Bool[Array, 'x y z']) → float

Estimate the maximum cutoff energy corresponding to a frequency mask.

Given a boolean mask in reciprocal space, calculates the maximum kinetic energy of the G-vectors that are included in the mask. This is useful for determining the effective energy cutoff of a given frequency mask, particularly when using non-spherical masks.

Parameters:

• cell_vectors (Float[Array, '3 3']) – Real space lattice vectors of the unit cell. A (3,3) matrix where each row is a lattice vector.

• mask (Bool[Array, 'x y z']) – Boolean mask indicating which G-vectors are included.

Returns:

The maximum kinetic energy (in the same units as the reciprocal lattice vectors) of any G-vector included in the mask.

Return type:

float

jrystal.grid.g2cell_vectors(g_vector_grid: Float[Array, 'x y z 3']) → Float[Array, '3 3']

Compute real space cell vectors from a G-vector grid.

Determines the real space lattice vectors by solving a linear system that relates the G-vector grid to the standard reciprocal space basis. This is useful when only the G-vectors are known and the corresponding real space cell vectors are needed.

Parameters:

g_vector_grid (Float[Array, 'x y z 3']) – Grid of G-vectors in reciprocal space.

Returns:

Real space lattice vectors as a (3,3) matrix where each row is a lattice vector.

Return type:

Float[Array, ‘3 3’]

jrystal.grid.g2r_vector_grid(g_vector_grid: Float[Array, 'x y z 3'], cell_vectors: Float[Array, '3 3'] | None = None) → Float[Array, 'x y z 3']

Transform a G-vector grid to the corresponding R-vector grid.

Converts a grid of vectors in reciprocal space (G-vectors) to the corresponding grid in real space (R-vectors). If cell vectors are not provided, they are computed from the G-vector grid.

This transformation is useful when switching between reciprocal and real space representations of crystal quantities.

Parameters:

• g_vector_grid (Float[Array, 'x y z 3']) – Grid of G-vectors in reciprocal space.

• cell_vectors (Optional[Float[Array, '3 3']], optional) – Real space lattice vectors. If None, they will be computed from the G-vector grid. Defaults to None.

Returns:

Grid of R-vectors in real space with the same shape as the input.

Return type:

Float[Array, ‘x y z 3’]

jrystal.grid.g_vectors(cell_vectors: Float[Array, '3 3'], grid_sizes: Tuple | List | Int[Array, 'd']) → Float[Array, 'x y z 3']

Generate G-vectors (reciprocal space vectors) for a given crystal cell.

Given the real space lattice vectors of a unit cell, computes the G-vectors in reciprocal space on a discrete grid. The G-vectors are fundamental for plane-wave calculations and Fourier transforms in periodic systems.

The G-vectors are defined as:

\[G_{ijk} = i\mathbf{b}_1 + j\mathbf{b}_2 + k\mathbf{b}_3\]

where \(\mathbf{b}_i\) are the reciprocal lattice vectors and \(i,j,k \in [0, n_i-1]\) for grid sizes \(n_i\).

Parameters:

• cell_vectors (Float[Array, '3 3']) – Real space lattice vectors of the unit cell. A (3,3) matrix where each row is a lattice vector.

• grid_sizes (Union[Tuple, List, Int[Array, 'd']]) – Number of grid points along each axis.

Returns:

A tensor with shape (*grid_sizes, 3) containing the G-vectors.

Return type:

Float[Array, ‘x y z 3’]

jrystal.grid.grid_vector_radius(grid_vector: Float[Array, 'x y z 3'])

Calculate the magnitude (radius) of vectors at each grid point.

Computes the Euclidean norm of vectors at each point in a grid. The function is vectorized to efficiently handle arbitrary grid shapes.

Parameters:

grid_vector (Float[Array, 'x y z 3']) – Array of vectors where the last dimension contains the vector components, and earlier dimensions are grid dimensions or batch dimensions.

Returns:

Array of vector magnitudes with shape matching all but the last dimension of the input.

Return type:

Float[Array, ‘x y z’]

jrystal.grid.k_vectors(cell_vectors: Float[Array, '3 3'], grid_sizes: Tuple | List | Int[Array, '3']) → Float[Array, 'kpt 3']

Generate k-vectors for Brillouin zone sampling.

Creates a uniform grid of \(k\)-points in reciprocal space using the Monkhorst-Pack scheme. This sampling is essential for integrating periodic functions over the Brillouin zone in electronic structure calculations.

Warning

This function is not differentiable as it uses monkhorst_pack from ase. In future, we will implement a custom differentiable version.

Parameters:

• cell_vectors (Float[Array, '3 3']) – Real space lattice vectors of the unit cell. A (3,3) matrix where each row is a lattice vector.

• grid_sizes (Union[Tuple, List, Int[Array, '3']]) – Number of \(k\)-points along each reciprocal lattice vector direction.

Returns:

An array of shape (\(n\), 3) containing k-vectors, where \(n\) is the total number of \(k\)-points (product of grid_sizes).

Return type:

Float[Array, ‘kpt 3’]

jrystal.grid.proper_grid_size(grid_sizes: Int | Int[Array, 'd'] | Tuple | List) → Array

Optimize grid sizes for efficient FFT operations.

Converts input grid dimensions to values that are well-suited for FFT computations by factoring them into products of small primes (2, 3, 5, 7). This optimization can significantly improve FFT performance.

Parameters:

grid_sizes (Union[Int, Int[Array, 'd'], Tuple, List]) – Input grid dimensions. Can be: - A single integer (same size for all dimensions) - A sequence of integers (size for each dimension) - A numpy array of integers

Returns:

A numpy array containing the optimized grid sizes.

Return type:

Array

Raises:

TypeError – If grid_sizes is not a valid numeric type.

jrystal.grid.r2g_vector_grid(r_vector_grid: Float[Array, 'x y z 3'], cell_vectors: Float[Array, '3 3'] | None = None) → Float[Array, 'x y z 3']

Transform an R-vector grid to the corresponding G-vector grid.

Converts a grid of vectors in real space (R-vectors) to the corresponding grid in reciprocal space (G-vectors). This is the inverse operation of g2r_vector_grid.

Parameters:

• r_vector_grid (Float[Array, 'x y z 3']) – Grid of R-vectors in real space.

• cell_vectors (Float[Array, '3 3']) – Real space lattice vectors.

Returns:

Grid of G-vectors in reciprocal space with the same shape as the input.

Return type:

Float[Array, ‘x y z 3’]

jrystal.grid.r_vectors(cell_vectors: Float[Array, '3 3'], grid_sizes: Tuple | List | Int[Array, '3']) → Float[Array, 'x y z 3']

Generate R-vectors (real space position vectors) for a given crystal cell.

Given the real space lattice vectors of a unit cell, computes the position vectors R on a discrete grid within the unit cell. These vectors define the sampling points where real-space quantities (like electron density) are evaluated.

The R-vectors are defined as:

\[R_{ijk} = \frac{i}{n_x}\mathbf{a}_1 + \frac{j}{n_y}\mathbf{a}_2 + \frac{k}{n_z}\mathbf{a}_3\]

where \(\mathbf{a}_i\) are the real space lattice vectors and \(i,j,k \in [0, n_i-1]\) for grid sizes \(n_i\).

Parameters:

• cell_vectors (Float[Array, '3 3']) – Real space lattice vectors of the unit cell. A (3,3) matrix where each row is a lattice vector.

• grid_sizes (Union[Tuple, List, Int[Array, '3']]) – Number of grid points along each axis.

Returns:

A tensor with shape (*grid_sizes, 3) containing the R-vectors.

Return type:

Float[Array, ‘x y z 3’]

jrystal.grid.spherical_mask(cell_vectors: Float[Array, '3 3'], grid_sizes: List | Array, cutoff_energy: float) → Bool[Array, 'x y z']

Create a spherical mask for frequency cutoff in reciprocal space.

Generates a boolean mask that selects G-vectors satisfying the energy cutoff condition:

\[\frac{\|G\|^2}{2} \leq E_\text{cutoff}\]

This mask is commonly used in plane-wave calculations to limit the basis set size while maintaining accuracy. G-vectors with kinetic energy above the cutoff are excluded.

Parameters:

• cell_vectors (Float[Array, '3 3']) – Real space lattice vectors of the unit cell. A (3,3) matrix where each row is a lattice vector.

• grid_sizes (Union[List, jax.Array]) – Grid dimensions in reciprocal space.

• cutoff_energy – Energy cutoff for the G-vectors.

Returns:

A boolean array of shape (*grid_sizes) where True indicates G-vectors within the energy cutoff sphere.

Return type:

Bool[Array, ‘x y z’]

jrystal.grid.translation_vectors(cell_vectors: Float[Array, '3 3'], cutoff: Float[Array, '3'] | float = 10000.0) → Float[Array, 'num 3']

Generate translation vectors for Ewald summation.

Creates a grid of translation vectors used in Ewald summation methods for computing long-range interactions in periodic systems. The grid extends to a distance determined by the cutoff parameter.

Parameters:

• cell_vectors (Float[Array, '3 3']) – Real space lattice vectors of the unit cell. A (3,3) matrix where each row is a lattice vector.

• cutoff (Union[Float[Array, '3'], float]) – Real space cutoff distance. Larger values give more precise Ewald summation results but increase computational cost. Default is 1e4.

Returns:

An array of shape (\(n\), 3) containing translation vectors, where \(n\) is determined by the cutoff distance.

Return type:

Float[Array, ‘num 3’]