modified the fit class

scripts/HoV6Sn6_contour.py

@@ -14,7 +14,7 @@
 tavi.new_file(tavi_file_name)
 
 nexus_file_name = "./test_data/nexus_exp1031.h5"
-tavi.get_nexus_data_from_disk(nexus_file_name)
+tavi.load_nexus_data_from_disk(nexus_file_name)
 dataset = tavi.data["IPTS32912_HB1A_exp1031"]
 
 # -------------- 001 s1 scans ------------

scripts/MnTe_contour.py

@@ -15,7 +15,7 @@
 
 tavi_file_name = "./test_data/tavi_test_exp0813.h5"
 tavi.new_file(tavi_file_name)
-tavi.get_nexus_data_from_disk(nexus_folder)
+tavi.load_nexus_data_from_disk(nexus_folder)
 dataset = tavi.data["IPTS34735_HB3_exp0813"]
 
 # -------------- H0L const Q scans ------------

scripts/test_fit_group.py

@@ -45,7 +45,7 @@ def test_fit_group(tavi):
     tavi.new_file(tavi_file_name)
 
     nexus_file_name = "./test_data/nexus_exp1031.h5"
-    tavi.get_nexus_data_from_disk(nexus_file_name)
+    tavi.load_nexus_data_from_disk(nexus_file_name)
     dataset = tavi.data["IPTS32912_HB1A_exp1031"]
 
     test_fit_group(tavi)

scripts/test_scan.py

@@ -25,7 +25,7 @@ def test_scan_group_contour_exp710():
     tavi.new_file(tavi_file_name)
 
     nexus_file_name = "./tests/test_data_folder/nexus_exp710.h5"
-    tavi.get_nexus_data_from_disk(nexus_file_name)
+    tavi.load_nexus_data_from_disk(nexus_file_name)
 
     scan_list = (
         [tavi.data[f"scan{i:04}"] for i in range(214, 225, 1)]

src/tavi/data/fit.py

@@ -1,9 +1,16 @@
+from typing import Optional
+
 import numpy as np
 from lmfit import Parameters, models
 
+from tavi.data.plotter import Plot1D
+
+
+class Fit1D(object):
+    """Fit a 1d curve
 
-class Fit(object):
-    """Save information about fits"""
+    Attributes:
+        NUM_PTS (int): number of points for the fit curve"""
 
     models = {
         # ---------- peak models ---------------
@@ -25,47 +32,35 @@ class Fit(object):
         "Spline": models.SplineModel,
     }
 
-    def __init__(self, x, y, err=None, fit_range=None):
-        """
-        initialize a fit model
+    def __init__(self, plot1d: Plot1D):
+        """initialize a fit model"""
+        self.NUM_PTS: int = 100
+        self.x = plot1d.x
+        self.y = plot1d.y
+        self.yerr = plot1d.yerr
 
-        Args:
-            x (list)
-            y (list)
-            err (list | None)
-            fit_range (tuple)
-            NUM_PTS (int): number of points for the fit curve
-        """
-        self.NUM_PTS = 100
-        self.range = fit_range
-        self.x = np.array(x)
-        self.y = np.array(y)
-        self.err = err
-
-        # trim the range
-        if fit_range is not None:
-            fit_min, fit_max = fit_range
-            mask = np.bitwise_and(x >= fit_min, x <= fit_max)
-            self.x = self.x[mask]
-            self.y = self.y[mask]
-            if self.err is not None:
-                self.err = self.err[mask]
-
-        self.x_plot = np.linspace(
-            self.x.min(),
-            self.x.max(),
-            num=self.NUM_PTS,
-        )
-        self.y_plot = None
-
-        self.background_models = []
-        self.signal_models = []
+        self.background_models: models = []
+        self.signal_models: models = []
         self.pars = Parameters()
         self.num_backgrounds = 0
         self.num_signals = 0
         self.fit_result = None
 
         self.PLOT_SEPARATELY = False
+        self.fit_plot: Optional[Plot1D] = None
+
+    def set_range(self, fit_min, fit_max):
+        """set the range used for fitting"""
+
+        mask = np.bitwise_and(self.x >= fit_min, self.x <= fit_max)
+        self.x = self.x[mask]
+        self.y = self.y[mask]
+        if self.yerr is not None:
+            self.yerr = self.yerr[mask]
+
+    @property
+    def x_plot(self):
+        return np.linspace(self.x.min(), self.x.max(), num=self.NUM_PTS)
 
     def add_background(
         self,
@@ -94,7 +89,7 @@ def add_background(
             prefix = f"b{self.num_backgrounds}_"
         else:
             prefix = ""
-        model = Fit.models[model](prefix=prefix, nan_policy="propagate")
+        model = Fit1D.models[model](prefix=prefix, nan_policy="propagate")
         param_names = model.param_names
         # guess initials
         pars = model.guess(self.y, x=self.x)
@@ -150,7 +145,7 @@ def add_signal(
         """
         self.num_signals += 1
         prefix = f"s{self.num_signals}_"
-        model = Fit.models[model](prefix=prefix, nan_policy="propagate")
+        model = Fit1D.models[model](prefix=prefix, nan_policy="propagate")
         param_names = model.param_names
         # guess initials
         pars = model.guess(self.y, x=self.x)
@@ -184,19 +179,19 @@ def add_signal(
             self.pars.add(pars[param_name])
         self.signal_models.append(model)
 
-    def perform_fit(self):
+    def perform_fit(self) -> None:
         model = np.sum(self.signal_models)
 
         if self.num_backgrounds > 0:
             model += np.sum(self.background_models)
 
-        if self.err is None:
+        if self.yerr is None:
             out = model.fit(self.y, self.pars, x=self.x)
         else:
-            out = model.fit(self.y, self.pars, x=self.x, weights=self.err)
+            out = model.fit(self.y, self.pars, x=self.x, weights=self.yerr)
 
         self.result = out
         self.y_plot = model.eval(out.params, x=self.x_plot)
 
-        fit_report = out.fit_report(min_correl=0.25)
-        return fit_report
+        self.fit_report = out.fit_report(min_correl=0.25)
+        self.fit_plot = Plot1D(x=self.x_plot, y=self.y_plot)

src/tavi/data/tavi.py

@@ -52,7 +52,7 @@ def new_file(self, file_path: Optional[str] = None) -> None:
         except OSError:
             print(f"Cannot create tavi file at {self.file_path}")
 
-    def get_nexus_data_from_disk(self, path_to_hdf5_folder):
+    def load_nexus_data_from_disk(self, path_to_hdf5_folder):
         """Copy hdf5 data from path_to_hdf5_folder into tavi."""
         # validate path
         if path_to_hdf5_folder[-1] != "/":
@@ -63,6 +63,7 @@ def get_nexus_data_from_disk(self, path_to_hdf5_folder):
         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
@@ -73,11 +74,13 @@ def get_nexus_data_from_disk(self, path_to_hdf5_folder):
                     scan_file.copy(
                         source=scan_file["/" + scan_name], dest=tavi_file["/data/" + dataset_name], expand_soft=True
                     )
-
-        self.load_data()
+                    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})
 
     # TODO
-    def get_spice_data_from_disk(self, path_to_spice_folder):
+    def load_spice_data_from_disk(self, path_to_spice_folder):
         """Load hdf5 data from path_to_hdf5.
 
         Args:
@@ -89,7 +92,7 @@ def get_spice_data_from_disk(self, path_to_spice_folder):
         self.load_data()
 
     # TODO
-    def get_data_from_oncat(self, user_credentials, ipts_info, OVERWRITE=True):
+    def load_data_from_oncat(self, user_credentials, ipts_info, OVERWRITE=True):
         """Load data from ONCat based on user_credentials and ipts_info.
 
         Args:
@@ -100,8 +103,9 @@ def get_data_from_oncat(self, user_credentials, ipts_info, OVERWRITE=True):
         """
         self.load_data()
 
+    # TODO load processed_data
     def load_data(self):
-        """Load tavi data into memory"""
+        """Load data, processed_data, fits and plots in a tavi file into memory"""
 
         with h5py.File(self.file_path, "r") as tavi_file:
             for dataset_name in (tavi_data := tavi_file["data"]):
@@ -110,7 +114,6 @@ def load_data(self):
                     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})
 
     def open_file(self, tavi_file_path):
@@ -121,7 +124,6 @@ def open_file(self, tavi_file_path):
 
         # TODO load processed_data fits, and plots
 
-    # TODO
     def save_file(self, path_to_hdf5: Optional[str] = None):
         """Save current data to a hdf5 on disk at path_to_hdf5
 
@@ -131,37 +133,24 @@ def save_file(self, path_to_hdf5: Optional[str] = None):
         if path_to_hdf5 is not None:
             self.file_path = path_to_hdf5
 
-        # TODO save processed_data fits and plots
-
-    def delete_data_entry(self, entry_id):
-        """Delete a date entry from file based on entry_id.
-
-        If an data entry if deleted, also delete the fits, scan_groups, plots
-        that contains the conresponding entry.
-
-        Args:
-            entry_id (str): ID of entry to be deteled.
-        """
-        pass
-
-    def select_entries(self, conditions):
-        """Return a list of data entry IDs which satisfis the conditions.
-
-        All data points in a scan have to satisfy all conditions.
+        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
 
-        Args:
-            conditions (dict): scan[key] equals value
-                or in the range of (vaule[0], value[1])
-                for example: {"IPTS":1234, "temperature":(2,5)}
+                if "/processed_data" not in tavi_file:
+                    tavi_file.create_group("processed_data", track_order=True)
 
-        Returns:
-            tuple: entry IDs
+                # TODO save processed_data fits and plots
+                if "/fits" not in tavi_file:
+                    tavi_file.create_group("fits", track_order=True)
 
-        """
-        pass
+                if "/plots" not in tavi_file:
+                    tavi_file.create_group("plots", track_order=True)
 
-    def get_selected(self):
-        return "scan0001"
+        except OSError:
+            print(f"Cannot create tavi file at {self.file_path}")
 
     def generate_scan_group(
         self,

tests/test_fit.py

@@ -2,8 +2,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from tavi.data.fit import Fit
-from tavi.data.plotter import Plot1DManager
+from tavi.data.fit import Fit1D
 from tavi.data.scan import Scan
 
 
@@ -13,14 +12,16 @@ def test_fit_single_peak():
     _, s1 = Scan.from_nexus_file(nexus_file_name)
 
     plot1d = s1.generate_curve(norm_channel="mcu", norm_val=30)
-    f1 = Fit(x=plot1d.x, y=plot1d.y, err=plot1d.yerr, fit_range=(0.5, 4.0))
+    f1 = Fit1D(plot1d)
+    f1.set_range(0.5, 4.0)
     f1.add_background(values=(0.7,))
     f1.add_signal(values=(None, 3.5, None), vary=(True, True, True))
-    fit_report = f1.perform_fit()
-    # assert np.allclose(f1.result.params["s1_center"].value, 3.54, atol=0.01)
-    # assert np.allclose(f1.result.params["s1_fwhm"].value, 0.39, atol=0.01)
+    f1.perform_fit()
+    assert np.allclose(f1.result.params["s1_center"].value, 3.54, atol=0.01)
+    assert np.allclose(f1.result.params["s1_fwhm"].value, 0.39, atol=0.01)
     fig, ax = plt.subplots()
     plot1d.plot_curve(ax)
+    f1.fit_plot.plot_curve(ax)
     plt.show()
 
 
@@ -29,8 +30,9 @@ def test_fit_two_peak():
     nexus_file_name = "./test_data/IPTS32124_CG4C_exp0424/scan0042.h5"
     _, s1 = Scan.from_nexus_file(nexus_file_name)
 
-    x, y, yerr, xlabel, ylabel, title, label = s1.generate_curve(norm_channel="mcu", norm_val=30)
-    f1 = Fit(x=x, y=y, err=yerr, fit_range=(0.0, 4.0))
+    plot1d = s1.generate_curve(norm_channel="mcu", norm_val=30)
+    f1 = Fit1D(plot1d)
+    f1.set_range(0.0, 4.0)
     f1.add_background(values=(0.7,))
     f1.add_signal(values=(None, 3.5, 0.29), vary=(True, True, True))
     f1.add_signal(
@@ -40,11 +42,11 @@ def test_fit_two_peak():
         mins=(0, 0, 0.1),
         maxs=(None, 0.1, 0.3),
     )
-    fit_report = f1.perform_fit()
+    f1.perform_fit()
     assert np.allclose(f1.result.params["s1_center"].value, 3.54, atol=0.01)
     assert np.allclose(f1.result.params["s1_fwhm"].value, 0.40, atol=0.01)
 
-    p1 = Plot1DManager()
-    p1.plot_curve(x, y, yerr, xlabel, ylabel, title, label)
-    p1.plot_curve(f1.x_plot, f1.y_plot, fmt="-", xlabel=xlabel, ylabel=ylabel, title=title)
+    fig, ax = plt.subplots()
+    plot1d.plot_curve(ax)
+    f1.fit_plot.plot_curve(ax)
     plt.show()

tests/test_scan.py

@@ -12,7 +12,7 @@ def test_load_scan_from_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.get_nexus_data_from_disk(nexus_data_folder)
+    tavi.load_nexus_data_from_disk(nexus_data_folder)
 
     with h5py.File(tavi.file_path, "r") as tavi_file:
         dataset_name, scan = Scan.from_nexus_entry(tavi_file["/data/IPTS32124_CG4C_exp0424/scan0042"])

tests/test_scan_group.py

@@ -12,7 +12,7 @@ def test_contour_plot():
     tavi_file = Path(tavi_file_name)
     if not tavi_file.is_file():
         nexus_data_folder = "./test_data/IPTS32124_CG4C_exp0424"
-        tavi.get_nexus_data_from_disk(nexus_data_folder)
+        tavi.load_nexus_data_from_disk(nexus_data_folder)
         tavi.load_data()
     else:
         tavi.open_file(tavi_file_name)