Adding basic code needed for PDE tutorials

PDEs/Chebyshev.py

+import numpy as np
+
+
+def ChebyshevInterpolate(u, ToSpec, x_out):
+    """Interpolate the data u[i], assumed to be
+    real-space values of a chebyshev expansion,
+    to the grid-points in x_out.
+    ToSpec is the transformation matrix from
+    grid-point values to Chebyshev coefficients.
+    """
+
+    theta = np.arccos(x_out)  # avoid recomputation
+    uspec = ToSpec @ u
+    out = x_out * 0.0
+    for k, coef in enumerate(uspec):
+        out += coef * np.cos(k * theta)
+    return out
+
+
+def ChebyshevHelpers(N, a=-1.0, b=1.0):
+    """ "return several useful objects for dealing with Chebyshev polynomials
+    PARAMETERS:
+      N - order of returned objects
+      [a,b]  -- coordinate range to be considered.  I.e. the usual [-1,1] is linearly mapped to [a,b]
+
+    RETURNS
+      x  -- collocation points  np.array of length N covering [a,b]
+            xi[0] is at the lower bound a
+      ToSpEC -- transformation matrix from real-space values (at xi) to spectral coefficients
+      ToPhys -- transformation matrix from spectral coefficients to real space values (at xi)
+      D1     -- transformation matrix of first derivative (real space -> real space)
+      D2     -- transformation matrix of second derivative (real space -> real space)
+    """
+
+    # linear mapping
+    #  X in [1, 1]  ->   x in [a,b]:  x = A X + B
+    A = 0.5 * (b - a)
+    B = 0.5 * (a + b)
+
+    def MappedCoords(X):
+        return A * X + B
+
+    # collocation points
+    # Kidder(2000), Eq. (3.14)
+    # '-' to make Xi[0] = lower boundary
+    Xi = -np.cos(np.pi * np.arange(N) / (N - 1))
+    xi = MappedCoords(Xi)
+
+    # helper arrays
+    # Kidder(2000), Eq. (3.13)
+    c = np.ones(N)
+    c[0] = 2.0
+    # Kidder (2000). Eq. (3.16)
+    cbar = np.ones(N)
+    cbar[0] = 2.0
+    cbar[N - 1] = 2.0
+
+    # trafo matrix from spectral to grid-points
+    # (this simply evaluates the Chebyshev polynomials
+    #  T_k = cos(k*arccos(X)) at the collocation points X)
+    ToPhys = np.ndarray((N, N))
+    for k in range(N):
+        ToPhys[:, k] = np.cos(k * np.arccos(Xi))
+
+    # trafo matrix from grid-points to spectral coefficients
+    #    rationale for (*) below:  if multiplied by another
+    #    ToPhys a.k.a.  T_k(x_n), then the left-hand-side collapses
+    #    to identity, and one has Eq. (3.15) from Kidder.
+    ToSpec = np.ndarray((N, N))
+    for k in range(N):
+        ToSpec[k, :] = 2.0 / ((N - 1) * cbar[k] * cbar[:]) * ToPhys[:, k]  # (*)
+
+    # print(ToSpec @ ToPhys)
+    # print(ToPhys @ ToSpec)
+
+    # spectral differentiaton matrix
+    # Kidder (2000) Eq. (3.12) as a matrix
+    Dtilde = np.ndarray((N, N))
+    Dtilde.fill(0.0)
+    # I am sure there's a better way
+    for k in reversed(range(N - 1)):
+        if k <= N - 3:
+            Dtilde[k, :] = 1 / c[k] * Dtilde[k + 2, :]
+        Dtilde[k, k + 1] += 2 * (k + 1) / c[k]
+    # multiply by -1 to account for our choice that the first collocation
+    # point is at the __lower__ boundary
+    # so far Dtilde is in primitive coordinates (X), to get
+    # d/dx = dX/dx d/dX,  must multiply by dX/dx = 1/A
+    Dtilde *= 1.0 / A
+
+    # full first derivative physical-to-physical differentiaton matrix
+    D1 = ToPhys @ Dtilde @ ToSpec
+
+    # full second derivative physical-to-physical differentiaton matrix
+    D2 = ToPhys @ Dtilde @ Dtilde @ ToSpec
+
+    return xi, ToSpec, ToPhys, D1, D2

PDEs/Fourier.py

+import numpy as np
+
+def RealToFourier(u):
+    """Given data u, at uniformly distributed grid-points
+    0<=x<1,  calculates Fourier coefficients A[k] such that
+
+    u = 0.5*A[0] + Sum_k(  A[k].real * cos(2 pi k x)
+                         - A[k].imag * sin(2 pi k x)
+                        )
+
+    returns complex array A
+    """
+    A=np.fft.rfft(u)
+    # adjust signs and magnitudes
+    A *= (2./len(u))
+    return A 
+
+def FourierToReal(A):
+    """Given spectral coefficients A[k], reconstruct the function
+   0.5*A[0] + Sum_k(  A[k].real * cos(2 pi k x)
+                    - A[k].imag * sin(2 pi k x)
+                   )
+    """
+    u = np.fft.irfft(A)
+    u = u * 0.5*len(u)
+    return u
+
+def dudx_Fourier(u):
+    """Compute the first derivative dudx, using Foutier spectral methods.
+    Assumes data is equally spaced in interval 0<=x<2pi and
+    periodic."""
+    A=RealToFourier(u)
+    # multiply A[k] by  j k, where j is the imagimnary unit
+    k= np.linspace(0, len(A)-1, len(A))
+    A *= 1j*k
+    # convert back tp real-space values
+    dudx=FourierToReal(A)
+    return dudx
+
+
+
+
+def FourierHelpers_Complex(N, a=0, b=2*np.pi):
+    """
+    computes useful objects for coding of spectral methods for
+    a real-valued Fourier-series truncated at order N/2:
+
+        u(x) = Re (\sum_{n=0}^{N/2} c_n e^{inx})     (1)
+
+    CONVENTION:
+      storage of the coefficients is in one complex array of length N/2+1
+         coefs = [a0, a_1-j*b_1, a_2-j*b_2, ... a_(N/2-1)-j*b_{N/2-1}, a_(N/2)]
+      where a_n is the coefficient of cos(n*x), and b_n the coefficient of sin(n*x)
+         
+      That means coefs[0] and coefs[N/2] are real, thus we have the full N 
+      degrees of freedom. 
+"""
+    if N%2!=0: raise "FourierHelpers works only for even N"
+
+     #  X in [0, 2*pi]  ->   x in [a,b]:  x = A X + B 
+    A=(b-a)/(2*np.pi)
+    B=a
+    def MappedCoords(X):
+        return A*X + B
+
+    M = int(N/2)  # biggest represented mode
+    X = np.linspace(0., 2*np.pi, N, endpoint=False)
+    x = MappedCoords(X)
+    ToSpec = np.zeros( (M+1, N), dtype=complex)
+    # cast to coefficients
+    cbar=np.ones(M)
+    cbar[0]=cbar[-1]=2.
+    for j in range(0,M+1):  
+        ToSpec[j,:]=2./N*np.exp(-1j * j * X)
+    ToSpec[ 0,:] /= 2
+    ToSpec[-1,:] /= 2
+
+    ToPhys = np.zeros( (N,M+1), dtype=complex)
+    # contains coefficients of u(x) = \sum_{n=0}^{N/2} a_n cos(nx) + \sum_{n=1}^{N/2-1} b_n sin(nx)
+    for j in range(0,M+1):  
+        ToPhys[:,j]=np.exp(1j * j * X)
+
+    Dtilde = np.zeros( (M+1,M+1), dtype=complex )
+    for k in range(0,M):
+        Dtilde[k,   k ] =  1j*k
+        # Dspec[M,M]=0, b /c cos(Mx)'=sin(Mx), which cannot be represented
+    # so far Dtilde is in primitive coordinates (X), to get
+    # d/dx = dX/dx d/dX,  must multiply by dX/dx = 1/A
+    Dtilde *= 1./A
+
+    # for grid-to-grid differentiation matrices 
+    # only real part matters, owing to the Re im Eq. (1) above
+    D1 = np.real(ToPhys @ Dtilde @ ToSpec)
+    D2 = np.real(ToPhys @ Dtilde @ Dtilde @ ToSpec)
+
+    return x, ToSpec, ToPhys, D1, D2
+
+
+
+
+def FourierHelpers_Real(N, a=0, b=2*np.pi):
+    """
+    computes various real-space -> real-space transformation matrices
+    for a Fourier basis, for real-valued functions.  In total we have 
+    N degrees of freedom (about half of them being cosine and sine-modes).
+       
+    We define our spectral coefficients a_n and b_n such that 
+
+        u(x) = \sum_{n=0}^{N/2} a_n cos(nx) + \sum_{n=1}^{N/2-1} b_n sin(nx)
+
+    We require N to be even;  b_0 and b_N/2 do not exist, because the bass-function
+    is zero at all collocation points.  
+
+    CONVENTION:
+      storage of the coefficients is in one real array of length N:
+         coefs = [a0, ... a_(N/2), b1, ... b_(N/2-1)]
+      This is a compact output format w/o any zeros, but makes for awkward indexing
+      if one wants to find a specific mode. 
+
+    RETURNS
+      x  -- collocation points  np.array of length N covering [a,b]
+            x[0] is at the lower bound a
+      ToSpEC -- transformation matrix from real-space values (at xi) to spectral coefficients
+      ToPhys -- transformation matrix from spectral coefficients to real space values (at xi)
+      D1     -- transformation matrix of first derivative (real space -> real space)
+      D2     -- transformation matrix of second derivative (real space -> real space)
+
+      x is a np.array of length(N)
+      all other quantities are NxN matrices
+"""
+    if N%2!=0: raise "FourierHelpers works only for even N"
+
+    M = int(N/2)  # biggest represented mode
+
+    #  X in [0, 2*pi]  ->   x in [a,b]:  x = A X + B 
+    A=(b-a)/(2*np.pi)
+    B=a
+    def MappedCoords(X):
+        return A*X + B
+
+    X = np.linspace(0., 2*np.pi, N, endpoint=False)
+    x=MappedCoords(X)
+    A=np.zeros( (M+1, N))
+    B=np.zeros( (M-1, N))
+    A[0,:]=1/N   #  a_0 = 1/N sum_j u_j
+    for j in range(1,M+1):
+        A[j,:]=2/N*np.cos(j*X)      #  a_j = 2/N sum_k u_k cos(j*x_k)
+    A[M,:] /= 2
+    for j in range(1,M):
+        B[j-1,:]=2/N*np.sin(j*X)  #  b_j = 2/N sum_k bu_k sin(j*x_k)
+    
+    # combine into one:
+    ToSpec = np.zeros( (N,N) )
+    ToSpec[0:M+1,:] = A
+    ToSpec[M+1:N,:] = B
+
+
+    ToPhys = np.zeros( (N,N) )
+    # contains coefficients of u(x) = \sum_{n=0}^{N/2} a_n cos(nx) + \sum_{n=1}^{N/2-1} b_n sin(nx)
+    for j in range(0,M+1):  # cos(N/2) is present
+        ToPhys[:,j]=np.cos(j*X)
+    for j in range(1,M): # sin(N/2 x) is _not_ present, so stop at M
+        ToPhys[:,M+j]=np.sin(j*X)
+
+    Dtilde = np.zeros( (N,N) )
+    # a_0 -> 0, so don't touch Dspec[0,:]
+    #   a_k cos(kx) + b_k sin(k x) ->   k*b_k cos(kx) - k*a_k sin(k x)
+    for k in range(1,M):
+        Dtilde[k,   M+k ] =  k
+        Dtilde[M+k, k   ] = -k
+    # highest a coefficient
+    #   cos(N/2 x) -> -sin(N/2 x) ==>>  ZERO, because the sin(N/2 x) is not represented
+    # this is accounted for because D[N/2,:] is never touched
+
+    # so far Dtilde is in primitive coordinates (X), to get
+    # d/dx = dX/dx d/dX,  must multiply by dX/dx = 1/A
+    Dtilde *= 1./A
+
+    D1 = ToPhys @ Dtilde @ ToSpec
+    D2 = ToPhys @ Dtilde @ Dtilde @ ToSpec
+
+    return x, ToSpec, ToPhys, D1, D2
\ No newline at end of file

PDEs/TimeSteppers.py

+import math
+
+
+def FE_Step(t, u, F, dt, info):
+    """Perform one Forward Euler step
+       - t current time
+       - u current solution
+       - F function that computes the right-hand-side 'F' in
+            du/dt = F
+         The caling sequence is F(t, u, info)
+       - dt time-step
+       - info extra information needed by the right-hand-side function
+         (e.g. mass-matrices for DG)
+    returns
+       t_new, u_new
+    """
+    u_new = u + dt * F(t, u, info)
+    return t + dt, u_new
+
+
+def RK2_Step(t, u, F, dt, info):
+    """ "Perform one Runge-Kutta 2 step
+       - t current time
+       - u current solution
+       - F function that computes the right-hand-side 'F' in
+            du/dt = F
+         The caling sequence is F(t, u, info)
+       - dt time-step
+       - info extra information needed by the right-hand-side function
+         (e.g. mass-matrices for DG)
+    returns
+       t_new, u_new
+    """
+
+    k1 = dt * F(t, u, info)
+    k2 = dt * F(t + 0.5 * dt, u + 0.5 * k1, info)
+    return (t + dt, u + k2)
+
+
+def RK4_Step(t, u, F, dt, info):
+    """ "Perform one Runge-Kutta 4 step
+       - t current time
+       - u current solution
+       - F function that computes the right-hand-side 'F' in
+            du/dt = F
+         The caling sequence is F(t, u, info)
+       - dt time-step
+       - info extra information needed by the right-hand-side function
+         (e.g. mass-matrices for DG)
+    returns
+       t_new, u_new
+    """
+
+    k1 = dt * F(t, u, info)
+    k2 = dt * F(t + 0.5 * dt, u + 0.5 * k1, info)
+    k3 = dt * F(t + 0.5 * dt, u + 0.5 * k2, info)
+    k4 = dt * F(t + dt, u + k3, info)
+    return (t + dt, u + k1 / 6.0 + k2 / 3.0 + k3 / 3.0 + k4 / 6.0)
+
+
+def Evolve(t, t_final, u, F, Tstepper, CF, info):
+    """Evolve the evolution equations represented by right-hand-side 'F'
+    with time-stepper Tstepper until final time 't_final'.
+      - t current time
+      - t_final final time
+      - u solution at current time 't'
+      - F right-hand-side of evolution equations, w/ calling sequence F(t, u, info)
+      - Tstepper time-stepper with calling sequence Tstepper(t, u, F, dt, info)
+      - CF courant-factor, i.e dt will be chosen such that dt < C dxmin
+      - info class holding any additional information necessary.  It is passed into 'rhs'.
+        It is also assumed that info.dxmin returns the minimal grid-spacing.
+    returns
+      t_final, u_final"""
+
+    # time-step
+    dt = CF * info.dxmin
+    Nsteps = math.ceil((t_final - t) / dt)  # round up number of steps
+    dt = (t_final - t) / Nsteps  # adjust dt to precisely reach t_final
+    for i in range(Nsteps):
+        t, u = Tstepper(t, u, F, dt, info)
+    return t, u