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

zea/ops.py

@@ -11,39 +11,62 @@ Operations can be run on their own:
 
 Examples
 ^^^^^^^^
-.. code-block:: python
+.. doctest::
 
-    data = np.random.randn(2000, 128, 1)
-    # static arguments are passed in the constructor
-    envelope_detect = EnvelopeDetect(axis=-1)
-    # other parameters can be passed here along with the data
-    envelope_data = envelope_detect(data=data)
+    >>> import numpy as np
+    >>> from zea.ops import EnvelopeDetect
+    >>> data = np.random.randn(2000, 128, 1)
+    >>> # static arguments are passed in the constructor
+    >>> envelope_detect = EnvelopeDetect(axis=-1)
+    >>> # other parameters can be passed here along with the data
+    >>> envelope_data = envelope_detect(data=data)
 
 Using a pipeline
 ----------------
 
 You can initialize with a default pipeline or create your own custom pipeline.
 
-.. code-block:: python
+.. doctest::
 
-    pipeline = Pipeline.from_default()
+    >>> from zea.ops import Pipeline, EnvelopeDetect, Normalize, LogCompress
+    >>> pipeline = Pipeline.from_default()
 
-    operations = [
-        EnvelopeDetect(),
-        Normalize(),
-        LogCompress(),
-    ]
-    pipeline_custom = Pipeline(operations)
+    >>> operations = [
+    ...     EnvelopeDetect(),
+    ...     Normalize(),
+    ...     LogCompress(),
+    ... ]
+    >>> pipeline_custom = Pipeline(operations)
 
 One can also load a pipeline from a config or yaml/json file:
 
-.. code-block:: python
+.. doctest::
+
+    >>> from zea import Pipeline
+
+    >>> # From JSON string
+    >>> json_string = '{"operations": ["identity"]}'
+    >>> pipeline = Pipeline.from_json(json_string)
+
+    >>> # from yaml file
+    >>> import yaml
+    >>> from zea import Config
+    >>> # Create a sample pipeline YAML file
+    >>> pipeline_dict = {
+    ...     "operations": [
+    ...         {"name": "identity"},
+    ...     ]
+    ... }
+    >>> with open("pipeline.yaml", "w") as f:
+    ...     yaml.dump(pipeline_dict, f)
+    >>> yaml_file = "pipeline.yaml"
+    >>> pipeline = Pipeline.from_yaml(yaml_file)
+
+.. testcleanup::
 
-    json_string = '{"operations": ["identity"]}'
-    pipeline = Pipeline.from_json(json_string)
+    import os
 
-    yaml_file = "pipeline.yaml"
-    pipeline = Pipeline.from_yaml(yaml_file)
+    os.remove("pipeline.yaml")
 
 Example of a yaml file:
 
@@ -56,8 +79,6 @@ Example of a yaml file:
           params:
             operations:
               - name: tof_correction
-                params:
-                  apply_phase_rotation: true
               - name: pfield_weighting
               - name: delay_and_sum
             num_patches: 100
@@ -79,7 +100,7 @@ import numpy as np
 import scipy
 import yaml
 from keras import ops
-from keras.src.layers.preprocessing.tf_data_layer import TFDataLayer
+from keras.src.layers.preprocessing.data_layer import DataLayer
 
 from zea import log
 from zea.backend import jit
@@ -99,8 +120,12 @@ from zea.internal.registry import ops_registry
 from zea.probes import Probe
 from zea.scan import Scan
 from zea.simulator import simulate_rf
-from zea.tensor_ops import batched_map, patched_map, resample, reshape_axis
-from zea.utils import FunctionTimer, deep_compare, map_negative_indices, translate
+from zea.tensor_ops import resample, reshape_axis, translate, vmap
+from zea.utils import (
+    FunctionTimer,
+    deep_compare,
+    map_negative_indices,
+)
 
 
 def get_ops(ops_name):
@@ -125,6 +150,7 @@ class Operation(keras.Operation):
         with_batch_dim: bool = True,
         jit_kwargs: dict | None = None,
         jittable: bool = True,
+        additional_output_keys: List[str] = None,
         **kwargs,
     ):
         """
@@ -143,6 +169,10 @@ class Operation(keras.Operation):
             with_batch_dim: Whether operations should expect a batch dimension in the input
             jit_kwargs: Additional keyword arguments for the JIT compiler
             jittable: Whether the operation can be JIT compiled
+            additional_output_keys: A list of additional output keys produced by the operation.
+                These are used to track if all keys are available for downstream operations.
+                If the operation has a conditional output, it is best to add all possible
+                output keys here.
         """
         super().__init__(**kwargs)
 
@@ -153,6 +183,7 @@ class Operation(keras.Operation):
         self.output_key = output_key  # Key for output data
         if self.output_key is None:
             self.output_key = self.key
+        self.additional_output_keys = additional_output_keys or []
 
         self.inputs = []  # Source(s) of input data (name of a previous operation)
         self.allow_multiple_inputs = False  # Only single input allowed by default
@@ -165,8 +196,6 @@ class Operation(keras.Operation):
         self._output_cache = {}
 
         # Obtain the input signature of the `call` method
-        self._input_signature = None
-        self._valid_keys = None  # Keys valid for the `call` method
         self._trace_signatures()
 
         if jit_kwargs is None:
@@ -186,6 +215,11 @@ class Operation(keras.Operation):
         with log.set_level("ERROR"):
             self.set_jit(jit_compile)
 
+    @property
+    def output_keys(self) -> List[str]:
+        """Get the output keys of the operation."""
+        return [self.output_key] + self.additional_output_keys
+
     @property
     def static_params(self):
         """Get the static parameters of the operation."""
@@ -207,10 +241,15 @@ class Operation(keras.Operation):
         self._valid_keys = set(self._input_signature.parameters.keys())
 
     @property
-    def valid_keys(self):
+    def valid_keys(self) -> set:
         """Get the valid keys for the `call` method."""
         return self._valid_keys
 
+    @property
+    def needs_keys(self) -> set:
+        """Get a set of all input keys needed by the operation."""
+        return self.valid_keys
+
     @property
     def jittable(self):
         """Check if the operation can be JIT compiled."""
@@ -408,8 +447,6 @@ class Pipeline:
         """
         self._call_pipeline = self.call
         self.name = name
-        self.timer = FunctionTimer()
-        self.timed = timed
 
         self._pipeline_layers = operations
 
@@ -419,6 +456,24 @@ class Pipeline:
         self.with_batch_dim = with_batch_dim
         self._validate_flag = validate
 
+        # Setup timer
+        if jit_options == "pipeline" and timed:
+            raise ValueError(
+                "timed=True cannot be used with jit_options='pipeline' as the entire "
+                "pipeline is compiled into a single function. Try setting jit_options to "
+                "'ops' or None."
+            )
+        if timed:
+            log.warning(
+                "Timer has been initialized for the pipeline. To get an accurate timing estimate, "
+                "the `block_until_ready()` is used, which will slow down the execution, so "
+                "do not use for regular processing!"
+            )
+            self._callable_layers = self._get_timed_operations()
+        else:
+            self._callable_layers = self._pipeline_layers
+        self._timed = timed
+
         if validate:
             self.validate()
         else:
@@ -427,19 +482,30 @@ class Pipeline:
         if jit_kwargs is None:
             jit_kwargs = {}
 
-        if keras.backend.backend() == "jax" and self.static_params:
+        if keras.backend.backend() == "jax" and self.static_params != []:
             jit_kwargs = {"static_argnames": self.static_params}
 
         self.jit_kwargs = jit_kwargs
         self.jit_options = jit_options  # will handle the jit compilation
 
     def needs(self, key) -> bool:
-        """Check if the pipeline needs a specific key."""
-        return key in self.valid_keys
+        """Check if the pipeline needs a specific key at the input."""
+        return key in self.needs_keys
+
+    @property
+    def output_keys(self) -> set:
+        """All output keys the pipeline guarantees to produce."""
+        output_keys = set()
+        for operation in self.operations:
+            output_keys.update(operation.output_keys)
+        return output_keys
 
     @property
     def valid_keys(self) -> set:
-        """Get a set of valid keys for the pipeline."""
+        """Get a set of valid keys for the pipeline.
+
+        This is all keys that can be passed to the pipeline as input.
+        """
         valid_keys = set()
         for operation in self.operations:
             valid_keys.update(operation.valid_keys)
@@ -453,8 +519,26 @@ class Pipeline:
             static_params.extend(operation.static_params)
         return list(set(static_params))
 
+    @property
+    def needs_keys(self) -> set:
+        """Get a set of all input keys needed by the pipeline.
+
+        Will keep track of keys that are already provided by previous operations.
+        """
+        needs = set()
+        has_so_far = set()
+        previous_operation = None
+        for operation in self.operations:
+            if previous_operation is not None:
+                has_so_far.update(previous_operation.output_keys)
+            needs.update(operation.needs_keys - has_so_far)
+            previous_operation = operation
+        return needs
+
     @classmethod
-    def from_default(cls, num_patches=100, baseband=False, pfield=False, **kwargs) -> "Pipeline":
+    def from_default(
+        cls, num_patches=100, baseband=False, pfield=False, timed=False, **kwargs
+    ) -> "Pipeline":
         """Create a default pipeline.
 
         Args:
@@ -465,6 +549,7 @@ class Pipeline:
                 so input signal has a single channel dim and is still on carrier frequency.
             pfield (bool): If True, apply Pfield weighting. Defaults to False.
                 This will calculate pressure field and only beamform the data to those locations.
+            timed (bool, optional): Whether to time each operation. Defaults to False.
             **kwargs: Additional keyword arguments to be passed to the Pipeline constructor.
 
         """
@@ -476,7 +561,7 @@ class Pipeline:
 
         # Get beamforming ops
         beamforming = [
-            TOFCorrection(apply_phase_rotation=True),
+            TOFCorrection(),
             DelayAndSum(),
         ]
         if pfield:
@@ -495,7 +580,7 @@ class Pipeline:
             Normalize(),
             LogCompress(),
         ]
-        return cls(operations, **kwargs)
+        return cls(operations, timed=timed, **kwargs)
 
     def copy(self) -> "Pipeline":
         """Create a copy of the pipeline."""
@@ -506,57 +591,60 @@ class Pipeline:
             jit_kwargs=self.jit_kwargs,
             name=self.name,
             validate=self._validate_flag,
+            timed=self._timed,
+        )
+
+    def reinitialize(self):
+        """Reinitialize the pipeline in place."""
+        self.__init__(
+            self._pipeline_layers,
+            with_batch_dim=self.with_batch_dim,
+            jit_options=self.jit_options,
+            jit_kwargs=self.jit_kwargs,
+            name=self.name,
+            validate=self._validate_flag,
+            timed=self._timed,
         )
 
     def prepend(self, operation: Operation):
         """Prepend an operation to the pipeline."""
         self._pipeline_layers.insert(0, operation)
-        self.copy()
+        self.reinitialize()
 
     def append(self, operation: Operation):
         """Append an operation to the pipeline."""
         self._pipeline_layers.append(operation)
-        self.copy()
+        self.reinitialize()
 
     def insert(self, index: int, operation: Operation):
         """Insert an operation at a specific index in the pipeline."""
         if index < 0 or index > len(self._pipeline_layers):
             raise IndexError("Index out of bounds for inserting operation.")
         self._pipeline_layers.insert(index, operation)
-        return self.copy()
+        self.reinitialize()
 
     @property
     def operations(self):
         """Alias for self.layers to match the zea naming convention"""
         return self._pipeline_layers
 
-    def timed_call(self, **inputs):
-        """Process input data through the pipeline."""
+    def reset_timer(self):
+        """Reset the timer for timed operations."""
+        if self._timed:
+            self._callable_layers = self._get_timed_operations()
+        else:
+            log.warning(
+                "Timer has not been initialized. Set timed=True when initializing the pipeline."
+            )
 
-        for op in self._pipeline_layers:
-            timed_op = self.timer(op, name=op.__class__.__name__)
-            try:
-                outputs = timed_op(**inputs)
-            except KeyError as exc:
-                raise KeyError(
-                    f"[zea.Pipeline] Operation '{op.__class__.__name__}' "
-                    f"requires input key '{exc.args[0]}', "
-                    "but it was not provided in the inputs.\n"
-                    "Check whether the objects (such as `zea.Scan`) passed to "
-                    "`pipeline.prepare_parameters()` contain all required keys.\n"
-                    f"Current list of all passed keys: {list(inputs.keys())}\n"
-                    f"Valid keys for this pipeline: {self.valid_keys}"
-                ) from exc
-            except Exception as exc:
-                raise RuntimeError(
-                    f"[zea.Pipeline] Error in operation '{op.__class__.__name__}': {exc}"
-                ) from exc
-            inputs = outputs
-        return outputs
+    def _get_timed_operations(self):
+        """Get a list of timed operations."""
+        self.timer = FunctionTimer()
+        return [self.timer(op, name=op.__class__.__name__) for op in self._pipeline_layers]
 
     def call(self, **inputs):
         """Process input data through the pipeline."""
-        for operation in self._pipeline_layers:
+        for operation in self._callable_layers:
             try:
                 outputs = operation(**inputs)
             except KeyError as exc:
@@ -579,14 +667,9 @@ class Pipeline:
     def __call__(self, return_numpy=False, **inputs):
         """Process input data through the pipeline."""
 
-        if any(key in inputs for key in ["probe", "scan", "config"]):
-            raise ValueError(
-                "Probe, Scan and Config objects should be first processed with "
-                "`Pipeline.prepare_parameters` before calling the pipeline. "
-                "e.g. inputs = Pipeline.prepare_parameters(probe, scan, config)"
-            )
-
-        if any(isinstance(arg, ZEAObject) for arg in inputs.values()):
+        if any(key in inputs for key in ["probe", "scan", "config"]) or any(
+            isinstance(arg, ZEAObject) for arg in inputs.values()
+        ):
             raise ValueError(
                 "Probe, Scan and Config objects should be first processed with "
                 "`Pipeline.prepare_parameters` before calling the pipeline. "
@@ -640,18 +723,13 @@ class Pipeline:
                 if operation.jittable and operation._jit_compile:
                     operation.set_jit(value == "ops")
 
-    @property
-    def _call_fn(self):
-        """Get the call function of the pipeline."""
-        return self.call if not self.timed else self.timed_call
-
     def jit(self):
         """JIT compile the pipeline."""
-        self._call_pipeline = jit(self._call_fn, **self.jit_kwargs)
+        self._call_pipeline = jit(self.call, **self.jit_kwargs)
 
     def unjit(self):
         """Un-JIT compile the pipeline."""
-        self._call_pipeline = self._call_fn
+        self._call_pipeline = self.call
 
     @property
     def jittable(self):
@@ -835,16 +913,17 @@ class Pipeline:
             Must have a ``pipeline`` key with a subkey ``operations``.
 
         Example:
-            .. code-block:: python
-
-                config = Config(
-                    {
-                        "operations": [
-                            "identity",
-                        ],
-                    }
-                )
-                pipeline = Pipeline.from_config(config)
+            .. doctest::
+
+                >>> from zea import Config, Pipeline
+                >>> config = Config(
+                ...     {
+                ...         "operations": [
+                ...             "identity",
+                ...         ],
+                ...     }
+                ... )
+                >>> pipeline = Pipeline.from_config(config)
         """
         return pipeline_from_config(Config(config), **kwargs)
 
@@ -860,9 +939,20 @@ class Pipeline:
             Must have the a `pipeline` key with a subkey `operations`.
 
         Example:
-        ```python
-        pipeline = Pipeline.from_yaml("pipeline.yaml")
-        ```
+            .. doctest::
+
+                >>> import yaml
+                >>> from zea import Config
+                >>> # Create a sample pipeline YAML file
+                >>> pipeline_dict = {
+                ...     "operations": [
+                ...         "identity",
+                ...     ],
+                ... }
+                >>> with open("pipeline.yaml", "w") as f:
+                ...     yaml.dump(pipeline_dict, f)
+                >>> from zea.ops import Pipeline
+                >>> pipeline = Pipeline.from_yaml("pipeline.yaml", jit_options=None)
         """
         return pipeline_from_yaml(file_path, **kwargs)
 
@@ -963,7 +1053,7 @@ class Pipeline:
             assert isinstance(scan, Scan), (
                 f"Expected an instance of `zea.scan.Scan`, got {type(scan)}"
             )
-            scan_dict = scan.to_tensor(include=self.valid_keys, keep_as_is=self.static_params)
+            scan_dict = scan.to_tensor(include=self.needs_keys, keep_as_is=self.static_params)
 
         if config is not None:
             assert isinstance(config, Config), (
@@ -1004,15 +1094,17 @@ def make_operation_chain(
         list: List of operations to be performed.
 
     Example:
-        .. code-block:: python
-
-            chain = make_operation_chain(
-                [
-                    "envelope_detect",
-                    {"name": "normalize", "params": {"output_range": (0, 1)}},
-                    SomeCustomOperation(),
-                ]
-            )
+        .. doctest::
+
+            >>> from zea.ops import make_operation_chain, LogCompress
+            >>> SomeCustomOperation = LogCompress  # just for demonstration
+            >>> chain = make_operation_chain(
+            ...     [
+            ...         "envelope_detect",
+            ...         {"name": "normalize", "params": {"output_range": (0, 1)}},
+            ...         SomeCustomOperation(),
+            ...     ]
+            ... )
     """
     chain = []
     for operation in operation_chain:
@@ -1228,11 +1320,23 @@ class PatchedGrid(Pipeline):
         for operation in self.operations:
             operation.with_batch_dim = False
 
+    @property
+    def _extra_keys(self):
+        return {"flatgrid", "grid_size_x", "grid_size_z"}
+
     @property
     def valid_keys(self) -> set:
-        """Get a set of valid keys for the pipeline. Adds the parameters that PatchedGrid itself
-        operates on (even if not used by operations inside it)."""
-        return super().valid_keys.union({"flatgrid", "grid_size_x", "grid_size_z"})
+        """Get a set of valid keys for the pipeline.
+        Adds the parameters that PatchedGrid itself operates on (even if not used by operations
+        inside it)."""
+        return super().valid_keys.union(self._extra_keys)
+
+    @property
+    def needs_keys(self) -> set:
+        """Get a set of all input keys needed by the pipeline.
+        Adds the parameters that PatchedGrid itself operates on (even if not used by operations
+        inside it)."""
+        return super().needs_keys.union(self._extra_keys)
 
     def call_item(self, inputs):
         """Process data in patches."""
@@ -1242,35 +1346,22 @@ class PatchedGrid(Pipeline):
         grid_size_z = inputs["grid_size_z"]
         flatgrid = inputs.pop("flatgrid")
 
-        # TODO: maybe using n_tx and n_el from kwargs is better but these are tensors now
-        # and this is not supported in broadcast_to
-        n_tx = inputs[self.key].shape[0]
-        n_pix = flatgrid.shape[0]
-        n_el = inputs[self.key].shape[2]
-        inputs["rx_apo"] = ops.broadcast_to(inputs.get("rx_apo", 1.0), (n_tx, n_pix, n_el))
-        inputs["rx_apo"] = ops.swapaxes(inputs["rx_apo"], 0, 1)  # put n_pix first
-
         # Define a list of keys to look up for patching
-        patch_keys = ["flat_pfield", "rx_apo"]
+        flat_pfield = inputs.pop("flat_pfield", None)
 
-        patch_arrays = {}
-        for key in patch_keys:
-            if key in inputs:
-                patch_arrays[key] = inputs.pop(key)
-
-        def patched_call(flatgrid, **patch_kwargs):
-            patch_args = {k: v for k, v in patch_kwargs.items() if v is not None}
-            patch_args["rx_apo"] = ops.swapaxes(patch_args["rx_apo"], 0, 1)
-            out = super(PatchedGrid, self).call(flatgrid=flatgrid, **patch_args, **inputs)
+        def patched_call(flatgrid, flat_pfield):
+            out = super(PatchedGrid, self).call(
+                flatgrid=flatgrid, flat_pfield=flat_pfield, **inputs
+            )
             return out[self.output_key]
 
-        out = patched_map(
+        out = vmap(
             patched_call,
-            flatgrid,
-            self.num_patches,
-            **patch_arrays,
-            jit=bool(self.jit_options),
-        )
+            chunks=self.num_patches,
+            fn_supports_batch=True,
+            disable_jit=not bool(self.jit_options),
+        )(flatgrid, flat_pfield)
+
         return ops.reshape(out, (grid_size_z, grid_size_x, *ops.shape(out)[1:]))
 
     def jittable_call(self, **inputs):
@@ -1309,7 +1400,7 @@ class Identity(Operation):
 
     def call(self, **kwargs) -> Dict:
         """Returns the input as is."""
-        return kwargs
+        return {}
 
 
 @ops_registry("merge")
@@ -1399,6 +1490,7 @@ class Simulate(Operation):
     def __init__(self, **kwargs):
         super().__init__(
             output_data_type=DataTypes.RAW_DATA,
+            additional_output_keys=["n_ch"],
             **kwargs,
         )
 
@@ -1451,18 +1543,16 @@ class TOFCorrection(Operation):
     STATIC_PARAMS = [
         "f_number",
         "apply_lens_correction",
-        "apply_phase_rotation",
         "grid_size_x",
         "grid_size_z",
     ]
 
-    def __init__(self, apply_phase_rotation=True, **kwargs):
+    def __init__(self, **kwargs):
         super().__init__(
             input_data_type=DataTypes.RAW_DATA,
             output_data_type=DataTypes.ALIGNED_DATA,
             **kwargs,
         )
-        self.apply_phase_rotation = apply_phase_rotation
 
     def call(
         self,
@@ -1477,6 +1567,8 @@ class TOFCorrection(Operation):
         tx_apodizations,
         initial_times,
         probe_geometry,
+        t_peak,
+        tx_waveform_indices,
         apply_lens_correction=None,
         lens_thickness=None,
         lens_sound_speed=None,
@@ -1497,6 +1589,9 @@ class TOFCorrection(Operation):
             tx_apodizations (ops.Tensor): Transmit apodizations
             initial_times (ops.Tensor): Initial times
             probe_geometry (ops.Tensor): Probe element positions
+            t_peak (float): Time to peak of the transmit pulse
+            tx_waveform_indices (ops.Tensor): Index of the transmit waveform for each
+                transmit. (All zero if there is only one waveform)
             apply_lens_correction (bool): Whether to apply lens correction
             lens_thickness (float): Lens thickness
             lens_sound_speed (float): Sound speed in the lens
@@ -1507,29 +1602,30 @@ class TOFCorrection(Operation):
 
         raw_data = kwargs[self.key]
 
-        kwargs = {
+        tof_kwargs = {
             "flatgrid": flatgrid,
-            "sound_speed": sound_speed,
-            "angles": polar_angles,
-            "focus_distances": focus_distances,
-            "sampling_frequency": sampling_frequency,
-            "fnum": f_number,
-            "apply_phase_rotation": self.apply_phase_rotation,
-            "demodulation_frequency": demodulation_frequency,
             "t0_delays": t0_delays,
             "tx_apodizations": tx_apodizations,
-            "initial_times": initial_times,
+            "sound_speed": sound_speed,
             "probe_geometry": probe_geometry,
+            "initial_times": initial_times,
+            "sampling_frequency": sampling_frequency,
+            "demodulation_frequency": demodulation_frequency,
+            "f_number": f_number,
+            "polar_angles": polar_angles,
+            "focus_distances": focus_distances,
+            "t_peak": t_peak,
+            "tx_waveform_indices": tx_waveform_indices,
             "apply_lens_correction": apply_lens_correction,
             "lens_thickness": lens_thickness,
             "lens_sound_speed": lens_sound_speed,
         }
 
         if not self.with_batch_dim:
-            tof_corrected = tof_correction(raw_data, **kwargs)
+            tof_corrected = tof_correction(raw_data, **tof_kwargs)
         else:
             tof_corrected = ops.map(
-                lambda data: tof_correction(data, **kwargs),
+                lambda data: tof_correction(data, **tof_kwargs),
                 raw_data,
             )
 
@@ -1587,44 +1683,35 @@ class DelayAndSum(Operation):
         **kwargs,
     ):
         super().__init__(
-            input_data_type=None,
+            input_data_type=DataTypes.ALIGNED_DATA,
             output_data_type=DataTypes.BEAMFORMED_DATA,
             **kwargs,
         )
         self.reshape_grid = reshape_grid
 
-    def process_image(self, data, rx_apo):
+    def process_image(self, data):
         """Performs DAS beamforming on tof-corrected input.
 
         Args:
             data (ops.Tensor): The TOF corrected input of shape `(n_tx, n_pix, n_el, n_ch)`
-            rx_apo (ops.Tensor): Receive apodization window of shape `(n_tx, n_pix, n_el, n_ch)`.
 
         Returns:
             ops.Tensor: The beamformed data of shape `(n_pix, n_ch)`
         """
         # Sum over the channels, i.e. DAS
-        data = ops.sum(rx_apo * data, -2)
+        data = ops.sum(data, -2)
 
         # Sum over transmits, i.e. Compounding
         data = ops.sum(data, 0)
 
         return data
 
-    def call(
-        self,
-        rx_apo=None,
-        grid=None,
-        **kwargs,
-    ):
+    def call(self, grid=None, **kwargs):
         """Performs DAS beamforming on tof-corrected input.
 
         Args:
             tof_corrected_data (ops.Tensor): The TOF corrected input of shape
                 `(n_tx, grid_size_z*grid_size_x, n_el, n_ch)` with optional batch dimension.
-            rx_apo (ops.Tensor): Receive apodization window
-                of shape `(n_tx, grid_size_z*grid_size_x, n_el)`
-                with optional batch dimension. Defaults to 1.0.
 
         Returns:
             dict: Dictionary containing beamformed_data
@@ -1634,17 +1721,11 @@ class DelayAndSum(Operation):
         """
         data = kwargs[self.key]
 
-        if rx_apo is None:
-            rx_apo = ops.ones(1, dtype=ops.dtype(data))
-        rx_apo = ops.broadcast_to(rx_apo[..., None], data.shape)
-
         if not self.with_batch_dim:
-            beamformed_data = self.process_image(data, rx_apo)
+            beamformed_data = self.process_image(data)
         else:
             # Apply process_image to each item in the batch
-            beamformed_data = batched_map(
-                lambda data, rx_apo: self.process_image(data, rx_apo), data, rx_apo=rx_apo
-            )
+            beamformed_data = ops.map(self.process_image, data)
 
         if self.reshape_grid:
             beamformed_data = reshape_axis(
@@ -1654,6 +1735,46 @@ class DelayAndSum(Operation):
         return {self.output_key: beamformed_data}
 
 
+def envelope_detect(data, axis=-3):
+    """Envelope detection of RF signals.
+
+    If the input data is real, it first applies the Hilbert transform along the specified axis
+    and then computes the magnitude of the resulting complex signal.
+    If the input data is complex, it computes the magnitude directly.
+
+    Args:
+        - data (Tensor): The beamformed data of shape (..., grid_size_z, grid_size_x, n_ch).
+        - axis (int): Axis along which to apply the Hilbert transform. Defaults to -3.
+
+    Returns:
+        - envelope_data (Tensor): The envelope detected data
+            of shape (..., grid_size_z, grid_size_x).
+    """
+    if data.shape[-1] == 2:
+        data = channels_to_complex(data)
+    else:
+        n_ax = ops.shape(data)[axis]
+        n_ax_float = ops.cast(n_ax, "float32")
+
+        # Calculate next power of 2: M = 2^ceil(log2(n_ax))
+        # see https://github.com/tue-bmd/zea/discussions/147
+        log2_n_ax = ops.log2(n_ax_float)
+        M = ops.cast(2 ** ops.ceil(log2_n_ax), "int32")
+
+        data = hilbert(data, N=M, axis=axis)
+        indices = ops.arange(n_ax)
+
+        data = ops.take(data, indices, axis=axis)
+        data = ops.squeeze(data, axis=-1)
+
+    # data = ops.abs(data)
+    real = ops.real(data)
+    imag = ops.imag(data)
+    data = ops.sqrt(real**2 + imag**2)
+    data = ops.cast(data, "float32")
+    return data
+
+
 @ops_registry("envelope_detect")
 class EnvelopeDetect(Operation):
     """Envelope detection of RF signals."""
@@ -1680,23 +1801,7 @@ class EnvelopeDetect(Operation):
         """
         data = kwargs[self.key]
 
-        if data.shape[-1] == 2:
-            data = channels_to_complex(data)
-        else:
-            n_ax = data.shape[self.axis]
-            M = 2 ** int(np.ceil(np.log2(n_ax)))
-            # data = scipy.signal.hilbert(data, N=M, axis=self.axis)
-            data = hilbert(data, N=M, axis=self.axis)
-            indices = ops.arange(n_ax)
-
-            data = ops.take(data, indices, axis=self.axis)
-            data = ops.squeeze(data, axis=-1)
-
-        # data = ops.abs(data)
-        real = ops.real(data)
-        imag = ops.imag(data)
-        data = ops.sqrt(real**2 + imag**2)
-        data = ops.cast(data, "float32")
+        data = envelope_detect(data, axis=self.axis)
 
         return {self.output_key: data}
 
@@ -1734,19 +1839,29 @@ class UpMix(Operation):
         return {self.output_key: data}
 
 
+def log_compress(data, eps=1e-16):
+    """Apply logarithmic compression to data."""
+    eps = ops.convert_to_tensor(eps, dtype=data.dtype)
+    data = ops.where(data == 0, eps, data)  # Avoid log(0)
+    return 20 * keras.ops.log10(data)
+
+
 @ops_registry("log_compress")
 class LogCompress(Operation):
     """Logarithmic compression of data."""
 
-    def __init__(
-        self,
-        **kwargs,
-    ):
+    def __init__(self, clip: bool = True, **kwargs):
+        """Initialize the LogCompress operation.
+
+        Args:
+            clip (bool): Whether to clip the output to a dynamic range. Defaults to True.
+        """
         super().__init__(
             input_data_type=DataTypes.ENVELOPE_DATA,
             output_data_type=DataTypes.IMAGE,
             **kwargs,
         )
+        self.clip = clip
 
     def call(self, dynamic_range=None, **kwargs):
         """Apply logarithmic compression to data.
@@ -1763,20 +1878,43 @@ class LogCompress(Operation):
             dynamic_range = ops.array(DEFAULT_DYNAMIC_RANGE)
         dynamic_range = ops.cast(dynamic_range, data.dtype)
 
-        small_number = ops.convert_to_tensor(1e-16, dtype=data.dtype)
-        data = ops.where(data == 0, small_number, data)
-        compressed_data = 20 * ops.log10(data)
-        compressed_data = ops.clip(compressed_data, dynamic_range[0], dynamic_range[1])
+        compressed_data = log_compress(data)
+        if self.clip:
+            compressed_data = ops.clip(compressed_data, dynamic_range[0], dynamic_range[1])
 
         return {self.output_key: compressed_data}
 
 
+def normalize(data, output_range, input_range=None):
+    """Normalize data to a given range.
+
+    Equivalent to `translate` with clipping.
+
+    Args:
+        data (ops.Tensor): Input data to normalize.
+        output_range (tuple): Range to which data should be mapped, e.g., (0, 1).
+        input_range (tuple, optional): Range of input data.
+            If None, the range will be computed from the data.
+            Defaults to None.
+    """
+    if input_range is None:
+        input_range = (None, None)
+    minval, maxval = input_range
+    if minval is None:
+        minval = ops.min(data)
+    if maxval is None:
+        maxval = ops.max(data)
+    data = ops.clip(data, minval, maxval)
+    normalized_data = translate(data, (minval, maxval), output_range)
+    return normalized_data
+
+
 @ops_registry("normalize")
 class Normalize(Operation):
     """Normalize data to a given range."""
 
     def __init__(self, output_range=None, input_range=None, **kwargs):
-        super().__init__(**kwargs)
+        super().__init__(additional_output_keys=["minval", "maxval"], **kwargs)
         if output_range is None:
             output_range = (0, 1)
         self.output_range = self.to_float32(output_range)
@@ -1821,11 +1959,9 @@ class Normalize(Operation):
         if maxval is None:
             maxval = ops.max(data)
 
-        # Clip the data to the input range
-        data = ops.clip(data, minval, maxval)
-
-        # Map the data to the output range
-        normalized_data = translate(data, (minval, maxval), self.output_range)
+        normalized_data = normalize(
+            data, output_range=self.output_range, input_range=(minval, maxval)
+        )
 
         return {self.output_key: normalized_data, "minval": minval, "maxval": maxval}
 
@@ -1855,6 +1991,18 @@ class ScanConvert(Operation):
             input_data_type=DataTypes.IMAGE,
             output_data_type=DataTypes.IMAGE_SC,
             jittable=jittable,
+            additional_output_keys=[
+                "resolution",
+                "x_lim",
+                "y_lim",
+                "z_lim",
+                "rho_range",
+                "theta_range",
+                "phi_range",
+                "d_rho",
+                "d_theta",
+                "d_phi",
+            ],
             **kwargs,
         )
         self.order = order
@@ -2081,6 +2229,7 @@ class Demodulate(Operation):
             input_data_type=DataTypes.RAW_DATA,
             output_data_type=DataTypes.RAW_DATA,
             jittable=True,
+            additional_output_keys=["demodulation_frequency", "center_frequency", "n_ch"],
             **kwargs,
         )
         self.axis = axis
@@ -2145,7 +2294,7 @@ class Lambda(Operation):
 
 
 @ops_registry("pad")
-class Pad(Operation, TFDataLayer):
+class Pad(Operation, DataLayer):
     """Pad layer for padding tensors to a specified shape."""
 
     def __init__(
@@ -2330,6 +2479,7 @@ class Downsample(Operation):
 
     def __init__(self, factor: int = 1, phase: int = 0, axis: int = -3, **kwargs):
         super().__init__(
+            additional_output_keys=["sampling_frequency", "n_ax"],
             **kwargs,
         )
         self.factor = factor
@@ -3095,6 +3245,50 @@ def demodulate(data, center_frequency, sampling_frequency, axis=-3):
     iq_data_signal_complex = analytical_signal * ops.exp(phasor_exponent)
 
     # Split the complex signal into two channels
-    iq_data_two_channel = complex_to_channels(iq_data_signal_complex[..., 0])
+    iq_data_two_channel = complex_to_channels(ops.squeeze(iq_data_signal_complex, axis=-1))
 
     return iq_data_two_channel
+
+
+def compute_time_to_peak_stack(waveforms, center_frequencies, waveform_sampling_frequency=250e6):
+    """Compute the time of the peak of each waveform in a stack of waveforms.
+
+    Args:
+        waveforms (ndarray): The waveforms of shape (n_waveforms, n_samples).
+        center_frequencies (ndarray): The center frequencies of the waveforms in Hz of shape
+            (n_waveforms,) or a scalar if all waveforms have the same center frequency.
+        waveform_sampling_frequency (float): The sampling frequency of the waveforms in Hz.
+
+    Returns:
+        ndarray: The time to peak for each waveform in seconds.
+    """
+    t_peak = []
+    center_frequencies = center_frequencies * ops.ones((waveforms.shape[0],))
+    for waveform, center_frequency in zip(waveforms, center_frequencies):
+        t_peak.append(compute_time_to_peak(waveform, center_frequency, waveform_sampling_frequency))
+    return ops.stack(t_peak)
+
+
+def compute_time_to_peak(waveform, center_frequency, waveform_sampling_frequency=250e6):
+    """Compute the time of the peak of the waveform.
+
+    Args:
+        waveform (ndarray): The waveform of shape (n_samples).
+        center_frequency (float): The center frequency of the waveform in Hz.
+        waveform_sampling_frequency (float): The sampling frequency of the waveform in Hz.
+
+    Returns:
+        float: The time to peak for the waveform in seconds.
+    """
+    n_samples = waveform.shape[0]
+    if n_samples == 0:
+        raise ValueError("Waveform has zero samples.")
+
+    waveforms_iq_complex_channels = demodulate(
+        waveform[..., None], center_frequency, waveform_sampling_frequency, axis=-1
+    )
+    waveforms_iq_complex = channels_to_complex(waveforms_iq_complex_channels)
+    envelope = ops.abs(waveforms_iq_complex)
+    peak_idx = ops.argmax(envelope, axis=-1)
+    t_peak = ops.cast(peak_idx, dtype="float32") / waveform_sampling_frequency
+    return t_peak