agd.ExportedCode.Notebooks_NonDiv.MongeAmpere

151def newton_root(f,x0,fargs=tuple(),stop="Default",
152	relax=None,damping=None,ad="Sparse",solver=None,in_place=False):
153	"""
154	Newton's method, for finding a root of a given function.
155	f : function to be solved
156	x0 : initial guess for the root
157	stop : stopping criterion, presented as either
158	   - keyword "Default" for using the stop_default class
159	   - a dict for using the stop_class with specified initialization arguments
160	   - a callable,  
161	relax : added to the jacobian before inversion
162	damping : criterion for step reduction
163	ad : is either 
164	   - keyword "Sparse" for using Sparse AD (Default)
165	   - keyword "Dense" for using Dense AD
166	   - a shape_bound given as a tuple, for Dense AD with few independent variables
167	"""
168
169	if stop == "Default": 	stop = stop_default()
170	elif isinstance(stop,dict): stop = stop_default(**stop)
171	if ad == "Dense":		ad = tuple() # Tuple represents shape_bound
172
173	# Create a variable featuring AD information
174	def ad_var(x0):
175		if not in_place: x0=np.copy(x0)
176
177		if ad == "Sparse":		   return Sparse.identity(constant=x0)
178		elif isinstance(ad,tuple): return Dense.identity( constant=x0,shape_bound=ad)						
179		assert False
180
181	# Find Newton's descent direction
182	def descent_direction(residue):
183		if relax is not None: residue+=relax # Relax is usually a multiple of identity
184
185		if solver is not None:     return solver(residue)
186		elif ad == "Sparse":       return residue.solve() # Sparse ad
187		elif isinstance(ad,tuple): return residue.solve(shape_bound=ad) # Dense ad
188
189	x=ad_var(x0)
190	for niter in itertools.count():
191		residue = f(x,*fargs)
192		if stop(residue,niter): return remove_ad(x)
193
194		direction = descent_direction(residue)
195
196		step = 1. if damping is None else damping(remove_ad(x),direction,*fargs)
197		x += step*direction

18def norm(arr,ord=2,axis=None,keepdims=False,averaged=False):
19	"""
20	Returns L^p norm of array, seen as a vector, w.r.t. weights.
21	Defined as : (sum_i x[i]^p)^(1/p)
22
23	Remark : not a matrix operator norm
24	
25	Inputs:
26	 - ord : exponent p
27	 - axis : int or None, axis along which to compute the norm. 
28	 - keepdims : wether to keep singleton dimensions.
29	 - averaged : wether to introduce a normalization factor, so that norm(ones(...))=1
30
31	Compatible with automatic differentiation.
32	"""
33	arr = ad_generic.array(arr)
34	if ord==np.inf or ord%2!=0:
35		try: arr = np.abs(arr)
36		except TypeError: arr = np.vectorize(np.abs)(arr)
37	if ord==np.inf: return np.max(arr,axis=axis,keepdims=keepdims)
38	sum_pow = np.sum(arr**ord,axis=axis,keepdims=keepdims)
39	
40	if averaged:
41		size = arr.size if axis is None else arr.shape[axis]
42		sum_pow/=size
43
44	return sum_pow**(1./ord)

18def SchemeNonMonotone(u,f,bc):
19    # Compute the hessian matrix of u
20    uxx = bc.Diff2(u,(1,0))
21    uyy = bc.Diff2(u,(0,1))
22    uxy = 0.25*(bc.Diff2(u,(1,1)) - bc.Diff2(u,(1,-1)))
23    
24    # Numerical scheme
25    det = uxx*uyy-uxy**2
26    residue = f - det
27    
28    # Boundary conditions
29    return np.where(bc.interior,residue,u-bc.grid_values)

31def MALBR_H(d2u):
32    a,b,c = np.sort(np.maximum(0.,d2u), axis=0)
33
34    # General formula, handling infinite values separately
35    A,B,C = (np.where(e==np.inf,0.,e) for e in (a,b,c))
36    result = 0.5*(A*B+B*C+C*A)-0.25*(A**2+B**2+C**2)
37    
38    pos_inf = np.logical_or.reduce(d2u==np.inf)    
39    result[pos_inf]=np.inf
40    
41    pos_ineq = a+b<c
42    result[pos_ineq] = (A*B)[pos_ineq]
43        
44    return result

46def SchemeMALBR(u,SB,f,bc):
47    # Compute the finite differences along the superbase directions
48    d2u = bc.Diff2(u,SB)
49    d2u[...,bc.not_interior] = 0. # Replace NaNs with arbitrary values to silence warnings
50    
51    # Numerical scheme
52    residue = f-MALBR_H(d2u).min(axis=0)
53    
54    # Boundary conditions
55    return np.where(bc.interior,residue,u-bc.grid_values)

57def InvalidMALBR(u,SB,f,bc):
58    residue = SchemeMALBR(u,SB,f,bc)
59    return np.any(residue[bc.interior]>=f/2)

61def SchemeMALBR_OptInner(u,SB,bc,oracle=None):
62    # If the active superbases are known, then take only these
63    if not(oracle is None):
64        SB = np.take_along_axis(SB,np.broadcast_to(oracle,SB.shape[:2]+(1,)+oracle.shape),axis=2)
65                
66    d2u = bc.Diff2(u,SB)
67    d2u[...,bc.not_interior] = 0. # Placeholder value to silent NaN warnings
68    # Evaluate the complex non-linear function using dense - sparse composition
69    result = ad.apply(MALBR_H,d2u,shape_bound=u.shape)
70    
71    return ad.min_argmin(result,axis=0)

73def SchemeMALBR_Opt(u,SB,f,bc):
74    
75    # Evaluate using the envelope theorem
76    result,_ = ad.apply(SchemeMALBR_OptInner, u,bc.as_field(SB),bc, envelope=True)
77        
78    # Boundary conditions
79    return np.where(bc.interior, f - result, u-bc.grid_values)

 81def ConstrainedMaximize(Q,l,m):
 82    dim = l.shape[0]
 83    if dim==1:
 84        return (l[0]+np.sqrt(Q[0,0]))/m[0]
 85    
 86    # Discard infinite values, handled afterwards
 87    pos_bad = l.min(axis=0)==-np.inf
 88    L = l.copy(); L[:,pos_bad]=0
 89    
 90    # Solve the quadratic equation
 91    A = lp.inverse(Q)    
 92    lAl = lp.dot_VAV(L,A,L)
 93    lAm = lp.dot_VAV(L,A,m)
 94    mAm = lp.dot_VAV(m,A,m)
 95    
 96    delta = lAm**2 - (lAl-1.)*mAm
 97    pos_bad = np.logical_or(pos_bad,delta<=0)
 98    delta[pos_bad] = 1.
 99    
100    mu = (lAm + np.sqrt(delta))/mAm
101    
102    # Check the positivity
103#    v = dot_AV(A,mu*m-L)
104    rm_ad = np.array
105    v = lp.dot_AV(rm_ad(A),rm_ad(mu)*rm_ad(m)-rm_ad(L))
106    pos_bad = np.logical_or(pos_bad,np.any(v<0,axis=0))
107    
108    result = mu
109    result[pos_bad] = -np.inf
110    
111    # Solve the lower dimensional sub-problems
112    # We could restrict to the bad positions, and avoid repeating computations
113    for i in range(dim):             
114        axes = np.full((dim),True); axes[i]=False
115        res = ConstrainedMaximize(Q[axes][:,axes],l[axes],m[axes])
116        result = np.maximum(result,res)
117    return result

119def SchemeUniform(u,SB,f,bc):
120    # Compute the finite differences along the superbase directions
121    d2u = bc.Diff2(u,SB) 
122    d2u[...,bc.not_interior] = 0. # Placeholder value to silent NaN warnings
123    
124    # Generate the parameters for the low dimensional optimization problem
125    Q = 0.5*np.array([[0,1,1],[1,0,1],[1,1,0]])
126    l = -d2u
127    m = lp.dot_VV(SB,SB)
128    
129    # Evaluate the numerical scheme
130    m = bc.as_field(m)
131    from agd.FiniteDifferences import as_field
132    Q = as_field(Q,m.shape[1:])
133    
134    dim = 2
135    alpha = dim * f**(1/dim)
136    mask= (alpha==0)
137
138    Q = Q* np.where(mask,1.,alpha**2)
139    residue = ConstrainedMaximize(Q,l,m).max(axis=0)
140    residue[mask] = np.max(l/m,axis=0).max(axis=0)[mask]
141    
142    # Boundary conditions
143    return np.where(bc.interior,residue,u-bc.grid_values)

145def SchemeUniform_OptInner(u,SB,f,bc,oracle=None):
146    # Use the oracle, if available, to select the active superbases only
147    if not(oracle is None):
148        SB = np.take_along_axis(SB,np.broadcast_to(oracle,SB.shape[:2]+(1,)+oracle.shape),axis=2)
149
150    d2u = bc.Diff2(u,SB) 
151    d2u[...,bc.not_interior] = 0. # Placeholder value to silent NaN warnings
152    
153    # Generate the parameters for the low dimensional optimization problem
154    Q = 0.5*np.array([[0,1,1],[1,0,1],[1,1,0]])
155    dim = 2
156    l = -d2u
157    m = lp.dot_VV(SB,SB)
158    
159    m = bc.as_field(m)
160    from agd.FiniteDifferences import as_field
161    Q = as_field(Q,m.shape[1:])
162    
163    dim = 2
164    alpha = dim * f**(1/dim)
165    mask= (alpha==0)
166
167    Q = Q* np.where(mask,1.,alpha**2)
168    # Evaluate the non-linear functional using dense-sparse composition
169    residue = ad.apply(ConstrainedMaximize,Q,l,m,shape_bound=u.shape).copy()
170    residue[:,mask] = np.max(l/m,axis=0)[:,mask]
171    
172    return ad.max_argmax(residue,axis=0)

174def SchemeUniform_Opt(u,SB,f,bc):
175    
176    # Evaluate the maximum over the superbases using the envelope theorem
177    residue,_ = ad.apply(SchemeUniform_OptInner, u,bc.as_field(SB),f,bc, envelope=True)
178    
179    return np.where(bc.interior,residue,u-bc.grid_values)

181def Hessian_ad(u,x):
182    x_ad = ad.Dense2.identity(constant=x,shape_free=(2,))
183    return u(x_ad).hessian()

184def MongeAmpere_ad(u,x):
185    return lp.det(Hessian_ad(u,x))