diff --git a/.gitignore b/.gitignore index 79d452a0..6f245877 100644 --- a/.gitignore +++ b/.gitignore @@ -45,3 +45,7 @@ __pycache__ # autogenerated by setuptools-scm _version.py + +# Benchmark/plot artifacts +benchmarks/results/ +benchmarks/plots/ diff --git a/README.md b/README.md index d247c6e9..eb11e5c1 100644 --- a/README.md +++ b/README.md @@ -73,6 +73,13 @@ python -m pip install -e .[test] pytest ``` +### Benchmarks +Pytest-based benchmark suites live under `benchmarks/` and can be run with +``` +python -m pytest benchmarks -m benchmark +``` +Results are written to `benchmarks/results` as JSON and CSV artifacts. + ## Citation diff --git a/benchmarks/.gitignore b/benchmarks/.gitignore new file mode 100644 index 00000000..8ce69486 --- /dev/null +++ b/benchmarks/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +results/ +plots/ diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 00000000..d98b6f18 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,76 @@ +# CPPE Benchmarks + +This directory contains pytest-based benchmark suites that compare CPPE backends +and track scaling with the number of polarizable sites. + +## Run + +From a configured development environment: + +```bash +python -m pytest benchmarks -m benchmark +``` + +Single-thread parity (recommended for fair comparisons): + +```bash +OMP_NUM_THREADS=1 NUMBA_NUM_THREADS=1 python -m pytest benchmarks -m benchmark +``` + +Optional runtime controls: + +- `CPPE_BENCH_REPEAT` (default: `3`) +- `CPPE_BENCH_WARMUP` (default: `1`) +- `--benchmark-output-dir` (default: `benchmarks/results`) + +Example: + +```bash +CPPE_BENCH_REPEAT=5 CPPE_BENCH_WARMUP=2 \ +python -m pytest benchmarks -m benchmark --benchmark-output-dir benchmarks/results +``` + +## Outputs + +Each run writes both JSON and CSV artifacts: + +- `benchmarks/results/latest.json` +- `benchmarks/results/latest.csv` +- timestamped copies for history + +## Plotting + +Generate runtime and speedup plots from benchmark CSV data: + +```bash +python benchmarks/plot_benchmarks.py \ + --input benchmarks/results/latest.csv \ + --output-dir benchmarks/plots \ + --format png +``` + +By default, this writes **compact overview plots** only. + +Optional: + +- add additional output formats: `--format svg` +- disable log y-axis: `--no-log-y` +- generate per-benchmark detailed plots: `--detailed` + +The script writes: + +- `benchmarks/plots/summary_aggregated.csv` +- `benchmarks/plots/summary_speedup.csv` +- `benchmarks/plots/overview_runtime.*` +- `benchmarks/plots/overview_speedup.*` +- `benchmarks/plots/runtime_*.png` and `speedup_*.png` (only with `--detailed`) +- `benchmarks/plots/REPORT.md` + +## Large synthetic systems + +Large synthetic waterbox-like benchmarks are provided separately and are not run in +the default benchmark marker selection. + +```bash +OMP_NUM_THREADS=1 NUMBA_NUM_THREADS=1 python -m pytest benchmarks -m benchmark_large +``` diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/benchmarks/conftest.py b/benchmarks/conftest.py new file mode 100644 index 00000000..c905e8ab --- /dev/null +++ b/benchmarks/conftest.py @@ -0,0 +1,76 @@ +import csv +import json +import os +from datetime import datetime, timezone +from pathlib import Path + +import pytest + +os.environ.setdefault("OMP_NUM_THREADS", "1") +os.environ.setdefault("NUMBA_NUM_THREADS", "1") + + +def pytest_addoption(parser): + parser.addoption( + "--benchmark-output-dir", + action="store", + default="benchmarks/results", + help="Directory for benchmark result artifacts.", + ) + + +def pytest_configure(config): + config.addinivalue_line( + "markers", + "benchmark: marks benchmark-style tests that collect timing metrics", + ) + config.addinivalue_line( + "markers", + "benchmark_large: marks large synthetic benchmark tests", + ) + + +@pytest.fixture(scope="session", autouse=True) +def benchmark_thread_control(): + from numba import set_num_threads + + threads = int(os.environ.get("NUMBA_NUM_THREADS", "1")) + set_num_threads(threads) + + +@pytest.fixture(scope="session") +def benchmark_recorder(request): + records = [] + + def _record(**kwargs): + records.append(kwargs) + + yield _record + + output_dir = Path(request.config.getoption("benchmark_output_dir")) + output_dir.mkdir(parents=True, exist_ok=True) + + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + payload = { + "created_utc": timestamp, + "repeat": int(os.environ.get("CPPE_BENCH_REPEAT", "3")), + "warmup": int(os.environ.get("CPPE_BENCH_WARMUP", "1")), + "omp_num_threads": int(os.environ.get("OMP_NUM_THREADS", "1")), + "numba_num_threads": int(os.environ.get("NUMBA_NUM_THREADS", "1")), + "records": records, + } + + latest_json = output_dir / "latest.json" + latest_json.write_text(json.dumps(payload, indent=2), encoding="utf-8") + stamped_json = output_dir / f"benchmarks-{timestamp}.json" + stamped_json.write_text(json.dumps(payload, indent=2), encoding="utf-8") + + if records: + fieldnames = sorted({k for rec in records for k in rec.keys()}) + latest_csv = output_dir / "latest.csv" + stamped_csv = output_dir / f"benchmarks-{timestamp}.csv" + for csv_path in (latest_csv, stamped_csv): + with csv_path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(records) diff --git a/benchmarks/plot_benchmarks.py b/benchmarks/plot_benchmarks.py new file mode 100644 index 00000000..ab8d3457 --- /dev/null +++ b/benchmarks/plot_benchmarks.py @@ -0,0 +1,328 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +from pathlib import Path + +import matplotlib.pyplot as plt +import pandas as pd + + +def _parse_args(): + parser = argparse.ArgumentParser( + description="Plot CPPE benchmark results and generate summary artifacts." + ) + parser.add_argument( + "--input", + default="benchmarks/results/latest.csv", + help="Path to benchmark CSV produced by benchmarks/conftest.py", + ) + parser.add_argument( + "--output-dir", + default="benchmarks/plots", + help="Directory for generated plots and summary tables.", + ) + parser.add_argument( + "--format", + dest="formats", + action="append", + default=["png"], + help="Output image format (repeatable), e.g. --format png --format svg", + ) + parser.add_argument( + "--dpi", type=int, default=160, help="Image DPI for raster plots." + ) + parser.add_argument( + "--no-log-y", + action="store_true", + help="Disable logarithmic y-axis on runtime plots.", + ) + parser.add_argument( + "--detailed", + action="store_true", + help="Also generate per-benchmark detailed plots.", + ) + return parser.parse_args() + + +def _slug(text: str) -> str: + return "".join(c if c.isalnum() else "_" for c in text).strip("_").lower() + + +def _save_fig(fig, out_base: Path, formats, dpi: int): + for fmt in formats: + fig.savefig(out_base.with_suffix(f".{fmt}"), dpi=dpi, bbox_inches="tight") + + +def _cleanup_detailed_outputs(out_dir: Path): + for pattern in ("runtime_*", "speedup_*"): + for path in out_dir.glob(pattern): + if path.name.startswith("overview_"): + continue + if path.is_file(): + path.unlink() + + +def _aggregate(df: pd.DataFrame) -> pd.DataFrame: + grouped = ( + df.groupby(["benchmark", "dataset", "backend", "n_polsites"], as_index=False) + .agg( + mean_s=("mean_s", "mean"), + std_s=("mean_s", "std"), + min_s=("mean_s", "min"), + max_s=("mean_s", "max"), + samples=("mean_s", "count"), + ) + .fillna({"std_s": 0.0}) + ) + return grouped + + +def _speedup_table(agg: pd.DataFrame) -> pd.DataFrame: + pivot = agg.pivot_table( + index=["benchmark", "dataset", "n_polsites"], + columns="backend", + values="mean_s", + ).reset_index() + if "cpp" in pivot.columns and "python" in pivot.columns: + pivot["python_over_cpp"] = pivot["python"] / pivot["cpp"] + pivot["cpp_over_python"] = pivot["cpp"] / pivot["python"] + return pivot + + +def _plot_runtime_curves( + agg: pd.DataFrame, out_dir: Path, formats, dpi: int, log_y: bool +): + for (benchmark, dataset), sub in agg.groupby(["benchmark", "dataset"]): + fig, ax = plt.subplots(figsize=(8.0, 5.0)) + for backend, bdf in sub.groupby("backend"): + bdf = bdf.sort_values("n_polsites") + ax.errorbar( + bdf["n_polsites"], + bdf["mean_s"], + yerr=bdf["std_s"], + marker="o", + linewidth=1.8, + capsize=3, + label=backend, + ) + + ax.set_xscale("log", base=2) + if log_y: + ax.set_yscale("log") + ax.set_xlabel("Number of polarizable sites (n_polsites)") + ax.set_ylabel("Runtime [s]") + ax.set_title(f"{benchmark} ({dataset})") + ax.grid(alpha=0.3) + ax.legend() + + out_base = out_dir / f"runtime_{_slug(benchmark)}_{_slug(dataset)}" + _save_fig(fig, out_base, formats, dpi) + plt.close(fig) + + +def _plot_speedups(speedup: pd.DataFrame, out_dir: Path, formats, dpi: int): + needed = {"cpp", "python", "python_over_cpp"} + if not needed.issubset(speedup.columns): + return + + for (benchmark, dataset), sub in speedup.groupby(["benchmark", "dataset"]): + sub = sub.sort_values("n_polsites") + fig, ax = plt.subplots(figsize=(8.0, 4.5)) + ax.plot( + sub["n_polsites"], + sub["python_over_cpp"], + marker="o", + linewidth=1.8, + color="#1f77b4", + ) + ax.axhline(1.0, color="black", linestyle="--", linewidth=1.0) + ax.set_xscale("log", base=2) + ax.set_xlabel("Number of polarizable sites (n_polsites)") + ax.set_ylabel("Python / C++ runtime") + ax.set_title(f"Speedup ratio: {benchmark} ({dataset})") + ax.grid(alpha=0.3) + + out_base = out_dir / f"speedup_{_slug(benchmark)}_{_slug(dataset)}" + _save_fig(fig, out_base, formats, dpi) + plt.close(fig) + + +def _plot_overview_runtime( + agg: pd.DataFrame, out_dir: Path, formats, dpi: int, log_y: bool +): + benchmarks = sorted(agg["benchmark"].unique()) + ncols = 2 + nrows = (len(benchmarks) + ncols - 1) // ncols + fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(12, 4.4 * nrows)) + axes = axes.flatten() if hasattr(axes, "flatten") else [axes] + + for idx, benchmark in enumerate(benchmarks): + ax = axes[idx] + sub = agg[agg["benchmark"] == benchmark] + for (backend, dataset), bdf in sub.groupby(["backend", "dataset"]): + bdf = bdf.sort_values("n_polsites") + label = ( + backend if sub["dataset"].nunique() == 1 else f"{backend} ({dataset})" + ) + ax.plot( + bdf["n_polsites"], bdf["mean_s"], marker="o", linewidth=1.7, label=label + ) + + ax.set_title(benchmark) + ax.set_xscale("log", base=2) + if log_y: + ax.set_yscale("log") + ax.set_xlabel("n_polsites") + ax.set_ylabel("runtime [s]") + ax.grid(alpha=0.3) + ax.legend(fontsize=8) + + for idx in range(len(benchmarks), len(axes)): + fig.delaxes(axes[idx]) + + fig.suptitle("CPPE Benchmark Runtime Overview", fontsize=14) + _save_fig(fig, out_dir / "overview_runtime", formats, dpi) + plt.close(fig) + + +def _plot_overview_speedup(speedup: pd.DataFrame, out_dir: Path, formats, dpi: int): + needed = {"benchmark", "dataset", "n_polsites", "python_over_cpp"} + if not needed.issubset(speedup.columns): + return + + benchmarks = sorted(speedup["benchmark"].unique()) + ncols = 2 + nrows = (len(benchmarks) + ncols - 1) // ncols + fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(12, 4.2 * nrows)) + axes = axes.flatten() if hasattr(axes, "flatten") else [axes] + + for idx, benchmark in enumerate(benchmarks): + ax = axes[idx] + sub = speedup[speedup["benchmark"] == benchmark] + for dataset, bdf in sub.groupby("dataset"): + bdf = bdf.sort_values("n_polsites") + label = dataset if sub["dataset"].nunique() > 1 else "python/cpp" + ax.plot( + bdf["n_polsites"], + bdf["python_over_cpp"], + marker="o", + linewidth=1.7, + label=label, + ) + + ax.axhline(1.0, color="black", linestyle="--", linewidth=1.0) + ax.set_title(benchmark) + ax.set_xscale("log", base=2) + ax.set_xlabel("n_polsites") + ax.set_ylabel("python / cpp") + ax.grid(alpha=0.3) + ax.legend(fontsize=8) + + for idx in range(len(benchmarks), len(axes)): + fig.delaxes(axes[idx]) + + fig.suptitle("CPPE Benchmark Speedup Overview", fontsize=14) + _save_fig(fig, out_dir / "overview_speedup", formats, dpi) + plt.close(fig) + + +def _write_report(out_dir: Path, agg: pd.DataFrame, speedup: pd.DataFrame): + def _markdown_table(df: pd.DataFrame) -> str: + headers = list(df.columns) + sep = ["---"] * len(headers) + rows = ["| " + " | ".join(headers) + " |", "| " + " | ".join(sep) + " |"] + for _, row in df.iterrows(): + cells = [str(row[h]) for h in headers] + rows.append("| " + " | ".join(cells) + " |") + return "\n".join(rows) + + lines = [ + "# Benchmark Report", + "", + "## Artifacts", + "- `summary_aggregated.csv`: aggregated runtime statistics", + "- `summary_speedup.csv`: backend speedup ratios", + "- `overview_runtime.*`: consolidated runtime overview", + "- `overview_speedup.*`: consolidated speedup overview", + "- `runtime_*.png|svg`: detailed runtime curves (with --detailed)", + "- `speedup_*.png|svg`: detailed speedup curves (with --detailed)", + "", + "## Fastest mean runtime by benchmark/dataset", + "", + ] + + fastest = ( + agg.sort_values("mean_s") + .groupby(["benchmark", "dataset"], as_index=False) + .first()[["benchmark", "dataset", "backend", "n_polsites", "mean_s"]] + ) + lines.append(_markdown_table(fastest)) + + if {"cpp", "python", "python_over_cpp"}.issubset(speedup.columns): + lines.extend( + [ + "", + "## Median Python/C++ ratio by benchmark", + "", + ] + ) + ratio = ( + speedup.groupby(["benchmark", "dataset"], as_index=False)["python_over_cpp"] + .median() + .rename(columns={"python_over_cpp": "median_python_over_cpp"}) + ) + lines.append(_markdown_table(ratio)) + + (out_dir / "REPORT.md").write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def main(): + args = _parse_args() + in_path = Path(args.input) + out_dir = Path(args.output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + if not args.detailed: + _cleanup_detailed_outputs(out_dir) + + df = pd.read_csv(in_path) + required = {"benchmark", "dataset", "backend", "n_polsites", "mean_s"} + missing = required.difference(df.columns) + if missing: + raise ValueError(f"Input CSV is missing required columns: {sorted(missing)}") + + df["n_polsites"] = pd.to_numeric(df["n_polsites"]) + df["mean_s"] = pd.to_numeric(df["mean_s"]) + + agg = _aggregate(df) + speedup = _speedup_table(agg) + + agg.to_csv(out_dir / "summary_aggregated.csv", index=False) + speedup.to_csv(out_dir / "summary_speedup.csv", index=False) + + _plot_overview_runtime( + agg, + out_dir, + formats=args.formats, + dpi=args.dpi, + log_y=not args.no_log_y, + ) + _plot_overview_speedup(speedup, out_dir, formats=args.formats, dpi=args.dpi) + + if args.detailed: + _plot_runtime_curves( + agg, + out_dir, + formats=args.formats, + dpi=args.dpi, + log_y=not args.no_log_y, + ) + _plot_speedups(speedup, out_dir, formats=args.formats, dpi=args.dpi) + _write_report(out_dir, agg, speedup) + + print(f"Wrote benchmark plots and summaries to: {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/synthetic_systems.py b/benchmarks/synthetic_systems.py new file mode 100644 index 00000000..2705e717 --- /dev/null +++ b/benchmarks/synthetic_systems.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import cppe + +from tests.cache import cache + + +def _clone_potential_with_shift(potential, index, shift): + px = potential.x + shift[0] + py = potential.y + shift[1] + pz = potential.z + shift[2] + cloned = cppe.Potential(px, py, pz, potential.element, index) + + for multipole in potential.multipoles: + m = cppe.Multipole(multipole.k) + for value in multipole.values: + m.add_value(float(value)) + cloned.add_multipole(m) + + if potential.is_polarizable: + cloned.set_polarizability( + cppe.Polarizability(list(potential.polarizability.values)) + ) + + return cloned + + +def replicated_waterbox_system(repeats=(4, 4, 4), spacing_bohr=18.0): + base_potentials = cppe.cpp.PotfileReader("tests/potfiles/pna_6w.pot").read() + base_molecule = cache.molecule["pna"] + + nx, ny, nz = repeats + all_potentials = [] + all_atoms = cppe.Molecule() + + base_n_sites = len(base_potentials) + for ix in range(nx): + for iy in range(ny): + for iz in range(nz): + shift = (ix * spacing_bohr, iy * spacing_bohr, iz * spacing_bohr) + pot_offset = len(all_potentials) + + for i, pot in enumerate(base_potentials): + all_potentials.append( + _clone_potential_with_shift(pot, pot_offset + i, shift) + ) + + for i, pot in enumerate(base_potentials): + for excl in pot.exclusions: + all_potentials[pot_offset + i].add_exclusion(pot_offset + excl) + + for atom in base_molecule: + all_atoms.append( + cppe.Atom( + atom.atomic_number, + atom.x + shift[0], + atom.y + shift[1], + atom.z + shift[2], + ) + ) + + n_polsites = len(cppe.get_polarizable_sites(all_potentials)) + label = f"waterbox_{nx}x{ny}x{nz}_{n_polsites}sites" + return { + "label": label, + "molecule": all_atoms, + "potentials": all_potentials, + "n_polsites": n_polsites, + "base_sites": base_n_sites, + } diff --git a/benchmarks/test_bmatrix_apply.py b/benchmarks/test_bmatrix_apply.py new file mode 100644 index 00000000..1348059e --- /dev/null +++ b/benchmarks/test_bmatrix_apply.py @@ -0,0 +1,47 @@ +import numpy as np +import pytest + +import cppe + +from .utils import benchmark_settings, time_callable + + +_LARGE_POTFILE = "tests/potfiles/loprop_solvated_20.pot" +_NSITES = [32, 64, 128, 256] + + +def _subset_polsites(n_polsites): + potentials = cppe.cpp.PotfileReader(_LARGE_POTFILE).read() + polsites = cppe.get_polarizable_sites(potentials) + if len(polsites) < n_polsites: + raise RuntimeError(f"Need at least {n_polsites} polarizable sites.") + return polsites[:n_polsites] + + +@pytest.mark.benchmark +@pytest.mark.parametrize("backend", ["cpp", "python"]) +@pytest.mark.parametrize("n_polsites", _NSITES) +def test_bmatrix_apply_scaling(benchmark_recorder, backend, n_polsites): + cppe.set_backend(backend) + polsites = _subset_polsites(n_polsites) + options = { + "potfile": _LARGE_POTFILE, + "summation_induced_fields": "direct", + "tree_ncrit": 64, + "tree_expansion_order": 5, + } + bmatrix = cppe.BMatrix(polsites, options) + vec = np.linspace(1.0, 2.0, 3 * n_polsites) + + repeat, warmup = benchmark_settings() + timings = time_callable(lambda: bmatrix.apply(vec), repeat=repeat, warmup=warmup) + + benchmark_recorder( + benchmark="bmatrix_apply", + backend=backend, + dataset="loprop_solvated_20.pot", + n_polsites=n_polsites, + **timings, + ) + + cppe.set_backend("cpp") diff --git a/benchmarks/test_induced_moments.py b/benchmarks/test_induced_moments.py new file mode 100644 index 00000000..8621373b --- /dev/null +++ b/benchmarks/test_induced_moments.py @@ -0,0 +1,49 @@ +import numpy as np +import pytest + +import cppe + +from .utils import benchmark_settings, time_callable + + +_LARGE_POTFILE = "tests/potfiles/loprop_solvated_20.pot" +_NSITES = [32, 64, 128, 256] + + +def _subset_polsites(n_polsites): + potentials = cppe.cpp.PotfileReader(_LARGE_POTFILE).read() + polsites = cppe.get_polarizable_sites(potentials) + if len(polsites) < n_polsites: + raise RuntimeError(f"Need at least {n_polsites} polarizable sites.") + return polsites[:n_polsites] + + +@pytest.mark.benchmark +@pytest.mark.parametrize("backend", ["cpp", "python"]) +@pytest.mark.parametrize("n_polsites", _NSITES) +def test_induced_moments_solver_scaling(benchmark_recorder, backend, n_polsites): + cppe.set_backend(backend) + polsites = _subset_polsites(n_polsites) + options = { + "potfile": _LARGE_POTFILE, + "summation_induced_fields": "direct", + "induced_thresh": 1e-8, + "maxiter": 50, + } + solver = cppe.InducedMoments(polsites, options) + rhs = np.linspace(0.1, 0.8, 3 * n_polsites) + + repeat, warmup = benchmark_settings() + timings = time_callable( + lambda: solver.compute(rhs, True), repeat=repeat, warmup=warmup + ) + + benchmark_recorder( + benchmark="induced_moments_solver", + backend=backend, + dataset="loprop_solvated_20.pot", + n_polsites=n_polsites, + **timings, + ) + + cppe.set_backend("cpp") diff --git a/benchmarks/test_multipole_expansion.py b/benchmarks/test_multipole_expansion.py new file mode 100644 index 00000000..11772011 --- /dev/null +++ b/benchmarks/test_multipole_expansion.py @@ -0,0 +1,42 @@ +import pytest + +import cppe + +from tests.cache import cache + +from .utils import benchmark_settings, time_callable + + +_LARGE_POTFILE = "tests/potfiles/loprop_solvated_20.pot" +_NSITES = [32, 64, 128, 256] + + +def _subset_potentials(n_polsites): + potentials = cppe.cpp.PotfileReader(_LARGE_POTFILE).read() + polsites = cppe.get_polarizable_sites(potentials) + if len(polsites) < n_polsites: + raise RuntimeError(f"Need at least {n_polsites} polarizable sites.") + return polsites[:n_polsites] + + +@pytest.mark.benchmark +@pytest.mark.parametrize("backend", ["cpp", "python"]) +@pytest.mark.parametrize("n_polsites", _NSITES) +def test_multipole_interaction_scaling(benchmark_recorder, backend, n_polsites): + cppe.set_backend(backend) + potentials = _subset_potentials(n_polsites) + molecule = cache.molecule["nilered"] + mexp = cppe.MultipoleExpansion(molecule, potentials) + + repeat, warmup = benchmark_settings() + timings = time_callable(mexp.interaction_energy, repeat=repeat, warmup=warmup) + + benchmark_recorder( + benchmark="multipole_interaction_energy", + backend=backend, + dataset="loprop_solvated_20.pot", + n_polsites=n_polsites, + **timings, + ) + + cppe.set_backend("cpp") diff --git a/benchmarks/test_multipole_fields.py b/benchmarks/test_multipole_fields.py new file mode 100644 index 00000000..b029d859 --- /dev/null +++ b/benchmarks/test_multipole_fields.py @@ -0,0 +1,45 @@ +import pytest + +import cppe + +from .utils import benchmark_settings, time_callable + + +_LARGE_POTFILE = "tests/potfiles/loprop_solvated_20.pot" +_NSITES = [32, 64, 128, 256] + + +def _subset_potentials(n_polsites): + potentials = cppe.cpp.PotfileReader(_LARGE_POTFILE).read() + polsites = cppe.get_polarizable_sites(potentials) + if len(polsites) < n_polsites: + raise RuntimeError(f"Need at least {n_polsites} polarizable sites.") + return polsites[:n_polsites] + + +@pytest.mark.benchmark +@pytest.mark.parametrize("backend", ["cpp", "python"]) +@pytest.mark.parametrize("n_polsites", _NSITES) +def test_multipole_fields_scaling(benchmark_recorder, backend, n_polsites): + cppe.set_backend(backend) + potentials = _subset_potentials(n_polsites) + options = { + "potfile": _LARGE_POTFILE, + "summation_induced_fields": "direct", + "tree_ncrit": 64, + "tree_expansion_order": 5, + } + fields = cppe.MultipoleFields(potentials, options) + + repeat, warmup = benchmark_settings() + timings = time_callable(fields.compute, repeat=repeat, warmup=warmup) + + benchmark_recorder( + benchmark="multipole_fields", + backend=backend, + dataset="loprop_solvated_20.pot", + n_polsites=n_polsites, + **timings, + ) + + cppe.set_backend("cpp") diff --git a/benchmarks/test_nuclear_fields.py b/benchmarks/test_nuclear_fields.py new file mode 100644 index 00000000..1a7e04e7 --- /dev/null +++ b/benchmarks/test_nuclear_fields.py @@ -0,0 +1,42 @@ +import pytest + +import cppe + +from tests.cache import cache + +from .utils import benchmark_settings, time_callable + + +_LARGE_POTFILE = "tests/potfiles/loprop_solvated_20.pot" +_NSITES = [32, 64, 128, 256] + + +def _subset_potentials(n_polsites): + potentials = cppe.cpp.PotfileReader(_LARGE_POTFILE).read() + polsites = cppe.get_polarizable_sites(potentials) + if len(polsites) < n_polsites: + raise RuntimeError(f"Need at least {n_polsites} polarizable sites.") + return polsites[:n_polsites] + + +@pytest.mark.benchmark +@pytest.mark.parametrize("backend", ["cpp", "python"]) +@pytest.mark.parametrize("n_polsites", _NSITES) +def test_nuclear_fields_scaling(benchmark_recorder, backend, n_polsites): + cppe.set_backend(backend) + potentials = _subset_potentials(n_polsites) + molecule = cache.molecule["nilered"] + fields = cppe.NuclearFields(molecule, potentials) + + repeat, warmup = benchmark_settings() + timings = time_callable(fields.compute, repeat=repeat, warmup=warmup) + + benchmark_recorder( + benchmark="nuclear_fields", + backend=backend, + dataset="loprop_solvated_20.pot", + n_polsites=n_polsites, + **timings, + ) + + cppe.set_backend("cpp") diff --git a/benchmarks/test_potfile_reader.py b/benchmarks/test_potfile_reader.py new file mode 100644 index 00000000..559acbf4 --- /dev/null +++ b/benchmarks/test_potfile_reader.py @@ -0,0 +1,38 @@ +from pathlib import Path + +import cppe +import pytest + +from .utils import benchmark_settings, time_callable + + +DATASETS = [ + "pna_6w.pot", + "pna_6w_isopol.pot", + "loprop_solvated_20.pot", +] + + +@pytest.mark.benchmark +@pytest.mark.parametrize("backend", ["cpp", "python"]) +@pytest.mark.parametrize("dataset", DATASETS) +def test_potfile_reader(benchmark_recorder, backend, dataset): + cppe.set_backend(backend) + potfile = Path("tests") / "potfiles" / dataset + + def _run(): + return cppe.PotfileReader(str(potfile)).read() + + repeat, warmup = benchmark_settings() + timings = time_callable(_run, repeat=repeat, warmup=warmup) + n_polsites = len(cppe.get_polarizable_sites(_run())) + + benchmark_recorder( + benchmark="potfile_read", + backend=backend, + dataset=dataset, + n_polsites=n_polsites, + **timings, + ) + + cppe.set_backend("cpp") diff --git a/benchmarks/test_state_cycle.py b/benchmarks/test_state_cycle.py new file mode 100644 index 00000000..82e85918 --- /dev/null +++ b/benchmarks/test_state_cycle.py @@ -0,0 +1,38 @@ +import numpy as np +import pytest + +import cppe + +from tests.cache import cache + +from .utils import benchmark_settings, time_callable + + +@pytest.mark.benchmark +@pytest.mark.parametrize("backend", ["cpp", "python"]) +@pytest.mark.parametrize("dataset", ["pna_6w.pot", "pna_6w_isopol.pot"]) +def test_state_cycle(benchmark_recorder, backend, dataset): + cppe.set_backend(backend) + options = {"potfile": f"tests/potfiles/{dataset}", "induced_thresh": 1e-10} + molecule = cache.molecule["pna"] + + def _run(): + state = cppe.CppeState(options, molecule, lambda _: None) + state.calculate_static_energies_and_fields() + zeros = np.zeros_like(state.static_fields) + state.update_induced_moments(zeros, False) + return state + + repeat, warmup = benchmark_settings() + timings = time_callable(_run, repeat=repeat, warmup=warmup) + n_polsites = _run().get_polarizable_site_number() + + benchmark_recorder( + benchmark="state_cycle", + backend=backend, + dataset=dataset, + n_polsites=n_polsites, + **timings, + ) + + cppe.set_backend("cpp") diff --git a/benchmarks/test_synthetic_large.py b/benchmarks/test_synthetic_large.py new file mode 100644 index 00000000..46757d87 --- /dev/null +++ b/benchmarks/test_synthetic_large.py @@ -0,0 +1,63 @@ +import numpy as np +import pytest + +import cppe + +from .synthetic_systems import replicated_waterbox_system +from .utils import benchmark_settings, time_callable + + +_REPEATS = [ + (4, 4, 4), + (6, 6, 6), +] + + +@pytest.mark.benchmark_large +@pytest.mark.parametrize("backend", ["cpp", "python"]) +@pytest.mark.parametrize("repeats", _REPEATS) +def test_synthetic_bmatrix_apply(benchmark_recorder, backend, repeats): + system = replicated_waterbox_system(repeats=repeats) + options = { + "potfile": "tests/potfiles/pna_6w.pot", + "summation_induced_fields": "direct", + } + polsites = cppe.get_polarizable_sites(system["potentials"]) + vec = np.linspace(0.1, 0.7, 3 * len(polsites)) + + cppe.set_backend(backend) + bmat = cppe.BMatrix(polsites, options) + repeat, warmup = benchmark_settings() + timings = time_callable(lambda: bmat.apply(vec), repeat=repeat, warmup=warmup) + + benchmark_recorder( + benchmark="synthetic_bmatrix_apply", + backend=backend, + dataset=system["label"], + n_polsites=system["n_polsites"], + **timings, + ) + + cppe.set_backend("cpp") + + +@pytest.mark.benchmark_large +@pytest.mark.parametrize("backend", ["cpp", "python"]) +@pytest.mark.parametrize("repeats", _REPEATS) +def test_synthetic_nuclear_fields(benchmark_recorder, backend, repeats): + system = replicated_waterbox_system(repeats=repeats) + + cppe.set_backend(backend) + nf = cppe.NuclearFields(system["molecule"], system["potentials"]) + repeat, warmup = benchmark_settings() + timings = time_callable(nf.compute, repeat=repeat, warmup=warmup) + + benchmark_recorder( + benchmark="synthetic_nuclear_fields", + backend=backend, + dataset=system["label"], + n_polsites=system["n_polsites"], + **timings, + ) + + cppe.set_backend("cpp") diff --git a/benchmarks/utils.py b/benchmarks/utils.py new file mode 100644 index 00000000..66fe43d6 --- /dev/null +++ b/benchmarks/utils.py @@ -0,0 +1,31 @@ +import os +import statistics +import time + + +def benchmark_settings(): + repeat = int(os.environ.get("CPPE_BENCH_REPEAT", "3")) + warmup = int(os.environ.get("CPPE_BENCH_WARMUP", "1")) + return repeat, warmup + + +def time_callable(func, *, repeat, warmup): + for _ in range(warmup): + func() + + samples = [] + for _ in range(repeat): + t0 = time.perf_counter() + func() + samples.append(time.perf_counter() - t0) + + mean_s = statistics.mean(samples) + std_s = statistics.pstdev(samples) if len(samples) > 1 else 0.0 + return { + "repeat": repeat, + "warmup": warmup, + "mean_s": mean_s, + "std_s": std_s, + "min_s": min(samples), + "max_s": max(samples), + } diff --git a/pyproject.toml b/pyproject.toml index 787f8faa..5b3287b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,9 @@ dynamic = ["version"] description = "C++ and Python library for Polarizable Embedding" readme = "README.md" requires-python = ">=3.8" +dependencies = [ + "numba", +] authors = [ { name = "Maximilian Scheurer", email = "maximilian.scheurer@iwr.uni-heidelberg.de" }, { name = "Peter Reinholdt" }, diff --git a/src/cppe/__init__.py b/src/cppe/__init__.py index 95dbce0d..ac847ff6 100644 --- a/src/cppe/__init__.py +++ b/src/cppe/__init__.py @@ -1,6 +1,64 @@ -from .pycppe import * -from .pycppe import __build_type__, __parallel__ +from . import cpp +from .cpp import * +from .cpp import __build_type__, __parallel__ from .pycppe.tensors import * +from .backend import get_backend, set_backend +from . import potfile as _potfile + + +def PotfileReader(*args, **kwargs): + if get_backend() == "python": + return _potfile.PotfileReader(*args, **kwargs) + return cpp.PotfileReader(*args, **kwargs) + + +def NuclearFields(*args, **kwargs): + if get_backend() == "python": + from .fields import NuclearFields as _PythonNuclearFields + + return _PythonNuclearFields(*args, **kwargs) + return cpp.NuclearFields(*args, **kwargs) + + +def BMatrix(*args, **kwargs): + if get_backend() == "python": + from .solver import BMatrix as _PythonBMatrix + + return _PythonBMatrix(*args, **kwargs) + return cpp.BMatrix(*args, **kwargs) + + +def InducedMoments(*args, **kwargs): + if get_backend() == "python": + from .solver import InducedMoments as _PythonInducedMoments + + return _PythonInducedMoments(*args, **kwargs) + return cpp.InducedMoments(*args, **kwargs) + + +def MultipoleExpansion(*args, **kwargs): + if get_backend() == "python": + from .multipole import MultipoleExpansion as _PythonMultipoleExpansion + + return _PythonMultipoleExpansion(*args, **kwargs) + return cpp.MultipoleExpansion(*args, **kwargs) + + +def MultipoleFields(*args, **kwargs): + if get_backend() == "python": + from .multipole_fields import MultipoleFields as _PythonMultipoleFields + + return _PythonMultipoleFields(*args, **kwargs) + return cpp.MultipoleFields(*args, **kwargs) + + +def CppeState(*args, **kwargs): + if get_backend() == "python": + from .state import CppeState as _PythonCppeState + + return _PythonCppeState(*args, **kwargs) + return cpp.CppeState(*args, **kwargs) + try: from ._version import version as __version__ @@ -15,5 +73,15 @@ all = [ "__build_type__", "__parallel__", + "cpp", + "get_backend", + "set_backend", + "PotfileReader", + "NuclearFields", + "BMatrix", + "InducedMoments", + "MultipoleExpansion", + "MultipoleFields", + "CppeState", "__version__", ] diff --git a/src/cppe/backend.py b/src/cppe/backend.py new file mode 100644 index 00000000..ae0932ef --- /dev/null +++ b/src/cppe/backend.py @@ -0,0 +1,35 @@ +import os + + +_ALLOWED_BACKENDS = {"cpp", "python"} +_DEFAULT_BACKEND = "cpp" + + +def _normalize_backend(name: str) -> str: + normalized = name.strip().lower() + if normalized not in _ALLOWED_BACKENDS: + allowed = ", ".join(sorted(_ALLOWED_BACKENDS)) + raise ValueError(f"Invalid backend '{name}'. Valid backends are: {allowed}.") + return normalized + + +def _backend_from_env() -> str: + raw = os.environ.get("CPPE_BACKEND", "") + if not raw: + return _DEFAULT_BACKEND + try: + return _normalize_backend(raw) + except ValueError: + return _DEFAULT_BACKEND + + +_backend = _backend_from_env() + + +def get_backend() -> str: + return _backend + + +def set_backend(name: str) -> None: + global _backend + _backend = _normalize_backend(name) diff --git a/src/cppe/cpp.py b/src/cppe/cpp.py new file mode 100644 index 00000000..6aab7914 --- /dev/null +++ b/src/cppe/cpp.py @@ -0,0 +1,3 @@ +from .pycppe import * +from .pycppe import __build_type__, __parallel__ +from .pycppe import tensors diff --git a/src/cppe/fields.py b/src/cppe/fields.py new file mode 100644 index 00000000..4d97123a --- /dev/null +++ b/src/cppe/fields.py @@ -0,0 +1,118 @@ +from __future__ import annotations + +import numpy as np +from numba import njit, prange + + +@njit(cache=True, parallel=True) +def _nuclear_fields_kernel(site_positions, atom_positions, atom_charges): + n_sites = site_positions.shape[0] + n_atoms = atom_positions.shape[0] + fields = np.zeros((n_sites, 3), dtype=np.float64) + + for i in prange(n_sites): + sx = site_positions[i, 0] + sy = site_positions[i, 1] + sz = site_positions[i, 2] + fx = 0.0 + fy = 0.0 + fz = 0.0 + for j in range(n_atoms): + dx = sx - atom_positions[j, 0] + dy = sy - atom_positions[j, 1] + dz = sz - atom_positions[j, 2] + r2 = dx * dx + dy * dy + dz * dz + inv_r = 1.0 / np.sqrt(r2) + inv_r3 = inv_r / r2 + q = atom_charges[j] + fx += q * dx * inv_r3 + fy += q * dy * inv_r3 + fz += q * dz * inv_r3 + fields[i, 0] = fx + fields[i, 1] = fy + fields[i, 2] = fz + + return fields.reshape(3 * n_sites) + + +@njit(cache=True, parallel=True) +def _nuclear_gradient_kernel(site_positions, atom_positions, atom_charges): + n_sites = site_positions.shape[0] + n_atoms = atom_positions.shape[0] + grad = np.zeros((3 * n_atoms, 3 * n_sites), dtype=np.float64) + + for i in prange(n_sites): + sx = site_positions[i, 0] + sy = site_positions[i, 1] + sz = site_positions[i, 2] + col = 3 * i + + for ai in range(n_atoms): + dx = sx - atom_positions[ai, 0] + dy = sy - atom_positions[ai, 1] + dz = sz - atom_positions[ai, 2] + + r2 = dx * dx + dy * dy + dz * dz + inv_r = 1.0 / np.sqrt(r2) + inv_r5 = (inv_r * inv_r * inv_r) / r2 + q = atom_charges[ai] + + txx = (3.0 * dx * dx - r2) * inv_r5 + txy = 3.0 * dx * dy * inv_r5 + txz = 3.0 * dx * dz * inv_r5 + tyy = (3.0 * dy * dy - r2) * inv_r5 + tyz = 3.0 * dy * dz * inv_r5 + tzz = (3.0 * dz * dz - r2) * inv_r5 + + row = 3 * ai + grad[row + 0, col + 0] = q * txx + grad[row + 0, col + 1] = q * txy + grad[row + 0, col + 2] = q * txz + + grad[row + 1, col + 0] = q * txy + grad[row + 1, col + 1] = q * tyy + grad[row + 1, col + 2] = q * tyz + + grad[row + 2, col + 0] = q * txz + grad[row + 2, col + 1] = q * tyz + grad[row + 2, col + 2] = q * tzz + + return grad + + +class NuclearFields: + def __init__(self, molecule, potentials): + self._molecule = molecule + self._polsites = [p for p in potentials if p.is_polarizable] + + @staticmethod + def _as_atom_arrays(molecule): + n_atoms = len(molecule) + atom_positions = np.empty((n_atoms, 3), dtype=np.float64) + atom_charges = np.empty(n_atoms, dtype=np.float64) + for i, atom in enumerate(molecule): + atom_positions[i, 0] = atom.x + atom_positions[i, 1] = atom.y + atom_positions[i, 2] = atom.z + atom_charges[i] = atom.charge + return atom_positions, atom_charges + + @staticmethod + def _as_site_positions(polsites): + n_sites = len(polsites) + site_positions = np.empty((n_sites, 3), dtype=np.float64) + for i, site in enumerate(polsites): + site_positions[i, 0] = site.x + site_positions[i, 1] = site.y + site_positions[i, 2] = site.z + return site_positions + + def compute(self): + site_positions = self._as_site_positions(self._polsites) + atom_positions, atom_charges = self._as_atom_arrays(self._molecule) + return _nuclear_fields_kernel(site_positions, atom_positions, atom_charges) + + def nuclear_gradient(self): + site_positions = self._as_site_positions(self._polsites) + atom_positions, atom_charges = self._as_atom_arrays(self._molecule) + return _nuclear_gradient_kernel(site_positions, atom_positions, atom_charges) diff --git a/src/cppe/multipole.py b/src/cppe/multipole.py new file mode 100644 index 00000000..c03f3002 --- /dev/null +++ b/src/cppe/multipole.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +import numpy as np +from numba import njit, prange + + +@njit(cache=True, parallel=True) +def _interaction_energy_kernel( + site_positions, + atom_positions, + atom_charges, + charges, + dipoles, + quadrupoles, + use_dipoles, + use_quadrupoles, +): + n_sites = site_positions.shape[0] + n_atoms = atom_positions.shape[0] + site_energies = np.zeros(n_sites, dtype=np.float64) + + for i in prange(n_sites): + sx = site_positions[i, 0] + sy = site_positions[i, 1] + sz = site_positions[i, 2] + + q0 = charges[i] + mux = dipoles[i, 0] + muy = dipoles[i, 1] + muz = dipoles[i, 2] + + qxx = quadrupoles[i, 0] + qxy = quadrupoles[i, 1] + qxz = quadrupoles[i, 2] + qyy = quadrupoles[i, 3] + qyz = quadrupoles[i, 4] + qzz = quadrupoles[i, 5] + + acc = 0.0 + for a in range(n_atoms): + dx = atom_positions[a, 0] - sx + dy = atom_positions[a, 1] - sy + dz = atom_positions[a, 2] - sz + + r2 = dx * dx + dy * dy + dz * dz + inv_r = 1.0 / np.sqrt(r2) + inv_r3 = inv_r / r2 + + e = q0 * inv_r + if use_dipoles: + e += (mux * dx + muy * dy + muz * dz) * inv_r3 + + if use_quadrupoles: + inv_r5 = inv_r3 / r2 + txx = (3.0 * dx * dx - r2) * inv_r5 + txy = 3.0 * dx * dy * inv_r5 + txz = 3.0 * dx * dz * inv_r5 + tyy = (3.0 * dy * dy - r2) * inv_r5 + tyz = 3.0 * dy * dz * inv_r5 + tzz = (3.0 * dz * dz - r2) * inv_r5 + e += ( + 0.5 * qxx * txx + + qxy * txy + + qxz * txz + + 0.5 * qyy * tyy + + qyz * tyz + + 0.5 * qzz * tzz + ) + + acc += atom_charges[a] * e + site_energies[i] = acc + + return site_energies.sum() + + +@njit(cache=True, parallel=True) +def _nuclear_gradient_kernel( + site_positions, + atom_positions, + atom_charges, + charges, + dipoles, + quadrupoles, + has_dipole, + has_quadrupole, +): + n_atoms = atom_positions.shape[0] + n_sites = site_positions.shape[0] + grad = np.zeros((n_atoms, 3), dtype=np.float64) + + for ai in prange(n_atoms): + ax = atom_positions[ai, 0] + ay = atom_positions[ai, 1] + az = atom_positions[ai, 2] + q_atom = atom_charges[ai] + + gx = 0.0 + gy = 0.0 + gz = 0.0 + + for j in range(n_sites): + dx = ax - site_positions[j, 0] + dy = ay - site_positions[j, 1] + dz = az - site_positions[j, 2] + + r2 = dx * dx + dy * dy + dz * dz + inv_r = 1.0 / np.sqrt(r2) + inv_r3 = inv_r / r2 + + q = charges[j] + gx += q * dx * inv_r3 + gy += q * dy * inv_r3 + gz += q * dz * inv_r3 + + if has_dipole[j]: + mx = dipoles[j, 0] + my = dipoles[j, 1] + mz = dipoles[j, 2] + mdotr = mx * dx + my * dy + mz * dz + inv_r5 = inv_r3 / r2 + gx += 3.0 * dx * mdotr * inv_r5 - mx * inv_r3 + gy += 3.0 * dy * mdotr * inv_r5 - my * inv_r3 + gz += 3.0 * dz * mdotr * inv_r5 - mz * inv_r3 + + if has_quadrupole[j]: + qxx = quadrupoles[j, 0] + qxy = quadrupoles[j, 1] + qxz = quadrupoles[j, 2] + qyy = quadrupoles[j, 3] + qyz = quadrupoles[j, 4] + qzz = quadrupoles[j, 5] + qrx = qxx * dx + qxy * dy + qxz * dz + qry = qxy * dx + qyy * dy + qyz * dz + qrz = qxz * dx + qyz * dy + qzz * dz + rqrr = dx * qrx + dy * qry + dz * qrz + inv_r5 = inv_r3 / r2 + inv_r7 = inv_r5 / r2 + gx += 7.5 * dx * rqrr * inv_r7 - 3.0 * qrx * inv_r5 + gy += 7.5 * dy * rqrr * inv_r7 - 3.0 * qry * inv_r5 + gz += 7.5 * dz * rqrr * inv_r7 - 3.0 * qrz * inv_r5 + + grad[ai, 0] = -q_atom * gx + grad[ai, 1] = -q_atom * gy + grad[ai, 2] = -q_atom * gz + + return grad + + +class MultipoleExpansion: + def __init__(self, molecule, potentials): + self._molecule = molecule + self._potentials = potentials + + n_atoms = len(molecule) + n_sites = len(potentials) + self._atom_positions = np.empty((n_atoms, 3), dtype=np.float64) + self._atom_charges = np.empty(n_atoms, dtype=np.float64) + for i, atom in enumerate(molecule): + self._atom_positions[i, 0] = atom.x + self._atom_positions[i, 1] = atom.y + self._atom_positions[i, 2] = atom.z + self._atom_charges[i] = atom.charge + + self._site_positions = np.empty((n_sites, 3), dtype=np.float64) + self._charges = np.zeros(n_sites, dtype=np.float64) + self._dipoles = np.zeros((n_sites, 3), dtype=np.float64) + self._quadrupoles = np.zeros((n_sites, 6), dtype=np.float64) + + self._max_order = 0 + self._has_dipole = np.zeros(n_sites, dtype=np.int8) + self._has_quadrupole = np.zeros(n_sites, dtype=np.int8) + for i, site in enumerate(potentials): + self._site_positions[i, 0] = site.x + self._site_positions[i, 1] = site.y + self._site_positions[i, 2] = site.z + for multipole in site.multipoles: + k = multipole.k + if k > self._max_order: + self._max_order = k + values = np.asarray(multipole.values, dtype=np.float64) + if k == 0: + self._charges[i] = values[0] + elif k == 1: + self._has_dipole[i] = 1 + self._dipoles[i, :] = values + elif k == 2: + self._has_quadrupole[i] = 1 + self._quadrupoles[i, :] = values + if self._max_order > 2: + raise NotImplementedError( + "Python MultipoleExpansion currently supports multipoles up to order 2." + ) + + def interaction_energy(self): + return _interaction_energy_kernel( + self._site_positions, + self._atom_positions, + self._atom_charges, + self._charges, + self._dipoles, + self._quadrupoles, + self._max_order > 0, + self._max_order > 1, + ) + + def nuclear_gradient(self): + return _nuclear_gradient_kernel( + self._site_positions, + self._atom_positions, + self._atom_charges, + self._charges, + self._dipoles, + self._quadrupoles, + self._has_dipole, + self._has_quadrupole, + ) diff --git a/src/cppe/multipole_fields.py b/src/cppe/multipole_fields.py new file mode 100644 index 00000000..380bc6c8 --- /dev/null +++ b/src/cppe/multipole_fields.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +import numpy as np +from numba import njit, prange + + +@njit(cache=True, parallel=True) +def _multipole_fields_direct_kernel( + positions, + include_mask, + charges, + dipoles, + quadrupoles, + has_dipole, + has_quadrupole, +): + n_sites = positions.shape[0] + out = np.zeros(3 * n_sites, dtype=np.float64) + + for i in prange(n_sites): + ix = positions[i, 0] + iy = positions[i, 1] + iz = positions[i, 2] + fx = 0.0 + fy = 0.0 + fz = 0.0 + + for j in range(n_sites): + if include_mask[i, j] == 0: + continue + + dx = ix - positions[j, 0] + dy = iy - positions[j, 1] + dz = iz - positions[j, 2] + + r2 = dx * dx + dy * dy + dz * dz + inv_r = 1.0 / np.sqrt(r2) + inv_r3 = inv_r / r2 + + q = charges[j] + fx += q * dx * inv_r3 + fy += q * dy * inv_r3 + fz += q * dz * inv_r3 + + if has_dipole[j]: + mx = dipoles[j, 0] + my = dipoles[j, 1] + mz = dipoles[j, 2] + mdotr = mx * dx + my * dy + mz * dz + inv_r5 = inv_r3 / r2 + fx += 3.0 * dx * mdotr * inv_r5 - mx * inv_r3 + fy += 3.0 * dy * mdotr * inv_r5 - my * inv_r3 + fz += 3.0 * dz * mdotr * inv_r5 - mz * inv_r3 + + if has_quadrupole[j]: + qxx = quadrupoles[j, 0] + qxy = quadrupoles[j, 1] + qxz = quadrupoles[j, 2] + qyy = quadrupoles[j, 3] + qyz = quadrupoles[j, 4] + qzz = quadrupoles[j, 5] + + qrx = qxx * dx + qxy * dy + qxz * dz + qry = qxy * dx + qyy * dy + qyz * dz + qrz = qxz * dx + qyz * dy + qzz * dz + + rqrr = dx * qrx + dy * qry + dz * qrz + + inv_r5 = inv_r3 / r2 + inv_r7 = inv_r5 / r2 + fx += 7.5 * dx * rqrr * inv_r7 - 3.0 * qrx * inv_r5 + fy += 7.5 * dy * rqrr * inv_r7 - 3.0 * qry * inv_r5 + fz += 7.5 * dz * rqrr * inv_r7 - 3.0 * qrz * inv_r5 + + i3 = 3 * i + out[i3 + 0] = fx + out[i3 + 1] = fy + out[i3 + 2] = fz + + return out + + +class MultipoleFields: + def __init__(self, potentials, options): + self._potentials = list(potentials) + self._options = dict(options) + if self._options.get("summation_induced_fields", "direct") != "direct": + raise NotImplementedError( + "Python MultipoleFields currently supports summation_induced_fields='direct' only." + ) + if self._options.get("damp_multipole", False): + raise NotImplementedError( + "Python MultipoleFields currently does not support damp_multipole=True." + ) + + def compute(self): + n_sites = len(self._potentials) + positions = np.empty((n_sites, 3), dtype=np.float64) + include_mask = np.zeros((n_sites, n_sites), dtype=np.int8) + charges = np.zeros(n_sites, dtype=np.float64) + dipoles = np.zeros((n_sites, 3), dtype=np.float64) + quadrupoles = np.zeros((n_sites, 6), dtype=np.float64) + has_dipole = np.zeros(n_sites, dtype=np.int8) + has_quadrupole = np.zeros(n_sites, dtype=np.int8) + + for i, target in enumerate(self._potentials): + positions[i, 0] = target.x + positions[i, 1] = target.y + positions[i, 2] = target.z + + for j, source in enumerate(self._potentials): + if i == j: + continue + if target.excludes_site(source.index): + continue + include_mask[i, j] = 1 + + for multipole in target.multipoles: + values = np.asarray(multipole.values, dtype=np.float64) + if multipole.k == 0: + charges[i] = values[0] + elif multipole.k == 1: + has_dipole[i] = 1 + dipoles[i, 0] = values[0] + dipoles[i, 1] = values[1] + dipoles[i, 2] = values[2] + elif multipole.k == 2: + has_quadrupole[i] = 1 + quadrupoles[i, :] = values + + return _multipole_fields_direct_kernel( + positions, + include_mask, + charges, + dipoles, + quadrupoles, + has_dipole, + has_quadrupole, + ) diff --git a/src/cppe/potfile.py b/src/cppe/potfile.py new file mode 100644 index 00000000..02ebb2bf --- /dev/null +++ b/src/cppe/potfile.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from pathlib import Path + +from .pycppe import Multipole, Polarizability, Potential, multipole_components + + +_ANG2BOHR = 1.8897261246 + + +class PotfileReader: + def __init__(self, potfile_name: str): + self._potfile = Path(potfile_name) + if not self._potfile.is_file(): + raise RuntimeError("Potential file does not exist.") + + def read(self) -> list[Potential]: + with self._potfile.open("r", encoding="utf-8") as handle: + lines = handle.readlines() + + potentials: list[Potential] = [] + num_sites = 0 + i = 0 + + while i < len(lines): + line = lines[i] + if "@COORDINATES" in line: + i += 1 + num_sites = int(lines[i].split()[0]) + i += 1 + unit = lines[i].strip() + if unit == "AA": + conversion = _ANG2BOHR + elif unit == "AU": + conversion = 1.0 + else: + raise RuntimeError("Invalid unit for potential file.") + + for site_idx in range(num_sites): + i += 1 + tokens = lines[i].split() + if len(tokens) < 4: + raise RuntimeError("Invalid coordinate line in potential file.") + element = tokens[0] + x = float(tokens[1]) * conversion + y = float(tokens[2]) * conversion + z = float(tokens[3]) * conversion + potentials.append(Potential(x, y, z, element, site_idx)) + + elif "ORDER" in line: + tokens = line.split() + if len(tokens) == 2: + order = int(tokens[1]) + i += 1 + num_multipoles = int(lines[i].split()[0]) + site_before = -1 + + for n_mul in range(num_multipoles): + i += 1 + tokens = lines[i].split() + site_num = int(tokens[0]) - 1 + + if site_num != site_before + 1: + diff = site_num - site_before + for d in range(1, diff): + mul = Multipole(order) + for _ in range(multipole_components(order)): + mul.add_value(0.0) + potentials[site_before + d].add_multipole(mul) + + mul = Multipole(order) + expected = multipole_components(order) + if len(tokens) < expected + 1: + raise RuntimeError( + "Invalid multipole line in potential file." + ) + for value_token in tokens[1 : expected + 1]: + mul.add_value(float(value_token)) + mul.remove_trace() + potentials[site_num].add_multipole(mul) + site_before = site_num + + if n_mul == num_multipoles - 1 and site_num != (num_sites - 1): + diff = num_sites - site_num + for d in range(1, diff): + missing_mul = Multipole(order) + for _ in range(multipole_components(order)): + missing_mul.add_value(0.0) + potentials[site_num + d].add_multipole(missing_mul) + + elif len(tokens) == 3: + order_1 = int(tokens[1]) + order_2 = int(tokens[2]) + if order_1 != 1 or order_2 != 1: + raise RuntimeError( + "Only dipole-dipole polarizabilities are currently supported." + ) + i += 1 + num_polarizabilities = int(lines[i].split()[0]) + for _ in range(num_polarizabilities): + i += 1 + tokens = lines[i].split() + site_num = int(tokens[0]) - 1 + expected = multipole_components(order_1 + order_2) + if len(tokens) < expected + 1: + raise RuntimeError( + "Invalid polarizability line in potential file." + ) + values = [float(token) for token in tokens[1 : expected + 1]] + if potentials[site_num].is_polarizable: + raise RuntimeError( + "Potential can only have one polarizability." + ) + potentials[site_num].set_polarizability(Polarizability(values)) + + else: + raise RuntimeError("Invalid number in potfile ORDER.") + + elif "EXCLISTS" in line: + i += 1 + num_exclusions = int(lines[i].split()[0]) + for _ in range(num_exclusions): + i += 1 + tokens = lines[i].split() + if not tokens: + continue + site_num = int(tokens[0]) - 1 + for token in tokens[1:]: + excl_site_number = int(token) - 1 + if excl_site_number < 0: + continue + potentials[site_num].add_exclusion(excl_site_number) + + i += 1 + + return potentials diff --git a/src/cppe/solver.py b/src/cppe/solver.py new file mode 100644 index 00000000..cf6ebd8d --- /dev/null +++ b/src/cppe/solver.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +import numpy as np +from numba import njit, prange + + +def _triangle_to_mat(values): + ret = np.zeros((3, 3), dtype=np.float64) + tri = np.triu_indices(3) + ret[tri] = values + ret[(tri[1], tri[0])] = values + return ret + + +@njit(cache=True, parallel=True) +def _induced_fields_direct_kernel(positions, include_mask, induced_moments): + n_sites = positions.shape[0] + out = np.zeros(3 * n_sites, dtype=np.float64) + + for i in prange(n_sites): + ix = positions[i, 0] + iy = positions[i, 1] + iz = positions[i, 2] + fx = 0.0 + fy = 0.0 + fz = 0.0 + for j in range(n_sites): + if include_mask[i, j] == 0: + continue + + dx = positions[j, 0] - ix + dy = positions[j, 1] - iy + dz = positions[j, 2] - iz + + r2 = dx * dx + dy * dy + dz * dz + inv_r = 1.0 / np.sqrt(r2) + inv_r5 = (inv_r * inv_r * inv_r) / r2 + + txx = (3.0 * dx * dx - r2) * inv_r5 + txy = 3.0 * dx * dy * inv_r5 + txz = 3.0 * dx * dz * inv_r5 + tyy = (3.0 * dy * dy - r2) * inv_r5 + tyz = 3.0 * dy * dz * inv_r5 + tzz = (3.0 * dz * dz - r2) * inv_r5 + + j3 = 3 * j + sx = induced_moments[j3 + 0] + sy = induced_moments[j3 + 1] + sz = induced_moments[j3 + 2] + + fx -= txx * sx + txy * sy + txz * sz + fy -= txy * sx + tyy * sy + tyz * sz + fz -= txz * sx + tyz * sy + tzz * sz + + i3 = 3 * i + out[i3 + 0] = fx + out[i3 + 1] = fy + out[i3 + 2] = fz + + return out + + +@njit(cache=True, parallel=True) +def _apply_diagonal_kernel(alpha_inverse, induced_moments): + n_sites = alpha_inverse.shape[0] + out = np.zeros(3 * n_sites, dtype=np.float64) + for i in prange(n_sites): + i3 = 3 * i + sx = induced_moments[i3 + 0] + sy = induced_moments[i3 + 1] + sz = induced_moments[i3 + 2] + a = alpha_inverse[i] + out[i3 + 0] = a[0, 0] * sx + a[0, 1] * sy + a[0, 2] * sz + out[i3 + 1] = a[1, 0] * sx + a[1, 1] * sy + a[1, 2] * sz + out[i3 + 2] = a[2, 0] * sx + a[2, 1] * sy + a[2, 2] * sz + return out + + +@njit(cache=True, parallel=True) +def _apply_diagonal_inverse_kernel(alpha, in_vector): + n_sites = alpha.shape[0] + out = np.zeros(3 * n_sites, dtype=np.float64) + for i in prange(n_sites): + i3 = 3 * i + sx = in_vector[i3 + 0] + sy = in_vector[i3 + 1] + sz = in_vector[i3 + 2] + a = alpha[i] + out[i3 + 0] = a[0, 0] * sx + a[0, 1] * sy + a[0, 2] * sz + out[i3 + 1] = a[1, 0] * sx + a[1, 1] * sy + a[1, 2] * sz + out[i3 + 2] = a[2, 0] * sx + a[2, 1] * sy + a[2, 2] * sz + return out + + +class BMatrix: + def __init__(self, polsites, options): + self._polsites = polsites + self._options = dict(options) + self._n_polsites = len(polsites) + + self._positions = np.empty((self._n_polsites, 3), dtype=np.float64) + self._alpha = np.empty((self._n_polsites, 3, 3), dtype=np.float64) + self._alpha_inverse = np.empty((self._n_polsites, 3, 3), dtype=np.float64) + self._include_mask = np.zeros( + (self._n_polsites, self._n_polsites), dtype=np.int8 + ) + + for i, pot in enumerate(polsites): + self._positions[i, 0] = pot.x + self._positions[i, 1] = pot.y + self._positions[i, 2] = pot.z + + alpha = _triangle_to_mat( + np.asarray(pot.polarizability.values, dtype=np.float64) + ) + self._alpha[i] = alpha + self._alpha_inverse[i] = np.linalg.inv(alpha) + + for j, pot_other in enumerate(polsites): + if i == j: + continue + if pot.excludes_site(pot_other.index): + continue + self._include_mask[i, j] = 1 + + scheme = self._options.get("summation_induced_fields", "direct") + if scheme != "direct": + raise NotImplementedError( + "Python BMatrix currently supports summation_induced_fields='direct' only." + ) + if self._options.get("damp_induced", False): + raise NotImplementedError( + "Python BMatrix currently does not support damp_induced=True." + ) + + def apply(self, induced_moments): + vec = np.asarray(induced_moments, dtype=np.float64) + induced_fields = _induced_fields_direct_kernel( + self._positions, self._include_mask, vec + ) + return induced_fields + _apply_diagonal_kernel(self._alpha_inverse, vec) + + def apply_diagonal_inverse(self, in_vector): + vec = np.asarray(in_vector, dtype=np.float64) + return _apply_diagonal_inverse_kernel(self._alpha, vec) + + def apply_diagonal(self, in_vector): + vec = np.asarray(in_vector, dtype=np.float64) + return _apply_diagonal_kernel(self._alpha_inverse, vec) + + def to_dense_matrix(self): + n = self._n_polsites + dense = np.zeros((3 * n, 3 * n), dtype=np.float64) + for i in range(n): + i3 = 3 * i + dense[i3 : i3 + 3, i3 : i3 + 3] = self._alpha_inverse[i] + xi, yi, zi = self._positions[i] + for j in range(n): + if self._include_mask[i, j] == 0: + continue + j3 = 3 * j + dx = self._positions[j, 0] - xi + dy = self._positions[j, 1] - yi + dz = self._positions[j, 2] - zi + r2 = dx * dx + dy * dy + dz * dz + inv_r = 1.0 / np.sqrt(r2) + inv_r5 = (inv_r * inv_r * inv_r) / r2 + txx = (3.0 * dx * dx - r2) * inv_r5 + txy = 3.0 * dx * dy * inv_r5 + txz = 3.0 * dx * dz * inv_r5 + tyy = (3.0 * dy * dy - r2) * inv_r5 + tyz = 3.0 * dy * dz * inv_r5 + tzz = (3.0 * dz * dz - r2) * inv_r5 + dense[i3 + 0, j3 + 0] = -txx + dense[i3 + 0, j3 + 1] = -txy + dense[i3 + 0, j3 + 2] = -txz + dense[i3 + 1, j3 + 0] = -txy + dense[i3 + 1, j3 + 1] = -tyy + dense[i3 + 1, j3 + 2] = -tyz + dense[i3 + 2, j3 + 0] = -txz + dense[i3 + 2, j3 + 1] = -tyz + dense[i3 + 2, j3 + 2] = -tzz + return dense + + def direct_inverse(self): + return np.linalg.inv(self.to_dense_matrix()) + + +class InducedMoments: + def __init__(self, potentials, options): + self._potentials = list(potentials) + self._polsites = [p for p in self._potentials if p.is_polarizable] + self._options = dict(options) + self._printer = lambda s: None + + def set_print_callback(self, printer): + self._printer = printer + + def compute(self, rhs, guess_or_make_guess, make_guess=None): + rhs_vec = np.asarray(rhs, dtype=np.float64) + if make_guess is None: + make_guess = bool(guess_or_make_guess) + guess = np.zeros_like(rhs_vec) + else: + guess = np.asarray(guess_or_make_guess, dtype=np.float64) + + bmat = BMatrix(self._polsites, self._options) + + if make_guess: + x = bmat.apply_diagonal_inverse(rhs_vec) + else: + x = guess.copy() + + r = rhs_vec - bmat.apply(x) + z = bmat.apply_diagonal_inverse(r) + p = z.copy() + + maxiter = int(self._options.get("maxiter", 50)) + thresh = float(self._options.get("induced_thresh", 1e-8)) + + for k in range(maxiter): + ap = bmat.apply(p) + alpha = float(np.dot(r, z) / np.dot(p, ap)) + x = x + alpha * p + r_next = r - alpha * ap + + rnorm = float(np.linalg.norm(r_next)) + self._printer(f"{k} --- Norm: {rnorm:.12f}") + if rnorm < thresh: + return x + + z_next = bmat.apply_diagonal_inverse(r_next) + beta = float(np.dot(z_next, r_next) / np.dot(z, r)) + p = z_next + beta * p + r = r_next + z = z_next + + raise RuntimeError("Failed to converge induced dipole moments.") diff --git a/src/cppe/state.py b/src/cppe/state.py new file mode 100644 index 00000000..b092a0e8 --- /dev/null +++ b/src/cppe/state.py @@ -0,0 +1,214 @@ +from __future__ import annotations + +import numpy as np + +from . import cpp + + +_DEFAULT_OPTIONS = { + "potfile": "potential.pot", + "iso_pol": False, + "induced_thresh": 1e-8, + "maxiter": 50, + "damp_induced": False, + "damping_factor_induced": 2.1304, + "damp_multipole": False, + "damping_factor_multipole": 2.1304, + "summation_induced_fields": "direct", + "tree_ncrit": 64, + "tree_expansion_order": 5, + "theta": 0.5, + "pe_border": False, + "border_type": "remove", + "border_rmin": 2.2, + "border_nredist": -1, + "border_redist_order": 1, + "border_redist_pol": False, +} + + +def _polarizable_sites(potentials): + return [p for p in potentials if p.is_polarizable] + + +class CppeState: + def __init__(self, options=None, molecule=None, printer=None): + if options is None: + options = {} + if molecule is None: + molecule = cpp.Molecule() + if printer is None: + printer = lambda _: None + + unknown = set(options).difference(_DEFAULT_OPTIONS) + if unknown: + key = sorted(unknown)[0] + raise ValueError(f"Option key '{key}' is invalid.") + + self._options = dict(_DEFAULT_OPTIONS) + self._options.update(options) + self._molecule = molecule + self._printer = printer + + if self._options["pe_border"]: + raise NotImplementedError( + "Python backend currently does not support pe_border manipulation." + ) + + from . import PotfileReader + + self._potentials = list(PotfileReader(self._options["potfile"]).read()) + if self._options["iso_pol"]: + for potential in self._potentials: + if potential.is_polarizable: + potential.polarizability.make_isotropic() + + self._nuc_fields = np.zeros(3 * len(self._potentials), dtype=np.float64) + self._multipole_fields = np.zeros(3 * len(self._potentials), dtype=np.float64) + self._polsites = _polarizable_sites(self._potentials) + self._induced_moments = np.zeros(3 * len(self._polsites), dtype=np.float64) + self._make_guess = True + + self._energies = { + "Electrostatic": {"Electronic": 0.0, "Nuclear": 0.0, "Multipoles": 0.0}, + "Polarization": {"Electronic": 0.0, "Nuclear": 0.0, "Multipoles": 0.0}, + } + + @property + def options(self): + return self._options + + @property + def energies(self): + return self._energies + + @property + def potentials(self): + return self._potentials + + @property + def static_fields(self): + return self._nuc_fields + self._multipole_fields + + @property + def nuclear_fields(self): + return self._nuc_fields + + @property + def multipole_fields(self): + return self._multipole_fields + + @property + def total_energy(self): + ele = self.energies["Electrostatic"] + pol = self.energies["Polarization"] + return sum(ele.values()) + sum(pol.values()) + + @property + def summary_string(self): + return ( + "Polarizable Embedding Summary\n" + f"Electrostatic total: {sum(self.energies['Electrostatic'].values()):.12f}\n" + f"Polarization total: {sum(self.energies['Polarization'].values()):.12f}\n" + f"Total energy: {self.total_energy:.12f}" + ) + + @property + def positions(self): + pos = np.zeros((len(self._potentials), 3), dtype=np.float64) + for i, p in enumerate(self._potentials): + pos[i, 0] = p.x + pos[i, 1] = p.y + pos[i, 2] = p.z + return pos + + @property + def positions_polarizable(self): + pos = np.zeros((len(self._polsites), 3), dtype=np.float64) + for i, p in enumerate(self._polsites): + pos[i, 0] = p.x + pos[i, 1] = p.y + pos[i, 2] = p.z + return pos + + def get_polarizable_site_number(self): + return len(self._polsites) + + def set_potentials(self, potentials): + self._potentials = list(potentials) + self._polsites = _polarizable_sites(self._potentials) + self._induced_moments = np.zeros(3 * len(self._polsites), dtype=np.float64) + self._nuc_fields = np.zeros(3 * len(self._potentials), dtype=np.float64) + self._multipole_fields = np.zeros(3 * len(self._potentials), dtype=np.float64) + self._make_guess = True + + def calculate_static_energies_and_fields(self): + from . import MultipoleExpansion, MultipoleFields, NuclearFields + + multipole_expansion = MultipoleExpansion(self._molecule, self._potentials) + self.energies["Electrostatic"]["Nuclear"] = ( + multipole_expansion.interaction_energy() + ) + self._nuc_fields = NuclearFields(self._molecule, self._potentials).compute() + self._multipole_fields = MultipoleFields( + self._potentials, self._options + ).compute() + + def nuclear_interaction_energy_gradient(self): + from . import MultipoleExpansion + + return MultipoleExpansion(self._molecule, self._potentials).nuclear_gradient() + + def nuclear_field_gradient(self): + from . import NuclearFields + + return NuclearFields(self._molecule, self._potentials).nuclear_gradient() + + def update_induced_moments(self, elec_fields, elec_only=False): + from . import InducedMoments + + elec_fields = np.asarray(elec_fields, dtype=np.float64) + if elec_only: + total_fields = elec_fields + else: + total_fields = elec_fields + self._nuc_fields + self._multipole_fields + + ind = InducedMoments(self._potentials, self._options) + ind.set_print_callback(self._printer) + self._induced_moments = ind.compute( + total_fields, self._induced_moments, self._make_guess + ) + + if self._make_guess: + self._make_guess = False + + epol_elec = -0.5 * float(np.dot(self._induced_moments, elec_fields)) + self.energies["Polarization"]["Electronic"] = epol_elec + if not elec_only: + epol_nuclear = -0.5 * float(np.dot(self._induced_moments, self._nuc_fields)) + epol_multipoles = -0.5 * float( + np.dot(self._induced_moments, self._multipole_fields) + ) + self.energies["Polarization"]["Nuclear"] = epol_nuclear + self.energies["Polarization"]["Multipoles"] = epol_multipoles + + def induced_moments_eef(self): + from . import InducedMoments + + n_pol = len(self._polsites) + ret = np.zeros((3 * n_pol, 3), dtype=np.float64) + fdn = np.zeros((3 * n_pol, 3), dtype=np.float64) + for s in range(n_pol): + l = 3 * s + fdn[l + 0, 0] = 1.0 + fdn[l + 1, 1] = 1.0 + fdn[l + 2, 2] = 1.0 + + ind = InducedMoments(self._potentials, self._options) + ind.set_print_callback(self._printer) + for a in range(3): + ret[:, a] = ind.compute(fdn[:, a], True) + return ret + + def get_induced_moments(self): + return self._induced_moments.tolist() diff --git a/src/python_iface/export_potential.cc b/src/python_iface/export_potential.cc index cbf92b27..d3eb6a23 100644 --- a/src/python_iface/export_potential.cc +++ b/src/python_iface/export_potential.cc @@ -20,9 +20,11 @@ void export_multipole(py::module& m) { // libcppe::Polarizability py::class_ polarizability(m, "Polarizability"); polarizability.def(py::init<>()) + .def(py::init>()) .def_property_readonly("values", &libcppe::Polarizability::get_values_vec) .def_property_readonly("isotropic_value", - &libcppe::Polarizability::get_isotropic_value); + &libcppe::Polarizability::get_isotropic_value) + .def("make_isotropic", &libcppe::Polarizability::make_isotropic); // libcppe::Potential py::class_ pot(m, "Potential", "Potential (Site)"); @@ -31,6 +33,9 @@ void export_multipole(py::module& m) { .def_readwrite("y", &libcppe::Potential::m_y, "y coordinate") .def_readwrite("z", &libcppe::Potential::m_z, "z coordinate") .def_readwrite("element", &libcppe::Potential::m_element, "element") + .def("add_multipole", &libcppe::Potential::add_multipole) + .def("add_exclusion", &libcppe::Potential::add_exclusion) + .def("set_polarizability", &libcppe::Potential::set_polarizability) .def_property_readonly("is_polarizable", &libcppe::Potential::is_polarizable) .def("excludes_site", &libcppe::Potential::excludes_site) .def_property_readonly("exclusions", &libcppe::Potential::get_exclusions) diff --git a/src/python_iface/export_utils.cc b/src/python_iface/export_utils.cc index ffafa2f0..fef96099 100644 --- a/src/python_iface/export_utils.cc +++ b/src/python_iface/export_utils.cc @@ -17,4 +17,4 @@ void export_utils(py::module& m) { py::class_ pot_manip(m, "PotManipulator", "Manipulator for potentials"); pot_manip.def(py::init, libcppe::Molecule>()); -} \ No newline at end of file +} diff --git a/tests/test_backend.py b/tests/test_backend.py new file mode 100644 index 00000000..fc096ee6 --- /dev/null +++ b/tests/test_backend.py @@ -0,0 +1,71 @@ +import importlib +from pathlib import Path + +import pytest + + +def _reload_backend_module(): + import cppe.backend as backend + + return importlib.reload(backend) + + +def test_default_backend_is_cpp(monkeypatch): + monkeypatch.delenv("CPPE_BACKEND", raising=False) + backend = _reload_backend_module() + assert backend.get_backend() == "cpp" + + +def test_backend_from_env(monkeypatch): + monkeypatch.setenv("CPPE_BACKEND", "python") + backend = _reload_backend_module() + assert backend.get_backend() == "python" + + +def test_invalid_backend_env_falls_back_to_cpp(monkeypatch): + monkeypatch.setenv("CPPE_BACKEND", "invalid") + backend = _reload_backend_module() + assert backend.get_backend() == "cpp" + + +def test_set_backend_changes_active_backend(monkeypatch): + monkeypatch.delenv("CPPE_BACKEND", raising=False) + backend = _reload_backend_module() + + backend.set_backend("python") + assert backend.get_backend() == "python" + + backend.set_backend("cpp") + assert backend.get_backend() == "cpp" + + +def test_set_backend_rejects_invalid_name(monkeypatch): + monkeypatch.delenv("CPPE_BACKEND", raising=False) + backend = _reload_backend_module() + + with pytest.raises(ValueError, match="Invalid backend"): + backend.set_backend("fortran") + + +def test_cpp_namespace_is_available(): + import cppe + + assert hasattr(cppe, "cpp") + assert hasattr(cppe.cpp, "PotfileReader") + + +def test_potfile_reader_dispatches_by_backend(): + import cppe + import cppe.potfile as py_potfile + + potfile_path = Path(__file__).parent / "potfiles" / "pna_6w.pot" + + cppe.set_backend("cpp") + cpp_reader = cppe.PotfileReader(str(potfile_path)) + assert isinstance(cpp_reader, cppe.cpp.PotfileReader) + + cppe.set_backend("python") + py_reader = cppe.PotfileReader(str(potfile_path)) + assert isinstance(py_reader, py_potfile.PotfileReader) + + cppe.set_backend("cpp") diff --git a/tests/test_bmatrix_python.py b/tests/test_bmatrix_python.py new file mode 100644 index 00000000..69b807c6 --- /dev/null +++ b/tests/test_bmatrix_python.py @@ -0,0 +1,74 @@ +import numpy as np + +import cppe + + +def _polsites(): + potentials = cppe.cpp.PotfileReader("tests/potfiles/pna_6w.pot").read() + return cppe.get_polarizable_sites(potentials) + + +def test_bmatrix_apply_python_matches_cpp_direct(): + polsites = _polsites() + options = { + "potfile": "tests/potfiles/pna_6w.pot", + "summation_induced_fields": "direct", + } + vec = np.linspace(-0.2, 0.7, 3 * len(polsites)) + + cppe.set_backend("cpp") + ref = cppe.BMatrix(polsites, options).apply(vec) + + cppe.set_backend("python") + got = cppe.BMatrix(polsites, options).apply(vec) + cppe.set_backend("cpp") + + np.testing.assert_allclose(got, ref, atol=1e-12) + + +def test_bmatrix_dense_python_matches_cpp_direct(): + polsites = _polsites() + options = { + "potfile": "tests/potfiles/pna_6w.pot", + "summation_induced_fields": "direct", + } + + cppe.set_backend("cpp") + ref = cppe.BMatrix(polsites, options).to_dense_matrix() + + cppe.set_backend("python") + got = cppe.BMatrix(polsites, options).to_dense_matrix() + cppe.set_backend("cpp") + + np.testing.assert_allclose(got, ref, atol=1e-12) + + +def test_bmatrix_python_numba_nopython_compiles(): + from cppe import solver + + polsites = _polsites() + options = { + "potfile": "tests/potfiles/pna_6w.pot", + "summation_induced_fields": "direct", + } + vec = np.linspace(0.1, 1.1, 3 * len(polsites)) + + cppe.set_backend("python") + bmat = cppe.BMatrix(polsites, options) + bmat.apply(vec) + bmat.apply_diagonal(vec) + bmat.apply_diagonal_inverse(vec) + cppe.set_backend("cpp") + + assert solver._induced_fields_direct_kernel.nopython_signatures + assert solver._apply_diagonal_kernel.nopython_signatures + assert solver._apply_diagonal_inverse_kernel.nopython_signatures + + +def test_triangle_to_mat_matches_cpp_math_helper(): + from cppe import solver + + tri = np.array([1.1, -0.2, 0.7, 2.5, -0.9, 3.2]) + got = solver._triangle_to_mat(tri) + ref = cppe.triangle_to_mat(tri.tolist()) + np.testing.assert_allclose(got, ref, atol=1e-14) diff --git a/tests/test_induced_moments_python.py b/tests/test_induced_moments_python.py new file mode 100644 index 00000000..605bdbd6 --- /dev/null +++ b/tests/test_induced_moments_python.py @@ -0,0 +1,25 @@ +import numpy as np + +import cppe + +from .cache import cache + + +def test_induced_moments_python_matches_cpp_direct(): + mol = cache.molecule["pna"] + options = {"potfile": "tests/potfiles/pna_6w.pot", "induced_thresh": 1e-12} + + cppe.set_backend("cpp") + state = cppe.CppeState(options, mol, lambda _: None) + state.calculate_static_energies_and_fields() + rhs = state.static_fields + potentials = state.potentials + + cppe.set_backend("cpp") + ref = cppe.InducedMoments(potentials, options).compute(rhs, True) + + cppe.set_backend("python") + got = cppe.InducedMoments(potentials, options).compute(rhs, True) + + cppe.set_backend("cpp") + np.testing.assert_allclose(got, ref, atol=1e-10) diff --git a/tests/test_multipole_expansion_python.py b/tests/test_multipole_expansion_python.py new file mode 100644 index 00000000..1ef30352 --- /dev/null +++ b/tests/test_multipole_expansion_python.py @@ -0,0 +1,46 @@ +import numpy as np + +import cppe + +from .cache import cache + + +def test_multipole_interaction_energy_python_matches_cpp(): + mol = cache.molecule["pna"] + potentials = cppe.cpp.PotfileReader("tests/potfiles/pna_6w.pot").read() + + cppe.set_backend("cpp") + e_cpp = cppe.MultipoleExpansion(mol, potentials).interaction_energy() + + cppe.set_backend("python") + e_py = cppe.MultipoleExpansion(mol, potentials).interaction_energy() + + cppe.set_backend("cpp") + np.testing.assert_allclose(e_py, e_cpp, atol=1e-12) + + +def test_multipole_nuclear_gradient_python_matches_cpp(): + mol = cache.molecule["pna"] + potentials = cppe.cpp.PotfileReader("tests/potfiles/pna_6w.pot").read() + + cppe.set_backend("cpp") + grad_cpp = cppe.MultipoleExpansion(mol, potentials).nuclear_gradient() + + cppe.set_backend("python") + grad_py = cppe.MultipoleExpansion(mol, potentials).nuclear_gradient() + + cppe.set_backend("cpp") + np.testing.assert_allclose(grad_py, grad_cpp, atol=1e-12) + + +def test_multipole_interaction_numba_nopython_compiles(): + from cppe import multipole + + mol = cache.molecule["pna"] + potentials = cppe.cpp.PotfileReader("tests/potfiles/pna_6w.pot").read() + + cppe.set_backend("python") + cppe.MultipoleExpansion(mol, potentials).interaction_energy() + cppe.set_backend("cpp") + + assert multipole._interaction_energy_kernel.nopython_signatures diff --git a/tests/test_multipole_fields_python.py b/tests/test_multipole_fields_python.py new file mode 100644 index 00000000..218f88b7 --- /dev/null +++ b/tests/test_multipole_fields_python.py @@ -0,0 +1,36 @@ +import numpy as np + +import cppe + + +def test_multipole_fields_python_matches_cpp_direct(): + potentials = cppe.cpp.PotfileReader("tests/potfiles/pna_6w.pot").read() + options = { + "potfile": "tests/potfiles/pna_6w.pot", + "summation_induced_fields": "direct", + } + + cppe.set_backend("cpp") + ref = cppe.MultipoleFields(potentials, options).compute() + + cppe.set_backend("python") + got = cppe.MultipoleFields(potentials, options).compute() + + cppe.set_backend("cpp") + np.testing.assert_allclose(got, ref, atol=1e-12) + + +def test_multipole_fields_numba_nopython_compiles(): + from cppe import multipole_fields + + potentials = cppe.cpp.PotfileReader("tests/potfiles/pna_6w.pot").read() + options = { + "potfile": "tests/potfiles/pna_6w.pot", + "summation_induced_fields": "direct", + } + + cppe.set_backend("python") + cppe.MultipoleFields(potentials, options).compute() + cppe.set_backend("cpp") + + assert multipole_fields._multipole_fields_direct_kernel.nopython_signatures diff --git a/tests/test_nuclear_fields_python.py b/tests/test_nuclear_fields_python.py new file mode 100644 index 00000000..1ec73612 --- /dev/null +++ b/tests/test_nuclear_fields_python.py @@ -0,0 +1,48 @@ +import numpy as np + +import cppe + +from .cache import cache + + +def test_nuclear_fields_python_matches_cpp(): + mol = cache.molecule["pna"] + potentials = cppe.cpp.PotfileReader("tests/potfiles/pna_6w.pot").read() + + cppe.set_backend("cpp") + fields_cpp = cppe.NuclearFields(mol, potentials).compute() + + cppe.set_backend("python") + fields_py = cppe.NuclearFields(mol, potentials).compute() + + cppe.set_backend("cpp") + np.testing.assert_allclose(fields_py, fields_cpp, atol=1e-14) + + +def test_nuclear_gradient_python_matches_cpp(): + mol = cache.molecule["pna"] + potentials = cppe.cpp.PotfileReader("tests/potfiles/pna_6w.pot").read() + + cppe.set_backend("cpp") + grad_cpp = cppe.NuclearFields(mol, potentials).nuclear_gradient() + + cppe.set_backend("python") + grad_py = cppe.NuclearFields(mol, potentials).nuclear_gradient() + + cppe.set_backend("cpp") + np.testing.assert_allclose(grad_py, grad_cpp, atol=1e-13) + + +def test_nuclear_fields_numba_nopython_compiles(): + from cppe import fields + + mol = cache.molecule["pna"] + potentials = cppe.cpp.PotfileReader("tests/potfiles/pna_6w.pot").read() + + cppe.set_backend("python") + cppe.NuclearFields(mol, potentials).compute() + cppe.NuclearFields(mol, potentials).nuclear_gradient() + cppe.set_backend("cpp") + + assert fields._nuclear_fields_kernel.nopython_signatures + assert fields._nuclear_gradient_kernel.nopython_signatures diff --git a/tests/test_potfile_reader_python.py b/tests/test_potfile_reader_python.py new file mode 100644 index 00000000..c7a079ad --- /dev/null +++ b/tests/test_potfile_reader_python.py @@ -0,0 +1,76 @@ +from pathlib import Path + +import numpy as np +import pytest + +import cppe + + +def _serialize_potential(potential): + multipoles = [] + for multipole in potential.multipoles: + multipoles.append((multipole.k, np.array(multipole.values, dtype=float))) + + polarizability = None + if potential.is_polarizable: + polarizability = np.array(potential.polarizability.values, dtype=float) + + return { + "index": potential.index, + "element": potential.element, + "position": np.array(potential.position, dtype=float), + "exclusions": list(potential.exclusions), + "multipoles": multipoles, + "polarizability": polarizability, + } + + +@pytest.mark.parametrize( + "filename", + [ + "pna_6w.pot", + "pna_6w_isopol.pot", + "loprop_solvated_20.pot", + ], +) +def test_python_potfile_reader_matches_cpp(filename): + potfile = Path(__file__).parent / "potfiles" / filename + + cpp_potentials = cppe.PotfileReader(str(potfile)).read() + cppe.set_backend("python") + py_potentials = cppe.PotfileReader(str(potfile)).read() + cppe.set_backend("cpp") + + assert len(cpp_potentials) == len(py_potentials) + + for cpp_potential, py_potential in zip(cpp_potentials, py_potentials): + cpp_data = _serialize_potential(cpp_potential) + py_data = _serialize_potential(py_potential) + + assert cpp_data["index"] == py_data["index"] + assert cpp_data["element"] == py_data["element"] + assert cpp_data["exclusions"] == py_data["exclusions"] + np.testing.assert_allclose( + cpp_data["position"], py_data["position"], atol=1e-14 + ) + + assert len(cpp_data["multipoles"]) == len(py_data["multipoles"]) + for (cpp_k, cpp_values), (py_k, py_values) in zip( + cpp_data["multipoles"], py_data["multipoles"] + ): + assert cpp_k == py_k + np.testing.assert_allclose(cpp_values, py_values, atol=1e-14) + + if cpp_data["polarizability"] is None: + assert py_data["polarizability"] is None + else: + np.testing.assert_allclose( + cpp_data["polarizability"], py_data["polarizability"], atol=1e-14 + ) + + +def test_python_potfile_reader_missing_file_raises(): + cppe.set_backend("python") + with pytest.raises(RuntimeError, match="Potential file does not exist"): + cppe.PotfileReader("does-not-exist.pot") + cppe.set_backend("cpp") diff --git a/tests/test_state_python.py b/tests/test_state_python.py new file mode 100644 index 00000000..5d714d79 --- /dev/null +++ b/tests/test_state_python.py @@ -0,0 +1,90 @@ +import numpy as np + +import cppe + +from .cache import cache + + +def test_state_static_fields_python_matches_cpp(): + mol = cache.molecule["pna"] + options = {"potfile": "tests/potfiles/pna_6w.pot"} + + cppe.set_backend("cpp") + state_cpp = cppe.CppeState(options, mol, lambda _: None) + state_cpp.calculate_static_energies_and_fields() + + cppe.set_backend("python") + state_py = cppe.CppeState(options, mol, lambda _: None) + state_py.calculate_static_energies_and_fields() + + cppe.set_backend("cpp") + np.testing.assert_allclose( + state_py.nuclear_fields, state_cpp.nuclear_fields, atol=1e-13 + ) + np.testing.assert_allclose( + state_py.multipole_fields, state_cpp.multipole_fields, atol=1e-13 + ) + np.testing.assert_allclose( + state_py.static_fields, state_cpp.static_fields, atol=1e-13 + ) + np.testing.assert_allclose( + state_py.energies["Electrostatic"]["Nuclear"], + state_cpp.energies["Electrostatic"]["Nuclear"], + atol=1e-12, + ) + + +def test_state_gradient_methods_python_matches_cpp(): + mol = cache.molecule["pna"] + options = {"potfile": "tests/potfiles/pna_6w.pot"} + + cppe.set_backend("cpp") + state_cpp = cppe.CppeState(options, mol, lambda _: None) + grad_e_cpp = state_cpp.nuclear_interaction_energy_gradient() + grad_f_cpp = state_cpp.nuclear_field_gradient() + + cppe.set_backend("python") + state_py = cppe.CppeState(options, mol, lambda _: None) + grad_e_py = state_py.nuclear_interaction_energy_gradient() + grad_f_py = state_py.nuclear_field_gradient() + + cppe.set_backend("cpp") + np.testing.assert_allclose(grad_e_py, grad_e_cpp, atol=1e-12) + np.testing.assert_allclose(grad_f_py, grad_f_cpp, atol=1e-12) + + +def test_state_update_induced_moments_python_matches_cpp(): + mol = cache.molecule["pna"] + options = {"potfile": "tests/potfiles/pna_6w.pot", "induced_thresh": 1e-12} + + cppe.set_backend("cpp") + state_cpp = cppe.CppeState(options, mol, lambda _: None) + state_cpp.calculate_static_energies_and_fields() + zeros_cpp = np.zeros_like(state_cpp.static_fields) + state_cpp.update_induced_moments(zeros_cpp, False) + + cppe.set_backend("python") + state_py = cppe.CppeState(options, mol, lambda _: None) + state_py.calculate_static_energies_and_fields() + zeros_py = np.zeros_like(state_py.static_fields) + state_py.update_induced_moments(zeros_py, False) + + cppe.set_backend("cpp") + np.testing.assert_allclose( + state_py.get_induced_moments(), state_cpp.get_induced_moments(), atol=1e-10 + ) + np.testing.assert_allclose( + state_py.energies["Polarization"]["Electronic"], + state_cpp.energies["Polarization"]["Electronic"], + atol=1e-12, + ) + np.testing.assert_allclose( + state_py.energies["Polarization"]["Nuclear"], + state_cpp.energies["Polarization"]["Nuclear"], + atol=1e-12, + ) + np.testing.assert_allclose( + state_py.energies["Polarization"]["Multipoles"], + state_cpp.energies["Polarization"]["Multipoles"], + atol=1e-12, + )