From 60810e5e0abebf8ae6d1ab671d5b33e995963cd6 Mon Sep 17 00:00:00 2001
From: Armin Shayesteh Zadeh <67717991+arminshzd@users.noreply.github.com>
Date: Mon, 8 Dec 2025 09:01:20 -0600
Subject: [PATCH 1/6] Add BayesWHAM driver and configuration examples

---
 README.md                  |  34 +++
 examples/bwham_config.yaml |  14 +
 pyproject.toml             |   1 +
 src/pywham/__init__.py     |   4 +
 src/pywham/bwham.py        | 530 +++++++++++++++++++++++++++++++++++++
 5 files changed, 583 insertions(+)
 create mode 100644 examples/bwham_config.yaml
 create mode 100644 src/pywham/bwham.py

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..5251c63 100644
--- a/src/pywham/__init__.py
+++ b/src/pywham/__init__.py
@@ -1,8 +1,12 @@
 """Python translation of the WHAM C implementations."""
 
+from .bwham import BayesWHAM, BayesWhamConfig, build_config as build_bwham_config
 from .visualization import plot_free_energy_1d, plot_free_energy_2d, save_free_energy_plots
 
 __all__ = [
+    "BayesWHAM",
+    "BayesWhamConfig",
+    "build_bwham_config",
     "plot_free_energy_1d",
     "plot_free_energy_2d",
     "save_free_energy_plots",
diff --git a/src/pywham/bwham.py b/src/pywham/bwham.py
new file mode 100644
index 0000000..9723981
--- /dev/null
+++ b/src/pywham/bwham.py
@@ -0,0 +1,530 @@
+"""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
+
+
+@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
+
+    def _posterior_samples(self) -> int:
+        usable = max(self.config.num_samples - self.config.burn_in, 0)
+        if usable == 0:
+            return 0
+        return (usable + self.config.thinning - 1) // self.config.thinning
+
+    # --- 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]:
+        lines = wham.config.metadata_path.read_text(encoding="utf-8").splitlines()
+        num_windows = wham.get_numwindows(lines)
+        print(f"#Number of windows = {num_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:
+            for i in range(hist_group.num_windows):
+                hist_group.temperatures[i] = wham.config.kT
+        return hist_group, have_temp, entries
+
+    def _run_wham_1d(self, wham: Wham1D, hist_group: HistGroup1D, have_energy: bool) -> tuple[List[float], List[float]]:
+        probabilities = [0.0 for _ in range(wham.config.num_bins)]
+        iteration = 0
+        first = True
+        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 iteration >= 100000:
+                print(f"Too many iterations: {iteration}")
+                break
+        total = sum(probabilities)
+        if total:
+            probabilities = [p / total for p in probabilities]
+        free_energy, _ = wham.calc_free(probabilities)
+        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:
+            return np.empty((0, wham.config.num_bins)), np.empty((0, wham.config.num_bins)), np.empty((0, base_group.num_windows))
+        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):
+            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
+            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)
+            idx += 1
+        return free_samples[:idx, :], prob_samples[:idx, :], window_samples[:idx, :]
+
+    def _format_output_1d(self, wham: Wham1D, result: BayesResult1D) -> None:
+        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):
+        lines = wham.config.metadata_path.read_text(encoding="utf-8").splitlines()
+        num_windows = wham.get_numwindows(lines)
+        print(f"#Number of windows = {num_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:
+            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
+        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
+    ) -> tuple[List[List[float]], List[List[float]]]:
+        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
+        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 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
+        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:
+            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))
+        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):
+            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
+            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)
+            idx += 1
+        return free_samples[:idx, :, :], prob_samples[:idx, :, :], window_samples[:idx, :]
+
+    def _format_output_2d(self, wham: Wham2D, result: BayesResult2D) -> None:
+        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:
+        if self.config.dimension == 1:
+            wham = Wham1D(self.config.base_config)  # type: ignore[arg-type]
+            hist_group, have_energy, _ = self._load_1d(wham)
+            map_free, map_prob = self._run_wham_1d(wham, hist_group, have_energy)
+            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._format_output_1d(wham, result)
+            return
+
+        wham2d = Wham2D(self.config.base_config)  # type: ignore[arg-type]
+        hist_group2d, have_energy2d, _, mask = self._load_2d(wham2d)
+        map_free_2d, map_prob_2d = self._run_wham_2d(wham2d, hist_group2d, have_energy2d, mask)
+        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._format_output_2d(wham2d, result2d)
+
+
+# --- configuration ------------------------------------------------------
+
+
+def build_config(yaml_path: Path) -> BayesWhamConfig:
+    if not yaml_path.exists():
+        raise FileNotFoundError(f"YAML configuration file not found: {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:
+        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}")
+        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=Path(config_raw["metadata_file"]),
+            freefile_path=Path(config_raw["freefile"]),
+            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=None,
+            aux_data_path=Path(config_raw["aux_data_file"]) if "aux_data_file" in config_raw else None,
+        )
+        return BayesWhamConfig(
+            dimension=2,
+            base_config=base,
+            num_samples=num_samples,
+            burn_in=burn_in,
+            thinning=thinning,
+            dirichlet_alpha=dirichlet_alpha,
+        )
+
+    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}")
+    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=Path(config_raw["metadata_file"]),
+        freefile_path=Path(config_raw["freefile"]),
+        periodic=periodic,
+        period=period,
+        k_B=k_B,
+        num_mc_runs=0,
+        mc_seed=None,
+        mc_workers=None,
+        ingest_workers=None,
+        aux_data_path=None,
+    )
+    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()

From 5705927afc90c900bcf917d23f568fe601afaa85 Mon Sep 17 00:00:00 2001
From: Armin Shayesteh Zadeh <67717991+arminshzd@users.noreply.github.com>
Date: Mon, 8 Dec 2025 10:30:35 -0600
Subject: [PATCH 2/6] Add aux data export for BayesWHAM

---
 src/pywham/bwham.py |  45 +++++++++++++----
 tests/test_bwham.py | 114 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 150 insertions(+), 9 deletions(-)
 create mode 100644 tests/test_bwham.py

diff --git a/src/pywham/bwham.py b/src/pywham/bwham.py
index 9723981..980dc65 100644
--- a/src/pywham/bwham.py
+++ b/src/pywham/bwham.py
@@ -327,7 +327,7 @@ def _format_output_2d(self, wham: Wham2D, result: BayesResult2D) -> None:
     def run(self) -> None:
         if self.config.dimension == 1:
             wham = Wham1D(self.config.base_config)  # type: ignore[arg-type]
-            hist_group, have_energy, _ = self._load_1d(wham)
+            hist_group, have_energy, entries = self._load_1d(wham)
             map_free, map_prob = self._run_wham_1d(wham, hist_group, have_energy)
             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)
@@ -351,10 +351,11 @@ def run(self) -> None:
                 std_window_free=std_window.tolist(),
             )
             self._format_output_1d(wham, result)
+            wham._write_aux_data(entries, hist_group, result.map_window_free, window_samples.tolist())
             return
 
         wham2d = Wham2D(self.config.base_config)  # type: ignore[arg-type]
-        hist_group2d, have_energy2d, _, mask = self._load_2d(wham2d)
+        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)
         samples_free_2d, samples_prob_2d, window_samples_2d = self._bayes_samples_2d(
             wham2d, hist_group2d, have_energy2d, mask
@@ -395,11 +396,19 @@ def run(self) -> None:
             std_window_free=std_window_2d.tolist(),
         )
         self._format_output_2d(wham2d, result2d)
+        wham2d._write_aux_data(entries2d, hist_group2d, result2d.map_window_free, window_samples_2d.tolist())
 
 
 # --- 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}")
@@ -434,6 +443,13 @@ def build_config(yaml_path: Path) -> BayesWhamConfig:
         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"]),
@@ -444,8 +460,8 @@ def build_config(yaml_path: Path) -> BayesWhamConfig:
             tolerance=float(config_raw["tolerance"]),
             temperature=float(config_raw["temperature"]),
             numpad=int(config_raw["numpad"]),
-            metadata_path=Path(config_raw["metadata_file"]),
-            freefile_path=Path(config_raw["freefile"]),
+            metadata_path=metadata_path,
+            freefile_path=freefile_path,
             use_mask=bool(config_raw["use_mask"]),
             periodic_x=periodic_x,
             period_x=period_x,
@@ -457,8 +473,12 @@ def build_config(yaml_path: Path) -> BayesWhamConfig:
             num_mc_runs=0,
             mc_seed=None,
             mc_workers=None,
-            freefile_error_path=None,
-            aux_data_path=Path(config_raw["aux_data_file"]) if "aux_data_file" in config_raw else 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,
         )
         return BayesWhamConfig(
             dimension=2,
@@ -484,6 +504,13 @@ def build_config(yaml_path: Path) -> BayesWhamConfig:
     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"]),
@@ -491,8 +518,8 @@ def build_config(yaml_path: Path) -> BayesWhamConfig:
         tolerance=float(config_raw["tolerance"]),
         temperature=float(config_raw["temperature"]),
         numpad=int(config_raw["numpad"]),
-        metadata_path=Path(config_raw["metadata_file"]),
-        freefile_path=Path(config_raw["freefile"]),
+        metadata_path=metadata_path,
+        freefile_path=freefile_path,
         periodic=periodic,
         period=period,
         k_B=k_B,
@@ -500,7 +527,7 @@ def build_config(yaml_path: Path) -> BayesWhamConfig:
         mc_seed=None,
         mc_workers=None,
         ingest_workers=None,
-        aux_data_path=None,
+        aux_data_path=aux_data_path,
     )
     return BayesWhamConfig(
         dimension=1,
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")
+

From f879f1b606beeaf154385e641225c46a8b31823c Mon Sep 17 00:00:00 2001
From: Armin Shayesteh Zadeh <arminshzd@gmail.com>
Date: Tue, 16 Dec 2025 12:53:03 -0600
Subject: [PATCH 3/6] add logging to bwham. expose bwham to cli.

---
 src/pywham/__init__.py | 28 ++++++++++++++++++++++++++--
 src/pywham/bwham.py    | 27 +++++++++++++++++++++++----
 2 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/src/pywham/__init__.py b/src/pywham/__init__.py
index 5251c63..fd4c13b 100644
--- a/src/pywham/__init__.py
+++ b/src/pywham/__init__.py
@@ -1,7 +1,7 @@
 """Python translation of the WHAM C implementations."""
 
-from .bwham import BayesWHAM, BayesWhamConfig, build_config as build_bwham_config
-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",
@@ -11,3 +11,27 @@
     "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
index 980dc65..f402640 100644
--- a/src/pywham/bwham.py
+++ b/src/pywham/bwham.py
@@ -114,7 +114,9 @@ def _load_1d(self, wham: Wham1D) -> tuple[HistGroup1D, bool, list]:
                 hist_group.temperatures[i] = wham.config.kT
         return hist_group, have_temp, entries
 
-    def _run_wham_1d(self, wham: Wham1D, hist_group: HistGroup1D, have_energy: bool) -> tuple[List[float], List[float]]:
+    def _run_wham_1d(
+        self, wham: Wham1D, hist_group: HistGroup1D, have_energy: bool, log_iterations: bool = False
+    ) -> tuple[List[float], List[float]]:
         probabilities = [0.0 for _ in range(wham.config.num_bins)]
         iteration = 0
         first = True
@@ -123,6 +125,12 @@ def _run_wham_1d(self, wham: Wham1D, hist_group: HistGroup1D, have_energy: bool)
             wham.save_free(hist_group)
             wham.wham_iteration(hist_group, probabilities, have_energy)
             iteration += 1
+            if log_iterations and iteration % 10 == 0:
+                error = wham.average_diff(hist_group)
+                print(f"# Iteration {iteration:8d} | error {error:12.6e}")
+            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
@@ -234,7 +242,12 @@ def _load_2d(self, wham: Wham2D):
         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
+        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]]]:
         dtype = np.float32 if wham.config.use_float32 else np.float64
         x_grid, y_grid = wham._coordinate_grids(dtype)
@@ -270,6 +283,12 @@ def _run_wham_2d(
                 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.average_diff(logged_current, logged_previous)
+                print(f"# Iteration {iteration:8d} | error {error:12.6e}")
+            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
@@ -328,7 +347,7 @@ def run(self) -> None:
         if self.config.dimension == 1:
             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)
+            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)
@@ -356,7 +375,7 @@ def run(self) -> None:
 
         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)
+        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
         )

From b9c9a89f3c7342bad2cc1701d4da5c92fee0adb4 Mon Sep 17 00:00:00 2001
From: Armin Shayesteh Zadeh <arminshzd@gmail.com>
Date: Wed, 17 Dec 2025 08:36:01 -0600
Subject: [PATCH 4/6] new output and error reporting

---
 src/pywham/bwham.py  | 74 +++++++++++++++++++++++++++++++++++++++++---
 src/pywham/wham1d.py |  9 +++++-
 src/pywham/wham2d.py | 10 +++++-
 tests/test_wham1d.py |  1 +
 tests/test_wham2d.py |  1 +
 5 files changed, 88 insertions(+), 7 deletions(-)

diff --git a/src/pywham/bwham.py b/src/pywham/bwham.py
index f402640..20e65ac 100644
--- a/src/pywham/bwham.py
+++ b/src/pywham/bwham.py
@@ -16,6 +16,10 @@
 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."""
@@ -68,12 +72,26 @@ class BayesWHAM:
 
     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
-        return (usable + self.config.thinning - 1) // self.config.thinning
+        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:
@@ -103,41 +121,51 @@ def _clone_group_1d(self, base: HistGroup1D) -> HistGroup1D:
         )
 
     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.average_diff(hist_group)
-                print(f"# Iteration {iteration:8d} | error {error:12.6e}")
+                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(
@@ -148,7 +176,9 @@ def _bayes_samples_1d(
     ) -> 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)
@@ -157,18 +187,23 @@ def _bayes_samples_1d(
         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(
@@ -222,9 +257,11 @@ def _clone_group_2d(self, base: HistGroup2D) -> HistGroup2D:
         )
 
     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:
@@ -234,11 +271,13 @@ def _load_2d(self, wham: Wham2D):
         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(
@@ -249,6 +288,7 @@ def _run_wham_2d(
         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)
 
@@ -266,6 +306,7 @@ def _run_wham_2d(
         iteration = 0
         first = True
         converged = False
+        printed_iteration = False
         while not converged or first:
             first = False
             wham.save_free(hist_group)
@@ -284,8 +325,9 @@ def _run_wham_2d(
             ]
             converged = wham.is_converged(hist_group, logged_current, logged_previous)
             if log_iterations and iteration % 10 == 0:
-                error = wham.average_diff(logged_current, logged_previous)
-                print(f"# Iteration {iteration:8d} | error {error:12.6e}")
+                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)
@@ -297,6 +339,9 @@ def _run_wham_2d(
         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(
@@ -304,9 +349,11 @@ def _bayes_samples_2d(
     ) -> 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)
@@ -315,18 +362,23 @@ def _bayes_samples_2d(
         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):
@@ -344,7 +396,9 @@ def _format_output_2d(self, wham: Wham2D, result: BayesResult2D) -> None:
 
     # --- 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)
@@ -369,10 +423,13 @@ def run(self) -> None:
                 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)
@@ -414,8 +471,10 @@ def run(self) -> None:
             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 ------------------------------------------------------
@@ -432,6 +491,7 @@ 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")
@@ -442,6 +502,7 @@ def build_config(yaml_path: Path) -> BayesWhamConfig:
     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")
@@ -499,6 +560,7 @@ def build_config(yaml_path: Path) -> BayesWhamConfig:
             ),
             aux_data_path=aux_data_path,
         )
+        _log_message("Finished assembling 2D BayesWHAM configuration.")
         return BayesWhamConfig(
             dimension=2,
             base_config=base,
@@ -508,6 +570,7 @@ def build_config(yaml_path: Path) -> BayesWhamConfig:
             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 = [
@@ -548,6 +611,7 @@ def build_config(yaml_path: Path) -> BayesWhamConfig:
         ingest_workers=None,
         aux_data_path=aux_data_path,
     )
+    _log_message("Finished assembling 1D BayesWHAM configuration.")
     return BayesWhamConfig(
         dimension=1,
         base_config=base,
diff --git a/src/pywham/wham1d.py b/src/pywham/wham1d.py
index 8b6806e..5e47080 100644
--- a/src/pywham/wham1d.py
+++ b/src/pywham/wham1d.py
@@ -310,6 +310,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 +427,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)
diff --git a/src/pywham/wham2d.py b/src/pywham/wham2d.py
index 3f589dc..3ef4554 100644
--- a/src/pywham/wham2d.py
+++ b/src/pywham/wham2d.py
@@ -395,6 +395,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 +624,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_wham1d.py b/tests/test_wham1d.py
index 25eca46..d4d716a 100644
--- a/tests/test_wham1d.py
+++ b/tests/test_wham1d.py
@@ -152,6 +152,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..1ed0d50 100644
--- a/tests/test_wham2d.py
+++ b/tests/test_wham2d.py
@@ -222,6 +222,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)

From 080469322e8272c7d1256cad5e1462d820a2a9f6 Mon Sep 17 00:00:00 2001
From: Armin Shayesteh Zadeh <arminshzd@gmail.com>
Date: Thu, 18 Dec 2025 10:47:22 -0600
Subject: [PATCH 5/6] handle zero weight in reweight. handle edge samples in
 wham outputs. The aux output now tracks the discarded sampels as well to
 prevent mismatch between the number of samples between the original and the
 aux cv trajs.

---
 src/pywham/reweight.py | 77 ++++++++++++++++++++++++++++++++++++------
 src/pywham/wham1d.py   | 26 +++++++++++---
 src/pywham/wham2d.py   | 24 ++++++++++---
 tests/test_wham1d.py   | 13 ++++++-
 tests/test_wham2d.py   |  5 ++-
 5 files changed, 123 insertions(+), 22 deletions(-)

diff --git a/src/pywham/reweight.py b/src/pywham/reweight.py
index 54b5eac..ce665e4 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
@@ -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,6 +448,17 @@ 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):
diff --git a/src/pywham/wham1d.py b/src/pywham/wham1d.py
index 5e47080..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 = {
@@ -577,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
@@ -587,6 +601,8 @@ class WindowLoadResult:
     min_nonzero: int
     max_nonzero: int
     warnings: list[str]
+    raw_samples: int
+    dropped_samples: list[int]
 
 
 def _build_histogram(
@@ -643,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}")
 
@@ -658,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 3ef4554..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
diff --git a/tests/test_wham1d.py b/tests/test_wham1d.py
index d4d716a..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")
diff --git a/tests/test_wham2d.py b/tests/test_wham2d.py
index 1ed0d50..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]]

From 4bab401e6e2785f4b4654ed2dea22ae23db2d109 Mon Sep 17 00:00:00 2001
From: Armin Shayesteh Zadeh <arminshzd@gmail.com>
Date: Fri, 19 Dec 2025 18:21:02 -0600
Subject: [PATCH 6/6] fix to the periodic behavior

---
 src/pywham/reweight.py | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/pywham/reweight.py b/src/pywham/reweight.py
index ce665e4..65aa659 100644
--- a/src/pywham/reweight.py
+++ b/src/pywham/reweight.py
@@ -265,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)
@@ -463,8 +463,16 @@ def _locate_bin(values: np.ndarray, edges: Sequence[np.ndarray]) -> tuple[int, .
     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)