Make RhoStatisticsPlot.rhoPlots functionally configurable

lsst · Apr 3, 2024 · ef99026 · ef99026
1 parent a444243
commit ef99026
Showing 1 changed file with 31 additions and 21 deletions.
diff --git a/python/lsst/analysis/tools/actions/plot/rhoStatisticsPlot.py b/python/lsst/analysis/tools/actions/plot/rhoStatisticsPlot.py
@@ -84,15 +84,17 @@ def getInputSchema(self) -> KeyedDataSchema:
 
     def getOutputNames(self) -> Iterable[str]:
         # Docstring inherited
-        return ("rho3alt", "rho1", "rho2", "rho3", "rho4", "rho5")
+        for key in self.rhoPlots.keys():
+            yield key
 
     def __call__(self, data: KeyedData, **kwargs) -> TypedPlotMapping:
         self._validateInput(data)
         return self.makePlot(data, **kwargs)
 
     def _validateInput(self, data: KeyedData) -> None:
-        if not set(("rho3alt", "rho1", "rho2", "rho3", "rho4", "rho5")).issubset(data.keys()):
-            raise ValueError("Input data must contain rho3alt, rho1, rho2, rho3, rho4, and rho5.")
+        required = set(self.rhoPlots.keys())
+        if not required.issubset(data.keys()):
+            raise ValueError(f"Input data must contain {', '.join(self.rhoPlots.keys())}")
 
     def makePlot(
         self, data: KeyedData, plotInfo: Mapping[str, str] | None = None, **kwargs: Any
@@ -134,28 +136,36 @@ def makePlot(
         :ref:`getting started guide<analysis-tools-getting-started>`.
         """
         fig_dict: dict[str, TypedPlot] = {}
-        for rho_name in ("rho1", "rho2", "rho3", "rho4", "rho5"):
-            rho: XYPlot = self.rhoPlots[rho_name]
+        for rho_name in self.rhoPlots.keys():
+            if rho_name != "rho3alt":
+                rho: XYPlot = self.rhoPlots[rho_name]
+
+                subdata = {
+                    "x": data[rho_name].meanr,  # type: ignore
+                    "y": data[rho_name].xip,  # type: ignore
+                    "yerr": np.sqrt(data[rho_name].varxip),  # type: ignore
+                    "xerr": None,
+                }
+                fig = rho(subdata, **kwargs)
+                if plotInfo is not None:
+                    fig_dict[rho_name] = TypedPlot(
+                        plot=addPlotInfo(fig, plotInfo),
+                        plotType=rho.getPlotName(),
+                    )
 
+        # rho3alt is handled differently because its attributes differ.
+        if "rho3alt" in self.rhoPlots:
             subdata = {
-                "x": data[rho_name].meanr,  # type: ignore
-                "y": data[rho_name].xip,  # type: ignore
-                "yerr": np.sqrt(data[rho_name].varxip),  # type: ignore
+                "x": data["rho3alt"].meanr,  # type: ignore
+                "y": data["rho3alt"].xi,  # type: ignore
+                "yerr": np.sqrt(data["rho3alt"].varxi),  # type: ignore
                 "xerr": None,
             }
-            fig = rho(subdata, **kwargs)
+            fig = self.rhoPlots["rho3alt"](subdata, **kwargs)  # type: ignore[misc]
             if plotInfo is not None:
-                fig_dict[rho_name] = TypedPlot(plot=addPlotInfo(fig, plotInfo), plotType=rho.getPlotName())
-
-        # rho3alt is handled differently because its attributes differ.
-        subdata = {
-            "x": data["rho3alt"].meanr,  # type: ignore
-            "y": data["rho3alt"].xi,  # type: ignore
-            "yerr": np.sqrt(data["rho3alt"].varxi),  # type: ignore
-            "xerr": None,
-        }
-        fig = self.rhoPlots["rho3alt"](subdata, **kwargs)  # type: ignore[misc]
-        if plotInfo is not None:
-            fig_dict["rho3alt"] = TypedPlot(plot=addPlotInfo(fig, plotInfo), plotType=rho.getPlotName())
+                fig_dict["rho3alt"] = TypedPlot(
+                    plot=addPlotInfo(fig, plotInfo),
+                    plotType=rho.getPlotName(),
+                )
 
         return fig_dict