def __init__( self, stencil_factory: StencilFactory, grid_data: GridData, ): grid_indexing = stencil_factory.grid_indexing self._dx = grid_data.dx self._dy = grid_data.dy self._rdxa = grid_data.rdxa self._rdya = grid_data.rdya self._cosa_u = grid_data.cosa_u self._cosa_v = grid_data.cosa_v self._rsin_u = grid_data.rsin_u self._rsin_v = grid_data.rsin_v self._sin_sg1 = grid_data.sin_sg1 self._sin_sg2 = grid_data.sin_sg2 self._sin_sg3 = grid_data.sin_sg3 self._sin_sg4 = grid_data.sin_sg4 origin = grid_indexing.origin_full() domain = grid_indexing.domain_full() ax_offsets = axis_offsets(grid_indexing, origin, domain) kwargs = {"externals": ax_offsets, "origin": origin, "domain": domain} origin_corners = grid_indexing.origin_full(add=(1, 1, 0)) domain_corners = grid_indexing.domain_full(add=(-1, -1, 0)) corner_offsets = axis_offsets(grid_indexing, origin_corners, domain_corners) kwargs_corners = { "externals": corner_offsets, "origin": origin_corners, "domain": domain_corners, } self._main_uc_contra_stencil = stencil_factory.from_origin_domain( main_uc_contra, **kwargs) self._main_vc_contra_stencil = stencil_factory.from_origin_domain( main_vc_contra, **kwargs) self._uc_contra_y_edge_stencil = stencil_factory.from_origin_domain( uc_contra_y_edge, **kwargs) self._vc_contra_y_edge_stencil = stencil_factory.from_origin_domain( vc_contra_y_edge, **kwargs) self._uc_contra_x_edge_stencil = stencil_factory.from_origin_domain( uc_contra_x_edge, **kwargs) self._vc_contra_x_edge_stencil = stencil_factory.from_origin_domain( vc_contra_x_edge, **kwargs) self._uc_contra_corners_stencil = stencil_factory.from_origin_domain( uc_contra_corners, **kwargs_corners) self._vc_contra_corners_stencil = stencil_factory.from_origin_domain( vc_contra_corners, **kwargs_corners) self._fxadv_fluxes_stencil = stencil_factory.from_origin_domain( fxadv_fluxes_stencil, **kwargs)
def __init__( self, stencil_factory: StencilFactory, grid_data: GridData, grid_type: int, jord: int, ): if jord not in (5, 6, 7, 8): raise NotImplementedError( "Currently ytp_v is only supported for hord_mt == 5,6,7,8") assert grid_type < 3 grid_indexing = stencil_factory.grid_indexing origin = grid_indexing.origin_compute() domain = grid_indexing.domain_compute(add=(1, 1, 0)) self._dy = grid_data.dy self._dya = grid_data.dya self._rdy = grid_data.rdy ax_offsets = axis_offsets(grid_indexing, origin, domain) self.stencil = stencil_factory.from_origin_domain( ytp_v_stencil_defn, externals={ "jord": jord, "mord": jord, "yt_minmax": False, **ax_offsets, }, origin=origin, domain=domain, skip_passes=("GreedyMerging", ), )
def __init__( self, stencil_factory: StencilFactory, dims: Sequence[str], y_field, ) -> None: """ Args: stencil_factory: creates stencils dims: dimensionality of the data to be copied y_field: 3D gt4py storage to use for y-differenceable field (x-differenceable field uses same memory as base field) """ grid_indexing = stencil_factory.grid_indexing origin, domain = grid_indexing.get_origin_domain( dims=dims, halos=(grid_indexing.n_halo, grid_indexing.n_halo)) self._y_field = y_field ax_offsets = axis_offsets(grid_indexing, origin, domain) self._copy_corners_xy = stencil_factory.from_origin_domain( func=copy_corners_xy_stencil_defn, origin=origin, domain=domain, externals={ **ax_offsets, }, )
def __init__(self, direction: str, stencil_factory: StencilFactory) -> None: """The grid for this stencil""" grid_indexing = stencil_factory.grid_indexing n_halo = grid_indexing.n_halo origin, domain = grid_indexing.get_origin_domain( dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], halos=(n_halo, n_halo)) ax_offsets = axis_offsets(grid_indexing, origin, domain) if direction == "x": self._copy_corners = stencil_factory.from_origin_domain( func=copy_corners_x_stencil_defn, origin=origin, domain=domain, externals={ **ax_offsets, }, ) elif direction == "y": self._copy_corners = stencil_factory.from_origin_domain( func=copy_corners_y_stencil_defn, origin=origin, domain=domain, externals={ **ax_offsets, }, ) else: raise ValueError("Direction must be either 'x' or 'y'")
def __init__(self, grid): super().__init__(grid) self.in_vars["data_vars"] = { "pe": { "istart": grid.is_ - 1, "iend": grid.ie + 1, "jstart": grid.js - 1, "jend": grid.je + 1, "kend": grid.npz + 1, "kaxis": 1, }, "delp": {}, } self.in_vars["parameters"] = ["ptop"] self.out_vars = {"pe": self.in_vars["data_vars"]["pe"]} ax_offsets_pe = axis_offsets( self.grid.grid_indexing, self.grid.grid_indexing.origin_full(), self.grid.grid_indexing.domain_full(add=(0, 0, 1)), ) self.compute_func = spec.grid.stencil_factory.from_origin_domain( pe_halo.edge_pe, origin=self.grid.grid_indexing.origin_full(), domain=self.grid.grid_indexing.domain_full(add=(0, 0, 1)), externals={**ax_offsets_pe}, )
def __init__(self, grid): super().__init__(grid) self.in_vars["data_vars"] = { "uc": {}, "vc": {}, "ut": {}, "ub": grid.compute_dict_buffer_2d(), } self.in_vars["parameters"] = ["dt5", "dt4"] self.out_vars = {"ub": grid.compute_dict_buffer_2d()} origin = self.grid.compute_origin() domain = self.grid.domain_shape_compute(add=(1, 1, 0)) ax_offsets = axis_offsets(self.grid, origin, domain) self.compute_func = self.grid.stencil_factory.from_origin_domain( ubke, externals=ax_offsets, origin=origin, domain=domain)
def __init__(self, stencil_factory: StencilFactory): grid_indexing = stencil_factory.grid_indexing origin = grid_indexing.origin_full() domain = grid_indexing.domain_full(add=(0, 0, 1)) ax_offsets = axis_offsets(grid_indexing, origin, domain) self._edge_pe_update = stencil_factory.from_origin_domain( func=edge_pe_update, externals={ **ax_offsets, }, origin=origin, domain=domain, ) shape_2D = grid_indexing.domain_full(add=(1, 1, 1))[0:2] self._pe_tmp = utils.make_storage_from_shape( shape_2D, grid_indexing.origin_full())
def test_axis_offsets( domain: Index3D, n_halo: int, south_edge: bool, north_edge: bool, west_edge: bool, east_edge: bool, origin_offset: Index3D, domain_offset: Index3D, i_start, i_end, j_start, j_end, ): grid = fv3core.utils.grid.GridIndexing( domain=domain, n_halo=n_halo, south_edge=south_edge, north_edge=north_edge, west_edge=west_edge, east_edge=east_edge, ) origin = (n_halo, n_halo, 0) call_origin = tuple(compute + offset for (compute, offset) in zip(origin, origin_offset)) call_domain = tuple(compute + offset for (compute, offset) in zip(domain, domain_offset)) axis_offsets = grid.axis_offsets(call_origin, call_domain) if west_edge: assert axis_offsets["i_start"] == i_start else: assert axis_offsets["i_start"] == gtscript.I[0] - np.iinfo( np.int16).max if east_edge: assert axis_offsets["i_end"] == i_end else: assert axis_offsets["i_end"] == gtscript.I[-1] + np.iinfo(np.int16).max if south_edge: assert axis_offsets["j_start"] == j_start else: assert axis_offsets["j_start"] == gtscript.J[0] - np.iinfo( np.int16).max if north_edge: assert axis_offsets["j_end"] == j_end else: assert axis_offsets["j_end"] == gtscript.J[-1] + np.iinfo(np.int16).max
def __init__( self, direction: str, stencil_factory: StencilFactory, temporary_field=None, origin=None, domain=None, ) -> None: n_halo = stencil_factory.grid_indexing.n_halo ( default_origin, default_domain, ) = stencil_factory.grid_indexing.get_origin_domain( dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_INTERFACE_DIM], halos=(n_halo, n_halo), ) if origin is None: origin = default_origin """The origin for the corner computation""" if domain is None: domain = default_domain """The full domain required to do corner computation everywhere""" if temporary_field is not None: self._corner_tmp = temporary_field else: self._corner_tmp = utils.make_storage_from_shape( stencil_factory.grid_indexing.max_shape, origin=origin) if direction == "x": defn = fill_corners_bgrid_x_defn elif direction == "y": defn = fill_corners_bgrid_y_defn else: raise ValueError("Direction must be either 'x' or 'y'") externals = axis_offsets(stencil_factory.grid_indexing, origin=origin, domain=domain) self._fill_corners_bgrid = stencil_factory.from_origin_domain( func=defn, origin=origin, domain=domain, externals=externals)
def compute(self, inputs): self.make_storage_data_input_vars(inputs) origin = self.grid.full_origin() domain = self.grid.domain_shape_full() axes_offsets = axis_offsets(self.grid, origin, domain) if inputs["dir"] == 1: stencil = self.grid.stencil_factory.from_origin_domain( corners.fill_corners_2cells_x_stencil, externals=axes_offsets, origin=origin, domain=domain, ) elif inputs["dir"] == 2: stencil = self.grid.stencil_factory.from_origin_domain( corners.fill_corners_2cells_y_stencil, externals=axes_offsets, origin=origin, domain=domain, ) stencil(inputs["q4c"], inputs["q4c"]) return self.slice_output(inputs, {"q4c": inputs["q4c"]})
def compute(self, inputs): nord_column = inputs["nord_col"][:] self.make_storage_data_input_vars(inputs) for nord in utils.unique(nord_column): if nord != 0: ki = [ k for k in range(self.grid.npz) if nord_column[0, 0, k] == nord ] origin = (self.grid.isd, self.grid.jsd, ki[0]) domain = (self.grid.nid + 1, self.grid.njd + 1, len(ki)) axes_offsets = axis_offsets(self.grid.grid_indexing, origin, domain) vector_corner_fill = self.grid.stencil_factory.from_origin_domain( corners.fill_corners_dgrid_defn, externals=axes_offsets, origin=origin, domain=domain, ) vector_corner_fill(inputs["vc"], inputs["vc"], inputs["uc"], inputs["uc"], -1.0) return self.slice_output(inputs)
def __init__(self, stencil_factory: StencilFactory, area): grid_indexing = stencil_factory.grid_indexing self._area = area largest_possible_shape = grid_indexing.domain_full(add=(1, 1, 1)) self._gz_x = gt4py_utils.make_storage_from_shape( largest_possible_shape, grid_indexing.origin_compute(add=(0, -grid_indexing.n_halo, 0)), ) self._gz_y = gt4py_utils.make_storage_from_shape( largest_possible_shape, grid_indexing.origin_compute(add=(0, -grid_indexing.n_halo, 0)), ) full_origin = grid_indexing.origin_full() full_domain = grid_indexing.domain_full(add=(0, 0, 1)) self._double_copy_stencil = stencil_factory.from_origin_domain( double_copy, origin=full_origin, domain=full_domain, ) ax_offsets = axis_offsets(grid_indexing, full_origin, full_domain) self._fill_corners_x_stencil = stencil_factory.from_origin_domain( corners.fill_corners_2cells_x_stencil, externals=ax_offsets, origin=full_origin, domain=full_domain, ) self._fill_corners_y_stencil = stencil_factory.from_origin_domain( corners.fill_corners_2cells_y_stencil, externals=ax_offsets, origin=full_origin, domain=full_domain, ) self._update_dz_c = stencil_factory.from_origin_domain( update_dz_c, origin=grid_indexing.origin_compute(add=(-1, -1, 0)), domain=grid_indexing.domain_compute(add=(2, 2, 1)), )
def __init__( self, stencil_factory: StencilFactory, damping_coefficients: DampingCoefficients, rarea, nord, nk: Optional[int] = None, ): """ nord sets the order of damping to apply: nord = 0: del-2 nord = 1: del-4 nord = 2: del-6 """ grid_indexing = stencil_factory.grid_indexing self._del6_u = damping_coefficients.del6_u self._del6_v = damping_coefficients.del6_v self._rarea = rarea if max(nord[:]) > 3: raise ValueError("nord must be less than 3") if not np.all(n in [0, 2, 3] for n in nord[:]): raise NotImplementedError("nord must have values 0, 2, or 3") self._nmax = int(max(nord[:])) i1 = grid_indexing.isc - 1 - self._nmax i2 = grid_indexing.iec + 1 + self._nmax j1 = grid_indexing.jsc - 1 - self._nmax j2 = grid_indexing.jec + 1 + self._nmax if nk is None: nk = grid_indexing.domain[2] nk = nk origin_d2 = (i1, j1, 0) domain_d2 = (i2 - i1 + 1, j2 - j1 + 1, nk) f1_ny = grid_indexing.jec - grid_indexing.jsc + 1 + 2 * self._nmax f1_nx = grid_indexing.iec - grid_indexing.isc + 2 + 2 * self._nmax fx_origin = (grid_indexing.isc - self._nmax, grid_indexing.jsc - self._nmax, 0) self._nord = nord if nk <= 3: raise NotImplementedError("nk must be more than 3 for DelnFluxNoSG") self._k_bounds = [1, 1, 1, nk - 3] preamble_ax_offsets = axis_offsets(grid_indexing, origin_d2, domain_d2) fx_ax_offsets = axis_offsets(grid_indexing, fx_origin, (f1_nx, f1_ny, nk)) fy_ax_offsets = axis_offsets( grid_indexing, fx_origin, (f1_nx - 1, f1_ny + 1, nk) ) origins_d2 = [] domains_d2 = [] origins_flux = [] domains_fx = [] domains_fy = [] for n in range(self._nmax): nt = self._nmax - 1 - n nt_ny = grid_indexing.jec - grid_indexing.jsc + 3 + 2 * nt nt_nx = grid_indexing.iec - grid_indexing.isc + 3 + 2 * nt origins_d2.append( (grid_indexing.isc - nt - 1, grid_indexing.jsc - nt - 1, 0) ) domains_d2.append((nt_nx, nt_ny, nk)) origins_flux.append((grid_indexing.isc - nt, grid_indexing.jsc - nt, 0)) domains_fx.append((nt_nx - 1, nt_ny - 2, nk)) domains_fy.append((nt_nx - 2, nt_ny - 1, nk)) nord_dictionary = { "nord0": nord[0], "nord1": nord[1], "nord2": nord[2], "nord3": nord[3], } self._d2_damp = stencil_factory.from_origin_domain( d2_damp_interval, externals={ **nord_dictionary, **preamble_ax_offsets, }, origin=origin_d2, domain=domain_d2, ) self._copy_stencil_interval = stencil_factory.from_origin_domain( copy_stencil_interval, externals={ **nord_dictionary, **preamble_ax_offsets, }, origin=origin_d2, domain=domain_d2, ) self._d2_stencil = get_stencils_with_varied_bounds( d2_highorder_stencil, origins_d2, domains_d2, stencil_factory=stencil_factory, externals={**nord_dictionary}, ) self._column_conditional_fx_calculation = get_stencils_with_varied_bounds( fx_calc_stencil_column, origins_flux, domains_fx, stencil_factory=stencil_factory, externals={**nord_dictionary}, ) self._column_conditional_fy_calculation = get_stencils_with_varied_bounds( fy_calc_stencil_column, origins_flux, domains_fy, stencil_factory=stencil_factory, externals={**nord_dictionary}, ) self._fx_calc_stencil = stencil_factory.from_origin_domain( fx_calc_stencil_nord, externals={**fx_ax_offsets, **nord_dictionary}, origin=fx_origin, domain=(f1_nx, f1_ny, nk), ) self._fy_calc_stencil = stencil_factory.from_origin_domain( fy_calc_stencil_nord, externals={**fy_ax_offsets, **nord_dictionary}, origin=fx_origin, domain=(f1_nx - 1, f1_ny + 1, nk), ) corner_origin, corner_domain = grid_indexing.get_origin_domain( dims=[X_DIM, Y_DIM, Z_DIM], halos=(grid_indexing.n_halo, grid_indexing.n_halo), ) corner_domain = corner_domain[:2] + (nk,) corner_axis_offsets = axis_offsets(grid_indexing, corner_origin, corner_domain) self._corner_tmp = utils.make_storage_from_shape( corner_domain, origin=corner_origin ) self._copy_corners_x_nord = stencil_factory.from_origin_domain( copy_corners_x_nord, externals={**corner_axis_offsets, **nord_dictionary}, origin=corner_origin, domain=corner_domain, ) self._copy_corners_y_nord = stencil_factory.from_origin_domain( copy_corners_y_nord, externals={**corner_axis_offsets, **nord_dictionary}, origin=corner_origin, domain=corner_domain, )
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 __init__( self, stencil_factory: StencilFactory, config: RemappingConfig, area_64, nq, pfull, ): grid_indexing = stencil_factory.grid_indexing if config.kord_tm >= 0: raise NotImplementedError( "map ppm, untested mode where kord_tm >= 0") hydrostatic = config.hydrostatic if hydrostatic: raise NotImplementedError("Hydrostatic is not implemented") shape_kplus = grid_indexing.domain_full(add=(0, 0, 1)) self._t_min = 184.0 self._nq = nq # do_omega = hydrostatic and last_step # TODO pull into inputs self._domain_jextra = ( grid_indexing.domain[0], grid_indexing.domain[1] + 1, grid_indexing.domain[2] + 1, ) self._pe1 = utils.make_storage_from_shape(shape_kplus) self._pe2 = utils.make_storage_from_shape(shape_kplus) self._dp2 = utils.make_storage_from_shape(shape_kplus) self._pn2 = utils.make_storage_from_shape(shape_kplus) self._pe0 = utils.make_storage_from_shape(shape_kplus) self._pe3 = utils.make_storage_from_shape(shape_kplus) self._gz: FloatField = utils.make_storage_from_shape( shape_kplus, grid_indexing.origin_compute()) self._cvm: FloatField = utils.make_storage_from_shape( shape_kplus, grid_indexing.origin_compute()) self._init_pe = stencil_factory.from_origin_domain( init_pe, origin=grid_indexing.origin_compute(), domain=self._domain_jextra) self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( moist_cv_pt_pressure, externals={ "kord_tm": config.kord_tm, "hydrostatic": hydrostatic }, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(add=(0, 0, 1)), ) self._moist_cv_pkz = stencil_factory.from_origin_domain( moist_cv.moist_pkz, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(), ) self._pn2_pk_delp = stencil_factory.from_origin_domain( pn2_pk_delp, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(), ) self._kord_tm = abs(config.kord_tm) self._map_single_pt = MapSingle( stencil_factory, self._kord_tm, 1, grid_indexing.isc, grid_indexing.iec, grid_indexing.jsc, grid_indexing.jec, ) self._mapn_tracer = MapNTracer( stencil_factory, abs(config.kord_tr), nq, grid_indexing.isc, grid_indexing.iec, grid_indexing.jsc, grid_indexing.jec, fill=config.fill, ) self._kord_wz = config.kord_wz self._map_single_w = MapSingle( stencil_factory, self._kord_wz, -2, grid_indexing.isc, grid_indexing.iec, grid_indexing.jsc, grid_indexing.jec, ) self._map_single_delz = MapSingle( stencil_factory, self._kord_wz, 1, grid_indexing.isc, grid_indexing.iec, grid_indexing.jsc, grid_indexing.jec, ) self._undo_delz_adjust_and_copy_peln = stencil_factory.from_origin_domain( undo_delz_adjust_and_copy_peln, origin=grid_indexing.origin_compute(), domain=( grid_indexing.domain[0], grid_indexing.domain[1], grid_indexing.domain[2] + 1, ), ) self._pressures_mapu = stencil_factory.from_origin_domain( pressures_mapu, origin=grid_indexing.origin_compute(), domain=self._domain_jextra, ) self._kord_mt = config.kord_mt self._map_single_u = MapSingle( stencil_factory, self._kord_mt, -1, grid_indexing.isc, grid_indexing.iec, grid_indexing.jsc, grid_indexing.jec + 1, ) self._pressures_mapv = stencil_factory.from_origin_domain( pressures_mapv, origin=grid_indexing.origin_compute(), domain=( grid_indexing.domain[0] + 1, grid_indexing.domain[1], grid_indexing.domain[2] + 1, ), ) self._map_single_v = MapSingle( stencil_factory, self._kord_mt, -1, grid_indexing.isc, grid_indexing.iec + 1, grid_indexing.jsc, grid_indexing.jec, ) ax_offsets_jextra = axis_offsets( grid_indexing, grid_indexing.origin_compute(), grid_indexing.domain_compute(add=(0, 1, 0)), ) self._update_ua = stencil_factory.from_origin_domain( update_ua, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(add=(0, 1, 0)), externals={**ax_offsets_jextra}, ) self._copy_from_below_stencil = stencil_factory.from_origin_domain( copy_from_below, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(), ) self._moist_cv_last_step_stencil = stencil_factory.from_origin_domain( moist_pt_last_step, origin=(grid_indexing.isc, grid_indexing.jsc, 0), domain=( grid_indexing.domain[0], grid_indexing.domain[1], grid_indexing.domain[2] + 1, ), ) self._basic_adjust_divide_stencil = stencil_factory.from_origin_domain( adjust_divide_stencil, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(), ) self._do_sat_adjust = config.do_sat_adj self.kmp = grid_indexing.domain[2] - 1 for k in range(pfull.shape[0]): if pfull[k] > 10.0e2: self.kmp = k break self._saturation_adjustment = SatAdjust3d(stencil_factory, config.sat_adjust, area_64, self.kmp) self._sum_te_stencil = stencil_factory.from_origin_domain( sum_te, origin=(grid_indexing.isc, grid_indexing.jsc, self.kmp), domain=( grid_indexing.domain[0], grid_indexing.domain[1], grid_indexing.domain[2] - self.kmp, ), )
def __init__( self, stencil_factory: StencilFactory, grid_data: GridData, nested: bool, grid_type: int, dord4: bool, ): grid_indexing = stencil_factory.grid_indexing self._cosa_s = grid_data.cosa_s self._cosa_u = grid_data.cosa_u self._cosa_v = grid_data.cosa_v self._rsin_u = grid_data.rsin_u self._rsin_v = grid_data.rsin_v self._rsin2 = grid_data.rsin2 self._dxa = grid_data.dxa self._dya = grid_data.dya self._sin_sg1 = grid_data.sin_sg1 self._sin_sg2 = grid_data.sin_sg2 self._sin_sg3 = grid_data.sin_sg3 self._sin_sg4 = grid_data.sin_sg4 if grid_type >= 3: raise NotImplementedError("unimplemented grid_type >= 3") self._big_number = 1e30 # 1e8 if 32 bit nx = grid_indexing.iec + 1 # grid.npx + 2 ny = grid_indexing.jec + 1 # grid.npy + 2 i1 = grid_indexing.isc - 1 j1 = grid_indexing.jsc - 1 id_ = 1 if dord4 else 0 pad = 2 + 2 * id_ npt = 4 if not nested else 0 if npt > grid_indexing.domain[0] - 1 or npt > grid_indexing.domain[ 1] - 1: npt = 0 self._utmp = utils.make_storage_from_shape( grid_indexing.max_shape, grid_indexing.origin_full(), ) self._vtmp = utils.make_storage_from_shape(grid_indexing.max_shape, grid_indexing.origin_full()) js1 = npt + OFFSET if grid_indexing.south_edge else grid_indexing.jsc - 1 je1 = ny - npt if grid_indexing.north_edge else grid_indexing.jec + 1 is1 = npt + OFFSET if grid_indexing.west_edge else grid_indexing.isd ie1 = nx - npt if grid_indexing.east_edge else grid_indexing.ied is2 = npt + OFFSET if grid_indexing.west_edge else grid_indexing.isc - 1 ie2 = nx - npt if grid_indexing.east_edge else grid_indexing.iec + 1 js2 = npt + OFFSET if grid_indexing.south_edge else grid_indexing.jsd je2 = ny - npt if grid_indexing.north_edge else grid_indexing.jed ifirst = (grid_indexing.isc + 2 if grid_indexing.west_edge else grid_indexing.isc - 1) ilast = (grid_indexing.iec - 1 if grid_indexing.east_edge else grid_indexing.iec + 2) idiff = ilast - ifirst + 1 jfirst = (grid_indexing.jsc + 2 if grid_indexing.south_edge else grid_indexing.jsc - 1) jlast = (grid_indexing.jec - 1 if grid_indexing.north_edge else grid_indexing.jec + 2) jdiff = jlast - jfirst + 1 self._set_tmps = stencil_factory.from_dims_halo( func=set_tmps, compute_dims=[X_DIM, Y_DIM, Z_DIM], compute_halos=(3, 3), ) self._lagrange_interpolation_y_p1 = stencil_factory.from_origin_domain( func=lagrange_interpolation_y_p1, origin=(is1, js1, 0), domain=(ie1 - is1 + 1, je1 - js1 + 1, grid_indexing.domain[2]), ) self._lagrange_interpolation_x_p1 = stencil_factory.from_origin_domain( func=lagrange_interpolation_x_p1, origin=(is2, js2, 0), domain=(ie2 - is2 + 1, je2 - js2 + 1, grid_indexing.domain[2]), ) origin = grid_indexing.origin_full() domain = grid_indexing.domain_full() ax_offsets = axis_offsets(grid_indexing, origin, domain) if npt == 0: d2a2c_avg_offset = -1 else: d2a2c_avg_offset = 3 self._avg_box = stencil_factory.from_dims_halo( func=avg_box, externals={"D2A2C_AVG_OFFSET": d2a2c_avg_offset}, compute_dims=[X_DIM, Y_DIM, Z_DIM], compute_halos=(3, 3), ) self._contravariant_components = stencil_factory.from_origin_domain( func=contravariant_components, origin=grid_indexing.origin_compute(add=(-1 - id_, -1 - id_, 0)), domain=grid_indexing.domain_compute(add=(pad, pad, 0)), ) origin_edges = grid_indexing.origin_compute(add=(-3, -3, 0)) domain_edges = grid_indexing.domain_compute(add=(6, 6, 0)) ax_offsets_edges = axis_offsets(grid_indexing, origin_edges, domain_edges) self._fill_corners_x = stencil_factory.from_origin_domain( func=fill_corners_x, externals=ax_offsets_edges, origin=origin_edges, domain=domain_edges, ) self._ut_main = stencil_factory.from_origin_domain( func=ut_main, origin=(ifirst, j1, 0), domain=(idiff, grid_indexing.domain[1] + 2, grid_indexing.domain[2]), ) self._east_west_edges = stencil_factory.from_origin_domain( func=east_west_edges, externals={ "i_end": ax_offsets_edges["i_end"], "i_start": ax_offsets_edges["i_start"], "local_je": ax_offsets_edges["local_je"], "local_js": ax_offsets_edges["local_js"], }, origin=origin_edges, domain=domain_edges, ) # Ydir: self._fill_corners_y = stencil_factory.from_origin_domain( func=fill_corners_y, externals={ **ax_offsets_edges, }, origin=origin_edges, domain=domain_edges, ) self._north_south_edges = stencil_factory.from_origin_domain( func=north_south_edges, externals={ "j_end": ax_offsets_edges["j_end"], "j_start": ax_offsets_edges["j_start"], "local_ie": ax_offsets_edges["local_ie"], "local_is": ax_offsets_edges["local_is"], "local_je": ax_offsets_edges["local_je"], "local_js": ax_offsets_edges["local_js"], }, origin=origin_edges, domain=domain_edges, ) self._vt_main = stencil_factory.from_origin_domain( func=vt_main, origin=(i1, jfirst, 0), domain=(grid_indexing.domain[0] + 2, jdiff, grid_indexing.domain[2]), )