def make_grid_storage(self, pygrid): shape = pygrid.domain_shape_full(add=(1, 1, 1)) for k in TranslateGrid.composite_grid_vars: if k in self.data: self.make_composite_var_storage(k, self.data[k], shape) del self.data[k] for k, axis in TranslateGrid.edge_var_axis.items(): if k in self.data: self.data[k] = utils.make_storage_data( self.data[k], shape, start=(0, 0, pygrid.halo), axis=axis, read_only=True, ) for k, v in self.data.items(): if type(v) is np.ndarray: # TODO: when grid initialization model exists, may want to use # it to inform this istart, jstart = pygrid.horizontal_starts_from_shape(v.shape) logger.debug("Storage for Grid variable {}, {}, {}, {}".format( k, istart, jstart, v.shape)) origin = (istart, jstart, 0) self.data[k] = utils.make_storage_data(v, shape, origin=origin, start=origin, read_only=True)
def _initialize_interpolation_constants(self, stencil_factory: StencilFactory, grid_indexing: GridIndexing): # because stencils only work on 3D at the moment, need to compute in 3D # and then make these 1D gk_3d = utils.make_storage_from_shape( (1, 1, grid_indexing.domain[2] + 1), (0, 0, 0)) gamma_3d = utils.make_storage_from_shape( (1, 1, grid_indexing.domain[2] + 1), (0, 0, 0)) beta_3d = utils.make_storage_from_shape( (1, 1, grid_indexing.domain[2] + 1), (0, 0, 0)) _cubic_spline_interpolation_constants = stencil_factory.from_origin_domain( cubic_spline_interpolation_constants, origin=(0, 0, 0), domain=(1, 1, grid_indexing.domain[2] + 1), ) _cubic_spline_interpolation_constants(self._dp0, gk_3d, beta_3d, gamma_3d) self._gk = utils.make_storage_data(gk_3d[0, 0, :], gk_3d.shape[2:], (0, )) self._beta = utils.make_storage_data(beta_3d[0, 0, :], beta_3d.shape[2:], (0, )) self._gamma = utils.make_storage_data(gamma_3d[0, 0, :], gamma_3d.shape[2:], (0, )) self._copy_corners = PreAllocatedCopiedCornersFactory( stencil_factory=stencil_factory, dims=[ fv3gfs.util.X_DIM, fv3gfs.util.Y_DIM, fv3gfs.util.Z_INTERFACE_DIM ], y_temporary=None, )
def compute_parallel(self, inputs, communicator): # ak, bk, pfull, and phis are numpy arrays at this point and # must be converted into gt4py storages for name in ("ak", "bk", "pfull", "phis"): inputs[name] = utils.make_storage_data( inputs[name], inputs[name].shape, len(inputs[name].shape) * (0, )) grid_data = spec.grid.grid_data grid_data.ak = inputs["ak"] grid_data.bk = inputs["bk"] self._base.compute_func = dyn_core.AcousticDynamics( communicator, spec.grid.stencil_factory, grid_data, spec.grid.damping_coefficients, spec.grid.grid_type, spec.grid.nested, spec.grid.stretched_grid, spec.namelist.dynamical_core.acoustic_dynamics, inputs["ak"], inputs["bk"], inputs["pfull"], inputs["phis"], ) return super().compute_parallel(inputs, communicator)
def make_storage_data( self, array: np.ndarray, istart: int = 0, jstart: int = 0, kstart: int = 0, dummy_axes: Optional[Tuple[int, int, int]] = None, axis: int = 2, names_4d: Optional[List[str]] = None, ) -> Dict[str, type(Field)]: use_shape = list(self.maxshape) if dummy_axes: for axis in dummy_axes: use_shape[axis] = 1 use_shape = tuple(use_shape) start = (istart, jstart, kstart) if names_4d: return utils.make_storage_dict( array, use_shape, start=start, origin=start, dummy=dummy_axes, axis=axis, names=names_4d, ) else: return utils.make_storage_data( array, use_shape, start=start, origin=start, dummy=dummy_axes, axis=axis, )
def compute_parallel(self, inputs, communicator): # ak, bk, and phis are numpy arrays at this point and # must be converted into gt4py storages for name in ("ak", "bk", "phis"): inputs[name] = utils.make_storage_data( inputs[name], inputs[name].shape, len(inputs[name].shape) * (0,) ) inputs["comm"] = communicator state = self.state_from_inputs(inputs) self.dycore = fv_dynamics.DynamicalCore( comm=communicator, grid_data=spec.grid.grid_data, stencil_factory=spec.grid.stencil_factory, damping_coefficients=spec.grid.damping_coefficients, config=spec.namelist.dynamical_core, ak=state["atmosphere_hybrid_a_coordinate"], bk=state["atmosphere_hybrid_b_coordinate"], phis=state["surface_geopotential"], ) self.dycore.step_dynamics( state, inputs["consv_te"], inputs["do_adiabatic_init"], inputs["bdt"], inputs["ptop"], inputs["n_split"], inputs["ks"], ) outputs = self.outputs_from_state(state) for name, value in outputs.items(): outputs[name] = self.subset_output(name, value) return outputs
def __init__( self, stencil_factory: StencilFactory, damping_coefficients: DampingCoefficients, rarea, nord: FloatFieldK, damp_c: FloatFieldK, ): """ nord sets the order of damping to apply: nord = 0: del-2 nord = 1: del-4 nord = 2: del-6 nord and damp_c define the damping coefficient used in DelnFluxNoSG """ self._no_compute = False if (damp_c <= 1e-4).all(): self._no_compute = True elif (damp_c[:-1] <= 1e-4).any(): raise NotImplementedError( "damp_c currently must be always greater than 10^-4 for delnflux" ) grid_indexing = stencil_factory.grid_indexing nk = grid_indexing.domain[2] self._origin = grid_indexing.origin_full() shape = grid_indexing.max_shape k_shape = (1, 1, nk) self._damp_3d = utils.make_storage_from_shape(k_shape) # fields must be 3d to assign to them self._fx2 = utils.make_storage_from_shape(shape) self._fy2 = utils.make_storage_from_shape(shape) self._d2 = utils.make_storage_from_shape(grid_indexing.domain_full()) damping_factor_calculation = stencil_factory.from_origin_domain( calc_damp, origin=(0, 0, 0), domain=k_shape ) self._add_diffusive_stencil = stencil_factory.from_dims_halo( func=add_diffusive_component, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], ) self._diffusive_damp_stencil = stencil_factory.from_dims_halo( func=diffusive_damp, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM] ) damping_factor_calculation( self._damp_3d, nord, damp_c, damping_coefficients.da_min ) self._damp = utils.make_storage_data(self._damp_3d[0, 0, :], (nk,), (0,)) self.delnflux_nosg = DelnFluxNoSG( stencil_factory, damping_coefficients, rarea, nord, nk=nk )
def compute(qvapor, qliquid, qrain, qsnow, qice, qgraupel, qcld, pt, delp, delz, peln): grid = spec.grid i_ext = grid.domain_shape_compute()[0] j_ext = grid.domain_shape_compute()[1] k_ext = grid.domain_shape_compute()[2] if spec.namelist.check_negative: raise Exception("Unimplemented namelist value check_negative=True") if spec.namelist.hydrostatic: d0_vap = constants.CP_VAP - constants.C_LIQ raise Exception("Unimplemented namelist hydrostatic=True") else: d0_vap = constants.CV_VAP - constants.C_LIQ lv00 = constants.HLV - d0_vap * constants.TICE fix_neg_water( pt, delp, delz, qvapor, qliquid, qrain, qsnow, qice, qgraupel, lv00, d0_vap, origin=grid.compute_origin(), domain=grid.domain_shape_compute(), ) fillq(qgraupel, delp, grid) fillq(qrain, delp, grid) # fix_water_vapor(delp, qvapor, origin=grid.compute_origin(), # domain=grid.domain_shape_compute()) # fix_water_vapor_nonstencil(grid, qvapor, delp) # fix_water_vapor_bottom(grid, qvapor, delp) upper_fix = utils.make_storage_from_shape(qvapor.shape, origin=(0, 0, 0)) lower_fix = utils.make_storage_from_shape(qvapor.shape, origin=(0, 0, 0)) bot_dp = delp[:, :, grid.npz - 1] full_bot_arr = utils.repeat(bot_dp[:, :, np.newaxis], k_ext + 1, axis=2) dp_bot = utils.make_storage_data(full_bot_arr) fix_water_vapor_down( qvapor, delp, upper_fix, lower_fix, dp_bot, origin=grid.compute_origin(), domain=grid.domain_shape_compute(), ) qvapor[:, :, grid.npz] = upper_fix[:, :, 0] fix_neg_cloud(delp, qcld, origin=grid.compute_origin(), domain=grid.domain_shape_compute())
def compute(self, inputs): self.make_storage_data_input_vars(inputs) inputs["index"] = utils.make_storage_data( np.arange(satadjust.QS_LENGTH), self.maxshape, origin=(0, 0, 0) ) kwargs = {"origin": (0, 0, 0), "domain": self.maxshape} self.compute_func(**inputs, **kwargs) for k, v in inputs.items(): if v.shape == self.maxshape: inputs[k] = np.squeeze(v) return inputs
def compute(self, inputs): self.make_storage_data_input_vars(inputs) index = np.arange(satadjust.QS_LENGTH) inputs["index"] = utils.make_storage_data(index, self.maxshape, origin=(0, 0, 0), read_only=False) self._compute_q_tables_stencil(**inputs) utils.device_sync() for k, v in inputs.items(): if v.shape == self.maxshape: inputs[k] = np.squeeze(v) return inputs
def make_grid_storage(self, pygrid): shape = pygrid.domain_shape_buffer_1cell() for k in TranslateGrid.composite_grid_vars: if k in self.data: self.make_composite_var_storage(k, self.data[k], shape) del self.data[k] for k, axis in TranslateGrid.edge_var_axis.items(): if k in self.data: edge_offset = pygrid.local_to_global_indices(pygrid.isd, pygrid.jsd)[ axis ] if axis == 0: edge_offset = pygrid.global_isd width = pygrid.subtile_width_x else: edge_offset = pygrid.global_jsd width = pygrid.subtile_width_y edge_slice = slice(int(edge_offset), int(edge_offset + width + 1)) self.data[k] = utils.make_storage_data( self.data[k][edge_slice], shape, start=(0, 0, pygrid.halo), axis=axis, ) for k, v in self.data.items(): if type(v) is np.ndarray: # TODO: when grid initialization model exists, may want to use # it to inform this istart, jstart = pygrid.horizontal_starts_from_shape(v.shape) logger.debug( "Storage for Grid variable {}, {}, {}, {}".format( k, istart, jstart, v.shape ) ) origin = (istart, jstart, 0) self.data[k] = utils.make_storage_data( v, shape, origin=origin, start=origin )
def compute(pt, pkz, heat_source, delz, delp, cappa, n_con, bdt): grid = spec.grid delt_column = np.ones(delz.shape[2]) * abs(bdt * spec.namelist.delt_max) delt_column[0] *= 0.1 delt_column[1] *= 0.5 delt = utils.make_storage_data( delt_column, delz.shape, origin=grid.default_origin() ) compute_pkz_tempadjust( delp, delz, cappa, heat_source, delt, pt, pkz, origin=grid.compute_origin(), domain=(grid.nic, grid.njc, n_con), )
def make_storage_data( self, array: np.ndarray, istart: int = 0, jstart: int = 0, kstart: int = 0, dummy_axes: Optional[Tuple[int, int, int]] = None, axis: int = 2, names_4d: Optional[List[str]] = None, read_only: bool = False, full_shape: bool = False, ) -> Dict[str, "Field"]: use_shape = list(self.maxshape) if dummy_axes: for axis in dummy_axes: use_shape[axis] = 1 elif not full_shape and len( array.shape) < 3 and axis == len(array.shape) - 1: use_shape[1] = 1 start = (istart, jstart, kstart) if names_4d: return utils.make_storage_dict( array, use_shape, start=start, origin=start, dummy=dummy_axes, axis=axis, names=names_4d, ) else: return utils.make_storage_data( array, use_shape, start=start, origin=start, dummy=dummy_axes, axis=axis, read_only=read_only, )
def __init__( self, stencil_factory: StencilFactory, grid_data: GridData, damping_coefficients: DampingCoefficients, column_namelist, nested: bool, stretched_grid: bool, config: DGridShallowWaterLagrangianDynamicsConfig, ): self._f0 = spec.grid.f0 self.grid = grid_data self.grid_indexing = stencil_factory.grid_indexing assert config.grid_type < 3, "ubke and vbke only implemented for grid_type < 3" assert not config.inline_q, "inline_q not yet implemented" assert ( config.d_ext <= 0 ), "untested d_ext > 0. need to call a2b_ord2, not yet implemented" assert (column_namelist["damp_vt"] > dcon_threshold).all() # TODO: in theory, we should check if damp_vt > 1e-5 for each k-level and # only compute delnflux for k-levels where this is true assert (column_namelist["damp_w"] > dcon_threshold).all() # TODO: in theory, we should check if damp_w > 1e-5 for each k-level and # only compute delnflux for k-levels where this is true # only compute for k-levels where this is true self.hydrostatic = config.hydrostatic self._tmp_heat_s = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._vort_x_delta = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._vort_y_delta = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_ke = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_vort = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._uc_contra = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._vc_contra = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_ut = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_vt = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_fx = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_fy = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_gx = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_gy = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_dw = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_wk = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_fx2 = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_fy2 = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._tmp_damp_3d = utils.make_storage_from_shape( (1, 1, self.grid_indexing.domain[2]) ) self._advected_u = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._advected_v = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._ub_contra = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._vb_contra = utils.make_storage_from_shape(self.grid_indexing.max_shape) self._column_namelist = column_namelist self.delnflux_nosg_w = DelnFluxNoSG( stencil_factory, damping_coefficients, grid_data.rarea, self._column_namelist["nord_w"], ) self.delnflux_nosg_v = DelnFluxNoSG( stencil_factory, damping_coefficients, grid_data.rarea, self._column_namelist["nord_v"], ) self.fvtp2d_dp = FiniteVolumeTransport( stencil_factory=stencil_factory, grid_data=grid_data, damping_coefficients=damping_coefficients, grid_type=config.grid_type, hord=config.hord_dp, nord=self._column_namelist["nord_v"], damp_c=self._column_namelist["damp_vt"], ) self.fvtp2d_dp_t = FiniteVolumeTransport( stencil_factory=stencil_factory, grid_data=grid_data, damping_coefficients=damping_coefficients, grid_type=config.grid_type, hord=config.hord_dp, nord=self._column_namelist["nord_t"], damp_c=self._column_namelist["damp_t"], ) self.fvtp2d_tm = FiniteVolumeTransport( stencil_factory=stencil_factory, grid_data=grid_data, damping_coefficients=damping_coefficients, grid_type=config.grid_type, hord=config.hord_tm, nord=self._column_namelist["nord_v"], damp_c=self._column_namelist["damp_vt"], ) self.fvtp2d_vt_nodelnflux = FiniteVolumeTransport( stencil_factory=stencil_factory, grid_data=grid_data, damping_coefficients=damping_coefficients, grid_type=config.grid_type, hord=config.hord_vt, ) self.fv_prep = FiniteVolumeFluxPrep( stencil_factory=stencil_factory, grid_data=grid_data, ) self.divergence_damping = DivergenceDamping( stencil_factory, grid_data, damping_coefficients, nested, stretched_grid, config.dddmp, config.d4_bg, config.nord, config.grid_type, column_namelist["nord"], column_namelist["d2_divg"], ) self._apply_pt_delp_fluxes = stencil_factory.from_dims_halo( func=apply_pt_delp_fluxes_stencil_defn, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], externals={ "inline_q": config.inline_q, }, ) self._kinetic_energy_update_part_1 = stencil_factory.from_dims_halo( func=kinetic_energy_update_part_1, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], externals={ "iord": config.hord_mt, "jord": config.hord_mt, "mord": config.hord_mt, "xt_minmax": False, "yt_minmax": False, }, skip_passes=("GreedyMerging",), ) self._kinetic_energy_update_part_2 = stencil_factory.from_dims_halo( func=kinetic_energy_update_part_2, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], externals={ "iord": config.hord_mt, "jord": config.hord_mt, "mord": config.hord_mt, "xt_minmax": False, "yt_minmax": False, }, skip_passes=("GreedyMerging",), ) self._flux_adjust_stencil = stencil_factory.from_dims_halo( func=flux_adjust, compute_dims=[X_DIM, Y_DIM, Z_DIM] ) self._flux_capacitor_stencil = stencil_factory.from_dims_halo( func=flux_capacitor, compute_dims=[X_DIM, Y_DIM, Z_DIM], compute_halos=(3, 3), ) self._vort_differencing_stencil = stencil_factory.from_dims_halo( func=vort_differencing, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], ) self._u_and_v_from_ke_stencil = stencil_factory.from_dims_halo( func=u_and_v_from_ke, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM] ) self._compute_vort_stencil = stencil_factory.from_dims_halo( func=compute_vort, externals={ "radius": constants.RADIUS, "do_f3d": config.do_f3d, "hydrostatic": self.hydrostatic, }, compute_dims=[X_DIM, Y_DIM, Z_DIM], compute_halos=(3, 3), ) self._adjust_w_and_qcon_stencil = stencil_factory.from_dims_halo( func=adjust_w_and_qcon, compute_dims=[X_DIM, Y_DIM, Z_DIM], ) self._heat_diss_stencil = stencil_factory.from_dims_halo( func=heat_diss, compute_dims=[X_DIM, Y_DIM, Z_DIM], ) self._heat_source_from_vorticity_damping_stencil = ( stencil_factory.from_dims_halo( func=heat_source_from_vorticity_damping, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], externals={ "do_skeb": config.do_skeb, "d_con": config.d_con, }, ) ) self._compute_vorticity_stencil = stencil_factory.from_dims_halo( compute_vorticity, compute_dims=[X_DIM, Y_DIM, Z_DIM], compute_halos=(3, 3), ) self._update_u_and_v_stencil = stencil_factory.from_dims_halo( update_u_and_v, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM] ) damping_factor_calculation_stencil = stencil_factory.from_origin_domain( delnflux.calc_damp, origin=(0, 0, 0), domain=(1, 1, stencil_factory.grid_indexing.domain[2]), ) damping_factor_calculation_stencil( self._tmp_damp_3d, self._column_namelist["nord_v"], self._column_namelist["damp_vt"], damping_coefficients.da_min_c, ) self._delnflux_damp_vt = utils.make_storage_data( self._tmp_damp_3d[0, 0, :], (self.grid_indexing.domain[2],), (0,) ) damping_factor_calculation_stencil( self._tmp_damp_3d, self._column_namelist["nord_w"], self._column_namelist["damp_w"], damping_coefficients.da_min_c, ) self._delnflux_damp_w = utils.make_storage_data( self._tmp_damp_3d[0, 0, :], (self.grid_indexing.domain[2],), (0,) ) y_temporary = utils.make_storage_from_shape( shape=self.grid_indexing.max_shape, origin=self.grid_indexing.origin ) self._copy_corners = PreAllocatedCopiedCornersFactory( stencil_factory, dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], y_temporary=y_temporary, )
def __init__( self, comm: fv3gfs.util.CubedSphereCommunicator, stencil_factory: StencilFactory, grid_data: GridData, damping_coefficients: DampingCoefficients, grid_type, nested, stretched_grid, config: AcousticDynamicsConfig, # TODO: move ak, bk, pfull, and phis into GridData ak: FloatFieldK, bk: FloatFieldK, pfull: FloatFieldK, phis: FloatFieldIJ, ): """ Args: comm: object for cubed sphere inter-process communication stencil_factory: creates stencils grid_data: metric terms defining the grid damping_coefficients: damping configuration grid_type: ??? nested: ??? stretched_grid: ??? config: configuration settings ak: atmosphere hybrid a coordinate (Pa) bk: atmosphere hybrid b coordinate (dimensionless) pfull: atmospheric Eulerian grid reference pressure (Pa) phis: surface geopotential height """ grid_indexing = stencil_factory.grid_indexing self.comm = comm self.config = config assert config.d_ext == 0, "d_ext != 0 is not implemented" assert config.beta == 0, "beta != 0 is not implemented" assert not config.use_logp, "use_logp=True is not implemented" self._da_min = damping_coefficients.da_min self.grid_data = grid_data self._pfull = pfull self._nk_heat_dissipation = get_nk_heat_dissipation( config.d_grid_shallow_water, npz=grid_indexing.domain[2], ) self.nonhydrostatic_pressure_gradient = ( nh_p_grad.NonHydrostaticPressureGradient(stencil_factory, grid_data, config.grid_type)) self._temporaries = dyncore_temporaries(grid_indexing) self._temporaries["gz"][:] = HUGE_R if not config.hydrostatic: self._temporaries["pk3"][:] = HUGE_R column_namelist = d_sw.get_column_namelist(config.d_grid_shallow_water, grid_indexing.domain[2]) if not config.hydrostatic: # To write lower dimensional storages, these need to be 3D # then converted to lower dimensional dp_ref_3d = utils.make_storage_from_shape(grid_indexing.max_shape) zs_3d = utils.make_storage_from_shape(grid_indexing.max_shape) dp_ref_stencil = stencil_factory.from_origin_domain( dp_ref_compute, origin=grid_indexing.origin_full(), domain=grid_indexing.domain_full(add=(0, 0, 1)), ) dp_ref_stencil( ak, bk, phis, dp_ref_3d, zs_3d, 1.0 / constants.GRAV, ) # After writing, make 'dp_ref' a K-field and 'zs' an IJ-field self._dp_ref = utils.make_storage_data(dp_ref_3d[0, 0, :], (dp_ref_3d.shape[2], ), (0, )) self._zs = utils.make_storage_data(zs_3d[:, :, 0], zs_3d.shape[0:2], (0, 0)) self.update_height_on_d_grid = updatedzd.UpdateHeightOnDGrid( stencil_factory, damping_coefficients, grid_data, grid_type, config.hord_tm, self._dp_ref, column_namelist, d_sw.k_bounds(), ) self.riem_solver3 = RiemannSolver3(stencil_factory, config.riemann) self.riem_solver_c = RiemannSolverC(stencil_factory, p_fac=config.p_fac) origin, domain = grid_indexing.get_origin_domain( [X_DIM, Y_DIM, Z_INTERFACE_DIM], halos=(2, 2)) self._compute_geopotential_stencil = stencil_factory.from_origin_domain( compute_geopotential, origin=origin, domain=domain, ) self.dgrid_shallow_water_lagrangian_dynamics = ( d_sw.DGridShallowWaterLagrangianDynamics( stencil_factory, grid_data, damping_coefficients, column_namelist, nested, stretched_grid, config.d_grid_shallow_water, )) self.cgrid_shallow_water_lagrangian_dynamics = CGridShallowWaterDynamics( stencil_factory, grid_data, nested, config.grid_type, config.nord, ) self._set_gz = stencil_factory.from_origin_domain( set_gz, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(add=(0, 0, 1)), ) self._set_pem = stencil_factory.from_origin_domain( set_pem, origin=grid_indexing.origin_compute(add=(-1, -1, 0)), domain=grid_indexing.domain_compute(add=(2, 2, 0)), ) self._p_grad_c = stencil_factory.from_origin_domain( p_grad_c_stencil, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(add=(1, 1, 0)), externals={"hydrostatic": config.hydrostatic}, ) self.update_geopotential_height_on_c_grid = ( updatedzc.UpdateGeopotentialHeightOnCGrid(stencil_factory, grid_data.area)) self._zero_data = stencil_factory.from_origin_domain( zero_data, origin=grid_indexing.origin_full(), domain=grid_indexing.domain_full(), ) ax_offsets_pe = axis_offsets( grid_indexing, grid_indexing.origin_full(), grid_indexing.domain_full(add=(0, 0, 1)), ) self._edge_pe_stencil = stencil_factory.from_origin_domain( pe_halo.edge_pe, origin=grid_indexing.origin_full(), domain=grid_indexing.domain_full(add=(0, 0, 1)), externals={**ax_offsets_pe}, ) """The stencil object responsible for updating the interface pressure""" self._do_del2cubed = self._nk_heat_dissipation != 0 and config.d_con > 1.0e-5 if self._do_del2cubed: nf_ke = min(3, config.nord + 1) self._hyperdiffusion = HyperdiffusionDamping(stencil_factory, damping_coefficients, grid_data.rarea, nmax=nf_ke) if config.rf_fast: self._rayleigh_damping = ray_fast.RayleighDamping( stencil_factory, rf_cutoff=config.rf_cutoff, tau=config.tau, hydrostatic=config.hydrostatic, ) self._compute_pkz_tempadjust = stencil_factory.from_origin_domain( temperature_adjust.compute_pkz_tempadjust, origin=grid_indexing.origin_compute(), domain=grid_indexing.restrict_vertical( nk=self._nk_heat_dissipation).domain_compute(), ) self._pk3_halo = PK3Halo(stencil_factory) self._copy_stencil = stencil_factory.from_origin_domain( basic.copy_defn, origin=grid_indexing.origin_full(), domain=grid_indexing.domain_full(add=(0, 0, 1)), ) # Halo updaters self._halo_updaters = AcousticDynamics._HaloUpdaters( self.comm, grid_indexing)
def compute(state, comm): # u, v, w, delz, delp, pt, pe, pk, phis, wsd, omga, ua, va, uc, vc, mfxd, # mfyd, cxd, cyd, pkz, peln, q_con, ak, bk, diss_estd, cappa, mdt, n_split, # akap, ptop, pfull, n_map, comm): grid = spec.grid init_step = state.n_map == 1 end_step = state.n_map == spec.namelist.k_split akap = state.akap # peln1 = math.log(ptop) # ptk = ptop**akap dt = state.mdt / state.n_split dt2 = 0.5 * dt hydrostatic = spec.namelist.hydrostatic rgrav = 1.0 / constants.GRAV n_split = state.n_split # TODO: Put defaults into code. # m_split = 1. + abs(dt_atmos)/real(k_split*n_split*abs(p_split)) # n_split = nint( real(n0split)/real(k_split*abs(p_split)) * stretch_fac + 0.5 ) ms = max(1, spec.namelist.m_split / 2.0) shape = state.delz.shape # NOTE: In Fortran model the halo update starts happens in fv_dynamics, not here. reqs = {} for halovar in [ "q_con_quantity", "cappa_quantity", "delp_quantity", "pt_quantity" ]: reqs[halovar] = comm.start_halo_update(state.__getattribute__(halovar), n_points=utils.halo) reqs_vector = comm.start_vector_halo_update(state.u_quantity, state.v_quantity, n_points=utils.halo) reqs["q_con_quantity"].wait() reqs["cappa_quantity"].wait() state.__dict__.update(dyncore_temporaries(shape)) if init_step: state.gz[:-1, :-1, :] = HUGE_R state.diss_estd[grid.slice_dict(grid.compute_dict())] = 0.0 if not hydrostatic: state.pk3[:-1, :-1, :] = HUGE_R state.mfxd[grid.slice_dict(grid.x3d_compute_dict())] = 0.0 state.mfyd[grid.slice_dict(grid.y3d_compute_dict())] = 0.0 state.cxd[grid.slice_dict(grid.x3d_compute_domain_y_dict())] = 0.0 state.cyd[grid.slice_dict(grid.y3d_compute_domain_x_dict())] = 0.0 if not hydrostatic: # k1k = akap / (1.0 - akap) # TODO: Is really just a column... when different shapes are supported # perhaps change this. state.dp_ref = utils.make_storage_from_shape(state.ak.shape, grid.default_origin()) state.zs = utils.make_storage_from_shape(state.ak.shape, grid.default_origin()) dp_ref_compute( state.ak, state.bk, state.phis, state.dp_ref, state.zs, rgrav, origin=grid.default_origin(), domain=grid.domain_shape_standard(add=(0, 0, 1)), ) n_con = get_n_con() for it in range(n_split): remap_step = False if spec.namelist.breed_vortex_inline or (it == n_split - 1): remap_step = True if not hydrostatic: reqs["w_quantity"] = comm.start_halo_update(state.w_quantity, n_points=utils.halo) if it == 0: set_gz( state.zs, state.delz, state.gz, origin=grid.compute_origin(), domain=(grid.nic, grid.njc, grid.npz + 1), ) reqs["gz_quantity"] = comm.start_halo_update( state.gz_quantity, n_points=utils.halo) if it == 0: reqs["delp_quantity"].wait() reqs["pt_quantity"].wait() beta_d = 0 else: beta_d = spec.namelist.beta last_step = False if it == n_split - 1 and end_step: last_step = True if it == n_split - 1 and end_step: if spec.namelist.use_old_omega: # apparently True set_pem( state.delp, state.pem, state.ptop, origin=(grid.is_ - 1, grid.js - 1, 0), domain=(grid.nic + 2, grid.njc + 2, grid.npz), ) reqs_vector.wait() if not hydrostatic: reqs["w_quantity"].wait() state.delpc, state.ptc = c_sw.compute( state.delp, state.pt, state.u, state.v, state.w, state.uc, state.vc, state.ua, state.va, state.ut, state.vt, state.divgd, state.omga, dt2, ) if spec.namelist.nord > 0: reqs["divgd_quantity"] = comm.start_halo_update( state.divgd_quantity, n_points=utils.halo) if not hydrostatic: if it == 0: reqs["gz_quantity"].wait() copy_stencil( state.gz, state.zh, origin=grid.default_origin(), domain=grid.domain_shape_buffer_k(), ) else: copy_stencil( state.gz, state.zh, origin=grid.default_origin(), domain=grid.domain_shape_buffer_k(), ) if not hydrostatic: state.gz, state.ws3 = updatedzc.compute(state.dp_ref, state.zs, state.ut, state.vt, state.gz, state.ws3, dt2) # TODO: This is really a 2d field. state.ws3 = utils.make_storage_data(state.ws3[:, :, -1], shape, origin=(0, 0, 0)) riem_solver_c.compute( ms, dt2, akap, state.cappa, state.ptop, state.phis, state.omga, state.ptc, state.q_con, state.delpc, state.gz, state.pkc, state.ws3, ) pgradc.compute(state.uc, state.vc, state.delpc, state.pkc, state.gz, dt2) reqc_vector = comm.start_vector_halo_update(state.uc_quantity, state.vc_quantity, n_points=utils.halo) if spec.namelist.nord > 0: reqs["divgd_quantity"].wait() reqc_vector.wait() state.nord_v, state.damp_vt = d_sw.compute( state.vt, state.delp, state.ptc, state.pt, state.u, state.v, state.w, state.uc, state.vc, state.ua, state.va, state.divgd, state.mfxd, state.mfyd, state.cxd, state.cyd, state.crx, state.cry, state.xfx, state.yfx, state.q_con, state.zh, state.heat_source, state.diss_estd, dt, ) for halovar in ["delp_quantity", "pt_quantity", "q_con_quantity"]: comm.halo_update(state.__getattribute__(halovar), n_points=utils.halo) # Not used unless we implement other betas and alternatives to nh_p_grad # if spec.namelist.d_ext > 0: # raise 'Unimplemented namelist option d_ext > 0' # else: # divg2 = utils.make_storage_from_shape(delz.shape, grid.compute_origin()) if not hydrostatic: updatedzd.compute( state.nord_v, state.damp_vt, state.dp_ref, state.zs, state.zh, state.crx, state.cry, state.xfx, state.yfx, state.wsd, dt, ) # TODO: This is really a 2d field. state.wsd = utils.make_storage_data(state.wsd[:, :, -1], shape, origin=grid.compute_origin()) riem_solver3.compute( remap_step, dt, akap, state.cappa, state.ptop, state.zs, state.w, state.delz, state.q_con, state.delp, state.pt, state.zh, state.pe, state.pkc, state.pk3, state.pk, state.peln, state.wsd, ) reqs["zh_quantity"] = comm.start_halo_update(state.zh_quantity, n_points=utils.halo) if grid.npx == grid.npy: reqs["pkc_quantity"] = comm.start_halo_update( state.pkc_quantity, n_points=2) else: reqs["pkc_quantity"] = comm.start_halo_update( state.pkc_quantity, n_points=utils.halo) if remap_step: pe_halo.compute(state.pe, state.delp, state.ptop) if spec.namelist.use_logp: raise Exception("unimplemented namelist option use_logp=True") else: pk3_halo.compute(state.pk3, state.delp, state.ptop, akap) if not hydrostatic: reqs["zh_quantity"].wait() if grid.npx != grid.npy: reqs["pkc_quantity"].wait() if not hydrostatic: basic.multiply_constant( state.zh, state.gz, constants.GRAV, origin=(grid.is_ - 2, grid.js - 2, 0), domain=(grid.nic + 4, grid.njc + 4, grid.npz + 1), ) if grid.npx == grid.npy: reqs["pkc_quantity"].wait() if spec.namelist.beta != 0: raise Exception( "Unimplemented namelist option -- we only support beta=0") if not hydrostatic: nh_p_grad.compute( state.u, state.v, state.pkc, state.gz, state.pk3, state.delp, dt, state.ptop, akap, ) if spec.namelist.rf_fast: # TODO: Pass through ks, or remove, inconsistent representation vs Fortran. ray_fast.compute( state.u, state.v, state.w, state.dp_ref, state.pfull, dt, state.ptop, state.ks, ) if it != n_split - 1: reqs_vector = comm.start_vector_halo_update(state.u_quantity, state.v_quantity, n_points=utils.halo) else: if spec.namelist.grid_type < 4: comm.synchronize_vector_interfaces(state.u_quantity, state.v_quantity) if n_con != 0 and spec.namelist.d_con > 1.0e-5: nf_ke = min(3, spec.namelist.nord + 1) comm.halo_update(state.heat_source_quantity, n_points=utils.halo) cd = constants.CNST_0P20 * grid.da_min del2cubed.compute(state.heat_source, nf_ke, cd, grid.npz) if not hydrostatic: temperature_adjust.compute( state.pt, state.pkz, state.heat_source, state.delz, state.delp, state.cappa, n_con, dt, )
def make_composite_var_storage(self, varname, data3d, shape): for s in range(9): self.data[varname + str(s + 1)] = utils.make_storage_data( np.squeeze(data3d[:, :, s]), shape, origin=(0, 0, 0))