zea 0.0.8__tar.gz → 0.0.10__tar.gz

{zea-0.0.8 → zea-0.0.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: zea
-Version: 0.0.8
+Version: 0.0.10
 Summary: A Toolbox for Cognitive Ultrasound Imaging. Provides a set of tools for processing of ultrasound data, all built in your favorite machine learning framework.
 License-File: LICENSE
 Keywords: ultrasound,machine learning,beamforming
@@ -44,8 +44,6 @@ Requires-Dist: jax ; extra == "backends"
 Requires-Dist: jax[cuda12-pip] (>=0.4.26) ; extra == "jax"
 Requires-Dist: keras (>=3.12)
 Requires-Dist: matplotlib (>=3.8)
-Requires-Dist: mock ; extra == "dev"
-Requires-Dist: mock ; extra == "docs"
 Requires-Dist: myst-parser ; extra == "dev"
 Requires-Dist: myst-parser ; extra == "docs"
 Requires-Dist: nbsphinx ; extra == "dev"
@@ -72,10 +70,10 @@ Requires-Dist: schema (>=0.7)
 Requires-Dist: scikit-image (>=0.23)
 Requires-Dist: scikit-learn (>=1.4)
 Requires-Dist: scipy (>=1.13)
+Requires-Dist: simpleitk (>=2.2.1) ; extra == "dev"
+Requires-Dist: simpleitk (>=2.2.1) ; extra == "tests"
 Requires-Dist: sphinx ; extra == "dev"
 Requires-Dist: sphinx ; extra == "docs"
-Requires-Dist: sphinx-argparse ; extra == "dev"
-Requires-Dist: sphinx-argparse ; extra == "docs"
 Requires-Dist: sphinx-autobuild ; extra == "dev"
 Requires-Dist: sphinx-autobuild ; extra == "docs"
 Requires-Dist: sphinx-autodoc-typehints ; extra == "dev"
@@ -86,6 +84,8 @@ Requires-Dist: sphinx-reredirects ; extra == "dev"
 Requires-Dist: sphinx-reredirects ; extra == "docs"
 Requires-Dist: sphinx_design ; extra == "dev"
 Requires-Dist: sphinx_design ; extra == "docs"
+Requires-Dist: sphinxcontrib-autoprogram ; extra == "dev"
+Requires-Dist: sphinxcontrib-autoprogram ; extra == "docs"
 Requires-Dist: sphinxcontrib-bibtex ; extra == "dev"
 Requires-Dist: sphinxcontrib-bibtex ; extra == "docs"
 Requires-Dist: tensorflow ; extra == "backends"

{zea-0.0.8 → zea-0.0.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "zea"
-version = "0.0.8"
+version = "0.0.10"
 description = "A Toolbox for Cognitive Ultrasound Imaging. Provides a set of tools for processing of ultrasound data, all built in your favorite machine learning framework."
 authors = [
     { name = "Tristan Stevens", email = "t.s.w.stevens@tue.nl" },
@@ -54,6 +54,7 @@ dev = [
     "papermill>=2.4",
     "ipykernel>=6.29.5",
     "cloudpickle>=3.1.1",
+    "simpleitk>=2.2.1",
     "ipywidgets",
     "pre-commit",
     "ruff",
@@ -64,10 +65,9 @@ dev = [
     "sphinx-autodoc-typehints",
     "sphinx-copybutton",
     "sphinx_design",
-    "sphinx-argparse",
+    "sphinxcontrib-autoprogram",
     "sphinx-reredirects",
     "sphinxcontrib-bibtex",
-    "mock",
     "myst-parser",
     "nbsphinx",
     "furo",
@@ -84,6 +84,7 @@ tests = [
     "papermill>=2.4",
     "ipykernel>=6.29.5",
     "cloudpickle>=3.1.1",
+    "simpleitk>=2.2.1",
     "ipywidgets",
     "pre-commit",
     "ruff",
@@ -95,10 +96,9 @@ docs = [
     "sphinx-autodoc-typehints",
     "sphinx-copybutton",
     "sphinx_design",
-    "sphinx-argparse",
+    "sphinxcontrib-autoprogram",
     "sphinx-reredirects",
     "sphinxcontrib-bibtex",
-    "mock",
     "myst-parser",
     "nbsphinx",
     "furo",

{zea-0.0.8 → zea-0.0.10}/zea/__init__.py

@@ -2,12 +2,16 @@
 
 import importlib.util
 import os
+from importlib.metadata import PackageNotFoundError, version
 
 from . import log
 
-# dynamically add __version__ attribute (see pyproject.toml)
-# __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.8"
+try:
+    # dynamically add __version__ attribute (see pyproject.toml)
+    __version__ = version("zea")
+except PackageNotFoundError:
+    # Package is not installed (e.g., running from source)
+    __version__ = "dev"
 
 
 def _bootstrap_backend():
@@ -80,8 +84,10 @@ def _bootstrap_backend():
     log.info(f"Using backend {keras_backend()!r}")
 
 
-# call and clean up namespace
-_bootstrap_backend()
+# Skip backend bootstrap when building on ReadTheDocs
+if os.environ.get("READTHEDOCS") != "True":
+    _bootstrap_backend()
+
 del _bootstrap_backend
 
 from . import (
@@ -89,12 +95,12 @@ from . import (
     beamform,
     data,
     display,
+    func,
     io_lib,
-    keras_ops,
     metrics,
     models,
+    ops,
     simulator,
-    tensor_ops,
     utils,
     visualize,
 )

{zea-0.0.8 → zea-0.0.10}/zea/agent/masks.py

@@ -4,13 +4,15 @@ Mask generation utilities.
 These masks are used as a measurement operator for focused scan-line subsampling.
 """
 
+from __future__ import annotations
+
 from typing import List
 
 import keras
 from keras import ops
 
-from zea import tensor_ops
 from zea.agent.gumbel import hard_straight_through
+from zea.func.tensor import nonzero
 
 _DEFAULT_DTYPE = "bool"
 
@@ -56,7 +58,7 @@ def k_hot_to_indices(selected_lines, n_actions: int, fill_value=-1):
 
     # Find nonzero indices for each frame
     def get_nonzero(row):
-        return tensor_ops.nonzero(row > 0, size=n_actions, fill_value=fill_value)[0]
+        return nonzero(row > 0, size=n_actions, fill_value=fill_value)[0]
 
     indices = ops.vectorized_map(get_nonzero, selected_lines)
     return indices
@@ -117,11 +119,21 @@ def initial_equispaced_lines(
         Tensor: k-hot-encoded line vector of shape (n_possible_actions).
             Needs to be converted to image size.
     """
+    assert n_actions > 0, "Number of actions must be > 0."
+    assert n_possible_actions > 0, "Number of possible actions must be > 0."
+    assert n_actions <= n_possible_actions, (
+        "Number of actions must be less than or equal to number of possible actions."
+    )
+
     if assert_equal_spacing:
         _assert_equal_spacing(n_actions, n_possible_actions)
-        selected_indices = ops.arange(0, n_possible_actions, n_possible_actions // n_actions)
-    else:
-        selected_indices = ops.linspace(0, n_possible_actions - 1, n_actions, dtype="int32")
+
+    # Distribute indices as evenly as possible
+    # This approach ensures spacing differs by at most 1 when not divisible
+    step = n_possible_actions / n_actions
+    selected_indices = ops.cast(
+        ops.round(ops.arange(0, n_actions, dtype="float32") * step), "int32"
+    )
 
     return indices_to_k_hot(selected_indices, n_possible_actions, dtype=dtype)
 

{zea-0.0.8 → zea-0.0.10}/zea/agent/selection.py

@@ -16,9 +16,9 @@ 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.func import tensor
 from zea.internal.registry import action_selection_registry
 
 
@@ -96,6 +96,7 @@ class GreedyEntropy(LinesActionModel):
         std_dev: float = 1,
         num_lines_to_update: int = 5,
         entropy_sigma: float = 1.0,
+        average_entropy_across_batch: bool = False,
     ):
         """Initialize the GreedyEntropy action selection model.
 
@@ -110,6 +111,10 @@ class GreedyEntropy(LinesActionModel):
                 to update. Must be odd.
             entropy_sigma (float, optional): The standard deviation of the Gaussian
                 Mixture components used to approximate the posterior.
+            average_entropy_across_batch (bool, optional): Whether to average entropy
+                across the batch when selecting lines. This can be useful when
+                selecting planes in 3D imaging, where the batch dimension represents
+                a third spatial dimension. Defaults to False.
         """
         super().__init__(n_actions, n_possible_actions, img_width, img_height)
 
@@ -117,6 +122,7 @@ class GreedyEntropy(LinesActionModel):
         # of the selected line is set to 0 once it's been selected.
         assert num_lines_to_update % 2 == 1, "num_samples must be odd."
         self.num_lines_to_update = num_lines_to_update
+        self.average_entropy_across_batch = average_entropy_across_batch
 
         # see here what I mean by upside_down_gaussian:
         # https://colab.research.google.com/drive/1CQp_Z6nADzOFsybdiH5Cag0vtVZjjioU?usp=sharing
@@ -153,7 +159,7 @@ class GreedyEntropy(LinesActionModel):
         assert particles.shape[1] > 1, "The entropy cannot be approximated using a single particle."
 
         if n_possible_actions is None:
-            n_possible_actions = particles.shape[-1]
+            n_possible_actions = ops.shape(particles)[-1]
 
         # TODO: I think we only need to compute the lower triangular
         # of this matrix, since it's symmetric
@@ -164,7 +170,8 @@ class GreedyEntropy(LinesActionModel):
         # Vertically stack all columns corresponding with the same line
         # This way we can just sum across the height axis and get the entropy
         # for each pixel in a given line
-        batch_size, n_particles, _, height, _ = gaussian_error_per_pixel_i_j.shape
+        batch_size, n_particles, _, height, _ = ops.shape(gaussian_error_per_pixel_i_j)
+
         gaussian_error_per_pixel_stacked = ops.transpose(
             ops.reshape(
                 ops.transpose(gaussian_error_per_pixel_i_j, (0, 1, 2, 4, 3)),
@@ -274,6 +281,8 @@ class GreedyEntropy(LinesActionModel):
 
         pixelwise_entropy = self.compute_pixelwise_entropy(particles)
         linewise_entropy = ops.sum(pixelwise_entropy, axis=1)
+        if self.average_entropy_across_batch:
+            linewise_entropy = ops.expand_dims(ops.mean(linewise_entropy, axis=0), axis=0)
 
         # Greedily select best line, reweight entropies, and repeat
         all_selected_lines = []
@@ -334,7 +343,7 @@ class EquispacedLines(LinesActionModel):
         n_possible_actions: int,
         img_width: int,
         img_height: int,
-        assert_equal_spacing=True,
+        assert_equal_spacing: bool = True,
     ):
         super().__init__(n_actions, n_possible_actions, img_width, img_height)
 
@@ -462,7 +471,7 @@ class CovarianceSamplingLines(LinesActionModel):
         particles = ops.reshape(particles, shape)
 
         # [batch_size, rows * stack_n_cols, n_possible_actions, n_possible_actions]
-        cov_matrix = tensor_ops.batch_cov(particles)
+        cov_matrix = tensor.batch_cov(particles)
 
         # Sum over the row dimension [batch_size, n_possible_actions, n_possible_actions]
         cov_matrix = ops.sum(cov_matrix, axis=1)
@@ -477,7 +486,7 @@ class CovarianceSamplingLines(LinesActionModel):
         # Subsample the covariance matrix with random lines
         def subsample_with_mask(mask):
             """Subsample the covariance matrix with a single mask."""
-            subsampled_cov_matrix = tensor_ops.boolean_mask(
+            subsampled_cov_matrix = tensor.boolean_mask(
                 cov_matrix, mask, size=batch_size * self.n_actions**2
             )
             return ops.reshape(subsampled_cov_matrix, [batch_size, self.n_actions, self.n_actions])

{zea-0.0.8 → zea-0.0.10}/zea/backend/__init__.py

@@ -59,8 +59,21 @@ def _import_torch():
         return None
 
 
+def _get_backend():
+    try:
+        backend_result = keras.backend.backend()
+        if isinstance(backend_result, str):
+            return backend_result
+        else:
+            # to handle mocked backends during testing
+            return None
+    except Exception:
+        return None
+
+
 tf_mod = _import_tf()
 jax_mod = _import_jax()
+backend = _get_backend()
 
 
 def tf_function(func=None, jit_compile=False, **kwargs):
@@ -131,7 +144,7 @@ class on_device:
         .. code-block:: python
 
             with zea.backend.on_device("gpu:3"):
-                pipeline = zea.Pipeline([zea.keras_ops.Abs()])
+                pipeline = zea.Pipeline([zea.ops.Abs()])
                 output = pipeline(data=keras.random.normal((10, 10)))  # output is on "cuda:3"
     """
 
@@ -184,7 +197,7 @@ class on_device:
         self._context.__exit__(exc_type, exc_val, exc_tb)
 
 
-if keras.backend.backend() in ["tensorflow", "jax", "numpy"]:
+if backend in [None, "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.
@@ -199,7 +212,8 @@ if keras.backend.backend() in ["tensorflow", "jax", "numpy"]:
         """
         with on_device(device):
             return func(*args, **kwargs)
-elif keras.backend.backend() == "torch":
+
+elif backend == "torch":
     from zea.backend.torch import func_on_device
 else:
-    raise ValueError(f"Unsupported backend: {keras.backend.backend()}")
+    raise ValueError(f"Unsupported backend: {backend}")

{zea-0.0.8 → zea-0.0.10}/zea/backend/tensorflow/dataloader.py

@@ -12,8 +12,8 @@ from keras.src.trainers.data_adapters import TFDatasetAdapter
 
 from zea.data.dataloader import H5Generator
 from zea.data.layers import Resizer
+from zea.func.tensor import translate
 from zea.internal.utils import find_methods_with_return_type
-from zea.tensor_ops import translate
 
 METHODS_THAT_RETURN_DATASET = find_methods_with_return_type(tf.data.Dataset, "DatasetV2")