implemented NetCDF4 writer class; outputs data structure according to…

… MS-GWaM's IO module

runs/icon_merit_regional.py

@@ -43,9 +43,13 @@ def autoreload():
 reader = io.ncdata(padding=params.padding, padding_tol=(60 - params.padding))
 
 # writer object
-# writer = io.writer(params.output_fn, params.rect_set, debug=params.debug_writer)
+writer = io.nc_writer(params)
 
 reader.read_dat(params.fn_grid, grid)
+
+clat_rad = np.copy(grid.clat)
+clon_rad = np.copy(grid.clon)
+
 grid.apply_f(utils.rad2deg)
 
 # we only keep the topography that is inside this lat-lon extent.
@@ -99,7 +103,7 @@ def autoreload():
 
 
 # %%
-
+autoreload()
 idxs = []
 pmfs = []
 
@@ -170,6 +174,8 @@ def autoreload():
 
     cell.topo = topo_orig
 
+    writer.output(tri_idx, clat_rad[tri_idx], clon_rad[tri_idx], cell.analysis)
+
     cell.uw = uw
 
     if params.plot:
@@ -216,8 +222,11 @@ def autoreload():
         plt.savefig("../output/T%i.pdf" % tri_idx)
         plt.show()
 
-        ideal = physics.ideal_pmf(U=params.U, V=params.V)
-        uw_comp = ideal.compute_uw_pmf(cell.analysis)
+    ideal = physics.ideal_pmf(U=params.U, V=params.V)
+    uw_comp = ideal.compute_uw_pmf(cell.analysis)
 
-        idxs.append(tri_idx)
-        pmfs.append(uw_comp)
+    idxs.append(tri_idx)
+    pmfs.append(uw_comp)
+
+
+# %%

src/io.py

@@ -544,6 +544,77 @@ def populate(self, idx, name, data):
         file.close()
 
 
+class nc_writer(object):
+
+    def __init__(self, params):
+
+        self.fn = params.output_fn
+
+        if self.fn[-3:] != ".nc":
+            self.fn += '.nc'
+
+        self.path = params.output_path
+        self.rect_set = params.rect_set
+        self.debug = params.debug_writer
+
+        rootgrp = nc.Dataset(self.path + self.fn, "w", format="NETCDF4")
+
+        _ = rootgrp.createDimension("nspec", params.n_modes)
+
+        self.n_modes = params.n_modes
+        rootgrp.close()
+
+    def output(self, id, clat, clon, analysis):
+
+        rootgrp = nc.Dataset(self.path + self.fn, "a", format="NETCDF4")
+
+        grp = rootgrp.createGroup(str(id))
+
+        is_land_var = grp.createVariable("is_land","i4")
+        is_land_var[:] = 1
+
+        clat_var = grp.createVariable("clat","f8")
+        clat_var[:] = clat
+        clon_var = grp.createVariable("clon","f8")
+        clon_var[:] = clon
+
+        dk_var = grp.createVariable("dk","f8")
+        dk_var[:] = analysis.dk
+        dl_var = grp.createVariable("dl","f8")
+        dl_var[:] = analysis.dl
+
+        pick_idx = np.where(analysis.ampls > 0)
+
+        H_spec_var = grp.createVariable("H_spec","f8", ("nspec",))
+        H_spec_var[:] = self.pad_zeros(analysis.ampls[pick_idx], self.n_modes)
+
+        kks_var = grp.createVariable("kks","f8", ("nspec",))
+        kks_var[:] = self.pad_zeros(analysis.kks[pick_idx], self.n_modes)
+
+        lls_var = grp.createVariable("lls","f8", ("nspec",))
+        lls_var[:] = self.pad_zeros(analysis.lls[pick_idx], self.n_modes)
+
+
+
+
+        rootgrp.close()
+
+
+    @staticmethod
+    def pad_zeros(lst, n_modes):
+
+        if lst.size < n_modes:
+            pad_len = n_modes - lst.size
+        else:
+            pad_len = 0
+
+        return np.concatenate((lst, np.zeros((pad_len))))
+
+
+
+
+
+
 class reader(object):
     """Simple reader class to read HDF5 output written by :class:`src.io.writer`"""