feat: direct index for plot data

jiangyi15 · Apr 24, 2024 · 81a2a86 · 81a2a86
1 parent 9824803
commit 81a2a86
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 4 deletions.
diff --git a/tf_pwa/config_loader/config_loader.py b/tf_pwa/config_loader/config_loader.py
@@ -1226,6 +1226,8 @@ def __init__(self, plot_config, decay_struct):
             self.params.append(i)
         for i in self.get_angle_vars(True):
             self.params.append(i)
+        for i in self.get_index_vars():
+            self.params.append(i)
         for i in self.get_extra_vars():
             self.params.append(i)
 
@@ -1423,6 +1425,30 @@ def readdata(x):
                 "readdata": readdata,
             }
 
+    def get_index_vars(self):
+
+        dic = self.config.get("index", {})
+
+        for k, v in dic.items():
+            idx = self.get_data_index("index", k)
+            id_ = v.get("id", k)
+            display = v.get("display", str(k))
+            trans = v.get("trans", None)
+            if trans is None:
+                trans = lambda x: x
+            else:
+                trans = sy.sympify(trans)
+                x = sy.symbols("x")
+                trans = sy.lambdify(x, trans, modules="numpy")
+            common_config = self.read_plot_config(v)
+            readdata = ReadData(idx, trans)
+            yield {
+                **common_config,
+                "name": k,
+                "display": display,
+                "readdata": readdata,
+            }
+
     def get_params(self, params=None):
         if params is None:
             return self.params

diff --git a/tf_pwa/config_loader/plotter.py b/tf_pwa/config_loader/plotter.py
@@ -159,16 +159,14 @@ def get_all_frame(self):
         name = conf.get("name")
         display = conf.get("display", name)
         upper_ylim = conf.get("upper_ylim", None)
-        idx = conf.get("idx")
-        trans = conf.get("trans", lambda x: x)
+        readdata = conf.get("readdata")
         has_legend = conf.get("legend", False)
         xrange = conf.get("range", None)
         bins = conf.get("bins", None)
         units = conf.get("units", "")
         yscale = conf.get("yscale", "linear")
         ret[name] = Frame(
-            idx,
-            trans=trans,
+            readdata,
             name=name,
             display=display,
             x_range=xrange,