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 import OptogenSIM, defocused_gaussian_beam
 from cleo.light.light_dependence import (
     LightDependent,
     cubic_interpolator,

cleo/light/light.py

@@ -44,6 +44,9 @@
 class LightModel(ABC):
     """Defines how light propagates given a source location and direction."""
 
+    wavelength: Quantity = field(default=473 * nmeter, kw_only=True)
+    """Light wavelength with unit (usually nmeter)."""
+
     @abstractmethod
     def transmittance(
         self,
@@ -225,6 +228,7 @@ def fiber473nm(
     K=0.125 / mm,  # absorbance coefficient
     S=7.37 / mm,  # scattering coefficient
     ntis=1.36,  # tissue index of refraction
+    wavelength=473 * nmeter, 
 ) -> OpticFiber:
     """Returns an :class:`OpticFiber` model with parameters for 473 nm light.
 
@@ -235,6 +239,7 @@ def fiber473nm(
         K=K,
         S=S,
         ntis=ntis,
+        wavelength=wavelength,
     )
 
 
@@ -313,8 +318,6 @@ class Light(Stimulator):
     Can also be an nx3 array for multiple sources.
     """
 
-    wavelength: Quantity = field(default=473 * nmeter, kw_only=True)
-    """light wavelength with unit (usually nmeter)"""
 
     @coords.validator
     def _check_coords(self, attribute, value):
@@ -369,6 +372,15 @@ def connect_to_neuron_group(
             )
         registry.register_light(self, neuron_group)
 
+    @property
+    def wavelength(self) -> Quantity:
+        """Light wavelength, delegated to the LightModel"""
+        return self.light_model.wavelength
+
+    @wavelength.setter
+    def wavelength(self, value: Quantity):
+        self.light_model.wavelength = value
+
     @property
     def n(self):
         """Number of light sources"""
@@ -402,6 +414,9 @@ def power_(self) -> Quantity:
         """Returns history of light power without units (/mwatt)."""
         return self.power / mwatt
 
+
+
+
     def add_self_to_plot(self, ax, axis_scale_unit, **kwargs) -> list[PathCollection]:
         # show light with point field, assigning r and z coordinates
         # to all points

cleo/light/optogensim.py

@@ -0,0 +1,113 @@
+import xarray as xr
+import numpy as np
+from functools import cached_property
+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(slots=False)
+class OptogenSIM(LightModel):
+    """Light model from OptogenSIM Monte Carlo simulations."""
+
+    beam_radius: Quantity = 100 * um
+    """Beam (1/e^2) radius. Must be within the simulated grid range."""
+    data: xr.DataArray = field(default=None, repr=False)
+    """The 4D (wavelength, beam_size, r, z) dataset. Defaults to the dataset
+    packaged with Cleo if not provided."""
+
+    def __attrs_post_init__(self):
+        if self.data is None:
+            path = str(files("cleo.light.data") / "light_model_4d.nc.gz")
+            self.data = xr.open_dataarray(path, engine="scipy")
+
+    @cached_property
+    def _rz_slice(self):
+        rz = self.data.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")])
+        return rz.assign_coords(z=rz.z - z_source)
+
+    @cached_property
+    def _r_range(self):
+        return (float(self._rz_slice.r.min()), float(self._rz_slice.r.max()))
+
+    @cached_property
+    def _z_range(self):
+        return (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
+
+
+def defocused_gaussian_beam(
+    wavelength: Quantity = 473 * nmeter,
+    beam_radius: Quantity = 100 * um,
+    data_path: str = None,
+) -> OptogenSIM:
+    """Construct an :class:`OptogenSIM` light model from a dataset file.
+
+    ``wavelength`` and ``beam_radius`` are as documented on
+    :class:`OptogenSIM`. ``data_path`` is an optional path to a 4D dataset;
+    if None, the dataset packaged with Cleo is used.
+    """
+    if data_path is None:
+        data_path = str(files("cleo.light.data") / "light_model_4d.nc.gz")
+    data = xr.open_dataarray(data_path, engine="scipy")
+    return OptogenSIM(wavelength=wavelength, beam_radius=beam_radius, data=data)

cleo/light/two_photon.py

@@ -114,8 +114,7 @@ def tp_light_from_scope(scope, wavelength=1060 * nmeter, **kwargs) -> Light:
     light = Light(
         coords=coords,
         direction=scope.direction,
-        light_model=GaussianEllipsoid(),
-        wavelength=wavelength,
+        light_model=GaussianEllipsoid(wavelength=wavelength),
         **kwargs,
     )
     return light

docs/overview.md

@@ -133,7 +133,6 @@ This requires injection of both a light source and an opsin---see the {doc}`tuto
 fiber = cleo.light.Light(
     coords=(0, 0, 0.5)*b2.mm,
     light_model=cleo.light.fiber473nm(),
-    wavelength=473*b2.nmeter
 )
 chr2 = cleo.opto.chr2_4s()
 sim.inject(fiber, ng).inject(chr2, ng)

notebooks/light_profile_explorer.py

@@ -0,0 +1,118 @@
+import marimo
+
+__generated_with = "0.23.14"
+app = marimo.App(width="medium")
+
+
+@app.cell
+def _():
+    return
+
+
+@app.cell
+def _(mo):
+    mo.md("""
+    # OptogenSIM profile explorer
+
+    Interactive 2D cross-sections of the OptogenSIM Monte Carlo light model.
+    Drag the sliders to set a wavelength and beam radius and see how light
+    propagates through gray matter. The heatmap shows log₁₀ transmittance in
+    the (r, z) plane, mirrored across the beam axis. z is depth from the
+    source; r is radial distance.
+
+    *Note: values between simulated grid points (wavelengths every 30 nm;
+    beam radii 10, 20, 100, 200, 400, 800 µm) are linearly interpolated.
+    """)
+    return
+
+
+@app.cell
+def _():
+    import numpy as np
+    import xarray as xr
+    import plotly.graph_objects as go
+
+    return go, np, xr
+
+
+@app.cell
+def _(xr):
+    from importlib.resources import files
+    _path = str(files("cleo.light.data") / "light_model_4d.nc.gz")
+    da = xr.open_dataarray(_path, engine="scipy")
+    da.load()
+    return (da,)
+
+
+@app.cell
+def _(da, mo):
+    wl_vals = da.wavelength.values
+    bs_vals = da.beam_size.values
+    wavelength = mo.ui.slider(
+        start=float(wl_vals.min()),
+        stop=float(wl_vals.max()),
+        step=5,
+        value=470.0,
+        label="Wavelength (nm)",
+        show_value=True,
+    )
+    beam = mo.ui.slider(
+        start=float(bs_vals.min()),
+        stop=float(bs_vals.max()),
+        step=10,
+        value=200.0,
+        label="Beam radius (µm)",
+        show_value=True,
+    )
+    mo.hstack([wavelength, beam], justify="start")
+    return beam, wavelength
+
+
+@app.cell
+def _(beam, da, go, np, wavelength):
+    sl = da.interp(wavelength=wavelength.value, beam_size=beam.value)
+
+    r = sl.r.values * 10.0   # cm -> mm
+    z = sl.z.values * 10.0   # cm -> mm
+
+    on_axis = sl.isel(r=0).values
+    z0 = z[int(np.argmax(on_axis))]
+    z_rel = z - z0
+
+    T = sl.values
+    T = T / np.nanmax(T)
+    logT = np.log10(np.clip(T, 1e-4, None))
+
+    r_full = np.concatenate([-r[::-1], r])
+    logT_full = np.vstack([logT[::-1, :], logT])
+
+    fig = go.Figure(
+        data=go.Heatmap(
+            x=z_rel,
+            y=r_full,
+            z=logT_full,
+            colorscale="Viridis",
+            zmin=-4, zmax=0,
+            colorbar=dict(title="log₁₀ T/peak"),
+        )
+    )
+    fig.update_layout(
+        title=f"{wavelength.value:.0f} nm, {beam.value:.0f} µm beam radius",
+        xaxis_title="z, depth from source (mm)",
+        yaxis_title="r, radial distance (mm)",
+        width=750, height=450,
+    )
+    fig.update_xaxes(range=[float(z_rel.min()), min(float(z_rel.max()), 3)])
+    fig
+    return
+
+
+@app.cell
+def _():
+    import marimo as mo
+
+    return (mo,)
+
+
+if __name__ == "__main__":
+    app.run()

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",