Repository: pymc-devs/pymc
Pull request: BUG: Mermaid nodes with spaces in name
<!-- !! Thank your for opening a PR !! -->

<!--- Provide a self-contained summary of your changes in the Title above -->
<!--- This is what will be shown in the automatic release notes: https://github.com/pymc-devs/pymc/releases -->

## Description
<!--- Describe your changes in detail -->

Example: plant growth will become plant_growth for the node name but will still be displayed as "plant growth"

## Related Issue
<!--- It is good practice to first open an issue explaining the bug / new feature that is addressed by this PR -->
<!--- Please type an `x` in one of the boxes below and provide the issue number after the # sign: -->

- [ ] Closes #7829
- [ ] Related to #

## Checklist
<!--- Make sure you have completed the following steps before submitting your PR -->
<!--- Feel free to type an `x` in all the boxes below to let us know you have completed the steps: -->
- [ ] Checked that [the pre-commit linting/style checks pass](https://docs.pymc.io/en/latest/contributing/python_style.html)
- [ ] 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](https://wiki.openstack.org/wiki/GitCommitMessages#Structural_split_of_changes)
<!--- You may find this guide helpful: https://mainmatter.com/blog/2021/05/26/keeping-a-clean-git-history/ -->

## Type of change
<!--- Select one of the categories below by typing an `x` in the box -->
- [ ] New feature / enhancement
- [ ] Bug fix
- [ ] Documentation
- [ ] Maintenance
- [ ] Other (please specify):
<!--- Additionally, if you are a maintainer or reviewer, please make sure that the appropriate labels are added to this PR -->


<!-- readthedocs-preview pymc start -->
----
📚 Documentation preview 📚: https://pymc--7835.org.readthedocs.build/en/7835/

<!-- readthedocs-preview pymc end -->

Problem statement:
BUG: `pm.model_to_mermaid` with spaces in variable names

From example in the README

```python
import pymc as pm
# Taking draws from a normal distribution
seed = 42
x_dist = pm.Normal.dist(shape=(100, 3))
x_data = pm.draw(x_dist, random_seed=seed)

# Define coordinate values for all dimensions of the data
coords={
 "trial": range(100),
 "features": ["sunlight hours", "water amount", "soil nitrogen"],
}

# Define generative model
with pm.Model(coords=coords) as generative_model:
   x = pm.Data("x", x_data, dims=["trial", "features"])

   # Model parameters
   betas = pm.Normal("betas", dims="features")
   sigma = pm.HalfNormal("sigma")

   # Linear model
   mu = x @ betas

   # Likelihood
   # Assuming we measure deviation of each plant from baseline
   plant_growth = pm.Normal("plant growth", mu, sigma, dims="trial")

print(pm.model_to_mermaid(generative_model))
```

Hints from issue discussion:

Commit message: add test with space in variable name

Patch to review:
diff --git a/pymc/model_graph.py b/pymc/model_graph.py
index 9ec8a502d..291e14582 100644
--- a/pymc/model_graph.py
+++ b/pymc/model_graph.py
@@ -795,36 +795,42 @@ def model_to_graphviz(
     )
 
 
+def _create_mermaid_node_name(name: str) -> str:
+    return name.replace(":", "_").replace(" ", "_")
+
+
 def _build_mermaid_node(node: NodeInfo) -> list[str]:
     var = node.var
     node_type = node.node_type
+    name = var.name
+    node_name = _create_mermaid_node_name(name)
     if node_type == NodeType.DATA:
         return [
-            f"{var.name}[{var.name} ~ Data]",
-            f"{var.name}@{{ shape: db }}",
+            f"{node_name}[{var.name} ~ Data]",
+            f"{node_name}@{{ shape: db }}",
         ]
     elif node_type == NodeType.OBSERVED_RV:
         return [
-            f"{var.name}([{var.name} ~ {random_variable_symbol(var)}])",
-            f"{var.name}@{{ shape: rounded }}",
-            f"style {var.name} fill:#757575",
+            f"{node_name}([{name} ~ {random_variable_symbol(var)}])",
+            f"{node_name}@{{ shape: rounded }}",
+            f"style {node_name} fill:#757575",
         ]
 
     elif node_type == NodeType.FREE_RV:
         return [
-            f"{var.name}([{var.name} ~ {random_variable_symbol(var)}])",
-            f"{var.name}@{{ shape: rounded }}",
+            f"{node_name}([{name} ~ {random_variable_symbol(var)}])",
+            f"{node_name}@{{ shape: rounded }}",
         ]
     elif node_type == NodeType.DETERMINISTIC:
         return [
-            f"{var.name}([{var.name} ~ Deterministic])",
-            f"{var.name}@{{ shape: rect }}",
+            f"{node_name}([{name} ~ Deterministic])",
+            f"{node_name}@{{ shape: rect }}",
         ]
     elif node_type == NodeType.POTENTIAL:
         return [
-            f"{var.name}([{var.name} ~ Potential])",
-            f"{var.name}@{{ shape: diam }}",
-            f"style {var.name} fill:#f0f0f0",
+            f"{node_name}([{name} ~ Potential])",
+            f"{node_name}@{{ shape: diam }}",
+            f"style {node_name} fill:#f0f0f0",
         ]
 
     return []
@@ -842,8 +848,8 @@ def _build_mermaid_edges(edges) -> list[str]:
     """Return a list of Mermaid edge definitions."""
     edge_lines = []
     for child, parent in edges:
-        child_id = str(child).replace(":", "_")
-        parent_id = str(parent).replace(":", "_")
+        child_id = _create_mermaid_node_name(child)
+        parent_id = _create_mermaid_node_name(parent)
         edge_lines.append(f"{parent_id} --> {child_id}")
     return edge_lines
 
diff --git a/tests/test_model_graph.py b/tests/test_model_graph.py
index 35a202215..6ffbce5b3 100644
--- a/tests/test_model_graph.py
+++ b/tests/test_model_graph.py
@@ -536,6 +536,14 @@ def simple_model() -> pm.Model:
     return model
 
 
+@pytest.fixture(scope="module")
+def variable_with_space():
+    with pm.Model() as model:
+        pm.Normal("plant growth")
+
+    return model
+
+
 def test_unknown_node_type(simple_model):
     with pytest.raises(ValueError, match="Node formatters must be of type NodeType."):
         model_to_graphviz(simple_model, node_formatters={"Unknown Node Type": "dummy"})
@@ -652,3 +660,17 @@ def test_model_to_mermaid(simple_model):
     %% Plates:
     """)
     assert model_to_mermaid(simple_model) == expected_mermaid_string.strip()
+
+
+def test_model_to_mermaid_with_variable_with_space(variable_with_space):
+    expected_mermaid_string = dedent("""
+    graph TD
+    %% Nodes:
+    plant_growth([plant growth ~ Normal])
+    plant_growth@{ shape: rounded }
+
+    %% Edges:
+
+    %% Plates:
+    """)
+    assert model_to_mermaid(variable_with_space) == expected_mermaid_string.strip()


Code review comments:
diff_hunk: @@ -536,6 +536,14 @@ def simple_model() -> pm.Model:
     return model
 
 
+@pytest.fixture(scope="module")
path: tests/test_model_graph.py
text: @user1:
What's the point of this fixture? Just put it in the test?

@author:
Removed. Thought I was going to do something else.