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

aspire/transforms.py

@@ -1,12 +1,17 @@
+import importlib
 import logging
 import math
-from typing import Any
+from typing import Any, Callable
 
+import h5py
 from array_api_compat import device as get_device
 from array_api_compat import is_torch_namespace
+from array_api_compat.common._typing import Array
 
 from .flows import get_flow_wrapper
 from .utils import (
+    asarray,
+    convert_dtype,
     copy_array,
     logit,
     sigmoid,
@@ -32,6 +37,10 @@ class BaseTransform:
         self.xp = xp
         if is_torch_namespace(self.xp) and dtype is None:
             dtype = self.xp.get_default_dtype()
+        elif isinstance(dtype, str):
+            from .utils import resolve_dtype
+
+            dtype = resolve_dtype(dtype, self.xp)
         self.dtype = dtype
 
     def fit(self, x):
@@ -44,6 +53,74 @@ class BaseTransform:
     def inverse(self, y):
         raise NotImplementedError("Subclasses must implement inverse method.")
 
+    def config_dict(self):
+        """Return the configuration of the transform as a dictionary."""
+        return {
+            "xp": self.xp.__name__,
+            "dtype": str(self.dtype) if self.dtype else None,
+        }
+
+    def save(self, h5_file: h5py.File, path: str = "data_transform"):
+        """Save config + any fitted state into an HDF5 file."""
+        from .utils import encode_dtype, recursively_save_to_h5_file
+
+        # store class name for reconstruction
+        grp = h5_file.create_group(path)
+        grp.attrs["class"] = self.__class__.__name__
+        # store config as JSON
+        config = self.config_dict()
+        config["dtype"] = encode_dtype(self.xp, config["dtype"])
+        recursively_save_to_h5_file(grp, "config", config)
+        # store any fitted arrays
+        self._save_state(grp)
+
+    @classmethod
+    def load(
+        cls,
+        h5_file: h5py.File,
+        path: str = "data_transform",
+        strict: bool = False,
+    ):
+        """Reconstruct transform from file.
+
+        Parameters
+        ----------
+        h5_file : h5py.File
+            The HDF5 file to load from.
+        path : str, optional
+            The path in the HDF5 file where the transform is stored.
+        strict : bool, optional
+            If True, raise an error if the class in the file does not match cls.
+            If False, load the class specified in the file. Default is False.
+        """
+        from .utils import decode_dtype, load_from_h5_file
+
+        grp = h5_file[path]
+        class_name = grp.attrs["class"]
+        if class_name != cls.__name__:
+            if strict:
+                raise ValueError(
+                    f"Expected class {cls.__name__}, got {class_name}."
+                )
+            else:
+                cls = getattr(importlib.import_module(__name__), class_name)
+                logger.info(
+                    f"Loading class {class_name} instead of {cls.__name__}."
+                )
+
+        config = load_from_h5_file(grp, "config")
+        config["xp"] = importlib.import_module(config["xp"])
+        config["dtype"] = decode_dtype(config["xp"], config["dtype"])
+        obj = cls(**config)
+        obj._load_state(grp)
+        return obj
+
+    def _save_state(self, h5_file: h5py.File):
+        pass
+
+    def _load_state(self, h5_file: h5py.File):
+        pass
+
 
 class IdentityTransform(BaseTransform):
     """Identity transform that does nothing to the data."""
@@ -138,6 +215,7 @@ class CompositeTransform(BaseTransform):
                 lower=lower_bounds[self.periodic_mask],
                 upper=upper_bounds[self.periodic_mask],
                 xp=self.xp,
+                dtype=self.dtype,
             )
         if self.bounded_parameters:
             logger.info(f"Bounded parameters: {self.bounded_parameters}")
@@ -158,11 +236,14 @@ class CompositeTransform(BaseTransform):
                 upper=upper_bounds[self.bounded_mask],
                 xp=self.xp,
                 eps=self.eps,
+                dtype=self.dtype,
             )
 
         if self.affine_transform:
             logger.info(f"Affine transform applied to: {self.parameters}")
-            self._affine_transform = AffineTransform(xp=self.xp)
+            self._affine_transform = AffineTransform(
+                xp=self.xp, dtype=self.dtype
+            )
         else:
             self._affine_transform = None
 
@@ -238,7 +319,13 @@ class CompositeTransform(BaseTransform):
 
         return x, log_abs_det_jacobian
 
-    def new_instance(self, xp=None):
+    def new_instance(self, xp=None, dtype: Any = None):
+        if xp is None:
+            xp = self.xp
+        if dtype is None:
+            dtype = self.dtype
+        dtype = convert_dtype(dtype, xp)
+
         return self.__class__(
             parameters=self.parameters,
             periodic_parameters=self.periodic_parameters,
@@ -248,8 +335,31 @@ class CompositeTransform(BaseTransform):
             device=self.device,
             xp=xp or self.xp,
             eps=self.eps,
+            dtype=dtype,
         )
 
+    def _save_state(self, h5_file):
+        if self.affine_transform:
+            affine_grp = h5_file.create_group("affine_transform")
+            self._affine_transform._save_state(affine_grp)
+
+    def _load_state(self, h5_file):
+        if self.affine_transform:
+            affine_grp = h5_file["affine_transform"]
+            self._affine_transform._load_state(affine_grp)
+
+    def config_dict(self):
+        return super().config_dict() | {
+            "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,
+            "affine_transform": self.affine_transform,
+            "eps": self.eps,
+            "device": self.device,
+        }
+
 
 class FlowTransform(CompositeTransform):
     """Subclass of CompositeTransform that uses a Flow for transformations.
@@ -293,6 +403,13 @@ class FlowTransform(CompositeTransform):
             eps=self.eps,
         )
 
+    def config_dict(self):
+        cfg = super().config_dict()
+        cfg.pop(
+            "periodic_parameters", None
+        )  # Remove periodic_parameters from config
+        return cfg
+
 
 class PeriodicTransform(BaseTransform):
     name: str = "periodic"
@@ -300,9 +417,9 @@ class PeriodicTransform(BaseTransform):
 
     def __init__(self, lower, upper, xp, dtype=None):
         super().__init__(xp=xp, dtype=dtype)
-        self.lower = xp.asarray(lower, dtype=dtype)
-        self.upper = xp.asarray(upper, dtype=dtype)
-        self._width = upper - lower
+        self.lower = xp.asarray(lower, dtype=self.dtype)
+        self.upper = xp.asarray(upper, dtype=self.dtype)
+        self._width = self.upper - self.lower
         self._shift = None
 
     def fit(self, x):
@@ -316,80 +433,190 @@ class PeriodicTransform(BaseTransform):
         x = self.lower + (y - self.lower) % self._width
         return x, self.xp.zeros(x.shape[0], device=get_device(x))
 
+    def config_dict(self):
+        return super().config_dict() | {
+            "lower": self.lower.tolist(),
+            "upper": self.upper.tolist(),
+        }
+
+
+class BoundedTransform(BaseTransform):
+    """Base class for bounded transforms.
+
+    Maps from [lower, upper] to [0, 1] and vice versa using a linear scaling.
+    If the interval [lower, upper] is too small, it will shift by the midpoint.
+
+    Must be subclassed to implement specific transforms (e.g., Probit, Logit).
+
+    Parameters
+    ----------
+    lower : Array
+        The lower bound of the interval.
+    upper : Array
+        The upper bound of the interval.
+    xp : Callable
+        The array API namespace to use (e.g., numpy, torch).
+    dtype : Any, optional
+        The data type to use for the transform. If not provided, defaults to
+        the default dtype of the array API namespace if available.
+    """
+
+    name: str = "bounded"
+    requires_prior_bounds: bool = True
+
+    def __init__(
+        self, lower: Array, upper: Array, xp: Callable, dtype: Any = None
+    ):
+        super().__init__(xp=xp, dtype=dtype)
+        self.lower = xp.atleast_1d(xp.asarray(lower, dtype=self.dtype))
+        self.upper = xp.atleast_1d(xp.asarray(upper, dtype=self.dtype))
+
+        self.interval_check(self.lower, self.upper)
+
+        self._denom = self.upper - self.lower
+        self._scale_log_abs_det_jacobian = -xp.log(self._denom).sum()
+
+    def to_unit_interval(self, x: Array) -> tuple[Array, Array]:
+        """Map from [lower, upper] to [0, 1].
+
+        Parameters
+        ----------
+        x : Array
+            The input array to be mapped.
+
+        Returns
+        -------
+        tuple[Array, Array]
+            A tuple containing the mapped array and the log absolute determinant Jacobian.
+        """
+        y = (x - self.lower) / self._denom
+        log_j = self._scale_log_abs_det_jacobian * self.xp.ones(
+            y.shape[0], device=get_device(y)
+        )
+        return y, log_j
+
+    def from_unit_interval(self, y: Array) -> tuple[Array, Array]:
+        """Map from [0, 1] to [lower, upper].
+
+        Parameters
+        ----------
+        y : Array
+            The input array to be mapped.
+
+        Returns
+        -------
+        tuple[Array, Array]
+            A tuple containing the mapped array and the log absolute determinant Jacobian.
+        """
+        x = self._denom * y + self.lower
+        log_j = -self._scale_log_abs_det_jacobian * self.xp.ones(
+            x.shape[0], device=get_device(x)
+        )
+        return x, log_j
 
-class ProbitTransform(BaseTransform):
+    def interval_check(self, lower: Array, upper: Array) -> bool:
+        """Check if the interval [lower, upper] is too small"""
+        if any((upper - lower) == 0.0):
+            raise ValueError(
+                f"Current floating precision ({self.dtype}) is too small for specified parameter ranges"
+            )
+
+    def fit(self, x):
+        return self.forward(x)[0]
+
+    def forward(self, x):
+        raise NotImplementedError("Subclasses must implement forward method.")
+
+    def inverse(self, y):
+        raise NotImplementedError("Subclasses must implement inverse method.")
+
+    def config_dict(self):
+        return super().config_dict() | {
+            "lower": self.lower.tolist(),
+            "upper": self.upper.tolist(),
+        }
+
+
+class ProbitTransform(BoundedTransform):
     name: str = "probit"
     requires_prior_bounds: bool = True
 
     def __init__(self, lower, upper, xp, eps=1e-6, dtype=None):
-        self.lower = xp.asarray(lower, dtype=dtype)
-        self.upper = xp.asarray(upper, dtype=dtype)
-        self._scale_log_abs_det_jacobian = -xp.log(upper - lower).sum()
+        super().__init__(xp=xp, dtype=dtype, lower=lower, upper=upper)
         self.eps = eps
-        self.xp = xp
 
-    def fit(self, x):
+    def fit(self, x: Array) -> Array:
         return self.forward(x)[0]
 
-    def forward(self, x):
+    def forward(self, x: Array) -> tuple[Array, Array]:
         from scipy.special import erfinv
 
-        y = (x - self.lower) / (self.upper - self.lower)
+        y, log_j_unit = self.to_unit_interval(x)
         y = self.xp.clip(y, self.eps, 1.0 - self.eps)
         y = erfinv(2 * y - 1) * math.sqrt(2)
-        log_abs_det_jacobian = (
-            0.5 * (math.log(2 * math.pi) + y**2).sum(-1)
-            + self._scale_log_abs_det_jacobian
-        )
+        log_abs_det_jacobian = 0.5 * (math.log(2 * math.pi) + y**2).sum(-1)
+        log_abs_det_jacobian = log_abs_det_jacobian + log_j_unit
         return y, log_abs_det_jacobian
 
-    def inverse(self, y):
+    def inverse(self, y: Array) -> tuple[Array, Array]:
         from scipy.special import erf
 
-        log_abs_det_jacobian = (
-            -(0.5 * (math.log(2 * math.pi) + y**2)).sum(-1)
-            - self._scale_log_abs_det_jacobian
-        )
+        log_abs_det_jacobian = -(0.5 * (math.log(2 * math.pi) + y**2)).sum(-1)
         x = 0.5 * (1 + erf(y / math.sqrt(2)))
-        x = (self.upper - self.lower) * x + self.lower
+        x, log_j_unit = self.from_unit_interval(x)
+        log_abs_det_jacobian = log_abs_det_jacobian + log_j_unit
         return x, log_abs_det_jacobian
 
+    def config_dict(self):
+        return super().config_dict() | {
+            "eps": self.eps,
+        }
+
 
-class LogitTransform(BaseTransform):
+class LogitTransform(BoundedTransform):
     name: str = "logit"
     requires_prior_bounds: bool = True
 
-    def __init__(self, lower, upper, xp, eps=1e-6, dtype=None):
-        self.lower = xp.asarray(lower, dtype=dtype)
-        self.upper = xp.asarray(upper, dtype=dtype)
-        self._scale_log_abs_det_jacobian = -xp.log(upper - lower).sum()
+    def __init__(
+        self,
+        lower: Array,
+        upper: Array,
+        xp: Callable,
+        eps: float = 1e-6,
+        dtype: Any = None,
+    ):
+        super().__init__(xp=xp, dtype=dtype, lower=lower, upper=upper)
         self.eps = eps
-        self.xp = xp
 
-    def fit(self, x):
+    def fit(self, x: Array) -> Array:
         return self.forward(x)[0]
 
-    def forward(self, x):
-        y = (x - self.lower) / (self.upper - self.lower)
+    def forward(self, x: Array) -> tuple[Array, Array]:
+        y, log_j_unit = self.to_unit_interval(x)
         y, log_abs_det_jacobian = logit(y, eps=self.eps)
-        log_abs_det_jacobian += self._scale_log_abs_det_jacobian
+        log_abs_det_jacobian = log_abs_det_jacobian + log_j_unit
         return y, log_abs_det_jacobian
 
-    def inverse(self, y):
+    def inverse(self, y: Array) -> tuple[Array, Array]:
         x, log_abs_det_jacobian = sigmoid(y)
-        log_abs_det_jacobian -= self._scale_log_abs_det_jacobian
-        x = (self.upper - self.lower) * x + self.lower
+        x, log_j_unit = self.from_unit_interval(x)
+        log_abs_det_jacobian = log_abs_det_jacobian + log_j_unit
         return x, log_abs_det_jacobian
 
+    def config_dict(self) -> dict[str, Any]:
+        return super().config_dict() | {
+            "eps": self.eps,
+        }
+
 
 class AffineTransform(BaseTransform):
     name: str = "affine"
     requires_prior_bounds: bool = False
 
-    def __init__(self, xp):
+    def __init__(self, xp, dtype=None):
+        super().__init__(xp=xp, dtype=dtype)
         self._mean = None
         self._std = None
-        self.xp = xp
 
     def fit(self, x):
         self._mean = x.mean(0)
@@ -409,6 +636,18 @@ class AffineTransform(BaseTransform):
             y.shape[0], device=get_device(y)
         )
 
+    def config_dict(self):
+        return super().config_dict()
+
+    def _save_state(self, h5_file):
+        h5_file.create_dataset("mean", data=self._mean)
+        h5_file.create_dataset("std", data=self._std)
+
+    def _load_state(self, h5_file):
+        self._mean = asarray(h5_file["mean"][()], xp=self.xp)
+        self._std = asarray(h5_file["std"][()], xp=self.xp)
+        self.log_abs_det_jacobian = -self.xp.log(self.xp.abs(self._std)).sum()
+
 
 class FlowPreconditioningTransform(BaseTransform):
     def __init__(
@@ -440,8 +679,10 @@ class FlowPreconditioningTransform(BaseTransform):
         self.device = device or "cpu"
         self.flow_backend = flow_backend
         self.flow_matching = flow_matching
-        self.flow_kwargs = flow_kwargs or {}
-        self.fit_kwargs = fit_kwargs or {}
+        self.flow_kwargs = dict(flow_kwargs or {})
+        if dtype is not None:
+            self.flow_kwargs.setdefault("dtype", dtype)
+        self.fit_kwargs = dict(fit_kwargs or {})
 
         FlowClass = get_flow_wrapper(
             backend=flow_backend, flow_matching=flow_matching
@@ -479,7 +720,14 @@ class FlowPreconditioningTransform(BaseTransform):
     def inverse(self, y):
         return self.flow.inverse(y, xp=self.xp)
 
-    def new_instance(self, xp=None):
+    def new_instance(self, xp=None, dtype: Any = None):
+        if xp is None:
+            xp = self.xp
+        if dtype is None:
+            dtype = self.dtype
+
+        dtype = convert_dtype(dtype, xp)
+
         return self.__class__(
             parameters=self.parameters,
             periodic_parameters=self.periodic_parameters,
@@ -488,11 +736,16 @@ class FlowPreconditioningTransform(BaseTransform):
             bounded_transform=self.bounded_transform,
             affine_transform=self.affine_transform,
             device=self.device,
-            xp=xp or self.xp,
+            xp=xp,
             eps=self.eps,
-            dtype=self.dtype,
+            dtype=dtype,
             flow_backend=self.flow_backend,
             flow_matching=self.flow_matching,
             flow_kwargs=self.flow_kwargs,
             fit_kwargs=self.fit_kwargs,
         )
+
+    def save(self, h5_file, path="data_transform"):
+        raise NotImplementedError(
+            "FlowPreconditioningTransform does not support save method yet."
+        )