diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index 0c17461..df65570 100644 --- a/src/loch/_sampler.py +++ b/src/loch/_sampler.py @@ -2397,6 +2397,12 @@ def _initialise_gpu_memory(self): "Could not find the GhostNonGhostNonbondedForce in the system" ) + # Resolve the unit conversions once. Doing this per atom, by + # formatting and re-parsing a string, dominates the setup time for + # large systems. + nm_to_angstrom = _sr.u("1 nm").to("angstrom") + kj_per_mol_to_kcal_per_mol = _sr.u("1 kJ/mol").to("kcal/mol") + # Get the parameters for the GhostNonGhostNonbondedForce. charges = _np.zeros(self._num_atoms, dtype=_np.float32) sigmas = _np.zeros(self._num_atoms, dtype=_np.float32) @@ -2410,10 +2416,8 @@ def _initialise_gpu_memory(self): # Charge in |e|, sigma in nm, epsilon in kJ/mol. charges[i] = q # Rescale and convert units. - sigmas[i] = _sr.u(f"{2.0 * half_sigma} nm").to("angstrom") - epsilons[i] = _sr.u(f"{(0.5 * two_sqrt_epsilon) ** 2} kJ/mol").to( - "kcal/mol" - ) + sigmas[i] = 2.0 * half_sigma * nm_to_angstrom + epsilons[i] = (0.5 * two_sqrt_epsilon) ** 2 * kj_per_mol_to_kcal_per_mol # Store the softening parameter. alphas[i] = alpha