Fix geo embeding

[SarahAlidoost]

closes #91

The main changes in this PR:

• fixing the calculation of geo embeding for each patch in dataset class
• removing the SpatialTransfoemr
• Other changes (minor) in the class st_encoder_decoder is mainly improving model performance.

Todo:

• re-run hourly notebook when levante is back
• update the scripts when Run training on the three years dataset #25 is merged --> done ✔️

[SarahAlidoost]

@meiertgrootes and @rogerkuou This PR fixes #91. Regarding the model performance see #93. Please let me know if something isnot clear.

[rogerkuou]

Hi @SarahAlidoost , I had a look at the changes and they all look good! I haven't executed notebooks and I am planning to do it when Levante is back online.

I just have one tiny comment on the notwbook. Feel free to ignore.

Another general comment is, is this PR also aim to rename the daily/hourly variables to input? If this is the case, there are still many instances in dataset.py and st_encoder_decode.py, like these lines, but it seems you also get this tracked in your thirs point in issue #93

[SarahAlidoost]

Hi @SarahAlidoost , I had a look at the changes and they all look good! I haven't executed notebooks and I am planning to do it when Levante is back online.

@rogerkuou Hi, Thanks!, you can run the daily example notebook locally, I made the spatial subset even smaller, so it takes a few seconds for training loop to complete. Can you please run the notebook and see whether results, and the workflow are all good?

I just have one tiny comment on the notwbook. Feel free to ignore.

Another general comment is, is this PR also aim to rename the daily/hourly variables to input? If this is the case, there are still many instances in dataset.py and st_encoder_decode.py, like these lines, but it seems you also get this tracked in your thirs point in issue #93

Correct! in this PR I didn't want to rename daily/hourly. Renaming variables is mentioned in issue #67 which I think it is not the priority now. But the third point in issue #93 is about the type of input (i.e. SST or anomaly), whether daily or hourly.

[meiertgrootes]

Looks good to me. Also waiting on Levante to run, but the adaptations make sense. Removing the spatial transformer does potentially loose some spatial correlation information, however, as we are temporally aggregating beforehand anyway that is likely not so critical. And it does alleviate the crop size dependence.

One very minor point on the tuning configuration.

Go ahead and merge when you see fit.

[meiertgrootes]

Please ignore if I just overlooked it, but I can't find the adapted keywords for the data_config in tuning.py

[SarahAlidoost]

Yes, that's right. I was waiting for #25 to be merged and then update the scripts, see my Todo list here.

[SarahAlidoost]

Looks good to me. Also waiting on Levante to run, but the adaptations make sense. Removing the spatial transformer does potentially loose some spatial correlation information, however, as we are temporally aggregating beforehand anyway that is likely not so critical. And it does alleviate the crop size dependence.

Yes, that's correct. I think training on the whole globe should help the model learn the spatial correlations. Once it's trained, we can check this and maybe add a spatial mixing layer in the decoder if needed 🤔 I'll add it to issue #93 so we don't forget.

One very minor point on the tuning configuration.

Go ahead and merge when you see fit.

[rogerkuou]

Hi @SarahAlidoost , I had a look at the changes and they all look good! I haven't executed notebooks and I am planning to do it when Levante is back online.

@rogerkuou Hi, Thanks!, you can run the daily example notebook locally, I made the spatial subset even smaller, so it takes a few seconds for training loop to complete. Can you please run the notebook and see whether results, and the workflow are all good?

I just have one tiny comment on the notwbook. Feel free to ignore.
Another general comment is, is this PR also aim to rename the daily/hourly variables to input? If this is the case, there are still many instances in dataset.py and st_encoder_decode.py, like these lines, but it seems you also get this tracked in your thirs point in issue #93

Correct! in this PR I didn't want to rename daily/hourly. Renaming variables is mentioned in issue #67 which I think it is not the priority now. But the third point in issue #93 is about the type of input (i.e. SST or anomaly), whether daily or hourly.

Hi @SarahAlidoost , I did a local run with the hourly training example notebook, see this GH gist.

My observations:

• The training runs smoothly, with a similar result as yours. (see the end errors and histogram)
• However I only managed to run ~1/8 of your AoI. it runs much slower on my local comparing with your time profiling results in the example. I increased the training workers to 8 and data loading workers to 4. This took ~1hr.
• Batch size 10 is about to the max my me can hold.
• I used month 202101 for both training and validation. And 202202 for testing. since I only got two months file on my local

With this I think this PR is ready for another training run. The sloweness issue is likely to be caused by my local infra.