🐛[BUG]: Atlas checkerboard artefacts

[lhm-incom]

Version

0.12.1

On which installation method(s) does this occur?

Atlas.

Describe the issue

Thanks for sharing a wealth of weather modelling features!

Issue

I'm playing around with the Atlas model. Upon inspection of it's predictions, I observe systematic checkerboard patterns.
It's often over land (see 2t example below)

For some channels also apparent over ocean (see tp example below).

Potential causes

The artefacts don't seem present in the latent state (which is also of lower resolution). My best guess is that the decode operation self.autoencoders[0] (NattenCombinDiT) produces the artefacts.

• Perhaps the "autoencoder_0/data_processor_best.mdlus" weights from HuggingFace weren't correct.
• Perhaps the newest natten version (I used "0.21.5") introduced an issue. (same results with 0.21.1).
• Perhaps it's an inherent model artefact.

[pzharrington]

Hi @lhm-incom , thanks for reporting this. What forecast lead times are the above plots made for? I believe the model has not been extensively tested beyond 15 days, as it is focused on medium-range weather, so it may eventually accumulate errors enough to lead to artifacts like this late in a rollout. It would also be helpful if you can share the initialization date(s) for these plots.

[lhm-incom]

The above plots are from the first prediction step (6h lead time) using ECMWF analysis initial conditions from 2026-02-10T00:00:00.

[pzharrington]

The above plots are from the first prediction step (6h lead time) using ECMWF analysis initial conditions

I'll note that using analysis state ICs would be a sight distribution shift for the model compared to ERA5, though I'd expect it to work reasonably well. Are you sure all the variables needed by the model are provided? Where are you getting the analysis initial condition for the SST, and how are you filling the values over landmasses? If they are not NaN, the model expects landmass SSTs to be set to 0 K.

[lhm-incom]

The earth2studio atlas impl handles NaNs internally (sets them to zero). I suspected the same issue early in debugging and tried with a few approaches setting NaNs in the SST channel to both kelvin zero, celcius zero (the paper just mentioned imputation with "zero") and SST nan_mean (which would be the normalized zero). These did not fix the checkerboard pattern.

I just rechecked the variables on my side: They match the list in earth2studio/models/px/atlas.py VARIABLES.

I think that it would be valuable to provide a usage example using for instance Arco ERA5 data (c.f. the related issue) with a prediction visualisation. That would make it easy to compare private implementation with a proven reference code.

[pzharrington]

Hmm, ok. Just for clarification, are you using the IFS data source to download your IC data, or something else? I have added the request for an explicit Atlas example here: #695

[lhm-incom]

Our team receives IFS data directly from ECMWF. The pipeline was validated for use with other models on prior occasions.

[swpowell]

I also get the checkerboard pattern when using GFS data source in the package. Pattern appears at f+006 and is worse in some variables than others. Tha pattern is obvious in t2m, for example, at f+006. Less obvious in z500 at f+006 but becomes noticeable to the point of making the model output unusable by 10 days.

Edit: Also, the t2m checkerboarding shows up at f+006 mostly over water and coastal areas only. Then it shows up over land later.

Edit 2: Adding image.

[pzharrington]

Thanks @swpowell for the additional report/info. The artifacts concentrating over water definitely suggests some problem or discrepancy with the land-sea mask. For reference, here is a quick example script using ARCO to compare against:

import numpy as np
import torch

from earth2studio.data.arco import ARCO
from earth2studio.data.utils import fetch_data
from earth2studio.models.px.atlas import Atlas

def main():

    time = np.array([np.datetime64("2020-01-02T00:00")])
    n_steps = 40
    device = torch.cuda.current_device()

    # Load model
    package = Atlas.load_default_package()
    model = Atlas.load_model(package).to(device)

    # Model input coordinates
    prognostic_ic = model.input_coords()
    lead_time = prognostic_ic["lead_time"]
    variable = prognostic_ic["variable"]

    # Fetch initial conditions from ARCO
    data = ARCO(cache=True, verbose=True)
    x, coords = fetch_data(
        source=data,
        time=time,
        variable=variable,
        lead_time=lead_time,
        device=device,
    )

    # Add a batch dim
    bs = 1
    x = x.unsqueeze(0).repeat(bs, 1, 1, 1, 1, 1)
    coords["batch"] = np.arange(bs)
    coords.move_to_end("batch", last=False)

    # Run model forward pass
    y, ycoords = model(x, coords)

if __name__ == "__main__":
    main()

[swpowell]

Thanks @pzharrington for the code snippet using ARCO data. I can confirm this works as expected and I don't get a checkerboard pattern in the output. I agree this looks like a data source-dependent problem with land/sea masking.

Edit: OK it's not land/sea masking exactly, and the likely solution requires non-trivial effort to line up SST data with the data source analysis. Atlas expects 'sst' as a variable. This downloads fine from the ARCO dataset, but loading the GFS dataset (and presumably some other datasets too) loads all zeroes into the input tensor for 'sst'. Currently GFS data is only loaded from s3://noaa-gfs-bdp-gfs/gfs.YYYYMDD/... which does not contain SST. Therefore, 'sst' isn't even in the GFS lexicon.

GFS also has SST data but it lives at s3://noaa-gfs-bdp-gfs/nsst.YYYYMMDD/ after 1 Jan 2023 and s3://noaa-gfs-bdp-gfs/sst.YYYYMMDD before that back to 11 Feb 2021. However, that data is only daily and lags atmospheric analysis by 1 day. For example, on 18 Feb, I can get GFS SST data for 17 Feb. Also, there is 1/12 deg and 1/2 deg SST data, so the 1/12 deg data would need to be degraded. I suspect SST data could also be used from a different source.

Bottom line is, for now, users probably need to modify the input tensor after grabbing GFS analysis but before inference to replace the zeroes with real SST data from some source. Updating earth2studio will require some consensus on where to get near real-time SST data from.

[pzharrington]

Yes, thanks for your investigation @swpowell , I can confirm what you're reporting about GFS sst being empty. I do see a warning from the data source at least when trying to download it. For a real-time initial condition for SST, I believe you should be able to pull from copernicusmarine but it likely will require some regridding as you mentioned. There is not (yet) an earth2studio data source for copernicusmarine though, it would have to be a manual data pipeline for that variable.