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_returns_simple_value(self): input_state = { 'air_temperature': DataArray( np.zeros([2, 2, 4]), dims=['x', 'y', 'z'], attrs={'units': 'degK'}, ) } input_properties = { 'air_temperature': { 'dims': ['x', 'y', 'z'], 'units': 'degK', } } raw_arrays = get_numpy_arrays_with_properties(input_state, input_properties) raw_arrays = {key + '_tendency': value for key, value in raw_arrays.items()} output_properties = { 'air_temperature_tendency': { 'dims': ['x', 'y', 'z'], 'units': 'degK/s', } } return_value = restore_data_arrays_with_properties( raw_arrays, output_properties, input_state, input_properties ) assert isinstance(return_value, dict) assert len(return_value.keys()) == 1 assert isinstance(return_value['air_temperature_tendency'], DataArray) assert return_value['air_temperature_tendency'].attrs['units'] is 'degK/s' assert np.byte_bounds( return_value['air_temperature_tendency'].values) == np.byte_bounds( input_state['air_temperature'].values) assert (return_value['air_temperature_tendency'].values.base is input_state['air_temperature'].values) assert return_value['air_temperature_tendency'].shape == (2, 2, 4)
def test_restores_scalar_array(self): T_array = np.array(0.) input_properties = { 'surface_temperature': { 'units': 'degK', 'dims': ['*'], }, } input_state = { 'surface_temperature': DataArray( T_array, dims=[], attrs={'units': 'degK'}, ), } raw_arrays = get_numpy_arrays_with_properties(input_state, input_properties) output_properties = { 'surface_temperature': { 'units': 'degK', } } return_value = restore_data_arrays_with_properties( raw_arrays, output_properties, input_state, input_properties ) assert len(return_value.keys()) == 1 assert 'surface_temperature' in return_value.keys() assert len(return_value['surface_temperature'].values.shape) == 0 assert return_value['surface_temperature'].attrs['units'] == 'degK'
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_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_raises_if_dims_property_not_specified(self): property_dictionary = { 'air_temperature': { 'units': 'degK', }, } state = { 'air_temperature': DataArray( np.zeros([4], dtype=np.float64), dims=['z'], attrs={'units': 'degK'}, ), } try: get_numpy_arrays_with_properties(state, property_dictionary) except InvalidPropertyDictError: pass else: raise AssertionError('should have raised ValueError')
def test_raises_when_quantity_has_extra_dim_and_unmatched_wildcard(self): input_state = { 'air_temperature': DataArray( np.zeros([2, 4]), dims=['foo', 'z'], attrs={'units': 'degK'}, ) } input_properties = { 'air_temperature': { 'dims': ['y', 'z'], 'units': 'degK', } } try: get_numpy_arrays_with_properties(input_state, input_properties) except InvalidStateError: pass else: raise AssertionError('should have raised InvalidStateError')
def test_returns_numpy_array_alias_doesnt_apply_to_state(self): T_array = np.zeros([2, 3, 4], dtype=np.float64) + 280. property_dictionary = { 'air_temperature': { 'units': 'degK', 'dims': ['x', 'y', 'z'], 'alias': 'T', }, } state = { 'T': DataArray( T_array, dims=['x', 'y', 'z'], attrs={'units': 'degK'}, ), } try: get_numpy_arrays_with_properties( state, property_dictionary) except KeyError: pass else: raise AssertionError('should have raised KeyError')
def test_unit_conversion_doesnt_modify_input(self): property_dictionary = { 'air_temperature': { 'dims': ['x', 'y', 'z'], 'units': 'degC', }, } state = { 'air_temperature': DataArray( np.zeros([4], dtype=np.float64), dims=['z'], attrs={'units': 'degK'}, ), } return_value = get_numpy_arrays_with_properties(state, property_dictionary) assert np.all(state['air_temperature'].values == 0.) assert state['air_temperature'].attrs['units'] is 'degK'
def test_raises_if_state_quantity_units_undefined(self): property_dictionary = { 'air_temperature': { 'dims': ['x', 'y', 'z'], 'units': 'degK', }, } state = { 'air_temperature': DataArray( np.zeros([4], dtype=np.float64), dims=['z'], ), } try: return_value = get_numpy_arrays_with_properties(state, property_dictionary) except InvalidStateError: pass else: raise AssertionError('should have raised InvalidStateError')
def test_converting_units_maintains_float32_dtype(self): property_dictionary = { 'air_temperature': { 'dims': ['x', 'y', 'z'], 'units': 'degC', }, } state = { 'air_temperature': DataArray( np.zeros([4], dtype=np.float32), dims=['z'], attrs={'units': 'degK'}, ), } return_value = get_numpy_arrays_with_properties(state, property_dictionary) assert isinstance(return_value, dict) assert len(return_value.keys()) == 1 assert 'air_temperature' in return_value.keys() assert return_value['air_temperature'].dtype is np.dtype('float32')
def test_returns_scalar_array(self): T_array = np.array(0.) property_dictionary = { 'air_temperature': { 'units': 'degK', 'dims': [], }, } state = { 'air_temperature': DataArray( T_array, dims=[], attrs={'units': 'degK'}, ), } return_value = get_numpy_arrays_with_properties(state, property_dictionary) assert isinstance(return_value, dict) assert len(return_value.keys()) == 1 assert isinstance(return_value['air_temperature'], np.ndarray) assert np.byte_bounds(return_value['air_temperature']) == np.byte_bounds( T_array)
def test_collects_horizontal_dimensions(self): random = np.random.RandomState(0) T_array = random.randn(3, 2, 4) input_state = { 'air_temperature': DataArray( T_array, dims=['x', 'y', 'z'], attrs={'units': 'degK'}, ) } input_properties = { 'air_temperature': { 'dims': ['z', '*'], 'units': 'degK', } } return_value = get_numpy_arrays_with_properties(input_state, input_properties) assert return_value['air_temperature'].shape == (4, 6) for i in range(3): for j in range(2): for k in range(4): assert return_value['air_temperature'][k, j+2*i] == T_array[i, j, k]
def test_returns_numpy_array(self): T_array = np.zeros([2, 3, 4], dtype=np.float64) + 280. property_dictionary = { 'air_temperature': { 'units': 'degK', 'dims': ['x', 'y', 'z'], }, } state = { 'air_temperature': DataArray( T_array, dims=['x', 'y', 'z'], attrs={'units': 'degK'}, ), } return_value = get_numpy_arrays_with_properties(state, property_dictionary) assert isinstance(return_value, dict) assert len(return_value.keys()) == 1 assert isinstance(return_value['air_temperature'], np.ndarray) assert np.byte_bounds(return_value['air_temperature']) == np.byte_bounds( T_array) assert return_value['air_temperature'].base is T_array
def __call__(self, state): """ Gets tendencies and diagnostics from the passed model state. Copied from sympl develop branch (to-be v0.3.3), ignoring checks. Args ---- state : dict A model state dictionary. Returns ------- tendencies : 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. diagnostics : dict 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 ------ KeyError If a required quantity is missing from the state. InvalidStateError If state is not a valid input for the Prognostic instance. """ raw_state = get_numpy_arrays_with_properties(state, self.input_properties) raw_state['time'] = state['time'] raw_tendencies, raw_diagnostics = self.array_call(raw_state) tendencies = restore_data_arrays_with_properties( raw_tendencies, self.tendency_properties, state, self.input_properties) diagnostics = restore_data_arrays_with_properties( raw_diagnostics, self.diagnostic_properties, state, self.input_properties) return tendencies, diagnostics