Upload 2 files

config.py

@@ -0,0 +1,86 @@
+"""
+Central configuration for Rabi Oscillation Quality Classifier.
+V2: Calibrated to real data distributions.
+"""
+import torch
+import os
+
+# ─── Paths ───────────────────────────────────────────────────────────────────
+PROJECT_DIR = os.path.dirname(os.path.abspath(__file__))
+DATA_DIR = os.path.join(PROJECT_DIR, "data copy")
+UI_TEST_DIR = os.path.join(PROJECT_DIR, "ui_test_samples")
+MODEL_PATH = os.path.join(PROJECT_DIR, "best_model.pt")
+CORRECTIONS_DIR = os.path.join(PROJECT_DIR, "data_corrections")
+FAILED_PRED_DIR = os.path.join(PROJECT_DIR, "failed_predictions")
+
+# ─── Device (MPS > CUDA > CPU) ──────────────────────────────────────────────
+if torch.backends.mps.is_available():
+    DEVICE = torch.device("mps")
+elif torch.cuda.is_available():
+    DEVICE = torch.device("cuda")
+else:
+    DEVICE = torch.device("cpu")
+
+# ─── Signal Preprocessing ───────────────────────────────────────────────────
+SEQ_LEN = 256           # Increased from 128 for finer resolution
+NUM_FIT_PARAMS = 4      # amplitude, T, phase, offset
+
+# ─── Model ───────────────────────────────────────────────────────────────────
+NUM_CLASSES = 3         # Classes per head: 0=Bad, 1=Acceptable, 2=Perfect
+IN_CHANNELS = 3         # raw signal + fit curve + residual (signal - fit)
+RESNET_DIM = 256        # ResNet backbone output dim
+DENSE_DIM = 64          # Dense branch output dim
+FUSED_DIM = RESNET_DIM + DENSE_DIM  # 320
+
+# ─── Training ────────────────────────────────────────────────────────────────
+BATCH_SIZE = 256
+LR = 3e-4               # Lower LR for better convergence
+WEIGHT_DECAY = 1e-4
+EPOCHS = 80              # More epochs
+TRAIN_SIZE = 200_000     # Doubled for better generalization
+VAL_SIZE = 40_000
+PATIENCE = 15            # More patience
+FOCAL_GAMMA = 2.0
+CLASS_WEIGHTS = [1.2, 0.9, 0.9]  # Stronger upweight for class 0
+IGNORE_INDEX = -1
+
+# ─── Synthetic Generator — Calibrated to Real Data ──────────────────────────
+# Real data statistics (from analysis of 2736 JSON files):
+#   amplitude: mean=0.090, std=0.083, range=[~0, 0.40], p25=0.0002
+#   T:         mean=0.101 but p75=0.062, most values 0.047-0.110
+#   phase:     range=[-π, π]
+#   offset:    mean=0.185, std=0.166, range=[~0, 0.98]
+#   x_range:   mean=0.106, range=[0.006, 3.58]
+#   n_points:  mean=145, range=[80, 1000]
+
+# Data Quality SNR thresholds (dB)
+SNR_PERFECT_MIN = 15.0       # Loosened from 20 (real data is noisier)
+SNR_ACCEPTABLE_MIN = 5.0     # Loosened from 8
+SNR_BORDERLINE_MIN = 3.0
+
+# Fit Quality error thresholds (fraction)
+FIT_PERFECT_MAX_ERR = 0.03
+FIT_ACCEPTABLE_MAX_ERR = 0.20
+FIT_BAD_MIN_ERR = 0.20
+FIT_BORDERLINE_ERR_RANGE = (0.15, 0.30)
+
+# Rabi oscillation parameter ranges — CALIBRATED TO REAL DATA
+AMPLITUDE_RANGE = (0.0001, 0.42)   # Real: [~0, 0.40]
+PERIOD_RANGE = (0.025, 0.80)       # Real: [0.025, 0.78] for good fits
+PHASE_RANGE = (-3.1416, 3.1416)    # Real: [-π, π]
+OFFSET_RANGE = (0.0001, 0.98)      # Real: [~0, 0.98]
+DECAY_RANGE = (0.0, 30.0)          # Reduced max decay
+X_MAX_RANGE = (0.005, 0.50)        # Real: [0.006, 3.58] but most < 0.2
+
+# Extended ranges for diversity
+X_MAX_EXTENDED = (0.50, 3.60)      # For occasional long sweeps
+PERIOD_BAD_RANGE = (0.002, 19.0)   # Real bad fits have T up to 19
+
+# Per-sample x-axis variation
+N_POINTS_RANGE = (80, 500)         # Real: [80, 1000]
+
+# Proportion of borderline cases in Class 0
+BORDERLINE_FRACTION = 0.5          # Increased from 0.4
+
+# Proportion of near-zero amplitude samples (mimics ~25% of real data)
+NEAR_ZERO_AMP_FRACTION = 0.20

preprocessor.py

@@ -0,0 +1,148 @@
+"""
+Preprocessor for Rabi oscillation JSON experiment files.
+V3: Outputs 3 channels — signal + fit + residual (signal − fit).
+"""
+import json
+import numpy as np
+import torch
+from config import SEQ_LEN
+
+
+def load_json(path: str) -> dict:
+    """Load a JSON experiment file."""
+    with open(path, 'r') as f:
+        return json.load(f)
+
+
+def extract_fit_params(data: dict) -> dict:
+    """Extract the 4 fit parameters from JSON data.
+    
+    Returns dict with keys: amplitude, T, phase, offset.
+    Returns zeros if fitted_data is missing or None.
+    """
+    defaults = {'amplitude': 0.0, 'T': 0.05, 'phase': 0.0, 'offset': 0.0}
+    
+    fitted = data.get('fitted_data')
+    if fitted is None:
+        return defaults
+    
+    params_list = fitted.get('parameters')
+    if params_list is None:
+        return defaults
+    
+    params = dict(defaults)
+    for p in params_list:
+        if p.get('name') in params:
+            params[p['name']] = float(p.get('value', 0.0))
+    return params
+
+
+def rabi_oscillation(x, amplitude, T, phase, offset):
+    """Compute Rabi oscillation curve (without decay, matching JSON fit)."""
+    return amplitude * np.cos((2 * np.pi / T) * x + phase) + offset
+
+
+def reconstruct_fit_curve(x: np.ndarray, params: dict) -> np.ndarray:
+    """Reconstruct the fit curve on the original X-axis using JSON parameters."""
+    return rabi_oscillation(
+        x,
+        amplitude=params['amplitude'],
+        T=params['T'],
+        phase=params['phase'],
+        offset=params['offset']
+    )
+
+
+def preprocess_sample(data: dict):
+    """
+    Preprocess a single JSON experiment sample.
+    
+    V3: Returns 3-channel tensor (signal, fit, residual).
+    
+    Returns:
+        signal_tensor: (3, SEQ_LEN) — [signal, fit, residual]
+        params_tensor: (4,) fit parameters [amplitude, T, phase, offset]
+        question_id: str ('q1' or 'q2')
+        score: float (0.0 or 1.0)
+    """
+    # Extract raw data
+    x_raw = np.array(data['measured_data']['x_values'], dtype=np.float64)
+    y_raw = np.array(data['measured_data']['y_values'], dtype=np.float64)
+    
+    # Extract fit parameters
+    params = extract_fit_params(data)
+    
+    # Reconstruct fit curve on original x-axis
+    y_fit = reconstruct_fit_curve(x_raw, params)
+    
+    # Create uniform interpolation grid
+    x_uniform = np.linspace(x_raw.min(), x_raw.max(), SEQ_LEN)
+    
+    # Interpolate both curves to fixed length
+    y_raw_interp = np.interp(x_uniform, x_raw, y_raw)
+    y_fit_interp = np.interp(x_uniform, x_raw, y_fit)
+    
+    # Normalize both with same min-max scaling (from raw signal)
+    y_min = y_raw_interp.min()
+    y_max = y_raw_interp.max()
+    y_range = y_max - y_min
+    if y_range < 1e-10:
+        y_range = 1.0  # Avoid division by zero for flat signals
+    
+    y_raw_norm = (y_raw_interp - y_min) / y_range
+    y_fit_norm = (y_fit_interp - y_min) / y_range
+    
+    # Residual channel: signal − fit, normalized to [-1, 1]
+    residual = y_raw_norm - y_fit_norm
+    res_absmax = np.abs(residual).max()
+    if res_absmax < 1e-10:
+        res_absmax = 1.0
+    res_norm = residual / res_absmax
+    
+    # Build 3-channel tensor
+    signal_tensor = torch.tensor(
+        np.stack([y_raw_norm, y_fit_norm, res_norm], axis=0),
+        dtype=torch.float32
+    )
+    params_tensor = torch.tensor(
+        [params['amplitude'], params['T'], params['phase'], params['offset']],
+        dtype=torch.float32
+    )
+    
+    # Extract labels
+    question_id = data.get('question_id', 'q1')
+    score = float(data.get('score', 0.0))
+    
+    return signal_tensor, params_tensor, question_id, score
+
+
+def load_test_dataset(data_dir: str) -> list:
+    """
+    Load and preprocess all JSON files from the test data directory.
+    
+    Returns list of dicts with keys:
+        signal, params, question_id, score, filename
+    """
+    import os
+    samples = []
+    json_files = sorted([f for f in os.listdir(data_dir) if f.endswith('.json')])
+    
+    for fname in json_files:
+        path = os.path.join(data_dir, fname)
+        try:
+            data = load_json(path)
+            signal_3ch, params, qid, score = preprocess_sample(data)
+            samples.append({
+                'signal': signal_3ch,       # (3, SEQ_LEN)
+                'params': params,
+                'question_id': qid,
+                'score': score,
+                'filename': fname,
+                'raw_data': data,
+            })
+        except Exception as e:
+            print(f"Warning: Failed to preprocess {fname}: {e}")
+            continue
+    
+    print(f"Loaded {len(samples)} samples from {data_dir}")
+    return samples