zea 0.0.8__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.8"
+__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")
 

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/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):

zea/data/convert/__main__.py

@@ -1,86 +1,127 @@
 import argparse
 
 
-def get_parser():
-    """
-    Build and parse command-line arguments for converting raw datasets to a zea dataset.
-
-    Returns:
-        argparse.Namespace: Parsed arguments with the following attributes:
-            dataset (str): One of "echonet", "echonetlvh", "camus", "picmus", "verasonics".
-            src (str): Source folder path.
-            dst (str): Destination folder path.
-            split_path (str|None): Optional path to a split.yaml to copy dataset splits.
-            no_hyperthreading (bool): Disable hyperthreading for multiprocessing.
-            frames (list[str]): MATLAB frames spec (e.g., ["all"], integers, or ranges like "4-8").
-            no_rejection (bool): EchonetLVH flag to skip manual_rejections.txt filtering.
-            batch (str|None): EchonetLVH Batch directory to process (e.g., "Batch2").
-            convert_measurements (bool): EchonetLVH flag to convert only measurements CSV.
-            convert_images (bool): EchonetLVH flag to convert only image files.
-            max_files (int|None): EchonetLVH maximum number of files to process.
-            force (bool): EchonetLVH flag to force recomputation even if parameters exist.
-    """
-    parser = argparse.ArgumentParser(description="Convert raw data to a zea dataset.")
-    parser.add_argument(
-        "dataset",
-        choices=["echonet", "echonetlvh", "camus", "picmus", "verasonics"],
-        help="Raw dataset to convert",
-    )
-    parser.add_argument("src", type=str, help="Source folder path")
-    parser.add_argument("dst", type=str, help="Destination folder path")
-    parser.add_argument(
+def _add_parser_args_echonet(subparsers):
+    """Add Echonet specific arguments to the parser."""
+    echonet_parser = subparsers.add_parser("echonet", help="Convert Echonet dataset")
+    echonet_parser.add_argument("src", type=str, help="Source folder path")
+    echonet_parser.add_argument("dst", type=str, help="Destination folder path")
+    echonet_parser.add_argument(
         "--split_path",
         type=str,
         help="Path to the split.yaml file containing the dataset split if a split should be copied",
     )
-    parser.add_argument(
+    echonet_parser.add_argument(
         "--no_hyperthreading",
         action="store_true",
         help="Disable hyperthreading for multiprocessing",
     )
-    # Dataset specific arguments:
 
-    # verasonics
-    parser.add_argument(
-        "--frames",
-        default=["all"],
-        type=str,
-        nargs="+",
-        help="verasonics: The frames to add to the file. This can be a list of integers, a range "
-        "of integers (e.g. 4-8), or 'all'.",
+
+def _add_parser_args_camus(subparsers):
+    """Add CAMUS specific arguments to the parser."""
+    camus_parser = subparsers.add_parser("camus", help="Convert CAMUS dataset")
+    camus_parser.add_argument("src", type=str, help="Source folder path")
+    camus_parser.add_argument("dst", type=str, help="Destination folder path")
+    camus_parser.add_argument(
+        "--no_hyperthreading",
+        action="store_true",
+        help="Disable hyperthreading for multiprocessing",
     )
-    # ECHONET_LVH
-    parser.add_argument(
+
+
+def _add_parser_args_echonetlvh(subparsers):
+    """Add EchonetLVH specific arguments to the parser."""
+    echonetlvh_parser = subparsers.add_parser("echonetlvh", help="Convert EchonetLVH dataset")
+    echonetlvh_parser.add_argument("src", type=str, help="Source folder path")
+    echonetlvh_parser.add_argument("dst", type=str, help="Destination folder path")
+    echonetlvh_parser.add_argument(
         "--no_rejection",
         action="store_true",
-        help="EchonetLVH: Do not reject sequences in manual_rejections.txt",
+        help="Do not reject sequences in manual_rejections.txt",
     )
-
-    parser.add_argument(
+    echonetlvh_parser.add_argument(
         "--batch",
         type=str,
         default=None,
-        help="EchonetLVH: Specify which BatchX directory to process, e.g. --batch=Batch2",
+        help="Specify which BatchX directory to process, e.g. --batch=Batch2",
     )
-    parser.add_argument(
+    echonetlvh_parser.add_argument(
         "--convert_measurements",
         action="store_true",
-        help="EchonetLVH: Only convert measurements CSV file",
+        help="Only convert measurements CSV file",
     )
-    parser.add_argument(
-        "--convert_images", action="store_true", help="EchonetLVH: Only convert image files"
+    echonetlvh_parser.add_argument(
+        "--convert_images",
+        action="store_true",
+        help="Only convert image files",
     )
-    parser.add_argument(
+    echonetlvh_parser.add_argument(
         "--max_files",
         type=int,
         default=None,
-        help="EchonetLVH: Maximum number of files to process (for testing)",
+        help="Maximum number of files to process (for testing)",
     )
-    parser.add_argument(
+    echonetlvh_parser.add_argument(
         "--force",
         action="store_true",
-        help="EchonetLVH: Force recomputation even if parameters already exist",
+        help="Force recomputation even if parameters already exist",
     )
+    echonetlvh_parser.add_argument(
+        "--no_hyperthreading",
+        action="store_true",
+        help="Disable hyperthreading for multiprocessing",
+    )
+
+
+def _add_parser_args_picmus(subparsers):
+    """Add PICMUS specific arguments to the parser."""
+    picmus_parser = subparsers.add_parser("picmus", help="Convert PICMUS dataset")
+    picmus_parser.add_argument("src", type=str, help="Source folder path")
+    picmus_parser.add_argument("dst", type=str, help="Destination folder path")
+
+
+def _add_parser_args_verasonics(subparsers):
+    verasonics_parser = subparsers.add_parser(
+        "verasonics", help="Convert Verasonics data to zea dataset"
+    )
+    verasonics_parser.add_argument("src", type=str, help="Source folder path")
+    verasonics_parser.add_argument("dst", type=str, help="Destination folder path")
+    verasonics_parser.add_argument(
+        "--frames",
+        default=["all"],
+        type=str,
+        nargs="+",
+        help="The frames to add to the file. This can be a list of integers, a range "
+        "of integers (e.g. 4-8), or 'all'.",
+    )
+    verasonics_parser.add_argument(
+        "--allow_accumulate",
+        action="store_true",
+        help=(
+            "Sometimes, some transmits are already accumulated on the Verasonics system "
+            "(e.g. harmonic imaging through pulse inversion). In this case, the mode in the "
+            "Receive structure is set to 1 (accumulate). If this flag is set, such files "
+            "will be processed. Otherwise, an error is raised when such a mode is detected."
+        ),
+    )
+    verasonics_parser.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        help="Device to use for conversion (e.g., 'cpu' or 'gpu:0').",
+    )
+
+
+def get_parser():
+    """Build and parse command-line arguments for converting raw datasets to a zea dataset."""
+    parser = argparse.ArgumentParser(description="Convert raw data to a zea dataset.")
+    subparsers = parser.add_subparsers(dest="dataset", required=True)
+    _add_parser_args_echonet(subparsers)
+    _add_parser_args_echonetlvh(subparsers)
+    _add_parser_args_camus(subparsers)
+    _add_parser_args_picmus(subparsers)
+    _add_parser_args_verasonics(subparsers)
     return parser
 
 

zea/data/convert/camus.py

@@ -25,8 +25,8 @@ from tqdm import tqdm
 from zea import log
 from zea.data.convert.utils import unzip
 from zea.data.data_format import generate_zea_dataset
+from zea.func.tensor import translate
 from zea.internal.utils import find_first_nonzero_index
-from zea.tensor_ops import translate
 
 
 def transform_sc_image_to_polar(image_sc, output_size=None, fit_outline=True):
@@ -133,7 +133,13 @@ def sitk_load(filepath: str | Path) -> Tuple[np.ndarray, Dict[str, Any]]:
         - Collection of metadata.
     """
     # Load image and save info
-    import SimpleITK as sitk
+    try:
+        import SimpleITK as sitk
+    except ImportError as exc:
+        raise ImportError(
+            "SimpleITK is not installed. "
+            "Please install it with `pip install SimpleITK` to convert CAMUS dataset."
+        ) from exc
 
     image = sitk.ReadImage(str(filepath))
 

zea/data/convert/echonet.py

@@ -24,7 +24,7 @@ from tqdm import tqdm
 from zea import log
 from zea.data import generate_zea_dataset
 from zea.data.convert.utils import load_avi, unzip
-from zea.tensor_ops import translate
+from zea.func.tensor import translate
 
 
 def segment(tensor, number_erasing=0, min_clip=0):

zea/data/convert/echonetlvh/__init__.py

@@ -30,7 +30,7 @@ from zea.data.convert.echonet import H5Processor
 from zea.data.convert.echonetlvh.precompute_crop import precompute_cone_parameters
 from zea.data.convert.utils import load_avi, unzip
 from zea.display import cartesian_to_polar_matrix
-from zea.tensor_ops import translate
+from zea.func.tensor import translate
 
 
 def overwrite_splits(source_dir):