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

import numpy as np
import xarray as xr
from mpas_tools.io import write_netcdf
from mpas_tools.transects import lon_lat_to_cartesian
from mpas_tools.vector import Vector

from polaris import Step
from polaris.ocean.vertical import init_vertical_coord


[docs]class Init(Step): """ 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='graph.info', work_dir_target=f'{base_mesh.path}/graph.info') self.add_output_file(filename='initial_state.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 = xr.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 = 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) 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. 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) ds['fCell'] = xr.zeros_like(ds_mesh.xCell) ds['fEdge'] = xr.zeros_like(ds_mesh.xEdge) ds['fVertex'] = xr.zeros_like(ds_mesh.xVertex) write_netcdf(ds, 'initial_state.nc')
[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)))