Add minimal gamma correction transform for experimenting with whether changing numerical ranges help training

src/virtual_stain_flow/transforms/__init__.py

@@ -3,13 +3,15 @@
 """
 
 from .channelwise import ChannelwiseTransform
+from .gamma import ContinuousGammaTransform
 from .normalizations import (
     MaxScaleNormalize,
     ZScoreNormalize,
 )
 
 __all__ = [
     "ChannelwiseTransform",
+    "ContinuousGammaTransform",
     "MaxScaleNormalize",
     "ZScoreNormalize",
 ]

src/virtual_stain_flow/transforms/gamma.py

@@ -0,0 +1,164 @@
+"""Fixed, reversible gamma correction of normalized microscopy intensities."""
+
+from numbers import Real
+
+import numpy as np
+
+from .base_transform import LoggableTransform
+
+
+class ContinuousGammaTransform(LoggableTransform):
+	"""
+	Apply a continuous, strictly increasing power curve on a fixed interval.
+
+	For ``lower < x < upper``, with ``t = (x - lower) / (upper - lower)``,
+	the output is ``lower + (upper - lower) * t**gamma``. All other values
+	are unchanged. The interval endpoints are fixed, making the mapping
+	continuous (without jumps in pixel values). 
+
+	``gamma < 1`` expands contrast near ``lower`` and compresses it near
+	``upper``; ``gamma > 1`` does the reverse; ``gamma == 1`` is identity.
+	No image statistics are estimated. The inverse uses ``1 / gamma`` on
+	the same interval. Extreme powers and finite precision can lose detail.
+
+	:param gamma: Positive, finite correction exponent.
+	:param lower: Lower fixed point of the curve.
+	:param upper: Upper fixed point of the curve.
+	"""
+
+	def __init__(
+		self,
+		gamma: float,
+		lower: float = 0.0,
+		upper: float = 1.0,
+		name: str = "ContinuousGammaTransform",
+		p: float = 1.0,
+	):
+		"""
+		Initialize the continuous gamma transform with the given parameters.
+
+		:param gamma: Positive, finite correction exponent.
+		:param lower: Lower fixed point of the curve.
+		:param upper: Upper fixed point of the curve.
+		:param name: Name used for logging.
+		:param p: Must be 1.0 for deterministic, reversible correction.
+		"""
+		gamma = self._finite_real(gamma, "gamma")
+		lower = self._finite_real(lower, "lower")
+		upper = self._finite_real(upper, "upper")
+		p = self._finite_real(p, "p")
+		if gamma <= 0:
+			raise ValueError("gamma must be positive with finite reciprocal values.")
+		inverse_gamma = 1.0 / gamma
+		with np.errstate(over="ignore"):
+			inverse_inverse_gamma = 1.0 / inverse_gamma
+		if (
+			not np.isfinite(inverse_gamma)
+			or not np.isfinite(inverse_inverse_gamma)
+		):
+			raise ValueError("gamma must be positive with finite reciprocal values.")
+		if not 0.0 <= lower < upper <= 1.0:
+			raise ValueError("Expected 0 <= lower < upper <= 1.")
+		if p != 1.0:
+			raise ValueError("p must be 1.0 for deterministic, reversible correction.")
+
+		super().__init__(name=name, p=p)
+		self._gamma = gamma
+		self._lower = lower
+		self._upper = upper
+
+	@staticmethod
+	def _finite_real(value: float, name: str) -> float:
+		"""
+		Ensure value is a finite real number.
+
+		:param value: The value to check.
+		:param name: Name of the parameter (for error messages).
+		:return: The validated finite real number.
+		"""
+		if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real):
+			raise TypeError(f"{name} must be a real number.")
+		value = float(value)
+		if not np.isfinite(value):
+			raise ValueError(f"{name} must be finite.")
+		return value
+
+	@property
+	def gamma(self) -> float:
+		"""Power applied within the interval."""
+		return self._gamma
+
+	@property
+	def lower(self) -> float:
+		"""Lower fixed point of the curve."""
+		return self._lower
+
+	@property
+	def upper(self) -> float:
+		"""Upper fixed point of the curve."""
+		return self._upper
+
+	def apply(self, img: np.ndarray, **params) -> np.ndarray:
+		"""Correct normalized intensities without clipping or estimating statistics."""
+		if not isinstance(img, np.ndarray):
+			raise TypeError("Expected input image to be a NumPy array.")
+		if not np.issubdtype(img.dtype, np.floating):
+			raise TypeError("Expected a floating-point image; apply MaxScaleNormalize first.")
+		if not np.all(np.isfinite(img)):
+			raise ValueError("Input image must contain only finite values.")
+		if np.any(img < 0.0) or np.any(img > 1.0):
+			raise ValueError("Input image must be normalized to [0, 1]; values are not clipped.")
+
+		output = img.astype(np.result_type(img.dtype, np.float32), copy=True)
+		if self.gamma == 1.0:
+			return output
+
+		# Compute only the open interval, preserving fixed points and all
+		# unaffected values exactly; fractional powers never see negatives.
+		work = img.astype(np.result_type(img.dtype, np.float64), copy=False)
+		# Use representable endpoints in both directions. Otherwise a rounded
+		# float32 lower endpoint can be mistaken for an interior value, whose
+		# tiny offset is greatly amplified by a fractional power.
+		lower = output.dtype.type(self.lower).item()
+		upper = output.dtype.type(self.upper).item()
+		active = (work > lower) & (work < upper)
+		if not np.any(active):
+			return output
+		width = upper - lower
+		t = (work[active] - lower) / width
+		output[active] = lower + width * np.power(t, self.gamma)
+		return output
+
+	def inverse(self) -> "ContinuousGammaTransform":
+		"""
+		Return the deterministic analytic inverse, callable like any transform.
+
+		This undoes gamma correction only. To undo preceding max scaling,
+		multiply the recovered normalized image by its normalization factor.
+		Values clipped during max scaling cannot be recovered.
+		"""
+		return type(self)(
+			gamma=1.0 / self.gamma,
+			lower=self.lower,
+			upper=self.upper,
+			name=f"{self.name}_inverse",
+		)
+
+	def __repr__(self) -> str:
+		return (
+			f"{self.__class__.__name__}(name={self.name}, gamma={self.gamma}, "
+			f"lower={self.lower}, upper={self.upper}, p={self.p})"
+		)
+
+	def to_config(self) -> dict:
+		"""Return JSON-serializable parameters for inherited from_config()."""
+		return {
+			"class": self.__class__.__name__,
+			"name": self.name,
+			"params": {
+				"gamma": self.gamma,
+				"lower": self.lower,
+				"upper": self.upper,
+				"p": self.p,
+			},
+		}

tests/transforms/test_gamma.py

@@ -0,0 +1,116 @@
+"""Tests for fixed interval gamma correction and its analytic inverse."""
+
+
+import numpy as np
+import pytest
+
+from virtual_stain_flow.transforms import (
+    ChannelwiseTransform,
+    ContinuousGammaTransform,
+)
+
+
+def test_defaults_and_repr():
+    transform = ContinuousGammaTransform(gamma=0.5)
+    assert (transform.gamma, transform.lower, transform.upper, transform.p) == (
+        0.5, 0.0, 1.0, 1.0
+    )
+    assert repr(transform) == (
+        "ContinuousGammaTransform(name=ContinuousGammaTransform, gamma=0.5, "
+        "lower=0.0, upper=1.0, p=1.0)"
+    )
+
+
+def test_rejects_gamma_whose_inverse_cannot_be_inverted():
+    with pytest.raises(ValueError, match="finite reciprocal"):
+        ContinuousGammaTransform(gamma=np.finfo(np.float64).max)
+
+
+def test_gamma_serialization_round_trip():
+    transform = ContinuousGammaTransform(
+        gamma=0.5,
+        lower=0.1,
+        upper=0.9,
+        name="gamma_custom",
+    )
+    config = transform.to_config()
+    restored = ContinuousGammaTransform.from_config(config)
+
+    assert isinstance(restored, ContinuousGammaTransform)
+    assert restored.name == transform.name
+    assert (restored.gamma, restored.lower, restored.upper, restored.p) == (
+        transform.gamma,
+        transform.lower,
+        transform.upper,
+        transform.p,
+    )
+
+    image = np.linspace(0, 1, 101).reshape(1, 101)
+    np.testing.assert_allclose(restored.apply(image), transform.apply(image), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("gamma", [0.25, 0.5, 1.0, 2.0, 4.0])
+def test_default_interval_is_power_law(gamma):
+    image = np.linspace(0, 1, 101).reshape(1, 101)
+    actual = ContinuousGammaTransform(gamma=gamma)(image=image)["image"]
+    np.testing.assert_allclose(actual, image**gamma, rtol=1e-14, atol=0)
+
+
+def test_selected_interval_and_known_values():
+    image = np.array([0, 0.1, 0.125, 0.2, 0.325, 0.5, 0.8, 1.0])
+    transform = ContinuousGammaTransform(gamma=0.5, lower=0.1, upper=0.5)
+    actual = transform.apply(image)
+    np.testing.assert_allclose(actual, [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.8, 1])
+    np.testing.assert_array_equal(actual[[0, 1, 5, 6, 7]], image[[0, 1, 5, 6, 7]])
+
+
+@pytest.mark.parametrize("gamma", [0.5, 1.0, 2.0])
+def test_strict_monotonicity_continuity_and_fixed_points(gamma):
+    transform = ContinuousGammaTransform(gamma=gamma, lower=0.1, upper=0.7)
+    x = np.linspace(0, 1, 10001)
+    y = transform.apply(x)
+    assert np.all(np.diff(y) > 0)
+    assert y.min() == 0 and y.max() == 1
+    for boundary in (0.1, 0.7):
+        adjacent = np.array([
+            np.nextafter(boundary, 0), boundary, np.nextafter(boundary, 1)
+        ])
+        result = transform.apply(adjacent)
+        assert result[1] == boundary
+        np.testing.assert_allclose(result, boundary, atol=1e-8, rtol=0)
+
+
+@pytest.mark.parametrize("gamma", [0.25, 0.5, 1.0, 2.0, 4.0])
+@pytest.mark.parametrize("bounds", [(0.0, 1.0), (0.1, 0.7)])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_inverse_round_trip_both_directions(gamma, bounds, dtype):
+    # Stay away from a nonzero lower endpoint for large powers: rounding there
+    # can discard the tiny offset, a finite-precision rather than analytic limit.
+    image = np.array([0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.9, 1], dtype=dtype)
+    transform = ContinuousGammaTransform(gamma, lower=bounds[0], upper=bounds[1])
+    inverse = transform.inverse()
+    assert inverse.gamma == 1 / gamma
+    assert (inverse.lower, inverse.upper, inverse.p) == (*bounds, 1.0)
+    tolerance = 2e-6 if dtype == np.float32 else 1e-13
+    np.testing.assert_allclose(
+        inverse.apply(transform.apply(image)), image, atol=tolerance, rtol=0
+    )
+    np.testing.assert_allclose(
+        transform.apply(inverse.apply(image)), image, atol=tolerance, rtol=0
+    )
+
+
+@pytest.mark.parametrize("value", [np.nan, np.inf, -np.inf, -0.01, 1.01])
+def test_invalid_pixels_rejected_not_clipped(value):
+    with pytest.raises(ValueError, match="finite values|normalized"):
+        ContinuousGammaTransform(0.5).apply(np.array([0.0, value, 1.0]))
+
+
+def test_channelwise_correction_and_unaffected_channel():
+    image = np.linspace(0, 1, 48).reshape(3, 4, 4)
+    gamma = ContinuousGammaTransform(0.5)
+    wrapper = ChannelwiseTransform([gamma, None, ContinuousGammaTransform(2)])
+    corrected = wrapper.apply(image)
+    np.testing.assert_allclose(corrected[0], np.sqrt(image[0]))
+    np.testing.assert_array_equal(corrected[1], image[1])
+    np.testing.assert_allclose(corrected[2], image[2] ** 2)