The MTW tensor
The MTW tensor
This is a Mathematica notebook designed to allow us to compute the MTW tensor for various cost functions. Following the work of Ma, Trudinger, Wang, and many others, this tensor is important to determining if the transport maps are smooth for these cost functions.
For this, we will start with a cost function (the default uses the logarithmic cost), with an eye to building a more general notebook that can handle a greater variety of cost functions.
First clear any values that may already have been assigned to the names of the various objects to be calculated. The names of the coordinates that you use are also cleared.
In[]:=
Clear[xcoord,ycoord,Hess,InverseHess,costx,costy,costxxy,costxyy,costxxyy,MTW,QuadMTW,n]
First we set the dimension. We will assume that dim X=dim Y, which is necessary for our analysis.
In[]:=
n=2
Out[]=
2
Note that it is important not to use the symbols, i, j, k, l, p, q, r, s or n as constants or coordinates in the costs and coordinates. The reason is that the first five of those symbols are used as summation or table indices in the calculations done below, and n is the dimension of the space.
In[]:=
xcoord={x1,x2}
Out[]=
{x1,x2}
In[]:=
ycoord={y1,y2}
Out[]=
{y1,y2}
Here we define the cost function. As an example, we will use the free energy for an exponential family.
In[]:=
cost=Log[1+Sum[Exp[(xcoord[[i]]-ycoord[[i]])],{i,1,n}]]
Out[]=
Log[1++]
x1-y1
x2-y2
Computing the Hessian and its inverse
Computing the Hessian and its inverse
This is where we define the Hessian and its inverse.
In[]:=
Hess:=Hess=Simplify[Table[D[cost,xcoord[[i]],ycoord[[j]]],{i,1,n},{j,1,n}]]
In[]:=
InverseHess:=InverseHess=Simplify[Inverse[Hess]]
If you want to see the Hess and InverseHess explicitly, you can run the following lines.
In[]:=
MatrixForm[Hess]
Out[]//MatrixForm=
- x1+y1+y2 x2 y2 2 ( x2+y1 x1+y2 y1+y2 | x1+x2+y1+y2 2 ( x2+y1 x1+y2 y1+y2 |
x1+x2+y1+y2 2 ( x2+y1 x1+y2 y1+y2 | - x2+y1+y2 x1 y1 2 ( x2+y1 x1+y2 y1+y2 |
In[]:=
MatrixForm[InverseHess]
Out[]//MatrixForm=
- -x1-y1-y2 x1 y1 x2+y1 x1+y2 y1+y2 | -1- x1-y1 x2-y2 |
-1- x1-y1 x2-y2 | - -x2-y1-y2 x2 y2 x2+y1 x1+y2 y1+y2 |
It is also worthwhile to make sure that the Hessian is actually non-singular, as this is important for the regularity theory.
In[]:=
Det[Hess]
Out[]=
x1+2x2+3y1+2y2
4
(++)
x2+y1
x1+y2
y1+y2
2x1+x2+2y1+3y2
4
(++)
x2+y1
x1+y2
y1+y2
x1+x2+3y1+3y2
4
(++)
x2+y1
x1+y2
y1+y2
Computing the third and fourth derivatives
Computing the third and fourth derivatives
In order to define the MTW tensor, we must also define two tensors of third derivatives.
In[]:=
costxxy:=costxxy=Simplify[Table[D[cost,xcoord[[i]],xcoord[[j]],ycoord[[k]]],{i,1,n},{j,1,n},{k,1,n}]]
In[]:=
costxyy:=costxyy=Simplify[Table[D[cost,xcoord[[i]],ycoord[[j]],ycoord[[k]]],{i,1,n},{j,1,n},{k,1,n}]]
We must also define a tensor of fourth order derivatives.
In[]:=
costxxyy:=costxxyy=Simplify[Table[D[cost,xcoord[[i]],xcoord[[j]],ycoord[[k]],ycoord[[l]]],{i,1,n},{j,1,n},{k,1,n},{l,1,n}]]
It is possible to display cxxy,cxyy, and cxxyy, but for most practical purposes, they will be completely unwieldy.
Computing the MTW tensor
Computing the MTW tensor
With all of these, we can define the MTW tensor, which is a fourth order tensor.
In[]:=
MTW=Simplify[Table[Sum[(Sum[costxxy[[i,j,p]]*InverseHess[[p,q]]*costxyy[[q,r,s]],{p,1,n},{q,1,n}]-costxxyy[[i,j,r,s]])*InverseHess[[r,k]]*InverseHess[[s,l]],{r,1,n},{s,1,n}],{i,1,n},{j,1,n},{k,1,n},{l,1,n}]]
Out[]=
{{{{2,0},{0,0}},{{0,1},{1,0}}},{{{0,1},{1,0}},{{0,0},{0,2}}}}
Having the full fourth order MTW tensor is not particularly useful, and the relevant quantity for the regularity theory is the quadratic tensor MTW(v,v,w,w) for two orthogonal vectors v and w.
If you would like to specify what v and w are exactly, you can use the following two lines to specify them exactly. Otherwise, you can leave them as is to obtain the full quadratic curvature tensor.
In[]:=
v={v1,v2}
Out[]=
{v1,v2}
In[]:=
w={w1,w2}
Out[]=
{w1,w2}
In[]:=
QuadMTW=Simplify[Sum[MTW[[i,j,k,l]]*v[[i]]*v[[j]]*w[[k]]*w[[l]],{i,1,n},{j,1,n},{k,1,n},{l,1,n}]]
Out[]=
2
2
(v1w1+v2w2)
If this is non-negative definite, then the cost satisfies the MTW(0) condition.
Computing the Convexity
Computing the Convexity
For computing the MTW tensor, Mathematica is quite useful. Computing the cost-convexity requires less computation power, but we have provided the calculation nonetheless.
In[]:=
costx:=costx=Simplify[Table[D[cost,xcoord[[i]]],{i,1,n}]]
In[]:=
costy:=costy=Simplify[Table[D[cost,ycoord[[i]]],{i,1,n}]]
In order to prove regularity of optimal transport, one must show that for all x and y (respectively), costx and costy are injective as functions of y and x, respectively. One must also be able to invert these in order to understand the notion of cost-convexity.
Acknowledgements
Acknowledgements
This notebook was written by Gabriel Khan. It is adapted from James B. Hartle’s notebook for computing curvature of a Riemannian metric.
This notebook was written by Gabriel Khan. It is adapted from James B. Hartle’s notebook for computing curvature of a Riemannian metric.