agd.Eikonal.DictIn

  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"""
  5The Eikonal.dictIn class is used to hold the input parameters to the eikonal solvers of 
  6the HFM library, CPU based or GPU based.
  7"""
  8
  9import numpy as np
 10from collections.abc import MutableMapping
 11from .. import AutomaticDifferentiation as ad
 12from .. import FiniteDifferences as fd
 13from .. import LinearParallel as lp
 14from .. import Metrics
 15from . import run_detail
 16from . import DictIn_detail
 17
 18_array_float_fields = {
 19	'origin','dims','gridScale','gridScales','values',
 20	'seeds','seeds_Unoriented','tips','tips_Unoriented',
 21	'seedValues','seedValues_Unoriented','seedValueVariation','seedFlags',
 22	'cost','speed','costVariation',
 23	'inspectSensitivity','inspectSensitivityWeights','inspectSensitivityLengths',
 24	'exportVoronoiFlags','factoringIndexShift',
 25	'chart_mapping',
 26	'control','controls','state_transition_costs',
 27}
 28
 29# Alternative key for setting or getting a single element
 30_singleIn = {
 31	'seed':'seeds','seedValue':'seedValues',
 32	'seed_Unoriented':'seeds_Unoriented','seedValue_Unoriented':'seedValues_Unoriented',
 33	'tip':'tips','tip_Unoriented':'tips_Unoriented',
 34	'seedFlag':'seedFlags','seedFlag_Unoriented':'seedFlags_Unoriented',
 35}
 36
 37_readonlyIn = {
 38	'float_t'
 39}
 40
 41_array_module = {
 42	'cpu':'numpy','cpu_raw':'numpy','gpu_transfer':'numpy',
 43	'gpu':'cupy','cpu_transfer':'cupy',
 44}
 45
 46_singleOut = {
 47	'geodesic':'geodesics','geodesic_Unoriented':'geodesics_Unoriented',
 48	'geodesic_euclideanLength':'geodesics_euclideanLength',
 49}
 50
 51SEModels = {'ReedsShepp2','ReedsSheppForward2','Elastica2','Dubins2',
 52'ReedsSheppExt2','ReedsSheppForwardExt2','ElasticaExt2','ElasticaExt2_5','DubinsExt2',
 53'ConvexReedsSheppForward2','ConvexElastica2','ConvexDubins2',
 54'ReedsShepp3','ReedsSheppForward3'}
 55
 56# These models do not follow the usual dimension naming convention (physical dimension last)
 57dimModels = {'ElasticaExt2_5':3,'Riemann3_Periodic':3,'ReedsSheppGPU3':5} 
 58
 59class dictOut(MutableMapping):
 60	"""
 61	A dictionnary like structure used as output of the Eikonal solvers. 
 62	"""
 63
 64	def __init__(self,store=None):
 65		self.store=store
 66
 67	def __copy__(self): return dictOut(self.store.copy())
 68	def copy(self): return self.__copy__()
 69
 70	def __repr__(self):
 71		return f"dictOut({self.store})"
 72
 73	def __setitem__(self, key, value):
 74		if key in _singleOut: key = _singleOut[key]; value = [value]
 75		self.store[key] = value
 76
 77	def __getitem__(self, key): 
 78		if key in _singleOut:
 79			values = self.store[_singleOut[key]]
 80			if len(values)!=1: 
 81				raise ValueError(f"Found {len(values)} values for key {key}")
 82			return values[0]
 83		return self.store[key]
 84
 85	def __delitem__(self, key): 
 86		key = _singleOut.get(key,key)
 87		del self.store[key]
 88
 89	def __iter__(self): return iter(self.store)
 90	def __len__(self):  return len(self.store)
 91	def keys(self): 
 92		"""
 93		The keys of the dictionary-like structure.
 94		"""
 95		return self.store.keys()
 96
 97def CenteredLinspace(a,b,n):
 98	"""
 99	Returns a linspace shifted by half a node length.
100	Inputs : 
101	 - a,b : interval endpoints
102	 - n : number of points
103	"""
104	n_=int(n); assert(n==n_) #Allow floats for convenience
105	r,dr=np.linspace(a,b,n_,endpoint=False,retstep=True)
106	if np.any(np.isnan(dr)): assert n==1; dr=b-a #Cupy 8.6 bug
107	return r+dr/2
108
109
110class dictIn(MutableMapping):
111	"""
112	A dictionary like structure used as input of the Eikonal solvers.
113
114	__init__ arguments: 
115	- store : a dictionary, used for initialization.
116
117	See dictIn().RunHelp() for details on the eikonal solver inputs.
118	"""
119
120	default_mode = 'cpu' # class attribute
121
122	def __init__(self, store=None):
123		if store is None: store=dict()
124		self.store = {'arrayOrdering':'RowMajor'}
125		if 'mode' in store:
126			mode = store['mode']
127			self.store['mode']=mode
128		else:
129			mode = dictIn.default_mode
130		assert mode in ('cpu','cpu_raw','cpu_transfer','gpu','gpu_transfer')
131		self._mode = mode
132		if self.mode in ('gpu','cpu_transfer'):
133			import cupy as cp
134			self.xp = cp
135			caster = lambda x : cp.asarray(x,dtype=float_t)
136			self.array_float_caster = lambda x : ad.asarray(x,caster=caster) #cupy >= 8.6
137		else: 
138			self.xp = np
139			self.array_float_caster = lambda x : np.asarray(x,dtype=float_t)
140		
141		if self.mode in ('gpu','cpu_transfer','gpu_transfer'): float_t=np.float32
142		else: float_t=np.float64
143
144		if 'float_t' in store:
145			float_t = store['float_t']
146			self.store['float_t']=float_t
147
148		self._float_t = float_t
149		if store: self.update(store)
150
151	def __copy__(self):     return dictIn(self.store.copy())
152	def __deepcopy__(self): return dictIn(self.store.deepcopy())
153	def copy(self): 
154		"""
155		Returns a shallow copy of the structure.
156		"""
157		return self.__copy__()
158
159	@property
160	def mode(self): 
161		"""
162		The running mode of the eikonal solver, see the Run method.
163		The input data must be provided using a compatible array module: 
164		numpy in 'cpu' mode, cupy in 'gpu' mode.
165
166		Supported running mode for the eikonal solver : 
167		- 'cpu' : Run algorithm on host, store data on host
168		- 'cpu_transfer' : Run algorithm on host, store data on device
169		- 'cpu_raw' : Raw call to the HFM CPU library (debug purposes)
170		- 'gpu' : Run algorithm on device, store data on device
171		- 'gpu_transfer' : Run algorithm on device, store data on host
172		"""
173		return self._mode
174	@property
175	def float_t(self):
176		"""
177		The floating point type of the data arrays. Typically np.float64 in 'cpu' mode, 
178		and np.float32 in 'gpu' mode.
179		"""
180		return self._float_t
181	
182	def __repr__(self): 
183		return f"dictIn({self.store})"
184
185	def __setitem__(self, key, value):
186		"""
187		Set a key:value pair in the dictionary like structure.
188
189		Special treatment of keys and values:
190		- The values associated to some keys are converted, if needed, to numeric 
191		arrays of floats for the numpy or cupy module. (key in _array_float_fields) 
192		- The key value pair 'seed':x is converted into 'seeds':[x], (as for setting 
193		 multiple seeds, but their number is just one). (key in _singleIn)
194		- Some keys are readonly, changes will be rejected (key in _readonlyIn.)
195		"""
196		if key=='mode':
197			if _array_module[value]!=_array_module[self.mode]:
198				raise ValueError('Switching between modes with distinct array storage')
199			else: self._mode = value
200		if key in _readonlyIn and self.store[key]!=value: 
201			raise ValueError(f"Key {key} is readonly (set at init)") 
202		if key in _singleIn: 
203			key = _singleIn[key]; value = [value]
204		if key in _array_float_fields and not ad.isndarray(value):
205			value = self.array_float_caster(value)
206		self.store[key] = value
207
208	def __getitem__(self, key): 
209		"""
210		Get a value associated to a given key, in the dictionary like structure.
211
212		Special treatment of keys and values:
213		- self['seed'] = self['seeds'][0], if self['seeds'] has length one, 
214		and fails otherwise. (key in _singleIn)
215		"""
216		if key in _singleIn:
217			values = self.store[_singleIn[key]]
218			if len(values)!=1: 
219				raise ValueError(f"Found {len(values)} values for key {key}")
220			return values[0]
221		return self.store[key]
222
223	def __delitem__(self, key): 
224		key = _singleIn.get(key,key)
225		del self.store[key]
226
227	def __iter__(self): return iter(self.store)
228	
229	def __len__(self):  return len(self.store)
230
231	def keys(self): 
232		"""
233		The keys of this dictionary structure.
234		"""
235		return self.store.keys()
236
237	def RunHelp(self,mode=None):
238		"""
239		Help on the eikonal solver, depending on the running mode.
240		"""
241		if mode is None: mode = self.mode
242		if mode in ('cpu','cpu_raw','cpu_transfer'): 
243			help(run_detail.RunSmart)
244		else:
245			from . import HFM_CUDA
246			help(HFM_CUDA.RunGPU)
247
248	def Run(self,join=None,**kwargs):
249		"""
250		Calls the HFM library, prints log and returns output.
251		Inputs : 
252		- join (optional) : join the dictionary with these additional entries before running.
253		- **kwargs (optional) : passed to the run_detail.RunSmart or HFM_CUDA.RunGPU methods.
254		
255		See dictIn().RunHelp() for additional details, depending on the running mode.
256		"""
257		if join is not None:
258			other = self.copy()
259			other.update(join)
260			return other.Run(**kwargs)
261
262		if self['arrayOrdering']!='RowMajor': 
263			raise ValueError("Unsupported array ordering")
264		def to_dictOut(out):
265			if isinstance(out,tuple): return (dictOut(out[0]),) + out[1:]
266			else: return dictOut(out)
267			
268		if   self.mode=='cpu': return to_dictOut(run_detail.RunSmart(self,**kwargs))
269		elif self.mode=='cpu_raw': return to_dictOut(run_detail.RunRaw(self.store,**kwargs))
270		elif self.mode=='cpu_transfer':
271			cpuIn = ad.cupy_generic.cupy_get(self,dtype64=True,iterables=(dictIn,Metrics.Base))
272			cpuIn.xp = np; cpuIn._mode = 'cpu'; cpuIn['mode'] = 'cpu'
273			for key in list(cpuIn.keys()): 
274				if key.startswith('traits'): cpuIn.pop(key)
275			return to_dictOut(run_detail.RunSmart(cpuIn,**kwargs))
276		
277		from . import HFM_CUDA
278		if   self.mode=='gpu': return to_dictOut(HFM_CUDA.RunGPU(self,**kwargs))
279		elif self.mode=='gpu_transfer':
280			gpuStoreIn = ad.cupy_generic.cupy_set(self.store, # host->device
281				dtype32=True, iterables=(dict,Metrics.Base))
282			gpuIn = dictIn({**gpuStoreIn,'mode':'gpu'})
283			gpuOut = HFM_CUDA.RunGPU(gpuIn)
284			cpuOut = ad.cupy_generic.cupy_get(gpuOut,iterables=(dict,list))
285			return to_dictOut(cpuOut) # device->host
286
287	# ------- Grid related functions ------
288
289	@property
290	def shape(self):
291		"""
292		The shape of the discretization grid.
293		"""
294		dims = self['dims']
295		if ad.cupy_generic.from_cupy(dims): dims = dims.get()
296		return tuple(dims.astype(int))
297
298	@property
299	def size(self): 
300		"""
301		The number of points in the discretization grid.
302		"""
303		return np.prod(self.shape)
304	
305	@property
306	def SE(self):
307		"""
308		Wether the model is based on the Special Euclidean group.
309		True for curvature penalized models.
310		"""
311		return self['model'] in SEModels	
312	
313	@property
314	def vdim(self):
315		"""
316		The dimension of the ambient vector space.
317		"""
318		model = self['model']
319		if model in dimModels: return dimModels[model]
320		dim = int(model[-1])
321		return (2*dim-1) if self.SE else dim
322
323	@property
324	def nTheta(self):
325		"""
326		Number of points for discretizing the interval [0,2 pi], in the angular space 
327		discretization, for the SE models (a.k.a. curvature penalized models).
328		"""
329		if not self.SE: raise ValueError("Not an SE model")
330		shape = self.shape
331		if self.vdim!=len(self.shape): raise ValueError("Angular resolution not set")
332		n = shape[-1]
333		return (2*n) if self.get('projective',False) and self.vdim==3 else n
334
335	@nTheta.setter
336	def nTheta(self,value):
337		if not self.SE: raise ValueError("Not an SE model")
338		shape = self.shape
339		vdim = self.vdim
340		projective = self.get('projective',False)
341		if vdim==len(shape): shape=shape[:int((vdim+1)/2)] #raise ValueError("Angular resolution already set")
342		if   vdim==3: self['dims'] = (*shape, value/2 if projective else value) 
343		elif vdim==5: self['dims'] = (*shape, value/4 if projective  else value/2, value) 
344
345	@property
346	def gridScales(self):
347		"""
348		The discretization grid scale along each axis.
349		"""
350		if self.SE:
351			h = self['gridScale']
352			hTheta = 2.*np.pi / self.nTheta
353			if self.vdim==3: return self.array_float_caster( (h,h,hTheta) )
354			else: return self.array_float_caster( (h,h,h,hTheta,hTheta) )
355		elif 'gridScales' in self: return self['gridScales']
356		else: return self.array_float_caster((self['gridScale'],)*self.vdim)	
357		
358	@property
359	def corners(self):
360		"""
361		Returns the extreme points grid[:,0,...,0] and grid[:,-1,...,-1] of the 
362		discretization grid.
363		"""
364		gridScales = self.gridScales
365		dims = self['dims']
366		physdim = (len(dims)+1)//2 if self.SE else len(dims)
367		origin = self.get('origin',self.xp.zeros_like(dims[:physdim]))
368		if self.SE: 
369			tail = self.array_float_caster((-gridScales[-1]/2,)*(len(dims)-len(origin)))
370			origin = np.concatenate((origin,tail))
371		return (origin,origin+gridScales*dims)
372
373	def Axes(self,dims=None):
374		"""
375		The discretization points used along each coordinate axis.
376		"""
377		bottom,top = self.corners
378		if dims is None: dims=self['dims']
379		return [self.array_float_caster(CenteredLinspace(b,t,d)) 
380			for b,t,d in zip(bottom,top,dims)]
381
382	def Grid(self,dims=None):
383		"""
384		Returns a grid of coordinates, containing all the discretization points of the domain.
385		Similar to np.meshgrid(*self.Axes(),indexing='ij')
386		Inputs : 
387		- dims(optional) : use a different sampling of the domain
388		"""
389		axes = self.Axes(dims);
390		ordering = self['arrayOrdering']
391		if ordering=='RowMajor': return ad.array(np.meshgrid(*axes,indexing='ij',copy=False))
392		elif ordering=='YXZ_RowMajor': return ad.array(np.meshgrid(*axes,copy=False))
393		else: raise ValueError('Unsupported arrayOrdering : '+ordering)
394
395	def SetUniformTips(self,dims):
396		"""
397		Place regularly spaced tip points all over the domain, 
398		from which to backtrack minimal geodesics.
399		Inputs : 
400		- dims : number of tips to use along each dimension.
401		"""
402		self['tips'] = self.Grid(dims).reshape(self.vdim,-1).T
403
404	def SetRect(self,sides,sampleBoundary=False,gridScale=None,gridScales=None,
405		dimx=None,dims=None):
406		"""
407		Defines a box domain, for the HFM library.
408		Inputs:
409		- sides, e.g. ((a,b),(c,d),(e,f)) for the domain [a,b]x[c,d]x[e,f]
410		- sampleBoundary : switch between sampling at the pixel centers, and sampling including the boundary
411		- gridScale, gridScales : side h>0 of each pixel (alt : axis dependent)
412		- dimx, dims : number of points along the first axis (alt : along all axes)
413		"""
414		# Ok to set a new domain, or completely replace the domain
415		domain_count = sum(e in self for e in ('gridScale','gridScales','dims','origin'))
416		if domain_count not in (0,3): raise ValueError("Domain already partially set")
417		
418		caster = self.array_float_caster
419		corner0,corner1 = caster(sides).T
420		dim = len(corner0)
421		sb=float(sampleBoundary)
422		width = corner1-corner0
423		if gridScale is not None: 
424			gridScales=[gridScale]*dim; self['gridScale']=gridScale
425		elif gridScales is not None:
426			self['gridScales']=gridScales
427		elif dimx is not None:
428			gridScale=width[0]/(dimx-sb); gridScales=[gridScale]*dim; self['gridScale']=gridScale
429		elif dims is not None:
430			gridScales=width/(caster(dims)-sb); self['gridScales']=gridScales
431		else: 
432			raise ValueError('Missing argument gridScale, gridScales, dimx, or dims')
433
434		h=caster(gridScales)
435		ratios = (corner1-corner0)/h + sb
436		dims = np.round(ratios)
437		assert(np.min(dims)>0)
438		origin = corner0 + (ratios-dims-sb)*h/2
439		self['dims']   = dims
440		self['origin'] = origin
441
442	def PointFromIndex(self,index,to=False):
443		"""
444		Turns an index into a point.
445		Optional argument to: if true, inverse transformation, turning a point into a continuous index
446		"""
447		bottom,_ = self.corners
448		scale = self.gridScales
449		start = bottom +0.5*scale
450		index = self.array_float_caster(index)
451		assert index.shape[-1]==self.vdim
452		if not to: return start+scale*index
453		else: return (index-start)/scale
454
455	def IndexFromPoint(self,point):
456		"""
457		Returns the index that yields the position closest to a point, and the error.
458		"""
459		point = self.array_float_caster(point)
460		continuousIndex = self.PointFromIndex(point,to=True)
461		index = np.round(continuousIndex)
462		return index.astype(int),(continuousIndex-index)
463
464	def OrientedPoints(self,pointU):
465		"""
466		Appends all possible orientations to the point coordinates.
467		"""
468		pointU = self.array_float_caster(pointU)
469		if self['model'] not in SEModels: 
470			raise ValueError("OrientedPoints only makes sense for models SE space.")
471		if self.vdim!=3:
472			raise ValueError("Sorry, oriented point not implemented for SE(3) models.")
473		pdim = int((self.vdim+1)/2) # Number of physical dimensions
474		if pointU.shape[-1]!=pdim:
475			raise ValueError(f"Unoriented points expected to have {pdim} dimensions, "
476				f"found {pointU.shape[-1]}.")
477		theta = self.Axes()[2]
478		point = self.xp.full((len(theta),*pointU.shape[:-1],self.vdim),np.nan)
479		for i,t in enumerate(theta):
480			point[i,...,:pdim]=pointU
481			point[i,...,pdim:]=t
482		return point
483
484	def VectorFromOffset(self,offset,to=False):
485		"""
486		Turns a finite difference offset into a vector, by multiplying by the gridScale.
487		Inputs : 
488		- offset : the offset to convert.
489		- to (optional) : if True, produces an offset from a vector (reverse operation).
490		"""
491		offset = self.array_float_caster(offset)
492		if not to: return offset*self.gridScales
493		else: return offset/self.gridScales  
494
495	def GridNeighbors(self,point,gridRadius):
496		"""
497		Returns the neighbors around a point on the grid. 
498		Geometry last convention
499		Inputs: 
500		- point (array): geometry last
501		- gridRadius (scalar): given in pixels
502		"""
503		xp = self.xp
504		point = self.array_float_caster(point)
505		point_cindex = self.PointFromIndex(point,to=True)
506		aX = [xp.arange(int(np.floor(ci-gridRadius)),int(np.ceil(ci+gridRadius)+1),
507			dtype=self.float_t) for ci in point_cindex]
508		neigh_index =  xp.stack(xp.meshgrid( *aX, indexing='ij'),axis=-1)
509		neigh_index = neigh_index.reshape(-1,neigh_index.shape[-1])
510
511		# Check which neighbors are close enough
512		offset = neigh_index-point_cindex
513		close = np.sum(offset**2,axis=-1) < gridRadius**2
514
515		# Check which neighbors are in the domain (periodicity omitted)
516		neigh = self.PointFromIndex(neigh_index)
517		bottom,top = self.corners
518		inRange = np.all(np.logical_and(bottom<neigh,neigh<top),axis=-1)
519
520		return neigh[np.logical_and(close,inRange)]
521
522	SetFactor = DictIn_detail.SetFactor
523	SetSphere = DictIn_detail.SetSphere
524	@property
525	def factoringPointChoice(self): return DictIn_detail.factoringPointChoice(self)