10 dataloader

[rogerkuou]

close #10

Added dataset.py module handling patching SST dataset.
Updated example notebook.
added test

[rogerkuou]

Hi @SarahAlidoost This PR is ready for review. Can you have a look when you have time?

The only thing is, after updating the example notebook with the patching, I could not reproduce the prediction results. I could not find the problem within 15 mins so I found it might be good to have a look together.

[SarahAlidoost]

@rogerkuou thanks for the implementation 👍 . I was able to run the notebook and reproduce the prediction issue. While I am looking into it, perhaps you can have a look at other issues that I commented here.

• The lsm_mask from era5 doesn't have lat/lon coordinates name, but latitude/longitude. This should be either fixed in the example notebook or in the util function. Can you please fix it in this PR?

• Also, in the notebook, in SSTDataset(), you are using lsm_mask and not lsm_subset. We can fix this at the end.

I left a few suggestions on the code mainly to improve the readability and making the code generic. Also, I added a code block taking care of the month since we are going to have a data with more than one month (model supports this in Batch dimension only, see issue #23). But since torch DataLoader cannot handle different lengths of days in a month, we need to add a padding to shorter months, filled them with zeroes. I found that is possible passing collate_fn to Dataloader. We can create a custom function to support padding. Can you please add that function? We can implement this in another PR.

I am reviewing the notebook and will come with comments later. Please let me now if something is unclear.

[SarahAlidoost]

@rogerkuou Hi, I had a look at the notebook and found a few issues with the shapes of arrays and how the loss is calculated. With the changes that I suggested for the Dataset class, and these changes in the notebook, results should be correct. With this SSTDataset implementation, we can remove the function prepare_spatiotemporal_batch from st_encoder_decoder.py:

for epoch in range(101):
    for batch in dataloader:

        # Initialize gradients
        optimizer.zero_grad()

        # Get batch data
        daily_batch = batch["daily_mask_patch"]
        daily_mask = batch["daily_mask_patch"]
        monthly_target = batch["monthly_patch"]
        land_mask = batch["land_mask_patch"][0,...]  # same for all batches

        # Batch prediction
        pred = model(daily_batch, daily_mask, land_mask)
        # Compute loss
        ocean = ~land_mask
        ocean_broadcasted = ocean.unsqueeze(0)  # [1,H,W]

        loss = torch.nn.functional.l1_loss(
            pred*ocean, monthly_target*ocean, reduction='sum'
        )  # For SST there is only one channel
        loss = loss / ocean_broadcasted.sum() 
        loss.backward()
        optimizer.step()

later:

dataset_pred = SSTDataset(
    daily_ds=daily_subset["ts"],
    monthly_ds=monthly_subset["ts"],
    land_mask=lsm_subset["lsm"],
    patch_size=(daily_subset.sizes["lat"], daily_subset.sizes["lon"]),
)
dataloader_pred = DataLoader(
    dataset_pred,
    batch_size=len(dataset_pred),
    pin_memory=False,
)
full_batch = next(iter(dataloader_pred))

daily_batch = full_batch["daily_mask_patch"]
daily_mask = full_batch["daily_mask_patch"]
monthly_target = full_batch["monthly_patch"]
land_mask = full_batch["land_mask_patch"][0,...]  # same for all batches
print(daily_batch.shape, daily_mask.shape, monthly_target.shape, land_mask.shape)

model.eval()
with torch.no_grad():
    pred = model(daily_batch, daily_mask, land_mask)


from st_encoder_decoder import pred_to_numpy
monthly_prediction = pred_to_numpy(pred, land_mask=land_mask)

[rogerkuou]

Hi @SarahAlidoost, thanks for the review! I adpated all of your comments, updated the notebook and removed prepare_spatiotemporal_batch . Can you have another look? Thanks!

[rogerkuou]

Hi @SarahAlidoost, thanks for the review! I adpated all of your comments, updated the notebook and removed prepare_spatiotemporal_batch . Can you have another look? Thanks!

[SarahAlidoost]

@rogerkuou thanks for addressing the comments! looks very good! 👍 The only thing I see that the notebook doesnot include the outputs, can you please push the complete notebook? Also, this lsm_mask = lsm_mask.isel(time=0) is not needed in the notebook. After this change, let's merge this PR.

[rogerkuou]

@rogerkuou thanks for addressing the comments! looks very good! 👍 The only thing I see that the notebook doesnot include the outputs, can you please push the complete notebook? Also, this lsm_mask = lsm_mask.isel(time=0) is not needed in the notebook. After this change, let's merge this PR.

Added model output and prediction output to the example notebook. removed the mentioned line. I will merge it now.

[SarahAlidoost]

Added model output and prediction output to the example notebook. removed the mentioned line. I will merge it now.

@rogerkuou In your commit add output to notebook the notebook still doesn't include the outputs cells (e.g. the plots). It seems they are not pushed. It is okay, I will fix it in other PR.

[rogerkuou]

Hi @SarahAlidoost, my apologies! I misunderstood your comments. I thought by outputs you meant saving the trained model/prediction, but now I see that you meant the cell outputs. Actually I was not sure if this it's a good idea since the plots in notebooks can consume quite some volumes in GitHub history. So far we are only investigating with a small dataset. Maybe let's only save the outputs when we reach a more stable version?

[SarahAlidoost]

Hi @SarahAlidoost, my apologies! I misunderstood your comments. I thought by outputs you meant saving the trained model/prediction, but now I see that you meant the cell outputs.

No problem! We can fix it in another PR.

Actually I was not sure if this it's a good idea since the plots in notebooks can consume quite some volumes in GitHub history. So far we are only investigating with a small dataset. Maybe let's only save the outputs when we reach a more stable version?

The notebooks will be a part of documentation and they are also used in our update meetings. Another reason that github doesn't have a good interface to render the changes in the notebook between branches. Sometimes we want to compare the output with the main branch. So, let's keep the outputs. 🙏