ensemble.sampling.initialize_simulator_template

ensemble.sampling.initialize_simulator_template#

initialize_simulator_template(init_simulator_fn, shift_fn, nbrs=None, init_with_PRNGKey=True, extra_simulator_kwargs=None)[source]#

Initializes the simulator template and reference state.

Parameters:

  • init_simulator_fn – Function returning a (init_fn, apply_fn) tuple when provided with a potential energy function.

  • shift_fn (Callable) – Function shifting positions back into the simulation box after a simulation update step.

  • nbrs (NeighborList) – Neighbor list to allocate new neighbor list based on particle positions.

  • init_with_PRNGKey (bool) – Whether simulator init function takes an PRNGKey, should be set to False e.g. for the Gradient Descend energy minimization routine.

  • extra_simulator_kwargs (Mapping) – Additional arguments when initializing the simulator state.

Returns:

Returns a function to initialize the simulator state and the corresponding simulator template.

Contents