agd.ODE.hamiltonian_base

  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
  4r"""
  5This module implements some basic functionality for solving ODEs derived from a 
  6Hamiltonian in a manner compatible with automatic differentiation.
  7(Flow computation, symplectic schemes, etc)
  8
  9Recall that Hamilton's equations read 
 10$$
 11\frac {dq}{dt} = \frac {\partial H}{\partial p},
 12\quad
 13\frac {dp}{dt} = - \frac {\partial H}{\partial q}.
 14$$
 15"""
 16
 17import numpy as np
 18from copy import copy
 19
 20from .. import AutomaticDifferentiation as ad
 21
 22def fixedpoint(f,x,tol=1e-9,nitermax=100):
 23	"""
 24	Iterates the function f on the data x until a fixed point is found, 
 25	up to prescribed tolerance, or the maximum number of iterations is reached.
 26	"""
 27	norm_infinity = ad.Optimization.norm_infinity
 28	x_old = x
 29	for i in range(nitermax):
 30		x=f(x)
 31		if norm_infinity(x-x_old)<tol: break
 32		x_old = x
 33	return x
 34
 35damp_None = 0 # Placeholder for the damping factors damp_q and damp_p of class HamiltonianBase
 36def read_None(H,x): pass # Placeholder for the data read functions read_q and read_p
 37def incr_None(H,x): pass # Placeholder for the data write functions incr_q and incr_p
 38
 39class HamiltonianBase:
 40	"""
 41	Base class for Hamiltonians.
 42
 43	Default initialized fields : 
 44	- impl_solver : Fixed point solver to be used for the implicit time steps
 45	- preserve_q, preserve_p : Wether updates modify q and p in place
 46	- damp_q, damp_p : damping factor between scheme iterations
 47	- read_q, read_p : called before each damping step, for reading q and p
 48	- incr_q, incr_p : called after each damping step, as a source term for q and p
 49	"""
 50	def __init__(self):
 51		self.impl_solver = fixedpoint 
 52		self.preserve_q = True; self.preserve_p = True 
 53		self.Impl2_p_merged = True 
 54		self.damp_q = damp_None; self.damp_p = damp_None 
 55		self.read_q = read_None; self.read_p = read_None
 56		self.incr_q = incr_None; self.incr_p = incr_None
 57		self.current_iter = None 
 58
 59	def H(self,q,p):
 60		"""Evaluates the Hamiltonian, at a given position and impulsion."""
 61		raise NotImplementedError
 62
 63	def DqH(q,p):
 64		"""Differentiates the Hamiltonian, w.r.t. position."""
 65		raise NotImplementedError
 66
 67	def DpH(q,p):
 68		"""Differentiates the Hamiltonian, w.r.t. impulsion."""
 69		raise NotImplementedError
 70
 71	def flow(self,q,p):
 72		"""
 73		Symplectic gradient of the Hamiltonian.
 74		"""
 75		return (self.DpH(q,p),-self.DqH(q,p))
 76
 77	def flow_cat(self,qp,t=None):
 78		"""
 79		Symplectic gradient of the hamiltonian, intended for odeint. 
 80
 81		Input : 
 82		 - qp : position q, impulsion p, flattened and concatenated.
 83		 - t : ignored parameter (compatibility with scipy.integrate.odeint)
 84		 - shape : how to reshape q and p for calling DpH and DqH, if needed
 85
 86		Output :
 87		 - symplectic gradient, concatenated and flattened.
 88		"""
 89		d = len(qp)//2
 90		q = qp[:d]; p = qp[d:]
 91		if hasattr(self,'shape_free') and self.shape_free is not None:
 92			q = q.reshape((*self.shape_free,-1))
 93			p = p.reshape((*self.shape_free,-1))
 94		return np.concatenate(self.flow(q,p),axis=0).reshape(-1)
 95
 96	def integrate(self,q,p,scheme,niter=None,dt=None,T=None,path=False):
 97		"""
 98		Solves Hamilton's equations by running the scheme niter times.
 99
100		Inputs : 
101		 - q,p : Initial position and impulsion.
102		 - scheme : ODE integration scheme. (string or callable)
103		 - niter,dt,T : number of steps, time step, and total time
104			  (exactly two among the three must be specified)
105		 - path : wether to return the intermediate steps. 
106			   (If a positive number, period of intermediate steps to return)
107		
108		Output : 
109		 - q,p if path is False. 
110		  Otherwise np.stack([q0,...,qn],axis=-1),np.stack([p0,...,pn],axis=-1),[t0,..tn], 
111		  	with n=niter, tn=T.
112		"""
113		if isinstance(scheme,str):
114			schemes = self.nonsymplectic_schemes()
115			schemes.update(self.symplectic_schemes())
116			scheme = schemes[scheme]
117
118		if (niter is None) + (T is None) + (dt is None) != 1: 
119			raise ValueError("Exactly two of niter, dt and T must be specified")
120		if T is None:    T  = niter*dt
121		elif dt is None: dt = T/niter
122		elif niter is None: 
123			niter = int(np.ceil(T/dt))
124			dt = T/niter # slightly decreased time step
125
126		q,p = copy(q),copy(p)
127		if path: Q,P = [copy(q)],[copy(p)]
128
129		for i in range(niter):
130			q,p = scheme(q,p,dt)
131			if path and not i%path: Q.append(copy(q)); P.append(copy(p))
132
133		if path: return np.stack(Q,axis=-1),np.stack(P,axis=-1),np.linspace(0,T,niter+1)
134		else: return q,p
135
136	def nonsymplectic_schemes(self):
137		"""
138		Standard ODE integration schemes
139		"""
140		def Euler(q,p,dt):
141			dq,dp = self.flow(q,p)
142			return q+dt*dq, p+dt*dp
143
144		def RK2(q,p,dt):
145			dq1,dp1 = self.flow(q, p)
146			dq2,dp2 = self.flow(q+0.5*dt*dq1, p+0.5*dt*dp1)
147			return q+dt*dq2, p+dt*dp2
148
149		def RK4(q,p,dt):
150			dq1,dp1 = self.flow(q, p)
151			dq2,dp2 = self.flow(q+0.5*dt*dq1, p+0.5*dt*dp1)
152			dq3,dp3 = self.flow(q+0.5*dt*dq2, p+0.5*dt*dp2)
153			dq4,dp4 = self.flow(q+dt*dq3, p+dt*dp3)
154			return q+dt*(dq1+2*dq2+2*dq3+dq4)/6., p+dt*(dp1+2*dp2+2*dp3+dp4)/6.
155
156		return {"Euler":Euler,"Runge-Kutta-2":RK2,"Runge-Kutta-4":RK4}
157
158	# -------------- sub-steps of schemes --------------
159
160	def _mk_copy(self,x): 
161		"""Returns a copy of variable x. Specialize if needed (for GPU c_contiguity...)."""
162		return x.copy()
163
164	def Expl_q(self,q,p,δ):
165		"""Explicit time step for the position q."""
166		if self.preserve_q: q = self._mk_copy(q)
167		q += δ*self.DpH(q,p)
168		return q 
169	def Expl_p(self,q,p,δ):
170		"""Explicit time step for the impulsion p."""
171		if self.preserve_p: p = self._mk_copy(p)
172		p -= δ*self.DqH(q,p)
173		return p 
174
175	def Damp_qp(self,q,p,δ):
176		"""
177		Optional damping step, interleaved in between Verlet and Ruth4 steps, equivalent to : 
178		p *= exp(-δ*damp_p); q *= exp(-δ*damp_q)
179		Please set damp_q and damp_p as appropriate.
180		"""
181		if self.damp_q is not damp_None:
182			if self.preserve_q: q = self._mk_copy(q)
183			q *= np.exp(-δ*self.damp_q) # Maybe we can avoid recomputing these exponentials
184		if self.damp_p is not damp_None:
185			if self.preserve_p: p = self._mk_copy(p)
186			p *= np.exp(-δ*self.damp_p)
187		return q,p
188
189	def Impl_p(self,q,p,δ):
190		"""
191		Implicit time step for the impulsion p.
192		"""
193		return self.impl_solver(lambda p_:p-δ*self.DqH(q,p_), p)
194
195	def Impl_q(self,q,p,δ):
196		"""
197		Implicit time step for the position q.
198		"""
199		return self.impl_solver(lambda q_:q-δ*self.DqH(q_,p), q)
200
201	def Impl2_p(self,q,p,δ_before,δ_total,δ_after):
202		"""
203		Merge two implicit time steps for the impulsion p, with a damping step in between.
204		"""
205		#Default implementation, without optimization (i.e. merging of the Impl_p steps)
206		p   = self.Impl_p(q,p,δ_before)
207		self.read_q(self,q); self.read_p(self,p) # Read position and momentum before damp
208		q,p = self.Damp_qp(q,p,δ_total)
209		self.incr_q(self,q); self.incr_p(self,p) # Modify position and momentum after damp
210		p   = self.Impl_p(q,p,δ_after)
211		return q,p
212
213	# ---------------- Symplectic schemes ---------------
214
215	def Euler_p(self,q,p,δ,niter=1):
216		"""niter time steps of the symplectic Euler scheme, starting with impulsion p update."""
217		for i in range(niter):
218			self.current_iter = i
219			p=self.Impl_p(q,p,δ)
220			q=self.Expl_q(q,p,δ)
221			if i!=niter-1:
222				self.read_q(self,q); self.read_p(self,p) 
223				self.Damp_qp(q,p,δ) 
224				self.incr_q(self,q); self.incr_p(self,p)
225		return q,p
226
227	def Euler_q(self,q,p,δ,niter=1):
228		"""niter time steps of the symplectic Euler scheme, starting with position q update."""
229		for i in range(niter):
230			self.current_iter = i
231			q=self.Impl_q(q,p,δ)
232			p=self.Expl_p(q,p,δ)
233			if i!=niter-1:
234				self.read_q(self,q); self.read_p(self,p) 
235				self.Damp_qp(q,p,δ) 
236				self.incr_q(self,q); self.incr_p(self,p)
237		return q,p
238
239	def Verlet_p(self,q,p,δ,niter=1):
240		"""
241		niter time steps of the symplectic Verlet scheme.
242		Optional damping steps interleaved.
243		"""
244		hδ=δ/2
245		for i in range(niter):
246			self.current_iter = i
247			if i==0: p=self.Impl_p(q,p,hδ)
248			q=self.Expl_q(q,p,δ)
249			if i==niter-1: p=self.Impl_p(q,p,hδ)
250			else: q,p=self.Impl2_p(q,p,hδ,δ,hδ)
251		return q,p
252
253	def Ruth4_p(self,q,p,δ,niter=1):
254		"""
255		niter time steps of the Ruth 1983 4th order symplectic scheme, as of Wikipedia.
256		Optional damping steps interleaved.
257		"""
258		t = 2.**(1./3.)
259		c1 = 1/(2*(2-t)); c2 = (1-t)*c1; 
260		d1 = 2*c1; d2 = -t*d1;
261		for i in range(niter):
262			self.current_iter = i
263			if i==0: p=self.Impl_p(q,p,c1*δ)
264			q=self.Expl_q(q,p,d1*δ)
265			p=self.Impl_p(q,p,c2*δ)
266			q=self.Expl_q(q,p,d2*δ)
267			p=self.Impl_p(q,p,c2*δ)
268			q=self.Expl_q(q,p,d1*δ)
269			if i==niter-1: p=self.Impl_p(q,p,c1*δ) 
270			else: q,p=self.Impl2_p(q,p,c1*δ,δ,c1*δ)
271		return q,p
272
273	def Sympl_p(self,q,p,δ,niter=1,order=2):
274		"""
275		niter steps of the Euler_p, Verlet_p or Ruth4_p symplectic scheme, 
276		depending on the order parameter.
277		"""
278		if order==1: return self.Euler_p(q,p,δ,niter)
279		if order==2: return self.Verlet_p(q,p,δ,niter)
280		if order==4: return self.Ruth4_p( q,p,δ,niter)
281		raise ValueError(f"Found {order=}, while expecting 1,2 or 4.")
282
283	def symplectic_schemes(self):
284		return {'Euler-p':self.Euler_p,'Euler-q':self.Euler_q,'Verlet-p':self.Verlet_p,
285		'Ruth4-p':self.Ruth4_p}
286
287	# ------ Routines for backpropagation ------
288
289	def _Sympl_p(self,q,p,δ,order=2):
290		"""
291		Returns a callable, which computes any given iteration of Sympl_p, and saves appropriate 
292		keypoints for computational and memory efficiency.
293		"""
294		# TODO : remove since probably useless ? (Only use case is backprop in seismogram)
295
296		# The main objective of this routine is to preserve accuracy when damp_p and damp_q are set.
297		# Otherwise, the iterations may be reversed by iterating the symplectic scheme 
298		# in negative time, providing a result with similar accuracy. 
299		def next(qp): # A single time step, including a preliminary damping
300			qp = self.Damp_qp(*qp,δ)
301			return self.Sympl_p(*qp,δ,order=order)
302		from .backtrack import RecurseRewind
303		return RecurseRewind(next,self.Damp_qp(q,p,-δ)) # A single negative damping step should be fine...
304
305	def seismogram(self,q,p,δ,niter,order=2,qh_ind=None,ph_ind=None):
306
307		H_fwd = copy(self)
308		H_fwd.read_q = H_fwd.read_p = read_None; H_fwd.incr_q = H_fwd.incr_p = incr_None
309		qh = []; ph = []; initial=True 
310		if qh_ind is not None: H_fwd.read_q = lambda _,q : qh.append(q.reshape(-1)[qh_ind])
311		if ph_ind is not None: H_fwd.read_p = lambda _,p : ph.append(p.reshape(-1)[ph_ind])
312		qf,pf = H_fwd.Sympl_p(q,p,δ,niter,order)
313
314		return qf,pf,ad.asarray(qh),ad.asarray(ph)