diff --git a/dpgen/generator/arginfo.py b/dpgen/generator/arginfo.py index 0f8abb865..2c46dd1b0 100644 --- a/dpgen/generator/arginfo.py +++ b/dpgen/generator/arginfo.py @@ -1,7 +1,8 @@ import textwrap -from typing import Union +from typing import Optional, Union from dargs import Argument, Variant +from dargs.dargs import ArgumentValueError from dpgen.arginfo import general_mdata_arginfo @@ -987,6 +988,155 @@ def fp_style_custom_args() -> list[Argument]: ] +def fp_style_pwmat_args() -> list[Argument]: + """Return first-principles arguments for PWmat labeling.""" + required_generated_keys = { + "node1", + "node2", + "in.atom", + "ecut", + "e_error", + "rho_error", + "kspacing", + "flag_symm", + } + + def has_required_generated_keys(params): + return required_generated_keys.issubset(params) + + generated_args = [ + Argument("node1", int, optional=False, doc="First PWmat node-grid size."), + Argument("node2", int, optional=False, doc="Second PWmat node-grid size."), + Argument( + "in.atom", + str, + optional=False, + doc="Atom-configuration filename written to the PWmat input.", + ), + Argument( + "ecut", + [int, float], + optional=False, + doc="Plane-wave energy cutoff.", + ), + Argument( + "e_error", + [int, float], + optional=False, + doc="Electronic-energy convergence threshold.", + ), + Argument( + "rho_error", + [int, float], + optional=False, + doc="Charge-density convergence threshold.", + ), + Argument( + "kspacing", + [int, float], + optional=False, + doc="Reciprocal-space spacing used to generate MP_N123.", + ), + Argument( + "flag_symm", + [int, str], + optional=True, + default="NONE", + doc="PWmat symmetry flag: 0, 1, 2, 3, or 'NONE'.", + ), + Argument( + "icmix", + [int, float], + optional=True, + doc="SCF mixing parameter used to build scf_iter0_2.", + ), + Argument( + "smearing", + int, + optional=True, + doc="PWmat smearing method written to SCF iteration settings.", + ), + Argument( + "sigma", + [int, float], + optional=True, + doc="Smearing width written to SCF iteration settings.", + ), + Argument( + "user_pwmat_params", + dict, + optional=True, + doc="Arbitrary PWmat keys overriding the generated input dictionary.", + ), + ] + + return [ + Argument( + "fp_pp_path", + str, + optional=False, + doc="Directory containing PWmat pseudopotential files.", + ), + Argument( + "fp_pp_files", + list[str], + optional=False, + doc="Pseudopotential filenames ordered consistently with type_map.", + ), + Argument( + "fp_incar", + str, + optional=True, + doc="Existing etot.input template; this takes highest priority.", + ), + Argument( + "user_fp_params", + dict, + optional=True, + extra_check=has_required_generated_keys, + extra_check_errmsg=( + "user_fp_params must define node1, node2, in.atom, ecut, " + "e_error, rho_error, kspacing, and flag_symm" + ), + doc=( + "Compatibility input mapping. The current generator consumes " + "node1, node2, in.atom, ecut, e_error, rho_error, kspacing, and " + "flag_symm, regenerates etot.input, and ignores other keys." + ), + ), + Argument( + "fp_params", + dict, + optional=True, + sub_fields=generated_args, + doc=( + "Parameters used by make_pwmat_input_user_dict when neither " + "fp_incar nor user_fp_params is supplied." + ), + ), + ] + + +class _FpStyleVariant(Variant): + """Validate cross-field requirements for first-principles backends.""" + + def get_choice(self, argdict: dict, path: Optional[list[str]] = None) -> Argument: + """Return the selected backend after validating PWmat input sources. + + Dargs flattens variant fields into their parent mapping, so a regular + field-level extra check cannot express this at-least-one constraint. + """ + choice = super().get_choice(argdict, path) + input_sources = {"fp_incar", "user_fp_params", "fp_params"} + if choice.name == "pwmat" and input_sources.isdisjoint(argdict): + raise ArgumentValueError( + path, + "PWmat requires at least one input source: fp_incar, " + "user_fp_params, or fp_params.", + ) + return choice + + def fp_style_variant_type_args() -> Variant: doc_fp_style = "Software for First Principles." doc_amber_diff = ( @@ -1001,8 +1151,12 @@ def fp_style_variant_type_args() -> Variant: "The command argument in the machine file should be the script to run custom FP codes. " "The extra forward and backward files can be defined in the machine file." ) + doc_pwmat = ( + "PWmat density-functional labeling. The machine command should invoke " + "the site-specific PWmat executable." + ) - return Variant( + return _FpStyleVariant( "fp_style", [ Argument("vasp", dict, fp_style_vasp_args()), @@ -1013,7 +1167,7 @@ def fp_style_variant_type_args() -> Variant: Argument( "amber/diff", dict, fp_style_amber_diff_args(), doc=doc_amber_diff ), - Argument("pwmat", dict, [], doc="TODO: add doc"), + Argument("pwmat", dict, fp_style_pwmat_args(), doc=doc_pwmat), Argument("pwscf", dict, fp_style_pwscf_args()), Argument("cpx", dict, fp_style_cpx_args()), Argument("custom", dict, fp_style_custom_args(), doc=doc_custom), diff --git a/tests/test_pwmat_arginfo.py b/tests/test_pwmat_arginfo.py new file mode 100644 index 000000000..910c65cfc --- /dev/null +++ b/tests/test_pwmat_arginfo.py @@ -0,0 +1,94 @@ +"""Validate PWmat first-principles parameter documentation.""" + +import unittest + +from dargs import Argument +from dargs.dargs import ArgumentValueError + +from dpgen.generator.arginfo import fp_style_variant_type_args + + +class TestPWmatArginfo(unittest.TestCase): + def setUp(self): + self.arginfo = Argument("fp", dict, sub_variants=[fp_style_variant_type_args()]) + self.pseudopotentials = { + "fp_style": "pwmat", + "fp_pp_path": ".", + "fp_pp_files": ["C.UPF", "H.UPF"], + } + + def check(self, data): + normalized = self.arginfo.normalize_value(data) + self.arginfo.check_value(normalized, strict=True) + + def test_existing_input_file(self): + self.check({**self.pseudopotentials, "fp_incar": "etot.input"}) + + def test_requires_an_input_source(self): + """A valid PWmat configuration must select one supported input path.""" + with self.assertRaisesRegex(ArgumentValueError, "at least one input source"): + self.check(self.pseudopotentials) + + def test_generated_fp_params(self): + self.check( + { + **self.pseudopotentials, + "fp_params": { + "node1": 4, + "node2": 1, + "in.atom": "atom.config", + "ecut": 50, + "e_error": 1e-4, + "rho_error": 1e-4, + "kspacing": 0.1, + "flag_symm": "NONE", + "icmix": 1.0, + "smearing": 2, + "sigma": 0.025, + "user_pwmat_params": {"job": "SCF", "out.wg": False}, + }, + } + ) + + def test_generated_fp_params_default_symmetry(self): + """PWmat accepts generated inputs without an explicit symmetry override.""" + self.check( + { + **self.pseudopotentials, + "fp_params": { + "node1": 4, + "node2": 1, + "in.atom": "atom.config", + "ecut": 50, + "e_error": 1e-4, + "rho_error": 1e-4, + "kspacing": 0.1, + }, + } + ) + + def test_compatibility_user_fp_params(self): + self.check( + { + **self.pseudopotentials, + "user_fp_params": { + "node1": 4, + "node2": 1, + "job": "SCF", + "in.atom": "atom.config", + "in.psp1": "C.UPF", + "in.psp2": "H.UPF", + "ecut": 50, + "flag_symm": 2, + "e_error": 1e-4, + "rho_error": 1e-4, + "scf_iter0_1": "6 4 3 0.0000 0.025 2", + "xcfunctional": "PBE", + "kspacing": 0.1, + }, + } + ) + + +if __name__ == "__main__": + unittest.main()