Skip to content

Commit 36166e7

Browse files
xylarxylar
authored
Merge pull request #184 from xylar/cached_init
Add the capability to use cached meshes and ICs
2 parents 1f2616f + 0fa5a4c commit 36166e7

File tree

14 files changed

+555
-31
lines changed

14 files changed

+555
-31
lines changed

‎compass/__main__.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@
22

33
import sys
44
import argparse
5+
import os
56

67
import compass
7-
from compass import list, setup, clean, suite, run
8+
from compass import list, setup, clean, suite, run, cache
89

910

1011
def main():
@@ -40,11 +41,18 @@ def main():
4041

4142
args = parser.parse_args(sys.argv[1:2])
4243

44+
# only allow the "compass cache" command if we're on Anvil or Chrysalis
45+
allow_cache = ('COMPASS_MACHINE' in os.environ and
46+
os.environ['COMPASS_MACHINE'] in ['anvil', 'chrysalis'])
47+
4348
commands = {'list': list.main,
4449
'setup': setup.main,
4550
'clean': clean.main,
4651
'suite': suite.main,
4752
'run': run.main}
53+
if allow_cache:
54+
commands['cache'] = cache.main
55+
4856
if args.command not in commands:
4957
print('Unrecognized command {}'.format(args.command))
5058
parser.print_help()

‎compass/cache.py‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
import argparse
2+
import json
3+
import sys
4+
from datetime import datetime
5+
import os
6+
from importlib import resources
7+
import configparser
8+
import shutil
9+
import pickle
10+
11+
from compass.config import add_config
12+
13+
14+
def update_cache(step_paths, date_string=None, dry_run=False):
15+
"""
16+
Cache one or more compass output files for use in a cached variant of the
17+
test case or step
18+
19+
Parameters
20+
----------
21+
step_paths : list of str
22+
The relative path of the original (uncached) steps from the base work
23+
directory
24+
25+
date_string : str, optional
26+
The datestamp (YYMMDD) to use on the files. Default is today's date.
27+
28+
dry_run : bool, optional
29+
Whether this is a dry run (producing the json file but not copying
30+
files to the LCRC server)
31+
"""
32+
if 'COMPASS_MACHINE' not in os.environ:
33+
machine = None
34+
invalid = True
35+
else:
36+
machine = os.environ['COMPASS_MACHINE']
37+
invalid = machine not in ['anvil', 'chrysalis']
38+
39+
if invalid:
40+
raise ValueError('You must cache files from either Anvil or Chrysalis')
41+
42+
config = configparser.ConfigParser(
43+
interpolation=configparser.ExtendedInterpolation())
44+
add_config(config, 'compass.machines', '{}.cfg'.format(machine))
45+
46+
if date_string is None:
47+
date_string = datetime.now().strftime("%y%m%d")
48+
49+
# make a dictionary with MPAS cores as keys, and lists of steps as values
50+
steps = dict()
51+
for path in step_paths:
52+
with open(f'{path}/step.pickle', 'rb') as handle:
53+
_, step = pickle.load(handle)
54+
55+
mpas_core = step.mpas_core.name
56+
57+
if mpas_core in steps:
58+
steps[mpas_core].append(step)
59+
else:
60+
steps[mpas_core] = [step]
61+
62+
# now, iterate over cores and steps
63+
for mpas_core in steps:
64+
database_root = config.get('paths', f'{mpas_core}_database_root')
65+
cache_root = f'{database_root}/compass_cache'
66+
67+
package = f'compass.{mpas_core}'
68+
try:
69+
with open(f'{mpas_core}_cached_files.json') as data_file:
70+
cached_files = json.load(data_file)
71+
except FileNotFoundError:
72+
# we don't have a local version of the file yet, let's see if
73+
# there's a remote one for this MPAS core
74+
try:
75+
with resources.path(package, 'cached_files.json') as path:
76+
with open(path) as data_file:
77+
cached_files = json.load(data_file)
78+
except FileNotFoundError:
79+
# no cached files yet for this core
80+
cached_files = dict()
81+
82+
for step in steps[mpas_core]:
83+
# load the step from its pickle file
84+
85+
step_path = step.path
86+
87+
for output in step.outputs:
88+
output = os.path.basename(output)
89+
out_filename = os.path.join(step_path, output)
90+
# remove the MPAS core from the file path
91+
target = out_filename[len(mpas_core)+1:]
92+
path, ext = os.path.splitext(target)
93+
target = f'{path}.{date_string}{ext}'
94+
cached_files[out_filename] = target
95+
96+
print(out_filename)
97+
print(f' ==> {target}')
98+
output_path = f'{cache_root}/{target}'
99+
print(f' copy to: {output_path}')
100+
print()
101+
if not dry_run:
102+
directory = os.path.dirname(output_path)
103+
try:
104+
os.makedirs(directory)
105+
except FileExistsError:
106+
pass
107+
shutil.copyfile(out_filename, output_path)
108+
109+
out_filename = f'{mpas_core}_cached_files.json'
110+
with open(out_filename, 'w') as data_file:
111+
json.dump(cached_files, data_file, indent=4)
112+
113+
114+
def main():
115+
parser = argparse.ArgumentParser(
116+
description='Cache the output files from one or more steps for use in '
117+
'a cached variant of the step',
118+
prog='compass cache')
119+
parser.add_argument("-i", "--orig_steps", nargs='+', dest="orig_steps",
120+
type=str,
121+
help="The relative path of the original (uncached) "
122+
"steps from the base work directory",
123+
metavar="STEP")
124+
parser.add_argument("-d", "--date_string", dest="date_string", type=str,
125+
help="The datestamp (YYMMDD) to use on the files. "
126+
"Default is today's date.",
127+
metavar="DATE")
128+
parser.add_argument("-r", "--dry_run", dest="dry_run",
129+
help="Whether this is a dry run (producing the json "
130+
"file but not copying files to the LCRC server).",
131+
action="store_true")
132+
133+
args = parser.parse_args(sys.argv[2:])
134+
update_cache(step_paths=args.orig_steps, date_string=args.date_string,
135+
dry_run=args.dry_run)

‎compass/mpas_core.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
from importlib import resources
2+
import json
3+
4+
15
class MpasCore:
26
"""
37
The base class for housing all the tests for a given MPAS core, such as
@@ -10,6 +14,11 @@ class MpasCore:
1014
1115
test_groups : dict
1216
A dictionary of test groups for the MPAS core with their names as keys
17+
18+
cached_files : dict
19+
A dictionary that maps from output file names in test cases to cached
20+
files in the ``compass_cache`` database for the MPAS core. These
21+
file mappings are read in from ``cached_files.json`` in the MPAS core.
1322
"""
1423

1524
def __init__(self, name):
@@ -26,6 +35,9 @@ def __init__(self, name):
2635
# test groups are added with add_test_groups()
2736
self.test_groups = dict()
2837

38+
self.cached_files = dict()
39+
self._read_cached_files()
40+
2941
def add_test_group(self, test_group):
3042
"""
3143
Add a test group to the MPAS core
@@ -36,3 +48,16 @@ def add_test_group(self, test_group):
3648
the test group to add
3749
"""
3850
self.test_groups[test_group.name] = test_group
51+
52+
def _read_cached_files(self):
53+
""" Read in the dictionary of cached files from cached_files.json """
54+
55+
package = f'compass.{self.name}'
56+
filename = 'cached_files.json'
57+
try:
58+
with resources.path(package, filename) as path:
59+
with open(path) as data_file:
60+
self.cached_files = json.load(data_file)
61+
except FileNotFoundError:
62+
# no cached files for this core
63+
pass

‎compass/ocean/cached_files.json‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
{
2+
"ocean/global_ocean/QU240/mesh/mesh/culled_mesh.nc": "global_ocean/QU240/mesh/mesh/culled_mesh.210809.nc",
3+
"ocean/global_ocean/QU240/mesh/mesh/culled_graph.info": "global_ocean/QU240/mesh/mesh/culled_graph.210809.info",
4+
"ocean/global_ocean/QU240/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/QU240/mesh/mesh/critical_passages_mask_final.210809.nc",
5+
"ocean/global_ocean/QU240/PHC/init/initial_state/initial_state.nc": "global_ocean/QU240/PHC/init/initial_state/initial_state.210809.nc",
6+
"ocean/global_ocean/QU240/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/QU240/PHC/init/initial_state/init_mode_forcing_data.210809.nc",
7+
"ocean/global_ocean/QU240/PHC/init/initial_state/graph.info": "global_ocean/QU240/PHC/init/initial_state/graph.210809.info",
8+
"ocean/global_ocean/QUwISC240/mesh/mesh/culled_mesh.nc": "global_ocean/QUwISC240/mesh/mesh/culled_mesh.210809.nc",
9+
"ocean/global_ocean/QUwISC240/mesh/mesh/culled_graph.info": "global_ocean/QUwISC240/mesh/mesh/culled_graph.210809.info",
10+
"ocean/global_ocean/QUwISC240/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/QUwISC240/mesh/mesh/critical_passages_mask_final.210809.nc",
11+
"ocean/global_ocean/QUwISC240/PHC/init/initial_state/initial_state.nc": "global_ocean/QUwISC240/PHC/init/initial_state/initial_state.210809.nc",
12+
"ocean/global_ocean/QUwISC240/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/QUwISC240/PHC/init/initial_state/init_mode_forcing_data.210809.nc",
13+
"ocean/global_ocean/QUwISC240/PHC/init/initial_state/graph.info": "global_ocean/QUwISC240/PHC/init/initial_state/graph.210809.info",
14+
"ocean/global_ocean/QUwISC240/PHC/init/ssh_adjustment/adjusted_init.nc": "global_ocean/QUwISC240/PHC/init/ssh_adjustment/adjusted_init.210809.nc",
15+
"ocean/global_ocean/EC30to60/mesh/mesh/culled_mesh.nc": "global_ocean/EC30to60/mesh/mesh/culled_mesh.210809.nc",
16+
"ocean/global_ocean/EC30to60/mesh/mesh/culled_graph.info": "global_ocean/EC30to60/mesh/mesh/culled_graph.210809.info",
17+
"ocean/global_ocean/EC30to60/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/EC30to60/mesh/mesh/critical_passages_mask_final.210809.nc",
18+
"ocean/global_ocean/EC30to60/PHC/init/initial_state/initial_state.nc": "global_ocean/EC30to60/PHC/init/initial_state/initial_state.210809.nc",
19+
"ocean/global_ocean/EC30to60/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/EC30to60/PHC/init/initial_state/init_mode_forcing_data.210809.nc",
20+
"ocean/global_ocean/EC30to60/PHC/init/initial_state/graph.info": "global_ocean/EC30to60/PHC/init/initial_state/graph.210809.info",
21+
"ocean/global_ocean/WC14/mesh/mesh/culled_mesh.nc": "global_ocean/WC14/mesh/mesh/culled_mesh.210809.nc",
22+
"ocean/global_ocean/WC14/mesh/mesh/culled_graph.info": "global_ocean/WC14/mesh/mesh/culled_graph.210809.info",
23+
"ocean/global_ocean/WC14/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/WC14/mesh/mesh/critical_passages_mask_final.210809.nc",
24+
"ocean/global_ocean/WC14/PHC/init/initial_state/initial_state.nc": "global_ocean/WC14/PHC/init/initial_state/initial_state.210809.nc",
25+
"ocean/global_ocean/WC14/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/WC14/PHC/init/initial_state/init_mode_forcing_data.210809.nc",
26+
"ocean/global_ocean/WC14/PHC/init/initial_state/graph.info": "global_ocean/WC14/PHC/init/initial_state/graph.210809.info",
27+
"ocean/global_ocean/ECwISC30to60/mesh/mesh/culled_mesh.nc": "global_ocean/ECwISC30to60/mesh/mesh/culled_mesh.210809.nc",
28+
"ocean/global_ocean/ECwISC30to60/mesh/mesh/culled_graph.info": "global_ocean/ECwISC30to60/mesh/mesh/culled_graph.210809.info",
29+
"ocean/global_ocean/ECwISC30to60/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/ECwISC30to60/mesh/mesh/critical_passages_mask_final.210809.nc",
30+
"ocean/global_ocean/ECwISC30to60/PHC/init/initial_state/initial_state.nc": "global_ocean/ECwISC30to60/PHC/init/initial_state/initial_state.210809.nc",
31+
"ocean/global_ocean/ECwISC30to60/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/ECwISC30to60/PHC/init/initial_state/init_mode_forcing_data.210809.nc",
32+
"ocean/global_ocean/ECwISC30to60/PHC/init/ssh_adjustment/adjusted_init.nc": "global_ocean/ECwISC30to60/PHC/init/ssh_adjustment/adjusted_init.210809.nc",
33+
"ocean/global_ocean/ECwISC30to60/PHC/init/initial_state/graph.info": "global_ocean/ECwISC30to60/PHC/init/initial_state/graph.210809.info",
34+
"ocean/global_ocean/SOwISC12to60/mesh/mesh/culled_mesh.nc": "global_ocean/SOwISC12to60/mesh/mesh/culled_mesh.210810.nc",
35+
"ocean/global_ocean/SOwISC12to60/mesh/mesh/culled_graph.info": "global_ocean/SOwISC12to60/mesh/mesh/culled_graph.210810.info",
36+
"ocean/global_ocean/SOwISC12to60/mesh/mesh/critical_passages_mask_final.nc": "global_ocean/SOwISC12to60/mesh/mesh/critical_passages_mask_final.210810.nc",
37+
"ocean/global_ocean/SOwISC12to60/PHC/init/initial_state/initial_state.nc": "global_ocean/SOwISC12to60/PHC/init/initial_state/initial_state.210810.nc",
38+
"ocean/global_ocean/SOwISC12to60/PHC/init/initial_state/init_mode_forcing_data.nc": "global_ocean/SOwISC12to60/PHC/init/initial_state/init_mode_forcing_data.210810.nc",
39+
"ocean/global_ocean/SOwISC12to60/PHC/init/ssh_adjustment/adjusted_init.nc": "global_ocean/SOwISC12to60/PHC/init/ssh_adjustment/adjusted_init.210810.nc",
40+
"ocean/global_ocean/SOwISC12to60/PHC/init/initial_state/graph.info": "global_ocean/SOwISC12to60/PHC/init/initial_state/graph.210810.info",
41+
"ocean/global_convergence/cosine_bell/QU60/mesh/mesh.nc": "global_convergence/cosine_bell/QU60/mesh/mesh.210803.nc",
42+
"ocean/global_convergence/cosine_bell/QU60/mesh/graph.info": "global_convergence/cosine_bell/QU60/mesh/graph.210803.info",
43+
"ocean/global_convergence/cosine_bell/QU60/init/namelist.ocean": "global_convergence/cosine_bell/QU60/init/namelist.210803.ocean",
44+
"ocean/global_convergence/cosine_bell/QU60/init/initial_state.nc": "global_convergence/cosine_bell/QU60/init/initial_state.210803.nc",
45+
"ocean/global_convergence/cosine_bell/QU90/mesh/mesh.nc": "global_convergence/cosine_bell/QU90/mesh/mesh.210803.nc",
46+
"ocean/global_convergence/cosine_bell/QU90/mesh/graph.info": "global_convergence/cosine_bell/QU90/mesh/graph.210803.info",
47+
"ocean/global_convergence/cosine_bell/QU90/init/namelist.ocean": "global_convergence/cosine_bell/QU90/init/namelist.210803.ocean",
48+
"ocean/global_convergence/cosine_bell/QU90/init/initial_state.nc": "global_convergence/cosine_bell/QU90/init/initial_state.210803.nc",
49+
"ocean/global_convergence/cosine_bell/QU120/mesh/mesh.nc": "global_convergence/cosine_bell/QU120/mesh/mesh.210803.nc",
50+
"ocean/global_convergence/cosine_bell/QU120/mesh/graph.info": "global_convergence/cosine_bell/QU120/mesh/graph.210803.info",
51+
"ocean/global_convergence/cosine_bell/QU120/init/namelist.ocean": "global_convergence/cosine_bell/QU120/init/namelist.210803.ocean",
52+
"ocean/global_convergence/cosine_bell/QU120/init/initial_state.nc": "global_convergence/cosine_bell/QU120/init/initial_state.210803.nc",
53+
"ocean/global_convergence/cosine_bell/QU180/mesh/mesh.nc": "global_convergence/cosine_bell/QU180/mesh/mesh.210803.nc",
54+
"ocean/global_convergence/cosine_bell/QU180/mesh/graph.info": "global_convergence/cosine_bell/QU180/mesh/graph.210803.info",
55+
"ocean/global_convergence/cosine_bell/QU180/init/namelist.ocean": "global_convergence/cosine_bell/QU180/init/namelist.210803.ocean",
56+
"ocean/global_convergence/cosine_bell/QU180/init/initial_state.nc": "global_convergence/cosine_bell/QU180/init/initial_state.210803.nc",
57+
"ocean/global_convergence/cosine_bell/QU210/mesh/mesh.nc": "global_convergence/cosine_bell/QU210/mesh/mesh.210803.nc",
58+
"ocean/global_convergence/cosine_bell/QU210/mesh/graph.info": "global_convergence/cosine_bell/QU210/mesh/graph.210803.info",
59+
"ocean/global_convergence/cosine_bell/QU210/init/namelist.ocean": "global_convergence/cosine_bell/QU210/init/namelist.210803.ocean",
60+
"ocean/global_convergence/cosine_bell/QU210/init/initial_state.nc": "global_convergence/cosine_bell/QU210/init/initial_state.210803.nc",
61+
"ocean/global_convergence/cosine_bell/QU240/mesh/mesh.nc": "global_convergence/cosine_bell/QU240/mesh/mesh.210803.nc",
62+
"ocean/global_convergence/cosine_bell/QU240/mesh/graph.info": "global_convergence/cosine_bell/QU240/mesh/graph.210803.info",
63+
"ocean/global_convergence/cosine_bell/QU240/init/namelist.ocean": "global_convergence/cosine_bell/QU240/init/namelist.210803.ocean",
64+
"ocean/global_convergence/cosine_bell/QU240/init/initial_state.nc": "global_convergence/cosine_bell/QU240/init/initial_state.210803.nc",
65+
"ocean/global_convergence/cosine_bell/QU150/mesh/mesh.nc": "global_convergence/cosine_bell/QU150/mesh/mesh.210803.nc",
66+
"ocean/global_convergence/cosine_bell/QU150/mesh/graph.info": "global_convergence/cosine_bell/QU150/mesh/graph.210803.info",
67+
"ocean/global_convergence/cosine_bell/QU150/init/namelist.ocean": "global_convergence/cosine_bell/QU150/init/namelist.210803.ocean",
68+
"ocean/global_convergence/cosine_bell/QU150/init/initial_state.nc": "global_convergence/cosine_bell/QU150/init/initial_state.210803.nc"
69+
}

‎compass/ocean/suites/cosine_bell_cached_init.txt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
ocean/global_convergence/cosine_bell
2+
cached: QU60_mesh QU60_init QU90_mesh QU90_init QU120_mesh QU120_init
3+
cached: QU150_mesh QU150_init QU180_mesh QU180_init QU210_mesh QU210_init
4+
cached: QU240_mesh QU240_init

‎compass/ocean/suites/nightly.txt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,12 @@ ocean/global_ocean/QU240/EN4_1900/performance_test
2222
ocean/global_ocean/QU240/PHC_BGC/init
2323
ocean/global_ocean/QU240/PHC_BGC/performance_test
2424

25+
ocean/global_ocean/QUwISC240/mesh
26+
cached
27+
ocean/global_ocean/QUwISC240/PHC/init
28+
cached
29+
ocean/global_ocean/QUwISC240/PHC/performance_test
30+
2531
ocean/ice_shelf_2d/5km/z-star/restart_test
2632
ocean/ice_shelf_2d/5km/z-level/restart_test
2733

0 commit comments

Comments
 (0)