# Define reusable helpers for the Render The Base DAG section.
base_labels = {
"need": "Need\nscore",
"intent": "Intent\nsignal",
"match": "Match\nquality",
"engagement": "Engagement",
"renewal": "Renewal\nvalue",
"support": "Support\nload",
}
base_positions = {
"need": (0.10, 0.76),
"intent": (0.10, 0.24),
"match": (0.34, 0.52),
"engagement": (0.66, 0.52),
"renewal": (0.90, 0.72),
"support": (0.90, 0.30),
}
base_node_colors = {
"need": "#eef2ff",
"intent": "#eef2ff",
"match": "#e0f2fe",
"engagement": "#e0f2fe",
"renewal": "#dcfce7",
"support": "#dcfce7",
}
base_edge_radii = {
("need", "match"): -0.04,
("intent", "match"): 0.04,
("match", "engagement"): 0.00,
("intent", "renewal"): 0.18,
("engagement", "renewal"): -0.04,
("engagement", "support"): 0.04,
}
def draw_teaching_style_dag(edge_table, labels, positions, node_colors, title, path, edge_radii=None):
"""
Draw a small DAG using the shared tutorial visual style.
Parameters
----------
edge_table : pd.DataFrame
Table of source-target edges and endpoint marks to convert, draw, or score.
labels : dict[str, str]
Display labels used for graph nodes in the rendered figure.
positions : dict
Node positions used to draw the graph.
node_colors : dict[str, str] or None
Optional mapping from node names to display colors.
title : str
Title shown above the plot.
path : str or pathlib.Path
Optional output path for the figure.
edge_radii : dict or None
Optional curvature settings for individual plotted edges.
Returns
-------
pathlib.Path
Path to the saved teaching DAG figure.
"""
edge_radii = edge_radii or {}
fig, ax = plt.subplots(figsize=(12, 6))
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_axis_off()
for row in edge_table.itertuples(index=False):
source = row.source
target = row.target
edge_type = getattr(row, "edge_type", "directed")
ax.annotate(
"",
xy=positions[target],
xytext=positions[source],
arrowprops=dict(
arrowstyle="-|>",
color="#334155",
linewidth=1.5,
mutation_scale=18,
shrinkA=34,
shrinkB=46,
linestyle="--" if edge_type == "latent" else "-",
connectionstyle=f"arc3,rad={edge_radii.get((source, target), 0.0)}",
),
zorder=1,
)
for node, (x, y) in positions.items():
ax.text(
x,
y,
labels[node],
ha="center",
va="center",
fontsize=10.5,
fontweight="bold",
bbox=dict(
boxstyle="round,pad=0.45",
facecolor=node_colors.get(node, "#e0f2fe"),
edgecolor="#334155",
linewidth=1.2,
),
zorder=2,
)
ax.set_title(title, pad=18)
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(path, dpi=160, bbox_inches="tight")
plt.show()
return path
base_dag_path = FIGURE_DIR / f"{NOTEBOOK_PREFIX}_base_teaching_dag.png"
draw_teaching_style_dag(
base_edge_table,
base_labels,
base_positions,
base_node_colors,
"Base Teaching DAG",
base_dag_path,
edge_radii=base_edge_radii,
)