def fit_transport_plans(self, point_clouds, masses=None, rescaling=True,
maxerr=1e-6, maxiter=20, t_init=1.,
*args, **kwargs):
"""
Fits transport plans to point clouds.
Args:
point_clouds (dict): dictionary of point clouds
masses (dict): dictionary of masses assigned to each point cloud
(default will be uniform masses for each point cloud)
rescaling (bool): rescale or not the coordinates of point clouds
to fit in [0, 1]² to make computations easier
(rescaling undone when returning transport plans)
maxerr (float): threshold error under which Newton's algo stops
maxiter (int): threshold iteration under which Newton's algo stops
t_init (float): inital value of t for Newton's algorithm
"""
nb_clouds = len(point_clouds)
cloud_keys = list(point_clouds.keys())
# compute the transport plan of each cloud
self.potentials = {}
self.transport_plans = np.zeros((nb_clouds, 2 * self.grid_size**2))
for c in range(nb_clouds):
# get point cloud
sample = point_clouds[cloud_keys[c]].astype('float64')
N = len(sample)
if rescaling:
# rescale to [0, 1]²
sample_min, sample_max = np.min(sample, axis=0), np.max(sample, axis=0)
sample = (sample - sample_min) / (sample_max - sample_min)
# build target measure
if masses is not None:
nu = masses[cloud_keys[c]]
else:
nu = np.ones(N)/N
# compute OT between source and target (get Kantorovich potentials)
self.potentials[c] = newton_ot(self.domain, sample, nu, psi0=None, verbose=False,
maxerr=maxerr, maxiter=maxiter, t_init=t_init)
# build the *discretized* transport plan associated to these potentials
pd = PowerDiagram(sample, -self.potentials[c], self.domain)
img = pd.image_integrals([0, 0], [1, 1], [self.grid_size, self.grid_size])
img /= np.expand_dims(img[ :, :, 2], -1)
if rescaling:
# undo rescaling
img[:, :, :2] = (sample_max - sample_min) * img[:, :, :2] + sample_min
# save transport plan
self.transport_plans[c] = img[:, :, :2].flatten()
from pysdot.radial_funcs import RadialFuncInBall from pysdot import OptimalTransport from pysdot import PowerDiagram import pylab as plt import numpy as np import unittest domain = ConvexPolyhedraAssembly() domain.add_box([0, 0], [1, 1]) positions = [ [ 0.25, 0.5 ], [ 0.75, 0.5 ] ] weights = [ 0.25**2, 0.25**2 ] pd = PowerDiagram( positions, weights, domain, radial_func=RadialFuncInBall()) img = pd.image_integrals( [ 0, 0 ], [ 1, 1 ], [ 100, 100 ] ) img[ :, :, 0 ] *= 1e4 img[ :, :, 1 ] *= 1e4 # for d in range( 3 ): # plt.subplot( 1, 3, d + 1 ) # plt.imshow( img[ :, :, d ] ) # plt.colorbar() # plt.show() # w = 0.1 # print( np.pi * w ) # print( np.pi / 2 * w**2 )