opt

Optimize molecule geometrics in internal coordinates.

Notes

• This module is heavily based off of the optimize and step modules of geomeTRIC. See the LICENSE-3RD-PARTY for license information.

Classes:

• Step –

  The outputs of step in the optimization process.

• ConvergenceCriteria –

  The convergence criteria for the optimization process.

• Params –

  Parameters for the optimization process.

Functions:

• optimize –

  Optimize the geometry of a molecule.

Attributes:

• CONVERGENCE_CRITERIA (dict[str, ConvergenceCriteria]) –

  Default convergence criteria for optimization.

CONVERGENCE_CRITERIA module-attribute

CONVERGENCE_CRITERIA: dict[str, ConvergenceCriteria] = {
    "GAU": {
        "energy": 1e-06,
        "rms_grad": 0.0003,
        "max_grad": 0.00045,
        "rms_disp": 0.0012 * _ANGSTROM_TO_BOHR,
        "max_disp": 0.0018 * _ANGSTROM_TO_BOHR,
    },
    "GAU_LOOSE": {
        "energy": 1e-06,
        "rms_grad": 0.0017,
        "max_grad": 0.0025,
        "rms_disp": 0.0067 * _ANGSTROM_TO_BOHR,
        "max_disp": 0.01 * _ANGSTROM_TO_BOHR,
    },
    "GAU_TIGHT": {
        "energy": 1e-06,
        "rms_grad": 1e-05,
        "max_grad": 1.5e-05,
        "rms_disp": 4e-05 * _ANGSTROM_TO_BOHR,
        "max_disp": 6e-05 * _ANGSTROM_TO_BOHR,
    },
}

Default convergence criteria for optimization.

Step

Bases: NamedTuple

The outputs of step in the optimization process.

Attributes:

• coords_x (Tensor) –

  The cartesian coordinates [a0].

• grad_x (Tensor) –

  The gradients [Eh/a0] in cartesian coordinates.

• coords_q (Tensor) –

  The internal coordinates.

• grad_q (Tensor) –

  The gradients in internal coordinates.

coords_x instance-attribute

coords_x: Tensor

The cartesian coordinates [a0].

grad_x instance-attribute

grad_x: Tensor

The gradients [Eh/a0] in cartesian coordinates.

coords_q instance-attribute

coords_q: Tensor

The internal coordinates.

grad_q instance-attribute

grad_q: Tensor

The gradients in internal coordinates.

ConvergenceCriteria

Bases: TypedDict

The convergence criteria for the optimization process.

Attributes:

• energy (float) –

  The energy [Eh] convergence criteria.

• rms_grad (float) –

  The root mean square of the gradients [Eh/a0] convergence criteria.

• max_grad (float) –

  The maximum gradient [Eh/a0] convergence criteria.

• rms_disp (float) –

  The root mean square of the displacements [a0] convergence criteria.

• max_disp (float) –

  The maximum displacement [a0] convergence criteria.

energy instance-attribute

energy: float

The energy [Eh] convergence criteria.

rms_grad instance-attribute

rms_grad: float

The root mean square of the gradients [Eh/a0] convergence criteria.

max_grad instance-attribute

max_grad: float

The maximum gradient [Eh/a0] convergence criteria.

rms_disp instance-attribute

rms_disp: float

The root mean square of the displacements [a0] convergence criteria.

max_disp instance-attribute

max_disp: float

The maximum displacement [a0] convergence criteria.

Params dataclass

Parameters for the optimization process.

Attributes:

• max_steps (int) –

  The maximum number of steps to take.

• epsilon (float) –

  The minimum eigenvalue of the Hessian.

• trust (float) –

  The initial trust radius [a0]. This will be adjusted during the optimization

• trust_min (float) –

  The lower bound of the trust radius [a0].

• trust_max (float) –

  The upper bound of the trust radius [a0].

• criteria (ConvergenceCriteria) –

  The convergence criteria to use.

max_steps class-attribute instance-attribute

max_steps: int = 300

The maximum number of steps to take.

epsilon class-attribute instance-attribute

epsilon: float = 1e-05

The minimum eigenvalue of the Hessian.

trust class-attribute instance-attribute

trust: float = 0.1 * _ANGSTROM_TO_BOHR

The initial trust radius [a0]. This will be adjusted during the optimization process.

trust_min class-attribute instance-attribute

trust_min: float = 0.0012 * _ANGSTROM_TO_BOHR

The lower bound of the trust radius [a0].

trust_max class-attribute instance-attribute

trust_max: float = 0.3 * _ANGSTROM_TO_BOHR

The upper bound of the trust radius [a0].

criteria class-attribute instance-attribute

criteria: ConvergenceCriteria = field(
    default_factory=lambda: {
        None: CONVERGENCE_CRITERIA["GAU"]
    }
)

The convergence criteria to use.

optimize

optimize(
    coords_x: Tensor,
    ic: IC,
    energy_fn: EnergyFn,
    atomic_nums: Tensor,
    params: Params | None = None,
) -> tuple[list[Step], bool]

Optimize the geometry of a molecule.

Parameters:

• coords_x (Tensor) –

  The initial cartesian coordinates [a0].

• ic (IC) –

  The internal coordinate representation of the molecule.

• energy_fn (EnergyFn) –

  A function that computes the energy and gradients of the molecule. It should take the cartesian coordinates and return the energy [Eh] and gradients [Eh/a0] in atomic units.

• atomic_nums (Tensor) –

  The atomic numbers of the atoms in the molecule.

• params (Params | None, default: None ) –

  The parameters for the optimization.

Returns:

• tuple[list[Step], bool] –

  The history of the optimization process and whether the optimization converged or not.

Source code in tico/opt.py

def optimize(
    coords_x: torch.Tensor,
    ic: tico.ic.IC,
    energy_fn: EnergyFn,
    atomic_nums: torch.Tensor,
    params: Params | None = None,
) -> tuple[list[Step], bool]:
    """Optimize the geometry of a molecule.

    Args:
        coords_x: The initial cartesian coordinates [a0].
        ic: The internal coordinate representation of the molecule.
        energy_fn: A function that computes the energy and gradients of the molecule.
            It should take the cartesian coordinates and return the energy [Eh] and
            gradients [Eh/a0] in atomic units.
        atomic_nums: The atomic numbers of the atoms in the molecule.
        params: The parameters for the optimization.

    Returns:
        The history of the optimization process and whether the optimization converged
        or not.
    """
    params = params if params is not None else Params()

    trust = params.trust
    history = []

    reject_step = False

    coords_x = coords_x.double()
    coords_q = ic.compute_q(coords_x)

    hess_q = ic.guess_hess_q(coords_x, atomic_nums)
    energy, grad_x = energy_fn(coords_x)

    for _ in range(params.max_steps):
        grad_q = _compute_grad_q(coords_x, grad_x, ic)

        if not reject_step:
            history.append(Step(coords_x, grad_x, coords_q, grad_q))
            hess_q = tico.hess.update_hess_q(hess_q, history, ic)

        min_eig_val = sorted(torch.linalg.eigh(hess_q)[0])[0].real
        step_size = (
            params.epsilon - min_eig_val if min_eig_val < params.epsilon else 0.0
        )

        dq, _ = _compute_dq(step_size, coords_x, grad_q, hess_q, ic)
        coords_x_new, coords_q_new, q_converged = ic.dq_to_x(coords_x, dq)

        if tico.utils.compute_rmsd(coords_x_new, coords_x)[0] > 1.1 * trust:
            dq = _line_search(trust, step_size, coords_x, grad_q, hess_q, dq, ic)
            coords_x_new, coords_q_new, q_converged = ic.dq_to_x(coords_x, dq)

        energy_new, grad_x_new = energy_fn(coords_x_new)

        rms_grad, max_grad = _compute_grad_norm(coords_x_new, grad_x_new, ic)
        rms_disp, max_disp = tico.utils.compute_rmsd(coords_x_new, coords_x)

        if _has_converged(
            energy, energy_new, rms_grad, max_grad, rms_disp, max_disp, params
        ):
            grad_q_new = _compute_grad_q(coords_x, grad_x, ic)

            history = history + [
                Step(coords_x_new, grad_x_new, coords_q_new, grad_q_new)
            ]
            return history, True

        expected_improve = float(
            0.5 * torch.linalg.multi_dot([dq.unsqueeze(0), hess_q, dq])
            + torch.dot(dq, grad_q)
        )
        step_quality = (energy_new - energy) / expected_improve

        constr_violation = (
            isinstance(ic, tico.ic.DLC)
            and ic.constr is not None
            and torch.max(torch.abs(ic.compute_constr_delta(coords_x_new))) > 1e-1
        )

        trust, reject_step = _update_trust(
            step_quality,
            rms_disp,
            torch.abs(energy_new - energy) < params.criteria["energy"],
            trust,
            params.trust_min,
            params.trust_max,
            constr_violation,
        )

        if not reject_step:
            coords_x, coords_q = coords_x_new, coords_q_new
            energy, grad_x = energy_new, grad_x_new

    return history, False