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