Upload 4 files

ai_model.py

@@ -0,0 +1,166 @@
+"""
+AI Model for Rabi Oscillation Parameter Estimation.
+Transformer-based model that predicts fit parameters (A, T, phi, C) from raw signal data.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import math
+
+DEVICE = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
+FIXED_LEN = 256
+A_MIN, A_MAX = 1e-5, 0.5
+T_MIN, T_MAX = 0.02, 0.3
+C_MIN, C_MAX = -0.1, 1.0
+LOG_A_MIN, LOG_A_MAX = math.log(A_MIN), math.log(A_MAX)
+LOG_T_MIN, LOG_T_MAX = math.log(T_MIN), math.log(T_MAX)
+AUX_SCALES = np.array([0.1, 0.1, 0.1, 0.05, 0.1, 2.0], dtype=np.float32)
+
+
+def preprocess_sample(x_raw, y_raw):
+    xu = np.linspace(x_raw[0], x_raw[-1], FIXED_LEN)
+    yr = np.interp(xu, x_raw, y_raw)
+    ymin, ymax = float(yr.min()), float(yr.max())
+    ymean, ystd = float(yr.mean()), float(yr.std())
+    xrange = float(x_raw[-1] - x_raw[0])
+    ydm = yr - ymean
+    zc = np.where(np.diff(np.sign(ydm)))[0]
+    nper = len(zc) / 2.0 if len(zc) >= 1 else 0.5
+    span = ymax - ymin
+    if span < 1e-15:
+        span = 1.0
+    yn = ((yr - ymin) / span).astype(np.float32)
+    fft = np.fft.rfft(yn)
+    fft_mag = np.abs(fft[:64])
+    fft_ang = np.angle(fft[:64])
+    if len(fft_mag) < 64:
+        fft_mag = np.pad(fft_mag, (0, 64 - len(fft_mag)))
+        fft_ang = np.pad(fft_ang, (0, 64 - len(fft_ang)))
+    fft_feats = np.concatenate([fft_mag, fft_ang]).astype(np.float32)
+    aux = np.array([ymin, ymax, ymean, ystd, xrange, nper], dtype=np.float32) / AUX_SCALES
+    return yn, fft_feats, aux
+
+
+def target_to_params(t):
+    nA, nT, sp, cp, nC = t[..., 0], t[..., 1], t[..., 2], t[..., 3], t[..., 4]
+    A = np.exp(nA * (LOG_A_MAX - LOG_A_MIN) + LOG_A_MIN)
+    T = np.exp(nT * (LOG_T_MAX - LOG_T_MIN) + LOG_T_MIN)
+    phi = np.arctan2(sp, cp)
+    C = nC * (C_MAX - C_MIN) + C_MIN
+    if t.ndim == 1:
+        return np.array([A, T, phi, C], dtype=np.float64)
+    return np.stack([A, T, phi, C], -1).astype(np.float64)
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=512):
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        return x + self.pe[:x.size(1), :]
+
+
+class RabiEstimator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Sequential(nn.Conv1d(1, 64, kernel_size=7, stride=2, padding=3), nn.BatchNorm1d(64), nn.GELU())
+        self.conv2 = nn.Sequential(nn.Conv1d(64, 128, kernel_size=5, stride=2, padding=2), nn.BatchNorm1d(128), nn.GELU())
+        self.conv3 = nn.Sequential(nn.Conv1d(128, 256, kernel_size=5, stride=2, padding=2), nn.BatchNorm1d(256), nn.GELU())
+        self.d_model = 256
+        self.pos_enc = PositionalEncoding(self.d_model)
+        encoder_layer = nn.TransformerEncoderLayer(d_model=self.d_model, nhead=8, dim_feedforward=512, dropout=0.1, batch_first=True)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=6)
+        self.fft_net = nn.Sequential(nn.Linear(128, 256), nn.GELU(), nn.Dropout(0.1), nn.Linear(256, 128), nn.GELU())
+        self.aux_net = nn.Sequential(nn.Linear(6, 64), nn.GELU(), nn.Linear(64, 128), nn.GELU())
+        self.fusion = nn.Sequential(nn.Linear(256 + 128 + 128, 512), nn.BatchNorm1d(512), nn.GELU(), nn.Dropout(0.1))
+        self.head_atc = nn.Sequential(nn.Linear(512, 256), nn.GELU(), nn.Linear(256, 3))
+        self.head_ph = nn.Sequential(nn.Linear(512, 256), nn.GELU(), nn.Linear(256, 2))
+
+    def forward(self, wave, fft_f, aux):
+        x = wave.unsqueeze(1)
+        x = self.conv3(self.conv2(self.conv1(x)))
+        x = x.permute(0, 2, 1)
+        x = self.pos_enc(x)
+        x = self.transformer(x)
+        x_pool = x.mean(dim=1)
+        f = self.fft_net(fft_f)
+        a = self.aux_net(aux)
+        feat = self.fusion(torch.cat([x_pool, f, a], dim=1))
+        atc = torch.sigmoid(self.head_atc(feat))
+        ph = torch.tanh(self.head_ph(feat))
+        return torch.cat([atc[:, :2], ph, atc[:, 2:3]], 1)
+
+
+# ═══════════════════════════════════════════════════════════════════
+# Helper functions
+# ═══════════════════════════════════════════════════════════════════
+
+def rabi_formula(x, A, T, phi, C):
+    """Compute Rabi oscillation: A * cos(2*pi/T * x + phi) + C"""
+    return A * np.cos((2 * np.pi / T) * x + phi) + C
+
+
+def load_model(model_path):
+    """Load trained RabiEstimator weights from disk."""
+    model = RabiEstimator().to(DEVICE)
+    checkpoint = torch.load(model_path, map_location=DEVICE, weights_only=False)
+    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+        model.load_state_dict(checkpoint['model_state_dict'])
+    else:
+        model.load_state_dict(checkpoint)
+    model.eval()
+    return model
+
+
+def predict_ai_fit(x_raw, y_raw, model):
+    """
+    Predict Rabi parameters using the AI model and return the absolute fit curve.
+
+    Args:
+        x_raw: np.array of x values (time/amplitude axis)
+        y_raw: np.array of y values (signal)
+        model: loaded RabiEstimator model
+
+    Returns:
+        y_fit: np.array of predicted fit values evaluated on x_raw
+        params: dict with keys 'amplitude', 'T', 'phase', 'offset'
+    """
+    x_raw = np.asarray(x_raw, dtype=np.float64)
+    y_raw = np.asarray(y_raw, dtype=np.float64)
+
+    # Ensure eval mode (BatchNorm1d requires it for batch_size=1)
+    model.eval()
+
+    # Sort by x for interpolation
+    sort_idx = np.argsort(x_raw)
+    x_sorted = x_raw[sort_idx]
+    y_sorted = y_raw[sort_idx]
+
+    yn, fft_feats, aux = preprocess_sample(x_sorted, y_sorted)
+
+    wave_t = torch.tensor(yn, dtype=torch.float32).unsqueeze(0).to(DEVICE)
+    fft_t = torch.tensor(fft_feats, dtype=torch.float32).unsqueeze(0).to(DEVICE)
+    aux_t = torch.tensor(aux, dtype=torch.float32).unsqueeze(0).to(DEVICE)
+
+    with torch.no_grad():
+        pred = model(wave_t, fft_t, aux_t)
+
+    pred_np = pred.cpu().numpy()[0]
+    A, T, phi, C = target_to_params(pred_np)
+
+    # Evaluate fit on the original x values (unsorted order preserved)
+    y_fit = rabi_formula(x_raw, float(A), float(T), float(phi), float(C))
+
+    return y_fit, {
+        'amplitude': float(A),
+        'T': float(T),
+        'phase': float(phi),
+        'offset': float(C),
+    }

app.py

@@ -0,0 +1,333 @@
+"""
+Human-in-the-Loop Feedback Tool for Rabi Oscillation Classification.
+Production-ready Gradio interface with batch processing, AI model integration,
+dynamic graphing based on curve overlap, and MongoDB feedback storage.
+"""
+import gradio as gr
+import torch
+import numpy as np
+import plotly.graph_objects as go
+import json
+import os
+from datetime import datetime, timezone
+from pymongo import MongoClient
+import certifi
+
+from ai_model import (
+    RabiEstimator, predict_ai_fit, load_model, rabi_formula, DEVICE
+)
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Configuration
+# ═══════════════════════════════════════════════════════════════════════════
+MONGO_URI = "mongodb+srv://Admin:zeyl3RK8oEu6Qhz3@cluster0.iwlxabj.mongodb.net/?appName=Cluster0"
+DB_NAME = "rabi_classifier"
+COLLECTION_NAME = "physicist_feedback"
+MODEL_WEIGHTS = "rabi_model_best.pt"
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Global state
+# ═══════════════════════════════════════════════════════════════════════════
+AI_MODEL = None
+MONGO_COLLECTION = None
+
+def init_globals():
+    """Initialize AI model and MongoDB connection at startup."""
+    global AI_MODEL, MONGO_COLLECTION
+
+    # Load AI model
+    if os.path.exists(MODEL_WEIGHTS):
+        try:
+            AI_MODEL = load_model(MODEL_WEIGHTS)
+            print(f"✓ AI model loaded from {MODEL_WEIGHTS} (device: {DEVICE})")
+        except Exception as e:
+            print(f"⚠ Failed to load AI model: {e}")
+    else:
+        print(f"⚠ Model weights not found at {MODEL_WEIGHTS}")
+
+    # Connect to MongoDB (non-blocking — app works without it)
+    try:
+        client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=3000,
+                            tlsCAFile=certifi.where())
+        client.admin.command('ping')
+        db = client[DB_NAME]
+        MONGO_COLLECTION = db[COLLECTION_NAME]
+        print(f"✓ MongoDB connected: {DB_NAME}.{COLLECTION_NAME}")
+    except Exception as e:
+        MONGO_COLLECTION = None
+        print(f"⚠ MongoDB not available (will retry on submit): {e}")
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# JSON helpers
+# ═══════════════════════════════════════════════════════════════════════════
+
+def extract_fit_params(data):
+    """Extract the 4 fit parameters from JSON fitted_data."""
+    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
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Plotting
+# ═══════════════════════════════════════════════════════════════════════════
+
+def create_plot(data, filename, regular_params, ai_params, curves_overlap):
+    """Create Plotly plot with raw data, Regular Fit (blue), AI Fit (green)."""
+    x = np.array(data['measured_data']['x_values'])
+    y = np.array(data['measured_data']['y_values'])
+    x_dense = np.linspace(x.min(), x.max(), 500)
+
+    y_regular = rabi_formula(
+        x_dense, regular_params['amplitude'], regular_params['T'],
+        regular_params['phase'], regular_params['offset']
+    )
+
+    fig = go.Figure()
+
+    # Raw data points
+    fig.add_trace(go.Scatter(
+        x=x, y=y, mode='markers', name='Raw Data',
+        marker=dict(color='rgba(52, 152, 219, 0.7)', size=6),
+    ))
+
+    # Regular Fit — Blue
+    fig.add_trace(go.Scatter(
+        x=x_dense, y=y_regular, mode='lines', name='Regular Fit (JSON)',
+        line=dict(color='#2980b9', width=2.5),
+    ))
+
+    # AI Fit — Green
+    if ai_params is not None:
+        y_ai = rabi_formula(
+            x_dense, ai_params['amplitude'], ai_params['T'],
+            ai_params['phase'], ai_params['offset']
+        )
+        if curves_overlap:
+            fig.add_trace(go.Scatter(
+                x=x_dense, y=y_ai, mode='lines',
+                name='AI Fit (overlapping)',
+                line=dict(color='#27ae60', width=3.5, dash='dot'),
+            ))
+        else:
+            fig.add_trace(go.Scatter(
+                x=x_dense, y=y_ai, mode='lines', name='AI Fit',
+                line=dict(color='#27ae60', width=2.5),
+            ))
+
+    title = f"Rabi Oscillation — {filename}"
+    if curves_overlap:
+        title += " (OVERLAP)"
+
+    fig.update_layout(
+        title=title,
+        xaxis_title='Drive Amplitude (a.u.)',
+        yaxis_title='Signal (a.u.)',
+        height=550,
+        margin=dict(l=50, r=30, t=60, b=50),
+        legend=dict(x=0.01, y=0.99),
+        font=dict(size=14),
+    )
+    return fig
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Core processing
+# ═══════════════════════════════════════════════════════════════════════════
+
+def process_experiment(data, filename):
+    x = np.array(data['measured_data']['x_values'])
+    y = np.array(data['measured_data']['y_values'])
+
+    regular_params = extract_fit_params(data)
+    ai_params = None
+    overlap = True
+    model_status_msg = "Awaiting model..."
+
+    if AI_MODEL is not None:
+        try:
+            _, ai_params = predict_ai_fit(x, y, AI_MODEL)
+            y_reg = rabi_formula(x, regular_params['amplitude'], regular_params['T'],
+                                 regular_params['phase'], regular_params['offset'])
+            y_ai = rabi_formula(x, ai_params['amplitude'], ai_params['T'],
+                                ai_params['phase'], ai_params['offset'])
+            overlap = np.allclose(y_reg, y_ai, rtol=1e-3, atol=1e-3)
+            model_status_msg = "✅ AI Prediction successful."
+        except Exception as e:
+            model_status_msg = f"❌ AI prediction failed: {e}"
+            ai_params = None
+            overlap = True
+    else:
+        model_status_msg = "⚠ No AI model loaded. Only regular fit shown."
+
+    fig = create_plot(data, filename, regular_params, ai_params, overlap)
+    return fig, overlap, regular_params, ai_params, model_status_msg
+
+def display_experiment(batch_data, idx):
+    if not batch_data or idx < 0 or idx >= len(batch_data):
+        return (
+            None, "📂 Upload JSON files to begin", "", "Waiting for data...",
+            gr.update(visible=True), gr.update(visible=False),
+            None, "", None, None, None, "", "", "",
+            batch_data, idx, False
+        )
+
+    exp = batch_data[idx]
+    fig, overlap, reg_p, ai_p, mod_stat = process_experiment(exp['data'], exp['filename'])
+    progress = f"Experiment {idx + 1} / {len(batch_data)} — {exp['filename']}"
+
+    if overlap:
+        overlap_info = "🔗 Curves overlap (Regular and AI solutions match). Rate once."
+    else:
+        overlap_info = "↔️ Curves diverge (Regular = Blue, AI = Green). Rate each separately."
+
+    return (
+        fig, progress, overlap_info, mod_stat,
+        gr.update(visible=overlap),
+        gr.update(visible=not overlap),
+        None, "", None, None, None, "", "", "",
+        batch_data, idx, overlap
+    )
+
+def on_upload(files):
+    if not files: return display_experiment([], 0)
+    batch_data = []
+    for f in files:
+        try:
+            fpath = f.name if hasattr(f, 'name') else str(f)
+            with open(fpath, 'r') as fp: data = json.load(fp)
+            batch_data.append({'data': data, 'filename': os.path.basename(fpath)})
+        except: pass
+    return display_experiment(batch_data, 0)
+
+def on_prev(batch, idx): return display_experiment(batch, max(0, idx - 1))
+def on_next(batch, idx): return display_experiment(batch, min(len(batch) - 1, idx + 1) if batch else 0)
+
+def on_submit(batch, idx, overlap, data_rate, data_comment, fit_sing, fit_reg, fit_ai, c_sing, c_reg, c_ai):
+    if not batch: return ("⚠ No experiment.",) + display_experiment([], 0)
+    if data_rate is None: return ("⚠ Rate Data Quality.",) + display_experiment(batch, idx)
+    if overlap and fit_sing is None: return ("⚠ Rate Fit Quality.",) + display_experiment(batch, idx)
+    if not overlap and (fit_reg is None or fit_ai is None): return ("⚠ Rate both fits.",) + display_experiment(batch, idx)
+
+    exp = batch[idx]
+    x = np.array(exp['data']['measured_data']['x_values'])
+    y = np.array(exp['data']['measured_data']['y_values'])
+    ai_params = None
+    if AI_MODEL:
+        try: _, ai_params = predict_ai_fit(x, y, AI_MODEL)
+        except: pass
+
+    doc = {
+        'filename': exp['filename'],
+        'raw_data': exp['data'],
+        'ai_prediction': ai_params,
+        'curves_overlap': bool(overlap),
+        'data_quality_rating': data_rate,
+        'data_comment': data_comment or '',
+        'timestamp': datetime.now(timezone.utc).isoformat(),
+    }
+    if overlap:
+        doc['fit_quality_rating'] = {'combined': fit_sing}
+        doc['fit_comments'] = {'combined': c_sing or ''}
+    else:
+        doc['fit_quality_rating'] = {'regular_fit': fit_reg, 'ai_fit': fit_ai}
+        doc['fit_comments'] = {'regular_fit': c_reg or '', 'ai_fit': c_ai or ''}
+
+    status = "Feedback saved!"
+    if MONGO_COLLECTION is not None:
+        try: MONGO_COLLECTION.insert_one(doc)
+        except Exception as e: status = f"MongoDB error: {e}"
+    else: status = "⚠ MongoDB disconnected. Not saved."
+
+    nxt = idx + 1 if idx < len(batch) - 1 else idx
+    status += " Loaded next." if idx < len(batch) - 1 else " Last experiment!"
+    return (status,) + display_experiment(batch, nxt)
+
+
+RATING_CHOICES = ["Bad (0)", "Borderline (1)", "Perfect (2)"]
+
+def build_ui():
+    init_globals()
+    with gr.Blocks(title="Rabi Feedback") as demo:
+        batch_data = gr.State([])
+        batch_index = gr.State(0)
+        overlap_state = gr.State(False)
+
+        gr.Markdown("# ⚛️ Rabi Oscillation — Physicist Feedback Tool")
+        gr.Markdown("Upload JSON experiments, review the quality of the data and fits, and submit your ratings.")
+        
+        with gr.Row():
+            file_upload = gr.File(label="Upload JSON Files", file_count="multiple", file_types=[".json"])
+        
+        with gr.Row():
+            prev_btn = gr.Button("← Previous")
+            progress_md = gr.Markdown("📂 Upload JSON files to begin")
+            next_btn = gr.Button("Next →")
+
+        plot = gr.Plot()
+        
+        with gr.Row():
+            overlap_info = gr.Markdown("")
+            model_status = gr.Markdown("Awaiting model...")
+
+        gr.Markdown("---")
+        
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### 1. Data Quality")
+                data_rating = gr.Radio(choices=RATING_CHOICES, label="Rate raw data visually:")
+                comment_data = gr.Textbox(label="Data notes", lines=2)
+            
+            with gr.Column(visible=True) as fit_single_col:
+                gr.Markdown("### 2. Fit Quality (Combined)")
+                fit_single_rating = gr.Radio(choices=RATING_CHOICES, label="Rate the unified fit:")
+                comment_single = gr.Textbox(label="Fit notes", lines=2)
+                
+        with gr.Column(visible=False) as fit_dual_col:
+            gr.Markdown("### 2. Fit Quality (Diverging)")
+            with gr.Row():
+                with gr.Column():
+                    fit_regular_rating = gr.Radio(choices=RATING_CHOICES, label="Regular Fit (Blue):")
+                    comment_regular = gr.Textbox(label="Notes on blue curve", lines=2)
+                with gr.Column():
+                    fit_ai_rating = gr.Radio(choices=RATING_CHOICES, label="AI Fit (Green):")
+                    comment_ai = gr.Textbox(label="Notes on green curve", lines=2)
+
+        gr.Markdown("---")
+        submit_btn = gr.Button("Submit Feedback & Next", variant="primary")
+        status_md = gr.Markdown("")
+
+        outputs = [
+            plot, progress_md, overlap_info, model_status,
+            fit_single_col, fit_dual_col,
+            data_rating, comment_data,
+            fit_single_rating, fit_regular_rating, fit_ai_rating,
+            comment_single, comment_regular, comment_ai,
+            batch_data, batch_index, overlap_state
+        ]
+
+        file_upload.change(on_upload, [file_upload], outputs)
+        prev_btn.click(on_prev, [batch_data, batch_index], outputs)
+        next_btn.click(on_next, [batch_data, batch_index], outputs)
+        submit_btn.click(on_submit, [
+            batch_data, batch_index, overlap_state,
+            data_rating, comment_data,
+            fit_single_rating, fit_regular_rating, fit_ai_rating,
+            comment_single, comment_regular, comment_ai
+        ], [status_md] + outputs)
+
+    return demo
+
+if __name__ == '__main__':
+    demo = build_ui()
+    demo.launch()

rabi_model_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c9062ef5e9c2a954869e825ce7dc8c33975a2644320639b186b0fd95703f28f
+size 16459501

requirements.txt

@@ -0,0 +1,8 @@
+torch
+gradio
+pymongo
+dnspython
+scipy
+numpy
+plotly
+certifi