ossify 0.2.4__tar.gz → 0.2.6__tar.gz

{ossify-0.2.4 → ossify-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ossify
-Version: 0.2.4
+Version: 0.2.6
 Summary: Mesh and skeleton analysis
 Author-email: Casey Schneider-Mizell <caseysm@gmail.com>
 License-File: LICENSE

{ossify-0.2.4 → ossify-0.2.6}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 allow-direct-references = true
 [project]
 name = "ossify"
-version = "0.2.4"
+version = "0.2.6"
 description = "Mesh and skeleton analysis"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -95,7 +95,7 @@ default-groups = ["dev", "docs", "lint", "profile", "viz"]
 
 
 [tool.bumpversion]
-current_version = "0.2.4"
+current_version = "0.2.6"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 regex = false

{ossify-0.2.4 → ossify-0.2.6}/src/ossify/__init__.py

@@ -5,7 +5,7 @@ from .base import *
 from .file_io import *
 from .translate import *
 
-__version__ = "0.2.4"
+__version__ = "0.2.6"
 
 
 def __getattr__(name):

{ossify-0.2.4 → ossify-0.2.6}/src/ossify/compartments.py

@@ -26,6 +26,7 @@ from dataclasses import dataclass, field
 from typing import Callable, List, Literal, Optional, Union
 
 import numpy as np
+import orjson
 import pandas as pd
 from scipy import sparse
 
@@ -38,6 +39,7 @@ from .algorithms import (
 from .base import Cell
 from .structured_prediction import (
     TransitionSchema,
+    _ossify_version,
     absorb_small_compartments,
     decode_tree,
 )
@@ -954,6 +956,10 @@ class ProbaVertexModel(SkeletonVertexModel):
         See :func:`make_skel_prop_df`. Defaults to :data:`DEFAULT_FEATURE_SPEC`.
     """
 
+    #: Bumped when :meth:`to_dict`'s shape changes incompatibly. Independent of
+    #: the ossify package version.
+    _SCHEMA_VERSION = 1
+
     def __init__(
         self,
         estimator,
@@ -974,6 +980,11 @@ class ProbaVertexModel(SkeletonVertexModel):
         self._feature_columns = list(feature_columns)
         self._smooth_alpha = smooth_alpha
         self._unary_clip = float(unary_clip)
+        #: The ``config`` argument the instance was built from via
+        #: :meth:`from_config` (a bundled model name, a user config dict, or a
+        #: config file path); ``None`` when built directly from an estimator.
+        #: Lets :meth:`to_dict` round-trip without re-serializing the estimator.
+        self._config_ref: Optional[Union[dict, str]] = None
         super().__init__(
             downstream_hops=downstream_hops,
             upstream_hops=upstream_hops,
@@ -981,6 +992,12 @@ class ProbaVertexModel(SkeletonVertexModel):
             feature_spec=feature_spec,
         )
 
+    @staticmethod
+    def _resolve_config(config: Optional[Union[dict, str]]) -> dict:
+        if not isinstance(config, dict):
+            config = load_model_config(config, dir=None)
+        return config
+
     @classmethod
     def from_config(
         cls,
@@ -997,22 +1014,125 @@ class ProbaVertexModel(SkeletonVertexModel):
         the config's ``spread_alpha``; ``unary_clip`` to the config's ``unary_clip``
         key, falling back to :data:`DEFAULT_UNARY_CLIP`.
         """
-        if not isinstance(config, dict):
-            config = load_model_config(config, dir=None)
+        config_ref = config if config is not None else DEFAULT_MODEL
+        resolved = cls._resolve_config(config)
         if smooth_alpha is None:
-            smooth_alpha = config.get("spread_alpha")
+            smooth_alpha = resolved.get("spread_alpha")
         if unary_clip is None:
-            unary_clip = config.get("unary_clip", DEFAULT_UNARY_CLIP)
-        return cls(
-            _load_xgboost_classifier(config["model_file"]),
-            config["feature_columns"],
-            downstream_hops=config.get("downstream_hops", 0),
-            upstream_hops=config.get("upstream_hops", 0),
-            bidirectional_hops=config.get("bidirectional_hops", 0),
+            unary_clip = resolved.get("unary_clip", DEFAULT_UNARY_CLIP)
+        obj = cls(
+            _load_xgboost_classifier(resolved["model_file"]),
+            resolved["feature_columns"],
+            downstream_hops=resolved.get("downstream_hops", 0),
+            upstream_hops=resolved.get("upstream_hops", 0),
+            bidirectional_hops=resolved.get("bidirectional_hops", 0),
             smooth_alpha=smooth_alpha,
             unary_clip=unary_clip,
             feature_spec=feature_spec,
         )
+        obj._config_ref = config_ref
+        return obj
+
+    def to_dict(self, *, model_file: Optional[Union[str, pathlib.Path]] = None) -> dict:
+        """Plain-data representation, round-trippable via :meth:`from_dict`.
+
+        When built via :meth:`from_config`, the ``config`` it was given (a
+        bundled model name, config file path, or dict) is stored directly, so
+        reloading re-resolves the same estimator without touching it.
+
+        Otherwise (built directly from an ``estimator``), the estimator has no
+        known source file: pass ``model_file`` to save it there (requires an
+        estimator with a ``save_model(path)`` method, e.g. XGBoost) and a fresh
+        config dict pointing at it is built automatically. Without
+        ``model_file``, this raises -- there is nothing to persist besides the
+        estimator itself, and this is deliberately not a pickle.
+
+        Raises
+        ------
+        ValueError
+            If ``feature_spec`` is non-default (it may contain callables that
+            cannot be serialized) or the estimator's source can't be
+            determined.
+        """
+        if self._feature_spec is not None:
+            raise ValueError(
+                "This ProbaVertexModel has a custom feature_spec, which may "
+                "contain callables and cannot be serialized. Reattach it "
+                "explicitly when reloading: "
+                "ProbaVertexModel.from_dict(d, feature_spec=...)."
+            )
+        config_ref = self._config_ref
+        if config_ref is None:
+            if model_file is None:
+                raise ValueError(
+                    "This ProbaVertexModel was not built via from_config(...), "
+                    "so its estimator has no known source file. Pass "
+                    "model_file=<path> to save the fitted estimator there "
+                    "(requires an estimator with a save_model(path) method, "
+                    "e.g. XGBoost), or save/load the estimator yourself and "
+                    "build a config dict for from_dict()."
+                )
+            if not hasattr(self._estimator, "save_model"):
+                raise TypeError(
+                    f"estimator of type {type(self._estimator).__name__} has no "
+                    "save_model(path) method; save it yourself and pass a "
+                    "config dict to from_dict() instead."
+                )
+            self._estimator.save_model(model_file)
+            config_ref = {
+                "model_file": str(model_file),
+                "feature_columns": self._feature_columns,
+                "downstream_hops": self._downstream_hops,
+                "upstream_hops": self._upstream_hops,
+                "bidirectional_hops": self._bidirectional_hops,
+            }
+        elif isinstance(config_ref, dict):
+            config_ref = {
+                k: (str(v) if isinstance(v, pathlib.Path) else v)
+                for k, v in config_ref.items()
+            }
+        return {
+            "type": "ProbaVertexModel",
+            "version": self._SCHEMA_VERSION,
+            "ossify_version": _ossify_version(),
+            "config": config_ref,
+            "smooth_alpha": self._smooth_alpha,
+            "unary_clip": self._unary_clip,
+        }
+
+    @classmethod
+    def from_dict(
+        cls,
+        d: dict,
+        *,
+        feature_spec: Optional[List[FeatureDef]] = None,
+    ) -> "ProbaVertexModel":
+        """Reconstruct a :class:`ProbaVertexModel` from :meth:`to_dict` output.
+
+        Unlike :meth:`from_config`, ``smooth_alpha``/``unary_clip`` are taken
+        as-is from ``d`` (not re-defaulted from the config file), so a
+        deliberate ``smooth_alpha: null`` (no smoothing) round-trips correctly.
+        """
+        version = d.get("version", 1)
+        if version > cls._SCHEMA_VERSION:
+            raise ValueError(
+                f"ProbaVertexModel dict has schema version {version}, newer "
+                f"than the {cls._SCHEMA_VERSION} this ossify release supports; "
+                "upgrade ossify to load it."
+            )
+        resolved = cls._resolve_config(d["config"])
+        obj = cls(
+            _load_xgboost_classifier(resolved["model_file"]),
+            resolved["feature_columns"],
+            downstream_hops=resolved.get("downstream_hops", 0),
+            upstream_hops=resolved.get("upstream_hops", 0),
+            bidirectional_hops=resolved.get("bidirectional_hops", 0),
+            smooth_alpha=d.get("smooth_alpha"),
+            unary_clip=d.get("unary_clip", DEFAULT_UNARY_CLIP),
+            feature_spec=feature_spec,
+        )
+        obj._config_ref = d["config"]
+        return obj
 
     @property
     def estimator(self):
@@ -1059,6 +1179,11 @@ class ProbaVertexModel(SkeletonVertexModel):
 
 # --- Structured compartment labeler -----------------------------------------
 
+#: Registry of encoder types :class:`StructuredLabeler` can (de)serialize.
+#: Extend when a new :class:`SkeletonVertexModel` subclass grows a
+#: to_dict/from_dict pair.
+_ENCODER_TYPES = {"ProbaVertexModel": ProbaVertexModel}
+
 
 class StructuredLabeler:
     """Compose an encoder + transition priors into a tree-decoded labeler.
@@ -1095,6 +1220,10 @@ class StructuredLabeler:
         this or less.
     """
 
+    #: Bumped when :meth:`to_dict`'s shape changes incompatibly. Independent of
+    #: the ossify package version.
+    _SCHEMA_VERSION = 1
+
     def __init__(
         self,
         encoder: SkeletonVertexModel,
@@ -1198,3 +1327,96 @@ class StructuredLabeler:
         """
         labels, edge_cost = self._decode(cell)
         return self._as_labels(labels, return_labels_as), edge_cost
+
+    # --- Serialization -------------------------------------------------------
+
+    def to_dict(self, *, model_file: Optional[Union[str, pathlib.Path]] = None) -> dict:
+        """Plain-data representation, round-trippable via :meth:`from_dict`.
+
+        Composes :meth:`TransitionSchema.to_dict` and the encoder's own
+        ``to_dict`` -- so a labeler built from a bundled model config (the
+        common case) serializes down to a small JSON-safe dict with no model
+        weights inlined. ``model_file`` is forwarded to the encoder's
+        ``to_dict`` for the case where its estimator has no known source file
+        (see :meth:`ProbaVertexModel.to_dict`). Includes a ``version`` (this
+        dict shape, bumped on breaking changes) and ``ossify_version`` (the
+        release that wrote it, for provenance/debugging).
+        """
+        encoder_type = type(self._encoder).__name__
+        if encoder_type not in _ENCODER_TYPES:
+            raise TypeError(
+                f"Serialization is not supported for encoder type "
+                f"{encoder_type!r}; supported types: {sorted(_ENCODER_TYPES)}."
+            )
+        return {
+            "type": "StructuredLabeler",
+            "version": self._SCHEMA_VERSION,
+            "ossify_version": _ossify_version(),
+            "schema": self._schema.to_dict(),
+            "encoder": self._encoder.to_dict(model_file=model_file),
+            "absorb_min_size": self._absorb_min_size,
+            "absorb_min_weight": self._absorb_min_weight,
+        }
+
+    @classmethod
+    def from_dict(
+        cls,
+        d: dict,
+        *,
+        feature_spec: Optional[List[FeatureDef]] = None,
+    ) -> "StructuredLabeler":
+        """Reconstruct a :class:`StructuredLabeler` from :meth:`to_dict` output.
+
+        ``feature_spec`` is forwarded to the encoder's ``from_dict`` and only
+        needed when the original encoder was built with a non-default
+        ``feature_spec`` (see :meth:`ProbaVertexModel.to_dict`).
+        """
+        version = d.get("version", 1)
+        if version > cls._SCHEMA_VERSION:
+            raise ValueError(
+                f"StructuredLabeler dict has schema version {version}, newer "
+                f"than the {cls._SCHEMA_VERSION} this ossify release supports; "
+                "upgrade ossify to load it."
+            )
+        encoder_dict = d["encoder"]
+        encoder_type = encoder_dict.get("type")
+        encoder_cls = _ENCODER_TYPES.get(encoder_type)
+        if encoder_cls is None:
+            raise ValueError(
+                f"Unknown or unsupported encoder type: {encoder_type!r}; "
+                f"supported types: {sorted(_ENCODER_TYPES)}."
+            )
+        encoder = encoder_cls.from_dict(encoder_dict, feature_spec=feature_spec)
+        schema = TransitionSchema.from_dict(d["schema"])
+        return cls(
+            encoder,
+            schema,
+            absorb_min_size=d.get("absorb_min_size"),
+            absorb_min_weight=d.get("absorb_min_weight"),
+        )
+
+    def save_config(
+        self,
+        filename: Union[str, pathlib.Path],
+        *,
+        model_file: Optional[Union[str, pathlib.Path]] = None,
+    ) -> None:
+        """Write :meth:`to_dict` to ``filename`` as JSON."""
+        with open(filename, "wb") as f:
+            f.write(
+                orjson.dumps(
+                    self.to_dict(model_file=model_file), option=orjson.OPT_INDENT_2
+                )
+            )
+
+    @classmethod
+    def load_config(
+        cls,
+        filename: Union[str, pathlib.Path],
+        *,
+        feature_spec: Optional[List[FeatureDef]] = None,
+    ) -> "StructuredLabeler":
+        """Load a :class:`StructuredLabeler` from a JSON file written by :meth:`save_config`."""
+        with open(filename, "rb") as f:
+            d = orjson.loads(f.read())
+        return cls.from_dict(d, feature_spec=feature_spec)

{ossify-0.2.4 → ossify-0.2.6}/src/ossify/file_io.py

@@ -1,6 +1,7 @@
 import io
 import os
 import tarfile
+import warnings
 from enum import IntEnum
 from pathlib import Path
 from typing import TYPE_CHECKING, BinaryIO, Optional, Union
@@ -464,7 +465,16 @@ def load_dataframe(tinfo, tf) -> pd.DataFrame:
         Loaded DataFrame
     """
     df_buf = pa.BufferReader(tf.extractfile(tinfo).read())
-    return pd.read_feather(df_buf)
+    with warnings.catch_warnings():
+        # pandas' read_feather wraps pyarrow.feather.read_feather, which pyarrow
+        # >=24 deprecated in favor of pyarrow.ipc.open_file; this is purely an
+        # upstream implementation detail pandas hasn't addressed yet.
+        warnings.filterwarnings(
+            "ignore",
+            message=r"pyarrow\.feather\.read_feather is deprecated",
+            category=FutureWarning,
+        )
+        return pd.read_feather(df_buf)
 
 
 def load_array(tinfo, tf) -> np.ndarray:

{ossify-0.2.4 → ossify-0.2.6}/src/ossify/plot.py

@@ -968,7 +968,7 @@ def plot_points(
     # which dominates small synapse markers).
     if "linewidths" not in scatter_kws:
         scatter_kws["linewidths"] = 0
-    if isinstance(palette, str):
+    if isinstance(palette, str) and colors is not None and not isinstance(colors, str):
         scatter_kws["cmap"] = palette
         if color_norm is not None:
             scatter_kws["vmin"], scatter_kws["vmax"] = color_norm

{ossify-0.2.4 → ossify-0.2.6}/src/ossify/structured_prediction.py

@@ -26,6 +26,13 @@ __all__ = [
 ]
 
 
+def _ossify_version() -> str:
+    """The installed ossify release, for provenance in serialized dicts."""
+    from . import __version__
+
+    return __version__
+
+
 class TransitionSchema:
     """Declarative label set with soft parent->child transition costs.
 
@@ -51,6 +58,10 @@ class TransitionSchema:
         label changes over the tree.
     """
 
+    #: Bumped when :meth:`to_dict`'s shape changes incompatibly. Independent of
+    #: the ossify package version.
+    _SCHEMA_VERSION = 1
+
     def __init__(
         self,
         classes: Sequence,
@@ -105,6 +116,42 @@ class TransitionSchema:
             mask[self._index[c]] = True
         return mask
 
+    def to_dict(self) -> dict:
+        """Plain-data representation, round-trippable via :meth:`from_dict`.
+
+        ``transitions`` is written as a list of ``[parent, child, cost]``
+        triples rather than a ``{(parent, child): cost}`` dict, since JSON/YAML
+        cannot use tuples as object keys. Includes a ``version`` (this dict
+        shape, bumped on breaking changes) and ``ossify_version`` (the release
+        that wrote it, for provenance/debugging) alongside the parameters.
+        """
+        return {
+            "type": "TransitionSchema",
+            "version": self._SCHEMA_VERSION,
+            "ossify_version": _ossify_version(),
+            "classes": list(self.classes),
+            "transitions": [[a, b, cost] for (a, b), cost in self.transitions.items()],
+            "root_classes": list(self.root_classes),
+            "default_cost": self.default_cost,
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "TransitionSchema":
+        """Reconstruct a :class:`TransitionSchema` from :meth:`to_dict` output."""
+        version = d.get("version", 1)
+        if version > cls._SCHEMA_VERSION:
+            raise ValueError(
+                f"TransitionSchema dict has schema version {version}, newer than "
+                f"the {cls._SCHEMA_VERSION} this ossify release supports; "
+                "upgrade ossify to load it."
+            )
+        return cls(
+            classes=d["classes"],
+            transitions={(a, b): cost for a, b, cost in d.get("transitions", [])},
+            root_classes=d.get("root_classes"),
+            default_cost=d.get("default_cost", 0.0),
+        )
+
 
 def tree_map_decode(
     parent_array: np.ndarray,