Fix blocky effect

[SarahAlidoost]

closes #17

This PR adds :

• a documentation on how to select the patch sizes, read this carefully when setting up training.

This PR fixes main things as:

• a bug in the optimization loop in the notebook, how loss is calculated
• a bug in the optimization loop in the notebook, how land mask is applying
• a bug in reshaping arrays inside the model
• a bug in the notebook in inference (prediction cell).

This PR also improves:

• the spatial refinemnet
• the dataset class to cover the whole input data
• disable overlapping data batches created by dataset because it affects the loss

Results:
Note: As can be seen in the final plots, there is not spatial artifacts. I also run the notebook on residuals in this branch. The results also are without artifacts.

Note: The notebooks in this PR can be run on local systems. If the input size or data patch size are changed, the model patch size and overlap should be adjusted accordingly.

[SarahAlidoost]

@meiertgrootes please also see the top comment when reviewing.

[meiertgrootes]

Moving to lat,lon based encoding this would be less relevant, but how does this match to the total grid size given theglobal input data?

[SarahAlidoost]

Here the max_H and max_W in SpatialPositionalEncoding2D is the maximum number of patches in the height and width directions. For example, if the image is 1024×1024 pixels with a patch size of (1, 4, 4), we have 256 patches in H and W. The max_H and max_W arguments can be changed by users if needed. However, about global data, we don't usually give the large image to the model, but we subset it and embed the subsets in the batch dimensions.

[meiertgrootes]

Ok, thanks. So max_H, max_W are basically the total input dimensions quantised by the patch size. That is fine, of course, and I also understand that we aren't passing the full global image in one go. My point was, that as it stands, the 2d positional encodings are pre-calculated on the basis of these sizes, i.e. the sine and cosine values derive from the max_H and max_W values. Should one change input dimensions/resolution this may be a challenge to generalisation. Other approaches also have their own difficulties, this is just sth. to remain aware of.

[meiertgrootes]

The batch handling look a lot cleaner and the carry over og lat long information will elp with global spatial embedding. This may make the chocie of max size (see comment) less critical.