def geodesic(nTime,
nameFileD,
mub0,
mub1,
cCongestion,
eps,
Nit,
detailStudy=False,
verbose=False,
tol=1e-6):
"""Implementation of the algorithm for computing the geodesic in the Wasserstein space.
Arguments.
nTime: Number of discretization points in time
nameFileD: Name of the .off file where the triangle mesh is stored
mub0: initial probability distribution
mub1: final probability distribution
cCongestion: constant for the intensity of the regularization
(alpha in the article)
eps: regularization parameter for the linear system inversion
Nit: Number of iterations
detailStudy: True if the value of the objective functional (i.e. the
Lagrangian) and the residuals are computed at every time step (slow),
false if computed every 10 iterations (fast)
Output.
phi,mu,A,E,B: values of the relevant quantities in the interpolation,
cf. the article for more details (Note: E is the momentum, denoted by m
in the article) objectiveValue: evolution (in term of the number of
iterations of the ADMM) of the objective value, ie the Lagrangian
primalResidual,dualResidual: evolution (in term of the number of
iterations of the ADMM) of primal and dual residuals, cf the article for
details of their computation.
"""
startImport = time.time()
# Boolean value saying wether there is congestion or not
isCongestion = cCongestion >= 10**(-10)
if verbose:
print(15 * "-" + " Parameters for the computation of the geodesic " +
15 * "-")
print("Number of discretization points in time: {}".format(nTime))
print("Name of the mesh file: {}".format(nameFileD))
if isCongestion:
print("Congestion parameter: {}\n".format(cCongestion))
else:
print("No regularization\n")
# Time domain: staggered grid
xTimeS = np.linspace(0, 1, nTime + 1)
# Step Time
DeltaTime = xTimeS[1] - xTimeS[0]
# Time domain: centered grid
xTimeC = np.linspace(DeltaTime / 2, 1 - DeltaTime / 2, nTime)
# Domain D: call the routines
Vertices, Triangles, Edges = read_off.readOff(nameFileD)
areaTriangles, angleTriangles, baseFunction = geometricQuantities(
Vertices, Triangles, Edges)
gradientDMatrix, divergenceDMatrix, LaplacianDMatrix = geometricMatrices(
Vertices, Triangles, Edges, areaTriangles, angleTriangles,
baseFunction)
originTriangles, areaVertices, vertexTriangles = trianglesToVertices(
Vertices, Triangles, areaTriangles)
# Size of the domain D
nVertices = Vertices.shape[0]
nTriangles = Triangles.shape[0]
nEdges = Edges.shape[0]
# Vectorized quantities
# Enable to call in parallel on [0,1] x D something which is only defined on D.
# Vectorized arrays
areaVectorized = np.kron(np.kron(np.ones(6 * nTime), areaTriangles),
np.ones(3)).reshape(nTime, 2, 3, nTriangles, 3)
areaVerticesGlobal = np.kron(np.ones(nTime), areaVertices).reshape(
(nTime, nVertices))
areaVerticesGlobalStaggerred = np.kron(np.ones(nTime + 1),
areaVertices).reshape(
(nTime + 1, nVertices))
# Vectorized matrices
vertexTrianglesGlobal = scsp.kron(scsp.eye(nTime), vertexTriangles)
originTrianglesGlobal = scsp.kron(scsp.eye(nTime), originTriangles)
# Data structure with all the relevant informations. To be used as an argument of functions
geomDic = {
"nTime": nTime,
"DeltaTime": DeltaTime,
"Vertices": Vertices,
"Triangles": Triangles,
"Edges": Edges,
"areaTriangles": areaTriangles,
"gradientDMatrix": gradientDMatrix,
"divergenceDMatrix": divergenceDMatrix,
"LaplacianDMatrix": LaplacianDMatrix,
"originTriangles": originTriangles,
"areaVertices": areaVertices,
"vertexTriangles": vertexTriangles,
"nVertices": nVertices,
"nTriangles": nTriangles,
"nEdges": nEdges,
"areaVerticesGlobal": areaVerticesGlobal,
"areaVectorized": areaVectorized,
}
# Build the Laplacian matrix in space time and its inverse
LaplacianInvert = laplacian_inverse.buildLaplacianMatrix(geomDic, eps)
# Variable initialization
# Primal variable phi.
# Staggerred in Time, defined on the vertices of D
phi = np.zeros((nTime + 1, nVertices))
# Lagrange multiplier associated to mu.
# Centered in Time and lives on the vertices of D
mu = np.zeros((nTime, nVertices))
# Momentum E, lagrange mutliplier.
# Centered in Time, the second component is indicating on which side of
# the temporal it comes from. Third component indicates the origine of
# the triangle on which it is. Fourth component is the triangle. Last
# component corresponds to the fact that we are looking at a vector of R^3.
E = np.zeros((nTime, 2, 3, nTriangles, 3))
# Primal Variable A, corresponds to d_t phi.
# Same staggering pattern as mu
A = np.zeros((nTime, nVertices))
# Primal variable B, same pattern as E
B = np.zeros((nTime, 2, 3, nTriangles, 3))
# Lagrange multiplier associated to the congestion. If there is no
# congestion, the value of this parameter will always stay at 0.
lambdaC = np.zeros((nTime, nVertices))
# Making sure the boundary values are normalized
mub0 /= np.sum(mub0)
mub1 /= np.sum(mub1)
# Build the boundary term
BT = np.zeros((nTime + 1, nVertices))
BT[0, :] = -mub0
BT[-1, :] = mub1
# ADMM iterations
# Value of the "augmentation parameter" for the augmented Lagragian problem (update dynamically)
r = 1.
# Initialize the array which will contain the values of the objective functional
if detailStudy:
objectiveValue = np.zeros(3 * Nit)
else:
objectiveValue = np.zeros((Nit // 10))
# Initialize the arry which will contain the residuals
primalResidual = np.zeros(Nit)
dualResidual = np.zeros(Nit)
# Main Loop
for counterMain in range(Nit):
if verbose:
print(30 * "-" + " Iteration " + str(counterMain + 1) + " " +
30 * "-")
if detailStudy:
objectiveValue[3 * counterMain] = objectiveFunctional(
phi, mu, A, E, B, lambdaC, BT, geomDic, r, cCongestion,
isCongestion)
elif (counterMain % 10) == 0:
objectiveValue[counterMain // 10] = objectiveFunctional(
phi, mu, A, E, B, lambdaC, BT, geomDic, r, cCongestion,
isCongestion)
# Laplace problem
startLaplace = time.time()
# Build the RHS
RHS = np.zeros((nTime + 1, nVertices))
RHS -= BT * nTime
RHS -= gradATime(mu, geomDic)
RHS += np.multiply(r * gradATime(A + lambdaC, geomDic),
areaVerticesGlobalStaggerred / 3)
# We take the adjoint wrt tp the scalar product weighted by areas, hence the multiplication by areaVectorized
RHS -= divergenceD(E, geomDic)
RHS += r * divergenceD(np.multiply(B, areaVectorized), geomDic)
# Solve the system
phi = 1. / r * LaplacianInvert(RHS)
endLaplace = time.time()
if verbose:
print("Solving the Laplace system: {}s.".format(
round(endLaplace - startLaplace, 2)))
if detailStudy:
objectiveValue[3 * counterMain + 1] = objectiveFunctional(
phi, mu, A, E, B, lambdaC, BT, geomDic, r, cCongestion,
isCongestion)
# Projection over a convex set ---------------------------------------------------------
# It projects on the set A + 1/2 |B|^2 <= 0. We reduce to a 1D projection, then use a Newton method with a fixed number of iteration.
startProj = time.time()
# Computing the derivatives of phi
dTphi = gradTime(phi, geomDic)
dDphi = gradientD(phi, geomDic)
# Computing what there is to project
toProjectA = dTphi + 3. / r * np.divide(mu, areaVerticesGlobal)
toProjectB = dDphi + 1. / r * np.divide(E, areaVectorized)
# bSquaredArray will contain
# (sum_{a ~ v} |a| |B_{a,v}|**2) / (4*sum_{a ~ v} |a|)
# for each vertex v
bSquaredArray = np.zeros((nTime, nVertices))
# square and sum to compute in account the eulcidean norm and the temporal average
squareAux = np.sum(np.square(toProjectB), axis=(1, 4))
# average wrt triangles
bSquaredArray = originTrianglesGlobal.dot(
squareAux.reshape(nTime * 3 * nTriangles)).reshape(
(nTime, nVertices))
# divide by the sum of the areas of the neighboring triangles
bSquaredArray = np.divide(bSquaredArray, 4 * areaVerticesGlobal)
# Value of the objective functional. For the points not in the convex, we want it to vanish.
projObjective = toProjectA + bSquaredArray
# projDiscriminating is 1 is the point needs to be projected, 0 if it is already in the convex
projDiscriminating = (np.greater(
projObjective, 10**(-16) * np.ones(
(nTime, nVertices)))).astype(float)
# Newton method iteration
# Value of the Lagrange multiplier. Initialized at 0, not updated if already in the convex set
xProj = np.zeros((nTime, nVertices))
counterProj = 0
# Newton's loop
while np.max(projObjective) > 10**(-10) and counterProj < 20:
# Objective functional
projObjective = (toProjectA + 6 * (1. + cCongestion * r) * xProj +
np.divide(bSquaredArray, np.square(1 - xProj)))
# Derivative of the ojective functional
dProjObjective = 6 * (1. + cCongestion * r) - 2. * np.divide(
bSquaredArray, np.power(xProj - 1, 3))
# Update of xProj
xProj -= np.divide(np.multiply(projDiscriminating, projObjective),
dProjObjective)
counterProj += 1
# Update of A
A = toProjectA + 6 * (1. + cCongestion * r) * xProj
# Update of lambda
lambdaC = -6 * cCongestion * r * xProj
# Transfer xProj, which is defined on vertices into something which is defined on triangles
xProjTriangles = np.kron(
vertexTrianglesGlobal.dot(xProj.reshape(
nTime * nVertices)).reshape((nTime, 3, nTriangles)),
np.ones(3),
).reshape((nTime, 3, nTriangles, 3))
# Update of B
B[:, 0, :, :, :] = np.divide(toProjectB[:, 0, :, :, :],
1. - xProjTriangles)
B[:, 1, :, :, :] = np.divide(toProjectB[:, 1, :, :, :],
1. - xProjTriangles)
# Print the info
endProj = time.time()
if verbose:
print("Pointwise projection: {}s.".format(
str(round(endProj - startProj, 2))))
print("{} iterations needed; error committed: {}.".format(
counterProj, np.max(projObjective)))
if detailStudy:
objectiveValue[3 * counterMain + 2] = objectiveFunctional(
phi, mu, A, E, B, lambdaC, BT, geomDic, r, cCongestion,
isCongestion)
# Gradient descent in (E,muTilde), i.e. in the dual
# No need to recompute the derivatives of phi
# Update for mu
mu -= r / 3 * np.multiply(areaVerticesGlobal, A + lambdaC - dTphi)
# Update for E
E -= r * np.multiply(areaVectorized, B - dDphi)
# Compute the residuals
# For the primal residual, just what was updated in the dual
primalResidual[counterMain] = sqrt((scalarProductFun(
A + lambdaC - dTphi,
np.multiply(A + lambdaC - dTphi, areaVerticesGlobal / 3.0),
geomDic,
) + scalarProductTriangles(B - dDphi, B - dDphi, geomDic)) /
np.sum(areaTriangles))
# For the residual, take the RHS of the Laplace system and conserve only
# BT and the dual variables mu, E
dualResidualAux = np.zeros((nTime + 1, nVertices))
dualResidualAux += BT / DeltaTime
dualResidualAux += gradATime(mu, geomDic)
dualResidualAux += divergenceD(E, geomDic)
dualResidual[counterMain] = r * sqrt(
scalarProductFun(
dualResidualAux,
np.multiply(dualResidualAux,
areaVerticesGlobalStaggerred / 3.0),
geomDic,
) / np.sum(areaTriangles))
# Break early if residuals are small
if primalResidual[counterMain] < tol and dualResidual[
counterMain] < tol:
break
# Update the parameter r
# cf. Boyd et al. for an explanantion of the rule
if primalResidual[counterMain] >= 10 * dualResidual[counterMain]:
r *= 2
elif 10 * primalResidual[counterMain] <= dualResidual[counterMain]:
r /= 2
# Printing some results
if verbose:
if detailStudy:
print("Maximizing in phi, should go up: {}".format(
objectiveValue[3 * counterMain + 1] -
objectiveValue[3 * counterMain]))
print("Maximizing in A,B, should go up: {}".format(
objectiveValue[3 * counterMain + 2] -
objectiveValue[3 * counterMain + 1]))
if counterMain >= 1:
print("Dual update: should go down: {}".format(
objectiveValue[3 * counterMain] -
objectiveValue[3 * counterMain - 1]))
print("Values of phi: {}/{}\n".format(np.max(phi), np.min(phi)))
print("Values of A: {}/{}\n".format(np.max(A), np.min(A)))
print("Values of mu: {}/{}\n".format(np.max(mu), np.min(mu)))
print("Values of E: {}/{}\n".format(np.max(E), np.min(E)))
if isCongestion:
print("Congestion")
print(
scalarProductFun(lambdaC, mu, geomDic) - 1 /
(2. * cCongestion) * scalarProductFun(
lambdaC, np.multiply(lambdaC, areaVerticesGlobal /
3.), geomDic))
print(cCongestion / 2. *
np.sum(np.divide(np.square(mu), 1 / 3. * areaVertices)) /
nTime)
# Print some informations at the end of the loop
if verbose:
print("Final value of the augmenting parameter: {}".format(r))
# Integral of mu wrt space (depends on time), should sum up to 1.
intMu = np.sum(mu, axis=(-1))
print("Minimal and maximal value of int mu: {}/{}".format(
np.min(intMu), np.max(intMu)))
print("Maximal and minimal value of mu: {}/{}".format(
np.min(mu), np.max(mu)))
dTphi = gradTime(phi, geomDic)
dDphi = gradientD(phi, geomDic)
print("Agreement between nabla_t,D and (A,B)")
print(np.max(np.abs(dTphi - A)))
print(np.max(np.abs(dDphi - B)))
endProgramm = time.time()
print("Primal/dual residuals at end: {}/{}".format(
primalResidual[counterMain], dualResidual[counterMain]))
print("Congestion norm: {}".format(
np.linalg.norm(lambdaC - cCongestion * (mu / (areaVertices / 3)))))
print("Objective value at end: {}".format(objectiveValue[counterMain //
10]))
print("Total number of iterations: {}".format(counterMain))
print("Total time taken by the computation of the geodesic: {}".format(
round(endProgramm - startImport, 2)))
return phi, mu, A, E, B, objectiveValue, primalResidual, dualResidual
# Value for the congestion parameter (alpha in the article)
cCongestion = 0.1
# -----------------------------------------------------------------------------------------------
# Read the .off file
# -----------------------------------------------------------------------------------------------
# Extract Vertices, Triangles, Edges
Vertices, Triangles, Edges = read_off.readOff(nameFileD)
# Compute areas of Triangles
areaTriangles, angleTriangles, baseFunction = surface_pre_computations.geometricQuantities(
Vertices, Triangles, Edges)
# Compute the areas of the Vertices
originTriangles, areaVertices, vertexTriangles = surface_pre_computations.trianglesToVertices(
Vertices, Triangles, areaTriangles)
# -----------------------------------------------------------------------------------------------
# Define the boundary conditions
# -----------------------------------------------------------------------------------------------
nVertices = Vertices.shape[0]
mub0 = np.zeros(nVertices)
mub1 = np.zeros(nVertices)
lengthScale = 0.1
# Center of the "blobs" for mub0 and mub1
center0 = Vertices[4492, :]
center1 = Vertices[4225, :]
def __init__(self,
vertices,
edges,
triangles,
boundaries=None,
normal=None):
self.vertices = vertices
self.edges = edges
self.triangles = triangles
self.n_vertices = vertices.shape[0]
self.n_edges = edges.shape[0]
self.n_triangles = triangles.shape[0]
t_geom = time()
print("Building geometric quantities...", end=' ')
self.areaTriangles, self.angleTriangles, self.baseFunction = geometricQuantities(
vertices, triangles, edges)
print(f"(in {time() - t_geom:.1f}s)")
t_mat = time()
print("Building geometric matrices...", end=' ')
self.gradient, self.divergence, self.laplacian = geometricMatrices(
vertices, triangles, edges, self.areaTriangles,
self.angleTriangles, self.baseFunction)
print(f"(in {time() - t_mat:.1f}s)")
t_TtoV = time()
print("Building triangles to vertices...", end=' ')
self.originTriangles, self.areaVertices, self.vertexTriangles = trianglesToVertices(
self.vertices, self.triangles, self.areaTriangles)
print(f"(in {time() - t_TtoV:.1f}s)")
t_mean = time()
print("Building adjacency matrix...", end=' ')
self.mean_triangles = mean_of_triangles(vertices, triangles,
self.areaTriangles)
print(f"(in {time() - t_mean:.1f}s)")
# Allow to sum on R3 components
self.triangles_R3_to_R = np.kron(np.eye(self.n_triangles),
np.ones((3, 1)))
self.has_boundaries = False
if boundaries is not None or normal is not None:
assert boundaries is not None and normal is not None, "Boundaries and normal: provide either both or " \
"neither"
assert boundaries.ndim == 1 and normal.ndim == 2 and normal.shape[
1] == 3
assert boundaries.shape[0] == normal.shape[0], f"Boundaries and normal should have same length. " \
f"Got {boundaries.shape[0]}, {normal.shape[0]}."
self.has_boundaries = True
self.boundaries = boundaries
self.normal = normal
self.n_boundaries = boundaries.shape[0] if self.has_boundaries else 0
# self.normal_product makes the product w.r.t the normal.
if self.has_boundaries:
k = np.kron(np.eye(self.n_boundaries), np.ones(3))
k[k == 1] = self.normal.flatten()
self.normal_product = k
assert self.normal_product.shape == (self.n_boundaries,
3 * self.n_boundaries)
assert (self.normal_product[1, 3:6] == self.normal[1]).all()