[Bug]: WellDataModule normalization not applied to val and test data

[radka-j]

Describe the issue:

Hi!

I noticed that in the WellDataModule the use_normalization and normalization_type parameters are only passed to WellDataset when creating the train data but not when creating the validation or test datasets. This seems to have been removed in this commit.

Is this a bug or is this handled somewhere else that I might have missed?

Code to reproduce the issue:

See implementation of the `WellDataModule` [here](https://github.com/PolymathicAI/the_well/blob/master/the_well/data/datamodule.py). Specifically the `__init__` method where `self.train_dataset`, `self.val_dataset`, `self.rollout_val_dataset` and `self.test_dataset` are created.

Version

Latest available on GitHub

Environment

n/a

Context for the issue:

n/a

[mikemccabe210]

Hi Radka,

Thanks for bringing this to our attention. Looking at it, I think it's probably a pattern that could use significant improvement for generalizability to other code bases, but I don't think it's a bug within the provided training/eval code itself. It looks like the commit where this was merged changed the training loop to use the normalization from the train dataset object within the model rollout loop which is unnecessarily confusing, but I think should still be accurate.

If you're running into this using the datamodule outside of the Well code, we can prioritize it for a faster fix. Otherwise, we'll just note it down as a needed improvement for the next release. Let us know!

Best,
Mike

EDIT: This is pretty messy, so I'm just going to hotfix it so behavior in the datamodule is consistent. Should be in within a day or two. We'll revisit the overall pattern in a later release.

[radka-j]

Thank you for pointing me to the right part of the codebase! I was trying to figure out if the normalisation happens elsewhere but couldn't find it easily so thought I'd just check here.

[maxxxzdn]

Hi folks, thanks for addressing the issue!

I just trained an FNO on active_matter with the default hyperparameters, and it scored 0.15 VRMSE pretty comfortably after just 1h of training. Therefore, I wonder, am I doing something wrong, or are the values reported in Table 2 wrong?

[mikemccabe210]

Hi @maxxxzdn - the repo has had a number of changes and bug fixes since it was originally released so the table is outdated and what you're reporting is consistent with what we're seeing in current runs. We have a team member gradually re-running things on the current version of the repo to account for the changes/fixes and we're planning to update the issue (#49), HF, and the table on arxiv once he's finished to get everything in final sync for new users.

In general, now that the data has been out long enough to see real use, even once those are fixed, I'd mostly only recommend the baselines for sanity checking. The original baselines were run in an intentionally naive time-boxed setting due to having to budget for HP tuning across so many datasets, and most of the work we've seen using the data since then has usually employed much higher tuning/training budgets for the Well datasets they've used.