Permalink
Newer
100644
124 lines (113 sloc)
5.31 KB
|
||
| 1 | from scipy.integrate import ode | |
| 2 | import numpy as np | |
| 3 | ||
| 4 | # A default callback function | |
| 5 | # the geodesic can be halted by the callback returning False | |
| 6 | def callback_func(geo): | |
| 7 | return True | |
| 8 | ||
| 9 | # We construct a geodesic class that inherits from scipy's ode class | |
| 10 | class geodesic(ode): | |
| 11 | ||
| 12 | def __init__(self, r, j, Avv, M, N, x, v, lam = 0.0, dtd = None, atol = 1e-6, rtol = 1e-6, callback = None, parameterspacenorm = False,invSVD=False): | |
| 13 | """ | |
| 14 | Construct an instance of the geodesic object | |
| 15 | r = function for calculating the model. This is not needed for the geodesic, but is used to save the geodesic path in data space. Called as r(x) | |
| 16 | j = function for calculating the jacobian of the model with respect to parameters. Called as j(x) | |
| 17 | Avv = function for calculating the second directional derivative of the model in a direction v. Called as Avv(x,v) | |
| 18 | M = Number of model predictions | |
| 19 | N = Number of model parameters | |
| 20 | x = vector of initial parameter values | |
| 21 | v = vector of initial velocity | |
| 22 | lam = set to nonzero to calculate geodesic on the model graph. | |
| 23 | dtd = metric for the parameter space contribution to the model graph | |
| 24 | atol = absolute tolerance for solving the geodesic | |
| 25 | rtol = relative tolerance for solving geodesic | |
| 26 | callback = function called after each geodesic step. Called as callback(geo) where geo is the current instance of the geodesic class. | |
| 27 | parameterspacenorm = Set to True to reparameterize the geodesic to have a constant parameter space norm. (Default is to have a constant data space norm.) | |
| 28 | invSVD = Set to true to use the singular value decomposition to calculate the inverse metric. This is slower, but can help with nearly singular FIM. | |
| 29 | """ | |
| 30 | self.r, self.j, self.Avv = r, j, Avv | |
| 31 | self.M, self.N = M, N | |
| 32 | self.lam = lam | |
| 33 | if dtd is None: | |
| 34 | self.dtd = np.eye(N) | |
| 35 | else: | |
| 36 | self.dtd = dtd | |
| 37 | self.atol = atol | |
| 38 | self.rtol = rtol | |
| 39 | ode.__init__(self, self.geodesic_rhs, jac = None) | |
| 40 | self.set_initial_value(x, v) | |
| 41 | ode.set_integrator(self, 'vode', atol = atol, rtol = rtol) | |
| 42 | if callback is None: | |
| 43 | self.callback = callback_func | |
| 44 | else: | |
| 45 | self.callback = callback | |
| 46 | self.parameterspacenorm = parameterspacenorm | |
| 47 | self.invSVD = False | |
| 48 | ||
| 49 | def geodesic_rhs(self, t, xv): | |
| 50 | """ | |
| 51 | This function implements the RHS of the geodesic equation | |
| 52 | This should not need to edited or called directly by the user | |
| 53 | t = "time" of the geodesic (tau) | |
| 54 | xv = vector of current parameters and velocities | |
| 55 | """ | |
| 56 | x = xv[:self.N] | |
| 57 | v = xv[self.N:] | |
| 58 | j = self.j(x) | |
| 59 | g = np.dot(j.T, j) + self.lam*self.dtd | |
| 60 | Avv = self.Avv(x, v) | |
| 61 | if self.invSVD: | |
| 62 | u,s,vh = np.linalg.svd(j,0) | |
| 63 | a = - np.dot( vh.T, np.dot(u.T, Avv)/s ) | |
| 64 | else: | |
| 65 | a = -np.linalg.solve(g, np.dot(j.T, Avv) ) | |
| 66 | if self.parameterspacenorm: | |
| 67 | a -= np.dot(a,v)*v/np.dot(v,v) | |
| 68 | return np.append(v, a) | |
| 69 | ||
| 70 | def set_initial_value(self, x, v): | |
| 71 | """ | |
| 72 | Set the initial parameter values and velocities | |
| 73 | x = vector of initial parameter values | |
| 74 | v = vector of initial velocity | |
| 75 | """ | |
| 76 | self.xs = np.array([x]) | |
| 77 | self.vs = np.array([v]) | |
| 78 | self.ts = np.array([0.0]) | |
| 79 | self.rs = np.array([ self.r(x) ] ) | |
| 80 | self.vels = np.array([ np.dot(self.j(x), v) ] ) | |
| 81 | ode.set_initial_value( self, np.append(x, v), 0.0 ) | |
| 82 | ||
| 83 | def step(self, dt = 1.0): | |
| 84 | """ | |
| 85 | Integrate the geodesic for one step | |
| 86 | dt = target time step to use | |
| 87 | """ | |
| 88 | ode.integrate(self, self.t + dt, step = 1) | |
| 89 | self.xs = np.append(self.xs, [self.y[:self.N]], axis = 0) | |
| 90 | self.vs = np.append(self.vs, [self.y[self.N:]], axis = 0 ) | |
| 91 | self.rs = np.append(self.rs, [self.r(self.xs[-1])], axis = 0) | |
| 92 | self.vels = np.append(self.vels, [np.dot(self.j(self.xs[-1]), self.vs[-1])], axis = 0) | |
| 93 | self.ts = np.append(self.ts, self.t) | |
| 94 | ||
| 95 | def integrate(self, tmax, maxsteps = 500): | |
| 96 | """ | |
| 97 | Integrate the geodesic up to a fixed maximimum time (tau) or number of steps. | |
| 98 | After each step, the callback is called and the calculation will end if the callback returns False | |
| 99 | tmax = maximum time (tau) for integration | |
| 100 | maxsteps = maximum number of steps | |
| 101 | """ | |
| 102 | cont = True | |
| 103 | while self.successful() and len(self.xs) < maxsteps and self.t < tmax and cont: | |
| 104 | self.step(tmax - self.t) | |
| 105 | cont = self.callback(self) | |
| 106 | ||
| 107 | def InitialVelocity(x, jac, Avv): | |
| 108 | """ | |
| 109 | Routine for calculating the initial velocity | |
| 110 | x: Initial Parameter Values | |
| 111 | jac: function for calculating the jacobian. Called as jac(x) | |
| 112 | Avv: function for calculating a direction second derivative. Called as Avv(x,v) | |
| 113 | """ | |
| 114 | ||
| 115 | j = jac(x) | |
| 116 | u,s,vh = np.linalg.svd(j) | |
| 117 | v = vh[-1] | |
| 118 | a = -np.linalg.solve( np.dot(j.T, j), np.dot(j.T, Avv(x,v) )) | |
| 119 | # We choose the direction in which the velocity will increase, since this is most likely to find a singularity quickest. | |
| 120 | if np.dot(v, a) < 0: | |
| 121 | v *= -1 | |
| 122 | return v | |
| 123 |