Merge pull request #39 from ilabcode/dev

Dev

Author: PTWaade

src/ActionModels_variations/utils/give_inputs.jl

@@ -46,7 +46,7 @@ function ActionModels.give_inputs!(hgf::HGF, inputs::Array)
     #Take each row in the array
     for input in eachrow(inputs)
         #Input it to the hgf
-        update_hgf!(hgf, input)
+        update_hgf!(hgf, Vector(input))
     end
 
     return nothing

src/ActionModels_variations/utils/reset.jl

@@ -7,28 +7,60 @@ function ActionModels.reset!(hgf::HGF)
 
         #For categorical state nodes
         if node isa CategoricalStateNode
-            #Reset the posterior to all 0's
-            node.states.posterior .= zero(Real)
-            #Set to missing
-            node.states.prediction = missing
-            node.states.value_prediction_error = missing
+            #Set states to vectors of missing
+            node.states.posterior .= missing
+            node.states.value_prediction_error .= missing
+            #Empty prediction state
+            empty!(node.states.prediction)
 
-        #For other nodes
+            #For binary input nodes
+        elseif node isa BinaryInputNode
+            #Set states to missing 
+            node.states.value_prediction_error .= missing
+            node.states.input_value = missing
+
+            #For continuous state nodes
+        elseif node isa ContinuousStateNode
+            #Set posterior to initial belief
+            node.states.posterior_mean = node.params.initial_mean
+            node.states.posterior_precision = node.params.initial_precision
+            #For other states
+            for state_name in [
+                :value_prediction_error,
+                :volatility_prediction_error,
+                :prediction_mean,
+                :prediction_volatility,
+                :prediction_precision,
+                :auxiliary_prediction_precision,
+            ]
+                #Set the state to missing
+                setfield!(node.states, state_name, missing)
+            end
+
+            #For continuous input nodes
+        elseif node isa ContinuousInputNode
+
+            #For all states except auxiliary prediction precision
+            for state_name in [
+                :input_value,
+                :value_prediction_error,
+                :volatility_prediction_error,
+                :prediction_volatility,
+                :prediction_precision,
+            ]
+                #Set the state to missing
+                setfield!(node.states, state_name, missing)
+            end
+
+            #For other nodes
         else
             #For each state
             for state_name in fieldnames(typeof(node.states))
-                #Set the state to first value in history
+                #Set the state to missing
                 setfield!(node.states, state_name, missing)
             end
         end
 
-        #For continuous state nodes
-        if node isa ContinuousStateNode
-            #Set the initial posterior
-            node.states.posterior_mean = node.params.initial_mean
-            node.states.posterior_precision = node.params.initial_precision
-        end
-
         #For each state in the history
         for state_name in fieldnames(typeof(node.history))
 

src/create_hgf/create_premade_hgf.jl

@@ -16,6 +16,8 @@ function premade_hgf(model_name::String, config::Dict = Dict(); verbose = true)
         "binary_3level" => premade_binary_3level,           #The standard binary input 3 level HGF
         "JGET" => premade_JGET,                             #The JGET model
         "categorical_3level" => premade_categorical_3level, #The standard categorical input 3 level HGF
+        "categorical_3level_state_transitions" =>
+            premade_categorical_3level_state_transitions,   #Categorical 3 level HGF for learning state transitions
     )
 
     #Check that the specified model is in the list of keys

src/create_hgf/init_hgf.jl

@@ -282,8 +282,13 @@ function init_hgf(;
                 push!(node.category_parent_order, parent.name)
             end
 
-            #Set posterior to node
-            node.states.posterior = zeros(length(node.value_parents))
+            #Set posterior to vector of zeros equal to the number of categories
+            node.states.posterior = Vector{Union{Real,Missing}}(missing,length(node.value_parents))
+            push!(node.history.posterior, node.states.posterior)
+
+            #Set posterior to vector of missing equal to the number of categories
+            node.states.value_prediction_error = node.states.posterior
+            push!(node.history.value_prediction_error, node.states.value_prediction_error)
 
             #For other nodes
         else

src/premade_models/premade_hgfs.jl

@@ -324,7 +324,7 @@ end
 function premade_categorical_3level(config::Dict; verbose::Bool = true)
 
     #Defaults
-    spec_defaults = Dict(
+    defaults = Dict(
         "n_categories" => 4,
         ("x2", "evolution_rate") => 0,
         ("x2", "initial_mean") => 0,
@@ -338,18 +338,18 @@ function premade_categorical_3level(config::Dict; verbose::Bool = true)
 
     #Warn the user about used defaults and misspecified keys
     if verbose
-        warn_premade_defaults(spec_defaults, config)
+        warn_premade_defaults(defaults, config)
     end
 
     #Merge to overwrite defaults
-    config = merge(spec_defaults, config)
+    config = merge(defaults, config)
 
 
     ##Prep category node parent names
     #Vector for category node binary parent names
-    category_binary_parent_names = []
+    category_binary_parent_names = Vector{String}()
     #Vector for binary node continuous parent names
-    binary_continuous_parent_names = []
+    binary_continuous_parent_names = Vector{String}()
     #Populate the above vectors with node names
     for category_number = 1:config["n_categories"]
         push!(category_binary_parent_names, "x1_" * string(category_number))
@@ -361,7 +361,7 @@ function premade_categorical_3level(config::Dict; verbose::Bool = true)
     input_nodes = Dict("name" => "u", "type" => "categorical")
 
     ##List of state nodes
-    state_nodes =[Dict{String, Any}("name" => "x1", "type" => "categorical")]
+    state_nodes = [Dict{String,Any}("name" => "x1", "type" => "categorical")]
 
     #Add category node binary parents
     for node_name in category_binary_parent_names
@@ -431,4 +431,184 @@ function premade_categorical_3level(config::Dict; verbose::Bool = true)
         edges = edges,
         verbose = false,
     )
-end
+end
+
+function premade_categorical_3level_state_transitions(config::Dict; verbose::Bool = true)
+
+    #Defaults
+    defaults = Dict(
+        "n_categories" => 4,
+        ("x2", "evolution_rate") => 0,
+        ("x2", "initial_mean") => 0,
+        ("x2", "initial_precision") => 1,
+        ("x3", "evolution_rate") => 0,
+        ("x3", "initial_mean") => 0,
+        ("x3", "initial_precision") => 1,
+        ("x1", "x2", "value_coupling") => 1,
+        ("x2", "x3", "volatility_coupling") => 1,
+    )
+
+    #Warn the user about used defaults and misspecified keys
+    if verbose
+        warn_premade_defaults(defaults, config)
+    end
+
+    #Merge to overwrite defaults
+    config = merge(defaults, config)
+
+
+    ##Prepare node names
+    #Empty lists
+    categorical_input_node_names = Vector{String}()
+    categorical_state_node_names = Vector{String}()
+    categorical_node_binary_parent_names = Vector{String}()
+    binary_node_continuous_parent_names = Vector{String}()
+
+    #Go through each category that the transition may have been from
+    for category_from = 1:config["n_categories"]
+        #One input node and its state node parent for each                             
+        push!(categorical_input_node_names, "u" * string(category_from))
+        push!(categorical_state_node_names, "x1_" * string(category_from))
+        #Go through each category that the transition may have been to
+        for category_to = 1:config["n_categories"]
+            #Each categorical state node has a binary parent for each
+            push!(
+                categorical_node_binary_parent_names,
+                "x1_" * string(category_from) * "_" * string(category_to),
+            )
+            #And each binary parent has a continuous parent of its own
+            push!(
+                binary_node_continuous_parent_names,
+                "x2_" * string(category_from) * "_" * string(category_to),
+            )
+        end
+    end
+
+    ##Create input nodes
+    #Initialize list
+    input_nodes = Vector{Dict}()
+
+    #For each categorical input node
+    for node_name in categorical_input_node_names
+        #Add it to the list
+        push!(input_nodes, Dict("name" => node_name, "type" => "categorical"))
+    end
+
+    ##Create state nodes
+    #Initialize list
+    state_nodes = Vector{Dict}()
+
+    #For each cateogrical state node
+    for node_name in categorical_state_node_names
+        #Add it to the list
+        push!(state_nodes, Dict("name" => node_name, "type" => "categorical"))
+    end
+
+    #For each categorical node binary parent
+    for node_name in categorical_node_binary_parent_names
+        #Add it to the list                                    
+        push!(state_nodes, Dict("name" => node_name, "type" => "binary"))
+    end
+
+    #For each binary node continuous parent
+    for node_name in binary_node_continuous_parent_names
+        #Add it to the list, with parameter settings from the config
+        push!(
+            state_nodes,
+            Dict(
+                "name" => node_name,
+                "type" => "continuous",
+                "evolution_rate" => config[("x2", "evolution_rate")],
+                "initial_mean" => config[("x2", "initial_mean")],
+                "initial_precision" => config[("x2", "initial_precision")],
+            ),
+        )
+    end
+
+    #Add the shared volatility parent of the continuous nodes
+    push!(
+        state_nodes,
+        Dict(
+            "name" => "x3",
+            "type" => "continuous",
+            "evolution_rate" => config[("x3", "evolution_rate")],
+            "initial_mean" => config[("x3", "initial_mean")],
+            "initial_precision" => config[("x3", "initial_precision")],
+        ),
+    )
+
+    ##Create child-parent relations
+    #Initialize list                                                     
+    edges = Vector{Dict}()
+
+    #For each categorical input node and its corresponding state node parent
+    for (child_name, parent_name) in
+        zip(categorical_input_node_names, categorical_state_node_names)
+        #Add their relation to the list
+        push!(edges, Dict("child" => child_name, "value_parents" => parent_name))
+    end
+
+    #For each categorical state node
+    for child_node_name in categorical_state_node_names
+        #Get the category it represents transitions from
+        (child_supername, child_category_from) = split(child_node_name, "_")
+
+        #For each potential parent node
+        for parent_node_name in categorical_node_binary_parent_names
+            #Get the category it represents transitions from
+            (parent_supername, parent_category_from, parent_category_to) =
+                split(parent_node_name, "_")
+
+            #If these match
+            if parent_category_from == child_category_from
+                #Add the parent as parent of the child
+                push!(
+                    edges,
+                    Dict("child" => child_node_name, "value_parents" => parent_node_name),
+                )
+            end
+        end
+    end
+
+    #For each binary parent of categorical nodes and their corresponding continuous parents
+    for (child_name, parent_name) in
+        zip(categorical_node_binary_parent_names, binary_node_continuous_parent_names)
+        #Add their relations to the list, with the same value coupling
+        push!(
+            edges,
+            Dict(
+                "child" => child_name,
+                "value_parents" => (parent_name, config[("x1", "x2", "value_coupling")]),
+            ),
+        )
+    end
+
+    #Add the shared continuous node volatility parent to the continuous nodes
+    for child_name in binary_node_continuous_parent_names
+        push!(
+            edges,
+            Dict(
+                "child" => child_name,
+                "volatility_parents" => ("x3", config[("x2", "x3", "volatility_coupling")]),
+            ),
+        )
+    end
+
+    #Initialize the HGF
+    init_hgf(
+        input_nodes = input_nodes,
+        state_nodes = state_nodes,
+        edges = edges,
+        verbose = false,
+    )
+end
+
+
+
+
+
+
+
+
+
+