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

zea/__init__.py

@@ -7,10 +7,10 @@ from . import log
 
 # dynamically add __version__ attribute (see pyproject.toml)
 # __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.4"
+__version__ = "0.0.5"
 
 
-def setup():
+def _bootstrap_backend():
     """Setup function to initialize the zea package."""
 
     def _check_backend_installed():
@@ -40,14 +40,14 @@ def setup():
 
     _check_backend_installed()
 
-    import keras
+    from keras.backend import backend as keras_backend
 
-    log.info(f"Using backend {keras.backend.backend()!r}")
+    log.info(f"Using backend {keras_backend()!r}")
 
 
 # call and clean up namespace
-setup()
-del setup
+_bootstrap_backend()
+del _bootstrap_backend
 
 from . import (
     agent,
@@ -55,6 +55,7 @@ from . import (
     data,
     display,
     io_lib,
+    keras_ops,
     metrics,
     models,
     simulator,
@@ -68,7 +69,7 @@ from .data.file import File, load_file
 from .datapaths import set_data_paths
 from .interface import Interface
 from .internal.device import init_device
-from .internal.setup_zea import set_backend, setup, setup_config
+from .internal.setup_zea import setup, setup_config
 from .ops import Pipeline
 from .probes import Probe
 from .scan import Scan

zea/__main__.py

@@ -9,30 +9,22 @@ import argparse
 import sys
 from pathlib import Path
 
-from zea import log
 from zea.visualize import set_mpl_style
 
 
-def get_args():
+def get_parser():
     """Command line argument parser"""
-    parser = argparse.ArgumentParser(description="Process ultrasound data.")
-    parser.add_argument("-c", "--config", type=str, default=None, help="path to config file.")
+    parser = argparse.ArgumentParser(
+        description="Load and process ultrasound data based on a configuration file."
+    )
+    parser.add_argument("-c", "--config", type=str, default=None, help="path to the config file.")
     parser.add_argument(
         "-t",
         "--task",
         default="view",
         choices=["view"],
         type=str,
-        help="which task to run",
-    )
-    parser.add_argument(
-        "--backend",
-        default=None,
-        type=str,
-        help=(
-            "Keras backend to use. Default is the one set by the environment "
-            "variable KERAS_BACKEND."
-        ),
+        help="Which task to run. Currently only 'view' is supported.",
     )
     parser.add_argument(
         "--skip_validate_file",
@@ -40,27 +32,18 @@ def get_args():
         action="store_true",
         help="Skip zea file integrity checks. Use with caution.",
     )
-    parser.add_argument("--gui", default=False, action=argparse.BooleanOptionalAction)
-    args = parser.parse_args()
-    return args
+    return parser
 
 
 def main():
     """main entrypoint for zea"""
-    args = get_args()
+    args = get_parser().parse_args()
 
     set_mpl_style()
 
-    if args.backend:
-        from zea.internal.setup_zea import set_backend
-
-        set_backend(args.backend)
-
     wd = Path(__file__).parent.resolve()
     sys.path.append(str(wd))
 
-    import keras
-
     from zea.interface import Interface
     from zea.internal.setup_zea import setup
 
@@ -72,7 +55,6 @@ def main():
             validate_file=not args.skip_validate_file,
         )
 
-        log.info(f"Using {keras.backend.backend()} backend")
         cli.run(plot=True)
     else:
         raise ValueError(f"Unknown task {args.task}, see `zea --help` for available tasks.")

zea/data/__main__.py

@@ -9,8 +9,8 @@ import argparse
 from zea import Folder
 
 
-def main():
-    parser = argparse.ArgumentParser(description="Copy a zea.Folder to a new location.")
+def get_parser():
+    parser = argparse.ArgumentParser(description="Copy a :class:`zea.Folder` to a new location.")
     parser.add_argument("src", help="Source folder path")
     parser.add_argument("dst", help="Destination folder path")
     parser.add_argument("key", help="Key to access in the hdf5 files")
@@ -20,8 +20,11 @@ def main():
         choices=["a", "w", "r+", "x"],
         help="Mode in which to open the destination files (default: 'a')",
     )
+    return parser
+
 
-    args = parser.parse_args()
+def main():
+    args = get_parser().parse_args()
 
     src_folder = Folder(args.src, args.key, validate=False)
     src_folder.copy(args.dst, args.key, mode=args.mode)

zea/data/file.py

@@ -290,7 +290,7 @@ class File(h5py.File):
         """
         scan_parameters = {}
         if "scan" in self:
-            scan_parameters = recursively_load_dict_contents_from_group(self, "scan")
+            scan_parameters = self.recursively_load_dict_contents_from_group("scan")
         elif "event" in list(self.keys())[0]:
             if event is None:
                 raise ValueError(
@@ -305,38 +305,17 @@ class File(h5py.File):
                 f"Found number of events: {len(self.keys())}."
             )
 
-            scan_parameters = recursively_load_dict_contents_from_group(self, f"event_{event}/scan")
+            scan_parameters = self.recursively_load_dict_contents_from_group(f"event_{event}/scan")
         else:
             log.warning("Could not find scan parameters in file.")
 
         return scan_parameters
 
     def get_scan_parameters(self, event=None) -> dict:
-        """Returns a dictionary of default parameters to initialize a scan
-        object that works with the file.
+        """Returns a dictionary of scan parameters stored in the file."""
+        return self.get_parameters(event)
 
-        Returns:
-            dict: The default parameters (the keys are identical to the
-                __init__ parameters of the Scan class).
-        """
-        file_scan_parameters = self.get_parameters(event)
-
-        scan_parameters = {}
-        for parameter, value in file_scan_parameters.items():
-            if parameter in Scan.VALID_PARAMS:
-                param_type = Scan.VALID_PARAMS[parameter]["type"]
-                if param_type in (bool, int, float):
-                    scan_parameters[parameter] = param_type(value)
-                elif isinstance(param_type, tuple) and float in param_type:
-                    scan_parameters[parameter] = float(value)
-                else:
-                    scan_parameters[parameter] = value
-
-        if len(scan_parameters) == 0:
-            log.info(f"Could not find proper scan parameters in {self}.")
-        return scan_parameters
-
-    def scan(self, event=None, **kwargs) -> Scan:
+    def scan(self, event=None, safe=True, **kwargs) -> Scan:
         """Returns a Scan object initialized with the parameters from the file.
 
         Args:
@@ -348,6 +327,9 @@ class File(h5py.File):
                     ...
 
                 Defaults to None. In that case no event structure is expected.
+            safe (bool, optional): If True, will only use parameters that are
+                defined in the Scan class. If False, will use all parameters
+                from the file. Defaults to True.
             **kwargs: Additional keyword arguments to pass to the Scan object.
                 These will override the parameters from the file if they are
                 present in the file.
@@ -355,7 +337,7 @@ class File(h5py.File):
         Returns:
             Scan: The scan object.
         """
-        return Scan.merge(self.get_scan_parameters(event), kwargs)
+        return Scan.merge(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
@@ -388,21 +370,24 @@ class File(h5py.File):
         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:
+    def recursively_load_dict_contents_from_group(self, path: str) -> dict:
         """Load dict from contents of group
 
         Values inside the group are converted to numpy arrays
-        or primitive types (int, float, str). Single element
-        arrays are converted to the corresponding primitive type (if squeeze=True)
+        or primitive types (int, float, str).
 
         Args:
             path (str): path to group
-            squeeze (bool, optional): squeeze arrays with single element.
-                Defaults to False.
         Returns:
             dict: dictionary with contents of group
         """
-        return recursively_load_dict_contents_from_group(self, path, squeeze)
+        ans = {}
+        for key, item in self[path].items():
+            if isinstance(item, h5py.Dataset):
+                ans[key] = item[()]
+            elif isinstance(item, h5py.Group):
+                ans[key] = self.recursively_load_dict_contents_from_group(path + "/" + key + "/")
+        return ans
 
     @classmethod
     def get_shape(cls, path: str, key: str) -> tuple:
@@ -519,54 +504,14 @@ def load_file(
         # the number of selected transmits
         if data_type in ["raw_data", "aligned_data"]:
             indices = File._prepare_indices(indices)
-            n_tx = data.shape[1]
             if isinstance(indices, tuple) and len(indices) > 1:
-                tx_idx = indices[1]
-                transmits = np.arange(n_tx)[tx_idx]
-                scan_kwargs["selected_transmits"] = transmits
+                scan_kwargs["selected_transmits"] = indices[1]
 
         scan = file.scan(**scan_kwargs)
 
         return data, scan, probe
 
 
-def recursively_load_dict_contents_from_group(
-    h5file: h5py._hl.files.File, path: str, squeeze: bool = False
-) -> dict:
-    """Load dict from contents of group
-
-    Values inside the group are converted to numpy arrays
-    or primitive types (int, float, str). Single element
-    arrays are converted to the corresponding primitive type (if squeeze=True)
-
-    Args:
-        h5file (h5py._hl.files.File): h5py file object
-        path (str): path to group
-        squeeze (bool, optional): squeeze arrays with single element.
-            Defaults to False.
-    Returns:
-        dict: dictionary with contents of group
-    """
-    ans = {}
-    for key, item in h5file[path].items():
-        if isinstance(item, h5py._hl.dataset.Dataset):
-            ans[key] = item[()]
-            # all ones in shape
-            if squeeze:
-                if ans[key].shape == () or all(i == 1 for i in ans[key].shape):
-                    # check for strings
-                    if isinstance(ans[key], str):
-                        ans[key] = str(ans[key])
-                    # check for integers
-                    elif int(ans[key]) == float(ans[key]):
-                        ans[key] = int(ans[key])
-                    else:
-                        ans[key] = float(ans[key])
-        elif isinstance(item, h5py._hl.group.Group):
-            ans[key] = recursively_load_dict_contents_from_group(h5file, path + "/" + key + "/")
-    return ans
-
-
 def _print_hdf5_attrs(hdf5_obj, prefix=""):
     """Recursively prints all keys, attributes, and shapes in an HDF5 file.
 

zea/display.py

@@ -3,7 +3,6 @@
 from functools import partial
 from typing import Tuple, Union
 
-import keras
 import numpy as np
 import scipy
 from keras import ops
@@ -342,6 +341,7 @@ def map_coordinates(inputs, coordinates, order, fill_mode="constant", fill_value
     """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)
         out = scipy.ndimage.map_coordinates(
             inputs, coordinates, order=order, mode=fill_mode, cval=fill_value
         )
@@ -359,10 +359,6 @@ def map_coordinates(inputs, coordinates, order, fill_mode="constant", fill_value
 def _interpolate_batch(images, coordinates, fill_value=0.0, order=1, vectorize=True):
     """Interpolate a batch of images."""
 
-    # TODO: figure out why tensorflow map_coordinates is broken
-    if keras.backend.backend() == "tensorflow":
-        assert order > 1, "Some bug in tensorflow in map_coordinates, set order > 1 to use scipy."
-
     image_shape = images.shape
     num_image_dims = coordinates.shape[0]
 

zea/doppler.py

@@ -0,0 +1,75 @@
+"""Doppler functions for processing I/Q ultrasound data."""
+
+import numpy as np
+from keras import ops
+
+from zea import tensor_ops
+
+
+def color_doppler(
+    data,
+    center_frequency,
+    pulse_repetition_frequency,
+    sound_speed,
+    hamming_size=None,
+    lag=1,
+):
+    """Compute Color Doppler from packet of I/Q Data.
+
+    Args:
+        data (ndarray): I/Q complex data of shape (n_frames, grid_size_z, grid_size_x).
+            n_frames corresponds to the ensemble length used to compute
+            the Doppler signal.
+        center_frequency (float): Center frequency of the ultrasound probe in Hz.
+        pulse_repetition_frequency (float): Pulse repetition frequency in Hz.
+        sound_speed (float): Speed of sound in the medium in m/s.
+        hamming_size (int or tuple, optional): Size of the Hamming window to apply
+            for spatial averaging. If None, no window is applied.
+            If an integer, it is applied to both dimensions. If a tuple, it should
+            contain two integers for the row and column dimensions.
+        lag (int, optional): Lag for the auto-correlation computation.
+            Defaults to 1, meaning Doppler is computed from the current frame
+            and the next frame.
+
+    Returns:
+        doppler_velocities (ndarray): Doppler velocity map of shape (grid_size_z, grid_size_x) in
+            meters/second.
+
+    """
+    assert data.ndim == 3, "Data must be a 3-D array"
+    if not (isinstance(lag, int) and lag >= 1):
+        raise ValueError("lag must be an integer >= 1")
+    n_frames = data.shape[0]
+    assert n_frames > lag, "Data must have more frames than the lag"
+
+    if hamming_size is None:
+        hamming_size = np.array([1, 1], dtype=int)
+    elif np.isscalar(hamming_size):
+        hamming_size = np.array([int(hamming_size), int(hamming_size)], dtype=int)
+    else:
+        assert len(hamming_size) == 2, "hamming_size must be an integer or a tuple of two integers"
+        hamming_size = np.array(hamming_size, dtype=int)
+    if not np.all(hamming_size > 0):
+        raise ValueError("hamming_size must contain integers > 0")
+
+    # Auto-correlation method
+    iq1 = data[: n_frames - lag]
+    iq2 = data[lag:]
+    autocorr = ops.sum(iq1 * ops.conj(iq2), axis=0)  # Ensemble auto-correlation
+
+    # Spatial weighted average
+    if hamming_size[0] != 1 and hamming_size[1] != 1:
+        h_row = np.hamming(hamming_size[0])
+        h_col = np.hamming(hamming_size[1])
+        autocorr = tensor_ops.apply_along_axis(
+            lambda x: tensor_ops.correlate(x, h_row, mode="same"), 0, autocorr
+        )
+        autocorr = tensor_ops.apply_along_axis(
+            lambda x: tensor_ops.correlate(x, h_col, mode="same"), 1, autocorr
+        )
+
+    # Doppler velocity
+    nyquist_velocity = sound_speed * pulse_repetition_frequency / (4 * center_frequency * lag)
+    phase = ops.arctan2(ops.imag(autocorr), ops.real(autocorr))
+    doppler_velocities = -nyquist_velocity * phase / np.pi
+    return doppler_velocities

zea/internal/_generate_keras_ops.py

@@ -0,0 +1,125 @@
+"""This file creates a :class:`zea.Operation` for all unary :mod:`keras.ops`
+and :mod:`keras.ops.image` functions.
+
+They can be used in zea pipelines like any other :class:`zea.Operation`, for example:
+
+.. code-block:: python
+
+    from zea.keras_ops import Squeeze
+
+    op = Squeeze(axis=1)
+"""
+
+import inspect
+import shutil
+import tempfile
+from pathlib import Path
+
+import keras
+
+
+def _filter_funcs_by_first_arg(funcs, arg_name):
+    """Filter a list of (name, func) tuples to those whose first argument matches arg_name."""
+    filtered = []
+    for name, func in funcs:
+        try:
+            sig = inspect.signature(func)
+            params = list(sig.parameters.keys())
+            if params and params[0] == arg_name:
+                filtered.append((name, func))
+        except (ValueError, TypeError):
+            # Skip functions that can't be inspected
+            continue
+    return filtered
+
+
+def _functions_from_namespace(namespace):
+    """Get all functions from a given namespace."""
+    return [(name, obj) for name, obj in inspect.getmembers(namespace) if inspect.isfunction(obj)]
+
+
+def _unary_functions_from_namespace(namespace, arg_name="x"):
+    """Get all unary functions from a given namespace."""
+    funcs = _functions_from_namespace(namespace)
+    return _filter_funcs_by_first_arg(funcs, arg_name)
+
+
+def _snake_to_pascal(name):
+    """Convert a snake_case name to PascalCase."""
+    return "".join(word.capitalize() for word in name.split("_"))
+
+
+def _generate_operation_class_code(name, namespace):
+    """Generate Python code for a zea.Operation class for a given keras.ops function."""
+    class_name = _snake_to_pascal(name)
+    module_path = f"{namespace.__name__}.{name}"
+    doc = f"Operation wrapping {module_path}."
+
+    return f'''
+@ops_registry("{module_path}")
+class {class_name}(Lambda):
+    """{doc}"""
+
+    def __init__(self, **kwargs):
+        try:
+            super().__init__(func={module_path}, **kwargs)
+        except AttributeError as e:
+            raise MissingKerasOps("{class_name}", "{module_path}") from e
+'''
+
+
+def _generate_ops_file():
+    """Generate a .py file with all operation class definitions."""
+
+    # File header with version info
+    content = f'''"""Auto-generated :class:`zea.Operation` for all unary :mod:`keras.ops`
+and :mod:`keras.ops.image` functions.
+
+They can be used in zea pipelines like any other :class:`zea.Operation`, for example:
+
+.. code-block:: python
+
+    from zea.keras_ops import Squeeze
+
+    op = Squeeze(axis=1)
+
+This file is generated automatically. Do not edit manually.
+Generated with Keras {keras.__version__}
+"""
+
+import keras
+
+from zea.internal.registry import ops_registry
+from zea.ops import Lambda
+
+class MissingKerasOps(ValueError):
+    def __init__(self, class_name: str, func: str):
+        super().__init__(
+            f"Failed to create {{class_name}} with {{func}}. " +
+            "This may be due to an incompatible version of `keras`. " +
+            "Please try to upgrade `keras` to the latest version by running " +
+            "`pip install --upgrade keras`."
+        )
+
+'''
+
+    for name, _ in _unary_functions_from_namespace(keras.ops, "x"):
+        content += _generate_operation_class_code(name, keras.ops)
+
+    for name, _ in _unary_functions_from_namespace(keras.ops.image, "images"):
+        content += _generate_operation_class_code(name, keras.ops.image)
+
+    # Write to a temporary file first, then move to final location
+    target_path = Path(__file__).parent.parent / "keras_ops.py"
+    with tempfile.NamedTemporaryFile("w", delete=False, encoding="utf-8") as tmp_file:
+        tmp_file.write(content)
+        temp_path = Path(tmp_file.name)
+
+    # Atomic move to avoid partial writes
+    shutil.move(temp_path, target_path)
+
+    print("Done generating `keras_ops.py`.")
+
+
+if __name__ == "__main__":
+    _generate_ops_file()

zea/internal/core.py

@@ -119,11 +119,18 @@ class Object:
         return cls(**reduced_params)
 
     @classmethod
-    def merge(cls, obj1: dict, obj2: dict):
-        """Merge multiple objects and safely initialize a new object."""
+    def merge(cls, obj1: dict, obj2: dict, safe: bool = False):
+        """Merge multiple objects and safely initialize a new object.
+
+        Optionally can safely initialize the object, which removes any invalid
+        arguments.
+        """
         # TODO: support actual zea.core.Objects, now we only support dictionaries
         params = update_dictionary(obj1, obj2)
-        return cls.safe_initialize(**params)
+        if not safe:
+            return cls(**params)
+        else:
+            return cls.safe_initialize(**params)
 
     @classmethod
     def _tree_unflatten(cls, aux, children):

zea/internal/device.py

@@ -64,24 +64,43 @@ def get_gpu_memory(verbose=True):
 
     Returns:
         memory_free_values: list of available memory for each gpu in MiB.
+        Returns empty list if nvidia-smi is not available.
     """
     if not check_nvidia_smi():
         log.warning(
             "nvidia-smi is not available. Please install nvidia-utils. "
-            "Cannot retrieve GPU memory. Falling back to CPU.."
+            "Cannot retrieve GPU memory. Falling back to CPU."
         )
-        return None
+        return []
 
     def _output_to_list(x):
         return x.decode("ascii").split("\n")[:-1]
 
-    COMMAND = "nvidia-smi --query-gpu=memory.free --format=csv"
+    COMMAND = [
+        "nvidia-smi",
+        "--query-gpu=memory.free",
+        "--format=csv,noheader,nounits",
+    ]
+    # Fail-safe timeout (seconds). Override with ZEA_NVIDIA_SMI_TIMEOUT; set <=0 to disable.
+    smi_timeout = float(os.getenv("ZEA_NVIDIA_SMI_TIMEOUT", "30"))
     try:
-        memory_free_info = _output_to_list(sp.check_output(COMMAND.split()))[1:]
-    except Exception as e:
-        print(f"An error occurred: {e}")
+        if smi_timeout > 0:
+            raw = sp.check_output(COMMAND, timeout=smi_timeout)
+        else:
+            raw = sp.check_output(COMMAND)
+        memory_free_info = _output_to_list(raw)
+    except sp.TimeoutExpired:
+        log.warning(f"nvidia-smi timed out after {smi_timeout}s. Falling back to CPU.")
+        return []
+    except sp.SubprocessError as e:
+        log.warning(f"Failed to retrieve GPU memory: {e}")
+        return []
+
+    memory_free_values = [int(x) for x in memory_free_info]
 
-    memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
+    if verbose:
+        header = "GPU settings"
+        print("-" * 2 + header.center(50 - 4, "-") + "-" * 2)
 
     # only show enabled devices
     if os.environ.get("CUDA_VISIBLE_DEVICES", "") != "":
@@ -89,10 +108,12 @@ def get_gpu_memory(verbose=True):
         gpus = [int(gpu) for gpu in gpus.split(",")][: len(memory_free_values)]
         if verbose:
             # Report the number of disabled GPUs out of the total
-            num_disabled_gpus = len(memory_free_values) - len(gpus)
             num_gpus = len(memory_free_values)
-
-            print(f"{num_disabled_gpus / num_gpus} GPUs were disabled")
+            num_disabled_gpus = num_gpus - len(gpus)
+            if num_gpus > 0:
+                print(f"{num_disabled_gpus}/{num_gpus} GPUs were disabled")
+            else:
+                print("No GPUs detected by nvidia-smi.")
 
         memory_free_values = [memory_free_values[gpu] for gpu in gpus]
 
@@ -253,15 +274,11 @@ def get_device(device="auto:1", verbose=True, hide_others=True):
             os.environ["CUDA_VISIBLE_DEVICES"] = ""
         # returns None to indicate CPU
 
-    if device.lower() == "cpu":
+    if isinstance(device, str) and device.lower() == "cpu":
         return _cpu_case()
 
-    if verbose:
-        header = "GPU settings"
-        print("-" * 2 + header.center(50 - 4, "-") + "-" * 2)
-
     memory = get_gpu_memory(verbose=verbose)
-    if memory is None:  # nvidia-smi not working, fallback to CPU
+    if len(memory) == 0:  # nvidia-smi not working, fallback to CPU
         return _cpu_case()
 
     gpu_ids = list(range(len(memory)))

zea/internal/notebooks.py

@@ -0,0 +1,39 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib import animation
+
+from zea import Scan
+
+
+def animate_images(
+    images, path, scan: Scan = None, interval=100, cmap="gray", figsize=(5, 4), dpi=80
+):
+    """Helper function to animate a list of images."""
+    if interval <= 0:
+        raise ValueError("interval must be a positive integer (milliseconds).")
+    if len(images) == 0:
+        raise ValueError("images must be a non-empty sequence.")
+    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
+    if scan is not None:
+        extent = scan.extent * 1e3 if getattr(scan, "extent") is not None else None
+    else:
+        extent = None
+    im = ax.imshow(np.array(images[0]), animated=True, cmap=cmap, extent=extent)
+    ax.set_xlabel("X (mm)")
+    ax.set_ylabel("Z (mm)")
+
+    def update(frame):
+        im.set_array(np.array(images[frame]))
+        return [im]
+
+    ani = animation.FuncAnimation(
+        fig,
+        update,
+        frames=len(images),
+        blit=True,
+        interval=interval,
+    )
+    plt.close(fig)
+    fps = max(1, 1000 // interval)
+
+    ani.save(path, writer="imagemagick", fps=fps)