docs/ARCHITECTURE.md

# System Architecture (Hardware Agnostic)

Deep technical reference for the SciMLx autonomous research loop components, optimized for both NVIDIA GPUs (PyTorch) and Apple Silicon (MLX).

---

## 3-Tier Scientific Implementation (SI) Layer

SciMLx utilizes a modular SI layer in `core/` to decouple scientific logic from underlying hardware and compute frameworks. The layer is organized into three tiers:

### Tier 1: Foundations
Core math and device abstractions that provide a stable, hardware-agnostic base.
- **`device.py` (Hardware Agnostic Dispatch)**: Automatically detects the best available backend (CUDA, MLX, MPS, or CPU) and provides a single API for framework-agnostic tensor creation (`to_array()`) and device placement (`to_device()`).
- **`units.py` (`SciMLTensor`)**: Ensures mathematical and physical groundedness by performing dimensional analysis on every operation using the `pint` unit registry.
- **`lie_math.py` & `heat_kernels.py` (Geometric Math)**: Implements Lie Algebra foundations and mesh-based heat kernel signatures for geometry-aware operators.
- **`oracle_constants.py` (Buckingham Pi Theorem)**: Identifies dimensionless groups (like Reynolds or Péclet numbers) to aid in feature discovery and similarity analysis.

### Tier 2: Models
Production-grade neural operators and the infrastructure to build them.
- **Dual-Backend Operators**: All new models are implemented with both PyTorch (`_torch.py`) and MLX (`_mlx.py`) backends, managed by a central dispatcher (e.g., `models/mff.py`).
- **`scaffold.py` (Automated Scaffolding)**: Generates dual-backend model stubs from research proposals, ensuring feature parity across hardware.
- **`losses.py` (Physics-Informed Operators)**: Framework-agnostic implementations of Sobolev ($H^1$, $H^2$) and Spectral losses that penalize unphysical oscillations.
- **`spectral_governor.py` (Frequency Governance)**: Dynamically monitors the Fourier spectrum of residuals across backends and adjusts loss weighting to ensure high-frequency features are captured.

### Tier 3: Production
Systems for scaling, deploying, and automating the research cycle.
- **`deployment.py`**: Integration with Google Cloud Platform (Vertex AI, Compute Engine) for serverless GPU training.
- **`model_versioning.py`**: A centralized model registry and lineage tracking system for managing champion models.
- **`hpo.py` (Bayesian Optimization)**: Automated hyperparameter search that adapts to hardware-specific constraints.
- **`dp_federated.py` (Differential Privacy)**: Logic for secure, privacy-preserving federated training of scientific models.
- **`arxiv_agent.py` (ASIL Pipeline)**: Orchestrates the Agentic Scientist Ideation Loop, from literature review to automated model scaffolding.

---

## Two-Mode Operation

The system supports two operating modes that can be mixed within a session:

**Mode A — Human-Guided**  
A human (or AI agent) reads `RESEARCH_BRAIN.md`, interprets results, edits `experiments.yaml` directly, and invokes `autorun.py` to execute the queue. The system handles execution, retry, and logging; the human handles strategy.

**Mode B — Fully Autonomous (`agent_loop.py`)**  
`agent_loop.py` performs one full autonomous cycle:
1. Calls `tracker.analyze_lineage()` to build per-benchmark summaries.
2. Calls `HypothesisEngine.analyze_benchmark()` for each priority benchmark.
3. Calls `BayesianHPO.ask()` to sample hyperparameters.
4. Generates new `ExperimentConfig` entries and appends them to `experiments.yaml`.
5. Triggers `autorun.py` to process the queue on NVIDIA GPUs.

---

## Unified Trainer (`core/trainer.py`)

The trainer is designed to be high-performance while remaining flexible across backends:

### Compute Optimizations
- **NVIDIA/PyTorch**: Utilizes `torch.compile()` for kernel fusion and `torch.amp` for mixed precision training.
- **Apple/MLX**: Leverages MLX's lazy evaluation and unified memory for efficient processing on M-series chips.
- **Precision Management**: Configurable precision levels (float32, bfloat16) mapped to hardware-specific best practices.

### Training Logic
- **EMA (Exponential Moving Average)**: Maintains a shadow copy of model weights for more stable evaluation.
- **Dynamic Budget Extension**: Automatically extends training time by 20% if the loss is still decreasing significantly at the end of the budget.
- **Snapshot Ensembling**: Optionally saves and averages multiple model states throughout the run.

---

## Hardware-Accelerated PDE Solvers (`data/simulations/`)

All PDE solvers are implemented using framework-native spectral methods to ensure high-speed simulation on the active device:
- **Spectral Methods**: Utilize fast Fourier transforms (`torch.fft` or `mlx.fft`) for high-speed spectral derivatives and integration.
- **Zero-Copy Data**: Solvers execute directly on the `DEVICE`, producing tensors that never leave high-speed device memory during training.
- **Batch Processing**: All simulations are vectorized to solve multiple initial conditions in parallel, maximizing device throughput.

---

## High-Throughput Data Pipeline

The I/O bottleneck is eliminated through:
1. **`PDEDataset`**: An `IterableDataset` that interfaces with cached `.npz` files or on-the-fly solvers.
2. **`DataLoader`**: Standard PyTorch implementation with:
   - `pin_memory=True`: For faster Host-to-Device transfer.
   - `num_workers > 0`: For multi-process data pre-fetching.
   - `prefetch_factor`: To keep the GPU saturated.

---

## HypothesisEngine (`core/hypothesis.py`)

The engine classifies experiment outcomes to guide follow-up logic:

| Mode | Detection Logic |
|---|---|
| `gradient_collapse` | `val_l2_rel ≥ 1.0`, NaN loss, or CUDA launch errors. |
| `spectral_bias` | High-frequency error > 0.3 in spectral diagnostic. |
| `capacity_limited` | Small model (`hidden_dim < 64`) with high error. |
| `cuda_oom` | Log analysis detects "Out of Memory" on GPU. |

---

## Retry Escalation (`autorun.py`)

Escalates through recovery levels for NVIDIA environments:
- **r1 — `smart_fix()`**: Detects CUDA OOM and automatically halves `hidden_dim`.
- **r2**: Aggressive reduction of `hidden_dim`, `n_layers`, and `n_modes`.
- **r3**: Minimal viable fallback (`h=32, l=2, lr=1e-4`).

---

## Cloud Infrastructure (GCP)

Configured for project `gdpr-494411`:
- **Vertex AI**: Custom container execution using the project's Artifact Registry.
- **Compute Engine**: G2-standard instances with NVIDIA L4 GPUs for development.