diff --git a/docs/source/apidoc/janus_core.rst b/docs/source/apidoc/janus_core.rst index 7570e295..984bd207 100644 --- a/docs/source/apidoc/janus_core.rst +++ b/docs/source/apidoc/janus_core.rst @@ -4,6 +4,16 @@ janus\_core package Submodules ---------- +janus\_core.cli module +---------------------- + +.. automodule:: janus_core.cli + :members: + :special-members: + :private-members: + :undoc-members: + :show-inheritance: + janus\_core.geom\_opt module ---------------------------- diff --git a/janus_core/cli.py b/janus_core/cli.py new file mode 100644 index 00000000..03a4347d --- /dev/null +++ b/janus_core/cli.py @@ -0,0 +1,143 @@ +"""Set up commandline interface.""" + +import ast +from typing import Annotated + +import typer + +from janus_core.single_point import SinglePoint + +app = typer.Typer() + + +class TyperDict: # pylint: disable=too-few-public-methods + """ + Custom dictionary for typer. + + Parameters + ---------- + value : str + Value of string representing a dictionary. + """ + + def __init__(self, value: str): + """ + Initialise class. + + Parameters + ---------- + value : str + Value of string representing a dictionary. + """ + self.value = value + + def __str__(self): + """ + String representation of class. + + Returns + ------- + str + Class name and value of string representing a dictionary. + """ + return f"" + + +def parse_dict_class(value: str): + """ + Convert string input into a dictionary. + + Parameters + ---------- + value : str + String representing dictionary to be parsed. + + Returns + ------- + TyperDict + Parsed string as a dictionary. + """ + return TyperDict(ast.literal_eval(value)) + + +@app.command() +def singlepoint( + structure: Annotated[ + str, typer.Option(help="Path to structure to perform calculations") + ], + architecture: Annotated[ + str, typer.Option("--arch", help="MLIP architecture to use for calculations") + ] = "mace_mp", + device: Annotated[str, typer.Option(help="Device to run calculations on")] = "cpu", + properties: Annotated[ + list[str], + typer.Option( + "--property", + help="Properties to calculate. If not specified, 'energy', 'forces', and 'stress' will be returned.", + ), + ] = None, + read_kwargs: Annotated[ + TyperDict, + typer.Option( + parser=parse_dict_class, + help="Keyword arguments to pass to ase.io.read [default: {}]", + metavar="DICT", + ), + ] = None, + calc_kwargs: Annotated[ + TyperDict, + typer.Option( + parser=parse_dict_class, + help="Keyword arguments to pass to selected calculator [default: {}]", + metavar="DICT", + ), + ] = None, +): + """ + Perform single point calculations. + + Parameters + ---------- + structure : str + Structure to simulate. + architecture : Optional[str] + MLIP architecture to use for single point calculations. + Default is "mace_mp". + device : Optional[str] + Device to run model on. Default is "cpu". + properties : Optional[str] + Physical properties to calculate. Default is "energy". + read_kwargs : Optional[dict[str, Any]] + Keyword arguments to pass to ase.io.read. Default is {}. + calc_kwargs : Optional[dict[str, Any]] + Keyword arguments to pass to the selected calculator. Default is {}. + """ + read_kwargs = read_kwargs.value if read_kwargs else {} + calc_kwargs = calc_kwargs.value if calc_kwargs else {} + + if not isinstance(read_kwargs, dict): + raise ValueError("read_kwargs must be a dictionary") + if not isinstance(calc_kwargs, dict): + raise ValueError("calc_kwargs must be a dictionary") + + s_point = SinglePoint( + structure=structure, + architecture=architecture, + device=device, + read_kwargs=read_kwargs, + calc_kwargs=calc_kwargs, + ) + print(s_point.run_single_point(properties=properties)) + + +@app.command() +def test(name: str): + """ + Dummy alternative CLI command. + + Parameters + ---------- + name : str + Name of person. + """ + print(f"Hello, {name}!") diff --git a/janus_core/single_point.py b/janus_core/single_point.py index cdd4663c..5de81564 100644 --- a/janus_core/single_point.py +++ b/janus_core/single_point.py @@ -15,8 +15,8 @@ class SinglePoint: Parameters ---------- - system : str - System to simulate. + structure : str + Structure to simulate. architecture : Literal[architectures] MLIP architecture to use for single point calculations. Default is "mace_mp". @@ -24,43 +24,47 @@ class SinglePoint: Device to run model on. Default is "cpu". read_kwargs : Optional[dict[str, Any]] Keyword arguments to pass to ase.io.read. Default is {}. - **kwargs - Additional keyword arguments passed to the selected calculator. + calc_kwargs : Optional[dict[str, Any]] + Keyword arguments to pass to the selected calculator. Default is {}. Attributes ---------- architecture : Literal[architectures] MLIP architecture to use for single point calculations. - system : str - System to simulate. + structure : str + Path of structure to simulate. device : Literal[devices] Device to run MLIP model on. + struct : Union[Atoms, list[Atoms] + ASE Atoms or list of Atoms structures to simulate. + structname : str + Name of structure from its filename. Methods ------- - read_system(**kwargs) - Read system and system name. + read_structure(**kwargs) + Read structure and structure name. set_calculator(**kwargs) - Configure calculator and attach to system. + Configure calculator and attach to structure. run_single_point(properties=None) Run single point calculations. """ def __init__( self, - system: str, + structure: str, architecture: Literal[architectures] = "mace_mp", device: Literal[devices] = "cpu", read_kwargs: Optional[dict[str, Any]] = None, - **kwargs, + calc_kwargs: Optional[dict[str, Any]] = None, ) -> None: """ - Read the system being simulated and attach an MLIP calculator. + Read the structure being simulated and attach an MLIP calculator. Parameters ---------- - system : str - System to simulate. + structure : str + Path of structure to simulate. architecture : Literal[architectures] MLIP architecture to use for single point calculations. Default is "mace_mp". @@ -68,21 +72,22 @@ def __init__( Device to run MLIP model on. Default is "cpu". read_kwargs : Optional[dict[str, Any]] Keyword arguments to pass to ase.io.read. Default is {}. - **kwargs - Additional keyword arguments passed to the selected calculator. + calc_kwargs : Optional[dict[str, Any]] + Keyword arguments to pass to the selected calculator. Default is {}. """ self.architecture = architecture self.device = device - self.system = system + self.structure = structure - # Read system and get calculator + # Read structure and get calculator read_kwargs = read_kwargs if read_kwargs else {} - self.read_system(**read_kwargs) - self.set_calculator(**kwargs) + calc_kwargs = calc_kwargs if calc_kwargs else {} + self.read_structure(**read_kwargs) + self.set_calculator(**calc_kwargs) - def read_system(self, **kwargs) -> None: + def read_structure(self, **kwargs) -> None: """ - Read system and system name. + Read structure and structure name. If the file contains multiple structures, only the last configuration will be read by default. @@ -92,14 +97,14 @@ def read_system(self, **kwargs) -> None: **kwargs Keyword arguments passed to ase.io.read. """ - self.sys = read(self.system, **kwargs) - self.sysname = pathlib.Path(self.system).stem + self.struct = read(self.structure, **kwargs) + self.structname = pathlib.Path(self.structure).stem def set_calculator( self, read_kwargs: Optional[dict[str, Any]] = None, **kwargs ) -> None: """ - Configure calculator and attach to system. + Configure calculator and attach to structure. Parameters ---------- @@ -113,15 +118,15 @@ def set_calculator( device=self.device, **kwargs, ) - if self.sys is None: + if self.struct is None: read_kwargs = read_kwargs if read_kwargs else {} - self.read_system(**read_kwargs) + self.read_structure(**read_kwargs) - if isinstance(self.sys, list): - for sys in self.sys: - sys.calc = calculator + if isinstance(self.struct, list): + for struct in self.struct: + struct.calc = calculator else: - self.sys.calc = calculator + self.struct.calc = calculator def _get_potential_energy(self) -> Union[float, list[float]]: """ @@ -130,12 +135,12 @@ def _get_potential_energy(self) -> Union[float, list[float]]: Returns ------- Union[float, list[float]] - Potential energy of system(s). + Potential energy of structure(s). """ - if isinstance(self.sys, list): - return [sys.get_potential_energy() for sys in self.sys] + if isinstance(self.struct, list): + return [struct.get_potential_energy() for struct in self.struct] - return self.sys.get_potential_energy() + return self.struct.get_potential_energy() def _get_forces(self) -> Union[ndarray, list[ndarray]]: """ @@ -144,12 +149,12 @@ def _get_forces(self) -> Union[ndarray, list[ndarray]]: Returns ------- Union[ndarray, list[ndarray]] - Forces of system(s). + Forces of structure(s). """ - if isinstance(self.sys, list): - return [sys.get_forces() for sys in self.sys] + if isinstance(self.struct, list): + return [struct.get_forces() for struct in self.struct] - return self.sys.get_forces() + return self.struct.get_forces() def _get_stress(self) -> Union[ndarray, list[ndarray]]: """ @@ -158,12 +163,12 @@ def _get_stress(self) -> Union[ndarray, list[ndarray]]: Returns ------- Union[ndarray, list[ndarray]] - Stress of system(s). + Stress of structure(s). """ - if isinstance(self.sys, list): - return [sys.get_stress() for sys in self.sys] + if isinstance(self.struct, list): + return [struct.get_stress() for struct in self.struct] - return self.sys.get_stress() + return self.struct.get_stress() def run_single_point( self, properties: Optional[Union[str, list[str]]] = None diff --git a/pyproject.toml b/pyproject.toml index dd7da6e2..0f8fa2e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,10 +22,14 @@ classifiers = [ repository = "https://github.com/stfc/janus-core/" documentation = "https://stfc.github.io/janus-core/" +[tool.poetry.scripts] +janus = "janus_core.cli:app" + [tool.poetry.dependencies] python = "^3.9" ase = "^3.22.1" mace-torch = "^0.3.4" +typer = "^0.9.0" [tool.poetry.group.dev.dependencies] coverage = {extras = ["toml"], version = "^7.4.1"} diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 00000000..39457695 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,92 @@ +"""Test commandline interface.""" + +from pathlib import Path + +from typer.testing import CliRunner + +from janus_core.cli import app + +DATA_PATH = Path(__file__).parent / "data" + +runner = CliRunner() + + +def test_janus_help(): + """Test calling `janus --help`.""" + result = runner.invoke(app, ["--help"]) + assert result.exit_code == 0 + # Command is returned as "root" + assert "Usage: root [OPTIONS] COMMAND [ARGS]..." in result.stdout + + +def test_singlepoint_help(): + """Test calling `janus singlepoint --help`.""" + result = runner.invoke(app, ["singlepoint", "--help"]) + assert result.exit_code == 0 + # Command is returned as "root" + assert "Usage: root singlepoint [OPTIONS]" in result.stdout + + +def test_singlepoint(): + """Test singlepoint calculation.""" + result = runner.invoke(app, ["singlepoint", "--structure", DATA_PATH / "NaCl.cif"]) + assert result.exit_code == 0 + assert "{'energy': -" in result.stdout + assert "'forces': array" in result.stdout + assert "'stress': array" in result.stdout + + +def test_singlepoint_properties(): + """Test properties for singlepoint calculation.""" + result = runner.invoke( + app, + [ + "singlepoint", + "--structure", + DATA_PATH / "NaCl.cif", + "--property", + "energy", + "--property", + "forces", + ], + ) + assert result.exit_code == 0 + assert "{'energy': -" in result.stdout + assert "'forces': array" in result.stdout + assert "'stress': array" not in result.stdout + + +def test_singlepoint_read_kwargs(): + """Test setting read_kwargs for singlepoint calculation.""" + result = runner.invoke( + app, + [ + "singlepoint", + "--structure", + DATA_PATH / "benzene-traj.xyz", + "--read-kwargs", + "{'index': ':'}", + "--property", + "energy", + ], + ) + assert result.exit_code == 0 + assert "'energy': [" in result.stdout + + +def test_singlepoint_calc_kwargs(): + """Test setting calc_kwargs for singlepoint calculation.""" + result = runner.invoke( + app, + [ + "singlepoint", + "--structure", + DATA_PATH / "NaCl.cif", + "--calc-kwargs", + "{'default_dtype': 'float32'}", + "--property", + "energy", + ], + ) + assert result.exit_code == 0 + assert "Using float32 for MACECalculator" in result.stdout diff --git a/tests/test_geom_opt.py b/tests/test_geom_opt.py index 7348478a..d40e9477 100644 --- a/tests/test_geom_opt.py +++ b/tests/test_geom_opt.py @@ -16,53 +16,38 @@ MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model" test_data = [ - ("mace", "NaCl.cif", MODEL_PATH, -27.046359959669214, {}), - ("mace", "NaCl.cif", MODEL_PATH, -27.04636199814088, {"fmax": 0.001}), + ("mace", "NaCl.cif", -27.046359959669214, {}), + ("mace", "NaCl.cif", -27.04636199814088, {"fmax": 0.001}), + ("mace", "NaCl.cif", -27.0463392211678, {"filter_func": UnitCellFilter}), + ("mace", "H2O.cif", -14.051389496520015, {"filter_func": None}), ( "mace", "NaCl.cif", - MODEL_PATH, - -27.0463392211678, - {"filter_func": UnitCellFilter}, - ), - ("mace", "H2O.cif", MODEL_PATH, -14.051389496520015, {"filter_func": None}), - ( - "mace", - "NaCl.cif", - MODEL_PATH, -27.04631447979008, {"filter_func": UnitCellFilter, "filter_kwargs": {"scalar_pressure": 0.0001}}, ), + ("mace", "NaCl.cif", -27.046353221978332, {"opt_kwargs": {"alpha": 100}}), ( "mace", "NaCl.cif", - MODEL_PATH, - -27.046353221978332, - {"opt_kwargs": {"alpha": 100}}, - ), - ( - "mace", - "NaCl.cif", - MODEL_PATH, -27.03561540212425, {"filter_func": UnitCellFilter, "dyn_kwargs": {"steps": 1}}, ), ] -@pytest.mark.parametrize( - "architecture, structure, model_path, expected, kwargs", test_data -) -def test_optimize(architecture, structure, model_path, expected, kwargs): +@pytest.mark.parametrize("architecture, structure, expected, kwargs", test_data) +def test_optimize(architecture, structure, expected, kwargs): """Test optimizing geometry using MACE.""" - data_path = DATA_PATH / structure single_point = SinglePoint( - system=data_path, architecture=architecture, model_paths=model_path + structure=DATA_PATH / structure, + architecture=architecture, + calc_kwargs={"model_paths": MODEL_PATH}, ) init_energy = single_point.run_single_point("energy")["energy"] - atoms = optimize(single_point.sys, **kwargs) + atoms = optimize(single_point.struct, **kwargs) assert atoms.get_potential_energy() < init_energy assert atoms.get_potential_energy() == pytest.approx(expected) @@ -70,16 +55,18 @@ def test_optimize(architecture, structure, model_path, expected, kwargs): def test_saving_struct(tmp_path): """Test saving optimized structure.""" - data_path = DATA_PATH / "NaCl.cif" struct_path = tmp_path / "NaCl.xyz" + single_point = SinglePoint( - system=data_path, architecture="mace", model_paths=MODEL_PATH + structure=DATA_PATH / "NaCl.cif", + architecture="mace", + calc_kwargs={"model_paths": MODEL_PATH}, ) init_energy = single_point.run_single_point("energy")["energy"] optimize( - single_point.sys, + single_point.struct, struct_kwargs={"filename": struct_path, "format": "extxyz"}, ) opt_struct = read(struct_path) @@ -89,27 +76,31 @@ def test_saving_struct(tmp_path): def test_saving_traj(tmp_path): """Test saving optimization trajectory output.""" - data_path = DATA_PATH / "NaCl.cif" single_point = SinglePoint( - system=data_path, architecture="mace", model_paths=MODEL_PATH + structure=DATA_PATH / "NaCl.cif", + architecture="mace", + calc_kwargs={"model_paths": MODEL_PATH}, + ) + optimize( + single_point.struct, opt_kwargs={"trajectory": str(tmp_path / "NaCl.traj")} ) - optimize(single_point.sys, opt_kwargs={"trajectory": str(tmp_path / "NaCl.traj")}) traj = read(tmp_path / "NaCl.traj", index=":") assert len(traj) == 3 def test_traj_reformat(tmp_path): """Test saving optimization trajectory in different format.""" - data_path = DATA_PATH / "NaCl.cif" - traj_path_binary = tmp_path / "NaCl.traj" - traj_path_xyz = tmp_path / "NaCl-traj.xyz" - single_point = SinglePoint( - system=data_path, architecture="mace", model_paths=MODEL_PATH + structure=DATA_PATH / "NaCl.cif", + architecture="mace", + calc_kwargs={"model_paths": MODEL_PATH}, ) + traj_path_binary = tmp_path / "NaCl.traj" + traj_path_xyz = tmp_path / "NaCl-traj.xyz" + optimize( - single_point.sys, + single_point.struct, opt_kwargs={"trajectory": str(traj_path_binary)}, traj_kwargs={"filename": traj_path_xyz}, ) @@ -120,10 +111,11 @@ def test_traj_reformat(tmp_path): def test_missing_traj_kwarg(tmp_path): """Test saving optimization trajectory in different format.""" - data_path = DATA_PATH / "NaCl.cif" - traj_path = tmp_path / "NaCl-traj.xyz" single_point = SinglePoint( - system=data_path, architecture="mace", model_paths=MODEL_PATH + structure=DATA_PATH / "NaCl.cif", + architecture="mace", + calc_kwargs={"model_paths": MODEL_PATH}, ) + traj_path = tmp_path / "NaCl-traj.xyz" with pytest.raises(ValueError): - optimize(single_point.sys, traj_kwargs={"filename": traj_path}) + optimize(single_point.struct, traj_kwargs={"filename": traj_path}) diff --git a/tests/test_single_point.py b/tests/test_single_point.py index 44b77ec4..957d4bc3 100644 --- a/tests/test_single_point.py +++ b/tests/test_single_point.py @@ -10,35 +10,32 @@ MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model" test_data = [ - (DATA_PATH / "benzene.xyz", -76.0605725422795, "energy", "energy", {}), + (DATA_PATH / "benzene.xyz", -76.0605725422795, "energy", "energy", {}, None), ( DATA_PATH / "benzene.xyz", -76.06057739257812, ["energy"], "energy", {"default_dtype": "float32"}, + None, ), - ( - DATA_PATH / "benzene.xyz", - -0.0360169762840179, - ["forces"], - "forces", - {"idx": [0, 1]}, - ), - (DATA_PATH / "NaCl.cif", -0.004783275999053424, ["stress"], "stress", {"idx": [0]}), + (DATA_PATH / "benzene.xyz", -0.0360169762840179, ["forces"], "forces", {}, [0, 1]), + (DATA_PATH / "NaCl.cif", -0.004783275999053424, ["stress"], "stress", {}, [0]), ] -@pytest.mark.parametrize("system, expected, properties, prop_key, kwargs", test_data) -def test_potential_energy(system, expected, properties, prop_key, kwargs): +@pytest.mark.parametrize( + "structure, expected, properties, prop_key, calc_kwargs, idx", test_data +) +def test_potential_energy(structure, expected, properties, prop_key, calc_kwargs, idx): """Test single point energy using MACE calculators.""" + calc_kwargs["model_paths"] = MODEL_PATH single_point = SinglePoint( - system=system, architecture="mace", model_paths=MODEL_PATH, **kwargs + structure=structure, architecture="mace", calc_kwargs=calc_kwargs ) results = single_point.run_single_point(properties)[prop_key] # Check correct values returned - idx = kwargs.pop("idx", None) if idx is not None: if len(idx) == 1: assert results[idx[0]] == pytest.approx(expected) @@ -52,10 +49,10 @@ def test_potential_energy(system, expected, properties, prop_key, kwargs): def test_single_point_none(): """Test single point stress using MACE calculator.""" - data_path = DATA_PATH / "NaCl.cif" - model_path = MODEL_PATH single_point = SinglePoint( - system=data_path, architecture="mace", model_paths=model_path + structure=DATA_PATH / "NaCl.cif", + architecture="mace", + calc_kwargs={"model_paths": MODEL_PATH}, ) results = single_point.run_single_point() @@ -65,16 +62,14 @@ def test_single_point_none(): def test_single_point_traj(): """Test single point stress using MACE calculator.""" - data_path = DATA_PATH / "benzene-traj.xyz" - model_path = MODEL_PATH single_point = SinglePoint( - system=data_path, + structure=DATA_PATH / "benzene-traj.xyz", architecture="mace", - model_paths=model_path, read_kwargs={"index": ":"}, + calc_kwargs={"model_paths": MODEL_PATH}, ) - assert len(single_point.sys) == 2 + assert len(single_point.struct) == 2 results = single_point.run_single_point("energy") assert results["energy"][0] == pytest.approx(-76.0605725422795) assert results["energy"][1] == pytest.approx(-74.80419118083256)