geolip_image_encoder_conv.py

"""
GeoLIP Image Encoder - CONV variation
=====================================
Complete trainable model: conv encoder → S^(d-1) → magnitude → constellation → classify.

Classes:
  ConvEncoder:                8-layer conv → D-dim projection
  InternalConstellationCore:  Three-domain head (external + geometric + internal)
  GeoLIPImageEncoder:         Full pipeline: encoder + MagnitudeFlow + core

Usage:
    from geolip_encoder import GeoLIPImageEncoder

    model = GeoLIPImageEncoder(num_classes=100, output_dim=384, n_anchors=2048)
    out = model.forward_paired(v1, v2)
    loss, ld = model.compute_loss(out, targets)

Author: AbstractPhil + Claude Opus 4.6
License: Apache 2.0
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

from geolip_core import (
    Constellation, Patchwork, MagnitudeFlow,
    make_activation, param_count, model_summary,
)
from geolip_losses import (
    cv_loss, cv_metric, spread_loss, attraction_loss,
    nce_loss, ce_loss_paired, bridge_loss_paired,
    assign_bce_loss, assign_nce_loss, knn_accuracy,
)


# ══════════════════════════════════════════════════════════════════
# CONV ENCODER — 8-layer, proven on CIFAR-100
# ══════════════════════════════════════════════════════════════════

class ConvEncoder(nn.Module):
    """8-layer conv → D-dim projection on S^(d-1).

    Architecture: 4 blocks of (conv-BN-GELU, conv-BN-GELU, MaxPool)
    Channels: 64 → 128 → 256 → 384
    Output: (B, output_dim) after linear + LayerNorm

    Note: L2 normalization is NOT applied here — the caller decides
    when to normalize (preserving raw magnitude for MagnitudeFlow).
    """

    def __init__(self, output_dim=256):
        super().__init__()
        self.output_dim = output_dim
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(),
            nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.GELU(),
            nn.Conv2d(128, 128, 3, padding=1), nn.BatchNorm2d(128), nn.GELU(),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1), nn.BatchNorm2d(256), nn.GELU(),
            nn.Conv2d(256, 256, 3, padding=1), nn.BatchNorm2d(256), nn.GELU(),
            nn.MaxPool2d(2),
            nn.Conv2d(256, 384, 3, padding=1), nn.BatchNorm2d(384), nn.GELU(),
            nn.Conv2d(384, 384, 3, padding=1), nn.BatchNorm2d(384), nn.GELU(),
            nn.MaxPool2d(2),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
        )
        self.proj = nn.Sequential(
            nn.Linear(384, output_dim),
            nn.LayerNorm(output_dim),
        )

    def forward(self, x):
        """Returns: (B, output_dim) unnormalized features."""
        return self.proj(self.features(x))


# ══════════════════════════════════════════════════════════════════
# INTERNAL CONSTELLATION CORE — three-domain head
# ══════════════════════════════════════════════════════════════════

class InternalConstellationCore(nn.Module):
    """Constellation with independent internal + external objectives.

    The constellation discovers its own structure. The task head reads it.

    Three domains:
      EXTERNAL:   CE + embedding NCE → task_head, patchwork, encoder
      GEOMETRIC:  patchwork NCE + bridge → patchwork, encoder, anchors
      INTERNAL:   assign + tri NCE + attract + CV + spread → anchors, encoder

    Args:
        num_classes: classification targets
        dim: embedding dimension
        n_anchors: anchors on S^(dim-1)
        n_comp: patchwork compartments
        d_comp: hidden dim per compartment
        anchor_drop: training anchor dropout
        activation: activation function name
        cv_target: target CV for geometric loss
        infonce_temp: embedding NCE temperature
        assign_temp: assignment temperature
        assign_sharpness: BCE target sharpness
    """

    def __init__(
        self,
        num_classes=100,
        dim=256,
        n_anchors=128,
        n_comp=8,
        d_comp=64,
        anchor_drop=0.15,
        activation='squared_relu',
        cv_target=0.22,
        infonce_temp=0.07,
        assign_temp=0.1,
        assign_sharpness=5.0,
    ):
        super().__init__()
        self.num_classes = num_classes
        self.dim = dim
        self.n_anchors = n_anchors
        self.cv_target = cv_target
        self.infonce_temp = infonce_temp
        self.assign_temp = assign_temp
        self.assign_sharpness = assign_sharpness

        self.config = {k: v for k, v in locals().items()
                       if k != 'self' and not k.startswith('_')}

        # Constellation — owns its own geometry
        self.constellation = Constellation(n_anchors, dim, anchor_drop)

        # Patchwork — interprets distance patterns
        self.patchwork = Patchwork(n_anchors, n_comp, d_comp, activation)
        pw_dim = self.patchwork.output_dim

        # Bridge: patchwork predicts constellation's assignment
        self.bridge = nn.Sequential(nn.Linear(pw_dim, n_anchors))

        # Task head: reads assignment + patchwork + embedding
        total_feat = n_anchors + pw_dim + dim
        self.task_head = nn.Sequential(
            nn.Linear(total_feat, pw_dim),
            make_activation(activation),
            nn.LayerNorm(pw_dim),
            nn.Dropout(0.1),
            nn.Linear(pw_dim, num_classes),
        )

        # Buffers
        self.register_buffer('anchor_classes', torch.zeros(n_anchors, dtype=torch.long))
        self.register_buffer('class_centroids', torch.zeros(num_classes, dim))

    def _triangulate(self, emb):
        """emb → (cos, tri, nearest, soft_assign)."""
        anchors_n = F.normalize(self.constellation.anchors, dim=-1)
        cos = emb @ anchors_n.T
        tri = 1.0 - cos
        _, nearest = cos.max(dim=-1)
        soft_assign = F.softmax(cos / self.assign_temp, dim=-1)
        return cos, tri, nearest, soft_assign

    def forward_paired(self, emb1, emb2, mag1=None, mag2=None):
        """Paired forward for training. Returns dict with all intermediates."""
        cos1, tri1, nearest1, assign1 = self._triangulate(emb1)
        cos2, tri2, nearest2, assign2 = self._triangulate(emb2)

        # Magnitude weighting
        tri1_w = tri1 * mag1 if mag1 is not None else tri1
        tri2_w = tri2 * mag2 if mag2 is not None else tri2

        # Patchwork
        pw1 = self.patchwork(tri1_w)
        pw2 = self.patchwork(tri2_w)

        # Bridge
        bridge1 = self.bridge(pw1)
        bridge2 = self.bridge(pw2)

        # Task head
        feat1 = torch.cat([assign1, pw1, emb1], dim=-1)
        feat2 = torch.cat([assign2, pw2, emb2], dim=-1)
        logits1 = self.task_head(feat1)
        logits2 = self.task_head(feat2)

        return {
            'embedding': emb1, 'embedding_aug': emb2,
            'mag1': mag1, 'mag2': mag2,
            'cos1': cos1, 'cos2': cos2,
            'tri1': tri1, 'tri2': tri2,
            'nearest': nearest1,
            'assign1': assign1, 'assign2': assign2,
            'patchwork1': pw1, 'patchwork1_aug': pw2,
            'bridge1': bridge1, 'bridge2': bridge2,
            'logits': logits1, 'logits_aug': logits2,
        }

    def forward(self, emb, mag=None):
        """Single view for eval."""
        out = self.forward_paired(emb, emb, mag, mag)
        return {
            'logits': out['logits'],
            'embedding': emb,
            'magnitude': mag,
            'triangulation': out['tri1'],
            'cos_to_anchors': out['cos1'],
            'nearest': out['nearest'],
            'assignment': out['assign1'],
            'patchwork': out['patchwork1'],
        }

    def compute_loss(self, output, targets,
                     w_ce=1.0, w_nce_emb=0.5,
                     w_nce_pw=1.0, w_bridge=1.0,
                     w_assign=0.5, w_assign_nce=0.25,
                     w_nce_tri=0.5, w_attract=0.25,
                     w_cv=0.01, w_spread=0.01,
                     cv_batched=True):
        """Three-domain cooperative loss.

        Returns:
            total_loss, loss_dict
        """
        ld = {}
        emb1, emb2 = output['embedding'], output['embedding_aug']

        # ── EXTERNAL ──
        l_ce, acc = ce_loss_paired(output['logits'], output['logits_aug'], targets)
        ld['ce'], ld['acc'] = l_ce, acc

        l_nce_emb, nce_emb_acc = nce_loss(emb1, emb2, self.infonce_temp, normalize=False)
        ld['nce_emb'], ld['nce_emb_acc'] = l_nce_emb, nce_emb_acc

        # ── GEOMETRIC ──
        l_nce_pw, nce_pw_acc = nce_loss(
            output['patchwork1'], output['patchwork1_aug'], self.assign_temp, normalize=True)
        ld['nce_pw'], ld['nce_pw_acc'] = l_nce_pw, nce_pw_acc

        l_bridge, bridge_acc = bridge_loss_paired(
            output['bridge1'], output['bridge2'],
            output['assign1'], output['assign2'])
        ld['bridge'], ld['bridge_acc'] = l_bridge, bridge_acc

        # ── INTERNAL ──
        l_assign, assign_ent = assign_bce_loss(output['assign1'], output['cos1'])
        ld['assign'], ld['assign_entropy'] = l_assign, assign_ent

        l_assign_nce, assign_nce_acc = assign_nce_loss(
            output['assign1'], output['assign2'], self.assign_temp)
        ld['assign_nce'], ld['assign_nce_acc'] = l_assign_nce, assign_nce_acc

        l_nce_tri, nce_tri_acc = nce_loss(
            output['tri1'], output['tri2'], 0.1, normalize=True)
        ld['nce_tri'], ld['nce_tri_acc'] = l_nce_tri, nce_tri_acc

        l_attract, nearest_cos = attraction_loss(output['cos1'])
        ld['attract'], ld['nearest_cos'] = l_attract, nearest_cos

        l_cv = cv_loss(emb1, target=self.cv_target, batched=cv_batched)
        ld['cv'] = l_cv

        l_spread = spread_loss(self.constellation.anchors)
        ld['spread'] = l_spread

        # ── kNN ──
        ld['knn_acc'] = knn_accuracy(emb1, targets)

        # ── TOTAL ──
        loss_external = w_ce * l_ce + w_nce_emb * l_nce_emb
        loss_geometric = w_nce_pw * l_nce_pw + w_bridge * l_bridge
        loss_internal = (w_assign * l_assign + w_assign_nce * l_assign_nce
                         + w_nce_tri * l_nce_tri + w_attract * l_attract
                         + w_cv * l_cv + w_spread * l_spread)

        loss = loss_external + loss_geometric + loss_internal

        ld['loss_external'] = loss_external.item()
        ld['loss_geometric'] = loss_geometric.item()
        ld['loss_internal'] = loss_internal.item()
        ld['t_ce'] = l_ce.item()
        ld['t_nce_emb'] = l_nce_emb.item()
        ld['t_nce_pw'] = l_nce_pw.item()
        ld['t_bridge'] = l_bridge.item()
        ld['t_assign'] = l_assign.item()
        ld['t_assign_nce'] = l_assign_nce.item()
        ld['t_nce_tri'] = l_nce_tri.item()
        ld['t_attract'] = l_attract.item()
        ld['total'] = loss

        return loss, ld


# ══════════════════════════════════════════════════════════════════
# GEOLIP IMAGE ENCODER — full pipeline
# ══════════════════════════════════════════════════════════════════

class GeoLIPImageEncoder(nn.Module):
    """Complete GeoLIP model: ConvEncoder → S^(d-1) → MagnitudeFlow → Core.

    Args:
        num_classes: classification targets
        output_dim: embedding dimension on S^(d-1)
        n_anchors: constellation anchors
        n_comp: patchwork compartments
        d_comp: per-compartment hidden dim
        anchor_drop: training anchor dropout
        activation: activation function name
        cv_target: CV loss target
        infonce_temp: embedding NCE temperature
        assign_temp: assignment temperature
        assign_sharpness: BCE sharpness
        mag_hidden: magnitude relay patchwork hidden dim
        mag_heads: unused (API compat)
        mag_layers: relay layers in MagnitudeFlow
        mag_min: minimum magnitude
        mag_max: maximum magnitude
    """

    def __init__(
        self,
        num_classes=100,
        output_dim=384,
        n_anchors=512,
        n_comp=8,
        d_comp=64,
        anchor_drop=0.15,
        activation='squared_relu',
        cv_target=0.22,
        infonce_temp=0.07,
        assign_temp=0.1,
        assign_sharpness=5.0,
        mag_hidden=64,
        mag_heads=4,
        mag_layers=2,
        mag_min=0.1,
        mag_max=5.0,
    ):
        super().__init__()
        self.output_dim = output_dim
        self.config = {k: v for k, v in locals().items()
                       if k != 'self' and not k.startswith('_')}

        self.encoder = ConvEncoder(output_dim)

        self.mag_flow = MagnitudeFlow(
            dim=output_dim, n_anchors=n_anchors,
            hidden_dim=mag_hidden, n_heads=mag_heads, n_layers=mag_layers,
            mag_min=mag_min, mag_max=mag_max, n_comp=n_comp,
        )

        self.core = InternalConstellationCore(
            num_classes=num_classes, dim=output_dim,
            n_anchors=n_anchors, n_comp=n_comp, d_comp=d_comp,
            anchor_drop=anchor_drop, activation=activation,
            cv_target=cv_target, infonce_temp=infonce_temp,
            assign_temp=assign_temp, assign_sharpness=assign_sharpness,
        )

        self._init_encoder_weights()

    def _init_encoder_weights(self):
        for m in self.encoder.modules():
            if isinstance(m, nn.Linear):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None: nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None: nn.init.zeros_(m.bias)
            elif isinstance(m, (nn.BatchNorm2d, nn.LayerNorm)):
                nn.init.ones_(m.weight); nn.init.zeros_(m.bias)

    def _encode(self, x):
        """Pixels → S^(d-1) + per-anchor magnitude."""
        feat = self.encoder(x)
        raw_mag = feat.norm(dim=-1, keepdim=True)
        emb = F.normalize(feat, dim=-1)

        anchors_n = F.normalize(self.core.constellation.anchors, dim=-1)
        tri = emb @ anchors_n.T
        mag, mag_comp = self.mag_flow(emb, tri, raw_mag)

        return emb, mag, mag_comp

    def forward_paired(self, v1, v2):
        """Training: two views → full pipeline."""
        emb1, mag1, mc1 = self._encode(v1)
        emb2, mag2, mc2 = self._encode(v2)
        out = self.core.forward_paired(emb1, emb2, mag1, mag2)
        out['mag_comp1'] = mc1
        out['mag_comp2'] = mc2
        return out

    def forward(self, x):
        """Eval: single view → classify."""
        emb, mag, mag_comp = self._encode(x)
        out = self.core(emb, mag)
        out['mag_comp'] = mag_comp
        return out

    def compute_loss(self, output, targets, **kwargs):
        """Delegate to core's three-domain loss."""
        return self.core.compute_loss(output, targets, **kwargs)

    def get_anchor_param_ids(self):
        """Return set of param ids that should have weight_decay=0.

        Includes constellation anchors + all relay layer anchors.
        """
        ids = set(id(p) for p in self.core.constellation.parameters())
        for relay in self.mag_flow.relays:
            ids.add(id(relay.anchors))
        return ids

    def make_optimizer(self, lr=3e-4, weight_decay=0.05):
        """Build AdamW with proper anchor exclusion from weight decay."""
        anchor_ids = self.get_anchor_param_ids()
        decay = [p for p in self.parameters() if id(p) not in anchor_ids]
        nodecay = [p for p in self.parameters() if id(p) in anchor_ids]
        return torch.optim.AdamW([
            {'params': decay, 'weight_decay': weight_decay},
            {'params': nodecay, 'weight_decay': 0.0},
        ], lr=lr)

    def summary(self):
        """Print parameter breakdown."""
        print("GeoLIPImageEncoder Summary")
        print("=" * 50)
        param_count(self.encoder, "encoder")
        param_count(self.mag_flow, "mag_flow")
        param_count(self.core.constellation, "constellation")
        param_count(self.core.patchwork, "patchwork")
        param_count(self.core.bridge, "bridge")
        param_count(self.core.task_head, "task_head")
        print("-" * 50)
        total = model_summary(self)
        print(f"\n  Config: {self.config}")
        return total