Add the PseudoDojoFamily

aiida_pseudo/cli/install.py

@@ -15,6 +15,7 @@
 from .root import cmd_root
 
 URL_SSSP_BASE = 'https://legacy-archive.materialscloud.org/file/2018.0001/v4/'
+URL_PSEUDODOJO_BASE = 'http://www.pseudo-dojo.org/pseudos/'
 
 
 @cmd_root.group('install')
@@ -145,3 +146,112 @@ def cmd_install_sssp(version, functional, protocol, traceback):
         family.set_cutoffs({'normal': cutoffs})
 
         echo.echo_success(f'installed `{label}` containing {family.count()} pseudo potentials')
+
+
+@cmd_install.command('pseudo-dojo')
+@options.VERSION(type=click.Choice(['0.3', '0.4', '1.0']), default='0.4')
+@options.FUNCTIONAL(type=click.Choice(['pbe', 'pbesol', 'pw']), default='pbe')
+@options.RELATIVISTIC(type=click.Choice(['sr', 'sr3plus', 'fr']), default='sr')
+@options.PROTOCOL(type=click.Choice(['standard', 'stringent']), default='standard')
+@options.PSEUDO_FORMAT(type=click.Choice(['psp8', 'upf', 'psml', 'jthxml']), default='psp8')
+@options.DEFAULT_STRINGENCY(type=click.Choice(['low', 'normal', 'high']), default='normal')
+@options.TRACEBACK()
+@decorators.with_dbenv()
+def cmd_install_pseudo_dojo(version, functional, relativistic, protocol, pseudo_format, default_stringency, traceback):
+    """Install a PseudoDojo configuration.
+
+    The PseudoDojo configuration will be automatically downloaded from pseudo-dojo.org to create a new
+    `PseudoDojoFamily` subclass instance based on the specified pseudopotential format.
+    """
+    # pylint: disable=too-many-locals,too-many-arguments,too-many-statements
+    import requests
+
+    from aiida.common.files import md5_file
+    from aiida.orm import Group, QueryBuilder
+    from aiida_pseudo import __version__
+    from aiida_pseudo.data.pseudo import JthXmlData, Psp8Data, PsmlData, UpfData
+    from aiida_pseudo.groups.family import PseudoDojoConfiguration, PseudoDojoFamily
+
+    from .utils import attempt, create_family_from_archive
+
+    pseudo_type_mapping = {
+        'jthxml': JthXmlData,
+        'psp8': Psp8Data,
+        'psml': PsmlData,
+        'upf': UpfData,
+    }
+
+    try:
+        pseudo_type = pseudo_type_mapping[pseudo_format]
+    except KeyError:
+        echo.echo_critical(f'{pseudo_format} is not a valid PseudoDojo pseudopotential format')
+
+    configuration = PseudoDojoConfiguration(version, functional, relativistic, protocol)
+    label = PseudoDojoFamily.format_configuration_label(configuration)
+    description = f'PseudoDojo v{version} {functional} {relativistic} {protocol} {pseudo_format} ' + \
+        f'installed with aiida-pseudo v{__version__}'
+
+    if configuration not in PseudoDojoFamily.valid_configurations:
+        echo.echo_critical(
+            f'{version} {functional} {relativistic} {protocol} {pseudo_format} is not a valid PseudoDojo '
+            f'configuration for {PseudoDojoFamily}'
+        )
+
+    if QueryBuilder().append(PseudoDojoFamily, filters={'label': label}).first():
+        echo.echo_critical(f'{PseudoDojoFamily.__name__}<{label}> is already installed')
+
+    url_versions = {'0.4': '04', '0.3': None, '1.0': None}
+
+    with tempfile.TemporaryDirectory() as dirpath:
+        if relativistic == 'sr3plus':
+            url_archive = f'{URL_PSEUDODOJO_BASE}/nc-sr-{url_versions[version]}-3plus_{functional}_{protocol}_' + \
+                f'{pseudo_format}.tgz'
+            url_metadata = f'{URL_PSEUDODOJO_BASE}/nc-sr-{url_versions[version]}-3plus_{functional}_{protocol}' + \
+                '_djrepo.tgz'
+        elif pseudo_format == 'xml':
+            url_archive = f'{URL_PSEUDODOJO_BASE}/paw_{functional}_{protocol}_{pseudo_format}.tgz'
+            url_metadata = f'{URL_PSEUDODOJO_BASE}/paw_{functional}_{protocol}_djrepo.tgz'
+        elif version == '03':
+            url_archive = f'{URL_PSEUDODOJO_BASE}/nc-{relativistic}_{functional}_{protocol}_{pseudo_format}.tgz'
+            url_metadata = f'{URL_PSEUDODOJO_BASE}/nc-{relativistic}_{functional}_{protocol}_djrepo.tgz'
+        else:
+            url_archive = f'{URL_PSEUDODOJO_BASE}/nc-{relativistic}-{url_versions[version]}_{functional}_' + \
+                f'{protocol}_{pseudo_format}.tgz'
+            url_metadata = f'{URL_PSEUDODOJO_BASE}/nc-{relativistic}-{url_versions[version]}_{functional}_' + \
+                f'{protocol}_djrepo.tgz'
+
+        filepath_archive = os.path.join(dirpath, 'archive.tgz')
+        filepath_metadata = os.path.join(dirpath, 'metadata.tgz')
+
+        with attempt('downloading selected pseudo potentials archive... ', include_traceback=traceback):
+            response = requests.get(url_archive)
+            response.raise_for_status()
+            with open(filepath_archive, 'wb') as handle:
+                handle.write(response.content)
+                handle.flush()
+                description += f'\nArchive pseudos md5: {md5_file(filepath_archive)}'
+
+        with attempt('unpacking archive and parsing pseudos... ', include_traceback=traceback):
+            family = create_family_from_archive(PseudoDojoFamily, label, filepath_archive, pseudo_type=pseudo_type)
+
+        with attempt('downloading selected pseudo potentials metadata archive... ', include_traceback=traceback):
+            response = requests.get(url_metadata)
+            response.raise_for_status()
+            with open(filepath_metadata, 'wb') as handle:
+                handle.write(response.content)
+                handle.flush()
+                description += f'\nPseudo metadata archive md5: {md5_file(filepath_metadata)}'
+
+        with attempt('unpacking metadata archive and parsing metadata...', include_traceback=traceback):
+            md5s, cutoffs = PseudoDojoFamily.parse_djrepos_from_archive(filepath_metadata, pseudo_type=pseudo_type)
+
+        for element, md5 in md5s.items():
+            if family.get_pseudo(element).md5 != md5:
+                Group.objects.delete(family.pk)
+                msg = f'md5 of pseudo for element {element} does not match that of the metadata {md5}'
+                echo.echo_critical(msg)
+
+        family.description = description
+        family.set_cutoffs(cutoffs, default_stringency=default_stringency)
+
+        echo.echo_success(f'installed `{label}` containing {family.count()} pseudo potentials`')

aiida_pseudo/cli/params/options.py

@@ -7,18 +7,50 @@
 from aiida.cmdline.params.options import OverridableOption
 from .types import PseudoPotentialFamilyTypeParam
 
-__all__ = ('VERSION', 'FUNCTIONAL', 'PROTOCOL', 'TRACEBACK', 'FAMILY_TYPE', 'ARCHIVE_FORMAT')
+__all__ = (
+    'VERSION', 'FUNCTIONAL', 'RELATIVISTIC', 'PROTOCOL', 'PSEUDO_FORMAT', 'DEFAULT_STRINGENCY', 'TRACEBACK',
+    'FAMILY_TYPE', 'ARCHIVE_FORMAT'
+)
 
 VERSION = OverridableOption(
-    '-v', '--version', type=click.STRING, required=False, help='Select the version of the SSSP configuration.'
+    '-v', '--version', type=click.STRING, required=False, help='Select the version of the installed configuration.'
 )
 
 FUNCTIONAL = OverridableOption(
-    '-f', '--functional', type=click.STRING, required=False, help='Select the functional of the SSSP configuration.'
+    '-x',
+    '--functional',
+    type=click.STRING,
+    required=False,
+    help='Select the functional of the installed configuration.'
+)
+
+RELATIVISTIC = OverridableOption(
+    '-r',
+    '--relativistic',
+    type=click.STRING,
+    required=False,
+    help='Select the type of relativistic effects included in the installed configuration.'
 )
 
 PROTOCOL = OverridableOption(
-    '-p', '--protocol', type=click.STRING, required=False, help='Select the protocol of the SSSP configuration.'
+    '-p', '--protocol', type=click.STRING, required=False, help='Select the protocol of the installed configuration.'
+)
+
+PSEUDO_FORMAT = OverridableOption(
+    '-f',
+    '--pseudo-format',
+    type=click.STRING,
+    required=True,
+    help='Select the pseudopotential file format of the installed configuration.'
+)
+
+DEFAULT_STRINGENCY = OverridableOption(
+    '-s',
+    '--default-stringency',
+    type=click.STRING,
+    required=False,
+    help='Select the default stringency level for the installed configuration. See the documentation for valid '
+    'options.'
 )
 
 TRACEBACK = OverridableOption(

aiida_pseudo/cli/utils.py

@@ -31,16 +31,18 @@ def attempt(message, exception_types=Exception, include_traceback=False):
         echo.echo_highlight(' [OK]', color='success', bold=True)
 
 
-def create_family_from_archive(cls, label, filepath_archive, fmt=None):
+def create_family_from_archive(cls, label, filepath_archive, fmt=None, pseudo_type=None):
     """Construct a new pseudo family instance from a tar.gz archive.
 
     .. warning:: the archive should not contain any subdirectories, but just the pseudo potential files.
 
-    :param cls: the class to use, e.g. ``SsspFamily``
+    :param cls: the pseudopotential family class to use, e.g. ``SsspFamily``
     :param label: the label for the new family
-    :param filepath: absolute filepath to the .tar.gz archive containing the pseudo potentials.
-    :param filepath: optional absolute filepath to the .json file containing the pseudo potentials metadata.
+    :param filepath_archive: absolute filepath to the .tar.gz archive containing the pseudo potentials
     :param fmt: the format of the archive, if not specified will attempt to guess based on extension of ``filepath``
+    :param pseudo_type: subclass of ``PseudoPotentialData`` to be used for the parsed pseudos. If not specified and
+        the family only defines a single supported pseudo type in ``_pseudo_types`` then that will be used otherwise
+        a ``ValueError`` is raised.
     :return: newly created family
     :raises OSError: if the archive could not be unpacked or pseudos in it could not be parsed into a family
     """
@@ -55,7 +57,7 @@ def create_family_from_archive(cls, label, filepath_archive, fmt=None):
             raise OSError(f'failed to unpack the archive `{filepath_archive}`: {exception}') from exception
 
         try:
-            family = cls.create_from_folder(dirpath, label)
+            family = cls.create_from_folder(dirpath, label, pseudo_type=pseudo_type)
         except ValueError as exception:
             raise OSError(f'failed to parse pseudos from `{dirpath}`: {exception}') from exception
 

aiida_pseudo/common/units.py

@@ -2,3 +2,4 @@
 """Module with constants for unit conversions."""
 
 RY_TO_EV = 13.6056917253  # Taken from `qe_tools.constants` v2.0
+HA_TO_EV = RY_TO_EV / 2

aiida_pseudo/data/pseudo/__init__.py

@@ -6,5 +6,6 @@
 from .psml import *
 from .psp8 import *
 from .upf import *
+from .jthxml import *
 
-__all__ = (pseudo.__all__ + psf.__all__ + psml.__all__ + psp8.__all__ + upf.__all__)
+__all__ = (pseudo.__all__ + psf.__all__ + psml.__all__ + psp8.__all__ + upf.__all__ + jthxml.__all__)

aiida_pseudo/data/pseudo/jthxml.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+"""Module for data plugin to represent a pseudo potential in JTH XML format."""
+import re
+from typing import BinaryIO
+
+from .pseudo import PseudoPotentialData
+
+__all__ = ('JthXmlData',)
+
+REGEX_ELEMENT = re.compile(r"""\s*symbol\s*=\s*['"]\s*(?P<element>[a-zA-Z]{1,2})\s*['"].*""")
+
+
+def parse_element(stream: BinaryIO):
+    """Parse the content of the JTH XML file to determine the element.
+
+    :param stream: a filelike object with the binary content of the file.
+    :return: the symbol of the element following the IUPAC naming standard.
+    """
+    from xml.dom.minidom import parse
+
+    try:
+        xml = parse(stream)
+        element = xml.getElementsByTagName('atom')[0].attributes['symbol'].value
+    except (AttributeError, IndexError, KeyError) as exception:
+        raise ValueError(f'could not parse the element from the XML content: {exception}') from exception
+
+    return element.capitalize()
+
+
+class JthXmlData(PseudoPotentialData):
+    """Data plugin to represent a pseudo potential in JTH XML format."""
+
+    def set_file(self, stream: BinaryIO, filename: str = None, **kwargs):  # pylint: disable=arguments-differ
+        """Set the file content.
+
+        :param stream: a filelike object with the binary content of the file.
+        :param filename: optional explicit filename to give to the file stored in the repository.
+        :raises ValueError: if the element symbol is invalid.
+        """
+        stream.seek(0)
+        self.element = parse_element(stream)
+        stream.seek(0)
+        super().set_file(stream, filename, **kwargs)

aiida_pseudo/groups/family/__init__.py

@@ -2,6 +2,7 @@
 # pylint: disable=undefined-variable
 """Module with group plugins to represent pseudo potential families."""
 from .pseudo import *
+from .pseudo_dojo import *
 from .sssp import *
 
-__all__ = (pseudo.__all__ + sssp.__all__)
+__all__ = (pseudo.__all__ + pseudo_dojo.__all__ + sssp.__all__)