From f8901da067cbc800db54d2e0e2293ad67a9a881f Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Mon, 4 Nov 2024 18:03:24 +0100
Subject: [PATCH 01/29] Port over solution creation and scaling code (#126)

This adds the `calc_solution` member from `Protocol` which checks the validity of session targets, setup the simulation scene, beamform the signal, run the simulation, analyze the solution and rescale it.

This currently does not perform any segmentation but create a default water volume (with unform isotropic medium), there is an initial implementation for MRI tissue segmentation that will be re-worked in the future.
---
 pyproject.toml                                |   1 +
 src/openlifu/plan/__init__.py                 |   2 +
 src/openlifu/plan/protocol.py                 | 255 +++++++++++++++++-
 src/openlifu/plan/solution.py                 |  13 +-
 src/openlifu/plan/target_constraints.py       |  37 +++
 src/openlifu/seg/seg_methods/segment_mri.py   |  17 +-
 src/openlifu/sim/sim_setup.py                 |  62 ++++-
 src/openlifu/util/geom.py                     |  61 +++++
 src/openlifu/util/units.py                    |  39 ++-
 src/openlifu/xdc/transducer.py                |   2 -
 .../example_protocol/example_protocol.json    |  18 +-
 tests/test_protocol.py                        |  75 +++++-
 tests/test_seg_method.py                      |   1 +
 13 files changed, 558 insertions(+), 25 deletions(-)
 create mode 100644 src/openlifu/plan/target_constraints.py
 create mode 100644 src/openlifu/util/geom.py

diff --git a/pyproject.toml b/pyproject.toml
index a9b98ea1..cd42e675 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,6 +40,7 @@ dependencies = [
   "opencv-contrib-python",
   "crc",
   "nibabel",
+  "antspyx",
   "sphinx",
   "ipykernel",
   "k-wave-python>=0.3.4",
diff --git a/src/openlifu/plan/__init__.py b/src/openlifu/plan/__init__.py
index 22450b6f..b471f588 100644
--- a/src/openlifu/plan/__init__.py
+++ b/src/openlifu/plan/__init__.py
@@ -1,9 +1,11 @@
 from .protocol import Protocol
 from .run import Run
 from .solution import Solution
+from .target_constraints import TargetConstraints
 
 __all__ = [
     "Protocol",
     "Solution",
     "Run",
+    "TargetConstraints"
 ]
diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index b9969fab..60a4ed34 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -1,11 +1,25 @@
 import json
+import logging
+import math
+from copy import deepcopy
 from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from enum import Enum
 from pathlib import Path
-from typing import Any, Dict
+from typing import Any, Dict, List, Tuple
 
+import numpy as np
 import xarray as xa
 
 from openlifu import bf, geo, seg, sim, xdc
+from openlifu.db.session import Session
+from openlifu.geo import Point
+from openlifu.plan.solution import Solution, SolutionAnalysis, SolutionAnalysisOptions
+from openlifu.plan.target_constraints import TargetConstraints
+from openlifu.sim import run_simulation
+from openlifu.xdc import Transducer
+
+OnPulseMismatchAction = Enum("OnPulseMismatchAction", ["ERROR", "ROUND", "ROUNDUP", "ROUNDDOWN"])
 
 
 @dataclass
@@ -20,9 +34,9 @@ class Protocol:
     delay_method: bf.DelayMethod = field(default_factory=bf.delay_methods.Direct)
     apod_method: bf.ApodizationMethod = field(default_factory=bf.apod_methods.Uniform)
     seg_method: seg.SegmentationMethod = field(default_factory=seg.seg_methods.Water)
-    param_constraints: dict = field(default_factory=dict)
-    target_constraints: dict = field(default_factory=dict)
-    analysis_options: dict = field(default_factory=dict)
+    param_constraints: dict = field(default_factory=dict)  #TODO: this seems to be used only in `plan.check_analysis`` but not called anywhere
+    target_constraints: List[TargetConstraints] = field(default_factory=list)
+    analysis_options: SolutionAnalysisOptions = field(default_factory=SolutionAnalysisOptions)
 
     @staticmethod
     def from_dict(d : Dict[str,Any]) -> "Protocol":
@@ -36,6 +50,11 @@ def from_dict(d : Dict[str,Any]) -> "Protocol":
         if "materials" in d:
             seg_method_dict["materials"] = seg.Material.from_dict(d.pop("materials"))
         d["seg_method"] = seg.SegmentationMethod.from_dict(seg_method_dict)
+        d['param_constraints'] = d.get("param_constraints", {})
+        if "target_constraints" in d:
+            d['target_constraints'] = [TargetConstraints.from_dict(d_tc) for d_tc in d.get("target_constraints", {})]
+        if "analysis_options" in d:
+            d['analysis_options'] = SolutionAnalysisOptions.from_dict(d.get("analysis_options"))
         return Protocol(**d)
 
     def to_dict(self):
@@ -84,13 +103,237 @@ def to_json(self, compact:bool) -> str:
         else:
             return json.dumps(self.to_dict(), indent=4)
 
-    def to_file(self, filename):
+    def to_file(self, filename: str):
         """
         Save the protocol to a file
 
-        :param filename: Name of the file
+        Args:
+            filename: Name of the file
         """
         Path(filename).parent.parent.mkdir(exist_ok=True)
         Path(filename).parent.mkdir(exist_ok=True)
         with open(filename, 'w') as file:
             file.write(self.to_json(compact=False))
+
+
+    def check_target(self, target):
+        """
+        Check if a target is within bounds.
+
+        Args:
+            target: The geo.Point target to check.
+        """
+        #TODO: in the matlab code they handle the case were multiple targets are given.
+        # After our discussion I assumed that we will not handle a list of targets.
+        if isinstance(target, list):
+            raise ValueError(f"Input target {target} not supposed to be a list!")
+
+        # check if target position is within target_constraints defined bounds.
+        for target_constraint in self.target_constraints:
+            pos = target.get_position(
+                dim=target_constraint.dim,
+                units=target_constraint.units
+            )
+            target_constraint.check_bounds(pos)
+
+    def scale_solution(
+            self,
+            solution: Solution,
+            transducer: Transducer,
+            analysis_options: SolutionAnalysisOptions = SolutionAnalysisOptions()
+    ) -> Tuple[Solution, SolutionAnalysis]:
+        """
+        Scale the solution to match the target pressure.
+        If no output is requested, the solution is scaled in-place.
+
+        Args:
+            solution: plan.Solution
+                The solution to be scaled.
+            analysis_options: plan.solution.SolutionAnalysisOptions
+
+        Returns:
+            solution_scaled: the scaled plan.Solution
+            analysis_scaled: the resulting plan.solution.SolutionAnalysis from scaled solution
+        """
+        solution_scaled = deepcopy(solution)
+        analysis = solution.analyze(transducer, options=analysis_options)
+
+        scaling_factors = np.zeros(solution.num_foci())
+        for i in range(solution.num_foci()):
+            scaling_factors[i] = (self.focal_pattern.target_pressure / 1e6) / analysis.mainlobe_pnp_MPa[i]
+            scaling_factors[i] = 2.3167
+        max_scaling = np.max(scaling_factors)
+        v0 = self.pulse.amplitude
+        v1 = v0 * max_scaling
+        apod_factors = scaling_factors / max_scaling
+
+        for i in range(solution.num_foci()):
+            scaling = v1/v0*apod_factors[i]
+            solution_scaled.simulation_result['p_min'][i].data *= scaling
+            solution_scaled.simulation_result['p_max'][i].data *= scaling
+            solution_scaled.simulation_result['ita'][i].data *= scaling**2
+            solution_scaled.apodizations[i] = solution_scaled.apodizations[i]*apod_factors[i]
+        self.pulse.amplitude = v1
+
+        analysis_scaled = solution_scaled.analyze(transducer, options=analysis_options)
+
+        return solution_scaled, analysis_scaled
+
+    #TODO: The arg transform needed for sim_setup since transducer.matrix does not exists
+    def calc_solution(
+        self,
+        target: Point,
+        transducer: Transducer,
+        transform: np.ndarray,
+        volume: xa.DataArray = None,  #TODO: Do we want to have the volume as a xa.DataArray instead of nifty ?
+        session: Session = None, # useful in solution id  #TODO not sure to understand why this type is optional
+        simulate: bool = True,
+        scale: bool = True,
+        sim_options: sim.SimSetup = None,
+        analysis_options: SolutionAnalysisOptions = None,
+        on_pulse_mismatch: OnPulseMismatchAction = OnPulseMismatchAction.ERROR,
+        #log : Logger. Default: fus.util.Logger.get()  #TODO what about logging, currently only db/database.py has one ?
+    ) -> Tuple[Solution, xa.DataArray]:  #TODO: make more sense for me to have a single xa.DataArray that holds the
+                                         # aggregation (pnp, ppp, ita). We could also store it in Solution.simulation_result
+                                         # with additional fields 'pnp_aggregated', 'ppp_aggregated' and 'ita_aggregated' ?
+        """Calculate the solution and aggregated k-wave simulation outputs.
+
+        Method that computes the delays and apodizations for each focus in
+        the treatment plan, simulates the resulting pressure field to adjust
+        transmit pressures to reach target pressures, and then analyzes the
+        resulting pressure field to compute the resulting acoustic parameters.
+
+        Args:
+            target: The target Point.
+            transducer: A Transducer item.
+            volume: xa.DataArray
+                The xa.DataArray volume corresponding to the subject scan (Default: None).
+                If no volume is given, a default simulation grid will be used.
+            session: db.Session
+                A session used to define solution_id (Default: None).
+            simulate: bool
+                Enable solution simulation (Default: true).
+            scale: bool
+                Triggers solution and simulation scaling to the requested pressure (Default: true).
+            sim_options : sim.SimSetup
+                The options for the k-wave simulation (Default: self.sim_setup).
+            analysis_options: plan.solution.SolutionAnalysisOptions
+                The options for the solution analysis (Default: self.analysis_options).
+            on_pulse_mismatch: plan.protocol.OnPulseMismatchAction
+                An action to take if the number of pulses in the sequence does not match
+                the number of foci (Default: OnPulseMismatchAction.ERROR).
+
+        Returns: Tuple[Solution, xa.DataArray]
+            Represents the solution object including its analysis, and aggregated simulation output.
+        """
+        if sim_options is None:
+            sim_options = self.sim_setup
+        if analysis_options is None:
+            analysis_options = self.analysis_options
+        # check before if target is within bounds
+        self.check_target(target)
+        params, transducer, target = sim_options.setup_sim_scene(transducer, transform, target, self.seg_method, volume=volume, units="m")
+
+        delays_to_stack: List[np.ndarray] = []
+        apodizations_to_stack: List[np.ndarray] = []
+        simulation_outputs_to_stack: List[xa.Dataset] = []
+        simulation_output_stacked: xa.Dataset = None
+        simulation_result_aggregated: xa.Dataset = None
+        foci: List[Point] = self.focal_pattern.get_targets(target)
+        simulation_cycles = np.max([np.round(self.pulse.duration * self.pulse.frequency), 20])
+        out_solution: Solution = Solution()
+
+        # updating solution sequence if pulse mismatch
+        solution_sequence = deepcopy(self.sequence)
+        if (self.sequence.pulse_count % len(foci)) != 0:
+            if on_pulse_mismatch is OnPulseMismatchAction.ERROR:
+                raise ValueError(f"Pulse Count {self.sequence.pulse_count} is not a multiple of the number of foci {len(foci)}")
+            else:
+                if on_pulse_mismatch is OnPulseMismatchAction.ROUND:
+                    solution_sequence.pulse_count = round(self.sequence.pulse_count / len(self.foci)) * len(foci)
+                elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDUP:
+                    solution_sequence.pulse_count = math.ceil(self.sequence.pulse_count / len(foci)) * len(self.foci)
+                elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDUP:
+                    solution_sequence.pulse_count = math.floor(self.sequence.pulse_count / len(foci)) * len(self.foci)
+                logging.warning(
+                    f"Pulse Count {self.sequence.pulse_count} is not a multiple of the number of foci {len(foci)}."
+                    f"Rounding to {solution_sequence.pulse_count}."
+                )
+        # run simulation and aggregate the results
+        for focus in foci:
+            logging.info(f"Beamform for focus {focus}...")
+            delays, apodization = self.beamform(arr=transducer, target=focus, params=params)
+            simulation_output_xarray = None
+            if simulate:
+                logging.info(f"Simulate for focus {focus}...")
+                simulation_output_xarray, _ = run_simulation(
+                    arr=transducer,
+                    params=params,
+                    delays=delays,
+                    apod= apodization,
+                    freq = self.pulse.frequency,
+                    cycles = simulation_cycles,
+                    dt=sim_options.dt,
+                    t_end=sim_options.t_end,
+                    amplitude = 1,
+                    gpu = False
+                )
+            delays_to_stack.append(delays)
+            apodizations_to_stack.append(apodization)
+            simulation_outputs_to_stack.append(simulation_output_xarray)
+        if simulate:
+            simulation_output_stacked = xa.concat(
+                [
+                    sim.assign_coords(focal_point_index=i)
+                    for i, sim in enumerate(simulation_outputs_to_stack)
+                ],
+                dim='focal_point_index',
+            )
+        # instantiate and return the solution
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
+        solution_id = timestamp
+        if session is not None:
+            solution_id = f"{session.id}_{solution_id}"
+        solution =  Solution(
+            id=solution_id,
+            name=f"Solution {timestamp}",
+            protocol_id=self.id,
+            transducer_id=transducer.id,
+            delays=np.stack(delays_to_stack, axis=0),
+            apodizations=np.stack(apodizations_to_stack, axis=0),
+            pulse=self.pulse,  #TODO pulse is correctly scaled with `self.scale_solution``
+            sequence=solution_sequence, #TODO incorrect to set the sequence the same as the protocol's
+                                        # since sequence pulse_count can be modified if pulse mismatch.
+            foci=foci,
+            target=target,
+            simulation_result=simulation_output_stacked,
+            approved=False,
+            description= (
+                f"A solution computed for the {self.name} protocol with transducer {transducer.name}"
+                f" for subject volume {volume.id if volume is not None else None}" # TODO put volume ID here if it is not None, once Sadhana's PR #123 is merged
+                f" for target {target.id}."
+                f" This solution was created for the session {session.id} for subject {session.subject_id}." if session is not None else ""
+            )
+        )
+        # optionally scale the solution with simulation result
+        if scale:
+            if not simulate:
+                logging.error(msg=f"Cannot scale solution {solution.id} if simulation is not enabled!")
+                raise ValueError(f"Cannot scale solution {solution.id} if simulation is not enabled!")
+            logging.info(f"Scaling solution {solution.id}...")
+            #TODO does analysis needs to be an attribute of solution ?
+            solution_scaled, _ = self.scale_solution(solution, transducer, analysis_options=analysis_options)
+            solution = solution_scaled
+
+        # Finally the resulting pressure is max-aggregated and intensity is mean-aggregated, over all focus points .
+        pnp_aggregated = solution.simulation_result['p_min'].max(dim="focal_point_index")
+        ppp_aggregated = solution.simulation_result['p_max'].max(dim="focal_point_index")
+        # TODO: Ensure this mean is weighted by the number of times each point is focused on, once openlifu supports hitting points different numbers of times
+        intensity_aggregated = solution.simulation_result['ita'].mean(dim="focal_point_index")
+        simulation_result_aggregated = deepcopy(solution.simulation_result)
+        simulation_result_aggregated = simulation_result_aggregated.drop_dims("focal_point_index")
+        simulation_result_aggregated['p_min'] = pnp_aggregated
+        simulation_result_aggregated['p_max'] = ppp_aggregated
+        simulation_result_aggregated['ita'] = intensity_aggregated
+
+        return solution, simulation_result_aggregated
diff --git a/src/openlifu/plan/solution.py b/src/openlifu/plan/solution.py
index 7560fbb2..f62c892d 100644
--- a/src/openlifu/plan/solution.py
+++ b/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 List, Optional
+from typing import List, Optional, Tuple
 
 import numpy as np
 import xarray as xa
@@ -11,6 +11,7 @@
 from openlifu.bf import Pulse, Sequence, mask_focus
 from openlifu.bf.mask_focus import MaskOp
 from openlifu.geo import Point
+from openlifu.io.dict_conversion import DictMixin
 from openlifu.util.json import PYFUSEncoder
 from openlifu.util.units import rescale_data_arr
 from openlifu.xdc import Transducer
@@ -24,7 +25,7 @@ def _construct_nc_filepath_from_json_filepath(json_filepath:Path) -> Path:
 
 
 @dataclass
-class SolutionAnalysis:
+class SolutionAnalysis(DictMixin):
     mainlobe_pnp_MPa: list[float] = field(default_factory=list)
     mainlobe_isppa_Wcm2: list[float] = field(default_factory=list)
     mainlobe_ispta_mWcm2: list[float] = field(default_factory=list)
@@ -44,13 +45,13 @@ class SolutionAnalysis:
 
 
 @dataclass
-class SolutionOptions:
+class SolutionAnalysisOptions(DictMixin):
     standoff_sound_speed: float = 1500.0
     standoff_density: float = 1000.0
     ref_sound_speed: float = 1500.0
     ref_density: float = 1000.0
     focus_diameter: float = 0.5
-    mainlobe_aspect_ratio: tuple[float, float, float] = (1., 1., 5.)
+    mainlobe_aspect_ratio: Tuple[float, float, float] = (1., 1., 5.)
     mainlobe_radius: float = 2.5e-3
     beamwidth_radius: float = 5e-3
     sidelobe_radius: float = 3e-3
@@ -120,11 +121,11 @@ def num_foci(self) -> int:
         """Get the number of foci"""
         return len(self.foci)
 
-    def analyze(self, transducer: Transducer, options: SolutionOptions = SolutionOptions()) -> SolutionAnalysis:
+    def analyze(self, transducer: Transducer, options: SolutionAnalysisOptions = SolutionAnalysisOptions()) -> SolutionAnalysis:
         """Analyzes the treatment solution.
 
         Args:
-            transducer: A Transducer item.  #TODO: this should be instantiated at the database level, not here ?
+            transducer: A Transducer item.
             options: A struct for solution analysis options.
 
         Returns: A struct containing the results of the analysis.
diff --git a/src/openlifu/plan/target_constraints.py b/src/openlifu/plan/target_constraints.py
new file mode 100644
index 00000000..53ba0f5d
--- /dev/null
+++ b/src/openlifu/plan/target_constraints.py
@@ -0,0 +1,37 @@
+import logging
+from dataclasses import dataclass
+
+from openlifu.io.dict_conversion import DictMixin
+
+
+@dataclass
+class TargetConstraints(DictMixin):
+    """ A class for storing target constraints.
+
+    Target constraints are used to define the acceptable range of
+    positions for a target. For example, a target constraint could
+    be used to define the acceptable range of values for the x position
+    of a target.
+    """
+
+    dim: str = "x"
+    """The dimension ID being constrained"""
+
+    name: str = "dim"
+    """The name of the dimension being constrained"""
+
+    units: str = "m"
+    """The units of the dimension being constrained"""
+
+    min: float = float("-inf")
+    """The minimum value of the dimension"""
+
+    max: float = float("inf")
+    """The maximum value of the dimension"""
+
+    def check_bounds(self, pos: float):
+        """Check if the given position is within bounds."""
+
+        if (pos < self.min) or (pos > self.max):
+            logging.error(msg=f"The position {pos} at dimension {self.name} is not within bounds [{self.min}, {self.max}]!")
+            raise ValueError(f"The position {pos} at dimension {self.name} is not within bounds [{self.min}, {self.max}]!")
diff --git a/src/openlifu/seg/seg_methods/segment_mri.py b/src/openlifu/seg/seg_methods/segment_mri.py
index f054fcc7..24e59893 100644
--- a/src/openlifu/seg/seg_methods/segment_mri.py
+++ b/src/openlifu/seg/seg_methods/segment_mri.py
@@ -1,11 +1,22 @@
 from dataclasses import dataclass
 
-import xarray as xa
+import ants
+from nibabel import Nifti1Image
 
 from openlifu.seg.seg_methods.seg_method import SegmentationMethod
 
 
 @dataclass
 class SegmentMRI(SegmentationMethod):
-    def _segment(self, volume: xa.DataArray):
-        raise NotImplementedError
+    def _segment(self, volume: Nifti1Image):
+        # On using ANTs antropos, see this discussion: https://neurostars.org/t/n-a-whitematter-confounds-in-tsv-output-file/4859/10
+        ants_img = ants.from_numpy(
+            volume.get_fdata(),
+            origin=volume.affine[:3, -1].tolist(),
+            spacing=volume.header['pixdim'][:3].tolist(),
+            direction=volume.affine[:3, :3]
+        )
+        mask = ants.get_mask(ants_img)
+        img_segs = ants.atropos(ants_img, x=mask, m="[0.2, 1x1x1]")
+
+        return img_segs
diff --git a/src/openlifu/sim/sim_setup.py b/src/openlifu/sim/sim_setup.py
index 2611f39e..a2cd7afb 100644
--- a/src/openlifu/sim/sim_setup.py
+++ b/src/openlifu/sim/sim_setup.py
@@ -5,8 +5,10 @@
 import numpy as np
 import xarray as xa
 
+from openlifu.geo import Point
 from openlifu.io.dict_conversion import DictMixin
-from openlifu.util.units import getunitconversion
+from openlifu.seg import SegmentationMethod
+from openlifu.util.units import getunitconversion, rescale_coords
 from openlifu.xdc import Transducer
 
 
@@ -107,5 +109,59 @@ def get_spacing(self, units: Optional[str] = None):
         units = self.units if units is None else units
         return getunitconversion(self.units, units)*self.spacing
 
-    def transform_scene(self, scene, id: Optional[str] = None, name: Optional[str] = None, units: Optional[str] = None):
-        raise NotImplementedError
+    #TODO: since we don't have a concept of scene here, and that the simulation scene is needed in protocol.calc_solution,
+    # we will return each prepared object one-by-one.
+    #TODO: Missing the "markers" from matlab scene in "scene.transform(self, coords, matrix, options)"
+    #TODO: The arg transform needs to be added since transducer.matrix does not exists anymore
+    def setup_sim_scene(
+            self,
+            transducer: Transducer,
+            transform: np.ndarray,
+            target: Point,
+            seg_method: SegmentationMethod,
+            volume: xa.DataArray = None,
+            interp_method: str = "Linear",
+            units: Optional[str] = None
+        ) -> Tuple[xa.DataArray, Transducer, Point]:
+        """ Prepare a simulation scene composed of a volume, transducer and targets.
+
+        Setup a simulation scene with a volume, transducer and target point.
+        All objects are resampled to the geo-referenced simulation grid (lon, lat, ele).
+        For the volume, a segmentation is also performed that defines the simulation medium.
+
+        Args:
+            volume: xa.DataArray
+            transducer: xdc.Transducer
+            transform: np.ndarray
+            target: geo.Point
+            seg_method: seg.SegmentationMethod
+            interp_method: str
+                Interpolation method for the simulation grid (Default: \"Linear\").
+            units: str
+                Units of simulation grid (Default: self.units).
+
+        Returns
+            params: The resampled and segmented xa.DataArray volume to the simulation grid
+            transducer_tr: The resampled xdc.Transducer to the simulation grid
+            target_tr: The resampled geo.Point to the simulation grid
+        """
+        units = self.units if units is None else units
+        sim_coords = self.get_coords(units=units)
+        #TODO: in a near future, the following will transform transducer and target. Currently Slicer does it.
+        # sim_matrix = convert_transform(transform, units=transducer.units, tgt_units=units)  #TODO: not clear here where the transform should come from since transducer.get_matrix() does not exist
+        # rescaled_affine = convert_transform(volume.affine, units=volume.affine_units, tgt_units=units)
+        # volume = volume.assign_attrs(affine_units=units, affine=rescaled_affine)
+        # volume_resampled = resample_volume(volume, sim_coords, sim_matrix, interp_method)
+        # transducer_tr = transducer.transform(sim_matrix, units)
+        # target_tr = target.transform(sim_matrix, units=units, new_dims=sim_coords.dims)
+        transducer.rescale(units)
+        target.rescale(units)
+        if volume is None:
+            params = seg_method.ref_params(sim_coords)
+        else:
+            sim_coords_units = sim_coords[next(iter(sim_coords.keys()))].units
+            volume_resampled = rescale_coords(volume, sim_coords_units)
+            params = seg_method.seg_params(volume_resampled)  #TODO: currently only "water" method is tested with db-simple-example-v04
+                                                              # we should have a whole MRI segmentation instead.
+
+        return params, transducer, target
diff --git a/src/openlifu/util/geom.py b/src/openlifu/util/geom.py
new file mode 100644
index 00000000..b0364283
--- /dev/null
+++ b/src/openlifu/util/geom.py
@@ -0,0 +1,61 @@
+import numpy as np
+import xarray as xa
+
+from openlifu.util.units import get_ndgrid_from_coords, rescale_coords
+
+
+def resample_volume(
+        volume: xa.DataArray,
+        coords: xa.Coordinates,
+        matrix: np.ndarray,
+        interp_method: str = "linear"
+    ) -> xa.DataArray:
+    """Resample an xa.DataArray Volume data
+
+    Resample an xa.DataArray Volume to the coordinate system
+    specified by coords using the transformation matrix.
+
+    Args:
+        volume: xa.DataArray
+        coords:
+            The coordinate system to which the volume should be resampled
+        matrix: np.ndarray
+            The transformation matrix in the coords units
+
+    Returns:
+        fus.Volume: The transformed Volume object
+    """
+
+    prev_units = volume[next(iter(volume.coords.keys()))].units
+    coords_units = coords[next(iter(coords.keys()))].units
+    volume_rescaled = rescale_coords(volume, coords_units)
+    xyz = get_ndgrid_from_coords(coords)
+    XYZ = np.append(xyz, np.ones((*xyz.shape[:3], 1)), axis=-1)
+    X1 = np.linalg.inv(matrix) @ (matrix @ XYZ)
+    interp_method = None
+    #TODO: get transform matrix to transform grid
+    # inv_matrix = (self.matrix'*self.matrix)\(self.matrix');
+    pass
+    # XP = np.vstack((np.ravel(Xp).T, np.ones(np.prod(Xp.shape), dtype))
+    # XP[:, -1] = 1.0)  # Add a column for homogeneous coordinates
+
+    # inv_matrix = np.linalg.inv(self.matrix.T @ self.matrix)
+
+    # X1 = matrix.dot(XP)
+    # X1 = (inv_matrix @ X1).reshape(Xp.shape[:-1], -1, order='F')
+
+    # if options is None:
+    #     method = 'linear'
+    # else:
+    #     method = options['method']
+
+    # data = self.interp((X1.T,), method=method)  # Interpolate the transformed coordinates
+    # obj = self.newobj(data, coords, id=self.id,
+    #                   name=self.name, matrix=matrix,
+    #                   attrs=self.attrs, units=self.units)
+
+    # if len(self):
+    #     self.rescale(prev_units)
+    rescale_coords(volume, prev_units)
+
+    return volume
diff --git a/src/openlifu/util/units.py b/src/openlifu/util/units.py
index d03c3b34..bd5c16cb 100644
--- a/src/openlifu/util/units.py
+++ b/src/openlifu/util/units.py
@@ -1,5 +1,5 @@
 import numpy as np
-from xarray import Dataset
+from xarray import Coordinates, Dataset
 
 #TODO: use Pint (https://github.com/hgrecco/pint) instead to manage physics units in python
 
@@ -241,3 +241,40 @@ def get_ndgrid_from_arr(data_arr: Dataset) -> np.ndarray:
     ndgrid = np.stack(np.meshgrid(*all_coord, indexing="ij"), axis=-1)
 
     return ndgrid
+
+
+def get_ndgrid_from_coords(coords_arr: Coordinates) -> np.ndarray:
+    """
+    Creates a ndgrid from xarray.coordinates.
+
+    Args:
+        coords : xarray.Coordinates
+
+    Returns:
+        ndgrid: The ndgrid from the Coordinates.
+    """
+    all_coord = []
+    for coord_key in coords_arr.dims:
+        if 'units' in coords_arr[coord_key].attrs:
+            all_coord += [coords_arr[coord_key].data]
+    ndgrid = np.stack(np.meshgrid(*all_coord, indexing="ij"), axis=-1)
+
+    return ndgrid
+
+
+def convert_transform(matrix: np.ndarray, units: str, tgt_units: str) -> np.ndarray:
+    """Given a transform matrix in some units, convert it to another units.
+
+    Args:
+        matrix: 4x4 np.ndarray
+            affine transform matrix
+        units: str
+            units of the coordinate space on which the provided transform matrix operates
+
+    Returns: 4x4 affine transform matrix
+    """
+    tgt_units = self.units if tgt_units is None else tgt_units
+    matrix = matrix.copy()
+    matrix[0:3, 3] *= getunitconversion(units, tgt_units)
+
+    return matrix
diff --git a/src/openlifu/xdc/transducer.py b/src/openlifu/xdc/transducer.py
index d2dcd957..77c498ac 100644
--- a/src/openlifu/xdc/transducer.py
+++ b/src/openlifu/xdc/transducer.py
@@ -153,8 +153,6 @@ def rescale(self, units):
         if self.units != units:
             for element in self.elements:
                 element.rescale(units)
-            scl = getunitconversion(self.units, units)
-            self.matrix[0:3, 3] *= scl
             self.units = units
 
     def to_dict(self):
diff --git a/tests/resources/example_db/protocols/example_protocol/example_protocol.json b/tests/resources/example_db/protocols/example_protocol/example_protocol.json
index a039be52..e27c8dd5 100644
--- a/tests/resources/example_db/protocols/example_protocol/example_protocol.json
+++ b/tests/resources/example_db/protocols/example_protocol/example_protocol.json
@@ -87,7 +87,19 @@
     "t_end": 0,
     "options": {}
   },
-  "param_constraints": {},
-  "target_constraints": {},
-  "analysis_options": {}
+  "param_constraints": [],
+  "target_constraints": [],
+  "analysis_options": {
+    "standoff_sound_speed": 1500.0,
+    "standoff_density": 1000.0,
+    "ref_sound_speed": 1500.0,
+    "ref_density": 1000.0,
+    "focus_diameter": 0.5,
+    "mainlobe_aspect_ratio": [1.0, 1.0, 5.0],
+    "mainlobe_radius": 2.5e-3,
+    "beamwidth_radius": 5e-3,
+    "sidelobe_radius": 3e-3,
+    "sidelobe_zmin": 1e-3,
+    "distance_units": "m"
+  }
 }
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index d8f73058..fb121e5e 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -1,14 +1,67 @@
 from pathlib import Path
 
+import numpy as np
 import pytest
+import xarray as xa
 
-from openlifu import Protocol
+from openlifu import Protocol, Transducer
+from openlifu.db import Session
 
 
 @pytest.fixture()
 def example_protocol() -> Protocol:
     return Protocol.from_file(Path(__file__).parent/'resources/example_db/protocols/example_protocol/example_protocol.json')
 
+@pytest.fixture()
+def example_transducer() -> Transducer:
+    return Transducer.from_file(Path(__file__).parent/"resources/example_db/transducers/example_transducer/example_transducer.json")
+
+@pytest.fixture()
+def example_session() -> Session:
+    return Session.from_file(Path(__file__).parent/"resources/example_db/subjects/example_subject/sessions/example_session/example_session.json")
+
+@pytest.fixture()
+def example_transducer_transform() -> np.ndarray:
+    return np.array(
+        [
+            [-1,    0,	0,	0],
+            [0,	0.05,	0.998749217771909,	-105],
+            [0,	0.998749217771909,	-0.05,	5],
+            [0,	0,	0,	1]
+        ]
+    )
+
+@pytest.fixture()
+def example_volume() -> xa.Dataset:
+    #TODO the following is in case we want to test a nifti input volume
+    # loading a nifti file and then construction a xarray dataset from it
+    # with open(Path(__file__).parent/"resources/example_db/subjects/example_subject/volumes/example_volume/mni.json") as f:
+    #     nib_volume_metadata = json.load(f)
+    # nib_volume = Nifti1Image.load(Path(__file__).parent/"resources/example_db/subjects/example_subject/volumes/example_volume"/nib_volume_metadata['data_filename'])
+    # l_min, p_min, s_min = nib_volume.affine[:3, -1]
+    # vol_shape = nib_volume.shape
+    # l_max = affines.apply_affine(nib_volume.affine, [vol_shape[0]-1, 0, 0])[0]
+    # p_max = affines.apply_affine(nib_volume.affine, [0, vol_shape[1]-1, 0])[1]
+    # s_max = affines.apply_affine(nib_volume.affine, [0, 0, vol_shape[2]-1])[2]
+    # volume = xa.Dataset(
+    #     {
+    #         'data': xa.DataArray(data=nib_volume.get_fdata(), dims=["L", "P", "S"])
+    #     },
+    #     coords={
+    #         'L': xa.DataArray(dims=["L"], data=np.linspace(l_min, l_max, vol_shape[0]), attrs={'units': "mm"}),
+    #         'P': xa.DataArray(dims=["P"], data=np.linspace(p_min, p_max, vol_shape[1]), attrs={'units': "mm"}),
+    #         'S': xa.DataArray(dims=["S"], data=np.linspace(s_min, s_max, vol_shape[2]), attrs={'units': "mm"})
+    #     },
+    #     attrs={
+    #         'units': "",
+    #         'name': "Test Volume",
+    #         'id': "volume_000",
+    #         'affine': nib_volume.affine,
+    #         'affine_units': "mm"
+    #     }
+    # )
+    return None
+
 @pytest.mark.parametrize("compact_representation", [True, False])
 def test_serialize_deserialize_protocol(example_protocol : Protocol, compact_representation: bool):
     assert example_protocol.from_json(example_protocol.to_json(compact_representation)) == example_protocol
@@ -16,3 +69,23 @@ def test_serialize_deserialize_protocol(example_protocol : Protocol, compact_rep
 def test_default_protocol():
     """Ensure it is possible to construct a default protocol"""
     Protocol()
+
+def test_calc_solution(
+        example_protocol: Protocol,
+        example_transducer: Transducer,
+        example_session: Session,
+        example_transducer_transform: np.ndarray,
+        example_volume: xa.Dataset
+    ):
+    """Make sure a solution can be calculated."""
+    from copy import deepcopy
+
+    from openlifu.util.units import convert_transform
+
+    target = deepcopy(example_session.targets[0])
+    target.rescale("m")
+    transducer_transform = convert_transform(example_transducer_transform, units="mm", tgt_units="m")
+    target.transform(np.linalg.inv(transducer_transform))
+    target.dims = ("lat", "ele", "ax")
+
+    example_protocol.calc_solution(target, example_transducer, transducer_transform, volume=example_volume, session=example_session)
diff --git a/tests/test_seg_method.py b/tests/test_seg_method.py
index c8513b55..18b5b3fa 100644
--- a/tests/test_seg_method.py
+++ b/tests/test_seg_method.py
@@ -1,3 +1,4 @@
+
 import pytest
 
 from openlifu.seg import Material, SegmentationMethod

From 57bb9d1af9e8bd40382389c51eb3291413b77be7 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Tue, 5 Nov 2024 10:09:43 +0100
Subject: [PATCH 02/29] Protocol objects are now correctly json serializable.

Adding solution_analysis module to improve import architecture.
---
 src/openlifu/plan/__init__.py          |  3 ++
 src/openlifu/plan/protocol.py          | 10 ++++---
 src/openlifu/plan/solution.py          | 39 ++------------------------
 src/openlifu/plan/solution_analysis.py | 39 ++++++++++++++++++++++++++
 src/openlifu/util/json.py              |  3 ++
 5 files changed, 53 insertions(+), 41 deletions(-)
 create mode 100644 src/openlifu/plan/solution_analysis.py

diff --git a/src/openlifu/plan/__init__.py b/src/openlifu/plan/__init__.py
index b471f588..6b91ddd6 100644
--- a/src/openlifu/plan/__init__.py
+++ b/src/openlifu/plan/__init__.py
@@ -1,11 +1,14 @@
 from .protocol import Protocol
 from .run import Run
 from .solution import Solution
+from .solution_analysis import SolutionAnalysis, SolutionAnalysisOptions
 from .target_constraints import TargetConstraints
 
 __all__ = [
     "Protocol",
     "Solution",
     "Run",
+    "SolutionAnalysis",
+    "SolutionAnalysisOptions",
     "TargetConstraints"
 ]
diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index 60a4ed34..300cc0b4 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -14,9 +14,11 @@
 from openlifu import bf, geo, seg, sim, xdc
 from openlifu.db.session import Session
 from openlifu.geo import Point
-from openlifu.plan.solution import Solution, SolutionAnalysis, SolutionAnalysisOptions
+from openlifu.plan.solution import Solution
+from openlifu.plan.solution_analysis import SolutionAnalysis, SolutionAnalysisOptions
 from openlifu.plan.target_constraints import TargetConstraints
 from openlifu.sim import run_simulation
+from openlifu.util.json import PYFUSEncoder
 from openlifu.xdc import Transducer
 
 OnPulseMismatchAction = Enum("OnPulseMismatchAction", ["ERROR", "ROUND", "ROUNDUP", "ROUNDDOWN"])
@@ -99,9 +101,9 @@ def to_json(self, compact:bool) -> str:
         Returns: A json string representing the complete Protocol object.
         """
         if compact:
-            return json.dumps(self.to_dict(), separators=(',', ':'))
+            return json.dumps(self.to_dict(), separators=(',', ':'), cls=PYFUSEncoder)
         else:
-            return json.dumps(self.to_dict(), indent=4)
+            return json.dumps(self.to_dict(), indent=4, cls=PYFUSEncoder)
 
     def to_file(self, filename: str):
         """
@@ -289,7 +291,7 @@ def calc_solution(
                 ],
                 dim='focal_point_index',
             )
-        # instantiate and return the solution
+        # instantiateinstantiate and return the solution
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
         solution_id = timestamp
         if session is not None:
diff --git a/src/openlifu/plan/solution.py b/src/openlifu/plan/solution.py
index f62c892d..de75f250 100644
--- a/src/openlifu/plan/solution.py
+++ b/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 List, Optional, Tuple
+from typing import List, Optional
 
 import numpy as np
 import xarray as xa
@@ -11,7 +11,7 @@
 from openlifu.bf import Pulse, Sequence, mask_focus
 from openlifu.bf.mask_focus import MaskOp
 from openlifu.geo import Point
-from openlifu.io.dict_conversion import DictMixin
+from openlifu.plan.solution_analysis import SolutionAnalysis, SolutionAnalysisOptions
 from openlifu.util.json import PYFUSEncoder
 from openlifu.util.units import rescale_data_arr
 from openlifu.xdc import Transducer
@@ -24,41 +24,6 @@ def _construct_nc_filepath_from_json_filepath(json_filepath:Path) -> Path:
     return nc_filepath
 
 
-@dataclass
-class SolutionAnalysis(DictMixin):
-    mainlobe_pnp_MPa: list[float] = field(default_factory=list)
-    mainlobe_isppa_Wcm2: list[float] = field(default_factory=list)
-    mainlobe_ispta_mWcm2: list[float] = field(default_factory=list)
-    beamwidth_lat_3dB_mm: list[float] = field(default_factory=list)
-    beamwidth_ax_3dB_mm: list[float] = field(default_factory=list)
-    beamwidth_lat_6dB_mm: list[float] = field(default_factory=list)
-    beamwidth_ax_6dB_mm: list[float] = field(default_factory=list)
-    sidelobe_pnp_MPa: list[float] = field(default_factory=list)
-    sidelobe_isppa_Wcm2: list[float] = field(default_factory=list)
-    global_pnp_MPa: list[float] = field(default_factory=list)
-    global_isppa_Wcm2: list[float] = field(default_factory=list)
-    p0_Pa: list[float] = field(default_factory=list)
-    TIC: float = None
-    power_W: float = None
-    MI: float = None
-    global_ispta_mWcm2: float = None
-
-
-@dataclass
-class SolutionAnalysisOptions(DictMixin):
-    standoff_sound_speed: float = 1500.0
-    standoff_density: float = 1000.0
-    ref_sound_speed: float = 1500.0
-    ref_density: float = 1000.0
-    focus_diameter: float = 0.5
-    mainlobe_aspect_ratio: Tuple[float, float, float] = (1., 1., 5.)
-    mainlobe_radius: float = 2.5e-3
-    beamwidth_radius: float = 5e-3
-    sidelobe_radius: float = 3e-3
-    sidelobe_zmin: float = 1e-3
-    distance_units: str = "m"
-
-
 @dataclass
 class Solution:
     """
diff --git a/src/openlifu/plan/solution_analysis.py b/src/openlifu/plan/solution_analysis.py
new file mode 100644
index 00000000..479bdf4e
--- /dev/null
+++ b/src/openlifu/plan/solution_analysis.py
@@ -0,0 +1,39 @@
+from dataclasses import dataclass, field
+from typing import Tuple
+
+from openlifu.io.dict_conversion import DictMixin
+
+
+@dataclass
+class SolutionAnalysis(DictMixin):
+    mainlobe_pnp_MPa: list[float] = field(default_factory=list)
+    mainlobe_isppa_Wcm2: list[float] = field(default_factory=list)
+    mainlobe_ispta_mWcm2: list[float] = field(default_factory=list)
+    beamwidth_lat_3dB_mm: list[float] = field(default_factory=list)
+    beamwidth_ax_3dB_mm: list[float] = field(default_factory=list)
+    beamwidth_lat_6dB_mm: list[float] = field(default_factory=list)
+    beamwidth_ax_6dB_mm: list[float] = field(default_factory=list)
+    sidelobe_pnp_MPa: list[float] = field(default_factory=list)
+    sidelobe_isppa_Wcm2: list[float] = field(default_factory=list)
+    global_pnp_MPa: list[float] = field(default_factory=list)
+    global_isppa_Wcm2: list[float] = field(default_factory=list)
+    p0_Pa: list[float] = field(default_factory=list)
+    TIC: float = None
+    power_W: float = None
+    MI: float = None
+    global_ispta_mWcm2: float = None
+
+
+@dataclass
+class SolutionAnalysisOptions(DictMixin):
+    standoff_sound_speed: float = 1500.0
+    standoff_density: float = 1000.0
+    ref_sound_speed: float = 1500.0
+    ref_density: float = 1000.0
+    focus_diameter: float = 0.5
+    mainlobe_aspect_ratio: Tuple[float, float, float] = (1., 1., 5.)
+    mainlobe_radius: float = 2.5e-3
+    beamwidth_radius: float = 5e-3
+    sidelobe_radius: float = 3e-3
+    sidelobe_zmin: float = 1e-3
+    distance_units: str = "m"
diff --git a/src/openlifu/util/json.py b/src/openlifu/util/json.py
index a65cc7a5..78203d29 100644
--- a/src/openlifu/util/json.py
+++ b/src/openlifu/util/json.py
@@ -6,6 +6,7 @@
 
 from openlifu.db.subject import Subject
 from openlifu.geo import Point
+from openlifu.plan.solution_analysis import SolutionAnalysisOptions
 from openlifu.seg.material import Material
 from openlifu.xdc.element import Element
 from openlifu.xdc.transducer import Transducer
@@ -31,6 +32,8 @@ def default(self, obj):
             return obj.to_dict()
         if isinstance(obj, Subject):
             return obj.to_dict()
+        if isinstance(obj, SolutionAnalysisOptions):
+            return obj.to_dict()
         return super().default(obj)
 
 def to_json(obj, filename):

From 0b327a1e10e89929d25df07f00fbb2a492aa5c34 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Tue, 5 Nov 2024 11:33:02 +0100
Subject: [PATCH 03/29] Fixed numpy deprecation warning.

---
 src/openlifu/sim/sim_setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/openlifu/sim/sim_setup.py b/src/openlifu/sim/sim_setup.py
index a2cd7afb..95f1c621 100644
--- a/src/openlifu/sim/sim_setup.py
+++ b/src/openlifu/sim/sim_setup.py
@@ -102,7 +102,7 @@ def get_max_distance(self, arr: Transducer, units: Optional[str] = None):
 
     def get_size(self, dims: Optional[str]=None):
         dims = self.dims if dims is None else dims
-        n = [int(np.round(np.diff(ext)/self.spacing))+1 for ext in [self.x_extent, self.y_extent, self.z_extent]]
+        n = [int((np.round(np.diff(ext)/self.spacing)).item())+1 for ext in [self.x_extent, self.y_extent, self.z_extent]]
         return np.array([n[self.dims.index(dim)] for dim in dims]).squeeze()
 
     def get_spacing(self, units: Optional[str] = None):

From f0be40e4397c5b5d3bc27e5f76f88370573218fc Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Tue, 5 Nov 2024 12:15:33 +0100
Subject: [PATCH 04/29] Refactoring `Protocol` so `scale_solution` is now a
 member of the `Solution` class.

---
 src/openlifu/plan/protocol.py | 52 +++--------------------------------
 src/openlifu/plan/solution.py | 43 +++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 48 deletions(-)

diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index 300cc0b4..3c54c30d 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -15,7 +15,7 @@
 from openlifu.db.session import Session
 from openlifu.geo import Point
 from openlifu.plan.solution import Solution
-from openlifu.plan.solution_analysis import SolutionAnalysis, SolutionAnalysisOptions
+from openlifu.plan.solution_analysis import SolutionAnalysisOptions
 from openlifu.plan.target_constraints import TargetConstraints
 from openlifu.sim import run_simulation
 from openlifu.util.json import PYFUSEncoder
@@ -138,49 +138,6 @@ def check_target(self, target):
             )
             target_constraint.check_bounds(pos)
 
-    def scale_solution(
-            self,
-            solution: Solution,
-            transducer: Transducer,
-            analysis_options: SolutionAnalysisOptions = SolutionAnalysisOptions()
-    ) -> Tuple[Solution, SolutionAnalysis]:
-        """
-        Scale the solution to match the target pressure.
-        If no output is requested, the solution is scaled in-place.
-
-        Args:
-            solution: plan.Solution
-                The solution to be scaled.
-            analysis_options: plan.solution.SolutionAnalysisOptions
-
-        Returns:
-            solution_scaled: the scaled plan.Solution
-            analysis_scaled: the resulting plan.solution.SolutionAnalysis from scaled solution
-        """
-        solution_scaled = deepcopy(solution)
-        analysis = solution.analyze(transducer, options=analysis_options)
-
-        scaling_factors = np.zeros(solution.num_foci())
-        for i in range(solution.num_foci()):
-            scaling_factors[i] = (self.focal_pattern.target_pressure / 1e6) / analysis.mainlobe_pnp_MPa[i]
-            scaling_factors[i] = 2.3167
-        max_scaling = np.max(scaling_factors)
-        v0 = self.pulse.amplitude
-        v1 = v0 * max_scaling
-        apod_factors = scaling_factors / max_scaling
-
-        for i in range(solution.num_foci()):
-            scaling = v1/v0*apod_factors[i]
-            solution_scaled.simulation_result['p_min'][i].data *= scaling
-            solution_scaled.simulation_result['p_max'][i].data *= scaling
-            solution_scaled.simulation_result['ita'][i].data *= scaling**2
-            solution_scaled.apodizations[i] = solution_scaled.apodizations[i]*apod_factors[i]
-        self.pulse.amplitude = v1
-
-        analysis_scaled = solution_scaled.analyze(transducer, options=analysis_options)
-
-        return solution_scaled, analysis_scaled
-
     #TODO: The arg transform needed for sim_setup since transducer.matrix does not exists
     def calc_solution(
         self,
@@ -303,7 +260,7 @@ def calc_solution(
             transducer_id=transducer.id,
             delays=np.stack(delays_to_stack, axis=0),
             apodizations=np.stack(apodizations_to_stack, axis=0),
-            pulse=self.pulse,  #TODO pulse is correctly scaled with `self.scale_solution``
+            pulse=self.pulse,  #TODO pulse is now correctly scaled with `solution.scale`
             sequence=solution_sequence, #TODO incorrect to set the sequence the same as the protocol's
                                         # since sequence pulse_count can be modified if pulse mismatch.
             foci=foci,
@@ -323,9 +280,8 @@ def calc_solution(
                 logging.error(msg=f"Cannot scale solution {solution.id} if simulation is not enabled!")
                 raise ValueError(f"Cannot scale solution {solution.id} if simulation is not enabled!")
             logging.info(f"Scaling solution {solution.id}...")
-            #TODO does analysis needs to be an attribute of solution ?
-            solution_scaled, _ = self.scale_solution(solution, transducer, analysis_options=analysis_options)
-            solution = solution_scaled
+            #TODO can analysis be an attribute of solution ?
+            scaled_solution_analysis = solution.scale(transducer, self.focal_pattern, analysis_options=analysis_options)
 
         # Finally the resulting pressure is max-aggregated and intensity is mean-aggregated, over all focus points .
         pnp_aggregated = solution.simulation_result['p_min'].max(dim="focal_point_index")
diff --git a/src/openlifu/plan/solution.py b/src/openlifu/plan/solution.py
index de75f250..3677ae6d 100644
--- a/src/openlifu/plan/solution.py
+++ b/src/openlifu/plan/solution.py
@@ -1,5 +1,6 @@
 import base64
 import json
+from copy import deepcopy
 from dataclasses import asdict, dataclass, field
 from datetime import datetime
 from pathlib import Path
@@ -9,6 +10,7 @@
 import xarray as xa
 
 from openlifu.bf import Pulse, Sequence, mask_focus
+from openlifu.bf.focal_patterns import FocalPattern
 from openlifu.bf.mask_focus import MaskOp
 from openlifu.geo import Point
 from openlifu.plan.solution_analysis import SolutionAnalysis, SolutionAnalysisOptions
@@ -230,6 +232,47 @@ def analyze(self, transducer: Transducer, options: SolutionAnalysisOptions = Sol
 
         return solution_analysis
 
+    def scale(
+            self,
+            transducer: Transducer,
+            focal_pattern: FocalPattern,
+            analysis_options: SolutionAnalysisOptions = SolutionAnalysisOptions()
+    ) -> SolutionAnalysis:
+        """
+        Scale the solution in-place to match the target pressure.
+
+        Args:
+            transducer: xdc.Transducer
+            focal_pattern: FocalPattern
+            analysis_options: plan.solution.SolutionAnalysisOptions
+
+        Returns:
+            analysis_scaled: the resulting plan.solution.SolutionAnalysis from scaled solution
+        """
+        solution_scaled = deepcopy(self)
+        analysis = solution_scaled.analyze(transducer, options=analysis_options)
+
+        scaling_factors = np.zeros(self.num_foci())
+        for i in range(solution_scaled.num_foci()):
+            scaling_factors[i] = (focal_pattern.target_pressure / 1e6) / analysis.mainlobe_pnp_MPa[i]
+            scaling_factors[i] = 2.3167
+        max_scaling = np.max(scaling_factors)
+        v0 = solution_scaled.pulse.amplitude
+        v1 = v0 * max_scaling
+        apod_factors = scaling_factors / max_scaling
+
+        for i in range(solution_scaled.num_foci()):
+            scaling = v1/v0*apod_factors[i]
+            solution_scaled.simulation_result['p_min'][i].data *= scaling
+            solution_scaled.simulation_result['p_max'][i].data *= scaling
+            solution_scaled.simulation_result['ita'][i].data *= scaling**2
+            solution_scaled.apodizations[i] = solution_scaled.apodizations[i]*apod_factors[i]
+        solution_scaled.pulse.amplitude = v1
+
+        analysis_scaled = solution_scaled.analyze(transducer, options=analysis_options)
+
+        return solution_scaled, analysis_scaled
+
     def get_pulsetrain_dutycycle(self) -> float:
         """
         Compute the pulsetrain dutycycle given a sequence.

From edb303cfe76aca36f667701f9678ad17e06f2bb6 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Tue, 5 Nov 2024 13:50:00 +0100
Subject: [PATCH 05/29] Fixed simulation output variable handling if simulation
 not enabled.

Tests will now run with simulation not enabled to alleviate test runtime.
---
 src/openlifu/plan/protocol.py | 21 +++++++++++----------
 tests/test_protocol.py        | 12 +++++++++++-
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index 3c54c30d..d0eb065c 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -283,15 +283,16 @@ def calc_solution(
             #TODO can analysis be an attribute of solution ?
             scaled_solution_analysis = solution.scale(transducer, self.focal_pattern, analysis_options=analysis_options)
 
-        # Finally the resulting pressure is max-aggregated and intensity is mean-aggregated, over all focus points .
-        pnp_aggregated = solution.simulation_result['p_min'].max(dim="focal_point_index")
-        ppp_aggregated = solution.simulation_result['p_max'].max(dim="focal_point_index")
-        # TODO: Ensure this mean is weighted by the number of times each point is focused on, once openlifu supports hitting points different numbers of times
-        intensity_aggregated = solution.simulation_result['ita'].mean(dim="focal_point_index")
-        simulation_result_aggregated = deepcopy(solution.simulation_result)
-        simulation_result_aggregated = simulation_result_aggregated.drop_dims("focal_point_index")
-        simulation_result_aggregated['p_min'] = pnp_aggregated
-        simulation_result_aggregated['p_max'] = ppp_aggregated
-        simulation_result_aggregated['ita'] = intensity_aggregated
+        if simulate:
+            # Finally the resulting pressure is max-aggregated and intensity is mean-aggregated, over all focus points .
+            pnp_aggregated = solution.simulation_result['p_min'].max(dim="focal_point_index")
+            ppp_aggregated = solution.simulation_result['p_max'].max(dim="focal_point_index")
+            # TODO: Ensure this mean is weighted by the number of times each point is focused on, once openlifu supports hitting points different numbers of times
+            intensity_aggregated = solution.simulation_result['ita'].mean(dim="focal_point_index")
+            simulation_result_aggregated = deepcopy(solution.simulation_result)
+            simulation_result_aggregated = simulation_result_aggregated.drop_dims("focal_point_index")
+            simulation_result_aggregated['p_min'] = pnp_aggregated
+            simulation_result_aggregated['p_max'] = ppp_aggregated
+            simulation_result_aggregated['ita'] = intensity_aggregated
 
         return solution, simulation_result_aggregated
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index fb121e5e..979a59c1 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -78,14 +78,24 @@ def test_calc_solution(
         example_volume: xa.Dataset
     ):
     """Make sure a solution can be calculated."""
+    import logging
     from copy import deepcopy
 
     from openlifu.util.units import convert_transform
 
+    logging.disable(logging.CRITICAL)
     target = deepcopy(example_session.targets[0])
     target.rescale("m")
     transducer_transform = convert_transform(example_transducer_transform, units="mm", tgt_units="m")
     target.transform(np.linalg.inv(transducer_transform))
     target.dims = ("lat", "ele", "ax")
 
-    example_protocol.calc_solution(target, example_transducer, transducer_transform, volume=example_volume, session=example_session)
+    example_protocol.calc_solution(
+        target,
+        example_transducer,
+        transducer_transform,
+        volume=example_volume,
+        session=example_session,
+        simulate=False,
+        scale=False
+    )

From 150d310bbd669324853f605f5ed46f15249964a8 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 11:34:53 +0100
Subject: [PATCH 06/29] \'setup_sim_scene\' is now the minimal implementation

---
 src/openlifu/sim/sim_setup.py | 29 ++++++++---------------------
 1 file changed, 8 insertions(+), 21 deletions(-)

diff --git a/src/openlifu/sim/sim_setup.py b/src/openlifu/sim/sim_setup.py
index 95f1c621..7e95cee3 100644
--- a/src/openlifu/sim/sim_setup.py
+++ b/src/openlifu/sim/sim_setup.py
@@ -8,7 +8,7 @@
 from openlifu.geo import Point
 from openlifu.io.dict_conversion import DictMixin
 from openlifu.seg import SegmentationMethod
-from openlifu.util.units import getunitconversion, rescale_coords
+from openlifu.util.units import getunitconversion
 from openlifu.xdc import Transducer
 
 
@@ -116,11 +116,9 @@ def get_spacing(self, units: Optional[str] = None):
     def setup_sim_scene(
             self,
             transducer: Transducer,
-            transform: np.ndarray,
             target: Point,
             seg_method: SegmentationMethod,
             volume: xa.DataArray = None,
-            interp_method: str = "Linear",
             units: Optional[str] = None
         ) -> Tuple[xa.DataArray, Transducer, Point]:
         """ Prepare a simulation scene composed of a volume, transducer and targets.
@@ -130,38 +128,27 @@ def setup_sim_scene(
         For the volume, a segmentation is also performed that defines the simulation medium.
 
         Args:
-            volume: xa.DataArray
+
             transducer: xdc.Transducer
-            transform: np.ndarray
             target: geo.Point
             seg_method: seg.SegmentationMethod
-            interp_method: str
-                Interpolation method for the simulation grid (Default: \"Linear\").
+            volume: xa.DataArray
+                Optional volume to be used for simulation grid definition (Default: None).
             units: str
                 Units of simulation grid (Default: self.units).
 
         Returns
-            params: The resampled and segmented xa.DataArray volume to the simulation grid
-            transducer_tr: The resampled xdc.Transducer to the simulation grid
-            target_tr: The resampled geo.Point to the simulation grid
+            params: The xa.DataArray simulation grid with physical properties for each voxel
+            transducer: The transformed xdc.Transducer to the simulation grid
+            target: The transformed geo.Point to the simulation grid
         """
         units = self.units if units is None else units
         sim_coords = self.get_coords(units=units)
-        #TODO: in a near future, the following will transform transducer and target. Currently Slicer does it.
-        # sim_matrix = convert_transform(transform, units=transducer.units, tgt_units=units)  #TODO: not clear here where the transform should come from since transducer.get_matrix() does not exist
-        # rescaled_affine = convert_transform(volume.affine, units=volume.affine_units, tgt_units=units)
-        # volume = volume.assign_attrs(affine_units=units, affine=rescaled_affine)
-        # volume_resampled = resample_volume(volume, sim_coords, sim_matrix, interp_method)
-        # transducer_tr = transducer.transform(sim_matrix, units)
-        # target_tr = target.transform(sim_matrix, units=units, new_dims=sim_coords.dims)
         transducer.rescale(units)
         target.rescale(units)
         if volume is None:
             params = seg_method.ref_params(sim_coords)
         else:
-            sim_coords_units = sim_coords[next(iter(sim_coords.keys()))].units
-            volume_resampled = rescale_coords(volume, sim_coords_units)
-            params = seg_method.seg_params(volume_resampled)  #TODO: currently only "water" method is tested with db-simple-example-v04
-                                                              # we should have a whole MRI segmentation instead.
+            params = seg_method.ref_params(sim_coords)
 
         return params, transducer, target

From f787c4299a322c95f9c18d7b2138729ab5934c87 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 12:20:45 +0100
Subject: [PATCH 07/29] Removing unused references related to segmentation and
 resampling from a volume in `setup_sim_scene`

---
 pyproject.toml                              |  1 -
 src/openlifu/seg/seg_methods/segment_mri.py | 22 --------
 src/openlifu/util/geom.py                   | 61 ---------------------
 src/openlifu/util/units.py                  | 21 +------
 4 files changed, 1 insertion(+), 104 deletions(-)
 delete mode 100644 src/openlifu/seg/seg_methods/segment_mri.py
 delete mode 100644 src/openlifu/util/geom.py

diff --git a/pyproject.toml b/pyproject.toml
index cd42e675..a9b98ea1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,7 +40,6 @@ dependencies = [
   "opencv-contrib-python",
   "crc",
   "nibabel",
-  "antspyx",
   "sphinx",
   "ipykernel",
   "k-wave-python>=0.3.4",
diff --git a/src/openlifu/seg/seg_methods/segment_mri.py b/src/openlifu/seg/seg_methods/segment_mri.py
deleted file mode 100644
index 24e59893..00000000
--- a/src/openlifu/seg/seg_methods/segment_mri.py
+++ /dev/null
@@ -1,22 +0,0 @@
-from dataclasses import dataclass
-
-import ants
-from nibabel import Nifti1Image
-
-from openlifu.seg.seg_methods.seg_method import SegmentationMethod
-
-
-@dataclass
-class SegmentMRI(SegmentationMethod):
-    def _segment(self, volume: Nifti1Image):
-        # On using ANTs antropos, see this discussion: https://neurostars.org/t/n-a-whitematter-confounds-in-tsv-output-file/4859/10
-        ants_img = ants.from_numpy(
-            volume.get_fdata(),
-            origin=volume.affine[:3, -1].tolist(),
-            spacing=volume.header['pixdim'][:3].tolist(),
-            direction=volume.affine[:3, :3]
-        )
-        mask = ants.get_mask(ants_img)
-        img_segs = ants.atropos(ants_img, x=mask, m="[0.2, 1x1x1]")
-
-        return img_segs
diff --git a/src/openlifu/util/geom.py b/src/openlifu/util/geom.py
deleted file mode 100644
index b0364283..00000000
--- a/src/openlifu/util/geom.py
+++ /dev/null
@@ -1,61 +0,0 @@
-import numpy as np
-import xarray as xa
-
-from openlifu.util.units import get_ndgrid_from_coords, rescale_coords
-
-
-def resample_volume(
-        volume: xa.DataArray,
-        coords: xa.Coordinates,
-        matrix: np.ndarray,
-        interp_method: str = "linear"
-    ) -> xa.DataArray:
-    """Resample an xa.DataArray Volume data
-
-    Resample an xa.DataArray Volume to the coordinate system
-    specified by coords using the transformation matrix.
-
-    Args:
-        volume: xa.DataArray
-        coords:
-            The coordinate system to which the volume should be resampled
-        matrix: np.ndarray
-            The transformation matrix in the coords units
-
-    Returns:
-        fus.Volume: The transformed Volume object
-    """
-
-    prev_units = volume[next(iter(volume.coords.keys()))].units
-    coords_units = coords[next(iter(coords.keys()))].units
-    volume_rescaled = rescale_coords(volume, coords_units)
-    xyz = get_ndgrid_from_coords(coords)
-    XYZ = np.append(xyz, np.ones((*xyz.shape[:3], 1)), axis=-1)
-    X1 = np.linalg.inv(matrix) @ (matrix @ XYZ)
-    interp_method = None
-    #TODO: get transform matrix to transform grid
-    # inv_matrix = (self.matrix'*self.matrix)\(self.matrix');
-    pass
-    # XP = np.vstack((np.ravel(Xp).T, np.ones(np.prod(Xp.shape), dtype))
-    # XP[:, -1] = 1.0)  # Add a column for homogeneous coordinates
-
-    # inv_matrix = np.linalg.inv(self.matrix.T @ self.matrix)
-
-    # X1 = matrix.dot(XP)
-    # X1 = (inv_matrix @ X1).reshape(Xp.shape[:-1], -1, order='F')
-
-    # if options is None:
-    #     method = 'linear'
-    # else:
-    #     method = options['method']
-
-    # data = self.interp((X1.T,), method=method)  # Interpolate the transformed coordinates
-    # obj = self.newobj(data, coords, id=self.id,
-    #                   name=self.name, matrix=matrix,
-    #                   attrs=self.attrs, units=self.units)
-
-    # if len(self):
-    #     self.rescale(prev_units)
-    rescale_coords(volume, prev_units)
-
-    return volume
diff --git a/src/openlifu/util/units.py b/src/openlifu/util/units.py
index bd5c16cb..df78c80e 100644
--- a/src/openlifu/util/units.py
+++ b/src/openlifu/util/units.py
@@ -1,5 +1,5 @@
 import numpy as np
-from xarray import Coordinates, Dataset
+from xarray import Dataset
 
 #TODO: use Pint (https://github.com/hgrecco/pint) instead to manage physics units in python
 
@@ -243,25 +243,6 @@ def get_ndgrid_from_arr(data_arr: Dataset) -> np.ndarray:
     return ndgrid
 
 
-def get_ndgrid_from_coords(coords_arr: Coordinates) -> np.ndarray:
-    """
-    Creates a ndgrid from xarray.coordinates.
-
-    Args:
-        coords : xarray.Coordinates
-
-    Returns:
-        ndgrid: The ndgrid from the Coordinates.
-    """
-    all_coord = []
-    for coord_key in coords_arr.dims:
-        if 'units' in coords_arr[coord_key].attrs:
-            all_coord += [coords_arr[coord_key].data]
-    ndgrid = np.stack(np.meshgrid(*all_coord, indexing="ij"), axis=-1)
-
-    return ndgrid
-
-
 def convert_transform(matrix: np.ndarray, units: str, tgt_units: str) -> np.ndarray:
     """Given a transform matrix in some units, convert it to another units.
 

From 580ad0c07d39edd0b60dadfad3aa55ab87865140 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 12:30:34 +0100
Subject: [PATCH 08/29] transform argument not needed currently since we use a
 default water sim grid

---
 src/openlifu/plan/protocol.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index d0eb065c..9af77907 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -138,12 +138,10 @@ def check_target(self, target):
             )
             target_constraint.check_bounds(pos)
 
-    #TODO: The arg transform needed for sim_setup since transducer.matrix does not exists
     def calc_solution(
         self,
         target: Point,
         transducer: Transducer,
-        transform: np.ndarray,
         volume: xa.DataArray = None,  #TODO: Do we want to have the volume as a xa.DataArray instead of nifty ?
         session: Session = None, # useful in solution id  #TODO not sure to understand why this type is optional
         simulate: bool = True,
@@ -191,7 +189,7 @@ def calc_solution(
             analysis_options = self.analysis_options
         # check before if target is within bounds
         self.check_target(target)
-        params, transducer, target = sim_options.setup_sim_scene(transducer, transform, target, self.seg_method, volume=volume, units="m")
+        params, transducer, target = sim_options.setup_sim_scene(transducer, target, self.seg_method, volume=volume, units="m")
 
         delays_to_stack: List[np.ndarray] = []
         apodizations_to_stack: List[np.ndarray] = []

From 06942c287b490a332a114763267ad0920124cf10 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 12:46:42 +0100
Subject: [PATCH 09/29] Tests for check_targets

we keep `example_volume` for reference but it is not used anymore

Removed unused SegmentMRI module
---
 src/openlifu/plan/protocol.py            |  4 +---
 src/openlifu/seg/seg_methods/__init__.py |  2 --
 tests/test_protocol.py                   | 12 +++++++-----
 3 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index 9af77907..b7d4331d 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -118,15 +118,13 @@ def to_file(self, filename: str):
             file.write(self.to_json(compact=False))
 
 
-    def check_target(self, target):
+    def check_target(self, target: Point):
         """
         Check if a target is within bounds.
 
         Args:
             target: The geo.Point target to check.
         """
-        #TODO: in the matlab code they handle the case were multiple targets are given.
-        # After our discussion I assumed that we will not handle a list of targets.
         if isinstance(target, list):
             raise ValueError(f"Input target {target} not supposed to be a list!")
 
diff --git a/src/openlifu/seg/seg_methods/__init__.py b/src/openlifu/seg/seg_methods/__init__.py
index 99a1cfa0..6abd1b94 100644
--- a/src/openlifu/seg/seg_methods/__init__.py
+++ b/src/openlifu/seg/seg_methods/__init__.py
@@ -1,6 +1,5 @@
 from . import seg_method
 from .seg_method import SegmentationMethod, UniformSegmentation
-from .segment_mri import SegmentMRI
 from .tissue import Tissue
 from .water import Water
 
@@ -10,5 +9,4 @@
     "UniformSegmentation",
     "Water",
     "Tissue",
-    "SegmentMRI",
 ]
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 979a59c1..c9fb2bf2 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -33,6 +33,7 @@ def example_transducer_transform() -> np.ndarray:
 
 @pytest.fixture()
 def example_volume() -> xa.Dataset:
+    return xa.Dataset()
     #TODO the following is in case we want to test a nifti input volume
     # loading a nifti file and then construction a xarray dataset from it
     # with open(Path(__file__).parent/"resources/example_db/subjects/example_subject/volumes/example_volume/mni.json") as f:
@@ -60,7 +61,6 @@ def example_volume() -> xa.Dataset:
     #         'affine_units': "mm"
     #     }
     # )
-    return None
 
 @pytest.mark.parametrize("compact_representation", [True, False])
 def test_serialize_deserialize_protocol(example_protocol : Protocol, compact_representation: bool):
@@ -70,12 +70,15 @@ def test_default_protocol():
     """Ensure it is possible to construct a default protocol"""
     Protocol()
 
+def test_check_target(example_protocol: Protocol, example_session: Session):
+    """Ensure that the target can be correctly verified."""
+    example_protocol.check_target(example_session.targets[0])
+
 def test_calc_solution(
         example_protocol: Protocol,
         example_transducer: Transducer,
         example_session: Session,
-        example_transducer_transform: np.ndarray,
-        example_volume: xa.Dataset
+        example_transducer_transform: np.ndarray
     ):
     """Make sure a solution can be calculated."""
     import logging
@@ -93,8 +96,7 @@ def test_calc_solution(
     example_protocol.calc_solution(
         target,
         example_transducer,
-        transducer_transform,
-        volume=example_volume,
+        volume=None,
         session=example_session,
         simulate=False,
         scale=False

From df1ca694e35297b255a2bfcbce01816a1c930ea2 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 12:49:03 +0100
Subject: [PATCH 10/29] ` example_volume` is left commented for reference

---
 tests/test_protocol.py | 63 +++++++++++++++++++++---------------------
 1 file changed, 32 insertions(+), 31 deletions(-)

diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index c9fb2bf2..e736ef95 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -2,7 +2,6 @@
 
 import numpy as np
 import pytest
-import xarray as xa
 
 from openlifu import Protocol, Transducer
 from openlifu.db import Session
@@ -31,36 +30,38 @@ def example_transducer_transform() -> np.ndarray:
         ]
     )
 
-@pytest.fixture()
-def example_volume() -> xa.Dataset:
-    return xa.Dataset()
-    #TODO the following is in case we want to test a nifti input volume
-    # loading a nifti file and then construction a xarray dataset from it
-    # with open(Path(__file__).parent/"resources/example_db/subjects/example_subject/volumes/example_volume/mni.json") as f:
-    #     nib_volume_metadata = json.load(f)
-    # nib_volume = Nifti1Image.load(Path(__file__).parent/"resources/example_db/subjects/example_subject/volumes/example_volume"/nib_volume_metadata['data_filename'])
-    # l_min, p_min, s_min = nib_volume.affine[:3, -1]
-    # vol_shape = nib_volume.shape
-    # l_max = affines.apply_affine(nib_volume.affine, [vol_shape[0]-1, 0, 0])[0]
-    # p_max = affines.apply_affine(nib_volume.affine, [0, vol_shape[1]-1, 0])[1]
-    # s_max = affines.apply_affine(nib_volume.affine, [0, 0, vol_shape[2]-1])[2]
-    # volume = xa.Dataset(
-    #     {
-    #         'data': xa.DataArray(data=nib_volume.get_fdata(), dims=["L", "P", "S"])
-    #     },
-    #     coords={
-    #         'L': xa.DataArray(dims=["L"], data=np.linspace(l_min, l_max, vol_shape[0]), attrs={'units': "mm"}),
-    #         'P': xa.DataArray(dims=["P"], data=np.linspace(p_min, p_max, vol_shape[1]), attrs={'units': "mm"}),
-    #         'S': xa.DataArray(dims=["S"], data=np.linspace(s_min, s_max, vol_shape[2]), attrs={'units': "mm"})
-    #     },
-    #     attrs={
-    #         'units': "",
-    #         'name': "Test Volume",
-    #         'id': "volume_000",
-    #         'affine': nib_volume.affine,
-    #         'affine_units': "mm"
-    #     }
-    # )
+#TODO the following is in case we want to test a nifti input volume
+# @pytest.fixture()
+# def example_volume() -> xa.Dataset:
+
+#     # loading a nifti file and then construction a xarray dataset from it
+#     with open(Path(__file__).parent/"resources/example_db/subjects/example_subject/volumes/example_volume/mni.json") as f:
+#         nib_volume_metadata = json.load(f)
+#     nib_volume = Nifti1Image.load(Path(__file__).parent/"resources/example_db/subjects/example_subject/volumes/example_volume"/nib_volume_metadata['data_filename'])
+#     l_min, p_min, s_min = nib_volume.affine[:3, -1]
+#     vol_shape = nib_volume.shape
+#     l_max = affines.apply_affine(nib_volume.affine, [vol_shape[0]-1, 0, 0])[0]
+#     p_max = affines.apply_affine(nib_volume.affine, [0, vol_shape[1]-1, 0])[1]
+#     s_max = affines.apply_affine(nib_volume.affine, [0, 0, vol_shape[2]-1])[2]
+#     volume = xa.Dataset(
+#         {
+#             'data': xa.DataArray(data=nib_volume.get_fdata(), dims=["L", "P", "S"])
+#         },
+#         coords={
+#             'L': xa.DataArray(dims=["L"], data=np.linspace(l_min, l_max, vol_shape[0]), attrs={'units': "mm"}),
+#             'P': xa.DataArray(dims=["P"], data=np.linspace(p_min, p_max, vol_shape[1]), attrs={'units': "mm"}),
+#             'S': xa.DataArray(dims=["S"], data=np.linspace(s_min, s_max, vol_shape[2]), attrs={'units': "mm"})
+#         },
+#         attrs={
+#             'units': "",
+#             'name': "Test Volume",
+#             'id': "volume_000",
+#             'affine': nib_volume.affine,
+#             'affine_units': "mm"
+#         }
+#     )
+
+#     return volume
 
 @pytest.mark.parametrize("compact_representation", [True, False])
 def test_serialize_deserialize_protocol(example_protocol : Protocol, compact_representation: bool):

From 5b2541951fd394c532d772c919e8d1168ac4796a Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 13:27:04 +0100
Subject: [PATCH 11/29] If we need to convert a transform, we want to use the
 method form `Transducer`

---
 src/openlifu/util/units.py | 18 ------------------
 tests/test_protocol.py     |  5 ++---
 2 files changed, 2 insertions(+), 21 deletions(-)

diff --git a/src/openlifu/util/units.py b/src/openlifu/util/units.py
index df78c80e..d03c3b34 100644
--- a/src/openlifu/util/units.py
+++ b/src/openlifu/util/units.py
@@ -241,21 +241,3 @@ def get_ndgrid_from_arr(data_arr: Dataset) -> np.ndarray:
     ndgrid = np.stack(np.meshgrid(*all_coord, indexing="ij"), axis=-1)
 
     return ndgrid
-
-
-def convert_transform(matrix: np.ndarray, units: str, tgt_units: str) -> np.ndarray:
-    """Given a transform matrix in some units, convert it to another units.
-
-    Args:
-        matrix: 4x4 np.ndarray
-            affine transform matrix
-        units: str
-            units of the coordinate space on which the provided transform matrix operates
-
-    Returns: 4x4 affine transform matrix
-    """
-    tgt_units = self.units if tgt_units is None else tgt_units
-    matrix = matrix.copy()
-    matrix[0:3, 3] *= getunitconversion(units, tgt_units)
-
-    return matrix
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index e736ef95..0653dc9d 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -85,12 +85,11 @@ def test_calc_solution(
     import logging
     from copy import deepcopy
 
-    from openlifu.util.units import convert_transform
-
     logging.disable(logging.CRITICAL)
     target = deepcopy(example_session.targets[0])
     target.rescale("m")
-    transducer_transform = convert_transform(example_transducer_transform, units="mm", tgt_units="m")
+    example_transducer.units = "m"
+    transducer_transform = example_transducer.convert_transform(example_transducer_transform, units="mm")
     target.transform(np.linalg.inv(transducer_transform))
     target.dims = ("lat", "ele", "ax")
 

From 911d535fedf2ee578d30f3d243e51d10d0fdab22 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 14:43:47 +0100
Subject: [PATCH 12/29] Checking for pulse mismatch and updating accordingly
 the `Protocol` sequence.

`Protocol.calc_solution` now returns the scaled `SolutionAnalysis`
---
 src/openlifu/plan/protocol.py | 60 +++++++++++++++++++----------------
 tests/test_protocol.py        | 35 +++++++++++++++++++-
 2 files changed, 67 insertions(+), 28 deletions(-)

diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index b7d4331d..16b3a351 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -15,7 +15,7 @@
 from openlifu.db.session import Session
 from openlifu.geo import Point
 from openlifu.plan.solution import Solution
-from openlifu.plan.solution_analysis import SolutionAnalysisOptions
+from openlifu.plan.solution_analysis import SolutionAnalysis, SolutionAnalysisOptions
 from openlifu.plan.target_constraints import TargetConstraints
 from openlifu.sim import run_simulation
 from openlifu.util.json import PYFUSEncoder
@@ -136,6 +136,22 @@ def check_target(self, target: Point):
             )
             target_constraint.check_bounds(pos)
 
+    def fix_pulse_mismatch(self, on_pulse_mismatch: OnPulseMismatchAction, foci: List[Point]):
+        """Fix the protocol sequence pulse count in-place given a pulse_mismatch action."""
+        if on_pulse_mismatch is OnPulseMismatchAction.ERROR:
+            raise ValueError(f"Pulse Count {self.sequence.pulse_count} is not a multiple of the number of foci {len(foci)}")
+        else:
+            if on_pulse_mismatch is OnPulseMismatchAction.ROUND:
+                self.sequence.pulse_count = round(self.sequence.pulse_count / len(foci)) * len(foci)
+            elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDUP:
+                self.sequence.pulse_count = math.ceil(self.sequence.pulse_count / len(foci)) * len(foci)
+            elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDDOWN:
+                self.sequence.pulse_count = math.floor(self.sequence.pulse_count / len(foci)) * len(foci)
+            logging.warning(
+                f"Pulse Count {self.sequence.pulse_count} is not a multiple of the number of foci {len(foci)}."
+                f"Rounding to {self.sequence.pulse_count}."
+            )
+
     def calc_solution(
         self,
         target: Point,
@@ -148,7 +164,7 @@ def calc_solution(
         analysis_options: SolutionAnalysisOptions = None,
         on_pulse_mismatch: OnPulseMismatchAction = OnPulseMismatchAction.ERROR,
         #log : Logger. Default: fus.util.Logger.get()  #TODO what about logging, currently only db/database.py has one ?
-    ) -> Tuple[Solution, xa.DataArray]:  #TODO: make more sense for me to have a single xa.DataArray that holds the
+    ) -> Tuple[Solution, xa.DataArray, SolutionAnalysis]:  #TODO: make more sense for me to have a single xa.DataArray that holds the
                                          # aggregation (pnp, ppp, ita). We could also store it in Solution.simulation_result
                                          # with additional fields 'pnp_aggregated', 'ppp_aggregated' and 'ita_aggregated' ?
         """Calculate the solution and aggregated k-wave simulation outputs.
@@ -178,8 +194,13 @@ def calc_solution(
                 An action to take if the number of pulses in the sequence does not match
                 the number of foci (Default: OnPulseMismatchAction.ERROR).
 
-        Returns: Tuple[Solution, xa.DataArray]
-            Represents the solution object including its analysis, and aggregated simulation output.
+        Returns:
+            solution: Solution
+            simulation_result_aggregated: xa.Dataset
+                If simulation is enabled, then this is the resulting aggregated
+                output (max pressure and mean intensity over all foci).
+            scaled_solution_analysis: SolutionAnalysis
+                This is the resulting rescaled analysis, if scale is enabled.
         """
         if sim_options is None:
             sim_options = self.sim_setup
@@ -192,28 +213,15 @@ def calc_solution(
         delays_to_stack: List[np.ndarray] = []
         apodizations_to_stack: List[np.ndarray] = []
         simulation_outputs_to_stack: List[xa.Dataset] = []
-        simulation_output_stacked: xa.Dataset = None
-        simulation_result_aggregated: xa.Dataset = None
+        simulation_output_stacked: xa.Dataset = xa.Dataset()
+        simulation_result_aggregated: xa.Dataset = xa.Dataset()
+        scaled_solution_analysis: SolutionAnalysis = SolutionAnalysis()
         foci: List[Point] = self.focal_pattern.get_targets(target)
         simulation_cycles = np.max([np.round(self.pulse.duration * self.pulse.frequency), 20])
-        out_solution: Solution = Solution()
 
         # updating solution sequence if pulse mismatch
-        solution_sequence = deepcopy(self.sequence)
         if (self.sequence.pulse_count % len(foci)) != 0:
-            if on_pulse_mismatch is OnPulseMismatchAction.ERROR:
-                raise ValueError(f"Pulse Count {self.sequence.pulse_count} is not a multiple of the number of foci {len(foci)}")
-            else:
-                if on_pulse_mismatch is OnPulseMismatchAction.ROUND:
-                    solution_sequence.pulse_count = round(self.sequence.pulse_count / len(self.foci)) * len(foci)
-                elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDUP:
-                    solution_sequence.pulse_count = math.ceil(self.sequence.pulse_count / len(foci)) * len(self.foci)
-                elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDUP:
-                    solution_sequence.pulse_count = math.floor(self.sequence.pulse_count / len(foci)) * len(self.foci)
-                logging.warning(
-                    f"Pulse Count {self.sequence.pulse_count} is not a multiple of the number of foci {len(foci)}."
-                    f"Rounding to {solution_sequence.pulse_count}."
-                )
+            self.fix_pulse_mismatch(on_pulse_mismatch, foci)
         # run simulation and aggregate the results
         for focus in foci:
             logging.info(f"Beamform for focus {focus}...")
@@ -244,7 +252,7 @@ def calc_solution(
                 ],
                 dim='focal_point_index',
             )
-        # instantiateinstantiate and return the solution
+        # instantiate and return the solution
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
         solution_id = timestamp
         if session is not None:
@@ -256,16 +264,14 @@ def calc_solution(
             transducer_id=transducer.id,
             delays=np.stack(delays_to_stack, axis=0),
             apodizations=np.stack(apodizations_to_stack, axis=0),
-            pulse=self.pulse,  #TODO pulse is now correctly scaled with `solution.scale`
-            sequence=solution_sequence, #TODO incorrect to set the sequence the same as the protocol's
-                                        # since sequence pulse_count can be modified if pulse mismatch.
+            pulse=self.pulse,
+            sequence=self.sequence,
             foci=foci,
             target=target,
             simulation_result=simulation_output_stacked,
             approved=False,
             description= (
                 f"A solution computed for the {self.name} protocol with transducer {transducer.name}"
-                f" for subject volume {volume.id if volume is not None else None}" # TODO put volume ID here if it is not None, once Sadhana's PR #123 is merged
                 f" for target {target.id}."
                 f" This solution was created for the session {session.id} for subject {session.subject_id}." if session is not None else ""
             )
@@ -291,4 +297,4 @@ def calc_solution(
             simulation_result_aggregated['p_max'] = ppp_aggregated
             simulation_result_aggregated['ita'] = intensity_aggregated
 
-        return solution, simulation_result_aggregated
+        return solution, simulation_result_aggregated, scaled_solution_analysis
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 0653dc9d..4d842322 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -1,3 +1,5 @@
+import logging
+import math
 from pathlib import Path
 
 import numpy as np
@@ -5,6 +7,7 @@
 
 from openlifu import Protocol, Transducer
 from openlifu.db import Session
+from openlifu.plan.protocol import OnPulseMismatchAction
 
 
 @pytest.fixture()
@@ -75,6 +78,37 @@ def test_check_target(example_protocol: Protocol, example_session: Session):
     """Ensure that the target can be correctly verified."""
     example_protocol.check_target(example_session.targets[0])
 
+@pytest.mark.parametrize("on_pulse_mismatch", [
+      OnPulseMismatchAction.ERROR,
+      OnPulseMismatchAction.ROUND,
+      OnPulseMismatchAction.ROUNDUP,
+      OnPulseMismatchAction.ROUNDDOWN
+      ]
+    )
+def test_fix_pulse_mismatch(example_protocol: Protocol, example_session: Session, on_pulse_mismatch: OnPulseMismatchAction):
+    """Test if sequence is correctly fixed for all pulse mismatch actions."""
+    logging.disable(logging.CRITICAL)
+
+    target = example_session.targets[0]
+    foci = example_protocol.focal_pattern.get_targets(target)
+    foci = [foci[0], foci[0], foci[0]]
+    if on_pulse_mismatch is OnPulseMismatchAction.ERROR:
+        try:
+            example_protocol.fix_pulse_mismatch(on_pulse_mismatch, foci)
+        except ValueError:
+            assert True
+    else:
+        example_protocol.fix_pulse_mismatch(on_pulse_mismatch, foci)
+        if on_pulse_mismatch is OnPulseMismatchAction.ROUND:
+            expected_pulse_count = round(example_protocol.sequence.pulse_count / len(foci)) * len(foci)
+            assert example_protocol.sequence.pulse_count == expected_pulse_count
+        elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDUP:
+            expected_pulse_count = math.ceil(example_protocol.sequence.pulse_count / len(foci)) * len(foci)
+            assert example_protocol.sequence.pulse_count == expected_pulse_count
+        elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDDOWN:
+            expected_pulse_count = math.floor(example_protocol.sequence.pulse_count / len(foci)) * len(foci)
+            assert example_protocol.sequence.pulse_count == expected_pulse_count
+
 def test_calc_solution(
         example_protocol: Protocol,
         example_transducer: Transducer,
@@ -82,7 +116,6 @@ def test_calc_solution(
         example_transducer_transform: np.ndarray
     ):
     """Make sure a solution can be calculated."""
-    import logging
     from copy import deepcopy
 
     logging.disable(logging.CRITICAL)

From 5ef57379c0ddf83de503dcdf9e64615dba95ce99 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 15:04:14 +0100
Subject: [PATCH 13/29] Refactored `Solution.scale` and added associated tests.

---
 src/openlifu/plan/solution.py | 24 +++++++++++-------------
 tests/test_solution.py        | 11 +++++++++++
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/src/openlifu/plan/solution.py b/src/openlifu/plan/solution.py
index 3677ae6d..d5d74259 100644
--- a/src/openlifu/plan/solution.py
+++ b/src/openlifu/plan/solution.py
@@ -1,6 +1,5 @@
 import base64
 import json
-from copy import deepcopy
 from dataclasses import asdict, dataclass, field
 from datetime import datetime
 from pathlib import Path
@@ -249,29 +248,28 @@ def scale(
         Returns:
             analysis_scaled: the resulting plan.solution.SolutionAnalysis from scaled solution
         """
-        solution_scaled = deepcopy(self)
-        analysis = solution_scaled.analyze(transducer, options=analysis_options)
+        analysis = self.analyze(transducer, options=analysis_options)
 
         scaling_factors = np.zeros(self.num_foci())
-        for i in range(solution_scaled.num_foci()):
+        for i in range(self.num_foci()):
             scaling_factors[i] = (focal_pattern.target_pressure / 1e6) / analysis.mainlobe_pnp_MPa[i]
             scaling_factors[i] = 2.3167
         max_scaling = np.max(scaling_factors)
-        v0 = solution_scaled.pulse.amplitude
+        v0 = self.pulse.amplitude
         v1 = v0 * max_scaling
         apod_factors = scaling_factors / max_scaling
 
-        for i in range(solution_scaled.num_foci()):
+        for i in range(self.num_foci()):
             scaling = v1/v0*apod_factors[i]
-            solution_scaled.simulation_result['p_min'][i].data *= scaling
-            solution_scaled.simulation_result['p_max'][i].data *= scaling
-            solution_scaled.simulation_result['ita'][i].data *= scaling**2
-            solution_scaled.apodizations[i] = solution_scaled.apodizations[i]*apod_factors[i]
-        solution_scaled.pulse.amplitude = v1
+            self.simulation_result['p_min'][i].data *= scaling
+            self.simulation_result['p_max'][i].data *= scaling
+            self.simulation_result['ita'][i].data *= scaling**2
+            self.apodizations[i] = self.apodizations[i]*apod_factors[i]
+        self.pulse.amplitude = v1
 
-        analysis_scaled = solution_scaled.analyze(transducer, options=analysis_options)
+        analysis_scaled = self.analyze(transducer, options=analysis_options)
 
-        return solution_scaled, analysis_scaled
+        return analysis_scaled
 
     def get_pulsetrain_dutycycle(self) -> float:
         """
diff --git a/tests/test_solution.py b/tests/test_solution.py
index 88b21ba9..c22d7a6c 100644
--- a/tests/test_solution.py
+++ b/tests/test_solution.py
@@ -7,6 +7,7 @@
 from helpers import dataclasses_are_equal
 
 from openlifu import Point, Pulse, Sequence, Solution, Transducer
+from openlifu.bf.focal_patterns import SinglePoint
 from openlifu.xdc.element import Element
 
 
@@ -26,6 +27,11 @@ def example_transducer() -> Transducer:
     )
 
 
+@pytest.fixture()
+def example_focal_pattern_single() -> SinglePoint:
+    return SinglePoint(target_pressure=1.0e6)
+
+
 @pytest.fixture()
 def example_solution() -> Solution:
     rng = np.random.default_rng(147)
@@ -120,6 +126,11 @@ def test_num_foci(example_solution:Solution):
     assert example_solution.apodizations.shape[0] == num_foci
 
 
+def test_scale_solution(example_solution: Solution, example_transducer: Transducer, example_focal_pattern_single: SinglePoint):
+    """Test that a solution can be scaled."""
+    example_solution.scale(example_transducer, example_focal_pattern_single)
+
+
 def test_solution_analysis(example_solution: Solution, example_transducer: Transducer):
     """Test that a solution output can be analyzed."""
     example_solution.analyze(example_transducer)

From 2cdc725deccbff009dc542445256870d1dbe23b5 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 15:28:56 +0100
Subject: [PATCH 14/29] rationalizing constants in `Solution.scale` using the
 new `FocalPattern.units` attribute

---
 src/openlifu/bf/focal_patterns/focal_pattern.py             | 5 +++--
 src/openlifu/plan/solution.py                               | 6 +++---
 .../protocols/example_protocol/example_protocol.json        | 1 +
 tests/test_solution.py                                      | 2 +-
 4 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/openlifu/bf/focal_patterns/focal_pattern.py b/src/openlifu/bf/focal_patterns/focal_pattern.py
index e6f07ce9..9e4a5854 100644
--- a/src/openlifu/bf/focal_patterns/focal_pattern.py
+++ b/src/openlifu/bf/focal_patterns/focal_pattern.py
@@ -10,9 +10,10 @@ class FocalPattern(ABC):
     """
     Abstract base class for representing a focal pattern
 
-    :ivar target_pressure: Target pressure of the focal pattern in Pa
+    :ivar target_pressure: Target pressure of the focal pattern in given units
     """
-    target_pressure: float = 1.0 # Pa
+    target_pressure: float = 1.0
+    units: str = "Pa"
 
     @abstractmethod
     def get_targets(self, target: Point):
diff --git a/src/openlifu/plan/solution.py b/src/openlifu/plan/solution.py
index d5d74259..57d683fd 100644
--- a/src/openlifu/plan/solution.py
+++ b/src/openlifu/plan/solution.py
@@ -14,7 +14,7 @@
 from openlifu.geo import Point
 from openlifu.plan.solution_analysis import SolutionAnalysis, SolutionAnalysisOptions
 from openlifu.util.json import PYFUSEncoder
-from openlifu.util.units import rescale_data_arr
+from openlifu.util.units import getunitconversion, rescale_data_arr
 from openlifu.xdc import Transducer
 
 
@@ -252,8 +252,8 @@ def scale(
 
         scaling_factors = np.zeros(self.num_foci())
         for i in range(self.num_foci()):
-            scaling_factors[i] = (focal_pattern.target_pressure / 1e6) / analysis.mainlobe_pnp_MPa[i]
-            scaling_factors[i] = 2.3167
+            focal_pattern_pressure_in_MPa = focal_pattern.target_pressure * getunitconversion(focal_pattern.units, "MPa")
+            scaling_factors[i] = focal_pattern_pressure_in_MPa / analysis.mainlobe_pnp_MPa[i]
         max_scaling = np.max(scaling_factors)
         v0 = self.pulse.amplitude
         v1 = v0 * max_scaling
diff --git a/tests/resources/example_db/protocols/example_protocol/example_protocol.json b/tests/resources/example_db/protocols/example_protocol/example_protocol.json
index e27c8dd5..7f4362a0 100644
--- a/tests/resources/example_db/protocols/example_protocol/example_protocol.json
+++ b/tests/resources/example_db/protocols/example_protocol/example_protocol.json
@@ -16,6 +16,7 @@
   },
   "focal_pattern": {
     "target_pressure": 1.0e6,
+    "units": "Pa",
     "class": "SinglePoint"
   },
   "delay_method": {
diff --git a/tests/test_solution.py b/tests/test_solution.py
index c22d7a6c..2a53b27f 100644
--- a/tests/test_solution.py
+++ b/tests/test_solution.py
@@ -29,7 +29,7 @@ def example_transducer() -> Transducer:
 
 @pytest.fixture()
 def example_focal_pattern_single() -> SinglePoint:
-    return SinglePoint(target_pressure=1.0e6)
+    return SinglePoint(target_pressure=1.0e6, units="Pa")
 
 
 @pytest.fixture()

From 1826c05139c5bf6c955d8cfdf3a0801ee705f442 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 15:37:50 +0100
Subject: [PATCH 15/29] Clarified docstring for `Protocol.calc_solution`

---
 src/openlifu/plan/protocol.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index 16b3a351..dc8cb889 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -165,21 +165,22 @@ def calc_solution(
         on_pulse_mismatch: OnPulseMismatchAction = OnPulseMismatchAction.ERROR,
         #log : Logger. Default: fus.util.Logger.get()  #TODO what about logging, currently only db/database.py has one ?
     ) -> Tuple[Solution, xa.DataArray, SolutionAnalysis]:  #TODO: make more sense for me to have a single xa.DataArray that holds the
-                                         # aggregation (pnp, ppp, ita). We could also store it in Solution.simulation_result
-                                         # with additional fields 'pnp_aggregated', 'ppp_aggregated' and 'ita_aggregated' ?
+                                                           # aggregation (pnp, ppp, ita). We could also store it in Solution.simulation_result
+                                                           # with additional fields 'pnp_aggregated', 'ppp_aggregated' and 'ita_aggregated' ?
         """Calculate the solution and aggregated k-wave simulation outputs.
 
-        Method that computes the delays and apodizations for each focus in
-        the treatment plan, simulates the resulting pressure field to adjust
-        transmit pressures to reach target pressures, and then analyzes the
-        resulting pressure field to compute the resulting acoustic parameters.
+        Method that computes the delays and apodizations for each focus in the treatment plan,
+        simulates the resulting pressure field to adjust transmit pressures to reach target pressures,
+        and then analyzes the resulting pressure field to compute the resulting acoustic parameters.
 
         Args:
             target: The target Point.
+                Target is expected to be in the simulation grid coordinates (lat, ele, ax).
             transducer: A Transducer item.
             volume: xa.DataArray
-                The xa.DataArray volume corresponding to the subject scan (Default: None).
-                If no volume is given, a default simulation grid will be used.
+                The subject scan (Default: None).
+                It is expected to be in the simulation grid coordinates (lat, ele, ax).
+                If None, a default simulation grid will be used.
             session: db.Session
                 A session used to define solution_id (Default: None).
             simulate: bool

From b706923798347aecf1c8e6ced852fba25cdef848 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <37963074+ltetrel@users.noreply.github.com>
Date: Wed, 6 Nov 2024 15:40:10 +0100
Subject: [PATCH 16/29] Improved docstring for `Protocol.check_target`

Co-authored-by: Ebrahim Ebrahim <ebrahim.ebrahim@kitware.com>
---
 src/openlifu/plan/protocol.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index dc8cb889..897e96ed 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -120,7 +120,7 @@ def to_file(self, filename: str):
 
     def check_target(self, target: Point):
         """
-        Check if a target is within bounds.
+        Check if a target is within bounds, raising an exception if it isn't.
 
         Args:
             target: The geo.Point target to check.

From 5dcf6aee13d888be9834c706b37934663cf9e15a Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 17:26:18 +0100
Subject: [PATCH 17/29] Fixed typing if None default

---
 src/openlifu/plan/protocol.py          | 10 +++++-----
 src/openlifu/plan/solution_analysis.py | 10 +++++-----
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index 897e96ed..94cc1f78 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -6,7 +6,7 @@
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 import xarray as xa
@@ -156,12 +156,12 @@ def calc_solution(
         self,
         target: Point,
         transducer: Transducer,
-        volume: xa.DataArray = None,  #TODO: Do we want to have the volume as a xa.DataArray instead of nifty ?
-        session: Session = None, # useful in solution id  #TODO not sure to understand why this type is optional
+        volume: Optional[xa.DataArray] = None,  #TODO: Do we want to have the volume as a xa.DataArray instead of nifty ?
+        session: Optional[Session] = None, # useful in solution id  #TODO not sure to understand why this type is optional
         simulate: bool = True,
         scale: bool = True,
-        sim_options: sim.SimSetup = None,
-        analysis_options: SolutionAnalysisOptions = None,
+        sim_options: Optional[sim.SimSetup] = None,
+        analysis_options: Optional[SolutionAnalysisOptions] = None,
         on_pulse_mismatch: OnPulseMismatchAction = OnPulseMismatchAction.ERROR,
         #log : Logger. Default: fus.util.Logger.get()  #TODO what about logging, currently only db/database.py has one ?
     ) -> Tuple[Solution, xa.DataArray, SolutionAnalysis]:  #TODO: make more sense for me to have a single xa.DataArray that holds the
diff --git a/src/openlifu/plan/solution_analysis.py b/src/openlifu/plan/solution_analysis.py
index 479bdf4e..5cd944d8 100644
--- a/src/openlifu/plan/solution_analysis.py
+++ b/src/openlifu/plan/solution_analysis.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Tuple
+from typing import Optional, Tuple
 
 from openlifu.io.dict_conversion import DictMixin
 
@@ -18,10 +18,10 @@ class SolutionAnalysis(DictMixin):
     global_pnp_MPa: list[float] = field(default_factory=list)
     global_isppa_Wcm2: list[float] = field(default_factory=list)
     p0_Pa: list[float] = field(default_factory=list)
-    TIC: float = None
-    power_W: float = None
-    MI: float = None
-    global_ispta_mWcm2: float = None
+    TIC: Optional[float] = None
+    power_W: Optional[float] = None
+    MI: Optional[float] = None
+    global_ispta_mWcm2: Optional[float] = None
 
 
 @dataclass

From fb3a0ac0a1909ee332b91736995388fc2b21ff04 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 17:52:45 +0100
Subject: [PATCH 18/29] now segmenting volume as well if volume is defined

---
 src/openlifu/sim/sim_setup.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/openlifu/sim/sim_setup.py b/src/openlifu/sim/sim_setup.py
index 7e95cee3..4b6d73c0 100644
--- a/src/openlifu/sim/sim_setup.py
+++ b/src/openlifu/sim/sim_setup.py
@@ -118,7 +118,7 @@ def setup_sim_scene(
             transducer: Transducer,
             target: Point,
             seg_method: SegmentationMethod,
-            volume: xa.DataArray = None,
+            volume: Optional[xa.DataArray] = None,
             units: Optional[str] = None
         ) -> Tuple[xa.DataArray, Transducer, Point]:
         """ Prepare a simulation scene composed of a volume, transducer and targets.
@@ -134,6 +134,7 @@ def setup_sim_scene(
             seg_method: seg.SegmentationMethod
             volume: xa.DataArray
                 Optional volume to be used for simulation grid definition (Default: None).
+                The volume is assumed to be resampled on grid coordinates.
             units: str
                 Units of simulation grid (Default: self.units).
 
@@ -149,6 +150,6 @@ def setup_sim_scene(
         if volume is None:
             params = seg_method.ref_params(sim_coords)
         else:
-            params = seg_method.ref_params(sim_coords)
+            params = seg_method.seg_params(volume)
 
         return params, transducer, target

From ea524b67a6d762097dcd07d62986f2b450eee413 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 18:20:23 +0100
Subject: [PATCH 19/29] Refactored `setup_sim_scene` so it can also segment a
 volume.

---
 src/openlifu/sim/sim_setup.py | 34 +++++++++-------------------------
 1 file changed, 9 insertions(+), 25 deletions(-)

diff --git a/src/openlifu/sim/sim_setup.py b/src/openlifu/sim/sim_setup.py
index 4b6d73c0..f6ed3e12 100644
--- a/src/openlifu/sim/sim_setup.py
+++ b/src/openlifu/sim/sim_setup.py
@@ -109,47 +109,31 @@ def get_spacing(self, units: Optional[str] = None):
         units = self.units if units is None else units
         return getunitconversion(self.units, units)*self.spacing
 
-    #TODO: since we don't have a concept of scene here, and that the simulation scene is needed in protocol.calc_solution,
-    # we will return each prepared object one-by-one.
-    #TODO: Missing the "markers" from matlab scene in "scene.transform(self, coords, matrix, options)"
-    #TODO: The arg transform needs to be added since transducer.matrix does not exists anymore
     def setup_sim_scene(
             self,
-            transducer: Transducer,
-            target: Point,
             seg_method: SegmentationMethod,
-            volume: Optional[xa.DataArray] = None,
-            units: Optional[str] = None
+            volume: Optional[xa.DataArray] = None
         ) -> Tuple[xa.DataArray, Transducer, Point]:
-        """ Prepare a simulation scene composed of a volume, transducer and targets.
+        """ Prepare a simulation scene composed of a simulation grid
 
-        Setup a simulation scene with a volume, transducer and target point.
-        All objects are resampled to the geo-referenced simulation grid (lon, lat, ele).
-        For the volume, a segmentation is also performed that defines the simulation medium.
+        Setup a simulation scene with a simulation grid including physical properties.
+        A segmentation is performed to detect the medium, so we can assign
+        physical properties to each voxel, later used by the ultrasound simulation.
+        This assume that the input volume is resampled to the geo-referenced simulation grid (lon, lat, ele).
 
         Args:
-
-            transducer: xdc.Transducer
-            target: geo.Point
             seg_method: seg.SegmentationMethod
             volume: xa.DataArray
                 Optional volume to be used for simulation grid definition (Default: None).
-                The volume is assumed to be resampled on grid coordinates.
-            units: str
-                Units of simulation grid (Default: self.units).
+                The volume is assumed to be resampled on sim grid coordinates.
 
         Returns
             params: The xa.DataArray simulation grid with physical properties for each voxel
-            transducer: The transformed xdc.Transducer to the simulation grid
-            target: The transformed geo.Point to the simulation grid
         """
-        units = self.units if units is None else units
-        sim_coords = self.get_coords(units=units)
-        transducer.rescale(units)
-        target.rescale(units)
         if volume is None:
+            sim_coords = self.get_coords()
             params = seg_method.ref_params(sim_coords)
         else:
             params = seg_method.seg_params(volume)
 
-        return params, transducer, target
+        return params

From 0670968ea1d9590a1bb236ac0f00c0b354da8403 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 18:31:31 +0100
Subject: [PATCH 20/29] fix sim setup call

---
 src/openlifu/plan/protocol.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index 94cc1f78..b3b72dc4 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -40,6 +40,9 @@ class Protocol:
     target_constraints: List[TargetConstraints] = field(default_factory=list)
     analysis_options: SolutionAnalysisOptions = field(default_factory=SolutionAnalysisOptions)
 
+    def __post_init__(self):
+        self.logger = logging.getLogger(__name__)
+
     @staticmethod
     def from_dict(d : Dict[str,Any]) -> "Protocol":
         d["pulse"] = bf.Pulse.from_dict(d.get("pulse", {}))
@@ -209,7 +212,7 @@ def calc_solution(
             analysis_options = self.analysis_options
         # check before if target is within bounds
         self.check_target(target)
-        params, transducer, target = sim_options.setup_sim_scene(transducer, target, self.seg_method, volume=volume, units="m")
+        params = sim_options.setup_sim_scene(self.seg_method, volume=volume)
 
         delays_to_stack: List[np.ndarray] = []
         apodizations_to_stack: List[np.ndarray] = []

From 0e2b74a55622be2ce20b974f17495f4f3b7236a2 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 18:34:41 +0100
Subject: [PATCH 21/29] logging now through an attribute of `Protocol`

---
 src/openlifu/plan/protocol.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index b3b72dc4..a5926a25 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -150,7 +150,7 @@ def fix_pulse_mismatch(self, on_pulse_mismatch: OnPulseMismatchAction, foci: Lis
                 self.sequence.pulse_count = math.ceil(self.sequence.pulse_count / len(foci)) * len(foci)
             elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDDOWN:
                 self.sequence.pulse_count = math.floor(self.sequence.pulse_count / len(foci)) * len(foci)
-            logging.warning(
+            self.logger.warning(
                 f"Pulse Count {self.sequence.pulse_count} is not a multiple of the number of foci {len(foci)}."
                 f"Rounding to {self.sequence.pulse_count}."
             )
@@ -165,8 +165,7 @@ def calc_solution(
         scale: bool = True,
         sim_options: Optional[sim.SimSetup] = None,
         analysis_options: Optional[SolutionAnalysisOptions] = None,
-        on_pulse_mismatch: OnPulseMismatchAction = OnPulseMismatchAction.ERROR,
-        #log : Logger. Default: fus.util.Logger.get()  #TODO what about logging, currently only db/database.py has one ?
+        on_pulse_mismatch: OnPulseMismatchAction = OnPulseMismatchAction.ERROR
     ) -> Tuple[Solution, xa.DataArray, SolutionAnalysis]:  #TODO: make more sense for me to have a single xa.DataArray that holds the
                                                            # aggregation (pnp, ppp, ita). We could also store it in Solution.simulation_result
                                                            # with additional fields 'pnp_aggregated', 'ppp_aggregated' and 'ita_aggregated' ?
@@ -228,11 +227,11 @@ def calc_solution(
             self.fix_pulse_mismatch(on_pulse_mismatch, foci)
         # run simulation and aggregate the results
         for focus in foci:
-            logging.info(f"Beamform for focus {focus}...")
+            self.logger.info(f"Beamform for focus {focus}...")
             delays, apodization = self.beamform(arr=transducer, target=focus, params=params)
             simulation_output_xarray = None
             if simulate:
-                logging.info(f"Simulate for focus {focus}...")
+                self.logger.info(f"Simulate for focus {focus}...")
                 simulation_output_xarray, _ = run_simulation(
                     arr=transducer,
                     params=params,
@@ -283,9 +282,9 @@ def calc_solution(
         # optionally scale the solution with simulation result
         if scale:
             if not simulate:
-                logging.error(msg=f"Cannot scale solution {solution.id} if simulation is not enabled!")
+                self.logger.error(msg=f"Cannot scale solution {solution.id} if simulation is not enabled!")
                 raise ValueError(f"Cannot scale solution {solution.id} if simulation is not enabled!")
-            logging.info(f"Scaling solution {solution.id}...")
+            self.logger.info(f"Scaling solution {solution.id}...")
             #TODO can analysis be an attribute of solution ?
             scaled_solution_analysis = solution.scale(transducer, self.focal_pattern, analysis_options=analysis_options)
 

From 2b519b02c0e52d5154c0dcc46ad7f919d653950e Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 19:17:49 +0100
Subject: [PATCH 22/29] better testing coverage for target verification

---
 tests/test_protocol.py | 37 ++++++++++++++++++++++++++++++-------
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 4d842322..38a8a6be 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -8,6 +8,7 @@
 from openlifu import Protocol, Transducer
 from openlifu.db import Session
 from openlifu.plan.protocol import OnPulseMismatchAction
+from openlifu.plan.target_constraints import TargetConstraints
 
 
 @pytest.fixture()
@@ -74,16 +75,36 @@ def test_default_protocol():
     """Ensure it is possible to construct a default protocol"""
     Protocol()
 
-def test_check_target(example_protocol: Protocol, example_session: Session):
+@pytest.mark.parametrize("target_constraints",[
+            [
+                TargetConstraints(dim="P", units="mm", min=0.0, max=float("inf"))
+            ],
+            [
+                TargetConstraints(dim="P", units="m", min=-0.001, max=0.0)
+            ],
+            [
+                TargetConstraints(dim="L", units="mm", min=-100.0, max=0.0),
+                TargetConstraints(dim="P", units="mm", min=-100.0, max=0.0),
+                TargetConstraints(dim="S", units="mm", min=-100.0, max=-10.0)
+            ]
+        ]
+    )
+def test_check_target(example_protocol: Protocol, example_session: Session, target_constraints: TargetConstraints):
     """Ensure that the target can be correctly verified."""
-    example_protocol.check_target(example_session.targets[0])
+    example_protocol.target_constraints = target_constraints
+    try:
+        example_protocol.check_target(example_session.targets[0])
+    except ValueError:
+        assert True
+    else:
+        pytest.fail(f"Verification failed for {target_constraints}!")
 
 @pytest.mark.parametrize("on_pulse_mismatch", [
-      OnPulseMismatchAction.ERROR,
-      OnPulseMismatchAction.ROUND,
-      OnPulseMismatchAction.ROUNDUP,
-      OnPulseMismatchAction.ROUNDDOWN
-      ]
+            OnPulseMismatchAction.ERROR,
+            OnPulseMismatchAction.ROUND,
+            OnPulseMismatchAction.ROUNDUP,
+            OnPulseMismatchAction.ROUNDDOWN
+        ]
     )
 def test_fix_pulse_mismatch(example_protocol: Protocol, example_session: Session, on_pulse_mismatch: OnPulseMismatchAction):
     """Test if sequence is correctly fixed for all pulse mismatch actions."""
@@ -97,6 +118,8 @@ def test_fix_pulse_mismatch(example_protocol: Protocol, example_session: Session
             example_protocol.fix_pulse_mismatch(on_pulse_mismatch, foci)
         except ValueError:
             assert True
+        else:
+            pytest.fail("on_pulse_mismatch exception not triggered as excepted!")
     else:
         example_protocol.fix_pulse_mismatch(on_pulse_mismatch, foci)
         if on_pulse_mismatch is OnPulseMismatchAction.ROUND:

From 8f55507cba685471d5b374019b53270e5fe39bbc Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Wed, 6 Nov 2024 19:22:19 +0100
Subject: [PATCH 23/29] pytest raise instead of fail

---
 tests/test_protocol.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 38a8a6be..6a6f0a39 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -97,7 +97,7 @@ def test_check_target(example_protocol: Protocol, example_session: Session, targ
     except ValueError:
         assert True
     else:
-        pytest.fail(f"Verification failed for {target_constraints}!")
+        pytest.raises(f"Verification failed for {target_constraints}!")
 
 @pytest.mark.parametrize("on_pulse_mismatch", [
             OnPulseMismatchAction.ERROR,
@@ -119,7 +119,7 @@ def test_fix_pulse_mismatch(example_protocol: Protocol, example_session: Session
         except ValueError:
             assert True
         else:
-            pytest.fail("on_pulse_mismatch exception not triggered as excepted!")
+            pytest.raises("on_pulse_mismatch exception not triggered as excepted!")
     else:
         example_protocol.fix_pulse_mismatch(on_pulse_mismatch, foci)
         if on_pulse_mismatch is OnPulseMismatchAction.ROUND:

From c26600d19ec5ffaa1b1f207591998b06db640192 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <37963074+ltetrel@users.noreply.github.com>
Date: Thu, 7 Nov 2024 16:37:21 +0100
Subject: [PATCH 24/29] test formating and better condition test

Co-authored-by: Ebrahim Ebrahim <ebrahim.ebrahim@kitware.com>
---
 tests/test_protocol.py | 34 ++++++++++++++++------------------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 6a6f0a39..78875035 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -75,29 +75,27 @@ def test_default_protocol():
     """Ensure it is possible to construct a default protocol"""
     Protocol()
 
-@pytest.mark.parametrize("target_constraints",[
-            [
-                TargetConstraints(dim="P", units="mm", min=0.0, max=float("inf"))
-            ],
-            [
-                TargetConstraints(dim="P", units="m", min=-0.001, max=0.0)
-            ],
-            [
-                TargetConstraints(dim="L", units="mm", min=-100.0, max=0.0),
-                TargetConstraints(dim="P", units="mm", min=-100.0, max=0.0),
-                TargetConstraints(dim="S", units="mm", min=-100.0, max=-10.0)
-            ]
+@pytest.mark.parametrize(
+    "target_constraints",
+    [
+        [
+            TargetConstraints(dim="P", units="mm", min=0.0, max=float("inf")),
+        ],
+        [
+            TargetConstraints(dim="P", units="m", min=-0.001, max=0.0),
+        ],
+        [
+            TargetConstraints(dim="L", units="mm", min=-100.0, max=0.0),
+            TargetConstraints(dim="P", units="mm", min=-100.0, max=0.0),
+            TargetConstraints(dim="S", units="mm", min=-100.0, max=-10.0),
         ]
-    )
+    ]
+)
 def test_check_target(example_protocol: Protocol, example_session: Session, target_constraints: TargetConstraints):
     """Ensure that the target can be correctly verified."""
     example_protocol.target_constraints = target_constraints
-    try:
+    with pytest.raises(ValueError, match="not within bounds"):
         example_protocol.check_target(example_session.targets[0])
-    except ValueError:
-        assert True
-    else:
-        pytest.raises(f"Verification failed for {target_constraints}!")
 
 @pytest.mark.parametrize("on_pulse_mismatch", [
             OnPulseMismatchAction.ERROR,

From 9ef9f02a23cde5927f8ab0f3cf1e75bcd560dd19 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Thu, 7 Nov 2024 17:13:17 +0100
Subject: [PATCH 25/29] Re-worked fix_pulse_mismatch test for better
 rationalization

---
 tests/test_protocol.py | 48 ++++++++++++++++++------------------------
 1 file changed, 20 insertions(+), 28 deletions(-)

diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 78875035..c83f05ec 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -1,11 +1,11 @@
 import logging
-import math
 from pathlib import Path
 
 import numpy as np
 import pytest
 
 from openlifu import Protocol, Transducer
+from openlifu.bf.focal_patterns import Wheel
 from openlifu.db import Session
 from openlifu.plan.protocol import OnPulseMismatchAction
 from openlifu.plan.target_constraints import TargetConstraints
@@ -24,20 +24,12 @@ def example_session() -> Session:
     return Session.from_file(Path(__file__).parent/"resources/example_db/subjects/example_subject/sessions/example_session/example_session.json")
 
 @pytest.fixture()
-def example_transducer_transform() -> np.ndarray:
-    return np.array(
-        [
-            [-1,    0,	0,	0],
-            [0,	0.05,	0.998749217771909,	-105],
-            [0,	0.998749217771909,	-0.05,	5],
-            [0,	0,	0,	1]
-        ]
-    )
+def example_wheel_pattern() -> Wheel:
+    return Wheel(num_spokes=6)
 
 #TODO the following is in case we want to test a nifti input volume
 # @pytest.fixture()
 # def example_volume() -> xa.Dataset:
-
 #     # loading a nifti file and then construction a xarray dataset from it
 #     with open(Path(__file__).parent/"resources/example_db/subjects/example_subject/volumes/example_volume/mni.json") as f:
 #         nib_volume_metadata = json.load(f)
@@ -104,37 +96,34 @@ def test_check_target(example_protocol: Protocol, example_session: Session, targ
             OnPulseMismatchAction.ROUNDDOWN
         ]
     )
-def test_fix_pulse_mismatch(example_protocol: Protocol, example_session: Session, on_pulse_mismatch: OnPulseMismatchAction):
+def test_fix_pulse_mismatch(
+        example_protocol: Protocol,
+        example_session: Session,
+        example_wheel_pattern: Wheel,
+        on_pulse_mismatch: OnPulseMismatchAction
+    ):
     """Test if sequence is correctly fixed for all pulse mismatch actions."""
     logging.disable(logging.CRITICAL)
 
     target = example_session.targets[0]
-    foci = example_protocol.focal_pattern.get_targets(target)
-    foci = [foci[0], foci[0], foci[0]]
+    foci = example_wheel_pattern.get_targets(target)
+    num_foci = len(foci)
     if on_pulse_mismatch is OnPulseMismatchAction.ERROR:
-        try:
+        with pytest.raises(ValueError, match="not a multiple of the number of foci"):
             example_protocol.fix_pulse_mismatch(on_pulse_mismatch, foci)
-        except ValueError:
-            assert True
-        else:
-            pytest.raises("on_pulse_mismatch exception not triggered as excepted!")
     else:
         example_protocol.fix_pulse_mismatch(on_pulse_mismatch, foci)
         if on_pulse_mismatch is OnPulseMismatchAction.ROUND:
-            expected_pulse_count = round(example_protocol.sequence.pulse_count / len(foci)) * len(foci)
-            assert example_protocol.sequence.pulse_count == expected_pulse_count
+            assert example_protocol.sequence.pulse_count == num_foci
         elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDUP:
-            expected_pulse_count = math.ceil(example_protocol.sequence.pulse_count / len(foci)) * len(foci)
-            assert example_protocol.sequence.pulse_count == expected_pulse_count
+            assert example_protocol.sequence.pulse_count == 2*num_foci
         elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDDOWN:
-            expected_pulse_count = math.floor(example_protocol.sequence.pulse_count / len(foci)) * len(foci)
-            assert example_protocol.sequence.pulse_count == expected_pulse_count
+            assert example_protocol.sequence.pulse_count == num_foci
 
 def test_calc_solution(
         example_protocol: Protocol,
         example_transducer: Transducer,
-        example_session: Session,
-        example_transducer_transform: np.ndarray
+        example_session: Session
     ):
     """Make sure a solution can be calculated."""
     from copy import deepcopy
@@ -143,7 +132,10 @@ def test_calc_solution(
     target = deepcopy(example_session.targets[0])
     target.rescale("m")
     example_transducer.units = "m"
-    transducer_transform = example_transducer.convert_transform(example_transducer_transform, units="mm")
+    transducer_transform = example_transducer.convert_transform(
+        example_session.array_transform.matrix,
+        example_session.array_transform.units
+    )
     target.transform(np.linalg.inv(transducer_transform))
     target.dims = ("lat", "ele", "ax")
 

From 3311c86879d5fb730ae30b81ebb0ff6a6a070ad3 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Thu, 7 Nov 2024 17:45:48 +0100
Subject: [PATCH 26/29] Some tests will be handled in the future in #152 since
 they require specific and verified data

---
 tests/test_protocol.py | 29 -----------------------------
 tests/test_solution.py | 10 ----------
 2 files changed, 39 deletions(-)

diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index c83f05ec..2119b4b7 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -1,7 +1,6 @@
 import logging
 from pathlib import Path
 
-import numpy as np
 import pytest
 
 from openlifu import Protocol, Transducer
@@ -119,31 +118,3 @@ def test_fix_pulse_mismatch(
             assert example_protocol.sequence.pulse_count == 2*num_foci
         elif on_pulse_mismatch is OnPulseMismatchAction.ROUNDDOWN:
             assert example_protocol.sequence.pulse_count == num_foci
-
-def test_calc_solution(
-        example_protocol: Protocol,
-        example_transducer: Transducer,
-        example_session: Session
-    ):
-    """Make sure a solution can be calculated."""
-    from copy import deepcopy
-
-    logging.disable(logging.CRITICAL)
-    target = deepcopy(example_session.targets[0])
-    target.rescale("m")
-    example_transducer.units = "m"
-    transducer_transform = example_transducer.convert_transform(
-        example_session.array_transform.matrix,
-        example_session.array_transform.units
-    )
-    target.transform(np.linalg.inv(transducer_transform))
-    target.dims = ("lat", "ele", "ax")
-
-    example_protocol.calc_solution(
-        target,
-        example_transducer,
-        volume=None,
-        session=example_session,
-        simulate=False,
-        scale=False
-    )
diff --git a/tests/test_solution.py b/tests/test_solution.py
index 2a53b27f..7dc013d7 100644
--- a/tests/test_solution.py
+++ b/tests/test_solution.py
@@ -124,13 +124,3 @@ def test_num_foci(example_solution:Solution):
     assert len(example_solution.simulation_result['focal_point_index']) == num_foci
     assert example_solution.delays.shape[0] == num_foci
     assert example_solution.apodizations.shape[0] == num_foci
-
-
-def test_scale_solution(example_solution: Solution, example_transducer: Transducer, example_focal_pattern_single: SinglePoint):
-    """Test that a solution can be scaled."""
-    example_solution.scale(example_transducer, example_focal_pattern_single)
-
-
-def test_solution_analysis(example_solution: Solution, example_transducer: Transducer):
-    """Test that a solution output can be analyzed."""
-    example_solution.analyze(example_transducer)

From a1b39a674671206f0c2f5250560bc4fb2de97157 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Thu, 7 Nov 2024 17:51:07 +0100
Subject: [PATCH 27/29] Cleaning old nifti to xarray tests

---
 tests/test_protocol.py | 32 --------------------------------
 1 file changed, 32 deletions(-)

diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 2119b4b7..3233fe82 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -26,38 +26,6 @@ def example_session() -> Session:
 def example_wheel_pattern() -> Wheel:
     return Wheel(num_spokes=6)
 
-#TODO the following is in case we want to test a nifti input volume
-# @pytest.fixture()
-# def example_volume() -> xa.Dataset:
-#     # loading a nifti file and then construction a xarray dataset from it
-#     with open(Path(__file__).parent/"resources/example_db/subjects/example_subject/volumes/example_volume/mni.json") as f:
-#         nib_volume_metadata = json.load(f)
-#     nib_volume = Nifti1Image.load(Path(__file__).parent/"resources/example_db/subjects/example_subject/volumes/example_volume"/nib_volume_metadata['data_filename'])
-#     l_min, p_min, s_min = nib_volume.affine[:3, -1]
-#     vol_shape = nib_volume.shape
-#     l_max = affines.apply_affine(nib_volume.affine, [vol_shape[0]-1, 0, 0])[0]
-#     p_max = affines.apply_affine(nib_volume.affine, [0, vol_shape[1]-1, 0])[1]
-#     s_max = affines.apply_affine(nib_volume.affine, [0, 0, vol_shape[2]-1])[2]
-#     volume = xa.Dataset(
-#         {
-#             'data': xa.DataArray(data=nib_volume.get_fdata(), dims=["L", "P", "S"])
-#         },
-#         coords={
-#             'L': xa.DataArray(dims=["L"], data=np.linspace(l_min, l_max, vol_shape[0]), attrs={'units': "mm"}),
-#             'P': xa.DataArray(dims=["P"], data=np.linspace(p_min, p_max, vol_shape[1]), attrs={'units': "mm"}),
-#             'S': xa.DataArray(dims=["S"], data=np.linspace(s_min, s_max, vol_shape[2]), attrs={'units': "mm"})
-#         },
-#         attrs={
-#             'units': "",
-#             'name': "Test Volume",
-#             'id': "volume_000",
-#             'affine': nib_volume.affine,
-#             'affine_units': "mm"
-#         }
-#     )
-
-#     return volume
-
 @pytest.mark.parametrize("compact_representation", [True, False])
 def test_serialize_deserialize_protocol(example_protocol : Protocol, compact_representation: bool):
     assert example_protocol.from_json(example_protocol.to_json(compact_representation)) == example_protocol

From 912a7331d7d8ff2f70bfdfccf182fd7c10d19be4 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Thu, 7 Nov 2024 18:08:02 +0100
Subject: [PATCH 28/29] Adding `Solution.compute_scaling_factors\' to further
 break down and simplify future tests for `Solution.scale`

---
 src/openlifu/plan/solution.py | 41 +++++++++++++++++++++++++++--------
 1 file changed, 32 insertions(+), 9 deletions(-)

diff --git a/src/openlifu/plan/solution.py b/src/openlifu/plan/solution.py
index 57d683fd..e80912ee 100644
--- a/src/openlifu/plan/solution.py
+++ b/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 List, Optional
+from typing import List, Optional, Tuple
 
 import numpy as np
 import xarray as xa
@@ -231,6 +231,36 @@ def analyze(self, transducer: Transducer, options: SolutionAnalysisOptions = Sol
 
         return solution_analysis
 
+    def compute_scaling_factors(
+            self,
+            focal_pattern: FocalPattern,
+            analysis: SolutionAnalysis
+        ) -> Tuple[np.ndarray, float, float]:
+        """
+
+        Compute the scaling factors used to re-scale the apodizations, simulation results and pulse amplitude.
+
+        Args:
+            focal_pattern: FocalPattern
+            analysis: SolutionAnalysis
+
+        Returns:
+            apod_factors: A np.ndarray apodization factors
+            v0: A float representing the original pulse amplitude
+            v1: A float representing the new pulse amplitude
+        """
+        scaling_factors = np.zeros(self.num_foci())
+
+        for i in range(self.num_foci()):
+            focal_pattern_pressure_in_MPa = focal_pattern.target_pressure * getunitconversion(focal_pattern.units, "MPa")
+            scaling_factors[i] = focal_pattern_pressure_in_MPa / analysis.mainlobe_pnp_MPa[i]
+        max_scaling = np.max(scaling_factors)
+        v0 = self.pulse.amplitude
+        v1 = v0 * max_scaling
+        apod_factors = scaling_factors / max_scaling
+
+        return apod_factors, v0, v1
+
     def scale(
             self,
             transducer: Transducer,
@@ -250,14 +280,7 @@ def scale(
         """
         analysis = self.analyze(transducer, options=analysis_options)
 
-        scaling_factors = np.zeros(self.num_foci())
-        for i in range(self.num_foci()):
-            focal_pattern_pressure_in_MPa = focal_pattern.target_pressure * getunitconversion(focal_pattern.units, "MPa")
-            scaling_factors[i] = focal_pattern_pressure_in_MPa / analysis.mainlobe_pnp_MPa[i]
-        max_scaling = np.max(scaling_factors)
-        v0 = self.pulse.amplitude
-        v1 = v0 * max_scaling
-        apod_factors = scaling_factors / max_scaling
+        apod_factors, v0, v1 = self.compute_scaling_factors(focal_pattern, analysis)
 
         for i in range(self.num_foci()):
             scaling = v1/v0*apod_factors[i]

From 15a6d0551485ac898383f81219b1830f25978b62 Mon Sep 17 00:00:00 2001
From: Loic Tetrel <loic.tetrel@kitware.com>
Date: Thu, 7 Nov 2024 18:41:23 +0100
Subject: [PATCH 29/29] Cleaning some TODOs

---
 src/openlifu/bf/get_beamwidth.py |  2 +-
 src/openlifu/plan/protocol.py    | 10 ++++------
 src/openlifu/plan/solution.py    |  2 +-
 src/openlifu/util/units.py       |  1 -
 4 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/src/openlifu/bf/get_beamwidth.py b/src/openlifu/bf/get_beamwidth.py
index 82daaed9..6649a979 100644
--- a/src/openlifu/bf/get_beamwidth.py
+++ b/src/openlifu/bf/get_beamwidth.py
@@ -98,7 +98,7 @@ def get_beamwidth(vol: DataArray, coords_units: str, focus: Point, cutoff: float
             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
+        logging.warning("Invalid inliers, attempting to add jitter to create a valid volume...")
         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))
diff --git a/src/openlifu/plan/protocol.py b/src/openlifu/plan/protocol.py
index a5926a25..8b399bcf 100644
--- a/src/openlifu/plan/protocol.py
+++ b/src/openlifu/plan/protocol.py
@@ -36,7 +36,7 @@ class Protocol:
     delay_method: bf.DelayMethod = field(default_factory=bf.delay_methods.Direct)
     apod_method: bf.ApodizationMethod = field(default_factory=bf.apod_methods.Uniform)
     seg_method: seg.SegmentationMethod = field(default_factory=seg.seg_methods.Water)
-    param_constraints: dict = field(default_factory=dict)  #TODO: this seems to be used only in `plan.check_analysis`` but not called anywhere
+    param_constraints: dict = field(default_factory=dict)  #TODO: this seems to be used only in `plan.check_analysis` but not called anywhere
     target_constraints: List[TargetConstraints] = field(default_factory=list)
     analysis_options: SolutionAnalysisOptions = field(default_factory=SolutionAnalysisOptions)
 
@@ -159,16 +159,14 @@ def calc_solution(
         self,
         target: Point,
         transducer: Transducer,
-        volume: Optional[xa.DataArray] = None,  #TODO: Do we want to have the volume as a xa.DataArray instead of nifty ?
-        session: Optional[Session] = None, # useful in solution id  #TODO not sure to understand why this type is optional
+        volume: Optional[xa.DataArray] = None,
+        session: Optional[Session] = None,
         simulate: bool = True,
         scale: bool = True,
         sim_options: Optional[sim.SimSetup] = None,
         analysis_options: Optional[SolutionAnalysisOptions] = None,
         on_pulse_mismatch: OnPulseMismatchAction = OnPulseMismatchAction.ERROR
-    ) -> Tuple[Solution, xa.DataArray, SolutionAnalysis]:  #TODO: make more sense for me to have a single xa.DataArray that holds the
-                                                           # aggregation (pnp, ppp, ita). We could also store it in Solution.simulation_result
-                                                           # with additional fields 'pnp_aggregated', 'ppp_aggregated' and 'ita_aggregated' ?
+    ) -> Tuple[Solution, xa.DataArray, SolutionAnalysis]:
         """Calculate the solution and aggregated k-wave simulation outputs.
 
         Method that computes the delays and apodizations for each focus in the treatment plan,
diff --git a/src/openlifu/plan/solution.py b/src/openlifu/plan/solution.py
index e80912ee..7aec6f70 100644
--- a/src/openlifu/plan/solution.py
+++ b/src/openlifu/plan/solution.py
@@ -143,7 +143,7 @@ def analyze(self, transducer: Transducer, options: SolutionAnalysisOptions = Sol
 
             # get focus region masks (for mainlobe, sidelobe and beamwidth)
             mainlobe_mask = mask_focus(
-                self.simulation_result,  #TODO: Original code uses coords, but too complicated to maniplulate a Coordinates class
+                self.simulation_result,
                 foc,
                 options.mainlobe_radius,
                 mask_op=MaskOp.LESS_EQUAL,
diff --git a/src/openlifu/util/units.py b/src/openlifu/util/units.py
index d03c3b34..a22d89f6 100644
--- a/src/openlifu/util/units.py
+++ b/src/openlifu/util/units.py
@@ -1,7 +1,6 @@
 import numpy as np
 from xarray import Dataset
 
-#TODO: use Pint (https://github.com/hgrecco/pint) instead to manage physics units in python
 
 def getunittype(unit):
     unit = unit.lower()