diff options
-rw-r--r-- | examples/demo_barycenter_1D.py | 2 | ||||
-rw-r--r-- | ot/plot.py | 42 | ||||
-rw-r--r-- | ot/utils.py | 2 |
3 files changed, 43 insertions, 3 deletions
diff --git a/examples/demo_barycenter_1D.py b/examples/demo_barycenter_1D.py index 47b6be8..db1ff4a 100644 --- a/examples/demo_barycenter_1D.py +++ b/examples/demo_barycenter_1D.py @@ -37,7 +37,7 @@ for i in range(nbd): pl.plot(x,A[:,i]) pl.title('Distributions') -#%% barucenter computation +#%% barycenter computation # l2bary bary_l2=A.mean(1) @@ -9,7 +9,25 @@ from matplotlib import gridspec def plot1D_mat(a,b,M,title=''): - """ Plot matrix M with the source and target 1D distribution """ + """ Plot matrix M with the source and target 1D distribution + + Creates a subplot with the source distribution a on the left and + target distribution b on the tot. The matrix M is shown in between. + + + Parameters + ---------- + + a : np.array (na,) + Source distribution + b : np.array (nb,) + Target distribution + M : np.array (na,nb) + Matrix to plot + + + + """ na=M.shape[0] nb=M.shape[1] @@ -43,7 +61,27 @@ def plot1D_mat(a,b,M,title=''): def plot2D_samples_mat(xs,xt,G,thr=1e-8,**kwargs): - """ Plot matrix M in 2D with lines using alpha values""" + """ Plot matrix M in 2D with lines using alpha values + + Plot lines between source and target 2D samples with a color + proportional to the value of the matrix G between samples. + + + Parameters + ---------- + + xs : np.array (ns,2) + Source samples positions + b : np.array (nt,2) + Target samples positions + G : np.array (na,nb) + OT matrix + thr : float, optional + threshold above which the line is drawn + **kwargs : dict + paameters given to the plot functions (default color is black if nothing given) + + """ if ('color' not in kwargs) and ('c' not in kwargs): kwargs['color']='k' mx=G.max() diff --git a/ot/utils.py b/ot/utils.py index 2110c01..24f65a8 100644 --- a/ot/utils.py +++ b/ot/utils.py @@ -45,6 +45,7 @@ def dist(x1,x2=None,metric='sqeuclidean'): Returns ------- + M : np.array (n1,n2) distance matrix computed with given metric @@ -70,6 +71,7 @@ def dist0(n,method='lin_square'): Returns ------- + M : np.array (n1,n2) distance matrix computed with given metric |