zea 0.0.3__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.3"
+__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/agent/selection.py

@@ -155,7 +155,9 @@ 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
@@ -176,33 +178,35 @@ class GreedyEntropy(LinesActionModel):
         # [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.
@@ -260,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)
 

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/data_format.py

@@ -468,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:

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