From a275ac76f12f79639e067f5b1bf111eff7648bd6 Mon Sep 17 00:00:00 2001 From: Tobias Wicky-Pfund Date: Thu, 28 Aug 2025 22:04:10 +0200 Subject: [PATCH 1/5] math is not a thing --- pyfv3/stencils/delnflux.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyfv3/stencils/delnflux.py b/pyfv3/stencils/delnflux.py index f80d210b..8be8a883 100644 --- a/pyfv3/stencils/delnflux.py +++ b/pyfv3/stencils/delnflux.py @@ -120,7 +120,7 @@ def d2_highorder_stencil( ): with computation(PARALLEL), interval(...): if nord > current_nord: - d2 = ((fx - fx[1, 0, 0]) + (fy - fy[0, 1, 0])) * rarea + d2 = (fx - fx[1, 0, 0] + fy - fy[0, 1, 0]) * rarea def d2_damp_interval( From 4f06cbd86e6a11b3c26b330e5a2963a1829ec9b6 Mon Sep 17 00:00:00 2001 From: Tobias Wicky-Pfund Date: Tue, 2 Sep 2025 18:27:22 +0200 Subject: [PATCH 2/5] no more corner codes --- pyfv3/stencils/delnflux.py | 226 ++++++++++++++++++++++++------------- 1 file changed, 149 insertions(+), 77 deletions(-) diff --git a/pyfv3/stencils/delnflux.py b/pyfv3/stencils/delnflux.py index 8be8a883..aac8263f 100644 --- a/pyfv3/stencils/delnflux.py +++ b/pyfv3/stencils/delnflux.py @@ -194,77 +194,43 @@ def copy_corners_y_nord( with computation(PARALLEL), interval(...): if nord > current_nord: - with horizontal( - region[i_start - 3, j_start - 3], region[i_start - 3, j_end + 3] - ): - q_out = q_in[5, 0, 0] - with horizontal( - region[i_start - 2, j_start - 3], region[i_start - 3, j_end + 2] - ): - q_out = q_in[4, 1, 0] - with horizontal( - region[i_start - 1, j_start - 3], region[i_start - 3, j_end + 1] - ): - q_out = q_in[3, 2, 0] - with horizontal( - region[i_start - 3, j_start - 2], region[i_start - 2, j_end + 3] - ): - q_out = q_in[4, -1, 0] - with horizontal( - region[i_start - 2, j_start - 2], region[i_start - 2, j_end + 2] - ): - q_out = q_in[3, 0, 0] - with horizontal( - region[i_start - 1, j_start - 2], region[i_start - 2, j_end + 1] - ): - q_out = q_in[2, 1, 0] - with horizontal( - region[i_start - 3, j_start - 1], region[i_start - 1, j_end + 3] - ): - q_out = q_in[3, -2, 0] - with horizontal( - region[i_start - 2, j_start - 1], region[i_start - 1, j_end + 2] - ): - q_out = q_in[2, -1, 0] - with horizontal( - region[i_start - 1, j_start - 1], region[i_start - 1, j_end + 1] - ): - q_out = q_in[1, 0, 0] - with horizontal( - region[i_end + 1, j_start - 3], region[i_end + 3, j_end + 1] - ): - q_out = q_in[-3, 2, 0] - with horizontal( - region[i_end + 2, j_start - 3], region[i_end + 3, j_end + 2] - ): - q_out = q_in[-4, 1, 0] - with horizontal( - region[i_end + 3, j_start - 3], region[i_end + 3, j_end + 3] - ): - q_out = q_in[-5, 0, 0] - with horizontal( - region[i_end + 1, j_start - 2], region[i_end + 2, j_end + 1] - ): - q_out = q_in[-2, 1, 0] - with horizontal( - region[i_end + 2, j_start - 2], region[i_end + 2, j_end + 2] - ): - q_out = q_in[-3, 0, 0] - with horizontal( - region[i_end + 3, j_start - 2], region[i_end + 2, j_end + 3] - ): - q_out = q_in[-4, -1, 0] - with horizontal( - region[i_end + 1, j_start - 1], region[i_end + 1, j_end + 1] - ): + # with horizontal(region[i_start - 3, j_end + 3]): + # q_out = q_in[5, 0, 0] + # with horizontal(region[i_start - 3, j_end + 2]): + # q_out = q_in[4, 1, 0] + # with horizontal(region[i_start - 3, j_end + 1]): + # q_out = q_in[3, 2, 0] + # with horizontal(region[i_start - 2, j_end + 3]): + # q_out = q_in[4, -1, 0] + # with horizontal(region[i_start - 2, j_end + 2]): + # q_out = q_in[3, 0, 0] + # with horizontal(region[i_start - 2, j_end + 1]): + # q_out = q_in[2, 1, 0] + # with horizontal(region[i_start - 1, j_end + 3]): + # q_out = q_in[3, -2, 0] + # with horizontal(region[i_start - 1, j_end + 2]): + # q_out = q_in[2, -1, 0] + # with horizontal(region[i_start - 1, j_end + 1]): + # q_out = q_in[1, 0, 0] + # with horizontal(region[i_end + 3, j_end + 1]): + # q_out = q_in[-3, 2, 0] + # with horizontal(region[i_end + 3, j_end + 2]): + # q_out = q_in[-4, 1, 0] + # with horizontal(region[i_end + 3, j_end + 3]): + # q_out = q_in[-5, 0, 0] + + # with horizontal(region[i_end + 2, j_end + 1]): + # q_out = q_in[-2, 1, 0] + # with horizontal(region[i_end + 2, j_end + 2]): + # q_out = q_in[-3, 0, 0] + # with horizontal(region[i_end + 2, j_end + 3]): + # q_out = q_in[-4, -1, 0] + + with horizontal(region[i_end + 1, j_end + 1]): q_out = q_in[-1, 0, 0] - with horizontal( - region[i_end + 2, j_start - 1], region[i_end + 1, j_end + 2] - ): + with horizontal(region[i_end + 1, j_end + 2]): q_out = q_in[-2, -1, 0] - with horizontal( - region[i_end + 3, j_start - 1], region[i_end + 1, j_end + 3] - ): + with horizontal(region[i_end + 1, j_end + 3]): q_out = q_in[-3, -2, 0] @@ -621,6 +587,108 @@ def __init__( skip_passes=("UnreachableStmtPruning",), ) + def numpy_corner_y(self, d2, nord_data: Quantity): + for k in range(nord_data.shape[0]): + if nord_data.data[k] > 0: + d2[0, 0, k] = d2[5, 0, k] + d2[1, 0, k] = d2[5, 1, k] + d2[2, 0, k] = d2[5, 2, k] + + d2[0, 1, k] = d2[4, 0, k] + d2[1, 1, k] = d2[4, 1, k] + d2[2, 1, k] = d2[4, 2, k] + + d2[0, 2, k] = d2[3, 0, k] + d2[1, 2, k] = d2[3, 1, k] + d2[2, 2, k] = d2[3, 2, k] + + d2[-4, 0, k] = d2[-7, 2, k] + d2[-3, 0, k] = d2[-7, 1, k] + d2[-2, 0, k] = d2[-7, 0, k] + + d2[-4, 1, k] = d2[-6, 2, k] + d2[-3, 1, k] = d2[-6, 1, k] + d2[-2, 1, k] = d2[-6, 0, k] + + d2[-4, 2, k] = d2[-5, 2, k] + d2[-3, 2, k] = d2[-5, 1, k] + d2[-2, 2, k] = d2[-5, 0, k] + + d2[0, -2, k] = d2[5, -2, k] + d2[0, -3, k] = d2[4, -2, k] + d2[0, -4, k] = d2[3, -2, k] + + d2[1, -2, k] = d2[5, -3, k] + d2[1, -3, k] = d2[4, -3, k] + d2[1, -4, k] = d2[3, -3, k] + + d2[2, -2, k] = d2[5, -4, k] + d2[2, -3, k] = d2[4, -4, k] + d2[2, -4, k] = d2[3, -4, k] + + d2[-2, -4, k] = d2[-5, -2, k] + d2[-2, -3, k] = d2[-6, -2, k] + d2[-2, -2, k] = d2[-7, -2, k] + + d2[-3, -4, k] = d2[-5, -3, k] + d2[-3, -3, k] = d2[-6, -3, k] + d2[-3, -2, k] = d2[-7, -3, k] + + d2[-4, -4, k] = d2[-5, -4, k] + d2[-4, -3, k] = d2[-6, -4, k] + d2[-4, -2, k] = d2[-7, -4, k] + + def numpy_corner_x(self, d2, nord_data: Quantity): + for k in range(nord_data.shape[0]): + if nord_data.data[k] > 0: + d2[0, 0, k] = d2[0, 5, k] + d2[1, 0, k] = d2[0, 4, k] + d2[2, 0, k] = d2[0, 3, k] + + d2[0, 1, k] = d2[1, 5, k] + d2[1, 1, k] = d2[1, 4, k] + d2[2, 1, k] = d2[1, 3, k] + + d2[0, 2, k] = d2[2, 5, k] + d2[1, 2, k] = d2[2, 4, k] + d2[2, 2, k] = d2[2, 3, k] + + d2[0, -4, k] = d2[2, -7, k] + d2[1, -4, k] = d2[2, -6, k] + d2[2, -4, k] = d2[2, -5, k] + + d2[0, -3, k] = d2[1, -7, k] + d2[1, -3, k] = d2[1, -6, k] + d2[2, -3, k] = d2[1, -5, k] + + d2[0, -2, k] = d2[0, -7, k] + d2[1, -2, k] = d2[0, -6, k] + d2[2, -2, k] = d2[0, -5, k] + + d2[-2, 0, k] = d2[-2, 5, k] + d2[-2, 1, k] = d2[-3, 5, k] + d2[-2, 2, k] = d2[-4, 5, k] + + d2[-3, 0, k] = d2[-2, 4, k] + d2[-3, 1, k] = d2[-3, 4, k] + d2[-3, 2, k] = d2[-4, 4, k] + + d2[-4, 0, k] = d2[-2, 3, k] + d2[-4, 1, k] = d2[-3, 3, k] + d2[-4, 2, k] = d2[-4, 3, k] + + d2[-4, -2, k] = d2[-2, -5, k] + d2[-4, -3, k] = d2[-3, -5, k] + d2[-4, -4, k] = d2[-4, -5, k] + + d2[-3, -2, k] = d2[-2, -6, k] + d2[-3, -3, k] = d2[-3, -6, k] + d2[-3, -4, k] = d2[-4, -6, k] + + d2[-2, -2, k] = d2[-2, -7, k] + d2[-2, -3, k] = d2[-3, -7, k] + d2[-2, -4, k] = d2[-4, -7, k] + def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): """ Computes flux fields which would apply del-n damping to q, @@ -643,11 +711,13 @@ def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): else: self._copy_stencil_interval(q_in=q, q_out=d2, nord=self._nord) - self._copy_corners_x_nord(q_in=d2, q_out=d2, nord=self._nord, current_nord=0) + # self._copy_corners_x_nord(q_in=d2, q_out=d2, nord=self._nord, current_nord=0) + self.numpy_corner_x(d2, self._nord) self._fx_calc_stencil(q=d2, del6_v=self._del6_v, fx=fx2, nord=self._nord) - self._copy_corners_y_nord(q_in=d2, q_out=d2, nord=self._nord, current_nord=0) + # self._copy_corners_y_nord(q_in=d2, q_out=d2, nord=self._nord, current_nord=0) + self.numpy_corner_y(d2, self._nord) self._fy_calc_stencil(q=d2, del6_u=self._del6_u, fy=fy2, nord=self._nord) @@ -661,17 +731,19 @@ def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): current_nord=n, ) - self._copy_corners_x_nord( - q_in=d2, q_out=d2, nord=self._nord, current_nord=n - ) + # self._copy_corners_x_nord( + # q_in=d2, q_out=d2, nord=self._nord, current_nord=n + # ) + self.numpy_corner_x(d2, self._nord) self._column_conditional_fx_calculation[n]( q=d2, del6_v=self._del6_v, fx=fx2, nord=self._nord, current_nord=n ) - self._copy_corners_y_nord( - q_in=d2, q_out=d2, nord=self._nord, current_nord=n - ) + # self._copy_corners_y_nord( + # q_in=d2, q_out=d2, nord=self._nord, current_nord=n + # ) + self.numpy_corner_y(d2, self._nord) self._column_conditional_fy_calculation[n]( q=d2, del6_u=self._del6_u, fy=fy2, nord=self._nord, current_nord=n From 6ec66b9f36ae3fb3b54a02f197c8e4369d81c743 Mon Sep 17 00:00:00 2001 From: Tobias Wicky-Pfund Date: Tue, 2 Sep 2025 18:29:28 +0200 Subject: [PATCH 3/5] make blocks red --- pyfv3/stencils/delnflux.py | 167 ------------------------------------- 1 file changed, 167 deletions(-) diff --git a/pyfv3/stencils/delnflux.py b/pyfv3/stencils/delnflux.py index aac8263f..6e7300bf 100644 --- a/pyfv3/stencils/delnflux.py +++ b/pyfv3/stencils/delnflux.py @@ -182,144 +182,6 @@ def diffusive_damp( fy = fy + 0.5 * damp * (mass[0, -1, 0] + mass) * fy2 -def copy_corners_y_nord( - q_in: FloatField, q_out: FloatField, nord: FloatFieldK, current_nord: int -): - """ - Args: - q_in (in): - q_out (out): - """ - from __externals__ import i_end, i_start, j_end, j_start - - with computation(PARALLEL), interval(...): - if nord > current_nord: - # with horizontal(region[i_start - 3, j_end + 3]): - # q_out = q_in[5, 0, 0] - # with horizontal(region[i_start - 3, j_end + 2]): - # q_out = q_in[4, 1, 0] - # with horizontal(region[i_start - 3, j_end + 1]): - # q_out = q_in[3, 2, 0] - # with horizontal(region[i_start - 2, j_end + 3]): - # q_out = q_in[4, -1, 0] - # with horizontal(region[i_start - 2, j_end + 2]): - # q_out = q_in[3, 0, 0] - # with horizontal(region[i_start - 2, j_end + 1]): - # q_out = q_in[2, 1, 0] - # with horizontal(region[i_start - 1, j_end + 3]): - # q_out = q_in[3, -2, 0] - # with horizontal(region[i_start - 1, j_end + 2]): - # q_out = q_in[2, -1, 0] - # with horizontal(region[i_start - 1, j_end + 1]): - # q_out = q_in[1, 0, 0] - # with horizontal(region[i_end + 3, j_end + 1]): - # q_out = q_in[-3, 2, 0] - # with horizontal(region[i_end + 3, j_end + 2]): - # q_out = q_in[-4, 1, 0] - # with horizontal(region[i_end + 3, j_end + 3]): - # q_out = q_in[-5, 0, 0] - - # with horizontal(region[i_end + 2, j_end + 1]): - # q_out = q_in[-2, 1, 0] - # with horizontal(region[i_end + 2, j_end + 2]): - # q_out = q_in[-3, 0, 0] - # with horizontal(region[i_end + 2, j_end + 3]): - # q_out = q_in[-4, -1, 0] - - with horizontal(region[i_end + 1, j_end + 1]): - q_out = q_in[-1, 0, 0] - with horizontal(region[i_end + 1, j_end + 2]): - q_out = q_in[-2, -1, 0] - with horizontal(region[i_end + 1, j_end + 3]): - q_out = q_in[-3, -2, 0] - - -def copy_corners_x_nord( - q_in: FloatField, q_out: FloatField, nord: FloatFieldK, current_nord: int -): - """ - Args: - q_in (in): - q_out (out): - """ - from __externals__ import i_end, i_start, j_end, j_start - - with computation(PARALLEL), interval(...): - if nord > current_nord: - with horizontal( - region[i_start - 3, j_start - 3], region[i_end + 3, j_start - 3] - ): - q_out = q_in[0, 5, 0] - with horizontal( - region[i_start - 2, j_start - 3], region[i_end + 3, j_start - 2] - ): - q_out = q_in[-1, 4, 0] - with horizontal( - region[i_start - 1, j_start - 3], region[i_end + 3, j_start - 1] - ): - q_out = q_in[-2, 3, 0] - with horizontal( - region[i_start - 3, j_start - 2], region[i_end + 2, j_start - 3] - ): - q_out = q_in[1, 4, 0] - with horizontal( - region[i_start - 2, j_start - 2], region[i_end + 2, j_start - 2] - ): - q_out = q_in[0, 3, 0] - with horizontal( - region[i_start - 1, j_start - 2], region[i_end + 2, j_start - 1] - ): - q_out = q_in[-1, 2, 0] - with horizontal( - region[i_start - 3, j_start - 1], region[i_end + 1, j_start - 3] - ): - q_out = q_in[2, 3, 0] - with horizontal( - region[i_start - 2, j_start - 1], region[i_end + 1, j_start - 2] - ): - q_out = q_in[1, 2, 0] - with horizontal( - region[i_start - 1, j_start - 1], region[i_end + 1, j_start - 1] - ): - q_out = q_in[0, 1, 0] - with horizontal( - region[i_start - 3, j_end + 1], region[i_end + 1, j_end + 3] - ): - q_out = q_in[2, -3, 0] - with horizontal( - region[i_start - 2, j_end + 1], region[i_end + 1, j_end + 2] - ): - q_out = q_in[1, -2, 0] - with horizontal( - region[i_start - 1, j_end + 1], region[i_end + 1, j_end + 1] - ): - q_out = q_in[0, -1, 0] - with horizontal( - region[i_start - 3, j_end + 2], region[i_end + 2, j_end + 3] - ): - q_out = q_in[1, -4, 0] - with horizontal( - region[i_start - 2, j_end + 2], region[i_end + 2, j_end + 2] - ): - q_out = q_in[0, -3, 0] - with horizontal( - region[i_start - 1, j_end + 2], region[i_end + 2, j_end + 1] - ): - q_out = q_in[-1, -2, 0] - with horizontal( - region[i_start - 3, j_end + 3], region[i_end + 3, j_end + 3] - ): - q_out = q_in[0, -5, 0] - with horizontal( - region[i_start - 2, j_end + 3], region[i_end + 3, j_end + 2] - ): - q_out = q_in[-1, -4, 0] - with horizontal( - region[i_start - 1, j_end + 3], region[i_end + 3, j_end + 1] - ): - q_out = q_in[-2, -3, 0] - - class DelnFlux: """ Fortran name is deln_flux @@ -565,27 +427,6 @@ def __init__( 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 = grid_indexing.axis_offsets(corner_origin, corner_domain) - - self._copy_corners_x_nord = stencil_factory.from_origin_domain( - copy_corners_x_nord, - externals={**corner_axis_offsets}, - origin=corner_origin, - domain=corner_domain, - skip_passes=("UnreachableStmtPruning",), - ) - self._copy_corners_y_nord = stencil_factory.from_origin_domain( - copy_corners_y_nord, - externals={**corner_axis_offsets}, - origin=corner_origin, - domain=corner_domain, - skip_passes=("UnreachableStmtPruning",), - ) def numpy_corner_y(self, d2, nord_data: Quantity): for k in range(nord_data.shape[0]): @@ -711,12 +552,10 @@ def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): else: self._copy_stencil_interval(q_in=q, q_out=d2, nord=self._nord) - # self._copy_corners_x_nord(q_in=d2, q_out=d2, nord=self._nord, current_nord=0) self.numpy_corner_x(d2, self._nord) self._fx_calc_stencil(q=d2, del6_v=self._del6_v, fx=fx2, nord=self._nord) - # self._copy_corners_y_nord(q_in=d2, q_out=d2, nord=self._nord, current_nord=0) self.numpy_corner_y(d2, self._nord) self._fy_calc_stencil(q=d2, del6_u=self._del6_u, fy=fy2, nord=self._nord) @@ -731,18 +570,12 @@ def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): current_nord=n, ) - # self._copy_corners_x_nord( - # q_in=d2, q_out=d2, nord=self._nord, current_nord=n - # ) self.numpy_corner_x(d2, self._nord) self._column_conditional_fx_calculation[n]( q=d2, del6_v=self._del6_v, fx=fx2, nord=self._nord, current_nord=n ) - # self._copy_corners_y_nord( - # q_in=d2, q_out=d2, nord=self._nord, current_nord=n - # ) self.numpy_corner_y(d2, self._nord) self._column_conditional_fy_calculation[n]( From e0ba21cd6da969136a1540499a525b3cb644e7dc Mon Sep 17 00:00:00 2001 From: Tobias Wicky-Pfund Date: Thu, 4 Sep 2025 16:07:00 +0200 Subject: [PATCH 4/5] map update --- pyfv3/stencils/delnflux.py | 220 +++++++++++++++++++------------------ 1 file changed, 112 insertions(+), 108 deletions(-) diff --git a/pyfv3/stencils/delnflux.py b/pyfv3/stencils/delnflux.py index 6e7300bf..beb99dbf 100644 --- a/pyfv3/stencils/delnflux.py +++ b/pyfv3/stencils/delnflux.py @@ -8,6 +8,8 @@ from ndsl.dsl.stencil import get_stencils_with_varied_bounds from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK from ndsl.grid import DampingCoefficients +from dace.frontend.python.interface import nounroll as dace_nounroll +import dace def calc_damp(damp_c: Quantity, da_min: Float, nord: Quantity) -> Quantity: @@ -282,9 +284,9 @@ def __call__( # fx2 and fy2 are local variables containing the diffusive flux, which # gets added to the base flux below if d2 is None: - self.delnflux_nosg(q, self._fx2, self._fy2, self._damp, self._d2, mass) + self.delnflux_nosg(q, self._fx2, self._fy2, self._damp, self._d2.data, mass) else: - self.delnflux_nosg(q, self._fx2, self._fy2, self._damp, d2, mass) + self.delnflux_nosg(q, self._fx2, self._fy2, self._damp, self._d2.data, mass) if mass is None: self._add_diffusive_stencil(fx, self._fx2, fy, self._fy2) @@ -299,6 +301,110 @@ def __call__( return fx, fy +def copy_corner_x(nord, d2): + for k in dace.map[nord.data.shape[0]]: + if nord.data[k] > 0: + d2[0, 0, k] = d2[0, 5, k] + d2[1, 0, k] = d2[0, 4, k] + d2[2, 0, k] = d2[0, 3, k] + + d2[0, 1, k] = d2[1, 5, k] + d2[1, 1, k] = d2[1, 4, k] + d2[2, 1, k] = d2[1, 3, k] + + d2[0, 2, k] = d2[2, 5, k] + d2[1, 2, k] = d2[2, 4, k] + d2[2, 2, k] = d2[2, 3, k] + + d2[0, -4, k] = d2[2, -7, k] + d2[1, -4, k] = d2[2, -6, k] + d2[2, -4, k] = d2[2, -5, k] + + d2[0, -3, k] = d2[1, -7, k] + d2[1, -3, k] = d2[1, -6, k] + d2[2, -3, k] = d2[1, -5, k] + + d2[0, -2, k] = d2[0, -7, k] + d2[1, -2, k] = d2[0, -6, k] + d2[2, -2, k] = d2[0, -5, k] + + d2[-2, 0, k] = d2[-2, 5, k] + d2[-2, 1, k] = d2[-3, 5, k] + d2[-2, 2, k] = d2[-4, 5, k] + + d2[-3, 0, k] = d2[-2, 4, k] + d2[-3, 1, k] = d2[-3, 4, k] + d2[-3, 2, k] = d2[-4, 4, k] + + d2[-4, 0, k] = d2[-2, 3, k] + d2[-4, 1, k] = d2[-3, 3, k] + d2[-4, 2, k] = d2[-4, 3, k] + + d2[-4, -2, k] = d2[-2, -5, k] + d2[-4, -3, k] = d2[-3, -5, k] + d2[-4, -4, k] = d2[-4, -5, k] + + d2[-3, -2, k] = d2[-2, -6, k] + d2[-3, -3, k] = d2[-3, -6, k] + d2[-3, -4, k] = d2[-4, -6, k] + + d2[-2, -2, k] = d2[-2, -7, k] + d2[-2, -3, k] = d2[-3, -7, k] + d2[-2, -4, k] = d2[-4, -7, k] + + +def copy_corner_y(nord, d2): + for k in dace.map[nord.data.shape[0]]: + if nord.data[k] > 0: + d2[0, 0, k] = d2[5, 0, k] + d2[1, 0, k] = d2[5, 1, k] + d2[2, 0, k] = d2[5, 2, k] + + d2[0, 1, k] = d2[4, 0, k] + d2[1, 1, k] = d2[4, 1, k] + d2[2, 1, k] = d2[4, 2, k] + + d2[0, 2, k] = d2[3, 0, k] + d2[1, 2, k] = d2[3, 1, k] + d2[2, 2, k] = d2[3, 2, k] + + d2[-4, 0, k] = d2[-7, 2, k] + d2[-3, 0, k] = d2[-7, 1, k] + d2[-2, 0, k] = d2[-7, 0, k] + + d2[-4, 1, k] = d2[-6, 2, k] + d2[-3, 1, k] = d2[-6, 1, k] + d2[-2, 1, k] = d2[-6, 0, k] + + d2[-4, 2, k] = d2[-5, 2, k] + d2[-3, 2, k] = d2[-5, 1, k] + d2[-2, 2, k] = d2[-5, 0, k] + + d2[0, -2, k] = d2[5, -2, k] + d2[0, -3, k] = d2[4, -2, k] + d2[0, -4, k] = d2[3, -2, k] + + d2[1, -2, k] = d2[5, -3, k] + d2[1, -3, k] = d2[4, -3, k] + d2[1, -4, k] = d2[3, -3, k] + + d2[2, -2, k] = d2[5, -4, k] + d2[2, -3, k] = d2[4, -4, k] + d2[2, -4, k] = d2[3, -4, k] + + d2[-2, -4, k] = d2[-5, -2, k] + d2[-2, -3, k] = d2[-6, -2, k] + d2[-2, -2, k] = d2[-7, -2, k] + + d2[-3, -4, k] = d2[-5, -3, k] + d2[-3, -3, k] = d2[-6, -3, k] + d2[-3, -2, k] = d2[-7, -3, k] + + d2[-4, -4, k] = d2[-5, -4, k] + d2[-4, -3, k] = d2[-6, -4, k] + d2[-4, -2, k] = d2[-7, -4, k] + + class DelnFluxNoSG: """ This contains the mechanics of del6_vt and some of deln_flux from @@ -428,108 +534,6 @@ def __init__( domain=(f1_nx - 1, f1_ny + 1, nk), ) - def numpy_corner_y(self, d2, nord_data: Quantity): - for k in range(nord_data.shape[0]): - if nord_data.data[k] > 0: - d2[0, 0, k] = d2[5, 0, k] - d2[1, 0, k] = d2[5, 1, k] - d2[2, 0, k] = d2[5, 2, k] - - d2[0, 1, k] = d2[4, 0, k] - d2[1, 1, k] = d2[4, 1, k] - d2[2, 1, k] = d2[4, 2, k] - - d2[0, 2, k] = d2[3, 0, k] - d2[1, 2, k] = d2[3, 1, k] - d2[2, 2, k] = d2[3, 2, k] - - d2[-4, 0, k] = d2[-7, 2, k] - d2[-3, 0, k] = d2[-7, 1, k] - d2[-2, 0, k] = d2[-7, 0, k] - - d2[-4, 1, k] = d2[-6, 2, k] - d2[-3, 1, k] = d2[-6, 1, k] - d2[-2, 1, k] = d2[-6, 0, k] - - d2[-4, 2, k] = d2[-5, 2, k] - d2[-3, 2, k] = d2[-5, 1, k] - d2[-2, 2, k] = d2[-5, 0, k] - - d2[0, -2, k] = d2[5, -2, k] - d2[0, -3, k] = d2[4, -2, k] - d2[0, -4, k] = d2[3, -2, k] - - d2[1, -2, k] = d2[5, -3, k] - d2[1, -3, k] = d2[4, -3, k] - d2[1, -4, k] = d2[3, -3, k] - - d2[2, -2, k] = d2[5, -4, k] - d2[2, -3, k] = d2[4, -4, k] - d2[2, -4, k] = d2[3, -4, k] - - d2[-2, -4, k] = d2[-5, -2, k] - d2[-2, -3, k] = d2[-6, -2, k] - d2[-2, -2, k] = d2[-7, -2, k] - - d2[-3, -4, k] = d2[-5, -3, k] - d2[-3, -3, k] = d2[-6, -3, k] - d2[-3, -2, k] = d2[-7, -3, k] - - d2[-4, -4, k] = d2[-5, -4, k] - d2[-4, -3, k] = d2[-6, -4, k] - d2[-4, -2, k] = d2[-7, -4, k] - - def numpy_corner_x(self, d2, nord_data: Quantity): - for k in range(nord_data.shape[0]): - if nord_data.data[k] > 0: - d2[0, 0, k] = d2[0, 5, k] - d2[1, 0, k] = d2[0, 4, k] - d2[2, 0, k] = d2[0, 3, k] - - d2[0, 1, k] = d2[1, 5, k] - d2[1, 1, k] = d2[1, 4, k] - d2[2, 1, k] = d2[1, 3, k] - - d2[0, 2, k] = d2[2, 5, k] - d2[1, 2, k] = d2[2, 4, k] - d2[2, 2, k] = d2[2, 3, k] - - d2[0, -4, k] = d2[2, -7, k] - d2[1, -4, k] = d2[2, -6, k] - d2[2, -4, k] = d2[2, -5, k] - - d2[0, -3, k] = d2[1, -7, k] - d2[1, -3, k] = d2[1, -6, k] - d2[2, -3, k] = d2[1, -5, k] - - d2[0, -2, k] = d2[0, -7, k] - d2[1, -2, k] = d2[0, -6, k] - d2[2, -2, k] = d2[0, -5, k] - - d2[-2, 0, k] = d2[-2, 5, k] - d2[-2, 1, k] = d2[-3, 5, k] - d2[-2, 2, k] = d2[-4, 5, k] - - d2[-3, 0, k] = d2[-2, 4, k] - d2[-3, 1, k] = d2[-3, 4, k] - d2[-3, 2, k] = d2[-4, 4, k] - - d2[-4, 0, k] = d2[-2, 3, k] - d2[-4, 1, k] = d2[-3, 3, k] - d2[-4, 2, k] = d2[-4, 3, k] - - d2[-4, -2, k] = d2[-2, -5, k] - d2[-4, -3, k] = d2[-3, -5, k] - d2[-4, -4, k] = d2[-4, -5, k] - - d2[-3, -2, k] = d2[-2, -6, k] - d2[-3, -3, k] = d2[-3, -6, k] - d2[-3, -4, k] = d2[-4, -6, k] - - d2[-2, -2, k] = d2[-2, -7, k] - d2[-2, -3, k] = d2[-3, -7, k] - d2[-2, -4, k] = d2[-4, -7, k] - def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): """ Computes flux fields which would apply del-n damping to q, @@ -552,11 +556,11 @@ def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): else: self._copy_stencil_interval(q_in=q, q_out=d2, nord=self._nord) - self.numpy_corner_x(d2, self._nord) + # copy_corner_x(self._nord, d2) self._fx_calc_stencil(q=d2, del6_v=self._del6_v, fx=fx2, nord=self._nord) - self.numpy_corner_y(d2, self._nord) + # copy_corner_y(self._nord, d2) self._fy_calc_stencil(q=d2, del6_u=self._del6_u, fy=fy2, nord=self._nord) @@ -570,13 +574,13 @@ def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): current_nord=n, ) - self.numpy_corner_x(d2, self._nord) + # copy_corner_x(self._nord, d2) self._column_conditional_fx_calculation[n]( q=d2, del6_v=self._del6_v, fx=fx2, nord=self._nord, current_nord=n ) - self.numpy_corner_y(d2, self._nord) + # copy_corner_y(self._nord, d2) self._column_conditional_fy_calculation[n]( q=d2, del6_u=self._del6_u, fy=fy2, nord=self._nord, current_nord=n From 83ef7ea287647d051e28034766e601c8765333ec Mon Sep 17 00:00:00 2001 From: Tobias Wicky-Pfund Date: Fri, 5 Sep 2025 17:15:57 +0200 Subject: [PATCH 5/5] 2d idea --- pyfv3/stencils/delnflux.py | 101 +++---------------------------------- pyfv3/stencils/fvtp2d.py | 8 +-- 2 files changed, 8 insertions(+), 101 deletions(-) diff --git a/pyfv3/stencils/delnflux.py b/pyfv3/stencils/delnflux.py index beb99dbf..6675c93b 100644 --- a/pyfv3/stencils/delnflux.py +++ b/pyfv3/stencils/delnflux.py @@ -8,8 +8,11 @@ from ndsl.dsl.stencil import get_stencils_with_varied_bounds from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK from ndsl.grid import DampingCoefficients -from dace.frontend.python.interface import nounroll as dace_nounroll import dace +from ndsl.stencils.corners import ( + two_dimensional_corner_copy_x, + two_dimensional_corner_copy_y, +) def calc_damp(damp_c: Quantity, da_min: Float, nord: Quantity) -> Quantity: @@ -304,105 +307,13 @@ def __call__( def copy_corner_x(nord, d2): for k in dace.map[nord.data.shape[0]]: if nord.data[k] > 0: - d2[0, 0, k] = d2[0, 5, k] - d2[1, 0, k] = d2[0, 4, k] - d2[2, 0, k] = d2[0, 3, k] - - d2[0, 1, k] = d2[1, 5, k] - d2[1, 1, k] = d2[1, 4, k] - d2[2, 1, k] = d2[1, 3, k] - - d2[0, 2, k] = d2[2, 5, k] - d2[1, 2, k] = d2[2, 4, k] - d2[2, 2, k] = d2[2, 3, k] - - d2[0, -4, k] = d2[2, -7, k] - d2[1, -4, k] = d2[2, -6, k] - d2[2, -4, k] = d2[2, -5, k] - - d2[0, -3, k] = d2[1, -7, k] - d2[1, -3, k] = d2[1, -6, k] - d2[2, -3, k] = d2[1, -5, k] - - d2[0, -2, k] = d2[0, -7, k] - d2[1, -2, k] = d2[0, -6, k] - d2[2, -2, k] = d2[0, -5, k] - - d2[-2, 0, k] = d2[-2, 5, k] - d2[-2, 1, k] = d2[-3, 5, k] - d2[-2, 2, k] = d2[-4, 5, k] - - d2[-3, 0, k] = d2[-2, 4, k] - d2[-3, 1, k] = d2[-3, 4, k] - d2[-3, 2, k] = d2[-4, 4, k] - - d2[-4, 0, k] = d2[-2, 3, k] - d2[-4, 1, k] = d2[-3, 3, k] - d2[-4, 2, k] = d2[-4, 3, k] - - d2[-4, -2, k] = d2[-2, -5, k] - d2[-4, -3, k] = d2[-3, -5, k] - d2[-4, -4, k] = d2[-4, -5, k] - - d2[-3, -2, k] = d2[-2, -6, k] - d2[-3, -3, k] = d2[-3, -6, k] - d2[-3, -4, k] = d2[-4, -6, k] - - d2[-2, -2, k] = d2[-2, -7, k] - d2[-2, -3, k] = d2[-3, -7, k] - d2[-2, -4, k] = d2[-4, -7, k] + two_dimensional_corner_copy_x(d2[:, :, k]) def copy_corner_y(nord, d2): for k in dace.map[nord.data.shape[0]]: if nord.data[k] > 0: - d2[0, 0, k] = d2[5, 0, k] - d2[1, 0, k] = d2[5, 1, k] - d2[2, 0, k] = d2[5, 2, k] - - d2[0, 1, k] = d2[4, 0, k] - d2[1, 1, k] = d2[4, 1, k] - d2[2, 1, k] = d2[4, 2, k] - - d2[0, 2, k] = d2[3, 0, k] - d2[1, 2, k] = d2[3, 1, k] - d2[2, 2, k] = d2[3, 2, k] - - d2[-4, 0, k] = d2[-7, 2, k] - d2[-3, 0, k] = d2[-7, 1, k] - d2[-2, 0, k] = d2[-7, 0, k] - - d2[-4, 1, k] = d2[-6, 2, k] - d2[-3, 1, k] = d2[-6, 1, k] - d2[-2, 1, k] = d2[-6, 0, k] - - d2[-4, 2, k] = d2[-5, 2, k] - d2[-3, 2, k] = d2[-5, 1, k] - d2[-2, 2, k] = d2[-5, 0, k] - - d2[0, -2, k] = d2[5, -2, k] - d2[0, -3, k] = d2[4, -2, k] - d2[0, -4, k] = d2[3, -2, k] - - d2[1, -2, k] = d2[5, -3, k] - d2[1, -3, k] = d2[4, -3, k] - d2[1, -4, k] = d2[3, -3, k] - - d2[2, -2, k] = d2[5, -4, k] - d2[2, -3, k] = d2[4, -4, k] - d2[2, -4, k] = d2[3, -4, k] - - d2[-2, -4, k] = d2[-5, -2, k] - d2[-2, -3, k] = d2[-6, -2, k] - d2[-2, -2, k] = d2[-7, -2, k] - - d2[-3, -4, k] = d2[-5, -3, k] - d2[-3, -3, k] = d2[-6, -3, k] - d2[-3, -2, k] = d2[-7, -3, k] - - d2[-4, -4, k] = d2[-5, -4, k] - d2[-4, -3, k] = d2[-6, -4, k] - d2[-4, -2, k] = d2[-7, -4, k] + two_dimensional_corner_copy_y(d2[:, :, k]) class DelnFluxNoSG: diff --git a/pyfv3/stencils/fvtp2d.py b/pyfv3/stencils/fvtp2d.py index 8906227e..6d0fa5ae 100644 --- a/pyfv3/stencils/fvtp2d.py +++ b/pyfv3/stencils/fvtp2d.py @@ -179,9 +179,7 @@ def make_quantity(): # self.delnflux = None self._do_delnflux = False - self._copy_corners_y: corners.CopyCorners = corners.CopyCorners( - "y", stencil_factory - ) + self._copy_corners_y: corners.CopyCorners = corners.CopyCorners("y") self.y_piecewise_parabolic_inner = YPiecewiseParabolic( stencil_factory=stencil_factory, dya=grid_data.dya, @@ -204,9 +202,7 @@ def make_quantity(): domain=idx.domain_compute(add=(1, 1, 1)), ) - self._copy_corners_x: corners.CopyCorners = corners.CopyCorners( - "x", stencil_factory - ) + self._copy_corners_x: corners.CopyCorners = corners.CopyCorners("x") self.x_piecewise_parabolic_inner = XPiecewiseParabolic( stencil_factory=stencil_factory, dxa=grid_data.dxa,