Add periodicity to UNet using CircularPad2d

[ma595]

This PR aims to address the issue observed by @TomMelt, where a vertical 'seam' or fold is observed at the meridian line (lon = 0). We hadn't noticed this before as we were plotting data (lons 0-360).

This is illustrated here in the raw flux outputs of the network. Tom has already verified that the inputs are smooth.

This issue appears because the Unet does not enforce longitudinal periodicity. We can fix this by using the CircularPad2d function where convolution is done, i.e. in classes Conv_block, Upsample and Attention_block.

I also recommend a replication pad at the poles using nn.ReplicationPad2d. This will eliminate potential steep gradients there.

self.pad_layer = nn.Sequential(
            nn.CircularPad2d((padding, padding, 0, 0)),
            nn.ReplicationPad2d((0, 0, padding, padding))
        )

TODO:

• Zero pad or replication pad in latitude?
• Light testcase to ensure training will work.
• Add some docs to code for clarity

Closes #56

[TomMelt]

Thanks for raising this @ma595 . Would be good to see if we can make the output continuous across the 0 meridian line.

[amangupta2]

@ma595 Looks like a nice strainghtforward fix. Is it possible to add this during inference itself? Or do we need to retain the model with the ReplicationPad2d layer?

[ma595]

@ma595 Looks like a nice strainghtforward fix. Is it possible to add this during inference itself? Or do we need to retain the model with the ReplicationPad2d layer?

I think we'd need to retrain the model.

[TomMelt]

just a reminder, can you remove these lines completely. The default for Conv2d is padding=0. https://docs.pytorch.org/docs/stable/generated/torch.nn.Conv2d.html