ebm_binary.py

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
import time
import torch
import torch.nn.functional as F
from torch import Tensor, nn
from torch.utils.data import DataLoader


@dataclass
class MarkovChainState:
    """State of the negative-phase Markov chains."""

    visible: Tensor
    auxiliaries: dict[str, Tensor] = field(default_factory=dict)


@dataclass
class TrainingHistory:
    """Container for training diagnostics."""

    energy_gaps: list[float] = field(default_factory=list)
    energy_data: list[float] = field(default_factory=list)
    energy_model: list[float] = field(default_factory=list)
    log_pseudo_likelihood: list[float] = field(default_factory=list)
    acceptance_rate: list[float] = field(default_factory=list)


class BinaryEnergyFunction(nn.Module, ABC):
    """Base class for binary-visible energy functions."""

    def __init__(self, num_visibles: int) -> None:
        super().__init__()
        if num_visibles <= 0:
            raise ValueError("`num_visibles` must be positive.")
        self.num_visibles = num_visibles

    @property
    def device(self) -> torch.device:
        """Return the current parameter device."""

        return next(self.parameters()).device

    @property
    def dtype(self) -> torch.dtype:
        """Return the current parameter dtype."""

        return next(self.parameters()).dtype

    @abstractmethod
    def energy(self, visible: Tensor) -> Tensor:
        """Return the energy of a batch of visible configurations."""

    @torch.no_grad()
    def initialize_chain_state(self, num_chains: int) -> MarkovChainState:
        """Default random initialization for binary visible chains."""

        visible = torch.bernoulli(
            torch.full(
                (num_chains, self.num_visibles),
                0.5,
                device=self.device,
                dtype=self.dtype,
            )
        )
        return MarkovChainState(visible=visible)

    def forward(self, visible: Tensor) -> Tensor:
        """Alias of :meth:`energy` for nn.Module interoperability."""

        return self.energy(visible)


class BinaryTransitionKernel(ABC):
    """Base class for binary-state transition kernels."""

    @torch.no_grad()
    def initialize_state(self, energy_function: BinaryEnergyFunction, num_chains: int) -> MarkovChainState:
        """Initialize the chains before running the negative phase."""

        return energy_function.initialize_chain_state(num_chains=num_chains)

    @abstractmethod
    @torch.no_grad()
    def step(
        self,
        energy_function: BinaryEnergyFunction,
        state: MarkovChainState,
    ) -> tuple[MarkovChainState, float | None]:
        """Run one transition step."""

    @torch.no_grad()
    def run(
        self,
        energy_function: BinaryEnergyFunction,
        state: MarkovChainState,
        num_steps: int,
    ) -> tuple[MarkovChainState, float | None]:
        """Run several transition steps and aggregate acceptance if available."""

        if num_steps < 0:
            raise ValueError("`num_steps` must be non-negative.")

        acceptance_rates: list[float] = []
        for _ in range(num_steps):
            state, acceptance_rate = self.step(energy_function=energy_function, state=state)
            if acceptance_rate is not None:
                acceptance_rates.append(float(acceptance_rate))

        mean_acceptance = sum(acceptance_rates) / len(acceptance_rates) if acceptance_rates else None
        return state, mean_acceptance


class BinaryEnergyBasedModel(nn.Module):
    """Generic binary EBM trained with a pluggable transition kernel.

    The training loop only depends on two interchangeable pieces:

    - a visible-space energy function ``E(x)``
    - a transition kernel that samples under ``exp(-E(x))``
    """

    def __init__(
        self,
        energy_function: BinaryEnergyFunction,
        transition_kernel: BinaryTransitionKernel,
    ) -> None:
        super().__init__()
        self.energy_function = energy_function
        self.transition_kernel = transition_kernel

    @property
    def num_visibles(self) -> int:
        """Return the visible dimension of the energy function."""

        return self.energy_function.num_visibles

    @property
    def device(self) -> torch.device:
        """Return the current parameter device."""

        return self.energy_function.device

    @property
    def dtype(self) -> torch.dtype:
        """Return the current parameter dtype."""

        return self.energy_function.dtype

    def energy(self, visible: Tensor) -> Tensor:
        """Compute energies through the wrapped energy function."""

        return self.energy_function.energy(visible)

    @torch.no_grad()
    def initialize_state(self, num_chains: int) -> MarkovChainState:
        """Initialize a Markov-chain state with the configured kernel."""

        return self.transition_kernel.initialize_state(self.energy_function, num_chains=num_chains)

    @torch.no_grad()
    def update_state(
        self,
        state: MarkovChainState,
        num_steps: int,
    ) -> tuple[MarkovChainState, float | None]:
        """Advance a chain state with the configured kernel."""

        return self.transition_kernel.run(
            energy_function=self.energy_function,
            state=state,
            num_steps=num_steps,
        )

    @torch.no_grad()
    def negative_phase_step(
        self,
        negative_phase: str,
        state: MarkovChainState | None,
        num_steps: int,
        num_model_chains: int,
    ) -> tuple[Tensor, MarkovChainState, float | None]:
        """Generate model samples for one optimization step.

        modes:
            - ``pcd``: persistent chains are reused across minibatches.
            - ``rdm``: chains are reinitialized from scratch each minibatch.
        """

        if negative_phase == "pcd":
            if state is None or state.visible.shape[0] != num_model_chains:
                state = self.initialize_state(num_chains=num_model_chains)
        elif negative_phase == "rdm":
            state = self.initialize_state(num_chains=num_model_chains)
        else:
            raise ValueError("`negative_phase` must be either 'pcd' or 'rdm'.")

        state, acceptance_rate = self.update_state(state=state, num_steps=num_steps)
        return state.visible.detach(), state, acceptance_rate

    @torch.no_grad()
    def log_pseudo_likelihood(self, visible: Tensor) -> Tensor:
        """Estimate the stochastic log pseudo-likelihood of a minibatch."""

        visible = _prepare_visible_tensor(visible, device=self.device, dtype=self.dtype)
        bit_index = torch.randint(
            low=0,
            high=self.num_visibles,
            size=(visible.shape[0],),
            device=visible.device,
        )

        visible_corrupted = visible.clone()
        visible_corrupted[torch.arange(visible.shape[0], device=visible.device), bit_index] = 1.0 - visible_corrupted[
            torch.arange(visible.shape[0], device=visible.device),
            bit_index,
        ]

        energy_original = self.energy(visible)
        energy_corrupted = self.energy(visible_corrupted)
        return self.num_visibles * F.logsigmoid(energy_corrupted - energy_original).mean()

    @torch.no_grad()
    def sample(
        self,
        num_samples: int,
        num_mcmc_steps: int = 100,
        initial_visible: Tensor | None = None,
    ) -> Tensor:
        """Generate visible samples with the configured kernel."""

        if initial_visible is None:
            state = self.initialize_state(num_chains=num_samples)
        else:
            visible = _prepare_visible_tensor(initial_visible, device=self.device, dtype=self.dtype)
            if visible.shape[1] != self.num_visibles:
                raise ValueError(
                    f"Expected flattened visible dimension {self.num_visibles}, got {visible.shape[1]}."
                )
            state = MarkovChainState(visible=visible)

        state, _ = self.update_state(state=state, num_steps=num_mcmc_steps)
        return state.visible

    def fit(
        self,
        dataloader: DataLoader,
        num_epochs: int,
        learning_rate: float = 1e-3,
        num_negative_phase_steps: int = 1,
        num_persistent_chains: int = 128,
        weight_decay: float = 0.0,
        compute_log_pseudo_likelihood: bool = True,
        negative_phase: str = "pcd",
        show_progress: bool = False,
        progress_label: str | None = None,
        progress_every: int = 1,
    ) -> tuple[TrainingHistory, MarkovChainState]:
        """Train the EBM with either persistent or random-start chains."""

        if num_epochs <= 0:
            raise ValueError("`num_epochs` must be positive.")
        if progress_every <= 0:
            raise ValueError("`progress_every` must be positive.")

        optimizer = torch.optim.SGD(self.parameters(), lr=learning_rate, weight_decay=weight_decay)
        state = self.initialize_state(num_chains=num_persistent_chains)
        history = TrainingHistory()
        progress_prefix = f"[{progress_label}] " if progress_label else ""
        try:
            batches_per_epoch = len(dataloader)
        except TypeError:
            batches_per_epoch = None
        total_expected_updates = batches_per_epoch * num_epochs if batches_per_epoch is not None else None

        self.train()
        train_start_time = time.perf_counter()
        total_updates = 0
        for epoch_index in range(num_epochs):
            saw_batch = False
            epoch_start_energy_gap_index = len(history.energy_gaps)
            epoch_start_log_pll_index = len(history.log_pseudo_likelihood)
            epoch_start_accept_index = len(history.acceptance_rate)
            epoch_start_time = time.perf_counter()
            for batch in dataloader:
                saw_batch = True
                visible_data = _extract_visible_batch(batch, device=self.device, dtype=self.dtype)
                if visible_data.shape[1] != self.num_visibles:
                    raise ValueError(
                        f"Expected flattened visible dimension {self.num_visibles}, got {visible_data.shape[1]}."
                    )

                visible_model, state, acceptance_rate = self.negative_phase_step(
                    negative_phase=negative_phase,
                    state=state,
                    num_steps=num_negative_phase_steps,
                    num_model_chains=num_persistent_chains,
                )

                positive_energy = self.energy(visible_data).mean()
                negative_energy = self.energy(visible_model).mean()
                energy_gap = positive_energy - negative_energy

                optimizer.zero_grad(set_to_none=True)
                energy_gap.backward()
                optimizer.step()

                history.energy_gaps.append(float(energy_gap.detach().cpu()))
                history.energy_data.append(float(positive_energy.detach().cpu()))
                history.energy_model.append(float(negative_energy.detach().cpu()))

                if compute_log_pseudo_likelihood:
                    log_pseudo_likelihood = self.log_pseudo_likelihood(visible_data)
                    history.log_pseudo_likelihood.append(float(log_pseudo_likelihood.detach().cpu()))
                if acceptance_rate is not None:
                    history.acceptance_rate.append(float(acceptance_rate))
                total_updates += 1

            if not saw_batch:
                raise ValueError("The dataloader produced no batches.")

            if show_progress and ((epoch_index + 1) % progress_every == 0 or epoch_index + 1 == num_epochs):
                epoch_seconds = time.perf_counter() - epoch_start_time
                elapsed_seconds = time.perf_counter() - train_start_time
                mean_energy_gap = sum(
                    history.energy_gaps[epoch_start_energy_gap_index:]
                ) / max(
                    1,
                    len(history.energy_gaps) - epoch_start_energy_gap_index,
                )

                status_parts = [
                    f"{progress_prefix}epoch {epoch_index + 1}/{num_epochs}",
                    f"elapsed={elapsed_seconds:.1f}s",
                    f"epoch_time={epoch_seconds:.1f}s",
                    f"updates={total_updates}",
                    f"mean_energy_gap={mean_energy_gap:.4f}",
                ]

                if compute_log_pseudo_likelihood and len(history.log_pseudo_likelihood) > epoch_start_log_pll_index:
                    mean_log_pll = sum(history.log_pseudo_likelihood[epoch_start_log_pll_index:]) / max(
                        1,
                        len(history.log_pseudo_likelihood) - epoch_start_log_pll_index,
                    )
                    status_parts.append(f"mean_log_pll={mean_log_pll:.4f}")

                if len(history.acceptance_rate) > epoch_start_accept_index:
                    mean_acceptance = sum(history.acceptance_rate[epoch_start_accept_index:]) / max(
                        1,
                        len(history.acceptance_rate) - epoch_start_accept_index,
                    )
                    status_parts.append(f"mean_accept={mean_acceptance:.3f}")

                if total_expected_updates is not None and total_updates > 0:
                    seconds_per_update = elapsed_seconds / total_updates
                    eta_seconds = seconds_per_update * (total_expected_updates - total_updates)
                    status_parts.append(f"eta={eta_seconds:.1f}s")

                print(" | ".join(status_parts))

        return history, state


def _prepare_visible_tensor(
    visible: Tensor,
    *,
    device: torch.device,
    dtype: torch.dtype,
) -> Tensor:
    """Move a visible tensor to the requested device/dtype and flatten it if needed."""

    visible = visible.to(device=device, dtype=dtype)
    if visible.ndim == 1:
        visible = visible.unsqueeze(0)
    elif visible.ndim > 2:
        visible = visible.reshape(visible.shape[0], -1)
    return visible


def _extract_visible_batch(
    batch: Tensor | tuple[Tensor, ...] | list[Tensor],
    device: torch.device,
    dtype: torch.dtype,
) -> Tensor:
    """Extract the visible tensor from a batch returned by a dataloader."""

    if isinstance(batch, Tensor):
        visible = batch
    elif isinstance(batch, (tuple, list)) and len(batch) >= 1:
        visible = batch[0]
    else:
        raise TypeError("Unsupported batch type returned by the dataloader.")
    return _prepare_visible_tensor(visible, device=device, dtype=dtype)


def covariance_matrix(visible: Tensor) -> Tensor:
    """Compute the empirical covariance matrix of visible configurations."""

    centered = visible - visible.mean(dim=0, keepdim=True)
    return centered.t() @ centered / visible.shape[0]


def make_binary_dataloader(
    data: Tensor,
    batch_size: int,
    shuffle: bool = True,
) -> DataLoader:
    """Create a simple dataloader for already-binary visible data."""

    return DataLoader(data, batch_size=batch_size, shuffle=shuffle, drop_last=True)