-
Notifications
You must be signed in to change notification settings - Fork 38
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add
files_for_e3sm steps for ocean and sea-ice meshes
These also need to be added to `inputdata` in `shared/meshes/mpas`
- Loading branch information
Showing
6 changed files
with
242 additions
and
27 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
115 changes: 115 additions & 0 deletions
115
compass/ocean/tests/global_ocean/files_for_e3sm/ocean_mesh.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,115 @@ | ||
| import os | ||
|
|
||
| import numpy as np | ||
| import xarray as xr | ||
| 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.vertical import ( | ||
| compute_cell_mask, | ||
| compute_ssh_from_layer_thickness, | ||
| compute_zmid_from_layer_thickness, | ||
| ) | ||
|
|
||
|
|
||
| class OceanMesh(FilesForE3SMStep): | ||
| """ | ||
| A step for creating an MPAS-Ocean mesh from variables from an MPAS-Ocean | ||
| initial state file | ||
| """ | ||
| 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='ocean_mesh') | ||
|
|
||
| # 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 run(self): | ||
| """ | ||
| Run this step of the testcase | ||
| """ | ||
| super().run() | ||
|
|
||
| dest_filename = f'{self.mesh_short_name}.{self.creation_date}.nc' | ||
|
|
||
| with xr.open_dataset('initial_state.nc') as ds: | ||
|
|
||
| keep_vars = self.mesh_vars + [ | ||
| 'refBottomDepth', 'vertCoordMovementWeights', | ||
| 'bottomDepth', 'maxLevelCell', | ||
| 'layerThickness', 'restingThickness' | ||
| ] | ||
|
|
||
| if 'minLevelCell' in ds: | ||
| keep_vars.append('minLevelCell') | ||
|
|
||
| if self.with_ice_shelf_cavities: | ||
| keep_vars = keep_vars + [ | ||
| 'landIceMask', 'landIceDraft', 'landIceFraction' | ||
| ] | ||
|
|
||
| ds = ds[keep_vars] | ||
| ds.load() | ||
|
|
||
| for attr in list(ds.attrs): | ||
| # drop config options from global attributes | ||
| if attr.startswith('config_'): | ||
| ds.attrs.pop(attr) | ||
|
|
||
| ref_bot_depth = ds.refBottomDepth.values | ||
| interfaces = np.append([0], ref_bot_depth) | ||
|
|
||
| if 'minLevelCell' not in ds: | ||
| ds['minLevelCell'] = xr.ones_like(ds.maxLevelCell) | ||
| ds.minLevelCell.attrs['long_name'] = \ | ||
| 'Index to the first active ocean cell in each column.' | ||
|
|
||
| ds['refTopDepth'] = ('nVertLevels', interfaces[0:-1]) | ||
| ds.refTopDepth.attrs['units'] = 'm' | ||
| ds.refTopDepth.attrs['long_name'] = \ | ||
| "Reference depth of ocean for each vertical level. Used in " \ | ||
| "'z-level' type runs." | ||
| ds['refZMid'] = ('nVertLevels', | ||
| -0.5 * (interfaces[1:] + interfaces[0:-1])) | ||
| ds.refZMid.attrs['units'] = 'm' | ||
| ds.refZMid.attrs['long_name'] = \ | ||
| 'Reference mid z-coordinate of ocean for each vertical ' \ | ||
| 'level. This has a negative value.' | ||
| ds['refLayerThickness'] = ('nVertLevels', | ||
| interfaces[1:] - interfaces[0:-1]) | ||
| ds.refLayerThickness.attrs['units'] = 'm' | ||
| ds.refLayerThickness.attrs['long_name'] = \ | ||
| 'Reference layer thickness of ocean for each vertical level.' | ||
|
|
||
| ds['cellMask'] = compute_cell_mask( | ||
| ds.minLevelCell - 1, ds.maxLevelCell - 1, | ||
| ds.sizes['nVertLevels']) | ||
| ds.cellMask.attrs['long_name'] = \ | ||
| 'Mask on cells that determines if computations should be ' \ | ||
| 'done on cells.' | ||
| ds['ssh'] = compute_ssh_from_layer_thickness( | ||
| ds.layerThickness, ds.bottomDepth, ds.cellMask) | ||
| ds.ssh.attrs['units'] = 'm' | ||
| ds.ssh.attrs['long_name'] = 'sea surface height' | ||
| ds['zMid'] = compute_zmid_from_layer_thickness( | ||
| ds.layerThickness, ds.ssh, ds.cellMask) | ||
| ds.zMid.attrs['units'] = 'm' | ||
| ds.zMid.attrs['long_name'] = \ | ||
| 'z-coordinate of the mid-depth of the layer' | ||
|
|
||
| write_netcdf(ds, dest_filename) | ||
|
|
||
| symlink(os.path.abspath(dest_filename), | ||
| f'{self.ocean_mesh_dir}/{dest_filename}') |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
49 changes: 49 additions & 0 deletions
49
compass/ocean/tests/global_ocean/files_for_e3sm/seaice_mesh.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,49 @@ | ||
| import os | ||
|
|
||
| import xarray as xr | ||
| 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, | ||
| ) | ||
|
|
||
|
|
||
| class SeaiceMesh(FilesForE3SMStep): | ||
| """ | ||
| A step for creating an MPAS-Seaice mesh from variables from an MPAS-Ocean | ||
| initial state file | ||
| """ | ||
| 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_mesh') | ||
|
|
||
| # 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 run(self): | ||
| """ | ||
| Run this step of the testcase | ||
| """ | ||
| super().run() | ||
|
|
||
| dest_filename = f'{self.mesh_short_name}.{self.creation_date}.nc' | ||
|
|
||
| with xr.open_dataset('initial_state.nc') as ds: | ||
|
|
||
| keep_vars = self.mesh_vars | ||
| ds = ds[keep_vars] | ||
| ds.load() | ||
| write_netcdf(ds, dest_filename) | ||
|
|
||
| symlink(os.path.abspath(dest_filename), | ||
| f'{self.seaice_mesh_dir}/{dest_filename}') |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters