Listing: Listing

def first_completion(correct, window=5, required=4):
    correct = correct.bool()
    horizon = len(correct)
    if window < 1 or not 0 <= required <= window:
        raise ValueError("invalid recovery pattern")
    for end in range(window, horizon + 1):
        if int(correct[end-window:end].sum()) >= required:
            return {
                "first_completion": end,
                "horizon": horizon,
                "censored": False,
                "restricted_time": end,
            }
    return {
        "first_completion": None,
        "horizon": horizon,
        "censored": True,
        "restricted_time": horizon,
    }

def sequence_metrics(points, rules, predictions, post_window=10,
                     stable_margin=5, recovery_horizon=50):
    if points.ndim != 2 or points.shape[1] != 2:
        raise ValueError("one [T, 2] sequence required")
    length = len(points)
    if length < 1 or rules.shape != (length,) \
            or predictions.shape != (length,):
        raise ValueError("aligned nonempty sequence required")
    if rules.dtype != torch.long or predictions.dtype != torch.long:
        raise ValueError("integer labels required")
    if not bool(((rules >= 0) & (rules < 4)).all()) \
            or not bool(((predictions >= 0) & (predictions < 4)).all()):
        raise ValueError("rule or prediction outside label range")
    if post_window < 1 or stable_margin < 0 or recovery_horizon < 1:
        raise ValueError("invalid window controls")

    mappings = rule_table(points)
    labels = mappings.gather(-1, rules[:, None]).squeeze(-1)
    correct = predictions == labels
    switches = (torch.nonzero(rules[1:] != rules[:-1]).flatten() + 1).tolist()
    boundaries = [0] + switches + [length]
    semantic = effective_rule_ids(rules)

    old_count = new_count = other_count = informative = 0
    repeated_errors = post_errors = 0
    early_correct = early_exposure = 0
    function_switches = 0
    recovery = []

    for j, start in enumerate(switches):
        next_start = boundaries[j + 2]
        stop = min(start + post_window, next_start, length)
        old_id, new_id = int(rules[start-1]), int(rules[start])
        old_label = mappings[start:stop, old_id]
        new_label = mappings[start:stop, new_id]
        pred = predictions[start:stop]
        info = old_label != new_label
        old = info & (pred == old_label)
        new = info & (pred == new_label)
        other = info & ~old & ~new

        informative += int(info.sum())
        old_count += int(old.sum())
        new_count += int(new.sum())
        other_count += int(other.sum())

        errors = ~correct[start:stop]
        repeat = predictions[start:stop] == predictions[start-1:stop-1]
        post_errors += int(errors.sum())
        repeated_errors += int((errors & repeat).sum())
        early_correct += int(correct[start:stop].sum())
        early_exposure += stop - start

        if int(semantic[start-1]) != int(semantic[start]):
            function_switches += 1
            end = min(start + recovery_horizon, next_start, length)
            record = first_completion(correct[start:end])
            record["switch_index"] = start
            recovery.append(record)

    stable_correct = stable_exposure = 0
    for left, right in zip(boundaries[:-1], boundaries[1:]):
        lo, hi = left + stable_margin, right - stable_margin
        if lo < hi:
            stable_correct += int(correct[lo:hi].sum())
            stable_exposure += hi - lo

    assert old_count + new_count + other_count == informative
    return {
        "accuracy": float(correct.double().mean()),
        "id_switches": len(switches),
        "function_switches": function_switches,
        "informative_trials": informative,
        "old_responses": old_count,
        "new_responses": new_count,
        "other_errors": other_count,
        "old_fraction_sample": (
            old_count / informative if informative else None
        ),
        "repeated_error_fraction_sample": (
            repeated_errors / post_errors if post_errors else None
        ),
        "p_old": old_count / length,
        "n_info": informative / length,
        "p_repeat": repeated_errors / length,
        "n_error": post_errors / length,
        "c_early": early_correct / length,
        "n_early": early_exposure / length,
        "c_stable": stable_correct / length,
        "n_stable": stable_exposure / length,
        "recovery_records": recovery,
        "normalization": "counts divided by complete sequence length",
    }

def ratio_certificate(rows, numerator, denominator, delta=0.05,
                      number_of_ratios=1, upper_one=True):
    if not rows or not 0 < delta < 1 or number_of_ratios < 1:
        raise ValueError("invalid certification controls")
    p = torch.tensor([r[numerator] for r in rows], dtype=torch.float64)
    n = torch.tensor([r[denominator] for r in rows], dtype=torch.float64)
    if not bool(((p >= 0) & (p <= 1) & (n >= 0) & (n <= 1)).all()):
        raise ValueError("bounded sequence summaries required")
    if upper_one and not bool((p <= n).all()):
        raise ValueError("fraction numerator must not exceed denominator")

    # Two means for each of the predeclared ratios.
    radius = math.sqrt(
        math.log(4.0 * number_of_ratios / delta) / (2.0 * len(rows))
    )
    pm, nm = float(p.mean()), float(n.mean())
    pl, pu = max(0.0, pm-radius), min(1.0, pm+radius)
    nl, nu = max(0.0, nm-radius), min(1.0, nm+radius)
    result = {
        "numerator_interval": (pl, pu),
        "denominator_interval": (nl, nu),
        "radius": radius,
        "target": "ratio of population sequence-summary means",
        "estimate": pm / nm if nm > 0 else None,
    }
    if nl == 0.0:
        result.update(status="unresolved_denominator", interval=None)
    else:
        lower, upper = pl / nu, pu / nl
        if upper_one:
            upper = min(1.0, upper)
        result.update(status="certified_interval", interval=(lower, upper))
    return result