agd.Interpolation

  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 file implements some spline interpolation methods, on uniform grids,
  6in a manner compatible with automatic differentiation.
  7
  8If none of the involved arrays use automatic differentiation, and if the options are 
  9compatible, then a bypass through ndimage may be used.
 10"""
 11
 12
 13import numpy as np
 14import itertools
 15import scipy.linalg
 16
 17from . import AutomaticDifferentiation as ad
 18
 19#TODO : use a smaller stencil (3 pts instead of 4) for 2nd degree interpolation 
 20# when far from the boundary, in non-periodic mode. (Introduce interior nodes.)
 21
 22def origin_scale_shape(grid):
 23	"""
 24	This function is indended for extracting the origin, scale, and shape,
 25	of a uniform coordinate system provided as a meshgrid.
 26	"""
 27	def _orig_step_len(a,axis):
 28		a = ad.asarray(a)
 29		assert a.ndim==1 or a.ndim>=axis
 30		if a.ndim>1:a=a.__getitem__((0,)*axis+(slice(None,),)+(0,)*(a.ndim-axis-1))
 31		return a[0],a[1]-a[0],len(a)
 32	origin_scale_shape = [_orig_step_len(a,axis) for axis,a in enumerate(grid)]
 33	origin,scale,shape = [list(l) for l in zip(*origin_scale_shape)]
 34	caster = ad.cupy_generic.array_float_caster(grid,iterables=(list,tuple))
 35	return caster(origin),caster(scale),tuple(shape)
 36
 37
 38def _append_dims(x,ndim):
 39	"""Appends specified number of singleton dimensions"""
 40	return np.reshape(x,x.shape+(1,)*ndim)
 41
 42def map_coordinates(input,coordinates,*args,
 43	grid=None,origin=None,scale=None,depth=0,order=1,**kwargs):
 44	"""
 45	Thin wrapper over the ndimage.map_coordinates function, which adds the possibility of 
 46	rescaling the coordinates using a reference grid, and interpolating tensors.
 47	Will dispatch to cupyx.scipy.ndimage.map_coordinates if needed.
 48
 49	Additional inputs : 
 50	- grid (optional) : reference coordinate system, which must be uniform
 51	- origin,scale (optional) : obtained from origin_scale_shape(grid)
 52	- depth : depth of interpolated objects 0->scalar, 1->vector, 2->matrix
 53	- order (optional) : set default 1 for better cupy/numpy reproducibility
 54	"""
 55	kwargs['order']=order
 56	if ad.cupy_generic.from_cupy(input):
 57		from cupyx.scipy.ndimage import map_coordinates as mc
 58		def map_coordinates(arr,x,*args,**kwargs):
 59			# Cupy (version 7.8) requires the coordinates array to be flattened
 60			shape = x.shape[1:]
 61			return mc(arr,x.reshape((len(x),-1)),*args,**kwargs).reshape(shape)
 62	else: from scipy.ndimage import map_coordinates
 63
 64	if grid is not None: origin,scale,_ = origin_scale_shape(grid)
 65	assert (origin is None) == (scale is None)
 66	if origin is None: return map_coordinates(input,coordinates,*args,**kwargs)
 67	origin,scale = (_append_dims(e,coordinates.ndim-1) for e in (origin,scale))
 68	y = (coordinates-origin)/scale
 69
 70	if depth==0: return map_coordinates(input,y,*args,**kwargs)
 71	oshape = input.shape[:depth]
 72	input = input.reshape((-1,)+input.shape[depth:])
 73	out = ad.array([map_coordinates(input_,y,*args,**kwargs) for input_ in input])
 74	out.reshape(oshape+y.shape[1:])
 75	return out
 76
 77class _spline_univariate:
 78	"""
 79	A univariate spline of a given order, with not-a-knot boundary conditions.
 80	"""
 81	def __init__(self,order,shape,periodic):
 82		assert isinstance(order,int)
 83		self.order=order
 84		if not (order>=1 and order<=3):
 85			raise ValueError("Unsupported spline order")
 86
 87		assert isinstance(shape,int) and shape>=0
 88		self.shape = shape
 89
 90		assert isinstance(periodic,bool)
 91		self.periodic=periodic
 92
 93
 94	def __call__(self,xa,xs):
 95		"""
 96		Weight at absolute position xa, of of spline centered at xs.
 97		"""
 98		xa=ad.asarray(xa)
 99		xs=ad.asarray(xs)
100		if   self.order==1: return self._call1(xa,xs)
101		elif self.order==2: return self._call2(xa,xs)
102		elif self.order==3: return self._call3(xa,xs)
103		else: assert False
104
105	def nodes(self,interior):
106		"""
107		Returns range(a,b) where [x+a,x+b] contains the support
108		 of the spline centered at some point x (interior or boundary) 
109		"""
110		if   self.order==1: return range(-1,1)
111		elif self.order==2: #return range(-1,2)
112			return range(-1,2) if interior else range(-2,2)
113		elif self.order==3: 
114			return range(-2,2) if interior else range(-3,4)
115		assert False
116
117	def interior(self,x,tol=None):
118		"""
119		Wether the interior nodes can be used, or one should fall back to boundary nodes.
120		"""
121		if tol is None: tol = ad.precision(x) * np.max(self.shape)
122		if np.any(x<-tol) or np.any(x>self.shape-1+tol):
123			raise ValueError("Interpolating data outside domain")
124		if self.order==1 or self.periodic:
125			return np.full_like(x,True,dtype='bool')
126		elif self.order==2:
127			return x>=1
128		elif self.order==3:
129			return np.logical_and(x>=1,x<=self.shape-2)
130
131	def _call1(self,xa,xs):
132		"""
133		A piecewise linear spline, defined over [-1,1].
134		"""
135		x=xa-xs.astype(xa.dtype) # Avoid float32 + int32 -> float64 cast on GPU
136		result = np.zeros_like(x)
137		seg=ad.asarray(np.floor(x+1))
138		
139		# relative interval [-1,0[
140		pos = seg==0
141		result[pos] = 1+x[pos]
142
143		# relative interval [0,1[
144		pos = seg==1
145		result[pos] = 1-x[pos]
146
147		return result
148
149		# Rejected because derivative at node points is inconsistent accross nodes
150		# return np.maximum(0.,1.-np.abs(x))
151
152	def _call2(self,xa,xs):
153		"""
154		A piecewise quadratic spline function, defined over [-1,2]
155		"""
156		x=ad.asarray(xa-xs)
157		result = np.zeros_like(x)
158
159
160		# Which spline segment to use
161		seg = ad.asarray(np.floor(x+1))
162
163		if not self.periodic:
164			# Implements the not-a-knot boundary condition
165			pos = np.logical_and(xa<=1,xs<=3)
166			seg[pos] = 2-xs[pos]
167
168		# relative interval [-1,0[
169		pos = seg==0
170		x_ = x[pos]+1
171		result[pos] = x_**2
172
173		# relative interval [0,1[
174		pos = seg==1
175		x_ = x[pos]-0.5
176		result[pos] = 1.5 - 2.*x_**2
177
178		# relative interval [1,2[
179		pos = seg==2
180		x_ = 2.-x[pos]
181		result[pos] = x_**2 
182
183		return result
184
185	def _call3(self,xa,xs):
186		"""
187		A piecewise cubic spline function, defined over [-2,2]
188		"""
189		x=ad.asarray(xa-xs)
190		result = np.zeros_like(x)
191		s=self.shape-1
192
193		# Which spline segment to use
194		seg = ad.asarray(np.floor(x+2))
195		if not self.periodic:
196			# Implements the not-a-knot boundary condition
197			pos = np.logical_and(xa<=1,xs<=3)
198			seg[pos] = 3-xs[pos]
199			pos = np.logical_and(xa>=s-1,xs>=s-3)
200			seg[pos] = self.shape-1-xs[pos]
201
202		def f0(y): return y**3
203		def f1(y): return 1+3*y+3*y**2-3.*y**3
204
205		# relative interval [-2,-1[
206		pos = seg==0
207		result[pos] = f0(x[pos]+2.)
208
209		# relative interval [-1,0[
210		pos = seg==1
211		result[pos] = f1(x[pos]+1.)
212
213		# relative interval [0,1[
214		pos = seg==2
215		result[pos] = f1(1.-x[pos])
216
217		# relative interval [1,2[
218		pos = seg==3
219		result[pos] = f0(2.-x[pos])
220		
221		"""
222		if not self.periodic:
223			# End of not-a-knot boundary condition
224			def g(y): return 4.-6.*y**2+(50./27.)*y**3
225			xa=ad.broadcast_to(xa,x.shape)
226			pos = np.logical_and(xa<=3,xs==3)
227			result[pos] = g(3.-xa[pos])
228			pos = np.logical_and(xa>s-3,xs==s-3)
229			result[pos] = g(xa[pos]-(s-3))
230		"""
231
232		return result
233
234	def _band(self):
235		rg = np.arange(self.shape+0.)
236		if self.order==2:
237			band_tr = np.stack((self(rg,rg-1),self(rg,rg),self(rg,rg+1),self(rg,rg+2)),axis=0)
238			return _banded_transpose((2,1),band_tr)
239		elif self.order==3:
240			band_tr = np.stack((self(rg,rg-3),self(rg,rg-2),self(rg,rg-1),
241				self(rg,rg),self(rg,rg+1),self(rg,rg+2),self(rg,rg+3)),axis=0)
242			return _banded_transpose((3,3),band_tr)
243		else: assert False
244
245
246	def make_coefs(self,values,overwrite_values=False):
247		"""
248		Produces the node coefficients corresponding to given values.
249		!! Call convention : interpolation is along the first axis. !!
250		"""
251		if self.order==1: return values
252		if self.periodic: raise ValueError("Periodic interpolation is not supported for degree > 1")
253		assert len(values)==self.shape
254
255		def solver(vals): return scipy.linalg.solve_banded(*self._band(),vals,
256				overwrite_ab=True,overwrite_b=overwrite_values)
257
258		if ad.is_ad(values): return values.apply_linear_operator(solver)
259#			raise ValueError("AD interpolation is not supported for degree > 1")
260
261		return solver(values)
262		
263def _banded_transpose(lu,t):
264	l,u=lu
265	return (u,l),np.stack(tuple(np.roll(t[i],i-u) for i in reversed(range(l+u+1))))
266#	return np.stack((np.roll(t[2],1),t[1],np.roll(t[0],-1)),axis=0)
267
268def _banded_densify(lu,t):
269	"""Turn banded matrix in to dense matrix (inefficient, for testing)"""
270	l,u=lu
271	n = t.shape[1]
272	mat = np.zeros((n,n))
273	for i in range(n):
274		for j in range(n):
275			k=u+i-j
276			if 0<=k and k<len(t):
277				mat[i,j]=t[k,j]
278	return mat
279
280class _spline_tensor:
281	"""
282	A tensor product of univariate splines.
283	"""
284	def __init__(self,orders,shape,periodic):
285		assert all(isinstance(t,tuple) for t in (orders,shape,periodic))
286		self.splines = tuple(_spline_univariate(order,s,per) 
287			for (order,s,per) in zip(orders,shape,periodic))
288		assert all(len(t)==self.vdim for t in (orders,shape,periodic))
289
290	@property
291	def order(self):
292		return tuple(spline.order for spline in self.splines)
293	@property
294	def shape(self):
295		return tuple(spline.shape for spline in self.splines)
296	@property
297	def periodic(self):
298		return tuple(spline.periodic for spline in self.splines)
299	
300	@property
301	def vdim(self):
302		return len(self.splines)
303
304	def __call__(self,xa,xs):
305		"""
306		Weight at absolute position xa, of of spline centered at xs.
307		"""
308		return np.prod( ad.asarray(tuple(spline(xai,xsi) 
309			for (xai,xsi,spline) in zip(xa,xs,self.splines))), axis=0)
310
311	def nodes(self,interior):
312		"""
313		The spline at x (interior or boundary) is supported 
314		on the union of x+node+[0,1]**vdim
315		"""
316		_nodes = tuple(spline.nodes(interior) for spline in self.splines)
317		return np.asarray(tuple(itertools.product(*_nodes))).T
318
319	def interior(self,x):
320		"""
321		Wether the interior nodes can be used.
322		"""
323		return ad.asarray(tuple(spline.interior(xi) 
324			for (spline,xi) in zip(self.splines,x))).all(axis=0)
325#		return np.logical_and.reduce(tuple(spline.interior(xi) # cupy has no reduce
326#			for (spline,xi) in zip(self.splines,x)))
327
328	def make_coefs(self,values,overwrite_values=False):
329		"""
330		Produces the node coefficients corresponding to given values.
331		!! Call convention : interpolation is along the first axes. 
332		"""
333		for i,spline in enumerate(self.splines): #reversed(list(enumerate(self.splines))):
334			values = np.moveaxis(spline.make_coefs(
335				np.moveaxis(values,i,0),overwrite_values=overwrite_values),0,i)
336		return values
337
338class UniformGridInterpolation:
339	"""
340	Interpolates values on a uniform grid, in arbitrary dimension, using splines of 
341	a given order.
342	"""
343
344	def __init__(self,grid,values=None,order=1,periodic=False,check_grid=True):
345		"""
346		- grid (ndarray) : must be a uniform grid. E.g. np.meshgrid(aX,aY,indexing='ij')
347		 where aX,aY have uniform spacing. Alternatively, provide only the axes.
348		- values (ndarray) : interpolated values.
349		- order (int, tuple of ints) : spline interpolation order (<=3), along each axis.
350		- periodic (bool, tuple of bool) : wether periodic interpolation, along each axis.
351		"""
352		if isinstance(grid,dict):
353			self.shape  = grid['shape']
354			self.origin = ad.asarray(grid['origin'])
355			self.scale  = ad.asarray(grid['scale'])
356			if grid.get('cell_centered',False): 
357				self.origin += self.scale/2 # Convert to node_centered origin
358		else:
359			self.origin,self.scale,self.shape = origin_scale_shape(grid)
360			if check_grid and grid[0].ndim>1:
361				assert np.allclose(grid,self._grid(),atol=1e-5) #Atol allows float32 types
362
363		if order is None: order = 1
364		if isinstance(order,int): order = (order,)*self.vdim
365
366		if periodic is None: periodic=False
367		if isinstance(periodic,bool): periodic = (periodic,)*self.vdim
368		self.periodic=periodic
369
370		self.spline = _spline_tensor(order,self.shape,periodic)
371		assert self.spline.vdim == self.vdim
372		self.interior_nodes = self.spline.nodes(interior=True)
373		self.boundary_nodes = self.spline.nodes(interior=False)
374
375		self.coef = None if values is None else self.make_coefs(values)
376
377	@property
378	def vdim(self):
379		"""
380		Vector dimension of the interpolation points.
381		"""
382		return len(self.shape)
383	@property
384	def oshape(self):
385		"""
386		Number of dimensions of the interpolated objects.
387		"""
388		return self.coef.shape[self.vdim:]
389
390	def __call__(self,x,interior=None):
391		"""
392		Interpolates the data at the position x.
393		"""
394		x=ad.asarray(x)
395		assert len(x)==self.vdim
396		
397		pdim = x.ndim-1 # Number of dimensions of position
398		origin,scale = (_append_dims(e,pdim) for e in (self.origin,self.scale))
399
400		# Separate treatment of interior and boundary points
401		if interior is None:
402			y = (ad.remove_ad(x) - origin)/scale
403			interior_x = self.spline.interior(y)
404			boundary_x = np.logical_not(interior_x)
405			interior_result = self(x[:,interior_x],True)
406			boundary_result = self(x[:,boundary_x],False)
407
408			result_shape = self.oshape+x.shape[1:]
409			if result_shape==tuple(): #numpy zeros_like has a bug for empty shapes
410				some_result = interior_result if interior_result.size>0 else boundary_result
411				result = np.zeros_like(some_result.reshape(-1)[0])
412			else: result = np.zeros_like(interior_result,shape=result_shape)
413
414			try:
415				result[...,interior_x] = interior_result
416				result[...,boundary_x] = boundary_result
417			except (ValueError,IndexError):
418				# Some cupy versions do not handle Ellipsis correctly
419				ellipsis = (slice(None),)*len(self.oshape)
420				result.__setitem__((*ellipsis,interior_x),interior_result)
421				result.__setitem__((*ellipsis,boundary_x),boundary_result)
422
423			return result
424
425		# Rescale the coordinates in reference rectangle
426		y = np.expand_dims((x - origin)/scale,axis=1)
427		# Bottom left pixel
428		yf = np.floor(y).astype(int)
429		# All pixels supporting an active spline
430		nodes = self.interior_nodes if interior else self.boundary_nodes
431		ys = yf - _append_dims(nodes,pdim)
432		
433		# Spline coefficients, taking care of out of domain 
434		ys_ = ys.copy()
435		out = np.full_like(ad.remove_ad(x[0]),False,dtype=bool)
436		for i,(d,per) in enumerate(zip(self.shape,self.periodic)):
437			if per: 
438				ys_[i] = ys_[i]%d
439			else: 
440				bad = np.logical_or(ys_[i]<0,ys_[i]>=d)
441				out = np.logical_or(out,bad)
442				try: 
443					ys_[i,bad] = 0 
444				except ValueError: # Old cupy versions do not support such slices
445					ys_i = ys_[i]
446					ys_i[bad] = 0
447					ys_[i] = ys_i
448
449		coef = self.coef[tuple(ys_)]
450		coef[out]=0.
451		ondim = len(self.oshape)
452		coef = np.moveaxis(coef,range(-ondim,0),range(ondim))
453		
454		# Spline weights
455		weight = self.spline(y,ys)
456		return (coef*weight).sum(axis=ondim)
457
458	def set_values(self,values):
459		self.coef = self.make_coefs(values)
460
461	def make_coefs(self,values,overwrite_values=False):
462		values = ad.asarray(values)
463		xp = ad.cupy_generic.get_array_module(values)
464		self.interior_nodes = xp.array(self.interior_nodes)
465		self.boundary_nodes = xp.array(self.boundary_nodes)
466
467		ondim = values.ndim - self.vdim
468		# Internally, interpolation is along the first axes.
469		# (Contrary to external interface)
470		val = np.moveaxis(values,range(ondim),range(-ondim,0))
471
472		return self.spline.make_coefs(val,overwrite_values=overwrite_values)
473
474	def _grid(self):
475		xp = ad.cupy_generic.get_array_module(self.origin)
476		dtype = self.origin.dtype
477		return ad.asarray(np.meshgrid(*(o+h*xp.arange(s+0.,dtype=dtype) 
478			for (o,h,s) in zip(self.origin,self.scale,self.shape)), indexing='ij'))