Merge pull request #107 from ilabcode/dev

version 0.5.1

Project.toml

@@ -4,7 +4,7 @@ authors = ["Peter Thestrup Waade ptw@cas.au.dk",
             "Anna Hedvig Møller hedvig.2808@gmail.com",
             "Jacopo Comoglio jacopo.comoglio@gmail.com",
             "Christoph Mathys chmathys@cas.au.dk"]
-version = "0.5.0"
+version = "0.5.1"
 
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"

src/fitting/create_model.jl

@@ -143,24 +143,38 @@ Create a Turing model object used for fitting an ActionModels agent.
                 #If there is only a single action
                 if !multiple_actions
 
-                    #If the action isn't missing, or if missing actions are to be imputed
-                    if !ismissing(actions[group][timestep]) || impute_missing_actions
-                        #Sample the action from the probability distribution
-                        actions[group][timestep] ~ action_distribution
+                    # If missing actions are to be imputed, we do not care if the action exists
+                    if !impute_missing_actions
+                        # If we should not impute missing actions, we need to check if the action exists
+                        @inbounds action_exists = !ismissing(actions[group][timestep])
+                        # if the action doesn't exist, we skip this timestep
+                        if !action_exists
+                            continue
+                        end
                     end
 
+                    #Sample the action from the probability distribution
+                    actions[group][timestep] ~ action_distribution
+
+
                     #If there are multiple actions
                 else
                     #Go through each separate action
                     for (action_idx, single_distribution) in enumerate(action_distribution)
 
-                        #If the action isn't missing, or if missing actions are to be imputed
-                        if !ismissing(actions[group][timestep, action_idx]) ||
-                           impute_missing_actions
-
-                            #Sample the action from the probability distribution
-                            actions[group][timestep, action_idx] ~ single_distribution
+                        # If missing actions are to be imputed, we do not care if the action exists
+                        if !impute_missing_actions
+                            # If we should not impute missing actions, we need to check if the action exists
+                            @inbounds action_exists = !ismissing(actions[group][timestep, action_idx])
+                            # if the action doesn't exist, we skip this timestep
+                            if !action_exists
+                                continue
+                            end
                         end
+
+                        #Sample the action from the probability distribution
+                        @inbounds actions[group][timestep, action_idx] ~ single_distribution
+
                     end
                 end
             end

src/fitting/fit_model.jl

@@ -46,7 +46,7 @@ function fit_model(
     input_cols::Vector = [:input],
     action_cols::Vector = [:action],
     fixed_parameters::Dict = Dict(),
-    sampler::Union{Missing, DynamicPPL.Sampler} = missing,
+    sampler::Turing.Inference.InferenceAlgorithm = NUTS(-1, 0.65; adtype=AutoReverseDiff(true)),
     n_cores::Integer = 1,
     n_iterations::Integer = 1000,
     n_chains::Integer = 2,
@@ -56,10 +56,7 @@ function fit_model(
     sampler_kwargs...,
 )
     ### SETUP ###
-    # If no sampler has been specified, use NUTS
-    if ismissing(sampler)
-        sampler = NUTS(-1, 0.65; adtype=AutoReverseDiff(true))
-    end
+
     #Convert column names to symbols
     independent_group_cols = Symbol.(independent_group_cols)
     multilevel_group_cols = Symbol.(multilevel_group_cols)
@@ -326,7 +323,7 @@ function fit_model(
     inputs::Array,
     actions::Array;
     fixed_parameters::Dict = Dict(),
-    sampler::Union{Missing, DynamicPPL.Sampler} = missing,
+    sampler::Turing.Inference.InferenceAlgorithm = NUTS(-1, 0.65; adtype=AutoReverseDiff(true)),
     n_cores::Integer = 1,
     n_iterations::Integer = 1000,
     n_chains = 2,

test/testsuite/fitting_tests.jl

@@ -149,3 +149,23 @@ end
     )
 
 end
+
+
+@testset "ensure parameters are reset after fitting" begin
+
+    agent = premade_agent("binary_rescorla_wagner_softmax", verbose = false)
+
+    initial_parameters = get_parameters(agent)
+
+    param_priors = Dict("learning_rate" => Uniform(0, 1))
+
+    inputs = [1, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0]
+
+    actions = give_inputs!(agent, inputs)
+
+    chains =
+        fit_model(agent, param_priors, inputs, actions, n_chains = 1, n_iterations = 10, verbose = false)
+
+    @test get_parameters(agent) == initial_parameters
+
+end