Source code for emc2.simulator.subcolumn

import numpy as np
import xarray as xr


[docs]def set_convective_sub_col_frac(model, hyd_type, N_columns=None): """ Sets the hydrometeor fraction due to convection in each subcolumn. Parameters ---------- model: :py:func: `emc2.core.Model` The model we are generating the subcolumns of convective fraction for. hyd_type: str The hydrometeor type to generate the fraction for. N_columns: int or None The number of subcolumns to generate. Specifying this will set the number of subcolumns in the model parameter when the first subcolumns are generated. Therefore, after those are generated this must either be equal to None or the number of subcolumns in the model. Setting this to None will use the number of subcolumns in the model parameter. Returns ------- model: :py:func: `emc2.core.Model` The Model object with the convective fraction in each subcolumn added. """ if model.num_subcolumns == 0 and N_columns is None: raise RuntimeError("The number of subcolumns must be specified in the model!") if model.num_subcolumns != N_columns and model.num_subcolumns != 0 and N_columns is not None: raise ValueError("The number of subcolumns has already been specified (%d) and != %d" % (model.num_subcolumns, N_columns)) if model.num_subcolumns == 0: model.ds['subcolumn'] = xr.DataArray(np.arange(0, N_columns), dims='subcolumn') data_frac = np.round(model.ds[model.conv_frac_names[hyd_type]].values * model.num_subcolumns) # In case we only have one time step if len(data_frac.shape) == 1: data_frac = data_frac[np.newaxis, :] conv_profs = np.zeros((model.num_subcolumns, data_frac.shape[0], data_frac.shape[1]), dtype=bool) for i in range(1, model.num_subcolumns + 1): for k in range(data_frac.shape[1]): mask = np.where(data_frac[:, k] == i)[0] conv_profs[0:i, mask, k] = True my_dims = model.ds[model.conv_frac_names[hyd_type]].dims model.ds[("conv_frac_subcolumns_" + hyd_type)] = xr.DataArray( conv_profs, dims=('subcolumn', my_dims[0], my_dims[1])) model.ds[("conv_frac_subcolumns_" + hyd_type)].attrs["units"] = "boolean" model.ds[("conv_frac_subcolumns_" + hyd_type)].attrs["long_name"] = ( "Is there hydrometeors of type %s in each subcolumn?" % hyd_type) return model
[docs]def set_stratiform_sub_col_frac(model): """ Sets the hydrometeor fraction due to stratiform cloud particles in each subcolumn. Parameters ---------- model: :py:func: `emc2.core.Model` The model we are generating the subcolumns of stratiform fraction for. Returns ------- model: :py:func: `emc2.core.Model` The Model object with the stratiform hydrometeor fraction in each subcolumn added. """ if "conv_frac_subcolumns_cl" not in model.ds.variables.keys(): raise KeyError("You have to generate the convective fraction in each subcolumn " + "before the stratiform fraction in each subcolumn is generated.") if "conv_frac_subcolumns_ci" not in model.ds.variables.keys(): raise KeyError("You have to generate the convective fraction in each subcolumn " + "before the stratiform fraction in each subcolumn is generated.") conv_profs1 = model.ds["conv_frac_subcolumns_cl"] conv_profs2 = model.ds["conv_frac_subcolumns_ci"] N_columns = len(model.ds["subcolumn"]) subcolumn_dims = conv_profs1.dims data_frac1 = model.ds[model.strat_frac_names["cl"]] data_frac1 = np.round(data_frac1.values * N_columns).astype(int) data_frac2 = model.ds[model.strat_frac_names["ci"]] data_frac2 = np.round(data_frac2.values * N_columns).astype(int) strat_profs1 = np.zeros((N_columns, data_frac1.shape[0], data_frac1.shape[1]), dtype=bool) strat_profs2 = np.zeros_like(strat_profs1, dtype=bool) conv_profs = np.logical_or(conv_profs1.values, conv_profs2.values) is_cloud = np.logical_or(data_frac1 > 0, data_frac2 > 0) is_cloud_one_above = np.roll(is_cloud, -1) is_cloud_one_above[:, -1] = False overlapping_cloud = np.logical_and(is_cloud, is_cloud_one_above) cld_2_assigns = np.stack([data_frac1, data_frac2], axis=0) cld_2_assign_min = cld_2_assigns.min(axis=0) cld_2_assign_max = cld_2_assigns.max(axis=0) I_min = np.argmin(cld_2_assigns, axis=0) I_max = np.argmax(cld_2_assigns, axis=0) for tt in range(data_frac1.shape[0]): for j in range(data_frac1.shape[1] - 1, -1, -1): cld_2_assign = np.array([data_frac1[tt, j], data_frac2[tt, j]]) I_min = np.argmin(cld_2_assign) I_max = np.argmax(cld_2_assign) if cld_2_assign[I_max] == 0: continue if overlapping_cloud[tt, j]: overlying_locs1 = np.where(np.logical_and(strat_profs1[:, tt, j + 1], ~conv_profs[:, tt, j]))[0] overlying_locs2 = np.where(np.logical_and(strat_profs2[:, tt, j + 1], ~conv_profs[:, tt, j]))[0] overlying_num = np.array([len(overlying_locs1), len(overlying_locs2)], dtype=int) over_diff = abs(overlying_num[1] - overlying_num[0]) Iover_min = np.argmin(overlying_num) Iover_max = np.argmax(overlying_num) over_unique_lo = _setxor(overlying_locs1, overlying_locs2) if overlying_num[Iover_min] > cld_2_assign[I_max]: if cld_2_assign[I_max] > 0: rand_locs = _randperm(overlying_num.min(), size=cld_2_assign[I_max]) inds = locals()["overlying_locs%d" % (Iover_min + 1)][rand_locs[0:cld_2_assign[I_min]]] locals()['strat_profs%d' % (I_min + 1)][inds, tt, j] = True inds = locals()["overlying_locs%d" % (Iover_min + 1)][rand_locs] locals()['strat_profs%d' % (I_max + 1)][inds, tt, j] = True cld_2_assign = np.zeros(2) elif overlying_num[Iover_min] > cld_2_assign[I_min]: if cld_2_assign[I_min] > 0: rand_locs = _randperm(overlying_num.min(), size=cld_2_assign[I_min]) inds = locals()["overlying_locs%d" % (Iover_min + 1)][rand_locs] locals()['strat_profs%d' % (I_min + 1)][inds, tt, j] = True inds = locals()["overlying_locs%d" % (Iover_min + 1)] locals()['strat_profs%d' % (I_max + 1)][inds, tt, j] = True cld_2_assign[I_min] = 0 cld_2_assign[I_max] -= overlying_num[Iover_min] if cld_2_assign[I_max] > 0 and over_diff > 0: rand_locs = _randperm(over_diff, size=cld_2_assign[I_max]) inds = over_unique_lo[rand_locs] locals()['strat_profs%d' % (I_max + 1)][inds, tt, j] = True cld_2_assign[I_max] = 0. else: locals()['strat_profs%d' % (I_max + 1)][over_unique_lo, tt, j] = True cld_2_assign[I_max] -= over_diff elif overlying_num[Iover_max] > cld_2_assign[I_min]: inds = locals()["overlying_locs%d" % (Iover_min + 1)] locals()['strat_profs%d' % (I_min + 1)][inds, tt, j] = True locals()['strat_profs%d' % (I_max + 1)][inds, tt, j] = True cld_2_assign -= overlying_num[Iover_min] if over_diff > cld_2_assign[I_max]: rand_locs = _randperm(over_diff, size=cld_2_assign[I_max]) inds = over_unique_lo[rand_locs[1:cld_2_assign[I_min]]] locals()['strat_profs%d' % (I_min + 1)][inds, tt, j] = True inds = over_unique_lo[rand_locs] locals()['strat_profs%d' % (I_max + 1)][inds, tt, j] = True cld_2_assign = np.zeros(2) else: if cld_2_assign[I_min] > 0: rand_locs = _randperm(over_diff, size=cld_2_assign[I_min]) inds = over_unique_lo[rand_locs] locals()['strat_profs%d' % (I_min + 1)][inds, tt, j] = True cld_2_assign[I_min] = 0 locals()['strat_profs%d' % (I_max + 1)][over_unique_lo, tt, j] = True cld_2_assign[I_max] -= over_diff else: inds = locals()["overlying_locs%d" % (Iover_max + 1)] locals()['strat_profs%d' % (I_min + 1)][inds, tt, j] = True locals()['strat_profs%d' % (I_max + 1)][inds, tt, j] = True cld_2_assign -= overlying_num[Iover_max] if cld_2_assign[I_max] > 0: sprof = locals()["strat_profs%d" % (I_max + 1)] free_locs_max = np.where(np.logical_and(~sprof[:, tt, j], ~conv_profs[:, tt, j]))[0] free_num = len(free_locs_max) rand_locs = _randperm(free_num, size=int(cld_2_assign[I_max])) locals()["strat_profs%d" % (I_max + 1)][free_locs_max[rand_locs], tt, j] = True if cld_2_assign[I_min] > 0.: locals()["strat_profs%d" % (I_min + 1)][free_locs_max[rand_locs[0:cld_2_assign[I_min]]], tt, j] = True model.ds['strat_frac_subcolumns_cl'] = xr.DataArray(strat_profs1, dims=(subcolumn_dims[0], subcolumn_dims[1], subcolumn_dims[2])) model.ds['strat_frac_subcolumns_ci'] = xr.DataArray(strat_profs2, dims=(subcolumn_dims[0], subcolumn_dims[1], subcolumn_dims[2])) model.ds['strat_frac_subcolumns_cl'].attrs["long_name"] = \ "Liquid cloud particles present? [stratiform]" model.ds['strat_frac_subcolumns_cl'].attrs["units"] = "0 = no, 1 = yes" model.ds['strat_frac_subcolumns_ci'].attrs["long_name"] = \ "Liquid cloud particles present? [stratiform]" model.ds['strat_frac_subcolumns_ci'].attrs["units"] = "0 = no, 1 = yes" return model
[docs]def set_precip_sub_col_frac(model, convective=True, N_columns=None): """ Sets the hydrometeor fraction due to precipitation in each subcolumn. This module works for both stratiform and convective precipitation. Parameters ---------- model: :py:func: `emc2.core.Model` The model we are generating the subcolumns of stratiform fraction for. convective: bool Set to True to generate subcolumns for convective precipitation. Set to False to generate subcolumns for stratiform precipitation. N_columns: int or None Use this to set the number of subcolumns in the model. This can only be set once. After the number of subcolumns is set, use None to make EMC2 automatically detect the number of subcolumns. Returns ------- model: :py:func: `emc2.core.Model` The Model object with the stratiform hydrometeor fraction in each subcolumn added. """ if model.num_subcolumns == 0 and N_columns is None: raise RuntimeError("The number of subcolumns must be specified in the model!") if model.num_subcolumns != N_columns and model.num_subcolumns != 0 and N_columns is not None: raise ValueError("The number of subcolumns has already been specified (%d) and != %d" % (model.num_subcolumns, N_columns)) if model.num_subcolumns == 0: model.ds['subcolumn'] = xr.DataArray(np.arange(0, N_columns), dims='subcolumn') if convective: data_frac1 = model.ds[model.conv_frac_names['pl']] data_frac2 = model.ds[model.conv_frac_names['pi']] out_prof_name1 = 'conv_frac_subcolumns_pl' out_prof_name2 = 'conv_frac_subcolumns_pi' out_prof_long_name1 = 'Liquid precipitation present? [convective]' out_prof_long_name2 = 'Ice precipitation present? [convective]' in_prof_cloud_name_liq = 'conv_frac_subcolumns_cl' in_prof_cloud_name_ice = 'conv_frac_subcolumns_ci' else: data_frac1 = model.ds[model.strat_frac_names['pl']] data_frac2 = model.ds[model.strat_frac_names['pi']] out_prof_name1 = 'strat_frac_subcolumns_pl' out_prof_name2 = 'strat_frac_subcolumns_pi' in_prof_cloud_name_liq = 'strat_frac_subcolumns_cl' in_prof_cloud_name_ice = 'strat_frac_subcolumns_ci' out_prof_long_name1 = 'Liquid precipitation present? [stratiform]' out_prof_long_name2 = 'Ice precipitation present? [stratiform]' if in_prof_cloud_name_liq not in model.ds.variables.keys(): raise KeyError("%s is not a variable in the model object. Please ensure that you have" + "generated the cloud particle subcolumns before running this routine." % in_prof_cloud_name_liq) if in_prof_cloud_name_ice not in model.ds.variables.keys(): raise KeyError("%s is not a variable in the model object. Please ensure that you have" + "generated the cloud particle subcolumns before running this routine." % in_prof_cloud_name_ice) data_frac1 = np.round(data_frac1 * model.num_subcolumns) data_frac2 = np.round(data_frac2 * model.num_subcolumns) p_strat_profs = np.zeros((model.num_subcolumns, data_frac1.shape[0], data_frac1.shape[1], 2), dtype=bool) strat_profs = np.logical_or(model.ds[in_prof_cloud_name_ice].values, model.ds[in_prof_cloud_name_liq].values) subcolumn_dims = model.ds[in_prof_cloud_name_ice].dims is_cloud = np.logical_or(data_frac1 > 0, data_frac2 > 0) is_cloud_one_above = np.roll(is_cloud, -1) is_cloud_one_above[-1, :] = False overlapping_cloud = np.logical_and(is_cloud, is_cloud_one_above) precip_exist = np.stack([np.where(data_frac1 > 0, True, False), np.where(data_frac2 > 0, True, False)]) PF_val = np.max(np.stack([data_frac1, data_frac2]), axis=0) cond = [strat_profs, ~strat_profs] for tt in range(data_frac1.shape[0]): for j in range(data_frac1.shape[1] - 1, -1, -1): if overlapping_cloud[tt, j]: overlying_locs = np.where(np.any(p_strat_profs[:, tt, j + 1, :], axis=1))[0] overlying_num = len(overlying_locs) if overlying_num > PF_val[tt, j]: rand_locs = _randperm(overlying_num, PF_val[tt, j]) for i in range(2): if precip_exist[i, tt, j]: p_strat_profs[overlying_locs[rand_locs], tt, j, i] = True PF_val[tt, j] = 0. else: for i in range(2): if precip_exist[i, tt, j]: p_strat_profs[overlying_locs, tt, j, i] = True PF_val[tt, j] -= overlying_num for ii in range(2): free_locs = np.where(np.logical_and( ~np.all(p_strat_profs[:, tt, j, :], axis=1), cond[ii][:, tt, j]))[0] free_num = len(free_locs) if free_num > 0: if free_num > PF_val[tt, j]: rand_locs = _randperm(free_num, PF_val[tt, j]) for i in range(2): if precip_exist[i, tt, j]: p_strat_profs[free_locs[rand_locs], tt, j, i] = True PF_val[tt, j] = 0 else: for i in range(2): if precip_exist[i, tt, j]: p_strat_profs[free_locs, tt, j, i] = True PF_val[tt, j] -= free_num model.ds[out_prof_name1] = xr.DataArray(p_strat_profs[:, :, :, 0], dims=(subcolumn_dims[0], subcolumn_dims[1], subcolumn_dims[2])) model.ds[out_prof_name2] = xr.DataArray(p_strat_profs[:, :, :, 1], dims=(subcolumn_dims[0], subcolumn_dims[1], subcolumn_dims[2])) model.ds[out_prof_name1].attrs["long_name"] = out_prof_long_name1 model.ds[out_prof_name1].attrs["units"] = "0 = no, 1 = yes" model.ds[out_prof_name2].attrs["long_name"] = out_prof_long_name2 model.ds[out_prof_name2].attrs["units"] = "0 = no, 1 = yes" return model
[docs]def set_q_n(model, hyd_type, is_conv=True, qc_flag=True, inv_rel_var=1): """ This function distributes the mixing ratio and number concentration into the subcolumns. For :math:`q_c`, the horizontal distribution follows Equation 8 of Morrison and Gettelman (2008). Parameters ---------- model: :func:`emc2.core.Model` The model to calculate the mixing ratio in each subcolumn for. hyd_type: str The hydrometeor type. is_conv: bool Set to True to calculate the mixing ratio assuming convective clouds. qc_flag: bool Set to True to horizontally distribute the mixing ratio according to Morrison and Gettleman (2008) inv_rel_var: float The inverse of the relative variance for the p.d.f. in Morrison and Gettleman (2008) Returns ------- model: :func:`emc2.core.Model` The model with mixing ratio calculated in each subcolumn. References ---------- Morrison, H. and A. Gettelman, 2008: A New Two-Moment Bulk Stratiform Cloud Microphysics Scheme in the Community Atmosphere Model, Version 3 (CAM3). Part I: Description and Numerical Tests. J. Climate, 21, 3642–3659, https://doi.org/10.1175/2008JCLI2105.1 """ if model.num_subcolumns == 0: raise RuntimeError("The number of subcolumns must be specified in the model!") if not is_conv: hyd_profs = model.ds[model.strat_frac_names[hyd_type]] N_profs = model.ds[model.N_field[hyd_type]] N_profs = N_profs / hyd_profs.values sub_hyd_profs = model.ds['strat_frac_subcolumns_%s' % hyd_type].values N_profs = np.tile(N_profs.values, (model.num_subcolumns, 1, 1)) N_profs = np.where(sub_hyd_profs, N_profs, 0) q_array = model.ds[model.q_names_stratiform[hyd_type]] q_name = "strat_q_subcolumns_%s" % hyd_type n_name = "strat_n_subcolumns_%s" % hyd_type else: hyd_profs = model.ds[model.conv_frac_names[hyd_type]] sub_hyd_profs = hyd_profs.values q_array = model.ds[model.q_names_convective[hyd_type]] q_name = "conv_q_subcolumns_%s" % hyd_type q_ic_mean = q_array / hyd_profs if qc_flag: tot_hyd_in_sub = sub_hyd_profs.sum(axis=0) q_profs = np.zeros_like(sub_hyd_profs, dtype=float) for tt in range(hyd_profs.shape[1]): for j in range(hyd_profs.shape[0] - 1, -1, -1): hyd_in_sub_loc = np.where(sub_hyd_profs[:, tt, j])[0] if tot_hyd_in_sub[tt, j] == 1: q_profs[hyd_in_sub_loc, tt, j] = q_ic_mean[tt, j] elif tot_hyd_in_sub[tt, j] > 1: alpha = inv_rel_var / q_ic_mean[tt, j] a = inv_rel_var b = 1 / alpha randlocs = np.random.permutation(tot_hyd_in_sub[tt, j]) rand_gamma_vals = np.random.gamma(a, b, tot_hyd_in_sub[tt, j] - 1) valid_vals = False counter_4_valid = 0 while not valid_vals: counter_4_valid += 1 valid_vals = (q_ic_mean[tt, j] * tot_hyd_in_sub[tt, j] - rand_gamma_vals[0:-counter_4_valid].sum()) > 0 q_profs[hyd_in_sub_loc[ randlocs[:-(counter_4_valid + 1)]], tt, j] = rand_gamma_vals[:-counter_4_valid] q_profs[hyd_in_sub_loc[randlocs[-counter_4_valid:]], tt, j] = ( q_ic_mean[tt, j] * tot_hyd_in_sub[tt, j] - np.sum(rand_gamma_vals[:-counter_4_valid])) / \ (1 + counter_4_valid) else: q_profs = q_array / hyd_profs q_profs = np.tile(q_profs.values, (model.num_subcolumns, 1, 1)) q_profs = np.where(sub_hyd_profs, q_profs, 0) model.ds[q_name] = xr.DataArray(q_profs, dims=('subcolumn', model.time_dim, model.height_dim)) if not is_conv: model.ds[n_name] = xr.DataArray(N_profs, dims=('subcolumn', model.time_dim, model.height_dim)) return model
def _randperm(x, size=None): if size is None: size = len(x) return np.random.permutation(x)[0:int(size)].astype(int) def _setxor(x, y): first_set = np.setdiff1d(x, y) second_set = np.setdiff1d(y, x) return np.concatenate([first_set, second_set])