Comutação de três estados, modulação axis-map(wave) por par de blocos de cubo e persistência em XML

cubosai.py

# cubosai.py ~ comutação de três estados, modulação axis-map(wave) por par de blocos de cubo e persistência em XML

from typing import Tuple, Dict, Any, Optional, List
import xml.etree.ElementTree as ET
import torch.optim as optim
import torch.nn as nn
import numpy as np
import torch
import math
import io

# Sinalizador global padrão ~ classes individuals aceitam o parâmetro @use_complex
DEFAULT_USE_COMPLEX = True

def to_complex_tensor(real: torch.Tensor, imag: torch.Tensor) -> torch.Tensor:

    return torch.stack([real, imag], dim=-1)

def complex_from_tensor(z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

    return z[..., 0], z[..., 1]

def complex_add(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:

    ar, ai = complex_from_tensor(a)
    br, bi = complex_from_tensor(b)
    return to_complex_tensor(ar+br, ai+bi)

def complex_mul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:

    ar, ai = complex_from_tensor(a)
    br, bi = complex_from_tensor(b)
    return to_complex_tensor(ar*br-ai*bi, ar*bi+ai*br)

def complex_abs2(a: torch.Tensor) -> torch.Tensor:

    ar, ai = complex_from_tensor(a)
    return ar*ar+ai*ai

def complex_conj(a: torch.Tensor) -> torch.Tensor:

    r, i = complex_from_tensor(a)
    return to_complex_tensor(r, -i)

def complex_cube(a: torch.Tensor) -> torch.Tensor:

    return complex_mul(complex_mul(a, a), a)

def complex_real_part(a: torch.Tensor) -> torch.Tensor:

    r, _ = complex_from_tensor(a)
    return r

class CubosaiAxisMap:

    # Armazena parâmetros de onda para um par de blocos(modulação de onda)

    def __init__(self, freq: float = 1.0,
                       phase: float = 0.0,
                       amp: float = 1.0,
                       decay: float = 0.0,
                       tau: float = 0.5,
                       scale: float = 0.1,
                       stochastic_gate: bool = False):
        # Parâmetros de onda ~ cee
        self.freq, self.phase, self.amp, self.decay = freq, phase, amp, decay
        # Parâmetros de probabilidade de gating ~ p = sigmoid((tau-ncc)/scale)
        self.tau, self.scale = tau, scale
        # Permitir influência estocástica opcional no gating(~p para ajustar/empurrar cee ~gate)
        self.stochastic_gate = stochastic_gate
        # Uma modulação final
        self.last_mod = 1.0

    def wave(self, t: float) -> float:

        # Valor de regressão em onda cee ~ (cosseno com decaimento)
        return self.amp*math.exp(-self.decay*t)*math.cos(2*math.pi*self.freq*t+self.phase)

    def prob_from_ncc(self, ncc: float) -> float:

        return 1.0/(1.0+math.exp(-((self.tau-ncc)/max(1e-8, self.scale))))

    def compute_modulation(self, t: float, ncc: float) -> float:

        p, w = self.prob_from_ncc(ncc), self.wave(t)
        self.last_mod = float(p*w)
        return self.last_mod

    def to_dict(self) -> Dict[str, float]:

        return {'freq': self.freq,
                'phase': self.phase,
                'amp': self.amp,
                'decay': self.decay,
                'tau': self.tau,
                'scale': self.scale,
                'stochastic_gate': self.stochastic_gate}

    @classmethod
    def from_dict(cls, d: Dict[str, float]):

        return cls(freq=float(d.get('freq', 1.0)),
                   phase=float(d.get('phase', 0.0)),
                   amp=float(d.get('amp', 1.0)),
                   decay=float(d.get('decay', 0.0)),
                   tau=float(d.get('tau', 0.5)),
                   scale=float(d.get('scale', 0.1)),
                   stochastic_gate=bool(d.get('stochastic_gate', False)))

class CubosaiBlock(nn.Module):

    # CubeBlock: 6 faces, cada face é um tensor complexo 4D ou tensor real dependendo do uso de @use_complex
    # Formato das faces ~ (6, d0, d1, d2) ~ cee, armazenadas como parâmetros reais/imaginários cee complexas

    def __init__(self, d0: int, d1: int, d2: int, device=None, use_complex: bool=DEFAULT_USE_COMPLEX, gate_smooth_tau: float=2.0):

        super().__init__()
        self.sides_shape, self.use_complex, dev = (6, d0, d1, d2), use_complex, device
        # Lados aprendíveis
        if self.use_complex:
            self.side_r = nn.Parameter(torch.randn(*self.sides_shape, device=dev)*0.01)
            self.side_i = nn.Parameter(torch.randn(*self.sides_shape, device=dev)*0.01)
        else:
            # Quando apenas real cee, armazenar um único parâmetro e tratar a ~parte imaginária como zero
            self.side_r, self.side_i = nn.Parameter(torch.randn(*self.sides_shape, device=dev)*0.01), None
        # Estados da máscara de gating: 0(desligado), 1(ligado) <~> 2(transicionalj)
        self.register_buffer('gate', torch.zeros(*self.sides_shape, dtype=torch.uint8, device=dev))
        # Valores suavizados dos gates; float 0..1
        self.register_buffer('gate_smooth', torch.zeros(*self.sides_shape, dtype=torch.float32, device=dev))
        # Informações de centróide/mediana por face
        self.register_buffer('centroid_r', torch.zeros(6, device=dev))
        self.register_buffer('centroid_i', torch.zeros(6, device=dev))
        # **Parâmetros de histerese cee
        self.hyst_on = 0.6  # limiar relativo, ~ligar
        self.hyst_off = 0.4 # limiar relativo, !desligar
        # **Constante de tempo de suavização para ~suavização ~exponencial do @gate_smooth(maior ~! mais lento)
        self.gate_smooth_tau, self._eps = gate_smooth_tau, 1e-8

    def forward(self) -> torch.Tensor:

        if self.use_complex:
            return to_complex_tensor(self.side_r, self.side_i)
        else:
            # Criar tensor complexo com zeros na parte imaginária ~ para compatibilidade
            imag = torch.zeros_like(self.side_r)
            return to_complex_tensor(self.side_r, imag)

    def compute_centroids(self):

        z = self.forward()
        r, i = complex_from_tensor(z)
        self.centroid_r[:] = r.view(6, -1).mean(dim=1).detach()
        self.centroid_i[:] = i.view(6, -1).mean(dim=1).detach()
        return self.centroid_r, self.centroid_i

    def update_gates_by_ncc(self, base_threshold: float=0.1,
                                  adapt_per_side: bool=True,
                                  median_scale: float=0.5,
                                  dt: float=1.0,
                                  axis_map: Optional[CubosaiAxisMap]=None):

        # Calcular limiar adaptativo por lado com base na magnitude mediana
        # De histerese: separar limiares de ligado/desligado dependendo:
        #  ~Suivização exponencial rumo ao estado discreto do gate:
        #   com empurrão opcional [AxisMap.stochastic_gate]
        r = self.side_r
        if self.use_complex:
            i = self.side_i
            # mag = sqrt(r^2v+i^2+{∃}) é cee complexa regularizada
            mag = torch.sqrt(r*r+i*i+self._eps)
        else: mag = torch.abs(r)+self._eps
        mag_flat = mag.view(6, -1)
        med = mag_flat.median(dim=1).values.view(6, 1, 1, 1)
        mmax = mag_flat.amax(dim=1).view(6, 1, 1, 1).clamp(min=self._eps)
        if adapt_per_side:
            thr = base_threshold*(med*median_scale+mmax*(1.0-median_scale))
        else: thr = base_threshold*mmax
        # Normalizar a mag em relação a mmax para decisão por histerese
        mag_norm = mag/mmax
        # Calcular o estado discreto desejado do gate com a histerese
        prev_gate = self.gate.clone()
        # O transitório padrão
        desired = torch.full_like(self.gate, 2, dtype=torch.uint8)
        on_cond = mag_norm >= self.hyst_on
        off_cond = mag_norm <= self.hyst_off
        desired = torch.where(on_cond, torch.tensor(1, dtype=torch.uint8, device=self.gate.device), desired)
        desired = torch.where(off_cond, torch.tensor(0, dtype=torch.uint8, device=self.gate.device), desired)
        # Aplicar empurrão estocástico opcional a partir da probabilidade de @axis_map
        if axis_map is not None and axis_map.stochastic_gate:
            # Cee, computar @ncc escalar po lado como a distância média aos cantos normalizada:
            ncc = (1.0-(med.view(6)/(mmax.view(6)+self._eps))).cpu().numpy()
            p = axis_map.prob_from_ncc(float(np.mean(ncc)))
            # Empurrar alguns elementos transitórios de forma probabilística: p in (0, 1)
            rand = torch.rand_like(self.gate, dtype=torch.float32)
            nudge_on = (rand < float(p)) & (desired == 2)
            desired = torch.where(nudge_on, torch.tensor(1, dtype=torch.uint8, device=self.gate.device), desired)
        # **Atualizar o gate discreto com suavização das transições:
        #   Quando mudar para ligado/desligado, permitir estado
        #   cee; transitório por uma atualização igual a ~ (2)
        new_gate = prev_gate.clone()
        # Posições onde desired != ~ definir ~ 2 primeiro, exceto se igual
        changed = desired != prev_gate
        new_gate = torch.where(changed, torch.tensor(2, dtype=torch.uint8, device=self.gate.device), prev_gate)
        stable = ~changed
        new_gate = torch.where(stable, desired, new_gate)
        # Confirmando nos buffers
        self.gate[:] = new_gate
        # Atualizar suavidade: movimento exponencial em direção ao alvo, cee ~ (0/1)
        target = (desired == 1).to(torch.float32)
        alpha = 1.0-math.exp(-dt/max(1e-6, self.gate_smooth_tau))
        self.gate_smooth[:] = (1.0-alpha)*self.gate_smooth+alpha*target
        # Agora retornar métrica resumo por lado: mag média normalizada e média normalizada ~ do gate suavizado
        per_side_mag = mag_flat.mean(dim=1).detach().cpu().numpy().tolist()
        per_side_gate_smooth = self.gate_smooth.view(6, -1).mean(dim=1).detach().cpu().numpy().tolist()
        return per_side_mag, per_side_gate_smooth

    def to_dict(self) -> Dict[str, Any]:

        d = {}
        d['side_r'] = self.side_r.detach().cpu().numpy().tolist()
        if self.use_complex:
            d['side_i'] = self.side_i.detach().cpu().numpy().tolist()
        else: d['side_i'] = None
        d['gate'] = self.gate.cpu().numpy().tolist()
        d['gate_smooth'] = self.gate_smooth.cpu().numpy().tolist()
        d['centroid_r'] = self.centroid_r.cpu().numpy().tolist()
        d['centroid_i'] = self.centroid_i.cpu().numpy().tolist()
        d['use_complex'] = bool(self.use_complex)
        return d

    def load_dict(self, d: Dict[str, Any]):
        self.side_r.data = torch.tensor(d['side_r'], device=self.side_r.device, dtype=self.side_r.dtype)
        if d.get('side_i') is not None and self.use_complex:
            self.side_i.data = torch.tensor(d['side_i'], device=self.side_i.device, dtype=self.side_i.dtype)
        elif d.get('side_i') is not None and not self.use_complex:
            # Cee ~ atualmente for apenas real, então ignorar a parte imag armazenada
            pass
        self.gate.data = torch.tensor(d['gate'], device=self.gate.device, dtype=self.gate.dtype)
        self.gate_smooth.data = torch.tensor(d.get('gate_smooth', self.gate_smooth.cpu().numpy().tolist()), device=self.gate_smooth.device, dtype=self.gate_smooth.dtype)
        self.centroid_r.data = torch.tensor(d['centroid_r'], device=self.centroid_r.device, dtype=self.centroid_r.dtype)
        self.centroid_i.data = torch.tensor(d['centroid_i'], device=self.centroid_i.device, dtype=self.centroid_i.dtype)

class AxisProjector(nn.Module):

    # Projeções lineares complexas aprendidas:

    def __init__(self, in_dim: int, out_dim: int, device=None, use_complex: bool = DEFAULT_USE_COMPLEX):

        super().__init__()
        dev = device
        self.use_complex = use_complex
        # Cee forem entradas, tratá-las como vetores complexos achatados de
        # comprimento @in_dim ~ (última dimensão 2) ~ armazenar pesos para
        # 'matmul-complexo' como partes real/imag separadas:
        self.Wr = nn.Parameter(torch.randn(out_dim, in_dim, device=dev)*0.02)
        self.Wi = nn.Parameter(torch.randn(out_dim, in_dim, device=dev)*0.02)
        self.br = nn.Parameter(torch.zeros(out_dim, device=dev))
        self.bi = nn.Parameter(torch.zeros(out_dim, device=dev))

    def forward(self, z: torch.Tensor) -> torch.Tensor:

        # Z-shapo ~ (..., in_dim, 2) | (in_dim,2)
        r, i = complex_from_tensor(z)
        # Garantir que os formatos de [~r] e [~i] sejam compatíveis com as expectativas do matmul(in_dim, ...):
        # Cee .....?(out_dim, in_dim) ~ @(in_dim, N) => (out_dim, N) ~~~~~~
        # (in_dim, 1) | (in_dim,) com 'matmul' vai broadcastar corretamente
        real = torch.matmul(self.Wr, r)-torch.matmul(self.Wi, i)+self.br.unsqueeze(-1)
        imag = torch.matmul(self.Wr, i)+torch.matmul(self.Wi, r)+self.bi.unsqueeze(-1)
        return to_complex_tensor(real, imag)

class CubosaiNetwork(nn.Module):

    # Montar os blocos, projetores, axis maps e classificador:

    def __init__(self, block_count: int, side_dims: Tuple[int,int,int], proj_dim: int, num_classes: int, device=None, use_complex: bool = DEFAULT_USE_COMPLEX):

        super().__init__()
        d0, d1, d2 = side_dims
        self.device = device
        self.use_complex = use_complex
        self.blocks = nn.ModuleList([CubosaiBlock(d0,d1,d2, device=device, use_complex=use_complex) for _ in range(block_count)])
        flat_side_len = 6*d0*d1*d2
        self.proj = nn.ModuleList()
        for _ in range(block_count-1):
            self.proj.append(AxisProjector(flat_side_len*2, proj_dim, device=device, use_complex=use_complex))
        self.axis_maps: List[CubosaiAxisMap] = [CubosaiAxisMap() for _ in range(block_count-1)]
        # Classificador mapeando magnitudes concatenadas dos projetores para logits:
        self.classifier_r = nn.Linear((block_count-1)*proj_dim, num_classes, device=device)
        self.classifier_i = nn.Linear((block_count-1)*proj_dim, num_classes, device=device)
        # Armazenar a última modulação por projetor
        self.last_mods = [1.0 for _ in range(block_count-1)]
        # Incrementar o contador de passos para os ~ciclos
        self._t = 0.0

    def step_cycle(self, dt: float=1.0, gate_threshold: float=0.1):

        # Avançar o tempo, atualizar centróides e calcular modulações do axis map
        self._t+=dt
        per_pair_ncc, per_block_ncc, per_block_gate_smooth = [], [], []
        for idx, blk in enumerate(self.blocks):
            blk.compute_centroids()
            # Passar o @axis_map associado para empurrão estocástico ~ cee disponível(usar média dos vizinhos)
            axis_map = None
            # Escolher um @axis_map vizinho ~ cee existir, para influenciar este bloco:
            if len(self.axis_maps) > 0:
                if idx == 0: axis_map = self.axis_maps[0]
                elif idx == len(self.blocks)-1: axis_map = self.axis_maps[-1]
                else: axis_map = self.axis_maps[idx-1]
            mag, gate_smooth = blk.update_gates_by_ncc(base_threshold=gate_threshold, axis_map=axis_map, dt=dt)
            per_block_ncc.append(mag)
            per_block_gate_smooth.append(gate_smooth)
        # Calcular resumo NCC por par e modular os axis maps:
        for i in range(len(self.proj)):
            # Usar a NCC média entre os lados de ambos os blocks como @ncc de resumo
            ncc_A, ncc_B = np.mean(per_block_ncc[i]), np.mean(per_block_ncc[i+1])
            ncc_summary = float(0.5*(ncc_A+ncc_B))
            mod = self.axis_maps[i].compute_modulation(self._t, ncc_summary)
            self.last_mods[i] = mod
            per_pair_ncc.append(ncc_summary)
        return per_block_ncc, per_block_gate_smooth, per_pair_ncc, list(self.last_mods)

    def forward(self, x: Optional[torch.Tensor] = None) -> torch.Tensor:

        # Calcular projeções entre blocos adjacentes para aplicar fatores de modulação:
        projs = []
        for i in range(len(self.proj)):
            # Cee ~ (6, d0, d1, d2, 2)
            A, B = self.blocks[i].forward(), self.blocks[i+1].forward()
            # Achatar A e B em vetores
            Af, Bf = A.view(-1, 2), B.view(-1, 2)
            vec = torch.cat([Af, Bf], dim=0)  # (in_dim*2, 2)
            # ...(in_dim, 1, 2) | ...(in_dim, 2) aceitável?
            z = vec.unsqueeze(0)
            # AxisProjector.forward ~ espera última dimensão 2; fornecer (in_dim, 1, 2) e pegar a primeira fatia
            p = self.proj[i](z[0])
            # Garantir que p tenha formato (out_dim, 1, 2) | (out_dim, N, 2); normalizar para(out_dim, N, 2)
            if p.ndim == 2: p = p.unsqueeze(-1)
            # Magnitude ao quadrado por @out_dim cee: somar sobre os canais real/imag
            mag2 = complex_abs2(p).sum(dim=-2).squeeze(-1)
            # Finalmente aplicar a modulação, última modulação escalar
            mod = float(self.last_mods[i]) if i < len(self.last_mods) else 1.0
            mag2 = mag2*mod
            projs.append(mag2)
        if len(projs) == 0:
            feat = torch.zeros(1, self.classifier_r.in_features, device=self.classifier_r.weight.device)
        else: feat = torch.cat(projs, dim=0).unsqueeze(0)
        rlogits, _ilogits = self.classifier_r(feat), self.classifier_i(feat)
        logits = rlogits
        return logits

    def save_to_xml(self, path: str):

        root = ET.Element('CubosaiNetwork')
        meta = ET.SubElement(root, 'Meta')
        meta_use_complex = ET.SubElement(meta, 'UseComplex')
        meta_use_complex.text = repr(bool(self.use_complex))
        blocks_el = ET.SubElement(root, 'Blocks')
        for idx, blk in enumerate(self.blocks):
            be = ET.SubElement(blocks_el, 'Block', idx=str(idx))
            bd = blk.to_dict()
            for k,v in bd.items():
                sub = ET.SubElement(be, k)
                sub.text = repr(v)
        maps_el = ET.SubElement(root, 'AxisMaps')
        for idx, am in enumerate(self.axis_maps):
            me, md = ET.SubElement(maps_el, 'AxisMap', idx=str(idx)), am.to_dict()
            for k,v in md.items():
                sub = ET.SubElement(me, k)
                sub.text = repr(v)
        cls_el = ET.SubElement(root, 'Classifier')
        cls_el_r_w = ET.SubElement(cls_el, 'Wr')
        cls_el_r_w.text = repr(self.classifier_r.weight.detach().cpu().numpy().tolist())
        cls_el_r_b = ET.SubElement(cls_el, 'Br')
        cls_el_r_b.text = repr(self.classifier_r.bias.detach().cpu().numpy().tolist())
        cls_el_i_w = ET.SubElement(cls_el, 'Wi')
        cls_el_i_w.text = repr(self.classifier_i.weight.detach().cpu().numpy().tolist())
        cls_el_i_b = ET.SubElement(cls_el, 'Bi')
        cls_el_i_b.text = repr(self.classifier_i.bias.detach().cpu().numpy().tolist())
        tree = ET.ElementTree(root)
        tree.write(path)

    def load_from_xml(self, path: str):

        tree = ET.parse(path)
        root = tree.getroot()
        meta = root.find('Meta')
        if meta is not None:
            uc = meta.find('UseComplex')
            if uc is not None: loaded_uc = eval(uc.text)
        blocks_el = root.find('Blocks')
        if blocks_el is None:
            raise RuntimeError('Missing Blocks in XML')
        for be in blocks_el.findall('Block'):
            idx = int(be.attrib['idx'])
            d = {}
            for k in ['side_r','side_i','gate','gate_smooth','centroid_r','centroid_i','use_complex']:
                node = be.find(k)
                if node is None:
                    continue
                d[k] = eval(node.text)
            self.blocks[idx].load_dict(d)
        maps_el = root.find('AxisMaps')
        if maps_el is not None:
            ams = []
            for me in maps_el.findall('AxisMap'):
                md = {}
                for k in ['freq','phase','amp','decay','tau','scale','stochastic_gate']:
                    node = me.find(k)
                    if node is None:
                        continue
                    md[k] = eval(node.text)
                ams.append(AxisMap.from_dict(md))
            if len(ams) == len(self.axis_maps):
                self.axis_maps = ams
        cls_el = root.find('Classifier')
        if cls_el is not None:
            Wr = eval(cls_el.find('Wr').text)
            Br = eval(cls_el.find('Br').text)
            Wi = eval(cls_el.find('Wi').text)
            Bi = eval(cls_el.find('Bi').text)
            self.classifier_r.weight.data = torch.tensor(Wr, device=self.classifier_r.weight.device, dtype=self.classifier_r.weight.dtype)
            self.classifier_r.bias.data = torch.tensor(Br, device=self.classifier_r.bias.device, dtype=self.classifier_r.bias.dtype)
            self.classifier_i.weight.data = torch.tensor(Wi, device=self.classifier_i.weight.device, dtype=self.classifier_i.weight.dtype)
            self.classifier_i.bias.data = torch.tensor(Bi, device=self.classifier_i.bias.device, dtype=self.classifier_i.bias.dtype)

def cubosai_run():
    device = torch.device('cpu')
    net = CubosaiNetwork(block_count=3, side_dims=(2,2,2), proj_dim=8, num_classes=3, device=device, use_complex=True)
    opt = optim.Adam(net.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()
    net.step_cycle(dt=0.5, gate_threshold=0.2)
    for epoch in range(361):
        net.train()
        net.step_cycle(dt=0.1, gate_threshold=0.2)
        logits = net()
        target = torch.tensor([1], dtype=torch.long)
        loss = criterion(logits, target)
        opt.zero_grad()
        loss.backward()
        opt.step()
        print(f"epoch {epoch} loss {loss.item():.4f} mods {net.last_mods}")
    #net.save_to_xml('cubosai_modal.xml')
    #net.load_from_xml('cubosai_modal.xml')
cubosai_run()