agd.FiniteDifferences

  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 finite differences operations, as well as related array 
  6reshaping and broadcasting operations. The main functions are the following, see the 
  7the detailed help below.
  8
  9Elementary finite differences:
 10- DiffUpwind
 11- DiffCentered
 12- Diff2
 13
 14Array broadcasting:
 15- as_field
 16- common_field
 17
 18Block based array reshaping:
 19- block_expand
 20- block_squeeze
 21
 22"""
 23
 24import numpy as np
 25import itertools
 26from . import AutomaticDifferentiation as ad
 27import functools
 28import operator
 29
 30# --- Domain shape related methods --------
 31
 32def as_field(u,shape,conditional=True,depth=None):
 33	"""
 34	Checks if the last dimensions of u match the given shape. 
 35	If not, u is extended with these additional dimensions.
 36	conditional : if False, reshaping is always done
 37	depth (optional) : the depth of the geometrical tensor field (1: vectors, 2: matrices)
 38	"""
 39	u=ad.asarray(u)
 40	ndim = len(shape)
 41	def as_is():
 42		if not conditional: return False
 43		elif depth is None: return u.ndim>=ndim and u.shape[-ndim:]==shape
 44		else: assert u.shape[depth:] in (tuple(),shape); return u.shape[depth:]==shape
 45	if as_is(): return u
 46	else: return np.broadcast_to(u.reshape(u.shape+(1,)*ndim), u.shape+shape)
 47
 48def common_shape(arrays,depths):
 49	"""Finds a common shape to the arrays for broadcasting"""
 50	arrays = tuple(None if arr is None else ad.asarray(arr) for arr in arrays)
 51	assert len(arrays)==len(depths)
 52	for arr,depth in zip(arrays,depths):
 53		if arr is None: continue
 54		shape = arr.shape[depth:]
 55		if shape!=tuple(): break
 56	for arr,depth in zip(arrays,depths): 
 57		assert arr is None or arr.shape[depth:] in (tuple(),shape)
 58	return shape
 59
 60def common_field(arrays,depths,shape=None):
 61	"""
 62	Adds trailing dimensions, and broadcasts the given arrays, for suitable interoperation.
 63	
 64	Inputs: 
 65	- arrays : a list [a_1,...,a_n], or iterable, of numeric arrays such that
 66	 a_i.shape = shape_i + shape, or a_i.shape = shape_i, for each 1<=i<=n.
 67	- depths : defined as [len(shape_i) for 1<=i<=n]
 68	- shape (optional) : the trailing shape.
 69	
 70	Output:
 71	- the arrays, with added trailing and dimensions broadcasting so that
 72	 a_i.shape = shape_i + shape for each 1<=i<=n.
 73
 74	"""
 75	arrays = tuple(None if arr is None else ad.asarray(arr) for arr in arrays)
 76	assert len(arrays)==len(depths)
 77	if shape is None: shape = common_shape(arrays,depths)
 78	return tuple(None if arr is None else as_field(arr,shape,depth=depth) 
 79		for (arr,depth) in zip(arrays,depths))
 80
 81def round_up_ratio(num,den):
 82	"""
 83	Returns the least multiple of den after num.
 84	num and den must be integers, with den>0. 
 85	"""
 86	num,den = np.asarray(num),np.asarray(den)
 87	return (num+den-1)//den
 88
 89def block_expand(arr,shape_i,shape_o=None,**kwargs):
 90	"""
 91	Reshape an array so as to factor shape_i (the inner shape),
 92	and move these axes last.
 93	Inputs : 
 94	 - **kwargs : passed to np.pad
 95	Output : 
 96	 - padded and reshaped array
 97	"""
 98	ndim = len(shape_i)
 99	if ndim==0: return arr # Empty block
100	shape_pre = arr.shape[:-ndim]
101	ndim_pre = len(shape_pre)
102	shape_tot=np.array(arr.shape[-ndim:])
103	shape_i = np.array(shape_i)
104
105	# Extend data
106	if shape_o is None: shape_o = round_up_ratio(shape_tot,shape_i)
107	shape_pad = (0,)*ndim_pre + tuple(shape_o*shape_i - shape_tot)
108	arr = np.pad(arr, tuple( (0,s) for s in shape_pad), **kwargs) 
109
110	# Reshape
111	shape_interleaved = tuple(np.stack( (shape_o,shape_i), axis=1).flatten())
112	arr = arr.reshape(shape_pre + shape_interleaved)
113
114	# Move axes
115	rg = np.arange(ndim)
116	axes_interleaved = ndim_pre + 1+2*rg
117	axes_split = ndim_pre + ndim+rg
118	arr = np.moveaxis(arr,axes_interleaved,axes_split)
119
120	return arr
121
122def block_squeeze(arr,shape):
123	"""
124	Inverse operation to block_expand.
125	"""
126	ndim = len(shape)
127	if ndim==0: return arr # Empty block
128	shape_pre = arr.shape[:-2*ndim]
129	ndim_pre = len(shape_pre)
130	shape_o = arr.shape[(-2*ndim):-ndim]
131	shape_i = arr.shape[-ndim:]
132
133	# Move axes
134	rg = np.arange(ndim)
135	axes_interleaved = ndim_pre + 1+2*rg
136	axes_split = ndim_pre + ndim+rg
137	arr = np.moveaxis(arr,axes_split,axes_interleaved)
138
139	# Reshape
140	arr = arr.reshape(shape_pre
141		+tuple(s_o*s_i for (s_o,s_i) in zip(shape_o,shape_i)) )
142
143	# Extract subdomain
144	region = tuple(slice(0,s) for s in (shape_pre+shape))
145	arr = arr.__getitem__(region)
146
147	return arr
148
149def block_neighbors(shape_i,top=True):
150	"""
151	The first layer of neighbors to a block, along the top or bottom, ordered 
152	by increasing indices. Used for efficient data load when a implementing 
153	gradients, divergences, etc, with a compact scheme
154	"""
155	shape_i = np.array(shape_i); vdim=len(shape_i); size_i = np.prod(shape_i)
156	aX = [np.arange(s+1) if top else np.arange(s-1,2*s) for s in shape_i]
157	x_e = np.array(np.meshgrid(*aX)).astype(int)
158	x_e = np.moveaxis(x_e.reshape((vdim,-1)),0,-1)
159	x_e = np.array([x for x in x_e if np.any(x==(shape_i if top else shape_i-1))])
160	ind = np.arange(2**vdim*size_i).reshape((2,)*vdim+tuple(shape_i))
161	ind = block_squeeze(ind,tuple(2*shape_i))
162	ind_x = np.array([ind[tuple(x)] for x in x_e])
163	x_e = x_e[np.argsort(ind_x)]
164	#print(np.array([ind[tuple(x)] for x in x_e]))
165	if not top: x_e-=shape_i
166	return x_e
167
168
169# ----- Utilities for finite differences ------
170
171def BoundedSlices(slices,shape):
172	"""
173	Returns the input slices with None replaced with the upper bound
174	from the given shape
175	"""
176	if slices[-1]==Ellipsis:
177		slices=slices[:-1]+(slice(None,None,None),)*(len(shape)-len(slices)+1)
178	def BoundedSlice(s,n):
179		if not isinstance(s,slice):
180			return slice(s,s+1)
181		else:
182			return slice(s.start, n if s.stop is None else s.stop, s.step)
183	return tuple(BoundedSlice(s,n) for s,n in zip(slices,shape))
184
185def OffsetToIndex(shape,offset, mode='clip', uniform=None, where=(Ellipsis,)):
186	"""
187	Returns the index corresponding to position + offset, 
188	and a boolean for wether it falls inside the domain.
189	"""
190	ndim = len(shape) # Domain dimension
191	assert(offset.shape[0]==ndim)
192	if uniform is None: # Uniform = True iff offsets are independent of the position in the domain
193		uniform = not ((offset.ndim > ndim) and (offset.shape[-ndim:]==shape))
194
195	odim = (offset.ndim-1) if uniform else (offset.ndim - ndim-1) # Dimensions for distinct offsets 
196	everywhere = where==(Ellipsis,)
197	xp=ad.cupy_generic.get_array_module(offset)
198
199	grid = (xp.mgrid[tuple(slice(n) for n in shape)]
200		if everywhere else
201		np.squeeze(xp.mgrid[BoundedSlices(where,shape)],
202		tuple(1+i for i,s in enumerate(where) if not isinstance(s,slice)) ) )
203	grid = grid.reshape( (ndim,) + (1,)*odim+grid.shape[1:])
204
205	if not everywhere and not uniform:
206		offset = offset[(slice(None),)*(1+odim)+where]
207
208	neigh = grid + (offset.reshape(offset.shape + (1,)*ndim) if uniform else offset)
209	bound = xp.array(shape,dtype=neigh.dtype).reshape((ndim,)+(1,)*(neigh.ndim-1))
210
211	if mode=='wrap': neigh = np.mod(neigh,bound); inside=True
212	else: inside = np.logical_and(np.all(neigh>=0,axis=0),np.all(neigh<bound,axis=0))
213
214	neighIndex = np.ravel_multi_index(neigh, shape, mode=mode)
215	return neighIndex, inside
216
217def TakeAtOffset(u,offset, padding=np.nan, **kwargs):
218	"""
219	- padding: outside fill value. Set padding=None for periodic boundary conditions
220	- **kwargs : passed to OffsetToIndex
221	"""
222	mode = 'wrap' if padding is None else 'clip'
223	neighIndex, inside = OffsetToIndex(u.shape,offset,mode=mode, **kwargs)
224
225	values = u.reshape(-1)[neighIndex]
226	if padding is not None:
227		if ad.isndarray(values):
228			values[np.logical_not(inside)] = padding
229		elif not inside:
230			values = padding
231	return values
232
233def AlignedSum(u,offset,multiples,weights,**kwargs):
234	"""
235	Returns sum along the direction offset, with specified multiples and weights.
236	Input : 
237	 - kwargs : passed to TakeAtOffset
238	"""
239	return sum(TakeAtOffset(u,mult*ad.asarray(offset),**kwargs)*weight 
240		for mult,weight in zip(multiples,weights))
241
242
243# --------- Degenerate elliptic finite differences -------
244
245def Diff2(u,offset,gridScale=1.,order=2,**kwargs):
246	"""
247	Approximates <offset, (d^2 u) offset> with specified accuracy order.
248	When order=2, corresponds to the (negated) degenerate elliptic scheme
249	$$
250		<e,(d^2 u) e> = (u(x+h*e)-2*u(x)+u(x-h*e))/h^2 + O(h^2),
251	$$
252	where e = offset and h=gridScale.
253	Input : 
254	 - u : array of shape (n1,...,nd) where n1,...,nd are the domain dimensions
255	 - offset : array with integer entries, of shape (d,k1,...,kr,n1,...,nd) or (d,k1,...,kr)
256	   where d is the number of dimensions in the domain, and k1,...,kr are arbitrary.
257	 - gridscale (optional) : scale of the discretization grid used in the finite differences
258	 - order (optional) : approximation order of the finite differences
259	 - kwargs : passed to AlignedSum
260	"""
261	if   order<=2: multiples,weights = (1,0,-1),(1.,-2.,1.)
262	elif order<=4: multiples,weights = (2,1,0,-1,-2),(-1/12.,4/3.,-15/6.,4/3.,-1/12.)
263	else: raise ValueError("Unsupported order") 
264	return AlignedSum(u,offset,multiples,np.array(weights)/gridScale**2,**kwargs)
265
266
267def DiffUpwind(u,offset,gridScale=1.,order=1,**kwargs):
268	"""
269	Approximates <grad u, offset> with specified accuracy order.
270	When order=1, corresponds the (negated) degenerate elliptic scheme
271	$$
272		 <grad u, e> = (u(x+h*e)-u(x))/h + O(h)
273	$$
274	where e = offset and h=gridScale.
275	Input : 
276	 - u : array of shape (n1,...,nd) where n1,...,nd are the domain dimensions
277	 - offset : array with integer entries, of shape (d,k1,...,kr,n1,...,nd) or (d,k1,...,kr)
278	   where d is the number of dimensions in the domain, and k1,...,kr are arbitrary.
279	 - gridscale (optional) : scale of the discretization grid used in the finite differences
280	 - order (optional) : approximation order of the finite differences
281	 - kwargs : passed to AlignedSum
282	"""
283	if   order==1: multiples,weights = (1,0),	( 1.,-1.)
284	elif order==2: multiples,weights = (2,1,0),  (-0.5,2.,-1.5)
285	elif order==3: multiples,weights = (3,2,1,0),(1./3.,-1.5,3.,-11./6.)
286	else: raise ValueError("Unsupported order")
287	return AlignedSum(u,offset,multiples,np.array(weights)/gridScale,**kwargs)
288
289# --------- Non-Degenerate elliptic finite differences ---------
290
291def DiffCentered(u,offset,gridScale=1.,order=2,**kwargs):
292	"""
293	Approximates <d u, offset> with specified accuracy order.
294	When order=2, corresponds to the centered finite difference
295	$$
296		 <grad u, e> = (u(x+h*e)-u(x-h*e))/(2h) + O(h^2)
297	$$
298	where e = offset and h=gridScale.
299		Input : 
300	 - u : array of shape (n1,...,nd) where n1,...,nd are the domain dimensions
301	 - offset : array with integer entries, of shape (d,k1,...,kr,n1,...,nd) or (d,k1,...,kr)
302	   where d is the number of dimensions in the domain, and k1,...,kr are arbitrary.
303	 - gridscale (optional) : scale of the discretization grid used in the finite differences
304	 - order (optional) : approximation order of the finite differences
305	 - kwargs : passed to AlignedSum
306	"""
307	if   order<=2: multiples,weights = ( 1,-1),( 1.,-1.)
308	elif order<=4: multiples,weights = ( 2, 1,-1,-2),(-1/6., 4/3.,-4/3., 1/6.)
309	else: raise ValueError("Unsupported order")
310	return AlignedSum(u,offset,multiples,np.array(weights)/(2*gridScale),**kwargs)
311
312def DiffCross(u,offset0,offset1,gridScale=1.,order=2,**kwargs):
313	"""
314	Approximates <offsets0, (d^2 u) offset1> with second order accuracy.
315	Centered finite differences scheme, but lacking the degenerate ellipticity property.
316	"""
317	if   order<=2: multiples,weights = ( 1,-1),(1.,1.)
318	elif order<=4: multiples,weights = ( 2, 1,-1,-2),(-1/12.,4/3.,4/3.,-1/12.)
319	else: raise ValueError("Unsupported order")
320	weights = np.array(weights)/(4*gridScale**2)
321	return (AlignedSum(u,offset0+offset1,multiples,weights,**kwargs) 
322		- AlignedSum(u,offset0-offset1,multiples,weights,**kwargs) )
323
324# ------------ Composite finite differences ----------
325
326def AxesOffsets(u=None,offsets=None,dimension=None):
327	"""
328	Returns the offsets corresponding to the axes.
329		Inputs : 
330	 - offsets (optional). Defaults to np.eye(dimension)
331	 - dimension (optional). Defaults to u.ndim
332	"""
333	if offsets is None:
334		if dimension is None:
335			dimension = u.ndim
336		offsets = np.eye(dimension).astype(int)
337	return offsets
338
339
340
341def DiffHessian(u,offsets=None,dimension=None,**kwargs):
342	"""
343	Approximates the matrix (<offsets[i], (d^2 u) offsets[j]> )_{ij}, using AxesOffsets as offsets.
344	Centered and cross finite differences are used, thus lacking the degenerate ellipticity property.
345	"""
346	from . import Metrics
347	offsets=AxesOffsets(u,offsets,dimension)
348	return Metrics.misc.expand_symmetric_matrix([
349		Diff2(u,offsets[i],**kwargs) if i==j else DiffCross(u,offsets[i],offsets[j],**kwargs) 
350		for i in range(len(offsets)) for j in range(i+1)])
351
352def DiffGradient(u,offsets=None,dimension=None,**kwargs):
353	"""
354	Approximates the vector (<d u, offsets[i]>)_i, using AxesOffsets as offsets
355	Centered finite differences are used, thus lacking the degerate ellipticity property.
356	"""
357	return DiffCentered(u,AxesOffsets(u,offsets,dimension),**kwargs)
358
359def DiffEll(u,offset,gridScale=1.0,order=2,α=0.,**kwargs):
360	"""
361	Helper function for discretizing (<grad u,e> - α)^2, which appears in elliptic energies. Returns 
362	$$
363	    [u(x)-u(x-h*e)-αh, u(x+h*e)-u(x)-αh]/(h sqrt(2))
364	$$
365	if order=2. Sum the squares to approximate (<grad u,e>-α)^2. Also accepts order=4.
366	"""
367	assert order in (2,4); s2 = np.sqrt(2); s6 = np.sqrt(6)
368	if order==2: return ad.array([(-DiffUpwind(u,-offset,gridScale,**kwargs)-α)/s2,
369	                              ( DiffUpwind(u, offset,gridScale,**kwargs)-α)/s2])
370	if order==4: return ad.array([(-DiffUpwind(u,-offset,gridScale,order=2,**kwargs)-α)/s6,
371	                              (DiffCentered(u,offset,gridScale,        **kwargs)-α)*(2/s6),
372                                  ( DiffUpwind(u, offset,gridScale,order=2,**kwargs)-α)/s6])
373
374# ----------- Interpolation ---------
375
376def UniformGridInterpolator1D(bounds,values,mode='clip',axis=-1):
377	"""
378	Interpolation on a uniform grid. mode is in ('clip','wrap', ('fill',fill_value) )
379	"""
380	val = values.swapaxes(axis,0)
381	fill_value = None
382	if isinstance(mode,tuple):
383		mode,fill_value = mode		
384	def interp(position):
385		endpoint=not (mode=='wrap')
386		size = val.size
387		index_continuous = (size-int(endpoint))*(position-bounds[0])/(bounds[-1]-bounds[0])
388		index0 = np.floor(index_continuous).astype(int)
389		index1 = np.ceil(index_continuous).astype(int)
390		index_rem = index_continuous-index0
391		
392		fill_indices=False
393		if mode=='wrap':
394			index0=index0%size
395			index1=index1%size
396		else: 
397			if mode=='fill':
398				 fill_indices = np.logical_or(index0<0, index1>=size)
399			index0 = np.clip(index0,0,size-1) 
400			index1 = np.clip(index1,0,size-1)
401		
402		index_rem = index_rem.reshape(index_rem.shape+(1,)*(val.ndim-1))
403		result = val[index0]*(1.-index_rem) + val[index1]*index_rem
404		if mode=='fill': result[fill_indices] = fill_value
405		result = np.moveaxis(result,range(position.ndim),range(-position.ndim,0))
406		return result
407	return interp
408
409# def AxesOrderingBounds(grid):
410# 	dim = len(grid)
411# 	lbounds = grid.__getitem__((slice(None),)+(0,)*dim)
412# 	ubounds = grid.__getitem__((slice(None),)+(-1,)*dim)
413
414# 	def active(i):
415# 		di = grid.__getitem__((slice(None),)+(0,)*i+(1,)+(0,)*(dim-1-i))
416# 		return np.argmax(np.abs(di-lbounds))
417# 	axes = tuple(active(i) for i in range(dim))
418
419# 	return axes,lbounds,ubounds