this is your model! great stuff btw i learned a lot from reading your paper on your tennis model!
i did something similar but didn’t have the init_dist part and just used the sigma_s. it is sampling extremely slow however but definitely faster with the init_dist. that might be expected though since the dataset has almost 10000 matches in it