4dvarnet l63 ronan

config/__init__.py

@@ -1,7 +1,5 @@
 from hydra.core.config_store import ConfigStore
 from omegaconf import OmegaConf
-from pathlib import Path
-from itertools import product
 
 OmegaConf.register_new_resolver(
     "_singleton",
@@ -18,6 +16,11 @@
 )
 
 
+OmegaConf.register_new_resolver("add", lambda x,y: int(x+y), replace=True)
+OmegaConf.register_new_resolver("subtract", lambda x,y: int(x-y), replace=True)
+OmegaConf.register_new_resolver("mul", lambda x,y: int(x)*y, replace=True)
+OmegaConf.register_new_resolver("div", lambda x,y: int(x/y+0.5), replace=True)
+
 class SingletonStore:
     STORE = dict()
 
@@ -34,15 +37,18 @@ def clear(cls):
 
 domains = {
     "eNATL": dict(lon=[-100, 42], lat=[7, 69]),
+    "ceNATL": dict(lon=[-61, -9], lat=[12, 64]),
     "NATL": dict(lon=[-77, 5], lat=[27, 64]),
     "cNATL": dict(lon=[-51, -9], lat=[32, 54]),
     "osmosis": dict(lon=[-22.5, -10.5], lat=[44, 56]),
     "gf": dict(lon=[-66, -54], lat=[32, 44]),
+    "fgf": dict(lon=[-66, -54], lat=[33, 45]),
     "2gf": dict(lon=[-71., -49.], lat=[32, 44]),
     "4gf": dict(lon=[-71., -29.], lat=[32, 44]),
     "calm": dict(lon=[-41., -29.], lat=[32, 44]),
     "qnatl": dict(lon=[-77., 0.], lat=[27., 64.]),
-    "canaries": dict(lon=[-30, -15], lat=[32, 47])
+    "canaries": dict(lon=[-31, -14], lat=[33, 46]),
+    "canaries_t": dict(lon=[-29, -17], lat=[33, 45]),
 }
 
 for n, d in domains.items():

config/xp/base_l63.yaml

@@ -0,0 +1,164 @@
+# @package _global_
+
+xp_name: baseline
+
+
+datamodule:
+  _target_: src.data_l63.BaseDataModule
+  param_datamodule:
+    dl_kw: {batch_size: 128, num_workers: 1}
+  input_data:
+    _target_: src.data_l63.create_l63_datasets
+    param_dataset:
+      path_l63_dataset: '../../Dataset4DVarNet/dataset_l63_normalized_20230717.nc' #'../../Dataset4DVarNet/dataset_L63_with_noise.nc'# '../../Dataset4DVarNet/dataset_L63_with_noise.nc'
+      genSuffixObs: 'JamesData'
+      flag_load_all_data: False
+      flag_generate_L63_data: False #True #
+      NbTraining: 10000
+      NbTest: 1000
+      time_step: 1
+      dT: 200
+      varNoise: 2.0
+      sampling_step: ??
+      flagTypeMissData: 2
+
+model:
+  _target_: src.models_l63.Lit4dVarNet_L63
+  params:
+    k_n_grad: ??
+    n_grad: 10
+    lr: 1e-3 
+    batch_size: 128
+    DimAE: 10 # ?? #
+#    phi_param: 'unet'
+#    solver: ?? #'4dvarnet-with-rnd' #?? #''4dvarnet-with-state-and-rnd' # #'4dvarnet-with-subgradients' #'4dvarnet-with-rnd-grad' #''4dvarnet-with-state-and-rnd' #
+    dim_grad_solver: 100 #25
+    dropout: 0.2
+    alpha_prior: 0.1e3
+    alpha_mse: 1.e3
+    alpha_gmse: 0. #50.
+    alpha_var_cost_grad: 0.00e3
+    lr_grad: ?? #1e-2 #1.e3
+    lr_rnd: 0.e-3
+    sig_rnd_init: 0.e-2
+    sig_lstm_init: 0.e-2
+    sig_obs_noise: 0. #1e-1
+    init_state: 'obs_interp' # ?? #'zeros' #
+    suffix_exp: ?? #'-lstm+grad' #'
+    UsePeriodicBoundary: False
+    degradation_operator: 'no-degradation'  #'gaussian_smoothing' #??  # 
+    sig_perturbation_grad: 1e-2
+    alpha_perturbation_grad: 0.9
+    gamma_degradation: 1e-1
+    type_step_lstm: 'linear-relu' #'constant'  #'linear'#
+    param_lstm_step: 50 #20 #-1 #
+    dt_mse_test: 10
+    post_projection: False #True #  ?? # 
+    post_median_filter: False # True # 
+    median_filter_width: 3
+    shapeData:
+      - 3
+      - ${datamodule.input_data.param_dataset.dT}
+      - 1
+    shapeData_modGrad: ${model.params.shapeData}
+    #  - ${mul:${datamodule.input_data.param_dataset.dT},3} 
+#  ckpt: 'outputs/2023-06-10/23-30-37/base_l63/checkpoints/model-l63-lstmonly+gobs+grad-4dvarnet-with-state-and-rnd-JamesData-Obs20-Noise02-unet-igrad10_08-dgrad100-val_loss=8.5268-epoch=216.ckpt'
+  patch_weight:
+    _target_: src.models_l63.get_constant_crop_l63
+    patch_dims: ${model.params.shapeData}
+    crop: 
+      - 0
+      - 10
+      - 0
+  Phi:
+    _target_: src.models_l63.Phi_unet_like_bilin #src.models_l63.Phi_unet_1_layer_bis #src.models_l63.UNet_3_layers #
+    shapeData: ${model.params.shapeData}
+    DimAE: ${model.params.DimAE}
+#    rateDropout: 0.2
+    name: 'unet-bilin' # 'odenet' # 
+  m_NormObs:
+    _target_: src.solver_l63.Model_WeightedL2Norm #src.solver_l63.Model_WeightedTrainableL2Norm #src.solver_l63.Model_WeightedGL2Norm # src.solver_l63.Model_WeightedL2Norm #
+    #kernel_x: 5
+    #kernel_y: 1
+    #sigma: 1.
+  m_NormPhi:
+    _target_: src.solver_l63.Model_WeightedL2Norm
+  mod_H:
+    _target_: src.models_l63.Model_H7 #src.models_l63.Model_H # src.models_l63.Model_H_with_relu # src.models_l63.Model_HwithSSTBN_nolin_tanh #src.models_l63.Model_HwithTrainableLocalisation #src.models_l63.Model_HwithLocalisation
+    shape_data: ${model.params.shapeData}
+    #dim: 20
+    #sampling: ${datamodule.input_data.param_dataset.sampling_step}
+  mod_Grad:
+    _target_: src.solver_l63.model_Grad_with_lstm
+    ShapeData: ${model.params.shapeData_modGrad}
+    DimLSTM: ${model.params.dim_grad_solver} 
+    rateDropout: ${model.params.dropout} 
+    sig_lstm_init: ${model.params.sig_lstm_init}
+    name: 'lstm-grad'
+
+
+hydra:
+  runtime:
+    output_dir: ./outputs/${hydra.job.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
+
+trainer:
+  _target_: pytorch_lightning.Trainer
+  inference_mode: False
+  #gradient_clip_val: 0.5
+  accelerator: gpu
+  devices: 1
+  logger: 
+    _target_: pytorch_lightning.loggers.CSVLogger
+    save_dir: ${hydra:runtime.output_dir}
+    name: ${hydra:runtime.choices.xp}
+    version: ''
+  max_epochs: 800
+  callbacks:
+    - _target_: src.versioning_cb.VersioningCallback
+    - _target_: pytorch_lightning.callbacks.LearningRateMonitor
+    - _target_: pytorch_lightning.callbacks.ModelCheckpoint
+      monitor: val_mse
+      save_top_k: 3
+      filename: 
+        _target_: src.models_l63.create_filename_ckpt
+        suffix: '-{val_loss:.4f}-{epoch:03d}'#'{val_mse:.4f}-{epoch:03d}'
+        params_data: ${datamodule.input_data.param_dataset}
+        params_model: ${model.params}
+        name_solver: ${model.mod_Grad.name}
+        name_phi: ${model.Phi.name}
+
+entrypoints:
+  - _target_: pytorch_lightning.seed_everything
+    seed: 333
+#  - _target_: src.train.base_training  # src.train.fine_tuning # 
+#    trainer: ${trainer}
+#    lit_mod: ${model}
+#    dm: ${datamodule}
+#    update_params: False
+#    ckpt: 'outputs/2023-07-21/18-36-53/base_l63/checkpoints/model-l63-interp-x1only-dT200--convlstm-grad-JamesData-Obs20-Noise02-odenet_10-igrad10_08-dgrad100-val_loss=32.4257-epoch=101.ckpt'
+#    ckpt: 'outputs/2023-06-08/23-38-13/base_l63/checkpoints/model-l63-lstmonly+gobs+grad-4dvarnet-with-rnd-JamesData-Obs20-Noise02-unet-igrad10_04-dgrad100-val_loss=15.7092-epoch=014.ckpt'
+#    ckpt: 'outputs/2023-06-08/23-23-08/base_l63/checkpoints/model-l63-lstmonly+gobs-4dvarnet-with-rnd-JamesData-Obs20-Noise02-unet-igrad10_04-dgrad100-val_loss=9.7803-epoch=092.ckpt'
+#    ckpt: 'outputs/2023-06-08/18-32-56/base_l63/checkpoints/model-l63-lstmonly+gobs-4dvarnet-with-rnd-JamesData-Obs20-Noise02-unet-igrad10_08-dgrad100-val_loss=8.8370-epoch=222.ckpt'
+#    ckpt: 'outputs/2023-05-24/16-55-55/base_l63/checkpoints/model-l63-4dvarnet-with-state-and-rnd-JamesData-Obs08-Noise02-unet-igrad10_04-dgrad100-val_loss=3.2830-epoch=179.ckpt'
+#    ckpt: 'outputs/2023-05-23/22-14-34/base_l63/checkpoints/model-l63-4dvarnet-with-state-and-rnd-JamesData-Obs08-Noise02-unet-igrad10_04-dgrad100-val_loss=3.0530-epoch=337.ckpt'
+#    ckpt: 'outputs/2023-06-10/23-30-37/base_l63/checkpoints/model-l63-lstmonly+gobs+grad-4dvarnet-with-state-and-rnd-JamesData-Obs20-Noise02-unet-igrad10_08-dgrad100-val_loss=8.5268-epoch=216.ckpt'
+  - _target_: src.test.base_testing #
+    trainer: ${trainer}
+    lit_mod: ${model}
+    dm: ${datamodule}
+    num_members: 1
+    ckpt: 'outputs/2023-07-24/08-42-27/base_l63/checkpoints/model-l63-interp-x1onlylstm+grad-modH7-dT200--convlstm-grad-JamesData-Obs20-Noise02-odenet_10-igrad10_04-dgrad100-val_loss=13.5492-epoch=774.ckpt'
+#    ckpt: 'outputs/2023-07-23/22-23-55/base_l63/checkpoints/model-l63-interp-x1onlylstm+grad-modH6-50-dT200--convlstm-grad-JamesData-Obs20-Noise02-odenet_10-igrad10_04-dgrad100-val_loss=14.4306-epoch=298.ckpt'
+#    ckpt: 'outputs/2023-07-23/20-08-33/base_l63/checkpoints/model-l63-interp-x1onlylstm+grad-modH5-20-dT200--convlstm-grad-JamesData-Obs20-Noise02-odenet_10-igrad10_04-dgrad100-val_loss=13.4317-epoch=415.ckpt'
+#    ckpt: 'outputs/2023-07-23/10-37-19/base_l63/checkpoints/model-l63-interp-x1onlylstm+grad-modH2bis-20-dT200--convlstm-grad-JamesData-Obs20-Noise02-odenet_10-igrad10_04-dgrad100-val_loss=13.5894-epoch=412.ckpt'
+#    ckpt: 'outputs/2023-07-21/18-11-32/base_l63/checkpoints/model-l63-interp-x1only-dT200--convlstm-grad-JamesData-Obs20-Noise02-odenet_10-igrad10_04-dgrad100-val_loss=13.8485-epoch=380.ckpt'
+#    ckpt: 'outputs/2023-07-21/17-55-45/base_l63/checkpoints/model-l63-interp-x1only-dT200--convlstm-grad-JamesData-Obs20-Noise02-odenet_10-igrad10_04-dgrad100-val_loss=14.8211-epoch=218.ckpt'
+#    ckpt: 'outputs/2023-07-22/23-20-51/base_l63/checkpoints/model-l63-interp-x1onlylstm+grad-modH2-50-dT200--convlstm-grad-JamesData-Obs20-Noise02-odenet_10-igrad10_04-dgrad100-val_loss=16.6168-epoch=375.ckpt'
+#    test_fn:
+#      _target_: src.utils.diagnostics
+#      _partial_: true
+#      test_domain:
+#        time: {_target_: builtins.slice, _args_: ["2012-10-22", "2012-12-02"]}
+#        lat: ${domain.test.lat}
+#        lon: ${domain.test.lat}
+