Adding shape validation for add_matrix_to_collection (#4258)

ktsitsi · web-flow · commit 2d91da66948f · 2025-10-21T16:23:04.000-04:00
* Adding shape validation for add_matrix_to_collection

* Adding test for add_X_layer

* Add context to the validation function

* Schema validation optional default False

* PR fixes

* PR fixes vol 2

* Adding all versions in tests

* Adding full docstrings

* Add entry in HISTORY.md

* Minor fixes for PR

* Opening internal var dataframe in read mode when writing &amp; text fix

* PR fixes - timestamp, closed property

* Add unit tests for previous versions

* PR comments

* PR comments

* Add warning instead of log msg

* Rebase with main

* Final misc

* Match argument to all tests
diff --git a/apis/python/HISTORY.md b/apis/python/HISTORY.md
@@ -10,6 +10,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Changed
 
+- \[[#4258](https://github.com/single-cell-data/TileDB-SOMA/pull/4258)\] Add optional schema validation to `tiledbsoma.io.add_X_layer`, which will generate an error if the provided matrix does not match the Experiment shape. Validation is enabled by default, for any dataset created with SOMA 1.15 or later.
+
 ### Deprecated
 
 ### Removed
diff --git a/apis/python/remote_tests/test_02_analysis.py b/apis/python/remote_tests/test_02_analysis.py
@@ -79,8 +79,11 @@ def test_write_with_updates(conftest_context, conftest_namespace, conftest_defau
         assert sorted(list(exp.ms["RNA"].X.keys())) == ["data", "logcounts"]
 
     # Add dimensional-reduction results
+    # identify highly variable genes but don't subset the data
+    # https://github.com/single-cell-data/TileDB-SOMA/pull/4258#discussion_r2417729246
     sc.pp.highly_variable_genes(adata, inplace=True)
-    adata = adata[:, adata.var.highly_variable]
+
+    # Run PCA on the full dataset, scanpy handles HVG selection
     sc.pp.scale(adata)
     sc.tl.pca(adata, use_highly_variable=True, n_comps=5)
 
diff --git a/apis/python/src/tiledbsoma/io/ingest.py b/apis/python/src/tiledbsoma/io/ingest.py
@@ -1654,6 +1654,90 @@ def update_matrix(
     )
 
 
+def _validate_matrix_to_collection(
+    exp: Experiment,
+    meas: Measurement,
+    collection_name: str,
+    matrix_name: str,
+    matrix_data: Matrix | h5py.Dataset,
+) -> None:
+    """Validates against the Experiment shape, as defined by the axis dataframes.
+    Skips validation for pre-1.15 schemas till issue is resolved
+    https://github.com/single-cell-data/TileDB-SOMA/pull/4258#discussion_r2429676202.
+
+    Args:
+        exp: The experiment object
+        measurement_name: Name of the measurement
+        collection_name: Name of the collection
+        matrix_name: Name of the matrix
+        matrix_data: The matrix data to validate
+        meas: The already-opened measurement object
+
+    Raises:
+        ValueError: If the matrix shape is incompatible.
+    """
+    if not exp.obs.tiledbsoma_has_upgraded_domain:
+        warnings.warn(
+            f"Skipped validation for arrays with pre-1.15 schema\
+                       {matrix_name}",
+            stacklevel=2,
+        )
+        return
+
+    obs_soma_joinid_domain = exp.obs._maybe_soma_joinid_shape
+
+    # Try to access the var DataFrame directly, fallback to explicit open if needed
+    if meas.var.closed:
+        with DataFrame.open(meas.var.uri, "r", context=meas.context, tiledb_timestamp=meas.tiledb_timestamp) as var_df:
+            var_soma_joinid_domain = var_df._maybe_soma_joinid_shape
+    else:
+        var_soma_joinid_domain = meas.var._maybe_soma_joinid_shape
+
+    target_obs_size = obs_soma_joinid_domain  # O: Observation (row) size
+    target_var_size = var_soma_joinid_domain  # V: Variable (column) size
+    target_shape = (target_obs_size, target_var_size)
+    # Check if the dimensionality is the same
+    if len(matrix_data.shape) != len(target_shape):
+        raise SOMAError(
+            f"Matrix '{matrix_name}' has {len(matrix_data.shape)} dimensions, "
+            f"but the existing matrix has {len(target_shape)}. Shapes must match."
+        )
+
+    if collection_name == "X":
+        # X: (O, V) - Must match both dimensions exactly
+        if matrix_data.shape != target_shape:
+            raise SOMAError(
+                f"Matrix 'X' shape {matrix_data.shape} is incompatible with target shape {target_shape}. "
+                "Both dimensions must match exactly."
+            )
+    elif collection_name == "obsm":
+        if matrix_data.shape[0] != target_obs_size:
+            raise SOMAError(
+                f"Matrix '{matrix_name}' must match the observation size (O). "
+                f"Found {matrix_data.shape[0]}, expected {target_obs_size}."
+            )
+    elif collection_name == "varm":
+        if matrix_data.shape[0] != target_var_size:
+            raise SOMAError(
+                f"Matrix '{matrix_name}' must match the variable size (V). "
+                f"Found {matrix_data.shape[0]}, expected {target_var_size}."
+            )
+    elif collection_name == "obsp":
+        expected_shape = (target_obs_size, target_obs_size)
+        if matrix_data.shape[0:2] != expected_shape:
+            raise SOMAError(
+                f"Matrix '{matrix_name}' must match the observation size. "
+                f"Found {matrix_data.shape[0:2]}, expected {(target_obs_size, target_obs_size)}."
+            )
+    elif collection_name == "varp":
+        expected_shape = (target_var_size, target_var_size)
+        if matrix_data.shape[0:2] != expected_shape:
+            raise SOMAError(
+                f"Matrix '{matrix_name}' must match the observation size (O). "
+                f"Found {matrix_data.shape[0:2]}, expected {(target_var_size, target_var_size)}."
+            )
+
+
 def add_X_layer(
     exp: Experiment,
     measurement_name: str,
@@ -1663,11 +1747,34 @@ def add_X_layer(
     ingest_mode: IngestMode = "write",
     use_relative_uri: bool | None = None,
     context: SOMATileDBContext | None = None,
+    schema_validation: bool = True,
 ) -> None:
-    """This is useful for adding X data, for example from
+    """Add a new X layer to a measurement in the experiment.
+
+    This is useful for adding X data, for example from
     `Scanpy <https://scanpy.readthedocs.io/>`_'s ``scanpy.pp.normalize_total``,
     ``scanpy.pp.log1p``, etc.
 
+    Args:
+        exp: The experiment to add the X layer to. Must be open for writing.
+        measurement_name: Name of the measurement within the experiment.
+        X_layer_name: Name for the new X layer (e.g., "normalized", "log1p").
+        X_layer_data: Matrix data to store in the X layer. Can be a sparse matrix
+            (e.g., scipy.csr_matrix) or an h5py.Dataset.
+        ingest_mode: Mode for ingestion. Defaults to "write".
+        use_relative_uri: Whether to use relative URIs when storing references.
+            If None, the default behavior is used.
+        context: TileDB context for the operation. If None, uses the default context.
+        schema_validation: Whether to validate the matrix schema before adding.
+            Defaults to True.
+
+    Returns:
+        None
+
+    Raises:
+        ValueError: If the experiment is not open for writing or if schema
+            validation fails.
+
     Lifecycle:
         Maturing.
     """
@@ -1682,6 +1789,7 @@ def add_X_layer(
         ingest_mode=ingest_mode,
         use_relative_uri=use_relative_uri,
         context=context,
+        schema_validation=schema_validation,
     )
 
 
@@ -1695,10 +1803,39 @@ def add_matrix_to_collection(
     ingest_mode: IngestMode = "write",
     use_relative_uri: bool | None = None,
     context: SOMATileDBContext | None = None,
+    schema_validation: bool = True,
 ) -> None:
-    """This is useful for adding X/obsp/varm/etc data, for example from
+    """Add a matrix to a specified collection within a measurement.
+
+    This is useful for adding X/obsp/varm/etc data, for example from
     Scanpy's ``scanpy.pp.normalize_total``, ``scanpy.pp.log1p``, etc.
 
+    Args:
+        exp: The experiment to add the matrix to. Must be open for writing.
+        measurement_name: Name of the measurement within the experiment.
+        collection_name: Name of the collection to add the matrix to
+            (e.g., "X", "obsp", "varm").
+        matrix_name: Name for the new matrix within the collection.
+        matrix_data: Matrix data to store. Can be a sparse matrix
+            (e.g., scipy.csr_matrix) or an h5py.Dataset.
+        ingest_mode: Mode for ingestion. Defaults to "write".
+        use_relative_uri: Whether to use relative URIs when storing references.
+            If None, the default behavior is used.
+        context: TileDB context for the operation. If None, uses the default context.
+        schema_validation: Whether to validate the matrix schema before adding.
+            Defaults to True.
+
+    Returns:
+        None
+
+    Raises:
+        ValueError: If schema validation fails or if the collection structure
+            is invalid.
+
+    Notes:
+        - If the collection doesn't exist, it will be created automatically.
+        - The matrix is stored as a SparseNDArray with identity axis mappings.
+
     Lifecycle:
         Maturing.
     """
@@ -1711,13 +1848,15 @@ def add_matrix_to_collection(
     # tiledb://namespace/uuid.  When the caller passes a creation URI (which
     # they must) via exp.uri, we need to follow that.
     extend_creation_uri = exp.uri.startswith("tiledb://")
-
     with exp.ms[measurement_name] as meas:
         if extend_creation_uri:
             coll_uri = f"{exp.uri}/ms/{_util.sanitize_key(measurement_name)}/{_util.sanitize_key(collection_name)}"
         else:
             coll_uri = _util.uri_joinpath(meas.uri, _util.sanitize_key(collection_name))
 
+        if schema_validation:
+            _validate_matrix_to_collection(exp, meas, collection_name, matrix_name, matrix_data)
+
         if collection_name in meas:
             coll = cast("Collection[RawHandle]", meas[collection_name])
         else:
@@ -1728,6 +1867,7 @@ def add_matrix_to_collection(
                 context=context,
             )
             _maybe_set(meas, collection_name, coll, use_relative_uri=use_relative_uri)
+
         with coll:
             matrix_uri = _util.uri_joinpath(coll_uri, _util.sanitize_key(matrix_name))
 
diff --git a/apis/python/tests/test_basic_anndata_io.py b/apis/python/tests/test_basic_anndata_io.py
@@ -4,6 +4,7 @@
 import gc
 import json
 import random
+import shutil
 from copy import deepcopy
 from pathlib import Path
 
@@ -15,6 +16,7 @@
 import scipy
 import somacore
 from anndata import AnnData
+from packaging.version import Version
 from scipy.sparse import csr_matrix
 
 import tiledbsoma
@@ -23,7 +25,7 @@
 from tiledbsoma._soma_object import SOMAObject
 from tiledbsoma.io._common import _TILEDBSOMA_TYPE, UnsDict, UnsMapping
 
-from ._util import TESTDATA, assert_adata_equal, make_pd_df
+from ._util import ROOT_DATA_DIR, TESTDATA, assert_adata_equal, make_pd_df
 
 
 @pytest.fixture
@@ -472,6 +474,115 @@ def test_add_matrix_to_collection(conftest_pbmc_small, tmp_path):
         assert sorted(list(exp_r.ms["RNA"]["newthing"].keys())) == ["X_pcd"]
 
 
+@pytest.mark.parametrize(
+    "matrix_type, offset",
+    [
+        ("obsm", (1, 0)),
+        ("obsm", (-1, 0)),
+        ("obsm", (0, 0)),  # Should succeed without raising errors
+        ("varm", (1, 0)),
+        ("varm", (-1, 0)),
+        ("varm", (0, 0)),
+        ("obsp", (1, 0)),
+        ("obsp", (0, 1)),
+        ("obsp", (-1, 0)),
+        ("obsp", (0, 0)),
+        ("varp", (1, 0)),
+        ("varp", (0, 1)),
+        ("varp", (-1, 0)),
+        ("varp", (0, 0)),
+    ],
+)
+@pytest.mark.parametrize("version", ["1.14.5", "1.7.3", "1.12.3", "1.14.5", "1.15.0", "1.15.7", "1.16.1"])
+@pytest.mark.parametrize(
+    "name_and_expected_shape",
+    [["pbmc3k_unprocessed", (2700, 13714)], ["pbmc3k_processed", (2638, 1838)]],
+)
+@pytest.mark.medium_runner
+def test_matrix_shape_enforcement(soma_tiledb_context, version, name_and_expected_shape, tmp_path, matrix_type, offset):
+    """Test shape validation for all matrix collection types."""
+    output_path = tmp_path.as_posix()
+    name, expected_shape = name_and_expected_shape
+    original_path = ROOT_DATA_DIR / "soma-experiment-versions-2025-04-04" / version / name
+    uri = str(original_path)
+    if not Path(uri).is_dir():
+        raise RuntimeError(
+            f"Missing '{uri}' directory. Try running `make data` from the TileDB-SOMA project root directory.",
+        )
+
+    shutil.copytree(uri, output_path, dirs_exist_ok=True)
+    n_obs = expected_shape[0]
+    n_vars = expected_shape[1]
+
+    # Determine expected shape based on matrix type
+    shape_map = {
+        "obsm": (n_obs, 5),
+        "varm": (n_vars, 5),
+        "obsp": (n_obs, n_obs),
+        "varp": (n_vars, n_vars),
+    }
+    base_shape = shape_map[matrix_type]
+    bad_shape = (max(1, base_shape[0] + offset[0]), max(1, base_shape[1] + offset[1]))
+
+    # Create invalid matrix
+    bad_matrix = csr_matrix((bad_shape[0], bad_shape[1])) if matrix_type in ["obsp", "varp"] else np.ones(bad_shape)
+
+    if Version(version) >= Version("1.15"):
+        if offset != (0, 0):
+            # Should raise error of invalid shape for observation/variable size
+            with (
+                _factory.open(output_path, "w") as exp,
+                pytest.raises(tiledbsoma.SOMAError, match=r"Matrix .* must match the .* size"),
+            ):
+                tiledbsoma.io.add_matrix_to_collection(exp, "RNA", matrix_type, "bad", bad_matrix)
+        else:
+            # Should pass the validation without errors
+            with _factory.open(output_path, "w") as exp:
+                tiledbsoma.io.add_matrix_to_collection(exp, "RNA", matrix_type, "bad", bad_matrix)
+    else:
+        # pre-1.15 versions are not being checked for schema validation.
+        if offset == (0, 0):
+            with _factory.open(output_path, "w") as exp:
+                tiledbsoma.io.add_matrix_to_collection(exp, "RNA", matrix_type, "bad", bad_matrix)
+        else:
+            # TODO: _maybe_soma_joinid_shape returns max-domain for pre-1.14 instead of None
+            # https://github.com/single-cell-data/TileDB-SOMA/pull/4258#discussion_r2429676202
+            pass
+
+
+@pytest.mark.parametrize(
+    "offset",
+    [
+        (1, 0),
+        (0, 1),
+        (1, 1),
+    ],
+)
+def test_matrix_X_layer_enforcement(conftest_pbmc_small, tmp_path, offset):
+    """Test shape validation for all matrix collection types."""
+    output_path = tmp_path.as_posix()
+    original = conftest_pbmc_small.copy()
+
+    # Create SOMA experiment
+    tiledbsoma.io.from_anndata(output_path, conftest_pbmc_small, measurement_name="RNA")
+    assert_adata_equal(original, conftest_pbmc_small)
+
+    # Determine expected shape based on matrix type
+    shape_map = conftest_pbmc_small.n_obs, conftest_pbmc_small.n_vars
+    base_shape = shape_map
+    bad_shape = (max(1, base_shape[0] + offset[0]), max(1, base_shape[1] + offset[1]))
+
+    # Create invalid matrix
+    bad_matrix = np.ones(bad_shape)
+
+    # Should raise error for invalid shape
+    with (
+        _factory.open(output_path, "w") as exp,
+        pytest.raises(tiledbsoma.SOMAError, match=r"Matrix .* shape .* is incompatible with target shape *."),
+    ):
+        tiledbsoma.io.add_X_layer(exp, "RNA", "bad", bad_matrix, schema_validation=True)
+
+
 def test_export_anndata(conftest_pbmc_small, tmp_path):
     output_path = tmp_path.as_posix()
 
