def test_berger_insolation_with_bad_solar_constant(): berger = BergerSolarInsolation() nx = 5 ny = 10 input_state = { 'time': datetime(2016, 12, 20, 6), 'longitude': DataArray(np.linspace(-90, 90, nx, endpoint=False), dims=['longitude'], attrs={'units': 'degree_E'}), 'latitude': DataArray(np.linspace(-180., 180., num=ny), dims=['latitude'], attrs={'units': 'degrees_north'}), 'solar_constant': DataArray(1364. * np.ones((2)), dims=['latitude'], attrs={'units': 'W m^-2'}), } with pytest.raises(ValueError) as excinfo: berger(input_state) assert 'Solar constant should' in str(excinfo.value)
def test_constant_diagnostic_diagnostic_properties(): diagnostics = { 'diag1': DataArray( np.zeros([10]), dims=['dim1'], attrs={'units': 'm'}, ), 'diag2': DataArray( np.zeros([2, 2]), dims=['dim2', 'dim3'], attrs={'units': 'degK'}, ) } diagnostic = ConstantDiagnosticComponent(diagnostics) assert diagnostic.diagnostic_properties == { 'diag1': { 'dims': ('dim1', ), 'units': 'm', }, 'diag2': { 'dims': ('dim2', 'dim3'), 'units': 'degK', } } assert diagnostic.input_properties == {}
def test_packs_two_tracers(self): np.random.seed(0) register_tracer('tracer1', 'g/m^3') register_tracer('tracer2', 'kg') input_state = { 'time': timedelta(0), 'tracer1': DataArray( np.random.randn(5), dims=['dim1'], attrs={'units': 'g/m^3'}, ), 'tracer2': DataArray( np.random.randn(5), dims=['dim1'], attrs={'units': 'kg'} ), } unpacked = self.call_component(input_state) packed = self.component.state_given['tracers'] assert isinstance(packed, np.ndarray) assert packed.shape == (2, 5) assert np.all(packed[0, :] == input_state['tracer1'].values) assert np.all(packed[1, :] == input_state['tracer2'].values) assert len(unpacked) == 2 assert np.all(unpacked['tracer1'].values == input_state['tracer1'].values) assert np.all(unpacked['tracer2'].values == input_state['tracer2'].values)
def test_hs_without_latitude(): hs = HeldSuarez() random = np.random.RandomState(0) input_state = { 'air_pressure': DataArray( random.rand(2, 3, 6), dims=['longitude', 'latitude', 'mid_levels'], attrs={'units': 'hPa'}, ), 'surface_air_pressure': DataArray( random.rand(2, 3), dims=['longitude', 'latitude'], attrs={'units': 'hPa'}, ), 'air_temperature': DataArray(270. + random.randn(2, 3, 6), dims=['longitude', 'latitude', 'mid_levels'], attrs={'units': 'degK'}), 'eastward_wind': DataArray(random.randn(2, 3, 6), dims=['longitude', 'latitude', 'mid_levels'], attrs={'units': 'm/s'}), 'northward_wind': DataArray(random.randn(2, 3, 6), dims=['longitude', 'latitude', 'mid_levels'], attrs={'units': 'm/s'}) } with pytest.raises(IndexError) as excinfo: hs(input_state) assert 'quantity labeled' in str(excinfo.value)
def test_match_dims_like_hardcoded_dimensions_non_matching_lengths(): input_state = { 'air_temperature': DataArray( np.zeros([2, 3, 4]), dims=['alpha', 'beta', 'gamma'], attrs={'units': 'degK'}, ), 'air_pressure': DataArray( np.zeros([4, 2, 3]), dims=['alpha', 'beta', 'gamma'], attrs={'units': 'Pa'}, ), } input_properties = { 'air_temperature': { 'dims': ['alpha', 'beta', 'gamma'], 'units': 'degK', 'match_dims_like': 'air_pressure', }, 'air_pressure': { 'dims': ['alpha', 'beta', 'gamma'], 'units': 'Pa', }, } try: raw_arrays = get_numpy_arrays_with_properties(input_state, input_properties) except InvalidStateError: pass else: raise AssertionError('should have raised InvalidStateError')
def test_match_dims_like_wildcard_dimensions_use_same_ordering(): input_state = { 'air_temperature': DataArray( np.random.randn(2, 3, 4), dims=['alpha', 'beta', 'gamma'], attrs={'units': 'degK'}, ), 'air_pressure': DataArray( np.zeros([4, 2, 3]), dims=['gamma', 'alpha', 'beta'], attrs={'units': 'Pa'}, ), } for i in range(4): input_state['air_pressure'][i, :, :] = input_state['air_temperature'][:, :, i] input_properties = { 'air_temperature': { 'dims': ['*'], 'units': 'degK', 'match_dims_like': 'air_pressure', }, 'air_pressure': { 'dims': ['*'], 'units': 'Pa', }, } raw_arrays = get_numpy_arrays_with_properties(input_state, input_properties) assert np.all(raw_arrays['air_temperature'] == raw_arrays['air_pressure'])
def test_match_dims_like_partly_hardcoded_dimensions_matching_lengths(): input_state = { 'air_temperature': DataArray( np.zeros([2, 3, 4]), dims=['lat', 'lon', 'mid_levels'], attrs={'units': 'degK'}, ), 'air_pressure': DataArray( np.zeros([2, 3, 4]), dims=['lat', 'lon', 'interface_levels'], attrs={'units': 'Pa'}, ), } input_properties = { 'air_temperature': { 'dims': ['*', 'mid_levels'], 'units': 'degK', 'match_dims_like': 'air_pressure', }, 'air_pressure': { 'dims': ['*', 'interface_levels'], 'units': 'Pa', }, } raw_arrays = get_numpy_arrays_with_properties(input_state, input_properties) assert np.byte_bounds(input_state['air_temperature'].values) == np.byte_bounds(raw_arrays['air_temperature']) assert np.byte_bounds(input_state['air_pressure'].values) == np.byte_bounds(raw_arrays['air_pressure'])
def test_array_addition_keeps_left_attr(): a = DataArray(np.array([1., 2., 3.]), attrs={'units': 'K'}) b = DataArray(np.array([2., 1., 3.]), attrs={'units': 'm/s', 'foo': 'bar'}) result = a + b assert (result.values == np.array([3., 3., 6.])).all() assert len(result.attrs) == 1 assert result.attrs['units'] == 'K'
def test_match_dims_like_star_z_matching_lengths(self): set_direction_names(x=['lat'], y=['lon'], z=['mid_levels', 'interface_levels']) input_state = { 'air_temperature': DataArray( np.zeros([2, 3, 4]), dims=['lat', 'lon', 'interface_levels'], attrs={'units': 'degK'}, ), 'air_pressure': DataArray( np.zeros([2, 3, 4]), dims=['lat', 'lon', 'interface_levels'], attrs={'units': 'Pa'}, ), } input_properties = { 'air_temperature': { 'dims': ['*', 'z'], 'units': 'degK', 'match_dims_like': 'air_pressure', }, 'air_pressure': { 'dims': ['*', 'z'], 'units': 'Pa', }, } raw_arrays = get_numpy_arrays_with_properties(input_state, input_properties) assert np.byte_bounds(raw_arrays['air_temperature']) == np.byte_bounds( input_state['air_temperature'].values) assert np.byte_bounds(raw_arrays['air_pressure']) == np.byte_bounds( input_state['air_pressure'].values)
def test_packs_prepended_and_normal_tracers_register_after_init(self): self.component = self.component.__class__( prepend_tracers=[('tracer1', 'g/m^3')]) register_tracer('tracer2', 'J/m^3') input_state = { 'time': timedelta(0), 'tracer1': DataArray( np.random.randn(5), dims=['dim1'], attrs={'units': 'g/m^3'}, ), 'tracer2': DataArray( np.random.randn(5), dims=['dim1'], attrs={'units': 'J/m^3'}, ), } unpacked = self.call_component(input_state) packed = self.component.state_given['tracers'] assert isinstance(packed, np.ndarray) assert packed.shape == (2, 5) assert np.all(packed[0, :] == input_state['tracer1'].values) assert np.all(packed[1, :] == input_state['tracer2'].values) assert len(unpacked) == 2 assert np.all( unpacked['tracer1'].values == input_state['tracer1'].values) assert np.all( unpacked['tracer2'].values == input_state['tracer2'].values)
def test_array_unit_conversion_different_units(): a = DataArray(np.array([1., 2., 3.]), attrs={'units': 'km', 'foo': 'bar'}) result = a.to_units('m') assert (result.values == np.array([1000., 2000., 3000.])).all() assert len(result.attrs) == 2 assert result.attrs['units'] == 'm' assert result.attrs['foo'] == 'bar'
def setUp(self): self.array_1d = DataArray(np.zeros((2,)), dims=['lon']) self.array_2d = DataArray(np.zeros((2, 2)), dims=['lat', 'lon']) self.array_3d = DataArray(np.zeros((2, 2, 2)), dims=['lon', 'lat', 'interface_levels']) set_direction_names( x=['lon'], y=['lat'], z=['mid_levels', 'interface_levels'])
def get_3d_input_state(self): nx, ny, nz = 2, 3, 10 p_interface = DataArray( np.linspace(1e5, 0, nz + 1), dims=['interface_levels'], attrs={'units': 'Pa'}, ) p = DataArray(0.5 * (p_interface.values[1:] + p_interface.values[:-1]), dims=['mid_levels'], attrs={'units': 'Pa'}) random = np.random.RandomState(0) return { 'air_pressure': p, 'air_temperature': DataArray(270. + random.randn(nx, ny, nz), dims=['longitude', 'latitude', 'mid_levels'], attrs={'units': 'degK'}), 'specific_humidity': DataArray(random.rand(nx, ny, nz) * 15., dims=['longitude', 'latitude', 'mid_levels'], attrs={'units': 'kg/kg'}), 'air_pressure_on_interface_levels': p_interface, }
def test_dataarray_three_steps(self, mock_prognostic_call): mock_prognostic_call.return_value = ( {'air_temperature': DataArray(np.ones((3, 3)), attrs={'units': 'K/s'})}, {}) state = {'air_temperature': DataArray(np.ones((3, 3))*273., attrs={'units': 'K'})} timestep = timedelta(seconds=1.) time_stepper = self.timestepper_class([MockPrognostic()]) 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 len(state['air_temperature'].attrs) == 1 assert state['air_temperature'].attrs['units'] == 'K' assert list(new_state.keys()) == ['air_temperature'] assert (new_state['air_temperature'] == np.ones((3, 3))*274.).all() state = new_state assert len(state['air_temperature'].attrs) == 1 assert state['air_temperature'].attrs['units'] == 'K' diagnostics, new_state = time_stepper.__call__(new_state, timestep) assert list(state.keys()) == ['air_temperature'] assert (state['air_temperature'] == np.ones((3, 3))*274.).all() assert list(new_state.keys()) == ['air_temperature'] assert (new_state['air_temperature'] == np.ones((3, 3))*275.).all() state = new_state assert len(state['air_temperature'].attrs) == 1 assert state['air_temperature'].attrs['units'] == 'K' diagnostics, new_state = time_stepper.__call__(state, timestep) assert list(state.keys()) == ['air_temperature'] assert (state['air_temperature'] == np.ones((3, 3))*275.).all() assert list(new_state.keys()) == ['air_temperature'] assert (new_state['air_temperature'] == np.ones((3, 3))*276.).all() assert len(new_state['air_temperature'].attrs) == 1 assert new_state['air_temperature'].attrs['units'] == 'K'
def test_inputs_one_scaling_with_unit_conversion(self): self.input_properties = { 'input1': { 'dims': ['dim1'], 'units': 'm', }, 'input2': { 'dims': ['dim1'], 'units': 'm', }, } state = { 'time': timedelta(0), 'input1': DataArray(np.ones([10]), dims=['dim1'], attrs={'units': 'km'}), 'input2': DataArray(np.ones([10]), dims=['dim1'], attrs={'units': 'm'}), } base_component = self.get_component() component = ScalingWrapper(base_component, input_scale_factors={'input1': 0.5}) assert isinstance(component, self.component_type) self.call_component(component, state) assert base_component.state_given.keys() == state.keys() assert np.all(base_component.state_given['input1'] == 500.) assert np.all(base_component.state_given['input2'] == 1.)
def test_expands_named_dimension_with_wildcard_present(self): random = np.random.RandomState(0) T_array = random.randn(3) input_state = { 'air_pressure': DataArray( np.zeros([3, 4]), dims=['dim1', 'dim2'], attrs={'units': 'Pa'}, ), 'air_temperature': DataArray( T_array, dims=['dim1'], attrs={'units': 'degK'}, ) } input_properties = { 'air_pressure': { 'dims': ['*', 'dim2'], 'units': 'Pa', }, 'air_temperature': { 'dims': ['*', 'dim2'], 'units': 'degK', }, } return_value = get_numpy_arrays_with_properties(input_state, input_properties) assert return_value['air_temperature'].shape == (3, 4) assert np.all(return_value['air_temperature'] == T_array[:, None])
def test_all_requested_properties_are_returned(self): property_dictionary = { 'air_temperature': { 'dims': ['x', 'y', 'z'], 'units': 'degK', }, 'air_pressure': { 'dims': ['x', 'y', 'z'], 'units': 'Pa', }, } state = { 'air_temperature': DataArray( np.zeros([4], dtype=np.float64), dims=['z'], attrs={'units': 'degK'}, ), 'air_pressure': DataArray( np.zeros([2,2,4], dtype=np.float64), dims=['x', 'y', 'z'], attrs={'units': 'Pa'}, ), 'eastward_wind': DataArray( np.zeros([2,2,4], dtype=np.float64), attrs={'units': 'm/s'} ), } return_value = get_numpy_arrays_with_properties(state, property_dictionary) assert isinstance(return_value, dict) assert len(return_value.keys()) == 2 assert 'air_temperature' in return_value.keys() assert 'air_pressure' in return_value.keys() assert np.all(return_value['air_temperature'] == 0.) assert np.all(return_value['air_pressure'] == 0.)
def test_match_dims_like_wildcard_dimensions_matching_lengths(): input_state = { 'air_temperature': DataArray( np.zeros([2, 3, 4]), dims=['alpha', 'beta', 'gamma'], attrs={'units': 'degK'}, ), 'air_pressure': DataArray( np.zeros([2, 3, 4]), dims=['alpha', 'beta', 'gamma'], attrs={'units': 'Pa'}, ), } input_properties = { 'air_temperature': { 'dims': ['*'], 'units': 'degK', 'match_dims_like': 'air_pressure', }, 'air_pressure': { 'dims': ['*'], 'units': 'Pa', }, } raw_arrays = get_numpy_arrays_with_properties(input_state, input_properties)
def test_constant_prognostic_diagnostic_properties(): tendencies = {} diagnostics = { 'diag1': DataArray( np.zeros([10]), dims=['dim1'], attrs={'units': 'm'}, ), 'diag2': DataArray( np.zeros([2, 2]), dims=['dim2', 'dim3'], attrs={'units': 'degK'}, ) } prog = ConstantTendencyComponent(tendencies, diagnostics) assert prog.diagnostic_properties == { 'diag1': { 'dims': ('dim1', ), 'units': 'm', }, 'diag2': { 'dims': ('dim2', 'dim3'), 'units': 'degK', } } assert prog.tendency_properties == {} assert prog.input_properties == {}
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 = {'eastward_wind': DataArray(1., attrs={'units': 'm/s'})} timestep = timedelta(seconds=1.) time_stepper = self.timestepper_class([MockPrognostic()]) diagnostics, new_state = time_stepper.__call__(state, timestep) assert state == {'eastward_wind': DataArray(1., attrs={'units': 'm/s'})} assert new_state == {'eastward_wind': DataArray(21., attrs={'units': 'm/s'})}
def test_array_unit_conversion_different_units_doesnt_modify_original(): a = DataArray(np.array([1., 2., 3.]), attrs={'units': 'km', 'foo': 'bar'}) a.to_units('m') assert (a.values == np.array([1., 2., 3.])).all() assert len(a.attrs) == 2 assert a.attrs['foo'] == 'bar' assert a.attrs['units'] == 'km'
def test_array_subtraction_keeps_left_attrs(): a = DataArray(np.array([1., 2., 3.]), attrs={'units': 'm/s', 'foo': 'bar'}) b = DataArray(np.array([2., 1., 3.]), attrs={'units': 'K'}) result = a - b assert (result.values == np.array([-1., 1., 0.])).all() assert len(result.attrs) == 2 assert result.attrs['units'] == 'm/s' assert result.attrs['foo'] == 'bar'
def __call__(self, state, timestep): self._num_updates += 1 return ({ 'num_updates': DataArray([self._num_updates], attrs={'units': ''}) }, { 'value': DataArray([1], attrs={'units': 'm'}) })
def update_radiative_state(self, atmosphere, surface, state0, sw=True): """Update CliMT formatted atmospheric state using parameters from our model. Parameters: atmosphere (konrad.atmosphere.Atmosphere): Atmosphere model. surface (konrad.surface): Surface model. state0 (dictionary): atmospheric state in the format for climt sw (bool): Toggle between shortwave and longwave calculations. Returns: dictionary: updated state """ state0["air_temperature"] = DataArray(atmosphere["T"][0, :], dims=("mid_levels", ), attrs={"units": "degK"}) vmr_h2o = atmosphere["H2O"][0, :] specific_humidity = vmr2specific_humidity(vmr_h2o) state0["specific_humidity"] = DataArray(specific_humidity, dims=("mid_levels", ), attrs={"units": "g/g"}) # CliMT/konrad name mapping gas_name_mapping = [ ("mole_fraction_of_methane_in_air", "CH4"), ("mole_fraction_of_carbon_dioxide_in_air", "CO2"), ("mole_fraction_of_nitrous_oxide_in_air", "N2O"), ("mole_fraction_of_ozone_in_air", "O3"), ("mole_fraction_of_cfc11_in_air", "CFC11"), ("mole_fraction_of_cfc12_in_air", "CFC12"), ("mole_fraction_of_cfc22_in_air", "CFC22"), ("mole_fraction_of_carbon_tetrachloride_in_air", "CCl4"), ("mole_fraction_of_oxygen_in_air", "O2"), ] for climt_key, konrad_key in gas_name_mapping: vmr = atmosphere.get(konrad_key, default=0, keepdims=False) state0[climt_key] = DataArray(vmr, dims=("mid_levels", ), attrs={"units": "mole/mole"}) # Surface quantities state0["surface_temperature"] = DataArray(np.array( surface["temperature"][-1]), attrs={"units": "degK"}) if sw: # properties required only for shortwave state0["zenith_angle"] = DataArray( np.array(np.deg2rad(self.current_solar_angle)), attrs={"units": "radians"}, ) return state0
def __call__(self, state): """ Get the Held-Suarez forcing tendencies Args: state (dict): A model state dictionary. Returns: tendencies (dict), diagnostics (dict): * A dictionary whose keys are strings indicating state quantities and values are the time derivative of those quantities in units/second at the time of the input state. * A dictionary whose keys are strings indicating state quantities and values are the value of those quantities at the time of the input state. Raises: IndexError: if the input state does not contain the key :code:`latitude`. """ if 'latitude' not in state: raise IndexError('state must contain a quantity labeled "latitude"') sigma = state['air_pressure']/state['surface_air_pressure'] sigma.attrs['units'] = '' Teq = self._get_Teq(state['latitude'], state['air_pressure']) k_t = self._get_k_t(state['latitude'], sigma) k_v = self._get_k_v(sigma) input_arrays = self.get_numpy_arrays_from_state('_climt_inputs', state) tendencies = { 'eastward_wind': DataArray( - k_v.values * input_arrays['eastward_wind'], dims=combine_dimensions( [k_v, state['eastward_wind']], out_dims=('x', 'y', 'z')), attrs={'units': 'm s^-2'}).squeeze(), 'northward_wind': DataArray( - k_v.values * input_arrays['northward_wind'], dims=combine_dimensions( [k_v, state['northward_wind']], out_dims=('x', 'y', 'z')), attrs={'units': 'm s^-2'}).squeeze(), 'air_temperature': DataArray( - k_t.values * (input_arrays['air_temperature'] - Teq.values), dims=combine_dimensions( [k_t, state['air_temperature']], out_dims=('x', 'y', 'z')), attrs={'units': 'K s^-1'}).squeeze() } return tendencies, {}
def update_radiative_state(self, atmosphere, surface, state0, sw=True): """ Update CliMT formatted atmospheric state using parameters from our model. Parameters: atmosphere (konrad.atmosphere.Atmosphere): Atmosphere model. surface (konrad.surface): Surface model. state0 (dictionary): atmospheric state in the format for climt sw (bool): Toggle between shortwave and longwave calculations. Returns: dictionary: updated state """ state0['air_temperature'] = DataArray(atmosphere['T'][0, :], dims=('mid_levels', ), attrs={'units': 'degK'}) vmr_h2o = atmosphere['H2O'][0, :] specific_humidity = vmr2specific_humidity(vmr_h2o) state0['specific_humidity'] = DataArray(specific_humidity, dims=('mid_levels', ), attrs={'units': 'g/g'}) # CliMT/konrad name mapping gas_name_mapping = [ ('mole_fraction_of_methane_in_air', 'CH4'), ('mole_fraction_of_carbon_dioxide_in_air', 'CO2'), ('mole_fraction_of_nitrous_oxide_in_air', 'N2O'), ('mole_fraction_of_ozone_in_air', 'O3'), ('mole_fraction_of_cfc11_in_air', 'CFC11'), ('mole_fraction_of_cfc12_in_air', 'CFC12'), ('mole_fraction_of_cfc22_in_air', 'CFC22'), ('mole_fraction_of_carbon_tetrachloride_in_air', 'CCl4'), ('mole_fraction_of_oxygen_in_air', 'O2'), ] for climt_key, konrad_key in gas_name_mapping: state0[climt_key] = DataArray(atmosphere.get(konrad_key, default=0, keepdims=False), dims=('mid_levels', ), attrs={'units': 'mole/mole'}) # Surface quantities state0['surface_temperature'] = DataArray(np.array( surface['temperature'][-1]), attrs={'units': 'degK'}) if sw: # properties required only for shortwave state0['zenith_angle'] = DataArray(np.array( np.deg2rad(self.current_solar_angle)), attrs={'units': 'radians'}) return state0
def get_3d_input_state(self): state = { 'latitude': DataArray(np.linspace(-90, 90, num=10), dims=['latitude'], attrs={'units': 'degrees_N'}), 'sigma_on_interface_levels': DataArray(np.linspace(0., 1., num=6), dims=['interface_levels'], attrs={'units': ''}), } return state
def test_relaxation_prognostic_caching_timescale_and_equilibrium(): prognostic = RelaxationPrognostic( 'quantity', relaxation_timescale=DataArray(np.array([1., 1., 1.]), attrs={'units': 's'}), equilibrium_value=DataArray(np.array([1., 3., 5.]), attrs={'units': 'degK'})) state = { 'quantity': DataArray(np.array([0., 1., 2.]), attrs={'units': 'degK'}), } tendencies, diagnostics = prognostic(state) assert np.all(tendencies['quantity'].values == np.array([1., 2., 3.]))
def test_relaxation_prognostic_with_change(): prognostic = RelaxationPrognostic('quantity') state = { '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([1., 3., 5.]), attrs={'units': 'degK'}), } tendencies, diagnostics = prognostic(state) assert np.all(tendencies['quantity'].values == np.array([1., 2., 3.]))
def test_packs_two_tracers(self): np.random.seed(0) dims = ['tracer', '*'] register_tracer('tracer1', 'g/m^3') register_tracer('tracer2', 'kg') state = { 'tracer1': DataArray(np.random.randn(5), dims=['dim1'], attrs={'units': 'g/m^3'}), 'tracer2': DataArray(np.random.randn(5), dims=['dim1'], attrs={'units': 'kg'}) } packer = TracerPacker(self.component, dims) packed = packer.pack(state) assert isinstance(packed, np.ndarray) assert packed.shape == (2, 5)