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

zea/probes.py

@@ -16,14 +16,14 @@ Example usage
 
 We can initialize a generic probe with the following code:
 
-.. code-block:: python
+.. doctest::
 
-    import zea
+    >>> from zea import Probe
 
-    probe = zea.Probe.from_name("generic")
-    print(probe.get_parameters())
-
-"""
+    >>> probe = Probe.from_name("generic")
+    >>> print(probe.get_parameters())
+    {'probe_geometry': None, 'center_frequency': None, 'sampling_frequency': None, 'xlims': None, 'zlims': None}
+"""  # noqa: E501
 
 import numpy as np
 

zea/scan.py

@@ -43,38 +43,38 @@ Comparison to ``zea.Config`` and ``zea.Probe``
 Example Usage
 ^^^^^^^^^^^^^
 
-.. code-block:: python
-
-    from zea import Config, Probe, Scan
-
-    # Initialize Scan from a Probe's parameters
-    probe = Probe.from_name("verasonics_l11_4v")
-    scan = Scan(**probe.get_parameters(), grid_size_z=256)
-
-    # Or initialize from a Config object
-    config = Config.from_hf("zeahub/configs", "config_picmus_rf.yaml", repo_type="dataset")
-    scan = Scan(**config.scan, n_tx=11)
-
-    # Or manually specify parameters
-    scan = Scan(
-        grid_size_x=128,
-        grid_size_z=256,
-        xlims=(-0.02, 0.02),
-        zlims=(0.0, 0.06),
-        center_frequency=6.25e6,
-        sound_speed=1540.0,
-        sampling_frequency=25e6,
-        n_el=128,
-        n_tx=11,
-    )
-
-    # Access a derived property (computed lazily)
-    grid = scan.grid  # shape: (grid_size_z, grid_size_x, 3)
-
-    # Select a subset of transmit events
-    scan.set_transmits(3)  # Use 3 evenly spaced transmits
-    scan.set_transmits([0, 2, 4])  # Use specific transmit indices
-    scan.set_transmits("all")  # Use all transmits
+.. doctest::
+
+    >>> from zea import Config, Probe, Scan
+
+    >>> # Initialize Scan from a Probe's parameters
+    >>> probe = Probe.from_name("verasonics_l11_4v")
+    >>> scan = Scan(**probe.get_parameters(), grid_size_z=256, n_tx=11)
+
+    >>> # Or initialize from a Config object
+    >>> config = Config.from_hf("zeahub/configs", "config_picmus_rf.yaml", repo_type="dataset")
+    >>> scan = Scan(n_tx=11, **config.scan)
+
+    >>> # Or manually specify parameters
+    >>> scan = Scan(
+    ...     grid_size_x=128,
+    ...     grid_size_z=256,
+    ...     xlims=(-0.02, 0.02),
+    ...     zlims=(0.0, 0.06),
+    ...     center_frequency=6.25e6,
+    ...     sound_speed=1540.0,
+    ...     sampling_frequency=25e6,
+    ...     n_el=128,
+    ...     n_tx=11,
+    ... )
+
+    >>> # Access a derived property (computed lazily)
+    >>> grid = scan.grid  # shape: (grid_size_z, grid_size_x, 3)
+
+    >>> # Select a subset of transmit events
+    >>> _ = scan.set_transmits(3)  # Use 3 evenly spaced transmits
+    >>> _ = scan.set_transmits([0, 2, 4])  # Use specific transmit indices
+    >>> _ = scan.set_transmits("all")  # Use all transmits
 
 """
 
@@ -83,7 +83,11 @@ from keras import ops
 
 from zea import log
 from zea.beamform.pfield import compute_pfield
-from zea.beamform.pixelgrid import cartesian_pixel_grid, check_for_aliasing, polar_pixel_grid
+from zea.beamform.pixelgrid import (
+    cartesian_pixel_grid,
+    check_for_aliasing,
+    polar_pixel_grid,
+)
 from zea.display import (
     compute_scan_convert_2d_coordinates,
     compute_scan_convert_3d_coordinates,
@@ -130,6 +134,15 @@ class Scan(Parameters):
         demodulation_frequency (float, optional): Demodulation frequency in Hz.
         time_to_next_transmit (np.ndarray): The time between subsequent
             transmit events of shape (n_frames, n_tx).
+        tgc_gain_curve (np.ndarray): Time gain compensation (TGC) curve of shape (n_ax,).
+        waveforms_one_way (np.ndarray): The one-way transmit waveforms of shape
+            (n_waveforms, n_samples).
+        waveforms_two_way (np.ndarray): The two-way transmit waveforms of shape
+            (n_waveforms, n_samples).
+        tx_waveform_indices (np.ndarray): Indices of the waveform used for each
+            transmit event of shape (n_tx,).
+        t_peak (np.ndarray, optional): The time of the peak of the pulse of every transmit waveform
+            of shape (n_waveforms,).
         pixels_per_wavelength (int, optional): Number of pixels per wavelength.
             Defaults to 4.
         element_width (float, optional): Width of each transducer element in meters.
@@ -160,6 +173,7 @@ class Scan(Parameters):
             Can be "cartesian" or "polar". Defaults to "cartesian".
         dynamic_range (tuple, optional): Dynamic range for image display.
             Defined in dB as (min_dB, max_dB). Defaults to (-60, 0).
+
     """
 
     VALID_PARAMS = {
@@ -200,6 +214,11 @@ class Scan(Parameters):
         "focus_distances": {"type": np.ndarray},
         "initial_times": {"type": np.ndarray},
         "time_to_next_transmit": {"type": np.ndarray},
+        "tgc_gain_curve": {"type": np.ndarray},
+        "waveforms_one_way": {"type": np.ndarray},
+        "waveforms_two_way": {"type": np.ndarray},
+        "tx_waveform_indices": {"type": np.ndarray},
+        "t_peak": {"type": np.ndarray},
         # scan conversion parameters
         "theta_range": {"type": (tuple, list)},
         "phi_range": {"type": (tuple, list)},
@@ -209,8 +228,9 @@ class Scan(Parameters):
     }
 
     def __init__(self, **kwargs):
-        # Store the current selection state before initialization
-        selected_transmits_input = kwargs.pop("selected_transmits", None)
+        # Ensure that selected_transmits is present and set to None by default
+        selected_transmits_input = kwargs.get("selected_transmits", None)
+        kwargs["selected_transmits"] = None
 
         # Initialize parent class
         super().__init__(**kwargs)
@@ -240,7 +260,10 @@ class Scan(Parameters):
             )
         elif self.grid_type == "cartesian":
             return cartesian_pixel_grid(
-                self.xlims, self.zlims, grid_size_z=self.grid_size_z, grid_size_x=self.grid_size_x
+                self.xlims,
+                self.zlims,
+                grid_size_z=self.grid_size_z,
+                grid_size_x=self.grid_size_x,
             )
         else:
             raise ValueError(
@@ -301,7 +324,10 @@ class Scan(Parameters):
                 radius * np.cos(-np.pi / 2 + self.polar_limits[1]),
             )
             xlims_plane = (self.probe_geometry[0, 0], self.probe_geometry[-1, 0])
-            xlims = min(xlims_polar[0], xlims_plane[0]), max(xlims_polar[1], xlims_plane[1])
+            xlims = (
+                min(xlims_polar[0], xlims_plane[0]),
+                max(xlims_polar[1], xlims_plane[1]),
+            )
         return xlims
 
     @cache_with_dependencies("sound_speed", "sampling_frequency", "n_ax")
@@ -344,6 +370,11 @@ class Scan(Parameters):
         """The total number of transmits in the full dataset."""
         return self._params["n_tx"]
 
+    @property
+    def n_tx_selected(self):
+        """The number of currently selected transmits."""
+        return len(self.selected_transmits)
+
     @cache_with_dependencies("selected_transmits")
     def n_tx(self):
         """The number of currently selected transmits."""
@@ -387,14 +418,12 @@ class Scan(Parameters):
         if selection is None or selection == "all":
             self._selected_transmits = None
             self._invalidate("selected_transmits")
-            self._invalidate_dependents("selected_transmits")
             return self
 
         # Handle "center" - use center transmit
         if selection == "center":
             self._selected_transmits = [n_tx_total // 2]
             self._invalidate("selected_transmits")
-            self._invalidate_dependents("selected_transmits")
             return self
 
         # Handle integer - select evenly spaced transmits
@@ -416,7 +445,6 @@ class Scan(Parameters):
                 self._selected_transmits = list(np.rint(tx_indices).astype(int))
 
             self._invalidate("selected_transmits")
-            self._invalidate_dependents("selected_transmits")
             return self
 
         # Handle slice - convert to list of indices
@@ -436,7 +464,6 @@ class Scan(Parameters):
                 int(i) for i in selection
             ]  # Convert numpy integers to Python ints
             self._invalidate("selected_transmits")
-            self._invalidate_dependents("selected_transmits")
             return self
 
         # Aliasing check
@@ -481,7 +508,7 @@ class Scan(Parameters):
         value = self._params.get("azimuth_angles")
         if value is None:
             log.warning("No azimuth angles provided, using zeros")
-            value = np.zeros(self.n_tx_total)
+            value = np.zeros(self.n_tx_selected)
 
         return value[self.selected_transmits]
 
@@ -492,7 +519,7 @@ class Scan(Parameters):
         value = self._params.get("t0_delays")
         if value is None:
             log.warning("No transmit delays provided, using zeros")
-            return np.zeros((self.n_tx_total, self.n_el))
+            return np.zeros((self.n_tx_selected, self.n_el))
 
         return value[self.selected_transmits]
 
@@ -502,7 +529,7 @@ class Scan(Parameters):
         value = self._params.get("tx_apodizations")
         if value is None:
             log.warning("No transmit apodizations provided, using ones")
-            value = np.ones((self.n_tx_total, self.n_el))
+            value = np.ones((self.n_tx_selected, self.n_el))
 
         return value[self.selected_transmits]
 
@@ -512,7 +539,7 @@ class Scan(Parameters):
         value = self._params.get("focus_distances")
         if value is None:
             log.warning("No focus distances provided, using zeros")
-            value = np.zeros(self.n_tx_total)
+            value = np.zeros(self.n_tx_selected)
 
         return value[self.selected_transmits]
 
@@ -522,10 +549,19 @@ class Scan(Parameters):
         value = self._params.get("initial_times")
         if value is None:
             log.warning("No initial times provided, using zeros")
-            value = np.zeros(self.n_tx_total)
+            value = np.zeros(self.n_tx_selected)
 
         return value[self.selected_transmits]
 
+    @property
+    def t_peak(self):
+        """The time of the peak of the pulse in seconds of shape (n_waveforms,)."""
+        t_peak = self._params.get("t_peak")
+        if t_peak is None:
+            t_peak = np.array([1 / self.center_frequency])
+
+        return t_peak
+
     @cache_with_dependencies("selected_transmits")
     def time_to_next_transmit(self):
         """The time between subsequent transmit events of shape (n_frames, n_tx)."""
@@ -536,6 +572,23 @@ class Scan(Parameters):
         selected = self.selected_transmits
         return value[:, selected]
 
+    @cache_with_dependencies("n_ax")
+    def tgc_gain_curve(self):
+        """Time gain compensation (TGC) curve of shape (n_ax,)."""
+        value = self._params.get("tgc_gain_curve")
+        if value is None:
+            return np.ones(self.n_ax)
+        return value[: self.n_ax]
+
+    @cache_with_dependencies("selected_transmits")
+    def tx_waveform_indices(self):
+        """Indices of the waveform used for each transmit event of shape (n_tx,)."""
+        value = self._params.get("tx_waveform_indices")
+        if value is None:
+            return np.zeros(self.n_tx_selected, dtype=int)
+
+        return value[self.selected_transmits]
+
     @cache_with_dependencies(
         "sound_speed",
         "center_frequency",
@@ -545,8 +598,9 @@ class Scan(Parameters):
         "tx_apodizations",
         "grid",
         "t0_delays",
+        "pfield_kwargs",
     )
-    def pfield(self):
+    def pfield(self) -> np.ndarray:
         """Compute or return the pressure field (pfield) for weighting."""
         pfield = compute_pfield(
             sound_speed=self.sound_speed,
@@ -596,7 +650,12 @@ class Scan(Parameters):
         return coords
 
     @cache_with_dependencies(
-        "rho_range", "theta_range", "phi_range", "resolution", "grid_size_z", "grid_size_x"
+        "rho_range",
+        "theta_range",
+        "phi_range",
+        "resolution",
+        "grid_size_z",
+        "grid_size_x",
     )
     def coordinates_3d(self):
         """The coordinates for scan conversion."""
@@ -644,6 +703,7 @@ class Scan(Parameters):
         """The width of each transducer element in meters."""
         value = self._params.get("element_width")
         if value is None:
+            # assume uniform spacing
             return np.linalg.norm(self.probe_geometry[1] - self.probe_geometry[0])
         return value
 
@@ -651,5 +711,5 @@ class Scan(Parameters):
         if key == "selected_transmits":
             # If setting selected_transmits, call set_transmits to handle logic
             self.set_transmits(value)
-            return
+            return super().__setattr__(key, self.selected_transmits)
         return super().__setattr__(key, value)

zea/simulator.py

@@ -14,27 +14,30 @@ Example usage
 A simple example of simulating RF data with a single scatterer at the center of the probe. For a
 more in depth example see the notebook: :doc:`../notebooks/data/zea_simulation_example`.
 
-.. code-block:: python
-
-    raw_data = simulate_rf(
-        scatterer_positions=np.array([[0, 0, 20e-3]]),
-        scatterer_magnitudes=np.array([1.0]),
-        probe_geometry=np.stack(
-            [np.linspace(-20e-3, 20e-3, 64), np.zeros(64), np.zeros(64)], axis=-1
-        ),
-        apply_lens_correction=True,
-        lens_thickness=1e-3,
-        lens_sound_speed=1000,
-        sound_speed=1540,
-        n_ax=1024,
-        center_frequency=5e6,
-        sampling_frequency=20e6,
-        t0_delays=np.zeros((1, 64)),
-        initial_times=np.zeros(1),
-        element_width=0.2e-3,
-        attenuation_coef=0.5,
-        tx_apodizations=np.ones((1, 64)),
-    )
+.. doctest::
+
+    >>> from zea.simulator import simulate_rf
+    >>> import numpy as np
+
+    >>> raw_data = simulate_rf(
+    ...     scatterer_positions=np.array([[0, 0, 20e-3]]),
+    ...     scatterer_magnitudes=np.array([1.0]),
+    ...     probe_geometry=np.stack(
+    ...         [np.linspace(-20e-3, 20e-3, 64), np.zeros(64), np.zeros(64)], axis=-1
+    ...     ),
+    ...     apply_lens_correction=True,
+    ...     lens_thickness=1e-3,
+    ...     lens_sound_speed=1000,
+    ...     sound_speed=1540,
+    ...     n_ax=1024,
+    ...     center_frequency=5e6,
+    ...     sampling_frequency=20e6,
+    ...     t0_delays=np.zeros((1, 64)),
+    ...     initial_times=np.zeros(1),
+    ...     element_width=0.2e-3,
+    ...     attenuation_coef=0.5,
+    ...     tx_apodizations=np.ones((1, 64)),
+    ... )
 
 """