Source code for polaris.ocean.tasks.single_column.init

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') 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, _ = xr.broadcast( xr.zeros_like(ds_mesh.xEdge), 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') 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 write_netcdf(ds_forcing, 'forcing.nc')