import numpy as np
import xarray as xr
from mpas_tools.io import open_dataset
from mpas_tools.transects import lon_lat_to_cartesian
from mpas_tools.vector import Vector
from polaris.ocean.coriolis import add_coriolis_to_dataset
from polaris.ocean.model import OceanIOStep
from polaris.ocean.vertical import init_vertical_coord
[docs]
class Init(OceanIOStep):
"""
A step for an initial condition for for the cosine bell test case
"""
[docs]
def __init__(self, component, name, subdir, base_mesh):
"""
Create the step
Parameters
----------
component : polaris.Component
The component the step belongs to
name : str
The name of the step
subdir : str
The subdirectory for the step
base_mesh : polaris.Step
The base mesh step
"""
super().__init__(component=component, name=name, subdir=subdir)
self.add_input_file(
filename='mesh.nc',
work_dir_target=f'{base_mesh.path}/base_mesh.nc',
)
self.add_input_file(
filename='reconstruction_weights.nc',
work_dir_target=f'{base_mesh.path}/reconstruction_weights.nc',
)
def setup(self):
super().setup()
self.add_output_files_for_ocean_model_input(
horiz_mesh_filename='culled_mesh.nc',
)
[docs]
def run(self):
"""
Run this step of the task
"""
config = self.config
section = config['cosine_bell']
temperature = section.getfloat('temperature')
salinity = section.getfloat('salinity')
lat_center = section.getfloat('lat_center')
lon_center = section.getfloat('lon_center')
radius = section.getfloat('radius')
psi0 = section.getfloat('psi0')
vel_pd = section.getfloat('vel_pd')
section = config['vertical_grid']
bottom_depth = section.getfloat('bottom_depth')
ds_mesh = open_dataset('mesh.nc')
angleEdge = ds_mesh.angleEdge
latCell = ds_mesh.latCell
latEdge = ds_mesh.latEdge
lonCell = ds_mesh.lonCell
sphere_radius = ds_mesh.sphere_radius
ds_mesh = add_coriolis_to_dataset(config, ds_mesh)
self.write_horiz_mesh_dataset(ds_mesh, 'culled_mesh.nc', config)
ds = ds_mesh.copy()
ds['bottomDepth'] = bottom_depth * xr.ones_like(latCell)
ds['ssh'] = xr.zeros_like(latCell)
init_vertical_coord(config, ds)
temperature_array = temperature * xr.ones_like(ds_mesh.latCell)
temperature_array, _ = xr.broadcast(temperature_array, ds.refZMid)
ds['temperature'] = temperature_array.expand_dims(dim='Time', axis=0)
ds['salinity'] = salinity * xr.ones_like(ds.temperature)
# temperature and salinity must be set before this call:
# write_vert_coord_dataset converts restingThickness to
# RefPseudoThickness via pseudothickness_from_ds, which requires T/S
self.write_vert_coord_dataset(ds, 'vert_coord.nc', config)
x_center, y_center, z_center = lon_lat_to_cartesian(
lon_center, lat_center, sphere_radius, degrees=False
)
x_cells, y_cells, z_cells = lon_lat_to_cartesian(
lonCell, latCell, sphere_radius, degrees=False
)
xyz_center = Vector(x_center, y_center, z_center)
xyz_cells = Vector(x_cells, y_cells, z_cells)
ang_dist_from_center = xyz_cells.angular_distance(xyz_center)
distance_from_center = ang_dist_from_center * sphere_radius
bell_value = cosine_bell(psi0, distance_from_center, radius)
debug_tracers = xr.where(
distance_from_center < radius, bell_value, 0.0
)
debug_tracers_array, _ = xr.broadcast(debug_tracers, ds.refZMid)
ds['tracer1'] = debug_tracers_array.expand_dims(dim='Time', axis=0)
ds['tracer2'] = ds.tracer1
ds['tracer3'] = ds.tracer1
# Initialize velocity
s_per_hour = 3600.0
velocity = (
2.0
* np.pi
* np.cos(angleEdge)
* sphere_radius
* np.cos(latEdge)
/ (s_per_hour * vel_pd)
)
velocity_array, _ = xr.broadcast(velocity, ds.refZMid)
ds['normalVelocity'] = velocity_array.expand_dims(dim='Time', axis=0)
self.write_initial_state_dataset(ds, 'init.nc', config)
[docs]
def cosine_bell(max_value, ri, r):
"""
Compute values according to cosine bell function
Parameters
----------
max_value : float
Maximum value of the cosine bell function
ri : np.ndarray of type float
Distance from the center of the cosine bell in meters
r : float
Radius of the cosine bell in meters
Returns
-------
f : np.ndarray of type float
Cosine bell tracer values
"""
return max_value / 2.0 * (1.0 + np.cos(np.pi * np.divide(ri, r)))