Port over minimal solution analysis code from matlab (#104)

src/openlifu/__init__.py

@@ -15,6 +15,9 @@
     apod_methods,
     delay_methods,
     focal_patterns,
+    get_beamwidth,
+    mask_focus,
+    offset_grid,
 )
 from openlifu.db import Database
 
@@ -59,6 +62,9 @@
     "focal_patterns",
     "delay_methods",
     "apod_methods",
+    "get_beamwidth",
+    "mask_focus",
+    "offset_grid",
     "SimSetup",
     "Database",
     "__version__",

src/openlifu/bf/__init__.py

@@ -1,6 +1,9 @@
 from .apod_methods import ApodizationMethod
 from .delay_methods import DelayMethod
 from .focal_patterns import FocalPattern, SinglePoint, Wheel
+from .get_beamwidth import get_beamwidth
+from .mask_focus import mask_focus
+from .offset_grid import offset_grid
 from .pulse import Pulse
 from .sequence import Sequence
 
@@ -12,4 +15,7 @@
     "SinglePoint",
     "Pulse",
     "Sequence",
+    "offset_grid",
+    "mask_focus",
+    "get_beamwidth"
 ]

src/openlifu/bf/get_beamwidth.py

@@ -0,0 +1,138 @@
+import logging
+from dataclasses import dataclass
+from typing import Any, Dict
+
+import numpy as np
+from scipy.spatial import ConvexHull, QhullError  # pylint: disable = no-name-in-module
+from xarray import DataArray
+
+from openlifu.bf import offset_grid
+from openlifu.geo import Point
+from openlifu.util.units import getunitconversion
+
+
+@dataclass
+class BeamwidthResult:
+    dims: tuple
+    beamwidth: float
+    units: str
+    inlier_mask: np.ndarray = None
+    fit_mask: np.ndarray = None
+    inlier_points: np.ndarray = None
+    fit_points: np.ndarray = None
+    inlier_hull: np.ndarray = None
+    fit_hull: np.ndarray = None
+
+
+def get_beamwidth(vol: DataArray, coords_units: str, focus: Point, cutoff: float, options: Dict[str, Any]) -> BeamwidthResult:
+    """
+    Calculates the beam width of a volume at a given point.
+
+    Args:
+        vol: xarray.DataArray
+            The input DataArray.
+        coords_units: str
+            The unit of data coordinates.
+        focus : Point
+            The focus point to be used as reference.
+        cutoff: float
+            Cutoff value for the volume data.
+        options: dictionary
+            Additional parameters:
+                - 'units': str
+                    Target units. Default focus.units.
+                - 'dims': tuple
+                    Dimension ids to operate on. Defaults (0, 1).
+                - 'mask': np.ndarray[bool]
+                    Search mask.
+                - 'hulls': bool
+                    Whether to use convex hulls. Defaults False.
+                - 'points': bool
+                    Whether to output a list of points for the beam width. Defaults False.
+                - 'masks': np.ndarray[bool]
+                    Whether to output a mask for the beam width. Defaults False.
+                - 'simplify_hulls': bool
+                    Whether to simplify the convex hull using Delaunay. Defaults True.
+
+    Returns:
+        A dictionary containing beam width, units and optionally other data.
+    """
+    # Define default values for options if not provided
+    if not hasattr(options, 'units') or options['units'] is None:
+        options['units'] = focus.units
+    if not hasattr(options, 'dims') or options['dims'] is None:
+        options['dims'] = (0, 1)
+    if not hasattr(options, 'hulls') or options['hulls'] is None:
+        options['hulls'] = False
+    if not hasattr(options, 'points') or options['points'] is None:
+        options['points'] = False
+    if not hasattr(options, 'masks') or options['masks'] is None:
+        options['masks'] = False
+    if not hasattr(options, 'simplify_hulls') or options['simplify_hulls'] is None:
+        options['simplify_hulls'] = True
+
+    # Get coordinates of volume
+    coords = vol.coords
+
+    scale = getunitconversion(coords_units, options['units'])
+    coords_rescaled = [coord * scale for coord in coords]  # equivalent of coords.rescale(args["units"])
+
+    if options['mask'] is None:
+        search_mask = np.ones(coords_rescaled.shape)
+    else:
+        search_mask = options['mask']
+
+    ngrid0 = np.meshgrid(*coords, indexing='ij')
+    mdata = search_mask * vol.data
+    inlier_mask = mdata > cutoff
+    ogrid = offset_grid(coords, focus)
+    omask = [ogrid[..., ii][inlier_mask] for ii in range(ogrid.shape[-1])]  #TODO: better to use vectorization here omask = ogrid[inlier_mask]
+    inlier_points = [ngrid0[ii][inlier_mask] for ii in range(len(ngrid0))]
+    inlier_points = np.stack(inlier_points, axis=-1)
+
+    try:
+        # Create convex hull(s) from the set of inlier points
+        if options['simplify_hulls']:
+            inlier_hull = ConvexHull(inlier_points)  #TODO: simplification not implemented in scipy ConvexHull
+        else:
+            inlier_hull = ConvexHull(inlier_points)
+    except QhullError:
+        # If convex hull creation fails (e.g., too few points), add jitter and try again
+        logging.warning("Invalid inliers, attempting to add jitter to create a valid volume...")  #TODO: should be using self.logger
+        minmax_coords = np.array([(np.min(coords[i]), np.max(coords[i])) for i in range(len(coords))])  #TODO: min-max should be from coords.extent
+        coords_shape = tuple([len(coords[i]) for i in range(len(coords))])
+        dx = np.mean(np.diff(minmax_coords) / (np.array(coords_shape) - 1))
+        rng = np.random.default_rng()
+        inlier_points = inlier_points + (rng.random(inlier_points.shape) - 0.5)*dx/2
+        # Create convex hull(s) from the set of inlier points
+        if options['simplify_hulls']:
+            inlier_hull = ConvexHull(inlier_points)  #TODO: simplification not implemented in scipy ConvexHull
+            # inlier_hull = Delaunay(inlier_points)
+        else:
+            inlier_hull = ConvexHull(inlier_points)
+
+    hull_indices = np.unique(inlier_hull.simplices)
+    hull_points = np.stack([omask[i][inlier_hull.vertices] for i in range(len(omask))], axis=-1)
+    omat = [hull_points[:, i:(i+1)] - hull_points[:, i:(i+1)].T for i in range(hull_points.shape[-1])]
+    omat = np.stack([omat[i] for i in options['dims']], axis=-1)
+    dists = np.sqrt(np.sum(omat**2, axis=-1))
+    beamwidth = np.max(dists)
+    d_dims = np.sqrt(np.sum(ogrid[..., options['dims']]**2, axis=-1))
+    fit_mask = d_dims <= (beamwidth/2)
+
+    res = BeamwidthResult(
+        dims=options['dims'],
+        beamwidth=beamwidth,
+        units=options['units']
+    )
+    if options['masks']:
+        res.inlier_mask = inlier_mask
+        res.fit_mask = fit_mask
+    if options['points'] or options['hulls']:
+        res.inlier_points = inlier_points
+        res.fit_points = [ngrid0[ii][fit_mask] for ii in range(len(ngrid0))]
+    if options['hulls']:
+        res.inlier_hull = inlier_hull
+        res.fit_hull = ConvexHull(np.stack(res.fit_points, axis=-1))  #TODO: simplification not implemented in scipy ConvexHull
+
+    return res

src/openlifu/bf/mask_focus.py

@@ -0,0 +1,74 @@
+from enum import Enum
+from math import inf
+from typing import Tuple
+
+import numpy as np
+from xarray import Dataset
+
+from openlifu.bf.offset_grid import offset_grid
+from openlifu.geo import Point
+from openlifu.util.units import get_ndgrid_from_arr, rescale_coords
+
+MaskOp = Enum("MaskOp", ["GREATER", "GREATER_EQUAL", "LESS", "LESS_EQUAL"])
+
+
+def mask_focus(
+        data_arr: Dataset,
+        focus: Point,
+        distance: float,
+        mask_op: MaskOp = MaskOp.LESS_EQUAL,
+        units: str = "m",
+        aspect_ratio: Tuple[float, float, float] = (1., 1., 10.),
+        z_min: float = -inf
+        ) -> np.ndarray:
+    """
+    Creates a mask for points within a (scaled) distance from the focus point.
+
+    Args:
+        data_arr : xarray.Dataset
+        focus : fus.Point object
+            The focus point to be used as reference for masking.
+        distance : float
+            The maximum allowed distance in units defined by 'units' option.
+        units : str
+            Distance units (Default: "m").
+        mask_op : str
+            Operation to perform on the mask (Default: "LESS_EQUAL").
+        aspect_ratio : Tuple[float, float, float]
+            Aspect ratio to calculate distance (Default: (1, 1, 10)).
+        z_min : float
+            Minimum z-value for masking points below it (Default: -inf).
+
+    Returns:
+        A boolean array indicating which points are within the specified range.
+    """
+    # Rescale coordinates and focus to the specified units
+    data_arr_rescaled = rescale_coords(data_arr, units)
+    focus.rescale(units)
+
+    # Calculate distances from the focus for each point in coords and compare with distance limit
+    # The distance is calculated by first transforming the coordinate system so that the focus point is on
+    # the z' axis, adjusting the x and y axes to be orthogonal to the z' axis, and then calculating the distance
+    # e.g. d = sqrt(((x'-x0')/ax)^2 + ((y'-y0')/ay)^2 + ((z'-z0')/az)^2). This is useful for calculating how far
+    # away from an oblong focal spot each point is.
+    ogrid = offset_grid(data_arr_rescaled, focus)
+    ogrid_aspect_corrected = ogrid*aspect_ratio
+    m = np.sqrt(np.sum(ogrid_aspect_corrected**2, axis=-1))
+
+    # mask based on distance
+    if mask_op is MaskOp.GREATER:
+        mask = np.greater(m, distance)
+    elif mask_op is MaskOp.GREATER_EQUAL:
+        mask = np.greater_equal(m, distance)
+    elif mask_op is MaskOp.LESS:
+        mask = np.less(m, distance)
+    elif mask_op is MaskOp.LESS_EQUAL:
+        mask = np.less_equal(m, distance)
+    else:
+        raise ValueError(f"Mask operation {mask_op} is not defined!")
+    if z_min > -inf:
+        XYZ = get_ndgrid_from_arr(data_arr_rescaled)
+        zmask = XYZ[..., 2] > z_min
+        mask &= zmask
+
+    return mask

src/openlifu/bf/offset_grid.py

@@ -0,0 +1,30 @@
+import numpy as np
+from xarray import Dataset
+
+from openlifu.geo import Point
+from openlifu.util.units import get_ndgrid_from_arr
+
+
+def offset_grid(data_arr: Dataset, focus: Point):
+    """
+    Calculates the distance from the focus point for each point in the coordinate system.
+
+    Distances are returned from a coordinate system rotated in azimuth,
+    then elevation, so that the 'z' axis points at the focus.
+
+    Args:
+        data_arr: xarray.Dataset
+        focus : fus.Point object
+            The focus point to be used as reference.
+
+    Returns:
+        A list of distance arrays offsets 'dx', 'dy', 'dz' for each point in the coordinate system.
+    """
+    m = focus.get_matrix()
+
+    # Create a grid of homogeneous points in the coordinate system
+    xyz = get_ndgrid_from_arr(data_arr)
+    XYZ = np.append(xyz, np.ones((*xyz.shape[:3], 1)), axis=-1)
+    ogrid = XYZ @ np.linalg.inv(m).T[np.newaxis, np.newaxis, ...]
+
+    return ogrid[..., :3]

src/openlifu/geo.py

@@ -41,14 +41,16 @@ def get_matrix(self, origin: np.ndarray = np.eye(4), center_on_point: bool = Fal
             center = pos
         else:
             center = np.zeros(3)
-        zvec = pos / np.linalg.norm(pos)
+        zvec = np.array([0., 0., 1.])
+        if np.linalg.norm(pos) != 0:
+            zvec = pos / np.linalg.norm(pos)
         az = -np.arctan2(zvec[0], zvec[2])
         xvec = np.array([np.cos(az), 0.0, np.sin(az)])
         yvec = np.cross(zvec, xvec)
         m = np.array([[xvec[0], yvec[0], zvec[0], center[0]],
-                        [xvec[1], yvec[1], zvec[1], center[1]],
-                        [xvec[2], yvec[2], zvec[2], center[2]],
-                        [0.0, 0.0, 0.0, 1.0]])
+                      [xvec[1], yvec[1], zvec[1], center[1]],
+                      [xvec[2], yvec[2], zvec[2], center[2]],
+                      [0.0, 0.0, 0.0, 1.0]])
         if not local:
             m = np.dot(origin, m)
         return m