Fix the Zonly filter bug (#111)

* change sgnn mdanalysis dependence to mdtraj dependence

* Change the way generators.admp and generators.classical are imported
So the backends can still be accessible via dmff.admp and dmff.classical

* Change jax and jaxmd version requirements

* Allow jax 0.4.1

* Add the DMC example files from Wei Feng

* update DMC example files

* Change the API of OpenMMSamplerState, make it more flexible

* Delete prm.xml

* Fix the ADMP PME with only monopoles

* Fix the bug in admp/recip.py for noncubic box.

* Improve OpenMMSampleState interface, fix the running platform in test

* Update the version of the dependence packages in github test worflow.

* Update ut.yml

* Update ut.yml

* Hotfix: the Zonly atoms missing filter

* Change test python version to 3.9

dmff/admp/spatial.py

@@ -101,7 +101,8 @@ def construct_local_frames(positions, box):
         vec_y = jnp.zeros((n_sites, 3))
         # Z-Only
         x_of_vec_z = jnp.round(jnp.abs(vec_z[:,0]))
-        vec_x_Zonly = jnp.array([1.-x_of_vec_z, x_of_vec_z, jnp.zeros_like(x_of_vec_z)]).T
+        vec_x_Zonly = jnp.array([1.-x_of_vec_z, x_of_vec_z, jnp.zeros_like(x_of_vec_z)]).T[Zonly_filter]
+
         vec_x = vec_x.at[Zonly_filter].set(vec_x_Zonly)
         # for those that are not Z-Only, get normalized vecX
         vec_x_not_Zonly = positions[x_atoms[not_Zonly_filter]] - positions[not_Zonly_filter]

dmff/settings.py

@@ -16,4 +16,4 @@ def update_jax_precision(precision):
 
 update_jax_precision(PRECISION)
 
-__all__ = ['PRECISION', 'DO_JIT', 'DEBUG', "update_jax_precision"]
+__all__ = ['PRECISION', 'DO_JIT', 'DEBUG', "update_jax_precision"]