diff --git a/README.md b/README.md index d19f705..7f84481 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,13 @@ For development and tests: python -m pip install -e ".[test,lint]" ``` +For a Conda-based development environment with DFTB+ included: + +```bash +conda env create -f environment.yml +conda activate thermoscreening +``` + ## DFTB+ Setup DFTB+ calculations require two external pieces: @@ -34,13 +41,31 @@ DFTB+ calculations require two external pieces: 1. The `dftb+` and `modes` executables on `PATH`. 2. Slater-Koster parameter files downloaded separately from DFTB.org. -ThermoScreening does not vendor Slater-Koster files. Point the calculator to a parameter directory in one of two ways: +Install DFTB+ with Conda if it is not already available: + +```bash +conda install -c conda-forge dftbplus +``` + +Download the default `3ob-3-1` Slater-Koster files into a user-local directory: + +```bash +thermo setup-dftb +``` + +The command prints the `DFTB_PREFIX` export needed by DFTB+ and ThermoScreening: + +```bash +export DFTB_PREFIX="$HOME/.local/share/thermoscreening/slakos/3ob-3-1/" +``` + +Add that line to your shell configuration for persistent use. Verify the setup with: ```bash -export DFTB_PREFIX=/path/to/3ob-3-1/ +thermo doctor ``` -or pass `slako_dir` explicitly: +ThermoScreening does not vendor Slater-Koster files. For custom installations, point the calculator to a parameter directory with `DFTB_PREFIX` or pass `slako_dir` explicitly: ```python from ThermoScreening.thermo.api import dftbplus_thermo diff --git a/ThermoScreening/__main__.py b/ThermoScreening/__main__.py index f40d09a..fc3786f 100644 --- a/ThermoScreening/__main__.py +++ b/ThermoScreening/__main__.py @@ -6,8 +6,10 @@ Example: $ python -m ThermoScreening """ +import sys + from .cli import main if __name__ == "__main__": - main() + sys.exit(main()) diff --git a/ThermoScreening/cli/dftb_setup.py b/ThermoScreening/cli/dftb_setup.py new file mode 100644 index 0000000..acb0a91 --- /dev/null +++ b/ThermoScreening/cli/dftb_setup.py @@ -0,0 +1,179 @@ +""" +Helpers for installing and validating external DFTB+ dependencies. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +import os +import shutil +import tarfile +import tempfile +from urllib.request import urlopen + + +DEFAULT_PARAMETER_SET = "3ob-3-1" +DEFAULT_SLAKO_URL = ( + "https://github.com/dftbparams/3ob/releases/latest/download/" + f"{DEFAULT_PARAMETER_SET}.tar.xz" +) +REQUIRED_PARAMETER_FILE = "C-C.skf" + + +@dataclass(frozen=True) +class Diagnostic: + """ + Result for one DFTB+ setup check. + """ + + name: str + ok: bool + detail: str + + +def default_install_root() -> Path: + """ + Return the default user-local Slater-Koster install directory. + """ + + return Path.home() / ".local" / "share" / "thermoscreening" / "slakos" + + +def default_parameter_dir(install_root: str | Path | None = None) -> Path: + """ + Return the default 3ob parameter directory under an install root. + """ + + root = Path(install_root).expanduser() if install_root is not None else default_install_root() + return root / DEFAULT_PARAMETER_SET + + +def _download_file(url: str, destination: Path, timeout: int = 60) -> None: + """ + Download a URL to a local path. + """ + + with urlopen(url, timeout=timeout) as response: # nosec B310 - user-visible setup helper + with destination.open("wb") as output_file: + shutil.copyfileobj(response, output_file) + + +def _safe_extract_tar(archive_path: Path, destination: Path) -> None: + """ + Extract a tar archive while rejecting path traversal entries. + """ + + destination_resolved = destination.resolve() + + with tarfile.open(archive_path, "r:xz") as archive: + for member in archive.getmembers(): + member_path = destination_resolved / member.name + if not member_path.resolve().is_relative_to(destination_resolved): + raise ValueError(f"Unsafe archive member path: {member.name}") + + archive.extractall(destination_resolved, filter="data") + + +def install_slakos( + install_root: str | Path | None = None, + url: str = DEFAULT_SLAKO_URL, + force: bool = False, +) -> Path: + """ + Download and extract the default Slater-Koster parameter set. + """ + + root = Path(install_root).expanduser() if install_root is not None else default_install_root() + parameter_dir = default_parameter_dir(root) + marker_file = parameter_dir / REQUIRED_PARAMETER_FILE + + if marker_file.exists() and not force: + return parameter_dir.resolve() + + root.mkdir(parents=True, exist_ok=True) + + with tempfile.TemporaryDirectory(prefix="thermoscreening-dftb-") as tmp_dir: + archive_path = Path(tmp_dir) / f"{DEFAULT_PARAMETER_SET}.tar.xz" + _download_file(url, archive_path) + _safe_extract_tar(archive_path, root) + + if not marker_file.exists(): + raise FileNotFoundError( + f"Downloaded parameter set is missing {REQUIRED_PARAMETER_FILE}: " + f"{parameter_dir}" + ) + + return parameter_dir.resolve() + + +def dftb_prefix_export(parameter_dir: str | Path) -> str: + """ + Return the shell export line for a Slater-Koster parameter directory. + """ + + return f'export DFTB_PREFIX="{Path(parameter_dir).expanduser().resolve()}{os.sep}"' + + +def check_dftb_setup(env: dict[str, str] | None = None) -> list[Diagnostic]: + """ + Check whether DFTB+ executables and parameters are available. + """ + + current_env = os.environ if env is None else env + prefix = current_env.get("DFTB_PREFIX") + diagnostics = [ + Diagnostic( + "dftb+", + shutil.which("dftb+") is not None, + shutil.which("dftb+") or "not found on PATH", + ), + Diagnostic( + "modes", + shutil.which("modes") is not None, + shutil.which("modes") or "not found on PATH", + ), + ] + + if not prefix: + diagnostics.extend( + [ + Diagnostic("DFTB_PREFIX", False, "not set"), + Diagnostic(REQUIRED_PARAMETER_FILE, False, "DFTB_PREFIX is not set"), + ] + ) + return diagnostics + + parameter_dir = Path(prefix).expanduser() + parameter_file = parameter_dir / REQUIRED_PARAMETER_FILE + diagnostics.extend( + [ + Diagnostic( + "DFTB_PREFIX", + parameter_dir.is_dir(), + str(parameter_dir.resolve()) if parameter_dir.is_dir() else "directory not found", + ), + Diagnostic( + REQUIRED_PARAMETER_FILE, + parameter_file.is_file(), + str(parameter_file.resolve()) if parameter_file.is_file() else "not found", + ), + ] + ) + + return diagnostics + + +def format_diagnostics(diagnostics: list[Diagnostic]) -> str: + """ + Format setup diagnostics for terminal output. + """ + + lines = [] + width = max(len(item.name) for item in diagnostics) + + for item in diagnostics: + status = "found" if item.ok else "missing" + lines.append(f"{item.name:<{width}} {status:<7} {item.detail}") + + return "\n".join(lines) diff --git a/ThermoScreening/cli/thermo.py b/ThermoScreening/cli/thermo.py index 33e454f..d68b902 100644 --- a/ThermoScreening/cli/thermo.py +++ b/ThermoScreening/cli/thermo.py @@ -1,10 +1,63 @@ from argparse import ArgumentParser -from ThermoScreening.thermo.api import execute +import sys import time + +from ThermoScreening.cli.dftb_setup import ( + DEFAULT_SLAKO_URL, + check_dftb_setup, + dftb_prefix_export, + format_diagnostics, + install_slakos, +) +from ThermoScreening.thermo.api import execute from ThermoScreening.version import __version__ -def parse_args(): +DFTB_COMMANDS = {"setup-dftb", "doctor"} + + +def _run_parser(): + parser = ArgumentParser(description="ThermoScreening") + parser.add_argument("input_file", type=str, help="Input file") + parser.add_argument( + "-v", "--verbose", action="store_true", help="Verbose output", default=True + ) + return parser + + +def _command_parser(): + parser = ArgumentParser(description="ThermoScreening") + subparsers = parser.add_subparsers(dest="command", required=True) + + setup_parser = subparsers.add_parser( + "setup-dftb", + help="Download the default DFTB+ Slater-Koster parameter set.", + ) + setup_parser.add_argument( + "--install-root", + default=None, + help="Directory where Slater-Koster parameter sets are installed.", + ) + setup_parser.add_argument( + "--url", + default=DEFAULT_SLAKO_URL, + help="Archive URL for the default Slater-Koster parameter set.", + ) + setup_parser.add_argument( + "--force", + action="store_true", + help="Download and extract even when the parameter set already exists.", + ) + + subparsers.add_parser( + "doctor", + help="Check DFTB+ executables and Slater-Koster parameter configuration.", + ) + + return parser + + +def parse_args(argv=None): """ Parse command line arguments @@ -13,16 +66,49 @@ def parse_args(): args : argparse.Namespace Command line arguments """ - - parser = ArgumentParser(description="ThermoScreening") - parser.add_argument("input_file", type=str, help="Input file") - parser.add_argument( - "-v", "--verbose", action="store_true", help="Verbose output", default=True - ) - args = parser.parse_args() + if argv is None: + argv = sys.argv[1:] + + if argv and argv[0] in DFTB_COMMANDS: + parser = _command_parser() + args = parser.parse_args(argv) + else: + parser = _run_parser() + args = parser.parse_args(argv) + args.command = "run" + return args +def run_setup_dftb(parser_args): + """ + Download the default DFTB+ Slater-Koster parameter set. + """ + + parameter_dir = install_slakos( + install_root=parser_args.install_root, + url=parser_args.url, + force=parser_args.force, + ) + + print("Slater-Koster files: ", parameter_dir) + print("Shell configuration:") + print(dftb_prefix_export(parameter_dir)) + + return 0 + + +def run_doctor(): + """ + Check whether DFTB+ executables and parameters are available. + """ + + diagnostics = check_dftb_setup() + print(format_diagnostics(diagnostics)) + + return 0 if all(item.ok for item in diagnostics) else 1 + + def main(): """ Main function to run the thermo cli. It parses the command line arguments @@ -38,13 +124,22 @@ def main(): """ parser_args = parse_args() + + command = getattr(parser_args, "command", "run") + + if command == "setup-dftb": + return run_setup_dftb(parser_args) + + if command == "doctor": + return run_doctor() + input_file = parser_args.input_file verbose = parser_args.verbose if verbose: print("Input file: ", input_file) print("Verbose: ", verbose) - + # start timer start = time.time() diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..a434486 --- /dev/null +++ b/environment.yml @@ -0,0 +1,9 @@ +name: thermoscreening +channels: + - conda-forge +dependencies: + - python>=3.12 + - dftbplus + - pip + - pip: + - "-e .[test,lint]" diff --git a/pyproject.toml b/pyproject.toml index 3170a29..a4eaf70 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ docs = [ ] [tool.setuptools.packages.find] -exclude = ["external"] +exclude = ["external", "external.*", "tests", "tests.*"] [tool.setuptools_scm] version_file = "ThermoScreening/__version__.py" diff --git a/tests/cli/test_dftb_setup.py b/tests/cli/test_dftb_setup.py new file mode 100644 index 0000000..1161008 --- /dev/null +++ b/tests/cli/test_dftb_setup.py @@ -0,0 +1,134 @@ +import io +import tarfile + +import pytest + +from ThermoScreening.cli import dftb_setup +from ThermoScreening.cli.dftb_setup import ( + DEFAULT_PARAMETER_SET, + REQUIRED_PARAMETER_FILE, + Diagnostic, + check_dftb_setup, + dftb_prefix_export, + format_diagnostics, + install_slakos, +) + + +def _parameter_archive(tmp_path, files=None): + files = files or {REQUIRED_PARAMETER_FILE: "parameter data"} + source_dir = tmp_path / "source" / DEFAULT_PARAMETER_SET + source_dir.mkdir(parents=True) + + for name, content in files.items(): + target = source_dir / name + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(content, encoding="utf-8") + + archive_path = tmp_path / f"{DEFAULT_PARAMETER_SET}.tar.xz" + with tarfile.open(archive_path, "w:xz") as archive: + archive.add(source_dir, arcname=DEFAULT_PARAMETER_SET) + + return archive_path + + +def test_install_slakos_extracts_archive(tmp_path): + archive_path = _parameter_archive(tmp_path) + + installed_dir = install_slakos( + install_root=tmp_path / "install", + url=archive_path.as_uri(), + ) + + assert installed_dir == (tmp_path / "install" / DEFAULT_PARAMETER_SET).resolve() + assert (installed_dir / REQUIRED_PARAMETER_FILE).read_text( + encoding="utf-8" + ) == "parameter data" + + +def test_install_slakos_reuses_existing_directory(monkeypatch, tmp_path): + marker_file = tmp_path / "install" / DEFAULT_PARAMETER_SET / REQUIRED_PARAMETER_FILE + marker_file.parent.mkdir(parents=True) + marker_file.write_text("existing", encoding="utf-8") + + def fail_download(*args, **kwargs): + raise AssertionError("download should not run") + + monkeypatch.setattr(dftb_setup, "_download_file", fail_download) + + installed_dir = install_slakos(install_root=tmp_path / "install") + + assert installed_dir == marker_file.parent.resolve() + assert marker_file.read_text(encoding="utf-8") == "existing" + + +def test_install_slakos_requires_marker_file(tmp_path): + archive_path = _parameter_archive(tmp_path, files={"H-H.skf": "parameter data"}) + + with pytest.raises(FileNotFoundError, match=REQUIRED_PARAMETER_FILE): + install_slakos(install_root=tmp_path / "install", url=archive_path.as_uri()) + + +def test_install_slakos_rejects_unsafe_archive_member(tmp_path): + archive_path = tmp_path / f"{DEFAULT_PARAMETER_SET}.tar.xz" + data = b"unsafe" + + with tarfile.open(archive_path, "w:xz") as archive: + member = tarfile.TarInfo("../unsafe.txt") + member.size = len(data) + archive.addfile(member, io.BytesIO(data)) + + with pytest.raises(ValueError, match="Unsafe archive member path"): + install_slakos(install_root=tmp_path / "install", url=archive_path.as_uri()) + + assert not (tmp_path / "unsafe.txt").exists() + + +def test_dftb_prefix_export_adds_trailing_separator(tmp_path): + expected = f'export DFTB_PREFIX="{tmp_path.resolve()}/"' + + assert dftb_prefix_export(tmp_path) == expected + + +def test_check_dftb_setup_reports_ready_environment(monkeypatch, tmp_path): + marker_file = tmp_path / REQUIRED_PARAMETER_FILE + marker_file.write_text("parameter data", encoding="utf-8") + + def fake_which(command): + return f"/usr/bin/{command}" + + monkeypatch.setattr(dftb_setup.shutil, "which", fake_which) + + diagnostics = check_dftb_setup({"DFTB_PREFIX": str(tmp_path)}) + + assert [(item.name, item.ok) for item in diagnostics] == [ + ("dftb+", True), + ("modes", True), + ("DFTB_PREFIX", True), + (REQUIRED_PARAMETER_FILE, True), + ] + + +def test_check_dftb_setup_reports_missing_environment(monkeypatch): + monkeypatch.setattr(dftb_setup.shutil, "which", lambda command: None) + + diagnostics = check_dftb_setup({}) + + assert [(item.name, item.ok) for item in diagnostics] == [ + ("dftb+", False), + ("modes", False), + ("DFTB_PREFIX", False), + (REQUIRED_PARAMETER_FILE, False), + ] + + +def test_format_diagnostics_aligns_statuses(): + output = format_diagnostics( + [ + Diagnostic("dftb+", True, "/usr/bin/dftb+"), + Diagnostic("DFTB_PREFIX", False, "not set"), + ] + ) + + assert "dftb+ found" in output + assert "DFTB_PREFIX missing" in output diff --git a/tests/thermo/test_main.py b/tests/thermo/test_main.py index c28c8a8..110a688 100644 --- a/tests/thermo/test_main.py +++ b/tests/thermo/test_main.py @@ -76,5 +76,68 @@ def test_main_executes_input_file_without_verbose(monkeypatch, capsys): assert calls == ["thermo.in"] assert capsys.readouterr().out == "" + +def test_parse_args_setup_dftb_command(): + args = thermo.parse_args( + [ + "setup-dftb", + "--install-root", + "/tmp/slakos", + "--url", + "file:///tmp/3ob-3-1.tar.xz", + "--force", + ] + ) + + assert args.command == "setup-dftb" + assert args.install_root == "/tmp/slakos" + assert args.url == "file:///tmp/3ob-3-1.tar.xz" + assert args.force is True + + +def test_parse_args_doctor_command(): + args = thermo.parse_args(["doctor"]) + + assert args.command == "doctor" + + +def test_main_runs_setup_dftb(monkeypatch, tmp_path, capsys): + monkeypatch.setattr( + thermo, + "parse_args", + lambda: argparse.Namespace( + command="setup-dftb", + install_root=str(tmp_path), + url="file:///tmp/3ob-3-1.tar.xz", + force=True, + ), + ) + monkeypatch.setattr( + thermo, + "install_slakos", + lambda install_root, url, force: tmp_path / "3ob-3-1", + ) + + assert thermo.main() == 0 + + output = capsys.readouterr().out + assert "Slater-Koster files:" in output + assert "export DFTB_PREFIX=" in output + + +def test_main_runs_doctor(monkeypatch, capsys): + diagnostic = type("DiagnosticStub", (), {"ok": False})() + + monkeypatch.setattr( + thermo, + "parse_args", + lambda: argparse.Namespace(command="doctor"), + ) + monkeypatch.setattr(thermo, "check_dftb_setup", lambda: [diagnostic]) + monkeypatch.setattr(thermo, "format_diagnostics", lambda diagnostics: "not ready") + + assert thermo.main() == 1 + assert capsys.readouterr().out == "not ready\n" + if __name__ == '__main__': unittest.main()