reorganized repo

src/tavi/__init__.py

@@ -2,6 +2,7 @@
 Contains the entry point for the application
 """
 
+# TODO template
 try:
     from ._version import __version__  # noqa: F401
 except ImportError:

src/tavi/data/fit.py

@@ -0,0 +1,155 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from lmfit import models
+
+
+class Fit(object):
+    """Save information about fits"""
+
+    models = {
+        # ---------- peak models ---------------
+        "Gaussian": models.GaussianModel,
+        "Lorentzian": models.LorentzianModel,
+        "Voigt": models.VoigtModel,
+        "PseudoVoigt": models.PseudoVoigtModel,
+        "DampedOscillator": models.DampedOscillatorModel,
+        "DampedHarmonicOscillator": models.DampedHarmonicOscillatorModel,
+        # ---------- background models ------------
+        "Constant": models.ConstantModel,
+        "Linear": models.LinearModel,
+        "Quadratic": models.QuadraticModel,
+        "Polynomial": models.PolynomialModel,
+        "Exponential": models.ExponentialModel,
+        "PowerLaw": models.PowerLawModel,
+        # --------- expression ---------------
+        "Expression": models.ExpressionModel,
+        "Spline": models.SplineModel,
+    }
+
+    def __init__(self, x, y, err=None, fit_range=None):
+        """
+        initialize a fit model
+
+        Args:
+            x (list)
+            y (list)
+            err (list | None)
+            fit_range (tuple)
+        """
+
+        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.background_models = []
+        self.signal_models = []
+        self.num_backgrounds = 0
+        self.num_signals = 0
+        self.FIT_STATUS = None
+        self.chi_squred = 0
+        self.PLOT_SEPARATELY = False
+
+    def add_background(
+        self,
+        model="Constant",
+        p0=None,
+        min=None,
+        max=None,
+        fixed=None,
+        expr=None,
+    ):
+        """Set the model for background
+
+        Args:
+            model (str): Constant, Linear, Quadratic, Polynomial,
+                         Exponential, PowerLaw
+            p0 (tuple | None): inital parameters
+            min (tuple | None): minimum
+            max (tuple | None): maximum
+            fixed (tuple | None): tuple of flags
+            expr (tuple| None ): constraint expressions
+        """
+        self.num_backgrounds += 1
+
+        # add prefix if more than one background
+        if self.num_backgrounds > 1:
+            prefix = f"b{self.num_backgrounds}_"
+        else:
+            prefix = ""
+        model = Fit.models[model](prefix=prefix, nan_policy="propagate")
+
+        pass
+
+        # pars = model.guess(self.y, x=self.x)
+        # pars["c"].set(value=0.7, vary=True, expr="")
+
+        self.background_models.append(model)
+
+    def add_signal(
+        self,
+        model="Gaussian",
+        p0=None,
+        min=None,
+        max=None,
+        expr=None,
+    ):
+        """Set the model for background
+
+        Args:
+            model (str): Constant, Linear, Quadratic, Polynomial,
+                         Exponential, PowerLaw
+            p0 (tuple | None): inital parameters
+            min (tuple | None): minimum
+            max (tuple | None): maximum
+            expr (str| None ): constraint expression
+        """
+        self.num_signals += 1
+        prefix = f"s{self.num_signals}_"
+        model = Fit.models[model](prefix=prefix, nan_policy="propagate")
+        print(model.param_names)
+        self.signal_models.append(model)
+        pars = model.guess(self.y, x=self.x)
+        pars["c"].set(value=0.7, vary=True, expr="")
+
+    def perform_fit(self):
+
+        model = np.sum(self.signal_models)
+
+        if self.num_backgrounds > 0:
+            model += np.sum(self.background_models)
+
+        if self.err is None:
+            # pars = model.guess(self.y, x=self.x)
+            out = model.fit(self.y, pars, x=self.x)
+        else:
+            pars = model.fit(self.y, x=self.x, weights=self.err)
+
+        # out = model.fit(self.y, pars, x=self.x)
+        print(out.fit_report(min_correl=0.25))
+
+
+# # plot fitting results
+# fig, ax = plt.subplots()
+# if std is None:
+#     ax.plot(self.scan_data[x_str], self.scan_data[y_str], "o")
+# else:
+#     ax.errorbar(x, y, yerr=std, fmt="o")
+# ax.plot(x, out.best_fit, "-")
+
+# if "scan_title" in self.scan_params:
+#     ax.set_title(self.scan_params["scan_title"])
+# ax.set_xlabel(x_str)
+# ax.set_ylabel(y_str)
+# ax.grid(alpha=0.6)
+# ax.set_ylim(bottom=0)
+# plt.tight_layout()

src/tavi/data/nexus_reader.py

@@ -0,0 +1,127 @@
+import numpy as np
+
+
+def nexus_to_dict(nexus_entry):
+    """Reads a nexus entry, convert to dictionaries of data and meta_data
+
+    Args:
+        nexus entry
+
+    Returns:
+        meta_data (dict)
+        data (dict)
+    """
+
+    def dataset_to_string(ds):
+        return str(ds.asstr()[...])
+
+    scan_info = {
+        "scan": int(nexus_entry.name[-4:]),  # last 4 digits are scan number
+        "time": dataset_to_string(nexus_entry["start_time"]),
+        "scan_title": dataset_to_string(nexus_entry["title"]),
+        # "preset_type": "normal",
+        "preset_channel": dataset_to_string(nexus_entry["monitor/mode"]),
+        "preset_value": float(nexus_entry["monitor/preset"][...]),
+        "def_y": nexus_entry["data"].attrs["signal"],
+        "def_x": nexus_entry["data"].attrs["axes"],
+    }
+
+    sample_ub_info = {
+        "samplename": dataset_to_string(nexus_entry["sample/name"]),
+        "lattice_constants": nexus_entry["sample/unit_cell"][...],
+        "ub_matrix": np.reshape(nexus_entry["sample/orientation_matrix"][...], (3, 3)),
+        # "mode": # UB mode
+        "palne_normal": nexus_entry["sample/plane_normal"][...],
+        "ubconf": dataset_to_string(nexus_entry["sample/ub_conf"]),
+    }
+
+    instrument = dataset_to_string(nexus_entry["instrument/name"])
+
+    instrument_info = {
+        "instrument": instrument,
+        # "monochromator":
+        # "analyzer":
+        "sense": dataset_to_string(nexus_entry["instrument/monochromator/sense"])
+        + dataset_to_string(nexus_entry["sample/sense"])
+        + dataset_to_string(nexus_entry["instrument/analyser/sense"]),
+        "collimation": nexus_entry["instrument/collimator/divergence_x"][...],
+    }
+
+    num = np.size(nexus_entry["data/" + scan_info["def_x"]])
+
+    data = {
+        "Pt.": np.arange(start=1, stop=num + 1, step=1),
+        # detector
+        "detector": nexus_entry["instrument/detector/data"][...],
+        # monitor
+        "time": nexus_entry["monitor/time"][...],
+        "monitor": nexus_entry["monitor/monitor"][...],
+        "mcu": nexus_entry["monitor/mcu"][...],
+        # sample
+        "s1": nexus_entry["sample/s1"][...],
+        "s2": nexus_entry["sample/s2"][...],
+        "sgl": nexus_entry["sample/sgl"][...],
+        "sgu": nexus_entry["sample/sgu"][...],
+        "stl": nexus_entry["sample/stl"][...],
+        "stu": nexus_entry["sample/stu"][...],
+        "q": nexus_entry["sample/q"][...],
+        "qh": nexus_entry["sample/qh"][...],
+        "qk": nexus_entry["sample/qk"][...],
+        "ql": nexus_entry["sample/ql"][...],
+        "en": nexus_entry["sample/en"][...],
+    }
+
+    # analyser
+    ana_str = (
+        ("a1", "a2", "afocus", "ef") + tuple([f"qm{i+1}" for i in range(8)]) + tuple([f"xm{i+1}" for i in range(8)])
+    )
+    nexus_ana_str = nexus_entry["instrument/analyser"].keys()
+    for key in ana_str:
+        if key in nexus_ana_str:
+            data.update({key: nexus_entry["instrument/analyser/" + key][...]})
+
+    # monochromator
+    mono_str = ("m1", "m2", "ei", "focal_length", "mfocus", "marc", "mtrans")
+    nexus_mono_str = nexus_entry["instrument/monochromator"].keys()
+    for key in mono_str:
+        if key in nexus_mono_str:
+            data.update({key: nexus_entry["instrument/monochromator/" + key][...]})
+
+    # slits
+    slits_str1 = ("bat", "bab", "bal", "bar", "bbt", "bbb", "bbl", "bbr")
+    slits_str2 = ("slita_lf", "slita_rt", "slita_tp", "slitb_bt", "slitb_lf", "slitb_rt", "slitb_tp")
+    slits_str3 = ("slit_pre_bt", "slit_pre_lf", "slit_pre_rt", "slit_pre_tp")
+    slit_str = slits_str1 + slits_str2 + slits_str3
+
+    nexus_slit_str = nexus_entry["instrument/slit"].keys()
+    for key in slit_str:
+        if key in nexus_slit_str:
+            data.update({key: nexus_entry["instrument/slit/" + key][...]})
+
+    # temprature
+    temperatue_str = (
+        "coldtip",
+        "tsample",
+        "temp_a",
+        "temp_2",
+        "vti",
+        "sample",
+        "temp",
+        "dr_tsample",
+        "dr_temp",
+    )
+    for t in nexus_entry["sample"].keys():
+        if t in temperatue_str:
+            data.update({t: nexus_entry["sample/" + t][...]})
+
+    return scan_info, sample_ub_info, instrument_info, data
+
+
+def nexus_to_SPICE(nexus_entry):
+    """Reads a nexus entry, convert to a SPICE scan file
+
+    Args:
+        nexus entry
+
+    """
+    pass