Merge pull request #35 from couxxuoc/fix/33-check-positions-jaxmd_utils

Fix checking the non-fractional positions (fixes #33)

Author: kabylda

so3lr/jaxmd_utils.py

@@ -47,7 +47,8 @@ def featurize(R, neighbor, neighbor_lr, **kwargs):
 
         if 'perturbation' in kwargs:
             pert = kwargs.get('perturbation')
-            positions = transform(pert, positions)
+            if positions is not None:
+                positions = transform(pert, positions)
             box = transform(pert, box)
 
         d = jax.vmap(partial(displacement_fn, **kwargs))