|
| 1 | +""" |
| 2 | +Network group process |
| 3 | +===================== |
| 4 | +""" |
| 5 | + |
| 6 | +import re |
| 7 | +from datetime import datetime |
| 8 | + |
| 9 | +from matplotlib import pyplot as plt |
| 10 | +from matplotlib.animation import FuncAnimation |
| 11 | +from matplotlib.colors import ListedColormap |
| 12 | +from numba import njit |
| 13 | +from numpy import arange, array, empty, ones |
| 14 | + |
| 15 | +from py_eddy_tracker import data |
| 16 | +from py_eddy_tracker.generic import flatten_line_matrix |
| 17 | +from py_eddy_tracker.observations.network import Network |
| 18 | +from py_eddy_tracker.observations.observation import EddiesObservations |
| 19 | + |
| 20 | + |
| 21 | +# %% |
| 22 | +class VideoAnimation(FuncAnimation): |
| 23 | + def _repr_html_(self, *args, **kwargs): |
| 24 | + """To get video in html and have a player""" |
| 25 | + content = self.to_html5_video() |
| 26 | + return re.sub( |
| 27 | + 'width="[0-9]*"\sheight="[0-9]*"', 'width="100%" height="100%"', content |
| 28 | + ) |
| 29 | + |
| 30 | + def save(self, *args, **kwargs): |
| 31 | + if args[0].endswith("gif"): |
| 32 | + # In this case gif is use to create thumbnail which are not use but consume same time than video |
| 33 | + # So we create an empty file, to save time |
| 34 | + with open(args[0], "w") as _: |
| 35 | + pass |
| 36 | + return |
| 37 | + return super().save(*args, **kwargs) |
| 38 | + |
| 39 | + |
| 40 | +# %% |
| 41 | +NETWORK_GROUPS = list() |
| 42 | + |
| 43 | + |
| 44 | +@njit(cache=True) |
| 45 | +def apply_replace(x, x0, x1): |
| 46 | + nb = x.shape[0] |
| 47 | + for i in range(nb): |
| 48 | + if x[i] == x0: |
| 49 | + x[i] = x1 |
| 50 | + |
| 51 | + |
| 52 | +# %% |
| 53 | +# Modified class to catch group process at each step in order to illustrate processing |
| 54 | +class MyNetwork(Network): |
| 55 | + def get_group_array(self, results, nb_obs): |
| 56 | + """With a loop on all pair of index, we will label each obs with a group |
| 57 | + number |
| 58 | + """ |
| 59 | + nb_obs = array(nb_obs, dtype="u4") |
| 60 | + day_start = nb_obs.cumsum() - nb_obs |
| 61 | + gr = empty(nb_obs.sum(), dtype="u4") |
| 62 | + gr[:] = self.NOGROUP |
| 63 | + |
| 64 | + id_free = 1 |
| 65 | + for i, j, ii, ij in results: |
| 66 | + gr_i = gr[slice(day_start[i], day_start[i] + nb_obs[i])] |
| 67 | + gr_j = gr[slice(day_start[j], day_start[j] + nb_obs[j])] |
| 68 | + # obs with no groups |
| 69 | + m = (gr_i[ii] == self.NOGROUP) * (gr_j[ij] == self.NOGROUP) |
| 70 | + nb_new = m.sum() |
| 71 | + gr_i[ii[m]] = gr_j[ij[m]] = arange(id_free, id_free + nb_new) |
| 72 | + id_free += nb_new |
| 73 | + # associate obs with no group with obs with group |
| 74 | + m = (gr_i[ii] != self.NOGROUP) * (gr_j[ij] == self.NOGROUP) |
| 75 | + gr_j[ij[m]] = gr_i[ii[m]] |
| 76 | + m = (gr_i[ii] == self.NOGROUP) * (gr_j[ij] != self.NOGROUP) |
| 77 | + gr_i[ii[m]] = gr_j[ij[m]] |
| 78 | + # case where 2 obs have a different group |
| 79 | + m = gr_i[ii] != gr_j[ij] |
| 80 | + if m.any(): |
| 81 | + # Merge of group, ref over etu |
| 82 | + for i_, j_ in zip(ii[m], ij[m]): |
| 83 | + g0, g1 = gr_i[i_], gr_j[j_] |
| 84 | + apply_replace(gr, g0, g1) |
| 85 | + NETWORK_GROUPS.append((i, j, gr.copy())) |
| 86 | + return gr |
| 87 | + |
| 88 | + |
| 89 | +# %% |
| 90 | +# Movie period |
| 91 | +t0 = (datetime(2005, 5, 1) - datetime(1950, 1, 1)).days |
| 92 | +t1 = (datetime(2005, 6, 1) - datetime(1950, 1, 1)).days |
| 93 | + |
| 94 | +# %% |
| 95 | +# Get data from period and area |
| 96 | +e = EddiesObservations.load_file(data.get_path("Anticyclonic_seg.nc")) |
| 97 | +e = e.extract_with_mask((e.time >= t0) * (e.time < t1)).extract_with_area( |
| 98 | + dict(llcrnrlon=25, urcrnrlon=35, llcrnrlat=31, urcrnrlat=37.5) |
| 99 | +) |
| 100 | +# %% |
| 101 | +# Reproduce individual daily identification(for demonstration) |
| 102 | +EDDIES_BY_DAYS = list() |
| 103 | +for i, b0, b1 in e.iter_on("time"): |
| 104 | + EDDIES_BY_DAYS.append(e.index(i)) |
| 105 | +# need for display |
| 106 | +e = EddiesObservations.concatenate(EDDIES_BY_DAYS) |
| 107 | + |
| 108 | +# %% |
| 109 | +# Run network building group to intercept every step |
| 110 | +n = MyNetwork.from_eddiesobservations(EDDIES_BY_DAYS, window=7) |
| 111 | +_ = n.group_observations(minimal_area=True) |
| 112 | + |
| 113 | + |
| 114 | +# %% |
| 115 | +def update(frame): |
| 116 | + i_current, i_match, gr = NETWORK_GROUPS[frame] |
| 117 | + current = EDDIES_BY_DAYS[i_current] |
| 118 | + x = flatten_line_matrix(current.contour_lon_e) |
| 119 | + y = flatten_line_matrix(current.contour_lat_e) |
| 120 | + current_contour.set_data(x, y) |
| 121 | + match = EDDIES_BY_DAYS[i_match] |
| 122 | + x = flatten_line_matrix(match.contour_lon_e) |
| 123 | + y = flatten_line_matrix(match.contour_lat_e) |
| 124 | + matched_contour.set_data(x, y) |
| 125 | + groups.set_array(gr) |
| 126 | + txt.set_text(f"Day {i_current} match with day {i_match}") |
| 127 | + s = 80 * ones(gr.shape) |
| 128 | + s[gr == 0] = 4 |
| 129 | + groups.set_sizes(s) |
| 130 | + |
| 131 | + |
| 132 | +# %% |
| 133 | +# Anim |
| 134 | +# ---- |
| 135 | +fig = plt.figure(figsize=(16, 9), dpi=50) |
| 136 | +ax = fig.add_axes([0, 0, 1, 1]) |
| 137 | +ax.set_aspect("equal"), ax.grid(), ax.set_xlim(26, 34), ax.set_ylim(32, 36.5) |
| 138 | +groups = ax.scatter( |
| 139 | + e.lon, |
| 140 | + e.lat, |
| 141 | + c=NETWORK_GROUPS[0][2], |
| 142 | + cmap=ListedColormap(["gray", *e.COLORS[:-1]], name="from_list", N=30), |
| 143 | + vmin=0, |
| 144 | + vmax=30, |
| 145 | +) |
| 146 | +current_contour = ax.plot([], [], "k", lw=2, label="Current contour")[0] |
| 147 | +matched_contour = ax.plot([], [], "r", lw=1, ls="--", label="Candidate contour")[0] |
| 148 | +txt = ax.text(31.5, 35.5, "", fontsize=25) |
| 149 | +ax.legend(fontsize=25) |
| 150 | +ani = VideoAnimation(fig, update, frames=len(NETWORK_GROUPS), interval=220) |
0 commit comments