diff --git a/pyproject.toml b/pyproject.toml index 21ea08f..eb57bec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,8 @@ classifiers=[ "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", ] @@ -34,8 +36,19 @@ demo = [ "matplotlib", "jupyter" ] -dev = ["twine"] -test = ["pytest"] +dev = [ + "flake8", + "mccabe", + "mypy", + "pylint", + "twine" +] + +test = [ + "coverage", + "pytest", + "tox" +] [project.urls] homepage = "https://github.com/whalenpt/rkstiff" diff --git a/rkstiff/__init__.py b/rkstiff/__init__.py index 76ac8d2..84d9548 100644 --- a/rkstiff/__init__.py +++ b/rkstiff/__init__.py @@ -1,6 +1,6 @@ try: import importlib_metadata as metadata -except: +except ModuleNotFoundError: import importlib.metadata as metadata __version__ = metadata.version("rkstiff") diff --git a/rkstiff/derivatives.py b/rkstiff/derivatives.py index 0c21238..b706b5b 100644 --- a/rkstiff/derivatives.py +++ b/rkstiff/derivatives.py @@ -1,8 +1,8 @@ - import numpy as np -def dx_rfft(kx,u,n=1): - """ Takes derivative(s) of a real valued array in spectral space + +def dx_rfft(kx, u, n=1): + """Takes derivative(s) of a real valued array in spectral space INPUTS kx - wavenumbers of the spectral grid u - real valued array @@ -10,23 +10,24 @@ def dx_rfft(kx,u,n=1): OUTPUTS uxn - derivative of u to the nth power """ - if not isinstance(n,int): - raise TypeError('derivative order n must be an integer, it is {}'.format(n)) + if not isinstance(n, int): + raise TypeError("derivative order n must be an integer, it is {}".format(n)) if n < 0: - raise ValueError('derivative order n must non-negative, it is {}'.format(n)) + raise ValueError("derivative order n must non-negative, it is {}".format(n)) if n == 0: return u uFFT = np.fft.rfft(u) if n == 1: - uxn = np.fft.irfft(1j*kx*uFFT) + uxn = np.fft.irfft(1j * kx * uFFT) else: - uxn = np.fft.irfft(np.power(1j*kx,n)*uFFT) + uxn = np.fft.irfft(np.power(1j * kx, n) * uFFT) return uxn -def dx_fft(kx,u,n=1): - """ Takes derivative(s) of a complex valued array in spectral space + +def dx_fft(kx, u, n=1): + """Takes derivative(s) of a complex valued array in spectral space INPUTS kx - wavenumbers of the spectral grid u - complex valued array @@ -35,19 +36,16 @@ def dx_fft(kx,u,n=1): uxn - derivative of u to the nth power """ - if not isinstance(n,int): - raise TypeError('derivative order n must be an integer, it is {}'.format(n)) + if not isinstance(n, int): + raise TypeError("derivative order n must be an integer, it is {}".format(n)) if n < 0: - raise ValueError('derivative order n must non-negative, it is {}'.format(n)) + raise ValueError("derivative order n must non-negative, it is {}".format(n)) if n == 0: return u uFFT = np.fft.fft(u) if n == 1: - uxn = np.fft.ifft(1j*kx*uFFT) + uxn = np.fft.ifft(1j * kx * uFFT) else: - uxn = np.fft.ifft(np.power(1j*kx,n)*uFFT) + uxn = np.fft.ifft(np.power(1j * kx, n) * uFFT) return uxn - - - diff --git a/rkstiff/etd.py b/rkstiff/etd.py index d199f30..77f163e 100644 --- a/rkstiff/etd.py +++ b/rkstiff/etd.py @@ -1,33 +1,36 @@ - import numpy as np -from rkstiff.solver import StiffSolverAS,StiffSolverCS +from rkstiff.solver import StiffSolverAS, StiffSolverCS + def phi1(z): - """ Computes RKETD psi-function of the first order. + """Computes RKETD psi-function of the first order. INPUTS - z - real or complex-valued input array + z - real or complex-valued input array OUTPUT return (exp(z)-1)/z -> real or complex-valued RKETD function of first order """ - return (np.exp(z)-1)/z + return (np.exp(z) - 1) / z + def phi2(z): - """ Computes RKETD psi-function of the second order. + """Computes RKETD psi-function of the second order. INPUTS - z - real or complex-valued input array + z - real or complex-valued input array OUTPUT return 2!*(exp(z)-1-z/2)/z^2 -> real or complex-valued RKETD function of second order """ - return 2*(np.exp(z)-1-z)/z**2 + return 2 * (np.exp(z) - 1 - z) / z**2 + def phi3(z): - """ Computes RKETD psi-function of the third order. + """Computes RKETD psi-function of the third order. INPUTS - z - real or complex-valued input array + z - real or complex-valued input array OUTPUT return 3!*(exp(z)-1-z/2-z^3/6)/z^3 -> real or complex-valued RKETD function of third order """ - return 6*(np.exp(z)-1-z-z**2/2)/z**3 + return 6 * (np.exp(z) - 1 - z - z**2 / 2) / z**3 + class ETDAS(StiffSolverAS): """ @@ -47,25 +50,60 @@ class ETDAS(StiffSolverAS): """ - def __init__(self,linop,NLfunc,epsilon=1e-4,incrF = 1.25, decrF = 0.85, safetyF = 0.8,\ - adapt_cutoff = 0.01, minh = 1e-16, modecutoff = 0.01, contour_points = 32,\ - contour_radius = 1.0): - super().__init__(linop,NLfunc,epsilon=epsilon,incrF=incrF,decrF=decrF,\ - safetyF=safetyF,adapt_cutoff=adapt_cutoff,minh=minh) + def __init__( + self, + linop, + NLfunc, + epsilon=1e-4, + incrF=1.25, + decrF=0.85, + safetyF=0.8, + adapt_cutoff=0.01, + minh=1e-16, + modecutoff=0.01, + contour_points=32, + contour_radius=1.0, + ): + super().__init__( + linop, + NLfunc, + epsilon=epsilon, + incrF=incrF, + decrF=decrF, + safetyF=safetyF, + adapt_cutoff=adapt_cutoff, + minh=minh, + ) self.modecutoff = modecutoff if (self.modecutoff > 1.0) or (self.modecutoff <= 0): - raise ValueError('modecutoff must be between 0.0 and 1.0 but is {}'.format(self.modecutoff)) + raise ValueError( + "modecutoff must be between 0.0 and 1.0 but is {}".format( + self.modecutoff + ) + ) self.contour_points = contour_points - if not isinstance(self.contour_points,int): - raise TypeError('contour_points must be an integer but is {}'.format(self.contour_points)) + if not isinstance(self.contour_points, int): + raise TypeError( + "contour_points must be an integer but is {}".format( + self.contour_points + ) + ) if self.contour_points <= 1: - raise ValueError('contour_points must be an integer greater than 1 but is {}'.format(self.contour_points)) + raise ValueError( + "contour_points must be an integer greater than 1 but is {}".format( + self.contour_points + ) + ) self.contour_radius = contour_radius if self.contour_radius <= 0: - raise ValueError('contour_radius must greater than 0 but is {}'.format(self.contour_radius)) + raise ValueError( + "contour_radius must greater than 0 but is {}".format( + self.contour_radius + ) + ) self._h_coeff = None - + class ETDCS(StiffSolverCS): """ @@ -85,21 +123,35 @@ class ETDCS(StiffSolverCS): """ - def __init__(self,linop,NLfunc,modecutoff = 0.01,contour_points = 32,contour_radius = 1.0): - super().__init__(linop,NLfunc) + def __init__( + self, linop, NLfunc, modecutoff=0.01, contour_points=32, contour_radius=1.0 + ): + super().__init__(linop, NLfunc) self.modecutoff = modecutoff if (self.modecutoff > 1.0) or (self.modecutoff <= 0): - raise ValueError('modecutoff must be between 0.0 and 1.0 but is {}'.format(self.modecutoff)) + raise ValueError( + "modecutoff must be between 0.0 and 1.0 but is {}".format( + self.modecutoff + ) + ) self.contour_points = contour_points - if not isinstance(self.contour_points,int): - raise TypeError('contour_points must be an integer but is {}'.format(self.contour_points)) + if not isinstance(self.contour_points, int): + raise TypeError( + "contour_points must be an integer but is {}".format( + self.contour_points + ) + ) if self.contour_points <= 1: - raise ValueError('contour_points must be an integer greater than 1 but is {}'.format(self.contour_points)) + raise ValueError( + "contour_points must be an integer greater than 1 but is {}".format( + self.contour_points + ) + ) self.contour_radius = contour_radius if self.contour_radius <= 0: - raise ValueError('contour_radius must greater than 0 but is {}'.format(self.contour_radius)) + raise ValueError( + "contour_radius must greater than 0 but is {}".format( + self.contour_radius + ) + ) self._h_coeff = None - - - - diff --git a/rkstiff/etd34.py b/rkstiff/etd34.py index ec36ba2..1c72a0a 100644 --- a/rkstiff/etd34.py +++ b/rkstiff/etd34.py @@ -1,205 +1,249 @@ - -from rkstiff.etd import ETDAS,phi1,phi2,phi3 +from rkstiff.etd import ETDAS, phi1, phi2, phi3 import numpy as np -from typing import Callable +from typing import Callable, Union from scipy.linalg import expm + class _ETD34_Diagonal: """ - ETD34 diagonal system strategy for ETD34 solver + ETD34 diagonal system strategy for ETD34 solver """ - def __init__(self,linop,NLfunc,contourM,contourR,modecutoff): + def __init__(self, linop, NLfunc, contourM, contourR, modecutoff): self.linop = linop self.NLfunc = NLfunc self.M = contourM self.R = contourR self.modecutoff = modecutoff - + N = linop.shape[0] - self._EL, self._EL2 = [np.zeros(N,dtype=np.complex128) for _ in range(2)] - self._a21, self._a31, self._a32, self._a41, self._a43, self._a51,\ - self._a52, self._a54 = [np.zeros(N,dtype=np.complex128) for _ in range(8)] - self._NL1, self._NL2, self._NL3,self._NL4, self._NL5 = [np.zeros(N,dtype=np.complex128) for _ in range(5)] - self._k = np.zeros(N,dtype=np.complex128) - self._err = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): - z = h*self.linop - self._updateCoeffsDiagonal(h,z) - - def _updateCoeffsDiagonal(self,h,z): + self._EL, self._EL2 = [np.zeros(N, dtype=np.complex128) for _ in range(2)] + ( + self._a21, + self._a31, + self._a32, + self._a41, + self._a43, + self._a51, + self._a52, + self._a54, + ) = [np.zeros(N, dtype=np.complex128) for _ in range(8)] + self._NL1, self._NL2, self._NL3, self._NL4, self._NL5 = [ + np.zeros(N, dtype=np.complex128) for _ in range(5) + ] + self._k = np.zeros(N, dtype=np.complex128) + self._err = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): + z = h * self.linop + self._updateCoeffsDiagonal(h, z) + + def _updateCoeffsDiagonal(self, h, z): self._EL = np.exp(z) - self._EL2 = np.exp(z/2) - + self._EL2 = np.exp(z / 2) + smallmode_idx = np.abs(z) < self.modecutoff - zb = z[~smallmode_idx] # z big + zb = z[~smallmode_idx] # z big # compute big mode coeffs - phi1_12 = h*phi1(zb/2) - phi2_12 = h*phi2(zb/2) - phi1_1 = h*phi1(zb) - phi2_1 = h*phi2(zb) - phi3_1 = h*phi3(zb) - - self._a21[~smallmode_idx] = 0.5*phi1_12 - self._a31[~smallmode_idx] = 0.5*(phi1_12-phi2_12) - self._a32[~smallmode_idx] = 0.5*phi2_12 + phi1_12 = h * phi1(zb / 2) + phi2_12 = h * phi2(zb / 2) + phi1_1 = h * phi1(zb) + phi2_1 = h * phi2(zb) + phi3_1 = h * phi3(zb) + + self._a21[~smallmode_idx] = 0.5 * phi1_12 + self._a31[~smallmode_idx] = 0.5 * (phi1_12 - phi2_12) + self._a32[~smallmode_idx] = 0.5 * phi2_12 self._a41[~smallmode_idx] = phi1_1 - phi2_1 self._a43[~smallmode_idx] = phi2_1 - self._a51[~smallmode_idx] = phi1_1 - (3.0/2)*phi2_1+(2.0/3)*phi3_1 - self._a52[~smallmode_idx] = phi2_1 - (2.0/3)*phi3_1 - self._a54[~smallmode_idx] = -(1.0/2)*phi2_1 + (2.0/3)*phi3_1 - + self._a51[~smallmode_idx] = phi1_1 - (3.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1 + self._a52[~smallmode_idx] = phi2_1 - (2.0 / 3) * phi3_1 + self._a54[~smallmode_idx] = -(1.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1 + # compute small mode coeffs - zs = z[smallmode_idx] # z small - r = self.R*np.exp(2j*np.pi*np.arange(0.5,self.M)/self.M) - rr, zz = np.meshgrid(r,zs) - Z = zz+rr - - phi1_12 = h*np.sum(phi1(Z/2),axis=1)/self.M - phi2_12 = h*np.sum(phi2(Z/2),axis=1)/self.M - phi1_1 = h*np.sum(phi1(Z),axis=1)/self.M - phi2_1 = h*np.sum(phi2(Z),axis=1)/self.M - phi3_1 = h*np.sum(phi3(Z),axis=1)/self.M - - self._a21[smallmode_idx] = 0.5*phi1_12 - self._a31[smallmode_idx] = 0.5*(phi1_12-phi2_12) - self._a32[smallmode_idx] = 0.5*phi2_12 + zs = z[smallmode_idx] # z small + r = self.R * np.exp(2j * np.pi * np.arange(0.5, self.M) / self.M) + rr, zz = np.meshgrid(r, zs) + Z = zz + rr + + phi1_12 = h * np.sum(phi1(Z / 2), axis=1) / self.M + phi2_12 = h * np.sum(phi2(Z / 2), axis=1) / self.M + phi1_1 = h * np.sum(phi1(Z), axis=1) / self.M + phi2_1 = h * np.sum(phi2(Z), axis=1) / self.M + phi3_1 = h * np.sum(phi3(Z), axis=1) / self.M + + self._a21[smallmode_idx] = 0.5 * phi1_12 + self._a31[smallmode_idx] = 0.5 * (phi1_12 - phi2_12) + self._a32[smallmode_idx] = 0.5 * phi2_12 self._a41[smallmode_idx] = phi1_1 - phi2_1 self._a43[smallmode_idx] = phi2_1 - self._a51[smallmode_idx] = phi1_1 - (3.0/2)*phi2_1+(2.0/3)*phi3_1 - self._a52[smallmode_idx] = phi2_1 - (2.0/3)*phi3_1 - self._a54[smallmode_idx] = -(1.0/2)*phi2_1 + (2.0/3)*phi3_1 - - def _N1_init(self,u): + self._a51[smallmode_idx] = phi1_1 - (3.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1 + self._a52[smallmode_idx] = phi2_1 - (2.0 / 3) * phi3_1 + self._a54[smallmode_idx] = -(1.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1 + + def _N1_init(self, u): self._NL1 = self.NLfunc(u) - - def _updateStages(self,u,h,accept): + + def _updateStages(self, u, h, accept): # Use First is same as last principle (FSAL) -> k5 stage is input u for next step if accept: - self._NL1 = self._NL5.copy() - # If not accept, then step failed, reuse previously computed N1 - self._k = self._EL2*u + self._a21*self._NL1 + self._NL1 = self._NL5.copy() + # If not accept, then step failed, reuse previously computed N1 + self._k = self._EL2 * u + self._a21 * self._NL1 self._NL2 = self.NLfunc(self._k) - self._k = self._EL2*u + self._a31*self._NL1 + self._a32*self._NL2 + self._k = self._EL2 * u + self._a31 * self._NL1 + self._a32 * self._NL2 self._NL3 = self.NLfunc(self._k) - self._k = self._EL*u + self._a41*self._NL1 + self._a43*self._NL3 + self._k = self._EL * u + self._a41 * self._NL1 + self._a43 * self._NL3 self._NL4 = self.NLfunc(self._k) - self._k = self._EL*u + self._a51*self._NL1 + self._a52*(self._NL2+self._NL3) \ - + self._a54*self._NL4 + self._k = ( + self._EL * u + + self._a51 * self._NL1 + + self._a52 * (self._NL2 + self._NL3) + + self._a54 * self._NL4 + ) self._NL5 = self.NLfunc(self._k) - self._err = self._a54*(self._NL4-self._NL5) + self._err = self._a54 * (self._NL4 - self._NL5) return self._k, self._err + class _ETD34_Diagonalized(_ETD34_Diagonal): """ - ETD34 non-diagonal system with eigenvector diagonalization strategy for ETD34 solver + ETD34 non-diagonal system with eigenvector diagonalization strategy for ETD34 solver """ - def __init__(self,linop,NLfunc,contourM,contourR,modecutoff): - super().__init__(linop,NLfunc,contourM,contourR,modecutoff) + + def __init__(self, linop, NLfunc, contourM, contourR, modecutoff): + super().__init__(linop, NLfunc, contourM, contourR, modecutoff) if len(linop.shape) == 1: - raise Exception('Cannot diagonalize a 1D system') + raise Exception("Cannot diagonalize a 1D system") linop_cond = np.linalg.cond(linop) if linop_cond > 1e16: - raise Exception('Cannot diagonalize a non-invertible linear operator L') + raise Exception("Cannot diagonalize a non-invertible linear operator L") if linop_cond > 1000: - print('''Warning: linear matrix array has a large condition number of {:.2f}, - method may be unstable'''.format(linop_cond)) + print( + """Warning: linear matrix array has a large condition number of {:.2f}, + method may be unstable""".format( + linop_cond + ) + ) self._eig_vals, self._S = np.linalg.eig(linop) self._Sinv = np.linalg.inv(self._S) self._v = np.zeros(linop.shape[0]) - - def _updateCoeffs(self,h): - z = h*self._eig_vals - self._updateCoeffsDiagonal(h,z) - - def _N1_init(self,u): + + def _updateCoeffs(self, h): + z = h * self._eig_vals + self._updateCoeffsDiagonal(h, z) + + def _N1_init(self, u): self._NL1 = self._Sinv.dot(self.NLfunc(u)) self._v = self._Sinv.dot(u) - - def _updateStages(self,u,h,accept): + + def _updateStages(self, u, h, accept): # Use First is same as last principle (FSAL) -> k5 stage is input u for next step if accept: - self._NL1 = self._NL5.copy() + self._NL1 = self._NL5.copy() self._v = self._Sinv.dot(u) - - self._k = self._EL2*self._v + self._a21*self._NL1 + + self._k = self._EL2 * self._v + self._a21 * self._NL1 self._NL2 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._k = self._EL2*self._v + self._a31*self._NL1 + self._a32*self._NL2 + self._k = self._EL2 * self._v + self._a31 * self._NL1 + self._a32 * self._NL2 self._NL3 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._k = self._EL*self._v + self._a41*self._NL1 + self._a43*self._NL3 + self._k = self._EL * self._v + self._a41 * self._NL1 + self._a43 * self._NL3 self._NL4 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._k = self._EL*self._v + self._a51*self._NL1 + self._a52*(self._NL2+self._NL3) \ - + self._a54*self._NL4 + self._k = ( + self._EL * self._v + + self._a51 * self._NL1 + + self._a52 * (self._NL2 + self._NL3) + + self._a54 * self._NL4 + ) self._NL5 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._err = self._a54*(self._NL4-self._NL5) + self._err = self._a54 * (self._NL4 - self._NL5) return self._S.dot(self._k), self._err class _ETD34_NonDiagonal: """ - ETD34 non-diagonal system strategy for ETD34 solver + ETD34 non-diagonal system strategy for ETD34 solver """ - def __init__(self,linop,NLfunc,contourM,contourR): + + def __init__(self, linop, NLfunc, contourM, contourR): self.linop = linop self.NLfunc = NLfunc self.M = contourM self.R = contourR - + N = linop.shape[0] - self._EL, self._EL2 = [np.zeros(shape=linop.shape,dtype=np.complex128) for _ in range(2)] - self._a21, self._a31, self._a32, self._a41, self._a43, self._a51,\ - self._a52, self._a54 = [np.zeros(shape=linop.shape,dtype=np.complex128) for _ in range(8)] - self._NL1, self._NL2, self._NL3,self._NL4, self._NL5 = [np.zeros(N,dtype=np.complex128) for _ in range(5)] - self._k = np.zeros(N,dtype=np.complex128) - self._err = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): - z = h*self.linop + self._EL, self._EL2 = [ + np.zeros(shape=linop.shape, dtype=np.complex128) for _ in range(2) + ] + ( + self._a21, + self._a31, + self._a32, + self._a41, + self._a43, + self._a51, + self._a52, + self._a54, + ) = [np.zeros(shape=linop.shape, dtype=np.complex128) for _ in range(8)] + self._NL1, self._NL2, self._NL3, self._NL4, self._NL5 = [ + np.zeros(N, dtype=np.complex128) for _ in range(5) + ] + self._k = np.zeros(N, dtype=np.complex128) + self._err = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): + z = h * self.linop self._EL = expm(z) - self._EL2 = expm(z/2) + self._EL2 = expm(z / 2) + + contour_points = self.R * np.exp(2j * np.pi * np.arange(0.5, self.M) / self.M) - contour_points = self.R*np.exp(2j*np.pi*np.arange(0.5,self.M)/self.M) - - phi1_12, phi2_12,phi1_1,phi2_1,phi3_1 = [np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(5)] + phi1_12, phi2_12, phi1_1, phi2_1, phi3_1 = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(5) + ] for point in contour_points: - Q = np.linalg.inv(point*np.eye(*self.linop.shape)-z) - Q2 = np.linalg.inv(point*np.eye(*self.linop.shape)-z/2) - phi1_12 += point*phi1(point)*Q2/self.M - phi2_12 += point*phi2(point)*Q2/self.M - phi1_1 += point*phi1(point)*Q/self.M - phi2_1 += point*phi2(point)*Q/self.M - phi3_1 += point*phi3(point)*Q/self.M - - self._a21 = 0.5*h*phi1_12 - self._a31 = 0.5*h*(phi1_12-phi2_12) - self._a32 = 0.5*h*phi2_12 - self._a41 = h*(phi1_1 - phi2_1) - self._a43 = h*phi2_1 - self._a51 = h*(phi1_1 - (3.0/2)*phi2_1+(2.0/3)*phi3_1) - self._a52 = h*(phi2_1 - (2.0/3)*phi3_1) - self._a54 = h*(-(1.0/2)*phi2_1 + (2.0/3)*phi3_1) - - def _N1_init(self,u): + Q = np.linalg.inv(point * np.eye(*self.linop.shape) - z) + Q2 = np.linalg.inv(point * np.eye(*self.linop.shape) - z / 2) + phi1_12 += point * phi1(point) * Q2 / self.M + phi2_12 += point * phi2(point) * Q2 / self.M + phi1_1 += point * phi1(point) * Q / self.M + phi2_1 += point * phi2(point) * Q / self.M + phi3_1 += point * phi3(point) * Q / self.M + + self._a21 = 0.5 * h * phi1_12 + self._a31 = 0.5 * h * (phi1_12 - phi2_12) + self._a32 = 0.5 * h * phi2_12 + self._a41 = h * (phi1_1 - phi2_1) + self._a43 = h * phi2_1 + self._a51 = h * (phi1_1 - (3.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1) + self._a52 = h * (phi2_1 - (2.0 / 3) * phi3_1) + self._a54 = h * (-(1.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1) + + def _N1_init(self, u): self._NL1 = self.NLfunc(u) - - def _updateStages(self,u,h,accept): + + def _updateStages(self, u, h, accept): # Use First is same as last principle (FSAL) -> k5 stage is input u for next step if accept: - self._NL1 = self._NL5.copy() - + self._NL1 = self._NL5.copy() + self._k = self._EL2.dot(u) + self._a21.dot(self._NL1) self._NL2 = self.NLfunc(self._k) self._k = self._EL2.dot(u) + self._a31.dot(self._NL1) + self._a32.dot(self._NL2) self._NL3 = self.NLfunc(self._k) self._k = self._EL.dot(u) + self._a41.dot(self._NL1) + self._a43.dot(self._NL3) self._NL4 = self.NLfunc(self._k) - self._k = self._EL.dot(u) + self._a51.dot(self._NL1) + self._a52.dot(self._NL2+self._NL3) \ - + self._a54.dot(self._NL4) + self._k = ( + self._EL.dot(u) + + self._a51.dot(self._NL1) + + self._a52.dot(self._NL2 + self._NL3) + + self._a54.dot(self._NL4) + ) self._NL5 = self.NLfunc(self._k) - self._err = self._a54.dot(self._NL4-self._NL5) + self._err = self._a54.dot(self._NL4 - self._NL5) return self._k, self._err + class ETD34(ETDAS): """ Exponential time-differencing adaptive step solver of 4th order with 3rd order embedding @@ -209,30 +253,41 @@ class ETD34(ETDAS): linop : np.array NLfunc : function - t : time-array stored with evolve function call - u : output-array stored with evolve function call + t : time-array stored with evolve function call + u : output-array stored with evolve function call logs : array of info stored related to the solver ETD Parameters (see ETDAS class in etd module) ______________ modecutoff : float contour_points : int - contour_radius : float + contour_radius : float StiffSolverAS Parameters (see StiffSolverAS class in solver module) ________________________ epsilon : float incrF : float - decrF : float + decrF : float safetyF : float adapt_cutoff : float - minh : float + minh : float """ - - def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray],epsilon : float = 1e-4,incrF : float = 1.25,\ - decrF : float = 0.85, safetyF : float = 0.8, adapt_cutoff : float = 0.01,\ - minh : float = 1e-16, modecutoff : float = 0.01, contour_points : int = 32, - contour_radius : float = 1.0, diagonalize : bool = False): + + def __init__( + self, + linop: np.ndarray, + NLfunc: Callable[[np.ndarray], np.ndarray], + epsilon: float = 1e-4, + incrF: float = 1.25, + decrF: float = 0.85, + safetyF: float = 0.8, + adapt_cutoff: float = 0.01, + minh: float = 1e-16, + modecutoff: float = 0.01, + contour_points: int = 32, + contour_radius: float = 1.0, + diagonalize: bool = False, + ): """ INPUTS ______ @@ -241,60 +296,74 @@ def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray], Linear operator (L) in the system dtU = LU + NL(U). Can be either a 2D numpy array (matrix) or a 1D array (diagonal system). L can be either real-valued or complex-valued. - NLfunc : function + NLfunc : function Nonlinear function (NL(U)) in the system dtU = LU + NL(U). Can be a complex or real-valued function. diagonalize : bool, optional - Diagonalize the linear operator (matrix) and solve the diagonalized system + Diagonalize the linear operator (matrix) and solve the diagonalized system ETDAS variables: modecutoff, contour_points, contour_radius (see ETDAS documentation from etd module) - StiffSolverAS variables: epsilon, incrF, decrF, safetyF, adapt_cutoff, minh + StiffSolverAS variables: epsilon, incrF, decrF, safetyF, adapt_cutoff, minh (see StiffSolverAS documentation from solver module) """ - super().__init__(linop,NLfunc,epsilon=epsilon,incrF=incrF,decrF=decrF,\ - safetyF=safetyF,adapt_cutoff=adapt_cutoff,minh=minh,\ - modecutoff=modecutoff,contour_points=contour_points,\ - contour_radius=contour_radius) - self._method = None + super().__init__( + linop, + NLfunc, + epsilon=epsilon, + incrF=incrF, + decrF=decrF, + safetyF=safetyF, + adapt_cutoff=adapt_cutoff, + minh=minh, + modecutoff=modecutoff, + contour_points=contour_points, + contour_radius=contour_radius, + ) + self._method = Union[_ETD34_Diagonal, _ETD34_Diagonalized] if self._diag: - self._method = _ETD34_Diagonal(linop,NLfunc,self.contour_points,\ - self.contour_radius,self.modecutoff) + self._method = _ETD34_Diagonal( + linop, NLfunc, self.contour_points, self.contour_radius, self.modecutoff + ) else: if diagonalize: - self._method = _ETD34_Diagonalized(linop,NLfunc,self.contour_points,\ - self.contour_radius,self.modecutoff) + self._method = _ETD34_Diagonalized( + linop, + NLfunc, + self.contour_points, + self.contour_radius, + self.modecutoff, + ) else: - self._method = _ETD34_NonDiagonal(linop,NLfunc,self.contour_points,\ - self.contour_radius) + self._method = _ETD34_NonDiagonal( + linop, NLfunc, self.contour_points, self.contour_radius + ) self.__N1_init = False self._accept = False def _reset(self): - #Resets solver to its initial state + # Resets solver to its initial state self.__N1_init = False self._h_coeff = None self._accept = False - - def _updateCoeffs(self,h): + + def _updateCoeffs(self, h): # Update coefficients if step size h changed if h == self._h_coeff: return self._h_coeff = h self._method._updateCoeffs(h) self.logs.append("ETD34 coefficients updated") - - def _updateStages(self,u,h): + + def _updateStages(self, u, h): # Computes u_{n+1} from u_{n} through one RK passthrough self._updateCoeffs(h) if not self.__N1_init: self._method._N1_init(u) self.__N1_init = True - return self._method._updateStages(u,h,self._accept) - + return self._method._updateStages(u, h, self._accept) + def _q(self): # Order variable for computing suggested step size (embedded order + 1) - return 4 - - + return 4 diff --git a/rkstiff/etd35.py b/rkstiff/etd35.py index a618e7f..9cd1fc9 100644 --- a/rkstiff/etd35.py +++ b/rkstiff/etd35.py @@ -1,276 +1,390 @@ - -from rkstiff.etd import ETDAS,phi1,phi2,phi3 +from rkstiff.etd import ETDAS, phi1, phi2, phi3 import numpy as np -from typing import Callable +from typing import Callable, Union from scipy.linalg import expm + class _ETD35_Diagonal: """ - ETD35 diagonal system strategy for ETD35 solver + ETD35 diagonal system strategy for ETD35 solver """ - def __init__(self,linop,NLfunc,contourM,contourR,modecutoff): + + def __init__(self, linop, NLfunc, contourM, contourR, modecutoff): self.linop = linop self.NLfunc = NLfunc self.M = contourM self.R = contourR self.modecutoff = modecutoff - + N = linop.shape[0] - self._EL14, self._EL12, self._EL34, self._EL = [np.zeros(N,dtype=np.complex128) for _ in range(4)] - self._a21, self._a31, self._a32, self._a41, self._a43, self._a51,\ - self._a52,self._a54, self._a61, self._a62, self._a63,\ - self._a64, self._a65, self._a71, self._a73, self._a74,\ - self._a75, self._a76 = [np.zeros(N,dtype=np.complex128) for _ in range(18)] - self._NL1, self._NL2, self._NL3,self._NL4, self._NL5,\ - self._NL6 = [np.zeros(N,dtype=np.complex128) for _ in range(6)] - self._k = np.zeros(N,dtype=np.complex128) - self._err = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): - z = h*self.linop - self._updateCoeffsDiagonal(h,z) - - def _updateCoeffsDiagonal(self,h,z): + self._EL14, self._EL12, self._EL34, self._EL = [ + np.zeros(N, dtype=np.complex128) for _ in range(4) + ] + ( + self._a21, + self._a31, + self._a32, + self._a41, + self._a43, + self._a51, + self._a52, + self._a54, + self._a61, + self._a62, + self._a63, + self._a64, + self._a65, + self._a71, + self._a73, + self._a74, + self._a75, + self._a76, + ) = [np.zeros(N, dtype=np.complex128) for _ in range(18)] + self._NL1, self._NL2, self._NL3, self._NL4, self._NL5, self._NL6 = [ + np.zeros(N, dtype=np.complex128) for _ in range(6) + ] + self._k = np.zeros(N, dtype=np.complex128) + self._err = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): + z = h * self.linop + self._updateCoeffsDiagonal(h, z) + + def _updateCoeffsDiagonal(self, h, z): # diagonal system -> L is 1D array (of independent modes) - self._EL14 = np.exp(z/4.0) - self._EL12 = np.exp(z/2.0) - self._EL34 = np.exp(3.0*z/4.0) + self._EL14 = np.exp(z / 4.0) + self._EL12 = np.exp(z / 2.0) + self._EL34 = np.exp(3.0 * z / 4.0) self._EL = np.exp(z) smallmode_idx = np.abs(z) < self.modecutoff # compute big mode coeffs - zb = z[~smallmode_idx] # z big - phi1_14 = h*phi1(zb/4) - phi2_14 = h*phi2(zb/4) - phi1_12 = h*phi1(zb/2) - phi2_12 = h*phi2(zb/2) - phi1_34 = h*phi1(3*zb/4) - phi2_34 = h*phi2(3*zb/4) - phi1_1 = h*phi1(zb) - phi2_1 = h*phi2(zb) - phi3_1 = h*phi3(zb) - - self._a21[~smallmode_idx] = phi1_14/4. - self._a31[~smallmode_idx] = (phi1_14-phi2_14/2.)/4. - self._a32[~smallmode_idx] = phi2_14/8. - self._a41[~smallmode_idx] = (phi1_12 - phi2_12)/2. - self._a43[~smallmode_idx] = phi2_12/2. - self._a51[~smallmode_idx] = 3.0*(phi1_34 - 3.*phi2_34/4.)/4.0 - self._a52[~smallmode_idx] = -3*phi1_34/8. - self._a54[~smallmode_idx] = 9*phi2_34/16. - self._a61[~smallmode_idx] = (-77*phi1_1+59*phi2_1)/42. - self._a62[~smallmode_idx] = 8*phi1_1/7. - self._a63[~smallmode_idx] = (111*phi1_1-87*phi2_1)/28. - self._a65[~smallmode_idx] = (-47*phi1_1+143*phi2_1)/84. - self._a71[~smallmode_idx] = 7*(257*phi1_1 - 497*phi2_1 + 270*phi3_1)/2700 + zb = z[~smallmode_idx] # z big + phi1_14 = h * phi1(zb / 4) + phi2_14 = h * phi2(zb / 4) + phi1_12 = h * phi1(zb / 2) + phi2_12 = h * phi2(zb / 2) + phi1_34 = h * phi1(3 * zb / 4) + phi2_34 = h * phi2(3 * zb / 4) + phi1_1 = h * phi1(zb) + phi2_1 = h * phi2(zb) + phi3_1 = h * phi3(zb) + + self._a21[~smallmode_idx] = phi1_14 / 4.0 + self._a31[~smallmode_idx] = (phi1_14 - phi2_14 / 2.0) / 4.0 + self._a32[~smallmode_idx] = phi2_14 / 8.0 + self._a41[~smallmode_idx] = (phi1_12 - phi2_12) / 2.0 + self._a43[~smallmode_idx] = phi2_12 / 2.0 + self._a51[~smallmode_idx] = 3.0 * (phi1_34 - 3.0 * phi2_34 / 4.0) / 4.0 + self._a52[~smallmode_idx] = -3 * phi1_34 / 8.0 + self._a54[~smallmode_idx] = 9 * phi2_34 / 16.0 + self._a61[~smallmode_idx] = (-77 * phi1_1 + 59 * phi2_1) / 42.0 + self._a62[~smallmode_idx] = 8 * phi1_1 / 7.0 + self._a63[~smallmode_idx] = (111 * phi1_1 - 87 * phi2_1) / 28.0 + self._a65[~smallmode_idx] = (-47 * phi1_1 + 143 * phi2_1) / 84.0 + self._a71[~smallmode_idx] = ( + 7 * (257 * phi1_1 - 497 * phi2_1 + 270 * phi3_1) / 2700 + ) # Paper has error in a73/b3 phi2 coefficient (states this is 497 but it is actually 467) - self._a73[~smallmode_idx] = (1097*phi1_1 - 467*phi2_1 - 150*phi3_1)/1350 - self._a74[~smallmode_idx] = 2*(-49*phi1_1 + 199*phi2_1 - 135*phi3_1)/225 - self._a75[~smallmode_idx] = (-313*phi1_1 + 883*phi2_1 - 90*phi3_1)/1350 - self._a76[~smallmode_idx] = (509*phi1_1 - 2129*phi2_1 + 1830*phi3_1)/2700 - + self._a73[~smallmode_idx] = (1097 * phi1_1 - 467 * phi2_1 - 150 * phi3_1) / 1350 + self._a74[~smallmode_idx] = ( + 2 * (-49 * phi1_1 + 199 * phi2_1 - 135 * phi3_1) / 225 + ) + self._a75[~smallmode_idx] = (-313 * phi1_1 + 883 * phi2_1 - 90 * phi3_1) / 1350 + self._a76[~smallmode_idx] = ( + 509 * phi1_1 - 2129 * phi2_1 + 1830 * phi3_1 + ) / 2700 + # compute small mode coeffs - zs = z[smallmode_idx] # z small - r = self.R*np.exp(2j*np.pi*np.arange(0.5,self.M)/self.M) - rr, zz = np.meshgrid(r,zs) - Z = zz+rr - - phi1_14 = h*np.sum(phi1(Z/4),axis=1)/self.M - phi2_14 = h*np.sum(phi2(Z/4),axis=1)/self.M - phi1_12 = h*np.sum(phi1(Z/2),axis=1)/self.M - phi2_12 = h*np.sum(phi2(Z/2),axis=1)/self.M - phi1_34 = h*np.sum(phi1(3*Z/4),axis=1)/self.M - phi2_34 = h*np.sum(phi2(3*Z/4),axis=1)/self.M - phi1_1 = h*np.sum(phi1(Z),axis=1)/self.M - phi2_1 = h*np.sum(phi2(Z),axis=1)/self.M - phi3_1 = h*np.sum(phi3(Z),axis=1)/self.M - - self._a21[smallmode_idx] = phi1_14/4. - self._a31[smallmode_idx] = (phi1_14-phi2_14/2.)/4. - self._a32[smallmode_idx] = phi2_14/8. - self._a41[smallmode_idx] = (phi1_12 - phi2_12)/2. - self._a43[smallmode_idx] = phi2_12/2. - self._a51[smallmode_idx] = 3.0*(phi1_34 - 3.*phi2_34/4.)/4.0 - self._a52[smallmode_idx] = -3*phi1_34/8. - self._a54[smallmode_idx] = 9*phi2_34/16. - self._a61[smallmode_idx] = (-77*phi1_1+59*phi2_1)/42. - self._a62[smallmode_idx] = 8*phi1_1/7. - self._a63[smallmode_idx] = (111*phi1_1-87*phi2_1)/28. - self._a65[smallmode_idx] = (-47*phi1_1+143*phi2_1)/84. - self._a71[smallmode_idx] = 7*(257*phi1_1 - 497*phi2_1 + 270*phi3_1)/2700 + zs = z[smallmode_idx] # z small + r = self.R * np.exp(2j * np.pi * np.arange(0.5, self.M) / self.M) + rr, zz = np.meshgrid(r, zs) + Z = zz + rr + + phi1_14 = h * np.sum(phi1(Z / 4), axis=1) / self.M + phi2_14 = h * np.sum(phi2(Z / 4), axis=1) / self.M + phi1_12 = h * np.sum(phi1(Z / 2), axis=1) / self.M + phi2_12 = h * np.sum(phi2(Z / 2), axis=1) / self.M + phi1_34 = h * np.sum(phi1(3 * Z / 4), axis=1) / self.M + phi2_34 = h * np.sum(phi2(3 * Z / 4), axis=1) / self.M + phi1_1 = h * np.sum(phi1(Z), axis=1) / self.M + phi2_1 = h * np.sum(phi2(Z), axis=1) / self.M + phi3_1 = h * np.sum(phi3(Z), axis=1) / self.M + + self._a21[smallmode_idx] = phi1_14 / 4.0 + self._a31[smallmode_idx] = (phi1_14 - phi2_14 / 2.0) / 4.0 + self._a32[smallmode_idx] = phi2_14 / 8.0 + self._a41[smallmode_idx] = (phi1_12 - phi2_12) / 2.0 + self._a43[smallmode_idx] = phi2_12 / 2.0 + self._a51[smallmode_idx] = 3.0 * (phi1_34 - 3.0 * phi2_34 / 4.0) / 4.0 + self._a52[smallmode_idx] = -3 * phi1_34 / 8.0 + self._a54[smallmode_idx] = 9 * phi2_34 / 16.0 + self._a61[smallmode_idx] = (-77 * phi1_1 + 59 * phi2_1) / 42.0 + self._a62[smallmode_idx] = 8 * phi1_1 / 7.0 + self._a63[smallmode_idx] = (111 * phi1_1 - 87 * phi2_1) / 28.0 + self._a65[smallmode_idx] = (-47 * phi1_1 + 143 * phi2_1) / 84.0 + self._a71[smallmode_idx] = ( + 7 * (257 * phi1_1 - 497 * phi2_1 + 270 * phi3_1) / 2700 + ) # Paper has error in a73/b3 phi2 coefficient (states this is 497 but it is actually 467) - self._a73[smallmode_idx] = (1097*phi1_1 - 467*phi2_1 - 150*phi3_1)/1350 - self._a74[smallmode_idx] = 2*(-49*phi1_1 + 199*phi2_1 - 135*phi3_1)/225 - self._a75[smallmode_idx] = (-313*phi1_1 + 883*phi2_1 - 90*phi3_1)/1350 - self._a76[smallmode_idx] = (509*phi1_1 - 2129*phi2_1 + 1830*phi3_1)/2700 + self._a73[smallmode_idx] = (1097 * phi1_1 - 467 * phi2_1 - 150 * phi3_1) / 1350 + self._a74[smallmode_idx] = ( + 2 * (-49 * phi1_1 + 199 * phi2_1 - 135 * phi3_1) / 225 + ) + self._a75[smallmode_idx] = (-313 * phi1_1 + 883 * phi2_1 - 90 * phi3_1) / 1350 + self._a76[smallmode_idx] = (509 * phi1_1 - 2129 * phi2_1 + 1830 * phi3_1) / 2700 - def _new_step_init(self,u): + def _new_step_init(self, u): # Need to recompute NL1 if step was accepted (if not accepted use previously computed value) self._NL1 = self.NLfunc(u) - - def _updateStages(self,u,h,accept): - # One passthrough of RK solver for diagonal system + + def _updateStages(self, u, h, accept): + # One passthrough of RK solver for diagonal system if accept: self._new_step_init(u) - self._k = self._EL14*u + self._a21*self._NL1 + self._k = self._EL14 * u + self._a21 * self._NL1 self._NL2 = self.NLfunc(self._k) - self._k = self._EL14*u + self._a31*self._NL1 + self._a32*self._NL2 + self._k = self._EL14 * u + self._a31 * self._NL1 + self._a32 * self._NL2 self._NL3 = self.NLfunc(self._k) - self._k = self._EL12*u + self._a41*self._NL1 + self._a43*self._NL3 + self._k = self._EL12 * u + self._a41 * self._NL1 + self._a43 * self._NL3 self._NL4 = self.NLfunc(self._k) - self._k = self._EL34*u + self._a51*self._NL1 + self._a52*(self._NL2 - self._NL3) \ - + self._a54*self._NL4 + self._k = ( + self._EL34 * u + + self._a51 * self._NL1 + + self._a52 * (self._NL2 - self._NL3) + + self._a54 * self._NL4 + ) self._NL5 = self.NLfunc(self._k) - self._k = self._EL*u + self._a61*self._NL1 + self._a62*(self._NL2 - 3*self._NL4/2.0) \ - + self._a63*self._NL3 + self._a65*self._NL5 + self._k = ( + self._EL * u + + self._a61 * self._NL1 + + self._a62 * (self._NL2 - 3 * self._NL4 / 2.0) + + self._a63 * self._NL3 + + self._a65 * self._NL5 + ) self._NL6 = self.NLfunc(self._k) - self._k = self._EL*u + self._a71*self._NL1 + self._a73*self._NL3 + self._a74*self._NL4 \ - + self._a75*self._NL5 + self._a76*self._NL6 - self._err = self._a75*(-self._NL1 + 4*self._NL3 - 6*self._NL4 + 4*self._NL5 - self._NL6) + self._k = ( + self._EL * u + + self._a71 * self._NL1 + + self._a73 * self._NL3 + + self._a74 * self._NL4 + + self._a75 * self._NL5 + + self._a76 * self._NL6 + ) + self._err = self._a75 * ( + -self._NL1 + 4 * self._NL3 - 6 * self._NL4 + 4 * self._NL5 - self._NL6 + ) return self._k, self._err + class _ETD35_Diagonalized(_ETD35_Diagonal): """ - ETD35 non-diagonal system with eigenvector diagonalization strategy for ETD35 solver + ETD35 non-diagonal system with eigenvector diagonalization strategy for ETD35 solver """ - def __init__(self,linop,NLfunc,contourM,contourR,modecutoff): - super().__init__(linop,NLfunc,contourM,contourR,modecutoff) + + def __init__(self, linop, NLfunc, contourM, contourR, modecutoff): + super().__init__(linop, NLfunc, contourM, contourR, modecutoff) if len(linop.shape) == 1: - raise Exception('Cannot diagonalize a 1D system') + raise Exception("Cannot diagonalize a 1D system") linop_cond = np.linalg.cond(linop) if linop_cond > 1e16: - raise Exception('Cannot diagonalize a non-invertible linear operator L') + raise Exception("Cannot diagonalize a non-invertible linear operator L") if linop_cond > 1000: - print('''Warning: linear matrix array has a large condition number of {:.2f}, - method may be unstable'''.format(linop_cond)) + print( + """Warning: linear matrix array has a large condition number of {:.2f}, + method may be unstable""".format( + linop_cond + ) + ) self._eig_vals, self._S = np.linalg.eig(linop) self._Sinv = np.linalg.inv(self._S) self._v = np.zeros(linop.shape[0]) - - def _updateCoeffs(self,h): - z = h*self._eig_vals - self._updateCoeffsDiagonal(h,z) - - def _new_step_init(self,u): + + def _updateCoeffs(self, h): + z = h * self._eig_vals + self._updateCoeffsDiagonal(h, z) + + def _new_step_init(self, u): self._NL1 = self._Sinv.dot(self.NLfunc(u)) self._v = self._Sinv.dot(u) - - def _updateStages(self,u,h,accept): - # One passthrough of RK solver for diagonalized non-diagonal system + + def _updateStages(self, u, h, accept): + # One passthrough of RK solver for diagonalized non-diagonal system if accept: self._new_step_init(u) - self._k = self._EL14*self._v + self._a21*self._NL1 + self._k = self._EL14 * self._v + self._a21 * self._NL1 self._NL2 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._k = self._EL14*self._v + self._a31*self._NL1 + self._a32*self._NL2 + self._k = self._EL14 * self._v + self._a31 * self._NL1 + self._a32 * self._NL2 self._NL3 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._k = self._EL12*self._v + self._a41*self._NL1 + self._a43*self._NL3 + self._k = self._EL12 * self._v + self._a41 * self._NL1 + self._a43 * self._NL3 self._NL4 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._k = self._EL34*self._v + self._a51*self._NL1 + self._a52*(self._NL2 - self._NL3) \ - + self._a54*self._NL4 + self._k = ( + self._EL34 * self._v + + self._a51 * self._NL1 + + self._a52 * (self._NL2 - self._NL3) + + self._a54 * self._NL4 + ) self._NL5 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._k = self._EL*self._v + self._a61*self._NL1 + self._a62*(self._NL2 - 3*self._NL4/2.0) \ - + self._a63*self._NL3 + self._a65*self._NL5 + self._k = ( + self._EL * self._v + + self._a61 * self._NL1 + + self._a62 * (self._NL2 - 3 * self._NL4 / 2.0) + + self._a63 * self._NL3 + + self._a65 * self._NL5 + ) self._NL6 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._k = self._EL*self._v + self._a71*self._NL1 + self._a73*self._NL3 + self._a74*self._NL4 \ - + self._a75*self._NL5 + self._a76*self._NL6 - self._err = self._a75*(-self._NL1 + 4*self._NL3 - 6*self._NL4 + 4*self._NL5 - self._NL6) + self._k = ( + self._EL * self._v + + self._a71 * self._NL1 + + self._a73 * self._NL3 + + self._a74 * self._NL4 + + self._a75 * self._NL5 + + self._a76 * self._NL6 + ) + self._err = self._a75 * ( + -self._NL1 + 4 * self._NL3 - 6 * self._NL4 + 4 * self._NL5 - self._NL6 + ) return self._S.dot(self._k), self._err - + + class _ETD35_NonDiagonal: """ - ETD35 non-diagonal system strategy for ETD35 solver + ETD35 non-diagonal system strategy for ETD35 solver """ - def __init__(self,linop,NLfunc,contourM,contourR): + + def __init__(self, linop, NLfunc, contourM, contourR): self.linop = linop self.NLfunc = NLfunc self.M = contourM self.R = contourR - + N = linop.shape[0] - self._EL14, self._EL12, self._EL34, self._EL = [np.zeros(shape=linop.shape,dtype=np.complex128) for _ in range(4)] - self._a21, self._a31, self._a32, self._a41, self._a43, self._a51,\ - self._a52,self._a54, self._a61, self._a62, self._a63,\ - self._a64, self._a65, self._a71, self._a73, self._a74,\ - self._a75, self._a76 = [np.zeros(shape=linop.shape,dtype=np.complex128) for _ in range(18)] - self._NL1, self._NL2, self._NL3,self._NL4, self._NL5,\ - self._NL6 = [np.zeros(N,dtype=np.complex128) for _ in range(6)] - self._k = np.zeros(N,dtype=np.complex128) - self._err = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): + self._EL14, self._EL12, self._EL34, self._EL = [ + np.zeros(shape=linop.shape, dtype=np.complex128) for _ in range(4) + ] + ( + self._a21, + self._a31, + self._a32, + self._a41, + self._a43, + self._a51, + self._a52, + self._a54, + self._a61, + self._a62, + self._a63, + self._a64, + self._a65, + self._a71, + self._a73, + self._a74, + self._a75, + self._a76, + ) = [np.zeros(shape=linop.shape, dtype=np.complex128) for _ in range(18)] + self._NL1, self._NL2, self._NL3, self._NL4, self._NL5, self._NL6 = [ + np.zeros(N, dtype=np.complex128) for _ in range(6) + ] + self._k = np.zeros(N, dtype=np.complex128) + self._err = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): # Update RK coefficients for 'matrix' np.array operator linop - z = h*self.linop + z = h * self.linop # Use expm matrix exponential function from scipy - self._EL14 = expm(z/4.0) - self._EL12 = expm(z/2.0) - self._EL34 = expm(3.0*z/4.0) + self._EL14 = expm(z / 4.0) + self._EL12 = expm(z / 2.0) + self._EL34 = expm(3.0 * z / 4.0) self._EL = expm(z) # Use contour integral evaluation for psi etd functions - contour_points = self.R*np.exp(2j*np.pi*np.arange(0.5,self.M)/self.M) - phi1_14,phi2_14,phi1_12,phi2_12,phi1_34,phi2_34 = [\ - np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(6)] - phi1_1,phi2_1,phi3_1 = [np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(3)] + contour_points = self.R * np.exp(2j * np.pi * np.arange(0.5, self.M) / self.M) + phi1_14, phi2_14, phi1_12, phi2_12, phi1_34, phi2_34 = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(6) + ] + phi1_1, phi2_1, phi3_1 = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(3) + ] for point in contour_points: - Q14 = np.linalg.inv(point*np.eye(*self.linop.shape)-z/4.) - Q12 = np.linalg.inv(point*np.eye(*self.linop.shape)-z/2.) - Q34 = np.linalg.inv(point*np.eye(*self.linop.shape)-3*z/4.) - Q = np.linalg.inv(point*np.eye(*self.linop.shape)-z) - phi1_14 += point*phi1(point)*Q14/self.M - phi2_14 += point*phi2(point)*Q14/self.M - phi1_12 += point*phi1(point)*Q12/self.M - phi2_12 += point*phi2(point)*Q12/self.M - phi1_34 += point*phi1(point)*Q34/self.M - phi2_34 += point*phi2(point)*Q34/self.M - phi1_1 += point*phi1(point)*Q/self.M - phi2_1 += point*phi2(point)*Q/self.M - phi3_1 += point*phi3(point)*Q/self.M - - - self._a21 = h*phi1_14/4. - self._a31 = h*(phi1_14-phi2_14/2.)/4. - self._a32 = h*phi2_14/8. - self._a41 = h*(phi1_12 - phi2_12)/2. - self._a43 = h*phi2_12/2. - self._a51 = h*3.0*(phi1_34 - 3.*phi2_34/4.)/4.0 - self._a52 = -3*h*phi1_34/8. - self._a54 = h*9*phi2_34/16. - self._a61 = h*(-77*phi1_1+59*phi2_1)/42. - self._a62 = h*8*phi1_1/7. - self._a63 = h*(111*phi1_1-87*phi2_1)/28. - self._a65 = h*(-47*phi1_1+143*phi2_1)/84. - self._a71 = h*7*(257*phi1_1 - 497*phi2_1 + 270*phi3_1)/2700 - self._a73 = h*(1097*phi1_1 - 467*phi2_1 - 150*phi3_1)/1350 - self._a74 = h*2*(-49*phi1_1 + 199*phi2_1 - 135*phi3_1)/225 - self._a75 = h*(-313*phi1_1 + 883*phi2_1 - 90*phi3_1)/1350 - self._a76 = h*(509*phi1_1 - 2129*phi2_1 + 1830*phi3_1)/2700 - - def _new_step_init(self,u): + Q14 = np.linalg.inv(point * np.eye(*self.linop.shape) - z / 4.0) + Q12 = np.linalg.inv(point * np.eye(*self.linop.shape) - z / 2.0) + Q34 = np.linalg.inv(point * np.eye(*self.linop.shape) - 3 * z / 4.0) + Q = np.linalg.inv(point * np.eye(*self.linop.shape) - z) + phi1_14 += point * phi1(point) * Q14 / self.M + phi2_14 += point * phi2(point) * Q14 / self.M + phi1_12 += point * phi1(point) * Q12 / self.M + phi2_12 += point * phi2(point) * Q12 / self.M + phi1_34 += point * phi1(point) * Q34 / self.M + phi2_34 += point * phi2(point) * Q34 / self.M + phi1_1 += point * phi1(point) * Q / self.M + phi2_1 += point * phi2(point) * Q / self.M + phi3_1 += point * phi3(point) * Q / self.M + + self._a21 = h * phi1_14 / 4.0 + self._a31 = h * (phi1_14 - phi2_14 / 2.0) / 4.0 + self._a32 = h * phi2_14 / 8.0 + self._a41 = h * (phi1_12 - phi2_12) / 2.0 + self._a43 = h * phi2_12 / 2.0 + self._a51 = h * 3.0 * (phi1_34 - 3.0 * phi2_34 / 4.0) / 4.0 + self._a52 = -3 * h * phi1_34 / 8.0 + self._a54 = h * 9 * phi2_34 / 16.0 + self._a61 = h * (-77 * phi1_1 + 59 * phi2_1) / 42.0 + self._a62 = h * 8 * phi1_1 / 7.0 + self._a63 = h * (111 * phi1_1 - 87 * phi2_1) / 28.0 + self._a65 = h * (-47 * phi1_1 + 143 * phi2_1) / 84.0 + self._a71 = h * 7 * (257 * phi1_1 - 497 * phi2_1 + 270 * phi3_1) / 2700 + self._a73 = h * (1097 * phi1_1 - 467 * phi2_1 - 150 * phi3_1) / 1350 + self._a74 = h * 2 * (-49 * phi1_1 + 199 * phi2_1 - 135 * phi3_1) / 225 + self._a75 = h * (-313 * phi1_1 + 883 * phi2_1 - 90 * phi3_1) / 1350 + self._a76 = h * (509 * phi1_1 - 2129 * phi2_1 + 1830 * phi3_1) / 2700 + + def _new_step_init(self, u): self._NL1 = self.NLfunc(u) - - def _updateStages(self,u,h,accept): + + def _updateStages(self, u, h, accept): if accept: self._new_step_init(u) - + self._k = self._EL14.dot(u) + self._a21.dot(self._NL1) self._NL2 = self.NLfunc(self._k) - self._k = self._EL14.dot(u) + self._a31.dot(self._NL1) + self._a32.dot(self._NL2) + self._k = ( + self._EL14.dot(u) + self._a31.dot(self._NL1) + self._a32.dot(self._NL2) + ) self._NL3 = self.NLfunc(self._k) - self._k = self._EL12.dot(u) + self._a41.dot(self._NL1) + self._a43.dot(self._NL3) + self._k = ( + self._EL12.dot(u) + self._a41.dot(self._NL1) + self._a43.dot(self._NL3) + ) self._NL4 = self.NLfunc(self._k) - self._k = self._EL34.dot(u) + self._a51.dot(self._NL1) + self._a52.dot(self._NL2 - self._NL3) \ - + self._a54.dot(self._NL4) + self._k = ( + self._EL34.dot(u) + + self._a51.dot(self._NL1) + + self._a52.dot(self._NL2 - self._NL3) + + self._a54.dot(self._NL4) + ) self._NL5 = self.NLfunc(self._k) - self._k = self._EL.dot(u) + self._a61.dot(self._NL1) + self._a62.dot(self._NL2 - 3*self._NL4/2.0) \ - + self._a63.dot(self._NL3) + self._a65.dot(self._NL5) + self._k = ( + self._EL.dot(u) + + self._a61.dot(self._NL1) + + self._a62.dot(self._NL2 - 3 * self._NL4 / 2.0) + + self._a63.dot(self._NL3) + + self._a65.dot(self._NL5) + ) self._NL6 = self.NLfunc(self._k) - self._k = self._EL.dot(u) + self._a71.dot(self._NL1) + self._a73.dot(self._NL3) + self._a74.dot(self._NL4) \ - + self._a75.dot(self._NL5) + self._a76.dot(self._NL6) - self._err = self._a75.dot(-self._NL1 + 4*self._NL3 - 6*self._NL4 + 4*self._NL5 - self._NL6) + self._k = ( + self._EL.dot(u) + + self._a71.dot(self._NL1) + + self._a73.dot(self._NL3) + + self._a74.dot(self._NL4) + + self._a75.dot(self._NL5) + + self._a76.dot(self._NL6) + ) + self._err = self._a75.dot( + -self._NL1 + 4 * self._NL3 - 6 * self._NL4 + 4 * self._NL5 - self._NL6 + ) return self._k, self._err @@ -283,30 +397,41 @@ class ETD35(ETDAS): linop : np.array NLfunc : function - t : time-array stored with evolve function call - u : output-array stored with evolve function call + t : time-array stored with evolve function call + u : output-array stored with evolve function call logs : array of info stored related to the solver ETD Parameters (see ETDAS class in etd module) ______________ modecutoff : float contour_points : int - contour_radius : float + contour_radius : float StiffSolverAS Parameters (see StiffSolverAS class in solver module) ________________________ epsilon : float incrF : float - decrF : float + decrF : float safetyF : float adapt_cutoff : float - minh : float + minh : float """ - def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray],epsilon : float = 1e-4,incrF : float = 1.25,\ - decrF : float = 0.85, safetyF : float = 0.8, adapt_cutoff : float = 0.01,\ - minh : float = 1e-16, modecutoff : float = 0.01, contour_points : int = 32, - contour_radius : float = 1.0, diagonalize : bool = False): + def __init__( + self, + linop: np.ndarray, + NLfunc: Callable[[np.ndarray], np.ndarray], + epsilon: float = 1e-4, + incrF: float = 1.25, + decrF: float = 0.85, + safetyF: float = 0.8, + adapt_cutoff: float = 0.01, + minh: float = 1e-16, + modecutoff: float = 0.01, + contour_points: int = 32, + contour_radius: float = 1.0, + diagonalize: bool = False, + ): """ INPUTS ______ @@ -315,61 +440,74 @@ def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray], Linear operator (L) in the system dtU = LU + NL(U). Can be either a 2D numpy array (matrix) or a 1D array (diagonal system). L can be either real-valued or complex-valued. - NLfunc : function + NLfunc : function Nonlinear function (NL(U)) in the system dtU = LU + NL(U). Can be a complex or real-valued function. diagonalize : bool, optional - Diagonalize the linear operator (matrix) and solve the diagonalized system + Diagonalize the linear operator (matrix) and solve the diagonalized system ETDAS variables: modecutoff, contour_points, contour_radius (see ETDAS documentation from etd module) - StiffSolverAS variables: epsilon, incrF, decrF, safetyF, adapt_cutoff, minh + StiffSolverAS variables: epsilon, incrF, decrF, safetyF, adapt_cutoff, minh (see StiffSolverAS documentation from solver module) """ - super().__init__(linop,NLfunc,epsilon=epsilon,incrF=incrF,decrF=decrF,\ - safetyF=safetyF,adapt_cutoff=adapt_cutoff,minh=minh,\ - modecutoff=modecutoff,contour_points=contour_points,\ - contour_radius=contour_radius) - self._method = None + super().__init__( + linop, + NLfunc, + epsilon=epsilon, + incrF=incrF, + decrF=decrF, + safetyF=safetyF, + adapt_cutoff=adapt_cutoff, + minh=minh, + modecutoff=modecutoff, + contour_points=contour_points, + contour_radius=contour_radius, + ) + self._method = Union[_ETD35_Diagonal, _ETD35_Diagonalized] if self._diag: - self._method = _ETD35_Diagonal(linop,NLfunc,self.contour_points,\ - self.contour_radius,self.modecutoff) + self._method = _ETD35_Diagonal( + linop, NLfunc, self.contour_points, self.contour_radius, self.modecutoff + ) else: if diagonalize: - self._method = _ETD35_Diagonalized(linop,NLfunc,self.contour_points,\ - self.contour_radius,self.modecutoff) + self._method = _ETD35_Diagonalized( + linop, + NLfunc, + self.contour_points, + self.contour_radius, + self.modecutoff, + ) else: - self._method = _ETD35_NonDiagonal(linop,NLfunc,self.contour_points,\ - self.contour_radius) + self._method = _ETD35_NonDiagonal( + linop, NLfunc, self.contour_points, self.contour_radius + ) self._stages_init = False self._accept = False def _reset(self): - #Resets solver to its initial state + # Resets solver to its initial state self._stages_init = False self._h_coeff = None self._accept = False - - def _updateCoeffs(self,h): + + def _updateCoeffs(self, h): # Update coefficients if step size h changed if h == self._h_coeff: return self._h_coeff = h self._method._updateCoeffs(h) self.logs.append("ETD35 coefficients updated") - - def _updateStages(self,u,h): + + def _updateStages(self, u, h): # Computes u_{n+1} from u_{n} through one RK passthrough self._updateCoeffs(h) if not self._stages_init: self._method._new_step_init(u) self._stages_init = True - return self._method._updateStages(u,h,self._accept) - + return self._method._updateStages(u, h, self._accept) + def _q(self): # Order variable for computing suggested step size (embedded order + 1) - return 4 - - - + return 4 diff --git a/rkstiff/etd4.py b/rkstiff/etd4.py index e7ce837..73e16f5 100644 --- a/rkstiff/etd4.py +++ b/rkstiff/etd4.py @@ -1,142 +1,169 @@ - -from rkstiff.etd import ETDCS,phi1,phi2,phi3 +from rkstiff.etd import ETDCS, phi1, phi2, phi3 import numpy as np -from typing import Callable +from typing import Callable, Union from scipy.linalg import expm + class _ETD4_Diagonal: - """ ETD4 diagonal system strategy for ETD4 solver """ - def __init__(self,linop,NLfunc,contourM,contourR,modecutoff): + """ETD4 diagonal system strategy for ETD4 solver""" + + def __init__(self, linop, NLfunc, contourM, contourR, modecutoff): self.linop = linop self.NLfunc = NLfunc self.M = contourM self.R = contourR self.modecutoff = modecutoff - + N = linop.shape[0] - self._EL, self._EL2 = [np.zeros(N,dtype=np.complex128) for _ in range(2)] - self._a21, self._a31, self._a32, self._a41, self._a43 = [np.zeros(N,dtype=np.complex128) for _ in range(5)] - self._b1, self._b2, self._b4 = [np.zeros(N,dtype=np.complex128) for _ in range(3)] - self._NL1, self._NL2, self._NL3,self._NL4 = [np.zeros(N,dtype=np.complex128) for _ in range(4)] - self._k = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): - z = h*self.linop + self._EL, self._EL2 = [np.zeros(N, dtype=np.complex128) for _ in range(2)] + self._a21, self._a31, self._a32, self._a41, self._a43 = [ + np.zeros(N, dtype=np.complex128) for _ in range(5) + ] + self._b1, self._b2, self._b4 = [ + np.zeros(N, dtype=np.complex128) for _ in range(3) + ] + self._NL1, self._NL2, self._NL3, self._NL4 = [ + np.zeros(N, dtype=np.complex128) for _ in range(4) + ] + self._k = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): + z = h * self.linop self._EL = np.exp(z) - self._EL2 = np.exp(z/2) - + self._EL2 = np.exp(z / 2) + smallmode_idx = np.abs(z) < self.modecutoff - zb = z[~smallmode_idx] # z big + zb = z[~smallmode_idx] # z big # compute big mode coeffs - phi1_12 = h*phi1(zb/2) - phi2_12 = h*phi2(zb/2) - phi1_1 = h*phi1(zb) - phi2_1 = h*phi2(zb) - phi3_1 = h*phi3(zb) - - self._a21[~smallmode_idx] = 0.5*phi1_12 - self._a31[~smallmode_idx] = 0.5*(phi1_12-phi2_12) - self._a32[~smallmode_idx] = 0.5*phi2_12 + phi1_12 = h * phi1(zb / 2) + phi2_12 = h * phi2(zb / 2) + phi1_1 = h * phi1(zb) + phi2_1 = h * phi2(zb) + phi3_1 = h * phi3(zb) + + self._a21[~smallmode_idx] = 0.5 * phi1_12 + self._a31[~smallmode_idx] = 0.5 * (phi1_12 - phi2_12) + self._a32[~smallmode_idx] = 0.5 * phi2_12 self._a41[~smallmode_idx] = phi1_1 - phi2_1 self._a43[~smallmode_idx] = phi2_1 - self._b1[~smallmode_idx] = phi1_1 - (3.0/2)*phi2_1+(2.0/3)*phi3_1 - self._b2[~smallmode_idx] = phi2_1 - (2.0/3)*phi3_1 - self._b4[~smallmode_idx] = -(1.0/2)*phi2_1 + (2.0/3)*phi3_1 - + self._b1[~smallmode_idx] = phi1_1 - (3.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1 + self._b2[~smallmode_idx] = phi2_1 - (2.0 / 3) * phi3_1 + self._b4[~smallmode_idx] = -(1.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1 + # compute small mode coeffs - zs = z[smallmode_idx] # z small - r = self.R*np.exp(2j*np.pi*np.arange(0.5,self.M)/self.M) - rr, zz = np.meshgrid(r,zs) - Z = zz+rr - - phi1_12 = h*np.sum(phi1(Z/2),axis=1)/self.M - phi2_12 = h*np.sum(phi2(Z/2),axis=1)/self.M - phi1_1 = h*np.sum(phi1(Z),axis=1)/self.M - phi2_1 = h*np.sum(phi2(Z),axis=1)/self.M - phi3_1 = h*np.sum(phi3(Z),axis=1)/self.M - - self._a21[smallmode_idx] = 0.5*phi1_12 - self._a31[smallmode_idx] = 0.5*(phi1_12-phi2_12) - self._a32[smallmode_idx] = 0.5*phi2_12 + zs = z[smallmode_idx] # z small + r = self.R * np.exp(2j * np.pi * np.arange(0.5, self.M) / self.M) + rr, zz = np.meshgrid(r, zs) + Z = zz + rr + + phi1_12 = h * np.sum(phi1(Z / 2), axis=1) / self.M + phi2_12 = h * np.sum(phi2(Z / 2), axis=1) / self.M + phi1_1 = h * np.sum(phi1(Z), axis=1) / self.M + phi2_1 = h * np.sum(phi2(Z), axis=1) / self.M + phi3_1 = h * np.sum(phi3(Z), axis=1) / self.M + + self._a21[smallmode_idx] = 0.5 * phi1_12 + self._a31[smallmode_idx] = 0.5 * (phi1_12 - phi2_12) + self._a32[smallmode_idx] = 0.5 * phi2_12 self._a41[smallmode_idx] = phi1_1 - phi2_1 self._a43[smallmode_idx] = phi2_1 - self._b1[smallmode_idx] = phi1_1 - (3.0/2)*phi2_1+(2.0/3)*phi3_1 - self._b2[smallmode_idx] = phi2_1 - (2.0/3)*phi3_1 - self._b4[smallmode_idx] = -(1.0/2)*phi2_1 + (2.0/3)*phi3_1 - - def _N1_init(self,u): + self._b1[smallmode_idx] = phi1_1 - (3.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1 + self._b2[smallmode_idx] = phi2_1 - (2.0 / 3) * phi3_1 + self._b4[smallmode_idx] = -(1.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1 + + def _N1_init(self, u): self._NL1 = self.NLfunc(u) - - def _updateStages(self,u): - # Use First is same as last principle (FSAL) - self._k = self._EL2*u + self._a21*self._NL1 + + def _updateStages(self, u): + # Use First is same as last principle (FSAL) + self._k = self._EL2 * u + self._a21 * self._NL1 self._NL2 = self.NLfunc(self._k) - self._k = self._EL2*u + self._a31*self._NL1 + self._a32*self._NL2 + self._k = self._EL2 * u + self._a31 * self._NL1 + self._a32 * self._NL2 self._NL3 = self.NLfunc(self._k) - self._k = self._EL*u + self._a41*self._NL1 + self._a43*self._NL3 + self._k = self._EL * u + self._a41 * self._NL1 + self._a43 * self._NL3 self._NL4 = self.NLfunc(self._k) - self._k = self._EL*u + self._b1*self._NL1 + self._b2*(self._NL2+self._NL3) \ - + self._b4*self._NL4 + self._k = ( + self._EL * u + + self._b1 * self._NL1 + + self._b2 * (self._NL2 + self._NL3) + + self._b4 * self._NL4 + ) self._NL1 = self.NLfunc(self._k) return self._k + class _ETD4_NonDiagonal: - """ ETD4 non-diagonal system strategy for ETD4 solver """ - def __init__(self,linop,NLfunc,contourM,contourR): + """ETD4 non-diagonal system strategy for ETD4 solver""" + + def __init__(self, linop, NLfunc, contourM, contourR): self.linop = linop self.NLfunc = NLfunc self.M = contourM self.R = contourR - + N = linop.shape[0] - self._EL, self._EL2 = [np.zeros(shape=linop.shape,dtype=np.complex128) for _ in range(2)] - self._a21, self._a31, self._a32, self._a41, self._a43 = [\ - np.zeros(shape=linop.shape,dtype=np.complex128) for _ in range(5)] - self._b1,self._b2, self._b4 = [np.zeros(shape=linop.shape,dtype=np.complex128) for _ in range(3)] - self._NL1, self._NL2, self._NL3,self._NL4 = [np.zeros(N,dtype=np.complex128) for _ in range(4)] - self._k = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): - z = h*self.linop + self._EL, self._EL2 = [ + np.zeros(shape=linop.shape, dtype=np.complex128) for _ in range(2) + ] + self._a21, self._a31, self._a32, self._a41, self._a43 = [ + np.zeros(shape=linop.shape, dtype=np.complex128) for _ in range(5) + ] + self._b1, self._b2, self._b4 = [ + np.zeros(shape=linop.shape, dtype=np.complex128) for _ in range(3) + ] + self._NL1, self._NL2, self._NL3, self._NL4 = [ + np.zeros(N, dtype=np.complex128) for _ in range(4) + ] + self._k = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): + z = h * self.linop self._EL = expm(z) - self._EL2 = expm(z/2) + self._EL2 = expm(z / 2) - contour_points = self.R*np.exp(2j*np.pi*np.arange(0.5,self.M)/self.M) - - phi1_12, phi2_12,phi1_1,phi2_1,phi3_1 = [np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(5)] + contour_points = self.R * np.exp(2j * np.pi * np.arange(0.5, self.M) / self.M) + + phi1_12, phi2_12, phi1_1, phi2_1, phi3_1 = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(5) + ] for point in contour_points: - Q = np.linalg.inv(point*np.eye(*self.linop.shape)-z) - Q2 = np.linalg.inv(point*np.eye(*self.linop.shape)-z/2) - phi1_12 += point*phi1(point)*Q2/self.M - phi2_12 += point*phi2(point)*Q2/self.M - phi1_1 += point*phi1(point)*Q/self.M - phi2_1 += point*phi2(point)*Q/self.M - phi3_1 += point*phi3(point)*Q/self.M - - self._a21 = 0.5*h*phi1_12 - self._a31 = 0.5*h*(phi1_12-phi2_12) - self._a32 = 0.5*h*phi2_12 - self._a41 = h*(phi1_1 - phi2_1) - self._a43 = h*phi2_1 - self._b1 = h*(phi1_1 - (3.0/2)*phi2_1+(2.0/3)*phi3_1) - self._b2 = h*(phi2_1 - (2.0/3)*phi3_1) - self._b4 = h*(-(1.0/2)*phi2_1 + (2.0/3)*phi3_1) - - def _N1_init(self,u): + Q = np.linalg.inv(point * np.eye(*self.linop.shape) - z) + Q2 = np.linalg.inv(point * np.eye(*self.linop.shape) - z / 2) + phi1_12 += point * phi1(point) * Q2 / self.M + phi2_12 += point * phi2(point) * Q2 / self.M + phi1_1 += point * phi1(point) * Q / self.M + phi2_1 += point * phi2(point) * Q / self.M + phi3_1 += point * phi3(point) * Q / self.M + + self._a21 = 0.5 * h * phi1_12 + self._a31 = 0.5 * h * (phi1_12 - phi2_12) + self._a32 = 0.5 * h * phi2_12 + self._a41 = h * (phi1_1 - phi2_1) + self._a43 = h * phi2_1 + self._b1 = h * (phi1_1 - (3.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1) + self._b2 = h * (phi2_1 - (2.0 / 3) * phi3_1) + self._b4 = h * (-(1.0 / 2) * phi2_1 + (2.0 / 3) * phi3_1) + + def _N1_init(self, u): self._NL1 = self.NLfunc(u) - - def _updateStages(self,u): + + def _updateStages(self, u): self._k = self._EL2.dot(u) + self._a21.dot(self._NL1) self._NL2 = self.NLfunc(self._k) self._k = self._EL2.dot(u) + self._a31.dot(self._NL1) + self._a32.dot(self._NL2) self._NL3 = self.NLfunc(self._k) self._k = self._EL.dot(u) + self._a41.dot(self._NL1) + self._a43.dot(self._NL3) self._NL4 = self.NLfunc(self._k) - self._k = self._EL.dot(u) + self._b1.dot(self._NL1) + self._b2.dot(self._NL2+self._NL3) \ - + self._b4.dot(self._NL4) - self._NL1 = self.NLfunc(self._k) # Use First is same as last principle (FSAL) + self._k = ( + self._EL.dot(u) + + self._b1.dot(self._NL1) + + self._b2.dot(self._NL2 + self._NL3) + + self._b4.dot(self._NL4) + ) + self._NL1 = self.NLfunc(self._k) # Use First is same as last principle (FSAL) return self._k + class ETD4(ETDCS): """ Exponential time-differencing constant step solver of 4th order (Krogstad) @@ -146,20 +173,26 @@ class ETD4(ETDCS): linop : np.array NLfunc : function - t : time-array stored with evolve function call - u : output-array stored with evolve function call + t : time-array stored with evolve function call + u : output-array stored with evolve function call logs : array of info stored related to the solver ETD Parameters (see ETDAS class in etd module) ______________ modecutoff : float contour_points : int - contour_radius : float + contour_radius : float """ - - def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray],\ - modecutoff : float = 0.01, contour_points : int = 32,contour_radius : float = 1.0): + + def __init__( + self, + linop: np.ndarray, + NLfunc: Callable[[np.ndarray], np.ndarray], + modecutoff: float = 0.01, + contour_points: int = 32, + contour_radius: float = 1.0, + ): """ INPUTS ______ @@ -168,44 +201,50 @@ def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray], Linear operator (L) in the system dtU = LU + NL(U). Can be either a 2D numpy array (matrix) or a 1D array (diagonal system). L can be either real-valued or complex-valued. - NLfunc : function + NLfunc : function Nonlinear function (NL(U)) in the system dtU = LU + NL(U). Can be a complex or real-valued function. diagonalize : bool, optional - Diagonalize the linear operator (matrix) and solve the diagonalized system + Diagonalize the linear operator (matrix) and solve the diagonalized system ETDCS variables: modecutoff, contour_points, contour_radius (see ETDAS documentation from etd module) """ - super().__init__(linop,NLfunc,modecutoff=modecutoff,\ - contour_points=contour_points,contour_radius=contour_radius) - self._method = None + super().__init__( + linop, + NLfunc, + modecutoff=modecutoff, + contour_points=contour_points, + contour_radius=contour_radius, + ) + self._method = Union[_ETD4_Diagonal, _ETD4_NonDiagonal] if self._diag: - self._method = _ETD4_Diagonal(linop,NLfunc,self.contour_points,\ - self.contour_radius,self.modecutoff) + self._method = _ETD4_Diagonal( + linop, NLfunc, self.contour_points, self.contour_radius, self.modecutoff + ) else: - self._method = _ETD4_NonDiagonal(linop,NLfunc,self.contour_points,\ - self.contour_radius) + self._method = _ETD4_NonDiagonal( + linop, NLfunc, self.contour_points, self.contour_radius + ) self.__N1_init = False def _reset(self): - """ Resets solver to its initial state """ + """Resets solver to its initial state""" self.__N1_init = False self._h_coeff = None - def _updateCoeffs(self,h): - """ Update coefficients if step size h changed """ + def _updateCoeffs(self, h): + """Update coefficients if step size h changed""" if h == self._h_coeff: return self._h_coeff = h self._method._updateCoeffs(h) self.logs.append("ETD4 coefficients updated") - - def _updateStages(self,u,h): - """ Computes u_{n+1} from u_{n} in one RK passthrough """ + + def _updateStages(self, u, h): + """Computes u_{n+1} from u_{n} in one RK passthrough""" self._updateCoeffs(h) if not self.__N1_init: self._method._N1_init(u) self.__N1_init = True return self._method._updateStages(u) - diff --git a/rkstiff/etd5.py b/rkstiff/etd5.py index fc9e90b..cc386b8 100644 --- a/rkstiff/etd5.py +++ b/rkstiff/etd5.py @@ -1,218 +1,294 @@ - -from rkstiff.etd import ETDCS,phi1,phi2,phi3 +from rkstiff.etd import ETDCS, phi1, phi2, phi3 import numpy as np -from typing import Callable +from typing import Callable, Union from scipy.linalg import expm + class _ETD5_Diagonal: """ - ETD5 diagonal system strategy for ETD5 solver + ETD5 diagonal system strategy for ETD5 solver """ - def __init__(self,linop,NLfunc,contourM,contourR,modecutoff): + + def __init__(self, linop, NLfunc, contourM, contourR, modecutoff): self.linop = linop self.NLfunc = NLfunc self.M = contourM self.R = contourR self.modecutoff = modecutoff - + N = linop.shape[0] - self._EL14, self._EL12, self._EL34, self._EL = [\ - np.zeros(N,dtype=np.complex128) for _ in range(4)] - self._a21, self._a31, self._a32, self._a41, self._a43, self._a51,\ - self._a52,self._a54, self._a61, self._a62, self._a63,\ - self._a64, self._a65 = [np.zeros(N,dtype=np.complex128) for _ in range(13)] - self._b1, self._b3, self._b4,self._b5,self._b6 = [\ - np.zeros(N,dtype=np.complex128) for _ in range(5)] - self._NL1, self._NL2, self._NL3,self._NL4, self._NL5,\ - self._NL6 = [np.zeros(N,dtype=np.complex128) for _ in range(6)] - self._k = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): - z = h*self.linop + self._EL14, self._EL12, self._EL34, self._EL = [ + np.zeros(N, dtype=np.complex128) for _ in range(4) + ] + ( + self._a21, + self._a31, + self._a32, + self._a41, + self._a43, + self._a51, + self._a52, + self._a54, + self._a61, + self._a62, + self._a63, + self._a64, + self._a65, + ) = [np.zeros(N, dtype=np.complex128) for _ in range(13)] + self._b1, self._b3, self._b4, self._b5, self._b6 = [ + np.zeros(N, dtype=np.complex128) for _ in range(5) + ] + self._NL1, self._NL2, self._NL3, self._NL4, self._NL5, self._NL6 = [ + np.zeros(N, dtype=np.complex128) for _ in range(6) + ] + self._k = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): + z = h * self.linop # diagonal system -> L is 1D array (of independent modes) - self._EL14 = np.exp(z/4.0) - self._EL12 = np.exp(z/2.0) - self._EL34 = np.exp(3.0*z/4.0) + self._EL14 = np.exp(z / 4.0) + self._EL12 = np.exp(z / 2.0) + self._EL34 = np.exp(3.0 * z / 4.0) self._EL = np.exp(z) smallmode_idx = np.abs(z) < self.modecutoff # compute big mode coeffs - zb = z[~smallmode_idx] # z big - phi1_14 = h*phi1(zb/4) - phi2_14 = h*phi2(zb/4) - phi1_12 = h*phi1(zb/2) - phi2_12 = h*phi2(zb/2) - phi1_34 = h*phi1(3*zb/4) - phi2_34 = h*phi2(3*zb/4) - phi1_1 = h*phi1(zb) - phi2_1 = h*phi2(zb) - phi3_1 = h*phi3(zb) - - self._a21[~smallmode_idx] = phi1_14/4. - self._a31[~smallmode_idx] = (phi1_14-phi2_14/2.)/4. - self._a32[~smallmode_idx] = phi2_14/8. - self._a41[~smallmode_idx] = (phi1_12 - phi2_12)/2. - self._a43[~smallmode_idx] = phi2_12/2. - self._a51[~smallmode_idx] = 3.0*(phi1_34 - 3.*phi2_34/4.)/4.0 - self._a52[~smallmode_idx] = -3*phi1_34/8. - self._a54[~smallmode_idx] = 9*phi2_34/16. - self._a61[~smallmode_idx] = (-77*phi1_1+59*phi2_1)/42. - self._a62[~smallmode_idx] = 8*phi1_1/7. - self._a63[~smallmode_idx] = (111*phi1_1-87*phi2_1)/28. - self._a65[~smallmode_idx] = (-47*phi1_1+143*phi2_1)/84. - self._b1[~smallmode_idx] = 7*(257*phi1_1 - 497*phi2_1 + 270*phi3_1)/2700 + zb = z[~smallmode_idx] # z big + phi1_14 = h * phi1(zb / 4) + phi2_14 = h * phi2(zb / 4) + phi1_12 = h * phi1(zb / 2) + phi2_12 = h * phi2(zb / 2) + phi1_34 = h * phi1(3 * zb / 4) + phi2_34 = h * phi2(3 * zb / 4) + phi1_1 = h * phi1(zb) + phi2_1 = h * phi2(zb) + phi3_1 = h * phi3(zb) + + self._a21[~smallmode_idx] = phi1_14 / 4.0 + self._a31[~smallmode_idx] = (phi1_14 - phi2_14 / 2.0) / 4.0 + self._a32[~smallmode_idx] = phi2_14 / 8.0 + self._a41[~smallmode_idx] = (phi1_12 - phi2_12) / 2.0 + self._a43[~smallmode_idx] = phi2_12 / 2.0 + self._a51[~smallmode_idx] = 3.0 * (phi1_34 - 3.0 * phi2_34 / 4.0) / 4.0 + self._a52[~smallmode_idx] = -3 * phi1_34 / 8.0 + self._a54[~smallmode_idx] = 9 * phi2_34 / 16.0 + self._a61[~smallmode_idx] = (-77 * phi1_1 + 59 * phi2_1) / 42.0 + self._a62[~smallmode_idx] = 8 * phi1_1 / 7.0 + self._a63[~smallmode_idx] = (111 * phi1_1 - 87 * phi2_1) / 28.0 + self._a65[~smallmode_idx] = (-47 * phi1_1 + 143 * phi2_1) / 84.0 + self._b1[~smallmode_idx] = ( + 7 * (257 * phi1_1 - 497 * phi2_1 + 270 * phi3_1) / 2700 + ) # Paper has error in b3 phi2 coefficient (states this is 497 but it is actually 467) - self._b3[~smallmode_idx] = (1097*phi1_1 - 467*phi2_1 - 150*phi3_1)/1350 - self._b4[~smallmode_idx] = 2*(-49*phi1_1 + 199*phi2_1 - 135*phi3_1)/225 - self._b5[~smallmode_idx] = (-313*phi1_1 + 883*phi2_1 - 90*phi3_1)/1350 - self._b6[~smallmode_idx] = (509*phi1_1 - 2129*phi2_1 + 1830*phi3_1)/2700 - + self._b3[~smallmode_idx] = (1097 * phi1_1 - 467 * phi2_1 - 150 * phi3_1) / 1350 + self._b4[~smallmode_idx] = ( + 2 * (-49 * phi1_1 + 199 * phi2_1 - 135 * phi3_1) / 225 + ) + self._b5[~smallmode_idx] = (-313 * phi1_1 + 883 * phi2_1 - 90 * phi3_1) / 1350 + self._b6[~smallmode_idx] = (509 * phi1_1 - 2129 * phi2_1 + 1830 * phi3_1) / 2700 + # compute small mode coeffs - zs = z[smallmode_idx] # z small - r = self.R*np.exp(2j*np.pi*np.arange(0.5,self.M)/self.M) - rr, zz = np.meshgrid(r,zs) - Z = zz+rr - - phi1_14 = h*np.sum(phi1(Z/4),axis=1)/self.M - phi2_14 = h*np.sum(phi2(Z/4),axis=1)/self.M - phi1_12 = h*np.sum(phi1(Z/2),axis=1)/self.M - phi2_12 = h*np.sum(phi2(Z/2),axis=1)/self.M - phi1_34 = h*np.sum(phi1(3*Z/4),axis=1)/self.M - phi2_34 = h*np.sum(phi2(3*Z/4),axis=1)/self.M - phi1_1 = h*np.sum(phi1(Z),axis=1)/self.M - phi2_1 = h*np.sum(phi2(Z),axis=1)/self.M - phi3_1 = h*np.sum(phi3(Z),axis=1)/self.M - - self._a21[smallmode_idx] = phi1_14/4. - self._a31[smallmode_idx] = (phi1_14-phi2_14/2.)/4. - self._a32[smallmode_idx] = phi2_14/8. - self._a41[smallmode_idx] = (phi1_12 - phi2_12)/2. - self._a43[smallmode_idx] = phi2_12/2. - self._a51[smallmode_idx] = 3.0*(phi1_34 - 3.*phi2_34/4.)/4.0 - self._a52[smallmode_idx] = -3*phi1_34/8. - self._a54[smallmode_idx] = 9*phi2_34/16. - self._a61[smallmode_idx] = (-77*phi1_1+59*phi2_1)/42. - self._a62[smallmode_idx] = 8*phi1_1/7. - self._a63[smallmode_idx] = (111*phi1_1-87*phi2_1)/28. - self._a65[smallmode_idx] = (-47*phi1_1+143*phi2_1)/84. - self._b1[smallmode_idx] = 7*(257*phi1_1 - 497*phi2_1 + 270*phi3_1)/2700 + zs = z[smallmode_idx] # z small + r = self.R * np.exp(2j * np.pi * np.arange(0.5, self.M) / self.M) + rr, zz = np.meshgrid(r, zs) + Z = zz + rr + + phi1_14 = h * np.sum(phi1(Z / 4), axis=1) / self.M + phi2_14 = h * np.sum(phi2(Z / 4), axis=1) / self.M + phi1_12 = h * np.sum(phi1(Z / 2), axis=1) / self.M + phi2_12 = h * np.sum(phi2(Z / 2), axis=1) / self.M + phi1_34 = h * np.sum(phi1(3 * Z / 4), axis=1) / self.M + phi2_34 = h * np.sum(phi2(3 * Z / 4), axis=1) / self.M + phi1_1 = h * np.sum(phi1(Z), axis=1) / self.M + phi2_1 = h * np.sum(phi2(Z), axis=1) / self.M + phi3_1 = h * np.sum(phi3(Z), axis=1) / self.M + + self._a21[smallmode_idx] = phi1_14 / 4.0 + self._a31[smallmode_idx] = (phi1_14 - phi2_14 / 2.0) / 4.0 + self._a32[smallmode_idx] = phi2_14 / 8.0 + self._a41[smallmode_idx] = (phi1_12 - phi2_12) / 2.0 + self._a43[smallmode_idx] = phi2_12 / 2.0 + self._a51[smallmode_idx] = 3.0 * (phi1_34 - 3.0 * phi2_34 / 4.0) / 4.0 + self._a52[smallmode_idx] = -3 * phi1_34 / 8.0 + self._a54[smallmode_idx] = 9 * phi2_34 / 16.0 + self._a61[smallmode_idx] = (-77 * phi1_1 + 59 * phi2_1) / 42.0 + self._a62[smallmode_idx] = 8 * phi1_1 / 7.0 + self._a63[smallmode_idx] = (111 * phi1_1 - 87 * phi2_1) / 28.0 + self._a65[smallmode_idx] = (-47 * phi1_1 + 143 * phi2_1) / 84.0 + self._b1[smallmode_idx] = ( + 7 * (257 * phi1_1 - 497 * phi2_1 + 270 * phi3_1) / 2700 + ) # Paper has error in b3 phi2 coefficient (states this is 497 but it is actually 467) - self._b3[smallmode_idx] = (1097*phi1_1 - 467*phi2_1 - 150*phi3_1)/1350 - self._b4[smallmode_idx] = 2*(-49*phi1_1 + 199*phi2_1 - 135*phi3_1)/225 - self._b5[smallmode_idx] = (-313*phi1_1 + 883*phi2_1 - 90*phi3_1)/1350 - self._b6[smallmode_idx] = (509*phi1_1 - 2129*phi2_1 + 1830*phi3_1)/2700 + self._b3[smallmode_idx] = (1097 * phi1_1 - 467 * phi2_1 - 150 * phi3_1) / 1350 + self._b4[smallmode_idx] = 2 * (-49 * phi1_1 + 199 * phi2_1 - 135 * phi3_1) / 225 + self._b5[smallmode_idx] = (-313 * phi1_1 + 883 * phi2_1 - 90 * phi3_1) / 1350 + self._b6[smallmode_idx] = (509 * phi1_1 - 2129 * phi2_1 + 1830 * phi3_1) / 2700 - def _N1_init(self,u): - """ Initialize N1 before first call to updateStages """ + def _N1_init(self, u): + """Initialize N1 before first call to updateStages""" self._NL1 = self.NLfunc(u) - - def _updateStages(self,u): - """ One passthrough of RK solver for diagonal system """ - self._k = self._EL14*u + self._a21*self._NL1 + + def _updateStages(self, u): + """One passthrough of RK solver for diagonal system""" + self._k = self._EL14 * u + self._a21 * self._NL1 self._NL2 = self.NLfunc(self._k) - self._k = self._EL14*u + self._a31*self._NL1 + self._a32*self._NL2 + self._k = self._EL14 * u + self._a31 * self._NL1 + self._a32 * self._NL2 self._NL3 = self.NLfunc(self._k) - self._k = self._EL12*u + self._a41*self._NL1 + self._a43*self._NL3 + self._k = self._EL12 * u + self._a41 * self._NL1 + self._a43 * self._NL3 self._NL4 = self.NLfunc(self._k) - self._k = self._EL34*u + self._a51*self._NL1 + self._a52*(self._NL2 - self._NL3) \ - + self._a54*self._NL4 + self._k = ( + self._EL34 * u + + self._a51 * self._NL1 + + self._a52 * (self._NL2 - self._NL3) + + self._a54 * self._NL4 + ) self._NL5 = self.NLfunc(self._k) - self._k = self._EL*u + self._a61*self._NL1 + self._a62*(self._NL2 - 3*self._NL4/2.0) \ - + self._a63*self._NL3 + self._a65*self._NL5 - self._NL6 = self.NLfunc(self._k) - self._k = self._EL*u + self._b1*self._NL1 + self._b3*self._NL3 + self._b4*self._NL4 \ - + self._b5*self._NL5 + self._b6*self._NL6 - self._NL1 = self._NL6 # FSAL principle + self._k = ( + self._EL * u + + self._a61 * self._NL1 + + self._a62 * (self._NL2 - 3 * self._NL4 / 2.0) + + self._a63 * self._NL3 + + self._a65 * self._NL5 + ) + self._NL6 = self.NLfunc(self._k) + self._k = ( + self._EL * u + + self._b1 * self._NL1 + + self._b3 * self._NL3 + + self._b4 * self._NL4 + + self._b5 * self._NL5 + + self._b6 * self._NL6 + ) + self._NL1 = self._NL6 # FSAL principle return self._k + class _ETD5_NonDiagonal: - """ ETD5 non-diagonal system strategy for ETD5 solver """ - def __init__(self,linop,NLfunc,contourM,contourR): + """ETD5 non-diagonal system strategy for ETD5 solver""" + + def __init__(self, linop, NLfunc, contourM, contourR): self.linop = linop self.NLfunc = NLfunc self.M = contourM self.R = contourR - + N = linop.shape[0] - self._EL14, self._EL12, self._EL34, self._EL = [np.zeros(shape=linop.shape,dtype=np.complex128) for _ in range(4)] - self._a21, self._a31, self._a32, self._a41, self._a43, self._a51,\ - self._a52,self._a54, self._a61, self._a62, self._a63,\ - self._a64, self._a65 = [np.zeros(shape=linop.shape,dtype=np.complex128) for _ in range(13)] - self._b1, self._b3, self._b4,self._b5, self._b6 = [\ - np.zeros(shape=linop.shape,dtype=np.complex128) for _ in range(5)] - self._NL1, self._NL2, self._NL3,self._NL4, self._NL5,\ - self._NL6 = [np.zeros(N,dtype=np.complex128) for _ in range(6)] - self._k = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): + self._EL14, self._EL12, self._EL34, self._EL = [ + np.zeros(shape=linop.shape, dtype=np.complex128) for _ in range(4) + ] + ( + self._a21, + self._a31, + self._a32, + self._a41, + self._a43, + self._a51, + self._a52, + self._a54, + self._a61, + self._a62, + self._a63, + self._a64, + self._a65, + ) = [np.zeros(shape=linop.shape, dtype=np.complex128) for _ in range(13)] + self._b1, self._b3, self._b4, self._b5, self._b6 = [ + np.zeros(shape=linop.shape, dtype=np.complex128) for _ in range(5) + ] + self._NL1, self._NL2, self._NL3, self._NL4, self._NL5, self._NL6 = [ + np.zeros(N, dtype=np.complex128) for _ in range(6) + ] + self._k = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): # Update RK coefficients for 'matrix' np.array operator linop - z = h*self.linop + z = h * self.linop # Use expm matrix exponential function from scipy - self._EL14 = expm(z/4.0) - self._EL12 = expm(z/2.0) - self._EL34 = expm(3.0*z/4.0) + self._EL14 = expm(z / 4.0) + self._EL12 = expm(z / 2.0) + self._EL34 = expm(3.0 * z / 4.0) self._EL = expm(z) # Use contour integral evaluation for psi etd functions - contour_points = self.R*np.exp(2j*np.pi*np.arange(0.5,self.M)/self.M) - phi1_14,phi2_14,phi1_12,phi2_12,phi1_34,phi2_34 = [\ - np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(6)] - phi1_1,phi2_1,phi3_1 = [np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(3)] + contour_points = self.R * np.exp(2j * np.pi * np.arange(0.5, self.M) / self.M) + phi1_14, phi2_14, phi1_12, phi2_12, phi1_34, phi2_34 = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(6) + ] + phi1_1, phi2_1, phi3_1 = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(3) + ] for point in contour_points: - Q14 = np.linalg.inv(point*np.eye(*self.linop.shape)-z/4.) - Q12 = np.linalg.inv(point*np.eye(*self.linop.shape)-z/2.) - Q34 = np.linalg.inv(point*np.eye(*self.linop.shape)-3*z/4.) - Q = np.linalg.inv(point*np.eye(*self.linop.shape)-z) - phi1_14 += point*phi1(point)*Q14/self.M - phi2_14 += point*phi2(point)*Q14/self.M - phi1_12 += point*phi1(point)*Q12/self.M - phi2_12 += point*phi2(point)*Q12/self.M - phi1_34 += point*phi1(point)*Q34/self.M - phi2_34 += point*phi2(point)*Q34/self.M - phi1_1 += point*phi1(point)*Q/self.M - phi2_1 += point*phi2(point)*Q/self.M - phi3_1 += point*phi3(point)*Q/self.M - - - self._a21 = h*phi1_14/4. - self._a31 = h*(phi1_14-phi2_14/2.)/4. - self._a32 = h*phi2_14/8. - self._a41 = h*(phi1_12 - phi2_12)/2. - self._a43 = h*phi2_12/2. - self._a51 = h*3.0*(phi1_34 - 3.*phi2_34/4.)/4.0 - self._a52 = -3*h*phi1_34/8. - self._a54 = h*9*phi2_34/16. - self._a61 = h*(-77*phi1_1+59*phi2_1)/42. - self._a62 = h*8*phi1_1/7. - self._a63 = h*(111*phi1_1-87*phi2_1)/28. - self._a65 = h*(-47*phi1_1+143*phi2_1)/84. - self._b1 = h*7*(257*phi1_1 - 497*phi2_1 + 270*phi3_1)/2700 - self._b3 = h*(1097*phi1_1 - 467*phi2_1 - 150*phi3_1)/1350 - self._b4 = h*2*(-49*phi1_1 + 199*phi2_1 - 135*phi3_1)/225 - self._b5 = h*(-313*phi1_1 + 883*phi2_1 - 90*phi3_1)/1350 - self._b6 = h*(509*phi1_1 - 2129*phi2_1 + 1830*phi3_1)/2700 - - def _N1_init(self,u): - """ Initialize N1 before first call to updateStages """ + Q14 = np.linalg.inv(point * np.eye(*self.linop.shape) - z / 4.0) + Q12 = np.linalg.inv(point * np.eye(*self.linop.shape) - z / 2.0) + Q34 = np.linalg.inv(point * np.eye(*self.linop.shape) - 3 * z / 4.0) + Q = np.linalg.inv(point * np.eye(*self.linop.shape) - z) + phi1_14 += point * phi1(point) * Q14 / self.M + phi2_14 += point * phi2(point) * Q14 / self.M + phi1_12 += point * phi1(point) * Q12 / self.M + phi2_12 += point * phi2(point) * Q12 / self.M + phi1_34 += point * phi1(point) * Q34 / self.M + phi2_34 += point * phi2(point) * Q34 / self.M + phi1_1 += point * phi1(point) * Q / self.M + phi2_1 += point * phi2(point) * Q / self.M + phi3_1 += point * phi3(point) * Q / self.M + + self._a21 = h * phi1_14 / 4.0 + self._a31 = h * (phi1_14 - phi2_14 / 2.0) / 4.0 + self._a32 = h * phi2_14 / 8.0 + self._a41 = h * (phi1_12 - phi2_12) / 2.0 + self._a43 = h * phi2_12 / 2.0 + self._a51 = h * 3.0 * (phi1_34 - 3.0 * phi2_34 / 4.0) / 4.0 + self._a52 = -3 * h * phi1_34 / 8.0 + self._a54 = h * 9 * phi2_34 / 16.0 + self._a61 = h * (-77 * phi1_1 + 59 * phi2_1) / 42.0 + self._a62 = h * 8 * phi1_1 / 7.0 + self._a63 = h * (111 * phi1_1 - 87 * phi2_1) / 28.0 + self._a65 = h * (-47 * phi1_1 + 143 * phi2_1) / 84.0 + self._b1 = h * 7 * (257 * phi1_1 - 497 * phi2_1 + 270 * phi3_1) / 2700 + self._b3 = h * (1097 * phi1_1 - 467 * phi2_1 - 150 * phi3_1) / 1350 + self._b4 = h * 2 * (-49 * phi1_1 + 199 * phi2_1 - 135 * phi3_1) / 225 + self._b5 = h * (-313 * phi1_1 + 883 * phi2_1 - 90 * phi3_1) / 1350 + self._b6 = h * (509 * phi1_1 - 2129 * phi2_1 + 1830 * phi3_1) / 2700 + + def _N1_init(self, u): + """Initialize N1 before first call to updateStages""" self._NL1 = self.NLfunc(u) - - def _updateStages(self,u): + + def _updateStages(self, u): self._k = self._EL14.dot(u) + self._a21.dot(self._NL1) self._NL2 = self.NLfunc(self._k) - self._k = self._EL14.dot(u) + self._a31.dot(self._NL1) + self._a32.dot(self._NL2) + self._k = ( + self._EL14.dot(u) + self._a31.dot(self._NL1) + self._a32.dot(self._NL2) + ) self._NL3 = self.NLfunc(self._k) - self._k = self._EL12.dot(u) + self._a41.dot(self._NL1) + self._a43.dot(self._NL3) + self._k = ( + self._EL12.dot(u) + self._a41.dot(self._NL1) + self._a43.dot(self._NL3) + ) self._NL4 = self.NLfunc(self._k) - self._k = self._EL34.dot(u) + self._a51.dot(self._NL1) + self._a52.dot(self._NL2 - self._NL3) \ - + self._a54.dot(self._NL4) + self._k = ( + self._EL34.dot(u) + + self._a51.dot(self._NL1) + + self._a52.dot(self._NL2 - self._NL3) + + self._a54.dot(self._NL4) + ) self._NL5 = self.NLfunc(self._k) - self._k = self._EL.dot(u) + self._a61.dot(self._NL1) + self._a62.dot(self._NL2 - 3*self._NL4/2.0) \ - + self._a63.dot(self._NL3) + self._a65.dot(self._NL5) + self._k = ( + self._EL.dot(u) + + self._a61.dot(self._NL1) + + self._a62.dot(self._NL2 - 3 * self._NL4 / 2.0) + + self._a63.dot(self._NL3) + + self._a65.dot(self._NL5) + ) self._NL6 = self.NLfunc(self._k) - self._k = self._EL.dot(u) + self._b1.dot(self._NL1) + self._b3.dot(self._NL3) + self._b4.dot(self._NL4) \ - + self._b5.dot(self._NL5) + self._b6.dot(self._NL6) - self._NL1 = self._NL6 # FSAL principle + self._k = ( + self._EL.dot(u) + + self._b1.dot(self._NL1) + + self._b3.dot(self._NL3) + + self._b4.dot(self._NL4) + + self._b5.dot(self._NL5) + + self._b6.dot(self._NL6) + ) + self._NL1 = self._NL6 # FSAL principle return self._k @@ -225,21 +301,26 @@ class ETD5(ETDCS): linop : np.array NLfunc : function - t : time-array stored with evolve function call - u : output-array stored with evolve function call + t : time-array stored with evolve function call + u : output-array stored with evolve function call logs : array of info stored related to the solver ETD Parameters (see ETDAS class in etd module) ______________ modecutoff : float contour_points : int - contour_radius : float + contour_radius : float """ - - def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray],\ - modecutoff : float = 0.01, contour_points : int = 32,\ - contour_radius : float = 1.0): + + def __init__( + self, + linop: np.ndarray, + NLfunc: Callable[[np.ndarray], np.ndarray], + modecutoff: float = 0.01, + contour_points: int = 32, + contour_radius: float = 1.0, + ): """ INPUTS ______ @@ -248,44 +329,50 @@ def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray], Linear operator (L) in the system dtU = LU + NL(U). Can be either a 2D numpy array (matrix) or a 1D array (diagonal system). L can be either real-valued or complex-valued. - NLfunc : function + NLfunc : function Nonlinear function (NL(U)) in the system dtU = LU + NL(U). Can be a complex or real-valued function. diagonalize : bool, optional - Diagonalize the linear operator (matrix) and solve the diagonalized system + Diagonalize the linear operator (matrix) and solve the diagonalized system ETDCS variables: modecutoff, contour_points, contour_radius (see ETDAS documentation from etd module) """ - super().__init__(linop,NLfunc,modecutoff=modecutoff,\ - contour_points=contour_points,contour_radius=contour_radius) - self._method = None + super().__init__( + linop, + NLfunc, + modecutoff=modecutoff, + contour_points=contour_points, + contour_radius=contour_radius, + ) + self._method = Union[_ETD5_Diagonal, _ETD5_NonDiagonal] if self._diag: - self._method = _ETD5_Diagonal(linop,NLfunc,self.contour_points,\ - self.contour_radius,self.modecutoff) + self._method = _ETD5_Diagonal( + linop, NLfunc, self.contour_points, self.contour_radius, self.modecutoff + ) else: - self._method = _ETD5_NonDiagonal(linop,NLfunc,self.contour_points,\ - self.contour_radius) + self._method = _ETD5_NonDiagonal( + linop, NLfunc, self.contour_points, self.contour_radius + ) self.__N1_init = False def _reset(self): - """ Resets solver to its initial state """ + """Resets solver to its initial state""" self.__N1_init = False self._h_coeff = None - def _updateCoeffs(self,h): - """ Update coefficients if step size h changed """ + def _updateCoeffs(self, h): + """Update coefficients if step size h changed""" if h == self._h_coeff: return self._h_coeff = h self._method._updateCoeffs(h) self.logs.append("ETD5 coefficients updated") - - def _updateStages(self,u,h): - """ Computes u_{n+1} from u_{n} in one RK passthrough """ + + def _updateStages(self, u, h): + """Computes u_{n+1} from u_{n} in one RK passthrough""" self._updateCoeffs(h) if not self.__N1_init: self._method._N1_init(u) self.__N1_init = True return self._method._updateStages(u) - diff --git a/rkstiff/grids.py b/rkstiff/grids.py index 14c2324..19b306e 100644 --- a/rkstiff/grids.py +++ b/rkstiff/grids.py @@ -1,9 +1,10 @@ - import numpy as np -import scipy.special as sp +import scipy.special as sp # type: ignore +from typing import Optional + -def construct_x_kx_rfft(N : int,a : float=0.0,b :float=2*np.pi): - """ Constructs a uniform 1D spatial grid and rfft spectral wavenumbers for real-valued functions +def construct_x_kx_rfft(N: int, a: float = 0.0, b: float = 2 * np.pi): + """Constructs a uniform 1D spatial grid and rfft spectral wavenumbers for real-valued functions INPUTS N - even integer greater than 2 a - left endpoint in spatial grid @@ -13,21 +14,25 @@ def construct_x_kx_rfft(N : int,a : float=0.0,b :float=2*np.pi): kx - spectral wavenumber grid """ - if not isinstance(N,int): - raise TypeError('Number of grid points N must be an integer, it is {}'.format(N)) + if not isinstance(N, int): + raise TypeError( + "Number of grid points N must be an integer, it is {}".format(N) + ) if N <= 2: - raise ValueError('Number of grid points N must be larger than 2, it is {}'.format(N)) + raise ValueError( + "Number of grid points N must be larger than 2, it is {}".format(N) + ) if (N % 2) != 0: - raise ValueError('Integer N in construct_x_kx_rfft must be an even number') + raise ValueError("Integer N in construct_x_kx_rfft must be an even number") - dx = (b-a)/N - x = np.arange(a,b,dx) - kx = 2*np.pi*np.fft.rfftfreq(N,d=dx) - return x,kx + dx = (b - a) / N + x = np.arange(a, b, dx) + kx = 2 * np.pi * np.fft.rfftfreq(N, d=dx) + return x, kx -def construct_x_kx_fft(N : int,a : float=0.0,b : float=2*np.pi): - """ Constructs a uniform 1D spatial grid and fft spectral wavenumbers for complex-valued functions +def construct_x_kx_fft(N: int, a: float = 0.0, b: float = 2 * np.pi): + """Constructs a uniform 1D spatial grid and fft spectral wavenumbers for complex-valued functions INPUTS N - even integer greater than 2 a - left endpoint in spatial grid @@ -37,20 +42,25 @@ def construct_x_kx_fft(N : int,a : float=0.0,b : float=2*np.pi): kx - spectral wavenumber grid """ - if not isinstance(N,int): - raise TypeError('Number of grid points N must be an integer, it is {}'.format(N)) + if not isinstance(N, int): + raise TypeError( + "Number of grid points N must be an integer, it is {}".format(N) + ) if N <= 2: - raise ValueError('Number of grid points N must be larger than 2, it is {}'.format(N)) + raise ValueError( + "Number of grid points N must be larger than 2, it is {}".format(N) + ) if (N % 2) != 0: - raise ValueError('Integer N in construct_x_kx_rfft must be an even number') + raise ValueError("Integer N in construct_x_kx_rfft must be an even number") + + dx = (b - a) / N + x = np.arange(a, b, dx) + kx = 2 * np.pi * np.fft.fftfreq(N, d=dx) + return x, kx - dx = (b-a)/N - x = np.arange(a,b,dx) - kx = 2*np.pi*np.fft.fftfreq(N,d=dx) - return x,kx -def construct_x_cheb(N : int,a : float=-1.0,b : float=1.0): - """ Constructs a 1D grid with Chebyshev spatial discretization +def construct_x_cheb(N: int, a: float = -1.0, b: float = 1.0): + """Constructs a 1D grid with Chebyshev spatial discretization INPUTS N - positive integer @@ -60,16 +70,23 @@ def construct_x_cheb(N : int,a : float=-1.0,b : float=1.0): x - grid discretized at Chebyshev points """ - if not isinstance(N,int): - raise TypeError('Max Chebyshev grid point number N must be an integer, it is {}'.format(N)) + if not isinstance(N, int): + raise TypeError( + "Max Chebyshev grid point number N must be an integer, it is {}".format(N) + ) if N < 2: - raise ValueError('Max Chebyshev grid point number N must be larger than 1, it is {}'.format(N)) - x = np.polynomial.chebyshev.chebpts2(N+1) - x = a+(b-a)*(x+1)/2. + raise ValueError( + "Max Chebyshev grid point number N must be larger than 1, it is {}".format( + N + ) + ) + x = np.polynomial.chebyshev.chebpts2(N + 1) + x = a + (b - a) * (x + 1) / 2.0 return x -def construct_x_Dx_cheb(N : int,a : float=-1,b : float=1): - """ Constructs a 1D grid with Chebyshev spatial discretization along with a + +def construct_x_Dx_cheb(N: int, a: float = -1, b: float = 1): + """Constructs a 1D grid with Chebyshev spatial discretization along with a differentiation matrix for functions sampled on this grid INPUTS @@ -80,13 +97,15 @@ def construct_x_Dx_cheb(N : int,a : float=-1,b : float=1): x - grid discretized at Chebyshev points Dx - derivative matrix for functions sampled at Chebyshev points """ - x = construct_x_cheb(N,a,b) - c = np.r_[2,np.ones(N-1),2]*np.power(-1,np.arange(0,N+1)) - X = np.tile(x.reshape(N+1,1),(1,N+1)) # put copies of x in columns (first row is x0) + x = construct_x_cheb(N, a, b) + c = np.r_[2, np.ones(N - 1), 2] * np.power(-1, np.arange(0, N + 1)) + X = np.tile( + x.reshape(N + 1, 1), (1, N + 1) + ) # put copies of x in columns (first row is x0) dX = X - X.T - Dx = np.outer(c,1./c)/(dX+np.eye(N+1)) + Dx = np.outer(c, 1.0 / c) / (dX + np.eye(N + 1)) Dx = Dx - np.diag(Dx.sum(axis=1)) - return x,Dx + return x, Dx class HankelTransform: @@ -101,7 +120,7 @@ class HankelTransform: hankel transform matrix is nr x nr. rmax : float - Maximum radius of sampled points. + Maximum radius of sampled points. r : np.array, dtype = float Radial points for a spectral grid suitable for the Hankel transform. @@ -115,7 +134,7 @@ class HankelTransform: _______ ht(f): - Computes a hankel transform of the function f sampled at the radial points + Computes a hankel transform of the function f sampled at the radial points specified by r. iht(g): @@ -124,7 +143,7 @@ class HankelTransform: """ - def __init__(self,nr,rmax=1.0): + def __init__(self, nr, rmax=1.0): """ Constructs a Hankel transform Matrix that is used in the forward hankel transform function ht and inverse hankel transform function iht @@ -148,81 +167,81 @@ def __init__(self,nr,rmax=1.0): self._bessel_zeros = None self._jN = None self._Y = None - self._setnr_setrmax(nr,rmax) + self._setnr_setrmax(nr, rmax) - def _setnr_setrmax(self,nr,rmax): + def _setnr_setrmax(self, nr, rmax): self._setbesselzeros(nr) self._setgrids(rmax) self._sethankelmatrix() self._rmax = rmax - self._nr = nr + self._nr = nr - def _setbesselzeros(self,nr): + def _setbesselzeros(self, nr): # find and save first nr bessel zeros in an array, also save the nr+1 bessel zero - bessel_zeros = sp.jn_zeros(0,nr+1) - self._bessel_zeros,self._jN = bessel_zeros[:-1], bessel_zeros[-1] + bessel_zeros = sp.jn_zeros(0, nr + 1) + self._bessel_zeros, self._jN = bessel_zeros[:-1], bessel_zeros[-1] - def _setgrids(self,rmax): + def _setgrids(self, rmax): # set the r and kr radial grids given a maximum radius rmax - self.r = self._bessel_zeros*rmax/self._jN - self.kr = self._bessel_zeros/rmax + self.r = self._bessel_zeros * rmax / self._jN + self.kr = self._bessel_zeros / rmax def _sethankelmatrix(self): # set the Hankel matrix used in the Hankel transform for the saved radial grid j1vec = sp.j1(self._bessel_zeros) - bessel_arg = np.outer(self._bessel_zeros,self._bessel_zeros)/self._jN - self._Y = 2*sp.j0(bessel_arg)/(self._jN*j1vec**2) + bessel_arg = np.outer(self._bessel_zeros, self._bessel_zeros) / self._jN + self._Y = 2 * sp.j0(bessel_arg) / (self._jN * j1vec**2) def _scalefactor(self): # factor used in transforming from real space to spectral space and vice_versa for internal use - return self.rmax**2/self._jN + return self.rmax**2 / self._jN def hankelMatrix(self): - """ Returns a copy of the Hankel transform matrix constructed by the class. """ + """Returns a copy of the Hankel transform matrix constructed by the class.""" return self._Y.copy() def besselZeros(self): - """ Returns a copy of the Bessel zeros used by the Hankel transform. """ + """Returns a copy of the Bessel zeros used by the Hankel transform.""" return self._bessel_zeros.copy() @property def nr(self): - """ Returns the number of radial points for the radial grid specified by the class. """ + """Returns the number of radial points for the radial grid specified by the class.""" return self._nr @nr.setter - def nr(self,nr): - """ Sets the number of radial points in the grid to be used by the Hankel transform. """ - if not isinstance(nr,int): - raise ValueError('nr must be an integer') - if nr < 4: - raise ValueError('nr must be greater than or equal to 4') + def nr(self, nr): + """Sets the number of radial points in the grid to be used by the Hankel transform.""" + if not isinstance(nr, int): + raise ValueError("nr must be an integer") + if nr < 4: + raise ValueError("nr must be greater than or equal to 4") self._setbesselzeros(nr) self._setgrids(self.rmax) self._sethankelmatrix() - self._nr = nr + self._nr = nr @property def rmax(self): - """ Returns the maximum radius of the radial grid used by the Hankel transform """ + """Returns the maximum radius of the radial grid used by the Hankel transform""" return self._rmax @rmax.setter - def rmax(self,rmax): - """ Sets the maximum radius of the radial grid used by the Hankel transform """ + def rmax(self, rmax): + """Sets the maximum radius of the radial grid used by the Hankel transform""" if rmax <= 0: - raise ValueError('rmax must be non-negative') + raise ValueError("rmax must be non-negative") self._setgrids(rmax) - self._rmax = rmax + self._rmax = rmax - def ht(self,f : np.ndarray) -> np.ndarray: + def ht(self, f: np.ndarray) -> np.ndarray: """ Computes a hankel transform of the function f sampled at the radial points specified by r - + INPUTS ______ - + f : np.array, dtype=float function sampled at the discretized points specified by r @@ -234,16 +253,16 @@ def ht(self,f : np.ndarray) -> np.ndarray: to the spectral space grid points kr """ - return self._scalefactor()*self._Y.dot(f) + return self._scalefactor() * self._Y.dot(f) - def iht(self,g : np.ndarray) -> np.ndarray: + def iht(self, g: np.ndarray) -> np.ndarray: """ Computes an inverse hankel transform of the function g sampled at the spectral space points specified by kr - + INPUTS ______ - + g : np.array, dtype=float spectral space function sampled at the discretized points specified by kr @@ -255,10 +274,10 @@ def iht(self,g : np.ndarray) -> np.ndarray: to values on the radial grid r """ - return self._Y.dot(g)/self._scalefactor() + return self._Y.dot(g) / self._scalefactor() -def mirrorGrid(r : np.ndarray,u : np.ndarray=None,axis : int=-1): +def mirrorGrid(r: np.ndarray, u: Optional[np.ndarray] = None, axis: int = -1): """ Converts r grid from [0,rmax] interval to [-rmax,rmax] interval and adjusts function output u accordingly @@ -272,7 +291,7 @@ def mirrorGrid(r : np.ndarray,u : np.ndarray=None,axis : int=-1): u : np.array function values specified at radial points given by r - axis : int + axis : int axis value determines how to mirror the u array (-1 -> stack horizontally, 0 -> stack vertically) @@ -280,26 +299,24 @@ def mirrorGrid(r : np.ndarray,u : np.ndarray=None,axis : int=-1): _______ rnew : np.array, dtype=float - 'radial' grid on the interval [-rmax,rmax] + 'radial' grid on the interval [-rmax,rmax] unew : np.array function values specified at 'radial' points given by rnew """ - rnew = np.hstack([-np.flipud(r),r]) + rnew = np.hstack([-np.flipud(r), r]) if u is None: return rnew if axis == -1: - unew = np.hstack([np.flipud(u),u]) + unew = np.hstack([np.flipud(u), u]) elif axis == 0: - unew = np.vstack([np.flipud(u),u]) + unew = np.vstack([np.flipud(u), u]) elif axis == 1: - unew = np.hstack([np.fliplr(u),u]) + unew = np.hstack([np.fliplr(u), u]) else: - raise ValueError('axis variable must be -1 or 0') - - return rnew,unew - + raise ValueError("axis variable must be -1 or 0") + return rnew, unew diff --git a/rkstiff/if34.py b/rkstiff/if34.py index 0086d2d..c93fd5e 100644 --- a/rkstiff/if34.py +++ b/rkstiff/if34.py @@ -1,133 +1,161 @@ - from rkstiff.solver import StiffSolverAS import numpy as np -from typing import Callable +from typing import Callable, Union from scipy.linalg import expm + class _IF34_Diagonal: """ - IF34 diagonal system strategy for IF34 solver + IF34 diagonal system strategy for IF34 solver """ - def __init__(self,linop,NLfunc): + + def __init__(self, linop, NLfunc): self.linop = linop self.NLfunc = NLfunc - + N = linop.shape[0] - self._EL, self._EL2 = [np.zeros(N,dtype=np.complex128) for _ in range(2)] - self._NL1, self._NL2, self._NL3,self._NL4, self._NL5 = [np.zeros(N,dtype=np.complex128) for _ in range(5)] - self._k = np.zeros(N,dtype=np.complex128) - self._err = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): - z = h*self.linop - self._updateCoeffsDiagonal(h,z) - - def _updateCoeffsDiagonal(self,h,z): + self._EL, self._EL2 = [np.zeros(N, dtype=np.complex128) for _ in range(2)] + self._NL1, self._NL2, self._NL3, self._NL4, self._NL5 = [ + np.zeros(N, dtype=np.complex128) for _ in range(5) + ] + self._k = np.zeros(N, dtype=np.complex128) + self._err = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): + z = h * self.linop + self._updateCoeffsDiagonal(h, z) + + def _updateCoeffsDiagonal(self, h, z): self._EL = np.exp(z) - self._EL2 = np.exp(z/2) - - def _N1_init(self,u): + self._EL2 = np.exp(z / 2) + + def _N1_init(self, u): self._NL1 = self.NLfunc(u) - - def _updateStages(self,u,h,accept): + + def _updateStages(self, u, h, accept): if accept: self._NL1 = self._NL5.copy() - - self._k = self._EL2*u + h*self._EL2*self._NL1/2.0 + + self._k = self._EL2 * u + h * self._EL2 * self._NL1 / 2.0 self._NL2 = self.NLfunc(self._k) - self._k = self._EL2*u + h*self._NL2/2. + self._k = self._EL2 * u + h * self._NL2 / 2.0 self._NL3 = self.NLfunc(self._k) - self._k = self._EL*u + h*self._EL2*self._NL3 + self._k = self._EL * u + h * self._EL2 * self._NL3 self._NL4 = self.NLfunc(self._k) - self._k = self._EL*u + h*(self._EL*self._NL1/6.0 + self._EL2*self._NL2/3.0 \ - + self._EL2*self._NL3/3.0 + self._NL4/6.0) + self._k = self._EL * u + h * ( + self._EL * self._NL1 / 6.0 + + self._EL2 * self._NL2 / 3.0 + + self._EL2 * self._NL3 / 3.0 + + self._NL4 / 6.0 + ) self._NL5 = self.NLfunc(self._k) - self._err = h*(self._NL4-self._NL5)/6. + self._err = h * (self._NL4 - self._NL5) / 6.0 return self._k, self._err + class _IF34_Diagonalized(_IF34_Diagonal): """ - IF34 non-diagonal system with eigenvector diagonalization strategy for IF34 solver + IF34 non-diagonal system with eigenvector diagonalization strategy for IF34 solver """ - def __init__(self,linop,NLfunc): - super().__init__(linop,NLfunc) + + def __init__(self, linop, NLfunc): + super().__init__(linop, NLfunc) if len(linop.shape) == 1: - raise Exception('Cannot diagonalize a 1D system') + raise Exception("Cannot diagonalize a 1D system") linop_cond = np.linalg.cond(linop) if linop_cond > 1e16: - raise Exception('Cannot diagonalize a non-invertible linear operator L') + raise Exception("Cannot diagonalize a non-invertible linear operator L") if linop_cond > 1000: - print('''Warning: linear matrix array has a large condition number of {:.2f}, - method may be unstable'''.format(linop_cond)) + print( + """Warning: linear matrix array has a large condition number of {:.2f}, + method may be unstable""".format( + linop_cond + ) + ) self._eig_vals, self._S = np.linalg.eig(linop) self._Sinv = np.linalg.inv(self._S) self._v = np.zeros(linop.shape[0]) - - def _updateCoeffs(self,h): - z = h*self._eig_vals - self._updateCoeffsDiagonal(h,z) - - def _N1_init(self,u): + + def _updateCoeffs(self, h): + z = h * self._eig_vals + self._updateCoeffsDiagonal(h, z) + + def _N1_init(self, u): self._NL1 = self._Sinv.dot(self.NLfunc(u)) self._v = self._Sinv.dot(u) - - def _updateStages(self,u,h,accept): + + def _updateStages(self, u, h, accept): # Use First is same as last principle (FSAL) -> k5 stage is input u for next step if accept: - self._NL1 = self._NL5.copy() + self._NL1 = self._NL5.copy() self._v = self._Sinv.dot(u) - self._k = self._EL2*self._v + h*self._EL2*self._NL1/2.0 + self._k = self._EL2 * self._v + h * self._EL2 * self._NL1 / 2.0 self._NL2 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._k = self._EL2*self._v + h*self._NL2/2. + self._k = self._EL2 * self._v + h * self._NL2 / 2.0 self._NL3 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._k = self._EL*self._v + h*self._EL2*self._NL3 + self._k = self._EL * self._v + h * self._EL2 * self._NL3 self._NL4 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._k = self._EL*self._v + h*(self._EL*self._NL1/6.0 + self._EL2*self._NL2/3.0 \ - + self._EL2*self._NL3/3.0 + self._NL4/6.0) + self._k = self._EL * self._v + h * ( + self._EL * self._NL1 / 6.0 + + self._EL2 * self._NL2 / 3.0 + + self._EL2 * self._NL3 / 3.0 + + self._NL4 / 6.0 + ) self._NL5 = self._Sinv.dot(self.NLfunc(self._S.dot(self._k))) - self._err = h*(self._NL4-self._NL5)/6. + self._err = h * (self._NL4 - self._NL5) / 6.0 return self._S.dot(self._k), self._err + class _IF34_NonDiagonal: """ - IF34 non-diagonal system strategy for IF34 solver + IF34 non-diagonal system strategy for IF34 solver """ - def __init__(self,linop,NLfunc): + + def __init__(self, linop, NLfunc): self.linop = linop self.NLfunc = NLfunc - + N = linop.shape[0] - self._EL, self._EL2 = [np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(2)] - self._NL1, self._NL2, self._NL3,self._NL4, self._NL5 = [np.zeros(N,dtype=np.complex128) for _ in range(5)] - self._k = np.zeros(N,dtype=np.complex128) - self._err = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): - z = h*self.linop + self._EL, self._EL2 = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(2) + ] + self._NL1, self._NL2, self._NL3, self._NL4, self._NL5 = [ + np.zeros(N, dtype=np.complex128) for _ in range(5) + ] + self._k = np.zeros(N, dtype=np.complex128) + self._err = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): + z = h * self.linop self._EL = expm(z) - self._EL2 = expm(z/2) - - def _N1_init(self,u): + self._EL2 = expm(z / 2) + + def _N1_init(self, u): self._NL1 = self.NLfunc(u) - - def _updateStages(self,u,h,accept): + + def _updateStages(self, u, h, accept): # Use First is same as last principle (FSAL) -> k5 stage is input u for next step if accept: - self._NL1 = self._NL5.copy() + self._NL1 = self._NL5.copy() - self._k = self._EL2.dot(u) + h*self._EL2.dot(self._NL1/2.0) + self._k = self._EL2.dot(u) + h * self._EL2.dot(self._NL1 / 2.0) self._NL2 = self.NLfunc(self._k) - self._k = self._EL2.dot(u) + h*self._NL2/2. + self._k = self._EL2.dot(u) + h * self._NL2 / 2.0 self._NL3 = self.NLfunc(self._k) - self._k = self._EL.dot(u) + h*self._EL2.dot(self._NL3) + self._k = self._EL.dot(u) + h * self._EL2.dot(self._NL3) self._NL4 = self.NLfunc(self._k) - self._k = self._EL.dot(u) + h*(self._EL.dot(self._NL1/6.0) + self._EL2.dot(self._NL2/3.0) \ - + self._EL2.dot(self._NL3/3.0) + self._NL4/6.0) + self._k = self._EL.dot(u) + h * ( + self._EL.dot(self._NL1 / 6.0) + + self._EL2.dot(self._NL2 / 3.0) + + self._EL2.dot(self._NL3 / 3.0) + + self._NL4 / 6.0 + ) self._NL5 = self.NLfunc(self._k) - self._err = h*(self._NL4-self._NL5)/6. + self._err = h * (self._NL4 - self._NL5) / 6.0 return self._k, self._err + class IF34(StiffSolverAS): """ Integrating factor adaptive step solver of 4th order with 3rd order embedding @@ -136,24 +164,32 @@ class IF34(StiffSolverAS): __________ linop : np.array NLfunc : function - t : time-array stored with evolve function call - u : output-array stored with evolve function call + t : time-array stored with evolve function call + u : output-array stored with evolve function call logs : array of info stored related to the solver StiffSolverAS Parameters (see StiffSolverAS class in solver module) ________________________ epsilon : float incrF : float - decrF : float + decrF : float safetyF : float adapt_cutoff : float - minh : float + minh : float """ - def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray],\ - epsilon : float = 1e-4,incrF : float = 1.25,\ - decrF : float = 0.85, safetyF : float = 0.8, adapt_cutoff : float = 0.01,\ - minh : float = 1e-16, diagonalize : bool = False): + def __init__( + self, + linop: np.ndarray, + NLfunc: Callable[[np.ndarray], np.ndarray], + epsilon: float = 1e-4, + incrF: float = 1.25, + decrF: float = 0.85, + safetyF: float = 0.8, + adapt_cutoff: float = 0.01, + minh: float = 1e-16, + diagonalize: bool = False, + ): """ INPUTS ______ @@ -162,53 +198,58 @@ def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray], Linear operator (L) in the system dtU = LU + NL(U). Can be either a 2D numpy array (matrix) or a 1D array (diagonal system). L can be either real-valued or complex-valued. - NLfunc : function + NLfunc : function Nonlinear function (NL(U)) in the system dtU = LU + NL(U). Can be a complex or real-valued function. diagonalize : bool, optional - Diagonalize the linear operator (matrix) and solve the diagonalized system + Diagonalize the linear operator (matrix) and solve the diagonalized system - StiffSolverAS variables: epsilon, incrF, decrF, safetyF, adapt_cutoff, minh + StiffSolverAS variables: epsilon, incrF, decrF, safetyF, adapt_cutoff, minh (see StiffSolverAS documentation from solver module) """ - super().__init__(linop,NLfunc,epsilon=epsilon,incrF=incrF,decrF=decrF,\ - safetyF=safetyF,adapt_cutoff=adapt_cutoff,minh=minh) - self._method = None + super().__init__( + linop, + NLfunc, + epsilon=epsilon, + incrF=incrF, + decrF=decrF, + safetyF=safetyF, + adapt_cutoff=adapt_cutoff, + minh=minh, + ) + self._method = Union[_IF34_Diagonal, _IF34_Diagonalized, _IF34_NonDiagonal] if self._diag: - self._method = _IF34_Diagonal(linop,NLfunc) + self._method = _IF34_Diagonal(linop, NLfunc) else: if diagonalize: - self._method = _IF34_Diagonalized(linop,NLfunc) + self._method = _IF34_Diagonalized(linop, NLfunc) else: - self._method = _IF34_NonDiagonal(linop,NLfunc) + self._method = _IF34_NonDiagonal(linop, NLfunc) self._reset() def _reset(self): - #Resets solver to its initial state + # Resets solver to its initial state self.__N1_init = False self._h_coeff = None self._accept = False - - def _updateCoeffs(self,h): + + def _updateCoeffs(self, h): # Update coefficients if step size h changed if h == self._h_coeff: return self._h_coeff = h self._method._updateCoeffs(h) self.logs.append("IF34 coefficients updated") - - def _updateStages(self,u,h): + + def _updateStages(self, u, h): # Computes u_{n+1} from u_{n} through one RK passthrough self._updateCoeffs(h) if not self.__N1_init: self._method._N1_init(u) self.__N1_init = True - return self._method._updateStages(u,h,self._accept) - + return self._method._updateStages(u, h, self._accept) + def _q(self): # Order variable for computing suggested step size (embedded order + 1) - return 4 - - - + return 4 diff --git a/rkstiff/if4.py b/rkstiff/if4.py index 6f91046..e708fb8 100644 --- a/rkstiff/if4.py +++ b/rkstiff/if4.py @@ -1,96 +1,114 @@ - from rkstiff.solver import StiffSolverCS import numpy as np -from typing import Callable +from typing import Callable, Union from scipy.linalg import expm + class _IF4_Diagonal: - """ IF4 diagonal system strategy for IF4 solver """ - def __init__(self,linop,NLfunc): + """IF4 diagonal system strategy for IF4 solver""" + + def __init__(self, linop, NLfunc): self.linop = linop self.NLfunc = NLfunc - + N = linop.shape[0] - self._EL, self._EL2 = [np.zeros(N,dtype=np.complex128) for _ in range(2)] - self._NL1, self._NL2, self._NL3,self._NL4 = [np.zeros(N,dtype=np.complex128) for _ in range(4)] - self._k = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): - z = h*self.linop + self._EL, self._EL2 = [np.zeros(N, dtype=np.complex128) for _ in range(2)] + self._NL1, self._NL2, self._NL3, self._NL4 = [ + np.zeros(N, dtype=np.complex128) for _ in range(4) + ] + self._k = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): + z = h * self.linop self._EL = np.exp(z) - self._EL2 = np.exp(z/2) - - def _N1_init(self,u): - """ Need to initialize N1 before first updateStage call """ + self._EL2 = np.exp(z / 2) + + def _N1_init(self, u): + """Need to initialize N1 before first updateStage call""" self._NL1 = self.NLfunc(u) - - def _updateStages(self,u,h): - """ One RK step """ - self._k = self._EL2*u + h*self._EL2*self._NL1/2.0 + + def _updateStages(self, u, h): + """One RK step""" + self._k = self._EL2 * u + h * self._EL2 * self._NL1 / 2.0 self._NL2 = self.NLfunc(self._k) - self._k = self._EL2*u + h*self._NL2/2. + self._k = self._EL2 * u + h * self._NL2 / 2.0 self._NL3 = self.NLfunc(self._k) - self._k = self._EL*u + h*self._EL2*self._NL3 + self._k = self._EL * u + h * self._EL2 * self._NL3 self._NL4 = self.NLfunc(self._k) - self._k = self._EL*u + h*(self._EL*self._NL1/6.0 + self._EL2*self._NL2/3.0 \ - + self._EL2*self._NL3/3.0 + self._NL4/6.0) - self._NL1 = self.NLfunc(self._k) # FSAL principle + self._k = self._EL * u + h * ( + self._EL * self._NL1 / 6.0 + + self._EL2 * self._NL2 / 3.0 + + self._EL2 * self._NL3 / 3.0 + + self._NL4 / 6.0 + ) + self._NL1 = self.NLfunc(self._k) # FSAL principle return self._k + class _IF4_NonDiagonal: - """ IF4 non-diagonal system strategy for IF4 solver """ - def __init__(self,linop,NLfunc): + """IF4 non-diagonal system strategy for IF4 solver""" + + def __init__(self, linop, NLfunc): self.linop = linop self.NLfunc = NLfunc - + N = linop.shape[0] - self._EL, self._EL2 = [np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(2)] - self._NL1, self._NL2, self._NL3,self._NL4 = [np.zeros(N,dtype=np.complex128) for _ in range(4)] - self._k = np.zeros(N,dtype=np.complex128) - - def _updateCoeffs(self,h): - z = h*self.linop + self._EL, self._EL2 = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(2) + ] + self._NL1, self._NL2, self._NL3, self._NL4 = [ + np.zeros(N, dtype=np.complex128) for _ in range(4) + ] + self._k = np.zeros(N, dtype=np.complex128) + + def _updateCoeffs(self, h): + z = h * self.linop self._EL = expm(z) - self._EL2 = expm(z/2) - - def _N1_init(self,u): + self._EL2 = expm(z / 2) + + def _N1_init(self, u): self._NL1 = self.NLfunc(u) - - def _updateStages(self,u,h): - self._k = self._EL2.dot(u) + h*self._EL2.dot(self._NL1/2.0) + + def _updateStages(self, u, h): + self._k = self._EL2.dot(u) + h * self._EL2.dot(self._NL1 / 2.0) self._NL2 = self.NLfunc(self._k) - self._k = self._EL2.dot(u) + h*self._NL2/2. + self._k = self._EL2.dot(u) + h * self._NL2 / 2.0 self._NL3 = self.NLfunc(self._k) - self._k = self._EL.dot(u) + h*self._EL2.dot(self._NL3) + self._k = self._EL.dot(u) + h * self._EL2.dot(self._NL3) self._NL4 = self.NLfunc(self._k) - self._k = self._EL.dot(u) + h*(self._EL.dot(self._NL1/6.0) + self._EL2.dot(self._NL2/3.0) \ - + self._EL2.dot(self._NL3/3.0) + self._NL4/6.0) - self._NL1 = self.NLfunc(self._k) # FSAL principle + self._k = self._EL.dot(u) + h * ( + self._EL.dot(self._NL1 / 6.0) + + self._EL2.dot(self._NL2 / 3.0) + + self._EL2.dot(self._NL3 / 3.0) + + self._NL4 / 6.0 + ) + self._NL1 = self.NLfunc(self._k) # FSAL principle return self._k + class IF4(StiffSolverCS): """ - Integrating factor constant step solver of 4th order + Integrating factor constant step solver of 4th order ATTRIBUTES __________ linop : np.array NLfunc : function - t : time-array stored with evolve function call - u : output-array stored with evolve function call + t : time-array stored with evolve function call + u : output-array stored with evolve function call logs : array of info stored related to the solver StiffSolverAS Parameters (see StiffSolverAS class in solver module) ________________________ epsilon : float incrF : float - decrF : float + decrF : float safetyF : float adapt_cutoff : float - minh : float + minh : float """ - def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray]): + def __init__(self, linop: np.ndarray, NLfunc: Callable[[np.ndarray], np.ndarray]): """ INPUTS ______ @@ -99,37 +117,36 @@ def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray]) Linear operator (L) in the system dtU = LU + NL(U). Can be either a 2D numpy array (matrix) or a 1D array (diagonal system). L can be either real-valued or complex-valued. - NLfunc : function + NLfunc : function Nonlinear function (NL(U)) in the system dtU = LU + NL(U). Can be a complex or real-valued function. """ - super().__init__(linop,NLfunc) - self._method = None + super().__init__(linop, NLfunc) + self._method = Union[_IF4_Diagonal, _IF4_NonDiagonal] if self._diag: - self._method = _IF4_Diagonal(linop,NLfunc) + self._method = _IF4_Diagonal(linop, NLfunc) else: - self._method = _IF4_NonDiagonal(linop,NLfunc) + self._method = _IF4_NonDiagonal(linop, NLfunc) self._reset() def _reset(self): - """ Resets solver to its initial state """ + """Resets solver to its initial state""" self.__N1_init = False self._h_coeff = None - - def _updateCoeffs(self,h): - """ Update coefficients if step size h changed """ + + def _updateCoeffs(self, h): + """Update coefficients if step size h changed""" if h == self._h_coeff: return self._h_coeff = h self._method._updateCoeffs(h) self.logs.append("IF4 coefficients updated") - - def _updateStages(self,u,h): - """ Computes u_{n+1} from u_{n} through one RK passthrough """ + + def _updateStages(self, u, h): + """Computes u_{n+1} from u_{n} through one RK passthrough""" self._updateCoeffs(h) if not self.__N1_init: self._method._N1_init(u) self.__N1_init = True - return self._method._updateStages(u,h) - + return self._method._updateStages(u, h) diff --git a/rkstiff/if45dp.py b/rkstiff/if45dp.py index 0d93db5..c19c4dc 100644 --- a/rkstiff/if45dp.py +++ b/rkstiff/if45dp.py @@ -1,8 +1,7 @@ - import numpy as np from typing import Callable from rkstiff.solver import StiffSolverAS -from scipy.linalg import expm + class IF45DP(StiffSolverAS): """ @@ -13,23 +12,32 @@ class IF45DP(StiffSolverAS): __________ linop : np.array NLfunc : function - t : time-array stored with evolve function call - u : output-array stored with evolve function call + t : time-array stored with evolve function call + u : output-array stored with evolve function call logs : array of info stored related to the solver StiffSolverAS Parameters (see StiffSolverAS class in solver module) ________________________ epsilon : float incrF : float - decrF : float + decrF : float safetyF : float adapt_cutoff : float - minh : float + minh : float """ - def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray],\ - epsilon : float = 1e-4,incrF : float = 1.25,\ - decrF : float = 0.85, safetyF : float = 0.8, adapt_cutoff : float = 0.01,\ - minh : float = 1e-16, diagonalize : bool = False): + + def __init__( + self, + linop: np.ndarray, + NLfunc: Callable[[np.ndarray], np.ndarray], + epsilon: float = 1e-4, + incrF: float = 1.25, + decrF: float = 0.85, + safetyF: float = 0.8, + adapt_cutoff: float = 0.01, + minh: float = 1e-16, + diagonalize: bool = False, + ): """ INPUTS ______ @@ -38,41 +46,69 @@ def __init__(self,linop : np.ndarray,NLfunc : Callable[[np.ndarray],np.ndarray], Linear operator (L) in the system dtU = LU + NL(U). Can be either a 2D numpy array (matrix) or a 1D array (diagonal system). L can be either real-valued or complex-valued. - NLfunc : function + NLfunc : function Nonlinear function (NL(U)) in the system dtU = LU + NL(U). Can be a complex or real-valued function. diagonalize : bool, optional - Diagonalize the linear operator (matrix) and solve the diagonalized system + Diagonalize the linear operator (matrix) and solve the diagonalized system - StiffSolverAS variables: epsilon, incrF, decrF, safetyF, adapt_cutoff, minh + StiffSolverAS variables: epsilon, incrF, decrF, safetyF, adapt_cutoff, minh (see StiffSolverAS documentation from solver module) """ - super().__init__(linop,NLfunc,epsilon=epsilon,incrF=incrF,decrF=decrF,\ - safetyF=safetyF,adapt_cutoff=adapt_cutoff,minh=minh) + super().__init__( + linop, + NLfunc, + epsilon=epsilon, + incrF=incrF, + decrF=decrF, + safetyF=safetyF, + adapt_cutoff=adapt_cutoff, + minh=minh, + ) if len(linop.shape) > 1: - raise Exception('IF45DP only handles 1D linear operators (diagonal systems): try IF34,ETD34, or ETD35') - self._EL15, self._EL310, self._EL45, self._EL89,self._EL = \ - [np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(5)] - self._NL1, self._NL2, self._NL3, self._NL4,self._NL5,self._NL6,\ - self._NL7 = [np.zeros(self.linop.shape[0],dtype=np.complex128) for _ in range(7)] - self._a21, self._a31, self._a32, self._a41, self._a42, self._a43,\ - self._a51, self._a52, self._a53, self._a54 = [np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(10)] - self._a61, self._a62, self._a63, self._a64, self._a65 = [np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(5)] - self._a71, self._a73, self._a74, self._a75 = [np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(4)] + raise Exception( + "IF45DP only handles 1D linear operators (diagonal systems): try IF34,ETD34, or ETD35" + ) + self._EL15, self._EL310, self._EL45, self._EL89, self._EL = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(5) + ] + self._NL1, self._NL2, self._NL3, self._NL4, self._NL5, self._NL6, self._NL7 = [ + np.zeros(self.linop.shape[0], dtype=np.complex128) for _ in range(7) + ] + ( + self._a21, + self._a31, + self._a32, + self._a41, + self._a42, + self._a43, + self._a51, + self._a52, + self._a53, + self._a54, + ) = [np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(10)] + self._a61, self._a62, self._a63, self._a64, self._a65 = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(5) + ] + self._a71, self._a73, self._a74, self._a75 = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(4) + ] self._a76 = 0.0 - self._r1, self._r3, self._r5, self._r5 = [np.zeros(shape=self.linop.shape,dtype=np.complex128) for _ in range(4)] + self._r1, self._r3, self._r5, self._r5 = [ + np.zeros(shape=self.linop.shape, dtype=np.complex128) for _ in range(4) + ] self._r6, self._r7 = 0.0, 0.0 - self._k = np.zeros(self.linop.shape[0],dtype=np.complex128) - self._err = np.zeros(self.linop.shape[0],dtype=np.complex128) + self._k = np.zeros(self.linop.shape[0], dtype=np.complex128) + self._err = np.zeros(self.linop.shape[0], dtype=np.complex128) self._h_coeff = None self.__N1_init = False - + def _reset(self): self._h_coeff = None self.__N1_init = False - - def _updateStages(self,u,h): + + def _updateStages(self, u, h): self._updateCoeffs(h) # First is same as last principle if not self.__N1_init: @@ -80,71 +116,96 @@ def _updateStages(self,u,h): self.__N1_init = True elif self._accept: self._NL1 = self._NL7.copy() - - self._k = self._EL15*u + self._a21*self._NL1 + + self._k = self._EL15 * u + self._a21 * self._NL1 self._NL2 = self.NLfunc(self._k) - self._k = self._EL310*u + self._a31*self._NL1 + self._a32*self._NL2 + self._k = self._EL310 * u + self._a31 * self._NL1 + self._a32 * self._NL2 self._NL3 = self.NLfunc(self._k) - self._k = self._EL45*u + self._a41*self._NL1 + self._a42*self._NL2 + self._a43*self._NL3 + self._k = ( + self._EL45 * u + + self._a41 * self._NL1 + + self._a42 * self._NL2 + + self._a43 * self._NL3 + ) self._NL4 = self.NLfunc(self._k) - self._k = self._EL89*u + self._a51*self._NL1 + self._a52*self._NL2 + self._a53*self._NL3 \ - + self._a54*self._NL4 + self._k = ( + self._EL89 * u + + self._a51 * self._NL1 + + self._a52 * self._NL2 + + self._a53 * self._NL3 + + self._a54 * self._NL4 + ) self._NL5 = self.NLfunc(self._k) - self._k = self._EL*u + self._a61*self._NL1 + self._a62*self._NL2 + self._a63*self._NL3 \ - + self._a64*self._NL4 + self._a65*self._NL5 - self._NL6 = self.NLfunc(self._k) - self._k = self._EL*u + self._a71*self._NL1 + self._a73*self._NL3 + self._a74*self._NL4 \ - + self._a75*self._NL5 + self._a76*self._NL6 + self._k = ( + self._EL * u + + self._a61 * self._NL1 + + self._a62 * self._NL2 + + self._a63 * self._NL3 + + self._a64 * self._NL4 + + self._a65 * self._NL5 + ) + self._NL6 = self.NLfunc(self._k) + self._k = ( + self._EL * u + + self._a71 * self._NL1 + + self._a73 * self._NL3 + + self._a74 * self._NL4 + + self._a75 * self._NL5 + + self._a76 * self._NL6 + ) self._NL7 = self.NLfunc(self._k) - self._err = self._r1*self._NL1 + self._r3*self._NL3 + self._r4*self._NL4 \ - + self._r5*self._NL5 + self._r6*self._NL6 + self._r7*self._NL7 + self._err = ( + self._r1 * self._NL1 + + self._r3 * self._NL3 + + self._r4 * self._NL4 + + self._r5 * self._NL5 + + self._r6 * self._NL6 + + self._r7 * self._NL7 + ) return self._k, self._err - - def _updateCoeffs(self,h): + + def _updateCoeffs(self, h): # Update coefficients if step size h changed if h == self._h_coeff: return self._h_coeff = h - z = h*self.linop - self._EL15 = np.exp(z/5) - self._EL310 = np.exp(3*z/10) - self._EL45 = np.exp(4*z/5) - self._EL89 = np.exp(8*z/9) + z = h * self.linop + self._EL15 = np.exp(z / 5) + self._EL310 = np.exp(3 * z / 10) + self._EL45 = np.exp(4 * z / 5) + self._EL89 = np.exp(8 * z / 9) self._EL = np.exp(z) - EL710 = np.exp(7*z/10) - EL19 = np.exp(z/9) - self._a21 = h*self._EL15/5.0 - self._a31 = 3*h*self._EL310/40.0 - self._a32 = 9*h*np.exp(z/10)/40.0 - self._a41 = 44*h*self._EL45/45.0 - self._a42 = -56*h*np.exp(3*z/5)/15.0 - self._a43 = 32*h*np.exp(z/2)/9.0 - self._a51 = 19372*h*self._EL89/6561.0 - self._a52 = -25360*h*np.exp(31*z/45)/2187.0 - self._a53 = 64448.0*h*np.exp(53*z/90)/6561.0 - self._a54 = -212*h*np.exp(4*z/45)/729.0 - self._a61 = 9017*h*self._EL/3168.0 - self._a62 = -355*h*self._EL45/33.0 - self._a63 = 46732*h*EL710/5247.0 - self._a64 = 49*h*self._EL15/176.0 - self._a65 = -5103*h*EL19/18656.0 - self._a71 = 35*h*self._EL/384.0 - self._a73 = 500*h*EL710/1113.0 - self._a75 = -2187*h*EL19/6784.0 - self._a74 = 125*h*self._EL15/192.0 - self._a76 = 11*h/84.0 - self._r1 = h*71*self._EL/57600.0 - self._r3 = -71*h*EL710/16695.0 - self._r4 = 17*h*self._EL15/1920.0 - self._r5 = -17253*h*EL19/339200.0 - self._r6 = 22*h/525.0 - self._r7 = -h/40.0 - + EL710 = np.exp(7 * z / 10) + EL19 = np.exp(z / 9) + self._a21 = h * self._EL15 / 5.0 + self._a31 = 3 * h * self._EL310 / 40.0 + self._a32 = 9 * h * np.exp(z / 10) / 40.0 + self._a41 = 44 * h * self._EL45 / 45.0 + self._a42 = -56 * h * np.exp(3 * z / 5) / 15.0 + self._a43 = 32 * h * np.exp(z / 2) / 9.0 + self._a51 = 19372 * h * self._EL89 / 6561.0 + self._a52 = -25360 * h * np.exp(31 * z / 45) / 2187.0 + self._a53 = 64448.0 * h * np.exp(53 * z / 90) / 6561.0 + self._a54 = -212 * h * np.exp(4 * z / 45) / 729.0 + self._a61 = 9017 * h * self._EL / 3168.0 + self._a62 = -355 * h * self._EL45 / 33.0 + self._a63 = 46732 * h * EL710 / 5247.0 + self._a64 = 49 * h * self._EL15 / 176.0 + self._a65 = -5103 * h * EL19 / 18656.0 + self._a71 = 35 * h * self._EL / 384.0 + self._a73 = 500 * h * EL710 / 1113.0 + self._a75 = -2187 * h * EL19 / 6784.0 + self._a74 = 125 * h * self._EL15 / 192.0 + self._a76 = 11 * h / 84.0 + self._r1 = h * 71 * self._EL / 57600.0 + self._r3 = -71 * h * EL710 / 16695.0 + self._r4 = 17 * h * self._EL15 / 1920.0 + self._r5 = -17253 * h * EL19 / 339200.0 + self._r6 = 22 * h / 525.0 + self._r7 = -h / 40.0 + self.logs.append("IF45DP coefficients updated") - + def _q(self): return 5 - - - diff --git a/rkstiff/models.py b/rkstiff/models.py index 3d5583d..f8fa21e 100644 --- a/rkstiff/models.py +++ b/rkstiff/models.py @@ -1,48 +1,51 @@ import numpy as np -def kdvSoliton(x,A=0.5,x0=0,t=0): - u0 = 0.5*A**2/(np.cosh(A*(x-x0-A**2*t)/2)**2) + +def kdvSoliton(x, A=0.5, x0=0, t=0): + u0 = 0.5 * A**2 / (np.cosh(A * (x - x0 - A**2 * t) / 2) ** 2) return u0 -def kdvMultiSoliton(x,A,x0,t=0): + +def kdvMultiSoliton(x, A, x0, t=0): assert len(x0) == len(A) M = len(A) N = len(x) - A = np.array(A).reshape(1,M) - x0 = np.array(x0).reshape(1,M) - u0 = 0.5*A**2/(np.cosh(A*(x.reshape(N,1)-x0-A**2*t)/2)**2) - u0 = np.sum(u0,axis=1) + A = np.array(A).reshape(1, M) + x0 = np.array(x0).reshape(1, M) + u0 = 0.5 * A**2 / (np.cosh(A * (x.reshape(N, 1) - x0 - A**2 * t) / 2) ** 2) + u0 = np.sum(u0, axis=1) return u0 -def kdvOps(kx): - L = 1j*kx**3 - def NL(uf): - u = np.fft.irfft(uf) - ux = np.fft.irfft(1j*kx*uf) - return -6*np.fft.rfft(u*ux) - return L,NL +def kdvOps(kx): + L = 1j * kx**3 -def burgersOps(kx,mu): - L = -mu*kx**2 def NL(uf): u = np.fft.irfft(uf) - ux = np.fft.irfft(1j*kx*uf) - return -np.fft.rfft(u*ux) - return L,NL + ux = np.fft.irfft(1j * kx * uf) + return -6 * np.fft.rfft(u * ux) -def allenCahnOps(x,Dx,epsilon): - epsilon = 0.01 - D2 = Dx.dot(Dx) - L = epsilon*D2 + np.eye(*D2.shape) - L = L[1:-1,1:-1] # remove boundary points - def NL(u): - return x[1:-1] - np.power(u+x[1:-1],3) - return L,NL + return L, NL +def burgersOps(kx, mu): + L = -mu * kx**2 + def NL(uf): + u = np.fft.irfft(uf) + ux = np.fft.irfft(1j * kx * uf) + return -np.fft.rfft(u * ux) + return L, NL +def allenCahnOps(x, Dx, epsilon): + epsilon = 0.01 + D2 = Dx.dot(Dx) + L = epsilon * D2 + np.eye(*D2.shape) + L = L[1:-1, 1:-1] # remove boundary points + + def NL(u): + return x[1:-1] - np.power(u + x[1:-1], 3) + return L, NL diff --git a/rkstiff/solver.py b/rkstiff/solver.py index d9f225c..ef8b976 100644 --- a/rkstiff/solver.py +++ b/rkstiff/solver.py @@ -1,7 +1,6 @@ - -from abc import ABC,abstractmethod +from abc import ABC, abstractmethod import numpy as np -from typing import Tuple,Optional +from typing import Tuple, Optional class StiffSolverAS(ABC): @@ -16,7 +15,7 @@ class StiffSolverAS(ABC): Linear operator (L) in the system dtU = LU + NL(U). Can be either a 2D numpy array (matrix) or a 1D array (diagonal system). L can be either real-valued or complex-valued. - NLfunc : function + NLfunc : function Nonlinear function (NL(U)) in the system dtU = LU + NL(U). Can be a complex or real-valued function. u : list @@ -33,14 +32,14 @@ class StiffSolverAS(ABC): a 'optimal' step size of h_opt is computed and compared vs the current step size h_current. If h_opt > incrF*h_current, then the solver will suggest h_opt as the next step size, otherwise it will continue using the current step size. This parameter is used such that very small changes to the step size are avoided and hence potentially expensive evaluations - of the RK coefficients are avoided. + of the RK coefficients are avoided. decrF : float, < 1.0 Decrement factor for decreasing the step size utilized in propagating a system. After each step in the propagation a 'optimal' step size of h_opt is computed and compared vs the current step size h_current. If h_opt > decrF*h_current and h_opt < h_current, then the solver will suggest h_new = decrF*h_current as the next step size. This parameter is used such that very small changes to the step size are avoided and hence potentially expensive evaluations - of the RK coefficients are avoided. + of the RK coefficients are avoided. epsilon : float Relative error tolerance for system. Solver will suggest step sizes in an attempt to keep the two-norm relative error of the system @@ -56,7 +55,7 @@ class StiffSolverAS(ABC): adapt_cutoff : float < 1 Limits values used in the computation of the suggested step size to those with |u| > adapt_cutoff*max(|u|). To include all - values in the calculation of step size: set this to a very small number. + values in the calculation of step size: set this to a very small number. minh : float Minimum step size that can be taken by the solver before throwing an exception @@ -66,29 +65,39 @@ class StiffSolverAS(ABC): _______ step(u,h_suggest): - Propagates a given array of u one step using an RK method for stiff PDEs. + Propagates a given array of u one step using an RK method for stiff PDEs. evolve(u,t0,tf,h_init=None,store_data=True,store_freq=1) Propagates an initial value (array) of u given at time t0 - until a final time tf is reached using a RK method for stiff PDEs + until a final time tf is reached using a RK method for stiff PDEs reset() Resets solver including erasing stored variables such as self.t, self.u, and self.logs """ + MAX_LOOPS = 50 - MAX_S = 4 # maximum step increase factor - MIN_S = 0.25 # minimum step decrease factor - - def __init__(self,linop,NLfunc,epsilon = 1e-4,incrF = 1.25, decrF = 0.85,\ - safetyF = 0.8, adapt_cutoff = 0.01, minh = 1e-16): + MAX_S = 4 # maximum step increase factor + MIN_S = 0.25 # minimum step decrease factor + + def __init__( + self, + linop, + NLfunc, + epsilon=1e-4, + incrF=1.25, + decrF=0.85, + safetyF=0.8, + adapt_cutoff=0.01, + minh=1e-16, + ): """ linop : np.array Linear operator (L) in the system dtU = LU + NL(U). Can be either a 2D numpy array (matrix) or a 1D array (diagonal system). L can be either real-valued or complex-valued. - NLfunc : function + NLfunc : function Nonlinear function (NL(U)) in the system dtU = LU + NL(U). Can be a complex or real-valued function. epsilon : float @@ -102,14 +111,14 @@ def __init__(self,linop,NLfunc,epsilon = 1e-4,incrF = 1.25, decrF = 0.85,\ a 'optimal' step size of h_opt is computed and compared vs the current step size h_current. If h_opt > incrF*h_current, then the solver will suggest h_opt as the next step size, otherwise it will continue using the current step size. This parameter is used such that very small changes to the step size are avoided and hence potentially expensive evaluations - of the RK coefficients are avoided. + of the RK coefficients are avoided. decrF : float, < 1.0 Decrement factor for decreasing the step size utilized in propagating a system. After each step in the propagation a 'optimal' step size of h_opt is computed and compared vs the current step size h_current. If h_opt > decrF*h_current and h_opt < h_current, then the solver will suggest h_new = decrF*h_current as the next step size. This parameter is used such that very small changes to the step size are avoided and hence potentially expensive evaluations - of the RK coefficients are avoided. + of the RK coefficients are avoided. safetyF : float Safety factor for adaptive stepping. The 'optimal' step size computed using the embedded RK methods is multiplied by @@ -119,7 +128,7 @@ def __init__(self,linop,NLfunc,epsilon = 1e-4,incrF = 1.25, decrF = 0.85,\ adapt_cutoff : float < 1 Limits values used in the computation of the suggested step size to those with |u| > adapt_cutoff*max(|u|). To include all - values in the calculation of step size: set this to a very small number. + values in the calculation of step size: set this to a very small number. minh : float Minimum step size that can be taken by the solver before throwing an exception @@ -133,70 +142,72 @@ def __init__(self,linop,NLfunc,epsilon = 1e-4,incrF = 1.25, decrF = 0.85,\ dims = linop.shape self._diag = True if len(dims) > 2: - raise ValueError('linop must be a 1D or 2D array') + raise ValueError("linop must be a 1D or 2D array") elif len(dims) == 2: - if (dims[0] != dims[1]): - raise ValueError('linop must be a square matrix') + if dims[0] != dims[1]: + raise ValueError("linop must be a square matrix") self._diag = False self.epsilon = epsilon if self.epsilon <= 0: - raise ValueError('epsilon must be positive but is {}'.format(self.epsilon)) + raise ValueError("epsilon must be positive but is {}".format(self.epsilon)) self.incrF = incrF if self.incrF <= 1.0: - raise ValueError('incrF must be > 1.0 but is {}'.format(self.incrF)) - + raise ValueError("incrF must be > 1.0 but is {}".format(self.incrF)) + self.decrF = decrF if self.decrF >= 1.0: - raise ValueError('decrF must be < 1.0 but is {}'.format(self.decrF)) - + raise ValueError("decrF must be < 1.0 but is {}".format(self.decrF)) + self.safetyF = safetyF if self.safetyF > 1.0: - raise ValueError('safetyF must be <= 1.0 but is {}'.format(self.safetyF)) + raise ValueError("safetyF must be <= 1.0 but is {}".format(self.safetyF)) self.adapt_cutoff = adapt_cutoff if self.adapt_cutoff >= 1.0: - raise ValueError('adapt_cutoff must be < 1.0 but is {}'.format(self.adapt_cutoff)) + raise ValueError( + "adapt_cutoff must be < 1.0 but is {}".format(self.adapt_cutoff) + ) self.minh = minh if self.minh <= 0: - raise ValueError('minh must be positive but is {}'.format(self.minh)) + raise ValueError("minh must be positive but is {}".format(self.minh)) - self.__t0, self.__tf, self.__tc = 0,0,0 + self.__t0, self.__tf, self.__tc = 0, 0, 0 self.__h_prev = 0.0 - self._accept = False + self._accept = False self.t, self.u, self.logs = [], [], [] def reset(self): - """ Resets solver to its initial state (such that its ready to call the functions evolve - or step on a new input). Erases stored variables such as self.t, self.u, and self.logs. + """Resets solver to its initial state (such that its ready to call the functions evolve + or step on a new input). Erases stored variables such as self.t, self.u, and self.logs. """ self.t, self.u, self.logs = [], [], [] - self.__t0, self.__tf, self.__tc = 0,0,0 + self.__t0, self.__tf, self.__tc = 0, 0, 0 self.__h_prev = 0.0 - self._accept = False + self._accept = False self._reset() # reset solver dependent variables of the subclass @abstractmethod - def _reset(): + def _reset(self): pass - + # update RK stages, returns u_{n+1} given u_{n}, to be overwritten by subclass @abstractmethod - def _updateStages(self,u,h): + def _updateStages(self, u, h): pass - + # q value in computation of suggested step size, to be overwritten by subclass @abstractmethod def _q(self): pass - - def step(self,u : np.ndarray,h_suggest : float) -> Tuple[np.ndarray,float,float]: - """ - Propagates a given array of u one step using an RK method for stiff PDEs. - + + def step(self, u: np.ndarray, h_suggest: float) -> Tuple[np.ndarray, float, float]: + """ + Propagates a given array of u one step using an RK method for stiff PDEs. + INPUTS: u : np.array input to propagate @@ -212,22 +223,22 @@ def step(self,u : np.ndarray,h_suggest : float) -> Tuple[np.ndarray,float,float] h_suggest : float step size the algorithm suggests you take for the next step """ - + h = h_suggest assert h >= 0.0 numloops = 0 - while(True): - unew,err = self._updateStages(u,h) - self.__h_prev = h + while True: + unew, err = self._updateStages(u, h) + self.__h_prev = h # Compute step size change factor s - s = self._computeS(unew,err) + s = self._computeS(unew, err) # If s is less than 1, inf, or nan, reject step and reduce step size - if np.isinf(s) or np.isnan(s) or s < 1.0: - h = self._rejectStepSize(s,h) + if np.isinf(s) or np.isnan(s) or s < 1.0: + h = self._rejectStepSize(s, h) # If s is bigger than 1 accept h and the step else: - h_suggest = self._acceptStepSize(s,h) - return unew, h, h_suggest + h_suggest = self._acceptStepSize(s, h) + return unew, h, h_suggest numloops += 1 if numloops > self.MAX_LOOPS: failure_str = """Solver failed: adaptive step made too many attempts to find a step @@ -235,105 +246,114 @@ def step(self,u : np.ndarray,h_suggest : float) -> Tuple[np.ndarray,float,float] self.logs.append(failure_str) raise Exception(failure_str) if h < self.minh: - failure_str = """Solver failed: adaptive step reached minimum step size """ + failure_str = ( + """Solver failed: adaptive step reached minimum step size """ + ) self.logs.append(failure_str) raise Exception(failure_str) - + raise Exception("Propagation Failed") - + # helper function for computing coefficient s used in generating suggested step size - def _computeS(self,u,err): + def _computeS(self, u, err): # Use adapt_cutoff to ignore small modes/values in the computation of the step size magu = np.abs(u) - idx = magu/magu.max() > self.adapt_cutoff - tol = self.epsilon*np.linalg.norm(u[idx]) - s = self.safetyF*np.power(tol/np.linalg.norm(err[idx]),1.0/self._q()) + idx = magu / magu.max() > self.adapt_cutoff + tol = self.epsilon * np.linalg.norm(u[idx]) + s = self.safetyF * np.power(tol / np.linalg.norm(err[idx]), 1.0 / self._q()) return s - + # helper function computing suggested step size if step is rejected - def _rejectStepSize(self,s,h): + def _rejectStepSize(self, s, h): self._accept = False # Check that s is a number if np.isinf(s) or np.isnan(s): - self.logs.append('inf or nan number encountered: reducing step size to {}!'.format(h)) - return self.MIN_S*h - - s = np.max([s,self.MIN_S]) # dont let s be too small - s = np.min([s,self.decrF]) # dont let s be too close to 1 - self.logs.append('step rejected with s = {:.2f}'.format(s)) - hnew = s*h - self.logs.append('reducing step size to {}'.format(hnew)) + self.logs.append( + "inf or nan number encountered: reducing step size to {}!".format(h) + ) + return self.MIN_S * h + + s = np.max([s, self.MIN_S]) # dont let s be too small + s = np.min([s, self.decrF]) # dont let s be too close to 1 + self.logs.append("step rejected with s = {:.2f}".format(s)) + hnew = s * h + self.logs.append("reducing step size to {}".format(hnew)) return hnew - + # helper function computing suggested step size if step is accepted - def _acceptStepSize(self,s,h): + def _acceptStepSize(self, s, h): self._accept = True - s = np.min([s,self.MAX_S]) # dont let s be too big - self.logs.append('step accepted with s = {:.2f}'.format(s)) + s = np.min([s, self.MAX_S]) # dont let s be too big + self.logs.append("step accepted with s = {:.2f}".format(s)) # if s much larger than 1, increase the step size if s > self.incrF: - h_suggest = s*h - self.logs.append('increasing step size to {}'.format(h_suggest)) + h_suggest = s * h + self.logs.append("increasing step size to {}".format(h_suggest)) return h_suggest return h - - def evolve(self,u : np.ndarray,t0 : float,tf : float,\ - h_init : Optional[float]=None,\ - store_data : bool=True, store_freq : int=1) -> np.ndarray: - """ + def evolve( + self, + u: np.ndarray, + t0: float, + tf: float, + h_init: Optional[float] = None, + store_data: bool = True, + store_freq: int = 1, + ) -> np.ndarray: + """ This function propagates an initial value (array) of u given at time t0 - until a final time tf is reached using a RK method for stiff PDEs - + until a final time tf is reached using a RK method for stiff PDEs + INPUTS: u : np.array - initial value input + initial value input t0 : float initial time at which u is evaluated tf : float end time at which propagation stops h_init : float, optional starting step size for RK method (default of (tf-t0)/100 if not specified) - store_data : bool, optional + store_data : bool, optional value that determines whether to keep track of the propagation array u at each step of the RK method. Values stored in self.u and self.t - store_freq : int, optional + store_freq : int, optional store propagation data in self.t and self.u after every [store_freq] step is taken OUTPUTS: - u : np.array - final value of the input propagated from t0 to tf + u : np.array + final value of the input propagated from t0 to tf """ - + self.reset() self.__t0, self.__tf, self.__tc = t0, tf, t0 - + if store_data: self.t.append(t0) self.u.append(u) - + # Set initial step size if none given if h_init is None: - h_init = (self.__tf-self.__t0)/100. + h_init = (self.__tf - self.__t0) / 100.0 h = h_init - + # Make sure step size isn't larger than entire propagation time - if self.__tc+h > self.__tf: + if self.__tc + h > self.__tf: h = self.__tf - self.__tc - + step_count = 0 while self.__tc < self.__tf: - u,h,h_suggest = self.step(u,h) + u, h, h_suggest = self.step(u, h) self.__tc += h step_count += 1 - if self.__tc+h_suggest > self.__tf: + if self.__tc + h_suggest > self.__tf: h = self.__tf - self.__tc else: h = h_suggest - + if store_data and (step_count % store_freq == 0): self.t.append(self.__tc) self.u.append(u) - + return u @@ -349,7 +369,7 @@ class StiffSolverCS(ABC): Linear operator (L) in the system dtU = LU + NL(U). Can be either a 2D numpy array (matrix) or a 1D array (diagonal system). L can be either real-valued or complex-valued. - NLfunc : function + NLfunc : function Nonlinear function (NL(U)) in the system dtU = LU + NL(U). Can be a complex or real-valued function. u : list @@ -362,25 +382,25 @@ class StiffSolverCS(ABC): _______ step(u,h): - Propagates a given array of u, taking a step of size h, using an RK method for stiff PDEs. + Propagates a given array of u, taking a step of size h, using an RK method for stiff PDEs. evolve(u,t0,tf,h,store_data=True,store_freq=1) Propagates an initial value (array) of u given at time t0 - until a final time tf is reached using a RK method for stiff PDEs. + until a final time tf is reached using a RK method for stiff PDEs. Solver will always take a step of size h reset() Resets solver including erasing stored variables such as self.t, self.u, and self.logs """ - - def __init__(self,linop,NLfunc): + + def __init__(self, linop, NLfunc): """ linop : np.array Linear operator (L) in the system dtU = LU + NL(U). Can be either a 2D numpy array (matrix) or a 1D array (diagonal system). L can be either real-valued or complex-valued. - NLfunc : function + NLfunc : function Nonlinear function (NL(U)) in the system dtU = LU + NL(U). Can be a complex or real-valued function. """ @@ -392,37 +412,37 @@ def __init__(self,linop,NLfunc): dims = linop.shape self._diag = True if len(dims) > 2: - raise ValueError('linop must be a 1D or 2D array') + raise ValueError("linop must be a 1D or 2D array") elif len(dims) == 2: - if (dims[0] != dims[1]): - raise ValueError('linop must be a square matrix') + if dims[0] != dims[1]: + raise ValueError("linop must be a square matrix") self._diag = False - self.__t0, self.__tf, self.__tc = 0,0,0 + self.__t0, self.__tf, self.__tc = 0, 0, 0 self.t, self.u, self.logs = [], [], [] def reset(self): - """ Resets solver to its initial state (such that its ready to call the functions evolve - or step on a new input). Erases stored variables such as self.t, self.u, and self.logs. + """Resets solver to its initial state (such that its ready to call the functions evolve + or step on a new input). Erases stored variables such as self.t, self.u, and self.logs. """ self.t, self.u, self.logs = [], [], [] - self.__t0, self.__tf, self.__tc = 0,0,0 + self.__t0, self.__tf, self.__tc = 0, 0, 0 self._reset() # reset solver dependent variables of the subclass @abstractmethod - def _reset(): + def _reset(self): pass - + # update RK stages, returns u_{n+1} given u_{n}, to be overwritten by subclass @abstractmethod - def _updateStages(self,u,h): + def _updateStages(self, u, h): pass - - def step(self,u : np.ndarray, h : float) -> np.ndarray: - """ - Propagates a given array of u one step using an RK method for stiff PDEs. - + + def step(self, u: np.ndarray, h: float) -> np.ndarray: + """ + Propagates a given array of u one step using an RK method for stiff PDEs. + INPUTS: u : np.array input to propagate @@ -433,20 +453,27 @@ def step(self,u : np.ndarray, h : float) -> np.ndarray: unew : np.array result of input u propagated one step in the RK algorithm """ - + assert h >= 0.0 - unew = self._updateStages(u,h) + unew = self._updateStages(u, h) return unew - - def evolve(self,u : np.ndarray,t0 : float,tf : float, h : float,\ - store_data : bool=True, store_freq : int=1) -> np.ndarray: - """ + + def evolve( + self, + u: np.ndarray, + t0: float, + tf: float, + h: float, + store_data: bool = True, + store_freq: int = 1, + ) -> np.ndarray: + """ This function propagates an initial value (array) of u given at time t0 - until a final time tf is reached using a RK method for stiff PDEs - + until a final time tf is reached using a RK method for stiff PDEs + INPUTS: u : np.array - initial value input + initial value input t0 : float initial time at which u is evaluated tf : float @@ -454,37 +481,38 @@ def evolve(self,u : np.ndarray,t0 : float,tf : float, h : float,\ greater than or equal to this value h : float step size for RK method - store_data : bool, optional + store_data : bool, optional value that determines whether to keep track of the propagation array u at each step of the RK method. Values stored in self.u and self.t - store_freq : int, optional + store_freq : int, optional store propagation data in self.t and self.u after every [store_freq] step is taken OUTPUTS: - u : np.array - final value of the input propagated from t0 to tf + u : np.array + final value of the input propagated from t0 to tf tf : float actual final time propagated (may be greater than specified tf) """ - + self.reset() self.__t0, self.__tf, self.__tc = t0, tf, t0 - + if store_data: self.t.append(t0) self.u.append(u) - + # Make sure step size isn't larger than entire propagation time - if self.__tc+h > self.__tf: - raise ValueError('Reduce step size h, it needs to be less than or equal to tf-t0') - + if self.__tc + h > self.__tf: + raise ValueError( + "Reduce step size h, it needs to be less than or equal to tf-t0" + ) + step_count = 0 while self.__tc < self.__tf: - u = self.step(u,h) + u = self.step(u, h) self.__tc += h step_count += 1 if store_data and (step_count % store_freq == 0): self.t.append(self.__tc) self.u.append(u) - - return u + return u diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py index a8347ba..aad0e01 100644 --- a/tests/test_derivatives.py +++ b/tests/test_derivatives.py @@ -1,121 +1,126 @@ - from rkstiff.grids import construct_x_kx_rfft, construct_x_kx_fft from rkstiff.grids import construct_x_Dx_cheb from rkstiff.derivatives import dx_rfft, dx_fft import numpy as np + def test_periodic_dx_rfft(): N = 100 - a, b = 0, 2*np.pi - x,kx = construct_x_kx_rfft(N,a,b) + a, b = 0, 2 * np.pi + x, kx = construct_x_kx_rfft(N, a, b) u = np.sin(x) ux_exact = np.cos(x) - ux_approx = dx_rfft(kx,u) - assert np.allclose(ux_exact,ux_approx) + ux_approx = dx_rfft(kx, u) + assert np.allclose(ux_exact, ux_approx) + def test_zeroboundaries_dx_rfft(): N = 400 - a, b = -30., 30. - x,kx = construct_x_kx_rfft(N,a,b) - u = 1./np.cosh(x) - ux_exact = -np.tanh(x)/np.cosh(x) - ux_approx = dx_rfft(kx,u) - assert np.allclose(ux_exact,ux_approx) + a, b = -30.0, 30.0 + x, kx = construct_x_kx_rfft(N, a, b) + u = 1.0 / np.cosh(x) + ux_exact = -np.tanh(x) / np.cosh(x) + ux_approx = dx_rfft(kx, u) + assert np.allclose(ux_exact, ux_approx) + def test_gauss_dx_rfft(): N = 128 - a,b = -10,10 - x,kx = construct_x_kx_rfft(N,a,b) - u = np.exp(-x**2) - ux_exact = -2*x*np.exp(-x**2) - ux_approx = dx_rfft(kx,u) - assert np.allclose(ux_exact,ux_approx) + a, b = -10, 10 + x, kx = construct_x_kx_rfft(N, a, b) + u = np.exp(-(x**2)) + ux_exact = -2 * x * np.exp(-(x**2)) + ux_approx = dx_rfft(kx, u) + assert np.allclose(ux_exact, ux_approx) def test_manydx_rfft(): - N = 128 - a, b = 0, 2*np.pi - x,kx = construct_x_kx_rfft(N,a,b) + N = 128 + a, b = 0, 2 * np.pi + x, kx = construct_x_kx_rfft(N, a, b) u = np.sin(x) ux_exact = np.sin(x) ux_approx = u.copy() for _ in range(4): - ux_approx = dx_rfft(kx,ux_approx) - rel_err = np.linalg.norm(ux_exact-ux_approx)/np.linalg.norm(ux_exact) + ux_approx = dx_rfft(kx, ux_approx) + rel_err = np.linalg.norm(ux_exact - ux_approx) / np.linalg.norm(ux_exact) assert rel_err < 1e-6 ux_approx = u.copy() - ux_approx = dx_rfft(kx,ux_approx,8) - rel_err = np.linalg.norm(ux_exact-ux_approx)/np.linalg.norm(ux_exact) + ux_approx = dx_rfft(kx, ux_approx, 8) + rel_err = np.linalg.norm(ux_exact - ux_approx) / np.linalg.norm(ux_exact) assert rel_err < 0.1 - - def test_manydx_fft(): - N = 128 - a, b = 0, 2*np.pi - x,kx = construct_x_kx_fft(N,a,b) + N = 128 + a, b = 0, 2 * np.pi + x, kx = construct_x_kx_fft(N, a, b) u = np.sin(x) ux_exact = np.sin(x) ux_approx = u.copy() for _ in range(4): - ux_approx = dx_fft(kx,ux_approx) - rel_err = np.linalg.norm(ux_exact-ux_approx)/np.linalg.norm(ux_exact) + ux_approx = dx_fft(kx, ux_approx) + rel_err = np.linalg.norm(ux_exact - ux_approx) / np.linalg.norm(ux_exact) assert rel_err < 1e-6 ux_approx = u.copy() - ux_approx = dx_fft(kx,ux_approx,8) - rel_err = np.linalg.norm(ux_exact-ux_approx)/np.linalg.norm(ux_exact) - assert rel_err < 0.1 + ux_approx = dx_fft(kx, ux_approx, 8) + rel_err = np.linalg.norm(ux_exact - ux_approx) / np.linalg.norm(ux_exact) + assert rel_err < 0.1 + def test_periodic_dx_fft(): N = 100 - a, b = 0, 2*np.pi - x,kx = construct_x_kx_fft(N,a,b) + a, b = 0, 2 * np.pi + x, kx = construct_x_kx_fft(N, a, b) u = np.sin(x) ux_exact = np.cos(x) - ux_approx = dx_fft(kx,u) - assert np.allclose(ux_exact,ux_approx) + ux_approx = dx_fft(kx, u) + assert np.allclose(ux_exact, ux_approx) + def test_zeroboundaries_dx_fft(): N = 400 - a, b = -30., 30. - x,kx = construct_x_kx_fft(N,a,b) - u = 1./np.cosh(x) - ux_exact = -np.tanh(x)/np.cosh(x) - ux_approx = dx_fft(kx,u) - assert np.allclose(ux_exact,ux_approx) + a, b = -30.0, 30.0 + x, kx = construct_x_kx_fft(N, a, b) + u = 1.0 / np.cosh(x) + ux_exact = -np.tanh(x) / np.cosh(x) + ux_approx = dx_fft(kx, u) + assert np.allclose(ux_exact, ux_approx) + def test_gauss_dx_fft(): N = 128 - a,b = -10,10 - x,kx = construct_x_kx_fft(N,a,b) - u = np.exp(-x**2) - ux_exact = -2*x*np.exp(-x**2) - ux_approx = dx_fft(kx,u) - assert np.allclose(ux_exact,ux_approx) + a, b = -10, 10 + x, kx = construct_x_kx_fft(N, a, b) + u = np.exp(-(x**2)) + ux_exact = -2 * x * np.exp(-(x**2)) + ux_approx = dx_fft(kx, u) + assert np.allclose(ux_exact, ux_approx) def test_exp_trig_x_Dx_cheb(): # standard interval [-1,1] - N = 20; a = -1; b = 1 - x,Dx = construct_x_Dx_cheb(N,-1,1) - u = np.exp(x)*np.sin(5*x) - Du_exact = np.exp(x)*(np.sin(5*x)+5*np.cos(5*x)) - Du_approx = Dx.dot(u) + N = 20 + a = -1 + b = 1 + x, Dx = construct_x_Dx_cheb(N, -1, 1) + u = np.exp(x) * np.sin(5 * x) + Du_exact = np.exp(x) * (np.sin(5 * x) + 5 * np.cos(5 * x)) + Du_approx = Dx.dot(u) error = Du_exact - Du_approx - assert np.linalg.norm(error)/np.linalg.norm(Du_exact) < 1e-8 + assert np.linalg.norm(error) / np.linalg.norm(Du_exact) < 1e-8 # non-standard interval [-3,3] - N = 30; a = -3; b = 3 - x,Dx = construct_x_Dx_cheb(N,a,b) - u = np.exp(x)*np.sin(5*x) - Du_exact = np.exp(x)*(np.sin(5*x)+5*np.cos(5*x)) - Du_approx = Dx.dot(u) + N = 30 + a = -3 + b = 3 + x, Dx = construct_x_Dx_cheb(N, a, b) + u = np.exp(x) * np.sin(5 * x) + Du_exact = np.exp(x) * (np.sin(5 * x) + 5 * np.cos(5 * x)) + Du_approx = Dx.dot(u) error = Du_exact - Du_approx - assert np.linalg.norm(error)/np.linalg.norm(Du_exact) < 1e-7 - - + assert np.linalg.norm(error) / np.linalg.norm(Du_exact) < 1e-7 diff --git a/tests/test_solvers.py b/tests/test_solvers.py index cb5fbd2..ba85878 100644 --- a/tests/test_solvers.py +++ b/tests/test_solvers.py @@ -1,4 +1,3 @@ - from rkstiff.grids import construct_x_kx_rfft, construct_x_Dx_cheb from rkstiff.etd34 import ETD34 from rkstiff.etd35 import ETD35 @@ -13,201 +12,236 @@ import numpy as np import pytest + # Check that soliton 'shape' is preserved through propagation def NLstub(u): - return + return + def Lstub(N): return np.ones(N) + def test_solver_input(): # Test that invalid linear operator raises ValueError with pytest.raises(ValueError): - L = np.zeros(shape=(4,2)) - solver = ETD34(linop=L,NLfunc=NLstub) + L = np.zeros(shape=(4, 2)) + solver = ETD34(linop=L, NLfunc=NLstub) # Test that invalid linear increment factor raises ValueError with pytest.raises(ValueError): - solver = ETD34(linop = Lstub(10),NLfunc=NLstub,incrF = 0.8) + solver = ETD34(linop=Lstub(10), NLfunc=NLstub, incrF=0.8) # Test that invalid linear decrement factor raises ValueError with pytest.raises(ValueError): - solver = ETD34(linop = Lstub(10),NLfunc=NLstub,decrF = 1.1) + solver = ETD34(linop=Lstub(10), NLfunc=NLstub, decrF=1.1) # Test that negative epsilon raises ValueError with pytest.raises(ValueError): - solver = ETD34(linop = Lstub(10),NLfunc=NLstub,epsilon = -1e-3) - # Test that invalid safetyFactor raises ValueError + solver = ETD34(linop=Lstub(10), NLfunc=NLstub, epsilon=-1e-3) + # Test that invalid safetyFactor raises ValueError with pytest.raises(ValueError): - solver = ETD34(linop = Lstub(10),NLfunc=NLstub,safetyF = 1.2) + solver = ETD34(linop=Lstub(10), NLfunc=NLstub, safetyF=1.2) # Test that invalid adapt_cutoff raises ValueError with pytest.raises(ValueError): - solver = ETD34(linop = Lstub(10),NLfunc=NLstub,adapt_cutoff = 1) + solver = ETD34(linop=Lstub(10), NLfunc=NLstub, adapt_cutoff=1) with pytest.raises(ValueError): - solver = ETD34(linop = Lstub(10),NLfunc=NLstub,minh = 0) + solver = ETD34(linop=Lstub(10), NLfunc=NLstub, minh=0) + def test_etdsolver_input(): # Test that modecutoff is valid with pytest.raises(ValueError): - solver = ETD35(linop = Lstub(10),NLfunc=NLstub,modecutoff=1.2) + solver = ETD35(linop=Lstub(10), NLfunc=NLstub, modecutoff=1.2) with pytest.raises(ValueError): - solver = ETD35(linop = Lstub(10),NLfunc=NLstub,modecutoff=0) + solver = ETD35(linop=Lstub(10), NLfunc=NLstub, modecutoff=0) # Test that contour_points is an integer with pytest.raises(TypeError): - solver = ETD35(linop = Lstub(10),NLfunc=NLstub,contour_points = 12.2) + solver = ETD35(linop=Lstub(10), NLfunc=NLstub, contour_points=12.2) # Test that contour_points is greater than 1 with pytest.raises(ValueError): - solver = ETD35(linop = Lstub(10),NLfunc=NLstub,contour_points = 1) + solver = ETD35(linop=Lstub(10), NLfunc=NLstub, contour_points=1) # Test that contour_radius is positive with pytest.raises(ValueError): - solver = ETD35(linop = Lstub(10),NLfunc=NLstub,contour_radius = 0) + solver = ETD35(linop=Lstub(10), NLfunc=NLstub, contour_radius=0) + def test_abc_error(): # Test that abstract base classes cannot be instantiated with pytest.raises(TypeError): - solver = ETDAS(linop = Lstub(10),NLfunc=NLstub) + solver = ETDAS(linop=Lstub(10), NLfunc=NLstub) with pytest.raises(TypeError): - solver = StiffSolverAS(linop = Lstub(10),NLfunc=NLstub) + solver = StiffSolverAS(linop=Lstub(10), NLfunc=NLstub) def allen_cahn_setup(): - N = 20 - a = -1; b = 1 - x,Dx = construct_x_Dx_cheb(N,a,b) + N = 20 + a = -1 + b = 1 + x, Dx = construct_x_Dx_cheb(N, a, b) epsilon = 0.01 - L,NL = models.allenCahnOps(x,Dx,epsilon) - u0 = 0.53*x + 0.47*np.sin(-1.5*np.pi*x) + L, NL = models.allenCahnOps(x, Dx, epsilon) + u0 = 0.53 * x + 0.47 * np.sin(-1.5 * np.pi * x) w0 = u0 - x - w0int = w0[1:-1]; u0int = u0[1:-1]; xint = x[1:-1] - return xint,u0int,w0int,L,NL + w0int = w0[1:-1] + u0int = u0[1:-1] + xint = x[1:-1] + return xint, u0int, w0int, L, NL + def burgers_setup(): N = 1024 - a,b = -np.pi, np.pi - x,kx = construct_x_kx_rfft(N,a,b) + a, b = -np.pi, np.pi + x, kx = construct_x_kx_rfft(N, a, b) mu = 0.0005 - L,NL = models.burgersOps(kx,mu) + L, NL = models.burgersOps(kx, mu) - u0 = np.exp(-10*np.sin(x/2)**2) + u0 = np.exp(-10 * np.sin(x / 2) ** 2) u0FFT = np.fft.rfft(u0) - return u0FFT,L,NL + return u0FFT, L, NL + def kdv_soliton_setup(): N = 256 - a,b = -30,30 - x,kx = construct_x_kx_rfft(N,a,b) - A, x0, t0 = 1., -5., 0 + a, b = -30, 30 + x, kx = construct_x_kx_rfft(N, a, b) + A, x0, t0 = 1.0, -5.0, 0 h = 0.025 steps = 200 - tf = h*steps + tf = h * steps - u0 = models.kdvSoliton(x,A=A,x0=x0,t=t0) + u0 = models.kdvSoliton(x, A=A, x0=x0, t=t0) u0FFT = np.fft.rfft(u0) - uexactFFT = np.fft.rfft(models.kdvSoliton(x,A=A,x0=x0,t=tf)) - L,NL = models.kdvOps(kx) + uexactFFT = np.fft.rfft(models.kdvSoliton(x, A=A, x0=x0, t=tf)) + L, NL = models.kdvOps(kx) + + return u0FFT, L, NL, uexactFFT, h, steps - return u0FFT,L,NL,uexactFFT,h,steps -def kdv_CS_step_eval(solver,u0FFT,uexactFFT,h,steps,tol): +def kdv_CS_step_eval(solver, u0FFT, uexactFFT, h, steps, tol): for _ in range(steps): - u0FFT = solver.step(u0FFT,h) - rel_err = np.linalg.norm(u0FFT-uexactFFT)/np.linalg.norm(uexactFFT) - assert rel_err < tol + u0FFT = solver.step(u0FFT, h) + rel_err = np.linalg.norm(u0FFT - uexactFFT) / np.linalg.norm(uexactFFT) + assert rel_err < tol + -def kdv_AS_step_eval(solver,u0FFT,uexactFFT,h,steps,tol): +def kdv_AS_step_eval(solver, u0FFT, uexactFFT, h, steps, tol): for _ in range(steps): - u0FFT,h_actual,h_suggest = solver.step(u0FFT,h) + u0FFT, h_actual, h_suggest = solver.step(u0FFT, h) assert (h_actual - h) < 1e-10 - rel_err = np.linalg.norm(u0FFT-uexactFFT)/np.linalg.norm(uexactFFT) - assert rel_err < tol + rel_err = np.linalg.norm(u0FFT - uexactFFT) / np.linalg.norm(uexactFFT) + assert rel_err < tol + + +def kdv_evolve_eval(solver, u0FFT, uexactFFT, h, tf, tol): + uFFT = solver.evolve(u0FFT, 0.0, tf, h, store_data=False) + rel_err = np.linalg.norm(uFFT - uexactFFT) / np.linalg.norm(uexactFFT) + assert rel_err < tol -def kdv_evolve_eval(solver,u0FFT,uexactFFT,h,tf,tol): - uFFT = solver.evolve(u0FFT,0.0,tf,h,store_data=False) - rel_err = np.linalg.norm(uFFT-uexactFFT)/np.linalg.norm(uexactFFT) - assert rel_err < tol def test_etd34_nondiag(): - xint,u0int,w0int,L,NL = allen_cahn_setup() - solver = ETD34(linop=L,NLfunc=NL,epsilon=1e-3,contour_points=64,contour_radius=20) - wfint = solver.evolve(w0int,t0=0,tf=60,store_data=False) + xint, u0int, w0int, L, NL = allen_cahn_setup() + solver = ETD34( + linop=L, NLfunc=NL, epsilon=1e-3, contour_points=64, contour_radius=20 + ) + wfint = solver.evolve(w0int, t0=0, tf=60, store_data=False) ufint = wfint.real + xint - assert np.abs(u0int[0]-ufint[0]) < 0.01 - assert np.abs(u0int[7]-ufint[7]) > 1 + assert np.abs(u0int[0] - ufint[0]) < 0.01 + assert np.abs(u0int[7] - ufint[7]) > 1 + def test_etd34(): - u0FFT,L,NL,uexactFFT,h,steps = kdv_soliton_setup() - solver = ETD34(linop=L,NLfunc=NL,epsilon=1e-1) # small epsilon -> step size isn't auto-reduced - kdv_AS_step_eval(solver,u0FFT,uexactFFT,h,steps,tol=1e-4) + u0FFT, L, NL, uexactFFT, h, steps = kdv_soliton_setup() + solver = ETD34( + linop=L, NLfunc=NL, epsilon=1e-1 + ) # small epsilon -> step size isn't auto-reduced + kdv_AS_step_eval(solver, u0FFT, uexactFFT, h, steps, tol=1e-4) solver.reset() solver.epsilon = 1e-4 - kdv_evolve_eval(solver,u0FFT,uexactFFT,h,tf=h*steps,tol=1e-4) + kdv_evolve_eval(solver, u0FFT, uexactFFT, h, tf=h * steps, tol=1e-4) + def test_etd35(): - u0FFT,L,NL,uexactFFT,h,steps = kdv_soliton_setup() - solver = ETD35(linop=L,NLfunc=NL,epsilon=1e-1) # small epsilon -> step size isn't auto-reduced - kdv_AS_step_eval(solver,u0FFT,uexactFFT,h,steps,tol=1e-4) + u0FFT, L, NL, uexactFFT, h, steps = kdv_soliton_setup() + solver = ETD35( + linop=L, NLfunc=NL, epsilon=1e-1 + ) # small epsilon -> step size isn't auto-reduced + kdv_AS_step_eval(solver, u0FFT, uexactFFT, h, steps, tol=1e-4) solver.reset() solver.epsilon = 1e-4 - kdv_evolve_eval(solver,u0FFT,uexactFFT,h,tf=h*steps,tol=1e-5) + kdv_evolve_eval(solver, u0FFT, uexactFFT, h, tf=h * steps, tol=1e-5) + def test_etd35_nondiag(): - xint,u0int,w0int,L,NL = allen_cahn_setup() - solver = ETD35(linop=L,NLfunc=NL,epsilon=1e-4,contour_points=32,contour_radius=10) - wfint = solver.evolve(w0int,t0=0,tf=60,store_data=False) + xint, u0int, w0int, L, NL = allen_cahn_setup() + solver = ETD35( + linop=L, NLfunc=NL, epsilon=1e-4, contour_points=32, contour_radius=10 + ) + wfint = solver.evolve(w0int, t0=0, tf=60, store_data=False) ufint = wfint.real + xint - assert np.abs(u0int[0]-ufint[0]) < 0.01 - assert np.abs(u0int[7]-ufint[7]) > 1 + assert np.abs(u0int[0] - ufint[0]) < 0.01 + assert np.abs(u0int[7] - ufint[7]) > 1 + def test_if34(): - u0FFT,L,NL = burgers_setup() - solver = IF34(linop=L,NLfunc=NL,epsilon=1e-4) - uFFT = solver.evolve(u0FFT,t0=0,tf=0.85,store_data=False) - rel_err = np.abs(np.linalg.norm(uFFT)-np.linalg.norm(u0FFT))/np.linalg.norm(u0FFT) + u0FFT, L, NL = burgers_setup() + solver = IF34(linop=L, NLfunc=NL, epsilon=1e-4) + uFFT = solver.evolve(u0FFT, t0=0, tf=0.85, store_data=False) + rel_err = np.abs(np.linalg.norm(uFFT) - np.linalg.norm(u0FFT)) / np.linalg.norm( + u0FFT + ) assert rel_err < 1e-2 + def test_if34_nondiag(): - xint,u0int,w0int,L,NL = allen_cahn_setup() - solver = IF34(linop=L,NLfunc=NL,epsilon=1e-3) - wfint = solver.evolve(w0int,t0=0,tf=60,store_data=False) + xint, u0int, w0int, L, NL = allen_cahn_setup() + solver = IF34(linop=L, NLfunc=NL, epsilon=1e-3) + wfint = solver.evolve(w0int, t0=0, tf=60, store_data=False) ufint = wfint.real + xint - assert np.abs(u0int[0]-ufint[0]) < 0.01 - assert np.abs(u0int[7]-ufint[7]) > 1 + assert np.abs(u0int[0] - ufint[0]) < 0.01 + assert np.abs(u0int[7] - ufint[7]) > 1 + def test_if45dp(): - u0FFT,L,NL = burgers_setup() - solver = IF45DP(linop=L,NLfunc=NL,epsilon=1e-3) - uFFT = solver.evolve(u0FFT,t0=0,tf=0.85,store_data=False) - rel_err = np.abs(np.linalg.norm(uFFT)-np.linalg.norm(u0FFT))/np.linalg.norm(u0FFT) + u0FFT, L, NL = burgers_setup() + solver = IF45DP(linop=L, NLfunc=NL, epsilon=1e-3) + uFFT = solver.evolve(u0FFT, t0=0, tf=0.85, store_data=False) + rel_err = np.abs(np.linalg.norm(uFFT) - np.linalg.norm(u0FFT)) / np.linalg.norm( + u0FFT + ) assert rel_err < 1e-2 def test_etd4_step(): - u0FFT,L,NL,uexactFFT,h,steps = kdv_soliton_setup() - solver = ETD4(linop=L,NLfunc=NL) - kdv_CS_step_eval(solver,u0FFT,uexactFFT,h,steps,1e-6) + u0FFT, L, NL, uexactFFT, h, steps = kdv_soliton_setup() + solver = ETD4(linop=L, NLfunc=NL) + kdv_CS_step_eval(solver, u0FFT, uexactFFT, h, steps, 1e-6) + def test_etd4_evolve(): - u0FFT,L,NL,uexactFFT,h,steps = kdv_soliton_setup() - solver = ETD4(linop=L,NLfunc=NL) - kdv_evolve_eval(solver,u0FFT,uexactFFT,h,tf=h*steps,tol=1e-6) + u0FFT, L, NL, uexactFFT, h, steps = kdv_soliton_setup() + solver = ETD4(linop=L, NLfunc=NL) + kdv_evolve_eval(solver, u0FFT, uexactFFT, h, tf=h * steps, tol=1e-6) + def test_etd5_step(): - u0FFT,L,NL,uexactFFT,h,steps = kdv_soliton_setup() - solver = ETD5(linop=L,NLfunc=NL) - kdv_CS_step_eval(solver,u0FFT,uexactFFT,h,steps,tol=1e-6) + u0FFT, L, NL, uexactFFT, h, steps = kdv_soliton_setup() + solver = ETD5(linop=L, NLfunc=NL) + kdv_CS_step_eval(solver, u0FFT, uexactFFT, h, steps, tol=1e-6) + def test_etd5_evolve(): - u0FFT,L,NL,uexactFFT,h,steps = kdv_soliton_setup() - solver = ETD5(linop=L,NLfunc=NL) - kdv_evolve_eval(solver,u0FFT,uexactFFT,h=h,tf=h*steps,tol=1e-6) + u0FFT, L, NL, uexactFFT, h, steps = kdv_soliton_setup() + solver = ETD5(linop=L, NLfunc=NL) + kdv_evolve_eval(solver, u0FFT, uexactFFT, h=h, tf=h * steps, tol=1e-6) -def test_if4_step(): - u0FFT,L,NL,uexactFFT,h,steps = kdv_soliton_setup() - solver = IF4(linop=L,NLfunc=NL) - kdv_CS_step_eval(solver,u0FFT,uexactFFT,h,steps,tol=1e-5) -def test_if4_evolve(): - u0FFT,L,NL,uexactFFT,h,steps = kdv_soliton_setup() - solver = IF4(linop=L,NLfunc=NL) - kdv_evolve_eval(solver,u0FFT,uexactFFT,h=h,tf=h*steps,tol=1e-5) +def test_if4_step(): + u0FFT, L, NL, uexactFFT, h, steps = kdv_soliton_setup() + solver = IF4(linop=L, NLfunc=NL) + kdv_CS_step_eval(solver, u0FFT, uexactFFT, h, steps, tol=1e-5) +def test_if4_evolve(): + u0FFT, L, NL, uexactFFT, h, steps = kdv_soliton_setup() + solver = IF4(linop=L, NLfunc=NL) + kdv_evolve_eval(solver, u0FFT, uexactFFT, h=h, tf=h * steps, tol=1e-5) diff --git a/tests/test_transforms.py b/tests/test_transforms.py index acbb755..1b09bb6 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1,37 +1,36 @@ -import rkstiff.grids +import rkstiff.grids import numpy as np + def test_hankel(): - HT = rkstiff.grids.HankelTransform(nr=50,rmax=2.0) + HT = rkstiff.grids.HankelTransform(nr=50, rmax=2.0) a = 4 - f1 = np.exp(-a**2*HT.r**2) - fsp1_ex = np.exp(-HT.kr**2/(4*a**2))/(2*a**2) + f1 = np.exp(-(a**2) * HT.r**2) + fsp1_ex = np.exp(-HT.kr**2 / (4 * a**2)) / (2 * a**2) fsp1 = HT.ht(f1) - error1 = np.linalg.norm(fsp1-fsp1_ex)/np.linalg.norm(fsp1_ex) + error1 = np.linalg.norm(fsp1 - fsp1_ex) / np.linalg.norm(fsp1_ex) assert error1 < 1e-10 HT.rmax = 4 sigma = 2 w = 0.5 - f2 = np.exp(-sigma*HT.r**2)*np.sin(w*HT.r**2) - omega = 1./(4*(sigma**2+w**2)) - fsp2_ex = -2*omega*np.exp(-sigma*omega*HT.kr**2)*(-w*np.cos(w*omega*HT.kr**2)+sigma*np.sin(w*omega*HT.kr**2)) + f2 = np.exp(-sigma * HT.r**2) * np.sin(w * HT.r**2) + omega = 1.0 / (4 * (sigma**2 + w**2)) + fsp2_ex = ( + -2 + * omega + * np.exp(-sigma * omega * HT.kr**2) + * (-w * np.cos(w * omega * HT.kr**2) + sigma * np.sin(w * omega * HT.kr**2)) + ) fsp2 = HT.ht(f2) - error2 = np.linalg.norm(fsp2-fsp2_ex)/np.linalg.norm(fsp2_ex) + error2 = np.linalg.norm(fsp2 - fsp2_ex) / np.linalg.norm(fsp2_ex) assert error2 < 1e-10 - HT.nr = 25 HT.rmax = 3 - f3 = np.exp(-a*HT.r) - fsp3_ex = a*np.power(a**2+HT.kr**2,-3./2) + f3 = np.exp(-a * HT.r) + fsp3_ex = a * np.power(a**2 + HT.kr**2, -3.0 / 2) fsp3 = HT.ht(f3) - error3 = np.linalg.norm(fsp3-fsp3_ex)/np.linalg.norm(fsp3_ex) + error3 = np.linalg.norm(fsp3 - fsp3_ex) / np.linalg.norm(fsp3_ex) assert error3 < 1e-2 - - - - - -