ossify 0.2.3__tar.gz → 0.2.5__tar.gz

{ossify-0.2.3 → ossify-0.2.5}/PKG-INFO

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

{ossify-0.2.3 → ossify-0.2.5}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 allow-direct-references = true
 [project]
 name = "ossify"
-version = "0.2.3"
+version = "0.2.5"
 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.3"
+current_version = "0.2.5"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 regex = false

{ossify-0.2.3 → ossify-0.2.5}/src/ossify/__init__.py

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

{ossify-0.2.3 → ossify-0.2.5}/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.3 → ossify-0.2.5}/src/ossify/file_io.py

@@ -913,6 +913,68 @@ def build_point_cloud(
         return cell
 
 
+def _exact_integer_id_coverage(link_ids: pd.Index, source_ids: pd.Index) -> bool:
+    """Test whether two integer-ID domains are exactly the same set.
+
+    Both arguments must already be canonicalized ``int64`` pandas ``Index``
+    objects. The comparison is order-independent and stays entirely in the
+    integer domain -- it never coerces identifiers through ``float64`` -- so it
+    is safe for IDs above ``2**53``.
+
+    Parameters
+    ----------
+    link_ids :
+        The linkage column for a candidate source layer.
+    source_ids :
+        That layer's vertex index.
+
+    Returns
+    -------
+    :
+        ``True`` iff the two contain exactly the same set of IDs.
+    """
+    if len(link_ids) != len(source_ids):
+        return False
+    # ``symmetric_difference`` is empty iff neither side has an ID the other
+    # lacks. Using pandas ``Index`` set ops keeps the comparison in int64 and
+    # avoids materializing Python ``set`` objects over large ID arrays.
+    return len(link_ids.symmetric_difference(source_ids)) == 0
+
+
+def _linkage_source_error(linkage_pair, link_df, cell, sample_size: int = 10) -> str:
+    """Build a concise diagnostic message for an unresolvable linkage source.
+
+    Reports, for each endpoint, the layer vertex count, the linkage row count,
+    the number of unique linkage IDs, and the missing/extra ID counts, plus a
+    small bounded sample of the offending IDs. This replaces the enormous
+    pandas ``KeyError`` that a bad ``.loc`` reindex would otherwise raise.
+    """
+    lines = [
+        f"Could not infer a valid source layer for linkage "
+        f"({linkage_pair[0]!r}, {linkage_pair[1]!r}) with {len(link_df)} rows.",
+        "A source endpoint must map exactly one linkage row to each of its "
+        "vertices (unique, complete coverage). Neither endpoint qualifies:",
+    ]
+    for source in linkage_pair:
+        source_ids = canonicalize_ids(
+            pd.Index(cell._all_objects[source].vertex_index), name=source
+        )
+        link_ids = canonicalize_ids(pd.Index(link_df[source]), name=source)
+        unique_link_ids = link_ids.unique()
+        missing = source_ids.difference(link_ids)
+        extra = link_ids.difference(source_ids)
+        lines.append(
+            f"  - {source!r}: vertices={len(source_ids)}, "
+            f"link_rows={len(link_ids)}, unique_link_ids={len(unique_link_ids)}, "
+            f"missing_source_ids={len(missing)}, extra_link_ids={len(extra)}"
+        )
+        if len(missing):
+            lines.append(f"      missing sample: {missing[:sample_size].tolist()}")
+        if len(extra):
+            lines.append(f"      extra sample: {extra[:sample_size].tolist()}")
+    return "\n".join(lines)
+
+
 def build_linkage(
     linkage_pair,
     tf,
@@ -922,23 +984,50 @@ def build_linkage(
     link_df = load_dataframe(f"{prefix}/linkage.feather", tf)
     # Legacy .osy files may store link columns with mixed int64/uint64 dtypes
     # (dtype optimization only downcast signed ints, leaving uint64 untouched).
-    # Canonicalize both columns to int64 before the label-based ``.loc`` reindex
-    # below so that lookup cannot collide two IDs above 2**53 through a float
-    # coercion of mismatched signed/unsigned keys.
+    # Canonicalize both columns to int64 before any label-based reindex below so
+    # that lookup cannot collide two IDs above 2**53 through a float coercion of
+    # mismatched signed/unsigned keys.
     for col in linkage_pair:
         if col in link_df.columns:
             link_df[col] = canonicalize_ids(link_df[col], name=col)
-    # Determine source based on the length of the vertices in the mapping and in the skeleton layer
-    if len(link_df) == len(cell._all_objects[linkage_pair[0]].nodes):
-        source_layer = linkage_pair[0]
-        target_layer = linkage_pair[1]
-    elif len(link_df) == len(cell._all_objects[linkage_pair[1]].nodes):
-        source_layer = linkage_pair[1]
-        target_layer = linkage_pair[0]
-    else:
-        raise ValueError("Linkage DataFrame does not match any layer.")
+
+    # The archive sorts the pair names, so serialized column order does not
+    # preserve the original source direction. Infer the source by exact
+    # source-domain validation rather than row count alone: a layer qualifies
+    # as source only when its linkage column has no missing IDs, is unique, has
+    # one row per vertex, and covers exactly that layer's vertex index. Row
+    # count alone is ambiguous whenever the two layers have equal cardinality
+    # (e.g. a many-to-one graph <-> annotation link where the counts coincide).
+    candidates = []
+    for source, target in (
+        (linkage_pair[0], linkage_pair[1]),
+        (linkage_pair[1], linkage_pair[0]),
+    ):
+        source_ids = canonicalize_ids(
+            pd.Index(cell._all_objects[source].vertex_index), name=source
+        )
+        link_ids = canonicalize_ids(pd.Index(link_df[source]), name=source)
+        if (
+            len(link_ids) == len(source_ids)
+            and link_ids.is_unique
+            and _exact_integer_id_coverage(link_ids, source_ids)
+        ):
+            candidates.append((source, target))
+
+    if not candidates:
+        raise ValueError(_linkage_source_error(linkage_pair, link_df, cell))
+
+    # Zero candidates raises above. One candidate is the unambiguous source.
+    # Two candidates means a true bijection: both columns are unique and cover
+    # their layers exactly, so the mapping is one-to-one. MorphSync stores every
+    # link in both directions (see ``add_link``), so either choice yields the
+    # same bidirectional link; pick the first deterministically.
+    source_layer, target_layer = candidates[0]
 
     layer = cell._all_objects[source_layer]
+    # Validation above guarantees the source column is unique and covers exactly
+    # the source vertex index, so this reindex is a total, collision-free
+    # reordering -- it can no longer raise a large KeyError.
     cell._all_objects[source_layer]._process_linkage(
         Link(
             link_df.set_index(source_layer).loc[layer.vertex_index][target_layer],

{ossify-0.2.3 → ossify-0.2.5}/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,