Math4AI

Publication-ready figures for mathematics and machine learning — one import away.

Math4AI is the companion library for the Mathematics for AI book and any technical project that generates chapter figures with Python. It gives you a shared visual language: consistent colors, matrix and vector layouts, reusable icons, and mindmaps driven by JSON — so every diagram in a long manuscript looks like it belongs to the same book.

Works on Linux x86_64, Python 3.9+ (3.12 recommended). Dependencies: matplotlib, numpy.

---

Why Math4AI?

Writing a math or ML book means hundreds of small plotting scripts — one per figure. Without a shared toolkit, colors drift, brace styles diverge, and mindmaps become copy-paste nightmares.

Math4AI fixes that:

• One palette across linear algebra, optimization, probability, and information theory chapters
• Matrix & vector helpers — dashed rows/cols, braces, piercing highlights, affine grids, arrows
• Dozens of built-in icons for mindmaps (SVD, Bellman, KL, Langevin, EM, …)
• JSON mindmaps — declare nodes and edges; render a publication-quality figure from mindmap.json
• Local icon hooks — add a custom shape in one script without forking the library
• save_figure_to_script_dir — output lands next to your script, ready for \includegraphics

Built for authors and educators who care how their PDFs look, not just that the code runs.

---

Get started

From PyPI (release wheel, Linux x86_64):

pip install math4ai

From source (editable dev install):

git clone https://github.com/xuanxinjushi/math4ai.git math4ai-plot
cd math4ai-plot
pip install -e ".[dev]"

Run tests:

pytest
# or: make test

Quick check:

python -c "from math4ai import plane_color; print(plane_color)"   # blue

Docker

Build image:

docker build -t math4ai:dev .

Run tests in container:

docker run --rm math4ai:dev

Generate docs icon gallery from container into local repo:

docker run --rm -v "$PWD":/app -w /app math4ai:dev make docs-icons

Open generated gallery at docs/icons.md.

---

What you can build

Area	Examples
Colors	plane_color, point_x_color, normal_vector_color, …
Matrices	Matrix, dashed rows/cols, braces, row/column coloring
Vectors	draw_vec, draw_affine_grid, arrow
Mindmaps	render_mindmap(__file__) + mindmap.json
Icons	QR, LU, SVD, gradient, Bayes, Markov, score field, …
Figures	save_figure_to_script_dir, 2D/3D grid styling, math fonts

---

A minute of code

Consistent scatter plot

from math4ai import point_x_color, save_figure_to_script_dir
import matplotlib.pyplot as plt

fig, ax = plt.subplots()
ax.scatter([1, 2, 3], [1, 4, 9], color=point_x_color)
save_figure_to_script_dir("example.png")

Mindmap from JSON

from math4ai import render_mindmap

if __name__ == "__main__":
    render_mindmap(__file__)   # reads mindmap.json beside this script

Custom mindmap icon (local hook)

def draw_my_topic_icon(ax, center_x, center_y, size=0.5):
    ...

Reference "icon": "my_topic" in mindmap.json — no library patch required.

---

Who is it for?

• Authors of textbooks and monographs with matplotlib figure pipelines
• Course materials that need repeatable, professional diagrams
• Researchers publishing papers who want book-grade consistency across chapters

Install once, import everywhere your scripts live.