allow any object with render method to be used for label distribution (#83)

* allow any object with render method to be used for label distribution

* make base public

Author: tlambert03

src/microsim/schema/sample/__init__.py

@@ -1,13 +1,15 @@
+from ._distributions._base import BaseDistribution
 from ._distributions.cosem import CosemLabel
 from ._distributions.matslines import MatsLines
 from .fluorophore import Fluorophore
-from .sample import Distribution, FluorophoreDistribution, Sample
+from .sample import AnyDistribution, FluorophoreDistribution, Sample
 
 __all__ = [
-    "MatsLines",
-    "Sample",
+    "AnyDistribution",
+    "BaseDistribution",
     "CosemLabel",
-    "Distribution",
-    "FluorophoreDistribution",
     "Fluorophore",
+    "FluorophoreDistribution",
+    "MatsLines",
+    "Sample",
 ]

src/microsim/schema/sample/_distributions/_base.py

@@ -1,16 +1,43 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Annotated, Any, Protocol
 
-from pydantic import BaseModel
+from pydantic import BaseModel, GetCoreSchemaHandler
+from pydantic_core import core_schema
+from typing_extensions import runtime_checkable
 
 if TYPE_CHECKING:
     from microsim._data_array import xrDataArray
     from microsim.schema.backend import NumpyAPI
 
 
-class _BaseDistribution(BaseModel, ABC):
+@runtime_checkable
+class Renderable(Protocol):
+    @abstractmethod
+    def render(self, space: xrDataArray, xp: NumpyAPI | None = None) -> xrDataArray:
+        """Render the distribution into the given space."""
+
+
+class _IsInstanceAnySer:
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        def _validate(obj: Any) -> Any:
+            if not isinstance(obj, source_type):
+                raise ValueError(f"Expected {source_type}, got {type(obj)}")
+            return obj
+
+        return core_schema.no_info_before_validator_function(
+            _validate, core_schema.any_schema()
+        )
+
+
+RenderableType = Annotated[Renderable, _IsInstanceAnySer()]
+
+
+class BaseDistribution(BaseModel, ABC):
     @classmethod
     def is_random(cls) -> bool:
         """Return True if this distribution generates randomized results."""

src/microsim/schema/sample/_distributions/cosem.py

@@ -7,7 +7,7 @@
 from microsim.cosem.models import CosemDataset, CosemImage
 from microsim.schema.backend import NumpyAPI
 
-from ._base import _BaseDistribution
+from ._base import BaseDistribution
 
 if TYPE_CHECKING:
     from microsim._data_array import xrDataArray
@@ -25,7 +25,7 @@ def _validate_dataset(v: Any) -> CosemDataset:
 Dataset = Annotated[CosemDataset, BeforeValidator(_validate_dataset)]
 
 
-class CosemLabel(_BaseDistribution):
+class CosemLabel(BaseDistribution):
     """Renders ground truth based on a specific layer from a COSEM dataset.
 
     Go to https://openorganelle.janelia.org/datasets/ to find a dataset, and then

src/microsim/schema/sample/_distributions/direct.py

@@ -4,13 +4,13 @@
 
 from microsim.schema.backend import NumpyAPI
 
-from ._base import _BaseDistribution
+from ._base import BaseDistribution
 
 if TYPE_CHECKING:
     from microsim._data_array import xrDataArray
 
 
-class FixedArrayTruth(_BaseDistribution):
+class FixedArrayTruth(BaseDistribution):
     type: Literal["fixed-array"] = "fixed-array"
     array: Any
 

src/microsim/schema/sample/_distributions/matslines.py

@@ -5,15 +5,15 @@
 import numpy as np
 
 from microsim.schema.backend import NumpyAPI
-from microsim.schema.sample._distributions._base import _BaseDistribution
+from microsim.schema.sample._distributions._base import BaseDistribution
 
 if TYPE_CHECKING:
     import numpy.typing as npt
 
     from microsim._data_array import xrDataArray
 
 
-class MatsLines(_BaseDistribution):
+class MatsLines(BaseDistribution):
     type: Literal["matslines"] = "matslines"
     density: float = 1
     length: int = 10

src/microsim/schema/sample/sample.py

@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from typing import Any, get_args
+from typing import Any
 
 import numpy as np
 from pydantic import Field, model_validator
@@ -9,14 +9,15 @@
 from microsim.schema.backend import NumpyAPI
 from microsim.schema.spectrum import Spectrum
 
+from ._distributions._base import Renderable, RenderableType
 from ._distributions.cosem import CosemLabel
 from ._distributions.direct import FixedArrayTruth
 from ._distributions.matslines import MatsLines
 from .fluorophore import Fluorophore
 
-Distribution = MatsLines | CosemLabel | FixedArrayTruth
-DistributionTypes = get_args(Distribution)
-
+AnyDistribution = MatsLines | CosemLabel | FixedArrayTruth | RenderableType
+# TODO: this feels like an unDRY hack
+DistributionTypes = MatsLines | CosemLabel | FixedArrayTruth | Renderable
 
 # This is a placeholder fluorophore for when no fluorophore is specified
 # it has broad excitation and emission spectra, high extinction coefficient.
@@ -34,7 +35,7 @@
 
 
 class FluorophoreDistribution(SimBaseModel):
-    distribution: Distribution = Field(...)
+    distribution: AnyDistribution = Field(union_mode="left_to_right")
     fluorophore: Fluorophore = MOCK_FLUOR
     # either a scalar that will be multiplied by the distribution
     # (e.g. to increase/decrease concentration of fluorophore)
@@ -47,7 +48,7 @@ def __hash__(self) -> int:
 
     def cache_path(self) -> tuple[str, ...] | None:
         if hasattr(self.distribution, "cache_path"):
-            return self.distribution.cache_path()
+            return self.distribution.cache_path()  # type: ignore [no-any-return]
         return None
 
     def render(self, space: xrDataArray, xp: NumpyAPI | None = None) -> xrDataArray:

src/microsim/schema/simulation.py

@@ -307,7 +307,7 @@ def _truth_cache_path(
         scale = f'scale{"_".join(str(x) for x in truth_space.scale)}'
         conc = f"conc{label.concentration}"
         truth_cache = truth_cache / shape / scale / conc
-        if label.distribution.is_random():
+        if hasattr(label.distribution, "is_random") and label.distribution.is_random():
             truth_cache = truth_cache / f"seed{seed}"
         return truth_cache
 

tests/test_simulation.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import pytest
+from pydantic import ValidationError
 
 import microsim.schema as ms
 from microsim.schema.optical_config.lib import FITC
@@ -120,6 +121,28 @@ def test_simulation_from_ground_truth() -> None:
     np.testing.assert_array_almost_equal(sim_truth, ground_truth)
 
 
+def test_simulation_custom_distribution() -> None:
+    """Test that we can use any class with a render method as a distribution."""
+
+    class BadDistribution:
+        def have_no_render_method(self, space, xp: ms.NumpyAPI | None = None):
+            return space
+
+    with pytest.raises(ValidationError):
+        ms.Sample(labels=[BadDistribution()])
+
+    class GoodDistribution:
+        def render(self, space, xp: ms.NumpyAPI | None = None):
+            return space
+
+    sim = ms.Simulation(
+        truth_space=ms.ShapeScaleSpace(shape=(64, 128, 128), scale=(0.2, 0.1, 0.1)),
+        output_space={"downscale": 1},
+        sample=ms.Sample(labels=[GoodDistribution()]),
+    )
+    sim.run()
+
+
 def test_pickle(sim1: ms.Simulation) -> None:
     pickled = pickle.dumps(sim1)
     assert pickle.loads(pickled) == sim1