Merge pull request #131 from LSSTDESC/ceci2

Update for ceci version 2

pyproject.toml

@@ -20,7 +20,7 @@ dependencies = [
     "pyyaml",
     "numpy",
     "click",
-    "ceci>=1.10.1",
+    "ceci>=2.0.1",
     "qp-prob>=0.8.3",
     "scipy>=1.9.0",
 ]

src/rail/core/stage.py

@@ -136,10 +136,10 @@ class RailStage(PipelineStage):
 
     data_store = DATA_STORE()
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Constructor:
         Do RailStage specific initialization"""
-        PipelineStage.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self._input_length = None
         self.io = StageIO(self)
 
@@ -288,7 +288,7 @@ def set_data(self, tag, data, path=None, do_read=True):
         if isinstance(data, DataHandle):
             aliased_tag = data.tag
             if tag in self.input_tags():
-                self.config.aliases[tag] = aliased_tag
+                self._aliases[tag] = aliased_tag
                 if data.has_path:
                     self._inputs[tag] = data.path
             arg_data = data.data

src/rail/creation/degrader.py

@@ -23,10 +23,6 @@ class Degrader(RailStage):  # pragma: no cover
     inputs = [("input", PqHandle)]
     outputs = [("output", PqHandle)]
 
-    def __init__(self, args, comm=None):
-        """Initialize Degrader that can degrade photometric data"""
-        RailStage.__init__(self, args, comm=comm)
-
     def __call__(self, sample, seed: int = None):
         """The main interface method for ``Degrader``.
 

src/rail/creation/degraders/addRandom.py

@@ -17,13 +17,13 @@ class AddColumnOfRandom(Noisifier):
         ),
     )
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """
         Constructor
 
         Does standard Noisifier initialization
         """
-        Noisifier.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
 
     def _initNoiseModel(self):  # pragma: no cover
         np.random.seed(self.config.seed)

src/rail/creation/degraders/quantityCut.py

@@ -17,13 +17,13 @@ class QuantityCut(Selector):
     config_options = Selector.config_options.copy()
     config_options.update(cuts=dict)
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Constructor.
 
         Performs standard Degrader initialization as well as defining the cuts
         to be applied.
         """
-        Selector.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self.cuts = None
         self.set_cuts(self.config["cuts"])
 

src/rail/creation/engine.py

@@ -19,9 +19,9 @@ class Modeler(RailStage):  # pragma: no cover
     inputs = [("input", DataHandle)]
     outputs = [("model", ModelHandle)]
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Initialize Modeler"""
-        RailStage.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self.model = None
 
     def fit_model(self):
@@ -56,9 +56,9 @@ class Creator(RailStage):  # pragma: no cover
     inputs = [("model", ModelHandle)]
     outputs = [("output", TableHandle)]
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Initialize Creator"""
-        RailStage.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self.model = None
         if not isinstance(args, dict):  # pragma: no cover
             args = vars(args)
@@ -144,9 +144,9 @@ class PosteriorCalculator(RailStage):  # pragma: no cover
     ]
     outputs = [("output", QPHandle)]
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Initialize PosteriorCalculator"""
-        RailStage.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self.model = None
         if not isinstance(args, dict):  # pragma: no cover
             args = vars(args)

src/rail/creation/noisifier.py

@@ -29,10 +29,6 @@ class Noisifier(RailStage):
     inputs = [("input", PqHandle)]
     outputs = [("output", PqHandle)]
 
-    def __init__(self, args, comm=None):
-        """Initialize Noisifier that can add noise to photometric data"""
-        RailStage.__init__(self, args, comm=comm)
-
     def _initNoiseModel(self):  # pragma: no cover
         raise NotImplementedError("Noisifier._initNoiseModel()")
 

src/rail/creation/selector.py

@@ -31,10 +31,6 @@ class Selector(RailStage):
     inputs = [("input", PqHandle)]
     outputs = [("output", PqHandle)]
 
-    def __init__(self, args, comm=None):
-        """Initialize Noisifier that can add noise to photometric data"""
-        RailStage.__init__(self, args, comm=comm)
-
     def __call__(self, sample):
         """The main interface method for ``Selector``.
 

src/rail/estimation/algos/equal_count.py

@@ -29,9 +29,6 @@ class EqualCountClassifier(PZClassifier):
     )
     outputs = [("output", TableHandle)]
 
-    def __init__(self, args, comm=None):
-        PZClassifier.__init__(self, args, comm=comm)
-
     def run(self):
         test_data = self.get_data("input")
         npdf = test_data.npdf

src/rail/estimation/algos/naive_stack.py

@@ -17,9 +17,6 @@ class NaiveStackInformer(PzInformer):
     name = "NaiveStackInformer"
     config_options = PzInformer.config_options.copy()
 
-    def __init__(self, args, comm=None):
-        PzInformer.__init__(self, args, comm=comm)
-
     def run(self):
         self.add_data("model", np.array([None]))
 
@@ -39,8 +36,8 @@ class NaiveStackSummarizer(PZSummarizer):
     inputs = [("input", QPHandle)]
     outputs = [("output", QPHandle), ("single_NZ", QPHandle)]
 
-    def __init__(self, args, comm=None):
-        PZSummarizer.__init__(self, args, comm=comm)
+    def __init__(self, args, **kwargs):
+        super().__init__(args, **kwargs)
         self.zgrid = None
 
     def run(self):

src/rail/estimation/algos/point_est_hist.py

@@ -17,9 +17,6 @@ class PointEstHistInformer(PzInformer):
     name = "PointEstHistInformer"
     config_options = PzInformer.config_options.copy()
 
-    def __init__(self, args, comm=None):
-        PzInformer.__init__(self, args, comm=comm)
-
     def run(self):
         self.add_data("model", np.array([None]))
 
@@ -40,8 +37,8 @@ class PointEstHistSummarizer(PZSummarizer):
     inputs = [("input", QPHandle)]
     outputs = [("output", QPHandle), ("single_NZ", QPHandle)]
 
-    def __init__(self, args, comm=None):
-        PZSummarizer.__init__(self, args, comm=comm)
+    def __init__(self, args, **kwargs):
+        super().__init__(args, **kwargs)
         self.zgrid = None
         self.bincents = None
 

src/rail/estimation/algos/random_gauss.py

@@ -20,9 +20,6 @@ class RandomGaussInformer(CatInformer):
     name = "RandomGaussInformer"
     config_options = CatInformer.config_options.copy()
 
-    def __init__(self, args, comm=None):
-        CatInformer.__init__(self, args, comm=comm)
-
     def run(self):
         self.add_data("model", np.array([None]))
 
@@ -47,10 +44,10 @@ class RandomGaussEstimator(CatEstimator):
         ),
     )
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Constructor:
         Do CatEstimator specific initialization"""
-        CatEstimator.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self.zgrid = None
 
     def _process_chunk(self, start, end, data, first):

src/rail/estimation/algos/train_z.py

@@ -36,9 +36,6 @@ class TrainZInformer(CatInformer):
         redshift_col=SHARED_PARAMS,
     )
 
-    def __init__(self, args, comm=None):
-        CatInformer.__init__(self, args, comm=comm)
-
     def run(self):
         if self.config.hdf5_groupname:
             training_data = self.get_data("input")[self.config.hdf5_groupname]
@@ -65,11 +62,11 @@ class TrainZEstimator(CatEstimator):
     config_options = CatEstimator.config_options.copy()
     config_options.update(zmin=SHARED_PARAMS, zmax=SHARED_PARAMS, nzbins=SHARED_PARAMS)
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         self.zgrid = None
         self.train_pdf = None
         self.zmode = None
-        CatEstimator.__init__(self, args, comm=comm)
+        CatEstimator.__init__(self, args, **kwargs)
 
     def open_model(self, **kwargs):
         CatEstimator.open_model(self, **kwargs)

src/rail/estimation/algos/true_nz.py

@@ -28,8 +28,8 @@ class TrueNZHistogrammer(RailStage):
     inputs = [("input", TableHandle), ("tomography_bins", TableHandle)]
     outputs = [("true_NZ", QPHandle)]
 
-    def __init__(self, args, comm=None):
-        RailStage.__init__(self, args, comm=comm)
+    def __init__(self, args, **kwargs):
+        super().__init__(args, **kwargs)
         self.zgrid = None
         self.bincents = None
 

src/rail/estimation/algos/uniform_binning.py

@@ -36,18 +36,6 @@ class UniformBinningClassifier(PZClassifier):
     )
     outputs = [("output", Hdf5Handle)]
 
-    def __init__(self, args, comm=None):
-        """Initialize the UniformBinningClassifier.
-
-        Parameters
-        ----------
-        args : dict
-            Configuration arguments for the classifier.
-        comm : MPI.Comm, optional
-            MPI communicator for parallel processing.
-        """
-        PZClassifier.__init__(self, args, comm=comm)
-
     def _process_chunk(self, start, end, data, first):
         """Process a chunk of data for uniform binning classification.
 

src/rail/estimation/algos/var_inf.py

@@ -20,9 +20,6 @@ class VarInfStackInformer(PzInformer):
     name = "VarInfStackInformer"
     config_options = PzInformer.config_options.copy()
 
-    def __init__(self, args, comm=None):
-        PzInformer.__init__(self, args, comm=comm)
-
     def run(self):
         self.add_data("model", np.array([None]))
 
@@ -63,8 +60,8 @@ class VarInfStackSummarizer(PZSummarizer):
     inputs = [("input", QPHandle)]
     outputs = [("output", QPHandle), ("single_NZ", QPHandle)]
 
-    def __init__(self, args, comm=None):
-        PZSummarizer.__init__(self, args, comm=comm)
+    def __init__(self, args, **kwargs):
+        super().__init__(args, **kwargs)
         self.zgrid = None
 
     def run(self):

src/rail/estimation/classifier.py

@@ -23,9 +23,9 @@ class CatClassifier(RailStage):  # pragma: no cover
     inputs = [("model", ModelHandle), ("input", TableHandle)]
     outputs = [("output", TableHandle)]
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Initialize Classifier"""
-        RailStage.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self._output_handle = None
         self.model = None
         if not isinstance(args, dict):  # pragma: no cover
@@ -106,7 +106,7 @@ class PZClassifier(RailStage):
     inputs = [("input", QPHandle)]
     outputs = [("output", Hdf5Handle)]
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Initialize the PZClassifier.
 
         Parameters
@@ -116,7 +116,7 @@ def __init__(self, args, comm=None):
         comm : MPI.Comm, optional
             MPI communicator for parallel processing.
         """
-        RailStage.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self._output_handle = None
 
     def classify(self, input_data):

src/rail/estimation/estimator.py

@@ -36,9 +36,9 @@ class CatEstimator(RailStage, PointEstimationMixin):
     inputs = [("model", ModelHandle), ("input", TableHandle)]
     outputs = [("output", QPHandle)]
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Initialize Estimator"""
-        RailStage.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self._output_handle = None
         self.model = None
 

src/rail/estimation/informer.py

@@ -32,9 +32,9 @@ class CatInformer(RailStage):
     inputs = [("input", TableHandle)]
     outputs = [("model", ModelHandle)]
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Initialize Informer that can inform models for redshift estimation"""
-        RailStage.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self.model = None
 
     def inform(self, training_data):
@@ -88,9 +88,9 @@ class PzInformer(RailStage):
     inputs = [("input", QPHandle)]
     outputs = [("model", ModelHandle)]
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Initialize Informer that can inform models for redshift estimation"""
-        RailStage.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self.model = None
 
     def inform(self, training_data):

src/rail/estimation/summarizer.py

@@ -26,10 +26,6 @@ class CatSummarizer(RailStage):
     inputs = [("input", TableHandle)]
     outputs = [("output", QPHandle)]
 
-    def __init__(self, args, comm=None):
-        """Initialize Summarizer"""
-        RailStage.__init__(self, args, comm=comm)
-
     def summarize(self, input_data):
         """The main run method for the summarization, should be implemented
         in the specific subclass.
@@ -74,10 +70,6 @@ class PZSummarizer(RailStage):
     inputs = [("model", ModelHandle), ("input", QPHandle)]
     outputs = [("output", QPHandle)]
 
-    def __init__(self, args, comm=None):
-        """Initialize Estimator that can sample galaxy data."""
-        RailStage.__init__(self, args, comm=comm)
-
     def summarize(self, input_data):
         """The main run method for the summarization, should be implemented
         in the specific subclass.
@@ -145,9 +137,9 @@ class SZPZSummarizer(RailStage):
     ]
     outputs = [("output", QPHandle)]
 
-    def __init__(self, args, comm=None):
+    def __init__(self, args, **kwargs):
         """Initialize Estimator that can sample galaxy data."""
-        RailStage.__init__(self, args, comm=comm)
+        super().__init__(args, **kwargs)
         self.model = None
         # NOTE: open model removed from init, need to put an
         # `open_model` call explicitly in the run method for