Using compiled log-likelihood functions

Compile log-likelihood function¶

In [1]:

import matplotlib.pyplot as plt
import numpy as np

import hssm

import matplotlib.pyplot as plt import numpy as np import hssm

Simulate Data¶

In [2]:

obs_ddm = hssm.simulate_data(
    theta={"v": 0.5, "a": 1.5, "t": 0.3, "z": 0.5},
    model="angle",
    size=500,
)

obs_ddm = hssm.simulate_data( theta={"v": 0.5, "a": 1.5, "t": 0.3, "z": 0.5}, model="angle", size=500, )

Basic HSSM model¶

In [3]:

model = hssm.HSSM(
    data=obs_ddm, loglik_kind="analytical", process_initvals=True, p_outlier=0
)

model = hssm.HSSM( data=obs_ddm, loglik_kind="analytical", process_initvals=True, p_outlier=0 )

Model initialized successfully.

In [4]:

model.graph()

model.graph()

Out[4]:

We can now use the compile_logp() method to compile the log-likelihood function created by the hssm. This illustrates the simplest use case, compile_logp() with no additional arguments. Check the documentation for more details on how to make compile_logp() work for more customized use cases.

In [5]:

logp_fun = model.compile_logp()  # msynth.pymc_model.compile_logp()
print(logp_fun(model.initial_point(transformed=False)))

logp_fun = model.compile_logp() # msynth.pymc_model.compile_logp() print(logp_fun(model.initial_point(transformed=False)))

-33054.449962780134

Note that logp_fun takes as input a dictionary of parameter values with keys corresponding to the parameters names created by hssm. It might be helpful to take a look at the initial_point() method to see how the parameters are passed.

In [6]:

print(model.initial_point(transformed=False))

print(model.initial_point(transformed=False))

{'a': array(2.), 'z': array(0.5), 't': array(2.), 'v': array(0.)}

Timing the compiled log-likelihood function¶

In [7]:

# time
import time

my_start_point = model.initial_point(transformed=False)
start_time = time.time()
for i in range(1000):
    logp_fun(my_start_point)
print((time.time() - start_time) / 1000)

# time import time my_start_point = model.initial_point(transformed=False) start_time = time.time() for i in range(1000): logp_fun(my_start_point) print((time.time() - start_time) / 1000)

7.621407508850098e-05

Using the compiled log-likelihood function¶

For illustration, we use our compile log-likelihood function with a different MCMC sampler, using the zeus package.

Note: zeus is an optional dependency of hssm and can be installed using the command pip install hssm[notebook]. Alternatively, you can install it directly with pip install zeus-mcmc.

Wrap the compiled log-likelihood to accomodate zeus¶

In [8]:

def mylogp(theta: list[float]) -> float:
    """Wrap function for compiled log probability function to work with zeus sampler.

    Args
    ----
        theta: List of model parameters [v, a, z, t] where:
            v: Drift rate
            a: Boundary separation
            z: Starting point
            t: Non-decision time

    Returns
    -------
        float: Log probability value for the given parameters
    """
    v, a, z, t = theta
    return logp_fun({"v": v, "a": a, "z": z, "t": t})

def mylogp(theta: list[float]) -> float: """Wrap function for compiled log probability function to work with zeus sampler. Args ---- theta: List of model parameters [v, a, z, t] where: v: Drift rate a: Boundary separation z: Starting point t: Non-decision time Returns ------- float: Log probability value for the given parameters """ v, a, z, t = theta return logp_fun({"v": v, "a": a, "z": z, "t": t})

Test sampling with zeus¶

In [9]:

import zeus

start = np.random.uniform(low=-0.2, high=0.2, size=(8, 4)) + np.tile(
    [0.5, 1.5, 0.5, 0.3], (8, 1)
)
sampler = zeus.EnsembleSampler(8, 4, mylogp)

import zeus start = np.random.uniform(low=-0.2, high=0.2, size=(8, 4)) + np.tile( [0.5, 1.5, 0.5, 0.3], (8, 1) ) sampler = zeus.EnsembleSampler(8, 4, mylogp)

In [10]:

sampler.run_mcmc(start, 1000)

sampler.run_mcmc(start, 1000)

Initialising ensemble of 8 walkers...

Sampling progress :   0%|          | 0/1000 [00:00<?, ?it/s]

Sampling progress :   2%|▏         | 20/1000 [00:00<00:04, 196.76it/s]

Sampling progress :   4%|▍         | 40/1000 [00:00<00:05, 174.40it/s]

Sampling progress :   6%|▌         | 58/1000 [00:00<00:05, 169.83it/s]

Sampling progress :   8%|▊         | 76/1000 [00:00<00:05, 163.47it/s]

Sampling progress :   9%|▉         | 93/1000 [00:00<00:05, 160.78it/s]

Sampling progress :  11%|█         | 110/1000 [00:00<00:05, 156.87it/s]

Sampling progress :  13%|█▎        | 126/1000 [00:00<00:05, 157.08it/s]

Sampling progress :  14%|█▍        | 142/1000 [00:00<00:05, 155.76it/s]

Sampling progress :  16%|█▌        | 158/1000 [00:00<00:05, 151.24it/s]

Sampling progress :  17%|█▋        | 174/1000 [00:01<00:05, 151.42it/s]

Sampling progress :  19%|█▉        | 191/1000 [00:01<00:05, 153.92it/s]

Sampling progress :  21%|██        | 207/1000 [00:01<00:05, 145.19it/s]

Sampling progress :  22%|██▏       | 222/1000 [00:01<00:05, 142.94it/s]

Sampling progress :  24%|██▎       | 237/1000 [00:01<00:05, 141.95it/s]

Sampling progress :  25%|██▌       | 252/1000 [00:01<00:05, 143.93it/s]

Sampling progress :  27%|██▋       | 267/1000 [00:01<00:05, 145.47it/s]

Sampling progress :  28%|██▊       | 284/1000 [00:01<00:04, 151.18it/s]

Sampling progress :  30%|███       | 301/1000 [00:01<00:04, 154.57it/s]

Sampling progress :  32%|███▏      | 318/1000 [00:02<00:04, 155.93it/s]

Sampling progress :  33%|███▎      | 334/1000 [00:02<00:04, 154.44it/s]

Sampling progress :  35%|███▌      | 350/1000 [00:02<00:04, 149.32it/s]

Sampling progress :  37%|███▋      | 366/1000 [00:02<00:04, 150.01it/s]

Sampling progress :  38%|███▊      | 382/1000 [00:02<00:04, 148.91it/s]

Sampling progress :  40%|███▉      | 399/1000 [00:02<00:03, 150.61it/s]

Sampling progress :  42%|████▏     | 415/1000 [00:02<00:04, 143.45it/s]

Sampling progress :  43%|████▎     | 430/1000 [00:02<00:04, 139.51it/s]

Sampling progress :  44%|████▍     | 445/1000 [00:02<00:04, 138.27it/s]

Sampling progress :  46%|████▌     | 460/1000 [00:03<00:03, 139.47it/s]

Sampling progress :  47%|████▋     | 474/1000 [00:03<00:03, 137.46it/s]

Sampling progress :  49%|████▉     | 489/1000 [00:03<00:03, 138.74it/s]

Sampling progress :  50%|█████     | 504/1000 [00:03<00:03, 141.26it/s]

Sampling progress :  52%|█████▏    | 521/1000 [00:03<00:03, 148.50it/s]

Sampling progress :  54%|█████▎    | 537/1000 [00:03<00:03, 150.07it/s]

Sampling progress :  55%|█████▌    | 553/1000 [00:03<00:03, 143.40it/s]

Sampling progress :  57%|█████▋    | 568/1000 [00:03<00:03, 143.68it/s]

Sampling progress :  58%|█████▊    | 583/1000 [00:03<00:02, 144.19it/s]

Sampling progress :  60%|█████▉    | 598/1000 [00:04<00:02, 143.59it/s]

Sampling progress :  61%|██████▏   | 613/1000 [00:04<00:02, 142.66it/s]

Sampling progress :  63%|██████▎   | 628/1000 [00:04<00:02, 143.65it/s]

Sampling progress :  64%|██████▍   | 643/1000 [00:04<00:02, 142.45it/s]

Sampling progress :  66%|██████▌   | 658/1000 [00:04<00:02, 142.06it/s]

Sampling progress :  67%|██████▋   | 674/1000 [00:04<00:02, 146.28it/s]

Sampling progress :  69%|██████▉   | 689/1000 [00:04<00:02, 146.79it/s]

Sampling progress :  70%|███████   | 704/1000 [00:04<00:02, 144.76it/s]

Sampling progress :  72%|███████▏  | 719/1000 [00:04<00:01, 140.90it/s]

Sampling progress :  73%|███████▎  | 734/1000 [00:04<00:01, 142.16it/s]

Sampling progress :  75%|███████▌  | 750/1000 [00:05<00:01, 143.48it/s]

Sampling progress :  76%|███████▋  | 765/1000 [00:05<00:01, 128.61it/s]

Sampling progress :  78%|███████▊  | 779/1000 [00:05<00:01, 125.01it/s]

Sampling progress :  79%|███████▉  | 793/1000 [00:05<00:01, 128.28it/s]

Sampling progress :  81%|████████  | 809/1000 [00:05<00:01, 135.92it/s]

Sampling progress :  82%|████████▏ | 824/1000 [00:05<00:01, 137.36it/s]

Sampling progress :  84%|████████▍ | 838/1000 [00:05<00:01, 132.83it/s]

Sampling progress :  85%|████████▌ | 852/1000 [00:05<00:01, 132.17it/s]

Sampling progress :  87%|████████▋ | 867/1000 [00:05<00:00, 135.19it/s]

Sampling progress :  88%|████████▊ | 882/1000 [00:06<00:00, 136.85it/s]

Sampling progress :  90%|████████▉ | 896/1000 [00:06<00:00, 136.71it/s]

Sampling progress :  91%|█████████ | 911/1000 [00:06<00:00, 140.31it/s]

Sampling progress :  93%|█████████▎| 926/1000 [00:06<00:00, 142.90it/s]

Sampling progress :  94%|█████████▍| 941/1000 [00:06<00:00, 144.79it/s]

Sampling progress :  96%|█████████▌| 957/1000 [00:06<00:00, 146.77it/s]

Sampling progress :  97%|█████████▋| 972/1000 [00:06<00:00, 146.14it/s]

Sampling progress :  99%|█████████▉| 989/1000 [00:06<00:00, 152.41it/s]

Sampling progress : 100%|██████████| 1000/1000 [00:06<00:00, 145.51it/s]

In [11]:

plt.figure(figsize=(16, 1.5 * 4))
for n in range(4):
    plt.subplot2grid((4, 1), (n, 0))
    plt.plot(sampler.get_chain()[:, :, n], alpha=0.5)
plt.tight_layout()
plt.show()

plt.figure(figsize=(16, 1.5 * 4)) for n in range(4): plt.subplot2grid((4, 1), (n, 0)) plt.plot(sampler.get_chain()[:, :, n], alpha=0.5) plt.tight_layout() plt.show()