Merge branch 'main' into normalization

Author: radka-j

configs/config.yaml configs/autoencoder.yamlconfigs/config.yaml renamed to configs/autoencoder.yaml

@@ -1,12 +1,16 @@
 defaults:
   - data: reaction_diffusion
+  # Or set to TheWell dataset:
+  # - data: the_well
   - model: ae
   - trainer: default
   - logging: wandb
   - _self_
 
 seed: 42
 experiment_name: autoencoder
+data:
+  autoencoder_mode: true
 output:
   checkpoint_name: autoencoder.ckpt
   save_config: true

configs/data/reaction_diffusion.yaml

@@ -1,5 +1,5 @@
-data_path: null
-use_simulator: true
+data_path: ${oc.env:AUTOCAST_DATASETS,./datasets}/reaction_diffusion
+use_simulator: false
 split:
   n_train: 4
   n_valid: 2

configs/data/the_well.yaml

@@ -0,0 +1,36 @@
+# TheWell dataset configuration
+# Configure TheWellDataModule for loading datasets from The Well benchmark
+_target_: autocast.data.datamodule.TheWellDataModule
+
+# Dataset name from The Well.
+# Set to `null` here so it must be provided or overridden from the CLI:
+# Example override: `python script.py data=the_well data.well_dataset_name=turbulent_radiative_layer_2D`
+well_dataset_name: turbulent_radiative_layer_2D
+
+# Path to The Well datasets. Read from `AUTOCAST_DATASETS` environment variable
+# Example: export AUTOCAST_DATASETS=/path/to/autocast/datasets
+well_base_path: ${oc.env:AUTOCAST_DATASETS,./datasets}
+
+# Temporal configuration
+n_steps_input: 1
+n_steps_output: 4
+min_dt_stride: 1
+max_dt_stride: 1
+
+# DataLoader configuration
+batch_size: 16
+
+# Normalization
+use_normalization: true
+normalization_path: ../stats.yaml  # defaults to ../stats.yaml relative to dataset
+
+# Mode configuration - can be overridden in autoencoder.yaml
+autoencoder_mode: false
+
+# Additional TheWell parameters
+max_rollout_steps: 100
+flatten_tensors: true
+cache_small: true
+max_cache_size: 1.0e9
+return_grid: true
+boundary_return_type: padding

configs/decoder/identity.yaml

@@ -0,0 +1 @@
+_target_: autocast.decoders.identity.IdentityDecoder

configs/encoder/identity.yaml

@@ -0,0 +1 @@
+_target_: autocast.encoders.identity.IdentityEncoder

configs/model/encoder_processor_decoder.yaml

@@ -0,0 +1,12 @@
+defaults:
+  - /encoder: identity
+  - /decoder: identity
+  - /processor: flow_matching
+  - _self_
+
+learning_rate: 0.001
+train_processor_only: false
+teacher_forcing_ratio: 0.5
+max_rollout_steps: 10
+loss_func:
+  _target_: torch.nn.MSELoss

configs/processor.yaml

@@ -1,8 +1,6 @@
 defaults:
   - data: reaction_diffusion
-  - encoder: permute_concat
-  - decoder: channels_last
-  - processor: fno
+  - model: encoder_processor_decoder
   - trainer: default
   - logging: wandb
   - _self_
@@ -14,18 +12,12 @@ output:
   save_config: true
 
 training:
-  n_steps_input: 4
+  n_steps_input: 1
   n_steps_output: 4
-  stride: 4
+  stride: null
   autoencoder_checkpoint: null
   freeze_autoencoder: false
 
-encoder_processor_decoder:
-  learning_rate: 0.001
-  train_processor_only: false
-  loss_func:
-    _target_: torch.nn.MSELoss
-
 hydra:
   run:
     dir: outputs/${experiment_name}/${now:%Y-%m-%d_%H-%M-%S}

configs/processor/diffusion.yaml

@@ -0,0 +1,22 @@
+_target_: autocast.processors.diffusion.DiffusionProcessor
+backbone:
+  _target_: autocast.nn.unet.TemporalUNetBackbone
+  in_channels: null
+  out_channels: null
+  cond_channels: null
+  mod_features: 256
+  hid_channels: [32, 64, 128]
+  hid_blocks: [2, 2, 2]
+  spatial: 2
+  periodic: false
+schedule:
+  _target_: azula.noise.VPSchedule
+denoiser_type: karras
+teacher_forcing_ratio: 0.0
+stride: ${training.stride}
+max_rollout_steps: ${training.n_steps_output}
+learning_rate: 0.0001
+n_steps_output: null
+n_channels_out: null
+sampler_steps: 50
+sampler: euler

configs/processor/flow_matching.yaml

@@ -0,0 +1,18 @@
+_target_: autocast.processors.flow_matching.FlowMatchingProcessor
+stride: ${training.stride}
+teacher_forcing_ratio: 0.0
+max_rollout_steps: ${training.n_steps_output}
+learning_rate: 0.0001
+flow_ode_steps: 4
+n_steps_output: null
+n_channels_out: null
+backbone:
+  _target_: autocast.nn.unet.TemporalUNetBackbone
+  in_channels: null
+  out_channels: null
+  cond_channels: null
+  mod_features: 256
+  hid_channels: [32, 64, 128]
+  hid_blocks: [2, 2, 2]
+  spatial: 2
+  periodic: false