One table, many questions

import os
import time

import numpy as np
from sklearn.datasets import fetch_openml
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split

from causilo_client import Causilo

cx = Causilo(os.environ["CAUSILO_ENDPOINT"], token=os.environ["CAUSILO_KEY"])

_spent = [(cx.usage() or {}).get("used", 0)]


def counted(meta) -> str:
    """Cells this call added to the month, from the server's own running total.

    `metadata.quota.used` is what the service has counted for this key, so the
    difference between two calls is what the second one cost. Reading it this
    way means the example cannot disagree with the ledger.
    """
    used = (meta.get("quota") or {}).get("used")
    if used is None:
        return "not metered"
    delta, _spent[0] = used - _spent[0], used
    return f"{delta:,} cells"
print("ready:", cx.health())

X, y = fetch_openml(data_id=1590, as_frame=True, return_X_y=True)
X_train, X_rest, y_train, y_rest = train_test_split(
    X, y, train_size=12_000, test_size=6_000, stratify=y, random_state=42
)
history = X_train.assign(income=y_train.values)

cuts = [len(X_rest) * i // 3 for i in range(4)]
batches = [X_rest.iloc[a:b] for a, b in zip(cuts, cuts[1:])]
labels = [y_rest.values[a:b] for a, b in zip(cuts, cuts[1:])]
print("history", history.shape, "batches", [b.shape for b in batches])


POSITIVE = ">50K"


def auc(proba, meta, truth):
    proba = np.array(proba)
    good = proba[:, meta["classes"].index(POSITIVE)]
    return round(roc_auc_score(truth == POSITIVE, good), 4)

print("\nwithout the cache")
plain = []
for i, (batch, truth) in enumerate(zip(batches, labels)):
    t = time.monotonic()
    proba, meta = cx.predict_with_metadata(
        history, batch, target="income", model="causilo-clf", output_type="probas"
    )
    plain.append(auc(proba, meta, truth))
    # The cells the server charged, read off the response rather than worked
    # out here. A number this script computed would be this script's opinion
    # of the billing rule, and an example that shows its own opinion of the
    # bill is worth nothing -- it agrees with itself when the rule changes.
    print(f"  batch {i}: {time.monotonic() - t:.2f}s  "
          f"counted {counted(meta)}  "
          f"AUC {plain[-1]}")

print("\nwith the cache")
t = time.monotonic()
proba, ctx = cx.predict_cached(
    history, batches[0], target="income", model="causilo-clf", output_type="probas"
)

print(f"  batch 0: {time.monotonic() - t:.2f}s  "
      f"counted {counted(ctx.metadata)}  (the context travels with this one)")
print("  handle:", ctx)

# The creating call's metadata rides on the handle.
cached = [auc(proba, ctx.metadata, labels[0])]
for i, (batch, truth) in enumerate(zip(batches[1:], labels[1:]), start=1):
    t = time.monotonic()
    proba, meta = ctx.predict_with_metadata(batch)
    cached.append(auc(proba, meta, truth))
    print(f"  batch {i}: {time.monotonic() - t:.2f}s  "
          f"counted {counted(meta)}  "
          f"AUC {cached[-1]}")

# What the server is holding, and what it costs.
for held in cx.list_contexts():
    print("\nheld:", held)

print("released:", ctx.release())
print("held after release:", cx.list_contexts())

print("\nAUC without the cache:", plain)
print("AUC with the cache:   ", cached)