Merge pull request #171 from ComputationalPsychiatry/dev

Dev

Author: PTWaade

Project.toml

@@ -1,7 +1,7 @@
 name = "ActionModels"
 uuid = "320cf53b-cc3b-4b34-9a10-0ecb113566a3"
 authors = ["Peter Thestrup Waade [email protected]", "Luke Ring [email protected]", "Malte Lau Møller", "Christoph Mathys [email protected]"]
-version = "0.7.2"
+version = "0.7.3"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

docs/julia_files/B_user_guide/4_model_fitting.jl

@@ -99,7 +99,7 @@ chns = sample_posterior!(model)
 # We can specify the number of samples and chains to sample with the `n_samples` and `n_chains` keyword arguments.
 # The `init_params` keyword argument can be used to specify how the initial parameters for the chains are set.
 # It can be set to `:MAP` or `:MLE` to use the maximum a posteriori or maximum likelihood estimates as the initial parameters, respectively.
-# It can be set to `:sample_prior` to draw a single sample from the prior distribution, or to `nothing` to use Turing's default of random values between -2 and 2 as the initial parameters.`
+# It can be set to `:sample_prior` to draw a single sample from the prior distribution, or to `nothing` to use Turing's default of random values between -2 and 2 as the initial parameters.
 # Finally, a vector of initial parameters can be passed, which will be used as the initial parameters for all chains.
 # Other arguments for the sampling can also be passed. This includes the autodifferentiation backend to use, which can be set with the `ad_type` keyword argument, and the sampler to use, which can be set with the `sampler` keyword argument.
 # Notably, `sample_posterior!` will return the already sampled `Chains` object if the posterior has already been sampled. Set `resample = true` to override the already sampled posterior.

src/fitting_models/turing_model/create_model.jl

@@ -175,9 +175,26 @@ function create_model(
         state_name => state.initial_value for
         (state_name, state) in pairs(action_model.states)
     )
+        
+    ## Create population data ##
+    #Remove action and observation columns
+    population_data = data[!,setdiff(Symbol.(names(data)), vcat(observation_cols, action_cols))]
+    #If there are session columns
+    if length(session_cols) > 0
+        #Only one row per session
+        population_data = unique(population_data, session_cols)
+        #Sort population data by session columns
+        population_data = sort(population_data, session_cols)
+    else
+        #If there are no session columns, just take the first row
+        population_data = population_data[1:1, :]
+    end
 
-    ## Group data by sessions ##
-    grouped_data = groupby(data, session_cols)
+    ## Create sessions data ##
+    #Only keep actions, observations and session columns
+    sessions_data = data[!, unique(vcat(collect(observation_cols), collect(action_cols), session_cols))]
+    #Group sessions data by session columns
+    sessions_data = groupby(sessions_data, session_cols, sort = true)
 
     ## Create IDs for each session ##
     session_ids = [
@@ -186,18 +203,18 @@ function create_model(
                 string(col_name) * id_column_separator * string(first(subdata)[col_name]) for col_name in session_cols
             ],
             id_separator,
-        ) for subdata in grouped_data
+        ) for subdata in sessions_data
     ]
 
     ## Extract observations and actions ##
     observations = Vector{Tuple{observation_types_data...}}[
         Tuple{observation_types_data...}.(
             eachrow(session_data[!, collect(observation_cols)]),
-        ) for session_data in grouped_data
+        ) for session_data in sessions_data
     ]
     actions = Vector{Tuple{action_types_data...}}[
         Tuple{action_types_data...}.(eachrow(session_data[!, collect(action_cols)])) for
-        session_data in grouped_data
+        session_data in sessions_data
     ]
 
     ### CREATE MODEL ###
@@ -220,15 +237,6 @@ function create_model(
         initial_states,
     )
 
-    ## Extract population data ##
-    if length(session_cols) == 0
-        #If there are no session columns, the population data is the same as the data
-        population_data = data[1:1, :]
-    else
-        #Otherwise, extract the population data from the grouped data
-        population_data = unique(data, session_cols)
-    end
-
     return ModelFit(
         model = model,
         population_model_type = population_model_type,

src/fitting_models/turing_model/population_models/glm_population_model.jl

@@ -88,6 +88,12 @@ function create_model(
         regressions = F[regressions]
     end
 
+    #If only a symbol was specified for session cols
+    if session_cols isa Symbol
+        #Convert single session column to vector
+        session_cols = [session_cols]
+    end
+
     #Make sure that single formulas are made into Regression objects
     regressions = [
         regression isa Regression ? regression : Regression(regression) for
@@ -107,8 +113,21 @@ function create_model(
         kwargs...,
     )
 
-    #Extract just the data needed for the linear regression
-    population_data = unique(data, session_cols)
+    ## Create population data ##
+    #Remove action and observation columns
+    population_data =
+        data[!, setdiff(Symbol.(names(data)), vcat(observation_cols, action_cols))]
+    #If there are session columns
+    if length(session_cols) > 0
+        #Only one row per session
+        population_data = unique(population_data, session_cols)
+        #Sort population data by session columns
+        population_data = sort(population_data, session_cols)
+    else
+        #If there are no session columns, just take the first row
+        population_data = population_data[1:1, :]
+    end
+
     #Extract number of sessions
     n_sessions = nrow(population_data)