def get_1d_input_state(self, component=None):
if component is None:
component = self.get_component_instance()
return climt.get_default_state([component],
grid_state=get_grid(nx=None,
ny=None,
nz=30))
def test_get_default_grid(self):
grid = get_grid()
self.assert_grid_quantities_present(grid)
self.assert_grid_quantities_have_dimensions(grid, [
'lat', 'lon', 'mid_levels', 'interface_levels',
'ice_interface_levels'
])
def test_enthalpy_and_water_conservation():
conv_adj = climt.DryConvectiveAdjustment()
state = climt.get_default_state([conv_adj],
grid_state=climt.get_grid(nz=35))
dp = (state['air_pressure_on_interface_levels'][:-1] -
state['air_pressure_on_interface_levels'][1:]) / (
state['air_pressure_on_interface_levels'][0] -
state['air_pressure_on_interface_levels'][-1])
state['air_temperature'][:1] += 10
state['specific_humidity'][0] = 0.5
initial_enthalpy = np.sum(
heat_capacity(state['specific_humidity']) * state['air_temperature'] *
dp)
initial_water = np.sum(state['specific_humidity'].values * dp.values)
diag, output = conv_adj(state, timedelta(hours=1))
final_water = np.sum(output['specific_humidity'].values * dp.values)
final_enthalpy = np.sum(
heat_capacity(output['specific_humidity']) *
output['air_temperature'] * dp)
assert np.isclose(initial_water, final_water)
assert np.isclose(initial_enthalpy, final_enthalpy)
def test_dcmip_options():
state = climt.get_default_state([DcmipInitialConditions()],
grid_state=get_grid(nx=64, ny=64, nz=10))
dry_state = DcmipInitialConditions(moist=False)(state)
moist_state = DcmipInitialConditions(moist=True)(state)
not_perturbed_state = DcmipInitialConditions(moist=False,
add_perturbation=False)(state)
tropical_cyclone_state = DcmipInitialConditions(
moist=True, condition_type='tropical_cyclone')(state)
assert not np.all(
np.isclose(dry_state['specific_humidity'].values,
moist_state['specific_humidity'].values))
assert not np.all(
np.isclose(dry_state['eastward_wind'].values,
not_perturbed_state['eastward_wind'].values))
assert not np.all(
np.isclose(
tropical_cyclone_state['surface_air_pressure'].values - 1.015e5,
np.zeros(
not_perturbed_state['surface_air_pressure'].values.shape)))
def get_3d_input_state(self, component=None):
if component is None:
component = self.get_component_instance()
state = climt.get_default_state([component],
grid_state=climt.get_grid(nx=10, ny=5))
state['cloud_area_fraction_in_atmosphere_layer'][16:19] = 0.5
state['mass_content_of_cloud_ice_in_atmosphere_layer'][16:19] = 0.3
return state
def __init__(self,
dt_seconds=1800,
nx=64,
ny=32,
nz=10,
state=None,
input_fields_to_store=input_vars,
output_fields_to_store=output_vars,
input_save_fn=None,
output_save_fn=None,
save_interval=6,
convection=None,
extra_components=[]):
"""
Initialize model. Uses SSTs from Andersen and Kuang 2012.
Creates initial state unless state is given.
"""
climt.set_constants_from_dict(
{'stellar_irradiance': {
'value': 200,
'units': 'W m^-2'
}})
self.model_time_step = timedelta(seconds=dt_seconds)
self.step_counter = 0
self.save_interval = save_interval
# Create components
if convection is None:
convection = climt.EmanuelConvection(
tendencies_in_diagnostics=True)
simple_physics = TimeDifferencingWrapper(
climt.SimplePhysics(tendencies_in_diagnostics=True))
radiation = climt.GrayLongwaveRadiation(tendencies_in_diagnostics=True)
components = [simple_physics, radiation, convection] + extra_components
self.dycore = climt.GFSDynamicalCore(components,
number_of_damped_levels=2)
grid = climt.get_grid(nx=nx, ny=ny, nz=nz)
if state is None:
self.create_initial_state(grid)
else:
self.state = state
if not input_save_fn is None:
self.input_netcdf_monitor = NetCDFMonitor(
input_save_fn,
write_on_store=True,
store_names=input_fields_to_store)
if not output_save_fn is None:
self.output_netcdf_monitor = NetCDFMonitor(
output_save_fn,
write_on_store=True,
store_names=output_fields_to_store)
def test_get_1d_vertical_grid(self):
grid = get_grid(nz=20)
self.assert_grid_quantities_present(grid)
self.assert_grid_quantities_have_dimensions(
grid, ['mid_levels', 'interface_levels', 'ice_interface_levels'])
self.assert_grid_dimensions_have_lengths(grid, {
'mid_levels': 20,
'interface_levels': 21
})
def get_3d_input_state(self):
state = climt.get_default_state([self.get_component_instance()],
grid_state=get_grid(nx=32,
ny=32,
nz=28))
# state = super(TestGFSDycoreWithDcmipInitialConditions, self).get_3d_input_state()
state.update(
climt.DcmipInitialConditions(add_perturbation=True)(state))
return state
def test_3d_initialization_is_full_based_on_wildcard():
grid_state = get_grid(nx=10, ny=10, nz=20)
rrtmg_shortwave = RRTMGShortwave()
instellation = Instellation()
state = get_default_state([rrtmg_shortwave, instellation],
grid_state=grid_state)
for quantity_name, properties in rrtmg_shortwave.input_properties.items():
if '*' in properties['dims']:
assert len(
state[quantity_name].dims) == len(properties['dims']) + 1
if tuple(properties['dims']) == ('*', ):
assert set(state[quantity_name].dims) == {'lat', 'lon'}
elif tuple(properties['dims']) == ('mid_levels', '*'):
assert state[quantity_name].dims[0] == 'mid_levels'
assert set(state[quantity_name].dims[1:]) == {'lat', 'lon'}
def test_get_3d_grid_custom_dim_names(self):
grid = get_grid(nx=3, ny=8, nz=20, x_name='name1', y_name='name2')
self.assert_grid_quantities_present(grid,
latitude=True,
longitude=True)
self.assert_grid_quantities_have_dimensions(grid, [
'mid_levels', 'interface_levels', 'name1', 'name2',
'ice_interface_levels'
])
self.assert_grid_dimensions_have_lengths(grid, {
'mid_levels': 20,
'interface_levels': 21,
'name1': 3,
'name2': 8
})
def test_get_3d_grid(self):
grid = get_grid(nx=4, ny=6, nz=20)
self.assert_grid_quantities_present(grid,
latitude=True,
longitude=True)
self.assert_grid_quantities_have_dimensions(grid, [
'mid_levels', 'interface_levels', 'lat', 'lon',
'ice_interface_levels'
])
self.assert_grid_dimensions_have_lengths(grid, {
'mid_levels': 20,
'interface_levels': 21,
'lat': 6,
'lon': 4
})
def test_get_1d_grid_custom_surface_pressure(self):
grid = get_grid(nz=20, p_surf_in_Pa=0.9e5)
self.assert_grid_quantities_present(grid)
self.assert_grid_quantities_have_dimensions(
grid, ['mid_levels', 'interface_levels', 'ice_interface_levels'])
self.assert_grid_dimensions_have_lengths(grid, {
'mid_levels': 20,
'interface_levels': 21
})
p = grid['air_pressure'].to_units('Pa')
p_interface = grid['air_pressure_on_interface_levels'].to_units('Pa')
assert grid['surface_air_pressure'].to_units('Pa') == 0.9e5
assert np.isclose(p_interface[0], 0.9e5)
assert np.all(p_interface[1:].values < p_interface[:-1].values)
assert np.all(p[1:].values < p[:-1].values)
assert np.all(p[:].values < p_interface[:-1].values)
assert np.all(p[:].values > p_interface[1:].values)
def test_component_3d_grid(self):
grid = get_grid(nx=16, ny=16, nz=32)
component = self.get_component_instance()
state = get_default_state([component], grid_state=grid)
call_component(component, state)
def test_GFSDynamicalCore(self):
grid = get_grid(nx=12, ny=16, nz=28)
component = GFSDynamicalCore()
state = get_default_state([component], grid_state=grid)
call_component(component, state)
def test_component_3d_grid(self):
grid = get_grid(nx=3, ny=4, nz=10)
component = cls()
state = get_default_state([component], grid_state=grid)
assert_state_is_full(state, component)
call_component(component, state)
################################################################################
### IMPORTS
################################################################################
import numpy as np
import matplotlib.pyplot as plt
import climt
from metpy.calc import moist_lapse
from metpy.units import units
from metpy.plots import SkewT
RH = 0.8
# to get same pressure levels as CliMT
N_levels = 30
grid = climt.get_grid(nx=1, ny=1, nz=N_levels)
state = climt.get_default_state([], grid_state=grid)
pressures = state['air_pressure'].values[:, 0, 0]
# setup sample Ts points
# N_sample_pts = 20 # test
N_sample_pts = 200 # full run
minT = 100
maxT = 400
T_surf_sample = np.linspace(minT, maxT, N_sample_pts)
T_data = np.zeros((N_sample_pts, N_levels))
T_data[:, 0] = T_surf_sample
print('Calculating moist adiabats...')
for i in range(N_sample_pts):
import climt
from sympl import PlotFunctionMonitor
from datetime import timedelta
def plot_function(fig, state):
ax = fig.add_subplot(1, 1, 1)
ax.contourf(state['longitude'], state['latitude'], state['zenith_angle'])
ax.set_xlabel('Longitude')
ax.set_ylabel('Latitude')
fig.suptitle('Zenith Angle at time: ' + str(state['time']))
monitor = PlotFunctionMonitor(plot_function)
instellation = climt.Instellation()
state = climt.get_default_state([instellation],
grid_state=climt.get_grid(
nx=100, ny=100, latitude_grid='regular'))
time_step = timedelta(hours=6)
for i in range(8000):
diag = instellation(state)
state.update(diag)
monitor.store(state)
state['time'] += time_step
def get_3d_input_state(self):
state = climt.get_default_state([self.get_component_instance()],
grid_state=get_grid(nx=16,
ny=16,
nz=28))
return state
def plot_function(fig, state):
ax = fig.add_subplot(1, 1, 1)
CS = ax.contourf(state['longitude'], state['latitude'],
state['surface_air_pressure'].to_units('mbar'))
plt.colorbar(CS)
ax.set_title('Surface Pressure at: '+str(state['time']))
monitor = PlotFunctionMonitor(plot_function)
set_constant('reference_air_pressure', value=1e5, units='Pa')
dycore = climt.GFSDynamicalCore()
dcmip = climt.DcmipInitialConditions(add_perturbation=True)
grid = climt.get_grid(nx=128, ny=64, nz=20)
my_state = climt.get_default_state([dycore], grid_state=grid)
timestep = timedelta(minutes=10)
out = dcmip(my_state)
my_state.update(out)
for i in range(1000):
diag, output = dycore(my_state, timestep)
monitor.store(my_state)
my_state.update(output)
my_state['time'] += timestep
ax = fig.add_subplot(1, 2, 1)
state['air_temperature'].mean(dim='lon').plot.contourf(ax=ax, levels=16)
ax.set_title('Temperature')
ax = fig.add_subplot(1, 2, 2)
state['eastward_wind'].mean(dim='lon').plot.contourf(ax=ax, levels=16)
ax.set_title('Zonal Wind')
plt.suptitle('Time: ' + str(state['time']))
model_time_step = timedelta(seconds=600)
monitor = PlotFunctionMonitor(plot_function)
grid = climt.get_grid(nx=128, ny=62)
held_suarez = climt.HeldSuarez()
dycore = climt.GFSDynamicalCore([held_suarez])
my_state = climt.get_default_state([dycore], grid_state=grid)
my_state['eastward_wind'].values[:] = np.random.randn(
*my_state['eastward_wind'].shape)
for i in range(10000):
diag, output = dycore(my_state, model_time_step)
if (my_state['time'].hour % 2 == 0 and my_state['time'].minute == 0):
print('max. zonal wind: ', np.amax(my_state['eastward_wind'].values))
monitor.store(my_state)
my_state.update(output)
my_state['time'] += model_time_step
ax.set_ylim(1e3, 10.)
ax.set_title('Temperature')
ax.grid()
ax.set_xlabel('K')
ax.set_yticklabels([])
plt.suptitle('Radiative Eq. with RRTMG')
monitor = PlotFunctionMonitor(plot_function)
rad_sw = RRTMGShortwave()
rad_lw = RRTMGLongwave()
time_stepper = AdamsBashforth([rad_sw, rad_lw])
timestep = timedelta(hours=3)
grid = get_grid(nx=1, ny=1, nz=60)
state = get_default_state([rad_sw, rad_lw], grid_state=grid)
tp_profiles = np.load('thermodynamic_profiles.npz')
mol_profiles = np.load('molecule_profiles.npz')
state['air_pressure'].values[:] = tp_profiles['air_pressure'][:, np.newaxis,
np.newaxis]
state['air_temperature'].values[:] = tp_profiles['air_temperature'][:,
np.newaxis,
np.newaxis]
state['air_pressure_on_interface_levels'].values[:] = tp_profiles[
'interface_pressures'][:, np.newaxis, np.newaxis]
state['specific_humidity'].values[:] = mol_profiles[
'specific_humidity'][:, np.newaxis, np.newaxis] * 1e-3
constant_duration = 6
radiation_lw = UpdateFrequencyWrapper(
climt.RRTMGLongwave(), constant_duration*model_time_step)
radiation_sw = UpdateFrequencyWrapper(
climt.RRTMGShortwave(), constant_duration*model_time_step)
slab_surface = climt.SlabSurface()
dycore = climt.GFSDynamicalCore(
[simple_physics, slab_surface, radiation_sw,
radiation_lw, convection], number_of_damped_levels=5
)
grid = climt.get_grid(nx=NUM_COLS, ny=NUM_ROWS)
# Create model state
my_state = climt.get_default_state([dycore], grid_state=grid)
# Set initial/boundary conditions
latitudes = my_state['latitude'].values
longitudes = my_state['longitude'].values
zenith_angle = np.radians(latitudes)
surface_shape = [len(longitudes), len(latitudes)]
my_state['zenith_angle'].values = zenith_angle
my_state['eastward_wind'].values[:] = np.random.randn(
*my_state['eastward_wind'].shape)
my_state['ocean_mixed_layer_thickness'].values[:] = 50
ax.plot(state['air_temperature'].values.flatten(),
state['air_pressure'].values.flatten() / 100, '-o')
ax.axes.invert_yaxis()
ax.set_yscale('log')
ax.set_ylim(1e3, 10.)
ax.set_title('Temperature')
ax.grid()
ax.set_xlabel('K')
ax.set_yticklabels([])
plt.suptitle('Radiative Eq. with RRTMG')
monitor = PlotFunctionMonitor(plot_function)
rad_sw = RRTMGShortwave()
rad_lw = RRTMGLongwave()
time_stepper = AdamsBashforth([rad_sw, rad_lw])
timestep = timedelta(hours=3)
grid = get_grid(nx=1, ny=1, nz=30)
state = get_default_state([rad_sw, rad_lw], grid_state=grid)
for i in range(100000):
diagnostics, new_state = time_stepper(state, timestep)
state.update(diagnostics)
if i % 2 == 0:
monitor.store(state)
state = new_state
