Change Solution.focus to a list of foci (#125)

This includes
* changing Solution.focus to Solution.foci and making that be a *list*
  of Points rather than a Point.
* changing the xarray simulation data to have an additional coord
  "focal_point_index"

Author: ebrahimebrahim

src/openlifu/plan/solution.py

@@ -3,7 +3,7 @@
 from dataclasses import asdict, dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
 import xarray
@@ -51,8 +51,11 @@ class Solution:
     """Pulse to send to the transducer when running sonication"""
     sequence: Sequence = field(default_factory=Sequence)
     """Pulse sequence to use when running sonication"""
-    focus: Optional[Point] = None
-    """Point that is being focused on in this Solution; part of the focal pattern of the target"""
+    foci: List[Point] = field(default_factory=list)
+    """Points that are focused on in this Solution due to the focal pattern around the target.
+    Each item in this list is a unique point from the focal pattern, and the pulse sequence is
+    what determines how many times each point will be used.
+    """
 
     # there was "target_id" in the matlab software, but here we do not have the concept of a target ID.
     # I believe this was only needed in the matlab software because solutions were organized by target rather
@@ -109,8 +112,10 @@ def from_json(json_string : str, simulation_result: Optional[xarray.Dataset]=Non
             solution_dict["apodizations"] = np.array(solution_dict["apodizations"])
         solution_dict["pulse"] = Pulse.from_dict(solution_dict["pulse"])
         solution_dict["sequence"] = Sequence.from_dict(solution_dict["sequence"])
-        if solution_dict["focus"] is not None:
-            solution_dict["focus"] = Point.from_dict(solution_dict["focus"])
+        solution_dict["foci"] = [
+            Point.from_dict(focus_dict)
+            for focus_dict in solution_dict["foci"]
+        ]
         if solution_dict["target"] is not None:
             solution_dict["target"] = Point.from_dict(solution_dict["target"])
 

tests/resources/example_db/subjects/example_subject/sessions/example_session/solutions/example_solution/example_solution.json

@@ -44,15 +44,17 @@
     "pulse_train_interval": 1,
     "pulse_train_count": 1
   },
-  "focus": {
-    "id": "example_target",
-    "name": "Example Target",
-    "color": [1.0, 0.0, 0.0],
-    "radius": 0.001,
-    "position": [0.0, -0.0022437460888595447, 0.05518120697745499],
-    "dims": ["lat", "ele", "ax"],
-    "units": "m"
-  },
+  "foci": [
+    {
+      "id": "example_target",
+      "name": "Example Target",
+      "color": [1.0, 0.0, 0.0],
+      "radius": 0.001,
+      "position": [0.0, -0.0022437460888595447, 0.05518120697745499],
+      "dims": ["lat", "ele", "ax"],
+      "units": "m"
+    }
+  ],
   "target": {
     "id": "example_target",
     "name": "Example Target",

tests/resources/example_db/subjects/example_subject/sessions/example_session/solutions/example_solution/example_solution.nc

@@ -178,6 +178,7 @@ def test_load_solution(example_database:Database, example_session:Session):
     example_solution = example_database.load_solution(example_session, "example_solution")
     assert example_solution.name == "Example Solution"
     assert "p_min" in example_solution.simulation_result.data_vars # ensure the xarray dataset got loaded too
+    assert len(example_solution.simulation_result['focal_point_index']) == len(example_solution.foci) # ensure simulation data was loaded for all foci
 
 def test_write_solution(example_database:Database, example_session:Session):
     solution = Solution(name="bleh", id='new_solution')

tests/test_database.py

@@ -23,11 +23,11 @@ def example_solution() -> Solution:
         apodizations=np.array([0.5, 0.75, 1.0, 0.85]),
         pulse=Pulse(frequency=42),
         sequence=Sequence(pulse_count=27),
-        focus=Point(id = "test_focus_point"),
+        foci=[Point(id = "test_focus_point")],
         target=Point(id = "test_target_point"),
         simulation_result=xarray.Dataset(
-            {'pnp': (['x', 'y', 'z'], rng.random((3, 2, 3)))},
-            coords={'x': np.linspace(0, 1, 3), 'y': np.linspace(0, 1, 2), 'z': np.linspace(0, 1, 3)},
+            {'p_min': (['focal_point_index', 'x', 'y', 'z'], rng.random((1, 3, 2, 3)))},
+            coords={'x': np.linspace(0, 1, 3), 'y': np.linspace(0, 1, 2), 'z': np.linspace(0, 1, 3), 'focal_point_index': [0]},
         ),
     )
 
@@ -67,3 +67,8 @@ def test_save_load_solution_custom_dataset_filepath(example_solution:Solution, t
     nc_filepath = tmp_path/"some_other_directory"/"sim_output.nc"
     example_solution.to_files(json_filepath, nc_filepath)
     assert dataclasses_are_equal(Solution.from_files(json_filepath, nc_filepath), example_solution)
+
+def test_num_foci(example_solution:Solution):
+    """Ensure that the number of foci in the test solution matches the number of foci provided in its simulation data.
+    (This is more checking correctness of the test example rather than correctness of code, but it is important.)"""
+    assert len(example_solution.simulation_result['focal_point_index']) == len(example_solution.foci)