Use refactored _print_name to display SymbolicRVs

Author: larryshamalama

pymc/model_graph.py

@@ -181,14 +181,11 @@ def _make_node(self, var_name, graph, *, nx=False, cluster=False, formatting: st
                 style = "filled"
             else:
                 style = None
-            symbol = v.owner.op.__class__.__name__
-            if symbol == "MarginalMixtureRV":
+            symbol = v.owner.op._print_name[0]
+            if symbol == "MarginalMixture":
                 components = v.owner.inputs[2:]
                 if len(components) == 2:
-                    component_names = [
-                        var.owner.op.__class__.__name__.replace("Unmeasurable", "")[:-2]
-                        for var in components
-                    ]
+                    component_names = [var.owner.op._print_name[0] for var in components]
                     if check_zip_graph_from_components(components):
                         # ZeroInflated distribution
                         component_names.remove("DiracDelta")
@@ -198,17 +195,21 @@ def _make_node(self, var_name, graph, *, nx=False, cluster=False, formatting: st
                         symbol = f"{'-'.join(component_names)}Mixture"
                 elif len(components) == 1:
                     # single component dispatch mixture
-                    symbol = f"{components[0].owner.op.__class__.__name__.replace('Unmeasurable', '')[:-2]}Mixture"
+                    symbol = f"{components[0].owner.op._print_name[0]}Mixture"
                 else:
                     symbol = symbol[:-2]  # just MarginalMixture
-            elif symbol == "CensoredRV":
+            elif symbol == "Censored":
                 censored_dist = v.owner.inputs[0]
-                symbol = symbol[:-2] + censored_dist.owner.op.__class__.__name__[:-2]
-            elif symbol == "TruncatedRV":
+                symbol = symbol + censored_dist.owner.op._print_name[0]
+            elif symbol == "Truncated":
                 truncated_dist = v.owner.op.base_rv_op
-                symbol = symbol[:-2] + truncated_dist.__class__.__name__[:-2]
-            elif symbol.endswith("RV"):
-                symbol = symbol[:-2]
+                symbol = symbol + truncated_dist._print_name[0]
+            elif symbol == "RandomWalk":
+                innovation_dist = v.owner.inputs[1].owner.op._print_name[0]
+                if innovation_dist == "Normal":
+                    symbol = "Gaussian" + symbol
+                else:
+                    symbol = innovation_dist + symbol
             label = f"{var_name}\n~\n{symbol}"
         else:
             shape = "box"

pymc/tests/test_model_graph.py

@@ -28,6 +28,7 @@
     GaussianRandomWalk,
     Mixture,
     Normal,
+    RandomWalk,
     StudentT,
     Truncated,
     ZeroInflatedPoisson,
@@ -399,8 +400,21 @@ def test_subgraph(self, var_names, vars_to_plot, compute_graph):
             },
             "MarginalMixture",
         ),
-        (GaussianRandomWalk, {"init_dist": Normal.dist(0.0, 5.0), "steps": 10}, "RandomWalk"),
+        (
+            GaussianRandomWalk,
+            {"init_dist": Normal.dist(0.0, 5.0), "steps": 10},
+            "GaussianRandomWalk",
+        ),
         (Truncated, {"dist": StudentT.dist(7), "upper": 3.0}, "TruncatedStudentT"),
+        (
+            RandomWalk,
+            {
+                "innovation_dist": pm.StudentT.dist(7),
+                "init_dist": pm.Normal.dist(0, 1),
+                "steps": 10,
+            },
+            "StudentTRandomWalk",
+        ),
     ],
 )
 def test_symbolic_distribution_display(symbolic_dist, dist_kwargs, display_name):