Source code for sisl_toolbox.siesta.atom._atom
# 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""" Atom input/output writer
Developer: Nick Papior
Contact: nickpapior <at> gmail.com
sisl-version: >0.10.0
This enables reading input files for atom and also helps in parsing output files.
To read in an input file one can do:
>>> atom = AtomInput.from_input("INP")
>>> atom.pg("OUT")
which will read and write the same file.
One can also plot output from `atom` using:
>>> atom.plot()
which will show 4 plots for different sections. A command-line tool is also
made available through the `stoolbox`.
"""
import sys
from functools import reduce
from pathlib import Path
import numpy as np
try:
from scipy.integrate import trapezoid
except ImportError:
from scipy.integrate import trapz as trapezoid
from scipy.interpolate import interp1d
import sisl as si
from sisl._internal import set_module
from sisl.utils import NotNonePropertyDict, PropertyDict
__all__ = ["AtomInput", "atom_plot_cli"]
_script = Path(sys.argv[0]).name
# We don't use the Siesta units here...
_Ang2Bohr = si.units.convert("Ang", "Bohr")
# Atom/Siesta uses this order of occupation for
# core electrons
# idx | shell | electrons | cum. electrons | noble gas
# 1 | 1s | 2 | 2 | [He]
# 2 | 2s | 2 | 4 | [He]2s2
# 3 | 2p | 6 | 10 | [Ne]
# 4 | 3s | 2 | 12 | [Ne]3s
# 5 | 3p | 6 | 18 | [Ar]
# 6 | 3d | 10 | 28 | [Ar]3d10
# 7 | 4s | 2 | 30 | [Ar]3d10 2s2
# 8 | 4p | 6 | 36 | [Kr]
# 9 | 4d | 10 | 46 | [Kr]4d10
# 10 | 5s | 2 | 48 | [Kr]4d10 5s2
# 11 | 5p | 6 | 54 | [Xe]
# 12 | 4f | 14 | 68 | [Xe]4f14
# 13 | 5d | 10 | 78 | [Xe]4f14 5d10
# 14 | 6s | 2 | 80 | [Xe]4f14 5d10 6s2
# 15 | 6p | 6 | 86 | [Rn]
# 16 | 7s | 2 | 88 | [Rn]7s2
# 17 | 5f | 14 | 102 | [Rn]7s2 5f14
# 18 | 6d | 10 | 112 | [Rn]7s2 5f14 6d10
# 19 | 7p | 6 | 118 | [Og]
_shell_order = [
# Occupation shell order
"1s", # [He]
"2s",
"2p", # [Ne]
"3s",
"3p", # [Ar]
"3d",
"4s",
"4p", # [Kr]
"4d",
"5s",
"5p", # [Xe]
"4f",
"5d",
"6s",
"6p", # [Rn]
"7s",
"5f",
"6d",
"7p", # [Og]
]
_spdfgh = "spdfgh"
@set_module("sisl_toolbox.siesta.atom")
class AtomInput:
"""Input for the ``atom`` program see [AtomLicense]_
This class enables the construction of the ``INP`` file to be fed to ``atom``.
::
# Example input for ATOM
#
# Comments allowed here
#
# ae Si ground state all-electron
# Si car
# 0.0
# 3 2
# 3 0 2.00 0.00
# 3 1 2.00 0.00
#
# Comments allowed here
#
#2345678901234567890123456789012345678901234567890 Ruler
References
----------
.. [AtomLicense] https://siesta.icmab.es/SIESTA_MATERIAL/Pseudos/atom_licence.html
"""
def __init__(
self, atom, define=("NEW_CC", "FREE_FORMAT_RC_INPUT", "NO_PS_CUTOFFS"), **opts
):
# opts = {
# "flavor": "tm2",
# "xc": "pb",
# optionally libxc
# "equation": "r",
# "logr": 2.
# "cc": False,
# "rcore": 2.
# }
self.atom = atom
assert isinstance(atom, si.Atom)
if "." in self.atom.tag:
raise ValueError("The atom 'tag' must not contain a '.'!")
# We need to check that atom has 4 orbitals, with increasing l
# We don't care about n or any other stuff, so these could be
# SphericalOrbital, for that matter
l = 0
for orb in self.atom:
if orb.l != l:
raise ValueError(
f"{self.__class__.__name__} atom argument does not have "
f"increasing l quantum number index {l} has l={orb.l}"
)
l += 1
if l != 4:
raise ValueError(
f"{self.__class__.__name__} atom argument must have 4 orbitals. "
f"One for each s-p-d-f shell"
)
self.opts = PropertyDict(**opts)
# Check options passed and define defaults
self.opts.setdefault("equation", "r")
if self.opts.equation not in " rs":
# ' ' == non-polarized
# s == polarized
# r == relativistic
raise ValueError(
f"{self.__class__.__name__} failed to initialize; opts{'equation': <v>} has wrong value, should be [ rs]."
)
if self.opts.equation == "s":
raise NotImplementedError(
f"{self.__class__.__name__} does not implement spin-polarized option (use relativistic)"
)
self.opts.setdefault("flavor", "tm2")
if self.opts.flavor not in ("hsc", "ker", "tm2"):
# hsc == Hamann-Schluter-Chiang
# ker == Kerker
# tm2 == Troullier-Martins
raise ValueError(
f"{self.__class__.__name__} failed to initialize; opts{'flavor': <v>} has wrong value, should be [hsc|ker|tm2]."
)
self.opts.setdefault("logr", 2.0)
# default to true if set
self.opts.setdefault("cc", "rcore" in self.opts)
# rcore only used if cc is True
self.opts.setdefault("rcore", 0.0)
self.opts.setdefault("xc", "pb")
# Read in the core valence shells for this atom
# figure out what the default value is.
# We do this my finding the minimum index of valence shells
# in the _shell_order list, then we use that as the default number
# of core-shells occpupied
# e.g if the minimum valence shell is 2p, it would mean that
# _shell_order.index("2p") == 2
# which has 1s and 2s occupied.
try:
core = reduce(
min,
(_shell_order.index(f"{orb.n}{_spdfgh[orb.l]}") for orb in atom),
len(_shell_order),
)
except Exception:
core = -1
self.opts.setdefault("core", core)
if self.opts.core == -1:
raise ValueError(
f"Default value for {self.atom.symbol} not added, please add core= at instantiation"
)
# Store the defined names
if define is None:
define = []
elif isinstance(define, str):
define = [define]
# store
self.define = define
[docs]
@classmethod
def from_input(cls, inp):
"""Return atom object respecting the input
Parameters
----------
inp : list or str
create `AtomInput` from the content of `inp`
"""
def _get_content(f):
if f.is_file():
return open(f, "r").readlines()
return None
if isinstance(inp, (tuple, list)):
# it is already in correct format
pass
elif isinstance(inp, (str, Path)):
# convert to path
inp = Path(inp)
# Check if it is a path or an input
content = _get_content(inp)
if content is None:
content = _get_content(inp / "INP")
if content is None:
raise ValueError(
f"Could not find any input file in {str(inp)} or {str(inp / 'INP')}"
)
inp = content
else:
raise ValueError(f"Unknown input format inp={inp}?")
# Now read lines
defines = []
opts = PropertyDict()
def bypass_comments(inp):
if inp[0].startswith("#"):
inp.pop(0)
bypass_comments(inp)
def bypass(inp, defines):
bypass_comments(inp)
if inp[0].startswith("%define"):
line = inp.pop(0)
defines.append(line.split()[1].strip())
bypass(inp, defines)
bypass(inp, defines)
# Now prepare reading
# First line has to contain the *type* of calculation
# pg|pe|ae|pt <comment>
line = inp.pop(0).strip()
if line.startswith("pg"):
opts.cc = False
elif line.startswith("pe"):
opts.cc = True
# <flavor> logr?
line = inp.pop(0).strip().split()
opts.flavor = line[0]
if len(line) >= 2:
opts.logr = float(line[1]) / _Ang2Bohr
# <element> <xc>' rs'?
line = inp.pop(0)
symbol = line.split()[0]
# now get xc equation
if len(line) >= 11:
opts.equation = line[10:10]
opts.xc = line[:10].split()[1]
line = line.split()
if len(line) >= 3:
opts.libxc = int(line[2])
# currently not used line
inp.pop(0)
# core, valence
core, valence = inp.pop(0).split()
opts.core = int(core)
valence = int(valence)
orbs = []
for _ in range(valence):
n, l, *occ = inp.pop(0).split()
orb = PropertyDict()
orb.n = int(n)
orb.l = int(l)
# currently we don't distinguish between up/down
orb.q0 = sum(map(float, occ))
orbs.append(orb)
# now we read the line with rc's and core-correction
rcs = inp.pop(0).split()
if len(rcs) >= 6:
# core-correction
opts.rcore = float(rcs[5]) / _Ang2Bohr
for orb in orbs:
orb.R = float(rcs[orb.l]) / _Ang2Bohr
# Now create orbitals
orbs = [si.AtomicOrbital(**orb, m=0, zeta=1) for orb in orbs]
# now re-arrange ensuring we have correct order of l shells
orbs = sorted(orbs, key=lambda orb: orb.l)
atom = si.Atom(symbol, orbs)
return cls(atom, defines, **opts)
[docs]
@classmethod
def from_yaml(cls, file, nodes=()):
"""Parse the yaml file"""
from sisl_toolbox.siesta.minimizer._yaml_reader import parse_variable, read_yaml
dic = read_yaml(file, nodes)
element = dic["element"]
tag = dic.get("tag")
mass = dic.get("mass", None)
# get default options for pseudo
opts = NotNonePropertyDict()
pseudo = dic["pseudo"]
opts["logr"] = parse_variable(pseudo.get("log-radius"), unit="Ang").value
opts["rcore"] = parse_variable(pseudo.get("core-correction"), unit="Ang").value
opts["xc"] = pseudo.get("xc")
opts["equation"] = pseudo.get("equation")
opts["flavor"] = pseudo.get("flavor")
define = pseudo.get(
"define", ("NEW_CC", "FREE_FORMAT_RC_INPUT", "NO_PS_CUTOFFS")
)
# Now on to parsing the valence shells
orbs = []
for key in dic:
if key not in _shell_order:
continue
# Now we know the occupation is a shell
pseudo = dic[key].get("pseudo", {})
cutoff = parse_variable(pseudo.get("cutoff"), 2.1, "Ang").value
charge = parse_variable(pseudo.get("charge"), 0.0).value
orbs.append(si.AtomicOrbital(key, m=0, R=cutoff, q0=charge))
atom = si.Atom(element, orbs, mass=mass, tag=tag)
return cls(atom, define, **opts)
[docs]
def write_header(self, f):
f.write("# This file is generated by sisl pseudo\n")
# Define all names
for define in self.define:
f.write(f"%define {define.upper()}\n")
[docs]
def write_generation(self, f):
if self.opts.cc:
# use core corrections
pg = "pe"
else:
# do not use core corrections
pg = "pg"
logr = self.opts.logr * _Ang2Bohr
f.write(f" {pg:2s} {self.atom.symbol} pseudo potential\n")
if logr < 0.0:
f.write(f" {self.opts.flavor:3s}\n")
else:
f.write(f" {self.opts.flavor:3s}{logr:9.3f}\n")
xc = self.opts.xc
equation = self.opts.equation
rcore = self.opts.rcore * _Ang2Bohr
f.write(f" {self.atom.symbol:2s} {xc:2s}{equation:1s}")
if "libxc" in self.opts:
f.write(f" {self.opts.libxc:8d}")
f.write(f"\n {0.0:5.1f}\n")
# now extract the charges for each orbital
atom = self.atom
core = self.opts.core
valence = len(atom)
f.write(f"{core:5d}{valence:5d}\n")
Rs = [0.0] * 4 # always 4: s, p, d, f
for orb in sorted(atom.orbitals, key=lambda x: x.l):
# Write the configuration of this orbital
n, l = orb.n, orb.l
f.write(f"{n:5d}{l:5d}{orb.q0:10.3f}{0.0:10.3f}\n")
Rs[l] = orb.R * _Ang2Bohr
f.write(
f"{Rs[0]:10.7f} {Rs[1]:10.7f} {Rs[2]:10.7f} {Rs[3]:10.7f} {0.0:10.7f} {rcore:10.7f}\n"
)
[docs]
def write_all_electron(self, f, charges=(0.0,)):
q0 = self.atom.q0.sum()
xc = self.opts.xc
equation = self.opts.equation
core = self.opts.core
for charge in charges:
if isinstance(charge, dict):
out = self.excite(**charge)
else:
out = self.excite(charge)
q = out.atom.q0.sum()
f.write(f" ae {out.atom.symbol} # all-electron q={q0 - q}\n")
f.write(f" {out.atom.symbol:2s} {xc:2s}{equation:1s}")
if "libxc" in out.opts:
f.write(f" {out.opts.libxc:8d}")
f.write(f"\n {0.0:5.1f}\n")
# now extract the charges for each orbital
atom = out.atom
valence = len(atom)
f.write(f"{core:5d}{valence:5d}\n")
for orb in sorted(atom.orbitals, key=lambda x: x.l):
n, l = orb.n, orb.l
f.write(f"{n:5d}{l:5d}{orb.q0:10.3f}{0.0:10.3f}\n")
[docs]
def write_test(self, f, charges=(0,)):
q0 = self.atom.q0.sum()
xc = self.opts.xc
equation = self.opts.equation
core = self.opts.core
for charge in charges:
if isinstance(charge, dict):
out = self.excite(**charge)
else:
out = self.excite(charge)
q = out.atom.q0.sum()
f.write(f" pt {out.atom.symbol} # pseudo test q={q0 - q}\n")
f.write(f" {out.atom.symbol:2s} {xc:2s}{equation:1s}")
if "libxc" in out.opts:
f.write(f" {out.opts.libxc:8d}")
f.write(f"\n {0.0:5.1f}\n")
# now extract the charges for each orbital
atom = out.atom
valence = len(atom)
f.write(f"{core:5d}{valence:5d}\n")
for orb in sorted(atom.orbitals, key=lambda x: x.l):
n, l = orb.n, orb.l
f.write(f"{n:5d}{l:5d}{orb.q0:10.3f}{0.0:10.3f}\n")
def _get_out(self, path, filename):
if path is None:
return Path(filename)
return Path(path) / Path(filename)
[docs]
def __call__(self, filename="INP", path=None):
"""Open a file and return self"""
out = self._get_out(path, filename)
self._enter = open(out, "w"), self.atom
return self
def __enter__(self):
if not hasattr(self, "_enter"):
self()
return self._enter[0]
def __exit__(self, exc_type, exc_value, traceback):
if hasattr(self, "_enter"):
self.atom = self._enter[1]
del self._enter
[docs]
def ae(self, filename="INP", path=None):
with self(filename, path) as f:
self.write_header(f)
self.write_all_electron(f)
[docs]
def pg(self, filename="INP", path=None):
with self(filename, path) as f:
self.write_header(f)
self.write_generation(f)
[docs]
def excite(self, *charge, **lq):
"""Excite contained atom to another charge
Notes
-----
This is charge, *not* electrons.
"""
if len(charge) > 1:
raise ValueError(
f"{self.__class__.__name__}.excite takes only "
"a single argument or [spdf]=charge arguments"
)
elif len(charge) == 1:
charge = charge[0]
if "charge" in lq:
raise ValueError(
f"{self.__class__.__name__}.excite does not accept "
"both charge as argument and keyword argument"
)
else:
charge = 0
charge = lq.pop("charge", charge)
# get indices of the orders
shell_idx = [
_shell_order.index(f"{orb.n}{_spdfgh[orb.l]}") for orb in self.atom
]
def _charge(idx):
# while the g and h shells are never used, we just have them...
return {"s": 2, "p": 6, "d": 10, "f": 14, "g": 18, "h": 22}[
_shell_order[idx][1]
]
orig_charge = [_charge(idx) for idx in shell_idx]
atom = self.atom.copy()
# find order of orbitals (from highest index to lowest)
# depending on the sign of charge
shell_sort = shell_idx[:]
shell_sort.sort(reverse=charge > 0)
for l, q in lq.items():
l = _spdfgh.index(l)
for orb in atom:
if orb.l == l:
orb._q0 -= q
assert orb.q0 >= 0.0
# now finalize the charge
while abs(charge) > 1e-9:
for idx_sort in shell_sort:
idx = shell_idx.index(idx_sort)
orb = atom[idx]
if charge > 0 and orb.q0 > 0:
dq = min(orb.q0, charge)
elif charge < 0 and orb.q0 < orig_charge[idx]:
dq = max(orb.q0 - orig_charge[idx], charge)
else:
dq = 0.0
orb._q0 -= dq
charge -= dq
return self.__class__(atom, self.define, **self.opts)
[docs]
def plot(
self,
path=None,
plot=("wavefunction", "charge", "log", "potential"),
l="spdf",
show=True,
):
"""Plot everything related to this psf file
Parameters
----------
path : str or pathlib.Path, optional
from which directory should files be read
plot : list-like of str, optional
which data to plot
l : list-like, optional
which l-shells to plot (for those that have l-shell decompositions)
show : bool, optional
call `matplotlib.pyplot.show()` at the end
Returns
-------
fig : figure for axes
axs : axes used for plotting
"""
import matplotlib.pyplot as plt
if path is None:
path = Path.cwd()
else:
path = Path(path)
def get_xy(f, yfactors=None):
"""Return x, y data from file `f` with y being calculated as the factors between the columns"""
nonlocal path
f = path / f
if not f.is_file():
print(f"Could not find file: {str(f)}")
return None, None
data = np.loadtxt(f)
ncol = data.shape[1]
if yfactors is None:
yfactors = [0, 1]
yfactors = np.pad(yfactors, (0, ncol - len(yfactors)), constant_values=0.0)
x = data[:, 0]
y = (data * yfactors.reshape(1, -1)).sum(1)
return x, y
l2i = {
"s": 0,
0: 0,
"p": 1,
1: 1,
"d": 2,
2: 2,
"f": 3,
3: 3,
"g": 4,
4: 4, # never used
}
# Get this atoms default calculated binding length
# We use this one since there are many missing elements
# in vdw table.
# And convert to Bohr
atom_r = self.atom.radius("calc") * _Ang2Bohr
def plot_wavefunction(ax):
# somewhat similar to ae.gplot
ax.set_title("Wavefunction")
ax.set_xlabel("Radius [Bohr]")
for shell in l:
il = l2i[shell]
orb = self.atom.orbitals[il]
r, w = get_xy(f"AEWFNR{il}")
if not r is None:
p = ax.plot(r, w, label=f"AE {_spdfgh[il]}")
color = p[0].get_color()
ax.axvline(orb.R * _Ang2Bohr, color=color, alpha=0.5)
r, w = get_xy(f"PSWFNR{il}")
if not r is None:
ax.plot(r, w, "--", label=f"PS {_spdfgh[il]}")
ax.set_xlim(0, atom_r * 5)
ax.autoscale(enable=True, axis="y", tight=True)
ax.legend()
def plot_charge(ax):
ax.set_title("Charge")
ax.set_xlabel("Radius [Bohr]")
ax.set_ylabel("(4.pi.r^2) Charge [electrons/Bohr]")
# Get current core-correction length
ae_r, ae_cc = get_xy("AECHARGE", [0, 0, 0, 1])
_, ae_vc = get_xy("AECHARGE", [0, 1, 1, -1])
if not ae_cc is None:
p = ax.plot(ae_r, ae_cc, label=f"AE core")
color = p[0].get_color()
if self.opts.get("cc", False):
ax.axvline(self.opts.rcore * _Ang2Bohr, color=color, alpha=0.5)
ax.plot(ae_r, ae_vc, "--", label=f"AE valence")
ps_r, ps_cc = get_xy("PSCHARGE", [0, 0, 0, 1])
_, ps_vc = get_xy("PSCHARGE", [0, 1, 1])
if not ps_r is None:
ax.plot(ps_r, ps_cc, "--", label=f"PS core")
ax.plot(ps_r, ps_vc, ":", label=f"PS valence")
# Now determine the overlap between all-electron core-charge
# and the pseudopotential valence charge
if np.allclose(ae_r, ps_r):
# Determine dR
# dr = ae_r[1] - ae_r[0]
# Integrate number of core-electrons and valence electrons
core_c = trapezoid(ae_cc, ae_r)
valence_c = trapezoid(ps_vc, ps_r)
print(f"Total charge in atom: {core_c + valence_c:.5f}")
overlap_c = trapezoid(np.minimum(ae_cc, ps_vc), ae_r)
ax.set_title(f"Charge: int(min(AE_cc, PS_vc)) = {overlap_c:.3f} e")
# We will try and *guess-stimate* a good position for rc for core-corrections
# Javier Junquera's document says:
# r_pc has to be chosen such that the valence charge density is negligeable compared to
# the core one for r < r_pc.
# Tests show that it might be located where the core charge density is from 1 to 2 times
# larger than the valence charge density
with np.errstate(divide="ignore", invalid="ignore"):
fraction = ae_cc / ps_vc
np.nan_to_num(fraction, copy=False)
ax2 = (
ax.twinx()
) # instantiate a second axes that shares the same x-axis
ax2.plot(ae_r, fraction, "k", alpha=0.5, label="c/v")
marks = np.array([0.5, 1.0, 1.5])
min_x = (ae_r > 0.2).nonzero()[0].min()
max_x = (fraction[min_x:] > 0.0).nonzero()[0].max() + min_x
r_marks = interp1d(fraction[min_x:max_x], ae_r[min_x:max_x])(marks)
ax2.scatter(r_marks, marks, alpha=0.5)
ax2.set_ylim(0, 3)
print(f"Core-correction r_pc {marks}: {r_marks} Bohr")
ax.set_xlim(0, atom_r)
ax.set_ylim(0)
ax.legend()
def plot_log(ax):
ax.set_title("d-log of wavefunction")
ax.set_xlabel("Energy [Ry]")
ax.set_ylabel("Derivative of wavefunction")
for shell in l:
il = l2i[shell]
e, log = get_xy(f"AELOGD{il}")
emark = np.loadtxt(path / f"AEEV{il}")
if emark.ndim == 1:
emark.shape = (1, -1)
emark = emark[:, 0]
if not e is None:
p = ax.plot(e, log, label=f"AE {_spdfgh[il]}")
idx_mark = np.fabs(e.reshape(-1, 1) - emark.reshape(1, -1)).argmin(
axis=0
)
ax.scatter(emark, log[idx_mark], color=p[0].get_color(), alpha=0.5)
# And now PS
e, log = get_xy(f"PSLOGD{il}")
emark = np.loadtxt(path / f"PSEV{il}")
if emark.ndim == 1:
emark.shape = (1, -1)
emark = emark[:, 0]
if not e is None:
p = ax.plot(e, log, ":", label=f"PS {_spdfgh[il]}")
idx_mark = np.fabs(e.reshape(-1, 1) - emark.reshape(1, -1)).argmin(
axis=0
)
ax.scatter(emark, log[idx_mark], color=p[0].get_color(), alpha=0.5)
ax.legend()
def plot_potential(ax):
ax.set_title("Pseudopotential")
ax.set_xlabel("Radius [Bohr]")
ax.set_ylabel("Potential [Ry]")
for shell in l:
il = l2i[shell]
orb = self.atom.orbitals[il]
r, V = get_xy(f"PSPOTR{il}")
if not r is None:
p = ax.plot(r, V, label=f"PS {_spdfgh[il]}")
color = p[0].get_color()
ax.axvline(orb.R * _Ang2Bohr, color=color, alpha=0.5)
ax.set_xlim(0, atom_r * 3)
ax.legend()
nrows = len(l) // 2
ncols = len(l) // nrows
if nrows * ncols < len(l):
ncols += 1
fig, axs = plt.subplots(nrows, ncols, figsize=(11, 10))
def next_rc(ir, ic, nrows, ncols):
ic = ic + 1
if ic == ncols:
ic = 0
ir = ir + 1
return ir, ic
ir, ic = 0, 0
for this_plot in map(lambda x: x.lower(), plot):
if this_plot == "wavefunction":
plot_wavefunction(axs[ir][ic])
elif this_plot == "log":
plot_log(axs[ir][ic])
elif this_plot == "charge":
plot_charge(axs[ir][ic])
elif this_plot == "potential":
plot_potential(axs[ir][ic])
ir, ic = next_rc(ir, ic, nrows, ncols)
if show:
plt.show()
return fig, axs
def atom_plot_cli(subp=None):
"""Run plotting command for the output of atom"""
is_sub = not subp is None
title = "Plotting facility for atom output (run in the atom output directory)"
if is_sub:
global _script
_script = f"{_script} atom-plot"
p = subp.add_parser("atom-plot", description=title, help=title)
else:
import argparse
p = argparse.ArgumentParser(title)
p.add_argument(
"--plot",
"-P",
action="append",
type=str,
choices=("wavefunction", "charge", "log", "potential"),
help="""Determine what to plot""",
)
p.add_argument("-l", default="spdf", type=str, help="""Which l shells to plot""")
p.add_argument("--save", "-S", default=None, help="""Save output plots to file.""")
p.add_argument(
"--show",
default=False,
action="store_true",
help="""Force showing the plot (only if --save is specified)""",
)
p.add_argument(
"input", type=str, default="INP", help="""Input file name (default INP)"""
)
if is_sub:
p.set_defaults(runner=atom_plot)
else:
atom_plot(p.parse_args())
# Import object holding all the CLI
from sisl_toolbox.cli import register_toolbox_cli
register_toolbox_cli(atom_plot_cli)
@set_module("sisl_toolbox.siesta.atom")
def atom_plot(args):
import matplotlib.pyplot as plt
input = Path(args.input)
atom = AtomInput.from_input(input)
# If the specified input is a file, use the parent
# Otherwise use the input *as is*.
if input.is_file():
path = input.parent
else:
path = input
# if users have not specified what to plot, we plot everything
if args.plot is None:
args.plot = ("wavefunction", "charge", "log", "potential")
fig = atom.plot(path, plot=args.plot, l=args.l, show=False)[0]
if args.save is None:
plt.show()
else:
fig.savefig(args.save)
if args.show:
plt.show()