simgeometry 0.1.0__py3-none-any.whl

simgeometry/grid.py

@@ -0,0 +1,238 @@
+"""Grid helpers for geometry, grids, and simulation domains."""
+
+from __future__ import annotations
+
+import cunumpy as xp
+import numpy as np
+from maybempi import MPI
+from numpy import ndarray
+
+from simgeometry._normalize import ensure_dim_tuple
+from simgeometry._types import Array, ConfigValue, Printable
+from simgeometry.boundingbox import BoundingBox
+
+
+class Grid(BoundingBox):
+    """Represent grid coordinates for a bounded simulation domain."""
+
+    def __init__(
+        self,
+        ndim: int,
+        lower_bound: tuple,
+        upper_bound: tuple,
+        num_gridpoints: tuple,
+        num_ghost_points: int = 0,
+        comm: MPI.Comm | None = None,
+        domain_decomposition: ConfigValue = None,
+        periodic: bool | None = None,
+    ) -> None:
+        """
+        Initialize a grid.
+
+        Parameters:
+        ndim (int): Number of dimensions (e.g., 1, 2, or 3).
+        bounds (list of tuples): List of tuples specifying min and max for each dimension.
+        num_gridpoints (list of int or int in case of ndim==1): Number of grid points in each dimension.
+                                      For example, [num_points_x, num_points_y, num_points_z] for 3D.
+        """
+        super().__init__(
+            lower_bound=lower_bound,
+            upper_bound=upper_bound,
+            comm=comm,
+            domain_decomposition=domain_decomposition,
+            periodic=periodic,
+        )
+
+        num_gridpoints = ensure_dim_tuple(
+            num_gridpoints,
+            ndim=ndim,
+        )
+
+        self.num_gridpoints = num_gridpoints  # xp.array(num_gridpoints, dtype=int)
+        self._num_ghost_points = num_ghost_points
+        self.num_gridpoints_ndarray = np.array(num_gridpoints, dtype=int)
+        # On a bounded axis the first and last node sit on the boundary; on a
+        # periodic axis the last node is a neighbour of the first, not a copy
+        # of it, so there are ``n`` cells over the length rather than ``n-1``.
+        # This is the same convention as ``FieldConfig.node_spacing`` and
+        # the field solvers, so markers and the solver agree on where the
+        # nodes are.
+        periodic_axes = self.periodic or (False,) * ndim
+        self.grid_spacing = np.array(
+            [
+                (self.upper_bound[dir] - self.lower_bound[dir])
+                / (
+                    num_gridpoints[dir]
+                    if periodic_axes[dir]
+                    else num_gridpoints[dir] - 1
+                )
+                for dir in range(ndim)
+            ],
+            dtype=float,
+        )
+
+        self._proc_index_bounds = [
+            (int(start), int(end))
+            for start, end in self.domain_decomposition_class.get_index_bounds(
+                self.num_gridpoints
+            )
+        ]
+        self._num_gridpoints_local = [
+            end - start for start, end in self._proc_index_bounds
+        ]
+        # Same neighbours as the decomposition: periodic wrap, PROC_NULL at walls.
+        self._neighbor_ranks = list(self.domain_decomposition_class.neighbour_ranks)
+        self._num_gridpoints_local = tuple(self._num_gridpoints_local)
+
+        # Add attribute: number of local gridpoints including ghost points
+        self._num_gridpoints_local_with_ghosts = tuple(
+            n + 2 * self._num_ghost_points for n in self._num_gridpoints_local
+        )
+
+    def show(self, **kwargs):
+        """Display grid lines and optional nodes, or an axis-aligned slice."""
+        from simgeometry.collection import GeometryCollection
+
+        geometry = GeometryCollection({}, bounds=(self.lower_bound, self.upper_bound))
+        return geometry.show(grid=self, labels=False, **kwargs)
+
+    def show_3d(self, **kwargs):
+        """Display a three-dimensional grid and optional nodes."""
+        from simgeometry.collection import GeometryCollection
+
+        geometry = GeometryCollection({}, bounds=(self.lower_bound, self.upper_bound))
+        return geometry.show_3d(grid=self, labels=False, **kwargs)
+
+    @property
+    def indices(self) -> list[tuple[int, int]]:
+        """Return the indices."""
+        indices = []
+        if self.ndim == 1:
+            return list(range(self.num_gridpoints[0]))
+        elif self.ndim == 2:
+            for i in range(self.num_gridpoints[0]):
+                for j in range(self.num_gridpoints[1]):
+                    indices.append((i, j))
+            return indices
+        elif self.ndim == 3:
+            for i in range(self.num_gridpoints[0]):
+                for j in range(self.num_gridpoints[1]):
+                    for k in range(self.num_gridpoints[2]):
+                        indices.append((i, j, k))
+            return indices
+
+    def node_points(self) -> np.ndarray:
+        """Return the coordinates of every node as an ``(ndim, N)`` host array.
+
+        The nodes are in the order of :attr:`indices` (C order of the grid);
+        column ``n`` equals :meth:`node_coordinates` of the ``n``-th index.
+        """
+        lower = np.asarray(self.lower_bound, dtype=float)
+        spacing = np.asarray(xp.to_numpy(self.grid_spacing), dtype=float)
+        index = np.indices(tuple(int(n) for n in self.num_gridpoints)).reshape(
+            self.ndim, -1
+        )
+        return lower[:, None] + index.astype(float) * spacing[:, None]
+
+    def indices_from_mask(self, mask: Array) -> list[int] | list[tuple[int, ...]]:
+        """Return the :attr:`indices` of the nodes where the grid-shaped ``mask`` is set.
+
+        ``mask`` is a boolean array of shape :attr:`num_gridpoints` (or flat, in
+        C order); the result keeps the order of :attr:`indices`.
+        """
+        mask = np.asarray(xp.to_numpy(mask), dtype=bool).reshape(
+            tuple(int(n) for n in self.num_gridpoints)
+        )
+        if self.ndim == 1:
+            return np.flatnonzero(mask).tolist()
+        return [tuple(index) for index in np.argwhere(mask).tolist()]
+
+    def node_coordinates(self, index: int | tuple[int, ...]) -> np.ndarray:
+        """Return the physical coordinates of a grid node.
+
+        ``index`` is a tuple of integer node indices (a plain integer in 1D).
+        The result is ``lower_bound + index * grid_spacing`` as a float array of
+        length ``ndim``, on the host.
+        """
+        index_array = np.atleast_1d(np.asarray(index, dtype=float))
+        if index_array.shape != (self.ndim,):
+            raise ValueError(
+                f"Expected an index with {self.ndim} entries, got {index}.",
+            )
+        lower = np.asarray(self.lower_bound, dtype=float)
+        spacing = np.asarray(xp.to_numpy(self.grid_spacing), dtype=float)
+        return lower + index_array * spacing
+
+    def get_cell_indices(self, *coords: Printable) -> tuple[int, ...]:
+        """
+        Get the cell indices for a given position, overridden to include cell size calculations.
+
+        Parameters:
+        coords (float): Coordinates for each dimension.
+
+        Returns:
+        tuple: Indices of the cell containing the position.
+        """
+
+        if len(coords) != self.ndim:
+            raise ValueError(f"Expected {self.ndim} dimensions, got {len(coords)}.")
+
+        indices = []
+
+        if self.contains(coords):
+            for coord, min_bound, cell_size in zip(
+                coords,
+                self.lower_bound,
+                self.grid_spacing,
+            ):
+                index = int((coord - min_bound) / cell_size)
+                indices.append(index)
+
+        return tuple(indices)
+
+    def get_grid_points(self) -> list[ndarray]:
+        """Return grid points for one axis."""
+        grid_points_list = []
+        periodic_axes = self.periodic or (False,) * self.ndim
+        for dim in range(self.ndim):
+            min_bound = self.lower_bound[dim]
+            max_bound = self.upper_bound[dim]
+            num_points = self.num_gridpoints[dim]
+            if periodic_axes[dim]:
+                # ``num_points`` cells over the length: the last node is one
+                # spacing before the upper bound (see ``grid_spacing``).
+                grid_points = min_bound + xp.arange(num_points) * self.grid_spacing[dim]
+            else:
+                grid_points = xp.linspace(min_bound, max_bound, num_points)
+            grid_points_list.append(grid_points)
+        return grid_points_list
+
+    @property
+    def meshgrid(self) -> tuple[Array, ...]:
+        """Return the meshgrid."""
+        return xp.meshgrid(*self.get_grid_points())
+
+    @property
+    def interpolator(self) -> ConfigValue:
+        """Return the interpolator."""
+        return self._interpolator
+
+    @property
+    def num_gridpoints_local(self) -> tuple[int, ...]:
+        """Return the local grid point counts."""
+        return self._num_gridpoints_local
+
+    @property
+    def num_gridpoints_local_with_ghosts(self) -> tuple[int, ...]:
+        """Return local grid point counts including ghosts."""
+        return self._num_gridpoints_local_with_ghosts
+
+    @property
+    def proc_index_bounds(self) -> list[tuple[int, int]]:
+        """Return this rank global index bounds."""
+        return self._proc_index_bounds
+
+    @property
+    def neighbor_ranks(self) -> list[tuple[int, int]]:
+        """Return neighboring MPI ranks."""
+        return self._neighbor_ranks

simgeometry/kernels/__init__.py

@@ -0,0 +1 @@
+"""Compiled containment kernels for geometric objects."""

simgeometry/kernels/_build.py

@@ -0,0 +1,53 @@
+"""The options every kernel folder passes to :meth:`cunumpy.kernels.Kernel.from_folder`.
+
+The ``__init__.py`` of a kernel folder ``simgeometry/kernels/<name>/`` declares its
+kernel with these::
+
+    kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS)
+
+which finds ``<name>_pyccel.py``, ``<name>_numpy.py`` and ``<name>_cuda.cu`` in
+the folder.
+Nothing is compiled here: the Pyccel version is compiled on its first host call
+(cached on disk by :func:`~simgeometry.kernels.compile.compile_cached`), the CUDA
+version on its first launch. Importing this module needs neither Pyccel nor CuPy.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+#: CUDA threads per block of every kernel.
+BLOCK_SIZE = 256
+
+#: Default dynamic shared memory per block of a CUDA device, in bytes.
+SHARED_MEMORY_PER_BLOCK = 48 * 1024
+
+
+def shared_memory_fits(num_values: int, itemsize: int = 8) -> bool:
+    """Whether ``num_values`` values fit the default shared memory of a block."""
+    return num_values * itemsize <= SHARED_MEMORY_PER_BLOCK
+
+
+def _compile(module: Any) -> Any:
+    from simgeometry.kernels.compile import compile_cached
+
+    return compile_cached(module)
+
+
+def _marker_count(args: tuple[Any, ...]) -> int:
+    """Use the last axis of the first array as the CUDA thread count."""
+    return int(args[0].shape[-1])
+
+
+#: ``<name>_pyccel.py`` compiled with Pyccel on first use; CUDA for CuPy
+#: arguments, the host version for NumPy ones (also while CuPy is active);
+#: without ``n_threads`` one CUDA thread per entry of the last axis of the
+#: first array (the marker count of a component-major ``(ncomp, N)`` or
+#: ``(N,)`` marker array). CUDA sources include cunumpy's array-view header.
+KERNEL_OPTIONS: dict[str, Any] = {
+    "host_suffix": "_pyccel",
+    "compile_host": _compile,
+    "dispatch": "arrays",
+    "block_size": BLOCK_SIZE,
+    "n_threads_from": _marker_count,
+}

simgeometry/kernels/_metal.py

@@ -0,0 +1,41 @@
+"""Helpers of the Metal versions (``<name>/<name>_metal.py``) of the kernels.
+
+A Metal version is the body of a Metal Shading Language kernel run by
+:class:`cunumpy.kernels.MetalKernel` on the GPU of an Apple-silicon Mac, wrapped in
+a function with the arguments of the Pyccel version: NumPy arrays in, results
+written into the output arrays. The GPU has no float64, so the kernels compute in
+float32 (``float64="cast"``): expect a relative error of about 1e-5 against the
+CPU versions. Import a version only where :func:`metal_available` is true.
+"""
+
+from __future__ import annotations
+
+import cunumpy as xp
+
+#: Threads per threadgroup of every Metal kernel.
+THREADGROUP = 256
+
+
+def metal_available() -> bool:
+    """Whether the Metal versions can run here (MLX installed, Metal GPU present)."""
+    return bool(xp.kernels.metal_available())
+
+
+def make_kernel(
+    name: str,
+    source: str,
+    inputs: list[str],
+    outputs: list[str],
+    header: str = "",
+) -> xp.kernels.MetalKernel:
+    """A :class:`~cunumpy.kernels.MetalKernel` computing in float32, outputs zeroed."""
+    return xp.kernels.MetalKernel(
+        source,
+        inputs=inputs,
+        outputs=outputs,
+        name=name,
+        header=header,
+        threadgroup=THREADGROUP,
+        float64="cast",
+        init_value=0,
+    )

simgeometry/kernels/compile.py

@@ -0,0 +1,24 @@
+"""Compile simgeometry's Pyccel kernel modules once and reuse the builds across processes.
+
+The cache is :class:`cunumpy.kernels.PyccelCache`: a kernel module is built with
+``pyccel compile`` on its first use, keyed on its source, the Pyccel and Python
+versions and the platform, and stored under ``SIMGEOMETRY_KERNEL_CACHE`` (default
+``~/.cache/simgeometry/kernels``, or ``<tmp>/simgeometry-kernels`` if that is not writable).
+Later processes import the stored extension directly.
+
+Kernels are built with OpenMP, so the loops marked ``#$ omp parallel for`` run on
+all cores (limit them with ``OMP_NUM_THREADS``; under MPI set it to the cores per
+rank). ``SIMGEOMETRY_OPENMP=0`` builds them serial, except the kernels that call the OpenMP
+runtime; the setting is part of the cache key.
+"""
+
+from __future__ import annotations
+
+import cunumpy as xp
+
+#: The on-disk build cache of simgeometry's kernels.
+KERNEL_CACHE = xp.kernels.PyccelCache("simgeometry")
+
+#: Return the compiled form of a kernel module, building it on first use
+#: (:meth:`cunumpy.kernels.PyccelCache.compile`).
+compile_cached = KERNEL_CACHE.compile

simgeometry/kernels/contains_ball/__init__.py

@@ -0,0 +1,12 @@
+"""The contains_ball kernel: Pyccel, NumPy and CUDA versions with the same arguments.
+
+``kernel(...)`` runs the CUDA version when the arrays are CuPy arrays, and
+otherwise the Pyccel version, or the NumPy version if Pyccel is unavailable
+(see ``cunumpy.kernels.set_host_kernel_implementation`` to choose).
+"""
+
+import cunumpy as xp
+
+from simgeometry.kernels._build import KERNEL_OPTIONS
+
+kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS)

simgeometry/kernels/contains_ball/contains_ball_cuda.cu

@@ -0,0 +1,27 @@
+// The CUDA version (GPU) of the contains_ball kernel; contains_ball_pyccel.py next
+// to this file is the Pyccel version (CPU), with the same arguments.  Loaded by
+// picnic.kernels.catalog: one thread per marker; positions are (ndim, N).
+//
+// mask[marker] = |position - midpoint| <= radius (negated if inverted).
+
+#include "cunumpy/array_view.cuh"
+
+extern "C" __global__
+void contains_ball(
+    Array2D<double> positions,
+    Array1D<char> mask,
+    Array1D<double> midpoint,
+    double radius,
+    int inverted
+) {
+    const long long marker = blockDim.x * (long long)blockIdx.x + threadIdx.x;
+    if (marker >= positions.shape[1]) return;
+    double distance_squared = 0.0;
+    for (long long axis = 0; axis < positions.shape[0]; ++axis) {
+        const double delta = positions(axis, marker) - midpoint(axis);
+        distance_squared += delta * delta;
+    }
+    bool inside = distance_squared <= radius * radius;
+    if (inverted == 1) inside = !inside;
+    mask(marker) = inside ? 1 : 0;
+}

simgeometry/kernels/contains_ball/contains_ball_metal.py

@@ -0,0 +1,53 @@
+"""The Metal version (Apple-silicon GPU, float32) of the ``contains_ball`` kernel.
+
+``contains_ball_pyccel.py`` next to this file is the CPU version with the same
+arguments. One thread per marker (``positions`` is ``(ndim, N)``); see :mod:`simgeometry.kernels._metal`. The test is
+done in float32, so a marker within about 1e-7 (relative) of the boundary may be
+classified differently from the CPU version.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from simgeometry._types import Array
+from simgeometry.kernels._metal import make_kernel
+
+_SOURCE = """
+    uint marker = thread_position_in_grid.x;
+    uint n = positions_shape[1];
+    if (marker >= n) return;
+    uint ndim = positions_shape[0];
+    float distance_squared = 0.0f;
+    for (uint a = 0; a < ndim; ++a) {
+        float delta = positions[a * n + marker] - midpoint[a];
+        distance_squared += delta * delta;
+    }
+    bool inside = distance_squared <= params[0] * params[0];
+    if (params[1] == 1.0f) inside = !inside;
+    inside_flag[marker] = inside ? 1 : 0;
+"""
+
+_KERNEL = make_kernel(
+    "simgeometry_contains_ball",
+    _SOURCE,
+    inputs=["positions", "midpoint", "params"],
+    outputs=["inside_flag"],
+)
+
+
+def contains_ball(
+    positions: Array,
+    mask: Array,
+    midpoint: Array,
+    radius: float,
+    inverted: int,
+) -> None:
+    """Set ``mask[marker]`` to whether the marker lies within ``radius`` of ``midpoint``."""
+    if not positions.shape[1]:
+        return
+    flags = np.zeros(positions.shape[1], dtype=np.uint8)
+    _KERNEL(
+        positions, midpoint, np.array([radius, inverted], dtype=np.float32), out=flags
+    )
+    mask[:] = flags.astype(bool)

simgeometry/kernels/contains_ball/contains_ball_numpy.py

@@ -0,0 +1,23 @@
+"""The NumPy version of the ``contains_ball`` kernel, run when Pyccel is unavailable.
+
+Same arguments and result as ``contains_ball_pyccel.py``, vectorized over the markers.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def contains_ball(
+    positions: np.ndarray,
+    mask: np.ndarray,
+    midpoint: np.ndarray,
+    radius: float,
+    inverted: int,
+) -> None:
+    """Set ``mask`` to whether each marker lies within ``radius`` of ``midpoint``."""
+    inside = (
+        np.sum((positions - np.asarray(midpoint)[:, None]) ** 2, axis=0)
+        <= radius * radius
+    )
+    mask[...] = ~inside if inverted == 1 else inside

simgeometry/kernels/contains_ball/contains_ball_pyccel.py

@@ -0,0 +1,31 @@
+"""Mask of the markers inside a ball (disc in 2D, sphere in 3D, any dimension).
+
+The Pyccel version (CPU) of the ``contains_ball`` kernel; ``contains_ball_cuda.cu``
+next to this file is the CUDA version (GPU), with the same arguments.  ``positions`` is ``(ndim, N)`` (component-major).  A point
+on the surface counts as inside.  With ``inverted == 1`` the mask is negated.
+"""
+
+
+def contains_ball(
+    positions: "float[:,:]",
+    mask: "bool[:]",
+    midpoint: "float[:]",
+    radius: float,
+    inverted: int,
+) -> None:
+    """Set ``mask[marker]`` to whether the marker lies within ``radius`` of ``midpoint``."""
+    num_markers = positions.shape[1]
+    ndim = positions.shape[0]
+    radius_squared = radius * radius
+    # fmt: off
+    #$ omp parallel for private(distance_squared, axis, delta, inside) schedule(dynamic, 4096)
+    # fmt: on
+    for marker in range(num_markers):
+        distance_squared = 0.0
+        for axis in range(ndim):
+            delta = positions[axis, marker] - midpoint[axis]
+            distance_squared += delta * delta
+        inside = distance_squared <= radius_squared
+        if inverted == 1:
+            inside = not inside
+        mask[marker] = inside

simgeometry/kernels/contains_box/__init__.py

@@ -0,0 +1,12 @@
+"""The contains_box kernel: Pyccel, NumPy and CUDA versions with the same arguments.
+
+``kernel(...)`` runs the CUDA version when the arrays are CuPy arrays, and
+otherwise the Pyccel version, or the NumPy version if Pyccel is unavailable
+(see ``cunumpy.kernels.set_host_kernel_implementation`` to choose).
+"""
+
+import cunumpy as xp
+
+from simgeometry.kernels._build import KERNEL_OPTIONS
+
+kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS)

simgeometry/kernels/contains_box/contains_box_cuda.cu

@@ -0,0 +1,26 @@
+// The CUDA version (GPU) of the contains_box kernel; contains_box_pyccel.py next
+// to this file is the Pyccel version (CPU), with the same arguments.  Loaded by
+// picnic.kernels.catalog: one thread per marker; positions are (ndim, N).
+//
+// mask[marker] = marker inside the box [lower, upper] (negated if inverted).
+
+#include "cunumpy/array_view.cuh"
+
+extern "C" __global__
+void contains_box(
+    Array2D<double> positions,
+    Array1D<char> mask,
+    Array1D<double> lower,
+    Array1D<double> upper,
+    int inverted
+) {
+    const long long marker = blockDim.x * (long long)blockIdx.x + threadIdx.x;
+    if (marker >= positions.shape[1]) return;
+    bool inside = true;
+    for (long long axis = 0; axis < positions.shape[0]; ++axis) {
+        const double x = positions(axis, marker);
+        if (x < lower(axis) || x > upper(axis)) inside = false;
+    }
+    if (inverted == 1) inside = !inside;
+    mask(marker) = inside ? 1 : 0;
+}

simgeometry/kernels/contains_box/contains_box_metal.py

@@ -0,0 +1,50 @@
+"""The Metal version (Apple-silicon GPU, float32) of the ``contains_box`` kernel.
+
+``contains_box_pyccel.py`` next to this file is the CPU version with the same
+arguments. One thread per marker (``positions`` is ``(ndim, N)``); see :mod:`simgeometry.kernels._metal`. The test is
+done in float32, so a marker within about 1e-7 (relative) of the boundary may be
+classified differently from the CPU version.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from simgeometry._types import Array
+from simgeometry.kernels._metal import make_kernel
+
+_SOURCE = """
+    uint marker = thread_position_in_grid.x;
+    uint n = positions_shape[1];
+    if (marker >= n) return;
+    uint ndim = positions_shape[0];
+    bool inside = true;
+    for (uint a = 0; a < ndim; ++a) {
+        float x = positions[a * n + marker];
+        if (x < lower[a] || x > upper[a]) inside = false;
+    }
+    if (flags[0] == 1) inside = !inside;
+    inside_flag[marker] = inside ? 1 : 0;
+"""
+
+_KERNEL = make_kernel(
+    "simgeometry_contains_box",
+    _SOURCE,
+    inputs=["positions", "lower", "upper", "flags"],
+    outputs=["inside_flag"],
+)
+
+
+def contains_box(
+    positions: Array,
+    mask: Array,
+    lower: Array,
+    upper: Array,
+    inverted: int,
+) -> None:
+    """Set ``mask[marker]`` to whether the marker lies in the box ``[lower, upper]``."""
+    if not positions.shape[1]:
+        return
+    flags = np.zeros(positions.shape[1], dtype=np.uint8)
+    _KERNEL(positions, lower, upper, np.array([inverted], dtype=np.int32), out=flags)
+    mask[:] = flags.astype(bool)

simgeometry/kernels/contains_box/contains_box_numpy.py

@@ -0,0 +1,24 @@
+"""The NumPy version of the ``contains_box`` kernel, run when Pyccel is unavailable.
+
+Same arguments and result as ``contains_box_pyccel.py``, vectorized over the markers.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def contains_box(
+    positions: np.ndarray,
+    mask: np.ndarray,
+    lower: np.ndarray,
+    upper: np.ndarray,
+    inverted: int,
+) -> None:
+    """Set ``mask`` to whether each marker lies in the box ``[lower, upper]``."""
+    inside = np.all(
+        (positions >= np.asarray(lower)[:, None])
+        & (positions <= np.asarray(upper)[:, None]),
+        axis=0,
+    )
+    mask[...] = ~inside if inverted == 1 else inside

simgeometry/kernels/contains_box/contains_box_pyccel.py

@@ -0,0 +1,32 @@
+"""Mask of the markers inside an axis-aligned box (any dimension).
+
+The Pyccel version (CPU) of the ``contains_box`` kernel; ``contains_box_cuda.cu``
+next to this file is the CUDA version (GPU), with the same arguments.  ``positions`` is ``(ndim, N)`` (component-major).  A point
+on the box surface counts as inside.  With ``inverted == 1`` the mask is
+negated: the solid is everything outside the box.
+"""
+
+
+def contains_box(
+    positions: "float[:,:]",
+    mask: "bool[:]",
+    lower: "float[:]",
+    upper: "float[:]",
+    inverted: int,
+) -> None:
+    """Set ``mask[marker]`` to whether the marker lies in the box ``[lower, upper]``."""
+    num_markers = positions.shape[1]
+    ndim = positions.shape[0]
+    # fmt: off
+    #$ omp parallel for private(inside, axis) schedule(dynamic, 4096)
+    # fmt: on
+    for marker in range(num_markers):
+        inside = True
+        for axis in range(ndim):
+            if positions[axis, marker] < lower[axis]:
+                inside = False
+            if positions[axis, marker] > upper[axis]:
+                inside = False
+        if inverted == 1:
+            inside = not inside
+        mask[marker] = inside

simgeometry/kernels/contains_polygon_2d/__init__.py

@@ -0,0 +1,12 @@
+"""The contains_polygon_2d kernel: Pyccel, NumPy and CUDA versions with the same arguments.
+
+``kernel(...)`` runs the CUDA version when the arrays are CuPy arrays, and
+otherwise the Pyccel version, or the NumPy version if Pyccel is unavailable
+(see ``cunumpy.kernels.set_host_kernel_implementation`` to choose).
+"""
+
+import cunumpy as xp
+
+from simgeometry.kernels._build import KERNEL_OPTIONS
+
+kernel = xp.kernels.Kernel.from_folder(__name__, **KERNEL_OPTIONS)