nous 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

nous/explain/loo.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+import numpy as np
+import pandas as pd
+import torch
+from typing import Optional, Sequence
+from ..model import NousNet
+from .facts_desc import render_fact_descriptions
+from .aggregator import format_agg_mixture
+
+def rule_impact_df(
+    model: NousNet, x_sample, feature_names: Sequence[str], class_names: Optional[Sequence[str]] = None, vs_class: Optional[int] = None,
+    loo_mode: str = 'replacement',  # 'replacement' | 'frozen' | 'active_only'
+    top_m_rules: Optional[int] = None,
+    use_pruning: bool = False, pruning_threshold: float = 0.0, use_pre_norm: bool = False,
+    base_state=None, device=None
+) -> pd.DataFrame:
+    """
+    Honest leave-one-out at the rule level with recomputed gating (before top-k),
+    supporting multiple modes (replacement, frozen, active_only).
+    """
+    device = device or next(model.parameters()).device
+    task = model.config['task_type']
+
+    if base_state is None:
+        base_probas, base_logits, base_internals = model.forward_explain(
+            x_sample, apply_pruning=use_pruning, pruning_threshold=pruning_threshold, device=device
+        )
+    else:
+        base_probas, base_logits, base_internals = base_state
+
+    if task == "classification":
+        base_logits = np.array(base_logits)
+        pred_idx = int(np.argmax(base_probas))
+        if vs_class is None:
+            runner_up_idx = int(np.argsort(base_logits)[-2]) if base_logits.size > 1 else pred_idx
+        else:
+            runner_up_idx = int(vs_class)
+        base_margin = float(base_logits[pred_idx] - base_logits[runner_up_idx])
+    else:
+        base_pred = float(base_logits[0])
+
+    restrict_masks = None
+    if loo_mode in ('frozen', 'active_only'):
+        # Freeze current active set by extracting gate masks and reusing as restrict masks.
+        restrict_masks = []
+        block_keys = sorted([k for k in base_internals.keys() if k.startswith("block_")], key=lambda s: int(s.split("_")[1]))
+        for key in block_keys:
+            gm = base_internals[key]['gate_mask']
+            if isinstance(gm, torch.Tensor):
+                gm = gm.squeeze(0)
+            restrict_masks.append((gm > 0).float())
+
+    fact_desc = render_fact_descriptions(model, feature_names)
+
+    rows = []
+    block_keys = sorted([k for k in base_internals.keys() if k.startswith("block_")], key=lambda s: int(s.split("_")[1]))
+    for b_idx, key in enumerate(block_keys):
+        details = base_internals[key]
+        metric_tensor = details['pre_norm_sum'] if (use_pre_norm and 'pre_norm_sum' in details) else details['gated_activations']
+        ga_np = metric_tensor.squeeze(0).abs().cpu().numpy() if isinstance(metric_tensor, torch.Tensor) else np.abs(metric_tensor)
+        active_rules = np.where(ga_np > 1e-12)[0]
+        if len(active_rules) == 0:
+            continue
+
+        if top_m_rules is not None and len(active_rules) > top_m_rules:
+            order = np.argsort(-ga_np[active_rules])
+            active_rules = active_rules[order[:top_m_rules]]
+
+        agg_w = details['aggregator_weights']
+        if isinstance(agg_w, torch.Tensor):
+            agg_w = agg_w.cpu()
+
+        facts_used = details.get('facts_used', None)
+        if isinstance(facts_used, torch.Tensor):
+            facts_used = facts_used.cpu().numpy()
+
+        for r in active_rules:
+            drop_spec = (b_idx, int(r))
+            restr = None if loo_mode == 'replacement' else restrict_masks
+
+            drop_probas, drop_logits, _ = model.forward_explain(
+                x_sample, drop_rule_spec=drop_spec, restrict_masks=restr,
+                apply_pruning=use_pruning, pruning_threshold=pruning_threshold, device=device
+            )
+
+            if agg_w is not None:
+                aggs = format_agg_mixture(agg_w[r])
+            else:
+                aggs = "AND (fixed)"
+
+            facts_str = "—"
+            if facts_used is not None and facts_used.shape[0] > r:
+                used = facts_used[r]
+                if np.ndim(used) == 0:
+                    used = [int(used)]
+                facts_str = ", ".join([f"F{int(fid)+1}" for fid in used])
+                facts_str += "  →  " + " | ".join([f"[F{int(fid)+1}] {fact_desc[int(fid)]}" for fid in used])
+
+            if task == "classification":
+                drop_logits = np.array(drop_logits)
+                dlogit = float(base_logits[pred_idx] - drop_logits[pred_idx])
+                drop_margin = float(drop_logits[pred_idx] - drop_logits[runner_up_idx])
+                dmargin = float(base_margin - drop_margin)
+                critical = (np.argmax(drop_probas) != pred_idx)
+                supports = dmargin > 0
+                rows.append({
+                    "block": b_idx+1, "rule": int(r+1),
+                    "aggregators": aggs, "facts": facts_str,
+                    "Δlogit(pred)": dlogit,
+                    f"Δmargin(vs {class_names[runner_up_idx] if class_names else runner_up_idx})": dmargin,
+                    "supports_pred": supports, "critical_flip": critical
+                })
+            else:
+                drop_pred = float(drop_logits[0])
+                d_pred = base_pred - drop_pred
+                rows.append({
+                    "block": b_idx+1, "rule": int(r+1),
+                    "aggregators": aggs, "facts": facts_str,
+                    "Δprediction": d_pred
+                })
+
+    df = pd.DataFrame(rows)
+    if not df.empty:
+        if task == "classification":
+            key = [c for c in df.columns if c.startswith("Δmargin(")][0]
+            df['abs_impact'] = df[key].abs()
+        else:
+            df['abs_impact'] = df['Δprediction'].abs()
+        df = df.sort_values(by='abs_impact', ascending=False).drop(columns=['abs_impact']).reset_index(drop=True)
+    return df

nous/explain/mse.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+import numpy as np
+import torch
+from typing import Optional, Sequence, Dict, Any, List
+from ..model import NousNet
+from .loo import rule_impact_df
+
+def minimal_sufficient_explanation(
+    model: NousNet, x_sample, feature_names: Sequence[str], class_names: Optional[Sequence[str]] = None,
+    margin_tolerance: float = 0.0, prob_tolerance: float = 0.0, pred_tolerance_reg: float = 0.02,
+    freeze_non_active: bool = True, device=None
+) -> Dict[str, Any]:
+    """
+    Greedy backward elimination among originally active rules.
+    Preserves the prediction (class unchanged with margins within tolerance; regression within abs tol).
+    """
+    device = device or next(model.parameters()).device
+    task = model.config['task_type']
+
+    base_probas, base_logits, base_internals = model.forward_explain(x_sample, device=device)
+    if task == "classification":
+        pred_idx = int(np.argmax(base_probas))
+        runner_up = int(np.argsort(base_logits)[-2]) if base_logits.size > 1 else pred_idx
+        base_margin = float(base_logits[pred_idx] - base_logits[runner_up])
+        base_conf = float(base_probas[pred_idx])
+    else:
+        base_pred = float(base_logits[0])
+
+    # Extract active set (frozen)
+    active_masks: List[torch.Tensor] = []
+    block_keys = sorted([k for k in base_internals.keys() if k.startswith("block_")], key=lambda s: int(s.split("_")[1]))
+    for key in block_keys:
+        gm = base_internals[key]['gate_mask']
+        if isinstance(gm, torch.Tensor):
+            gm = gm.squeeze(0)
+        active_masks.append((gm > 0).float())
+    current_masks = [m.clone() for m in active_masks]
+
+    # Rank removal candidates by (absolute) impact ascending
+    imp = rule_impact_df(
+        model, x_sample, feature_names, class_names=class_names,
+        loo_mode='frozen', device=device
+    )
+    if imp.empty:
+        return {
+            "kept_masks": current_masks, "kept": [],
+            "removed": [], "pred_preserved": True,
+            "size": sum(int(m.sum().item()) for m in current_masks)
+        }
+
+    if task == "classification":
+        margin_col = [c for c in imp.columns if c.startswith("Δmargin(")][0]
+        imp = imp.assign(order=imp[margin_col].abs()).sort_values("order", ascending=True).drop(columns=["order"])
+    else:
+        imp = imp.assign(order=imp["Δprediction"].abs()).sort_values("order", ascending=True).drop(columns=["order"])
+
+    removed = []
+    for _, row in imp.iterrows():
+        b = int(row["block"]) - 1
+        r = int(row["rule"]) - 1
+        if current_masks[b][r] == 0:
+            continue
+
+        proposal = [m.clone() for m in current_masks]
+        proposal[b][r] = 0.0
+        probas2, logits2, _ = model.forward_explain(x_sample, restrict_masks=proposal, device=device)
+
+        if task == "classification":
+            new_pred = int(np.argmax(probas2))
+            if new_pred != pred_idx:
+                continue
+            new_margin = float(logits2[pred_idx] - logits2[runner_up])
+            new_conf = float(probas2[pred_idx])
+            if new_margin < base_margin - margin_tolerance:
+                continue
+            if prob_tolerance > 0.0 and (new_conf < base_conf - prob_tolerance):
+                continue
+        else:
+            new_pred = float(logits2[0])
+            if abs(new_pred - base_pred) > pred_tolerance_reg:
+                continue
+
+        current_masks = proposal
+        removed.append((b+1, r+1))
+
+    kept = [(i+1, int(idx.item())+1) for i, m in enumerate(current_masks) for idx in torch.where(m > 0)[0]]
+    return {
+        "kept_masks": current_masks,
+        "kept": kept,
+        "removed": removed,
+        "pred_preserved": True,
+        "size": sum(int(m.sum().item()) for m in current_masks)
+    }

nous/explain/pruning.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+import numpy as np
+from typing import Optional
+from sklearn.metrics import mean_absolute_error
+from ..model import NousNet
+
+def select_pruning_threshold_global(
+    model: NousNet, X_val, target_fidelity: float = 0.99, task_type: Optional[str] = None,
+    metric_reg: str = "mae", tol_reg: float = 0.05, pointwise_tol: Optional[float] = None,
+    max_samples: int = 500, device=None
+) -> float:
+    """
+    Grid selection of activation threshold (post-gating) keeping fidelity (classification) or MAE (regression).
+    """
+    device = device or next(model.parameters()).device
+    task = task_type or model.config['task_type']
+
+    n = min(len(X_val), max_samples)
+    Xv = X_val[:n]
+
+    base_preds = []
+    acts = []
+    for i in range(n):
+        _, logit_b, internals = model.forward_explain(Xv[i], device=device)
+        base_preds.append(int(np.argmax(logit_b)) if task == "classification" else float(logit_b[0]))
+        for key in [k for k in internals.keys() if k.startswith("block_")]:
+            ga = internals[key]['gated_activations'].abs().cpu().numpy().ravel()
+            acts.extend(list(ga))
+    acts = np.array(acts)
+    if acts.size == 0:
+        return 0.0
+
+    qs = np.linspace(0.5, 0.999, 25)
+    candidates = np.unique(np.quantile(acts, qs))
+
+    best_t = 0.0
+    for t in candidates:
+        if task == "classification":
+            agree = 0
+            for i in range(n):
+                _, logit_p, _ = model.forward_explain(Xv[i], apply_pruning=True, pruning_threshold=float(t), device=device)
+                agree += int(int(np.argmax(logit_p)) == base_preds[i])
+            fidelity = agree / n
+            if fidelity >= target_fidelity:
+                best_t = float(t)
+        else:
+            preds_p, preds_b = [], []
+            max_abs = 0.0
+            for i in range(n):
+                _, logit_p, _ = model.forward_explain(Xv[i], apply_pruning=True, pruning_threshold=float(t), device=device)
+                _, logit_b, _ = model.forward_explain(Xv[i], device=device)
+                pv, bv = float(logit_p[0]), float(logit_b[0])
+                preds_p.append(pv); preds_b.append(bv)
+                max_abs = max(max_abs, abs(pv - bv))
+            mae_p = mean_absolute_error(preds_b, preds_p)
+            ok_mae = (mae_p <= tol_reg)
+            ok_point = True if pointwise_tol is None else (max_abs <= pointwise_tol)
+            if ok_mae and ok_point:
+                best_t = float(t)
+    return best_t
+
+
+def select_pruning_threshold_global_bs(
+    model: NousNet, X_val, target_fidelity: float = 0.99, task_type: Optional[str] = None,
+    metric_reg: str = "mae", tol_reg: float = 0.05, pointwise_tol: Optional[float] = None,
+    max_samples: int = 500, device=None
+) -> float:
+    """
+    Binary search selection of activation threshold (post-gating) with fidelity/MAE constraints.
+    """
+    device = device or next(model.parameters()).device
+    task = task_type or model.config['task_type']
+    n = min(len(X_val), max_samples)
+    Xv = X_val[:n]
+
+    base_preds = []
+    acts = []
+    for i in range(n):
+        _, logit_b, internals = model.forward_explain(Xv[i], device=device)
+        base_preds.append(int(np.argmax(logit_b)) if task == "classification" else float(logit_b[0]))
+        for key in [k for k in internals.keys() if k.startswith("block_")]:
+            ga = internals[key]['gated_activations'].abs().cpu().numpy().ravel()
+            acts.extend(list(ga))
+    acts = np.array(acts)
+    if acts.size == 0:
+        return 0.0
+
+    lo, hi = 0.0, float(np.quantile(acts, 0.999))
+    best_t = 0.0
+    for _ in range(14):
+        mid = (lo + hi) / 2.0
+        if task == "classification":
+            agree = 0
+            for i in range(n):
+                _, logit_p, _ = model.forward_explain(Xv[i], apply_pruning=True, pruning_threshold=float(mid), device=device)
+                agree += int(int(np.argmax(logit_p)) == base_preds[i])
+            fidelity = agree / n
+            if fidelity >= target_fidelity:
+                best_t = mid; lo = mid
+            else:
+                hi = mid
+        else:
+            preds_p, preds_b, max_abs = [], [], 0.0
+            for i in range(n):
+                _, logit_p, _ = model.forward_explain(Xv[i], apply_pruning=True, pruning_threshold=float(mid), device=device)
+                _, logit_b, _ = model.forward_explain(Xv[i], device=device)
+                pv, bv = float(logit_p[0]), float(logit_b[0])
+                preds_p.append(pv); preds_b.append(bv)
+                max_abs = max(max_abs, abs(pv - bv))
+            mae_p = mean_absolute_error(preds_b, preds_p)
+            ok_mae = (mae_p <= tol_reg)
+            ok_point = True if pointwise_tol is None else (max_abs <= pointwise_tol)
+            if ok_mae and ok_point:
+                best_t = mid; lo = mid
+            else:
+                hi = mid
+    return float(best_t)

nous/explain/stability.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+import numpy as np
+import torch
+from typing import Optional, Sequence, Dict, Any
+from ..model import NousNet
+from .loo import rule_impact_df
+
+def explanation_stability(
+    model: NousNet, x_sample, feature_names: Sequence[str], class_names: Optional[Sequence[str]] = None,
+    k_top: int = 5, sigma: float = 0.05, trials: int = 20, loo_mode: str = 'frozen',
+    use_pruning: bool = False, pruning_threshold: float = 0.0, device=None
+) -> Dict[str, Any]:
+    """
+    Stability of top-k rule explanations under small input perturbations.
+    Returns mean/std Jaccard overlap vs base top-k selection.
+    """
+    device = device or next(model.parameters()).device
+    base_imp = rule_impact_df(
+        model, x_sample, feature_names, class_names=class_names,
+        loo_mode=loo_mode, use_pruning=use_pruning, pruning_threshold=pruning_threshold
+    )
+    base_set = set((int(r.block), int(r.rule)) for _, r in base_imp.head(k_top).iterrows()) if not base_imp.empty else set()
+
+    overlaps = []
+    for _ in range(trials):
+        x = torch.tensor(x_sample, dtype=torch.float32, device=device)
+        noise = torch.randn_like(x) * sigma
+        x_noisy = (x + noise).cpu().numpy()
+
+        imp = rule_impact_df(
+            model, x_noisy, feature_names, class_names=class_names,
+            loo_mode=loo_mode, use_pruning=use_pruning, pruning_threshold=pruning_threshold
+        )
+        cur_set = set((int(r.block), int(r.rule)) for _, r in imp.head(k_top).iterrows()) if not imp.empty else set()
+        inter = len(base_set & cur_set)
+        union = len(base_set | cur_set) if (base_set | cur_set) else 1
+        overlaps.append(inter / union)
+    return {
+        "mean_jaccard": float(np.mean(overlaps)),
+        "std_jaccard": float(np.std(overlaps)),
+        "base_top_rules": list(base_set)
+    }

nous/explain/traces.py

@@ -0,0 +1,285 @@
+from __future__ import annotations
+import numpy as np
+import torch
+import torch.nn.functional as F
+from typing import Optional, Sequence, List, Dict, Any
+from ..model import NousNet
+from ..prototypes import ScaledPrototypeLayer
+from ..explain.aggregator import AGG_NAMES
+from ..explain.facts_desc import render_fact_descriptions
+
+@torch.no_grad()
+def get_last_block_static_metadata(model: NousNet, top_k_facts_override: Optional[int] = None):
+    """
+    Returns:
+      - agg_weights: torch.Tensor [R, A] or None
+      - facts_used: np.ndarray [R, K] (indices of inputs to the last block) or None
+    """
+    blk = model.blocks[-1]
+    agg = None
+    facts_used = None
+
+    if hasattr(blk, "aggregator_logits"):
+        agg = F.softmax(blk.aggregator_logits, dim=1).detach().cpu()
+
+    # Fact indices used per rule depend on block type
+    from ..rules.blocks import SimpleNousBlock
+    from ..rules.softmax import SoftmaxRuleLayer
+    from ..rules.sparse import SparseRuleLayer
+
+    if isinstance(blk, SimpleNousBlock):
+        facts_used = blk.rule.idx.detach().cpu().numpy()  # [R, 2]
+    elif isinstance(blk, SoftmaxRuleLayer):
+        k = top_k_facts_override or getattr(blk, "top_k_facts", 2)
+        fl = F.softmax(blk.fact_logits, dim=1)
+        _, topk = torch.topk(fl, k=min(k, blk.input_dim), dim=1)
+        facts_used = topk.detach().cpu().numpy()  # [R, k]
+    elif isinstance(blk, SparseRuleLayer):
+        k = top_k_facts_override or getattr(blk, "top_k_facts", 2)
+        prob = blk.hard_concrete.get_proba().detach()
+        _, topk = torch.topk(prob, k=min(k, blk.input_dim), dim=1)
+        facts_used = topk.cpu().numpy()
+    return agg, facts_used
+
+@torch.no_grad()
+def _block_fact_mapping(block, top_k_per_rule: int = 2):
+    """
+    Returns np.ndarray [R, k] — input indices of the given block for each rule.
+    """
+    from ..rules.blocks import SimpleNousBlock
+    from ..rules.softmax import SoftmaxRuleLayer
+    from ..rules.sparse import SparseRuleLayer
+
+    if isinstance(block, SimpleNousBlock):
+        return block.rule.idx.detach().cpu().numpy()
+    elif isinstance(block, SoftmaxRuleLayer):
+        fl = F.softmax(block.fact_logits, dim=1)
+        _, topk = torch.topk(fl, k=min(top_k_per_rule, block.input_dim), dim=1)
+        return topk.detach().cpu().numpy()
+    elif isinstance(block, SparseRuleLayer):
+        prob = block.hard_concrete.get_proba().detach()
+        _, topk = torch.topk(prob, k=min(top_k_per_rule, block.input_dim), dim=1)
+        return topk.cpu().numpy()
+    else:
+        return None
+
+@torch.no_grad()
+def trace_rule_to_base_facts(model: NousNet, rule_idx_last: int, top_k_per_step: int = 2) -> List[int]:
+    """
+    Trace a rule from the last block down to base β-facts.
+    """
+    selected = {int(rule_idx_last)}
+    for b in reversed(range(len(model.blocks))):
+        blk = model.blocks[b]
+        mapping = _block_fact_mapping(blk, top_k_per_rule=top_k_per_step)
+        if mapping is None:
+            return []
+        prev_sel = set()
+        for r in selected:
+            if r < 0 or r >= mapping.shape[0]:
+                continue
+            prev_sel.update([int(i) for i in mapping[r].tolist()])
+        selected = prev_sel
+    return sorted(selected)
+
+@torch.no_grad()
+def prototype_top_rules(model: NousNet, proto_id: int, top_k_rules: int = 10):
+    """
+    Returns list of (rule_idx, weight) from the prototype vector sorted by |weight|.
+    """
+    assert isinstance(model.head, ScaledPrototypeLayer), "Prototypes head is not enabled."
+    P = model.head.prototypes.detach().cpu().numpy()
+    v = P[proto_id]
+    idx = np.argsort(-np.abs(v))[:top_k_rules]
+    return [(int(i), float(v[i])) for i in idx]
+
+@torch.no_grad()
+def prototype_contribution_df(model: NousNet, x_sample, class_names: Optional[Sequence[str]] = None, top_k: int = 5, device=None):
+    """
+    For a single example — contribution of prototypes to the predicted class.
+    """
+    import pandas as pd
+
+    assert isinstance(model.head, ScaledPrototypeLayer), "Prototypes head is not enabled."
+    device = device or next(model.parameters()).device
+
+    h = model.encode(np.array([x_sample]), device=device)
+    d, act = model.head.compute_dist_act(h.to(device))
+    act = act.squeeze(0).cpu().numpy()
+    d = d.squeeze(0).cpu().numpy()
+
+    W = model.head.prototype_class.detach().cpu().numpy()  # [M, C]
+    probs, logits, _ = model.forward_explain(x_sample, device=device)
+    c_hat = int(np.argmax(probs))
+    w_c = W[:, c_hat]
+    contrib = act * w_c
+
+    Wsm = F.softmax(model.head.prototype_class, dim=1).detach().cpu().numpy()
+    ent = (-Wsm * np.clip(np.log(Wsm + 1e-12), -50, 50)).sum(axis=1)
+    primary = np.argmax(Wsm, axis=1)
+
+    order = np.argsort(-np.abs(contrib))[:top_k]
+    rows = []
+    for j in order:
+        rows.append({
+            "proto": int(j),
+            "distance": float(d[j]),
+            "activation": float(act[j]),
+            "w_c": float(w_c[j]),
+            "contribution": float(contrib[j]),
+            "primary_class": (class_names[primary[j]] if class_names is not None else int(primary[j])),
+            "entropy": float(ent[j])
+        })
+    import pandas as pd
+    return pd.DataFrame(rows)
+
+@torch.no_grad()
+def prototype_report_global(
+    model: NousNet, X, y: Optional[Sequence[int]] = None, class_names: Optional[Sequence[str]] = None,
+    top_k_rules: int = 8, top_k_facts_per_rule: int = 2, trace_to_base: bool = True,
+    k_neighbors: int = 10, chunk_size: int = 2048, device=None
+):
+    """
+    Global prototype report:
+      - primary_class probabilities and entropy,
+      - average activation, top1 frequency,
+      - top rules and (optionally) tracing to base β-facts,
+      - nearest training examples and labels (if y provided).
+    """
+    import pandas as pd
+    assert isinstance(model.head, ScaledPrototypeLayer), "Prototypes head is not enabled."
+    device = device or next(model.parameters()).device
+    model.eval()
+
+    H = model.encode(X, device=device)
+    Hn = F.normalize(H, p=2, dim=1)
+    P = model.head.get_params()
+    Pn = P["prototypes_norm"]
+    Wsm = P["class_probs"].cpu().numpy()
+    Wraw = P["class_weights"].cpu().numpy()
+
+    M = Pn.shape[0]
+    N = Hn.shape[0]
+    top1_count = np.zeros(M, dtype=int)
+    mean_act = np.zeros(M, dtype=float)
+    neigh_idx: list[list[int]] = [[] for _ in range(M)]
+
+    for i in range(0, N, chunk_size):
+        h_chunk = Hn[i:i+chunk_size].to(device)
+        d = torch.cdist(h_chunk, Pn.to(device))
+        act = torch.exp(-F.softplus(model.head.temperature) * d).cpu().numpy()
+        top_idx = np.argmax(act, axis=1)
+        for t in top_idx:
+            top1_count[int(t)] += 1
+        mean_act += act.sum(axis=0)
+        k = min(k_neighbors, act.shape[0])
+        for m in range(M):
+            order = np.argsort(d[:, m].cpu().numpy())[:k]
+            neigh_idx[m].extend((i + order).tolist())
+
+    mean_act /= max(1, N)
+
+    agg_last, facts_last = get_last_block_static_metadata(model, top_k_facts_override=top_k_facts_per_rule)
+    fd = render_fact_descriptions(model, model.config['feature_names'])
+
+    def agg_row_to_str(aw_row):
+        if aw_row is None:
+            return "AND (fixed)"
+        A = int(aw_row.shape[0]) if hasattr(aw_row, "shape") else len(AGG_NAMES)
+        names = AGG_NAMES[:A]
+        parts = []
+        for i, a in enumerate(aw_row.tolist()):
+            if i < len(names) and a > 1e-6:
+                parts.append(f"{float(a):.2f} {names[i]}")
+        return " + ".join(parts) if parts else "∅"
+
+    rule_meta = []
+    for m in range(M):
+        top_rules = prototype_top_rules(model, m, top_k_rules=top_k_rules)
+        rules_desc = []
+        for (r, w) in top_rules:
+            r = int(r)
+            agg_str = "AND (fixed)"
+            if agg_last is not None and 0 <= r < agg_last.shape[0]:
+                aw = agg_last[r]
+                agg_str = agg_row_to_str(aw)
+            facts_str = ""
+            if trace_to_base:
+                base_facts = trace_rule_to_base_facts(model, r, top_k_per_step=top_k_facts_per_rule)
+                if base_facts:
+                    facts_str = " | ".join([f"[F{int(fid)+1}] {fd[int(fid)]}" for fid in base_facts[:top_k_facts_per_rule]])
+            elif facts_last is not None and r < facts_last.shape[0]:
+                prev_units = facts_last[r]
+                facts_str = ", ".join([f"Unit{int(u)}" for u in prev_units])
+            rules_desc.append(dict(rule=int(r), weight=float(w), aggregators=agg_str, facts=facts_str))
+        rule_meta.append(rules_desc)
+
+    rows = []
+    for m in range(M):
+        primary = int(np.argmax(Wsm[m]))
+        rows.append({
+            "proto": m,
+            "primary_class": (class_names[primary] if class_names is not None else primary),
+            "class_probs": Wsm[m].tolist(),
+            "class_weights": Wraw[m].tolist(),
+            "class_entropy": float((-Wsm[m] * np.log(Wsm[m] + 1e-12)).sum()),
+            "mean_activation": float(mean_act[m]),
+            "top1_count": int(top1_count[m]),
+            "neighbors": neigh_idx[m][:k_neighbors],
+            "neighbor_labels": [int(y[idx]) for idx in neigh_idx[m][:k_neighbors]] if y is not None else None,
+            "top_rules_meta": rule_meta[m]
+        })
+    return pd.DataFrame(rows)
+
+@torch.no_grad()
+def describe_prototype(model: NousNet, proto_id: int, feature_names: Sequence[str], class_names: Optional[Sequence[str]] = None,
+                       top_k_rules: int = 8, top_k_facts_per_rule: int = 2, trace_to_base: bool = True) -> str:
+    """
+    Human-readable description for a single prototype:
+      - class affinity (weights),
+      - top rules and aggregator mixtures,
+      - base β-facts (if traced).
+    """
+    assert isinstance(model.head, ScaledPrototypeLayer), "Prototypes head is not enabled."
+
+    P = model.head.get_params()
+    Wsm = P["class_probs"].cpu().numpy()
+    m = int(proto_id)
+    primary = int(np.argmax(Wsm[m]))
+    primary_name = class_names[primary] if class_names is not None else f"Class {primary}"
+
+    lines = []
+    lines.append(f"Prototype #{m}")
+    lines.append(f"  Primary class: {primary_name} | probs={np.round(Wsm[m], 3).tolist()} | entropy={float((-Wsm[m]*np.log(Wsm[m]+1e-12)).sum()):.3f}")
+    lines.append("  Top rules:")
+    agg_last, _ = get_last_block_static_metadata(model, top_k_facts_override=top_k_facts_per_rule)
+    fd = render_fact_descriptions(model, feature_names)
+
+    def agg_row_to_str(aw_row):
+        if aw_row is None:
+            return "AND (fixed)"
+        A = int(aw_row.shape[0]) if hasattr(aw_row, "shape") else len(AGG_NAMES)
+        names = AGG_NAMES[:A]
+        parts = []
+        for i, a in enumerate(aw_row.tolist()):
+            if i < len(names) and a > 1e-6:
+                parts.append(f"{float(a):.2f} {names[i]}")
+        return " + ".join(parts) if parts else "∅"
+
+    for (r, w) in prototype_top_rules(model, m, top_k_rules=top_k_rules):
+        r = int(r)
+        agg_str = "AND (fixed)"
+        if agg_last is not None and 0 <= r < agg_last.shape[0]:
+            aw = agg_last[r]
+            agg_str = agg_row_to_str(aw)
+
+        facts_str = ""
+        if trace_to_base:
+            base_facts = trace_rule_to_base_facts(model, r, top_k_per_step=top_k_facts_per_rule)
+            if base_facts:
+                facts_str = " | ".join([f"[F{int(fid)+1}] {fd[int(fid)]}" for fid in base_facts[:top_k_facts_per_rule]])
+
+        lines.append(f"    • R{r+1}: weight={w:+.3f} | {agg_str}")
+        if facts_str:
+            lines.append(f"        {facts_str}")
+    return "\n".join(lines)

nous/explain/utils.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+import torch
+
+def extract_active_masks(internals) -> list[torch.Tensor]:
+    """
+    Extract per-block active gate masks (0/1) from internals produced by forward_explain.
+    """
+    masks = []
+    keys = sorted([k for k in internals.keys() if k.startswith("block_")], key=lambda s: int(s.split("_")[1]))
+    for key in keys:
+        gate_mask = internals[key]['gate_mask']
+        if isinstance(gate_mask, torch.Tensor):
+            gate_mask = gate_mask.squeeze(0)
+        masks.append((gate_mask > 0).float())
+    return masks