Consider pulling real data from cloud in tutorials with zarr

[nicholasloveday]

Some of the tutorials use real data. These don't work as well in binder as you need to download the data and save it to disk.

An alternative approach would be to pull the data from the cloud into memory.

E.g.,

hres = xr.open_zarr('gs://weatherbench2/datasets/hres/2016-2022-0012-1440x721.zarr')
hres["2m_temperature"].sel(time="2020-01-01T00:00:00", prediction_timedelta=pd.Timedelta("1 days")).plot()

takes 1.5 seconds to pull the ECMWF forecast down from the cloud and plot it.

The dependencies that would need to be added for the tutorials are zarr and gcsfs

There is also reanalysis data and data driven models that can be pulled down from the cloud (see https://weatherbench2.readthedocs.io/en/latest/data-guide.html).

You can get data on the same grid so it makes verification with scores super easy!

Something to discuss

[nicholasloveday]

Okay, it looks like it's slower (~20s) to establish the initial connection with the Google storage and then pulling data down is fast

[nicholasloveday]

Here's an example of verifying GraphCast against ERA5 with scores

from scores.continuous import mse
import xarray as xr
import pandas as pd

graphcast = xr.open_zarr(
    "gs://weatherbench2/datasets/graphcast/2020/date_range_2019-11-16_2021-02-01_12_hours_derived.zarr"
)
fcst = graphcast["2m_temperature"].sel(
    time="2020-01-01T00:00:00", prediction_timedelta=pd.Timedelta("5 days")
)
era5 = xr.open_zarr("gs://weatherbench2/datasets/era5-forecasts/2020-1440x721.zarr")
obs = era5["2m_temperature"].sel(
    time="2020-01-01T00:00:00", prediction_timedelta=pd.Timedelta("0 days")
)
fcst = fcst.compute()
obs = obs.compute()
obs = obs.rename({"latitude": "lat", "longitude": "lon"})
result = mse(fcst, obs, preserve_dims="all")
result.plot(vmax=100)

which also takes ~1.5 seconds after an initial connection to the cloud storage

[tennlee]

If that works well in binder, that's good to know. I'd like some testing to be done before we proceed, and it might be a few weeks before I could do this myself. I'd be happy to see a new notebook created on a branch which we can then develop and test. It might be nice to have an ML-focused notebook which goes into some new areas, possibly looking at evaluation more than the use of individual scores. Thanks very much for putting your example together.

[nicholasloveday]

Yes - I agree that we need to test this in binder. It worked well on my laptop, so it may be an improvement for people running it locally on their computer.

A ML focused notebook sounds like a great idea. A scores + weatherbench2 tutorial would be quite nice.

[tennlee]

I've started a branch for this on my fork, based on your sample code. I won't be able to do this very quickly, so if you want to do something more quickly, feel free. It required the packages "zarr" and "gcsfs" to be installed. Zarr is a common format to want, so that's fine to add to the tutorial requirements, and I think there will be enough interest in these datasets to justify adding gcsfs to the tutorial requirements also. Another option would be to simply document those requirements in the notebook itself. But if the goal is to have this work nicely in binder, it's probably more reliable to add them to the requirements. If you'd like to, I'm happy to add you to my fork as well and you can push directly to the feature branch if you want to.

[nikeethr]

Good idea to use zarr for tutorials with low disk space, but somewhat reasonable RAM. However, it should be limited to that specific use case.

It may be tempting to use zarr "in general" but I'd like to raise some comments especially citing what you observed.

Okay, it looks like it's slower (~20s) to establish the initial connection with the Google storage and then pulling data down is fast

This is the crux of the problem and probably fine for a notebook. But for an actual operational system there are a few issues:

The download of an initial metadata structure is a likely reason why you see that initial delay, which based on the chunking strategy, compression, and the extents & dimensions required, may or may not be big. Note that zarr will have to be used "as-is" i.e. any re-chunking will also probably incur speed penalties as it is non-trivial to re-chunk without downloading additional data - though if the rechunking happens on server side (i.e. cloud) and is distributed it may be justifiable (though I'd be pleasantly surprised if this were the case and be done for free, because it'd either result in a large amount of compute or a large amount of data bloat or both).

It then follows that zarr makes more sense working cloud-native (or native to whichever premises it is stored in and linked via ethernet) - there's less speed hinderance/rate-limiting that way. e.g. if this were AWS, launching a lambda function or a ECS job in the same region as the storage to retrieve the data into a temporary store before downloading it directly will likely be faster. Unless of course as stated above, there is some process in the upstream server that already does (again doubt it).

This naturally should prompt the question, "why use zarr in a native environment when I can access netcdf directly (and with less storage overhead)?". That's a valid question, and in fact a crucial one that can invalidate the purpose of zarr. My opinion on the purpose of `zarr is that it isn't for speed or for efficiency, especially not in production systems - it's for convenience or ease of access via compatible protocols that also support fsspec.

Notebooks, experiments and prototyping are imo the ideal use-cases. This is a long winded way of saying that zarr can and probably should be added to notebooks, but not to the scores package itself without proper benchmarking (or disabling it by default/making it opt-in).

[nicholasloveday]

The download of an initial metadata structure is a likely reason why you see that initial delay, which based on the chunking strategy, compression, and the extents & dimensions required, may or may not be big.

Thanks for the explanation.

@nikeethr I used zarr data in the twCRPS tutorial https://scores.readthedocs.io/en/stable/tutorials/Threshold_Weighted_CRPS_for_Ensembles.html.

It seems to work pretty well as the user doesn't have to store the data on disk.

[nikeethr]

Yea that's one of the ideal use-cases. I also see that zarr has only been added to "tutorial" in pyproject.toml - which in line with what I was thinking - so all good :).

The important part is probably that the select and compute/load are done as early as possible which is also evident in the tutorial.

[nicholasloveday]

Another thing that I have found Zarr useful for is if I only want to pull down data from the cloud for one forecast parameter rather than a netCDF that contains every forecast parameter and is many, many times bigger (only to have to throw away the data from all the other parameters)

[tennlee]

I wonder if there would be value in creating some areas in the general documentation for this kind of useful information. https://scores.readthedocs.io/en/stable/data.html might be one place that this particular topic could go. I'd be happy to see a PR to that effect. In theory, it seems to me that local disk access "should" be faster with NetCDF than Zarr, but in practise, I've found zarr can be more efficient even on local disk. Presumably, that's because of how the actual files happen to be zipped/chunked on disk, as there are multiple compression and chunking options. So on-disk performance may be deeper than choosing a file format.

[nikeethr]

I don’t think thats totally surprising that you've noticed zarr seem faster on a local disk - that is my experience when using it in some applications in the past when zarr was still very experimental. Due to my “too good to be true” inkling - since Zarr was relatively barebones still - I benchmarked it back then and found there actually isn’t any special reason for this.

On the other hand, there are several reasons the default usage of xarray was giving resulting in an unfair comparison. There are actually lot of variables that come into play but I won’t get into detail here. Though as you mention how the file system is configured is definitely one of them. In fact between that early benchmarking and now, I’ve used other methods of accessing data including a different language - Rust, or bypassing xarray altogether (using netcd4 Python bindings directly) and found despite the data being unchanged, found operating closer to the underlying netcdf C library seemed to perform much better. This isn’t to discredit xarray - more so that its default loader settings were probably not optimal for my use case and I didn’t think of trying to optimize it.

Again for notebooks it’s a no brainer since the typical use case is to access smaller slices of data for analysis. So just plopping Zarr as the engine or using it directly might be fine or even recommended.

Operational systems however, especially when processing terabytes of data definitely need a controlled benchmark using the scores package prior to any recommendations. Which is actually not very trivial without some sysadmin knowledge of the underlying system e.g. HPC/cloud. A basic benchmarking and be done on a users own storage device e.g. SSD/HDD to start with - though it probably still wouldn’t make it sufficiently trivial. Noting that there are dimensions (figuratively and literally) to this analysis e.g. depending on how much data is actually being processed/requested for each operation.

If Zarr is orders of magnitude faster (e.g. not just a small 5% - or whatever is deemed insignificant - speedup that can be attributed for variables unaccounted for when setting up the experiment) without any obvious tradeoffs then that might be sufficient evidence.

And then there’s also the potential of NcZarr which I haven’t read up completely but could be used to retain a more coupled structure to the underlying netcdf object that may mitigate some inefficiencies

EDIT: wrote this on my phone so a combination of autocorrect and fat thumb seemed to have obfuscated some of the grammar. I’ll rewrite it more concisely when I get a chance but hopefully the essence is conveyed.