P/segmethod cleanup

db_dvc.dvc

@@ -1,6 +1,6 @@
 outs:
-- md5: bf109e310734c61c0eecff737562b425.dir
-  nfiles: 257
+- md5: b7b8aea2669121775f0de4df5d046d6e.dir
+  nfiles: 797
   hash: md5
   path: db_dvc
-  size: 933880511
+  size: 2357261179

src/openlifu/plan/protocol.py

@@ -62,7 +62,7 @@ class Protocol:
     apod_method: Annotated[bf.ApodizationMethod, OpenLIFUFieldData("Apodization method", "The method used to calculate transmit apodizations. By default, apodizations are uniform")] = field(default_factory=bf.apod_methods.Uniform)
     """The method used to calculate transmit apodizations. By default, apodizations are uniform"""
 
-    seg_method: Annotated[seg.SegmentationMethod, OpenLIFUFieldData("Segmentation method", "The method used to segment the subject's MRI for delay calculation. By default, the entire field is assumed to be water")] = field(default_factory=seg.seg_methods.Water)
+    seg_method: Annotated[seg.SegmentationMethod, OpenLIFUFieldData("Segmentation method", "The method used to segment the subject's MRI for delay calculation. By default, the entire field is assumed to be water")] = field(default_factory=seg.seg_methods.UniformWater)
     """The method used to segment the subject's MRI for delay calculation. By default, the entire field is assumed to be water"""
 
     param_constraints: Annotated[dict, OpenLIFUFieldData("Parameter constraints", "The constraints on the analysis parameters. If computed parameters are outside of the ranges defined here, warnings or errors may be flagged to reject the solution")] = field(default_factory=dict)

src/openlifu/seg/__init__.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+from . import seg_methods
 from .material import AIR, MATERIALS, SKULL, STANDOFF, TISSUE, WATER, Material
-from .seg_methods import SegmentationMethod
+from .seg_method import SegmentationMethod
 
 __all__ = [
     "Material",
@@ -12,4 +13,5 @@
     "AIR",
     "STANDOFF",
     "SegmentationMethod",
+    "seg_methods",
 ]

src/openlifu/seg/material.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Annotated
+from typing import Annotated, Any
 
 from openlifu.util.annotations import OpenLIFUFieldData
 
@@ -22,9 +22,6 @@
                                         "units": "W/m/K"}}
 @dataclass
 class Material:
-    id: Annotated[str, OpenLIFUFieldData("Material ID", "The unique identifier of the material")] = "material"
-    """The unique identifier of the material"""
-
     name: Annotated[str, OpenLIFUFieldData("Material name", "Name for the material")] = "Material"
     """Name for the material"""
 
@@ -43,6 +40,16 @@ class Material:
     thermal_conductivity: Annotated[float, OpenLIFUFieldData("Thermal conductivity (W/m/K)", "Thermal conductivity of the material (W/m/K)")] = 0.598  # W/m/K
     """Thermal conductivity of the material (W/m/K)"""
 
+    def to_dict(self):
+        return {
+            "name": self.name,
+            "sound_speed": self.sound_speed,
+            "density": self.density,
+            "attenuation": self.attenuation,
+            "specific_heat": self.specific_heat,
+            "thermal_conductivity": self.thermal_conductivity
+        }
+
     @classmethod
     def param_info(cls, param_id: str):
         if param_id not in PARAM_INFO:
@@ -55,65 +62,38 @@ def get_param(self, param_id: str):
         return self.__getattribute__(param_id)
 
     @staticmethod
-    def get_materials(material_id="all", as_dict=True):
-        material_id = ("water", "tissue", "skull", "air", "standoff") if material_id == "all" else material_id
-        if isinstance(material_id, (list, tuple)):
-            materials = {m: Material.get_materials(m, as_dict=False) for m in material_id}
-        elif material_id in MATERIALS:
-            materials = MATERIALS[material_id]
-        else:
-            raise ValueError(f"Material {material_id} not found.")
-        if as_dict:
-            return {materials.id: materials}
-        else:
-            return materials
+    def from_dict(d: dict[str, Any]):
+        return Material(**d)
 
-    @staticmethod
-    def from_dict(d):
-        if isinstance(d, (list, tuple)):
-            return {dd['id']: Material.from_dict(dd) for dd in d}
-        elif isinstance(d, str):
-            return Material.get_materials(d, as_dict=False)
-        elif isinstance(d, Material):
-            return d
-        else:
-            return Material(**d)
-
-
-WATER = Material(id="water",
-                 name="water",
+WATER = Material(name="water",
                  sound_speed=1500.0,
                  density=1000.0,
                  attenuation=0.0,
                  specific_heat=4182.0,
                  thermal_conductivity=0.598)
 
-TISSUE = Material(id="tissue",
-                  name="tissue",
+TISSUE = Material(name="tissue",
                   sound_speed=1540.0,
                   density=1000.0,
                   attenuation=0.0,
                   specific_heat=3600.0,
                   thermal_conductivity=0.5)
 
-SKULL = Material(id="skull",
-                 name="skull",
+SKULL = Material(name="skull",
                  sound_speed=4080.0,
                  density=1900.0,
                  attenuation=0.0,
                  specific_heat=1100.0,
                  thermal_conductivity=0.3)
 
-AIR = Material(id="air",
-               name="air",
+AIR = Material(name="air",
                sound_speed=344.0,
                density=1.25,
                attenuation=0.0,
                specific_heat=1012.0,
                thermal_conductivity=0.025)
 
-STANDOFF = Material(id="standoff",
-                    name="standoff",
+STANDOFF = Material(name="standoff",
                     sound_speed=1420.0,
                     density=1000.0,
                     attenuation=1.0,

src/openlifu/seg/seg_methods/seg_method.py → src/openlifu/seg/seg_method.py

@@ -1,52 +1,63 @@
 from __future__ import annotations
 
 import copy
-from abc import abstractmethod
-from dataclasses import asdict, dataclass, field
-from typing import Annotated
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Annotated, Any
 
 import numpy as np
 import xarray as xa
 
-from openlifu.seg import seg_methods
 from openlifu.seg.material import MATERIALS, PARAM_INFO, Material
 from openlifu.util.annotations import OpenLIFUFieldData
 
 
 @dataclass
-class SegmentationMethod:
+class SegmentationMethod(ABC):
     materials: Annotated[dict[str, Material], OpenLIFUFieldData("Segmentation materials", "Dictionary mapping of label names to material definitions used during segmentation")] = field(default_factory=lambda: MATERIALS.copy())
     """Dictionary mapping of label names to material definitions used during segmentation"""
 
     ref_material: Annotated[str, OpenLIFUFieldData("Reference material", "Reference material ID to use")] = "water"
     """Reference material ID to use"""
 
     def __post_init__(self):
+        if self.materials is None:
+            self.materials = MATERIALS.copy()
         if self.ref_material not in self.materials:
             raise ValueError(f"Reference material {self.ref_material} not found.")
+
     @abstractmethod
-    def _segment(self, volume: xa.DataArray):
+    def _segment(self, volume: xa.DataArray) -> xa.DataArray:
         pass
 
+    def to_dict(self) -> dict[str, Any]:
+        d = self.__dict__.copy()
+        d['materials'] = { k: v.to_dict() for k, v in self.materials.items() }
+        d['class'] = self.__class__.__name__
+        return d
+
     @staticmethod
-    def from_dict(d):
-        if isinstance(d, str):
-            if d == "water":
-                return seg_methods.Water()
-            elif d == "tissue":
-                return seg_methods.Tissue()
-            elif d == "segmented":
-                return seg_methods.SegmentMRI()
-        else:
-            d = copy.deepcopy(d)
-            short_classname = d.pop("class")
-            if "materials" in d:
-                for material_key, material_definition in d["materials"].items():
-                    if not isinstance(material_definition, Material): # if it is given as a dict rather than a fully hydrated object
-                        d["materials"][material_key] = Material.from_dict(material_definition)
-            module_dict = seg_methods.__dict__
-            class_constructor = module_dict[short_classname]
-            return class_constructor(**d)
+    def from_dict(d: dict) -> SegmentationMethod:
+        from openlifu.seg import seg_methods
+        if not isinstance(d, dict):  # previous implementations might pass str
+            raise TypeError(f"Expected dict for from_dict, got {type(d).__name__}")
+
+        d = copy.deepcopy(d)
+        short_classname = d.pop("class")
+
+        # Recursively construct Material instances
+        materials_dict = d.get("materials")
+        if materials_dict is not None:
+            d["materials"] = {
+                k: v if isinstance(v, Material) else Material.from_dict(v)
+                for k, v in materials_dict.items()
+            }
+
+        # Ignore ref_material if class is `UniformWater` or `UniformTissue`
+        if short_classname in ["UniformWater", "UniformTissue"]:
+            d.pop("ref_material")
+        class_constructor = getattr(seg_methods, short_classname)
+        return class_constructor(**d)
 
     def _material_indices(self, materials: dict | None = None):
         materials = self.materials if materials is None else materials
@@ -84,13 +95,3 @@ def _ref_segment(self, coords: xa.Coordinates):
         sz = list(coords.sizes.values())
         seg = xa.DataArray(np.full(sz, m_idx, dtype=int), coords=coords)
         return seg
-
-    def to_dict(self):
-        d = asdict(self)
-        d['class'] = self.__class__.__name__
-        return d
-
-@dataclass
-class UniformSegmentation(SegmentationMethod):
-    def _segment(self, vol: xa.DataArray):
-        return self._ref_segment(vol.coords)

src/openlifu/seg/seg_methods/__init__.py

@@ -1,14 +1,9 @@
 from __future__ import annotations
 
-from . import seg_method
-from .seg_method import SegmentationMethod, UniformSegmentation
-from .tissue import Tissue
-from .water import Water
+from .uniform import UniformSegmentation, UniformTissue, UniformWater
 
 __all__ = [
-    "seg_method",
-    "SegmentationMethod",
     "UniformSegmentation",
-    "Water",
-    "Tissue",
+    "UniformWater",
+    "UniformTissue",
 ]

src/openlifu/seg/seg_methods/uniform.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import xarray as xa
+
+from openlifu.seg.material import MATERIALS, Material
+from openlifu.seg.seg_method import SegmentationMethod
+
+
+class UniformSegmentation(SegmentationMethod):
+    def _segment(self, volume: xa.DataArray):
+        return self._ref_segment(volume.coords)
+
+class UniformTissue(UniformSegmentation):
+    """ Assigns the tissue material to all voxels in the volume. """
+    def __init__(self, materials: dict[str, Material] | None = None):
+        if materials is None:
+            materials = MATERIALS.copy()
+        super().__init__(materials=materials, ref_material="tissue")
+
+class UniformWater(UniformSegmentation):
+    """ Assigns the water material to all voxels in the volume. """
+    def __init__(self, materials: dict[str, Material] | None = None):
+        if materials is None:
+            materials = MATERIALS.copy()
+        super().__init__(materials=materials, ref_material="water")