zea 0.0.2__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.2"
+__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
 
@@ -58,7 +59,7 @@ def tof_correction(
     demodulation_frequency,
     fnum,
     angles,
-    vfocus,
+    focus_distances,
     apply_phase_rotation=False,
     apply_lens_correction=False,
     lens_thickness=1e-3,
@@ -84,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
@@ -133,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,
@@ -487,7 +488,7 @@ def fnumber_mask(flatgrid, probe_geometry, f_number, fnum_window_fn):
     # 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))
+    max_alpha = ops.arctan(1 / (2 * f_number + keras.backend.epsilon()))
 
     normalized_angle = alpha / max_alpha
     mask = fnum_window_fn(normalized_angle)

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

@@ -16,52 +16,52 @@ def check_for_aliasing(scan):
     depth = scan.zlims[1] - scan.zlims[0]
     wvln = scan.wavelength
 
-    if width / scan.Nx > wvln / 2:
+    if width / scan.grid_size_x > wvln / 2:
         log.warning(
-            f"width/Nx = {width / scan.Nx:.7f} < wavelength/2 = {wvln / 2}. "
-            f"Consider either increasing scan.Nx to {int(np.ceil(width / (wvln / 2)))} or more, or "
-            "increasing scan.pixels_per_wavelength to 2 or more."
+            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.Nz > wvln / 2:
+    if depth / scan.grid_size_z > wvln / 2:
         log.warning(
-            f"depth/Nz = {depth / scan.Nz:.7f} < wavelength/2 = {wvln / 2:.7f}. "
-            f"Consider either increasing scan.Nz to {int(np.ceil(depth / (wvln / 2)))} or more, or "
-            "increasing scan.pixels_per_wavelength to 2 or more."
+            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, Nx=None, Nz=None, dx=None, dz=None):
+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")
@@ -102,29 +102,30 @@ def radial_pixel_grid(rlims, dr, oris, dirs):
     return grid
 
 
-def polar_pixel_grid(polar_limits, zlims, Nz: int, Nr: int):
+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): Azimuthal limits of pixel grid ([azimuth_min, azimuth_max])
+        polar_limits (tuple): Polar limits of pixel grid ([polar_min, polar_max])
         zlims (tuple): Depth limits of pixel grid ([zmin, zmax])
-        Nz (int, optional): Number of depth pixels.
-        Nr (int, optional): Number of azimuthal pixels.
+        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 (Nz, Nr, 3) in Cartesian coordinates (x, y, z)
+        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]) / Nz
+    dr = (zlims[1] - zlims[0]) / num_radial_pixels
 
     oris = np.array([0, 0, 0])
-    oris = np.tile(oris, (Nr, 1))
-    dirs_az = np.linspace(*polar_limits, Nr)
+    oris = np.tile(oris, (num_polar_pixels, 1))
+    dirs_az = np.linspace(*polar_limits, num_polar_pixels)
 
-    dirs_el = np.zeros(Nr)
+    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/data_format.py

@@ -42,8 +42,8 @@ def generate_example_dataset(
     sound_speed=1540,
     center_frequency=7e6,
     sampling_frequency=40e6,
-    n_z=512,
-    n_x=512,
+    grid_size_z=512,
+    grid_size_x=512,
 ):
     """Generates an example dataset that contains all the necessary fields.
     Note: This dataset does not contain actual data, but is filled with random
@@ -60,7 +60,7 @@ def generate_example_dataset(
     # creating some fake raw and image data
     raw_data = np.ones((n_frames, n_tx, n_ax, n_el, n_ch))
     # image data is in dB
-    image = np.ones((n_frames, n_z, n_x)) * -40
+    image = np.ones((n_frames, grid_size_z, grid_size_x)) * -40
 
     # creating some fake scan parameters
     t0_delays = np.zeros((n_tx, n_el), dtype=np.float32)
@@ -82,8 +82,8 @@ def generate_example_dataset(
 
     if add_optional_dtypes:
         aligned_data = np.ones((n_frames, n_tx, n_ax, n_el, n_ch))
-        envelope_data = np.ones((n_frames, n_z, n_x))
-        beamformed_data = np.ones((n_frames, n_z, n_x, n_ch))
+        envelope_data = np.ones((n_frames, grid_size_z, grid_size_x))
+        beamformed_data = np.ones((n_frames, grid_size_z, grid_size_x, n_ch))
         image_sc = np.ones_like(image)
     else:
         aligned_data = None
@@ -123,10 +123,11 @@ def validate_input_data(raw_data, aligned_data, envelope_data, beamformed_data,
         aligned_data (np.ndarray): The aligned data of the ultrasound measurement of
             shape (n_frames, n_tx, n_ax, n_el, n_ch).
         envelope_data (np.ndarray): The envelope data of the ultrasound measurement of
-            shape (n_frames, n_z, n_x).
+            shape (n_frames, grid_size_z, grid_size_x).
         beamformed_data (np.ndarray): The beamformed data of the ultrasound measurement of
-            shape (n_frames, n_z, n_x).
-        image (np.ndarray): The ultrasound images to be saved of shape (n_frames, n_z, n_x).
+            shape (n_frames, grid_size_z, grid_size_x).
+        image (np.ndarray): The ultrasound images to be saved
+            of shape (n_frames, grid_size_z, grid_size_x).
         image_sc (np.ndarray): The scan converted ultrasound images to be saved
             of shape (n_frames, output_size_z, output_size_x).
     """
@@ -254,7 +255,7 @@ def _write_datasets(
         group_name=data_group_name,
         name="envelope_data",
         data=_convert_datatype(envelope_data),
-        description="The envelope_data of shape (n_frames, n_z, n_x).",
+        description="The envelope_data of shape (n_frames, grid_size_z, grid_size_x).",
         unit="unitless",
     )
 
@@ -262,7 +263,7 @@ def _write_datasets(
         group_name=data_group_name,
         name="beamformed_data",
         data=_convert_datatype(beamformed_data),
-        description="The beamformed_data of shape (n_frames, n_z, n_x).",
+        description="The beamformed_data of shape (n_frames, grid_size_z, grid_size_x).",
         unit="unitless",
     )
 
@@ -271,7 +272,7 @@ def _write_datasets(
         name="image",
         data=_convert_datatype(image),
         unit="unitless",
-        description="The images of shape [n_frames, n_z, n_x]",
+        description="The images of shape [n_frames, grid_size_z, grid_size_x]",
     )
 
     _add_dataset(
@@ -546,10 +547,11 @@ def generate_zea_dataset(
         aligned_data (np.ndarray): The aligned data of the ultrasound measurement of
             shape (n_frames, n_tx, n_ax, n_el, n_ch).
         envelope_data (np.ndarray): The envelope data of the ultrasound measurement of
-            shape (n_frames, n_z, n_x).
+            shape (n_frames, grid_size_z, grid_size_x).
         beamformed_data (np.ndarray): The beamformed data of the ultrasound measurement of
-            shape (n_frames, n_z, n_x, n_ch).
-        image (np.ndarray): The ultrasound images to be saved of shape (n_frames, n_z, n_x).
+            shape (n_frames, grid_size_z, grid_size_x, n_ch).
+        image (np.ndarray): The ultrasound images to be saved
+            of shape (n_frames, grid_size_z, grid_size_x).
         image_sc (np.ndarray): The scan converted ultrasound images to be saved
             of shape (n_frames, output_size_z, output_size_x).
         probe_geometry (np.ndarray): The probe geometry of shape (n_el, 3).

zea/data/file.py

@@ -366,9 +366,7 @@ class File(h5py.File):
         """
         file_scan_parameters = self.get_parameters(event)
 
-        probe_parameters = reduce_to_signature(
-            Probe.from_name("generic").__init__, file_scan_parameters
-        )
+        probe_parameters = reduce_to_signature(Probe.__init__, file_scan_parameters)
         return probe_parameters
 
     def probe(self, event=None) -> Probe:
@@ -387,21 +385,8 @@ class File(h5py.File):
         Returns:
             Probe: The probe object.
         """
-        probe_parameters = self.get_probe_parameters(event)
-        if self.probe_name == "generic":
-            return Probe.from_name(self.probe_name, **probe_parameters)
-        else:
-            probe = Probe.from_name(self.probe_name)
-
-            probe_geometry = probe_parameters.get("probe_geometry", None)
-            if not np.allclose(probe_geometry, probe.probe_geometry):
-                probe.probe_geometry = probe_geometry
-                log.warning(
-                    "The probe geometry in the data file does not "
-                    "match the probe geometry of the probe. The probe "
-                    "geometry has been updated to match the data file."
-                )
-        return probe
+        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:
         """Load dict from contents of group

zea/data/preset_utils.py

@@ -50,60 +50,74 @@ def _hf_download(repo_id, filename, cache_dir=HF_DATASETS_DIR):
     )
 
 
-def _hf_get_snapshot_dir(repo_id, cache_dir=HF_DATASETS_DIR):
-    repo_cache_dir = Path(cache_dir) / f"datasets--{repo_id.replace('/', '--')}"
-    snapshots_dir = repo_cache_dir / "snapshots"
-    if not snapshots_dir.exists() or not any(snapshots_dir.iterdir()):
-        # Try to trigger a download to populate the cache
-        files = _hf_list_files(repo_id)
-        # Pick the first file (prefer .h5/.hdf5 if possible)
-        h5_files = [f for f in files if f.endswith(".h5") or f.endswith(".hdf5")]
-        target_file = h5_files[0] if h5_files else files[0]
-        _hf_download(repo_id, target_file, cache_dir)
-        # Now try again
-        if not snapshots_dir.exists() or not any(snapshots_dir.iterdir()):
-            raise FileNotFoundError(
-                f"No snapshots found in Hugging Face cache for {repo_id} after download attempt"
-            )
-    snapshot_hashes = sorted(snapshots_dir.iterdir(), key=lambda p: p.stat().st_mtime, reverse=True)
-    if not snapshot_hashes:
-        raise FileNotFoundError(f"No snapshot found for {repo_id} in cache.")
-    return snapshot_hashes[0]
+def _get_snapshot_dir_from_downloaded_file(downloaded_file_path: str | Path) -> Path:
+    """Extract the snapshot directory from a downloaded file's path.
+
+    HF Hub downloads to: cache_dir/datasets--org--repo/snapshots/{hash}/path/to/filename
+    This navigates up to find the {hash} directory (the snapshot directory).
+    """
+    file_path = Path(downloaded_file_path)
+
+    # Navigate up the path until we find the snapshots directory
+    current = file_path.parent
+    while current.name != "snapshots" and current.parent != current:
+        current = current.parent
+
+    if current.name == "snapshots":
+        # Return the snapshot hash directory (first subdirectory of snapshots)
+        snapshot_dirs = [d for d in current.iterdir() if d.is_dir()]
+        if snapshot_dirs:
+            # Return the most recent snapshot directory
+            return max(snapshot_dirs, key=lambda p: p.stat().st_mtime)
+
+    raise FileNotFoundError(f"Could not find snapshot directory for {downloaded_file_path}")
+
+
+def _download_files_in_path(
+    repo_id: str, files: list, path_filter: str = None, cache_dir=HF_DATASETS_DIR
+) -> list[str]:
+    """Download all files matching the path filter."""
+    downloaded_files = []
+    for f in files:
+        if path_filter is None or f.startswith(path_filter):
+            downloaded_path = _hf_download(repo_id, f, cache_dir)
+            downloaded_files.append(downloaded_path)
+
+    return downloaded_files
 
 
 def _hf_resolve_path(hf_path: str, cache_dir=HF_DATASETS_DIR):
-    """Download a file or directory from Hugging Face Hub to a local cache directory.
-    Returns the local path to the downloaded file or directory.
+    """Resolve a Hugging Face path to a local cache directory path.
+
+    Downloads files from a HuggingFace dataset repository and returns
+    the local path where they are cached. Handles:
+    - hf://org/repo/subdir/ - Downloads all files in subdirectory
+    - hf://org/repo/file.h5 - Downloads specific file
+    - hf://org/repo - Downloads all files in repo
     """
     repo_id, subpath = _hf_parse_path(hf_path)
     files = _hf_list_files(repo_id)
-    snapshot_dir = _hf_get_snapshot_dir(repo_id, cache_dir)
-
-    def is_h5(f):
-        return f.endswith(".h5") or f.endswith(".hdf5")
 
     if subpath:
-        # Directory
+        # Directory case
         if any(f.startswith(subpath + "/") for f in files):
-            local_dir = snapshot_dir / subpath
-            for f in files:
-                if f.startswith(subpath + "/") and is_h5(f):
-                    _hf_download(repo_id, f, cache_dir)
-            if not local_dir.exists():
-                raise FileNotFoundError(f"Directory {local_dir} not found after download.")
-            return local_dir
-        # File
-        elif any(f == subpath for f in files) and is_h5(subpath):
-            _hf_download(repo_id, subpath, cache_dir)
-            local_file = snapshot_dir / subpath
-            if not local_file.exists():
-                raise FileNotFoundError(f"File {local_file} not found after download.")
-            return local_file
+            downloaded_files = _download_files_in_path(repo_id, files, subpath + "/", cache_dir)
+            if not downloaded_files:
+                raise FileNotFoundError(f"No files found in directory {subpath}")
+
+            snapshot_dir = _get_snapshot_dir_from_downloaded_file(downloaded_files[0])
+            return snapshot_dir / subpath
+
+        # File case
+        elif subpath in files:
+            downloaded_file = _hf_download(repo_id, subpath, cache_dir)
+            return Path(downloaded_file)
         else:
             raise FileNotFoundError(f"{subpath} not found in {repo_id}")
     else:
-        # All .h5/.hdf5 files in repo
-        for f in files:
-            if is_h5(f):
-                _hf_download(repo_id, f, cache_dir)
-        return snapshot_dir
+        # All files in repo
+        downloaded_files = _download_files_in_path(repo_id, files, None, cache_dir)
+        if not downloaded_files:
+            raise FileNotFoundError(f"No files found in repository {repo_id}")
+
+        return _get_snapshot_dir_from_downloaded_file(downloaded_files[0])

zea/datapaths.py

@@ -257,13 +257,13 @@ def set_data_paths(
 
     Returns:
         dict: Absolute paths to location of data. Stores the following parameters:
-            ``data_root``, ``repo_root``, ``output``, ``system``, ``username``, ``hostname``
+            ``data_root``, ``zea_root``, ``output``, ``system``, ``username``, ``hostname``
 
     """
     username = getpass.getuser()
     system = platform.system().lower()
     hostname = socket.gethostname()
-    repo_root = importlib.resources.files("zea")  # ultrasound-toolbox/zea
+    zea_root = importlib.resources.files("zea")
 
     # If user_config is None, use the default users.yaml file
     if isinstance(user_config, type(None)):
@@ -304,7 +304,7 @@ def set_data_paths(
 
     data_path = {
         "data_root": Path(data_root),
-        "repo_root": repo_root,
+        "zea_root": zea_root,
         "output": Path(output),
         "system": system,
         "username": username,

zea/display.py

@@ -130,7 +130,7 @@ def scan_convert_2d(
 
     Returns:
         ndarray: The scan-converted 2D ultrasound image in Cartesian coordinates.
-            Has dimensions (n_z, n_x). Coordinates outside the input image
+            Has dimensions (grid_size_z, grid_size_x). Coordinates outside the input image
             ranges are filled with NaNs.
         parameters (dict): A dictionary containing information about the scan conversion.
             Contains the resolution, x, and z limits, rho and theta ranges.
@@ -269,7 +269,7 @@ def scan_convert_3d(
 
     Returns:
         ndarray: The scan-converted 3D ultrasound image in Cartesian coordinates.
-            Has dimensions (n_z, n_x, n_y). Coordinates outside the input image
+            Has dimensions (grid_size_z, grid_size_x, n_y). Coordinates outside the input image
             ranges are filled with NaNs.
         parameters (dict): A dictionary containing information about the scan conversion.
             Contains the resolution, x, y, and z limits, rho, theta, and phi ranges.