From 02aea0aed5a39e8fc842e3d06208342835821cac Mon Sep 17 00:00:00 2001
From: Zain Sohail <zainsohail04@gmail.com>
Date: Mon, 16 Sep 2024 16:03:51 +0200
Subject: [PATCH] fix some config problems

---
 sed/config/config_model.py     | 21 +++++++++++++--------
 sed/loader/mpes/loader.py      |  2 +-
 tests/calibrator/test_delay.py |  7 +++----
 3 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/sed/config/config_model.py b/sed/config/config_model.py
index ae2ef64b..2356d4dd 100644
--- a/sed/config/config_model.py
+++ b/sed/config/config_model.py
@@ -8,7 +8,6 @@
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import DirectoryPath
-from pydantic import Field
 from pydantic import field_validator
 from pydantic import FilePath
 from pydantic import HttpUrl
@@ -87,7 +86,7 @@ class subChannel(BaseModel):
 class DataframeModel(BaseModel):
     columns: ColumnsModel
     units: Optional[dict[str, str]] = None
-    channels: dict[str, ChannelModel] = Field(default_factory=dict)
+    channels: Optional[dict[str, ChannelModel]] = None
     # other settings
     tof_binwidth: float
     tof_binning: int
@@ -214,12 +213,18 @@ class Calibration(BaseModel):
     calibration: Optional[Calibration] = None
 
     class Offsets(BaseModel):
-        creation_date: datetime
-        constant: float
-        flip_delay_axis: bool
-        weight: float
-        preserve_mean: bool
-        reduction: str
+        creation_date: Optional[datetime] = None
+        constant: Optional[float] = None
+        flip_delay_axis: Optional[bool] = False
+
+        ## This seems rather complicated way to define offsets,
+        # inconsistent in how args vs config are for add_offsets
+        class OffsetColumn(BaseModel):
+            weight: float
+            preserve_mean: bool
+            reduction: Optional[str] = None
+
+        columns: Optional[dict[str, OffsetColumn]] = None
 
     offsets: Optional[Offsets] = None
 
diff --git a/sed/loader/mpes/loader.py b/sed/loader/mpes/loader.py
index ebc9a08c..eaf296f0 100644
--- a/sed/loader/mpes/loader.py
+++ b/sed/loader/mpes/loader.py
@@ -72,7 +72,7 @@ def hdf5_to_dataframe(
 
     electron_channels = []
     column_names = []
-
+    print("Print values: ", channels)
     for name, channel in channels.items():
         if channel["format"] == "per_electron":
             if channel["dataset_key"] in test_proc:
diff --git a/tests/calibrator/test_delay.py b/tests/calibrator/test_delay.py
index d7677296..a81542e2 100644
--- a/tests/calibrator/test_delay.py
+++ b/tests/calibrator/test_delay.py
@@ -136,9 +136,8 @@ def test_delay_parameters_from_delay_range_mm() -> None:
         "offsets": {
             "constant": 1,
             "flip_delay_axis": True,
-            "bam": {
-                "weight": 0.001,
-                "preserve_mean": False,
+            "columns": {
+                "bam": {"weight": 0.001, "preserve_mean": False},
             },
         },
     },
@@ -201,7 +200,7 @@ def test_add_offset_from_dict(df=test_dataframe) -> None:
     """test that the timing offset is corrected for correctly from config"""
     cfg_ = cfg.copy()
     offsets = cfg["delay"]["offsets"]  # type:ignore
-    offsets["bam"].pop("weight")
+    offsets["columns"]["bam"].pop("weight")
     offsets["flip_delay_axis"] = False
     cfg_.pop("delay")
     config = parse_config(