list of metadata_fns

abacusorg · Oct 5, 2023 · e9b3b42 · e9b3b42
1 parent d2d5523
commit e9b3b42
Showing 1 changed file with 12 additions and 8 deletions.
diff --git a/abacusnbody/metadata/abacussummit.py b/abacusnbody/metadata/abacussummit.py
@@ -9,7 +9,7 @@
 import msgpack
 
 metadata = None
-metadata_fn = 'abacussummit_headers_compressed.asdf'
+metadata_fns = ['abacussummit_headers_compressed.asdf']
 
 def get_meta(simname, redshift=None):
     '''
@@ -30,22 +30,26 @@ def get_meta(simname, redshift=None):
         The time-independent parameters and, if `redshift` is given,
         the time-dependent state values.
     '''
-
+    # TODO: could add cases: if DESI2 in simname, if png in simname, if Summit in simname
     if not simname.startswith('Abacus'):
         simname = 'Abacus' + simname
 
     global metadata
     if metadata is None:
-        with importlib.resources.open_binary('abacusnbody.metadata', metadata_fn) as fp, asdf.open(fp) as af:
-            metadata = dict(af.tree)
-            del metadata['asdf_library'], metadata['history']
-            for sim in metadata:
-                metadata[sim]['param'] = msgpack.loads(metadata[sim]['param'].data, strict_map_key=False)
-                metadata[sim]['state'] = msgpack.loads(metadata[sim]['state'].data, strict_map_key=False)
+        metadata = {}
+        for i, metadata_fn in enumerate(metadata_fns):
+            with importlib.resources.open_binary('abacusnbody.metadata', metadata_fn) as fp, asdf.open(fp) as af:
+                af_tree = dict(af.tree)
+                del af_tree['asdf_library'], af_tree['history']
+                for sim in af_tree:
+                    metadata[sim] = {}
+                    metadata[sim]['param'] = msgpack.loads(af_tree[sim]['param'].data, strict_map_key=False)
+                    metadata[sim]['state'] = msgpack.loads(af_tree[sim]['state'].data, strict_map_key=False)
 
     if simname not in metadata:
         raise ValueError(f'Simulation "{simname}" is not in metadata file "{metadata_fn}"')
 
+
     res = dict(metadata[simname]['param'])
     if 'CLASS_power_spectrum' in metadata[simname]:
         res['CLASS_power_spectrum'] = metadata[simname]['CLASS_power_spectrum']