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

zea/__init__.py

@@ -7,10 +7,10 @@ from . import log
 
 # dynamically add __version__ attribute (see pyproject.toml)
 # __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.4"
+__version__ = "0.0.6"
 
 
-def setup():
+def _bootstrap_backend():
     """Setup function to initialize the zea package."""
 
     def _check_backend_installed():
@@ -40,14 +40,14 @@ def setup():
 
     _check_backend_installed()
 
-    import keras
+    from keras.backend import backend as keras_backend
 
-    log.info(f"Using backend {keras.backend.backend()!r}")
+    log.info(f"Using backend {keras_backend()!r}")
 
 
 # call and clean up namespace
-setup()
-del setup
+_bootstrap_backend()
+del _bootstrap_backend
 
 from . import (
     agent,
@@ -55,6 +55,7 @@ from . import (
     data,
     display,
     io_lib,
+    keras_ops,
     metrics,
     models,
     simulator,
@@ -68,7 +69,7 @@ from .data.file import File, load_file
 from .datapaths import set_data_paths
 from .interface import Interface
 from .internal.device import init_device
-from .internal.setup_zea import set_backend, setup, setup_config
+from .internal.setup_zea import setup, setup_config
 from .ops import Pipeline
 from .probes import Probe
 from .scan import Scan

zea/__main__.py

@@ -9,30 +9,22 @@ import argparse
 import sys
 from pathlib import Path
 
-from zea import log
 from zea.visualize import set_mpl_style
 
 
-def get_args():
+def get_parser():
     """Command line argument parser"""
-    parser = argparse.ArgumentParser(description="Process ultrasound data.")
-    parser.add_argument("-c", "--config", type=str, default=None, help="path to config file.")
+    parser = argparse.ArgumentParser(
+        description="Load and process ultrasound data based on a configuration file."
+    )
+    parser.add_argument("-c", "--config", type=str, default=None, help="path to the config file.")
     parser.add_argument(
         "-t",
         "--task",
         default="view",
         choices=["view"],
         type=str,
-        help="which task to run",
-    )
-    parser.add_argument(
-        "--backend",
-        default=None,
-        type=str,
-        help=(
-            "Keras backend to use. Default is the one set by the environment "
-            "variable KERAS_BACKEND."
-        ),
+        help="Which task to run. Currently only 'view' is supported.",
     )
     parser.add_argument(
         "--skip_validate_file",
@@ -40,27 +32,18 @@ def get_args():
         action="store_true",
         help="Skip zea file integrity checks. Use with caution.",
     )
-    parser.add_argument("--gui", default=False, action=argparse.BooleanOptionalAction)
-    args = parser.parse_args()
-    return args
+    return parser
 
 
 def main():
     """main entrypoint for zea"""
-    args = get_args()
+    args = get_parser().parse_args()
 
     set_mpl_style()
 
-    if args.backend:
-        from zea.internal.setup_zea import set_backend
-
-        set_backend(args.backend)
-
     wd = Path(__file__).parent.resolve()
     sys.path.append(str(wd))
 
-    import keras
-
     from zea.interface import Interface
     from zea.internal.setup_zea import setup
 
@@ -72,7 +55,6 @@ def main():
             validate_file=not args.skip_validate_file,
         )
 
-        log.info(f"Using {keras.backend.backend()} backend")
         cli.run(plot=True)
     else:
         raise ValueError(f"Unknown task {args.task}, see `zea --help` for available tasks.")

zea/agent/selection.py

@@ -11,11 +11,14 @@ For a comprehensive example usage, see: :doc:`../notebooks/agent/agent_example`
 All strategies are stateless, meaning that they do not maintain any internal state.
 """
 
+from typing import Callable
+
 import keras
 from keras import ops
 
 from zea import tensor_ops
 from zea.agent import masks
+from zea.backend.autograd import AutoGrad
 from zea.internal.registry import action_selection_registry
 
 
@@ -493,3 +496,166 @@ class CovarianceSamplingLines(LinesActionModel):
         best_mask = ops.squeeze(best_mask, axis=0)
 
         return best_mask, self.lines_to_im_size(best_mask)
+
+
+class TaskBasedLines(GreedyEntropy):
+    """Task-based line selection for maximizing information gain.
+
+    This action selection strategy chooses lines to maximize information gain with respect
+    to a downstream task outcome. It uses gradient-based saliency to identify which image
+    regions contribute most to task uncertainty, then selects lines accordingly.
+    """
+
+    def __init__(
+        self,
+        n_actions: int,
+        n_possible_actions: int,
+        img_width: int,
+        img_height: int,
+        downstream_task_function: Callable,
+        mean: float = 0,
+        std_dev: float = 1,
+        num_lines_to_update: int = 5,
+        **kwargs,
+    ):
+        """Initialize the TaskBasedLines action selection model.
+
+        Args:
+            n_actions (int): The number of actions the agent can take.
+            n_possible_actions (int): The number of possible actions (line positions).
+            img_width (int): The width of the input image.
+            img_height (int): The height of the input image.
+            downstream_task_function (Callable): A differentiable function that takes a
+                batch of inputs and produces scalar outputs. This represents the downstream
+                task for which information gain should be maximized.
+            mean (float, optional): The mean of the RBF used for reweighting. Defaults to 0.
+            std_dev (float, optional): The standard deviation of the RBF used for reweighting.
+                Defaults to 1.
+            num_lines_to_update (int, optional): The number of lines around the selected line
+                to update during reweighting. Must be odd. Defaults to 5.
+            **kwargs: Additional keyword arguments passed to the parent class.
+        """
+        super().__init__(
+            n_actions,
+            n_possible_actions,
+            img_width,
+            img_height,
+            mean,
+            std_dev,
+            num_lines_to_update,
+        )
+        self.downstream_task_function = downstream_task_function
+
+    def compute_output_and_saliency_propagation(self, particles):
+        """Compute saliency-weighted posterior variance for task-based selection.
+
+        This method computes how much each pixel contributes to the variance of the
+        downstream task output. It uses automatic differentiation to compute gradients
+        of the task function with respect to each particle, then weights the posterior
+        variance by the squared mean gradient.
+
+        Args:
+            particles (Tensor): Particles of shape (batch_size, n_particles, height, width)
+                representing the posterior distribution over images.
+
+        Returns:
+            Tensor: Pixelwise contribution to downstream task variance,
+                of shape (batch_size, height, width). Higher values indicate pixels
+                that contribute more to task uncertainty.
+        """
+        autograd = AutoGrad()
+
+        autograd.set_function(self.downstream_task_function)
+        downstream_grad_and_value_fn = autograd.get_gradient_and_value_jit_fn()
+        jacobian, _ = ops.vectorized_map(
+            lambda p: ops.vectorized_map(
+                downstream_grad_and_value_fn,
+                p,
+            ),
+            particles,
+        )
+
+        posterior_variance = ops.var(particles, axis=1)
+        mean_jacobian = ops.mean(jacobian, axis=1)
+        return posterior_variance * (mean_jacobian**2)
+
+    def sum_neighbouring_columns_into_n_possible_actions(self, full_linewise_salience):
+        """Aggregate column-wise saliency into line-wise saliency scores.
+
+        This method groups neighboring columns together to create saliency scores
+        for each possible line action. Since each line action may correspond to
+        multiple image columns, this aggregation is necessary to match the action space.
+
+        Args:
+            full_linewise_salience (Tensor): Saliency values for each column,
+                of shape (batch_size, full_image_width).
+
+        Returns:
+            Tensor: Aggregated saliency scores for each possible action,
+                of shape (batch_size, n_possible_actions).
+
+        Raises:
+            AssertionError: If the image width is not evenly divisible by n_possible_actions.
+        """
+        batch_size = ops.shape(full_linewise_salience)[0]
+        full_image_width = ops.shape(full_linewise_salience)[1]
+        assert full_image_width % self.n_possible_actions == 0, (
+            "n_possible_actions must divide evenly into image width"
+        )
+        cols_per_action = full_image_width // self.n_possible_actions
+        stacked_linewise_salience = ops.reshape(
+            full_linewise_salience,
+            (batch_size, self.n_possible_actions, cols_per_action),
+        )
+        return ops.sum(stacked_linewise_salience, axis=2)
+
+    def sample(self, particles):
+        """Sample actions using task-based information gain maximization.
+
+        This method computes which lines would provide the most information about
+        the downstream task by:
+        1. Computing pixelwise contribution to task variance using gradients
+        2. Aggregating contributions into line-wise scores
+        3. Greedily selecting lines with highest contribution scores
+        4. Reweighting scores around selected lines (inherited from GreedyEntropy)
+
+        Args:
+            particles (Tensor): Particles representing the posterior distribution,
+                of shape (batch_size, n_particles, height, width).
+
+        Returns:
+            Tuple[Tensor, Tensor, Tensor]:
+                - selected_lines_k_hot: Selected lines as k-hot vectors,
+                  shaped (batch_size, n_possible_actions)
+                - masks: Binary masks of shape (batch_size, img_height, img_width)
+                - pixelwise_contribution_to_var_dst: Pixelwise contribution to downstream
+                  task variance, of shape (batch_size, height, width)
+
+        Note:
+            Unlike the parent GreedyEntropy class, this method returns an additional
+            tensor containing the pixelwise contribution scores for analysis.
+        """
+        pixelwise_contribution_to_var_dst = self.compute_output_and_saliency_propagation(particles)
+        linewise_contribution_to_var_dst = ops.sum(pixelwise_contribution_to_var_dst, axis=1)
+        actionwise_contribution_to_var_dst = self.sum_neighbouring_columns_into_n_possible_actions(
+            linewise_contribution_to_var_dst
+        )
+
+        # Greedily select best line, reweight entropies, and repeat
+        all_selected_lines = []
+        for _ in range(self.n_actions):
+            max_contribution_line, actionwise_contribution_to_var_dst = ops.vectorized_map(
+                self.select_line_and_reweight_entropy,
+                actionwise_contribution_to_var_dst,
+            )
+            all_selected_lines.append(max_contribution_line)
+
+        selected_lines_k_hot = ops.any(
+            ops.one_hot(all_selected_lines, self.n_possible_actions, dtype=masks._DEFAULT_DTYPE),
+            axis=0,
+        )
+        return (
+            selected_lines_k_hot,
+            self.lines_to_im_size(selected_lines_k_hot),
+            pixelwise_contribution_to_var_dst,
+        )

zea/backend/__init__.py

@@ -25,6 +25,8 @@ Key Features
 
 """
 
+from contextlib import nullcontext
+
 import keras
 
 from zea import log
@@ -114,3 +116,90 @@ def _jit_compile(func, jax=True, tensorflow=True, **kwargs):
         log.warning("Initialize zea.Pipeline with jit_options=None to suppress this warning.")
         log.warning("Falling back to non-compiled mode.")
         return func
+
+
+class on_device:
+    """Context manager to set the device regardless of backend.
+
+    For the `torch` backend, you need to manually move the model and data to the device before
+    using this context manager.
+
+    Args:
+        device (str): Device string, e.g. ``'cuda'``, ``'gpu'``, or ``'cpu'``.
+
+    Example:
+        .. code-block:: python
+
+            with zea.backend.on_device("gpu:3"):
+                pipeline = zea.Pipeline([zea.keras_ops.Abs()])
+                output = pipeline(data=keras.random.normal((10, 10)))  # output is on "cuda:3"
+    """
+
+    def __init__(self, device: str):
+        self.device = self.get_device(device)
+        self._context = self.get_context(self.device)
+
+    def get_context(self, device):
+        if device is None:
+            return nullcontext()
+
+        if keras.backend.backend() == "tensorflow":
+            import tensorflow as tf
+
+            return tf.device(device)
+
+        if keras.backend.backend() == "jax":
+            import jax
+
+            return jax.default_device(device)
+        if keras.backend.backend() == "torch":
+            import torch
+
+            return torch.device(device)
+
+        return nullcontext()
+
+    def get_device(self, device: str):
+        if device is None:
+            return None
+
+        device = device.lower()
+
+        if keras.backend.backend() == "tensorflow":
+            return device.replace("cuda", "gpu")
+
+        if keras.backend.backend() == "jax":
+            from zea.backend.jax import str_to_jax_device
+
+            device = device.replace("cuda", "gpu")
+            return str_to_jax_device(device)
+
+        if keras.backend.backend() == "torch":
+            return device.replace("gpu", "cuda")
+
+    def __enter__(self):
+        self._context.__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._context.__exit__(exc_type, exc_val, exc_tb)
+
+
+if keras.backend.backend() in ["tensorflow", "jax", "numpy"]:
+
+    def func_on_device(func, device, *args, **kwargs):
+        """Moves all tensor arguments of a function to a specified device before calling it.
+
+        Args:
+            func (callable): Function to be called.
+            device (str): Device to move tensors to.
+            *args: Positional arguments to be passed to the function.
+            **kwargs: Keyword arguments to be passed to the function.
+        Returns:
+            The output of the function.
+        """
+        with on_device(device):
+            return func(*args, **kwargs)
+elif keras.backend.backend() == "torch":
+    from zea.backend.torch import func_on_device
+else:
+    raise ValueError(f"Unsupported backend: {keras.backend.backend()}")

zea/backend/jax/__init__.py

@@ -1,47 +1,21 @@
 """Jax utilities for zea."""
 
 import jax
-import numpy as np
 
 
-def on_device_jax(func, inputs, device, return_numpy=False, **kwargs):
-    """Applies a JAX function to inputs on a specified device.
-
+def str_to_jax_device(device):
+    """Convert a device string to a JAX device.
     Args:
-        func (callable): The function to apply.
-        inputs (ndarray): Input array.
-        device (str): Device string, e.g. ``'cuda'``, ``'gpu'``, or ``'cpu'``.
-        return_numpy (bool, optional): Whether to convert output
-            data back to numpy. Defaults to False.
-        **kwargs: Additional keyword arguments to be passed to the ``func``.
-
+        device (str): Device string, e.g. ``'gpu:0'``, or ``'cpu:0'``.
     Returns:
-        jax.numpy.DeviceArray or ndarray: The output data.
-
-    Raises:
-        AssertionError: If ``func`` is not a function from the JAX library.
-
-    Note:
-        This function converts the ``inputs`` array to a JAX array and moves
-        it to the specified ``device``. It then applies the ``func`` function to the inputs
-        and returns the output data. If the output is a dictionary, it extracts the first value
-        from the dictionary. If ``return_numpy`` is True, it converts the output data back to a
-        numpy array before returning.
-
-    Example:
-        .. code-block:: python
-
-            import jax.numpy as jnp
-
+        jax.Device: The corresponding JAX device.
+    """
 
-            def square(x):
-                return x**2
+    if not isinstance(device, str):
+        raise ValueError(f"Device must be a string, got {type(device)}")
 
+    device = device.lower().replace("cuda", "gpu")
 
-            inputs = [1, 2, 3, 4, 5]
-            device = "gpu"
-            output = on_device_jax(square, inputs, device)
-    """
     device = device.split(":")
     if len(device) == 2:
         device_type, device_number = device
@@ -51,20 +25,9 @@ def on_device_jax(func, inputs, device, return_numpy=False, **kwargs):
         device_type = device[0]
         device_number = 0
 
-    if device_number > len(jax.devices(device_type)):
-        raise ValueError(
-            f"Device {device} is not available from JAX devices: {jax.devices(device_type)}"
-        )
-
-    jax_device = jax.devices(device_type)[device_number]
-
-    with jax.default_device(jax_device):
-        outputs = func(inputs, **kwargs)
-
-    if isinstance(outputs, dict):
-        outputs = list(outputs.values())[0]
-
-    if return_numpy:
-        outputs = np.array(outputs)
-
-    return outputs
+    available = jax.devices(device_type)
+    if len(available) == 0:
+        raise ValueError(f"No JAX devices available for type '{device_type}'.")
+    if device_number < 0 or device_number >= len(available):
+        raise ValueError(f"Device '{device}' is not available; JAX devices found: {available}")
+    return available[device_number]

zea/backend/tensorflow/__init__.py

@@ -15,52 +15,3 @@ sys.path = [str(p) if isinstance(p, PosixPath) else p for p in sys.path]
 import tensorflow as tf  # noqa: E402
 
 from .dataloader import make_dataloader  # noqa: E402
-
-
-def on_device_tf(func, inputs, device, return_numpy=False, **kwargs):
-    """Applies a Tensorflow function to inputs on a specified device.
-
-    Args:
-        func (function): Function to apply to the input data.
-        inputs (ndarray): Input array.
-        device (str): Device string, e.g. ``'cuda'``, ``'gpu'``, or ``'cpu'``.
-        return_numpy (bool, optional): Whether to convert output
-            data back to numpy. Defaults to False.
-        **kwargs: Additional keyword arguments to be passed to the ``func``.
-
-    Returns:
-        tf.Tensor or ndarray: The output data.
-
-    Raises:
-        AssertionError: If ``func`` is not a function from the tensorflow library.
-
-    Note:
-        This function converts the ``inputs`` array to a tf.Tensor and moves
-        it to the specified ``device``. It then applies the ``func`` function to the inputs
-        and returns the output data. If the output is a dictionary, it extracts the first value
-        from the dictionary. If ``return_numpy`` is True, it converts the output data back to a
-        numpy array before returning.
-
-    Example:
-        .. code-block:: python
-
-            import tensorflow as tf
-
-
-            def square(x):
-                return x**2
-
-
-            inputs = [1, 2, 3, 4, 5]
-            device = "cuda"
-            output = on_device_tf(square, inputs, device)
-    """
-    device = device.replace("cuda", "gpu")
-
-    with tf.device(device):
-        outputs = func(inputs, **kwargs)
-
-    if return_numpy:
-        if not isinstance(outputs, np.ndarray):
-            outputs = outputs.numpy()
-    return outputs

zea/backend/torch/__init__.py

@@ -3,72 +3,37 @@
 Initialize modules for registries.
 """
 
-import numpy as np
 import torch
 
 
-def on_device_torch(func, inputs, device, return_numpy=False, **kwargs):
-    """Applies a Torch function to inputs on a specified device.
+def func_on_device(func, device, *args, **kwargs):
+    """Moves all tensor arguments of a function to a specified device before calling it.
 
     Args:
-        func (function): Function to apply to the input data.
-        inputs (ndarray): Input array.
-        device (str): Device string, e.g. ``'cuda'``, ``'gpu'``, or ``'cpu'``.
-        return_numpy (bool, optional): Whether to convert output
-            data back to numpy. Defaults to False.
-        **kwargs: Additional keyword arguments to be passed to the ``func``.
-
+        func (callable): Function to be called.
+        device (str or torch.device): Device to move tensors to.
+        *args: Positional arguments to be passed to the function.
+        **kwargs: Keyword arguments to be passed to the function.
     Returns:
-        torch.Tensor or ndarray: The output data.
-
-    Raises:
-        AssertionError: If ``func`` is not a function from the torch library.
-
-    Note:
-        This function converts the ``inputs`` array to a torch.Tensor and moves it to
-        the specified ``device``. It then applies the ``func`` function to the inputs and
-        returns the output data. If the output is a dictionary, it extracts the first value
-        from the dictionary. If ``return_numpy`` is True, it converts the output data back to a
-        numpy array before returning.
-
-    Example:
-        .. code-block:: python
-
-            import torch
-
-
-            def square(x):
-                return x**2
-
-
-            inputs = [1, 2, 3, 4, 5]
-            device = "cuda"
-            output = on_device_torch(square, inputs, device)
-            print(output)
+        The output of the function.
     """
-    device = device.replace("gpu", "cuda")
-
-    if not isinstance(inputs, torch.Tensor):
-        inputs = torch.tensor(inputs)
-
-    inputs = inputs.to(device)
-
-    # check that function is a function from torch library
-    # assert "torch" in str(type(func)), f"func: {func} should be a torch function"
-    if hasattr(func, "to"):
-        func = func.to(device)
-
-    with torch.device(device):
-        outputs = func(inputs, **kwargs)
-
-    if isinstance(outputs, dict):
-        # depends a bit how flexible we want to be...
-        # but for now quick and dirty solution
-        key = list(outputs.keys())[0]
-        outputs = outputs[key]
-
-    if return_numpy:
-        if not isinstance(outputs, np.ndarray):
-            outputs = outputs.cpu().numpy()
-
-    return outputs
+    if device is None:
+        return func(*args, **kwargs)
+
+    if isinstance(device, str):
+        device = torch.device(device)
+
+    def move_to_device(x):
+        if isinstance(x, torch.Tensor):
+            return x.to(device)
+        elif isinstance(x, (list, tuple)):
+            return type(x)(move_to_device(i) for i in x)
+        elif isinstance(x, dict):
+            return {k: move_to_device(v) for k, v in x.items()}
+        else:
+            return x
+
+    args = move_to_device(args)
+    kwargs = move_to_device(kwargs)
+
+    return func(*args, **kwargs)

zea/data/__main__.py

@@ -9,8 +9,8 @@ import argparse
 from zea import Folder
 
 
-def main():
-    parser = argparse.ArgumentParser(description="Copy a zea.Folder to a new location.")
+def get_parser():
+    parser = argparse.ArgumentParser(description="Copy a :class:`zea.Folder` to a new location.")
     parser.add_argument("src", help="Source folder path")
     parser.add_argument("dst", help="Destination folder path")
     parser.add_argument("key", help="Key to access in the hdf5 files")
@@ -20,8 +20,11 @@ def main():
         choices=["a", "w", "r+", "x"],
         help="Mode in which to open the destination files (default: 'a')",
     )
+    return parser
+
 
-    args = parser.parse_args()
+def main():
+    args = get_parser().parse_args()
 
     src_folder = Folder(args.src, args.key, validate=False)
     src_folder.copy(args.dst, args.key, mode=args.mode)