[forward] add pmcxcl and user-defined input

environment_dev.yml

@@ -54,6 +54,6 @@ dependencies:
       - pywavefront==1.3
       - setuptools-scm
       - snirf==0.8
-      - pmcx==0.3.3
+      - pmcx==0.3.4
+      - pmcxcl==0.2.0
       - sphinxcontrib-bibtex
-

environment_doc.yml

@@ -55,4 +55,5 @@ dependencies:
       - setuptools-scm==7.1.0
       - snirf==0.7.4
       - pmcx
+      - pmcxcl
       - open3d==0.16

src/cedalion/dataclasses/__init__.py

@@ -9,6 +9,7 @@
     TrimeshSurface,
     VTKSurface,
     affine_transform_from_numpy,
+    Voxels,
 )
 from .schemas import (
     build_labeled_points,

src/cedalion/dataclasses/geometry.py

@@ -96,6 +96,63 @@ def snap(self, points: cdt.LabeledPointCloud):
         return snapped
 
 
+@dataclass
+class Voxels():
+    """3D voxels represented by a np.array.
+
+    Attributes:
+        voxels (np.ndarray): The voxels.
+        crs (str): The coordinate reference system of the voxels.
+        units (pint.Unit): The units of the voxels.
+    """
+    voxels: np.ndarray
+    crs: str
+    units: pint.Unit
+
+    @property
+    def vertices(self) -> cdt.LabeledPointCloud:
+        result = xr.DataArray(
+            self.voxels,
+            dims=["label", self.crs],
+            coords={"label": np.arange(len(self.voxels))},
+            attrs={"units": self.units},
+        )
+        result = result.pint.quantify()
+
+        return result
+
+    @property
+    def nvertices(self) -> int:
+        return len(self.voxels)
+
+    def apply_transform(self, transform: cdt.AffineTransform) -> "Voxels":
+        # convert to homogeneous coordinates
+        num, dim = self.voxels.shape
+        hom = np.ones((num,dim+1))
+        hom[:,:3] = self.voxels
+        # apply transformation
+        hom = (transform.pint.dequantify().values.dot(hom.T)).T
+        # backtransformation
+        transformed = np.array([hom[i,:3] / hom[i,3] for i in range(hom.shape[0])])
+
+        new_units = self.units * transform.pint.units
+        new_crs = transform.dims[0]
+
+        return Voxels(transformed, new_crs, new_units)
+
+    def _build_kdtree(self):
+        self._kdtree = KDTree(self.voxels)
+
+    def __post_init__(self):
+        self._kdtree = None
+
+    @property
+    def kdtree(self):
+        if self._kdtree is None:
+            self._build_kdtree()
+        return self._kdtree
+
+
 @dataclass
 class TrimeshSurface(Surface):
     """A surface represented by a trimesh object.

src/cedalion/geometry/segmentation.py

@@ -12,6 +12,42 @@
 import cedalion.dataclasses as cdc
 
 
+def voxels_from_segmentation(
+    segmentation_mask: xr.DataArray,
+    segmentation_types: List[str],
+    isovalue=0.9,
+    fill_holes_in_mask=False,
+) -> cdc.Voxels:
+    """Generate voxels from a segmentation mask.
+
+    Args:
+        segmentation_mask : xr.DataArray
+            Segmentation mask.
+        segmentation_types : List[str]
+            List of segmentation types.
+        isovalue : float, optional
+            Isovalue for marching cubes, by default 0.9.
+        fill_holes_in_mask : bool, optional
+            Fill holes in the mask, by default False.
+
+    Returns:
+        cdc.Voxels
+            Voxels in voxel space.
+    """
+    combined_mask = (
+        segmentation_mask.sel(segmentation_type=segmentation_types)
+        .any("segmentation_type")
+        .values
+    )
+
+    if fill_holes_in_mask:
+        combined_mask = binary_fill_holes(combined_mask).astype(combined_mask.dtype)
+
+    voxels = np.argwhere(combined_mask)
+
+    return cdc.Voxels(voxels, "ijk", cedalion.units.Unit("1"))
+
+
 def surface_from_segmentation(
     segmentation_mask: xr.DataArray,
     segmentation_types: List[str],
@@ -32,6 +68,7 @@ def surface_from_segmentation(
     Returns:
         A cedalion.Surface object.
     """
+
     combined_mask = (
         segmentation_mask.sel(segmentation_type=segmentation_types)
         .any("segmentation_type")
@@ -59,6 +96,7 @@ def cell_coordinates(volume, flat: bool = False):
     Returns:
         xr.DataArray: A DataArray with the coordinates of the cells in the volume.
     """
+
     # coordinates in voxel space
     i = np.arange(volume.shape[0])
     j = np.arange(volume.shape[1])

src/cedalion/imagereco/forward_model.py

@@ -647,7 +647,7 @@ def _get_unitinmm(self):
         length = xrutils.norm(pts_ras[1] - pts_ras[0], pts_ras.points.crs)
         return length.pint.magnitude.item()
 
-    def _get_fluence_from_mcx(self, i_optode: int, nphoton: int):
+    def _get_fluence_from_mcx(self, i_optode: int, **kwargs):
         """Run MCX simulation to get fluence for one optode.
 
         Args:
@@ -658,8 +658,10 @@ def _get_fluence_from_mcx(self, i_optode: int, nphoton: int):
             np.ndarray: Fluence in each voxel.
         """
 
+        kwargs.setdefault("nphoton", 1e8)
+
         cfg = {
-            "nphoton": nphoton,
+            "nphoton": kwargs['nphoton'],
             "vol": self.volume,
             "tstart": 0,
             "tend": 5e-9,
@@ -670,16 +672,24 @@ def _get_fluence_from_mcx(self, i_optode: int, nphoton: int):
             "issrcfrom0": 1,
             "isnormalized": 1,
             "outputtype": "fluence",
-            "seed": int(np.floor(np.random.rand() * 1e7)),
-            "issavedet": 1,
+            "issavedet": 0,
             "unitinmm": self.unitinmm,
         }
 
-        import pmcx
-        result = pmcx.run(cfg)
+        # merging default cfg with additional positional arguments
+
+        cfg = { **cfg, **kwargs }
+
+        # if pmcx fails, try pmcxcl
+
+        if "cuda" in cfg and cfg["cuda"]:
+            import pmcx
+            result = pmcx.run(cfg)
+        else:
+            import pmcxcl
+            result = pmcxcl.run(cfg)
 
         fluence = result["flux"][:, :, :, 0]  # there is only one time bin
-        fluence = fluence * cfg["tstep"] / result["stat"]["normalizer"]
 
         return fluence
 
@@ -722,11 +732,13 @@ def _fluence_at_optodes(self, fluence, emitting_opt):
 
         return result
 
-    def compute_fluence_mcx(self, nphoton: int = 1e8):
+    def compute_fluence_mcx(self, **kwargs):
         """Compute fluence for each channel and wavelength using MCX package.
 
         Args:
             nphoton (int): Number of photons to simulate.
+            along with other pmcx/pmcxcl accepted input fields,
+            see https://pypi.org/project/pmcx/
 
         Returns:
             xr.DataArray: Fluence in each voxel for each channel and wavelength.
@@ -763,9 +775,9 @@ def compute_fluence_mcx(self, nphoton: int = 1e8):
             label = self.optode_pos.label.values[i_opt]
             print(f"simulating fluence for {label}. {i_opt+1} / {n_optodes}")
 
-            # run MCX
+            # run MCX or MCXCL
             # shape: [i,j,k]
-            fluence = self._get_fluence_from_mcx(i_opt, nphoton=nphoton)
+            fluence = self._get_fluence_from_mcx(i_opt, **kwargs)
 
             # FIXME shortcut: currently tissue props are wavelength independent -> copy
             for i_wl in range(n_wavelength):

src/cedalion/io/forward_model.py

@@ -44,7 +44,7 @@ def load_Adot(fn: str):
 def save_fluence(fn : str, fluence_all, fluence_at_optodes):
     """Save forward model computation results.
 
-    This method uses a lossy compressions algorithm to reduce file size.
+    This method uses a lossy compression algorithm to reduce file size.
     """
 
     with h5py.File(fn, "w") as f: