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

zea/internal/core.py

@@ -1,6 +1,7 @@
 """Base classes for the toolbox"""
 
 import enum
+import hashlib
 import json
 import pickle
 from copy import deepcopy
@@ -8,7 +9,8 @@ from copy import deepcopy
 import keras
 import numpy as np
 
-from zea.utils import reduce_to_signature, update_dictionary
+from zea.internal.utils import reduce_to_signature
+from zea.utils import update_dictionary
 
 CONVERT_TO_KERAS_TYPES = (np.ndarray, int, float, list, tuple, bool)
 BASE_FLOAT_PRECISION = "float32"
@@ -76,7 +78,7 @@ class Object:
             attributes.pop(
                 "_serialized", None
             )  # Remove the cached serialized attribute to avoid recursion
-            self._serialized = pickle.dumps(attributes)
+            self._serialized = serialize_elements([attributes])
         return self._serialized
 
     def __setattr__(self, name: str, value):
@@ -167,9 +169,7 @@ def _skip_to_tensor(value):
     # Skip str (because JIT does not support it)
     # Skip methods and functions
     # Skip byte strings
-    if isinstance(value, str) or callable(value) or isinstance(value, bytes):
-        return True
-    return False
+    return isinstance(value, str) or callable(value) or isinstance(value, bytes)
 
 
 def dict_to_tensor(dictionary, keep_as_is=None):
@@ -184,8 +184,9 @@ def dict_to_tensor(dictionary, keep_as_is=None):
         # Get the value from the dictionary
         value = dictionary[key]
 
-        if isinstance(value, Object):
+        if isinstance(value, Object) and hasattr(value, "to_tensor"):
             snapshot[key] = value.to_tensor(keep_as_is=keep_as_is)
+            continue
 
         # Skip certain types
         if _skip_to_tensor(value):
@@ -288,3 +289,65 @@ class ZEADecoderJSON(json.JSONDecoder):
                 obj[key] = self._MOD_TYPES_MAP[value] if value is not None else None
 
         return obj
+
+
+def serialize_elements(key_elements: list) -> str:
+    """Serialize elements of a list to a string.
+
+    Generally, uses the pickle representation of the elements.
+
+    Args:
+        key_elements (list): List of elements to serialize. Can be nested lists
+            or tuples. In this case the elements are serialized recursively.
+
+    Returns:
+        str: A serialized string representation of the elements, joined by underscores.
+    """
+
+    def _serialize(element) -> str:
+        return pickle.dumps(element).hex()
+
+    def _serialize_element(element) -> str:
+        if isinstance(element, (list, tuple)):
+            # If element is a list or tuple, serialize its elements recursively
+            element = serialize_elements(element)
+        elif isinstance(element, Object) and hasattr(element, "serialized"):
+            # Use the serialized attribute if it exists
+            element = str(element.serialized)
+        elif isinstance(element, keras.random.SeedGenerator):
+            # If element is a SeedGenerator, use the state
+            element = keras.ops.convert_to_numpy(element.state.value)
+            element = _serialize(element)
+        elif isinstance(element, dict):
+            # If element is a dictionary, sort its keys and serialize its values recursively.
+            # This is needed to ensure the internal state and ordering of the dictionary does
+            # not affect the serialization.
+            keys = list(sorted(element.keys()))
+            values = [element[k] for k in keys]
+            keys = serialize_elements(keys)
+            values = serialize_elements(values)
+            element = f"k_{keys}_v_{values}"
+        else:
+            # Otherwise, serialize the element directly
+            element = _serialize(element)
+
+        return element
+
+    serialized_elements = []
+    for element in key_elements:
+        serialized_elements.append(_serialize_element(element))
+
+    return "_".join(serialized_elements)
+
+
+def hash_elements(key_elements: list) -> str:
+    """Generate an MD5 hash of the elements.
+
+    Args:
+        key_elements (list): List of elements to serialize and hash.
+
+    Returns:
+        str: An MD5 hash of the serialized elements.
+    """
+    serialized = serialize_elements(key_elements)
+    return hashlib.md5(serialized.encode()).hexdigest()

zea/internal/device.py

@@ -377,7 +377,11 @@ def init_device(
     allow_preallocate: bool = True,
     verbose: bool = True,
 ):
-    """Selects a GPU or CPU device based on the config.
+    """Automatically selects a GPU or CPU device.
+
+    Useful to call at the start of a script to set the device for
+    tensorflow, jax or pytorch. The function will select a GPU based
+    on available memory, or fall back to CPU if no GPU is available.
 
     Args:
         backend (str): String indicating which backend to use. Can be
@@ -412,7 +416,7 @@ def init_device(
     elif backend in ["numpy", "cpu"]:
         device = "cpu"
     else:
-        raise ValueError(f"Unknown backend ({backend}) in config.")
+        raise ValueError(f"Unknown backend ({backend}).")
 
     # Early exit if device is CPU
     if device == "cpu":

zea/internal/dummy_scan.py

@@ -0,0 +1,330 @@
+"""Module to create dummy Scan objects for testing and simulation purposes."""
+
+import numpy as np
+
+from zea.beamform.delays import compute_t0_delays_focused, compute_t0_delays_planewave
+from zea.probes import Probe
+from zea.scan import Scan
+
+
+def _get_linear_probe():
+    """Returns a probe for ultrasound simulation tests."""
+    n_el = 128
+    aperture = 30e-3
+    probe_geometry = np.stack(
+        [
+            np.linspace(-aperture / 2, aperture / 2, n_el),
+            np.zeros(n_el),
+            np.zeros(n_el),
+        ],
+        axis=1,
+    )
+
+    return Probe(
+        probe_geometry=probe_geometry,
+        center_frequency=2.5e6,
+        sampling_frequency=10e6,
+    )
+
+
+def _get_phased_array_probe():
+    """Returns a probe for ultrasound simulation tests."""
+    n_el = 80
+    aperture = 20e-3
+    probe_geometry = np.stack(
+        [
+            np.linspace(-aperture / 2, aperture / 2, n_el),
+            np.zeros(n_el),
+            np.zeros(n_el),
+        ],
+        axis=1,
+    )
+
+    return Probe(
+        probe_geometry=probe_geometry,
+        center_frequency=3.12e6,
+        sampling_frequency=12.5e6,
+    )
+
+
+def _get_n_ax(ultrasound_probe):
+    """Returns the number of ax for ultrasound simulation tests based on the center
+    frequency. A probe with a higher center frequency needs more samples to cover
+    the image depth.
+    """
+    is_low_frequency_probe = ultrasound_probe.center_frequency < 4e6
+
+    if is_low_frequency_probe:
+        return 510
+
+    return 1024
+
+
+def _get_probe(kind) -> Probe:
+    if kind == "linear":
+        return _get_linear_probe()
+    elif kind == "phased_array":
+        return _get_phased_array_probe()
+    else:
+        raise ValueError(f"Unknown probe kind: {kind}")
+
+
+def _get_constant_scan_kwargs():
+    return {
+        "lens_sound_speed": 1000,
+        "lens_thickness": 1e-3,
+        "n_ch": 1,
+        "selected_transmits": "all",
+        "sound_speed": 1540.0,
+        "apply_lens_correction": False,
+        "attenuation_coef": 0.0,
+    }
+
+
+def _get_lims_and_gridsize(center_frequency, sound_speed):
+    """Returns the limits and gridsize for ultrasound simulation tests."""
+    xlims, zlims = (-20e-3, 20e-3), (0, 35e-3)
+    width, height = xlims[1] - xlims[0], zlims[1] - zlims[0]
+    wavelength = sound_speed / center_frequency
+    gridsize = (
+        int(width / (0.5 * wavelength)) + 1,
+        int(height / (0.5 * wavelength)) + 1,
+    )
+    return {"xlims": xlims, "zlims": zlims, "grid_size_x": gridsize[0], "grid_size_z": gridsize[1]}
+
+
+def _get_planewave_scan(ultrasound_probe, grid_type, **kwargs):
+    """Returns a scan for ultrasound simulation tests."""
+    constant_scan_kwargs = _get_constant_scan_kwargs()
+    n_el = ultrasound_probe.n_el
+    n_tx = 8
+
+    tx_apodizations = np.ones((n_tx, n_el)) * np.hanning(n_el)[None]
+    probe_geometry = ultrasound_probe.probe_geometry
+
+    angles = np.linspace(10, -10, n_tx) * np.pi / 180
+
+    sound_speed = constant_scan_kwargs["sound_speed"]
+    focus_distances = np.ones(n_tx) * np.inf
+    t0_delays = compute_t0_delays_planewave(
+        probe_geometry=probe_geometry, polar_angles=angles, sound_speed=sound_speed
+    )
+
+    return Scan(
+        n_tx=n_tx,
+        n_el=n_el,
+        center_frequency=ultrasound_probe.center_frequency,
+        sampling_frequency=ultrasound_probe.sampling_frequency,
+        probe_geometry=probe_geometry,
+        t0_delays=t0_delays,
+        tx_apodizations=tx_apodizations,
+        element_width=np.linalg.norm(probe_geometry[1] - probe_geometry[0]),
+        focus_distances=focus_distances,
+        polar_angles=angles,
+        initial_times=np.ones(n_tx) * 1e-6,
+        n_ax=_get_n_ax(ultrasound_probe),
+        grid_type=grid_type,
+        **_get_lims_and_gridsize(ultrasound_probe.center_frequency, sound_speed),
+        **constant_scan_kwargs,
+        **kwargs,
+    )
+
+
+def _get_multistatic_scan(ultrasound_probe, grid_type, **kwargs):
+    n_el = ultrasound_probe.n_el
+    n_tx = 8
+
+    tx_apodizations = np.zeros((n_tx, n_el))
+    for n, idx in enumerate(np.linspace(0, n_el - 1, n_tx, dtype=int)):
+        tx_apodizations[n, idx] = 1
+    probe_geometry = ultrasound_probe.probe_geometry
+
+    focus_distances = np.zeros(n_tx)
+    t0_delays = np.zeros((n_tx, n_el))
+
+    constant_scan_kwargs = _get_constant_scan_kwargs()
+
+    return Scan(
+        n_tx=n_tx,
+        n_el=n_el,
+        center_frequency=ultrasound_probe.center_frequency,
+        sampling_frequency=ultrasound_probe.sampling_frequency,
+        probe_geometry=probe_geometry,
+        t0_delays=t0_delays,
+        tx_apodizations=tx_apodizations,
+        element_width=np.linalg.norm(probe_geometry[1] - probe_geometry[0]),
+        focus_distances=focus_distances,
+        polar_angles=np.zeros(n_tx),
+        initial_times=np.ones(n_tx) * 1e-6,
+        n_ax=_get_n_ax(ultrasound_probe),
+        grid_type=grid_type,
+        **_get_lims_and_gridsize(
+            ultrasound_probe.center_frequency, constant_scan_kwargs["sound_speed"]
+        ),
+        **constant_scan_kwargs,
+        **kwargs,
+    )
+
+
+def _get_diverging_scan(ultrasound_probe, grid_type, **kwargs):
+    """Returns a scan for ultrasound simulation tests."""
+    constant_scan_kwargs = _get_constant_scan_kwargs()
+    n_el = ultrasound_probe.n_el
+    n_tx = 8
+
+    tx_apodizations = np.ones((n_tx, n_el)) * np.hanning(n_el)[None]
+
+    angles = np.linspace(10, -10, n_tx) * np.pi / 180
+
+    sound_speed = constant_scan_kwargs["sound_speed"]
+    focus_distances = np.ones(n_tx) * -15e-3
+    t0_delays = compute_t0_delays_focused(
+        origins=np.zeros((n_tx, 3)),
+        focus_distances=focus_distances,
+        probe_geometry=ultrasound_probe.probe_geometry,
+        polar_angles=angles,
+        sound_speed=sound_speed,
+    )
+    element_width = np.linalg.norm(
+        ultrasound_probe.probe_geometry[1] - ultrasound_probe.probe_geometry[0]
+    )
+
+    return Scan(
+        n_tx=n_tx,
+        n_el=n_el,
+        center_frequency=ultrasound_probe.center_frequency,
+        sampling_frequency=ultrasound_probe.sampling_frequency,
+        probe_geometry=ultrasound_probe.probe_geometry,
+        t0_delays=t0_delays,
+        tx_apodizations=tx_apodizations,
+        element_width=element_width,
+        focus_distances=focus_distances,
+        polar_angles=angles,
+        initial_times=np.ones(n_tx) * 1e-6,
+        n_ax=_get_n_ax(ultrasound_probe),
+        grid_type=grid_type,
+        **_get_lims_and_gridsize(ultrasound_probe.center_frequency, sound_speed),
+        **constant_scan_kwargs,
+        **kwargs,
+    )
+
+
+def _get_focused_scan(ultrasound_probe, grid_type, **kwargs):
+    """Returns a scan for ultrasound simulation tests."""
+    constant_scan_kwargs = _get_constant_scan_kwargs()
+    n_el = ultrasound_probe.n_el
+    n_tx = 8
+
+    tx_apodizations = np.ones((n_tx, n_el)) * np.hanning(n_el)[None]
+
+    angles = np.linspace(30, -30, n_tx) * np.pi / 180
+
+    sound_speed = constant_scan_kwargs["sound_speed"]
+    focus_distances = np.ones(n_tx) * 15e-3
+    t0_delays = compute_t0_delays_focused(
+        origins=np.zeros((n_tx, 3)),
+        focus_distances=focus_distances,
+        probe_geometry=ultrasound_probe.probe_geometry,
+        polar_angles=angles,
+        sound_speed=sound_speed,
+    )
+    element_width = np.linalg.norm(
+        ultrasound_probe.probe_geometry[1] - ultrasound_probe.probe_geometry[0]
+    )
+
+    return Scan(
+        n_tx=n_tx,
+        n_el=n_el,
+        center_frequency=ultrasound_probe.center_frequency,
+        sampling_frequency=ultrasound_probe.sampling_frequency,
+        probe_geometry=ultrasound_probe.probe_geometry,
+        t0_delays=t0_delays,
+        tx_apodizations=tx_apodizations,
+        element_width=element_width,
+        focus_distances=focus_distances,
+        polar_angles=angles,
+        initial_times=np.ones(n_tx) * 1e-6,
+        n_ax=_get_n_ax(ultrasound_probe),
+        grid_type=grid_type,
+        **_get_lims_and_gridsize(ultrasound_probe.center_frequency, sound_speed),
+        **constant_scan_kwargs,
+        **kwargs,
+    )
+
+
+def _get_linescan_scan(ultrasound_probe, grid_type, **kwargs):
+    """Returns a scan for ultrasound simulation tests."""
+    constant_scan_kwargs = _get_constant_scan_kwargs()
+    n_el = ultrasound_probe.n_el
+    n_tx = 8
+
+    center_elements = np.linspace(0, n_el + 1, n_tx + 2, dtype=int)
+    center_elements = center_elements[1:-1]
+    tx_apodizations = np.zeros((n_tx, n_el))
+    aperture_size_elements = 24
+
+    # Define subapertures
+    origins = []
+    for n, idx in enumerate(center_elements):
+        el0 = np.clip(idx - aperture_size_elements // 2, 0, n_el)
+        el1 = np.clip(idx + aperture_size_elements // 2, 0, n_el)
+        tx_apodizations[n, el0:el1] = np.hanning(el1 - el0)[None]
+        origins.append(ultrasound_probe.probe_geometry[idx])
+    origins = np.stack(origins, axis=0)
+
+    # All angles should be zero because each line fires straight ahead
+    angles = np.zeros(n_tx)
+
+    sound_speed = constant_scan_kwargs["sound_speed"]
+
+    focus_distances = np.ones(n_tx) * 15e-3
+    t0_delays = compute_t0_delays_focused(
+        origins=origins,
+        focus_distances=focus_distances,
+        probe_geometry=ultrasound_probe.probe_geometry,
+        polar_angles=angles,
+        sound_speed=sound_speed,
+    )
+    element_width = np.linalg.norm(
+        ultrasound_probe.probe_geometry[1] - ultrasound_probe.probe_geometry[0]
+    )
+
+    return Scan(
+        n_tx=n_tx,
+        n_el=n_el,
+        center_frequency=ultrasound_probe.center_frequency,
+        sampling_frequency=ultrasound_probe.sampling_frequency,
+        probe_geometry=ultrasound_probe.probe_geometry,
+        t0_delays=t0_delays,
+        tx_apodizations=tx_apodizations,
+        element_width=element_width,
+        focus_distances=focus_distances,
+        polar_angles=angles,
+        initial_times=np.ones(n_tx) * 1e-6,
+        n_ax=_get_n_ax(ultrasound_probe),
+        grid_type=grid_type,
+        **_get_lims_and_gridsize(ultrasound_probe.center_frequency, sound_speed),
+        **constant_scan_kwargs,
+        **kwargs,
+    )
+
+
+def _get_scan(ultrasound_probe, kind, grid_type="cartesian", **kwargs) -> Scan:
+    if kind == "planewave":
+        return _get_planewave_scan(ultrasound_probe, grid_type, **kwargs)
+    elif kind == "multistatic":
+        return _get_multistatic_scan(ultrasound_probe, grid_type, **kwargs)
+    elif kind == "diverging":
+        return _get_diverging_scan(ultrasound_probe, grid_type, **kwargs)
+    elif kind == "focused":
+        return _get_focused_scan(ultrasound_probe, grid_type, **kwargs)
+    elif kind == "linescan":
+        return _get_linescan_scan(ultrasound_probe, grid_type, **kwargs)
+    else:
+        raise ValueError(f"Unknown scan kind: {kind}")
+
+
+def get_scan(kind="planewave", probe_kind="linear", grid_type="cartesian", **kwargs) -> Scan:
+    """Returns a scan for ultrasound simulation tests."""
+    ultrasound_probe = _get_probe(probe_kind)
+    return _get_scan(ultrasound_probe, kind, grid_type, **kwargs)

zea/internal/operators.py

@@ -6,6 +6,7 @@ Handles task-dependent operations (A) and noises (n) to simulate a measurement y
 
 import abc
 
+import numpy as np
 from keras import ops
 
 from zea.internal.core import Object
@@ -26,13 +27,16 @@ class Operator(abc.ABC, Object):
 
     """
 
-    sigma = 0.0
-
     @abc.abstractmethod
     def forward(self, data, *args, **kwargs):
         """Implements the forward operator A: x -> y."""
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def transpose(self, data, *args, **kwargs):
+        """Implements the transpose (or adjoint) of the operator A^T: y -> x."""
+        raise NotImplementedError
+
     @abc.abstractmethod
     def __str__(self):
         """String representation of the operator."""
@@ -78,5 +82,113 @@ class InpaintingOperator(Operator):
         # return self.mask * data
         return ops.where(mask, data, self.min_val)
 
+    def transpose(self, data, mask):
+        # masking operation is diagonal --> A.T = A
+        return self.forward(data, mask)
+
     def __str__(self):
         return "y = Ax + n, where A = I * M"
+
+
+@operator_registry(name="fourier_blur")
+class FourierBlurOperator(Operator):
+    """Fourier-domain blurring operator class.
+
+    Applies blurring by masking high frequencies in the Fourier domain.
+
+    Formally defined as:
+        y = F^(-1)(M * F(x))
+
+    where F is the FFT, F^(-1) is the inverse FFT and M is the frequency mask.
+    """
+
+    def __init__(self, shape, cutoff_freq=0.5, smooth=True, **kwargs):
+        """Initialize the Fourier blur operator.
+
+        Args:
+            shape: Shape of the input data (H, W), (H, W, C), or (B, H, W, C).
+            cutoff_freq: Cutoff frequency as fraction of Nyquist frequency (0.0 to 1.0).
+            smooth: If True, use Gaussian rolloff; otherwise use hard cutoff.
+            **kwargs: Additional arguments.
+        """
+        super().__init__(**kwargs)
+        self.cutoff_freq = cutoff_freq
+        self.shape = shape
+
+        # Precompute frequency mask
+        self.freq_mask = self.make_lowpass_mask(shape=shape, cutoff_freq=cutoff_freq, smooth=smooth)
+
+    def make_lowpass_mask(self, shape, cutoff_freq=0.1, smooth=True):
+        """
+        Create a low-pass Fourier mask of given shape.
+        cutoff: relative frequency radius (0 < cutoff < 0.5)
+        smooth: if True, use Gaussian rolloff
+        """
+        # Accept (H, W), (H, W, C) or (B, H, W, C)
+        if len(shape) == 2:
+            H, W = shape
+        elif len(shape) >= 3:
+            H, W = shape[-3], shape[-2]
+        else:
+            raise ValueError(f"Invalid shape {shape}. Expected (H, W), (H, W, C), or (B, H, W, C).")
+        fy = np.fft.fftfreq(H)
+        fx = np.fft.fftfreq(W)
+        FX, FY = np.meshgrid(fx, fy)
+        R = np.sqrt(FX**2 + FY**2)
+
+        if smooth:
+            sigma = cutoff_freq / np.sqrt(2 * np.log(2))
+            mask = np.exp(-(R**2) / (2 * sigma**2))
+        else:
+            mask = (R < cutoff_freq).astype(np.float32)
+
+        # Shift DC to top-left to match fft2 conventions
+        # mask = np.fft.ifftshift(mask)
+        return ops.convert_to_tensor(mask)
+
+    def forward(self, data):
+        """Apply Fourier-domain blurring.
+
+        Args:
+            data: Input tensor of shape (B, H, W, C)
+
+        Returns:
+            Blurred data tensor.
+        """
+        # Convert to float32 for FFT
+        data_real = ops.cast(data, "float32")
+        data_imag = ops.zeros_like(data_real)
+
+        # fft2 calculates the 2D FFT on the last two dims, so we want
+        # H, W to be at the end
+        data_real = ops.transpose(data_real, (0, 3, 1, 2))
+        data_imag = ops.transpose(data_imag, (0, 3, 1, 2))
+
+        # Apply FFT - expects tuple (real, imag), returns tuple (real, imag)
+        fft_real, fft_imag = ops.fft2((data_real, data_imag))
+
+        # Apply frequency mask to both real and imaginary parts
+        mask_real = ops.real(self.freq_mask)  # Extract real part of complex mask
+        masked_fft_real = fft_real * mask_real
+        masked_fft_imag = fft_imag * mask_real
+
+        # Apply inverse FFT
+        blurred_real, blurred_imag = ops.ifft2((masked_fft_real, masked_fft_imag))
+
+        # transpose back to original shape
+        blurred_real = ops.transpose(blurred_real, (0, 2, 3, 1))
+
+        # Take real part (imaginary should be ~0 for real input)
+        blurred_data = ops.cast(blurred_real, data.dtype)
+
+        return blurred_data
+
+    def transpose(self, data):
+        """
+        transpose = forward because A^* (F^{-1} M F)^* = (F^* M^* (F^{-1})^*) = F^{-1} M F = A
+        i.e. this is a self-adjoint operator.
+        """
+        return self.forward(data)
+
+    def __str__(self):
+        return f"y = F^(-1)(M * F(x)) filter at {self.cutoff_freq}"