Source code for compass.ocean.tests.global_ocean.files_for_e3sm.seaice_graph_partition

import os
from glob import glob

import numpy as np
import xarray as xr
from mpas_tools.io import write_netcdf
from mpas_tools.logging import check_call
from pyremap import MpasCellMeshDescriptor, Remapper

from compass.io import symlink
from compass.ocean.tests.global_ocean.files_for_e3sm.files_for_e3sm_step import (  # noqa: E501
    FilesForE3SMStep,
)
from compass.ocean.tests.global_ocean.files_for_e3sm.graph_partition import (
    get_core_list,
)


[docs] class SeaiceGraphPartition(FilesForE3SMStep): """ A step for creating graph partition files for the sea-ice mesh """
[docs] def __init__(self, test_case): """ Create a new step Parameters ---------- test_case : compass.ocean.tests.global_ocean.files_for_e3sm.FilesForE3SM The test case this step belongs to """ # noqa: E501 super().__init__(test_case, name='seaice_graph_partition', ntasks=36, min_tasks=1) for filename in ['icePresent_QU60km_polar.nc', 'seaice_QU60km_polar.nc']: self.add_input_file(filename=filename, target=filename, database='partition', database_component='seaice')
# for now, we won't define any outputs because they include the mesh # short name, which is not known at setup time. Currently, this is # safe because no other steps depend on the outputs of this one. def setup(self): """ setup input files based on config options """ super().setup() graph_filename = self.config.get('files_for_e3sm', 'graph_filename') if graph_filename != 'autodetect': graph_filename = os.path.normpath(os.path.join( self.test_case.work_dir, graph_filename)) self.add_input_file(filename='graph.info', target=graph_filename)
[docs] def run(self): """ Run this step of the testcase """ super().run() logger = self.logger config = self.config plotting = config.getboolean('files_for_e3sm', 'plot_seaice_partitions') creation_date = self.creation_date mesh_filename = 'restart.nc' with xr.open_dataset(mesh_filename) as ds: ncells = ds.sizes['nCells'] # cullCell causes problems when the sea-ice partition tool tries # to add it again if 'cullCell' in ds: ds = ds.drop_vars(['cullCell']) write_netcdf(ds, 'mesh.nc') mesh_filename = 'mesh.nc' max_cells_per_core = config.getint('files_for_e3sm', 'max_cells_per_core') min_cells_per_core = config.getint('files_for_e3sm', 'min_cells_per_core') cores = get_core_list(ncells=ncells, max_cells_per_core=max_cells_per_core, min_cells_per_core=min_cells_per_core) logger.info(f'Creating graph files between {np.amin(cores)} and ' f'{np.amax(cores)}') if 1 in cores: args = ['touch', f'mpas-seaice.graph.info.{creation_date}.part.1'] check_call(args, logger) mapping_filename = _make_mapping_file( in_mesh_filename='seaice_QU60km_polar.nc', in_mesh_name='QU60km', out_mesh_filename=mesh_filename, out_mesh_name=self.mesh_short_name, ntasks=self.ntasks, config=self.config, logger=logger, method='bilinear') args = ['prepare_seaice_partitions', '-i', 'seaice_QU60km_polar.nc', '-p', 'icePresent_QU60km_polar.nc', '-m', mesh_filename, '-o', '.', '-w', mapping_filename] check_call(args, logger) args = ['create_seaice_partitions', '-m', mesh_filename, '-o', '.', '-p', f'mpas-seaice.graph.info.{creation_date}', '-g', 'gpmetis', '-n'] if plotting: args.append('--plotting') for ncores in cores: if ncores > ncells: raise ValueError('Can\t have more tasks than cells in a ' 'partition file.') if ncores <= 0: raise ValueError(f'Can\t make a partition for {ncores} tasks.') out_filename = \ f'mpas-seaice.graph.info.{creation_date}.part.{ncores}' if os.path.exists(out_filename): continue if ncores > 1: args.append(f'{ncores}') check_call(args, logger) # create link in assembled files directory inputdata_dir = os.path.join(self.seaice_inputdata_dir, 'partitions') try: os.makedirs(inputdata_dir) except FileExistsError: pass files = glob('mpas-seaice.graph.info.*') for file in files: symlink(os.path.abspath(file), f'{inputdata_dir}/{file}')
def _make_mapping_file(in_mesh_filename, in_mesh_name, out_mesh_filename, out_mesh_name, ntasks, config, logger, method): parallel_executable = config.get('parallel', 'parallel_executable') mapping_file_name = f'map_{in_mesh_name}_to_{out_mesh_name}_{method}.nc' in_descriptor = MpasCellMeshDescriptor(in_mesh_filename, in_mesh_name) out_descriptor = MpasCellMeshDescriptor(out_mesh_filename, out_mesh_name) remapper = Remapper(in_descriptor, out_descriptor, mapping_file_name) remapper.build_mapping_file(method=method, mpiTasks=ntasks, tempdir='.', logger=logger, esmf_parallel_exec=parallel_executable) return mapping_file_name