Incorporate virtual fit results into Session (#179)

Virtual fit results take the form of a dictionary mapping target IDs
to pairs `(approval, transforms)`, where:

  `approval` is a boolean indicating whether the virtual fit for that target has been approved, and
  `transforms` is a list of transducer transforms resulting from the virtual fit for that target.

The idea is that the list of transforms would be ordered from best to worst, and should of course
contain at least one transform.

Author: ebrahimebrahim

src/openlifu/db/database.py

@@ -153,15 +153,17 @@ def write_session(self, subject:Subject, session:Session, on_conflict=OnConflict
         if session.subject_id != subject.id:
             raise ValueError("IDs do not match between the given subject and the subject referenced in the session.")
 
-        # Validate the approved fit target ID
-        if (
-            session.virtual_fit_approval_for_target_id is not None
-            and session.virtual_fit_approval_for_target_id not in [target.id for target in session.targets]
-        ):
-            raise ValueError(
-                f"The provided virtual_fit_approval_for_target_id of {session.virtual_fit_approval_for_target_id} is not"
-                " in this session's list of targets."
-            )
+        # Validate the virtual fit results
+        for target_id, (_, transforms) in session.virtual_fit_results.items():
+            if target_id not in [target.id for target in session.targets]:
+                raise ValueError(
+                    f"The virtual_fit_results of session {session.id} references a target {target_id} that is not"
+                    " in the session's list of targets."
+                )
+            if len(transforms)<1:
+                raise ValueError(
+                    f"The virtual_fit_results of session {session.id} provides no transforms for target {target_id}."
+                )
 
         # Check if the session already exists in the database
         session_ids = self.get_session_ids(subject.id)

src/openlifu/db/session.py

@@ -1,8 +1,9 @@
+import copy
 import json
 from dataclasses import asdict, dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 
@@ -22,7 +23,7 @@ class ArrayTransform:
 
     units : str
     """The units of the space on which to apply the transform matrix , e.g. "mm"
-    (In order to apply the transform to transducer points,.
+    (In order to apply the transform to transducer points,
     first represent the points in these units.)"""
 
 @dataclass
@@ -68,9 +69,15 @@ class Session:
     attrs: dict = field(default_factory=dict)
     """Dictionary of additional custom attributes to save to the session"""
 
-    virtual_fit_approval_for_target_id: Optional[str] = None
-    """Approval state of virtual fit. `None` if there is no approval, otherwise this is the ID
-    of the target for which virtual fitting has been marked approved."""
+    virtual_fit_results: Dict[str,Tuple[bool,List[ArrayTransform]]] = field(default_factory=dict)
+    """Virtual fit results. This is a dictionary mapping target IDs to pairs `(approval, transforms)`,
+    where:
+        `approval` is a boolean indicating whether the virtual fit for that target has been approved, and
+        `transforms` is a list of transducer transforms resulting from the virtual fit for that target.
+
+    The idea is that the list of transforms would be ordered from best to worst, and should of course
+    contain at least one transform.
+    """
 
     transducer_tracking_approved: Optional[bool] = False
     """Approval state of transducer tracking. `True` means the user has provided some kind of
@@ -128,6 +135,12 @@ def from_dict(d:Dict):
             d['targets'] = [Point.from_dict(d['targets'])]
         elif isinstance(d['targets'], Point):
             d['targets'] = [d['targets']]
+        if 'virtual_fit_results' in d:
+            for target_id,(approval,transforms) in d['virtual_fit_results'].items():
+                d['virtual_fit_results'][target_id] = (
+                    approval,
+                    [ArrayTransform(t_dict["matrix"], t_dict["units"]) for t_dict in transforms],
+                )
         if isinstance(d['markers'], list):
             if len(d['markers'])>0 and isinstance(d['markers'][0], dict):
                 d['markers'] = [Point.from_dict(p) for p in d['markers']]
@@ -143,14 +156,20 @@ def to_dict(self):
 
         :returns: Dictionary of session parameters
         """
-        d = self.__dict__.copy()
+        d = copy.deepcopy(self.__dict__) # Deep copy needed so that we don't modify the internals of self below
         d['date_created'] = d['date_created'].isoformat()
         d['date_modified'] = d['date_modified'].isoformat()
         d['targets'] = [p.to_dict() for p in d['targets']]
         d['markers'] = [p.to_dict() for p in d['markers']]
 
         d['array_transform'] = asdict(d['array_transform'])
 
+        for target_id,(approval,transforms) in d['virtual_fit_results'].items():
+            d['virtual_fit_results'][target_id] = (
+                approval,
+                [asdict(t) for t in transforms],
+            )
+
         return d
 
     @staticmethod

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

@@ -27,6 +27,21 @@
   },
   "attrs": {},
   "date_modified": "2024-04-09 11:27:30",
-  "virtual_fit_approval_for_target_id": "example_target",
+  "virtual_fit_results": {
+    "example_target": [
+      true,
+      [
+        {
+          "matrix": [
+            [1.1, 0, 0, 0],
+            [0, 1.2, 0, 0],
+            [0, 0, 1.3, 0],
+            [0, 0.05, 0, 1.4]
+          ],
+          "units": "mm"
+        }
+      ]
+    ]
+  },
   "transducer_tracking_approved": false
 }

tests/test_database.py

@@ -3,16 +3,18 @@
 from contextlib import nullcontext as does_not_raise
 from datetime import datetime, timedelta
 from pathlib import Path
-from typing import Optional
+from typing import List
 from unittest.mock import patch
 
+import numpy as np
 import pytest
 from helpers import dataclasses_are_equal
 from vtk import vtkImageData, vtkPolyData
 
 from openlifu import Point, Solution
 from openlifu.db import Session, Subject, User
 from openlifu.db.database import Database, OnConflictOpts
+from openlifu.db.session import ArrayTransform
 from openlifu.photoscan import Photoscan
 from openlifu.plan import Protocol, Run
 from openlifu.xdc import Transducer
@@ -347,26 +349,37 @@ def test_write_session_mismatched_id(example_database: Database, example_subject
         example_database.write_session(example_subject, session)
 
 @pytest.mark.parametrize(
-    ("virtual_fit_approval_for_target_id", "expectation"),
+    ("target_ids", "numbers_of_transforms", "expectation"),
     [
-        (None, does_not_raise()), # see https://docs.pytest.org/en/6.2.x/example/parametrize.html#parametrizing-conditional-raising
-        ("an_existing_target_id", does_not_raise()),
-        ("bogus_target_id", pytest.raises(ValueError, match="virtual_fit_approval_for_target_id.*not in")),
+        # see https://docs.pytest.org/en/6.2.x/example/parametrize.html#parametrizing-conditional-raising
+        ([], [], does_not_raise()),
+        (["an_existing_target_id"], [1], does_not_raise()),
+        (["an_existing_target_id"], [2], does_not_raise()),
+        (["bogus_target_id"], [1], pytest.raises(ValueError, match="references a target")),
+        (["an_existing_target_id", "bogus_target_id"], [1,1], pytest.raises(ValueError, match="references a target")),
+        (["an_existing_target_id"], [0], pytest.raises(ValueError, match="provides no transforms")),
     ]
 )
 def test_write_session_with_invalid_fit_approval(
     example_database: Database,
     example_subject: Subject,
-    virtual_fit_approval_for_target_id: Optional[str],
+    target_ids: List[str],
+    numbers_of_transforms: List[int],
     expectation,
 ):
-    """Verify that writing a session with fit approval raises the invalid target error if and only if the
-    target being approved does not exist."""
+    """Verify that write_session complains appropriately about invalid virtual fit results"""
+    rng = np.random.default_rng()
     session = Session(
         id="unique_id_2764592837465",
         subject_id=example_subject.id,
         targets=[Point(id="an_existing_target_id")],
-        virtual_fit_approval_for_target_id=virtual_fit_approval_for_target_id,
+        virtual_fit_results={
+            target_id : (
+                True,
+                [ArrayTransform(matrix=rng.random(size=(4,4)),units="mm") for _ in range(num_transforms)],
+            )
+            for target_id, num_transforms in zip(target_ids, numbers_of_transforms)
+        },
     )
     with expectation:
         example_database.write_session(example_subject, session)