def test_leapfrog_array_two_steps_filtered_williams(mock_prognostic_call):
"""Test that the Asselin filter is being correctly applied with a
Williams factor of alpha=0.5"""
mock_prognostic_call.return_value = ({
'air_temperature':
np.ones((3, 3)) * 0.
}, {})
state = {'air_temperature': np.ones((3, 3)) * 273.}
timestep = timedelta(seconds=1.)
time_stepper = Leapfrog([MockPrognostic()],
asselin_strength=0.5,
alpha=0.5)
diagnostics, new_state = time_stepper.__call__(state, timestep)
assert list(state.keys()) == ['air_temperature']
assert (state['air_temperature'] == np.ones((3, 3)) * 273.).all()
assert list(new_state.keys()) == ['air_temperature']
assert (new_state['air_temperature'] == np.ones((3, 3)) * 273.).all()
state = new_state
mock_prognostic_call.return_value = ({
'air_temperature':
np.ones((3, 3)) * 2.
}, {})
diagnostics, new_state = time_stepper.__call__(state, timestep)
# Asselin filter modifies the current state
assert list(state.keys()) == ['air_temperature']
assert (state['air_temperature'] == np.ones((3, 3)) * 273.5).all()
assert list(new_state.keys()) == ['air_temperature']
assert (new_state['air_temperature'] == np.ones((3, 3)) * 276.5).all()
def test_leapfrog_array_two_steps_filtered(mock_prognostic_call):
"""Test that the Asselin filter is being correctly applied"""
mock_prognostic_call.return_value = ({
'air_temperature':
np.ones((3, 3)) * 0.
}, {})
state = {'time': timedelta(0), 'air_temperature': np.ones((3, 3)) * 273.}
timestep = timedelta(seconds=1.)
time_stepper = Leapfrog(MockEmptyTendencyComponent(),
asselin_strength=0.5,
alpha=1.)
diagnostics, new_state = time_stepper.__call__(state, timestep)
assert same_list(state.keys(), ['time', 'air_temperature'])
assert (state['air_temperature'] == np.ones((3, 3)) * 273.).all()
assert same_list(new_state.keys(), ['time', 'air_temperature'])
assert (new_state['air_temperature'] == np.ones((3, 3)) * 273.).all()
state = new_state
mock_prognostic_call.return_value = ({
'air_temperature':
np.ones((3, 3)) * 2.
}, {})
diagnostics, new_state = time_stepper.__call__(state, timestep)
# Asselin filter modifies the current state
assert same_list(state.keys(), ['time', 'air_temperature'])
assert (state['air_temperature'] == np.ones((3, 3)) * 274.).all()
assert same_list(new_state.keys(), ['time', 'air_temperature'])
assert (new_state['air_temperature'] == np.ones((3, 3)) * 277.).all()
def test_leapfrog_requires_same_timestep(mock_prognostic_call):
"""Test that the Asselin filter is being correctly applied"""
mock_prognostic_call.return_value = ({'air_temperature': 0.}, {})
state = {'air_temperature': 273.}
time_stepper = Leapfrog([MockPrognostic()], asselin_strength=0.5)
diagnostics, state = time_stepper.__call__(state, timedelta(seconds=1.))
try:
time_stepper.__call__(state, timedelta(seconds=2.))
except ValueError:
pass
except Exception as err:
raise err
else:
raise AssertionError('Leapfrog must require timestep to be constant')
def test_leapfrog_float_two_steps_filtered(mock_prognostic_call):
"""Test that the Asselin filter is being correctly applied"""
mock_prognostic_call.return_value = ({'air_temperature': 0.}, {})
state = {'air_temperature': 273.}
timestep = timedelta(seconds=1.)
time_stepper = Leapfrog([MockPrognostic()], asselin_strength=0.5, alpha=1.)
diagnostics, new_state = time_stepper.__call__(state, timestep)
assert state == {'air_temperature': 273.}
assert new_state == {'air_temperature': 273.}
state = new_state
mock_prognostic_call.return_value = ({'air_temperature': 2.}, {})
diagnostics, new_state = time_stepper.__call__(state, timestep)
# Asselin filter modifies the current state
assert state == {'air_temperature': 274.}
assert new_state == {'air_temperature': 277.}
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.))
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
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: