jax_md_mod.custom_quantity.energy_wrapper

Contents

jax_md_mod.custom_quantity.energy_wrapper#

energy_wrapper(energy_fn_template, fixed_energy_params=None)[source]#

Wrapper around energy_fn to allow energy computation via traj_util.quantity_traj.

Parameters:
  • energy_fn_template – Function creating an energy function when called with energy parameters.

  • fixed_energy_params – Always use the energy function obtained when using the fixed energy params. If not given, the function uses dynamially specified parameters.