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

import os
import xarray
import numpy as np
from glob import glob

from mpas_tools.logging import check_call

from import symlink
from compass.step import Step

[docs]class OceanGraphPartition(Step): """ A step for creating an E3SM ocean initial condition from the results of a dynamic-adjustment process to dissipate fast waves """
[docs] def __init__(self, test_case, mesh, restart_filename): """ Create a new step Parameters ---------- test_case : compass.ocean.tests.global_ocean.files_for_e3sm.FilesForE3SM The test case this step belongs to mesh : compass.ocean.tests.global_ocean.mesh.Mesh The test case that creates the mesh used by this test case restart_filename : str A restart file from the end of the dynamic adjustment test case to use as the basis for an E3SM initial condition """ super().__init__(test_case, name='ocean_graph_partition', cores=1, min_cores=1, threads=1) self.add_input_file(filename='README', target='../README') self.add_input_file(filename='', target='../{}'.format(restart_filename)) mesh_path = mesh.mesh_step.path self.add_input_file( filename='', work_dir_target='{}/'.format(mesh_path))
# 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.
[docs] def run(self): """ Run this step of the testcase """ logger = self.logger with xarray.open_dataset('') 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('', '{}'.format(creation_date)) nCells = sum(1 for _ in open('')) min_graph_size = int(nCells / 6000) max_graph_size = int(nCells / 100)'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', '{}'.format(creation_date), '{}'.format(n[index])] check_call(args, logger) # create link in assembled files directory files = glob('*') 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))