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

zea/__init__.py

@@ -7,36 +7,71 @@ from . import log
 
 # dynamically add __version__ attribute (see pyproject.toml)
 # __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.6"
+__version__ = "0.0.7"
 
 
 def _bootstrap_backend():
     """Setup function to initialize the zea package."""
 
     def _check_backend_installed():
-        """Assert that at least one ML backend (torch, tensorflow, jax) is installed.
-        If not, raise an AssertionError with a helpful install message.
-        """
-
-        ml_backends = ["torch", "tensorflow", "jax"]
-        for backend in ml_backends:
-            if importlib.util.find_spec(backend) is not None:
-                return
+        """Verify that the required ML backend is installed.
 
-        backend_env = os.environ.get("KERAS_BACKEND", "numpy")
-        install_guide_urls = {
+        Raises ImportError if:
+        1. No ML backend (torch, tensorflow, jax) is installed
+        2. KERAS_BACKEND points to a backend that is not installed
+        """
+        ML_BACKENDS = ["torch", "tensorflow", "jax"]
+        INSTALL_URLS = {
             "torch": "https://pytorch.org/get-started/locally/",
             "tensorflow": "https://www.tensorflow.org/install",
             "jax": "https://docs.jax.dev/en/latest/installation.html",
         }
-        guide_url = install_guide_urls.get(backend_env, "https://keras.io/getting_started/")
-        raise ImportError(
-            "No ML backend (torch, tensorflow, jax) installed in current environment. "
-            f"Please install at least one ML backend before importing {__package__} or "
-            f"any other library. Current KERAS_BACKEND is set to '{backend_env}', "
-            f"please install it first, see: {guide_url}. One simple alternative is to "
-            f"install with default backend: `pip install {__package__}[jax]`."
-        )
+        KERAS_DEFAULT_BACKEND = "tensorflow"
+        DOCS_URL = "https://zea.readthedocs.io/en/latest/installation.html"
+
+        # Determine which backend Keras will try to use
+        backend_env = os.environ.get("KERAS_BACKEND")
+        effective_backend = backend_env or KERAS_DEFAULT_BACKEND
+
+        # Find all installed ML backends
+        installed_backends = [
+            backend for backend in ML_BACKENDS if importlib.util.find_spec(backend) is not None
+        ]
+
+        # Error if no backends are installed
+        if not installed_backends:
+            if backend_env:
+                backend_status = f"KERAS_BACKEND is set to '{backend_env}'"
+            else:
+                backend_status = f"KERAS_BACKEND is not set (defaults to '{KERAS_DEFAULT_BACKEND}')"
+            install_url = INSTALL_URLS.get(effective_backend, "https://keras.io/getting_started/")
+            raise ImportError(
+                f"No ML backend (torch, tensorflow, jax) installed in current "
+                f"environment. Please install at least one ML backend before importing "
+                f"{__package__}. {backend_status}, please install it first, see: "
+                f"{install_url}. One simple alternative is to install with default "
+                f"backend: `pip install {__package__}[jax]`. For more information, "
+                f"see: {DOCS_URL}"
+            )
+
+        # Error if the effective backend is not installed
+        # (skip numpy which doesn't need installation)
+        if effective_backend not in ["numpy"] and effective_backend not in installed_backends:
+            if backend_env:
+                backend_status = f"KERAS_BACKEND environment variable is set to '{backend_env}'"
+            else:
+                backend_status = (
+                    f"KERAS_BACKEND is not set, which defaults to '{KERAS_DEFAULT_BACKEND}'"
+                )
+            install_url = INSTALL_URLS.get(effective_backend, "https://keras.io/getting_started/")
+            raise ImportError(
+                f"{backend_status}, but this backend is not installed. "
+                f"Installed backends: {', '.join(installed_backends)}. "
+                f"Please either install '{effective_backend}' (see: {install_url}) "
+                f"or set KERAS_BACKEND to one of the installed backends "
+                f"(e.g., export KERAS_BACKEND={installed_backends[0]}). "
+                f"For more information, see: {DOCS_URL}"
+            )
 
     _check_backend_installed()
 

zea/agent/__init__.py

@@ -7,21 +7,21 @@ For a practical example, see :doc:`../notebooks/agent/agent_example`.
 Example usage
 ^^^^^^^^^^^^^
 
-.. code-block:: python
+.. doctest::
 
-    import zea
-    import numpy as np
+    >>> import zea
+    >>> import numpy as np
 
-    agent = zea.agent.selection.GreedyEntropy(
-        n_actions=7,
-        n_possible_actions=112,
-        img_width=112,
-        img_height=112,
-    )
+    >>> agent = zea.agent.selection.GreedyEntropy(
+    ...     n_actions=7,
+    ...     n_possible_actions=112,
+    ...     img_width=112,
+    ...     img_height=112,
+    ... )
 
-    # (batch, samples, height, width)
-    particles = np.random.rand(1, 10, 112, 112)
-    lines, mask = agent.sample(particles)
+    >>> # (batch, samples, height, width)
+    >>> particles = np.random.rand(1, 10, 112, 112)
+    >>> lines, mask = agent.sample(particles)  # doctest: +SKIP
 """
 
 from . import masks, selection

zea/agent/masks.py

@@ -91,7 +91,8 @@ def random_uniform_lines(
 
 def _assert_equal_spacing(n_actions, n_possible_actions):
     assert n_possible_actions % n_actions == 0, (
-        "Number of actions must divide evenly into possible actions to use equispaced sampling."
+        "Number of actions must divide evenly into possible actions to use equispaced sampling. "
+        "If you do not care about equal spacing, set `assert_equal_spacing=False`."
     )
 
 

zea/backend/tensorflow/dataloader.py

@@ -12,7 +12,8 @@ from keras.src.trainers.data_adapters import TFDatasetAdapter
 
 from zea.data.dataloader import H5Generator
 from zea.data.layers import Resizer
-from zea.utils import find_methods_with_return_type, 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 safe_vectorize
+from zea.tensor_ops import vmap
 
 
 def fnum_window_fn_rect(normalized_angle):
@@ -57,10 +57,11 @@ def tof_correction(
     initial_times,
     sampling_frequency,
     demodulation_frequency,
-    fnum,
-    angles,
+    f_number,
+    polar_angles,
     focus_distances,
-    apply_phase_rotation=False,
+    t_peak,
+    tx_waveform_indices,
     apply_lens_correction=False,
     lens_thickness=1e-3,
     lens_sound_speed=1000,
@@ -73,21 +74,19 @@ def tof_correction(
         flatgrid (ops.Tensor): Pixel locations x, y, z of shape `(n_pix, 3)`
         t0_delays (ops.Tensor): Times at which the elements fire shifted such
             that the first element fires at t=0 of shape `(n_tx, n_el)`
-        tx_apodizations (ops.Tensor): Transmit apodizations of shape
-            `(n_tx, n_el)`
+        tx_apodizations (ops.Tensor): Transmit apodizations of shape `(n_tx, n_el)`
         sound_speed (float): Speed-of-sound.
-        probe_geometry (ops.Tensor): Element positions x, y, z of shape
-        (num_samples, 3)
-        initial_times (ops.Tensor): Time-ofsampling_frequencyet per transmission of shape
-            `(n_tx,)`.
+        probe_geometry (ops.Tensor): Element positions x, y, z of shape (n_el, 3)
+        initial_times (Tensor): The probe transmit time offsets of shape `(n_tx,)`.
         sampling_frequency (float): Sampling frequency.
         demodulation_frequency (float): Demodulation frequency.
-        fnum (int, optional): Focus number. Defaults to 1.
-        angles (ops.Tensor): The angles of the plane waves in radians of shape
-            `(n_tx,)`
+        f_number (float): Focus number (ratio of focal depth to aperture size).
+        polar_angles (ops.Tensor): The angles of the waves in radians of shape `(n_tx,)`
         focus_distances (ops.Tensor): The focus distance of shape `(n_tx,)`
-        apply_phase_rotation (bool, optional): Whether to apply phase rotation to
-            time-of-flights. Defaults to False.
+        t_peak (ops.Tensor): Time of the peak of the pulse in seconds.
+            Shape `(n_waveforms,)`.
+        tx_waveform_indices (ops.Tensor): The indices of the waveform used for each
+            transmit of shape `(n_tx,)`.
         apply_lens_correction (bool, optional): Whether to apply lens correction to
             time-of-flights. This makes it slower, but more accurate in the near-field.
             Defaults to False.
@@ -101,14 +100,12 @@ def tof_correction(
 
     Returns:
         (ops.Tensor): time-of-flight corrected data
-        with shape: `(n_tx, n_pix, n_el, num_rf_iq_channels)`.
+        with shape: `(n_tx, n_pix, n_el, n_ch)`.
     """
 
     assert len(data.shape) == 4, (
         "The input data should have 4 dimensions, "
-        f"namely num_transmits, num_elements, num_samples, "
-        f"num_rf_iq_channels, got {len(data.shape)} dimensions: ."
-        f"{data.shape}"
+        f"namely n_tx, n_ax, n_el, n_ch, got {len(data.shape)} dimensions: {data.shape}"
     )
 
     n_tx, n_ax, n_el, _ = ops.shape(data)
@@ -135,19 +132,33 @@ def tof_correction(
         n_tx,
         n_el,
         focus_distances,
-        angles,
+        polar_angles,
+        t_peak=t_peak,
+        tx_waveform_indices=tx_waveform_indices,
         lens_thickness=lens_thickness,
         lens_sound_speed=lens_sound_speed,
     )
 
     n_pix = ops.shape(flatgrid)[0]
     mask = ops.cond(
-        fnum == 0,
+        f_number == 0,
         lambda: ops.ones((n_pix, n_el, 1)),
-        lambda: fnumber_mask(flatgrid, probe_geometry, fnum, fnum_window_fn=fnum_window_fn),
+        lambda: fnumber_mask(flatgrid, probe_geometry, f_number, fnum_window_fn=fnum_window_fn),
     )
 
     def _apply_delays(data_tx, txdel):
+        """Applies the delays to TOF correct a single transmit.
+
+        Args:
+            data_tx (ops.Tensor): The RF/IQ data for a single transmit of shape
+                `(n_ax, n_el, n_ch)`.
+            txdel (ops.Tensor): The transmit delays for a single transmit in samples
+                (not in seconds) of shape `(n_pix, 1)`.
+
+        Returns:
+            ops.Tensor: The time-of-flight corrected data of shape
+            `(n_pix, n_el, n_ch)`.
+        """
         # data_tx is of shape (num_elements, num_samples, 1 or 2)
 
         # Take receive delays and add the transmit delays for this transmit
@@ -164,22 +175,22 @@ def tof_correction(
         # Apply the mask
         tof_tx = tof_tx * mask
 
-        # Phase correction
+        # Apply phase rotation if using IQ data
+        # This is needed because interpolating the IQ data without phase rotation
+        # is not equivalent to interpolating the RF data and then IQ demodulating
+        # See the docstring from complex_rotate for more details
+        apply_phase_rotation = data_tx.shape[-1] == 2
         if apply_phase_rotation:
-            tshift = delays[:, :] / sampling_frequency
-            tdemod = flatgrid[:, None, 2] * 2 / sound_speed
-            theta = 2 * np.pi * demodulation_frequency * (tshift - tdemod)
-            tof_tx = _complex_rotate(tof_tx, theta)
+            total_delay_seconds = delays[:, :] / sampling_frequency
+            theta = 2 * np.pi * demodulation_frequency * total_delay_seconds
+            tof_tx = complex_rotate(tof_tx, theta)
         return tof_tx
 
     # Reshape to (n_tx, n_pix, 1)
     txdel = ops.moveaxis(txdel, 1, 0)
     txdel = txdel[..., None]
 
-    return safe_vectorize(
-        _apply_delays,
-        signature="(n_samples,n_el,n_ch),(n_pix,1)->(n_pix,n_el,n_ch)",
-    )(data, txdel)
+    return vmap(_apply_delays)(data, txdel)
 
 
 def calculate_delays(
@@ -194,6 +205,8 @@ def calculate_delays(
     n_el,
     focus_distances,
     polar_angles,
+    t_peak,
+    tx_waveform_indices,
     **kwargs,
 ):
     """Calculates the delays in samples to every pixel in the grid.
@@ -225,12 +238,18 @@ def calculate_delays(
             assume plane wave transmission.
         polar_angles (Tensor): The polar angles of the plane waves in radians
             of shape `(n_tx,)`.
+        t_peak (Tensor): Time of the peak of the pulse in seconds of shape
+            `(n_waveforms,)`.
+        tx_waveform_indices (Tensor): The indices of the waveform used for each
+            transmit of shape `(n_tx,)`.
+
 
     Returns:
-        transmit_delays (Tensor): The tensor of transmit delays to every pixel,
-            shape `(n_pix, n_tx)`.
+        transmit_delays (Tensor): The tensor of transmit delays to every pixel
+            in samples (not in seconds), of shape `(n_pix, n_tx)`.
         receive_delays (Tensor): The tensor of receive delays from every pixel
-            back to the transducer element, shape `(n_pix, n_el)`.
+            back to the transducer element in samples (not in seconds), of shape
+            `(n_pix, n_el)`.
     """
 
     def _tx_distances(polar_angles, t0_delays, tx_apodizations, focus_distances):
@@ -244,10 +263,7 @@ def calculate_delays(
             sound_speed,
         )
 
-    tx_distances = safe_vectorize(
-        _tx_distances,
-        signature="(),(n_el),(n_el),()->(n_pix)",
-    )(polar_angles, t0_delays, tx_apodizations, focus_distances)
+    tx_distances = vmap(_tx_distances)(polar_angles, t0_delays, tx_apodizations, focus_distances)
     tx_distances = ops.transpose(tx_distances, (1, 0))
     # tx_distances shape is now (n_pix, n_tx)
 
@@ -255,14 +271,18 @@ def calculate_delays(
     def _rx_distances(probe_geometry):
         return distance_Rx(grid, probe_geometry)
 
-    rx_distances = safe_vectorize(_rx_distances, signature="(3)->(n_pix)")(probe_geometry)
+    rx_distances = vmap(_rx_distances)(probe_geometry)
     rx_distances = ops.transpose(rx_distances, (1, 0))
     # rx_distances shape is now (n_pix, n_el)
 
     # Compute the delays [in samples] from the distances
     # The units here are ([m]/[m/s]-[s])*[1/s] resulting in a unitless quantity
     # TODO: Add pulse width to transmit delays
-    tx_delays = (tx_distances / sound_speed - initial_times[None]) * sampling_frequency
+    tx_delays = (
+        tx_distances / sound_speed
+        - initial_times[None]
+        + ops.take(t_peak, tx_waveform_indices)[None]
+    ) * sampling_frequency
     rx_delays = (rx_distances / sound_speed) * sampling_frequency
 
     return tx_delays, rx_delays
@@ -288,11 +308,11 @@ def apply_delays(data, delays, clip_min: int = -1, clip_max: int = -1):
 
     Returns:
         ops.Tensor: The samples received by each transducer element corresponding to the
-            reflections of each pixel in the image of shape `(n_el, n_pix, n_ch)`.
+            reflections of each pixel in the image of shape `(n_pix, n_el, n_ch)`.
     """
 
     # Add a dummy channel dimension to the delays tensor to ensure it has the
-    # same number of dimensions as the data. The new shape is (1, n_el, n_pix)
+    # same number of dimensions as the data. The new shape is (n_pix, n_el, 1)
     delays = delays[..., None]
 
     # Get the integer values above and below the exact delay values
@@ -318,7 +338,7 @@ def apply_delays(data, delays, clip_min: int = -1, clip_max: int = -1):
 
     # Gather pixel values
     # Here we extract for each transducer element the sample containing the
-    # reflection from each pixel. These are of shape `(n_el, n_pix, n_ch)`.
+    # reflection from each pixel. These are of shape `(n_pix, n_el, n_ch)`.
     data0 = ops.take_along_axis(data, d0, 0)
     data1 = ops.take_along_axis(data, d1, 0)
 
@@ -332,7 +352,7 @@ def apply_delays(data, delays, clip_min: int = -1, clip_max: int = -1):
     return reflection_samples
 
 
-def _complex_rotate(iq, theta):
+def complex_rotate(iq, theta):
     """Performs a simple phase rotation of I and Q component.
 
     Args:
@@ -341,11 +361,41 @@ def _complex_rotate(iq, theta):
 
     Returns:
         Tensor: The rotated tensor of shape `(..., 2)`.
+
+    .. dropdown:: Explanation
+
+        The IQ data is related to the RF data as follows:
+
+        .. math::
+
+            x(t) &= I(t)\\cos(\\omega_c t) + Q(t)\\cos(\\omega_c t + \\pi/2)\\\\
+            &= I(t)\\cos(\\omega_c t) - Q(t)\\sin(\\omega_c t)
+
+
+        If we want to delay the RF data `x(t)` by `Δt` we can substitute in
+        :math:`t=t+\\Delta t`. We also define :math:`I'(t) = I(t + \\Delta t)`,
+        :math:`Q'(t) = Q(t + \\Delta t)`, and :math:`\\theta=\\omega_c\\Delta t`.
+        This gives us:
+
+        .. math::
+
+            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)\\\\
+            &- \\overbrace{(Q'(t)\\cos(\\theta)
+            + I'(t)\\sin(\\theta))}^{Q_\\Delta(t)} \\sin(\\omega_c t)
+
+        This means that to correctly interpolate the IQ data to the new components
+        :math:`I_\\Delta(t)` and :math:`Q_\\Delta(t)`, it is not sufficient to just
+        interpolate the I- and Q-channels independently. We also need to rotate the
+        I- and Q-channels by the angle :math:`\\theta`. This function performs this
+        rotation.
     """
-    # assert iq.shape[-1] == 2, (
-    #     "The last dimension of the input tensor should be 2, "
-    #     f"got {iq.shape[-1]} dimensions and shape {iq.shape}."
-    # )
+    assert iq.shape[-1] == 2, (
+        "The last dimension of the input tensor should be 2, "
+        f"got {iq.shape[-1]} dimensions and shape {iq.shape}."
+    )
     # Select i and q channels
     i = iq[..., 0]
     q = iq[..., 1]
@@ -485,8 +535,8 @@ def fnumber_mask(flatgrid, probe_geometry, f_number, fnum_window_fn):
 
     alpha = ops.arccos(grid_relative_to_probe_z)
 
-    # The f-number is fnum = z/aperture = 1/(2 * tan(alpha))
-    # Rearranging gives us alpha = arctan(1/(2 * fnum))
+    # The f-number is f_number = z/aperture = 1/(2 * tan(alpha))
+    # Rearranging gives us alpha = arctan(1/(2 * f_number))
     # We can use this to compute the maximum angle alpha that is allowed
     max_alpha = ops.arctan(1 / (2 * f_number + keras.backend.epsilon()))
 

zea/beamform/lens_correction.py

@@ -15,6 +15,8 @@ def calculate_lens_corrected_delays(
     n_el,
     focus_distances,
     polar_angles,
+    t_peak,
+    tx_waveform_indices,
     lens_sound_speed=1000,
     lens_thickness=1e-3,
     n_iter=2,
@@ -33,6 +35,9 @@ def calculate_lens_corrected_delays(
         n_el (int): The number of elements.
         focus_distances (ndarray): The focus distances of shape (n_tx,).
         polar_angles (ndarray): The polar angles of shape (n_tx,).
+        t_peak (ndarray): The time of the peak of the pulse in seconds of shape (n_waveforms,).
+        tx_waveform_indices (ndarray): The indices of the waveforms used for each transmit
+            of shape (n_tx,).
         lens_sound_speed (float): The speed of sound in the lens in m/s.
         lens_thickness (float): The thickness of the lens in meters.
         n_iter (int): The number of iterations to run the Newton-Raphson method.
@@ -59,9 +64,11 @@ def calculate_lens_corrected_delays(
         # Add a large offset to elements that are not used in the transmit to
         # diqualify them from being the closest element
         apod_offset = ops.where(tx_apodizations[tx] == 0, 10.0, 0)
-        tx_min = ops.min(rx_delays + t0_delays[tx] + apod_offset, axis=-1) + initial_times[tx]
+        tx_min = ops.min(rx_delays + t0_delays[tx] + apod_offset, axis=-1) - initial_times[tx]
         tx_delays.append(tx_min)
-    tx_delays = ops.stack(tx_delays, axis=-1)
+    tx_delays = ops.stack(tx_delays, axis=-1) + ops.take(t_peak, tx_waveform_indices)[None]
+    tx_delays = ops.nan_to_num(tx_delays, nan=0.0, posinf=0.0, neginf=0.0)
+    rx_delays = ops.nan_to_num(rx_delays, nan=0.0, posinf=0.0, neginf=0.0)
 
     tx_delays *= sampling_frequency
     rx_delays *= sampling_frequency

zea/beamform/pfield.py

@@ -43,7 +43,7 @@ def compute_pfield(
     grid,
     t0_delays,
     frequency_step=4,
-    db_thresh=-1,
+    db_thresh=-1.0,
     downsample=10,
     downmix=4,
     alpha=1,
@@ -65,7 +65,7 @@ def compute_pfield(
         t0_delays (array): Transmit delays for each transmit event.
         frequency_step (int, optional): Frequency step. Default is 4.
             Higher is faster but less accurate.
-        db_thresh (int, optional): dB threshold. Default is -1.
+        db_thresh (float, optional): dB threshold. Default is -1.0
             Higher is faster but less accurate.
         downsample (int, optional): Downsample the grid for faster computation.
             Default is 10. Higher is faster but less accurate.
@@ -85,6 +85,13 @@ def compute_pfield(
     # medium params
     alpha_db = 0  # currently we ignore attenuation in the compounding
 
+    # cast to float32
+    sound_speed = ops.cast(sound_speed, "float32")
+    center_frequency = ops.cast(center_frequency, "float32")
+    bandwidth_percent = ops.cast(bandwidth_percent, "float32")
+    alpha_db = ops.cast(alpha_db, "float32")
+    db_thresh = ops.cast(db_thresh, "float32")
+
     # probe params
     center_frequency = center_frequency / downmix  # downmixing the frequency
 

zea/config.py

@@ -16,23 +16,30 @@ Features
 Example Usage
 ^^^^^^^^^^^^^
 
-.. code-block:: python
+.. doctest::
 
-    from zea import Config
+    >>> from zea import Config
 
-    # Load from YAML
-    config = Config.from_yaml("config.yaml")
-    # Load from HuggingFace Hub
-    config = Config.from_hf("zea/diffusion-echonet-dynamic", "train_config.yaml")
+    >>> # Load from YAML
+    >>> config = Config.from_yaml("../configs/config_echonet.yaml")
+    >>> # Load from HuggingFace Hub
+    >>> config = Config.from_hf("zeahub/configs", "config_picmus_rf.yaml", repo_type="dataset")
 
-    # Access attributes with dot notation
-    print(config.model.name)
+    >>> # Access attributes with dot notation
+    >>> print(config.data.dtype)
+    raw_data
 
-    # Update recursively
-    config.update_recursive({"model": {"name": "new_model"}})
+    >>> # Update recursively
+    >>> config.update_recursive({"data": {"dtype": "raw_data"}})
 
-    # Save to YAML
-    config.save_to_yaml("new_config.yaml")
+    >>> # Save to YAML
+    >>> config.save_to_yaml("new_config.yaml")
+
+.. testcleanup::
+
+    import os
+
+    os.remove("new_config.yaml")
 
 """
 
@@ -166,14 +173,14 @@ class Config(dict):
         each element is updated recursively if it is a Config, otherwise replaced.
 
         Example:
+            .. doctest::
 
-        .. code-block:: python
-
-            config = Config({"a": 1, "b": {"c": 2, "d": 3}})
-            config.update_recursive({"a": 4, "b": {"c": 5}})
-            print(config)
-            # <Config {'a': 4, 'b': {'c': 5, 'd': 3}}>
-            # Notice how "d" is kept and only "c" is updated.
+                >>> from zea import Config
+                >>> config = Config({"a": 1, "b": {"c": 2, "d": 3}})
+                >>> config.update_recursive({"a": 4, "b": {"c": 5}})
+                >>> # Notice how "d" is kept and only "c" is updated.
+                >>> print(config)
+                <Config {'a': 4, 'b': {'c': 5, 'd': 3}}>
 
         Args:
             dictionary (dict, optional): Dictionary to update from.
@@ -453,12 +460,6 @@ class Config(dict):
     def from_hf(cls, repo_id, path, **kwargs):
         """Load config object from huggingface hub.
 
-        Example:
-
-        .. code-block:: python
-
-            config = Config.from_hf("zeahub/configs", "config_camus.yaml", repo_type="dataset")
-
         Args:
             repo_id (str): huggingface hub repo id.
                 For example: "zeahub/configs"
@@ -471,6 +472,14 @@ class Config(dict):
 
         Returns:
             Config: config object.
+
+        Example:
+            .. doctest::
+
+                >>> from zea import Config
+                >>> config = Config.from_hf(
+                ...     "zeahub/configs", "config_camus.yaml", repo_type="dataset"
+                ... )
         """
         local_path = hf_hub_download(repo_id, path, **kwargs)
         return _load_config_from_yaml(local_path, config_class=cls)

zea/data/__init__.py

@@ -15,22 +15,28 @@ See the data notebook for a more detailed example: :doc:`../notebooks/data/zea_d
 Examples usage
 ^^^^^^^^^^^^^^
 
-.. code-block:: python
-
-    from zea import File, Dataset
-
-    # Open a single zea data file
-    with File("path/to/file.hdf5", mode="r") as file:
-        file.summary()
-        data = file.load_data("raw_data", indices=[0])
-        scan = file.scan()
-        probe = file.probe()
-
-    # Work with a dataset (folder or list of files)
-    dataset = Dataset("path/to/folder", key="raw_data")
-    for file in dataset:
-        print(file)
-    dataset.close()
+.. doctest::
+
+    >>> from zea import File, Dataset
+
+    >>> # Work with a single 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:
+    ...     # file.summary()
+    ...     data = file.load_data("raw_data", indices=[0])
+    ...     scan = file.scan()
+    ...     probe = file.probe()
+
+    >>> # Work with a dataset (folder or list of files)
+    >>> dataset = Dataset("hf://zeahub/picmus", key="raw_data")
+    >>> files = []
+    >>> for file in dataset:
+    ...     files.append(file)  # process each file as needed
+    >>> dataset.close()
 
 Subpackage layout
 -----------------

zea/data/convert/camus.py

@@ -20,7 +20,8 @@ from tqdm import tqdm
 # from zea.display import transform_sc_image_to_polar
 from zea import log
 from zea.data.data_format import generate_zea_dataset
-from zea.utils import find_first_nonzero_index, 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):