@@ -259,3 +259,62 @@ def bivariate(self, data, i=None, j=None):
259259 data .barycenter [self ._mode ][(j , i )] = data .barycenter [self ._mode ][(i , j )]
260260
261261 return self ._statfn (bc )
262+
263+
264+ class GromovWasserstainTau (Undirected , Unsigned ):
265+ """Gromov-Wasserstain distance (GWTau)"""
266+
267+ name = "Gromov-Wasserstain Distance"
268+ identifier = "gwtau"
269+ labels = ["unsigned" , "distance" , "unordered" , "nonlinear" , "undirected" ]
270+
271+ @staticmethod
272+ def vec_geo_dist (x ):
273+ diffs = np .diff (x , axis = 0 )
274+ distances = np .linalg .norm (diffs , axis = 1 )
275+ return np .cumsum (distances )
276+
277+ @staticmethod
278+ def wass_sorted (x1 , x2 ):
279+ x1 = np .sort (x1 )[::- 1 ] # sort in descending order
280+ x2 = np .sort (x2 )[::- 1 ]
281+
282+ if len (x1 ) == len (x2 ):
283+ res = np .sqrt (np .mean ((x1 - x2 ) ** 2 ))
284+ else :
285+ N , M = len (x1 ), len (x2 )
286+ i_ratios = np .arange (1 , N + 1 ) / N
287+ j_ratios = np .arange (1 , M + 1 ) / M
288+
289+
290+ min_values = np .minimum .outer (i_ratios , j_ratios )
291+ max_values = np .maximum .outer (i_ratios - 1 / N , j_ratios - 1 / M )
292+
293+ lam = np .where (min_values > max_values , min_values - max_values , 0 )
294+
295+ diffs_squared = (x1 [:, None ] - x2 ) ** 2
296+ my_sum = np .sum (lam * diffs_squared )
297+
298+ res = np .sqrt (my_sum )
299+
300+ return res
301+
302+ @staticmethod
303+ def gwtau (xi , xj ):
304+ timei = np .arange (len (xi ))
305+ timej = np .arange (len (xj ))
306+ traji = np .column_stack ([timei , xi ])
307+ trajj = np .column_stack ([timej , xj ])
308+
309+ vi = GromovWasserstainTau .vec_geo_dist (traji )
310+ vj = GromovWasserstainTau .vec_geo_dist (trajj )
311+ gw = GromovWasserstainTau .wass_sorted (vi , vj )
312+
313+ return gw
314+
315+ @parse_bivariate
316+ def bivariate (self , data , i = None , j = None ):
317+ x , y = data .to_numpy ()[[i , j ]]
318+ # insert compute SPI code here (computes on x and y)
319+ stat = self .gwtau (x , y )
320+ return stat
0 commit comments