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,135 @@
+from functools import cached_property
+from importlib.resources import files
+
+import numpy as np
+import xarray as xr
+from attrs import define, field
+from brian2.units import Quantity, cm, nmeter, um
+
+from cleo.coords import coords_from_xyz
+from cleo.light.light import LightModel
+from cleo.utilities import uniform_cylinder_rθz, xyz_from_rθz
+
+
+@define
+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 transmittance dataset. Defaults to the dataset packaged with Cleo.
+
+    To supply your own, pass an :class:`xarray.DataArray` with dimensions
+    ``(wavelength, beam_size, r, z)`` in that order, with coordinate values in
+    these units: ``wavelength`` in nm, ``beam_size`` (the beam radius) in µm,
+    and ``r`` and ``z`` in cm, where ``z`` is measured from the top of the
+    simulated tissue volume. Values are transmittance in [0, 1]. The packaged
+    dataset spans wavelengths 380-740 nm (30 nm steps) and beam radii of 10, 20,
+    100, 200, 400, and 800 µm; it was generated with a modified version of the
+    OptogenSIM Monte Carlo model (https://github.com/akhilsnarra-ui/optogensim-cleo) sweeping over
+    wavelength and beam radius."""
+
+    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=1.5e5, **kwargs
+    ):
+        r_thresh, zc_thresh, zc_back = self._find_rz_thresholds(T_threshold)
+        # cylinder spans from behind the source (zc_back < 0) through zc_thresh
+        total_length = zc_thresh - zc_back
+        r, theta, zc = uniform_cylinder_rθz(n_points_per_source, r_thresh, total_length)
+        start = coords + zc_back * direction
+        end = coords + zc_thresh * direction
+        x, y, z = xyz_from_rθz(r, theta, zc, start, end)
+        density_factor = 3
+        cyl_vol = np.pi * r_thresh**2 * total_length
+        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
+
+        # Forward extent (z >= 0): first z where T drops below threshold.
+        fwd = z_vals >= 0
+        z_pos, T_pos = z_vals[fwd], T_z[fwd]
+        below = np.where(T_pos < thresh)[0]
+        zc_thresh = (z_pos[below[0]] if len(below) > 0 else z_pos[-1]) * cm
+
+        # Backward extent (z < 0): most-negative z still above threshold,
+        # so backscatter behind the source is included in the visualization.
+        back = z_vals < 0
+        z_neg, T_neg = z_vals[back], T_z[back]
+        above_back = z_neg[T_neg >= thresh]
+        zc_back = (float(above_back.min()) if len(above_back) > 0 else 0.0) * cm
+
+        # r threshold at the forward midpoint.
+        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, zc_back
+
+
+def defocused_gaussian_beam(
+    wavelength: Quantity = 473 * nmeter,
+    beam_radius: Quantity = 100 * um,
+) -> OptogenSIM:
+    """Construct an :class:`OptogenSIM` model for a defocused Gaussian beam.
+
+    This is the default OptogenSIM profile: light emitted from a defocused
+    Gaussian beam source, propagating through gray matter, using the Monte
+    Carlo dataset packaged with Cleo. ``wavelength`` and ``beam_radius`` are
+    as documented on :class:`OptogenSIM`. To use a custom dataset, construct
+    :class:`OptogenSIM` directly and pass your own ``data`` array.
+    """
+    return OptogenSIM(wavelength=wavelength, beam_radius=beam_radius)

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

cleo/registry.py

@@ -122,6 +122,8 @@ def _add_brian_object(self, obj):
     def _remove_brian_object(self, obj):
         self.brian_objects.remove(obj)
         self.sim.network.remove(obj)
+        # brian2's BrianObject.__del__ warns if _network is None when the object is garbage-collected ("never included in a network"). This object was in the network and is being deliberately discarded (replaced by a combined light source), so set _network to the network's id to suppress the spurious warning.
+        obj._network = self.sim.network.id
 
     def _get_or_create_light_prop_syn(
         self, ldd: "LightDependent", ng: NeuronGroup

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,126 @@
+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 plotly.graph_objects as go
+    from cleo.light import OptogenSIM
+    from brian2 import mm, nmeter, um
+
+    return OptogenSIM, go, mm, nmeter, np, um
+
+
+@app.cell
+def _(OptogenSIM):
+    _m = OptogenSIM()  # default instance, just to read the packaged data grid
+    wl_vals = _m.data.wavelength.values
+    bs_vals = _m.data.beam_size.values
+    return bs_vals, wl_vals
+
+
+@app.cell
+def _(bs_vals, mo, wl_vals):
+    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 _(OptogenSIM, beam, go, mm, nmeter, np, um, wavelength):
+    model = OptogenSIM(
+        wavelength=wavelength.value * nmeter,
+        beam_radius=beam.value * um,
+    )
+
+    # build an (r, z) grid of target points, as a user would: real coordinates
+    # with brian2 units. Source at origin pointing in +z.
+    r_mm = np.linspace(0, 3, 120)
+    z_mm = np.linspace(-1, 6, 200)
+    R, Z = np.meshgrid(r_mm, z_mm, indexing="ij")
+    grid_coords = np.stack([R.ravel(), np.zeros(R.size), Z.ravel()], axis=-1) * mm
+
+    source = np.array([0, 0, 0]) * mm
+    direction = np.array([0, 0, 1.0])
+
+    T = model.transmittance(source, direction, grid_coords).reshape(R.shape)
+
+    # normalize to peak for display (colorbar is log10 T/peak)
+
+    logT = np.log10(np.clip(T, 1e-4, None))
+
+    # mirror across the beam axis for a symmetric view
+    r_full = np.concatenate([-r_mm[::-1], r_mm])
+    logT_full = np.vstack([logT[::-1, :], logT])
+
+    fig = go.Figure(
+        data=go.Heatmap(
+            x=z_mm,
+            y=r_full,
+            z=logT_full,
+            colorscale="Viridis",
+            zmin=-4, zmax=0,
+            colorbar=dict(title="log₁₀ T"),
+        )
+    )
+    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,
+        yaxis=dict(scaleanchor="x", scaleratio=1),
+    )
+    fig.update_xaxes(range=[float(z_mm.min()), min(float(z_mm.max()), 3)])
+    fig
+    return
+
+
+@app.cell
+def _():
+    import marimo as mo
+
+    return (mo,)
+
+
+if __name__ == "__main__":
+    app.run()

pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
     "scipy>=1.7.2,<1.14.1; python_version<'3.13'",
     "scipy>=1.14.1; python_version>='3.13'",
     "numpy >= 2.1.0,<3.0.0; python_version>='3.13'",
+    "xarray>=2025.1.0",
     "tklfp>=0.3,<0.4",
     "attrs>=23.0,<24.0",
     "neo>=0.12.0,<0.13.0",
@@ -40,11 +41,13 @@ dev = [
     "ruff>=0.1.15,<1.0.0",
     "pytest-xdist>=3.5.0,<4.0.0",
     "marimo>=0.11.20,<1.0.0",
+    "plotly>=5.0.0,<7.0.0",
     "cvxpy>=1.6.5,<2.0.0",
 ]
 
 [tool.poetry]
 packages = [{ include = "cleo" }]
+include = [{ path = "cleo/light/data/*.nc.gz", format = ["sdist", "wheel"] }]
 classifiers = [
     "Development Status :: 4 - Beta",
     "Intended Audience :: Science/Research",