Quantum Machine Learning Hands-On: Extract Quantum Features with Rimay

# task1_payload.py: engineer the physics, build the 9-feature payload.
import json
import numpy as np
import pandas as pd
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

df = fetch_openml("ai4i2020", version=1, as_frame=True).frame
df["type_code"] = df["Type"].map({"L": 0, "M": 1, "H": 2}).astype(float)

df["temp_diff"] = ...     # TODO
df["power"] = ...         # TODO
df["overstrain"] = ...    # TODO

# Column ORDER is the feature-to-qubit map and decides the pair topology.
CTRL9 = [
    "Air temperature [K]",
    "Process temperature [K]",
    "Rotational speed [rpm]",
    "Torque [Nm]",
    "Tool wear [min]",
    "temp_diff",
    "power",
    "overstrain",
    "type_code",
]

X = df[CTRL9].astype(float)
y = df["Machine failure"].astype(int)

# Keep every failure, subsample non-failures up to the 3000-row cap.
rng = np.random.RandomState(42)
pos = y[y == 1].index.to_numpy()                       # 339 rows
neg = rng.choice(y[y == 0].index.to_numpy(), 3000 - len(pos), replace=False)
keep = np.concatenate([pos, neg])
rng.shuffle(keep)

Xk, yk = X.loc[keep], y.loc[keep]
X_train, X_test, y_train, y_test = train_test_split(
    Xk, yk, test_size=0.2, random_state=42, stratify=yk
)

# Fit the scaler on the training rows only. Rimay applies its own minmax scaler
# on top, which is monotone per column, so this does not change the encoding.
sc = StandardScaler().fit(X_train)
X_train = pd.DataFrame(sc.transform(X_train), columns=CTRL9)
X_test = pd.DataFrame(sc.transform(X_test), columns=CTRL9)
y_train = pd.Series(y_train.to_numpy(), name="label")
y_test = pd.Series(y_test.to_numpy(), name="label")

payload = {
    "X_train": X_train.to_dict(),
    "y_train": y_train.to_frame("label").to_dict(),
    "X_test": X_test.to_dict(),
    "y_test": y_test.to_frame("label").to_dict(),
}

# Row keys must agree between the feature block and the label block, per split.
for xk_, yk_ in (("X_train", "y_train"), ("X_test", "y_test")):
    xkeys = list(next(iter(payload[xk_].values())).keys())
    assert list(payload[yk_]["label"].keys()) == xkeys, xk_
    for col in payload[xk_].values():
        assert list(col.keys()) == xkeys, xk_

with open("data.json", "w") as fh:
    json.dump(payload, fh)

print(len(X_train), len(X_test), X_train.shape[1], int(y_train.sum() + y_test.sum()))

# task2_pools.py
import io, json, os
from dotenv import load_dotenv
from qhub.api.platform import HubPlatformClient

load_dotenv()
platform = HubPlatformClient(api_key=os.getenv("KQH_PERSONAL_ACCESS_TOKEN"))

input_dp = platform.data_pools.create_data_pool(name="Rimay Lab Input")
output_dp = platform.data_pools.create_data_pool(name="Rimay Lab Output")

# A reused input pool may already hold a data.json. Replace it rather than adding a second.
for f in platform.data_pools.get_data_pool_files(id=input_dp.id):
    if f.name == "data.json":
        platform.data_pools.delete_data_pool_file(id=input_dp.id, file_id=f.id)

body = open("data.json", "rb").read()
platform.data_pools.add_data_pool_file(
    id=input_dp.id, file=("data.json", io.BytesIO(body))
)

print("input ", input_dp.id)
print("output", output_dp.id)
print(f"uploaded data.json ({len(body) / 1024:.0f} KB)")

# task3_submit.py
import os
import uuid
from dotenv import load_dotenv
from qhub.service.client import HubServiceClient

load_dotenv()

INPUT_POOL_ID = "paste-the-input-pool-id-task2-printed"
OUTPUT_POOL_ID = "paste-the-output-pool-id-task2-printed"

for _name, _val in (("INPUT_POOL_ID", INPUT_POOL_ID), ("OUTPUT_POOL_ID", OUTPUT_POOL_ID)):
    try:
        uuid.UUID(_val)
    except ValueError:
        raise SystemExit(f"{_name} is still the placeholder: Replace with the Pool-ID from Task 2")

service = HubServiceClient(
    service_endpoint=(
        "https://gateway.hub.kipu-quantum.com/kipu-quantum/"
        "rimay---quantum-feature-extraction---simulator/1.0.0"
    ),
    access_key_id=os.getenv("KQH_ACCESS_KEY_ID"),
    secret_access_key=os.getenv("KQH_SECRET_ACCESS_KEY"),
)

execution = service.run(request={
    "input_data_pool": {"id": INPUT_POOL_ID, "ref": "DATAPOOL"},
    "output_data_pool": {"id": OUTPUT_POOL_ID, "ref": "DATAPOOL"},
    "_mode": "fit_transform",   # fit on train, transform train and test in one run
    "_fit_reference": "None",   # the LITERAL STRING, not Python None; see the caution
    "num_shots": 2000,          # not 500: see the shot-count note above
    "num_runs": 1,
})

print("execution id:", execution.id)
execution.wait_for_final_state(timeout=900, wait=10)
print("status:", execution.status)

if execution.status == "FAILED":
    for log in execution.logs():
        print(log)

# task4_fetch.py
import os
import uuid
import numpy as np
from pathlib import Path
from dotenv import load_dotenv
from qhub.api.platform import HubPlatformClient

load_dotenv()
platform = HubPlatformClient(api_key=os.getenv("KQH_PERSONAL_ACCESS_TOKEN"))

OUTPUT_POOL_ID = "paste-the-output-pool-id-task2-printed"

try:
    uuid.UUID(OUTPUT_POOL_ID)
except ValueError:
    raise SystemExit("OUTPUT_POOL_ID is still the placeholder: paste the id Task 2 printed")

out_dir = Path("rimay_output")
out_dir.mkdir(exist_ok=True)

def download_verified(dp_id, f, dest, attempts=3):
    """Stream a pool file and verify the full content_length actually arrived."""
    for attempt in range(1, attempts + 1):
        try:
            stream = platform.data_pools.get_data_pool_file(id=dp_id, file_id=f.id)
            n = 0
            with open(dest, "wb") as fp:
                for chunk in stream:
                    n += len(chunk)
                    fp.write(chunk)
            if f.content_length and n != f.content_length:
                raise IOError(f"received {n}/{f.content_length} bytes")
            return n
        except Exception as e:
            if attempt == attempts:
                raise RuntimeError(f"failed to download {f.name}: {e}") from e

for f in platform.data_pools.get_data_pool_files(id=OUTPUT_POOL_ID):
    if f.name.endswith(".npy"):
        n = download_verified(OUTPUT_POOL_ID, f, out_dir / f.name)
        print(f"  {f.name} ({n} bytes)")

Xc_train = np.load(out_dir / "Xc_train.npy")
Xq_train = np.load(out_dir / "Xq_train_0.npy")
n_c, n_q = Xc_train.shape[1], Xq_train.shape[1]
print(f"classical {n_c}, quantum {n_q} = {n_c} per-feature + {n_q - n_c} pairs")
assert n_q == 2 * n_c - 1
assert n_c == 9, "expected the 9-feature extraction; re-run Tasks 1 to 3"

classical 9, quantum 17 = 9 per-feature + 8 pairs

# lab_common.py, started in Task 5
import numpy as np
from pathlib import Path
from sklearn.base import BaseEstimator, TransformerMixin

import numpy as np
from pathlib import Path
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import GridSearchCV, StratifiedKFold
from sklearn.naive_bayes import GaussianNB

# The raw sensor columns inside the 9-column table: five sensors plus type code.
RAW = [0, 1, 2, 3, 4, 8]

# Kipu brand colors (tokens.css)
GOLD, MUTED, FG, BG = "#b38f12", "#6b6b6b", "#0f1319", "#ffffff"

def load_arrays(out_dir="rimay_output"):
    """The six row-aligned arrays Task 4 downloaded, in one call."""
    d = Path(out_dir)
    return (np.load(d / "Xc_train.npy"), np.load(d / "Xq_train_0.npy"),
            np.load(d / "y_train.npy").ravel().astype(int),
            np.load(d / "Xc_test.npy"), np.load(d / "Xq_test_0.npy"),
            np.load(d / "y_test.npy").ravel().astype(int))

def fisher_scores(X, y):
    """Per-column Fisher ratio: squared class-mean gap over summed class variance."""
    a, b = X[y == 1], X[y == 0]
    return (a.mean(0) - b.mean(0))**2 / (a.var(0) + b.var(0) + 1e-12)

class FisherGate(BaseEstimator, TransformerMixin):
    """Adaptive cutoff instead of a fixed top-k. Keep all n_raw classical
    columns; keep a quantum column only if its Fisher score clears the median
    of the classical columns' own scores. In words: a quantum column has to be
    at least as informative, on its own, as a typical column you already had.
    Computed inside fit, so it sees the training fold only."""
    def __init__(self, n_raw=9, q=0.5):
        self.n_raw, self.q = n_raw, q
    def fit(self, X, y):
        f = fisher_scores(X, y)
        thr = np.quantile(f[:self.n_raw], self.q)
        self.keep_ = np.r_[np.arange(self.n_raw),
                           self.n_raw + np.where(f[self.n_raw:] >= thr)[0]]
        return self
    def transform(self, X):
        return X[:, self.keep_]

# task5_compare.py: the progression, raw -> expert -> expert + Rimay.
import numpy as np
from scipy.stats import wilcoxon
from sklearn.base import clone
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import StratifiedKFold, GridSearchCV
from sklearn.naive_bayes import GaussianNB
from sklearn.model_selection import StratifiedShuffleSplit, RepeatedStratifiedKFold
from sklearn.metrics import average_precision_score
from sklearn.linear_model import LogisticRegression
from sklearn.svm import LinearSVC
from lab_common import load_arrays, FisherGate, RAW

# RBF SVM and gradient boosting sit out of this panel for training time; the
# study's Holm-corrected readings for both are quoted below the expected output.
MODELS = {
    "naive bayes":         (GaussianNB(), {"m__var_smoothing": [1e-9, 1e-7, 1e-5, 1e-3]}),
    "logistic regression": (LogisticRegression(max_iter=5000), {"m__C": [0.01, 0.1, 1, 10, 100]}),
    "linear SVM":          (LinearSVC(max_iter=20000, dual="auto"), {"m__C": [0.01, 0.1, 1, 10]}),
}

def holm(pvals):
    """Holm step-down correction over a dict of name -> raw p-value."""
    items = sorted(pvals.items(), key=lambda kv: kv[1])
    m, out, running = len(items), {}, 0.0
    for i, (name, p) in enumerate(items):
        running = max(running, (m - i) * p)
        out[name] = min(1.0, running)
    return out

Xc_train, Xq_train, y_train, Xc_validate, Xq_validate, y_validate = load_arrays()
Xc = np.vstack([Xc_train, Xc_validate])
Xq = np.vstack([Xq_train, Xq_validate])
Y = np.concatenate([y_train, y_validate])
assert Xc.shape[1] == 9, "rimay_output holds an older run; re-run Tasks 1 to 4"

# The progression: each arm is read against the one it has to beat.
# TODO: the gate's knob. A quantum column must clear this quantile of the
# classical columns' own Fisher scores. Where does a quantum column earn its
# place: above the weakest classical column, the typical one, the best one?
GATE_QUANTILE = ...

ARMS = {
    "raw":          (Xc[:, RAW], None),   # sensors only, no extraction needed
    "expert":       (..., None),          # TODO: which columns?
    "expert+all":   (..., None),          # TODO: which columns?
    "expert+gated": (..., FisherGate(q=GATE_QUANTILE)),  # TODO: which columns feed the gate?
}
# TODO: each arm is read against the arm it has to beat. Which one is that?
BASELINE = {"expert": ..., "expert+all": ..., "expert+gated": ...}

def fit_score(X, est, grid, sel, tr, te, seed):
    steps = [("sc", StandardScaler())]
    if sel is not None:
        steps.insert(0, ("sel", clone(sel)))
    steps.append(("m", clone(est)))
    gs = GridSearchCV(Pipeline(steps), grid, scoring="average_precision",
                      cv=StratifiedKFold(3, shuffle=True, random_state=seed), n_jobs=1)
    gs.fit(X[tr], Y[tr])
    b = gs.best_estimator_
    s = (b.decision_function(X[te]) if hasattr(b, "decision_function")
         else b.predict_proba(X[te])[:, 1])
    return average_precision_score(Y[te], s)

# Two evaluation protocols over the same arms. split: 15 random 80/20 draws.
# oof: 5 folds x 3 repeats, every row tested exactly once per repeat, the
# fold-level pairing the Step 7 curves argue for.
PROTOCOLS = {
    "split": list(StratifiedShuffleSplit(n_splits=15, test_size=0.2,
                                         random_state=7).split(Xc, Y)),
    "oof":   list(RepeatedStratifiedKFold(n_splits=5, n_repeats=3,
                                          random_state=7).split(Xc, Y)),
}

res = {}
for proto, splits in PROTOCOLS.items():
    out = {(m, a): [] for m in MODELS for a in ARMS}
    for seed, (tr, te) in enumerate(splits):
        for m, (est, grid) in MODELS.items():
            for a, (X, sel) in ARMS.items():
                out[(m, a)].append(fit_score(X, est, grid, sel, tr, te, seed))
    res[proto] = out

# Absolute AP per arm, delta against the arm's own baseline, and a Wilcoxon
# signed-rank p per step, Holm corrected across the three comparisons within
# each model, once per protocol.
for m in MODELS:
    ph = {}
    for proto in ("split", "oof"):
        ph[proto] = holm({a: wilcoxon(np.array(res[proto][(m, a)]),
                                      np.array(res[proto][(m, BASELINE[a])])).pvalue
                          for a in BASELINE})
    line = f"{m:20s} raw {np.mean(res['split'][(m, 'raw')]):.4f}"
    for a in ("expert", "expert+all", "expert+gated"):
        v = np.mean(res["split"][(m, a)])
        d = v - np.mean(res["split"][(m, BASELINE[a])])
        line += (f" | {a} {v:.4f} ({d:+.4f} vs {BASELINE[a]},"
                 f" p {ph['split'][a]:.4f}, oof p {ph['oof'][a]:.4f})")
    print(line)

# lab_common.py, appended in Task 6
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import GridSearchCV, StratifiedKFold
from sklearn.naive_bayes import GaussianNB

def nb_search(seed, gate=False, n_jobs=1):
    """The Task 5 naive-Bayes protocol as one reusable estimator: scaler plus
    GaussianNB, smoothing grid chosen by inner CV, optionally behind the
    FisherGate. Steps 6 to 8 reuse it."""
    steps = [("sc", StandardScaler()), ("m", GaussianNB())]
    if gate:
        steps.insert(0, ("sel", FisherGate()))
    grid = {"m__var_smoothing": [1e-9, 1e-7, 1e-5, 1e-3]}
    return GridSearchCV(Pipeline(steps), grid, scoring="average_precision",
                        cv=StratifiedKFold(3, shuffle=True, random_state=seed),
                        n_jobs=n_jobs)

# task6_prcurve.py: ROC and PR for the progression, one model, three arms.
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import (precision_recall_curve, roc_curve,
                             average_precision_score, roc_auc_score)
from lab_common import load_arrays, nb_search, RAW, GOLD, MUTED, FG, BG

Xc_train, Xq_train, y_train, Xc_validate, Xq_validate, y_validate = load_arrays()
assert Xc_train.shape[1] == 9, "rimay_output holds an older run; re-run Tasks 1 to 4"

def fit_score(Xtr, Xte, gate=False):
    gs = nb_search(seed=0, gate=gate)
    gs.fit(Xtr, y_train)
    return gs.best_estimator_.predict_proba(Xte)[:, 1]

# The extraction already handed you one fixed split: 2,400 train, 600 validation.
CURVES = [
    # TODO: three (name, color, scores) entries, the progression from Task 5:
    # raw 6 in MUTED, expert 9 in FG, expert + Rimay (gated) in GOLD
]

fig, (axr, axp) = plt.subplots(1, 2, figsize=(14, 6), dpi=100)
fig.patch.set_facecolor(BG)
prev = y_validate.mean()

for name, col, s in CURVES:
    fpr, tpr, _ = roc_curve(y_validate, s)
    axr.plot(fpr, tpr, color=col, lw=3, solid_capstyle="round",
             label=f"{name}  AUC {roc_auc_score(y_validate, s):.3f}")
    p, r, _ = precision_recall_curve(y_validate, s)
    axp.plot(r, p, color=col, lw=3, solid_capstyle="round",
             label=f"{name}  AP {average_precision_score(y_validate, s):.3f}")

axr.plot([0, 1], [0, 1], color=MUTED, lw=1, ls="--")           # chance: the diagonal
axp.axhline(prev, color=MUTED, lw=1, ls="--")                  # chance: the prevalence
axr.set_xlabel("FALSE POSITIVE RATE"); axr.set_ylabel("TRUE POSITIVE RATE")
axp.set_xlabel("RECALL"); axp.set_ylabel("PRECISION")
axr.set_title("ROC: the gaps compress", color=FG, fontweight="bold", loc="left")
axp.set_title("PR: the picture your stakeholder lives in", color=FG, fontweight="bold", loc="left")
for ax in (axr, axp):
    ax.set_facecolor(BG)
    for sp in ("top", "right"): ax.spines[sp].set_visible(False)
    for sp in ("left", "bottom"): ax.spines[sp].set_color(MUTED)
    ax.xaxis.label.set_fontfamily("monospace"); ax.yaxis.label.set_fontfamily("monospace")
    ax.xaxis.label.set_color(MUTED); ax.yaxis.label.set_color(MUTED)
    ax.tick_params(colors=MUTED)
    ax.set_xlim(0, 1); ax.set_ylim(0, 1.02)
    ax.legend(frameon=False, fontsize=9, loc="lower right")
fig.tight_layout()
fig.savefig("pr-vs-roc.png", facecolor=BG, bbox_inches="tight")
plt.show()

print(f"prevalence {prev:.3f}, positives in validation {int(y_validate.sum())}")
for name, _, s in CURVES:
    print(f"{name:24s} ROC AUC {roc_auc_score(y_validate, s):.4f}"
          f"  AP {average_precision_score(y_validate, s):.4f}")

# task7_oof_pr.py: the same comparison scored two ways.
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import RepeatedStratifiedKFold
from sklearn.metrics import precision_recall_curve, average_precision_score
from lab_common import load_arrays, nb_search, GOLD, MUTED, FG, BG

Xc_train, Xq_train, y_train, Xc_validate, Xq_validate, y_validate = load_arrays()

# Pool the split back together: out-of-fold scores every one of the 3,000 rows.
Xc = np.vstack([Xc_train, Xc_validate])
Xq = np.vstack([Xq_train, Xq_validate])
Y = np.concatenate([y_train, y_validate])
X_hyb = np.hstack([Xc, Xq])

N_SPLITS = ...    # TODO: how many folds, so that every fold still holds enough failures?
N_REPEATS = ...   # TODO: how many repeats, to smooth where the fold boundaries fell?

def oof_scores(X, gate=False):
    """Every row scored N_REPEATS times by a model that never saw it."""
    acc, cnt = np.zeros(len(Y)), np.zeros(len(Y))
    rkf = RepeatedStratifiedKFold(n_splits=N_SPLITS, n_repeats=N_REPEATS, random_state=7)
    for tr, te in rkf.split(X, Y):
        gs = nb_search(seed=7, gate=gate, n_jobs=6)
        gs.fit(X[tr], Y[tr])
        acc[te] += gs.best_estimator_.predict_proba(X[te])[:, 1]
        cnt[te] += 1
    assert cnt.min() == N_REPEATS
    return acc / cnt

def single_split_scores(X_tr, X_te, gate=False):
    """The one fixed split the extraction handed you: 2,400 train, 600 validation."""
    gs = nb_search(seed=0, gate=gate)
    gs.fit(X_tr, y_train)
    return gs.best_estimator_.predict_proba(X_te)[:, 1]

panels = [
    ("One fixed split, 600 rows", y_validate,
     [("expert 9", FG, single_split_scores(Xc_train, Xc_validate)),
      ("expert + Rimay (gated)", GOLD,
       single_split_scores(np.hstack([Xc_train, Xq_train]),
                           np.hstack([Xc_validate, Xq_validate]), gate=True))]),
    ("Out of fold, all 3,000 rows", Y,
     [("expert 9", FG, oof_scores(Xc)),
      ("expert + Rimay (gated)", GOLD, oof_scores(X_hyb, gate=True))]),
]

fig, axes = plt.subplots(1, 2, figsize=(14, 6), dpi=100)
fig.patch.set_facecolor(BG)
for ax, (title, y_true, curves) in zip(axes, panels):
    for name, col, s in curves:
        p, r, _ = precision_recall_curve(y_true, s)
        ax.plot(r, p, color=col, lw=3, solid_capstyle="round",
                label=f"{name}  AP {average_precision_score(y_true, s):.3f}")
    ax.axhline(y_true.mean(), color=MUTED, lw=1, ls="--")
    ax.set_facecolor(BG)
    for sp in ("top", "right"): ax.spines[sp].set_visible(False)
    for sp in ("left", "bottom"): ax.spines[sp].set_color(MUTED)
    ax.set_xlabel("RECALL"); ax.set_ylabel("PRECISION")
    ax.xaxis.label.set_fontfamily("monospace"); ax.yaxis.label.set_fontfamily("monospace")
    ax.xaxis.label.set_color(MUTED); ax.yaxis.label.set_color(MUTED)
    ax.tick_params(colors=MUTED)
    ax.set_xlim(0, 1); ax.set_ylim(0, 1.02)
    ax.set_title(title, color=FG, fontweight="bold", loc="left")
    ax.legend(frameon=False, fontsize=9, loc="lower left")
fig.tight_layout()
fig.savefig("pr-single-vs-oof.png", facecolor=BG, bbox_inches="tight")
plt.show()

for title, y_true, curves in panels:
    for name, _, s in curves:
        print(f"{title:32s} {name:24s} AP {average_precision_score(y_true, s):.4f}")