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

zea/__init__.py

@@ -7,7 +7,7 @@ from . import log
 
 # dynamically add __version__ attribute (see pyproject.toml)
 # __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.3"
+__version__ = "0.0.4"
 
 
 def setup():

zea/agent/selection.py

@@ -155,7 +155,9 @@ class GreedyEntropy(LinesActionModel):
         # TODO: I think we only need to compute the lower triangular
         # of this matrix, since it's symmetric
         squared_l2_error_matrices = (particles[:, :, None, ...] - particles[:, None, :, ...]) ** 2
-        gaussian_error_per_pixel_i_j = ops.exp((squared_l2_error_matrices) / (2 * entropy_sigma**2))
+        gaussian_error_per_pixel_i_j = ops.exp(
+            -(squared_l2_error_matrices) / (2 * entropy_sigma**2)
+        )
         # 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
@@ -176,33 +178,35 @@ class GreedyEntropy(LinesActionModel):
         # [n_particles, n_particles, batch, height, width]
         return gaussian_error_per_pixel_stacked
 
-    def compute_gmm_entropy_per_line(self, particles):
-        """Compute the entropy for each line using a Gaussian Mixture Model.
-
+    def compute_pixelwise_entropy(self, particles):
+        """
         This function computes the entropy for each line using a Gaussian Mixture Model
         approximation of the posterior distribution.
-        For more details see Section 4 here: https://arxiv.org/abs/2406.14388
+        For more details see Section VI. B here: https://arxiv.org/pdf/2410.13310
 
         Args:
             particles (Tensor): Particles of shape (batch_size, n_particles, height, width)
 
         Returns:
-            Tensor: batch of entropies per line, of shape (batch, n_possible_actions)
+            Tensor: batch of entropies per pixel, of shape (batch, height, width)
         """
-        gaussian_error_per_pixel_stacked = GreedyEntropy.compute_pairwise_pixel_gaussian_error(
+        n_particles = ops.shape(particles)[1]
+        gaussian_error_per_pixel_stacked = self.compute_pairwise_pixel_gaussian_error(
             particles,
             self.stack_n_cols,
             self.n_possible_actions,
             self.entropy_sigma,
         )
-        gaussian_error_per_line = ops.sum(gaussian_error_per_pixel_stacked, axis=3)
         # sum out first dimension of (n_particles x n_particles) error matrix
-        # [n_particles, batch, n_possible_actions]
-        entropy_per_line_i = ops.sum(gaussian_error_per_line, axis=1)
+        # [n_particles, batch, height, width]
+        pixelwise_entropy_sum_j = ops.sum(
+            (1 / n_particles) * gaussian_error_per_pixel_stacked, axis=1
+        )
+        log_pixelwise_entropy_sum_j = ops.log(pixelwise_entropy_sum_j)
         # sum out second dimension of (n_particles x n_particles) error matrix
-        # [batch, n_possible_actions]
-        entropy_per_line = ops.sum(entropy_per_line_i, axis=1)
-        return entropy_per_line
+        # [batch, height, width]
+        pixelwise_entropy = -ops.sum((1 / n_particles) * log_pixelwise_entropy_sum_j, axis=1)
+        return pixelwise_entropy
 
     def select_line_and_reweight_entropy(self, entropy_per_line):
         """Select the line with maximum entropy and reweight the entropies.
@@ -260,17 +264,19 @@ class GreedyEntropy(LinesActionModel):
             particles (Tensor): Particles of shape (batch_size, n_particles, height, width)
 
         Returns:
-           Tuple[Tensor, Tensor]:
+            Tuple[Tensor, Tensor]:
                 - Newly selected lines as k-hot vectors, shaped (batch_size, n_possible_actions)
-                - Masks of shape (batch_size, img_height, img_width)
+            - Masks of shape (batch_size, img_height, img_width)
         """
-        entropy_per_line = self.compute_gmm_entropy_per_line(particles)
+
+        pixelwise_entropy = self.compute_pixelwise_entropy(particles)
+        linewise_entropy = ops.sum(pixelwise_entropy, axis=1)
 
         # Greedily select best line, reweight entropies, and repeat
         all_selected_lines = []
         for _ in range(self.n_actions):
-            max_entropy_line, entropy_per_line = ops.vectorized_map(
-                self.select_line_and_reweight_entropy, entropy_per_line
+            max_entropy_line, linewise_entropy = ops.vectorized_map(
+                self.select_line_and_reweight_entropy, linewise_entropy
             )
             all_selected_lines.append(max_entropy_line)
 

zea/data/data_format.py

@@ -468,32 +468,34 @@ def _write_datasets(
             ),
             unit="-",
         )
-        n_waveforms = len(waveforms_one_way)
-        for n, waveform_1way, waveform_2way in zip(
-            range(n_waveforms), waveforms_one_way, waveforms_two_way
-        ):
-            _add_dataset(
-                group_name=scan_group_name + "/waveforms_one_way",
-                name=f"waveform_{str(n).zfill(3)}",
-                data=waveform_1way,
-                description=(
-                    "One-way waveform as simulated by the Verasonics system, "
-                    "sampled at 250MHz. This is the waveform after being filtered "
-                    "by the tranducer bandwidth once."
-                ),
-                unit="V",
-            )
-            _add_dataset(
-                group_name=scan_group_name + "/waveforms_two_way",
-                name=f"waveform_{str(n).zfill(3)}",
-                data=waveform_2way,
-                description=(
-                    "Two-way waveform as simulated by the Verasonics system, "
-                    "sampled at 250MHz. This is the waveform after being filtered "
-                    "by the tranducer bandwidth twice."
-                ),
-                unit="V",
-            )
+
+        if waveforms_one_way is not None:
+            for n in range(len(waveforms_one_way)):
+                _add_dataset(
+                    group_name=scan_group_name + "/waveforms_one_way",
+                    name=f"waveform_{str(n).zfill(3)}",
+                    data=waveforms_one_way[n],
+                    description=(
+                        "One-way waveform as simulated by the Verasonics system, "
+                        "sampled at 250MHz. This is the waveform after being filtered "
+                        "by the tranducer bandwidth once."
+                    ),
+                    unit="V",
+                )
+
+        if waveforms_two_way is not None:
+            for n in range(len(waveforms_two_way)):
+                _add_dataset(
+                    group_name=scan_group_name + "/waveforms_two_way",
+                    name=f"waveform_{str(n).zfill(3)}",
+                    data=waveforms_two_way[n],
+                    description=(
+                        "Two-way waveform as simulated by the Verasonics system, "
+                        "sampled at 250MHz. This is the waveform after being filtered "
+                        "by the tranducer bandwidth twice."
+                    ),
+                    unit="V",
+                )
 
     # Add additional elements
     if additional_elements is not None:

zea/internal/core.py

@@ -185,7 +185,7 @@ def dict_to_tensor(dictionary, keep_as_is=None):
             continue
 
         # Convert the value to a tensor
-        snapshot[key] = _to_tensor(key, value)
+        snapshot[key] = _to_tensor(key, value, keep_as_is=keep_as_is)
 
     return snapshot
 

zea/internal/viewer.py

@@ -28,7 +28,7 @@ def plt_window_has_been_closed(fig):
     return not plt.fignum_exists(fig.number)
 
 
-def filename_from_window_dialog(window_name=None, filetypes=None, initialdir=None):
+def filename_from_window_dialog(window_name=None, filetypes=None, initialdir=None) -> Path:
     """Get filename through dialog window
     Args:
         window_name: string with name of window

zea/log.py

@@ -289,6 +289,14 @@ def critical(message, *args, **kwargs):
     return message
 
 
+def number_to_str(number, decimals=2):
+    """Formats a number to a string with the given number of decimals."""
+    if isinstance(number, (int, float)):
+        return f"{number:.{decimals}f}"
+    else:
+        raise ValueError(f"Expected a number, got {type(number)}: {number}")
+
+
 def set_file_logger_directory(directory):
     """Sets the log level of the logger."""
     global LOG_DIR, file_logger

zea/models/diffusion.py

@@ -161,16 +161,15 @@ class DiffusionModel(DeepGenerativeModel):
                     f"DiffusionGuidance object, got {guidance}"
                 )
 
-    def call(self, inputs, training=False, **kwargs):
+    def call(self, inputs, training=False, network=None, **kwargs):
+        """Calls the score network.
+
+        If network is not provided, will use the exponential moving
+        average network if training is False, otherwise the regular network.
         """
-        Calls the score network.
+        if network is None:
+            network = self.network if training else self.ema_network
 
-        Will use the exponential moving average network if training is False,
-        otherwise the regular network."""
-        if training:
-            network = self.network
-        else:
-            network = self.ema_network
         return network(inputs, training=training, **kwargs)
 
     def sample(self, n_samples=1, n_steps=20, seed=None, **kwargs):
@@ -367,18 +366,27 @@ class DiffusionModel(DeepGenerativeModel):
 
     def linear_diffusion_schedule(self, diffusion_times):
         """Create a linear diffusion schedule"""
-        compute_alpha_t = lambda t: ops.prod(
-            1 - diffusion_times[:t], axis=diffusion_times.shape[1:]
-        )
-        alphas = ops.vectorized_map(compute_alpha_t, ops.arange(len(diffusion_times)))
+
+        def _compute_alpha_t(t):
+            """Compute alpha_t for linear diffusion schedule"""
+            return ops.prod(1 - diffusion_times[:t], axis=diffusion_times.shape[1:])
+
+        alphas = ops.vectorized_map(_compute_alpha_t, ops.arange(len(diffusion_times)))
         signal_rates = ops.sqrt(alphas)
         noise_rates = ops.sqrt(1 - alphas)
         return signal_rates, noise_rates
 
-    def denoise(self, noisy_images, noise_rates, signal_rates, training):
-        """Predict noise component and calculate the image component using it"""
+    def denoise(
+        self,
+        noisy_images,
+        noise_rates,
+        signal_rates,
+        training,
+        network=None,
+    ):
+        """Predict noise component and calculate the image component using it."""
 
-        pred_noises = self([noisy_images, noise_rates**2], training=training)
+        pred_noises = self([noisy_images, noise_rates**2], training=training, network=network)
         pred_images = (noisy_images - noise_rates * pred_noises) / signal_rates
 
         return pred_noises, pred_images
@@ -435,6 +443,9 @@ class DiffusionModel(DeepGenerativeModel):
         seed: keras.random.SeedGenerator | None = None,
         verbose: bool = True,
         track_progress_type: Literal[None, "x_0", "x_t"] = "x_0",
+        disable_jit: bool = False,
+        training: bool = False,
+        network_type: Literal[None, "main", "ema"] = None,
     ):
         """Reverse diffusion process to generate images from noise.
 
@@ -447,6 +458,10 @@ class DiffusionModel(DeepGenerativeModel):
             seed: Random seed generator.
             verbose: Whether to show a progress bar.
             track_progress_type: Type of progress tracking ("x_0" or "x_t").
+            disable_jit: Whether to disable JIT compilation.
+            training: Whether to use the training mode of the network.
+            network_type: Which network to use ("main" or "ema"). If None, uses the
+                network based on the `training` argument.
 
         Returns:
             Generated images.
@@ -478,8 +493,19 @@ class DiffusionModel(DeepGenerativeModel):
             next_noise_rates, next_signal_rates = self.diffusion_schedule(next_diffusion_times)
 
             # denoise
+            if network_type == "ema":
+                network = self.ema_network
+            elif network_type == "main":
+                network = self.network
+            else:
+                network = None
+
             pred_noises, pred_images = self.denoise(
-                noisy_images, noise_rates, signal_rates, training=False
+                noisy_images,
+                noise_rates,
+                signal_rates,
+                training=training,
+                network=network,
             )
 
             seed, seed1 = split_seed(seed, 2)
@@ -515,7 +541,7 @@ class DiffusionModel(DeepGenerativeModel):
                 seed,
             ),
             # can't jit this with progbar or tracking intermediate values
-            disable_jit=verbose or track_progress_type,
+            disable_jit=verbose or track_progress_type or disable_jit,
         )
 
         return pred_images

zea/ops.py

@@ -100,7 +100,7 @@ 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 deep_compare, map_negative_indices, translate
+from zea.utils import FunctionTimer, deep_compare, map_negative_indices, translate
 
 
 def get_ops(ops_name):
@@ -378,6 +378,7 @@ class Pipeline:
         jit_kwargs: dict | None = None,
         name="pipeline",
         validate=True,
+        timed: bool = False,
     ):
         """Initialize a pipeline
 
@@ -399,6 +400,8 @@ class Pipeline:
         """
         self._call_pipeline = self.call
         self.name = name
+        self.timer = FunctionTimer()
+        self.timed = timed
 
         self._pipeline_layers = operations
 
@@ -519,6 +522,30 @@ class Pipeline:
         """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."""
+
+        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 call(self, **inputs):
         """Process input data through the pipeline."""
         for operation in self._pipeline_layers:
@@ -605,13 +632,18 @@ 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, **self.jit_kwargs)
+        self._call_pipeline = jit(self._call_fn, **self.jit_kwargs)
 
     def unjit(self):
         """Un-JIT compile the pipeline."""
-        self._call_pipeline = self.call
+        self._call_pipeline = self._call_fn
 
     @property
     def jittable(self):
@@ -2090,7 +2122,7 @@ class Demodulate(Operation):
 
 @ops_registry("lambda")
 class Lambda(Operation):
-    """Use any funcion as an operation."""
+    """Use any function as an operation."""
 
     def __init__(self, func, func_kwargs=None, **kwargs):
         super().__init__(**kwargs)

zea/utils.py

@@ -662,3 +662,34 @@ class FunctionTimer:
             yaml.dump(cropped_timings, f, default_flow_style=False)
 
         self.last_append = len(self.timings[func_name])
+
+    def print(self, drop_first: bool | int = False):
+        """Print timing statistics for all recorded functions using formatted output."""
+
+        # Print title
+        print(log.bold("Function Timing Statistics"))
+        header = (
+            f"{log.cyan('Function'):<30} {log.green('Mean'):<22} "
+            f"{log.green('Median'):<22} {log.green('Std Dev'):<22} "
+            f"{log.yellow('Min'):<22} {log.yellow('Max'):<22} {log.magenta('Count'):<18}"
+        )
+        length = len(log.remove_color_escape_codes(header))
+        print("=" * length)
+
+        # Print header
+        print(header)
+        print("-" * length)
+
+        # Print data rows
+        for func_name in self.timings.keys():
+            stats = self.get_stats(func_name, drop_first=drop_first)
+            row = (
+                f"{log.cyan(func_name):<30} "
+                f"{log.green(log.number_to_str(stats['mean'], 6)):<22} "
+                f"{log.green(log.number_to_str(stats['median'], 6)):<22} "
+                f"{log.green(log.number_to_str(stats['std_dev'], 6)):<22} "
+                f"{log.yellow(log.number_to_str(stats['min'], 6)):<22} "
+                f"{log.yellow(log.number_to_str(stats['max'], 6)):<22} "
+                f"{log.magenta(str(stats['count'])):<18}"
+            )
+            print(row)

zea/visualize.py

@@ -169,7 +169,7 @@ def plot_image_grid(
     return fig, fig_contents
 
 
-def plot_quadrants(ax, array, fixed_coord, cmap, slice_index, stride=1, centroid=None):
+def plot_quadrants(ax, array, fixed_coord, cmap, slice_index, stride=1, centroid=None, **kwargs):
     """
     For a given 3D array, plot a plane with fixed_coord using four individual quadrants.
 
@@ -183,6 +183,7 @@ def plot_quadrants(ax, array, fixed_coord, cmap, slice_index, stride=1, centroid
         stride (int, optional): The stride step for plotting. Defaults to 1.
         centroid (tuple, optional): centroid around which to break the quadrants.
             If None, the middle of the image is used.
+        **kwargs: Additional keyword arguments for the plot_surface method.
 
     Returns:
         matplotlib.axes.Axes3DSubplot: The axis with the plotted quadrants.
@@ -238,6 +239,7 @@ def plot_quadrants(ax, array, fixed_coord, cmap, slice_index, stride=1, centroid
                 cstride=stride,
                 facecolors=facecolors,
                 shade=False,
+                **kwargs,
             )
         elif fixed_coord == "y":
             X, Z = np.mgrid[: quadrant.shape[0] + 1, : quadrant.shape[1] + 1]
@@ -252,6 +254,7 @@ def plot_quadrants(ax, array, fixed_coord, cmap, slice_index, stride=1, centroid
                 cstride=stride,
                 facecolors=facecolors,
                 shade=False,
+                **kwargs,
             )
         elif fixed_coord == "z":
             X, Y = np.mgrid[: quadrant.shape[0] + 1, : quadrant.shape[1] + 1]
@@ -266,6 +269,7 @@ def plot_quadrants(ax, array, fixed_coord, cmap, slice_index, stride=1, centroid
                 cstride=stride,
                 facecolors=facecolors,
                 shade=False,
+                **kwargs,
             )
     return ax
 
@@ -281,6 +285,7 @@ def plot_biplanes(
     show_axes=None,
     fig=None,
     ax=None,
+    **kwargs,
 ):
     """
     Plot three intersecting planes from a 3D volume in 3D space.
@@ -301,6 +306,7 @@ def plot_biplanes(
             Defaults to None. Can be used to reuse the figure in a loop.
         ax (matplotlib.axes.Axes3DSubplot, optional): Matplotlib 3D axes object.
             Defaults to None. Can be used to reuse the axes in a loop.
+        **kwargs: Additional keyword arguments for the plot_surface method.
 
     Returns:
         tuple: A tuple containing the figure and axes objects (fig, ax).
@@ -340,11 +346,11 @@ def plot_biplanes(
         ax.zaxis.pane.fill = False
 
     if slice_x is not None:
-        plot_quadrants(ax, volume, "x", cmap=cmap, slice_index=slice_x, stride=stride)
+        plot_quadrants(ax, volume, "x", cmap=cmap, slice_index=slice_x, stride=stride, **kwargs)
     if slice_y is not None:
-        plot_quadrants(ax, volume, "y", cmap=cmap, slice_index=slice_y, stride=stride)
+        plot_quadrants(ax, volume, "y", cmap=cmap, slice_index=slice_y, stride=stride, **kwargs)
     if slice_z is not None:
-        plot_quadrants(ax, volume, "z", cmap=cmap, slice_index=slice_z, stride=stride)
+        plot_quadrants(ax, volume, "z", cmap=cmap, slice_index=slice_z, stride=stride, **kwargs)
 
     # Optionally show axes
     if show_axes:

{zea-0.0.3.dist-info → zea-0.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zea
-Version: 0.0.3
+Version: 0.0.4
 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.
 Keywords: ultrasound,machine learning,beamforming
 Author: Tristan Stevens

{zea-0.0.3.dist-info → zea-0.0.4.dist-info}/RECORD

@@ -1,9 +1,9 @@
-zea/__init__.py,sha256=tJNjKpJfAMmUrUdL4Ob9yuP1VZltNoFJ45Nliv20r-E,2274
+zea/__init__.py,sha256=-et4Iz8oJRRaQahurztVr7Rvz9Do7WkXfq5D5CuhX7w,2274
 zea/__main__.py,sha256=QQaxECJxTUo99q43OobV1Tc9m4POtmjRbeKqpkCZ4Lo,2020
 zea/agent/__init__.py,sha256=uJjMiPvvCXmUxC2mkSkh1Q9Ege0Vf7WazKX1_Ul80GY,924
 zea/agent/gumbel.py,sha256=WbvSrM8meXtZmDLkyXDGUyRzjB4dWZmNO_qzloRxB_s,3770
 zea/agent/masks.py,sha256=qdSGbTs9449hUxcX6gAl_s47mrs1FKI6U_T2KjS-iz8,6581
-zea/agent/selection.py,sha256=JPMlkJ1Du9KSUtRROi0g7VRYYXBLZ5NjY_LORWOjmDg,19650
+zea/agent/selection.py,sha256=j9FejT28jRGnRIEL6F1YZYdl22G0XxKOkua9qhKavIg,19764
 zea/backend/__init__.py,sha256=9evPyuxJtIeleZlX_eAggJtxxKRofLQiD_ZV3yARg2Y,4227
 zea/backend/autograd.py,sha256=buWy19ctDFsAoktZaLm5qyLN9nQRicI91-V3ND-t3f4,6868
 zea/backend/jax/__init__.py,sha256=kU89ZZeVQOm9tdrlGSLLlkrJzhUNW0wOet5lnU5BMPU,2168
@@ -43,7 +43,7 @@ zea/data/convert/echonetlvh/precompute_crop.py,sha256=A-dSbGngEvL5QfAid3d9OoZRnx
 zea/data/convert/images.py,sha256=8a2oHqAA_fLhFuNnqBsNCjKTdX5pADX-ntZ6BHUoZdQ,5229
 zea/data/convert/matlab.py,sha256=nep-n-CK7OoTQQytDZuMx3dTa7AngmYRodE55ZTCCvI,41992
 zea/data/convert/picmus.py,sha256=1TK1vDWzfHmnFErtt0ad36TInIiy4fm3nMmfx83G5nw,6203
-zea/data/data_format.py,sha256=Lb20YhiZaZ5ZvXNElkaI2DjNEiSHhscnS5hfKie-gdE,25569
+zea/data/data_format.py,sha256=n4OTQKl1newgvTamlPprAsRubAVt3bPeNYHJXNc89G0,25684
 zea/data/dataloader.py,sha256=ZAlWbZB6F43EtLvRcUzfeHVtMhYViJTvyaRb9CmpEFc,14865
 zea/data/datasets.py,sha256=XcNlJVmkfNo7TXp80y4Si_by75Ix2ULIU_iEiNEhl6Q,24492
 zea/data/file.py,sha256=dG6if-hofBzfOTr_39_X_OFlTx7t8s7ggrYR7XwyhI8,29387
@@ -59,22 +59,22 @@ zea/internal/config/create.py,sha256=f8dy6YcAUZbROxwCvORI26ACMRO4Q7EL188KyTL8Bpk
 zea/internal/config/parameters.py,sha256=3YE8aBeiosxTuIT0RAEEMkJPmkykSyZZ7hbmmyGnsuc,6602
 zea/internal/config/validation.py,sha256=vPGFWC-UVtNh8jqfQ4V8HTUIjfnOtsJrqcG2R2ae5hg,6522
 zea/internal/convert.py,sha256=EDg6vY73vQht-AwwrYUJ2XrE2L64qL3JBXxnft7Xw70,4596
-zea/internal/core.py,sha256=2Cg6Xe1E_ZTaHM56F1VPTkt_AUUS4-coCxuB-y5T15o,9364
+zea/internal/core.py,sha256=QAiN1yFgIm8PGTn8RyE3LyZd5YWpbxuSGwbjfbEYJzU,9387
 zea/internal/device.py,sha256=pcgKHWudolflxe6umQh5DVNDz8QtQBZ8tDiWDNwTays,15175
 zea/internal/git_info.py,sha256=vEeN7cdppNIJPRkC69pUQqtAfTdwCN-tpfo702xpGzY,1040
 zea/internal/operators.py,sha256=0CMv4s0k8MIlD0uMKcxiHq4YryIfacklBFnCqf7RO0A,1837
 zea/internal/parameters.py,sha256=MbYYQQuESQtzXA1b7S3psESdfTbntzQCgljkV_zTlLw,17857
 zea/internal/registry.py,sha256=lQsJbYUz1S2eKzE-7shRYWUBulX2TjHTRN37evrYIGA,7884
 zea/internal/setup_zea.py,sha256=P8dmHk_0hwfukaf-DfPxeDofOH4Kj1s5YvUs8yeFqAQ,7800
-zea/internal/viewer.py,sha256=RYEHJqLncAHhnRMUVmzS8gcpF2XsmIetseE3EADOaIA,15737
+zea/internal/viewer.py,sha256=VsDehOQgqvVp7gkw0kz2Y2v_My8P8FshJEOCM4E0RIk,15745
 zea/io_lib.py,sha256=z65zGUDaRsxeYnddgy9LtEkbHfKhW3O3wPumqozjZvg,12527
-zea/log.py,sha256=eougfMaJaCPvIb7m7GDqSQiV3-NM5LO7ciOV6MW26XM,10257
+zea/log.py,sha256=UJIL91lHUgWc-vrlJWOn1PX60fX3TFQ5slIs-b5EWEQ,10540
 zea/metrics.py,sha256=22BonPRPs6sfpDqtSEQGC1icNhsSywBak4iaf_1cook,4894
 zea/models/__init__.py,sha256=gBW1pXrD01beK10dlV6GeY6H17Ym_6wrzMDXrbUs1dk,4183
 zea/models/base.py,sha256=_l1RlXIYS2V5py-en5llJpX1oU0IXK_hzLhfCYybzHg,7121
 zea/models/carotid_segmenter.py,sha256=qxv4xSSbwLQ3AWeP8xoVCFhpPLOqsN-4dNq1ECkG3FM,5401
 zea/models/dense.py,sha256=EwrSawfSTi9oWE55Grr7jtwLXC9MNqEOO3su4oYHqfg,4067
-zea/models/diffusion.py,sha256=ehka0B0OjlJNYLLIZG3kGBURNeN6R5jvg5ftJY0iRsw,30951
+zea/models/diffusion.py,sha256=DJhrV_NGgRFK8p3kIBelM-0mqvmuDbdxRyWugaftv-0,31841
 zea/models/echonet.py,sha256=toLw6QjpdROvGrm6VEuqHuxxKkiAWtX-__5YzrE-LJA,6346
 zea/models/generative.py,sha256=iujicyFDuCD7NEk_cZ8thlZ2Rl3Qa8LfkwPsZdWYpR0,2625
 zea/models/gmm.py,sha256=6YYoiizsD7BtISOToEXkki6Cc8iXMkgmPH_rMFQKs3E,8324
@@ -85,7 +85,7 @@ zea/models/presets.py,sha256=drTQF2XCAxMFSqZtu_6m-kpPtcG6io4DrXsBcdFss-4,2623
 zea/models/taesd.py,sha256=Vab5jywo4uxXPXMQ-8VdpOMhVwQsEc_PNztctDzFZfk,8403
 zea/models/unet.py,sha256=B_600WLn3jj_BuhDdRrj5I1XQIHylaQQN8ISq-vt7zM,6392
 zea/models/utils.py,sha256=My6VY27S4udOn7xatIM22Qgn8jED1FmnA2yZw1mcuVw,2015
-zea/ops.py,sha256=At97g7ibM4X66_S2FauALulgDwAiJ1fer-_Wd1pIzwU,108080
+zea/ops.py,sha256=wgm_p1rbQFFQ5856u2SIYGGTv3dFaJR1ogdMkiteChc,109464
 zea/probes.py,sha256=991X4ilpMablekNxAHwBU6DoBIZHrmzoJgs2C6zfU0U,7632
 zea/scan.py,sha256=dp-PiQZWk4HnRROeu6cwrGcM8FNjezDzkF05ZJcacVI,26935
 zea/simulator.py,sha256=KziYNRNAwIyKNE4nzVp2cbNyLATaaUkF-Ity6O1fx20,11279
@@ -95,11 +95,11 @@ zea/tools/fit_scan_cone.py,sha256=Pw1kiXUAXenCVIO1nfqWVne2i0eWWXHFmoAQ7l_Qzro,24
 zea/tools/hf.py,sha256=ibuTGaLitK1vOlQdZJuVmDBoUkYRrEfCctj3kiiOQGs,5477
 zea/tools/selection_tool.py,sha256=dNksdVw_sNcelAyszbrseOTQRDxJpwlapabdh7feNoQ,29848
 zea/tools/wndb.py,sha256=8XY056arnDKpVV7k-B5PrMa-RANur3ldPSR4GW3jgS4,666
-zea/utils.py,sha256=0T98hhK0O_GjeA-8oFL4U7Sm0qbb40EJ7I3hDfonA-A,22041
-zea/visualize.py,sha256=fB3-NskzBCHAd_EZTaUcRzcEIMUnvRViNLB7e7yjGu8,23621
+zea/utils.py,sha256=8gxZaojsC0D2S63gs6lN0ulcSJ1jry3AzOMFvBdl4H4,23351
+zea/visualize.py,sha256=4-bEHRYG3hUukmBkx7x4VBoa3oAJffO3HnnnzTfduDE,23905
 zea/zea_darkmode.mplstyle,sha256=wHTXkgy00tLEbRmr8GZULb5zIzU0MTMn9xC0Z3WT7Bo,42141
-zea-0.0.3.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-zea-0.0.3.dist-info/METADATA,sha256=8FDMbWb-pzmaH4xrAoEm8SqvFWACh9K2cmE1ZoM3PHs,6489
-zea-0.0.3.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-zea-0.0.3.dist-info/entry_points.txt,sha256=hQcQYCHdMu2LRM1PGZuaGU5EwAjTGErC-QakgwZKZeo,41
-zea-0.0.3.dist-info/RECORD,,
+zea-0.0.4.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+zea-0.0.4.dist-info/METADATA,sha256=VL_XaEx_hcxAPQTU2EcGEaz1RfYxbZ5p_5XRfmqOA6Q,6489
+zea-0.0.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+zea-0.0.4.dist-info/entry_points.txt,sha256=hQcQYCHdMu2LRM1PGZuaGU5EwAjTGErC-QakgwZKZeo,41
+zea-0.0.4.dist-info/RECORD,,