tracking.py

from ..tracking import TrackingTechnique
from abc import ABC, abstractmethod
from enstools.feature.identification._proto_gen import template_pb2
import google.protobuf
from dask import delayed, compute
import numpy as np
import xarray as xr


class TrackingTemplate(TrackingTechnique):
    """
    Template for a default "fall-back" tracking technique.
    There is also a more specific template_object_compare, which provides a template for techniques which are solely
    based on comparing the features of the object descriptions.
    """

    def track(self, timesteps: google.protobuf.pyext._message.RepeatedCompositeContainer, subset: xr.Dataset):
        """
        Implementation of track() for general tracking techniques.
        In this template, we iterate over the valid_times of this forecast/reanalysis and compare all objects of two
        consecutive timestamps. If they correspond according to some metric, their IDs should be added to the list.
        See also the template_object_compare for a more specific case.

        Parameters
        ----------
        timesteps : iterable of identification_pb2.Timestep
                The timesteps, a list of timestamps of this (re)forecast.
        subset : xarray.Dataset
                The subset (i.e., current member) to track.

        Returns
        -------
        List of tuples of object IDs, representing the pairwise connections between objects.
        """

        connections = []

        from enstools.misc import get_time_dim
        time_dim = get_time_dim(subset)

        for t in range(0, len(timesteps) - 1):
            t1 = timesteps[t]
            data_t1 = subset.sel({time_dim: t1.valid_time})
            t2 = timesteps[t + 1]
            data_t2 = subset.sel({time_dim: t2.valid_time})

            for o1 in t1.objects:
                for o2 in t2.objects:
                    # if same:
                    #   connections.append((o1.id, o2.id))
                    pass
        # TODO let user just return pairs of objects?
        # TODO provide overlap tracking as base method
        return connections