Create elliptic.data.py

elliptic.data.py

@@ -0,0 +1,94 @@
+# elliptic_data.py — Loader do Elliptic Bitcoin Dataset via PyG
+import torch
+import numpy as np
+import pandas as pd
+from torch_geometric.datasets import EllipticBitcoinDataset
+from torch_geometric.loader import NeighborLoader
+from torch_geometric.transforms import NormalizeFeatures
+import os
+
+def carregar_elliptic(root='/tmp/elliptic', normalize=True):
+    """
+    Carrega o Elliptic Bitcoin Dataset via PyG.
+    
+    Estatísticas reais:
+    - 203,769 nós (transações Bitcoin)
+    - 234,355 arestas (fluxo de Bitcoin)
+    - 166 features por nó (94 locais + 72 agregadas)
+    - 2 classes: ilícito (lavagem) / lícito
+    - 49 timesteps (jan 2017 - set 2018)
+    - ~21% rotulados, ~79% desconhecidos
+    
+    Split temporal (como no paper):
+    - Treino: timesteps 1-34
+    - Teste:  timesteps 35-49
+    """
+    transform = NormalizeFeatures() if normalize else None
+    
+    try:
+        dataset = EllipticBitcoinDataset(root=root, transform=transform)
+        data    = dataset[0]
+        return data, True
+    except Exception as e:
+        return None, str(e)
+
+
+def preparar_splits(data):
+    """
+    Split temporal como descrito no paper original:
+    Treino nos primeiros timesteps, teste nos últimos.
+    Máscara 'unknown' (classe 2) excluída do treino/teste.
+    """
+    # PyG já fornece máscaras train/test no Elliptic
+    # Classe 0 = ilícito, 1 = lícito, 2 = desconhecido
+    
+    # Filtra apenas nós rotulados
+    labeled_mask = data.y != 2
+    train_mask   = data.train_mask & labeled_mask
+    test_mask    = data.test_mask  & labeled_mask
+    
+    # Estatísticas
+    y_train = data.y[train_mask]
+    y_test  = data.y[test_mask]
+    
+    stats = {
+        'n_nos':        data.x.shape[0],
+        'n_arestas':    data.edge_index.shape[1],
+        'n_features':   data.x.shape[1],
+        'n_rotulados':  int(labeled_mask.sum()),
+        'n_train':      int(train_mask.sum()),
+        'n_test':       int(test_mask.sum()),
+        'n_ilicito_train': int((y_train == 0).sum()),
+        'n_licito_train':  int((y_train == 1).sum()),
+        'n_ilicito_test':  int((y_test  == 0).sum()),
+        'n_licito_test':   int((y_test  == 1).sum()),
+        'taxa_fraude_train': float((y_train==0).sum()/len(y_train)),
+        'taxa_fraude_test':  float((y_test ==0).sum()/len(y_test)),
+    }
+    
+    data.train_mask_labeled = train_mask
+    data.test_mask_labeled  = test_mask
+    
+    return data, stats
+
+
+def criar_loaders(data, num_neighbors=[10, 5], batch_size=512):
+    """
+    Mini-batch com NeighborLoader para GraphSAGE inductive.
+    Amostra vizinhos em vez de usar o grafo completo.
+    """
+    train_loader = NeighborLoader(
+        data,
+        num_neighbors=num_neighbors,
+        batch_size=batch_size,
+        input_nodes=data.train_mask_labeled,
+        shuffle=True,
+    )
+    test_loader = NeighborLoader(
+        data,
+        num_neighbors=num_neighbors,
+        batch_size=batch_size,
+        input_nodes=data.test_mask_labeled,
+        shuffle=False,
+    )
+    return train_loader, test_loader