CEA

cea.py

#//////////////////////////////////////////////////////////////////////////////////////////////


# cea.py / !!! Co-Evolutionary Auto-Encoder !!!


#//////////////////////////////////////////////////////////////////////////////////////////////

import os
import json
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from typing import Dict, List, Tuple
from collections import defaultdict
import matplotlib.pyplot as plt
from datetime import datetime
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#//////////////////////////////////////////////////////////////////////////////////////////////

# >• Compressor / Auto-encoder
class Compressor(nn.Module):

    def __init__(self, input_dim: int, bottleneck: int):
        super().__init__()
        self.enc = nn.Sequential(
            nn.Linear(input_dim, max(32, input_dim * 2)),
            nn.ReLU(),
            nn.Linear(max(32, input_dim * 2), bottleneck))
        self.dec = nn.Sequential(
            nn.Linear(bottleneck, max(32, input_dim * 2)),
            nn.ReLU(),
            nn.Linear(max(32, input_dim * 2), input_dim))

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        z = self.enc(x)
        recon = self.dec(z)
        return z, recon

# >• Channel allowance for gating & gradient dampening
class ChannelAllowance(nn.Module):

    def __init__(self, latent_dim: int, channels: int, hidden: int = 64, eps: float = 1e-6):
        super().__init__()
        self.attn = nn.Sequential(
            nn.Linear(latent_dim, max(hidden, latent_dim)),
            nn.ReLU(),
            nn.Linear(max(hidden, latent_dim), channels),
            nn.Sigmoid())
        self.eps = eps

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        if z.dim() == 1:
            z = z.unsqueeze(0)
        gates = self.attn(z)
        return gates.squeeze(0)

    def dampen_grad_hook(self, gate_vec: torch.Tensor, attenuation: float = 0.9):
        gate_vec = gate_vec.detach()
        def hook(grad: torch.Tensor):
            damping = 1.0 - (attenuation * (1.0 - gate_vec.to(grad.device)))
            if damping.numel() == grad.shape[0]:
                view = [grad.shape[0]] + [1] * (grad.dim() - 1)
                return grad * damping.view(*view)
            return grad * damping
        return hook

# >• NodeState as parameterized module
class NodeState(nn.Module):

    def __init__(self, coeff_dim: int):
        super().__init__()
        self.coeff_dim = coeff_dim
        self.a = nn.Parameter(torch.zeros(coeff_dim, device=device))
        self.register_buffer("diag_sigma", torch.ones(coeff_dim, device=device) * 1e-2)
        self.w = nn.Parameter(torch.tensor(0.0, device=device)) 
        self.hook_handle = None

    def pack_feature(self, detach_params: bool = False) -> torch.Tensor:
        if detach_params:
            return torch.cat([self.a.detach(), self.diag_sigma, self.w.detach().unsqueeze(0)])
        return torch.cat([self.a, self.diag_sigma, self.w.unsqueeze(0)])

# >• MetaController
class MetaController(nn.Module):

    def __init__(self,
                 coeff_dim: int,
                 compressor_bottleneck: int = 8,
                 max_members: int = 24,
                 allowance_hidden: int = 64,
                 compression_input_dim: int = None):
        super().__init__()
        self.coeff_dim = coeff_dim
        input_dim = (coeff_dim + coeff_dim + 1) if compression_input_dim is None else compression_input_dim
        self.compressor = Compressor(input_dim=input_dim, bottleneck=compressor_bottleneck).to(device)
        self.allowance = ChannelAllowance(compressor_bottleneck, channels=coeff_dim, hidden=allowance_hidden).to(device)
        self.projector = nn.Sequential(
            nn.Linear(compressor_bottleneck, max(32, compressor_bottleneck * 2)),
            nn.ReLU(),
            nn.Linear(max(32, compressor_bottleneck * 2), coeff_dim)).to(device)
        self.max_members = max_members
        self.members: Dict[int, NodeState] = {}
        self.optim = optim.Adam(self.parameters(), lr=1e-3)

    def attach(self, node_id: int, node: NodeState) -> bool:
        if len(self.members) >= self.max_members: 
            return False
        self.members[node_id] = node
        return True

    def aggregate_features(self, detach_params: bool = True) -> Tuple[torch.Tensor, List[int]]:
        if not self.members:
            return None, []
        feats = [self.members[nid].pack_feature(detach_params=detach_params) for nid in list(self.members.keys())]
        X = torch.stack(feats).to(device)
        z, recon = self.compressor(X)
        mean_z = z.mean(dim=0)
        return mean_z, list(self.members.keys())

    def refine(self, metric_fn_cpu, ea_trigger_delta: float = 1e-3):
        mean_z, ids = self.aggregate_features(detach_params=True)
        if mean_z is None: 
            return
        
        def eval_fn_np(z_np):
            return metric_fn_cpu(z_np)
            
        z_opt, z_score = micro_ea_optimize(eval_fn_np, dim=mean_z.shape[0], pop_size=48, generations=40)
        z_tensor = torch.tensor(z_opt, device=device, dtype=torch.float32)
        with torch.no_grad():
            gates = self.allowance(z_tensor)
            dec = self.compressor.dec(z_tensor.unsqueeze(0))[0]
            latent_proxy = dec[:self.coeff_dim]
            gated = latent_proxy * gates[:self.coeff_dim]
            delta = self.projector(gated.unsqueeze(0)).squeeze(0)
        for nid in ids:
            node = self.members[nid]
            with torch.no_grad():
                node.a.add_(0.01 * delta)
                node.diag_sigma.mul_(0.995)
            if node.hook_handle is not None:
                node.hook_handle.remove()
            if hasattr(node.a, "register_hook"):
                hook_fn = self.allowance.dampen_grad_hook(gates[:self.coeff_dim], attenuation=0.9)
                node.hook_handle = node.a.register_hook(hook_fn)

    def save_state(self, path: str):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        torch.save(self.state_dict(), path + ".pth")
        manifest = {
            "coeff_dim": self.coeff_dim,
            "max_members": self.max_members,
            "member_ids": list(self.members.keys()),
            "members": {str(nid): {
                "a": self.members[nid].a.detach().cpu().tolist(),
                "diag_sigma": self.members[nid].diag_sigma.detach().cpu().tolist(),
                "w": float(self.members[nid].w.detach().cpu().item())
            } for nid in self.members}}
        with open(path + ".json", "w") as f:
            json.dump(manifest, f)

# >• Recursive least-squares updating
def rls_update(node: NodeState, x_vec: torch.Tensor, y: float, lam: float = 0.99):
    P_diag = node.diag_sigma  
    Px = P_diag * x_vec**2
    denom = lam + Px.sum()
    K = (P_diag * x_vec) / denom  
    y_pred = (node.a * x_vec).sum()
    err = y - y_pred
    with torch.no_grad():
        node.a.add_(K * err)
        node.diag_sigma.copy_((P_diag - K * x_vec * P_diag) / lam)

# >• Micro EA
def micro_ea_optimize(eval_fn, dim: int, pop_size: int = 48, generations: int = 40, rng=None):
    if rng is None: 
        rng = np.random.default_rng()
    pop = rng.normal(scale=0.1, size=(pop_size, dim))
    fitness = np.array([eval_fn(ind) for ind in pop])
    for _ in range(generations):
        for i in range(pop_size):
            idxs = [j for j in range(pop_size) if j != i]
            a, b, c = rng.choice(idxs, 3, replace=False)
            F = 0.6
            trial = pop[i] + F * (pop[a] - pop[b]) + F * (pop[c] - pop[a])
            tfit = eval_fn(trial)
            if tfit > fitness[i]:
                pop[i] = trial
                fitness[i] = tfit
    best_idx = np.argmax(fitness)
    return pop[best_idx], fitness[best_idx]

# >• Task metric
def posterior_metric(z_np: np.ndarray) -> float:
    return -float(np.linalg.norm(z_np))

# >• Metrics Tracker
class MetricsTracker:

    def __init__(self):
        self.metrics = defaultdict(list)
        self.step = 0
        
    def log(self, key: str, value: float):
        self.metrics[key].append(value)
    
    def log_dict(self, metrics_dict: Dict[str, float]):
        for key, value in metrics_dict.items():
            self.log(key, value)
    
    def get_latest(self, key: str) -> float:
        if key in self.metrics and self.metrics[key]:
            return self.metrics[key][-1]
        return 0.0
    
    def save(self, path: str):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        metrics_json = {k: v for k, v in self.metrics.items()}
        with open(path, "w") as f:
            json.dump(metrics_json, f, indent=2)
    
    def load(self, path: str):
        if os.path.exists(path):
            with open(path, "r") as f:
                self.metrics = defaultdict(list, json.load(f))

# >• Dashboard Display
def display_dashboard(tracker: MetricsTracker, step: int, interval: int = 50):
    if step % interval != 0:
        return
    print("\n" + "="*80)
    print(f">• EVOLUTION DASHBOARD | Step {step}")
    print("="*80)
    recon_loss = tracker.get_latest("recon_loss")
    print(f"  Reconstruction Loss:      {recon_loss:.6f}")
    attached = tracker.get_latest("attached_nodes")
    mean_weight = tracker.get_latest("mean_node_weight")
    print(f"  Attached Nodes:           {int(attached)} / max capacity")
    print(f"  Mean Node Weight:         {mean_weight:.6f}")
    ea_fitness = tracker.get_latest("ea_fitness")
    ea_improvement = tracker.get_latest("ea_improvement")
    print(f"  EA Best Fitness:          {ea_fitness:.6f}")
    print(f"  EA Improvement:           {ea_improvement:.6f}")
    avg_grad_norm = tracker.get_latest("avg_grad_norm")
    gate_sparsity = tracker.get_latest("gate_sparsity")
    print(f"  Avg Gradient Norm:        {avg_grad_norm:.6f}")
    print(f"  Gate Sparsity (inactive): {gate_sparsity:.2%}")
    latent_norm = tracker.get_latest("mean_latent_norm")
    print(f"  Mean Latent Norm:         {latent_norm:.6f}")
    print("="*80 + "\n")

# >• Plot Generation
def save_plots(tracker: MetricsTracker, path: str = "./saved/plots.png"):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    keys_to_plot = [
        ("recon_loss", "Reconstruction Loss"),
        ("ea_fitness", "EA Fitness Score"),
        ("mean_node_weight", "Mean Node Weight"),
        ("gate_sparsity", "Gate Sparsity"),
        ("avg_grad_norm", "Gradient Norm"),
        ("attached_nodes", "Attached Nodes")]
    fig, axes = plt.subplots(2, 3, figsize=(15, 8))
    axes = axes.flatten()
    for idx, (key, title) in enumerate(keys_to_plot):
        if key in tracker.metrics and tracker.metrics[key]:
            ax = axes[idx]
            values = tracker.metrics[key]
            ax.plot(values, linewidth=1.5, color='steelblue')
            ax.set_title(title, fontsize=11, fontweight='bold')
            ax.set_xlabel('Step (×interval)')
            ax.grid(alpha=0.3)
            ax.set_yscale('log' if 'loss' in key.lower() else 'linear')
    plt.tight_layout()
    plt.savefig(path, dpi=120)
    print(f">• Plots saved to {path}")
    plt.close()

# >• Load MetaController
def load_meta_controller(meta_path: str, 
                         coeff_dim: int, 
                         compressor_bottleneck: int = 8,
                         max_members: int = 24) -> MetaController:
    meta = MetaController(
        coeff_dim=coeff_dim,
        compressor_bottleneck=compressor_bottleneck,
        max_members=max_members).to(device)
    pth_file = meta_path + ".pth"
    if os.path.exists(pth_file):
        meta.load_state_dict(torch.load(pth_file, map_location=device))
        print(f">• Loaded MetaController weights from {pth_file}")
    else:
        print(f">• Warning: {pth_file} not found")
        return None
    json_file = meta_path + ".json"
    if os.path.exists(json_file):
        with open(json_file, "r") as f:
            manifest = json.load(f)
        for node_id_str, node_data in manifest.get("members", {}).items():
            node_id = int(node_id_str)
            node = NodeState(coeff_dim).to(device)
            node.a.data = torch.tensor(node_data["a"], device=device, dtype=torch.float32)
            node.diag_sigma.data = torch.tensor(node_data["diag_sigma"], device=device, dtype=torch.float32)
            node.w.data = torch.tensor(node_data["w"], device=device, dtype=torch.float32)
            meta.attach(node_id, node)
        print(f">• Loaded {len(meta.members)} node states from {json_file}")
    return meta

# >• Load Node States
def load_node_states(json_path: str, coeff_dim: int) -> Dict[int, NodeState]:
    nodes = {}
    if os.path.exists(json_path):
        with open(json_path, "r") as f:
            manifest = json.load(f)
        for node_id_str, node_data in manifest.get("members", {}).items():
            node_id = int(node_id_str)
            node = NodeState(coeff_dim).to(device)
            node.a.data = torch.tensor(node_data["a"], device=device, dtype=torch.float32)
            node.diag_sigma.data = torch.tensor(node_data["diag_sigma"], device=device, dtype=torch.float32)
            node.w.data = torch.tensor(node_data["w"], device=device, dtype=torch.float32)
            nodes[node_id] = node
        print(f">• Loaded {len(nodes)} nodes from {json_path}")
    else:
        print(f">• Warning: {json_path} not found")
    return nodes

# >• Resume from Checkpoint
def resume_from_checkpoint(checkpoint_path: str, 
                           coeff_dim: int = 6,
                           resume_steps: int = 100):
    print(f">• Loading checkpoint from {checkpoint_path}")
    meta = load_meta_controller(checkpoint_path, coeff_dim=coeff_dim)
    if meta is None:
        print(">• Failed to load checkpoint")
        return None, None
    metrics_path = checkpoint_path.replace(".pth", "") + "_metrics.json"
    tracker = MetricsTracker()
    if os.path.exists(metrics_path):
        tracker.load(metrics_path)
        print(f">• Loaded metrics history from {metrics_path}")
    num_steps = len(tracker.metrics.get('recon_loss', []))
    print(f">• Resuming from step {num_steps}")
    print(f">• MetaController has {len(meta.members)} active members")
    return meta, tracker

# >• TRAINING LOOP
def training_loop_monitored(num_nodes: int = 64, 
                            coeff_dim: int = 6, 
                            steps: int = 400, 
                            attach_threshold: float = 0.5,
                            checkpoint_interval: int = 200,
                            dashboard_interval: int = 50):
    print(">• Initializing training environment...")
    nodes: Dict[int, NodeState] = {i: NodeState(coeff_dim).to(device) for i in range(num_nodes)}
    input_dim = coeff_dim + coeff_dim + 1
    meta = MetaController(coeff_dim=coeff_dim,
                          compressor_bottleneck=8,
                          max_members=32,
                          compression_input_dim=input_dim).to(device)
    all_node_params = [n.a for n in nodes.values()] + [n.w for n in nodes.values()]
    node_optim = optim.Adam(all_node_params, lr=1e-3)
    rng = np.random.default_rng(42)
    tracker = MetricsTracker()
    prev_ea_fitness = 0.0
    print(f">• Initialized {num_nodes} nodes with coeff_dim={coeff_dim}")
    print(f">• Device: {device}")
    print(f">• Starting {steps} training steps\n")
    for t in range(steps):
        # ** Measurement Collection & RLS Updating **
        for i, node in nodes.items():
            x_vec = torch.randn(coeff_dim, device=device)
            y = float((node.a.detach() @ x_vec).cpu().numpy() + 0.05 * rng.normal())
            rls_update(node, x_vec, y, lam=0.995)
            with torch.no_grad():
                s = torch.tensor(abs(y), device=device)
                alpha = 0.96
                node.w.mul_(alpha).add_((1.0 - alpha) * s)
            if node.w.item() > attach_threshold and (i not in meta.members):
                meta.attach(i, node)
        if t % 20 == 0:
            tracker.log("attached_nodes", float(len(meta.members)))
            mean_w = float(np.mean([n.w.item() for n in meta.members.values()]) if meta.members else 0.0)
            tracker.log("mean_node_weight", mean_w)
        # ** Macro Space Periodic Black-Box Refinement **
        if t % 8 == 0:
            meta.refine(posterior_metric)
            mean_z, _ = meta.aggregate_features(detach_params=True)
            if mean_z is not None:
                ea_fitness = posterior_metric(mean_z.detach().cpu().numpy())
                improvement = abs(ea_fitness - prev_ea_fitness)
                tracker.log("ea_fitness", float(ea_fitness))
                tracker.log("ea_improvement", float(improvement))
                prev_ea_fitness = ea_fitness
        # ** Structural Backpropagation; Auto-Encoding Consistency **
        if t % 4 == 0 and len(meta.members) > 0:
            feats = [meta.members[nid].pack_feature(detach_params=False) for nid in meta.members]
            X = torch.stack(feats).to(device)
            z, recon = meta.compressor(X)
            recon_loss = nn.functional.mse_loss(recon, X)
            reg = sum(p.norm() * 1e-4 for p in meta.projector.parameters())
            loss = recon_loss + reg
            meta.optim.zero_grad()
            node_optim.zero_grad()
            loss.backward()
            all_grads = []
            for p in meta.parameters():
                if p.grad is not None:
                    all_grads.append(p.grad.detach().norm().item())
            for p in all_node_params:
                if p.grad is not None:
                    all_grads.append(p.grad.detach().norm().item())
            avg_grad_norm = np.mean(all_grads) if all_grads else 0.0
            tracker.log("avg_grad_norm", float(avg_grad_norm))
            tracker.log("recon_loss", float(recon_loss.item()))
            with torch.no_grad():
                mean_z, _ = meta.aggregate_features(detach_params=True)
                if mean_z is not None:
                    gates = meta.allowance(mean_z)
                    inactive = float((gates < 0.1).float().mean().item())
                    tracker.log("gate_sparsity", inactive)
                    tracker.log("mean_latent_norm", float(mean_z.norm().item()))
            torch.nn.utils.clip_grad_norm_(meta.parameters(), max_norm=1.0)
            torch.nn.utils.clip_grad_norm_(all_node_params, max_norm=0.5)
            meta.optim.step()
            node_optim.step()
        display_dashboard(tracker, t, interval=dashboard_interval)
        if t % checkpoint_interval == 0 and t > 0:
            meta.save_state(f"./saved/meta_{t}")
            tracker.save(f"./saved/metrics_{t}.json")
            print(f">• Checkpoint saved at step {t}")
    print("\n>• Finalizing training...")
    meta.save_state("./saved/meta_final")
    tracker.save("./saved/metrics_final.json")
    save_plots(tracker, "./saved/evolution_plots.png")
    print(">• Training complete!")
    return meta, nodes, tracker

# >• EVALUATION & INFERENCE
def evaluate_loaded_model(checkpoint_path: str, coeff_dim: int = 6):
    print(">• Starting model evaluation...\n")
    meta = load_meta_controller(checkpoint_path, coeff_dim=coeff_dim)
    if meta is None:
        return
    print(f"\n>• Attached Nodes: {len(meta.members)}")
    print("─" * 60)
    for node_id, node in sorted(meta.members.items()):
        a_norm = float(node.a.norm().item())
        w_val = float(node.w.item())
        print(f"  Node {node_id:3d} | a_norm: {a_norm:.6f} | weight: {w_val:.6f}")
    mean_z, ids = meta.aggregate_features(detach_params=True)
    if mean_z is not None:
        print(f"\n>• Latent Space Summary:")
        print(f"  Mean latent norm:    {float(mean_z.norm().item()):.6f}")
        print(f"  Latent dimensions:   {mean_z.shape[0]}")
        gates = meta.allowance(mean_z)
        print(f"\n>• Gate Activation:")
        print(f"  Active gates:        {int((gates > 0.5).sum().item())} / {gates.shape[0]}")
        print(f"  Mean gate value:     {float(gates.mean().item()):.6f}")
        print(f"  Gate range:          [{float(gates.min().item()):.4f}, {float(gates.max().item()):.4f}]")
    print("\n" + "="*60 + "\n")

#//////////////////////////////////////////////////////////////////////////////////////////////

if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "--resume":
        if len(sys.argv) > 2:
            checkpoint = sys.argv[2]
            print(f"\n>• RESUMING FROM CHECKPOINT: {checkpoint}\n")
            meta, tracker = resume_from_checkpoint(checkpoint, coeff_dim=6)
            if meta is not None:
                # Continue training
                print("\n(Checkpoint loaded. Ready for continued training or evaluation.)\n")
        else:
            print("Usage: python full_module.py --resume <checkpoint_path>")
    elif len(sys.argv) > 1 and sys.argv[1] == "--evaluate":
        if len(sys.argv) > 2:
            checkpoint = sys.argv[2]
            print(f"\n>• EVALUATING CHECKPOINT: {checkpoint}\n")
            evaluate_loaded_model(checkpoint, coeff_dim=6)
        else:
            print("Usage: python full_module.py --evaluate <checkpoint_path>")
    else:
        print(">• BEGIN FRESH TRAINING\n")
        meta, nodes, tracker = training_loop_monitored(
            num_nodes=48, 
            coeff_dim=6, 
            steps=200,
            checkpoint_interval=50,
            dashboard_interval=50)
        print("\n>• Training finished successfully!")
        print(f">• Metrics saved to ./saved/metrics_final.json")
        print(f">• Plots saved to ./saved/evolution_plots.png")
        print(f">• Model saved to ./saved/meta_final.pth & .json")
        print(f"\n>• TIP: Resume training with:")
        print(f"   python full_module.py --resume ./saved/meta_final")
        print(f"\n>• TIP: Evaluate model with:")
        print(f"   python full_module.py --evaluate ./saved/meta_final\n")

Oh! of course peace, with islamic country that loves supporting terrorism and just also happens to be Persia.