zea 0.0.6__py3-none-any.whl → 0.0.8__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
@@ -67,10 +88,10 @@ Example of a yaml file:
 
 """
 
-import copy
 import hashlib
 import inspect
 import json
+import uuid
 from functools import partial
 from typing import Any, Dict, List, Union
 
@@ -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,20 @@ 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 (
+    apply_along_axis,
+    correlate,
+    extend_n_dims,
+    resample,
+    reshape_axis,
+    translate,
+    vmap,
+)
+from zea.utils import (
+    FunctionTimer,
+    deep_compare,
+    map_negative_indices,
+)
 
 
 def get_ops(ops_name):
@@ -125,6 +158,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 +177,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 +191,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 +204,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 +223,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."""
@@ -204,13 +246,18 @@ class Operation(keras.Operation):
         Analyze and store the input/output signatures of the `call` method.
         """
         self._input_signature = inspect.signature(self.call)
-        self._valid_keys = set(self._input_signature.parameters.keys())
+        self._valid_keys = set(self._input_signature.parameters.keys()) | {self.key}
 
     @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."""
@@ -366,6 +413,36 @@ class Operation(keras.Operation):
         return True
 
 
+class ImageOperation(Operation):
+    """
+    Base class for image processing operations.
+
+    This class extends the Operation class to provide a common interface
+    for operations that process image data, with shape (batch, height, width, channels)
+    or (height, width, channels) if batch dimension is not present.
+
+    Subclasses should implement the `call` method to define the image processing logic, and call
+    ``super().call(**kwargs)`` to validate the input data shape.
+    """
+
+    def call(self, **kwargs):
+        """
+        Validate input data shape for image operations.
+
+        Args:
+            **kwargs: Keyword arguments containing input data.
+
+        Raises:
+            AssertionError: If input data does not have the expected number of dimensions.
+        """
+        data = kwargs[self.key]
+
+        if self.with_batch_dim:
+            assert ops.ndim(data) == 4, "Input data must have 4 dimensions (b, h, w, c)."
+        else:
+            assert ops.ndim(data) == 3, "Input data must have 3 dimensions (h, w, c)."
+
+
 @ops_registry("pipeline")
 class Pipeline:
     """Pipeline class for processing ultrasound data through a series of operations."""
@@ -408,8 +485,6 @@ class Pipeline:
         """
         self._call_pipeline = self.call
         self.name = name
-        self.timer = FunctionTimer()
-        self.timed = timed
 
         self._pipeline_layers = operations
 
@@ -419,6 +494,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 +520,40 @@ 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
 
+        self._logged_difference_keys = False
+
+        # Do not log again for nested pipelines
+        for nested_pipeline in self._nested_pipelines:
+            nested_pipeline._logged_difference_keys = True
+
     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 _nested_pipelines(self):
+        return [operation for operation in self.operations if isinstance(operation, Pipeline)]
+
+    @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 +567,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 +597,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 +609,7 @@ class Pipeline:
 
         # Get beamforming ops
         beamforming = [
-            TOFCorrection(apply_phase_rotation=True),
+            TOFCorrection(),
             DelayAndSum(),
         ]
         if pfield:
@@ -491,11 +624,12 @@ class Pipeline:
 
         # Add display ops
         operations += [
+            ReshapeGrid(),
             EnvelopeDetect(),
             Normalize(),
             LogCompress(),
         ]
-        return cls(operations, **kwargs)
+        return cls(operations, timed=timed, **kwargs)
 
     def copy(self) -> "Pipeline":
         """Create a copy of the pipeline."""
@@ -506,57 +640,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 +716,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. "
@@ -599,6 +731,16 @@ class Pipeline:
                 "Please ensure all inputs are convertible to tensors."
             )
 
+        if not self._logged_difference_keys:
+            difference_keys = set(inputs.keys()) - self.valid_keys
+            if difference_keys:
+                log.debug(
+                    f"[zea.Pipeline] The following input keys are not used by the pipeline: "
+                    f"{difference_keys}. Make sure this is intended. "
+                    "This warning will only be shown once."
+                )
+                self._logged_difference_keys = True
+
         ## PROCESSING
         outputs = self._call_pipeline(**inputs)
 
@@ -640,18 +782,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):
@@ -726,59 +863,9 @@ class Pipeline:
             return params
 
     def __str__(self):
-        """String representation of the pipeline.
-
-        Will print on two parallel pipeline lines if it detects a splitting operations
-        (such as multi_bandpass_filter)
-        Will merge the pipeline lines if it detects a stacking operation (such as stack)
-        """
-        split_operations = []
-        merge_operations = ["Stack"]
-
+        """String representation of the pipeline."""
         operations = [operation.__class__.__name__ for operation in self.operations]
         string = " -> ".join(operations)
-
-        if any(operation in split_operations for operation in operations):
-            # a second line is needed with same length as the first line
-            split_line = " " * len(string)
-            # find the splitting operation and index and print \-> instead of -> after
-            split_detected = False
-            merge_detected = False
-            split_operation = None
-            for operation in operations:
-                if operation in split_operations:
-                    index = string.index(operation)
-                    index = index + len(operation)
-                    split_line = split_line[:index] + "\\->" + split_line[index + len("\\->") :]
-                    split_detected = True
-                    merge_detected = False
-                    split_operation = operation
-                    continue
-
-                if operation in merge_operations:
-                    index = string.index(operation)
-                    index = index - 4
-                    split_line = split_line[:index] + "/" + split_line[index + 1 :]
-                    split_detected = False
-                    merge_detected = True
-                    continue
-
-                if split_detected:
-                    # print all operations in the second line
-                    index = string.index(operation)
-                    split_line = (
-                        split_line[:index]
-                        + operation
-                        + " -> "
-                        + split_line[index + len(operation) + len(" -> ") :]
-                    )
-            assert merge_detected is True, log.error(
-                "Pipeline was never merged back together (with Stack operation), even "
-                f"though it was split with {split_operation}. "
-                "Please properly define your operation chain."
-            )
-            return f"\n{string}\n{split_line}\n"
-
         return string
 
     def __repr__(self):
@@ -835,16 +922,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 +948,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 +1062,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 +1103,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:
@@ -1093,7 +1194,7 @@ def pipeline_from_config(config: Config, **kwargs) -> Pipeline:
     operations = make_operation_chain(config.operations)
 
     # merge pipeline config without operations with kwargs
-    pipeline_config = copy.deepcopy(config)
+    pipeline_config = config.copy()
     pipeline_config.pop("operations")
 
     kwargs = {**pipeline_config, **kwargs}
@@ -1164,33 +1265,134 @@ def pipeline_to_yaml(pipeline: Pipeline, file_path: str) -> None:
         yaml.dump(pipeline_dict, f, Dumper=yaml.Dumper, indent=4)
 
 
-@ops_registry("patched_grid")
-class PatchedGrid(Pipeline):
+@ops_registry("map")
+class Map(Pipeline):
     """
-    With this class you can form a pipeline that will be applied to patches of the grid.
-    This is useful to avoid OOM errors when processing large grids.
-
-    Some things to NOTE about this class:
-
-    - The ops have to use flatgrid and flat_pfield as inputs, these will be patched.
+    A pipeline that maps its operations over specified input arguments.
 
-    - Changing anything other than `self.output_data_type` in the dict will not be propagated!
+    This can be used to reduce memory usage by processing data in chunks.
 
+    Notes
+    -----
+    - When `chunks` and `batch_size` are both None (default), this behaves like a normal Pipeline.
+    - Changing anything other than ``self.output_key`` in the dict will not be propagated.
     - Will be jitted as a single operation, not the individual operations.
-
     - This class handles the batching.
 
+    For more information on how to use ``in_axes``, ``out_axes``, `see the documentation for
+    jax.vmap <https://docs.jax.dev/en/latest/_autosummary/jax.vmap.html>`_.
+
+    Example
+    -------
+        .. doctest::
+
+            >>> from zea.ops import Map, Pipeline, Demodulate, TOFCorrection
+
+            >>> # apply operations in batches of 8
+            >>> # in this case, over the first axis of "data"
+            >>> # or more specifically, process 8 transmits at a time
+
+            >>> pipeline_mapped = Map(
+            ...     [
+            ...         Demodulate(),
+            ...         TOFCorrection(),
+            ...     ],
+            ...     argnames="data",
+            ...     batch_size=8,
+            ... )
+
+            >>> # you can also map a subset of the operations
+            >>> # for example, demodulate in 4 chunks
+            >>> # or more specifically, split the transmit axis into 4 parts
+
+            >>> pipeline_mapped = Pipeline(
+            ...     [
+            ...         Map([Demodulate()], argnames="data", chunks=4),
+            ...         TOFCorrection(),
+            ...     ],
+            ... )
     """
 
-    def __init__(self, *args, num_patches=10, **kwargs):
-        super().__init__(*args, name="patched_grid", **kwargs)
-        self.num_patches = num_patches
+    def __init__(
+        self,
+        operations: List[Operation],
+        argnames: List[str] | str,
+        in_axes: List[Union[int, None]] | int = 0,
+        out_axes: List[Union[int, None]] | int = 0,
+        chunks: int | None = None,
+        batch_size: int | None = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            operations (list): List of operations to be performed.
+            argnames (str or list): List of argument names (or keys) to map over.
+                Can also be a single string if only one argument is mapped over.
+            in_axes (int or list): Axes to map over for each argument.
+                If a single int is provided, it is used for all arguments.
+            out_axes (int or list): Axes to map over for each output.
+                If a single int is provided, it is used for all outputs.
+            chunks (int, optional): Number of chunks to split the input data into.
+                If None, no chunking is performed. Mutually exclusive with ``batch_size``.
+            batch_size (int, optional): Size of batches to process at once.
+                If None, no batching is performed. Mutually exclusive with ``chunks``.
+        """
+        super().__init__(operations, **kwargs)
 
-        for operation in self.operations:
-            if isinstance(operation, DelayAndSum):
-                operation.reshape_grid = False
+        if batch_size is not None and chunks is not None:
+            raise ValueError(
+                "batch_size and chunks are mutually exclusive. Please specify only one."
+            )
 
-        self._jittable_call = self.jittable_call
+        if batch_size is not None and batch_size <= 0:
+            raise ValueError("batch_size must be a positive integer.")
+
+        if chunks is not None and chunks <= 0:
+            raise ValueError("chunks must be a positive integer.")
+
+        if isinstance(argnames, str):
+            argnames = [argnames]
+
+        self.argnames = argnames
+        self.in_axes = in_axes
+        self.out_axes = out_axes
+        self.chunks = chunks
+        self.batch_size = batch_size
+
+        if chunks is None and batch_size is None:
+            log.warning(
+                "[zea.ops.Map] Both `chunks` and `batch_size` are None. "
+                "This will behave like a normal Pipeline. "
+                "Consider setting one of them to process data in chunks or batches."
+            )
+
+        def call_item(**inputs):
+            """Process data in patches."""
+            mapped_args = []
+            for argname in argnames:
+                mapped_args.append(inputs.pop(argname, None))
+
+            def patched_call(*args):
+                mapped_kwargs = [(k, v) for k, v in zip(argnames, args)]
+                out = super(Map, self).call(**dict(mapped_kwargs), **inputs)
+
+                # TODO: maybe it is possible to output everything?
+                # e.g. prepend a empty dimension to all inputs and just map over everything?
+                return out[self.output_key]
+
+            out = vmap(
+                patched_call,
+                in_axes=in_axes,
+                out_axes=out_axes,
+                chunks=chunks,
+                batch_size=batch_size,
+                fn_supports_batch=True,
+                disable_jit=not bool(self.jit_options),
+            )(*mapped_args)
+
+            return out
+
+        self.call_item = call_item
 
     @property
     def jit_options(self):
@@ -1230,59 +1432,28 @@ class PatchedGrid(Pipeline):
 
     @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"})
-
-    def call_item(self, inputs):
-        """Process data in patches."""
-        # Extract necessary parameters
-        # make sure to add those as valid keys above!
-        grid_size_x = inputs["grid_size_x"]
-        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"]
-
-        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)
-            return out[self.output_key]
-
-        out = patched_map(
-            patched_call,
-            flatgrid,
-            self.num_patches,
-            **patch_arrays,
-            jit=bool(self.jit_options),
-        )
-        return ops.reshape(out, (grid_size_z, grid_size_x, *ops.shape(out)[1:]))
+        """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.argnames)
+
+    @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.argnames)
 
     def jittable_call(self, **inputs):
         """Process input data through the pipeline."""
         if self._with_batch_dim:
             input_data = inputs.pop(self.key)
             output = ops.map(
-                lambda x: self.call_item({self.key: x, **inputs}),
+                lambda x: self.call_item(**{self.key: x, **inputs}),
                 input_data,
             )
         else:
-            output = self.call_item(inputs)
+            output = self.call_item(**inputs)
 
         return {self.output_key: output}
 
@@ -1295,11 +1466,61 @@ class PatchedGrid(Pipeline):
     def get_dict(self):
         """Get the configuration of the pipeline."""
         config = super().get_dict()
-        config.update({"name": "patched_grid"})
+        config["params"].update(
+            {
+                "argnames": self.argnames,
+                "in_axes": self.in_axes,
+                "out_axes": self.out_axes,
+                "chunks": self.chunks,
+                "batch_size": self.batch_size,
+            }
+        )
+        return config
+
+
+@ops_registry("patched_grid")
+class PatchedGrid(Map):
+    """
+    A pipeline that maps its operations over `flatgrid` and `flat_pfield` keys.
+
+    This can be used to reduce memory usage by processing data in chunks.
+
+    For more information and flexibility, see :class:`zea.ops.Map`.
+    """
+
+    def __init__(self, *args, num_patches=10, **kwargs):
+        super().__init__(*args, argnames=["flatgrid", "flat_pfield"], chunks=num_patches, **kwargs)
+        self.num_patches = num_patches
+
+    def get_dict(self):
+        """Get the configuration of the pipeline."""
+        config = super().get_dict()
+        config["params"].pop("argnames")
+        config["params"].pop("chunks")
         config["params"].update({"num_patches": self.num_patches})
         return config
 
 
+@ops_registry("reshape_grid")
+class ReshapeGrid(Operation):
+    """Reshape flat grid data to grid shape."""
+
+    def __init__(self, axis=0, **kwargs):
+        super().__init__(**kwargs)
+        self.axis = axis
+
+    def call(self, grid, **kwargs):
+        """
+        Args:
+            - data (Tensor): The flat grid data of shape (..., n_pix, ...).
+        Returns:
+            - reshaped_data (Tensor): The reshaped data of shape (..., grid.shape, ...).
+        """
+        data = kwargs[self.key]
+        reshaped_data = reshape_axis(data, grid.shape[:-1], self.axis + int(self.with_batch_dim))
+        return {self.output_key: reshaped_data}
+
+
 ## Base Operations
 
 
@@ -1309,7 +1530,7 @@ class Identity(Operation):
 
     def call(self, **kwargs) -> Dict:
         """Returns the input as is."""
-        return kwargs
+        return {}
 
 
 @ops_registry("merge")
@@ -1332,21 +1553,6 @@ class Merge(Operation):
         return merged
 
 
-@ops_registry("split")
-class Split(Operation):
-    """Operation that splits an input dictionary  n copies."""
-
-    def __init__(self, n: int, **kwargs):
-        super().__init__(**kwargs)
-        self.n = n
-
-    def call(self, **kwargs) -> List[Dict]:
-        """
-        Splits the input dictionary into n copies.
-        """
-        return [kwargs.copy() for _ in range(self.n)]
-
-
 @ops_registry("stack")
 class Stack(Operation):
     """Stack multiple data arrays along a new axis.
@@ -1399,6 +1605,7 @@ class Simulate(Operation):
     def __init__(self, **kwargs):
         super().__init__(
             output_data_type=DataTypes.RAW_DATA,
+            additional_output_keys=["n_ch"],
             **kwargs,
         )
 
@@ -1451,18 +1658,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 +1682,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 +1704,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 +1717,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,
             )
 
@@ -1556,7 +1767,7 @@ class PfieldWeighting(Operation):
         Returns:
             dict: Dictionary containing weighted data
         """
-        data = kwargs[self.key]
+        data = kwargs[self.key]  # must start with ((batch_size,) n_tx, n_pix, ...)
 
         if flat_pfield is None:
             return {self.output_key: data}
@@ -1564,14 +1775,16 @@ class PfieldWeighting(Operation):
         # Swap (n_pix, n_tx) to (n_tx, n_pix)
         flat_pfield = ops.swapaxes(flat_pfield, 0, 1)
 
-        # Perform element-wise multiplication with the pressure weight mask
-        # Also add the required dimensions for broadcasting
+        # Add batch dimension if needed
         if self.with_batch_dim:
             pfield_expanded = ops.expand_dims(flat_pfield, axis=0)
         else:
             pfield_expanded = flat_pfield
 
-        pfield_expanded = pfield_expanded[..., None, None]
+        append_n_dims = ops.ndim(data) - ops.ndim(pfield_expanded)
+        pfield_expanded = extend_n_dims(pfield_expanded, axis=-1, n_dims=append_n_dims)
+
+        # Perform element-wise multiplication with the pressure weight mask
         weighted_data = data * pfield_expanded
 
         return {self.output_key: weighted_data}
@@ -1581,79 +1794,93 @@ class PfieldWeighting(Operation):
 class DelayAndSum(Operation):
     """Sums time-delayed signals along channels and transmits."""
 
-    def __init__(
-        self,
-        reshape_grid=True,
-        **kwargs,
-    ):
+    def __init__(self, **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
-                of shape `(grid_size_z*grid_size_x, n_ch)` when reshape_grid is False
-                or `(grid_size_z, grid_size_x, n_ch)` when reshape_grid is True,
+                of shape `(grid_size_z*grid_size_x, n_ch)`
                 with optional batch dimension.
         """
         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
-            )
-
-        if self.reshape_grid:
-            beamformed_data = reshape_axis(
-                beamformed_data, grid.shape[:2], axis=int(self.with_batch_dim)
-            )
+            beamformed_data = ops.map(self.process_image, data)
 
         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 +1907,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 +1945,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 +1984,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 +2065,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 +2097,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
@@ -1913,7 +2167,7 @@ class ScanConvert(Operation):
 
 
 @ops_registry("gaussian_blur")
-class GaussianBlur(Operation):
+class GaussianBlur(ImageOperation):
     """
     GaussianBlur is an operation that applies a Gaussian blur to an input image.
     Uses scipy.ndimage.gaussian_filter to create a kernel.
@@ -1963,6 +2217,13 @@ class GaussianBlur(Operation):
         return ops.convert_to_tensor(kernel)
 
     def call(self, **kwargs):
+        """Apply a Gaussian filter to the input data.
+
+        Args:
+            data (ops.Tensor): Input image data of shape (height, width, channels) with
+                optional batch dimension if ``self.with_batch_dim``.
+        """
+        super().call(**kwargs)
         data = kwargs[self.key]
 
         # Add batch dimension if not present
@@ -1997,7 +2258,7 @@ class GaussianBlur(Operation):
 
 
 @ops_registry("lee_filter")
-class LeeFilter(Operation):
+class LeeFilter(ImageOperation):
     """
     The Lee filter is a speckle reduction filter commonly used in synthetic aperture radar (SAR)
     and ultrasound image processing. It smooths the image while preserving edges and details.
@@ -2027,7 +2288,7 @@ class LeeFilter(Operation):
             pad_mode=self.pad_mode,
             with_batch_dim=self.with_batch_dim,
             jittable=self._jittable,
-            key=self.key,
+            key="data",
         )
 
     @property
@@ -2043,24 +2304,29 @@ class LeeFilter(Operation):
             self.gaussian_blur.with_batch_dim = value
 
     def call(self, **kwargs):
-        data = kwargs[self.key]
+        """Apply the Lee filter to the input data.
+
+        Args:
+            data (ops.Tensor): Input image data of shape (height, width, channels) with
+                optional batch dimension if ``self.with_batch_dim``.
+        """
+        super().call(**kwargs)
+        data = kwargs.pop(self.key)
 
         # Apply Gaussian blur to get local mean
-        img_mean = self.gaussian_blur.call(**kwargs)[self.gaussian_blur.output_key]
+        img_mean = self.gaussian_blur.call(data=data, **kwargs)[self.gaussian_blur.output_key]
 
         # Apply Gaussian blur to squared data to get local squared mean
-        data_squared = data**2
-        kwargs[self.gaussian_blur.key] = data_squared
-        img_sqr_mean = self.gaussian_blur.call(**kwargs)[self.gaussian_blur.output_key]
+        img_sqr_mean = self.gaussian_blur.call(
+            data=data**2,
+            **kwargs,
+        )[self.gaussian_blur.output_key]
 
         # Calculate local variance
         img_variance = img_sqr_mean - img_mean**2
 
         # Calculate global variance (per channel)
-        if self.with_batch_dim:
-            overall_variance = ops.var(data, axis=(-3, -2), keepdims=True)
-        else:
-            overall_variance = ops.var(data, axis=(-2, -1), keepdims=True)
+        overall_variance = ops.var(data, axis=(-3, -2), keepdims=True)
 
         # Calculate adaptive weights
         img_weights = img_variance / (img_variance + overall_variance)
@@ -2081,6 +2347,11 @@ 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
@@ -2106,6 +2377,121 @@ class Demodulate(Operation):
         }
 
 
+@ops_registry("fir_filter")
+class FirFilter(Operation):
+    """Apply a FIR filter to the input signal using convolution.
+
+    Looks for the filter taps in the input dictionary using the specified ``filter_key``.
+    """
+
+    def __init__(
+        self,
+        axis: int,
+        complex_channels: bool = False,
+        filter_key: str = "fir_filter_taps",
+        **kwargs,
+    ):
+        """
+        Args:
+            axis (int): Axis along which to apply the filter. Cannot be the batch dimension.
+                When using ``complex_channels=True``, the complex channels are removed to convert
+                to complex numbers before filtering, so adjust the ``axis`` accordingly!
+            complex_channels (bool): Whether the last dimension of the input signal represents
+                complex channels (real and imaginary parts). When True, it will convert the signal
+                to ``complex`` dtype before filtering and convert it back to two channels
+                after filtering.
+            filter_key (str): Key in the input dictionary where the FIR filter taps are stored.
+                Default is "fir_filter_taps".
+        """
+        super().__init__(**kwargs)
+        self._check_axis(axis)
+
+        self.axis = axis
+        self.complex_channels = complex_channels
+        self.filter_key = filter_key
+
+    def _check_axis(self, axis, ndim=None):
+        """Check if the axis is valid."""
+        if ndim is not None:
+            if axis < -ndim or axis >= ndim:
+                raise ValueError(f"Axis {axis} is out of bounds for array of dimension {ndim}.")
+
+        if self.with_batch_dim and (axis == 0 or (ndim is not None and axis == -ndim)):
+            raise ValueError("Cannot apply FIR filter along batch dimension.")
+
+    @property
+    def valid_keys(self):
+        """Get the valid keys for the `call` method."""
+        return self._valid_keys.union({self.filter_key})
+
+    def call(self, **kwargs):
+        signal = kwargs[self.key]
+        fir_filter_taps = kwargs[self.filter_key]
+
+        if self.complex_channels:
+            signal = channels_to_complex(signal)
+
+        self._check_axis(self.axis, ndim=ops.ndim(signal))
+
+        def _convolve(signal):
+            """Apply the filter to the signal using correlation."""
+            return correlate(signal, fir_filter_taps[::-1], mode="same")
+
+        filtered_signal = apply_along_axis(_convolve, self.axis, signal)
+
+        if self.complex_channels:
+            filtered_signal = complex_to_channels(filtered_signal)
+
+        return {self.output_key: filtered_signal}
+
+
+@ops_registry("low_pass_filter")
+class LowPassFilter(FirFilter):
+    """Apply a low-pass FIR filter to the input signal using convolution.
+
+    It is recommended to use :class:`FirFilter` with pre-computed filter taps for jittable
+    operations. The :class:`LowPassFilter` operation itself is not jittable and is provided
+    for convenience only.
+
+    Uses :func:`get_low_pass_iq_filter` to compute the filter taps.
+    """
+
+    def __init__(self, axis: int, complex_channels: bool = False, num_taps: int = 128, **kwargs):
+        """Initialize the LowPassFilter operation.
+
+        Args:
+            axis (int): Axis along which to apply the filter. Cannot be the batch dimension.
+                When using ``complex_channels=True``, the complex channels are removed to convert
+                to complex numbers before filtering, so adjust the ``axis`` accordingly.
+            complex_channels (bool): Whether the last dimension of the input signal represents
+                complex channels (real and imaginary parts). When True, it will convert the signal
+                to ``complex`` dtype before filtering and convert it back to two channels
+                after filtering.
+            num_taps (int): Number of taps in the FIR filter. Default is 128.
+        """
+        self._random_suffix = str(uuid.uuid4())
+        kwargs.pop("filter_key", None)
+        kwargs.pop("jittable", None)
+        super().__init__(
+            axis=axis,
+            complex_channels=complex_channels,
+            filter_key=f"low_pass_{self._random_suffix}",
+            jittable=False,
+            **kwargs,
+        )
+        self.num_taps = num_taps
+
+    def call(self, bandwidth, sampling_frequency, center_frequency, **kwargs):
+        lpf = get_low_pass_iq_filter(
+            self.num_taps,
+            ops.convert_to_numpy(sampling_frequency).item(),
+            ops.convert_to_numpy(center_frequency).item(),
+            ops.convert_to_numpy(bandwidth).item(),
+        )
+        kwargs[self.filter_key] = lpf
+        return super().call(**kwargs)
+
+
 @ops_registry("lambda")
 class Lambda(Operation):
     """Use any function as an operation."""
@@ -2140,12 +2526,15 @@ class Lambda(Operation):
 
     def call(self, **kwargs):
         data = kwargs[self.key]
-        data = self.func(data)
+        if self.with_batch_dim:
+            data = ops.map(self.func, data)
+        else:
+            data = self.func(data)
         return {self.output_key: data}
 
 
 @ops_registry("pad")
-class Pad(Operation, TFDataLayer):
+class Pad(Operation, DataLayer):
     """Pad layer for padding tensors to a specified shape."""
 
     def __init__(
@@ -2330,6 +2719,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
@@ -2894,7 +3284,7 @@ def get_band_pass_filter(num_taps, sampling_frequency, f1, f2):
     return bpf
 
 
-def get_low_pass_iq_filter(num_taps, sampling_frequency, f, bw):
+def get_low_pass_iq_filter(num_taps, sampling_frequency, center_frequency, bandwidth):
     """Design complex low-pass filter.
 
     The filter is a low-pass FIR filter modulated to the center frequency.
@@ -2902,16 +3292,16 @@ def get_low_pass_iq_filter(num_taps, sampling_frequency, f, bw):
     Args:
         num_taps (int): number of taps in filter.
         sampling_frequency (float): sample frequency.
-        f (float): center frequency.
-        bw (float): bandwidth in Hz.
+        center_frequency (float): center frequency.
+        bandwidth (float): bandwidth in Hz.
 
     Raises:
-        ValueError: if cutoff frequency (bw / 2) is not within (0, sampling_frequency / 2)
+        ValueError: if cutoff frequency (bandwidth / 2) is not within (0, sampling_frequency / 2)
 
     Returns:
         ndarray: Complex-valued low-pass filter
     """
-    cutoff = bw / 2
+    cutoff = bandwidth / 2
     if not (0 < cutoff < sampling_frequency / 2):
         raise ValueError(
             f"Cutoff frequency must be within (0, sampling_frequency / 2), "
@@ -2921,7 +3311,7 @@ def get_low_pass_iq_filter(num_taps, sampling_frequency, f, bw):
     lpf = scipy.signal.firwin(num_taps, cutoff, pass_zero=True, fs=sampling_frequency)
     # Modulate to center frequency to make it complex
     time_points = np.arange(num_taps) / sampling_frequency
-    lpf_complex = lpf * np.exp(1j * 2 * np.pi * f * time_points)
+    lpf_complex = lpf * np.exp(1j * 2 * np.pi * center_frequency * time_points)
     return lpf_complex
 
 
@@ -3095,6 +3485,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