Install pseudos from downloaded files (aiidalab#524)

fixes aiidalab#490

In this commit, we pinging the aiida-pseudo version to use the
--download-only option to download the pseudo archive to the container
for quick installation when the container first time start.
Note that from the test all 8 pseudos seems to only save ~10s but in the place where the network condition is bad, this is very helpful.

Further more Refactoring the methods in setup_pseudos.py to support flexible parameter passing.

docker/Dockerfile

@@ -11,8 +11,14 @@ USER ${NB_USER}
 RUN cd ${PREINSTALL_APP_FOLDER} && \
      # Remove all untracked files and directories. For example the setup lock flag file.
      git clean -fx && \
-     # It is important to install from `aiidalab install` to mimic the same installation operation from the app store.
-     aiidalab install --yes --python "/opt/conda/bin/python" "quantum-espresso@file://${PREINSTALL_APP_FOLDER}" && \
+     # It is important to install from `aiidalab install` to mimic the exact installation operation as
+     # from the app store.
+     # The command wil first install the dependencies from list by parsing setup config files,
+     # (for `aiidalab/aiidalab<23.03.2` the `setup.py` should be in the root folder of the app https://github.com/aiidalab/aiidalab/pull/382).
+     # and then the app and restart the daemon in the end.
+     # But since the aiida profile not yet exists, the daemon restart will fail but it is not a problem.
+     # Because we only need the dependencies to be installed.
+     aiidalab install --yes --python ${CONDA_DIR}/bin/python  "quantum-espresso@file://${PREINSTALL_APP_FOLDER}" && \
      fix-permissions "${CONDA_DIR}" && \
      fix-permissions "/home/${NB_USER}"
 
@@ -28,6 +34,11 @@ RUN mamba create -p /opt/conda/envs/quantum-espresso --yes \
      fix-permissions "${CONDA_DIR}" && \
      fix-permissions "/home/${NB_USER}"
 
+# Download the QE pseudopotentials to the folder for afterware installation.
+ENV PSEUDO_FOLDER ${CONDA_DIR}/pseudo
+RUN mkdir -p ${PSEUDO_FOLDER} && \
+    python -m aiidalab_qe download-pseudos --dest ${PSEUDO_FOLDER}
+
 COPY before-notebook.d/* /usr/local/bin/before-notebook.d/
 
 WORKDIR "/home/${NB_USER}"

docker/before-notebook.d/71_install-qeapp.sh

@@ -13,11 +13,9 @@ else
 fi
 
 # Install the pseudo libraries if not already installed.
-# This can be simplified and accelerated once the following PR is merged:
-# https://github.com/aiidateam/aiida-pseudo/pull/135
 if aiida-pseudo list | grep -q "no pseudo potential families"; then
     echo "Installing pseudo potential families."
-    python -m aiidalab_qe install-pseudos
+    python -m aiidalab_qe install-pseudos --source ${PSEUDO_FOLDER}
 else
     echo "Pseudo potential families are already installed."
 fi

setup.cfg

@@ -28,6 +28,7 @@ install_requires =
     Jinja2~=3.0
     aiida-quantumespresso~=4.3.0
     aiidalab-widgets-base==2.1.0a0
+    aiida-pseudo~=1.4
     filelock~=3.8
     importlib-resources~=5.2
     widget-bandsplot~=0.5.1

src/aiidalab_qe/__main__.py

@@ -1,12 +1,16 @@
 """For running the app from the command line used for post_install script.
 """
 
+from pathlib import Path
+
 import click
 from aiida import load_profile
 
 from aiidalab_qe.common.setup_codes import codes_are_setup
 from aiidalab_qe.common.setup_codes import install as install_qe_codes
-from aiidalab_qe.common.setup_pseudos import install as setup_pseudos
+
+# The default profile name of AiiDAlab container.
+_DEFAULT_PROFILE = "default"
 
 
 @click.group()
@@ -16,7 +20,7 @@ def cli():
 
 @cli.command()
 @click.option("-f", "--force", is_flag=True)
-@click.option("-p", "--profile", default="default")
+@click.option("-p", "--profile", default=_DEFAULT_PROFILE)
 def install_qe(force, profile):
     load_profile(profile)
     try:
@@ -29,16 +33,51 @@ def install_qe(force, profile):
 
 
 @cli.command()
-@click.option("-p", "--profile", default="default")
-def install_pseudos(profile):
+@click.option("-p", "--profile", default=_DEFAULT_PROFILE, help="AiiDA profile name.")
+@click.option(
+    "source",
+    "--source",
+    default=None,
+    help="The source folder to install from local.",
+    type=click.Path(exists=True, path_type=Path, resolve_path=True),
+)
+def install_pseudos(profile, source):
+    """Install pseudopotentials from a local folder if source is specified,
+    otherwise download from remote repositories.
+    """
+    from aiidalab_qe.common.setup_pseudos import install
+
     load_profile(profile)
+
     try:
-        for msg, _ in setup_pseudos():
+        for msg, _ in install(cwd=source, download_only=False):
             click.echo(msg)
         click.secho("Pseudopotentials are installed!", fg="green")
     except Exception as error:
         raise click.ClickException(f"Failed to set up pseudo potentials: {error}")
 
 
+@cli.command()
+@click.option(
+    "dest",
+    "--dest",
+    default=None,
+    help="The dest folder where to download the pseudos.",
+    type=click.Path(exists=True, path_type=Path, resolve_path=True),
+)
+def download_pseudos(dest):
+    from aiidalab_qe.common.setup_pseudos import EXPECTED_PSEUDOS, _install_pseudos
+
+    try:
+        for progress in _install_pseudos(
+            EXPECTED_PSEUDOS, download_only=True, cwd=dest
+        ):
+            click.echo(progress)
+        click.secho("Pseudopotentials are downloaded!", fg="green")
+
+    except Exception as error:
+        raise click.ClickException(f"Failed to download pseudo potentials: {error}")
+
+
 if __name__ == "__main__":
     cli()

src/aiidalab_qe/app/configuration/advanced.py

@@ -16,6 +16,7 @@
 
 from aiidalab_qe.app.parameters import DEFAULT_PARAMETERS
 from aiidalab_qe.common.panel import Panel
+from aiidalab_qe.common.setup_pseudos import PseudoFamily
 
 from .pseudos import PseudoFamilySelector, PseudoSetter
 
@@ -262,8 +263,9 @@ def set_panel_value(self, parameters):
         """Set the panel value from the given parameters."""
 
         if "pseudo_family" in parameters:
+            pseudo_family_string = parameters["pseudo_family"]
             self.pseudo_family_selector.load_from_pseudo_family(
-                parameters.get("pseudo_family")
+                PseudoFamily.from_string(pseudo_family_string)
             )
         if "pseudos" in parameters["pw"]:
             self.pseudo_setter.set_pseudos(parameters["pw"]["pseudos"], {})
@@ -295,11 +297,16 @@ def reset(self):
         with self.hold_trait_notifications():
             # Reset protocol dependent settings
             self._update_settings_from_protocol(self.protocol)
-            self.pseudo_family_selector.load_from_pseudo_family(
-                DEFAULT_PARAMETERS["advanced"]["pseudo_family"]
-            )
+
+            # reset the pseudo family
+            pseudo_family_dict = DEFAULT_PARAMETERS["advanced"]["pseudo_family"]
+            pseudo_family = PseudoFamily(**pseudo_family_dict)
+
+            self.pseudo_family_selector.load_from_pseudo_family(pseudo_family)
+
             # reset total charge
             self.total_charge.value = DEFAULT_PARAMETERS["advanced"]["tot_charge"]
+
             # reset the override checkbox
             self.override.value = False
             self.smearing.reset()

src/aiidalab_qe/app/configuration/pseudos.py

@@ -3,7 +3,6 @@
 
 import io
 import re
-from dataclasses import dataclass
 
 import ipywidgets as ipw
 import traitlets as tl
@@ -14,58 +13,18 @@
 from aiidalab_widgets_base.utils import StatusHTML
 
 from aiidalab_qe.app.parameters import DEFAULT_PARAMETERS
+from aiidalab_qe.common.setup_pseudos import (
+    PSEUDODOJO_VERSION,
+    SSSP_VERSION,
+    PseudoFamily,
+)
 
 UpfData = DataFactory("pseudo.upf")
 SsspFamily = GroupFactory("pseudo.family.sssp")
 PseudoDojoFamily = GroupFactory("pseudo.family.pseudo_dojo")
 CutoffsPseudoPotentialFamily = GroupFactory("pseudo.family.cutoffs")
 
 
-@dataclass(frozen=True)
-class PseudoFamily:
-    library: str
-    version: str
-    functional: str
-    accuracy: str
-    relativistic: str | None = None
-    file_type: str | None = None
-
-    @classmethod
-    def from_string(cls, pseudo_family_string: str) -> PseudoFamily:
-        """Initialize from a pseudo family string."""
-        # We support two pseudo families: SSSP and PseudoDojo
-        # They are formatted as follows:
-        # SSSP: SSSP/<version>/<functional>/<accuracy>
-        # PseudoDojo: PseudoDojo/<version>/<functional>/<relativistic>/<accuracy>/<file_type>
-        # where <relativistic> is either 'SR' or 'FR' and <file_type> is either 'upf' or 'psml'
-        # Before we unify the format of family strings, the conditions below are necessary
-        # to distinguish between the two families
-        library = pseudo_family_string.split("/")[0]
-        if library == "SSSP":
-            version, functional, accuracy = pseudo_family_string.split("/")[1:]
-            relativistic = None
-            file_type = None
-        elif library == "PseudoDojo":
-            (
-                version,
-                functional,
-                relativistic,
-                accuracy,
-                file_type,
-            ) = pseudo_family_string.split("/")[1:]
-        else:
-            raise ValueError(f"Unknown pseudo family {pseudo_family_string}")
-
-        return cls(
-            library=library,
-            version=version,
-            functional=functional,
-            accuracy=accuracy,
-            relativistic=relativistic,
-            file_type=file_type,
-        )
-
-
 class PseudoFamilySelector(ipw.VBox):
     title = ipw.HTML(
         """<div style="padding-top: 0px; padding-bottom: 10px">
@@ -189,12 +148,13 @@ def set_value(self, _=None):
         """
         library, accuracy = self.library_selection.value.split()
         functional = self.dft_functional.value
+        # XXX (jusong.yu): a validator is needed to check the family string is consistent with the list of pseudo families defined in the setup_pseudos.py
         if library == "PseudoDojo":
-            pseudo_family_string = f"PseudoDojo/0.4/{functional}/SR/{accuracy}/upf"
+            pseudo_family_string = (
+                f"PseudoDojo/{PSEUDODOJO_VERSION}/{functional}/SR/{accuracy}/upf"
+            )
         elif library == "SSSP":
-            # XXX (unkcpz): the version is hard coded here which bring the
-            # the risk of inconsistency when the version is changed in the future, we want to have a centralized place to store the information.
-            pseudo_family_string = f"SSSP/1.2/{functional}/{accuracy}"
+            pseudo_family_string = f"SSSP/{SSSP_VERSION}/{functional}/{accuracy}"
         else:
             raise ValueError(
                 f"Unknown pseudo family {self.override_protocol_pseudo_family.value}"
@@ -240,13 +200,12 @@ def _update_settings_from_protocol(self, protocol):
         pseudo_family_string = PwBaseWorkChain.get_protocol_inputs(protocol)[
             "pseudo_family"
         ]
+        pseudo_family = PseudoFamily.from_string(pseudo_family_string)
 
-        self.load_from_pseudo_family(pseudo_family_string)
+        self.load_from_pseudo_family(pseudo_family)
 
-    def load_from_pseudo_family(self, pseudo_family_string: str):
+    def load_from_pseudo_family(self, pseudo_family: PseudoFamily):
         """Reload the widget from the given pseudo family string."""
-        pseudo_family = PseudoFamily.from_string(pseudo_family_string)
-
         with self.hold_trait_notifications():
             # will trigger the callback to set the value of widgets
             self.library_selection.value = (

src/aiidalab_qe/app/parameters/qeapp.yaml

@@ -15,7 +15,11 @@ workchain:
 
 ## Advanced pw settings
 advanced:
-    pseudo_family: SSSP/1.2/PBEsol/efficiency
+    pseudo_family:
+        library: SSSP
+        version: 1.2
+        functional: PBEsol
+        accuracy: efficiency
     tot_charge: 0
 
 ## Codes

src/aiidalab_qe/app/submission/__init__.py

@@ -23,12 +23,6 @@
 
 from .resource import ParallelizationSettings, ResourceSelectionWidget
 
-PROTOCOL_PSEUDO_MAP = {
-    "fast": "SSSP/1.2/PBE/efficiency",
-    "moderate": "SSSP/1.2/PBE/efficiency",
-    "precise": "SSSP/1.2/PBE/precision",
-}
-
 
 class SubmitQeAppWorkChainStep(ipw.VBox, WizardAppWidgetStep):
     """Step for submission of a bands workchain."""