Stopped on #52, #51, #50,#49,#48, #47

notebooks/FourTracks/analysis/fit.py

@@ -1,15 +1,18 @@
+import lmfit
+import cmath
+from lmfit.parameter import Parameters
 import numpy as np
 import pandas as pd
 from particle import Particle
 from typing import Union
+from lmfit import Model
+from lmfit.model import ModelResult
 
 Rho0 = Particle.from_pdgid(113)
 Pi0 = Particle.from_pdgid(111)
 PiPlus = Particle.from_pdgid(211)
 pi_pl_mass = PiPlus.mass / 1000
 
-bckg_y = None
-
 
 def bw(
     x: Union[np.ndarray, pd.Series], M: float, G: float, amp: float
@@ -30,33 +33,97 @@ def bw(
     return fff
 
 
-def bw_bckg(
-    x: Union[np.ndarray, pd.Series], M: float, G: float, amp: float, amp_bckg: float
-) -> Union[np.ndarray, pd.Series]:
-    return bw(x=x, M=M, G=G, amp=amp) + amp_bckg * bckg_y
+def polyn(x, c0, c1, c2, c3, c4, c5, c6, c7, amp):
+    """
+    c - array of polynomial coefficients (including coefficients equal to zero) from highest degree to the constant term
+    """
+    return amp * np.polyval([c7, c6, c5, c4, c3, c2, c1, c0], x)
 
 
-def sod(
-    x: Union[np.ndarray, pd.Series], M: float, G: float, amp: float, bterm: float
-) -> Union[np.ndarray, pd.Series]:
+# def ff(
+#     x, y, functions: list, params: Parameters, numpoints=100, **fit_kws
+# ) -> tuple(ModelResult, pd.DataFrame):
+#     """
+#     x,
+#     y,
+#     functions - list of functions used in Model,
+#     params - initial Parameters for Model,
+#     numpoints - number of points for model evaluation
+#     Returns:
+#     pd.Dataframe with initial data and evaluation of fit functions
+#     """
+#     mod = Model(functions[0], prefix=functions[0].__name__ + "_")
+
+#     for f in functions[1:]:
+#         mod += Model(f, prefix=f.__name__ + "_")
+
+#     result = mod.fit(y, params, x=x, weights=1 / np.sqrt(y), *fit_kws)
+#     xx = np.linspace(x.min(), x.max(), numpoints)
+#     xx = np.linspace(x.min(), x.max(), 100)
+
+#     ff = result.eval(x=xx)
+
+#     fbw1 = fit.bw(
+#         x=xx,
+#         M=result.best_values["bw1_M"],
+#         G=result.best_values["bw1_G"],
+#         amp=result.best_values["bw1_amp"],
+#     )
+#     fbw2 = fit.bw(
+#         x=xx,
+#         M=result.best_values["bw2_M"],
+#         G=result.best_values["bw2_G"],
+#         amp=result.best_values["bw2_amp"],
+#     )
+#     fbckg = fit.polyn(
+#         xx,
+#         result.best_values["bckg_c0"],
+#         result.best_values["bckg_c1"],
+#         result.best_values["bckg_c2"],
+#         result.best_values["bckg_c3"],
+#         result.best_values["bckg_c4"],
+#         0,
+#         0,
+#         0,
+#         result.best_values["bckg_amp"],
+#     )
+#     df = pd.DataFrame([xx, ff, fbw1, fbw2, fbckg]).T
+#     df.columns = ["x", "fit", "bw1", "bw2", "bkgr"]
+#     df = pd.DataFrame([xx, ff, fbw, fbckg]).T
+#     df.columns = ["x", "fit", "bw", "bckg"]
+#     return pd.DataFrame([])
+
 
-    # rho mass-dependent width
-    val = ((x * 2 - 4 * pi_pl_mass ** 2) / (M ** 2 - 4 * pi_pl_mass ** 2)) ** 1.5
+def fTerm(
+    x: Union[np.ndarray, pd.Series], M: float, G: float, amp: float
+) -> Union[np.ndarray, pd.Series]:
+    """
+    """
+    # rho mass-dependent  width
+    val = ((x ** 2 - 4 * pi_pl_mass ** 2) / (M ** 2 - 4 * pi_pl_mass ** 2)) ** 1.5
     mdwidth = G * (M / x) * val  # mass dependent width
-    ix = complex(0, 1)
-    denom = x ** 2 - M ** 2 + ix * M ** 2 * mdwidth
-    aterm = amp * np.sqrt(x * M * mdwidth) / denom
-    total = (aterm + bterm) ** 2
-    return np.abs(total)
+
+    denom = (x ** 2 - M ** 2) + (M * mdwidth * 1j)
+    return amp * np.sqrt(x * M * mdwidth) / denom
 
 
-def sod_bckg(
+def fInterBWs(
     x: Union[np.ndarray, pd.Series],
-    M: float,
-    G: float,
+    M1: float,
+    G1: float,
+    amp1: float,
+    M2: float,
+    G2: float,
+    amp2: float,
     amp: float,
-    bterm: float,
-    amp_bckg: float,
 ) -> Union[np.ndarray, pd.Series]:
-    return sod(x=x, M=M, G=G, amp=amp, bterm=bterm) + amp_bckg * bckg_y
+    """
+    """
+    aterm = fTerm(x, M1, G1, amp1)
+    bterm = fTerm(x, M2, G2, amp2)
+
+    amptot = np.abs((aterm + bterm) ** 2)
+    amptwo = np.abs(aterm) ** 2 + np.abs(bterm) ** 2
+
+    return amp * (amptot - amptwo)
 

notebooks/FourTracks/data/cuts/kinematics.py

@@ -6,4 +6,4 @@ def GetTracksWithPtCut(
     tracks: pd.DataFrame, maxPt: float = 0.15, minPt: float = 0
 ) -> pd.Series:
     pt = pt_events(tracks)
-    return tracks.loc[pt[(pt < maxPt) & (pt >= minPt)].reset_index().entry]
+    return tracks.loc[pt[(pt < maxPt) & (pt >= minPt)].index]

notebooks/FourTracks/data/cuts/trigger.py

@@ -5,17 +5,24 @@
 
 
 def GetTriggeredEvents(
-    tracks: pd.DataFrame, fileName: str = "", saveToFile: bool = False
+    tracks: pd.DataFrame,
+    dataPath: str = r"D:\GoogleDrive\Job\cern\Alice\analysis\data\RhoPrime",
+    year: int = 2015,
+    fileName: str = "",
+    rewrite: bool = False,
 ) -> pd.DataFrame:
     """
-    Method get events with four tracks and check fake trigger events
+    Method gets events with four tracks and checks fake trigger events
     """
 
-    pth = join(r"D:\GoogleDrive\Job\cern\Alice\analysis\data\RhoPrime\2015", fileName)
+    pth = join(dataPath, str(year), fileName)
 
-    if exists(pth) and fileName:
+    if exists(pth) and not rewrite and fileName:
         return pd.read_parquet(pth)
 
+    if not fileName:
+        raise NameError("The name of file should not be empty for saving purposes")
+
     # fired FORs numbers for 4 zq tracks
     for_sensors = pd.DataFrame(
         orig_events.FORChip[pd.unique(tracks.reset_index().entry)]
@@ -73,9 +80,14 @@ def GetTriggeredEvents(
     # df_dbg = df.copy()
 
     # take only matched tracks and fill vPhi arrays for inner and outer
-    # df = df[(df.Inner_matched.apply(any) | df.Outer_matched.apply(any))][['entry', 'vPhiInner', 'vPhiOuter']].groupby('entry').sum()
+    # df = (
+    #     df[(df.Inner_matched.apply(any) | df.Outer_matched.apply(any))][
+    #         ["entry", "vPhiInner", "vPhiOuter"]
+    #     ]
+    #     .groupby("entry")
+    #     .sum()
+    # )
     df = df[["entry", "vPhiInner", "vPhiOuter"]].groupby("entry").sum()
-
     df["triggered"] = False
 
     # check incorrect topology
@@ -100,10 +112,13 @@ def GetTriggeredEvents(
                 ):
                     df.at[t, "triggered"] = True
 
-    tdf = pd.DataFrame(df.index[~df.triggered], columns=["Untriggered"]).join(
-        pd.DataFrame(df.index[df.triggered], columns=["Triggered"]), how="left"
-    )
+    tdf = pd.DataFrame.from_dict(
+        {
+            "triggered": df.triggered[df.triggered].index.values,
+            "untriggered": pd.Series(df.triggered[~df.triggered].index.values),
+        },
+        orient="index",
+    ).T
 
-    if saveToFile:
-        tdf.to_parquet(pth)
+    tdf.to_parquet(pth)
     return tdf

notebooks/FourTracks/vis/ErrorHist.py

@@ -1,77 +1,44 @@
-from matplotlib.pyplot import xlabel
 from FourTracks import pd, np, plt, hep
-from typing import Sized, Union
 
 
-def DrawErrorHists(
+def DrawErrorHist(
     data: list,
     bins: int,
     range: tuple,
     title: str,
-    label: list,
-    x_label: list,
-    y_label: list,
-    color: list,
+    x_label: str,
+    y_label: str,
     showBinWidth: bool = True,
-    xerr: bool = True,
+    MeV: bool = True,
+    **histplot_kwargs,
 ) -> pd.DataFrame:
     """
+    Function draws histogram list with errors bars.
     """
 
-    # if (
-    #     pd.DataFrame([data, label, x_label, y_label, color]).count(axis=1).nunique()
-    #     != 1
-    # ):
-    #     raise IndexError("Input lists have different sizes")
-
     plt.style.use(hep.style.ROOT)
     fig = plt.figure()
     ax = fig.add_subplot()
     fig.suptitle(title)
 
-    for i in np.arange(len(data)):
-        hep.histplot(
-            np.histogram(data[i], bins=bins, range=range),
-            yerr=True,
-            xerr=xerr,
-            color=color[i],
-            histtype="errorbar",
-            label=label[i],
-        )
+    val = (range[1] - range[0]) / bins
 
-        val = (range[1] - range[0]) * 1000 // bins
-        ax.set_xlabel(x_label[i])
-        ylabel = y_label[i]
-        if showBinWidth:
-            ylabel = ylabel + " per " + f"{val:.0f} MeV/c"
-        ax.set_ylabel(ylabel)
-
-    _ = plt.legend()
-    return pd.DataFrame(map(pd.DataFrame.describe, data), index=label)
+    hep.histplot(
+        [np.histogram(d, bins=bins, range=range) for d in data],
+        yerr=True,
+        xerr=[0, val/2][showBinWidth],
+        **histplot_kwargs,
+    )
 
+    ax.set_xlabel(x_label)
+    ylabel = y_label
+    if showBinWidth:
+        ylabel = (
+            ylabel + " per " + f"{[1,1000][MeV] * val:.0f} {['GeV/c','MeV/c'][MeV]}"
+        )
+    ax.set_ylabel(ylabel)
 
-def DrawErrorHist(
-    data: Union[list, pd.Series, np.ndarray],
-    bins: int,
-    range: tuple,
-    title: str,
-    label: str,
-    x_label: str,
-    y_label: str,
-    color: str,
-    showBinWidth: bool = True,
-    xerr: bool = True,
-):
-    return DrawErrorHists(
-        [data],
-        bins=bins,
-        range=range,
-        title=title,
-        label=[label],
-        x_label=[x_label],
-        y_label=[y_label],
-        color=[color],
-        showBinWidth=showBinWidth,
-        xerr=xerr,
+    _ = plt.legend()
+    return pd.DataFrame(
+        map(pd.DataFrame.describe, data), index=histplot_kwargs["label"]
     )
-