lqg/cpp_data_fit.py at main · RothkopfLab/lqg · GitHub

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
import argparse
import numpyro
from numpyro.infer import MCMC, NUTS, init_to_median
from jax import random
import arviz as az


from lqg.io import load_tracking_data
from lqg.infer.models import shared_params_lqg_model
from lqg.infer import get_model_params
from lqg import tracking

numpyro.set_host_device_count(4)


def parse_args(args=None, namespace=None):
    parser = argparse.ArgumentParser(description="Continuous Psychophysics")
    parser.add_argument(
        "--delay",
        type=int,
        default=12,
        help="Time delay, by which target and mouse position are shifted",
    )
    parser.add_argument(
        "--clip",
        type=int,
        default=180,
        help="Clip the initial n time steps of the data",
    )
    parser.add_argument(
        "--nsamp", type=int, default=5_000, help="Number of samples drawn by NUTS"
    )
    parser.add_argument(
        "--nburnin", type=int, default=1_500, help="Number of burn-in samples."
    )
    parser.add_argument("--nchain", type=int, default=4)
    parser.add_argument("--model", type=str, default="BoundedActor", help="Model type")
    parser.add_argument("--seed", type=int, default=1, help="Random seed (for NUTS)")
    parser.add_argument(
        "--shared_params",
        type=str,
        nargs="*",
        default=[
            "action_variability",
            "action_cost",
            "sigma_cursor",
            "subj_noise",
            "subj_vel_noise",
        ],
        help="Parameters of the model to be shared across conditions",
    )
    args = parser.parse_args(args=args, namespace=namespace)

    model_params = get_model_params(getattr(tracking, args.model)).keys()

    args.shared_params = [p for p in args.shared_params if p in list(model_params)]
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    data, bws = load_tracking_data(
        delay=args.delay, clip=args.clip, subtract_mean=False
    )

    print(data.shape)

    nuts_kernel = NUTS(shared_params_lqg_model, init_strategy=init_to_median)

    mcmc = MCMC(
        nuts_kernel,
        num_warmup=args.nburnin,
        num_samples=args.nsamp,
        num_chains=args.nchain,
    )
    mcmc.run(
        random.PRNGKey(args.seed),
        data,
        getattr(tracking, args.model),
        shared_params=args.shared_params,
    )

    inference_data = az.convert_to_inference_data(mcmc)
    inference_data.to_netcdf(f"data/processed/{args.model}-{args.seed}.nc")