supply_chain_network_inference.py

"""
Supply Chain Network Inference from Inventory Snapshots
========================================================
Companion code for "Chain Reactions Are Delayed" — Affective Analytics
https://frankcorrigan.substack.com

Reproduces the three-method comparison from the article:
  1. Correlation
  2. Lag Correlation
  3. Convergent Cross-Mapping (NIPS paper method)

Paper: Singhal B, Kiss IZ, Li J-S.
       "Network Inference with Partial State measurements using forced-delay embedding."
       PNAS Nexus, Vol. 5, Issue 1, January 2026.

Requirements: numpy, scipy, matplotlib
  pip install numpy scipy matplotlib
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from scipy.stats import pearsonr

# ── RANDOM SEED ───────────────────────────────────────────────────────────────
np.random.seed(42)
rng = np.random.default_rng(42)


# ══════════════════════════════════════════════════════════════════════════════
# 1. GROUND TRUTH NETWORK
# ══════════════════════════════════════════════════════════════════════════════

NODE_NAMES = ["Supplier", "DC1", "DC2", "Store1", "Store2", "Store3"]
N = len(NODE_NAMES)

# True directed edges: (source, destination, lead_time_days)
TRUE_EDGES = [
    {"from": 0, "to": 1, "lag": 7},   # Supplier → DC1,    7 days
    {"from": 0, "to": 2, "lag": 5},   # Supplier → DC2,    5 days
    {"from": 1, "to": 3, "lag": 3},   # DC1      → Store1, 3 days
    {"from": 1, "to": 4, "lag": 4},   # DC1      → Store2, 4 days
    {"from": 2, "to": 5, "lag": 2},   # DC2      → Store3, 2 days
]
TRUE_SET = set((e["from"], e["to"]) for e in TRUE_EDGES)


# ══════════════════════════════════════════════════════════════════════════════
# 2. SIMULATE INVENTORY
# ══════════════════════════════════════════════════════════════════════════════

def simulate_inventory(T=200, seed=42):
    """
    Event-driven, mass-conserving inventory simulation.

    Each node holds inventory. Downstream nodes consume (demand).
    When a node drops below its reorder point, the upstream node
    ships a replenishment that arrives after the true lead time.
    The supplier is replenished from outside when it runs low.

    Returns: inv array of shape (N, T)
    """
    rng = np.random.default_rng(seed)

    BASE_DEMAND   = [0, 6, 5, 3, 3, 3]      # units consumed per day
    REORDER_POINT = [150, 60, 60, 30, 30, 30]
    ORDER_QTY     = [100, 50, 50, 20, 20, 20]
    SUPPLIER_REORDER = 150

    inv = np.zeros((N, T))
    inv[:, 0] = [200, 100, 100, 50, 50, 50]

    # In-transit pipeline queues: one array per edge of length = lag
    # pipelines[i][t] = units arriving at destination in t more days
    pipes = {(e["from"], e["to"], e["lag"]): np.zeros(e["lag"]) for e in TRUE_EDGES}

    for t in range(1, T):
        inflows = np.zeros(N)
        for (s, d, lag), pipe in pipes.items():
            inflows[d] += pipe[-1]

        demand = np.array([
            max(0, BASE_DEMAND[i] + rng.normal(0, 1.2))
            for i in range(N)
        ])

        ext = SUPPLIER_REORDER + rng.normal(0, 8) if inv[0, t-1] < SUPPLIER_REORDER else 0

        for i in range(N):
            inv[i, t] = max(5, inv[i, t-1] - demand[i] + inflows[i] + (ext if i == 0 else 0))

        for (s, d, lag), pipe in pipes.items():
            pipes[(s, d, lag)] = np.roll(pipe, 1)
            pipes[(s, d, lag)][0] = 0

            if inv[d, t] < REORDER_POINT[d] and inv[s, t] > REORDER_POINT[s] * 0.5:
                qty = max(5, min(
                    ORDER_QTY[d] + rng.normal(0, 3),
                    inv[s, t] - REORDER_POINT[s] * 0.4
                ))
                inv[s, t] -= qty
                pipes[(s, d, lag)][0] = qty

    return inv


# ══════════════════════════════════════════════════════════════════════════════
# 3. INFERENCE METHODS
# ══════════════════════════════════════════════════════════════════════════════

def method_correlation(inv, threshold=0.25):
    """
    Method 1: Pearson correlation between all pairs of inventory time series.
    Direction is ambiguous — all pairs above threshold treated as bidirectional,
    then deduplicated by keeping i→j if corr(i,j) > corr(j,i).
    """
    edges = []
    scores = np.zeros((N, N))

    for i in range(N):
        for j in range(N):
            if i == j:
                continue
            r, _ = pearsonr(inv[i], inv[j])
            scores[i, j] = abs(r)

    for i in range(N):
        for j in range(N):
            if i != j and scores[i, j] > threshold and scores[i, j] >= scores[j, i]:
                edges.append({"from": i, "to": j, "weight": round(scores[i, j], 3)})

    return edges


def method_lag_correlation(inv, max_lag=15, threshold=0.30):
    """
    Method 2: Cross-correlation with lag sweep.
    For each ordered pair (i, j), find the lag where correlation peaks.
    If lag > 0 and i→j is stronger than j→i, infer i drives j.
    """
    best_r   = np.zeros((N, N))
    best_lag = np.zeros((N, N), dtype=int)

    for i in range(N):
        for j in range(N):
            if i == j:
                continue
            for lag in range(1, max_lag + 1):
                r, _ = pearsonr(inv[i, :-lag], inv[j, lag:])
                if abs(r) > abs(best_r[i, j]):
                    best_r[i, j]   = r
                    best_lag[i, j] = lag

    edges = []
    for i in range(N):
        for j in range(N):
            if i == j:
                continue
            if abs(best_r[i, j]) > threshold and abs(best_r[i, j]) > abs(best_r[j, i]):
                edges.append({
                    "from": i, "to": j,
                    "lag":  int(best_lag[i, j]),
                    "weight": round(float(abs(best_r[i, j])), 3)
                })

    return edges


def delay_embed(x, E, tau):
    """
    Construct Takens delay-embedding matrix.
    Row t: [x(t), x(t - tau), x(t - 2*tau), ..., x(t - (E-1)*tau)]
    Shape: (len(x) - (E-1)*tau,  E)
    """
    start = (E - 1) * tau
    return np.array([
        [x[t - e * tau] for e in range(E)]
        for t in range(start, len(x))
    ])


def ccm_score(x, y, E=3, tau=1):
    """
    Convergent Cross-Mapping: does X drive Y?

    Tests whether Y's shadow manifold (delay embedding of y) can
    reconstruct X's states via nearest-neighbor prediction.
    High skill → X drives Y.

    Uses full library (single score, not convergence test).
    For the convergence test, call with increasing lib_size subsets.
    """
    My = delay_embed(y, E, tau)
    n  = len(My)
    x_aligned = x[(E - 1) * tau : (E - 1) * tau + n]

    x_pred = []
    x_true = []

    for i in range(n):
        dists = np.sqrt(((My - My[i]) ** 2).sum(axis=1))
        dists[i] = np.inf  # exclude self

        nn_idx   = np.argsort(dists)[:E + 1]
        nn_dists = dists[nn_idx]

        if nn_dists[0] == 0:
            w = np.zeros(E + 1); w[0] = 1.0
        else:
            w = np.exp(-nn_dists / nn_dists[0])
            w /= w.sum()

        x_pred.append((w * x_aligned[nn_idx]).sum())
        x_true.append(x_aligned[i])

    r, _ = pearsonr(x_true, x_pred)
    return r


def method_ccm(inv, max_lag=12, threshold=0.35, E=3):
    """
    Method 3: Convergent Cross-Mapping with forced lag (NIPS paper approach).

    For each ordered pair (i, j) and each candidate lag, embed j's time series
    at tau = candidate_lag. Test if j's manifold can reconstruct i.
    High CCM score at a specific lag → i drives j with that lead time.

    Note: CCM was designed for smooth continuous dynamical systems.
    Threshold-triggered inventory data (event-driven, short time series)
    is a domain mismatch — see article for full discussion.
    """
    best_score = np.zeros((N, N))
    best_lag   = np.zeros((N, N), dtype=int)

    for i in range(N):
        for j in range(N):
            if i == j:
                continue
            for lag in range(1, max_lag + 1):
                score = ccm_score(inv[i], inv[j], E=E, tau=lag)
                if score > best_score[i, j]:
                    best_score[i, j] = score
                    best_lag[i, j]   = lag

    edges = []
    for i in range(N):
        for j in range(N):
            if i == j:
                continue
            if best_score[i, j] > threshold and best_score[i, j] > best_score[j, i]:
                edges.append({
                    "from": i, "to": j,
                    "lag":  int(best_lag[i, j]),
                    "weight": round(float(best_score[i, j]), 3)
                })

    return edges


# ══════════════════════════════════════════════════════════════════════════════
# 4. SCORING
# ══════════════════════════════════════════════════════════════════════════════

def score(edges, label):
    found = set((e["from"], e["to"]) for e in edges)
    tp = len(found & TRUE_SET)
    fp = len(found - TRUE_SET)
    fn = len(TRUE_SET - found)
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall    = tp / (tp + fn) if (tp + fn) > 0 else 0
    print(f"  {label:<30} edges={len(edges):2d}  TP={tp}  FP={fp}  FN={fn}"
          f"  Precision={precision:.0%}  Recall={recall:.0%}")
    return {"label": label, "edges": edges, "tp": tp, "fp": fp, "fn": fn,
            "precision": precision, "recall": recall}


# ══════════════════════════════════════════════════════════════════════════════
# 5. VISUALIZATION
# ══════════════════════════════════════════════════════════════════════════════

# Node positions for network diagrams
NODE_POS = {
    0: (0.5, 0.90),   # Supplier
    1: (0.22, 0.50),  # DC1
    2: (0.78, 0.50),  # DC2
    3: (0.08, 0.10),  # Store1
    4: (0.38, 0.10),  # Store2
    5: (0.78, 0.10),  # Store3
}

NODE_COLORS_PLOT = ["#222", "#444", "#444", "#999", "#999", "#999"]


def draw_network(ax, edges, edge_color, title, stats_str, highlight_true=True):
    ax.set_xlim(-0.05, 1.05)
    ax.set_ylim(-0.05, 1.05)
    ax.set_aspect("equal")
    ax.axis("off")
    ax.set_title(title, fontsize=9, fontfamily="monospace", pad=6)

    for e in edges:
        s, d = NODE_POS[e["from"]], NODE_POS[e["to"]]
        is_true = (e["from"], e["to"]) in TRUE_SET
        color = "#2a7a4b" if (is_true and highlight_true) else edge_color
        alpha = 0.85 if is_true else 0.20
        lw    = 1.8  if is_true else 0.8
        ax.annotate(
            "", xy=d, xytext=s,
            arrowprops=dict(
                arrowstyle="-|>", color=color,
                lw=lw, alpha=alpha,
                connectionstyle="arc3,rad=0.08",
                mutation_scale=10
            )
        )

    for i, (x, y) in NODE_POS.items():
        ax.plot(x, y, "o", color=NODE_COLORS_PLOT[i], markersize=10, zorder=5)
        ax.text(x, y - 0.10, NODE_NAMES[i], ha="center", va="top",
                fontsize=7, fontfamily="monospace", color="#555")

    ax.text(0.5, -0.04, stats_str, ha="center", va="top",
            fontsize=7.5, fontfamily="monospace", color="#888",
            transform=ax.transAxes)


def plot_inventory(inv, T_plot=60):
    colors = ["#111", "#b84a2e", "#2563a8", "#b8872e", "#7a2e8a", "#2a7a4b"]
    fig, ax = plt.subplots(figsize=(9, 3.2))
    for i in range(N):
        ax.plot(inv[i, :T_plot], color=colors[i], lw=1.8 if i == 0 else 1.2,
                alpha=1 if i == 0 else 0.7, label=NODE_NAMES[i])
    ax.set_xlabel("Day", fontfamily="monospace", fontsize=9)
    ax.set_ylabel("Inventory", fontfamily="monospace", fontsize=9)
    ax.set_title("What we observe — inventory snapshots (60 days)",
                 fontfamily="monospace", fontsize=10)
    ax.legend(fontsize=8, loc="upper right")
    ax.spines[["top", "right"]].set_visible(False)
    ax.tick_params(labelsize=8)
    fig.tight_layout()
    return fig


def plot_networks(results):
    configs = [
        ("Ground Truth",          TRUE_EDGES,           "#111",     False,
         "5 edges  |  Precision: 100%  |  Recall: 100%"),
        (results[0]["label"],     results[0]["edges"],   "#b84a2e",  True,
         f"{len(results[0]['edges'])} edges  |  "
         f"Precision: {results[0]['precision']:.0%}  |  Recall: {results[0]['recall']:.0%}"),
        (results[1]["label"],     results[1]["edges"],   "#2563a8",  True,
         f"{len(results[1]['edges'])} edges  |  "
         f"Precision: {results[1]['precision']:.0%}  |  Recall: {results[1]['recall']:.0%}"),
        (results[2]["label"],     results[2]["edges"],   "#2a7a4b",  True,
         f"{len(results[2]['edges'])} edges  |  "
         f"Precision: {results[2]['precision']:.0%}  |  Recall: {results[2]['recall']:.0%}"),
    ]

    fig, axes = plt.subplots(1, 4, figsize=(14, 4))
    fig.suptitle("Network Recovery — Side by Side",
                 fontfamily="monospace", fontsize=11, y=1.01)

    for ax, (title, edges, color, ht, stats) in zip(axes, configs):
        draw_network(ax, edges, color, title, stats, highlight_true=ht)

    # Legend for true edges
    true_patch  = mpatches.Patch(color="#2a7a4b", label="True edge")
    false_patch = mpatches.Patch(color="#bbb",    label="False positive")
    fig.legend(handles=[true_patch, false_patch], loc="lower center",
               ncol=2, fontsize=8, frameon=False, bbox_to_anchor=(0.5, -0.06))

    fig.tight_layout()
    return fig


# ══════════════════════════════════════════════════════════════════════════════
# 6. MAIN
# ══════════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":

    print("=" * 60)
    print("SUPPLY CHAIN NETWORK INFERENCE")
    print("Chain Reactions Are Delayed — Affective Analytics")
    print("=" * 60)

    # ── Simulate ──────────────────────────────────────────────────
    print("\n[1] Simulating inventory...")
    inv = simulate_inventory(T=200)

    for i, name in enumerate(NODE_NAMES):
        print(f"    {name:<10} mean={inv[i].mean():.1f}  "
              f"std={inv[i].std():.1f}  "
              f"min={inv[i].min():.1f}  "
              f"max={inv[i].max():.1f}")

    # ── Run inference ──────────────────────────────────────────────
    print("\n[2] Running inference methods...")
    true_edge_str = [(NODE_NAMES[e["from"]], NODE_NAMES[e["to"]], f"lag={e['lag']}") for e in TRUE_EDGES]
    print(f"    True edges: {true_edge_str}\n")

    corr_edges = method_correlation(inv)
    lag_edges  = method_lag_correlation(inv)

    print("    Running CCM (this takes ~30 seconds)...")
    ccm_edges  = method_ccm(inv)

    # ── Score ──────────────────────────────────────────────────────
    print("\n[3] Results:")
    results = [
        score(corr_edges, "Correlation"),
        score(lag_edges,  "Lag Correlation"),
        score(ccm_edges,  "Convergent Cross-Mapping"),
    ]

    # ── Plot ───────────────────────────────────────────────────────
    print("\n[4] Generating plots...")

    fig_inv = plot_inventory(inv)
    fig_inv.savefig("inventory_snapshots.png", dpi=150, bbox_inches="tight")
    print("    Saved: inventory_snapshots.png")

    fig_net = plot_networks(results)
    fig_net.savefig("network_recovery.png", dpi=150, bbox_inches="tight")
    print("    Saved: network_recovery.png")

    plt.show()

    print("\nDone.")
    print("Read the article: https://frankcorrigan.substack.com")