diff --git a/runs/icon_merit_regional.py b/runs/icon_merit_regional.py index 95da513..f3a28d1 100644 --- a/runs/icon_merit_regional.py +++ b/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) \ No newline at end of file + idxs.append(tri_idx) + pmfs.append(uw_comp) + + +# %% diff --git a/src/io.py b/src/io.py index 8f40cd8..bda5636 100644 --- a/src/io.py +++ b/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`"""