Iinitial commit

.gitignore

+*.nc

coarsen_weights.py

+#!/usr/bin/env python3
+#SBATCH --account=bd1154
+#SBATCH --partition=compute
+#SBATCH --constraint=512G
+#SBATCH --time=02:00:00
+
+import argparse
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+
+def coarsen_weight_file(inweights, outweights, nbisections=1):
+    coarsen = 4**nbisections
+    weights = xr.open_dataset(inweights)
+
+    w = weights[["src_address", "dst_address", "remap_matrix"]].to_dataframe()
+    w.dst_address = ((w.dst_address - 1) // coarsen) + 1
+    w.remap_matrix = w.remap_matrix / coarsen
+
+    new_w = w.groupby(["src_address", "dst_address"]).sum().reset_index()
+    new_w.index.name = "num_links"
+    new_w = new_w.to_xarray()
+
+    new_w["src_grid_imask"] = ("src_grid_size",), np.ones(weights.src_grid_size.size)
+    new_w["src_grid_imask"].attrs = {"units": "unitless"}
+    new_w["src_grid_dims"] = ("src_grid_rank",), [weights.src_grid_size.size]
+    new_w["dst_grid_imask"] = ("dst_grid_size",), np.ones(
+        weights.dst_grid_size.size // coarsen
+    )
+    new_w["dst_grid_imask"].attrs = {"units": "unitless"}
+    new_w["dst_grid_dims"] = ("dst_grid_rank",), [weights.dst_grid_size.size]
+
+    new_w["src_grid_center_lat"] = ("src_grid_size",), np.zeros(
+        new_w.src_grid_size.size
+    )
+    new_w["src_grid_center_lat"].attrs = {"units": "radians"}
+    new_w["src_grid_center_lon"] = ("src_grid_size",), np.zeros(
+        new_w.src_grid_size.size
+    )
+    new_w["src_grid_center_lon"].attrs = {"units": "radians"}
+    new_w["src_grid_frac"] = ("src_grid_size",), np.zeros(new_w.src_grid_size.size)
+    new_w["src_grid_frac"].attrs = {"units": "unitless"}
+
+    new_w["dst_grid_center_lat"] = ("dst_grid_size",), np.zeros(
+        new_w.dst_grid_size.size
+    )
+    new_w["dst_grid_center_lat"].attrs = {"units": "radians"}
+    new_w["dst_grid_center_lon"] = ("dst_grid_size",), np.zeros(
+        new_w.dst_grid_size.size
+    )
+    new_w["dst_grid_center_lon"].attrs = {"units": "radians"}
+    new_w["dst_grid_frac"] = ("dst_grid_size",), np.ones(new_w.dst_grid_size.size)
+    new_w["dst_grid_frac"].attrs = {"units": "unitless"}
+
+    new_w.attrs = weights.attrs
+    new_w["remap_matrix"] = new_w["remap_matrix"].expand_dims("num_wgts", axis=-1)
+    new_w.to_netcdf(outweights)
+
+
+def _main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("infile", type=str)
+    parser.add_argument("outfile", type=str)
+    parser.add_argument("--nbisections", "-n", type=int, default=1)
+    args = parser.parse_args()
+
+    coarsen_weight_file(args.infile, args.outfile, args.nbisections)
+
+
+if __name__ == "__main__":
+    _main()

gen_weights.sh

+#!/bin/bash
+
+grid_id=0058
+grids_folder=/pool/data/ICON/grids/public/mpim/${grid_id}
+grid_refinement=R02B08
+grid_label=O
+grid_name=icon_grid_${grid_id}_${grid_refinement}_${grid_label}
+zoom=13
+
+cdo -P 4 -gennn,hpz${zoom} -const,0,${grids_folder}/icon_grid_${grid_id}_${grid_refinement}_${grid_label}.nc weights_${grid_refinement}_${grid_label}_hpz${zoom}.nc