Update example to allow for more than 1 analysis

ninakuni · Sep 5, 2023 · 0755ef1 · 0755ef1
1 parent b4fd564
commit 0755ef1
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 19 deletions.
diff --git a/example_files/yaml/doc_example.yaml b/example_files/yaml/doc_example.yaml
@@ -1,7 +1,7 @@
-Me2017:
+Me2017_1:
   model: Me2017
   svd-path: ./svdmodels
-  outdir: outdir2
+  outdir: outdir_1
   label: injection
   prior: priors/Me2017.prior
   tmin: 0.1
@@ -12,7 +12,30 @@ Me2017:
   Ebv-max: 0
   injection: ./injection.json
   injection-num: 0
-  injection-outfile: outdir/lc.csv
+  injection-outfile: outdir_1/lc.csv
+  generation-seed: 42
+  filters: ztfg,ztfr,ztfi
+  plot: True
+  remove-nondetections: True
+  ztf-uncertainties: True
+  ztf-sampling: True
+  ztf-ToO: 180
+
+Me2017_2:
+  model: Me2017
+  svd-path: ./svdmodels
+  outdir: outdir_2
+  label: injection
+  prior: priors/Me2017.prior
+  tmin: 0.1
+  tmax: 20
+  dt: 0.5
+  error-budget: 1
+  nlive: 512
+  Ebv-max: 0
+  injection: ./injection.json
+  injection-num: 0
+  injection-outfile: outdir_2/lc.csv
   generation-seed: 42
   filters: ztfg,ztfr,ztfi
   plot: True

diff --git a/nmma/em/analysis.py b/nmma/em/analysis.py
@@ -368,22 +368,7 @@ def get_parser():
     return parser
 
 
-def main(args=None):
-    if args is None:
-        parser = get_parser()
-        args = parser.parse_args()
-        if args.config is not None:
-            yaml_dict = yaml.safe_load(Path(args.config).read_text())
-            if not len(yaml_dict.keys()) == 1:
-                print("WARNING: yaml file can only contain one model for now.")
-                exit()
-            params = yaml_dict[list(yaml_dict.keys())[0]]
-            for key, value in params.items():
-                key = key.replace("-", "_")
-                if key not in args:
-                    print(f"{key} not a known argument... please remove")
-                    exit()
-                setattr(args, key, value)
+def analysis(args):
 
     if args.sampler == "pymultinest":
         if len(args.outdir) > 64:
@@ -903,3 +888,24 @@ def main(args=None):
         plt.tight_layout()
         plt.savefig(plotName)
         plt.close()
+
+
+def main(args=None):
+    if args is None:
+        parser = get_parser()
+        args = parser.parse_args()
+        if args.config is not None:
+            yaml_dict = yaml.safe_load(Path(args.config).read_text())
+            for analysis_set in yaml_dict.keys():
+                params = yaml_dict[analysis_set]
+                for key, value in params.items():
+                    key = key.replace("-", "_")
+                    if key not in args:
+                        print(f"{key} not a known argument... please remove")
+                        exit()
+                    setattr(args, key, value)
+                analysis(args)
+        else:
+            analysis(args)
+    else:
+        analysis(args)