pyic_intake.py: First commit.

tools/pyic_intake.py

+#!/usr/bin/env python
+import os
+import glob
+import yaml
+import argparse
+
+def generate_master_catalog(catalog_dir='default'):
+    data_all = dict()
+    data_all["sources"] = dict()
+    
+    catlist = glob.glob(f"{catalog_dir}/cat_*.yaml")
+    catlist.sort()
+    
+    # assuming that every catalog only contains one simulation
+    for fname in catlist:
+        with open(fname, "r") as file:
+            print(f"Adding {fname}")
+            data = yaml.safe_load(file)
+            key = list(data["sources"])[0]
+            data_all["sources"][key] = data["sources"][key]
+    
+    if catalog_dir=='default':
+        catalog_dir = os.path.expanduser("~") + '/intake_pyicon/'
+    with open(f"{catalog_dir}/main_catalog.yaml", "w") as file:
+        yaml.dump(data_all, file, default_flow_style=False)
+    return
+
+def generate_simulation_catalog(run, 
+                                path_data='./', 
+                                project='MPI-M ICON',
+                                description='',
+                                tstr='auto',
+                                fname_out='auto',
+                                catalog_dir='default',
+                                verbose=False,
+                               ):
+
+    path_data = os.path.abspath(path_data)
+    
+    if tstr=='auto':
+        flist = glob.glob(f'{path_data}/*.nc')
+        tstr = flist[0].split('_')[-1].split('.nc')[0]
+        print(f'Using tstr={tstr}')
+
+    searchstr = f"{path_data}/{run}*{tstr}*.nc"
+    print(f"Looking in the following path for tags: {searchstr}")
+    fnames = glob.glob(searchstr)
+    tags = [0]*len(fnames)
+    for nn, fname in enumerate(fnames):
+        tags[nn] = fname.split('/')[-1].split(f'{run}_')[-1].split(f'_{tstr}.nc')[0]
+    print(f"Found the following tags: {tags}")
+
+    Dmaster = dict()
+    Dmaster['sources'] = dict()
+    Dmaster['sources'][run] = dict()
+    Dmaster['sources'][run]['args'] = dict(urlpath = path_data)
+    Dmaster['sources'][run]['driver'] = 'netcdf'
+    Dmaster['sources'][run]['parameters'] = dict()
+    Dmaster['sources'][run]['parameters']['tag'] = dict()
+    Dmaster['sources'][run]['parameters']['tag']['allowed'] = tags
+    Dmaster['sources'][run]['parameters']['tag']['default'] = tags[0]
+    Dmaster['sources'][run]['parameters']['tag']['type'] = 'str'
+    
+    Dmaster['sources'][run]['metadata'] = dict()
+    Dmaster['sources'][run]['metadata']['project'] = project
+    Dmaster['sources'][run]['metadata']['simulation_id'] = run
+    Dmaster['sources'][run]['metadata']['description'] = description
+
+    if verbose:
+        print('----')
+        print(Dmaster)
+        print('----')
+
+    if fname_out=='auto':
+        fname_out = f"cat_{run}.yaml"
+
+    if catalog_dir=='default':
+        catalog_dir = os.path.expanduser("~") + '/intake_pyicon/'
+    if not os.path.exists(catalog_dir):
+      os.makedirs(catalog_dir)
+
+    fpath_out = catalog_dir+'/'+fname_out
+    print(f"Writing catalog file {fpath_out}")
+    with open(fpath_out, 'w') as file:
+        yaml.dump(Dmaster, file, default_flow_style=False)
+
+    generate_master_catalog(catalog_dir)
+    return
+
+if __name__ == "__main__":
+    
+    parser = argparse.ArgumentParser(description=generate_simulation_catalog.__doc__)
+    parser.add_argument("run", type=str, 
+        help="Help is about come -- but not just now.")
+    parser.add_argument("--path_data", type=str, default='./',
+        help="Help is about come -- but not just now.")
+    parser.add_argument("--project", type=str, default='MPI-M ICON',
+        help="Help is about come -- but not just now.")
+    parser.add_argument("--description", type=str, default='',
+        help="Help is about come -- but not just now.")
+    parser.add_argument("--tstr", type=str, default='auto',
+        help="Help is about come -- but not just now.")
+    parser.add_argument("--fname_out", type=str, default='auto',
+        help="Help is about come -- but not just now.")
+    parser.add_argument("--catalog_dir", type=str, default='default',
+        help="Help is about come -- but not just now.")  
+
+    # parse the arguments
+    args = parser.parse_args()
+    
+    # call the function with parsed arguments
+    generate_simulation_catalog(
+        args.run, 
+        path_data=args.path_data,
+        description=args.description, 
+        tstr=args.tstr, 
+        fname_out=args.fname_out, 
+        catalog_dir=args.catalog_dir,
+    )