|
| 1 | +from functools import cached_property |
| 2 | +from importlib.resources import files |
| 3 | + |
| 4 | +import numpy as np |
| 5 | +import xarray as xr |
| 6 | +import warnings |
| 7 | +from attrs import define, field |
| 8 | +from brian2.units import Quantity, cm, nmeter, um |
| 9 | + |
| 10 | +from cleo.coords import coords_from_xyz |
| 11 | +from cleo.light.light import LightModel |
| 12 | +from cleo.utilities import uniform_cylinder_rθz, xyz_from_rθz |
| 13 | + |
| 14 | + |
| 15 | +@define |
| 16 | +class OptogenSIM(LightModel): |
| 17 | + """Light model from OptogenSIM Monte Carlo simulations.""" |
| 18 | + |
| 19 | + beam_radius: Quantity = 100 * um |
| 20 | + """Beam (1/e^2) radius. Must be within the simulated grid range.""" |
| 21 | + data: xr.DataArray = field(default=None, repr=False) |
| 22 | + """The 4D transmittance dataset. Defaults to the dataset packaged with Cleo. |
| 23 | +
|
| 24 | + To supply your own, pass an :class:`xarray.DataArray` with dimensions |
| 25 | + ``(wavelength, beam_size, r, z)`` in that order, with coordinate values in |
| 26 | + these units: ``wavelength`` in nm, ``beam_size`` (the beam radius) in µm, |
| 27 | + and ``r`` and ``z`` in cm, where ``z`` is measured from the top of the |
| 28 | + simulated tissue volume. Values are transmittance in [0, 1]. The packaged |
| 29 | + dataset spans wavelengths 380-740 nm (30 nm steps) and beam radii of 10, 20, |
| 30 | + 100, 200, 400, and 800 µm; it was generated with a modified version of the |
| 31 | + OptogenSIM Monte Carlo model (https://github.com/akhilsnarra-ui/optogensim-cleo) sweeping over |
| 32 | + wavelength and beam radius.""" |
| 33 | + |
| 34 | + def __attrs_post_init__(self): |
| 35 | + if self.data is None: |
| 36 | + path = str(files("cleo.light.data") / "light_model_4d.nc.gz") |
| 37 | + self.data = xr.open_dataarray(path, engine="scipy") |
| 38 | + |
| 39 | + @cached_property |
| 40 | + def _rz_slice(self): |
| 41 | + rz = self.data.interp( |
| 42 | + wavelength=self.wavelength / nmeter, beam_size=self.beam_radius / um |
| 43 | + ) |
| 44 | + # data measures z from the atlas top; re-zero to the source |
| 45 | + z_source = float(rz.isel(r=0).z[rz.isel(r=0).argmax("z")]) |
| 46 | + return rz.assign_coords(z=rz.z - z_source) |
| 47 | + |
| 48 | + @cached_property |
| 49 | + def _r_range(self): |
| 50 | + return (float(self._rz_slice.r.min()), float(self._rz_slice.r.max())) |
| 51 | + |
| 52 | + @cached_property |
| 53 | + def _z_range(self): |
| 54 | + return (float(self._rz_slice.z.min()), float(self._rz_slice.z.max())) |
| 55 | + |
| 56 | + def transmittance(self, source_coords, source_dir_uvec, target_coords): |
| 57 | + assert np.allclose(np.linalg.norm(source_dir_uvec, axis=-1), 1) |
| 58 | + r, z = self._get_rz_for_xyz(source_coords, source_dir_uvec, target_coords) |
| 59 | + r_cm = np.asarray(r / cm) |
| 60 | + z_cm = np.asarray(z / cm) |
| 61 | + |
| 62 | + # Warn if targets fall outside the simulated data range. Targets with |
| 63 | + # r beyond the far edge or z outside the data range get transmittance |
| 64 | + # 0. Targets nearer the axis than the first radial bin (r < r_min) |
| 65 | + # are clipped to the edge, i.e. take the near-axis peak value. |
| 66 | + # Data coverage may be extended in a future re-run. |
| 67 | + outside = ( |
| 68 | + (r_cm > self._r_range[1]) |
| 69 | + | (z_cm < self._z_range[0]) |
| 70 | + | (z_cm > self._z_range[1]) |
| 71 | + ) |
| 72 | + if np.any(outside): |
| 73 | + warnings.warn( |
| 74 | + "Some target coordinates fall outside the OptogenSIM data range " |
| 75 | + f"(r: {self._r_range} cm, z: {self._z_range} cm). Transmittance " |
| 76 | + "is set to 0 outside the data range." |
| 77 | + ) |
| 78 | + |
| 79 | + T = self._rz_slice.interp( |
| 80 | + r=xr.DataArray(np.clip(r_cm, *self._r_range)), |
| 81 | + z=xr.DataArray(np.clip(z_cm, *self._z_range)), |
| 82 | + ).values |
| 83 | + T = np.nan_to_num(T, nan=0.0) |
| 84 | + T[outside] = 0 # zero outside data range; keep real backscatter within it |
| 85 | + return T |
| 86 | + |
| 87 | + @property |
| 88 | + def area0(self): |
| 89 | + return np.pi * self.beam_radius**2 |
| 90 | + |
| 91 | + def viz_params( |
| 92 | + self, coords, direction, T_threshold, n_points_per_source=1.5e5, **kwargs |
| 93 | + ): |
| 94 | + r_thresh, zc_thresh, zc_back = self._find_rz_thresholds(T_threshold) |
| 95 | + # cylinder spans from behind the source (zc_back < 0) through zc_thresh |
| 96 | + total_length = zc_thresh - zc_back |
| 97 | + r, theta, zc = uniform_cylinder_rθz(n_points_per_source, r_thresh, total_length) |
| 98 | + start = coords + zc_back * direction |
| 99 | + end = coords + zc_thresh * direction |
| 100 | + x, y, z = xyz_from_rθz(r, theta, zc, start, end) |
| 101 | + density_factor = 3 |
| 102 | + cyl_vol = np.pi * r_thresh**2 * total_length |
| 103 | + markersize = (cyl_vol / n_points_per_source * density_factor) ** (1 / 3) |
| 104 | + intensity_scale = 1.5 * (4e3 / n_points_per_source) ** (1 / 3) |
| 105 | + return coords_from_xyz(x, y, z), markersize, intensity_scale |
| 106 | + |
| 107 | + def _find_rz_thresholds(self, thresh): |
| 108 | + """Find r and z extents for visualization at the given transmittance threshold.""" |
| 109 | + # On-axis (r[0]) profile along z (re-zeroed so source at z=0). |
| 110 | + on_axis = self._rz_slice.isel(r=0) |
| 111 | + z_vals = on_axis.z.values # cm, source at 0 |
| 112 | + T_z = on_axis.values |
| 113 | + |
| 114 | + # Forward extent (z >= 0): first z where T drops below threshold. |
| 115 | + fwd = z_vals >= 0 |
| 116 | + z_pos, T_pos = z_vals[fwd], T_z[fwd] |
| 117 | + below = np.where(T_pos < thresh)[0] |
| 118 | + zc_thresh = (z_pos[below[0]] if len(below) > 0 else z_pos[-1]) * cm |
| 119 | + |
| 120 | + # Backward extent (z < 0): most-negative z still above threshold, |
| 121 | + # so backscatter behind the source is included in the visualization. |
| 122 | + back = z_vals < 0 |
| 123 | + z_neg, T_neg = z_vals[back], T_z[back] |
| 124 | + above_back = z_neg[T_neg >= thresh] |
| 125 | + zc_back = (float(above_back.min()) if len(above_back) > 0 else 0.0) * cm |
| 126 | + |
| 127 | + # r threshold at the forward midpoint. |
| 128 | + z_mid = float(zc_thresh / cm) / 2 |
| 129 | + z_mid = np.clip( |
| 130 | + z_mid, float(self._rz_slice.z.min()), float(self._rz_slice.z.max()) |
| 131 | + ) |
| 132 | + r_vals = self._rz_slice.r.values |
| 133 | + T_r = self._rz_slice.interp(z=float(z_mid)).values |
| 134 | + below_r = np.where(T_r < thresh)[0] |
| 135 | + r_thresh = (r_vals[below_r[0]] if len(below_r) > 0 else r_vals[-1]) * cm |
| 136 | + |
| 137 | + return r_thresh * 1.2, zc_thresh, zc_back |
| 138 | + |
| 139 | + |
| 140 | +def defocused_gaussian_beam( |
| 141 | + wavelength: Quantity = 473 * nmeter, |
| 142 | + beam_radius: Quantity = 100 * um, |
| 143 | +) -> OptogenSIM: |
| 144 | + """Construct an :class:`OptogenSIM` model for a defocused Gaussian beam. |
| 145 | +
|
| 146 | + This is the default OptogenSIM profile: light emitted from a defocused |
| 147 | + Gaussian beam source, propagating through gray matter, using the Monte |
| 148 | + Carlo dataset packaged with Cleo. ``wavelength`` and ``beam_radius`` are |
| 149 | + as documented on :class:`OptogenSIM`. To use a custom dataset, construct |
| 150 | + :class:`OptogenSIM` directly and pass your own ``data`` array. |
| 151 | + """ |
| 152 | + return OptogenSIM(wavelength=wavelength, beam_radius=beam_radius) |
0 commit comments