special treatment for constants

Author: miranov25

UTILS/dfextensions/AliasDataFrame.py

@@ -15,17 +15,6 @@ class AliasDataFrame:
     """
     A wrapper for pandas DataFrame that supports on-demand computed columns (aliases)
     with dependency tracking and persistence.
-    Example usage:
-    >>> df = pd.DataFrame({"x": [1, 2, 3], "y": [10, 20, 30]})
-    >>> adf = AliasDataFrame(df)
-    >>> adf.add_alias("z", "x + y")
-    >>> adf.add_alias("w", "z * 2")
-    >>> adf.materialize_all()
-    >>> print(adf.df)
-    You can also save and load the dataframe along with aliases:
-    >>> adf.save("mydata")
-    >>> adf2 = AliasDataFrame.load("mydata")
-    >>> adf2.describe_aliases()
     """
 
     def __init__(self, df):
@@ -35,10 +24,6 @@ def __init__(self, df):
         self.constant_aliases = set()  # Optional set of constants that should not be materialized
 
     def add_alias(self, name, expression, dtype=None, is_constant=False):
-        """
-        Add an alias expression to the DataFrame.
-        Optionally specify output dtype and whether it's a constant (scalar-only).
-        """
         try:
             dummy_env = {k: 1 for k in list(self.df.columns) + list(self.aliases.keys())}
             dummy_env.update(self._default_functions())
@@ -142,6 +127,12 @@ def materialize_alias0(self, name, dtype=None):
         if name in self.aliases:
             local_env = {col: self.df[col] for col in self.df.columns}
             local_env.update({k: self.df[k] for k in self.aliases if k in self.df})
+            for cname in self.constant_aliases:
+                try:
+                    val = eval(self.aliases[cname], self._default_functions())
+                    local_env[cname] = val
+                except Exception as e:
+                    print(f"[Alias constant] Failed to evaluate constant '{cname}': {e}")
             result = eval(self.aliases[name], self._default_functions(), local_env)
             result_dtype = dtype or self.alias_dtypes.get(name)
             if result_dtype is not None:
@@ -177,6 +168,12 @@ def visit(n):
                 continue
             local_env = {col: self.df[col] for col in self.df.columns}
             local_env.update({k: self.df[k] for k in self.aliases if k in self.df})
+            for cname in self.constant_aliases:
+                try:
+                    val = eval(self.aliases[cname], self._default_functions())
+                    local_env[cname] = val
+                except Exception as e:
+                    print(f"[Alias constant] Failed to evaluate constant '{cname}': {e}")
             try:
                 result = eval(self.aliases[alias], self._default_functions(), local_env)
                 result_dtype = dtype or self.alias_dtypes.get(alias)
@@ -202,6 +199,9 @@ def materialize_all(self, dtype=None):
             try:
                 local_env = {col: self.df[col] for col in self.df.columns}
                 local_env.update({k: self.df[k] for k in self.df.columns if k in self.aliases})
+                for cname in self.constant_aliases:
+                    val = eval(self.aliases[cname], self._default_functions())
+                    local_env[cname] = val
                 result = eval(self.aliases[name], self._default_functions(), local_env)
                 result_dtype = dtype or self.alias_dtypes.get(name)
                 if result_dtype is not None:
@@ -222,19 +222,18 @@ def save(self, path_prefix, dropAliasColumns=True):
         else:
             cols = list(self.df.columns)
 
-        # Save Parquet with metadata
         table = pa.Table.from_pandas(self.df[cols])
         metadata = {
             "aliases": json.dumps(self.aliases),
-            "dtypes": json.dumps({k: v.__name__ for k, v in self.alias_dtypes.items()})
+            "dtypes": json.dumps({k: v.__name__ for k, v in self.alias_dtypes.items()}),
+            "constants": json.dumps(list(self.constant_aliases))
         }
         existing_meta = table.schema.metadata or {}
         combined_meta = existing_meta.copy()
         combined_meta.update({k.encode(): v.encode() for k, v in metadata.items()})
         table = table.replace_schema_metadata(combined_meta)
         pq.write_table(table, f"{path_prefix}.parquet", compression="zstd")
 
-        # Also write JSON file for explicit tracking
         with open(f"{path_prefix}.aliases.json", "w") as f:
             json.dump(metadata, f, indent=2)
 
@@ -245,17 +244,19 @@ def load(path_prefix):
         df = table.to_pandas()
         adf = AliasDataFrame(df)
 
-        # Try metadata first
         meta = table.schema.metadata or {}
         if b"aliases" in meta and b"dtypes" in meta:
             adf.aliases = json.loads(meta[b"aliases"].decode())
             adf.alias_dtypes = {k: getattr(np, v) for k, v in json.loads(meta[b"dtypes"].decode()).items()}
+            if b"constants" in meta:
+                adf.constant_aliases = set(json.loads(meta[b"constants"].decode()))
         else:
-            # Fallback to JSON
             with open(f"{path_prefix}.aliases.json") as f:
                 data = json.load(f)
                 adf.aliases = json.loads(data["aliases"])
                 adf.alias_dtypes = {k: getattr(np, v) for k, v in json.loads(data["dtypes"]).items()}
+                if "constants" in data:
+                    adf.constant_aliases = set(json.loads(data["constants"]))
 
         return adf
 
@@ -272,7 +273,13 @@ def export_tree(self, filename, treename="tree", dropAliasColumns=True):
         f = ROOT.TFile.Open(filename, "UPDATE")
         tree = f.Get(treename)
         for alias, expr in self.aliases.items():
-            tree.SetAlias(alias, expr)
+            expr_str = expr
+            try:
+                val = float(expr)
+                expr_str = f"({val}+0)"
+            except Exception:
+                pass
+            tree.SetAlias(alias, expr_str)
         tree.Write("", ROOT.TObject.kOverwrite)
         f.Close()
 

UTILS/dfextensions/AliasDataFrameTest.py

@@ -0,0 +1,49 @@
+import unittest
+import pandas as pd
+import numpy as np
+from AliasDataFrame import AliasDataFrame  # Adjust this if you're using a different import method
+
+class TestAliasDataFrame(unittest.TestCase):
+    def setUp(self):
+        df = pd.DataFrame({
+            "x": np.arange(5),
+            "y": np.arange(5, 10),
+            "CTPLumi_countsFV0": np.array([2000, 2100, 2200, 2300, 2400])
+        })
+        self.adf = AliasDataFrame(df)
+
+    def test_basic_alias(self):
+        self.adf.add_alias("z", "x + y")
+        self.adf.materialize_alias("z")
+        expected = self.adf.df["x"] + self.adf.df["y"]
+        pd.testing.assert_series_equal(self.adf.df["z"], expected, check_names=False)
+
+    def test_dtype(self):
+        self.adf.add_alias("z", "x + y", dtype=np.float16)
+        self.adf.materialize_alias("z")
+        self.assertEqual(self.adf.df["z"].dtype, np.float16)
+
+    def test_constant(self):
+        self.adf.add_alias("c", "42.0", dtype=np.float32, is_constant=True)
+        self.adf.add_alias("z", "x + c")
+        self.adf.materialize_alias("z")
+        expected = self.adf.df["x"] + 42.0
+        pd.testing.assert_series_equal(self.adf.df["z"], expected, check_names=False)
+
+    def test_dependency_order(self):
+        self.adf.add_alias("a", "x + y")
+        self.adf.add_alias("b", "a * 2")
+        self.adf.materialize_alias("b")
+        expected = (self.adf.df["x"] + self.adf.df["y"]) * 2
+        pd.testing.assert_series_equal(self.adf.df["b"], expected, check_names=False)
+
+    def test_log_rate_with_constant(self):
+        median = self.adf.df["CTPLumi_countsFV0"].median()
+        self.adf.add_alias("countsFV0_median", f"{median}", dtype=np.float16, is_constant=True)
+        self.adf.add_alias("logRate", "log(CTPLumi_countsFV0/countsFV0_median)", dtype=np.float16)
+        self.adf.materialize_alias("logRate")
+        expected = np.log(self.adf.df["CTPLumi_countsFV0"] / median).astype(np.float16)
+        pd.testing.assert_series_equal(self.adf.df["logRate"], expected, check_names=False)
+
+if __name__ == "__main__":
+    unittest.main()