zea 0.0.7__py3-none-any.whl → 0.0.9__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.7"
+__version__ = "0.0.9"
 
 
 def _bootstrap_backend():
@@ -89,12 +89,12 @@ from . import (
     beamform,
     data,
     display,
+    func,
     io_lib,
-    keras_ops,
     metrics,
     models,
+    ops,
     simulator,
-    tensor_ops,
     utils,
     visualize,
 )

zea/agent/masks.py

@@ -9,8 +9,8 @@ from typing import List
 import keras
 from keras import ops
 
-from zea import tensor_ops
 from zea.agent.gumbel import hard_straight_through
+from zea.func.tensor import nonzero
 
 _DEFAULT_DTYPE = "bool"
 
@@ -56,7 +56,7 @@ def k_hot_to_indices(selected_lines, n_actions: int, fill_value=-1):
 
     # Find nonzero indices for each frame
     def get_nonzero(row):
-        return tensor_ops.nonzero(row > 0, size=n_actions, fill_value=fill_value)[0]
+        return nonzero(row > 0, size=n_actions, fill_value=fill_value)[0]
 
     indices = ops.vectorized_map(get_nonzero, selected_lines)
     return indices

zea/agent/selection.py

@@ -16,9 +16,9 @@ from typing import Callable
 import keras
 from keras import ops
 
-from zea import tensor_ops
 from zea.agent import masks
 from zea.backend.autograd import AutoGrad
+from zea.func import tensor
 from zea.internal.registry import action_selection_registry
 
 
@@ -462,7 +462,7 @@ class CovarianceSamplingLines(LinesActionModel):
         particles = ops.reshape(particles, shape)
 
         # [batch_size, rows * stack_n_cols, n_possible_actions, n_possible_actions]
-        cov_matrix = tensor_ops.batch_cov(particles)
+        cov_matrix = tensor.batch_cov(particles)
 
         # Sum over the row dimension [batch_size, n_possible_actions, n_possible_actions]
         cov_matrix = ops.sum(cov_matrix, axis=1)
@@ -477,7 +477,7 @@ class CovarianceSamplingLines(LinesActionModel):
         # Subsample the covariance matrix with random lines
         def subsample_with_mask(mask):
             """Subsample the covariance matrix with a single mask."""
-            subsampled_cov_matrix = tensor_ops.boolean_mask(
+            subsampled_cov_matrix = tensor.boolean_mask(
                 cov_matrix, mask, size=batch_size * self.n_actions**2
             )
             return ops.reshape(subsampled_cov_matrix, [batch_size, self.n_actions, self.n_actions])

zea/backend/__init__.py

@@ -131,7 +131,7 @@ class on_device:
         .. code-block:: python
 
             with zea.backend.on_device("gpu:3"):
-                pipeline = zea.Pipeline([zea.keras_ops.Abs()])
+                pipeline = zea.Pipeline([zea.ops.Abs()])
                 output = pipeline(data=keras.random.normal((10, 10)))  # output is on "cuda:3"
     """
 

zea/backend/tensorflow/dataloader.py

@@ -12,8 +12,8 @@ from keras.src.trainers.data_adapters import TFDatasetAdapter
 
 from zea.data.dataloader import H5Generator
 from zea.data.layers import Resizer
+from zea.func.tensor import translate
 from zea.internal.utils import find_methods_with_return_type
-from zea.tensor_ops import translate
 
 METHODS_THAT_RETURN_DATASET = find_methods_with_return_type(tf.data.Dataset, "DatasetV2")
 
@@ -155,10 +155,6 @@ def make_dataloader(
     Mimics the native TF function ``tf.keras.utils.image_dataset_from_directory``
     but for .hdf5 files.
 
-    Saves a dataset_info.yaml file in the directory with information about the dataset.
-    This file is used to load the dataset later on, which speeds up the initial loading
-    of the dataset for very large datasets.
-
     Does the following in order to load a dataset:
 
         - Find all .hdf5 files in the director(ies)

zea/beamform/beamformer.py

@@ -5,7 +5,7 @@ import numpy as np
 from keras import ops
 
 from zea.beamform.lens_correction import calculate_lens_corrected_delays
-from zea.tensor_ops import vmap
+from zea.func.tensor import vmap
 
 
 def fnum_window_fn_rect(normalized_angle):
@@ -379,7 +379,7 @@ def complex_rotate(iq, theta):
 
         .. math::
 
-            x(t + \\Delta t) &= I'(t) \\cos(\\omega_c (t + \\Delta t)) 
+            x(t + \\Delta t) &= I'(t) \\cos(\\omega_c (t + \\Delta t))
             - Q'(t) \\sin(\\omega_c (t + \\Delta t))\\\\
             &=  \\overbrace{(I'(t)\\cos(\\theta)
             - Q'(t)\\sin(\\theta) )}^{I_\\Delta(t)} \\cos(\\omega_c t)\\\\
@@ -452,6 +452,8 @@ def distance_Tx_generic(
             `(n_el,)`.
         probe_geometry (ops.Tensor): The positions of the transducer elements of shape
             `(n_el, 3)`.
+        focus_distance (float): The focus distance in meters.
+        polar_angle (float): The polar angle in radians.
         sound_speed (float): The speed of sound in m/s. Defaults to 1540.
 
     Returns:

zea/beamform/pfield.py

@@ -24,8 +24,8 @@ import numpy as np
 from keras import ops
 
 from zea import log
+from zea.func.tensor import sinc
 from zea.internal.cache import cache_output
-from zea.tensor_ops import sinc
 
 
 def _abs_sinc(x):
@@ -101,7 +101,7 @@ def compute_pfield(
     # array params
     probe_geometry = ops.convert_to_tensor(probe_geometry, dtype="float32")
 
-    pitch = probe_geometry[1, 0] - probe_geometry[0, 0]  # element pitch
+    pitch = ops.abs(probe_geometry[1, 0] - probe_geometry[0, 0])  # element pitch
 
     kerf = 0.1 * pitch  # for now this is hardcoded
     element_width = pitch - kerf

zea/beamform/pixelgrid.py

@@ -45,7 +45,7 @@ def cartesian_pixel_grid(xlims, zlims, grid_size_x=None, grid_size_z=None, dx=No
         ValueError: Either grid_size_x and grid_size_z or dx and dz must be defined.
 
     Returns:
-        grid (np.ndarray): Pixel grid of size (grid_size_z, nx, 3) in
+        grid (np.ndarray): Pixel grid of size (grid_size_z, grid_size_x, 3) in
             Cartesian coordinates (x, y, z)
     """
     assert (bool(grid_size_x) and bool(grid_size_z)) ^ (bool(dx) and bool(dz)), (

zea/data/__init__.py

@@ -38,15 +38,6 @@ Examples usage
     ...     files.append(file)  # process each file as needed
     >>> dataset.close()
 
-Subpackage layout
------------------
-
-- ``file.py``: Implements :class:`zea.File` and related file utilities.
-- ``datasets.py``: Implements :class:`zea.Dataset` and folder management.
-- ``dataloader.py``: Data loading utilities for batching and shuffling.
-- ``data_format.py``: Data validation and example dataset generation.
-- ``convert/``: Data conversion tools (e.g., from external formats).
-
 """  # noqa: E501
 
 from .convert.camus import sitk_load

zea/data/augmentations.py

@@ -4,7 +4,7 @@ import keras
 import numpy as np
 from keras import layers, ops
 
-from zea.tensor_ops import is_jax_prng_key, split_seed
+from zea.func.tensor import is_jax_prng_key, split_seed
 
 
 class RandomCircleInclusion(layers.Layer):
@@ -30,7 +30,7 @@ class RandomCircleInclusion(layers.Layer):
 
     def __init__(
         self,
-        radius: int,
+        radius: int | tuple[int, int],
         fill_value: float = 1.0,
         circle_axes: tuple[int, int] = (1, 2),
         with_batch_dim=True,
@@ -38,25 +38,70 @@ class RandomCircleInclusion(layers.Layer):
         recovery_threshold=0.1,
         randomize_location_across_batch=True,
         seed=None,
+        width_range: tuple[int, int] = None,
+        height_range: tuple[int, int] = None,
         **kwargs,
     ):
         """
         Initialize RandomCircleInclusion.
 
         Args:
-            radius (int): Radius of the circle to include.
+            radius (int or tuple[int, int]): Radius of the circle/ellipse to include.
             fill_value (float): Value to fill inside the circle.
-            circle_axes (tuple[int, int]): Axes along which to draw the circle (height, width).
+            circle_axes (tuple[int, int]): Axes along which to draw the circle
+                (height, width).
             with_batch_dim (bool): Whether input has a batch dimension.
             return_centers (bool): Whether to return circle centers along with images.
             recovery_threshold (float): Threshold for considering a pixel as recovered.
-            randomize_location_across_batch (bool): If True, randomize circle location
-                per batch element.
+            randomize_location_across_batch (bool): If True (and with_batch_dim=True),
+                each batch element gets a different random center. If False, all batch
+                elements share the same center.
             seed (Any): Optional random seed for reproducibility.
+            width_range (tuple[int, int], optional): Range (min, max) for circle
+                center x (width axis).
+            height_range (tuple[int, int], optional): Range (min, max) for circle
+                center y (height axis).
             **kwargs: Additional keyword arguments for the parent Layer.
+
+        Example:
+            .. doctest::
+
+                >>> from zea.data.augmentations import RandomCircleInclusion
+                >>> from keras import ops
+
+                >>> layer = RandomCircleInclusion(
+                ...     radius=5,
+                ...     circle_axes=(1, 2),
+                ...     with_batch_dim=True,
+                ... )
+                >>> image = ops.zeros((1, 28, 28), dtype="float32")
+                >>> out = layer(image)  # doctest: +SKIP
+
         """
         super().__init__(**kwargs)
-        self.radius = radius
+
+        # Validate randomize_location_across_batch only makes sense with batch dim
+        if not with_batch_dim and not randomize_location_across_batch:
+            raise ValueError(
+                "randomize_location_across_batch=False is only applicable when "
+                "with_batch_dim=True. When with_batch_dim=False, there is no batch "
+                "to randomize across."
+            )
+        # Convert radius to tuple if int, else validate tuple
+        if isinstance(radius, int):
+            if radius <= 0:
+                raise ValueError(f"radius must be a positive integer, got {radius}.")
+            self.radius = (radius, radius)
+        elif isinstance(radius, tuple) and len(radius) == 2:
+            rx, ry = radius
+            if not all(isinstance(r, int) for r in (rx, ry)):
+                raise TypeError(f"radius tuple must contain two integers, got {radius!r}.")
+            if rx <= 0 or ry <= 0:
+                raise ValueError(f"radius components must be positive, got {radius!r}.")
+            self.radius = (rx, ry)
+        else:
+            raise TypeError("radius must be an int or a tuple of two ints")
+
         self.fill_value = fill_value
         self.circle_axes = circle_axes
         self.with_batch_dim = with_batch_dim
@@ -64,6 +109,8 @@ class RandomCircleInclusion(layers.Layer):
         self.recovery_threshold = recovery_threshold
         self.randomize_location_across_batch = randomize_location_across_batch
         self.seed = seed
+        self.width_range = width_range
+        self.height_range = height_range
         self._axis1 = None
         self._axis2 = None
         self._perm = None
@@ -116,6 +163,43 @@ class RandomCircleInclusion(layers.Layer):
         self._static_w = int(permuted_shape[-1])
         self._static_shape = tuple(permuted_shape)
 
+        # Validate that ellipse can fit within image bounds
+        rx, ry = self.radius
+        min_required_width = 2 * rx + 1
+        min_required_height = 2 * ry + 1
+
+        if self._static_w < min_required_width:
+            raise ValueError(
+                f"Image width ({self._static_w}) is too small for radius {rx}. "
+                f"Minimum required width: {min_required_width}"
+            )
+        if self._static_h < min_required_height:
+            raise ValueError(
+                f"Image height ({self._static_h}) is too small for radius {ry}. "
+                f"Minimum required height: {min_required_height}"
+            )
+
+        # Validate width_range and height_range if provided
+        if self.width_range is not None:
+            min_x, max_x = self.width_range
+            if min_x >= max_x:
+                raise ValueError(f"width_range must have min < max, got {self.width_range}")
+            if min_x < rx or max_x > self._static_w - rx:
+                raise ValueError(
+                    f"width_range {self.width_range} would place circle outside image bounds. "
+                    f"Valid range: [{rx}, {self._static_w - rx})"
+                )
+
+        if self.height_range is not None:
+            min_y, max_y = self.height_range
+            if min_y >= max_y:
+                raise ValueError(f"height_range must have min < max, got {self.height_range}")
+            if min_y < ry or max_y > self._static_h - ry:
+                raise ValueError(
+                    f"height_range {self.height_range} would place circle outside image bounds. "
+                    f"Valid range: [{ry}, {self._static_h - ry})"
+                )
+
         super().build(input_shape)
 
     def compute_output_shape(self, input_shape):
@@ -165,7 +249,7 @@ class RandomCircleInclusion(layers.Layer):
             centers (Tensor): Tensor of shape (batch, 2) with circle centers.
             h (int): Height of the image.
             w (int): Width of the image.
-            radius (int): Radius of the circle.
+            radius (tuple[int, int]): Radii of the ellipse (rx, ry).
             dtype (str or dtype): Data type for the mask.
 
         Returns:
@@ -176,12 +260,12 @@ class RandomCircleInclusion(layers.Layer):
         Y, X = ops.meshgrid(Y, X, indexing="ij")
         Y = ops.expand_dims(Y, 0)  # (1, h, w)
         X = ops.expand_dims(X, 0)  # (1, h, w)
-        # cx = ops.cast(centers[:, 0], "float32")[:, None, None]
-        # cy = ops.cast(centers[:, 1], "float32")[:, None, None]
         cx = centers[:, 0][:, None, None]
         cy = centers[:, 1][:, None, None]
-        dist2 = (X - cx) ** 2 + (Y - cy) ** 2
-        mask = ops.cast(dist2 <= radius**2, dtype)
+        rx, ry = radius
+        # Ellipse equation: ((X-cx)/rx)^2 + ((Y-cy)/ry)^2 <= 1
+        dist = ((X - cx) / rx) ** 2 + ((Y - cy) / ry) ** 2
+        mask = ops.cast(dist <= 1, dtype)
         return mask
 
     def call(self, x, seed=None):
@@ -197,9 +281,17 @@ class RandomCircleInclusion(layers.Layer):
                 centers if return_centers is True.
         """
         if keras.backend.backend() == "jax" and not is_jax_prng_key(seed):
-            raise NotImplementedError(
-                "jax.random.key() is not supported, please use jax.random.PRNGKey()"
-            )
+            if isinstance(seed, keras.random.SeedGenerator):
+                raise ValueError(
+                    "When using JAX backend, please provide a jax.random.PRNGKey as seed, "
+                    "instead of keras.random.SeedGenerator."
+                )
+            else:
+                raise TypeError(
+                    f"When using JAX backend, seed must be a JAX PRNG key (created with "
+                    f"jax.random.PRNGKey()), but got {type(seed)}. Note: jax.random.key() "
+                    f"keys are not currently supported."
+                )
         seed = seed if seed is not None else self.seed
 
         if self.with_batch_dim:
@@ -209,22 +301,33 @@ class RandomCircleInclusion(layers.Layer):
                     imgs, centers = ops.map(lambda arg: self._call(arg, seed), x)
                 else:
                     raise NotImplementedError(
-                        "You cannot fix circle locations across while using"
+                        "You cannot fix circle locations across batch while using "
                         + "RandomCircleInclusion as a dataset augmentation, "
                         + "since samples in a batch are handled independently."
                     )
             else:
+                batch_size = ops.shape(x)[0]
                 if self.randomize_location_across_batch:
-                    batch_size = ops.shape(x)[0]
                     seeds = split_seed(seed, batch_size)
-                    if all(seed is seeds[0] for seed in seeds):
+                    if all(s is seeds[0] for s in seeds):
                         imgs, centers = ops.map(lambda arg: self._call(arg, seeds[0]), x)
                     else:
                         imgs, centers = ops.map(
                             lambda args: self._call(args[0], args[1]), (x, seeds)
                         )
                 else:
-                    imgs, centers = ops.map(lambda arg: self._call(arg, seed), x)
+                    # Generate one random center that will be used for all batch elements
+                    img0, center0 = self._call(x[0], seed)
+
+                    # Apply the same center to all batch elements
+                    imgs_list, centers_list = [img0], [center0]
+                    for i in range(1, batch_size):
+                        img_aug, center_out = self._call_with_fixed_center(x[i], center0)
+                        imgs_list.append(img_aug)
+                        centers_list.append(center_out)
+
+                    imgs = ops.stack(imgs_list, axis=0)
+                    centers = ops.stack(centers_list, axis=0)
         else:
             imgs, centers = self._call(x, seed)
 
@@ -248,17 +351,28 @@ class RandomCircleInclusion(layers.Layer):
         flat, flat_batch_size, h, w = self._flatten_batch_and_other_dims(x)
 
         def _draw_circle_2d(img2d):
+            rx, ry = self.radius
+            # Determine allowed ranges for center
+            if self.width_range is not None:
+                min_x, max_x = self.width_range
+            else:
+                min_x, max_x = rx, w - rx
+            if self.height_range is not None:
+                min_y, max_y = self.height_range
+            else:
+                min_y, max_y = ry, h - ry
+            # Ensure the ellipse fits within the allowed region
             cx = ops.cast(
-                keras.random.uniform((), self.radius, w - self.radius, seed=seed),
+                keras.random.uniform((), min_x, max_x, seed=seed),
                 "int32",
             )
             new_seed, _ = split_seed(seed, 2)  # ensure that cx and cy are independent
             cy = ops.cast(
-                keras.random.uniform((), self.radius, h - self.radius, seed=new_seed),
+                keras.random.uniform((), min_y, max_y, seed=new_seed),
                 "int32",
             )
             mask = self._make_circle_mask(
-                ops.stack([cx, cy])[None, :], h, w, self.radius, img2d.dtype
+                ops.stack([cx, cy])[None, :], h, w, (rx, ry), img2d.dtype
             )[0]
             img_aug = img2d * (1 - mask) + self.fill_value * mask
             center = ops.stack([cx, cy])
@@ -271,6 +385,67 @@ class RandomCircleInclusion(layers.Layer):
         centers = ops.reshape(centers, centers_shape)
         return (aug_imgs, centers)
 
+    def _apply_circle_mask(self, flat, center, h, w):
+        """Apply circle mask to flattened image data.
+
+        Args:
+            flat (Tensor): Flattened image data of shape (flat_batch, h, w).
+            center (Tensor): Center coordinates, either (2,) or (flat_batch, 2).
+            h (int): Height of images.
+            w (int): Width of images.
+
+        Returns:
+            Tensor: Augmented images with circle applied.
+        """
+        rx, ry = self.radius
+
+        # Ensure center has batch dimension for broadcasting
+        if len(center.shape) == 1:
+            # Single center (2,) -> broadcast to all slices
+            center_batched = ops.tile(ops.reshape(center, [1, 2]), [flat.shape[0], 1])
+        else:
+            # Already batched (flat_batch, 2)
+            center_batched = center
+
+        # Create masks for all slices using vectorized_map or broadcasting
+        masks = self._make_circle_mask(center_batched, h, w, (rx, ry), flat.dtype)
+
+        # Apply masks
+        aug_imgs = flat * (1 - masks) + self.fill_value * masks
+        return aug_imgs
+
+    def _call_with_fixed_center(self, x, fixed_center):
+        """Apply augmentation using a pre-determined center.
+
+        Args:
+            x (Tensor): Input image tensor.
+            fixed_center (Tensor): Pre-determined center coordinates, either (2,)
+                for a single center or (flat_batch, 2) for per-slice centers.
+
+        Returns:
+            tuple: (augmented image, center coordinates).
+        """
+        x = self._permute_axes_to_circle_last(x)
+        flat, flat_batch_size, h, w = self._flatten_batch_and_other_dims(x)
+
+        # Apply circle mask with fixed center (handles both single and batched centers)
+        aug_imgs = self._apply_circle_mask(flat, fixed_center, h, w)
+        aug_imgs = ops.reshape(aug_imgs, x.shape)
+        aug_imgs = ops.transpose(aug_imgs, axes=self._inv_perm)
+
+        # Return centers matching the expected shape
+        if len(fixed_center.shape) == 1:
+            # Single center (2,) -> broadcast to match flat_batch_size
+            if flat_batch_size == 1:
+                centers = fixed_center
+            else:
+                centers = ops.tile(ops.reshape(fixed_center, [1, 2]), [flat_batch_size, 1])
+        else:
+            # Already batched centers (flat_batch, 2)
+            centers = fixed_center
+
+        return (aug_imgs, centers)
+
     def get_config(self):
         """
         Get layer configuration for serialization.
@@ -285,6 +460,8 @@ class RandomCircleInclusion(layers.Layer):
                 "fill_value": self.fill_value,
                 "circle_axes": self.circle_axes,
                 "return_centers": self.return_centers,
+                "width_range": self.width_range,
+                "height_range": self.height_range,
             }
         )
         return cfg
@@ -293,7 +470,8 @@ class RandomCircleInclusion(layers.Layer):
         self, images, centers, recovery_threshold, fill_value=None
     ):
         """
-        Evaluate the percentage of the true circle that has been recovered in the images.
+        Evaluate the percentage of the true circle that has been recovered in the images,
+        and return a mask of the detected part of the circle.
 
         Args:
             images (Tensor): Tensor of images (any shape, with circle axes as specified).
@@ -302,8 +480,12 @@ class RandomCircleInclusion(layers.Layer):
             fill_value (float, optional): Optionally override fill_value for cases
                 where image range has changed.
 
-        Returns:
-            Tensor: Percentage recovered for each circle (shape: [num_circles]).
+         Returns:
+            Tuple[Tensor, Tensor]:
+                - percent_recovered: [batch] - average recovery percentage per batch element,
+                  averaged across all non-batch dimensions (e.g., frames, slices)
+                - recovered_masks: [batch, flat_batch, h, w] or [batch, h, w] or [flat_batch, h, w]
+                   depending on input shape - binary mask of detected circle regions
         """
         fill_value = fill_value or self.fill_value
 
@@ -318,12 +500,23 @@ class RandomCircleInclusion(layers.Layer):
             recovered_sum = ops.sum(recovered, axis=[1, 2])
             mask_sum = ops.sum(mask, axis=[1, 2])
             percent_recovered = recovered_sum / (mask_sum + 1e-8)
-            return percent_recovered
+            # recovered_mask: binary mask of detected part of the circle
+            recovered_mask = ops.cast(recovered > 0, flat_image.dtype)
+            return percent_recovered, recovered_mask
 
         if self.with_batch_dim:
-            return ops.vectorized_map(
+            results = ops.vectorized_map(
                 lambda args: _evaluate_recovered_circle_accuracy(args[0], args[1]),
                 (images, centers),
-            )[..., 0]
+            )
+            percent_recovered, recovered_masks = results
+            # If there are multiple circles per batch element (e.g., multiple frames/slices),
+            # take the mean across all non-batch dimensions to get one value per batch element
+            if len(percent_recovered.shape) > 1:
+                # Average over all axes except the batch dimension (axis 0)
+                axes_to_reduce = tuple(range(1, len(percent_recovered.shape)))
+                percent_recovered = ops.mean(percent_recovered, axis=axes_to_reduce)
+            return percent_recovered, recovered_masks
         else:
-            return _evaluate_recovered_circle_accuracy(images, centers)
+            percent_recovered, recovered_mask = _evaluate_recovered_circle_accuracy(images, centers)
+            return percent_recovered, recovered_mask

zea/data/convert/__init__.py

@@ -1,6 +1 @@
-"""All functions in this module are used to convert different datasets to the zea format."""
-
-from .camus import convert_camus
-from .images import convert_image_dataset
-from .matlab import zea_from_matlab_raw
-from .picmus import convert_picmus
+"""Data conversion of datasets to the ``zea`` data format."""