fullwave25 1.3.0.dev0__py3-none-any.whl → 1.3.2.dev0__py3-none-any.whl

fullwave/__init__.py

@@ -1,10 +1,13 @@
 """fullwave module."""
 
 import logging
+import os
 import platform
 import time
 from importlib.metadata import PackageNotFoundError, version
 
+import numexpr
+
 from .grid import Grid
 from .medium import Medium, MediumExponentialAttenuation, MediumRelaxationMaps
 from .sensor import Sensor
@@ -63,7 +66,17 @@ try:
     __version__ = version("fullwave")
 except PackageNotFoundError:
     # Update via bump-my-version, not manually
-    __version__ = "1.3.0-dev0"
+    __version__ = "1.3.2-dev0"
 
 VERSION = __version__  # for convenience
 logger.info("Fullwave version: %s", __version__)
+
+# numexpr keeps 16 threads unless told otherwise, and the medium, the absorbing
+# layer and the coefficients run through it on grids of hundreds of millions of
+# cells. A thread count set in the environment is left alone.
+if not any(
+    os.environ.get(name)
+    for name in ("NUMEXPR_MAX_THREADS", "NUMEXPR_NUM_THREADS", "OMP_NUM_THREADS")
+):
+    numexpr.set_num_threads(min(numexpr.detect_number_of_cores(), numexpr.MAX_THREADS))
+logger.info("numexpr threads: %d", numexpr.get_num_threads())

fullwave/medium.py

@@ -19,6 +19,7 @@ from fullwave.solver.utils import initialize_relaxation_param_dict
 from fullwave.utils import check_functions, plot_utils
 from fullwave.utils.coordinates import coords_to_map, map_to_coords
 from fullwave.utils.relaxation_parameters import (
+    _array_module,
     band_scaled_sound_speed,
     generate_relaxation_params,
 )
@@ -200,7 +201,8 @@ def _make_the_lossless_voxels_lossless(
     """
     if not relaxation_param_dict:
         return
-    coefficient = np.asarray(alpha_coeff)
+    xp = np if np.ndim(alpha_coeff) == 0 else _array_module(alpha_coeff)
+    coefficient = xp.asarray(alpha_coeff)
     if coefficient.ndim == 0:
         if float(coefficient) != 0.0:
             return
@@ -273,6 +275,7 @@ class MediumRelaxationMaps:
         n_jobs: int = -1,
         dtype: type = np.float64,
         use_gpu: bool = False,
+        compute_coefficients: bool = True,
     ) -> None:
         """Medium class for Fullwave.
 
@@ -330,6 +333,10 @@ class MediumRelaxationMaps:
         use_gpu : bool, optional
             If True, use CuPy for GPU-accelerated computation (default is False).
             Requires CuPy to be installed. Falls back to CPU if CuPy is unavailable.
+        compute_coefficients : bool, optional
+            If True (the default), compute ``relaxation_param_dict_for_fw2``
+            here. False leaves it empty for a caller that writes every entry
+            itself, as the PML builder does.
 
         """
         check_functions.check_compatible_value(
@@ -354,7 +361,6 @@ class MediumRelaxationMaps:
         try:
             self.relaxation_param_dict = initialize_relaxation_param_dict(
                 n_relaxation_mechanisms=n_relaxation_mechanisms,
-                value=xp.zeros_like(xp.asarray(sound_speed), dtype=self.dtype),
             )
             self.grid = grid
             self.is_3d = grid.is_3d
@@ -379,7 +385,6 @@ class MediumRelaxationMaps:
             xp = np
             self.relaxation_param_dict = initialize_relaxation_param_dict(
                 n_relaxation_mechanisms=n_relaxation_mechanisms,
-                value=np.zeros_like(np.asarray(sound_speed), dtype=self.dtype),
             )
             self.grid = grid
             self.is_3d = grid.is_3d
@@ -404,12 +409,23 @@ class MediumRelaxationMaps:
             relaxation_param_updates=relaxation_param_dict,
         )
         self.use_isotropic_relaxation = use_isotropic_relaxation
-        self.relaxation_param_dict_for_fw2 = self._calc_relaxation_param_dict_for_fw2(
-            use_isotropic_relaxation=self.use_isotropic_relaxation,
+        self.relaxation_param_dict_for_fw2 = (
+            self._calc_relaxation_param_dict_for_fw2(
+                use_isotropic_relaxation=self.use_isotropic_relaxation,
+            )
+            if compute_coefficients
+            else {}
         )
         self.check_fields()
         logger.debug("MediumRelaxationMaps instance created.")
 
+    def build_coefficients(self) -> None:
+        """Compute ``relaxation_param_dict_for_fw2`` if the constructor was told to skip it."""
+        if not self.relaxation_param_dict_for_fw2:
+            self.relaxation_param_dict_for_fw2 = self._calc_relaxation_param_dict_for_fw2(
+                use_isotropic_relaxation=self.use_isotropic_relaxation,
+            )
+
     def _update_relaxation_param_dict(
         self,
         relaxation_param_updates: dict[str, NDArray[np.float64]],
@@ -1627,12 +1643,18 @@ class Medium:
         parameters = self._relaxation_parameters_of(sound_speed, alpha_coeff, alpha_power)
         return parameters, sound_speed
 
-    def build(self) -> MediumRelaxationMaps:
+    def build(self, *, compute_coefficients: bool = True) -> MediumRelaxationMaps:
         """Retrieve the relaxation parameters from alpha and power maps.
 
         it uses the relaxation parameters look up table
         to generate the relaxation parameters.
 
+        Parameters
+        ----------
+        compute_coefficients : bool, optional
+            Passed to `MediumRelaxationMaps`. False skips the solver
+            coefficients, which a caller that extends the maps first does not use.
+
         Returns
         -------
             MediumRelaxationMaps: An instance of MediumRelaxationMaps
@@ -1670,6 +1692,7 @@ class Medium:
             n_jobs=self.n_jobs,
             dtype=self.dtype,
             use_gpu=self.use_gpu,
+            compute_coefficients=compute_coefficients,
         )
 
     def _db_mhz_cm_to_a_exp(

fullwave/solver/binary_manager.py

@@ -28,7 +28,7 @@ CACHE_DIR = Path.home() / ".cache" / "fullwave25" / "bins"
 
 # Pinned release tag for the solver binaries.
 # Update this only when new binaries are uploaded to a GitHub release.
-BINARY_RELEASE_TAG = "fullwave_bin_v1.4.1"
+BINARY_RELEASE_TAG = "fullwave_bin_v1.4.2"
 
 
 def _download_url(filename: str, tag: str) -> str:

fullwave/solver/input_file_writer.py

@@ -56,7 +56,7 @@ class InputFileWriter:
         path_fullwave_simulation_bin: Path = Path(__file__).parent / "bins" / "fullwave_solver_gpu",
         validate_input: bool = True,
         use_exponential_attenuation: bool = False,
-        use_isotropic_relaxation: bool = False,
+        use_isotropic_relaxation: bool = True,
         release_after_write: bool = False,
         pml_thickness: int = 0,
         exponential_attenuation_pml_thickness_px: int = 0,
@@ -94,6 +94,7 @@ class InputFileWriter:
             to reduce memory usage while retaining accuracy.
             For 2D it will reduce the memory usage by approximately 15%.
             For 3D it will reduce the memory usage by approximately 25%.
+            Only True is supported, and False raises NotImplementedError.
             This option omits the anisotropic relaxation mechanisms to model the attenuation.
             We usually recommend using isotropic relaxation mechanisms
             unless the anisotropic attenuation is required for the simulation.
@@ -126,6 +127,12 @@ class InputFileWriter:
         """
         logger.debug("Initializing InputFileWriter instance.")
 
+        if use_isotropic_relaxation is False:
+            error_msg = (
+                "Currently, only isotropic relaxation is supported. "
+                "Please set use_isotropic_relaxation to True for the simulation."
+            )
+            raise NotImplementedError(error_msg)
         self._work_dir = Path(work_dir)
         self.path_fullwave_simulation_bin = path_fullwave_simulation_bin
         self.use_isotropic_relaxation = use_isotropic_relaxation
@@ -1228,16 +1235,8 @@ class InputFileWriter:
         mechanisms = range(1, self.medium.n_relaxation_mechanisms + 1)
         for nu in mechanisms:
             var_name_list.extend([f"apmlu{nu}", f"bpmlu{nu}", f"apmlx{nu}", f"bpmlx{nu}"])
-        if not self.use_isotropic_relaxation:
-            var_name_list.extend(["kappay", "kappaw"])
-            for nu in mechanisms:
-                var_name_list.extend([f"apmlw{nu}", f"apmly{nu}", f"bpmlw{nu}", f"bpmly{nu}"])
         if self.is_3d:
-            var_name_list.append("modZ")
-        if self.is_3d and not self.use_isotropic_relaxation:
-            var_name_list.extend(["nZ", "dZ", "kappaz", "kappav"])
-            for nu in mechanisms:
-                var_name_list.extend([f"apmlz{nu}", f"apmlv{nu}", f"bpmlz{nu}", f"bpmlv{nu}"])
+            var_name_list.extend(["modZ", "nZ", "dZ"])
         for var_name in var_name_list:
             src_data = src_dir / f"{var_name}.dat"
             dst_data = dst_dir / f"{var_name}.dat"
@@ -1409,6 +1408,29 @@ class InputFileWriter:
             dst.unlink()
         dst.symlink_to(self.path_fullwave_simulation_bin.resolve())
 
+    @staticmethod
+    def _write_in_slabs(
+        variable_mat: np.ndarray,
+        dtype: DTypeLike,
+        save_path: str | Path,
+        slab_bytes: int = 256 * 1024**2,
+    ) -> None:
+        """Write an array in C order, cast one slab of its first axis at a time.
+
+        The bytes are those of one whole cast buffer, and the largest buffer
+        held is one slab, so many maps can be written at once.
+        """
+        dtype = np.dtype(dtype)
+        if variable_mat.ndim == 0:
+            np.ascontiguousarray(variable_mat, dtype=dtype).tofile(save_path)
+            return
+        row_bytes = max(1, dtype.itemsize * (variable_mat.size // max(variable_mat.shape[0], 1)))
+        rows_per_slab = max(1, slab_bytes // row_bytes)
+        with Path(save_path).open("wb") as handle:
+            for start in range(0, variable_mat.shape[0], rows_per_slab):
+                slab = variable_mat[start : start + rows_per_slab]
+                np.ascontiguousarray(slab, dtype=dtype).tofile(handle)
+
     @staticmethod
     def _write_ic(
         fname: str | Path,
@@ -1418,11 +1440,7 @@ class InputFileWriter:
         logger.debug("Writing initial condition matrix to %s", fname, stacklevel=2)
 
         t0 = time.perf_counter()
-        out = np.ascontiguousarray(
-            icmat.T,
-            dtype=np.float32,
-        )  # does transpose+cast into one new buffer
-        out.tofile(fname)
+        InputFileWriter._write_in_slabs(icmat.T, np.float32, fname)
         t1 = time.perf_counter()
 
         logger.debug("Initial condition matrix written in %.2e seconds", t1 - t0, stacklevel=2)
@@ -1486,8 +1504,7 @@ class InputFileWriter:
         if variable_mat.dtype == dtype and variable_mat.flags.c_contiguous:
             variable_mat.tofile(save_path)  # writes in C order
         else:
-            out = np.ascontiguousarray(variable_mat, dtype=dtype)  # at most one new buffer
-            out.tofile(save_path)
+            InputFileWriter._write_in_slabs(variable_mat, dtype, save_path)
 
         t1 = time.perf_counter()
         logger.debug("Matrix written in %.2e seconds", t1 - t0, stacklevel=2)

fullwave/solver/pml_builder.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import concurrent.futures
 import gc
 import logging
+import os
 from collections import OrderedDict
 from dataclasses import dataclass, field
 from functools import cached_property
@@ -312,12 +313,20 @@ class PMLBuilder:
         logger.debug("building extended medium for pml...")
         medium_as_given = self.medium_org
         if not isinstance(self.medium_org, fullwave.MediumRelaxationMaps):
-            self.medium_org = self.medium_org.build()
+            logger.info(
+                "Building the relaxation maps of %d cells...", int(self.medium_org.sound_speed.size)
+            )
+            self.medium_org = self.medium_org.build(compute_coefficients=False)
         base_attrs = ["sound_speed", "density", "beta"]
         original_alpha_coeff = getattr(self.medium_org, "alpha_coeff", None)
         if original_alpha_coeff is not None and np.ndim(original_alpha_coeff) != 0:
             base_attrs.append("alpha_coeff")
         relax_attrs = list(self.medium_org.relaxation_param_dict.keys())
+        logger.info(
+            "Extending %d maps by %d cells on each side...",
+            len(base_attrs) + len(relax_attrs),
+            self.num_boundary_points,
+        )
 
         if self.xp is not np:
             # Pass CuPy arrays directly — multi-GPU extension uses D2D copy (NVLink)
@@ -366,7 +375,11 @@ class PMLBuilder:
             n_jobs=self.medium_org.n_jobs,
             dtype=getattr(self.medium_org, "dtype", np.float64),
             use_gpu=False,
+            compute_coefficients=False,
         )
+        # The extended maps carry everything the absorbing layer reads, so the
+        # interior maps are released here rather than held through the run.
+        self.medium_org = medium_as_given
         logger.debug("building extended medium for pml...done")
 
         logger.debug("building extended source for pml...")
@@ -861,10 +874,16 @@ class PMLBuilder:
 
         eps = xp.finfo(xp.float64).eps
 
-        rate = d_x / kappa_x + alpha_x
         two_over_dt = 2.0 / dt
-        b = (two_over_dt - rate) / (two_over_dt + rate)
-        a = -(d_x / (kappa_x**2 + eps)) / (rate + two_over_dt)
+        if xp is np:
+            rate = ne.evaluate("d_x / kappa_x + alpha_x")
+            b = ne.evaluate("(two_over_dt - rate) / (two_over_dt + rate)")
+            a = ne.evaluate("-(d_x / (kappa_x ** 2 + eps)) / (rate + two_over_dt)")
+        else:
+            rate = d_x / kappa_x + alpha_x
+            b = (two_over_dt - rate) / (two_over_dt + rate)
+            a = -(d_x / (kappa_x**2 + eps)) / (rate + two_over_dt)
+        del rate
 
         if output_dtype is not None and output_dtype != xp.float64:
             a = a.astype(output_dtype, copy=False)
@@ -886,6 +905,10 @@ class PMLBuilder:
         """
         logger.debug("Running PML builder...")
         if use_pml:
+            logger.info(
+                "Building the absorbing layer on %d cells...",
+                int(self.extended_medium.sound_speed.size),
+            )
             extended_medium: fullwave.MediumRelaxationMaps = self.extended_medium.build()
             if self.is_3d:
                 return self._apply_pml_3d(
@@ -901,30 +924,54 @@ class PMLBuilder:
             )
 
         extended_medium: fullwave.MediumRelaxationMaps = self.extended_medium.build()
+        extended_medium.build_coefficients()
         return extended_medium
 
     def _medium_copy(self, relaxation_param_dict: dict, key: str) -> NDArray[np.float64]:
-        """Return a fresh float64 copy of one medium map."""
-        return self.xp.array(relaxation_param_dict[key], dtype=self.xp.float64, copy=True)
+        """Return a fresh float64 copy of one medium map.
+
+        On the host the copy is split along the first axis across threads,
+        because one thread copying a map of the whole padded grid is the slowest
+        step of the absorbing layer.
+        """
+        source = relaxation_param_dict[key]
+        if self.xp is not np or np.ndim(source) == 0:
+            return self.xp.array(source, dtype=self.xp.float64, copy=True)
+        copied = np.empty(np.shape(source), dtype=np.float64)
+        n_rows = copied.shape[0]
+        n_slabs = max(1, min(n_rows, os.cpu_count() or 1, ne.MAX_THREADS))
+        edges = np.linspace(0, n_rows, n_slabs + 1).astype(int)
+
+        def copy_slab(start: int, stop: int) -> None:
+            np.copyto(copied[start:stop], source[start:stop], casting="unsafe")
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=n_slabs) as pool:
+            list(pool.map(copy_slab, edges[:-1], edges[1:]))
+        return copied
 
     def _ramp_on_every_axis(
-        self, field: NDArray[np.float64], **kwargs: object
+        self, field: NDArray[np.float64], *, in_place: bool = False, **kwargs: object
     ) -> NDArray[np.float64]:
-        """Apply one ramp along each grid axis in turn, in two dimensions or three."""
+        """Apply one ramp along each grid axis in turn, in two dimensions or three.
+
+        The ramp writes only the layers at both ends of each axis. With
+        `in_place` the field is written itself, otherwise one copy is written.
+        """
         is_3d = field.ndim == self.three_dimensions
-        out = field
+        out = field if in_place else self.xp.array(field, copy=True)
         for axis_index in range(field.ndim):
             out = self._apply_transition_and_pml(
-                out, array_shape=field.shape, axis=axis_index, is_3d=is_3d, **kwargs
+                out, array_shape=field.shape, axis=axis_index, is_3d=is_3d, in_place=True, **kwargs
             )
         return out
 
     def _empty_damping_in_transition_layer(
-        self, damping: NDArray[np.float64]
+        self, damping: NDArray[np.float64], *, in_place: bool = False
     ) -> NDArray[np.float64]:
         """Return the damping strength taken to zero across the transition layer."""
         return self._ramp_on_every_axis(
             damping,
+            in_place=in_place,
             value_target=0.0,
             transition_type="cosine",
             transit_within_transition_layer=True,
@@ -935,10 +982,13 @@ class PMLBuilder:
         damping: NDArray[np.float64],
         n_polynomial: float,
         d_target_pml: float,
+        *,
+        in_place: bool = False,
     ) -> NDArray[np.float64]:
         """Return the damping strength raised to the PML target across the PML layer."""
         return self._ramp_on_every_axis(
             damping,
+            in_place=in_place,
             value_target=d_target_pml,
             n_polynomial=n_polynomial,
             transition_type="polynomial",
@@ -949,27 +999,36 @@ class PMLBuilder:
         self,
         relaxation_frequency: NDArray[np.float64],
         entrance: float,
+        *,
+        in_place: bool = False,
     ) -> NDArray[np.float64]:
         """Return the relaxation frequency replaced by the CFS ramp inside the PML layer.
 
         Replaced rather than ramped into, because ramping a relaxation frequency
         through the evaluation frequency is what confines the damping to the
-        outer edge.
+        outer edge. The PML layer is the slab of `edge` cells at both ends of
+        each axis, so only those slabs are copied from the ramp.
         """
         xp = self.xp
         edge = self.m_spatial_order + self.n_pml_layer
         ramp = self._ramp_on_every_axis(
             xp.full_like(relaxation_frequency, entrance),
+            in_place=True,
             value_target=0.0,
             transition_type="linear",
             transit_within_pml_layer=True,
         )
-        inside_pml_layer = xp.zeros(relaxation_frequency.shape, dtype=bool)
-        for axis_index in range(relaxation_frequency.ndim):
-            view = xp.moveaxis(inside_pml_layer, axis_index, 0)
-            view[:edge] = True
-            view[view.shape[0] - edge :] = True
-        return xp.where(inside_pml_layer, xp.asarray(ramp), xp.asarray(relaxation_frequency))
+        spliced = (
+            xp.asarray(relaxation_frequency)
+            if in_place
+            else xp.array(relaxation_frequency, copy=True)
+        )
+        for axis_index in range(spliced.ndim):
+            target = xp.moveaxis(spliced, axis_index, 0)
+            source = xp.moveaxis(ramp, axis_index, 0)
+            target[:edge] = source[:edge]
+            target[target.shape[0] - edge :] = source[source.shape[0] - edge :]
+        return spliced
 
     def _build_decoupled(
         self,
@@ -1002,16 +1061,16 @@ class PMLBuilder:
                 relaxation_frequency = self._medium_copy(
                     relaxation_param_dict, rename_dict[f"alpha_{letter}_nu{nu}"]
                 )
-                emptied = self._empty_damping_in_transition_layer(damping)
+                emptied = self._empty_damping_in_transition_layer(damping, in_place=True)
                 if nu > 1:
                     out_dict[f"d_{letter}_nu{nu}"] = emptied
                     out_dict[f"alpha_{letter}_nu{nu}"] = relaxation_frequency
                     continue
                 out_dict[f"d_{letter}_nu{nu}"] = self._raise_damping_in_pml_layer(
-                    emptied, n_polynomial, d_target_pml
+                    emptied, n_polynomial, d_target_pml, in_place=True
                 )
                 out_dict[f"alpha_{letter}_nu{nu}"] = self._splice_entrance_frequency_into_pml_layer(
-                    relaxation_frequency, entrance
+                    relaxation_frequency, entrance, in_place=True
                 )
         return out_dict
 
@@ -1124,9 +1183,11 @@ class PMLBuilder:
             found = np.asarray(relaxation_param_dict[key], dtype=np.float64)
             if found.ndim == 0 or np.allclose(found, 1.0):
                 continue
+            found = self.xp.array(found, copy=True)
             for axis in range(3 if is_3d else 2):
                 found = self._apply_transition_and_pml(
                     found,
+                    in_place=True,
                     value_target=1.0,
                     array_shape=found.shape,
                     axis=axis,
@@ -1213,13 +1274,9 @@ class PMLBuilder:
             # Return keys + values so parent can update dict safely
             return (f"a_pml_{axis}{nu}", a, f"b_pml_{axis}{nu}", b)
 
-        if self.xp is not np:
-            # GPU path: run sequentially
-            results = [_worker(nu, axis) for nu, axis in tasks]
-        else:
-            with concurrent.futures.ThreadPoolExecutor() as executor:
-                futures = [executor.submit(_worker, nu, axis) for nu, axis in tasks]
-                results = [f.result() for f in futures]
+        # One task at a time. numexpr threads each expression, and a pool of
+        # tasks held every task's temporaries of the whole grid at once.
+        results = [_worker(nu, axis) for nu, axis in tasks]
 
         for a_key, a_val, b_key, b_val in results:
             out_dict[a_key] = a_val
@@ -1312,13 +1369,9 @@ class PMLBuilder:
             # Return keys + values so parent can update dict safely
             return (f"a_pml_{axis}{nu}", a, f"b_pml_{axis}{nu}", b)
 
-        if self.xp is not np:
-            # GPU path: run sequentially
-            results = [_worker(nu, axis) for nu, axis in tasks]
-        else:
-            with concurrent.futures.ThreadPoolExecutor() as executor:
-                futures = [executor.submit(_worker, nu, axis) for nu, axis in tasks]
-                results = [f.result() for f in futures]
+        # One task at a time. numexpr threads each expression, and a pool of
+        # tasks held every task's temporaries of the whole grid at once.
+        results = [_worker(nu, axis) for nu, axis in tasks]
 
         for a_key, a_val, b_key, b_val in results:
             out_dict[a_key] = a_val
@@ -1347,6 +1400,7 @@ class PMLBuilder:
         transit_within_pml_layer: bool = False,
         disable_the_transition_and_pml: bool = False,
         is_3d: bool = False,
+        in_place: bool = False,
     ) -> NDArray[np.float64]:
         if transit_within_transition_layer and transit_within_pml_layer:
             error_msg = (
@@ -1428,8 +1482,11 @@ class PMLBuilder:
         if use_gpu:
             input_array = xp.asarray(input_array)
 
-        # Move axis to 0 for uniform processing
-        working_array = xp.moveaxis(input_array, axis, 0).copy()
+        # Move axis to 0 for uniform processing. Only the layers at both ends
+        # are written, so an array the caller owns needs no copy.
+        working_array = xp.moveaxis(input_array, axis, 0)
+        if not in_place:
+            working_array = working_array.copy()
 
         # Apply boundary conditions
         working_array[: m_offset + layer_thickness] = value_target

fullwave/solver/solver.py

@@ -1197,11 +1197,21 @@ class Solver:
 
         n_source_timesteps = source.icmat.shape[1]
 
-        gb = 1024.0**3
+        gib = 1024.0**3
 
         base_depth = depth // n_gpus
         remainder = depth % n_gpus
 
+        box, hole = (None, None)
+        if not self.use_exponential_attenuation:
+            box, hole = self._relaxation_regions(medium)
+        # Every GPU holds every source matrix and every zero pressure coordinate.
+        source_points = sum(
+            getattr(source, name, 0) or 0
+            for name in ("n_sources", "n_sources_add", "n_sources_u", "n_sources_v", "n_sources_w")
+        )
+        body_start = 0
+
         for rank, dev_id in enumerate(device_ids):
             depth_this = base_depth + (1 if rank < remainder else 0)
 
@@ -1213,10 +1223,12 @@ class Solver:
                 n_halo_sides = 2
             local_depth = depth_this + n_halo_sides * halo_depth
             slab = local_depth * lateral
+            low = max(body_start - (halo_depth if rank > 0 else 0), 0)
+            high = min(body_start + depth_this + (halo_depth if rank < n_gpus - 1 else 0), depth)
+            body_start += depth_this
 
             n_sources = max(source.n_sources // n_gpus, 0)
             n_sensors = max(sensor.n_sensors // n_gpus, 0)
-            n_air_local = max(medium.n_air // n_gpus, 0)
 
             if self.use_exponential_attenuation:
                 total = self._mem_exponential(
@@ -1231,29 +1243,32 @@ class Solver:
                     is_3d=self.is_3d,
                 )
             else:
+                band, interior = self._band_and_interior(low, high, grid.shape, box, hole)
                 total = self._mem_relaxation(
                     slab,
                     n_deriv_levels,
-                    n_sources,
+                    source_points,
                     n_source_timesteps,
                     save_gpu_memory=self.save_gpu_memory,
-                    n_air=n_air_local,
+                    n_air=medium.n_air,
                     n_sensors=n_sensors,
                     n_relax=self.n_relax_mechanisms,
                     float_bytes=float_bytes,
                     int_bytes=int_bytes,
                     is_3d=self.is_3d,
+                    band=band,
+                    interior=interior,
                 )
 
             mode = "exponential" if self.use_exponential_attenuation else "relaxation"
             saving = ", save_gpu_memory=True" if self.save_gpu_memory else ""
             logger.info(
                 "GPU memory estimate [GPU %s] (%s mode%s): "
-                "%.2f GB  (depth=%d +%d halo, lateral=%d)",
+                "%.2f GiB  (depth=%d +%d halo, lateral=%d)",
                 dev_id.strip(),
                 mode,
                 saving,
-                total / gb,
+                total / gib,
                 depth_this,
                 n_halo_sides * halo_depth,
                 lateral,
@@ -1337,8 +1352,20 @@ class Solver:
         float_bytes: int,
         int_bytes: int,
         is_3d: bool,
+        band: int | None = None,
+        interior: int | None = None,
     ) -> int:
-        """Return estimated GPU bytes for relaxation (power-law) solver.
+        """Return estimated GPU bytes for the relaxation (power-law) solver.
+
+        The count follows the allocation of the region trimmed kernel, the same
+        count as ``nrelax_device_bytes`` in the solver's ``gpu_memory_nrelax.cuh``.
+        Over the whole slab the card holds the pressure and one velocity for each
+        axis at two time levels, ``kappa_x2``, ``rho``, ``K``, the two finite
+        amplitude maps and ``dcmap``, and mechanism 1 of the pressure memory, its
+        ``a``, ``b`` and one memory variable for each axis. The mechanisms above
+        the first hold the same over the interior box alone, and the momentum
+        absorbing layer holds its ``a``, ``b`` and one memory variable for each
+        axis over the band around the hole alone.
 
         Parameters
         ----------
@@ -1347,7 +1374,7 @@ class Solver:
         n_deriv_levels : int
             Number of derivative-map levels.
         n_sources : int
-            Approximate source count on this GPU.
+            Source points this GPU holds, summed over every source kind.
         n_source_timesteps : int
             Number of source time steps.
         save_gpu_memory : bool
@@ -1362,8 +1389,13 @@ class Solver:
             Bytes per float (4).
         int_bytes : int
             Bytes per int (4).
-        is_3d: bool,
-            Whether the simulation is 3D (affects sensor memory).
+        is_3d : bool
+            Whether the simulation is 3D.
+        band : int | None
+            Grid points of the momentum layer band on this GPU. None takes the slab.
+        interior : int | None
+            Grid points of this GPU's share of the interior box of the mechanisms
+            above the first. None takes the slab.
 
         Returns
         -------
@@ -1373,36 +1405,152 @@ class Solver:
         """
         fb = float_bytes
         ib = int_bytes
-        ndim = 3 if is_3d else 2
-        n_fields = 4 if is_3d else 3  # p, u, [v], w
-
-        # wave fields: n_fields pairs x 2 time levels
-        mem = n_fields * 2 * slab * fb
-        # relaxation psi:
-        mem += 2 * (ndim * n_relax * 2 * slab * fb)
-        # material: rho + K + beta
-        mem += 3 * slab * fb
-        # kappa: 2 arrays (kappa_x1, kappa_x2)
-        mem += 2 * slab * fb
-        # PML: pml_x1 + pml_x2, each has 2 * n_relax arrays
-        mem += 2 * (2 * n_relax) * slab * fb
-        # (dmap + dcmap)
-        mem += 9 * 2 * n_deriv_levels * fb + slab * ib
-        # source (icmat + coords)
+        axes = 3 if is_3d else 2
+        band = slab if band is None else band
+        interior = slab if interior is None else interior
+        per_mechanism = 2 + axes
+
+        mem = 2 * (axes + 1) * slab * fb
+        mem += 5 * slab * fb
+        mem += slab * ib
+        mem += per_mechanism * slab * fb
+        mem += per_mechanism * max(n_relax - 1, 0) * interior * fb
+        mem += per_mechanism * band * fb
+        mem += 9 * 2 * n_deriv_levels * fb
         if n_sources > 0:
-            mem += n_sources * fb if save_gpu_memory else n_sources * n_source_timesteps * fb
-            mem += ndim * n_sources * ib
-
-        # air
+            samples = n_sources if save_gpu_memory else n_sources * n_source_timesteps
+            mem += samples * fb + axes * n_sources * ib
         if n_air > 0:
-            mem += ndim * n_air * ib
-        # sensor
+            mem += axes * n_air * ib
         if n_sensors > 0:
-            mem += n_sensors * fb
-            mem += (ndim + 1) * n_sensors * ib
-            mem += n_sensors * ib
+            mem += n_sensors * (fb + (axes + 2) * ib)
         return mem
 
+    @staticmethod
+    def _center_line_span(values: NDArray, axis: int, *, zeros: bool) -> tuple[int, int]:
+        """Return the first and one past the last index of a map's center line along one axis.
+
+        Parameters
+        ----------
+        values : NDArray
+            The map, on the host or on the GPU.
+        axis : int
+            The axis the line runs along. The other axes sit at their centers.
+        zeros : bool
+            True spans the zero cells, False the nonzero cells.
+
+        Returns
+        -------
+        tuple[int, int]
+            The span, empty as ``(0, 0)`` when no cell matches.
+
+        """
+        index = tuple(
+            slice(None) if a == axis else size // 2 for a, size in enumerate(values.shape)
+        )
+        line = values[index]
+        line = line.get() if hasattr(line, "get") else np.asarray(line)
+        found = np.flatnonzero(line == 0 if zeros else line != 0)
+        if found.size == 0:
+            return (0, 0)
+        return (int(found[0]), int(found[-1]) + 1)
+
+    def _relaxation_regions(
+        self,
+        medium: fullwave.MediumRelaxationMaps,
+    ) -> tuple[list[tuple[int, int]] | None, list[tuple[int, int]] | None]:
+        """Return the interior box and the momentum layer hole the kernel will measure.
+
+        The kernel stores every mechanism above the first over the box where its
+        ``a`` is nonzero, and the momentum absorbing layer over the band around the
+        box where the first momentum ``a`` is zero. Both are read here along the
+        center lines of each axis, which is how the kernel finds the hole, and a
+        clean box shape is assumed for both.
+
+        Parameters
+        ----------
+        medium : fullwave.MediumRelaxationMaps
+            The extended medium after the absorbing layer is built.
+
+        Returns
+        -------
+        tuple
+            The box and the hole as one ``(start, end)`` span for each axis. None
+            stands for a map the medium does not carry, which counts as the whole
+            slab.
+
+        """
+        maps = medium.relaxation_param_dict_for_fw2
+        box_maps = [maps.get(f"a_pml_u{nu}") for nu in range(2, self.n_relax_mechanisms + 1)]
+        box_maps = [values for values in box_maps if values is not None]
+        box = None
+        if self.n_relax_mechanisms > 1 and len(box_maps) == self.n_relax_mechanisms - 1:
+            box = []
+            for axis in range(box_maps[0].ndim):
+                spans = [self._center_line_span(values, axis, zeros=False) for values in box_maps]
+                spans = [span for span in spans if span[1] > span[0]]
+                box.append(
+                    (min(s for s, _ in spans), max(e for _, e in spans)) if spans else (0, 0)
+                )
+        hole = None
+        if maps.get("a_pml_x1") is not None:
+            values = maps["a_pml_x1"]
+            hole = [self._center_line_span(values, axis, zeros=True) for axis in range(values.ndim)]
+            if any(end <= start for start, end in hole):
+                hole = [(0, 0)] * values.ndim
+        return box, hole
+
+    @staticmethod
+    def _band_and_interior(
+        low: int,
+        high: int,
+        shape: tuple[int, ...],
+        box: list[tuple[int, int]] | None,
+        hole: list[tuple[int, int]] | None,
+    ) -> tuple[int, int]:
+        """Return one GPU's band and interior box cells, as the kernel counts them.
+
+        Parameters
+        ----------
+        low : int
+            The first grid row the GPU holds, its halo included.
+        high : int
+            One past the last grid row the GPU holds, its halo included.
+        shape : tuple[int, ...]
+            The extended grid shape.
+        box : list[tuple[int, int]] | None
+            The interior box, one span for each axis, or None for the whole slab.
+        hole : list[tuple[int, int]] | None
+            The momentum layer hole, one span for each axis, or None for the whole slab.
+
+        Returns
+        -------
+        tuple[int, int]
+            The band cells and the interior cells of this GPU.
+
+        """
+        rows = high - low
+        lateral = int(np.prod(shape[1:]))
+        if box is None:
+            interior = rows * lateral
+        else:
+            box_rows = max(0, min(high, box[0][1]) - max(low, box[0][0]))
+            interior = box_rows * int(np.prod([end - start for start, end in box[1:]]))
+        if hole is None:
+            return rows * lateral, interior
+        hole_low = min(max(hole[0][0] - low, 0), rows)
+        hole_high = max(min(max(hole[0][1] - low, 0), rows), hole_low)
+        hole_rows = hole_high - hole_low
+        band = (rows - hole_rows) * lateral
+        hole_columns = hole[1][1] - hole[1][0]
+        strip_columns = shape[1] - hole_columns
+        if len(shape) == 3:
+            strip_depths = shape[2] - (hole[2][1] - hole[2][0])
+            band += hole_rows * (strip_columns * shape[2] + hole_columns * strip_depths)
+        else:
+            band += hole_rows * strip_columns
+        return band, interior
+
     def print_info(self) -> None:
         """Print the Solver instance information."""
         print(str(self))

fullwave/transmit.py

@@ -272,12 +272,28 @@ def signal_of(
     layers = layer_of_each_pixel(coords, identifiers)
     center_hertz = grid.f0 if pulse.center_hertz is None else pulse.center_hertz
 
-    signal = np.zeros((len(identifiers), grid.nt))
-    for index, element in enumerate(identifiers):
+    # Pixels whose element weight, element delay, pixel delay and layer carry the
+    # same bits get the same row, so each row is built once, from one of its
+    # pixels, with the arguments that pixel gives.
+    element_index = identifiers - 1
+    columns = [
+        np.asarray(weights, dtype=np.float64)[element_index].view(np.uint64),
+        np.asarray(delays_s, dtype=np.float64)[element_index].view(np.uint64),
+        np.asarray(layers, dtype=np.uint64),
+    ]
+    if pixel_delays_s is not None:
+        columns.append(np.ascontiguousarray(pixel_delays_s, dtype=np.float64).view(np.uint64))
+    _, first_pixel, row_of_pixel = np.unique(
+        np.stack(columns, axis=1), axis=0, return_index=True, return_inverse=True
+    )
+
+    rows = np.zeros((len(first_pixel), grid.nt))
+    for row, index in enumerate(first_pixel):
+        element = identifiers[index]
         weight = float(weights[element - 1])
         if weight == 0.0:
             continue
-        signal[index] = pulse_utils.gaussian_modulated_sinusoidal_signal(
+        rows[row] = pulse_utils.gaussian_modulated_sinusoidal_signal(
             nt=grid.nt,
             f0=center_hertz,
             duration=grid.duration,
@@ -290,4 +306,4 @@ def signal_of(
             delay_sec=float(delays_s[element - 1])
             + (0.0 if pixel_delays_s is None else float(pixel_delays_s[index])),
         )
-    return signal
+    return rows[row_of_pixel.reshape(-1)]

fullwave/utils/relaxation_parameters.py

@@ -9,6 +9,7 @@ import time
 from pathlib import Path
 
 import numba as nb
+import numexpr as ne
 import numpy as np
 from numpy.typing import NDArray
 from scipy.io import loadmat
@@ -128,10 +129,17 @@ def _first_order_attenuation(
     exceeds one and the recursion grows without bound, so that is refused.
     """
     scaled = dict(relaxation_param_dict)
+    on_the_host = xp is np
     for direction in ("x1", "x2"):
         stretching = f"kappa_{direction}"
-        moved = 1.0 + ratio * (relaxation_param_dict[stretching] - 1.0)
-        if bool(xp.any(moved <= 0.0)):
+        kappa = relaxation_param_dict[stretching]
+        if on_the_host:
+            moved = ne.evaluate("1.0 + ratio * (kappa - 1.0)")
+            collapsed = bool(ne.evaluate("moved <= 0.0").any())
+        else:
+            moved = 1.0 + ratio * (kappa - 1.0)
+            collapsed = bool(xp.any(moved <= 0.0))
+        if collapsed:
             error_msg = (
                 f"scaling the attenuation by this factor carries {stretching} to "
                 f"{float(xp.min(moved)):.6g}, at or below zero, which makes the memory "
@@ -143,7 +151,8 @@ def _first_order_attenuation(
         scaled[stretching] = moved
         for i_relax in range(n_relaxation_mechanisms):
             strength = f"d_{direction}_nu{i_relax + 1}"
-            scaled[strength] = ratio * relaxation_param_dict[strength]
+            values = relaxation_param_dict[strength]
+            scaled[strength] = ne.evaluate("ratio * values") if on_the_host else ratio * values
     return scaled
 
 

{fullwave25-1.3.0.dev0.dist-info → fullwave25-1.3.2.dev0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: fullwave25
-Version: 1.3.0.dev0
+Version: 1.3.2-dev0
 Summary: Fullwave 2.5: Ultrasound wave propagation simulation with heterogeneous power law attenuation modelling capabilities
 Author: Masashi Sode, Gianmarco Pinton
 Maintainer: Masashi Sode
@@ -13,6 +13,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Scientific/Engineering :: Physics
 Requires-Python: >=3.10
+Requires-Dist: cupy-cuda12x>=13.6.0
 Requires-Dist: matplotlib>=3.10.7
 Requires-Dist: numba>=0.63.1
 Requires-Dist: numexpr>=2.14.1

{fullwave25-1.3.0.dev0.dist-info → fullwave25-1.3.2.dev0.dist-info}/RECORD

@@ -1,10 +1,10 @@
-fullwave/__init__.py,sha256=W_rJYZITUpH7crE6e87RdQCJQBvsDXDEm4g8k6mNm8I,1861
+fullwave/__init__.py,sha256=p-kSxLN7kRdorciAXD-u5Ffx4l8d6XHDpBuHRoRblPk,2382
 fullwave/grid.py,sha256=YbxAAgdYKR8AXwobtTZzHm9C4KR4D5--51zgbs0aQGY,7435
-fullwave/medium.py,sha256=vlGFkk-GZCuJGLfYRPi7Hc8VamI_uC4NQdxSCOFAhAA,67621
+fullwave/medium.py,sha256=Tz8Cwda6zLF8UFxjWFkRcAdko_KmF4l1zfmjWaqxu-M,68626
 fullwave/sensor.py,sha256=McHyaSKxM11_Kfb_2hfnOtV9Z1UMBCLkeovncLT8_bA,9308
 fullwave/source.py,sha256=2mJ9djoc_PZAdGbkpJBTPqiF58rxe4YcyUIyxAtjAWA,21143
 fullwave/transducer.py,sha256=NL_GxX4ViYfICeY-mrV1MrkWao_sEhgZWCQ4Lvhpuao,75263
-fullwave/transmit.py,sha256=6TK0ow86PM4ivgiU9DmD9pjy6WLqOu6JoHt7CrYi07I,8907
+fullwave/transmit.py,sha256=BAD_G-J7Q1pMppXw1yjLrQDUdkhPq-3Q4eNq4jaf7Ss,9693
 fullwave/beamformer/__init__.py,sha256=tNPq_5lfLXfAdueyfQtdTxLTVEV6nEDkeSGAeTf7Q5s,88
 fullwave/beamformer/beamformer.py,sha256=0jE_B5UlxDDBsUcQSHa5d1q3uSTSCyqgNiUJrU10juI,5601
 fullwave/constants/__init__.py,sha256=WOmkhLsSe86D71gwKsa4Tz8QPeOdayLGgvBCh7UEOPk,104
@@ -21,13 +21,13 @@ fullwave/medium_builder/presets/domain_water_gel.py,sha256=nFiL0CkmDeJWo7aR2Khn0
 fullwave/medium_builder/presets/data/.keep,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 fullwave/medium_builder/presets/data/abdominal_wall/i2365f_etfw1.mat,sha256=N3JFFxc9HW1S-NaTPcPlFgXWGIzlrr9X40JbDapTo4g,67348
 fullwave/solver/__init__.py,sha256=E0Jg5akBpicqTjku1puw9rAIoKiw7TwtYNDPiv8G5IY,30
-fullwave/solver/binary_manager.py,sha256=olqi2JGqLybQu8iThIo8D-KYtZaErk3clFpkT2LAKwg,5041
+fullwave/solver/binary_manager.py,sha256=VhI6VgkdfIwAeX4RkM5S0oHm_RWoo3Ph4BsNlz1vETM,5041
 fullwave/solver/cuda_utils.py,sha256=n7p3TJYbLAXmIzSMS5Q_8m4OWg5U9VG5_3hAxa0uTYI,11407
-fullwave/solver/input_file_writer.py,sha256=Z-Ha-1ceuAcT0CJM5iElsUVBKedwDRgWiLIaA-0Sm4M,57249
+fullwave/solver/input_file_writer.py,sha256=JamQHuA3-YrdLpdMFT6Ymw7noiCrtMr-DNo6LFxCvmM,57993
 fullwave/solver/launcher.py,sha256=zh-3HY4iKGmjreALu9WczDpWX3d74nOcIns9xVvnUIg,10531
-fullwave/solver/pml_builder.py,sha256=ClLInmYQSbhkno2-003dn-15V3cDy9x-F4NCMiBWlCA,84751
+fullwave/solver/pml_builder.py,sha256=Pfy3AklZT-p6TX61A5ntFmpt2mJb5TDaV-gdtpjIwNo,87341
 fullwave/solver/shipped_database.py,sha256=lPqlgY-ShDyRqtH65m5RyNIMVlqswECVDNh5XCk1Yzg,3194
-fullwave/solver/solver.py,sha256=val60dmnrcOmqMdBBVZw7QEHT8wDMp4Q1aFJMbkPc5c,58774
+fullwave/solver/solver.py,sha256=oQOrL-9YSi5C9vP6rOyfZW8soTHnWBLusi1rMIwI978,65346
 fullwave/solver/source_type.py,sha256=li4-4tpu4bd1i9DqyDXZKOhkzW4IV6OcxB97C13c5MU,11180
 fullwave/solver/utils.py,sha256=0Q9PKl-4hohWXGzuldqo2XubI-LEHMjs0kYgb4oTLC0,5561
 fullwave/solver/bins/_gpu/2d/num_relax=2/fullwave2_2d_2_relax_multi_gpu_cuda118,sha256=bIzyHFzr_bsIBSlCcRA3M9LlTSRrBVn6ASJHyozvcwU,1245432
@@ -49,11 +49,11 @@ fullwave/utils/memory_tempfile.py,sha256=BzMRNvINzLYdtV8NcSNjrTr2F9Qow8vYl1o7RUX
 fullwave/utils/numerical.py,sha256=khdi2KE3y_Ky9cEnmGXvVHjeYuKfb-8gQTV1a8nx9uA,3260
 fullwave/utils/plot_utils.py,sha256=SoByd2fEr3VdvMM9D2kNq8VgaAQ-fh82bDGRFdeq2R8,39327
 fullwave/utils/pulse.py,sha256=aqnCdBxivACouIGlot1aE2k4V9nctUxNCEPzdBpsupg,2811
-fullwave/utils/relaxation_parameters.py,sha256=uPN31h7YzgsxUBhHbEoxLapRplbpNITKJQPgbGoUtk4,45800
+fullwave/utils/relaxation_parameters.py,sha256=wHUNMtJGAJwgCyUYd24_qtQZflMnkQTYhxSJwjE3vA4,46151
 fullwave/utils/scatterer.py,sha256=CLqD085f9pqX4cF1YiT2IkzT5P0o1xp45Do7lf5BLnM,24049
 fullwave/utils/signal_filter.py,sha256=_Tf42aBygGtIUrcSPybYqvZoAW2o_4dd-hHTEvZ8--E,5219
 fullwave/utils/signal_process.py,sha256=MP38r37JSPHN4sT-lIReXS5BT2GT_tkRxD0Y1wQ-lNs,5153
-fullwave25-1.3.0.dev0.dist-info/METADATA,sha256=mgwujcZWq0A_pFpxRittcPsWh6GVO_xuKdTsQbqyXU0,21711
-fullwave25-1.3.0.dev0.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
-fullwave25-1.3.0.dev0.dist-info/licenses/LICENSE,sha256=46mU2C5kSwOnkqkw9XQAJlhBL2JAf1_uCD8lVcXyMRg,7652
-fullwave25-1.3.0.dev0.dist-info/RECORD,,
+fullwave25-1.3.2.dev0.dist-info/METADATA,sha256=wYeq9o1hmkqstMoW0B6zQrbOUQnQAIDBVbhkqGwHJKM,21747
+fullwave25-1.3.2.dev0.dist-info/WHEEL,sha256=6zicbvgMWfHSD_sSydPAPL1UNY8-QHIQA4J8fOxyAc8,87
+fullwave25-1.3.2.dev0.dist-info/licenses/LICENSE,sha256=46mU2C5kSwOnkqkw9XQAJlhBL2JAf1_uCD8lVcXyMRg,7652
+fullwave25-1.3.2.dev0.dist-info/RECORD,,

{fullwave25-1.3.0.dev0.dist-info → fullwave25-1.3.2.dev0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.32.0
+Generator: hatchling 1.32.1
 Root-Is-Purelib: true
 Tag: py3-none-any