From 40add50a6bbe43254436482484742d6b598c3573 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 4 May 2021 15:50:47 -0700 Subject: [PATCH 01/45] remerge ppm_volume_mean_x and y back into a2b_interpolation --- fv3core/stencils/a2b_ord4.py | 65 +++++++----------------------------- 1 file changed, 12 insertions(+), 53 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 6fc023866..0b817e1d6 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -235,14 +235,20 @@ def lagrange_x_func(qy): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) -def ppm_volume_mean_x( +def a2b_interpolation( qin: FloatField, + qout: FloatField, qx: FloatField, + qy: FloatField, + qxx: FloatField, + qyy: FloatField, dxa: FloatFieldIJ, + dya: FloatFieldIJ ): - from __externals__ import i_end, i_start + from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): + # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) with horizontal(region[i_start, :]): qx = qx_edge_west(qin, dxa) @@ -252,16 +258,7 @@ def ppm_volume_mean_x( qx = qx_edge_east(qin, dxa) with horizontal(region[i_end, :]): qx = qx_edge_east2(qin, dxa, qx) - - -def ppm_volume_mean_y( - qin: FloatField, - qy: FloatField, - dya: FloatFieldIJ, -): - from __externals__ import j_end, j_start - - with computation(PARALLEL), interval(...): + # ppm_volume_mean_y qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) with horizontal(region[:, j_start]): qy = qy_edge_south(qin, dya) @@ -272,17 +269,6 @@ def ppm_volume_mean_y( with horizontal(region[:, j_end]): qy = qy_edge_north2(qin, dya, qy) - -def a2b_interpolation( - qout: FloatField, - qx: FloatField, - qy: FloatField, - qxx: FloatField, - qyy: FloatField, -): - from __externals__ import i_end, i_start, j_end, j_start - - with computation(PARALLEL), interval(...): qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + qxx[0, 1, 0]) @@ -465,26 +451,6 @@ def __init__( self._qout_y_edge_north = FrozenStencil( qout_y_edge, origin=(is2, self.grid.je + 1, kstart), domain=(di2, 1, nk) ) - origin_x = (self.grid.is_, self.grid.js - 2, kstart) - domain_x = (self.grid.nic + 1, self.grid.njc + 4, nk) - ax_offsets_x = axis_offsets( - self.grid, - origin_x, - domain_x, - ) - self._ppm_volume_mean_x_stencil = FrozenStencil( - ppm_volume_mean_x, externals=ax_offsets_x, origin=origin_x, domain=domain_x - ) - origin_y = (self.grid.is_ - 2, self.grid.js, kstart) - domain_y = (self.grid.nic + 4, self.grid.njc + 1, nk) - ax_offsets_y = axis_offsets( - self.grid, - origin_y, - domain_y, - ) - self._ppm_volume_mean_y_stencil = FrozenStencil( - ppm_volume_mean_y, externals=ax_offsets_y, origin=origin_y, domain=domain_y - ) js = self.grid.js + 1 if self.grid.south_edge else self.grid.js je = self.grid.je if self.grid.north_edge else self.grid.je + 1 @@ -549,22 +515,15 @@ def __call__(self, qin: FloatField, qout: FloatField): ) self._compute_qout_edges(qin, qout) - self._ppm_volume_mean_x_stencil( - qin, - self._tmp_qx, - self.grid.dxa, - ) - self._ppm_volume_mean_y_stencil( - qin, - self._tmp_qy, - self.grid.dya, - ) self._a2b_interpolation_stencil( + qin, qout, self._tmp_qx, self._tmp_qy, self._tmp_qxx, self._tmp_qyy, + self.grid.dxa, + self.grid.dya, ) if self.replace: self._copy_stencil( From 8be8be8cbc682bbcdf83df6c61b9bf0ef53e6729 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 4 May 2021 16:32:59 -0700 Subject: [PATCH 02/45] is using extra variables any faster --- fv3core/stencils/a2b_ord4.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 0b817e1d6..55933b929 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -239,36 +239,43 @@ def a2b_interpolation( qin: FloatField, qout: FloatField, qx: FloatField, + qx_copy: FloatField, qy: FloatField, + qy_copy: FloatField, qxx: FloatField, qyy: FloatField, dxa: FloatFieldIJ, - dya: FloatFieldIJ + dya: FloatFieldIJ, ): from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) + qx_copy = qx with horizontal(region[i_start, :]): qx = qx_edge_west(qin, dxa) + qx_copy = qx with horizontal(region[i_start + 1, :]): - qx = qx_edge_west2(qin, dxa, qx) + qx = qx_edge_west2(qin, dxa, qx_copy) with horizontal(region[i_end + 1, :]): qx = qx_edge_east(qin, dxa) + qx_copy = qx with horizontal(region[i_end, :]): - qx = qx_edge_east2(qin, dxa, qx) + qx = qx_edge_east2(qin, dxa, qx_copy) # ppm_volume_mean_y qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) + qy_copy = qy with horizontal(region[:, j_start]): qy = qy_edge_south(qin, dya) + qy_copy = qy with horizontal(region[:, j_start + 1]): - qy = qy_edge_south2(qin, dya, qy) + qy = qy_edge_south2(qin, dya, qy_copy) with horizontal(region[:, j_end + 1]): qy = qy_edge_north(qin, dya) + qy_copy = qy with horizontal(region[:, j_end]): - qy = qy_edge_north2(qin, dya, qy) - + qy = qy_edge_north2(qin, dya, qy_copy) qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + qxx[0, 1, 0]) @@ -399,7 +406,8 @@ def __init__( self._tmp_qy = utils.make_storage_from_shape(shape) self._tmp_qxx = utils.make_storage_from_shape(shape) self._tmp_qyy = utils.make_storage_from_shape(shape) - + self._tmp_qx_copy = utils.make_storage_from_shape(shape) + self._tmp_qy_copy = utils.make_storage_from_shape(shape) if nk is None: nk = self.grid.npz - kstart corner_domain = (1, 1, nk) @@ -519,7 +527,9 @@ def __call__(self, qin: FloatField, qout: FloatField): qin, qout, self._tmp_qx, + self._tmp_qx_copy, self._tmp_qy, + self._tmp_qy_copy, self._tmp_qxx, self._tmp_qyy, self.grid.dxa, From 2b01c893f5a6e9e50425f1a419ef1ee03ea0ddf8 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Fri, 7 May 2021 09:05:03 -0700 Subject: [PATCH 03/45] Revert "is using extra variables any faster" This reverts commit 8be8be8cbc682bbcdf83df6c61b9bf0ef53e6729. --- fv3core/stencils/a2b_ord4.py | 24 +++++++----------------- 1 file changed, 7 insertions(+), 17 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 55933b929..0b817e1d6 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -239,43 +239,36 @@ def a2b_interpolation( qin: FloatField, qout: FloatField, qx: FloatField, - qx_copy: FloatField, qy: FloatField, - qy_copy: FloatField, qxx: FloatField, qyy: FloatField, dxa: FloatFieldIJ, - dya: FloatFieldIJ, + dya: FloatFieldIJ ): from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) - qx_copy = qx with horizontal(region[i_start, :]): qx = qx_edge_west(qin, dxa) - qx_copy = qx with horizontal(region[i_start + 1, :]): - qx = qx_edge_west2(qin, dxa, qx_copy) + qx = qx_edge_west2(qin, dxa, qx) with horizontal(region[i_end + 1, :]): qx = qx_edge_east(qin, dxa) - qx_copy = qx with horizontal(region[i_end, :]): - qx = qx_edge_east2(qin, dxa, qx_copy) + qx = qx_edge_east2(qin, dxa, qx) # ppm_volume_mean_y qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) - qy_copy = qy with horizontal(region[:, j_start]): qy = qy_edge_south(qin, dya) - qy_copy = qy with horizontal(region[:, j_start + 1]): - qy = qy_edge_south2(qin, dya, qy_copy) + qy = qy_edge_south2(qin, dya, qy) with horizontal(region[:, j_end + 1]): qy = qy_edge_north(qin, dya) - qy_copy = qy with horizontal(region[:, j_end]): - qy = qy_edge_north2(qin, dya, qy_copy) + qy = qy_edge_north2(qin, dya, qy) + qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + qxx[0, 1, 0]) @@ -406,8 +399,7 @@ def __init__( self._tmp_qy = utils.make_storage_from_shape(shape) self._tmp_qxx = utils.make_storage_from_shape(shape) self._tmp_qyy = utils.make_storage_from_shape(shape) - self._tmp_qx_copy = utils.make_storage_from_shape(shape) - self._tmp_qy_copy = utils.make_storage_from_shape(shape) + if nk is None: nk = self.grid.npz - kstart corner_domain = (1, 1, nk) @@ -527,9 +519,7 @@ def __call__(self, qin: FloatField, qout: FloatField): qin, qout, self._tmp_qx, - self._tmp_qx_copy, self._tmp_qy, - self._tmp_qy_copy, self._tmp_qxx, self._tmp_qyy, self.grid.dxa, From 5f4a5f53967e6f5494ef678e65a484744ddad5ab Mon Sep 17 00:00:00 2001 From: Rhea George Date: Fri, 7 May 2021 09:05:37 -0700 Subject: [PATCH 04/45] Revert "remerge ppm_volume_mean_x and y back into a2b_interpolation" This reverts commit 40add50a6bbe43254436482484742d6b598c3573. --- fv3core/stencils/a2b_ord4.py | 65 +++++++++++++++++++++++++++++------- 1 file changed, 53 insertions(+), 12 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 0b817e1d6..6fc023866 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -235,20 +235,14 @@ def lagrange_x_func(qy): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) -def a2b_interpolation( +def ppm_volume_mean_x( qin: FloatField, - qout: FloatField, qx: FloatField, - qy: FloatField, - qxx: FloatField, - qyy: FloatField, dxa: FloatFieldIJ, - dya: FloatFieldIJ ): - from __externals__ import i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start with computation(PARALLEL), interval(...): - # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) with horizontal(region[i_start, :]): qx = qx_edge_west(qin, dxa) @@ -258,7 +252,16 @@ def a2b_interpolation( qx = qx_edge_east(qin, dxa) with horizontal(region[i_end, :]): qx = qx_edge_east2(qin, dxa, qx) - # ppm_volume_mean_y + + +def ppm_volume_mean_y( + qin: FloatField, + qy: FloatField, + dya: FloatFieldIJ, +): + from __externals__ import j_end, j_start + + with computation(PARALLEL), interval(...): qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) with horizontal(region[:, j_start]): qy = qy_edge_south(qin, dya) @@ -269,6 +272,17 @@ def a2b_interpolation( with horizontal(region[:, j_end]): qy = qy_edge_north2(qin, dya, qy) + +def a2b_interpolation( + qout: FloatField, + qx: FloatField, + qy: FloatField, + qxx: FloatField, + qyy: FloatField, +): + from __externals__ import i_end, i_start, j_end, j_start + + with computation(PARALLEL), interval(...): qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + qxx[0, 1, 0]) @@ -451,6 +465,26 @@ def __init__( self._qout_y_edge_north = FrozenStencil( qout_y_edge, origin=(is2, self.grid.je + 1, kstart), domain=(di2, 1, nk) ) + origin_x = (self.grid.is_, self.grid.js - 2, kstart) + domain_x = (self.grid.nic + 1, self.grid.njc + 4, nk) + ax_offsets_x = axis_offsets( + self.grid, + origin_x, + domain_x, + ) + self._ppm_volume_mean_x_stencil = FrozenStencil( + ppm_volume_mean_x, externals=ax_offsets_x, origin=origin_x, domain=domain_x + ) + origin_y = (self.grid.is_ - 2, self.grid.js, kstart) + domain_y = (self.grid.nic + 4, self.grid.njc + 1, nk) + ax_offsets_y = axis_offsets( + self.grid, + origin_y, + domain_y, + ) + self._ppm_volume_mean_y_stencil = FrozenStencil( + ppm_volume_mean_y, externals=ax_offsets_y, origin=origin_y, domain=domain_y + ) js = self.grid.js + 1 if self.grid.south_edge else self.grid.js je = self.grid.je if self.grid.north_edge else self.grid.je + 1 @@ -515,15 +549,22 @@ def __call__(self, qin: FloatField, qout: FloatField): ) self._compute_qout_edges(qin, qout) - self._a2b_interpolation_stencil( + self._ppm_volume_mean_x_stencil( qin, + self._tmp_qx, + self.grid.dxa, + ) + self._ppm_volume_mean_y_stencil( + qin, + self._tmp_qy, + self.grid.dya, + ) + self._a2b_interpolation_stencil( qout, self._tmp_qx, self._tmp_qy, self._tmp_qxx, self._tmp_qyy, - self.grid.dxa, - self.grid.dya, ) if self.replace: self._copy_stencil( From 674e9195675dc37aeb773b47e89a36b07326cb8f Mon Sep 17 00:00:00 2001 From: Rhea George Date: Fri, 7 May 2021 09:11:45 -0700 Subject: [PATCH 05/45] merge ppm_volume_mean stencils into a2b_interpolation --- fv3core/stencils/a2b_ord4.py | 65 +++++++----------------------------- 1 file changed, 12 insertions(+), 53 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 6fc023866..0b817e1d6 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -235,14 +235,20 @@ def lagrange_x_func(qy): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) -def ppm_volume_mean_x( +def a2b_interpolation( qin: FloatField, + qout: FloatField, qx: FloatField, + qy: FloatField, + qxx: FloatField, + qyy: FloatField, dxa: FloatFieldIJ, + dya: FloatFieldIJ ): - from __externals__ import i_end, i_start + from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): + # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) with horizontal(region[i_start, :]): qx = qx_edge_west(qin, dxa) @@ -252,16 +258,7 @@ def ppm_volume_mean_x( qx = qx_edge_east(qin, dxa) with horizontal(region[i_end, :]): qx = qx_edge_east2(qin, dxa, qx) - - -def ppm_volume_mean_y( - qin: FloatField, - qy: FloatField, - dya: FloatFieldIJ, -): - from __externals__ import j_end, j_start - - with computation(PARALLEL), interval(...): + # ppm_volume_mean_y qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) with horizontal(region[:, j_start]): qy = qy_edge_south(qin, dya) @@ -272,17 +269,6 @@ def ppm_volume_mean_y( with horizontal(region[:, j_end]): qy = qy_edge_north2(qin, dya, qy) - -def a2b_interpolation( - qout: FloatField, - qx: FloatField, - qy: FloatField, - qxx: FloatField, - qyy: FloatField, -): - from __externals__ import i_end, i_start, j_end, j_start - - with computation(PARALLEL), interval(...): qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + qxx[0, 1, 0]) @@ -465,26 +451,6 @@ def __init__( self._qout_y_edge_north = FrozenStencil( qout_y_edge, origin=(is2, self.grid.je + 1, kstart), domain=(di2, 1, nk) ) - origin_x = (self.grid.is_, self.grid.js - 2, kstart) - domain_x = (self.grid.nic + 1, self.grid.njc + 4, nk) - ax_offsets_x = axis_offsets( - self.grid, - origin_x, - domain_x, - ) - self._ppm_volume_mean_x_stencil = FrozenStencil( - ppm_volume_mean_x, externals=ax_offsets_x, origin=origin_x, domain=domain_x - ) - origin_y = (self.grid.is_ - 2, self.grid.js, kstart) - domain_y = (self.grid.nic + 4, self.grid.njc + 1, nk) - ax_offsets_y = axis_offsets( - self.grid, - origin_y, - domain_y, - ) - self._ppm_volume_mean_y_stencil = FrozenStencil( - ppm_volume_mean_y, externals=ax_offsets_y, origin=origin_y, domain=domain_y - ) js = self.grid.js + 1 if self.grid.south_edge else self.grid.js je = self.grid.je if self.grid.north_edge else self.grid.je + 1 @@ -549,22 +515,15 @@ def __call__(self, qin: FloatField, qout: FloatField): ) self._compute_qout_edges(qin, qout) - self._ppm_volume_mean_x_stencil( - qin, - self._tmp_qx, - self.grid.dxa, - ) - self._ppm_volume_mean_y_stencil( - qin, - self._tmp_qy, - self.grid.dya, - ) self._a2b_interpolation_stencil( + qin, qout, self._tmp_qx, self._tmp_qy, self._tmp_qxx, self._tmp_qyy, + self.grid.dxa, + self.grid.dya, ) if self.replace: self._copy_stencil( From ed9b5e110de41c230ccab5cb1cd26e329f48b61c Mon Sep 17 00:00:00 2001 From: Rhea George Date: Fri, 7 May 2021 09:16:05 -0700 Subject: [PATCH 06/45] using an extra variable --- fv3core/stencils/a2b_ord4.py | 55 +++++++++++++++++++++++++++++------- 1 file changed, 45 insertions(+), 10 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 0b817e1d6..a6d719133 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -239,36 +239,68 @@ def a2b_interpolation( qin: FloatField, qout: FloatField, qx: FloatField, + qx_copy: FloatField, qy: FloatField, + qy_copy: FloatField, qxx: FloatField, qyy: FloatField, dxa: FloatFieldIJ, - dya: FloatFieldIJ + dya: FloatFieldIJ, ): from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) + qx_copy = qx with horizontal(region[i_start, :]): - qx = qx_edge_west(qin, dxa) + #qx = qx_edge_west(qin, dxa) + qx = 0.5 * ( + ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) + + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + dxa[-2, 0] / dxa[-1, 0]) + ) + + qx_copy = qx with horizontal(region[i_start + 1, :]): - qx = qx_edge_west2(qin, dxa, qx) + #qx = qx_edge_west2(qin, dxa, qx_copy) + qx = ( + 3.0 * (dxa / dxa[-1, 0] * qin[-1, 0, 0] + qin) - (dxa / dxa[-1, 0] * qx_copy[-1, 0, 0] + qx_copy[1, 0, 0]) + ) / (2.0 + 2.0 * dxa / dxa[-1, 0]) with horizontal(region[i_end + 1, :]): - qx = qx_edge_east(qin, dxa) + #qx = qx_edge_east(qin, dxa) + qx = 0.5 * ( + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + dxa[-2, 0] / dxa[-1, 0]) + + ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) + ) + + qx_copy = qx with horizontal(region[i_end, :]): - qx = qx_edge_east2(qin, dxa, qx) + #qx = qx_edge_east2(qin, dxa, qx_copy) + qx = ( + 3.0 * (qin[-1, 0, 0] + dxa[-1, 0] / dxa * qin) - (dxa[-1, 0] / dxa * qx_copy[1, 0, 0] + qx_copy[-1, 0, 0]) + ) / (2.0 + 2.0 * dxa[-1, 0] / dxa) + + # ppm_volume_mean_y qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) + qy_copy = qy with horizontal(region[:, j_start]): - qy = qy_edge_south(qin, dya) + #qy = qy_edge_south(qin, dya) + g_in = dya[0, 1] / dya + g_ou = dya[0, -2] / dya[0, -1] + qy= 0.5 * ( + ((2.0 + g_in) * qin - qin[0, 1, 0]) / (1.0 + g_in) + + ((2.0 + g_ou) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + g_ou) + ) + + qy_copy = qy with horizontal(region[:, j_start + 1]): - qy = qy_edge_south2(qin, dya, qy) + qy = qy_edge_south2(qin, dya, qy_copy) with horizontal(region[:, j_end + 1]): qy = qy_edge_north(qin, dya) + qy_copy = qy with horizontal(region[:, j_end]): - qy = qy_edge_north2(qin, dya, qy) - + qy = qy_edge_north2(qin, dya, qy_copy) qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + qxx[0, 1, 0]) @@ -399,7 +431,8 @@ def __init__( self._tmp_qy = utils.make_storage_from_shape(shape) self._tmp_qxx = utils.make_storage_from_shape(shape) self._tmp_qyy = utils.make_storage_from_shape(shape) - + self._tmp_qx_copy = utils.make_storage_from_shape(shape) + self._tmp_qy_copy = utils.make_storage_from_shape(shape) if nk is None: nk = self.grid.npz - kstart corner_domain = (1, 1, nk) @@ -519,7 +552,9 @@ def __call__(self, qin: FloatField, qout: FloatField): qin, qout, self._tmp_qx, + self._tmp_qx_copy, self._tmp_qy, + self._tmp_qy_copy, self._tmp_qxx, self._tmp_qyy, self.grid.dxa, From 46f1b9352dc75b3498b261e496a1f3f6da42f0c3 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 11 May 2021 20:49:20 -0700 Subject: [PATCH 07/45] intermediate qout edges commit --- fv3core/stencils/a2b_ord4.py | 150 +++++++++++++++++++++++++++++------ 1 file changed, 125 insertions(+), 25 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index a6d719133..eb1702a8f 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -13,11 +13,12 @@ import fv3core._config as spec import fv3core.utils.gt4py_utils as utils -from fv3core.decorators import FrozenStencil +from fv3core.decorators import FrozenStencil,gtstencil from fv3core.stencils.basic_operations import copy_defn from fv3core.utils import axis_offsets from fv3core.utils.typing import FloatField, FloatFieldI, FloatFieldIJ - +import gt4py +from typing import Tuple # comact 4-pt cubic interpolation c1 = 2.0 / 3.0 @@ -37,7 +38,7 @@ def great_circle_dist(p1a, p1b, p2a, p2b): tb = sin((p1b - p2b) / 2.0) ** 2.0 ta = sin((p1a - p2a) / 2.0) ** 2.0 return asin(sqrt(tb + cos(p1b) * cos(p2b) * ta)) * 2.0 - +#return asin(sqrt(sin((p1b - p2b) / 2.0) ** 2.0 + cos(p1b) * cos(p2b) * sin((p1a - p2a) / 2.0) ** 2.0)) * 2.0 @gtscript.function def extrap_corner( @@ -50,11 +51,25 @@ def extrap_corner( qa, qb, ): + #x1 = asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0 + #x2 = asin(sqrt(sin((p2b - p0b) / 2.0) ** 2.0 + cos(p2b) * cos(p0b) * sin((p2a - p0a) / 2.0) ** 2.0)) * 2.0 + #return qa + (asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) / (asin(sqrt(sin((p2b - p0b) / 2.0) ** 2.0 + cos(p2b) * cos(p0b) * sin((p2a - p0a) / 2.0) ** 2.0)) * 2.0 -asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) * (qa - qb) x1 = great_circle_dist(p1a, p1b, p0a, p0b) x2 = great_circle_dist(p2a, p2b, p0a, p0b) return qa + x1 / (x2 - x1) * (qa - qb) - - + +""" +def _a2b_corners(qin, qout, agrid1, agrid2, bgrid1, bgrid2): + with computation(PARALLEL), interval(...): + from __externals__ import i_start, j_start + + with horizontal(region[i_start, j_start]): + qout = _sw_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) + qout = _se_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) + qout = _ne_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) + qout = _nw_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) + return qout +""" def _sw_corner( qin: FloatField, qout: FloatField, @@ -250,6 +265,7 @@ def a2b_interpolation( from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): + # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) qx_copy = qx @@ -405,6 +421,38 @@ def qy_edge_north2(qin: FloatField, dya: FloatFieldIJ, qy: FloatField): 3.0 * (qin[0, -1, 0] + g_in * qin) - (g_in * qy[0, 1, 0] + qy[0, -1, 0]) ) / (2.0 + 2.0 * g_in) +def qout_edges_x(q2: FloatField, + qin: FloatField, + qout: FloatField, + dxa: FloatFieldIJ, + edge_w: FloatFieldIJ, + edge_e: FloatFieldIJ, +): + from __externals__ import i_end, i_start + + with computation(PARALLEL), interval(...): + q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + with horizontal(region[i_start, :]): + qout = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 + with horizontal(region[i_end + 1, :]): + qout = edge_e * q2[0, -1, 0] + (1.0 - edge_e) * q2 + +def qout_edges_y( + q1: FloatField, + qin: FloatField, + qout: FloatField, + dya: FloatFieldIJ, + edge_s: FloatFieldI, + edge_n: FloatFieldI, +): + from __externals__ import j_end, j_start + + with computation(PARALLEL), interval(...): + q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + with horizontal(region[:, j_start]): + qout = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 + with horizontal(region[:, j_end + 1]): + qout = edge_n * q1[-1, 0, 0] + (1.0 - edge_n) * q1 class AGrid2BGridFourthOrder: """ @@ -457,38 +505,65 @@ def __init__( origin=(self.grid.is_, self.grid.je + 1, kstart), domain=corner_domain, ) - js2 = self.grid.js + 1 if self.grid.south_edge else self.grid.js - je1 = self.grid.je if self.grid.north_edge else self.grid.je + 1 - dj2 = je1 - js2 + 1 - + #js2 = self.grid.js + 1 if self.grid.south_edge else self.grid.js + #je1 = self.grid.je if self.grid.north_edge else self.grid.je + 1 + #dj2 = je1 - js2 + 1 + js = self.grid.js + 1 if self.grid.south_edge else self.grid.js + je = self.grid.je if self.grid.north_edge else self.grid.je + 1 + is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ + ie = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 + dj = je - js + 1 + di = ie - is_ + 1 # edge_w is singleton in the I-dimension to work around gt4py not yet # supporting J-fields. As a result, the origin has to be zero for # edge_w, anything higher is outside its index range + self._qout_x_edge_west = FrozenStencil( qout_x_edge, - origin={"_all_": (self.grid.is_, js2, kstart), "edge_w": (0, js2)}, - domain=(1, dj2, nk), + origin={"_all_": (self.grid.is_, js, kstart), "edge_w": (0, js)}, + domain=(1, dj, nk), ) self._qout_x_edge_east = FrozenStencil( qout_x_edge, - origin={"_all_": (self.grid.ie + 1, js2, kstart), "edge_w": (0, js2)}, - domain=(1, dj2, nk), + origin={"_all_": (self.grid.ie + 1, js, kstart), "edge_w": (0, js)}, + domain=(1, dj, nk), ) - - is2 = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ - ie1 = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 - di2 = ie1 - is2 + 1 + + + self._edge_e = self._j_storage_repeat_over_i(self.grid.edge_e, shape[0:2]) + self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) + x_edge_origin = (self.grid.is_, js, kstart) + x_edge_domain = (self.grid.nic +2 , dj, nk) + x_offsets = axis_offsets(self.grid, x_edge_origin, x_edge_domain) + self._qout_edges_x = FrozenStencil(qout_edges_x,externals=x_offsets, + origin=x_edge_origin, + + domain=x_edge_domain, + ) + self._tmp_q2 = utils.make_storage_from_shape(shape) + self._tmp_q1 = utils.make_storage_from_shape(shape) + #is2 = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ + #ie1 = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 + #di2 = ie1 - is2 + 1 + y_edge_origin = (is_, self.grid.js, kstart) + y_edge_domain = (di, self.grid.njc + 2, nk) + y_offsets = axis_offsets(self.grid, y_edge_origin, y_edge_domain) + self._qout_edges_y = FrozenStencil(qout_edges_y,externals=y_offsets, + origin=y_edge_origin, + domain=y_edge_domain, + ) + self._qout_y_edge_south = FrozenStencil( - qout_y_edge, origin=(is2, self.grid.js, kstart), domain=(di2, 1, nk) + qout_y_edge, origin=(is_, self.grid.js, kstart), domain=(di, 1, nk) ) self._qout_y_edge_north = FrozenStencil( - qout_y_edge, origin=(is2, self.grid.je + 1, kstart), domain=(di2, 1, nk) + qout_y_edge, origin=(is_, self.grid.je + 1, kstart), domain=(di, 1, nk) ) - - js = self.grid.js + 1 if self.grid.south_edge else self.grid.js - je = self.grid.je if self.grid.north_edge else self.grid.je + 1 - is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ - ie = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 + + #js = self.grid.js + 1 if self.grid.south_edge else self.grid.js + #je = self.grid.je if self.grid.north_edge else self.grid.je + 1 + #is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ + #ie = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 origin = (is_, js, kstart) domain = (ie - is_ + 1, je - js + 1, nk) ax_offsets = axis_offsets( @@ -504,6 +579,13 @@ def __init__( origin=(self.grid.is_, self.grid.js, kstart), domain=(self.grid.nic + 1, self.grid.njc + 1, nk), ) + # TODO + # within regions, the edge_w and edge_w variables that are singleton in the + # I dimension error, workaround is repeating the data, but the longterm + # fix should happen in regions + def _j_storage_repeat_over_i(self, grid_array: gt4py.storage.Storage, shape: Tuple[int, int]): + dup = utils.repeat(grid_array, shape[1], axis=0) + return utils.make_storage_data(dup, shape, (0, 0)) def __call__(self, qin: FloatField, qout: FloatField): """Converts qin from A-grid to B-grid in qout. @@ -567,6 +649,16 @@ def __call__(self, qin: FloatField, qout: FloatField): ) def _compute_qout_edges(self, qin: FloatField, qout: FloatField): + + self._qout_edges_x(self._tmp_q2, + qin, + qout, + self.grid.dxa,self._edge_w,self._edge_e + ) + + """ + + if self.grid.west_edge: self._qout_x_edge_west( qin, @@ -581,7 +673,14 @@ def _compute_qout_edges(self, qin: FloatField, qout: FloatField): self.grid.edge_e, qout, ) - + """ + self._qout_edges_y( + self._tmp_q1, + qin, + qout, + self.grid.dya, self.grid.edge_s, self.grid.edge_n, + ) + """ if self.grid.south_edge: self._qout_y_edge_south( qin, @@ -596,3 +695,4 @@ def _compute_qout_edges(self, qin: FloatField, qout: FloatField): self.grid.edge_n, qout, ) + """ From 5e28d5b8a075feea3c60038c7818136d38b9cf6e Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 11 May 2021 20:51:39 -0700 Subject: [PATCH 08/45] qout edges cleanup --- fv3core/stencils/a2b_ord4.py | 56 +++--------------------------------- 1 file changed, 4 insertions(+), 52 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index eb1702a8f..7f48a563d 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -505,9 +505,7 @@ def __init__( origin=(self.grid.is_, self.grid.je + 1, kstart), domain=corner_domain, ) - #js2 = self.grid.js + 1 if self.grid.south_edge else self.grid.js - #je1 = self.grid.je if self.grid.north_edge else self.grid.je + 1 - #dj2 = je1 - js2 + 1 + js = self.grid.js + 1 if self.grid.south_edge else self.grid.js je = self.grid.je if self.grid.north_edge else self.grid.je + 1 is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ @@ -537,14 +535,11 @@ def __init__( x_offsets = axis_offsets(self.grid, x_edge_origin, x_edge_domain) self._qout_edges_x = FrozenStencil(qout_edges_x,externals=x_offsets, origin=x_edge_origin, - domain=x_edge_domain, ) self._tmp_q2 = utils.make_storage_from_shape(shape) self._tmp_q1 = utils.make_storage_from_shape(shape) - #is2 = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ - #ie1 = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 - #di2 = ie1 - is2 + 1 + y_edge_origin = (is_, self.grid.js, kstart) y_edge_domain = (di, self.grid.njc + 2, nk) y_offsets = axis_offsets(self.grid, y_edge_origin, y_edge_domain) @@ -552,18 +547,7 @@ def __init__( origin=y_edge_origin, domain=y_edge_domain, ) - - self._qout_y_edge_south = FrozenStencil( - qout_y_edge, origin=(is_, self.grid.js, kstart), domain=(di, 1, nk) - ) - self._qout_y_edge_north = FrozenStencil( - qout_y_edge, origin=(is_, self.grid.je + 1, kstart), domain=(di, 1, nk) - ) - #js = self.grid.js + 1 if self.grid.south_edge else self.grid.js - #je = self.grid.je if self.grid.north_edge else self.grid.je + 1 - #is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ - #ie = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 origin = (is_, js, kstart) domain = (ie - is_ + 1, je - js + 1, nk) ax_offsets = axis_offsets( @@ -656,43 +640,11 @@ def _compute_qout_edges(self, qin: FloatField, qout: FloatField): self.grid.dxa,self._edge_w,self._edge_e ) - """ - - - if self.grid.west_edge: - self._qout_x_edge_west( - qin, - self.grid.dxa, - self.grid.edge_w, - qout, - ) - if self.grid.east_edge: - self._qout_x_edge_east( - qin, - self.grid.dxa, - self.grid.edge_e, - qout, - ) - """ + self._qout_edges_y( self._tmp_q1, qin, qout, self.grid.dya, self.grid.edge_s, self.grid.edge_n, ) - """ - if self.grid.south_edge: - self._qout_y_edge_south( - qin, - self.grid.dya, - self.grid.edge_s, - qout, - ) - if self.grid.north_edge: - self._qout_y_edge_north( - qin, - self.grid.dya, - self.grid.edge_n, - qout, - ) - """ + From 17f45f9f96701db91f7897a519fa5c0b48c660a4 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 12 May 2021 14:34:04 -0700 Subject: [PATCH 09/45] put edges back into a2b interpolation --- fv3core/stencils/a2b_ord4.py | 88 +++++++++++++----------------------- 1 file changed, 32 insertions(+), 56 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 7f48a563d..03c341646 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -259,13 +259,30 @@ def a2b_interpolation( qy_copy: FloatField, qxx: FloatField, qyy: FloatField, + q1: FloatField, + q2: FloatField, dxa: FloatFieldIJ, dya: FloatFieldIJ, + edge_w: FloatFieldIJ, + edge_e: FloatFieldIJ, + edge_s: FloatFieldI, + edge_n: FloatFieldI, ): from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): - + #qout x edges + q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + with horizontal(region[i_start, j_start + 1 : j_end + 1]): + qout = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 + with horizontal(region[i_end + 1, j_start + 1 : j_end + 1]): + qout = edge_e * q2[0, -1, 0] + (1.0 - edge_e) * q2 + # qout y edges + q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + with horizontal(region[i_start + 1 : i_end + 1, j_start]): + qout = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 + with horizontal(region[i_start + 1 : i_end + 1, j_end + 1]): + qout = edge_n * q1[-1, 0, 0] + (1.0 - edge_n) * q1 # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) qx_copy = qx @@ -481,6 +498,9 @@ def __init__( self._tmp_qyy = utils.make_storage_from_shape(shape) self._tmp_qx_copy = utils.make_storage_from_shape(shape) self._tmp_qy_copy = utils.make_storage_from_shape(shape) + self._tmp_q2 = utils.make_storage_from_shape(shape) + self._tmp_q1 = utils.make_storage_from_shape(shape) + if nk is None: nk = self.grid.npz - kstart corner_domain = (1, 1, nk) @@ -510,43 +530,10 @@ def __init__( je = self.grid.je if self.grid.north_edge else self.grid.je + 1 is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ ie = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 - dj = je - js + 1 - di = ie - is_ + 1 - # edge_w is singleton in the I-dimension to work around gt4py not yet - # supporting J-fields. As a result, the origin has to be zero for - # edge_w, anything higher is outside its index range - - self._qout_x_edge_west = FrozenStencil( - qout_x_edge, - origin={"_all_": (self.grid.is_, js, kstart), "edge_w": (0, js)}, - domain=(1, dj, nk), - ) - self._qout_x_edge_east = FrozenStencil( - qout_x_edge, - origin={"_all_": (self.grid.ie + 1, js, kstart), "edge_w": (0, js)}, - domain=(1, dj, nk), - ) - - + self._edge_e = self._j_storage_repeat_over_i(self.grid.edge_e, shape[0:2]) - self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) - x_edge_origin = (self.grid.is_, js, kstart) - x_edge_domain = (self.grid.nic +2 , dj, nk) - x_offsets = axis_offsets(self.grid, x_edge_origin, x_edge_domain) - self._qout_edges_x = FrozenStencil(qout_edges_x,externals=x_offsets, - origin=x_edge_origin, - domain=x_edge_domain, - ) - self._tmp_q2 = utils.make_storage_from_shape(shape) - self._tmp_q1 = utils.make_storage_from_shape(shape) - - y_edge_origin = (is_, self.grid.js, kstart) - y_edge_domain = (di, self.grid.njc + 2, nk) - y_offsets = axis_offsets(self.grid, y_edge_origin, y_edge_domain) - self._qout_edges_y = FrozenStencil(qout_edges_y,externals=y_offsets, - origin=y_edge_origin, - domain=y_edge_domain, - ) + self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) + origin = (is_, js, kstart) domain = (ie - is_ + 1, je - js + 1, nk) @@ -566,7 +553,9 @@ def __init__( # TODO # within regions, the edge_w and edge_w variables that are singleton in the # I dimension error, workaround is repeating the data, but the longterm - # fix should happen in regions + # fix should happen in regions + # Setting the origin to 0 does not work when the domain size > 1 + # in that dimension def _j_storage_repeat_over_i(self, grid_array: gt4py.storage.Storage, shape: Tuple[int, int]): dup = utils.repeat(grid_array, shape[1], axis=0) return utils.make_storage_data(dup, shape, (0, 0)) @@ -613,7 +602,7 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid2, ) - self._compute_qout_edges(qin, qout) + #self._compute_qout_edges(qin, qout) self._a2b_interpolation_stencil( qin, qout, @@ -623,8 +612,10 @@ def __call__(self, qin: FloatField, qout: FloatField): self._tmp_qy_copy, self._tmp_qxx, self._tmp_qyy, + self._tmp_q1, + self._tmp_q2, self.grid.dxa, - self.grid.dya, + self.grid.dya,self._edge_w,self._edge_e,self.grid.edge_s, self.grid.edge_n, ) if self.replace: self._copy_stencil( @@ -632,19 +623,4 @@ def __call__(self, qin: FloatField, qout: FloatField): qin, ) - def _compute_qout_edges(self, qin: FloatField, qout: FloatField): - - self._qout_edges_x(self._tmp_q2, - qin, - qout, - self.grid.dxa,self._edge_w,self._edge_e - ) - - - self._qout_edges_y( - self._tmp_q1, - qin, - qout, - self.grid.dya, self.grid.edge_s, self.grid.edge_n, - ) - + From 02ca216cb7db8800a68e7b9474046d06b750ec1d Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 12 May 2021 15:10:27 -0700 Subject: [PATCH 10/45] put replace into a2b interpolation --- fv3core/stencils/a2b_ord4.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 03c341646..0fd8f2a8c 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -268,7 +268,7 @@ def a2b_interpolation( edge_s: FloatFieldI, edge_n: FloatFieldI, ): - from __externals__ import i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start, j_end, j_start, replace with computation(PARALLEL), interval(...): #qout x edges @@ -325,7 +325,6 @@ def a2b_interpolation( ((2.0 + g_in) * qin - qin[0, 1, 0]) / (1.0 + g_in) + ((2.0 + g_ou) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + g_ou) ) - qy_copy = qy with horizontal(region[:, j_start + 1]): qy = qy_edge_south2(qin, dya, qy_copy) @@ -344,8 +343,10 @@ def a2b_interpolation( qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[-1, 0, 0] + qyy[1, 0, 0]) with horizontal(region[i_end, :]): qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[1, 0, 0] + qyy[-1, 0, 0]) - qout = 0.5 * (qxx + qyy) - + with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): + qout = 0.5 * (qxx + qyy) + if __INLINED(replace): + qin = qout def qout_x_edge( qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ, qout: FloatField @@ -488,7 +489,6 @@ def __init__( """ assert namelist.grid_type < 3 self.grid = spec.grid - self.replace = replace shape = self.grid.domain_shape_full(add=(1, 1, 1)) full_origin = (self.grid.isd, self.grid.jsd, kstart) @@ -526,10 +526,10 @@ def __init__( domain=corner_domain, ) - js = self.grid.js + 1 if self.grid.south_edge else self.grid.js - je = self.grid.je if self.grid.north_edge else self.grid.je + 1 - is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ - ie = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 + js = self.grid.js #self.grid.js + 1 if self.grid.south_edge else self.grid.js + je = self.grid.je + 1 #self.grid.je if self.grid.north_edge else self.grid.je + 1 + is_ = self.grid.is_ #self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ + ie = self.grid.ie + 1 #self.grid.ie if self.grid.east_edge else self.grid.ie + 1 self._edge_e = self._j_storage_repeat_over_i(self.grid.edge_e, shape[0:2]) self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) @@ -543,7 +543,7 @@ def __init__( domain, ) self._a2b_interpolation_stencil = FrozenStencil( - a2b_interpolation, externals=ax_offsets, origin=origin, domain=domain + a2b_interpolation, externals={"replace": replace, **ax_offsets}, origin=origin, domain=domain ) self._copy_stencil = FrozenStencil( copy_defn, @@ -617,10 +617,10 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.dxa, self.grid.dya,self._edge_w,self._edge_e,self.grid.edge_s, self.grid.edge_n, ) - if self.replace: - self._copy_stencil( - qout, - qin, - ) + #if self.replace: + # self._copy_stencil( + # qout, + # qin, + # ) From ce5b61241a098823011824c587679a3d9c16ef15 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 12 May 2021 16:04:40 -0700 Subject: [PATCH 11/45] Revert "put replace into a2b interpolation" This reverts commit 02ca216cb7db8800a68e7b9474046d06b750ec1d. --- fv3core/stencils/a2b_ord4.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 0fd8f2a8c..03c341646 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -268,7 +268,7 @@ def a2b_interpolation( edge_s: FloatFieldI, edge_n: FloatFieldI, ): - from __externals__ import i_end, i_start, j_end, j_start, replace + from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): #qout x edges @@ -325,6 +325,7 @@ def a2b_interpolation( ((2.0 + g_in) * qin - qin[0, 1, 0]) / (1.0 + g_in) + ((2.0 + g_ou) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + g_ou) ) + qy_copy = qy with horizontal(region[:, j_start + 1]): qy = qy_edge_south2(qin, dya, qy_copy) @@ -343,10 +344,8 @@ def a2b_interpolation( qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[-1, 0, 0] + qyy[1, 0, 0]) with horizontal(region[i_end, :]): qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[1, 0, 0] + qyy[-1, 0, 0]) - with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): - qout = 0.5 * (qxx + qyy) - if __INLINED(replace): - qin = qout + qout = 0.5 * (qxx + qyy) + def qout_x_edge( qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ, qout: FloatField @@ -489,6 +488,7 @@ def __init__( """ assert namelist.grid_type < 3 self.grid = spec.grid + self.replace = replace shape = self.grid.domain_shape_full(add=(1, 1, 1)) full_origin = (self.grid.isd, self.grid.jsd, kstart) @@ -526,10 +526,10 @@ def __init__( domain=corner_domain, ) - js = self.grid.js #self.grid.js + 1 if self.grid.south_edge else self.grid.js - je = self.grid.je + 1 #self.grid.je if self.grid.north_edge else self.grid.je + 1 - is_ = self.grid.is_ #self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ - ie = self.grid.ie + 1 #self.grid.ie if self.grid.east_edge else self.grid.ie + 1 + js = self.grid.js + 1 if self.grid.south_edge else self.grid.js + je = self.grid.je if self.grid.north_edge else self.grid.je + 1 + is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ + ie = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 self._edge_e = self._j_storage_repeat_over_i(self.grid.edge_e, shape[0:2]) self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) @@ -543,7 +543,7 @@ def __init__( domain, ) self._a2b_interpolation_stencil = FrozenStencil( - a2b_interpolation, externals={"replace": replace, **ax_offsets}, origin=origin, domain=domain + a2b_interpolation, externals=ax_offsets, origin=origin, domain=domain ) self._copy_stencil = FrozenStencil( copy_defn, @@ -617,10 +617,10 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.dxa, self.grid.dya,self._edge_w,self._edge_e,self.grid.edge_s, self.grid.edge_n, ) - #if self.replace: - # self._copy_stencil( - # qout, - # qin, - # ) + if self.replace: + self._copy_stencil( + qout, + qin, + ) From 728b7ce986e33294246978749acc5c8576a9fe23 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 12 May 2021 21:36:15 -0700 Subject: [PATCH 12/45] Revert "put edges back into a2b interpolation" This reverts commit 17f45f9f96701db91f7897a519fa5c0b48c660a4. --- fv3core/stencils/a2b_ord4.py | 88 +++++++++++++++++++++++------------- 1 file changed, 56 insertions(+), 32 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 03c341646..7f48a563d 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -259,30 +259,13 @@ def a2b_interpolation( qy_copy: FloatField, qxx: FloatField, qyy: FloatField, - q1: FloatField, - q2: FloatField, dxa: FloatFieldIJ, dya: FloatFieldIJ, - edge_w: FloatFieldIJ, - edge_e: FloatFieldIJ, - edge_s: FloatFieldI, - edge_n: FloatFieldI, ): from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): - #qout x edges - q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - with horizontal(region[i_start, j_start + 1 : j_end + 1]): - qout = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 - with horizontal(region[i_end + 1, j_start + 1 : j_end + 1]): - qout = edge_e * q2[0, -1, 0] + (1.0 - edge_e) * q2 - # qout y edges - q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - with horizontal(region[i_start + 1 : i_end + 1, j_start]): - qout = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 - with horizontal(region[i_start + 1 : i_end + 1, j_end + 1]): - qout = edge_n * q1[-1, 0, 0] + (1.0 - edge_n) * q1 + # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) qx_copy = qx @@ -498,9 +481,6 @@ def __init__( self._tmp_qyy = utils.make_storage_from_shape(shape) self._tmp_qx_copy = utils.make_storage_from_shape(shape) self._tmp_qy_copy = utils.make_storage_from_shape(shape) - self._tmp_q2 = utils.make_storage_from_shape(shape) - self._tmp_q1 = utils.make_storage_from_shape(shape) - if nk is None: nk = self.grid.npz - kstart corner_domain = (1, 1, nk) @@ -530,10 +510,43 @@ def __init__( je = self.grid.je if self.grid.north_edge else self.grid.je + 1 is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ ie = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 - + dj = je - js + 1 + di = ie - is_ + 1 + # edge_w is singleton in the I-dimension to work around gt4py not yet + # supporting J-fields. As a result, the origin has to be zero for + # edge_w, anything higher is outside its index range + + self._qout_x_edge_west = FrozenStencil( + qout_x_edge, + origin={"_all_": (self.grid.is_, js, kstart), "edge_w": (0, js)}, + domain=(1, dj, nk), + ) + self._qout_x_edge_east = FrozenStencil( + qout_x_edge, + origin={"_all_": (self.grid.ie + 1, js, kstart), "edge_w": (0, js)}, + domain=(1, dj, nk), + ) + + self._edge_e = self._j_storage_repeat_over_i(self.grid.edge_e, shape[0:2]) - self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) - + self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) + x_edge_origin = (self.grid.is_, js, kstart) + x_edge_domain = (self.grid.nic +2 , dj, nk) + x_offsets = axis_offsets(self.grid, x_edge_origin, x_edge_domain) + self._qout_edges_x = FrozenStencil(qout_edges_x,externals=x_offsets, + origin=x_edge_origin, + domain=x_edge_domain, + ) + self._tmp_q2 = utils.make_storage_from_shape(shape) + self._tmp_q1 = utils.make_storage_from_shape(shape) + + y_edge_origin = (is_, self.grid.js, kstart) + y_edge_domain = (di, self.grid.njc + 2, nk) + y_offsets = axis_offsets(self.grid, y_edge_origin, y_edge_domain) + self._qout_edges_y = FrozenStencil(qout_edges_y,externals=y_offsets, + origin=y_edge_origin, + domain=y_edge_domain, + ) origin = (is_, js, kstart) domain = (ie - is_ + 1, je - js + 1, nk) @@ -553,9 +566,7 @@ def __init__( # TODO # within regions, the edge_w and edge_w variables that are singleton in the # I dimension error, workaround is repeating the data, but the longterm - # fix should happen in regions - # Setting the origin to 0 does not work when the domain size > 1 - # in that dimension + # fix should happen in regions def _j_storage_repeat_over_i(self, grid_array: gt4py.storage.Storage, shape: Tuple[int, int]): dup = utils.repeat(grid_array, shape[1], axis=0) return utils.make_storage_data(dup, shape, (0, 0)) @@ -602,7 +613,7 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid2, ) - #self._compute_qout_edges(qin, qout) + self._compute_qout_edges(qin, qout) self._a2b_interpolation_stencil( qin, qout, @@ -612,10 +623,8 @@ def __call__(self, qin: FloatField, qout: FloatField): self._tmp_qy_copy, self._tmp_qxx, self._tmp_qyy, - self._tmp_q1, - self._tmp_q2, self.grid.dxa, - self.grid.dya,self._edge_w,self._edge_e,self.grid.edge_s, self.grid.edge_n, + self.grid.dya, ) if self.replace: self._copy_stencil( @@ -623,4 +632,19 @@ def __call__(self, qin: FloatField, qout: FloatField): qin, ) - + def _compute_qout_edges(self, qin: FloatField, qout: FloatField): + + self._qout_edges_x(self._tmp_q2, + qin, + qout, + self.grid.dxa,self._edge_w,self._edge_e + ) + + + self._qout_edges_y( + self._tmp_q1, + qin, + qout, + self.grid.dya, self.grid.edge_s, self.grid.edge_n, + ) + From e3437baa07192d8899958cc2b79442ee18a07593 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 10:00:28 -0700 Subject: [PATCH 13/45] sw corner experiment --- fv3core/stencils/a2b_ord4.py | 86 +++++++++++++++++++----------------- 1 file changed, 46 insertions(+), 40 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index b71f782e3..3299aa343 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -53,10 +53,10 @@ def extrap_corner( ): #x1 = asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0 #x2 = asin(sqrt(sin((p2b - p0b) / 2.0) ** 2.0 + cos(p2b) * cos(p0b) * sin((p2a - p0a) / 2.0) ** 2.0)) * 2.0 - #return qa + (asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) / (asin(sqrt(sin((p2b - p0b) / 2.0) ** 2.0 + cos(p2b) * cos(p0b) * sin((p2a - p0a) / 2.0) ** 2.0)) * 2.0 -asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) * (qa - qb) - x1 = great_circle_dist(p1a, p1b, p0a, p0b) - x2 = great_circle_dist(p2a, p2b, p0a, p0b) - return qa + x1 / (x2 - x1) * (qa - qb) + return qa + (asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) / (asin(sqrt(sin((p2b - p0b) / 2.0) ** 2.0 + cos(p2b) * cos(p0b) * sin((p2a - p0a) / 2.0) ** 2.0)) * 2.0 -asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) * (qa - qb) +# x1 = great_circle_dist(p1a, p1b, p0a, p0b) +# x2 = great_circle_dist(p2a, p2b, p0a, p0b) +# return qa + x1 / (x2 - x1) * (qa - qb) """ def _a2b_corners(qin, qout, agrid1, agrid2, bgrid1, bgrid2): @@ -78,40 +78,42 @@ def _sw_corner( bgrid1: FloatFieldIJ, bgrid2: FloatFieldIJ, ): - + from __externals__ import i_start, j_start with computation(PARALLEL), interval(...): - ec1 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[0, 0], - agrid2[0, 0], - agrid1[1, 1], - agrid2[1, 1], - qin[0, 0, 0], - qin[1, 1, 0], - ) - ec2 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[-1, 0], - agrid2[-1, 0], - agrid1[-2, 1], - agrid2[-2, 1], - qin[-1, 0, 0], - qin[-2, 1, 0], - ) - ec3 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[0, -1], - agrid2[0, -1], - agrid1[1, -2], - agrid2[1, -2], - qin[0, -1, 0], - qin[1, -2, 0], - ) + tmp = 0.0 + with horizontal(region[i_start, j_start]): + ec1 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[0, 0], + agrid2[0, 0], + agrid1[1, 1], + agrid2[1, 1], + qin[0, 0, 0], + qin[1, 1, 0], + ) + ec2 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[-1, 0], + agrid2[-1, 0], + agrid1[-2, 1], + agrid2[-2, 1], + qin[-1, 0, 0], + qin[-2, 1, 0], + ) + ec3 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[0, -1], + agrid2[0, -1], + agrid1[1, -2], + agrid2[1, -2], + qin[0, -1, 0], + qin[1, -2, 0], + ) - qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) + qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) def _nw_corner( @@ -484,11 +486,15 @@ def __init__( if nk is None: nk = self.grid.npz - kstart corner_domain = (1, 1, nk) - + full_ax_offsets = axis_offsets( + self.grid, + self.grid.full_origin(), + self.grid.domain_shape_full(), + ) self._sw_corner_stencil = FrozenStencil( - _sw_corner, - origin=(self.grid.is_, self.grid.js, kstart), - domain=corner_domain, + _sw_corner, externals=full_ax_offsets, + origin=(self.grid.isd, self.grid.jsd, kstart), + domain=(self.grid.nid, self.grid.njd, nk), ) self._nw_corner_stencil = FrozenStencil( _nw_corner, From 6707db32a03fb094bc9faf9775aa65d7813928e7 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 11:09:52 -0700 Subject: [PATCH 14/45] maybe can still use all the functions calling functions --- fv3core/stencils/a2b_ord4.py | 42 ++++++++++++++++++++++++++++++------ 1 file changed, 36 insertions(+), 6 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 3299aa343..d6f0a5462 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -53,10 +53,10 @@ def extrap_corner( ): #x1 = asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0 #x2 = asin(sqrt(sin((p2b - p0b) / 2.0) ** 2.0 + cos(p2b) * cos(p0b) * sin((p2a - p0a) / 2.0) ** 2.0)) * 2.0 - return qa + (asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) / (asin(sqrt(sin((p2b - p0b) / 2.0) ** 2.0 + cos(p2b) * cos(p0b) * sin((p2a - p0a) / 2.0) ** 2.0)) * 2.0 -asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) * (qa - qb) -# x1 = great_circle_dist(p1a, p1b, p0a, p0b) -# x2 = great_circle_dist(p2a, p2b, p0a, p0b) -# return qa + x1 / (x2 - x1) * (qa - qb) + #return qa + (asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) / (asin(sqrt(sin((p2b - p0b) / 2.0) ** 2.0 + cos(p2b) * cos(p0b) * sin((p2a - p0a) / 2.0) ** 2.0)) * 2.0 -asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) * (qa - qb) + x1 = great_circle_dist(p1a, p1b, p0a, p0b) + x2 = great_circle_dist(p2a, p2b, p0a, p0b) + return qa + x1 / (x2 - x1) * (qa - qb) """ def _a2b_corners(qin, qout, agrid1, agrid2, bgrid1, bgrid2): @@ -112,9 +112,39 @@ def _sw_corner( qin[0, -1, 0], qin[1, -2, 0], ) - +with horizontal(region[i_end + 1, j_start]): qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) - + ec1 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[-1, 0], + agrid2[-1, 0], + agrid1[-2, 1], + agrid2[-2, 1], + qin[-1, 0, 0], + qin[-2, 1, 0], + ) + ec2 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[-1, -1], + agrid2[-1, -1], + agrid1[-2, -2], + agrid2[-2, -2], + qin[-1, -1, 0], + qin[-2, -2, 0], + ) + ec3 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[0, 0], + agrid2[0, 0], + agrid1[1, 1], + agrid2[1, 1], + qin[0, 0, 0], + qin[1, 1, 0], + ) + qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) def _nw_corner( qin: FloatField, From adbd45e258e1069382faac6a80dd58572251b528 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 11:20:25 -0700 Subject: [PATCH 15/45] put the nw corner with the sw corner --- fv3core/stencils/a2b_ord4.py | 92 ++++++++++++++++++------------------ 1 file changed, 46 insertions(+), 46 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index d6f0a5462..d4c0ae1e6 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -78,7 +78,7 @@ def _sw_corner( bgrid1: FloatFieldIJ, bgrid2: FloatFieldIJ, ): - from __externals__ import i_start, j_start + from __externals__ import i_start, i_end, j_start with computation(PARALLEL), interval(...): tmp = 0.0 with horizontal(region[i_start, j_start]): @@ -112,39 +112,39 @@ def _sw_corner( qin[0, -1, 0], qin[1, -2, 0], ) -with horizontal(region[i_end + 1, j_start]): qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) - ec1 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[-1, 0], - agrid2[-1, 0], - agrid1[-2, 1], - agrid2[-2, 1], - qin[-1, 0, 0], - qin[-2, 1, 0], - ) - ec2 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[-1, -1], - agrid2[-1, -1], - agrid1[-2, -2], - agrid2[-2, -2], - qin[-1, -1, 0], - qin[-2, -2, 0], - ) - ec3 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[0, 0], - agrid2[0, 0], - agrid1[1, 1], - agrid2[1, 1], - qin[0, 0, 0], - qin[1, 1, 0], - ) - qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) + with horizontal(region[i_end + 1, j_start]): + ec1 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[-1, 0], + agrid2[-1, 0], + agrid1[-2, 1], + agrid2[-2, 1], + qin[-1, 0, 0], + qin[-2, 1, 0], + ) + ec2 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[-1, -1], + agrid2[-1, -1], + agrid1[-2, -2], + agrid2[-2, -2], + qin[-1, -1, 0], + qin[-2, -2, 0], + ) + ec3 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[0, 0], + agrid2[0, 0], + agrid1[1, 1], + agrid2[1, 1], + qin[0, 0, 0], + qin[1, 1, 0], + ) + qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) def _nw_corner( qin: FloatField, @@ -526,11 +526,11 @@ def __init__( origin=(self.grid.isd, self.grid.jsd, kstart), domain=(self.grid.nid, self.grid.njd, nk), ) - self._nw_corner_stencil = FrozenStencil( - _nw_corner, - origin=(self.grid.ie + 1, self.grid.js, kstart), - domain=corner_domain, - ) + #self._nw_corner_stencil = FrozenStencil( + # _nw_corner, + # origin=(self.grid.ie + 1, self.grid.js, kstart), + # domain=corner_domain, + #) self._ne_corner_stencil = FrozenStencil( _ne_corner, origin=(self.grid.ie + 1, self.grid.je + 1, kstart), @@ -624,14 +624,14 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid2, ) - self._nw_corner_stencil( - qin, - qout, - self.grid.agrid1, - self.grid.agrid2, - self.grid.bgrid1, - self.grid.bgrid2, - ) + #self._nw_corner_stencil( + # qin, + # qout, + # self.grid.agrid1, + # self.grid.agrid2, + # self.grid.bgrid1, + # self.grid.bgrid2, + #) self._ne_corner_stencil( qin, qout, From 8281a01285eb2578dcf09256267251981e673665 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 13:18:18 -0700 Subject: [PATCH 16/45] all corners --- fv3core/stencils/a2b_ord4.py | 213 +++++++++++++---------------------- 1 file changed, 77 insertions(+), 136 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index d4c0ae1e6..817c0e9e6 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -78,7 +78,7 @@ def _sw_corner( bgrid1: FloatFieldIJ, bgrid2: FloatFieldIJ, ): - from __externals__ import i_start, i_end, j_start + from __externals__ import i_start, i_end, j_start, j_end with computation(PARALLEL), interval(...): tmp = 0.0 with horizontal(region[i_start, j_start]): @@ -146,130 +146,70 @@ def _sw_corner( ) qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) -def _nw_corner( - qin: FloatField, - qout: FloatField, - agrid1: FloatFieldIJ, - agrid2: FloatFieldIJ, - bgrid1: FloatFieldIJ, - bgrid2: FloatFieldIJ, -): - with computation(PARALLEL), interval(...): - ec1 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[-1, 0], - agrid2[-1, 0], - agrid1[-2, 1], - agrid2[-2, 1], - qin[-1, 0, 0], - qin[-2, 1, 0], - ) - ec2 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[-1, -1], - agrid2[-1, -1], - agrid1[-2, -2], - agrid2[-2, -2], - qin[-1, -1, 0], - qin[-2, -2, 0], - ) - ec3 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[0, 0], - agrid2[0, 0], - agrid1[1, 1], - agrid2[1, 1], - qin[0, 0, 0], - qin[1, 1, 0], - ) - qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) - - -def _ne_corner( - qin: FloatField, - qout: FloatField, - agrid1: FloatFieldIJ, - agrid2: FloatFieldIJ, - bgrid1: FloatFieldIJ, - bgrid2: FloatFieldIJ, -): - with computation(PARALLEL), interval(...): - ec1 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[-1, -1], - agrid2[-1, -1], - agrid1[-2, -2], - agrid2[-2, -2], - qin[-1, -1, 0], - qin[-2, -2, 0], - ) - ec2 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[0, -1], - agrid2[0, -1], - agrid1[1, -2], - agrid2[1, -2], - qin[0, -1, 0], - qin[1, -2, 0], - ) - ec3 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[-1, 0], - agrid2[-1, 0], - agrid1[-2, 1], - agrid2[-2, 1], - qin[-1, 0, 0], - qin[-2, 1, 0], - ) - qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) - - -def _se_corner( - qin: FloatField, - qout: FloatField, - agrid1: FloatFieldIJ, - agrid2: FloatFieldIJ, - bgrid1: FloatFieldIJ, - bgrid2: FloatFieldIJ, -): - with computation(PARALLEL), interval(...): - ec1 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[0, -1], - agrid2[0, -1], - agrid1[1, -2], - agrid2[1, -2], - qin[0, -1, 0], - qin[1, -2, 0], - ) - ec2 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[-1, -1], - agrid2[-1, -1], - agrid1[-2, -2], - agrid2[-2, -2], - qin[-1, -1, 0], - qin[-2, -2, 0], - ) - ec3 = extrap_corner( - bgrid1[0, 0], - bgrid2[0, 0], - agrid1[0, 0], - agrid2[0, 0], - agrid1[1, 1], - agrid2[1, 1], - qin[0, 0, 0], - qin[1, 1, 0], - ) - qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) + with horizontal(region[i_end + 1, j_end + 1]): + ec1 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[-1, -1], + agrid2[-1, -1], + agrid1[-2, -2], + agrid2[-2, -2], + qin[-1, -1, 0], + qin[-2, -2, 0], + ) + ec2 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[0, -1], + agrid2[0, -1], + agrid1[1, -2], + agrid2[1, -2], + qin[0, -1, 0], + qin[1, -2, 0], + ) + ec3 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[-1, 0], + agrid2[-1, 0], + agrid1[-2, 1], + agrid2[-2, 1], + qin[-1, 0, 0], + qin[-2, 1, 0], + ) + qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) + with horizontal(region[i_start, j_end + 1]): + ec1 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[0, -1], + agrid2[0, -1], + agrid1[1, -2], + agrid2[1, -2], + qin[0, -1, 0], + qin[1, -2, 0], + ) + ec2 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[-1, -1], + agrid2[-1, -1], + agrid1[-2, -2], + agrid2[-2, -2], + qin[-1, -1, 0], + qin[-2, -2, 0], + ) + ec3 = extrap_corner( + bgrid1[0, 0], + bgrid2[0, 0], + agrid1[0, 0], + agrid2[0, 0], + agrid1[1, 1], + agrid2[1, 1], + qin[0, 0, 0], + qin[1, 1, 0], + ) + qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) @gtscript.function @@ -531,16 +471,16 @@ def __init__( # origin=(self.grid.ie + 1, self.grid.js, kstart), # domain=corner_domain, #) - self._ne_corner_stencil = FrozenStencil( - _ne_corner, - origin=(self.grid.ie + 1, self.grid.je + 1, kstart), - domain=corner_domain, - ) - self._se_corner_stencil = FrozenStencil( - _se_corner, - origin=(self.grid.is_, self.grid.je + 1, kstart), - domain=corner_domain, - ) + #self._ne_corner_stencil = FrozenStencil( + # _ne_corner, + # origin=(self.grid.ie + 1, self.grid.je + 1, kstart), + # domain=corner_domain, + #) + #self._se_corner_stencil = FrozenStencil( + # _se_corner, + # origin=(self.grid.is_, self.grid.je + 1, kstart), + # domain=corner_domain, + #) js = self.grid.js + 1 if self.grid.south_edge else self.grid.js je = self.grid.je if self.grid.north_edge else self.grid.je + 1 @@ -632,6 +572,7 @@ def __call__(self, qin: FloatField, qout: FloatField): # self.grid.bgrid1, # self.grid.bgrid2, #) + """ self._ne_corner_stencil( qin, qout, @@ -648,7 +589,7 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid1, self.grid.bgrid2, ) - + """ self._compute_qout_edges(qin, qout) self._a2b_interpolation_stencil( qin, From a52633a41fef7114dce3b410ed1a017d1bfd88c7 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 16:37:03 -0700 Subject: [PATCH 17/45] excessive recursion reached, split 4 corners into 2 stencils. also working on removing copy variables --- fv3core/stencils/a2b_ord4.py | 71 +++++++++++++++++++++++++----------- 1 file changed, 49 insertions(+), 22 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 817c0e9e6..a8511efb4 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -70,7 +70,7 @@ def _a2b_corners(qin, qout, agrid1, agrid2, bgrid1, bgrid2): qout = _nw_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) return qout """ -def _sw_corner( +def _south_corners( qin: FloatField, qout: FloatField, agrid1: FloatFieldIJ, @@ -145,7 +145,18 @@ def _sw_corner( qin[1, 1, 0], ) qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) - +def _north_corners( + qin: FloatField, + qout: FloatField, + agrid1: FloatFieldIJ, + agrid2: FloatFieldIJ, + bgrid1: FloatFieldIJ, + bgrid2: FloatFieldIJ, +): + from __externals__ import i_start, i_end, j_start, j_end + with computation(PARALLEL), interval(...): + tmp = 0.0 + with horizontal(region[i_end + 1, j_end + 1]): ec1 = extrap_corner( bgrid1[0, 0], @@ -240,32 +251,30 @@ def a2b_interpolation( # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) - qx_copy = qx with horizontal(region[i_start, :]): - #qx = qx_edge_west(qin, dxa) qx = 0.5 * ( ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + dxa[-2, 0] / dxa[-1, 0]) ) - qx_copy = qx with horizontal(region[i_start + 1, :]): - #qx = qx_edge_west2(qin, dxa, qx_copy) - qx = ( - 3.0 * (dxa / dxa[-1, 0] * qin[-1, 0, 0] + qin) - (dxa / dxa[-1, 0] * qx_copy[-1, 0, 0] + qx_copy[1, 0, 0]) - ) / (2.0 + 2.0 * dxa / dxa[-1, 0]) + qx = ( + 3.0 * (dxa / dxa[-1, 0] * qin[-1, 0, 0] + qin) - (dxa / dxa[-1, 0] * (0.5 * ( + ((2.0 + dxa[0, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin) / (1.0 + dxa[0, 0] / dxa[-1, 0]) + + ((2.0 + dxa[-3, 0] / dxa[-2, 0]) * qin[-2, 0, 0] - qin[-3, 0, 0]) / (1.0 + dxa[-3, 0] / dxa[-2, 0]) + )) + (b2 * (qin[-1, 0, 0] + qin[2, 0, 0]) + b1 * (qin + qin[1, 0, 0]))) + ) / (2.0 + 2.0 * dxa / dxa[-1, 0]) with horizontal(region[i_end + 1, :]): - #qx = qx_edge_east(qin, dxa) qx = 0.5 * ( ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + dxa[-2, 0] / dxa[-1, 0]) + ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) ) - - qx_copy = qx with horizontal(region[i_end, :]): - #qx = qx_edge_east2(qin, dxa, qx_copy) qx = ( - 3.0 * (qin[-1, 0, 0] + dxa[-1, 0] / dxa * qin) - (dxa[-1, 0] / dxa * qx_copy[1, 0, 0] + qx_copy[-1, 0, 0]) + 3.0 * (qin[-1, 0, 0] + dxa[-1, 0] / dxa * qin) - (dxa[-1, 0] / dxa * ( 0.5 * ( + ((2.0 + dxa[-1, 0] / dxa) * qin - qin[-1, 0, 0]) / (1.0 + dxa[-1, 0] / dxa) + + ((2.0 + dxa[2, 0] / dxa[1, 0]) * qin[1, 0, 0] - qin[2, 0, 0]) / (1.0 + dxa[2, 0] / dxa[1, 0]) + )) + ( b2 * (qin[-3, 0, 0] + qin) + b1 * (qin[-2, 0, 0] + qin[-1, 0, 0]))) ) / (2.0 + 2.0 * dxa[-1, 0] / dxa) @@ -273,7 +282,6 @@ def a2b_interpolation( qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) qy_copy = qy with horizontal(region[:, j_start]): - #qy = qy_edge_south(qin, dya) g_in = dya[0, 1] / dya g_ou = dya[0, -2] / dya[0, -1] qy= 0.5 * ( @@ -283,12 +291,19 @@ def a2b_interpolation( qy_copy = qy with horizontal(region[:, j_start + 1]): - qy = qy_edge_south2(qin, dya, qy_copy) + qy = (3.0 * (dya / dya[0, -1] * qin[0, -1, 0] + qin) - (dya / dya[0, -1] * qy[0, -1, 0] + qy[0, 1, 0]) + ) / (2.0 + 2.0 * dya / dya[0, -1]) with horizontal(region[:, j_end + 1]): - qy = qy_edge_north(qin, dya) + qy = 0.5 * ( + ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + dya[0, -2] / dya[0, -1]) + + ((2.0 + dya[0, 1] / dya) * qin - qin[0, 1, 0]) / (1.0 + dya[0, 1] / dya) + ) qy_copy = qy with horizontal(region[:, j_end]): - qy = qy_edge_north2(qin, dya, qy_copy) + qy = ( + 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) - (dya[0, -1] / dya * qy[0, 1, 0] + qy[0, -1, 0]) + ) / (2.0 + 2.0 * dya[0, -1] / dya) + qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + qxx[0, 1, 0]) @@ -461,8 +476,13 @@ def __init__( self.grid.full_origin(), self.grid.domain_shape_full(), ) - self._sw_corner_stencil = FrozenStencil( - _sw_corner, externals=full_ax_offsets, + self._south_corners_stencil = FrozenStencil( + _south_corners, externals=full_ax_offsets, + origin=(self.grid.isd, self.grid.jsd, kstart), + domain=(self.grid.nid, self.grid.njd, nk), + ) + self._north_corners_stencil = FrozenStencil( + _north_corners, externals=full_ax_offsets, origin=(self.grid.isd, self.grid.jsd, kstart), domain=(self.grid.nid, self.grid.njd, nk), ) @@ -555,7 +575,15 @@ def __call__(self, qin: FloatField, qout: FloatField): qout: Output on B-grid (inout) """ - self._sw_corner_stencil( + self._south_corners_stencil( + qin, + qout, + self.grid.agrid1, + self.grid.agrid2, + self.grid.bgrid1, + self.grid.bgrid2, + ) + self._north_corners_stencil( qin, qout, self.grid.agrid1, @@ -563,7 +591,6 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid1, self.grid.bgrid2, ) - #self._nw_corner_stencil( # qin, # qout, From 99e7cebfc59d28a168ff4e4fe69b6141ca8569e9 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 16:48:50 -0700 Subject: [PATCH 18/45] remove copy vars --- fv3core/stencils/a2b_ord4.py | 266 +++++++++++++++++++---------------- 1 file changed, 143 insertions(+), 123 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index a8511efb4..7a07a983d 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -1,3 +1,6 @@ +from typing import Tuple + +import gt4py import gt4py.gtscript as gtscript from gt4py.gtscript import ( PARALLEL, @@ -13,12 +16,11 @@ import fv3core._config as spec import fv3core.utils.gt4py_utils as utils -from fv3core.decorators import FrozenStencil,gtstencil +from fv3core.decorators import FrozenStencil from fv3core.stencils.basic_operations import copy_defn from fv3core.utils import axis_offsets from fv3core.utils.typing import FloatField, FloatFieldI, FloatFieldIJ -import gt4py -from typing import Tuple + # comact 4-pt cubic interpolation c1 = 2.0 / 3.0 @@ -38,7 +40,7 @@ def great_circle_dist(p1a, p1b, p2a, p2b): tb = sin((p1b - p2b) / 2.0) ** 2.0 ta = sin((p1a - p2a) / 2.0) ** 2.0 return asin(sqrt(tb + cos(p1b) * cos(p2b) * ta)) * 2.0 -#return asin(sqrt(sin((p1b - p2b) / 2.0) ** 2.0 + cos(p1b) * cos(p2b) * sin((p1a - p2a) / 2.0) ** 2.0)) * 2.0 + @gtscript.function def extrap_corner( @@ -51,25 +53,25 @@ def extrap_corner( qa, qb, ): - #x1 = asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0 - #x2 = asin(sqrt(sin((p2b - p0b) / 2.0) ** 2.0 + cos(p2b) * cos(p0b) * sin((p2a - p0a) / 2.0) ** 2.0)) * 2.0 - #return qa + (asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) / (asin(sqrt(sin((p2b - p0b) / 2.0) ** 2.0 + cos(p2b) * cos(p0b) * sin((p2a - p0a) / 2.0) ** 2.0)) * 2.0 -asin(sqrt(sin((p1b - p0b) / 2.0) ** 2.0 + cos(p1b) * cos(p0b) * sin((p1a - p0a) / 2.0) ** 2.0)) * 2.0) * (qa - qb) x1 = great_circle_dist(p1a, p1b, p0a, p0b) x2 = great_circle_dist(p2a, p2b, p0a, p0b) return qa + x1 / (x2 - x1) * (qa - qb) - + + """ def _a2b_corners(qin, qout, agrid1, agrid2, bgrid1, bgrid2): with computation(PARALLEL), interval(...): from __externals__ import i_start, j_start - + with horizontal(region[i_start, j_start]): qout = _sw_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) qout = _se_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) qout = _ne_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) qout = _nw_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) - return qout -""" + return qout +""" + + def _south_corners( qin: FloatField, qout: FloatField, @@ -78,7 +80,8 @@ def _south_corners( bgrid1: FloatFieldIJ, bgrid2: FloatFieldIJ, ): - from __externals__ import i_start, i_end, j_start, j_end + from __externals__ import i_end, i_start, j_end, j_start + with computation(PARALLEL), interval(...): tmp = 0.0 with horizontal(region[i_start, j_start]): @@ -145,6 +148,8 @@ def _south_corners( qin[1, 1, 0], ) qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) + + def _north_corners( qin: FloatField, qout: FloatField, @@ -153,10 +158,11 @@ def _north_corners( bgrid1: FloatFieldIJ, bgrid2: FloatFieldIJ, ): - from __externals__ import i_start, i_end, j_start, j_end + from __externals__ import i_end, i_start, j_end, j_start + with computation(PARALLEL), interval(...): tmp = 0.0 - + with horizontal(region[i_end + 1, j_end + 1]): ec1 = extrap_corner( bgrid1[0, 0], @@ -237,9 +243,7 @@ def a2b_interpolation( qin: FloatField, qout: FloatField, qx: FloatField, - qx_copy: FloatField, qy: FloatField, - qy_copy: FloatField, qxx: FloatField, qyy: FloatField, dxa: FloatFieldIJ, @@ -248,62 +252,112 @@ def a2b_interpolation( from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): - + # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) with horizontal(region[i_start, :]): - qx = 0.5 * ( - ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) - + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + dxa[-2, 0] / dxa[-1, 0]) + qx = 0.5 * ( + ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) + + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) + / (1.0 + dxa[-2, 0] / dxa[-1, 0]) ) with horizontal(region[i_start + 1, :]): qx = ( - 3.0 * (dxa / dxa[-1, 0] * qin[-1, 0, 0] + qin) - (dxa / dxa[-1, 0] * (0.5 * ( - ((2.0 + dxa[0, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin) / (1.0 + dxa[0, 0] / dxa[-1, 0]) - + ((2.0 + dxa[-3, 0] / dxa[-2, 0]) * qin[-2, 0, 0] - qin[-3, 0, 0]) / (1.0 + dxa[-3, 0] / dxa[-2, 0]) - )) + (b2 * (qin[-1, 0, 0] + qin[2, 0, 0]) + b1 * (qin + qin[1, 0, 0]))) + 3.0 * (dxa / dxa[-1, 0] * qin[-1, 0, 0] + qin) + - ( + dxa + / dxa[-1, 0] + * ( + 0.5 + * ( + ((2.0 + dxa[0, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin) + / (1.0 + dxa[0, 0] / dxa[-1, 0]) + + ( + (2.0 + dxa[-3, 0] / dxa[-2, 0]) * qin[-2, 0, 0] + - qin[-3, 0, 0] + ) + / (1.0 + dxa[-3, 0] / dxa[-2, 0]) + ) + ) + + (b2 * (qin[-1, 0, 0] + qin[2, 0, 0]) + b1 * (qin + qin[1, 0, 0])) + ) ) / (2.0 + 2.0 * dxa / dxa[-1, 0]) with horizontal(region[i_end + 1, :]): qx = 0.5 * ( - ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + dxa[-2, 0] / dxa[-1, 0]) - + ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) + / (1.0 + dxa[-2, 0] / dxa[-1, 0]) + + ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) + / (1.0 + dxa[1, 0] / dxa) ) with horizontal(region[i_end, :]): qx = ( - 3.0 * (qin[-1, 0, 0] + dxa[-1, 0] / dxa * qin) - (dxa[-1, 0] / dxa * ( 0.5 * ( - ((2.0 + dxa[-1, 0] / dxa) * qin - qin[-1, 0, 0]) / (1.0 + dxa[-1, 0] / dxa) - + ((2.0 + dxa[2, 0] / dxa[1, 0]) * qin[1, 0, 0] - qin[2, 0, 0]) / (1.0 + dxa[2, 0] / dxa[1, 0]) - )) + ( b2 * (qin[-3, 0, 0] + qin) + b1 * (qin[-2, 0, 0] + qin[-1, 0, 0]))) + 3.0 * (qin[-1, 0, 0] + dxa[-1, 0] / dxa * qin) + - ( + dxa[-1, 0] + / dxa + * ( + 0.5 + * ( + ((2.0 + dxa[-1, 0] / dxa) * qin - qin[-1, 0, 0]) + / (1.0 + dxa[-1, 0] / dxa) + + ( + (2.0 + dxa[2, 0] / dxa[1, 0]) * qin[1, 0, 0] + - qin[2, 0, 0] + ) + / (1.0 + dxa[2, 0] / dxa[1, 0]) + ) + ) + + ( + b2 * (qin[-3, 0, 0] + qin) + + b1 * (qin[-2, 0, 0] + qin[-1, 0, 0]) + ) + ) ) / (2.0 + 2.0 * dxa[-1, 0] / dxa) - # ppm_volume_mean_y qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) - qy_copy = qy with horizontal(region[:, j_start]): - g_in = dya[0, 1] / dya - g_ou = dya[0, -2] / dya[0, -1] - qy= 0.5 * ( - ((2.0 + g_in) * qin - qin[0, 1, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + g_ou) + qy = 0.5 * ( + ((2.0 + dya[0, 1] / dya) * qin - qin[0, 1, 0]) / (1.0 + dya[0, 1] / dya) + + ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) + / (1.0 + dya[0, -2] / dya[0, -1]) ) - qy_copy = qy with horizontal(region[:, j_start + 1]): - qy = (3.0 * (dya / dya[0, -1] * qin[0, -1, 0] + qin) - (dya / dya[0, -1] * qy[0, -1, 0] + qy[0, 1, 0]) - ) / (2.0 + 2.0 * dya / dya[0, -1]) + qy = ( + 3.0 * (dya / dya[0, -1] * qin[0, -1, 0] + qin) + - ( + dya + / dya[0, -1] + * ( + 0.5 + * ( + ((2.0 + dya / dya[0, -1]) * qin[0, -1, 0] - qin) + / (1.0 + dya / dya[0, -1]) + + ( + (2.0 + dya[0, -3] / dya[0, -2]) * qin[0, -2, 0] + - qin[0, -3, 0] + ) + / (1.0 + dya[0, -3] / dya[0, -2]) + ) + ) + + (b2 * (qin[0, -1, 0] + qin[0, 2, 0]) + b1 * (qin + qin[0, 1, 0])) + ) + ) / (2.0 + 2.0 * dya / dya[0, -1]) with horizontal(region[:, j_end + 1]): - qy = 0.5 * ( - ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + dya[0, -2] / dya[0, -1]) - + ((2.0 + dya[0, 1] / dya) * qin - qin[0, 1, 0]) / (1.0 + dya[0, 1] / dya) + qy = 0.5 * ( + ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) + / (1.0 + dya[0, -2] / dya[0, -1]) + + ((2.0 + dya[0, 1] / dya) * qin - qin[0, 1, 0]) + / (1.0 + dya[0, 1] / dya) ) - qy_copy = qy with horizontal(region[:, j_end]): - qy = ( - 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) - (dya[0, -1] / dya * qy[0, 1, 0] + qy[0, -1, 0]) + qy = ( + 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) + - (dya[0, -1] / dya * qy[0, 1, 0] + qy[0, -1, 0]) ) / (2.0 + 2.0 * dya[0, -1] / dya) - + qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + qxx[0, 1, 0]) @@ -408,7 +462,9 @@ def qy_edge_north2(qin: FloatField, dya: FloatFieldIJ, qy: FloatField): 3.0 * (qin[0, -1, 0] + g_in * qin) - (g_in * qy[0, 1, 0] + qy[0, -1, 0]) ) / (2.0 + 2.0 * g_in) -def qout_edges_x(q2: FloatField, + +def qout_edges_x( + q2: FloatField, qin: FloatField, qout: FloatField, dxa: FloatFieldIJ, @@ -424,6 +480,7 @@ def qout_edges_x(q2: FloatField, with horizontal(region[i_end + 1, :]): qout = edge_e * q2[0, -1, 0] + (1.0 - edge_e) * q2 + def qout_edges_y( q1: FloatField, qin: FloatField, @@ -441,6 +498,7 @@ def qout_edges_y( with horizontal(region[:, j_end + 1]): qout = edge_n * q1[-1, 0, 0] + (1.0 - edge_n) * q1 + class AGrid2BGridFourthOrder: """ Fortran name is a2b_ord4, test module is A2B_Ord4 @@ -466,8 +524,6 @@ def __init__( self._tmp_qy = utils.make_storage_from_shape(shape) self._tmp_qxx = utils.make_storage_from_shape(shape) self._tmp_qyy = utils.make_storage_from_shape(shape) - self._tmp_qx_copy = utils.make_storage_from_shape(shape) - self._tmp_qy_copy = utils.make_storage_from_shape(shape) if nk is None: nk = self.grid.npz - kstart corner_domain = (1, 1, nk) @@ -477,31 +533,18 @@ def __init__( self.grid.domain_shape_full(), ) self._south_corners_stencil = FrozenStencil( - _south_corners, externals=full_ax_offsets, + _south_corners, + externals=full_ax_offsets, origin=(self.grid.isd, self.grid.jsd, kstart), domain=(self.grid.nid, self.grid.njd, nk), ) self._north_corners_stencil = FrozenStencil( - _north_corners, externals=full_ax_offsets, + _north_corners, + externals=full_ax_offsets, origin=(self.grid.isd, self.grid.jsd, kstart), domain=(self.grid.nid, self.grid.njd, nk), ) - #self._nw_corner_stencil = FrozenStencil( - # _nw_corner, - # origin=(self.grid.ie + 1, self.grid.js, kstart), - # domain=corner_domain, - #) - #self._ne_corner_stencil = FrozenStencil( - # _ne_corner, - # origin=(self.grid.ie + 1, self.grid.je + 1, kstart), - # domain=corner_domain, - #) - #self._se_corner_stencil = FrozenStencil( - # _se_corner, - # origin=(self.grid.is_, self.grid.je + 1, kstart), - # domain=corner_domain, - #) - + js = self.grid.js + 1 if self.grid.south_edge else self.grid.js je = self.grid.je if self.grid.north_edge else self.grid.je + 1 is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ @@ -511,7 +554,7 @@ def __init__( # edge_w is singleton in the I-dimension to work around gt4py not yet # supporting J-fields. As a result, the origin has to be zero for # edge_w, anything higher is outside its index range - + self._qout_x_edge_west = FrozenStencil( qout_x_edge, origin={"_all_": (self.grid.is_, js, kstart), "edge_w": (0, js)}, @@ -522,28 +565,31 @@ def __init__( origin={"_all_": (self.grid.ie + 1, js, kstart), "edge_w": (0, js)}, domain=(1, dj, nk), ) - - + self._edge_e = self._j_storage_repeat_over_i(self.grid.edge_e, shape[0:2]) self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) x_edge_origin = (self.grid.is_, js, kstart) - x_edge_domain = (self.grid.nic +2 , dj, nk) + x_edge_domain = (self.grid.nic + 2, dj, nk) x_offsets = axis_offsets(self.grid, x_edge_origin, x_edge_domain) - self._qout_edges_x = FrozenStencil(qout_edges_x,externals=x_offsets, - origin=x_edge_origin, - domain=x_edge_domain, - ) + self._qout_edges_x = FrozenStencil( + qout_edges_x, + externals=x_offsets, + origin=x_edge_origin, + domain=x_edge_domain, + ) self._tmp_q2 = utils.make_storage_from_shape(shape) self._tmp_q1 = utils.make_storage_from_shape(shape) - + y_edge_origin = (is_, self.grid.js, kstart) y_edge_domain = (di, self.grid.njc + 2, nk) y_offsets = axis_offsets(self.grid, y_edge_origin, y_edge_domain) - self._qout_edges_y = FrozenStencil(qout_edges_y,externals=y_offsets, - origin=y_edge_origin, - domain=y_edge_domain, - ) - + self._qout_edges_y = FrozenStencil( + qout_edges_y, + externals=y_offsets, + origin=y_edge_origin, + domain=y_edge_domain, + ) + origin = (is_, js, kstart) domain = (ie - is_ + 1, je - js + 1, nk) ax_offsets = axis_offsets( @@ -559,11 +605,14 @@ def __init__( origin=(self.grid.is_, self.grid.js, kstart), domain=(self.grid.nic + 1, self.grid.njc + 1, nk), ) - # TODO + + # TODO # within regions, the edge_w and edge_w variables that are singleton in the # I dimension error, workaround is repeating the data, but the longterm - # fix should happen in regions - def _j_storage_repeat_over_i(self, grid_array: gt4py.storage.Storage, shape: Tuple[int, int]): + # fix should happen in regions + def _j_storage_repeat_over_i( + self, grid_array: gt4py.storage.Storage, shape: Tuple[int, int] + ): dup = utils.repeat(grid_array, shape[1], axis=0) return utils.make_storage_data(dup, shape, (0, 0)) @@ -591,40 +640,13 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid1, self.grid.bgrid2, ) - #self._nw_corner_stencil( - # qin, - # qout, - # self.grid.agrid1, - # self.grid.agrid2, - # self.grid.bgrid1, - # self.grid.bgrid2, - #) - """ - self._ne_corner_stencil( - qin, - qout, - self.grid.agrid1, - self.grid.agrid2, - self.grid.bgrid1, - self.grid.bgrid2, - ) - self._se_corner_stencil( - qin, - qout, - self.grid.agrid1, - self.grid.agrid2, - self.grid.bgrid1, - self.grid.bgrid2, - ) - """ + self._compute_qout_edges(qin, qout) self._a2b_interpolation_stencil( qin, qout, self._tmp_qx, - self._tmp_qx_copy, self._tmp_qy, - self._tmp_qy_copy, self._tmp_qxx, self._tmp_qyy, self.grid.dxa, @@ -637,18 +659,16 @@ def __call__(self, qin: FloatField, qout: FloatField): ) def _compute_qout_edges(self, qin: FloatField, qout: FloatField): - - self._qout_edges_x(self._tmp_q2, - qin, - qout, - self.grid.dxa,self._edge_w,self._edge_e + + self._qout_edges_x( + self._tmp_q2, qin, qout, self.grid.dxa, self._edge_w, self._edge_e ) - - + self._qout_edges_y( self._tmp_q1, - qin, + qin, qout, - self.grid.dya, self.grid.edge_s, self.grid.edge_n, + self.grid.dya, + self.grid.edge_s, + self.grid.edge_n, ) - From be7da0d77d9c21df097c8204295759fced97e34d Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 19:08:05 -0700 Subject: [PATCH 19/45] merged qout edges and replace into a2b_interpolation, eliminated q1 and q2 as separate variables --- fv3core/stencils/a2b_ord4.py | 205 +++++++---------------------------- 1 file changed, 38 insertions(+), 167 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 7a07a983d..2be35b812 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -248,10 +248,25 @@ def a2b_interpolation( qyy: FloatField, dxa: FloatFieldIJ, dya: FloatFieldIJ, + edge_w: FloatFieldIJ, + edge_e: FloatFieldIJ, + edge_s: FloatFieldI, + edge_n: FloatFieldI, ): - from __externals__ import i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start, j_end, j_start, REPLACE with computation(PARALLEL), interval(...): + # qout_edges_x + + with horizontal(region[i_start, j_start + 1 : j_end + 1]): + qout = edge_w * ((qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / (dxa[-1, -1] + dxa[0, -1])) + (1.0 - edge_w) * ((qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa)) + with horizontal(region[i_end + 1, j_start + 1 : j_end + 1]): + qout = edge_e * ((qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / (dxa[-1, -1] + dxa[0, -1])) + (1.0 - edge_e) * ((qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa)) + + with horizontal(region[i_start + 1 : i_end + 1, j_start]): + qout = edge_s * ((qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) / (dya[-1, -1] + dya[-1, 0])) + (1.0 - edge_s) * ((qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya)) + with horizontal(region[i_start + 1 : i_end + 1, j_end + 1]): + qout = edge_n * ((qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) / (dya[-1, -1] + dya[-1, 0])) + (1.0 - edge_n) * ((qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya)) # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) @@ -368,7 +383,10 @@ def a2b_interpolation( qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[-1, 0, 0] + qyy[1, 0, 0]) with horizontal(region[i_end, :]): qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[1, 0, 0] + qyy[-1, 0, 0]) - qout = 0.5 * (qxx + qyy) + with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): + qout = 0.5 * (qxx + qyy) + if __INLINED(REPLACE): + qin = qout def qout_x_edge( @@ -387,81 +405,6 @@ def qout_y_edge( qout[0, 0, 0] = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 -@gtscript.function -def qx_edge_west(qin: FloatField, dxa: FloatFieldIJ): - g_in = dxa[1, 0] / dxa - g_ou = dxa[-2, 0] / dxa[-1, 0] - return 0.5 * ( - ((2.0 + g_in) * qin - qin[1, 0, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + g_ou) - ) - # This does not work, due to access of qx that is changing above - - # qx[1, 0, 0] = (3.0 * (g_in * qin + qin[1, 0, 0]) - # - (g_in * qx + qx[2, 0, 0])) / (2.0 + 2.0 * g_in) - - -@gtscript.function -def qx_edge_west2(qin: FloatField, dxa: FloatFieldIJ, qx: FloatField): - g_in = dxa / dxa[-1, 0] - return ( - 3.0 * (g_in * qin[-1, 0, 0] + qin) - (g_in * qx[-1, 0, 0] + qx[1, 0, 0]) - ) / (2.0 + 2.0 * g_in) - - -@gtscript.function -def qx_edge_east(qin: FloatField, dxa: FloatFieldIJ): - g_in = dxa[-2, 0] / dxa[-1, 0] - g_ou = dxa[1, 0] / dxa - return 0.5 * ( - ((2.0 + g_in) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin - qin[1, 0, 0]) / (1.0 + g_ou) - ) - - -@gtscript.function -def qx_edge_east2(qin: FloatField, dxa: FloatFieldIJ, qx: FloatField): - g_in = dxa[-1, 0] / dxa - return ( - 3.0 * (qin[-1, 0, 0] + g_in * qin) - (g_in * qx[1, 0, 0] + qx[-1, 0, 0]) - ) / (2.0 + 2.0 * g_in) - - -@gtscript.function -def qy_edge_south(qin: FloatField, dya: FloatFieldIJ): - g_in = dya[0, 1] / dya - g_ou = dya[0, -2] / dya[0, -1] - return 0.5 * ( - ((2.0 + g_in) * qin - qin[0, 1, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + g_ou) - ) - - -@gtscript.function -def qy_edge_south2(qin: FloatField, dya: FloatFieldIJ, qy: FloatField): - g_in = dya / dya[0, -1] - return ( - 3.0 * (g_in * qin[0, -1, 0] + qin) - (g_in * qy[0, -1, 0] + qy[0, 1, 0]) - ) / (2.0 + 2.0 * g_in) - - -@gtscript.function -def qy_edge_north(qin: FloatField, dya: FloatFieldIJ): - g_in = dya[0, -2] / dya[0, -1] - g_ou = dya[0, 1] / dya - return 0.5 * ( - ((2.0 + g_in) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin - qin[0, 1, 0]) / (1.0 + g_ou) - ) - - -@gtscript.function -def qy_edge_north2(qin: FloatField, dya: FloatFieldIJ, qy: FloatField): - g_in = dya[0, -1] / dya - return ( - 3.0 * (qin[0, -1, 0] + g_in * qin) - (g_in * qy[0, 1, 0] + qy[0, -1, 0]) - ) / (2.0 + 2.0 * g_in) - def qout_edges_x( q2: FloatField, @@ -524,86 +467,34 @@ def __init__( self._tmp_qy = utils.make_storage_from_shape(shape) self._tmp_qxx = utils.make_storage_from_shape(shape) self._tmp_qyy = utils.make_storage_from_shape(shape) + self._edge_e = self._j_storage_repeat_over_i(self.grid.edge_e, shape[0:2]) + self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) + if nk is None: nk = self.grid.npz - kstart corner_domain = (1, 1, nk) - full_ax_offsets = axis_offsets( + origin=(self.grid.is_, self.grid.js, kstart) + domain=(self.grid.nic + 1, self.grid.njc + 1, nk) + ax_offsets = axis_offsets( self.grid, - self.grid.full_origin(), - self.grid.domain_shape_full(), + origin, + domain ) self._south_corners_stencil = FrozenStencil( _south_corners, - externals=full_ax_offsets, - origin=(self.grid.isd, self.grid.jsd, kstart), - domain=(self.grid.nid, self.grid.njd, nk), + externals=ax_offsets, + origin=origin, + domain=domain, ) self._north_corners_stencil = FrozenStencil( _north_corners, - externals=full_ax_offsets, - origin=(self.grid.isd, self.grid.jsd, kstart), - domain=(self.grid.nid, self.grid.njd, nk), - ) - - js = self.grid.js + 1 if self.grid.south_edge else self.grid.js - je = self.grid.je if self.grid.north_edge else self.grid.je + 1 - is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ - ie = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 - dj = je - js + 1 - di = ie - is_ + 1 - # edge_w is singleton in the I-dimension to work around gt4py not yet - # supporting J-fields. As a result, the origin has to be zero for - # edge_w, anything higher is outside its index range - - self._qout_x_edge_west = FrozenStencil( - qout_x_edge, - origin={"_all_": (self.grid.is_, js, kstart), "edge_w": (0, js)}, - domain=(1, dj, nk), - ) - self._qout_x_edge_east = FrozenStencil( - qout_x_edge, - origin={"_all_": (self.grid.ie + 1, js, kstart), "edge_w": (0, js)}, - domain=(1, dj, nk), - ) - - self._edge_e = self._j_storage_repeat_over_i(self.grid.edge_e, shape[0:2]) - self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) - x_edge_origin = (self.grid.is_, js, kstart) - x_edge_domain = (self.grid.nic + 2, dj, nk) - x_offsets = axis_offsets(self.grid, x_edge_origin, x_edge_domain) - self._qout_edges_x = FrozenStencil( - qout_edges_x, - externals=x_offsets, - origin=x_edge_origin, - domain=x_edge_domain, - ) - self._tmp_q2 = utils.make_storage_from_shape(shape) - self._tmp_q1 = utils.make_storage_from_shape(shape) - - y_edge_origin = (is_, self.grid.js, kstart) - y_edge_domain = (di, self.grid.njc + 2, nk) - y_offsets = axis_offsets(self.grid, y_edge_origin, y_edge_domain) - self._qout_edges_y = FrozenStencil( - qout_edges_y, - externals=y_offsets, - origin=y_edge_origin, - domain=y_edge_domain, - ) - - origin = (is_, js, kstart) - domain = (ie - is_ + 1, je - js + 1, nk) - ax_offsets = axis_offsets( - self.grid, - origin, - domain, + externals=ax_offsets, + origin=origin, + domain=domain, ) + self._a2b_interpolation_stencil = FrozenStencil( - a2b_interpolation, externals=ax_offsets, origin=origin, domain=domain - ) - self._copy_stencil = FrozenStencil( - copy_defn, - origin=(self.grid.is_, self.grid.js, kstart), - domain=(self.grid.nic + 1, self.grid.njc + 1, nk), + a2b_interpolation, externals={'REPLACE': replace, **ax_offsets}, origin=origin, domain=domain ) # TODO @@ -641,7 +532,6 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid2, ) - self._compute_qout_edges(qin, qout) self._a2b_interpolation_stencil( qin, qout, @@ -650,25 +540,6 @@ def __call__(self, qin: FloatField, qout: FloatField): self._tmp_qxx, self._tmp_qyy, self.grid.dxa, - self.grid.dya, - ) - if self.replace: - self._copy_stencil( - qout, - qin, - ) - - def _compute_qout_edges(self, qin: FloatField, qout: FloatField): - - self._qout_edges_x( - self._tmp_q2, qin, qout, self.grid.dxa, self._edge_w, self._edge_e - ) - - self._qout_edges_y( - self._tmp_q1, - qin, - qout, - self.grid.dya, - self.grid.edge_s, - self.grid.edge_n, + self.grid.dya, self._edge_w, self._edge_e,self.grid.edge_s, self.grid.edge_n, ) + From 975de174c053161b50388ca06aa6246a67abfe37 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 19:30:31 -0700 Subject: [PATCH 20/45] qxx and qyy regions do not rely on value in main domain --- fv3core/stencils/a2b_ord4.py | 23 +++++++++-------------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 2be35b812..c9edc0d5f 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -257,7 +257,6 @@ def a2b_interpolation( with computation(PARALLEL), interval(...): # qout_edges_x - with horizontal(region[i_start, j_start + 1 : j_end + 1]): qout = edge_w * ((qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / (dxa[-1, -1] + dxa[0, -1])) + (1.0 - edge_w) * ((qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa)) with horizontal(region[i_end + 1, j_start + 1 : j_end + 1]): @@ -375,14 +374,14 @@ def a2b_interpolation( qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): - qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + qxx[0, 1, 0]) + qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0]))) with horizontal(region[:, j_end]): - qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, 1, 0] + qxx[0, -1, 0]) + qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, 1, 0] + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0]))) qyy = a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) with horizontal(region[i_start + 1, :]): - qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[-1, 0, 0] + qyy[1, 0, 0]) + qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[-1, 0, 0] + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0]))) with horizontal(region[i_end, :]): - qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[1, 0, 0] + qyy[-1, 0, 0]) + qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[1, 0, 0] + ( a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]))) with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): qout = 0.5 * (qxx + qyy) if __INLINED(REPLACE): @@ -459,20 +458,16 @@ def __init__( """ assert grid_type < 3 self.grid = spec.grid - self.replace = replace + shape = self.grid.domain_shape_full(add=(1, 1, 1)) - full_origin = (self.grid.isd, self.grid.jsd, kstart) - + if nk is None: + nk = self.grid.npz - kstart + self._edge_e = self._j_storage_repeat_over_i(self.grid.edge_e, shape[0:2]) + self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) self._tmp_qx = utils.make_storage_from_shape(shape) self._tmp_qy = utils.make_storage_from_shape(shape) self._tmp_qxx = utils.make_storage_from_shape(shape) self._tmp_qyy = utils.make_storage_from_shape(shape) - self._edge_e = self._j_storage_repeat_over_i(self.grid.edge_e, shape[0:2]) - self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) - - if nk is None: - nk = self.grid.npz - kstart - corner_domain = (1, 1, nk) origin=(self.grid.is_, self.grid.js, kstart) domain=(self.grid.nic + 1, self.grid.njc + 1, nk) ax_offsets = axis_offsets( From c81a67aa60042cd7a6c7eb20163e78f11fb6eb78 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 19:32:27 -0700 Subject: [PATCH 21/45] lint --- fv3core/stencils/a2b_ord4.py | 85 +++++++++++++++++++++++++----------- 1 file changed, 59 insertions(+), 26 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index c9edc0d5f..927a7d1b2 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -3,6 +3,7 @@ import gt4py import gt4py.gtscript as gtscript from gt4py.gtscript import ( + __INLINED, PARALLEL, asin, computation, @@ -17,7 +18,6 @@ import fv3core._config as spec import fv3core.utils.gt4py_utils as utils from fv3core.decorators import FrozenStencil -from fv3core.stencils.basic_operations import copy_defn from fv3core.utils import axis_offsets from fv3core.utils.typing import FloatField, FloatFieldI, FloatFieldIJ @@ -80,7 +80,7 @@ def _south_corners( bgrid1: FloatFieldIJ, bgrid2: FloatFieldIJ, ): - from __externals__ import i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start, j_start with computation(PARALLEL), interval(...): tmp = 0.0 @@ -158,7 +158,7 @@ def _north_corners( bgrid1: FloatFieldIJ, bgrid2: FloatFieldIJ, ): - from __externals__ import i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start, j_end with computation(PARALLEL), interval(...): tmp = 0.0 @@ -253,19 +253,39 @@ def a2b_interpolation( edge_s: FloatFieldI, edge_n: FloatFieldI, ): - from __externals__ import i_end, i_start, j_end, j_start, REPLACE + from __externals__ import REPLACE, i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): # qout_edges_x with horizontal(region[i_start, j_start + 1 : j_end + 1]): - qout = edge_w * ((qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / (dxa[-1, -1] + dxa[0, -1])) + (1.0 - edge_w) * ((qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa)) + qout = edge_w * ( + (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) + / (dxa[-1, -1] + dxa[0, -1]) + ) + (1.0 - edge_w) * ( + (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + ) with horizontal(region[i_end + 1, j_start + 1 : j_end + 1]): - qout = edge_e * ((qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / (dxa[-1, -1] + dxa[0, -1])) + (1.0 - edge_e) * ((qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa)) - + qout = edge_e * ( + (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) + / (dxa[-1, -1] + dxa[0, -1]) + ) + (1.0 - edge_e) * ( + (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + ) + with horizontal(region[i_start + 1 : i_end + 1, j_start]): - qout = edge_s * ((qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) / (dya[-1, -1] + dya[-1, 0])) + (1.0 - edge_s) * ((qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya)) + qout = edge_s * ( + (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) + / (dya[-1, -1] + dya[-1, 0]) + ) + (1.0 - edge_s) * ( + (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + ) with horizontal(region[i_start + 1 : i_end + 1, j_end + 1]): - qout = edge_n * ((qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) / (dya[-1, -1] + dya[-1, 0])) + (1.0 - edge_n) * ((qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya)) + qout = edge_n * ( + (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) + / (dya[-1, -1] + dya[-1, 0]) + ) + (1.0 - edge_n) * ( + (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + ) # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) @@ -374,14 +394,26 @@ def a2b_interpolation( qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): - qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0]))) + qxx = c1 * (qx[0, -1, 0] + qx) + c2 * ( + qout[0, -1, 0] + + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) + ) with horizontal(region[:, j_end]): - qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, 1, 0] + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0]))) + qxx = c1 * (qx[0, -1, 0] + qx) + c2 * ( + qout[0, 1, 0] + + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0])) + ) qyy = a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) with horizontal(region[i_start + 1, :]): - qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[-1, 0, 0] + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0]))) + qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * ( + qout[-1, 0, 0] + + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) + ) with horizontal(region[i_end, :]): - qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[1, 0, 0] + ( a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]))) + qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * ( + qout[1, 0, 0] + + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) + ) with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): qout = 0.5 * (qxx + qyy) if __INLINED(REPLACE): @@ -404,7 +436,6 @@ def qout_y_edge( qout[0, 0, 0] = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 - def qout_edges_x( q2: FloatField, qin: FloatField, @@ -458,7 +489,7 @@ def __init__( """ assert grid_type < 3 self.grid = spec.grid - + shape = self.grid.domain_shape_full(add=(1, 1, 1)) if nk is None: nk = self.grid.npz - kstart @@ -468,13 +499,9 @@ def __init__( self._tmp_qy = utils.make_storage_from_shape(shape) self._tmp_qxx = utils.make_storage_from_shape(shape) self._tmp_qyy = utils.make_storage_from_shape(shape) - origin=(self.grid.is_, self.grid.js, kstart) - domain=(self.grid.nic + 1, self.grid.njc + 1, nk) - ax_offsets = axis_offsets( - self.grid, - origin, - domain - ) + origin = (self.grid.is_, self.grid.js, kstart) + domain = (self.grid.nic + 1, self.grid.njc + 1, nk) + ax_offsets = axis_offsets(self.grid, origin, domain) self._south_corners_stencil = FrozenStencil( _south_corners, externals=ax_offsets, @@ -487,9 +514,12 @@ def __init__( origin=origin, domain=domain, ) - + self._a2b_interpolation_stencil = FrozenStencil( - a2b_interpolation, externals={'REPLACE': replace, **ax_offsets}, origin=origin, domain=domain + a2b_interpolation, + externals={"REPLACE": replace, **ax_offsets}, + origin=origin, + domain=domain, ) # TODO @@ -535,6 +565,9 @@ def __call__(self, qin: FloatField, qout: FloatField): self._tmp_qxx, self._tmp_qyy, self.grid.dxa, - self.grid.dya, self._edge_w, self._edge_e,self.grid.edge_s, self.grid.edge_n, + self.grid.dya, + self._edge_w, + self._edge_e, + self.grid.edge_s, + self.grid.edge_n, ) - From b102c2c0897ebbfcec157a73e21533a95d7e86fb Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 20:13:26 -0700 Subject: [PATCH 22/45] a2b not validating again, splitting apart qx, qxx and final qout --- fv3core/stencils/a2b_ord4.py | 106 +++++++++++++++++------------------ 1 file changed, 50 insertions(+), 56 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 927a7d1b2..8d609fc99 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -239,13 +239,11 @@ def lagrange_x_func(qy): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) -def a2b_interpolation( +def a2b_interpolation_prep( qin: FloatField, qout: FloatField, qx: FloatField, qy: FloatField, - qxx: FloatField, - qyy: FloatField, dxa: FloatFieldIJ, dya: FloatFieldIJ, edge_w: FloatFieldIJ, @@ -253,7 +251,7 @@ def a2b_interpolation( edge_s: FloatFieldI, edge_n: FloatFieldI, ): - from __externals__ import REPLACE, i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): # qout_edges_x @@ -391,6 +389,17 @@ def a2b_interpolation( 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) - (dya[0, -1] / dya * qy[0, 1, 0] + qy[0, -1, 0]) ) / (2.0 + 2.0 * dya[0, -1] / dya) + +def a2b_interpolation_2nd( + qout: FloatField, + qx: FloatField, + qy: FloatField, + qxx: FloatField, + qyy: FloatField, +): + from __externals__ import i_end, i_start, j_end, j_start + + with computation(PARALLEL), interval(...): qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): @@ -414,62 +423,22 @@ def a2b_interpolation( qout[1, 0, 0] + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) ) - with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): - qout = 0.5 * (qxx + qyy) - if __INLINED(REPLACE): - qin = qout - - -def qout_x_edge( - qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ, qout: FloatField -): - with computation(PARALLEL), interval(...): - q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - qout[0, 0, 0] = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 - -def qout_y_edge( - qin: FloatField, dya: FloatFieldIJ, edge_s: FloatFieldI, qout: FloatField -): - with computation(PARALLEL), interval(...): - q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - qout[0, 0, 0] = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 - - -def qout_edges_x( - q2: FloatField, +def a2b_interpolation( qin: FloatField, qout: FloatField, - dxa: FloatFieldIJ, - edge_w: FloatFieldIJ, - edge_e: FloatFieldIJ, + qxx: FloatField, + qyy: FloatField, ): - from __externals__ import i_end, i_start + from __externals__ import REPLACE, i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): - q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - with horizontal(region[i_start, :]): - qout = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 - with horizontal(region[i_end + 1, :]): - qout = edge_e * q2[0, -1, 0] + (1.0 - edge_e) * q2 - + with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): + qout = 0.5 * (qxx + qyy) + if __INLINED(REPLACE): + qin = qout -def qout_edges_y( - q1: FloatField, - qin: FloatField, - qout: FloatField, - dya: FloatFieldIJ, - edge_s: FloatFieldI, - edge_n: FloatFieldI, -): - from __externals__ import j_end, j_start - with computation(PARALLEL), interval(...): - q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - with horizontal(region[:, j_start]): - qout = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 - with horizontal(region[:, j_end + 1]): - qout = edge_n * q1[-1, 0, 0] + (1.0 - edge_n) * q1 class AGrid2BGridFourthOrder: @@ -514,7 +483,18 @@ def __init__( origin=origin, domain=domain, ) - + self._a2b_interpolation_prep_stencil = FrozenStencil( + a2b_interpolation_prep, + externals=ax_offsets, + origin=origin, + domain=domain, + ) + self._a2b_interpolation_2nd_stencil = FrozenStencil( + a2b_interpolation_2nd, + externals=ax_offsets, + origin=origin, + domain=domain, + ) self._a2b_interpolation_stencil = FrozenStencil( a2b_interpolation, externals={"REPLACE": replace, **ax_offsets}, @@ -557,13 +537,11 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid2, ) - self._a2b_interpolation_stencil( + self._a2b_interpolation_prep_stencil( qin, qout, self._tmp_qx, self._tmp_qy, - self._tmp_qxx, - self._tmp_qyy, self.grid.dxa, self.grid.dya, self._edge_w, @@ -571,3 +549,19 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.edge_s, self.grid.edge_n, ) + + self._a2b_interpolation_2nd_stencil( + qout, + self._tmp_qx, + self._tmp_qy, + self._tmp_qxx, + self._tmp_qyy, + ) + + self._a2b_interpolation_stencil( + qin, + qout, + self._tmp_qxx, + self._tmp_qyy, + + ) From 8abb14c07ad47e9822b7a191f48b0c2357dda11f Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 20:54:59 -0700 Subject: [PATCH 23/45] Revert "a2b not validating again, splitting apart qx, qxx and final qout" This reverts commit b102c2c0897ebbfcec157a73e21533a95d7e86fb. --- fv3core/stencils/a2b_ord4.py | 106 ++++++++++++++++++----------------- 1 file changed, 56 insertions(+), 50 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 8d609fc99..927a7d1b2 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -239,11 +239,13 @@ def lagrange_x_func(qy): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) -def a2b_interpolation_prep( +def a2b_interpolation( qin: FloatField, qout: FloatField, qx: FloatField, qy: FloatField, + qxx: FloatField, + qyy: FloatField, dxa: FloatFieldIJ, dya: FloatFieldIJ, edge_w: FloatFieldIJ, @@ -251,7 +253,7 @@ def a2b_interpolation_prep( edge_s: FloatFieldI, edge_n: FloatFieldI, ): - from __externals__ import i_end, i_start, j_end, j_start + from __externals__ import REPLACE, i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): # qout_edges_x @@ -389,17 +391,6 @@ def a2b_interpolation_prep( 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) - (dya[0, -1] / dya * qy[0, 1, 0] + qy[0, -1, 0]) ) / (2.0 + 2.0 * dya[0, -1] / dya) - -def a2b_interpolation_2nd( - qout: FloatField, - qx: FloatField, - qy: FloatField, - qxx: FloatField, - qyy: FloatField, -): - from __externals__ import i_end, i_start, j_end, j_start - - with computation(PARALLEL), interval(...): qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): @@ -423,22 +414,62 @@ def a2b_interpolation_2nd( qout[1, 0, 0] + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) ) + with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): + qout = 0.5 * (qxx + qyy) + if __INLINED(REPLACE): + qin = qout -def a2b_interpolation( + +def qout_x_edge( + qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ, qout: FloatField +): + with computation(PARALLEL), interval(...): + q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + qout[0, 0, 0] = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 + + +def qout_y_edge( + qin: FloatField, dya: FloatFieldIJ, edge_s: FloatFieldI, qout: FloatField +): + with computation(PARALLEL), interval(...): + q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + qout[0, 0, 0] = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 + + +def qout_edges_x( + q2: FloatField, qin: FloatField, qout: FloatField, - qxx: FloatField, - qyy: FloatField, + dxa: FloatFieldIJ, + edge_w: FloatFieldIJ, + edge_e: FloatFieldIJ, ): - from __externals__ import REPLACE, i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start with computation(PARALLEL), interval(...): - with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): - qout = 0.5 * (qxx + qyy) - if __INLINED(REPLACE): - qin = qout + q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + with horizontal(region[i_start, :]): + qout = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 + with horizontal(region[i_end + 1, :]): + qout = edge_e * q2[0, -1, 0] + (1.0 - edge_e) * q2 + +def qout_edges_y( + q1: FloatField, + qin: FloatField, + qout: FloatField, + dya: FloatFieldIJ, + edge_s: FloatFieldI, + edge_n: FloatFieldI, +): + from __externals__ import j_end, j_start + with computation(PARALLEL), interval(...): + q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + with horizontal(region[:, j_start]): + qout = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 + with horizontal(region[:, j_end + 1]): + qout = edge_n * q1[-1, 0, 0] + (1.0 - edge_n) * q1 class AGrid2BGridFourthOrder: @@ -483,18 +514,7 @@ def __init__( origin=origin, domain=domain, ) - self._a2b_interpolation_prep_stencil = FrozenStencil( - a2b_interpolation_prep, - externals=ax_offsets, - origin=origin, - domain=domain, - ) - self._a2b_interpolation_2nd_stencil = FrozenStencil( - a2b_interpolation_2nd, - externals=ax_offsets, - origin=origin, - domain=domain, - ) + self._a2b_interpolation_stencil = FrozenStencil( a2b_interpolation, externals={"REPLACE": replace, **ax_offsets}, @@ -537,11 +557,13 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid2, ) - self._a2b_interpolation_prep_stencil( + self._a2b_interpolation_stencil( qin, qout, self._tmp_qx, self._tmp_qy, + self._tmp_qxx, + self._tmp_qyy, self.grid.dxa, self.grid.dya, self._edge_w, @@ -549,19 +571,3 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.edge_s, self.grid.edge_n, ) - - self._a2b_interpolation_2nd_stencil( - qout, - self._tmp_qx, - self._tmp_qy, - self._tmp_qxx, - self._tmp_qyy, - ) - - self._a2b_interpolation_stencil( - qin, - qout, - self._tmp_qxx, - self._tmp_qyy, - - ) From b6744926441c60b2a5f3d8a63138f2913ef55420 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 21:53:33 -0700 Subject: [PATCH 24/45] cleanup unused methods --- fv3core/stencils/a2b_ord4.py | 51 +----------------------------------- 1 file changed, 1 insertion(+), 50 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 927a7d1b2..fd8bf2424 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -420,57 +420,8 @@ def a2b_interpolation( qin = qout -def qout_x_edge( - qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ, qout: FloatField -): - with computation(PARALLEL), interval(...): - q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - qout[0, 0, 0] = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 - - -def qout_y_edge( - qin: FloatField, dya: FloatFieldIJ, edge_s: FloatFieldI, qout: FloatField -): - with computation(PARALLEL), interval(...): - q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - qout[0, 0, 0] = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 - - -def qout_edges_x( - q2: FloatField, - qin: FloatField, - qout: FloatField, - dxa: FloatFieldIJ, - edge_w: FloatFieldIJ, - edge_e: FloatFieldIJ, -): - from __externals__ import i_end, i_start - - with computation(PARALLEL), interval(...): - q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - with horizontal(region[i_start, :]): - qout = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 - with horizontal(region[i_end + 1, :]): - qout = edge_e * q2[0, -1, 0] + (1.0 - edge_e) * q2 -def qout_edges_y( - q1: FloatField, - qin: FloatField, - qout: FloatField, - dya: FloatFieldIJ, - edge_s: FloatFieldI, - edge_n: FloatFieldI, -): - from __externals__ import j_end, j_start - - with computation(PARALLEL), interval(...): - q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - with horizontal(region[:, j_start]): - qout = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 - with horizontal(region[:, j_end + 1]): - qout = edge_n * q1[-1, 0, 0] + (1.0 - edge_n) * q1 - class AGrid2BGridFourthOrder: """ @@ -556,7 +507,7 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid1, self.grid.bgrid2, ) - + self._a2b_interpolation_stencil( qin, qout, From d62e2fda5e16b6cc046a397603442c00c0ea3b8c Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 22:29:21 -0700 Subject: [PATCH 25/45] Revert "cleanup unused methods" This reverts commit b6744926441c60b2a5f3d8a63138f2913ef55420. --- fv3core/stencils/a2b_ord4.py | 51 +++++++++++++++++++++++++++++++++++- 1 file changed, 50 insertions(+), 1 deletion(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index fd8bf2424..927a7d1b2 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -420,8 +420,57 @@ def a2b_interpolation( qin = qout +def qout_x_edge( + qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ, qout: FloatField +): + with computation(PARALLEL), interval(...): + q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + qout[0, 0, 0] = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 + + +def qout_y_edge( + qin: FloatField, dya: FloatFieldIJ, edge_s: FloatFieldI, qout: FloatField +): + with computation(PARALLEL), interval(...): + q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + qout[0, 0, 0] = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 + + +def qout_edges_x( + q2: FloatField, + qin: FloatField, + qout: FloatField, + dxa: FloatFieldIJ, + edge_w: FloatFieldIJ, + edge_e: FloatFieldIJ, +): + from __externals__ import i_end, i_start + + with computation(PARALLEL), interval(...): + q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + with horizontal(region[i_start, :]): + qout = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 + with horizontal(region[i_end + 1, :]): + qout = edge_e * q2[0, -1, 0] + (1.0 - edge_e) * q2 +def qout_edges_y( + q1: FloatField, + qin: FloatField, + qout: FloatField, + dya: FloatFieldIJ, + edge_s: FloatFieldI, + edge_n: FloatFieldI, +): + from __externals__ import j_end, j_start + + with computation(PARALLEL), interval(...): + q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + with horizontal(region[:, j_start]): + qout = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 + with horizontal(region[:, j_end + 1]): + qout = edge_n * q1[-1, 0, 0] + (1.0 - edge_n) * q1 + class AGrid2BGridFourthOrder: """ @@ -507,7 +556,7 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid1, self.grid.bgrid2, ) - + self._a2b_interpolation_stencil( qin, qout, From 07d2fa81c13cdbe80d0337e674844f6402800942 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Tue, 18 May 2021 22:29:33 -0700 Subject: [PATCH 26/45] Revert "Revert "a2b not validating again, splitting apart qx, qxx and final qout"" This reverts commit 8abb14c07ad47e9822b7a191f48b0c2357dda11f. --- fv3core/stencils/a2b_ord4.py | 106 +++++++++++++++++------------------ 1 file changed, 50 insertions(+), 56 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 927a7d1b2..8d609fc99 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -239,13 +239,11 @@ def lagrange_x_func(qy): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) -def a2b_interpolation( +def a2b_interpolation_prep( qin: FloatField, qout: FloatField, qx: FloatField, qy: FloatField, - qxx: FloatField, - qyy: FloatField, dxa: FloatFieldIJ, dya: FloatFieldIJ, edge_w: FloatFieldIJ, @@ -253,7 +251,7 @@ def a2b_interpolation( edge_s: FloatFieldI, edge_n: FloatFieldI, ): - from __externals__ import REPLACE, i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): # qout_edges_x @@ -391,6 +389,17 @@ def a2b_interpolation( 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) - (dya[0, -1] / dya * qy[0, 1, 0] + qy[0, -1, 0]) ) / (2.0 + 2.0 * dya[0, -1] / dya) + +def a2b_interpolation_2nd( + qout: FloatField, + qx: FloatField, + qy: FloatField, + qxx: FloatField, + qyy: FloatField, +): + from __externals__ import i_end, i_start, j_end, j_start + + with computation(PARALLEL), interval(...): qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) with horizontal(region[:, j_start + 1]): @@ -414,62 +423,22 @@ def a2b_interpolation( qout[1, 0, 0] + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) ) - with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): - qout = 0.5 * (qxx + qyy) - if __INLINED(REPLACE): - qin = qout - - -def qout_x_edge( - qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ, qout: FloatField -): - with computation(PARALLEL), interval(...): - q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - qout[0, 0, 0] = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 - -def qout_y_edge( - qin: FloatField, dya: FloatFieldIJ, edge_s: FloatFieldI, qout: FloatField -): - with computation(PARALLEL), interval(...): - q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - qout[0, 0, 0] = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 - - -def qout_edges_x( - q2: FloatField, +def a2b_interpolation( qin: FloatField, qout: FloatField, - dxa: FloatFieldIJ, - edge_w: FloatFieldIJ, - edge_e: FloatFieldIJ, + qxx: FloatField, + qyy: FloatField, ): - from __externals__ import i_end, i_start + from __externals__ import REPLACE, i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): - q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - with horizontal(region[i_start, :]): - qout = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 - with horizontal(region[i_end + 1, :]): - qout = edge_e * q2[0, -1, 0] + (1.0 - edge_e) * q2 - + with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): + qout = 0.5 * (qxx + qyy) + if __INLINED(REPLACE): + qin = qout -def qout_edges_y( - q1: FloatField, - qin: FloatField, - qout: FloatField, - dya: FloatFieldIJ, - edge_s: FloatFieldI, - edge_n: FloatFieldI, -): - from __externals__ import j_end, j_start - with computation(PARALLEL), interval(...): - q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - with horizontal(region[:, j_start]): - qout = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 - with horizontal(region[:, j_end + 1]): - qout = edge_n * q1[-1, 0, 0] + (1.0 - edge_n) * q1 class AGrid2BGridFourthOrder: @@ -514,7 +483,18 @@ def __init__( origin=origin, domain=domain, ) - + self._a2b_interpolation_prep_stencil = FrozenStencil( + a2b_interpolation_prep, + externals=ax_offsets, + origin=origin, + domain=domain, + ) + self._a2b_interpolation_2nd_stencil = FrozenStencil( + a2b_interpolation_2nd, + externals=ax_offsets, + origin=origin, + domain=domain, + ) self._a2b_interpolation_stencil = FrozenStencil( a2b_interpolation, externals={"REPLACE": replace, **ax_offsets}, @@ -557,13 +537,11 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid2, ) - self._a2b_interpolation_stencil( + self._a2b_interpolation_prep_stencil( qin, qout, self._tmp_qx, self._tmp_qy, - self._tmp_qxx, - self._tmp_qyy, self.grid.dxa, self.grid.dya, self._edge_w, @@ -571,3 +549,19 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.edge_s, self.grid.edge_n, ) + + self._a2b_interpolation_2nd_stencil( + qout, + self._tmp_qx, + self._tmp_qy, + self._tmp_qxx, + self._tmp_qyy, + ) + + self._a2b_interpolation_stencil( + qin, + qout, + self._tmp_qxx, + self._tmp_qyy, + + ) From 451ea0de8932af1d8576f6a56f42a27d0e21fd85 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 19 May 2021 12:12:29 -0700 Subject: [PATCH 27/45] the edges of qxx and qyy recompute qout rather than waiting for the qouts to be computed, also removed some redundant methods --- fv3core/stencils/a2b_ord4.py | 161 +++++++++++++++-------------------- 1 file changed, 68 insertions(+), 93 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 8d609fc99..5161222e1 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -239,11 +239,13 @@ def lagrange_x_func(qy): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) -def a2b_interpolation_prep( +def a2b_interpolation( qin: FloatField, qout: FloatField, qx: FloatField, qy: FloatField, + qxx: FloatField, + qyy: FloatField, dxa: FloatFieldIJ, dya: FloatFieldIJ, edge_w: FloatFieldIJ, @@ -251,40 +253,10 @@ def a2b_interpolation_prep( edge_s: FloatFieldI, edge_n: FloatFieldI, ): - from __externals__ import i_end, i_start, j_end, j_start + from __externals__ import REPLACE, i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): - # qout_edges_x - with horizontal(region[i_start, j_start + 1 : j_end + 1]): - qout = edge_w * ( - (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) - / (dxa[-1, -1] + dxa[0, -1]) - ) + (1.0 - edge_w) * ( - (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - ) - with horizontal(region[i_end + 1, j_start + 1 : j_end + 1]): - qout = edge_e * ( - (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) - / (dxa[-1, -1] + dxa[0, -1]) - ) + (1.0 - edge_e) * ( - (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - ) - - with horizontal(region[i_start + 1 : i_end + 1, j_start]): - qout = edge_s * ( - (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) - / (dya[-1, -1] + dya[-1, 0]) - ) + (1.0 - edge_s) * ( - (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - ) - with horizontal(region[i_start + 1 : i_end + 1, j_end + 1]): - qout = edge_n * ( - (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) - / (dya[-1, -1] + dya[-1, 0]) - ) + (1.0 - edge_n) * ( - (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - ) - + # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) with horizontal(region[i_start, :]): @@ -389,50 +361,78 @@ def a2b_interpolation_prep( 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) - (dya[0, -1] / dya * qy[0, 1, 0] + qy[0, -1, 0]) ) / (2.0 + 2.0 * dya[0, -1] / dya) - -def a2b_interpolation_2nd( - qout: FloatField, - qx: FloatField, - qy: FloatField, - qxx: FloatField, - qyy: FloatField, -): - from __externals__ import i_end, i_start, j_end, j_start - - with computation(PARALLEL), interval(...): - + # qout_edges_x + with horizontal(region[i_start, j_start + 1 : j_end + 1]): + qout = edge_w * ( + (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) + / (dxa[-1, -1] + dxa[0, -1]) + ) + (1.0 - edge_w) * ( + (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + ) + with horizontal(region[i_end + 1, j_start + 1 : j_end + 1]): + qout = edge_e * ( + (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) + / (dxa[-1, -1] + dxa[0, -1]) + ) + (1.0 - edge_e) * ( + (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + ) + # qout edges y + with horizontal(region[i_start + 1 : i_end + 1, j_start]): + qout = edge_s * ( + (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) + / (dya[-1, -1] + dya[-1, 0]) + ) + (1.0 - edge_s) * ( + (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + ) + with horizontal(region[i_start + 1 : i_end + 1, j_end + 1]): + qout = edge_n * ( + (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) + / (dya[-1, -1] + dya[-1, 0]) + ) + (1.0 - edge_n) * ( + (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + ) qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) - with horizontal(region[:, j_start + 1]): + with horizontal(region[i_start + 1 : i_end + 1, j_start + 1]): qxx = c1 * (qx[0, -1, 0] + qx) + c2 * ( - qout[0, -1, 0] + (edge_s * ( + (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) + / (dya[-1, -2] + dya[-1, -1]) + ) + (1.0 - edge_s) * ( + (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) / (dya[0, -2] + dya[0, -1]) + )) + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) ) - with horizontal(region[:, j_end]): + with horizontal(region[i_start + 1 : i_end + 1, j_end]): qxx = c1 * (qx[0, -1, 0] + qx) + c2 * ( - qout[0, 1, 0] + ( edge_n * ( + (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) + / (dya[-1, 0] + dya[-1, 1]) + ) + (1.0 - edge_n) * ( + (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) / (dya[0, 0] + dya[0, 1]) + )) + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0])) ) qyy = a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) - with horizontal(region[i_start + 1, :]): + with horizontal(region[i_start + 1, j_start + 1 : j_end + 1]): qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * ( - qout[-1, 0, 0] + (edge_w * ( + (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) + / (dxa[-2, -1] + dxa[-1, -1]) + ) + (1.0 - edge_w) * ( + (qin[-2, 0, 0] * dxa[-1, 0] + qin[-1, 0, 0] * dxa[-2, 0]) / (dxa[-2, 0] + dxa[-1, 0]) + )) + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) ) - with horizontal(region[i_end, :]): + with horizontal(region[i_end, j_start + 1 : j_end + 1]): qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * ( - qout[1, 0, 0] + (edge_e * ( + (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) + / (dxa[0, -1] + dxa[1, -1]) + ) + (1.0 - edge_e) * ( + (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) / (dxa[0, 0] + dxa[1, 0]) + )) + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) ) - -def a2b_interpolation( - qin: FloatField, - qout: FloatField, - qxx: FloatField, - qyy: FloatField, -): - from __externals__ import REPLACE, i_end, i_start, j_end, j_start - - with computation(PARALLEL), interval(...): with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): qout = 0.5 * (qxx + qyy) if __INLINED(REPLACE): @@ -440,7 +440,6 @@ def a2b_interpolation( - class AGrid2BGridFourthOrder: """ Fortran name is a2b_ord4, test module is A2B_Ord4 @@ -483,18 +482,7 @@ def __init__( origin=origin, domain=domain, ) - self._a2b_interpolation_prep_stencil = FrozenStencil( - a2b_interpolation_prep, - externals=ax_offsets, - origin=origin, - domain=domain, - ) - self._a2b_interpolation_2nd_stencil = FrozenStencil( - a2b_interpolation_2nd, - externals=ax_offsets, - origin=origin, - domain=domain, - ) + self._a2b_interpolation_stencil = FrozenStencil( a2b_interpolation, externals={"REPLACE": replace, **ax_offsets}, @@ -537,11 +525,13 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid2, ) - self._a2b_interpolation_prep_stencil( + self._a2b_interpolation_stencil( qin, qout, self._tmp_qx, self._tmp_qy, + self._tmp_qxx, + self._tmp_qyy, self.grid.dxa, self.grid.dya, self._edge_w, @@ -549,19 +539,4 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.edge_s, self.grid.edge_n, ) - - self._a2b_interpolation_2nd_stencil( - qout, - self._tmp_qx, - self._tmp_qy, - self._tmp_qxx, - self._tmp_qyy, - ) - - self._a2b_interpolation_stencil( - qin, - qout, - self._tmp_qxx, - self._tmp_qyy, - - ) + From 8d6eea139763a49b87ade583fa3f757b2c1db5d6 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 19 May 2021 14:36:23 -0700 Subject: [PATCH 28/45] failing 54 ranks but will it pass 6 ranks with gt backend --- fv3core/stencils/a2b_ord4.py | 51 ++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 5161222e1..ef9aa8e2d 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -439,6 +439,57 @@ def a2b_interpolation( qin = qout +def qout_x_edge( + qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ, qout: FloatField +): + with computation(PARALLEL), interval(...): + q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + qout[0, 0, 0] = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 + + +def qout_y_edge( + qin: FloatField, dya: FloatFieldIJ, edge_s: FloatFieldI, qout: FloatField +): + with computation(PARALLEL), interval(...): + q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + qout[0, 0, 0] = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 + + +def qout_edges_x( + q2: FloatField, + qin: FloatField, + qout: FloatField, + dxa: FloatFieldIJ, + edge_w: FloatFieldIJ, + edge_e: FloatFieldIJ, +): + from __externals__ import i_end, i_start + + with computation(PARALLEL), interval(...): + q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + with horizontal(region[i_start, :]): + qout = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 + with horizontal(region[i_end + 1, :]): + qout = edge_e * q2[0, -1, 0] + (1.0 - edge_e) * q2 + + +def qout_edges_y( + q1: FloatField, + qin: FloatField, + qout: FloatField, + dya: FloatFieldIJ, + edge_s: FloatFieldI, + edge_n: FloatFieldI, +): + from __externals__ import j_end, j_start + + with computation(PARALLEL), interval(...): + q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + with horizontal(region[:, j_start]): + qout = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 + with horizontal(region[:, j_end + 1]): + qout = edge_n * q1[-1, 0, 0] + (1.0 - edge_n) * q1 + class AGrid2BGridFourthOrder: """ From 7af235ab205307787f9432683d71505de9854eaa Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 19 May 2021 15:26:37 -0700 Subject: [PATCH 29/45] failing 54 ranks but will it pass 6 ranks with gt backend --- fv3core/stencils/a2b_ord4.py | 117 +++++++++++++++-------------------- 1 file changed, 51 insertions(+), 66 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index ef9aa8e2d..bea38c96c 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -239,24 +239,18 @@ def lagrange_x_func(qy): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) -def a2b_interpolation( +def a2b_interpolation_qxqy( qin: FloatField, qout: FloatField, qx: FloatField, qy: FloatField, - qxx: FloatField, - qyy: FloatField, dxa: FloatFieldIJ, dya: FloatFieldIJ, - edge_w: FloatFieldIJ, - edge_e: FloatFieldIJ, - edge_s: FloatFieldI, - edge_n: FloatFieldI, ): - from __externals__ import REPLACE, i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): - + #qout = qin # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) with horizontal(region[i_start, :]): @@ -361,13 +355,30 @@ def a2b_interpolation( 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) - (dya[0, -1] / dya * qy[0, 1, 0] + qy[0, -1, 0]) ) / (2.0 + 2.0 * dya[0, -1] / dya) +def a2b_interpolation( + qin: FloatField, + qout: FloatField, + qx: FloatField, + qy: FloatField, + qxx: FloatField, + qyy: FloatField, + dxa: FloatFieldIJ, + dya: FloatFieldIJ, + edge_w: FloatFieldIJ, + edge_e: FloatFieldIJ, + edge_s: FloatFieldI, + edge_n: FloatFieldI, +): + from __externals__ import REPLACE, i_end, i_start, j_end, j_start, local_is, local_js, local_ie,local_je + + with computation(PARALLEL), interval(...): # qout_edges_x - with horizontal(region[i_start, j_start + 1 : j_end + 1]): + with horizontal(region[i_start,j_start + 1 : j_end + 1]): qout = edge_w * ( (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / (dxa[-1, -1] + dxa[0, -1]) ) + (1.0 - edge_w) * ( - (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + (qin[-1, 0, 0] * dxa[0, 0] + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa[0, 0]) ) with horizontal(region[i_end + 1, j_start + 1 : j_end + 1]): qout = edge_e * ( @@ -413,7 +424,7 @@ def a2b_interpolation( + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0])) ) qyy = a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) - with horizontal(region[i_start + 1, j_start + 1 : j_end + 1]): + with horizontal(region[i_start + 1,j_start + 1 : j_end + 1]): qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * ( (edge_w * ( (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) @@ -439,57 +450,6 @@ def a2b_interpolation( qin = qout -def qout_x_edge( - qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ, qout: FloatField -): - with computation(PARALLEL), interval(...): - q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - qout[0, 0, 0] = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 - - -def qout_y_edge( - qin: FloatField, dya: FloatFieldIJ, edge_s: FloatFieldI, qout: FloatField -): - with computation(PARALLEL), interval(...): - q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - qout[0, 0, 0] = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 - - -def qout_edges_x( - q2: FloatField, - qin: FloatField, - qout: FloatField, - dxa: FloatFieldIJ, - edge_w: FloatFieldIJ, - edge_e: FloatFieldIJ, -): - from __externals__ import i_end, i_start - - with computation(PARALLEL), interval(...): - q2 = (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - with horizontal(region[i_start, :]): - qout = edge_w * q2[0, -1, 0] + (1.0 - edge_w) * q2 - with horizontal(region[i_end + 1, :]): - qout = edge_e * q2[0, -1, 0] + (1.0 - edge_e) * q2 - - -def qout_edges_y( - q1: FloatField, - qin: FloatField, - qout: FloatField, - dya: FloatFieldIJ, - edge_s: FloatFieldI, - edge_n: FloatFieldI, -): - from __externals__ import j_end, j_start - - with computation(PARALLEL), interval(...): - q1 = (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - with horizontal(region[:, j_start]): - qout = edge_s * q1[-1, 0, 0] + (1.0 - edge_s) * q1 - with horizontal(region[:, j_end + 1]): - qout = edge_n * q1[-1, 0, 0] + (1.0 - edge_n) * q1 - class AGrid2BGridFourthOrder: """ @@ -533,7 +493,22 @@ def __init__( origin=origin, domain=domain, ) - + origin_prep = (self.grid.is_ - 1, self.grid.js - 1, kstart) + domain_prep = (self.grid.nic + 3, self.grid.njc + 3, nk) + ax_offsets_prep = axis_offsets(self.grid, origin_prep, domain_prep) + self._a2b_interpolation_qxqy_stencil = FrozenStencil( + a2b_interpolation_qxqy, + externals=ax_offsets_prep, + origin=origin_prep, + domain=domain_prep, + ) + #js = self.grid.js + 1 if self.grid.south_edge else self.grid.js + #je = self.grid.je if self.grid.north_edge else self.grid.je + 1 + #is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ + #ie = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 + #origin = (is_, js, kstart) + #domain = (ie - is_ + 1, je - js + 1, nk) + #ax_offsets = axis_offsets(self.grid, origin, domain) self._a2b_interpolation_stencil = FrozenStencil( a2b_interpolation, externals={"REPLACE": replace, **ax_offsets}, @@ -575,7 +550,15 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid1, self.grid.bgrid2, ) - + self._a2b_interpolation_qxqy_stencil( + qin, + qout, + self._tmp_qx, + self._tmp_qy, + self.grid.dxa, + self.grid.dya, + ) + print('going into interpolation', qout[3,4,3]) self._a2b_interpolation_stencil( qin, qout, @@ -590,4 +573,6 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.edge_s, self.grid.edge_n, ) - + print('after interpolation', qout[3,4,3]) + #after 3.4542028400165024e-06 + # 3.700316608845214e-06 From 40f619fc8460ece61967f14e7c02329278802a1a Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 19 May 2021 16:06:03 -0700 Subject: [PATCH 30/45] split qx and qy stencils to they can go wider --- fv3core/stencils/a2b_ord4.py | 73 ++++++++++++++++++++++-------------- 1 file changed, 44 insertions(+), 29 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index bea38c96c..05880e5cb 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -239,18 +239,14 @@ def lagrange_x_func(qy): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) -def a2b_interpolation_qxqy( +def a2b_interpolation_qx( qin: FloatField, - qout: FloatField, qx: FloatField, - qy: FloatField, dxa: FloatFieldIJ, - dya: FloatFieldIJ, ): from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): - #qout = qin # ppm_volume_mean_x qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) with horizontal(region[i_start, :]): @@ -313,6 +309,15 @@ def a2b_interpolation_qxqy( ) ) / (2.0 + 2.0 * dxa[-1, 0] / dxa) +def a2b_interpolation_qy( + qin: FloatField, + qy: FloatField, + dya: FloatFieldIJ, +): + from __externals__ import i_end, i_start, j_end, j_start + + with computation(PARALLEL), interval(...): + # ppm_volume_mean_y qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) with horizontal(region[:, j_start]): @@ -355,6 +360,7 @@ def a2b_interpolation_qxqy( 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) - (dya[0, -1] / dya * qy[0, 1, 0] + qy[0, -1, 0]) ) / (2.0 + 2.0 * dya[0, -1] / dya) + def a2b_interpolation( qin: FloatField, qout: FloatField, @@ -369,7 +375,7 @@ def a2b_interpolation( edge_s: FloatFieldI, edge_n: FloatFieldI, ): - from __externals__ import REPLACE, i_end, i_start, j_end, j_start, local_is, local_js, local_ie,local_je + from __externals__ import REPLACE, i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): # qout_edges_x @@ -444,6 +450,8 @@ def a2b_interpolation( )) + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) ) + + with computation(PARALLEL), interval(...): with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): qout = 0.5 * (qxx + qyy) if __INLINED(REPLACE): @@ -493,29 +501,33 @@ def __init__( origin=origin, domain=domain, ) - origin_prep = (self.grid.is_ - 1, self.grid.js - 1, kstart) - domain_prep = (self.grid.nic + 3, self.grid.njc + 3, nk) - ax_offsets_prep = axis_offsets(self.grid, origin_prep, domain_prep) - self._a2b_interpolation_qxqy_stencil = FrozenStencil( - a2b_interpolation_qxqy, - externals=ax_offsets_prep, - origin=origin_prep, - domain=domain_prep, + origin_prep_x = (self.grid.is_ - 1, self.grid.js - 2, kstart) + domain_prep_x = (self.grid.nic + 3, self.grid.njc + 4, nk) + ax_offsets_prep_x = axis_offsets(self.grid, origin_prep_x, domain_prep_x) + self._a2b_interpolation_qx_stencil = FrozenStencil( + a2b_interpolation_qx, + externals=ax_offsets_prep_x, + origin=origin_prep_x, + domain=domain_prep_x, ) - #js = self.grid.js + 1 if self.grid.south_edge else self.grid.js - #je = self.grid.je if self.grid.north_edge else self.grid.je + 1 - #is_ = self.grid.is_ + 1 if self.grid.west_edge else self.grid.is_ - #ie = self.grid.ie if self.grid.east_edge else self.grid.ie + 1 - #origin = (is_, js, kstart) - #domain = (ie - is_ + 1, je - js + 1, nk) - #ax_offsets = axis_offsets(self.grid, origin, domain) + origin_prep_y = (self.grid.is_ - 2, self.grid.js - 1, kstart) + domain_prep_y = (self.grid.nic + 4, self.grid.njc + 3, nk) + ax_offsets_prep_y = axis_offsets(self.grid, origin_prep_y, domain_prep_y) + self._a2b_interpolation_qy_stencil = FrozenStencil( + a2b_interpolation_qy, + externals=ax_offsets_prep_y, + origin=origin_prep_y, + domain=domain_prep_y, + ) + self._a2b_interpolation_stencil = FrozenStencil( a2b_interpolation, externals={"REPLACE": replace, **ax_offsets}, origin=origin, domain=domain, ) - + + # TODO # within regions, the edge_w and edge_w variables that are singleton in the # I dimension error, workaround is repeating the data, but the longterm @@ -550,15 +562,19 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid1, self.grid.bgrid2, ) - self._a2b_interpolation_qxqy_stencil( + + self._a2b_interpolation_qx_stencil( qin, - qout, self._tmp_qx, - self._tmp_qy, self.grid.dxa, + ) + + self._a2b_interpolation_qy_stencil( + qin, + self._tmp_qy, self.grid.dya, ) - print('going into interpolation', qout[3,4,3]) + self._a2b_interpolation_stencil( qin, qout, @@ -573,6 +589,5 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.edge_s, self.grid.edge_n, ) - print('after interpolation', qout[3,4,3]) - #after 3.4542028400165024e-06 - # 3.700316608845214e-06 + + From d99b98ddd42739769e1e8740093e5af930c5bff3 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 19 May 2021 16:31:24 -0700 Subject: [PATCH 31/45] wow, need to split it out even more, this is insane --- fv3core/stencils/a2b_ord4.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 05880e5cb..365f9446d 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -375,7 +375,7 @@ def a2b_interpolation( edge_s: FloatFieldI, edge_n: FloatFieldI, ): - from __externals__ import REPLACE, i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): # qout_edges_x @@ -450,7 +450,12 @@ def a2b_interpolation( )) + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) ) - +def a2b_interpolation_final( + qout: FloatField, + qxx: FloatField, + qyy: FloatField, +): + from __externals__ import REPLACE, i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): qout = 0.5 * (qxx + qyy) @@ -522,6 +527,12 @@ def __init__( self._a2b_interpolation_stencil = FrozenStencil( a2b_interpolation, + externals=ax_offsets, + origin=origin, + domain=domain, + ) + self._a2b_interpolation_final_stencil = FrozenStencil( + a2b_interpolation_final, externals={"REPLACE": replace, **ax_offsets}, origin=origin, domain=domain, @@ -589,5 +600,9 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.edge_s, self.grid.edge_n, ) - + self._a2b_interpolation_final_stencil( + qout, + self._tmp_qxx, + self._tmp_qyy, + ) From 40b9b96b6b6321ec2c1ba8977939d2214197f09c Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 19 May 2021 17:07:37 -0700 Subject: [PATCH 32/45] fix for nh_p_grad --- fv3core/stencils/a2b_ord4.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 365f9446d..c373781c6 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -451,6 +451,7 @@ def a2b_interpolation( + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) ) def a2b_interpolation_final( + qin: FloatField, qout: FloatField, qxx: FloatField, qyy: FloatField, @@ -601,6 +602,7 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.edge_n, ) self._a2b_interpolation_final_stencil( + qin, qout, self._tmp_qxx, self._tmp_qyy, From 4b5236023e4f45663eccad98fa69a926dad1f9a0 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 19 May 2021 17:40:26 -0700 Subject: [PATCH 33/45] remove qxx and qyy temporaries from a2b_ord4 --- fv3core/stencils/a2b_ord4.py | 188 +++++++++++++++++++++-------------- 1 file changed, 112 insertions(+), 76 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index c373781c6..6fb1a3b3a 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -58,20 +58,6 @@ def extrap_corner( return qa + x1 / (x2 - x1) * (qa - qb) -""" -def _a2b_corners(qin, qout, agrid1, agrid2, bgrid1, bgrid2): - with computation(PARALLEL), interval(...): - from __externals__ import i_start, j_start - - with horizontal(region[i_start, j_start]): - qout = _sw_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) - qout = _se_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) - qout = _ne_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) - qout = _nw_corner(qin, qout, agrid1, agrid2, bgrid1, bgrid2) - return qout -""" - - def _south_corners( qin: FloatField, qout: FloatField, @@ -366,8 +352,6 @@ def a2b_interpolation( qout: FloatField, qx: FloatField, qy: FloatField, - qxx: FloatField, - qyy: FloatField, dxa: FloatFieldIJ, dya: FloatFieldIJ, edge_w: FloatFieldIJ, @@ -375,7 +359,7 @@ def a2b_interpolation( edge_s: FloatFieldI, edge_n: FloatFieldI, ): - from __externals__ import i_end, i_start, j_end, j_start + from __externals__ import REPLACE, i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): # qout_edges_x @@ -408,61 +392,129 @@ def a2b_interpolation( ) + (1.0 - edge_n) * ( (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) ) - qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) - with horizontal(region[i_start + 1 : i_end + 1, j_start + 1]): - qxx = c1 * (qx[0, -1, 0] + qx) + c2 * ( + + # combined (qxx and qyy folded in) + with horizontal(region[i_start + 1, j_start + 1]): + qout = 0.5 * (( c1 * (qx[0, -1, 0] + qx) + c2 * ( + (edge_s * ( + (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) + / (dya[-1, -2] + dya[-1, -1]) + ) + (1.0 - edge_s) * ( + (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) / (dya[0, -2] + dya[0, -1]) + )) + + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) + )) + (c1 * (qy[-1, 0, 0] + qy) + c2 * ( + (edge_w * ( + (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) + / (dxa[-2, -1] + dxa[-1, -1]) + ) + (1.0 - edge_w) * ( + (qin[-2, 0, 0] * dxa[-1, 0] + qin[-1, 0, 0] * dxa[-2, 0]) / (dxa[-2, 0] + dxa[-1, 0]) + )) + + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) + ))) + with horizontal(region[i_end, j_start + 1]): + qout = 0.5 * ((c1 * (qx[0, -1, 0] + qx) + c2 * ( (edge_s * ( - (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) - / (dya[-1, -2] + dya[-1, -1]) - ) + (1.0 - edge_s) * ( - (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) / (dya[0, -2] + dya[0, -1]) - )) + (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) + / (dya[-1, -2] + dya[-1, -1]) + ) + (1.0 - edge_s) * ( + (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) / (dya[0, -2] + dya[0, -1]) + )) + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) - ) - with horizontal(region[i_start + 1 : i_end + 1, j_end]): - qxx = c1 * (qx[0, -1, 0] + qx) + c2 * ( + )) + + (c1 * (qy[-1, 0, 0] + qy) + c2 * ( + (edge_e * ( + (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) + / (dxa[0, -1] + dxa[1, -1]) + ) + (1.0 - edge_e) * ( + (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) / (dxa[0, 0] + dxa[1, 0]) + )) + + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) + ))) + with horizontal(region[i_start+2:i_end, j_start + 1]): + qout = 0.5 * ((c1 * (qx[0, -1, 0] + qx) + c2 * ( + (edge_s * ( + (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) + / (dya[-1, -2] + dya[-1, -1]) + ) + (1.0 - edge_s) * ( + (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) / (dya[0, -2] + dya[0, -1]) + )) + + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) + )) + + (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy))) + with horizontal(region[i_end, j_start + 2:j_end]): + qout = 0.5 * ((a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) + ( c1 * (qy[-1, 0, 0] + qy) + c2 * ( + (edge_e * ( + (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) + / (dxa[0, -1] + dxa[1, -1]) + ) + (1.0 - edge_e) * ( + (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) / (dxa[0, 0] + dxa[1, 0]) + )) + + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) + ))) + with horizontal(region[i_start + 1, j_end]): + qout = 0.5 * (( c1 * (qx[0, -1, 0] + qx) + c2 * ( ( edge_n * ( - (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) - / (dya[-1, 0] + dya[-1, 1]) - ) + (1.0 - edge_n) * ( - (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) / (dya[0, 0] + dya[0, 1]) - )) + (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) + / (dya[-1, 0] + dya[-1, 1]) + ) + (1.0 - edge_n) * ( + (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) / (dya[0, 0] + dya[0, 1]) + )) + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0])) - ) - qyy = a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) - with horizontal(region[i_start + 1,j_start + 1 : j_end + 1]): - qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * ( + )) + ( c1 * (qy[-1, 0, 0] + qy) + c2 * ( (edge_w * ( - (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) - / (dxa[-2, -1] + dxa[-1, -1]) + (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) + / (dxa[-2, -1] + dxa[-1, -1]) ) + (1.0 - edge_w) * ( (qin[-2, 0, 0] * dxa[-1, 0] + qin[-1, 0, 0] * dxa[-2, 0]) / (dxa[-2, 0] + dxa[-1, 0]) )) + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) - ) - with horizontal(region[i_end, j_start + 1 : j_end + 1]): - qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * ( - (edge_e * ( - (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) - / (dxa[0, -1] + dxa[1, -1]) - ) + (1.0 - edge_e) * ( - (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) / (dxa[0, 0] + dxa[1, 0]) + ))) + with horizontal(region[i_end, j_end]): + qout = 0.5 * (( c1 * (qx[0, -1, 0] + qx) + c2 * ( + ( edge_n * ( + (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) + / (dya[-1, 0] + dya[-1, 1]) + ) + (1.0 - edge_n) * ( + (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) / (dya[0, 0] + dya[0, 1]) + )) + + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0])) )) - + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) - ) -def a2b_interpolation_final( - qin: FloatField, - qout: FloatField, - qxx: FloatField, - qyy: FloatField, -): - from __externals__ import REPLACE, i_end, i_start, j_end, j_start + + ( c1 * (qy[-1, 0, 0] + qy) + c2 * ( + (edge_e * ( + (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) + / (dxa[0, -1] + dxa[1, -1]) + ) + (1.0 - edge_e) * ( + (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) / (dxa[0, 0] + dxa[1, 0]) + )) + + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) + ))) + with horizontal(region[i_start + 2:i_end, j_end]): + qout = 0.5 * (( c1 * (qx[0, -1, 0] + qx) + c2 * ( + ( edge_n * ( + (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) + / (dya[-1, 0] + dya[-1, 1]) + ) + (1.0 - edge_n) * ( + (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) / (dya[0, 0] + dya[0, 1]) + )) + + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0])) + )) + ( a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy))) + with horizontal(region[i_start + 1, j_start + 2:j_end]): + qout = 0.5 * (( a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) + ( c1 * (qy[-1, 0, 0] + qy) + c2 * ( + (edge_w * ( + (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) + / (dxa[-2, -1] + dxa[-1, -1]) + ) + (1.0 - edge_w) * ( + (qin[-2, 0, 0] * dxa[-1, 0] + qin[-1, 0, 0] * dxa[-2, 0]) / (dxa[-2, 0] + dxa[-1, 0]) + )) + + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) + ))) + with horizontal(region[i_start + 2 : i_end, j_start + 2:j_end]): + qout = 0.5 * ((a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) + (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy))) with computation(PARALLEL), interval(...): - with horizontal(region[i_start + 1 : i_end + 1, j_start + 1 : j_end + 1]): - qout = 0.5 * (qxx + qyy) if __INLINED(REPLACE): qin = qout - + class AGrid2BGridFourthOrder: @@ -490,8 +542,6 @@ def __init__( self._edge_w = self._j_storage_repeat_over_i(self.grid.edge_w, shape[0:2]) self._tmp_qx = utils.make_storage_from_shape(shape) self._tmp_qy = utils.make_storage_from_shape(shape) - self._tmp_qxx = utils.make_storage_from_shape(shape) - self._tmp_qyy = utils.make_storage_from_shape(shape) origin = (self.grid.is_, self.grid.js, kstart) domain = (self.grid.nic + 1, self.grid.njc + 1, nk) ax_offsets = axis_offsets(self.grid, origin, domain) @@ -528,12 +578,6 @@ def __init__( self._a2b_interpolation_stencil = FrozenStencil( a2b_interpolation, - externals=ax_offsets, - origin=origin, - domain=domain, - ) - self._a2b_interpolation_final_stencil = FrozenStencil( - a2b_interpolation_final, externals={"REPLACE": replace, **ax_offsets}, origin=origin, domain=domain, @@ -592,8 +636,6 @@ def __call__(self, qin: FloatField, qout: FloatField): qout, self._tmp_qx, self._tmp_qy, - self._tmp_qxx, - self._tmp_qyy, self.grid.dxa, self.grid.dya, self._edge_w, @@ -601,10 +643,4 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.edge_s, self.grid.edge_n, ) - self._a2b_interpolation_final_stencil( - qin, - qout, - self._tmp_qxx, - self._tmp_qyy, - ) - + From ebab32e77f4d94032adadaff7c20d930ba5fd131 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 19 May 2021 20:08:59 -0700 Subject: [PATCH 34/45] gt backend not validating, seeing if this is from the replace nonsense --- fv3core/stencils/a2b_ord4.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 6fb1a3b3a..146275b78 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -20,7 +20,7 @@ from fv3core.decorators import FrozenStencil from fv3core.utils import axis_offsets from fv3core.utils.typing import FloatField, FloatFieldI, FloatFieldIJ - +from fv3core.stencils.basic_operations import copy_defn # comact 4-pt cubic interpolation c1 = 2.0 / 3.0 @@ -511,9 +511,9 @@ def a2b_interpolation( ))) with horizontal(region[i_start + 2 : i_end, j_start + 2:j_end]): qout = 0.5 * ((a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) + (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy))) - with computation(PARALLEL), interval(...): - if __INLINED(REPLACE): - qin = qout + #with computation(PARALLEL), interval(...): + # if __INLINED(REPLACE): + # qin = qout @@ -534,7 +534,7 @@ def __init__( """ assert grid_type < 3 self.grid = spec.grid - + self._replace = replace shape = self.grid.domain_shape_full(add=(1, 1, 1)) if nk is None: nk = self.grid.npz - kstart @@ -582,7 +582,11 @@ def __init__( origin=origin, domain=domain, ) - + self._copy_stencil = FrozenStencil( + copy_defn, + origin=origin, + domain=domain, + ) # TODO # within regions, the edge_w and edge_w variables that are singleton in the @@ -644,3 +648,5 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.edge_n, ) + if self._replace: + self._copy_stencil(qout, qin) From 5d49b685f712fe384f32cc31972439e4bb632a00 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Wed, 19 May 2021 21:35:17 -0700 Subject: [PATCH 35/45] lint --- fv3core/stencils/a2b_ord4.py | 479 +++++++++++++++++++++++++---------- 1 file changed, 347 insertions(+), 132 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 146275b78..63f2ffd04 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -3,7 +3,6 @@ import gt4py import gt4py.gtscript as gtscript from gt4py.gtscript import ( - __INLINED, PARALLEL, asin, computation, @@ -18,9 +17,10 @@ import fv3core._config as spec import fv3core.utils.gt4py_utils as utils from fv3core.decorators import FrozenStencil +from fv3core.stencils.basic_operations import copy_defn from fv3core.utils import axis_offsets from fv3core.utils.typing import FloatField, FloatFieldI, FloatFieldIJ -from fv3core.stencils.basic_operations import copy_defn + # comact 4-pt cubic interpolation c1 = 2.0 / 3.0 @@ -230,7 +230,7 @@ def a2b_interpolation_qx( qx: FloatField, dxa: FloatFieldIJ, ): - from __externals__ import i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start with computation(PARALLEL), interval(...): # ppm_volume_mean_x @@ -295,12 +295,13 @@ def a2b_interpolation_qx( ) ) / (2.0 + 2.0 * dxa[-1, 0] / dxa) + def a2b_interpolation_qy( qin: FloatField, qy: FloatField, dya: FloatFieldIJ, ): - from __externals__ import i_end, i_start, j_end, j_start + from __externals__ import j_end, j_start with computation(PARALLEL), interval(...): @@ -347,6 +348,7 @@ def a2b_interpolation_qy( - (dya[0, -1] / dya * qy[0, 1, 0] + qy[0, -1, 0]) ) / (2.0 + 2.0 * dya[0, -1] / dya) + def a2b_interpolation( qin: FloatField, qout: FloatField, @@ -359,16 +361,17 @@ def a2b_interpolation( edge_s: FloatFieldI, edge_n: FloatFieldI, ): - from __externals__ import REPLACE, i_end, i_start, j_end, j_start + from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): # qout_edges_x - with horizontal(region[i_start,j_start + 1 : j_end + 1]): + with horizontal(region[i_start, j_start + 1 : j_end + 1]): qout = edge_w * ( (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / (dxa[-1, -1] + dxa[0, -1]) ) + (1.0 - edge_w) * ( - (qin[-1, 0, 0] * dxa[0, 0] + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa[0, 0]) + (qin[-1, 0, 0] * dxa[0, 0] + qin * dxa[-1, 0]) + / (dxa[-1, 0] + dxa[0, 0]) ) with horizontal(region[i_end + 1, j_start + 1 : j_end + 1]): qout = edge_e * ( @@ -392,129 +395,341 @@ def a2b_interpolation( ) + (1.0 - edge_n) * ( (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) ) - + # combined (qxx and qyy folded in) with horizontal(region[i_start + 1, j_start + 1]): - qout = 0.5 * (( c1 * (qx[0, -1, 0] + qx) + c2 * ( - (edge_s * ( - (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) - / (dya[-1, -2] + dya[-1, -1]) - ) + (1.0 - edge_s) * ( - (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) / (dya[0, -2] + dya[0, -1]) - )) - + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) - )) + (c1 * (qy[-1, 0, 0] + qy) + c2 * ( - (edge_w * ( - (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) - / (dxa[-2, -1] + dxa[-1, -1]) - ) + (1.0 - edge_w) * ( - (qin[-2, 0, 0] * dxa[-1, 0] + qin[-1, 0, 0] * dxa[-2, 0]) / (dxa[-2, 0] + dxa[-1, 0]) - )) - + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) - ))) + qout = 0.5 * ( + ( + c1 * (qx[0, -1, 0] + qx) + + c2 + * ( + ( + edge_s + * ( + ( + qin[-1, -2, 0] * dya[-1, -1] + + qin[-1, -1, 0] * dya[-1, -2] + ) + / (dya[-1, -2] + dya[-1, -1]) + ) + + (1.0 - edge_s) + * ( + ( + qin[0, -2, 0] * dya[0, -1] + + qin[0, -1, 0] * dya[0, -2] + ) + / (dya[0, -2] + dya[0, -1]) + ) + ) + + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) + ) + ) + + ( + c1 * (qy[-1, 0, 0] + qy) + + c2 + * ( + ( + edge_w + * ( + ( + qin[-2, -1, 0] * dxa[-1, -1] + + qin[-1, -1, 0] * dxa[-2, -1] + ) + / (dxa[-2, -1] + dxa[-1, -1]) + ) + + (1.0 - edge_w) + * ( + ( + qin[-2, 0, 0] * dxa[-1, 0] + + qin[-1, 0, 0] * dxa[-2, 0] + ) + / (dxa[-2, 0] + dxa[-1, 0]) + ) + ) + + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) + ) + ) + ) with horizontal(region[i_end, j_start + 1]): - qout = 0.5 * ((c1 * (qx[0, -1, 0] + qx) + c2 * ( - (edge_s * ( - (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) - / (dya[-1, -2] + dya[-1, -1]) - ) + (1.0 - edge_s) * ( - (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) / (dya[0, -2] + dya[0, -1]) - )) - + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) - )) + - (c1 * (qy[-1, 0, 0] + qy) + c2 * ( - (edge_e * ( - (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) - / (dxa[0, -1] + dxa[1, -1]) - ) + (1.0 - edge_e) * ( - (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) / (dxa[0, 0] + dxa[1, 0]) - )) - + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) - ))) - with horizontal(region[i_start+2:i_end, j_start + 1]): - qout = 0.5 * ((c1 * (qx[0, -1, 0] + qx) + c2 * ( - (edge_s * ( - (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) - / (dya[-1, -2] + dya[-1, -1]) - ) + (1.0 - edge_s) * ( - (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) / (dya[0, -2] + dya[0, -1]) - )) - + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) - )) + - (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy))) - with horizontal(region[i_end, j_start + 2:j_end]): - qout = 0.5 * ((a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) + ( c1 * (qy[-1, 0, 0] + qy) + c2 * ( - (edge_e * ( - (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) - / (dxa[0, -1] + dxa[1, -1]) - ) + (1.0 - edge_e) * ( - (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) / (dxa[0, 0] + dxa[1, 0]) - )) - + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) - ))) + qout = 0.5 * ( + ( + c1 * (qx[0, -1, 0] + qx) + + c2 + * ( + ( + edge_s + * ( + ( + qin[-1, -2, 0] * dya[-1, -1] + + qin[-1, -1, 0] * dya[-1, -2] + ) + / (dya[-1, -2] + dya[-1, -1]) + ) + + (1.0 - edge_s) + * ( + ( + qin[0, -2, 0] * dya[0, -1] + + qin[0, -1, 0] * dya[0, -2] + ) + / (dya[0, -2] + dya[0, -1]) + ) + ) + + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) + ) + ) + + ( + c1 * (qy[-1, 0, 0] + qy) + + c2 + * ( + ( + edge_e + * ( + ( + qin[0, -1, 0] * dxa[1, -1] + + qin[1, -1, 0] * dxa[0, -1] + ) + / (dxa[0, -1] + dxa[1, -1]) + ) + + (1.0 - edge_e) + * ( + (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) + / (dxa[0, 0] + dxa[1, 0]) + ) + ) + + ( + a2 * (qy[-3, 0, 0] + qy) + + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) + ) + ) + ) + ) + with horizontal(region[i_start + 2 : i_end, j_start + 1]): + qout = 0.5 * ( + ( + c1 * (qx[0, -1, 0] + qx) + + c2 + * ( + ( + edge_s + * ( + ( + qin[-1, -2, 0] * dya[-1, -1] + + qin[-1, -1, 0] * dya[-1, -2] + ) + / (dya[-1, -2] + dya[-1, -1]) + ) + + (1.0 - edge_s) + * ( + ( + qin[0, -2, 0] * dya[0, -1] + + qin[0, -1, 0] * dya[0, -2] + ) + / (dya[0, -2] + dya[0, -1]) + ) + ) + + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) + ) + ) + + (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy)) + ) + with horizontal(region[i_end, j_start + 2 : j_end]): + qout = 0.5 * ( + (a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) + + ( + c1 * (qy[-1, 0, 0] + qy) + + c2 + * ( + ( + edge_e + * ( + ( + qin[0, -1, 0] * dxa[1, -1] + + qin[1, -1, 0] * dxa[0, -1] + ) + / (dxa[0, -1] + dxa[1, -1]) + ) + + (1.0 - edge_e) + * ( + (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) + / (dxa[0, 0] + dxa[1, 0]) + ) + ) + + ( + a2 * (qy[-3, 0, 0] + qy) + + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) + ) + ) + ) + ) with horizontal(region[i_start + 1, j_end]): - qout = 0.5 * (( c1 * (qx[0, -1, 0] + qx) + c2 * ( - ( edge_n * ( - (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) - / (dya[-1, 0] + dya[-1, 1]) - ) + (1.0 - edge_n) * ( - (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) / (dya[0, 0] + dya[0, 1]) - )) - + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0])) - )) + ( c1 * (qy[-1, 0, 0] + qy) + c2 * ( - (edge_w * ( - (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) - / (dxa[-2, -1] + dxa[-1, -1]) - ) + (1.0 - edge_w) * ( - (qin[-2, 0, 0] * dxa[-1, 0] + qin[-1, 0, 0] * dxa[-2, 0]) / (dxa[-2, 0] + dxa[-1, 0]) - )) - + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) - ))) + qout = 0.5 * ( + ( + c1 * (qx[0, -1, 0] + qx) + + c2 + * ( + ( + edge_n + * ( + ( + qin[-1, 0, 0] * dya[-1, 1] + + qin[-1, 1, 0] * dya[-1, 0] + ) + / (dya[-1, 0] + dya[-1, 1]) + ) + + (1.0 - edge_n) + * ( + (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) + / (dya[0, 0] + dya[0, 1]) + ) + ) + + ( + a2 * (qx[0, -3, 0] + qx) + + a1 * (qx[0, -2, 0] + qx[0, -1, 0]) + ) + ) + ) + + ( + c1 * (qy[-1, 0, 0] + qy) + + c2 + * ( + ( + edge_w + * ( + ( + qin[-2, -1, 0] * dxa[-1, -1] + + qin[-1, -1, 0] * dxa[-2, -1] + ) + / (dxa[-2, -1] + dxa[-1, -1]) + ) + + (1.0 - edge_w) + * ( + ( + qin[-2, 0, 0] * dxa[-1, 0] + + qin[-1, 0, 0] * dxa[-2, 0] + ) + / (dxa[-2, 0] + dxa[-1, 0]) + ) + ) + + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) + ) + ) + ) with horizontal(region[i_end, j_end]): - qout = 0.5 * (( c1 * (qx[0, -1, 0] + qx) + c2 * ( - ( edge_n * ( - (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) - / (dya[-1, 0] + dya[-1, 1]) - ) + (1.0 - edge_n) * ( - (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) / (dya[0, 0] + dya[0, 1]) - )) - + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0])) - )) - + ( c1 * (qy[-1, 0, 0] + qy) + c2 * ( - (edge_e * ( - (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) - / (dxa[0, -1] + dxa[1, -1]) - ) + (1.0 - edge_e) * ( - (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) / (dxa[0, 0] + dxa[1, 0]) - )) - + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) - ))) - with horizontal(region[i_start + 2:i_end, j_end]): - qout = 0.5 * (( c1 * (qx[0, -1, 0] + qx) + c2 * ( - ( edge_n * ( - (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) - / (dya[-1, 0] + dya[-1, 1]) - ) + (1.0 - edge_n) * ( - (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) / (dya[0, 0] + dya[0, 1]) - )) - + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0])) - )) + ( a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy))) - with horizontal(region[i_start + 1, j_start + 2:j_end]): - qout = 0.5 * (( a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) + ( c1 * (qy[-1, 0, 0] + qy) + c2 * ( - (edge_w * ( - (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) - / (dxa[-2, -1] + dxa[-1, -1]) - ) + (1.0 - edge_w) * ( - (qin[-2, 0, 0] * dxa[-1, 0] + qin[-1, 0, 0] * dxa[-2, 0]) / (dxa[-2, 0] + dxa[-1, 0]) - )) - + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) - ))) - with horizontal(region[i_start + 2 : i_end, j_start + 2:j_end]): - qout = 0.5 * ((a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) + (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy))) - #with computation(PARALLEL), interval(...): - # if __INLINED(REPLACE): - # qin = qout - + qout = 0.5 * ( + ( + c1 * (qx[0, -1, 0] + qx) + + c2 + * ( + ( + edge_n + * ( + ( + qin[-1, 0, 0] * dya[-1, 1] + + qin[-1, 1, 0] * dya[-1, 0] + ) + / (dya[-1, 0] + dya[-1, 1]) + ) + + (1.0 - edge_n) + * ( + (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) + / (dya[0, 0] + dya[0, 1]) + ) + ) + + ( + a2 * (qx[0, -3, 0] + qx) + + a1 * (qx[0, -2, 0] + qx[0, -1, 0]) + ) + ) + ) + + ( + c1 * (qy[-1, 0, 0] + qy) + + c2 + * ( + ( + edge_e + * ( + ( + qin[0, -1, 0] * dxa[1, -1] + + qin[1, -1, 0] * dxa[0, -1] + ) + / (dxa[0, -1] + dxa[1, -1]) + ) + + (1.0 - edge_e) + * ( + (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) + / (dxa[0, 0] + dxa[1, 0]) + ) + ) + + ( + a2 * (qy[-3, 0, 0] + qy) + + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) + ) + ) + ) + ) + with horizontal(region[i_start + 2 : i_end, j_end]): + qout = 0.5 * ( + ( + c1 * (qx[0, -1, 0] + qx) + + c2 + * ( + ( + edge_n + * ( + ( + qin[-1, 0, 0] * dya[-1, 1] + + qin[-1, 1, 0] * dya[-1, 0] + ) + / (dya[-1, 0] + dya[-1, 1]) + ) + + (1.0 - edge_n) + * ( + (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) + / (dya[0, 0] + dya[0, 1]) + ) + ) + + ( + a2 * (qx[0, -3, 0] + qx) + + a1 * (qx[0, -2, 0] + qx[0, -1, 0]) + ) + ) + ) + + (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy)) + ) + with horizontal(region[i_start + 1, j_start + 2 : j_end]): + qout = 0.5 * ( + (a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) + + ( + c1 * (qy[-1, 0, 0] + qy) + + c2 + * ( + ( + edge_w + * ( + ( + qin[-2, -1, 0] * dxa[-1, -1] + + qin[-1, -1, 0] * dxa[-2, -1] + ) + / (dxa[-2, -1] + dxa[-1, -1]) + ) + + (1.0 - edge_w) + * ( + ( + qin[-2, 0, 0] * dxa[-1, 0] + + qin[-1, 0, 0] * dxa[-2, 0] + ) + / (dxa[-2, 0] + dxa[-1, 0]) + ) + ) + + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) + ) + ) + ) + with horizontal(region[i_start + 2 : i_end, j_start + 2 : j_end]): + qout = 0.5 * ( + (a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) + + (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy)) + ) class AGrid2BGridFourthOrder: @@ -575,10 +790,10 @@ def __init__( origin=origin_prep_y, domain=domain_prep_y, ) - + self._a2b_interpolation_stencil = FrozenStencil( a2b_interpolation, - externals={"REPLACE": replace, **ax_offsets}, + externals=ax_offsets, origin=origin, domain=domain, ) @@ -587,7 +802,7 @@ def __init__( origin=origin, domain=domain, ) - + # TODO # within regions, the edge_w and edge_w variables that are singleton in the # I dimension error, workaround is repeating the data, but the longterm @@ -622,19 +837,19 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.bgrid1, self.grid.bgrid2, ) - + self._a2b_interpolation_qx_stencil( qin, self._tmp_qx, self.grid.dxa, ) - + self._a2b_interpolation_qy_stencil( qin, self._tmp_qy, self.grid.dya, ) - + self._a2b_interpolation_stencil( qin, qout, @@ -647,6 +862,6 @@ def __call__(self, qin: FloatField, qout: FloatField): self.grid.edge_s, self.grid.edge_n, ) - + if self._replace: self._copy_stencil(qout, qin) From 0952fb7c4ec98ded3b55cc95eaca2aab75e0c89b Mon Sep 17 00:00:00 2001 From: Rhea George Date: Thu, 20 May 2021 15:46:49 -0700 Subject: [PATCH 36/45] will using gtscript functions upset the gt backend validation --- fv3core/stencils/a2b_ord4.py | 81 +++++++++++++++++++++++------------- 1 file changed, 53 insertions(+), 28 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 63f2ffd04..3fae2d472 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -224,7 +224,55 @@ def lagrange_y_func(qx): def lagrange_x_func(qy): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) +@gtscript.function +def _ppm_volume_mean_x(qin: FloatField): + return b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) + + +@gtscript.function +def _ppm_volume_mean_y(qin: FloatField): + return b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) + +@gtscript.function +def _qx_edge_west(qin: FloatField, dxa: FloatFieldIJ): + g_in = dxa[1, 0] / dxa + g_ou = dxa[-2, 0] / dxa[-1, 0] + return 0.5 * ( + ((2.0 + g_in) * qin - qin[1, 0, 0]) / (1.0 + g_in) + + ((2.0 + g_ou) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + g_ou) + ) + +@gtscript.function +def _qx_edge_west2(qin: FloatField, dxa: FloatFieldIJ): + # TODO: should be able to use 2d variable with offset: + # qxleft = _qx_edge_west(qin[-1, 0, 0], dxa[-1, 0]) + # TODO this seemsed to work for a bit, and then stopped + #qxright = _ppm_volume_mean_x(qin[1, 0, 0]) + #g_in = dxa / dxa[-1, 0] + #return ( + # 3.0 * (g_in * qin[-1, 0, 0] + qin) - (g_in * qxleft + qxright) + #) / (2.0 + 2.0 * g_in) + return ( + 3.0 * (dxa / dxa[-1, 0] * qin[-1, 0, 0] + qin) + - ( + dxa + / dxa[-1, 0] + * ( + 0.5 + * ( + ((2.0 + dxa[0, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin) + / (1.0 + dxa[0, 0] / dxa[-1, 0]) + + ( + (2.0 + dxa[-3, 0] / dxa[-2, 0]) * qin[-2, 0, 0] + - qin[-3, 0, 0] + ) + / (1.0 + dxa[-3, 0] / dxa[-2, 0]) + ) + ) + + (b2 * (qin[-1, 0, 0] + qin[2, 0, 0]) + b1 * (qin + qin[1, 0, 0])) + ) + ) / (2.0 + 2.0 * dxa / dxa[-1, 0]) def a2b_interpolation_qx( qin: FloatField, qx: FloatField, @@ -233,36 +281,14 @@ def a2b_interpolation_qx( from __externals__ import i_end, i_start with computation(PARALLEL), interval(...): - # ppm_volume_mean_x - qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) + qx = _ppm_volume_mean_x(qin) + with horizontal(region[i_start, :]): - qx = 0.5 * ( - ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) - + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) - / (1.0 + dxa[-2, 0] / dxa[-1, 0]) - ) + qx = _qx_edge_west(qin, dxa) with horizontal(region[i_start + 1, :]): - qx = ( - 3.0 * (dxa / dxa[-1, 0] * qin[-1, 0, 0] + qin) - - ( - dxa - / dxa[-1, 0] - * ( - 0.5 - * ( - ((2.0 + dxa[0, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin) - / (1.0 + dxa[0, 0] / dxa[-1, 0]) - + ( - (2.0 + dxa[-3, 0] / dxa[-2, 0]) * qin[-2, 0, 0] - - qin[-3, 0, 0] - ) - / (1.0 + dxa[-3, 0] / dxa[-2, 0]) - ) - ) - + (b2 * (qin[-1, 0, 0] + qin[2, 0, 0]) + b1 * (qin + qin[1, 0, 0])) - ) - ) / (2.0 + 2.0 * dxa / dxa[-1, 0]) + qx = _qx_edge_west2(qin, dxa) + with horizontal(region[i_end + 1, :]): qx = 0.5 * ( ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) @@ -731,7 +757,6 @@ def a2b_interpolation( + (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy)) ) - class AGrid2BGridFourthOrder: """ Fortran name is a2b_ord4, test module is A2B_Ord4 From 1423cc1361c6c136326d51389882f0d786160af5 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Thu, 20 May 2021 16:38:48 -0700 Subject: [PATCH 37/45] qx and qy in functions --- fv3core/stencils/a2b_ord4.py | 205 +++++++++++++++++++++-------------- 1 file changed, 124 insertions(+), 81 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 3fae2d472..3f5813d19 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -247,71 +247,67 @@ def _qx_edge_west(qin: FloatField, dxa: FloatFieldIJ): def _qx_edge_west2(qin: FloatField, dxa: FloatFieldIJ): # TODO: should be able to use 2d variable with offset: # qxleft = _qx_edge_west(qin[-1, 0, 0], dxa[-1, 0]) - # TODO this seemsed to work for a bit, and then stopped - #qxright = _ppm_volume_mean_x(qin[1, 0, 0]) - #g_in = dxa / dxa[-1, 0] - #return ( + # TODO this seemed to work for a bit, and then stopped + # qxright = _ppm_volume_mean_x(qin[1, 0, 0]) + g_in = dxa / dxa[-1, 0] + # return ( # 3.0 * (g_in * qin[-1, 0, 0] + qin) - (g_in * qxleft + qxright) - #) / (2.0 + 2.0 * g_in) + # ) / (2.0 + 2.0 * g_in) + g_ou = dxa[-3, 0] / dxa[-2, 0] return ( - 3.0 * (dxa / dxa[-1, 0] * qin[-1, 0, 0] + qin) + 3.0 * (g_in * qin[-1, 0, 0] + qin) - ( - dxa - / dxa[-1, 0] + g_in * ( 0.5 * ( - ((2.0 + dxa[0, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin) - / (1.0 + dxa[0, 0] / dxa[-1, 0]) + ((2.0 + g_in) * qin[-1, 0, 0] - qin) + / (1.0 + g_in) + ( - (2.0 + dxa[-3, 0] / dxa[-2, 0]) * qin[-2, 0, 0] + (2.0 + g_ou) * qin[-2, 0, 0] - qin[-3, 0, 0] ) - / (1.0 + dxa[-3, 0] / dxa[-2, 0]) + / (1.0 + g_ou) ) ) + (b2 * (qin[-1, 0, 0] + qin[2, 0, 0]) + b1 * (qin + qin[1, 0, 0])) ) - ) / (2.0 + 2.0 * dxa / dxa[-1, 0]) -def a2b_interpolation_qx( - qin: FloatField, - qx: FloatField, - dxa: FloatFieldIJ, -): - from __externals__ import i_end, i_start + ) / (2.0 + 2.0 * g_in) - with computation(PARALLEL), interval(...): - qx = _ppm_volume_mean_x(qin) - - with horizontal(region[i_start, :]): - qx = _qx_edge_west(qin, dxa) +@gtscript.function +def _qx_edge_east(qin: FloatField, dxa: FloatFieldIJ): + g_in = dxa[-2, 0] / dxa[-1, 0] + g_ou = dxa[1, 0] / dxa + return 0.5 * ( + ((2.0 + g_in) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + g_in) + + ((2.0 + g_ou) * qin - qin[1, 0, 0]) / (1.0 + g_ou) + ) - with horizontal(region[i_start + 1, :]): - qx = _qx_edge_west2(qin, dxa) - - with horizontal(region[i_end + 1, :]): - qx = 0.5 * ( - ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) - / (1.0 + dxa[-2, 0] / dxa[-1, 0]) - + ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) - / (1.0 + dxa[1, 0] / dxa) - ) - with horizontal(region[i_end, :]): - qx = ( - 3.0 * (qin[-1, 0, 0] + dxa[-1, 0] / dxa * qin) + +@gtscript.function +def _qx_edge_east2(qin: FloatField, dxa: FloatFieldIJ): + # TODO when possible + # qxright = _qx_edge_east(qin[1, 0, 0], dxa[1, 0]) + # qxleft = _ppm_volume_mean_x(qin[-1, 0, 0]) + g_in = dxa[-1, 0] / dxa + # return ( + # 3.0 * (qin[-1, 0, 0] + g_in * qin) - (g_in * qxright + qxleft) + # ) / (2.0 + 2.0 * g_in) + g_ou = dxa[2, 0] / dxa[1, 0] + return ( + 3.0 * (qin[-1, 0, 0] + g_in * qin) - ( - dxa[-1, 0] - / dxa + g_in * ( 0.5 * ( - ((2.0 + dxa[-1, 0] / dxa) * qin - qin[-1, 0, 0]) - / (1.0 + dxa[-1, 0] / dxa) + ((2.0 + g_in) * qin - qin[-1, 0, 0]) + / (1.0 + g_in) + ( - (2.0 + dxa[2, 0] / dxa[1, 0]) * qin[1, 0, 0] + (2.0 + g_ou) * qin[1, 0, 0] - qin[2, 0, 0] ) - / (1.0 + dxa[2, 0] / dxa[1, 0]) + / (1.0 + g_ou) ) ) + ( @@ -319,60 +315,107 @@ def a2b_interpolation_qx( + b1 * (qin[-2, 0, 0] + qin[-1, 0, 0]) ) ) - ) / (2.0 + 2.0 * dxa[-1, 0] / dxa) - - -def a2b_interpolation_qy( - qin: FloatField, - qy: FloatField, - dya: FloatFieldIJ, -): - from __externals__ import j_end, j_start + ) / (2.0 + 2.0 * g_in) - with computation(PARALLEL), interval(...): +@gtscript.function +def _qy_edge_south(qin: FloatField, dya: FloatFieldIJ): + g_in = dya[0, 1] / dya + g_ou = dya[0, -2] / dya[0, -1] + return 0.5 * ( + ((2.0 + g_in) * qin - qin[0, 1, 0]) / (1.0 + g_in) + + ((2.0 + g_ou) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + g_ou) + ) - # ppm_volume_mean_y - qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) - with horizontal(region[:, j_start]): - qy = 0.5 * ( - ((2.0 + dya[0, 1] / dya) * qin - qin[0, 1, 0]) / (1.0 + dya[0, 1] / dya) - + ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) - / (1.0 + dya[0, -2] / dya[0, -1]) - ) - with horizontal(region[:, j_start + 1]): - qy = ( - 3.0 * (dya / dya[0, -1] * qin[0, -1, 0] + qin) +@gtscript.function +def _qy_edge_south2(qin: FloatField, dya: FloatFieldIJ): + g_in = dya / dya[0, -1] + #return ( + # 3.0 * (g_in * qin[0, -1, 0] + qin) - (g_in * qy[0, -1, 0] + qy[0, 1, 0]) + #) / (2.0 + 2.0 * g_in) + g_ou = dya[0, -3] / dya[0, -2] + return ( + 3.0 * (g_in * qin[0, -1, 0] + qin) - ( - dya - / dya[0, -1] + g_in * ( 0.5 * ( - ((2.0 + dya / dya[0, -1]) * qin[0, -1, 0] - qin) - / (1.0 + dya / dya[0, -1]) + ((2.0 + g_in) * qin[0, -1, 0] - qin) + / (1.0 + g_in) + ( - (2.0 + dya[0, -3] / dya[0, -2]) * qin[0, -2, 0] + (2.0 + g_ou) * qin[0, -2, 0] - qin[0, -3, 0] ) - / (1.0 + dya[0, -3] / dya[0, -2]) + / (1.0 + g_ou) ) ) + (b2 * (qin[0, -1, 0] + qin[0, 2, 0]) + b1 * (qin + qin[0, 1, 0])) ) - ) / (2.0 + 2.0 * dya / dya[0, -1]) + ) / (2.0 + 2.0 * g_in) + +@gtscript.function +def _qy_edge_north(qin: FloatField, dya: FloatFieldIJ): + g_in = dya[0, -2] / dya[0, -1] + g_ou = dya[0, 1] / dya + return 0.5 * ( + ((2.0 + g_in) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + g_in) + + ((2.0 + g_ou) * qin - qin[0, 1, 0]) / (1.0 + g_ou) + ) + + +@gtscript.function +def _qy_edge_north2(qin: FloatField, dya: FloatFieldIJ, qy): + # TODO when possible: + # qylower = _ppm_volume_mean_y(qin[0, -1, 0]) + # qyupper = _qy_edge_north(qin[0, 1, 0], dya[0, 1]) + g_in = dya[0, -1] / dya + #return ( + # 3.0 * (qin[0, -1, 0] + g_in * qin) - (g_in * qyupper + qlower) + #) / (2.0 + 2.0 * g_in) + g_ou = dya[0, 2] / dya[0, 1] + return (3.0 * (qin[0, -1, 0] + g_in * qin) - (g_in * (0.5 * ( + ((2.0 + g_in) * qin - qin[0, -1, 0]) / (1.0 + g_in) + + ((2.0 + g_ou) * qin[0, 1, 0] - qin[0, 2, 0]) / (1.0 + g_ou) + )) + (b2 * (qin[0, -3, 0] + qin[0, 0, 0]) + b1 * (qin[0, -2, 0] + qin[0, -1, 0]))) + ) / (2.0 + 2.0 * g_in) + +def a2b_interpolation_qx( + qin: FloatField, + qx: FloatField, + dxa: FloatFieldIJ, +): + from __externals__ import i_end, i_start + + with computation(PARALLEL), interval(...): + qx = _ppm_volume_mean_x(qin) + with horizontal(region[i_start, :]): + qx = _qx_edge_west(qin, dxa) + with horizontal(region[i_start + 1, :]): + qx = _qx_edge_west2(qin, dxa) + with horizontal(region[i_end + 1, :]): + qx = _qx_edge_east(qin, dxa) + with horizontal(region[i_end, :]): + qx = _qx_edge_east2(qin, dxa) + +def a2b_interpolation_qy( + qin: FloatField, + qy: FloatField, + dya: FloatFieldIJ, +): + from __externals__ import j_end, j_start + + with computation(PARALLEL), interval(...): + # ppm_volume_mean_y + qy = _ppm_volume_mean_y(qin) + with horizontal(region[:, j_start]): + qy =_qy_edge_south(qin, dya) + with horizontal(region[:, j_start + 1]): + qy = _qy_edge_south2(qin, dya) with horizontal(region[:, j_end + 1]): - qy = 0.5 * ( - ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) - / (1.0 + dya[0, -2] / dya[0, -1]) - + ((2.0 + dya[0, 1] / dya) * qin - qin[0, 1, 0]) - / (1.0 + dya[0, 1] / dya) - ) + qy =_qy_edge_north(qin, dya) with horizontal(region[:, j_end]): - qy = ( - 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) - - (dya[0, -1] / dya * qy[0, 1, 0] + qy[0, -1, 0]) - ) / (2.0 + 2.0 * dya[0, -1] / dya) + qy = _qy_edge_north2(qin, dya, qy) def a2b_interpolation( From b3a177629b51efb7cce30364108bad1dedea0e1c Mon Sep 17 00:00:00 2001 From: Rhea George Date: Thu, 20 May 2021 17:25:35 -0700 Subject: [PATCH 38/45] yaz --- fv3core/stencils/a2b_ord4.py | 59 +++++++++++++++++++----------------- 1 file changed, 31 insertions(+), 28 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 3f5813d19..293c8ed8f 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -379,7 +379,31 @@ def _qy_edge_north2(qin: FloatField, dya: FloatFieldIJ, qy): + ((2.0 + g_ou) * qin[0, 1, 0] - qin[0, 2, 0]) / (1.0 + g_ou) )) + (b2 * (qin[0, -3, 0] + qin[0, 0, 0]) + b1 * (qin[0, -2, 0] + qin[0, -1, 0]))) ) / (2.0 + 2.0 * g_in) - + +@gtscript.function +def _dxa_weighted_left_average_q(qin, dxa): + return (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + +@gtscript.function +def _dya_weighted_lower_average_q(qin, dya): + return (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) + +@gtscript.function +def _qout_x_edge(qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ): + # TODO when possible: + #q 2lower = _dxa_weighted_left_average_q(qin[0, -1, 0], dxa[0, -1]) + q2lower = (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / (dxa[-1, -1] + dxa[0, -1]) + q2 = _dxa_weighted_left_average_q(qin, dxa) + return edge_w * q2lower + (1.0 - edge_w) * q2 + + +@gtscript.function +def _qout_y_edge(qin: FloatField, dya: FloatFieldIJ,edge_s: FloatFieldI): + q1left = (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) / (dya[-1, -1] + dya[-1, 0]) + q1 = _dya_weighted_lower_average_q(qin, dya) + return edge_s * q1left + (1.0 - edge_s) * q1 + + def a2b_interpolation_qx( qin: FloatField, qx: FloatField, @@ -433,37 +457,16 @@ def a2b_interpolation( from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): - # qout_edges_x + tmp = 0.0 with horizontal(region[i_start, j_start + 1 : j_end + 1]): - qout = edge_w * ( - (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) - / (dxa[-1, -1] + dxa[0, -1]) - ) + (1.0 - edge_w) * ( - (qin[-1, 0, 0] * dxa[0, 0] + qin * dxa[-1, 0]) - / (dxa[-1, 0] + dxa[0, 0]) - ) + qout = _qout_x_edge(qin, dxa, edge_w) with horizontal(region[i_end + 1, j_start + 1 : j_end + 1]): - qout = edge_e * ( - (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) - / (dxa[-1, -1] + dxa[0, -1]) - ) + (1.0 - edge_e) * ( - (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) - ) - # qout edges y + qout = _qout_x_edge(qin, dxa, edge_e) with horizontal(region[i_start + 1 : i_end + 1, j_start]): - qout = edge_s * ( - (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) - / (dya[-1, -1] + dya[-1, 0]) - ) + (1.0 - edge_s) * ( - (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - ) + qout = _qout_y_edge(qin, dya, edge_s) + with horizontal(region[i_start + 1 : i_end + 1, j_end + 1]): - qout = edge_n * ( - (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) - / (dya[-1, -1] + dya[-1, 0]) - ) + (1.0 - edge_n) * ( - (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - ) + qout = _qout_y_edge(qin, dya, edge_n) # combined (qxx and qyy folded in) with horizontal(region[i_start + 1, j_start + 1]): From 0c5ce3be59540419892058ee0ff18577fbde6902 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Thu, 20 May 2021 20:47:11 -0700 Subject: [PATCH 39/45] start to qxx and qyy --- fv3core/stencils/a2b_ord4.py | 94 ++++++++++++++++-------------------- 1 file changed, 42 insertions(+), 52 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 293c8ed8f..13fd197bd 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -403,6 +403,33 @@ def _qout_y_edge(qin: FloatField, dya: FloatFieldIJ,edge_s: FloatFieldI): q1 = _dya_weighted_lower_average_q(qin, dya) return edge_s * q1left + (1.0 - edge_s) * q1 +@gtscript.function +def lagrange_y(qx: FloatField): + return a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) + + +@gtscript.function +def lagrange_x(qy: FloatField): + return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) + +@gtscript.function +def _cubic_interpolation_south(qx, qoutlower, qxxupper): + #return c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + qxx[0, 1, 0]) + return c1 * (qx[0, -1, 0] + qx) + c2 * (qoutlower + qxxupper) + +@gtscript.function +def _cubic_interpolation_north(qx: FloatField, qout: FloatField, qxx: FloatField): + return c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, 1, 0] + qxx[0, -1, 0]) + + +@gtscript.function +def _cubic_interpolation_west(qy, qoutleft, qyyright): + return c1 * (qy[-1, 0, 0] + qy) + c2 * (qoutleft + qyyright) + + +@gtscript.function +def _cubic_interpolation_east(qy: FloatField, qout: FloatField, qyy: FloatField): + return c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[1, 0, 0] + qyy[-1, 0, 0]) def a2b_interpolation_qx( qin: FloatField, @@ -470,58 +497,21 @@ def a2b_interpolation( # combined (qxx and qyy folded in) with horizontal(region[i_start + 1, j_start + 1]): - qout = 0.5 * ( - ( - c1 * (qx[0, -1, 0] + qx) - + c2 - * ( - ( - edge_s - * ( - ( - qin[-1, -2, 0] * dya[-1, -1] - + qin[-1, -1, 0] * dya[-1, -2] - ) - / (dya[-1, -2] + dya[-1, -1]) - ) - + (1.0 - edge_s) - * ( - ( - qin[0, -2, 0] * dya[0, -1] - + qin[0, -1, 0] * dya[0, -2] - ) - / (dya[0, -2] + dya[0, -1]) - ) - ) - + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) - ) - ) - + ( - c1 * (qy[-1, 0, 0] + qy) - + c2 - * ( - ( - edge_w - * ( - ( - qin[-2, -1, 0] * dxa[-1, -1] - + qin[-1, -1, 0] * dxa[-2, -1] - ) - / (dxa[-2, -1] + dxa[-1, -1]) - ) - + (1.0 - edge_w) - * ( - ( - qin[-2, 0, 0] * dxa[-1, 0] - + qin[-1, 0, 0] * dxa[-2, 0] - ) - / (dxa[-2, 0] + dxa[-1, 0]) - ) - ) - + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) - ) - ) - ) + #qxxupper = lagrange_y(qx[0, 1, 0]) + qxxupper = a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0]) + #qoutlower = _qout_y_edge(qin[0, -1, 0], dya[0, -1, 0], edge_s[0, -1]) + qoutlower = (edge_s* ((qin[-1, -2, 0] * dya[-1, -1]+ qin[-1, -1, 0] * dya[-1, -2])/ (dya[-1, -2] + dya[-1, -1])) + + (1.0 - edge_s)* ((qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2])/ (dya[0, -2] + dya[0, -1]))) + qxx = _cubic_interpolation_south(qx, qoutlower, qxxupper) + + #qyyupper = lagrange_x(qy[0, 1, 0]) + qyyright = a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0]) + #qoutleft = _qout_x_edge(qin[-1, 0, 0], dxa[-1, 0], edge_w) + qoutleft = (edge_w * ((qin[-2, -1, 0] * dxa[-1, -1]+ qin[-1, -1, 0] * dxa[-2, -1])/ (dxa[-2, -1] + dxa[-1, -1])) + + (1.0 - edge_w)* ((qin[-2, 0, 0] * dxa[-1, 0]+ qin[-1, 0, 0] * dxa[-2, 0])/ (dxa[-2, 0] + dxa[-1, 0]))) + qyy = _cubic_interpolation_west(qy, qoutleft, qyyright) + qout = 0.5 *(qxx + qyy) + with horizontal(region[i_end, j_start + 1]): qout = 0.5 * ( ( From c7ab0e4947f1e619c6e0cb5f121348bb8cb7d51a Mon Sep 17 00:00:00 2001 From: Rhea George Date: Thu, 20 May 2021 21:13:13 -0700 Subject: [PATCH 40/45] more functions --- fv3core/stencils/a2b_ord4.py | 132 ++++++++++------------------------- 1 file changed, 38 insertions(+), 94 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 13fd197bd..39cc4a851 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -428,9 +428,26 @@ def _cubic_interpolation_west(qy, qoutleft, qyyright): @gtscript.function -def _cubic_interpolation_east(qy: FloatField, qout: FloatField, qyy: FloatField): - return c1 * (qy[-1, 0, 0] + qy) + c2 * (qout[1, 0, 0] + qyy[-1, 0, 0]) +def _cubic_interpolation_east(qy, qoutright, qyyleft): + return c1 * (qy[-1, 0, 0] + qy) + c2 * (qoutright + qyyleft) +@gtscript.function +def qxx_south(qin, qx, dya, edge_s): + # qxxupper = lagrange_y(qx[0, 1, 0]) + qxxupper = a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0]) + # qoutlower = _qout_y_edge(qin[0, -1, 0], dya[0, -1, 0], edge_s[0, -1]) + qoutlower = (edge_s* ((qin[-1, -2, 0] * dya[-1, -1]+ qin[-1, -1, 0] * dya[-1, -2])/ (dya[-1, -2] + dya[-1, -1])) + + (1.0 - edge_s)* ((qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2])/ (dya[0, -2] + dya[0, -1]))) + return _cubic_interpolation_south(qx, qoutlower, qxxupper) +@gtscript.function +def qyy_west(qin, qy, dxa, edge_w): + # qyyright = lagrange_x(qy[0, 1, 0]) + qyyright = a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0]) + # qoutleft = _qout_x_edge(qin[-1, 0, 0], dxa[-1, 0], edge_w) + qoutleft = (edge_w * ((qin[-2, -1, 0] * dxa[-1, -1]+ qin[-1, -1, 0] * dxa[-2, -1])/ (dxa[-2, -1] + dxa[-1, -1])) + + (1.0 - edge_w)* ((qin[-2, 0, 0] * dxa[-1, 0]+ qin[-1, 0, 0] * dxa[-2, 0])/ (dxa[-2, 0] + dxa[-1, 0]))) + return _cubic_interpolation_west(qy, qoutleft, qyyright) + def a2b_interpolation_qx( qin: FloatField, qx: FloatField, @@ -497,103 +514,30 @@ def a2b_interpolation( # combined (qxx and qyy folded in) with horizontal(region[i_start + 1, j_start + 1]): - #qxxupper = lagrange_y(qx[0, 1, 0]) + qxx = qxx_south(qin, qx, dya, edge_s) + qyy = qyy_west(qin, qy, dxa, edge_w) + qout = 0.5 *(qxx + qyy) + + with horizontal(region[i_end, j_start + 1]): + # qxxupper = lagrange_y(qx[0, 1, 0]) qxxupper = a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0]) - #qoutlower = _qout_y_edge(qin[0, -1, 0], dya[0, -1, 0], edge_s[0, -1]) + # qoutlower = _qout_y_edge(qin[0, -1, 0], dya[0, -1, 0], edge_s[0, -1]) qoutlower = (edge_s* ((qin[-1, -2, 0] * dya[-1, -1]+ qin[-1, -1, 0] * dya[-1, -2])/ (dya[-1, -2] + dya[-1, -1])) + (1.0 - edge_s)* ((qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2])/ (dya[0, -2] + dya[0, -1]))) qxx = _cubic_interpolation_south(qx, qoutlower, qxxupper) - - #qyyupper = lagrange_x(qy[0, 1, 0]) - qyyright = a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0]) - #qoutleft = _qout_x_edge(qin[-1, 0, 0], dxa[-1, 0], edge_w) - qoutleft = (edge_w * ((qin[-2, -1, 0] * dxa[-1, -1]+ qin[-1, -1, 0] * dxa[-2, -1])/ (dxa[-2, -1] + dxa[-1, -1])) - + (1.0 - edge_w)* ((qin[-2, 0, 0] * dxa[-1, 0]+ qin[-1, 0, 0] * dxa[-2, 0])/ (dxa[-2, 0] + dxa[-1, 0]))) - qyy = _cubic_interpolation_west(qy, qoutleft, qyyright) - qout = 0.5 *(qxx + qyy) - - with horizontal(region[i_end, j_start + 1]): - qout = 0.5 * ( - ( - c1 * (qx[0, -1, 0] + qx) - + c2 - * ( - ( - edge_s - * ( - ( - qin[-1, -2, 0] * dya[-1, -1] - + qin[-1, -1, 0] * dya[-1, -2] - ) - / (dya[-1, -2] + dya[-1, -1]) - ) - + (1.0 - edge_s) - * ( - ( - qin[0, -2, 0] * dya[0, -1] - + qin[0, -1, 0] * dya[0, -2] - ) - / (dya[0, -2] + dya[0, -1]) - ) - ) - + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) - ) - ) - + ( - c1 * (qy[-1, 0, 0] + qy) - + c2 - * ( - ( - edge_e - * ( - ( - qin[0, -1, 0] * dxa[1, -1] - + qin[1, -1, 0] * dxa[0, -1] - ) - / (dxa[0, -1] + dxa[1, -1]) - ) - + (1.0 - edge_e) - * ( - (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) - / (dxa[0, 0] + dxa[1, 0]) - ) - ) - + ( - a2 * (qy[-3, 0, 0] + qy) - + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) - ) - ) - ) - ) + + # qyyleft = lagrange_x(qy[-1, 0, 0]) + qyyleft = a2 * (qy[-3, 0, 0] + qy)+ a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) + # qoutright = _qout_x_edge(qin[1, 0, 0], dxa[1, 0], edge_e) + qoutright = (edge_e* ((qin[0, -1, 0] * dxa[1, -1]+ qin[1, -1, 0] * dxa[0, -1])/ (dxa[0, -1] + dxa[1, -1])) + + (1.0 - edge_e)* ((qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0])/ (dxa[0, 0] + dxa[1, 0]))) + qyy = _cubic_interpolation_east(qy, qoutright, qyyleft) + qout = 0.5 * (qxx + qyy) with horizontal(region[i_start + 2 : i_end, j_start + 1]): - qout = 0.5 * ( - ( - c1 * (qx[0, -1, 0] + qx) - + c2 - * ( - ( - edge_s - * ( - ( - qin[-1, -2, 0] * dya[-1, -1] - + qin[-1, -1, 0] * dya[-1, -2] - ) - / (dya[-1, -2] + dya[-1, -1]) - ) - + (1.0 - edge_s) - * ( - ( - qin[0, -2, 0] * dya[0, -1] - + qin[0, -1, 0] * dya[0, -2] - ) - / (dya[0, -2] + dya[0, -1]) - ) - ) - + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) - ) - ) - + (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy)) - ) + qxx = qxx_south(qin, qx, dya, edge_s) + qyy = lagrange_x(qy) + qout = 0.5 * (qxx + qyy) + with horizontal(region[i_end, j_start + 2 : j_end]): qout = 0.5 * ( (a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) From 535f99218df3b2c1aafad0f856c38f4028424174 Mon Sep 17 00:00:00 2001 From: Rhea George Date: Thu, 20 May 2021 21:38:01 -0700 Subject: [PATCH 41/45] all functions --- fv3core/stencils/a2b_ord4.py | 255 +++++------------------------------ 1 file changed, 34 insertions(+), 221 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 39cc4a851..3ea5ed1f3 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -418,8 +418,8 @@ def _cubic_interpolation_south(qx, qoutlower, qxxupper): return c1 * (qx[0, -1, 0] + qx) + c2 * (qoutlower + qxxupper) @gtscript.function -def _cubic_interpolation_north(qx: FloatField, qout: FloatField, qxx: FloatField): - return c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, 1, 0] + qxx[0, -1, 0]) +def _cubic_interpolation_north(qx, qoutupper, qxxlower): + return c1 * (qx[0, -1, 0] + qx) + c2 * (qoutupper + qxxlower) @gtscript.function @@ -439,6 +439,16 @@ def qxx_south(qin, qx, dya, edge_s): qoutlower = (edge_s* ((qin[-1, -2, 0] * dya[-1, -1]+ qin[-1, -1, 0] * dya[-1, -2])/ (dya[-1, -2] + dya[-1, -1])) + (1.0 - edge_s)* ((qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2])/ (dya[0, -2] + dya[0, -1]))) return _cubic_interpolation_south(qx, qoutlower, qxxupper) + +@gtscript.function +def qxx_north(qin, qx, dya, edge_n): + # qxxlower = lagrange_y(qx[0, -1, 0]) + qxxlower = a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0]) + # qoutupper = _qout_y_edge(qin[0, 1, 0], dya[0, 1, 0], edge_s[0, 1]) + qoutupper = (edge_n* ((qin[-1, 0, 0] * dya[-1, 1]+ qin[-1, 1, 0] * dya[-1, 0])/ (dya[-1, 0] + dya[-1, 1])) + + (1.0 - edge_n) * ((qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0])/ (dya[0, 0] + dya[0, 1]))) + return _cubic_interpolation_north(qx, qoutupper, qxxlower) + @gtscript.function def qyy_west(qin, qy, dxa, edge_w): # qyyright = lagrange_x(qy[0, 1, 0]) @@ -447,7 +457,16 @@ def qyy_west(qin, qy, dxa, edge_w): qoutleft = (edge_w * ((qin[-2, -1, 0] * dxa[-1, -1]+ qin[-1, -1, 0] * dxa[-2, -1])/ (dxa[-2, -1] + dxa[-1, -1])) + (1.0 - edge_w)* ((qin[-2, 0, 0] * dxa[-1, 0]+ qin[-1, 0, 0] * dxa[-2, 0])/ (dxa[-2, 0] + dxa[-1, 0]))) return _cubic_interpolation_west(qy, qoutleft, qyyright) - + +@gtscript.function +def qyy_east(qin, qy, dxa, edge_e): + # qyyleft = lagrange_x(qy[-1, 0, 0]) + qyyleft = a2 * (qy[-3, 0, 0] + qy)+ a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) + # qoutright = _qout_x_edge(qin[1, 0, 0], dxa[1, 0], edge_e) + qoutright = (edge_e* ((qin[0, -1, 0] * dxa[1, -1]+ qin[1, -1, 0] * dxa[0, -1])/ (dxa[0, -1] + dxa[1, -1])) + + (1.0 - edge_e)* ((qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0])/ (dxa[0, 0] + dxa[1, 0]))) + return _cubic_interpolation_east(qy, qoutright, qyyleft) + def a2b_interpolation_qx( qin: FloatField, qx: FloatField, @@ -501,6 +520,7 @@ def a2b_interpolation( from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): + # TODO: remove when regions horizontalIf bug fixed tmp = 0.0 with horizontal(region[i_start, j_start + 1 : j_end + 1]): qout = _qout_x_edge(qin, dxa, edge_w) @@ -508,234 +528,27 @@ def a2b_interpolation( qout = _qout_x_edge(qin, dxa, edge_e) with horizontal(region[i_start + 1 : i_end + 1, j_start]): qout = _qout_y_edge(qin, dya, edge_s) - with horizontal(region[i_start + 1 : i_end + 1, j_end + 1]): qout = _qout_y_edge(qin, dya, edge_n) - # combined (qxx and qyy folded in) with horizontal(region[i_start + 1, j_start + 1]): - qxx = qxx_south(qin, qx, dya, edge_s) - qyy = qyy_west(qin, qy, dxa, edge_w) - qout = 0.5 *(qxx + qyy) - + qout = 0.5 *(qxx_south(qin, qx, dya, edge_s) + qyy_west(qin, qy, dxa, edge_w)) with horizontal(region[i_end, j_start + 1]): - # qxxupper = lagrange_y(qx[0, 1, 0]) - qxxupper = a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0]) - # qoutlower = _qout_y_edge(qin[0, -1, 0], dya[0, -1, 0], edge_s[0, -1]) - qoutlower = (edge_s* ((qin[-1, -2, 0] * dya[-1, -1]+ qin[-1, -1, 0] * dya[-1, -2])/ (dya[-1, -2] + dya[-1, -1])) - + (1.0 - edge_s)* ((qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2])/ (dya[0, -2] + dya[0, -1]))) - qxx = _cubic_interpolation_south(qx, qoutlower, qxxupper) - - # qyyleft = lagrange_x(qy[-1, 0, 0]) - qyyleft = a2 * (qy[-3, 0, 0] + qy)+ a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) - # qoutright = _qout_x_edge(qin[1, 0, 0], dxa[1, 0], edge_e) - qoutright = (edge_e* ((qin[0, -1, 0] * dxa[1, -1]+ qin[1, -1, 0] * dxa[0, -1])/ (dxa[0, -1] + dxa[1, -1])) - + (1.0 - edge_e)* ((qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0])/ (dxa[0, 0] + dxa[1, 0]))) - qyy = _cubic_interpolation_east(qy, qoutright, qyyleft) - qout = 0.5 * (qxx + qyy) - with horizontal(region[i_start + 2 : i_end, j_start + 1]): - qxx = qxx_south(qin, qx, dya, edge_s) - qyy = lagrange_x(qy) - qout = 0.5 * (qxx + qyy) - - with horizontal(region[i_end, j_start + 2 : j_end]): - qout = 0.5 * ( - (a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) - + ( - c1 * (qy[-1, 0, 0] + qy) - + c2 - * ( - ( - edge_e - * ( - ( - qin[0, -1, 0] * dxa[1, -1] - + qin[1, -1, 0] * dxa[0, -1] - ) - / (dxa[0, -1] + dxa[1, -1]) - ) - + (1.0 - edge_e) - * ( - (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) - / (dxa[0, 0] + dxa[1, 0]) - ) - ) - + ( - a2 * (qy[-3, 0, 0] + qy) - + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) - ) - ) - ) - ) + qout = 0.5 * (qxx_south(qin, qx, dya, edge_s) + qyy_east(qin, qy, dxa, edge_e)) with horizontal(region[i_start + 1, j_end]): - qout = 0.5 * ( - ( - c1 * (qx[0, -1, 0] + qx) - + c2 - * ( - ( - edge_n - * ( - ( - qin[-1, 0, 0] * dya[-1, 1] - + qin[-1, 1, 0] * dya[-1, 0] - ) - / (dya[-1, 0] + dya[-1, 1]) - ) - + (1.0 - edge_n) - * ( - (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) - / (dya[0, 0] + dya[0, 1]) - ) - ) - + ( - a2 * (qx[0, -3, 0] + qx) - + a1 * (qx[0, -2, 0] + qx[0, -1, 0]) - ) - ) - ) - + ( - c1 * (qy[-1, 0, 0] + qy) - + c2 - * ( - ( - edge_w - * ( - ( - qin[-2, -1, 0] * dxa[-1, -1] - + qin[-1, -1, 0] * dxa[-2, -1] - ) - / (dxa[-2, -1] + dxa[-1, -1]) - ) - + (1.0 - edge_w) - * ( - ( - qin[-2, 0, 0] * dxa[-1, 0] - + qin[-1, 0, 0] * dxa[-2, 0] - ) - / (dxa[-2, 0] + dxa[-1, 0]) - ) - ) - + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) - ) - ) - ) + qout = 0.5 * ( qxx_north(qin, qx, dya, edge_n) + qyy_west(qin, qy, dxa, edge_w)) with horizontal(region[i_end, j_end]): - qout = 0.5 * ( - ( - c1 * (qx[0, -1, 0] + qx) - + c2 - * ( - ( - edge_n - * ( - ( - qin[-1, 0, 0] * dya[-1, 1] - + qin[-1, 1, 0] * dya[-1, 0] - ) - / (dya[-1, 0] + dya[-1, 1]) - ) - + (1.0 - edge_n) - * ( - (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) - / (dya[0, 0] + dya[0, 1]) - ) - ) - + ( - a2 * (qx[0, -3, 0] + qx) - + a1 * (qx[0, -2, 0] + qx[0, -1, 0]) - ) - ) - ) - + ( - c1 * (qy[-1, 0, 0] + qy) - + c2 - * ( - ( - edge_e - * ( - ( - qin[0, -1, 0] * dxa[1, -1] - + qin[1, -1, 0] * dxa[0, -1] - ) - / (dxa[0, -1] + dxa[1, -1]) - ) - + (1.0 - edge_e) - * ( - (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) - / (dxa[0, 0] + dxa[1, 0]) - ) - ) - + ( - a2 * (qy[-3, 0, 0] + qy) - + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) - ) - ) - ) - ) + qout = 0.5 * (qxx_north(qin, qx, dya, edge_n)+ qyy_east(qin, qy, dxa, edge_e)) + with horizontal(region[i_start + 2 : i_end, j_start + 1]): + qout = 0.5 * ( qxx_south(qin, qx, dya, edge_s) + lagrange_x(qy)) + with horizontal(region[i_end, j_start + 2 : j_end]): + qout = 0.5 * (lagrange_y(qx)+ qyy_east(qin, qy, dxa, edge_e)) with horizontal(region[i_start + 2 : i_end, j_end]): - qout = 0.5 * ( - ( - c1 * (qx[0, -1, 0] + qx) - + c2 - * ( - ( - edge_n - * ( - ( - qin[-1, 0, 0] * dya[-1, 1] - + qin[-1, 1, 0] * dya[-1, 0] - ) - / (dya[-1, 0] + dya[-1, 1]) - ) - + (1.0 - edge_n) - * ( - (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) - / (dya[0, 0] + dya[0, 1]) - ) - ) - + ( - a2 * (qx[0, -3, 0] + qx) - + a1 * (qx[0, -2, 0] + qx[0, -1, 0]) - ) - ) - ) - + (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy)) - ) + qout = 0.5 * (qxx_north(qin, qx, dya, edge_n) + lagrange_x(qy)) with horizontal(region[i_start + 1, j_start + 2 : j_end]): - qout = 0.5 * ( - (a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) - + ( - c1 * (qy[-1, 0, 0] + qy) - + c2 - * ( - ( - edge_w - * ( - ( - qin[-2, -1, 0] * dxa[-1, -1] - + qin[-1, -1, 0] * dxa[-2, -1] - ) - / (dxa[-2, -1] + dxa[-1, -1]) - ) - + (1.0 - edge_w) - * ( - ( - qin[-2, 0, 0] * dxa[-1, 0] - + qin[-1, 0, 0] * dxa[-2, 0] - ) - / (dxa[-2, 0] + dxa[-1, 0]) - ) - ) - + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) - ) - ) - ) + qout = 0.5 * (lagrange_y(qx)+ qyy_west(qin, qy, dxa, edge_w)) with horizontal(region[i_start + 2 : i_end, j_start + 2 : j_end]): - qout = 0.5 * ( - (a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)) - + (a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy)) - ) + qout = 0.5 * (lagrange_y(qx) + lagrange_x(qy)) class AGrid2BGridFourthOrder: """ From 074f77b9b6a429def1f7288fa97bbdd1b7548e7b Mon Sep 17 00:00:00 2001 From: Rhea George Date: Thu, 20 May 2021 21:42:16 -0700 Subject: [PATCH 42/45] a little more --- fv3core/stencils/a2b_ord4.py | 25 ++++++++----------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 3ea5ed1f3..83f9f4bee 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -413,24 +413,15 @@ def lagrange_x(qy: FloatField): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) @gtscript.function -def _cubic_interpolation_south(qx, qoutlower, qxxupper): - #return c1 * (qx[0, -1, 0] + qx) + c2 * (qout[0, -1, 0] + qxx[0, 1, 0]) - return c1 * (qx[0, -1, 0] + qx) + c2 * (qoutlower + qxxupper) - -@gtscript.function -def _cubic_interpolation_north(qx, qoutupper, qxxlower): - return c1 * (qx[0, -1, 0] + qx) + c2 * (qoutupper + qxxlower) +def _cubic_interpolation_y(qx, qout_offset_y, qxx_offset_y): + return c1 * (qx[0, -1, 0] + qx) + c2 * (qout_offset_y + qxx_offset_y) @gtscript.function -def _cubic_interpolation_west(qy, qoutleft, qyyright): - return c1 * (qy[-1, 0, 0] + qy) + c2 * (qoutleft + qyyright) +def _cubic_interpolation_x(qy, qout_offset_x, qyy_offset_x): + return c1 * (qy[-1, 0, 0] + qy) + c2 * (qout_offset_x + qyy_offset_x) -@gtscript.function -def _cubic_interpolation_east(qy, qoutright, qyyleft): - return c1 * (qy[-1, 0, 0] + qy) + c2 * (qoutright + qyyleft) - @gtscript.function def qxx_south(qin, qx, dya, edge_s): # qxxupper = lagrange_y(qx[0, 1, 0]) @@ -438,7 +429,7 @@ def qxx_south(qin, qx, dya, edge_s): # qoutlower = _qout_y_edge(qin[0, -1, 0], dya[0, -1, 0], edge_s[0, -1]) qoutlower = (edge_s* ((qin[-1, -2, 0] * dya[-1, -1]+ qin[-1, -1, 0] * dya[-1, -2])/ (dya[-1, -2] + dya[-1, -1])) + (1.0 - edge_s)* ((qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2])/ (dya[0, -2] + dya[0, -1]))) - return _cubic_interpolation_south(qx, qoutlower, qxxupper) + return _cubic_interpolation_y(qx, qoutlower, qxxupper) @gtscript.function def qxx_north(qin, qx, dya, edge_n): @@ -447,7 +438,7 @@ def qxx_north(qin, qx, dya, edge_n): # qoutupper = _qout_y_edge(qin[0, 1, 0], dya[0, 1, 0], edge_s[0, 1]) qoutupper = (edge_n* ((qin[-1, 0, 0] * dya[-1, 1]+ qin[-1, 1, 0] * dya[-1, 0])/ (dya[-1, 0] + dya[-1, 1])) + (1.0 - edge_n) * ((qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0])/ (dya[0, 0] + dya[0, 1]))) - return _cubic_interpolation_north(qx, qoutupper, qxxlower) + return _cubic_interpolation_y(qx, qoutupper, qxxlower) @gtscript.function def qyy_west(qin, qy, dxa, edge_w): @@ -456,7 +447,7 @@ def qyy_west(qin, qy, dxa, edge_w): # qoutleft = _qout_x_edge(qin[-1, 0, 0], dxa[-1, 0], edge_w) qoutleft = (edge_w * ((qin[-2, -1, 0] * dxa[-1, -1]+ qin[-1, -1, 0] * dxa[-2, -1])/ (dxa[-2, -1] + dxa[-1, -1])) + (1.0 - edge_w)* ((qin[-2, 0, 0] * dxa[-1, 0]+ qin[-1, 0, 0] * dxa[-2, 0])/ (dxa[-2, 0] + dxa[-1, 0]))) - return _cubic_interpolation_west(qy, qoutleft, qyyright) + return _cubic_interpolation_x(qy, qoutleft, qyyright) @gtscript.function def qyy_east(qin, qy, dxa, edge_e): @@ -465,7 +456,7 @@ def qyy_east(qin, qy, dxa, edge_e): # qoutright = _qout_x_edge(qin[1, 0, 0], dxa[1, 0], edge_e) qoutright = (edge_e* ((qin[0, -1, 0] * dxa[1, -1]+ qin[1, -1, 0] * dxa[0, -1])/ (dxa[0, -1] + dxa[1, -1])) + (1.0 - edge_e)* ((qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0])/ (dxa[0, 0] + dxa[1, 0]))) - return _cubic_interpolation_east(qy, qoutright, qyyleft) + return _cubic_interpolation_x(qy, qoutright, qyyleft) def a2b_interpolation_qx( qin: FloatField, From 217aeb9a8f3e0bfc2df28f2d1fefa3e909a354bf Mon Sep 17 00:00:00 2001 From: Rhea George Date: Thu, 20 May 2021 21:43:20 -0700 Subject: [PATCH 43/45] lint --- fv3core/stencils/a2b_ord4.py | 224 +++++++++++++++++++++-------------- 1 file changed, 132 insertions(+), 92 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index 83f9f4bee..e4bc9d9db 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -224,6 +224,7 @@ def lagrange_y_func(qx): def lagrange_x_func(qy): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) + @gtscript.function def _ppm_volume_mean_x(qin: FloatField): return b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin) @@ -233,6 +234,7 @@ def _ppm_volume_mean_x(qin: FloatField): def _ppm_volume_mean_y(qin: FloatField): return b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin) + @gtscript.function def _qx_edge_west(qin: FloatField, dxa: FloatFieldIJ): g_in = dxa[1, 0] / dxa @@ -261,19 +263,15 @@ def _qx_edge_west2(qin: FloatField, dxa: FloatFieldIJ): * ( 0.5 * ( - ((2.0 + g_in) * qin[-1, 0, 0] - qin) - / (1.0 + g_in) - + ( - (2.0 + g_ou) * qin[-2, 0, 0] - - qin[-3, 0, 0] - ) - / (1.0 + g_ou) + ((2.0 + g_in) * qin[-1, 0, 0] - qin) / (1.0 + g_in) + + ((2.0 + g_ou) * qin[-2, 0, 0] - qin[-3, 0, 0]) / (1.0 + g_ou) ) ) + (b2 * (qin[-1, 0, 0] + qin[2, 0, 0]) + b1 * (qin + qin[1, 0, 0])) ) ) / (2.0 + 2.0 * g_in) + @gtscript.function def _qx_edge_east(qin: FloatField, dxa: FloatFieldIJ): g_in = dxa[-2, 0] / dxa[-1, 0] @@ -286,7 +284,7 @@ def _qx_edge_east(qin: FloatField, dxa: FloatFieldIJ): @gtscript.function def _qx_edge_east2(qin: FloatField, dxa: FloatFieldIJ): - # TODO when possible + # TODO when possible # qxright = _qx_edge_east(qin[1, 0, 0], dxa[1, 0]) # qxleft = _ppm_volume_mean_x(qin[-1, 0, 0]) g_in = dxa[-1, 0] / dxa @@ -294,28 +292,21 @@ def _qx_edge_east2(qin: FloatField, dxa: FloatFieldIJ): # 3.0 * (qin[-1, 0, 0] + g_in * qin) - (g_in * qxright + qxleft) # ) / (2.0 + 2.0 * g_in) g_ou = dxa[2, 0] / dxa[1, 0] - return ( + return ( 3.0 * (qin[-1, 0, 0] + g_in * qin) - - ( - g_in - * ( - 0.5 - * ( - ((2.0 + g_in) * qin - qin[-1, 0, 0]) - / (1.0 + g_in) - + ( - (2.0 + g_ou) * qin[1, 0, 0] - - qin[2, 0, 0] - ) - / (1.0 + g_ou) - ) - ) - + ( - b2 * (qin[-3, 0, 0] + qin) - + b1 * (qin[-2, 0, 0] + qin[-1, 0, 0]) - ) + - ( + g_in + * ( + 0.5 + * ( + ((2.0 + g_in) * qin - qin[-1, 0, 0]) / (1.0 + g_in) + + ((2.0 + g_ou) * qin[1, 0, 0] - qin[2, 0, 0]) / (1.0 + g_ou) ) - ) / (2.0 + 2.0 * g_in) + ) + + (b2 * (qin[-3, 0, 0] + qin) + b1 * (qin[-2, 0, 0] + qin[-1, 0, 0])) + ) + ) / (2.0 + 2.0 * g_in) + @gtscript.function def _qy_edge_south(qin: FloatField, dya: FloatFieldIJ): @@ -330,29 +321,25 @@ def _qy_edge_south(qin: FloatField, dya: FloatFieldIJ): @gtscript.function def _qy_edge_south2(qin: FloatField, dya: FloatFieldIJ): g_in = dya / dya[0, -1] - #return ( + # return ( # 3.0 * (g_in * qin[0, -1, 0] + qin) - (g_in * qy[0, -1, 0] + qy[0, 1, 0]) - #) / (2.0 + 2.0 * g_in) - g_ou = dya[0, -3] / dya[0, -2] + # ) / (2.0 + 2.0 * g_in) + g_ou = dya[0, -3] / dya[0, -2] return ( - 3.0 * (g_in * qin[0, -1, 0] + qin) - - ( - g_in - * ( - 0.5 - * ( - ((2.0 + g_in) * qin[0, -1, 0] - qin) - / (1.0 + g_in) - + ( - (2.0 + g_ou) * qin[0, -2, 0] - - qin[0, -3, 0] - ) - / (1.0 + g_ou) - ) - ) - + (b2 * (qin[0, -1, 0] + qin[0, 2, 0]) + b1 * (qin + qin[0, 1, 0])) + 3.0 * (g_in * qin[0, -1, 0] + qin) + - ( + g_in + * ( + 0.5 + * ( + ((2.0 + g_in) * qin[0, -1, 0] - qin) / (1.0 + g_in) + + ((2.0 + g_ou) * qin[0, -2, 0] - qin[0, -3, 0]) / (1.0 + g_ou) ) - ) / (2.0 + 2.0 * g_in) + ) + + (b2 * (qin[0, -1, 0] + qin[0, 2, 0]) + b1 * (qin + qin[0, 1, 0])) + ) + ) / (2.0 + 2.0 * g_in) + @gtscript.function def _qy_edge_north(qin: FloatField, dya: FloatFieldIJ): @@ -370,39 +357,59 @@ def _qy_edge_north2(qin: FloatField, dya: FloatFieldIJ, qy): # qylower = _ppm_volume_mean_y(qin[0, -1, 0]) # qyupper = _qy_edge_north(qin[0, 1, 0], dya[0, 1]) g_in = dya[0, -1] / dya - #return ( + # return ( # 3.0 * (qin[0, -1, 0] + g_in * qin) - (g_in * qyupper + qlower) - #) / (2.0 + 2.0 * g_in) + # ) / (2.0 + 2.0 * g_in) g_ou = dya[0, 2] / dya[0, 1] - return (3.0 * (qin[0, -1, 0] + g_in * qin) - (g_in * (0.5 * ( - ((2.0 + g_in) * qin - qin[0, -1, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin[0, 1, 0] - qin[0, 2, 0]) / (1.0 + g_ou) - )) + (b2 * (qin[0, -3, 0] + qin[0, 0, 0]) + b1 * (qin[0, -2, 0] + qin[0, -1, 0]))) - ) / (2.0 + 2.0 * g_in) + return ( + 3.0 * (qin[0, -1, 0] + g_in * qin) + - ( + g_in + * ( + 0.5 + * ( + ((2.0 + g_in) * qin - qin[0, -1, 0]) / (1.0 + g_in) + + ((2.0 + g_ou) * qin[0, 1, 0] - qin[0, 2, 0]) / (1.0 + g_ou) + ) + ) + + ( + b2 * (qin[0, -3, 0] + qin[0, 0, 0]) + + b1 * (qin[0, -2, 0] + qin[0, -1, 0]) + ) + ) + ) / (2.0 + 2.0 * g_in) + @gtscript.function def _dxa_weighted_left_average_q(qin, dxa): return (qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa) + @gtscript.function def _dya_weighted_lower_average_q(qin, dya): return (qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya) - + + @gtscript.function -def _qout_x_edge(qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ): +def _qout_x_edge(qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ): # TODO when possible: - #q 2lower = _dxa_weighted_left_average_q(qin[0, -1, 0], dxa[0, -1]) - q2lower = (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / (dxa[-1, -1] + dxa[0, -1]) - q2 = _dxa_weighted_left_average_q(qin, dxa) + # q 2lower = _dxa_weighted_left_average_q(qin[0, -1, 0], dxa[0, -1]) + q2lower = (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / ( + dxa[-1, -1] + dxa[0, -1] + ) + q2 = _dxa_weighted_left_average_q(qin, dxa) return edge_w * q2lower + (1.0 - edge_w) * q2 @gtscript.function -def _qout_y_edge(qin: FloatField, dya: FloatFieldIJ,edge_s: FloatFieldI): - q1left = (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) / (dya[-1, -1] + dya[-1, 0]) - q1 = _dya_weighted_lower_average_q(qin, dya) +def _qout_y_edge(qin: FloatField, dya: FloatFieldIJ, edge_s: FloatFieldI): + q1left = (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) / ( + dya[-1, -1] + dya[-1, 0] + ) + q1 = _dya_weighted_lower_average_q(qin, dya) return edge_s * q1left + (1.0 - edge_s) * q1 + @gtscript.function def lagrange_y(qx: FloatField): return a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx) @@ -412,6 +419,7 @@ def lagrange_y(qx: FloatField): def lagrange_x(qy: FloatField): return a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy) + @gtscript.function def _cubic_interpolation_y(qx, qout_offset_y, qxx_offset_y): return c1 * (qx[0, -1, 0] + qx) + c2 * (qout_offset_y + qxx_offset_y) @@ -427,36 +435,58 @@ def qxx_south(qin, qx, dya, edge_s): # qxxupper = lagrange_y(qx[0, 1, 0]) qxxupper = a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0]) # qoutlower = _qout_y_edge(qin[0, -1, 0], dya[0, -1, 0], edge_s[0, -1]) - qoutlower = (edge_s* ((qin[-1, -2, 0] * dya[-1, -1]+ qin[-1, -1, 0] * dya[-1, -2])/ (dya[-1, -2] + dya[-1, -1])) - + (1.0 - edge_s)* ((qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2])/ (dya[0, -2] + dya[0, -1]))) - return _cubic_interpolation_y(qx, qoutlower, qxxupper) + qoutlower = edge_s * ( + (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) + / (dya[-1, -2] + dya[-1, -1]) + ) + (1.0 - edge_s) * ( + (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) + / (dya[0, -2] + dya[0, -1]) + ) + return _cubic_interpolation_y(qx, qoutlower, qxxupper) + @gtscript.function def qxx_north(qin, qx, dya, edge_n): # qxxlower = lagrange_y(qx[0, -1, 0]) qxxlower = a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0]) # qoutupper = _qout_y_edge(qin[0, 1, 0], dya[0, 1, 0], edge_s[0, 1]) - qoutupper = (edge_n* ((qin[-1, 0, 0] * dya[-1, 1]+ qin[-1, 1, 0] * dya[-1, 0])/ (dya[-1, 0] + dya[-1, 1])) - + (1.0 - edge_n) * ((qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0])/ (dya[0, 0] + dya[0, 1]))) - return _cubic_interpolation_y(qx, qoutupper, qxxlower) + qoutupper = edge_n * ( + (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) + / (dya[-1, 0] + dya[-1, 1]) + ) + (1.0 - edge_n) * ( + (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) / (dya[0, 0] + dya[0, 1]) + ) + return _cubic_interpolation_y(qx, qoutupper, qxxlower) + @gtscript.function def qyy_west(qin, qy, dxa, edge_w): # qyyright = lagrange_x(qy[0, 1, 0]) qyyright = a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0]) # qoutleft = _qout_x_edge(qin[-1, 0, 0], dxa[-1, 0], edge_w) - qoutleft = (edge_w * ((qin[-2, -1, 0] * dxa[-1, -1]+ qin[-1, -1, 0] * dxa[-2, -1])/ (dxa[-2, -1] + dxa[-1, -1])) - + (1.0 - edge_w)* ((qin[-2, 0, 0] * dxa[-1, 0]+ qin[-1, 0, 0] * dxa[-2, 0])/ (dxa[-2, 0] + dxa[-1, 0]))) + qoutleft = edge_w * ( + (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) + / (dxa[-2, -1] + dxa[-1, -1]) + ) + (1.0 - edge_w) * ( + (qin[-2, 0, 0] * dxa[-1, 0] + qin[-1, 0, 0] * dxa[-2, 0]) + / (dxa[-2, 0] + dxa[-1, 0]) + ) return _cubic_interpolation_x(qy, qoutleft, qyyright) + @gtscript.function def qyy_east(qin, qy, dxa, edge_e): # qyyleft = lagrange_x(qy[-1, 0, 0]) - qyyleft = a2 * (qy[-3, 0, 0] + qy)+ a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) + qyyleft = a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) # qoutright = _qout_x_edge(qin[1, 0, 0], dxa[1, 0], edge_e) - qoutright = (edge_e* ((qin[0, -1, 0] * dxa[1, -1]+ qin[1, -1, 0] * dxa[0, -1])/ (dxa[0, -1] + dxa[1, -1])) - + (1.0 - edge_e)* ((qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0])/ (dxa[0, 0] + dxa[1, 0]))) - return _cubic_interpolation_x(qy, qoutright, qyyleft) + qoutright = edge_e * ( + (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) + / (dxa[0, -1] + dxa[1, -1]) + ) + (1.0 - edge_e) * ( + (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) / (dxa[0, 0] + dxa[1, 0]) + ) + return _cubic_interpolation_x(qy, qoutright, qyyleft) + def a2b_interpolation_qx( qin: FloatField, @@ -470,11 +500,12 @@ def a2b_interpolation_qx( with horizontal(region[i_start, :]): qx = _qx_edge_west(qin, dxa) with horizontal(region[i_start + 1, :]): - qx = _qx_edge_west2(qin, dxa) + qx = _qx_edge_west2(qin, dxa) with horizontal(region[i_end + 1, :]): - qx = _qx_edge_east(qin, dxa) + qx = _qx_edge_east(qin, dxa) with horizontal(region[i_end, :]): - qx = _qx_edge_east2(qin, dxa) + qx = _qx_edge_east2(qin, dxa) + def a2b_interpolation_qy( qin: FloatField, @@ -485,13 +516,13 @@ def a2b_interpolation_qy( with computation(PARALLEL), interval(...): # ppm_volume_mean_y - qy = _ppm_volume_mean_y(qin) + qy = _ppm_volume_mean_y(qin) with horizontal(region[:, j_start]): - qy =_qy_edge_south(qin, dya) + qy = _qy_edge_south(qin, dya) with horizontal(region[:, j_start + 1]): qy = _qy_edge_south2(qin, dya) with horizontal(region[:, j_end + 1]): - qy =_qy_edge_north(qin, dya) + qy = _qy_edge_north(qin, dya) with horizontal(region[:, j_end]): qy = _qy_edge_north2(qin, dya, qy) @@ -514,32 +545,41 @@ def a2b_interpolation( # TODO: remove when regions horizontalIf bug fixed tmp = 0.0 with horizontal(region[i_start, j_start + 1 : j_end + 1]): - qout = _qout_x_edge(qin, dxa, edge_w) + qout = _qout_x_edge(qin, dxa, edge_w) with horizontal(region[i_end + 1, j_start + 1 : j_end + 1]): - qout = _qout_x_edge(qin, dxa, edge_e) + qout = _qout_x_edge(qin, dxa, edge_e) with horizontal(region[i_start + 1 : i_end + 1, j_start]): qout = _qout_y_edge(qin, dya, edge_s) with horizontal(region[i_start + 1 : i_end + 1, j_end + 1]): qout = _qout_y_edge(qin, dya, edge_n) # combined (qxx and qyy folded in) with horizontal(region[i_start + 1, j_start + 1]): - qout = 0.5 *(qxx_south(qin, qx, dya, edge_s) + qyy_west(qin, qy, dxa, edge_w)) + qout = 0.5 * ( + qxx_south(qin, qx, dya, edge_s) + qyy_west(qin, qy, dxa, edge_w) + ) with horizontal(region[i_end, j_start + 1]): - qout = 0.5 * (qxx_south(qin, qx, dya, edge_s) + qyy_east(qin, qy, dxa, edge_e)) + qout = 0.5 * ( + qxx_south(qin, qx, dya, edge_s) + qyy_east(qin, qy, dxa, edge_e) + ) with horizontal(region[i_start + 1, j_end]): - qout = 0.5 * ( qxx_north(qin, qx, dya, edge_n) + qyy_west(qin, qy, dxa, edge_w)) + qout = 0.5 * ( + qxx_north(qin, qx, dya, edge_n) + qyy_west(qin, qy, dxa, edge_w) + ) with horizontal(region[i_end, j_end]): - qout = 0.5 * (qxx_north(qin, qx, dya, edge_n)+ qyy_east(qin, qy, dxa, edge_e)) + qout = 0.5 * ( + qxx_north(qin, qx, dya, edge_n) + qyy_east(qin, qy, dxa, edge_e) + ) with horizontal(region[i_start + 2 : i_end, j_start + 1]): - qout = 0.5 * ( qxx_south(qin, qx, dya, edge_s) + lagrange_x(qy)) + qout = 0.5 * (qxx_south(qin, qx, dya, edge_s) + lagrange_x(qy)) with horizontal(region[i_end, j_start + 2 : j_end]): - qout = 0.5 * (lagrange_y(qx)+ qyy_east(qin, qy, dxa, edge_e)) + qout = 0.5 * (lagrange_y(qx) + qyy_east(qin, qy, dxa, edge_e)) with horizontal(region[i_start + 2 : i_end, j_end]): - qout = 0.5 * (qxx_north(qin, qx, dya, edge_n) + lagrange_x(qy)) + qout = 0.5 * (qxx_north(qin, qx, dya, edge_n) + lagrange_x(qy)) with horizontal(region[i_start + 1, j_start + 2 : j_end]): - qout = 0.5 * (lagrange_y(qx)+ qyy_west(qin, qy, dxa, edge_w)) + qout = 0.5 * (lagrange_y(qx) + qyy_west(qin, qy, dxa, edge_w)) with horizontal(region[i_start + 2 : i_end, j_start + 2 : j_end]): - qout = 0.5 * (lagrange_y(qx) + lagrange_x(qy)) + qout = 0.5 * (lagrange_y(qx) + lagrange_x(qy)) + class AGrid2BGridFourthOrder: """ From c72d7b98599316b7296cf2046d743467bde8cbdf Mon Sep 17 00:00:00 2001 From: Rhea George Date: Fri, 21 May 2021 21:47:15 +0200 Subject: [PATCH 44/45] expanded some of the functions so less temporaries and do not run out of memory, but is slower, 5 seconds for some reason --- fv3core/stencils/a2b_ord4.py | 125 ++++++++++++++--------------------- 1 file changed, 49 insertions(+), 76 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index e4bc9d9db..a7f92556e 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -237,11 +237,9 @@ def _ppm_volume_mean_y(qin: FloatField): @gtscript.function def _qx_edge_west(qin: FloatField, dxa: FloatFieldIJ): - g_in = dxa[1, 0] / dxa - g_ou = dxa[-2, 0] / dxa[-1, 0] return 0.5 * ( - ((2.0 + g_in) * qin - qin[1, 0, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + g_ou) + ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) + + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + dxa[-2, 0] / dxa[-1, 0]) ) @@ -251,34 +249,30 @@ def _qx_edge_west2(qin: FloatField, dxa: FloatFieldIJ): # qxleft = _qx_edge_west(qin[-1, 0, 0], dxa[-1, 0]) # TODO this seemed to work for a bit, and then stopped # qxright = _ppm_volume_mean_x(qin[1, 0, 0]) - g_in = dxa / dxa[-1, 0] # return ( # 3.0 * (g_in * qin[-1, 0, 0] + qin) - (g_in * qxleft + qxright) # ) / (2.0 + 2.0 * g_in) - g_ou = dxa[-3, 0] / dxa[-2, 0] return ( - 3.0 * (g_in * qin[-1, 0, 0] + qin) + 3.0 * (dxa / dxa[-1, 0] * qin[-1, 0, 0] + qin) - ( - g_in + dxa / dxa[-1, 0] * ( 0.5 * ( - ((2.0 + g_in) * qin[-1, 0, 0] - qin) / (1.0 + g_in) - + ((2.0 + g_ou) * qin[-2, 0, 0] - qin[-3, 0, 0]) / (1.0 + g_ou) + ((2.0 + dxa / dxa[-1, 0]) * qin[-1, 0, 0] - qin) / (1.0 + dxa / dxa[-1, 0]) + + ((2.0 + dxa[-3, 0] / dxa[-2, 0]) * qin[-2, 0, 0] - qin[-3, 0, 0]) / (1.0 + dxa[-3, 0] / dxa[-2, 0]) ) ) + (b2 * (qin[-1, 0, 0] + qin[2, 0, 0]) + b1 * (qin + qin[1, 0, 0])) ) - ) / (2.0 + 2.0 * g_in) + ) / (2.0 + 2.0 * dxa / dxa[-1, 0]) @gtscript.function def _qx_edge_east(qin: FloatField, dxa: FloatFieldIJ): - g_in = dxa[-2, 0] / dxa[-1, 0] - g_ou = dxa[1, 0] / dxa return 0.5 * ( - ((2.0 + g_in) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin - qin[1, 0, 0]) / (1.0 + g_ou) + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + dxa[-2, 0] / dxa[-1, 0]) + + ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) ) @@ -287,67 +281,59 @@ def _qx_edge_east2(qin: FloatField, dxa: FloatFieldIJ): # TODO when possible # qxright = _qx_edge_east(qin[1, 0, 0], dxa[1, 0]) # qxleft = _ppm_volume_mean_x(qin[-1, 0, 0]) - g_in = dxa[-1, 0] / dxa # return ( # 3.0 * (qin[-1, 0, 0] + g_in * qin) - (g_in * qxright + qxleft) # ) / (2.0 + 2.0 * g_in) - g_ou = dxa[2, 0] / dxa[1, 0] return ( - 3.0 * (qin[-1, 0, 0] + g_in * qin) + 3.0 * (qin[-1, 0, 0] + dxa[-1, 0] / dxa * qin) - ( - g_in + dxa[-1, 0] / dxa * ( 0.5 * ( - ((2.0 + g_in) * qin - qin[-1, 0, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin[1, 0, 0] - qin[2, 0, 0]) / (1.0 + g_ou) + ((2.0 + dxa[-1, 0] / dxa) * qin - qin[-1, 0, 0]) / (1.0 + dxa[-1, 0] / dxa) + + ((2.0 + dxa[2, 0] / dxa[1, 0]) * qin[1, 0, 0] - qin[2, 0, 0]) / (1.0 + dxa[2, 0] / dxa[1, 0]) ) ) + (b2 * (qin[-3, 0, 0] + qin) + b1 * (qin[-2, 0, 0] + qin[-1, 0, 0])) ) - ) / (2.0 + 2.0 * g_in) + ) / (2.0 + 2.0 * dxa[-1, 0] / dxa) @gtscript.function def _qy_edge_south(qin: FloatField, dya: FloatFieldIJ): - g_in = dya[0, 1] / dya - g_ou = dya[0, -2] / dya[0, -1] return 0.5 * ( - ((2.0 + g_in) * qin - qin[0, 1, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + g_ou) + ((2.0 + dya[0, 1] / dya) * qin - qin[0, 1, 0]) / (1.0 + dya[0, 1] / dya) + + ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + dya[0, -2] / dya[0, -1]) ) @gtscript.function def _qy_edge_south2(qin: FloatField, dya: FloatFieldIJ): - g_in = dya / dya[0, -1] # return ( # 3.0 * (g_in * qin[0, -1, 0] + qin) - (g_in * qy[0, -1, 0] + qy[0, 1, 0]) # ) / (2.0 + 2.0 * g_in) - g_ou = dya[0, -3] / dya[0, -2] return ( - 3.0 * (g_in * qin[0, -1, 0] + qin) + 3.0 * (dya / dya[0, -1] * qin[0, -1, 0] + qin) - ( - g_in + dya / dya[0, -1] * ( 0.5 * ( - ((2.0 + g_in) * qin[0, -1, 0] - qin) / (1.0 + g_in) - + ((2.0 + g_ou) * qin[0, -2, 0] - qin[0, -3, 0]) / (1.0 + g_ou) + ((2.0 + dya / dya[0, -1]) * qin[0, -1, 0] - qin) / (1.0 + dya / dya[0, -1]) + + ((2.0 + dya[0, -3] / dya[0, -2]) * qin[0, -2, 0] - qin[0, -3, 0]) / (1.0 + dya[0, -3] / dya[0, -2]) ) ) + (b2 * (qin[0, -1, 0] + qin[0, 2, 0]) + b1 * (qin + qin[0, 1, 0])) ) - ) / (2.0 + 2.0 * g_in) + ) / (2.0 + 2.0 * dya / dya[0, -1]) @gtscript.function def _qy_edge_north(qin: FloatField, dya: FloatFieldIJ): - g_in = dya[0, -2] / dya[0, -1] - g_ou = dya[0, 1] / dya return 0.5 * ( - ((2.0 + g_in) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin - qin[0, 1, 0]) / (1.0 + g_ou) + ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + dya[0, -2] / dya[0, -1]) + + ((2.0 + dya[0, 1] / dya) * qin - qin[0, 1, 0]) / (1.0 + dya[0, 1] / dya) ) @@ -356,20 +342,19 @@ def _qy_edge_north2(qin: FloatField, dya: FloatFieldIJ, qy): # TODO when possible: # qylower = _ppm_volume_mean_y(qin[0, -1, 0]) # qyupper = _qy_edge_north(qin[0, 1, 0], dya[0, 1]) - g_in = dya[0, -1] / dya + # g_in = dya[0, -1] / dya # return ( # 3.0 * (qin[0, -1, 0] + g_in * qin) - (g_in * qyupper + qlower) # ) / (2.0 + 2.0 * g_in) - g_ou = dya[0, 2] / dya[0, 1] return ( - 3.0 * (qin[0, -1, 0] + g_in * qin) + 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) - ( - g_in + dya[0, -1] / dya * ( 0.5 * ( - ((2.0 + g_in) * qin - qin[0, -1, 0]) / (1.0 + g_in) - + ((2.0 + g_ou) * qin[0, 1, 0] - qin[0, 2, 0]) / (1.0 + g_ou) + ((2.0 + dya[0, -1] / dya) * qin - qin[0, -1, 0]) / (1.0 + dya[0, -1] / dya) + + ((2.0 + dya[0, 2] / dya[0, 1]) * qin[0, 1, 0] - qin[0, 2, 0]) / (1.0 + dya[0, 2] / dya[0, 1]) ) ) + ( @@ -377,7 +362,7 @@ def _qy_edge_north2(qin: FloatField, dya: FloatFieldIJ, qy): + b1 * (qin[0, -2, 0] + qin[0, -1, 0]) ) ) - ) / (2.0 + 2.0 * g_in) + ) / (2.0 + 2.0 * dya[0, -1] / dya) @gtscript.function @@ -393,21 +378,17 @@ def _dya_weighted_lower_average_q(qin, dya): @gtscript.function def _qout_x_edge(qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ): # TODO when possible: - # q 2lower = _dxa_weighted_left_average_q(qin[0, -1, 0], dxa[0, -1]) - q2lower = (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / ( - dxa[-1, -1] + dxa[0, -1] - ) - q2 = _dxa_weighted_left_average_q(qin, dxa) - return edge_w * q2lower + (1.0 - edge_w) * q2 - + # q2lower = _dxa_weighted_left_average_q(qin[0, -1, 0], dxa[0, -1]) + # q2 = _dxa_weighted_left_average_q(qin, dxa) + # return edge_w * q2lower + (1.0 - edge_w) * q2 + return edge_w * ((qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / (dxa[-1, -1] + dxa[0, -1])) + (1.0 - edge_w) * ((qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa)) @gtscript.function def _qout_y_edge(qin: FloatField, dya: FloatFieldIJ, edge_s: FloatFieldI): - q1left = (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) / ( - dya[-1, -1] + dya[-1, 0] - ) - q1 = _dya_weighted_lower_average_q(qin, dya) - return edge_s * q1left + (1.0 - edge_s) * q1 + # q1left = _dya_weighted_lower_average_q(qin[-1, 0, 0], dya[-1, 0]) + # q1 = _dya_weighted_lower_average_q(qin, dya) + # return edge_s * q1left + (1.0 - edge_s) * q1 + return edge_s * ((qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) / (dya[-1, -1] + dya[-1, 0])) + (1.0 - edge_s) * ((qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya)) @gtscript.function @@ -433,60 +414,52 @@ def _cubic_interpolation_x(qy, qout_offset_x, qyy_offset_x): @gtscript.function def qxx_south(qin, qx, dya, edge_s): # qxxupper = lagrange_y(qx[0, 1, 0]) - qxxupper = a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0]) # qoutlower = _qout_y_edge(qin[0, -1, 0], dya[0, -1, 0], edge_s[0, -1]) - qoutlower = edge_s * ( + #return _cubic_interpolation_y(qx, qoutlower, qxxupper) + return c1 * (qx[0, -1, 0] + qx) + c2 * ((edge_s * ( (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) / (dya[-1, -2] + dya[-1, -1]) ) + (1.0 - edge_s) * ( (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) / (dya[0, -2] + dya[0, -1]) - ) - return _cubic_interpolation_y(qx, qoutlower, qxxupper) - + )) + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0]))) @gtscript.function def qxx_north(qin, qx, dya, edge_n): # qxxlower = lagrange_y(qx[0, -1, 0]) - qxxlower = a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0]) # qoutupper = _qout_y_edge(qin[0, 1, 0], dya[0, 1, 0], edge_s[0, 1]) - qoutupper = edge_n * ( + #return _cubic_interpolation_y(qx, qoutupper, qxxlower) + return c1 * (qx[0, -1, 0] + qx) + c2 * ((edge_n * ( (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) / (dya[-1, 0] + dya[-1, 1]) ) + (1.0 - edge_n) * ( (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) / (dya[0, 0] + dya[0, 1]) - ) - return _cubic_interpolation_y(qx, qoutupper, qxxlower) - + )) + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0]))) @gtscript.function def qyy_west(qin, qy, dxa, edge_w): # qyyright = lagrange_x(qy[0, 1, 0]) - qyyright = a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0]) # qoutleft = _qout_x_edge(qin[-1, 0, 0], dxa[-1, 0], edge_w) - qoutleft = edge_w * ( + #return _cubic_interpolation_x(qy, qoutleft, qyyright) + return c1 * (qy[-1, 0, 0] + qy) + c2 * ((edge_w * ( (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) / (dxa[-2, -1] + dxa[-1, -1]) ) + (1.0 - edge_w) * ( (qin[-2, 0, 0] * dxa[-1, 0] + qin[-1, 0, 0] * dxa[-2, 0]) / (dxa[-2, 0] + dxa[-1, 0]) - ) - return _cubic_interpolation_x(qy, qoutleft, qyyright) - + )) + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0]))) @gtscript.function def qyy_east(qin, qy, dxa, edge_e): # qyyleft = lagrange_x(qy[-1, 0, 0]) - qyyleft = a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]) # qoutright = _qout_x_edge(qin[1, 0, 0], dxa[1, 0], edge_e) - qoutright = edge_e * ( + #return _cubic_interpolation_x(qy, qoutright, qyyleft) + return c1 * (qy[-1, 0, 0] + qy) + c2 * ((edge_e * ( (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) / (dxa[0, -1] + dxa[1, -1]) ) + (1.0 - edge_e) * ( (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) / (dxa[0, 0] + dxa[1, 0]) - ) - return _cubic_interpolation_x(qy, qoutright, qyyleft) - + )) + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]))) def a2b_interpolation_qx( qin: FloatField, From 801cb9b9ff4cce7782ad825eb6d38bb8316e561b Mon Sep 17 00:00:00 2001 From: Rhea George Date: Fri, 21 May 2021 12:47:55 -0700 Subject: [PATCH 45/45] lint --- fv3core/stencils/a2b_ord4.py | 159 ++++++++++++++++++++++++----------- 1 file changed, 110 insertions(+), 49 deletions(-) diff --git a/fv3core/stencils/a2b_ord4.py b/fv3core/stencils/a2b_ord4.py index a7f92556e..ca2e5927b 100644 --- a/fv3core/stencils/a2b_ord4.py +++ b/fv3core/stencils/a2b_ord4.py @@ -239,7 +239,8 @@ def _ppm_volume_mean_y(qin: FloatField): def _qx_edge_west(qin: FloatField, dxa: FloatFieldIJ): return 0.5 * ( ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) - + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + dxa[-2, 0] / dxa[-1, 0]) + + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) + / (1.0 + dxa[-2, 0] / dxa[-1, 0]) ) @@ -255,12 +256,15 @@ def _qx_edge_west2(qin: FloatField, dxa: FloatFieldIJ): return ( 3.0 * (dxa / dxa[-1, 0] * qin[-1, 0, 0] + qin) - ( - dxa / dxa[-1, 0] + dxa + / dxa[-1, 0] * ( 0.5 * ( - ((2.0 + dxa / dxa[-1, 0]) * qin[-1, 0, 0] - qin) / (1.0 + dxa / dxa[-1, 0]) - + ((2.0 + dxa[-3, 0] / dxa[-2, 0]) * qin[-2, 0, 0] - qin[-3, 0, 0]) / (1.0 + dxa[-3, 0] / dxa[-2, 0]) + ((2.0 + dxa / dxa[-1, 0]) * qin[-1, 0, 0] - qin) + / (1.0 + dxa / dxa[-1, 0]) + + ((2.0 + dxa[-3, 0] / dxa[-2, 0]) * qin[-2, 0, 0] - qin[-3, 0, 0]) + / (1.0 + dxa[-3, 0] / dxa[-2, 0]) ) ) + (b2 * (qin[-1, 0, 0] + qin[2, 0, 0]) + b1 * (qin + qin[1, 0, 0])) @@ -271,7 +275,8 @@ def _qx_edge_west2(qin: FloatField, dxa: FloatFieldIJ): @gtscript.function def _qx_edge_east(qin: FloatField, dxa: FloatFieldIJ): return 0.5 * ( - ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) / (1.0 + dxa[-2, 0] / dxa[-1, 0]) + ((2.0 + dxa[-2, 0] / dxa[-1, 0]) * qin[-1, 0, 0] - qin[-2, 0, 0]) + / (1.0 + dxa[-2, 0] / dxa[-1, 0]) + ((2.0 + dxa[1, 0] / dxa) * qin - qin[1, 0, 0]) / (1.0 + dxa[1, 0] / dxa) ) @@ -287,12 +292,15 @@ def _qx_edge_east2(qin: FloatField, dxa: FloatFieldIJ): return ( 3.0 * (qin[-1, 0, 0] + dxa[-1, 0] / dxa * qin) - ( - dxa[-1, 0] / dxa + dxa[-1, 0] + / dxa * ( 0.5 * ( - ((2.0 + dxa[-1, 0] / dxa) * qin - qin[-1, 0, 0]) / (1.0 + dxa[-1, 0] / dxa) - + ((2.0 + dxa[2, 0] / dxa[1, 0]) * qin[1, 0, 0] - qin[2, 0, 0]) / (1.0 + dxa[2, 0] / dxa[1, 0]) + ((2.0 + dxa[-1, 0] / dxa) * qin - qin[-1, 0, 0]) + / (1.0 + dxa[-1, 0] / dxa) + + ((2.0 + dxa[2, 0] / dxa[1, 0]) * qin[1, 0, 0] - qin[2, 0, 0]) + / (1.0 + dxa[2, 0] / dxa[1, 0]) ) ) + (b2 * (qin[-3, 0, 0] + qin) + b1 * (qin[-2, 0, 0] + qin[-1, 0, 0])) @@ -304,7 +312,8 @@ def _qx_edge_east2(qin: FloatField, dxa: FloatFieldIJ): def _qy_edge_south(qin: FloatField, dya: FloatFieldIJ): return 0.5 * ( ((2.0 + dya[0, 1] / dya) * qin - qin[0, 1, 0]) / (1.0 + dya[0, 1] / dya) - + ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + dya[0, -2] / dya[0, -1]) + + ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) + / (1.0 + dya[0, -2] / dya[0, -1]) ) @@ -316,12 +325,15 @@ def _qy_edge_south2(qin: FloatField, dya: FloatFieldIJ): return ( 3.0 * (dya / dya[0, -1] * qin[0, -1, 0] + qin) - ( - dya / dya[0, -1] + dya + / dya[0, -1] * ( 0.5 * ( - ((2.0 + dya / dya[0, -1]) * qin[0, -1, 0] - qin) / (1.0 + dya / dya[0, -1]) - + ((2.0 + dya[0, -3] / dya[0, -2]) * qin[0, -2, 0] - qin[0, -3, 0]) / (1.0 + dya[0, -3] / dya[0, -2]) + ((2.0 + dya / dya[0, -1]) * qin[0, -1, 0] - qin) + / (1.0 + dya / dya[0, -1]) + + ((2.0 + dya[0, -3] / dya[0, -2]) * qin[0, -2, 0] - qin[0, -3, 0]) + / (1.0 + dya[0, -3] / dya[0, -2]) ) ) + (b2 * (qin[0, -1, 0] + qin[0, 2, 0]) + b1 * (qin + qin[0, 1, 0])) @@ -332,7 +344,8 @@ def _qy_edge_south2(qin: FloatField, dya: FloatFieldIJ): @gtscript.function def _qy_edge_north(qin: FloatField, dya: FloatFieldIJ): return 0.5 * ( - ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) / (1.0 + dya[0, -2] / dya[0, -1]) + ((2.0 + dya[0, -2] / dya[0, -1]) * qin[0, -1, 0] - qin[0, -2, 0]) + / (1.0 + dya[0, -2] / dya[0, -1]) + ((2.0 + dya[0, 1] / dya) * qin - qin[0, 1, 0]) / (1.0 + dya[0, 1] / dya) ) @@ -349,12 +362,15 @@ def _qy_edge_north2(qin: FloatField, dya: FloatFieldIJ, qy): return ( 3.0 * (qin[0, -1, 0] + dya[0, -1] / dya * qin) - ( - dya[0, -1] / dya + dya[0, -1] + / dya * ( 0.5 * ( - ((2.0 + dya[0, -1] / dya) * qin - qin[0, -1, 0]) / (1.0 + dya[0, -1] / dya) - + ((2.0 + dya[0, 2] / dya[0, 1]) * qin[0, 1, 0] - qin[0, 2, 0]) / (1.0 + dya[0, 2] / dya[0, 1]) + ((2.0 + dya[0, -1] / dya) * qin - qin[0, -1, 0]) + / (1.0 + dya[0, -1] / dya) + + ((2.0 + dya[0, 2] / dya[0, 1]) * qin[0, 1, 0] - qin[0, 2, 0]) + / (1.0 + dya[0, 2] / dya[0, 1]) ) ) + ( @@ -381,14 +397,21 @@ def _qout_x_edge(qin: FloatField, dxa: FloatFieldIJ, edge_w: FloatFieldIJ): # q2lower = _dxa_weighted_left_average_q(qin[0, -1, 0], dxa[0, -1]) # q2 = _dxa_weighted_left_average_q(qin, dxa) # return edge_w * q2lower + (1.0 - edge_w) * q2 - return edge_w * ((qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) / (dxa[-1, -1] + dxa[0, -1])) + (1.0 - edge_w) * ((qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa)) + return edge_w * ( + (qin[-1, -1, 0] * dxa[0, -1] + qin[0, -1, 0] * dxa[-1, -1]) + / (dxa[-1, -1] + dxa[0, -1]) + ) + (1.0 - edge_w) * ((qin[-1, 0, 0] * dxa + qin * dxa[-1, 0]) / (dxa[-1, 0] + dxa)) + @gtscript.function def _qout_y_edge(qin: FloatField, dya: FloatFieldIJ, edge_s: FloatFieldI): - # q1left = _dya_weighted_lower_average_q(qin[-1, 0, 0], dya[-1, 0]) + # q1left = _dya_weighted_lower_average_q(qin[-1, 0, 0], dya[-1, 0]) # q1 = _dya_weighted_lower_average_q(qin, dya) # return edge_s * q1left + (1.0 - edge_s) * q1 - return edge_s * ((qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) / (dya[-1, -1] + dya[-1, 0])) + (1.0 - edge_s) * ((qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya)) + return edge_s * ( + (qin[-1, -1, 0] * dya[-1, 0] + qin[-1, 0, 0] * dya[-1, -1]) + / (dya[-1, -1] + dya[-1, 0]) + ) + (1.0 - edge_s) * ((qin[0, -1, 0] * dya + qin * dya[0, -1]) / (dya[0, -1] + dya)) @gtscript.function @@ -415,51 +438,89 @@ def _cubic_interpolation_x(qy, qout_offset_x, qyy_offset_x): def qxx_south(qin, qx, dya, edge_s): # qxxupper = lagrange_y(qx[0, 1, 0]) # qoutlower = _qout_y_edge(qin[0, -1, 0], dya[0, -1, 0], edge_s[0, -1]) - #return _cubic_interpolation_y(qx, qoutlower, qxxupper) - return c1 * (qx[0, -1, 0] + qx) + c2 * ((edge_s * ( - (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) - / (dya[-1, -2] + dya[-1, -1]) - ) + (1.0 - edge_s) * ( - (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) - / (dya[0, -2] + dya[0, -1]) - )) + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0]))) + # return _cubic_interpolation_y(qx, qoutlower, qxxupper) + return c1 * (qx[0, -1, 0] + qx) + c2 * ( + ( + edge_s + * ( + (qin[-1, -2, 0] * dya[-1, -1] + qin[-1, -1, 0] * dya[-1, -2]) + / (dya[-1, -2] + dya[-1, -1]) + ) + + (1.0 - edge_s) + * ( + (qin[0, -2, 0] * dya[0, -1] + qin[0, -1, 0] * dya[0, -2]) + / (dya[0, -2] + dya[0, -1]) + ) + ) + + (a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx + qx[0, 1, 0])) + ) + @gtscript.function def qxx_north(qin, qx, dya, edge_n): # qxxlower = lagrange_y(qx[0, -1, 0]) # qoutupper = _qout_y_edge(qin[0, 1, 0], dya[0, 1, 0], edge_s[0, 1]) - #return _cubic_interpolation_y(qx, qoutupper, qxxlower) - return c1 * (qx[0, -1, 0] + qx) + c2 * ((edge_n * ( - (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) - / (dya[-1, 0] + dya[-1, 1]) - ) + (1.0 - edge_n) * ( - (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) / (dya[0, 0] + dya[0, 1]) - )) + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0]))) + # return _cubic_interpolation_y(qx, qoutupper, qxxlower) + return c1 * (qx[0, -1, 0] + qx) + c2 * ( + ( + edge_n + * ( + (qin[-1, 0, 0] * dya[-1, 1] + qin[-1, 1, 0] * dya[-1, 0]) + / (dya[-1, 0] + dya[-1, 1]) + ) + + (1.0 - edge_n) + * ( + (qin[0, 0, 0] * dya[0, 1] + qin[0, 1, 0] * dya[0, 0]) + / (dya[0, 0] + dya[0, 1]) + ) + ) + + (a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] + qx[0, -1, 0])) + ) + @gtscript.function def qyy_west(qin, qy, dxa, edge_w): # qyyright = lagrange_x(qy[0, 1, 0]) # qoutleft = _qout_x_edge(qin[-1, 0, 0], dxa[-1, 0], edge_w) - #return _cubic_interpolation_x(qy, qoutleft, qyyright) - return c1 * (qy[-1, 0, 0] + qy) + c2 * ((edge_w * ( - (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) - / (dxa[-2, -1] + dxa[-1, -1]) - ) + (1.0 - edge_w) * ( - (qin[-2, 0, 0] * dxa[-1, 0] + qin[-1, 0, 0] * dxa[-2, 0]) - / (dxa[-2, 0] + dxa[-1, 0]) - )) + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0]))) + # return _cubic_interpolation_x(qy, qoutleft, qyyright) + return c1 * (qy[-1, 0, 0] + qy) + c2 * ( + ( + edge_w + * ( + (qin[-2, -1, 0] * dxa[-1, -1] + qin[-1, -1, 0] * dxa[-2, -1]) + / (dxa[-2, -1] + dxa[-1, -1]) + ) + + (1.0 - edge_w) + * ( + (qin[-2, 0, 0] * dxa[-1, 0] + qin[-1, 0, 0] * dxa[-2, 0]) + / (dxa[-2, 0] + dxa[-1, 0]) + ) + ) + + (a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy + qy[1, 0, 0])) + ) + @gtscript.function def qyy_east(qin, qy, dxa, edge_e): # qyyleft = lagrange_x(qy[-1, 0, 0]) # qoutright = _qout_x_edge(qin[1, 0, 0], dxa[1, 0], edge_e) - #return _cubic_interpolation_x(qy, qoutright, qyyleft) - return c1 * (qy[-1, 0, 0] + qy) + c2 * ((edge_e * ( - (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) - / (dxa[0, -1] + dxa[1, -1]) - ) + (1.0 - edge_e) * ( - (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) / (dxa[0, 0] + dxa[1, 0]) - )) + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0]))) + # return _cubic_interpolation_x(qy, qoutright, qyyleft) + return c1 * (qy[-1, 0, 0] + qy) + c2 * ( + ( + edge_e + * ( + (qin[0, -1, 0] * dxa[1, -1] + qin[1, -1, 0] * dxa[0, -1]) + / (dxa[0, -1] + dxa[1, -1]) + ) + + (1.0 - edge_e) + * ( + (qin[0, 0, 0] * dxa[1, 0] + qin[1, 0, 0] * dxa[0, 0]) + / (dxa[0, 0] + dxa[1, 0]) + ) + ) + + (a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] + qy[-1, 0, 0])) + ) + def a2b_interpolation_qx( qin: FloatField,