Merge branch 'main' into vp-flow

Author: Olllom

.github/workflows/CI.yml

@@ -41,7 +41,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install pytest "${{ matrix.cfg.torch-version }}" numpy nflows torchdiffeq einops
+          pip install pytest "${{ matrix.cfg.torch-version }}" numpy nflows torchdiffeq einops netCDF4
       - name: Install package
         run: |
           python setup.py install

README.md

@@ -139,6 +139,7 @@ pytest
   * [OpenMM](https://github.com/openmm/openmm) (for molecular examples)
   * [torchdiffeq](https://github.com/rtqichen/torchdiffeq) (for neural ODEs)
   * [ANODE](https://github.com/amirgholami/anode) (for neural ODEs)
+  * [netCDF4](https://unidata.github.io/netcdf4-python/) (for the `ReplayBufferReporter`)
   * [jax](https://github.com/google/jax) (for smooth flows / implicit backprop)
   * [jax2torch](https://github.com/lucidrains/jax2torch) (for smooth flows / implicit backprop)
 

bgflow/bg.py

@@ -2,6 +2,13 @@
 
 from .distribution.energy import Energy
 from .distribution.sampling import Sampler
+from .utils.types import pack_tensor_in_tuple
+
+__all__ = [
+    "BoltzmannGenerator", "unnormalized_kl_div", "unormalized_nll",
+    "sampling_efficiency", "effective_sample_size", "log_weights",
+    "log_weights_given_latent"
+]
 
 __all__ = [
     "BoltzmannGenerator", "unnormalized_kl_div", "unormalized_nll",
@@ -12,8 +19,7 @@
 
 def unnormalized_kl_div(prior, flow, target, n_samples, temperature=1.0):
     z = prior.sample(n_samples, temperature=temperature)
-    if isinstance(z, torch.Tensor):
-        z = (z,)
+    z = pack_tensor_in_tuple(z)
     *x, dlogp = flow(*z, temperature=temperature)
     return target.energy(*x, temperature=temperature) - dlogp
 
@@ -31,10 +37,8 @@ def log_weights(*x, prior, flow, target, temperature=1.0, normalize=True):
 
 
 def log_weights_given_latent(x, z, dlogp, prior, target, temperature=1.0, normalize=True):
-    if isinstance(x, torch.Tensor):
-        x = (x,)
-    if isinstance(z, torch.Tensor):
-        z = (z,)
+    x = pack_tensor_in_tuple(x)
+    z = pack_tensor_in_tuple(z)
     logw = (
         prior.energy(*z, temperature=temperature)
         + dlogp
@@ -94,8 +98,7 @@ def sample(
         with_weights=False,
     ):
         z = self._prior.sample(n_samples, temperature=temperature)
-        if isinstance(z, torch.Tensor):
-            z = (z,)
+        z = pack_tensor_in_tuple(z)
         *x, dlogp = self._flow(*z, temperature=temperature)
         results = list(x)
 

bgflow/distribution/energy/double_well.py

@@ -14,8 +14,8 @@ def __init__(self, dim, a=0, b=-4.0, c=1.0):
         self._c = c
 
     def _energy(self, x):
-        d = x[:, [0]]
-        v = x[:, 1:]
+        d = x[..., [0]]
+        v = x[..., 1:]
         e1 = self._a * d + self._b * d.pow(2) + self._c * d.pow(4)
         e2 = 0.5 * v.pow(2).sum(dim=-1, keepdim=True)
         return e1 + e2

bgflow/distribution/sampling/__init__.py

@@ -1,3 +1,5 @@
 from .base import *
 from .mcmc import *
-from .dataset import *
+from .dataset import *
+from .buffer import *
+from .iterative import *

bgflow/distribution/sampling/_iterative_helpers.py

@@ -0,0 +1,97 @@
+"""Helper classes and functions for iterative samplers."""
+
+import torch
+
+
+__all__ = ["AbstractSamplerState", "default_set_samples_hook", "default_extract_sample_hook"]
+
+
+class AbstractSamplerState:
+    """Defines the interface for implementations of the internal state of iterative samplers."""
+
+    def as_dict(self):
+        """Return a dictionary representing this instance. The dictionary has to define the
+        keys that are used within `SamplerStep`s of an `IterativeSampler`, such as "samples", "energies", ...
+        """
+        raise NotImplementedError()
+
+    def _replace(self, **kwargs):
+        """Return a new object with changed fields.
+        This function has to support all the keys that are used
+        within `SamplerStep`s of an `IterativeSampler` as well as the keys "energies_up_to_date" and
+        "forces_up_to_date"
+        """
+        raise NotImplementedError()
+
+    def evaluate_energy_force(self, energy_model, evaluate_energies=True, evaluate_forces=True):
+        """Return a new state with updated energies/forces."""
+        state = self.as_dict()
+        evaluate_energies = evaluate_energies and not state["energies_up_to_date"]
+        energies = energy_model.energy(*state["samples"])[..., 0] if evaluate_energies else state["energies"]
+
+        evaluate_forces = evaluate_forces and not state["forces_up_to_date"]
+        forces = energy_model.force(*state["samples"]) if evaluate_forces else state["forces"]
+        return self.replace(energies=energies, forces=forces)
+
+    def replace(self, **kwargs):
+        """Return a new state with updated fields."""
+
+        # keep track of energies and forces
+        state_dict = self.as_dict()
+        if "energies" in kwargs:
+            kwargs = {**kwargs, "energies_up_to_date": True}
+        elif "samples" in kwargs:
+            kwargs = {**kwargs, "energies_up_to_date": False}
+        if "forces" in kwargs:
+            kwargs = {**kwargs, "forces_up_to_date": True}
+        elif "samples" in kwargs:
+            kwargs = {**kwargs, "forces_up_to_date": False}
+
+        # map to primary unit cell
+        box_vectors = None
+        if "box_vectors" in kwargs:
+            box_vectors = kwargs["box_vectors"]
+        elif "box_vectors" in state_dict:
+            box_vectors = state_dict["box_vectors"]
+        if "samples" in kwargs and box_vectors is not None:
+            kwargs = {
+                **kwargs,
+                "samples": tuple(
+                    _map_to_primary_cell(x, cell)
+                    for x, cell in zip(kwargs["samples"], box_vectors)
+                )
+            }
+        return self._replace(**kwargs)
+
+
+def default_set_samples_hook(x):
+    """by default, use samples as is"""
+    return x
+
+
+def default_extract_sample_hook(state: AbstractSamplerState):
+    """Default extraction of samples from a SamplerState."""
+    return state.as_dict()["samples"]
+
+
+def _bmv(m, bv):
+    """Batched matrix-vector multiply."""
+    return torch.einsum("ij,...j->...i", m, bv)
+
+
+def _map_to_primary_cell(x, cell):
+    """Map coordinates to the primary unit cell of a periodic lattice.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        n-dimensional coordinates of shape (..., n), where n is the spatial dimension and ... denote an
+        arbitrary number of batch dimensions.
+    cell : torch.Tensor
+        Lattice vectors (column-wise). Has to be upper triangular.
+    """
+    if cell is None:
+        return x
+    n = _bmv(torch.inverse(cell), x)
+    n = torch.floor(n)
+    return x - _bmv(cell, n)

bgflow/distribution/sampling/base.py

@@ -1,13 +1,27 @@
-import torch
 
+from typing import Tuple
+import torch
+from ...utils.types import unpack_tensor_tuple, pack_tensor_in_list
 
 __all__ = ["Sampler"]
 
 
 class Sampler(torch.nn.Module):
-    
-    def __init__(self, **kwargs):
+    """Abstract base class for samplers.
+
+    Parameters
+    ----------
+    return_hook : Callable, optional
+        A function to postprocess the samples. This can (for example) be used to
+        only return samples at a selected thermodynamic state of a replica exchange sampler
+        or to combine the batch and sample dimension.
+        The function takes a list of tensors and should return a list of tensors.
+        Each tensor contains a batch of samples.
+    """
+
+    def __init__(self, return_hook=lambda x: x, **kwargs):
         super().__init__(**kwargs)
+        self.return_hook = return_hook
 
     def _sample_with_temperature(self, n_samples, temperature, *args, **kwargs):
         raise NotImplementedError()
@@ -16,7 +30,39 @@ def _sample(self, n_samples, *args, **kwargs):
         raise NotImplementedError()
 
     def sample(self, n_samples, temperature=1.0, *args, **kwargs):
+        """Create a number of samples.
+
+        Parameters
+        ----------
+        n_samples : int
+            The number of samples to be created.
+        temperature : float, optional
+            The relative temperature at which to create samples.
+            Only available for sampler that implement `_sample_with_temperature`.
+
+        Returns
+        -------
+        samples : Union[torch.Tensor, Tuple[torch.Tensor, ...]]
+            If this sampler reflects a joint distribution of multiple tensors,
+            it returns a tuple of tensors, each of which have length n_samples.
+            Otherwise it returns a single tensor of length n_samples.
+        """
         if isinstance(temperature, float) and temperature == 1.0:
-            return self._sample(n_samples, *args, **kwargs)
+            samples = self._sample(n_samples, *args, **kwargs)
         else:
-            return self._sample_with_temperature(n_samples, temperature, *args, **kwargs)
+            samples = self._sample_with_temperature(n_samples, temperature, *args, **kwargs)
+        samples = pack_tensor_in_list(samples)
+        return unpack_tensor_tuple(self.return_hook(samples))
+
+    def sample_to_cpu(self, n_samples, batch_size=64, *args,  **kwargs):
+        """A utility method for creating many samples that might not fit into GPU memory."""
+        with torch.no_grad():
+            samples = self.sample(min(n_samples, batch_size), *args, **kwargs)
+            samples = pack_tensor_in_list(samples)
+            samples = [tensor.detach().cpu() for tensor in samples]
+            while len(samples[0]) < n_samples:
+                new_samples = self.sample(min(n_samples-len(samples[0]), batch_size), *args, **kwargs)
+                new_samples = pack_tensor_in_list(new_samples)
+                for i, new in enumerate(new_samples):
+                    samples[i] = torch.cat([samples[i], new.detach().cpu()], dim=0)
+        return unpack_tensor_tuple(samples)