Add some benchmark for logdensity computations (#172)

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

.github/workflows/Tests.yml

@@ -53,8 +53,11 @@ jobs:
         uses: julia-actions/julia-runtest@v1
         env:
           TEST_GROUP: "compilation"
-      - name: Profile compiler passes
-        run: julia --color=yes --project=. -e "import Pkg; Pkg.instantiate(); Pkg.test(;coverage=true);" >> $GITHUB_STEP_SUMMARY
+      - name: Run profile (Only on ubuntu-latest and Julia version 1)
+        if: matrix.os == 'ubuntu-latest' && matrix.version == '1'
+        run: |
+          julia --color=yes --project=. -e 'import Pkg; Pkg.instantiate(); Pkg.test(;coverage=true);'
+        shell: bash
         env:
           TEST_GROUP: "profile"
       - name: Running `gibbs` tests

src/BUGSExamples/Volume_1/14_LSAT.jl

@@ -1,7 +1,7 @@
 name = "LSAT: item response"
 
 model_def = @bugs begin
-    # Calculate individual (binary) responses to each test from multinomial data
+    # Calculate individual (binary) responses to each test from Multinomial data
     for j in 1:culm[1]
         for k in 1:T
             r[j, k] = response[1, k]
@@ -30,7 +30,7 @@ model_def = @bugs begin
         alpha[k] ~ dnorm(0, 0.0001)
         a[k] = alpha[k] - mean(alpha[:])
     end
-    beta ~ truncated(dflat(), 0, nothing)
+    beta ~ dunif(0, 1000)
 end
 
 original = """
@@ -63,7 +63,7 @@ model {
         alpha[k] ~ dnorm(0, 0.0001)
         a[k] <- alpha[k] - mean(alpha[])
     }
-    beta ~ dflat()T(0, )
+    beta ~ dunif(0, 1000)
 }
 """
 

src/BUGSExamples/Volume_1/14_LSAT.md

@@ -0,0 +1,29 @@
+# Section 6 of the Law School Aptitude Test (LSAT)
+
+Section 6 of the Law School Aptitude Test (LSAT) is a 5-item multiple choice test; students score 1 on each item for the correct answer and 0 otherwise, giving R = 32 possible response patterns. Boch and Lieberman (1970) present data on LSAT for N = 1000 students, part of which is shown below.
+
+| Response Pattern | Item 1 | Item 2 | Item 3 | Item 4 | Item 5 | Frequency |
+|------------------|--------|--------|--------|--------|--------|-----------|
+| 1                | 0      | 0      | 0      | 0      | 0      | 3         |
+| 2                | 0      | 0      | 0      | 0      | 1      | 6         |
+| 3                | 0      | 0      | 0      | 1      | 0      | 2         |
+| .                | .      | .      | .      | .      | .      | .         |
+| .                | .      | .      | .      | .      | .      | .         |
+| .                | .      | .      | .      | .      | .      | .         |
+| 30               | 1      | 1      | 1      | 0      | 1      | 61        |
+| 31               | 1      | 1      | 1      | 1      | 0      | 28        |
+| 32               | 1      | 1      | 1      | 1      | 1      | 298       |
+
+The above data may be analysed using the one-parameter Rasch model (see Andersen (1980), pp.253-254; Boch and Aitkin (1981)). The probability $p_{jk}$ that student $j$ responds correctly to item $k$ is assumed to follow a logistic function parameterized by an 'item difficulty' or threshold parameter $a_k$ and a latent variable $q_j$ representing the student's underlying ability. The ability parameters are assumed to have a Normal distribution in the population of students. That is:
+
+$$logit(p_{jk}) = q_j - a_k, j = 1,...,1000; k = 1,...,5$$
+
+$$q_j \sim Normal(0, \tau)$$
+
+The above model is equivalent to the following random effects logistic regression:
+
+$$logit(p_{jk}) = \beta q_j - a_k, j = 1,...,1000; k = 1,...,5$$
+
+$$q_j \sim Normal(0, 1)$$
+
+where $\beta$ corresponds to the scale parameter ($\beta^2 = \tau$) of the latent ability distribution. We assume a half-normal distribution with small precision for $\beta$; this represents vague prior information but constrains $\beta$ to be positive. Standard vague normal priors are assumed for the $a_k$'s. Note that the location of the $a_k$'s depend upon the mean of the prior distribution for $q_j$ which we have arbitrarily fixed to be zero. Alternatively, Boch and Aitkin ensure identifiability by imposing a sum-to-zero constraint on the $a_k$'s. Hence we calculate $a_k = a_k - \bar{a}$ to enable comparision of the BUGS posterior parameter estimates with the Boch and Aitkin marginal maximum likelihood estimates.

src/compiler_pass.jl

@@ -769,7 +769,7 @@ function analyze_statement(pass::AddEdges, expr::Expr, loop_vars::NamedTuple)
         vertex_code = filter(
             !iszero, vertex_code isa AbstractArray ? vertex_code : [vertex_code]
         )
-        vertex_labels = map(x -> label_for(pass.g, x), vertex_code)
+        vertex_labels = [label_for(pass.g, code) for code in vertex_code]
         for r in vertex_labels
             if r != lhs_vn
                 add_edge!(pass.g, r, lhs_vn)

test/profiles/prof_compile_pass.jl

@@ -0,0 +1,40 @@
+using BenchmarkTools
+
+using JuliaBUGS
+using JuliaBUGS: BUGSExamples
+
+suite = BenchmarkGroup()
+
+for name in keys(BUGSExamples.VOLUME_1)
+    @info "Adding benchmark for $name"
+    model_def = BUGSExamples.VOLUME_1[name].model_def
+    data = BUGSExamples.VOLUME_1[name].data
+
+    non_data_scalars, non_data_array_sizes = JuliaBUGS.determine_array_sizes(
+        model_def, data
+    )
+
+    eval_env = JuliaBUGS.compute_data_transformation(
+        non_data_scalars, non_data_array_sizes, model_def, data
+    )
+
+    _suite = BenchmarkGroup()
+
+    _suite["CollectVariables"] = @benchmarkable JuliaBUGS.determine_array_sizes(
+        $model_def, $data
+    )
+
+    _suite["DataTransformation"] = @benchmarkable JuliaBUGS.compute_data_transformation(
+        $non_data_scalars, $non_data_array_sizes, $model_def, $data
+    )
+
+    model_def = JuliaBUGS.concretize_colon_indexing(model_def, eval_env)
+    _suite["GraphCreation"] = @benchmarkable JuliaBUGS.create_graph($model_def, $eval_env)
+
+    tune!(_suite)
+    suite[string(name)] = _suite
+end
+
+results = run(suite; verbose=false)
+result_dict = create_result_dict(results)
+print_markdown_table(result_dict)

test/profiles/prof_logdensity.jl

@@ -0,0 +1,36 @@
+using JuliaBUGS
+using LogDensityProblems, LogDensityProblemsAD
+using AdvancedHMC, AbstractMCMC, MCMCChains, ReverseDiff
+using BenchmarkTools
+
+suite = BenchmarkGroup()
+
+for m in keys(JuliaBUGS.BUGSExamples.VOLUME_1)
+    exempl = JuliaBUGS.BUGSExamples.VOLUME_1[m]
+    (; model_def, data, inits) = exempl
+    model = compile(model_def, data, inits)
+    ad_model = ADgradient(:ReverseDiff, model; compile=Val(false))
+    ad_model_compiled = ADgradient(:ReverseDiff, model; compile=Val(true))
+    θ = rand(LogDensityProblems.dimension(model))
+
+    sub_suite = BenchmarkGroup()
+
+    sub_suite["logdensity"] = @benchmarkable LogDensityProblems.logdensity($model, $θ)
+    sub_suite["AD logdensity"] = @benchmarkable LogDensityProblems.logdensity($ad_model, $θ)
+    sub_suite["AD logdensity_and_gradient"] = @benchmarkable LogDensityProblems.logdensity_and_gradient(
+        $ad_model, $θ
+    )
+    sub_suite["AD compiled logdensity"] = @benchmarkable LogDensityProblems.logdensity(
+        $ad_model_compiled, $θ
+    )
+    sub_suite["AD compiled logdensity_and_gradient"] = @benchmarkable LogDensityProblems.logdensity_and_gradient(
+        $ad_model_compiled, $θ
+    )
+
+    suite[m] = sub_suite
+end
+
+tune!(suite)
+results = run(suite; verbose=false)
+result_dict = create_result_dict(results)
+print_markdown_table(result_dict)