# SOFIA v7.1 — PyTorch implementation

Reference implementation of **SOFIA v7.1** CVGL student model targeting Jetson Orin NX with 500 MB – 1 GB VRAM after INT8 quantization.

Full design rationale: see `2_hypotesis/temp_hypotesis/HYP_SOFIA_v7_UltraDeep_дизайн.md`.

## Architecture

```
Input (256×256×3, sat or UAV)
   ↓
Stem: dual-conv (3→C_mid→C_stem_out), ×2 down
   ↓
Stage 1 (shared sat/UAV): StripDCN-lite × d1, ×2 down
   ↓
Stage 2 (shared): StripMixConv × d2, ×2 down
   ↓
Stage 3 (separate): MambaVision MV5 × d3, ×2 down
   ↓
Stage 4 (separate): MambaVision MV1 × d4, ×2 down (→ 8×8)
   ↓
Ultra-lite 1×1 Neck → F̃ [B, C_n, 8, 8]
   ↓
┌─ Sat-Head: GGeM → BN → Linear → L2 → g_sat
└─ UAV-Head: FiLM(altitude) → CHP (polar+FFT+mag) → BN → Linear → L2 → g_uav
```

## Presets

| Preset | Params | FLOPs | INT8 weights | FP16 weights | Target latency |
|--------|-------:|------:|-------------:|-------------:|---------------:|
| **M** (default) | ~500 M | ~132 G | ~500 MB | ~1 GB | ~18 ms |
| **L** | ~1 B | ~283 G | ~1 GB | ~2 GB | ~20 ms |
| Tiny | ~5 M | ~1.4 G | ~5 MB | ~10 MB | ~4 ms |

*Latency estimates на Jetson Orin NX 8GB (INT8 TRT mixed-precision для DCN/Mamba в FP16).*

## Dependencies

Required:
- Python ≥ 3.9
- PyTorch ≥ 2.0
- torchvision ≥ 0.15 (для `deform_conv2d`)

Optional (для production Mamba speed):
- [`mamba_ssm`](https://github.com/state-spaces/mamba) — ускоренные backends:
  - v1 `selective_scan_fn` для Mamba-1 (~5× vs Python loop)
  - v2 `Mamba2` модуль для Mamba-2 SSD dual form (~2–8× vs Mamba-1)
- [`causal-conv1d`](https://github.com/Dao-AILab/causal-conv1d) — ускоренный 1D conv для Mamba

## Mamba variants (приоритет по качеству/скорости)

| Variant | Описание | Speed | Quality | Когда использовать |
|---------|----------|:-----:|:-------:|--------------------|
| **`mamba2`** (default) | SSD dual form, scalar A per head | **2–8× faster** | best | Main choice если mamba_ssm v2 доступен |
| `mamba1` | Original selective scan, diagonal A | ref | +0% vs mamba2 | Legacy / reproducibility / если v2 недоступен |
| `efficient_vmamba` | Atrous scan (2 directions, no CUDA kernel) | 2–3× faster | ~−0.3% R@1 | Speed fallback без mamba_ssm |

### Выбор через config

```python
from code_sofia_v71 import sofia_m_config, SOFIAv71

cfg = sofia_m_config()

# Вариант 1: Mamba-2 (default, preferred) — если mamba_ssm v2 доступен
cfg.mamba_variant = "mamba2"
cfg.mamba_backend = "auto"   # falls back to torch if unavailable

# Вариант 2: EfficientVMamba — speed без зависимости от mamba_ssm
cfg.mamba_variant = "efficient_vmamba"

# Вариант 3: Mamba-1 — legacy / для сравнения
cfg.mamba_variant = "mamba1"

model = SOFIAv71(cfg)
```

### Проверка доступности backends

```python
from code_sofia_v71 import is_mamba_ssm_available, is_mamba2_available

print(f"Mamba-1 CUDA: {is_mamba_ssm_available()}")
print(f"Mamba-2 CUDA: {is_mamba2_available()}")
```

### Несовместимость параметров между variants

Mamba-1 и Mamba-2 используют **разную параметризацию $A$** (diagonal-per-channel vs scalar-per-head) — state_dict **НЕ совместим**. EfficientVMamba имеет N независимых scanners каждый с собственными параметрами Mamba-1.

**Следствие:** нельзя train в `mamba1` и load в `mamba2` checkpoint. Выбор variant нужно зафиксировать до training.

## Quick start

```python
import torch
from code_sofia_v71 import build_sofia

# Default = SOFIA-M (500 MB INT8 target)
model = build_sofia("M")
model.eval()

sat = torch.randn(2, 3, 256, 256)   # satellite image
uav = torch.randn(2, 3, 256, 256)   # UAV image
altitude = torch.tensor([120.0, 450.0])   # meters

out = model(sat=sat, uav=uav, altitude=altitude)
g_sat = out["g_sat"]   # [2, 512] L2-normalized
g_uav = out["g_uav"]   # [2, 512] L2-normalized

# Retrieval similarity
similarity = (g_sat * g_uav).sum(dim=-1)   # cosine
```

## Verification

```bash
# Smoke test (Tiny preset, CPU, fast)
python -m sofia_v71.verify

# Full SOFIA-M check with latency benchmark (GPU)
python -m sofia_v71.verify --preset M --device cuda --benchmark

# Maximum scale
python -m sofia_v71.verify --preset L --device cuda
```

Expected output for SOFIA-M:
```
Total params: ~500 M
...
FP16 weights: ~1000 MB  (~0.98 GB)
INT8 weights: ~500 MB   (~0.49 GB)
...
out[g_sat]: (1, 512)
out[g_uav]: (1, 512)
```

## Key novel components

### `layers.py`
- **`GGeM`** — per-channel learnable exponent Generalized Mean pooling (F11)
- **`CircularHarmonicPool`** — formally SO(2)-invariant UAV pool via polar → 1D FFT → magnitude (NOVEL NH2)
- **`AltitudeFiLM`** — telemetry-conditioned modulation (NOVEL NH4)
- **`RoPE2D`** — 2D rotary positional embedding

### `blocks.py`
- **`StripDCNLiteBlock`** — Strip DW MBConv с DCN offset на одной оси (NOVEL)
- **`StripMixConvBlock`** — Strip + MixConv (3/5/7 kernels) для multi-scale
- **`MambaVisionBlock`** — Mamba ∥ MHSA [∥ Strip] + FFN (MV5 variant — NOVEL 3-way)
- **`SimpleMambaBlock`** — reference pure-PyTorch Mamba-1 (replace with `mamba_ssm` в production)

### `model.py`
- **`SOFIAv71`** — full model with optional weight-sharing
- **`SatHead` vs `UAVHead`** — asymmetric physics-motivated design (NOVEL NH1)
- **`RingAuxHead`** — LPN Square-Ring training-only aux

## KD taps

Backbone forward returns features at stages s0/f1/f2/f3/f4. Использовать для
hierarchical knowledge distillation:

```python
out = model(sat=sat, uav=uav, return_features=True)
f3_sat = out["features_sat"]["f3"]   # для teacher feature alignment
```

## Production notes

### Mamba backend selection

Default `mamba_variant="mamba2"` с `backend="auto"`:
- Use `Mamba2` from `mamba_ssm.modules.mamba2` если установлено (CUDA, fast)
- Иначе fallback на pure-PyTorch simplified SSD scan (slow, работает)

Для Mamba-1 legacy путь — `mamba_variant="mamba1"`, backend так же резолвится.

Для EfficientVMamba — pure PyTorch, не нуждается в mamba_ssm.

## Quantization (`quant.py`)

Reusable utilities для INT8 PTQ/QAT в SOFIA.

### `OffsetClampSTE` — DCN-M2

Hard clamp DCN offsets к `[-k, +k]` со STE backward. Уже подключён в `StripDCNLiteBlock` через флаг `use_offset_clamp_ste=True` (default).

```python
from code_sofia_v71 import OffsetClampSTE, offset_clamp_ste

# Module form
clamp = OffsetClampSTE(kernel_size=7)
clamped = clamp(offsets)

# Functional form
clamped = offset_clamp_ste(offsets, kernel_size=7)
```

Backward — identity (gradient проходит насквозь). Можно ставить с epoch 1, не только в QAT.

### `KScaledFakeQuant` — k-scaled fake quantization

Multi-bin fake quant для long-tail distributions (Mamba Δ, y).

```python
from code_sofia_v71 import KScaledFakeQuant

fq = KScaledFakeQuant(num_bins=3)

# Calibration
fq.start_calibration()
with torch.no_grad():
    for batch in cal_loader:
        _ = model(batch)
fq.finalize_calibration(percentile=99.9)
# fq теперь в PTQ-режиме
```

### `KScaledMamba2Block` — drop-in для Mamba2Block

Подкласс `Mamba2Block` с k-scaled fake-quant узлами на критических путях
(`x_main`, `delta`, `y`). Внутреннее состояние scan'а `h_t` остаётся в
model dtype (R5 reparam principle).

```python
from code_sofia_v71 import KScaledMamba2Block

mamba = KScaledMamba2Block(
    channels=192,
    d_state=64,
    headdim=64,
    num_bins=3,
    targets=("x_main", "delta", "y"),    # subset
)

mamba.start_calibration()
# ... run calibration data ...
mamba.finalize_calibration(percentile=99.9)
```

**Constraint:** только `backend='torch'` поддерживается (mamba_ssm CUDA kernel
не имеет hooks для k-scaled fake-quant). Для full INT8 deploy на TRT нужен
custom plugin — отдельный deploy concern.

### Model-wide helpers

```python
from code_sofia_v71 import start_calibration, finalize_calibration, set_quant_enabled

# Начать калибровку всех KScaledFakeQuant и KScaledMamba2Block в модели
start_calibration(model)
with torch.no_grad():
    for batch in cal_loader:
        _ = model(batch)
finalize_calibration(model, percentile=99.9)

# Toggle on/off для FP-vs-INT8 ablation
set_quant_enabled(model, False)    # FP forward
y_fp = model(x)
set_quant_enabled(model, True)
y_q = model(x)
```

### Smoke test

```bash
python -m sofia_v71.quant
```

Прогонит unit-test на:
- DCN-M2 clamp (gradient pass-through)
- KScaledFakeQuant calibration + quant error
- KScaledMamba2Block FP-vs-INT8 diff

### DCN INT8 QAT
StripDCN использует `torchvision.ops.deform_conv2d`. Для INT8 deploy на TRT:
- Применить QAT с DCN-M1..M4 modifications (см. HYP Phase 7)
- Offset predictor: per-channel scale, offset clamping `[-k, +k]`
- Mask: FP16 micro-block внутри INT8 graph

### TensorRT export
```python
import torch
model.eval()
model_fuse = model  # apply fuse_reparam passes separately
torch.onnx.export(
    model,
    (sat, uav, altitude),
    "sofia_v71.onnx",
    input_names=["sat", "uav", "altitude"],
    output_names=["g_sat", "g_uav"],
    opset_version=17,
)
# Then:
# trtexec --onnx=sofia_v71.onnx --int8 --fp16 --saveEngine=sofia.plan
```

## Training

See `HYP_SOFIA_v7_UltraDeep_дизайн.md` Phase 8 for full training recipe:
- Loss: InfoNCE (Sample4Geo mining) + Ring aux + (opt) cross-view consistency
- Curriculum: PALW sigmoid warmup, 60 epochs, 4 phases
- Optimizer: AdamW, LR 3e-4 cosine, wd 0.05
- Temperature: τ 0.1 → 0.01 cosine decay

## File layout

```
code_sofia_v71/
├── __init__.py       — public exports
├── config.py         — SOFIAConfig + presets (M/L/Tiny)
├── layers.py         — GGeM, CHP, FiLM, RoPE2D, SE, LayerNorm2d
├── blocks.py         — StripDCN, StripMixConv, Mamba, MambaVision, Downsample
├── model.py          — Stem, Backbone, Neck, Heads, SOFIAv71
├── verify.py         — parameter counter + benchmark script
└── README.md         — this file
```

## References to design doc

Все architectural decisions обоснованы в
`2_hypotesis/temp_hypotesis/HYP_SOFIA_v7_UltraDeep_дизайн.md`:
- Phase 1: requirements R1–R13
- Phase 2': DCN / Strip / MambaVision operator catalog
- Phase 3': backbone design with 4 candidates (E/F/G)
- Phase 4'': CVGL-Aware Head v7.1-α (Asymmetric + CHP + FiLM)
- Phase 5': ablation matrix
- Phase 6: MambaVision MV5 operational details
- Phase 7: StripDCN QAT strategy
- Phase 8: training pipeline
