Grid specs

Quadrature choices for the XC integral. becke is the default atom-centered grid: Treutler-Ahlrichs radial shells at PySCF's level-3 counts (50 for H, 75 for Li-Ne, 80 for Na-Ar), 302-point Lebedev shells NWChem-pruned per radial region (prune=None for the full product grid), tail shells cut at r_max, and negligible-weight points dropped at build time (cutoff). points wraps an explicit grid. chunk controls XC streaming on either spec: "auto" (default) materializes the AO grid values only when they fit a memory budget and otherwise streams in O(chunk·nao) memory; None forces the materialized grid; an int streams with that chunk.

radial and radii_adjust choose the radial mapping and the cell sizes of the partition: Treutler-Ahlrichs' (the default, and PySCF's, which reaches a given accuracy with fewer radial shells) or Becke's (dftax's default through 0.8.1, which is becke(75, 302, radial="becke", radii_adjust="becke")). n_radial and lebedev take an int for every element or "pyscf" for PySCF's level-3 count per period. With becke(lebedev="pyscf") the grid is PySCF's default Grids, point for point; it differs from the default only from Na on, where PySCF uses 434 angular points.

screen is the large-system knob. With it, the grid is reordered into compact blocks and each block keeps only the shells that actually reach it, so the XC term costs ng·nsub² instead of ng·nao². It is off by default because the saving depends entirely on size: measured on a grad of the XC energy against the streamed dense path at def2-svp, it is a ~10% loss at 23 atoms, 1.53x at 53 and 3.11x at 153, with peak memory equal or lower and energies agreeing to 5e-12. Turn it on above ~50 atoms.

dftax.grid.becke

Radial quadratures, Bragg-Slater radii, and the fuzzy-Voronoi partition.

Two radial schemes (Becke's and Treutler-Ahlrichs' M4) and the atomic radii are coordinate-independent constants; the fuzzy-Voronoi partition is written in JAX so it is differentiable w.r.t. nuclear coordinates (needed for analytic forces).

BRAGG_BOHR module-attribute

ndarray(shape, dtype=None, buffer=None, offset=0, strides=None, order=None)

ndarray(shape, dtype=np.float64, buffer=None, offset=0, strides=None, order=None)

An array object represents a multidimensional, homogeneous array of fixed-size items. An associated data-type object describes the format of each element in the array (its byte-order, how many bytes it occupies in memory, whether it is an integer, a floating point number, or something else, etc.)

Arrays should be constructed using array, zeros or empty (refer to the See Also section below). The parameters given here refer to a low-level method (ndarray(...)) for instantiating an array.

For more information, refer to the numpy module and examine the methods and attributes of an array.

Parameters

(for the new method; see Notes below)

tuple of ints

Shape of created array.

dtype : data-type, optional Any object that can be interpreted as a numpy data type. Default is numpy.float64. buffer : object exposing buffer interface, optional Used to fill the array with data. offset : int, optional Offset of array data in buffer. strides : tuple of ints, optional Strides of data in memory. order : {'C', 'F'}, optional Row-major (C-style) or column-major (Fortran-style) order.

Attributes

T : ndarray Transpose of the array. data : buffer The array's elements, in memory. dtype : dtype object Describes the format of the elements in the array. flags : dict Dictionary containing information related to memory use, e.g., 'C_CONTIGUOUS', 'OWNDATA', 'WRITEABLE', etc. flat : numpy.flatiter object Flattened version of the array as an iterator. The iterator allows assignments, e.g., x.flat = 3 (See ndarray.flat for assignment examples; TODO). imag : ndarray Imaginary part of the array. real : ndarray Real part of the array. size : int Number of elements in the array. itemsize : int The memory use of each array element in bytes. nbytes : int The total number of bytes required to store the array data, i.e., itemsize * size. ndim : int The array's number of dimensions. shape : tuple of ints Shape of the array. strides : tuple of ints The step-size required to move from one element to the next in memory. For example, a contiguous (3, 4) array of type int16 in C-order has strides (8, 2). This implies that to move from element to element in memory requires jumps of 2 bytes. To move from row-to-row, one needs to jump 8 bytes at a time (2 * 4). ctypes : ctypes object Class containing properties of the array needed for interaction with ctypes. base : ndarray If the array is a view into another array, that array is its base (unless that array is also a view). The base array is where the array data is actually stored.

See Also

array : Construct an array. zeros : Create an array, each element of which is zero. empty : Create an array, but leave its allocated memory unchanged (i.e., it contains "garbage"). dtype : Create a data-type. numpy.typing.NDArray : An ndarray alias :term:generic <generic type> w.r.t. its dtype.type <numpy.dtype.type>.

Notes

There are two modes of creating an array using __new__:

1. If buffer is None, then only shape, dtype, and order are used.
2. If buffer is an object exposing the buffer interface, then all keywords are interpreted.

No __init__ method is needed because the array is fully initialized after the __new__ method.

Examples

These examples illustrate the low-level ndarray constructor. Refer to the See Also section above for easier ways of constructing an ndarray.

First mode, buffer is None:

import numpy as np np.ndarray(shape=(2,2), dtype=np.float64, order='F') array([[0.0e+000, 0.0e+000], # random [ nan, 2.5e-323]])

Second mode:

np.ndarray((2,), buffer=np.array([1,2,3]), ... offset=np.int_().itemsize, ... dtype=np.int_) # offset = 1*itemsize, i.e. skip first element array([2, 3])

TREUTLER_XI module-attribute

ndarray(shape, dtype=None, buffer=None, offset=0, strides=None, order=None)

ndarray(shape, dtype=np.float64, buffer=None, offset=0, strides=None, order=None)

An array object represents a multidimensional, homogeneous array of fixed-size items. An associated data-type object describes the format of each element in the array (its byte-order, how many bytes it occupies in memory, whether it is an integer, a floating point number, or something else, etc.)

Arrays should be constructed using array, zeros or empty (refer to the See Also section below). The parameters given here refer to a low-level method (ndarray(...)) for instantiating an array.

For more information, refer to the numpy module and examine the methods and attributes of an array.

Parameters

(for the new method; see Notes below)

tuple of ints

Shape of created array.

dtype : data-type, optional Any object that can be interpreted as a numpy data type. Default is numpy.float64. buffer : object exposing buffer interface, optional Used to fill the array with data. offset : int, optional Offset of array data in buffer. strides : tuple of ints, optional Strides of data in memory. order : {'C', 'F'}, optional Row-major (C-style) or column-major (Fortran-style) order.

Attributes

T : ndarray Transpose of the array. data : buffer The array's elements, in memory. dtype : dtype object Describes the format of the elements in the array. flags : dict Dictionary containing information related to memory use, e.g., 'C_CONTIGUOUS', 'OWNDATA', 'WRITEABLE', etc. flat : numpy.flatiter object Flattened version of the array as an iterator. The iterator allows assignments, e.g., x.flat = 3 (See ndarray.flat for assignment examples; TODO). imag : ndarray Imaginary part of the array. real : ndarray Real part of the array. size : int Number of elements in the array. itemsize : int The memory use of each array element in bytes. nbytes : int The total number of bytes required to store the array data, i.e., itemsize * size. ndim : int The array's number of dimensions. shape : tuple of ints Shape of the array. strides : tuple of ints The step-size required to move from one element to the next in memory. For example, a contiguous (3, 4) array of type int16 in C-order has strides (8, 2). This implies that to move from element to element in memory requires jumps of 2 bytes. To move from row-to-row, one needs to jump 8 bytes at a time (2 * 4). ctypes : ctypes object Class containing properties of the array needed for interaction with ctypes. base : ndarray If the array is a view into another array, that array is its base (unless that array is also a view). The base array is where the array data is actually stored.

See Also

array : Construct an array. zeros : Create an array, each element of which is zero. empty : Create an array, but leave its allocated memory unchanged (i.e., it contains "garbage"). dtype : Create a data-type. numpy.typing.NDArray : An ndarray alias :term:generic <generic type> w.r.t. its dtype.type <numpy.dtype.type>.

Notes

There are two modes of creating an array using __new__:

1. If buffer is None, then only shape, dtype, and order are used.
2. If buffer is an object exposing the buffer interface, then all keywords are interpreted.

No __init__ method is needed because the array is fully initialized after the __new__ method.

Examples

These examples illustrate the low-level ndarray constructor. Refer to the See Also section above for easier ways of constructing an ndarray.

First mode, buffer is None:

import numpy as np np.ndarray(shape=(2,2), dtype=np.float64, order='F') array([[0.0e+000, 0.0e+000], # random [ nan, 2.5e-323]])

Second mode:

np.ndarray((2,), buffer=np.array([1,2,3]), ... offset=np.int_().itemsize, ... dtype=np.int_) # offset = 1*itemsize, i.e. skip first element array([2, 3])

__cached__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__doc__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__file__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__name__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__package__ module-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

annotations module-attribute

bragg_radius(Z: int) -> float

Bragg-Slater radius (Bohr) for atomic number Z.

treutler_xi(Z: int) -> float

Treutler-Ahlrichs radial size parameter (Bohr) for atomic number Z.

becke_radial(n: int, scale: float) -> tuple[np.ndarray, np.ndarray]

Becke radial quadrature for ∫₀^∞ f(r) r² dr ≈ Σ w_i f(r_i).

Second-kind Gauss-Chebyshev nodes mapped by r = R (1+x)/(1-x).

Parameters:

Name	Type	Description	Default
n	int	number of radial points.	required
scale	float	radial size parameter R (Bohr).	required

treutler_radial(n: int, xi: float) -> tuple[np.ndarray, np.ndarray]

Treutler-Ahlrichs M4 radial quadrature for ∫₀^∞ f(r) r² dr.

The same second-kind Gauss-Chebyshev nodes as :func:becke_radial, mapped by r = (xi/ln 2) (1+x)^0.6 ln(2/(1-x)) (mapping M4 of JCP 102, 346 (1995), with alpha = 0.6). It puts fewer points in the far tail than Becke's mapping, so a coarse grid spends more of them where the density is.

Parameters:

Name	Type	Description	Default
n	int	number of radial points.	required
xi	float	radial size parameter (Bohr), :func:treutler_xi per element.	required

size_adjustment(bragg, radii_adjust: str | None)

Heteronuclear cell-size parameters a_AB (antisymmetric, natom²) from the atoms' Bragg radii bragg (natom,).

Becke 1988, App. A: a = u / (u² - 1) with u = (chi-1)/(chi+1). "becke" takes chi as the Bragg-radius ratio and bounds |a| by 0.45; "treutler" takes the square root of that ratio (JCP 102, 346) and bounds |a| by 1/2, as PySCF does; None leaves the cells unadjusted. In jnp, so under jit the natom² matrix is computed from the natom radii rather than embedded as a constant.

becke_partition(points, coords, Zs, radii_adjust: str | None = 'becke')

Unnormalized Becke cell functions P_A(point), shape (n_grid, n_atom).

Vectorized JAX (differentiable w.r.t. coords); includes the heteronuclear size adjustment selected by radii_adjust (see :func:size_adjustment). Points are processed in budget-derived chunks so the per-point (chunk, A, B) transient stays bounded for large molecules (checkpointed, so the backward pass used by the forces rematerializes each chunk instead of storing them all).

becke_weights(points, owners, coords, Zs, radii_adjust: str | None = 'becke')

Becke weight factors P_owner(p) / Σ_A P_A(p) of points p owned by atoms owners (int, (n_grid,)), shape (n_grid,).

One chunked pass over every atom's points with the pair quantities formed once, so a molecule's grid is one loop in the compiled program (a pass per atom was natom loops, each with its own natom² constants) and the (n_grid, n_atom) cell functions are never held. Differentiable w.r.t. coords, checkpointed like :func:becke_partition.

dftax.grid.points(coords, weights, *, chunk: int | str | None = 'auto') -> Points

An explicit quadrature grid (coords, weights) (e.g. from PySCF).

chunk streams the XC integral over grid chunks (see :func:becke); "auto" picks materialized vs streamed by memory budget.

dftax.grid.becke_grid(symbols: list[str], coords_bohr, n_radial: int | str = 'pyscf', lebedev: int | str = 302, prune: str | None = 'nwchem', r_max: float | None = 45.0, radial: str = 'treutler', radii_adjust: str | None = 'treutler')

Build a molecular quadrature grid (points in Bohr, weights).

Each atom contributes n_radial radial shells (mapped per radial) times a Lebedev angular grid whose order is pruned per radial region (see :func:becke); the per-atom grids are merged with Becke's fuzzy-Voronoi partition, its cells sized per radii_adjust.

Returns (coords, weights) as JAX arrays of shape (n_grid, 3) and (n_grid,), differentiable w.r.t. coords_bohr. Shapes depend only on symbols and the spec constants, so this can run under jit / vmap with traced coordinates.

Jitted, one compiled program per (symbols, spec) shape for every geometry. The partition is one chunked pass over all atoms' points (:func:~dftax.grid.becke.becke_weights), and each element's point template is one constant, so the program and its constants grow with the number of elements, not of atoms.

dftax.grid.becke_grid_size(symbols: list[str], n_radial: int | str = 'pyscf', lebedev: int | str = 302, prune: str | None = 'nwchem', r_max: float | None = 45.0, radial: str = 'treutler') -> int

Number of grid points :func:becke_grid will emit (static per spec).

Lets callers that build the grid under jit/vmap (forces, the batched solver) make eager, shape-dependent decisions, e.g. resolving the chunk="auto" streaming policy.