Add columns property on expression actions

lsst · Jun 9, 2021 · 2c829e9 · 2c829e9
1 parent ec24388
commit 2c829e9
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 4 deletions.
diff --git a/python/lsst/pipe/tasks/dataFrameActions/_actions.py b/python/lsst/pipe/tasks/dataFrameActions/_actions.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 __all__ = ("SingleColumnAction", "MultiColumnAction", "CoordColumn", "MagColumnDN", "SumColumns", "AddColumn",
-           "DivideColumns", "SubtractColumns", "MultiplyColumns",)
+           "DivideColumns", "SubtractColumns", "MultiplyColumns", "MagColumnNanoJansky",)
 
 from typing import Iterable
 
@@ -56,6 +56,16 @@ def __call__(self, df: pd.DataFrame, **kwargs):
             return -2.5 * np.log10(df[self.column] / fluxMag0)
 
 
+class MagColumnNanoJansky(SingleColumnAction):
+
+    def __call__(self, df: pd.DataFrame, **kwargs):
+
+        with np.warnings.catch_warnings():
+            np.warnings.filterwarnings('ignore', r'invalid value encountered')
+            np.warnings.filterwarnings('ignore', r'divide by zero')
+            return -2.5 * np.log10((df[self.column] * 1e-9) / 3631.0)
+
+
 class NanoJansky(SingleColumnAction):
     ab_flux_scale = Field(doc="Scaling of ab flux", dtype=float, default=(0*units.ABmag).to_value(units.nJy))
     coadd_zeropoint = Field(doc="Magnitude zero point", dtype=float, default=27)

diff --git a/python/lsst/pipe/tasks/dataFrameActions/_baseDataFrameActions.py b/python/lsst/pipe/tasks/dataFrameActions/_baseDataFrameActions.py
@@ -15,7 +15,8 @@ class DataFrameAction(ConfigurableAction):
                       " only works on frozen actions",
                   dtype=bool, default=False)
     cacheArgs = ListField(doc="If cache is True, this is a list of argument keys that will be used to "
-                          "compute the cache key in addition to the DataFrameId", dtype=str)
+                              "compute the cache key in addition to the DataFrameId",
+                          dtype=str, optional=True)
 
     def __init_subclass__(cls, **kwargs) -> None:
         cls._actionCache = {}

diff --git a/python/lsst/pipe/tasks/dataFrameActions/_evalColumnExpression.py b/python/lsst/pipe/tasks/dataFrameActions/_evalColumnExpression.py
@@ -26,7 +26,7 @@
 import ast
 import operator as op
 
-from typing import Mapping, MutableMapping, Set, Type, Union, Optional, Any
+from typing import Mapping, MutableMapping, Set, Type, Union, Optional, Any, Iterable
 
 from numpy import log10 as log
 from numpy import (cos, sin, cosh, sinh)
@@ -114,7 +114,12 @@ def __call__(self, df: pd.DataFrame, **kwargs) -> pd.Series:  # noqa: N807
         parser = ExpressionParser(**values_map)
         return parser.visit(node.body)
 
-    dct: MutableMapping[str, Any] = {"__call__": __call__}
+    # create the function to look up the columns for the dynamically created action
+    def columns(self) -> Iterable[str]:
+        for name in fields:
+            yield from getattr(self, name).columns
+
+    dct: MutableMapping[str, Any] = {"__call__": __call__, "columns": property(columns)}
     if docstring is not None:
         dct['__doc__'] = docstring
     dct.update(**fields)