def set_topography(self, vs):
with h5netcdf.File(DATA_FILES['topography'], 'r') as topography_file:
topo_x, topo_y, topo_z = (np.array(topography_file.variables[k],
dtype='float').T
for k in ('x', 'y', 'z'))
topo_z[topo_z > 0] = 0.
# smooth topography to match grid resolution
gaussian_sigma = (0.5 * len(topo_x) / vs.nx, 0.5 * len(topo_y) / vs.ny)
topo_z_smoothed = scipy.ndimage.gaussian_filter(topo_z,
sigma=gaussian_sigma)
topo_z_smoothed[topo_z >= -1] = 0
topo_x_shifted, topo_z_shifted = self._shift_longitude_array(
vs, topo_x, topo_z_smoothed)
coords = (vs.xt[2:-2], vs.yt[2:-2])
z_interp = allocate(vs, ('xt', 'yt'), local=False)
z_interp[2:-2, 2:-2] = veros.tools.interpolate(
(topo_x_shifted, topo_y),
topo_z_shifted,
coords,
kind='nearest',
fill=False)
depth_levels = 1 + np.argmin(np.abs(z_interp[:, :, np.newaxis] -
vs.zt[np.newaxis, np.newaxis, :]),
axis=2)
vs.kbot[2:-2, 2:-2] = np.where(z_interp < 0., depth_levels, 0)[2:-2,
2:-2]
vs.kbot *= vs.kbot < vs.nz
enforce_boundaries(vs, vs.kbot)
# remove marginal seas
# (dilate to close 1-cell passages, fill holes, undo dilation)
marginal = (scipy.ndimage.binary_erosion(
scipy.ndimage.binary_fill_holes(
scipy.ndimage.binary_dilation(vs.kbot == 0))))
vs.kbot[marginal] = 0
def set_initial_conditions(self, vs):
rpart_shortwave = 0.58
efold1_shortwave = 0.35
efold2_shortwave = 23.0
t_grid = (vs.xt[2:-2], vs.yt[2:-2], vs.zt)
xt_forc, yt_forc, zt_forc = (self._get_data(vs, k)
for k in ('xt', 'yt', 'zt'))
zt_forc = zt_forc[::-1]
# coordinates must be monotonous for this to work
assert np.diff(xt_forc).all() > 0
assert np.diff(yt_forc).all() > 0
# determine slice to read from forcing file
data_subset = (
slice(
max(0,
int(np.argmax(xt_forc >= vs.xt.min())) - 1),
len(xt_forc) -
max(0,
int(np.argmax(xt_forc[::-1] <= vs.xt.max())) - 1)),
slice(
max(0,
int(np.argmax(yt_forc >= vs.yt.min())) - 1),
len(yt_forc) -
max(0,
int(np.argmax(yt_forc[::-1] <= vs.yt.max())) - 1)),
Ellipsis)
xt_forc = xt_forc[data_subset[0]]
yt_forc = yt_forc[data_subset[1]]
# initial conditions
temp_raw = self._get_data(vs, 'temperature',
idx=data_subset)[..., ::-1]
temp_data = veros.tools.interpolate((xt_forc, yt_forc, zt_forc),
temp_raw, t_grid)
vs.temp[2:-2, 2:-2, :, 0] = temp_data * vs.maskT[2:-2, 2:-2, :]
vs.temp[2:-2, 2:-2, :, 1] = temp_data * vs.maskT[2:-2, 2:-2, :]
salt_raw = self._get_data(vs, 'salinity', idx=data_subset)[..., ::-1]
salt_data = veros.tools.interpolate((xt_forc, yt_forc, zt_forc),
salt_raw, t_grid)
vs.salt[2:-2, 2:-2, :, 0] = salt_data * vs.maskT[2:-2, 2:-2, :]
vs.salt[2:-2, 2:-2, :, 1] = salt_data * vs.maskT[2:-2, 2:-2, :]
# wind stress on MIT grid
time_grid = (vs.xt[2:-2], vs.yt[2:-2], np.arange(12))
taux_raw = self._get_data(vs, 'tau_x', idx=data_subset)
taux_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
taux_raw, time_grid)
vs.taux[2:-2, 2:-2, :] = taux_data / vs.rho_0
tauy_raw = self._get_data(vs, 'tau_y', idx=data_subset)
tauy_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
tauy_raw, time_grid)
vs.tauy[2:-2, 2:-2, :] = tauy_data / vs.rho_0
enforce_boundaries(vs, vs.taux)
enforce_boundaries(vs, vs.tauy)
# Qnet and dQ/dT and Qsol
qnet_raw = self._get_data(vs, 'q_net', idx=data_subset)
qnet_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
qnet_raw, time_grid)
vs.qnet[2:-2,
2:-2, :] = -qnet_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
qnec_raw = self._get_data(vs, 'dqdt', idx=data_subset)
qnec_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
qnec_raw, time_grid)
vs.qnec[2:-2,
2:-2, :] = qnec_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
qsol_raw = self._get_data(vs, 'swf', idx=data_subset)
qsol_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
qsol_raw, time_grid)
vs.qsol[2:-2,
2:-2, :] = -qsol_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
# SST and SSS
sst_raw = self._get_data(vs, 'sst', idx=data_subset)
sst_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
sst_raw, time_grid)
vs.t_star[2:-2,
2:-2, :] = sst_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
sss_raw = self._get_data(vs, 'sss', idx=data_subset)
sss_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
sss_raw, time_grid)
vs.s_star[2:-2,
2:-2, :] = sss_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
if vs.enable_idemix:
tidal_energy_raw = self._get_data(vs,
'tidal_energy',
idx=data_subset)
tidal_energy_data = veros.tools.interpolate(
(xt_forc, yt_forc), tidal_energy_raw, t_grid[:-1])
mask_x, mask_y = (i + 2 for i in np.indices((vs.nx, vs.ny)))
mask_z = np.maximum(0, vs.kbot[2:-2, 2:-2] - 1)
tidal_energy_data[:, :] *= vs.maskW[mask_x, mask_y,
mask_z] / vs.rho_0
vs.forc_iw_bottom[2:-2, 2:-2] = tidal_energy_data
"""
Initialize penetration profile for solar radiation and store divergence in divpen
note that pen is set to 0.0 at the surface instead of 1.0 to compensate for the
shortwave part of the total surface flux
"""
swarg1 = vs.zw / efold1_shortwave
swarg2 = vs.zw / efold2_shortwave
pen = rpart_shortwave * np.exp(swarg1) + (
1.0 - rpart_shortwave) * np.exp(swarg2)
pen[-1] = 0.
vs.divpen_shortwave[1:] = (pen[1:] - pen[:-1]) / vs.dzt[1:]
vs.divpen_shortwave[0] = pen[0] / vs.dzt[0]
def npzd(vs):
r"""
Main driving function for NPZD functionality
Computes transport terms and biological activity separately
\begin{equation}
\dfrac{\partial C_i}{\partial t} = T + S
\end{equation}
"""
if not vs.enable_npzd:
return
# TODO: Refactor transportation code to be defined only once and also used by thermodynamics
# TODO: Dissipation on W-grid if necessary
npzd_changes = biogeochemistry(vs)
"""
For vertical mixing
"""
a_tri = allocate(vs, ('xt', 'yt', 'zt'), include_ghosts=False)
b_tri = allocate(vs, ('xt', 'yt', 'zt'), include_ghosts=False)
c_tri = allocate(vs, ('xt', 'yt', 'zt'), include_ghosts=False)
d_tri = allocate(vs, ('xt', 'yt', 'zt'), include_ghosts=False)
delta = allocate(vs, ('xt', 'yt', 'zt'), include_ghosts=False)
ks = vs.kbot[2:-2, 2:-2] - 1
delta[:, :, :-1] = vs.dt_tracer / vs.dzw[np.newaxis, np.newaxis, :-1]\
* vs.kappaH[2:-2, 2:-2, :-1]
delta[:, :, -1] = 0
a_tri[:, :, 1:] = -delta[:, :, :-1] / vs.dzt[np.newaxis, np.newaxis, 1:]
b_tri[:, :,
1:] = 1 + (delta[:, :, 1:] +
delta[:, :, :-1]) / vs.dzt[np.newaxis, np.newaxis, 1:]
b_tri_edge = 1 + delta / vs.dzt[np.newaxis, np.newaxis, :]
c_tri[:, :, :-1] = -delta[:, :, :-1] / vs.dzt[np.newaxis, np.newaxis, :-1]
for tracer in vs.npzd_transported_tracers:
tracer_data = vs.npzd_tracers[tracer]
"""
Advection of tracers
"""
thermodynamics.advect_tracer(
vs, tracer_data[:, :, :, vs.tau],
vs.npzd_advection_derivatives[tracer][:, :, :, vs.tau])
# Adam-Bashforth timestepping
tracer_data[:, :, :, vs.taup1] = tracer_data[:, :, :, vs.tau] + vs.dt_tracer \
* ((1.5 + vs.AB_eps) * vs.npzd_advection_derivatives[tracer][:, :, :, vs.tau]
- (0.5 + vs.AB_eps) * vs.npzd_advection_derivatives[tracer][:, :, :, vs.taum1])\
* vs.maskT
"""
Diffusion of tracers
"""
if vs.enable_hor_diffusion:
horizontal_diffusion_change = np.zeros_like(tracer_data[:, :, :,
0])
diffusion.horizontal_diffusion(vs, tracer_data[:, :, :, vs.tau],
horizontal_diffusion_change)
tracer_data[:, :, :,
vs.taup1] += vs.dt_tracer * horizontal_diffusion_change
if vs.enable_biharmonic_mixing:
biharmonic_diffusion_change = np.empty_like(tracer_data[:, :, :,
0])
diffusion.biharmonic(vs, tracer_data[:, :, :, vs.tau],
np.sqrt(abs(vs.K_hbi)),
biharmonic_diffusion_change)
tracer_data[:, :, :,
vs.taup1] += vs.dt_tracer * biharmonic_diffusion_change
"""
Restoring zones
"""
# TODO add restoring zones to general tracers
"""
Isopycnal diffusion
"""
if vs.enable_neutral_diffusion:
dtracer_iso = np.zeros_like(tracer_data[..., 0])
isoneutral.isoneutral_diffusion_tracer(vs,
tracer_data,
dtracer_iso,
iso=True,
skew=False)
if vs.enable_skew_diffusion:
dtracer_skew = np.zeros_like(tracer_data[..., 0])
isoneutral.isoneutral_diffusion_tracer(vs,
tracer_data,
dtracer_skew,
iso=False,
skew=True)
"""
Vertical mixing of tracers
"""
d_tri[:, :, :] = tracer_data[2:-2, 2:-2, :, vs.taup1]
# TODO: surface flux?
# d_tri[:, :, -1] += surface_forcing
sol, mask = utilities.solve_implicit(vs,
ks,
a_tri,
b_tri,
c_tri,
d_tri,
b_edge=b_tri_edge)
tracer_data[2:-2, 2:-2, :, vs.taup1] = utilities.where(
vs, mask, sol, tracer_data[2:-2, 2:-2, :, vs.taup1])
# update by biogeochemical changes
for tracer, change in npzd_changes.items():
vs.npzd_tracers[tracer][:, :, :, vs.taup1] += change
# prepare next timestep with minimum tracer values
for tracer in vs.npzd_tracers.values():
tracer[:, :, :, vs.taup1] = np.maximum(tracer[:, :, :, vs.taup1],
vs.trcmin * vs.maskT)
for tracer in vs.npzd_tracers.values():
utilities.enforce_boundaries(vs, tracer)
def run(self, show_progress_bar=None):
"""Main routine of the simulation.
Note:
Make sure to call :meth:`setup` prior to this function.
Arguments:
show_progress_bar (:obj:`bool`, optional): Whether to show fancy progress bar via tqdm.
By default, only show if stdout is a terminal and Veros is running on a single process.
"""
vs = self.state
logger.info('\nStarting integration for {0[0]:.1f} {0[1]}'.format(time.format_time(vs.runlen)))
start_time, start_iteration = vs.time, vs.itt
profiler = None
pbar = progress.get_progress_bar(vs, use_tqdm=show_progress_bar)
with handlers.signals_to_exception():
try:
with pbar:
while vs.time - start_time < vs.runlen:
with vs.timers['diagnostics']:
diagnostics.write_restart(vs)
if vs.itt - start_iteration == 3 and rs.profile_mode and rst.proc_rank == 0:
# when using bohrium, most kernels should be pre-compiled by now
profiler = diagnostics.start_profiler()
with vs.timers['main']:
self.set_forcing(vs)
if vs.enable_idemix:
idemix.set_idemix_parameter(vs)
with vs.timers['eke']:
eke.set_eke_diffusivities(vs)
with vs.timers['tke']:
tke.set_tke_diffusivities(vs)
with vs.timers['momentum']:
momentum.momentum(vs)
with vs.timers['temperature']:
thermodynamics.thermodynamics(vs)
if vs.enable_eke or vs.enable_tke or vs.enable_idemix:
advection.calculate_velocity_on_wgrid(vs)
with vs.timers['eke']:
if vs.enable_eke:
eke.integrate_eke(vs)
with vs.timers['idemix']:
if vs.enable_idemix:
idemix.integrate_idemix(vs)
with vs.timers['tke']:
if vs.enable_tke:
tke.integrate_tke(vs)
utilities.enforce_boundaries(vs, vs.u[:, :, :, vs.taup1])
utilities.enforce_boundaries(vs, vs.v[:, :, :, vs.taup1])
if vs.enable_tke:
utilities.enforce_boundaries(vs, vs.tke[:, :, :, vs.taup1])
if vs.enable_eke:
utilities.enforce_boundaries(vs, vs.eke[:, :, :, vs.taup1])
if vs.enable_idemix:
utilities.enforce_boundaries(vs, vs.E_iw[:, :, :, vs.taup1])
momentum.vertical_velocity(vs)
with vs.timers['plugins']:
for plugin in self._plugin_interfaces:
with vs.timers[plugin.name]:
plugin.run_entrypoint(vs)
vs.itt += 1
vs.time += vs.dt_tracer
pbar.advance_time(vs.dt_tracer)
self.after_timestep(vs)
with vs.timers['diagnostics']:
if not diagnostics.sanity_check(vs):
raise RuntimeError('solution diverged at iteration {}'.format(vs.itt))
if vs.enable_neutral_diffusion and vs.enable_skew_diffusion:
isoneutral.isoneutral_diag_streamfunction(vs)
diagnostics.diagnose(vs)
diagnostics.output(vs)
# NOTE: benchmarks parse this, do not change / remove
logger.debug(' Time step took {:.2f}s', vs.timers['main'].get_last_time())
# permutate time indices
vs.taum1, vs.tau, vs.taup1 = vs.tau, vs.taup1, vs.taum1
except:
logger.critical('Stopping integration at iteration {}', vs.itt)
raise
else:
logger.success('Integration done\n')
finally:
diagnostics.write_restart(vs, force=True)
timing_summary = [
'',
'Timing summary:',
' setup time = {:.2f}s'.format(vs.timers['setup'].get_time()),
' main loop time = {:.2f}s'.format(vs.timers['main'].get_time()),
' momentum = {:.2f}s'.format(vs.timers['momentum'].get_time()),
' pressure = {:.2f}s'.format(vs.timers['pressure'].get_time()),
' friction = {:.2f}s'.format(vs.timers['friction'].get_time()),
' thermodynamics = {:.2f}s'.format(vs.timers['temperature'].get_time()),
' lateral mixing = {:.2f}s'.format(vs.timers['isoneutral'].get_time()),
' vertical mixing = {:.2f}s'.format(vs.timers['vmix'].get_time()),
' equation of state = {:.2f}s'.format(vs.timers['eq_of_state'].get_time()),
' EKE = {:.2f}s'.format(vs.timers['eke'].get_time()),
' IDEMIX = {:.2f}s'.format(vs.timers['idemix'].get_time()),
' TKE = {:.2f}s'.format(vs.timers['tke'].get_time()),
' diagnostics and I/O = {:.2f}s'.format(vs.timers['diagnostics'].get_time()),
' plugins = {:.2f}s'.format(vs.timers['plugins'].get_time()),
]
timing_summary.extend([
' {:<22} = {:.2f}s'.format(plugin.name, vs.timers[plugin.name].get_time())
for plugin in vs._plugin_interfaces
])
logger.debug('\n'.join(timing_summary))
if profiler is not None:
diagnostics.stop_profiler(profiler)
def set_initial_conditions(self, vs):
rpart_shortwave = 0.58
efold1_shortwave = 0.35
efold2_shortwave = 23.0
t_grid = (vs.xt[2:-2], vs.yt[2:-2], vs.zt)
xt_forc, yt_forc, zt_forc = (self._get_data(vs, k) for k in ('xt', 'yt', 'zt'))
zt_forc = zt_forc[::-1]
# initial conditions
temp_data = veros.tools.interpolate((xt_forc, yt_forc, zt_forc), self._get_data(vs, 'temperature')[:, :, ::-1],
t_grid, missing_value=0.)
vs.temp[2:-2, 2:-2, :, 0] = temp_data * vs.maskT[2:-2, 2:-2, :]
vs.temp[2:-2, 2:-2, :, 1] = temp_data * vs.maskT[2:-2, 2:-2, :]
salt_data = veros.tools.interpolate((xt_forc, yt_forc, zt_forc), self._get_data(vs, 'salinity')[:, :, ::-1],
t_grid, missing_value=0.)
vs.salt[2:-2, 2:-2, :, 0] = salt_data * vs.maskT[2:-2, 2:-2, :]
vs.salt[2:-2, 2:-2, :, 1] = salt_data * vs.maskT[2:-2, 2:-2, :]
# wind stress on MIT grid
time_grid = (vs.xt[2:-2], vs.yt[2:-2], np.arange(12))
taux_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
self._get_data(vs, 'tau_x'), time_grid,
missing_value=0.)
vs.taux[2:-2, 2:-2, :] = taux_data
mask = np.logical_and(vs.yt > self.so_wind_interval[0], vs.yt < self.so_wind_interval[1])[..., np.newaxis]
vs.taux *= 1. + mask * (self.so_wind_factor - 1.) * np.sin(np.pi * (vs.yt[np.newaxis, :, np.newaxis] - self.so_wind_interval[0]) \
/ (self.so_wind_interval[1] - self.so_wind_interval[0]))
tauy_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
self._get_data(vs, 'tau_y'), time_grid,
missing_value=0.)
vs.tauy[2:-2, 2:-2, :] = tauy_data
enforce_boundaries(vs, vs.taux)
enforce_boundaries(vs, vs.tauy)
# Qnet and dQ/dT and Qsol
qnet_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
self._get_data(vs, 'q_net'), time_grid, missing_value=0.)
vs.qnet[2:-2, 2:-2, :] = -qnet_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
qnec_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
self._get_data(vs, 'dqdt'), time_grid, missing_value=0.)
vs.qnec[2:-2, 2:-2, :] = qnec_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
qsol_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
self._get_data(vs, 'swf'), time_grid, missing_value=0.)
vs.qsol[2:-2, 2:-2, :] = -qsol_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
# SST and SSS
sst_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
self._get_data(vs, 'sst'), time_grid, missing_value=0.)
vs.t_star[2:-2, 2:-2, :] = sst_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
sss_data = veros.tools.interpolate((xt_forc, yt_forc, np.arange(12)),
self._get_data(vs, 'sss'), time_grid, missing_value=0.)
vs.s_star[2:-2, 2:-2, :] = sss_data * vs.maskT[2:-2, 2:-2, -1, np.newaxis]
if vs.enable_idemix:
tidal_energy_data = veros.tools.interpolate(
(xt_forc, yt_forc), self._get_data(vs, 'tidal_energy'), t_grid[:-1], missing_value=0.
)
mask_x, mask_y = (i + 2 for i in np.indices((vs.nx, vs.ny)))
mask_z = np.maximum(0, vs.kbot[2:-2, 2:-2] - 1)
tidal_energy_data[:, :] *= vs.maskW[mask_x, mask_y, mask_z] / vs.rho_0
vs.forc_iw_bottom[2:-2, 2:-2] = tidal_energy_data
# average variables in North Atlantic
na_average_vars = [vs.taux, vs.tauy, vs.qnet, vs.qnec, vs.qsol,
vs.t_star, vs.s_star, vs.salt, vs.temp]
for k in na_average_vars:
k[2:-2, 2:-2, ...] = self._fix_north_atlantic(vs, k[2:-2, 2:-2, ...])
"""
Initialize penetration profile for solar radiation and store divergence in divpen
note that pen is set to 0.0 at the surface instead of 1.0 to compensate for the
shortwave part of the total surface flux
"""
swarg1 = vs.zw / efold1_shortwave
swarg2 = vs.zw / efold2_shortwave
pen = rpart_shortwave * np.exp(swarg1) + (1.0 - rpart_shortwave) * np.exp(swarg2)
pen[-1] = 0.
vs.divpen_shortwave[1:] = (pen[1:] - pen[:-1]) / vs.dzt[1:]
vs.divpen_shortwave[0] = pen[0] / vs.dzt[0]
def carbon_flux(vs):
"""Calculates flux of CO2 over the ocean-atmosphere boundary
This is an adaptation of co2_calc_SWS from UVic ESCM
Note
----
This was written without an atmosphere component in veros.
Therefore an atmospheric pressure of 1 atm is assumed.
The concentration of CO2 (in units of ppmv) may be set in vs.atmospheric_co2
Note
----
This was written without an explicit sea ice component. Therefore a full
ice cover is assumed when temperature is below -1.8C and temperature forcing is negative
Returns
-------
numpy.ndarray(vs.nx, vs.ny) with flux in units of :math:`mmol / m^2 / s`
Positive indicates a flux into the ocean
"""
icemask = np.logical_and(vs.temp[:, :, -1, vs.tau] * vs.maskT[:, :, -1] < -1.8,
vs.forc_temp_surface < 0.0)
ao = np.logical_not(icemask)
atmospheric_pressure = 1 # atm NOTE: We don't have an atmosphere yet, hence constant pressure
# TODO get actual wind speed rather than deriving from wind stress
wind_speed = np.sqrt(np.abs(vs.surface_taux / vs.rho_0) + np.abs(vs.surface_tauy / vs.rho_0)) * 500
vs.wind_speed = wind_speed
# xconv is constant to convert piston_vel from cm/hr -> cm/s
# here it is 100.*a*xconv (100 => m to cm, a=0.337, xconv=1/3.6e+05)
xconv = 0.337 / 3.6e5
xconv *= 0.75 # NOTE: This seems like an approximation I don't know where they got it
vs.dco2star = co2calc_SWS(vs, vs.temp[:, :, -1, vs.tau], # [degree C]
vs.salt[:, :, -1, vs.tau], # [g/kg]
vs.dic[:, :, -1, vs.tau] * 1e-3, # [mmol -> mol]
vs.alkalinity[:, :, -1, vs.tau] * 1e-3, # [mmol -> mol]
vs.atmospheric_co2, # [ppmv]
atmospheric_pressure) # atm
# Schmidt number for CO2
# Wanninkhof, 1992, table A
scco2 = 2073.1 - 125.62 * vs.temp[:, :, -1, vs.tau] + 3.6276 * vs.temp[:, :, -1, vs.tau] ** 2 - 0.043219 * vs.temp[:, :, -1, vs.tau] ** 3
piston_vel = ao * xconv * (wind_speed) ** 2 * ((scco2/660.0)**(-0.5))
# NOTE: According to https://www.soest.hawaii.edu/oceanography/courses/OCN623/Spring%202015/Gas_Exchange_2015_one-lecture1.pdf there are 3 regimes we are looking at the wavy surface
# NOTE: https://aslopubs.onlinelibrary.wiley.com/doi/pdf/10.4319/lom.2014.12.351 uses the correct form for scco2
# Sweeney et al. 2007
# 1e3 added to convert to mmol / m^2 / s
co2_flux = piston_vel * vs.dco2star * vs.maskT[:, :, -1] * 1e3
utilities.enforce_boundaries(vs, co2_flux)
vs.cflux[...] = co2_flux
return vs.cflux
