Preserve declared grayscale volume and RGB TIFF semantics

src/polystore/config.py

@@ -133,44 +133,106 @@ def __new__(cls, value: str, tifffile_codec: str | None):
         return member
 
 
+class TiffPhotometric(Enum):
+    """Declared TIFF pixel interpretation, independent of array dimensions."""
+
+    MINISBLACK = "minisblack"
+    RGB = "rgb"
+
+
+class TiffPlanarConfig(Enum):
+    """Declared storage layout for TIFF sample dimensions."""
+
+    CONTIG = "contig"
+    SEPARATE = "separate"
+
+
 @dataclass(frozen=True)
 class TiffConfig:
-    """TIFF writer settings; defaults preserve the uncompressed file contract.
+    """TIFF codec and declared pixel axes; defaults preserve legacy writes.
 
     Args:
         compression: Lossless disk-TIFF codec. ``NONE`` preserves the existing
             uncompressed output; ``DEFLATE`` compresses TIFF pixels exactly.
         compression_level: Deflate level from 1 to 9. Ignored when compression
             is ``NONE``.
+        photometric: Explicit grayscale or RGB pixel interpretation.
+        axes: Dimension identities such as ``ZYX`` or ``YXS``. Requires
+            photometric; ``S`` declares samples rather than a spatial axis.
+        planarconfig: Sample layout. Declared axes select the layout when absent.
     """
 
     compression: TiffCompression = TiffCompression.NONE
     compression_level: int = 1
+    photometric: TiffPhotometric | None = None
+    axes: str | None = None
+    planarconfig: TiffPlanarConfig | None = None
 
     def __post_init__(self) -> None:
         if not isinstance(self.compression, TiffCompression):
             raise TypeError("compression must be a TiffCompression value")
-        if isinstance(self.compression_level, bool) or not isinstance(
-            self.compression_level, int
-        ) or not 1 <= self.compression_level <= 9:
+        if (
+            isinstance(self.compression_level, bool)
+            or not isinstance(self.compression_level, int)
+            or not 1 <= self.compression_level <= 9
+        ):
             raise ValueError("compression_level must be an integer from 1 to 9")
+        if self.photometric is not None and not isinstance(self.photometric, TiffPhotometric):
+            raise TypeError("photometric must be a TiffPhotometric value")
+        if self.planarconfig is not None and not isinstance(self.planarconfig, TiffPlanarConfig):
+            raise TypeError("planarconfig must be a TiffPlanarConfig value")
+        if self.axes is not None:
+            if (
+                not isinstance(self.axes, str)
+                or not self.axes
+                or not self.axes.isascii()
+                or not self.axes.isalpha()
+                or not self.axes.isupper()
+                or len(set(self.axes)) != len(self.axes)
+            ):
+                raise ValueError("axes must contain distinct uppercase dimension identities")
+            if self.photometric is None:
+                raise ValueError("axes require an explicit photometric interpretation")
+            sample_axis = self.axes.find("S")
+            if sample_axis >= 0:
+                if sample_axis == len(self.axes) - 1:
+                    layout = TiffPlanarConfig.CONTIG
+                elif sample_axis == len(self.axes) - 3:
+                    layout = TiffPlanarConfig.SEPARATE
+                else:
+                    raise ValueError("TIFF sample axis must be last or immediately before YX")
+                if self.planarconfig is not None and self.planarconfig is not layout:
+                    raise ValueError("planarconfig disagrees with the declared sample axis")
+                object.__setattr__(self, "planarconfig", layout)
+            elif self.photometric is TiffPhotometric.RGB or self.planarconfig is not None:
+                raise ValueError("RGB and planar TIFF layouts require a declared S sample axis")
+
+    def validate_shape(self, shape: Sequence[int]) -> None:
+        """Reject a declared axis roster that cannot describe the actual pixels."""
+        if self.axes is not None and len(self.axes) != len(shape):
+            raise ValueError(
+                f"Declared TIFF axes {self.axes!r} do not match pixel shape {tuple(shape)!r}"
+            )
 
     def tifffile_write_kwargs(self) -> dict[str, object]:
-        """Return only codec options accepted by tifffile.imwrite."""
-
+        """Encode codec options and declared pixel semantics for tifffile.imwrite."""
+        options: dict[str, object] = {}
         codec = self.compression.tifffile_codec
-        if codec is None:
-            return {}
-        return {
-            "compression": codec,
-            "compressionargs": {"level": self.compression_level},
-        }
+        if codec is not None:
+            options.update(compression=codec, compressionargs={"level": self.compression_level})
+        if self.photometric is not None:
+            options["photometric"] = self.photometric.value
+        if self.axes is not None:
+            options["metadata"] = {"axes": self.axes}
+        if self.planarconfig is not None:
+            options["planarconfig"] = self.planarconfig.value
+        return options
 
     def applies_to_path(self, path: str | Path) -> bool:
-        """Whether this configured codec applies to this disk output path."""
+        """Whether these configured TIFF options apply to this disk path."""
 
         return (
-            self.compression is not TiffCompression.NONE
+            bool(self.tifffile_write_kwargs())
             and Path(path).suffix.lower() in FileFormat.TIFF.extensions
         )
 
@@ -179,11 +241,11 @@ def tiff_write_batches(
     paths: Sequence[str | Path],
     config: TiffConfig | None,
 ) -> tuple[tuple[tuple[int, ...], TiffConfig | None], ...]:
-    """Partition consecutive writes by their exact TIFF codec requirement."""
+    """Partition consecutive writes by their exact TIFF configuration."""
 
     if not paths:
         return ()
-    if config is None or config.compression is TiffCompression.NONE:
+    if config is None or not config.tifffile_write_kwargs():
         return ((tuple(range(len(paths))), None),)
     batches: list[tuple[tuple[int, ...], TiffConfig | None]] = []
     start = 0

src/polystore/disk.py

@@ -201,17 +201,15 @@ def _tensorflow_reader(self, path, **kwargs):
         return tf.io.parse_tensor(tf.io.read_file(path.as_posix()), out_type=tf.dtypes.float32)
 
     def _tiff_writer(self, path, data, *, tiff_config: TiffConfig | None = None, **kwargs):
-        """Write TIFF pixels with optional lossless, per-call codec settings."""
+        """Write pixels with their declared TIFF semantics and codec options."""
 
         tifffile = FileFormat.TIFF.load_dependency()
         config = TiffConfig() if tiff_config is None else tiff_config
         if not isinstance(config, TiffConfig):
             raise TypeError("tiff_config must be a TiffConfig value")
-        tifffile.imwrite(
-            path,
-            storage_numpy_array(data),
-            **config.tifffile_write_kwargs(),
-        )
+        pixels = storage_numpy_array(data)
+        config.validate_shape(pixels.shape)
+        tifffile.imwrite(path, pixels, **config.tifffile_write_kwargs())
 
     def _tiff_reader(self, path):
         tifffile = FileFormat.TIFF.load_dependency()

tests/test_disk_backend.py

@@ -13,7 +13,13 @@
 
 from polystore import FileManager
 from polystore import disk as disk_module
-from polystore.config import TiffCompression, TiffConfig, tiff_write_batches
+from polystore.config import (
+    TiffCompression,
+    TiffConfig,
+    TiffPhotometric,
+    TiffPlanarConfig,
+    tiff_write_batches,
+)
 from polystore.disk import DiskBackend
 from polystore.exceptions import StorageResolutionError
 
@@ -137,15 +143,82 @@ def test_tiff_deflate_is_lossless_and_opt_in(tmp_path: Path) -> None:
     np.testing.assert_array_equal(tifffile.imread(compressed_path), pixels)
 
 
+@pytest.mark.parametrize("depth", [3, 4])
+def test_declared_grayscale_volume_does_not_become_rgb(tmp_path, depth) -> None:
+    pixels = np.arange(depth * 4 * 5, dtype=np.uint16).reshape(depth, 4, 5)
+    path = tmp_path / "volume.tif"
+    DiskBackend().save(
+        pixels,
+        path,
+        tiff_config=TiffConfig(photometric=TiffPhotometric.MINISBLACK, axes="ZYX"),
+    )
+
+    with tifffile.TiffFile(path) as image:
+        assert image.series[0].axes == "ZYX"
+        assert len(image.pages) == depth
+        assert image.pages[0].photometric.name == "MINISBLACK"
+        assert image.pages[0].samplesperpixel == 1
+        assert image.pages[0].compression.value == 1
+        np.testing.assert_array_equal(image.asarray(), pixels)
+
+
+@pytest.mark.parametrize(
+    ("axes", "shape", "layout"),
+    [
+        ("YXS", (4, 5, 3), TiffPlanarConfig.CONTIG),
+        ("SYX", (3, 4, 5), TiffPlanarConfig.SEPARATE),
+        ("ZSYX", (2, 3, 4, 5), TiffPlanarConfig.SEPARATE),
+    ],
+)
+def test_declared_rgb_preserves_samples_and_spatial_axes(tmp_path, axes, shape, layout) -> None:
+    pixels = np.arange(np.prod(shape), dtype=np.uint16).reshape(shape)
+    path = tmp_path / "rgb.tif"
+    DiskBackend().save(
+        pixels,
+        path,
+        tiff_config=TiffConfig(photometric=TiffPhotometric.RGB, axes=axes),
+    )
+
+    with tifffile.TiffFile(path) as image:
+        assert image.series[0].axes == axes
+        assert image.pages[0].photometric.name == "RGB"
+        assert image.pages[0].planarconfig.name == layout.name
+        assert image.pages[0].samplesperpixel == 3
+        np.testing.assert_array_equal(image.asarray(), pixels)
+
+
+def test_declared_axes_must_describe_actual_storage_layout(tmp_path) -> None:
+    path = tmp_path / "invalid.tif"
+    with pytest.raises(ValueError, match="do not match pixel shape"):
+        DiskBackend().save(
+            np.zeros((3, 4, 5)),
+            path,
+            tiff_config=TiffConfig(photometric=TiffPhotometric.MINISBLACK, axes="YX"),
+        )
+    assert not path.exists()
+    with pytest.raises(ValueError, match="disagrees"):
+        TiffConfig(
+            photometric=TiffPhotometric.RGB,
+            axes="SYX",
+            planarconfig=TiffPlanarConfig.CONTIG,
+        )
+    with pytest.raises(ValueError, match="sample axis"):
+        TiffConfig(photometric=TiffPhotometric.RGB, axes="SZYX")
+
+
 def test_tiff_write_batches_preserve_order_and_exclude_other_formats() -> None:
     config = TiffConfig(compression=TiffCompression.DEFLATE)
     assert tiff_write_batches(("a.csv", "b.tif", "c.tiff", "d.npy"), config) == (
         ((0,), None),
         ((1, 2), config),
         ((3,), None),
     )
-    assert tiff_write_batches(("a.csv", "b.tif"), TiffConfig()) == (
-        ((0, 1), None),
+    assert tiff_write_batches(("a.csv", "b.tif"), TiffConfig()) == (((0, 1), None),)
+    declared = TiffConfig(photometric=TiffPhotometric.MINISBLACK, axes="ZYX")
+    assert tiff_write_batches(("a.csv", "b.tif", "c.tiff", "d.npy"), declared) == (
+        ((0,), None),
+        ((1, 2), declared),
+        ((3,), None),
     )
     with pytest.raises(ValueError, match="compression_level"):
         TiffConfig(compression_level=10)