Move unused rescale_coordinates to attic

Author: maxentile

attic/jax_tricks/rescale.py

@@ -0,0 +1,18 @@
+import jax
+from jax import numpy as np
+import numpy as onp
+
+
+def rescale_coordinates(conf, indices, box, scales):
+    """Note: scales unused"""
+
+    mol_sizes = np.expand_dims(onp.bincount(indices), axis=1)
+    mol_centers = jax.ops.segment_sum(conf, indices) / mol_sizes
+
+    new_centers = mol_centers - box[2] * np.floor(np.expand_dims(mol_centers[..., 2], axis=-1) / box[2][2])
+    new_centers -= box[1] * np.floor(np.expand_dims(new_centers[..., 1], axis=-1) / box[1][1])
+    new_centers -= box[0] * np.floor(np.expand_dims(new_centers[..., 0], axis=-1) / box[0][0])
+
+    offset = new_centers - mol_centers
+
+    return conf + offset[indices]

attic/readme.md

@@ -2,6 +2,7 @@ Code, documentation, experiments we want to retain for reference, but that we're
 
 ### Contents
 * `docs/` -- write-up of initial vision for `timemachine`, involving efficient backpropagation through MD trajectories
+* `jax_tricks` -- misc. Jax functions
 * `thermo_deriv/` -- numerical experiments with "thermodynamic derivative" estimators, adjusting LJ parameters to match observables
     * note: currently missing dependencies `thermo_deriv.lj_non_periodic.lennard_jones`, `thermo_deriv.lj.lennard_jones`.
     * note: `langevin_coefficients` dependency has since changed -- some scripts rely on a version of `langevin_coefficients` prior to PR #459

timemachine/potentials/jax_utils.py

@@ -19,7 +19,6 @@
 
 from typing import Tuple
 
-import jax
 import jax.numpy as np
 import numpy as onp
 from jax import vmap
@@ -88,21 +87,6 @@ def convert_to_4d(x3, lamb, lambda_plane_idxs, lambda_offset_idxs, cutoff):
     return augment_dim(x3, w)
 
 
-def rescale_coordinates(conf, indices, box, scales):
-    """Note: scales unused"""
-
-    mol_sizes = np.expand_dims(onp.bincount(indices), axis=1)
-    mol_centers = jax.ops.segment_sum(conf, indices) / mol_sizes
-
-    new_centers = mol_centers - box[2] * np.floor(np.expand_dims(mol_centers[..., 2], axis=-1) / box[2][2])
-    new_centers -= box[1] * np.floor(np.expand_dims(new_centers[..., 1], axis=-1) / box[1][1])
-    new_centers -= box[0] * np.floor(np.expand_dims(new_centers[..., 0], axis=-1) / box[0][0])
-
-    offset = new_centers - mol_centers
-
-    return conf + offset[indices]
-
-
 def delta_r(ri, rj, box=None):
     diff = ri - rj  # this can be either N,N,3 or B,3