domain:
  train:
    lat:
      _target_: builtins.slice
      _args_:
      - 32
      - 44
    lon:
      _target_: builtins.slice
      _args_:
      - -66
      - -54
  test:
    lat:
      _target_: builtins.slice
      _args_:
      - 33
      - 43
    lon:
      _target_: builtins.slice
      _args_:
      - -65
      - -55
paths:
  natl_gf_5nads: data/natl_gf_w_5nadirs.nc
  enatl_gf_5nads: ../sla-data-registry/qdata/enatl_wo_tide.nc
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: 300
  callbacks:
  - _target_: src.versioning_cb.VersioningCallback
  - _target_: pytorch_lightning.callbacks.LearningRateMonitor
  - _target_: pytorch_lightning.callbacks.ModelCheckpoint
    monitor: val_rmse
    save_top_k: 10
    filename: '{val_rmse:.4f}-{epoch:03d}'
datamodule:
  _target_: src.data.BaseDataModule
  input_da:
    _target_: src.utils.load_altimetry_data
    path: ${paths.natl_gf_5nads}
  domains:
    train:
      time:
        _target_: builtins.slice
        _args_:
        - '2013-02-24'
        - '2013-09-30'
    val:
      time:
        _target_: builtins.slice
        _args_:
        - '2012-12-15'
        - '2013-02-24'
    test:
      time:
        _target_: builtins.slice
        _args_:
        - '2012-10-01'
        - '2012-12-20'
  xrds_kw:
    patch_dims:
      time: 15
      lat: 200
      lon: 200
    strides:
      time: 1
      lat: 200
      lon: 200
    domain_limits: ${domain.test}
  dl_kw:
    batch_size: 4
    num_workers: 1
  aug_factor: 5
  aug_only: false
model:
  _target_: src.models.Lit4dVarNet
  opt_fn:
    _target_: src.utils.cosanneal_lr_adam
    _partial_: true
    lr: 0.001
    T_max: ${trainer.max_epochs}
  rec_weight:
    _target_: src.utils.get_triang_time_wei
    patch_dims: ${datamodule.xrds_kw.patch_dims}
    crop:
      time: 0
      lat: 5
      lon: 5
    offset: 0
  solver:
    _target_: src.models.GradSolver
    n_step: 12
    lr_grad: 100
    prior_cost:
      _target_: src.models.BilinAEPriorCost
      dim_in: ${datamodule.xrds_kw.patch_dims.time}
      dim_hidden: 64
      downsamp: 2
    obs_cost:
      _target_: src.models.BaseObsCost
    grad_mod:
      _target_: src.models.ConvLstmGradModel
      dim_in: ${datamodule.xrds_kw.patch_dims.time}
      dim_hidden: 96
entrypoints:
- _target_: pytorch_lightning.seed_everything
  seed: 333
- _target_: src.train.base_training
  trainer: ${trainer}
  lit_mod: ${model}
  dm: ${datamodule}
  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}
