Skip to content

Add multi-fidelity/modality support#313

Draft
radka-j wants to merge 28 commits into
mainfrom
308-multi
Draft

Add multi-fidelity/modality support#313
radka-j wants to merge 28 commits into
mainfrom
308-multi

Conversation

@radka-j

@radka-j radka-j commented Apr 7, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Closes #308

We have 2 scripts in notebooks/ for the full workflow:

  • generate_reaction_diffusion.py
  • train_reaction_diffusion.py

QUESTIONS and TODOs:

  • at the moment we do end-to-end training on the encoder-processor-decoder but we might want to add an option to train the encoders independently, then freeze them and only train the rest of the model (attention + processor + decoder)

sgreenbury and others added 5 commits March 28, 2026 07:15
@sgreenbury
Refactor eval script: extract phase functions and fix Fabric handling
9084fd0 
- Break run_evaluation into focused phase functions (_run_test_metrics,
  _run_rollout_videos, _run_rollout_metrics, _run_rollout_per_timestep_metrics,
  _write_final_csvs) to reduce complexity and improve readability
- Extract _resolve_metric_fns to eliminate duplicated metric resolution logic
  between test and rollout metric setup
- Unwrap Fabric-wrapped model once and pass unwrapped model to functions that
  access model internals (parameter counting, benchmarks) instead of relying
  on Fabric proxy attribute access with pyright ignores
- Add depth limit to _unwrap_module to guard against circular wrapper chains
- Extract _setup_fabric, _load_model_checkpoint, and _resolve_eval_paths
  helpers from run_evaluation
@sgreenbury
Pass model instead of unwrapped_model
287cd97
@sgreenbury @radka-j @ContiPaolo
Begin initial multidataset
0ce61d1 
Co-authored-by: Radka Jersakova <r.jersakova@gmail.com>
Co-authored-by: Paolo Conti <ContiPaolo@users.noreply.github.com>
@sgreenbury @ContiPaolo @radka-j
Revise mask type and begin init
b8c0520 
Co-authored-by: Paolo Conti <ContiPaolo@users.noreply.github.com>
Co-authored-by: Radka Jersakova <r.jersakova@gmail.com>
@sgreenbury @ContiPaolo @radka-j
Begin implementing collation for ListSample
9d11ee4 
Co-authored-by: Paolo Conti <ContiPaolo@users.noreply.github.com>
Co-authored-by: Radka Jersakova <r.jersakova@gmail.com>
@radka-j radka-j marked this pull request as ready for review April 7, 2026 09:05
@radka-j radka-j marked this pull request as draft April 7, 2026 09:05
radka-j and others added 19 commits April 7, 2026 11:04
@radka-j @ContiPaolo
add notes
e4debf9 
Co-authored-by: Paolo Conti <ContiPaolo@users.noreply.github.com>
@radka-j
finish implementation of collate_list_batches
5ca959e
@ContiPaolo
Implement sequential and combinatorial multi-fidelity masks
21bcbe3 
Co-authored-by: Radka Jersakova r.jersakova@gmail.com
@radka-j
add initial MultiEncoder
90bba88
@ContiPaolo
add multifidelity transformer as attention mixer
49c75a3
@ContiPaolo
refactor(models): remove AttentionMixer pooling for shape parity and …
c559cc5 
…generalize MultiEncoder dimensional projections
@radka-j
finish implementing MultiSpatioTemporalDataset
dba0c08
@ContiPaolo
fixed consistency between encoder dimn and transformer dim
bde437a
@ContiPaolo
fix masking shapes for attention mixer
c891595
@ContiPaolo
fix multifidelity transformer logic to handle high-fidelity data
a5c8443
@radka-j
clean up MultiSpatioTemporalDataset implementation
549b009
@radka-j
mv ListSample and ListBatch to types
e85e1d6
@radka-j
add options to specify normalization per dataset
81ac888
@radka-j
add TODO
58f1d65
@radka-j
add tests for MultiSpatioTemporalDataset
9e85d6f
@ContiPaolo
add scripts to generate data and to train multifidelity trasnformer
59d1016
@ContiPaolo
fixed bug: now input_projection is applied to latent dimension and no…
9cbd3fd 
…t over channels
@ContiPaolo
fixed input projections
8ae4871
@radka-j
add scripts to generate data + train multifidelity transformer
e165c85
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Implement support for multi-fidelity and multi-modality

3 participants

@radka-j @sgreenbury @ContiPaolo
⚡