fixed errors in graphing

src/hssm/hssm.py

@@ -668,7 +668,7 @@ def graph(self, formatting="plain", name=None, figsize=None, dpi=300, fmt="png")
 
         graphviz = HSSMModelGraph(
             model=self.pymc_model, parent=self._parent_param
-        ).make_graph(formatting=formatting)
+        ).make_graph(formatting=formatting, response_str=self.response_str)
 
         width, height = (None, None) if figsize is None else figsize
 

src/hssm/utils.py

@@ -172,7 +172,10 @@ def __init__(self, model, parent):
         super().__init__(model)
 
     def make_graph(
-        self, var_names: Iterable[VarName] | None = None, formatting: str = "plain"
+        self,
+        var_names: Iterable[VarName] | None = None,
+        formatting: str = "plain",
+        response_str: str = "rt,response",
     ):
         """Make graphviz Digraph of PyMC model.
 
@@ -224,8 +227,8 @@ def make_graph(
                     label=f"{self.parent.name}\n~\nDeterministic",
                     shape="box",
                 )
-                shape = fast_eval(self.model[self.response_str].shape)
-                plate_label = f"{self.response_str}_obs({shape[0]})"
+                shape = fast_eval(self.model[response_str].shape)
+                plate_label = f"{response_str}_obs({shape[0]})"
 
                 sub.attr(
                     label=plate_label,
@@ -249,7 +252,7 @@ def make_graph(
                     graph.edge(parent.replace(":", "&"), child.replace(":", "&"))
 
         if self.parent.is_regression:
-            graph.edge(self.parent.name, self.response_str)
+            graph.edge(self.parent.name, response_str)
 
         return graph