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

aspire/utils.py

@@ -219,7 +219,7 @@ def logsumexp(x: Array, axis: int | None = None) -> Array:
 def to_numpy(x: Array, **kwargs) -> np.ndarray:
     """Convert an array to a numpy array.
 
-    This automatically moves the device to the CPU.
+    This automatically moves the array to the CPU.
 
     Parameters
     ----------
@@ -230,7 +230,7 @@ def to_numpy(x: Array, **kwargs) -> np.ndarray:
     """
     try:
         return np.asarray(to_device(x, "cpu"), **kwargs)
-    except ValueError:
+    except (ValueError, NotImplementedError):
         return np.asarray(x, **kwargs)
 
 
@@ -253,6 +253,135 @@ def asarray(x, xp: Any = None, **kwargs) -> Array:
         return xp.asarray(x, **kwargs)
 
 
+def resolve_dtype(dtype: Any | str | None, xp: Any) -> Any | None:
+    """Resolve a dtype specification into an XP-specific dtype.
+
+    Parameters
+    ----------
+    dtype : Any | str | None
+        The dtype specification. Can be None, a string, or a dtype-like object.
+    xp : module
+        The array API module that should interpret the dtype.
+
+    Returns
+    -------
+    Any | None
+        The resolved dtype object compatible with ``xp`` (or None if unspecified).
+    """
+    if dtype is None or xp is None:
+        return dtype
+
+    if isinstance(dtype, str):
+        dtype_name = _dtype_to_name(dtype)
+        if is_torch_namespace(xp):
+            resolved = getattr(xp, dtype_name, None)
+            if resolved is None:
+                raise ValueError(
+                    f"Unknown dtype '{dtype}' for namespace {xp.__name__}"
+                )
+            return resolved
+        try:
+            return xp.dtype(dtype_name)
+        except (AttributeError, TypeError, ValueError):
+            resolved = getattr(xp, dtype_name, None)
+            if resolved is not None:
+                return resolved
+            raise ValueError(
+                f"Unknown dtype '{dtype}' for namespace {getattr(xp, '__name__', xp)}"
+            )
+
+    if is_torch_namespace(xp):
+        return dtype
+
+    try:
+        return xp.dtype(dtype)
+    except (AttributeError, TypeError, ValueError):
+        return dtype
+
+
+def _dtype_to_name(dtype: Any | str | None) -> str | None:
+    """Extract a canonical (lowercase) name for a dtype-like object."""
+    if dtype is None:
+        return None
+    if isinstance(dtype, str):
+        name = dtype
+    elif hasattr(dtype, "name") and getattr(dtype, "name"):
+        name = dtype.name
+    elif hasattr(dtype, "__name__"):
+        name = dtype.__name__
+    else:
+        text = str(dtype)
+        if text.startswith("<class '") and text.endswith("'>"):
+            text = text.split("'")[1]
+        if text.startswith("dtype(") and text.endswith(")"):
+            inner = text[6:-1].strip("'\" ")
+            text = inner or text
+        name = text
+    name = name.split(".")[-1]
+    return name.strip(" '\"<>").lower()
+
+
+def convert_dtype(
+    dtype: Any | str | None,
+    target_xp: Any,
+    *,
+    source_xp: Any | None = None,
+) -> Any | None:
+    """Convert a dtype between array API namespaces.
+
+    Parameters
+    ----------
+    dtype : Any | str | None
+        The dtype to convert. Can be a dtype object, string, or None.
+    target_xp : module
+        The target array API namespace to convert the dtype into.
+    source_xp : module, optional
+        The source namespace of the dtype. Provided for API symmetry and future
+        use; currently unused but accepted.
+
+    Returns
+    -------
+    Any | None
+        The dtype object compatible with ``target_xp`` (or None if ``dtype`` is None).
+    """
+    if dtype is None:
+        return None
+    if target_xp is None:
+        raise ValueError("target_xp must be provided to convert dtype.")
+
+    target_name = getattr(target_xp, "__name__", "")
+    dtype_module = getattr(dtype, "__module__", "")
+    if dtype_module.startswith(target_name):
+        return dtype
+    if is_torch_namespace(target_xp) and str(dtype).startswith("torch."):
+        return dtype
+
+    name = _dtype_to_name(dtype)
+    if not name:
+        raise ValueError(f"Could not infer dtype name from {dtype!r}")
+
+    candidates = dict.fromkeys(
+        [name, name.lower(), name.upper(), name.capitalize()]
+    )
+    last_error: Exception | None = None
+    for candidate in candidates:
+        try:
+            return resolve_dtype(candidate, target_xp)
+        except ValueError as exc:
+            last_error = exc
+
+    # Fallback to direct attribute lookup
+    attr = getattr(target_xp, name, None) or getattr(
+        target_xp, name.lower(), None
+    )
+    if attr is not None:
+        return attr
+
+    raise ValueError(
+        f"Unable to convert dtype {dtype!r} to namespace {target_name}"
+    ) from last_error
+
+
 def copy_array(x, xp: Any = None) -> Array:
     """Copy an array based on the array API being used.
 
@@ -321,6 +450,59 @@ def disable_gradients(xp, inference: bool = True):
         yield
 
 
+def encode_dtype(xp, dtype):
+    """Encode a dtype for storage in an HDF5 file.
+
+    Parameters
+    ----------
+    xp : module
+        The array API module to use.
+    dtype : dtype
+        The dtype to encode.
+
+    Returns
+    -------
+    str
+        The encoded dtype.
+    """
+    if dtype is None:
+        return None
+    return {
+        "__dtype__": True,
+        "xp": xp.__name__,
+        "dtype": _dtype_to_name(dtype),
+    }
+
+
+def decode_dtype(xp, encoded_dtype):
+    """Decode a dtype from an HDF5 file.
+
+    Parameters
+    ----------
+    xp : module
+        The array API module to use.
+    encoded_dtype : dict
+        The encoded dtype.
+
+    Returns
+    -------
+    dtype
+        The decoded dtype.
+    """
+    if isinstance(encoded_dtype, dict) and encoded_dtype.get("__dtype__"):
+        if encoded_dtype["xp"] != xp.__name__:
+            raise ValueError(
+                f"Encoded dtype xp {encoded_dtype['xp']} does not match "
+                f"current xp {xp.__name__}"
+            )
+        if is_torch_namespace(xp):
+            return getattr(xp, encoded_dtype["dtype"].split(".")[-1])
+        else:
+            return xp.dtype(encoded_dtype["dtype"].split(".")[-1])
+    else:
+        return encoded_dtype
+
+
 def encode_for_hdf5(value: Any) -> Any:
     """Encode a value for storage in an HDF5 file.
 
@@ -328,6 +510,8 @@ def encode_for_hdf5(value: Any) -> Any:
     - None is replaced with "__none__"
     - Empty dictionaries are replaced with "__empty_dict__"
     """
+    if is_jax_array(value) or is_torch_array(value):
+        return to_numpy(value)
     if isinstance(value, CallHistory):
         return value.to_dict(list_to_dict=True)
     if isinstance(value, np.ndarray):
@@ -335,6 +519,9 @@ def encode_for_hdf5(value: Any) -> Any:
     if isinstance(value, (int, float, str)):
         return value
     if isinstance(value, (list, tuple)):
+        if all(isinstance(v, str) for v in value):
+            dt = h5py.string_dtype(encoding="utf-8")
+            return np.array(value, dtype=dt)
         return [encode_for_hdf5(v) for v in value]
     if isinstance(value, set):
         return {encode_for_hdf5(v) for v in value}
@@ -345,23 +532,89 @@ def encode_for_hdf5(value: Any) -> Any:
             return {k: encode_for_hdf5(v) for k, v in value.items()}
     if value is None:
         return "__none__"
+
     return value
 
 
-def recursively_save_to_h5_file(h5_file, path, dictionary):
-    """Recursively save a dictionary to an HDF5 file."""
-    for key, value in dictionary.items():
-        if isinstance(value, dict):
-            recursively_save_to_h5_file(h5_file, f"{path}/{key}", value)
-        else:
+def decode_from_hdf5(value: Any) -> Any:
+    """Decode a value loaded from an HDF5 file, reversing encode_for_hdf5."""
+    if isinstance(value, bytes):  # HDF5 may store strings as bytes
+        value = value.decode("utf-8")
+
+    if isinstance(value, str):
+        if value == "__none__":
+            return None
+        if value == "__empty_dict__":
+            return {}
+
+    if isinstance(value, np.ndarray):
+        # Try to collapse 0-D arrays into scalars
+        if value.shape == ():
+            return value.item()
+        if value.dtype.kind in {"S", "O"}:
             try:
-                h5_file.create_dataset(
-                    f"{path}/{key}", data=encode_for_hdf5(value)
-                )
-            except TypeError as error:
-                raise RuntimeError(
-                    f"Cannot save key {key} with value {value} to HDF5 file."
-                ) from error
+                return value.astype(str).tolist()
+            except Exception:
+                # fallback: leave as ndarray
+                return value
+        return value
+
+    if isinstance(value, list):
+        return [decode_from_hdf5(v) for v in value]
+    if isinstance(value, tuple):
+        return tuple(decode_from_hdf5(v) for v in value)
+    if isinstance(value, set):
+        return {decode_from_hdf5(v) for v in value}
+    if isinstance(value, dict):
+        return {
+            k.decode("utf-8"): decode_from_hdf5(v) for k, v in value.items()
+        }
+
+    # Fallback for ints, floats, strs, etc.
+    return value
+
+
+def recursively_save_to_h5_file(h5_file, path, dictionary):
+    """Save a dictionary to an HDF5 file with flattened keys under a given group path."""
+    # Ensure the group exists (or open it if already present)
+    group = h5_file.require_group(path)
+
+    def _save_flattened(g, prefix, d):
+        for key, value in d.items():
+            full_key = f"{prefix}.{key}" if prefix else key
+            if isinstance(value, dict):
+                _save_flattened(g, full_key, value)
+            else:
+                try:
+                    g.create_dataset(full_key, data=encode_for_hdf5(value))
+                except TypeError as error:
+                    try:
+                        # Try saving as a string
+                        dt = h5py.string_dtype(encoding="utf-8")
+                        g.create_dataset(
+                            full_key, data=np.array(str(value), dtype=dt)
+                        )
+                    except Exception:
+                        raise RuntimeError(
+                            f"Cannot save key {full_key} with value {value} to HDF5 file."
+                        ) from error
+
+    _save_flattened(group, "", dictionary)
+
+
+def load_from_h5_file(h5_file, path):
+    """Load a flattened dictionary from an HDF5 group and rebuild nesting."""
+    group = h5_file[path]
+    result = {}
+
+    for key, dataset in group.items():
+        parts = key.split(".")
+        d = result
+        for part in parts[:-1]:
+            d = d.setdefault(part, {})
+        d[parts[-1]] = decode_from_hdf5(dataset[()])
+
+    return result
 
 
 def get_package_version(package_name: str) -> str:
@@ -394,7 +647,15 @@ class AspireFile(h5py.File):
     def _set_aspire_metadata(self):
         from . import __version__ as aspire_version
 
-        self.attrs["aspire_version"] = aspire_version
+        if self.mode in {"w", "w-", "a", "r+"}:
+            self.attrs["aspire_version"] = aspire_version
+        else:
+            aspire_version = self.attrs.get("aspire_version", "unknown")
+            if aspire_version != "unknown":
+                logger.warning(
+                    f"Opened Aspire file created with version {aspire_version}. "
+                    f"Current version is {aspire_version}."
+                )
 
 
 def update_at_indices(x: Array, slc: Array, y: Array) -> Array:
@@ -403,6 +664,14 @@ def update_at_indices(x: Array, slc: Array, y: Array) -> Array:
     This is a workaround for the fact that array API does not support
     advanced indexing with all backends.
 
+    Examples
+    --------
+    >>> x = xp.array([[1, 2], [3, 4], [5, 6]])
+    >>> update_at_indices(x, (slice(None), 0), xp.array([10, 20, 30]))
+    [[10  2]
+     [20  4]
+     [30  6]]
+
     Parameters
     ----------
     x : Array

{aspire_inference-0.1.0a5.dist-info → aspire_inference-0.1.0a6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a5
+Version: 0.1.0a6
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -33,6 +33,7 @@ Requires-Dist: blackjax; extra == "blackjax"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-requires; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
 Dynamic: license-file
 
 # aspire: Accelerated Sequential Posterior Inference via REuse

aspire_inference-0.1.0a6.dist-info/RECORD

@@ -0,0 +1,28 @@
+aspire/__init__.py,sha256=45R0xWaLg0aJEPK5zoTK0aIek0KOwpHwQWS1jLCDhIE,365
+aspire/aspire.py,sha256=M5o-QxLthE_5daa1trgUfWxPz-g4rmpEUKimKosw4lw,17400
+aspire/history.py,sha256=l_j-riZKbTWK7Wz9zvvD_mTk9psNCKItiveYhr_pYv8,4313
+aspire/plot.py,sha256=oXwUDOb_953_ADm2KLk41JIfpE3JeiiQiSYKvUVwLqw,1423
+aspire/samples.py,sha256=lUn3cQdnN5gSHREuamZZteecdDIuasaWfcDXdVAvfZA,18662
+aspire/transforms.py,sha256=XMbf5MxK49elQeKDsmFraHN-0JeO1AciljdTk7k2ujk,24928
+aspire/utils.py,sha256=pj8O0chqfP6VS8bpW0wCw8W0P5JNQKvWRz1Rg9AYIhg,22525
+aspire/flows/__init__.py,sha256=3gGXF4HziMlZSmcEdJ_uHtrP-QEC6RXvylm4vtM-Xnk,1306
+aspire/flows/base.py,sha256=5UWKAiXDXLJ6Sg6a380ajLrGFaZSQyOnFEihQiiA4ko,2237
+aspire/flows/jax/__init__.py,sha256=7cmiY_MbEC8RDA8Cmi8HVnNJm0sqFKlBsDethdsy5lA,52
+aspire/flows/jax/flows.py,sha256=1HnVgQ1GUXNcvxiZqEV19H2QI9Th5bWX_QbNfGaUhuA,6625
+aspire/flows/jax/utils.py,sha256=5T6UrgpARG9VywC9qmTl45LjyZWuEdkW3XUladE6xJE,1518
+aspire/flows/torch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+aspire/flows/torch/flows.py,sha256=0_YkiMT49QolyQnEFsh28tfKLnURVF0Z6aTnaWLIUDI,11672
+aspire/samplers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+aspire/samplers/base.py,sha256=8slvgOBnacUrHXCVDAqo-3IZ_LB7-dS8wdMP55MI43Y,2907
+aspire/samplers/importance.py,sha256=opn_jY-V8snUz0JztLBtnaTT3WfrZ5OSpHBV5WAuM3M,676
+aspire/samplers/mcmc.py,sha256=ihHgzqvSoy1oxdFBjyqNUbCuRX7CqWjlshCUZcgEL5E,5151
+aspire/samplers/smc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+aspire/samplers/smc/base.py,sha256=66f_ORUvcKRqMIW35qjhUc-c0PFuY87lJa91MpSaTZI,10729
+aspire/samplers/smc/blackjax.py,sha256=4L4kgRKlaWl-knTWXXzdJTh-zZBh5BTpy5GaLDzT8Sc,11803
+aspire/samplers/smc/emcee.py,sha256=Wm0vvAlCcRhJMBt7_fU2ZnjDb8SN8jgUOTXLzNstRpA,2516
+aspire/samplers/smc/minipcn.py,sha256=ju1gcgyKHjodLEACPdL3eXA9ai8ZJ9_LwitD_Gmf1Rc,2765
+aspire_inference-0.1.0a6.dist-info/licenses/LICENSE,sha256=DN-eRtBfS9dZyT0Ds0Mdn2Y4nb-ZQ7h71vpASYBm5k4,1076
+aspire_inference-0.1.0a6.dist-info/METADATA,sha256=--Q4vjeyHU7-TcgV3HicQ9wxMuO_Vi1pba5wDUK2oD0,1617
+aspire_inference-0.1.0a6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aspire_inference-0.1.0a6.dist-info/top_level.txt,sha256=9FRIYEl2xwVFG7jSOBHsElHQ0y3_4fq01Cf4_OyMQn8,7
+aspire_inference-0.1.0a6.dist-info/RECORD,,

aspire_inference-0.1.0a5.dist-info/RECORD

@@ -1,28 +0,0 @@
-aspire/__init__.py,sha256=45R0xWaLg0aJEPK5zoTK0aIek0KOwpHwQWS1jLCDhIE,365
-aspire/aspire.py,sha256=AEkFUuOCF4F_iXUqRNst_4mucxozYRK4fG4V2wGrT4Q,15762
-aspire/history.py,sha256=l_j-riZKbTWK7Wz9zvvD_mTk9psNCKItiveYhr_pYv8,4313
-aspire/plot.py,sha256=oXwUDOb_953_ADm2KLk41JIfpE3JeiiQiSYKvUVwLqw,1423
-aspire/samples.py,sha256=mVyoHNcTPgqJXjKNeF2oA-KNV1asriWjToWX-T70E5k,16729
-aspire/transforms.py,sha256=fg2_UELJWjJ6gnqQi6X7s1CgKBhQ5hP7Ipil3tTjCeg,16566
-aspire/utils.py,sha256=fQeLMauCN3vAogKbVTVg9jfjW7nTEFi7V6Ot-BYNfxE,14301
-aspire/flows/__init__.py,sha256=3gGXF4HziMlZSmcEdJ_uHtrP-QEC6RXvylm4vtM-Xnk,1306
-aspire/flows/base.py,sha256=oTw2ZkxCsA5RZhnMuIu9M-2FPHvQG2TGFIEJZVK4a2g,1140
-aspire/flows/jax/__init__.py,sha256=7cmiY_MbEC8RDA8Cmi8HVnNJm0sqFKlBsDethdsy5lA,52
-aspire/flows/jax/flows.py,sha256=jZ93fnc7U7ZhZLVixGUTwyeDb6Vz0UWpYkkVHwirNug,2896
-aspire/flows/jax/utils.py,sha256=UlvXOOqC5fNsmVUnU4LSksliq7pLRm9NhOu0ZvVHqgc,1455
-aspire/flows/torch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aspire/flows/torch/flows.py,sha256=ZNnShj-FMr56ZbcY06fNQa0epolzMZBd8ok2TzKGZ8E,8996
-aspire/samplers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aspire/samplers/base.py,sha256=BZ5nY_wtvuOIpTaJWUYZflCFXPTDk24xB-qLirIn9qE,2835
-aspire/samplers/importance.py,sha256=3mY6JEqzdunHwAF6l3-CN-tBEdC_8J0LkhxD57DyHoY,609
-aspire/samplers/mcmc.py,sha256=uuCjHZeey5mqjntnYaisNytYBazIc0xuvRcXPHwtg0Y,5075
-aspire/samplers/smc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aspire/samplers/smc/base.py,sha256=GePA6tm8Dno_AjCeNuRX3KOaKnoKSFHSRAb-QWx9wJE,10531
-aspire/samplers/smc/blackjax.py,sha256=9w1ORzWTT1viwp99_ttLxnNMdgTO-VqAzsf-NhgG9vY,11722
-aspire/samplers/smc/emcee.py,sha256=ZXXyN2l1Bz5ZsCPEcswg-Kakiw41nNa2jEW1N8zGjuc,2498
-aspire/samplers/smc/minipcn.py,sha256=ZjeP4iHFR67G8WKEfMe0b1McrtPgQMNHyyy4vRx6WNE,2747
-aspire_inference-0.1.0a5.dist-info/licenses/LICENSE,sha256=DN-eRtBfS9dZyT0Ds0Mdn2Y4nb-ZQ7h71vpASYBm5k4,1076
-aspire_inference-0.1.0a5.dist-info/METADATA,sha256=Vq1jmMMrg6taHSFqJOCZJUex7wdakeoL7K6844VSlDs,1574
-aspire_inference-0.1.0a5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aspire_inference-0.1.0a5.dist-info/top_level.txt,sha256=9FRIYEl2xwVFG7jSOBHsElHQ0y3_4fq01Cf4_OyMQn8,7
-aspire_inference-0.1.0a5.dist-info/RECORD,,