path: root/sdil/rain_ep_adapter.py

"""No-autograd structured-bias adapter for the Rain EP implementation.

The adapter monkey-patches only the estimator's two-state parameter-gradient
measurement.  Rain's interaction objects already compute dense, convolutional
and bias energy derivatives by explicit local tensor operations, so no reverse
mode is introduced here.
"""

from __future__ import annotations

from types import MethodType
from typing import Iterable

import torch

from sdil.two_state_debias import (
    BatchedLocalAffineDebiaser,
    LocalAffineDebiaser,
)


Tensor = torch.Tensor


class LocalStructuredBias:
    """Fixed per-element bias affine in a local first-state measurement."""

    def __init__(self, ratio: float, seed: int = 1729) -> None:
        if ratio < 0.0:
            raise ValueError("bias ratio must be nonnegative")
        self.ratio = ratio
        self.seed = seed
        self._metadata: list[tuple[Tensor, Tensor, Tensor]] | None = None

    @torch.no_grad()
    def _initialize(self, local_states: list[Tensor]) -> None:
        self._metadata = []
        for tensor_index, state in enumerate(local_states):
            scale = state.square().mean().sqrt().clamp_min(1e-6)
            flat_index = torch.arange(
                state.numel(), dtype=state.dtype, device=state.device
            ).reshape(state.shape)
            phase = flat_index + float(self.seed + 97 * tensor_index)
            offset = torch.where(
                torch.remainder(phase, 2.0) < 1.0,
                torch.full_like(state, -0.5),
                torch.full_like(state, 0.5),
            )
            slope = 0.5 + torch.remainder(phase * 0.61803398875, 1.0)
            amplitude = self.ratio * scale
            self._metadata.append((scale, offset, amplitude * slope))

    @torch.no_grad()
    def measure(self, local_states: Iterable[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
        local_states = list(local_states)
        if any(state.requires_grad for state in local_states):
            raise ValueError("local bias state must be detached")
        if self._metadata is None:
            self._initialize(local_states)
        if len(local_states) != len(self._metadata):
            raise ValueError("parameter collection changed after bias initialization")
        bases = []
        biases = []
        for state, (scale, offset, scaled_slope) in zip(
            local_states, self._metadata
        ):
            basis = torch.tanh(state / scale)
            amplitude = self.ratio * scale
            bias = amplitude * offset + scaled_slope * basis
            bases.append(basis)
            biases.append(bias)
        return bases, biases


class RainGradientCorrector:
    """Apply clean/raw/constant/SDIL/oracle/noise measurement policies."""

    MODES = {
        "clean", "raw", "constant", "innovation", "oracle", "same_rms_noise"
    }

    def __init__(
        self,
        *,
        mode: str,
        bias_ratio: float,
        predictor_rate: float = 0.05,
        neutral_cadence: int = 1,
        seed: int = 1729,
    ) -> None:
        if mode not in self.MODES:
            raise ValueError(f"unrecognized correction mode {mode}")
        if neutral_cadence < 0:
            raise ValueError("neutral cadence must be nonnegative")
        self.mode = mode
        self.predictor_rate = predictor_rate
        self.neutral_cadence = neutral_cadence
        self.bias = LocalStructuredBias(bias_ratio, seed)
        self.debiaser: LocalAffineDebiaser | None = None
        self.steps = 0
        self.last_diagnostics: dict[str, float | int] = {}
        self._noise_generators: list[torch.Generator] | None = None
        self.seed = seed

    @torch.no_grad()
    def _initialize_debiaser(self, templates: list[Tensor]) -> None:
        self.debiaser = LocalAffineDebiaser(
            templates,
            feature_centers=[0.0] * len(templates),
            feature_scales=[1.0] * len(templates),
            affine=self.mode == "innovation",
        )

    @torch.no_grad()
    def observe_neutral(self, local_states: Iterable[Tensor]) -> None:
        """Fit the local bias field from one instruction-off observation."""
        if self.mode not in {"constant", "innovation"}:
            raise ValueError(
                "neutral predictor observations require constant or innovation mode")
        local_states = list(local_states)
        if any(value.requires_grad for value in local_states):
            raise ValueError("Rain adapter received a requires-grad tensor")
        bases, bias = self.bias.measure(local_states)
        if self.debiaser is None:
            self._initialize_debiaser(local_states)
        self.debiaser.update_neutral(bases, bias, self.predictor_rate)

    @torch.no_grad()
    def apply(self, clean: Iterable[Tensor], local_states: Iterable[Tensor]) -> list[Tensor]:
        clean = list(clean)
        local_states = list(local_states)
        if any(value.requires_grad for value in clean + local_states):
            raise ValueError("Rain adapter received a requires-grad tensor")
        if self.mode == "clean":
            return [value.clone() for value in clean]
        bases, bias = self.bias.measure(local_states)
        measured = [value + corruption for value, corruption in zip(clean, bias)]
        if self.mode == "raw":
            corrected = measured
        elif self.mode == "oracle":
            corrected = [value.clone() for value in clean]
        elif self.mode == "same_rms_noise":
            if self._noise_generators is None:
                self._noise_generators = []
                for index, template in enumerate(clean):
                    generator = torch.Generator(device=template.device)
                    generator.manual_seed(self.seed + 1009 * index)
                    self._noise_generators.append(generator)
            corrected = []
            for value, corruption, generator in zip(
                clean, bias, self._noise_generators
            ):
                noise = torch.randn(
                    value.shape, dtype=value.dtype, device=value.device,
                    generator=generator)
                noise.mul_(corruption.square().mean().sqrt())
                corrected.append(value + noise)
        else:
            if self.debiaser is None:
                self._initialize_debiaser(clean)
            if (
                self.neutral_cadence > 0
                and self.steps % self.neutral_cadence == 0
            ):
                self.debiaser.update_neutral(
                    bases, bias, self.predictor_rate)
            corrected = self.debiaser.residual(bases, measured)
        residual_bias = [
            value - target for value, target in zip(corrected, clean)
        ]
        total_elements = sum(value.numel() for value in bias)
        clean_square = sum(float(value.square().sum()) for value in clean)
        bias_square = sum(float(value.square().sum()) for value in bias)
        residual_square = sum(
            float(value.square().sum()) for value in residual_bias)
        clean_rms = (clean_square / total_elements) ** 0.5
        bias_rms = (bias_square / total_elements) ** 0.5
        residual_bias_rms = (residual_square / total_elements) ** 0.5
        self.last_diagnostics = {
            "step": self.steps,
            "clean_rms": clean_rms,
            "bias_rms": bias_rms,
            "residual_bias_rms": residual_bias_rms,
            "bias_to_clean_rms": bias_rms / max(clean_rms, 1e-30),
            "residual_to_clean_rms": residual_bias_rms / max(clean_rms, 1e-30),
            "neutral_observations": (
                0 if self.debiaser is None else self.debiaser.neutral_observations
            ),
        }
        self.steps += 1
        return corrected


def attach_to_rain_estimator(estimator, corrector: RainGradientCorrector):
    """Replace Rain's standard two-state measurement with a corrected one."""

    @torch.no_grad()
    def corrected_standard_param_grads(self, layers_first, layers_second):
        for layer in self._layers:
            layer.state = layers_first[layer.name]
        grads_first = [updater.grad() for updater in self._param_updaters]
        for layer in self._layers:
            layer.state = layers_second[layer.name]
        grads_second = [updater.grad() for updater in self._param_updaters]
        denominator = self._second_nudging - self._first_nudging
        clean = [
            (second - first) / denominator
            for first, second in zip(grads_first, grads_second)
        ]
        return corrector.apply(clean, grads_first)

    estimator._standard_param_grads = MethodType(
        corrected_standard_param_grads, estimator)
    estimator.sdil_corrector = corrector
    return estimator


class DillavouUpdateCorrector:
    """Fixed per-parameter update bias from Dillavou et al. Eq. (7).

    The corruption is added *after* the two-state EP estimator has formed its
    local parameter update.  It is therefore independent of the sign or
    centering of beta.  ``constant`` and ``innovation`` intentionally coincide
    when ``drift_ratio`` is zero: under the paper's strictly constant bias
    model, the Harnett-style predictor reduces to a local intercept estimate.

    ``drift_ratio`` is an optional state-dependent extension.  It is kept
    separate from the exact Dillavou model so reports cannot conflate them.
    """

    MODES = RainGradientCorrector.MODES

    def __init__(
        self,
        *,
        mode: str,
        bias_ratio: float,
        predictor_rate: float = 1.0,
        neutral_cadence: int = 1,
        drift_ratio: float = 0.0,
        seed: int = 1729,
    ) -> None:
        if mode not in self.MODES:
            raise ValueError(f"unrecognized correction mode {mode}")
        if bias_ratio < 0.0 or drift_ratio < 0.0:
            raise ValueError("Dillavou bias ratios must be nonnegative")
        if not 0.0 < predictor_rate <= 1.0:
            raise ValueError("predictor rate must lie in (0, 1]")
        if neutral_cadence < 0:
            raise ValueError("neutral cadence must be nonnegative")
        self.mode = mode
        self.bias_ratio = bias_ratio
        self.predictor_rate = predictor_rate
        self.neutral_cadence = neutral_cadence
        self.drift_ratio = drift_ratio
        self.seed = seed
        self.steps = 0
        self.debiaser: LocalAffineDebiaser | None = None
        self._offsets: list[Tensor] | None = None
        self._slopes: list[Tensor] | None = None
        self._centers: list[Tensor] | None = None
        self._scales: list[Tensor] | None = None
        self._noise_generators: list[torch.Generator] | None = None
        self.last_diagnostics: dict[str, float | int | str] = {}

    @staticmethod
    @torch.no_grad()
    def _unit_pattern(template: Tensor, phase: float) -> Tensor:
        index = torch.arange(
            template.numel(), dtype=template.dtype, device=template.device
        ).reshape(template.shape)
        pattern = torch.sin((index + phase) * 1.61803398875)
        pattern.add_(0.5 * torch.cos((index + phase) * 0.754877666))
        return pattern / pattern.square().mean().sqrt().clamp_min(1e-30)

    @torch.no_grad()
    def _initialize(
        self, clean: list[Tensor], parameter_states: list[Tensor]
    ) -> None:
        if len(clean) != len(parameter_states):
            raise ValueError("gradient and parameter collections disagree")
        self._offsets = []
        self._slopes = []
        self._centers = []
        self._scales = []
        feature_scales = []
        for index, (gradient, parameter) in enumerate(
            zip(clean, parameter_states)
        ):
            gradient_scale = gradient.square().mean().sqrt().clamp_min(1e-12)
            parameter_scale = parameter.square().mean().sqrt().clamp_min(1e-6)
            offset_pattern = self._unit_pattern(
                gradient, float(self.seed + 97 * index))
            slope_pattern = self._unit_pattern(
                gradient, float(self.seed + 193 * index + 41))
            self._offsets.append(
                self.bias_ratio * gradient_scale * offset_pattern)
            self._slopes.append(
                self.drift_ratio * gradient_scale * slope_pattern)
            self._centers.append(parameter.clone())
            self._scales.append(parameter_scale.clone())
            feature_scales.append(torch.ones_like(parameter_scale))
        self.debiaser = LocalAffineDebiaser(
            clean,
            feature_centers=[0.0] * len(clean),
            feature_scales=feature_scales,
            affine=self.mode == "innovation",
        )

    @torch.no_grad()
    def _measure(
        self, parameter_states: list[Tensor]
    ) -> tuple[list[Tensor], list[Tensor]]:
        features = []
        biases = []
        for parameter, center, scale, offset, slope in zip(
            parameter_states,
            self._centers,
            self._scales,
            self._offsets,
            self._slopes,
        ):
            feature = torch.tanh((parameter - center) / scale)
            features.append(feature)
            biases.append(offset + slope * feature)
        return features, biases

    @torch.no_grad()
    def apply(
        self, clean: Iterable[Tensor], parameter_states: Iterable[Tensor]
    ) -> list[Tensor]:
        clean = list(clean)
        parameter_states = list(parameter_states)
        if any(value.requires_grad for value in clean + parameter_states):
            raise ValueError("Dillavou adapter received a requires-grad tensor")
        if self.mode == "clean":
            return [value.clone() for value in clean]
        if self._offsets is None:
            self._initialize(clean, parameter_states)
        features, bias = self._measure(parameter_states)
        measured = [
            value + corruption for value, corruption in zip(clean, bias)
        ]
        if self.mode == "raw":
            corrected = measured
        elif self.mode == "oracle":
            corrected = [value.clone() for value in clean]
        elif self.mode == "same_rms_noise":
            if self._noise_generators is None:
                self._noise_generators = []
                for index, template in enumerate(clean):
                    generator = torch.Generator(device=template.device)
                    generator.manual_seed(self.seed + 1009 * index)
                    self._noise_generators.append(generator)
            corrected = []
            for value, corruption, generator in zip(
                clean, bias, self._noise_generators
            ):
                noise = torch.randn(
                    value.shape,
                    dtype=value.dtype,
                    device=value.device,
                    generator=generator,
                )
                noise.mul_(corruption.square().mean().sqrt())
                corrected.append(value + noise)
        else:
            if (
                self.neutral_cadence > 0
                and self.steps % self.neutral_cadence == 0
            ):
                self.debiaser.update_neutral(
                    features, bias, self.predictor_rate)
            corrected = self.debiaser.residual(features, measured)
        residual_bias = [
            value - target for value, target in zip(corrected, clean)
        ]
        total_elements = sum(value.numel() for value in clean)
        clean_square = sum(float(value.square().sum()) for value in clean)
        bias_square = sum(float(value.square().sum()) for value in bias)
        residual_square = sum(
            float(value.square().sum()) for value in residual_bias)
        clean_rms = (clean_square / total_elements) ** 0.5
        bias_rms = (bias_square / total_elements) ** 0.5
        residual_rms = (residual_square / total_elements) ** 0.5
        self.last_diagnostics = {
            "step": self.steps,
            "bias_model": (
                "dillavou_constant_update"
                if self.drift_ratio == 0.0
                else "dillavou_plus_local_state_drift"
            ),
            "clean_update_rms": clean_rms,
            "bias_update_rms": bias_rms,
            "residual_update_rms": residual_rms,
            "bias_to_clean_update_rms": (
                bias_rms / max(clean_rms, 1e-30)),
            "residual_to_clean_update_rms": (
                residual_rms / max(clean_rms, 1e-30)),
            "neutral_observations": self.debiaser.neutral_observations,
        }
        self.steps += 1
        return corrected


def attach_dillavou_to_rain_estimator(
    estimator, corrector: DillavouUpdateCorrector
):
    """Add the hardware update offset after Rain's local EP estimate."""

    @torch.no_grad()
    def corrected_standard_param_grads(self, layers_first, layers_second):
        for layer in self._layers:
            layer.state = layers_first[layer.name]
        grads_first = [updater.grad() for updater in self._param_updaters]
        for layer in self._layers:
            layer.state = layers_second[layer.name]
        grads_second = [updater.grad() for updater in self._param_updaters]
        denominator = self._second_nudging - self._first_nudging
        clean = [
            (second - first) / denominator
            for first, second in zip(grads_first, grads_second)
        ]
        parameter_states = [parameter.state for parameter in self._params]
        return corrector.apply(clean, parameter_states)

    estimator._standard_param_grads = MethodType(
        corrected_standard_param_grads, estimator)
    estimator.sdil_corrector = corrector
    return estimator


@torch.no_grad()
def observe_rain_neutral(estimator, corrector: RainGradientCorrector) -> None:
    """Expose one label-free Rain equilibrium to the local predictor."""
    local_states = [updater.grad() for updater in estimator._param_updaters]
    corrector.observe_neutral(local_states)


class RainLayerStateCorrector:
    """Correct a structured per-neuron bias before Rain's local EP update."""

    MODES = RainGradientCorrector.MODES

    def __init__(
        self,
        *,
        mode: str,
        bias_ratio: float,
        predictor_rate: float = 0.2,
        calibration_steps: int = 1,
        bias_normalization: str = "clean_difference",
        seed: int = 1729,
    ) -> None:
        if mode not in self.MODES:
            raise ValueError(f"unrecognized correction mode {mode}")
        if bias_ratio < 0.0:
            raise ValueError("bias ratio must be nonnegative")
        if calibration_steps < 0:
            raise ValueError("calibration steps must be nonnegative")
        if bias_normalization not in {"clean_difference", "first_state"}:
            raise ValueError("unrecognized layer-bias normalization")
        self.mode = mode
        self.bias_ratio = bias_ratio
        self.predictor_rate = predictor_rate
        self.calibration_steps = calibration_steps
        self.bias_normalization = bias_normalization
        self.seed = seed
        self.steps = 0
        self.last_diagnostics: dict[str, float | int] = {}
        self.debiaser: BatchedLocalAffineDebiaser | None = None
        self._metadata: list[tuple[Tensor, Tensor, Tensor]] | None = None
        self._noise_generators: list[torch.Generator] | None = None

    @torch.no_grad()
    def _initialize(
        self, first_states: list[Tensor], clean_differences: list[Tensor]
    ) -> None:
        self._metadata = []
        for index, (first, clean) in enumerate(zip(
            first_states, clean_differences
        )):
            scale = first.square().mean(dim=0).sqrt().clamp_min(1e-6)
            flat_index = torch.arange(
                scale.numel(), dtype=first.dtype, device=first.device
            ).reshape(scale.shape)
            phase = flat_index + float(self.seed + 97 * index)
            offset = torch.where(
                torch.remainder(phase, 2.0) < 1.0,
                torch.full_like(scale, -0.5),
                torch.full_like(scale, 0.5),
            )
            slope = 0.5 + torch.remainder(phase * 0.61803398875, 1.0)
            basis = torch.tanh(first / scale)
            source = offset.unsqueeze(0) + slope.unsqueeze(0) * basis
            source_rms = source.square().mean().sqrt().clamp_min(1e-30)
            reference_rms = (
                clean.square().mean().sqrt()
                if self.bias_normalization == "clean_difference"
                else first.square().mean().sqrt()
            )
            gain = self.bias_ratio * reference_rms / source_rms
            self._metadata.append(
                (scale, gain * offset, gain * slope)
            )
        self.debiaser = BatchedLocalAffineDebiaser(
            first_states,
            feature_centers=[0.0] * len(first_states),
            feature_scales=[1.0] * len(first_states),
            affine=self.mode == "innovation",
        )

    @torch.no_grad()
    def _measure(
        self, first_states: list[Tensor]
    ) -> tuple[list[Tensor], list[Tensor]]:
        bases = []
        biases = []
        for first, (scale, offset, slope) in zip(
            first_states, self._metadata
        ):
            basis = torch.tanh(first / scale)
            bases.append(basis)
            biases.append(
                offset.unsqueeze(0) + slope.unsqueeze(0) * basis)
        return bases, biases

    @torch.no_grad()
    def apply(
        self,
        layers_first: dict[str, Tensor],
        layers_second: dict[str, Tensor],
        layer_names: list[str],
    ) -> dict[str, Tensor]:
        if self.mode == "clean":
            return dict(layers_second)
        first = [layers_first[name] for name in layer_names]
        second = [layers_second[name] for name in layer_names]
        if any(value.requires_grad for value in first + second):
            raise ValueError("Rain layer adapter received a requires-grad tensor")
        clean = [after - before for before, after in zip(first, second)]
        if self._metadata is None:
            self._initialize(first, clean)
        bases, bias = self._measure(first)
        measured = [value + corruption for value, corruption in zip(clean, bias)]
        if self.mode == "raw":
            corrected = measured
        elif self.mode == "oracle":
            corrected = [value.clone() for value in clean]
        elif self.mode == "same_rms_noise":
            if self._noise_generators is None:
                self._noise_generators = []
                for index, template in enumerate(clean):
                    generator = torch.Generator(device=template.device)
                    generator.manual_seed(self.seed + 1009 * index)
                    self._noise_generators.append(generator)
            corrected = []
            for value, corruption, generator in zip(
                clean, bias, self._noise_generators
            ):
                noise = torch.randn(
                    value.shape, dtype=value.dtype, device=value.device,
                    generator=generator)
                noise.mul_(corruption.square().mean().sqrt())
                corrected.append(value + noise)
        else:
            if self.steps < self.calibration_steps:
                self.debiaser.update_neutral(
                    bases, bias, self.predictor_rate)
            corrected = self.debiaser.residual(bases, measured)
        residual_bias = [
            value - target for value, target in zip(corrected, clean)
        ]
        total_elements = sum(value.numel() for value in clean)
        clean_square = sum(float(value.square().sum()) for value in clean)
        bias_square = sum(float(value.square().sum()) for value in bias)
        residual_square = sum(
            float(value.square().sum()) for value in residual_bias)
        clean_rms = (clean_square / total_elements) ** 0.5
        bias_rms = (bias_square / total_elements) ** 0.5
        residual_rms = (residual_square / total_elements) ** 0.5
        self.last_diagnostics = {
            "step": self.steps,
            "clean_state_difference_rms": clean_rms,
            "bias_state_difference_rms": bias_rms,
            "residual_state_difference_rms": residual_rms,
            "bias_to_clean_state_difference_rms": (
                bias_rms / max(clean_rms, 1e-30)),
            "residual_to_clean_state_difference_rms": (
                residual_rms / max(clean_rms, 1e-30)),
            "neutral_observations": (
                0 if self.debiaser is None else self.debiaser.neutral_observations),
        }
        used_second = dict(layers_second)
        for name, before, difference in zip(layer_names, first, corrected):
            used_second[name] = before + difference
        self.steps += 1
        return used_second


def attach_layer_to_rain_estimator(
    estimator, corrector: RainLayerStateCorrector
):
    """Patch Rain immediately before its hand-written local parameter rule."""
    layer_names = [layer.name for layer in estimator._layers[1:]]

    @torch.no_grad()
    def corrected_standard_param_grads(self, layers_first, layers_second):
        used_second = corrector.apply(
            layers_first, layers_second, layer_names)
        for layer in self._layers:
            layer.state = layers_first[layer.name]
        grads_first = [updater.grad() for updater in self._param_updaters]
        for layer in self._layers:
            layer.state = used_second[layer.name]
        grads_second = [updater.grad() for updater in self._param_updaters]
        denominator = self._second_nudging - self._first_nudging
        return [
            (second - first) / denominator
            for first, second in zip(grads_first, grads_second)
        ]

    estimator._standard_param_grads = MethodType(
        corrected_standard_param_grads, estimator)
    estimator.sdil_corrector = corrector
    return estimator