JAX-MD Simulation Config

class mlip.simulation.configs.jax_md_config.JaxMDSimulationConfig(*, simulation_type: ~mlip.simulation.enums.SimulationType = SimulationType.MD, num_steps: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)], snapshot_interval: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 1, box: ~typing.Annotated[float, FieldInfo(annotation=NoneType, required=True, metadata=[Gt(gt=0)])] | ~typing.Annotated[list[~typing.Annotated[float, FieldInfo(annotation=NoneType, required=True, metadata=[Gt(gt=0)])]], FieldInfo(annotation=NoneType, required=True, metadata=[MinLen(min_length=3), MaxLen(max_length=3)])] | None = None, edge_capacity_multiplier: ~typing.Annotated[float, ~annotated_types.Ge(ge=1)] = 1.25, num_episodes: ~typing.Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Gt(gt=0)])] | None = None, timestep_fs: ~typing.Annotated[float, FieldInfo(annotation=NoneType, required=True, metadata=[Gt(gt=0)])] | None = 1.0, temperature_kelvin: ~typing.Annotated[float, FieldInfo(annotation=NoneType, required=True, metadata=[Gt(gt=0)])] | None = 300.0, temperature_schedule_config: ~mlip.simulation.configs.simulation_config.TemperatureScheduleConfig = TemperatureScheduleConfig(method=<TemperatureScheduleMethod.CONSTANT: 'constant'>, temperature=300.0, start_temperature=None, end_temperature=None, max_temperature=None, min_temperature=None, heating_period=None))

Configuration for the JAX-MD-based simulations.

Also includes the attributes of the parent class SimulationConfig.

The config is separated into three blocks: values that are used for both MD and minimization, and then the ones used exclusively for MD and for minimization, respectively.

num_episodes

Number of episodes to divide the simulation into. Each episode runs in a fully jitted way, and the loggers are only called after each episode. If not set, an appropriate value will be attempted to be select but it is possible that it may have to be manually set. For fewer than 1000 steps, num_episodes will be set so that the number of steps per episode will be 10. For more than 1000 steps, num_episodes will be set so that the number of steps per episode will be 1000. Therefore, if num_episodes is not set, it requires that num_steps be divisible by 1000 if greater than 1000 otherwise divisible by 10.

Type:

int | None

timestep_fs

The simulation timestep in femtoseconds. This is also used as the initial timestep in the FIRE minimization algorithm. The default is 1.0.

Type:

float | None

temperature_kelvin

The temperature in Kelvin, set to 300 by default. Must be set to None for energy minimizations.

Type:

float | None

temperature_schedule_config

The temperature schedule config to use for the simulation. Default is the constant schedule in which case temperature_kelvin will be applied.

Type:

mlip.simulation.configs.simulation_config.TemperatureScheduleConfig