Merge pull request #5 from LSSTDESC/issue/4/FlowCreator

Issue/4/flow creator
LSSTDESC · Sep 8, 2022 · a25333f · a25333f
2 parents b2e5acc + 2d0ab96
commit a25333f
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 8 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -20,8 +20,7 @@ classifiers = [
 ]
 dynamic = ["version"]
 dependencies = [
-    "pz-rail",
-    "pz-rail-hub",
+    "pz-rail-hub>=0.0.3",
 ]
 
 [project.optional-dependencies]

diff --git a/src/rail/pipelines/examples/goldenspike/goldenspike.py b/src/rail/pipelines/examples/goldenspike/goldenspike.py
@@ -30,8 +30,8 @@ def __init__(self):
         band_dict = {band:f'mag_{band}_lsst' for band in bands}
         rename_dict = {f'mag_{band}_lsst_err':f'mag_err_{band}_lsst' for band in bands}
 
-        self.flow_engine_train = FlowEngine.build(
-            flow=flow_file,
+        self.flow_engine_train = FlowCreator.build(
+            model=flow_file,
             n_samples=50,
             seed=1235,
             output=os.path.join(namer.get_data_dir(DataType.catalog, CatalogType.created), "output_flow_engine_train.pq"),
@@ -73,8 +73,8 @@ def __init__(self):
             output=os.path.join(namer.get_data_dir(DataType.catalog, CatalogType.degraded), "output_table_conv_train.hdf5"),
         )
 
-        self.flow_engine_test = FlowEngine.build(
-            flow=flow_file,
+        self.flow_engine_test = FlowCreator.build(
+            model=flow_file,
             n_samples=50,
             output=os.path.join(namer.get_data_dir(DataType.catalog, CatalogType.degraded), "output_flow_engine_test.pq"),
         )
@@ -176,5 +176,5 @@ def __init__(self):
 
 if __name__ == '__main__':    
     pipe = GoldenspikePipeline()
-    pipe.initialize(dict(flow=flow_file), dict(output_dir='.', log_dir='.', resume=False), None)
+    pipe.initialize(dict(model=flow_file), dict(output_dir='.', log_dir='.', resume=False), None)
     pipe.save('tmp_goldenspike.yml')
diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py
@@ -14,6 +14,6 @@ def test_golden():
     except OSError:  # pragma: no cover
         pass    
     pipe = GoldenspikePipeline()
-    pipe.initialize(dict(flow=flow_file), dict(output_dir=output_dir, log_dir=output_dir, resume=False), None)
+    pipe.initialize(dict(model=flow_file), dict(output_dir=output_dir, log_dir=output_dir, resume=False), None)
     pipe.save('tmp_goldenspike.yml')
     os.system(f"\\rm -rf {output_dir}")