Incorrect model graph

[Vutya]

Describe the bug I tried to instantiate the model from one of the HSSM tutorials, with drift rate v as a regression of other variables (x and y). Then I tried to plot the model (model.graph()), but for some reason it shows _vx and _vy as separate values, not as the parents of v. When I print the list of the parameters, it shows the correct formula. When I add the second parameter as a regression (e.g., decision boundary), it shows correct links between parameters for decision boundary, but not for drift rate. If I set only the decision boundary as the regression, then it is shown in a wrong way.

For now it looks like the problem might be in the __findparent() methods since it always sets the first parameter-regression in the _self._listparams, and it is always the drift rate. But I am not sure I understand 100% the meaning of this method and how I could modify it.

Versions HSSM==0.2.1, pymc==5.14.0, bambi==0.13.0, Python==3.11.9

To Reproduce

import hssm
import numpy as np
import pandas as pd
from ssms.basic_simulators.simulator import simulator

# Get mode simulations
# Set up trial by trial parameters
intercept = 0.3
x = np.random.uniform(-1, 1, size=1000)
y = np.random.uniform(-1, 1, size=1000)
v = intercept + (0.8 * x) + (0.3 * y)

true_values = np.column_stack(
    [v, np.repeat([[1.5, 0.5, 0.5]], axis=0, repeats=1000)]
)

obs_ddm_reg_v = simulator(true_values, model="ddm", n_samples=1)

dataset_reg_v = pd.DataFrame(
    {
        "rt": obs_ddm_reg_v["rts"].flatten(),
        "response": obs_ddm_reg_v["choices"].flatten(),
        "x": x,
        "y": y,
    }
)

model = hssm.HSSM(
    data=dataset_reg_v,
    model="ddm",
    include=[
        {
            "name": "v",
            "prior": {
                "Intercept": {
                    "name": "Uniform",
                    "lower": -3.0,
                    "upper": 3.0,
                    # "initval": 0   # optional --> set the initial value of the parameter (to e.g. avoid boundary violations at the intial sampling step)
                },
                "x": {
                    "name": "Uniform",
                    "lower": -1.0,
                    "upper": 1.0,
                },
                "y": {"name": "Uniform", "lower": -1.0, "upper": 1.0},
            },
            "formula": "v ~ 1 + x + y",
            "link": "identity",
        }
    ],
)

model.graph()

# Instantiate our hssm model
model1 = hssm.HSSM(
    data=dataset_reg_v,
    model="ddm",
    include=[
        {
            "name": "a",
            "prior": {
                "Intercept": {"name": "Uniform", "lower": 0.5, "upper": 3.0},
                "x": {"name": "Uniform", "lower": -1.0, "upper": 1.0},
                "y": {"name": "Uniform", "lower": -1.0, "upper": 1.0},
            },
            "formula": "a ~ 1 + x + y",
        },
        {
            "name": "v",
            "prior": {
                "Intercept": {"name": "Uniform", "lower": -3.0, "upper": 3.0},
                "x": {"name": "Uniform", "lower": -1.0, "upper": 1.0},
                "y": {"name": "Uniform", "lower": -1.0, "upper": 1.0},
            },
            "formula": "v ~ 1 + x + y",
        },
    ],
)

model1.graph()

Screenshots

model:

model1: