Miscellaneous fixes (#282)

This PR addresses the following small but unrelated issues:

 - [x] Missed edge cases not handled by #280 (71b5fa4).
 - [x] Some recent combination of changes on either or both CuPy and Jax makes `cupy.asarray(x)` for `x: JaxArray` return a copy instead of a view (134d085).
 - [x] Upcoming compatibility changes introduced by SSAGESLabs/hoomd-dlext/pull/23 (71991b5).
 - [x] Some fonts were intended to be loaded by #281, but currently fail (e6533cd).

Author: p-zubieta

docs/README.md

@@ -12,7 +12,8 @@ Make sure you have [GNU make](https://www.gpu.org/software/make/) installed.
 You can build the documentation from you local copy of the PySAGES repository as follows:
 
 ```shell
-pip install -r docs/requirements.txt
+cd docs
+pip install -r requirements.txt
 make html
 ```
 

docs/source/_templates/layout.html

@@ -1,6 +0,0 @@
-{% extends "!layout.html" %} {% block extrahead %}
-<link
-  rel="stylesheet"
-  href="https://fonts.googleapis.com/css?family=Atkinson Hyperlegible|Montserrat"
-/>
-{% endblock %}

docs/source/conf.py

@@ -100,6 +100,7 @@
 # # so a file named "default.css" will overwrite the builtin "default.css".
 html_static_path = ["_static"]
 html_css_files = [
+    "https://fonts.googleapis.com/css?family=Atkinson Hyperlegible|Montserrat",
     "https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/fontawesome.min.css",
     "https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/solid.min.css",
     "https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/brands.min.css",

pysages/backends/hoomd.py

@@ -20,7 +20,7 @@
 )
 from jax import jit
 from jax import numpy as np
-from jax.dlpack import from_dlpack as asarray
+from jax.dlpack import from_dlpack
 
 from pysages.backends.core import SamplingContext
 from pysages.backends.snapshot import (
@@ -61,11 +61,15 @@ def remove_half_step_hook(context):
         context.integrator.cpp_integrator.removeHalfStepHook()
 
 else:
+    if hasattr(hoomd.dlext, "__version__"):
+        SamplerBase = DLExtSampler
 
-    class SamplerBase(DLExtSampler, md.HalfStepHook):
-        def __init__(self, sysview, update, location, mode):
-            md.HalfStepHook.__init__(self)
-            DLExtSampler.__init__(self, sysview, update, location, mode)
+    else:
+
+        class SamplerBase(DLExtSampler, md.HalfStepHook):
+            def __init__(self, sysview, update, location, mode):
+                md.HalfStepHook.__init__(self)
+                DLExtSampler.__init__(self, sysview, update, location, mode)
 
     def is_on_gpu(context):
         return not isinstance(context.device, hoomd.device.CPU)
@@ -125,11 +129,11 @@ def snapshot_callback(positions, vel_mass, rtags, images, forces, n):
 
     def _pack_snapshot(self, positions, vel_mass, forces, rtags, images):
         return Snapshot(
-            asarray(positions),
-            asarray(vel_mass),
-            asarray(forces),
-            asarray(rtags),
-            asarray(images),
+            from_dlpack(positions),
+            from_dlpack(vel_mass),
+            from_dlpack(forces),
+            from_dlpack(rtags),
+            from_dlpack(images),
             self.box,
             self.dt,
         )
@@ -149,11 +153,11 @@ def default_location():
 def take_snapshot(sampling_context, location=default_location()):
     context = sampling_context.context
     sysview = sampling_context.view
-    positions = copy(asarray(positions_types(sysview, location, AccessMode.Read)))
-    vel_mass = copy(asarray(velocities_masses(sysview, location, AccessMode.Read)))
-    forces = copy(asarray(net_forces(sysview, location, AccessMode.ReadWrite)))
-    ids = copy(asarray(rtags(sysview, location, AccessMode.Read)))
-    imgs = copy(asarray(images(sysview, location, AccessMode.Read)))
+    positions = copy(from_dlpack(positions_types(sysview, location, AccessMode.Read)))
+    vel_mass = copy(from_dlpack(velocities_masses(sysview, location, AccessMode.Read)))
+    forces = copy(from_dlpack(net_forces(sysview, location, AccessMode.ReadWrite)))
+    ids = copy(from_dlpack(rtags(sysview, location, AccessMode.Read)))
+    imgs = copy(from_dlpack(images(sysview, location, AccessMode.Read)))
 
     check_device_array(positions)  # currently, we only support `DeviceArray`s
 
@@ -200,17 +204,14 @@ def masses(snapshot):
 
 
 def build_helpers(context, sampling_method):
+    utils = importlib.import_module(".utils", package="pysages.backends")
+
     # Depending on the device being used we need to use either cupy or numpy
     # (or numba) to generate a view of jax's DeviceArrays
     if is_on_gpu(context):
-        cupy = importlib.import_module("cupy")
-        view = cupy.asarray
-
-        def sync_forces():
-            cupy.cuda.get_current_stream().synchronize()
+        sync_forces, view = utils.cupy_helpers()
 
     else:
-        utils = importlib.import_module(".utils", package="pysages.backends")
         view = utils.view
 
         def sync_forces():

pysages/backends/lammps.py

@@ -14,7 +14,7 @@
 from jax import jit
 from jax import numpy as np
 from jax import vmap
-from jax.dlpack import from_dlpack as asarray
+from jax.dlpack import from_dlpack
 from lammps import dlext
 from lammps.dlext import ExecutionSpace, FixDLExt, LAMMPSView, has_kokkos_cuda_enabled
 
@@ -84,16 +84,16 @@ def update(timestep):
         self.set_callback(update)
 
     def _partial_snapshot(self, include_masses: bool = False):
-        positions = asarray(dlext.positions(self.view, self.location))
-        types = asarray(dlext.types(self.view, self.location))
-        velocities = asarray(dlext.velocities(self.view, self.location))
-        forces = asarray(dlext.forces(self.view, self.location))
-        tags_map = asarray(dlext.tags_map(self.view, self.location))
-        imgs = asarray(dlext.images(self.view, self.location))
+        positions = from_dlpack(dlext.positions(self.view, self.location))
+        types = from_dlpack(dlext.types(self.view, self.location))
+        velocities = from_dlpack(dlext.velocities(self.view, self.location))
+        forces = from_dlpack(dlext.forces(self.view, self.location))
+        tags_map = from_dlpack(dlext.tags_map(self.view, self.location))
+        imgs = from_dlpack(dlext.images(self.view, self.location))
 
         masses = None
         if include_masses:
-            masses = asarray(dlext.masses(self.view, self.location))
+            masses = from_dlpack(dlext.masses(self.view, self.location))
         vel_mass = (velocities, (masses, types))
 
         return Snapshot(positions, vel_mass, forces, tags_map, imgs, None, None)
@@ -127,19 +127,15 @@ def build_helpers(context, sampling_method, on_gpu, restore_fn):
     """
     Builds helper methods used for restoring snapshots and biasing a simulation.
     """
+    utils = importlib.import_module(".utils", package="pysages.backends")
     dim = context.extract_setting("dimension")
 
     # Depending on the device being used we need to use either cupy or numpy
     # (or numba) to generate a view of jax's DeviceArrays
     if on_gpu:
-        cupy = importlib.import_module("cupy")
-        view = cupy.asarray
-
-        def sync_forces():
-            cupy.cuda.get_current_stream().synchronize()
+        sync_forces, view = utils.cupy_helpers()
 
     else:
-        utils = importlib.import_module(".utils", package="pysages.backends")
         view = utils.view
 
         def sync_forces():

pysages/backends/openmm.py

@@ -8,7 +8,7 @@
 import openmm_dlext as dlext
 from jax import jit
 from jax import numpy as np
-from jax.dlpack import from_dlpack as asarray
+from jax.dlpack import from_dlpack
 from jax.lax import cond
 from openmm_dlext import ContextView, DeviceType, Force
 
@@ -60,15 +60,15 @@ def take_snapshot(sampling_context):
     context = sampling_context.context.context  # extra indirection for OpenMM
     context_view = sampling_context.view
 
-    positions = asarray(dlext.positions(context_view))
-    forces = asarray(dlext.forces(context_view))
-    ids = asarray(dlext.atom_ids(context_view))
+    positions = from_dlpack(dlext.positions(context_view))
+    forces = from_dlpack(dlext.forces(context_view))
+    ids = from_dlpack(dlext.atom_ids(context_view))
 
-    velocities = asarray(dlext.velocities(context_view))
+    velocities = from_dlpack(dlext.velocities(context_view))
     if is_on_gpu(context_view):
         vel_mass = velocities
     else:
-        inverse_masses = asarray(dlext.inverse_masses(context_view))
+        inverse_masses = from_dlpack(dlext.inverse_masses(context_view))
         vel_mass = (velocities, inverse_masses.reshape((-1, 1)))
 
     check_device_array(positions)  # currently, we only support `DeviceArray`s
@@ -126,29 +126,21 @@ def momenta(snapshot):
 
 
 def build_helpers(context, sampling_method):
+    utils = importlib.import_module(".utils", package="pysages.backends")
+
     # Depending on the device being used we need to use either cupy or numpy
     # (or numba) to generate a view of jax's DeviceArrays
     if is_on_gpu(context):
-        cupy = importlib.import_module("cupy")
-        view = cupy.asarray
-
         restore_vm = _restore_vm
-
-        def sync_forces():
-            cupy.cuda.get_current_stream().synchronize()
+        sync_forces, view = utils.cupy_helpers()
 
         @jit
         def adapt(biases):
             return np.int64(2**32 * biases.T)
 
     else:
-        utils = importlib.import_module(".utils", package="pysages.backends")
-        view = utils.view
-
         adapt = identity
-
-        def sync_forces():
-            pass
+        view = utils.view
 
         def restore_vm(view, snapshot, prev_snapshot):
             # TODO: Check if we can omit modifying the masses
@@ -158,6 +150,9 @@ def restore_vm(view, snapshot, prev_snapshot):
             velocities[:] = view(prev_snapshot.vel_mass[0])
             masses[:] = view(prev_snapshot.vel_mass[1])
 
+        def sync_forces():
+            pass
+
     def bias(snapshot, state, sync_backend):
         """Adds the computed bias to the forces."""
         if state.bias is None:

pysages/backends/utils.py

@@ -2,6 +2,7 @@
 # See LICENSE.md and CONTRIBUTORS.md at https://github.com/SSAGESLabs/PySAGES
 
 import ctypes
+import importlib
 
 import numba
 import numpy
@@ -11,6 +12,26 @@
 from pysages.utils import dispatch
 
 
+def cupy_helpers():
+    """Returns two methods:
+
+    `sync` -- for synchronizing the current CUDA stream
+    `view` -- to wrap a `JaxArray` as a `cupy.ndarray`
+    """
+    cupy = importlib.import_module("cupy")
+    dlpack = importlib.import_module("jax.dlpack")
+
+    def _sync():
+        """Synchronizes the current cupy's CUDA stream."""
+        cupy.cuda.get_current_stream().synchronize()
+
+    def _view(x: JaxArray):
+        """Wraps a view of `x: JaxArray` as a `cupy.ndarray`."""
+        return cupy.from_dlpack(dlpack.to_dlpack(x))
+
+    return _sync, _view
+
+
 @dispatch
 def view(array: JaxArray):
     """Return a writable view of a JAX DeviceArray."""

pysages/utils/compat.py

@@ -67,6 +67,7 @@ def has_method(fn, T, index):
 else:
     _bt = import_module("beartype.door")
     _typing = import_module("plum" if _plum_version_tuple < (2, 2, 1) else "typing")
+    _util = _typing.type if _plum_version_tuple < (2, 2, 1) else _typing
 
     def dispatch_table(dispatch):
         return dispatch.functions
@@ -75,8 +76,8 @@ def has_method(fn, T, index):
         types_at_index = set()
         for sig in fn.methods:
             typ = sig.types[index]
-            if _typing.get_origin(typ) is _typing.Union:
-                types_at_index.update(_typing.get_args(typ))
+            if _util.get_origin(typ) is _typing.Union:
+                types_at_index.update(_util.get_args(typ))
             else:
                 types_at_index.add(typ)
         return T in types_at_index