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