Asma/fix source writing/head

[sbAsma]

This PR is related to the issue " source and target not written the same way in BERT mode #38". The data is now accessed the same way. However, I keep having the following test failure:

0: =================================== FAILURES ===================================
0: ________________________________ test_datetime _________________________________
0: 
0: field = 'temperature', model_id = '1rj15n6y', BERT = True, epoch = 0
0: 
0:     def test_datetime(field, model_id, BERT, epoch = 0):
0:     
0:         """
0:         Check against ERA5 timestamps.
0:         Loop over all levels individually. 50 random samples for each level.
0:         """
0:     
0:         store = zarr.ZipStore(atmorep_target().format(model_id, model_id, str(epoch).zfill(5)))
0:         atmorep = zarr.group(store)
0:     
0:         nsamples = min(len(atmorep[field]), 50)
0:         samples = rnd.sample(range(len(atmorep[field])), nsamples)
0:         levels = [int(f.split("=")[1]) for f in atmorep[f"{field}/sample=00000"]] if BERT else atmorep[f"{field}/sample=00000"].ml[:]
0:     
0:         get_data = get_BERT if BERT else get_forecast
0:     
0:         for level in levels:
0:             #TODO: make it more elegant
0:             level_idx = level if BERT else np.where(levels == level)[0].tolist()[0]
0:     
0:             for s in samples:
0:                 data, datetime, lats, lons = get_data(atmorep, field, s, level_idx)
0:                 year, month = datetime.year, str(datetime.month).zfill(2)
0:     
0:                 era5_path = era5_fname().format(field, level, field, year, month, level)
0:                 if not os.path.isfile(era5_path):
0:                     warnings.warn(UserWarning((f"Timestamp {datetime} not found in ERA5. Skipping")))
0:                     continue
0:                 era5 = xr.open_dataset(era5_path, engine = "cfgrib")[grib_index(field)].sel(time = datetime, latitude = lats, longitude = lons)
0:     
0:                 #assert (data[0] == era5.values[0]).all(), "Mismatch between ERA5 and AtmoRep Timestamps"
0: >               assert np.isclose(data[0], era5.values[0],rtol=1e-04, atol=1e-07).all(), "Mismatch between ERA5 and AtmoRep Timestamps"
0: E               AssertionError: Mismatch between ERA5 and AtmoRep Timestamps
0: E               assert False
0: E                +  where False = <built-in method all of numpy.ndarray object at 0x147d686952f0>()
0: E                +    where <built-in method all of numpy.ndarray object at 0x147d686952f0> = array([False, False, False, False, False, False, False, False, False,\n       False, False, False, False, False, False, False, False, False,\n       False, False, False, False, False, False, False, False, False]).all
0: E                +      where array([False, False, False, False, False, False, False, False, False,\n       False, False, False, False, False, False, False, False, False,\n       False, False, False, False, False, False, False, False, False]) = <function isclose at 0x147e9c4e20b0>(array([267.93634, 267.91095, 267.89044, 267.92853, 267.97052, 268.0135 ,\n       268.02228, 267.88556, 267.6668 , 267.3...6.00275, 266.50568, 266.6678 , 266.28107, 265.575  , 265.5174 ,\n       265.79083, 265.90997, 266.15704], dtype=float32), array([265.21954, 266.00275, 266.50568, 266.6678 , 266.28107, 265.575  ,\n       265.5174 , 265.79083, 265.90997, 266.1...7.6668 , 267.33966, 266.92853, 266.47345, 266.09454, 265.84747,\n       265.57306, 265.21564, 264.91193], dtype=float32), rtol=0.0001, atol=1e-07)
0: E                +        where <function isclose at 0x147e9c4e20b0> = np.isclose
0: 
0: /p/home/jusers/semcheddine1/juwels/2_atmorep_dev/atmorep/atmorep/tests/validation_test.py:68: AssertionError

Might be related to the dataset not being loaded properly (I am unsure if this is what you meant @iluise during our first temporal interpolation meeting)

[iluise]

Hi,

Suggestion: please use git add <filename> instead of git add * so you can commit only the files you actually need.

About the issue on testing, it's weird that ERA5 and target now do not match anymore as you seem to have touched only source. Unless there's a problem with copying and memory in python. Are these the only changes wrt the develop branch or you introduced some other modifications?

[sbAsma]

Hi Ilaria,

Your suggested has been noted 👍

Regarding the issue on testing, I only changed source, yes. I am also reading the data from /p/scratch/atmo-rep/data/era5/new_structure/{}/ml{}/era5_{}_y{}_m{}_ml{}.grib during the test, if this helps.

[sbAsma]

Additional info here:
I did do individual commits, one for each file changed. However, many of these changes are related to the data partitions at JSC, or to my section of scratch, and therefore unnecessary to merge. I hope this helps. Of course, and from now on, I will only commit the files that are related to the issue being solved.

[clessig]

What do you mean by "data partitions at JSC"? Locations of files/paths at JSC?

[sbAsma]

I rebased to a new branch from develop, didn't do any changes besides output files partitions, and data links, and wanted to know how the link to ERA5 affects the tests. I have the following:

With the original link in test_utils.py:

def era5_fname():
    return "/gpfs/scratch/ehpc03/data/{}/ml{}/era5_{}_y{}_m{}_ml{}.grib"

I get:

0: atmorep/tests/validation_test.py::test_datetime
0:   /p/home/jusers/semcheddine1/juwels/2_atmorep_dev/atmorep/atmorep/tests/validation_test.py:63: UserWarning: Timestamp 2021-01-15 02:00:00 not found in ERA5. Skipping
0:     warnings.warn(UserWarning((f"Timestamp {datetime} not found in ERA5. Skipping")))
0: 
0: atmorep/tests/validation_test.py::test_datetime
0:   /p/home/jusers/semcheddine1/juwels/2_atmorep_dev/atmorep/atmorep/tests/validation_test.py:63: UserWarning: Timestamp 2021-02-28 06:00:00 not found in ERA5. Skipping
0:     warnings.warn(UserWarning((f"Timestamp {datetime} not found in ERA5. Skipping")))
0: 
...

Which I am understanding that it is the default output if ERA5 is not found

Now if I adapt the link to what I have access on JSC HPC systems:

def era5_fname():
    return "/p/scratch/atmo-rep/data/era5/new_structure/{}/ml{}/era5_{}_y{}_m{}_ml{}.grib"

I get:

=================================== FAILURES ===================================
0: ________________________________ test_datetime _________________________________
0: 
0: field = 'temperature', model_id = 'wndyj606', BERT = True, epoch = 0
0: 
0:     def test_datetime(field, model_id, BERT, epoch = 0):
0:     
0:         """
0:         Check against ERA5 timestamps.
0:         Loop over all levels individually. 50 random samples for each level.
0:         """
0:     
0:         store = zarr.ZipStore(atmorep_target().format(model_id, model_id, str(epoch).zfill(5)))
0:         atmorep = zarr.group(store)
0:     
0:         nsamples = min(len(atmorep[field]), 50)
0:         samples = rnd.sample(range(len(atmorep[field])), nsamples)
0:         levels = [int(f.split("=")[1]) for f in atmorep[f"{field}/sample=00000"]] if BERT else atmorep[f"{field}/sample=00000"].ml[:]
0:     
0:         get_data = get_BERT if BERT else get_forecast
0:     
0:         for level in levels:
0:             #TODO: make it more elegant
0:             level_idx = level if BERT else np.where(levels == level)[0].tolist()[0]
0:     
0:             for s in samples:
0:                 data, datetime, lats, lons = get_data(atmorep, field, s, level_idx)
0:                 year, month = datetime.year, str(datetime.month).zfill(2)
0:     
0:                 era5_path = era5_fname().format(field, level, field, year, month, level)
0:                 if not os.path.isfile(era5_path):
0:                     warnings.warn(UserWarning((f"Timestamp {datetime} not found in ERA5. Skipping")))
0:                     continue
0:                 era5 = xr.open_dataset(era5_path, engine = "cfgrib")[grib_index(field)].sel(time = datetime, latitude = lats, longitude = lons)
0:     
0:                 #assert (data[0] == era5.values[0]).all(), "Mismatch between ERA5 and AtmoRep Timestamps"
0: >               assert np.isclose(data[0], era5.values[0],rtol=1e-04, atol=1e-07).all(), "Mismatch between ERA5 and AtmoRep Timestamps"
0: E               AssertionError: Mismatch between ERA5 and AtmoRep Timestamps
0: E               assert False
0: E                +  where False = <built-in method all of numpy.ndarray object at 0x146bbc5eae50>()
0: E                +    where <built-in method all of numpy.ndarray object at 0x146bbc5eae50> = array([False, False, False, False, False, False, False, False, False,\n       False, False, False, False, False, False, False, False, False,\n       False, False,  True, False, False, False, False, False, False]).all
0: E                +      where array([False, False, False, False, False, False, False, False, False,\n       False, False, False, False, False, False, False, False, False,\n       False, False,  True, False, False, False, False, False, False]) = <function isclose at 0x146cd9895ff0>(array([255.32582, 255.68324, 256.57874, 257.1178 , 257.11487, 257.10315,\n       257.03772, 256.651  , 256.2174 , 255.8...3.20961, 253.25746, 253.26039, 253.28871, 253.30727, 253.21938,\n       252.95668, 252.62172, 252.41762], dtype=float32), array([253.28871, 253.30727, 253.21938, 252.95668, 252.62172, 252.41762,\n       255.32582, 255.68324, 256.57874, 257.1...4.0514 , 253.57191, 253.26234, 253.1725 , 253.13344, 253.13246,\n       253.20961, 253.25746, 253.26039], dtype=float32), rtol=0.0001, atol=1e-07)
0: E                +        where <function isclose at 0x146cd9895ff0> = np.isclose
0: 
0: /p/home/jusers/semcheddine1/juwels/2_atmorep_dev/atmorep/atmorep/tests/validation_test.py:68: AssertionError

Conclusion: The problem was here all along, and is not related to this "fix-source-writing" changes

[sbAsma]

What do you mean by "data partitions at JSC"? Locations of files/paths at JSC?

Apologies, I haven't seen this comment when I wrote the last one ^^'.

At JSC, data is under the path /p/scratch/atmo-rep/data/era5/new_structure. I am adapting my links to what I have access to here.