def combine_dimensions_1d_and_2d_extra_direction(self):
try:
combine_dimensions(
[self.array_1d, self.array_2d], out_dims=['y'])
except ValueError:
pass
except Exception as err:
raise err
else:
raise AssertionError('No exception raised but expected ValueError.')
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 __call__(self, state):
"""
Calculate longwave optical depth from input state.
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.
"""
lat = get_numpy_array(state['latitude'].to_units('degrees_north'),
out_dims=('x', 'y', 'z'))
sigma_interface = get_numpy_array(
state['sigma_on_interface_levels'].to_units(''),
out_dims=('x', 'y', 'z'))
tau = DataArray(
get_frierson_06_tau(lat, sigma_interface, self._tau0e, self._tau0p,
self._fl),
dims=combine_dimensions(
[state['latitude'], state['sigma_on_interface_levels']],
out_dims=('x', 'y', 'z')),
attrs={
'units': ''
}).squeeze()
return {
'longwave_optical_depth_on_interface_levels': tau,
}
def _get_k_v(self, sigma):
out_dims = combine_dimensions([sigma], out_dims=('x', 'y', 'z'))
sigma = get_numpy_array(sigma.to_units(''), out_dims=('x', 'y', 'z'))
return DataArray(
self._k_f * np.maximum(
0,
(sigma - self._sigma_b)/(1 - self._sigma_b)),
dims=out_dims,
attrs={'units': 's^-1'})
def _get_k_t(self, latitude, sigma):
out_dims = combine_dimensions([latitude, sigma], out_dims=('x', 'y', 'z'))
latitude = get_numpy_array(
latitude.to_units('degrees_N'), out_dims=('x', 'y', 'z'))
sigma = get_numpy_array(sigma, out_dims=('x', 'y', 'z'))
return DataArray(
self._k_a +
(self._k_s - self._k_a) *
np.maximum(0, (sigma - self._sigma_b)/(1 - self._sigma_b)) *
np.cos(np.radians(latitude))**4,
dims=out_dims,
attrs={'units': 's^-1'})
def _get_Teq(self, latitude, air_pressure):
out_dims = combine_dimensions(
[latitude, air_pressure], out_dims=('x', 'y', 'z'))
latitude = get_numpy_array(
latitude.to_units('degrees_N'), out_dims=['x', 'y', 'z'])
air_pressure = get_numpy_array(
air_pressure.to_units('Pa'), out_dims=['x', 'y', 'z'])
return DataArray(np.maximum(
200,
(315 - self._delta_T_y*np.sin(np.radians(latitude))**2 -
self._delta_theta_z*np.log(air_pressure/self._p0)*np.cos(np.radians(latitude))**2
) * (air_pressure/self._p0)**self._kappa),
dims=out_dims,
attrs={'units': 'degK'})
def __call__(self, state, timestep):
'''
Calculate surface and boundary layer tendencies.
Args:
state (dict):
The model state dictionary
timestep (timedelta):
The model timestep
Returns:
state (dict), diagnostics(dict) :
* The updated model state.
* diagnostics for Simple Physics
'''
if 'latitude' not in state.keys():
raise IndexError(
'Simple Physics: State must contain a quantity called latitude'
)
U = get_numpy_array(state['eastward_wind'].to_units('m/s'),
['x', 'y', 'z'])
V = get_numpy_array(state['northward_wind'].to_units('m/s'),
['x', 'y', 'z'])
P = get_numpy_array(state['air_pressure'].to_units('Pa'),
['x', 'y', 'z'])
Pint = get_numpy_array(
state['air_pressure_on_interface_levels'].to_units('Pa'),
['x', 'y', 'z'])
T = get_numpy_array(state['air_temperature'].to_units('degK'),
['x', 'y', 'z'])
q = get_numpy_array(state['specific_humidity'].to_units('g/g'),
['x', 'y', 'z'])
q_surface = get_numpy_array(
state['surface_specific_humidity'].to_units('g/g'), ['x', 'y'])
Ts = get_numpy_array(state['surface_temperature'].to_units('degK'),
['x', 'y'])
Ps = get_numpy_array(state['surface_air_pressure'].to_units('Pa'),
['x', 'y'])
lats = get_numpy_array(state['latitude'].to_units('degrees_north'),
['x', 'y'])
lats = np.asfortranarray(lats, dtype=np.double)
if lats.shape[0] == 1: # 1-d array only
num_longitudes = Ts.shape[0]
lat_list = [lats[0, :] for i in range(num_longitudes)]
lats = np.asfortranarray(np.stack(lat_list, axis=0))
dims_mid = combine_dimensions([
state['surface_temperature'], state['air_temperature'],
state['eastward_wind'], state['northward_wind']
], ['x', 'y', 'z'])
(t_out, u_out, v_out, q_out, precip_out, sensible_heat_flux,
latent_heat_flux) = phys.get_new_state(U, V, T, P, Pint, q, Ps, Ts,
q_surface, lats,
timestep.total_seconds())
latent_heat_flux[latent_heat_flux < 0] = 0
new_state = {
'eastward_wind':
DataArray(u_out, dims=dims_mid,
attrs=state['eastward_wind'].attrs),
'northward_wind':
DataArray(v_out,
dims=dims_mid,
attrs=state['northward_wind'].attrs),
'air_temperature':
DataArray(t_out,
dims=dims_mid,
attrs=state['air_temperature'].attrs),
'specific_humidity':
DataArray(q_out,
dims=dims_mid,
attrs=state['specific_humidity'].attrs),
}
diagnostics = self.create_state_dict_for('_climt_diagnostics', state)
diagnostics['stratiform_precipitation_rate'].values[:] = precip_out
diagnostics[
'surface_upward_sensible_heat_flux'].values[:] = sensible_heat_flux
diagnostics[
'surface_upward_latent_heat_flux'].values[:] = latent_heat_flux
return diagnostics, new_state
def test_combine_dimensions_2d_and_3d_z_y_x(self):
dims = combine_dimensions(
[self.array_2d, self.array_3d], out_dims=('z', 'y', 'x'))
assert same_list(dims, ['interface_levels', 'lat', 'lon'])
def combine_dimensions_1d_not_shared(self):
array_1d_x = DataArray(np.zeros((2,)), dims=['lon'])
array_1d_y = DataArray(np.zeros((2,)), dims=['lat'])
dims = combine_dimensions([array_1d_x, array_1d_y], out_dims=['x', 'y'])
assert same_list(dims, ['lon', 'lat'])
def combine_dimensions_1d_shared(self):
dims = combine_dimensions(
[self.array_1d, self.array_1d], out_dims=['x'])
assert same_list(dims, ['lon'])
def __call__(self, state, timestep):
"""
Gets diagnostics from the current model state and steps the state
forward in time according to the timestep.
Args:
state (dict): A model state dictionary. Will be updated with any
diagnostic quantities produced by this object for the time of
the input state.
Returns:
next_state (dict): A dictionary whose keys are strings indicating
state quantities and values are the value of those quantities
at the timestep after input state.
Raises:
KeyError: If a required quantity is missing from the state.
InvalidStateException: If state is not a valid input for the
Implicit instance for other reasons.
"""
T = get_numpy_array(state['air_temperature'].to_units('degK'),
out_dims=('x', 'y', 'z'))
q = get_numpy_array(state['specific_humidity'].to_units('kg/kg'),
out_dims=('x', 'y', 'z'))
p = get_numpy_array(state['air_pressure'].to_units('Pa'),
out_dims=('x', 'y', 'z'))
p_interface = get_numpy_array(
state['air_pressure_on_interface_levels'].to_units('Pa'),
out_dims=('x', 'y', 'z'))
q_sat = bolton_q_sat(T, p, self._Rd, self._Rh2O)
saturated = q > q_sat
dqsat_dT = bolton_dqsat_dT(T[saturated], self._Lv, self._Rh2O,
q_sat[saturated])
condensed_q = np.zeros_like(q)
condensed_q[saturated] = (q[saturated] - q_sat[saturated]) / (
1 + self._Lv / self._Cpd * dqsat_dT)
new_q = q.copy()
new_T = T.copy()
new_q[saturated] -= condensed_q[saturated]
new_T[saturated] += self._Lv / self._Cpd * condensed_q[saturated]
mass = (p_interface[:, :, 1:] - p_interface[:, :, :-1]) / (self._g *
self._rhow)
precipitation = np.sum(condensed_q * mass, axis=2)
dims_3d = combine_dimensions([
state['air_temperature'], state['specific_humidity'],
state['air_pressure']
],
out_dims=('x', 'y', 'z'))
dims_2d = dims_3d[:-1]
diagnostics = {
'column_integrated_precipitation_rate':
DataArray(precipitation / timestep.total_seconds(),
dims=dims_2d,
attrs={
'units': 'kg/s'
}).squeeze()
}
new_state = {
'air_temperature':
DataArray(new_T,
dims=dims_3d,
attrs=state['air_temperature'].attrs).squeeze(),
'specific_humidity':
DataArray(new_q,
dims=dims_3d,
attrs=state['specific_humidity'].attrs).squeeze(),
}
return new_state, diagnostics
