# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
from __future__ import annotations
r""" Hartree correction for FFT-Poisson solver for arbitrary electrode positions
Developer: Nick Papior
Contact: nickpapior <at> gmail.com
sisl-version: >=0.9.3
This Poisson solver uses pyamg to calculate an initial guess for the Poisson
solution to correct the FFT solution. It does this by setting up boundary
conditions on electrodes and then solving the Hartree potential using multi-grid
solvers.
It requires two inputs and has several optional flags.
- The siesta.TBT.nc file which contains the geometry that is to be calculated for
The reason for using the siesta.TBT.nc file is the ease of use:
The siesta.TBT.nc contains electrode atoms and device atoms. Hence it
becomes easy to read in the electrode atomic positions.
Note that since you'll always do a 0 V calculation it is easy to
do a 0 bias calculation first, then create a guess using this script,
and finally do bias calculations.
- The grid size of the simulation grid, this needs to be commensurate with the
actual Siesta grid used.
This script is a command-line utility with several options (please refer to
--help). There are a few important flags you should know about:
--tolerance [tol] specify the tolerance of the solution, the tighter the longer solution time
--shape-solver [nx ny nz] shape for which the solution is calculated (impacts speed)
--shape [nx ny nz] final shape of the solution, if shape-solver is not the same the solution will be interpolated (order=2)
--dtype [f|d] the data-type used to solve the Poisson equation
--out [file] any sisl compatible grid file, please at least do --out V.TSV.nc which is compatible with TranSiesta.
This tool requires the following packages:
- pyamg
Known problems:
- The pyamg solver requires quite a bit of memory, you should preferentially select
the largest grid (up to the actual grid size you want) possible.
- The Neumann implementation in the boundary conditions is not correct, hence
it may never converge (or produce nan's). If this happens please try another
boundary condition.
- It may not always converge which requires some fine-tuning of the tolerances,
secondly it may converge too fast so the solution is not really good.
"""
import argparse as argp
import os
import sys
from pathlib import Path
import numpy as np
import sisl as si
__all__ = ["pyamg_solve", "solve_poisson", "fftpoisson_fix_cli", "fftpoisson_fix_run"]
_BC = si.BoundaryCondition
# Base-script name
_script = Path(sys.argv[0]).name
_DEBUG = os.environ.get("SISL_TS_FFT_DEBUG", "False")
# Determine from env-var whether we should use debug mode
_DEBUG = _DEBUG.lower() in ("true", "t", "1", "y", "yes", "on")
def pyamg_solve(A, b, tolerance=1e-12, accel=None, title=""):
import pyamg
print(f"\nSetting up pyamg solver... {title}")
ml = pyamg.aggregation.smoothed_aggregation_solver(A, max_levels=1000)
del A
print(ml)
residuals = []
def callback(x):
# residuals calculated in the solve function is a pre-conditioned residual
# residuals.append(np.linalg.norm(b - A.dot(x)) ** 0.5)
print(
" {:4d} residual = {:.5e} x0-residual = {:.5e}".format(
len(residuals) - 1, residuals[-1], residuals[-1] / residuals[0]
)
)
x = ml.solve(
b,
tol=tolerance,
callback=callback,
residuals=residuals,
accel=accel,
cycle="W",
maxiter=1e7,
)
print("Done solving the Poisson equation!")
return x
[docs]
def solve_poisson(
geometry,
shape,
radius="empirical",
dtype=np.float64,
tolerance=1e-8,
accel=None,
boundary_fft=True,
device_val=None,
plot_boundary=False,
box=False,
boundary=None,
**elecs_V,
):
"""Solve Poisson equation"""
error = False
elecs = []
for name in geometry.names:
if ("+" in name) or (name in ["Buffer", "Device"]):
continue
# This is actually an electrode
elecs.append(name)
error = error or (name not in elecs_V)
if len(elecs) == 0:
raise ValueError(f"{_script}: Could not find any electrodes in the geometry.")
error = error or len(elecs) != len(elecs_V)
if error:
for name in elecs:
if not name in elecs_V:
print(f" missing electrode bias: {name}")
raise ValueError(
f"{_script}: Missing electrode arguments for specifying the bias."
)
if boundary is None:
bc = [[_BC.PERIODIC, _BC.PERIODIC] for _ in range(3)]
else:
bc = []
def bc2bc(s):
return {
"periodic": "PERIODIC",
"p": "PERIODIC",
_BC.PERIODIC: "PERIODIC",
"dirichlet": "DIRICHLET",
"d": "DIRICHLET",
_BC.DIRICHLET: "DIRICHLET",
"neumann": "NEUMANN",
"n": "NEUMANN",
_BC.NEUMANN: "NEUMANN",
}.get(s.lower(), s.upper())
for bottom, top in boundary:
bc.append([getattr(_BC, bc2bc(bottom)), getattr(_BC, bc2bc(top))])
if len(bc) != 3:
raise ValueError(
f"{_script}: Requires a 3x2 list input for the boundary conditions."
)
def _create_shape_tree(xyz, A, B=None):
"""Takes two lists A and B which returns a shape with a binary nesting
This makes further index handling much faster.
"""
if B is None or len(B) == 0:
return _create_shape_tree(xyz, *np.array_split(A, 2))
AA, BB = None, None
if len(A) == 1:
AA = si.Sphere(radius, xyz[A[0]])
if len(B) == 0:
return AA
if len(B) == 1:
BB = si.Sphere(radius, xyz[B[0]])
if len(A) == 0:
return BB
# Quick return if these are the final ones
if AA and BB:
return AA | BB
if not AA:
AA = _create_shape_tree(xyz, *np.array_split(A, 2))
if not BB:
BB = _create_shape_tree(xyz, *np.array_split(B, 2))
return AA | BB
# Create grid and specify boundary conditions
geometry.lattice.set_boundary_condition(bc)
grid = si.Grid(shape, geometry=geometry, dtype=dtype)
class _fake:
@property
def shape(self):
return shape
@property
def dtype(self):
return dtype
# Fake the grid to reduce memory requirement
grid.grid = _fake()
# Construct matrices we need to specify the boundary conditions on
A, b = grid.to.pyamg()
# Short-hand notation
xyz = geometry.xyz
if not device_val is None:
print(f"\nApplying device potential = {device_val}")
idx = geometry.names["Device"]
device = _create_shape_tree(xyz, idx)
idx = grid.index_truncate(grid.index(device))
idx = grid.pyamg_index(idx)
grid.pyamg_fix(A, b, idx, device_val)
# Apply electrode constants
print("\nApplying electrode potentials")
for i, elec in enumerate(elecs):
V = elecs_V[elec]
print(f" - {elec} = {V}")
idx = geometry.names[elec]
elec_shape = _create_shape_tree(xyz, idx)
idx = grid.index_truncate(grid.index(elec_shape))
idx = grid.pyamg_index(idx)
grid.pyamg_fix(A, b, idx, V)
del idx, elec_shape
# Now we have initialized both A and b with correct boundary conditions
# Lets solve the Poisson equation!
if box:
# No point in solving the boundary problem if requesting a box
boundary_fft = False
grid.grid = b.reshape(shape)
del A
else:
x = pyamg_solve(
A,
b,
tolerance=tolerance,
accel=accel,
title="solving electrode boundary conditions",
)
grid.grid = x.reshape(shape)
del A, b
if boundary_fft:
# Change boundaries to always use dirichlet
# This ensures that once we set the boundaries we don't
# get any side-effects
BC = si.BoundaryCondition
periodic = [
bc == BC.PERIODIC or geometry.nsc[i] > 1
for i, bc in enumerate(grid.lattice.boundary_condition[:, 0])
]
bc = np.repeat(np.array([BC.DIRICHLET], np.int32), 6).reshape(3, 2)
for i in (0, 1, 2):
if periodic[i]:
bc[i, :] = BC.PERIODIC
grid.lattice.set_boundary_condition(bc)
A, b = grid.to.pyamg()
# Solve only for the boundary fixed
def sl2idx(grid, sl):
return grid.pyamg_index(grid.mgrid(sl))
# Create slices
sl = [slice(0, g) for g in grid.shape]
# One boundary at a time
for i in (0, 1, 2):
if periodic[i]:
continue
new_sl = sl[:]
new_sl[i] = slice(0, 1)
idx = sl2idx(grid, new_sl)
grid.pyamg_fix(
A, b, idx, grid.grid[new_sl[0], new_sl[1], new_sl[2]].reshape(-1)
)
new_sl[i] = slice(grid.shape[i] - 1, grid.shape[i])
idx = sl2idx(grid, new_sl)
grid.pyamg_fix(
A, b, idx, grid.grid[new_sl[0], new_sl[1], new_sl[2]].reshape(-1)
)
if plot_boundary:
dat = b.reshape(*grid.shape)
# now plot every plane
import matplotlib.pyplot as plt
slicex3 = np.index_exp[:] * 3
axs = [
np.linspace(0, grid.lattice.length[ax], shape, endpoint=False)
for ax, shape in enumerate(grid.shape)
]
for i in (0, 1, 2):
idx = list(slicex3)
j = (i + 1) % 3
k = (i + 2) % 3
if i > j:
i, j = j, i
X, Y = np.meshgrid(axs[i], axs[j])
for v, head in ((0, "bottom"), (-1, "top")):
plt.figure()
plt.title(f"axis: {'ABC'[k]} ({head})")
idx[k] = v
plt.contourf(X, Y, dat[tuple(idx)].T)
plt.xlabel(f"Distance along {'ABC'[i]} [Ang]")
plt.ylabel(f"Distance along {'ABC'[j]} [Ang]")
plt.colorbar()
plt.show()
grid.grid = _fake()
x = pyamg_solve(
A,
b,
tolerance=tolerance,
accel=accel,
title="removing electrode boundaries and solving for edge fixing",
)
grid.grid = x.reshape(shape)
del A, b
return grid
def fftpoisson_fix_cli(subp=None):
is_sub = not subp is None
title = "FFT Poisson corrections for TranSiesta calculations for arbitrary number of electrodes."
if is_sub:
global _script
_script = f"{_script} ts-fft"
p = subp.add_parser("ts-fft", description=title, help=title)
else:
p = argp.ArgumentParser(title)
tuning = p.add_argument_group(
"tuning", "Tuning fine details of the Poisson calculation."
)
p.add_argument(
"--geometry",
"-G",
default="siesta.TBT.nc",
metavar="FILE",
help="siesta.TBT.nc file which contains the geometry and electrode information, currently we cannot read that from fdf-files.",
)
p.add_argument(
"--shape",
"-s",
nargs=3,
type=int,
required=True,
metavar=("A", "B", "C"),
help="Grid shape, this *has* to be conforming to the TranSiesta calculation, read from output: 'InitMesh: MESH = A x B x C'",
)
n = {"a": "first", "b": "second", "c": "third"}
for d in "abc":
p.add_argument(
f"--boundary-condition-{d}",
f"-bc-{d}",
nargs=2,
type=str,
default=["p", "p"],
metavar=("BOTTOM", "TOP"),
help=(
"Boundary condition along the {} lattice vector [periodic/p, neumann/n, dirichlet/d]. "
"Specify separate BC at the start and end of the lattice vector, respectively.".format(
n[d]
)
),
)
p.add_argument(
"--elec-V",
"-V",
action="append",
nargs=2,
metavar=("NAME", "V"),
default=[],
help="Specify chemical potential on electrode",
)
p.add_argument(
"--shape-solver",
"-ss",
"--pyamg-shape",
"-ps",
nargs=3,
type=int,
metavar=("A", "B", "C"),
default=None,
help="Grid used to solve the Poisson equation, if shape is different the Grid will be interpolated (order=2) after.",
)
p.add_argument(
"--device",
"-D",
type=float,
default=None,
metavar="VAL",
help="Fix the value of all device atoms to a value. In some cases this turns out to yield a better box boundary. The default is to *not* fix the potential on the device atoms.",
)
tuning.add_argument(
"--radius",
"-R",
type=float,
default=3.0,
metavar="R",
help=(
"Radius of atoms when figuring out the electrode sizes, this corresponds to the extend of "
"each electrode where boundary conditions are fixed. Should be tuned according to the atomic species [3 Ang]"
),
)
tuning.add_argument(
"--dtype",
"-d",
choices=["d", "f64", "f", "f32"],
default="d",
help="Precision of data (d/f64==double, f/f32==single)",
)
tuning.add_argument(
"--tolerance",
"-T",
type=float,
default=1e-7,
metavar="EPS",
help="Precision required for the pyamg solver. NOTE when using single precision arrays this should probably be on the order of 1e-5",
)
tuning.add_argument(
"--acceleration",
"-A",
dest="accel",
default="cg",
metavar="METHOD",
help="""Acceleration method for pyamg. May be useful if it fails to converge
Try one of: cg, gmres, fgmres, cr, cgnr, cgne, bicgstab, steepest_descent, minimal_residual""",
)
test = p.add_argument_group(
"testing",
"Options used for testing output. None of these options should be used for production runs!",
)
test.add_argument(
"--box",
dest="box",
action="store_true",
default=False,
help="Only store the initial box solution (i.e. do not run PyAMG)",
)
test.add_argument(
"--no-boundary-fft",
action="store_false",
dest="boundary_fft",
default=True,
help="Disable the 2nd solution with the boundaries fixed. Generally, this solver will first solve for the electrode boundary conditions, then for the fixed box boundary conditains (gotten from the first solution).",
)
if _DEBUG:
test.add_argument(
"--plot",
dest="plot",
default=None,
type=int,
help="Plot grid by averaging over the axis given as argument",
)
test.add_argument(
"--plot-boundary",
dest="plot_boundary",
action="store_true",
help="Plot all 6 edges of the box with their fixed values (just before 2nd pyamg solve step)",
)
p.add_argument(
"--out",
"-o",
action="append",
default=None,
help="Output file to store the resulting Poisson solution. It *has* to have TSV.nc file ending to make the file conforming with TranSiesta.",
)
if is_sub:
p.set_defaults(runner=fftpoisson_fix_run)
else:
fftpoisson_fix_run(p.parse_args())
def fftpoisson_fix_run(args):
if args.out is None:
print(
f">\n>\n>{_script}: No out-files has been specified, work will be carried out but not saved!\n>\n>\n"
)
# Read in geometry
geometry = si.get_sile(args.geometry).read_geometry()
# Figure out the electrodes
elecs_V = {}
if len(args.elec_V) == 0:
print(geometry.names)
raise ValueError(
f"{_script}: Please specify all electrode potentials using --elec-V"
)
for name, V in args.elec_V:
elecs_V[name] = float(V)
if args.dtype.lower() in ("f", "f32"):
dtype = np.float32
elif args.dtype.lower() in ("d", "f64"):
dtype = np.float64
# Now we can solve Poisson
if args.shape_solver is None:
shape = args.shape
else:
shape = args.shape_solver
# Create the boundary conditions
boundary = []
boundary.append(args.boundary_condition_a)
boundary.append(args.boundary_condition_b)
boundary.append(args.boundary_condition_c)
V = solve_poisson(
geometry,
shape,
radius=args.radius,
boundary=boundary,
dtype=dtype,
tolerance=args.tolerance,
box=args.box,
accel=args.accel,
boundary_fft=args.boundary_fft,
device_val=args.device,
plot_boundary=args.plot_boundary,
**elecs_V,
)
if _DEBUG:
if not args.plot is None:
dat = V.average(args.plot)
import matplotlib.pyplot as plt
axs = [
np.linspace(0, V.lattice.length[ax], shape, endpoint=False)
for ax, shape in enumerate(V.shape)
]
idx = list(range(3))
# Now plot data
del axs[args.plot]
del idx[args.plot]
X, Y = np.meshgrid(*axs)
plt.contourf(X, Y, np.squeeze(dat.grid).T)
plt.colorbar()
plt.title(f"Averaged over {'ABC'[args.plot]} axis")
plt.xlabel(f"Distance along {'ABC'[idx[0]]} [Ang]")
plt.ylabel(f"Distance along {'ABC'[idx[1]]} [Ang]")
plt.show()
if np.any(np.array(args.shape) != np.array(V.shape)):
print("\nInterpolating the solution...")
V = V.interp(args.shape, 2)
print("Done interpolating!")
print("")
# Write solution to the output
if not args.out is None:
for out in args.out:
print(f"Writing to file: {out}...")
V.write(out)
# Import object holding all the CLI
from sisl_toolbox.cli import register_toolbox_cli
register_toolbox_cli(fftpoisson_fix_cli)
if __name__ == "__main__":
fftpoisson_fix_cli()