import numpy as np
from pySDC.core.sweeper import Sweeper
[docs]
class boris_2nd_order(Sweeper):
"""
Custom sweeper class, implements Sweeper.py
Second-order sweeper using velocity-Verlet with Boris scheme as base integrator
Attributes:
S: node-to-node collocation matrix (first order)
SQ: node-to-node collocation matrix (second order)
ST: node-to-node trapezoidal matrix
Sx: node-to-node Euler half-step for position update
"""
def __init__(self, params):
"""
Initialization routine for the custom sweeper
Args:
params: parameters for the sweeper
"""
# call parent's initialization routine
if "QI" not in params:
params["QI"] = "IE"
if "QE" not in params:
params["QE"] = "EE"
super(boris_2nd_order, self).__init__(params)
# S- and SQ-matrices (derived from Q) and Sx- and ST-matrices for the integrator
[
self.S,
self.ST,
self.SQ,
self.Sx,
self.QQ,
self.QI,
self.QT,
self.Qx,
self.Q,
] = self.__get_Qd()
self.qQ = np.dot(self.coll.weights, self.coll.Qmat[1:, 1:])
def __get_Qd(self):
"""
Get integration matrices for 2nd-order SDC
Returns:
S: node-to-node collocation matrix (first order)
SQ: node-to-node collocation matrix (second order)
ST: node-to-node trapezoidal matrix
Sx: node-to-node Euler half-step for position update
"""
# set implicit and explicit Euler matrices (default, but can be changed)
QI = self.get_Qdelta_implicit(qd_type=self.params.QI)
QE = self.get_Qdelta_explicit(qd_type=self.params.QE)
# trapezoidal rule
QT = 1 / 2 * (QI + QE)
# Qx as in the paper
Qx = np.dot(QE, QT) + 1 / 2 * QE * QE
Sx = np.zeros(np.shape(self.coll.Qmat))
ST = np.zeros(np.shape(self.coll.Qmat))
S = np.zeros(np.shape(self.coll.Qmat))
# fill-in node-to-node matrices
Sx[0, :] = Qx[0, :]
ST[0, :] = QT[0, :]
S[0, :] = self.coll.Qmat[0, :]
for m in range(self.coll.num_nodes):
Sx[m + 1, :] = Qx[m + 1, :] - Qx[m, :]
ST[m + 1, :] = QT[m + 1, :] - QT[m, :]
S[m + 1, :] = self.coll.Qmat[m + 1, :] - self.coll.Qmat[m, :]
# SQ via dot-product, could also be done via QQ
SQ = np.dot(S, self.coll.Qmat)
# QQ-matrix via product of Q
QQ = np.dot(self.coll.Qmat, self.coll.Qmat)
return [S, ST, SQ, Sx, QQ, QI, QT, Qx, self.coll.Qmat]
[docs]
def update_nodes(self):
"""
Update the u- and f-values at the collocation nodes -> corresponds to a single sweep over all nodes
Returns:
None
"""
# get current level and problem description
L = self.level
P = L.prob
# only if the level has been touched before
assert L.status.unlocked
# get number of collocation nodes for easier access
M = self.coll.num_nodes
# initialize integral terms with zeros, will add stuff later
integral = [P.dtype_u(P.init, val=0.0) for l in range(M)]
# gather all terms which are known already (e.g. from the previous iteration)
# this corresponds to SF(u^k) - SdF(u^k) + tau (note: have integrals in pos and vel!)
for m in range(M):
for j in range(M + 1):
# build RHS from f-terms (containing the E field) and the B field
f = P.build_f(L.f[j], L.u[j], L.time + L.dt * self.coll.nodes[j - 1])
# add SQF(u^k) - SxF(u^k) for the position
integral[m].pos += L.dt * (L.dt * (self.SQ[m + 1, j] - self.Sx[m + 1, j]) * f)
# add SF(u^k) - STF(u^k) for the velocity
integral[m].vel += L.dt * (self.S[m + 1, j] - self.ST[m + 1, j]) * f
# add tau if associated
if L.tau[m] is not None:
integral[m] += L.tau[m]
# tau is 0-to-node, need to change it to node-to-node here
if m > 0:
integral[m] -= L.tau[m - 1]
# do the sweep
for m in range(0, M):
# build rhs, consisting of the known values from above and new values from previous nodes (at k+1)
tmp = P.dtype_u(integral[m])
for j in range(m + 1):
# build RHS from f-terms (containing the E field) and the B field
f = P.build_f(L.f[j], L.u[j], L.time + L.dt * self.coll.nodes[j - 1])
# add SxF(u^{k+1})
tmp.pos += L.dt * (L.dt * self.Sx[m + 1, j] * f)
# add pos at previous node + dt*v0
tmp.pos += L.u[m].pos + L.dt * self.coll.delta_m[m] * L.u[0].vel
# set new position, is explicit
L.u[m + 1].pos = tmp.pos
# get E field with new positions and compute mean
L.f[m + 1] = P.eval_f(L.u[m + 1], L.time + L.dt * self.coll.nodes[m])
ck = tmp.vel
# do the boris scheme
L.u[m + 1].vel = P.boris_solver(ck, L.dt * np.diag(self.QI)[m + 1], L.f[m], L.f[m + 1], L.u[m])
# indicate presence of new values at this level
L.status.updated = True
return None
[docs]
def integrate(self):
"""
Integrates the right-hand side
Returns:
list of dtype_u: containing the integral as values
"""
# get current level and problem description
L = self.level
P = L.prob
# create new instance of dtype_u, initialize values with 0
p = []
for m in range(1, self.coll.num_nodes + 1):
p.append(P.dtype_u(P.init, val=0.0))
# integrate RHS over all collocation nodes, RHS is here only f(x,v)!
for j in range(1, self.coll.num_nodes + 1):
f = P.build_f(L.f[j], L.u[j], L.time + L.dt * self.coll.nodes[j - 1])
p[-1].pos += L.dt * (L.dt * self.QQ[m, j] * f) + L.dt * self.coll.Qmat[m, j] * L.u[0].vel
p[-1].vel += L.dt * self.coll.Qmat[m, j] * f
return p
[docs]
def compute_end_point(self):
"""
Compute u at the right point of the interval
The value uend computed here is a full evaluation of the Picard formulation (always!)
Returns:
None
"""
# get current level and problem description
L = self.level
P = L.prob
# start with u0 and add integral over the full interval (using coll.weights)
L.uend = P.dtype_u(L.u[0])
for m in range(self.coll.num_nodes):
f = P.build_f(L.f[m + 1], L.u[m + 1], L.time + L.dt * self.coll.nodes[m])
L.uend.pos += L.dt * (L.dt * self.qQ[m] * f) + L.dt * self.coll.weights[m] * L.u[0].vel
L.uend.vel += L.dt * self.coll.weights[m] * f
# add up tau correction of the full interval (last entry)
if L.tau[-1] is not None:
L.uend += L.tau[-1]
return None
[docs]
def get_sweeper_mats(self):
"""
Returns the matrices Q, QQ, Qx, QT which define the sweeper.
"""
Q = self.Q[1:, 1:]
QQ = self.QQ[1:, 1:]
Qx = self.Qx[1:, 1:]
QT = self.QT[1:, 1:]
return Q, QQ, Qx, QT
[docs]
def get_scalar_problems_sweeper_mats(self, lambdas=None):
"""
This function returns the corresponding matrices of an SDC sweep matrix formulation
Args:
lambdas (numpy.narray): the first entry in lambdas is k-spring constant and the second is mu friction.
"""
Q, QQ, Qx, QT = self.get_sweeper_mats()
if lambdas is None:
pass
# should use lambdas from attached problem and make sure it is scalar SDC
raise NotImplementedError("At the moment, the values for the lambda have to be provided")
else:
k = lambdas[0]
mu = lambdas[1]
nnodes = self.coll.num_nodes
dt = self.level.dt
F = np.block(
[
[-k * np.eye(nnodes), -mu * np.eye(nnodes)],
[-k * np.eye(nnodes), -mu * np.eye(nnodes)],
]
)
C_coll = np.block([[np.eye(nnodes), dt * Q], [np.zeros([nnodes, nnodes]), np.eye(nnodes)]])
Q_coll = np.block(
[
[dt**2 * QQ, np.zeros([nnodes, nnodes])],
[np.zeros([nnodes, nnodes]), dt * Q],
]
)
Q_vv = np.block(
[
[dt**2 * Qx, np.zeros([nnodes, nnodes])],
[np.zeros([nnodes, nnodes]), dt * QT],
]
)
M_vv = np.eye(2 * nnodes) - np.dot(Q_vv, F)
return C_coll, Q_coll, Q_vv, M_vv, F
[docs]
def get_scalar_problems_manysweep_mats(self, nsweeps, lambdas=None):
"""
For a scalar problem, K sweeps of SDC can be written in matrix form.
Args:
nsweeps (int): number of sweeps
lambdas (numpy.ndarray): the first entry in lambdas is k-spring constant and the second is mu friction.
"""
nnodes = self.coll.num_nodes
C_coll, Q_coll, Q_vv, M_vv, F = self.get_scalar_problems_sweeper_mats(lambdas=lambdas)
K_sdc = np.dot(np.linalg.inv(M_vv), Q_coll - Q_vv) @ F
Keig, Kvec = np.linalg.eig(K_sdc)
Kp_sdc = np.linalg.matrix_power(K_sdc, nsweeps)
Kinv_sdc = np.linalg.inv(np.eye(2 * nnodes) - K_sdc)
Kdot_sdc = np.dot(np.eye(2 * nnodes) - Kp_sdc, Kinv_sdc)
MC = np.dot(np.linalg.inv(M_vv), C_coll)
Mat_sweep = Kp_sdc + np.dot(Kdot_sdc, MC)
return Mat_sweep, np.max(np.abs(Keig))
[docs]
def get_scalar_problems_picardsweep_mats(self, nsweeps, lambdas=None):
"""
For a scalar problem, K sweeps of SDC can be written in matrix form.
Args:
nsweeps (int): number of sweeps
lambdas (numpy.ndarray): the first entry in lambdas is k-spring constant and the second is mu friction.
"""
nnodes = self.coll.num_nodes
C_coll, Q_coll, Q_vv, M_vv, F = self.get_scalar_problems_sweeper_mats(lambdas=lambdas)
K_sdc = np.dot(Q_coll, F)
Keig, Kvec = np.linalg.eig(K_sdc)
Kp_sdc = np.linalg.matrix_power(K_sdc, nsweeps)
Kinv_sdc = np.linalg.inv(np.eye(2 * nnodes) - K_sdc)
Kdot_sdc = np.dot(np.eye(2 * nnodes) - Kp_sdc, Kinv_sdc)
Mat_sweep = Kp_sdc + np.dot(Kdot_sdc, C_coll)
return Mat_sweep, np.max(np.abs(Keig))