import numpy import pandas import xarray import myMPI import os import gc BASEDIR = "./" VAR_LIST = { "2m_temperature": "t2m", "total_precipitation": "tp", "2m_relative_humidity": "r2", "10m_wind_speed": "si10", } LOCATIONS = { "Larnaca": { "years": [2024, 2025], "lonlat": [33.63313, 34.90725], }, "Petrovaradin": { "years": [2015, 2016, 2017], "lonlat": [19.874949, 45.251838] } } def get_time_dim(ds): var = list(tmp.data_vars)[0] dims = tmp[var].dims # time_dims = [d for d in dims if d.lower() == "time" or ( d in ds.coords and numpy.issubdtype(ds[d].dtype, numpy.datetime64) ) ] # if len(time_dims) == 0: raise ValueError(f"No time dimension found in {filein}. Dimensions are: {dims}") # if len(time_dims) > 1: raise ValueError(f"More than one possible time dimension found: {time_dims}") # time_dim = time_dims[0] # return dims, time_dim def download_CERRA(var, year, download=True, process=True, clear=False): import cdsapi client = cdsapi.Client() # dr = f"{BASEDIR}/{var}" if os.path.exists(dr) is False: os.makedirs(dr) # if download: fileout = f"{dr}/reanalysis-cerra-single-levels-{var}-{year}.grib" print(f"Downloading {fileout}", flush=True) client.retrieve("reanalysis-cerra-single-levels", { "variable": [var], "level_type": "surface_or_atmosphere", "data_type": ["reanalysis"], "product_type": "forecast", "year": [str(year)], "month": ["%.2d" %n for n in range(13)[1:]], "day": ["%.2d" %n for n in range(32)[1:]], "time": ["%.2d:00" %(n*3) for n in range(8)], "leadtime_hour": ["1","2","3"], "data_format": "grib" }, fileout) print(f"DONE: {fileout}", flush=True) # if process: filein = f"{dr}/reanalysis-cerra-single-levels-{var}-{year}.grib" print(f"Processing {filein}", flush=True) tmp = xarray.open_dataset(filein, engine="cfgrib", decode_timedelta=False) # dims, time_dim = get_time_dim(tmp) new_order = tuple(d for d in dims if d != time_dim) + (time_dim,) tmp = tmp.transpose(*new_order) # for key in tmp.data_vars.keys(): for yr, ds_yr in tmp[[key]].groupby("time.year"): fileout = f"{dr}/reanalysis-cerra-single-levels-{var}-{key}-{yr}-time.nc" # ds_yr.to_netcdf( fileout, engine="netcdf4", compute=True, encoding={key: {"zlib": True, "complevel": 9}}, ) # print(f"DONE: {fileout}", flush=True) def fun_slave(mpi,task,opt): print("Requesting:", task, flush=True) rcode = task[0](*task[1:]) if rcode: print("Request failed with ", rcode, task, flush=True) else: print("Request done!", task, flush=True) return {} def fun_master(mpi,opt): download=True process=False clear=False # jobs = [] for var in VAR_LIST: for year in YEARS: jobs += [ [download_CERRA, var, year, download, process, clear] for var in VAR_LIST for year in YEARS ] # mpi.exec(jobs, multiple=False, verbose=True) def nearest_yx(lons, lats, lon, lat): # Handle longitude wrapping safely. dlon = ((lons - lon + 180) % 360) - 180 dlat = lats - lat # Approximate distance. Good enough for nearest-grid-cell selection. dist2 = (dlon * numpy.cos(numpy.deg2rad(lat)))**2 + dlat**2 # Find nearest y/x index. flat_index = int(numpy.argmin(dist2.values)) return numpy.unravel_index(flat_index, dist2.shape) def flatten_step_to_time(da): # Stack the two dimensions into one temporary dimension out = da.stack(datetime=("time", "step")) # Flatten valid_time in the same order valid_time = da["valid_time"].stack(datetime=("time", "step")) # Use valid_time as the real coordinate out = out.assign_coords(valid_time=("datetime", valid_time.data)) # Replace the stacked MultiIndex dimension with valid_time out = out.swap_dims({"datetime": "valid_time"}) # Clean up old coordinates out = out.drop_vars(["datetime", "time", "step"], errors="ignore") # Ensure chronological order out = out.sortby("valid_time") return out def add_location_dimension(da, loc): """ Convert scalar latitude/longitude coordinates into location-dependent coordinates. """ lat = float(da["latitude"]) lon = float(da["longitude"]) da = da.reset_coords(["latitude", "longitude"], drop=True) da = da.expand_dims(location=[loc]) da = da.assign_coords( latitude=("location", [lat]), longitude=("location", [lon]), ) return da def read_locations(years): arrays = {} for var in VAR_LIST: arrays[var] = {} for loc in LOCATIONS: arrays[var][loc] = [] for year in years: fname = f"{BASEDIR}/{var}/reanalysis-cerra-single-levels-{var}-{year}.grib" with xarray.open_dataset(fname, engine="cfgrib", decode_timedelta=False) as ds: da = ds[VAR_LIST[var]] # print(f"Processing {var}-{year}...", flush=True) # lons = da["longitude"].load() lats = da["latitude"].load() for loc in LOCATIONS: # print(f"Extracting {loc}...", flush=True) # lon, lat = LOCATIONS[loc]["lonlat"] iy, ix = nearest_yx(lons, lats, lon, lat) da_loc = da.isel(y=iy, x=ix) da_loc = flatten_step_to_time(da_loc) da_loc = add_location_dimension(da_loc, loc) # for coord in ["heightAboveGround"]: if coord in da_loc.coords: da_loc = da_loc.reset_coords(coord, drop=True) # da_loc = da_loc.rename(var) # Load while the file is still open da_loc = da_loc.load() arrays[var][loc].append(da_loc) # print("Combining arrays...", flush=True) # var_arrays = [] for var in VAR_LIST: loc_arrays = [] for loc in LOCATIONS: # concatenate years for this variable/location da_loc_all = xarray.concat(arrays[var][loc], dim="valid_time") da_loc_all = da_loc_all.sortby("valid_time") loc_arrays.append(da_loc_all) # concatenate locations for this variable da_var = xarray.concat(loc_arrays, dim="location") da_var = da_var.transpose("valid_time", "location") da_var = da_var.rename(var) var_arrays.append(da_var) # print("Merging variables...", flush=True) # ds_out = xarray.merge(var_arrays) df = ds_out.to_dataframe().reset_index() df.to_csv("cerra_timeseries_points.csv", index=False) # print("DONE", flush=True) if __name__ == "__main__": # Uncomment one of the lines below to perform the specific action # # Use this to download the CERRA datasets # mpi = myMPI.mpi(fun_master, fun_slave) # # Use this to extract data at specific locations # read_locations(numpy.arange(2014, 2026))