def test_tol_norm_called(self):
# Check that supplying tol_norm keyword to nonlin_solve works
self._tol_norm_used = False
def local_norm_func(x):
self._tol_norm_used = True
return np.absolute(x).max()
nonlin.newton_krylov(F, F.xin, f_tol=1e-2, maxiter=200, verbose=0,
tol_norm=local_norm_func)
assert_(self._tol_norm_used)
def test_tol_norm_called(self):
# Check that supplying tol_norm keyword to nonlin_solve works
self._tol_norm_used = False
def local_norm_func(x):
self._tol_norm_used = True
return np.absolute(x).max()
nonlin.newton_krylov(F, F.xin, f_tol=1e-2, maxiter=200, verbose=0,
tol_norm=local_norm_func)
assert_(self._tol_norm_used)
def SCFsqueeze(chi, chi_s, pdi, sigma, phi_b, segments, layers, disp=False):
""" Return a self consistent field within a specified number of layers.
"""
bs = BasicSystem()
parameters = (chi, chi_s, pdi, sigma, phi_b, segments)
u = bs.walk(parameters, disp).squeeze()
squeezing = layers - len(u) < 0
jac_solve_method = 'gmres'
while layers - len(u):
if squeezing:
u = np.delete(u, -1)
else:
u = np.append(u, u[-1])
u = newton_krylov(
bs.field_equations(parameters),
u[None, :],
verbose=bool(disp),
maxiter=30,
method=jac_solve_method,
).squeeze()
phi = bs.field_equations(parameters)(u, 1).squeeze()
return phi
def SCFsolve(field_equations, u_jz_guess, disp=False, maxiter=30):
"""Solve SCF equations using an initial guess and lattice parameters
This function finds a solution for the equations where the lattice size
is sufficiently large.
The Newton-Krylov solver really makes this one. With gmres, it was faster
than the other solvers by quite a lot.
"""
if disp: starttime = perf_counter()
# resizing loop variables
jac_solve_method = 'gmres'
lattice_too_small = True
# We tolerate up to 1 ppm deviation from bulk
# when counting layers_near_bulk
tol = 1e-6
while lattice_too_small:
if disp: print("Solving SCF equations")
try:
u_jz = newton_krylov(
field_equations,
u_jz_guess,
verbose=bool(disp),
maxiter=maxiter,
method=jac_solve_method,
)
except RuntimeError as e:
if str(e) == 'gmres is not re-entrant':
# Threads are racing to use gmres. Lose the race and use
# something slower but thread-safe.
jac_solve_method = 'lgmres'
continue
else:
raise
if disp:
print('lattice size:', u_jz.shape[1])
field_deviation = fabs(u_jz).sum(axis=0)
layers_near_bulk = field_deviation < tol
nbulk = layers_near_bulk.sum()
lattice_too_small = nbulk < MINBULK
if lattice_too_small:
# if there aren't enough layers_near_zero, grow the lattice 20%
newlayers = max(1, round((u_jz_guess.shape[1]) * 0.2)) + 2
if disp: print('Growing undersized lattice by', newlayers)
# find the layer closest to eqm with the bulk
i = field_deviation.argmin()
# make newlayers there via linear interpolation for each species
interpolation = [
np.linspace(field_z[i - 1], field_z[i], num=newlayers)
for field_z in u_jz
]
# then sandwich the interpolated layers between the originals
# interpolation includes the first and last points, so shift i
u_jz_guess = np.hstack(
(u_jz[:, :i - 1], interpolation, u_jz[:, i + 1:]))
# TODO: vectorize this interpolation and stacking?
if nbulk > 2 * MINBULK:
chop_end = np.diff(layers_near_bulk).nonzero()[0].max()
chop_start = chop_end - MINBULK
i = np.arange(u_jz.shape[1])
u_jz = u_jz[:, (i <= chop_start) | (i > chop_end)]
if disp:
print("SCFsolve execution time:", round(perf_counter() - starttime, 3),
"s")
print('lattice size:', u_jz.shape[1])
return u_jz
def run(
method,
eqn="advection",
limiter="weno",
cflnum=0.9,
r=8.0,
N=128,
IC="sin",
BC="periodic",
solver="krylov",
guesstype="BE",
ls="-k",
plotall=False,
squarenum=30,
):
"""
Run advection or Burgers equation test. Options:
method =
be
itrap
dwrk - optimal 2nd order downwind implicit SSP RK
fe - Forward Euler
eqn = 'advection' or 'burgers'
limiter = 'weno' or 'vanleer'
r is a parameter used to define the downwind RK method; if another method is used
then r is irrelevant (see the paper for details)
IC = sin, sinp, gaussian, or square (initial condition)
solver = fsolve or krylov (always tries fsolve as a last resort anyway)
guesstype = BE, BE_fine, itrap (different approaches to finding a good initial guess for the
nonlinear solve)
"""
# Set up grid
dx = 1.0 / N
nghost = 3
N2 = N + 2 * nghost
x = np.linspace(-(nghost - 0.5) * dx, 1.0 + (nghost - 0.5) * dx, N2)
t = 0.0
if IC == "sin":
q0 = np.sin(2 * np.pi * x)
myaxis = (0.0, 1.0, -1.1, 1.1)
elif IC == "sinp":
q0 = np.sin(2 * np.pi * x) + 1.5
myaxis = (0.0, 1.0, 0.4, 2.6)
elif IC == "gaussian":
q0 = np.exp(-200 * (x - 0.3) ** 2) # Smooth
elif IC == "square":
q0 = 1.0 * (x > 0.0) * (x < 0.5) # Discontinuous
myaxis = (0.0, 1.0, -0.1, 1.1)
elif IC == "BL-Riemann": # Buckley-Leverett Riemann problem
q0 = 1.0 * (x > 0.5) + 0.5 * (x < 0.0)
print q0[0]
myaxis = (0.0, 1.0, -0.1, 1.1)
else:
raise Exception("Unrecognized value of IC.")
q = q0.copy()
# pl.plot(x,q); pl.draw(); pl.hold(False)
if eqn == "advection":
flux = lambda q: q
dt = cflnum * dx
T = 1.0
elif eqn == "burgers":
flux = lambda q: 0.5 * q ** 2
dt = cflnum * dx / np.max(q)
T = 0.16
elif eqn == "buckley-leverett":
a = 1.0 / 3.0
flux = lambda q: q ** 2 / (q ** 2 + a * (1.0 - q) ** 2)
maxvel = 2 * a * 0.25
dt = cflnum * dx / maxvel
T = 0.25
else:
raise Exception("Unrecognized eqn")
while t < T:
# If we're nearly at the end, choose the step size to reach T exactly
if dt > T - t:
dt = T - t
# Take one step using the chosen method
if method == "be": # Backward Euler
nonlinearf = lambda (y): be(q, dt, dx, y, flux, limiter=limiter, BC=BC)
qnew = fsolve(nonlinearf, q)
q[:] = qnew[:]
elif method == "itrap": # Implicit Trapezoidal rule
nonlinearf = lambda (y): itrap_solve(q, dt, dx, y, flux, limiter=limiter, BC=BC)
y1 = fsolve(nonlinearf, q)
y1hat = np.empty([1, len(y1)])
y1hat[0, :] = y1
ql, qr = limit(y1hat, limiter)
q[1:] = q[1:] - dt / dx * (flux(qr[0, 1:]) - flux(qr[0, :-1]))
elif method == "dwrk": # 2nd-order Downwind Runge-Kutta
guess = setguess(q, dt, dx, flux, guesstype, r)
nonlinearf = lambda (y): dwrk_solve(q, dt, dx, y, flux, r, limiter=limiter, BC=BC)
if solver == "fsolve":
Y, info, ier, mesg = fsolve(nonlinearf, guess, full_output=1)
if ier > 1:
return ier, mesg
elif solver == "krylov":
try:
Y = newton_krylov(nonlinearf, guess, method="lgmres", maxiter=100)
except:
print "Trying fsolve as last resort"
Y, info, ier, mesg = fsolve(nonlinearf, guess, full_output=1)
if ier > 1:
return ier, mesg
Y = np.reshape(Y, (2, -1), "C")
y2 = np.empty([1, len(q)])
y2[0, :] = Y[1, :]
ql, qr = limit(y2, limiter)
q[1:] = y2[0, 1:] - dt / dx * (flux(qr[0, 1:]) - flux(qr[0, :-1])) / r
elif method == "fe":
qq = np.empty([1, len(q)])
qq[:, 0] = q
ql, qr = limit(qq, limiter)
q[1:] = q[1:] - dt / dx * (flux(qr[1:, 0]) - flux(qr[:-1, 0]))
else:
print "error: unrecognized method"
if BC == "periodic":
q[0:nghost] = q[-2 * nghost : -nghost] # Periodic boundary
q[-nghost:] = q[nghost : 2 * nghost] # Periodic boundary
# Else fixed Dirichlet; do nothing
t = t + dt
if plotall:
pl.plot(x, q, ls, linewidth=3, markersize=7, markevery=int(N / N))
pl.hold(True)
pl.axis(myaxis)
pl.title(str(t))
pl.draw()
pl.hold(False)
print t
pl.plot(x, q, ls, linewidth=3, markersize=7, markevery=int(N / squarenum))
pl.hold(True)
pl.axis(myaxis)
pl.title(str(t))
pl.draw()
