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
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$.
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.
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.
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)
Dimensions of the underlying domain.
Expected to be the empty tuple, or a tuple of length vdim
.
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
Anisotropy ratio of the norm
denoted $N$.
Defined as
$$
\max_{|u| = |v| = 1} \frac {N(u)}{N(v)}.
$$
Upper bound on $N(u)$, for any unit vector $u$, where $N$ is the norm
defined by the class.
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.
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.
Inverse of the flatten member function. Turns a suitable array into a metric.
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
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.
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.
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.
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