def _make_transect_masks(mesh_name, suffix, fcMask, logger, cores,
subdivision_threshold=10e3):
mesh_filename = 'restart.nc'
geojson_filename = '{}.geojson'.format(suffix)
mask_filename = '{}_{}.nc'.format(mesh_name, suffix)
fcMask.to_geojson(geojson_filename)
# these defaults may have been updated from config options -- pass them
# along to the subprocess
netcdf_format = mpas_tools.io.default_format
netcdf_engine = mpas_tools.io.default_engine
args = ['compute_mpas_transect_masks',
'-m', mesh_filename,
'-g', geojson_filename,
'-o', mask_filename,
'-t', 'edge',
'-s', '{}'.format(subdivision_threshold),
'--process_count', '{}'.format(cores),
'--add_edge_sign',
'--format', netcdf_format,
'--engine', netcdf_engine]
check_call(args, logger=logger)
# make links in output directory
output_dir = '../assembled_files/diagnostics/mpas_analysis/' \
'region_masks'
symlink('../../../../diagnostics_files/{}'.format(mask_filename),
'{}/{}'.format(output_dir, mask_filename))
def configure(self):
"""
Modify the configuration options for this test case
"""
with path('compass.landice.tests.enthalpy_benchmark', 'README') as \
target:
symlink(str(target), '{}/README'.format(self.work_dir))
def run(self):
"""
Run this step of the testcase
"""
with xarray.open_dataset('restart.nc') as ds:
mesh_short_name = ds.attrs['MPAS_Mesh_Short_Name']
mesh_prefix = ds.attrs['MPAS_Mesh_Prefix']
prefix = 'MPAS_Mesh_{}'.format(mesh_prefix)
creation_date = ds.attrs['{}_Version_Creation_Date'.format(prefix)]
try:
os.makedirs('../assembled_files/inputdata/ocn/mpas-o/{}'.format(
mesh_short_name))
except OSError:
pass
source_filename = 'restart.nc'
dest_filename = 'mpaso.{}.{}.nc'.format(mesh_short_name, creation_date)
with xarray.open_dataset(source_filename) as ds:
ds.load()
ds = ds.drop_vars('xtime')
write_netcdf(ds, dest_filename)
symlink(
'../../../../../ocean_initial_condition/{}'.format(dest_filename),
'../assembled_files/inputdata/ocn/mpas-o/{}/{}'.format(
mesh_short_name, dest_filename))
def configure(self):
"""
Modify the configuration options for this test case
"""
configure_global_ocean(test_case=self, mesh=self.mesh, init=self.init)
with path('compass.ocean.tests.global_ocean.files_for_e3sm',
'README') as target:
symlink(str(target), '{}/README'.format(self.work_dir))
def _make_moc_masks(mesh_short_name, logger, cores):
gf = GeometricFeatures()
mesh_filename = 'restart.nc'
function, prefix, date = get_aggregator_by_name('MOC Basins')
fcMask = function(gf)
suffix = '{}{}'.format(prefix, date)
geojson_filename = '{}.geojson'.format(suffix)
mask_filename = '{}_{}.nc'.format(mesh_short_name, suffix)
fcMask.to_geojson(geojson_filename)
# these defaults may have been updated from config options -- pass them
# along to the subprocess
netcdf_format = mpas_tools.io.default_format
netcdf_engine = mpas_tools.io.default_engine
args = ['compute_mpas_region_masks',
'-m', mesh_filename,
'-g', geojson_filename,
'-o', mask_filename,
'-t', 'cell',
'--process_count', '{}'.format(cores),
'--format', netcdf_format,
'--engine', netcdf_engine]
check_call(args, logger=logger)
mask_and_transect_filename = '{}_mocBasinsAndTransects{}.nc'.format(
mesh_short_name, date)
dsMesh = xarray.open_dataset(mesh_filename)
dsMask = xarray.open_dataset(mask_filename)
dsMasksAndTransects = add_moc_southern_boundary_transects(
dsMask, dsMesh, logger=logger)
write_netcdf(dsMasksAndTransects, mask_and_transect_filename,
char_dim_name='StrLen')
# make links in output directories (both inputdata and diagnostics)
output_dir = '../assembled_files/inputdata/ocn/mpas-o/{}'.format(
mesh_short_name)
symlink(
'../../../../../diagnostics_files/{}'.format(
mask_and_transect_filename),
'{}/{}'.format(output_dir, mask_and_transect_filename))
output_dir = '../assembled_files/diagnostics/mpas_analysis/' \
'region_masks'
symlink(
'../../../../diagnostics_files/{}'.format(
mask_and_transect_filename),
'{}/{}'.format(output_dir, mask_and_transect_filename))
def run(self):
"""
Run this step of the test case
"""
section = self.config['make_diagnostics_files']
mesh_name = section.get('mesh_name')
mesh_filename = section.get('mesh_filename')
cores = section.getint('cores')
with_ice_shelf_cavities = section.getboolean('with_ice_shelf_cavities')
symlink(os.path.join('..', mesh_filename), 'restart.nc')
make_diagnostics_files(self.config, self.logger, mesh_name,
with_ice_shelf_cavities, cores)
def run(self):
"""
Run this step of the testcase
"""
logger = self.logger
with xarray.open_dataset('restart.nc') as ds:
mesh_short_name = ds.attrs['MPAS_Mesh_Short_Name']
mesh_prefix = ds.attrs['MPAS_Mesh_Prefix']
prefix = 'MPAS_Mesh_{}'.format(mesh_prefix)
creation_date = ds.attrs['{}_Version_Creation_Date'.format(prefix)]
try:
os.makedirs('../assembled_files/inputdata/ocn/mpas-o/{}'.format(
mesh_short_name))
except OSError:
pass
symlink('graph.info', 'mpas-o.graph.info.{}'.format(creation_date))
nCells = sum(1 for _ in open('graph.info'))
min_graph_size = int(nCells / 6000)
max_graph_size = int(nCells / 100)
logger.info('Creating graph files between {} and {}'.format(
min_graph_size, max_graph_size))
n_power2 = 2**np.arange(1, 21)
n_multiples12 = 12 * np.arange(1, 9)
n = n_power2
for power10 in range(3):
n = np.concatenate([n, 10**power10 * n_multiples12])
for index in range(len(n)):
if min_graph_size <= n[index] <= max_graph_size:
args = [
'gpmetis', 'mpas-o.graph.info.{}'.format(creation_date),
'{}'.format(n[index])
]
check_call(args, logger)
# create link in assembled files directory
files = glob('mpas-o.graph.info.*')
dest_path = '../assembled_files/inputdata/ocn/mpas-o/{}'.format(
mesh_short_name)
for file in files:
symlink('../../../../../ocean_graph_partition/{}'.format(file),
'{}/{}'.format(dest_path, file))
def run(self):
"""
Run this step of the testcase
"""
with_ice_shelf_cavities = self.with_ice_shelf_cavities
with xarray.open_dataset('restart.nc') as ds:
mesh_short_name = ds.attrs['MPAS_Mesh_Short_Name']
mesh_prefix = ds.attrs['MPAS_Mesh_Prefix']
prefix = 'MPAS_Mesh_{}'.format(mesh_prefix)
creation_date = ds.attrs['{}_Version_Creation_Date'.format(prefix)]
try:
os.makedirs(
'../assembled_files/inputdata/ocn/mpas-seaice/{}'.format(
mesh_short_name))
except OSError:
pass
dest_filename = 'mpassi.{}.{}.nc'.format(mesh_short_name,
creation_date)
keep_vars = [
'areaCell', 'cellsOnCell', 'edgesOnCell', 'fCell', 'indexToCellID',
'latCell', 'lonCell', 'meshDensity', 'nEdgesOnCell',
'verticesOnCell', 'xCell', 'yCell', 'zCell', 'angleEdge',
'cellsOnEdge', 'dcEdge', 'dvEdge', 'edgesOnEdge', 'fEdge',
'indexToEdgeID', 'latEdge', 'lonEdge', 'nEdgesOnCell',
'nEdgesOnEdge', 'verticesOnEdge', 'weightsOnEdge', 'xEdge',
'yEdge', 'zEdge', 'areaTriangle', 'cellsOnVertex', 'edgesOnVertex',
'fVertex', 'indexToVertexID', 'kiteAreasOnVertex', 'latVertex',
'lonVertex', 'xVertex', 'yVertex', 'zVertex'
]
if with_ice_shelf_cavities:
keep_vars.append('landIceMask')
with xarray.open_dataset('restart.nc') as ds:
ds.load()
ds = ds[keep_vars]
write_netcdf(ds, dest_filename)
symlink(
'../../../../../seaice_initial_condition/{}'.format(dest_filename),
'../assembled_files/inputdata/ocn/mpas-seaice/{}/{}'.format(
mesh_short_name, dest_filename))
def run(self):
"""
Run this step of the testcase
"""
with_ice_shelf_cavities = self.with_ice_shelf_cavities
with xarray.open_dataset('restart.nc') as ds:
mesh_short_name = ds.attrs['MPAS_Mesh_Short_Name']
mesh_prefix = ds.attrs['MPAS_Mesh_Prefix']
prefix = 'MPAS_Mesh_{}'.format(mesh_prefix)
creation_date = ds.attrs['{}_Version_Creation_Date'.format(prefix)]
try:
os.makedirs('../assembled_files/inputdata/ocn/mpas-o/{}'.format(
mesh_short_name))
except OSError:
pass
if with_ice_shelf_cavities:
nomask_str = '.nomask'
else:
nomask_str = ''
scrip_filename = 'ocean.{}{}.scrip.{}.nc'.format(
mesh_short_name, nomask_str, creation_date)
scrip_from_mpas('restart.nc', scrip_filename)
symlink('../../../../../scrip/{}'.format(scrip_filename),
'../assembled_files/inputdata/ocn/mpas-o/{}/{}'.format(
mesh_short_name, scrip_filename))
if with_ice_shelf_cavities:
scrip_mask_filename = 'ocean.{}.mask.scrip.{}.nc'.format(
mesh_short_name, creation_date)
scrip_from_mpas('restart.nc', scrip_mask_filename,
useLandIceMask=True)
symlink(
'../../../../../scrip/{}'.format(
scrip_mask_filename),
'../assembled_files/inputdata/ocn/mpas-o/{}/{}'.format(
mesh_short_name, scrip_mask_filename))
def setup_case(path, test_case, config_file, machine, machine_info, work_dir,
baseline_dir, mpas_model_path, cached_steps):
"""
Set up one or more test cases
Parameters
----------
path : str
Relative path for a test cases to set up
test_case : compass.TestCase
A test case to set up
config_file : str
Configuration file with custom options for setting up and running test
cases
machine : str
The name of one of the machines with defined config options, which can
be listed with ``compass list --machines``
machine_info : mache.MachineInfo
Information about the machine, to be included in the config options
and passed along to each step
work_dir : str
A directory that will serve as the base for creating case directories
baseline_dir : str
Location of baseslines that can be compared to
mpas_model_path : str
The relative or absolute path to the root of a branch where the MPAS
model has been built
cached_steps : list of str
Which steps (if any) should be cached. If all steps should be cached,
the first entry is "_all"
"""
print(' {}'.format(path))
config = configparser.ConfigParser(
interpolation=configparser.ExtendedInterpolation())
# start with default compass config options
add_config(config, 'compass', 'default.cfg')
# add the E3SM config options from mache
merge_other_config(config, machine_info.config)
# add the compass machine config file
if machine is None:
machine = 'default'
add_config(config, 'compass.machines', '{}.cfg'.format(machine))
# add the config options for the MPAS core
mpas_core = test_case.mpas_core.name
add_config(config, 'compass.{}'.format(mpas_core),
'{}.cfg'.format(mpas_core))
# add the config options for the test group (if defined)
test_group = test_case.test_group.name
add_config(config, 'compass.{}.tests.{}'.format(mpas_core, test_group),
'{}.cfg'.format(test_group), exception=False)
# add the config options for the test case (if defined)
add_config(config, test_case.__module__,
'{}.cfg'.format(test_case.name), exception=False)
test_case_dir = os.path.join(work_dir, path)
try:
os.makedirs(test_case_dir)
except OSError:
pass
test_case.work_dir = test_case_dir
test_case.base_work_dir = work_dir
# add the custom config file once before calling configure() in case we
# need to use the config options from there
if config_file is not None:
config.read(config_file)
# add config options specific to the test case
test_case.config = config
test_case.configure()
# add the custom config file (again) last, so these options are the
# defaults
if config_file is not None:
config.read(config_file)
# add the baseline directory for this test case
if baseline_dir is not None:
test_case.baseline_dir = os.path.join(baseline_dir, path)
# set the mpas_model path from the command line if provided
if mpas_model_path is not None:
config.set('paths', 'mpas_model', mpas_model_path)
config.set('test_case', 'steps_to_run', ' '.join(test_case.steps_to_run))
# make sure all paths in the paths, namelists and streams sections are
# absolute paths
ensure_absolute_paths(config)
# write out the config file
test_case_config = '{}.cfg'.format(test_case.name)
test_case.config_filename = test_case_config
with open(os.path.join(test_case_dir, test_case_config), 'w') as f:
config.write(f)
if len(cached_steps) > 0 and cached_steps[0] == '_all':
cached_steps = list(test_case.steps.keys())
if len(cached_steps) > 0:
print_steps = ' '.join(cached_steps)
print(f' steps with cached outputs: {print_steps}')
for step_name in cached_steps:
test_case.steps[step_name].cached = True
# iterate over steps
for step in test_case.steps.values():
# make the step directory if it doesn't exist
step_dir = os.path.join(work_dir, step.path)
try:
os.makedirs(step_dir)
except OSError:
pass
symlink(os.path.join(test_case_dir, test_case_config),
os.path.join(step_dir, test_case_config))
step.work_dir = step_dir
step.base_work_dir = work_dir
step.config_filename = test_case_config
step.config = config
# set up the step
step.setup()
# process input, output, namelist and streams files
step.process_inputs_and_outputs()
# wait until we've set up all the steps before pickling because steps may
# need other steps to be set up
for step in test_case.steps.values():
# pickle the test case and step for use at runtime
pickle_filename = os.path.join(step.work_dir, 'step.pickle')
with open(pickle_filename, 'wb') as handle:
pickle.dump((test_case, step), handle,
protocol=pickle.HIGHEST_PROTOCOL)
# pickle the test case and step for use at runtime
pickle_filename = os.path.join(test_case.work_dir, 'test_case.pickle')
with open(pickle_filename, 'wb') as handle:
pickle.dump(test_case, handle, protocol=pickle.HIGHEST_PROTOCOL)
if 'LOAD_COMPASS_ENV' in os.environ:
script_filename = os.environ['LOAD_COMPASS_ENV']
# make a symlink to the script for loading the compass conda env.
symlink(script_filename, os.path.join(test_case_dir,
'load_compass_env.sh'))
def setup_cases(tests=None, numbers=None, config_file=None, machine=None,
work_dir=None, baseline_dir=None, mpas_model_path=None,
suite_name='custom', cached=None):
"""
Set up one or more test cases
Parameters
----------
tests : list of str, optional
Relative paths for a test cases to set up
numbers : list of str, optional
Case numbers to setup, as listed from ``compass list``, optionally with
a suffix ``c`` to indicate that all steps in that test case should be
cached
config_file : str, optional
Configuration file with custom options for setting up and running test
cases
machine : str, optional
The name of one of the machines with defined config options, which can
be listed with ``compass list --machines``
work_dir : str, optional
A directory that will serve as the base for creating case directories
baseline_dir : str, optional
Location of baseslines that can be compared to
mpas_model_path : str, optional
The relative or absolute path to the root of a branch where the MPAS
model has been built
suite_name : str, optional
The name of the test suite if tests are being set up through a test
suite or ``'custom'`` if not
cached : list of list of str, optional
For each test in ``tests``, which steps (if any) should be cached,
or a list with "_all" as the first entry if all steps in the test case
should be cached
Returns
-------
test_cases : dict of compass.TestCase
A dictionary of test cases, with the relative path in the work
directory as keys
"""
if machine is None and 'COMPASS_MACHINE' in os.environ:
machine = os.environ['COMPASS_MACHINE']
if machine is None:
machine = discover_machine()
if machine is None:
machine_info = MachineInfo(machine='unknown')
else:
machine_info = MachineInfo(machine=machine)
if config_file is None and machine is None:
raise ValueError('At least one of config_file and machine is needed.')
if tests is None and numbers is None:
raise ValueError('At least one of tests or numbers is needed.')
if cached is not None:
if tests is None:
warnings.warn('Ignoring "cached" argument because "tests" was '
'not provided')
elif len(cached) != len(tests):
raise ValueError('A list of cached steps must be provided for '
'each test in "tests"')
if work_dir is None:
work_dir = os.getcwd()
work_dir = os.path.abspath(work_dir)
mpas_cores = get_mpas_cores()
all_test_cases = dict()
for mpas_core in mpas_cores:
for test_group in mpas_core.test_groups.values():
for test_case in test_group.test_cases.values():
all_test_cases[test_case.path] = test_case
test_cases = dict()
cached_steps = dict()
if numbers is not None:
keys = list(all_test_cases)
for number in numbers:
cache_all = False
if number.endswith('c'):
cache_all = True
number = int(number[:-1])
else:
number = int(number)
if number >= len(keys):
raise ValueError('test number {} is out of range. There are '
'only {} tests.'.format(number, len(keys)))
path = keys[number]
if cache_all:
cached_steps[path] = ['_all']
else:
cached_steps[path] = list()
test_cases[path] = all_test_cases[path]
if tests is not None:
for index, path in enumerate(tests):
if path not in all_test_cases:
raise ValueError('Test case with path {} is not in '
'test_cases'.format(path))
if cached is not None:
cached_steps[path] = cached[index]
else:
cached_steps[path] = list()
test_cases[path] = all_test_cases[path]
# get the MPAS core of the first test case. We'll assume all tests are
# for this core
first_path = next(iter(test_cases))
mpas_core = test_cases[first_path].mpas_core.name
provenance.write(work_dir, test_cases, mpas_core=mpas_core,
config_filename=config_file,
mpas_model_path=mpas_model_path)
print('Setting up test cases:')
for path, test_case in test_cases.items():
setup_case(path, test_case, config_file, machine, machine_info,
work_dir, baseline_dir, mpas_model_path,
cached_steps=cached_steps[path])
test_suite = {'name': suite_name,
'test_cases': test_cases,
'work_dir': work_dir}
# pickle the test or step dictionary for use at runtime
pickle_file = os.path.join(test_suite['work_dir'],
'{}.pickle'.format(suite_name))
with open(pickle_file, 'wb') as handle:
pickle.dump(test_suite, handle, protocol=pickle.HIGHEST_PROTOCOL)
if 'LOAD_COMPASS_ENV' in os.environ:
script_filename = os.environ['LOAD_COMPASS_ENV']
# make a symlink to the script for loading the compass conda env.
symlink(script_filename, os.path.join(work_dir, 'load_compass_env.sh'))
max_cores, max_of_min_cores = _get_required_cores(test_cases)
print('target cores: {}'.format(max_cores))
print('minimum cores: {}'.format(max_of_min_cores))
return test_cases
def make_diagnostics_files(config, logger, mesh_short_name,
with_ice_shelf_cavities, cores):
"""
Run this step of the testcase
Parameters
----------
config : configparser.ConfigParser
Configuration options for this test case
logger : logging.Logger
A logger for output from the step
mesh_short_name : str
The E3SM short name of the mesh
with_ice_shelf_cavities : bool
Whether the mesh has ice-shelf cavities
cores : int
The number of cores to use to build mapping files
"""
for directory in [
'../assembled_files/inputdata/ocn/mpas-o/{}'.format(
mesh_short_name),
'../assembled_files/diagnostics/mpas_analysis/region_masks',
'../assembled_files/diagnostics/mpas_analysis/maps']:
try:
os.makedirs(directory)
except OSError:
pass
_make_moc_masks(mesh_short_name, logger, cores)
gf = GeometricFeatures()
region_groups = ['Antarctic Regions', 'Arctic Ocean Regions',
'Arctic Sea Ice Regions', 'Ocean Basins',
'Ocean Subbasins', 'ISMIP6 Regions']
if with_ice_shelf_cavities:
region_groups.append('Ice Shelves')
for region_group in region_groups:
function, prefix, date = get_aggregator_by_name(region_group)
suffix = '{}{}'.format(prefix, date)
fcMask = function(gf)
_make_region_masks(mesh_short_name, suffix=suffix, fcMask=fcMask,
logger=logger, cores=cores)
transect_groups = ['Transport Transects']
for transect_group in transect_groups:
function, prefix, date = get_aggregator_by_name(transect_group)
suffix = '{}{}'.format(prefix, date)
fcMask = function(gf)
_make_transect_masks(mesh_short_name, suffix=suffix, fcMask=fcMask,
logger=logger, cores=cores)
_make_analysis_lat_lon_map(config, mesh_short_name, cores, logger)
_make_analysis_polar_map(config, mesh_short_name,
projection='antarctic', cores=cores,
logger=logger)
_make_analysis_polar_map(config, mesh_short_name, projection='arctic',
cores=cores, logger=logger)
# make links in output directory
files = glob.glob('map_*')
# make links in output directory
output_dir = '../assembled_files/diagnostics/mpas_analysis/maps'
for filename in files:
symlink('../../../../diagnostics_files/{}'.format(filename),
'{}/{}'.format(output_dir, filename))
def process_inputs_and_outputs(self):
"""
Process the inputs to and outputs from a step added with
:py:meth:`compass.Step.add_input_file` and
:py:meth:`compass.Step.add_output_file`. This includes downloading
files, making symlinks, and converting relative paths to absolute
paths.
Also generates namelist and streams files
"""
mpas_core = self.mpas_core.name
step_dir = self.work_dir
config = self.config
# process the outputs first because cached outputs will add more inputs
if self.cached:
# forget about the inputs -- we won't used them, but we will add
# the cached outputs as inputs
self.input_data = list()
for output in self.outputs:
filename = os.path.join(self.path, output)
if filename not in self.mpas_core.cached_files:
raise ValueError(f'The file {filename} has not been added '
f'to the cache database')
target = self.mpas_core.cached_files[filename]
self.add_input_file(filename=output,
target=target,
database='compass_cache')
inputs = []
databases_with_downloads = set()
for entry in self.input_data:
filename = entry['filename']
target = entry['target']
database = entry['database']
url = entry['url']
work_dir_target = entry['work_dir_target']
package = entry['package']
copy = entry['copy']
if filename == '<<<model>>>':
model = self.config.get('executables', 'model')
filename = os.path.basename(model)
target = os.path.abspath(model)
if package is not None:
if target is None:
target = filename
with path(package, target) as package_path:
target = str(package_path)
if work_dir_target is not None:
target = os.path.join(self.base_work_dir, work_dir_target)
if target is not None:
download_target = target
else:
download_target = filename
download_path = None
if database is not None:
# we're downloading a file to a cache of a database (if it's
# not already there.
if url is None:
base_url = config.get('download', 'server_base_url')
core_path = config.get('download', 'core_path')
url = '{}/{}/{}'.format(base_url, core_path, database)
url = '{}/{}'.format(url, target)
database_root = config.get(
'paths', '{}_database_root'.format(mpas_core))
download_path = os.path.join(database_root, database,
download_target)
if not os.path.exists(download_path):
database_subdir = os.path.join(database_root, database)
databases_with_downloads.add(database_subdir)
elif url is not None:
download_path = download_target
if url is not None:
download_target = download(url, download_path, config)
if target is not None:
# this is the absolute path that we presumably want
target = download_target
if target is not None:
filepath = os.path.join(step_dir, filename)
if copy:
shutil.copy(target, filepath)
else:
symlink(target, filepath)
inputs.append(target)
else:
inputs.append(filename)
if len(databases_with_downloads) > 0:
self._fix_permissions(databases_with_downloads)
# convert inputs and outputs to absolute paths
self.inputs = [
os.path.abspath(os.path.join(step_dir, filename))
for filename in inputs
]
self.outputs = [
os.path.abspath(os.path.join(step_dir, filename))
for filename in self.outputs
]
self._generate_namelists()
self._generate_streams()
def adjust_ssh(variable, iteration_count, step):
"""
Adjust the sea surface height or land-ice pressure to be dynamically
consistent with one another. A series of short model runs are performed,
each with
Parameters
----------
variable : {'ssh', 'landIcePressure'}
The variable to adjust
iteration_count : int
The number of iterations of adjustment
step : compass.Step
the step for performing SSH or land-ice pressure adjustment
"""
cores = step.cores
config = step.config
logger = step.logger
out_filename = None
if variable not in ['ssh', 'landIcePressure']:
raise ValueError("Unknown variable to modify: {}".format(variable))
step.update_namelist_pio('namelist.ocean')
partition(cores, config, logger)
for iterIndex in range(iteration_count):
logger.info(" * Iteration {}/{}".format(iterIndex + 1,
iteration_count))
in_filename = 'adjusting_init{}.nc'.format(iterIndex)
out_filename = 'adjusting_init{}.nc'.format(iterIndex + 1)
symlink(in_filename, 'adjusting_init.nc')
logger.info(" * Running forward model")
run_model(step, update_pio=False, partition_graph=False)
logger.info(" - Complete")
logger.info(" * Updating SSH or land-ice pressure")
with xarray.open_dataset(in_filename) as ds:
# keep the data set with Time for output
ds_out = ds
ds = ds.isel(Time=0)
on_a_sphere = ds.attrs['on_a_sphere'].lower() == 'yes'
initSSH = ds.ssh
if 'minLevelCell' in ds:
minLevelCell = ds.minLevelCell - 1
else:
minLevelCell = xarray.zeros_like(ds.maxLevelCell)
with xarray.open_dataset('output_ssh.nc') as ds_ssh:
# get the last time entry
ds_ssh = ds_ssh.isel(Time=ds_ssh.sizes['Time'] - 1)
finalSSH = ds_ssh.ssh
topDensity = ds_ssh.density.isel(nVertLevels=minLevelCell)
mask = numpy.logical_and(ds.maxLevelCell > 0,
ds.modifyLandIcePressureMask == 1)
deltaSSH = mask * (finalSSH - initSSH)
# then, modify the SSH or land-ice pressure
if variable == 'ssh':
ssh = finalSSH.expand_dims(dim='Time', axis=0)
ds_out['ssh'] = ssh
# also update the landIceDraft variable, which will be used to
# compensate for the SSH due to land-ice pressure when
# computing sea-surface tilt
ds_out['landIceDraft'] = ssh
# we also need to stretch layerThickness to be compatible with
# the new SSH
stretch = ((finalSSH + ds.bottomDepth) /
(initSSH + ds.bottomDepth))
ds_out['layerThickness'] = ds_out.layerThickness * stretch
landIcePressure = ds.landIcePressure.values
else:
# Moving the SSH up or down by deltaSSH would change the
# land-ice pressure by density(SSH)*g*deltaSSH. If deltaSSH is
# positive (moving up), it means the land-ice pressure is too
# small and if deltaSSH is negative (moving down), it means
# land-ice pressure is too large, the sign of the second term
# makes sense.
gravity = constants['SHR_CONST_G']
deltaLandIcePressure = topDensity * gravity * deltaSSH
landIcePressure = numpy.maximum(
0.0, ds.landIcePressure + deltaLandIcePressure)
ds_out['landIcePressure'] = \
landIcePressure.expand_dims(dim='Time', axis=0)
finalSSH = initSSH
write_netcdf(ds_out, out_filename)
# Write the largest change in SSH and its lon/lat to a file
with open('maxDeltaSSH_{:03d}.log'.format(iterIndex), 'w') as \
log_file:
mask = landIcePressure > 0.
iCell = numpy.abs(deltaSSH.where(mask)).argmax().values
ds_cell = ds.isel(nCells=iCell)
if on_a_sphere:
coords = 'lon/lat: {:f} {:f}'.format(
numpy.rad2deg(ds_cell.lonCell.values),
numpy.rad2deg(ds_cell.latCell.values))
else:
coords = 'x/y: {:f} {:f}'.format(
1e-3 * ds_cell.xCell.values,
1e-3 * ds_cell.yCell.values)
string = 'deltaSSHMax: {:g}, {}'.format(
deltaSSH.isel(nCells=iCell).values, coords)
logger.info(' {}'.format(string))
log_file.write('{}\n'.format(string))
string = 'ssh: {:g}, landIcePressure: {:g}'.format(
finalSSH.isel(nCells=iCell).values,
landIcePressure.isel(nCells=iCell).values)
logger.info(' {}'.format(string))
log_file.write('{}\n'.format(string))
logger.info(" - Complete\n")
if out_filename is not None:
shutil.copy(out_filename, 'adjusted_init.nc')