aspire-inference 0.1.0a5__py3-none-any.whl → 0.1.0a6__py3-none-any.whl

aspire/aspire.py

@@ -6,6 +6,7 @@ from typing import Any, Callable
 import h5py
 
 from .flows import get_flow_wrapper
+from .flows.base import Flow
 from .history import History
 from .samples import Samples
 from .transforms import (
@@ -48,12 +49,17 @@ class Aspire:
     xp : Callable | None
         The array backend to use. If None, the default backend will be
         used.
+    flow : Flow | None
+        The flow object, if it already exists.
+        If None, a new flow will be created.
     flow_backend : str
         The backend to use for the flow. Options are 'zuko' or 'flowjax'.
     flow_matching : bool
         Whether to use flow matching.
     eps : float
         The epsilon value to use for data transforms.
+    dtype : Any | str | None
+        The data type to use for the samples, flow and transforms.
     **kwargs
         Keyword arguments to pass to the flow.
     """
@@ -71,9 +77,11 @@ class Aspire:
         bounded_transform: str = "logit",
         device: str | None = None,
         xp: Callable | None = None,
+        flow: Flow | None = None,
         flow_backend: str = "zuko",
         flow_matching: bool = False,
         eps: float = 1e-6,
+        dtype: Any | str | None = None,
         **kwargs,
     ) -> None:
         self.log_likelihood = log_likelihood
@@ -91,14 +99,20 @@ class Aspire:
         self.flow_backend = flow_backend
         self.flow_kwargs = kwargs
         self.xp = xp
+        self.dtype = dtype
 
-        self._flow = None
+        self._flow = flow
 
     @property
     def flow(self):
         """The normalizing flow object."""
         return self._flow
 
+    @flow.setter
+    def flow(self, flow: Flow):
+        """Set the normalizing flow object."""
+        self._flow = flow
+
     @property
     def sampler(self):
         """The sampler object."""
@@ -130,6 +144,7 @@ class Aspire:
             log_prior=log_prior,
             log_q=log_q,
             xp=xp,
+            dtype=self.dtype,
         )
 
         if evaluate:
@@ -159,6 +174,7 @@ class Aspire:
             device=self.device,
             xp=xp,
             eps=self.eps,
+            dtype=self.dtype,
         )
 
         # Check if FlowClass takes `parameters` as an argument
@@ -172,6 +188,7 @@ class Aspire:
             dims=self.dims,
             device=self.device,
             data_transform=data_transform,
+            dtype=self.dtype,
             **self.flow_kwargs,
         )
 
@@ -245,6 +262,7 @@ class Aspire:
                 periodic_parameters=self.periodic_parameters,
                 xp=self.xp,
                 device=self.device,
+                dtype=self.dtype,
                 **preconditioning_kwargs,
             )
         elif preconditioning == "flow":
@@ -259,6 +277,7 @@ class Aspire:
                 bounded_to_unbounded=self.bounded_to_unbounded,
                 prior_bounds=self.prior_bounds,
                 xp=self.xp,
+                dtype=self.dtype,
                 device=self.device,
                 **preconditioning_kwargs,
             )
@@ -271,6 +290,7 @@ class Aspire:
             dims=self.dims,
             prior_flow=self.flow,
             xp=self.xp,
+            dtype=self.dtype,
             preconditioning_transform=transform,
             **kwargs,
         )
@@ -397,17 +417,17 @@ class Aspire:
             method of the sampler.
         """
         config = {
-            # "log_likelihood": self.log_likelihood,
-            # "log_prior": self.log_prior,
+            "log_likelihood": self.log_likelihood.__name__,
+            "log_prior": self.log_prior.__name__,
             "dims": self.dims,
             "parameters": self.parameters,
             "periodic_parameters": self.periodic_parameters,
             "prior_bounds": self.prior_bounds,
             "bounded_to_unbounded": self.bounded_to_unbounded,
-            # "bounded_transform": self.bounded_transform,
+            "bounded_transform": self.bounded_transform,
             "flow_matching": self.flow_matching,
-            # "device": self.device,
-            # "xp": self.xp,
+            "device": self.device,
+            "xp": self.xp.__name__ if self.xp else None,
             "flow_backend": self.flow_backend,
             "flow_kwargs": self.flow_kwargs,
             "eps": self.eps,
@@ -437,6 +457,35 @@ class Aspire:
             self.config_dict(**kwargs),
         )
 
+    def save_flow(self, h5_file: h5py.File, path="flow") -> None:
+        """Save the flow to an HDF5 file.
+
+        Parameters
+        ----------
+        h5_file : h5py.File
+            The HDF5 file to save the flow to.
+        path : str
+            The path in the HDF5 file to save the flow to.
+        """
+        if self.flow is None:
+            raise ValueError("Flow has not been initialized.")
+        self.flow.save(h5_file, path=path)
+
+    def load_flow(self, h5_file: h5py.File, path="flow") -> None:
+        """Load the flow from an HDF5 file.
+
+        Parameters
+        ----------
+        h5_file : h5py.File
+            The HDF5 file to load the flow from.
+        path : str
+            The path in the HDF5 file to load the flow from.
+        """
+        FlowClass, xp = get_flow_wrapper(
+            backend=self.flow_backend, flow_matching=self.flow_matching
+        )
+        self._flow = FlowClass.load(h5_file, path=path)
+
     def save_config_to_json(self, filename: str) -> None:
         """Save the configuration to a JSON file."""
         import json

aspire/flows/base.py

@@ -1,3 +1,4 @@
+import inspect
 import logging
 from typing import Any
 
@@ -45,3 +46,39 @@ class Flow:
 
     def inverse_rescale(self, x):
         return self.data_transform.inverse(x)
+
+    def config_dict(self):
+        """Return a dictionary of the configuration of the flow.
+
+        This can be used to recreate the flow by passing the dictionary
+        as keyword arguments to the constructor.
+
+        This is automatically populated with the arguments passed to the
+        constructor.
+
+        Returns
+        -------
+        config : dict
+            The configuration dictionary.
+        """
+        return getattr(self, "_init_args", {})
+
+    def save(self, h5_file, path="flow"):
+        raise NotImplementedError
+
+    @classmethod
+    def load(cls, h5_file, path="flow"):
+        raise NotImplementedError
+
+    def __new__(cls, *args, **kwargs):
+        # Create instance
+        obj = super().__new__(cls)
+        # Inspect the subclass's __init__ signature
+        sig = inspect.signature(cls.__init__)
+        bound = sig.bind_partial(obj, *args, **kwargs)
+        bound.apply_defaults()
+        # Save args (excluding self)
+        obj._init_args = {
+            k: v for k, v in bound.arguments.items() if k != "self"
+        }
+        return obj

aspire/flows/jax/flows.py

@@ -1,10 +1,13 @@
 import logging
 from typing import Callable
 
+import jax
 import jax.numpy as jnp
 import jax.random as jrandom
 from flowjax.train import fit_to_data
 
+from ...transforms import IdentityTransform
+from ...utils import decode_dtype, encode_dtype, resolve_dtype
 from ..base import Flow
 from .utils import get_flow
 
@@ -14,11 +17,28 @@ logger = logging.getLogger(__name__)
 class FlowJax(Flow):
     xp = jnp
 
-    def __init__(self, dims: int, key=None, data_transform=None, **kwargs):
+    def __init__(
+        self,
+        dims: int,
+        key=None,
+        data_transform=None,
+        dtype=None,
+        **kwargs,
+    ):
         device = kwargs.pop("device", None)
         if device is not None:
             logger.warning("The device argument is not used in FlowJax. ")
+        resolved_dtype = (
+            resolve_dtype(dtype, jnp)
+            if dtype is not None
+            else jnp.dtype(jnp.float32)
+        )
+        if data_transform is None:
+            data_transform = IdentityTransform(self.xp, dtype=resolved_dtype)
+        elif getattr(data_transform, "dtype", None) is None:
+            data_transform.dtype = resolved_dtype
         super().__init__(dims, device=device, data_transform=data_transform)
+        self.dtype = resolved_dtype
         if key is None:
             key = jrandom.key(0)
             logger.warning(
@@ -33,14 +53,15 @@ class FlowJax(Flow):
         self._flow = get_flow(
             key=subkey,
             dims=self.dims,
+            dtype=self.dtype,
             **kwargs,
         )
 
     def fit(self, x, **kwargs):
         from ...history import FlowHistory
 
-        x = jnp.asarray(x)
-        x_prime = self.fit_data_transform(x)
+        x = jnp.asarray(x, dtype=self.dtype)
+        x_prime = jnp.asarray(self.fit_data_transform(x), dtype=self.dtype)
         self.key, subkey = jrandom.split(self.key)
         self._flow, losses = fit_to_data(subkey, self._flow, x_prime, **kwargs)
         return FlowHistory(
@@ -49,22 +70,27 @@ class FlowJax(Flow):
         )
 
     def forward(self, x, xp: Callable = jnp):
+        x = jnp.asarray(x, dtype=self.dtype)
         x_prime, log_abs_det_jacobian = self.rescale(x)
+        x_prime = jnp.asarray(x_prime, dtype=self.dtype)
         z, log_abs_det_jacobian_flow = self._flow.forward(x_prime)
         return xp.asarray(z), xp.asarray(
             log_abs_det_jacobian + log_abs_det_jacobian_flow
         )
 
     def inverse(self, z, xp: Callable = jnp):
-        z = jnp.asarray(z)
+        z = jnp.asarray(z, dtype=self.dtype)
         x_prime, log_abs_det_jacobian_flow = self._flow.inverse(z)
+        x_prime = jnp.asarray(x_prime, dtype=self.dtype)
         x, log_abs_det_jacobian = self.inverse_rescale(x_prime)
         return xp.asarray(x), xp.asarray(
             log_abs_det_jacobian + log_abs_det_jacobian_flow
         )
 
     def log_prob(self, x, xp: Callable = jnp):
+        x = jnp.asarray(x, dtype=self.dtype)
         x_prime, log_abs_det_jacobian = self.rescale(x)
+        x_prime = jnp.asarray(x_prime, dtype=self.dtype)
         log_prob = self._flow.log_prob(x_prime)
         return xp.asarray(log_prob + log_abs_det_jacobian)
 
@@ -80,3 +106,91 @@ class FlowJax(Flow):
         log_prob = self._flow.log_prob(x_prime)
         x, log_abs_det_jacobian = self.inverse_rescale(x_prime)
         return xp.asarray(x), xp.asarray(log_prob - log_abs_det_jacobian)
+
+    def save(self, h5_file, path="flow"):
+        import equinox as eqx
+        from array_api_compat import numpy as np
+
+        from ...utils import recursively_save_to_h5_file
+
+        grp = h5_file.require_group(path)
+
+        # ---- config ----
+        config = self.config_dict().copy()
+        config.pop("key", None)
+        config["key_data"] = jax.random.key_data(self.key)
+        dtype_value = config.get("dtype")
+        if dtype_value is None:
+            dtype_value = self.dtype
+        else:
+            dtype_value = jnp.dtype(dtype_value)
+        print(dtype_value)
+        config["dtype"] = encode_dtype(jnp, dtype_value)
+
+        data_transform = config.pop("data_transform", None)
+        if data_transform is not None:
+            data_transform.save(grp, "data_transform")
+
+        recursively_save_to_h5_file(grp, "config", config)
+
+        # ---- save arrays ----
+        arrays, _ = eqx.partition(self._flow, eqx.is_array)
+        leaves, _ = jax.tree_util.tree_flatten(arrays)
+
+        params_grp = grp.require_group("params")
+        # clear old datasets
+        for name in list(params_grp.keys()):
+            del params_grp[name]
+
+        for i, p in enumerate(leaves):
+            params_grp.create_dataset(str(i), data=np.asarray(p))
+
+    @classmethod
+    def load(cls, h5_file, path="flow"):
+        import equinox as eqx
+
+        from ...utils import load_from_h5_file
+
+        grp = h5_file[path]
+
+        # ---- config ----
+        config = load_from_h5_file(grp, "config")
+        config["dtype"] = decode_dtype(jnp, config.get("dtype"))
+        if "data_transform" in grp:
+            from ...transforms import BaseTransform
+
+            config["data_transform"] = BaseTransform.load(
+                grp,
+                "data_transform",
+                strict=False,
+            )
+
+        key_data = config.pop("key_data", None)
+        if key_data is not None:
+            config["key"] = jax.random.wrap_key_data(key_data)
+
+        kwargs = config.pop("kwargs", {})
+        config.update(kwargs)
+
+        # build object (will replace its _flow)
+        obj = cls(**config)
+
+        # ---- load arrays ----
+        params_grp = grp["params"]
+        loaded_params = [
+            jnp.array(params_grp[str(i)][:]) for i in range(len(params_grp))
+        ]
+
+        # rebuild template flow
+        kwargs.pop("device")
+        flow_template = get_flow(key=jrandom.key(0), dims=obj.dims, **kwargs)
+        arrays_template, static = eqx.partition(flow_template, eqx.is_array)
+
+        # use treedef from template
+        treedef = jax.tree_util.tree_structure(arrays_template)
+        arrays = jax.tree_util.tree_unflatten(treedef, loaded_params)
+
+        # recombine
+        obj._flow = eqx.combine(static, arrays)
+
+        return obj

aspire/flows/jax/utils.py

@@ -29,8 +29,11 @@ def get_flow(
     flow_type: str | Callable = "masked_autoregressive_flow",
     bijection_type: str | flowjax.bijections.AbstractBijection | None = None,
     bijection_kwargs: dict | None = None,
+    dtype=None,
     **kwargs,
 ) -> flowjax.distributions.Transformed:
+    dtype = dtype or jnp.float32
+
     if isinstance(flow_type, str):
         flow_type = get_flow_function_class(flow_type)
 
@@ -44,7 +47,7 @@ def get_flow(
     if bijection_kwargs is None:
         bijection_kwargs = {}
 
-    base_dist = flowjax.distributions.Normal(jnp.zeros(dims))
+    base_dist = flowjax.distributions.Normal(jnp.zeros(dims, dtype=dtype))
     key, subkey = jrandom.split(key)
     return flow_type(
         subkey,

aspire/flows/torch/flows.py

@@ -9,6 +9,8 @@ import zuko
 from array_api_compat import is_numpy_namespace, is_torch_array
 
 from ...history import FlowHistory
+from ...transforms import IdentityTransform
+from ...utils import decode_dtype, encode_dtype, resolve_dtype
 from ..base import Flow
 
 logger = logging.getLogger(__name__)
@@ -24,12 +26,23 @@ class BaseTorchFlow(Flow):
         seed: int = 1234,
         device: str = "cpu",
         data_transform=None,
+        dtype=None,
     ):
+        resolved_dtype = (
+            resolve_dtype(dtype, torch)
+            if dtype is not None
+            else torch.get_default_dtype()
+        )
+        if data_transform is None:
+            data_transform = IdentityTransform(self.xp, dtype=resolved_dtype)
+        elif getattr(data_transform, "dtype", None) is None:
+            data_transform.dtype = resolved_dtype
         super().__init__(
             dims,
             device=torch.device(device or "cpu"),
             data_transform=data_transform,
         )
+        self.dtype = resolved_dtype
         torch.manual_seed(seed)
         self.loc = None
         self.scale = None
@@ -41,12 +54,61 @@ class BaseTorchFlow(Flow):
     @flow.setter
     def flow(self, flow):
         self._flow = flow
-        self._flow.to(self.device)
+        self._flow.to(device=self.device, dtype=self.dtype)
         self._flow.compile()
 
     def fit(self, x) -> FlowHistory:
         raise NotImplementedError()
 
+    def save(self, h5_file, path="flow"):
+        """Save the weights of the flow to an HDF5 file."""
+        from ...utils import recursively_save_to_h5_file
+
+        flow_grp = h5_file.create_group(path)
+        # Save config
+        config = self.config_dict().copy()
+        data_transform = config.pop("data_transform", None)
+        dtype_value = config.get("dtype")
+        if dtype_value is None:
+            dtype_value = self.dtype
+        else:
+            dtype_value = resolve_dtype(dtype_value, torch)
+        config["dtype"] = encode_dtype(torch, dtype_value)
+        if data_transform is not None:
+            data_transform.save(flow_grp, "data_transform")
+        recursively_save_to_h5_file(flow_grp, "config", config)
+        # Save weights
+        weights_grp = flow_grp.create_group("weights")
+        for name, tensor in self._flow.state_dict().items():
+            weights_grp.create_dataset(name, data=tensor.cpu().numpy())
+
+    @classmethod
+    def load(self, h5_file, path="flow"):
+        """Load the weights of the flow from an HDF5 file."""
+        from ...utils import load_from_h5_file
+
+        flow_grp = h5_file[path]
+        # Load config
+        config = load_from_h5_file(flow_grp, "config")
+        config["dtype"] = decode_dtype(torch, config.get("dtype"))
+        if "data_transform" in flow_grp:
+            from ..transforms import BaseTransform
+
+            data_transform = BaseTransform.load(
+                flow_grp,
+                "data_transform",
+                strict=False,
+            )
+            config["data_transform"] = data_transform
+        obj = self(**config)
+        # Load weights
+        weights = {
+            name: torch.tensor(data[()])
+            for name, data in flow_grp["weights"].items()
+        }
+        obj._flow.load_state_dict(weights)
+        return obj
+
 
 class ZukoFlow(BaseTorchFlow):
     def __init__(
@@ -56,6 +118,7 @@ class ZukoFlow(BaseTorchFlow):
         data_transform=None,
         seed=1234,
         device: str = "cpu",
+        dtype=None,
         **kwargs,
     ):
         super().__init__(
@@ -63,6 +126,7 @@ class ZukoFlow(BaseTorchFlow):
             device=device,
             data_transform=data_transform,
             seed=seed,
+            dtype=dtype,
         )
 
         if isinstance(flow_class, str):
@@ -93,12 +157,10 @@ class ZukoFlow(BaseTorchFlow):
         from ...history import FlowHistory
 
         if not is_torch_array(x):
-            x = torch.tensor(
-                x, dtype=torch.get_default_dtype(), device=self.device
-            )
+            x = torch.tensor(x, dtype=self.dtype, device=self.device)
         else:
             x = torch.clone(x)
-            x = x.type(torch.get_default_dtype())
+            x = x.type(self.dtype)
             x = x.to(self.device)
         x_prime = self.fit_data_transform(x)
         indices = torch.randperm(x_prime.shape[0])
@@ -107,7 +169,7 @@ class ZukoFlow(BaseTorchFlow):
         n = x_prime.shape[0]
         x_train = torch.as_tensor(
             x_prime[: -int(validation_fraction * n)],
-            dtype=torch.get_default_dtype(),
+            dtype=self.dtype,
             device=self.device,
         )
 
@@ -117,13 +179,23 @@ class ZukoFlow(BaseTorchFlow):
         )
 
         if torch.isnan(x_train).any():
-            raise ValueError("Training data contains NaN values.")
+            dims_with_nan = (
+                torch.isnan(x_train).any(dim=0).nonzero(as_tuple=True)[0]
+            )
+            raise ValueError(
+                f"Training data contains NaN values in dimensions: {dims_with_nan.tolist()}"
+            )
         if not torch.isfinite(x_train).all():
-            raise ValueError("Training data contains infinite values.")
+            dims_with_inf = (
+                (~torch.isfinite(x_train)).any(dim=0).nonzero(as_tuple=True)[0]
+            )
+            raise ValueError(
+                f"Training data contains infinite values in dimensions: {dims_with_inf.tolist()}"
+            )
 
         x_val = torch.as_tensor(
             x_prime[-int(validation_fraction * n) :],
-            dtype=torch.get_default_dtype(),
+            dtype=self.dtype,
             device=self.device,
         )
         if torch.isnan(x_val).any():
@@ -207,18 +279,14 @@ class ZukoFlow(BaseTorchFlow):
         return xp.asarray(x)
 
     def log_prob(self, x, xp=torch_api):
-        x = torch.as_tensor(
-            x, dtype=torch.get_default_dtype(), device=self.device
-        )
+        x = torch.as_tensor(x, dtype=self.dtype, device=self.device)
         x_prime, log_abs_det_jacobian = self.rescale(x)
         return xp.asarray(
             self._flow().log_prob(x_prime) + log_abs_det_jacobian
         )
 
     def forward(self, x, xp=torch_api):
-        x = torch.as_tensor(
-            x, dtype=torch.get_default_dtype(), device=self.device
-        )
+        x = torch.as_tensor(x, dtype=self.dtype, device=self.device)
         x_prime, log_j_rescale = self.rescale(x)
         z, log_abs_det_jacobian = self._flow().transform.call_and_ladj(x_prime)
         if is_numpy_namespace(xp):
@@ -229,9 +297,7 @@ class ZukoFlow(BaseTorchFlow):
         return xp.asarray(z), xp.asarray(log_abs_det_jacobian + log_j_rescale)
 
     def inverse(self, z, xp=torch_api):
-        z = torch.as_tensor(
-            z, dtype=torch.get_default_dtype(), device=self.device
-        )
+        z = torch.as_tensor(z, dtype=self.dtype, device=self.device)
         with torch.no_grad():
             x_prime, log_abs_det_jacobian = (
                 self._flow().transform.inv.call_and_ladj(z)
@@ -253,6 +319,7 @@ class ZukoFlowMatching(ZukoFlow):
         seed=1234,
         device="cpu",
         eta: float = 1e-3,
+        dtype=None,
         **kwargs,
     ):
         kwargs.setdefault("hidden_features", 4 * [100])
@@ -262,6 +329,7 @@ class ZukoFlowMatching(ZukoFlow):
             device=device,
             data_transform=data_transform,
             flow_class="CNF",
+            dtype=dtype,
         )
         self.eta = eta
 

aspire/samplers/base.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Callable
+from typing import Any, Callable
 
 from ..flows.base import Flow
 from ..samples import Samples
@@ -36,6 +36,7 @@ class Sampler:
         dims: int,
         prior_flow: Flow,
         xp: Callable,
+        dtype: Any | str | None = None,
         parameters: list[str] | None = None,
         preconditioning_transform: Callable | None = None,
     ):
@@ -44,6 +45,7 @@ class Sampler:
         self.log_prior = log_prior
         self.dims = dims
         self.xp = xp
+        self.dtype = dtype
         self.parameters = parameters
         self.history = None
         self.n_likelihood_evaluations = 0

aspire/samplers/importance.py

@@ -8,7 +8,11 @@ class ImportanceSampler(Sampler):
     def sample(self, n_samples: int) -> Samples:
         x, log_q = self.prior_flow.sample_and_log_prob(n_samples)
         samples = Samples(
-            x, log_q=log_q, xp=self.xp, parameters=self.parameters
+            x,
+            log_q=log_q,
+            xp=self.xp,
+            parameters=self.parameters,
+            dtype=self.dtype,
         )
         samples.log_prior = samples.array_to_namespace(self.log_prior(samples))
         samples.log_likelihood = samples.array_to_namespace(

aspire/samplers/mcmc.py

@@ -15,7 +15,7 @@ class MCMCSampler(Sampler):
         while n_samples_drawn < n_samples:
             n_to_draw = n_samples - n_samples_drawn
             x, log_q = self.prior_flow.sample_and_log_prob(n_to_draw)
-            new_samples = Samples(x, xp=self.xp, log_q=log_q)
+            new_samples = Samples(x, xp=self.xp, log_q=log_q, dtype=self.dtype)
             new_samples.log_prior = new_samples.array_to_namespace(
                 self.log_prior(new_samples)
             )
@@ -44,7 +44,7 @@ class MCMCSampler(Sampler):
         Input samples are in the transformed space.
         """
         x, log_abs_det_jacobian = self.preconditioning_transform.inverse(z)
-        samples = Samples(x, xp=self.xp)
+        samples = Samples(x, xp=self.xp, dtype=self.dtype)
         samples.log_prior = self.log_prior(samples)
         samples.log_likelihood = self.log_likelihood(samples)
         log_prob = (
@@ -94,7 +94,9 @@ class Emcee(MCMCSampler):
         samples_evidence.log_likelihood = self.log_likelihood(samples_evidence)
         samples_evidence.compute_weights()
 
-        samples_mcmc = Samples(x, xp=self.xp, parameters=self.parameters)
+        samples_mcmc = Samples(
+            x, xp=self.xp, parameters=self.parameters, dtype=self.dtype
+        )
         samples_mcmc.log_prior = samples_mcmc.array_to_namespace(
             self.log_prior(samples_mcmc)
         )