_atom.py

# 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/.
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`.
"""
from typing import Optional, List
from typing_extensions import Annotated

import sys
from collections.abc import Iterable
from enum import Enum
from functools import reduce
from pathlib import Path

import numpy as np
from scipy.interpolate import interp1d

import sisl as si
from sisl.utils import NotNonePropertyDict, PropertyDict

from sisl_toolbox.cli._cli_arguments import CLIArgument, CLIOption

__all__ = ["AtomInput", "atom_plot"]


_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'


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

        # 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.)
        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

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

    @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-radii"), 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.).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)

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

    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.:
            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.] * 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")

    def write_all_electron(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"   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")

    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)

    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

    def ae(self, filename="INP", path=None):
        with self(filename, path) as f:
            self.write_header(f)
            self.write_all_electron(f)

    def pg(self, filename="INP", path=None):
        with self(filename, path) as f:
            self.write_header(f)
            self.write_generation(f)

    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.

        # 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.
                orb._q0 -= dq
                charge -= dq
        return self.__class__(atom, self.define, **self.opts)

    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.)
            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 = np.trapz(ae_cc, ae_r)
                valence_c = np.trapz(ps_vc, ps_r)
                print(f"Total charge in atom: {core_c + valence_c:.5f}")
                overlap_c = np.trapz(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., 1.5])
                min_x = (ae_r > 0.2).nonzero()[0].min()
                max_x = (fraction[min_x:] > 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


class AtomPlotOption(Enum):
    """Plotting options for atom"""
    wavefunction = 'wavefunction'
    charge = 'charge'
    log = 'log'
    potential = 'potential'

def atom_plot(
    plot: Annotated[Optional[List[AtomPlotOption]], CLIArgument()] = None, 
    input: Path = Path("INP"), 
    l: str = 'spdf', 
    save: Annotated[Optional[str], CLIOption("-S", "--save")] = None, 
    show: bool = False
):
    """Plotting facility for atom output (run in the atom output directory)

    Parameters
    ----------
    plot:
        Determine what to plot. If None is given, it plots everything.
    input:
        Input file name.
    l: 
        Which l shells to plot.
    save: 
        Save output plots to file.
    show: 
        Force showing the plot.
    """
    import matplotlib.pyplot as plt

    input_path = Path(input)
    atom = AtomInput.from_input(input_path)

    # If the specified input is a file, use the parent
    # Otherwise use the input *as is*.
    if input_path.is_file():
        path = input_path.parent
    else:
        path = input_path

    # if users have not specified what to plot, we plot everything
    if plot is None:
        plots = [p.value for p in AtomPlotOption]
    else:
        plots = [p.value if isinstance(p, AtomPlotOption) else p for p in plot ]

    fig = atom.plot(path, plot=plots, l=l, show=False)[0]

    if save is None:
        plt.show()
    else:
        fig.savefig(save)
        if show:
            plt.show()