def test_tendencies_in_diagnostics_one_tendency_with_component_name(self):
input_properties = {}
diagnostic_properties = {}
tendency_properties = {'output1': {'dims': ['dim1'], 'units': 'm/s'}}
diagnostic_output = {}
tendency_output = {
'output1': np.ones([10]) * 2.,
}
prognostic = self.prognostic_class(input_properties,
diagnostic_properties,
tendency_properties,
diagnostic_output, tendency_output)
stepper = self.timestepper_class(prognostic,
tendencies_in_diagnostics=True,
name='component')
state = {
'time':
timedelta(0),
'output1':
DataArray(np.ones([10]) * 10., dims=['dim1'], attrs={'units':
'm'}),
}
diagnostics, _ = stepper(state, timedelta(seconds=5))
assert 'output1_tendency_from_component' in diagnostics.keys()
assert len(diagnostics['output1_tendency_from_component'].dims) == 1
assert 'dim1' in diagnostics['output1_tendency_from_component'].dims
assert units_are_compatible(
diagnostics['output1_tendency_from_component'].attrs['units'],
'm s^-1')
assert np.allclose(
diagnostics['output1_tendency_from_component'].values, 2.)
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 = {'time': timedelta(0), 'air_temperature': np.ones((3, 3)) * 273.}
timestep = timedelta(seconds=1.)
time_stepper = Leapfrog(MockEmptyTendencyComponent(),
asselin_strength=0.5,
alpha=0.5)
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)) * 273.5).all()
assert same_list(new_state.keys(), ['time', 'air_temperature'])
assert (new_state['air_temperature'] == np.ones((3, 3)) * 276.5).all()
def test_array_four_steps(self, mock_prognostic_call):
mock_prognostic_call.return_value = ({
'air_temperature':
np.ones((3, 3)) * 1.
}, {})
state = {
'time': timedelta(0),
'air_temperature': np.ones((3, 3)) * 273.
}
timestep = timedelta(seconds=1.)
time_stepper = self.timestepper_class(MockEmptyTendencyComponent())
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)) * 274.).all()
state = new_state
diagnostics, new_state = time_stepper.__call__(new_state, timestep)
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)) * 275.).all()
state = new_state
diagnostics, new_state = time_stepper.__call__(state, timestep)
assert same_list(state.keys(), ['time', 'air_temperature'])
assert (state['air_temperature'] == np.ones((3, 3)) * 275.).all()
assert same_list(new_state.keys(), ['time', 'air_temperature'])
assert (new_state['air_temperature'] == np.ones((3, 3)) * 276.).all()
state = new_state
diagnostics, new_state = time_stepper.__call__(state, timestep)
assert same_list(state.keys(), ['time', 'air_temperature'])
assert (state['air_temperature'] == np.ones((3, 3)) * 276.).all()
assert same_list(new_state.keys(), ['time', 'air_temperature'])
assert (new_state['air_temperature'] == np.ones((3, 3)) * 277.).all()
def test_unused_quantities_carried_over(self):
state = {'time': timedelta(0), 'air_temperature': 273.}
time_stepper = self.timestepper_class(MockEmptyTendencyComponent())
timestep = timedelta(seconds=1.)
diagnostics, new_state = time_stepper.__call__(state, timestep)
assert state == {'time': timedelta(0), 'air_temperature': 273.}
assert new_state == {'time': timedelta(0), 'air_temperature': 273.}
def test_constant_prognostic_cannot_modify_through_output_dict():
prog = ConstantTendencyComponent({}, {})
tendencies, diagnostics = prog({'time': timedelta(0)})
tendencies['a'] = 'b'
diagnostics['c'] = 'd'
tendencies, diagnostics = prog({'time': timedelta(0)})
assert len(tendencies) == 0
assert len(diagnostics) == 0
def test_float_one_step(self, mock_prognostic_call):
mock_prognostic_call.return_value = ({'air_temperature': 1.}, {})
state = {'time': timedelta(0), 'air_temperature': 273.}
timestep = timedelta(seconds=1.)
time_stepper = self.timestepper_class(MockEmptyTendencyComponent())
diagnostics, new_state = time_stepper.__call__(state, timestep)
assert state == {'time': timedelta(0), 'air_temperature': 273.}
assert new_state == {'time': timedelta(0), 'air_temperature': 274.}
def test_tendencies_in_diagnostics_no_tendency(self):
input_properties = {}
diagnostic_properties = {}
tendency_properties = {}
diagnostic_output = {}
tendency_output = {}
prognostic = self.prognostic_class(input_properties,
diagnostic_properties,
tendency_properties,
diagnostic_output, tendency_output)
stepper = self.timestepper_class(prognostic,
tendencies_in_diagnostics=True)
state = {'time': timedelta(0)}
diagnostics, _ = stepper(state, timedelta(seconds=5))
assert diagnostics == {}
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 = {'time': timedelta(0), 'air_temperature': 273.}
time_stepper = Leapfrog([MockEmptyTendencyComponent()],
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_adams_bashforth_requires_same_timestep(mock_prognostic_call):
"""Test that the Asselin filter is being correctly applied"""
mock_prognostic_call.return_value = ({'air_temperature': 0.}, {})
state = {'time': timedelta(0), 'air_temperature': 273.}
time_stepper = AdamsBashforth(MockEmptyTendencyComponent())
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(
'AdamsBashforth must require timestep to be constant')
def test_constant_prognostic_empty_dicts():
prog = ConstantTendencyComponent({}, {})
tendencies, diagnostics = prog({'time': timedelta(0)})
assert isinstance(tendencies, dict)
assert isinstance(diagnostics, dict)
assert len(tendencies) == 0
assert len(diagnostics) == 0
def test_incrementing_years_using_days(self):
dt = datetime(1900, 1, 1, calendar=self.calendar)
dt_1950 = dt + timedelta(days=365 * 50)
assert dt_1950.year == 1950
assert dt_1950.month == 1
assert dt_1950.day == 1
assert dt_1950.hour == 0
def test_constant_diagnostic_cannot_modify_through_input_dict():
in_diagnostics = {}
diag = ConstantDiagnosticComponent(in_diagnostics)
in_diagnostics['a'] = 'b'
diagnostics = diag({'time': timedelta(0)})
assert isinstance(diagnostics, dict)
assert len(diagnostics) == 0
def get_era5_forcing(latent_filename, i_timestep, latent=True):
state = {}
ds = xr.open_dataset(latent_filename)
state['surface_latent_heat_flux'] = ds['lhf'][i_timestep] / 3600. # divide by one hour to go from J/m^2 to W/m^2
state['surface_latent_heat_flux'].attrs['units'] = 'W/m^2'
state['surface_sensible_heat_flux'] = ds['shf'][i_timestep] / 3600.
state['surface_sensible_heat_flux'].attrs['units'] = 'W/m^2'
state['surface_temperature'] = ds['sst'][i_timestep]
state['surface_temperature'].attrs['units'] = 'degK'
state['surface_air_pressure'] = ds['p_surface'][i_timestep]
state['surface_air_pressure'].attrs['units'] = 'Pa'
state['vertical_wind'] = ds['w'][i_timestep, :]
state['vertical_wind'].attrs['units'] = 'm/s'
state['liquid_water_static_energy_horizontal_advective_tendency'] = ds['sl_adv'][i_timestep, :]
state['total_water_mixing_ratio_horizontal_advective_tendency'] = ds['rt_adv'][i_timestep, :]
state['downwelling_shortwave_radiation_at_3km'] = ds['swdn_tod'][i_timestep]
state['downwelling_shortwave_radiation_at_3km'].attrs['units'] = 'W/m^2'
state['downwelling_shortwave_radiation_at_top_of_atmosphere'] = ds['swdn_toa'][i_timestep]
state['downwelling_shortwave_radiation_at_top_of_atmosphere'].attrs['units'] = 'W/m^2'
state['mid_cloud_fraction'] = ds['cldmid'][i_timestep]
state['mid_cloud_fraction'].attrs['units'] = ''
state['high_cloud_fraction'] = ds['cldhigh'][i_timestep]
state['high_cloud_fraction'].attrs['units'] = ''
state['total_water_mixing_ratio_at_3km'] = ds['rt'][i_timestep, -1]
state['total_water_mixing_ratio_at_3km'].attrs['units'] = 'kg/kg'
state['liquid_water_static_energy_at_3km'] = ds['sl'][i_timestep, -1]
state['liquid_water_static_energy_at_3km'].attrs['units'] = 'J/kg'
state['rain_water_mixing_ratio_at_3km'] = ds['rrain'][i_timestep, -1]
state['rain_water_mixing_ratio_at_3km'].attrs['units'] = 'kg/kg'
if latent:
state['total_water_mixing_ratio'] = ds['rt'][i_timestep, :]
state['total_water_mixing_ratio'].attrs['units'] = 'kg/kg'
state['liquid_water_static_energy'] = ds['sl'][i_timestep, :]
state['liquid_water_static_energy'].attrs['units'] = 'J/kg'
convert_dataarray_to_sympl(state)
if latent:
state['liquid_water_static_energy_components_horizontal_advective_tendency'] = \
sympl.DataArray(
convert_height_to_principal_components(
state['liquid_water_static_energy_horizontal_advective_tendency'],
basis_name='sl', subtract_mean=False
), dims=['sl_latent'], attrs={'units': 's^-1'}
)
state['total_water_mixing_ratio_components_horizontal_advective_tendency'] = \
sympl.DataArray(
convert_height_to_principal_components(
state['total_water_mixing_ratio_horizontal_advective_tendency'],
basis_name='rt', subtract_mean=False
), dims=['rt_latent'], attrs={'units': 's^-1'}
)
pc_state = {}
pc_state.update(state)
pc_state['time'] = sympl.timedelta(0)
pc_state = height_to_pc(pc_state)
pc_state.pop('total_water_mixing_ratio_components')
pc_state.pop('liquid_water_static_energy_components')
state.update(pc_state)
return state
def test_constant_prognostic_cannot_modify_through_input_dict():
in_tendencies = {}
in_diagnostics = {}
prog = ConstantTendencyComponent(in_tendencies, in_diagnostics)
in_tendencies['a'] = 'b'
in_diagnostics['c'] = 'd'
tendencies, diagnostics = prog({'time': timedelta(0)})
assert len(tendencies) == 0
assert len(diagnostics) == 0
def test_float_one_step_with_units(self, mock_prognostic_call):
mock_prognostic_call.return_value = ({
'eastward_wind':
DataArray(0.02, attrs={'units': 'km/s^2'})
}, {})
state = {
'time': timedelta(0),
'eastward_wind': DataArray(1., attrs={'units': 'm/s'})
}
timestep = timedelta(seconds=1.)
time_stepper = self.timestepper_class(MockEmptyTendencyComponent())
diagnostics, new_state = time_stepper.__call__(state, timestep)
assert state == {
'time': timedelta(0),
'eastward_wind': DataArray(1., attrs={'units': 'm/s'})
}
assert same_list(new_state.keys(), ['time', 'eastward_wind'])
assert np.allclose(new_state['eastward_wind'].values, 21.)
assert new_state['eastward_wind'].attrs['units'] == 'm/s'
def get_era5_state(latent_filename, latent=True, i_timestep=0):
state = {}
ds = xr.open_dataset(latent_filename)
state['total_water_mixing_ratio'] = ds['rt'][i_timestep, :]
state['total_water_mixing_ratio'].attrs['units'] = 'kg/kg'
state['liquid_water_static_energy'] = ds['sl'][i_timestep, :]
state['liquid_water_static_energy'].attrs['units'] = 'J/kg'
state['height'] = sympl.DataArray(
np.linspace(0, 3000., 20),
dims=['z_star'],
attrs={'units': 'm'},
)
state['time'] = sympl.timedelta(0)
if latent:
state['vertical_wind'] = ds['w'][i_timestep, :]
state['vertical_wind'].attrs['units'] = 'm/s'
convert_dataarray_to_sympl(state)
if latent:
state = height_to_pc(state)
state.pop('vertical_wind_components')
state['time'] = sympl.timedelta(0)
return state
def test_dataarray_one_step(self, mock_prognostic_call):
mock_prognostic_call.return_value = ({
'air_temperature':
DataArray(np.ones((3, 3)), attrs={'units': 'K/s'})
}, {})
state = {
'time':
timedelta(0),
'air_temperature':
DataArray(np.ones((3, 3)) * 273., attrs={'units': 'K'})
}
timestep = timedelta(seconds=1.)
time_stepper = self.timestepper_class(MockEmptyTendencyComponent())
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 len(state['air_temperature'].attrs) == 1
assert state['air_temperature'].attrs['units'] == 'K'
assert same_list(new_state.keys(), ['time', 'air_temperature'])
assert (new_state['air_temperature'] == np.ones((3, 3)) * 274.).all()
assert len(new_state['air_temperature'].attrs) == 1
assert new_state['air_temperature'].attrs['units'] == 'K'
def test_given_tendency_not_modified_with_two_components(self):
input_properties = {}
diagnostic_properties = {}
tendency_properties = {
'tend1': {
'dims': ['dim1'],
'units': 'm/s',
}
}
diagnostic_output = {}
tendency_output_1 = {'tend1': np.ones([10]) * 5.}
prognostic1 = self.prognostic_class(input_properties,
diagnostic_properties,
tendency_properties,
diagnostic_output,
tendency_output_1)
tendency_output_2 = {'tend1': np.ones([10]) * 5.}
prognostic2 = self.prognostic_class(input_properties,
diagnostic_properties,
tendency_properties,
diagnostic_output,
tendency_output_2)
stepper = self.timestepper_class(prognostic1,
prognostic2,
tendencies_in_diagnostics=True)
state = {
'time':
timedelta(0),
'tend1':
DataArray(
np.ones([10]),
dims=['dim1'],
attrs={'units': 'm'},
)
}
_, _ = stepper(state, timedelta(seconds=5))
assert np.all(tendency_output_1['tend1'] == 5.)
assert np.all(tendency_output_2['tend1'] == 5.)
def test_relaxation_prognostic_at_equilibrium():
prognostic = RelaxationTendencyComponent('quantity', 'degK')
state = {
'time':
timedelta(0),
'quantity':
DataArray(np.array([0., 1., 2.]), attrs={'units': 'degK'}),
'quantity_relaxation_timescale':
DataArray(np.array([1., 1., 1.]), attrs={'units': 's'}),
'equilibrium_quantity':
DataArray(np.array([0., 1., 2.]), attrs={'units': 'degK'}),
}
tendencies, diagnostics = prognostic(state)
assert np.all(tendencies['quantity'].values == 0.)
def test_relaxation_prognostic_with_change_different_equilibrium_units():
prognostic = RelaxationTendencyComponent('quantity', 'm')
state = {
'time':
timedelta(0),
'quantity':
DataArray(np.array([0., 1., 2.]), attrs={'units': 'm'}),
'quantity_relaxation_timescale':
DataArray(np.array([1., 2., 3.]), attrs={'units': 's'}),
'equilibrium_quantity':
DataArray(np.array([1., 3., 5.]) * 1e-3, attrs={'units': 'km'}),
}
tendencies, diagnostics = prognostic(state)
assert np.all(tendencies['quantity'].values == np.array([1., 1., 1.]))
def test_copies_untouched_quantities(self):
input_properties = {}
diagnostic_properties = {}
tendency_properties = {
'output1': {
'dims': ['dim1'],
'units': 'm/s'
},
}
diagnostic_output = {}
tendency_output = {
'output1': np.ones([10]) * 2.,
}
prognostic = self.prognostic_class(input_properties,
diagnostic_properties,
tendency_properties,
diagnostic_output, tendency_output)
stepper = self.timestepper_class(prognostic,
tendencies_in_diagnostics=True,
name='component')
untouched_quantity = DataArray(np.ones([10]) * 10.,
dims=['dim1'],
attrs={'units': 'J'})
state = {
'time':
timedelta(0),
'output1':
DataArray(np.ones([10]) * 10., dims=['dim1'], attrs={'units':
'm'}),
'input1':
untouched_quantity,
}
_, new_state = stepper(state, timedelta(seconds=5))
assert 'input1' in new_state.keys()
assert new_state['input1'].dims == untouched_quantity.dims
assert np.allclose(new_state['input1'].values, 10.)
assert new_state['input1'].attrs['units'] == 'J'
def get_test_state(pc_value=0.):
n_features = marble.components.marble.name_feature_counts
state = {
'time':
sp.timedelta(0),
'liquid_water_static_energy_components':
sp.DataArray(np.ones([n_features['sl']]) * pc_value,
dims=('sl_latent', ),
attrs={'units': ''}),
'total_water_mixing_ratio_components':
sp.DataArray(np.ones([n_features['rt']]) * pc_value,
dims=('rt_latent', ),
attrs={'units': ''}),
'cloud_water_mixing_ratio_components':
sp.DataArray(np.ones([n_features['rcld']]) * pc_value,
dims=('rcld_latent', ),
attrs={'units': ''}),
'rain_water_mixing_ratio_components':
sp.DataArray(np.ones([n_features['rrain']]) * pc_value,
dims=('rrain_latent', ),
attrs={'units': ''}),
'cloud_fraction_components':
sp.DataArray(np.ones([n_features['cld']]) * pc_value,
dims=('cld_latent', ),
attrs={'units': ''}),
'liquid_water_static_energy_components_horizontal_advective_tendency':
sp.DataArray(np.ones([n_features['sl']]) * pc_value,
dims=('sl_latent', ),
attrs={'units': ''}),
'total_water_mixing_ratio_components_horizontal_advective_tendency':
sp.DataArray(np.ones([n_features['sl']]) * pc_value,
dims=('rt_latent', ),
attrs={'units': ''}),
'vertical_wind_components':
sp.DataArray(np.ones([n_features['w']]) * pc_value,
dims=('w_latent', ),
attrs={'units': ''}),
}
return state
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 = {'time': timedelta(0), 'air_temperature': 273.}
timestep = timedelta(seconds=1.)
time_stepper = Leapfrog(MockEmptyTendencyComponent(),
asselin_strength=0.5,
alpha=1.)
diagnostics, new_state = time_stepper.__call__(state, timestep)
assert state == {'time': timedelta(0), 'air_temperature': 273.}
assert new_state == {'time': timedelta(0), '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 == {'time': timedelta(0), 'air_temperature': 274.}
assert new_state == {'time': timedelta(0), 'air_temperature': 277.}
def test_constant_diagnostic_empty_dict():
diag = ConstantDiagnosticComponent({})
diagnostics = diag({'time': timedelta(0)})
assert isinstance(diagnostics, dict)
assert len(diagnostics) == 0
def test_constant_diagnostic_cannot_modify_through_output_dict():
diag = ConstantDiagnosticComponent({})
diagnostics = diag({'time': timedelta(0)})
diagnostics['c'] = 'd'
diagnostics = diag({'time': timedelta(0)})
assert len(diagnostics) == 0