Adding missing unit tests for solution analysis

Removed `calc_dist_from_focus`, the algorithm is now in `mask_focus`

pre-commit hook

Author: ltetrel

src/openlifu/__init__.py

@@ -13,11 +13,9 @@
     Pulse,
     Sequence,
     apod_methods,
-    calc_dist_from_focus,
     delay_methods,
     focal_patterns,
     get_beamwidth,
-    get_focus_matrix,
     mask_focus,
     offset_grid,
 )
@@ -64,9 +62,7 @@
     "focal_patterns",
     "delay_methods",
     "apod_methods",
-    "calc_dist_from_focus",
     "get_beamwidth",
-    "get_focus_matrix",
     "mask_focus",
     "offset_grid",
     "SimSetup",

src/openlifu/bf/__init__.py

@@ -1,5 +1,4 @@
 from .apod_methods import ApodizationMethod
-from .calc_dist_from_focus import calc_dist_from_focus
 from .delay_methods import DelayMethod
 from .focal_patterns import FocalPattern, SinglePoint, Wheel
 from .get_beamwidth import get_beamwidth
@@ -17,7 +16,6 @@
     "Pulse",
     "Sequence",
     "offset_grid",
-    "calc_dist_from_focus",
     "mask_focus",
     "get_beamwidth"
 ]

src/openlifu/bf/calc_dist_from_focus.py

@@ -5,7 +5,7 @@
 import numpy as np
 from xarray import Dataset
 
-from openlifu.bf import calc_dist_from_focus
+from openlifu.bf.offset_grid import offset_grid
 from openlifu.geo import Point
 from openlifu.util.units import get_ndgrid_from_arr, rescale_coords
 
@@ -47,7 +47,13 @@ def mask_focus(
     focus.rescale(units)
 
     # Calculate distances from the focus for each point in coords and compare with distance limit
-    m = calc_dist_from_focus(data_arr_rescaled, focus, aspect_ratio)
+    # 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:

src/openlifu/bf/mask_focus.py

@@ -0,0 +1,50 @@
+import numpy as np
+import pytest
+from xarray import DataArray, Dataset
+
+from openlifu.bf import mask_focus
+from openlifu.geo import Point
+
+
+@pytest.fixture()
+def example_focus() -> DataArray:
+    return Point(
+        id="test_focus_point",
+        radius=1.,
+        position=np.array([0.5, 0.5, 0.5]),
+        dims=("x", "y", "z"),
+        units="mm"
+    )
+
+@pytest.fixture()
+def example_xarr() -> DataArray:
+    rng = np.random.default_rng(147)
+    return Dataset(
+            {
+                'p': DataArray(
+                    data=rng.random((3, 2, 3)),
+                    dims=["x", "y", "z"],
+                    attrs={'units': "Pa"}
+                )
+            },
+            coords={
+                'x': DataArray(dims=["x"], data=np.linspace(0, 1, 3), attrs={'units': "mm"}),
+                'y': DataArray(dims=["y"], data=np.linspace(0, 1, 2), attrs={'units': "mm"}),
+                'z': DataArray(dims=["z"], data=np.linspace(0, 1, 3), attrs={'units': "mm"})
+            }
+        )
+
+def test_mask_focus(example_xarr: Dataset, example_focus: Point):
+    """Test that the distance grid from the focus point is correct."""
+    expected = np.array(
+        [[[ True, True, True],
+          [ True, True, False]],
+
+         [[ True, True, True],
+          [True, False, False]],
+
+         [[ True, True, False],
+          [ False, False, False]]])
+    mask = mask_focus(example_xarr, example_focus, distance=0.01)
+
+    np.testing.assert_equal(mask, expected)

tests/test_mask_focus.py

@@ -0,0 +1,67 @@
+import numpy as np
+import pytest
+from xarray import DataArray, Dataset
+
+from openlifu.bf import offset_grid
+from openlifu.geo import Point
+
+
+@pytest.fixture()
+def example_focus() -> DataArray:
+    return Point(
+        id="test_focus_point",
+        radius=1.,
+        position=np.array([0.5, 0.5, 0.5]),
+        dims=("x", "y", "z"),
+        units="mm"
+    )
+
+@pytest.fixture()
+def example_xarr() -> DataArray:
+    rng = np.random.default_rng(147)
+    return Dataset(
+            {
+                'p': DataArray(
+                    data=rng.random((3, 2, 3)),
+                    dims=["x", "y", "z"],
+                    attrs={'units': "Pa"}
+                )
+            },
+            coords={
+                'x': DataArray(dims=["x"], data=np.linspace(0, 1, 3), attrs={'units': "mm"}),
+                'y': DataArray(dims=["y"], data=np.linspace(0, 1, 2), attrs={'units': "mm"}),
+                'z': DataArray(dims=["z"], data=np.linspace(0, 1, 3), attrs={'units': "mm"})
+            }
+        )
+
+def test_offset_grid(example_xarr: Dataset, example_focus: Point):
+    """Test that the distance grid from the focus point is correct."""
+    expected = np.array([
+       [[[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00],
+         [-3.53553391e-01, -2.04124145e-01,  2.88675135e-01],
+         [-7.07106781e-01, -4.08248290e-01,  5.77350269e-01]],
+
+        [[ 0.00000000e+00,  8.16496581e-01,  5.77350269e-01],
+         [-3.53553391e-01,  6.12372436e-01,  8.66025404e-01],
+         [-7.07106781e-01,  4.08248290e-01,  1.15470054e+00]]],
+
+
+       [[[ 3.53553391e-01, -2.04124145e-01,  2.88675135e-01],
+         [ 0.00000000e+00, -4.08248290e-01,  5.77350269e-01],
+         [-3.53553391e-01, -6.12372436e-01,  8.66025404e-01]],
+
+        [[ 3.53553391e-01,  6.12372436e-01,  8.66025404e-01],
+         [ 0.00000000e+00,  4.08248290e-01,  1.15470054e+00],
+         [-3.53553391e-01,  2.04124145e-01,  1.44337567e+00]]],
+
+
+       [[[ 7.07106781e-01, -4.08248290e-01,  5.77350269e-01],
+         [ 3.53553391e-01, -6.12372436e-01,  8.66025404e-01],
+         [ 0.00000000e+00, -8.16496581e-01,  1.15470054e+00]],
+
+        [[ 7.07106781e-01,  4.08248290e-01,  1.15470054e+00],
+         [ 3.53553391e-01,  2.04124145e-01,  1.44337567e+00],
+         [ 0.00000000e+00,  0.00000000e+00,  1.73205081e+00]]]])
+    offset = offset_grid(example_xarr, example_focus)
+
+    np.testing.assert_almost_equal(offset, expected)

tests/test_offset_grid.py

@@ -1,7 +1,36 @@
 import numpy as np
 import pytest
+from xarray import DataArray, Dataset
 
-from openlifu.util.units import getsiscale
+from openlifu.util.units import (
+    get_ndgrid_from_arr,
+    getsiscale,
+    rescale_coords,
+    rescale_data_arr,
+)
+
+
+@pytest.fixture()
+def example_xarr() -> Dataset:
+    rng = np.random.default_rng(147)
+    return Dataset(
+            {
+                'p': DataArray(
+                    data=rng.random((3, 2)),
+                    dims=["x", "y"],
+                    attrs={'units': "Pa"}
+                ),
+                'it': DataArray(
+                    data=rng.random((3, 2)),
+                    dims=["x", "y"],
+                    attrs={'units': "W/cm^2"}
+                )
+            },
+            coords={
+                'x': DataArray(dims=["x"], data=np.linspace(0, 1, 3), attrs={'units': "m"}),
+                'y': DataArray(dims=["y"], data=np.linspace(0, 1, 2), attrs={'units': "m"})
+            }
+        )
 
 
 def test_getsiscale():
@@ -20,3 +49,34 @@ def test_getsiscale():
     assert getsiscale('MHz', 'frequency') == 1e6
     assert getsiscale('GHz', 'frequency') == 1e9
     assert getsiscale('THz', 'frequency') == 1e12
+
+
+def test_rescale_data_arr(example_xarr: Dataset):
+    """Test that an xarray data can be correctly rescaled."""
+    expected_p = 1e-6 * example_xarr['p'].data
+    expected_it = 1e4 * example_xarr['it'].data
+    rescaled_p = rescale_data_arr(example_xarr['p'], units="MPa")
+    rescaled_it = rescale_data_arr(example_xarr['it'], units="W/m^2")
+
+    np.testing.assert_almost_equal(rescaled_p, expected_p)
+    np.testing.assert_almost_equal(rescaled_it, expected_it)
+
+
+def test_rescale_coords(example_xarr: Dataset):
+    """Test that an xarray coords can be correctly rescaled."""
+    expected_x = 1e3 * example_xarr['p'].coords['x'].data
+    expected_y = 1e3 * example_xarr['p'].coords['y'].data
+    rescaled = rescale_coords(example_xarr['p'], units="mm")
+
+    np.testing.assert_almost_equal(rescaled['x'].data, expected_x)
+    np.testing.assert_almost_equal(rescaled['y'].data, expected_y)
+
+
+def test_get_ndgrid_from_arr(example_xarr: Dataset):
+    """Test that an ndgrid can be constructed from the coordinates of an xarray."""
+    expected_X = np.array([[0., 0.], [0.5, 0.5], [1., 1.]])
+    expected_Y = np.array([[0., 1.], [0., 1.], [0., 1.]])
+    expected = np.stack([expected_X, expected_Y], axis=-1)
+    ndgrid = get_ndgrid_from_arr(example_xarr)
+
+    np.testing.assert_equal(ndgrid, expected)