lsurf 1.0.0__py3-none-any.whl

lsurf/geometry/cell_geometry.py

@@ -0,0 +1,247 @@
+# The Clear BSD License
+#
+# Copyright (c) 2026 Tobias Heibges
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted (subject to the limitations in the disclaimer
+# below) provided that the following conditions are met:
+#
+#      * Redistributions of source code must retain the above copyright notice,
+#      this list of conditions and the following disclaimer.
+#
+#      * Redistributions in binary form must reproduce the above copyright
+#      notice, this list of conditions and the following disclaimer in the
+#      documentation and/or other materials provided with the distribution.
+#
+#      * Neither the name of the copyright holder nor the names of its
+#      contributors may be used to endorse or promote products derived from this
+#      software without specific prior written permission.
+#
+# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
+# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
+# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
+# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
+# IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+
+"""
+Cell-Based Geometry Container
+
+Extends Geometry with explicit cell-based material assignment for
+complex geometries where surfaces intersect.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from .cell import Cell
+from .geometry import Geometry
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+    from ..materials import MaterialField
+
+
+@dataclass(frozen=True)
+class CellGeometry(Geometry):
+    """
+    Geometry with explicit cell-based material assignment.
+
+    Extends the standard Geometry with cells that define material regions
+    by the intersection of half-spaces. Use this when surfaces intersect
+    and the simple front/back model cannot describe all regions.
+
+    Parameters
+    ----------
+    surfaces : tuple of Surface
+        All optical/absorber surfaces in the geometry.
+    detectors : tuple of Surface
+        All detector surfaces in the geometry.
+    background_material : MaterialField
+        The background/ambient material.
+    media : dict
+        Mapping from medium name to MaterialField.
+    surface_names : dict
+        Mapping from surface name to index in surfaces tuple.
+    detector_names : dict
+        Mapping from detector name to index in detectors tuple.
+    cells : tuple of Cell
+        Cells defining material regions by half-space intersection.
+
+    Examples
+    --------
+    >>> # Created by GeometryBuilder in cell mode
+    >>> builder = GeometryBuilder()
+    >>> builder.register_medium("air", AIR)
+    >>> builder.register_medium("water", WATER)
+    >>> builder.set_background("air")
+    >>> builder.add_surface_only(plane_x)
+    >>> builder.add_surface_only(plane_y)
+    >>> builder.add_cell("air", ("plane_x", True), ("plane_y", True))
+    >>> builder.add_cell("water", ("plane_x", False), ("plane_y", True))
+    >>> geometry = builder.build()  # Returns CellGeometry
+    >>>
+    >>> # Query material at positions
+    >>> positions = np.array([[1, 1, 0], [-1, 1, 0]])
+    >>> materials = geometry.get_material_at(positions)
+    """
+
+    cells: tuple[Cell, ...]
+
+    def get_material_at(
+        self,
+        positions: NDArray[np.float64],
+    ) -> list[tuple[MaterialField, NDArray[np.bool_]]]:
+        """
+        Determine material for each position based on cell membership.
+
+        Parameters
+        ----------
+        positions : NDArray, shape (N, 3)
+            Positions to query.
+
+        Returns
+        -------
+        list of (MaterialField, NDArray[bool])
+            List of (material, mask) pairs where mask indicates which
+            positions have that material. A position may have no material
+            if it doesn't fall into any defined cell.
+
+        Notes
+        -----
+        Cells are checked in order. The first cell that contains a position
+        determines its material. Positions not in any cell get the background
+        material.
+        """
+        positions = np.asarray(positions, dtype=np.float64)
+        n_positions = len(positions)
+
+        # Compute signed distances for all surfaces
+        signed_distances: dict[str, NDArray[np.float64]] = {}
+        for surface in self.surfaces:
+            if surface.gpu_capable:
+                sd = surface.signed_distance(positions.astype(np.float32))
+                signed_distances[surface.name] = sd.astype(np.float64)
+
+        # Track which positions are assigned
+        assigned = np.zeros(n_positions, dtype=np.bool_)
+        results: list[tuple[MaterialField, NDArray[np.bool_]]] = []
+
+        # Check each cell
+        for cell in self.cells:
+            # Get material for this cell
+            material = self.media.get(cell.medium_name)
+            if material is None:
+                continue
+
+            # Find positions in this cell that aren't already assigned
+            in_cell = cell.contains(signed_distances)
+            new_in_cell = in_cell & ~assigned
+
+            if np.any(new_in_cell):
+                results.append((material, new_in_cell))
+                assigned |= new_in_cell
+
+        # Assign background material to remaining positions
+        remaining = ~assigned
+        if np.any(remaining):
+            results.append((self.background_material, remaining))
+
+        return results
+
+    def get_material_index_at(
+        self,
+        positions: NDArray[np.float64],
+    ) -> NDArray[np.int32]:
+        """
+        Get cell index for each position (-1 for background).
+
+        Parameters
+        ----------
+        positions : NDArray, shape (N, 3)
+            Positions to query.
+
+        Returns
+        -------
+        NDArray[np.int32], shape (N,)
+            Cell index for each position, or -1 for positions in
+            the background (not in any cell).
+        """
+        positions = np.asarray(positions, dtype=np.float64)
+        n_positions = len(positions)
+
+        # Compute signed distances for all surfaces
+        signed_distances: dict[str, NDArray[np.float64]] = {}
+        for surface in self.surfaces:
+            if surface.gpu_capable:
+                sd = surface.signed_distance(positions.astype(np.float32))
+                signed_distances[surface.name] = sd.astype(np.float64)
+
+        # Initialize with -1 (background)
+        cell_indices = np.full(n_positions, -1, dtype=np.int32)
+
+        # Check each cell in order
+        for i, cell in enumerate(self.cells):
+            in_cell = cell.contains(signed_distances)
+            # Only assign if not already assigned
+            unassigned_in_cell = in_cell & (cell_indices == -1)
+            cell_indices[unassigned_in_cell] = i
+
+        return cell_indices
+
+    def get_cell(self, index: int) -> Cell:
+        """
+        Get cell by index.
+
+        Parameters
+        ----------
+        index : int
+            Cell index.
+
+        Returns
+        -------
+        Cell
+            The cell at the given index.
+
+        Raises
+        ------
+        IndexError
+            If index is out of range.
+        """
+        return self.cells[index]
+
+    def get_cell_by_name(self, name: str) -> Cell | None:
+        """
+        Get cell by name.
+
+        Parameters
+        ----------
+        name : str
+            Cell name.
+
+        Returns
+        -------
+        Cell or None
+            The cell with the given name, or None if not found.
+        """
+        for cell in self.cells:
+            if cell.name == name:
+                return cell
+        return None
+
+    @property
+    def num_cells(self) -> int:
+        """Number of cells in the geometry."""
+        return len(self.cells)