Download scripts/figure1.py from ChatterjeeLab/PIVOT: direct link, hf CLI and curl.
- Browser
- Download file 4.07 kB
-
https://huggingface.co/ChatterjeeLab/PIVOT/resolve/main/scripts/figure1.py
- Command line
-
hf download hf://ChatterjeeLab/PIVOT/scripts/figure1.py
-
curl -L -o figure1.py https://huggingface.co/ChatterjeeLab/PIVOT/resolve/main/scripts/figure1.py
4.07 kB
| """Draw the PIVOT overview with editable vector paths and embedded fonts.""" | |
| from pathlib import Path | |
| import numpy as np | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| from matplotlib import font_manager as fm | |
| from matplotlib.patches import Ellipse, FancyArrowPatch, PathPatch, FancyBboxPatch | |
| from matplotlib.path import Path as MPath | |
| ROOT = Path(__file__).resolve().parents[1] | |
| for f in (ROOT / "assets/fonts").glob("*.ttf"): | |
| fm.fontManager.addfont(str(f)) | |
| plt.rcParams.update( | |
| { | |
| "font.family": "Ubuntu", | |
| "mathtext.fontset": "cm", | |
| "pdf.fonttype": 42, | |
| "ps.fonttype": 42, | |
| "font.size": 13, | |
| } | |
| ) | |
| blue = "#62AEDD" | |
| navy = "#245E84" | |
| red = "#9D2944" | |
| ink = "#243746" | |
| pale = "#E9F4FB" | |
| gray = "#9EB0BD" | |
| fig, ax = plt.subplots(figsize=(13.6, 4.8)) | |
| ax.set(xlim=(0, 13.6), ylim=(0, 4.8)) | |
| ax.axis("off") | |
| def text(x, y, s, size=13, color=ink, **kw): | |
| ax.text( | |
| x, | |
| y, | |
| s, | |
| fontsize=size, | |
| color=color, | |
| ha=kw.pop("ha", "center"), | |
| va="center", | |
| **kw, | |
| ) | |
| def arrow(start, end, color=blue, rad=0, lw=2.2): | |
| ax.add_patch( | |
| FancyArrowPatch( | |
| start, | |
| end, | |
| arrowstyle="-|>", | |
| mutation_scale=16, | |
| connectionstyle=f"arc3,rad={rad}", | |
| color=color, | |
| lw=lw, | |
| ) | |
| ) | |
| def cloud(x, y, w, h, color, seed): | |
| rng = np.random.default_rng(seed) | |
| ax.add_patch(Ellipse((x, y), w, h, facecolor=color, alpha=0.10, edgecolor="none")) | |
| z = rng.normal(size=(36, 2)) | |
| z = z[np.linalg.norm(z, axis=1) < 2] | |
| ax.scatter( | |
| x + z[:, 0] * w / 5, | |
| y + z[:, 1] * h / 5, | |
| s=rng.uniform(9, 24, len(z)), | |
| color=color, | |
| alpha=0.7, | |
| linewidth=0, | |
| ) | |
| cloud(1.45, 3.08, 2.1, 1.72, blue, 2) | |
| text(1.45, 4.19, "Control population", 18, navy, weight="bold") | |
| text(1.45, 2.02, r"$c_0\sim\rho_0$", 18) | |
| # One broad map across the figure, with faint parallel sample trajectories. | |
| for off in [-0.32, -0.12, 0.12, 0.32]: | |
| verts = [(2.55, 3.1 + off), (4.3, 3.8 + off), (6.7, 3.9 + off), (8.75, 3.15 + off)] | |
| ax.add_patch( | |
| PathPatch( | |
| MPath(verts, [MPath.MOVETO, MPath.CURVE4, MPath.CURVE4, MPath.CURVE4]), | |
| fill=False, | |
| edgecolor=blue, | |
| alpha=0.20, | |
| lw=10, | |
| ) | |
| ) | |
| arrow((2.6, 3.1), (8.8, 3.14), blue, 0.20, 3) | |
| text(5.6, 3.72, r"$X_\theta(0,1,c_0,e_u)$", 23) | |
| text(5.6, 3.17, "Predict the response", 18, navy, weight="bold") | |
| text(5.6, 4.38, r"$u=\{(g_j,o_j)\}_{j=1}^{M}\quad\longmapsto\quad e_u$", 17) | |
| arrow((5.6, 4.12), (5.6, 3.92), gray, 0, 1.2) | |
| cloud(9.4, 3.18, 1.75, 1.5, blue, 9) | |
| text(9.4, 4.19, "Predicted cells", 18, navy, weight="bold") | |
| cloud(11.97, 3.7, 1.65, 1.18, red, 3) | |
| text(11.97, 4.56, "Target cells", 18, red, weight="bold") | |
| text(12, 2.88, r"$c^\star\sim\rho^\star$", 18) | |
| arrow((10.2, 3.23), (11.22, 3.63), red, -0.08, 1.6) | |
| text(11, 2.57, "Endpoint reward", 17, red) | |
| # Reward gradients return to the intervention coordinates. | |
| arrow((11.3, 2.31), (5.22, 1.89), red, -0.17, 2.6) | |
| text(7.95, 1.36, "Optimize intervention embeddings", 17, red) | |
| text(3.05, 1.72, "Rank admissible interventions", 17, navy, weight="bold") | |
| arrow((5.01, 2.03), (3.13, 2.25), navy, -0.10, 1.5) | |
| text(3.45, 2.48, r"$e^{(L)}\;\longrightarrow\;u_{1:K}$", 18) | |
| ax.add_patch( | |
| FancyBboxPatch( | |
| (0.18, 0.12), | |
| 13.22, | |
| 1.03, | |
| boxstyle="round,pad=0.02,rounding_size=.08", | |
| facecolor=pale, | |
| edgecolor="none", | |
| ) | |
| ) | |
| text(0.51, 0.78, "PIVOT", 18, navy, ha="left", weight="bold") | |
| text( | |
| 7.8, | |
| 0.74, | |
| r"$\hat c_1=X_\theta(0,1,c_0,e)\qquad g_e=J_eX_\theta^{\mathsf{T}}\nabla_{\hat c_1}r\qquad e^+=e+\gamma\,\dfrac{g_e}{\|g_e\|_2+\epsilon}$", | |
| 20, | |
| ) | |
| text( | |
| 7.6, | |
| 0.29, | |
| "Endpoint prediction, reward gradients, and admissible interventions.", | |
| 16, | |
| ) | |
| fig.subplots_adjust(left=0, right=1, bottom=0, top=1) | |
| for ext in ["pdf", "png", "svg"]: | |
| fig.savefig(ROOT / "assets" / f"figure1.{ext}", dpi=250, facecolor="white") | |