Upload tests/test_v4.py

tests/test_v4.py

@@ -0,0 +1,245 @@
+"""
+V4 integration tests for Q-TensorFormer.
+
+Tests QKAN DARUAN activations, energy-aware training,
+and the combined v4 pipeline.
+"""
+
+import torch
+import sys
+import os
+
+# Add src to path for testing
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from src.qkan import DARUAN, QKANLayer, HQKANFFN, create_qkan_ffn
+from src.energy_v4 import (
+    EnergyEstimatorV4, ParetoTracker, HARDWARE_PROFILES,
+    estimate_model_energy, HardwareProfile
+)
+from src.config import ModelConfig, TrainingConfig, BudgetConfig
+
+
+def test_daruan_basic():
+    """Test DARUAN activation on scalar input."""
+    daruan = DARUAN(n_repeats=3, base_activation="silu")
+    x = torch.randn(10)
+    out = daruan(x)
+    assert out.shape == (10,), f"Expected (10,), got {out.shape}"
+    assert not torch.isnan(out).any(), "NaN in DARUAN output"
+    print("✓ DARUAN basic: PASSED")
+
+
+def test_daruan_batched():
+    """Test DARUAN on batched tensor."""
+    daruan = DARUAN(n_repeats=5, base_activation="gelu")
+    x = torch.randn(32, 128)
+    out = daruan(x)
+    assert out.shape == (32, 128), f"Expected (32, 128), got {out.shape}"
+    assert not torch.isnan(out).any(), "NaN in DARUAN output"
+    print("✓ DARUAN batched: PASSED")
+
+
+def test_qkan_layer():
+    """Test QKANLayer as drop-in for Linear + Activation."""
+    layer = QKANLayer(128, 256, n_repeats=3)
+    x = torch.randn(16, 128)
+    out = layer(x)
+    assert out.shape == (16, 256), f"Expected (16, 256), got {out.shape}"
+
+    params = layer.parameter_count()
+    dense_params = 128 * 256 + 256  # weight + bias
+    print(f"  QKAN params: {params} vs dense: {dense_params} ({(1 - params/dense_params)*100:.1f}% reduction)")
+    print("✓ QKANLayer: PASSED")
+
+
+def test_hqkan_ffn():
+    """Test HQKAN FFN as drop-in for transformer FFN."""
+    ffn = HQKANFFN(hidden_dim=128, ff_multiplier=4, n_repeats=3)
+    x = torch.randn(8, 64, 128)  # (batch, seq_len, d_model)
+    out = ffn(x)
+    assert out.shape == (8, 64, 128), f"Expected (8, 64, 128), got {out.shape}"
+    print(f"  HQKAN FFN params: {ffn.total_params}")
+    print("✓ HQKAN FFN: PASSED")
+
+
+def test_create_qkan_ffn():
+    """Test factory function for all QKAN FFN variants."""
+    # Standard
+    ffn_std = create_qkan_ffn(128, 4, n_repeats=3)
+    x = torch.randn(4, 32, 128)
+    out = ffn_std(x)
+    assert out.shape == (4, 32, 128)
+    print("✓ create_qkan_ffn (standard): PASSED")
+
+    # TT-QKAN hybrid
+    ffn_tt = create_qkan_ffn(128, 4, n_repeats=3, use_tt=True, tt_rank=4)
+    out = ffn_tt(x)
+    assert out.shape == (4, 32, 128), f"Expected (4, 32, 128), got {out.shape}"
+    print("✓ create_qkan_ffn (TT-hybrid): PASSED")
+
+
+def test_energy_estimator():
+    """Test hardware-aware energy estimator."""
+    est = EnergyEstimatorV4("cpu_intel_xeon")
+
+    # Compute energy for a model forward pass
+    flops = 1e9  # 1 GFLOP
+    energy = est.compute_energy(flops, batch_size=16, memory_gb=0.5)
+    assert energy > 0, f"Energy should be positive, got {energy}"
+    print(f"  Energy for 1 GFLOP on CPU: {energy:.2f} μJ")
+
+    # Carbon footprint
+    carbon = est.carbon_footprint(energy)
+    assert carbon > 0, f"Carbon should be positive"
+    print(f"  Carbon: {carbon:.6f} g CO2")
+    print("✓ EnergyEstimator: PASSED")
+
+
+def test_energy_all_hardware():
+    """Test energy estimation across all hardware targets."""
+    est = EnergyEstimatorV4()
+    flops = 1e9
+
+    print("  Hardware comparison (1 GFLOP):")
+    for hw_name in ["cpu_intel_xeon", "cpu_apple_m2", "gpu_a100", "edge_tpu", "edge_mobile"]:
+        est.set_hardware(hw_name)
+        energy = est.compute_energy(flops, batch_size=16)
+        print(f"    {HARDWARE_PROFILES[hw_name].name}: {energy:.4f} μJ")
+    print("✓ All hardware targets: PASSED")
+
+
+def test_quantum_energy():
+    """Test quantum circuit energy estimation."""
+    est = EnergyEstimatorV4("cpu_intel_xeon")
+    energy = est.quantum_energy(n_qubits=4, n_layers=2, n_tokens=100)
+    assert energy > 0
+    print(f"  Quantum energy (4 qubits, 2 layers, 100 tokens): {energy:.2f} μJ")
+    print("✓ Quantum energy estimation: PASSED")
+
+
+def test_training_energy():
+    """Test total training energy estimate."""
+    est = EnergyEstimatorV4("gpu_a100")
+    result = est.training_energy_estimate(
+        total_flops=1e9,
+        n_epochs=10,
+        batch_size=16,
+        dataset_size=10000,
+        quantum_tokens_per_batch=128,
+        n_qubits=4,
+        n_qlayers=2,
+    )
+    assert "total_energy_uj" in result
+    print(f"  Total training energy: {result['total_energy_j']:.4f} J")
+    print(f"  Carbon: {result['carbon_g']:.4f} g CO2")
+    print(f"  Equivalent smartphone charges: {result['equivalent_smartphone_charges']:.4f}")
+    print("✓ Training energy estimate: PASSED")
+
+
+def test_pareto_tracker():
+    """Test Pareto frontier tracking."""
+    tracker = ParetoTracker()
+
+    # Add some points
+    assert tracker.record(ppl=100, energy_uj=1000, step=0)  # First point always Pareto
+    assert tracker.record(ppl=80, energy_uj=900, step=1)    # Better both → Pareto
+    assert not tracker.record(ppl=90, energy_uj=950, step=2)  # Dominated by (80, 900)
+    assert tracker.record(ppl=75, energy_uj=1200, step=3)   # Better ppl, worse energy → Pareto
+
+    summary = tracker.summary()
+    assert summary["points"] == 3, f"Expected 3 Pareto points, got {summary['points']}"
+    print(f"  Pareto frontier: {summary['frontier']}")
+    print("✓ ParetoTracker: PASSED")
+
+
+def test_budget_integration():
+    """Test budget constraints with energy-aware optimization."""
+    config = ModelConfig(
+        d_model=64, n_layers=2, n_heads=4, tt_rank=4,
+        vocab_size=5000, use_quantum=False,
+    )
+    budget = BudgetConfig(
+        max_params=500000,
+        max_latency_ms=50.0,
+        max_energy_per_query=100.0,
+    )
+
+    # Validate configs
+    config.validate()
+    budget.validate()
+
+    print(f"  Model config: d={config.d_model}, layers={config.n_layers}")
+    print(f"  Budget: params≤{budget.max_params}, latency≤{budget.max_latency_ms}ms, energy≤{budget.max_energy_per_query}μJ")
+    print("✓ Budget integration: PASSED")
+
+
+def test_e2e_v4_pipeline():
+    """End-to-end v4 pipeline test."""
+    from src.models import create_model
+    from src.config import ModelConfig
+    from src.energy_v4 import estimate_model_energy, EnergyEstimatorV4
+
+    config = ModelConfig(
+        vocab_size=1000,
+        d_model=64,
+        n_layers=2,
+        n_heads=4,
+        tt_rank=4,
+        max_seq_len=64,
+        n_qubits=4,
+        use_quantum=False,  # Skip quantum for basic test
+    )
+
+    model = create_model(config, model_type="qtensor")
+
+    # Forward pass
+    x = torch.randint(0, 1000, (2, 16))
+    logits = model(x)
+    assert logits.shape == (2, 16, 1000), f"Expected (2, 16, 1000), got {logits.shape}"
+
+    # Energy estimate
+    est = EnergyEstimatorV4("cpu_apple_m2")
+    est_result = estimate_model_energy(model, est, seq_len=64, batch_size=2)
+    print(f"  E2E energy: {est_result['energy_uj']:.2f} μJ")
+    print(f"  E2E carbon: {est_result['carbon_per_query_ug']:.4f} μg CO2")
+    print(f"  E2E params: {est_result['params']}")
+    print("✓ End-to-end v4 pipeline: PASSED")
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("Q-TensorFormer v4 — Integration Tests")
+    print("=" * 60)
+
+    tests = [
+        ("DARUAN basic", test_daruan_basic),
+        ("DARUAN batched", test_daruan_batched),
+        ("QKANLayer", test_qkan_layer),
+        ("HQKAN FFN", test_hqkan_ffn),
+        ("create_qkan_ffn", test_create_qkan_ffn),
+        ("EnergyEstimator", test_energy_estimator),
+        ("All Hardware", test_energy_all_hardware),
+        ("Quantum Energy", test_quantum_energy),
+        ("Training Energy", test_training_energy),
+        ("ParetoTracker", test_pareto_tracker),
+        ("Budget Integration", test_budget_integration),
+        ("E2E v4 Pipeline", test_e2e_v4_pipeline),
+    ]
+
+    passed = 0
+    failed = 0
+    for name, test_fn in tests:
+        try:
+            test_fn()
+            passed += 1
+        except Exception as e:
+            print(f"✗ {name}: FAILED — {e}")
+            failed += 1
+
+    print(f"\n{'=' * 60}")
+    print(f"Results: {passed}/{passed + failed} tests passed")
+    if failed:
+        print(f"FAILED: {failed} test(s)")
+    else:
+        print("✅ ALL TESTS PASSED")