Basis optimization CLI

docs/toolbox/basis/basis.rst

@@ -0,0 +1,93 @@
+Basis optimization
+==================
+
+.. currentmodule:: sisl_toolbox.siesta.minimizer
+
+Optimizing a basis set for SIESTA is a cumbersome task.
+The goal of this toolbox is to allow users to **optimize a basis set with just one CLI command.**
+
+The commands and their options can be accessed like:
+
+.. code-block:: bash
+
+   stoolbox basis --help
+
+In summary, whenever one wants to optimize a basis set for a given system,
+the first step is to create a directory with the input files to run the
+calculation. This directory should contain, as usual:
+
+- The ``.fdf`` files with all the input parameters for the calculation.
+- The pseudopotentials (``.psf`` or ``.psml`` files).
+
+Then, one can directly run the optimization:
+
+.. code-block:: bash
+
+   stoolbox basis optim --geometry input.fdf
+
+with ``input.fdf`` being the input file for the calculation. This will use all the default
+values. Since there are many possible tweaks, we invite you to read carefully the help
+message from the CLI. Here we will just mention some important things that could go unnoticed.
+
+**Basis enthalpy:** The quantity that is minimized is the basis enthalpy. This is :math:`H = E + pV`
+with :math:`E` being the energy of the system, :math:`V` the volume of the basis and :math:`p` a "pressure" that
+is defined in the fdf file with the ```BasisPressure`` table. This "pressure" penalizes bigger
+basis, which result in more expensive calculations. It is the responsibility of the user to
+set this value. As a rule of thumb, we recommend to set it to ``0.02 GPa`` for the first two
+rows of the periodic table and ``0.2 GPa`` for the rest.
+
+**The SIESTA command:** There is a ``--siesta-cmd`` option to specify the way of executing SIESTA. By default, it
+is simply calling the ``siesta`` command, but you could set it for example to ``mpirun -n 4 siesta``
+so that SIESTA is ran in parallel.
+
+There is no problem with using this CLI in clusters inside a submitted job, for example.
+
+**A custom basis definition:** It may happen that the conventional optimizable parameters as well as their lower and
+upper bounds are not good for your case (e.g. you would like the upper bound for a cutoff
+radius to be higher). In that case, you can create a custom ``--basis-spec``. The best way
+to do it is by calling
+
+.. code-block:: bash
+
+   stoolbox basis build --geometry input.fdf
+
+which will generate a yaml file with a basis specification that you can tweak manually.
+Then, you can pass it directly to the optimization using the ``--config`` option:
+
+.. code-block:: bash
+
+   stoolbox basis optim --geometry input.fdf --config my_config.yaml
+
+**Installing the optimizers:** The default optimizer is BADS (https://github.com/acerbilab/bads)
+which is the one that we have found works best to optimize basis sets. The optimizer is however
+not installed by default. You can install it using pip:
+
+.. code-block:: bash
+
+   pip install pybads
+
+and the same would apply for other optimizers that you may want to use.
+
+**Output:** The output that appears on the terminal is left to the particular optimizer.
+However, sisl generates ``.dat`` files which contain information about each SIESTA execution.
+These files contain one column for each variable being optimized and one column for the
+metric to minimize.
+
+
+Python API
+----------
+
+The functions that do the work are also usable in python code by importing them:
+
+.. code-block:: python
+
+   from sisl_toolbox.siesta.minimizer import optimize_basis, write_basis_to_yaml
+
+Here is their documentation:
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+
+   optimize_basis
+   write_basis_to_yaml

docs/toolbox/index.rst

@@ -22,10 +22,12 @@ Toolboxes should be imported directly.
 
 
 The implemented toolboxes are listed here:
-   
+
 .. toctree::
    :maxdepth: 1
 
    transiesta/ts_fft
    siesta/atom_plot
    btd/btd
+   siesta/minimizer
+   basis/basis

docs/toolbox/siesta/minimizer.rst

@@ -0,0 +1,33 @@
+Minimizers
+=================
+
+.. currentmodule:: sisl_toolbox.siesta.minimizer
+
+In `sisl_toolbox.siesta.minimizer` there is a collection of minimizers
+that given some `variables`, a `runner` and a `metric` optimizes the
+variables to minimize the metric.
+
+These are the minimizer classes implemented:
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+
+   BaseMinimize
+   LocalMinimize
+   DualAnnealingMinimize
+   BADSMinimize
+   ParticleSwarmsMinimize
+
+For each of them, there is a subclass particularly tailored to optimize
+SIESTA runs:
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+
+   MinimizeSiesta
+   LocalMinimizeSiesta
+   DualAnnealingMinimizeSiesta
+   BADSMinimizeSiesta
+   ParticleSwarmsMinimizeSiesta

src/sisl/_lib/_argparse.py

@@ -1,11 +1,220 @@
 # 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/.
+
+import argparse
+import inspect
+import typing
+from typing import Any, Callable, Literal, Optional, Union
+
+from sisl._lib._docscrape import FunctionDoc
+
 try:
-    from rich_argparse import RichHelpFormatter
+    from rich_argparse import RawTextRichHelpFormatter
 
-    SislHelpFormatter = RichHelpFormatter
+    SislHelpFormatter = RawTextRichHelpFormatter
 except ImportError:
-    import argparse
-
     SislHelpFormatter = argparse.RawDescriptionHelpFormatter
+
+
+def is_optional(field):
+    """Check whether the annotation for a parameter is an Optional type."""
+    return typing.get_origin(field) is Union and type(None) in typing.get_args(field)
+
+
+def is_literal(field):
+    """Check whether the annotation for a parameter is a Literal type."""
+    return typing.get_origin(field) is Literal
+
+
+def get_optional_arg(field):
+    """Get the type of an optional argument from the typehint.
+
+    It only works if the annotation only has one type.
+
+    E.g.: Optional[int] -> int
+    E.g.: Optional[Union[int, str]] -> raises ValueError
+    """
+    if not is_optional(field):
+        return field
+
+    args = typing.get_args(field)
+    if len(args) > 2:
+        raise ValueError("Optional type must have at most 2 arguments")
+    for arg in args:
+        if arg is not type(None):
+            return arg
+
+    raise ValueError("No non-None type found in Union")
+
+
+def get_literal_args(field):
+    """Get the values of a literal.
+
+    E.g.: Literal[1, 2, 3] -> (1, 2, 3)
+    """
+    return typing.get_args(field)
+
+
+class NotPassedArg:
+    """Placeholder to use for arguments that have not been passed.
+
+    By setting this as the default value for an argument, we can
+    later check if the argument was passed through the CLI or not.
+    """
+
+    def __init__(self, val):
+        self.val = val
+
+    def __repr__(self):
+        return repr(self.val)
+
+    def __str__(self):
+        return str(self.val)
+
+    def __eq__(self, other):
+        return self.val == other
+
+    def __getattr__(self, name):
+        return getattr(self.val, name)
+
+
+def get_runner(func):
+    """Wraps a function to receive the args parsed from argparse"""
+
+    def _runner(args):
+
+        # Get the config argument. If present, load the arguments
+        # from the config (yaml) file.
+        config = getattr(args, "config", None)
+        config_args = {}
+        if config is not None:
+            import yaml
+
+            with open(args.config, "r") as f:
+                config_args = yaml.safe_load(f)
+
+        # Build the final arguments dictionary, using the arguments of the
+        # config file as defaults.
+        final_args = {}
+        for k, v in vars(args).items():
+            if k in ("runner", "config"):
+                continue
+            elif isinstance(v, NotPassedArg):
+                final_args[k] = config_args.get(k, v.val)
+            else:
+                final_args[k] = v
+
+        # And call the function
+        return func(**final_args)
+
+    return _runner
+
+
+def get_argparse_parser(
+    func: Callable,
+    name: Optional[str] = None,
+    subp=None,
+    parser_kwargs: dict[str, Any] = {},
+    arg_aliases: dict[str, str] = {},
+    defaults: dict[str, Any] = {},
+    add_config: bool = True,
+) -> argparse.ArgumentParser:
+    """Creates an argument parser from a function's signature and docstring.
+
+    The created argument parser just mimics the function. It is a CLI version
+    of the function.
+
+    Parameters
+    ----------
+    func :
+        The function to create the parser for.
+    name :
+        The name of the parser. If None, the function's name is used.
+    subp :
+        The subparser to add the parser to. If None, a new isolated
+        parser is created.
+    parser_kwargs :
+        Additional arguments to pass to the parser.
+    arg_aliases :
+        Dictionary holding aliases (shortcuts) for the arguments. The keys
+        of this dictionary are the argument names, and the values are the
+        aliases. For example, if the function has an argument called
+        `--size`, and you want to add a shortcut `-s`, you can pass
+        `arg_aliases={"size": "s"}`.
+    defaults :
+        Dictionary holding default values for the arguments. The keys
+        of this dictionary are the argument names, and the values are
+        the default values. The defaults are taken from the function's
+        signature if not specified.
+    add_config :
+        If True, adds a `--config` argument to the parser. This
+        argument accepts a path to a YAML file containing the
+        arguments for the function.
+    """
+
+    # Check if the function needs to be added as a subparser
+    is_sub = not subp is None
+
+    # Get the function's information
+    fdoc = FunctionDoc(func)
+    signature = inspect.signature(func)
+
+    # Initialize parser
+    title = "".join(fdoc["Summary"])
+    parser_help = "\n".join(fdoc["Extended Summary"])
+    if is_sub:
+        p = subp.add_parser(
+            name or func.__name__.replace("_", "-"),
+            description=parser_help,
+            help=title,
+            **parser_kwargs,
+        )
+    else:
+        p = argparse.ArgumentParser(title, **parser_kwargs)
+
+    # Add the config argument to load the arguments from a YAML file
+    if add_config:
+        p.add_argument(
+            "--config",
+            "-c",
+            type=str,
+            default=None,
+            help="Path to a YAML file containing the arguments for the command",
+        )
+
+    group = p.add_argument_group("Function arguments")
+
+    # Add all the function's parameters to the parser
+    parameters_help = {p.name: p.desc for p in fdoc["Parameters"]}
+    for param in signature.parameters.values():
+
+        arg_names = [f"--{param.name.replace('_', '-')}"]
+        if param.name in arg_aliases:
+            arg_names.append(f"-{arg_aliases[param.name]}")
+
+        annotation = param.annotation
+        if is_optional(annotation):
+            annotation = get_optional_arg(annotation)
+
+        choices = None
+        if is_literal(annotation):
+            choices = get_literal_args(annotation)
+            annotation = type(choices[0])
+
+        group.add_argument(
+            *arg_names,
+            type=annotation,
+            default=NotPassedArg(param.default),
+            choices=choices,
+            action=argparse.BooleanOptionalAction if annotation is bool else None,
+            required=param.default is inspect._empty,
+            help="\n".join(parameters_help.get(param.name, [])),
+        )
+
+    if is_sub:
+        defaults = {"runner": get_runner(func), **defaults}
+
+    p.set_defaults(**defaults)
+
+    return p

src/sisl_toolbox/cli/_cli.py

@@ -44,8 +44,7 @@ def register(self, setup):
         """
         self._cmds.append(setup)
 
-    def __call__(self, argv=None):
-
+    def get_parser(self):
         # Create command-line
         cmd = Path(sys.argv[0])
         p = argparse.ArgumentParser(
@@ -70,6 +69,12 @@ def __call__(self, argv=None):
         for cmd in self._cmds:
             cmd(subp, parser_kwargs=dict(formatter_class=p.formatter_class))
 
+        return p
+
+    def __call__(self, argv=None):
+
+        p = self.get_parser()
+
         args = p.parse_args(argv)
         args.runner(args)
 

src/sisl_toolbox/cli/_cli_imports.py

@@ -8,5 +8,6 @@
 __all__ = []
 
 
-import sisl_toolbox.siesta.atom  # noqa: F401
+import sisl_toolbox.siesta.atom._register_cli  # noqa: F401
+import sisl_toolbox.siesta.minimizer._register_cli  # noqa: F401
 import sisl_toolbox.transiesta.poisson  # noqa: F401