zea 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

zea/__init__.py

@@ -7,7 +7,7 @@ from . import log
 
 # dynamically add __version__ attribute (see pyproject.toml)
 # __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.2"
+__version__ = "0.0.4"
 
 
 def setup():

zea/agent/selection.py

@@ -92,6 +92,7 @@ class GreedyEntropy(LinesActionModel):
         mean: float = 0,
         std_dev: float = 1,
         num_lines_to_update: int = 5,
+        entropy_sigma: float = 1.0,
     ):
         """Initialize the GreedyEntropy action selection model.
 
@@ -104,6 +105,8 @@ class GreedyEntropy(LinesActionModel):
             std_dev (float, optional): The standard deviation of the RBF. Defaults to 1.
             num_lines_to_update (int, optional): The number of lines around the selected line
                 to update. Must be odd.
+            entropy_sigma (float, optional): The standard deviation of the Gaussian
+                Mixture components used to approximate the posterior.
         """
         super().__init__(n_actions, n_possible_actions, img_width, img_height)
 
@@ -124,7 +127,7 @@ class GreedyEntropy(LinesActionModel):
             self.num_lines_to_update,
         )
         self.upside_down_gaussian = upside_down_gaussian(points_to_evaluate)
-        self.entropy_sigma = 1
+        self.entropy_sigma = entropy_sigma
 
     @staticmethod
     def compute_pairwise_pixel_gaussian_error(
@@ -152,51 +155,58 @@ class GreedyEntropy(LinesActionModel):
         # TODO: I think we only need to compute the lower triangular
         # of this matrix, since it's symmetric
         squared_l2_error_matrices = (particles[:, :, None, ...] - particles[:, None, :, ...]) ** 2
-        gaussian_error_per_pixel_i_j = ops.exp((squared_l2_error_matrices) / (2 * entropy_sigma**2))
+        gaussian_error_per_pixel_i_j = ops.exp(
+            -(squared_l2_error_matrices) / (2 * entropy_sigma**2)
+        )
         # Vertically stack all columns corresponding with the same line
         # This way we can just sum across the height axis and get the entropy
         # for each pixel in a given line
         batch_size, n_particles, _, height, _ = gaussian_error_per_pixel_i_j.shape
-        gaussian_error_per_pixel_stacked = ops.reshape(
-            gaussian_error_per_pixel_i_j,
-            [
-                batch_size,
-                n_particles,
-                n_particles,
-                height * stack_n_cols,
-                n_possible_actions,
-            ],
+        gaussian_error_per_pixel_stacked = ops.transpose(
+            ops.reshape(
+                ops.transpose(gaussian_error_per_pixel_i_j, (0, 1, 2, 4, 3)),
+                [
+                    batch_size,
+                    n_particles,
+                    n_particles,
+                    n_possible_actions,
+                    height * stack_n_cols,
+                ],
+            ),
+            (0, 1, 2, 4, 3),
         )
         # [n_particles, n_particles, batch, height, width]
         return gaussian_error_per_pixel_stacked
 
-    def compute_gmm_entropy_per_line(self, particles):
-        """Compute the entropy for each line using a Gaussian Mixture Model.
-
+    def compute_pixelwise_entropy(self, particles):
+        """
         This function computes the entropy for each line using a Gaussian Mixture Model
         approximation of the posterior distribution.
-        For more details see Section 4 here: https://arxiv.org/abs/2406.14388
+        For more details see Section VI. B here: https://arxiv.org/pdf/2410.13310
 
         Args:
             particles (Tensor): Particles of shape (batch_size, n_particles, height, width)
 
         Returns:
-            Tensor: batch of entropies per line, of shape (batch, n_possible_actions)
+            Tensor: batch of entropies per pixel, of shape (batch, height, width)
         """
-        gaussian_error_per_pixel_stacked = GreedyEntropy.compute_pairwise_pixel_gaussian_error(
+        n_particles = ops.shape(particles)[1]
+        gaussian_error_per_pixel_stacked = self.compute_pairwise_pixel_gaussian_error(
             particles,
             self.stack_n_cols,
             self.n_possible_actions,
             self.entropy_sigma,
         )
-        gaussian_error_per_line = ops.sum(gaussian_error_per_pixel_stacked, axis=3)
         # sum out first dimension of (n_particles x n_particles) error matrix
-        # [n_particles, batch, n_possible_actions]
-        entropy_per_line_i = ops.sum(gaussian_error_per_line, axis=1)
+        # [n_particles, batch, height, width]
+        pixelwise_entropy_sum_j = ops.sum(
+            (1 / n_particles) * gaussian_error_per_pixel_stacked, axis=1
+        )
+        log_pixelwise_entropy_sum_j = ops.log(pixelwise_entropy_sum_j)
         # sum out second dimension of (n_particles x n_particles) error matrix
-        # [batch, n_possible_actions]
-        entropy_per_line = ops.sum(entropy_per_line_i, axis=1)
-        return entropy_per_line
+        # [batch, height, width]
+        pixelwise_entropy = -ops.sum((1 / n_particles) * log_pixelwise_entropy_sum_j, axis=1)
+        return pixelwise_entropy
 
     def select_line_and_reweight_entropy(self, entropy_per_line):
         """Select the line with maximum entropy and reweight the entropies.
@@ -254,17 +264,19 @@ class GreedyEntropy(LinesActionModel):
             particles (Tensor): Particles of shape (batch_size, n_particles, height, width)
 
         Returns:
-           Tuple[Tensor, Tensor]:
+            Tuple[Tensor, Tensor]:
                 - Newly selected lines as k-hot vectors, shaped (batch_size, n_possible_actions)
-                - Masks of shape (batch_size, img_height, img_width)
+            - Masks of shape (batch_size, img_height, img_width)
         """
-        entropy_per_line = self.compute_gmm_entropy_per_line(particles)
+
+        pixelwise_entropy = self.compute_pixelwise_entropy(particles)
+        linewise_entropy = ops.sum(pixelwise_entropy, axis=1)
 
         # Greedily select best line, reweight entropies, and repeat
         all_selected_lines = []
         for _ in range(self.n_actions):
-            max_entropy_line, entropy_per_line = ops.vectorized_map(
-                self.select_line_and_reweight_entropy, entropy_per_line
+            max_entropy_line, linewise_entropy = ops.vectorized_map(
+                self.select_line_and_reweight_entropy, linewise_entropy
             )
             all_selected_lines.append(max_entropy_line)
 
@@ -428,10 +440,15 @@ class CovarianceSamplingLines(LinesActionModel):
                 generation. Defaults to None.
 
         Returns:
-            Tensor: The mask of shape (batch_size, img_size, img_size)
+            Tuple[Tensor, Tensor]:
+                - Newly selected lines as k-hot vectors, shaped (batch_size, n_possible_actions)
+                - Masks of shape (batch_size, img_height, img_width)
         """
         batch_size, n_particles, rows, _ = ops.shape(particles)
 
+        # [batch_size, rows, cols, n_particles]
+        particles = ops.transpose(particles, (0, 2, 3, 1))
+
         # [batch_size, rows * stack_n_cols, n_possible_actions, n_particles]
         shape = [
             batch_size,
@@ -441,7 +458,7 @@ class CovarianceSamplingLines(LinesActionModel):
         ]
         particles = ops.reshape(particles, shape)
 
-        # [batch_size, rows, n_possible_actions, n_possible_actions]
+        # [batch_size, rows * stack_n_cols, n_possible_actions, n_possible_actions]
         cov_matrix = tensor_ops.batch_cov(particles)
 
         # Sum over the row dimension [batch_size, n_possible_actions, n_possible_actions]

zea/backend/__init__.py

@@ -1,4 +1,4 @@
-"""Backend subpackage for ``zea``.
+"""Backend-specific utilities.
 
 This subpackage provides backend-specific utilities for the ``zea`` library. Most backend logic is handled by Keras 3, but a few features require custom wrappers to ensure compatibility and performance across JAX, TensorFlow, and PyTorch.
 
@@ -9,7 +9,7 @@ Key Features
 ------------
 
 - **JIT Compilation** (:func:`zea.backend.jit`):
-  Provides a unified interface for just-in-time (JIT) compilation of functions, dispatching to the appropriate backend (JAX or TensorFlow) as needed. This enables accelerated execution of computationally intensive routines.
+  Provides a unified interface for just-in-time (JIT) compilation of functions, dispatching to the appropriate backend (JAX or TensorFlow) as needed. This enables accelerated execution of computationally intensive routines. Note that jit compilation is not yet supported when using the `torch` backend.
 
 - **Automatic Differentiation** (:class:`zea.backend.AutoGrad`):
   Offers a backend-agnostic wrapper for automatic differentiation, allowing gradient computation regardless of the underlying ML library.
@@ -108,7 +108,9 @@ def _jit_compile(func, jax=True, tensorflow=True, **kwargs):
         return func
     else:
         log.warning(
-            f"Unsupported backend: {backend}. Supported backends are 'tensorflow' and 'jax'."
+            f"JIT compilation not currently supported for backend {backend}. "
+            "Supported backends are 'tensorflow' and 'jax'."
         )
+        log.warning("Initialize zea.Pipeline with jit_options=None to suppress this warning.")
         log.warning("Falling back to non-compiled mode.")
         return func

zea/beamform/__init__.py

@@ -17,4 +17,4 @@ see the pipeline example notebook: :doc:`../notebooks/pipeline/zea_pipeline_exam
 
 """
 
-from . import beamformer, delays, lens_correction, pfield, pixelgrid
+from . import beamformer, delays, lens_correction, pfield, phantoms, pixelgrid

zea/beamform/beamformer.py

@@ -1,5 +1,6 @@
 """Main beamforming functions for ultrasound imaging."""
 
+import keras
 import numpy as np
 from keras import ops
 
@@ -58,7 +59,7 @@ def tof_correction(
     demodulation_frequency,
     fnum,
     angles,
-    vfocus,
+    focus_distances,
     apply_phase_rotation=False,
     apply_lens_correction=False,
     lens_thickness=1e-3,
@@ -84,7 +85,7 @@ def tof_correction(
         fnum (int, optional): Focus number. Defaults to 1.
         angles (ops.Tensor): The angles of the plane waves in radians of shape
             `(n_tx,)`
-        vfocus (ops.Tensor): The focus distance of shape `(n_tx,)`
+        focus_distances (ops.Tensor): The focus distance of shape `(n_tx,)`
         apply_phase_rotation (bool, optional): Whether to apply phase rotation to
             time-of-flights. Defaults to False.
         apply_lens_correction (bool, optional): Whether to apply lens correction to
@@ -133,7 +134,7 @@ def tof_correction(
         sound_speed,
         n_tx,
         n_el,
-        vfocus,
+        focus_distances,
         angles,
         lens_thickness=lens_thickness,
         lens_sound_speed=lens_sound_speed,
@@ -487,7 +488,7 @@ def fnumber_mask(flatgrid, probe_geometry, f_number, fnum_window_fn):
     # The f-number is fnum = z/aperture = 1/(2 * tan(alpha))
     # Rearranging gives us alpha = arctan(1/(2 * fnum))
     # We can use this to compute the maximum angle alpha that is allowed
-    max_alpha = ops.arctan(1 / (2 * f_number))
+    max_alpha = ops.arctan(1 / (2 * f_number + keras.backend.epsilon()))
 
     normalized_angle = alpha / max_alpha
     mask = fnum_window_fn(normalized_angle)

zea/beamform/pfield.py

@@ -60,7 +60,8 @@ def compute_pfield(
         n_el (int): Number of elements in the probe.
         probe_geometry (array): Geometry of the probe elements.
         tx_apodizations (array): Transmit apodization values.
-        grid (array): Grid points where the pressure field is computed of shape (Nz, Nx, 3).
+        grid (array): Grid points where the pressure field is computed
+            of shape (grid_size_z, grid_size_x, 3).
         t0_delays (array): Transmit delays for each transmit event.
         frequency_step (int, optional): Frequency step. Default is 4.
             Higher is faster but less accurate.
@@ -78,7 +79,8 @@ def compute_pfield(
         verbose (bool, optional): Whether to print progress.
 
     Returns:
-        ops.array: The (normalized) pressure field (across tx events) of shape (n_tx, Nz, Nx).
+        ops.array: The (normalized) pressure field (across tx events)
+            of shape (n_tx, grid_size_z, grid_size_x).
     """
     # medium params
     alpha_db = 0  # currently we ignore attenuation in the compounding
@@ -293,7 +295,8 @@ def normalize_pressure_field(pfield, alpha: float = 1.0, percentile: float = 10.
     Normalize the input array of intensities by zeroing out values below a given percentile.
 
     Args:
-        pfield (array): The unnormalized pressure field array of shape (n_tx, Nz, Nx).
+        pfield (array): The unnormalized pressure field array
+            of shape (n_tx, grid_size_z, grid_size_x).
         alpha (float, optional): Exponent to 'sharpen or smooth' the weighting.
             Higher values result in sharper weighting. Default is 1.0.
         percentile (int, optional): minimum percentile threshold to keep in the weighting.

zea/beamform/pixelgrid.py

@@ -16,52 +16,52 @@ def check_for_aliasing(scan):
     depth = scan.zlims[1] - scan.zlims[0]
     wvln = scan.wavelength
 
-    if width / scan.Nx > wvln / 2:
+    if width / scan.grid_size_x > wvln / 2:
         log.warning(
-            f"width/Nx = {width / scan.Nx:.7f} < wavelength/2 = {wvln / 2}. "
-            f"Consider either increasing scan.Nx to {int(np.ceil(width / (wvln / 2)))} or more, or "
-            "increasing scan.pixels_per_wavelength to 2 or more."
+            f"width/grid_size_x = {width / scan.grid_size_x:.7f} < wavelength/2 = {wvln / 2}. "
+            f"Consider either increasing scan.grid_size_x to {int(np.ceil(width / (wvln / 2)))} "
+            "or more, or increasing scan.pixels_per_wavelength to 2 or more."
         )
-    if depth / scan.Nz > wvln / 2:
+    if depth / scan.grid_size_z > wvln / 2:
         log.warning(
-            f"depth/Nz = {depth / scan.Nz:.7f} < wavelength/2 = {wvln / 2:.7f}. "
-            f"Consider either increasing scan.Nz to {int(np.ceil(depth / (wvln / 2)))} or more, or "
-            "increasing scan.pixels_per_wavelength to 2 or more."
+            f"depth/grid_size_z = {depth / scan.grid_size_z:.7f} < wavelength/2 = {wvln / 2:.7f}. "
+            f"Consider either increasing scan.grid_size_z to {int(np.ceil(depth / (wvln / 2)))} "
+            "or more, or increasing scan.pixels_per_wavelength to 2 or more."
         )
 
 
-def cartesian_pixel_grid(xlims, zlims, Nx=None, Nz=None, dx=None, dz=None):
+def cartesian_pixel_grid(xlims, zlims, grid_size_x=None, grid_size_z=None, dx=None, dz=None):
     """Generate a Cartesian pixel grid based on input parameters.
 
     Args:
         xlims (tuple): Azimuthal limits of pixel grid ([xmin, xmax])
         zlims (tuple): Depth limits of pixel grid ([zmin, zmax])
-        Nx (int): Number of azimuthal pixels, overrides dx and dz parameters
-        Nz (int): Number of depth pixels, overrides dx and dz parameters
+        grid_size_x (int): Number of azimuthal pixels, overrides dx and dz parameters
+        grid_size_z (int): Number of depth pixels, overrides dx and dz parameters
         dx (float): Pixel spacing in azimuth
         dz (float): Pixel spacing in depth
 
     Raises:
-        ValueError: Either Nx and Nz or dx and dz must be defined.
+        ValueError: Either grid_size_x and grid_size_z or dx and dz must be defined.
 
     Returns:
-        grid (np.ndarray): Pixel grid of size (nz, nx, 3) in
+        grid (np.ndarray): Pixel grid of size (grid_size_z, nx, 3) in
             Cartesian coordinates (x, y, z)
     """
-    assert (bool(Nx) and bool(Nz)) ^ (bool(dx) and bool(dz)), (
-        "Either Nx and Nz or dx and dz must be defined."
+    assert (bool(grid_size_x) and bool(grid_size_z)) ^ (bool(dx) and bool(dz)), (
+        "Either grid_size_x and grid_size_z or dx and dz must be defined."
     )
 
     # Determine the grid spacing
-    if Nx is not None and Nz is not None:
-        x = np.linspace(xlims[0], xlims[1] + eps, Nx)
-        z = np.linspace(zlims[0], zlims[1] + eps, Nz)
+    if grid_size_x is not None and grid_size_z is not None:
+        x = np.linspace(xlims[0], xlims[1] + eps, grid_size_x)
+        z = np.linspace(zlims[0], zlims[1] + eps, grid_size_z)
     elif dx is not None and dz is not None:
         sign = np.sign(xlims[1] - xlims[0])
         x = np.arange(xlims[0], xlims[1] + eps, sign * dx)
         z = np.arange(zlims[0], zlims[1] + eps, sign * dz)
     else:
-        raise ValueError("Either Nx and Nz or dx and dz must be defined.")
+        raise ValueError("Either grid_size_x and grid_size_z or dx and dz must be defined.")
 
     # Create the pixel grid
     z_grid, x_grid = np.meshgrid(z, x, indexing="ij")
@@ -102,29 +102,30 @@ def radial_pixel_grid(rlims, dr, oris, dirs):
     return grid
 
 
-def polar_pixel_grid(polar_limits, zlims, Nz: int, Nr: int):
+def polar_pixel_grid(polar_limits, zlims, num_radial_pixels: int, num_polar_pixels: int):
     """Generate a polar grid.
 
     Uses radial_pixel_grid but based on parameters that are present in the scan class.
 
     Args:
-        polar_limits (tuple): Azimuthal limits of pixel grid ([azimuth_min, azimuth_max])
+        polar_limits (tuple): Polar limits of pixel grid ([polar_min, polar_max])
         zlims (tuple): Depth limits of pixel grid ([zmin, zmax])
-        Nz (int, optional): Number of depth pixels.
-        Nr (int, optional): Number of azimuthal pixels.
+        num_radial_pixels (int, optional): Number of depth pixels.
+        num_polar_pixels (int, optional): Number of polar pixels.
 
     Returns:
-        grid (np.ndarray): Pixel grid of size (Nz, Nr, 3) in Cartesian coordinates (x, y, z)
+        grid (np.ndarray): Pixel grid of size (num_radial_pixels, num_polar_pixels, 3)
+        in Cartesian coordinates (x, y, z)
     """
     assert len(polar_limits) == 2, "polar_limits must be a tuple of length 2."
     assert len(zlims) == 2, "zlims must be a tuple of length 2."
 
-    dr = (zlims[1] - zlims[0]) / Nz
+    dr = (zlims[1] - zlims[0]) / num_radial_pixels
 
     oris = np.array([0, 0, 0])
-    oris = np.tile(oris, (Nr, 1))
-    dirs_az = np.linspace(*polar_limits, Nr)
+    oris = np.tile(oris, (num_polar_pixels, 1))
+    dirs_az = np.linspace(*polar_limits, num_polar_pixels)
 
-    dirs_el = np.zeros(Nr)
+    dirs_el = np.zeros(num_polar_pixels)
     dirs = np.vstack((dirs_az, dirs_el)).T
     return radial_pixel_grid(zlims, dr, oris, dirs).transpose(1, 0, 2)

zea/config.py

@@ -47,8 +47,9 @@ import yaml
 from huggingface_hub import hf_hub_download
 
 from zea import log
+from zea.data.preset_utils import HF_PREFIX, _hf_resolve_path
 from zea.internal.config.validation import config_schema
-from zea.internal.core import object_to_tensor
+from zea.internal.core import dict_to_tensor
 
 
 class Config(dict):
@@ -430,8 +431,22 @@ class Config(dict):
                         v._recursive_setattr(set_key, set_value)
 
     @classmethod
-    def from_yaml(cls, path, **kwargs):
-        """Load config object from yaml file"""
+    def from_path(cls, path, **kwargs):
+        """Load config object from a file path.
+
+        Args:
+            path (str or Path): The path to the config file.
+                Can be a string or a Path object. Additionally can be a string with
+                the prefix 'hf://', in which case it will be resolved to a
+                huggingface path.
+
+        Returns:
+            Config: config object.
+        """
+        if str(path).startswith(HF_PREFIX):
+            path = _hf_resolve_path(str(path))
+        if isinstance(path, str):
+            path = Path(path)
         return _load_config_from_yaml(path, config_class=cls, **kwargs)
 
     @classmethod
@@ -460,9 +475,14 @@ class Config(dict):
         local_path = hf_hub_download(repo_id, path, **kwargs)
         return _load_config_from_yaml(local_path, config_class=cls)
 
-    def to_tensor(self):
+    @classmethod
+    def from_yaml(cls, path, **kwargs):
+        """Load config object from yaml file."""
+        return cls.from_path(path, **kwargs)
+
+    def to_tensor(self, keep_as_is=None):
         """Convert the attributes in the object to keras tensors"""
-        return object_to_tensor(self)
+        return dict_to_tensor(self.serialize(), keep_as_is=keep_as_is)
 
 
 def check_config(config: Union[dict, Config], verbose: bool = False):

zea/data/data_format.py

@@ -42,8 +42,8 @@ def generate_example_dataset(
     sound_speed=1540,
     center_frequency=7e6,
     sampling_frequency=40e6,
-    n_z=512,
-    n_x=512,
+    grid_size_z=512,
+    grid_size_x=512,
 ):
     """Generates an example dataset that contains all the necessary fields.
     Note: This dataset does not contain actual data, but is filled with random
@@ -60,7 +60,7 @@ def generate_example_dataset(
     # creating some fake raw and image data
     raw_data = np.ones((n_frames, n_tx, n_ax, n_el, n_ch))
     # image data is in dB
-    image = np.ones((n_frames, n_z, n_x)) * -40
+    image = np.ones((n_frames, grid_size_z, grid_size_x)) * -40
 
     # creating some fake scan parameters
     t0_delays = np.zeros((n_tx, n_el), dtype=np.float32)
@@ -82,8 +82,8 @@ def generate_example_dataset(
 
     if add_optional_dtypes:
         aligned_data = np.ones((n_frames, n_tx, n_ax, n_el, n_ch))
-        envelope_data = np.ones((n_frames, n_z, n_x))
-        beamformed_data = np.ones((n_frames, n_z, n_x, n_ch))
+        envelope_data = np.ones((n_frames, grid_size_z, grid_size_x))
+        beamformed_data = np.ones((n_frames, grid_size_z, grid_size_x, n_ch))
         image_sc = np.ones_like(image)
     else:
         aligned_data = None
@@ -123,10 +123,11 @@ def validate_input_data(raw_data, aligned_data, envelope_data, beamformed_data,
         aligned_data (np.ndarray): The aligned data of the ultrasound measurement of
             shape (n_frames, n_tx, n_ax, n_el, n_ch).
         envelope_data (np.ndarray): The envelope data of the ultrasound measurement of
-            shape (n_frames, n_z, n_x).
+            shape (n_frames, grid_size_z, grid_size_x).
         beamformed_data (np.ndarray): The beamformed data of the ultrasound measurement of
-            shape (n_frames, n_z, n_x).
-        image (np.ndarray): The ultrasound images to be saved of shape (n_frames, n_z, n_x).
+            shape (n_frames, grid_size_z, grid_size_x).
+        image (np.ndarray): The ultrasound images to be saved
+            of shape (n_frames, grid_size_z, grid_size_x).
         image_sc (np.ndarray): The scan converted ultrasound images to be saved
             of shape (n_frames, output_size_z, output_size_x).
     """
@@ -254,7 +255,7 @@ def _write_datasets(
         group_name=data_group_name,
         name="envelope_data",
         data=_convert_datatype(envelope_data),
-        description="The envelope_data of shape (n_frames, n_z, n_x).",
+        description="The envelope_data of shape (n_frames, grid_size_z, grid_size_x).",
         unit="unitless",
     )
 
@@ -262,7 +263,7 @@ def _write_datasets(
         group_name=data_group_name,
         name="beamformed_data",
         data=_convert_datatype(beamformed_data),
-        description="The beamformed_data of shape (n_frames, n_z, n_x).",
+        description="The beamformed_data of shape (n_frames, grid_size_z, grid_size_x).",
         unit="unitless",
     )
 
@@ -271,7 +272,7 @@ def _write_datasets(
         name="image",
         data=_convert_datatype(image),
         unit="unitless",
-        description="The images of shape [n_frames, n_z, n_x]",
+        description="The images of shape [n_frames, grid_size_z, grid_size_x]",
     )
 
     _add_dataset(
@@ -467,32 +468,34 @@ def _write_datasets(
             ),
             unit="-",
         )
-        n_waveforms = len(waveforms_one_way)
-        for n, waveform_1way, waveform_2way in zip(
-            range(n_waveforms), waveforms_one_way, waveforms_two_way
-        ):
-            _add_dataset(
-                group_name=scan_group_name + "/waveforms_one_way",
-                name=f"waveform_{str(n).zfill(3)}",
-                data=waveform_1way,
-                description=(
-                    "One-way waveform as simulated by the Verasonics system, "
-                    "sampled at 250MHz. This is the waveform after being filtered "
-                    "by the tranducer bandwidth once."
-                ),
-                unit="V",
-            )
-            _add_dataset(
-                group_name=scan_group_name + "/waveforms_two_way",
-                name=f"waveform_{str(n).zfill(3)}",
-                data=waveform_2way,
-                description=(
-                    "Two-way waveform as simulated by the Verasonics system, "
-                    "sampled at 250MHz. This is the waveform after being filtered "
-                    "by the tranducer bandwidth twice."
-                ),
-                unit="V",
-            )
+
+        if waveforms_one_way is not None:
+            for n in range(len(waveforms_one_way)):
+                _add_dataset(
+                    group_name=scan_group_name + "/waveforms_one_way",
+                    name=f"waveform_{str(n).zfill(3)}",
+                    data=waveforms_one_way[n],
+                    description=(
+                        "One-way waveform as simulated by the Verasonics system, "
+                        "sampled at 250MHz. This is the waveform after being filtered "
+                        "by the tranducer bandwidth once."
+                    ),
+                    unit="V",
+                )
+
+        if waveforms_two_way is not None:
+            for n in range(len(waveforms_two_way)):
+                _add_dataset(
+                    group_name=scan_group_name + "/waveforms_two_way",
+                    name=f"waveform_{str(n).zfill(3)}",
+                    data=waveforms_two_way[n],
+                    description=(
+                        "Two-way waveform as simulated by the Verasonics system, "
+                        "sampled at 250MHz. This is the waveform after being filtered "
+                        "by the tranducer bandwidth twice."
+                    ),
+                    unit="V",
+                )
 
     # Add additional elements
     if additional_elements is not None:
@@ -546,10 +549,11 @@ def generate_zea_dataset(
         aligned_data (np.ndarray): The aligned data of the ultrasound measurement of
             shape (n_frames, n_tx, n_ax, n_el, n_ch).
         envelope_data (np.ndarray): The envelope data of the ultrasound measurement of
-            shape (n_frames, n_z, n_x).
+            shape (n_frames, grid_size_z, grid_size_x).
         beamformed_data (np.ndarray): The beamformed data of the ultrasound measurement of
-            shape (n_frames, n_z, n_x, n_ch).
-        image (np.ndarray): The ultrasound images to be saved of shape (n_frames, n_z, n_x).
+            shape (n_frames, grid_size_z, grid_size_x, n_ch).
+        image (np.ndarray): The ultrasound images to be saved
+            of shape (n_frames, grid_size_z, grid_size_x).
         image_sc (np.ndarray): The scan converted ultrasound images to be saved
             of shape (n_frames, output_size_z, output_size_x).
         probe_geometry (np.ndarray): The probe geometry of shape (n_el, 3).

zea/data/file.py

@@ -366,9 +366,7 @@ class File(h5py.File):
         """
         file_scan_parameters = self.get_parameters(event)
 
-        probe_parameters = reduce_to_signature(
-            Probe.from_name("generic").__init__, file_scan_parameters
-        )
+        probe_parameters = reduce_to_signature(Probe.__init__, file_scan_parameters)
         return probe_parameters
 
     def probe(self, event=None) -> Probe:
@@ -387,21 +385,8 @@ class File(h5py.File):
         Returns:
             Probe: The probe object.
         """
-        probe_parameters = self.get_probe_parameters(event)
-        if self.probe_name == "generic":
-            return Probe.from_name(self.probe_name, **probe_parameters)
-        else:
-            probe = Probe.from_name(self.probe_name)
-
-            probe_geometry = probe_parameters.get("probe_geometry", None)
-            if not np.allclose(probe_geometry, probe.probe_geometry):
-                probe.probe_geometry = probe_geometry
-                log.warning(
-                    "The probe geometry in the data file does not "
-                    "match the probe geometry of the probe. The probe "
-                    "geometry has been updated to match the data file."
-                )
-        return probe
+        probe_parameters_file = self.get_probe_parameters(event)
+        return Probe.from_parameters(self.probe_name, probe_parameters_file)
 
     def recursively_load_dict_contents_from_group(self, path: str, squeeze: bool = False) -> dict:
         """Load dict from contents of group