def test_put_prognostic_tendency_in_diagnostics_no_tendencies():
class MockPrognostic(Prognostic):
def __call__(self, state):
return {}, {}
prognostic = TendencyInDiagnosticsWrapper(MockPrognostic(), 'scheme')
tendencies, diagnostics = prognostic({})
assert len(tendencies) == 0
assert len(diagnostics) == 0
def test_get_component_aliases_with_two_component_args():
components = [
MockDiagnostic(),
MockImplicit(),
MockDiagnostic(),
TendencyInDiagnosticsWrapper(DummyPrognostic(), 'dummy')
]
for comp in components[:3]:
aliases = get_component_aliases(comp, components[-1])
assert type(aliases) == dict
assert len(aliases.keys()) == 2
for k in ['T', 'P']:
assert k in list(aliases.values())
def test_put_prognostic_tendency_in_diagnostics_one_tendency():
class MockPrognostic(Prognostic):
tendency_properties = {'quantity': {}}
def __call__(self, state):
return {'quantity': 1.}, {}
prognostic = TendencyInDiagnosticsWrapper(MockPrognostic(), 'scheme')
tendencies, diagnostics = prognostic({})
assert 'tendency_of_quantity_due_to_scheme' in prognostic.diagnostics
tendencies, diagnostics = prognostic({})
assert 'tendency_of_quantity_due_to_scheme' in diagnostics.keys()
assert len(diagnostics) == 1
assert tendencies['quantity'] == 1.
assert diagnostics['tendency_of_quantity_due_to_scheme'] == 1.
def test_get_component_aliases_with_single_component_arg():
components = [
MockPrognostic(),
MockImplicit(),
MockDiagnostic(),
TendencyInDiagnosticsWrapper(DummyPrognostic(), 'dummy')
]
for c, comp in enumerate(components):
aliases = get_component_aliases(comp)
assert type(aliases) == dict
if c == 3:
assert len(aliases.keys()) == 2
for k in ['T', 'P']:
assert k in list(aliases.values())
else:
assert len(aliases.keys()) == 0
def main():
# ============ Adjustable Variables ============
# Integration Options
dt = timedelta(minutes=15) # timestep
duration = '48_00:00' # run duration ('<days>_<hours>:<mins>')t
linearized = True
ncout_freq = 6 # netcdf write frequency (hours)
plot_freq = 6 # plot Monitor call frequency (hours)
ntrunc = 42 # triangular truncation for spharm (e.g., 21 --> T21)
# Diffusion Options
diff_on = True # Use diffusion?
k = 2.338e16 # Diffusion coefficient for del^4 hyperdiffusion
# Forcing Options
forcing_on = True # Apply vort. tendency forcing?
damp_ts = 14.7 # Damping timescale (in days)
# I/O Options
ncoutfile = os.path.join(os.path.dirname(__file__), 'sardeshmukh88.nc')
append_nc = False # Append to an existing netCDF file?
# ==============================================
start = time()
# Get the initial state
state = super_rotation(linearized=linearized, ntrunc=ntrunc)
# Set up the Timestepper with the desired Prognostics
if linearized:
dynamics_prog = LinearizedDynamics(ntrunc=ntrunc)
diffusion_prog = LinearizedDiffusion(k=k, ntrunc=ntrunc)
damping_prog = LinearizedDamping(tau=damp_ts)
else:
dynamics_prog = NonlinearDynamics(ntrunc=ntrunc)
diffusion_prog = NonlinearDiffusion(k=k, ntrunc=ntrunc)
damping_prog = NonlinearDamping(tau=damp_ts)
prognostics = [TendencyInDiagnosticsWrapper(dynamics_prog, 'dynamics')]
if diff_on:
prognostics.append(TendencyInDiagnosticsWrapper(diffusion_prog, 'diffusion'))
if forcing_on:
# Get our suptropical RWS forcing (from equatorial divergence)
rws, rlat, rlon = rws_from_tropical_divergence(state)
prognostics.append(TendencyInDiagnosticsWrapper(Forcing.from_numpy_array(rws, rlat, rlon, ntrunc=ntrunc,
linearized=linearized), 'forcing'))
prognostics.append(TendencyInDiagnosticsWrapper(damping_prog, 'damping'))
stepper = Leapfrog(prognostics)
# Create Monitors for plotting & storing data
plt_monitor = PlotFunctionMonitor(debug_plots.fourpanel)
if os.path.isfile(ncoutfile) and not append_nc:
os.remove(ncoutfile)
aliases = get_component_aliases(*prognostics)
nc_monitor = NetCDFMonitor(ncoutfile, write_on_store=True, aliases=aliases)
# Figure out the end date of this run
d, h, m = re.split('[_:]', duration)
end_date = state['time'] + timedelta(days=int(d), hours=int(h), minutes=int(m))
# Begin the integration loop
idate = state['time']
while state['time'] <= end_date:
# Get the state at the next timestep using our Timestepper
diagnostics, next_state = stepper(state, dt)
# Add any calculated diagnostics to our current state
state.update(diagnostics)
# Write state to netCDF every <ncout_freq> hours
fhour = (state['time'] - idate).days*24 + (state['time'] - idate).seconds/3600
if fhour % ncout_freq == 0:
print(state['time'])
nc_monitor.store(state)
# Make plot(s) every <plot_freq> hours
if fhour % plot_freq == 0:
plt_monitor.store(state)
# Advance the state to the next timestep
next_state['time'] = state['time'] + dt
state = next_state
print('TOTAL INTEGRATION TIME: {:.02f} min\n'.format((time()-start)/60.))
centerlocs=centerlocs,
amplitudes=amplitudes,
widths=widths,
ntrunc=ntrunc,
linearized=linearized)
# Set up the Timestepper with the desired Prognostics
if linearized:
dynamics_prog = LinearizedDynamics(ntrunc=ntrunc)
diffusion_prog = LinearizedDiffusion(k=k, ntrunc=ntrunc)
damping_prog = LinearizedDamping(tau=damp_ts)
else:
dynamics_prog = NonlinearDynamics(ntrunc=ntrunc)
diffusion_prog = NonlinearDiffusion(k=k, ntrunc=ntrunc)
damping_prog = NonlinearDamping(tau=damp_ts)
prognostics = [TendencyInDiagnosticsWrapper(dynamics_prog, 'dynamics')]
# Add diffusion
if diff_on:
prognostics.append(
TendencyInDiagnosticsWrapper(diffusion_prog, 'diffusion'))
# Add our forcing
if forcing_on:
# Get our suptropical RWS forcing (from equatorial divergence)
prognostics.append(TendencyInDiagnosticsWrapper(forcing_prog, 'forcing'))
prognostics.append(TendencyInDiagnosticsWrapper(damping_prog, 'damping'))
# Create Timestepper
stepper = Leapfrog(prognostics)
# Create Monitors for plotting & storing data
append_nc = False # Append to an existing netCDF file?
# ==============================================
start = time()
# Get the initial state
state = sinusoidal_perts_on_zonal_jet(linearized=linearized, ntrunc=ntrunc)
# Set up the Timestepper with the desired Prognostics
if linearized:
dynamics_prog = LinearizedDynamics(ntrunc=ntrunc)
diffusion_prog = LinearizedDiffusion(k=k, ntrunc=ntrunc)
else:
dynamics_prog = NonlinearDynamics(ntrunc=ntrunc)
diffusion_prog = NonlinearDiffusion(k=k, ntrunc=ntrunc)
prognostics = [TendencyInDiagnosticsWrapper(dynamics_prog, 'dynamics')]
# Add diffusion
if diff_on:
prognostics.append(TendencyInDiagnosticsWrapper(diffusion_prog, 'diffusion'))
# Create Timestepper
stepper = Leapfrog(prognostics)
# Create Monitors for plotting & storing data
plt_monitor = PlotFunctionMonitor(debug_plots.fourpanel)
if os.path.isfile(ncoutfile) and not append_nc:
os.remove(ncoutfile)
aliases = get_component_aliases(*prognostics)
nc_monitor = NetCDFMonitor(ncoutfile, write_on_store=True, aliases=aliases)