├── .DS_Store ├── .gitattributes ├── .gitignore ├── .ipynb_checkpoints ├── double-well-checkpoint.ipynb ├── dwutils-checkpoint.py ├── pitchfork-checkpoint.ipynb ├── sim_doublewell-checkpoint.jl ├── sim_pitchfork-checkpoint.jl └── utils-checkpoint.py ├── README.md ├── __pycache__ ├── dwutils.cpython-38.pyc ├── fpsolve.cpython-38.pyc └── utils.cpython-38.pyc ├── data ├── .DS_Store └── wake-cop.mat ├── double-well.ipynb ├── dwutils.py ├── fpsolve.py ├── mean-field-model ├── README.md ├── data │ ├── .gitattributes │ ├── monte_carlo.mat │ ├── order_parameters.mat │ └── wake_expt.mat ├── fpsolve.py ├── lindy.py ├── meanfield.py ├── symmetry-breaking.ipynb └── utils.py ├── pitchfork.ipynb ├── sim_doublewell.jl ├── sim_pitchfork.jl ├── utils.py └── wake.ipynb /.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dynamicslab/langevin-regression/35468293c3b133a79562924375d80b5155d763fc/.DS_Store -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | *.mat filter=lfs diff=lfs merge=lfs -text 2 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | .ipynb_checkpoints 3 | __pycache__ 4 | __pycache__/* 5 | -------------------------------------------------------------------------------- /.ipynb_checkpoints/dwutils-checkpoint.py: -------------------------------------------------------------------------------- 1 | """ 2 | Utility functions for stochastic modeling of the double-well potential 3 | 4 | Jared Callaham (2020) 5 | """ 6 | 7 | import numpy as np 8 | import sympy 9 | from scipy.optimize import curve_fit 10 | 11 | import utils 12 | import fpsolve 13 | 14 | import warnings 15 | warnings.simplefilter(action='ignore', category=FutureWarning) 16 | 17 | import julia 18 | jl = julia.Julia() 19 | jl.include("../sim_doublewell.jl") 20 | 21 | def switched_states(X, thresh=1): 22 | # Decide whether the bistable states are "up" or "down" based on threshold 23 | 24 | N = len(X) 25 | state = np.zeros((N)) 26 | 27 | # Step forward until X is either positive or negative (if simulation is initialized at 0) 28 | idx = 0 29 | while abs(X[idx]) < thresh: 30 | idx += 1 31 | 32 | # Consider everything up to this point to be in the first well 33 | cur_state = np.sign(X[idx]) 34 | state[:idx] = cur_state 35 | 36 | while idx thresh: 40 | cur_state = -cur_state 41 | 42 | state[idx] = cur_state 43 | idx += 1 44 | 45 | return state 46 | 47 | 48 | def dwell_times(states, dt=1): 49 | # Given a vector of states (i.e. from switched_states() ), return dwell time in each state 50 | N = len(states) 51 | switch_times = [] # List of dwell times 52 | idx = 0 53 | last_switch = idx 54 | cur_state = states[idx] 55 | 56 | while idx < N: 57 | if states[idx] != cur_state: 58 | switch_times.append( dt*(idx-last_switch) ) 59 | last_switch = idx 60 | cur_state = states[idx] 61 | 62 | idx += 1 63 | 64 | return switch_times 65 | 66 | def dwell_stats(X, thresh, dt): 67 | 68 | state = switched_states(X, thresh=thresh) # Categorize into "up" or "down" 69 | switch_times = dwell_times(state, dt=dt) # Compute list of dwell times in each metastable state 70 | 71 | if len(switch_times) > 0: 72 | return np.mean(switch_times), np.std(switch_times)/np.sqrt(len(switch_times)) 73 | else: 74 | return np.nan, np.nan 75 | 76 | 77 | def fit_pdf(X, edges, p_hist, dt, p0=None): 78 | # Mean-square displacement 79 | fit_start, fit_stop = 0.1, 3 80 | n_lags = int(fit_stop/dt) 81 | tau = dt*np.arange(1, n_lags) 82 | 83 | # Lagged mean square displacement 84 | msd = np.zeros((len(tau))) 85 | for i in range(1, n_lags): 86 | msd[i-1] = np.mean( (X[i:]-X[:-i])**2) 87 | 88 | # Linear fit for radial displacement 89 | to_fit = np.nonzero( (tau > fit_start) * (tau < fit_stop) )[0] 90 | p_rad = np.polyfit(tau[to_fit], msd[to_fit], deg=1) 91 | a_pdf = 0.5*p_rad[0] 92 | 93 | # Fit PDF 94 | p_model = lambda x, C, a, b: C*np.exp(a*x**2 + b*x**4) 95 | centers = 0.5*(edges[1:]+edges[:-1]) 96 | if p0 is not None: 97 | popt, pcov = curve_fit(p_model, centers, p_hist, p0=p0) 98 | else: 99 | popt, pcov = curve_fit(p_model, centers, p_hist) 100 | 101 | # Separate parameters in model 102 | sigma_pdf = np.sqrt(2*a_pdf) 103 | lamb_pdf = popt[1]*sigma_pdf**2 104 | mu_pdf = popt[2]*2*sigma_pdf**2 105 | 106 | return lamb_pdf, mu_pdf, sigma_pdf 107 | 108 | 109 | def langevin_regression(X, edges, p_hist, dt, stride=200, kl_reg=0): 110 | """ 111 | Wrapper for full Langevin regression so we can loop over it to explore variations with distance from bifurcation 112 | """ 113 | centers = 0.5*(edges[1:]+edges[:-1]) 114 | N = len(centers) 115 | 116 | # Kramers-Moyal average 117 | tau = stride*dt 118 | f_KM, a_KM, f_err, a_err = KM_avg(X, bins, stride=stride, dt=dt) 119 | 120 | # Initialize libraries 121 | x = sympy.symbols('x') 122 | 123 | f_expr = np.array([x**i for i in [1, 3]]) # Polynomial library for drift 124 | s_expr = np.array([x**i for i in [0]]) # Polynomial library for diffusion 125 | 126 | lib_f = np.zeros([len(f_expr), N]) 127 | for k in range(len(f_expr)): 128 | lamb_expr = sympy.lambdify(x, f_expr[k]) 129 | lib_f[k] = lamb_expr(centers) 130 | 131 | lib_s = np.zeros([len(s_expr), N]) 132 | for k in range(len(s_expr)): 133 | lamb_expr = sympy.lambdify(x, s_expr[k]) 134 | lib_s[k] = lamb_expr(centers) 135 | 136 | # Initialize Xi with plain least-squares (just helpf the optimization a bit) 137 | Xi0 = np.zeros((len(f_expr) + len(s_expr))) 138 | mask = np.nonzero(np.isfinite(f_KM))[0] 139 | Xi0[:len(f_expr)] = np.linalg.lstsq( lib_f[:, mask].T, f_KM[mask], rcond=None)[0] 140 | Xi0[len(f_expr):] = np.linalg.lstsq( lib_s[:,mask].T, np.sqrt(2*a_KM[mask]), rcond=None)[0] 141 | 142 | # Parameter dictionary for optimization 143 | W = np.array((f_err.flatten(), a_err.flatten())) 144 | W[np.less(abs(W), 1e-12, where=np.isfinite(W))] = 1e6 # Set zero entries to large weights 145 | W[np.logical_not(np.isfinite(W))] = 1e6 # Set NaN entries to large numbers (small weights) 146 | W = 1/W # Invert error for weights 147 | W = W/np.nansum(W.flatten()) 148 | 149 | # Adjoint solver 150 | afp = fpsolve.AdjFP(centers) 151 | 152 | # Forward solver 153 | fp = fpsolve.SteadyFP(N, centers[1]-centers[0]) 154 | 155 | params = {"W": W, "f_KM": f_KM, "a_KM": a_KM, "Xi0": Xi0, 156 | "f_expr": f_expr, "s_expr": s_expr, 157 | "lib_f": lib_f.T, "lib_s": lib_s.T, "N": N, 158 | "kl_reg": kl_reg, 159 | "fp": fp, "afp": afp, "p_hist": p_hist, "tau": tau, 160 | "radial": False} 161 | 162 | # Tune KL regularization automatically 163 | Xi, _ = utils.AFP_opt(utils.cost, params) 164 | return Xi 165 | 166 | def model_eval(eps, sigma, N, kl_reg): 167 | """ 168 | Construct and evaluate all models of the double-well 169 | 1. Analytic normal form model 170 | 2. PDF fitting without Kramers-Moyal average 171 | 3. Full Langevin regression 172 | """ 173 | ### Generate data 174 | x_eq = np.sqrt(eps) # Equilibrium value 175 | 176 | edges = np.linspace(-2*x_eq, 2*x_eq, N+1) 177 | centers = 0.5*(edges[:-1]+edges[1:]) 178 | dx = centers[1]-centers[0] 179 | 180 | dt = 1e-2 181 | tmax = int(1e5) 182 | t, X = jl.run_sim(eps, sigma, dt, tmax) 183 | X, V = X[0, :], X[1, :] 184 | 185 | # PDF of states 186 | p_hist = np.histogram(X, edges, density=True)[0] 187 | 188 | # Dwell-time slope 189 | b, b_err = dwell_stats(X, x_eq, dt) 190 | print("\tData: ", b, b_err) 191 | 192 | ### 1. Normal form 193 | lamb1 = -1 + np.sqrt(1 + eps) 194 | lamb2 = -1 - np.sqrt(1 + eps) 195 | h = -lamb1/lamb2 196 | mu = -(1+h)**2*lamb1/eps 197 | 198 | _, phi1 = jl.run_nf(lamb1, mu, sigma/(2*np.sqrt(1+eps)), dt, tmax) 199 | X_nf = (1+h)*phi1[0, :] 200 | 201 | # Statistics 202 | p_nf = np.histogram(X_nf, edges, density=True)[0] 203 | b_nf, b_nf_err = dwell_stats(X_nf, x_eq, dt) 204 | print("\tNormal form: ", b_nf, b_nf_err) 205 | 206 | ### 2. PDF fit 207 | Xi = fit_pdf(X, edges, p_hist, dt, p0=[1, lamb1/sigma**2, mu/sigma**2]) 208 | #print(Xi) 209 | 210 | # Monte Carlo evaluation 211 | _, X_pdf = jl.run_nf(Xi[0], Xi[1], Xi[2], dt, tmax) 212 | X_pdf = X_pdf[0, :] 213 | 214 | # Statistics 215 | p_pdf = np.histogram(X_pdf, edges, density=True)[0] 216 | b_pdf, b_pdf_err = dwell_stats(X_pdf, x_eq, dt) 217 | print("\tPDF fit: ", b_pdf, b_pdf_err) 218 | 219 | ### 3. Langevin regression 220 | Xi = langevin_regression(X, edges, p_hist, dt, stride=200, kl_reg=kl_reg) 221 | #print(Xi) 222 | 223 | # Monte Carlo evaluation 224 | _, X_lr = jl.run_nf(Xi[0], Xi[1], Xi[2], dt, tmax) 225 | X_lr = X_lr[0, :] 226 | 227 | # Statistics 228 | p_lr = np.histogram(X_lr, edges, density=True)[0] 229 | b_lr, b_lr_err = dwell_stats(X_lr, x_eq, dt) 230 | print("\tLangevin regression: ", b_lr, b_lr_err) 231 | 232 | ### KL-divergence of all models against true data 233 | KL_nf = utils.kl_divergence(p_hist, p_nf, dx=dx, tol=1e-6) 234 | KL_pdf = utils.kl_divergence(p_hist, p_pdf, dx=dx, tol=1e-6) 235 | KL_lr = utils.kl_divergence(p_hist, p_lr, dx=dx, tol=1e-6) 236 | print("\tKL div: ", KL_nf, KL_pdf, KL_lr) 237 | 238 | return [b, b_nf, b_pdf, b_lr], [KL_nf, KL_pdf, KL_lr] -------------------------------------------------------------------------------- /.ipynb_checkpoints/sim_doublewell-checkpoint.jl: -------------------------------------------------------------------------------- 1 | """ 2 | Julia code to simulate 1D particle in a double-well potential 3 | 4 | Jared Callaham (2020) 5 | """ 6 | using DifferentialEquations 7 | 8 | # Noisy double well potential 9 | 10 | ### Physical dynamics 11 | function f_fn(dx, x, p, t) 12 | ϵ, σ = p 13 | dx[1] = x[2] 14 | dx[2] = -2*x[2] + ϵ*x[1] - x[1]^3 15 | end 16 | 17 | function σ_fn(dx, x, p, t) 18 | ϵ, σ = p 19 | dx[1] = 0 20 | dx[2] = σ 21 | end 22 | 23 | ### Normal form dynamics (pitchfork bifurcation) 24 | function fn_fn(dx, x, p, t) 25 | λ, μ, σ = p 26 | dx[1] = λ*x[1] + μ*x[1]^3 27 | end 28 | 29 | function σn_fn(dx, x, p, t) 30 | λ, μ, σ = p 31 | dx[1] = σ 32 | end 33 | 34 | function run_nf(ϵ, μ, σ, dt, tmax) 35 | prob = SDEProblem(fn_fn, σn_fn, [0.0] , (0.0, tmax), [ϵ, μ, σ]) 36 | sol = solve(prob, SRIW1(), dt=dt, adaptive=false); 37 | return sol.t, sol[:, :] 38 | end 39 | 40 | function run_sim(η, σ, dt, tmax) 41 | prob = SDEProblem(f_fn,σ_fn, [0.0, 0.0], (0.0, tmax), [η, σ]) 42 | sol = solve(prob, SRIW1(), dt=dt, adaptive=false) 43 | return sol.t, sol[:, :] 44 | end -------------------------------------------------------------------------------- /.ipynb_checkpoints/sim_pitchfork-checkpoint.jl: -------------------------------------------------------------------------------- 1 | """ 2 | Julia code to simulate pitchfork bifurcation normal form forced by colored noise 3 | 4 | Jared Callaham (2020) 5 | """ 6 | using DifferentialEquations 7 | using MAT 8 | 9 | # Drift component 10 | μ = 1 # Fixed points at +/- 1 11 | λ = 1 12 | 13 | r0 = sqrt(λ/μ) # Equilibrium point 14 | α = Int(1e2) # Inverse autocorrelation time of the forcing (smaller than lambda for scale separation) 15 | σ = α*0.5 16 | 17 | f_ex(x) = λ*x - μ*x^3 18 | 19 | # Simulate SDE 20 | function f_fn(dx, x, p, t) 21 | dx[1] = f_ex(x[1]) + x[2] 22 | dx[2] = -α*x[2] 23 | end 24 | 25 | function σ_fn(dx, x, p, t) 26 | dx[1] = 0 27 | dx[2] = σ 28 | end 29 | 30 | dt = 0.001 31 | x0 = [0.0, 0.0] 32 | tspan = (0.0, Int(1e4)) 33 | prob = SDEProblem(f_fn, σ_fn, x0, tspan) 34 | 35 | # EM() or SRIW1() 36 | sol = solve(prob, EM(), dt=dt); 37 | t = sol.t 38 | X = sol[1, :] 39 | 40 | matwrite("./data/pitchfork.mat", Dict( 41 | "X" => X, 42 | "dt" => dt, 43 | "lamb" => λ, 44 | "mu" => μ 45 | ); compress = true) -------------------------------------------------------------------------------- /.ipynb_checkpoints/utils-checkpoint.py: -------------------------------------------------------------------------------- 1 | """ 2 | Utility functions for Langevin regression 3 | 4 | Jared Callaham (2020) 5 | """ 6 | 7 | import numpy as np 8 | from time import time 9 | from scipy.optimize import minimize 10 | 11 | # Return a single expression from a list of expressions and coefficients 12 | # Note this will give a SymPy expression and not a function 13 | def sindy_model(Xi, expr_list): 14 | return sum([Xi[i]*expr_list[i] for i in range(len(expr_list))]) 15 | 16 | 17 | def ntrapz(I, dx): 18 | if isinstance(dx, int) or isinstance(dx, float) or len(dx)==1: 19 | return np.trapz(I, dx=dx, axis=0) 20 | else: 21 | return np.trapz( ntrapz(I, dx[1:]), dx=dx[0]) 22 | 23 | 24 | def kl_divergence(p_in, q_in, dx=1, tol=None): 25 | """ 26 | Approximate Kullback-Leibler divergence for arbitrary dimensionality 27 | """ 28 | if tol==None: 29 | tol = max( min(p_in.flatten()), min(q_in.flatten())) 30 | q = q_in.copy() 31 | p = p_in.copy() 32 | q[q bins[i]) * (Y[:-1] < bins[i+1]) )[0] 51 | 52 | if len(mask) > 0: 53 | f_KM[i] = np.mean(dY[mask]) # Conditional average ~ drift 54 | a_KM[i] = 0.5*np.mean(dY2[mask]) # Conditional variance ~ diffusion 55 | 56 | # Estimate error by variance of samples in the bin 57 | a_KM[i] = 0.5*np.mean(dY2[mask]) # Conditional average 58 | a_err[i] = np.std(dY2[mask])/np.sqrt(len(mask)) 59 | 60 | else: 61 | f_KM[i] = np.nan 62 | f_err[i] = np.nan 63 | a_KM[i] = np.nan 64 | a_err[i] = np.nan 65 | 66 | return f_KM, a_KM, f_err, a_err 67 | 68 | 69 | # Return optimal coefficients for finite-time correctio 70 | def AFP_opt(cost, params): 71 | ### RUN OPTIMIZATION PROBLEM 72 | start_time = time() 73 | Xi0 = params["Xi0"] 74 | 75 | is_complex = np.iscomplex(Xi0[0]) 76 | 77 | if is_complex: 78 | Xi0 = np.concatenate((np.real(Xi0), np.imag(Xi0))) # Split vector in two for complex 79 | opt_fun = lambda Xi: cost(Xi[:len(Xi)//2] + 1j*Xi[len(Xi)//2:], params) 80 | 81 | else: 82 | opt_fun = lambda Xi: cost(Xi, params) 83 | 84 | res = minimize(opt_fun, Xi0, method='nelder-mead', 85 | options={'disp': False, 'maxfev':int(1e4)}) 86 | print('%%%% Optimization time: {0} seconds, Cost: {1} %%%%'.format(time() - start_time, res.fun) ) 87 | 88 | # Return coefficients and cost function 89 | if is_complex: 90 | # Return to complex number 91 | return res.x[:len(res.x)//2] + 1j*res.x[len(res.x)//2:], res.fun 92 | else: 93 | return res.x, res.fun 94 | 95 | 96 | 97 | def SSR_loop(opt_fun, params): 98 | """ 99 | Stepwise sparse regression: general function for a given optimization problem 100 | opt_fun should take the parameters and return coefficients and cost 101 | 102 | Requires a list of drift and diffusion expressions, 103 | (although these are just passed to the opt_fun) 104 | """ 105 | 106 | # Lists of candidate expressions... coefficients are optimized 107 | f_expr, s_expr = params['f_expr'].copy(), params['s_expr'].copy() 108 | lib_f, lib_s = params['lib_f'].copy(), params['lib_s'].copy() 109 | Xi0 = params['Xi0'].copy() 110 | 111 | m = len(f_expr) + len(s_expr) 112 | 113 | Xi = np.zeros((m, m-1), dtype=Xi0.dtype) # Output results 114 | V = np.zeros((m-1)) # Cost at each step 115 | 116 | # Full regression problem as baseline 117 | Xi[:, 0], V[0] = opt_fun(params) 118 | 119 | # Start with all candidates 120 | active = np.array([i for i in range(m)]) 121 | 122 | # Iterate and threshold 123 | for k in range(1, m-1): 124 | # Loop through remaining terms and find the one that increases the cost function the least 125 | min_idx = -1 126 | V[k] = 1e8 127 | for j in range(len(active)): 128 | tmp_active = active.copy() 129 | tmp_active = np.delete(tmp_active, j) # Try deleting this term 130 | 131 | # Break off masks for drift/diffusion 132 | f_active = tmp_active[tmp_active < len(f_expr)] 133 | s_active = tmp_active[tmp_active >= len(f_expr)] - len(f_expr) 134 | print(f_active) 135 | print(s_active) 136 | 137 | print(f_expr[f_active], s_expr[s_active]) 138 | params['f_expr'] = f_expr[f_active] 139 | params['s_expr'] = s_expr[s_active] 140 | params['lib_f'] = lib_f[:, f_active] 141 | params['lib_s'] = lib_s[:, s_active] 142 | params['Xi0'] = Xi0[tmp_active] 143 | 144 | # Ensure that there is at least one drift and diffusion term left 145 | if len(s_active) > 0 and len(f_active) > 0: 146 | tmp_Xi, tmp_V = opt_fun(params) 147 | 148 | # Keep minimum cost 149 | if tmp_V < V[k]: 150 | # Ensure that there is at least one drift and diffusion term left 151 | #if (IS_DRIFT and len(f_active)>1) or (not IS_DRIFT and len(a_active)>1): 152 | min_idx = j 153 | V[k] = tmp_V 154 | min_Xi = tmp_Xi 155 | 156 | print("Cost: {0}".format(V[k])) 157 | # Delete least important term 158 | active = np.delete(active, min_idx) # Remove inactive index 159 | Xi0[active] = min_Xi # Re-initialize with best results from previous 160 | Xi[active, k] = min_Xi 161 | print(Xi[:, k]) 162 | 163 | return Xi, V 164 | 165 | def cost(Xi, params): 166 | """ 167 | Least-squares cost function for optimization 168 | This version is only good in 1D, but could be extended pretty easily 169 | Xi - current coefficient estimates 170 | param - inputs to optimization problem: grid points, list of candidate expressions, regularizations 171 | W, f_KM, a_KM, x_pts, y_pts, x_msh, y_msh, f_expr, a_expr, l1_reg, l2_reg, kl_reg, p_hist, etc 172 | """ 173 | 174 | # Unpack parameters 175 | W = params['W'] # Optimization weights 176 | 177 | # Kramers-Moyal coefficients 178 | f_KM, a_KM = params['f_KM'].flatten(), params['a_KM'].flatten() 179 | 180 | fp, afp = params['fp'], params['afp'] # Fokker-Planck solvers 181 | lib_f, lib_s = params['lib_f'], params['lib_s'] 182 | N = params['N'] 183 | 184 | # Construct parameterized drift and diffusion functions from libraries and current coefficients 185 | f_vals = lib_f @ Xi[:lib_f.shape[1]] 186 | a_vals = 0.5*(lib_s @ Xi[lib_f.shape[1]:])**2 187 | 188 | # Solve AFP equation to find finite-time corrected drift/diffusion 189 | # corresponding to the current parameters Xi 190 | afp.precompute_operator(np.reshape(f_vals, N), np.reshape(a_vals, N)) 191 | f_tau, a_tau = afp.solve(params['tau']) 192 | 193 | # Histogram points without data have NaN values in K-M average - ignore these in the average 194 | mask = np.nonzero(np.isfinite(f_KM))[0] 195 | V = np.sum(W[0, mask]*abs(f_tau[mask] - f_KM[mask])**2) \ 196 | + np.sum(W[1, mask]*abs(a_tau[mask] - a_KM[mask])**2) 197 | 198 | # Include PDF constraint via Kullbeck-Leibler divergence regularization 199 | if params['kl_reg'] > 0: 200 | p_hist = params['p_hist'] # Empirical PDF 201 | p_est = fp.solve(f_vals, a_vals) # Solve Fokker-Planck equation for steady-state PDF 202 | kl = utils.kl_divergence(p_hist, p_est, dx=fp.dx, tol=1e-6) 203 | kl = max(0, kl) # Numerical integration can occasionally produce small negative values 204 | V += params['kl_reg']*kl 205 | return V 206 | 207 | 208 | # 1D Markov test 209 | def markov_test(X, lag, N=32, L=2): 210 | # Lagged time series 211 | X1 = X[:-2*lag:lag] 212 | X2 = X[lag:-lag:lag] 213 | X3 = X[2*lag::lag] 214 | 215 | # Two-time joint pdfs 216 | bins = np.linspace(-L, L, N+1) 217 | dx = bins[1]-bins[0] 218 | p12, _, _ = np.histogram2d(X1, X2, bins=[bins, bins], density=True) 219 | p23, _, _ = np.histogram2d(X2, X3, bins=[bins, bins], density=True) 220 | p2, _ = np.histogram(X2, bins=bins, density=True) 221 | p2[p2<1e-4] = 1e-4 222 | 223 | # Conditional PDF (Markov assumption) 224 | pcond_23 = p23.copy() 225 | for j in range(pcond_23.shape[1]): 226 | pcond_23[:, j] = pcond_23[:, j]/p2 227 | 228 | # Three-time PDFs 229 | p123, _ = np.histogramdd(np.array([X1, X2, X3]).T, bins=np.array([bins, bins, bins]), density=True) 230 | p123_markov = np.einsum('ij,jk->ijk',p12, pcond_23) 231 | 232 | # Chi^2 value 233 | #return utils.ntrapz( (p123 - p123_markov)**2, [dx, dx, dx] )/(np.var(p123.flatten()) + np.var(p123_markov.flatten())) 234 | return kl_divergence(p123, p123_markov, dx=[dx, dx, dx], tol=1e-6) 235 | 236 | 237 | 238 | ### FAST AUTOCORRELATION FUNCTION 239 | # From https://dfm.io/posts/autocorr/ 240 | 241 | def next_pow_two(n): 242 | i = 1 243 | while i < n: 244 | i = i << 1 245 | return i 246 | 247 | def autocorr_func_1d(x, norm=True): 248 | x = np.atleast_1d(x) 249 | if len(x.shape) != 1: 250 | raise ValueError("invalid dimensions for 1D autocorrelation function") 251 | n = next_pow_two(len(x)) 252 | 253 | # Compute the FFT and then (from that) the auto-correlation function 254 | f = np.fft.fft(x - np.mean(x), n=2*n) 255 | acf = np.fft.ifft(f * np.conjugate(f))[:len(x)].real 256 | acf /= 4*n 257 | 258 | # Optionally normalize 259 | if norm: 260 | acf /= acf[0] 261 | 262 | return acf -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Nonlinear stochastic modeling with Langevin regression 2 | ___________ 3 | __J. L. Callaham, J.-C. Loiseau, G. Rigas, and S. L. Brunton (2020)__ 4 | 5 | Code and details for the manuscript on data-driven stochastic modeling with Langevin regression. The repository should have everything needed to reproduce the results of the systems in the paper and demonstrate the basics of Langevin regression. The main ideas are: 6 | 7 | 1. Estimating drift and diffusion from data with Kramers-Moyal averaging 8 | 2. Correcting finite-time sampling effects with the adjoint Fokker-Planck equation 9 | 3. Enforcing consistency with the empirical PDF via the steady-state Fokker-Planck equation 10 | 11 | The full optimization problem combines all of these, along with a sparsification routine called SSR that was [previously proposed](https://arxiv.org/abs/1712.02432) to extend [SINDy](https://github.com/dynamicslab/pysindy) to stochastic systems. 12 | 13 | The repository several Python packages: 14 | 15 | * `utils.py`: A number of functions to do the main work of Langevin regression. There is code to compute the Kramers-Moyal average, the Langevin regression "cost function", a wrapper around the Nelder-Mead optimization, an implementation of SSR, etc. 16 | * `fpsolve.py`: A library of Fokker-Planck solvers (steady-state and adjoint). The details are in Appendix B of the paper. 17 | * `dwutils.py`: Several utility functions specifically for the notebook on the 1D double-well potential (see below). These do things like compute the dwell time distribution for the metastable states. 18 | 19 | To demonstrate, we include two notebooks for the simulated examples in the paper: 20 | 21 | ### 1. Pitchfork normal form driven by colored noise 22 | 23 | To explore the effects of correlated forcing and finite-time sampling, we look at the bistable normal form of a pitchfork bifurcation, driven by an Ornstein-Uhlenbeck process: 24 | $$ 25 | \dot{x} = \lambda x - \mu x^3 + \eta 26 | $$ 27 | $$ 28 | \dot{\eta} = - \alpha \eta + \sigma w(t), 29 | $$ 30 | where $w(t)$ is a white noise process (see paper for details). 31 | With coarse sampling rates $\tau \gg \alpha^{-1}$ and adjoint corrections, we show that a statistically consistent model can be identified with only white noise forcing. 32 | 33 | ### 2. One-dimensional particle in a double-well potential 34 | 35 | The true dynamics of the system are given by the second-order Langevin equation 36 | $$ 37 | \ddot{x} + \gamma \dot{x} + U'(x) = \sqrt{2 \gamma \kB T} w(t), 38 | $$ 39 | with double-well potential 40 | $$ 41 | U(x) = -\frac{\alpha}{2} x^2 + \frac{\beta}{4} x^4. 42 | $$ 43 | See the paper for nondimensionalization and further discussion. 44 | 45 | On short time scales, the dynamics of the state $x$ are smooth, since the direct white noise forcing is integrated twice. However, the macroscopic behavior displays metastable switching similar to the pitchfork normal form. In fact, for the weakly supercritical system we show in Appendix C that a stochastic normal form model can be derived analytically. 46 | 47 | Far from the bifurcation we can continue to model the macroscopic dynamics, but we have to resort to data-driven methods. 48 | This notebook demonstrates the use of Langevin regression to reduce the dynamics to a first-order system with consistent statistics over a wide range of parameters. 49 | 50 | ### 3. Turbulent axisymmetric wake 51 | 52 | For the third example we model a global integral quantity (the center of pressure) from experimental measurements of the wake behind an axisymmetric bluff body. We demonstrate that sparse stepwise regression (SSR) can clearly identify a model that balances accuracy and complexity from a library of candidate drift and diffusion functions. 53 | 54 | 55 | References 56 | ---------------------- 57 | - Jared L. Callaham, 58 | Jean-Christophe Loiseau, 59 | Georgios Rigas, 60 | and Steven L. Brunton 61 | *Nonlinear stochastic modeling with Langevin regression.* __Add arxiv reference__ 62 | 63 | - Steven L. Brunton, Joshua L. Proctor, and J. Nathan Kutz. 64 | *Discovering governing equations from data by sparse identification 65 | of nonlinear dynamical systems.* Proceedings of the National 66 | Academy of Sciences 113.15 (2016): 3932-3937. 67 | `[DOI] `__ 68 | 69 | - Lorenzo Boninsegna, Feliks Nüske, and Cecilia Clementi. *Sparse learning of stochastic dynamic equations.* Journal of Chemical Physics 148.24 (2018). `[DOI] __ 70 | 71 | - Complete bibliography in the manuscript -------------------------------------------------------------------------------- /__pycache__/dwutils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dynamicslab/langevin-regression/35468293c3b133a79562924375d80b5155d763fc/__pycache__/dwutils.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/fpsolve.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dynamicslab/langevin-regression/35468293c3b133a79562924375d80b5155d763fc/__pycache__/fpsolve.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dynamicslab/langevin-regression/35468293c3b133a79562924375d80b5155d763fc/__pycache__/utils.cpython-38.pyc -------------------------------------------------------------------------------- /data/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dynamicslab/langevin-regression/35468293c3b133a79562924375d80b5155d763fc/data/.DS_Store -------------------------------------------------------------------------------- /data/wake-cop.mat: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:010f3e8145a7015d018203e5c199102c7a6932d6bc545e4dc28f69d49026e602 3 | size 35640184 4 | -------------------------------------------------------------------------------- /dwutils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Utility functions for stochastic modeling of the double-well potential 3 | 4 | Jared Callaham (2020) 5 | """ 6 | 7 | import numpy as np 8 | import sympy 9 | from scipy.optimize import curve_fit 10 | 11 | import utils 12 | import fpsolve 13 | 14 | import warnings 15 | warnings.simplefilter(action='ignore', category=FutureWarning) 16 | 17 | import julia 18 | jl = julia.Julia() 19 | jl.include("../sim_doublewell.jl") 20 | 21 | def switched_states(X, thresh=1): 22 | # Decide whether the bistable states are "up" or "down" based on threshold 23 | 24 | N = len(X) 25 | state = np.zeros((N)) 26 | 27 | # Step forward until X is either positive or negative (if simulation is initialized at 0) 28 | idx = 0 29 | while abs(X[idx]) < thresh: 30 | idx += 1 31 | 32 | # Consider everything up to this point to be in the first well 33 | cur_state = np.sign(X[idx]) 34 | state[:idx] = cur_state 35 | 36 | while idx thresh: 40 | cur_state = -cur_state 41 | 42 | state[idx] = cur_state 43 | idx += 1 44 | 45 | return state 46 | 47 | 48 | def dwell_times(states, dt=1): 49 | # Given a vector of states (i.e. from switched_states() ), return dwell time in each state 50 | N = len(states) 51 | switch_times = [] # List of dwell times 52 | idx = 0 53 | last_switch = idx 54 | cur_state = states[idx] 55 | 56 | while idx < N: 57 | if states[idx] != cur_state: 58 | switch_times.append( dt*(idx-last_switch) ) 59 | last_switch = idx 60 | cur_state = states[idx] 61 | 62 | idx += 1 63 | 64 | return switch_times 65 | 66 | def dwell_stats(X, thresh, dt): 67 | 68 | state = switched_states(X, thresh=thresh) # Categorize into "up" or "down" 69 | switch_times = dwell_times(state, dt=dt) # Compute list of dwell times in each metastable state 70 | 71 | if len(switch_times) > 0: 72 | return np.mean(switch_times), np.std(switch_times)/np.sqrt(len(switch_times)) 73 | else: 74 | return np.nan, np.nan 75 | 76 | 77 | def fit_pdf(X, edges, p_hist, dt, p0=None): 78 | # Mean-square displacement 79 | fit_start, fit_stop = 0.1, 3 80 | n_lags = int(fit_stop/dt) 81 | tau = dt*np.arange(1, n_lags) 82 | 83 | # Lagged mean square displacement 84 | msd = np.zeros((len(tau))) 85 | for i in range(1, n_lags): 86 | msd[i-1] = np.mean( (X[i:]-X[:-i])**2) 87 | 88 | # Linear fit for radial displacement 89 | to_fit = np.nonzero( (tau > fit_start) * (tau < fit_stop) )[0] 90 | p_rad = np.polyfit(tau[to_fit], msd[to_fit], deg=1) 91 | a_pdf = 0.5*p_rad[0] 92 | 93 | # Fit PDF 94 | p_model = lambda x, C, a, b: C*np.exp(a*x**2 + b*x**4) 95 | centers = 0.5*(edges[1:]+edges[:-1]) 96 | if p0 is not None: 97 | popt, pcov = curve_fit(p_model, centers, p_hist, p0=p0) 98 | else: 99 | popt, pcov = curve_fit(p_model, centers, p_hist) 100 | 101 | # Separate parameters in model 102 | sigma_pdf = np.sqrt(2*a_pdf) 103 | lamb_pdf = popt[1]*sigma_pdf**2 104 | mu_pdf = popt[2]*2*sigma_pdf**2 105 | 106 | return lamb_pdf, mu_pdf, sigma_pdf 107 | 108 | 109 | def langevin_regression(X, edges, p_hist, dt, stride=200, kl_reg=0): 110 | """ 111 | Wrapper for full Langevin regression so we can loop over it to explore variations with distance from bifurcation 112 | """ 113 | centers = 0.5*(edges[1:]+edges[:-1]) 114 | N = len(centers) 115 | 116 | # Kramers-Moyal average 117 | tau = stride*dt 118 | f_KM, a_KM, f_err, a_err = KM_avg(X, bins, stride=stride, dt=dt) 119 | 120 | # Initialize libraries 121 | x = sympy.symbols('x') 122 | 123 | f_expr = np.array([x**i for i in [1, 3]]) # Polynomial library for drift 124 | s_expr = np.array([x**i for i in [0]]) # Polynomial library for diffusion 125 | 126 | lib_f = np.zeros([len(f_expr), N]) 127 | for k in range(len(f_expr)): 128 | lamb_expr = sympy.lambdify(x, f_expr[k]) 129 | lib_f[k] = lamb_expr(centers) 130 | 131 | lib_s = np.zeros([len(s_expr), N]) 132 | for k in range(len(s_expr)): 133 | lamb_expr = sympy.lambdify(x, s_expr[k]) 134 | lib_s[k] = lamb_expr(centers) 135 | 136 | # Initialize Xi with plain least-squares (just helpf the optimization a bit) 137 | Xi0 = np.zeros((len(f_expr) + len(s_expr))) 138 | mask = np.nonzero(np.isfinite(f_KM))[0] 139 | Xi0[:len(f_expr)] = np.linalg.lstsq( lib_f[:, mask].T, f_KM[mask], rcond=None)[0] 140 | Xi0[len(f_expr):] = np.linalg.lstsq( lib_s[:,mask].T, np.sqrt(2*a_KM[mask]), rcond=None)[0] 141 | 142 | # Parameter dictionary for optimization 143 | W = np.array((f_err.flatten(), a_err.flatten())) 144 | W[np.less(abs(W), 1e-12, where=np.isfinite(W))] = 1e6 # Set zero entries to large weights 145 | W[np.logical_not(np.isfinite(W))] = 1e6 # Set NaN entries to large numbers (small weights) 146 | W = 1/W # Invert error for weights 147 | W = W/np.nansum(W.flatten()) 148 | 149 | # Adjoint solver 150 | afp = fpsolve.AdjFP(centers) 151 | 152 | # Forward solver 153 | fp = fpsolve.SteadyFP(N, centers[1]-centers[0]) 154 | 155 | params = {"W": W, "f_KM": f_KM, "a_KM": a_KM, "Xi0": Xi0, 156 | "f_expr": f_expr, "s_expr": s_expr, 157 | "lib_f": lib_f.T, "lib_s": lib_s.T, "N": N, 158 | "kl_reg": kl_reg, 159 | "fp": fp, "afp": afp, "p_hist": p_hist, "tau": tau, 160 | "radial": False} 161 | 162 | # Tune KL regularization automatically 163 | Xi, _ = utils.AFP_opt(utils.cost, params) 164 | return Xi 165 | 166 | def model_eval(eps, sigma, N, kl_reg): 167 | """ 168 | Construct and evaluate all models of the double-well 169 | 1. Analytic normal form model 170 | 2. PDF fitting without Kramers-Moyal average 171 | 3. Full Langevin regression 172 | """ 173 | ### Generate data 174 | x_eq = np.sqrt(eps) # Equilibrium value 175 | 176 | edges = np.linspace(-2*x_eq, 2*x_eq, N+1) 177 | centers = 0.5*(edges[:-1]+edges[1:]) 178 | dx = centers[1]-centers[0] 179 | 180 | dt = 1e-2 181 | tmax = int(1e5) 182 | t, X = jl.run_sim(eps, sigma, dt, tmax) 183 | X, V = X[0, :], X[1, :] 184 | 185 | # PDF of states 186 | p_hist = np.histogram(X, edges, density=True)[0] 187 | 188 | # Dwell-time slope 189 | b, b_err = dwell_stats(X, x_eq, dt) 190 | print("\tData: ", b, b_err) 191 | 192 | ### 1. Normal form 193 | lamb1 = -1 + np.sqrt(1 + eps) 194 | lamb2 = -1 - np.sqrt(1 + eps) 195 | h = -lamb1/lamb2 196 | mu = -(1+h)**2*lamb1/eps 197 | 198 | _, phi1 = jl.run_nf(lamb1, mu, sigma/(2*np.sqrt(1+eps)), dt, tmax) 199 | X_nf = (1+h)*phi1[0, :] 200 | 201 | # Statistics 202 | p_nf = np.histogram(X_nf, edges, density=True)[0] 203 | b_nf, b_nf_err = dwell_stats(X_nf, x_eq, dt) 204 | print("\tNormal form: ", b_nf, b_nf_err) 205 | 206 | ### 2. PDF fit 207 | Xi = fit_pdf(X, edges, p_hist, dt, p0=[1, lamb1/sigma**2, mu/sigma**2]) 208 | #print(Xi) 209 | 210 | # Monte Carlo evaluation 211 | _, X_pdf = jl.run_nf(Xi[0], Xi[1], Xi[2], dt, tmax) 212 | X_pdf = X_pdf[0, :] 213 | 214 | # Statistics 215 | p_pdf = np.histogram(X_pdf, edges, density=True)[0] 216 | b_pdf, b_pdf_err = dwell_stats(X_pdf, x_eq, dt) 217 | print("\tPDF fit: ", b_pdf, b_pdf_err) 218 | 219 | ### 3. Langevin regression 220 | Xi = langevin_regression(X, edges, p_hist, dt, stride=200, kl_reg=kl_reg) 221 | #print(Xi) 222 | 223 | # Monte Carlo evaluation 224 | _, X_lr = jl.run_nf(Xi[0], Xi[1], Xi[2], dt, tmax) 225 | X_lr = X_lr[0, :] 226 | 227 | # Statistics 228 | p_lr = np.histogram(X_lr, edges, density=True)[0] 229 | b_lr, b_lr_err = dwell_stats(X_lr, x_eq, dt) 230 | print("\tLangevin regression: ", b_lr, b_lr_err) 231 | 232 | ### KL-divergence of all models against true data 233 | KL_nf = utils.kl_divergence(p_hist, p_nf, dx=dx, tol=1e-6) 234 | KL_pdf = utils.kl_divergence(p_hist, p_pdf, dx=dx, tol=1e-6) 235 | KL_lr = utils.kl_divergence(p_hist, p_lr, dx=dx, tol=1e-6) 236 | print("\tKL div: ", KL_nf, KL_pdf, KL_lr) 237 | 238 | return [b, b_nf, b_pdf, b_lr], [KL_nf, KL_pdf, KL_lr] -------------------------------------------------------------------------------- /fpsolve.py: -------------------------------------------------------------------------------- 1 | """ 2 | Package to solve Fokker-Planck equations 3 | 4 | * Steady-state Fourier-space solver for the PDF 5 | * Adjoint finite difference solver for first/second moments 6 | 7 | (Only 1D and 2D implemented so far) 8 | 9 | Jared Callaham (2020) 10 | """ 11 | 12 | import numpy as np 13 | from numpy.fft import fft, fftn, fftfreq, ifftn 14 | from scipy import linalg, sparse 15 | 16 | class SteadyFP: 17 | """ 18 | Solver object for steady-state Fokker-Planck equation 19 | 20 | Initializing this independently avoids having to re-initialize all of the indexing arrays 21 | for repeated loops with different drift and diffusion 22 | 23 | Jared Callaham (2020) 24 | """ 25 | 26 | def __init__(self, N, dx): 27 | """ 28 | ndim - number of dimensions 29 | N - array of ndim ints: grid resolution N[0] x N[1] x ... x N[ndim-1] 30 | dx - grid spacing (array of floats) 31 | """ 32 | 33 | if isinstance(N, int): 34 | self.ndim = 1 35 | else: 36 | self.ndim = len(N) 37 | 38 | self.N = N 39 | self.dx = dx 40 | 41 | # Set up indexing matrices for ndim=1, 2 42 | if self.ndim == 1: 43 | self.k = 2*np.pi*fftfreq(N, dx) 44 | self.idx = np.zeros((self.N, self.N), dtype=np.int32) 45 | for i in range(self.N): 46 | self.idx[i, :] = i-np.arange(N) 47 | 48 | elif self.ndim == 2: 49 | # Fourier frequencies 50 | self.k = [2*np.pi*fftfreq(N[i], dx[i]) for i in range(self.ndim)] 51 | self.idx = np.zeros((2, self.N[0], self.N[1], self.N[0], self.N[1]), dtype=np.int32) 52 | 53 | for m in range(N[0]): 54 | for n in range(N[1]): 55 | self.idx[0, m, n, :, :] = m-np.tile(np.arange(N[0]), [N[1], 1]).T 56 | self.idx[1, m, n, :, :] = n-np.tile(np.arange(N[1]), [N[0], 1]) 57 | 58 | else: 59 | print("WARNING: NOT IMPLEMENTED FOR HIGHER DIMENSIONS") 60 | 61 | self.A = None # Need to initialize with precompute_operator 62 | 63 | def precompute_operator(self, f, a): 64 | """ 65 | f - array of drift coefficients on domain (ndim x N[0] x N[1] x ... x N[ndim]) 66 | a - array of diffusion coefficients on domain (ndim x N[0] x N[1] x ... x N[ndim]) 67 | NOTE: To generalize to covariate noise, would need to add a dimension to a 68 | """ 69 | 70 | if self.ndim == 1: 71 | f_hat = self.dx*fftn(f) 72 | a_hat = self.dx*fftn(a) 73 | 74 | # Set up spectral projection operator 75 | self.A = np.einsum('i,ij->ij', -1j*self.k, f_hat[self.idx]) \ 76 | + np.einsum('i,ij->ij', -self.k**2, a_hat[self.idx]) 77 | 78 | if self.ndim == 2: 79 | # Initialize Fourier transformed coefficients 80 | f_hat = np.zeros(np.append([self.ndim], self.N), dtype=np.complex64) 81 | a_hat = np.zeros(f_hat.shape, dtype=np.complex64) 82 | for i in range(self.ndim): 83 | f_hat[i] = np.prod(self.dx)*fftn(f[i]) 84 | a_hat[i] = np.prod(self.dx)*fftn(a[i]) 85 | 86 | self.A = -1j*np.einsum('i,ijkl->ijkl', self.k[0], f_hat[0, self.idx[0], self.idx[1]]) \ 87 | -1j*np.einsum('j,ijkl->ijkl', self.k[1], f_hat[1, self.idx[0], self.idx[1]]) \ 88 | -np.einsum('i,ijkl->ijkl', self.k[0]**2, a_hat[0, self.idx[0], self.idx[1]]) \ 89 | -np.einsum('j,ijkl->ijkl', self.k[1]**2, a_hat[1, self.idx[0], self.idx[1]]) 90 | 91 | self.A = np.reshape(self.A, (np.prod(self.N), np.prod(self.N))) 92 | 93 | def solve(self, f, a): 94 | """ 95 | Solve Fokker-Planck equation from input drift coefficients 96 | """ 97 | self.precompute_operator(f, a) 98 | 99 | q_hat = np.linalg.lstsq(self.A[1:, 1:], -self.A[1:, 0], rcond=1e-6)[0] 100 | q_hat = np.append([1], q_hat) 101 | return np.real(ifftn( np.reshape(q_hat, self.N) ))/np.prod(self.dx) 102 | 103 | 104 | 105 | 106 | class AdjFP: 107 | """ 108 | Solver object for adjoint Fokker-Planck equation 109 | 110 | Jared Callaham (2020) 111 | """ 112 | 113 | 114 | # 1D derivative operators 115 | @staticmethod 116 | def derivs1d(x): 117 | N = len(x) 118 | dx = x[1]-x[0] 119 | one = np.ones((N)) 120 | 121 | # First derivative 122 | Dx = sparse.diags([one, -one], [1, -1], shape=(N, N)) 123 | Dx = sparse.lil_matrix(Dx) 124 | # Forward/backwards difference at boundaries 125 | Dx[0, :3] = [-3, 4, -1] 126 | Dx[-1, -3:] = [1, -4, 3] 127 | Dx = sparse.csr_matrix(Dx)/(2*dx) 128 | 129 | # Second derivative 130 | Dxx = sparse.diags([one, -2*one, one], [1, 0, -1], shape=(N, N)) 131 | Dxx = sparse.lil_matrix(Dxx) 132 | # Forwards/backwards differences (second-order accurate) 133 | Dxx[-1, -4:] = [1.25, -2.75, 1.75, -.25] 134 | Dxx[0, :4] = [-.25, 1.75, -2.75, 1.25] 135 | Dxx = sparse.csr_matrix(Dxx)/(dx**2) 136 | 137 | return Dx, Dxx 138 | 139 | @staticmethod 140 | def derivs2d(x, y): 141 | hx, hy = x[1]-x[0], y[1]-y[0] 142 | Nx, Ny = len(x), len(y) 143 | 144 | Dy = sparse.diags( [-1, 1], [-1, 1], shape=(Ny, Ny) ).toarray() 145 | 146 | # Second-order forward/backwards at boundaries 147 | Dy[0, :3] = np.array([-3, 4, -1]) 148 | Dy[-1, -3:] = np.array([1, -4, 3]) 149 | # Repeat for each x-location 150 | Dy = linalg.block_diag(*Dy.reshape(1, Ny, Ny).repeat(Nx,axis=0))/(2*hy) 151 | Dy = sparse.csr_matrix(Dy) 152 | 153 | Dx = sparse.diags( [-1, 1], [-Ny, Ny], shape=(Nx*Ny, Nx*Ny)).toarray() 154 | # Second-order forwards/backwards at boundaries 155 | for i in range(Ny): 156 | Dx[i, i] = -3 157 | Dx[i, Ny+i] = 4 158 | Dx[i, 2*Ny+i] = -1 159 | Dx[-(i+1), -(i+1)] = 3 160 | Dx[-(i+1), -(Ny+i+1)] = -4 161 | Dx[-(i+1), -(2*Ny+i+1)] = 1 162 | Dx = sparse.csr_matrix(Dx)/(2*hx) 163 | 164 | Dxx = sparse.csr_matrix(Dx @ Dx) 165 | Dyy = sparse.csr_matrix(Dy @ Dy) 166 | 167 | return Dx, Dy, Dxx, Dyy 168 | 169 | def __init__(self, x, ndim=1): 170 | """ 171 | x - uniform grid (array of floats) 172 | """ 173 | 174 | self.ndim = ndim 175 | 176 | if self.ndim == 1: 177 | self.N = [len(x)] 178 | self.dx = [x[1]-x[0]] 179 | self.x = [x] 180 | self.Dx, self.Dxx = AdjFP.derivs1d(x) 181 | self.precompute_operator = self.operator1d 182 | else: 183 | self.x = x 184 | self.N = [len(x[i]) for i in range(len(x))] 185 | self.dx = [x[i][1]-x[i][0] for i in range(len(x))] 186 | self.Dx, self.Dy, self.Dxx, self.Dyy = AdjFP.derivs2d(*x) 187 | self.precompute_operator = self.operator2d 188 | 189 | self.XX = np.meshgrid(*self.x, indexing='ij') 190 | self.precompute_moments() 191 | 192 | 193 | 194 | def precompute_moments(self): 195 | self.m1 = np.zeros([self.ndim, np.prod(self.N), np.prod(self.N)]) 196 | self.m2 = np.zeros([self.ndim, np.prod(self.N), np.prod(self.N)]) 197 | 198 | for d in range(self.ndim): 199 | for i in range(np.prod(self.N)): 200 | self.m1[d, i, :] = self.XX[d].flatten() - self.XX[d].flatten()[i] 201 | self.m2[d, i, :] = (self.XX[d].flatten() - self.XX[d].flatten()[i])**2 202 | 203 | 204 | def operator1d(self, f, a): 205 | self.L = sparse.diags(f) @ self.Dx + sparse.diags(a) @ self.Dxx 206 | 207 | 208 | def operator2d(self, f, a): 209 | self.L = sparse.diags(f[0]) @ self.Dx + sparse.diags(f[1]) @ self.Dy + \ 210 | sparse.diags(a[0]) @ self.Dxx + sparse.diags(a[1]) @ self.Dyy 211 | 212 | def solve(self, tau, d=0): 213 | if self.L is None: 214 | print("Need to initialize operator") 215 | return None 216 | 217 | L_tau = linalg.expm(self.L.todense()*tau) 218 | 219 | f_tau = np.einsum('ij,ij->i', L_tau, self.m1[d])/tau 220 | a_tau = np.einsum('ij,ij->i', L_tau, self.m2[d])/(2*tau) 221 | 222 | return f_tau, a_tau -------------------------------------------------------------------------------- /mean-field-model/README.md: -------------------------------------------------------------------------------- 1 | # Symmetry-breaking model 2 | ___________ 3 | 4 | Code and data to reproduce key results from "An empirical mean-field model of symmetry-breaking in a turbulent wake" by J. L. Callaham, G. Rigas, J.-C. Loiseau, and S. L. Brunton (2022) -------------------------------------------------------------------------------- /mean-field-model/data/.gitattributes: -------------------------------------------------------------------------------- 1 | wake_expt.mat filter=lfs diff=lfs merge=lfs -text 2 | -------------------------------------------------------------------------------- /mean-field-model/data/monte_carlo.mat: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:5031955dfe7de41b487d32ac82f462c2af6524f48d510d63e89d6920e65970b7 3 | size 64000248 4 | -------------------------------------------------------------------------------- /mean-field-model/data/order_parameters.mat: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:d51128dd513c07f419adafa4f8c7329edc1c7f0e4afe19028c289220812b05fc 3 | size 95040296 4 | -------------------------------------------------------------------------------- /mean-field-model/data/wake_expt.mat: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:951d1de076460b563336d4e43f3c721d2949782279359971d17aa12c3293c96d 3 | size 1077121304 4 | -------------------------------------------------------------------------------- /mean-field-model/fpsolve.py: -------------------------------------------------------------------------------- 1 | """ 2 | Package to solve Fokker-Planck equations 3 | 4 | * Steady-state Fourier-space solver for the PDF 5 | - Galerkin projection for 1D/2D 6 | - Arnoldi iteration for higher dimensions (slower but easier to scale) 7 | * Adjoint finite difference solver for first/second moments 8 | 9 | Jared Callaham (2020) 10 | """ 11 | 12 | import numpy as np 13 | from numpy.fft import fft, ifft, fftn, fftfreq, ifftn 14 | from scipy import linalg, sparse 15 | 16 | from scipy.sparse.linalg import LinearOperator, eigs 17 | import utils 18 | 19 | class SteadyFP: 20 | """ 21 | Solver object for steady-state Fokker-Planck equation 22 | 23 | Initializing this independently avoids having to re-initialize all of the indexing arrays 24 | for repeated loops with different drift and diffusion 25 | 26 | Jared Callaham (2020) 27 | """ 28 | 29 | def __init__(self, x, ndim=1, method=None): 30 | """ 31 | ndim - number of dimensions 32 | N - array of ndim ints: grid resolution N[0] x N[1] x ... x N[ndim-1] 33 | dx - grid spacing (array of floats) 34 | 35 | method - Galerkin or Arnoldi 36 | if not selected, defaults to Galerkin for ndim < 3, Arnoldi otherwise 37 | """ 38 | self.ndim = ndim 39 | 40 | if (method is None) and (ndim < 3): 41 | method = "galerkin" 42 | elif (method is None): 43 | method = "arnoldi" 44 | 45 | if self.ndim == 1: 46 | self.N = len(x) 47 | self.dx = x[1]-x[0] 48 | self.x = x 49 | self.k = 2*np.pi*fftfreq(self.N, self.dx) 50 | else: 51 | self.x = x 52 | self.N = [len(x[i]) for i in range(len(x))] 53 | self.dx = [x[i][1]-x[i][0] for i in range(len(x))] 54 | self.k = [2*np.pi*fftfreq(self.N[i], self.dx[i]) for i in range(self.ndim)] 55 | self.KK = np.meshgrid(*self.k, indexing='ij') 56 | 57 | if method=="arnoldi": 58 | self.solve = self.arnoldi_solve 59 | elif method=="galerkin": 60 | self.galerkin_init() 61 | self.solve = self.galerkin_solve 62 | else: 63 | print("Not implemented") 64 | 65 | def FP_operator(self, p, f, a): 66 | """ 67 | Linear Fokker-Planck operator with functions for f and a 68 | 69 | q = L*p 70 | """ 71 | 72 | p = np.reshape(p, self.N) 73 | q = np.sum(np.array([ ifft(-1j*self.KK[d]*fft(f[d]*p, axis=d) - self.KK[d]**2*fft(a[d]*p, axis=d), axis=d) 74 | for d in range(self.ndim) ]), axis=0) 75 | return q.flatten() 76 | 77 | def arnoldi_operator(self, f, a): 78 | self.L = LinearOperator((np.prod(self.N), np.prod(self.N)), 79 | matvec=lambda p: self.FP_operator(p, f, a)) 80 | 81 | def arnoldi_solve(self, f, a): 82 | """ 83 | Solve Fokker-Planck equation from input drift/diffusion coefficients 84 | 85 | Higher-dimensional version using Arnoldi iteration to estimate leading eigenvector 86 | """ 87 | self.arnoldi_operator(f, a) 88 | 89 | evals, evecs = eigs(self.L, k=1, which="LR") 90 | p_sol = np.reshape( np.real(evecs[:, 0]), self.N) 91 | p_sol /= utils.ntrapz(p_sol, self.dx) 92 | return p_sol 93 | 94 | 95 | def galerkin_init(self): 96 | # Set up indexing matrices for ndim=1, 2 97 | if self.ndim == 1: 98 | self.k = 2*np.pi*fftfreq(self.N, self.dx) 99 | self.idx = np.zeros((self.N, self.N), dtype=np.int32) 100 | for i in range(self.N): 101 | self.idx[i, :] = i-np.arange(self.N) 102 | 103 | elif self.ndim == 2: 104 | # Fourier frequencies 105 | self.k = [2*np.pi*fftfreq(self.N[i], self.dx[i]) for i in range(self.ndim)] 106 | self.idx = np.zeros((2, self.N[0], self.N[1], self.N[0], self.N[1]), dtype=np.int32) 107 | 108 | for m in range(self.N[0]): 109 | for n in range(self.N[1]): 110 | self.idx[0, m, n, :, :] = m-np.tile(np.arange(self.N[0]), [self.N[1], 1]).T 111 | self.idx[1, m, n, :, :] = n-np.tile(np.arange(self.N[1]), [self.N[0], 1]) 112 | 113 | else: 114 | print("WARNING: NOT IMPLEMENTED FOR HIGHER DIMENSIONS - USE ARNOLDI") 115 | 116 | def galerkin_operator(self, f, a): 117 | """ 118 | f - array of drift coefficients on domain (ndim x N[0] x N[1] x ... x N[ndim]) 119 | a - array of diffusion coefficients on domain (ndim x N[0] x N[1] x ... x N[ndim]) 120 | NOTE: To generalize to covariate noise, would need to add a dimension to a 121 | """ 122 | 123 | if self.ndim == 1: 124 | f_hat = self.dx*fftn(f) 125 | a_hat = self.dx*fftn(a) 126 | 127 | # Set up spectral projection operator 128 | self.L = np.einsum('i,ij->ij', -1j*self.k, f_hat[self.idx]) \ 129 | + np.einsum('i,ij->ij', -self.k**2, a_hat[self.idx]) 130 | 131 | """ 132 | # KEEP FOR REFERENCE: naive implementation with N^4 loops (~3:30 for N=64, ~14s for N=32) 133 | 134 | A = np.zeros((Nx, Ny, Nx, Ny), dtype=np.complex64) 135 | # First two loops are over projected variables (k') 136 | for m in range(Nx): 137 | for n in range(Ny): 138 | # Then loop over k 139 | for i in range(Nx): 140 | for j in range(Ny): 141 | A[m, n, i, j] = -1j*(kx[m]*f_hat[0, m-i, n-j] + ky[n]*f_hat[1, m-i, n-j]) 142 | A[m, n, i, j] -= (kx[m]**2*a_hat[0, m-i, n-j] + ky[n]**2*a_hat[1, m-i, n-j]) 143 | """ 144 | if self.ndim == 2: 145 | # Initialize Fourier transformed coefficients 146 | f_hat = np.zeros(np.append([self.ndim], self.N), dtype=np.complex64) 147 | a_hat = np.zeros(f_hat.shape, dtype=np.complex64) 148 | for i in range(self.ndim): 149 | f_hat[i] = np.prod(self.dx)*fftn(f[i]) 150 | a_hat[i] = np.prod(self.dx)*fftn(a[i]) 151 | 152 | self.L = -1j*np.einsum('i,ijkl->ijkl', self.k[0], f_hat[0, self.idx[0], self.idx[1]]) \ 153 | -1j*np.einsum('j,ijkl->ijkl', self.k[1], f_hat[1, self.idx[0], self.idx[1]]) \ 154 | -np.einsum('i,ijkl->ijkl', self.k[0]**2, a_hat[0, self.idx[0], self.idx[1]]) \ 155 | -np.einsum('j,ijkl->ijkl', self.k[1]**2, a_hat[1, self.idx[0], self.idx[1]]) 156 | 157 | self.L = np.reshape(self.L, (np.prod(self.N), np.prod(self.N))) 158 | 159 | 160 | def galerkin_solve(self, f, a): 161 | """ 162 | Solve Fokker-Planck equation from input drift coefficients 163 | 164 | Fourier-Galerkin projection for fast solving in 1 or 2 dimensions 165 | """ 166 | self.galerkin_operator(f, a) 167 | 168 | q_hat = np.linalg.lstsq(self.L[1:, 1:], -self.L[1:, 0], rcond=1e-6)[0] 169 | q_hat = np.append([1], q_hat) 170 | return np.real(ifftn( np.reshape(q_hat, self.N) ))/np.prod(self.dx) 171 | 172 | 173 | 174 | 175 | class AdjFP: 176 | """ 177 | Solver object for adjoint Fokker-Planck equation 178 | 179 | Jared Callaham (2020) 180 | """ 181 | 182 | @staticmethod 183 | def unit(N, idx): 184 | """ 185 | Construct an n-dimensional "unit matrix" 186 | where the only nonzero element is given by idx 187 | """ 188 | # Here idx should be a linear idx (output will be reshaped to matrix) 189 | e = np.zeros(np.prod(N)) 190 | e[idx] = 1 191 | return np.reshape(e, N) 192 | 193 | 194 | @staticmethod 195 | def deriv1(u, dx, axis=0, bc="extrapolate"): 196 | """ 197 | First derivative of u 198 | Extrapolated boundary conditions 199 | 200 | u = n-dimensional array 201 | dx = n-length delta-x values 202 | dim = dimension to differentiate 203 | bc - boundary conditions 204 | """ 205 | u = np.moveaxis(u, axis, 0) # Move so that the axis being differentiated is first 206 | du = np.zeros_like(u) 207 | 208 | du[1:-1] = (u[2:]-u[:-2]) # Centered 209 | 210 | if bc=="extrapolate": 211 | du[0] = -3*u[0]+4*u[1]-u[2] 212 | du[-1] = 3*u[-1]-4*u[-2]+u[-3] 213 | elif bc=="dirichlet": 214 | du[0] = np.ones_like(u[0]) 215 | du[-1] = np.ones_like(u[-1]) 216 | else: 217 | raise NotImplementedError 218 | 219 | u = np.moveaxis(u, 0, axis) # Move axis back 220 | du = np.moveaxis(du, 0, axis)/(2*dx[axis]) 221 | return du 222 | 223 | @staticmethod 224 | def deriv2(u, dx, axis=0, bc="extrapolate"): 225 | """ 226 | Second derivative of u 227 | Extrapolated boundary conditions 228 | 229 | u = n-dimensional array 230 | dx = n-length delta-x values 231 | dim = dimension to differentiate 232 | """ 233 | u = np.moveaxis(u, axis, 0) # Move so that the axis being differentiated is first 234 | du = np.zeros_like(u) 235 | 236 | du[1:-1] = (u[2:]-2*u[1:-1]+u[:-2]) # Centered 237 | 238 | if bc=="extrapolate": 239 | du[0] = 2*u[0]-5*u[1]+4*u[2]-u[3] 240 | du[-1] = 2*u[-1]-5*u[-2]+4*u[-3]-u[-4] 241 | elif bc=="dirichlet": 242 | du[0] = np.ones_like(u[0]) 243 | du[-1] = np.ones_like(u[-1]) 244 | else: 245 | raise NotImplementedError 246 | 247 | u = np.moveaxis(u, 0, axis) # Move axis back to original location 248 | du = np.moveaxis(du, 0, axis)/(dx[axis]**2) 249 | return du 250 | 251 | @staticmethod 252 | def fd_op(N, dx, deriv, axis, bc="extrapolate"): 253 | """ 254 | Compute sparse derivative operator by evaluating finite-differences on "unit matrices" 255 | deriv - finite-differencing function (deriv1 or deriv2) 256 | axis - axis along which to differentiate 257 | """ 258 | row_idx = [] 259 | col_idx = [] 260 | entries = [] 261 | 262 | for k in range(np.prod(N)): 263 | col = deriv(AdjFP.unit(N, k), dx, axis=axis, bc=bc).flatten() # "Unit vector" 264 | cur_idx = np.flatnonzero(col) 265 | row_idx = np.concatenate([row_idx, cur_idx]) 266 | col_idx = np.concatenate([col_idx, k*np.ones_like(cur_idx)]) # Current column for entries 267 | entries = np.concatenate([entries, col[cur_idx]]) 268 | 269 | D = sparse.coo_matrix((entries, (row_idx, col_idx)), shape=(np.prod(N), np.prod(N))) 270 | return sparse.csc_matrix(D) 271 | 272 | 273 | def __init__(self, x, ndim=1, method="step"): 274 | """ 275 | x - uniform grid (array of floats) 276 | """ 277 | self.ndim = ndim 278 | 279 | if self.ndim == 1: 280 | self.N = [len(x)] 281 | self.dx = [x[1]-x[0]] 282 | self.x = [x] 283 | else: 284 | self.x = x 285 | self.N = [len(x[i]) for i in range(len(x))] 286 | self.dx = [x[i][1]-x[i][0] for i in range(len(x))] 287 | 288 | # Precompute and store sparse finite-difference matrices 289 | self.precompute_matrices() 290 | 291 | # Compute grid 292 | self.XX = np.meshgrid(*self.x, indexing='ij') 293 | self.XX = [XX.flatten() for XX in self.XX] 294 | 295 | if method=="step": 296 | self.solve = self.step_solve # Euler time-stepping 297 | elif method=="exp": 298 | self.solve = self.exp_solve # Matrix exponential 299 | 300 | 301 | def precompute_matrices(self): 302 | self.Dx = [AdjFP.fd_op(self.N, self.dx, AdjFP.deriv1, axis) for axis in range(self.ndim)] 303 | self.Dxx = [AdjFP.fd_op(self.N, self.dx, AdjFP.deriv2, axis) for axis in range(self.ndim)] 304 | 305 | def precompute_operator(self, f, a): 306 | if self.ndim==1: 307 | f, a = [f], [a] 308 | self.L = np.sum([ sparse.diags(f[i]) @ self.Dx[i] + sparse.diags(a[i]) @ self.Dxx[i] 309 | for i in range(self.ndim)]) 310 | 311 | 312 | def exp_solve(self, tau, dt=None, d=0): 313 | """ 314 | Solve with matrix exponential 315 | """ 316 | if self.L is None: 317 | print("Need to initialize operator") 318 | return None 319 | 320 | L_tau = linalg.expm(self.L.todense()*tau) 321 | 322 | w1 = L_tau @ self.XX[d] 323 | w2 = L_tau @ self.XX[d]**2 324 | 325 | f_tau = (w1 - self.XX[d])/tau 326 | a_tau = (w2 - 2*self.XX[d]*w1 + self.XX[d]**2)/(2*tau) 327 | 328 | return f_tau, a_tau 329 | 330 | def step_solve(self, tau, dt, d=0): 331 | """ 332 | Solve with forward Euler time-stepping 333 | """ 334 | # Evolve observables (zero-centered moments) 335 | w1 = self.XX[d].copy() 336 | w2 = self.XX[d].copy()**2 337 | for i in range(int(tau//dt)): 338 | w1 += dt*(self.L @ w1) 339 | w2 += dt*(self.L @ w2) 340 | 341 | # Evaluate finite-time moments for drift/diffusion 342 | f_tau = (w1 - self.XX[d])/tau 343 | a_tau = (w2 - 2*self.XX[d]*w1 + self.XX[d]**2)/(2*tau) 344 | return f_tau, a_tau 345 | 346 | 347 | -------------------------------------------------------------------------------- /mean-field-model/lindy.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import sympy 3 | from time import time 4 | 5 | from scipy import sparse, linalg 6 | from scipy.optimize import minimize 7 | from numpy.linalg import lstsq 8 | 9 | import utils 10 | import fpsolve 11 | 12 | # freestream parameters for non-dimensionalization 13 | Uinf = 15 # freestream velocity [m/sec] 14 | Pinf = 0.5*1.2041*15**2 # rho_air = 1.2041 (20C) -> Pinf = 135.4612 15 | D = 196.5 # Diameter of the body [mm] 16 | 17 | # Sampling rate 18 | fs = 225 # Hz 19 | dt = 1/fs 20 | 21 | class LRModel(): 22 | def __init__(self, data, dim=1, dt=1/225): 23 | self.data = data 24 | self.dim = dim 25 | self.dt = dt 26 | 27 | def compute_KM(self, lag): 28 | self.f_KM = None 29 | self.a_KM = None 30 | self.W = None 31 | self.tau = lag*self.dt 32 | raise NotImplementedError 33 | 34 | def check_KM_initialized(self): 35 | return ( hasattr(self, 'f_KM') and hasattr(self, 'a_KM') and 36 | hasattr(self, 'W') and hasattr(self, 'tau') ) 37 | 38 | def compute_PDF(self): 39 | self.p_hist = None 40 | raise NotImplementedError 41 | 42 | def check_PDF_initialized(self): 43 | return hasattr(self, 'p_hist') 44 | 45 | def precompute_libraries(self, x_sym, f_expr, s_expr, x_eval): 46 | self.f_expr = f_expr 47 | self.s_expr = s_expr 48 | N = len(x_eval) 49 | self.lib_f = np.zeros([N, len(f_expr)], dtype=x_eval.dtype) 50 | for k in range(len(f_expr)): 51 | lamb_expr = sympy.lambdify(x_sym, f_expr[k]) 52 | self.lib_f[:, k] = lamb_expr(x_eval) 53 | 54 | self.lib_s = np.zeros([N, len(s_expr)]) 55 | for k in range(len(s_expr)): 56 | lamb_expr = sympy.lambdify(x_sym, s_expr[k]) 57 | self.lib_s[:, k] = lamb_expr(x_eval) 58 | 59 | def check_libs_initialized(self): 60 | return hasattr(self, 'lib_f') and hasattr(self, 'lib_s') 61 | 62 | def initial_coeffs(self): 63 | self.Xi0 = None 64 | raise NotImplementedError 65 | 66 | def check_coeffs_initialized(self): 67 | return hasattr(self, 'Xi0') 68 | 69 | def initialize_fpsolve(self, centers, dim=1): 70 | self.afp = fpsolve.AdjFP(centers, ndim=dim, method="exp") # Adjoint solver 71 | self.fp = fpsolve.SteadyFP(centers, ndim=dim, method="galerkin") # Forward solver 72 | 73 | def check_fpsolve_initialized(self): 74 | return hasattr(self, 'afp') and hasattr(self, 'fp') 75 | 76 | def initialized(self): 77 | return (self.check_KM_initialized() and 78 | self.check_PDF_initialized() and 79 | self.check_libs_initialized() and 80 | self.check_coeffs_initialized() and 81 | self.check_fpsolve_initialized() ) 82 | 83 | def cost(self, Xi, kl_reg): 84 | return NotImplementedError 85 | 86 | def fit(self, kl_reg=1e-3): 87 | assert self.initialized() 88 | print(f"KL Regularization: {kl_reg}") 89 | 90 | start_time = time() 91 | is_complex = np.iscomplex(self.Xi0[0]) 92 | 93 | if is_complex: 94 | Xi0 = np.concatenate((np.real(self.Xi0), np.imag(self.Xi0))) # Split vector in two for complex 95 | opt_fun = lambda Xi: self.cost(Xi[:len(Xi)//2] + 1j*Xi[len(Xi)//2:], kl_reg) 96 | 97 | else: 98 | Xi0 = self.Xi0 99 | opt_fun = lambda Xi: self.cost(Xi, kl_reg) 100 | 101 | res = minimize(opt_fun, Xi0, method='nelder-mead', 102 | options={'disp': False, 'maxfev':int(1e4)}) 103 | print('%%%% Optimization time: {0} seconds, Cost: {1} %%%%'.format(time() - start_time, res.fun) ) 104 | 105 | # Return coefficients and cost function 106 | if is_complex: 107 | return res.x[:len(res.x)//2] + 1j*res.x[len(res.x)//2:], res.fun 108 | else: 109 | return res.x, res.fun 110 | 111 | 112 | class A_model(LRModel): 113 | def __init__(self, data, dt=1/225, nbins=24, domain_width=2, lag=200): 114 | super().__init__(data, dim=2, dt=dt) 115 | self.data = data 116 | self.nbins = nbins 117 | self.domain_width = domain_width 118 | 119 | self.compute_PDF() 120 | self.compute_KM(lag) 121 | self.precompute_libraries() 122 | self.initial_coeffs() 123 | self.initialize_fpsolve() 124 | 125 | def compute_PDF(self): 126 | ######################################################### 127 | ### PROBABILITY DENSITIES 128 | ######################################################### 129 | A = self.data 130 | self.bins = np.linspace(-self.domain_width, self.domain_width, self.nbins+1) 131 | dx = self.bins[1]-self.bins[0] 132 | self.centers1d = (self.bins[:-1]+self.bins[1:])/2 133 | 134 | YY, XX = np.meshgrid(self.centers1d, self.centers1d) 135 | self.centers2d = [XX, YY] 136 | 137 | # Histogram: [real, imag] 138 | self.p_hist, _, _ = np.histogram2d(np.real(A), np.imag(A), bins=[self.bins, self.bins], density=True) 139 | 140 | 141 | def compute_KM(self, lag): 142 | self.tau = lag*self.dt 143 | A = self.data 144 | 145 | dA = (A[lag:] - A[:-lag])/self.tau # Cartesian step (finite-difference derivative estimate) 146 | dA2 = self.tau*(np.real(dA)**2 + 1j*np.imag(dA)**2) # Multivariate variance (assuming diagonal) 147 | 148 | N = len(self.bins)-1 149 | f_KM = np.zeros((N, N), dtype=np.complex64) 150 | f_err = np.zeros(f_KM.shape) 151 | a_KM = np.zeros((f_KM.shape), dtype=np.complex64) 152 | a_err = np.zeros((f_KM.shape)) 153 | 154 | for i in range(N): 155 | for j in range(N): 156 | # Find where signal falls into this bin 157 | mask = (np.real(A[:-lag]) > self.bins[i]) * (np.real(A[:-lag]) < self.bins[i+1]) * \ 158 | (np.imag(A[:-lag]) > self.bins[j]) * (np.imag(A[:-lag]) < self.bins[j+1]) 159 | 160 | mask_idx = np.nonzero(mask)[0] 161 | nmask = len(mask_idx) 162 | 163 | if nmask > 0: 164 | # Conditional mean 165 | f_KM[i, j] = np.mean(dA[mask_idx]) # Conditional average 166 | f_err[i, j] = np.std(abs(dA[mask_idx]))/np.sqrt(nmask) 167 | 168 | # Conditional variance (assumes diagonal forcing) 169 | a_KM[i, j] = 0.5*np.mean(dA2[mask_idx]) # Conditional average 170 | a_err[i, j] = np.std(dA2[mask_idx])/np.sqrt(nmask) 171 | 172 | else: 173 | f_KM[i, j] = np.nan 174 | f_err[i, j] = np.nan 175 | a_KM[i, j] = np.nan 176 | a_err[i, j] = np.nan 177 | 178 | self.f_KM = np.array([np.real(f_KM), np.imag(f_KM)]) 179 | self.a_KM = np.array([np.real(a_KM), np.imag(a_KM)]) 180 | f_err = np.array([np.real(f_err), np.imag(f_err)]) 181 | a_err = np.array([np.real(a_err), np.imag(a_err)]) 182 | 183 | self.W = utils.KM_weights(f_err, a_err) 184 | 185 | def precompute_libraries(self): 186 | # Initialize sympy expression 187 | z = sympy.symbols('z') # real, imag 188 | 189 | # Lists of candidate functions 190 | self.f_expr = np.array([z, z*abs(z)**2]) 191 | self.s_expr = np.array([z**0, abs(z)**2]) 192 | 193 | ZZ = self.centers2d[0] + 1j*self.centers2d[1] 194 | LRModel.precompute_libraries(self, z, self.f_expr, self.s_expr, ZZ.flatten()) 195 | 196 | 197 | def initial_coeffs(self): 198 | self.Xi0 = np.zeros((len(self.f_expr) + len(self.s_expr)), dtype=np.complex64) 199 | lhs = (self.f_KM[0, :, :] + 1j*self.f_KM[1, :, :]).flatten() 200 | mask = np.nonzero(np.isfinite(lhs))[0] 201 | self.Xi0[:len(self.f_expr)] = lstsq( self.lib_f[mask, :], lhs[mask], rcond=None)[0] 202 | 203 | lhs = np.sqrt(2*self.a_KM[0, :, :].flatten()) + 1j*np.sqrt(2*self.a_KM[1, :, :].flatten()) 204 | self.Xi0[len(self.f_expr):] = lstsq( self.lib_s[mask, :], lhs[mask], rcond=None)[0] 205 | 206 | def initialize_fpsolve(self): 207 | LRModel.initialize_fpsolve(self, centers=[self.centers1d, self.centers1d], dim=2) 208 | 209 | def cost(self, Xi, kl_reg): 210 | r"""Least-squares cost function for optimization""" 211 | f_KM, a_KM = self.f_KM[0].flatten(), self.a_KM[0].flatten() 212 | 213 | Xi_f = Xi[:self.lib_f.shape[1]] 214 | Xi_s = Xi[self.lib_f.shape[1]:] 215 | 216 | f_vals = self.lib_f @ Xi_f 217 | s_vals = self.lib_s @ Xi_s 218 | a_vals = 0.5*( np.real(s_vals)**2 + 1j*(np.imag(s_vals))**2 ) 219 | 220 | # Solve adjoint Fokker-Planck equation 221 | self.afp.precompute_operator([np.real(f_vals), np.imag(f_vals)], 222 | [np.real(a_vals), np.imag(a_vals)]) 223 | f_tau, a_tau = self.afp.solve(self.tau, d=0) # Assumes real/imag symmetry 224 | 225 | mask = np.nonzero(np.isfinite(f_KM))[0] 226 | V = np.sum(self.W[0, mask]*abs(f_tau[mask] - f_KM[mask])**2) \ 227 | + np.sum(self.W[1, mask]*abs(a_tau[mask] - a_KM[mask])**2) 228 | 229 | if kl_reg > 0: 230 | p_est = self.fp.solve( 231 | [np.reshape(np.real(f_vals), self.fp.N), np.reshape(np.imag(f_vals), self.fp.N)], 232 | [np.reshape(np.real(a_vals), self.fp.N), np.reshape(np.imag(a_vals), self.fp.N)] 233 | ) 234 | 235 | kl = utils.kl_divergence(self.p_hist, p_est, dx=self.fp.dx, tol=1e-6) 236 | kl = max(0, kl) 237 | V += kl_reg*kl 238 | 239 | if not np.isfinite(V): 240 | print('Error in cost function') 241 | print(Xi) 242 | print(f) 243 | print(a) 244 | return None 245 | 246 | return V 247 | 248 | 249 | class B_model(LRModel): 250 | def __init__(self, data, dt=1/225, nbins=40, domain_width=4, lag=200): 251 | super().__init__(data, dim=1, dt=dt) 252 | self.data = data 253 | self.nbins = nbins 254 | self.domain_width = domain_width 255 | 256 | self.compute_PDF() 257 | self.compute_KM(lag) 258 | self.precompute_libraries() 259 | self.initial_coeffs() 260 | self.initialize_fpsolve() 261 | 262 | def compute_PDF(self): 263 | self.bins = np.linspace(-self.domain_width, self.domain_width, self.nbins+1) 264 | dx = self.bins[1]-self.bins[0] 265 | self.centers1d = (self.bins[:-1]+self.bins[1:])/2 266 | 267 | # Histogram: [real, imag] 268 | self.p_hist = np.histogram(self.data, bins=self.bins, density=True)[0] 269 | 270 | 271 | def compute_KM(self, lag): 272 | self.tau = lag*self.dt 273 | B = self.data 274 | 275 | N = len(self.bins)-1 276 | self.f_KM = np.zeros((N)) 277 | f_err = np.zeros(self.f_KM.shape) 278 | self.a_KM = np.zeros((self.f_KM.shape)) 279 | a_err = np.zeros((self.f_KM.shape)) 280 | 281 | dB = np.real(B[lag:] - B[:-lag])/self.tau # Step (finite-difference derivative estimate) 282 | dB2 = self.tau*dB**2 # Variance 283 | 284 | for i in range(N): 285 | # Find where signal falls into this bin 286 | mask = (B[:-lag] > self.bins[i]) * (B[:-lag] < self.bins[i+1]) 287 | mask_idx = np.nonzero(mask)[0] 288 | 289 | if len(mask_idx) > 0: 290 | # Conditional mean 291 | self.f_KM[i] = np.mean(dB[mask_idx]) # Conditional average 292 | f_err[i] = np.std(dB[mask_idx])/np.sqrt(len(mask_idx)) 293 | 294 | # Conditional variance 295 | self.a_KM[i] = 0.5*np.mean(dB2[mask_idx]) # Conditional average 296 | a_err[i] = np.std(dB2[mask_idx])/np.sqrt(len(mask_idx)) 297 | 298 | else: 299 | self.f_KM[i] = np.nan 300 | f_err[i] = np.nan 301 | self.a_KM[i] = np.nan 302 | a_err[i] = np.nan 303 | 304 | self.W = utils.KM_weights(f_err, a_err) 305 | 306 | def precompute_libraries(self): 307 | # Initialize sympy expression 308 | z = sympy.symbols('z') 309 | 310 | # Lists of candidate functions 311 | self.f_expr = np.array([z]) 312 | self.s_expr = np.array([z**0, z**2]) 313 | 314 | ZZ = self.centers1d 315 | LRModel.precompute_libraries(self, z, self.f_expr, self.s_expr, ZZ) 316 | 317 | def initial_coeffs(self): 318 | self.Xi0 = np.zeros((len(self.f_expr) + len(self.s_expr))) 319 | lhs = self.f_KM 320 | mask = np.nonzero(np.isfinite(lhs))[0] 321 | self.Xi0[:len(self.f_expr)] = lstsq( self.lib_f[mask, :], lhs[mask], rcond=None)[0] 322 | 323 | lhs = np.sqrt(2*self.a_KM) 324 | self.Xi0[len(self.f_expr):] = lstsq( self.lib_s[mask, :], lhs[mask], rcond=None)[0] 325 | 326 | def initialize_fpsolve(self): 327 | LRModel.initialize_fpsolve(self, centers=self.centers1d, dim=1) 328 | 329 | def cost(self, Xi, kl_reg): 330 | f_KM, a_KM = self.f_KM, self.a_KM 331 | 332 | Xi_f = Xi[:self.lib_f.shape[1]] 333 | Xi_s = Xi[self.lib_f.shape[1]:] 334 | 335 | f_vals = self.lib_f @ Xi_f 336 | s_vals = self.lib_s @ Xi_s 337 | a_vals = 0.5*s_vals**2 338 | 339 | # Solve adjoint Fokker-Planck equation 340 | self.afp.precompute_operator(f_vals, a_vals) 341 | f_tau, a_tau = self.afp.solve(self.tau, d=0) 342 | 343 | mask = np.nonzero(np.isfinite(f_KM))[0] 344 | V = np.sum(self.W[0, mask]*abs(f_tau[mask] - f_KM[mask])**2) \ 345 | + np.sum(self.W[1, mask]*abs(a_tau[mask] - a_KM[mask])**2) 346 | 347 | if kl_reg > 0: 348 | p_est = self.fp.solve(f_vals, a_vals) 349 | 350 | kl = utils.kl_divergence(self.p_hist, p_est, dx=self.fp.dx, tol=1e-6) 351 | kl = max(0, kl) 352 | V += kl_reg*kl 353 | 354 | return V 355 | 356 | 357 | 358 | class cop_model(B_model): 359 | 360 | def compute_PDF(self): 361 | self.bins = np.linspace(0, self.domain_width, self.nbins+1) 362 | dx = self.bins[1]-self.bins[0] 363 | self.centers1d = (self.bins[:-1]+self.bins[1:])/2 364 | 365 | # Histogram: [real, imag] 366 | self.p_hist = np.histogram(self.data, bins=self.bins, density=True)[0] 367 | 368 | def precompute_libraries(self): 369 | # Initialize sympy expression 370 | z = sympy.symbols('z') 371 | 372 | # Lists of candidate functions 373 | self.f_expr = np.array([z**i for i in [1, 3]]) # Polynomial library for drift 374 | self.s_expr = np.array([z**i for i in [0, 2]]) # Polynomial library for diffusion 375 | 376 | ZZ = self.centers1d 377 | LRModel.precompute_libraries(self, z, self.f_expr, self.s_expr, ZZ) 378 | 379 | def cost(self, Xi, kl_reg): 380 | f_KM, a_KM = self.f_KM, self.a_KM 381 | 382 | Xi_f = Xi[:self.lib_f.shape[1]] 383 | Xi_s = Xi[self.lib_f.shape[1]:] 384 | 385 | f_vals = self.lib_f @ Xi_f 386 | 387 | s_vals = self.lib_s @ Xi_s 388 | a_vals = 0.5*s_vals**2 389 | f_vals += a_vals/self.centers1d # Diffusion-induced drift from polar change of variables 390 | 391 | # Solve adjoint Fokker-Planck equation 392 | self.afp.precompute_operator(f_vals, a_vals) 393 | f_tau, a_tau = self.afp.solve(self.tau, d=0) 394 | 395 | mask = np.nonzero(np.isfinite(f_KM))[0] 396 | V = np.sum(self.W[0, mask]*abs(f_tau[mask] - f_KM[mask])**2) \ 397 | + np.sum(self.W[1, mask]*abs(a_tau[mask] - a_KM[mask])**2) 398 | 399 | if kl_reg > 0: 400 | p_est = self.fp.solve(f_vals, a_vals) 401 | 402 | kl = utils.kl_divergence(self.p_hist, p_est, dx=self.fp.dx, tol=1e-6) 403 | kl = max(0, kl) 404 | V += kl_reg*kl 405 | 406 | return V -------------------------------------------------------------------------------- /mean-field-model/meanfield.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy.fft import fft, ifft, fftshift, fftfreq 3 | 4 | ### GLOBAL VARIABLES 5 | Uinf = 15 # freestream velocity [m/sec] 6 | Pinf = 0.5*1.2041*15**2 # rho_air = 1.2041 (20C) -> Pinf = 135.4612 7 | D = 196.5 # Diameter of the body [mm] 8 | 9 | # location of the 64 pressure taps 10 | rESP=np.linspace(11,88,8)/D # non-dimensional with D 11 | thetaESP= np.array([270, 315, 0, 45, 90, 135, 180, 225]) #[degrees] 12 | nr = len(rESP) 13 | nth = len(thetaESP) 14 | 15 | # Sampling rate 16 | fs = 225 # Hz 17 | dt = 1/fs 18 | 19 | 20 | # Mass matrix 21 | dth = 2*np.pi/8 22 | dr = rESP[1]-rESP[0] 23 | w = (rESP*dr) 24 | W = np.diag(w**0.5) # Sqrt of mass matrix 25 | W_inv = np.linalg.inv(W) 26 | ip = lambda p, q: (W @ p).T @ (W @ q).conj() 27 | 28 | # For full-field (nr*nth) 29 | W_full = np.repeat(rESP*dr*dth,8) 30 | W_full = np.diag(W_full**0.5) 31 | 32 | 33 | def order_parameter(q): 34 | """ 35 | Derive order parameter from integrated m=1 amplitudre 36 | q = [nr x nth x T] 37 | """ 38 | qhat = fft(q, axis=1, norm='ortho') 39 | A = np.trapz( rESP * qhat[:, 1, :].T, dx=dr, axis=1) 40 | return A 41 | 42 | def phase_align(q, A, rot_frame=False): 43 | """ 44 | Phase-align with order parameter 45 | q = [nr x nth x T] 46 | """ 47 | qhat = fft(q, axis=1, norm='ortho') 48 | m = fftfreq(nth, d=1/nth) # Azimuthal wavenumber 49 | qhat_rot = np.zeros_like(qhat) 50 | phi = np.angle(A) 51 | for k in range(nth): 52 | if rot_frame: 53 | # Align with rotating frame 54 | qhat_rot[:, k, :] = qhat[:, k, :]*np.exp(-1j*np.sign(m[k])*phi) 55 | else: 56 | # Full phase alignment 57 | qhat_rot[:, k, :] = qhat[:, k, :]*np.exp(-1j*m[k]*phi) 58 | return np.real(ifft(qhat_rot, axis=1, norm='ortho')) 59 | 60 | 61 | def cond_avg(q, A, edges): 62 | qd = np.zeros([nr, nth, len(edges)-1]) 63 | qd_err = np.zeros_like(qd) 64 | for i in range(len(edges)-1): 65 | mask = np.nonzero((abs(A)>edges[i]) * (abs(A)= len(f_expr)] - len(f_expr) 107 | print(f_active) 108 | print(s_active) 109 | 110 | print(f_expr[f_active], s_expr[s_active]) 111 | params['f_expr'] = f_expr[f_active] 112 | params['s_expr'] = s_expr[s_active] 113 | params['lib_f'] = lib_f[:, f_active] 114 | params['lib_s'] = lib_s[:, s_active] 115 | params['Xi0'] = Xi0[tmp_active] 116 | 117 | # Ensure that there is at least one drift and diffusion term left 118 | if len(s_active) > 0 and len(f_active) > 0: 119 | tmp_Xi, tmp_V = opt_fun(params) 120 | 121 | # Keep minimum cost 122 | if tmp_V < V[k]: 123 | # Ensure that there is at least one drift and diffusion term left 124 | #if (IS_DRIFT and len(f_active)>1) or (not IS_DRIFT and len(a_active)>1): 125 | min_idx = j 126 | V[k] = tmp_V 127 | min_Xi = tmp_Xi 128 | 129 | print("Cost: {0}".format(V[k])) 130 | # Delete least important term 131 | active = np.delete(active, min_idx) # Remove inactive index 132 | Xi0[active] = min_Xi # Re-initialize with best results from previous 133 | Xi[active, k] = min_Xi 134 | print(Xi[:, k]) 135 | 136 | return Xi, V 137 | 138 | 139 | # 1D Markov test 140 | def markov_test(X, lag, N=32, L=2): 141 | # Lagged time series 142 | X1 = X[:-2*lag:lag] 143 | X2 = X[lag:-lag:lag] 144 | X3 = X[2*lag::lag] 145 | 146 | # Two-time joint pdfs 147 | bins = np.linspace(-L, L, N+1) 148 | dx = bins[1]-bins[0] 149 | p12, _, _ = np.histogram2d(X1, X2, bins=[bins, bins], density=True) 150 | p23, _, _ = np.histogram2d(X2, X3, bins=[bins, bins], density=True) 151 | p2, _ = np.histogram(X2, bins=bins, density=True) 152 | p2[p2<1e-4] = 1e-4 153 | 154 | # Conditional PDF (Markov assumption) 155 | pcond_23 = p23.copy() 156 | for j in range(pcond_23.shape[1]): 157 | pcond_23[:, j] = pcond_23[:, j]/p2 158 | 159 | # Three-time PDFs 160 | p123, _ = np.histogramdd(np.array([X1, X2, X3]).T, bins=np.array([bins, bins, bins]), density=True) 161 | p123_markov = np.einsum('ij,jk->ijk',p12, pcond_23) 162 | 163 | # Chi^2 value 164 | #return utils.ntrapz( (p123 - p123_markov)**2, [dx, dx, dx] )/(np.var(p123.flatten()) + np.var(p123_markov.flatten())) 165 | return kl_divergence(p123, p123_markov, dx=[dx, dx, dx], tol=1e-6) 166 | 167 | 168 | 169 | ### FAST AUTOCORRELATION FUNCTION 170 | # From https://dfm.io/posts/autocorr/ 171 | 172 | def next_pow_two(n): 173 | i = 1 174 | while i < n: 175 | i = i << 1 176 | return i 177 | 178 | def autocorr_func_1d(x, norm=True): 179 | x = np.atleast_1d(x) 180 | if len(x.shape) != 1: 181 | raise ValueError("invalid dimensions for 1D autocorrelation function") 182 | n = next_pow_two(len(x)) 183 | 184 | # Compute the FFT and then (from that) the auto-correlation function 185 | f = np.fft.fft(x - np.mean(x), n=2*n) 186 | acf = np.fft.ifft(f * np.conjugate(f))[:len(x)].real 187 | acf /= 4*n 188 | 189 | # Optionally normalize 190 | if norm: 191 | acf /= acf[0] 192 | 193 | return acf 194 | 195 | 196 | def calc_cop(p): 197 | """ 198 | # Compute center of pressure (adapted from George's MATLAB code) 199 | Reads in a snapshot p: (azimuthal) x (radial) 200 | 201 | % Rxi = 1/Sum[Pi*dAi] * Sum [Pi*xi*dAi] 202 | % Ryi = 1/Sum[Pi*dAi] * Sum [Pi*yi*dAi] 203 | 204 | % xi = ri*cos(thetai) 205 | % yi = ri*sin(thetai) 206 | % dAi = ri*dr*dtheta 207 | """ 208 | # location of the 64 pressure taps 209 | D = 196.5 210 | rESP=np.linspace(11,88,8)/D # non-dimensional with D 211 | thetaESP= np.array([270, 315, 0, 45, 90, 135, 180, 225]) #[degrees] 212 | 213 | p = p.T # nr x nth 214 | nr, nth = p.shape 215 | 216 | # Find ri , dAi 217 | dAi = np.zeros(nr) 218 | ri = np.zeros(nr) 219 | dr = rESP[1]-rESP[0] 220 | dth = 2*np.pi/8 221 | for i in range(1, len(rESP)): 222 | ri[i] = 0.5*(rESP[i]+rESP[i-1]) 223 | dAi[i] = ri[i]*dr*dth 224 | 225 | # Construct P, xi,yi dAi 226 | Ptrapz = 0.5*(p[:,:-1]+p[:,1:] ) # 8 x 7 227 | dAi = dAi[1:] # 1x7 228 | ri = ri[1:] # 1x7 229 | 230 | xi= np.outer( np.cos(thetaESP*np.pi/180), ri) # 8 x 7 231 | yi= np.outer( np.sin(thetaESP*np.pi/180), ri) # 8 x 7 232 | 233 | dA = np.repeat(dAi[None, :], 8, axis=0) 234 | 235 | SumPA = sum(sum(Ptrapz*dA)) 236 | 237 | Rxi = 1/SumPA*sum(sum(Ptrapz*xi*dA)) 238 | Ryi = 1/SumPA*sum(sum(Ptrapz*yi*dA)) 239 | 240 | return (Rxi, Ryi) 241 | 242 | 243 | def load_new_expt(file): 244 | """Rearranges order of arrays to be as expected in python""" 245 | import scipy.io as sio 246 | data = sio.loadmat(file) 247 | P = np.concatenate(data['P'][0], axis=0).T 248 | P2 = P.copy() 249 | 250 | P2[24:32] = P[32:40] 251 | P2[32:40] = P[48:56] 252 | P2[40:48] = P[24:32] 253 | P2[48:56] = P[40:48] 254 | 255 | P2 = np.reshape(P2, [8, 8, -1], order='F') # azimuthal x radius 256 | P2 = np.flip(P2, axis=0) 257 | return P2 -------------------------------------------------------------------------------- /sim_doublewell.jl: -------------------------------------------------------------------------------- 1 | """ 2 | Julia code to simulate 1D particle in a double-well potential 3 | 4 | Jared Callaham (2020) 5 | """ 6 | using DifferentialEquations 7 | 8 | # Noisy double well potential 9 | 10 | ### Physical dynamics 11 | function f_fn(dx, x, p, t) 12 | ϵ, σ = p 13 | dx[1] = x[2] 14 | dx[2] = -2*x[2] + ϵ*x[1] - x[1]^3 15 | end 16 | 17 | function σ_fn(dx, x, p, t) 18 | ϵ, σ = p 19 | dx[1] = 0 20 | dx[2] = σ 21 | end 22 | 23 | ### Normal form dynamics (pitchfork bifurcation) 24 | function fn_fn(dx, x, p, t) 25 | λ, μ, σ = p 26 | dx[1] = λ*x[1] + μ*x[1]^3 27 | end 28 | 29 | function σn_fn(dx, x, p, t) 30 | λ, μ, σ = p 31 | dx[1] = σ 32 | end 33 | 34 | function run_nf(ϵ, μ, σ, dt, tmax) 35 | prob = SDEProblem(fn_fn, σn_fn, [0.0] , (0.0, tmax), [ϵ, μ, σ]) 36 | sol = solve(prob, SRIW1(), dt=dt, adaptive=false); 37 | return sol.t, sol[:, :] 38 | end 39 | 40 | function run_sim(η, σ, dt, tmax) 41 | prob = SDEProblem(f_fn,σ_fn, [0.0, 0.0], (0.0, tmax), [η, σ]) 42 | sol = solve(prob, SRIW1(), dt=dt, adaptive=false) 43 | return sol.t, sol[:, :] 44 | end -------------------------------------------------------------------------------- /sim_pitchfork.jl: -------------------------------------------------------------------------------- 1 | """ 2 | Julia code to simulate pitchfork bifurcation normal form forced by colored noise 3 | 4 | Jared Callaham (2020) 5 | """ 6 | using DifferentialEquations 7 | using MAT 8 | 9 | # Drift component 10 | μ = 1 # Fixed points at +/- 1 11 | λ = 1 12 | 13 | r0 = sqrt(λ/μ) # Equilibrium point 14 | α = Int(1e2) # Inverse autocorrelation time of the forcing (smaller than lambda for scale separation) 15 | σ = α*0.5 16 | 17 | f_ex(x) = λ*x - μ*x^3 18 | 19 | # Simulate SDE 20 | function f_fn(dx, x, p, t) 21 | dx[1] = f_ex(x[1]) + x[2] 22 | dx[2] = -α*x[2] 23 | end 24 | 25 | function σ_fn(dx, x, p, t) 26 | dx[1] = 0 27 | dx[2] = σ 28 | end 29 | 30 | dt = 0.001 31 | x0 = [0.0, 0.0] 32 | tspan = (0.0, Int(1e4)) 33 | prob = SDEProblem(f_fn, σ_fn, x0, tspan) 34 | 35 | # EM() or SRIW1() 36 | sol = solve(prob, EM(), dt=dt); 37 | t = sol.t 38 | X = sol[1, :] 39 | 40 | matwrite("./data/pitchfork.mat", Dict( 41 | "X" => X, 42 | "dt" => dt, 43 | "lamb" => λ, 44 | "mu" => μ 45 | ); compress = true) -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Utility functions for Langevin regression 3 | 4 | Jared Callaham (2020) 5 | """ 6 | 7 | import numpy as np 8 | from time import time 9 | from scipy.optimize import minimize 10 | 11 | # Return a single expression from a list of expressions and coefficients 12 | # Note this will give a SymPy expression and not a function 13 | def sindy_model(Xi, expr_list): 14 | return sum([Xi[i]*expr_list[i] for i in range(len(expr_list))]) 15 | 16 | 17 | def ntrapz(I, dx): 18 | if isinstance(dx, int) or isinstance(dx, float) or len(dx)==1: 19 | return np.trapz(I, dx=dx, axis=0) 20 | else: 21 | return np.trapz( ntrapz(I, dx[1:]), dx=dx[0]) 22 | 23 | 24 | def kl_divergence(p_in, q_in, dx=1, tol=None): 25 | """ 26 | Approximate Kullback-Leibler divergence for arbitrary dimensionality 27 | """ 28 | if tol==None: 29 | tol = max( min(p_in.flatten()), min(q_in.flatten())) 30 | q = q_in.copy() 31 | p = p_in.copy() 32 | q[q bins[i]) * (Y[:-1] < bins[i+1]) )[0] 51 | 52 | if len(mask) > 0: 53 | f_KM[i] = np.mean(dY[mask]) # Conditional average ~ drift 54 | a_KM[i] = 0.5*np.mean(dY2[mask]) # Conditional variance ~ diffusion 55 | 56 | # Estimate error by variance of samples in the bin 57 | a_KM[i] = 0.5*np.mean(dY2[mask]) # Conditional average 58 | a_err[i] = np.std(dY2[mask])/np.sqrt(len(mask)) 59 | 60 | else: 61 | f_KM[i] = np.nan 62 | f_err[i] = np.nan 63 | a_KM[i] = np.nan 64 | a_err[i] = np.nan 65 | 66 | return f_KM, a_KM, f_err, a_err 67 | 68 | 69 | # Return optimal coefficients for finite-time correctio 70 | def AFP_opt(cost, params): 71 | ### RUN OPTIMIZATION PROBLEM 72 | start_time = time() 73 | Xi0 = params["Xi0"] 74 | 75 | is_complex = np.iscomplex(Xi0[0]) 76 | 77 | if is_complex: 78 | Xi0 = np.concatenate((np.real(Xi0), np.imag(Xi0))) # Split vector in two for complex 79 | opt_fun = lambda Xi: cost(Xi[:len(Xi)//2] + 1j*Xi[len(Xi)//2:], params) 80 | 81 | else: 82 | opt_fun = lambda Xi: cost(Xi, params) 83 | 84 | res = minimize(opt_fun, Xi0, method='nelder-mead', 85 | options={'disp': False, 'maxfev':int(1e4)}) 86 | print('%%%% Optimization time: {0} seconds, Cost: {1} %%%%'.format(time() - start_time, res.fun) ) 87 | 88 | # Return coefficients and cost function 89 | if is_complex: 90 | # Return to complex number 91 | return res.x[:len(res.x)//2] + 1j*res.x[len(res.x)//2:], res.fun 92 | else: 93 | return res.x, res.fun 94 | 95 | 96 | 97 | def SSR_loop(opt_fun, params): 98 | """ 99 | Stepwise sparse regression: general function for a given optimization problem 100 | opt_fun should take the parameters and return coefficients and cost 101 | 102 | Requires a list of drift and diffusion expressions, 103 | (although these are just passed to the opt_fun) 104 | """ 105 | 106 | # Lists of candidate expressions... coefficients are optimized 107 | f_expr, s_expr = params['f_expr'].copy(), params['s_expr'].copy() 108 | lib_f, lib_s = params['lib_f'].copy(), params['lib_s'].copy() 109 | Xi0 = params['Xi0'].copy() 110 | 111 | m = len(f_expr) + len(s_expr) 112 | 113 | Xi = np.zeros((m, m-1), dtype=Xi0.dtype) # Output results 114 | V = np.zeros((m-1)) # Cost at each step 115 | 116 | # Full regression problem as baseline 117 | Xi[:, 0], V[0] = opt_fun(params) 118 | 119 | # Start with all candidates 120 | active = np.array([i for i in range(m)]) 121 | 122 | # Iterate and threshold 123 | for k in range(1, m-1): 124 | # Loop through remaining terms and find the one that increases the cost function the least 125 | min_idx = -1 126 | V[k] = 1e8 127 | for j in range(len(active)): 128 | tmp_active = active.copy() 129 | tmp_active = np.delete(tmp_active, j) # Try deleting this term 130 | 131 | # Break off masks for drift/diffusion 132 | f_active = tmp_active[tmp_active < len(f_expr)] 133 | s_active = tmp_active[tmp_active >= len(f_expr)] - len(f_expr) 134 | print(f_active) 135 | print(s_active) 136 | 137 | print(f_expr[f_active], s_expr[s_active]) 138 | params['f_expr'] = f_expr[f_active] 139 | params['s_expr'] = s_expr[s_active] 140 | params['lib_f'] = lib_f[:, f_active] 141 | params['lib_s'] = lib_s[:, s_active] 142 | params['Xi0'] = Xi0[tmp_active] 143 | 144 | # Ensure that there is at least one drift and diffusion term left 145 | if len(s_active) > 0 and len(f_active) > 0: 146 | tmp_Xi, tmp_V = opt_fun(params) 147 | 148 | # Keep minimum cost 149 | if tmp_V < V[k]: 150 | # Ensure that there is at least one drift and diffusion term left 151 | #if (IS_DRIFT and len(f_active)>1) or (not IS_DRIFT and len(a_active)>1): 152 | min_idx = j 153 | V[k] = tmp_V 154 | min_Xi = tmp_Xi 155 | 156 | print("Cost: {0}".format(V[k])) 157 | # Delete least important term 158 | active = np.delete(active, min_idx) # Remove inactive index 159 | Xi0[active] = min_Xi # Re-initialize with best results from previous 160 | Xi[active, k] = min_Xi 161 | print(Xi[:, k]) 162 | 163 | return Xi, V 164 | 165 | def cost(Xi, params): 166 | """ 167 | Least-squares cost function for optimization 168 | This version is only good in 1D, but could be extended pretty easily 169 | Xi - current coefficient estimates 170 | param - inputs to optimization problem: grid points, list of candidate expressions, regularizations 171 | W, f_KM, a_KM, x_pts, y_pts, x_msh, y_msh, f_expr, a_expr, l1_reg, l2_reg, kl_reg, p_hist, etc 172 | """ 173 | 174 | # Unpack parameters 175 | W = params['W'] # Optimization weights 176 | 177 | # Kramers-Moyal coefficients 178 | f_KM, a_KM = params['f_KM'].flatten(), params['a_KM'].flatten() 179 | 180 | fp, afp = params['fp'], params['afp'] # Fokker-Planck solvers 181 | lib_f, lib_s = params['lib_f'], params['lib_s'] 182 | N = params['N'] 183 | 184 | # Construct parameterized drift and diffusion functions from libraries and current coefficients 185 | f_vals = lib_f @ Xi[:lib_f.shape[1]] 186 | a_vals = 0.5*(lib_s @ Xi[lib_f.shape[1]:])**2 187 | 188 | # Solve AFP equation to find finite-time corrected drift/diffusion 189 | # corresponding to the current parameters Xi 190 | afp.precompute_operator(np.reshape(f_vals, N), np.reshape(a_vals, N)) 191 | f_tau, a_tau = afp.solve(params['tau']) 192 | 193 | # Histogram points without data have NaN values in K-M average - ignore these in the average 194 | mask = np.nonzero(np.isfinite(f_KM))[0] 195 | V = np.sum(W[0, mask]*abs(f_tau[mask] - f_KM[mask])**2) \ 196 | + np.sum(W[1, mask]*abs(a_tau[mask] - a_KM[mask])**2) 197 | 198 | # Include PDF constraint via Kullbeck-Leibler divergence regularization 199 | if params['kl_reg'] > 0: 200 | p_hist = params['p_hist'] # Empirical PDF 201 | p_est = fp.solve(f_vals, a_vals) # Solve Fokker-Planck equation for steady-state PDF 202 | kl = utils.kl_divergence(p_hist, p_est, dx=fp.dx, tol=1e-6) 203 | kl = max(0, kl) # Numerical integration can occasionally produce small negative values 204 | V += params['kl_reg']*kl 205 | return V 206 | 207 | 208 | # 1D Markov test 209 | def markov_test(X, lag, N=32, L=2): 210 | # Lagged time series 211 | X1 = X[:-2*lag:lag] 212 | X2 = X[lag:-lag:lag] 213 | X3 = X[2*lag::lag] 214 | 215 | # Two-time joint pdfs 216 | bins = np.linspace(-L, L, N+1) 217 | dx = bins[1]-bins[0] 218 | p12, _, _ = np.histogram2d(X1, X2, bins=[bins, bins], density=True) 219 | p23, _, _ = np.histogram2d(X2, X3, bins=[bins, bins], density=True) 220 | p2, _ = np.histogram(X2, bins=bins, density=True) 221 | p2[p2<1e-4] = 1e-4 222 | 223 | # Conditional PDF (Markov assumption) 224 | pcond_23 = p23.copy() 225 | for j in range(pcond_23.shape[1]): 226 | pcond_23[:, j] = pcond_23[:, j]/p2 227 | 228 | # Three-time PDFs 229 | p123, _ = np.histogramdd(np.array([X1, X2, X3]).T, bins=np.array([bins, bins, bins]), density=True) 230 | p123_markov = np.einsum('ij,jk->ijk',p12, pcond_23) 231 | 232 | # Chi^2 value 233 | #return utils.ntrapz( (p123 - p123_markov)**2, [dx, dx, dx] )/(np.var(p123.flatten()) + np.var(p123_markov.flatten())) 234 | return kl_divergence(p123, p123_markov, dx=[dx, dx, dx], tol=1e-6) 235 | 236 | 237 | 238 | ### FAST AUTOCORRELATION FUNCTION 239 | # From https://dfm.io/posts/autocorr/ 240 | 241 | def next_pow_two(n): 242 | i = 1 243 | while i < n: 244 | i = i << 1 245 | return i 246 | 247 | def autocorr_func_1d(x, norm=True): 248 | x = np.atleast_1d(x) 249 | if len(x.shape) != 1: 250 | raise ValueError("invalid dimensions for 1D autocorrelation function") 251 | n = next_pow_two(len(x)) 252 | 253 | # Compute the FFT and then (from that) the auto-correlation function 254 | f = np.fft.fft(x - np.mean(x), n=2*n) 255 | acf = np.fft.ifft(f * np.conjugate(f))[:len(x)].real 256 | acf /= 4*n 257 | 258 | # Optionally normalize 259 | if norm: 260 | acf /= acf[0] 261 | 262 | return acf --------------------------------------------------------------------------------