Skip to content

Commit

Permalink
CLI: add the aiida-pseudo family cutoffs set command
Browse files Browse the repository at this point in the history
This command allows to set recommended cutoffs for a given family. The
cutoffs are specified in a JSON file that is given as an argument and
the stringency is specified through the `-s/--stringency` option. If the
stringency already exists, it is simply overwritten.
  • Loading branch information
sphuber committed Apr 9, 2021
1 parent 051c014 commit c47c3b2
Show file tree
Hide file tree
Showing 6 changed files with 118 additions and 4 deletions.
1 change: 1 addition & 0 deletions aiida_pseudo/cli/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
58 changes: 58 additions & 0 deletions aiida_pseudo/cli/family.py
Original file line number Diff line number Diff line change
@@ -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}`.')
10 changes: 10 additions & 0 deletions aiida_pseudo/cli/params/arguments.py
Original file line number Diff line number Diff line change
@@ -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',))
)
4 changes: 2 additions & 2 deletions aiida_pseudo/cli/params/options.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
4 changes: 2 additions & 2 deletions aiida_pseudo/cli/show.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
45 changes: 45 additions & 0 deletions tests/cli/test_family.py
Original file line number Diff line number Diff line change
@@ -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]

0 comments on commit c47c3b2

Please sign in to comment.