Skip to content
Snippets Groups Projects
configuration.py 7.81 KiB
Newer Older
import xarray as xr
from enstools.feature.util.enstools_utils import get_latitude_dim, get_longitude_dim, get_vertical_dim
import logging
from enstools.feature.identification.pv_streamer.data_util import get_pv_field_name
import numpy as np

class PVConfig:
    # algorithm specific config
    def __init__(self):
        pass

    @staticmethod
    def get_default(dataset: xr.Dataset, vdim_name):
        config = PVConfig()

        # size of radius of stereographic projection (radius per hemisphere), therefore img=2rX4r
        config.proj_radius = 200
        # w_bar distance in km (Spr: 1500-2000)
        config.w_bar = 1000.0 # 2*wbar=maxwidth
        # min length of 2D streamer in km (only 2d mode)
        config.min_streamer_len = 1000.0
        # volume threshold of 3D streamers
        config.min_streamer_volume = 50000 # km2K
        # 3D bridges Kelvin Threshold
        config.bridge_K = 8 # xK bridges remove.
        config.exclude_isolated_cutoffs = False

        # preprocess remove holes so distance mask works
        config.remove_holes_maxarea = int(config.proj_radius)
        if vdim_name is None:
            # 2D
            config.dH = 1 * 125
            return config

        # 3D: distance between levels
        th_values = dataset.coords[vdim_name].data
        res_theta = abs(th_values[1] - th_values[0])

        # scale analysis. vertical extent in range 1000km x 1000km, typical height 10K -> stretch vertically by 100"km/K"
        # so extent in scale 1000x1000x1000. w_bar = 1000 as well. (max w 2000)
        config.dH = res_theta * 125

        return config


class DetectionConfig:
    # general detection config
    def __init__(self):
        pass

    @staticmethod
    def init_def_atm_table():
        # default atmosphere for ERAI from https://rda.ucar.edu/datasets/ds627.1/docs/Pressure_and_isentropic_levels/
        # translate z to theta and pressure so we have an approx z distance
        z_km = [0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5, 5.5, 6, 6.5, 7, 7.5, 8, 8.5, 9, 9.5, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
        th = [288.1, 289.8, 291.5, 293.2, 294.9, 296.7, 298.4, 300.3, 302.1, 304.1, 305.9, 308.0, 309.9, 312.0, 314.1, 316.2, 318.4, 320.5, 322.9, 325.2, 327.5, 332.4, 347.4, 363.4, 380.1, 397.5, 415.7, 434.8, 454.7, 475.6, 497.3]
        p = [1013.25, 954.61, 898.76, 845.59, 795.01, 746.91, 701.21, 657.80, 616.60, 577.52, 540.48, 505.39, 472.17, 440.75, 411.05, 382.99, 356.51, 331.54, 308.0, 285.84, 264.99, 226.99, 193.99, 165.79, 141.7, 121.11, 103.52, 88.497, 75.652, 64.674, 55.293]
        return np.asarray(z_km), np.asarray(th), np.asarray(p)

    @staticmethod
    def config_from_dataset(dataset: xr.Dataset):
        config = DetectionConfig()

        config.latitude_dim = get_latitude_dim(dataset)
        if config.latitude_dim is None:
            raise ValueError("Could not detect latitude dimension in dataset.")
        config.longitude_dim = get_longitude_dim(dataset)
        if config.longitude_dim is None:
            raise ValueError("Could not detect longitude dimension in dataset.")
        config.vertical_dim = get_vertical_dim(dataset)
        
        config.pv_dim = get_pv_field_name(dataset)
        if config.pv_dim is None:
            raise ValueError("Error, can't find PV in dataset")
            
        if config.vertical_dim not in dataset[config.pv_dim].dims:
            print("Executing for 2 dimensions.")
            config.dims = 2
        else:
            config.dims = 3

        if config.vertical_dim == 'th_level':
            config.vtype = 'theta'
        elif config.vertical_dim in ['level', 'lev']:
            config.vtype = 'pressure'
        else:
            config.vtype = None
            print("Unknown vtype, assume 2D")

        config.res_lon = abs(dataset.coords[config.longitude_dim].values[1] - dataset.coords[config.longitude_dim].values[0])
        config.res_lat = abs(dataset.coords[config.latitude_dim].values[1] - dataset.coords[config.latitude_dim].values[0])
        if config.vertical_dim is not None:
            config.res_z = abs(dataset.coords[config.vertical_dim].values[1] - dataset.coords[config.vertical_dim].values[0])
        else:
            config.res_z = -1

        # order of latitudes in dataset
        if (dataset.coords[config.latitude_dim][1].values - dataset.coords[config.latitude_dim][0].values) > 0:
            config.lat_order = 'asc'
        else:
            config.lat_order = 'desc'

        minLon = dataset.coords[config.longitude_dim].values[0]
        maxLon = dataset.coords[config.longitude_dim].values[-1]
        if minLon > maxLon:
            minLon, maxLon = maxLon, minLon

        minLat = dataset.coords[config.latitude_dim].values[0]
        maxLat = dataset.coords[config.latitude_dim].values[-1]
        if minLat > maxLat:
            minLat, maxLat = maxLat, minLat

        config.lon_min = minLon
        config.lon_max = maxLon
        config.lat_min = minLat
        config.lat_max = maxLat

        config.lon_indices_len = len(dataset.coords[config.longitude_dim].values)
        config.lat_indices_len = len(dataset.coords[config.latitude_dim].values)

        if maxLon - minLon < 360 - config.res_lon:
            raise ValueError("Error: Longitude range not 360 degrees, needed in this algorithm for feasible results.")

        if maxLat - minLat == 180:
            config.data_extent = 'global'
        elif minLat == 0 and maxLat == 90:
            config.data_extent = 'nh'
        elif minLat == -90 and maxLat == 0:
            config.data_extent = 'sh'
        else:
            raise ValueError("Unknown data extent: Either NH, SH or global.")

        logging.debug("Detected type: " + config.data_extent)

        if config.data_extent != 'nh':
            raise ValueError("Currently only northern hemispheric datasets usable. Use lat:0..90,lon:0..360 datasets.")

        config.tab_z, config.tab_th, config.tab_p = DetectionConfig.init_def_atm_table()

        ##############
        # from vertical data dim (theta,pressure...) to approx z height
        # for real-world volume
        ##############
        ######
        # scale analysis -> assume distance inbetween theta levels the same, need stretch for distances in algorithm
        # use mixed units for distances km^2*K, but stretch vertically so algorithm makes sense
        ######
        # includes config.alg.dH

            # 2D
            config.alg = PVConfig.get_default(dataset, config.vertical_dim)
            return config

        if config.vtype == 'theta':
            v_data = dataset.coords[config.vertical_dim].data

            idx = np.searchsorted(config.tab_th, v_data)
            if 0 in idx:
                print("todo idx == 0 ->> access -1")
            idx_bw = (v_data - config.tab_th[idx - 1]) / (config.tab_th[idx] - config.tab_th[idx - 1])
            config.v_data_z = config.tab_z[idx - 1] + idx_bw * (config.tab_z[idx] - config.tab_z[idx - 1])

        elif config.vtype == 'pressure':
            v_data = dataset.coords[config.vertical_dim].data

            idx = config.tab_p.size - np.searchsorted(config.tab_p[::-1], v_data, side='right') # tab_p is descending -> inverse and inverse indices
            if 0 in idx:
                print("todo idx == 0 ->> access -1")

            # where index == config.tab_p.size
            idx = np.where(idx == config.tab_p.size, config.tab_p.size - 1, idx)

            idx_bw = (v_data - config.tab_p[idx - 1]) / (config.tab_p[idx] - config.tab_p[idx - 1])
            config.v_data_z = config.tab_z[idx - 1] + idx_bw * (config.tab_z[idx] - config.tab_z[idx - 1])

        # distances between levels
        config.v_data_z_dist = np.abs(np.diff(config.v_data_z))

        if config.vtype != 'theta':
            print("only theta currently supported") # TODO
            exit()
        config.alg = PVConfig.get_default(dataset, config.vertical_dim)

        return config