zea 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

zea/__init__.py

@@ -7,7 +7,7 @@ from . import log
 
 # dynamically add __version__ attribute (see pyproject.toml)
 # __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.1"
+__version__ = "0.0.3"
 
 
 def setup():

zea/agent/selection.py

@@ -92,6 +92,7 @@ class GreedyEntropy(LinesActionModel):
         mean: float = 0,
         std_dev: float = 1,
         num_lines_to_update: int = 5,
+        entropy_sigma: float = 1.0,
     ):
         """Initialize the GreedyEntropy action selection model.
 
@@ -104,6 +105,8 @@ class GreedyEntropy(LinesActionModel):
             std_dev (float, optional): The standard deviation of the RBF. Defaults to 1.
             num_lines_to_update (int, optional): The number of lines around the selected line
                 to update. Must be odd.
+            entropy_sigma (float, optional): The standard deviation of the Gaussian
+                Mixture components used to approximate the posterior.
         """
         super().__init__(n_actions, n_possible_actions, img_width, img_height)
 
@@ -124,7 +127,7 @@ class GreedyEntropy(LinesActionModel):
             self.num_lines_to_update,
         )
         self.upside_down_gaussian = upside_down_gaussian(points_to_evaluate)
-        self.entropy_sigma = 1
+        self.entropy_sigma = entropy_sigma
 
     @staticmethod
     def compute_pairwise_pixel_gaussian_error(
@@ -157,15 +160,18 @@ class GreedyEntropy(LinesActionModel):
         # This way we can just sum across the height axis and get the entropy
         # for each pixel in a given line
         batch_size, n_particles, _, height, _ = gaussian_error_per_pixel_i_j.shape
-        gaussian_error_per_pixel_stacked = ops.reshape(
-            gaussian_error_per_pixel_i_j,
-            [
-                batch_size,
-                n_particles,
-                n_particles,
-                height * stack_n_cols,
-                n_possible_actions,
-            ],
+        gaussian_error_per_pixel_stacked = ops.transpose(
+            ops.reshape(
+                ops.transpose(gaussian_error_per_pixel_i_j, (0, 1, 2, 4, 3)),
+                [
+                    batch_size,
+                    n_particles,
+                    n_particles,
+                    n_possible_actions,
+                    height * stack_n_cols,
+                ],
+            ),
+            (0, 1, 2, 4, 3),
         )
         # [n_particles, n_particles, batch, height, width]
         return gaussian_error_per_pixel_stacked
@@ -428,10 +434,15 @@ class CovarianceSamplingLines(LinesActionModel):
                 generation. Defaults to None.
 
         Returns:
-            Tensor: The mask of shape (batch_size, img_size, img_size)
+            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)
         """
         batch_size, n_particles, rows, _ = ops.shape(particles)
 
+        # [batch_size, rows, cols, n_particles]
+        particles = ops.transpose(particles, (0, 2, 3, 1))
+
         # [batch_size, rows * stack_n_cols, n_possible_actions, n_particles]
         shape = [
             batch_size,
@@ -441,7 +452,7 @@ class CovarianceSamplingLines(LinesActionModel):
         ]
         particles = ops.reshape(particles, shape)
 
-        # [batch_size, rows, n_possible_actions, n_possible_actions]
+        # [batch_size, rows * stack_n_cols, n_possible_actions, n_possible_actions]
         cov_matrix = tensor_ops.batch_cov(particles)
 
         # Sum over the row dimension [batch_size, n_possible_actions, n_possible_actions]

zea/backend/__init__.py

@@ -1,4 +1,4 @@
-"""Backend subpackage for ``zea``.
+"""Backend-specific utilities.
 
 This subpackage provides backend-specific utilities for the ``zea`` library. Most backend logic is handled by Keras 3, but a few features require custom wrappers to ensure compatibility and performance across JAX, TensorFlow, and PyTorch.
 
@@ -9,7 +9,7 @@ Key Features
 ------------
 
 - **JIT Compilation** (:func:`zea.backend.jit`):
-  Provides a unified interface for just-in-time (JIT) compilation of functions, dispatching to the appropriate backend (JAX or TensorFlow) as needed. This enables accelerated execution of computationally intensive routines.
+  Provides a unified interface for just-in-time (JIT) compilation of functions, dispatching to the appropriate backend (JAX or TensorFlow) as needed. This enables accelerated execution of computationally intensive routines. Note that jit compilation is not yet supported when using the `torch` backend.
 
 - **Automatic Differentiation** (:class:`zea.backend.AutoGrad`):
   Offers a backend-agnostic wrapper for automatic differentiation, allowing gradient computation regardless of the underlying ML library.
@@ -108,7 +108,9 @@ def _jit_compile(func, jax=True, tensorflow=True, **kwargs):
         return func
     else:
         log.warning(
-            f"Unsupported backend: {backend}. Supported backends are 'tensorflow' and 'jax'."
+            f"JIT compilation not currently supported for backend {backend}. "
+            "Supported backends are 'tensorflow' and 'jax'."
         )
+        log.warning("Initialize zea.Pipeline with jit_options=None to suppress this warning.")
         log.warning("Falling back to non-compiled mode.")
         return func

zea/beamform/__init__.py

@@ -17,4 +17,4 @@ see the pipeline example notebook: :doc:`../notebooks/pipeline/zea_pipeline_exam
 
 """
 
-from . import beamformer, delays, lens_correction, pfield, pixelgrid
+from . import beamformer, delays, lens_correction, pfield, phantoms, pixelgrid

zea/beamform/beamformer.py

@@ -1,5 +1,6 @@
 """Main beamforming functions for ultrasound imaging."""
 
+import keras
 import numpy as np
 from keras import ops
 
@@ -7,6 +8,45 @@ from zea.beamform.lens_correction import calculate_lens_corrected_delays
 from zea.tensor_ops import safe_vectorize
 
 
+def fnum_window_fn_rect(normalized_angle):
+    """Rectangular window function for f-number masking."""
+    return ops.where(normalized_angle <= 1.0, 1.0, 0.0)
+
+
+def fnum_window_fn_hann(normalized_angle):
+    """Hann window function for f-number masking."""
+    # Use a Hann window function to smoothly transition the mask
+    return ops.where(
+        normalized_angle <= 1.0,
+        0.5 * (1 + ops.cos(np.pi * normalized_angle)),
+        0.0,
+    )
+
+
+def fnum_window_fn_tukey(normalized_angle, alpha=0.5):
+    """Tukey window function for f-number masking.
+
+    Args:
+        normalized_angle (ops.Tensor): Normalized angle values in the range [0, 1].
+        alpha (float, optional): The alpha parameter for the Tukey window. 0.0 corresponds to a
+            rectangular window, 1.0 corresponds to a Hann window. Defaults to 0.5.
+    """
+    # Use a Tukey window function to smoothly transition the mask
+    normalized_angle = ops.clip(ops.abs(normalized_angle), 0.0, 1.0)
+
+    beta = 1.0 - alpha
+
+    return ops.where(
+        normalized_angle < beta,
+        1.0,
+        ops.where(
+            normalized_angle < 1.0,
+            0.5 * (1 + ops.cos(np.pi * (normalized_angle - beta) / (ops.abs(alpha) + 1e-6))),
+            0.0,
+        ),
+    )
+
+
 def tof_correction(
     data,
     flatgrid,
@@ -19,11 +59,12 @@ def tof_correction(
     demodulation_frequency,
     fnum,
     angles,
-    vfocus,
+    focus_distances,
     apply_phase_rotation=False,
     apply_lens_correction=False,
     lens_thickness=1e-3,
     lens_sound_speed=1000,
+    fnum_window_fn=fnum_window_fn_rect,
 ):
     """Time-of-flight correction for a flat grid.
 
@@ -44,7 +85,7 @@ def tof_correction(
         fnum (int, optional): Focus number. Defaults to 1.
         angles (ops.Tensor): The angles of the plane waves in radians of shape
             `(n_tx,)`
-        vfocus (ops.Tensor): The focus distance 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.
         apply_lens_correction (bool, optional): Whether to apply lens correction to
@@ -54,6 +95,9 @@ def tof_correction(
             lens correction. Defaults to 1e-3.
         lens_sound_speed (float, optional): Speed of sound in the lens in m/s. Used
             for lens correction Defaults to 1000.
+        fnum_window_fn (callable, optional): F-number function to define the transition from
+            straight in front of the element (fn(0.0)) to the largest angle within the f-number cone
+            (fn(1.0)). The function should be zero for fn(x>1.0).
 
     Returns:
         (ops.Tensor): time-of-flight corrected data
@@ -90,7 +134,7 @@ def tof_correction(
         sound_speed,
         n_tx,
         n_el,
-        vfocus,
+        focus_distances,
         angles,
         lens_thickness=lens_thickness,
         lens_sound_speed=lens_sound_speed,
@@ -100,7 +144,7 @@ def tof_correction(
     mask = ops.cond(
         fnum == 0,
         lambda: ops.ones((n_pix, n_el, 1)),
-        lambda: apod_mask(flatgrid, probe_geometry, fnum),
+        lambda: fnumber_mask(flatgrid, probe_geometry, fnum, fnum_window_fn=fnum_window_fn),
     )
 
     def _apply_delays(data_tx, txdel):
@@ -408,64 +452,48 @@ def distance_Tx_generic(
     return dist
 
 
-def apod_mask(grid, probe_geometry, f_number):
+def fnumber_mask(flatgrid, probe_geometry, f_number, fnum_window_fn):
     """Apodization mask for the receive beamformer.
 
-    Computes a binary mask to disregard pixels outside of the vision cone of a
+    Computes a mask to disregard pixels outside of the vision cone of a
     transducer element. Transducer elements can only accurately measure
     signals within some range of incidence angles. Waves coming in from the
     side do not register correctly leading to a worse image.
 
     Args:
-        grid (ops.Tensor): The flattened image grid `(n_pix, 3)`.
+        flatgrid (ops.Tensor): The flattened image grid `(n_pix, 3)`.
         probe_geometry (ops.Tensor): The transducer element positions of shape
             `(n_el, 3)`.
         f_number (int): The receive f-number. Set to zero to not use masking and
             return 1. (The f-number is the  ratio between distance from the transducer
             and the size of the aperture below which transducer elements contribute to
             the signal for a pixel.).
+        fnum_window_fn (callable): F-number function to define the transition from
+            straight in front of the element (fn(0.0)) to the largest angle within the f-number cone
+            (fn(1.0)). The function should be zero for fn(x>1.0).
+
 
     Returns:
         Tensor: Mask of shape `(n_pix, n_el, 1)`
     """
-    n_pix = ops.shape(grid)[0]
-    n_el = ops.shape(probe_geometry)[0]
-
-    # Get the depth of every pixel
-    z_pixel = grid[:, 2]
-    # Get the lateral location of each pixel
-    x_pixel = grid[:, 0]
-    # Get the lateral location of each element
-    x_element = ops.cast(probe_geometry[:, 0], dtype="float32")
-
-    # Compute the aperture size for every pixel
-    # The f-number is by definition f=z/aperture
-    aperture = z_pixel / f_number
-
-    # Use matrix multiplication to expand aperture tensor, x_pixel tensor, and
-    # x_element tensor to shape (n_pix, n_el)
-    ones_aperture = ops.ones(
-        (1, n_el), dtype=ops.dtype(aperture)
-    )  # getting error here? pip install -U keras ;)
-    ones_x_pixel = ops.ones((1, n_el), dtype=ops.dtype(x_pixel))
-    ones_x_element = ops.ones((n_pix, 1), dtype=ops.dtype(x_element))
-
-    aperture = ops.matmul(aperture[..., None], ones_aperture)
-    expanded_x_pixel = ops.matmul(x_pixel[..., None], ones_x_pixel)
-    expanded_x_element = ops.matmul(ones_x_element, x_element[None])
-
-    # Compute the lateral distance between elements and pixels
-    # Of shape (n_pix, n_el)
-    distance = ops.abs(expanded_x_pixel - expanded_x_element)
-
-    # Compute binary mask for which the lateral pixel distance is less than
-    # half
-    # the aperture i.e. where the pixel lies within the vision cone of the
-    # element
-    mask = distance <= aperture / 2
-    mask = ops.cast(mask, "float32")
-
-    # Add dummy dimension for RF/IQ channel channel
+
+    grid_relative_to_probe = flatgrid[:, None] - probe_geometry[None]
+
+    grid_relative_to_probe_norm = ops.linalg.norm(grid_relative_to_probe, axis=-1)
+
+    grid_relative_to_probe_z = grid_relative_to_probe[..., 2] / (grid_relative_to_probe_norm + 1e-6)
+
+    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))
+    # We can use this to compute the maximum angle alpha that is allowed
+    max_alpha = ops.arctan(1 / (2 * f_number + keras.backend.epsilon()))
+
+    normalized_angle = alpha / max_alpha
+    mask = fnum_window_fn(normalized_angle)
+
+    # Add dummy channel dimension
     mask = mask[..., None]
 
     return mask

zea/beamform/pfield.py

@@ -60,7 +60,8 @@ def compute_pfield(
         n_el (int): Number of elements in the probe.
         probe_geometry (array): Geometry of the probe elements.
         tx_apodizations (array): Transmit apodization values.
-        grid (array): Grid points where the pressure field is computed of shape (Nz, Nx, 3).
+        grid (array): Grid points where the pressure field is computed
+            of shape (grid_size_z, grid_size_x, 3).
         t0_delays (array): Transmit delays for each transmit event.
         frequency_step (int, optional): Frequency step. Default is 4.
             Higher is faster but less accurate.
@@ -78,7 +79,8 @@ def compute_pfield(
         verbose (bool, optional): Whether to print progress.
 
     Returns:
-        ops.array: The (normalized) pressure field (across tx events) of shape (n_tx, Nz, Nx).
+        ops.array: The (normalized) pressure field (across tx events)
+            of shape (n_tx, grid_size_z, grid_size_x).
     """
     # medium params
     alpha_db = 0  # currently we ignore attenuation in the compounding
@@ -293,7 +295,8 @@ def normalize_pressure_field(pfield, alpha: float = 1.0, percentile: float = 10.
     Normalize the input array of intensities by zeroing out values below a given percentile.
 
     Args:
-        pfield (array): The unnormalized pressure field array of shape (n_tx, Nz, Nx).
+        pfield (array): The unnormalized pressure field array
+            of shape (n_tx, grid_size_z, grid_size_x).
         alpha (float, optional): Exponent to 'sharpen or smooth' the weighting.
             Higher values result in sharper weighting. Default is 1.0.
         percentile (int, optional): minimum percentile threshold to keep in the weighting.

zea/beamform/phantoms.py

@@ -0,0 +1,43 @@
+import numpy as np
+
+
+def fish():
+    """Returns a scatterer phantom for ultrasound simulation tests.
+
+    Returns:
+        ndarray: The scatterer positions of shape (n_scat, 3).
+    """
+    # The size is the height of the fish
+    size = 11e-3
+    z_offset = 2.0 * size
+
+    # See https://en.wikipedia.org/wiki/Fish_curve
+    def fish_curve(t, size=1):
+        x = size * (np.cos(t) - np.sin(t) ** 2 / np.sqrt(2))
+        y = size * np.cos(t) * np.sin(t)
+        return x, y
+
+    scat_x, scat_z = fish_curve(np.linspace(0, 2 * np.pi, 100), size=size)
+
+    scat_x = np.concatenate(
+        [
+            scat_x,
+            np.array([size * 0.7]),
+            np.array([size * 1.1]),
+            np.array([size * 1.4]),
+            np.array([size * 1.2]),
+        ]
+    )
+    scat_y = np.zeros_like(scat_x)
+    scat_z = np.concatenate(
+        [
+            scat_z,
+            np.array([-size * 0.1]),
+            np.array([-size * 0.25]),
+            np.array([-size * 0.6]),
+            np.array([-size * 1.0]),
+        ]
+    )
+
+    scat_z += z_offset
+    return np.stack([scat_x, scat_y, scat_z], axis=1)

zea/beamform/pixelgrid.py

@@ -7,101 +7,61 @@ from zea import log
 eps = 1e-10
 
 
-def get_grid(
-    xlims,
-    zlims,
-    Nx: int,
-    Nz: int,
-    sound_speed,
-    center_frequency,
-    pixels_per_wavelength,
-    verbose=False,
-):
-    """Creates a pixelgrid based on scan class parameters."""
-
-    if Nx and Nz:
-        grid = cartesian_pixel_grid(xlims, zlims, Nx=int(Nx), Nz=int(Nz))
-    else:
-        wvln = sound_speed / center_frequency
-        dx = wvln / pixels_per_wavelength
-        dz = dx
-        grid = cartesian_pixel_grid(xlims, zlims, dx=dx, dz=dz)
-        if verbose:
-            print(
-                f"Pixelgrid was set automatically to Nx: {grid.shape[1]}, Nz: {grid.shape[0]}, "
-                f"using {pixels_per_wavelength} pixels per wavelength."
-            )
-    return grid
-
-
 def check_for_aliasing(scan):
     """Checks if the scan class parameters will cause spatial aliasing due to a too low pixel
     density. If so, a warning is printed with a suggestion to increase the pixel density by either
     increasing the number of pixels, or decreasing the pixel spacing, depending on which parameter
     was set by the user."""
-    wvln = scan.sound_speed / scan.center_frequency
-    dx = wvln / scan.pixels_per_wavelength
-    dz = dx
-
     width = scan.xlims[1] - scan.xlims[0]
     depth = scan.zlims[1] - scan.zlims[0]
-
-    if scan.Nx and scan.Nz:
-        if width / scan.Nx > wvln / 2:
-            log.warning(
-                f"width/Nx = {width / scan.Nx:.7f} < wvln/2 = {wvln / 2}. "
-                f"Consider increasing scan.Nx to {int(np.ceil(width / (wvln / 2)))} or more."
-            )
-        if depth / scan.Nz > wvln / 2:
-            log.warning(
-                f"depth/Nz = {depth / scan.Nz:.7f} < wvln/2 = {wvln / 2:.7f}. "
-                f"Consider increasing scan.Nz to {int(np.ceil(depth / (wvln / 2)))} or more."
-            )
-    else:
-        if dx > wvln / 2:
-            log.warning(
-                f"dx = {dx:.7f} > wvln/2 = {wvln / 2:.7f}. "
-                f"Consider increasing scan.pixels_per_wavelength to 2 or more"
-            )
-        if dz > wvln / 2:
-            log.warning(
-                f"dz = {dz:.7f} > wvln/2 = {wvln / 2:.7f}. "
-                f"Consider increasing scan.pixels_per_wavelength to 2 or more"
-            )
-
-
-def cartesian_pixel_grid(xlims, zlims, Nx=None, Nz=None, dx=None, dz=None):
+    wvln = scan.wavelength
+
+    if width / scan.grid_size_x > wvln / 2:
+        log.warning(
+            f"width/grid_size_x = {width / scan.grid_size_x:.7f} < wavelength/2 = {wvln / 2}. "
+            f"Consider either increasing scan.grid_size_x to {int(np.ceil(width / (wvln / 2)))} "
+            "or more, or increasing scan.pixels_per_wavelength to 2 or more."
+        )
+    if depth / scan.grid_size_z > wvln / 2:
+        log.warning(
+            f"depth/grid_size_z = {depth / scan.grid_size_z:.7f} < wavelength/2 = {wvln / 2:.7f}. "
+            f"Consider either increasing scan.grid_size_z to {int(np.ceil(depth / (wvln / 2)))} "
+            "or more, or increasing scan.pixels_per_wavelength to 2 or more."
+        )
+
+
+def cartesian_pixel_grid(xlims, zlims, grid_size_x=None, grid_size_z=None, dx=None, dz=None):
     """Generate a Cartesian pixel grid based on input parameters.
 
     Args:
         xlims (tuple): Azimuthal limits of pixel grid ([xmin, xmax])
         zlims (tuple): Depth limits of pixel grid ([zmin, zmax])
-        Nx (int): Number of azimuthal pixels, overrides dx and dz parameters
-        Nz (int): Number of depth pixels, overrides dx and dz parameters
+        grid_size_x (int): Number of azimuthal pixels, overrides dx and dz parameters
+        grid_size_z (int): Number of depth pixels, overrides dx and dz parameters
         dx (float): Pixel spacing in azimuth
         dz (float): Pixel spacing in depth
 
     Raises:
-        ValueError: Either Nx and Nz or dx and dz must be defined.
+        ValueError: Either grid_size_x and grid_size_z or dx and dz must be defined.
 
     Returns:
-        grid (np.ndarray): Pixel grid of size (nz, nx, 3) in
+        grid (np.ndarray): Pixel grid of size (grid_size_z, nx, 3) in
             Cartesian coordinates (x, y, z)
     """
-    assert (bool(Nx) and bool(Nz)) ^ (bool(dx) and bool(dz)), (
-        "Either Nx and Nz or dx and dz must be defined."
+    assert (bool(grid_size_x) and bool(grid_size_z)) ^ (bool(dx) and bool(dz)), (
+        "Either grid_size_x and grid_size_z or dx and dz must be defined."
     )
 
     # Determine the grid spacing
-    if Nx is not None and Nz is not None:
-        x = np.linspace(xlims[0], xlims[1] + eps, Nx)
-        z = np.linspace(zlims[0], zlims[1] + eps, Nz)
+    if grid_size_x is not None and grid_size_z is not None:
+        x = np.linspace(xlims[0], xlims[1] + eps, grid_size_x)
+        z = np.linspace(zlims[0], zlims[1] + eps, grid_size_z)
     elif dx is not None and dz is not None:
         sign = np.sign(xlims[1] - xlims[0])
         x = np.arange(xlims[0], xlims[1] + eps, sign * dx)
         z = np.arange(zlims[0], zlims[1] + eps, sign * dz)
     else:
-        raise ValueError("Either Nx and Nz or dx and dz must be defined.")
+        raise ValueError("Either grid_size_x and grid_size_z or dx and dz must be defined.")
 
     # Create the pixel grid
     z_grid, x_grid = np.meshgrid(z, x, indexing="ij")
@@ -130,7 +90,7 @@ def radial_pixel_grid(rlims, dr, oris, dirs):
             Cartesian coordinates (x, y, z), with nr being the number of radial pixels.
     """
     # Get focusing positions in rho-theta coordinates
-    r = np.arange(rlims[0], rlims[1] + eps, dr)  # Depth rho
+    r = np.arange(rlims[0], rlims[1], dr)  # Depth rho
     t = dirs[:, 0]  # Use azimuthal angle theta (ignore elevation angle)
     tt, rr = np.meshgrid(t, r, indexing="ij")
 
@@ -140,3 +100,32 @@ def radial_pixel_grid(rlims, dr, oris, dirs):
     yy = 0 * xx
     grid = np.stack((xx, yy, zz), axis=-1)
     return grid
+
+
+def polar_pixel_grid(polar_limits, zlims, num_radial_pixels: int, num_polar_pixels: int):
+    """Generate a polar grid.
+
+    Uses radial_pixel_grid but based on parameters that are present in the scan class.
+
+    Args:
+        polar_limits (tuple): Polar limits of pixel grid ([polar_min, polar_max])
+        zlims (tuple): Depth limits of pixel grid ([zmin, zmax])
+        num_radial_pixels (int, optional): Number of depth pixels.
+        num_polar_pixels (int, optional): Number of polar pixels.
+
+    Returns:
+        grid (np.ndarray): Pixel grid of size (num_radial_pixels, num_polar_pixels, 3)
+        in Cartesian coordinates (x, y, z)
+    """
+    assert len(polar_limits) == 2, "polar_limits must be a tuple of length 2."
+    assert len(zlims) == 2, "zlims must be a tuple of length 2."
+
+    dr = (zlims[1] - zlims[0]) / num_radial_pixels
+
+    oris = np.array([0, 0, 0])
+    oris = np.tile(oris, (num_polar_pixels, 1))
+    dirs_az = np.linspace(*polar_limits, num_polar_pixels)
+
+    dirs_el = np.zeros(num_polar_pixels)
+    dirs = np.vstack((dirs_az, dirs_el)).T
+    return radial_pixel_grid(zlims, dr, oris, dirs).transpose(1, 0, 2)

zea/config.py

@@ -47,8 +47,9 @@ import yaml
 from huggingface_hub import hf_hub_download
 
 from zea import log
+from zea.data.preset_utils import HF_PREFIX, _hf_resolve_path
 from zea.internal.config.validation import config_schema
-from zea.internal.core import object_to_tensor
+from zea.internal.core import dict_to_tensor
 
 
 class Config(dict):
@@ -430,8 +431,22 @@ class Config(dict):
                         v._recursive_setattr(set_key, set_value)
 
     @classmethod
-    def from_yaml(cls, path, **kwargs):
-        """Load config object from yaml file"""
+    def from_path(cls, path, **kwargs):
+        """Load config object from a file path.
+
+        Args:
+            path (str or Path): The path to the config file.
+                Can be a string or a Path object. Additionally can be a string with
+                the prefix 'hf://', in which case it will be resolved to a
+                huggingface path.
+
+        Returns:
+            Config: config object.
+        """
+        if str(path).startswith(HF_PREFIX):
+            path = _hf_resolve_path(str(path))
+        if isinstance(path, str):
+            path = Path(path)
         return _load_config_from_yaml(path, config_class=cls, **kwargs)
 
     @classmethod
@@ -460,9 +475,14 @@ class Config(dict):
         local_path = hf_hub_download(repo_id, path, **kwargs)
         return _load_config_from_yaml(local_path, config_class=cls)
 
-    def to_tensor(self):
+    @classmethod
+    def from_yaml(cls, path, **kwargs):
+        """Load config object from yaml file."""
+        return cls.from_path(path, **kwargs)
+
+    def to_tensor(self, keep_as_is=None):
         """Convert the attributes in the object to keras tensors"""
-        return object_to_tensor(self)
+        return dict_to_tensor(self.serialize(), keep_as_is=keep_as_is)
 
 
 def check_config(config: Union[dict, Config], verbose: bool = False):

zea/data/__main__.py

@@ -0,0 +1,31 @@
+"""Command-line interface for copying a zea.Folder to a new location.
+
+Usage:
+    python -m zea.data <source_folder> <destination_folder> <key>
+"""
+
+import argparse
+
+from zea import Folder
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Copy a 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")
+    parser.add_argument(
+        "--mode",
+        default="a",
+        choices=["a", "w", "r+", "x"],
+        help="Mode in which to open the destination files (default: 'a')",
+    )
+
+    args = parser.parse_args()
+
+    src_folder = Folder(args.src, args.key, validate=False)
+    src_folder.copy(args.dst, args.key, mode=args.mode)
+
+
+if __name__ == "__main__":
+    main()