zea 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

zea/data/data_format.py

@@ -7,11 +7,12 @@ from dataclasses import dataclass
 from pathlib import Path
 
 import numpy as np
+from keras.utils import pad_sequences
 
 from zea import log
 from zea.data.file import File, validate_file
 from zea.internal.checks import _DATA_TYPES
-from zea.utils import first_not_none_item
+from zea.internal.utils import first_not_none_item
 
 
 @dataclass
@@ -470,32 +471,30 @@ def _write_datasets(
         )
 
         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",
-                )
+            _add_dataset(
+                group_name=scan_group_name,
+                name="waveforms_one_way",
+                data=pad_sequences(waveforms_one_way, dtype=np.float32, padding="post"),
+                description=(
+                    "One-way waveform as simulated by the Verasonics system, "
+                    "sampled at 250MHz. This is the waveform after being filtered "
+                    "by the transducer 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_dataset(
+                group_name=scan_group_name,
+                name="waveforms_two_way",
+                data=pad_sequences(waveforms_two_way, dtype=np.float32, padding="post"),
+                description=(
+                    "Two-way waveform as simulated by the Verasonics system, "
+                    "sampled at 250MHz. This is the waveform after being filtered "
+                    "by the transducer bandwidth twice."
+                ),
+                unit="V",
+            )
 
     # Add additional elements
     if additional_elements is not None:
@@ -539,6 +538,7 @@ def generate_zea_dataset(
     additional_elements=None,
     event_structure=False,
     cast_to_float=True,
+    overwrite=False,
 ):
     """Generates a dataset in the zea format.
 
@@ -585,10 +585,10 @@ def generate_zea_dataset(
             waveform was used for each transmit event.
         waveforms_one_way (list): List of one-way waveforms as simulated by the Verasonics
             system, sampled at 250MHz. This is the waveform after being filtered by the
-            tranducer bandwidth once. Every element in the list is a 1D numpy array.
+            transducer bandwidth once. Every element in the list is a 1D numpy array.
         waveforms_two_way (list): List of two-way waveforms as simulated by the Verasonics
             system, sampled at 250MHz. This is the waveform after being filtered by the
-            tranducer bandwidth twice. Every element in the list is a 1D numpy array.
+            transducer bandwidth twice. Every element in the list is a 1D numpy array.
         additional_elements (List[DatasetElement]): A list of additional dataset
             elements to be added to the dataset. Each element should be a DatasetElement
             object. The additional elements are added under the scan group.
@@ -598,6 +598,7 @@ def generate_zea_dataset(
             Instead of just a single data and scan group.
         cast_to_float (bool): Whether to store data as float32. You may want to set this
             to False if storing images.
+        overwrite (bool): Whether to overwrite the file if it already exists. Defaults to False.
 
     """
     # check if all args are lists
@@ -637,10 +638,10 @@ def generate_zea_dataset(
     # make sure input arguments of func is same length as data_and_parameters
     # except `path` and `event_structure` arguments and ofcourse `data_and_parameters` itself
     assert (
-        len(data_and_parameters) == len(inspect.signature(generate_zea_dataset).parameters) - 3
+        len(data_and_parameters) == len(inspect.signature(generate_zea_dataset).parameters) - 4
     ), (
         "All arguments should be put in data_and_parameters except "
-        "`path`, `event_structure`, and `cast_to_float` arguments."
+        "`path`, `event_structure`, `cast_to_float`, and `overwrite` arguments."
     )
 
     if event_structure:
@@ -682,7 +683,7 @@ def generate_zea_dataset(
     # Convert path to Path object
     path = Path(path)
 
-    if path.exists():
+    if path.exists() and not overwrite:
         raise FileExistsError(f"The file {path} already exists.")
 
     # Create the directory if it does not exist

zea/data/datasets.py

@@ -48,13 +48,15 @@ from zea.data.preset_utils import (
     _hf_resolve_path,
 )
 from zea.datapaths import format_data_path
+from zea.internal.utils import (
+    calculate_file_hash,
+    reduce_to_signature,
+)
 from zea.io_lib import search_file_tree
 from zea.tools.hf import HFPath
 from zea.utils import (
-    calculate_file_hash,
     date_string_to_readable,
     get_date_string,
-    reduce_to_signature,
 )
 
 _CHECK_MAX_DATASET_SIZE = 10000
@@ -241,7 +243,7 @@ class Folder:
             return
 
         num_frames_per_file = []
-        validated_succesfully = True
+        validated_successfully = True
         for file_path in tqdm.tqdm(
             self.file_paths,
             total=self.n_files,
@@ -253,9 +255,9 @@ class Folder:
                 validation_error_log.append(f"File {file_path} is not a valid zea dataset.\n{e}\n")
                 # convert into warning
                 log.warning(f"Error in file {file_path}.\n{e}")
-                validated_succesfully = False
+                validated_successfully = False
 
-        if not validated_succesfully:
+        if not validated_successfully:
             log.warning(
                 "Check warnings above for details. No validation file was created. "
                 f"See {validation_error_file_path} for details."

zea/data/file.py

@@ -6,6 +6,7 @@ from typing import List
 
 import h5py
 import numpy as np
+from keras.utils import pad_sequences
 
 from zea import log
 from zea.data.preset_utils import HF_PREFIX, _hf_resolve_path
@@ -15,9 +16,9 @@ from zea.internal.checks import (
     _REQUIRED_SCAN_KEYS,
     get_check,
 )
+from zea.internal.utils import reduce_to_signature
 from zea.probes import Probe
 from zea.scan import Scan
-from zea.utils import reduce_to_signature
 
 
 def assert_key(file: h5py.File, key: str):
@@ -219,6 +220,41 @@ class File(h5py.File):
     def load_data(self, data_type, indices: str | int | List[int] = "all"):
         """Load data from the file.
 
+        The indices parameter can be used to load a subset of the data. This can be
+
+        - 'all' to load all data
+
+        - an int to load a single frame
+
+        - a list of ints to load specific frames
+
+        - a tuple of lists, ranges or slices to index frames and transmits. Note that
+          indexing with lists of indices for both axes is not supported. In that case,
+          try to define one of the axes with a slice.
+
+        .. doctest::
+
+            >>> from zea import File
+
+            >>> path_to_file = (
+            ...     "hf://zeahub/picmus/database/experiments/contrast_speckle/"
+            ...     "contrast_speckle_expe_dataset_iq/contrast_speckle_expe_dataset_iq.hdf5"
+            ... )
+
+            >>> with File(path_to_file, mode="r") as file:
+            ...     # data has shape (n_frames, n_tx, n_el, n_ax, n_ch)
+            ...     data = file.load_data("raw_data")
+            ...     data.shape
+            ...     # load first frame only
+            ...     data = file.load_data("raw_data", indices=0)
+            ...     data.shape
+            ...     # load frame 0 and transmits 0, 2 and 4
+            ...     data = file.load_data("raw_data", indices=(0, [0, 2, 4]))
+            ...     data.shape
+            (1, 75, 832, 128, 2)
+            (75, 832, 128, 2)
+            (3, 832, 128, 2)
+
         Args:
             data_type (str): The type of data to load. Options are 'raw_data', 'aligned_data',
                 'beamformed_data', 'envelope_data', 'image' and 'image_sc'.
@@ -337,7 +373,7 @@ class File(h5py.File):
         Returns:
             Scan: The scan object.
         """
-        return Scan.merge(self.get_scan_parameters(event), kwargs, safe=safe)
+        return Scan.merge(_reformat_waveforms(self.get_scan_parameters(event)), kwargs, safe=safe)
 
     def get_probe_parameters(self, event=None) -> dict:
         """Returns a dictionary of probe parameters to initialize a probe
@@ -772,3 +808,69 @@ def _assert_unit_and_description_present(hdf5_file, _prefix=""):
             assert "description" in hdf5_file[key].attrs.keys(), (
                 f"The file {_prefix}/{key} does not have a description attribute."
             )
+
+
+def _reformat_waveforms(scan_kwargs: dict) -> dict:
+    """Reformat waveforms from dict to array if needed. This is for backwards compatibility and will
+    be removed in a future version of zea.
+
+    Args:
+        scan_kwargs (dict): The scan parameters.
+
+    Returns:
+        scan_kwargs (dict): The scan parameters with the keys waveforms_one_way and
+            waveforms_two_way reformatted to arrays if they were stored as dicts.
+    """
+
+    # TODO: remove this in a future version of zea
+    if "waveforms_one_way" in scan_kwargs and isinstance(scan_kwargs["waveforms_one_way"], dict):
+        log.warning(
+            "The waveforms_one_way parameter is stored as a dictionary in the file. "
+            "Converting to array. This will be deprecated in future versions of zea. "
+            "Please update your files to store waveforms as arrays of shape `(n_tx, n_samples)`."
+        )
+        scan_kwargs["waveforms_one_way"] = _waveforms_dict_to_array(
+            scan_kwargs["waveforms_one_way"]
+        )
+
+    if "waveforms_two_way" in scan_kwargs and isinstance(scan_kwargs["waveforms_two_way"], dict):
+        log.warning(
+            "The waveforms_two_way parameter is stored as a dictionary in the file. "
+            "Converting to array. This will be deprecated in future versions of zea. "
+            "Please update your files to store waveforms as arrays of shape `(n_tx, n_samples)`."
+        )
+        scan_kwargs["waveforms_two_way"] = _waveforms_dict_to_array(
+            scan_kwargs["waveforms_two_way"]
+        )
+    return scan_kwargs
+
+
+def _waveforms_dict_to_array(waveforms_dict: dict):
+    """Convert waveforms stored as a dictionary to a padded numpy array."""
+    waveforms = dict_to_sorted_list(waveforms_dict)
+    return pad_sequences(waveforms, dtype=np.float32, padding="post")
+
+
+def dict_to_sorted_list(dictionary: dict):
+    """Convert a dictionary with sortable keys to a sorted list of values.
+
+    .. note::
+
+        This function operates on the top level of the dictionary only.
+        If the dictionary contains nested dictionaries, those will not be sorted.
+
+    Example:
+        .. doctest::
+
+            >>> from zea.data.file import dict_to_sorted_list
+            >>> input_dict = {"number_000": 5, "number_001": 1, "number_002": 23}
+            >>> dict_to_sorted_list(input_dict)
+            [5, 1, 23]
+
+    Args:
+        dictionary (dict): The dictionary to convert. The keys must be sortable.
+
+    Returns:
+        list: The sorted list of values.
+    """
+    return [value for _, value in sorted(dictionary.items())]

zea/data/layers.py

@@ -2,7 +2,7 @@
 
 import keras
 import numpy as np
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
+from keras.src.layers.preprocessing.data_layer import DataLayer
 
 from zea.ops import Pad as PadOp
 from zea.utils import map_negative_indices
@@ -11,7 +11,7 @@ from zea.utils import map_negative_indices
 class Pad(PadOp):
     """Pad layer for padding tensors to a specified shape which can be used in tf.data pipelines."""
 
-    __call__ = TFDataLayer.__call__
+    __call__ = DataLayer.__call__
 
     def call(self, inputs):
         """
@@ -20,12 +20,12 @@ class Pad(PadOp):
         return super().call(data=inputs)["data"]
 
 
-class Resizer(TFDataLayer):
+class Resizer(DataLayer):
     """
     Resize layer for resizing images. Can deal with N-dimensional images.
     Can do resize, center_crop, random_crop and crop_or_pad.
 
-    Can be used in tf.data pipelines.
+    Can be used in `tf.data` pipelines.
     """
 
     def __init__(
@@ -36,7 +36,6 @@ class Resizer(TFDataLayer):
         seed: int | None = None,
         **resize_kwargs,
     ):
-        # noqa: E501
         """
         Initializes the data loader with the specified parameters.
 
@@ -47,7 +46,7 @@ class Resizer(TFDataLayer):
                 ['random_crop'](https://keras.io/api/layers/preprocessing_layers/image_augmentation/random_crop/),
                 ['resize'](https://keras.io/api/layers/preprocessing_layers/image_preprocessing/resizing/),
                 'crop_or_pad': resizes an image to a target width and height by either centrally
-                    cropping the image, padding it evenly with zeros or a combination of both.
+                cropping the image, padding it evenly with zeros or a combination of both.
             resize_axes (tuple | None, optional): The axes along which to resize.
                 Must be of length 2. Defaults to None. In that case, can only process
                 default tensors of shape (batch, height, width, channels), where the

zea/datapaths.py

@@ -11,12 +11,24 @@ to set up your local data paths.
 Example usage
 ^^^^^^^^^^^^^
 
-.. code-block:: python
+.. doctest::
 
-    from zea.datapaths import set_data_paths
+    >>> import yaml
+    >>> from zea.datapaths import set_data_paths
 
-    user = set_data_paths("users.yaml")
-    print(user.data_root)
+    >>> user_config = {"data_root": "/path/to/data", "output": "/path/to/output"}
+    >>> with open("users.yaml", "w", encoding="utf-8") as file:
+    ...     yaml.dump(user_config, file)
+
+    >>> user = set_data_paths("users.yaml")
+    >>> print(user.data_root)
+    /path/to/data
+
+.. testcleanup::
+
+    import os
+
+    os.remove("users.yaml")
 
 """
 

zea/display.py

@@ -9,8 +9,8 @@ from keras import ops
 from PIL import Image
 
 from zea import log
+from zea.tensor_ops import translate
 from zea.tools.fit_scan_cone import fit_and_crop_around_scan_cone
-from zea.utils import translate
 
 
 def to_8bit(image, dynamic_range: Union[None, tuple] = None, pillow: bool = True):
@@ -340,12 +340,14 @@ def scan_convert(
 def map_coordinates(inputs, coordinates, order, fill_mode="constant", fill_value=0):
     """map_coordinates using keras.ops or scipy.ndimage when order > 1."""
     if order > 1:
-        inputs = ops.convert_to_numpy(inputs)
-        coordinates = ops.convert_to_numpy(coordinates)
+        # Preserve original dtype before conversion
+        original_dtype = ops.dtype(inputs)
+        inputs_np = ops.convert_to_numpy(inputs).astype(np.float32)
+        coordinates_np = ops.convert_to_numpy(coordinates).astype(np.float32)
         out = scipy.ndimage.map_coordinates(
-            inputs, coordinates, order=order, mode=fill_mode, cval=fill_value
+            inputs_np, coordinates_np, order=order, mode=fill_mode, cval=fill_value
         )
-        return ops.convert_to_tensor(out)
+        return ops.convert_to_tensor(out.astype(original_dtype))
     else:
         return ops.image.map_coordinates(
             inputs,

zea/interface.py

@@ -3,15 +3,15 @@
 Example usage
 ^^^^^^^^^^^^^^
 
-.. code-block:: python
+.. doctest::
 
-    import zea
-    from zea.internal.setup_zea import setup_config
+    >>> import zea
+    >>> from zea.internal.setup_zea import setup_config
 
-    config = setup_config("hf://zeahub/configs/config_camus.yaml")
+    >>> config = setup_config("hf://zeahub/configs/config_camus.yaml")
 
-    interface = zea.Interface(config)
-    interface.run(plot=True)
+    >>> interface = zea.Interface(config)
+    >>> interface.run(plot=True)  # doctest: +SKIP
 
 """
 
@@ -31,15 +31,15 @@ from zea.data.file import File
 from zea.datapaths import format_data_path
 from zea.display import to_8bit
 from zea.internal.core import DataTypes
+from zea.internal.utils import keep_trying
 from zea.internal.viewer import (
     ImageViewerMatplotlib,
     ImageViewerOpenCV,
     filename_from_window_dialog,
     running_in_notebook,
 )
-from zea.io_lib import matplotlib_figure_to_numpy
+from zea.io_lib import matplotlib_figure_to_numpy, save_video
 from zea.ops import Pipeline
-from zea.utils import keep_trying, save_to_gif, save_to_mp4
 
 
 class Interface:
@@ -266,10 +266,11 @@ class Interface:
         save = self.config.plot.save
 
         if self.frame_no == "all":
-            if not asyncio.get_event_loop().is_running():
-                asyncio.run(self.run_movie(save))
-            else:
-                asyncio.create_task(self.run_movie(save))
+            try:
+                loop = asyncio.get_running_loop()
+                loop.create_task(self.run_movie(save))  # already running loop
+            except RuntimeError:
+                asyncio.run(self.run_movie(save))  # no loop yet
 
         else:
             if plot:
@@ -520,10 +521,7 @@ class Interface:
 
         fps = self.config.plot.fps
 
-        if self.config.plot.video_extension == "gif":
-            save_to_gif(images, path, fps=fps)
-        elif self.config.plot.video_extension == "mp4":
-            save_to_mp4(images, path, fps=fps)
+        save_video(images, path, fps=fps)
 
         if self.verbose:
             log.info(f"Video saved to {log.yellow(path)}")

zea/internal/_generate_keras_ops.py

@@ -3,11 +3,10 @@ and :mod:`keras.ops.image` functions.
 
 They can be used in zea pipelines like any other :class:`zea.Operation`, for example:
 
-.. code-block:: python
+.. doctest::
 
-    from zea.keras_ops import Squeeze
-
-    op = Squeeze(axis=1)
+    >>> from zea.keras_ops import Squeeze
+    >>> op = Squeeze(axis=1)
 """
 
 import inspect
@@ -77,11 +76,11 @@ and :mod:`keras.ops.image` functions.
 
 They can be used in zea pipelines like any other :class:`zea.Operation`, for example:
 
-.. code-block:: python
+.. doctest::
 
-    from zea.keras_ops import Squeeze
+    >>> from zea.keras_ops import Squeeze
 
-    op = Squeeze(axis=1)
+    >>> op = Squeeze(axis=1)
 
 This file is generated automatically. Do not edit manually.
 Generated with Keras {keras.__version__}

zea/internal/cache.py

@@ -21,10 +21,8 @@
 
 import ast
 import atexit
-import hashlib
 import inspect
 import os
-import pickle
 import tempfile
 import textwrap
 from pathlib import Path
@@ -33,6 +31,7 @@ import joblib
 import keras
 
 from zea import log
+from zea.internal.core import hash_elements
 
 _DEFAULT_ZEA_CACHE_DIR = Path.home() / ".cache" / "zea"
 
@@ -80,52 +79,6 @@ _CACHE_DIR = ZEA_CACHE_DIR / "cached_funcs"
 _CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
 
-def serialize_elements(key_elements: list, shorten: bool = False) -> str:
-    """Serialize elements of a list to generate a cache key.
-
-    In general uses the string representation of the elements unless
-    the element has a `serialized` attribute, in which case it uses that.
-    For instance this is useful for custom classes that inherit from `zea.core.Object`.
-
-    Args:
-        key_elements (list): List of elements to serialize. Can be nested lists
-            or tuples. In this case the elements are serialized recursively.
-        shorten (bool): If True, the serialized string is hashed to a shorter
-            representation using MD5. Defaults to False.
-
-    Returns:
-        str: A serialized string representation of the elements, joined by underscores.
-
-    """
-    serialized_elements = []
-    for element in key_elements:
-        if isinstance(element, (list, tuple)):
-            # If element is a list or tuple, serialize its elements recursively
-            serialized_elements.append(serialize_elements(element))
-        elif hasattr(element, "serialized"):
-            # Use the serialized attribute if it exists (e.g. for zea.core.Object)
-            serialized_elements.append(str(element.serialized))
-        elif isinstance(element, str):
-            # If element is a string, use it as is
-            serialized_elements.append(element)
-        elif isinstance(element, keras.random.SeedGenerator):
-            # If element is a SeedGenerator, use the state
-            element = keras.ops.convert_to_numpy(element.state.value)
-            element = pickle.dumps(element)
-            element = hashlib.md5(element).hexdigest()
-            serialized_elements.append(element)
-        else:
-            # Otherwise, serialize the element using pickle and hash it
-            element = pickle.dumps(element)
-            element = hashlib.md5(element).hexdigest()
-            serialized_elements.append(element)
-
-    serialized = "_".join(serialized_elements)
-    if shorten:
-        return hashlib.md5(serialized.encode()).hexdigest()
-    return serialized
-
-
 def get_function_source(func):
     """Recursively get the source code of a function and its nested functions."""
     try:
@@ -188,7 +141,7 @@ def generate_cache_key(func, args, kwargs, arg_names):
     # Add keras backend
     key_elements.append(keras.backend.backend())
 
-    return f"{func.__qualname__}_" + serialize_elements(key_elements, shorten=True)
+    return f"{func.__qualname__}_" + hash_elements(key_elements)
 
 
 def cache_output(*arg_names, verbose=False):

zea/internal/config/validation.py

@@ -15,9 +15,8 @@ from pathlib import Path
 
 from schema import And, Optional, Or, Regex, Schema
 
-import zea.metrics  # noqa: F401
 from zea.internal.checks import _DATA_TYPES
-from zea.internal.registry import metrics_registry
+from zea.metrics import metrics_registry
 
 # predefined checks, later used in schema to check validity of parameter
 any_number = Or(