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

import os

import xarray
from mpas_tools.io import write_netcdf

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.metadata import (
    add_mesh_and_init_metadata,
)


[docs] class OceanInitialCondition(FilesForE3SMStep): """ A step for creating an E3SM ocean initial condition from the results of a dynamic-adjustment process to dissipate fast waves Attributes ---------- add_metadata : bool Whether to add metadata to the ocean initial condition """
[docs] def __init__(self, test_case, add_metadata=False): """ Create a new step Parameters ---------- test_case : compass.ocean.tests.global_ocean.files_for_e3sm.FilesForE3SM The test case this step belongs to add_metadata : bool, optional Whether to add metadata to the ocean initial condition """ # noqa: E501 super().__init__(test_case, name='ocean_initial_condition') self.add_metadata = add_metadata
# 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 """ super().run() source_filename = 'restart.nc' dest_filename = f'mpaso.{self.mesh_short_name}.{self.creation_date}.nc' with xarray.open_dataset(source_filename) as ds: ds.load() if 'xtime' in ds.data_vars: ds = ds.drop_vars('xtime') write_netcdf(ds, dest_filename) if self.add_metadata: add_mesh_and_init_metadata([dest_filename], self.config, source_filename) symlink( os.path.abspath(dest_filename), f'{self.ocean_inputdata_dir}/{dest_filename}')