optimize

torch_sim.runners.optimize(system, model, *, optimizer, convergence_fn=None, trajectory_reporter=None, autobatcher=False, max_steps=10_000, steps_between_swaps=5, **optimizer_kwargs)[source]

Optimize a system using a model and optimizer.

Parameters:

  • system (StateLike) – Input system to optimize (ASE Atoms, Pymatgen Structure, or SimState)

  • model (ModelInterface) – Neural network model module

  • optimizer (Callable) – Optimization algorithm function

  • convergence_fn (Callable | None) – Condition for convergence, should return a boolean tensor of length n_batches

  • optimizer_kwargs (dict) – Additional keyword arguments for optimizer init function

  • trajectory_reporter (TrajectoryReporter | dict | None) – Optional reporter for tracking optimization trajectory. If a dict, will be passed to the TrajectoryReporter constructor.

  • autobatcher (InFlightAutoBatcher | bool) – Optional autobatcher to use. If False, the system will assume infinite memory and will not batch, but will still remove converged structures from the batch. If True, the system will estimate the memory available and batch accordingly. If a InFlightAutoBatcher, the system will use the provided autobatcher, but will reset the max_attempts to max_steps // steps_between_swaps.

  • max_steps (int) – Maximum number of total optimization steps

  • steps_between_swaps (int) – Number of steps to take before checking convergence and swapping out states.

Returns:

Optimized system state

Return type:

SimState