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()