import numpy as np from pathlib import Path from config import Config as cfg from scipy.ndimage import uniform_filter, maximum_filter import pandas as pd import os import re import pygrib as pgrb import sys from utils import uv2df, nint, select_period_bounds, write_mpz_to_nc from formatter import mpz_df_to_xls import time from concurrent.futures import ThreadPoolExecutor t0 = time.time() INIT_DATE_HH = sys.argv[1] datestamp = pd.Timestamp(INIT_DATE_HH[:-2]) SROK = INIT_DATE_HH[-2:] if SROK in ['12', '18']: datestamp = datestamp+pd.Timedelta(1,'d') # flows down to advection, summer aternoons,etc. Valid but maybe fix to match fortran (month from init_date there) OUTPUT_DATE = datestamp.strftime('%Y%m%d') if SROK == '12': hrs_grbs = [f"{h//24}{(h%24)//10}{(h%24)%10}" for h in range(12,61)] # 012...212 gribs elif SROK == '18': hrs_grbs = [f"{h//24}{(h%24)//10}{(h%24)%10}" for h in range(6,55)] # 006...206 gribs else: OUTPUT_DATE = datestamp.strftime('%Y%m%d') hrs_grbs = [f"{h//24}{(h%24)//10}{(h%24)%10}" for h in range(0,49)] # 000...200 gribs GRIBS_DIR = Path(cfg.path_gribs.format(INIT_DATE_HH)) OUTPUT_DIR = Path(f'{cfg.path_output}/{OUTPUT_DATE}/{SROK}') os.makedirs(OUTPUT_DIR / 'tables', exist_ok=True) NLAT, NLON, NH = cfg.nlat, cfg.nlon, cfg.nhours # 1000, 2000, 49 NCELLS = NLAT * NLON # 2 000 000 save_grib_vars_to_npz=cfg.save_grib_vars_to_npz test_city = 'Baikalsk' follow_city = True ##################################################################################################################################################### # Reading raw data from gribs ##################################################################################################################################################### adv_lvl = 850 if 2 <= datestamp.month <= 9 else 925 # Advection field t925 or t850 defined by season print(f'Advection field: {adv_lvl} (n_month={datestamp.month})') grbs_msgs = {sfx: [d for d in descs if not (sfx == 'a_pl' and d['levels'][0] != adv_lvl)] for sfx, descs in cfg.grbs_msgs_mpz.items()} grb_vars = { desc['name']: np.empty((NH, NLAT, NLON), dtype=np.float32) for descs in grbs_msgs.values() for desc in descs} grb_vars['tp'] = np.zeros((NH, NLAT, NLON), dtype=np.float32) # OLD per-hour reader: one task = one forecast hour, opening all 3 suffix files # sequentially inside the task. Replaced by the per-file task model below (mirrors # profiles.py). Kept for reference. # def read_gribs(hr_grb): # h = hrs_grbs.index(hr_grb) # print(f'Reading gribs {hr_grb}') # # for sfx, msgs in grbs_msgs.items(): # fl = str(GRIBS_DIR / f'igfi0{hr_grb}0000{sfx}.ENA6km.grb') # grbs = pgrb.open(fl) # # for idx, desc in msgs.items(): # if (sfx == 'a_pl') and (desc[1] != adv_lvl): # continue # if len(desc) > 1: # name = f'{desc[0]}_{desc[1]}' # msg = grbs.message(idx) # assert (msg.shortName == desc[0]) and (msg.level == desc[1]), f"Wrong message idx for {sfx}: {desc}" # else: # name = f'{desc[0]}' # if (name == 'tp') and (hr_grb == '000'): # continue # msg = grbs.message(idx) # assert msg.shortName == desc[0], f"Wrong message idx for {sfx}: {desc}" # grb_vars[name][h] = msg.values # grbs.close() def read_gribs_one_task(args): # One task = one file (hour × suffix), mirroring profiles.py: threads grab files # individually for finer load-balancing and parallel opens within an hour. t00=time.time() hr_grb, sfx, descs = args #print(f'Reading {sfx}_{hr_grb}') h = hrs_grbs.index(hr_grb) fl = str(GRIBS_DIR / f'igfi0{hr_grb}0000{sfx}.ENA6km.grb') grbs = pgrb.open(fl) for desc in descs: # all mpz descriptors are single-message msg_idx, var, lev, name = desc['messages'][0], desc['var'], desc['levels'][0], desc['name'] # (a_pl's unused advection level is already filtered out of grbs_msgs above.) # tp is accumulated precip; the first forecast hour is the zero baseline. # (hrs_grbs[0], not literal '000', so it also holds for the 12/18H windows.) if (name == 'tp') and (hr_grb == hrs_grbs[0]): continue msg = grbs.message(msg_idx) assert (msg.shortName == var) and (msg.level == lev), f"Wrong message idx for {fl}: {desc}" grb_vars[name][h] = msg.values grbs.close() print(f'Reading {sfx}_{hr_grb} done in {(time.time() - t00)/60:.1f}m') tasks = [(hr_grb, sfx, descs) for hr_grb in hrs_grbs for sfx, descs in grbs_msgs.items()] print('Reading gribs...') if Path(OUTPUT_DIR / "grib_vars.npz").is_file(): save_grib_vars_to_npz = False with np.load(OUTPUT_DIR / "grib_vars.npz") as npz: for k in grb_vars: grb_vars[k] = np.asarray(npz[k], dtype=np.float32) else: with ThreadPoolExecutor(max_workers=cfg.n_jobs_reading_gribs) as ex: list(ex.map(read_gribs_one_task, tasks)) print(f"Reading gribs done in: {(time.time() - t0)/60:.1f}m \n") t0 = time.time() if save_grib_vars_to_npz: np.savez_compressed(OUTPUT_DIR / f'grib_vars.npz', **grb_vars) for var in grb_vars: print(f'{var}: {grb_vars[var].mean(axis = (-1,-2))}') ##################################################################################################################################################### ### в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ Vars processing - stacking, smoothing, etc. в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ### ##################################################################################################################################################### def smooth_3_3(arr): a = np.pad(arr, ((0,0),(1,1),(1,1)), mode='edge') return ( a[:, 0:-2, 0:-2] + a[:, 0:-2, 1:-1] + a[:, 0:-2, 2:] + a[:, 1:-1, 0:-2] + a[:, 1:-1, 1:-1] + a[:, 1:-1, 2:] + a[:, 2:, 0:-2] + a[:, 2:, 1:-1] + a[:, 2:, 2:] ) / np.float32(9) #Or lise that: #def box_filter_spatial(arr, k=3): # p = k // 2 # padding size: for k=3, p=1 # Pad edges by 1 pixel on each spatial side, replicate border values # (T,H,W) --> (T, H+2, W+2) # a = np.pad(arr, ((0,0),(p,p),(p,p)), mode='edge') # Cumulative sum along rows, then columns # Each cell now contains sum of all elements above-left of it # cs = np.cumsum(np.cumsum(a, axis=1), axis=2) # Extract box sums using the 2D prefix sum formula: # sum of rectangle = bottom_right - bottom_left - top_right + top_left # Then divide by k*k to get the mean # return (cs[:,k:,k:] - cs[:,k:,:-k] - cs[:,:-k,k:] + cs[:,:-k,:-k]) / k**2 #t2m_grid = uniform_filter(np.stack(t2m_list, axis=0) - 273.16, size=(1, 3, 3), mode='nearest') # stacking hours --> smoothing 3*3 --> (49,1000,2000) #tz_grid = uniform_filter(np.stack(tz_list, axis=0) - 273.16, size=(1, 3, 3), mode='nearest') # stacking hours --> smoothing 3*3 --> (49,1000,2000) Pr_grid = np.diff(grb_vars['tp'], axis=0, prepend=0) # deaccumulate h0 = 0.0, h1 = h1, h2 = h2-h1 ... del grb_vars['tp'] #prec_grid = np.maximum(prec_grid, smooth_3_3(prec_grid)) # maximum between grid cell value and 3*3 window np.maximum(Pr_grid, smooth_3_3(Pr_grid), out=Pr_grid) # inplace Pr_grid[Pr_grid < cfg.pr_min] = 0.0 #prec_grid_bc = prec_grid[None] # (1, 49, 1000, 2000) — broadcasted to shape-match pertrubation variants # в”Ђв”Ђ Advection в”Ђв”Ђв”Ђв”Ђ tadv_grid = grb_vars[f't_{adv_lvl}'] #Adv_grid = np.diff(tadv_grid, axis=0, prepend=tadv_grid[[0]]) # not used at all, Adv for periods derived later #Adv_grid = maximum_filter(Adv_grid, size=(1, 3, 3), mode='nearest') del grb_vars[f't_{adv_lvl}'] # в”Ђв”Ђ Vz, Str pertrubated в”Ђв”Ђв”Ђв”Ђ Vz_grid, _ = uv2df(grb_vars['u_70'], grb_vars['v_70']) del grb_vars[f'u_70'], grb_vars[f'v_70'], _ #ddz_grid = nint(ddz_grid) Vz_grid = smooth_3_3(Vz_grid) # smoothed 3*3 np.maximum(Vz_grid, 1.0, out=Vz_grid) # 1.0 where Vz < 1.0, inplace Vz_grid = nint(Vz_grid) # rounding to int, nint rounding exactly like in Fortran #Vz_grid_pert = nint(np.maximum(Vz_grid[None] + dv_arr, 1.0)) # --> (9, 1, 1, 1) ensemble Vz variants --> where < 1 = 1 --> round to int # loading constants with np.load(cfg.constfields) as npz: dz = np.asarray(npz['dz'], dtype=np.float32) # layer thickness grid lons = np.asarray(npz['icon6_geo_lon'], dtype=np.float32) lats = np.asarray(npz['icon6_geo_lat'], dtype=np.float32) Str_grid = grb_vars['2t'] - grb_vars['t_70'] del grb_vars['2t'], grb_vars['t_70'] Str_grid /= dz Str_grid = smooth_3_3(Str_grid) #Str_grid_pert = Str_grid[None] + dStr_arr # --> (9, 1, 1, 1) ensemble Str variants prmsl = grb_vars['prmsl'] / 100 print(f"Processing done in: {(time.time() - t0)/60:.1f}m \n") t0 = time.time() ##################################################################################################################################################### ### в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ MPRZ decision tree в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ### ##################################################################################################################################################### # в”Ђв”Ђ Cities metadata в”Ђв”Ђв”Ђв”Ђв”Ђ cities = pd.read_csv(cfg.cities_csv) cities = cities[cities['note'] != 'add_coords'] i_idx = cities['i'].values # (n_cities,) 0-based column (lon) j_idx = cities['j'].values # (n_cities,) 0-based row (lat) utc_offset = cities['utc_offset'].values # (n_cities,) n_cities = len(cities) fo_names = cities['fo_name'].values cities_name_en = cities['name'].values t_c_idx = np.where(cities_name_en == test_city)[0] # в”Ђв”Ђ MPZ decision tree в”Ђв”Ђв”Ђв”Ђв”Ђ def get_typeMPZ_hour(Vz_grid, Str_grid, Pr_grid): conditions = [ Vz_grid > 6, ## Check Str_grid >= 0.9, Pr_grid >= 1.0, (Str_grid >= 0.5) & (Pr_grid != 0), (Str_grid >= 0.5) & (Vz_grid >= 5), Str_grid >= 0.5, (Str_grid >= 0.1) & (Pr_grid != 0), (Str_grid >= 0.1) & (Vz_grid >= 5), (Str_grid >= 0.1) & (Vz_grid > 3), Str_grid >= 0.1, Vz_grid > 5, Vz_grid > 4, ] choices = [3.3, 3.2, 3.1, 2.3, 2.2, 2.1, 2.3, 2.2, 2.1, 1.3, 2.1, 1.2] return np.select(conditions, choices, default=1.1).astype(np.float32) # --> (49, 1000, 2000) typeMPZ {1.1 - 3.3} typeMPZ_grid = get_typeMPZ_hour(Vz_grid, Str_grid, Pr_grid) # в”Ђв”Ђ Saving for maps в”Ђв”Ђв”Ђв”Ђв”Ђ #np.savez_compressed(OUTPUT_DIR / f'MPZ_{INIT_DATE_HH}.npz', # INIT_DATE, fix if needed # typeMPZ=typeMPZ_grid, # lons = lons, # lats = lats) # Writing to netCDF write_mpz_to_nc(typeMPZ = typeMPZ_grid, lats = lats, lons = lons, output_path = str(OUTPUT_DIR / f'mpz_Russia_{INIT_DATE_HH}.nc')) # Writing to .dat for grads = [typeMPZ, p0, Pr, u, v, parStr, parAdv] # Fortran, strict order with open(f'{cfg.path_grads}/mpz.dat', 'wb') as f: for h in range(NH): for fld in [typeMPZ_grid, prmsl]: fld[h].astype('{date_ctl}', text, count=1) ctl.write_text(text) n_stamped += n print(f'tdef date stamped ({date_ctl}) in {n_stamped} ctl files') print(f"MPZ grid done in: {(time.time() - t0)/60:.1f}m \n") t0 = time.time() # в”Ђв”Ђ Cities values extraction в”Ђв”Ђв”Ђв”Ђв”Ђ Vz_cities = Vz_grid[:, j_idx, i_idx].T # (n_cities, 49) Str_cities = Str_grid[:, j_idx, i_idx].T Pr_cities = Pr_grid[:, j_idx, i_idx].T tadv_cities = tadv_grid[:, j_idx, i_idx].T typeMPZ_cities = typeMPZ_grid[:, j_idx, i_idx].T del Vz_grid, Str_grid, Pr_grid, tadv_grid, typeMPZ_grid #per_city_vars['Vz'] = nint(per_city_vars['Vz']) # в”Ђв”Ђ Summer afternoon worsening в”Ђв”Ђв”Ђв”Ђ def apply_afternoon_worsening(typeMPZ_arr): if 3 <= datestamp.month <= 11:# ! check time_mect = nh-1 + dtimeCity(nc,n) ih_utc = np.arange(cfg.nhours) # (49,) local_hour = (ih_utc[np.newaxis, :] + utc_offset[:, np.newaxis]) % 24 # (n_cities, 49) is_afternoon = (local_hour >= 13) & (local_hour <= 18) # (n_cities, 49) can_worsen = (typeMPZ_arr < 3.0) if typeMPZ_arr.ndim == 2 else (typeMPZ_arr < 2.0) # different conditions for deterministic/pertubed pipelines --> (9, n_cities, 49) ia = is_afternoon if typeMPZ_arr.ndim == 2 else is_afternoon[None] # 2 dim for base case, 3 for pertrubed typeMPZ_arr = np.where( ia & can_worsen, typeMPZ_arr + 1.0, typeMPZ_arr).astype(np.float32) print(f'Summer afternoon worsening applied (n_month={datestamp.month})') else: print(f'Summer afternoon worsening not applied (n_month={datestamp.month})') return typeMPZ_arr typeMPZ_cities = apply_afternoon_worsening(typeMPZ_cities) print(f"Cities extraction done in: {(time.time() - t0)/60:.1f}m \n") t0 = time.time() ##################################################################################################################################################### ### в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ Periodic MPZ в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ### ##################################################################################################################################################### _season = select_period_bounds(datestamp) bounds = np.array(cfg.prd_bounds[_season], dtype=np.int32) # (nperiods, 3) adv_b = np.array(cfg.adv_bounds, dtype=np.int32) # (nperiods, 3) #ref_offset = int(cities.loc[cities["fo_name"]==cfg.fos_12H[0], "utc_offset"].iloc[0]) if SROK=="12" else 3 #print(f'SROK={SROK}, time_offset={ref_offset}') shift = utc_offset - 3 # - ref_offset # (n_cities,) relative to Moscow UTC+3 # why shift from moscow time? ih = np.arange(cfg.nhours) # (nhours,) # Here invalid negative indexing and masking for cities that out of forecast range # out_of_range cleaning in the very rnd of the block # в”Ђв”Ђ Str/Vz 3h aggregatioh в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ # Why aggregation over 3 middle hours in the period? # (n_cities, nperiods) starts of 3-hours aggregation periods "we should get 10AM, in Dvstk MSK[5] = Dvstk 10AM" str_start = bounds[np.newaxis, :, 0] - shift[:, np.newaxis] # contains negative vals for periods/cities out of range # (n_cities, nperiods, nhours), True where hour in aggregation 3h period str_mask = (ih[np.newaxis, np.newaxis, :] >= str_start[:, :, np.newaxis]) \ & (ih[np.newaxis, np.newaxis, :] <= str_start[:, :, np.newaxis] + 2) # contains some invalid Trues for periods/cities out of range Str_p = np.nanmean(np.where(str_mask, Str_cities[:, np.newaxis, :], np.nan), axis=2) #(n_cities, nperiods) Str aggregated values Vz_p = nint(np.nanmean(np.where(str_mask, Vz_cities[:, np.newaxis, :], np.nan), axis=2)).astype(np.float32) # float32 to preserve NaN for out-of-range city-period combos # в”Ђв”Ђ prec aggregatioh over the full period в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ pr_start = bounds[np.newaxis, :, 1] - shift[:, np.newaxis] pr_end = bounds[np.newaxis, :, 2] - shift[:, np.newaxis] # starts and ends of aggregation periods, with time shifts pr_mask = (ih[np.newaxis, np.newaxis, :] >= pr_start[:, :, np.newaxis]) \ & (ih[np.newaxis, np.newaxis, :] <= pr_end[:, :, np.newaxis]) # (n_cities, nperiods, nhours), True where hour in aggregation period Pr_p = np.nansum(np.where(pr_mask, Pr_cities[:, np.newaxis, :], np.nan), axis=2).astype(np.float32) # (ncities, n_periods) Pr_p = np.where(Pr_p < cfg.pr_min, 0.0, Pr_p) # в”Ђв”Ђ Advection: two 3-hour tadv tendencies, city-shifted в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ tadv = tadv_cities # (n_cities, nhours) dt1_utc = adv_b[np.newaxis, :, 0] - shift[:, np.newaxis] # (n_cities, nperiods): hours indices for adv vals #how advection is representative for the periods? dt2_utc = adv_b[np.newaxis, :, 1] - shift[:, np.newaxis] dt1_val = np.take_along_axis(tadv, dt1_utc, axis=1) \ - np.take_along_axis(tadv, dt1_utc - 3, axis=1) # (ncities, n_periods): advection in 3h period 1 dt2_val = np.take_along_axis(tadv, dt2_utc, axis=1) \ - np.take_along_axis(tadv, dt2_utc - 3, axis=1) # (ncities, n_periods): advection in 3h period 2 same_sign = ((dt1_val > 0) & (dt2_val > 0)) | ((dt1_val < 0) & (dt2_val < 0)) # (ncities, n_periods): bool True if advection in 3h periods the same or the opposite sign opp_sign = ((dt1_val > 0) & (dt2_val < 0)) | ((dt1_val < 0) & (dt2_val > 0)) Adv_p = np.where(same_sign, dt1_val + dt2_val, np.where(opp_sign, dt2_val, 0.0)) # (ncities, n_periods): advection vals used for correction later # в”Ђв”Ђ typeMPZ decision tree в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ #conditions = [ # Vz_p > 6, # Str_p >= 0.9, # Pr_p >= 1.0, # (Str_p >= 0.5) & (Pr_p != 0), # Str_p >= 0.5, # (Str_p >= 0.1) & (Vz_p > 5) & (Pr_p != 0), # (Str_p >= 0.1) & (Vz_p > 5), # Str_p >= 0.1, # Vz<=5, weak inv. в†’ 1.3 # Vz_p > 5, # no inv., strong wind # Vz_p > 4, # no inv., moderate wind #] #choices = [3.3, 3.2, 3.1, # 2.3, 2.2, # 2.3, 2.1, # 1.3, # 2.1, 1.2] #typeMPZ_p = np.select(conditions, choices, default=1.1).astype(np.float32) typeMPZ_p = get_typeMPZ_hour(Vz_p, Str_p, Pr_p) # Summer afternoon bump typeMPZ + 1 for periods 1, 5 in Mar–Nov. if 3 <= datestamp.month <= 11: p_afternoon = np.zeros((n_cities, cfg.nperiods), dtype=bool) p_afternoon[:, [1, 5]] = True # Fortran np=2 and np=6 bump = p_afternoon & (Vz_p <= 5) & (typeMPZ_p < 2.0) typeMPZ_p = np.where(bump, typeMPZ_p + 1.0, typeMPZ_p).astype(np.float32) # в”Ђв”Ђ Out of forecast range values, = NaN, for example msk-8h hours are out of period 0 в”Ђ out_of_range = (str_start < 0) | (pr_start < 0) | (pr_end < 0) | ((dt1_utc - 3) < 0) | ((dt2_utc - 3) < 0) # (n_cities, nperiods) typeMPZ_p = np.where(out_of_range, np.nan, typeMPZ_p) Str_p = np.where(out_of_range, np.nan, Str_p) Vz_p = np.where(out_of_range, np.nan, Vz_p) Pr_p = np.where(out_of_range, np.nan, Pr_p) Adv_p = np.where(out_of_range, np.nan, Adv_p) # в”Ђв”Ђ periodic typeMPZ correction from hourly typeMPZ в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ hourly = typeMPZ_cities # (n_cities, nhours) typeMPZ_p_res = np.empty((n_cities, cfg.nperiods), dtype=np.float32) # (n_cities, n_periods): empty, filled with typeMPZ {1.1...3.3} in next cycle for p in range(cfg.nperiods): tree_val = typeMPZ_p[:, p] #(n_cities, ) raw typeMPZ values # contains nans for out of range vals, they all False in any coditioning is_high = (tree_val >= 3.2) & (tree_val <= 3.3) # (n_cities, ): bool mask, correction not applied where typeMPZ = {3.1...3.3} p_window = pr_mask[:, p, :] #(n_cities, n_hours): bool, True for period full window any_31 = ((hourly == 3.1) & p_window).any(axis=1) # (n_cities, ): bool, any hour within the period window = 3.1 any_23 = ((hourly == 2.3) & p_window).any(axis=1) # (n_cities, ): bool, any hour within the period window = 2.3 prev_is_31 = (typeMPZ_p_res[:, p - 1] == 3.1) if p > 0 else np.zeros(n_cities, dtype=bool) # (n_cities, ): bool any hour within the previous period window = 3.1 # careful with nans propagation print(p, hourly.shape, any_31.shape, prev_is_31.shape, tree_val.shape) result = np.where(any_31, 3.1, np.where(any_23, 2.3, np.where(prev_is_31, 2.3, tree_val))).astype(np.float32) # (n_cities, ): typeMPZ {1.1...3.3} # в”Ђв”Ђ Advection correction в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ # Applied when typeMPZ {2.1...2.3} is_2x = (result > 2.0) & (result < 2.4) # (n_cities, ): bool # Replicates Fortran stale-variable bug: parAdv_p in the output loop retains ! saved bug # the value from the last iteration of the preceding loop (np=Nperiods), # so every period uses period-6 advection instead of its own. adv = Adv_p[:, -1] # (n_cities, ): (n_cities, ) bool result = np.where(is_2x & (adv < cfg.par_adv_minus_cr), 2.3, result) result = np.where(is_2x & (adv >= cfg.par_adv_plus_cr), result - 1, result) # (n_cities, ): typeMPZ {1.1...3.3} typeMPZ_p_res[:, p] = np.where(is_high, tree_val, result) # (n_cities, ): filling with typeMPZ {1.1...3.3} # в”Ђв”Ђ Integer MPZ class 1/2/3 в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ MPZ_p = np.where(out_of_range, np.nan, np.floor(typeMPZ_p_res)).astype(np.int32) typeMPZ_p_res= np.where(out_of_range, np.nan, typeMPZ_p_res) # cleaning out_of_range print(f"Periodic MPZ done in: {(time.time() - t0)/60:.1f}m \n") t0 = time.time() ##################################################################################################################################################### ### в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ Probablistic MPZ в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ### ##################################################################################################################################################### # IMPORTANT: summ_wors_hr must be True for typeMPZ_ensemble_cities to match Fortran's # subtypeMPZ_city (afternoon worsening applied for months 3–11). P_city = np.zeros((n_cities, cfg.nperiods, 3), dtype=np.int32) # (n_cities, nperiods, 3): dummy to fill # в”Ђв”Ђ 9 ensemple variants for Vz and Str в”Ђв”Ђв”Ђ dv_arr = np.array(cfg.dv, dtype=np.float32)[:, None, None] # pertrubation dVz variants for probablistic forecast dStr_arr = np.array(cfg.dstr, dtype=np.float32)[:, None, None] # pertrubation variants for probablistic forecast typeMPZ_cities_pert = get_typeMPZ_hour(Vz_cities[None] + dv_arr, Str_cities[None] + dStr_arr, Pr_cities[None]) typeMPZ_cities_pert = apply_afternoon_worsening(typeMPZ_cities_pert) for p in cfg.prds_to_write: # 0-based p=2..6 в†’ Fortran np=3..Nperiods # в”Ђв”Ђ Raw probabilities в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ nh0 = cfg.prob_srarts[p] - shift # (n_cities,) city-shifted prob window start hour nlen = int(cfg.prob_lens[p]) # (n_cities,) city-shifted prob window length prob_mask = (ih[np.newaxis, :] >= nh0[:, np.newaxis]) & \ (ih[np.newaxis, :] < nh0[:, np.newaxis] + nlen) # (n_cities, nhours): bool, True where window for prob aggregating # inclusive [nh1, nh1+N-1] in Fortran = N hours m = prob_mask[np.newaxis] # (1, n_cities, nhours): bool, mask broadcast over 9 variants t = typeMPZ_cities_pert # (9, n_cities, nhours): {1.1, ..., 3.3}, pertrubated vars # (n_cities,): int, count of 1 MPZ subtype per all pert vars and hours within a period cnt1 = np.sum((t < 2.0) & m, axis=(0, 2)).astype(np.int32)# ! changed .astype(np.float32) cnt2 = np.sum((t >= 2.0) & (t < 3.0) & m, axis=(0, 2)).astype(np.int32) #... cnt3 = np.sum((t >= 3.0) & m, axis=(0, 2)).astype(np.int32) N = cnt1 + cnt2 + cnt3 # (n_cities,): count of all MPZ subtypes P1 = nint(cnt1 * 100.0 / N).astype(np.int32) # (n_cities,): % of 1 MPZ type P2 = nint(cnt2 * 100.0 / N).astype(np.int32) # ... P3 = nint(cnt3 * 100.0 / N).astype(np.int32) # в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ Tie-breaking: when rounded percentages tie, add one count from the deterministic result. в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ # Fortran only breaks specific tie patterns (not all ties): # all_eq: P1==P2==P3 в†’ any det class eligible # p12_p3z: P1==P2, P3=0 в†’ only det==1 or det==2 # p1z_p23: P1=0, P2==P3 в†’ only det==2 or det==3 det = np.where(np.isnan(MPZ_p[:, p]), 0, MPZ_p[:, p]).astype(np.int32) # (n_cities,): intMPZ type, 0 for out-of-range (NaN) cities (nowhere because preiods 2..7) all_eq = (P1 == P2) & (P1 == P3) # (n_cities,): bool p12_p3z = (~all_eq) & (P1 == P2) & (P3 == 0) p1z_p23 = (~all_eq) & (P1 == 0) & (P2 == P3) cnt1 += ((all_eq | p12_p3z) & (det == 1)).astype(np.int32)# ! changed .astype(np.float32) cnt2 += ((all_eq | p12_p3z | p1z_p23) & (det == 2)).astype(np.int32) cnt3 += ((all_eq | p1z_p23) & (det == 3)).astype(np.int32) N += (all_eq | (p12_p3z & ((det == 1) | (det == 2))) | (p1z_p23 & ((det == 2) | (det == 3)))).astype(np.int32) P1 = nint(cnt1 * 100.0 / N).astype(np.int32) P2 = nint(cnt2 * 100.0 / N).astype(np.int32) P3 = nint(cnt3 * 100.0 / N).astype(np.int32) # в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ Normalize to sum=100: add remainder to the largest class; subtract from smallest if over. в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ # Fortran sequential if-chain в†’ first-wins for ties, matching np.argmax/argmin. Parr = np.stack([P1, P2, P3], axis=1) # (n_cities, 3) dP = (100 - Parr.sum(axis=1)) # (n_cities,): reminder of 100-sum target = np.where(dP > 0, np.argmax(Parr, axis=1), np.argmin(Parr, axis=1)) Parr[np.arange(n_cities), target] += dP P_city[:, p, :] = Parr # (n_cities, nperiods, 3): probability in %, filling dummy # в”Ђв”Ђ Probability-based typeMPZ_p_res correction в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ # Applied in Fortran's output table loop for np=3..Nperiods. # If one class dominates with P > 55 %, shift the integer part of typeMPZ_p_res to match it. #for p in cfg.prds_to_write: # ! here typeMPZ might change 2x, etc. 3.3 --> 1.1 # P_arr = P_city[:, p, :].astype(np.int32) # (n_cities, 3): probability in % # cand = np.where(P_arr > 55, P_arr, -1.0) # (n_cities, 3): probability in %, -1 where $ ≤55 # Pmax_val = cand.max(axis=1) # (n_cities,): probability in %, -1 when nothing exceeds 55 # Pmax_idx = np.argmax(cand, axis=1) + 1 # (n_cities,): {1-3} dominant class # has_dominant = (Pmax_val > 0) # (n_cities,): bool, True where it is dominant class # curr_int = np.floor(typeMPZ_p_res[:, p]).astype(np.int32) # (n_cities,): int current typeMPZ main class {1-3} # dtype_p = (curr_int - Pmax_idx).astype(np.float32) # (n_cities,): correction value {1-3} # typeMPZ_p_res[:, p] = np.where(has_dominant, # (n_cities,): correction only where it is dominant probablistic class # typeMPZ_p_res[:, p] - dtype_p, # (n_cities,): correction # typeMPZ_p_res[:, p]) # (n_cities,): leave as it is if no dominant probablistic class print(f"Probablistic MPZ done in: {(time.time() - t0)/60:.1f}m \n") t0 = time.time() ##################################################################################################################################################### ### в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ Output tables в”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђв”Ђ### ##################################################################################################################################################### periods = pd.DataFrame(typeMPZ_p_res[:, 2:], columns = [f'p_{i}_py' for i in cfg.prds_to_write], index = cities['name_rus_mpz_tables']) probs = pd.DataFrame(P_city[:,2:,:].reshape(n_cities,-1), columns = [f'p{x}_prob{i}' for x in range(5) for i in range(1,4)], index = cities['name_rus_mpz_tables']) full_tbl = pd.concat([periods, probs], axis = 1) full_tbl['fo'] = fo_names fos_to_write = cfg.fos_12H if SROK == '12' else np.unique(fo_names) for fo in fos_to_write: wb = mpz_df_to_xls(full_tbl[full_tbl['fo'] == fo], init_date=INIT_DATE_HH, output_date=OUTPUT_DATE) wb.save(OUTPUT_DIR / 'tables' / f'mpz_{fo}.xlsx') os.makedirs(OUTPUT_DIR / 'plots' / fo, exist_ok=True) # creating directories for plots from Grads print(f"All MPZ done in: {(time.time() - t0)/60:.1f}m \n")