Give a hand-built object a public way onto the API surface

CHANGELOG.md

@@ -19,4 +19,5 @@ and this project adheres to [Semantic Versioning][].
 - `mantispy.tl`: aggregation, consensus profiles, mAP and replicate retrieval, hit calling, effect sizes, dose response, mechanism-of-action retrieval and enrichment, differential features, transport across sites and single-cell heterogeneity
 - `mantispy.metrics`, `mantispy.get` and `mantispy.pl`, for judging a correction, reading results out and plotting them
 - `mantispy.settings`, holding the verbosity and the cache directory the datasets download into
+- `mantispy.io`: `stamp`, which puts an `AnnData` built elsewhere — a published h5ad, another pipeline's output, a matrix of learned embeddings — on the mantispy API surface
 - `mantispy.metrics`: `known_relationships`, the share of annotated perturbation pairs whose similarity falls in either tail of the distribution over all pairs, and `evaluate_correction(covariates=...)`, which reports what a representation spends its variance on besides the batch and the label

docs/api.md

@@ -27,6 +27,7 @@ The Stores column in the tables below lists the keys each function writes.
     io.read_jump
     io.read
     io.write
+    io.stamp
     io.validate
 ```
 
@@ -46,6 +47,7 @@ pip install 'mantispy[spatial]'
 | `io.read_plate` | returns `SpatialData`: fields of view as Images, segmentations as Labels, wells as Shapes; the `cells` and `wells` Tables of a gallery source follow the contract below |
 | `io.read_jump` | as `io.read_profiles`, plus `obs`: `Metadata_JCP2022`, `Metadata_Perturbation`, `Metadata_InChIKey`, `Metadata_Control` |
 | `io.write` | validates first, then writes h5ad (or zarr for a `.zarr` suffix) |
+| `io.stamp` | `uns["mantispy"]`: `schema_version`, `resolution`; `var` (the required annotation columns, empty where absent). Refuses an object whose `obs` lacks a column that resolution requires |
 
 `io.validate` returns a report rather than raising, unless `raise_on_error=True`:
 

docs/tutorials/10_differential_features.ipynb

@@ -604,7 +604,6 @@
     "import anndata as ad\n",
     "\n",
     "from mantispy._core._stats import benjamini_hochberg\n",
-    "from mantispy._core.schema import stamp\n",
     "\n",
     "controls = adata[adata.obs[\"Metadata_Control\"].to_numpy()].copy()\n",
     "control_values = np.asarray(controls.X, dtype=np.float64)\n",
@@ -750,7 +749,7 @@
     "    obs.iloc[picked, obs.columns.get_loc(\"Metadata_Perturbation\")] = \"compound\"\n",
     "    obs[\"Metadata_Control\"] = obs[\"Metadata_Perturbation\"] == \"DMSO\"\n",
     "    scratch = ad.AnnData(X=np.asarray(controls.X).copy(), obs=obs, var=controls.var.copy())\n",
-    "    stamp(scratch, resolution=\"well\")\n",
+    "    mt.io.stamp(scratch, resolution=\"well\")\n",
     "    if use_int:\n",
     "        mt.pp.rank_int(scratch)\n",
     "    mt.tl.differential_features(scratch, block=None, key_added=\"d\")\n",

src/mantispy/_core/features.py

@@ -227,6 +227,28 @@ def parse_feature_names(names: Sequence[str], channels: Sequence[str] | None = N
     return parsed
 
 
+def empty_annotation(names: Sequence[str] | pd.Index) -> pd.DataFrame:
+    """The annotation table for features whose names carry no CellProfiler structure.
+
+    Learned embeddings, cluster compositions and feature-family signatures all have columns that are features but are not measurements of a compartment in a channel. The schema still asks for the annotation columns, so they are supplied empty rather than guessed at: :func:`parse_feature_names` reads ``openphenom_nahualX_17`` as the ``nahualX`` group of the ``openphenom`` object, which would give such an object feature families named after the model's own tensors.
+
+    Args:
+        names: The feature names, which are used only as the index.
+
+    Returns:
+        A frame indexed by ``names`` with the columns of :data:`COLUMNS`, every annotation column null and ``is_feature`` true.
+        Text columns are ``category`` dtype, as :func:`parse_feature_names` returns them, so an entirely missing column survives an h5ad round trip.
+    """
+    index = pd.Index(names)
+    empty = pd.DataFrame(index=index, columns=COLUMNS, dtype=object)
+    for column in _TEXT_COLUMNS:
+        empty[column] = pd.Categorical([None] * len(index))
+    for column in _FLOAT_COLUMNS:
+        empty[column] = np.full(len(index), np.nan)
+    empty["is_feature"] = True
+    return empty
+
+
 def load_blocklist(name: str = "default") -> list[str]:
     """Return the feature names blocked by default (the pycytominer blocklist)."""
     if name != "default":

src/mantispy/io/__init__.py

@@ -5,12 +5,12 @@
 from mantispy._core.schema import validate
 
 from ._jump import read_jump
-from ._profiles import METADATA_PREFIXES, read, read_profiles, write
+from ._profiles import METADATA_PREFIXES, read, read_profiles, stamp, write
 
 if TYPE_CHECKING:
     from ._plate import read_plate
 
-__all__ = ["METADATA_PREFIXES", "read", "read_jump", "read_plate", "read_profiles", "validate", "write"]
+__all__ = ["METADATA_PREFIXES", "read", "read_jump", "read_plate", "read_profiles", "stamp", "validate", "write"]
 
 # read_plate needs the spatial extra, which importing mantispy must not.
 _LAZY = {"read_plate": "mantispy.io._plate"}

src/mantispy/io/_profiles.py

@@ -14,12 +14,21 @@
 import numpy as np
 import pandas as pd
 
-from mantispy._core.features import _infer_channels, parse_feature_names
+from mantispy._core.features import _infer_channels, empty_annotation, parse_feature_names
 from mantispy._core.frames import as_frame, categorize_metadata
 from mantispy._core.logging import get_logger, report_drop
 from mantispy._core.plate import normalize_well
 from mantispy._core.provenance import record_params
-from mantispy._core.schema import SCHEMA_VERSION, SUPPORTED_VERSIONS, migrate, stamp, validate
+from mantispy._core.schema import (
+    REQUIRED_OBS,
+    REQUIRED_VAR,
+    RESOLUTIONS,
+    SCHEMA_VERSION,
+    SUPPORTED_VERSIONS,
+    migrate,
+    validate,
+)
+from mantispy._core.schema import stamp as _record
 from mantispy.io._cellprofiler import export_prefix, read_export
 
 #: Column-name prefixes that mark metadata. Real accessions use all four.
@@ -211,7 +220,7 @@ def from_dataframe(
     obs.index = pd.Index([str(i) for i in range(len(obs))])
 
     adata = ad.AnnData(X=X, obs=obs, var=parsed.loc[feature_names])
-    stamp(adata, resolution=resolution)
+    _record(adata, resolution=resolution)
     # Recorded whether given or inferred, so the parse can be reproduced from the object.
     if vocabulary:
         adata.uns["mantispy"]["channels"] = list(vocabulary)
@@ -463,9 +472,63 @@ def write(adata: ad.AnnData, path: str | Path) -> None:
         ValueError: The object does not satisfy the schema, with every failure :func:`~mantispy.io.validate` found in the message.
     """
     path = Path(path)
-    stamp(adata)
+    _record(adata)
     validate(adata, raise_on_error=True)
     if path.suffix == ".zarr":
         adata.write_zarr(path)
     else:
         adata.write_h5ad(path)
+
+
+def stamp(adata: ad.AnnData, resolution: str | None = None, copy: bool = False) -> ad.AnnData | None:
+    """Mark an :class:`~anndata.AnnData` built elsewhere as a mantispy object.
+
+    Every reader here, and every tool that returns a new object, records this already. This is the entry point for an object that did not come from one of them: a published ``h5ad``, another pipeline's output, a subset assembled in a notebook, or a matrix of learned embeddings with its metadata alongside.
+
+    Args:
+        adata: The object to stamp.
+        resolution: What one row is: ``"cell"``, ``"well"`` or ``"perturbation"``. ``obs`` has to carry the columns that resolution requires. The default keeps whatever resolution the object already records, and falls back to ``"well"`` for an object that records none, so re-stamping a subset does not quietly demote it.
+        copy: Return a stamped copy instead of stamping in place.
+
+    Returns:
+        ``None``, or the stamped copy.
+        Writes the schema version and the resolution to ``uns["mantispy"]``, and adds the missing feature-annotation columns to ``var``.
+
+    Raises:
+        ValueError: ``resolution`` is not one of the three, or ``obs`` lacks a column that resolution requires.
+
+    Notes:
+        Only the ``obs`` columns the resolution requires are checked, because that is what the rest of the package dispatches on. :func:`validate` gives the full report, including what it warns about rather than blocks. ``X`` is one of the things it rather than this checks: the package stores features as ``float32``, and a matrix that came out of scikit-learn or :func:`numpy.load` is ``float64``, so an embedding usually wants ``adata.X = adata.X.astype("float32")`` before it is written.
+
+        Any of the feature-annotation columns the schema requires that ``var`` does not already have are added empty, and columns already present are left as they are. They are not filled by parsing the feature names: the parser finds structure in names that have none — it reads ``openphenom_nahualX_17`` as the ``nahualX`` group of an ``openphenom`` object — and an embedding would then carry feature families named after the model's own tensors. An object read by :func:`read_profiles` already has the parsed annotation and keeps it.
+
+    Examples:
+        Bringing in a matrix of learned embeddings, one row per well:
+
+        >>> import anndata as ad
+        >>> import mantispy as mt
+        >>> adata = ad.AnnData(embeddings, obs=metadata)  # doctest: +SKIP
+        >>> mt.io.stamp(adata, resolution="well")  # doctest: +SKIP
+    """
+    if resolution is None:
+        resolution = adata.uns.get("mantispy", {}).get("resolution", "well")
+    if resolution not in RESOLUTIONS:
+        raise ValueError(f"resolution must be one of {RESOLUTIONS}, got {resolution!r}")
+
+    # Checked before the copy, so a call that is going to be rejected does not duplicate X first.
+    missing_obs = [column for column in REQUIRED_OBS[resolution] if column not in adata.obs]
+    if missing_obs:
+        raise ValueError(
+            f"obs is missing {missing_obs}, which every {resolution}-resolution object needs. Add the "
+            "column(s), or stamp at a resolution whose requirements obs meets."
+        )
+
+    target = adata.copy() if copy else adata
+    absent = [column for column in REQUIRED_VAR if column not in target.var]
+    if absent:
+        empty = empty_annotation(target.var.index)
+        for column in absent:
+            target.var[column] = empty[column]
+
+    _record(target, resolution=resolution)
+    return target if copy else None

tests/test_io_profiles.py

@@ -384,3 +384,121 @@ def test_index_columns_name_the_observations(tmp_path):
         mt.io.read_profiles(path, index_columns=("Metadata_Plate",))
     with pytest.raises(KeyError, match="index columns not in metadata"):
         mt.io.read_profiles(path, index_columns=("Metadata_Nope",))
+
+
+@pytest.mark.parametrize("resolution", ["cell", "well", "perturbation"])
+def test_stamp_puts_a_hand_built_object_on_the_api_surface(resolution):
+    """An object from another pipeline, or a matrix of learned embeddings, arrives without the
+    stamp every reader here writes, and nothing public used to establish it."""
+    import anndata as ad
+
+    columns = {
+        "cell": {"Metadata_Plate": "P1", "Metadata_Well": "A01"},
+        "well": {"Metadata_Plate": "P1", "Metadata_Well": "A01"},
+        "perturbation": {"Metadata_Perturbation": "cmpd"},
+    }[resolution]
+    obs = pd.DataFrame({name: [value] * 4 for name, value in columns.items()}, index=list("abcd"))
+    adata = ad.AnnData(np.arange(20, dtype=np.float32).reshape(4, 5), obs=obs)
+
+    assert mt.io.stamp(adata, resolution=resolution) is None
+    assert adata.uns["mantispy"]["resolution"] == resolution
+    assert mt.io.validate(adata).ok
+
+
+def test_stamp_refuses_what_the_resolution_needs_and_obs_lacks():
+    """Stamping regardless would push the failure into whichever tool ran next."""
+    import anndata as ad
+
+    adata = ad.AnnData(np.zeros((3, 2), dtype=np.float32), obs=pd.DataFrame(index=list("abc")))
+    with pytest.raises(ValueError, match=r"Metadata_Plate.*Metadata_Well"):
+        mt.io.stamp(adata)
+    with pytest.raises(ValueError, match="resolution must be one of"):
+        mt.io.stamp(adata, resolution="plate")
+    assert "mantispy" not in adata.uns
+
+
+def test_stamp_can_leave_the_original_alone():
+    import anndata as ad
+
+    obs = pd.DataFrame({"Metadata_Perturbation": ["a", "b"]}, index=["x", "y"])
+    adata = ad.AnnData(np.zeros((2, 3), dtype=np.float32), obs=obs)
+
+    stamped = mt.io.stamp(adata, resolution="perturbation", copy=True)
+    assert stamped.uns["mantispy"]["resolution"] == "perturbation"
+    assert "mantispy" not in adata.uns
+    assert list(adata.var.columns) == []
+
+
+def test_stamp_keeps_the_resolution_the_object_already_records():
+    """A subset of a cell-resolution object is still cell-resolution, and the well default would
+    have demoted it silently: tl.aggregate then takes the non-cell branch and fills
+    Metadata_CellCount with NaN, which disables its min_cells filter."""
+    import anndata as ad
+
+    obs = pd.DataFrame({"Metadata_Plate": ["P1"] * 3, "Metadata_Well": ["A01"] * 3}, index=list("abc"))
+    adata = ad.AnnData(np.zeros((3, 2), dtype=np.float32), obs=obs)
+    mt.io.stamp(adata, resolution="cell")
+
+    mt.io.stamp(adata[:2].copy())
+    mt.io.stamp(adata)
+    assert adata.uns["mantispy"]["resolution"] == "cell"
+    assert mt.io.stamp(adata, resolution="well") is None
+    assert adata.uns["mantispy"]["resolution"] == "well"
+
+
+def test_stamp_lets_a_learned_embedding_be_written(tmp_path):
+    """An embedding has no CellProfiler feature names, so its var carries none of the annotation
+    the schema requires, and `io.write` validates before writing. Without the annotation columns
+    a stamped embedding failed on ten missing var columns and could not be written at all."""
+    import anndata as ad
+
+    obs = pd.DataFrame(
+        {"Metadata_Plate": ["P1"] * 4, "Metadata_Well": ["A01", "A02", "A03", "A04"]},
+        index=list("abcd"),
+    )
+    adata = ad.AnnData(np.arange(24, dtype=np.float32).reshape(4, 6), obs=obs)
+    # The names JUMP-Lite ships. Parsing them reads 'openphenom' as the object and 'nahualX' as
+    # the feature group, so an embedding stamped by the parser grew feature families named after
+    # the model's own tensors.
+    adata.var_names = [f"openphenom_nahualX_{index}" for index in range(6)]
+
+    mt.io.stamp(adata, resolution="well")
+    report = mt.io.validate(adata)
+    assert report.ok, str(report)
+    for column in ("object", "feature_group", "feature", "channel", "params"):
+        assert adata.var[column].isna().all(), column
+    assert adata.var["is_feature"].all()
+
+    path = tmp_path / "embedding.h5ad"
+    mt.io.write(adata, path)
+    assert mt.io.read(path).shape == (4, 6)
+
+
+def test_stamp_keeps_an_annotation_that_is_already_there():
+    """A profile object read by read_profiles carries the parsed annotation, and stamping it
+    again must not blank it."""
+    frame = _frame()
+    adata = from_dataframe(frame)
+    parsed = adata.var["feature_group"].copy()
+
+    mt.io.stamp(adata, resolution="well")
+    pd.testing.assert_series_equal(adata.var["feature_group"], parsed)
+
+
+def test_stamp_fills_only_the_annotation_columns_that_are_missing():
+    """An object hand-built with part of the annotation is the case where the merge can go wrong:
+    the columns that are there have to survive, and the ones added have to be categorical, because
+    an object array of NaN cannot be written to h5ad."""
+    import anndata as ad
+
+    obs = pd.DataFrame({"Metadata_Plate": ["P1"] * 2, "Metadata_Well": ["A01", "A02"]}, index=list("ab"))
+    adata = ad.AnnData(np.zeros((2, 3), dtype=np.float32), obs=obs)
+    adata.var["object"] = pd.Categorical(["Cells", "Nuclei", "Cells"])
+    adata.var["is_feature"] = [True, True, False]
+
+    mt.io.stamp(adata)
+    assert list(adata.var["object"]) == ["Cells", "Nuclei", "Cells"]
+    assert list(adata.var["is_feature"]) == [True, True, False]
+    assert adata.var["feature_group"].isna().all()
+    assert isinstance(adata.var["feature_group"].dtype, pd.CategoricalDtype)
+    assert mt.io.validate(adata).ok