density_optimization

compute_density_optimization_metrics(callback: ConvergenceCallback, ground_state_energy: Energies, convergence_criterion: str | ConvergenceCriterion) → dict

Compute metrics for the density optimization process.

density_optimization(sample: OFData, mol: Mole, optimizer: Optimizer, func_factory: FunctionalFactory, callback: ConvergenceCallback | None = None, initialization: str | Tensor = 'minao', max_xc_memory: int = 4000, normalize_initial_guess: bool = True, ks_basis: str = '6-31G(2df,p)', proj_minao_module: MLDFTLitModule | None = None, sad_guess_kwargs: dict | None = None, disable_printing: bool = False, disable_pbar: bool = None) → tuple[Energies, Tensor, bool, Callable]

Perform density optimization for a given sample and molecule.

Parameters:

• sample – OFData sample object containing required tensors for the functional.

• mol – Molecule object used for the initial guess and building the grid.

• optimizer – Optimizer used for the density optimization process.

• func_factory – FunctionalFactory object used to construct the energy functional.

• callback – ConvergenceCallback object to store the optimization history.

• initialization – Which initial guess to use.

• max_xc_memory – Maximum memory for the XC functional.

• normalize_initial_guess – Whether to normalize the initial guess to the correct number of electrons.

• ks_basis – Basis set to use for the initial guess.

• proj_minao_module – Lightning Module used for the proj_minao initial guess. Only required if initialization is ‘proj_minao’.

• sad_guess_kwargs – Dictionary of arguments for the SAD guesser. Only required if initialization is ‘sad’ or ‘sad_transformed’.

• disable_printing – Whether to disable printing (useful when running density optimization during training with the rich progress bar, as output becomes messy).

• disable_pbar – Whether to disable the progress bar. If None, it will use the value of disable_printing.

Returns:

Energies object containing the total energy and its components. torch.Tensor: Final coefficients in the AO basis. bool: Whether the optimization converged.

Return type:

Energies

density_optimization_with_label(sample: OFData, mol: Mole, optimizer: Optimizer, func_factory: FunctionalFactory, callback: ConvergenceCallback | None, max_xc_memory: int = 4000, initial_guess_str: str = 'minao', normalize_initial_guess: bool = True, ks_basis: str = '6-31G(2df,p)', proj_minao_module: MLDFTLitModule | None = None, sad_guess_kwargs: dict | None = None, convergence_criterion: str | ConvergenceCriterion = 'last_iter', disable_printing: bool = False) → tuple[dict, ConvergenceCallback, Energies, Callable]

Perform density optimization for a given sample and molecule.

Parameters:

• sample – OFData sample object containing required tensors for the functional.

• mol – Molecule object used for the initial guess and building the grid.

• optimizer – Optimizer used for the density optimization process.

• initial_guess_str – Which initial guess to use.

• normalize_initial_guess – Whether to normalize the initial guess to the correct number of electrons.

• ks_basis – Basis set to use for the initial guess.

• proj_minao_module – Lightning Module used for the proj_minao initial guess. Only required if initial_guess_str is ‘proj_minao’.

• sad_guess_kwargs – Dictionary of arguments for the SAD guesser. Only required if initial_guess_str is ‘sad’ or ‘sad_transformed’.

• convergence_criterion – Criterion to determine the ground state of the optimization process.

Returns:

Dictionary containing metrics of the density optimization. callback: ConvergenceCallback object containing the optimization history for plotting and further analysis. energies_label: The ground state label of the molecule. energy_functional: The constructed energy functional.

Return type:

metric_dict

initial_guess(sample: OFData, mol: Mole, initial_guess_str: str, normalize_initial_guess: bool = True, ks_basis: str = '6-31G(2df,p)', proj_minao_module: MLDFTLitModule | None = None, sad_guess_kwargs: dict | None = None) → Tensor

Return the initial guess for the coefficients.

Parameters:

• sample – The OFData object.

• mol – The molecule.

Returns:

The initial guess for the coefficients.

Return type:

torch.Tensor

Raises:

KeyError – If the initial guess is not in INITIAL_GUESS (e.g. implemented guess).