2d flow

projects/super_res/config.py

@@ -1,6 +1,5 @@
 from ml_collections import config_dict
 
-#batch_size = 4
 config = config_dict.ConfigDict()
 
 config.dim = 64
@@ -9,34 +8,36 @@
 config.random_fourier_features = True,
 config.learned_sinusoidal_dim = 32
 config.diffusion_steps = 1500
-config.sampling_steps = 6
-config.loss = "l1"
+config.sampling_steps = 20
+config.loss = "l2"
 config.objective = "pred_v"
-config.lr = 8e-5
-config.steps = 5000000
+config.lr = 1e-4
+config.steps = 700000
 config.grad_acc = 1
-config.val_num_of_batch = 1
-config.save_and_sample_every = 5000
+config.val_num_of_batch = 5
+config.save_and_sample_every = 20000
 config.ema_decay = 0.995
 config.amp = False
 config.split_batches = True
-config.additional_note = ""
+config.additional_note = "2d-nomulti-nols-ensemble"
 config.eval_folder = "./evaluate"
 config.results_folder = "./results"
 config.tensorboard_dir = "./tensorboard"
 config.milestone = 1
+config.rollout = None
+config.rollout_batch = None
 
 config.batch_size = 1
 config.data_config = config_dict.ConfigDict({
     "dataset_name": "c384",
     "length": 7,
-    #"channels": ["UGRD10m_coarse","VGRD10m_coarse"],
     "channels": ["PRATEsfc_coarse"],
-    #"img_channel": 2,
     "img_channel": 1,
     "img_size": 384,
-    "logscale": True,
-    "quick": True
+    "logscale": False,
+    "multi": False,
+    "flow": "2d",
+    "minipatch": False
 })
 
 config.data_name = f"{config.data_config['dataset_name']}-{config.data_config['channels']}-{config.objective}-{config.loss}-d{config.dim}-t{config.diffusion_steps}{config.additional_note}"

projects/super_res/config_focal.py

@@ -0,0 +1,44 @@
+from ml_collections import config_dict
+
+config = config_dict.ConfigDict()
+
+config.dim = 64
+config.dim_mults = (1, 1, 2, 2, 3, 4)
+config.learned_sinusoidal_cond = True,
+config.random_fourier_features = True,
+config.learned_sinusoidal_dim = 32
+config.diffusion_steps = 1500
+config.sampling_steps = 20
+config.loss = "focal"
+config.objective = "pred_v"
+config.lr = 1e-4
+config.steps = 700000
+config.grad_acc = 1
+config.val_num_of_batch = 5
+config.save_and_sample_every = 20000
+config.ema_decay = 0.995
+config.amp = False
+config.split_batches = True
+config.additional_note = "2d-multi-ls-focal-ensemble"
+config.eval_folder = "./evaluate"
+config.results_folder = "./results"
+config.tensorboard_dir = "./tensorboard"
+config.milestone = 1
+config.rollout = None
+config.rollout_batch = None
+
+config.batch_size = 1
+config.data_config = config_dict.ConfigDict({
+    "dataset_name": "c384",
+    "length": 7,
+    "channels": ["PRATEsfc_coarse"],
+    "img_channel": 1,
+    "img_size": 384,
+    "logscale": True,
+    "multi": True,
+    "flow": "2d",
+    "minipatch": False
+})
+
+config.data_name = f"{config.data_config['dataset_name']}-{config.data_config['channels']}-{config.objective}-{config.loss}-d{config.dim}-t{config.diffusion_steps}{config.additional_note}"
+config.model_name = f"c384-{config.data_config['channels']}-{config.objective}-{config.loss}-d{config.dim}-t{config.diffusion_steps}{config.additional_note}"

projects/super_res/config_infer.py

@@ -2,41 +2,43 @@
 
 config = config_dict.ConfigDict()
 
-
 config.dim = 64
-config.dim_mults = (1, 1, 2, 2, 4, 4)
+config.dim_mults = (1, 1, 2, 2, 3, 4)
 config.learned_sinusoidal_cond = True,
 config.random_fourier_features = True,
 config.learned_sinusoidal_dim = 32
 config.diffusion_steps = 1500
-config.sampling_steps = 6
-config.loss = "l1"
+config.sampling_steps = 20
+config.loss = "l2"
 config.objective = "pred_v"
-config.lr = 8e-5
-config.steps = 5000000
-config.grad_acc = 2
+config.lr = 1e-4
+config.steps = 700000
+config.grad_acc = 1
 config.val_num_of_batch = 5
-config.save_and_sample_every = 5000
+config.save_and_sample_every = 20000
 config.ema_decay = 0.995
 config.amp = False
 config.split_batches = True
-config.additional_note = ""
+config.additional_note = "2d-nomulti-nols-ensemble"
 config.eval_folder = "./evaluate"
 config.results_folder = "./results"
 config.tensorboard_dir = "./tensorboard"
-config.milestone = 1
+config.milestone = 2
+config.rollout = "partial"
+config.rollout_batch = 25
 
-config.batch_size = 4
+config.batch_size = 2
 config.data_config = config_dict.ConfigDict({
     "dataset_name": "c384",
     "length": 7,
-    #"channels": ["UGRD10m_coarse","VGRD10m_coarse"],
     "channels": ["PRATEsfc_coarse"],
-    #"img_channel": 2,
     "img_channel": 1,
     "img_size": 384,
-    "logscale": True
+    "logscale": False,
+    "multi": False,
+    "flow": "2d",
+    "minipatch": False
 })
 
-data_name = f"{config.data_config['dataset_name']}-{config.data_config['channels']}-{config.objective}-{config.loss}-d{config.dim}-t{config.diffusion_steps}{config.additional_note}"
-model_name = f"c384-{config.data_config['channels']}-{config.objective}-{config.loss}-d{config.dim}-t{config.diffusion_steps}{config.additional_note}"
+config.data_name = f"{config.data_config['dataset_name']}-{config.data_config['channels']}-{config.objective}-{config.loss}-d{config.dim}-t{config.diffusion_steps}{config.additional_note}"
+config.model_name = f"c384-{config.data_config['channels']}-{config.objective}-{config.loss}-d{config.dim}-t{config.diffusion_steps}{config.additional_note}"

projects/super_res/config_isr.py

@@ -0,0 +1,37 @@
+from ml_collections import config_dict
+
+config = config_dict.ConfigDict()
+
+config.lr = 1e-4
+config.steps = 700000
+config.grad_acc = 1
+config.val_num_of_batch = 5
+config.save_and_sample_every = 20000
+config.ema_decay = 0.995
+config.amp = False
+config.split_batches = True
+config.additional_note = "isr"
+config.eval_folder = "./evaluate"
+config.results_folder = "./results"
+config.tensorboard_dir = "./tensorboard"
+config.milestone = 1
+config.rollout = None
+config.rollout_batch = None
+
+config.batch_size = 1
+config.data_config = config_dict.ConfigDict({
+    "dataset_name": "c384",
+    "length": 7,
+    #"channels": ["UGRD10m_coarse","VGRD10m_coarse"],
+    "channels": ["PRATEsfc_coarse"],
+    #"img_channel": 2,
+    "img_channel": 1,
+    "img_size": 384,
+    "logscale": True,
+    "multi": True,
+    "flow": "2d",
+    "minipatch": False
+})
+
+config.data_name = f"{config.data_config['dataset_name']}-{config.data_config['channels']}-{config.additional_note}"
+config.model_name = f"c384-{config.data_config['channels']}-{config.additional_note}"

projects/super_res/config_isr_infer.py

@@ -0,0 +1,37 @@
+from ml_collections import config_dict
+
+config = config_dict.ConfigDict()
+
+config.lr = 1e-4
+config.steps = 700000
+config.grad_acc = 1
+config.val_num_of_batch = 5
+config.save_and_sample_every = 20000
+config.ema_decay = 0.995
+config.amp = False
+config.split_batches = True
+config.additional_note = "isr"
+config.eval_folder = "./evaluate"
+config.results_folder = "./results"
+config.tensorboard_dir = "./tensorboard"
+config.milestone = 2
+config.rollout = 'partial'
+config.rollout_batch = 25
+
+config.batch_size = 1
+config.data_config = config_dict.ConfigDict({
+    "dataset_name": "c384",
+    "length": 7,
+    #"channels": ["UGRD10m_coarse","VGRD10m_coarse"],
+    "channels": ["PRATEsfc_coarse"],
+    #"img_channel": 2,
+    "img_channel": 1,
+    "img_size": 384,
+    "logscale": True,
+    "multi": True,
+    "flow": "2d",
+    "minipatch": False
+})
+
+config.data_name = f"{config.data_config['dataset_name']}-{config.data_config['channels']}-{config.additional_note}"
+config.model_name = f"c384-{config.data_config['channels']}-{config.additional_note}"

projects/super_res/config_rvrt_full.py

@@ -0,0 +1,50 @@
+from ml_collections import config_dict
+
+#batch_size = 4
+config = config_dict.ConfigDict()
+
+config.dim = 120
+config.num_blocks = 6
+config.num_heads = 8
+config.depth = 8
+config.time_emb_dim = 32
+config.learned_sinusoidal_cond = True
+config.diffusion_steps = 1500
+config.sampling_steps = 20
+# config.loss = "l2"
+config.loss = "charbonnier"
+config.objective = "pred_x0"
+# config.lr = 8e-5
+config.lr = 1e-4
+# config.steps = 500000
+config.steps = 700000
+config.grad_acc = 1
+config.val_num_of_batch = 5
+config.save_and_sample_every = 20000
+config.ema_decay = 0.999
+config.amp = False
+config.split_batches = True
+config.additional_note = "rvrt_full"
+config.eval_folder = "./evaluate"
+config.results_folder = "./results"
+config.tensorboard_dir = "./tensorboard"
+config.milestone = 1
+config.rollout = None
+config.rollout_batch = None
+
+config.batch_size = 1
+config.data_config = config_dict.ConfigDict({
+    "dataset_name": "c384",
+    "length": 6,
+    #"channels": ["UGRD10m_coarse","VGRD10m_coarse"],
+    "channels": ["PRATEsfc_coarse"],
+    #"img_channel": 2,
+    "img_channel": 1,
+    "img_size": 384,
+    "logscale": True,
+    "multi": True,
+    "minipatch": False
+})
+
+config.data_name = f"{config.data_config['dataset_name']}-{config.data_config['channels']}-{config.additional_note}"
+config.model_name = f"c384-{config.data_config['channels']}-{config.additional_note}"

projects/super_res/config_rvrt_full_infer.py

@@ -0,0 +1,50 @@
+from ml_collections import config_dict
+
+#batch_size = 4
+config = config_dict.ConfigDict()
+
+config.dim = 120
+config.num_blocks = 6
+config.num_heads = 8
+config.depth = 8
+config.time_emb_dim = 32
+config.learned_sinusoidal_cond = True
+config.diffusion_steps = 1500
+config.sampling_steps = 20
+# config.loss = "l2"
+config.loss = "charbonnier"
+config.objective = "pred_x0"
+# config.lr = 8e-5
+config.lr = 1e-4
+# config.steps = 500000
+config.steps = 700000
+config.grad_acc = 1
+config.val_num_of_batch = 5
+config.save_and_sample_every = 20000
+config.ema_decay = 0.999
+config.amp = False
+config.split_batches = True
+config.additional_note = "rvrt_full"
+config.eval_folder = "./evaluate"
+config.results_folder = "./results"
+config.tensorboard_dir = "./tensorboard"
+config.milestone = 2
+config.rollout = 'partial'
+config.rollout_batch = 22
+
+config.batch_size = 1
+config.data_config = config_dict.ConfigDict({
+    "dataset_name": "c384",
+    "length": 6,
+    #"channels": ["UGRD10m_coarse","VGRD10m_coarse"],
+    "channels": ["PRATEsfc_coarse"],
+    #"img_channel": 2,
+    "img_channel": 1,
+    "img_size": 384,
+    "logscale": True,
+    "multi": True,
+    "minipatch": False
+})
+
+config.data_name = f"{config.data_config['dataset_name']}-{config.data_config['channels']}-{config.additional_note}"
+config.model_name = f"c384-{config.data_config['channels']}-{config.additional_note}"

projects/super_res/data/channel_data_gen.py

@@ -0,0 +1,29 @@
+import xarray as xr
+import numpy as np
+from pathlib import Path
+
+channel_folder = Path('./more_channels')
+channel_folder.mkdir(exist_ok = True, parents = True)
+
+c384 = xr.open_zarr("gs://vcm-ml-raw-flexible-retention/2021-07-19-PIRE/C3072-to-C384-res-diagnostics/pire_atmos_phys_3h_coarse.zarr").rename({"grid_xt_coarse": "x", "grid_yt_coarse": "y"})
+c48 = xr.open_zarr("gs://vcm-ml-intermediate/2021-10-12-PIRE-c48-post-spinup-verification/pire_atmos_phys_3h_coarse.zarr").rename({"grid_xt": "x", "grid_yt": "y"})
+
+channels = ["UGRD10m_coarse", "VGRD10m_coarse", "tsfc_coarse", "CPRATEsfc_coarse"]
+c384_np = np.stack([c384[channel].values for channel in channels], axis = 2)
+c48_np = np.stack([c48[channel].values for channel in channels], axis = 2)
+
+split = int(c384_np.shape[1] * 0.8)
+
+# compute statistics on training set
+c384_min, c384_max, c48_min, c48_max = c384_np[:, :split, :, :, :].min(axis=(0,1,3,4)).reshape(1,1,4,1,1), c384_np[:, :split, :, :, :].max(axis=(0,1,3,4)).reshape(1,1,4,1,1), c48_np[:, :split, :, :, :].min(axis=(0,1,3,4)).reshape(1,1,4,1,1), c48_np[:, :split, :, :, :].max(axis=(0,1,3,4)).reshape(1,1,4,1,1) 
+
+# normalize
+c384_norm= (c384_np - c384_min) / (c384_max - c384_min)
+c48_norm = (c48_np - c48_min) / (c48_max - c48_min)
+
+np.save('more_channels/c384_min.npy', c384_min)
+np.save('more_channels/c384_max.npy', c384_max)
+np.save('more_channels/c48_min.npy', c48_min)
+np.save('more_channels/c48_max.npy', c48_max)
+np.save('more_channels/c48_norm.npy', c48_norm)
+np.save('more_channels/c384_norm.npy', c384_norm)

projects/super_res/data/dataload.sh

@@ -0,0 +1,13 @@
+#! /bin/sh
+channel='c48_atmos_ave'
+file='atmos_8xdaily_ave_coarse.zarr'
+for member in $(seq -f "%04g" 1 11)
+do
+    mkdir -p /data/prakhars/ensemble/$channel/$member
+    gsutil -m cp -r gs://vcm-ml-raw-flexible-retention/2023-08-14-C384-reference-ensemble/ic_$member/diagnostics/$file /data/prakhars/ensemble/$channel/$member
+done
+# channel --> file
+# c384_precip_ave --> sfc_8xdaily_ave.zarr
+# c48_precip_plus_more_ave --> sfc_8xdaily_ave_coarse.zarr
+# c384_topo --> atmos_static.zarr
+# c48_atmos_ave --> atmos_8xdaily_ave_coarse.zarr