def sinkhorn_display_plan(gamma, X, Y, fig=None, tol=0.6):
fig = fig or plt.figure()
ax = plt.gca()
plot_2d_cloud(X, cmap="r", fig=fig)
plot_2d_cloud(Y, cmap="b", fig=fig)
if gamma is not None:
gamma /= gamma.max()
from matplotlib import collections as mc
lin, col = np.nonzero(gamma > tol)
src, dst = X[lin], Y[col]
src = src[:, 0:2]
dst = dst[:, 0:2]
ps = []
colors = []
for i in range(len(src)):
ps.append([src[i], dst[i]])
colors.append([0, 1, 0, 1])
lc = mc.LineCollection(ps, colors=colors)
ax.add_collection(lc)
plt.show()
def plot_empty_cells(mu, Y, ind, title=""):
fig = plt.figure()
if is_2d(Y):
plot_2d_cloud(Y[ind], fig=fig)
plot_2d_tri(mu.vertices, mu.triangles, fig=fig, c="r")
else:
plot_cloud(Y[ind])
plot_tri(mu.vertices, mu.triangles, fig=fig, c="r")
plt.title("{} empty cells{}".format(len(ind), title))
def plot(self, fig=None, as_img=False, colorbar=False):
import matplotlib.pyplot as plt
fig = fig or plt.figure()
if as_img:
# Only work with square densities
root = math.sqrt(self.values.size)
assert int(root + 0.5) ** 2 == self.values.size, "Plot as image only works with square densities"
n = int(root)
img = self.values.reshape(n, n)
plt.imshow(img, interpolation="none")
else:
from sdot.core.common import plot_2d_tri_func, plot_2d_cloud, plot_tri, plot_cloud, is_2d
if is_2d(self.vertices):
plot_2d_cloud(self.vertices, cmap=self.values, colorbar=colorbar, fig=fig)
# plot_2d_tri_func(self.vertices, self.triangles, self.values, fig=fig)
else:
plot_cloud(self.vertices, cmap=self.values, colorbar=colorbar, fig=fig)
mu_flat /= mu_flat.sum()
nu_flat /= nu_flat.sum()
fig = plt.figure()
if len(sys.argv) == 1:
ax = fig.add_subplot("111", projection="3d")
ax.plot_surface(X_mu, Y_mu, mu, color="r")
ax.plot_surface(X_nu, Y_nu, nu, color="b")
# import matplotlib as mpl
# legend1 = mpl.lines.Line2D([0],[0], linestyle="none", c='b', marker='o')
# legend2 = mpl.lines.Line2D([0],[0], linestyle="none", c='r', marker='o')
# ax.legend([legend1, legend2], ["Source", 'Target'], numpoints = 1)
elif len(sys.argv) in [3, 4]:
if is_2d(P_mu):
plot_2d_cloud(P_mu, fig=fig, cmap="r", labels=["source"])
plot_2d_cloud(P_nu, fig=fig, cmap="b", labels=["target"])
else:
plot_cloud(P_mu, fig=fig, cmap="r", labels=["source"])
plot_cloud(P_nu, fig=fig, cmap="b", labels=["target"])
ax = plt.gca()
ax.set_aspect("equal")
# plt.legend()
# plt.show()
# plt.imshow(M)
# plt.show()
# Normal
if discrete_init["normal"]:
with timeit("NORMAL") as timer:
# alpha = (b - a) / (M - m)
# beta = a - alpha * m
# return alpha * x + beta
# axis = 0
# m, M = np.min(Y[:, axis]), np.max(Y[:, axis])
# # min_nu = 0.1 # fandisk 1k
# min_nu = 0.3 # torus 1k
# nu = lambda y: linear_target(y[axis], min_nu, 1, m, M)
nu = np.apply_along_axis(nu, 1, Y)
fig = plt.figure()
if is_2d(X):
ax = plt.gca()
ax.set_aspect("equal")
plot_2d_cloud(Y, fig=fig, cmap="b", s=10, labels=["target"])
# plot_2d_hull(X[:, 0:2], fig=fig, c="r")
plot_2d_tri(X, T_X, fig=fig, c="r", label="source")
else:
ax = plt.gca(projection="3d")
plot_cloud(Y, ax=ax, fig=fig, cmap="b",labels=["target"], s=15)
plot_tri(X, T_X, fig=fig, color="r", label="source")
# plt.legend()
plt.show() # TODO
## Local
if tests_init["local"]:
start_local = time.time()
psi0_local = init_optimal_transport_3(mu, Y, nu, method="local",
eps_init=eps_init, display_empty=False,
save_results=save_results)