"""
This module provides some utilities to edit zea data files, either individually or in bulk.
Each operation is available both as a Python function and as a ``zea data`` command line
subcommand. See the :doc:`CLI documentation </cli>` for the available operations and their
command-line usage.
"""
import functools
from collections.abc import Sequence
from copy import copy as shallow_copy
from pathlib import Path
from typing import TypeVar
import numpy as np
import tyro
from keras import ops
from tqdm import tqdm
from zea import File, log
from zea.data.datasets import Dataset
from zea.data.spec import (
CONSISTENCY_DIMENSIONS,
FileSpec,
ProbeSpec,
ScanSpec,
Spec,
find_matched_shape,
is_lazy_array,
iter_slabs,
)
from zea.func.ultrasound import construct_acquisition_from_synthetic_aperture, decode_hadamard
from zea.internal.checks import _IMAGE_DATA_TYPES, _NON_IMAGE_DATA_TYPES
from zea.internal.preset_utils import HF_PREFIX, _hf_resolve_path
ALL_DATA_TYPES_EXCEPT_RAW = set(_IMAGE_DATA_TYPES + _NON_IMAGE_DATA_TYPES) - {"raw_data"}
SpecT = TypeVar("SpecT", bound=Spec)
# Data products stored log-compressed (in dB), which must be averaged in the linear domain.
_LOG_DOMAIN_FIELDS = frozenset({"image"})
# dB conversion factor: zea's ``log_compress`` uses ``20 * log10(x)`` (see
# ``zea.func.ultrasound.log_compress``), so the inverse is ``10 ** (dB / 20)``.
_DB_FACTOR = 20.0
OPERATION_NAMES = [
"sum",
"compound_frames",
"compound_transmits",
"resave",
"extract",
"summary",
"copy",
"decode_hadamard",
"sa_to_virtual_focus",
]
def _iter_folder_io(input_path: str | Path, output_path: str | Path):
"""Yields ``(input_file, output_file)`` path pairs for a folder operation.
Uses :class:`zea.Dataset` to iterate over every zea file in ``input_path``. The
output folder mirrors the structure of the input folder.
Args:
input_path (Path): Path to a folder containing zea data files.
output_path (Path): Path to the output folder.
Yields:
tuple[Path, Path]: Pairs of (input file, output file) paths.
"""
input_path, output_path = Path(input_path), Path(output_path)
# validate=False: this only enumerates paths — the operation validates each file
# when it opens it for the actual work.
with Dataset(input_path, validate=False) as dataset:
for file in dataset:
yield file.path, output_path / file.path.relative_to(input_path)
def _supports_folders(operation):
"""Decorator that lets a single-file operation also accept a folder as input.
When the decorated operation is called with a folder as ``input_path``, it is
applied to every zea file in that folder (iterated with :class:`zea.Dataset`),
writing the results to ``output_path`` and mirroring the input folder structure.
A single file is processed as before.
``input_path`` may also be an ``hf://`` path (pointing at a single file or a
folder in a Hugging Face dataset repo); it is downloaded via
:func:`zea.internal.preset_utils._hf_resolve_path` before dispatching.
Note that if a file operation does not need all the data, this might not be optimal since
it doesn't stream the data and instead downloads the whole file.
"""
@functools.wraps(operation)
def wrapper(input_path, output_path, *args, **kwargs):
input_path = str(input_path)
if input_path.startswith(HF_PREFIX):
input_path = _hf_resolve_path(input_path)
if not Path(input_path).is_dir():
return operation(input_path, output_path, *args, **kwargs)
output_path = Path(output_path)
if output_path.is_file():
raise NotADirectoryError(
f"Input {input_path} is a folder, so output {output_path} must be a "
"folder, but it is an existing file."
)
for in_path, out_path in tqdm(list(_iter_folder_io(input_path, output_path))):
operation(in_path, out_path, *args, **kwargs)
return None
return wrapper
def _data_arrays(track) -> dict:
"""Return every data-product array of a track, keyed by field name.
Map-based products (``image``, ``beamformed_data``, ...) store their array in
``values``; plain products (``raw_data``) are the array itself. The array is
returned together with the SCHEMA shape that describes it, so callers can
resolve named dimensions against it.
"""
arrays = {}
if track.data is None:
return arrays
for field_name, field_info in track.data.SCHEMA.items():
value = getattr(track.data, field_name)
if value is None:
continue
nested = field_info.get("spec")
if nested is None:
arrays[field_name] = (value, field_info["shape"])
else:
arrays[field_name] = (value.values, type(value).SCHEMA["values"]["shape"])
return arrays
def _set_data_array(track, field_name: str, array: np.ndarray) -> None:
"""Write ``array`` back into a track's data product (inverse of :func:`_data_arrays`)."""
value = getattr(track.data, field_name)
if isinstance(value, Spec):
value.values = array
else:
setattr(track.data, field_name, array)
def _dim_axis(array: np.ndarray, shape_spec, dim_name: str) -> int | None:
"""The single axis of ``array`` carrying ``dim_name``, or None if it has none."""
matched_shape = find_matched_shape(array, Spec._expected_shapes(shape_spec))
if matched_shape is None:
return None
axes = _named_axes(matched_shape, array.ndim, dim_name)
return axes[0] if axes else None
def _mean_over_axis(array, axis: int, log_domain: bool = False) -> np.ndarray:
"""Average ``array`` along ``axis``, keeping the axis, in the domain the data lives in.
float32 images are log-compressed (dB), so they are averaged in the linear domain;
uint8 images are averaged in float and clipped back into range.
The input is consumed slab by slab (see :func:`zea.data.spec.iter_slabs`), so a lazy
array is never read whole: only one slab and the (reduced) result are in memory. Slabs
are cut along the frame axis, so a reduction over frames accumulates across them while
any other reduction is done per slab and concatenated. Partial sums are accumulated in
float64 regardless of the input dtype, which keeps the result independent of how the
array happened to be split.
"""
dtype = np.dtype(array.dtype)
linear = log_domain and dtype == np.float32
# Widest accumulator the input can be cast to (IQ data products are complex64).
accumulate_as = np.complex128 if np.issubdtype(dtype, np.complexfloating) else np.float64
# Slabs tile the frame axis: reducing over it accumulates them into a single sum, any
# other axis reduces each slab on its own and keeps the results side by side.
total = None
partial_sums = []
for _, slab in iter_slabs(array, split_axes=False):
if linear:
# Undo the dB compression (10**(dB/20)) to average in the linear domain.
slab = np.power(10.0, slab / _DB_FACTOR)
# asarray: keepdims=True always gives an array, which the numpy stubs cannot say.
partial = np.asarray(np.sum(slab, axis=axis, keepdims=True, dtype=accumulate_as))
if axis != 0:
partial_sums.append(partial)
elif total is None:
total = partial
else:
total += partial
if total is None and not partial_sums:
# No slabs at all: the array is empty, which np.mean answers with an empty mean.
return np.mean(np.asarray(array), axis=axis, keepdims=True).astype(dtype)
if axis != 0:
total = np.concatenate(partial_sums, axis=0)
mean = total / array.shape[axis]
if linear:
return (_DB_FACTOR * np.log10(mean)).astype(np.float32)
if dtype == np.uint8:
return np.clip(mean, 0, 255).astype(np.uint8)
return mean.astype(dtype)
def _compound_named_dim(file_spec: FileSpec, dim_name: str) -> FileSpec:
"""Average every data product over ``dim_name``, collapsing that dimension to length 1.
Data products are averaged; everything else that carries the dimension (scan
parameters, timestamps, annotations, metrics) is reduced to its first entry, so
the whole spec stays consistent. Which fields carry the dimension is decided by
the SCHEMA, so products without it (e.g. ``image`` has no ``n_tx``) are left alone.
"""
# Averages are computed from the unsliced spec, then written into the reduced one.
means = []
for track in file_spec.tracks:
means.append(
{
field_name: _mean_over_axis(
array, axis, log_domain=field_name in _LOG_DOMAIN_FIELDS
)
for field_name, (array, shape_spec) in _data_arrays(track).items()
if (axis := _dim_axis(array, shape_spec, dim_name)) is not None
}
)
compounded = slice_spec_dims(file_spec, **{dim_name: [0]})
for track, track_means in zip(compounded.tracks, means):
for field_name, array in track_means.items():
_set_data_array(track, field_name, array)
compounded.__post_init__()
return compounded
[docs]
def sum_data(input_paths: Sequence[str | Path], output_path: str | Path, overwrite=False):
"""
Sums multiple raw data files and saves the result to a new file.
For images, this will actually average the images. If the images are uint8, it will average
directly. If the images are float32, we assume they are in the log-domain and we will do the
averaging in the linear domain.
Args:
input_paths (list[str, Path]): List of paths to the input raw data files. Each path
may be a single file, a folder, or an ``hf://`` path; folders (local or ``hf://``)
are expanded into all zea files they contain (using :class:`zea.Dataset`).
output_path (Path): Path to the output file where the summed data will be saved.
overwrite (bool, optional): Whether to overwrite the output file if it exists.
Defaults to False.
"""
_prepare_output_path(str(output_path), overwrite)
# validate=False: this only expands the paths; every file is opened (and so
# validated) again below.
with Dataset(input_paths, validate=False) as dataset:
input_paths = [file.path for file in dataset]
with File(input_paths[0]) as f:
file_spec = f._to_file_spec()
# Images are accumulated in the linear domain (float32 dB) or in float (uint8) and
# converted back to an average once every file has been added.
totals = []
for track in file_spec.tracks:
track_totals = {}
for field_name, (array, _) in _data_arrays(track).items():
if field_name in _LOG_DOMAIN_FIELDS:
track_totals[field_name] = (
np.power(10.0, array / _DB_FACTOR)
if array.dtype == np.float32
else array.astype(np.float32)
)
elif np.issubdtype(array.dtype, np.integer):
# Widen integer accumulators (raw_data may be int16) so summing many files
# cannot overflow; the sum is validated and cast back to the input dtype below.
track_totals[field_name] = array.astype(np.int64)
else:
track_totals[field_name] = array.copy()
totals.append(track_totals)
for path in input_paths[1:]:
with File(path) as f:
other = f._to_file_spec()
if len(other.tracks) != len(file_spec.tracks):
raise ValueError(
f"{path} has {len(other.tracks)} track(s), but {input_paths[0]} has "
f"{len(file_spec.tracks)}. Only files with matching tracks can be summed."
)
for i, (track, other_track) in enumerate(zip(file_spec.tracks, other.tracks)):
_assert_scans_equal(track.scan, other_track.scan, f"tracks[{i}].scan")
other_arrays = _data_arrays(other_track)
extra_fields = set(other_arrays) - set(totals[i])
if extra_fields:
raise ValueError(
f"{path} has field(s) {sorted(extra_fields)} not present in "
f"{input_paths[0]}. Only files with matching data products can be summed."
)
for field_name, total in totals[i].items():
if field_name not in other_arrays:
raise ValueError(
f"{path} is missing '{field_name}', present in {input_paths[0]}."
)
array, _ = other_arrays[field_name]
_assert_shapes_equal(total, array, field_name)
if field_name in _LOG_DOMAIN_FIELDS:
total += (
np.power(10.0, array / _DB_FACTOR)
if array.dtype == np.float32
else array.astype(np.float32)
)
else:
total += array
n_files = len(input_paths)
for track, track_totals in zip(file_spec.tracks, totals):
field_dtypes = {name: array.dtype for name, (array, _) in _data_arrays(track).items()}
for field_name, total in track_totals.items():
dtype = field_dtypes[field_name]
if field_name in _LOG_DOMAIN_FIELDS:
if dtype == np.float32:
# Back to dB (20*log10), and clamp: log-compressed images are <= 0 dB.
total = np.minimum(_DB_FACTOR * np.log10(total / n_files), 0.0).astype(
np.float32
)
else:
total = np.clip(total / n_files, 0, 255).astype(np.uint8)
elif np.issubdtype(dtype, np.integer):
# Cast the widened sum back to the input dtype, refusing silent overflow.
info = np.iinfo(dtype)
if total.min() < info.min or total.max() > info.max:
raise ValueError(
f"Summed '{field_name}' does not fit {np.dtype(dtype).name}: value "
f"range [{int(total.min())}, {int(total.max())}] exceeds "
f"[{info.min}, {info.max}]."
)
total = total.astype(dtype)
_set_data_array(track, field_name, total)
file_spec.__post_init__()
file_spec.save(str(output_path))
def _assert_shapes_equal(array0, array1, name="array"):
# Raise explicitly (not ``assert``) so validation survives ``python -O``.
shape0, shape1 = array0.shape, array1.shape
if shape0 != shape1:
raise ValueError(f"{name} shapes do not match. Got {shape0} and {shape1}.")
def _assert_scans_equal(scan, other_scan, name="scan"):
"""Check two ScanSpecs describe the same acquisition, field by field."""
# Raise explicitly (not ``assert``) so validation survives ``python -O``.
if scan is None or other_scan is None:
if scan is not other_scan:
raise ValueError(f"{name}: one file has scan parameters and the other does not.")
return
for field_name in scan.SCHEMA:
value, other_value = getattr(scan, field_name), getattr(other_scan, field_name)
if not np.array_equal(value, other_value):
raise ValueError(
f"{name}.{field_name} does not match across the summed files: "
f"{value} vs {other_value}."
)
[docs]
@_supports_folders
def compound_frames(input_path: str | Path, output_path: str | Path, overwrite=False):
"""
Compounds frames in a raw data file by averaging them.
Args:
input_path (str, Path): Path to the input raw data file, or a folder of files.
Also accepts an ``hf://`` path (file or folder).
output_path (Path): Path to the output file (or folder) where the compounded
data will be saved.
overwrite (bool, optional): Whether to overwrite the output file if it exists.
Defaults to False.
"""
_prepare_output_path(str(output_path), overwrite, input_path)
# Lazy: the average is accumulated slab by slab, so only the compounded frame (and one
# slab at a time) is ever in memory. The file stays open until the result is written.
with File(input_path) as f:
file_spec = _compound_named_dim(f._to_file_spec(lazy=True), "n_frames")
file_spec.save(str(output_path))
[docs]
@_supports_folders
def compound_transmits(input_path: str | Path, output_path: str | Path, overwrite=False):
"""
Compounds transmits in a raw data file by averaging them.
Note:
This function assumes that all transmits are identical. If this is not the case the
function will result in incorrect scan parameters.
Args:
input_path (str, Path): Path to the input raw data file, or a folder of files.
Also accepts an ``hf://`` path (file or folder).
output_path (Path): Path to the output file (or folder) where the compounded
data will be saved.
overwrite (bool, optional): Whether to overwrite the output file if it exists.
Defaults to False.
"""
_prepare_output_path(str(output_path), overwrite, input_path)
# Lazy, as in compound_frames: the data is reduced one slab at a time.
with File(input_path) as f:
file_spec = f._to_file_spec(lazy=True)
for i, track in enumerate(file_spec.tracks):
if track.scan is not None and not _all_tx_are_identical(track.scan):
log.warning(
f"Not all transmits in track {i} are identical. Compounding transmits may "
"lead to unexpected results."
)
file_spec = _compound_named_dim(file_spec, "n_tx")
file_spec.save(str(output_path))
def _all_tx_are_identical(scan: ScanSpec):
"""Checks if all transmits in a ScanSpec are identical."""
attributes_to_check = [
scan.polar_angles,
scan.azimuth_angles,
scan.t0_delays,
scan.tx_apodizations,
scan.focus_distances,
scan.transmit_origins,
scan.initial_times,
]
for attr in attributes_to_check:
if attr is not None and not _check_all_identical(attr, axis=0):
return False
return True
def _check_all_identical(array, axis=0):
"""Checks if all elements along a given axis are identical."""
first = array.take(0, axis=axis)
return np.all(np.equal(array, first), axis=axis).all()
[docs]
@_supports_folders
def resave(
input_path: str | Path,
output_path: str | Path,
overwrite=False,
**kwargs,
):
"""
Resaves a zea data file to a new location.
The data is never held in memory: every array is copied from the input to the output
one slab at a time (see :func:`zea.data.spec.iter_slabs`), so a file far larger than
RAM can be rewritten — to pick up a new zea version, a different chunking, or another
compression codec.
Args:
input_path (str, Path): Path to the input zea data file, or a folder of files.
Also accepts an ``hf://`` path (file or folder).
output_path (Path): Path to the output file (or folder) where the data will be saved.
overwrite (bool, optional): Whether to overwrite the output file if it exists.
Defaults to False.
chunk_axes (tuple, optional): Dimension names to chunk with size 1. Defaults
to ``("n_frames",)`` — one full frame per chunk — so partial and streamed
reads fetch only the requested frames; other axes stay at full extent. Use
``None``/``()`` for contiguous storage. See
:meth:`zea.data.spec.Spec._resolve_chunks`.
**kwargs: Passed to :meth:`zea.data.spec.FileSpec.save` (e.g. ``compression``).
"""
_prepare_output_path(str(output_path), overwrite, input_path)
if not Path(input_path).exists():
raise FileNotFoundError(f"Input path {input_path} does not exist.")
# The spec's arrays stay backed by the input file, so it has to stay open until the
# last slab has been copied across.
with File(input_path) as f:
file_spec = f._to_file_spec(lazy=True)
file_spec.save(str(output_path), **kwargs)
def _is_select_all(index) -> bool:
"""Whether ``index`` selects everything, so slicing along it is a no-op."""
return isinstance(index, slice) and index == slice(None)
def _named_axes(matched_shape: tuple, ndim: int, dim_name: str) -> list[int]:
"""Axes of a value's actual shape that carry ``dim_name`` in its SCHEMA shape.
Handles the ``"..."`` wildcard: dimensions before it are counted from the front
of the value's shape, dimensions after it from the back.
"""
if "..." not in matched_shape:
return [axis for axis, dim in enumerate(matched_shape) if dim == dim_name]
ellipsis_axis = matched_shape.index("...")
prefix, suffix = matched_shape[:ellipsis_axis], matched_shape[ellipsis_axis + 1 :]
return [axis for axis, dim in enumerate(prefix) if dim == dim_name] + [
ndim - len(suffix) + axis for axis, dim in enumerate(suffix) if dim == dim_name
]
def _as_indices(index, size: int) -> np.ndarray:
"""Normalize any supported index into a 1-D integer array.
A plain ``int`` is kept as a length-1 array so that the axis survives the
selection: dropping it would break the SCHEMA shape of the field.
"""
return np.atleast_1d(np.arange(size)[index])
def _take(value, indices: np.ndarray, axis: int) -> np.ndarray:
"""``np.take``, reading only the selected entries when ``value`` is still on disk.
This is what makes extracting a few frames from a large file cost a few frames of IO
and memory rather than the whole array. HDF5 only accepts a strictly increasing
selection, so the unique indices are read and then reordered (or repeated) in memory,
which is bounded by the size of the result.
"""
if not is_lazy_array(value):
return np.take(value, indices, axis=axis)
unique = np.unique(indices)
selection: list = [slice(None)] * value.ndim
selection[axis] = unique.tolist()
selected = value[tuple(selection)]
if unique.size == indices.size and np.array_equal(unique, indices):
return selected
return np.take(selected, np.searchsorted(unique, indices), axis=axis)
def _slice_map_coordinates(spec: Spec, sliced: Spec, index) -> None:
"""Slice a :class:`~zea.data.spec.Map`'s coordinates along its frame axis, if it has one.
Coordinates are optional and may either carry a leading ``n_frames`` axis or omit it
and broadcast across frames. Only the former is sliced; the shapes accepted here are
the frame-carrying ones from :meth:`Map.__post_init__`.
"""
coordinates = getattr(spec, "coordinates", None)
values = getattr(spec, "values", None)
if index is None or coordinates is None or values is None:
return
if coordinates.shape[:-1] in (values.shape, values.shape[:-1]):
indices = _as_indices(index, coordinates.shape[0])
setattr(sliced, "coordinates", _take(coordinates, indices, axis=0))
def _slice_spec_dims(spec: SpecT, dim_indices: dict) -> SpecT:
"""Recursive worker for :func:`slice_spec_dims` (no re-validation)."""
sliced = shallow_copy(spec)
sliced_timestamps = False
for field_name, field_info in spec.SCHEMA.items():
value = getattr(spec, field_name)
if value is None:
continue
if field_info.get("spec") is not None:
setattr(sliced, field_name, _slice_spec_dims(value, dim_indices))
continue
# Which of the SCHEMA's candidate shapes applies is decided by the value itself,
# so e.g. a scalar center_frequency is left alone while a per-transmit one is sliced.
matched_shape = find_matched_shape(value, Spec._expected_shapes(field_info["shape"]))
if matched_shape is None:
continue
for dim_name, index in dim_indices.items():
for axis in _named_axes(matched_shape, np.ndim(value), dim_name):
# One axis at a time, so multiple index *lists* do not broadcast
# against each other the way they would in a single fancy-index.
value = _take(value, _as_indices(index, value.shape[axis]), axis=axis)
sliced_timestamps |= field_name == "timestamps"
setattr(sliced, field_name, value)
# Map.coordinates is declared as ("...", 3): it names no dimension, and its leading
# frame axis is optional (coordinates may broadcast across frames). So whether it
# carries n_frames has to be decided against values, the way Map.__post_init__ does.
_slice_map_coordinates(spec, sliced, dim_indices.get("n_frames"))
# Map.timestamps must start at zero, so re-anchor the selection and push the
# dropped lead time into start_time_offset to keep the spec valid.
start_time_offset = getattr(sliced, "start_time_offset", None)
if sliced_timestamps and start_time_offset is not None:
timestamps = getattr(sliced, "timestamps")
setattr(sliced, "start_time_offset", start_time_offset + timestamps[0])
setattr(sliced, "timestamps", timestamps - timestamps[0])
# tracks sit outside SCHEMA (see FileSpec._SCHEMA_EXCLUDED_FIELDS), so recurse explicitly.
if isinstance(sliced, FileSpec):
sliced.tracks = [_slice_spec_dims(track, dim_indices) for track in sliced.tracks]
return sliced
[docs]
def slice_spec_dims(spec: SpecT, **dim_indices) -> SpecT:
"""Index a spec along named SCHEMA dimensions, wherever they occur.
Every field whose ``SCHEMA`` shape mentions one of the given dimension names is
indexed along that dimension, throughout the whole spec tree. So slicing
``n_frames`` also selects the matching entries of ``timestamps``, ``annotations``
and ``metrics``, and slicing ``n_tx`` also selects the matching ``t0_delays``,
``tx_apodizations`` and ``time_to_next_transmit`` — without this function needing
to know those fields exist.
Fields that do not carry the dimension are passed through untouched, and arrays
are shared with ``spec`` rather than copied where possible.
Note:
``custom`` elements are opaque to the schema and are never sliced.
Args:
spec (Spec): The spec to slice. Not modified.
**dim_indices: Index per named dimension, e.g. ``n_frames=[0, 2]``,
``n_tx=slice(0, 4)``. Values may be ints, lists, arrays, boolean masks
or slices, and unlike a single fancy-index expression, several
dimensions may be given as lists at the same time.
Returns:
Spec: A validated copy of ``spec`` with the requested dimensions indexed.
Example:
.. code-block:: python
sliced = slice_spec_dims(file_spec, n_frames=[0, 2], n_tx=slice(0, 4))
"""
unknown = set(dim_indices) - CONSISTENCY_DIMENSIONS
if unknown:
raise ValueError(
f"Unknown dimension(s) {sorted(unknown)}. Only dimensions that the schema keeps "
f"consistent across fields can be sliced: {sorted(CONSISTENCY_DIMENSIONS)}."
)
dim_indices = {dim: index for dim, index in dim_indices.items() if not _is_select_all(index)}
if not dim_indices:
return spec
sliced = _slice_spec_dims(spec, dim_indices)
# Re-run validation once, at the top: catches any dimension the slicing left
# inconsistent across the tree.
sliced.__post_init__()
return sliced
[docs]
def summary(input_path: str | Path):
"""Prints a summary of a zea data file to the console.
Args:
input_path (str, Path): Path to the zea data file. Also accepts an ``hf://`` path.
"""
with File(input_path) as f:
f.summary()
[docs]
def copy(src: str | Path, dst: str | Path, key: str, mode: str | None = None):
"""Copies zea files to a new location using :meth:`zea.Dataset.copy`.
Args:
src (str, Path): Source path. Can be a single file, a list of files, a folder,
or an ``hf://`` path.
dst (Path): Destination folder path.
key (str): Key to access in the HDF5 files. Use ``"all"`` or ``"*"`` to copy
everything.
mode (str, optional): HDF5 file mode for the destination files. Defaults to
None, which lets :meth:`zea.Dataset.copy` auto-select the mode (``"a"`` for a
single key, ``"w"`` when ``key`` is ``"all"``/``"*"``).
"""
dataset = Dataset(src)
dataset.copy(dst, key, mode=mode)
[docs]
@_supports_folders
def decode_hadamard_file_operation(input_path: Path, output_path: Path, overwrite=False):
"""Decodes Hadamard-encoded data in a zea file.
Args:
input_path: Path to the input zea data file, or a folder of files.
output_path: Path to the output file (or folder) where the decoded
data will be saved.
overwrite: Whether to overwrite the output file if it exists.
"""
_prepare_output_path(str(output_path), overwrite, input_path)
with File(input_path) as f:
file_spec = f._to_file_spec()
for track in file_spec.tracks:
try:
raw_data = getattr(track.data, "raw_data")
if raw_data is None:
raise ValueError("No raw_data found in the input file.")
tx_apodizations = track.scan.tx_apodizations
raw_data, tx_apodizations = decode_hadamard(raw_data, tx_apodizations)
track.scan.tx_apodizations = ops.convert_to_numpy(tx_apodizations)
_set_data_array(track, "raw_data", ops.convert_to_numpy(raw_data))
except ValueError:
log.warning(
f"Failed to decode Hadamard data for track {track}. "
"The raw_data will remain unchanged."
)
if overwrite:
_delete_file_if_exists(output_path)
file_spec.save(
path=str(output_path),
)
[docs]
@_supports_folders
def sa_to_virtual_focus(
input_path: Path,
output_path: Path,
polar_angle: float,
azimuth_angle: float,
focus_distance: float,
transmit_origin: tuple[float, float, float] = (0.0, 0.0, 0.0),
tx_apodization: str = "kaiser",
overwrite=False,
):
"""
Constructs a new zea file with the given transmit scheme from a synthetic aperture (SA)
acquisition. The new file will have a single transmit with the specified polar and azimuth
angles, focus distance, and transmit origin.
"""
_prepare_output_path(str(output_path), overwrite, input_path)
with File(input_path) as f:
file_spec = f._to_file_spec()
probe = file_spec.probe
if not isinstance(probe, ProbeSpec) or probe.probe_geometry is None:
raise ValueError("sa_to_virtual_focus requires a probe with a defined probe_geometry.")
probe_geometry = probe.probe_geometry
n_el = probe_geometry.shape[0]
for track in file_spec.tracks:
scan = track.scan
if tx_apodization is None:
tx_apodization = ops.ones((1, n_el), dtype="float32")
elif tx_apodization == "kaiser":
tx_apodization = ops.expand_dims(ops.cast(ops.kaiser(n_el, beta=5.0), "float32"), 0)
elif tx_apodization == "hanning":
tx_apodization = ops.expand_dims(ops.cast(ops.hanning(n_el), "float32"), 0)
raw_data, t0_delays = construct_acquisition_from_synthetic_aperture(
raw_data=getattr(track.data, "raw_data"),
probe_geometry=probe_geometry,
sampling_frequency=scan.sampling_frequency,
polar_angle=polar_angle,
azimuth_angle=azimuth_angle,
focus_distance=focus_distance,
transmit_origin=transmit_origin,
sound_speed=scan.sound_speed,
tx_apodization=tx_apodization,
)
_set_data_array(track, "raw_data", ops.convert_to_numpy(raw_data))
track.scan.t0_delays = np.asarray(t0_delays)
track.scan.tx_apodizations = ops.convert_to_numpy(tx_apodization)
track.scan.polar_angles = np.array([polar_angle])
track.scan.azimuth_angles = np.array([azimuth_angle])
track.scan.focus_distances = np.array([focus_distance])
track.scan.transmit_origins = np.array([transmit_origin])
if track.scan.initial_times is not None:
track.scan.initial_times = track.scan.initial_times[:1]
if track.scan.waveforms_one_way is not None:
track.scan.waveforms_one_way = track.scan.waveforms_one_way[:1]
if track.scan.waveforms_two_way is not None:
track.scan.waveforms_two_way = track.scan.waveforms_two_way[:1]
if track.scan.time_to_next_transmit is not None:
track.scan.time_to_next_transmit = track.scan.time_to_next_transmit[:, :1]
file_spec.save(
path=str(output_path),
)
def _delete_file_if_exists(path: Path):
"""Deletes a file if it exists."""
if path.exists():
path.unlink()
[docs]
def output_blocked(output_path: str | Path, overwrite: bool) -> bool:
"""Whether ``output_path`` exists and would be silently clobbered.
The single predicate behind every "refuse to overwrite an output" guard in
:mod:`zea`, so the definition of "blocked" (currently: the path exists and
``overwrite`` was not requested) only needs to change in one place.
"""
return Path(output_path).exists() and not overwrite
def _prepare_output_path(output_path: str, overwrite: bool, input_path: "str | Path | None" = None):
"""Guard the save target, matching :func:`resave`: refuse to clobber unless asked.
``FileSpec.save`` overwrites atomically and has no ``overwrite`` flag of its own, so
callers must enforce it here or an existing file is silently replaced.
Also refuses to save to an ``hf://`` path, which is read-only, and — when the caller
names its ``input_path`` — to write back over the file being read: an operation reads
its input while it writes the output, so rewriting in place would delete the data
before it has been copied.
"""
if output_path.startswith(HF_PREFIX):
raise ValueError(
f"Cannot save to an 'hf://' path: {output_path}. 'hf://' paths are read-only; "
"save to a local path instead."
)
if input_path is not None and Path(output_path).resolve() == Path(input_path).resolve():
raise ValueError(
f"Output path {output_path} is the input file. The input is read while the "
"output is written, so it cannot be rewritten in place: write to a new path "
"and replace the original afterwards."
)
if output_blocked(output_path, overwrite):
raise FileExistsError(
f"Output path {output_path} already exists. Use overwrite=True to overwrite."
)
if overwrite:
_delete_file_if_exists(Path(output_path))
def _interpret_index(input_str):
if "-" in input_str:
start, end = map(int, input_str.split("-"))
return list(range(start, end + 1))
else:
return [int(x) for x in input_str.split(" ")]
def _interpret_indices(input_str_list):
if input_str_list == "all" or input_str_list == ["all"]:
return slice(None)
if len(input_str_list) == 1 and "-" in input_str_list[0]:
start, end = map(int, input_str_list[0].split("-"))
return slice(start, end + 1)
indices = []
for part in input_str_list:
indices.extend(_interpret_index(part))
return indices
# ── Command line interface (tyro) ────────────────────────────────────────────
#
# The CLI subcommand dataclasses live in :mod:`zea.cli_args` (a light-import
# module) so that ``zea data …`` can be wired into the top-level ``zea`` CLI
# without importing this heavy module just to render ``--help``. The dataclasses
# there dispatch back to the operation functions above.
[docs]
def main():
"""Parse command line arguments and run the requested data operation.
Entry point for ``python -m zea.data``. This is equivalent to ``zea data``.
"""
from zea.cli_args import DataCommand, _run_data_command
args = tyro.cli(DataCommand)
_run_data_command(args)
if __name__ == "__main__":
main()