diff --git a/pyfv3/stencils/dyn_core.py b/pyfv3/stencils/dyn_core.py index 754c413c..903f6720 100644 --- a/pyfv3/stencils/dyn_core.py +++ b/pyfv3/stencils/dyn_core.py @@ -360,7 +360,11 @@ def __init__( ["pkc"], ) else: - self.pkc = comm.get_scalar_halo_updater([full_size_xyzi_halo_spec]) + self.pkc = WrappedHaloUpdater( + comm.get_scalar_halo_updater([full_size_xyzi_halo_spec]), + {"pkc": pkc}, + ["pkc"], + ) self.uc__vc = WrappedHaloUpdater( comm.get_vector_halo_updater( [full_size_xiyz_halo_spec], [full_size_xyiz_halo_spec] diff --git a/pyfv3/stencils/fxadv.py b/pyfv3/stencils/fxadv.py index c64c99fe..65c480a2 100644 --- a/pyfv3/stencils/fxadv.py +++ b/pyfv3/stencils/fxadv.py @@ -538,15 +538,9 @@ def __init__( origin = grid_indexing.origin_full() domain = grid_indexing.domain_full() ax_offsets = grid_indexing.axis_offsets(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 = grid_indexing.axis_offsets(origin_corners, domain_corners) - kwargs_corners = { - "externals": corner_offsets, - "origin": origin_corners, - "domain": domain_corners, - } self._main_uc_vc_contra_stencil = stencil_factory.from_origin_domain( main_uc_vc_contra, externals={"grid_type": grid_type, **ax_offsets}, @@ -555,25 +549,46 @@ def __init__( ) if self._grid_type < 3: self._uc_contra_y_edge_stencil = stencil_factory.from_origin_domain( - uc_contra_y_edge, **kwargs + uc_contra_y_edge, + externals=ax_offsets, + origin=origin, + domain=domain, ) self._vc_contra_y_edge_stencil = stencil_factory.from_origin_domain( - vc_contra_y_edge, **kwargs + vc_contra_y_edge, + externals=ax_offsets, + origin=origin, + domain=domain, ) self._vc_contra_x_edge_stencil = stencil_factory.from_origin_domain( - vc_contra_x_edge, **kwargs + vc_contra_x_edge, + externals=ax_offsets, + origin=origin, + domain=domain, ) self._uc_contra_x_edge_stencil = stencil_factory.from_origin_domain( - uc_contra_x_edge, **kwargs + uc_contra_x_edge, + externals=ax_offsets, + origin=origin, + domain=domain, ) self._uc_contra_corners_stencil = stencil_factory.from_origin_domain( - uc_contra_corners, **kwargs_corners + uc_contra_corners, + externals=corner_offsets, + origin=origin_corners, + domain=domain_corners, ) self._vc_contra_corners_stencil = stencil_factory.from_origin_domain( - vc_contra_corners, **kwargs_corners + vc_contra_corners, + externals=corner_offsets, + origin=origin_corners, + domain=domain_corners, ) self._fxadv_fluxes_stencil = stencil_factory.from_origin_domain( - fxadv_fluxes_stencil, **kwargs + fxadv_fluxes_stencil, + externals=ax_offsets, + origin=origin, + domain=domain, ) # self._set_nans = get_set_nan_func( # grid_indexing,