| """Copyright (c) Microsoft Corporation. Licensed under the MIT license.""" |
|
|
| import math |
| from functools import partial |
| from typing import Optional |
|
|
| import torch |
|
|
| __all__ = [ |
| "level_to_str", |
| "log_transform", |
| "log_untransform", |
| "normalise_surf_var", |
| "normalise_atmos_var", |
| "unnormalise_surf_var", |
| "unnormalise_atmos_var", |
| ] |
|
|
|
|
| def level_to_str(level: float) -> str: |
| """Convert a pressure level to a string in a consistent way. |
| |
| Args: |
| level (float): Pressure level. |
| |
| Returns: |
| str: Consistent string representation. |
| """ |
| level = round(float(level), 3) |
| |
| if level % 1 == 0: |
| level = int(level) |
| |
| return str(level).replace(".", "_") |
|
|
|
|
| def normalise_surf_var( |
| x: torch.Tensor, |
| name: str, |
| stats: Optional[dict[str, tuple[float, float]]] = None, |
| unnormalise: bool = False, |
| ) -> torch.Tensor: |
| """Normalise a surface-level variable.""" |
| if stats and name in stats: |
| location, scale = stats[name] |
| else: |
| location = locations[name] |
| scale = scales[name] |
| if unnormalise: |
| return x * scale + location |
| else: |
| return (x - location) / scale |
|
|
|
|
| def normalise_atmos_var( |
| x: torch.Tensor, |
| name: str, |
| atmos_levels: tuple[int | float, ...], |
| unnormalise: bool = False, |
| ) -> torch.Tensor: |
| """Normalise an atmospheric variable.""" |
| level_locations: list[int | float] = [] |
| level_scales: list[int | float] = [] |
| for level in atmos_levels: |
| level_locations.append(locations[f"{name}_{level_to_str(level)}"]) |
| level_scales.append(scales[f"{name}_{level_to_str(level)}"]) |
| location = torch.tensor(level_locations, dtype=x.dtype, device=x.device) |
| scale = torch.tensor(level_scales, dtype=x.dtype, device=x.device) |
|
|
| if unnormalise: |
| return x * scale[..., None, None] + location[..., None, None] |
| else: |
| return (x - location[..., None, None]) / scale[..., None, None] |
|
|
|
|
| unnormalise_surf_var = partial(normalise_surf_var, unnormalise=True) |
| unnormalise_atmos_var = partial(normalise_atmos_var, unnormalise=True) |
|
|
| _LOG_TRANSFORM_EPS = 1e-3 |
|
|
|
|
| def log_transform(x: torch.Tensor) -> torch.Tensor: |
| """Log-transform a variable: `log(x + eps) - log(eps)` with `eps = 1e-3`.""" |
| return torch.log(x + _LOG_TRANSFORM_EPS) - math.log(_LOG_TRANSFORM_EPS) |
|
|
|
|
| def log_untransform(x: torch.Tensor) -> torch.Tensor: |
| """Inverse of :func:`log_transform`: `eps * (exp(x) - 1)`. |
| |
| Use float32 or double to avoid numerical instability. |
| """ |
|
|
| return _LOG_TRANSFORM_EPS * (torch.exp(x) - 1) |
|
|
|
|
| locations: dict[str, float] = { |
| "z": -1.386496e03, |
| "lsm": 0.000000e00, |
| "slt": 0.000000e00, |
| "static_co": 0.000000e00, |
| "static_so2_log": 0.000000e00, |
| "static_ammonia_log": 0.000000e00, |
| "static_nox_log": 0.000000e00, |
| "static_so2": 0.000000e00, |
| "static_co_log": 0.000000e00, |
| "static_ammonia": 0.000000e00, |
| "static_nox": 0.000000e00, |
| "wmb": 0.000000e00, |
| "lat_mask": 0.000000e00, |
| "2t": 2.785140e02, |
| "10u": -5.135059e-02, |
| "10v": 1.891580e-01, |
| "msl": 1.009578e05, |
| "z_50": 1.993730e05, |
| "z_100": 1.576421e05, |
| "z_150": 1.331414e05, |
| "z_200": 1.153300e05, |
| "z_250": 1.012231e05, |
| "z_300": 8.941415e04, |
| "z_400": 6.998038e04, |
| "z_500": 5.411537e04, |
| "z_600": 4.064833e04, |
| "z_700": 2.892882e04, |
| "z_850": 1.374978e04, |
| "z_925": 7.015005e03, |
| "z_1000": 7.381545e02, |
| "u_50": 5.653076e00, |
| "u_100": 1.027951e01, |
| "u_150": 1.354061e01, |
| "u_200": 1.420915e01, |
| "u_250": 1.334584e01, |
| "u_300": 1.180173e01, |
| "u_400": 8.817291e00, |
| "u_500": 6.563273e00, |
| "u_600": 4.814521e00, |
| "u_700": 3.345237e00, |
| "u_850": 1.418379e00, |
| "u_925": 6.172657e-01, |
| "u_1000": -3.328723e-02, |
| "v_50": 4.226111e-03, |
| "v_100": 1.411897e-02, |
| "v_150": -3.697671e-02, |
| "v_200": -4.507801e-02, |
| "v_250": -2.980338e-02, |
| "v_300": -2.294770e-02, |
| "v_400": -1.771003e-02, |
| "v_500": -2.387986e-02, |
| "v_600": -2.716674e-02, |
| "v_700": 2.153583e-02, |
| "v_850": 1.428150e-01, |
| "v_925": 2.053480e-01, |
| "v_1000": 1.867637e-01, |
| "t_50": 2.124864e02, |
| "t_100": 2.084042e02, |
| "t_150": 2.133201e02, |
| "t_200": 2.180615e02, |
| "t_250": 2.227710e02, |
| "t_300": 2.288696e02, |
| "t_400": 2.421368e02, |
| "t_500": 2.529492e02, |
| "t_600": 2.611347e02, |
| "t_700": 2.674010e02, |
| "t_850": 2.745600e02, |
| "t_925": 2.773572e02, |
| "t_1000": 2.810130e02, |
| "q_50": 2.678180e-06, |
| "q_100": 2.633677e-06, |
| "q_150": 5.254625e-06, |
| "q_200": 1.940632e-05, |
| "q_250": 5.773618e-05, |
| "q_300": 1.273861e-04, |
| "q_400": 3.855659e-04, |
| "q_500": 8.529599e-04, |
| "q_600": 1.541429e-03, |
| "q_700": 2.431637e-03, |
| "q_850": 4.575618e-03, |
| "q_925": 6.033134e-03, |
| "q_1000": 7.030342e-03, |
| "w_50": -4.831830e-05, |
| "w_100": -1.158353e-05, |
| "w_150": 2.323354e-05, |
| "w_200": 2.460374e-05, |
| "w_250": 3.384949e-05, |
| "w_300": 1.310551e-04, |
| "w_400": 3.637421e-04, |
| "w_500": 4.369851e-04, |
| "w_600": 5.889935e-04, |
| "w_700": 3.267191e-03, |
| "w_850": 1.430376e-02, |
| "w_925": 2.180887e-02, |
| "w_1000": 2.735809e-02, |
| "co_100": 0.000000e00, |
| "co_1000": 0.000000e00, |
| "co_150": 0.000000e00, |
| "co_200": 0.000000e00, |
| "co_250": 0.000000e00, |
| "co_300": 0.000000e00, |
| "co_400": 0.000000e00, |
| "co_50": 0.000000e00, |
| "co_500": 0.000000e00, |
| "co_600": 0.000000e00, |
| "co_700": 0.000000e00, |
| "co_850": 0.000000e00, |
| "co_925": 0.000000e00, |
| "no2_100": 0.000000e00, |
| "no2_1000": 0.000000e00, |
| "no2_150": 0.000000e00, |
| "no2_200": 0.000000e00, |
| "no2_250": 0.000000e00, |
| "no2_300": 0.000000e00, |
| "no2_400": 0.000000e00, |
| "no2_50": 0.000000e00, |
| "no2_500": 0.000000e00, |
| "no2_600": 0.000000e00, |
| "no2_700": 0.000000e00, |
| "no2_850": 0.000000e00, |
| "no2_925": 0.000000e00, |
| "no_100": 0.000000e00, |
| "no_1000": 0.000000e00, |
| "no_150": 0.000000e00, |
| "no_200": 0.000000e00, |
| "no_250": 0.000000e00, |
| "no_300": 0.000000e00, |
| "no_400": 0.000000e00, |
| "no_50": 0.000000e00, |
| "no_500": 0.000000e00, |
| "no_600": 0.000000e00, |
| "no_700": 0.000000e00, |
| "no_850": 0.000000e00, |
| "no_925": 0.000000e00, |
| "go3_100": 0.000000e00, |
| "go3_1000": 0.000000e00, |
| "go3_150": 0.000000e00, |
| "go3_200": 0.000000e00, |
| "go3_250": 0.000000e00, |
| "go3_300": 0.000000e00, |
| "go3_400": 0.000000e00, |
| "go3_50": 0.000000e00, |
| "go3_500": 0.000000e00, |
| "go3_600": 0.000000e00, |
| "go3_700": 0.000000e00, |
| "go3_850": 0.000000e00, |
| "go3_925": 0.000000e00, |
| "pm10": 0.000000e00, |
| "pm1": 0.000000e00, |
| "pm2p5": 0.000000e00, |
| "so2_100": 0.000000e00, |
| "so2_1000": 0.000000e00, |
| "so2_150": 0.000000e00, |
| "so2_200": 0.000000e00, |
| "so2_250": 0.000000e00, |
| "so2_300": 0.000000e00, |
| "so2_400": 0.000000e00, |
| "so2_50": 0.000000e00, |
| "so2_500": 0.000000e00, |
| "so2_600": 0.000000e00, |
| "so2_700": 0.000000e00, |
| "so2_850": 0.000000e00, |
| "so2_925": 0.000000e00, |
| "tcco": 0.000000e00, |
| "tcno2": 0.000000e00, |
| "tc_no": 0.000000e00, |
| "gtco3": 0.000000e00, |
| "tcso2": 0.000000e00, |
| "swh": 2.53639e00, |
| "pp1d": 1.06605e01, |
| "mwd": 0.00000e00, |
| "mwp": 8.61902e00, |
| "shww": 1.17786e00, |
| "mdww": 0.00000e00, |
| "mpww": 4.22253e00, |
| "shts": 2.07783e00, |
| "mdts": 0.00000e00, |
| "mpts": 9.45161e00, |
| "swh1": 1.79443e00, |
| "mwd1": 0.00000e00, |
| "mwp1": 9.91430e00, |
| "swh2": 7.30807e-01, |
| "mwd2": 0.00000e00, |
| "mwp2": 8.95218e00, |
| "wind": 7.58267e00, |
| "dwi": 0.00000e00, |
| "10u_wave": 2.90337e-01, |
| "10v_wave": 1.26789e-01, |
| |
| "2d": 2.73876923e02, |
| "tcwv": 1.82695427e01, |
| "tcc": 6.7472446e-01, |
| "100u": 2.712852507829666e-02, |
| "100v": 1.9762367010116577e-01, |
| "sp": 9.6647375e04, |
| "lcc": 4.616926610469818e-01, |
| "mcc": 2.986297309398651e-01, |
| "hcc": 3.2651442289352417e-01, |
| "skt": 2.788576354980469e02, |
| "stl1": 2.7956781005859375e02, |
| "swvl1": 8.593245595693588e-02, |
| "ci": 1.1092530190944672e-01, |
| "scaled_sd": 1.3165518045425415e00, |
| |
| "i10fg": 9.003422737121582e00, |
| "blh": 6.290673217773438e02, |
| "uvb_1h": 6.96176e04, |
| "ssrd_1h": 5.889046e05, |
| "ttr_1h": -8.144387e05, |
| "scaled_tp_1h": 7.03631192445755e-02, |
| "scaled_sf_1h": 1.3227290473878384e-02, |
| |
| "insolation": 0.000000e00, |
| |
| "anor": 0.000000e00, |
| "isor": 0.000000e00, |
| "cvh": 0.000000e00, |
| "cl": 0.000000e00, |
| "dl": 0.000000e00, |
| "cvl": 0.000000e00, |
| "slor": 0.000000e00, |
| "slt_0": 0.000000e00, |
| "slt_1": 0.000000e00, |
| "slt_2": 0.000000e00, |
| "slt_3": 0.000000e00, |
| "slt_4": 0.000000e00, |
| "slt_5": 0.000000e00, |
| "slt_6": 0.000000e00, |
| "slt_7": 0.000000e00, |
| "sdfor": 0.000000e00, |
| "sdor": 0.000000e00, |
| "tvh_0": 0.000000e00, |
| "tvh_18": 0.000000e00, |
| "tvh_19": 0.000000e00, |
| "tvh_3": 0.000000e00, |
| "tvh_4": 0.000000e00, |
| "tvh_5": 0.000000e00, |
| "tvh_6": 0.000000e00, |
| "tvl_0": 0.000000e00, |
| "tvl_1": 0.000000e00, |
| "tvl_10": 0.000000e00, |
| "tvl_11": 0.000000e00, |
| "tvl_13": 0.000000e00, |
| "tvl_16": 0.000000e00, |
| "tvl_17": 0.000000e00, |
| "tvl_2": 0.000000e00, |
| "tvl_7": 0.000000e00, |
| "tvl_9": 0.000000e00, |
| } |
|
|
| scales: dict[str, float] = { |
| "z": 5.884467e04, |
| "lsm": 1.000000e00, |
| "slt": 7.000000e00, |
| "static_co": 1.000000e00, |
| "static_so2_log": 1.000000e00, |
| "static_ammonia_log": 1.000000e00, |
| "static_nox_log": 1.000000e00, |
| "static_so2": 1.000000e00, |
| "static_co_log": 1.000000e00, |
| "static_ammonia": 1.000000e00, |
| "static_nox": 1.000000e00, |
| "wmb": 1.000000e00, |
| "lat_mask": 1.000000e00, |
| "2t": 2.122036e01, |
| "10u": 5.547512e00, |
| "10v": 4.765339e00, |
| "msl": 1.332246e03, |
| "z_50": 5.875553e03, |
| "z_100": 5.510640e03, |
| "z_150": 5.823912e03, |
| "z_200": 5.820169e03, |
| "z_250": 5.536585e03, |
| "z_300": 5.091916e03, |
| "z_400": 4.150851e03, |
| "z_500": 3.353187e03, |
| "z_600": 2.695808e03, |
| "z_700": 2.136436e03, |
| "z_850": 1.470321e03, |
| "z_925": 1.228997e03, |
| "z_1000": 1.072307e03, |
| "u_50": 1.529281e01, |
| "u_100": 1.352611e01, |
| "u_150": 1.604335e01, |
| "u_200": 1.767630e01, |
| "u_250": 1.796710e01, |
| "u_300": 1.711917e01, |
| "u_400": 1.434276e01, |
| "u_500": 1.198419e01, |
| "u_600": 1.033421e01, |
| "u_700": 9.168821e00, |
| "u_850": 8.188043e00, |
| "u_925": 7.940808e00, |
| "u_1000": 6.141778e00, |
| "v_50": 7.058931e00, |
| "v_100": 7.479310e00, |
| "v_150": 9.571990e00, |
| "v_200": 1.188069e01, |
| "v_250": 1.338039e01, |
| "v_300": 1.334044e01, |
| "v_400": 1.122955e01, |
| "v_500": 9.181708e00, |
| "v_600": 7.803569e00, |
| "v_700": 6.871040e00, |
| "v_850": 6.264443e00, |
| "v_925": 6.470644e00, |
| "v_1000": 5.308203e00, |
| "t_50": 1.026284e01, |
| "t_100": 1.252901e01, |
| "t_150": 8.928709e00, |
| "t_200": 7.189547e00, |
| "t_250": 8.529282e00, |
| "t_300": 1.071679e01, |
| "t_400": 1.269102e01, |
| "t_500": 1.306447e01, |
| "t_600": 1.342046e01, |
| "t_700": 1.476523e01, |
| "t_850": 1.558880e01, |
| "t_925": 1.608798e01, |
| "t_1000": 1.713983e01, |
| "q_50": 3.571687e-07, |
| "q_100": 5.703754e-07, |
| "q_150": 3.794077e-06, |
| "q_200": 2.267534e-05, |
| "q_250": 7.446644e-05, |
| "q_300": 1.684361e-04, |
| "q_400": 5.078644e-04, |
| "q_500": 1.079294e-03, |
| "q_600": 1.769722e-03, |
| "q_700": 2.549169e-03, |
| "q_850": 4.112368e-03, |
| "q_925": 5.071058e-03, |
| "q_1000": 5.913548e-03, |
| "w_50": 2.031102e-02, |
| "w_100": 3.982386e-02, |
| "w_150": 7.233442e-02, |
| "w_200": 1.069730e-01, |
| "w_250": 1.431527e-01, |
| "w_300": 1.795840e-01, |
| "w_400": 2.365854e-01, |
| "w_500": 2.639899e-01, |
| "w_600": 2.768904e-01, |
| "w_700": 2.871598e-01, |
| "w_850": 2.900477e-01, |
| "w_925": 2.518952e-01, |
| "w_1000": 1.737641e-01, |
| "co_100": 1.388527e-07, |
| "co_1000": 1.811231e-05, |
| "co_150": 2.058023e-07, |
| "co_200": 2.400079e-07, |
| "co_250": 2.535403e-07, |
| "co_300": 2.594617e-07, |
| "co_400": 2.967479e-07, |
| "co_50": 3.389800e-08, |
| "co_500": 3.763697e-07, |
| "co_600": 6.095697e-07, |
| "co_700": 1.084358e-06, |
| "co_850": 3.787436e-06, |
| "co_925": 9.559742e-06, |
| "no2_100": 1.906433e-09, |
| "no2_1000": 1.931397e-07, |
| "no2_150": 2.686828e-09, |
| "no2_200": 1.976612e-09, |
| "no2_250": 1.646603e-09, |
| "no2_300": 1.461818e-09, |
| "no2_400": 1.455373e-09, |
| "no2_50": 1.184771e-09, |
| "no2_500": 1.834311e-09, |
| "no2_600": 1.515983e-08, |
| "no2_700": 2.405975e-08, |
| "no2_850": 9.421924e-08, |
| "no2_925": 1.378600e-07, |
| "no_100": 1.177009e-09, |
| "no_1000": 4.063175e-07, |
| "no_150": 1.501596e-09, |
| "no_200": 1.021288e-09, |
| "no_250": 7.026757e-10, |
| "no_300": 5.380344e-10, |
| "no_400": 4.131051e-10, |
| "no_50": 5.588770e-10, |
| "no_500": 1.647942e-09, |
| "no_600": 1.724509e-08, |
| "no_700": 2.791792e-08, |
| "no_850": 1.478467e-07, |
| "no_925": 2.512890e-07, |
| "go3_100": 3.564464e-06, |
| "go3_1000": 2.332223e-07, |
| "go3_150": 1.939059e-06, |
| "go3_200": 1.180154e-06, |
| "go3_250": 7.915061e-07, |
| "go3_300": 5.601596e-07, |
| "go3_400": 2.986021e-07, |
| "go3_50": 6.646655e-06, |
| "go3_500": 2.056183e-07, |
| "go3_600": 1.700490e-07, |
| "go3_700": 1.860511e-07, |
| "go3_850": 2.341993e-07, |
| "go3_925": 2.428035e-07, |
| "pm10": 6.662347e-06, |
| "pm1": 1.603840e-06, |
| "pm2p5": 2.336520e-06, |
| "so2_100": 1.636122e-09, |
| "so2_1000": 6.046371e-07, |
| "so2_150": 3.668129e-09, |
| "so2_200": 8.572589e-09, |
| "so2_250": 1.255334e-08, |
| "so2_300": 1.881883e-08, |
| "so2_400": 4.548033e-08, |
| "so2_50": 7.246278e-12, |
| "so2_500": 1.476261e-07, |
| "so2_600": 2.777121e-07, |
| "so2_700": 9.131610e-08, |
| "so2_850": 5.244673e-07, |
| "so2_925": 5.512759e-07, |
| "tcco": 5.207629e-03, |
| "tcno2": 6.141358e-05, |
| "tc_no": 2.865418e-05, |
| "gtco3": 9.886098e-03, |
| "tcso2": 1.574554e-04, |
| "swh": 1.38128e00, |
| "pp1d": 3.12242e00, |
| "mwd": 1.00000e00, |
| "mwp": 2.13391e00, |
| "shww": 1.23658e00, |
| "mdww": 1.00000e00, |
| "mpww": 2.08384e00, |
| "shts": 1.05034e00, |
| "mdts": 1.00000e00, |
| "mpts": 2.18806e00, |
| "swh1": 1.04481e00, |
| "mwd1": 1.00000e00, |
| "mwp1": 3.25726e00, |
| "swh2": 5.26008e-01, |
| "mwd2": 1.00000e00, |
| "mwp2": 4.38825e00, |
| "wind": 3.63699e00, |
| "dwi": 1.00000e00, |
| "10u_wave": 6.43740e00, |
| "10v_wave": 5.40235e00, |
| |
| "2d": 2.07733974e01, |
| "tcwv": 1.6429491e01, |
| "tcc": 3.65822792e-01, |
| "100u": 6.970843315124512e00, |
| "100v": 6.088253498077393e00, |
| "sp": 9.58669140625e03, |
| "lcc": 3.931414783000946e-01, |
| "mcc": 3.729339838027954e-01, |
| "hcc": 4.131758511066437e-01, |
| "skt": 2.2333234786987305e01, |
| "stl1": 2.144609832763672e01, |
| "swvl1": 1.4150479435920715e-01, |
| "ci": 2.9736775159835815e-01, |
| "scaled_sd": 3.0272440910339355e00, |
| |
| "i10fg": 4.655707359313965e00, |
| "blh": 4.7048150634765625e02, |
| "uvb_1h": 1.061103e05, |
| "ssrd_1h": 8.915032e05, |
| "ttr_1h": 1.771002e05, |
| "scaled_tp_1h": 1.8630525469779968e-01, |
| "scaled_sf_1h": 5.945565551519394e-02, |
| |
| "insolation": 1.000000e00, |
| |
| "anor": 1.000000e00, |
| "isor": 1.000000e00, |
| "cvh": 1.000000e00, |
| "cl": 1.000000e00, |
| "dl": 1.000000e00, |
| "cvl": 1.000000e00, |
| "slor": 1.000000e00, |
| "slt_0": 1.000000e00, |
| "slt_1": 1.000000e00, |
| "slt_2": 1.000000e00, |
| "slt_3": 1.000000e00, |
| "slt_4": 1.000000e00, |
| "slt_5": 1.000000e00, |
| "slt_6": 1.000000e00, |
| "slt_7": 1.000000e00, |
| "sdfor": 1.000000e00, |
| "sdor": 1.000000e00, |
| "tvh_0": 1.000000e00, |
| "tvh_18": 1.000000e00, |
| "tvh_19": 1.000000e00, |
| "tvh_3": 1.000000e00, |
| "tvh_4": 1.000000e00, |
| "tvh_5": 1.000000e00, |
| "tvh_6": 1.000000e00, |
| "tvl_0": 1.000000e00, |
| "tvl_1": 1.000000e00, |
| "tvl_10": 1.000000e00, |
| "tvl_11": 1.000000e00, |
| "tvl_13": 1.000000e00, |
| "tvl_16": 1.000000e00, |
| "tvl_17": 1.000000e00, |
| "tvl_2": 1.000000e00, |
| "tvl_7": 1.000000e00, |
| "tvl_9": 1.000000e00, |
| } |
|
|
