diff --git a/README.md b/README.md index aecc783..088c818 100644 --- a/README.md +++ b/README.md @@ -71,6 +71,39 @@ Each non-comment, non-empty line in `metadata_file` should contain: frames by `exp(-energy / kT)`; omit from all lines to use uniform weights. Lines may start with `#` for comments. Mixing lines with and without temperature is rejected. + +### BayesWHAM configuration + +Bayesian reconstruction uses the same metadata and trajectory layout as the WHAM +solvers and accepts the corresponding 1D or 2D histogram parameters. Append the +following fields to reuse existing YAML inputs: + +- `num_samples` (optional, default `200`): total Dirichlet-resampled histograms + to draw for posterior estimation. +- `burn_in` (optional, default `50`): number of initial samples to discard + before collecting posterior statistics. +- `thinning` (optional, default `1`): stride between retained samples. +- `dirichlet_alpha` (optional, default `1.0`): concentration added to each + histogram bin prior to sampling, mirroring the noninformative prior used in + the reference BayesWHAM script. + +Example 1D configuration (reuses the umbrella trajectories and metadata from the +WHAM example): + +```yaml +hist_min: -1.0 +hist_max: 1.0 +num_bins: 50 +tolerance: 1e-5 +temperature: 1.0 +numpad: 0 +metadata_file: examples/output_1D/umbrella_metadata.txt +freefile: examples/output_1D/bayes.free +num_samples: 400 +burn_in: 100 +thinning: 2 +dirichlet_alpha: 1.0 +``` All metadata-referenced trajectories (for WHAM or projection runs) may use relative paths; they are resolved against the directory that contains the metadata file, so you can keep each metadata bundle self-contained regardless of the working directory. @@ -145,6 +178,7 @@ Sample umbrella-sampling configuration files are included under `examples/`: # 1D trajectories and WHAM reconstruction python examples/langevin_umbrella.py --config examples/umbrella_config_1d.yaml python -m pywham.wham1d examples/wham1d_config.yaml +python -m pywham.bwham examples/bwham_config.yaml # 2D trajectories on the Müller-Brown surface and WHAM2D reconstruction python examples/langevin_umbrella_2d.py --config examples/umbrella_config_2d.yaml diff --git a/examples/bwham_config.yaml b/examples/bwham_config.yaml new file mode 100644 index 0000000..c6b6ace --- /dev/null +++ b/examples/bwham_config.yaml @@ -0,0 +1,14 @@ +# bwham-config.yml +hist_min: -1.0 +hist_max: 1.0 +num_bins: 50 +tolerance: 1e-5 +temperature: 1.0 +numpad: 0 +metadata_file: examples/output_1D/umbrella_metadata.txt +freefile: examples/output_1D/bayes.free +num_samples: 400 +burn_in: 100 +thinning: 2 +# optional Bayesian prior strength +# dirichlet_alpha: 1.0 diff --git a/pyproject.toml b/pyproject.toml index 461a649..720f1b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ build-backend = "setuptools.build_meta" wham = "pywham.wham1d:main" "wham-2d" = "pywham.wham2d:main" reweight = "pywham.reweight:main" +bwham = "pywham.bwham:main" [tool.setuptools.packages.find] where = ["src"] diff --git a/src/pywham/__init__.py b/src/pywham/__init__.py index 37bd5f3..fd4c13b 100644 --- a/src/pywham/__init__.py +++ b/src/pywham/__init__.py @@ -1,9 +1,37 @@ """Python translation of the WHAM C implementations.""" -from .visualization import plot_free_energy_1d, plot_free_energy_2d, save_free_energy_plots +from importlib import import_module +from typing import Any __all__ = [ + "BayesWHAM", + "BayesWhamConfig", + "build_bwham_config", "plot_free_energy_1d", "plot_free_energy_2d", "save_free_energy_plots", ] + +_LAZY_ATTRS = { + "BayesWHAM": (".bwham", "BayesWHAM"), + "BayesWhamConfig": (".bwham", "BayesWhamConfig"), + "build_bwham_config": (".bwham", "build_config"), + "plot_free_energy_1d": (".visualization", "plot_free_energy_1d"), + "plot_free_energy_2d": (".visualization", "plot_free_energy_2d"), + "save_free_energy_plots": (".visualization", "save_free_energy_plots"), +} + + +def __getattr__(name: str) -> Any: + try: + module_name, attr_name = _LAZY_ATTRS[name] + except KeyError: + raise AttributeError(f"module 'pywham' has no attribute {name!r}") from None + module = import_module(module_name, __name__) + attr = getattr(module, attr_name) + globals()[name] = attr + return attr + + +def __dir__() -> list[str]: + return sorted(__all__) diff --git a/src/pywham/bwham.py b/src/pywham/bwham.py new file mode 100644 index 0000000..20e65ac --- /dev/null +++ b/src/pywham/bwham.py @@ -0,0 +1,640 @@ +"""Bayesian WHAM driver built on the WHAM translations.""" + +from __future__ import annotations + +import math +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import List, Sequence + +import numpy as np +import yaml + +from .structures import HistGroup1D, HistGroup2D, Histogram1D, Histogram2D +from .wham1d import Wham1D, Wham1DConfig, parse_periodic as parse_periodic_1d, parse_units +from .wham2d import Wham2D, Wham2DConfig, parse_periodic as parse_periodic_2d + + +def _log_message(message: str) -> None: + print(f"[BayesWHAM] {message}") + + +@dataclass +class BayesWhamConfig: + """Configuration for running BayesWHAM.""" + + dimension: int + base_config: Wham1DConfig | Wham2DConfig + num_samples: int = 200 + burn_in: int = 50 + thinning: int = 1 + dirichlet_alpha: float = 1.0 + + +@dataclass +class BayesResult1D: + coordinate: List[float] + map_free_energy: List[float] + map_probabilities: List[float] + mean_free_energy: List[float] + std_free_energy: List[float] + mean_probabilities: List[float] + std_probabilities: List[float] + map_window_free: List[float] + mean_window_free: List[float] + std_window_free: List[float] + + +@dataclass +class BayesResult2D: + coordinate_x: List[List[float]] + coordinate_y: List[List[float]] + map_free_energy: List[List[float]] + map_probabilities: List[List[float]] + mean_free_energy: List[List[float]] + std_free_energy: List[List[float]] + mean_probabilities: List[List[float]] + std_probabilities: List[List[float]] + map_window_free: List[float] + mean_window_free: List[float] + std_window_free: List[float] + + +class BayesWHAM: + """Lightweight Bayesian wrapper over the WHAM translations. + + The implementation follows the workflow in the reference BayesWHAM script + by drawing Dirichlet-resampled histograms to approximate the posterior over + the free energy surface. We reuse the ingestion and iteration logic from the + WHAM translators to preserve trajectory handling and outputs. + """ + + def __init__(self, config: BayesWhamConfig): + self.config = config + self._log( + f"Initialized BayesWHAM for {config.dimension}D with {config.num_samples} samples " + f"(burn-in {config.burn_in}, thinning {config.thinning})." + ) + + def _log(self, message: str) -> None: + _log_message(message) + + def _posterior_samples(self) -> int: + self._log( + f"Preparing posterior sampling plan: total draws={self.config.num_samples}, " + f"burn-in={self.config.burn_in}, thinning={self.config.thinning}." + ) + usable = max(self.config.num_samples - self.config.burn_in, 0) + if usable == 0: + self._log("No usable samples remain after burn-in. Posterior sampling disabled.") + return 0 + sample_count = (usable + self.config.thinning - 1) // self.config.thinning + self._log(f"Posterior sampling will keep {sample_count} draws.") + return sample_count + + # --- 1D helpers ----------------------------------------------------- + def _resample_hist1d(self, hist: Histogram1D, rng: np.random.Generator) -> Histogram1D: + size = hist.last - hist.first + 1 + alpha = np.asarray(hist.data, dtype=float) + float(self.config.dirichlet_alpha) + weights = rng.dirichlet(alpha, size=1)[0] + counts = weights * float(hist.num_points) + return Histogram1D( + first=hist.first, + last=hist.last, + num_points=hist.num_points, + num_mc_samples=hist.num_mc_samples, + data=counts.tolist(), + cumulative=[0.0 for _ in range(size)], + ) + + def _clone_group_1d(self, base: HistGroup1D) -> HistGroup1D: + return HistGroup1D( + num_windows=base.num_windows, + bias_locations=list(base.bias_locations), + spring_constants=list(base.spring_constants), + free_energies=list(base.free_energies), + previous_free_energies=list(base.previous_free_energies), + temperatures=list(base.temperatures), + partitions=list(base.partitions), + histograms=[Histogram1D(0, 0, 0, 0) for _ in base.histograms], + ) + + def _load_1d(self, wham: Wham1D) -> tuple[HistGroup1D, bool, list]: + self._log(f"Loading 1D metadata from {wham.config.metadata_path}") + lines = wham.config.metadata_path.read_text(encoding="utf-8").splitlines() + num_windows = wham.get_numwindows(lines) + print(f"#Number of windows = {num_windows}") + self._log(f"Metadata reports {num_windows} windows.") + hist_group = wham.make_hist_group(num_windows) + count_windows, have_temp, entries = wham.read_metadata(lines, hist_group) + assert count_windows == hist_group.num_windows + if not have_temp: + self._log("No explicit temperatures detected; populating temperatures from configuration kT.") + for i in range(hist_group.num_windows): + hist_group.temperatures[i] = wham.config.kT + self._log("Histogram group populated for 1D run.") + return hist_group, have_temp, entries + + def _run_wham_1d( + self, wham: Wham1D, hist_group: HistGroup1D, have_energy: bool, log_iterations: bool = False + ) -> tuple[List[float], List[float]]: + self._log("Starting WHAM 1D iteration loop.") + probabilities = [0.0 for _ in range(wham.config.num_bins)] + iteration = 0 + first = True + printed_iteration = False + while not wham.is_converged(hist_group) or first: + first = False + wham.save_free(hist_group) + wham.wham_iteration(hist_group, probabilities, have_energy) + iteration += 1 + if log_iterations and iteration % 10 == 0: + error = wham.convergence_error(hist_group) + print(f"# Iteration {iteration:8d} | error {error:12.6e}", end="\r", flush=True) + printed_iteration = True + if log_iterations and iteration % 100 == 0: + free_snapshot, _ = wham.calc_free(probabilities) + wham._write_iteration_snapshot(iteration, free_snapshot, probabilities, hist_group.free_energies) + if iteration >= 100000: + print(f"Too many iterations: {iteration}") + break + if log_iterations and printed_iteration: + print() + total = sum(probabilities) + if total: + probabilities = [p / total for p in probabilities] + free_energy, _ = wham.calc_free(probabilities) + self._log(f"Completed WHAM 1D loop after {iteration} iterations.") + return free_energy, probabilities + + def _bayes_samples_1d( + self, + wham: Wham1D, + base_group: HistGroup1D, + have_energy: bool, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + sample_count = self._posterior_samples() + if sample_count == 0: + self._log("Skipping Bayesian sampling for 1D because no samples are requested.") + return np.empty((0, wham.config.num_bins)), np.empty((0, wham.config.num_bins)), np.empty((0, base_group.num_windows)) + self._log(f"Running Bayesian sampling for 1D histograms with {self.config.num_samples} total draws.") + rng = np.random.default_rng() + free_samples = np.zeros((sample_count, wham.config.num_bins), dtype=float) + prob_samples = np.zeros_like(free_samples) + window_samples = np.zeros((sample_count, base_group.num_windows), dtype=float) + + base_histograms = list(base_group.histograms) + idx = 0 + for draw in range(self.config.num_samples): + self._log(f"Dirichlet resampling draw {draw + 1}/{self.config.num_samples}.") + resampled = self._clone_group_1d(base_group) + resampled.histograms = [self._resample_hist1d(hist, rng) for hist in base_histograms] + free_energy, probabilities = self._run_wham_1d(wham, resampled, have_energy) + if draw < self.config.burn_in or (draw - self.config.burn_in) % self.config.thinning != 0: + continue + accepted_idx = idx + 1 + free_samples[idx, :] = np.asarray(free_energy, dtype=float) + prob_samples[idx, :] = np.asarray(probabilities, dtype=float) + window_samples[idx, :] = np.asarray(resampled.free_energies, dtype=float) + self._log(f"Accepted posterior sample {accepted_idx}/{sample_count}.") + idx += 1 + self._log("Completed Bayesian sampling for 1D histograms.") + return free_samples[:idx, :], prob_samples[:idx, :], window_samples[:idx, :] + + def _format_output_1d(self, wham: Wham1D, result: BayesResult1D) -> None: + self._log(f"Writing 1D BayesWHAM summary to {wham.config.freefile_path}") + with wham.config.freefile_path.open("w", encoding="utf-8") as freefile: + freefile.write("#Coor\tFree\tProb\tMeanFree\tStdFree\tMeanProb\tStdProb\n") + for coor, f_map, p_map, f_mean, f_std, p_mean, p_std in zip( + result.coordinate, + result.map_free_energy, + result.map_probabilities, + result.mean_free_energy, + result.std_free_energy, + result.mean_probabilities, + result.std_probabilities, + ): + freefile.write( + f"{coor:.6f}\t{f_map:.6f}\t{p_map:.6e}\t{f_mean:.6f}\t{f_std:.6f}\t{p_mean:.6e}\t{p_std:.6e}\n" + ) + freefile.write("\n#Window\tFree (free energy units)\tMean\tStd\n") + for idx, (f_map, f_mean, f_std) in enumerate( + zip(result.map_window_free, result.mean_window_free, result.std_window_free) + ): + freefile.write(f"#{idx}\t{f_map:.6f}\t{f_mean:.6f}\t{f_std:.6f}\n") + + # --- 2D helpers ----------------------------------------------------- + def _resample_hist2d(self, hist: Histogram2D, rng: np.random.Generator) -> Histogram2D: + size_x = hist.last_x - hist.first_x + 1 + size_y = hist.last_y - hist.first_y + 1 + flat = np.asarray(hist.data, dtype=float).reshape(size_x * size_y) + alpha = flat + float(self.config.dirichlet_alpha) + weights = rng.dirichlet(alpha, size=1)[0] + counts = (weights * float(hist.num_points)).reshape((size_x, size_y)) + return Histogram2D( + first_x=hist.first_x, + last_x=hist.last_x, + first_y=hist.first_y, + last_y=hist.last_y, + num_points=hist.num_points, + num_mc_samples=hist.num_mc_samples, + data=counts.tolist(), + cumulative=[0.0 for _ in range(size_x * size_y)], + ) + + def _clone_group_2d(self, base: HistGroup2D) -> HistGroup2D: + return HistGroup2D( + num_windows=base.num_windows, + bias_locations=[list(row) for row in base.bias_locations], + spring_x=list(base.spring_x), + spring_y=list(base.spring_y), + free_energies=list(base.free_energies), + previous_free_energies=list(base.previous_free_energies), + temperatures=list(base.temperatures), + partitions=list(base.partitions), + histograms=[Histogram2D(0, 0, 0, 0, 0, 0) for _ in base.histograms], + ) + + def _load_2d(self, wham: Wham2D): + self._log(f"Loading 2D metadata from {wham.config.metadata_path}") + lines = wham.config.metadata_path.read_text(encoding="utf-8").splitlines() + num_windows = wham.get_numwindows(lines) + print(f"#Number of windows = {num_windows}") + self._log(f"Metadata reports {num_windows} 2D windows.") + + mask = None + if wham.config.use_mask: + mask = [[0 for _ in range(wham.config.num_bins_y)] for _ in range(wham.config.num_bins_x)] + + hist_group = wham.make_hist_group(num_windows) + count_windows, have_temp, entries = wham.read_metadata(lines, hist_group, wham.config.use_mask, mask) + assert count_windows == hist_group.num_windows + if not have_temp: + self._log("No explicit window temperatures in metadata; applying config temperature.") + for i in range(hist_group.num_windows): + hist_group.temperatures[i] = wham.config.kT + for i in range(hist_group.num_windows): + hist_group.free_energies[i] = 1.0 + hist_group.previous_free_energies[i] = 1.0 + self._log("Histogram group populated for 2D run.") + return hist_group, have_temp, entries, mask + + def _run_wham_2d( + self, + wham: Wham2D, + hist_group: HistGroup2D, + have_energy: bool, + mask: list[list[int]] | None, + log_iterations: bool = False, + ) -> tuple[List[List[float]], List[List[float]]]: + self._log("Starting WHAM 2D iteration loop.") + dtype = np.float32 if wham.config.use_float32 else np.float64 + x_grid, y_grid = wham._coordinate_grids(dtype) + + num_lookup = np.zeros((wham.config.num_bins_x, wham.config.num_bins_y), dtype=dtype) + for hist in hist_group.histograms: + if hist.num_points == 0: + continue + x_slice = slice(hist.first_x, hist.last_x + 1) + y_slice = slice(hist.first_y, hist.last_y + 1) + num_lookup[x_slice, y_slice] += np.asarray(hist.data, dtype=dtype) + + bias_lookup = wham._build_bias_lookup(hist_group, x_grid, y_grid, dtype) + + prob = np.zeros((wham.config.num_bins_x, wham.config.num_bins_y), dtype=dtype) + iteration = 0 + first = True + converged = False + printed_iteration = False + while not converged or first: + first = False + wham.save_free(hist_group) + wham.wham_iteration(hist_group, prob, have_energy, wham.config.use_mask, mask, bias_lookup, num_lookup) + iteration += 1 + + epsilon = float(np.finfo(dtype).tiny) + hist_group.free_energies = [max(val, epsilon) for val in hist_group.free_energies] + hist_group.previous_free_energies = [max(val, epsilon) for val in hist_group.previous_free_energies] + logged_current = [ + hist_group.temperatures[i] * math.log(hist_group.free_energies[i]) for i in range(hist_group.num_windows) + ] + logged_previous = [ + hist_group.temperatures[i] * math.log(hist_group.previous_free_energies[i]) + for i in range(hist_group.num_windows) + ] + converged = wham.is_converged(hist_group, logged_current, logged_previous) + if log_iterations and iteration % 10 == 0: + error = wham.convergence_error(logged_current, logged_previous) + print(f"# Iteration {iteration:8d} | error {error:12.6e}", end="\r", flush=True) + printed_iteration = True + if log_iterations and iteration % 100 == 0: + free_snapshot = wham.calc_free(prob.tolist(), wham.config.use_mask, mask) + wham._write_iteration_snapshot(iteration, free_snapshot, prob, hist_group.free_energies) + if iteration >= 100000: + print(f"Too many iterations: {iteration}") + break + + free_energy = np.asarray(wham.calc_free(prob.tolist(), wham.config.use_mask, mask), dtype=prob.dtype) + total = float(np.sum(prob)) + if total > 0: + prob /= total + if log_iterations and printed_iteration: + print() + self._log(f"Completed WHAM 2D loop after {iteration} iterations.") + return free_energy.tolist(), prob.tolist() + + def _bayes_samples_2d( + self, wham: Wham2D, base_group: HistGroup2D, have_energy: bool, mask: list[list[int]] | None + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + sample_count = self._posterior_samples() + if sample_count == 0: + self._log("Skipping Bayesian sampling for 2D because no samples are requested.") + return np.empty((0, wham.config.num_bins_x, wham.config.num_bins_y)), np.empty( + (0, wham.config.num_bins_x, wham.config.num_bins_y) + ), np.empty((0, base_group.num_windows)) + self._log(f"Running Bayesian sampling for 2D histograms with {self.config.num_samples} total draws.") + rng = np.random.default_rng() + free_samples = np.zeros((sample_count, wham.config.num_bins_x, wham.config.num_bins_y), dtype=float) + prob_samples = np.zeros_like(free_samples) + window_samples = np.zeros((sample_count, base_group.num_windows), dtype=float) + + base_histograms = list(base_group.histograms) + idx = 0 + for draw in range(self.config.num_samples): + self._log(f"Dirichlet resampling draw {draw + 1}/{self.config.num_samples} for 2D histograms.") + resampled = self._clone_group_2d(base_group) + resampled.histograms = [self._resample_hist2d(hist, rng) for hist in base_histograms] + free_energy, probabilities = self._run_wham_2d(wham, resampled, have_energy, mask) + if draw < self.config.burn_in or (draw - self.config.burn_in) % self.config.thinning != 0: + continue + accepted_idx = idx + 1 + free_samples[idx, :, :] = np.asarray(free_energy, dtype=float) + prob_samples[idx, :, :] = np.asarray(probabilities, dtype=float) + window_samples[idx, :] = np.asarray(resampled.free_energies, dtype=float) + self._log(f"Accepted posterior sample {accepted_idx}/{sample_count} for 2D histograms.") + idx += 1 + self._log("Completed Bayesian sampling for 2D histograms.") + return free_samples[:idx, :, :], prob_samples[:idx, :, :], window_samples[:idx, :] + + def _format_output_2d(self, wham: Wham2D, result: BayesResult2D) -> None: + self._log(f"Writing 2D BayesWHAM summary to {wham.config.freefile_path}") + with wham.config.freefile_path.open("w", encoding="utf-8") as freefile: + freefile.write("#X\tY\tFree\tProb\tMeanFree\tStdFree\tMeanProb\tStdProb\n") + for i in range(wham.config.num_bins_x): + for j in range(wham.config.num_bins_y): + freefile.write( + f"{result.coordinate_x[i][j]:.6f}\t{result.coordinate_y[i][j]:.6f}\t{result.map_free_energy[i][j]:.6f}\t" + f"{result.map_probabilities[i][j]:.6e}\t{result.mean_free_energy[i][j]:.6f}\t{result.std_free_energy[i][j]:.6f}\t" + f"{result.mean_probabilities[i][j]:.6e}\t{result.std_probabilities[i][j]:.6e}\n" + ) + freefile.write("\n#Window\tFree (free energy units)\tMean\tStd\n") + for idx, (f_map, f_mean, f_std) in enumerate( + zip(result.map_window_free, result.mean_window_free, result.std_window_free) + ): + freefile.write(f"#{idx}\t{f_map:.6f}\t{f_mean:.6f}\t{f_std:.6f}\n") + + # --- User facing ---------------------------------------------------- + def run(self) -> None: + self._log("Starting BayesWHAM run.") + if self.config.dimension == 1: + self._log("Executing 1D workflow.") + wham = Wham1D(self.config.base_config) # type: ignore[arg-type] + hist_group, have_energy, entries = self._load_1d(wham) + map_free, map_prob = self._run_wham_1d(wham, hist_group, have_energy, log_iterations=True) + samples_free, samples_prob, window_samples = self._bayes_samples_1d(wham, hist_group, have_energy) + mean_free = samples_free.mean(axis=0) if samples_free.size else np.asarray(map_free, dtype=float) + std_free = samples_free.std(axis=0) if samples_free.size else np.zeros_like(map_free, dtype=float) + mean_prob = samples_prob.mean(axis=0) if samples_prob.size else np.asarray(map_prob, dtype=float) + std_prob = samples_prob.std(axis=0) if samples_prob.size else np.zeros_like(map_prob, dtype=float) + mean_window = window_samples.mean(axis=0) if window_samples.size else np.asarray(hist_group.free_energies) + std_window = window_samples.std(axis=0) if window_samples.size else np.zeros_like(hist_group.free_energies) + + coordinates = [wham.calc_coor(i) for i in range(wham.config.num_bins)] + result = BayesResult1D( + coordinate=coordinates, + map_free_energy=map_free, + map_probabilities=map_prob, + mean_free_energy=mean_free.tolist(), + std_free_energy=std_free.tolist(), + mean_probabilities=mean_prob.tolist(), + std_probabilities=std_prob.tolist(), + map_window_free=list(hist_group.free_energies), + mean_window_free=mean_window.tolist(), + std_window_free=std_window.tolist(), + ) + self._log("Writing outputs and auxiliary data for 1D run.") + self._format_output_1d(wham, result) + wham._write_aux_data(entries, hist_group, result.map_window_free, window_samples.tolist()) + self._log("1D BayesWHAM run completed.") + return + + self._log("Executing 2D workflow.") + wham2d = Wham2D(self.config.base_config) # type: ignore[arg-type] + hist_group2d, have_energy2d, entries2d, mask = self._load_2d(wham2d) + map_free_2d, map_prob_2d = self._run_wham_2d(wham2d, hist_group2d, have_energy2d, mask, log_iterations=True) + samples_free_2d, samples_prob_2d, window_samples_2d = self._bayes_samples_2d( + wham2d, hist_group2d, have_energy2d, mask + ) + + mean_free_2d = samples_free_2d.mean(axis=0) if samples_free_2d.size else np.asarray(map_free_2d, dtype=float) + std_free_2d = samples_free_2d.std(axis=0) if samples_free_2d.size else np.zeros_like(map_free_2d, dtype=float) + mean_prob_2d = samples_prob_2d.mean(axis=0) if samples_prob_2d.size else np.asarray(map_prob_2d, dtype=float) + std_prob_2d = samples_prob_2d.std(axis=0) if samples_prob_2d.size else np.zeros_like(map_prob_2d, dtype=float) + mean_window_2d = ( + window_samples_2d.mean(axis=0) if window_samples_2d.size else np.asarray(hist_group2d.free_energies, dtype=float) + ) + std_window_2d = window_samples_2d.std(axis=0) if window_samples_2d.size else np.zeros_like(hist_group2d.free_energies) + + x_coords = [] + y_coords = [] + for i in range(wham2d.config.num_bins_x): + row_x = [] + row_y = [] + for j in range(wham2d.config.num_bins_y): + coor = wham2d.calc_coor(i, j) + row_x.append(coor[0]) + row_y.append(coor[1]) + x_coords.append(row_x) + y_coords.append(row_y) + + result2d = BayesResult2D( + coordinate_x=x_coords, + coordinate_y=y_coords, + map_free_energy=map_free_2d, + map_probabilities=map_prob_2d, + mean_free_energy=mean_free_2d.tolist(), + std_free_energy=std_free_2d.tolist(), + mean_probabilities=mean_prob_2d.tolist(), + std_probabilities=std_prob_2d.tolist(), + map_window_free=list(hist_group2d.free_energies), + mean_window_free=mean_window_2d.tolist(), + std_window_free=std_window_2d.tolist(), + ) + self._log("Writing outputs and auxiliary data for 2D run.") + self._format_output_2d(wham2d, result2d) + wham2d._write_aux_data(entries2d, hist_group2d, result2d.map_window_free, window_samples_2d.tolist()) + self._log("2D BayesWHAM run completed.") + + +# --- configuration ------------------------------------------------------ + + +def _resolve_relative(path_raw: str, base_dir: Path) -> Path: + path = Path(path_raw) + if not path.is_absolute(): + path = (base_dir / path).resolve() + return path + + +def build_config(yaml_path: Path) -> BayesWhamConfig: + if not yaml_path.exists(): + raise FileNotFoundError(f"YAML configuration file not found: {yaml_path}") + + _log_message(f"Building BayesWHAM configuration from {yaml_path}") + config_raw = yaml.safe_load(yaml_path.read_text(encoding="utf-8")) + if not isinstance(config_raw, dict): + raise ValueError("YAML configuration must define a mapping of parameters") + + num_samples = int(config_raw.get("num_samples", 200)) + burn_in = int(config_raw.get("burn_in", 50)) + thinning = int(config_raw.get("thinning", 1)) + dirichlet_alpha = float(config_raw.get("dirichlet_alpha", 1.0)) + + if "hist_min_y" in config_raw: + _log_message("Detected 2D configuration in YAML.") + k_B = parse_units(config_raw.get("units")) + periodic_x, period_x = parse_periodic_2d(config_raw, "x") + periodic_y, period_y = parse_periodic_2d(config_raw, "y") + required_fields = [ + "hist_min_x", + "hist_max_x", + "num_bins_x", + "hist_min_y", + "hist_max_y", + "num_bins_y", + "tolerance", + "temperature", + "numpad", + "metadata_file", + "freefile", + "use_mask", + ] + for field in required_fields: + if field not in config_raw: + raise ValueError(f"Missing required configuration field: {field}") + metadata_path = _resolve_relative(str(config_raw["metadata_file"]), yaml_path.parent) + metadata_dir = metadata_path.parent + freefile_path = _resolve_relative(str(config_raw["freefile"]), metadata_dir) + aux_data_path = None + if "aux_data_file" in config_raw: + aux_data_path = _resolve_relative(str(config_raw["aux_data_file"]), metadata_dir) + + base = Wham2DConfig( + hist_min_x=float(config_raw["hist_min_x"]), + hist_max_x=float(config_raw["hist_max_x"]), + num_bins_x=int(config_raw["num_bins_x"]), + hist_min_y=float(config_raw["hist_min_y"]), + hist_max_y=float(config_raw["hist_max_y"]), + num_bins_y=int(config_raw["num_bins_y"]), + tolerance=float(config_raw["tolerance"]), + temperature=float(config_raw["temperature"]), + numpad=int(config_raw["numpad"]), + metadata_path=metadata_path, + freefile_path=freefile_path, + use_mask=bool(config_raw["use_mask"]), + periodic_x=periodic_x, + period_x=period_x, + periodic_y=period_y, + period_y=period_y, + k_B=k_B, + use_float32=bool(config_raw.get("use_float32", False)), + bias_chunk_size=(int(config_raw["bias_chunk_size"]) if "bias_chunk_size" in config_raw else None), + num_mc_runs=0, + mc_seed=None, + mc_workers=None, + freefile_error_path=( + _resolve_relative(str(config_raw["freefile_error"]), metadata_dir) + if "freefile_error" in config_raw + else None + ), + aux_data_path=aux_data_path, + ) + _log_message("Finished assembling 2D BayesWHAM configuration.") + return BayesWhamConfig( + dimension=2, + base_config=base, + num_samples=num_samples, + burn_in=burn_in, + thinning=thinning, + dirichlet_alpha=dirichlet_alpha, + ) + + _log_message("Detected 1D configuration in YAML. Assembling base parameters.") + k_B = parse_units(config_raw.get("units")) + periodic, period = parse_periodic_1d(config_raw) + required_fields = [ + "hist_min", + "hist_max", + "num_bins", + "tolerance", + "temperature", + "numpad", + "metadata_file", + "freefile", + ] + for field in required_fields: + if field not in config_raw: + raise ValueError(f"Missing required configuration field: {field}") + metadata_path = _resolve_relative(str(config_raw["metadata_file"]), yaml_path.parent) + metadata_dir = metadata_path.parent + freefile_path = _resolve_relative(str(config_raw["freefile"]), metadata_dir) + aux_data_path = None + if "aux_data_file" in config_raw: + aux_data_path = _resolve_relative(str(config_raw["aux_data_file"]), metadata_dir) + + base = Wham1DConfig( + hist_min=float(config_raw["hist_min"]), + hist_max=float(config_raw["hist_max"]), + num_bins=int(config_raw["num_bins"]), + tolerance=float(config_raw["tolerance"]), + temperature=float(config_raw["temperature"]), + numpad=int(config_raw["numpad"]), + metadata_path=metadata_path, + freefile_path=freefile_path, + periodic=periodic, + period=period, + k_B=k_B, + num_mc_runs=0, + mc_seed=None, + mc_workers=None, + ingest_workers=None, + aux_data_path=aux_data_path, + ) + _log_message("Finished assembling 1D BayesWHAM configuration.") + return BayesWhamConfig( + dimension=1, + base_config=base, + num_samples=num_samples, + burn_in=burn_in, + thinning=thinning, + dirichlet_alpha=dirichlet_alpha, + ) + + +# --- CLI --------------------------------------------------------------- + + +def main(argv: Sequence[str] | None = None) -> None: + args = sys.argv[1:] if argv is None else list(argv) + if len(args) != 1: + raise ValueError("bwham expects a single argument: path to a YAML configuration file") + yaml_path = Path(args[0]) + print(f"# Loading configuration from {yaml_path}") + config = build_config(yaml_path) + bwham = BayesWHAM(config) + bwham.run() + + +if __name__ == "__main__": + main() diff --git a/src/pywham/reweight.py b/src/pywham/reweight.py index 54b5eac..65aa659 100644 --- a/src/pywham/reweight.py +++ b/src/pywham/reweight.py @@ -4,7 +4,7 @@ import math import sys -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import List, Sequence @@ -18,6 +18,8 @@ class WindowRecord: bias_center: List[float] spring_constants: List[float] num_samples: int + raw_samples: int + dropped_samples: List[int] = field(default_factory=list) @dataclass @@ -108,6 +110,8 @@ def run(self) -> ReweightResult: p_map = np.zeros(total_bins_proj, dtype=float) p_mh = np.zeros((total_bins_proj, f_mh.shape[0]), dtype=float) + skipped_map = 0 + skipped_mh = 0 for sim_index, (traj_i, traj_proj_i) in enumerate(zip(traj, traj_proj), start=1): for sample_idx in range(traj_i.shape[0]): proj_coords = traj_proj_i[sample_idx] @@ -124,17 +128,26 @@ def run(self) -> ReweightResult: weights = bias_lookup[:, idx_umb] denom_map = float(np.dot(map_prefactors, weights)) if denom_map == 0.0: - raise ZeroDivisionError("Encountered zero denominator while computing MAP weight") - p_map[idx_proj] += 1.0 / denom_map + skipped_map += 1 + else: + p_map[idx_proj] += 1.0 / denom_map if f_mh.size > 0: denom_mh = mh_prefactors @ weights - if np.any(denom_mh == 0.0): - raise ZeroDivisionError("Encountered zero denominator while computing MH weights") - p_mh[idx_proj, :] += 1.0 / denom_mh + mask = denom_mh != 0.0 + skipped_mh += int(mask.size - np.count_nonzero(mask)) + if np.any(mask): + p_mh[idx_proj, mask] += 1.0 / denom_mh[mask] print(f"# Completed projection for simulation {sim_index}") + if skipped_map: + print(f"# Skipped {skipped_map} samples with zero MAP denominators; resulting bins will remain empty") + if skipped_mh: + print( + f"# Skipped {skipped_mh} MH weight evaluations with zero denominators; affected bins will remain empty" + ) + p_map = _normalize_vector(p_map) p_mh = _normalize_matrix_columns(p_mh) @@ -185,11 +198,17 @@ def _load_umbrella_trajectories(self) -> List[np.ndarray]: f"{record.trajectory} must contain at least {dim + 1} columns " "(an index/time column plus the umbrella coordinates)" ) - if record.num_samples > 0 and data.shape[0] != record.num_samples: + expected_raw = record.raw_samples if record.raw_samples > 0 else data.shape[0] + if data.shape[0] != expected_raw: + raise ValueError( + f"{record.trajectory} contains {data.shape[0]} samples, expected {expected_raw}" + ) + trimmed = _filter_samples(data[:, 1 : dim + 1], record.dropped_samples) + if record.num_samples > 0 and trimmed.shape[0] != record.num_samples: raise ValueError( - f"{record.trajectory} contains {data.shape[0]} samples, expected {record.num_samples}" + f"{record.trajectory} retains {trimmed.shape[0]} samples after filtering, expected {record.num_samples}" ) - trajectories.append(data[:, 1 : dim + 1]) + trajectories.append(trimmed) return trajectories def _load_projection_trajectories(self) -> List[np.ndarray]: @@ -200,10 +219,14 @@ def _load_projection_trajectories(self) -> List[np.ndarray]: lines = metadata_path.read_text(encoding="utf-8").splitlines() trajectories: List[np.ndarray] = [] dim_proj = len(self.aux.projection_hist_edges) + window_index = 0 for entry in lines: entry = entry.strip() if not entry or entry.startswith("#"): continue + if window_index >= len(self.aux.windows): + raise ValueError("Projection metadata specifies more trajectories than umbrella windows") + window_record = self.aux.windows[window_index] parts = entry.split() path = Path(parts[0]) if not path.is_absolute(): @@ -215,7 +238,16 @@ def _load_projection_trajectories(self) -> List[np.ndarray]: f"{path} must contain at least {dim_proj + 1} columns " "(an index/time column plus the projection coordinates)" ) - trajectories.append(data[:, 1 : dim_proj + 1]) + expected_raw = window_record.raw_samples if window_record.raw_samples > 0 else data.shape[0] + if data.shape[0] != expected_raw: + raise ValueError(f"{path} contains {data.shape[0]} samples, expected {expected_raw}") + filtered = _filter_samples(data[:, 1 : dim_proj + 1], window_record.dropped_samples) + if filtered.shape[0] != window_record.num_samples: + raise ValueError( + f"{path} retains {filtered.shape[0]} samples after filtering, expected {window_record.num_samples}" + ) + trajectories.append(filtered) + window_index += 1 if len(trajectories) != len(self.aux.windows): raise ValueError("Number of projection trajectories does not match number of umbrella windows") return trajectories @@ -233,10 +265,10 @@ def _bias_lookup( diff = delta[:, :, dim_index] if periodic: period = periods[dim_index] - diff = np.abs(diff) - delta[:, :, dim_index] = np.minimum.reduce( - [diff, np.abs(diff + period), np.abs(diff - period)] - ) + # Apply minimum image convention: wrap to [-period/2, period/2] + # This handles separations spanning multiple periods correctly + diff = diff - period * np.round(diff / period) + delta[:, :, dim_index] = np.abs(diff) else: delta[:, :, dim_index] = np.abs(diff) energy = 0.5 * np.sum(forces[:, None, :] * delta * delta, axis=2) @@ -281,7 +313,19 @@ def load_aux_data(yaml_path: Path) -> AuxData: bias_center = _ensure_float_list(entry.get("bias_center"), dim, "bias_center") springs = _ensure_float_list(entry.get("spring_constants"), dim, "spring_constants") num_samples = int(entry.get("num_samples", 0)) - windows.append(WindowRecord(trajectory, bias_center, springs, num_samples)) + raw_samples = int(entry.get("raw_samples", num_samples)) + dropped_raw = entry.get("dropped_samples", []) + dropped_samples = [int(idx) for idx in dropped_raw] if dropped_raw else [] + windows.append( + WindowRecord( + trajectory=trajectory, + bias_center=bias_center, + spring_constants=springs, + num_samples=num_samples, + raw_samples=raw_samples, + dropped_samples=dropped_samples, + ) + ) map_values_raw = config.get("map_values") if map_values_raw is None: @@ -404,12 +448,31 @@ def _bin_volumes(widths: Sequence[np.ndarray]) -> np.ndarray: return volume.reshape(-1) +def _filter_samples(data: np.ndarray, dropped: Sequence[int]) -> np.ndarray: + if not dropped: + return data + mask = np.ones(data.shape[0], dtype=bool) + for idx in dropped: + if idx < 0 or idx >= data.shape[0]: + raise ValueError(f"Dropped sample index {idx} is out of bounds for trajectory with {data.shape[0]} samples") + mask[idx] = False + return data[mask] + + def _locate_bin(values: np.ndarray, edges: Sequence[np.ndarray]) -> tuple[int, ...] | None: subs: List[int] = [] for coord, axis_edges in zip(values, edges): idx = int(np.searchsorted(axis_edges, coord, side="right") - 1) - if idx < 0 or idx >= len(axis_edges) - 1: + # Handle out of bounds: below minimum + if idx < 0: return None + # Handle upper edge: include values exactly at maximum in last bin + if idx >= len(axis_edges) - 1: + # If coord equals the upper edge exactly, include it in the last bin + if coord == axis_edges[-1] and idx == len(axis_edges) - 1: + idx = len(axis_edges) - 2 + else: + return None subs.append(idx) return tuple(subs) diff --git a/src/pywham/wham1d.py b/src/pywham/wham1d.py index 8b6806e..6698481 100644 --- a/src/pywham/wham1d.py +++ b/src/pywham/wham1d.py @@ -6,7 +6,7 @@ import math import os import sys -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Iterable, List, Tuple @@ -106,9 +106,13 @@ def get_numwindows(self, metadata: Iterable[str]) -> int: return sum(1 for line in metadata if self.is_metadata(line)) @staticmethod - def read_data(filename: Path, have_energy: bool, config: Wham1DConfig) -> Tuple[list[float], int]: + def read_data( + filename: Path, have_energy: bool, config: Wham1DConfig + ) -> Tuple[list[float], int, List[int], int]: histogram = [0.0 for _ in range(config.num_bins)] count = 0 + dropped_indices: list[int] = [] + total_samples = 0 with filename.open("r", encoding="utf-8") as handle: for raw in handle: if raw.startswith("#"): @@ -126,6 +130,8 @@ def read_data(filename: Path, have_energy: bool, config: Wham1DConfig) -> Tuple[ _, value_s = parts[:2] value = float(value_s) energy = 0.0 + sample_index = total_samples + total_samples += 1 if config.hist_min < value < config.hist_max: index = int((value - config.hist_min) / config.bin_width) if have_energy: @@ -133,7 +139,9 @@ def read_data(filename: Path, have_energy: bool, config: Wham1DConfig) -> Tuple[ else: histogram[index] += 1.0 count += 1 - return histogram, count + else: + dropped_indices.append(sample_index) + return histogram, count, dropped_indices, total_samples def read_metadata(self, lines: Iterable[str], hist_group: HistGroup1D) -> Tuple[int, bool, List[MetadataEntry]]: entries: list[MetadataEntry] = [] @@ -203,6 +211,8 @@ def read_metadata(self, lines: Iterable[str], hist_group: HistGroup1D) -> Tuple[ hist_group.temperatures[entry.index] = -1.0 hist_group.histograms[entry.index] = result.histogram hist_group.partitions[entry.index] = result.partition + entry.raw_samples = result.raw_samples + entry.dropped_samples = result.dropped_samples for warning in result.warnings: print(warning) @@ -265,6 +275,8 @@ def _write_aux_data( "bias_center": [float(entry.loc)], "spring_constants": [float(entry.spring)], "num_samples": int(histogram.num_points), + "raw_samples": int(entry.raw_samples) if entry.raw_samples else int(histogram.num_points), + "dropped_samples": list(entry.dropped_samples), } ) aux_data = { @@ -310,6 +322,13 @@ def average_diff(self, hist_group: HistGroup1D) -> float: previous = np.asarray(hist_group.previous_free_energies, dtype=float) return float(np.mean(np.abs(current - previous))) + def convergence_error(self, hist_group: HistGroup1D) -> float: + current = np.asarray(hist_group.free_energies, dtype=float) + previous = np.asarray(hist_group.previous_free_energies, dtype=float) + if current.size == 0: + return 0.0 + return float(np.max(np.abs(current - previous))) + def _write_iteration_snapshot( self, iteration: int, free_energy: list[float], probabilities: list[float], free_energies: list[float] ) -> None: @@ -420,7 +439,7 @@ def run(self) -> None: self.wham_iteration(hist_group, probabilities, have_energy) iteration += 1 if iteration % 10 == 0: - error = self.average_diff(hist_group) + error = self.convergence_error(hist_group) print(f"# Iteration {iteration:8d} | error {error:12.6e}") if iteration % 100 == 0: free_energy, _ = self.calc_free(probabilities) @@ -570,6 +589,8 @@ class MetadataEntry: spring: float correl_time: float temp: float + raw_samples: int = 0 + dropped_samples: List[int] = field(default_factory=list) @dataclass @@ -580,6 +601,8 @@ class WindowLoadResult: min_nonzero: int max_nonzero: int warnings: list[str] + raw_samples: int + dropped_samples: list[int] def _build_histogram( @@ -636,7 +659,7 @@ def _build_histogram( def _load_window_data(entry: MetadataEntry, have_energy: bool, config: Wham1DConfig) -> WindowLoadResult: - histogram, num_points = Wham1D.read_data(entry.filename, have_energy, config) + histogram, num_points, dropped_samples, raw_samples = Wham1D.read_data(entry.filename, have_energy, config) if num_points < 0: raise OSError(f"Error trying to read {entry.filename}") @@ -651,6 +674,8 @@ def _load_window_data(entry: MetadataEntry, have_energy: bool, config: Wham1DCon min_nonzero=trimmed.first, max_nonzero=trimmed.last, warnings=warnings, + raw_samples=raw_samples, + dropped_samples=dropped_samples, ) diff --git a/src/pywham/wham2d.py b/src/pywham/wham2d.py index 3f589dc..f800a92 100644 --- a/src/pywham/wham2d.py +++ b/src/pywham/wham2d.py @@ -6,7 +6,7 @@ import math import os import sys -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Iterable, List, Tuple @@ -71,6 +71,8 @@ class MetadataEntry2D: springy: float correl_time: float temp: float + raw_samples: int = 0 + dropped_samples: List[int] = field(default_factory=list) class Wham2D: @@ -187,6 +189,8 @@ def _write_aux_data( "bias_center": [float(entry.locx), float(entry.locy)], "spring_constants": [float(entry.springx), float(entry.springy)], "num_samples": int(histogram.num_points), + "raw_samples": int(entry.raw_samples) if entry.raw_samples else int(histogram.num_points), + "dropped_samples": list(entry.dropped_samples), } ) aux_data = { @@ -214,9 +218,13 @@ def _write_aux_data( path.write_text(yaml.safe_dump(aux_data, sort_keys=False), encoding="utf-8") print(f"# Wrote auxiliary data to {path}") - def read_data(self, filename: Path, have_energy: bool, use_mask: bool, mask: List[List[int]] | None) -> int: + def read_data( + self, filename: Path, have_energy: bool, use_mask: bool, mask: List[List[int]] | None + ) -> Tuple[int, List[int], int]: self.clear_histogram() count = 0 + dropped_indices: list[int] = [] + total_samples = 0 with filename.open("r", encoding="utf-8") as handle: for raw in handle: if raw.startswith("#"): @@ -231,11 +239,13 @@ def read_data(self, filename: Path, have_energy: bool, use_mask: bool, mask: Lis energy = float(energy_s) else: if len(parts) < 3: - continue + raise ValueError(f"Failure reading {filename}: missing coordinate value") _, value_x_s, value_y_s = parts[:3] value_x = float(value_x_s) value_y = float(value_y_s) energy = 0.0 + sample_index = total_samples + total_samples += 1 if ( self.config.hist_min_x < value_x < self.config.hist_max_x and self.config.hist_min_y < value_y < self.config.hist_max_y @@ -249,7 +259,9 @@ def read_data(self, filename: Path, have_energy: bool, use_mask: bool, mask: Lis count += 1 if use_mask and mask is not None: mask[index_x][index_y] = 1 - return count + else: + dropped_indices.append(sample_index) + return count, dropped_indices, total_samples def read_metadata( self, lines: Iterable[str], hist_group: HistGroup2D, use_mask: bool, mask: List[List[int]] | None @@ -301,7 +313,7 @@ def read_metadata( ) ) - num_points = self.read_data(filename, have_temp, use_mask, mask) + num_points, dropped_samples, raw_samples = self.read_data(filename, have_temp, use_mask, mask) if num_points < 0: raise OSError(f"Error trying to read {filename}") @@ -333,6 +345,8 @@ def read_metadata( trimmed.cumulative[bin_index] /= total hist_group.histograms[current_window] = trimmed hist_group.partitions[current_window] = total + entries[current_window].raw_samples = raw_samples + entries[current_window].dropped_samples = dropped_samples current_window += 1 return current_window, have_temp, entries @@ -395,6 +409,14 @@ def average_diff(self, logged_current: List[float], logged_previous: List[float] error += abs(cur - prev) return error / float(len(logged_current)) if logged_current else 0.0 + def convergence_error(self, logged_current: List[float], logged_previous: List[float]) -> float: + max_error = 0.0 + for cur, prev in zip(logged_current, logged_previous): + diff = abs(cur - prev) + if diff > max_error: + max_error = diff + return max_error + def _write_iteration_snapshot( self, iteration: int, free_energy: list[list[float]], probabilities: np.ndarray, free_energies: list[float] ) -> None: @@ -616,7 +638,7 @@ def run(self) -> None: ] converged = self.is_converged(hist_group, logged_current, logged_previous) if iteration % 10 == 0: - error = self.average_diff(logged_current, logged_previous) + error = self.convergence_error(logged_current, logged_previous) print(f"# Iteration {iteration:8d} | error {error:12.6e}") if iteration % 100 == 0: free_ene = self.calc_free(prob, self.config.use_mask, mask) diff --git a/tests/test_bwham.py b/tests/test_bwham.py new file mode 100644 index 0000000..25145a0 --- /dev/null +++ b/tests/test_bwham.py @@ -0,0 +1,114 @@ +import yaml +import pytest +from pathlib import Path + +from pywham.bwham import BayesWHAM, build_config + + +def _write_simple_metadata(base_dir: Path) -> Path: + base_dir.mkdir(parents=True, exist_ok=True) + data_file = base_dir / "traj.dat" + data_file.write_text("0 0.5 0\n1 1.5 0\n", encoding="utf-8") + metadata = base_dir / "metadata.txt" + metadata.write_text("traj.dat 0.0 1.0 1.0 300\n", encoding="utf-8") + return metadata + + +def test_build_config_resolves_paths_1d(tmp_path: Path) -> None: + metadata_path = _write_simple_metadata(tmp_path / "data1d") + config_path = tmp_path / "config1d.yaml" + config_path.write_text( + yaml.safe_dump( + { + "hist_min": 0.0, + "hist_max": 2.0, + "num_bins": 2, + "tolerance": 1e-6, + "temperature": 300.0, + "numpad": 0, + "metadata_file": "data1d/metadata.txt", + "freefile": "free/output.txt", + "aux_data_file": "aux/aux.yaml", + } + ), + encoding="utf-8", + ) + + config = build_config(config_path) + + metadata_resolved = metadata_path.resolve() + base_dir = metadata_resolved.parent + assert config.base_config.metadata_path == metadata_resolved + assert config.base_config.freefile_path == (base_dir / "free/output.txt").resolve() + assert config.base_config.aux_data_path == (base_dir / "aux/aux.yaml").resolve() + + +def test_build_config_resolves_paths_2d(tmp_path: Path) -> None: + metadata_path = _write_simple_metadata(tmp_path / "data2d") + config_path = tmp_path / "config2d.yaml" + config_path.write_text( + yaml.safe_dump( + { + "hist_min_x": 0.0, + "hist_max_x": 1.0, + "num_bins_x": 2, + "hist_min_y": 0.0, + "hist_max_y": 1.0, + "num_bins_y": 2, + "tolerance": 1e-6, + "temperature": 300.0, + "numpad": 0, + "metadata_file": "data2d/metadata.txt", + "freefile": "free2d.txt", + "use_mask": False, + "aux_data_file": "aux/out.yaml", + } + ), + encoding="utf-8", + ) + + config = build_config(config_path) + + metadata_resolved = metadata_path.resolve() + base_dir = metadata_resolved.parent + assert config.base_config.metadata_path == metadata_resolved + assert config.base_config.freefile_path == (base_dir / "free2d.txt").resolve() + assert config.base_config.aux_data_path == (base_dir / "aux/out.yaml").resolve() + + +def test_bwham_writes_aux_data(tmp_path: Path) -> None: + metadata_path = _write_simple_metadata(tmp_path / "bayes") + config_path = tmp_path / "bayes.yaml" + config_path.write_text( + yaml.safe_dump( + { + "hist_min": 0.0, + "hist_max": 2.0, + "num_bins": 2, + "tolerance": 1e-6, + "temperature": 300.0, + "numpad": 0, + "metadata_file": str(metadata_path.relative_to(config_path.parent)), + "freefile": "free.txt", + "aux_data_file": "aux/aux.yaml", + "num_samples": 2, + "burn_in": 0, + "thinning": 1, + } + ), + encoding="utf-8", + ) + + config = build_config(config_path) + bwham_runner = BayesWHAM(config) + bwham_runner.run() + + aux_path = config.base_config.aux_data_path + assert aux_path is not None and aux_path.exists() + + aux = yaml.safe_load(aux_path.read_text(encoding="utf-8")) + assert aux["metadata_file"] == str(config.base_config.metadata_path) + assert aux["map_values"] == pytest.approx([0.0]) + assert len(aux["mh_samples"]) == 2 + assert aux["windows"][0]["trajectory"].endswith("traj.dat") + diff --git a/tests/test_wham1d.py b/tests/test_wham1d.py index 25eca46..e4f9504 100644 --- a/tests/test_wham1d.py +++ b/tests/test_wham1d.py @@ -95,8 +95,10 @@ def test_is_metadata_and_numwindows(wham: Wham1D) -> None: def test_read_data_energy_and_range(tmp_path: Path, wham: Wham1D) -> None: datafile = tmp_path / "data.dat" datafile.write_text("0 0.5 0\n1 1.5 0\n#2 0.2 0\n", encoding="utf-8") - histogram, count = Wham1D.read_data(datafile, have_energy=True, config=wham.config) + histogram, count, dropped, raw = Wham1D.read_data(datafile, have_energy=True, config=wham.config) assert count == 2 + assert raw == 2 + assert dropped == [] assert histogram == [1.0, 1.0] assert Wham1D._find_range(histogram) == (0, 1) @@ -108,6 +110,15 @@ def test_read_data_invalid_columns(tmp_path: Path, wham: Wham1D) -> None: Wham1D.read_data(datafile, have_energy=True, config=wham.config) +def test_read_data_tracks_dropped_samples(tmp_path: Path, wham: Wham1D) -> None: + datafile = tmp_path / "mixed.dat" + datafile.write_text("0 -0.5\n1 0.5\n2 2.5\n", encoding="utf-8") + histogram, count, dropped, raw = Wham1D.read_data(datafile, have_energy=False, config=wham.config) + assert raw == 3 + assert count == 1 + assert dropped == [0, 2] + + def _prepare_metadata_files(tmp_path: Path) -> tuple[Path, Path]: datafile = tmp_path / "data.dat" datafile.write_text("0 0.5 0\n1 1.5 0\n", encoding="utf-8") @@ -152,6 +163,7 @@ def test_save_free_and_convergence(wham: Wham1D) -> None: group.free_energies[1] = 2.5 assert not wham.is_converged(group) assert wham.average_diff(group) == pytest.approx(0.25) + assert wham.convergence_error(group) == pytest.approx(0.5) def test_calc_free() -> None: diff --git a/tests/test_wham2d.py b/tests/test_wham2d.py index 33e1c90..aa73f51 100644 --- a/tests/test_wham2d.py +++ b/tests/test_wham2d.py @@ -122,14 +122,17 @@ def test_read_data_with_energy_and_mask(tmp_path: Path, basic_config: Wham2DConf 0 0.5 0.5 1.0 1 1.5 1.5 0.0 2 0.2 0.2 0.5 +3 2.5 2.5 0.1 """.strip()) wham = Wham2D(basic_config) mask = [[0, 0], [0, 0]] - count = wham.read_data(data_file, have_energy=True, use_mask=True, mask=mask) + count, dropped, raw = wham.read_data(data_file, have_energy=True, use_mask=True, mask=mask) expected_weight = math.exp(-1.0 / wham.config.kT) + math.exp(-0.5 / wham.config.kT) + assert raw == 4 assert count == 3 + assert dropped == [3] assert wham.histogram[0][0] == pytest.approx(expected_weight) assert wham.histogram[1][1] == pytest.approx(1.0) assert mask == [[1, 0], [0, 1]] @@ -222,6 +225,7 @@ def test_save_convergence_and_average_diff(basic_config: Wham2DConfig) -> None: logged_previous = [0.45, 0.55] assert wham.is_converged(group, logged_current, logged_previous) assert wham.average_diff(logged_current, logged_previous) == pytest.approx(0.05) + assert wham.convergence_error(logged_current, logged_previous) == pytest.approx(0.05) logged_current[1] = 0.0 assert not wham.is_converged(group, logged_current, logged_previous)