import numpy as np
import xarray as xr
from mpas_tools.io import write_netcdf
from mpas_tools.mesh.conversion import convert, cull
from mpas_tools.planar_hex import make_planar_hex_mesh
from polaris import Step
from polaris.ocean.vertical import init_vertical_coord
[docs]
class Init(Step):
"""
A step for creating a mesh and initial condition for single column
test cases
"""
[docs]
def __init__(self, component, indir, ideal_age=False):
"""
Create the step
Parameters
----------
component : polaris.Component
The component the step belongs to
indir : str
The subdirectory that the task belongs to, that this step will
go into a subdirectory of
ideal_age : bool, optional
Whether the initial condition should include the ideal age tracer
"""
super().__init__(component=component, name='init', indir=indir)
self.ideal_age = ideal_age
for file in [
'base_mesh.nc',
'culled_mesh.nc',
'culled_graph.info',
'initial_state.nc',
'forcing.nc',
]:
self.add_output_file(file)
[docs]
def run(self):
"""
Run this step of the test case
"""
logger = self.logger
config = self.config
section = config['single_column']
ideal_age = self.ideal_age
resolution = section.getfloat('resolution')
nx = section.getint('nx')
ny = section.getint('ny')
dc = 1e3 * resolution
ds_mesh = make_planar_hex_mesh(
nx=nx, ny=ny, dc=dc, nonperiodic_x=False, nonperiodic_y=False
)
write_netcdf(ds_mesh, 'base_mesh.nc')
ds_mesh = cull(ds_mesh, logger=logger)
ds_mesh = convert(
ds_mesh, graphInfoFileName='culled_graph.info', logger=logger
)
write_netcdf(ds_mesh, 'culled_mesh.nc')
ds = ds_mesh.copy()
x_cell = ds_mesh.xCell
bottom_depth = config.getfloat('vertical_grid', 'bottom_depth')
ds['bottomDepth'] = bottom_depth * xr.ones_like(x_cell)
ds['ssh'] = xr.zeros_like(x_cell)
init_vertical_coord(config, ds)
section = config['single_column']
surface_temperature = section.getfloat('surface_temperature')
temperature_gradient_mixed_layer = section.getfloat(
'temperature_gradient_mixed_layer'
)
temperature_difference_across_mixed_layer = section.getfloat(
'temperature_difference_across_mixed_layer'
)
temperature_gradient_interior = section.getfloat(
'temperature_gradient_interior'
)
mixed_layer_depth_temperature = section.getfloat(
'mixed_layer_depth_temperature'
)
surface_salinity = section.getfloat('surface_salinity')
salinity_gradient_mixed_layer = section.getfloat(
'salinity_gradient_mixed_layer'
)
salinity_difference_across_mixed_layer = section.getfloat(
'salinity_difference_across_mixed_layer'
)
salinity_gradient_interior = section.getfloat(
'salinity_gradient_interior'
)
mixed_layer_depth_salinity = section.getfloat(
'mixed_layer_depth_salinity'
)
coriolis_parameter = section.getfloat('coriolis_parameter')
u = section.getfloat('zonal_velocity')
v = section.getfloat('meridional_velocity')
z_mid = ds.refZMid
temperature_at_mixed_layer_depth = (
surface_temperature + temperature_difference_across_mixed_layer
)
temperature_vert = xr.where(
z_mid > -mixed_layer_depth_temperature,
surface_temperature + temperature_gradient_mixed_layer * z_mid,
temperature_at_mixed_layer_depth
+ temperature_gradient_interior
* (z_mid + mixed_layer_depth_temperature),
)
temperature_vert[0] = surface_temperature
temperature, _ = xr.broadcast(temperature_vert, x_cell)
temperature = temperature.transpose('nCells', 'nVertLevels')
temperature = temperature.expand_dims(dim='Time', axis=0)
salinity_at_mixed_layer_depth = (
surface_salinity + salinity_difference_across_mixed_layer
)
salinity_vert = xr.where(
z_mid > -mixed_layer_depth_salinity,
surface_salinity + salinity_gradient_mixed_layer * z_mid,
salinity_at_mixed_layer_depth
+ salinity_gradient_interior
* (z_mid + mixed_layer_depth_salinity),
)
salinity_vert[0] = surface_salinity
salinity, _ = xr.broadcast(salinity_vert, x_cell)
salinity = salinity.transpose('nCells', 'nVertLevels')
salinity = salinity.expand_dims(dim='Time', axis=0)
normal_velocity = u * np.cos(ds_mesh.angleEdge) + v * np.sin(
ds_mesh.angleEdge
)
normal_velocity, _ = xr.broadcast(normal_velocity, ds.refBottomDepth)
normal_velocity = normal_velocity.transpose('nEdges', 'nVertLevels')
normal_velocity = normal_velocity.expand_dims(dim='Time', axis=0)
if ideal_age:
ds['idealAgeTracers'] = xr.zeros_like(x_cell)
ds['temperature'] = temperature
ds['salinity'] = salinity
ds['normalVelocity'] = normal_velocity
ds['fCell'] = coriolis_parameter * xr.ones_like(x_cell)
ds['fEdge'] = coriolis_parameter * xr.ones_like(ds_mesh.xEdge)
ds['fVertex'] = coriolis_parameter * xr.ones_like(ds_mesh.xVertex)
ds.attrs['nx'] = nx
ds.attrs['ny'] = ny
ds.attrs['dc'] = dc
write_netcdf(ds, 'initial_state.nc')
# create forcing stream
ds_forcing = xr.Dataset()
forcing_array = xr.ones_like(temperature)
forcing_array_surface = xr.ones_like(ds.bottomDepth)
forcing_array_surface = forcing_array_surface.expand_dims(
dim='Time', axis=0
)
section = config['single_column_forcing']
temperature_piston_velocity = section.getfloat(
'temperature_piston_velocity'
)
salinity_piston_velocity = section.getfloat('salinity_piston_velocity')
temperature_surface_restoring_value = section.getfloat(
'temperature_surface_restoring_value'
)
salinity_surface_restoring_value = section.getfloat(
'salinity_surface_restoring_value'
)
temperature_interior_restoring_rate = section.getfloat(
'temperature_interior_restoring_rate'
)
salinity_interior_restoring_rate = section.getfloat(
'salinity_interior_restoring_rate'
)
latent_heat_flux = section.getfloat('latent_heat_flux')
sensible_heat_flux = section.getfloat('sensible_heat_flux')
shortwave_heat_flux = section.getfloat('shortwave_heat_flux')
evaporation_flux = section.getfloat('evaporation_flux')
rain_flux = section.getfloat('rain_flux')
river_runoff_flux = section.getfloat('river_runoff_flux')
ice_runoff_flux = section.getfloat('ice_runoff_flux')
subglacial_runoff_flux = section.getfloat('subglacial_runoff_flux')
iceberg_flux = section.getfloat('iceberg_flux')
wind_stress_zonal = section.getfloat('wind_stress_zonal')
wind_stress_meridional = section.getfloat('wind_stress_meridional')
ds_forcing['temperaturePistonVelocity'] = (
temperature_piston_velocity * forcing_array_surface
)
ds_forcing['salinityPistonVelocity'] = (
salinity_piston_velocity * forcing_array_surface
)
ds_forcing['temperatureSurfaceRestoringValue'] = (
temperature_surface_restoring_value * forcing_array_surface
)
ds_forcing['salinitySurfaceRestoringValue'] = (
salinity_surface_restoring_value * forcing_array_surface
)
ds_forcing['temperatureInteriorRestoringRate'] = (
temperature_interior_restoring_rate * forcing_array
)
ds_forcing['salinityInteriorRestoringRate'] = (
salinity_interior_restoring_rate * forcing_array
)
ds_forcing['temperatureInteriorRestoringValue'] = temperature
ds_forcing['salinityInteriorRestoringValue'] = salinity
ds_forcing['windStressZonal'] = (
wind_stress_zonal * forcing_array_surface
)
ds_forcing['windStressMeridional'] = (
wind_stress_meridional * forcing_array_surface
)
ds_forcing['latentHeatFlux'] = latent_heat_flux * forcing_array_surface
ds_forcing['sensibleHeatFlux'] = (
sensible_heat_flux * forcing_array_surface
)
ds_forcing['shortWaveHeatFlux'] = (
shortwave_heat_flux * forcing_array_surface
)
ds_forcing['evaporationFlux'] = (
evaporation_flux * forcing_array_surface
)
ds_forcing['rainFlux'] = rain_flux * forcing_array_surface
ds_forcing['riverRunoffFlux'] = (
river_runoff_flux * forcing_array_surface
)
ds_forcing['iceRunoffFlux'] = ice_runoff_flux * forcing_array_surface
ds_forcing['subglacialRunoffFlux'] = (
subglacial_runoff_flux * forcing_array_surface
)
ds_forcing['icebergFreshWaterFlux'] = (
iceberg_flux * forcing_array_surface
)
write_netcdf(ds_forcing, 'forcing.nc')