diff --git a/aiida_pseudo/cli/__init__.py b/aiida_pseudo/cli/__init__.py index a6e7718..bdeffa2 100644 --- a/aiida_pseudo/cli/__init__.py +++ b/aiida_pseudo/cli/__init__.py @@ -7,6 +7,7 @@ click_completion.init() from .root import cmd_root +from .family import cmd_family from .install import cmd_install, cmd_install_family, cmd_install_sssp, cmd_install_pseudo_dojo from .list import cmd_list from .show import cmd_show diff --git a/aiida_pseudo/cli/family.py b/aiida_pseudo/cli/family.py new file mode 100644 index 0000000..0b0d7e1 --- /dev/null +++ b/aiida_pseudo/cli/family.py @@ -0,0 +1,58 @@ +# -*- coding: utf-8 -*- +"""Commands to inspect or modify the contents of pseudo potential families.""" +import json + +import click + +from aiida.cmdline.utils import decorators, echo + +from ..groups.mixins import RecommendedCutoffMixin +from .params import arguments, options +from .root import cmd_root + + +@cmd_root.group('family') +def cmd_family(): + """Command group to inspect or modify the contents of pseudo potential families.""" + + +@cmd_family.group('cutoffs') +def cmd_family_cutoffs(): + """Command group to inspect or modify the recommended cutoffs of pseudo potential families.""" + + +@cmd_family_cutoffs.command('set') +@arguments.PSEUDO_POTENTIAL_FAMILY() +@click.argument('cutoffs', type=click.File(mode='rb')) +@options.STRINGENCY(required=True) +@decorators.with_dbenv() +def cmd_family_cutoffs_set(family, cutoffs, stringency): # noqa: D301 + """Set the recommended cutoffs for a pseudo potential family. + + The cutoffs should be provided as a JSON file through the argument `CUTOFFS` which should have the structure: + + \b + { + "Ag": { + "cutoff_wfc": 50.0, + "cutoff_rho": 200.0 + }, + ... + } + + where the cutoffs are expected to be in electronvolt. + """ + if not isinstance(family, RecommendedCutoffMixin): + raise click.BadParameter(f'family `{family}` does not support recommended cutoffs to be set.') + + try: + data = json.load(cutoffs) + except ValueError as exception: + raise click.BadParameter(f'`{cutoffs.name}` contains invalid JSON: {exception}', param_hint='CUTOFFS') + + try: + family.set_cutoffs({stringency: data}) + except ValueError as exception: + raise click.BadParameter(f'`{cutoffs.name}` contains invalid cutoffs: {exception}', param_hint='CUTOFFS') + + echo.echo_success(f'set cutoffs for `{family}` with the stringency `{stringency}`.') diff --git a/aiida_pseudo/cli/params/arguments.py b/aiida_pseudo/cli/params/arguments.py new file mode 100644 index 0000000..8db15a6 --- /dev/null +++ b/aiida_pseudo/cli/params/arguments.py @@ -0,0 +1,10 @@ +# -*- coding: utf-8 -*- +"""Reusable arguments for CLI commands.""" +from aiida.cmdline.params.arguments import OverridableArgument +from .types import PseudoPotentialFamilyParam + +__all__ = ('PSEUDO_POTENTIAL_FAMILY') + +PSEUDO_POTENTIAL_FAMILY = OverridableArgument( + 'family', type=PseudoPotentialFamilyParam(sub_classes=('aiida.groups:pseudo.family',)) +) diff --git a/aiida_pseudo/cli/params/options.py b/aiida_pseudo/cli/params/options.py index 62c14ff..c5824b3 100644 --- a/aiida_pseudo/cli/params/options.py +++ b/aiida_pseudo/cli/params/options.py @@ -8,8 +8,8 @@ from .types import PseudoPotentialFamilyTypeParam __all__ = ( - 'VERSION', 'FUNCTIONAL', 'RELATIVISTIC', 'PROTOCOL', 'PSEUDO_FORMAT', 'DEFAULT_STRINGENCY', 'TRACEBACK', - 'FAMILY_TYPE', 'ARCHIVE_FORMAT' + 'VERSION', 'FUNCTIONAL', 'RELATIVISTIC', 'PROTOCOL', 'PSEUDO_FORMAT', 'STRINGENCY', 'DEFAULT_STRINGENCY', + 'TRACEBACK', 'FAMILY_TYPE', 'ARCHIVE_FORMAT' ) VERSION = OverridableOption( diff --git a/aiida_pseudo/cli/show.py b/aiida_pseudo/cli/show.py index cc8cc8a..2d71e60 100644 --- a/aiida_pseudo/cli/show.py +++ b/aiida_pseudo/cli/show.py @@ -6,12 +6,12 @@ from aiida.cmdline.utils import decorators, echo from ..groups.mixins import RecommendedCutoffMixin -from .params import PseudoPotentialFamilyParam, options +from .params import arguments, options from .root import cmd_root @cmd_root.command('show') -@click.argument('family', type=PseudoPotentialFamilyParam(sub_classes=('aiida.groups:pseudo.family',))) +@arguments.PSEUDO_POTENTIAL_FAMILY() @options.STRINGENCY() @options_core.RAW() @decorators.with_dbenv() diff --git a/tests/cli/test_family.py b/tests/cli/test_family.py new file mode 100644 index 0000000..f7dc57f --- /dev/null +++ b/tests/cli/test_family.py @@ -0,0 +1,45 @@ +# -*- coding: utf-8 -*- +# pylint: disable=unused-argument +"""Tests for the command `aiida-pseudo show`.""" +import json +import pytest + +from aiida_pseudo.cli.family import cmd_family_cutoffs_set +from aiida_pseudo.groups.family import CutoffsFamily + + +@pytest.mark.usefixtures('clear_db') +def test_family_cutoffs_set(run_cli_command, get_pseudo_family, tmp_path): + """Test the `aiida-pseudo family cutoffs set` command.""" + family = get_pseudo_family(cls=CutoffsFamily) + cutoffs = { + 'normal': {}, + 'high': {}, + } + + for element in family.elements: + cutoffs['normal'][element] = {'cutoff_wfc': 1.0, 'cutoff_rho': 2.0} + cutoffs['high'][element] = {'cutoff_wfc': 3.0, 'cutoff_rho': 6.0} + + # Set only the normal cutoffs for the family + family.set_cutoffs({'normal': cutoffs['normal']}, 'normal') + + filepath = tmp_path / 'cutoffs.json' + + # Invalid JSON + filepath.write_text('invalid content') + result = run_cli_command(cmd_family_cutoffs_set, [family.label, str(filepath)], raises=True) + assert "Error: Missing option '-s' / '--stringency'" in result.output + + # Invalid cutoffs structure + filepath.write_text(json.dumps({'Ar': {'cutoff_rho': 300}})) + result = run_cli_command(cmd_family_cutoffs_set, [family.label, str(filepath), '-s', 'high'], raises=True) + assert 'Error: Invalid value for CUTOFFS:' in result.output + + # Set correct stringency + stringency = 'high' + filepath.write_text(json.dumps(cutoffs['high'])) + result = run_cli_command(cmd_family_cutoffs_set, [family.label, str(filepath), '-s', stringency]) + assert 'Success: set cutoffs for' in result.output + assert stringency in family.get_cutoff_stringencies() + assert family.get_cutoffs(stringency) == cutoffs[stringency]