# Define reusable helpers for the Helper Functions section.
def as_column(values):
"""
Return values as a two-dimensional column array expected by causal-learn ANM.
Parameters
----------
values : array-like
Numeric values transformed, standardized, summarized, or plotted by the helper.
Returns
-------
np.ndarray
Two-dimensional column vector version of the supplied values.
"""
return np.asarray(values).reshape(-1, 1)
def scalar_p_value(value):
"""
Convert a causal-learn p-value result into a plain Python float.
Parameters
----------
value : scalar or array-like
P-value object returned by ANM or PNL. Some causal-learn methods return
one-element arrays rather than bare scalars.
Returns
-------
float
First finite scalar p-value extracted from the result.
"""
array = np.asarray(value, dtype=float).reshape(-1)
finite_values = array[np.isfinite(array)]
if finite_values.size == 0:
return np.nan
return float(finite_values[0])
def run_anm_pair(data_x, data_y, x_name="x", y_name="y"):
"""
Run causal-learn ANM in both directions and return a tidy result row.
Parameters
----------
data_x : object
First variable sample used in directional functional-causal-model tests.
data_y : object
Second variable sample used in directional functional-causal-model tests.
x_name : str
Readable label for `x`.
y_name : str
Readable label for `y`.
Returns
-------
dict[str, object]
Directional ANM comparison row with p-values for both directions and the preferred causal direction.
"""
start = time.perf_counter()
p_forward, p_backward = ANM().cause_or_effect(as_column(data_x), as_column(data_y))
elapsed = time.perf_counter() - start
p_forward = scalar_p_value(p_forward)
p_backward = scalar_p_value(p_backward)
if p_forward > p_backward:
preferred_direction = f"{x_name} -> {y_name}"
preferred_source = x_name
preferred_target = y_name
else:
preferred_direction = f"{y_name} -> {x_name}"
preferred_source = y_name
preferred_target = x_name
return {
"x": x_name,
"y": y_name,
"p_value_x_to_y": p_forward,
"p_value_y_to_x": p_backward,
"preferred_direction": preferred_direction,
"preferred_source": preferred_source,
"preferred_target": preferred_target,
"p_value_margin": abs(p_forward - p_backward),
"elapsed_seconds": elapsed,
}
def fit_anm_residuals(data_x, data_y):
"""
Fit GP regressions in both directions and return fitted values plus residuals.
Parameters
----------
data_x : object
First variable sample used in directional functional-causal-model tests.
data_y : object
Second variable sample used in directional functional-causal-model tests.
Returns
-------
dict[str, np.ndarray]
Predicted values and residuals from fitting each direction of the additive-noise model.
"""
model = ANM()
x_col = as_column(data_x)
y_col = as_column(data_y)
y_hat = model.fit_gp(x_col, y_col)
x_hat = model.fit_gp(y_col, x_col)
return {
"y_hat_from_x": y_hat.ravel(),
"residual_y_from_x": (y_col - y_hat).ravel(),
"x_hat_from_y": x_hat.ravel(),
"residual_x_from_y": (x_col - x_hat).ravel(),
}
def direction_label(row):
"""
Create a compact text label for result tables.
Parameters
----------
row : object
Single DataFrame row or namedtuple being converted into a readable record.
Returns
-------
str
Readable direction label with the directional p-value margin.
"""
return f"{row['preferred_direction']} (margin={row['p_value_margin']:.3f})"
def edge_table_from_pairs(pair_table, run_label, min_margin=0.10):
"""
Convert pairwise ANM decisions into a graph edge table after a margin filter.
Parameters
----------
pair_table : pd.DataFrame
Table of `pair` values used by `edge_table_from_pairs`.
run_label : str
Readable label for `run`.
min_margin : object
Minimum directional p-value margin required to keep an edge candidate.
Returns
-------
pd.DataFrame
Edge table built from directional pairwise tests, including source, target, and p-value evidence.
"""
filtered = pair_table[pair_table["p_value_margin"] >= min_margin].copy()
edge_table = filtered.rename(columns={"preferred_source": "source", "preferred_target": "target"})
edge_table["run"] = run_label
edge_table["edge_type"] = "-->"
return edge_table[["run", "source", "edge_type", "target", "p_value_margin", "p_value_x_to_y", "p_value_y_to_x"]]
def directed_pairs(edge_df):
"""
Extract directed source-target pairs from a tidy edge table.
Parameters
----------
edge_df : pd.DataFrame
Learned or oracle edge table being summarized or drawn.
Returns
-------
set[tuple[str, str]]
Ordered source-target pairs representing directed learned edges.
"""
return set(zip(edge_df["source"], edge_df["target"])) if not edge_df.empty else set()
def summarize_pairwise_screen(edge_df, truth_df, label):
"""
Summarize pairwise screen output against direct truth edges.
Parameters
----------
edge_df : pd.DataFrame
Learned or oracle edge table being summarized or drawn.
truth_df : pd.DataFrame
Oracle edge table used as the synthetic ground-truth benchmark.
label : str
Readable label attached to the method, scenario, or output row.
Returns
-------
pd.DataFrame
Summary table for summarize pairwise screen, with one row per method, scenario, group, or diagnostic check.
"""
learned_directed = directed_pairs(edge_df)
true_directed = set(zip(truth_df["source"], truth_df["target"]))
true_skeleton = {frozenset(edge) for edge in true_directed}
learned_skeleton = {frozenset(edge) for edge in learned_directed}
correct_directed = learned_directed & true_directed
reversed_true = {(src, dst) for src, dst in true_directed if (dst, src) in learned_directed}
indirect_or_extra = learned_skeleton - true_skeleton
return pd.DataFrame(
[
{
"run": label,
"reported_edges": len(edge_df),
"true_direct_edges": len(true_directed),
"correct_direct_edges": len(correct_directed),
"reversed_true_edges": len(reversed_true),
"missing_true_adjacencies": len(true_skeleton - learned_skeleton),
"indirect_or_extra_adjacencies": len(indirect_or_extra),
}
]
)
GRAPH_POSITIONS = {
"need": (0.11, 0.72),
"intent": (0.11, 0.28),
"match": (0.39, 0.50),
"engagement": (0.62, 0.50),
"renewal": (0.89, 0.72),
"support": (0.89, 0.28),
}
NODE_LABELS = {name: name.title() for name in GRAPH_POSITIONS}
NODE_COLORS = {
"need": "#e0f2fe",
"intent": "#dbeafe",
"match": "#ecfccb",
"engagement": "#fef3c7",
"renewal": "#fee2e2",
"support": "#f3e8ff",
}
def trim_edge_to_box(start, end, box_w=0.145, box_h=0.095, gap=0.012):
"""
Return edge endpoints that stop just outside source and target boxes.
Parameters
----------
start : tuple[float, float]
Starting coordinate of the arrow or edge segment.
end : tuple[float, float]
Ending coordinate of the arrow or edge segment.
box_w : float
Width of the node box used to trim edge endpoints.
box_h : float
Height of the node box used to trim edge endpoints.
gap : float
Extra spacing between an edge endpoint and a node boundary.
Returns
-------
tuple[tuple[float, float], tuple[float, float]]
Start and end coordinates trimmed so the edge stops at box boundaries.
"""
x0, y0 = start
x1, y1 = end
dx = x1 - x0
dy = y1 - y0
length = float(np.hypot(dx, dy))
if length == 0:
return start, end
effective_w = box_w + 0.04
effective_h = box_h + 0.04
x_limit = (effective_w / 2) / abs(dx) if dx else np.inf
y_limit = (effective_h / 2) / abs(dy) if dy else np.inf
t = min(x_limit, y_limit) + gap / length
return (x0 + dx * t, y0 + dy * t), (x1 - dx * t, y1 - dy * t)
def draw_box_graph(edge_df, title, path, note=None):
"""
Draw a small DAG-style graph with rounded boxes and visible arrowheads.
Parameters
----------
edge_df : pd.DataFrame
Learned or oracle edge table being summarized or drawn.
title : str
Title shown above the plot.
path : str or pathlib.Path
Optional output path for the figure.
note : str or None
Optional explanatory note shown in the plot.
Returns
-------
None
Draws the graph diagram directly on the Matplotlib axes and saves it when requested.
"""
fig, ax = plt.subplots(figsize=(12, 6.2))
ax.set_axis_off()
ax.set_xlim(-0.02, 1.02)
ax.set_ylim(0.04, 0.96)
box_w, box_h = 0.145, 0.095
for row in edge_df.itertuples(index=False):
if row.source not in GRAPH_POSITIONS or row.target not in GRAPH_POSITIONS:
continue
start_raw = GRAPH_POSITIONS[row.source]
end_raw = GRAPH_POSITIONS[row.target]
start, end = trim_edge_to_box(start_raw, end_raw, box_w=box_w, box_h=box_h)
arrow = FancyArrowPatch(
start,
end,
arrowstyle="-|>",
mutation_scale=18,
linewidth=1.8,
color="#334155",
connectionstyle="arc3,rad=0.035",
zorder=2,
)
ax.add_patch(arrow)
for node, (x, y) in GRAPH_POSITIONS.items():
rect = FancyBboxPatch(
(x - box_w / 2, y - box_h / 2),
box_w,
box_h,
boxstyle="round,pad=0.018",
facecolor=NODE_COLORS[node],
edgecolor="#1f2937",
linewidth=1.1,
zorder=5,
)
ax.add_patch(rect)
ax.text(x, y, NODE_LABELS[node], ha="center", va="center", fontsize=10.5, fontweight="bold", zorder=6)
if note:
ax.text(0.50, 0.08, note, ha="center", va="center", fontsize=10, color="#475569")
ax.set_title(title, pad=18, fontsize=14, fontweight="bold")
fig.savefig(path, dpi=160, bbox_inches="tight")
plt.show()