Add OptogenSIMLight: Monte Carlo–based light model

.gitignore

@@ -153,4 +153,6 @@ dmypy.json
 .pytype/
 
 # Cython debug symbols
-cython_debug/
+cython_debug/
+.DS_Store
+**/.DS_Store

cleo/light/__init__.py

@@ -5,6 +5,7 @@
     OpticFiber,
     fiber473nm,
 )
+from cleo.light.optogensim_light import OptogenSIMLight
 from cleo.light.light_dependence import (
     LightDependent,
     cubic_interpolator,

cleo/light/optogensim_light.py

@@ -0,0 +1,91 @@
+import xarray as xr
+import numpy as np
+from attrs import define, field
+from brian2.units import Quantity, nmeter, um, cm, mm
+from jaxtyping import Float
+from cleo.light.light import LightModel
+from cleo.utilities import uniform_cylinder_rθz, xyz_from_rθz
+from cleo.coords import coords_from_xyz
+from importlib.resources import files
+
+
+@define
+class OptogenSIMLight(LightModel):
+    """Light model from OptogenSIM Monte Carlo simulations."""
+
+    wavelength: Quantity = 473 * nmeter
+    """Light wavelength. Must be within the simulated grid range."""
+    beam_radius: Quantity = 100 * um
+    """Beam (1/e^2) radius. Must be within the simulated grid range."""
+    data_path: str = None
+    """Path to the 4D dataset. If None, uses the dataset packaged with Cleo."""
+    _rz_slice: object = field(init=False, default=None, repr=False)
+    _r_range: object = field(init=False, default=None, repr=False)
+    _z_range: object = field(init=False, default=None, repr=False)
+
+    def __attrs_post_init__(self):
+        if self.data_path is None:
+            path = str(files("cleo.light.data") / "light_model_4d.nc.gz")
+        else:
+            path = self.data_path
+        rz = xr.open_dataarray(path, engine="scipy").interp(
+            wavelength=self.wavelength / nmeter, beam_size=self.beam_radius / um
+        )
+        # data measures z from the atlas top; re-zero to the source
+        z_source = float(rz.isel(r=0).z[rz.isel(r=0).argmax("z")])
+        self._rz_slice = rz.assign_coords(z=rz.z - z_source)
+        self._r_range = (float(self._rz_slice.r.min()), float(self._rz_slice.r.max()))
+        self._z_range = (float(self._rz_slice.z.min()), float(self._rz_slice.z.max()))
+
+    def transmittance(self, source_coords, source_dir_uvec, target_coords):
+        assert np.allclose(np.linalg.norm(source_dir_uvec, axis=-1), 1)
+        r, z = self._get_rz_for_xyz(source_coords, source_dir_uvec, target_coords)
+        r_cm = np.asarray(r / cm)
+        z_cm = np.asarray(z / cm)
+        T = self._rz_slice.interp(
+            r=xr.DataArray(np.clip(r_cm, *self._r_range)),
+            z=xr.DataArray(np.clip(z_cm, *self._z_range)),
+        ).values
+        T = np.nan_to_num(T, nan=0.0)
+        T[z_cm < self._z_range[0]] = 0  # zero only beyond data range; keep real backscatter within it
+        return T
+
+    @property
+    def area0(self):
+        return np.pi * self.beam_radius**2
+
+    def viz_params(self, coords, direction, T_threshold,
+                   n_points_per_source=4000, **kwargs):
+        r_thresh, zc_thresh = self._find_rz_thresholds(T_threshold)
+        r, theta, zc = uniform_cylinder_rθz(n_points_per_source, r_thresh, zc_thresh)
+        end = coords + zc_thresh * direction
+        x, y, z = xyz_from_rθz(r, theta, zc, coords, end)
+        density_factor = 3
+        cyl_vol = np.pi * r_thresh**2 * zc_thresh
+        markersize = (cyl_vol / n_points_per_source * density_factor) ** (1 / 3)
+        intensity_scale = 1.5 * (4e3 / n_points_per_source) ** (1 / 3)
+        return coords_from_xyz(x, y, z), markersize, intensity_scale
+
+    def _find_rz_thresholds(self, thresh):
+        """Find r and z extents for visualization at the given transmittance threshold."""
+        # On-axis (r[0]) profile along z (re-zeroed so source at z=0).
+        on_axis = self._rz_slice.isel(r=0)
+        z_vals = on_axis.z.values  # cm, source at 0
+        T_z = on_axis.values
+        # Only look in the forward (z >= 0) direction.
+        mask = z_vals >= 0
+        z_pos = z_vals[mask]
+        T_pos = T_z[mask]
+        below = np.where(T_pos < thresh)[0]
+        zc_thresh = (z_pos[below[0]] if len(below) > 0 else z_pos[-1]) * cm
+
+        # Find r threshold at z = zc_thresh / 2.
+        z_mid = float(zc_thresh / cm) / 2
+        z_mid = np.clip(z_mid, float(self._rz_slice.z.min()),
+                        float(self._rz_slice.z.max()))
+        r_vals = self._rz_slice.r.values
+        T_r = self._rz_slice.interp(z=float(z_mid)).values
+        below_r = np.where(T_r < thresh)[0]
+        r_thresh = (r_vals[below_r[0]] if len(below_r) > 0 else r_vals[-1]) * cm
+
+        return r_thresh * 1.2, zc_thresh

pyproject.toml

@@ -45,6 +45,7 @@ dev = [
 
 [tool.poetry]
 packages = [{ include = "cleo" }]
+include = [{ path = "cleo/light/data/*.nc.gz", format = ["sdist", "wheel"] }]
 classifiers = [
     "Development Status :: 4 - Beta",
     "Intended Audience :: Science/Research",

tests/light/test_light.py

@@ -4,7 +4,7 @@
 from brian2 import Network, NeuronGroup, asarray, mm, mm2, ms, mwatt, nmeter, np, um
 
 import cleo
-from cleo.light import GaussianEllipsoid, KoehlerBeam, Light, LightModel, fiber473nm
+from cleo.light import GaussianEllipsoid, KoehlerBeam, Light, LightModel, OptogenSIMLight, fiber473nm
 from cleo.utilities import normalize_coords, unit_safe_allclose
 
 
@@ -101,6 +101,43 @@ def test_OpticFiber():
     assert np.all(T == 0)
 
 
+def test_OptogenSIMLight():
+    from brian2 import mm2, cm
+
+    model = OptogenSIMLight()  # defaults: 473 nm, 100 um beam radius
+    source_coords = np.array([0, 0, 0]) * mm
+    source_direction = normalize_coords([0, 0, 1])  # pointing +z into tissue
+
+    # transmittance decreases with depth along the beam axis (within range)
+    depths = np.array([[0, 0, d] for d in [0.01, 0.05, 0.1, 0.2, 0.4]]) * mm
+    T_depth = model.transmittance(source_coords, source_direction, depths)
+    assert np.all(np.diff(T_depth.ravel()) < 0)  # strictly decreasing
+
+    # behind the source: real (small) backscatter within the data's negative-z
+    # range, but exactly 0 beyond it (the clip must not freeze a nonzero value)
+    peak = model.transmittance(
+        source_coords, source_direction, np.array([[0, 0, 0]]) * mm
+    )
+    near_behind = np.array([[0, 0, -0.05]]) * mm  # 0.05 mm behind source
+    T_near = model.transmittance(source_coords, source_direction, near_behind)
+    assert np.all(T_near >= 0)
+    assert np.all(T_near < peak)  # backscatter is weaker than the forward peak
+
+    z_back = model._z_range[0]  # cm, most-negative z with data (~ -0.29 cm)
+    far_behind = np.array([[0, 0, z_back - 0.05]]) * cm  # beyond the data range
+    T_far = model.transmittance(source_coords, source_direction, far_behind)
+    assert np.all(T_far == 0)
+
+    # transmittance does not increase with radial distance (at fixed depth)
+    radial = np.array([[r, 0, 0.05] for r in [0.0, 0.02, 0.05, 0.1, 0.2]]) * mm
+    T_radial = model.transmittance(source_coords, source_direction, radial)
+    assert np.all(np.diff(T_radial.ravel()) <= 0)  # non-increasing (r=0 plateau ok)
+
+    # area0 = pi * beam_radius^2
+    expected = np.pi * (100 * um) ** 2
+    assert np.isclose(float(model.area0 / mm2), float(expected / mm2))
+
+
 def test_reset():
     light = Light(light_model=fiber473nm())
     assert light.value == 0
@@ -152,7 +189,7 @@ def test_light_power_irradiance(n_coords, values, shape):
 
 
 @pytest.mark.parametrize(
-    "light_model", [fiber473nm(), GaussianEllipsoid(), KoehlerBeam(1 * mm)]
+    "light_model", [fiber473nm(), GaussianEllipsoid(), KoehlerBeam(1 * mm), OptogenSIMLight()]
 )
 @pytest.mark.parametrize(
     "m, squeeze_coords, squeeze_dir",