implemented load data, save and open_file

scripts/HoV6Sn6_contour.py

@@ -11,7 +11,7 @@
 tavi = TAVI()
 
 tavi_file_name = "./test_data/tavi_test_exp1031.h5"
-tavi.new_file(tavi_file_name)
+tavi.new(tavi_file_name)
 
 nexus_file_name = "./test_data/nexus_exp1031.h5"
 tavi.load_nexus_data_from_disk(nexus_file_name)

scripts/MnTe_contour.py

@@ -14,7 +14,7 @@
 tavi = TAVI()
 
 tavi_file_name = "./test_data/tavi_test_exp0813.h5"
-tavi.new_file(tavi_file_name)
+tavi.new(tavi_file_name)
 tavi.load_nexus_data_from_disk(nexus_folder)
 dataset = tavi.data["IPTS34735_HB3_exp0813"]
 

scripts/test_fit_group.py

@@ -42,7 +42,7 @@ def test_fit_group(tavi):
     tavi = TAVI()
 
     tavi_file_name = "./test_data/tavi_test_exp1031.h5"
-    tavi.new_file(tavi_file_name)
+    tavi.new(tavi_file_name)
 
     nexus_file_name = "./test_data/nexus_exp1031.h5"
     tavi.load_nexus_data_from_disk(nexus_file_name)

scripts/test_plotter.py

@@ -59,7 +59,7 @@ def test_2d_plot(tavi, tas):
     tavi = TAVI()
 
     tavi_file_name = "./test_data/tavi_test_exp424.h5"
-    tavi.open_file(tavi_file_name)
+    tavi.open(tavi_file_name)
 
     tas = CN()
     # tas.load_instrument_from_dicts(cg4c_config_params)

scripts/test_scan.py

@@ -22,7 +22,7 @@ def test_scan_group_contour_exp710():
     tavi = TAVI()
 
     tavi_file_name = "./tests/test_data_folder/tavi_test_exp710.h5"
-    tavi.new_file(tavi_file_name)
+    tavi.new(tavi_file_name)
 
     nexus_file_name = "./tests/test_data_folder/nexus_exp710.h5"
     tavi.load_nexus_data_from_disk(nexus_file_name)
@@ -91,7 +91,7 @@ def test_scan_group_waterfall(tavi):
     tavi = TAVI()
 
     tavi_file_name = "./test_data/tavi_test_exp424.h5"
-    tavi.open_file(tavi_file_name)
+    tavi.open(tavi_file_name)
 
     # test_scan_group_contour(tavi)
 

src/tavi/data/nxentry.py

@@ -191,14 +191,22 @@ def from_spice(
         return cls._dict_to_nexus_entry(nexus_dict)
 
     @classmethod
-    def from_nexus(cls, path_to_nexus: str, scan_num: Optional[int] = None):
+    def from_nexus(
+        cls,
+        path_to_nexus: str,
+        scan_num: Optional[int] = None,
+        prefix: Optional[str] = None,
+    ):
         """return a NexusEntry instance from loading a NeXus file
 
         Args:
             path_to_nexus (str): path to a NeXus file with the extension .h5
             scan_num (int): read all scans in file if not None
+            prefix (str)
         """
         with h5py.File(path_to_nexus, "r") as nexus_file:
+            if prefix is not None:
+                nexus_file = nexus_file[prefix]
             if scan_num is None:  # read all scans in file
                 nexus_dict = cls._read_recursively(nexus_file)
             else:  # read one scan only

src/tavi/data/tavi.py

@@ -4,7 +4,7 @@
 
 import h5py
 
-from tavi.data.scan import Scan
+from tavi.data.nxentry import NexusEntry
 from tavi.data.scan_group import ScanGroup
 
 
@@ -31,26 +31,16 @@ def __init__(self) -> None:
         self.fits: dict = {}
         self.plots: dict = {}
 
-    def new_file(self, file_path: Optional[str] = None) -> None:
-        """Create a new tavi file
-        Note:
-            create at file_path if given, otherwise create a temporary
-            file named tavi_temp.h5 at the current folder.
-        """
-        if file_path is None:
-            current_path = os.getcwd()
-            self.file_path = current_path + "/tavi_temp.h5"
-        else:
-            self.file_path = file_path
-
-        try:
-            with h5py.File(self.file_path, "w", track_order=True) as root:
-                root.create_group("data", track_order=True)
-                root.create_group("processed_data", track_order=True)
-                root.create_group("fits", track_order=True)
-                root.create_group("plots", track_order=True)
-        except OSError:
-            print(f"Cannot create tavi file at {self.file_path}")
+    @staticmethod
+    def load_data(scan_dict: dict):
+        data_dict: dict = {}
+        for scan_name, nxentry in scan_dict.items():
+            exp_id = nxentry["attrs"]["dataset_name"]
+            if not data_dict.get(exp_id):  # first entry
+                data_dict.update({exp_id: {scan_name: nxentry}})
+            else:
+                data_dict[exp_id].update({scan_name: nxentry})
+        return data_dict
 
     def load_nexus_data_from_disk(self, path_to_hdf5_folder):
         """Copy hdf5 data from path_to_hdf5_folder into tavi."""
@@ -62,93 +52,81 @@ def load_nexus_data_from_disk(self, path_to_hdf5_folder):
         scan_list = [scan for scan in scan_list if scan.startswith("scan")]
         scan_list.sort()
 
-        with h5py.File(self.file_path, "r+", track_order=True) as tavi_file:
-            scans = {}
-            for scan in scan_list:
-                scan_name = scan.split(".")[0]  # e.g. "scan0001"
-                scan_path = path_to_hdf5_folder + scan
-                with h5py.File(scan_path, "r") as scan_file:
-                    dataset_name = scan_file[scan_name].attrs["dataset_name"]
-                    if dataset_name not in tavi_file["/data"]:
-                        tavi_file.create_group("/data/" + dataset_name)
-                    scan_file.copy(
-                        source=scan_file["/" + scan_name], dest=tavi_file["/data/" + dataset_name], expand_soft=True
-                    )
-                    nexus_entry = scan_file[scan_name]
-                    _, scan_entry = Scan.from_nexus_entry(nexus_entry)
-                scans.update({scan_name: scan_entry})
-        self.data.update({dataset_name: scans})
+        scan_dict = {}
+        for scan_file in scan_list:
+            scan = NexusEntry.from_nexus(path_to_hdf5_folder + scan_file)
+            scan_dict.update(scan)
+        self.data = TAVI.load_data(scan_dict)
 
-    # TODO
     def load_spice_data_from_disk(self, path_to_spice_folder):
         """Load hdf5 data from path_to_hdf5.
 
         Args:
             path_to_spice_folder (str): path to spice folder
-            OVERWRITE (bool): overwrite exsiting data if Ture, oterwise append new scans in data.
-                Do not change processed_data, fit or plot
         """
-
-        self.load_data()
+        scan_dict = NexusEntry.from_spice(path_to_spice_folder)
+        self.data = TAVI.load_data(scan_dict)
 
     # TODO
-    def load_data_from_oncat(self, user_credentials, ipts_info, OVERWRITE=True):
+    def load_data_from_oncat(self, user_credentials, ipts_info):
         """Load data from ONCat based on user_credentials and ipts_info.
 
         Args:
             user_credentials (str): username and password.
             scipts_infoans (str): ipts number and exp number.
-            OVERWRITE (bool): overwrite exsiting data if Ture, oterwise append.
-                Do not change processed_data, fit or plot
         """
-        self.load_data()
-
-    # TODO load processed_data
-    def load_data(self):
-        """Load data, processed_data, fits and plots in a tavi file into memory"""
+        pass
 
+    def get_exp_id(self) -> list[str]:
+        "return a list of exp id"
+        exp_id_list = []
         with h5py.File(self.file_path, "r") as tavi_file:
-            for dataset_name in (tavi_data := tavi_file["data"]):
-                scans = {}
-                for scan_name in (dataset := tavi_data[dataset_name]):
-                    nexus_entry = dataset[scan_name]
-                    _, scan_entry = Scan.from_nexus_entry(nexus_entry)
-                    scans.update({scan_name: scan_entry})
-                self.data.update({dataset_name: scans})
+            for exp_id in tavi_file["data"].keys():
+                exp_id_list.append(exp_id)
+        return exp_id_list
 
     def open_file(self, tavi_file_path):
-        """Open existing tavi file"""
-        # TODO validate path
+        """Open an existing tavi file"""
+
+        if not os.path.exists(tavi_file_path):
+            raise ValueError(f"TAVI file does not exist at path {tavi_file_path}")
+
         self.file_path = tavi_file_path
-        self.load_data()
 
+        # load data
+        data_dict = {}
+        exp_id_list = self.get_exp_id()
+        for exp_id in exp_id_list:
+            scans = NexusEntry.from_nexus(tavi_file_path, prefix=f"/data/{exp_id}")
+            data_dict.update({exp_id: scans})
+        self.data = data_dict
         # TODO load processed_data fits, and plots
 
-    def save_file(self, path_to_hdf5: Optional[str] = None):
+    def save(self, file_path: Optional[str] = None):
         """Save current data to a hdf5 on disk at path_to_hdf5
 
         Args:
             path_to_hdf5 (str): path to hdf5 data file, ends with '.h5'
         """
-        if path_to_hdf5 is not None:
-            self.file_path = path_to_hdf5
-
-        try:  # mode a: Read/write if exists, create otherwise
-            with h5py.File(self.file_path, "a", track_order=True) as tavi_file:
-                if "/data" not in tavi_file:
-                    tavi_file.create_group("data", track_order=True)
-                pass
 
-                if "/processed_data" not in tavi_file:
-                    tavi_file.create_group("processed_data", track_order=True)
-
-                # TODO save processed_data fits and plots
-                if "/fits" not in tavi_file:
-                    tavi_file.create_group("fits", track_order=True)
-
-                if "/plots" not in tavi_file:
-                    tavi_file.create_group("plots", track_order=True)
+        if file_path is None:
+            current_path = os.getcwd()
+            self.file_path = current_path + "/tavi_temp.h5"
+        else:
+            self.file_path = file_path
 
+        try:
+            with h5py.File(self.file_path, "w", track_order=True) as root:
+                grp_data = root.create_group("data", track_order=True)
+                for exp_id, scans in self.data.items():
+                    grp_data.create_group(name=exp_id, track_order=True)
+                    for scan_name, nxentry in scans.items():
+                        name = f"/data/{exp_id}/{scan_name}"
+                        nxentry.to_nexus(file_path, name=name)
+                # TODO save pdata, fits and plots
+                grp_pdata = root.create_group("processed_data", track_order=True)
+                grp_fits = root.create_group("fits", track_order=True)
+                grp_plots = root.create_group("plots", track_order=True)
         except OSError:
             print(f"Cannot create tavi file at {self.file_path}")
 

tests/test_scan_group.py

@@ -15,7 +15,7 @@ def test_contour_plot():
         tavi.load_nexus_data_from_disk(nexus_data_folder)
         tavi.load_data()
     else:
-        tavi.open_file(tavi_file_name)
+        tavi.open(tavi_file_name)
     dataset = tavi.data["IPTS32124_CG4C_exp0424"]
 
     scan_list = list(range(42, 49, 1)) + list(range(70, 76, 1))

tests/test_tavi_data.py

@@ -7,71 +7,58 @@
 
 
 def test_new_tavi_file_at_given_path():
+    file_path = "./test_data/tavi_empty_test.h5"
     tavi = TAVI()
-    tavi_file_name = "./test_data/tavi_empty_test.h5"
-    tavi.new_file(tavi_file_name)
-    tavi.save_file()
+    tavi.save(file_path)
 
-    with h5py.File(tavi_file_name, "r") as f:
+    with h5py.File(file_path, "r") as f:
         keys = [key for key in f["/"].keys()]
     # check if groups preserves the order
     assert keys[0] == "data"
     assert keys[1] == "processed_data"
     assert keys[2] == "fits"
     assert keys[3] == "plots"
 
-    # delete file
-    os.remove(tavi_file_name)
+    os.remove(file_path)  # delete file
 
 
 def test_new_tavi_file_without_path():
     tavi = TAVI()
-    tavi.new_file()
-    tavi.save_file()
+    tavi.save()
 
+    assert tavi.file_path.split("/")[-1] == "tavi_temp.h5"
     with h5py.File(tavi.file_path, "r") as f:
         keys = [key for key in f["/"].keys()]
-    # check if groups preserves the order
     assert keys[1] == "processed_data"
-    # delete file
-    os.remove(tavi.file_path)
+    os.remove(tavi.file_path)  # delete file
 
 
-# TODO
 def test_load_spice_data_from_disk():
     tavi = TAVI()
-    tavi_file_name = "./test_data/tavi_test_exp424.h5"
-    tavi.new_file(tavi_file_name)
+    tavi.load_spice_data_from_disk("./test_data/exp424/")
+    assert len(tavi.data["IPTS32124_CG4C_exp0424"]) == 92
 
 
 def test_load_nexus_data_from_disk():
     tavi = TAVI()
-    tavi_file_name = "./test_data/tavi_test_exp424.h5"
-    tavi.new_file(tavi_file_name)
-
-    nexus_data_folder = "./test_data/IPTS32124_CG4C_exp0424"
-    tavi.load_nexus_data_from_disk(nexus_data_folder)
+    tavi.load_nexus_data_from_disk("./test_data/IPTS32124_CG4C_exp0424/")
     assert len(tavi.data["IPTS32124_CG4C_exp0424"]) == 92
 
-    # os.remove(tavi.file_path)
-
 
-def test_save_tavi_file():
+def test_save():
     tavi = TAVI()
-    tavi_file_name = "./test_data/tavi_test_exp424.h5"
-    tavi.new_file(tavi_file_name)
-    nexus_data_folder = "./test_data/IPTS32124_CG4C_exp0424"
-    tavi.load_nexus_data_from_disk(nexus_data_folder)
-    tavi.save_file()
+    path_to_spice_folder = "./test_data/exp424"
+    tavi.load_spice_data_from_disk(path_to_spice_folder)
+    tavi.save("./test_data/tavi_test_exp424.h5")
 
-    with h5py.File(tavi.file_path, "r") as tavi_file:
-        assert len(tavi_file["/data/IPTS32124_CG4C_exp0424"]) == 92
+    with h5py.File(tavi.file_path, "r") as f:
+        len(f["/data/IPTS32124_CG4C_exp0424"]) == 92
+
+    # os.remove(tavi.file_path)
 
 
 def test_open_tavi_file():
     tavi = TAVI()
-    tavi_file_name = "./test_data/tavi_test_exp424.h5"
-    tavi.open_file(tavi_file_name)
+    tavi.open_file("./test_data/tavi_test_exp424.h5")
 
-    assert len(scans := tavi.data["IPTS32124_CG4C_exp0424"]) == 92
-    assert scans["scan0001"].scan_info.scan_num == 1
+    assert len(tavi.data["IPTS32124_CG4C_exp0424"]) == 92