Unverified Commit d91e8d94 authored by Dion Häfner's avatar Dion Häfner Committed by GitHub

Merge pull request #65 from dionhaefner/overturning

Fix overturning diagnostic in distributed contexts
parents d25b552a 8525d099
......@@ -32,8 +32,9 @@ def test_progress_format(capsys):
assert 'Current iteration:' in captured_tqdm.out
def sanitize(prog):
# remove rates (inconsistent)
prog = re.sub(r'\d+\.\d{2}[smh]/\(model year\)', '?', prog)
# remove rates and ETA (inconsistent)
prog = re.sub(r'\d+\.\d{2}[smh]/\(model year\)', '?s/(model year)', prog)
prog = re.sub(r'\d+\.\d[smh] left', '? left', prog)
prog = prog.replace('\r', '\n')
prog = prog.strip()
return prog
......
......@@ -227,7 +227,7 @@ class FourDegreeTest(VerosPyOMSystemTest):
for s, (v1, v2) in differing_scalars.items():
print('{}, {}, {}'.format(s, v1, v2))
for a, (v1, v2) in differing_arrays.items():
if a in ('B1_gm', 'B2_gm', 'Ai_ez', 'Ai_nz', 'Ai_bx', 'Ai_by'):
if a in ('Ai_ez', 'Ai_nz', 'Ai_bx', 'Ai_by'):
# usually very small differences being amplified
continue
self.check_variable(a, atol=1e-4, data=(v1, v2))
......
......@@ -16,7 +16,8 @@ class ACC2NoEnergyConservationTest(VerosPyOMSystemTest):
for s, (v1, v2) in differing_scalars.items():
print('{}, {}, {}'.format(s, v1, v2))
for a, (v1, v2) in differing_arrays.items():
if 'salt' in a or a in ('B1_gm', 'B2_gm'): # salt and isoneutral streamfunctions aren't used by this example
if 'salt' in a:
# salt isn't used by this example
continue
self.check_variable(a, atol=1e-6, data=(v1, v2))
......
......@@ -180,7 +180,8 @@ class ACC2Test(VerosPyOMSystemTest):
for s, (v1, v2) in differing_scalars.items():
print('{}, {}, {}'.format(s, v1, v2))
for a, (v1, v2) in differing_arrays.items():
if 'salt' in a or a in ('B1_gm', 'B2_gm'): # salt and isoneutral streamfunctions aren't used by this example
if 'salt' in a:
# salt isn't used by this example
continue
self.check_variable(a, atol=1e-6, data=(v1, v2))
......
......@@ -153,24 +153,21 @@ def isoneutral_diag_streamfunction(vs):
calculate hor. components of streamfunction for eddy driven velocity
for diagnostics purpose only
"""
K_gm_pad = utilities.pad_z_edges(vs, vs.K_gm)
"""
meridional component at east face of 'T' cells
"""
K_gm_pad = utilities.pad_z_edges(vs, vs.K_gm)
diffloc = 0.25 * (K_gm_pad[1:-2, 2:-2, 1:-1] + K_gm_pad[1:-2, 2:-2, :-2] +
K_gm_pad[2:-1, 2:-2, 1:-1] + K_gm_pad[2:-1, 2:-2, :-2])
sumz = np.sum(diffloc[..., np.newaxis, np.newaxis] * vs.Ai_ez[1:-2, 2:-2, ...], axis=(3, 4))
vs.B2_gm[1:-2, 2:-2, :] = 0.25 * sumz
vs.B2_gm[1:-2, 2:-2, :] = 0.25 * diffloc * np.sum(vs.Ai_ez[1:-2, 2:-2, ...], axis=(3, 4))
"""
zonal component at north face of 'T' cells
"""
diffloc = 0.25 * (K_gm_pad[2:-2, 1:-2, 1:-1] + K_gm_pad[2:-2, 1:-2, :-2] +
K_gm_pad[2:-2, 2:-1, 1:-1] + K_gm_pad[2:-2, 2:-1, :-2])
sumz = np.sum(diffloc[..., np.newaxis, np.newaxis] * vs.Ai_nz[2:-2, 1:-2, ...], axis=(3, 4))
vs.B1_gm[2:-2, 1:-2, :] = -0.25 * sumz
vs.B1_gm[2:-2, 1:-2, :] = -0.25 * diffloc * np.sum(vs.Ai_nz[2:-2, 1:-2, ...], axis=(3, 4))
@veros_method(dist_safe=False, local_variables=[
......
......@@ -104,6 +104,8 @@ def calc_beta(vs):
vs.beta[:, 2:-2] = 0.5 * ((vs.coriolis_t[:, 3:-1] - vs.coriolis_t[:, 2:-2]) / vs.dyu[2:-2]
+ (vs.coriolis_t[:, 2:-2] - vs.coriolis_t[:, 1:-3]) / vs.dyu[1:-3])
utilities.enforce_boundaries(vs, vs.beta)
@veros_method
def calc_topo(vs):
......@@ -117,9 +119,10 @@ def calc_topo(vs):
vs.kbot[:, :2] = 0
vs.kbot[:, -2:] = 0
if vs.enable_cyclic_x:
utilities.enforce_boundaries(vs, vs.kbot)
else:
utilities.enforce_boundaries(vs, vs.kbot)
if not vs.enable_cyclic_x:
vs.kbot[:2, :] = 0
vs.kbot[-2:, :] = 0
......
......@@ -41,6 +41,11 @@ ISONEUTRAL_VARIABLES = OrderedDict([
])
@veros_method(inline=True)
def zonal_sum(vs, arr):
return global_sum(vs, np.sum(arr, axis=0), axis=0)
class Overturning(VerosDiagnostic):
"""Isopycnal overturning diagnostic. Computes and writes vertical streamfunctions
(zonally averaged).
......@@ -81,10 +86,10 @@ class Overturning(VerosDiagnostic):
self.sigma[...] = self.sigs + self.dsig * np.arange(self.nlevel)
# precalculate area below z levels
self.zarea[2:-2, :] = np.cumsum(global_sum(vs, np.sum(
self.zarea[2:-2, :] = np.cumsum(zonal_sum(vs,
vs.dxt[2:-2, np.newaxis, np.newaxis]
* vs.cosu[np.newaxis, 2:-2, np.newaxis]
* vs.maskV[2:-2, 2:-2, :], axis=0)) * vs.dzt[np.newaxis, :], axis=1)
* vs.maskV[2:-2, 2:-2, :]) * vs.dzt[np.newaxis, :], axis=1)
self.initialize_output(vs, self.variables,
var_data={'sigma': self.sigma},
......@@ -115,21 +120,18 @@ class Overturning(VerosDiagnostic):
trans = allocate(vs, ('yu', self.nlevel))
z_sig = allocate(vs, ('yu', self.nlevel))
fac = (vs.dxt[2:-2, np.newaxis, np.newaxis]
* vs.cosu[np.newaxis, 2:-2, np.newaxis]
* vs.dzt[np.newaxis, np.newaxis, :]
* vs.maskV[2:-2, 2:-2, :])
for m in range(self.nlevel):
# NOTE: vectorized version would be O(N^4) in memory
# consider cythonizing if performance-critical
mask = sig_loc_face > self.sigma[m]
trans[2:-2, m] = global_sum(vs, np.sum(
vs.v[2:-2, 2:-2, :, vs.tau]
* vs.dxt[2:-2, np.newaxis, np.newaxis]
* vs.cosu[np.newaxis, 2:-2, np.newaxis]
* vs.dzt[np.newaxis, np.newaxis, :]
* vs.maskV[2:-2, 2:-2, :] * mask, axis=(0, 2)))
z_sig[2:-2, m] = global_sum(vs, np.sum(
vs.dzt[np.newaxis, np.newaxis, :]
* vs.dxt[2:-2, np.newaxis, np.newaxis]
* vs.cosu[np.newaxis, 2:-2, np.newaxis]
* vs.maskV[2:-2, 2:-2, :] * mask, axis=(0, 2)))
trans[2:-2, m] = zonal_sum(vs, np.sum(vs.v[2:-2, 2:-2, :, vs.tau] * fac * mask, axis=2))
z_sig[2:-2, m] = zonal_sum(vs, np.sum(fac * mask, axis=2))
self.trans += trans
if vs.enable_neutral_diffusion and vs.enable_skew_diffusion:
......@@ -138,38 +140,37 @@ class Overturning(VerosDiagnostic):
for m in range(self.nlevel):
# NOTE: see above
mask = sig_loc_face > self.sigma[m]
bolus_trans[2:-2, m] = global_sum(vs,
bolus_trans[2:-2, m] = zonal_sum(vs,
np.sum(
(vs.B1_gm[2:-2, 2:-2, 1:] - vs.B1_gm[2:-2, 2:-2, :-1])
* vs.dxt[2:-2, np.newaxis, np.newaxis]
* vs.cosu[np.newaxis, 2:-2, np.newaxis]
* vs.maskV[2:-2, 2:-2, 1:]
* mask[:, :, 1:],
axis=(0, 2)
)
+ np.sum(
vs.B1_gm[2:-2, 2:-2, 0]
* vs.dxt[2:-2, np.newaxis]
* vs.cosu[np.newaxis, 2:-2]
* vs.maskV[2:-2, 2:-2, 0]
* mask[:, :, 0],
axis=0
axis=2
)
+
vs.B1_gm[2:-2, 2:-2, 0]
* vs.dxt[2:-2, np.newaxis]
* vs.cosu[np.newaxis, 2:-2]
* vs.maskV[2:-2, 2:-2, 0]
* mask[:, :, 0]
)
# streamfunction on geopotentials
self.vsf_depth[2:-2, :] += np.cumsum(global_sum(vs, np.sum(
self.vsf_depth[2:-2, :] += np.cumsum(zonal_sum(vs,
vs.dxt[2:-2, np.newaxis, np.newaxis]
* vs.cosu[np.newaxis, 2:-2, np.newaxis]
* vs.v[2:-2, 2:-2, :, vs.tau]
* vs.maskV[2:-2, 2:-2, :], axis=0)) * vs.dzt[np.newaxis, :], axis=1)
* vs.maskV[2:-2, 2:-2, :]) * vs.dzt[np.newaxis, :], axis=1)
if vs.enable_neutral_diffusion and vs.enable_skew_diffusion:
# streamfunction for eddy driven velocity on geopotentials
self.bolus_depth[2:-2, :] += global_sum(vs, np.sum(
self.bolus_depth[2:-2, :] += zonal_sum(vs,
vs.dxt[2:-2, np.newaxis, np.newaxis]
* vs.cosu[np.newaxis, 2:-2, np.newaxis]
* vs.B1_gm[2:-2, 2:-2, :], axis=0))
* vs.B1_gm[2:-2, 2:-2, :])
# interpolate from isopycnals to depth
self.vsf_iso[2:-2, :] += self._interpolate_along_axis(vs,
z_sig[2:-2, :], trans[2:-2, :],
......@@ -181,7 +182,6 @@ class Overturning(VerosDiagnostic):
self.nitts += 1
@veros_method
def _interpolate_along_axis(self, vs, coords, arr, interp_coords, axis=0):
# TODO: clean up this mess
......
......@@ -281,7 +281,15 @@ def exchange_cyclic_boundaries(vs, arr):
@dist_context_only
@veros_method(inline=True)
def _reduce(vs, arr, op):
def _reduce(vs, arr, op, axis=None):
if axis is None:
comm = rs.mpi_comm
else:
assert axis in (0, 1)
pi = proc_rank_to_index(rst.proc_rank)
other_axis = 1 - axis
comm = rs.mpi_comm.Split(pi[other_axis], rst.proc_rank)
if np.isscalar(arr):
squeeze = True
arr = np.array([arr])
......@@ -291,7 +299,7 @@ def _reduce(vs, arr, op):
arr = ascontiguousarray(arr)
res = np.empty_like(arr)
rs.mpi_comm.Allreduce(
comm.Allreduce(
get_array_buffer(vs, arr),
get_array_buffer(vs, res),
op=op
......@@ -305,37 +313,37 @@ def _reduce(vs, arr, op):
@dist_context_only
@veros_method
def global_and(vs, arr):
def global_and(vs, arr, axis=None):
from mpi4py import MPI
return _reduce(vs, arr, MPI.LAND)
return _reduce(vs, arr, MPI.LAND, axis=axis)
@dist_context_only
@veros_method
def global_or(vs, arr):
def global_or(vs, arr, axis=None):
from mpi4py import MPI
return _reduce(vs, arr, MPI.LOR)
return _reduce(vs, arr, MPI.LOR, axis=axis)
@dist_context_only
@veros_method
def global_max(vs, arr):
def global_max(vs, arr, axis=None):
from mpi4py import MPI
return _reduce(vs, arr, MPI.MAX)
return _reduce(vs, arr, MPI.MAX, axis=axis)
@dist_context_only
@veros_method
def global_min(vs, arr):
def global_min(vs, arr, axis=None):
from mpi4py import MPI
return _reduce(vs, arr, MPI.MIN)
return _reduce(vs, arr, MPI.MIN, axis=axis)
@dist_context_only
@veros_method
def global_sum(vs, arr):
def global_sum(vs, arr, axis=None):
from mpi4py import MPI
return _reduce(vs, arr, MPI.SUM)
return _reduce(vs, arr, MPI.SUM, axis=axis)
@dist_context_only
......
......@@ -96,6 +96,10 @@ class FancyProgressBar:
We only need TQDM to handle dynamic updates to the progress indicator.
"""
def __init__(self, *args, **kwargs):
kwargs.update(leave=True)
super().__init__(*args, **kwargs)
@property
def format_dict(other):
report_time = time.convert_time(self._time, 'seconds', self._time_unit)
......@@ -152,7 +156,6 @@ class FancyProgressBar:
def __exit__(self, *args, **kwargs):
logs.setup_logging(loglevel=rs.loglevel)
self._pbar.__exit__(*args, **kwargs)
def advance_time(self, amount):
self._iteration += 1
......@@ -160,7 +163,7 @@ class FancyProgressBar:
self.flush()
def flush(self):
self._pbar.update(0)
self._pbar.refresh()
def get_progress_bar(vs, use_tqdm=None):
......
......@@ -230,6 +230,9 @@ class VerosLegacy(veros.VerosSetup):
# diagnose vertical velocity at taup1
f.vertical_velocity()
# diagnose isoneutral streamfunction regardless of output settings
f.isoneutral_diag_streamfunction()
# shift time
m.itt += 1
vs.time += m.dt_tracer
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment