Batched evaluation¤
One vmapped solve over a stack of geometries. Orbital-sized outputs are
opt-in; mesh= shards the batch axis across devices.
dftax.ks.batched.scf_batched(mol, coords_batch, xc, *, spin: int | None = None, grid: Becke | None = None, coulomb: ExactSpec | DFSpec | None = None, dispersion=None, forces: bool = False, return_orbitals: bool = False, mesh: MeshSpec | None = None, max_iter: int = 128, e_tol: float = 1e-08, d_tol: float = 1e-06, diis_space: int = 8, level_shift: float = 0.0, guess=None) -> BatchedResult
¤
KS SCF over a batch of geometries coords_batch of shape
(B, n_atom, 3) (Bohr). mol supplies the (fixed) atoms + basis; only
the coordinates vary. spin follows the :class:~dftax.ks.energy.KS
rule (None → infer from mol; explicit → polarized channels).
With forces=True the result carries the analytic force tensor
(B, n_atom, 3); with return_orbitals=True it also carries
P/mo_coeff/mo_energy (O(B·nspin·nao²) memory, off by
default). mesh= shards the batch axis across a device mesh.
guess is a spec from :mod:dftax.ks.guess (core/sad/minao/
sap); it is resolved once (SAD/MinAO atomic densities are
geometry-independent) and applied per geometry inside the solve.
coulomb follows the KS default (density fitting): the auxiliary basis
is built once as a template and re-centered per geometry. Value-based
knobs are rejected here (Schwarz screen needs concrete integrals, but
the per-geometry build is traced), and batched forces=True with a
streamed hybrid RI-K is rejected (its vjp has no geometry gradients).
dftax.ks.batched.BatchedResult
¤
Per-geometry results of a batched KS solve (leading axis = batch B).
A distinct type from the single-solve :class:~dftax.ks.scf.KSResult, so
consumers written against per-geometry scalars (forces,
if res.converged) reject a batch loudly instead of mis-slicing it.
Orbital fields are None unless the run asked for them
(return_orbitals=True).
__annotations__
class-attribute
¤
dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)
__dataclass_fields__
class-attribute
¤
dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)
__dataclass_params__
class-attribute
¤
__dict__
class-attribute
¤
Read-only proxy of a mapping.
__doc__
class-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'.
__firstlineno__
class-attribute
¤
int([x]) -> integer int(x, base=10) -> integer
Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.
If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.
int('0b100', base=0) 4
__match_args__
class-attribute
¤
Built-in immutable sequence.
If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.
If the argument is a tuple, the return value is the same object.
__module__
class-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'.
__static_attributes__
class-attribute
¤
Built-in immutable sequence.
If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.
If the argument is a tuple, the return value is the same object.
__weakref__
property
¤
list of weak references to the object
__eq__(other)
¤
__init__(e_tot: Float[Array, B], e_elec: Float[Array, B], converged: Array, n_iter: Array, nocc: tuple[int, ...], forces: Float[Array, 'B n_atom 3'] | None = None, mo_energy: Float[Array, 'B nspin nmo'] | None = None, mo_coeff: Float[Array, 'B nspin nao nmo'] | None = None, P: Float[Array, 'B nspin nao nao'] | None = None) -> None
¤
__repr__()
¤
Batched direct minimization¤
Different molecules padded to common sizes and minimized in one compiled
program, closed and open shells, pure and hybrid. A system left unconverged
or in a non-aufbau state is re-solved alone by gdm
(BatchedGDMResult.fallback).
dftax.ks.batched_gdm.gdm_batched(kss: Sequence[KS], *, max_iter: int = 200, e_tol: float = 1e-08, d_tol: float = 1e-06, history: int = 20, max_step: float = 0.5, gap_floor: float = 0.05, guess: GuessSpec | jax.Array | None = None, bucket: Bucket | None = None, fallback: bool = True) -> BatchedGDMResult
¤
Minimize the KS energy of every system in kss at once. The
tolerances and step controls are those of :func:~dftax.ks.gdm.gdm.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
kss
|
Sequence[KS]
|
built :class: |
required |
bucket
|
Bucket | None
|
the padded sizes, for a batch of one spin type; default the smallest that fits each group. Passing a fixed bucket lets batches of similar molecules share one compiled program. |
None
|
guess
|
GuessSpec | Array | None
|
a guess spec as for :func: |
None
|
fallback
|
bool
|
re-solve one at a time, with :func: |
True
|
dftax.ks.batched_gdm.BatchedGDMResult
¤
Per-system results of :func:gdm_batched, in input order.
fallback marks the systems the batch did not settle (not converged,
or converged to a state whose occupied orbitals are not the lowest ones)
and that were re-solved one at a time by :func:~dftax.ks.gdm.gdm; their
n_iter is that solve's.
__annotations__
class-attribute
¤
dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)
__dataclass_fields__
class-attribute
¤
dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)
__dataclass_params__
class-attribute
¤
__dict__
class-attribute
¤
Read-only proxy of a mapping.
__doc__
class-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'.
__firstlineno__
class-attribute
¤
int([x]) -> integer int(x, base=10) -> integer
Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.
If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.
int('0b100', base=0) 4
__match_args__
class-attribute
¤
Built-in immutable sequence.
If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.
If the argument is a tuple, the return value is the same object.
__module__
class-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'.
__static_attributes__
class-attribute
¤
Built-in immutable sequence.
If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.
If the argument is a tuple, the return value is the same object.
__weakref__
property
¤
list of weak references to the object