Flow matching smc

.pre-commit-config.yaml

@@ -12,7 +12,7 @@ repos:
     rev: v1.1.396
     hooks:
     - id: pyright
-      additional_dependencies: [beartype, einops, jax, jaxtyping, pytest, typing_extensions, equinox, optax, tqdm]
+      additional_dependencies: [beartype, einops, jax, jaxtyping, pytest, typing_extensions, equinox, optax, tqdm, diffrax]
   - repo: https://github.com/nbQA-dev/nbQA
     rev: 1.9.1
     hooks:

docs/tutorials/custom_strategy.ipynb

@@ -29,7 +29,7 @@
    "source": [
     "import jax\n",
     "import jax.numpy as jnp\n",
-    "from flowMC.resource.local_kernel.Gaussian_random_walk import GaussianRandomWalk\n",
+    "from flowMC.resource.kernel.Gaussian_random_walk import GaussianRandomWalk\n",
     "from flowMC.resource.buffers import Buffer\n",
     "from flowMC.resource.logPDF import LogPDF\n",
     "from flowMC.strategy.optimization import AdamOptimization\n",

docs/tutorials/parallel_tempering.ipynb

@@ -19,7 +19,7 @@
     "from jaxtyping import Float, Array\n",
     "from typing import Any\n",
     "\n",
-    "from flowMC.resource.local_kernel.MALA import MALA\n",
+    "from flowMC.resource.kernel.MALA import MALA\n",
     "from flowMC.resource.buffers import Buffer\n",
     "from flowMC.resource.states import State\n",
     "from flowMC.strategy.take_steps import TakeSerialSteps\n",

docs/tutorials/train_normalizing_flow.ipynb

@@ -30,8 +30,8 @@
     "import equinox as eqx  # Equinox\n",
     "\n",
     "\n",
-    "from flowMC.resource.nf_model.realNVP import RealNVP\n",
-    "from flowMC.resource.nf_model.rqSpline import MaskedCouplingRQSpline"
+    "from flowMC.resource.model.nf_model.realNVP import RealNVP\n",
+    "from flowMC.resource.model.nf_model.rqSpline import MaskedCouplingRQSpline"
    ]
   },
   {

pyproject.toml

@@ -24,6 +24,7 @@ requires-python = ">=3.11"
 keywords = ["sampling", "inference", "machine learning", "normalizing", "autodiff", "jax"]
 dependencies = [
     "chex>=0.1.87",
+    "diffrax>=0.7.0",
     "equinox>=0.11.9",
     "jax[cpu]>=0.5.0",
     "jaxtyping>=0.2.36",
@@ -55,18 +56,17 @@ visualize = [
 cuda = [
     "jax[cuda12]>=0.5.0",
 ]
-codeqa = [
-    "coveralls>=4.0.1",
-    "pre-commit>=4.0.1",
-    "pyright>=1.1.389",
-    "pytest>=8.3.3",
-    "ruff>=0.8.0",
-]
 
 [dependency-groups]
 dev = [
     "flowMC",
     "ipykernel>=6.29.5",
+    "coveralls>=4.0.1",
+    "pre-commit>=4.0.1",
+    "pyright>=1.1.389",
+    "pytest>=8.3.3",
+    "ruff>=0.8.0",
+    "ipython>=8.30.0",
 ]
 
 [tool.uv.sources]

src/flowMC/resource/local_kernel/Gaussian_random_walk.py → src/flowMC/resource/kernel/Gaussian_random_walk.py

@@ -3,7 +3,7 @@
 from jaxtyping import Array, Float, Int, PRNGKeyArray, PyTree
 from typing import Callable
 
-from flowMC.resource.local_kernel.base import ProposalBase
+from flowMC.resource.kernel.base import ProposalBase
 from flowMC.resource.logPDF import LogPDF
 
 

src/flowMC/resource/local_kernel/HMC.py → src/flowMC/resource/kernel/HMC.py

@@ -4,7 +4,7 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Float, Int, PRNGKeyArray, PyTree
 
-from flowMC.resource.local_kernel.base import ProposalBase
+from flowMC.resource.kernel.base import ProposalBase
 from flowMC.resource.logPDF import LogPDF
 
 

src/flowMC/resource/local_kernel/MALA.py → src/flowMC/resource/kernel/MALA.py

@@ -5,7 +5,7 @@
 from typing import Callable
 
 from flowMC.resource.logPDF import LogPDF
-from flowMC.resource.local_kernel.base import ProposalBase
+from flowMC.resource.kernel.base import ProposalBase
 
 
 class MALA(ProposalBase):

src/flowMC/resource/nf_model/NF_proposal.py → src/flowMC/resource/kernel/NF_proposal.py

@@ -7,8 +7,8 @@
 from typing import Callable
 import equinox as eqx
 
-from flowMC.resource.nf_model.base import NFModel
-from flowMC.resource.local_kernel.base import ProposalBase
+from flowMC.resource.model.nf_model.base import NFModel
+from flowMC.resource.kernel.base import ProposalBase
 from flowMC.resource.logPDF import LogPDF
 
 

src/flowMC/resource/nf_model/common.py → src/flowMC/resource/model/common.py

@@ -1,11 +1,68 @@
-from typing import Callable, List, Tuple
+from typing import Callable, List, Tuple, Optional
 
 import equinox as eqx
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, Float, PRNGKeyArray
+from abc import abstractmethod
 
-from flowMC.resource.nf_model.base import Bijection, Distribution
+
+class Bijection(eqx.Module):
+    """Base class for bijective transformations.
+
+    This is an abstract template that should not be directly used.
+    """
+
+    @abstractmethod
+    def __init__(self):
+        raise NotImplementedError
+
+    def __call__(
+        self,
+        x: Float[Array, " n_dim"],
+        condition: Float[Array, " n_condition"],
+    ) -> tuple[Float[Array, " n_dim"], Float]:
+        return self.forward(x, condition)
+
+    @abstractmethod
+    def forward(
+        self,
+        x: Float[Array, " n_dim"],
+        condition: Float[Array, " n_condition"],
+    ) -> tuple[Float[Array, " n_dim"], Float]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def inverse(
+        self,
+        x: Float[Array, " n_dim"],
+        condition: Float[Array, " n_condition"],
+    ) -> tuple[Float[Array, " n_dim"], Float]:
+        raise NotImplementedError
+
+
+class Distribution(eqx.Module):
+    """Base class for probability distributions.
+
+    This is an abstract template that should not be directly used.
+    """
+
+    @abstractmethod
+    def __init__(self):
+        raise NotImplementedError
+
+    def __call__(self, x: Array, key: Optional[PRNGKeyArray] = None) -> Array:
+        return self.log_prob(x)
+
+    @abstractmethod
+    def log_prob(self, x: Array) -> Array:
+        raise NotImplementedError
+
+    @abstractmethod
+    def sample(
+        self, rng_key: PRNGKeyArray, n_samples: int
+    ) -> Float[Array, " n_samples n_features"]:
+        raise NotImplementedError
 
 
 class MLP(eqx.Module):