bfflearn 0.1.0.dev1__py3-none-any.whl

bff/bayes/utils.py

@@ -0,0 +1,333 @@
+from pathlib import Path
+from typing import Callable, Sequence, Tuple, Union
+
+import torch
+from torch.autograd.functional import hessian
+
+PathLike = Union[str, Path]
+
+
+def smape(y_true: torch.Tensor, y_pred: torch.Tensor) -> float:
+    """Compute the symmetric mean absolute percentage error."""
+
+    device = y_true.device if isinstance(y_true, torch.Tensor) else 'cpu'
+
+    y_true = check_tensor(y_true, device=device)
+    y_pred = check_tensor(y_pred, device=device)
+
+    abs_diff = torch.sum(torch.abs(y_true - y_pred), dim=1)
+    y_true_abs = torch.sum(torch.abs(y_true), dim=1)
+    y_pred_abs = torch.sum(torch.abs(y_pred), dim=1)
+    norm = y_true_abs + y_pred_abs
+
+    return float(torch.mean(abs_diff / norm).item())
+
+def initialize_walkers(
+    priors,
+    n_walkers: int,
+    constraint: Callable = None,
+    max_attempts: int | None = None,
+) -> torch.Tensor:
+
+    """
+    Initialize walkers for the MCMC sampler.
+
+    Parameters
+    ----------
+    priors : dict
+        Dict of prior distributions for each parameter.
+    n_walkers : int
+        Number of walkers to initialize.
+    specs : object, optional
+        Specifications object containing bounds and constraints.
+        If provided, walkers will be initialized within the bounds.
+
+    Returns
+    -------
+    torch.Tensor
+        A tensor of shape (n_walkers, n_params) containing the initial positions
+        of the walkers, sampled from the prior distributions.
+    """
+
+    if constraint is None:
+        means = torch.tensor([p.mean for p in priors], dtype=torch.float32)
+        stds = torch.tensor([p.scale for p in priors], dtype=torch.float32)
+        p0 = torch.normal(means.expand(n_walkers, -1), stds.expand(n_walkers, -1))
+    else:
+        n_params = constraint.n_params
+        n_dim = len(priors)
+        max_attempts = max_attempts or (1000 * n_walkers)
+
+        p0 = torch.empty((n_walkers, n_dim))
+        count = 0
+        attempts = 0
+        while count < n_walkers and attempts < max_attempts:
+            p0_trial = torch.tensor([p.sample().item() for p in priors])
+            if constraint(p0_trial[:n_params]):
+                p0[count] = p0_trial
+                count += 1
+            attempts += 1
+        if count < n_walkers:
+            raise RuntimeError(
+                "Failed to initialize constrained walkers from the priors. "
+                f"Accepted {count}/{n_walkers} samples after {attempts} attempts."
+            )
+    return p0
+
+
+def check_tensor(
+    x: Union[torch.Tensor, float, int, list, tuple],
+    device: str,
+    dtype: torch.dtype = torch.float32
+) -> torch.Tensor:
+    """Convert input to a torch tensor on the specified device."""
+    if not isinstance(x, torch.Tensor):
+        return torch.as_tensor(x, device=device, dtype=dtype)
+    return x.to(device, dtype=dtype)
+
+
+def check_device(device: str) -> None:
+    """Check if the specified device is available."""
+    if device.startswith("cuda"):
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA is not available.")
+    elif device.startswith("mps"):
+        if not torch.mps.is_available():
+            raise RuntimeError("MPS is not available.")
+
+
+@torch.no_grad()
+def nearest_positive_definite(A: torch.Tensor) -> torch.Tensor:
+    """Find the nearest positive definite matrix to A."""
+    # Symmetrize
+    A_sym = (A + A.T) / 2
+
+    # Check if already PD
+    if torch.all(torch.linalg.eigvalsh(A_sym) > 0):
+        return A_sym
+
+    # Eigen-decomposition
+    eigenvalues, eigenvectors = torch.linalg.eigh(A_sym)
+
+    # Shift eigenvalues minimally
+    min_eig = eigenvalues.min()
+    eps = 1e-8  # small positive shift
+
+    if min_eig < eps:
+        eigenvalues = eigenvalues - min_eig + eps
+
+    A_pd = (eigenvectors @ torch.diag(eigenvalues)) @ eigenvectors.T
+    return A_pd
+
+
+def train_test_split(
+    X: torch.Tensor, y: torch.Tensor, test_fraction: float = 0.2
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Split the dataset into training and testing sets."""
+
+    n = len(X)
+    if n != len(y):
+        raise ValueError("X and y must have the same length.")
+    if not (0 < test_fraction < 1):
+        raise ValueError("test_fraction must be between 0 and 1.")
+    if n < 2:
+        raise ValueError("X and y must have at least 2 samples.")
+
+    indices = torch.randperm(n)
+    test_size = int(n * test_fraction)
+    idx_train = indices[test_size:]
+    idx_test = indices[:test_size]
+    return X[idx_train], X[idx_test], y[idx_train], y[idx_test]
+
+
+# Global toggle for manual squared distance
+_MANUAL_MODE = False
+
+
+class enable_manual_dist:
+    """Enable Hessian-safe manual pairwise distances within a context."""
+    def __enter__(self) -> None:
+        global _MANUAL_MODE
+        self._prev = _MANUAL_MODE
+        _MANUAL_MODE = True
+
+    def __exit__(self, *args) -> None:
+        global _MANUAL_MODE
+        _MANUAL_MODE = self._prev
+
+
+def with_manual_sqdist_flag(fn: Callable) -> Callable:
+    """Inject the global manual-distance flag into kernel call sites."""
+    def wrapper(*args, manual_sqdist=False, **kwargs):
+        return fn(*args, manual_sqdist=manual_sqdist or _MANUAL_MODE, **kwargs)
+    return wrapper
+
+
+def find_max_stable_lr(
+    fn: Callable,
+    p0: torch.Tensor,
+    learning_rates: (
+        Sequence[Union[float, torch.Tensor]] | float | torch.Tensor | None
+    ) = None,
+    max_iter: int = 100,
+    param_bounds: Tuple[float, float] = (-7, 7),
+) -> Union[float, None]:
+    """Find the largest stable learning rate for gradient-based optimization.
+
+    Parameters
+    ----------
+    fn : Callable
+        Objective function returning a scalar tensor.
+    p0 : torch.Tensor
+        Initial parameter vector.
+    learning_rates : iterable of float, optional
+        Learning rates to test. Defaults to log-spaced values.
+    max_iter : int
+        Number of steps to test for each learning rate.
+    param_bounds : tuple of float
+        Bounds beyond which parameters are considered unstable.
+
+    Returns
+    -------
+    float or None
+        The largest stable learning rate found, or None if none were stable.
+    """
+    if learning_rates is None:
+        learning_rates = 10 ** torch.linspace(-1, -6, 6)
+    elif isinstance(learning_rates, torch.Tensor) and learning_rates.ndim == 0:
+        learning_rates = [float(learning_rates.item())]
+    elif isinstance(learning_rates, (int, float)):
+        learning_rates = [float(learning_rates)]
+    else:
+        learning_rates = list(learning_rates)
+    lower, upper = param_bounds
+    for lr in learning_rates:
+        x = p0.clone().detach().requires_grad_(True)
+        opt = torch.optim.SGD([x], lr=lr)
+
+        for i in range(max_iter):
+            if torch.any(x < lower) or torch.any(x > upper):
+                break
+            opt.zero_grad()
+            loss = -fn(x)
+            loss.backward()
+            opt.step()
+        else:
+            # Only gets executed if inner loop did not break (i.e., stable)
+            return lr
+
+
+def find_map(
+    fn: Callable,
+    x0: torch.Tensor,
+    lr: Union[float, torch.Tensor] = None,
+    max_iter: int = 10000,
+    tol_grad: float = 1e-2,
+    device: str = 'cpu',
+    logger: Callable = None
+) -> torch.Tensor:
+
+    """Find the maximum a posteriori (MAP) estimate
+    using gradient-based optimization.
+
+    Parameters
+    ----------
+    fn : Callable
+        Objective function to maximize (log-posterior).
+    x0 : torch.Tensor
+        Initial parameter vector.
+    lr : float or torch.Tensor, optional
+        Learning rate for optimization. If None, it will be determined by a search.
+    max_iter : int
+        Maximum number of optimization iterations.
+    tol_grad : float
+        Gradient norm threshold for convergence.
+    device : str
+        Device to perform computations on.
+    logger : Callable, optional
+        Logger for progress updates. Should accept a string and a level argument.
+
+    Returns
+    -------
+    torch.Tensor
+        The MAP estimate found by optimization.
+    """
+
+    if logger is not None:
+        logger.info("learning rate search: in progress...", level=2, overwrite=True)
+
+    lr_opt = find_max_stable_lr(fn, x0, learning_rates=lr)
+    if lr_opt is not None:
+        lr_opt = 0.5 * lr_opt
+        if logger is not None:
+            logger.info(f"learning rate search: Done. | lr = {lr_opt:.1e}", level=2)
+    else:
+        raise ValueError("No stable learning rate found.")
+
+    x0 = x0.clone().detach().to(device).requires_grad_(True)
+    optimizer = torch.optim.SGD([x0], lr=lr_opt)
+
+    for i in range(max_iter):
+        optimizer.zero_grad()
+        loss = -fn(x0)
+        loss.backward()
+        grad_norm = x0.grad.norm().item()
+        if i % 100 == 0 and logger is not None:
+            logger.info(
+                (
+                    f"MAP search: it. {i}/{max_iter} | "
+                    f"loss: {loss.item():.3f} | grad: {grad_norm:.3f}/{tol_grad}"
+                ),
+                level=2,
+                overwrite=True
+            )
+        optimizer.step()
+        if grad_norm < tol_grad:
+            if logger is not None:
+                logger.info("MAP search: Done.", level=2)
+            break
+
+    else:
+        if logger is not None:
+            logger.info(
+                "MAP search: Fail. | Max iterations reached without convergence.",
+                level=2,
+            )
+
+    return x0.detach()
+
+
+def laplace_approximation(
+    fn: Callable,
+    map_theta: torch.Tensor,
+    device: str = 'cpu'
+) -> torch.Tensor:
+    """
+    Perform Laplace approximation around MAP estimate.
+
+    Parameters
+    ----------
+    fn : Callable
+        Objective function to approximate (log-posterior).
+    map_theta : torch.Tensor
+        MAP estimate around which to perform the approximation.
+    device : str
+        Device to perform computations on.
+
+    Returns
+    -------
+    cov : torch.Tensor
+        The approximate posterior covariance (inverse Hessian).
+    """
+
+    map_theta = check_tensor(map_theta, device=device)
+    with enable_manual_dist():
+        H = -hessian(lambda th: fn(th).sum(), map_theta)
+
+    reg_eye = 1e-6 * torch.eye(H.shape[0], device=H.device)
+    cov = torch.linalg.inv(H + reg_eye)
+
+    # symmerize covariance matrix
+    cov = (cov + cov.T) / 2
+
+    return cov

bff/cli.py

@@ -0,0 +1,254 @@
+"""Command-line entry points for Bayesic Force Fields."""
+
+import os
+from pathlib import Path
+from typing import Callable
+
+import typer
+
+app = typer.Typer(
+    help="BayesicForceFields: Bayesian optimization of molecular force fields.",
+    no_args_is_help=True,
+    add_completion=False,
+)
+
+WorkflowMain = Callable[[Path], None]
+WORKFLOW_COMMANDS = (
+    "prepare",
+    "simulate",
+    "qoi",
+    "train",
+    "learn",
+    "validate",
+)
+
+_COMPLETION_ACTIVATE = f"""# Managed automatically by Bayesic Force Fields.
+if [ -n "${{BASH_VERSION-}}" ]; then
+    _bff_completion() {{
+        local cur prev words cword
+        if declare -F _init_completion >/dev/null 2>&1; then
+            _init_completion || return
+        else
+            COMPREPLY=()
+            cur="${{COMP_WORDS[COMP_CWORD]}}"
+            prev="${{COMP_WORDS[COMP_CWORD-1]}}"
+            words=("${{COMP_WORDS[@]}}")
+            cword=$COMP_CWORD
+        fi
+
+        local commands="{" ".join(WORKFLOW_COMMANDS)}"
+        if [[ $COMP_CWORD -eq 1 ]]; then
+            COMPREPLY=($(compgen -W "$commands" -- "$cur"))
+            return 0
+        fi
+
+        case "${{COMP_WORDS[1]}}" in
+            prepare|simulate|qoi|train|learn|validate)
+                COMPREPLY=($(compgen -f -- "$cur"))
+                ;;
+            *)
+                COMPREPLY=()
+                ;;
+        esac
+    }}
+
+    complete -o nosort -F _bff_completion bff
+elif [ -n "${{ZSH_VERSION-}}" ]; then
+    autoload -Uz compinit 2>/dev/null || true
+    compinit -i 2>/dev/null || true
+
+    _bff_completion() {{
+        local -a commands
+        commands=({" ".join(WORKFLOW_COMMANDS)})
+
+        if (( CURRENT == 2 )); then
+            compadd -- "${{commands[@]}}"
+            return 0
+        fi
+
+        case "${{words[2]}}" in
+            prepare|simulate|qoi|train|learn|validate)
+                _files
+                ;;
+        esac
+    }}
+
+    compdef _bff_completion bff
+fi
+"""
+
+_COMPLETION_DEACTIVATE = """# Managed automatically by Bayesic Force Fields.
+if [ -n "${BASH_VERSION-}" ]; then
+    complete -r bff 2>/dev/null || true
+elif [ -n "${ZSH_VERSION-}" ]; then
+    unfunction _bff_completion 2>/dev/null || true
+    compdef -d bff 2>/dev/null || true
+fi
+"""
+
+
+def _conda_completion_hook_paths() -> tuple[Path, Path] | None:
+    conda_prefix = os.environ.get("CONDA_PREFIX")
+    if not conda_prefix:
+        return None
+
+    conda_dir = Path(conda_prefix) / "etc" / "conda"
+    return (
+        conda_dir / "activate.d" / "bff-completion.sh",
+        conda_dir / "deactivate.d" / "bff-completion.sh",
+    )
+
+
+def _write_if_changed(path: Path, content: str) -> None:
+    if path.exists() and path.read_text() == content:
+        return
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(content)
+
+
+def _ensure_completion_hooks() -> None:
+    hook_paths = _conda_completion_hook_paths()
+    if hook_paths is None:
+        return
+
+    activate_hook, deactivate_hook = hook_paths
+    try:
+        _write_if_changed(activate_hook, _COMPLETION_ACTIVATE)
+        _write_if_changed(deactivate_hook, _COMPLETION_DEACTIVATE)
+    except OSError:
+        return
+
+
+def run_workflow(
+    fn_config: Path,
+    workflow_main: WorkflowMain,
+    workflow_name: str,
+) -> None:
+    """
+    Run a workflow main function with a config path.
+    """
+    try:
+        workflow_main(fn_config)
+    except FileNotFoundError as exc:
+        raise typer.BadParameter(str(exc), param_hint="fn_config") from exc
+    except ValueError as exc:
+        raise typer.BadParameter(str(exc), param_hint=workflow_name) from exc
+    except RuntimeError as exc:
+        typer.echo(str(exc), err=True)
+        raise typer.Exit(code=1) from exc
+
+
+def config_argument() -> Path:
+    return typer.Argument(
+        ...,
+        exists=True,
+        dir_okay=False,
+        help="Path to the configuration file.",
+    )
+
+
+@app.callback()
+def init_cli() -> None:
+    """
+    Initialize the CLI environment.
+    """
+    _ensure_completion_hooks()
+
+
+@app.command()
+def version() -> None:
+    """
+    Print package version.
+    """
+    from importlib import metadata
+
+    try:
+        version_str = metadata.version("bfflearn")
+    except metadata.PackageNotFoundError:
+        from . import __version__
+
+        version_str = __version__
+
+    typer.echo(f"BayesicForceFields version: {version_str}")
+
+
+@app.command()
+def prepare(fn_config: Path = config_argument()) -> None:
+    """
+    Prepare equilibration, training, and reference assets.
+    """
+    from bff.workflows.prepare import main as prepare_main
+
+    run_workflow(fn_config, prepare_main, "prepare")
+
+
+@app.command(name="simulate")
+def simulate(fn_config: Path = config_argument()) -> None:
+    """
+    Run a sampled molecular-dynamics campaign for training-set generation.
+    """
+    from bff.workflows.simulate import main as simulate_main
+
+    run_workflow(fn_config, simulate_main, "simulate")
+
+
+@app.command()
+def validate(fn_config: Path = config_argument()) -> None:
+    """
+    Run a validation campaign for provided parameter samples.
+    """
+    from bff.workflows.validate import main as validate_main
+
+    run_workflow(fn_config, validate_main, "validate")
+
+
+@app.command(hidden=True)
+def md(fn_config: Path = config_argument()) -> None:
+    """
+    Run molecular dynamics from a configuration file.
+    """
+    from bff.workflows.md import main as md_main
+
+    run_workflow(fn_config, md_main, "md")
+
+
+@app.command(name="qoi")
+def qoi(fn_config: Path = config_argument()) -> None:
+    """
+    Analyze quantities of interest from simulation data.
+    """
+    from bff.workflows.qoi import main as qoi_main
+
+    run_workflow(fn_config, qoi_main, "qoi")
+
+
+@app.command(name="analyze", hidden=True)
+def analyze_alias(fn_config: Path = config_argument()) -> None:
+    """
+    Hidden compatibility alias for the QoI workflow.
+    """
+    qoi(fn_config)
+
+
+@app.command()
+def train(fn_config: Path = config_argument()) -> None:
+    """
+    Train surrogate models from analyzed QoI datasets.
+    """
+    from bff.workflows.train import main as train_main
+
+    run_workflow(fn_config, train_main, "train")
+
+
+@app.command()
+def learn(fn_config: Path = config_argument()) -> None:
+    """
+    Run posterior learning from trained surrogate models.
+    """
+    from bff.workflows.learn import main as learn_main
+
+    run_workflow(fn_config, learn_main, "learn")
+
+
+if __name__ == "__main__":
+    app()

bff/data.py

@@ -0,0 +1,93 @@
+import numpy as np
+
+WATER_4SITE = np.array(
+    [[0.02909557,  0.03692916, -0.04588176],
+     [0.43909732, -0.54306981,  0.60411833],
+     [-0.90090283, -0.04307076,  0.12411808],
+     [-0.0309039,  -0.04307076,  0.05411814]]
+)
+
+WATER_3SITE = np.array(
+    [[0.02909557,  0.03692916, -0.04588176],
+     [0.43909732, -0.54306981,  0.60411833],
+     [-0.90090283, -0.04307076,  0.12411808]]
+)
+
+IONS = ['li', 'na', 'k', 'ca', 'cal', 'mg', 'f', 'cl', 'br', 'i']
+CATIONS = {'li', 'na', 'k', 'rb', 'ca', 'mg'}
+ANIONS = {'f', 'cl', 'br', 'i'}
+
+WATERS = ['SOL', 'water', 'HOH']
+
+CP2K_KIND_DEFAULTS = {
+    'h': {'basis_set': 'TZV2P-GTH-q1', 'potential': 'GTH-PBE-q1'},
+    'o': {'basis_set': 'TZV2P-GTH-q6', 'potential': 'GTH-PBE-q6'},
+    'c': {'basis_set': 'TZV2P-GTH-q4', 'potential': 'GTH-PBE-q4'},
+    'p': {'basis_set': 'TZV2P-GTH-q5', 'potential': 'GTH-PBE-q5'},
+    'n': {'basis_set': 'TZV2P-GTH-q5', 'potential': 'GTH-PBE-q5'},
+    'na': {'basis_set': 'TZV2P-GTH-q9', 'potential': 'GTH-PBE-q9'},
+    'ca': {'basis_set': 'TZV2P-MOLOPT-PBE-GTH-q10', 'potential': 'GTH-PBE-q10'},
+    'cl': {'basis_set': 'TZV2P-GTH-q7', 'potential': 'GTH-PBE-q7'},
+    's': {'basis_set': 'TZV2P-GTH-q6', 'potential': 'GTH-PBE-q6'},
+}
+
+CP2K_INPUT_TEMPLATE = {
+    'global': {
+        'project_name': 'project',
+        'print_level': 'LOW',
+        'run_type': 'MD',
+        'walltime': '03:50:00'
+    },
+    'motion': {
+        'md': {'temperature': 300.0, 'timestep': '[fs] 0.5', 'steps': 500},
+        'print': {
+            'forces': {
+                'format': 'DCD',
+                'each': {'md': 1}
+            },
+            'trajectory': {
+                'format': 'DCD',
+                'each': {'md': 1}
+            },
+            'restart': {
+                'backup_copies': 1,
+                'each': {'md': 1}
+            }
+        }
+    },
+    'force_eval': {
+        'method': 'Quickstep',
+        'subsys': {},
+        'dft': {
+            'basis_set_file_name': ['GTH_BASIS_SETS', 'BASIS_MOLOPT_UZH'],
+            'potential_file_name': 'GTH_POTENTIALS',
+            'scf': {
+                'max_scf': 20, 'eps_scf': 5e-07,
+                'outer_scf': {'max_scf': 20, 'eps_scf': 5e-07},
+                'ot': {
+                    'preconditioner': 'FULL_ALL',
+                    'energy_gap': -1.0,
+                    'minimizer': 'DIIS'}
+            },
+            'mgrid': {'cutoff': '[Ry] 400.0'},
+            'qs': {
+                'method': 'GPW',
+                'eps_default': 1e-12,
+                'extrapolation_order': 4,
+                'extrapolation': 'ASPC'
+            },
+            'xc': {
+                'xc_functional': {'pbe': {'parametrization': 'REVPBE'}},
+                'xc_grid': {'xc_deriv': 'NN50_SMOOTH'},
+                'vdw_potential': {
+                    'potential_type': 'PAIR_POTENTIAL',
+                    'pair_potential': {
+                        'type': 'DFTD3', 'r_cutoff': '[angstrom] 16.0',
+                        'reference_functional': 'revPBE',
+                        'parameter_file_name': 'dftd3.dat'
+                    }
+                }
+            }
+        }
+    }
+}

bff/domain/__init__.py

@@ -0,0 +1,12 @@
+from .specs import Bounds, ChargeConstraint, RandomParamsGenerator, Specs
+from .trainset import SimulationSystem, TrainSetInfo, TrajectorySet
+
+__all__ = [
+    "Bounds",
+    "ChargeConstraint",
+    "RandomParamsGenerator",
+    "SimulationSystem",
+    "Specs",
+    "TrainSetInfo",
+    "TrajectorySet",
+]