fix: support upscaling output space (tlambert03#121)

* fix: support upscaling output space

* refactor: update simulation parameters for upscale and downscale handling; improve test coverage

* change implementation

Author: tlambert03

examples/basic_confocal.py

@@ -1,9 +1,9 @@
 from microsim import schema as ms
-from microsim.util import ndview, ortho_plot
+from microsim.util import ortho_plot
 
 sim = ms.Simulation(
-    truth_space=ms.ShapeScaleSpace(shape=(64, 256, 256), scale=(0.04, 0.02, 0.02)),
-    output_space={"downscale": 4},
+    truth_space={"upscale": 4},
+    output_space=ms.ShapeScaleSpace(shape=(16, 64, 64), scale=(0.16, 0.08, 0.08)),
     sample=[
         ms.FluorophoreDistribution(
             distribution=ms.MatsLines(density=0.5, length=30, azimuth=5, max_r=1),
@@ -15,9 +15,7 @@
     modality=ms.Confocal(pinhole_au=0.5),
     settings=ms.Settings(random_seed=100, max_psf_radius_aus=8),
     detector=ms.CameraCCD(qe=0.82, read_noise=2, bit_depth=12),
-    # output_path="au1.tif",
 )
 
 result = sim.run()
-ortho_plot(result.data, show=True)
-ndview(result)
+ortho_plot(result)

pyproject.toml

@@ -155,6 +155,9 @@ init_typed = false
 module = ["microsim.xarray_jax.*"]
 ignore_errors = true
 
+[tool.pyright]
+reportArgumentType = false  # too hard with pydantic
+
 # https://docs.pytest.org/en/6.2.x/customize.html
 [tool.pytest.ini_options]
 minversion = "7.0"

src/microsim/schema/simulation.py

@@ -268,9 +268,9 @@ def digital_image(
             optical_image = self.optical_image()
         image = optical_image  # (C, Z, Y, X)
 
-        # downscale to output space
+        # rescale to output space
         # TODO: consider how we would integrate detector pixel size
-        # rather than a user-sepicified output space
+        # rather than a user-specified output space
         if self.output_space is not None:
             logger.info(f"Rescaling to output space {self.output_space.shape}")
             image = self.output_space.rescale(image)

src/microsim/schema/space.py

@@ -1,5 +1,5 @@
 from collections.abc import Callable, Sequence
-from typing import Any, Protocol, TypeVar
+from typing import Any, Protocol, TypeVar, runtime_checkable
 
 import numpy as np
 from pydantic import (
@@ -34,6 +34,7 @@ def _wrap_validator(
         return np.array(value, dtype=np.float64)
 
 
+@runtime_checkable
 class SpaceProtocol(Protocol):
     @property
     def axes(self) -> tuple[str, ...]: ...
@@ -72,6 +73,39 @@ def coords(self: SpaceProtocol) -> dict[str, FloatArray]:
 class _AxesSpace(_Space):
     axes: tuple[Axis, ...] = (Axis.Z, Axis.Y, Axis.X)
 
+    def _get_scale_ratios(self, img_space: Any) -> dict[str, int]:
+        """Calculate integer scale ratios between two spaces for coarsening."""
+        if not (
+            isinstance(img_space, SpaceProtocol) and isinstance(self, SpaceProtocol)
+        ):  # pragma: no cover
+            raise NotImplementedError(
+                f"Rescaling from {type(img_space)} to {type(self)} is not implemented."
+            )
+
+        if set(self.axes) != set(img_space.axes):  # pragma: no cover
+            raise ValueError(
+                f"Spaces must have the same axes. Got {self.axes} and {img_space.axes}."
+            )
+        # Create axis->scale mappings
+        self_scales = dict(zip(self.axes, self.scale, strict=True))
+        img_scales = dict(zip(img_space.axes, img_space.scale, strict=True))
+        return {
+            ax: int(self_scales[ax] / img_scales[ax])
+            for ax in self.axes
+            if ax in img_scales
+        }
+
+    def rescale(self, img: xrDataArray) -> xrDataArray:
+        if not (img_space := getattr(img, "space", None)):  # pragma: no cover
+            raise ValueError("Input image must have a 'space' attribute.")
+
+        dims = self._get_scale_ratios(img_space)
+        if any(d < 1 for d in dims.values()):  # pragma: no cover
+            raise NotImplementedError(
+                f"Can only downscale an image. Got downscale factors {dims}."
+            )
+        return img.coarsen(dims).sum()  # type: ignore
+
     @field_validator("axes", mode="before")
     def _cast_axes(cls, value: Any) -> tuple[Axis, ...]:
         return tuple(value)

tests/_util.py

@@ -1,9 +1,12 @@
+import os
 import socket
 from collections.abc import Callable
 
 import pytest
 
 try:
+    if os.getenv("MICROSIM_TEST_NO_INTERNET"):
+        raise OSError("Skipping internet test due to MICROSIM_TEST_NO_INTERNET")
     socket.create_connection(("8.8.8.8", 53), timeout=1)
     HAVE_INTERNET = True
 except OSError:

tests/test_scaling.py

@@ -0,0 +1,25 @@
+from microsim import schema as ms
+
+
+def test_downscale_truth() -> None:
+    sample = [
+        ms.FluorophoreDistribution(
+            distribution=ms.MatsLines(density=0.5, length=30, azimuth=5, max_r=1),
+        )
+    ]
+
+    sim_down = ms.Simulation(
+        truth_space=ms.ShapeScaleSpace(shape=(64, 256, 256), scale=(0.04, 0.02, 0.02)),
+        output_space={"downscale": 4},
+        sample=sample,
+        settings=ms.Settings(random_seed=100, max_psf_radius_aus=1),
+    )
+    sim_up = ms.Simulation(
+        truth_space={"upscale": 4},
+        output_space=ms.ShapeScaleSpace(shape=(16, 64, 64), scale=(0.16, 0.08, 0.08)),
+        sample=sample,
+        settings=ms.Settings(random_seed=100, max_psf_radius_aus=1),
+    )
+
+    assert sim_down.ground_truth().equals(sim_up.ground_truth())
+    assert sim_down.digital_image().identical(sim_up.digital_image())