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