agd.LinearParallel

  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
  4"""
  5This module implements some basic linear algebra routines, with the following characteristics.
  6- The geometry comes first, a.k.a vector has shape (vdim, n1,...,nk) where vdim is the 
  7 ambient vector space dimension, and n1,...,nk are arbitrary. 
  8 Note that *numpy uses the opposite convention*, putting vdim in last position.
  9- The routines are compatible with forward automatic differentiation, see module
 10 AutomaticDifferentiation.
 11"""
 12
 13import numpy as np
 14from . import AutomaticDifferentiation as ad
 15from . import FiniteDifferences as fd
 16
 17def identity(shape):
 18	dim = len(shape)
 19	a = np.full((dim,dim)+shape,0.)
 20	for i in range(dim):
 21		a[i,i]=1.
 22	return a
 23
 24def rotation(theta,axis=None):
 25	"""
 26	Dimension 2 : by a given angle.
 27	Dimension 3 : by a given angle, along a given axis.
 28	Three dimensional rotation matrix, with given axis and angle.
 29	Adapted from https://stackoverflow.com/a/6802723
 30	"""
 31	if axis is None:
 32		c,s=np.cos(theta),np.sin(theta)
 33		return ad.asarray([[c,-s],[s,c]])
 34	else:
 35		theta,axis = (ad.asarray(e) for e in (theta,axis))
 36		axis = axis / np.linalg.norm(axis,axis=0)
 37		theta,axis=fd.common_field((theta,axis),(0,1))
 38		a = np.cos(theta / 2.0)
 39		b, c, d = -axis * np.sin(theta / 2.0)
 40		aa, bb, cc, dd = a * a, b * b, c * c, d * d
 41		bc, ad_, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
 42		return ad.asarray([
 43			[aa + bb - cc - dd, 2 * (bc + ad_), 2 * (bd - ac)],
 44			[2 * (bc - ad_), aa + cc - bb - dd, 2 * (cd + ab)],
 45			[2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]])
 46#	return scipy.linalg.expm(np.cross(np.eye(3), axis/scipy.linalg.norm(axis)*theta)) # Alternative
 47
 48# Dot product (vector-vector, matrix-vector and matrix-matrix) in parallel
 49def dot_VV(v,w):
 50	"""
 51	Dot product <v,w> of two vectors.
 52	Inputs : 
 53	- v,w : arrays of shape (vdim, n1,...,nk), 
 54	 where vdim is the ambient vector space dimension
 55	"""
 56	v=ad.asarray(v); w=ad.asarray(w)
 57	if v.shape[0]!=w.shape[0]: raise ValueError('dot_VV : Incompatible shapes')
 58	return (v*w).sum(0)
 59
 60def dot_AV(a,v):
 61	"""
 62	Dot product a.v of a matrix and vector.
 63	Inputs : 
 64	- a : array of shape (wdim,vdim, n1,...,nk)
 65	- v : array of shape (vdim, n1,...,nk),
 66	 where vdim is the ambient vector space dimension
 67	"""
 68	if a.shape[1]!=v.shape[0]: raise ValueError("dot_AV : Incompatible shapes")
 69	return (a*np.expand_dims(v,axis=0)).sum(1)
 70
 71def dot_VA(v,a):
 72	"""
 73	Dot product v^T.a of a vector and matrix.
 74	Inputs : 
 75	- v : array of shape (vdim, n1,...,nk)
 76	- a : array of shape (vdim,wdim, n1,...,nk),
 77	 where vdim is the ambient vector space dimension
 78	"""
 79	m,n = a.shape[:2]
 80	bounds = a.shape[2:]
 81	if v.shape != (m,)+bounds:
 82		raise ValueError("dot_VA : Incompatible shapes")
 83
 84	return (v.reshape((m,1)+bounds)*a).sum(0)
 85
 86
 87def dot_AA(a,b):
 88	"""
 89	Dot product a.b of two matrices.
 90	Inputs : 
 91	- a: array of shape (vdim,wdim, n1,...,nk),
 92	- a: array of shape (wdim,xdim, n1,...,nk),	
 93	"""
 94	m,n=a.shape[:2]
 95	bounds = a.shape[2:]
 96	k = b.shape[1]
 97	if b.shape!=(n,k,)+bounds: raise ValueError("dot_AA error : Incompatible shapes")
 98	if ad.is_ad(a) or ad.is_ad(b): # Note : not very memory efficient
 99		return (a.reshape((m,n,1)+bounds)*b.reshape((1,n,k)+bounds)).sum(1)
100	else: 
101		return np.moveaxis(
102			np.moveaxis(a,(0,1),(-2,-1)) @ np.moveaxis(b,(0,1),(-2,-1)),(-2,-1),(0,1))
103
104def dot_VAV(v,a,w):
105	"""
106	Dot product <v,a.w> of two vectors and a matrix (usually symmetric).
107	Inputs (typical): 
108	- v: array of shape (vdim, n1,...,nk),	
109	- a: array of shape (vdim,vdim, n1,...,nk),
110	- w: array of shape (vdim, n1,...,nk),	
111	"""
112	return dot_VV(v,dot_AV(a,w))
113	
114def mult(k,x):
115	"""Multiplication by scalar, of a vector or matrix"""
116	# Quite useless, as it is directly handled by broadcasting
117	bounds = k.shape
118	dim = x.ndim-k.ndim
119	if x.shape[dim:]!=bounds:
120		raise ValueError("mult error : incompatible shapes")
121	return k.reshape((1,)*dim+bounds)*x
122	
123
124def perp(v):
125	"""
126	Rotates a vector by pi/2, producing [v[1],v[0]]
127	Inputs: 
128	- v: array of shape (2, n1,...,nk)
129	"""
130	if v.shape[0]!=2:
131		raise ValueError("perp error : Incompatible dimension")		
132	return ad.asarray( (-v[1],v[0]) )
133	
134def cross(v,w):
135	"""
136	Cross product v x w of two vectors.
137	Inputs: 
138	- v,w: arrays of shape (3, n1,...,nk)
139	"""
140	if v.shape[0]!=3 or v.shape!=w.shape:
141		raise ValueError("cross error : Incompatible dimensions")
142	return ad.asarray( (v[1]*w[2]-v[2]*w[1], \
143	v[2]*w[0]-v[0]*w[2], v[0]*w[1]-v[1]*w[0]) )
144	
145def outer(v,w):
146	"""
147	Outer product v w^T of two vectors.
148	Inputs : 
149	- v,w: arrays of shape (vdim, n1,...,nk),
150	 where vdim is the ambient vector space dimension
151	"""
152	if v.shape[1:] != w.shape[1:]:
153		raise ValueError("outer error : Incompatible dimensions")
154	m,n=v.shape[0],w.shape[0]
155	bounds = v.shape[1:]
156	return v.reshape((m,1)+bounds)*w.reshape((1,n)+bounds)
157
158def outer_self(v):
159	"""
160	Outer product v v^T of a vector with itself.
161	Inputs : 
162	- v: array of shape (vdim, n1,...,nk),
163	 where vdim is the ambient vector space dimension
164	"""
165	v=ad.asarray(v)
166	return outer(v,v)
167
168def transpose(a):
169	"""
170	Transpose a^T of a matrix.
171	Input : 
172	- a: array of shape (vdim,wdim, n1,...,nk),
173	"""
174	return a.transpose( (1,0,)+tuple(range(2,a.ndim)) )
175	
176def trace(a):
177	"""
178	Trace tr(a) of a square matrix, a.k.a sum of the diagonal elements.
179	Input : 
180	- a: array of shape (vdim,vdim, n1,...,nk),
181	 where vdim is the ambient vector space dimension
182	"""
183	vdim = a.shape[0]
184	if a.shape[1]!=vdim:
185		raise ValueError("trace error : incompatible dimensions")
186	return sum(a[i,i] for i in range(vdim))
187
188# Low dimensional special cases
189
190def det(a):
191	"""
192	Determinant of a square matrix.
193	Input : 
194	- a: array of shape (vdim,vdim, n1,...,nk),
195	 where vdim is the ambient vector space dimension
196	"""
197	a=ad.asarray(a)
198
199	dim = a.shape[0]
200	if a.shape[1]!=dim:
201		raise ValueError("inverse error : incompatible dimensions")
202	if dim==1:
203		return a[0,0]
204	elif dim==2:
205		return a[0,0]*a[1,1]-a[0,1]*a[1,0]
206	elif dim==3:
207		return a[0,0]*a[1,1]*a[2,2]+a[0,1]*a[1,2]*a[2,0]+a[0,2]*a[1,0]*a[2,1] \
208		- a[0,2]*a[1,1]*a[2,0] - a[1,2]*a[2,1]*a[0,0]- a[2,2]*a[0,1]*a[1,0]
209	else:
210		raise ValueError("det error : unsupported dimension") 
211
212# Suppressed due to extreme slowness, at least in cupy 6
213#	if not (ad.is_ad(a) or a.dtype==np.dtype('object')): 
214#		return np.linalg.det(np.moveaxis(a,(0,1),(-2,-1)))
215
216
217def inverse(a,avoid_np_linalg_inv=False):
218	"""
219	Inverse of a square matrix.
220	Input : 
221	- a: array of shape (vdim,vdim, n1,...,nk),
222	 where vdim is the ambient vector space dimension.
223	- avoid_np_linalg_inv: for unknown reasons, the np.linalg.inv routine can be very slow 
224	  when applied to large arrays of many small matrices. (numpy 1.23.3, mac OS, M1 processor)
225	  Use an explicit inversion instead, via Cramer's formula, implemented for shape (2,2) and (3,3)
226	"""
227	a=ad.asarray(a)
228	if not (ad.is_ad(a) or a.dtype==np.dtype('object') or avoid_np_linalg_inv):
229		try: return np.moveaxis(np.linalg.inv(np.moveaxis(a,(0,1),(-2,-1))),(-2,-1),(0,1))
230		except np.linalg.LinAlgError: pass # Old cupy versions do not support arrays of dimension>2
231
232	if isinstance(a,ad.Dense.denseAD):
233		b = inverse(a.value)
234		b_ = fd.as_field(b,(a.size_ad,),conditional=False) 
235		h = a.coef
236		return ad.Dense.denseAD(b,-dot_AA(b_,dot_AA(h,b_)))
237	elif isinstance(a,ad.Dense2.denseAD2):
238		b = inverse(a.value)
239		b1 = fd.as_field(b,(a.size_ad,),conditional=False)
240		h = a.coef1
241		h2 = a.coef2
242
243		bh = dot_AA(b1,h)
244		bhb = dot_AA(bh,b1)
245		bhbhb = dot_AA(np.broadcast_to(np.expand_dims(bh,-1),h2.shape),
246			np.broadcast_to(np.expand_dims(bhb,-2),h2.shape))
247
248		b2 = fd.as_field(b,(a.size_ad,a.size_ad),conditional=False)
249		bh2b = dot_AA(b2,dot_AA(h2,b2))
250		return ad.Dense2.denseAD2(b,-bhb,bhbhb+np.swapaxes(bhbhb,-1,-2)-bh2b)
251	elif ad.is_ad(a):
252		d=len(a)
253		return ad.apply(inverse,a,shape_bound=a.shape[2:])
254	elif a.shape[:2] == (2,2):
255		return ad.asarray([[a[1,1],-a[0,1]],[-a[1,0],a[0,0]]])/det(a)
256	elif a.shape[:2] == (3,3):
257		return ad.asarray([[
258			a[(i+1)%3,(j+1)%3]*a[(i+2)%3,(j+2)%3]-a[(i+1)%3,(j+2)%3]*a[(i+2)%3,(j+1)%3]
259			for i in range(3)] for j in range(3)])/det(a)
260	raise ValueError(f"Unsupported inverse for {type(a)} with dtype {a.dtype} and dimensions {a.shape}")
261
262def solve_AV(a,v):
263	"""
264	Solution to a linear system (preferably low dimensional).
265	Input : 
266	- a: array of shape (vdim,vdim, n1,...,nk),
267	- v: array of shape (vdim,vdim, n1,...,nk),
268	 where vdim is the ambient vector space dimension
269	"""
270	a,v=ad.asarray(a),ad.asarray(v)
271	if ad.is_ad(v) or a.dtype==np.dtype('object') or (
272		len(a)<=3 and ad.cupy_generic.from_cupy(a)): 
273		# Inefficient, but compatible with ndarray subclasses
274		# Also cupy.linalg.solve has a performance issue (cupy version 7.8) 
275		return dot_AV(inverse(a),v) 
276	return np.moveaxis(np.linalg.solve(np.moveaxis(a,(0,1),(-2,-1)),np.moveaxis(v,0,-1)),-1,0)