agd.ODE.hamiltonian

  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
  4import scipy.sparse
  5
  6from .. import AutomaticDifferentiation as ad
  7from .. import LinearParallel as lp
  8from .hamiltonian_base import HamiltonianBase,fixedpoint,damp_None,read_None,incr_None
  9import numpy as np
 10from copy import copy
 11from .backtrack import RecurseRewind
 12
 13
 14class MetricHamiltonian(HamiltonianBase):
 15	r"""
 16	Hamiltonian defined by an interpolated metric, which is dualized and interpolated.
 17	$$
 18		H(q,p) = \frac 1 2 F^*_q(p)^2
 19	$$
 20	__init__ arguments :
 21	- metric : dual defines the hamiltonian
 22	- **kwargs : passed to metric.dual().set_interpolation 
 23	"""
 24
 25	def __init__(self,metric,**kwargs):
 26		super(MetricHamiltonian,self).__init__()
 27		self._dualmetric = metric.dual()
 28		self._dualmetric.set_interpolation(**kwargs)
 29
 30	def H(self,q,p): return self._dualmetric.at(q).norm2(p)
 31	
 32	def DqH(self,q,p): 
 33		q_ad = ad.Dense.identity(constant=q,shape_free=(len(q),))
 34		return np.reshape(self._dualmetric.at(q_ad).norm2(p).gradient(),np.shape(p))
 35
 36	def DpH(self,q,p): return self._dualmetric.at(q).gradient2(p)
 37
 38class GenericHamiltonian(HamiltonianBase):
 39	r"""
 40	Hamiltonian defined by a arbitrary function $f$ of two variables, 
 41	the position $q$ and impulsion $p$, denoted
 42	$$
 43		H(q,p)
 44	$$
 45
 46	__init__ arguments : 
 47	- H : the hamiltonian, must take two arguments.
 48	- shape_free (optional) : shape of position and momentum variables, used for autodiff
 49	- disassociate_ad (optional) : hide AD information when calling $f$. (Use to avoid 
 50	conflicts if the definition of $f$ itself involves automatic differentiation.)
 51	- **kwargs : passed to HamiltonianBase
 52	"""
 53
 54	def __init__(self,H, shape_free=None, disassociate_ad=False):
 55		super(GenericHamiltonian,self).__init__()
 56		self.H = H
 57		self.shape_free = shape_free
 58		self.disassociate_ad = disassociate_ad
 59
 60	def _identity_ad(self,x,noad=None):
 61		x_ad = ad.Dense.identity(constant=x,shape_free=self.shape_free) 
 62		if self.disassociate_ad: 
 63			x_dis = ad.disassociate(x_ad,shape_free=self.shape_free)
 64			if noad is None: return x_dis
 65			else:return (x_dis,ad.disassociate(type(x_ad)(noad),shape_free=self.shape_free)) 
 66		else: 
 67			return x_ad if noad is None else (x_ad,noad)
 68
 69	def _gradient_ad(self,x):
 70		if self.disassociate_ad: x=ad.associate(x)
 71		return x.gradient()
 72
 73	def DqH(self,q,p):
 74		q_ad,p = self._identity_ad(q,noad=p)
 75		# If the reshaping fails, then consider setting shape_free.
 76		return np.reshape(self._gradient_ad(self.H(q_ad,p)), np.shape(q))
 77
 78	def DpH(self,q,p):
 79		p_ad,q = self._identity_ad(p,noad=q)
 80		# If the reshaping fails, then consider setting shape_free.
 81		return np.reshape(self._gradient_ad(self.H(q,p_ad)), np.shape(p))
 82
 83# -------------------------------- Separable -------------------------------
 84
 85class SeparableHamiltonianBase(HamiltonianBase):
 86	"""
 87	Base class for separable Hamiltonians, with generic form : H(q,p) = H_Q(q) + H_P(p).
 88	"""
 89	def __init__(self):
 90		super(SeparableHamiltonianBase,self).__init__()
 91
 92	# Redirect to single argument version of the partial derivatives
 93	def DqH(self,q,_): return self._DqH(q)
 94	def DpH(self,_,p): return self._DpH(p)
 95	def _DqH(q): 
 96		"""Derivative of the Hamiltonian w.r.t position"""
 97		raise NotImplementedError
 98	def _DpH(p): 
 99		"""Derivative of the Hamiltonian w.r.t impulsion"""
100		raise NotImplementedError
101
102	def Impl_p(self,q,p,δ): 
103		"""Time step for the impulsion p (implicit=explicit for a separable scheme"""
104		return self.Expl_p(q,p,δ)
105
106	def Impl_q(self,q,p,δ):
107		"""Time step for the position q (implicit=explicit for a separable scheme)"""
108		return self.Expl_q(q,p,δ)
109
110	def Impl2_p(self,q,p,δ_before,δ_total,δ_after):
111		"""
112		Merge two implicit time steps for the impulsion p, with a damping step in between.
113		"""
114		# If there is no damping, then the two explicit time steps can be merged
115		if (self.damp_p is damp_None and self.damp_q is damp_None and self.Impl2_p_merged 
116		and self.read_p is read_None and self.incr_q is incr_None): # Could allow read_p easily
117			self.read_q(self,q) # Read position before (non-existent) damp
118			p = self.Expl_p(q,p,δ_before+δ_after)
119			self.incr_p(self,p) # Modify momentum after (non-existent) damp
120			return q,p
121
122		return super().Impl2_p(q,p,δ_before,δ_total,δ_after)
123
124
125class SeparableHamiltonian(SeparableHamiltonianBase):
126	"""
127	Separable Hamiltonian defined by a pair of functions, differentiated with AD.
128	$$
129		H(q,p) = H_Q(q) + H_P(p).
130	$$
131	__init__ arguments : 
132	- Hq,Hp : the two functions $H_Q,H_P$, of a single argument, defining the hamiltonian
133	- shape_free (optional) : shape of position and momentum variables, used for autodiff
134	"""
135
136	def __init__(self,Hq,Hp,shape_free=None):
137		super(SeparableHamiltonian,self).__init__()
138		self.Hq = Hq
139		self.Hp = Hp
140		self.shape_free = shape_free
141
142	def H(self,q,p): return self.Hq(q) + self.Hp(p)
143
144	def _DqH(self,q): 
145		q_ad = ad.Dense.identity(constant=q,shape_free=self.shape_free)
146		return np.reshape(self.Hq(q_ad).gradient(),np.shape(q))
147
148	def _DpH(self,p): 
149		p_ad = ad.Dense.identity(constant=p,shape_free=self.shape_free)
150		return np.reshape(self.Hp(p_ad).gradient(),np.shape(p))
151
152# -------------------------------- Quadratic -------------------------------
153
154class QuadraticHamiltonianBase(SeparableHamiltonianBase):
155	"""
156	Base class for separable quadratic Hamiltonians. 
157	Implements the perturbed Hamiltonians which are preserved by the symplectic schemes.
158	"""
159	def __init__(self,shape_free=None):
160		super(QuadraticHamiltonianBase,self).__init__()
161		self.shape_free = shape_free
162
163	def flat(self,x):
164		"""Flattens the vector x, for e.g. product with a sparse matrix"""
165		if self.shape_free is None: return x.reshape(-1)
166#		assert x.shape[:len(self.shape_free)] == self.shape_free
167		return x.reshape((np.prod(self.shape_free),*x.shape[len(self.shape_free):])) 
168
169	def _dot(self,q,p):
170		"""Duality bracket between position and impulsion"""
171		return np.sum(q*p, 
172			axis=None if self.shape_free is None else tuple(range(len(self.shape_free))))
173
174	def _ABC(self,δ):
175		A = self._DpH
176		B = self._DqH
177		C = lambda q: δ**2 * A(B(q))
178		return A,B,C,self._dot
179
180	def HEuler_p(self,q,p,δ):
181		"""Modified Hamiltonian, preserved by the symplectic Euler_p scheme"""
182		A,B,_,dot = self._ABC(δ)
183		Ap = A(p)
184		return 0.5*dot(Ap,p) + 0.5*dot(B(q),q-δ*Ap)
185
186	def HVerlet_p(self,q,p,δ):
187		"""Modified Hamiltonian, preserved by the Verlet_p symplectic scheme"""
188		A,B,C,dot = self._ABC(δ)
189		return 0.5*dot(A(p),p) + 0.5*dot(B(q),q-0.25*C(q))
190
191	def HRuth4_p(self,q,p,δ):
192		"""Modified Hamiltonian, preserved by the Ruth4_p symplectic scheme"""
193		A,B,C,dot = self._ABC(δ)
194		Ap = A(p); a1,a2,a3,a4 = - 0.33349609375,-0.16400146484375,0.0319671630859375,0.009191930294036865
195		b1,b2,b3,b4,b5 = - 0.33349609375, -0.087890625, 0.06305503845214844,0.0053994059562683105,0.0041955700144171715
196		return 0.5*dot(p,Ap+C(a1*Ap+C(a2*Ap+C(a3*Ap)))) + 0.5*dot(B(q),q+C(b1*q+C(b2*q+C(b3*q+C(b4*q+C(b5*q))))))
197
198	def H_p(self,q,p,δ,order):
199		"""
200		Modified Hamiltonian, preserved by the Euler_p, Verlet_p, or Ruth4_p 
201		symplectic scheme, depending on the order parameter. (See method Sympl_p.)
202		"""
203		if order==1: return self.HEuler_p(q,p,δ)
204		if order==2: return self.HVerlet_p(q,p,δ)
205		if order==4: return self.HRuth4_p(q,p,δ)
206		raise ValueError(f"Found {order=}, while expecting 1,2 or 4.")
207
208
209	def Impl2_p(self,q,p,δ_before,δ_total,δ_after):
210		"""Merge two implicit time steps for the impulsion p, with a damping step in between."""
211
212		# We are computing α p - δ_after B β q - δ_before α B q, 
213		# Where α = exp(-δ_total damp_p) and β = exp(-δ_total damp_q) 
214
215		if (self.damp_q is damp_None and self.damp_p is not damp_None and self.Impl2_p_merged # β=1
216		and self.incr_q is incr_None and self.read_p is read_None): 
217			# Factorization : α p - (δ_after + δ_before α) B q
218			self.read_q(self,q) # Read position before damp
219			dp = self._DqH(q)
220			α = np.exp(-δ_total*self.damp_p)
221			p = α*p - (δ_after + δ_before*α) * dp
222			self.incr_p(self,p)
223			return q,p
224
225		if (self.damp_p is damp_None and self.damp_q is not damp_None and self.Impl2_p_merged # α=1
226		and self.read_p is read_None and self.incr_p is incr_None): 
227			# Factorization : p - B (δ_after β + δ_before) q
228			self.read_q(self,q)
229			β = np.exp(-δ_total*self.damp_q)
230			qnew = q*β
231			self.incr_q(self,qnew)
232			p = self.Expl_p(δ_after*qnew+δ_before*q,p,1) # Using Expl for reverse AD
233			self.incr_p(self,p) 
234			return qnew,p
235
236		# read and incr : We need to see what is the typical usage, for now likely too restrictive.
237
238		return super().Impl2_p(q,p,δ_before,δ_total,δ_after)
239
240	def seismogram_with_backprop(self,q,p,δ,niter,order=2,qh_ind=None,ph_ind=None):
241		"""
242		Computes niter time steps of a symplectic scheme, collects the values at given indices along
243		the way (the seismogram), and allows to backpropagate the results.
244		
245		Inputs : 
246		- qh_ind,ph_ind : indices at which to collect the values of q and p, 
247			in the flattened arrays. IMPORTANT : no duplicate values in either qh_ind or ph_ind.
248
249		Outputs : 
250		- (qf,pf) : final values of q and p
251		- (qh,ph) : history of q and p. The prescribed indices are extracted along the way, and 
252			concatenated into a "simogram". Last iteration is not included, use qf,pf.
253		- backprop : callable, which given (qf_coef,pf_coef) and (qh_coef,ph_coef), the gradient of
254			some objective(s) functional w.r.t the outputs, backpropagates the results to obtain the
255			gradients w.r.t the Hamiltonian parameters. 
256			qf_coef.shape == (*qf.shape,size_ad), and likewise pf,qh,ph.
257		"""
258
259		H_fwd = copy(self); H_rev = copy(self)
260		H_fwd.read_q = read_None; H_fwd.read_p = read_None; H_fwd.incr_q = incr_None; H_fwd.incr_p = incr_None
261		H_rev.read_q = read_None; H_rev.read_p = read_None; H_rev.incr_q = incr_None; H_rev.incr_p = incr_None
262
263		qh = []; ph = []; initial=True 
264#		if qh_ind is not None: H_fwd.read_q = lambda _,q : qh.append(q.reshape(-1)[qh_ind])
265#		if ph_ind is not None: H_fwd.read_p = lambda _,p : ph.append(p.reshape(-1)[ph_ind])
266
267		# We construct a callable, which will extract the desired seismogram, and also collect
268		# keypoint values of q and p.
269		# The main objective of this routine is to preserve accuracy when damp_p and damp_q are set.
270		# Otherwise, the iterations may be reversed by iterating the symplectic scheme 
271		# in negative time, providing a result with similar accuracy. 
272
273		# TODO : Define next(qp,niter) and modify RecurseRewind accordingly, so as to take advantage
274		# of Impl2_p merged steps (at best, expecting a factor 2 computational cost reduction)
275		def next(qp): # A single time step, including a preliminary damping
276			q,p = H_fwd.Damp_qp(*qp,δ)
277			q,p = H_fwd.Sympl_p(q,p,δ,order=order)
278			if initial and qh_ind is not None: qh.append(q.reshape(-1)[qh_ind].copy())
279			if initial and ph_ind is not None: ph.append(p.reshape(-1)[ph_ind].copy())
280			return q,p
281		# A single negative damping step should be fine...
282		qph_eval = RecurseRewind(next,self.Damp_qp(q,p,-δ))
283
284		# Perform forward propagation
285		qf,pf = qph_eval(niter)
286		# We drop the last element, for consistency with backprop. Recover it as qf[q_ind] 
287		qh = ad.array(qh[:-1]); ph = ad.array(ph[:-1]) 
288
289		# Remove seismogram extraction
290		initial = False; # H_fwd.read_q = read_None; H_fwd.read_p = read_None
291
292		def backprop(qf_grad=None,pf_grad=None,qh_grad=None,ph_grad=None,
293			check_val=False,check_ind=True):
294			"""
295			- qf_grad : gradient of objective functional(s) w.r.t qf, with shape (*qf.shape,size_ad)
296			- pf_grad, qh_grad, ph_grad : gradients w.r.t ph, qh, ph, similar to qf_grad
297			- check_val : check that the back propagation reconstructs values correctly
298			- check_ind : check that the seismogram indices do not contain duplicates
299			"""
300			# Data checks
301			if check_ind:
302				assert np.unique(qh_ind).size==qh_ind.size and np.unique(ph_ind).size==ph_ind.size
303			if ad.is_ad(qf) or ad.is_ad(pf) or ad.is_ad(qh) or ad.is_ad(ph): 
304				raise ValueError("Please choose between forward and reverse autodiff")
305			size_ad = max(x.size//y.size for x,y in 
306				((qf_grad,qf),(pf_grad,pf),(qh_grad,qh),(ph_grad,ph)) if x is not None and x.size>0)
307			for x_name,x_grad,x in (("qf",qf_grad,qf),("pf",pf_grad,pf),("qh",qh_grad,qh),("ph",ph_grad,ph)): 
308				if x_grad is not None and x_grad.shape!=(*x.shape,size_ad):
309					raise ValueError(f"Expecting shape {(*x.shape,size_ad)} for field {x_name}_grad, but found {x_grad.shape}")
310
311			# Insert the gradients in the backpropagation
312			if qf_grad is None: qf_grad = np.zeros_like(qf,shape=(*qf.shape,size_ad))
313			if pf_grad is None: pf_grad = np.zeros_like(pf,shape=(*pf.shape,size_ad))
314			qf_rev = ad.Dense.denseAD(qf, pf_grad) # Signs and exchange
315			pf_rev = ad.Dense.denseAD(pf,-qf_grad) # account for symplectic formalism
316
317			def incr_q(H,q):
318				rev_iter = niter-1-H.current_iter
319				q.coef.reshape((-1,size_ad))[ph_ind] += ph_grad[rev_iter-1] 
320				assert not check_val or np.allclose(q.value * np.exp(-δ*(H.damp_p+H.damp_q)), qph_eval(rev_iter)[0])
321				q.value = qph_eval(rev_iter)[0]
322			def incr_p(H,p):
323				rev_iter = niter-1-H.current_iter
324				p.coef.reshape((-1,size_ad))[qh_ind] -= qh_grad[rev_iter-1] 
325				assert not check_val or np.allclose(p.value * np.exp(-δ*(H.damp_p+H.damp_q)), qph_eval(rev_iter)[1])
326				p.value = qph_eval(rev_iter)[1]
327
328			if qh_ind is not None: H_rev.incr_p = incr_p
329			if ph_ind is not None: H_rev.incr_q = incr_q
330
331			# We do not reset the reverse AD accumulators, in case the user wants to accumulate 
332			# more data (e.g. seismograms with varying initial impulsions and frequencies) 
333			# H_rev.rev_reset() 
334
335			# Setup for time-reversed propagation
336			H_rev.damp_q,H_rev.damp_p = - H_fwd.damp_p,- H_fwd.damp_q
337			q0_rev,p0_rev = H_rev.Sympl_p(qf_rev,pf_rev,-δ,niter,order)
338
339			return -p0_rev.coef,q0_rev.coef
340
341		return qf,pf,ad.asarray(qh),ad.asarray(ph),backprop
342
343class QuadraticHamiltonian(QuadraticHamiltonianBase):
344	r"""
345	Quadratic Hamiltonian, defined by a pair of linear operators.
346	(Expected to be symmetric semi-definite.)
347	$$
348		H(q,p) = \frac 1 2 (< q, M_Q q > + < p, M_P p >).
349	$$
350
351	__init__ arguments : 
352	- Mq,Mp : positive semi-definite matrices $M_Q,M_P$, typically given in sparse form.
353	 Alternatively, define Mq,Mp as functions, and use the set_spmat
354	 to automatically generate the sparse matrices using automatic differentiation.
355	"""
356
357	def __init__(self,Mq,Mp,**kwargs):
358		super(QuadraticHamiltonian,self).__init__(**kwargs)
359		self.Mq = Mq
360		self.Mp = Mp
361		 # dMp, dMq are for reverse autodiff, see set_spmat
362		self.dMp_has = False
363		self.dMq_has = False
364	
365	def H(self,q,p):
366		A,B,C,dot = self._ABC(1) # δ is irrelevant here
367		return 0.5*dot(A(p),p) + 0.5*dot(B(q),q)
368
369	def _D_H(self,x,M,dM__has):
370		if ad.is_ad(x) and dM__has: raise ValueError("Cannot accumulate reverse AD without dt")
371		return np.reshape(ad.apply_linear_mapping(M,self.flat(x)),np.shape(x))
372
373	def _DqH(self,q): return self._D_H(q,self.Mq,self.dMq_has)
374	def _DpH(self,p): return self._D_H(p,self.Mp,self.dMp_has)
375
376	def rev_reset(self):
377		if self.dMq_has: self.dMq_acc[:]=0
378		if self.dMp_has: self.dMp_acc[:]=0
379
380	def Expl_q(self,q,p,δ):
381		"""Explicit time step for the position q."""
382		if self.preserve_q: q = self._mk_copy(q)
383		q += δ*self._D_H(p,self.Mp,False)
384		if ad.is_ad(p) and self.dMp_has: self.dMp_acc += δ*self.dMp_op(p)
385		return q 
386
387	def Expl_p(self,q,p,δ):
388		"""Explicit time step for the impulsion p."""
389		if self.preserve_p: p = self._mk_copy(p)
390		p -= δ*self._D_H(q,self.Mq,False)
391		if ad.is_ad(q) and self.dMq_has: self.dMq_acc -= δ*self.dMq_op(q)
392		return p 
393
394	def set_spmat(self,x,rev_ad=(None,None),**kwargs):
395		"""
396		Replaces Mq,Mp with suitable sparse matrices, generated by spmat,
397		if they are callables.
398
399		- x : Correctly shaped input for calling Mq,Mp.
400		- rev_ad (optional) : where to accumulate reverse autodiff.
401			See Eikonal.HFM_CUDA.AnisotropicWave.AcousticHamiltonian_Sparse for an example
402		"""
403		if self.shape_free is None: self.shape_free = x.shape
404		else: assert self.shape_free == x.shape
405		spmat = ad.Sparse2.hessian_operator
406		self.dMq_has, self.dMp_has = [y is not None for y in rev_ad]
407		if callable(self.Mq): self.Mq = spmat(self.Mq,x,**kwargs,rev_ad=self.dMq_has)
408		if callable(self.Mp): self.Mp = spmat(self.Mp,x,**kwargs,rev_ad=self.dMp_has)
409
410		if self.dMq_has: self.Mq,self.dMq_op,self.dMq_acc = *self.Mq,rev_ad[0]
411		if self.dMp_has: self.Mp,self.dMp_op,self.dMp_acc = *self.Mp,rev_ad[1]