Renamed mix field of Switch generative function to branches to more accurately reflect the pattern.

Author: femtomc

src/modeling_library/switch/assess.jl

@@ -9,7 +9,7 @@ function process!(gen_fn::Switch{C, N, K, T},
                   args::Tuple, 
                   choices::ChoiceMap, 
                   state::SwitchAssessState{T}) where {C, N, K, T}
-    (weight, retval) = assess(getindex(gen_fn.mix, index), args, choices)
+    (weight, retval) = assess(getindex(gen_fn.branches, index), args, choices)
     state.weight = weight
     state.retval = retval
 end

src/modeling_library/switch/generate.jl

@@ -14,7 +14,7 @@ function process!(gen_fn::Switch{C, N, K, T},
                   choices::ChoiceMap, 
                   state::SwitchGenerateState{T}) where {C, N, K, T}
 
-    (subtrace, weight) = generate(getindex(gen_fn.mix, index), args, choices)
+    (subtrace, weight) = generate(getindex(gen_fn.branches, index), args, choices)
     state.index = index
     state.subtrace = subtrace
     state.weight += weight

src/modeling_library/switch/propose.jl

@@ -10,7 +10,7 @@ function process!(gen_fn::Switch{C, N, K, T},
                   args::Tuple, 
                   state::SwitchProposeState{T}) where {C, N, K, T}
 
-    (submap, weight, retval) = propose(getindex(gen_fn.mix, index), args)
+    (submap, weight, retval) = propose(getindex(gen_fn.branches, index), args)
     state.choices = submap
     state.weight += weight
     state.retval = retval

src/modeling_library/switch/regenerate.jl

@@ -38,7 +38,7 @@ function process!(gen_fn::Switch{C, N, K, T},
                   kernel_argdiffs::Tuple,
                   selection::Selection, 
                   state::SwitchRegenerateState{T}) where {C, N, K, T}
-    branch_fn = getfield(gen_fn.mix, index)
+    branch_fn = getfield(gen_fn.branches, index)
     merged = regenerate_recurse_merge(get_choices(state.prev_trace), selection)
     new_trace, weight = generate(branch_fn, args, merged)
     retdiff = UnknownChange()

src/modeling_library/switch/simulate.jl

@@ -12,7 +12,7 @@ function process!(gen_fn::Switch{C, N, K, T},
                   args::Tuple, 
                   state::SwitchSimulateState{T}) where {C, N, K, T}
     local retval::T
-    subtrace = simulate(getindex(gen_fn.mix, index), args)
+    subtrace = simulate(getindex(gen_fn.branches, index), args)
     state.index = index
     state.noise += project(subtrace, EmptySelection())
     state.subtrace = subtrace

src/modeling_library/switch/switch.jl

@@ -1,5 +1,5 @@
 struct Switch{C, N, K, T} <: GenerativeFunction{T, Trace}
-    mix::NTuple{N, GenerativeFunction{T}}
+    branches::NTuple{N, GenerativeFunction{T}}
     cases::Dict{C, Int}
     function Switch(gen_fns::GenerativeFunction...)
         @assert !isempty(gen_fns)
@@ -14,10 +14,10 @@ struct Switch{C, N, K, T} <: GenerativeFunction{T, Trace}
 end
 export Switch
 
-has_argument_grads(switch_fn::Switch) = map(zip(map(has_argument_grads, switch_fn.mix)...)) do as
+has_argument_grads(switch_fn::Switch) = map(zip(map(has_argument_grads, switch_fn.branches)...)) do as
     all(as)
 end
-accepts_output_grad(switch_fn::Switch) = all(accepts_output_grad, switch_fn.mix)
+accepts_output_grad(switch_fn::Switch) = all(accepts_output_grad, switch_fn.branches)
 
 function (gen_fn::Switch)(index::Int, args...)
     (_, _, retval) = propose(gen_fn, (index, args...))

src/modeling_library/switch/update.jl

@@ -66,7 +66,7 @@ function process!(gen_fn::Switch{C, N, K, T},
 
     # Generate new trace.
     merged = update_recurse_merge(get_choices(state.prev_trace), choices)
-    branch_fn = getfield(gen_fn.mix, index)
+    branch_fn = getfield(gen_fn.branches, index)
     new_trace, weight = generate(branch_fn, args, merged)
     weight -= get_score(state.prev_trace)
     state.discard = update_discard(state.prev_trace, choices, new_trace)