boltzmann9 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

boltzmann9/project.py

@@ -0,0 +1,195 @@
+"""Project and run management utilities."""
+
+from __future__ import annotations
+
+import re
+import shutil
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict
+
+
+def _get_templates_dir() -> Path:
+    """Get the path to the templates directory."""
+    return Path(__file__).parent.parent / "templates"
+
+
+def _load_config_template(data_path: str) -> str:
+    """Load config.py template and substitute the data path.
+
+    Args:
+        data_path: Path to set for csv_path in the config.
+
+    Returns:
+        Config file content with substituted data path.
+    """
+    template_path = _get_templates_dir() / "config.py"
+    content = template_path.read_text()
+
+    # Replace the csv_path value (handles various quoting styles)
+    content = re.sub(
+        r'("csv_path":\s*)"[^"]*"',
+        f'\\1"{data_path}"',
+        content
+    )
+
+    return content
+
+
+def _load_generator_template() -> str:
+    """Load synthetic_generator.py template from templates directory.
+
+    Returns:
+        Generator script content.
+    """
+    template_path = _get_templates_dir() / "synthetic_generator.py"
+    return template_path.read_text()
+
+
+def create_project(project_path: str | Path) -> Path:
+    """Create a new BoltzmaNN9 project.
+
+    Args:
+        project_path: Path where the project will be created.
+
+    Returns:
+        Path to the created project.
+
+    Raises:
+        FileExistsError: If project directory already exists.
+    """
+    project_path = Path(project_path)
+
+    if project_path.exists():
+        raise FileExistsError(f"Project directory already exists: {project_path}")
+
+    # Create project structure
+    project_path.mkdir(parents=True)
+    data_dir = project_path / "data"
+    data_dir.mkdir()
+    output_dir = project_path / "output"
+    output_dir.mkdir()
+
+    # Create config.py from template with correct data path
+    data_csv_path = "data/data.csv"
+    config_content = _load_config_template(data_csv_path)
+    config_file = project_path / "config.py"
+    config_file.write_text(config_content)
+
+    # Create synthetic_generator.py in data folder from template
+    generator_content = _load_generator_template()
+    generator_file = data_dir / "synthetic_generator.py"
+    generator_file.write_text(generator_content)
+
+    print(f"Created project: {project_path}")
+    print(f"  - config.py")
+    print(f"  - data/")
+    print(f"    - synthetic_generator.py")
+    print(f"  - output/")
+    print(f"\nNext steps:")
+    print(f"  1. cd {project_path}/data && python synthetic_generator.py")
+    print(f"  2. Edit {project_path}/config.py as needed")
+    print(f"  3. python boltzmann.py train --project {project_path}")
+
+    return project_path
+
+
+def create_run(project_path: str | Path, run_name: str | None = None) -> Path:
+    """Create a new run directory within a project.
+
+    Args:
+        project_path: Path to the project.
+        run_name: Optional custom run name. If None, uses timestamp.
+
+    Returns:
+        Path to the created run directory.
+    """
+    project_path = Path(project_path)
+    output_dir = project_path / "output"
+
+    if not output_dir.exists():
+        output_dir.mkdir(parents=True)
+
+    # Generate run name with timestamp
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    if run_name:
+        run_dir_name = f"run_{timestamp}_{run_name}"
+    else:
+        run_dir_name = f"run_{timestamp}"
+
+    run_dir = output_dir / run_dir_name
+    run_dir.mkdir()
+
+    # Create subdirectories
+    (run_dir / "plots").mkdir()
+    (run_dir / "checkpoints").mkdir()
+
+    return run_dir
+
+
+def save_run_config(run_dir: Path, config_path: Path) -> None:
+    """Copy config file to run directory.
+
+    Args:
+        run_dir: Path to the run directory.
+        config_path: Path to the original config file.
+    """
+    dest = run_dir / "config.py"
+    shutil.copy2(config_path, dest)
+
+
+def get_run_paths(run_dir: Path) -> Dict[str, Path]:
+    """Get standard paths within a run directory.
+
+    Args:
+        run_dir: Path to the run directory.
+
+    Returns:
+        Dictionary with paths for model, log, plots, checkpoints, config.
+    """
+    return {
+        "model": run_dir / "model.pt",
+        "log": run_dir / "training.log",
+        "plots": run_dir / "plots",
+        "checkpoints": run_dir / "checkpoints",
+        "config": run_dir / "config.py",
+        "history": run_dir / "history.json",
+        "metrics": run_dir / "metrics.json",
+    }
+
+
+def find_latest_run(project_path: str | Path) -> Path | None:
+    """Find the most recent run in a project.
+
+    Args:
+        project_path: Path to the project.
+
+    Returns:
+        Path to the latest run directory, or None if no runs exist.
+    """
+    project_path = Path(project_path)
+    output_dir = project_path / "output"
+
+    if not output_dir.exists():
+        return None
+
+    runs = sorted(output_dir.glob("run_*"), reverse=True)
+    return runs[0] if runs else None
+
+
+def list_runs(project_path: str | Path) -> list[Path]:
+    """List all runs in a project.
+
+    Args:
+        project_path: Path to the project.
+
+    Returns:
+        List of run directory paths, sorted by date (newest first).
+    """
+    project_path = Path(project_path)
+    output_dir = project_path / "output"
+
+    if not output_dir.exists():
+        return []
+
+    return sorted(output_dir.glob("run_*"), reverse=True)

boltzmann9/run_utils.py

@@ -0,0 +1,262 @@
+"""Utilities for run logging and visualization."""
+
+from __future__ import annotations
+
+import json
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, TextIO
+
+
+class RunLogger:
+    """Logger that writes to both console and file."""
+
+    def __init__(self, log_path: Path):
+        """Initialize logger.
+
+        Args:
+            log_path: Path to log file.
+        """
+        self.log_path = log_path
+        self.log_file: TextIO | None = None
+        self._original_stdout = sys.stdout
+
+    def __enter__(self):
+        self.log_file = open(self.log_path, "w")
+        sys.stdout = self
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        sys.stdout = self._original_stdout
+        if self.log_file:
+            self.log_file.close()
+        return False
+
+    def write(self, message: str):
+        self._original_stdout.write(message)
+        if self.log_file:
+            self.log_file.write(message)
+            self.log_file.flush()
+
+    def flush(self):
+        self._original_stdout.flush()
+        if self.log_file:
+            self.log_file.flush()
+
+
+def save_history(history: Dict[str, list], path: Path) -> None:
+    """Save training history to JSON.
+
+    Args:
+        history: Training history dictionary.
+        path: Path to save JSON file.
+    """
+    # Convert any non-serializable types
+    serializable = {}
+    for key, values in history.items():
+        serializable[key] = [
+            float(v) if isinstance(v, (int, float)) else v
+            for v in values
+        ]
+
+    with open(path, "w") as f:
+        json.dump(serializable, f, indent=2)
+
+
+def load_history(path: Path) -> Dict[str, list]:
+    """Load training history from JSON.
+
+    Args:
+        path: Path to history JSON file.
+
+    Returns:
+        Training history dictionary.
+    """
+    with open(path) as f:
+        return json.load(f)
+
+
+def save_metrics(metrics: Dict[str, Any], path: Path) -> None:
+    """Save evaluation metrics to JSON.
+
+    Args:
+        metrics: Metrics dictionary.
+        path: Path to save JSON file.
+    """
+    # Convert nested structures
+    def convert(obj):
+        if isinstance(obj, dict):
+            return {k: convert(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [convert(v) for v in obj]
+        elif isinstance(obj, (int, float)):
+            return float(obj)
+        else:
+            return obj
+
+    with open(path, "w") as f:
+        json.dump(convert(metrics), f, indent=2)
+
+
+def load_metrics(path: Path) -> Dict[str, Any]:
+    """Load evaluation metrics from JSON.
+
+    Args:
+        path: Path to metrics JSON file.
+
+    Returns:
+        Metrics dictionary.
+    """
+    with open(path) as f:
+        return json.load(f)
+
+
+def save_plots(history: Dict[str, list], plots_dir: Path, model: Any = None) -> None:
+    """Save training plots to directory.
+
+    Args:
+        history: Training history dictionary.
+        plots_dir: Directory to save plots.
+        model: Optional RBM model instance. If provided, saves block diagram.
+    """
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        print("[save_plots] matplotlib not available, skipping plots")
+        return
+
+    import math
+
+    plots_dir.mkdir(parents=True, exist_ok=True)
+
+    # Save RBM block diagram if model is provided
+    if model is not None and hasattr(model, "draw_blocks"):
+        try:
+            model.draw_blocks(save_path=str(plots_dir / "RBM_blocks.png"), show=False)
+        except Exception as e:
+            print(f"[save_plots] Failed to save block diagram: {e}")
+
+    epochs = history.get(
+        "epoch", list(range(1, len(history.get("train_free_energy", [])) + 1))
+    )
+
+    def _is_all_nan(xs):
+        if not xs:
+            return True
+        return all(
+            (x is None) or (isinstance(x, float) and math.isnan(x)) for x in xs
+        )
+
+    # Plot 1: Free Energy
+    fig, ax = plt.subplots(figsize=(10, 6))
+    y_tr = history.get("train_free_energy", [])
+    y_va = history.get("val_free_energy", [])
+
+    ax.plot(epochs[:len(y_tr)], y_tr, label="train", linewidth=2)
+    if not _is_all_nan(y_va):
+        ax.plot(epochs[:len(y_va)], y_va, label="val", linewidth=2)
+
+    ax.set_title("Free Energy", fontsize=14)
+    ax.set_xlabel("Epoch", fontsize=12)
+    ax.set_ylabel("Mean FE (lower is better)", fontsize=12)
+    ax.legend()
+    ax.grid(True, alpha=0.3)
+    fig.tight_layout()
+    fig.savefig(plots_dir / "free_energy.png", dpi=150)
+    plt.close(fig)
+
+    # Plot 2: Reconstruction MSE
+    fig, ax = plt.subplots(figsize=(10, 6))
+    y_tr = history.get("train_recon_mse", [])
+    y_va = history.get("val_recon_mse", [])
+
+    ax.plot(epochs[:len(y_tr)], y_tr, label="train", linewidth=2)
+    if not _is_all_nan(y_va):
+        ax.plot(epochs[:len(y_va)], y_va, label="val", linewidth=2)
+
+    ax.set_title("Reconstruction MSE", fontsize=14)
+    ax.set_xlabel("Epoch", fontsize=12)
+    ax.set_ylabel("MSE", fontsize=12)
+    ax.legend()
+    ax.grid(True, alpha=0.3)
+    fig.tight_layout()
+    fig.savefig(plots_dir / "reconstruction_mse.png", dpi=150)
+    plt.close(fig)
+
+    # Plot 3: Reconstruction Bit Error
+    fig, ax = plt.subplots(figsize=(10, 6))
+    y_tr = history.get("train_recon_bit_error", [])
+    y_va = history.get("val_recon_bit_error", [])
+
+    ax.plot(epochs[:len(y_tr)], y_tr, label="train", linewidth=2)
+    if not _is_all_nan(y_va):
+        ax.plot(epochs[:len(y_va)], y_va, label="val", linewidth=2)
+
+    ax.set_title("Reconstruction Bit Error", fontsize=14)
+    ax.set_xlabel("Epoch", fontsize=12)
+    ax.set_ylabel("Fraction Mismatched", fontsize=12)
+    ax.legend()
+    ax.grid(True, alpha=0.3)
+    fig.tight_layout()
+    fig.savefig(plots_dir / "bit_error.png", dpi=150)
+    plt.close(fig)
+
+    # Plot 4: Learning Rate
+    if "lr" in history and history["lr"]:
+        fig, ax = plt.subplots(figsize=(10, 6))
+        lr_vals = history["lr"]
+        ax.plot(epochs[:len(lr_vals)], lr_vals, linewidth=2, color="green")
+        ax.set_title("Learning Rate Schedule", fontsize=14)
+        ax.set_xlabel("Epoch", fontsize=12)
+        ax.set_ylabel("Learning Rate", fontsize=12)
+        ax.set_yscale("log")
+        ax.grid(True, alpha=0.3)
+        fig.tight_layout()
+        fig.savefig(plots_dir / "learning_rate.png", dpi=150)
+        plt.close(fig)
+
+    # Combined plot
+    fig, axes = plt.subplots(1, 3, figsize=(16, 4))
+
+    # Free Energy
+    y_tr = history.get("train_free_energy", [])
+    y_va = history.get("val_free_energy", [])
+    axes[0].plot(epochs[:len(y_tr)], y_tr, label="train")
+    if not _is_all_nan(y_va):
+        axes[0].plot(epochs[:len(y_va)], y_va, label="val")
+    axes[0].set_title("Free Energy")
+    axes[0].set_xlabel("Epoch")
+    axes[0].set_ylabel("Mean FE")
+    axes[0].legend()
+    axes[0].grid(True, alpha=0.3)
+
+    # MSE
+    y_tr = history.get("train_recon_mse", [])
+    y_va = history.get("val_recon_mse", [])
+    axes[1].plot(epochs[:len(y_tr)], y_tr, label="train")
+    if not _is_all_nan(y_va):
+        axes[1].plot(epochs[:len(y_va)], y_va, label="val")
+    axes[1].set_title("Reconstruction MSE")
+    axes[1].set_xlabel("Epoch")
+    axes[1].set_ylabel("MSE")
+    axes[1].legend()
+    axes[1].grid(True, alpha=0.3)
+
+    # Bit Error
+    y_tr = history.get("train_recon_bit_error", [])
+    y_va = history.get("val_recon_bit_error", [])
+    axes[2].plot(epochs[:len(y_tr)], y_tr, label="train")
+    if not _is_all_nan(y_va):
+        axes[2].plot(epochs[:len(y_va)], y_va, label="val")
+    axes[2].set_title("Bit Error")
+    axes[2].set_xlabel("Epoch")
+    axes[2].set_ylabel("Fraction")
+    axes[2].legend()
+    axes[2].grid(True, alpha=0.3)
+
+    fig.tight_layout()
+    fig.savefig(plots_dir / "training_summary.png", dpi=150)
+    plt.close(fig)
+
+    print(f"Plots saved to: {plots_dir}")

boltzmann9/tester.py

@@ -0,0 +1,167 @@
+"""Testing utilities for RBM conditional probability evaluation."""
+
+from __future__ import annotations
+
+import math
+from collections import Counter
+from typing import Any, Dict, Sequence
+
+import torch
+from tqdm.auto import tqdm
+
+
+class RBMTester:
+    """Test RBM conditional probability estimation.
+
+    Evaluates how well the RBM can predict target visible units
+    given clamped (observed) visible units.
+
+    Args:
+        model: Trained RBM model.
+        test_dataloader: DataLoader for test data.
+        clamp_idx: Indices of visible units to clamp (condition on).
+        target_idx: Indices of visible units to predict.
+        device: Target device for computations.
+    """
+
+    def __init__(
+        self,
+        model,
+        test_dataloader,
+        clamp_idx: Sequence[int],
+        target_idx: Sequence[int],
+        *,
+        device=None,
+    ):
+        self.model = model
+        self.test_dataloader = test_dataloader
+        self.clamp_idx = list(clamp_idx)
+        self.target_idx = list(target_idx)
+        self.device = device or model.W.device
+
+    @staticmethod
+    def _bits_to_int(bits: torch.Tensor) -> int:
+        """Convert binary bits to integer (LSB-first)."""
+        weights = 2 ** torch.arange(bits.numel(), device=bits.device)
+        return int((bits * weights).sum().item())
+
+    @torch.no_grad()
+    def conditional_nll(
+        self,
+        *,
+        n_samples: int = 500,
+        burn_in: int = 300,
+        thin: int = 10,
+        laplace_alpha: float = 1.0,
+        log_every: int = 50,
+    ) -> Dict[str, Any]:
+        """Compute conditional negative log-likelihood over test set.
+
+        For each test sample, clamps the specified visible units and
+        samples from the conditional distribution p(target | clamp).
+
+        Computes per-bit NLL using empirical bit frequencies, which scales
+        to large target dimensions.
+
+        Args:
+            n_samples: Number of MCMC samples per test point.
+            burn_in: Burn-in steps for MCMC.
+            thin: Thinning interval between samples.
+            laplace_alpha: Laplace smoothing parameter for per-bit probabilities.
+            log_every: Log progress every N samples.
+
+        Returns:
+            Dictionary with:
+                - mean_nll_nats: Mean NLL in nats (sum over target bits).
+                - mean_nll_bits: Mean NLL in bits.
+                - mean_nll_per_bit: Mean NLL per target bit.
+                - nll_nats_per_sample: List of per-sample NLL in nats.
+                - nll_bits_per_sample: List of per-sample NLL in bits.
+        """
+        self.model.eval()
+
+        ln2 = math.log(2.0)
+
+        nlls_nats = []
+        nlls_bits = []
+
+        total_points = len(self.test_dataloader.dataset)
+        n_target_bits = len(self.target_idx)
+
+        outer_pbar = tqdm(
+            total=total_points,
+            desc="RBM conditional NLL",
+            leave=True,
+        )
+
+        for batch_idx, v in enumerate(self.test_dataloader):
+            v = v.to(self.device)
+
+            for i in range(v.size(0)):
+                # Clamp visible units
+                v_clamp = torch.zeros(
+                    self.model.nv,
+                    device=self.device,
+                    dtype=self.model.W.dtype,
+                )
+                v_clamp[self.clamp_idx] = v[i, self.clamp_idx]
+
+                # True target bits
+                true_bits = v[i, self.target_idx]
+
+                # Sample conditional
+                samples = self.model.sample_clamped(
+                    v_clamp=v_clamp,
+                    clamp_idx=self.clamp_idx,
+                    n_samples=n_samples,
+                    burn_in=burn_in,
+                    thin=thin,
+                )
+
+                # Get target bits from samples: (n_samples, n_target_bits)
+                sampled_bits = samples[:, self.target_idx]
+
+                # Compute per-bit empirical probability with Laplace smoothing
+                # For each target bit, estimate P(bit=true_value)
+                bit_matches = (sampled_bits == true_bits.unsqueeze(0)).float()
+                bit_probs = (bit_matches.sum(dim=0) + laplace_alpha) / (n_samples + 2 * laplace_alpha)
+
+                # NLL = -sum(log(p_i)) for each bit
+                # Clamp to avoid log(0)
+                bit_probs = bit_probs.clamp(min=1e-10, max=1 - 1e-10)
+                nll_nats = -torch.log(bit_probs).sum().item()
+                nll_bits = nll_nats / ln2
+
+                nlls_nats.append(nll_nats)
+                nlls_bits.append(nll_bits)
+
+                # Logging
+                if len(nlls_nats) % log_every == 0:
+                    mean_nats = sum(nlls_nats) / len(nlls_nats)
+                    mean_bits = sum(nlls_bits) / len(nlls_bits)
+                    outer_pbar.set_postfix(
+                        nll_nats=f"{mean_nats:.3f}",
+                        nll_bits=f"{mean_bits:.3f}",
+                        per_bit=f"{mean_nats / n_target_bits:.4f}",
+                    )
+
+                outer_pbar.update(1)
+
+        outer_pbar.close()
+
+        mean_nats = (
+            float(sum(nlls_nats) / len(nlls_nats)) if nlls_nats else float("nan")
+        )
+        mean_bits = (
+            float(sum(nlls_bits) / len(nlls_bits)) if nlls_bits else float("nan")
+        )
+        mean_per_bit = mean_nats / n_target_bits if n_target_bits > 0 else float("nan")
+
+        return {
+            "mean_nll_nats": mean_nats,
+            "mean_nll_bits": mean_bits,
+            "mean_nll_per_bit": mean_per_bit,
+            "n_target_bits": n_target_bits,
+            "nll_nats_per_sample": nlls_nats,
+            "nll_bits_per_sample": nlls_bits,
+        }

boltzmann9/utils.py

@@ -0,0 +1,42 @@
+"""Utility functions for device resolution and memory management."""
+
+from __future__ import annotations
+
+import torch
+
+
+def resolve_device(device_cfg: str | None) -> torch.device:
+    """Resolve device configuration to a torch.device.
+
+    Args:
+        device_cfg: Device configuration string. Can be:
+            - None or "auto": Auto-detect best available device
+            - "cpu": Use CPU
+            - "cuda:0", "cuda:1", etc.: Use specific CUDA device
+            - "mps": Use Apple Metal Performance Shaders
+
+    Returns:
+        torch.device instance for the resolved device.
+    """
+    if device_cfg is None or device_cfg == "auto":
+        if torch.cuda.is_available():
+            return torch.device("cuda:0")
+        if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+            return torch.device("mps")
+        return torch.device("cpu")
+    return torch.device(device_cfg)
+
+
+def resolve_pin_memory(pin_cfg, device: torch.device) -> bool:
+    """Resolve pin_memory configuration based on device type.
+
+    Args:
+        pin_cfg: Pin memory configuration. Can be "auto", True, or False.
+        device: The target torch device.
+
+    Returns:
+        Boolean indicating whether to pin memory.
+    """
+    if pin_cfg == "auto":
+        return device.type == "cuda"
+    return bool(pin_cfg)