agd.ODE.backtrack

  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
  5import itertools
  6import copy
  7import numpy as np
  8from ..Interpolation import map_coordinates,origin_scale_shape #UniformGridInterpolation
  9
 10def odeint_array(f,y,t,grid,t_delay=0,t_substeps=2,order=1,**kwargs):
 11	"""
 12	Solve an ODE where the vector field is interpolated from an array.
 13	The vector field is linearly interpolated, and the Euler midpoint scheme is used.
 14
 15	Inputs : 
 16	- f : generator of the vector fields to be interpolated, len(next(f))=vdim
 17	- y : initial value of the solution, len(y) = vdim 
 18	- t : time steps
 19	- grid : interpolation grid for the vector field f
 20	- t_delay : number of time steps to drop initially 
 21	- t_substeps : number of substeps
 22	- order : passed to odeint_array
 23	- **kwargs : passed to UniformGridInterpolation
 24	"""
 25	nt = len(t)
 26	solution = np.full_like(y,np.nan,shape=(nt,*y.shape))
 27	y0 = copy.copy(y)
 28	fit = iter(f)
 29	origin,scale,_ = origin_scale_shape(grid)
 30	def I(values,positions): # Interpolation
 31		return map_coordinates(values,positions,origin=origin,scale=scale,
 32			depth=1,order=order,**kwargs)
 33
 34	f1 = next(fit)
 35	for i in range(nt-1):
 36		f0,f1 = f1,next(fit) # Vector fields to be interpolated
 37		t0,t1 = t[i],t[i+1]
 38		dt = (t1-t0)/t_substeps
 39		solution[i][:,t_delay==i] = y[:,t_delay==i] # Insert initial values
 40		valid = t_delay <= i 
 41		y0 = solution[i] if np.all(valid) else solution[i][:,valid]
 42		for s in range(t_substeps):
 43			w = s/t_substeps
 44			y1 = y0 + 0.5*dt * ((1-w)*I(f0,y0) + w*I(f1,y0)) # Compute midpoint
 45			w = (s+0.5)/t_substeps
 46			y0 = y0 +     dt * ((1-w)*I(f0,y1) + w*I(f1,y1)) # Make step
 47		if np.all(valid): solution[i+1] = y0
 48		else: solution[i+1][:,valid]=y0 
 49
 50	return solution
 51
 52class RecurseRewind:
 53	r"""
 54	This class is designed to iterate a function, and then roll back the iterations, 
 55	with a limited memory usage. For that purpose appropriate keypoints are stored,
 56	 in a multilevel manner.
 57	__init__ args: 
 58	 - next : method to compute the next step.
 59	 - initial : initial value
 60	 - params : passed to next, in addition to current value
 61	 - base : base $b$ used to internally represent the iteration counter, which balances 
 62	 a tradeoff between memory usage and computational cost. Iterating $O(b^n)$ times
 63	  and then rewinding these iterations, has a storage cost of $O(n*b)$ and a 
 64	  computational cost of $O(n * b^n)$ function evaluations.
 65
 66	members:
 67	 - reversed : wether __next__ should advance or rewind the iterations
 68	"""
 69
 70	def __init__(self,next,initial,params=tuple(),base=5):
 71		self.next = next
 72		self.params = params
 73		self._base = base
 74		self._keys = {0:initial}
 75		self._index = 0
 76		self.reversed = False
 77		self._initial_value = True
 78
 79	def __call__(self,index=None):
 80		"""The value of the current iteration if index is None, otherwise moves to index"""
 81		if index is not None: self.moveto(index)
 82		return self.value
 83
 84	def __iter__(self):
 85		self._initial_value = True
 86		return self
 87
 88	def __next__(self):
 89		if self._initial_value: self._initial_value = False
 90		elif self.reversed: self.rewind()
 91		else: self.advance()
 92		return self.value
 93
 94
 95	def advance(self):
 96		"""Compute next iterate of f, and update keypoints"""
 97		self._keys[self.index+1] = self.next(self._keys[self.index],*self.params)
 98		self._index += 1
 99		# Delete previous keys to save storage
100		for i in range(self._basepow):
101			k = self.base**i
102			for j in range(1,self.base): 
103				self._keys.pop(self.index - j*k)
104
105	def rewind(self):
106		"""Compute previous iterate of f, and update keypoints"""
107		if self.index<=0: raise StopIteration
108		self._keys.pop(self.index)
109		i = self._basepow
110		if i:
111			k = self.base**i
112			self._index -= k
113			for _ in range(k-1):
114				self.advance()
115		else:
116			self._index -= 1
117
118	def moveto(self,index):
119		"""Compute prescribed iterate of f, and update keypoints"""
120		# Possible improvements here, especially for functions which can be iterated at a reduced 
121		# cost, e.g. by merging the final and initial substeps of symplectic schemes for wave PDEs
122		if index < self.index: 
123			for i in range(self.index-index): self.rewind()
124		if index > self.index:
125			for i in range(index-self.index): self.advance()
126
127	@property
128	def _basepow(self):
129		"""Returns the largest i such that self.base**i divides self.index"""
130		assert self.index!=0
131		basepow = self.base
132		i = 0
133		while self.index%basepow == 0:
134			basepow *= self.base
135			i+=1
136		return i
137
138	@property
139	def index(self):
140		"""The index of the current iteration"""
141		return self._index
142	@property
143	def value(self): 
144		"""The value of the current iteration"""
145		return self._keys[self.index]
146	@property
147	def base(self):return self._base
148
149