agd.Metrics.riemann

  1# Copyright 2020 Jean-Marie Mirebeau, University Paris-Sud, CNRS, University Paris-Saclay
  2# Distributed WITHOUT ANY WARRANTY. Licensed under the Apache License, Version 2.0, see http://www.apache.org/licenses/LICENSE-2.0
  3
  4import numpy as np
  5
  6from . import misc
  7from .base import Base
  8from .diagonal import Diagonal
  9from .. import LinearParallel as lp
 10from .. import AutomaticDifferentiation as ad
 11from .. import FiniteDifferences as fd
 12
 13class Riemann(Base):
 14	r"""
 15	A Riemann norm takes the form 
 16	$$
 17	F(x) = \sqrt{< x,m x>}, 
 18	$$
 19	where m is a given symmetric positive definite tensor.
 20
 21	Member fields and __init__ arguments : 
 22	- m : an array of shape $(d,d,n_1,..,n_k)$ where d=`vdim` is the ambient space dimension.
 23	The array must be symmetric, a.k.a m[i,j] = m[j,i] for all $0\leq i < j < d$. 
 24	"""
 25
 26	def __init__(self,m):
 27		self.m=ad.asarray(m)
 28
 29	def norm(self,v):
 30		v,m = fd.common_field((v,self.m),(1,2))
 31		return np.sqrt(lp.dot_VAV(v,m,v))
 32
 33	def gradient(self,v):
 34		v,m = fd.common_field((v,self.m),(1,2))
 35		mv = lp.dot_AV(m,v)
 36		return mv/np.sqrt(lp.dot_VV(v,mv))
 37
 38	def dual(self,**kwargs):
 39		"""kwargs : passed to lp.inverse (typical : avoid_np_linalg_inv=True, for speed)"""
 40		return Riemann(lp.inverse(self.m,**kwargs))
 41
 42	@property
 43	def vdim(self): return len(self.m)
 44
 45	@property
 46	def shape(self): return self.m.shape[2:]	
 47
 48	def eigvals(self):
 49		"""
 50		Eigenvalues of self.m
 51		"""
 52		try: return np.moveaxis(np.linalg.eigvalsh(np.moveaxis(self.m,(0,1),(-2,-1))),-1,0)
 53		except ValueError: 
 54			assert ad.cupy_generic.from_cupy(self.m)
 55			import cupy
 56			return cupy.asarray(Riemann(self.m.get()).eigvals())
 57	def is_definite(self):
 58		return self.eigvals().min(axis=0)>0
 59	def anisotropy(self):
 60		ev = self.eigvals()
 61		return np.sqrt(ev.max(axis=0)/ev.min(axis=0))
 62	def cost_bound(self):
 63		return np.sqrt(lp.trace(self.m))
 64
 65	def inv_transform(self,a):
 66		m,a = fd.common_field((self.m,a),depths=(2,2))
 67		return Riemann(lp.dot_AA(lp.transpose(a),lp.dot_AA(m,a)))
 68	def with_costs(self,costs):
 69		costs,m = fd.common_field((costs,self.m),depths=(1,2))
 70		return Riemann(m*lp.outer_self(costs))
 71
 72	def flatten(self):
 73		return misc.flatten_symmetric_matrix(self.m)
 74
 75	@classmethod
 76	def expand(cls,arr):
 77		return cls(misc.expand_symmetric_matrix(arr))
 78
 79	def model_HFM(self):
 80		return "Riemann"+str(self.vdim)
 81
 82	@classmethod
 83	def needle(cls,u,cost_parallel,cost_orthogonal,ret_u=False):
 84		"""
 85		Defines a Riemannian metric, with 
 86		- eigenvector u
 87		- eigenvalue cost_parallel**2 in the eigenspace spanned by u
 88		- eigenvalue cost_orthogonal**2 in the eigenspace orthogonal with u
 89
 90		The metric is 
 91		- needle-like if cost_parallel < cost_orthogonal
 92		- plate-like otherwise
 93
 94		Optional argument:
 95		- ret_u : wether to return the (normalized) vector u
 96		"""
 97		u,cost_parallel,cost_orthogonal = (ad.asarray(e) for e in (u,cost_parallel,cost_orthogonal))
 98		u,cost_parallel,cost_orthogonal = fd.common_field((u.copy(),cost_parallel,cost_orthogonal),(1,0,0))
 99		
100		# Eigenvector normalization
101		nu = ad.Optimization.norm(u,ord=2,axis=0)
102		mask = nu>0
103		u[:,mask] /= nu[mask]
104
105		xp = ad.cupy_generic.get_array_module(u)
106		ident = fd.as_field(xp.eye(len(u),dtype=u.dtype),cost_parallel.shape,conditional=False)
107
108		m = (cost_parallel**2-cost_orthogonal**2) * lp.outer_self(u) + cost_orthogonal**2 * ident
109		return (cls(m),u) if ret_u else cls(m)
110
111	@classmethod
112	def from_diagonal(cls,diag):
113		"""
114		Produces a Riemann norm whose tensors have the given diagonal.
115		"""
116		diag = ad.asarray(diag)
117		z = np.zeros_like(diag[0])
118		vdim = len(diag)
119		arr = ad.asarray([[z if i!=j else diag[i] for i in range(vdim)] for j in range(vdim)])
120		return cls(arr)
121
122	@classmethod
123	def from_cast(cls,metric):
124		if isinstance(metric,cls): return metric
125		diag = Diagonal.from_cast(metric)
126		return Riemann.from_diagonal(diag.costs**2)
127
128	@classmethod
129	def from_mapped_eigenvalues(cls,matrix,mapping):
130		"""
131		Defines a Riemannian metric which has the same eigenvectors as the provided 
132		matrix, but (possibly) distinct eigenvalues obtained by the provided mapping.
133
134		Inputs : 
135		- matrix: a symmetric matrix, with shape (dim,dim,...)
136		- mapping: a function, taking as input an array of shape (dim,...),
137			and returning a similarly shaped array. 
138			Called with the eigenvalues of matrix, sorted from smallest to largest.
139		"""
140
141		# Get eigenvalues and eigenvectors, caring that numpy puts physical axes last
142		m_ = np.moveaxis(matrix,(0,1),(-2,-1)) 
143		eVal_,eVec_ = np.linalg.eigh(m_) # Not compatible with AD.
144		eVal,eVec = np.moveaxis(eVal_,-1,0),np.moveaxis(eVec_,(-2,-1),(0,1))
145
146		# Apply provided mapping and construct new matrix
147		mVal = ad.asarray(mapping(eVal))
148		m = lp.outer(eVec,mVal*eVec).sum(axis=2)
149		return cls(m)
150
151	def __iter__(self):
152		yield self.m
153
154	def make_proj_dual(self,**kwargs):
155		"""
156		kwargs : passed to Diagonal.make_proj
157		"""
158		λ,v = np.linalg.eigh(np.moveaxis(self.m,(0,1),(-2,-1)))
159		λ = np.moveaxis(λ,-1,0) 
160		v_ = np.moveaxis(v,(-2,-1),(0,1))
161		proj_diag = Diagonal(λ).make_proj_dual(**kwargs)
162		def proj(x):
163			v,x = fd.common_field((v_,x),depths=(2,1))
164			return lp.dot_AV(v,proj_diag(lp.dot_VA(x,v)))
165		return proj
class Riemann(agd.Metrics.base.Base):
 14class Riemann(Base):
 15	r"""
 16	A Riemann norm takes the form 
 17	$$
 18	F(x) = \sqrt{< x,m x>}, 
 19	$$
 20	where m is a given symmetric positive definite tensor.
 21
 22	Member fields and __init__ arguments : 
 23	- m : an array of shape $(d,d,n_1,..,n_k)$ where d=`vdim` is the ambient space dimension.
 24	The array must be symmetric, a.k.a m[i,j] = m[j,i] for all $0\leq i < j < d$. 
 25	"""
 26
 27	def __init__(self,m):
 28		self.m=ad.asarray(m)
 29
 30	def norm(self,v):
 31		v,m = fd.common_field((v,self.m),(1,2))
 32		return np.sqrt(lp.dot_VAV(v,m,v))
 33
 34	def gradient(self,v):
 35		v,m = fd.common_field((v,self.m),(1,2))
 36		mv = lp.dot_AV(m,v)
 37		return mv/np.sqrt(lp.dot_VV(v,mv))
 38
 39	def dual(self,**kwargs):
 40		"""kwargs : passed to lp.inverse (typical : avoid_np_linalg_inv=True, for speed)"""
 41		return Riemann(lp.inverse(self.m,**kwargs))
 42
 43	@property
 44	def vdim(self): return len(self.m)
 45
 46	@property
 47	def shape(self): return self.m.shape[2:]	
 48
 49	def eigvals(self):
 50		"""
 51		Eigenvalues of self.m
 52		"""
 53		try: return np.moveaxis(np.linalg.eigvalsh(np.moveaxis(self.m,(0,1),(-2,-1))),-1,0)
 54		except ValueError: 
 55			assert ad.cupy_generic.from_cupy(self.m)
 56			import cupy
 57			return cupy.asarray(Riemann(self.m.get()).eigvals())
 58	def is_definite(self):
 59		return self.eigvals().min(axis=0)>0
 60	def anisotropy(self):
 61		ev = self.eigvals()
 62		return np.sqrt(ev.max(axis=0)/ev.min(axis=0))
 63	def cost_bound(self):
 64		return np.sqrt(lp.trace(self.m))
 65
 66	def inv_transform(self,a):
 67		m,a = fd.common_field((self.m,a),depths=(2,2))
 68		return Riemann(lp.dot_AA(lp.transpose(a),lp.dot_AA(m,a)))
 69	def with_costs(self,costs):
 70		costs,m = fd.common_field((costs,self.m),depths=(1,2))
 71		return Riemann(m*lp.outer_self(costs))
 72
 73	def flatten(self):
 74		return misc.flatten_symmetric_matrix(self.m)
 75
 76	@classmethod
 77	def expand(cls,arr):
 78		return cls(misc.expand_symmetric_matrix(arr))
 79
 80	def model_HFM(self):
 81		return "Riemann"+str(self.vdim)
 82
 83	@classmethod
 84	def needle(cls,u,cost_parallel,cost_orthogonal,ret_u=False):
 85		"""
 86		Defines a Riemannian metric, with 
 87		- eigenvector u
 88		- eigenvalue cost_parallel**2 in the eigenspace spanned by u
 89		- eigenvalue cost_orthogonal**2 in the eigenspace orthogonal with u
 90
 91		The metric is 
 92		- needle-like if cost_parallel < cost_orthogonal
 93		- plate-like otherwise
 94
 95		Optional argument:
 96		- ret_u : wether to return the (normalized) vector u
 97		"""
 98		u,cost_parallel,cost_orthogonal = (ad.asarray(e) for e in (u,cost_parallel,cost_orthogonal))
 99		u,cost_parallel,cost_orthogonal = fd.common_field((u.copy(),cost_parallel,cost_orthogonal),(1,0,0))
100		
101		# Eigenvector normalization
102		nu = ad.Optimization.norm(u,ord=2,axis=0)
103		mask = nu>0
104		u[:,mask] /= nu[mask]
105
106		xp = ad.cupy_generic.get_array_module(u)
107		ident = fd.as_field(xp.eye(len(u),dtype=u.dtype),cost_parallel.shape,conditional=False)
108
109		m = (cost_parallel**2-cost_orthogonal**2) * lp.outer_self(u) + cost_orthogonal**2 * ident
110		return (cls(m),u) if ret_u else cls(m)
111
112	@classmethod
113	def from_diagonal(cls,diag):
114		"""
115		Produces a Riemann norm whose tensors have the given diagonal.
116		"""
117		diag = ad.asarray(diag)
118		z = np.zeros_like(diag[0])
119		vdim = len(diag)
120		arr = ad.asarray([[z if i!=j else diag[i] for i in range(vdim)] for j in range(vdim)])
121		return cls(arr)
122
123	@classmethod
124	def from_cast(cls,metric):
125		if isinstance(metric,cls): return metric
126		diag = Diagonal.from_cast(metric)
127		return Riemann.from_diagonal(diag.costs**2)
128
129	@classmethod
130	def from_mapped_eigenvalues(cls,matrix,mapping):
131		"""
132		Defines a Riemannian metric which has the same eigenvectors as the provided 
133		matrix, but (possibly) distinct eigenvalues obtained by the provided mapping.
134
135		Inputs : 
136		- matrix: a symmetric matrix, with shape (dim,dim,...)
137		- mapping: a function, taking as input an array of shape (dim,...),
138			and returning a similarly shaped array. 
139			Called with the eigenvalues of matrix, sorted from smallest to largest.
140		"""
141
142		# Get eigenvalues and eigenvectors, caring that numpy puts physical axes last
143		m_ = np.moveaxis(matrix,(0,1),(-2,-1)) 
144		eVal_,eVec_ = np.linalg.eigh(m_) # Not compatible with AD.
145		eVal,eVec = np.moveaxis(eVal_,-1,0),np.moveaxis(eVec_,(-2,-1),(0,1))
146
147		# Apply provided mapping and construct new matrix
148		mVal = ad.asarray(mapping(eVal))
149		m = lp.outer(eVec,mVal*eVec).sum(axis=2)
150		return cls(m)
151
152	def __iter__(self):
153		yield self.m
154
155	def make_proj_dual(self,**kwargs):
156		"""
157		kwargs : passed to Diagonal.make_proj
158		"""
159		λ,v = np.linalg.eigh(np.moveaxis(self.m,(0,1),(-2,-1)))
160		λ = np.moveaxis(λ,-1,0) 
161		v_ = np.moveaxis(v,(-2,-1),(0,1))
162		proj_diag = Diagonal(λ).make_proj_dual(**kwargs)
163		def proj(x):
164			v,x = fd.common_field((v_,x),depths=(2,1))
165			return lp.dot_AV(v,proj_diag(lp.dot_VA(x,v)))
166		return proj

A Riemann norm takes the form $$ F(x) = \sqrt{< x,m x>}, $$ where m is a given symmetric positive definite tensor.

Member fields and __init__ arguments :

  • m : an array of shape $(d,d,n_1,..,n_k)$ where d=vdim is the ambient space dimension. The array must be symmetric, a.k.a m[i,j] = m[j,i] for all $0\leq i < j < d$.
Riemann(m)
27	def __init__(self,m):
28		self.m=ad.asarray(m)
m
def norm(self, v):
30	def norm(self,v):
31		v,m = fd.common_field((v,self.m),(1,2))
32		return np.sqrt(lp.dot_VAV(v,m,v))

Norm or quasi-norm defined by the class, often denoted $N$ in mathematical formulas. Unless incorrect data is provided, this member function obeys, for all vectors $u,v\in R^d$ and all $\alpha \geq 0$

  • $N(u+v) \leq N(u)+N(v)$
  • $N(\alpha u) = \alpha N(u)$
  • $N(u)\geq 0$ with equality iff $u=0$.

Broadcasting will occur depending on the shape of $v$ and of the class data.

def gradient(self, v):
34	def gradient(self,v):
35		v,m = fd.common_field((v,self.m),(1,2))
36		mv = lp.dot_AV(m,v)
37		return mv/np.sqrt(lp.dot_VV(v,mv))

Gradient of the norm defined by the class.

def dual(self, **kwargs):
39	def dual(self,**kwargs):
40		"""kwargs : passed to lp.inverse (typical : avoid_np_linalg_inv=True, for speed)"""
41		return Riemann(lp.inverse(self.m,**kwargs))

kwargs : passed to lp.inverse (typical : avoid_np_linalg_inv=True, for speed)

vdim
43	@property
44	def vdim(self): return len(self.m)

The ambient vector space dimension, often denoted $d$ in mathematical formulas.

shape
46	@property
47	def shape(self): return self.m.shape[2:]	

Dimensions of the underlying domain. Expected to be the empty tuple, or a tuple of length vdim.

def eigvals(self):
49	def eigvals(self):
50		"""
51		Eigenvalues of self.m
52		"""
53		try: return np.moveaxis(np.linalg.eigvalsh(np.moveaxis(self.m,(0,1),(-2,-1))),-1,0)
54		except ValueError: 
55			assert ad.cupy_generic.from_cupy(self.m)
56			import cupy
57			return cupy.asarray(Riemann(self.m.get()).eigvals())

Eigenvalues of self.m

def is_definite(self):
58	def is_definite(self):
59		return self.eigvals().min(axis=0)>0

Attempts to check wether the data defines a mathematically valid norm.

def anisotropy(self):
60	def anisotropy(self):
61		ev = self.eigvals()
62		return np.sqrt(ev.max(axis=0)/ev.min(axis=0))

Anisotropy ratio of the norm denoted $N$. Defined as $$ \max_{|u| = |v| = 1} \frac {N(u)}{N(v)}. $$

def cost_bound(self):
63	def cost_bound(self):
64		return np.sqrt(lp.trace(self.m))

Upper bound on $N(u)$, for any unit vector $u$, where $N$ is the norm defined by the class.

def inv_transform(self, a):
66	def inv_transform(self,a):
67		m,a = fd.common_field((self.m,a),depths=(2,2))
68		return Riemann(lp.dot_AA(lp.transpose(a),lp.dot_AA(m,a)))

Affine transformation of the norm. The new unit ball is the inverse image of the previous one.

def with_costs(self, costs):
69	def with_costs(self,costs):
70		costs,m = fd.common_field((costs,self.m),depths=(1,2))
71		return Riemann(m*lp.outer_self(costs))

Produces a norm $N'$ defined by $$ N'(x) = N(costs * x) $$ where the multiplication is elementwise.

def flatten(self):
73	def flatten(self):
74		return misc.flatten_symmetric_matrix(self.m)

Flattens and concatenate the member fields into a single array.

@classmethod
def expand(cls, arr):
76	@classmethod
77	def expand(cls,arr):
78		return cls(misc.expand_symmetric_matrix(arr))

Inverse of the flatten member function. Turns a suitable array into a metric.

def model_HFM(self):
80	def model_HFM(self):
81		return "Riemann"+str(self.vdim)

The name of the 'model' for parameter, as input to the HFM library.

@classmethod
def needle(cls, u, cost_parallel, cost_orthogonal, ret_u=False):
 83	@classmethod
 84	def needle(cls,u,cost_parallel,cost_orthogonal,ret_u=False):
 85		"""
 86		Defines a Riemannian metric, with 
 87		- eigenvector u
 88		- eigenvalue cost_parallel**2 in the eigenspace spanned by u
 89		- eigenvalue cost_orthogonal**2 in the eigenspace orthogonal with u
 90
 91		The metric is 
 92		- needle-like if cost_parallel < cost_orthogonal
 93		- plate-like otherwise
 94
 95		Optional argument:
 96		- ret_u : wether to return the (normalized) vector u
 97		"""
 98		u,cost_parallel,cost_orthogonal = (ad.asarray(e) for e in (u,cost_parallel,cost_orthogonal))
 99		u,cost_parallel,cost_orthogonal = fd.common_field((u.copy(),cost_parallel,cost_orthogonal),(1,0,0))
100		
101		# Eigenvector normalization
102		nu = ad.Optimization.norm(u,ord=2,axis=0)
103		mask = nu>0
104		u[:,mask] /= nu[mask]
105
106		xp = ad.cupy_generic.get_array_module(u)
107		ident = fd.as_field(xp.eye(len(u),dtype=u.dtype),cost_parallel.shape,conditional=False)
108
109		m = (cost_parallel**2-cost_orthogonal**2) * lp.outer_self(u) + cost_orthogonal**2 * ident
110		return (cls(m),u) if ret_u else cls(m)

Defines a Riemannian metric, with

  • eigenvector u
  • eigenvalue cost_parallel**2 in the eigenspace spanned by u
  • eigenvalue cost_orthogonal**2 in the eigenspace orthogonal with u

The metric is

  • needle-like if cost_parallel < cost_orthogonal
  • plate-like otherwise

Optional argument:

  • ret_u : wether to return the (normalized) vector u
@classmethod
def from_diagonal(cls, diag):
112	@classmethod
113	def from_diagonal(cls,diag):
114		"""
115		Produces a Riemann norm whose tensors have the given diagonal.
116		"""
117		diag = ad.asarray(diag)
118		z = np.zeros_like(diag[0])
119		vdim = len(diag)
120		arr = ad.asarray([[z if i!=j else diag[i] for i in range(vdim)] for j in range(vdim)])
121		return cls(arr)

Produces a Riemann norm whose tensors have the given diagonal.

@classmethod
def from_cast(cls, metric):
123	@classmethod
124	def from_cast(cls,metric):
125		if isinstance(metric,cls): return metric
126		diag = Diagonal.from_cast(metric)
127		return Riemann.from_diagonal(diag.costs**2)

Produces a metric by casting another metric of a compatible type.

@classmethod
def from_mapped_eigenvalues(cls, matrix, mapping):
129	@classmethod
130	def from_mapped_eigenvalues(cls,matrix,mapping):
131		"""
132		Defines a Riemannian metric which has the same eigenvectors as the provided 
133		matrix, but (possibly) distinct eigenvalues obtained by the provided mapping.
134
135		Inputs : 
136		- matrix: a symmetric matrix, with shape (dim,dim,...)
137		- mapping: a function, taking as input an array of shape (dim,...),
138			and returning a similarly shaped array. 
139			Called with the eigenvalues of matrix, sorted from smallest to largest.
140		"""
141
142		# Get eigenvalues and eigenvectors, caring that numpy puts physical axes last
143		m_ = np.moveaxis(matrix,(0,1),(-2,-1)) 
144		eVal_,eVec_ = np.linalg.eigh(m_) # Not compatible with AD.
145		eVal,eVec = np.moveaxis(eVal_,-1,0),np.moveaxis(eVec_,(-2,-1),(0,1))
146
147		# Apply provided mapping and construct new matrix
148		mVal = ad.asarray(mapping(eVal))
149		m = lp.outer(eVec,mVal*eVec).sum(axis=2)
150		return cls(m)

Defines a Riemannian metric which has the same eigenvectors as the provided matrix, but (possibly) distinct eigenvalues obtained by the provided mapping.

Inputs :

  • matrix: a symmetric matrix, with shape (dim,dim,...)
  • mapping: a function, taking as input an array of shape (dim,...), and returning a similarly shaped array. Called with the eigenvalues of matrix, sorted from smallest to largest.
def make_proj_dual(self, **kwargs):
155	def make_proj_dual(self,**kwargs):
156		"""
157		kwargs : passed to Diagonal.make_proj
158		"""
159		λ,v = np.linalg.eigh(np.moveaxis(self.m,(0,1),(-2,-1)))
160		λ = np.moveaxis(λ,-1,0) 
161		v_ = np.moveaxis(v,(-2,-1),(0,1))
162		proj_diag = Diagonal(λ).make_proj_dual(**kwargs)
163		def proj(x):
164			v,x = fd.common_field((v_,x),depths=(2,1))
165			return lp.dot_AV(v,proj_diag(lp.dot_VA(x,v)))
166		return proj

kwargs : passed to Diagonal.make_proj