Merge branch 'main' of github.com:NCAR/SWM

Author: supreethms1809

swm_python/lap_cartesian_vs_next.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# GT4Py - GridTools for Python
+# 
+# Copyright (c) 2014-2023, ETH Zurich
+# All rights reserved.
+# 
+# This file is part the GT4Py project and the GridTools framework.
+# GT4Py is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+# 
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+# # Demonstrates gt4py.cartesian with gt4py.next compatibility
+
+# Imports
+
+# In[19]:
+
+
+import numpy as np
+
+nx = 32
+ny = 32
+nz = 1
+dtype = np.float64
+
+
+# Storages
+# --
+# 
+# We create fields using the gt4py.next constructors. These fields are compatible with gt4py.cartesian when we use "I", "J", "K" as the dimension names.
+
+# In[8]:
+
+
+import gt4py.next as gtx
+
+#allocator = gtx.itir_python # should match the executor
+allocator = gtx.gtfn_cpu
+# allocator = gtx.gtfn_gpu
+
+# Note: for gt4py.next, names don't matter, for gt4py.cartesian they have to be "I", "J", "K"
+I = gtx.Dimension("I")
+J = gtx.Dimension("J")
+K = gtx.Dimension("K", kind=gtx.DimensionKind.VERTICAL)
+
+domain = gtx.domain({I: nx, J: ny, K: nz})
+
+inp = gtx.as_field(domain, np.fromfunction(lambda x, y, z: x**2+y**2, shape=(nx, ny, nz)), dtype, allocator=allocator)
+out_cartesian = gtx.zeros(domain, dtype, allocator=allocator)
+out_next = gtx.zeros(domain, dtype, allocator=allocator)
+
+
+# In[25]:
+
+
+#get_ipython().system('module load gcc')
+#get_ipython().system('export BOOST_HOME=/glade/derecho/scratch/haiyingx/boost_1_84_0/include/boost')
+
+
+# gt4py.cartesian
+# --
+
+# In[26]:
+
+
+import gt4py.cartesian.gtscript as gtscript
+
+#cartesian_backend = "numpy"
+#cartesian_backend = "gt:cpu_ifirst"
+cartesian_backend = "gt:gpu"
+
+@gtscript.stencil(backend=cartesian_backend)
+def lap_cartesian(
+    inp: gtscript.Field[dtype],
+    out: gtscript.Field[dtype],
+):
+    with computation(PARALLEL), interval(...):
+        out = -4.0 * inp[0, 0, 0] + inp[-1, 0, 0] + inp[1, 0, 0] + inp[0, -1, 0] + inp[0, 1, 0]
+
+#lap_cartesian(inp=inp, out=out_cartesian, origin=(1, 1, 0), domain=(nx-2, ny-2, nz))
+
+
+# In[ ]:
+
+
+from gt4py.next import Field
+
+#next_backend = gtx.itir_python
+# next_backend = gtx.gtfn_cpu
+next_backend = gtx.gtfn_gpu
+
+Ioff = gtx.FieldOffset("I", source=I, target=(I,))
+Joff = gtx.FieldOffset("J", source=J, target=(J,))
+
+@gtx.field_operator
+def lap_next(inp: Field[[I, J, K], dtype]) -> Field[[I, J, K], dtype]:
+    return -4.0 * inp + inp(Ioff[-1]) + inp(Ioff[1]) + inp(Joff[-1]) + inp(Joff[1])
+
+@gtx.program(backend=next_backend)
+def lap_next_program(inp: Field[[I, J, K], dtype], out: Field[[I, J, K], dtype]):
+    lap_next(inp, out=out[1:-1, 1:-1, :])
+
+lap_next_program(inp, out_next, offset_provider={"Ioff": I, "Joff": J})
+
+
+# In[ ]:
+
+
+assert np.allclose(out_cartesian.asnumpy(), out_next.asnumpy())
+
+
+# In[ ]:
+
+
+
+

swm_python/swm.ipynb

@@ -35,6 +35,7 @@
         "L_OUT = True # args.L_OUT\n",
         "VIS = True\n",
         "VIS_DT=100\n",
+        "VAL = True\n",
         "ITMAX = 1000\n",
         "dt = 90.\n",
         "tdt = dt\n",
@@ -98,6 +99,28 @@
         "    fig.suptitle(title)\n",
         "    #plt.xlabel('x')\n",
         "    #plt.ylabel('y')\n",
+        "    plt.show()\n",
+        "\n",
+        "def live_plot_val(fu, fv, fp, title=''):\n",
+        "    mxu = fu.max()\n",
+        "    mxv = fv.max()\n",
+        "    mxp = fp.max()\n",
+        "    clear_output(wait=True)\n",
+        "    fig, (ax1, ax2, ax3) = plt.subplots(figsize=(13, 3), ncols=3)\n",
+        "\n",
+        "    pos1 = ax1.imshow(fp, cmap='Blues', vmin=-mxp, vmax=mxp,interpolation='none')\n",
+        "    ax1.set_title('p')\n",
+        "    plt.colorbar(pos1,ax=ax1)\n",
+        "    pos2 = ax2.imshow(fu, cmap='Reds', vmin=-mxu, vmax=mxu,interpolation='none')\n",
+        "    ax2.set_title('u')\n",
+        "    plt.colorbar(pos2,ax=ax2)\n",
+        "    pos3 = ax3.imshow(fv, cmap='Greens',vmin=-mxv, vmax=mxv,interpolation='none')\n",
+        "    ax3.set_title('v')\n",
+        "    plt.colorbar(pos3, ax=ax3)\n",
+        "\n",
+        "    fig.suptitle(title)\n",
+        "    #plt.xlabel('x')\n",
+        "    #plt.ylabel('y')\n",
         "    plt.show()\n"
       ]
     },
@@ -293,7 +316,35 @@
       "id": "78ef5751-4589-4fbc-abfc-a316f672b3b9",
       "metadata": {},
       "outputs": [],
-      "source": []
+      "source": [
+        "if VAL:\n",
+        "\n",
+        "    u_val_f = 'ref/u.64.64.IT4000.txt'\n",
+        "    v_val_f = 'ref/v.64.64.IT4000.txt'\n",
+        "    p_val_f = 'ref/p.64.64.IT4000.txt'\n",
+        "    uval = np.zeros((M_LEN, N_LEN))\n",
+        "    vval = np.zeros((M_LEN, N_LEN))\n",
+        "    pval = np.zeros((M_LEN, N_LEN))\n",
+        "\n",
+        "    uref, vref, pref = read_arrays(v_val_f, u_val_f, p_val_f)\n",
+        "    uval = uref-unew\n",
+        "    vval = vref-vnew\n",
+        "    pval = pref-pnew\n",
+        "    \n",
+        "    uLinfN= np.linalg.norm(uval, np.inf)\n",
+        "    vLinfN= np.linalg.norm(vval, np.inf)\n",
+        "    pLinfN= np.linalg.norm(pval, np.inf)\n",
+        "\n",
+        "    \n",
+        "\n",
+        "    live_plot_val(uval, vval, pval, \"Val\")\n",
+        "    print(\"uLinfN: \", uLinfN)\n",
+        "    print(\"vLinfN: \", vLinfN)\n",
+        "    print(\"pLinfN: \", pLinfN)\n",
+        "    print(\"udiff max: \",uval.max())\n",
+        "    print(\"vdiff max: \",vval.max())\n",
+        "    print(\"pdiff max: \",pval.max())"
+      ]
     }
   ],
   "metadata": {