sad_guesser

class SADGuesser(sad_coeffs_per_basis_func: Tensor, normalization_mode: SADNormalizationMode, coeff_std_per_basis_func: Tensor | None = None)[source]

Class to generate SAD (Superposition of Atomic Densities) guess for samples or batches.

__init__(sad_coeffs_per_basis_func: Tensor, normalization_mode: SADNormalizationMode, coeff_std_per_basis_func: Tensor | None = None) None[source]

Initialize the SAD guesser directly by specifying the SAD coefficients per basis function.

Parameters:
  • sad_coeffs_per_basis_func – SAD coefficients per basis function. Shape (n_basis,).

  • normalization_mode – Normalization mode for the SAD guess, see SADNormalizationMode.

  • coeff_std_per_basis_func – Standard deviation of the coefficients per basis function. Required for weighted normalization.

forward(sample: OFData | OFBatch) Tensor[source]

Generate SAD guess for samples or batches. The guess will be cast to the dtype of the coeffs in the sample.

Parameters:

sample (OFData | OFBatch) – Sample or batch to generate SAD guess for

Returns:

SAD guess. Shape (n_samples, basis).

Return type:

torch.Tensor

classmethod from_dataset_statistics(dataset_statistics: DatasetStatistics, normalization_mode: SADNormalizationMode, basis_info: BasisInfo = None, weigher_key: str = 'ground_state_only', spherical_average: bool = True) SADGuesser[source]

Initialize the SAD guesser using a DatasetStatistics.

Parameters:
  • dataset_statistics – Dataset statistics object.

  • normalization_mode – Normalization mode for the SAD guess, see SADNormalizationMode.

  • basis_info – Basis information object.

  • weigher_key – Selects the sample weigher for the dataset statistics. Defaults to “ground_state_only”.

  • spherical_average – If True, the SAD coefficients are spherically averaged.

class SADNormalizationMode(*values)[source]

Enum for the normalization mode of the SAD guesser. Recommended is PER_ATOM_WEIGHTED.

NONE

No normalization.

PER_ATOM_SIMPLE

Normalize the SAD guess per atom, i.e. scale the coefficients for each atom by a single factor to make it so the number of electrons in the coefficients belonging to the atom in the guess match its atomic number.

PER_ATOM_WEIGHTED

Normalize the SAD guess per atom, i.e. adjust the coefficients for each atom such that the number of electrons in the coefficients belonging to the atom in the guess match its atomic number. The adjustment is weighted according to the standard deviation of the coefficients, more precisely, it is proportional to the standard deviation squared times the basis integral of the coefficients.

PER_MOLECULE_SIMPLE

Normalize the SAD guess per molecule, i.e. scale the SAD guess with a single factor for each molecule such that the number of electrons in the guess matches the number of electrons in the molecule.

PER_MOLECULE_WEIGHTED

Similar to PER_ATOM_WEIGHTED, but the adjustment is normalized per molecule, not per atom.

compute_spherical_average(tensor_per_basis_func: Tensor, l_per_basis_func: Tensor) Tensor[source]

Spherically average a tensor over the basis functions.

Parameters:
  • tensor_per_basis_func – Tensor to average. Shape (n_basis, *).

  • l_per_basis_func – See l_per_basis_func in BasisInfo. Shape (n_basis,).