"""Core data contracts shared across featurizers.
We have contracts for the data inputs and outputs of featurizers,
as well as for column names in the output DataFrame.
These contracts are enforced through a combination of Pydantic models,
Beartypes Pandera schemas.
Where inputs are the image path or arrays.
Outputs are the feature DataFrame and metadata dictionary/Dfs.
The package accepts:
- Single-channel 3D arrays shaped (z, y, x)
"""
from __future__ import annotations
from typing import Any
import numpy as np
import pandas as pd
import pandera.pandas as pa
from beartype import beartype
from pydantic import (
BaseModel,
ConfigDict,
Field,
field_validator,
model_validator,
)
from zedprofiler.exceptions import ContractError
EXPECTED_SPATIAL_DIMS = 3
TWO_DIMENSIONAL = 2
FOUR_DIMENSIONAL = 4
FIVE_OR_MORE_DIMENSIONS = 5
NON_METADATA_UNDERSCORE_SEPARATED_PARTS = 4
METADATA_UNDERSCORE_SEPARATED_PARTS = 3
REQUIRED_RETURN_KEYS = ("image_array", "features", "metadata")
FEATURES = [
"Colocalization",
"Granularity",
"Texture",
"Intensity",
"Neighbors",
"VolumeSizeShape",
]
# Pandera schema for validating numpy arrays with expected dimensionality
ImageArraySchema = pa.DataFrameSchema(
columns={},
checks=[
pa.Check(
lambda x: len(x.shape) in [3, 4, 5],
error="Array must have 3, 4, or 5 dimensions",
),
],
name="ImageArray",
)
[docs]
class ImageArrayModel(BaseModel):
"""Pydantic model for validating image arrays."""
array: np.ndarray
model_config = ConfigDict(arbitrary_types_allowed=True)
[docs]
@field_validator("array", mode="before")
@classmethod
def validate_array_type(_cls, v: object) -> np.ndarray:
"""Ensure array is a numpy array."""
if not isinstance(v, np.ndarray):
raise ValueError(f"Expected numpy array, got {type(v)}")
return v
[docs]
@field_validator("array", mode="after")
@classmethod
def validate_array_dtype_and_shape(_cls, arr: np.ndarray) -> np.ndarray:
"""Validate array dtype is numeric and shape is valid."""
if not np.issubdtype(arr.dtype, np.number):
raise ValueError(f"Array dtype must be numeric, got {arr.dtype}")
arr_shape = arr.shape
# Check dimensionality
if len(arr_shape) == TWO_DIMENSIONAL:
raise ValueError(
f"Input array has shape {arr_shape} with {TWO_DIMENSIONAL} "
f"dimensions. Expected {EXPECTED_SPATIAL_DIMS} dimensions.",
)
if len(arr_shape) == FOUR_DIMENSIONAL and arr_shape[0] > 1:
raise ValueError(
f"Input array has shape {arr_shape} with {FOUR_DIMENSIONAL} "
"dimensions, but the first dimension (channels) has size "
f"{arr_shape[0]}. Expected a single-channel 3D array.",
)
if (
len(arr_shape) >= FIVE_OR_MORE_DIMENSIONS
and arr_shape[0] > 1
and arr_shape[1] > 1
):
raise ValueError(
f"Input array has shape {arr_shape} with {len(arr_shape)} "
f"dimensions. Expected {EXPECTED_SPATIAL_DIMS} dimensions.",
)
# Check all dimensions are positive
for dim_size in arr_shape:
if dim_size <= 0:
raise ValueError(
f"Input array has shape {arr_shape} with non-positive "
"dimension size. All dimensions must have size greater than 0.",
)
# Check that the array is not a degenerate single-voxel volume
if sum(arr_shape) == len(arr_shape):
raise ValueError(
f"Input array has shape {arr_shape} with all dimensions equal to 1. "
"Expected all three dimensions to have size greater than 1.",
)
if not validate_image_array_shape_contracts(arr):
raise ValueError(
f"Input array with shape {arr_shape} failed shape contract validation.",
)
if not validate_image_array_type_contracts(arr):
raise ValueError(
f"Input array with dtype {arr.dtype} failed type contract validation.",
)
return arr
[docs]
class FeatureDictModel(BaseModel):
"""Pydantic model for validating feature dictionaries."""
features: dict[str, Any]
[docs]
@field_validator("features", mode="before")
@classmethod
def validate_is_dict(_cls, v: object) -> dict[str, Any]:
"""Ensure features is a dictionary."""
if not isinstance(v, dict):
raise ValueError(f"Expected dict, got {type(v)}")
return v
[docs]
class ReturnSchemaModel(BaseModel):
"""Pydantic model for validating return schema."""
result: dict[str, Any]
model_config = ConfigDict(arbitrary_types_allowed=True)
[docs]
@field_validator("result", mode="before")
@classmethod
def validate_is_dict(_cls, v: object) -> dict[str, Any]:
"""Ensure result is a dictionary."""
if not isinstance(v, dict):
raise ValueError(f"Expected dict, got {type(v)}")
return v
[docs]
@field_validator("result", mode="after")
@classmethod
def validate_keys_and_types(_cls, result: dict[str, Any]) -> dict[str, Any]:
"""Validate result has correct keys in correct order and correct types."""
actual_keys = tuple(result.keys())
if actual_keys != REQUIRED_RETURN_KEYS:
raise ValueError(
"Return result keys must match required deterministic order "
f"{REQUIRED_RETURN_KEYS}, got {actual_keys}.",
)
if not isinstance(result["image_array"], np.ndarray):
raise ValueError(
f"Return result key 'image_array' must be a numpy array, "
f"got {type(result['image_array'])}",
)
if not isinstance(result["features"], dict):
raise ValueError(
f"Return result key 'features' must be a dict, "
f"got {type(result['features'])}",
)
if not isinstance(result["metadata"], dict):
raise ValueError(
f"Return result key 'metadata' must be a dict, "
f"got {type(result['metadata'])}",
)
return result
[docs]
class ColumnNameModel(BaseModel):
"""Pydantic model for parsing and validating column names."""
column_name: str
compartment: str | None = None
channel: str | None = None
feature: str | None = None
[docs]
@field_validator("column_name", mode="before")
@classmethod
def validate_is_string(_cls, v: object) -> str:
"""Ensure column_name is a string."""
if not isinstance(v, str):
raise ValueError(f"Expected string, got {type(v)}")
return v
[docs]
@model_validator(mode="after")
def parse_column_name(self) -> ColumnNameModel:
"""Parse column name into components."""
parts = self.column_name.split("_")
if "Metadata" in self.column_name:
if len(parts) < METADATA_UNDERSCORE_SEPARATED_PARTS:
raise ValueError(
"Metadata column name must have at least "
f"{METADATA_UNDERSCORE_SEPARATED_PARTS} "
"parts separated by underscores, "
f"got {len(parts)} parts in '{self.column_name}'",
)
# Don't parse compartment/channel/feature for metadata columns
return self
if len(parts) < NON_METADATA_UNDERSCORE_SEPARATED_PARTS:
raise ValueError(
"Column name must have at least "
f"{NON_METADATA_UNDERSCORE_SEPARATED_PARTS} "
"parts separated by underscores, "
f"got {len(parts)} parts in '{self.column_name}'",
)
self.compartment = parts[0]
self.channel = parts[1]
self.feature = parts[2]
return self
[docs]
@beartype
def validate_image_array_shape_contracts(
arr: np.ndarray,
) -> bool:
"""Validate the input array for dimensionality
Parameters
----------
arr : np.ndarray
Input array to validate
Returns
-------
bool
The status of the validation
Raises
------
ContractError
If the input array does not meet the expected contract
"""
arr_shape = arr.shape
if len(arr_shape) == TWO_DIMENSIONAL:
raise ContractError(
f"Input array has shape {arr_shape} with {TWO_DIMENSIONAL} dimensions. "
f"Expected {EXPECTED_SPATIAL_DIMS} dimensions.",
)
if len(arr_shape) == FOUR_DIMENSIONAL and arr_shape[0] > 1:
raise ContractError(
f"Input array has shape {arr_shape} with {FOUR_DIMENSIONAL} dimensions, "
"but the first dimension (channels) has size "
f"{arr_shape[0]}. Expected a single-channel 3D array.",
)
if (
len(arr_shape) >= FIVE_OR_MORE_DIMENSIONS
and arr_shape[0] > 1
and arr_shape[1] > 1
):
raise ContractError(
f"Input array has shape {arr_shape} with {len(arr_shape)} dimensions. "
f"Expected {EXPECTED_SPATIAL_DIMS} dimensions.",
)
for dim_size in arr_shape:
if dim_size <= 0:
raise ContractError(
f"Input array has shape {arr_shape} with non-positive dimension size. "
"All dimensions must have size greater than 0.",
)
if sum(arr_shape) == len(arr_shape):
raise ContractError(
f"Input array has shape {arr_shape} with all dimensions equal to 1. "
"Expected all three dimensions to have size greater than 1.",
)
return True
[docs]
@beartype
def validate_image_array_type_contracts(
arr: np.ndarray,
) -> bool:
"""Validate the input array for type
Parameters
----------
arr : np.ndarray
Input array to validate
Returns
-------
bool
The status of the validation
Raises
------
ContractError
If the input array does not meet the expected contract
"""
if not isinstance(arr, np.ndarray):
raise ContractError(f"Input is of type {type(arr)}, expected a numpy array.")
# check for numeric dtype (int or float) in the array
if not np.issubdtype(arr.dtype, np.number):
raise ContractError(
f"Input array has dtype {arr.dtype}, expected a numeric dtype "
"(int or float).",
)
return True
[docs]
@beartype
def validate_return_with_pydantic(
result: dict[str, object],
) -> ReturnSchemaModel:
"""Validate return schema using Pydantic model.
Parameters
----------
result : dict[str, object]
Return result to validate
Returns
-------
ReturnSchemaModel
Validated return schema model
Raises
------
ContractError
If validation fails
"""
try:
return ReturnSchemaModel(result=result)
except Exception as e:
msg = (
"Return schema validation failed. Please ensure that the data "
f"fit the expected schema: {e}"
)
raise ContractError(msg)
[docs]
def validate_image_with_pydantic(arr: np.ndarray) -> ImageArrayModel:
"""Validate the input image array using Pydantic model.
Parameters
----------
arr : np.ndarray
Input image array to validate
Returns
-------
ImageArrayModel
Validated image array model
Raises
------
ContractError
If validation fails
"""
try:
return ImageArrayModel(array=arr)
except Exception as e:
msg = (
"Image array validation failed. Please ensure that the input "
f"array meets the expected contracts: {e}"
)
raise ContractError(msg)
[docs]
@beartype
def validate_column_name_with_pydantic(column_name: str) -> ColumnNameModel:
"""Validate column name using Pydantic model.
Parameters
----------
column_name : str
Column name to validate
Returns
-------
ColumnNameModel
Validated column name model with parsed components
Raises
------
ContractError
If validation fails
"""
try:
return ColumnNameModel(column_name=column_name)
except Exception as e:
raise ContractError(f"Column name validation failed: {e}")
[docs]
def create_image_array_schema() -> pa.SeriesSchema:
"""Create a Pandera schema for image array validation.
Returns
-------
pa.SeriesSchema
Pandera schema for numeric arrays
"""
# Use a single numeric dtype for the series schema; Pandera dtype
# objects should be instantiated rather than combined with `|`.
return pa.SeriesSchema(
dtype=pa.Float64(),
name="image_array",
checks=[pa.Check(lambda x: x is not None, error="Value cannot be None")],
)
[docs]
@beartype
def validate_return_schema_contract(
result: dict[str, object],
) -> bool:
"""Validate return schema keys, types, and deterministic key ordering."""
if not isinstance(result, dict):
raise ContractError(f"Return result must be a dict, got {type(result)}.")
actual_keys = tuple(result.keys())
if actual_keys != REQUIRED_RETURN_KEYS:
raise ContractError(
"Return result keys must match required deterministic order "
f"{REQUIRED_RETURN_KEYS}, got {actual_keys}.",
)
try:
ReturnSchemaModel(result=result)
except Exception as e:
raise ContractError(f"Return result validation failed: {e}")
return True
[docs]
class ExpectedFeatureNameValues(BaseModel):
"""Pydantic model for expected values in feature naming validation."""
compartments: list[str] | None = Field(default_factory=list)
channels: list[str] | None = Field(default_factory=list)
features: list[str] | None = Field(default_factory=list)
expected_values_dict: dict[str, list[str]] = Field(default_factory=dict)
model_config = ConfigDict(arbitrary_types_allowed=True)
def __init__(self, **data: object) -> None:
super().__init__(**data)
if self.compartments is not None:
self.compartments = list(set(self.compartments))
else:
raise ValueError("Compartments list cannot be None.")
if self.channels is not None:
# Add "NoChannel" to channels list
self.channels = list(set(self.channels) | {"NoChannel"})
else:
raise ValueError("Channels list cannot be None.")
if self.features is not None and len(self.features) > 0:
self.features = list(set(self.features))
else:
self.features = FEATURES
self.expected_values_dict = {
"compartments": self.compartments,
"channels": self.channels,
"features": self.features,
}
[docs]
@beartype
def validate_column_name_schema(
column_name: str,
channels: list[str],
compartments: list[str],
features: list[str] | None = None,
) -> bool:
"""Validate the column name schema for required fields and types
Parameters
----------
column_name : str
The column name to validate
channels : list[str]
List of valid channels for feature naming
compartments : list[str]
List of valid compartments for feature naming
features : list[str] | None, optional
List of valid features for feature naming, by default None
Returns
-------
bool
The status of the validation
Raises
------
ContractError
If the column name does not meet the expected schema
"""
non_metadata_underscore_separated_parts = NON_METADATA_UNDERSCORE_SEPARATED_PARTS
metadata_underscore_separated_parts = METADATA_UNDERSCORE_SEPARATED_PARTS
expected_values = ExpectedFeatureNameValues(
channels=channels,
compartments=compartments,
features=None,
).expected_values_dict
# check if the column name is a string
if not isinstance(column_name, str):
raise ContractError(f"Column name must be a string, got {type(column_name)}")
# check if the column name has at least 4 parts separated by underscores
parts = column_name.split("_")
if (
len(parts) < non_metadata_underscore_separated_parts
and "Metadata" not in column_name
):
msg = (
"Column name must have at least "
f"{non_metadata_underscore_separated_parts} "
"parts separated by underscores, "
f"got {len(parts)} parts in '{column_name}'"
)
raise ContractError(msg)
if "Metadata" in column_name:
if len(parts) < metadata_underscore_separated_parts:
raise ContractError(
"Metadata column name must have at least "
f"{metadata_underscore_separated_parts} "
"parts separated by "
f"underscores, got {len(parts)} parts in '{column_name}'",
)
return True
feature_components = pd.DataFrame(
[
{
"compartment": parts[0],
"channel": parts[1],
"feature": parts[2],
},
],
)
feature_component_schema = pa.DataFrameSchema(
{
"compartment": pa.Column(
str,
checks=pa.Check.isin(expected_values.get("compartments", [])),
nullable=False,
coerce=True,
),
"channel": pa.Column(
str,
checks=pa.Check.isin(expected_values.get("channels", [])),
nullable=False,
coerce=True,
),
"feature": pa.Column(
str,
checks=pa.Check.isin(expected_values.get("features", [])),
nullable=False,
coerce=True,
),
},
strict=True,
)
try:
feature_component_schema.validate(feature_components)
except (pa.errors.SchemaError, pa.errors.SchemaErrors) as e:
raise ContractError(f"Column name schema validation failed: {e}") from e
return True