data_workers: 10
name: Walrus-wella-delta-Isotr[Space-Adapt-]-AdamW-0.0002
automatic_setup: true
trainer:
  _target_: walrus.trainer.Trainer
  max_epoch: 200
  val_frequency: 10
  rollout_val_frequency: 10
  short_validation_length: 20
  max_rollout_steps: 200
  num_time_intervals: 5
  enable_amp: false
  loss_fn:
    _target_: the_well.benchmark.metrics.MAE
  formatter:
    _target_: hydra.utils.get_class
    path: walrus.data.well_to_multi_transformer.ChannelsFirstWithTimeFormatter
  revin:
    _target_: walrus.trainer.normalization_strat.SamplewiseRevNormalization
    _partial_: true
  prediction_type: delta
  grad_acc_steps: 4
  image_validation: true
  video_validation: true
  gradient_log_level: 0
  clip_gradient: 10
  log_interval: 200
  loss_multiplier: 100.0
  lr_scheduler_per_step: false
  skip_spectral_metrics: true
optimizer:
  _target_: torch.optim.AdamW
  weight_decay: 0.0001
  eps: 1.0e-10
  lr: 0.0002
lr_scheduler:
  _target_: walrus.optim.schedulers.InverseSqrtLinearWarmupSqrtCooldown
  warmup_epochs: 10
  cooldown_epochs: 10
  warmup_lr_factor: 0.1
  cooldown_lr_factor: 0.001
model:
  encoder:
    _partial_: true
    _target_: walrus.models.encoders.vstride_encoder.SpaceBagAdaptiveDVstrideEncoder
    learned_pad: true
    base_kernel_size1d:
    - - 4
      - 4
    base_kernel_size2d:
    - - 8
      - 4
    - - 8
      - 4
    base_kernel_size3d:
    - - 8
      - 4
    - - 8
      - 4
    - - 8
      - 4
    groups: 12
    kernel_scales_seq:
    - - 2
      - 2
    - - 4
      - 2
    - - 4
      - 4
    - - 8
      - 4
    variable_downsample: true
    variable_deterministic_ds: true
    activation:
      _partial_: true
      _target_: torch.nn.SiLU
  decoder:
    _partial_: true
    _target_: walrus.models.decoders.vstride_decoder.AdaptiveDVstrideDecoder
    learned_pad: true
    base_kernel_size1d:
    - - 4
      - 4
    base_kernel_size2d:
    - - 8
      - 4
    - - 8
      - 4
    base_kernel_size3d:
    - - 8
      - 4
    - - 8
      - 4
    - - 8
      - 4
    groups: 12
    activation:
      _partial_: true
      _target_: torch.nn.SiLU
  processor:
    space_mixing:
      _partial_: true
      _target_: walrus.models.spatial_blocks.full_attention.FullAttention
      num_heads: 16
      mlp_dim: null
    time_mixing:
      _partial_: true
      _target_: walrus.models.temporal_blocks.axial_time_attention.AxialTimeAttention
      num_heads: 16
      bias_type: rel
    channel_mixing:
      _partial_: true
      _target_: torch.nn.Identity
    _partial_: true
    _target_: walrus.models.spatiotemporal_blocks.space_time_split.SpaceTimeSplitBlock
  norm_layer:
    _partial_: true
    _target_: walrus.models.shared_utils.normalization.RMSGroupNorm
  _target_: walrus.models.IsotropicModel
  hidden_dim: 1408
  projection_dim: 48
  intermediate_dim: 352
  processor_blocks: 40
  drop_path: 0.05
  groups: 16
  max_d: 3
  static_axes: true
  weight_tied_axes: false
  causal_in_time: true
  include_d:
  - 2
  - 3
  override_dimensionality: 0
  jitter_patches: true
  gradient_checkpointing_freq: 2
  use_periodic_fixed_jitter: true
  input_field_drop: 0.0
data:
  field_index_map_override:
    closed_boundary: 0
    open_boundary: 1
    bias_correction: 2
    pressure: 3
    velocity_x: 4
    velocity_y: 5
    velocity_z: 6
    zeros_like_density: 7
    speed_of_sound: 8
    concentration: 9
    D_xx: 10
    D_xy: 11
    D_xz: 12
    D_yx: 13
    D_yy: 14
    D_yz: 15
    D_zx: 16
    D_zy: 17
    D_zz: 18
    E_xx: 19
    E_xy: 20
    E_xz: 21
    E_yx: 22
    E_yy: 23
    E_yz: 24
    E_zx: 25
    E_zy: 26
    E_zz: 27
    density: 28
    energy: 29
    velocity_r: 30
    velocity_theta: 31
    velocity_phi: 32
    momentum_x: 33
    momentum_y: 34
    momentum_z: 35
    pressure_re: 36
    pressure_im: 37
    mask: 38
    magnetic_field_x: 39
    magnetic_field_y: 40
    magnetic_field_z: 41
    A: 42
    B: 43
    height: 44
    internal_energy: 45
    temperature: 46
    electron_fraction: 47
    entropy: 48
    magnetic_field_log_r: 49
    magnetic_field_theta: 50
    magnetic_field_phi: 51
    velocity_log_r: 52
    buoyancy: 53
    tracer: 54
    log10_density: 55
    log10_temperature: 56
    c_zz: 57
    C_xx: 58
    C_xy: 59
    C_xz: 60
    C_yx: 61
    C_yy: 62
    C_yz: 63
    C_zx: 64
    C_zy: 65
    C_zz: 66
  transform:
    train:
      _target_: the_well.data.augmentation.RandomRotation90
      p: 1.0
  well_base_path: /mnt/gpuxl/polymathic/the_well/datasets/
  wandb_data_name: well_allmain_only
  module_parameters:
    _target_: walrus.data.MixedWellDataModule
    batch_size: 2
    n_steps_input: 6
    n_steps_output: 1
    min_dt_stride: 1
    max_dt_stride: 5
    max_samples: 2000
    well_dataset_info:
      active_matter:
        include_filters: []
        exclude_filters: []
      planetswe:
        include_filters: []
        exclude_filters: []
      acoustic_scattering_maze:
        include_filters: []
        exclude_filters: []
        field_transforms:
          density: torch.zeros_like
      acoustic_scattering_inclusions:
        include_filters: []
        exclude_filters: []
        field_transforms:
          density: torch.zeros_like
      acoustic_scattering_discontinuous:
        include_filters: []
        exclude_filters: []
        field_transforms:
          density: torch.zeros_like
      euler_multi_quadrants_openBC:
        include_filters: []
        exclude_filters: []
      euler_multi_quadrants_periodicBC:
        include_filters: []
        exclude_filters: []
      gray_scott_reaction_diffusion:
        include_filters: []
        exclude_filters: []
      rayleigh_benard:
        include_filters: []
        exclude_filters: []
      shear_flow:
        include_filters: []
        exclude_filters: []
      turbulent_radiative_layer_2D:
        include_filters: []
        exclude_filters: []
      helmholtz_staircase:
        include_filters: []
        exclude_filters: []
      viscoelastic_instability:
        include_filters: []
        exclude_filters: []
      supernova_explosion_128:
        include_filters: []
        exclude_filters: []
        step_downsample_factor: 0.5
        batch_downsample_factor: 0.5
        field_transforms:
          density: torch.log10
          temperature: torch.log10
      turbulence_gravity_cooling:
        include_filters: []
        exclude_filters: []
        step_downsample_factor: 0.5
        batch_downsample_factor: 0.5
        field_transforms:
          density: torch.log10
          temperature: torch.log10
      turbulent_radiative_layer_3D:
        include_filters: []
        exclude_filters: []
        step_downsample_factor: 0.5
        batch_downsample_factor: 0.5
        field_transforms:
          density: torch.log10
          temperature: torch.log10
      MHD_64:
        include_filters: []
        exclude_filters: []
        step_downsample_factor: 0.5
        batch_downsample_factor: 0.5
      rayleigh_taylor_instability:
        include_filters: []
        exclude_filters: []
        step_downsample_factor: 0.5
        batch_downsample_factor: 0.5
      flowbench_FPO_NS_2D_512x128_harmonics:
        include_filters: []
        exclude_filters: []
        path: /mnt/gpuxl/polymathic/WellFormattedExternalData/flowbench/flowbench_FPO_NS_2D_512x128_harmonics
auto_resume: true
folder_override: ''
checkpoint_override: ''
config_override: 
validation_mode: false
frozen_components:
- model
distribution:
  distribution_type: hsdp
  local_size: 4
logger:
  wandb: true
  wandb_project_name: walrus_Training_Attempts
checkpoint:
  _target_: walrus.trainer.checkpoints.CheckPointer
  save_dir: /mnt/home/polymathic/ceph/walrus_logging/runs/Walrus_ft_major_v2-wella-delta-Isotr[Space-Adapt-]-AdamW-0.0002/0/checkpoints
  load_checkpoint_path: null
  coalesced_checkpoint_path: null
  save_best: true
  checkpoint_frequency: 20
  align_fields: true
  load_chkpt_after_finetuning_expansion: false
finetuning_mods: {}
experiment_dir: /mnt/home/polymathic/ceph/walrus_logging/runs