agd.Eikonal.DictIn_detail

  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 additional functionality of the Eikonal.dictIn class, related to 
  6source factorization, and to solving PDEs on spheres.
  7"""
  8
  9import numpy as np
 10from .. import AutomaticDifferentiation as ad
 11from .. import FiniteDifferences as fd
 12from .. import LinearParallel as lp
 13from .. import Metrics
 14
 15def factoringPointChoice(self):
 16	if 'factoringPointChoice' in self: # 'Key' and 'Seed' are synonyms here
 17		assert self['factoringPointChoice'] in ('Both','Key','Seed') 
 18		return self['factoringPointChoice']
 19	elif 'factoringRadius' not in self: return None
 20	elif self.get('order',1)==3: return 'Both'
 21	else: return 'Seed'
 22
 23
 24def SetFactor(self,radius=None,value=None,gradient=None):
 25	"""
 26	This function setups additive factorization around the seeds.
 27	Inputs (optional): 
 28	- radius.
 29		Positive number -> approximate radius, in pixels, of source factorization.
 30		-1 -> Factorization over all the domain. 
 31		None -> source factorization over all the domain
 32	- value (optional).
 33	  callable, array -> approximate values of the solution. 
 34	  None -> reconstructed from the metric.
 35	- gradient (optional) 
 36	  callable, array -> approximate gradient of the solution.
 37	  Obtained from the values by automatic differentiation if unspecified.
 38
 39	Outputs : the subgrid used for factorization
 40	Side effect : sets 'factoringValues', 'factoringGradients', 
 41	   and in the case of a subgrid 'factoringIndexShift'
 42	"""
 43	# Set the factoring grid points
 44	if radius is None:
 45		radius = self.get('factoringRadius',10)
 46
 47	if radius<0 or radius>=np.max(self.shape):
 48		factGrid = self.Grid()
 49		self.pop('factoringIndexShift',None)
 50	else:
 51		seed_ci =  self.PointFromIndex(self['seed'],to=True)
 52		bottom = [max(0,int(np.floor(ci)-radius)) for ci in seed_ci]
 53		top = [min(si,int(np.ceil(ci)+radius)+1) for ci,si in zip(seed_ci,self.shape)]
 54		self['factoringIndexShift'] = bottom
 55		aX = [self.xp.arange(bi,ti,dtype=self.float_t) for bi,ti in zip(bottom,top)]
 56		factGrid = ad.array(np.meshgrid(*aX,indexing='ij'))
 57		factGrid = np.moveaxis(self.PointFromIndex(np.moveaxis(factGrid,0,-1)),-1,0)
 58#			raise ValueError("dictIn.SetFactor : unsupported radius type")
 59
 60	# Set the values
 61	if value is None:
 62		value = self.get('factoringPointChoice','Seed')
 63
 64	if isinstance(value,str):
 65		seed = self['seed']
 66		if 'metric' in self: metric = self['metric']
 67		elif self['model'].startswith('Isotropic'): metric = Metrics.Isotropic(1)
 68		else: raise ValueError("dictIn.SetFactor error : no metric found")
 69		if 'cost' in self: metric = metric.with_cost(self['cost'])
 70		fullGrid = factGrid if factGrid.shape[1:]==self.shape else self.Grid()
 71
 72		diff = lambda x : x-fd.as_field(seed,x.shape[1:],depth=1)
 73		if value in ('Key','Seed'):
 74			metric.set_interpolation(fullGrid) # Default order 1 interpolation suffices
 75			value = metric.at(seed)
 76		elif value=='Current': # Not really working ?
 77			# Strictly speaking, we are not interpolating the metric, 
 78			# since the point x at which it is evaluated lies on the grid
 79			metric.set_interpolation(fullGrid) 
 80			value = lambda x : metric.at(x).norm(diff(x))
 81			#Cheating : not differentiating w.r.t position, but should be enough here
 82			gradient = lambda x: metric.at(x).gradient(diff(x)) 
 83		elif value=='Both':
 84			# Pray that AD goes well ...
 85			metric.set_interpolation(fullGrid,order=3) # Need to differentiate w.r.t x
 86			value = lambda x : 0.5*(metric.at(seed).norm(diff(x)) + 
 87				metric.at(x).norm(diff(x)))
 88		else:
 89			raise ValueError(f"dictIn.SetFactor error : unsupported "
 90				"factoringPointChoice : {value} .")
 91
 92	if callable(value):
 93		if gradient is None:
 94			factGrid_ad = ad.Dense.identity(constant=factGrid, shape_free=(self.vdim,))
 95			value = value(factGrid_ad)
 96		else:
 97			value = value(factGrid)
 98
 99	if isinstance(value,Metrics.Base):
100		factDiff = factGrid - fd.as_field(self['seed'],factGrid.shape[1:],depth=1)
101		if gradient is None: 
102			gradient = value.gradient(factDiff)
103			# Avoids recomputation, but generates NaNs at the seed points.
104			value = lp.dot_VV(gradient,factDiff)
105			value[np.isnan(value)]=0 
106		else:
107			value = value.norm(factDiff)
108
109	if ad.is_ad(value):
110		gradient = value.gradient()
111		value = value.value
112
113	if not ad.isndarray(value): 
114		raise ValueError(f"dictIn.SetFactor : unsupported value type {type(value)}")
115
116	#Set the gradient
117	if callable(gradient):
118		gradient = gradient(factGrid)
119
120	if not ad.isndarray(gradient): 
121		raise ValueError(f"dictIn.SetFactor : unsupported gradient type {type(gradient)}")
122
123	self["factoringValues"] = value
124	self["factoringGradients"] = np.moveaxis(gradient,0,-1) # Geometry last in c++ code...
125
126	for key in ('factoringRadius', 'factoringPointChoice'): self.pop(key,None)
127
128	return factGrid
129
130#def proj_tocost(x): return 2./(1.+lp.dot_VV(x,x))
131def proj_tosphere(x):
132	"""
133	Maps a point of the equatorial plane to the sphere, by projection : 
134	x -> (2 x,1-|x∏^2)/(1+|x|^2)
135	""" 
136	sq = lp.dot_VV(x,x)
137	return np.concatenate((x,np.expand_dims((1.-sq)/2.,axis=0)),axis=0) * (2./(1.+sq))
138def proj_fromsphere(q): 
139	"""
140	Maps a point of the sphere to the equatorial plane, by projection.
141	(q,qz) -> q/(1+qz)
142	"""
143	return q[:-1]/(1.+q[-1])
144
145#def sphere_tocost(x,chart): return proj_tocost(x)
146def sphere_tosphere(x,chart): 
147	"""
148	See proj_tosphere. Last component is reversed according to chart.
149	"""
150	y = proj_tosphere(x)
151	y[-1] *= chart
152	return y
153def sphere_fromsphere(q,chart):
154	"""
155	See proj_fromsphere. Last component of q is reversed according to chart.
156	"""
157	q=q.copy()
158	q[-1] *= chart
159	return proj_fromsphere(q)
160
161def SetSphere(self,dimsp,separation=5,radius=1.1):
162	"""
163	Setups the manifold R^(d-k) x S^k, using the equatorial projection for the sphere S^k.
164	Only compatible with the GPU accelerated eikonal solver. 
165	Inputs : 
166		dimsp (int): the discretization of a half sphere involves dimsp^k pixels.
167		radius (optional, float > 1): each local chart has base domain [-radius,radius]^k.
168		  (This is NOT the radius of the sphere, which is 1, but of the parametrization.)
169		separation (optional, int) : number of pixels separating the two local charts.
170			Set to false to define a projective space using a single chart. 
171	Side effects : 
172		Sets chart_mapping,chart_jump, dimensions, origin, gridScales.
173	Output : 
174		- Conversion utilities, between the equatorial plane and the sphere 
175			(and grid in the non-projective case)
176	"""
177	if 'chart_mapping' in self: # Some space R^(d-k) x S^k is already set
178		vdim_rect = self.vdim - self['chart_mapping'].ndim+1
179		dims_rect = self['dims'][:vdim_rect]
180		origin_rect = self['origin'][:vdim_rect]
181		gridScales_rect = self['gridScales'][:vdim_rect]
182	elif 'dims' in self: # Some rectangle R^(d-k) is already set
183		dims_rect = self['dims']
184		origin_rect = self.get('origin',np.zeros_like(dims_rect))
185		if 'gridScales' in self: gridScales_rect = self['gridScales']
186		else: gridScales_rect = self['gridScale']*np.ones_like(dims_rect) 
187	else: # Nothing set. All coordinates are sphere like, d=k
188		dims_rect = self.xp.array(tuple(),dtype=self.float_t)
189		origin_rect = dims_rect.copy()
190		gridScales_rect = dims_rect.copy()
191
192	vdim_rect = len(dims_rect)
193	vdim_sphere = self.vdim-vdim_rect # Dimension of the sphere
194	gridScale_sphere = 2*radius/dimsp
195
196	if separation: # Sphere manifold, described using two charts
197		separation_radius = separation*gridScale_sphere /2
198		center_radius = radius + separation_radius
199
200		dims_sphere = (2*dimsp + separation,) + (dimsp,)*(vdim_sphere-1)
201		origin_sphere = (-2*radius-separation_radius,) + (-radius,)*(vdim_sphere-1)
202
203		def sphere_fromgrid(x):
204			"""
205			Produces a point of the equatorial plane, and a boolean indicating which projection to 
206			use (False:south pole, True:north pole), from a grid point.
207			"""
208			assert len(x)==vdim_sphere
209			x = x.copy()
210			chart = np.where(np.abs(x[0])>=separation_radius,np.sign(x[0]),np.nan)
211			x[0] -= center_radius*chart
212			return x,chart
213
214		def sphere_togrid(x,chart):
215			"""
216			Produces a point of the original grid, from a point of the equatorial plane 
217			and a boolean indicating the active chart.
218			"""
219			assert len(x)==vdim_sphere
220			x=x.copy()
221			mixed_chart = x[0]*chart < -radius # The two coordinate charts would get mixed
222			x[:,mixed_chart] = np.nan
223			x[0] += chart*center_radius
224			return x
225	else: # Projective space, described using a single chart
226		dims_sphere = (dimsp,)*vdim_sphere
227		origin_sphere = (-radius,)*vdim_sphere
228
229	dims_sphere,origin_sphere,gridScales_sphere = [self.array_float_caster(e) 
230		for e in (dims_sphere,origin_sphere, (gridScale_sphere,)*vdim_sphere)]
231
232	self['dims'] = np.concatenate((dims_rect,dims_sphere),axis=0)
233	self['gridScales'] = np.concatenate(
234		(gridScales_rect,gridScales_sphere),axis=0)
235	self.pop('gridScale',None)
236	self['origin'] = np.concatenate((origin_rect,origin_sphere),axis=0)
237
238	# Produce a coordinate system, and set the jumps, etc
239	aX = self.Axes()[vdim_rect:]
240	X = self.array_float_caster(np.meshgrid(*aX,indexing='ij'))
241	if separation: X,chart = sphere_fromgrid(X)
242	X2 = lp.dot_VV(X,X)
243
244	# Geodesics jump when the are sufficiently far away from the fundamental domain.
245	radius_jump = (1+radius)/2
246	self['chart_jump'] = X2 > radius_jump**2
247	if separation:
248		self['chart_mapping'] = sphere_togrid(
249			sphere_fromsphere(sphere_tosphere(X,chart),-chart),-chart)
250		return {'from_grid':sphere_fromgrid,'to_grid':sphere_togrid,
251			'from_sphere':sphere_fromsphere,'to_sphere':sphere_tosphere} 
252	else:
253		self['chart_mapping'] = proj_fromsphere(-proj_tosphere(X))
254		return {'from_sphere':proj_fromsphere,'to_sphere':proj_tosphere}