Source code for NEDAS.assim_tools.localization

from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from NEDAS.core import Context

"""Localization funcs needs to be pure functions (they need to be used in numba njit)"""

registry = {
    'gaspari_cohn': ('distance_based', 'gaspari_cohn_func'),
    'step': ('distance_based', 'step_func'),
    'exponential': ('distance_based', 'exponential_func'),
}



def get_localization_funcs(c: Context) -> dict:
    local_funcs = {}
    assert c.config.localization_def is not None, "c.localization_def needs to be defined in the config file"
    for key in ['horizontal', 'vertical', 'temporal']:
        assert key in c.config.localization_def, f"{key} needs to be defined in c.localization_def"
        if c.config.localization_def[key]:
            if 'type' not in c.config.localization_def[key]:
                raise KeyError(f"'type' needs to be specified for c.localization['{key}']")
            local_funcs[key] = get_localization_func_component(c.config.localization_def[key]['type'])
        else:
            local_funcs[key] = None
    return local_funcs


def get_localization_func_component(localization_types):
    localization_types = localization_types.lower().split(',')

    # distance-based localization schemes
    if 'gaspari_cohn' in localization_types:
        from .distance_based import gaspari_cohn_func as local_func
    elif 'step' in localization_types:
        from .distance_based import step_func as local_func
    elif 'exponential' in localization_types:
        from .distance_based import exponential_func as local_func

    # # correlation based localization schemes
    # elif 'SER' in localization_types:
    #     from NEDAS.assim_tools.localization.SER import CorrelationBasedLocalization as Localization
    # elif 'NICE' in localization_types:
    #     from NEDAS.assim_tools.localization.NICE import
    else:
        raise ValueError(f"Unknown localization type {localization_types}")
    return local_func

__all__ = ['registry', 'get_localization_funcs']