Spaces:
Running
Running
File size: 4,618 Bytes
eb52e7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | """
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
|