Fix cumulative and density function

src/compiler_pass.jl

@@ -631,27 +631,65 @@
         push!(args, (var.args[1], ()))
         push!(deps, (var.args[1], ()))
         return Expr(var.head, var.args[1], idxs...), deps, args
-    else # function call
-        fun_args = []
-        for i in 2:length(var.args)
-            e, d, a = evaluate_and_track_dependencies(var.args[i], env)
-            push!(fun_args, e)
+    elseif Meta.isexpr(var, :call)
+        if var.args[1] === :cumulative || var.args[1] === :density
+            arg1, arg2 = var.args[2:3]
+            arg1 = if arg1 isa Symbol
+                push!(deps, arg1)
+                # no need to add to arg, as the value doesn't matter
+                arg1
+            elseif Meta.isexpr(arg1, :ref)
+                v, indices... = arg1.args
+                for i in eachindex(indices)
+                    e, d, a = evaluate_and_track_dependencies(indices[i], env)
+                    union!(deps, d)
+                    union!(args, a)
+                    indices[i] = e
+                end
+                if any(!is_resolved, indices)
+                    error(
+                        "For now, the indices of the first argument to `cumulative` and `density` must be resolved, got $indices",
+                    )
+                end
+                push!(deps, (v, Tuple(indices)))
+                # no need to add to arg, as the value doesn't matter
+                Expr(:ref, v, indices...)
+            else
+                error(
+                    "First argument to `cumulative` and `density` must be variable, got $(arg1)",
+                )
+            end
+            e, d, a = evaluate_and_track_dependencies(arg2, env)
             union!(deps, d)
             union!(args, a)
-        end
+            return Expr(:call, var.args[1], arg1, e), deps, args
+        else
+            fun_args = []
+            for i in 2:length(var.args)
+                e, d, a = evaluate_and_track_dependencies(var.args[i], env)
+                push!(fun_args, e)
+                union!(deps, d)
+                union!(args, a)
+            end
 
-        for a in fun_args
-            a isa Symbol && a != :nothing && a != :(:) && (push!(deps, a); push!(args, a))
-        end
+            for a in fun_args
+                a isa Symbol &&
+                    a != :nothing &&
+                    a != :(:) &&
+                    (push!(deps, a); push!(args, a))
+            end
 
-        if (
-            var.args[1] ∈ BUGSPrimitives.BUGS_FUNCTIONS ||
-            var.args[1] ∈ (:+, :-, :*, :/, :^, :(:))
-        ) && all(is_resolved, args)
-            return getfield(JuliaBUGS, var.args[1])(fun_args...), deps, args
-        else
-            return Expr(var.head, var.args[1], fun_args...), deps, args
+            if (
+                var.args[1] ∈ BUGSPrimitives.BUGS_FUNCTIONS ||
+                var.args[1] ∈ (:+, :-, :*, :/, :^, :(:))
+            ) && all(is_resolved, args)
+                return getfield(JuliaBUGS, var.args[1])(fun_args...), deps, args
+            else
+                return Expr(var.head, var.args[1], fun_args...), deps, args
+            end
         end
+    else
+        error("Unexpected expression type: $var")
     end
 end
 
@@ -691,36 +729,63 @@
 
 _replace_constants_in_expr(x::Number, env) = x
 function _replace_constants_in_expr(x::Symbol, env)
-    if haskey(env, x)
-        if env[x] isa Number # only plug in scalar variables
-            return env[x]
-        else # if it's an array, raise error because array indexing should be explicit
-            error("$x")
-        end
+    if haskey(env, x) && env[x] isa Number
+        return env[x]
     end
     return x
 end
 function _replace_constants_in_expr(x::Expr, env)
-    if Meta.isexpr(x, :ref) && all(x -> x isa Number, x.args[2:end])
-        if haskey(env, x.args[1])
-            val = env[x.args[1]][try_cast_to_int.(x.args[2:end])...]
+    if Meta.isexpr(x, :ref)
+        v, indices... = x.args
+        if haskey(env, v) && all(x -> x isa Union{Int,Float64}, indices)
+            val = env[v][map(Int, indices)...]
             return ismissing(val) ? x : val
+        else
+            for i in eachindex(indices)
+                indices[i] = _replace_constants_in_expr(indices[i], env)
+            end
+            return Expr(:ref, v, indices...)
         end
-    else # don't try to eval the function, but try to simplify
-        x = deepcopy(x) # because we are mutating the args
-        for i in 2:length(x.args)
-            try
-                x.args[i] = _replace_constants_in_expr(x.args[i], env)
-            catch e
-                rethrow(
-                    ErrorException(
-                        "Array indexing in BUGS must be explicit. However, `$(e.msg)` is accessed as a scalar.",
-                    ),
+    elseif Meta.isexpr(x, :call)
+        if x.args[1] === :cumulative || x.args[1] === :density
+            if length(x.args) != 3
+                error(
+                    "`cumulative` and `density` are special functions in BUGS and takes two arguments, got $(length(x.args) - 1)",
                 )
             end
+            f, arg1, arg2 = x.args
+            if arg1 isa Symbol
+                return Expr(:call, f, arg1, _replace_constants_in_expr(arg2, env))
+            elseif Meta.isexpr(arg1, :ref)
+                v, indices... = arg1.args
+                for i in eachindex(indices)
+                    indices[i] = _replace_constants_in_expr(indices[i], env)
+                end
+                return Expr(
+                    :call,
+                    f,
+                    Expr(:ref, v, indices...),
+                    _replace_constants_in_expr(arg2, env),
+                )
+            else
+                error(
+                    "First argument to `cumulative` and `density` must be variable, got $(x.args[2])",
+                )
+            end
+        elseif x.args[1] === :deviance
+            @warn(
+                "`deviance` function is not supported in JuliaBUGS, `deviance` will be treated as a general function."
+            )
+        else
+            x = deepcopy(x) # because we are mutating the args
+            for i in 2:length(x.args)
+                x.args[i] = _replace_constants_in_expr(x.args[i], env)
+            end
+            return x
         end
+    else
+        error("Unexpected expression type: $x")
     end
-    return x
 end
 
 """

src/model.jl

@@ -36,6 +36,7 @@ struct BUGSModel <: AbstractBUGSModel
     transformed_var_lengths::Dict{VarName,Int} # TODO: store this as a delta from `untransformed_var_lengths`?
 
     varinfo::SimpleVarInfo
+    distributions::Dict{VarName,Distribution}
     parameters::Vector{VarName}
     sorted_nodes::Vector{VarName}
 
@@ -81,6 +82,7 @@ function BUGSModel(
 )
     vs = initialize_var_store(data, vars, array_sizes)
     vi = SimpleVarInfo(vs, 0.0)
+    dist_store = Dict{VarName,Distribution}()
     parameters = VarName[]
     untransformed_param_length = 0
     transformed_param_length = 0
@@ -91,11 +93,11 @@ function BUGSModel(
 
         ni = g[vn]
         @unpack node_type, node_function_expr, node_args = ni
-        args = Dict(getsym(arg) => vi[arg] for arg in node_args)
+        args = Dict(getsym(arg) => vi[arg] for arg in node_args) # TODO: get rid of this
         expr = node_function_expr.args[2]
         if node_type == JuliaBUGS.Logical
             value = try
-                _eval(expr, args)
+                _eval(expr, args, dist_store)
             catch e
                 rethrow(
                     # UninitializedVariableError(
@@ -108,14 +110,15 @@ function BUGSModel(
             vi = setindex!!(vi, value, vn)
         else
             dist = try
-                _eval(expr, args)
+                _eval(expr, args, dist_store)
             catch _
                 rethrow(
                     UninitializedVariableError(
                         "Encounter support error when evaluating the distribution of $vn. Try to initialize variables $(join(collect(keys(args)), ", ")) first if not yet.",
                     ),
                 )
             end
+            dist_store[vn] = dist
             value = evaluate(vn, data) # `evaluate(::VarName, env)` is defined in `src/utils.jl`
             if value isa Nothing # not observed
                 push!(parameters, vn)
@@ -155,6 +158,7 @@ function BUGSModel(
         untransformed_var_lengths,
         transformed_var_lengths,
         vi,
+        dist_store,
         parameters,
         sorted_nodes,
         g,
@@ -212,7 +216,7 @@ function get_params_varinfo(m::BUGSModel, vi::SimpleVarInfo)
             args = Dict(getsym(arg) => vi[arg] for arg in node_args)
             expr = node_function_expr.args[2]
             if vn in m.parameters
-                dist = _eval(expr, args)
+                dist = _eval(expr, args, m.distributions)
                 linked_val = DynamicPPL.link(dist, vi[vn])
                 d[vn] = linked_val
             end
@@ -252,7 +256,7 @@ function getparams(m::BUGSModel, vi::SimpleVarInfo; transformed::Bool=false)
         for v in m.parameters
             ni = m.g[v]
             args = (; (getsym(arg) => vi[arg] for arg in ni.node_args)...)
-            dist = _eval(ni.node_function_expr.args[2], args)
+            dist = _eval(ni.node_function_expr.args[2], args, m.distributions)
 
             link_vals = Bijectors.link(dist, vi[v])
             len = m.transformed_var_lengths[v]
@@ -288,7 +292,7 @@ function setparams!!(
     for v in m.parameters
         ni = m.g[v]
         args = (; (getsym(arg) => vi[arg] for arg in ni.node_args)...)
-        dist = _eval(ni.node_function_expr.args[2], args)
+        dist = _eval(ni.node_function_expr.args[2], args, m.distributions)
 
         len = if transformed
             m.transformed_var_lengths[v]
@@ -353,6 +357,7 @@ function AbstractPPL.condition(
         model.untransformed_var_lengths,
         model.transformed_var_lengths,
         varinfo,
+        model.distributions,
         new_parameters,
         sorted_blanket_with_vars,
         model.g,
@@ -377,6 +382,7 @@ function AbstractPPL.decondition(model::BUGSModel, var_group::Vector{<:VarName})
         model.untransformed_var_lengths,
         model.transformed_var_lengths,
         model.varinfo,
+        model.distributions,
         new_parameters,
         sorted_blanket_with_vars,
         model.g,
@@ -438,10 +444,10 @@ function AbstractPPL.evaluate!!(model::BUGSModel, ctx::SamplingContext)
         args = Dict(getsym(arg) => vi[arg] for arg in node_args)
         expr = node_function_expr.args[2]
         if node_type == JuliaBUGS.Logical
-            value = _eval(expr, args)
+            value = _eval(expr, args, model.distributions)
             vi = setindex!!(vi, value, vn)
         else
-            dist = _eval(expr, args)
+            dist = _eval(expr, args, model.distributions)
             value = rand(ctx.rng, dist) # just sample from the prior
             logp += logpdf(dist, value)
             vi = setindex!!(vi, value, vn)
@@ -464,10 +470,10 @@ function AbstractPPL.evaluate!!(model::BUGSModel, ::DefaultContext)
         args = Dict(getsym(arg) => vi[arg] for arg in node_args)
         expr = node_function_expr.args[2]
         if node_type == JuliaBUGS.Logical # be conservative -- always propagate values of logical nodes
-            value = _eval(expr, args)
+            value = _eval(expr, args, model.distributions)
             vi = setindex!!(vi, value, vn)
         else
-            dist = _eval(expr, args)
+            dist = _eval(expr, args, model.distributions)
             value = vi[vn]
             if model.transformed
                 # although the values stored in `vi` are in their original space, 
@@ -509,10 +515,10 @@ function AbstractPPL.evaluate!!(
         args = (; map(arg -> getsym(arg) => vi[arg], node_args)...)
         expr = node_function_expr.args[2]
         if node_type == JuliaBUGS.Logical
-            value = _eval(expr, args)
+            value = _eval(expr, args, model.distributions)
             vi = setindex!!(vi, value, vn)
         else
-            dist = _eval(expr, args)
+            dist = _eval(expr, args, model.distributions)
             if vn in model.parameters
                 l = var_lengths[vn]
                 if model.transformed

src/parser/bugs_macro.jl

@@ -1,71 +1,5 @@
-# handle `cumulative`, `density` and `deviance` functions
-# these are incorrect implementations, as it can only handle the case where the first argument is exactly the same as the LHS of the stochastic assignment
-# but can't handle cases like `cumulative(y[1], x)` where `y[i]` is defined in loops
-# other than that, `density` and `deviance` also require that the variable in place of first argument is observed
-# TODO: fix this
-function cumulative(expr::Expr)
-    return MacroTools.postwalk(expr) do sub_expr
-        if @capture(sub_expr, lhs_ = cumulative(s1_, s2_))
-            dist = find_tilde_rhs(expr, s1)
-            sub_expr.args[2].args[1] = :cdf
-            sub_expr.args[2].args[2] = dist
-            return sub_expr
-        else
-            return sub_expr
-        end
-    end
-end
-
-function density(expr::Expr)
-    return MacroTools.postwalk(expr) do sub_expr
-        if @capture(sub_expr, lhs_ = density(s1_, s2_))
-            dist = find_tilde_rhs(expr, s1)
-            sub_expr.args[2].args[1] = :pdf
-            sub_expr.args[2].args[2] = dist
-            return sub_expr
-        else
-            return sub_expr
-        end
-    end
-end
-
-function deviance(expr::Expr)
-    return MacroTools.postwalk(expr) do sub_expr
-        if @capture(sub_expr, lhs_ = deviance(s1_, s2_))
-            dist = find_tilde_rhs(expr, s1)
-            sub_expr.args[2].args[1] = :logpdf
-            sub_expr.args[2].args[2] = dist
-            sub_expr.args[2] = Expr(:call, :*, -2, sub_expr.args[2])
-            return sub_expr
-        else
-            return sub_expr
-        end
-    end
-end
-
-function find_tilde_rhs(expr::Expr, target::Union{Expr,Symbol})
-    dist = nothing
-    MacroTools.postwalk(expr) do sub_expr
-        if @capture(sub_expr, lhs_ ~ rhs_)
-            if lhs == target
-                isnothing(dist) || error("Exist two assignments to the same variable.")
-                dist = rhs
-            end
-        end
-        return sub_expr
-    end
-    isnothing(dist) && error(
-        "Error handling cumulative expression: can't find a stochastic assignment for $target.",
-    )
-    return dist
-end
-
-function handle_special_functions(expr::Expr)
-    return cumulative(density(deviance(expr)))
-end
-
-macro bugs(expr)
-    return Meta.quot(handle_special_functions(bugs_top(expr, __source__)))
+macro bugs(expr::Expr)
+    return Meta.quot(bugs_top(expr, __source__))
 end
 
 function bugs_top(@nospecialize(expr), __source__)
@@ -252,5 +186,5 @@ macro bugs(prog::String, replace_period=true, no_enclosure=false)
     if !isempty(error_container) # otherwise errors thrown in macro will be LoadError
         return :(throw(ErrorException(join($error_container, "\n"))))
     end
-    return Meta.quot(handle_special_functions(expr))
+    return Meta.quot(expr)
 end