Model to mermaid diagram with model_to_mermaid

[williambdean]

Description

Mermaid diagrams are in various places including built into GitHub! This allows users to visualize models with mermaid

Related Issue

• Closes #
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pymc--7826.org.readthedocs.build/en/7826/

[williambdean]

import pymc as pm
import numpy as np

seed = sum(map(ord, "Mermaid diagrams"))
rng = np.random.default_rng(seed)

X = rng.normal(size=(100, 3))

true_betas = np.array([0.5, -0.2, 0.1])
true_sigma = 1.0
y = X @ true_betas + rng.normal(0, true_sigma, size=X.shape[0])


coords = {"idx": range(len(X)), "features": ["feature1", "feature2", "feature3"]}
with pm.Model(coords=coords) as model:
    X_ = pm.Data("X", X, dims=("idx", "features"))
    y_ = pm.Data("y", y, dims="idx")
    beta = pm.Normal("beta", mu=0, sigma=1, dims="features")

    mu = pm.Deterministic("mu", pm.math.dot(X_, beta), dims="idx")
    sigma = pm.InverseGamma("sigma", alpha=1, beta=1)

    pm.Normal("obs", mu=mu, sigma=sigma, observed=y_, dims="idx")
    # Add a potential to penalize high values of beta
    pm.Potential("beta_penalty", -0.5 * pm.math.sum(beta**2))


print(pm.model_to_mermaid(model, include_dim_lengths=False))

graph TD
%% Nodes:
X[X ~ Data]
X@{ shape: db }
beta([beta ~ Normal])
beta@{ shape: rounded }
beta_penalty([beta_penalty ~ Potential])
beta_penalty@{ shape: diam }
style beta_penalty fill:#f0f0f0
mu([mu ~ Deterministic])
mu@{ shape: rect }
obs([obs ~ Normal])
obs@{ shape: rounded }
style obs fill:#757575
sigma([sigma ~ InvGamma])
sigma@{ shape: rounded }
y[y ~ Data]
y@{ shape: db }

%% Edges:
beta --> beta_penalty
X --> mu
beta --> mu
mu --> obs
sigma --> obs
obs --> y

%% Plates:
subgraph "features"
    beta
end
subgraph "idx"
    obs
    y
    mu
end
subgraph "idx x features"
    X
end

Loading

[codecov]

Codecov Report

Attention: Patch coverage is 73.58491% with 14 lines in your changes missing coverage. Please review.

Project coverage is 92.92%. Comparing base (cd2e1a3) to head (a7dd63a).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc/model_graph.py	73.07%	14 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7826      +/-   ##
==========================================
- Coverage   92.97%   92.92%   -0.06%     
==========================================
  Files         107      107              
  Lines       18228    18280      +52     
==========================================
+ Hits        16948    16986      +38     
- Misses       1280     1294      +14     

Files with missing lines	Coverage Δ
pymc/__init__.py	100.00% <100.00%> (ø)
pymc/model_graph.py	84.88% <73.07%> (-2.38%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ricardoV94]

Neat stuff!

[ricardoV94]

@williambdean how crazy would it be to add this for pytensor graphs? Would be easier to setup than d3viz stuff

[williambdean]

@williambdean how crazy would it be to add this for pytensor graphs? Would be easier to setup than d3viz stuff

Not crazy. Syntax is pretty easy. Let's discuss over there

[williambdean]

This is an another example model

import pymc as pm
import numpy as np

seed = sum(map(ord, "Something completely different"))
rng = np.random.default_rng(seed)

coords = {"idx": range(100)}
with pm.Model(coords=coords) as model:
    mu = pm.Normal("mu", mu=pm.Normal("mu_mu"), sigma=pm.HalfNormal("mu_sigma"), dims="idx")
    sigma = pm.InverseGamma(
        "sigma", 
        alpha=pm.HalfNormal("sigma_alpha"),
        beta=pm.HalfNormal("sigma_beta"),
        dims="idx",
    )

    pm.Normal("obs", mu=mu, sigma=sigma, dims="idx", observed=rng.normal(size=100))

print(pm.model_to_mermaid(model))

graph TD
%% Nodes:
mu([mu ~ Normal])
mu@{ shape: rounded }
mu_mu([mu_mu ~ Normal])
mu_mu@{ shape: rounded }
mu_sigma([mu_sigma ~ HalfNormal])
mu_sigma@{ shape: rounded }
obs([obs ~ Normal])
obs@{ shape: rounded }
style obs fill:#757575
sigma([sigma ~ InvGamma])
sigma@{ shape: rounded }
sigma_alpha([sigma_alpha ~ HalfNormal])
sigma_alpha@{ shape: rounded }
sigma_beta([sigma_beta ~ HalfNormal])
sigma_beta@{ shape: rounded }

%% Edges:
mu_mu --> mu
mu_sigma --> mu
mu --> obs
sigma --> obs
sigma_alpha --> sigma
sigma_beta --> sigma

%% Plates:
subgraph "idx (100)"
    obs
    sigma
    mu
end

Loading