intermediate commit for global run script

technically, this should work now, although I need to find a way to parallelise the embarrassing loop and possibly move the writing routine out. I will also need to implement the skipping of ocean grid cells. Finally, the south pole looks pretty okay.

Author: ray-chew

inputs/icon_global_run.py

@@ -0,0 +1,32 @@
+import numpy as np
+from src import var
+
+params = var.params()
+
+params.output_path = "/home/ray/git-projects/spec_appx/outputs/"
+params.output_fn = "icon_merit_reg"
+params.fn_grid = "../data/icon_compact.nc"
+params.fn_topo = "../data/topo_compact.nc"
+
+### South Pole
+params.lat_extent = None
+params.lon_extent = None
+
+params.tri_set = [13, 104, 105, 106]
+
+# Setup the Fourier parameters and object.
+params.nhi = 24
+params.nhj = 48
+
+params.n_modes = 50
+
+params.U, params.V = 10.0, 0.0
+
+params.rect = True
+
+params.debug = False
+params.dfft_first_guess = True
+params.refine = False
+params.verbose = False
+
+params.plot = True

runs/icon_merit_global.py

@@ -0,0 +1,222 @@
+# %%
+import sys
+
+# set system path to find local modules
+sys.path.append("..")
+
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+from src import io, var, utils, fourier, physics
+from wrappers import interface
+from vis import plotter, cart_plot
+
+from IPython import get_ipython
+
+ipython = get_ipython()
+
+if ipython is not None:
+    ipython.run_line_magic("load_ext", "autoreload")
+else:
+    print(ipython)
+
+def autoreload():
+    if ipython is not None:
+        ipython.run_line_magic("autoreload", "2")
+
+from sys import exit
+
+if __name__ != "__main__":
+    exit(0)
+# %%
+autoreload()
+from inputs.icon_regional_run import params
+
+if params.self_test():
+    params.print()
+
+print(params.path_compact_topo)
+
+grid = var.grid()
+
+# read grid
+reader = io.ncdata(padding=params.padding, padding_tol=(60 - params.padding))
+
+# writer object
+writer = io.nc_writer(params)
+
+reader.read_dat(params.path_compact_grid, grid)
+
+clat_rad = np.copy(grid.clat)
+clon_rad = np.copy(grid.clon)
+
+grid.apply_f(utils.rad2deg)
+
+n_cells = grid.clat.size
+
+for c_idx in range(n_cells)[:1]:
+    c_idx = 90
+    print(c_idx)
+
+    topo = var.topo_cell()
+    lat_verts = grid.clat_vertices[c_idx]
+    lon_verts = grid.clon_vertices[c_idx]
+
+    lat_extent = [lat_verts.min() - 1.0,lat_verts.min() - 1.0,lat_verts.max() + 1.0]
+    lon_extent = [lon_verts.min() - 1.0,lon_verts.min() - 1.0,lon_verts.max() + 1.0]
+    # we only keep the topography that is inside this lat-lon extent.
+    lat_verts = np.array(lat_extent)
+    lon_verts = np.array(lon_extent)
+
+    params.lat_extent = lat_extent
+    params.lon_extent = lon_extent
+
+    # read topography
+    if not params.enable_merit:
+        reader.read_dat(params.fn_topo, topo)
+        reader.read_topo(topo, topo, lon_verts, lat_verts)
+    else:
+        reader.read_merit_topo(topo, params)
+        topo.topo[np.where(topo.topo < -500.0)] = -500.0
+
+    topo.gen_mgrids()
+# %%
+    clon = np.array([grid.clon[c_idx]])
+    clat = np.array([grid.clat[c_idx]])
+    clon_vertices = np.array([grid.clon_vertices[c_idx]])
+    clat_vertices = np.array([grid.clat_vertices[c_idx]])
+
+    ncells = 1
+    nv = clon_vertices[0].size
+    # -- create the triangles
+    clon_vertices = np.where(clon_vertices < -180.0, clon_vertices + 360.0, clon_vertices)
+    clon_vertices = np.where(clon_vertices > 180.0, clon_vertices - 360.0, clon_vertices)
+
+    triangles = np.zeros((ncells, nv, 2), np.float32)
+
+    for i in range(0, ncells, 1):
+        triangles[i, :, 0] = np.array(clon_vertices[i, :])
+        triangles[i, :, 1] = np.array(clat_vertices[i, :])
+
+    print("--> triangles done")
+
+    cart_plot.lat_lon_icon(topo, triangles, ncells=ncells, clon=clon, clat=clat)
+
+
+# %%
+    tri_idx = 0
+    # initialise cell object
+    cell = var.topo_cell()
+
+    simplex_lon = triangles[tri_idx, :, 0]
+    simplex_lat = triangles[tri_idx, :, 1]
+
+    utils.get_lat_lon_segments(
+        simplex_lat, simplex_lon, cell, topo, rect=params.rect
+    )
+
+    topo_orig = np.copy(cell.topo)
+
+    if params.dfft_first_guess:
+        nhi = len(cell.lon)
+        nhj = len(cell.lat)
+
+    first_guess = interface.get_pmf(nhi, nhj, params.U, params.V)
+    fobj_tri = fourier.f_trans(nhi, nhj)
+
+    #######################################################
+    # do fourier...
+
+    if not params.dfft_first_guess:
+        freqs, uw_pmf_freqs, dat_2D_fg0 = first_guess.sappx(cell, lmbda=0.0)
+
+    #######################################################
+    # do fourier using DFFT
+
+    if params.dfft_first_guess:
+        ampls, uw_pmf_freqs, dat_2D_fg0, kls = first_guess.dfft(cell)
+        freqs = np.copy(ampls)
+
+    print("uw_pmf_freqs_sum:", uw_pmf_freqs.sum())
+
+    fq_cpy = np.copy(freqs)
+
+    indices = []
+    max_ampls = []
+
+    for ii in range(params.n_modes):
+        max_idx = np.unravel_index(fq_cpy.argmax(), fq_cpy.shape)
+        indices.append(max_idx)
+        max_ampls.append(fq_cpy[max_idx])
+        max_val = fq_cpy[max_idx]
+        fq_cpy[max_idx] = 0.0
+
+    utils.get_lat_lon_segments(
+        simplex_lat, simplex_lon, cell, topo, rect=False
+    )
+
+    k_idxs = [pair[1] for pair in indices]
+    l_idxs = [pair[0] for pair in indices]
+
+    second_guess = interface.get_pmf(nhi, nhj, params.U, params.V)
+
+    if params.dfft_first_guess:
+        second_guess.fobj.set_kls(
+            k_idxs, l_idxs, recompute_nhij=True, components="real"
+        )
+    else:
+        second_guess.fobj.set_kls(k_idxs, l_idxs, recompute_nhij=False)
+
+    freqs, uw, dat_2D_sg0 = second_guess.sappx(cell, lmbda=1e-1, updt_analysis=True)
+
+    cell.topo = topo_orig
+
+    writer.output(tri_idx, clat_rad[tri_idx], clon_rad[tri_idx], cell.analysis)
+    
+    cell.uw = uw
+
+    if params.plot:
+        fs = (15, 9.0)
+        v_extent = [dat_2D_sg0.min(), dat_2D_sg0.max()]
+
+        fig, axs = plt.subplots(2, 2, figsize=fs)
+
+        fig_obj = plotter.fig_obj(
+            fig, second_guess.fobj.nhar_i, second_guess.fobj.nhar_j
+        )
+        axs[0, 0] = fig_obj.phys_panel(
+            axs[0, 0],
+            dat_2D_sg0,
+            title="T%i: Reconstruction" % tri_idx,
+            xlabel="longitude [km]",
+            ylabel="latitude [km]",
+            extent=[cell.lon.min(), cell.lon.max(), cell.lat.min(), cell.lat.max()],
+            v_extent=v_extent,
+        )
+
+        axs[0, 1] = fig_obj.phys_panel(
+            axs[0, 1],
+            cell.topo * cell.mask,
+            title="T%i: Reconstruction" % tri_idx,
+            xlabel="longitude [km]",
+            ylabel="latitude [km]",
+            extent=[cell.lon.min(), cell.lon.max(), cell.lat.min(), cell.lat.max()],
+            v_extent=v_extent,
+        )
+
+        if params.dfft_first_guess:
+            axs[1, 0] = fig_obj.fft_freq_panel(
+                axs[1, 0], freqs, kls[0], kls[1], typ="real"
+            )
+            axs[1, 1] = fig_obj.fft_freq_panel(
+                axs[1, 1], uw, kls[0], kls[1], title="PMF spectrum", typ="real"
+            )
+        else:
+            axs[1, 0] = fig_obj.freq_panel(axs[1, 0], freqs)
+            axs[1, 1] = fig_obj.freq_panel(axs[1, 1], uw, title="PMF spectrum")
+
+        plt.tight_layout()
+        plt.savefig("%sT%i.pdf" % (params.path_output, tri_idx))
+        plt.show()
+# %%