Climatology is incorrect shape for forecasting with history

[jasonjewik]

Describe the bug The climatology dimension is incorrect.

To Reproduce Run the script here: https://gist.github.com/jasonjewik/c339325e4ae33c85e4cecc1356fdae38.

Expected behavior Instead of getting a 2D tensor of shape [32, 64], I'm getting a 3D tensor of shape [3, 32, 64].

Screenshots

Environment

• OS: Ubuntu 22.04
• Python version: 3.10
• Environment: only dependency is climate-learn

Additional context Even if I set history to 1, I still get a shape error.

[jasonjewik]

https://github.com/aditya-grover/climate-learn/blob/4b3d58821f889921014d7af2a1e49140b0f756a5/src/climate_learn/data/dataset/map_dataset.py#L102

If I change these lines to be

self.climatology: Data = {
    k: torch.mean(stacked_inp_data[k], dim=(0,1)) for k in stacked_inp_data
}

then the error goes away. This fixes the shape, but I'm not sure if it's correct. If I take the mean over dim 0 and dim 1, that means it's the mean over the entire data interval and over the per-sample sequences. We should probably just do it over the data interval. @prakhar6sharma, I'd appreciate your input here to double-check this.

[prakhar6sharma]

It's returning me a 2D tensor of shape [32, 64].

The climatology is taken for output data and not input data. Also, irrespective of history the each value in stacked_out_data should be a 3D tensor of size [n, 32, 64] (where n is the number of examples). Thus, torch.mean over the first dimension would result in a 2D tensor.

[jasonjewik]

Hmm okay. I forgot to mention this in the additional context section, but this behavior happens on my fork of the repo for a particular experiment I'm working on. It might be caused by some other change that I made which is unrelated to the main repo. Leaving this issue open for now until I can determine for sure.

[prakhar6sharma]

No problem. Let's keep the issue open.

Does the other experiments work for you? My best guess would be that it is related to something that you might have changed.