FEP Sampler Config¶

Configuration for FEP simulations managed by FEPSimulationSampler.

simulation_config¶

Config of the simulation to run for each lambda.

Type:: mlip.simulation.configs.jax_md_config.JaxMDSimulationConfig

alchemical_atom_indices¶

Indices of the atoms to treat as alchemical. Edges between the alchemical atoms and all other atoms are gradually switched off between adjacent simulations.

Type:: jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex | list[int]

lambda_edge_values¶

Alchemical edge weight to use for each simulation (1.0 = fully connected, 0.0 = fully decoupled). Must be the same length as lambda_repulsion_values, one entry per lambda window.

Type:: jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex | list[float]

lambda_repulsion_values¶

Repulsive potential weight between alchemical and non-alchemical atoms to use for each simulation. Must be the same length as lambda_edge_values, one entry per lambda window.

Type:: jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex | list[float]

use_alchemical_mlip¶

If True (default), uses an alchemical equivalent of the input force field, which weights alchemical edges in the message-passing equations. If False, uses the base model equations, and computes the alchemical potential by summing predictions on the two endstates.

Type:: bool

alchemical_energy_batch_size¶

Only used when use_alchemical_mlip is True. Controls the maximum batch size used when computing the per-lambda alchemical energy of each snapshot. Required by some models to prevent OOM errors. If None (default), computes all in a single batch.

Type:: int | None

repulsive_potential¶

Softcore repulsive potential applied to alchemical edges. Defaults to SoftcoreLennardJonesPotential().

Type:: mlip.simulation.fep.models.repulsive_potentials.SoftcoreRepulsivePotential

use_replica_exchange¶

Whether to use Hamiltonian Replica Exchange to swap adjacent lambdas between episodes, according to a Metropolis Criterion.

Type:: bool

num_equilibration_episodes¶

Number of episodes before Replica Exchange begins. Ignored if use_replica_exchange is False.

Type:: int

checkpoint_dir¶

Optional local directory, If set, a checkpoint is saved to this directory every checkpoint_interval_episodes episodes, enabling restoration. To also upload to remote storage, pass a checkpoint_dir_upload_fun to the sampler’s init. Note that to restore from a checkpoint, sampler.restore_checkpoint(...) must be used with a local directory before sampler.run(). If the checkpoint directory is in remote storage, it must be downloaded first.

Type:: pathlib.Path | None

checkpoint_interval_episodes¶

Number of episodes between checkpoint saves. Can be used to reduce upload frequency when a checkpoint_dir_upload_fun is provided (see checkpoint_dir above). Default is 1.

Type:: int

class mlip.simulation.fep.models.SoftcoreLennardJonesPotential(deterministic_scatter_ops: bool = False, sigma_map: dict[int, float]=<factory>, epsilon_map: dict[int, float]=<factory>)¶

Softcore Lennard-Jones (SCLJ) repulsive potential.

The pairwise SCLJ potential is:

\[U_{\mathrm{SCLJ}}(r, \lambda) = 4 \epsilon \lambda \left(D(r, \lambda)^{-2} - D(r, \lambda)^{-1}\right)\]

where \(\lambda\) is the repulsion_weight and:

\[D(r, \lambda) = 0.5 (1 - \lambda) + (r / \sigma)^6\]

The strength of repulsion increases with \(\lambda\) from 0 to 1, such that \(U_{\mathrm{SCLJ}}(r, 0) = 0\) and \(U_{\mathrm{SCLJ}}(r, 1)\) is the full Lennard-Jones potential.

deterministic_scatter_ops¶

Whether to use deterministic scatter operations. If True, uses a deterministic but slower alternative to scatter operations. Does not need to be set by the user; will be automatically set to True when needed inside a simulation.

Type:: bool

sigma_map¶

Dict specifying a sigma value (Angstrom) for each atomic number. The per-edge value is inferred using Lorentz-Berthelot mixing rules. Defaults to the maximum per-element value used by the GAFF-2.1 FF.

Type:: dict[int, float]

epsilon_map¶

Dict specifying an epsilon value (kJ/mol) for each atomic number. The per-edge value is inferred using Lorentz-Berthelot mixing rules. Defaults to a representative per-element value used by the GAFF-2.1 FF.

Type:: dict[int, float]

class mlip.simulation.fep.models.SoftcoreWCAPotential(deterministic_scatter_ops: bool = False, sigma_map: dict[int, float]=<factory>, epsilon_map: dict[int, float]=<factory>)¶

Softcore Weeks-Chandler-Andersen (SCWCA) repulsive potential.

Repulsion-only equivalent of the softcore Lennard-Jones (SCLJ) potential; shifted so that \(U = 0\) at \(\sigma\) and truncated to remove the attractive region (\(D(r, \lambda) > 2\)).