"""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) # Deal with rounding errors. # Return an integer representation whenever possible. if level % 1 == 0: level = int(level) # Replace the decimal separator with an underscore to avoid parameter name conflicts. 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, # Aurora 1.5 additional input surface variables "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, # Aurora 1.5 output-only surface variables "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 (not normalised during training) "insolation": 0.000000e00, # Static variables "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, # Aurora 1.5 additional input surface variables "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, # Aurora 1.5 output-only surface variables "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 (not normalised during training) "insolation": 1.000000e00, # Static variables "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, }