|  | 
|  | 1 | +from ..observations import EddiesObservations as Model | 
|  | 2 | +from ..observations import GridDataset | 
|  | 3 | +from numpy import where, meshgrid, bincount, ones, unique, bool_, arange | 
|  | 4 | +import logging | 
|  | 5 | +from os import path | 
|  | 6 | + | 
|  | 7 | + | 
|  | 8 | +class CheltonTracker(Model): | 
|  | 9 | +    GROUND = GridDataset(path.join(path.dirname(__file__), 'mask_old.nc'), 'lon', 'lat') | 
|  | 10 | + | 
|  | 11 | +    @staticmethod | 
|  | 12 | +    def cost_function(records_in, records_out, distance): | 
|  | 13 | +        """We minimize on distance between two obs | 
|  | 14 | +        """ | 
|  | 15 | +        return distance | 
|  | 16 | + | 
|  | 17 | +    def mask_function(self, other): | 
|  | 18 | +        """We mask link with ellips and ratio | 
|  | 19 | +        """ | 
|  | 20 | +        # Compute Parameter of ellips | 
|  | 21 | +        y = self.basic_formula_ellips_major_axis( | 
|  | 22 | +            self.obs['lat'], | 
|  | 23 | +            degrees=True) | 
|  | 24 | +        # mask from ellips | 
|  | 25 | +        mask = self.shifted_ellipsoid_degrees_mask( | 
|  | 26 | +            other, | 
|  | 27 | +            minor=1.5, | 
|  | 28 | +            major=y) | 
|  | 29 | + | 
|  | 30 | +        # We check ratio | 
|  | 31 | +        other_amplitude, self_amplitude = meshgrid( | 
|  | 32 | +            other.obs['amplitude'], self.obs['amplitude']) | 
|  | 33 | +        other_radius, self_radius = meshgrid( | 
|  | 34 | +            other.obs['radius_e'], self.obs['radius_e']) | 
|  | 35 | +        ratio_amplitude = other_amplitude / self_amplitude | 
|  | 36 | +        ratio_radius = other_radius / self_radius | 
|  | 37 | +        mask *= \ | 
|  | 38 | +            (ratio_amplitude < 2.5) * \ | 
|  | 39 | +            (ratio_amplitude > 1 / 2.5) * \ | 
|  | 40 | +            (ratio_radius < 2.5) * \ | 
|  | 41 | +            (ratio_radius > 1 / 2.5) | 
|  | 42 | +        indexs_closest = where(mask) | 
|  | 43 | +        mask[indexs_closest] = self.across_ground(self.obs[indexs_closest[0]], other.obs[indexs_closest[1]]) | 
|  | 44 | +        return mask | 
|  | 45 | + | 
|  | 46 | +    @classmethod | 
|  | 47 | +    def across_ground(cls, record0, record1): | 
|  | 48 | +        i, j, d_pix = cls.GROUND.compute_pixel_path( | 
|  | 49 | +            x0=record0['lon'], | 
|  | 50 | +            y0=record0['lat'], | 
|  | 51 | +            x1=record1['lon'], | 
|  | 52 | +            y1=record1['lat'], | 
|  | 53 | +            ) | 
|  | 54 | + | 
|  | 55 | +        data = cls.GROUND.grid('mask')[i, j] | 
|  | 56 | +        i_ground = unique(arange(len(record0)).repeat(d_pix + 1)[data==1]) | 
|  | 57 | +        mask = ones(record1.shape, dtype='bool') | 
|  | 58 | +        mask[i_ground] = False | 
|  | 59 | +        return mask | 
|  | 60 | + | 
|  | 61 | +    def solve_function(self, cost_matrix): | 
|  | 62 | +        """Give the best link for each self obs | 
|  | 63 | +        """ | 
|  | 64 | +        return where(self.solve_first(cost_matrix, multiple_link=True)) | 
|  | 65 | + | 
|  | 66 | +    def post_process_link(self, other, i_self, i_other): | 
|  | 67 | +        """When two self obs use the same other obs, we keep the self obs | 
|  | 68 | +        with amplitude max | 
|  | 69 | +        """ | 
|  | 70 | +        if unique(i_other).shape[0] != i_other.shape[0]: | 
|  | 71 | +            nb_link = bincount(i_other) | 
|  | 72 | +            mask = ones(i_self.shape, dtype=bool_) | 
|  | 73 | +            for i in where(nb_link > 1)[0]: | 
|  | 74 | +                m = i == i_other | 
|  | 75 | +                multiple_in = i_self[m] | 
|  | 76 | +                i_keep = self.obs['amplitude'][multiple_in].argmax() | 
|  | 77 | +                m[where(m)[0][i_keep]] = False | 
|  | 78 | +                mask[m] = False | 
|  | 79 | + | 
|  | 80 | +            i_self = i_self[mask] | 
|  | 81 | +            i_other = i_other[mask] | 
|  | 82 | +        return i_self, i_other | 
0 commit comments