# [xplus, converged, iter, k] = butcherintegration(x, f, h, a, b, [Niter=100])## Performs numerical integration of f starting from x for h time units.## f must be a function handle that returns dx/dt. If a defines an implicit# method, f may also return the Jacobian as a second output; if so, Newton's# method is used. Otherwise, fixed-point iteration is used.## a and b should be the coefficients from the Butcher tableau for the# integration scheme. Scalars are accepted for 1-stage methods. If the a# matrix is strictly lower triangular, the explicit formulas are used.## Niter is the maximum number of iterations to perform.importnumpyasnpfromnumpy.linalgimportnormfromscipy.linalgimportsolvedefbutcherintegration(x,func,h,a,b,Niter=100):""" Performs numerical integration of f starting from x for h time units. Returns (xplus, converged, iter, k). """x=np.asarray(x,dtype=float).reshape(-1)Nx=len(x)# Accept scalars for 1-stage methods.ifnp.isscalar(a):a=np.array([[a]],dtype=float)else:a=np.asarray(a,dtype=float)ifnp.isscalar(b):b=np.array([b],dtype=float)else:b=np.asarray(b,dtype=float)Ns=len(b)ifa.shape!=(Ns,Ns):raiseValueError('a must be square with one row/col per element of b!')b=h*ba=h*aifnp.count_nonzero(np.triu(a))==0:k=_butcher_explicit(x,func,a)converged=Trueniter=0else:# Probe whether func returns a Jacobian.test=func(x)has_jac=isinstance(test,tuple)ifhas_jac:k,converged,niter=_butcher_implicit_newton(x,func,a,Niter)else:k,converged,niter=_butcher_implicit_fixedpoint(x,func,a,Niter)xplus=x+k@breturnxplus,converged,niter,kdef_butcher_explicit(x,func,a):"""Explicit integration step."""Nx=len(x)Ns=a.shape[0]k=np.zeros((Nx,Ns))foriinrange(Ns):val=func(x+k@a[i,:])k[:,i]=np.asarray(val).reshape(-1)returnkdef_butcher_implicit_newton(x,func,a,Niter):"""Implicit step via Newton's method (func must return (f, J))."""Nx=len(x)Ns=a.shape[0]k=np.zeros((Nx,Ns))converged=Falseniter=0whilenotconvergedandniter<Niter:niter+=1dX=k@a.TifNs==1:f,J=func(x+dX.reshape(-1))f=np.asarray(f).reshape(-1)J=np.asarray(J)*a[0,0]else:f=np.zeros(Nx*Ns)J=np.zeros((Nx*Ns,Nx*Ns))foriinrange(Ns):I=slice(i*Nx,(i+1)*Nx)fi,Ji=func(x+dX[:,i])f[I]=np.asarray(fi).reshape(-1)J[I,:]=np.kron(a[i,:],np.asarray(Ji))f=f-k.flatten(order='F')J=J-np.eye(Nx*Ns)dk=-solve(J,f)k=k+dk.reshape(Nx,Ns,order='F')converged=norm(dk)<1e-8returnk,converged,niterdef_butcher_implicit_fixedpoint(x,func,a,Niter):"""Implicit step via fixed-point iteration (func returns only f)."""Nx=len(x)Ns=a.shape[0]k=np.zeros((Nx,Ns))converged=Falseniter=0whilenotconvergedandniter<Niter:niter+=1k_old=k.copy()dX=k@a.Tforiinrange(Ns):xi=(x+dX[:,i])ifNs>1else(x+dX.reshape(-1))val=func(xi)k[:,i]=np.asarray(val).reshape(-1)converged=norm(k-k_old)<1e-10returnk,converged,niter
# errorplot2d(x, T, [location])## Makes a phase plot of 2D trajectories. x is a dict mapping names to (2, N)# arrays. If x contains an 'exact' key it is plotted as a solid black line;# all others are plotted as dashed lines.importosimportmatplotlib.pyplotaspltdeferrorplot2d(x,T,location=None):ifos.getenv('OMIT_PLOTS')=='true':return""" Makes a phase plot of 2D trajectories. x must be a dict with string keys mapping to (2, N) arrays. The 'exact' key, if present, is plotted as a solid black line; others as dashed. """plt.figure()forkeyinx:style='-k'ifkey=='exact'else'--'plt.plot(x[key][0,:],x[key][1,:],style,label=key)plt.axis('equal')iflocation:plt.legend(loc=location)else:plt.legend()plt.xlabel('$x_1$')plt.ylabel('$x_2$',rotation=0)
# Example explicit numerical integration of a linear ODE.# [depends] functions/butcherintegration.py functions/errorplot2d.py# [makes] explicitint_example.matimportosimportsyssys.path.insert(0,os.path.join(os.path.dirname(os.path.abspath(__file__)),'functions'))importnumpyasnpimportscipy.iofrombutcherintegrationimportbutcherintegrationfromerrorplot2dimporterrorplot2d# Butcher tableaustabs={}tabs['euler']={'a':np.array([[0]]),'b':np.array([1])}tabs['heun']={'a':np.array([[0,0],[1,0]]),'b':np.array([0.5,0.5])}tabs['rk4']={'a':np.array([[0,0,0,0],[0.5,0,0,0],[0,0.5,0,0],[0,0,1,0]]),'b':np.array([1/6,1/3,1/3,1/6])}tabnames=list(tabs.keys())# Run each methodA=np.array([[0,1],[-1,0]])deff(x):returnA@x.reshape(-1,1)x0=np.array([[1],[0]])T=2*np.piMeval=32x={}fornintabnames:Nstep=int(np.ceil(Meval/len(tabs[n]['b'])))h=T/Nstepx[n]=np.zeros((2,Nstep+1))x[n][:,0]=x0.flatten()foriinrange(Nstep):x[n][:,i+1],*_=butcherintegration(x[n][:,i],f,h,tabs[n]['a'],tabs[n]['b'])# Add exact solutiontheta=np.linspace(0,2*np.pi,257)x['exact']=np.vstack((np.cos(theta),-np.sin(theta)))# Make ploterrorplot2d(x,T,'lower right')# Save datascipy.io.savemat('explicitint_example.mat',{'x':x,'T':T})