Using the do-operator to simulate data¶

This tutorial illustrates how we can use the underlying PyMC model of a given HSSM model, to forward simulate datasets that you may use for numerical studies.

In [1]:

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import hssm
import pymc as pm

import arviz as az import matplotlib.pyplot as plt import numpy as np import pandas as pd import hssm import pymc as pm

Simulate Data¶

We will simulate a simple dataset that contains two conditions.

In [2]:

# Simple model
data = hssm.simulate_data(
    model="ddm", theta=dict(v=0.5, a=1.5, z=0.5, t=0.1), size=500
)

# Simple model data = hssm.simulate_data( model="ddm", theta=dict(v=0.5, a=1.5, z=0.5, t=0.1), size=500 )

Specify HSSM Model¶

We will fit two separate models to the data.

In [3]:

# Model 1
hssm_model = hssm.HSSM(model="ddm", data=data)

# Model 1 hssm_model = hssm.HSSM(model="ddm", data=data)

Model initialized successfully.

Underlying PyMC model¶

In [4]:

pm.model_to_graphviz(hssm_model.pymc_model)

pm.model_to_graphviz(hssm_model.pymc_model)

Out[4]:

Using the do-operator¶

In [5]:

from pymc.model.transform.conditioning import do

synth_idata, synth_model = hssm_model.sample_do(params = {"v": 0,
                                                          "a": 1.5,
                                                          "z": 0.5,
                                                          "t": 0.1},
                                                draws = 100,
                                                return_model = True)

from pymc.model.transform.conditioning import do synth_idata, synth_model = hssm_model.sample_do(params = {"v": 0, "a": 1.5, "z": 0.5, "t": 0.1}, draws = 100, return_model = True)

Sampling: [rt,response]

Note that the process defines an new PyMC model, which sets the parameters we passed in our dictionary to Data objects.

In [6]:

pm.model_to_graphviz(synth_model)

pm.model_to_graphviz(synth_model)

Out[6]:

In [7]:

synth_idata

synth_idata

Out[7]:

arviz.InferenceData

• prior

  <xarray.Dataset> Size: 405kB
  Dimensions:  (chain: 1, draw: 100, __obs__: 500)
  Coordinates:
    * chain    (chain) int64 8B 0
    * draw     (draw) int64 800B 0 1 2 3 4 5 6 7 8 ... 91 92 93 94 95 96 97 98 99
    * __obs__  (__obs__) int64 4kB 0 1 2 3 4 5 6 7 ... 493 494 495 496 497 498 499
  Data variables:
      v_mean   (chain, draw, __obs__) float64 400kB 0.0 0.0 0.0 ... 0.0 0.0 0.0
  Attributes:
      created_at:                 2025-10-17T01:18:54.705842+00:00
      arviz_version:              0.22.0
      inference_library:          pymc
      inference_library_version:  5.25.1

  xarray.Dataset

  • Dimensions:
    • chain: 1
    • draw: 100
    • __obs__: 500
  • Coordinates: (3)
    • chain
      (chain)
      int64
      0

      array([0])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 6 ... 94 95 96 97 98 99

      array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
             18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
             36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
             54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
             72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89,
             90, 91, 92, 93, 94, 95, 96, 97, 98, 99])

    • __obs__
      (__obs__)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499], shape=(500,))

  • Data variables: (1)
    • v_mean
      (chain, draw, __obs__)
      float64
      0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0

      array([[[0., 0., 0., ..., 0., 0., 0.],
              [0., 0., 0., ..., 0., 0., 0.],
              [0., 0., 0., ..., 0., 0., 0.],
              ...,
              [0., 0., 0., ..., 0., 0., 0.],
              [0., 0., 0., ..., 0., 0., 0.],
              [0., 0., 0., ..., 0., 0., 0.]]], shape=(1, 100, 500))

  • Indexes: (3)
    • chain
      PandasIndex

      PandasIndex(Index([0], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
             18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
             36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
             54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
             72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89,
             90, 91, 92, 93, 94, 95, 96, 97, 98, 99],
            dtype='int64', name='draw'))

    • __obs__
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
            dtype='int64', name='__obs__', length=500))

  • Attributes: (4)

    created_at :

    2025-10-17T01:18:54.705842+00:00

    arviz_version :

    0.22.0

    inference_library :

    pymc

    inference_library_version :

    5.25.1

  
  
• prior_predictive

  <xarray.Dataset> Size: 805kB
  Dimensions:          (chain: 1, draw: 100, __obs__: 500, rt,response_dim: 2)
  Coordinates:
    * chain            (chain) int64 8B 0
    * draw             (draw) int64 800B 0 1 2 3 4 5 6 7 ... 93 94 95 96 97 98 99
    * __obs__          (__obs__) int64 4kB 0 1 2 3 4 5 ... 494 495 496 497 498 499
    * rt,response_dim  (rt,response_dim) int64 16B 0 1
  Data variables:
      rt,response      (chain, draw, __obs__, rt,response_dim) float64 800kB 1....
  Attributes:
      created_at:                 2025-10-17T01:18:54.708836+00:00
      arviz_version:              0.22.0
      inference_library:          pymc
      inference_library_version:  5.25.1

  xarray.Dataset

  • Dimensions:
    • chain: 1
    • draw: 100
    • __obs__: 500
    • rt,response_dim: 2
  • Coordinates: (4)
    • chain
      (chain)
      int64
      0

      array([0])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 6 ... 94 95 96 97 98 99

      array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
             18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
             36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
             54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
             72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89,
             90, 91, 92, 93, 94, 95, 96, 97, 98, 99])

    • __obs__
      (__obs__)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499], shape=(500,))

    • rt,response_dim
      (rt,response_dim)
      int64
      0 1

      array([0, 1])

  • Data variables: (1)
    • rt,response
      (chain, draw, __obs__, rt,response_dim)
      float64
      1.447 1.0 1.199 ... 1.0 1.531 1.0

      array([[[[ 1.44655442,  1.        ],
               [ 1.19892251,  1.        ],
               [ 0.89263052, -1.        ],
               ...,
               [ 0.71472335, -1.        ],
               [ 0.73178536, -1.        ],
               [ 2.36448574,  1.        ]],
      
              [[ 2.13502669,  1.        ],
               [ 2.94716907, -1.        ],
               [ 3.47942352,  1.        ],
               ...,
               [ 3.59060931,  1.        ],
               [ 3.22709524, -1.        ],
               [ 1.05840957, -1.        ]],
      
              [[ 3.17093587, -1.        ],
               [ 2.39028549,  1.        ],
               [ 2.4552834 , -1.        ],
               ...,
      ...
               ...,
               [ 2.50948834, -1.        ],
               [ 4.12036276,  1.        ],
               [ 1.79317844,  1.        ]],
      
              [[ 2.14002967, -1.        ],
               [ 2.82443666,  1.        ],
               [ 1.58908927, -1.        ],
               ...,
               [ 0.53689468,  1.        ],
               [ 1.74238849,  1.        ],
               [ 6.57336617, -1.        ]],
      
              [[ 0.77919495, -1.        ],
               [ 2.25651217,  1.        ],
               [ 2.38643312,  1.        ],
               ...,
               [ 1.18781722, -1.        ],
               [ 0.32062083,  1.        ],
               [ 1.53068888,  1.        ]]]], shape=(1, 100, 500, 2))

  • Indexes: (4)
    • chain
      PandasIndex

      PandasIndex(Index([0], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
             18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
             36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
             54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
             72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89,
             90, 91, 92, 93, 94, 95, 96, 97, 98, 99],
            dtype='int64', name='draw'))

    • __obs__
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
            dtype='int64', name='__obs__', length=500))

    • rt,response_dim
      PandasIndex

      PandasIndex(Index([0, 1], dtype='int64', name='rt,response_dim'))

  • Attributes: (4)

    created_at :

    2025-10-17T01:18:54.708836+00:00

    arviz_version :

    0.22.0

    inference_library :

    pymc

    inference_library_version :

    5.25.1

  
  
• observed_data

  <xarray.Dataset> Size: 12kB
  Dimensions:                  (__obs__: 500, rt,response_extra_dim_0: 2)
  Coordinates:
    * __obs__                  (__obs__) int64 4kB 0 1 2 3 4 ... 496 497 498 499
    * rt,response_extra_dim_0  (rt,response_extra_dim_0) int64 16B 0 1
  Data variables:
      rt,response              (__obs__, rt,response_extra_dim_0) float64 8kB 0...
  Attributes:
      created_at:                 2025-10-17T01:18:54.709714+00:00
      arviz_version:              0.22.0
      inference_library:          pymc
      inference_library_version:  5.25.1

  xarray.Dataset

  • Dimensions:
    • __obs__: 500
    • rt,response_extra_dim_0: 2
  • Coordinates: (2)
    • __obs__
      (__obs__)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499], shape=(500,))

    • rt,response_extra_dim_0
      (rt,response_extra_dim_0)
      int64
      0 1

      array([0, 1])

  • Data variables: (1)
    • rt,response
      (__obs__, rt,response_extra_dim_0)
      float64
      0.5801 1.0 2.519 ... 1.0 1.422 1.0

      array([[ 0.58008313,  1.        ],
             [ 2.51914597,  1.        ],
             [ 5.1725502 , -1.        ],
             [ 5.38408518,  1.        ],
             [ 2.92615318,  1.        ],
             [ 2.60508227,  1.        ],
             [ 1.60326254,  1.        ],
             [ 1.45287967,  1.        ],
             [ 1.62478471, -1.        ],
             [ 1.49670947,  1.        ],
             [ 2.34649038,  1.        ],
             [ 2.30925393,  1.        ],
             [ 2.05400562,  1.        ],
             [ 0.68057668,  1.        ],
             [ 0.88679963,  1.        ],
             [ 0.41172174,  1.        ],
             [ 2.29736328,  1.        ],
             [ 2.66417313,  1.        ],
             [ 1.93516386,  1.        ],
             [ 1.83781409,  1.        ],
      ...
             [ 4.18482494,  1.        ],
             [ 1.79381561,  1.        ],
             [ 0.59991771,  1.        ],
             [ 3.06441665,  1.        ],
             [ 2.87169862,  1.        ],
             [ 1.98977184,  1.        ],
             [ 2.80292845,  1.        ],
             [ 0.73743659,  1.        ],
             [ 3.1652298 ,  1.        ],
             [ 2.587255  ,  1.        ],
             [ 3.69039965,  1.        ],
             [ 0.96508873,  1.        ],
             [ 0.45725834, -1.        ],
             [ 0.81268889,  1.        ],
             [ 2.14377213,  1.        ],
             [ 1.3394289 ,  1.        ],
             [ 2.96452594,  1.        ],
             [ 3.14054394,  1.        ],
             [ 0.86679167,  1.        ],
             [ 1.42190468,  1.        ]])

  • Indexes: (2)
    • __obs__
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
            dtype='int64', name='__obs__', length=500))

    • rt,response_extra_dim_0
      PandasIndex

      PandasIndex(Index([0, 1], dtype='int64', name='rt,response_extra_dim_0'))

  • Attributes: (4)

    created_at :

    2025-10-17T01:18:54.709714+00:00

    arviz_version :

    0.22.0

    inference_library :

    pymc

    inference_library_version :

    5.25.1

  
  
• constant_data

  <xarray.Dataset> Size: 32B
  Dimensions:  ()
  Data variables:
      z        float64 8B 0.5
      a        float64 8B 1.5
      t        float64 8B 0.1
      v        float64 8B 0.0
  Attributes:
      created_at:                 2025-10-17T01:18:54.710879+00:00
      arviz_version:              0.22.0
      inference_library:          pymc
      inference_library_version:  5.25.1

  xarray.Dataset

  • Dimensions:
  • Coordinates: (0)
  • Data variables: (4)
    • z
      ()
      float64
      0.5

      array(0.5)

    • a
      ()
      float64
      1.5

      array(1.5)

    • t
      ()
      float64
      0.1

      array(0.1)

    • v
      ()
      float64
      0.0

      array(0.)

  • Indexes: (0)
  • Attributes: (4)

    created_at :

    2025-10-17T01:18:54.710879+00:00

    arviz_version :

    0.22.0

    inference_library :

    pymc

    inference_library_version :

    5.25.1

  
  

Observe two things:

1. A closer look at the prior group reveals that no parameters were actually sampled (it collects on the deterministic v_mean which, arguably, we can skip too)
2. The prior_predictive group is where we collected our synthetic data

Our synthetic data sits in the prior_predictive group, since, under the hood, were de facto called the sample_prior_predictive() method, on our new do-model.

If you prefer to get your outputs as a dataframe, you can use the predictive_idata_to_dataframe method from the hssm.utils module.

In [8]:

synth_df = hssm.utils.predictive_idata_to_dataframe(synth_idata,
                                                    predictive_group = "prior_predictive")

synth_df = hssm.utils.predictive_idata_to_dataframe(synth_idata, predictive_group = "prior_predictive")

In [9]:

synth_df

synth_df

Out[9]:

	chain	draw	__obs__	rt	response
0	0	0	0	1.446554	1.0
1	0	0	1	1.198923	1.0
2	0	0	2	0.892631	-1.0
3	0	0	3	3.382068	1.0
4	0	0	4	0.436959	-1.0
...	...	...	...	...	...
49995	0	99	495	3.510932	1.0
49996	0	99	496	1.799391	1.0
49997	0	99	497	1.187817	-1.0
49998	0	99	498	0.320621	1.0
49999	0	99	499	1.530689	1.0

50000 rows × 5 columns

From here you may re-attach the original data used to define the hssm_model (e.g.to re-attach covariates that you want as part of the model you are testing on your synthetic data).

The sample_do() function can be extremely helpful in setting up numerical experiments to test e.g. for parameter identifiability in scenarios that you wish to control precisely.