From 40d3650e26af91fefa3092b0b96b1340cfb193a6 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Tue, 3 Mar 2026 13:56:52 +0100 Subject: [PATCH 01/23] add mock macro --- Project.toml | 2 +- scripts/macro.jl | 22 ++++++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) create mode 100644 scripts/macro.jl diff --git a/Project.toml b/Project.toml index dc40b51..c6fb26c 100644 --- a/Project.toml +++ b/Project.toml @@ -24,5 +24,5 @@ FlexiChains = "0.3" LinearAlgebra = "1.10" PDMats = "0.11" Random = "1.10" -Turing = "0.40" +Turing = "0.40, 0.41, 0.42" julia = "1.10" diff --git a/scripts/macro.jl b/scripts/macro.jl new file mode 100644 index 0000000..edef810 --- /dev/null +++ b/scripts/macro.jl @@ -0,0 +1,22 @@ +macro brm(x) + dump(x) +end +# The example BRM model is supposed to work with three dataframes (dBMI, dpmean, dpsd), which IIUC, have to have the same numbers of rows in this example and could have been merged into one. +@brm model(dBMI, dpmean, dpsd) = begin + # BMI will be a model parameter + BMI ~ Normal(dBMI.BMI_measured, 1) # equivalently: BMI ~ Normal(BMI_measured, 1) |> (data=dBMI) + # Age_first, Age_second are functions of data, and would be computed/updated exactly once + Age_first, Age_second = ploynomial_expand(dpmean.Age; order=2) # equivalently: Age_first, Age_second = ploynomial_expand(Age; order=2) |> (data=dpmean) + # I'm assuming performance_mean is a function of data and model parameters - I think it kind of has to be + performance_mean ~ 1 + Age_first * Treatment + Age_second + (1 + Treatment | Subject) + (1 + Age_first | Experimenter) |> (data=dpmean) + # I'm assuming performance_sd is a function of data and model parameters - I think it kind of has to be + log(performance_sd) ~ 1 + AGE * BMI + max(Age, BMI) + (1 + Age * BMI | Subject) |> (data=dpsd) + # Peter didn't specify data here - but I think it would have to be specified? Or would it be added to the model via conditioning syntax? + Performance ~ Normal(performance_mean, performance_sd) # |> (data=observations_df) as an example + @defaults begin + gr(Subject, by=ClinicalGroup, + Block1=>[Treatment, Age:BMI], + Block2=>[Age, BMI] + ) + end +end \ No newline at end of file From 1d2ec14542b78258c719756f020f510375d65b53 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Tue, 3 Mar 2026 14:33:07 +0100 Subject: [PATCH 02/23] add single dataframe version --- scripts/macro.jl | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/scripts/macro.jl b/scripts/macro.jl index edef810..c053c3f 100644 --- a/scripts/macro.jl +++ b/scripts/macro.jl @@ -19,4 +19,19 @@ end Block2=>[Age, BMI] ) end -end \ No newline at end of file +end + +# With a single dataframe, the model definition could look as follows (which would lower to the "obvious" syntax): +model = @brm begin + BMI ~ Normal(BMI_measured, 1) + Age_first, Age_second = ploynomial_expand(Age; order=2) + performance_mean ~ 1 + Age_first * Treatment + Age_second + (1 + Treatment | Subject) + (1 + Age_first | Experimenter) + log(performance_sd) ~ 1 + AGE * BMI + max(Age, BMI) + (1 + Age * BMI | Subject) + Performance ~ Normal(performance_mean, performance_sd) + @defaults begin + gr(Subject, by=ClinicalGroup, + Block1=>[Treatment, Age:BMI], + Block2=>[Age, BMI] + ) + end +end \ No newline at end of file From 6baf94f58950c9380de51ae88ba64bfc2f32ba95 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Wed, 4 Mar 2026 13:12:51 +0100 Subject: [PATCH 03/23] first parsing prototype --- scripts/macro.jl | 132 ++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 130 insertions(+), 2 deletions(-) diff --git a/scripts/macro.jl b/scripts/macro.jl index c053c3f..dd9bc50 100644 --- a/scripts/macro.jl +++ b/scripts/macro.jl @@ -23,7 +23,7 @@ end # With a single dataframe, the model definition could look as follows (which would lower to the "obvious" syntax): model = @brm begin - BMI ~ Normal(BMI_measured, 1) + # BMI ~ Normal(BMI_measured, 1) Age_first, Age_second = ploynomial_expand(Age; order=2) performance_mean ~ 1 + Age_first * Treatment + Age_second + (1 + Treatment | Subject) + (1 + Age_first | Experimenter) log(performance_sd) ~ 1 + AGE * BMI + max(Age, BMI) + (1 + Age * BMI | Subject) @@ -34,4 +34,132 @@ model = @brm begin Block2=>[Age, BMI] ) end -end \ No newline at end of file +end + + +macro brm(x) + esc(_brm(x)) +end +begin +function ensurecols end +function maybedists end +# isxcall(x) = Meta.isexpr(x, :call) +isxcall(x, f) = Meta.isexpr(x, :call) && x.args[1] == f +fixcall(x) = x +fixcall(x::Expr) = if Meta.isexpr(x, :call) + f = x.args[1] + pargs = [] + args = [] + for arg in fixcall.(x.args[2:end]) + if Meta.isexpr(arg, :parameters) + append!(pargs, arg.args) + else + push!(args, arg) + end + end + if length(pargs) > 0 + Expr(x.head, f, Expr(:parameters, pargs...), args...) + else + Expr(x.head, f, args...) + end +else + Expr(x.head, fixcall.(x.args)...) +end +xensurecols(x::Expr) = if x.head == :call + Expr(:call, ensurecols, x.args...) |> fixcall +else + dump(x) + error("Don't know how to handle xensurecols($x)!") +end + parse!(x::LineNumberNode; info) = x +parse!(x::Expr; info) = if x.head == :block + Expr(:block, parse!.(x.args; info)...) +elseif x.head == :(=) + lhs, rhs = x.args + lhs = parse_assignment_lhs!(lhs; info) + rhs = parse_assignment_rhs!(rhs; info) + Expr(:(=), lhs, xensurecols(rhs)) +elseif isxcall(x, :~) + _, lhs, rhs = x.args + lhs = parse_sampling_lhs!(lhs; info) + rhs = parse_sampling_rhs!(rhs; info) + if isxcall(rhs, ensurecols) + :($lhs = $maybedists(;force=isdata($lhs))($(rhs.args[2:end]...))) + else + :($lhs = $maybedists(;force=isdata($lhs))($(rhs))) + end +else + dump(x) + error("Don't know how to handle parse!($x)!") +end +parse_assignment_lhs!(x::Symbol; info) = (get!(info.alllocals, x, :local); x) +parse_assignment_lhs!(x::Expr; info) = begin + @assert Meta.isexpr(x, (:tuple, :vect)) + args = parse_assignment_lhs!.(x.args; info) + Expr(x.head, args...) +end +parse_assignment_rhs!(x::Symbol; info) = (get!(info.alllocals, x, :nonlocal); x) +parse_assignment_rhs!(x::Expr; info) = if x.head == :call + Expr(:call, x.args[1], parse_assignment_rhs!.(x.args[2:end]; info)...) +elseif Meta.isexpr(x, :parameters) + x +else + dump(x) + error("Don't know how to handle parse_assignment_rhs!($x)!") +end +parse_sampling_lhs!(x::Symbol; info) = (get!(info.alllocals, x, :maybelocal); x) +parse_sampling_lhs!(x::Expr; info) = if x.head == :call + Expr(:call, x.args[1], parse_sampling_lhs!.(x.args[2:end]; info)...) +else + @assert Meta.isexpr(x, (:tuple, :vect)) + args = parse_sampling_lhs!.(x.args; info) + Expr(x.head, args...) +end +parse_sampling_rhs!(x::Number; info) = x +parse_sampling_rhs!(x::Symbol; info) = (get!(info.alllocals, x, :nonlocal); x) +parse_sampling_rhs!(x::Expr; info) = if x.head == :call + if x.args[1] in (:+, :*, :|) + Expr(:call, x.args[1], parse_sampling_rhs!.(x.args[2:end]; info)...) + else + xensurecols(parse_assignment_rhs!(x; info)) + end +elseif Meta.isexpr(x, :parameters) + x +else + dump(x) + error("Don't know how to handle parse_sampling_rhs!($x)!") +end +using OrderedCollections +_brm(x::Expr) = begin + @assert x.head == :block + alllocals = OrderedDict{Symbol,Symbol}() + info = (;alllocals) + x = parse!(x; info) + nonlocals = [key for (key, value) in pairs(alllocals) if value == :nonlocal] + maybelocals = [key for (key, value) in pairs(alllocals) if value == :maybelocal] + locals = [key for (key, value) in pairs(alllocals) if value == :local] + init = quote + (;$(nonlocals...)) = data(__df__) + (;$(maybelocals...)) = maybedata(__df__) + end + finalize = :(BRM(;$(keys(alllocals)...))) + Expr(:(=), :(model(__df__)), Expr(:block, init, x.args..., finalize)) +end +end +@macroexpand @brm begin + Age_first, Age_second = ploynomial_expand(Age; order=2) + performance_mean ~ 1 + Age_first * Treatment + Age_second + (1 + Treatment | Subject) + (1 + Age_first | Experimenter) + log(performance_sd) ~ 1 + Age * BMI + max(Age, BMI) + (1 + Age * BMI | Subject) + Performance ~ Normal(performance_mean, performance_sd) +end + + +# model(__df__) = begin +# (; Age, Treatment, Subject, Experimenter, BMI) = data(__df__) +# (; performance_mean, performance_sd, Performance) = maybedata(__df__) +# (Age_first, Age_second) = (ensurecols)(ploynomial_expand, Age; order=2) +# performance_mean = ((maybedists)(; force=isdata(performance_mean)))(1 + Age_first * Treatment + Age_second + ((1 + Treatment) | Subject) + ((1 + Age_first) | Experimenter)) +# log(performance_sd) = ((maybedists)(; force=isdata(log(performance_sd))))(1 + Age * BMI + (ensurecols)(max, Age, BMI) + ((1 + Age * BMI) | Subject)) +# Performance = ((maybedists)(; force=isdata(Performance)))(Normal, performance_mean, performance_sd) +# BRM(; Age_first, Age_second, Age, performance_mean, Treatment, Subject, Experimenter, performance_sd, BMI, Performance) +# end \ No newline at end of file From 799de284eba1791212fd915d087a979175bc5d8c Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Wed, 11 Mar 2026 11:27:10 +0100 Subject: [PATCH 04/23] add version which isn't doing the right thing and is 10 times slower than brms for the primal computation and 20 times (Enzyme) or 30 times (Mooncake) slower for the gradient computation --- scripts/Benchmarking/Project.toml | 6 + scripts/Benchmarking/main.jl | 90 ++++++++++ scripts/Project.toml | 11 +- scripts/examples/database.jl | 12 ++ scripts/macro.jl | 290 ++++++++++++++++++------------ scripts/vimpl.jl | 198 ++++++++++++++++++++ 6 files changed, 489 insertions(+), 118 deletions(-) create mode 100644 scripts/Benchmarking/Project.toml create mode 100644 scripts/Benchmarking/main.jl create mode 100644 scripts/examples/database.jl create mode 100644 scripts/vimpl.jl diff --git a/scripts/Benchmarking/Project.toml b/scripts/Benchmarking/Project.toml new file mode 100644 index 0000000..ddcc5ec --- /dev/null +++ b/scripts/Benchmarking/Project.toml @@ -0,0 +1,6 @@ +[deps] +Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" diff --git a/scripts/Benchmarking/main.jl b/scripts/Benchmarking/main.jl new file mode 100644 index 0000000..ef892d1 --- /dev/null +++ b/scripts/Benchmarking/main.jl @@ -0,0 +1,90 @@ +# Gemini claims you can do this! +using Pkg +Pkg.activate(@__DIR__) +insert!(LOAD_PATH, 2, joinpath(@__DIR__, "..")) + +include("../macro.jl") +include("../vimpl.jl") +include("../examples/database.jl") + +using Chairmarks, Random + +db = Database() +df = db.dataset(:bambi, :escs) +fdf = map(Vector{Float64}, (;df.drugs, df.o, df.c, df.e, df.a, df.n)) +# Using fdf is considerably faster than using df +tmp = @brm fdf """ + loc ~ o + c + e + a + n + log(err_scale) ~ 1 + drugs ~ Normal(loc, err_scale) +""" +display(@brm """ + loc ~ o + c + e + a + n + log(err_scale) ~ 1 + drugs ~ Normal(loc, err_scale) +""") +vtmp = VBRMI(tmp) +display(vtmp) +display(@be randn(LogDensityProblems.dimension(vtmp)) LogDensityProblems.logdensity($vtmp, _)) + +# import Reactant + +# rx = Reactant.to_rarray(randn(LogDensityProblems.dimension(vtmp))) +# errors: NoFieldMatchError(...) +# _rtmp = Reactant.to_rarray(vtmp) +# errors: MethodError: no method matching _copyto!(::SubArray{…}, ::Base.Broadcast.Broadcasted{…}) +# rtmp = Reactant.@compile LogDensityProblems.logdensity(vtmp, rx) +# display(@be randn(LogDensityProblems.dimension(vtmp)) LogDensityProblems.logdensity($rtmp, _)) +# error() + +using LogDensityProblemsAD, Mooncake, Enzyme, DifferentiationInterface +struct ADLogDensity{F, B, E} + f::F + backend::B + extras::E +end + +ADLogDensity(f, backend) = ADLogDensity( + f, + backend, + DifferentiationInterface.prepare_gradient( + Base.Fix1(LogDensityProblems.logdensity, f), + backend, + zeros(LogDensityProblems.dimension(f)) + ) +) + +LogDensityProblems.capabilities(::ADLogDensity) = LogDensityProblems.LogDensityOrder{1}() +LogDensityProblems.dimension(p::ADLogDensity) = LogDensityProblems.dimension(p.f) +LogDensityProblems.logdensity(p::ADLogDensity, x) = LogDensityProblems.logdensity(p.f, x) +LogDensityProblems.logdensity_and_gradient( + p::ADLogDensity, x +) = DifferentiationInterface.value_and_gradient( + Base.Fix1(LogDensityProblems.logdensity, p.f), p.extras, p.backend, x +) + + + +mvtmp1 = ADgradient(AutoMooncake(), vtmp) +mvtmp2 = ADLogDensity(vtmp, AutoMooncake()) +evtmp1 = ADgradient(AutoEnzyme(; mode = Enzyme.set_runtime_activity(Enzyme.Reverse)), vtmp) +evtmp2 = ADLogDensity(vtmp, AutoEnzyme(; + mode=Enzyme.set_runtime_activity(Enzyme.Reverse), + function_annotation=Enzyme.Duplicated +)) + +x = randn(LogDensityProblems.dimension(vtmp)) +display(mapreduce(hcat, (mvtmp1, mvtmp2, evtmp1, evtmp2)) do b + LogDensityProblems.logdensity_and_gradient(b, copy(x))[2] +end) + +@info "Primal" +display(@be randn(LogDensityProblems.dimension(vtmp)) LogDensityProblems.logdensity($vtmp, _)) +@info "Inefficient Mooncake skipped..." +# display(@be randn(LogDensityProblems.dimension(vtmp)) LogDensityProblems.logdensity_and_gradient($mvtmp1, _)) +@info "Mooncake" +display(@be randn(LogDensityProblems.dimension(vtmp)) LogDensityProblems.logdensity_and_gradient($mvtmp2, _)) +@warn "wrong Enzyme skipped..." +# display(@be randn(LogDensityProblems.dimension(vtmp)) LogDensityProblems.logdensity_and_gradient($evtmp1, _)) +@info "correct Enzyme" +display(@be randn(LogDensityProblems.dimension(vtmp)) LogDensityProblems.logdensity_and_gradient($evtmp2, _)) \ No newline at end of file diff --git a/scripts/Project.toml b/scripts/Project.toml index 4a3f7db..d5ac59f 100644 --- a/scripts/Project.toml +++ b/scripts/Project.toml @@ -1,9 +1,16 @@ [deps] BayesianRegressionModels = "cdd3e328-398d-47a6-a87b-6047aaf4b4bc" -CatalogServer = "a1b2c3d4-0000-0000-0000-000000000001" CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b" +Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +DynamicObjects = "23d02862-63fe-4c6e-8fdb-1d52cbbd39d5" +ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4" FlexiChains = "4a37a8b9-6e57-4b92-8664-298d46e639f7" +InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" +LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" +LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1" +LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Oxygen = "df9a0d86-3283-4920-82dc-4555fc0d1d8b" -YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" diff --git a/scripts/examples/database.jl b/scripts/examples/database.jl new file mode 100644 index 0000000..a98f505 --- /dev/null +++ b/scripts/examples/database.jl @@ -0,0 +1,12 @@ +include("all.jl") +using DynamicObjects + +@dynamicstruct "" serial struct Database + mod(m::Symbol) = _SOURCE_MODULES[m] + @cached dataset(m::Symbol, name::Symbol) = getproperty(mod(m), :load)(Val(name)) + # formula(m::Symbol, name::Symbol) = getproperty(mod(m), :examples)(Val(name))[1] + # example(m::Symbol, name::Symbol) = begin + # _formula, datakey = getproperty(mod(m), :examples)(Val(name)) + # brm(dataset(m, datakey), _formula) + # end +end \ No newline at end of file diff --git a/scripts/macro.jl b/scripts/macro.jl index dd9bc50..d6baf3c 100644 --- a/scripts/macro.jl +++ b/scripts/macro.jl @@ -1,49 +1,26 @@ +using OrderedCollections macro brm(x) - dump(x) + esc(_brm(x)) end -# The example BRM model is supposed to work with three dataframes (dBMI, dpmean, dpsd), which IIUC, have to have the same numbers of rows in this example and could have been merged into one. -@brm model(dBMI, dpmean, dpsd) = begin - # BMI will be a model parameter - BMI ~ Normal(dBMI.BMI_measured, 1) # equivalently: BMI ~ Normal(BMI_measured, 1) |> (data=dBMI) - # Age_first, Age_second are functions of data, and would be computed/updated exactly once - Age_first, Age_second = ploynomial_expand(dpmean.Age; order=2) # equivalently: Age_first, Age_second = ploynomial_expand(Age; order=2) |> (data=dpmean) - # I'm assuming performance_mean is a function of data and model parameters - I think it kind of has to be - performance_mean ~ 1 + Age_first * Treatment + Age_second + (1 + Treatment | Subject) + (1 + Age_first | Experimenter) |> (data=dpmean) - # I'm assuming performance_sd is a function of data and model parameters - I think it kind of has to be - log(performance_sd) ~ 1 + AGE * BMI + max(Age, BMI) + (1 + Age * BMI | Subject) |> (data=dpsd) - # Peter didn't specify data here - but I think it would have to be specified? Or would it be added to the model via conditioning syntax? - Performance ~ Normal(performance_mean, performance_sd) # |> (data=observations_df) as an example - @defaults begin - gr(Subject, by=ClinicalGroup, - Block1=>[Treatment, Age:BMI], - Block2=>[Age, BMI] - ) - end +macro brm(df, x) + esc(Expr(:call, _brm(x), df)) end - -# With a single dataframe, the model definition could look as follows (which would lower to the "obvious" syntax): -model = @brm begin - # BMI ~ Normal(BMI_measured, 1) - Age_first, Age_second = ploynomial_expand(Age; order=2) - performance_mean ~ 1 + Age_first * Treatment + Age_second + (1 + Treatment | Subject) + (1 + Age_first | Experimenter) - log(performance_sd) ~ 1 + AGE * BMI + max(Age, BMI) + (1 + Age * BMI | Subject) - Performance ~ Normal(performance_mean, performance_sd) - @defaults begin - gr(Subject, by=ClinicalGroup, - Block1=>[Treatment, Age:BMI], - Block2=>[Age, BMI] - ) - end +macro n(x) + esc(_n(x)) end - - -macro brm(x) - esc(_brm(x)) +macro x(x) + esc(_x(x)) +end +macro getproperty(x) + esc(_getproperty(x)) +end +_getproperty(x::Expr) = begin + @assert x.head == :(.) + @assert length(x.args) == 2 + lhs, qrhs = x.args + :(hasproperty($lhs, $qrhs) ? $x : $(qrhs.value)) end begin -function ensurecols end -function maybedists end -# isxcall(x) = Meta.isexpr(x, :call) isxcall(x, f) = Meta.isexpr(x, :call) && x.args[1] == f fixcall(x) = x fixcall(x::Expr) = if Meta.isexpr(x, :call) @@ -65,101 +42,182 @@ fixcall(x::Expr) = if Meta.isexpr(x, :call) else Expr(x.head, fixcall.(x.args)...) end -xensurecols(x::Expr) = if x.head == :call - Expr(:call, ensurecols, x.args...) |> fixcall -else - dump(x) - error("Don't know how to handle xensurecols($x)!") +function assign end +function doublepipe end +function gr end +_brm(x::AbstractString; kwargs...) = _brm(Meta.parse(""" +begin + $x end - parse!(x::LineNumberNode; info) = x +"""); kwargs...) +_brm(x::Expr; df=nothing) = begin + lhs, x = x.head == :(=) ? x.args : (:($(gensym("model"))(__df__)), x) + alllocals = OrderedDict{Symbol,Symbol}() + info = (;alllocals) + x = parse!(x; info) + nonlocals = [key for (key, value) in pairs(alllocals) if value == :nonlocal] + maybelocals = [key for (key, value) in pairs(alllocals) if value == :maybelocal] + init = quote + (;$(nonlocals...)) = data(__df__) + (;$(maybelocals...)) = maybedata(__df__) + end + finalize = quote + $BRMI(;$(keys(alllocals)...)) + end + if isnothing(df) + Expr(:(=), lhs, Expr(:block, init, x, finalize)) + else + Expr(:let, + Expr(:block, :(__df__ = $df), :(__ddf__ = $data(__df__)), [:($nonlocal = @getproperty __ddf__.$nonlocal) for nonlocal in nonlocals]...), + Expr(:block, :((;$(maybelocals...)) = maybedata(__df__)), x, finalize) + ) + end +end +brm(df, formula::AbstractString) = eval(_brm(formula; df)) +parse!(x; info) = x parse!(x::Expr; info) = if x.head == :block Expr(:block, parse!.(x.args; info)...) elseif x.head == :(=) lhs, rhs = x.args - lhs = parse_assignment_lhs!(lhs; info) - rhs = parse_assignment_rhs!(rhs; info) - Expr(:(=), lhs, xensurecols(rhs)) + parselocals!(rhs; info, val=:nonlocal) + parselocals!(lhs; info, val=:local) + :(@n $lhs = @x $assign($(xname(lhs)), $rhs)) elseif isxcall(x, :~) _, lhs, rhs = x.args - lhs = parse_sampling_lhs!(lhs; info) - rhs = parse_sampling_rhs!(rhs; info) - if isxcall(rhs, ensurecols) - :($lhs = $maybedists(;force=isdata($lhs))($(rhs.args[2:end]...))) - else - :($lhs = $maybedists(;force=isdata($lhs))($(rhs))) - end + parselocals!(rhs; info, val=:nonlocal) + parselocals!(lhs; info, val=:maybelocal) + :(@n $lhs = @x $x) else dump(x) error("Don't know how to handle parse!($x)!") end -parse_assignment_lhs!(x::Symbol; info) = (get!(info.alllocals, x, :local); x) -parse_assignment_lhs!(x::Expr; info) = begin - @assert Meta.isexpr(x, (:tuple, :vect)) - args = parse_assignment_lhs!.(x.args; info) - Expr(x.head, args...) -end -parse_assignment_rhs!(x::Symbol; info) = (get!(info.alllocals, x, :nonlocal); x) -parse_assignment_rhs!(x::Expr; info) = if x.head == :call - Expr(:call, x.args[1], parse_assignment_rhs!.(x.args[2:end]; info)...) -elseif Meta.isexpr(x, :parameters) - x +parselocals!(x; kwargs...) = x +parselocals!(x::Symbol; info, val) = get!(info.alllocals, x, val) +parselocals!(x::Expr; info, val) = if Meta.isexpr(x, (:call, :kw)) + parselocals!.(x.args[2:end]; info, val) else - dump(x) - error("Don't know how to handle parse_assignment_rhs!($x)!") -end -parse_sampling_lhs!(x::Symbol; info) = (get!(info.alllocals, x, :maybelocal); x) -parse_sampling_lhs!(x::Expr; info) = if x.head == :call - Expr(:call, x.args[1], parse_sampling_lhs!.(x.args[2:end]; info)...) -else - @assert Meta.isexpr(x, (:tuple, :vect)) - args = parse_sampling_lhs!.(x.args; info) - Expr(x.head, args...) -end -parse_sampling_rhs!(x::Number; info) = x -parse_sampling_rhs!(x::Symbol; info) = (get!(info.alllocals, x, :nonlocal); x) -parse_sampling_rhs!(x::Expr; info) = if x.head == :call - if x.args[1] in (:+, :*, :|) - Expr(:call, x.args[1], parse_sampling_rhs!.(x.args[2:end]; info)...) + parselocals!.(x.args; info, val) +end +_n(x::Expr) = begin + @assert x.head == :(=) + lhs, rhs = x.args + alhs = xassignable(lhs) + nlhs = xname(alhs) + :($alhs = $NamedColumn($nlhs, $rhs)) +end +xassignable(x::Symbol) = x +xassignable(x::Expr) = if Meta.isexpr(x, (:tuple, :vect)) + Expr(x.head, xassignable.(x.args)...) +elseif x.head == :call + if length(x.args) == 2 + xassignable(x.args[2]) else - xensurecols(parse_assignment_rhs!(x; info)) + @warn "Don't know how to handle xassignable($x)!" + Symbol(x) end -elseif Meta.isexpr(x, :parameters) - x else dump(x) - error("Don't know how to handle parse_sampling_rhs!($x)!") + error("Don't know how to handle xassignable($x)!") end -using OrderedCollections -_brm(x::Expr) = begin - @assert x.head == :block - alllocals = OrderedDict{Symbol,Symbol}() - info = (;alllocals) - x = parse!(x; info) - nonlocals = [key for (key, value) in pairs(alllocals) if value == :nonlocal] - maybelocals = [key for (key, value) in pairs(alllocals) if value == :maybelocal] - locals = [key for (key, value) in pairs(alllocals) if value == :local] - init = quote - (;$(nonlocals...)) = data(__df__) - (;$(maybelocals...)) = maybedata(__df__) - end - finalize = :(BRM(;$(keys(alllocals)...))) - Expr(:(=), :(model(__df__)), Expr(:block, init, x.args..., finalize)) +xname(x::Symbol) = Meta.quot(x) +xname(x::Expr) = if Meta.isexpr(x, (:tuple, :vect)) + Expr(x.head, xname.(x.args)...) +else + @warn "Don't know how to handle xassignable($x)!" + Symbol(x) + # dump(x) + # error("Don't know how to handle xname($x)!") end +_x(x) = x +_x(x::Symbol) = x +_x(x::Expr) = if x.head == :call + Expr(:call, ExprColumn, _x.(x.args)...) |> fixcall +elseif x.head == :|| + Expr(:call, ExprColumn, doublepipe, _x.(x.args)...) +else + Expr(x.head, _x.(x.args)...) +end +struct Data{P} + parent::P +end +Base.parent(d::Data) = getfield(d, :parent) +Base.hasproperty(d::Data, x::Symbol) = hasproperty(parent(d), x) +Base.getproperty(d::Data, x::Symbol) = NamedColumn(x, DataColumn(getproperty(parent(d), x))) +data(x) = Data(x) +struct MaybeData{P} + parent::P +end +Base.parent(d::MaybeData) = getfield(d, :parent) +Base.hasproperty(d::MaybeData, x::Symbol) = hasproperty(parent(d), x) +Base.getproperty(d::MaybeData, x::Symbol) = NamedColumn(x, hasproperty(d, x) ? DataColumn(getproperty(parent(d), x)) : MissingColumn()) +maybedata(x) = MaybeData(x) +abstract type AbstractColumn end +struct MissingColumn <: AbstractColumn end +struct DataColumn{P} <: AbstractColumn + parent::P end -@macroexpand @brm begin - Age_first, Age_second = ploynomial_expand(Age; order=2) - performance_mean ~ 1 + Age_first * Treatment + Age_second + (1 + Treatment | Subject) + (1 + Age_first | Experimenter) - log(performance_sd) ~ 1 + Age * BMI + max(Age, BMI) + (1 + Age * BMI | Subject) - Performance ~ Normal(performance_mean, performance_sd) +Base.parent(d::DataColumn) = getfield(d, :parent) +struct NamedColumn{N,P} <: AbstractColumn + name::N + parent::P +end +name(x::NamedColumn) = getfield(x, :name) +Base.parent(x::NamedColumn) = getfield(x, :parent) + +struct ExprColumn{F,A<:Tuple,K<:NamedTuple} <: AbstractColumn + f::F + args::A + kwargs::K + ExprColumn(f, args...; kwargs...) = new{typeof(f),typeof(args),typeof((;kwargs...))}(f,args,(;kwargs...)) + ExprColumn(f::Type, args...; kwargs...) = new{Type{f},typeof(args),typeof((;kwargs...))}(f,args,(;kwargs...)) +end +getf(x::ExprColumn) = getfield(x, :f) +getargs(x::ExprColumn) = getfield(x, :args) +getargs(x::ExprColumn, n) = (rv = getargs(x); @assert length(rv) == n; rv) +getargs(::typeof(+), x::ExprColumn{typeof(+)}) = getargs(x) +getargs(::typeof(+), x::ExprColumn) = (x,) +getargs(::typeof(+), x) = (x,) +getkwargs(x::ExprColumn) = getfield(x, :kwargs) +getop(x) = getf(x) +getop(::ExprColumn{typeof(doublepipe)}) = :|| +getop(::ExprColumn{typeof(assign)}) = :(=) + +struct LikelihoodColumn{P,R} <: AbstractColumn + parent::P + rhs::R end +Base.parent(d::LikelihoodColumn) = getfield(d, :parent) +rhs(d::LikelihoodColumn) = getfield(d, :rhs) +maybedists(lhs::AbstractColumn, x::AbstractColumn) = LikelihoodColumn(lhs, x) +struct BRMI{O<:NamedTuple} + operations::O +end +BRMI(;kwargs...) = BRMI((;kwargs...)) +Base.show(io::IO, (;operations)::BRMI) = begin + print(io, "BRMI:\n") + for (key, value::NamedColumn) in pairs(operations) + print(io, " ", key, ": ", parent(value), "\n") + end +end +Base.show(io::IO, d::DataColumn) = begin + print(io, "data (eltype=", eltype(parent(d)), ")") +end +Base.show(io::IO, x::ExprColumn{<:Union{typeof.((~,*,+,|,doublepipe,assign))...}}) = begin + print(io, "(", ) + join(io, getargs(x), " $(getop(x)) ") + print(io, ")") +end +nonemptyjoin(io::IO, iterator, args...; first) = if length(iterator) > 0 + print(io, first) + join(io, iterator, args...) +end +Base.show(io::IO, x::ExprColumn) = begin + print(io, getf(x), "(", ) + join(io, getargs(x), ", ") + nonemptyjoin(io, ["$key=$value" for (key, value) in pairs(getkwargs(x))], ", "; first="; ") + print(io, ")") +end +Base.show(io::IO, x::NamedColumn) = print(io, name(x)) -# model(__df__) = begin -# (; Age, Treatment, Subject, Experimenter, BMI) = data(__df__) -# (; performance_mean, performance_sd, Performance) = maybedata(__df__) -# (Age_first, Age_second) = (ensurecols)(ploynomial_expand, Age; order=2) -# performance_mean = ((maybedists)(; force=isdata(performance_mean)))(1 + Age_first * Treatment + Age_second + ((1 + Treatment) | Subject) + ((1 + Age_first) | Experimenter)) -# log(performance_sd) = ((maybedists)(; force=isdata(log(performance_sd))))(1 + Age * BMI + (ensurecols)(max, Age, BMI) + ((1 + Age * BMI) | Subject)) -# Performance = ((maybedists)(; force=isdata(Performance)))(Normal, performance_mean, performance_sd) -# BRM(; Age_first, Age_second, Age, performance_mean, Treatment, Subject, Experimenter, performance_sd, BMI, Performance) -# end \ No newline at end of file +end \ No newline at end of file diff --git a/scripts/vimpl.jl b/scripts/vimpl.jl new file mode 100644 index 0000000..9c5bceb --- /dev/null +++ b/scripts/vimpl.jl @@ -0,0 +1,198 @@ +using LogExpFunctions, InverseFunctions, Distributions, ElasticArrays, LogDensityProblems, LinearAlgebra + +struct VBRMI{P<:BRMI,M<:NamedTuple} + parent::P + meta::M +end +VBRMI(p::BRMI) = VBRMI(p, finalize(foldl(vmeta, p.operations; init=(;materialized=(;), blocks=(;))))) +finalize(x) = merge(x, (;block_data=map(x.blocks) do values + m, n = size(values) + (;L=zeros(n, n)) +end)) +rmerge(x::NamedTuple, y::NamedTuple) = begin + xykeys = (intersect(keys(x), keys(y))...,) + merge(x, y, map(rmerge, NamedTuple{xykeys}(x), NamedTuple{xykeys}(y))) +end +rmerge(::AbstractDict, ::AbstractDict) = error("Can only rmerge NamedTuples for now!") +rmerge(x, y) = y +vmeta(meta, x::NamedColumn) = begin + meta, m = vmeta(meta, parent(x))::Tuple + rmerge(meta, (;materialized=(;name(x)=>m))) +end +vmeta(meta, x::DataColumn) = meta, x +vmeta(meta, x::ExprColumn{typeof(assign)}) = vmeta_assignment(meta, getargs(x)...; getkwargs(x)...) +vmeta(meta, x::ExprColumn{typeof(~)}) = vmeta_sampling(meta, getargs(x)...; getkwargs(x)...) + +vmeta_assignment(meta, ::Symbol, x) = meta, vmaterialize(vbroadcasted(x; meta)) +vbroadcasted(;kwargs...) = (args...)->vbroadcasted(args...; kwargs...) +vbroadcasted(x::NamedColumn{<:Any,<:DataColumn}; meta) = parent(meta.materialized[name(x)]) +vbroadcasted(x::NamedColumn; meta) = meta.materialized[name(x)] +vbroadcasted(x::ExprColumn; meta) = Base.broadcasted(getf(x), map(vbroadcasted(;meta), getargs(x))...) +getinverse(x::ExprColumn{<:Any,<:Tuple{<:Any}}) = inverse(getf(x)) +getinverse(x::ExprColumn{<:Any,<:Tuple{<:ExprColumn}}) = inverse(getf(x)) ∘ getinverse(getargs(x, 1)[1]) +vmeta_sampling(meta, lhs::ExprColumn, rhs) = begin + meta, o = vmeta_sampling_rhs(meta, rhs; group=:__population__) + meta, vmaterialize(Base.broadcasted(getinverse(lhs), o)) +end +vmeta_sampling(meta, ::NamedColumn{<:Any,MissingColumn}, rhs) = begin + meta, o = vmeta_sampling_rhs(meta, rhs; group=:__population__) + meta, vmaterialize(o) +end +vmeta_sampling(meta, lhs::NamedColumn{<:Any,<:DataColumn}, rhs) = begin + meta, o = vmeta_sampling_rhs(meta, rhs; group=:__population__) + meta, LikelihoodColumn(parent(parent(lhs)), o) +end +vmeta_sampling_rhs(;kwargs...) = (args...)->vmeta_sampling_rhs(args...; kwargs...) +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(+)}; kwargs...) = begin + meta, args = foldl(getargs(x); init=(meta, ())) do (_meta, _args), _arg + _meta, _arg = vmeta_sampling_rhs(_meta, _arg; kwargs...) + _meta, (_args..., _arg) + end + meta, Base.broadcasted(+, args...) +end +vmeta_sampling_rhs(meta, x::ExprColumn; kwargs...) = vmeta_sampling_rhs(meta, vbroadcasted(x; meta); kwargs...) +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(*)}; kwargs...) = error("NOT IMPLEMENTED") +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(&)}; kwargs...) = error("NOT IMPLEMENTED") +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(|)}; kwargs...) = begin + lhs, rhs = getargs(x, 2) + vmeta_sampling_rhs(meta, lhs; group=rhs) +end +vmeta_sampling_rhs(meta, ::Int; group) = begin + meta, p = growblock!!(meta, group, 1) + meta, p +end +vmeta_sampling_rhs(meta, x::NamedColumn; kwargs...) = vmeta_sampling_rhs(meta, meta.materialized[name(x)]; kwargs...) +vmeta_sampling_rhs(meta, x::DataColumn; kwargs...) = vmeta_sampling_rhs(meta, parent(x); kwargs...) +vmeta_sampling_rhs(meta, x::AbstractVector{<:AbstractFloat}; group) = begin + meta, p = growblock!!(meta, group, 1) + meta, Base.broadcasted(*, x, p) +end +FBroadcasted{F,Style<:Union{Nothing, Base.Broadcast.BroadcastStyle},Axes} = Base.Broadcast.Broadcasted{Style,Axes,F} +vmeta_sampling_rhs(meta, x::FBroadcasted; group) = begin + meta, p = growblock!!(meta, group, 1) + meta, Base.broadcasted(*, x, p) +end +vmeta_sampling_rhs(meta, x::AbstractVector{<:Integer}; group) = begin + meta, p = growblock!!(meta, group, 1) + meta, Base.broadcasted(*, x, p) + # meta, p = growblock!!(meta, group, length(unique(parent(x)))-1) + # Base.broadcasted(*, x, p) +end +vmeta_sampling_rhs(meta, x::FBroadcasted{<:Type{<:Distribution}}; group) = meta, x +vmaterialize(x) = MaterializedColumn(Base.materialize(x), x) +struct MaterializedColumn{P,B} <: AbstractColumn + parent::P + broadcast::B +end +Base.parent(x::MaterializedColumn) = getfield(x, :parent) +getbroadcast(x::MaterializedColumn) = getfield(x, :broadcast) +Base.broadcastable(x::MaterializedColumn) = Base.broadcastable(parent(x)) + +n_levels(group::NamedColumn) = length(unique(parent(parent(group)))) +growblock!!(meta, group::Symbol, n) = growblock!!(meta, group, 1, n) +growblock!!(meta, group::Symbol, m, n) = begin + g = get(meta.blocks, group) do + ElasticMatrix(zeros(m, 0)) + end + idxs = (size(g, 2)+1):(size(g, 2)+n) + append!(g, zeros(m, n)) + rmerge(meta, (;blocks=(;group=>g))), view(g, :, 1)#idxs) +end +growblock!!(meta, group::NamedColumn, n) = growblock!!(meta, name(group), n_levels(group), n) + +Base.show(io::IO, (;parent, broadcast)::MaterializedColumn) = print(io, eltype(parent), "[...] .= ", broadcast) +Base.show(io::IO, (;parent, rhs)::LikelihoodColumn) = print(io, eltype(parent), "[...] .~ ", rhs) +Base.show(io::IO, vbrm::VBRMI) = begin + (;parent, meta) = vbrm + print(io, parent) + print(io, "dim: ", LogDensityProblems.dimension(vbrm), "\n") + print(io, "materialized:\n") + for (key, value) in pairs(meta.materialized) + print(io, " ", key, ": ", value, "\n") + end + print(io, "blocks (n_levels, n_params):\n") + for (key, value) in pairs(meta.blocks) + print(io, " ", key, ": ", size(value), "\n") + end +end +LogDensityProblems.dimension(vbrm::VBRMI) = hyperdim(vbrm) + directdim(vbrm) +hyperdim(vbrm::VBRMI) = sum(pairs(vbrm.meta.blocks)) do (k, v) + n = size(v, 2) + k == :__population__ ? 0 : n * (n+1) ÷ 2 +end +directdim(vbrm::VBRMI) = sum(length, vbrm.meta.blocks) +advance!!(x, pos) = x[pos+1], pos+1 +advance!!(x, pos, n) = view(x, pos+1:pos+n), pos+n +lprior!((;meta)::VBRMI, x::AbstractVector; init=(0., 0)) = foldl(pairs(meta.blocks); init) do (lprior, pos), (key, values) + m, n = size(values) + if key == :__population__ + xi, pos = advance!!(x, pos, n) + values[1, :] .= xi + lprior += sum(Base.Fix1(logpdf, Normal()), xi) + else + C = LinearAlgebra.Cholesky(meta.block_data[key].L, :L, 0) + lprior, pos = lprior!(C, x; init=(lprior, pos)) + for vi in eachrow(values) + xi, pos = advance!!(x, pos, n) + mul!(vi, C.L, xi) + lprior += sum(Base.Fix1(logpdf, Normal()), xi) + end + end + lprior, pos +end |> first +log_abs_tanh(x) = begin + z = -2*abs(x) + (log1mexp(z) - log1pexp(z)) +end +log_square_tanh(x) = 2 * log_abs_tanh(x) +"Either wrong or better LKJCholesky unconstraining + prior" +lprior!((;L)::Cholesky, x; init, eta=1.) = begin + lprior, pos = init + n = LinearAlgebra.checksquare(L) + log_scale, pos = advance!!(x, pos) + lprior += logpdf(Normal(), log_scale) + L[1, 1] = exp(log_scale) + for i in 2:n + log_scale, pos = advance!!(x, pos) + lprior += logpdf(Normal(), log_scale) + xi, pos = advance!!(x, pos) + tmp = log_abs_tanh(xi / sqrt(n-1)) + L[i, 1] = sign(xi) * exp(log_scale + tmp) + log_sos = 2 * tmp + lprior += log1mexp(log_sos) + for j in 2:i-1 + xi, pos = advance!!(x, pos) + tmp1 = .5 * log1mexp(log_sos) + lprior += tmp1 + tmp2 = log_abs_tanh(xi / sqrt(n-j)) + lprior += log1mexp(2*tmp2) + tmp = tmp1 + tmp2 + L[i, j] = sign(xi) * exp(log_scale + tmp) + log_sos = logaddexp(log_sos, 2*tmp) + end + L[i, i] = exp(log_scale + .5 * log1mexp(log_sos)) + lprior += (n - i + 2*eta-2) * .5 * log1mexp(log_sos) + end + lprior, pos +end +llikelihood!((;meta)::VBRMI) = foldl(meta.materialized; init=0.) do llikelihood, m + llikelihood + llikelihood!(m) +end +llikelihood!(::DataColumn) = 0. +llikelihood!(x::MaterializedColumn) = (Base.materialize!(parent(x), getbroadcast(x)); 0.) +# llikelihood!(x::LikelihoodColumn) = sum(Base.broadcasted(logpdf, rhs(x), parent(x)); init=0.) +# The below is faster for some reason? +llikelihood!(x::LikelihoodColumn) = ssum(Base.broadcasted(logpdf, rhs(x), parent(x)); init=0.) +llikelihood!(x) = error(typeof(x)) + +ssum(args...; kwargs...) = sum(args...; kwargs...) +ssum(x::Base.Broadcast.Broadcasted; init) = begin + rv = init + for xi in x + rv += xi + end + rv +end + +Distributions.logpdf(vbrmi::VBRMI, x::AbstractVector) = lprior!(vbrmi, x) + llikelihood!(vbrmi) +LogDensityProblems.logdensity(vbrmi::VBRMI, x::AbstractVector) = lprior!(vbrmi, x) + llikelihood!(vbrmi) \ No newline at end of file From 31b5d9bf47626359fe875ca327d89d276f1a9dc9 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Wed, 11 Mar 2026 11:45:25 +0100 Subject: [PATCH 05/23] remove Reactant because it always errors... --- scripts/Benchmarking/Project.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/scripts/Benchmarking/Project.toml b/scripts/Benchmarking/Project.toml index ddcc5ec..529e53a 100644 --- a/scripts/Benchmarking/Project.toml +++ b/scripts/Benchmarking/Project.toml @@ -3,4 +3,3 @@ Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" -Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" From 00b5ccaf4a63b2403784a69c958e24642cfd7df0 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Thu, 9 Apr 2026 13:37:07 +0200 Subject: [PATCH 06/23] move @brm macro/vimpl into a new BRMMacroWeb package source tree The @brm macro and the VBRMI implementation now live alongside the new BRMMacroWeb module under web-macro/src/, so Revise tracks edits without a restart. scripts/Benchmarking/main.jl is updated to include them from the new location; the other (untracked) entry points still on disk are left to fix themselves up. Co-Authored-By: Claude Opus 4.6 (1M context) --- scripts/Benchmarking/main.jl | 4 ++-- {scripts => web-macro/src}/macro.jl | 0 {scripts => web-macro/src}/vimpl.jl | 0 3 files changed, 2 insertions(+), 2 deletions(-) rename {scripts => web-macro/src}/macro.jl (100%) rename {scripts => web-macro/src}/vimpl.jl (100%) diff --git a/scripts/Benchmarking/main.jl b/scripts/Benchmarking/main.jl index ef892d1..311759e 100644 --- a/scripts/Benchmarking/main.jl +++ b/scripts/Benchmarking/main.jl @@ -3,8 +3,8 @@ using Pkg Pkg.activate(@__DIR__) insert!(LOAD_PATH, 2, joinpath(@__DIR__, "..")) -include("../macro.jl") -include("../vimpl.jl") +include("../../web-macro/src/macro.jl") +include("../../web-macro/src/vimpl.jl") include("../examples/database.jl") using Chairmarks, Random diff --git a/scripts/macro.jl b/web-macro/src/macro.jl similarity index 100% rename from scripts/macro.jl rename to web-macro/src/macro.jl diff --git a/scripts/vimpl.jl b/web-macro/src/vimpl.jl similarity index 100% rename from scripts/vimpl.jl rename to web-macro/src/vimpl.jl From a20abd70ae4da483f16ca30f8f2b2f1446a81ee1 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Thu, 9 Apr 2026 13:37:27 +0200 Subject: [PATCH 07/23] add BRMMacroWeb app + fix vimpl growblock!! / categorical / random-effects MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BRMMacroWeb is a small HTMXObjects app that walks the @brm pipeline stage by stage (Meta.parse → parse! → _brm let-block → eval → VBRMI → Chairmarks benchmark) against a synthetic dataset. Each stage renders incrementally and the VBRMI stage runs a finite-difference gradient sanity check that flags any parameter that fails to influence the log density — useful for catching the kinds of bugs fixed below. vimpl.jl fixes: - growblock!! used to return view(g, :, 1), aliasing every parameter in a block to the first column. Now returns view(g, :, idxs) so each growblock!! call gets its own freshly-appended slot. - _cat_lookup / _cat_re_lookup add treatment-coded categorical predictors at population level and as random slopes inside a grouping factor (e.g. (cohort | group)). - _re_lookup / _gc_idx wire random-effects views back to a length-N per-row lookup keyed by the row's group code, so (... | group) formulas materialize correctly instead of erroring with a DimensionMismatch. The default formula in BRMMacroWeb exercises Normal (with distributional regression on the scale), Poisson, Binomial (using a plain positional `n` argument instead of brms's trials() sidecar) and Bernoulli likelihoods in one model, sharing random-effects blocks across linear predictors. Co-Authored-By: Claude Opus 4.6 (1M context) --- web-macro/Project.toml | 25 ++++ web-macro/app/Project.toml | 8 + web-macro/app/main.jl | 8 + web-macro/src/BRMMacroWeb.jl | 279 +++++++++++++++++++++++++++++++++++ web-macro/src/vimpl.jl | 67 +++++++-- 5 files changed, 376 insertions(+), 11 deletions(-) create mode 100644 web-macro/Project.toml create mode 100644 web-macro/app/Project.toml create mode 100644 web-macro/app/main.jl create mode 100644 web-macro/src/BRMMacroWeb.jl diff --git a/web-macro/Project.toml b/web-macro/Project.toml new file mode 100644 index 0000000..7f26610 --- /dev/null +++ b/web-macro/Project.toml @@ -0,0 +1,25 @@ +name = "BRMMacroWeb" +uuid = "cd55decf-6be8-4ea5-907a-b2dc35e4cc14" +version = "0.1.0" + +[deps] +Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" +DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" +DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +DynamicObjects = "23d02862-63fe-4c6e-8fdb-1d52cbbd39d5" +ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4" +FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +FlexiChains = "4a37a8b9-6e57-4b92-8664-298d46e639f7" +HTMX = "27f3e1ef-6ef8-44dc-9e9a-2fb23ed44e83" +HTMXObjects = "b12ef442-5798-4353-80f3-9562b03a0cb6" +InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" +LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" +PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +TestModules = "63c02187-99fd-4e5c-aaf0-4d6bfebc181c" +Treebars = "e1e568c4-3a56-40a4-95fa-9b9c6c16fccb" +Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" diff --git a/web-macro/app/Project.toml b/web-macro/app/Project.toml new file mode 100644 index 0000000..bced045 --- /dev/null +++ b/web-macro/app/Project.toml @@ -0,0 +1,8 @@ +[deps] +BRMMacroWeb = "cd55decf-6be8-4ea5-907a-b2dc35e4cc14" +DynamicObjects = "23d02862-63fe-4c6e-8fdb-1d52cbbd39d5" +HTMX = "27f3e1ef-6ef8-44dc-9e9a-2fb23ed44e83" +HTMXObjects = "b12ef442-5798-4353-80f3-9562b03a0cb6" +Revise = "295af30f-e4ad-537b-8983-00126c2a3abe" +TestModules = "63c02187-99fd-4e5c-aaf0-4d6bfebc181c" +Treebars = "e1e568c4-3a56-40a4-95fa-9b9c6c16fccb" diff --git a/web-macro/app/main.jl b/web-macro/app/main.jl new file mode 100644 index 0000000..7df0546 --- /dev/null +++ b/web-macro/app/main.jl @@ -0,0 +1,8 @@ +using Revise +using BRMMacroWeb + +begin + BRMMacroWeb.terminate() + port = length(ARGS) >= 1 ? parse(Int, ARGS[1]) : 8121 + BRMMacroWeb.serve(; host="0.0.0.0", revise=:lazy, port, async=true) +end diff --git a/web-macro/src/BRMMacroWeb.jl b/web-macro/src/BRMMacroWeb.jl new file mode 100644 index 0000000..f02bdc0 --- /dev/null +++ b/web-macro/src/BRMMacroWeb.jl @@ -0,0 +1,279 @@ +module BRMMacroWeb + +using HTMXObjects +using Random +using Chairmarks +using DataFrames +using FiniteDifferences: FiniteDifferences, central_fdm + +# The @brm macro and the VBRMI implementation live alongside this module so +# Revise tracks them. The scripts/ entry points (parsing.jl, Benchmarking, +# StanBlocksImpl) include them via relative paths into here. +include("macro.jl") +include("vimpl.jl") + +# ── Default formula + synthetic data ──────────────────────────────────────── + +default_formula() = """loc1 ~ 1 + a + c1 + (1 + b + c1 | g1) + (1 | g2) +log(err1) ~ 1 + d +y1 ~ Normal(loc1, err1) + +log_rate ~ 1 + a + (1 | g3) +k1 ~ Poisson(exp(log_rate)) + +log_odds_bin ~ 1 + c2 + (1 | g2) +bin_succ ~ Binomial(bin_n, logistic(log_odds_bin)) + +log_odds_b ~ 1 + b +bin_y ~ Bernoulli(logistic(log_odds_b)) +""" + +# Synthetic data shaped roughly like the benchmarking example: continuous +# covariates plus a categorical `group` column. vimpl.jl currently materializes +# population-level terms only, so the default formula stays scalar — but the +# `group` column is kept around so users can experiment with `(1 | group)` once +# vimpl.jl supports it. +function synthetic_df(; n=64, seed=1) + rng = Xoshiro(seed) + # Continuous covariates + a = randn(rng, n) + b = randn(rng, n) + c = randn(rng, n) + d = randn(rng, n) + # Grouping factors with different numbers of levels — use these on the + # right-hand side of `(... | gN)` to test multiple random-effects blocks. + g1 = repeat(1:8, inner=cld(n, 8))[1:n] # 8 levels + g2 = repeat(1:4, inner=cld(n, 4))[1:n] # 4 levels + g3 = rand(rng, 1:6, n) # 6 levels, unordered + # Categorical predictors (treatment-coded) with small level counts. + c1 = rand(rng, 1:3, n) + c2 = rand(rng, 1:2, n) + c3 = rand(rng, 1:4, n) + # Continuous outcomes + eta1 = 0.5 .+ 1.2 .* a .- 0.7 .* b .+ 0.3 .* c .+ 0.1 .* d + y1 = eta1 .+ 0.3 .* randn(rng, n) + y2 = -0.2 .+ 0.6 .* a .+ 0.4 .* b .+ 0.2 .* randn(rng, n) + # Non-negative integer (count) outcomes — Poisson likelihoods + k1 = rand.(rng, Distributions.Poisson.(exp.(0.5 .* eta1))) + k2 = rand.(rng, Distributions.Poisson.(exp.(0.3 .+ 0.4 .* a))) + # Binomial-likelihood pair: variable trial counts plus successes + # (the `trials(size)` brms sidecar is unnecessary in our DSL — `size` + # is just another positional argument to `Binomial`). + bin_n = rand(rng, 5:30, n) + bin_p_true = @. 1 / (1 + exp(-(0.2 + 0.5 * a))) + bin_succ = [rand(rng, Distributions.Binomial(n_i, p_i)) + for (n_i, p_i) in zip(bin_n, bin_p_true)] + # Bernoulli-likelihood column (0/1) — for hierarchical-Bernoulli models + # like Kruschke's `therapeutic_touch`. + bin_y = [rand(rng, Distributions.Bernoulli(p_i)) ? 1 : 0 for p_i in bin_p_true] + DataFrame(; a, b, c, d, g1, g2, g3, c1, c2, c3, + y1, y2, k1, k2, bin_n, bin_succ, bin_y) +end + +# ── Pipeline stages ───────────────────────────────────────────────────────── +# +# Each stage is computed lazily on demand so the user can stop at any +# intermediate step (parsing, transforming, wrapping, eval'ing, materializing, +# benchmarking) and inspect the result without paying for the later stages. + +const STAGES = ( + :parse, # Meta.parse(formula) + :transform, # parse!(...) — rewrites = and ~ into @n/@x macro calls + :wrap, # _brm(formula; df) — full let-block ready to eval + :brmi, # eval(...) — BRMI value + :vbrmi, # VBRMI(brmi) — materialized action with blocks/dim + :bench, # Chairmarks @be primal logdensity +) + +stage_index(s::Symbol) = something(findfirst(==(s), STAGES), length(STAGES)) + +# Run the pipeline up to (and including) `stage`. Returns a NamedTuple +# carrying every intermediate value computed so far. +function pipeline(formula::AbstractString, stage::Symbol) + s = stage_index(stage) + df = synthetic_df() + out = (; df) + + s >= 1 || return out + raw = Meta.parse("begin\n$formula\nend") + out = merge(out, (; raw)) + + s >= 2 || return out + alllocals = OrderedDict{Symbol,Symbol}() + transformed = parse!(deepcopy(raw); info=(;alllocals)) + out = merge(out, (; transformed, alllocals)) + + s >= 3 || return out + wrapped = _brm(formula; df) + out = merge(out, (; wrapped)) + + s >= 4 || return out + brmi = eval(wrapped) + out = merge(out, (; brmi)) + + s >= 5 || return out + vbrmi = VBRMI(brmi) + dim = LogDensityProblems.dimension(vbrmi) + x0 = randn(Xoshiro(0), dim) + ldp = try + string(LogDensityProblems.logdensity(vbrmi, x0)) + catch e + "error: " * sprint(showerror, e) + end + grad = try + FiniteDifferences.grad( + central_fdm(5, 1), + Base.Fix1(LogDensityProblems.logdensity, vbrmi), + x0, + )[1] + catch e + e + end + out = merge(out, (; vbrmi, dim, ldp, x0, grad)) + + s >= 6 || return out + bench = try + @be randn(dim) LogDensityProblems.logdensity($vbrmi, _) + catch e + e + end + merge(out, (; bench)) +end + +# ── Rendering helpers ─────────────────────────────────────────────────────── + +_section(title, body) = (h.h3(title), h.pre(body)) + +function render_output(formula::AbstractString; stage::Symbol=:vbrmi) + sections = Vector{Any}[] # one entry per stage; rendered most-recent-first + try + out = pipeline(formula, stage) + + # Synthetic data always pinned at the top, collapsed by default so the + # macro pipeline output stays the focus. + data_section = Any[ + h.details( + h.summary("Synthetic data ($(nrow(out.df)) rows × $(ncol(out.df)) cols: " * + join(string.(names(out.df)), ", ") * ") — click to expand"), + render_table(out.df; sortable=false), + ), + ] + + if haskey(out, :raw) + push!(sections, Any[_section("1. Meta.parse — raw Julia AST", + sprint(show, out.raw))...]) + end + if haskey(out, :transformed) + push!(sections, Any[ + _section("2. parse! — rewritten AST (= → @n/@x assign, ~ → @n/@x ~)", + sprint(show, out.transformed))..., + h.h3(" locals classified by parse!"), + h.pre(sprint(show, out.alllocals)), + ]) + end + if haskey(out, :wrapped) + push!(sections, Any[_section("3. _brm — full let-block (df spliced as a literal)", + sprint(show, out.wrapped))...]) + end + if haskey(out, :brmi) + push!(sections, Any[_section("4. eval — BRMI value (parsed model)", + sprint(show, out.brmi))...]) + end + if haskey(out, :vbrmi) + vbrmi_children = Any[ + _section("5. VBRMI — materialized action (blocks, dim, columns)", + sprint(show, out.vbrmi))..., + h.h3(" logdensity at a fixed random point"), + h.p("dim = ", string(out.dim), ", logdensity = ", out.ldp), + h.h3(" finite-difference gradient sanity check"), + ] + if out.grad isa Exception + push!(vbrmi_children, + h.pre("gradient error: " * sprint(showerror, out.grad))) + else + tol = 1e-8 + live = findall(>(tol) ∘ abs, out.grad) + dead = findall(<=(tol) ∘ abs, out.grad) + summary_color = isempty(dead) ? "green" : "crimson" + push!(vbrmi_children, h.p( + "active params: ", h.strong("$(length(live))/$(out.dim)"), + " — ", + h.span(; style="color:$summary_color")( + isempty(dead) ? + "all parameters influence the logdensity ✓" : + "$(length(dead)) dead param(s) at indices $(dead)" + ), + )) + push!(vbrmi_children, + h.pre(sprint(show, MIME"text/plain"(), out.grad))) + end + push!(sections, vbrmi_children) + end + if haskey(out, :bench) + bench_body = out.bench isa Exception ? + h.pre("benchmark error: " * sprint(showerror, out.bench)) : + h.pre(sprint(show, MIME"text/plain"(), out.bench)) + push!(sections, Any[h.h3("6. Chairmarks @be — primal logdensity"), bench_body]) + end + + # Stages render most-recent-first; synthetic data sits at the very top. + children = reduce(vcat, reverse(sections); init=Any[]) + prepend!(children, data_section) + return h.div(; id="brm-macro-output")(children...) + catch e + return h.div(; id="brm-macro-output")( + h.h3("Error"), + h.pre(sprint(showerror, e, catch_backtrace())), + ) + end +end + +# ── Routes ────────────────────────────────────────────────────────────────── + +_stage_button(label, stage) = h.button(label; type="button", + hx_get="/stage/$stage", + hx_include="#brm-macro-form", + hx_target="#brm-macro-output", + hx_swap="outerHTML") + +@htmx struct AppContext + req = nothing + + @get index(; formula::String=default_formula()) = htmx(h.main(class="container")( + h.h1("BRM macro action"), + h.p( + "Enter a ", h.code("@brm"), " formula and step through the macro pipeline: ", + h.code("Meta.parse"), " → ", h.code("parse!"), " → ", h.code("_brm"), + " let-block → ", h.code("eval"), " → ", h.code("VBRMI"), " action → ", + h.code("Chairmarks"), " benchmark.", + ), + h.form(; id="brm-macro-form")( + h.label("Formula")( + h.textarea(formula; + name="formula", rows=8, + style="width:100%;font-family:monospace"), + ), + h.fieldset(; class="grid")( + _stage_button("1. Parse", :parse), + _stage_button("2. Transform", :transform), + _stage_button("3. Wrap", :wrap), + _stage_button("4. BRMI", :brmi), + _stage_button("5. VBRMI", :vbrmi), + _stage_button("6. Benchmark", :bench), + ), + ), + render_output(formula), + ); pico_version="2", extra_head=( + h.title("BRM macro action"), + h.style(":root { font-size: 87.5%; }"), + )) + + @get stage(name::AbstractString; formula::String=default_formula()) = + render_output(formula; stage=Symbol(name)) +end + +function __init__() + route!(AppContext()) +end + +end # module diff --git a/web-macro/src/vimpl.jl b/web-macro/src/vimpl.jl index 9c5bceb..bfc51f3 100644 --- a/web-macro/src/vimpl.jl +++ b/web-macro/src/vimpl.jl @@ -59,25 +59,70 @@ vmeta_sampling_rhs(meta, x::ExprColumn{typeof(|)}; kwargs...) = begin end vmeta_sampling_rhs(meta, ::Int; group) = begin meta, p = growblock!!(meta, group, 1) - meta, p -end + meta, _re_lookup(p, _gc_idx(group)) +end vmeta_sampling_rhs(meta, x::NamedColumn; kwargs...) = vmeta_sampling_rhs(meta, meta.materialized[name(x)]; kwargs...) vmeta_sampling_rhs(meta, x::DataColumn; kwargs...) = vmeta_sampling_rhs(meta, parent(x); kwargs...) vmeta_sampling_rhs(meta, x::AbstractVector{<:AbstractFloat}; group) = begin meta, p = growblock!!(meta, group, 1) - meta, Base.broadcasted(*, x, p) -end + meta, Base.broadcasted(*, x, _re_lookup(p, _gc_idx(group))) +end FBroadcasted{F,Style<:Union{Nothing, Base.Broadcast.BroadcastStyle},Axes} = Base.Broadcast.Broadcasted{Style,Axes,F} vmeta_sampling_rhs(meta, x::FBroadcasted; group) = begin meta, p = growblock!!(meta, group, 1) - meta, Base.broadcasted(*, x, p) -end + meta, Base.broadcasted(*, x, _re_lookup(p, _gc_idx(group))) +end vmeta_sampling_rhs(meta, x::AbstractVector{<:Integer}; group) = begin - meta, p = growblock!!(meta, group, 1) - meta, Base.broadcasted(*, x, p) - # meta, p = growblock!!(meta, group, length(unique(parent(x)))-1) - # Base.broadcasted(*, x, p) + # Treatment-coded categorical predictor: level 1 is the reference (drops out), + # remaining k-1 levels each get their own coefficient slot in the block. + # TODO: figure out where to cache `levels`/`level_map`/`dense`/`gc_idx` so we + # don't rebuild them on every VBRMI construction. Candidates: a new + # `meta.factor` NamedTuple keyed by column name, or attach it to the + # materialized entry for the source column — _gc_idx() does the same + # dense-mapping work for the grouping-factor side. + # TODO: investigate hooking into an existing "categorical values vector" + # abstraction instead of building the dense map by hand. CategoricalArrays.jl + # (used by DataFrames) already exposes `levels`, `levelcode`, and a `.refs` + # field that *is* the dense Int8/16 code vector — if the input column is + # already a CategoricalVector we'd avoid the Dict round-trip entirely. Same + # idea applies to PooledArrays.jl. Need to decide whether vimpl.jl should + # take a hard dep on CategoricalArrays or sniff for the duck-typed interface. + levels = sort(unique(x)) + level_map = Dict(l => i for (i, l) in enumerate(levels)) + dense = [level_map[l] for l in x] + meta, p = growblock!!(meta, group, length(levels) - 1) + meta, _cat_broadcast(p, _gc_idx(group), dense) +end +# Population case (gc_idx === nothing): p is (1, k-1), look up the (level-1)-th +# column via linear indexing. +_cat_broadcast(p, ::Nothing, dense) = + Base.broadcasted(_cat_lookup, Ref(p), dense) +# Grouped case: p is (n_levels, k-1), look up p[gc_idx[i], level - 1]. +_cat_broadcast(p, gc_idx::AbstractVector{Int}, dense) = + Base.broadcasted(_cat_re_lookup, Ref(p), gc_idx, dense) +_cat_lookup(p, level) = level == 1 ? zero(eltype(p)) : p[level - 1] +_cat_re_lookup(p, gc, level) = level == 1 ? zero(eltype(p)) : p[gc, level - 1] + +# Group code index: maps each row of the data to its row in the block matrix +# `values`. Returns `nothing` for the special :__population__ marker so that the +# population-level path can pass through `_re_lookup` unchanged. +_gc_idx(::Symbol) = nothing +function _gc_idx(group::NamedColumn) + raw = parent(parent(group)) + levels = sort(unique(raw)) + level_map = Dict(l => i for (i, l) in enumerate(levels)) + [level_map[l] for l in raw] end + +# Wrap a (m, 1) growblock view as a length-N broadcasted lookup keyed by +# `gc_idx`. For population groups (`gc_idx === nothing`) the (1, 1) view +# already broadcasts as a scalar against length-N data, so it passes through +# unchanged. Only handles n=1 (single column per growblock!! call) — multi-term +# random specs like `(1 + x | group)` decompose into separate growblock!! +# calls inside the parent `+` foldl, so each call hits this with n=1. +_re_lookup(p, ::Nothing) = p +_re_lookup(p, gc_idx::AbstractVector{Int}) = + Base.broadcasted(getindex, Ref(p), gc_idx, 1) vmeta_sampling_rhs(meta, x::FBroadcasted{<:Type{<:Distribution}}; group) = meta, x vmaterialize(x) = MaterializedColumn(Base.materialize(x), x) struct MaterializedColumn{P,B} <: AbstractColumn @@ -96,7 +141,7 @@ growblock!!(meta, group::Symbol, m, n) = begin end idxs = (size(g, 2)+1):(size(g, 2)+n) append!(g, zeros(m, n)) - rmerge(meta, (;blocks=(;group=>g))), view(g, :, 1)#idxs) + rmerge(meta, (;blocks=(;group=>g))), view(g, :, idxs) end growblock!!(meta, group::NamedColumn, n) = growblock!!(meta, name(group), n_levels(group), n) From 2a7cc96f5a8fbe271a44bfeadc9fd6a74b8a43f6 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Fri, 10 Apr 2026 09:44:27 +0200 Subject: [PATCH 08/23] web app: styled BRMI/VBRMI cards, TODO page with file-backed state, safety whitelist BRMMacroWeb.jl: - Styled HTML rendering for BRMI/VBRMI cards: deterministic per-symbol colors (HSL from hash), data columns normal-weight, parameters bold, likelihood statements underlined. Replaces the plain-text `sprint(show)` rendering with a recursive `_html_expr` tree walker that produces colored nests. - VBRMI card shows BRMI-level symbolic expressions (from the parent BRMI) with type/shape annotations, not raw Broadcasted internals. - Logdensity + finite-difference gradient check collapsed into a single
with a green/red summary line. - TODO page (/todo) with nav_sidebar navigation (HTMXObjects `page` property for auto fragment/full-page wrapping). - File-backed TODO state: each item is a .jl file under web-macro/todos/ with `# key: value` header + `#= markdown =#` body + raw formula. Status (open/done/deprioritized) and formula edits persist to disk. - Per-todo article cards with status-colored left border, done/deprioritize pills, collapsible details for done/deprioritized items. - Inline pipeline results: "Try in pipeline" posts to /stage/vbrmi and swaps the VBRMI output into a div inside the card, no page navigation. - 25 TODO items across 3 tiers with verification formulas where applicable. - Formula safety whitelist: _check_formula_safety! walks the parsed AST before eval and rejects any function call not in _ALLOWED_CALLS (math, distributions, DSL operators). Blocks macros, shell commands, eval, include, run, ccall, etc. vimpl.jl: - vbroadcasted(::Number; meta) fallback so literal numbers (e.g. the 2 in a^2) pass through to Base.broadcasted as scalars. macro.jl: - _show_top renamed from _show_op, _leaf_column helper for walking through link-function wrappers to find the innermost NamedColumn. Co-Authored-By: Claude Opus 4.6 (1M context) --- web-macro/src/BRMMacroWeb.jl | 1099 ++++++++++++++++- web-macro/src/macro.jl | 14 +- web-macro/src/vimpl.jl | 3 + .../1.1-verify-bernoulli-binomial-done.jl | 15 + ...xposure-already-works-without-a-wrapper.jl | 21 + .../todos/1.3-i-expr-likely-already-works.jl | 16 + web-macro/todos/1.4-scale-x-standardize-x.jl | 28 + ...1.5-zerocorr-independent-random-effects.jl | 30 + ...1.6-cache-levels-level_map-dense-gc_idx.jl | 17 + ...egoricalarrays-pooledarrays-integration.jl | 19 + web-macro/todos/2.1-interactions-a-b-a-b.jl | 23 + ...onfigurable-categorical-reference-level.jl | 19 + .../todos/2.3-per-parameter-prior-scales.jl | 28 + ...ed-non-centered-parameterization-toggle.jl | 18 + ...uped-random-effects-per-factor-variance.jl | 23 + .../2.6-multi-membership-random-effects-mm.jl | 20 + ...r-meta-analysis-and-weighted-regression.jl | 39 + .../3.1-multivariate-outcomes-cbind-y1-y2.jl | 39 + ...dirichlet-process-non-parametric-models.jl | 18 + .../3.11-zero-inflated-hurdle-likelihoods.jl | 17 + ...nferred-predictors-measurement-error-me.jl | 20 + ...ordinal-predictors-mo-monotonic-effects.jl | 18 + .../3.4-ordinal-outcomes-proportional-odds.jl | 15 + web-macro/todos/3.5-mixture-models.jl | 18 + .../3.6-splines-gp-submodels-s-bs-gp-t2.jl | 18 + .../3.7-autoregressive-submodels-ar-ar1.jl | 18 + .../3.8-decompositions-qr-orthogonal-polar.jl | 15 + .../3.9-spike-and-slab-horseshoe-priors.jl | 17 + 28 files changed, 1589 insertions(+), 56 deletions(-) create mode 100644 web-macro/todos/1.1-verify-bernoulli-binomial-done.jl create mode 100644 web-macro/todos/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl create mode 100644 web-macro/todos/1.3-i-expr-likely-already-works.jl create mode 100644 web-macro/todos/1.4-scale-x-standardize-x.jl create mode 100644 web-macro/todos/1.5-zerocorr-independent-random-effects.jl create mode 100644 web-macro/todos/1.6-cache-levels-level_map-dense-gc_idx.jl create mode 100644 web-macro/todos/1.7-categoricalarrays-pooledarrays-integration.jl create mode 100644 web-macro/todos/2.1-interactions-a-b-a-b.jl create mode 100644 web-macro/todos/2.2-configurable-categorical-reference-level.jl create mode 100644 web-macro/todos/2.3-per-parameter-prior-scales.jl create mode 100644 web-macro/todos/2.4-centered-non-centered-parameterization-toggle.jl create mode 100644 web-macro/todos/2.5-grouped-random-effects-per-factor-variance.jl create mode 100644 web-macro/todos/2.6-multi-membership-random-effects-mm.jl create mode 100644 web-macro/todos/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl create mode 100644 web-macro/todos/3.1-multivariate-outcomes-cbind-y1-y2.jl create mode 100644 web-macro/todos/3.10-dirichlet-process-non-parametric-models.jl create mode 100644 web-macro/todos/3.11-zero-inflated-hurdle-likelihoods.jl create mode 100644 web-macro/todos/3.2-inferred-predictors-measurement-error-me.jl create mode 100644 web-macro/todos/3.3-ordinal-predictors-mo-monotonic-effects.jl create mode 100644 web-macro/todos/3.4-ordinal-outcomes-proportional-odds.jl create mode 100644 web-macro/todos/3.5-mixture-models.jl create mode 100644 web-macro/todos/3.6-splines-gp-submodels-s-bs-gp-t2.jl create mode 100644 web-macro/todos/3.7-autoregressive-submodels-ar-ar1.jl create mode 100644 web-macro/todos/3.8-decompositions-qr-orthogonal-polar.jl create mode 100644 web-macro/todos/3.9-spike-and-slab-horseshoe-priors.jl diff --git a/web-macro/src/BRMMacroWeb.jl b/web-macro/src/BRMMacroWeb.jl index f02bdc0..97f46f7 100644 --- a/web-macro/src/BRMMacroWeb.jl +++ b/web-macro/src/BRMMacroWeb.jl @@ -49,6 +49,10 @@ function synthetic_df(; n=64, seed=1) c1 = rand(rng, 1:3, n) c2 = rand(rng, 1:2, n) c3 = rand(rng, 1:4, n) + # Positive exposure column — for Poisson-with-offset patterns where the + # log-rate gets a row-specific shift `+ log(exposure)` inside the + # likelihood expression directly (no `offset(...)` wrapper needed). + exposure = 0.5 .+ rand(rng, n) # Continuous outcomes eta1 = 0.5 .+ 1.2 .* a .- 0.7 .* b .+ 0.3 .* c .+ 0.1 .* d y1 = eta1 .+ 0.3 .* randn(rng, n) @@ -66,10 +70,103 @@ function synthetic_df(; n=64, seed=1) # Bernoulli-likelihood column (0/1) — for hierarchical-Bernoulli models # like Kruschke's `therapeutic_touch`. bin_y = [rand(rng, Distributions.Bernoulli(p_i)) ? 1 : 0 for p_i in bin_p_true] - DataFrame(; a, b, c, d, g1, g2, g3, c1, c2, c3, + DataFrame(; a, b, c, d, g1, g2, g3, c1, c2, c3, exposure, y1, y2, k1, k2, bin_n, bin_succ, bin_y) end +# ── Formula safety whitelist ──────────────────────────────────────────────── +# +# The formula textarea accepts arbitrary text that gets `Meta.parse`d then +# `eval`'d. To prevent arbitrary code execution when the web app is shared, +# we walk the parsed AST *before* any eval and reject any function call or +# expression type that isn't on the allowlist. The allowlist is deliberately +# generous for formula-writing (math, distributions, data-column references) +# but blocks I/O, shell, eval, include, ccall, macros, etc. + +const _ALLOWED_CALLS = Set{Symbol}([ + # DSL operators + :~, :(+), :(-), :(*), :(/), :(^), :(|), :(||), + # Comparison (may appear in ifelse-style expressions) + :(==), :(!=), :(<), :(>), :(<=), :(>=), + # Math + :log, :log2, :log10, :log1p, :exp, :exp2, :expm1, + :sqrt, :cbrt, :abs, :abs2, :sign, :floor, :ceil, :round, + :sin, :cos, :tan, :asin, :acos, :atan, + :min, :max, :clamp, :mod, :rem, :div, + :logistic, :logit, :softmax, :logsumexp, + :log_abs_tanh, :log_square_tanh, + # Distributions (Type constructors — the pass-through handles these) + :Normal, :Poisson, :Binomial, :Bernoulli, :Beta, :Gamma, + :Exponential, :Cauchy, :StudentT, :LogNormal, :Weibull, + :NegativeBinomial, :Geometric, :Laplace, :Uniform, + :MvNormal, :MixtureModel, :Dirichlet, + :InverseGamma, :InverseGaussian, :VonMises, :Pareto, + :OrderedLogistic, :Categorical, + # TODO stubs (not yet implemented but syntactically valid) + :scale, :center, :standardize, :factor, :offset, + :s, :bs, :t2, :gp, :ar, :ar1, :mo, + :cbind, :mvbind, :mm, :gr, :dp, :me, :centered, + :Horseshoe, :ZeroInflatedPoisson, :weighted, + # Data helpers + :length, :unique, :sort, :size, :eltype, :nrow, :ncol, +]) + +# Expression heads that are safe in a formula AST (literals, blocks, calls, …) +const _SAFE_HEADS = Set{Symbol}([ + :block, :call, :., :(=), :(||), :tuple, :vect, :ref, + :kw, :parameters, :(...), + # Comparison chains + :comparison, :&&, +]) + +struct FormulaSecurityError <: Exception + msg::String +end +Base.showerror(io::IO, e::FormulaSecurityError) = print(io, "FormulaSecurityError: ", e.msg) + +function _check_formula_safety!(x) + # Literals, symbols, line numbers — always safe + x isa Union{Number, AbstractString, Symbol, LineNumberNode, Nothing, Bool, QuoteNode} && return + x isa Expr || return + + # Reject dangerous expression types outright + if x.head == :macrocall + throw(FormulaSecurityError("macro calls are not allowed in formulas (got $(x.args[1]))")) + elseif x.head in (:cmd, :string) + throw(FormulaSecurityError("`\$(x.head)` expressions are not allowed in formulas")) + elseif x.head == :quote || x.head == :$ + throw(FormulaSecurityError("quote/interpolation expressions are not allowed in formulas")) + end + + # For :call expressions, check the function name is on the allowlist + if x.head == :call + fname = x.args[1] + if fname isa Symbol && fname ∉ _ALLOWED_CALLS + throw(FormulaSecurityError( + "function `$fname` is not in the formula allowlist. " * + "Allowed: arithmetic, math, distributions, DSL operators. " * + "See _ALLOWED_CALLS in BRMMacroWeb.jl for the full list.")) + end + # Also allow Type{...} constructors if the type name is allowed + if fname isa Expr && fname.head == :curly + tname = fname.args[1] + tname isa Symbol && tname ∉ _ALLOWED_CALLS && + throw(FormulaSecurityError("type constructor `$tname` is not in the formula allowlist")) + end + end + + # Check the expression head is expected + if x.head ∉ _SAFE_HEADS + throw(FormulaSecurityError( + "expression type `:$(x.head)` is not allowed in formulas")) + end + + # Recurse into children + for arg in x.args + _check_formula_safety!(arg) + end +end + # ── Pipeline stages ───────────────────────────────────────────────────────── # # Each stage is computed lazily on demand so the user can stop at any @@ -96,6 +193,10 @@ function pipeline(formula::AbstractString, stage::Symbol) s >= 1 || return out raw = Meta.parse("begin\n$formula\nend") + # Safety check: reject any AST node that isn't in the formula whitelist + # before handing the expression to eval. This blocks arbitrary code + # execution from the formula textarea. + _check_formula_safety!(raw) out = merge(out, (; raw)) s >= 2 || return out @@ -144,6 +245,203 @@ end _section(title, body) = (h.h3(title), h.pre(body)) +# ── Styled HTML rendering for BRMI / VBRMI cards ─────────────────────────── +# +# Each symbol gets a deterministic color, data columns are bold, parameters +# are italic, and likelihood statements are underlined. The tree walker +# (_html_expr) converts an ExprColumn AST into a nest of s. +# +# TODO: HTMX.jl should grow a generic `htmx_node(x)::Node` extension point +# so downstream packages can overload once and have every consumer dispatch +# automatically. For now these card functions are wired in by hand. + +# Deterministic HSL color per symbol (golden-ratio spread for visual variety). +_symbol_color(name::Symbol) = "hsl($(mod(hash(name) * 137, 360)), 60%, 40%)" + +# A colored with role-based font styling. +# Data columns are normal weight; parameters (latent/sampled) are bold. +_styled_name(name::Symbol, role::Symbol) = begin + s = "color:$(_symbol_color(name));" + role == :parameter && (s *= "font-weight:bold;") + h.span(string(name); style=s) +end + +# ── _html_expr: recursive ExprColumn → styled HTML ───────────────────────── + +_html_expr(x::NamedColumn{<:Any, <:DataColumn}) = _styled_name(name(x), :data) +_html_expr(x::NamedColumn{<:Any, MissingColumn}) = _styled_name(name(x), :parameter) +_html_expr(x::NamedColumn) = _styled_name(name(x), :derived) +_html_expr(x::Int) = h.span(string(x); style="color:#666") +_html_expr(x::Float64) = h.span(string(x); style="color:#666") +_html_expr(x::Number) = h.span(string(x); style="color:#666") +_html_expr(x::DataColumn) = h.span("data($(eltype(parent(x))))"; style="color:#999") +_html_expr(x::MaterializedColumn) = _html_expr(getbroadcast(x)) +_html_expr(x::LikelihoodColumn) = h.span( + _html_expr(parent(x)), h.span(" .~ "; style="color:#333"), _html_expr(rhs(x))) + +# Infix operators: always parenthesized so inner expressions like (1 + b | g1) +# keep their grouping. Top-level callers (_html_brmi_row) use _html_infix +# directly to skip the outermost parens. +_html_expr(x::ExprColumn{<:Union{typeof.((~,*,+,|,doublepipe,assign))...}}) = begin + h.span("(", _html_infix(x), ")") +end + +_html_infix(x::ExprColumn) = begin + op_str = " $(getop(x)) " + args = getargs(x) + parts = Any[] + for (i, arg) in enumerate(args) + i > 1 && push!(parts, h.span(op_str; style="color:#555")) + push!(parts, _html_expr(arg)) + end + h.span(parts...) +end + +# Function-call style: fname(args...; kwargs...) +_html_expr(x::ExprColumn) = begin + fname = getf(x) isa Function ? nameof(getf(x)) : + getf(x) isa Type ? nameof(getf(x)) : string(getf(x)) + args = getargs(x) + kw = getkwargs(x) + parts = Any[h.span(string(fname); style="color:#777"), "("] + for (i, arg) in enumerate(args) + i > 1 && push!(parts, ", ") + push!(parts, _html_expr(arg)) + end + if length(kw) > 0 + push!(parts, "; ") + for (i, (k, v)) in enumerate(pairs(kw)) + i > 1 && push!(parts, ", ") + push!(parts, "$k=", _html_expr(v)) + end + end + push!(parts, ")") + h.span(parts...) +end + +# Broadcasted objects (from VBRMI materialization): walk their inner structure +_html_expr(x::Base.Broadcast.Broadcasted) = begin + fname = x.f isa Function ? nameof(x.f) : + x.f isa Type ? nameof(x.f) : string(x.f) + args = x.args + # Infix for common operators + if x.f in (+, -, *, /) + parts = Any[] + for (i, arg) in enumerate(args) + i > 1 && push!(parts, h.span(" $(x.f) "; style="color:#555")) + push!(parts, _html_expr(arg)) + end + return h.span(parts...) + end + parts = Any[h.span(string(fname); style="color:#777"), "("] + for (i, arg) in enumerate(args) + i > 1 && push!(parts, ", ") + push!(parts, _html_expr(arg)) + end + push!(parts, ")") + h.span(parts...) +end + +# Arrays / views from block parameter slots: show as a compact shape description +_html_expr(x::SubArray) = h.span("param[$(join(size(x), "×"))]"; + style="font-style:italic;color:#888") +_html_expr(x::AbstractVector{<:Number}) = h.span("vec[$(length(x))]"; + style="font-weight:bold;color:#888") +_html_expr(x::Base.RefValue) = _html_expr(x[]) +_html_expr(x::AbstractMatrix) = h.span("mat[$(join(size(x), "×"))]"; + style="font-style:italic;color:#888") + +# Catch-all +_html_expr(x) = h.span(sprint(show, x; context=:compact=>true); style="color:#999") + +# ── BRMI card ─────────────────────────────────────────────────────────────── + +function brmi_card(brmi::BRMI) + rows = [_html_brmi_row(key, parent(value)) for (key, value) in pairs(brmi.operations)] + h.article(; style="margin:0.5rem 0")( + h.header(h.strong("BRMI"), + h.small(" — $(length(brmi.operations)) operations")), + h.div(; style="font-family:monospace;font-size:1.15em;line-height:1.8;padding:0.3rem 0")( + rows...), + ) +end + +_leaf_column(x::NamedColumn) = x +_leaf_column(x::ExprColumn) = _leaf_column(getargs(x)[1]) +_leaf_column(x) = x + +function _html_brmi_row(key, op::ExprColumn{typeof(~)}) + lhs_leaf = _leaf_column(getargs(op)[1]) + is_likelihood = lhs_leaf isa NamedColumn && parent(lhs_leaf) isa DataColumn + content = _html_infix(op) # no outer parens at top level + style = is_likelihood ? + "text-decoration:underline;text-decoration-color:#aaa;text-underline-offset:3px;" : "" + h.div(content; style) +end +_html_brmi_row(key, op::ExprColumn{typeof(assign)}) = h.div(_html_infix(op)) +_html_brmi_row(key, op) = h.div( + _styled_name(key, :data), + h.span(": "; style="color:#666"), + h.span(sprint(show, op); style="color:#999"), +) + +# ── VBRMI card ────────────────────────────────────────────────────────────── + +function vbrmi_card(vbrmi::VBRMI) + brmi = getfield(vbrmi, :parent) + (; meta) = vbrmi + n_dim = LogDensityProblems.dimension(vbrmi) + n_mat = length(meta.materialized) + n_blocks = length(meta.blocks) + + # Materialized columns: show the BRMI-level symbolic expression (styled) + # plus a compact description of the materialized value type. + mat_rows = [begin + nc = get(brmi.operations, key, nothing) + inner = nc !== nothing ? parent(nc) : nothing + if inner isa DataColumn + # Data column: styled name + eltype + h.div(_styled_name(key, :data), + h.span(": data($(eltype(parent(inner))))"; style="color:#999")) + elseif value isa LikelihoodColumn + # Likelihood: full expression, underlined + expr = inner !== nothing ? _html_infix(inner) : _styled_name(key, :data) + h.div(expr; style="text-decoration:underline;text-decoration-color:#aaa;text-underline-offset:3px") + elseif inner !== nothing + # Materialized (sampled/derived): symbolic expression + shape + expr = inner isa ExprColumn{<:Union{typeof(~),typeof(assign)}} ? + _html_infix(inner) : _html_expr(inner) + shape = "$(eltype(parent(value)))[$(length(parent(value)))]" + h.div(expr, h.span(" → $shape"; style="color:#999;font-size:0.85em")) + else + h.div(_styled_name(key, :derived), + h.span(": $(sprint(show, value))"; style="color:#999")) + end + end for (key, value) in pairs(meta.materialized)] + + # Blocks: show dimensions with styled block key + blocks_rows = [begin + m, n = size(value) + role = key == :__population__ ? :derived : :data + h.div( + _styled_name(key, role), + h.span(": (n_levels=$m, n_params=$n)"; style="color:#999"), + ) + end for (key, value) in pairs(meta.blocks)] + + h.article(; style="margin:0.5rem 0")( + h.header(h.strong("VBRMI"), + h.small(" — dim $n_dim, $n_mat materialized, $n_blocks blocks")), + h.h6(; style="margin-bottom:0.2rem")("materialized"), + h.div(; style="font-family:monospace;font-size:1.15em;line-height:1.8;margin-left:1rem")( + mat_rows...), + h.h6(; style="margin-bottom:0.2rem")("blocks"), + h.div(; style="font-family:monospace;font-size:1.15em;line-height:1.8;margin-left:1rem")( + blocks_rows...), + ) +end + + function render_output(formula::AbstractString; stage::Symbol=:vbrmi) sections = Vector{Any}[] # one entry per stage; rendered most-recent-first try @@ -176,38 +474,46 @@ function render_output(formula::AbstractString; stage::Symbol=:vbrmi) sprint(show, out.wrapped))...]) end if haskey(out, :brmi) - push!(sections, Any[_section("4. eval — BRMI value (parsed model)", - sprint(show, out.brmi))...]) + push!(sections, Any[ + h.h3("4. eval — BRMI value (parsed model)"), + brmi_card(out.brmi), + ]) end if haskey(out, :vbrmi) - vbrmi_children = Any[ - _section("5. VBRMI — materialized action (blocks, dim, columns)", - sprint(show, out.vbrmi))..., - h.h3(" logdensity at a fixed random point"), - h.p("dim = ", string(out.dim), ", logdensity = ", out.ldp), - h.h3(" finite-difference gradient sanity check"), - ] - if out.grad isa Exception - push!(vbrmi_children, - h.pre("gradient error: " * sprint(showerror, out.grad))) + # Build the FD check summary for the
toggle line + fd_summary = if out.grad isa Exception + h.span("logdensity + FD check: error"; style="color:crimson") + else + tol = 1e-8 + n_dead = count(<=(tol) ∘ abs, out.grad) + if n_dead == 0 + h.span("logdensity + FD check: $(out.dim)/$(out.dim) active ✓"; + style="color:green") + else + h.span("logdensity + FD check: $(n_dead) dead param(s)"; + style="color:crimson") + end + end + + fd_body = if out.grad isa Exception + h.pre("gradient error: " * sprint(showerror, out.grad)) else tol = 1e-8 - live = findall(>(tol) ∘ abs, out.grad) dead = findall(<=(tol) ∘ abs, out.grad) - summary_color = isempty(dead) ? "green" : "crimson" - push!(vbrmi_children, h.p( - "active params: ", h.strong("$(length(live))/$(out.dim)"), - " — ", - h.span(; style="color:$summary_color")( - isempty(dead) ? - "all parameters influence the logdensity ✓" : - "$(length(dead)) dead param(s) at indices $(dead)" - ), - )) - push!(vbrmi_children, - h.pre(sprint(show, MIME"text/plain"(), out.grad))) + h.div( + h.p("dim = ", string(out.dim), ", logdensity = ", out.ldp), + isempty(dead) ? "" : + h.p(; style="color:crimson")( + "dead param indices: ", string(dead)), + h.pre(sprint(show, MIME"text/plain"(), out.grad)), + ) end - push!(sections, vbrmi_children) + + push!(sections, Any[ + h.h3("5. VBRMI — materialized action (blocks, dim, columns)"), + vbrmi_card(out.vbrmi), + h.details(h.summary(fd_summary), fd_body), + ]) end if haskey(out, :bench) bench_body = out.bench isa Exception ? @@ -231,47 +537,732 @@ end # ── Routes ────────────────────────────────────────────────────────────────── _stage_button(label, stage) = h.button(label; type="button", + id="stage-$stage", hx_get="/stage/$stage", hx_include="#brm-macro-form", hx_target="#brm-macro-output", hx_swap="outerHTML") +# Pre-canned formulas. The ones above the divider exercise individual features +# in isolation; the last one stacks everything into a single multi-likelihood +# model. Click loads the formula into the textarea — user still hits "Render" +# (or any stage button) to advance the pipeline. +function presets() + [ + "min" => "loc ~ 1\nlog(err) ~ 1\ny1 ~ Normal(loc, err)\n", + "linear" => "loc ~ 1 + a\nlog(err) ~ 1\ny1 ~ Normal(loc, err)\n", + "multi-lin" => "loc ~ 1 + a + b + c + d\nlog(err) ~ 1\ny1 ~ Normal(loc, err)\n", + "categorical" => "loc ~ 1 + c1\nlog(err) ~ 1\ny1 ~ Normal(loc, err)\n", + "random intercept" => "loc ~ 1 + a + (1 | g1)\nlog(err) ~ 1\ny1 ~ Normal(loc, err)\n", + "random slope" => "loc ~ 1 + (1 + a | g1)\nlog(err) ~ 1\ny1 ~ Normal(loc, err)\n", + "categorical random slope" => "loc ~ 1 + (1 + c1 | g1)\nlog(err) ~ 1\ny1 ~ Normal(loc, err)\n", + "multiple groups" => "loc ~ 1 + a + (1 | g1) + (1 | g2)\nlog(err) ~ 1\ny1 ~ Normal(loc, err)\n", + "distributional" => "loc ~ 1 + a\nlog(err) ~ 1 + b\ny1 ~ Normal(loc, err)\n", + "Poisson" => "log_rate ~ 1 + a + (1 | g1)\nk1 ~ Poisson(exp(log_rate))\n", + "Binomial" => "log_odds ~ 1 + a + (1 | g1)\nbin_succ ~ Binomial(bin_n, logistic(log_odds))\n", + "Bernoulli" => "log_odds ~ 1 + a + (1 | g1)\nbin_y ~ Bernoulli(logistic(log_odds))\n", + # Joint smoke test for Peter's verified non-Normal examples: + # brms::cbpp_binomial → categorical + random intercept + Binomial with + # per-row trial counts; kruschke::therapeutic_touch → hierarchical + # Bernoulli. Both share the same grouping factor here so they hit the + # cross-likelihood block-sharing path too. + "cbpp + therapeutic touch" => """log_odds_bin ~ 1 + c1 + (1 | g1) +bin_succ ~ Binomial(bin_n, logistic(log_odds_bin)) + +log_odds_b ~ 1 + (1 | g1) +bin_y ~ Bernoulli(logistic(log_odds_b)) +""", + "everything" => default_formula(), + ] +end + +_preset_button(label, formula) = h.button(label; + type="button", + data_formula=formula, + onclick="document.querySelector('textarea[name=formula]').value = this.dataset.formula; document.getElementById('stage-vbrmi').click()", + style="font-size:0.8em;padding:0.2rem 0.5rem;margin:0") + +_index_body(formula::String) = h.div( + h.h1("BRM macro pipeline"), + h.p( + "Enter a ", h.code("@brm"), " formula and step through the macro pipeline: ", + h.code("Meta.parse"), " → ", h.code("parse!"), " → ", h.code("_brm"), + " let-block → ", h.code("eval"), " → ", h.code("VBRMI"), " action → ", + h.code("Chairmarks"), " benchmark.", + ), + h.details( + h.summary(h.small("Allowed functions in formulas")), + h.p(h.small( + join(sort(collect(string.(s) for s in _ALLOWED_CALLS)), ", "), + )), + ), + h.form(; id="brm-macro-form")( + h.label("Load preset"), + h.div(; style="display:flex;flex-wrap:wrap;gap:0.3rem;margin-bottom:0.6rem")( + [_preset_button(label, body) for (label, body) in presets()]..., + ), + h.label("Formula")( + h.textarea(formula; + name="formula", rows=8, + style="width:100%;font-family:monospace"), + ), + h.fieldset(; class="grid")( + _stage_button("1. Parse", :parse), + _stage_button("2. Transform", :transform), + _stage_button("3. Wrap", :wrap), + _stage_button("4. BRMI", :brmi), + _stage_button("5. VBRMI", :vbrmi), + _stage_button("6. Benchmark", :bench), + ), + ), + render_output(formula), +) + @htmx struct AppContext req = nothing - @get index(; formula::String=default_formula()) = htmx(h.main(class="container")( - h.h1("BRM macro action"), - h.p( - "Enter a ", h.code("@brm"), " formula and step through the macro pipeline: ", - h.code("Meta.parse"), " → ", h.code("parse!"), " → ", h.code("_brm"), - " let-block → ", h.code("eval"), " → ", h.code("VBRMI"), " action → ", - h.code("Chairmarks"), " benchmark.", - ), - h.form(; id="brm-macro-form")( - h.label("Formula")( - h.textarea(formula; - name="formula", rows=8, - style="width:100%;font-family:monospace"), - ), - h.fieldset(; class="grid")( - _stage_button("1. Parse", :parse), - _stage_button("2. Transform", :transform), - _stage_button("3. Wrap", :wrap), - _stage_button("4. BRMI", :brmi), - _stage_button("5. VBRMI", :vbrmi), - _stage_button("6. Benchmark", :bench), + # HTMXObjects auto-uses `page` to wrap any route's return value into a full + # page on direct browser navigation, while returning just the fragment for + # HTMX requests (see `_resolve_response` in HTMXObjects.jl). The sidebar's + # `hx-get` swaps target `#content` directly. + page(content) = htmx( + h.div(; style="display:flex;gap:1rem;align-items:flex-start")( + nav_sidebar([ + "Pipeline" => "/", + "TODO list" => "/todo", + ]), + h.main(; class="container", style="flex:1;min-width:0")( + h.div(; id="content")(content), ), + ); + pico_version="2", + extra_head=( + h.title("BRM macro action"), + h.style(":root { font-size: 87.5%; }"), ), - render_output(formula), - ); pico_version="2", extra_head=( - h.title("BRM macro action"), - h.style(":root { font-size: 87.5%; }"), - )) + ) + + @get index(; formula::String=default_formula(), label::String="") = begin + # If a TODO form posted us a (label, formula) pair, persist the edited + # formula to that TODO's .jl file so the next visit to the TODO page + # shows the user's edits instead of the seed default. + isempty(label) || _save_todo!(label; new_formula=formula) + _index_body(formula) + end + + @get mark(; label::String="", state::String="") = begin + isempty(label) && return "" + target = Symbol(state) + entry = _find_todo(label) + entry === nothing && return "" + next_status = entry.status == target ? :open : target + updated = _save_todo!(label; new_status=next_status) + # Re-render the whole card so the border + collapse state update + # together with the pill text. + _todo_card(updated) + end - @get stage(name::AbstractString; formula::String=default_formula()) = + @get stage(name::AbstractString; formula::String=default_formula(), label::String="") = begin + # When called from a TODO card's form, persist the (possibly edited) + # formula back to the todo's .jl file before rendering. + isempty(label) || _save_todo!(label; new_formula=formula) render_output(formula; stage=Symbol(name)) + end + + @get todo = begin + todos = _load_todos() + h.div( + h.h1("TODO — what's missing for full BRM coverage"), + h.p("Items grouped by tier. Each item has a sketch of what it is, why it matters, how to implement, and how to verify. Sourced from .jl files under ", h.code("web-macro/todos/"), "; status edits and edited formulas are written back to disk."), + h.h2("Tier 1 — cheap wins"), + [_todo_card(t) for t in todos if t.tier == 1]..., + h.h2("Tier 2 — moderate (one design decision each)"), + [_todo_card(t) for t in todos if t.tier == 2]..., + h.h2("Tier 3 — bigger features (real new infrastructure)"), + [_todo_card(t) for t in todos if t.tier == 3]..., + ) + end end +# ── TODO content (file-backed) ────────────────────────────────────────────── +# +# Each TODO is a `.jl` file under `web-macro/todos/`. File format: +# +# # label: 1.1 verify Bernoulli/Binomial — done +# # tier: 1 +# # status: open +# #= +# **Markdown body** with whatever explanation text you want. +# =# +# +# +# +# Header lines (`# key: value`) carry metadata. The `#= ... =#` block is the +# markdown body. Everything after the body block is the formula. The web app +# loads + parses these files on every render of the TODO page, and writes them +# back when the user toggles status or submits an edited formula. Reopening the +# server picks up exactly where the user left off — no in-memory state. + +struct TodoEntry + path::String + label::String + tier::Int + status::Symbol # :open | :done | :deprioritized + body::String # markdown + formula::Union{String,Nothing} +end + +_todos_dir() = joinpath(dirname(@__DIR__), "todos") +_slug(label::AbstractString) = lowercase(strip(replace(label, r"[^\w.]+" => "-"), '-')) + +const _todos_cache = Ref{Vector{TodoEntry}}() + +function _load_todos(; refresh::Bool=false) + if refresh || !isassigned(_todos_cache) + dir = _todos_dir() + isdir(dir) || _migrate_todos!() + files = sort(filter(endswith(".jl"), readdir(dir; join=true))) + _todos_cache[] = TodoEntry[_parse_todo_file(f) for f in files] + end + _todos_cache[] +end + +function _find_todo(label::AbstractString) + for t in _load_todos() + t.label == label && return t + end + nothing +end + +function _parse_todo_file(path::String) + lines = readlines(path) + header = Dict{String,String}() + i = 1 + + # Header: leading lines matching `# key: value`. Stop at the first + # non-matching line. Line-based avoids any UTF-8 byte-index footguns + # (labels routinely contain multi-byte characters like `✓`). + while i <= length(lines) + m = match(r"^# (\w+):\s*(.*)$", lines[i]) + m === nothing && break + header[m[1]] = m[2] + i += 1 + end + + # Body: optional `#= ... =#` block. Both delimiters live on their own + # lines (the writer guarantees this), so a simple line scan suffices. + body_lines = String[] + if i <= length(lines) && strip(lines[i]) == "#=" + i += 1 + while i <= length(lines) && strip(lines[i]) != "=#" + push!(body_lines, lines[i]) + i += 1 + end + i <= length(lines) && (i += 1) # consume `=#` + end + body = join(body_lines, '\n') + + # Formula: everything that's left, stripped. + formula_text = strip(join(lines[i:end], '\n')) + formula = isempty(formula_text) ? nothing : String(formula_text) + + label = get(header, "label", basename(path)) + tier = parse(Int, get(header, "tier", "1")) + status = Symbol(get(header, "status", "open")) + TodoEntry(path, label, tier, status, body, formula) +end + +function _write_todo_file(todo::TodoEntry) + io = IOBuffer() + println(io, "# label: ", todo.label) + println(io, "# tier: ", todo.tier) + println(io, "# status: ", todo.status) + if !isempty(todo.body) + println(io, "#=") + println(io, todo.body) + println(io, "=#") + end + if todo.formula !== nothing && !isempty(todo.formula) + println(io) + print(io, todo.formula) + endswith(todo.formula, "\n") || println(io) + end + write(todo.path, take!(io)) +end + +function _save_todo!(label::AbstractString; + new_status::Union{Symbol,Nothing}=nothing, + new_formula::Union{String,Nothing}=nothing) + todo = _find_todo(label) + todo === nothing && return nothing + updated = TodoEntry( + todo.path, todo.label, todo.tier, + something(new_status, todo.status), + todo.body, + new_formula === nothing ? todo.formula : new_formula, + ) + _write_todo_file(updated) + _load_todos(refresh=true) + updated +end + +# One-shot migration: dumps the in-source `_tier1()`/`_tier2()`/`_tier3()` +# seed lists to files on first run. Skips files that already exist, so user +# edits to existing files survive. Once everything is migrated, the in-source +# seed functions are dead code that can eventually be deleted. +function _migrate_todos!() + dir = _todos_dir() + isdir(dir) || mkpath(dir) + for (tier, items) in [(1, _tier1()), (2, _tier2()), (3, _tier3())] + for item in items + label, body, formula = item isa Pair ? + (first(item), last(item), nothing) : + (item.label, item.body, item.formula) + slug = _slug(label) + path = joinpath(dir, slug * ".jl") + isfile(path) && continue + _write_todo_file(TodoEntry(path, label, tier, :open, body, formula)) + end + end +end + +# ── Rendering: one Pico CSS article per TODO with status-colored border ──── + +const _STATUS_COLORS = ( + open = "#888", + done = "#2e7d32", + deprioritized = "#a05a2c", +) +_status_color(s::Symbol) = get(_STATUS_COLORS, s, "#888") + +function _todo_card(todo::TodoEntry) + border_color = _status_color(todo.status) + body_children = Any[HTMXObjects.md_to_node(todo.body)] + if todo.formula !== nothing + push!(body_children, _formula_form(todo.label, todo.formula)) + # Inline pipeline-result target — the form's hx_get fills this div + # with `render_output(formula; stage=:vbrmi)` so the user sees the + # VBRMI/finite-difference output right inside the card. + push!(body_children, h.div(; id="todo-result-$(hash(todo.label))", + style="margin-top:0.5rem")) + end + # `:open` status → expanded; `:done`/`:deprioritized` → collapsed by default. + # Pills sit inside the so they're always reachable, but their + # onclick stops propagation so clicking a pill doesn't also toggle the + # disclosure. + h.article(; + id="todo-card-$(hash(todo.label))", + style="border-left:6px solid $border_color;margin:0.8rem 0;padding:0.5rem 1rem", + )( + h.details(; open=todo.status == :open)( + h.summary(; style="cursor:pointer;list-style-position:outside")( + h.strong(todo.label), " ", _status_pills(todo.label, todo.status), + ), + h.div(; style="margin-top:0.5rem")(body_children...), + ), + ) +end + +function _formula_form(label::String, formula::String) + h.form(; + hx_get="/stage/vbrmi", + hx_target="#todo-result-$(hash(label))", + hx_swap="innerHTML", + style="margin:0.5rem 0", + )( + h.input(; type="hidden", name="label", value=label), + h.textarea(formula; + name="formula", + rows=max(3, count('\n', formula) + 1), + style="width:100%;font-family:monospace;font-size:0.85em"), + h.button("Try in pipeline ▶"; + type="submit", + style="font-size:0.85em;padding:0.3rem 0.8rem;margin:0.3rem 0 0 0"), + ) +end + +# Mutually-exclusive status pills wrapped in a single span so one toggle +# replaces both at once via `outerHTML` targeting `#status-{hash(label)}`. +function _status_pills(label::AbstractString, state::Symbol) + h.span(; id="status-$(hash(label))", + style="margin-left:0.5rem;display:inline-flex;gap:0.3rem;vertical-align:middle")( + _state_pill(label, :done, state, "✓ done", "mark done"), + _state_pill(label, :deprioritized, state, "✓ deprioritized", "deprioritize"), + ) +end + +# Convenience overload that fetches the current status from disk. +_status_pills(label::AbstractString) = _status_pills(label, + something(_find_todo(label), (status=:open,)).status) + +# Pill text and color reflect the *current* state. When the pill's target_state +# is currently active, it shows the active label (e.g. "✓ done") in the active +# color and clicking it toggles back to :open. Otherwise it shows the inactive +# action label (e.g. "mark done") in gray and clicking it sets the target state. +function _state_pill(label, target_state, current_state, active_text, inactive_text) + is_active = current_state == target_state + text = is_active ? active_text : inactive_text + bg = is_active ? _status_color(target_state) : "#888" + h.button(text; + type="button", + hx_get="/mark?label=$(HTTP.URIs.escapeuri(label))&state=$(target_state)", + hx_target="#todo-card-$(hash(label))", + hx_swap="outerHTML", + onclick="event.stopPropagation()", + style="font-size:0.7em;padding:0.15rem 0.6rem;border:none;border-radius:1rem;" * + "color:white;background:$bg;cursor:pointer", + ) +end + +_tier1() = [ + (label="1.1 verify Bernoulli/Binomial", + body=raw""" +**Status: done.** ✓ Confirmed that the existing `FBroadcasted{<:Type{<:Distribution}}` pass-through in `vimpl.jl` handles both Bernoulli and Binomial cleanly. + +**Verification.** The form below loads the **cbpp + therapeutic touch** model — a faithful translation of `brms::cbpp_binomial` (categorical predictor + random intercept + Binomial with per-row trial counts) and `kruschke::therapeutic_touch` (hierarchical Bernoulli) into one multi-likelihood model. brms's `incidence | trials(size) ~ ...` sidecar collapses to a plain positional argument: `bin_succ ~ Binomial(bin_n, logistic(η))`. If the gradient sanity check stays green every other `Distribution` family (Beta, Gamma, NegBinomial, …) should be free as well. +""", + formula="""log_odds_bin ~ 1 + c1 + (1 | g1) +bin_succ ~ Binomial(bin_n, logistic(log_odds_bin)) + +log_odds_b ~ 1 + (1 | g1) +bin_y ~ Bernoulli(logistic(log_odds_b)) +"""), + + (label="1.2 offset / fixed exposure — already works without a wrapper", + body=raw""" +**Status: already works without any new code.** brms needs `offset(z)` because R's formula syntax has no other way to put a "no-coefficient term" into the linear predictor — the only thing on the RHS of `~` is the formula DSL. In our DSL the linear predictor and the likelihood are *separate* `~` lines, and the second one (the likelihood) takes a free-form Julia expression. Anything inside that expression gets evaluated as plain code at materialization time via `vbroadcasted` — function calls dispatch to whatever Julia function the symbol resolves to, and data column references are pulled from the dataframe. + +So instead of `count ~ x + offset(log(exposure))`, you write the offset directly inside the likelihood: + +```julia +loc ~ 1 + a +k1 ~ Poisson(exp(loc + log(exposure))) +``` + +The `log(exposure)` here is just `Base.log` applied to the `exposure` data column, broadcasted across rows and added to `loc` (which is the materialized linear predictor). No parameter is allocated for it because `growblock!!` is never called for that branch — there's no `~` on the data side, just an argument to `Poisson(...)`. + +**Verification.** Form below loads exactly that model. The VBRMI dim should match the offset-free version (only the population intercept + slope on `a`); the gradient sanity check should stay green; and the materialized `k1` likelihood should incorporate the row-specific exposure shift. +""", + formula="""loc ~ 1 + a +k1 ~ Poisson(exp(loc + log(exposure))) +"""), + + (label="1.3 I(expr) — likely already works", + body=raw""" +**What it is.** brms's `I()` is a literal-escape: `I(x^2)` says "compute `x^2` from the data and treat it as a single column". brms needs it because `+`, `*`, `:`, `|`, … all have special meaning inside an R formula. + +**Why we probably don't need it.** Our DSL is parsed by Julia first, then walked by `_x`. `_x` recursively wraps every `Expr(:call, f, args...)` in an `ExprColumn`, regardless of whether `f` is special. So `loc ~ a + x^2` becomes `+(a, ^(x, 2))` → `ExprColumn(+, NamedColumn(:a), ExprColumn(^, NamedColumn(:x), 2))`. The `^` is just another function call, no special handling needed. + +The only operators that have DSL meaning in our system are `~` (sampling), `=` (assignment), and `|` / `||` inside random-effects specs. Everything else (`^`, `/`, `sqrt`, `log`, `exp`, `mod`, `min`, `max`, …) is a regular function call resolved at materialization time via `vbroadcasted`. + +**Verification.** The form below loads a model with three nonlinear terms (`a^2`, `sqrt(abs(b))`, `log(exposure)`) directly as population-level covariates. The VBRMI dim should match the number of distinct terms; the gradient sanity check should be all-active. If it works, that confirms `I()` is unnecessary because Julia function calls are first-class on the formula RHS. +""", + formula="""loc ~ 1 + a + a^2 + sqrt(abs(b)) + log(exposure) +y1 ~ Normal(loc, 1) +"""), + + "1.4 scale(x) / standardize(x)" => raw""" +**What it is.** brms's `scale(x)` z-transforms a column at parse time: `scale(x) = (x - mean(x)) / std(x)`. The model sees the standardized column. Crucial for default priors (which are scale-invariant only after standardization) and sampler stability (well-conditioned linear predictors). + +**Why it matters.** Most brms vignettes do `scale(x)` automatically as a convenience. Without it, every formula has to either manually z-transform the data or accept poorly-scaled coefficients. + +**Implementation.** +1. Add `function scale end` (and `function center end`, `function standardize end`) to `macro.jl`. +2. Add a `vmeta_sampling_rhs` overload in `vimpl.jl`: +```julia +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(scale)}; group) = begin + inner = vbroadcasted(only(getargs(x)); meta) + materialized = Base.materialize(inner) + z = (materialized .- Statistics.mean(materialized)) ./ Statistics.std(materialized) + vmeta_sampling_rhs(meta, z; group) +end +``` +The standardization happens once when the BRMI is materialized into a VBRMI. Composes with the existing dense-map caching TODO. +3. Add `Statistics` to `vimpl.jl`'s using-list (or vendor `mean`/`std` inline). + +**Verification.** Preset: `loc ~ 1 + scale(a) + scale(b); y1 ~ Normal(loc, 1)`. Compare against the unscaled version: same dim, different posterior geometry. The fitted coefficients should be ≈ the unscaled coefficients × std(x). +""", + "1.5 zerocorr — independent random effects" => raw""" +**What it is.** brms (via lme4 syntax) lets you opt out of the LKJ correlation between multiple random terms in the same group. `(1 + x || group)` (double bar) says "estimate the random intercept and the random slope independently — don't fit a 2×2 Cholesky factor between them". Useful when there isn't enough data to estimate the correlations, or when you have prior reason to believe the terms are uncorrelated. + +**Why it matters.** Multi-term random specs are common, and the LKJ correlation often dominates the prior cost without much identifiability. Letting users skip it is a meaningful sampling speedup and prior simplification. + +**Implementation.** Our `_x` walker already wraps `||` as `ExprColumn{typeof(doublepipe)}`. Add a `vmeta_sampling_rhs` overload that splits each term inside the `||` LHS into its own block (with a synthetic per-term key like `Symbol(group_name, :__nocor__, term_index)`): + +```julia +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(doublepipe)}; kwargs...) = begin + lhs, rhs = getargs(x, 2) + terms = lhs isa ExprColumn{typeof(+)} ? getargs(lhs) : (lhs,) + foldl(enumerate(terms); init=(meta, ())) do (m, args), (i, term) + nocor_key = NamedColumn(Symbol(name(rhs), :__nocor__, i), parent(rhs)) + m, arg = vmeta_sampling_rhs(m, term; group=nocor_key) + m, (args..., arg) + end |> ((m, args),) -> (m, Base.broadcasted(+, args...)) +end +``` + +Each per-term block ends up as 1×1 with one `log_scale` Cholesky parameter — `lprior!`'s existing single-column path handles this with no changes. + +**Verification.** Preset: `loc ~ 1 + (1 + a || g1); y1 ~ Normal(loc, 1)`. Compare its dim against the correlated `(1 + a | g1)` version: the correlated version has 3 Cholesky params (1+2/2 for a 2×2), the uncorrelated version has 2 (one log_scale per term). Same direct-parameter count (2 cols × 8 levels = 16) either way. +""", + "1.6 cache levels / level_map / dense / gc_idx" => raw""" +**What it is.** Stop rebuilding the dense level mapping (`Dict(level => row_index)`) and the gc_idx vector on every `VBRMI(brmi)` call. Cache them once per source data column. + +**Why it matters.** Today every `VBRMI` build re-traces the categorical / grouping columns, sorts unique values, builds a Dict, and walks the column to dense-encode it. For models with many categorical columns or many `VBRMI` rebuilds (e.g. during AD), this adds up. + +**Implementation.** Pick a storage layout for the per-column metadata. Two candidates: +- A new `meta.factor` NamedTuple keyed by source column name, holding `(; levels, level_map, dense, gc_idx)` per column. Built lazily on first reference, indexed via `name(column)`. +- Attach the metadata to `meta.materialized[column_name]` directly. More tightly coupled but avoids a parallel NamedTuple. + +The TODO already lives at `vimpl.jl:78–86`. Once a layout is picked, refactor `_gc_idx` and the inline dense map in the categorical path to read from the cache, falling back to a build-on-miss helper. + +**Verification.** No behavioral change — the gradient sanity check should stay green. Benchmark `VBRMI(brmi)` with Chairmarks before/after and confirm a measurable speedup on a model with multiple categorical/grouping columns. +""", + "1.7 CategoricalArrays / PooledArrays integration" => raw""" +**What it is.** When the input column is already a `CategoricalVector` or a `PooledArray`, the dense level mapping is already computed and stored in the column's `.refs` field. Use it directly instead of rebuilding via `Dict`. + +**Why it matters.** Most real-world DataFrames use `CategoricalArrays.jl` for factor columns. Skipping the rebuild eliminates allocation entirely for the common case and gets us "for free" interop with the standard categorical-data ecosystem. + +**Implementation.** Two design choices: +- **Hard dep**: add `CategoricalArrays` to vimpl.jl's deps, dispatch on `CategoricalVector`, read `levelcode.(col)` and `levels(col)` directly. +- **Duck-typed**: sniff for the `.refs` field and `levels` method without importing the package, falling back to the generic Dict path. + +Recommend hard dep — it's the standard for tabular Julia code, and the duck-type path is more code with no real win. Same for `PooledArrays`. + +The actual integration is small once the design is picked: a method specialization in `_gc_idx` and in the categorical-predictor path. Composes with the caching TODO above. + +**Verification.** Preset (or test) that builds a DataFrame with a `CategoricalVector` column and uses it as a grouping factor / categorical predictor. Confirm the gradient sanity check stays green and the per-VBRMI allocation count drops. +""", +] + +_tier2() = [ + "2.1 interactions a:b, a*b" => raw""" +**What it is.** brms's `a:b` is the elementwise interaction term (a single coefficient multiplying `a[i] * b[i]`). `a*b` is the "main effects + interaction" shorthand: it desugars to `a + b + a:b`. + +**Why it matters.** Interactions are the most commonly missed feature in regression DSLs. Without them, every model that needs `a:b` has to manually create the interaction column in the input DataFrame. + +**Implementation.** +1. **Parser side.** Add a `:` case to `_x` so that `a:b` becomes `ExprColumn(:, NamedColumn(:a), NamedColumn(:b))` instead of falling through to a Symbol/Range parse error. +2. **Materialization side.** Add `vmeta_sampling_rhs(meta, x::ExprColumn{typeof(:)}; group)` that elementwise-multiplies the operands and dispatches to the float-vector path. For continuous × continuous it's a single coefficient on `a .* b`; for categorical × continuous it's `(k-1)` coefficients (one per non-reference level of the categorical, multiplied by the continuous); for categorical × categorical it's `(k₁-1)*(k₂-1)` coefficients via a 2D `_cat_lookup`. +3. **`a*b` desugaring.** At parse time in `_x`, rewrite `*` between formula terms as `+(a, b, :(a:b))`. This needs care because `*` also means multiplication elsewhere (e.g. `Normal(0, 2*sigma)`); the rewrite should only apply at formula-RHS top-level. + +**Verification.** Presets exercising each interaction type: +- continuous×continuous: `loc ~ 1 + a + b + a:b; y1 ~ Normal(loc, 1)` → dim 4 +- continuous×categorical: `loc ~ 1 + a + c1 + a:c1; y1 ~ Normal(loc, 1)` → dim 6 (1 + 1 + 2 + 2) +- categorical×categorical: `loc ~ 1 + c1 + c2 + c1:c2; y1 ~ Normal(loc, 1)` → dim 5 (1 + 2 + 1 + 2) +- shorthand: `loc ~ 1 + a*b; y1 ~ Normal(loc, 1)` should match `loc ~ 1 + a + b + a:b` exactly. +""", + "2.2 configurable categorical reference level" => raw""" +**What it is.** Currently the reference level for treatment-coded categoricals is `sort(unique(x))[1]`. brms / lme4 let you override this via `factor(x, ref="some_level")` or by reordering the factor's levels. + +**Why it matters.** The reference level changes the interpretation of the intercept (it becomes "the mean for the reference level") and of the coefficients (each becomes "the difference from reference"). For some analyses, changing the reference is the only way to make the coefficients directly answer the research question. + +**Implementation.** +1. Add `function factor end` to `macro.jl`. +2. Either store the override at parse time (rewrite `factor(x, ref=:level3)` into a wrapper that the materializer recognizes) or at materialization time via a `meta.factor_ref` NamedTuple keyed by column name. +3. The categorical-predictor path's `levels = sort(unique(x))` becomes `levels = sort(unique(x), by=l -> l == ref ? -Inf : l)` so the chosen reference always sorts first. + +**Verification.** Preset: `loc ~ 1 + factor(c1, ref=2); y1 ~ Normal(loc, 1)`. Compare the fitted coefficients against the default-reference version — they should differ by the level-2-vs-level-1 mean shift but produce the same logdensity. +""", + "2.3 per-parameter prior scales" => raw""" +**What it is.** Currently every parameter is `Normal(0, 1)` in `lprior!`. brms / Stan-style models routinely set custom priors per coefficient: `b ~ Normal(0, 0.5)` for tight priors on slopes, `b ~ Cauchy(0, 1)` for heavy-tailed priors, etc. + +**Why it matters.** Default `Normal(0, 1)` is fine after standardization but poor on raw scales. Allowing per-parameter prior scales is the prerequisite for spike-and-slab, Horseshoe, and most prior sensitivity analyses. Without it, users have no way to express domain knowledge about parameter magnitudes. + +**Implementation.** This is the largest design decision in Tier 2 because it has knock-on effects for every other prior-related TODO. + +Two storage candidates: +- **Per-block scales**: extend `meta.block_data[group]` with a per-column scale vector. `lprior!` multiplies the standard-normal draw by the scale before storing. Simple but only handles Normal-with-scale priors. +- **Per-block prior distributions**: store a vector of `Distribution` objects per block. `lprior!` calls `logpdf(prior_i, xi)` for each parameter. More general; handles Cauchy, StudentT, Horseshoe, etc. + +Recommend the second — it's strictly more powerful and the runtime cost is identical (one `logpdf` call per parameter). Default value is `Normal(0, 1)` for backward compatibility. + +**Implementation sketch.** +1. Extend `meta.block_data` with a `priors` field per block. +2. The macro syntax `b ~ Normal(0, 0.5)` parses as a sampling statement with a Distribution-typed RHS. Currently this is reserved for likelihood declarations; it would need a new "is this a prior or a likelihood?" branch in `vmeta_sampling`. Likelihood: LHS is a data column. Prior: LHS is a maybelocal (parameter). +3. `lprior!` reads the per-column prior and calls `logpdf(prior, value)` instead of the hard-coded `logpdf(Normal(), value)`. + +**Verification.** Preset: `loc ~ 1 + a; b ~ Normal(0, 0.1); y1 ~ Normal(loc, 1)` — confirm the gradient is dampened on `b` compared to the default-prior version, and that the dead-param check still passes. +""", + "2.4 centered / non-centered parameterization toggle" => raw""" +**What it is.** Currently every random-effect block uses non-centered parameterization (we sample standard normals and apply `mul!(vi, C.L, xi)`). brms / Stan let you choose centered (sample directly from `Normal(0, σ)` per group) on a per-factor basis. + +**Why it matters.** Non-centered is the default for "weak data per group" cases (Neal's funnel pathology), but for "strong data per group" cases centered samples better. Letting users choose is a meaningful sampling speedup for the latter regime. + +**Implementation.** Small change to `lprior!` and `growblock!!`. Add a `centered::Bool` flag to `meta.block_data[group]`. In `lprior!`'s non-population branch, if the block is centered, sample directly from `Normal(0, exp(log_scale))` instead of `Normal(0, 1)` then multiplying by `L`. The Cholesky machinery for off-diagonal correlations still applies in the centered case — just on the column before the variance scaling rather than after. + +Once (2.3) is in place, the centered/non-centered choice could be encoded as `(1 | g) ~ Normal(0, σ)` (centered) vs the implicit non-centered default — but for now a per-block kwarg or a wrapper function (e.g. `centered((1 + a | g))`) is simpler. + +**Verification.** Same model with both parameterizations should produce the same logdensity at the same parameter values (after the appropriate change of variables). Sampling efficiency on a known-funnel dataset should differ. +""", + "2.5 grouped random effects (per-factor variance)" => raw""" +**What it is.** Peter's "different variance by diagnosis" pattern: `(1 | subject) gr(diagnosis)` says "the random intercept by subject has a different variance per diagnosis level". In brms this is a custom group structure where the variance hyperparameter itself depends on a second factor. + +**Why it matters.** Common in clinical data where treatment groups have intrinsically different between-subject variability. Without this, you have to fit separate models per diagnosis or accept a single pooled variance. + +**Implementation.** Bigger than it looks because the variance is no longer a single scalar but a length-`n_levels(diagnosis)` vector that needs its own prior and its own gradient. + +Proposed shape: +- A new `gr(group_factor)` wrapper recognized in the `~` RHS via a `function gr end` stub (already exists in `macro.jl`). +- The wrapped block stores `n_levels(group_factor)` log-scale parameters instead of one. `lprior!` walks them, multiplying each subject's random intercept by the diagnosis-specific scale. +- Requires the gc_idx for the inner factor (subject) AND for the outer factor (diagnosis) — both vectors of length N. + +This composes naturally with (1.6 caching) and (2.3 per-parameter priors). + +**Verification.** Preset against synthetic data with two grouping factors, one nested inside the other, with intentionally different per-outer-level variance. Confirm the fitted scales recover the synthetic values. +""", + "2.6 multi-membership random effects mm()" => raw""" +**What it is.** brms's `mm(g1, g2, ...)` lets one observation belong to **multiple** levels of the same random factor simultaneously, with weights summing to 1. Standard use: a student belongs to multiple schools across the year, and we want their random effect to be a weighted average of the per-school effects. + +**Why it matters.** Standard random effects assume each observation belongs to exactly one group. Multi-membership is the only clean way to handle observations that span groups (mobile students, patients seen by multiple clinicians, etc.). + +**Implementation.** `_gc_idx` would have to return a row-of-vectors instead of a single Int per row. Two paths: +- **Sparse design matrix**: replace the `gc_idx` lookup with a sparse `(N × n_levels)` matrix where each row's nonzero entries are the membership weights. The materialized random effect becomes `sparse_membership * random_effects_vector`. +- **Per-row lookup loop**: keep the row-major view but make `_re_lookup` iterate the membership list per row, summing weighted contributions. + +The sparse matrix approach is more memory-efficient and SIMD-friendly. Needs a new wrapper in the formula syntax: `(1 | mm(g1, g2; weights=...))`. + +**Verification.** Preset against synthetic data where each observation has 2 random group memberships with weights summing to 1. Compare against the equivalent "fully observed in primary group only" model. +""", + "2.7 se() / weights() for meta-analysis and weighted regression" => raw""" +**What it is.** brms's `y | se(sigma_y) ~ ...` lets each observation have its own known standard error (typical for meta-analysis where each `y` is itself a summary estimate). `y | weights(w) ~ ...` is observation-level weighting (typical for survey data or sample-size correction). + +**Why it matters.** Both are extremely common in applied work. Without them, meta-analysis can't be expressed at all in this DSL, and weighted regression has to be hacked via likelihood multiplication. + +**Implementation.** Both are sidecar modifiers on the LHS of `~`, so they need parser support similar to brms's `|` syntax. Or, more naturally for our DSL: pass them as positional arguments to the distribution itself. +- `y ~ Normal(mu, se_y)` for the meta-analysis case (already works! `se_y` is just another data column). +- For weights, define a `weighted` likelihood wrapper: `y ~ weighted(Normal(mu, sigma), w)` where `weights` multiplies the per-row logpdf by `w[i]`. Needs a new `vmeta_sampling_rhs` overload and a new `LikelihoodColumn`-like type with a per-row weight. + +Meta-analysis is essentially free (already works). Weights need ~15 lines. + +**Verification.** Preset for meta-analysis: `y ~ Normal(mu, se_y)` with `se_y` from synthetic data. Preset for weights: `y ~ weighted(Normal(mu, 1), weight); loc ~ 1 + a` confirming the gradient is rescaled per-row by `weight`. +""", +] + +_tier3() = [ + "3.1 multivariate outcomes cbind(y1, y2)" => raw""" +**What it is.** brms's `cbind(y1, y2) ~ x + (1 | g)` declares that `y1` and `y2` share the same linear predictor structure but have correlated residuals. The likelihood becomes multivariate normal (or multivariate-t) over `(y1, y2)` with a covariance matrix to estimate. + +**Why it matters.** Unblocks `mcelreath::waffle_divorce_multivariate` and any model where multiple outcomes share latent structure (joint pharmacology/efficacy, paired outcomes, mediation analysis). + +**Implementation.** Real new infrastructure: +1. New parser support for `cbind(...)` on the LHS of `~`. +2. New `LikelihoodColumn`-like type that holds a tuple of data columns and a multivariate distribution. +3. `llikelihood!` calls `logpdf(MvNormal(loc_vec, Σ), [y1[i], y2[i]])` per row. +4. `Σ` is a new parameter block: a Cholesky factor over the outcomes (separate from the random-effects Cholesky). + +**Verification.** Translate `mcelreath::waffle_divorce_multivariate` directly. Compare fitted parameters against the published reference. +""", + "3.2 inferred predictors / measurement error me()" => raw""" +**What it is.** brms's `me(x_obs, sd_x)` says "the predictor `x_obs` is itself measured with error of size `sd_x`; sample the latent true value during inference". The model sees both the observed value and the latent. + +**Why it matters.** Standard regression treats predictors as fixed/known. When predictors are themselves estimates (e.g. from a previous study or a noisy sensor), ignoring measurement error biases the slope estimates toward zero. `me()` is the principled fix. + +**Implementation.** Bigger architectural change: predictor columns become latent variables sampled during inference, not data columns evaluated once. Needs: +- A new column type analogous to `MissingColumn` but with an observation-driven prior `Normal(x_obs, sd_x)`. +- The latent column gets a slot in the population block (one parameter per row). +- `lprior!` adds the per-row Normal prior contribution. +- `vbroadcasted` resolves the column to the latent values, not the observed ones. + +**Verification.** Preset against synthetic data where the true `x` is known but only a noisy observed version is in the dataframe. Compare slope estimates with and without `me()`. +""", + "3.3 ordinal predictors mo() (monotonic effects)" => raw""" +**What it is.** brms's `mo(x)` for an ordinal predictor with K levels: instead of `K-1` independent treatment-coded coefficients, fit a single "total effect" β plus a `K-1`-dim simplex of inter-level shape. Forces the effect to be monotonic in the ordering of `x`'s levels. + +**Why it matters.** Likert-scale predictors and ordered categorical inputs (e.g. age groups) have a natural ordering that treatment coding ignores. `mo()` enforces the monotonicity prior, dramatically reducing the parameter count and tightening posterior inference. + +**Implementation.** New block layout: one β coefficient + one Dirichlet-distributed simplex of length `K-1`. `_cat_lookup`-style materialization but the per-level contribution is `β * cumulative_sum(simplex)[level]` instead of `coefficients[level]`. + +Needs a new prior block type (Dirichlet) in `lprior!`, plus new parser support for `mo(x)` and a new `vmeta_sampling_rhs` overload. + +**Verification.** Preset against synthetic data where the true effect is monotonic but the levels are unordered in the data. Compare fitted shape parameters against the synthetic monotonic curve. +""", + "3.4 ordinal outcomes (proportional odds)" => raw""" +**What it is.** When `y` is itself ordered categorical (Likert response, severity grades, …), use a cumulative-link model: `Pr(y ≤ k) = logistic(α_k - η)` where `α_k` are K-1 cutpoints and η is the linear predictor. The likelihood is the difference of consecutive CDFs. + +**Why it matters.** Ordinal outcomes are common in survey data, clinical scoring, and any "rating" task. Treating them as continuous is statistically wrong; treating them as nominal categorical loses the ordering information. + +**Implementation.** New likelihood family with a vector of cutpoints as additional parameters. `Distributions.jl` has `OrderedLogistic` already — the pass-through path should mostly handle it once the parser knows to extract cutpoints from a `cumulative` wrapper. + +The cutpoints need an ordered prior (e.g. ordered transform of unconstrained reals), which means a new prior block type — similar to (3.3)'s simplex. + +**Verification.** Preset against synthetic Likert data. Compare cutpoints against `polr` from R's MASS package. +""", + "3.5 mixture models" => raw""" +**What it is.** brms's `mixture(Normal, Normal)` lets the likelihood be a weighted mixture of K component distributions, with mixing weights estimated as parameters. + +**Why it matters.** Heterogeneous populations, latent class analysis, robust regression (Normal + heavy-tailed component), zero-inflated outcomes, … + +**Implementation.** Extends the existing `Distribution` pass-through. Need a new `MixtureModel` wrapper that holds component distributions plus a weights parameter block. `llikelihood!` uses `logsumexp(log_weights .+ logpdf.(components, y))` per row. + +Composes with (3.4 ordered priors) for the mixing weights' Dirichlet-like prior. + +**Verification.** Preset against synthetic two-component-Normal data. Confirm the recovered mixture weights and component parameters. +""", + "3.6 splines / GP submodels s(), bs(), gp(), t2()" => raw""" +**What it is.** Smoothers in the linear predictor: `s(x)` for a generic spline, `bs(x, knots=...)` for a B-spline basis, `gp(x)` for a Gaussian process, `t2(x, y)` for a tensor-product spline. brms / mgcv use these heavily. + +**Why it matters.** Nonlinear effects without committing to a specific functional form. The de facto way to model dose-response curves, time effects, growth curves, spatial trends, … + +**Implementation.** Each smoother is a basis-matrix builder that grows a population block by `n_basis` columns and stores the basis matrix as part of `meta`. Function stubs (`s`, `bs`, `t2`, `gp`) already exist in `scripts/parsing.jl` so the parser side is partly done. + +For each smoother type, the materialization is `basis_matrix * coefficients` (length-N output). The smoothness prior is a structured prior on the coefficients (typically a Gaussian prior with a banded or 2D-difference penalty matrix), which requires (2.3 per-parameter priors) as a prerequisite. + +**Verification.** Preset against synthetic curve data. Compare fitted smoother against `mgcv::gam`. +""", + "3.7 autoregressive submodels ar(), ar1()" => raw""" +**What it is.** Add an AR(p) structure to the residuals: `y_t = η_t + φ * (y_{t-1} - η_{t-1}) + ε_t`. brms's `ar(time, p=1)` specifies the order and the time variable. + +**Why it matters.** Time-series and repeated-measures data routinely have autocorrelated errors. Ignoring AR structure inflates the effective sample size and gives overconfident posteriors. + +**Implementation.** Bigger than it looks because the likelihood is no longer per-row independent — it's a chain. Need a new `LikelihoodColumn`-like type that holds the time index and walks the data in time order, accumulating the AR contribution row by row. + +Composes with (2.5 grouped random effects) for per-subject AR structure. + +**Verification.** Preset against synthetic AR(1) data. Recover φ. +""", + "3.8 decompositions (QR, orthogonal polar)" => raw""" +**What it is.** Numerical-stability transformations of the population design matrix. brms uses QR decomposition on the design matrix internally so the sampler sees an orthogonal-columns version, then transforms back at the end. Stan does the same. + +**Why it matters.** When population covariates are correlated (which is the norm), the unrotated design matrix gives a poorly-conditioned posterior that NUTS struggles with. QR fixes this with no statistical change. + +**Implementation.** Mostly orthogonal to the formula DSL — happens at `VBRMI` build time. Add a per-block transform: store both the original design matrix and the QR factor, run sampling on the rotated parameter space, transform back when extracting coefficients. + +Could be implemented today without any parser changes, as a `qr_transform=true` flag on `VBRMI`. Lift to a default-on once verified. + +**Verification.** Same model with and without QR should produce identical logdensity values, but the gradient should be better-conditioned (smaller condition number on the Hessian). +""", + "3.9 spike-and-slab / Horseshoe priors" => raw""" +**What it is.** Sparsity-inducing priors for high-dimensional regression. Spike-and-slab puts a delta-spike at zero plus a wide slab; Horseshoe uses a half-Cauchy hyperprior on a per-coefficient scale, producing a heavy-tailed shrinkage prior. + +**Why it matters.** Without sparsity priors, high-dimensional regressions overfit. These are the standard solution in Bayesian variable selection and high-dim genomics / finance. + +**Implementation.** Depends entirely on (2.3 per-parameter priors). Once that's in place, spike-and-slab is `prior = Mixture(Normal(0, ε), Normal(0, slab_scale))` per coefficient, and Horseshoe is `Normal(0, λ_i * τ)` with `λ_i ~ HalfCauchy(0, 1)` and `τ ~ HalfCauchy(0, 1)` — both expressible in the existing Distribution pass-through once per-coefficient priors are wired up. + +**Verification.** Preset on a sparse synthetic regression (mostly-zero true coefficients with a few large ones). Confirm the Horseshoe-fitted coefficients shrink the noise toward zero and preserve the signal. +""", + "3.10 Dirichlet process / non-parametric models" => raw""" +**What it is.** Models where the number of components / clusters / random-effect levels is itself inferred during sampling, via a Dirichlet process or stick-breaking prior. + +**Why it matters.** When you don't know how many clusters are in your data, fixing K is itself a strong assumption. DP priors let the model decide. + +**Implementation.** The heaviest item on the list. Needs sampling-time level inference, a stick-breaking parameter block, and a different `growblock!!` that grows during sampling rather than at `VBRMI` build time. Probably requires a different `lprior!` interface entirely. + +Defer until everything else is solid. + +**Verification.** Preset against synthetic data with an unknown number of latent clusters. Compare recovered K against the truth. +""", + "3.11 zero-inflated / hurdle likelihoods" => raw""" +**What it is.** ZI Poisson, ZI Negative Binomial, hurdle Poisson, hurdle Gamma, … — likelihoods that mix a point mass at zero (or a separate "is zero" Bernoulli) with a continuous/count distribution for the nonzero values. + +**Why it matters.** Count data with excess zeros (insurance claims, species abundance, healthcare utilization) is everywhere. Standard Poisson / NegBin underfits the zero count. + +**Implementation.** Should mostly work through the existing `Distribution` pass-through once we use `Distributions.jl`'s ZI distributions (or write small wrappers). The mixing weight needs its own linear predictor, which is just another distributional-regression-style `~` line. + +**Verification.** Preset against synthetic ZI Poisson data. Confirm the recovered zero-inflation probability matches the synthetic generator. +""", +] + function __init__() route!(AppContext()) end diff --git a/web-macro/src/macro.jl b/web-macro/src/macro.jl index d6baf3c..6f30dc6 100644 --- a/web-macro/src/macro.jl +++ b/web-macro/src/macro.jl @@ -194,12 +194,22 @@ struct BRMI{O<:NamedTuple} operations::O end BRMI(;kwargs...) = BRMI((;kwargs...)) -Base.show(io::IO, (;operations)::BRMI) = begin +Base.show(io::IO, (;operations)::BRMI) = begin print(io, "BRMI:\n") for (key, value::NamedColumn) in pairs(operations) - print(io, " ", key, ": ", parent(value), "\n") + print(io, " ") + _show_top(io, key, parent(value)) + print(io, "\n") end end +# Top-level entries in a BRMI listing are either ExprColumns whose own show +# already names the LHS (`loc ~ ...`, `(err = ...)`) or other column types +# (DataColumn, MaterializedColumn, ...) whose show prints just the value with +# no name. For the former we strip the outermost parens; for the latter we +# prefix with the operation key so the name doesn't get lost. +_show_top(io::IO, key, op) = print(io, key, ": ", op) +_show_top(io::IO, key, op::ExprColumn{<:Union{typeof(~),typeof(assign)}}) = + join(io, getargs(op), " $(getop(op)) ") Base.show(io::IO, d::DataColumn) = begin print(io, "data (eltype=", eltype(parent(d)), ")") end diff --git a/web-macro/src/vimpl.jl b/web-macro/src/vimpl.jl index bfc51f3..009e08a 100644 --- a/web-macro/src/vimpl.jl +++ b/web-macro/src/vimpl.jl @@ -28,6 +28,9 @@ vbroadcasted(;kwargs...) = (args...)->vbroadcasted(args...; kwargs...) vbroadcasted(x::NamedColumn{<:Any,<:DataColumn}; meta) = parent(meta.materialized[name(x)]) vbroadcasted(x::NamedColumn; meta) = meta.materialized[name(x)] vbroadcasted(x::ExprColumn; meta) = Base.broadcasted(getf(x), map(vbroadcasted(;meta), getargs(x))...) +# Literals (Int, Float, etc.) inside formula expressions like `a^2` pass +# through unchanged — they get broadcasted as scalars by Base.broadcasted. +vbroadcasted(x::Number; meta) = x getinverse(x::ExprColumn{<:Any,<:Tuple{<:Any}}) = inverse(getf(x)) getinverse(x::ExprColumn{<:Any,<:Tuple{<:ExprColumn}}) = inverse(getf(x)) ∘ getinverse(getargs(x, 1)[1]) vmeta_sampling(meta, lhs::ExprColumn, rhs) = begin diff --git a/web-macro/todos/1.1-verify-bernoulli-binomial-done.jl b/web-macro/todos/1.1-verify-bernoulli-binomial-done.jl new file mode 100644 index 0000000..6aa9a26 --- /dev/null +++ b/web-macro/todos/1.1-verify-bernoulli-binomial-done.jl @@ -0,0 +1,15 @@ +# label: 1.1 verify Bernoulli/Binomial +# tier: 1 +# status: done +#= +**Status: done.** ✓ Confirmed that the existing `FBroadcasted{<:Type{<:Distribution}}` pass-through in `vimpl.jl` handles both Bernoulli and Binomial cleanly. + +**Verification.** The form below loads the **cbpp + therapeutic touch** model — a faithful translation of `brms::cbpp_binomial` (categorical predictor + random intercept + Binomial with per-row trial counts) and `kruschke::therapeutic_touch` (hierarchical Bernoulli) into one multi-likelihood model. brms's `incidence | trials(size) ~ ...` sidecar collapses to a plain positional argument: `bin_succ ~ Binomial(bin_n, logistic(η))`. If the gradient sanity check stays green every other `Distribution` family (Beta, Gamma, NegBinomial, …) should be free as well. + +=# + +log_odds_bin ~ 1 + c1 + (1 | g1) +bin_succ ~ Binomial(bin_n, logistic(log_odds_bin)) + +log_odds_b ~ 1 + (1 | g1) +bin_y ~ Bernoulli(logistic(log_odds_b)) diff --git a/web-macro/todos/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl b/web-macro/todos/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl new file mode 100644 index 0000000..9ed9e36 --- /dev/null +++ b/web-macro/todos/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl @@ -0,0 +1,21 @@ +# label: 1.2 offset / fixed exposure — already works without a wrapper +# tier: 1 +# status: done +#= +**Status: already works without any new code.** brms needs `offset(z)` because R's formula syntax has no other way to put a "no-coefficient term" into the linear predictor — the only thing on the RHS of `~` is the formula DSL. In our DSL the linear predictor and the likelihood are *separate* `~` lines, and the second one (the likelihood) takes a free-form Julia expression. Anything inside that expression gets evaluated as plain code at materialization time via `vbroadcasted` — function calls dispatch to whatever Julia function the symbol resolves to, and data column references are pulled from the dataframe. + +So instead of `count ~ x + offset(log(exposure))`, you write the offset directly inside the likelihood: + +```julia +loc ~ 1 + a +k1 ~ Poisson(exp(loc + log(exposure))) +``` + +The `log(exposure)` here is just `Base.log` applied to the `exposure` data column, broadcasted across rows and added to `loc` (which is the materialized linear predictor). No parameter is allocated for it because `growblock!!` is never called for that branch — there's no `~` on the data side, just an argument to `Poisson(...)`. + +**Verification.** Form below loads exactly that model. The VBRMI dim should match the offset-free version (only the population intercept + slope on `a`); the gradient sanity check should stay green; and the materialized `k1` likelihood should incorporate the row-specific exposure shift. + +=# + +loc ~ 1 + a +k1 ~ Poisson(exp(loc + log(exposure))) diff --git a/web-macro/todos/1.3-i-expr-likely-already-works.jl b/web-macro/todos/1.3-i-expr-likely-already-works.jl new file mode 100644 index 0000000..a1aac08 --- /dev/null +++ b/web-macro/todos/1.3-i-expr-likely-already-works.jl @@ -0,0 +1,16 @@ +# label: 1.3 I(expr) — likely already works +# tier: 1 +# status: done +#= +**What it is.** brms's `I()` is a literal-escape: `I(x^2)` says "compute `x^2` from the data and treat it as a single column". brms needs it because `+`, `*`, `:`, `|`, … all have special meaning inside an R formula. + +**Why we probably don't need it.** Our DSL is parsed by Julia first, then walked by `_x`. `_x` recursively wraps every `Expr(:call, f, args...)` in an `ExprColumn`, regardless of whether `f` is special. So `loc ~ a + x^2` becomes `+(a, ^(x, 2))` → `ExprColumn(+, NamedColumn(:a), ExprColumn(^, NamedColumn(:x), 2))`. The `^` is just another function call, no special handling needed. + +The only operators that have DSL meaning in our system are `~` (sampling), `=` (assignment), and `|` / `||` inside random-effects specs. Everything else (`^`, `/`, `sqrt`, `log`, `exp`, `mod`, `min`, `max`, …) is a regular function call resolved at materialization time via `vbroadcasted`. + +**Verification.** The form below loads a model with three nonlinear terms (`a^2`, `sqrt(abs(b))`, `log(exposure)`) directly as population-level covariates. The VBRMI dim should match the number of distinct terms; the gradient sanity check should be all-active. If it works, that confirms `I()` is unnecessary because Julia function calls are first-class on the formula RHS. + +=# + +loc ~ 1 + a + a^2 + sqrt(abs(b)) + log(exposure) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/1.4-scale-x-standardize-x.jl b/web-macro/todos/1.4-scale-x-standardize-x.jl new file mode 100644 index 0000000..1a7218f --- /dev/null +++ b/web-macro/todos/1.4-scale-x-standardize-x.jl @@ -0,0 +1,28 @@ +# label: 1.4 scale(x) / standardize(x) +# tier: 1 +# status: deprioritized +#= +**What it is.** brms's `scale(x)` z-transforms a column at parse time: `scale(x) = (x - mean(x)) / std(x)`. The model sees the standardized column. Crucial for default priors (which are scale-invariant only after standardization) and sampler stability (well-conditioned linear predictors). + +**Why it matters.** Most brms vignettes do `scale(x)` automatically as a convenience. Without it, every formula has to either manually z-transform the data or accept poorly-scaled coefficients. + +**Implementation.** +1. Add `function scale end` (and `function center end`, `function standardize end`) to `macro.jl`. +2. Add a `vmeta_sampling_rhs` overload in `vimpl.jl`: +```julia +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(scale)}; group) = begin + inner = vbroadcasted(only(getargs(x)); meta) + materialized = Base.materialize(inner) + z = (materialized .- Statistics.mean(materialized)) ./ Statistics.std(materialized) + vmeta_sampling_rhs(meta, z; group) +end +``` +The standardization happens once when the BRMI is materialized into a VBRMI. Composes with the existing dense-map caching TODO. +3. Add `Statistics` to `vimpl.jl`'s using-list (or vendor `mean`/`std` inline). + +**Verification.** Preset: `loc ~ 1 + scale(a) + scale(b); y1 ~ Normal(loc, 1)`. Compare against the unscaled version: same dim, different posterior geometry. The fitted coefficients should be ≈ the unscaled coefficients × std(x). + +=# + +loc ~ 1 + scale(a) + scale(b) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/1.5-zerocorr-independent-random-effects.jl b/web-macro/todos/1.5-zerocorr-independent-random-effects.jl new file mode 100644 index 0000000..938c869 --- /dev/null +++ b/web-macro/todos/1.5-zerocorr-independent-random-effects.jl @@ -0,0 +1,30 @@ +# label: 1.5 zerocorr — independent random effects +# tier: 1 +# status: open +#= +**What it is.** brms (via lme4 syntax) lets you opt out of the LKJ correlation between multiple random terms in the same group. `(1 + x || group)` (double bar) says "estimate the random intercept and the random slope independently — don't fit a 2×2 Cholesky factor between them". Useful when there isn't enough data to estimate the correlations, or when you have prior reason to believe the terms are uncorrelated. + +**Why it matters.** Multi-term random specs are common, and the LKJ correlation often dominates the prior cost without much identifiability. Letting users skip it is a meaningful sampling speedup and prior simplification. + +**Implementation.** Our `_x` walker already wraps `||` as `ExprColumn{typeof(doublepipe)}`. Add a `vmeta_sampling_rhs` overload that splits each term inside the `||` LHS into its own block (with a synthetic per-term key like `Symbol(group_name, :__nocor__, term_index)`): + +```julia +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(doublepipe)}; kwargs...) = begin + lhs, rhs = getargs(x, 2) + terms = lhs isa ExprColumn{typeof(+)} ? getargs(lhs) : (lhs,) + foldl(enumerate(terms); init=(meta, ())) do (m, args), (i, term) + nocor_key = NamedColumn(Symbol(name(rhs), :__nocor__, i), parent(rhs)) + m, arg = vmeta_sampling_rhs(m, term; group=nocor_key) + m, (args..., arg) + end |> ((m, args),) -> (m, Base.broadcasted(+, args...)) +end +``` + +Each per-term block ends up as 1×1 with one `log_scale` Cholesky parameter — `lprior!`'s existing single-column path handles this with no changes. + +**Verification.** Preset: `loc ~ 1 + (1 + a || g1); y1 ~ Normal(loc, 1)`. Compare its dim against the correlated `(1 + a | g1)` version: the correlated version has 3 Cholesky params (1+2/2 for a 2×2), the uncorrelated version has 2 (one log_scale per term). Same direct-parameter count (2 cols × 8 levels = 16) either way. + +=# + +loc ~ 1 + (1 + a || g1) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/1.6-cache-levels-level_map-dense-gc_idx.jl b/web-macro/todos/1.6-cache-levels-level_map-dense-gc_idx.jl new file mode 100644 index 0000000..15447b8 --- /dev/null +++ b/web-macro/todos/1.6-cache-levels-level_map-dense-gc_idx.jl @@ -0,0 +1,17 @@ +# label: 1.6 cache levels / level_map / dense / gc_idx +# tier: 1 +# status: deprioritized +#= +**What it is.** Stop rebuilding the dense level mapping (`Dict(level => row_index)`) and the gc_idx vector on every `VBRMI(brmi)` call. Cache them once per source data column. + +**Why it matters.** Today every `VBRMI` build re-traces the categorical / grouping columns, sorts unique values, builds a Dict, and walks the column to dense-encode it. For models with many categorical columns or many `VBRMI` rebuilds (e.g. during AD), this adds up. + +**Implementation.** Pick a storage layout for the per-column metadata. Two candidates: +- A new `meta.factor` NamedTuple keyed by source column name, holding `(; levels, level_map, dense, gc_idx)` per column. Built lazily on first reference, indexed via `name(column)`. +- Attach the metadata to `meta.materialized[column_name]` directly. More tightly coupled but avoids a parallel NamedTuple. + +The TODO already lives at `vimpl.jl:78–86`. Once a layout is picked, refactor `_gc_idx` and the inline dense map in the categorical path to read from the cache, falling back to a build-on-miss helper. + +**Verification.** No behavioral change — the gradient sanity check should stay green. Benchmark `VBRMI(brmi)` with Chairmarks before/after and confirm a measurable speedup on a model with multiple categorical/grouping columns. + +=# diff --git a/web-macro/todos/1.7-categoricalarrays-pooledarrays-integration.jl b/web-macro/todos/1.7-categoricalarrays-pooledarrays-integration.jl new file mode 100644 index 0000000..0240eca --- /dev/null +++ b/web-macro/todos/1.7-categoricalarrays-pooledarrays-integration.jl @@ -0,0 +1,19 @@ +# label: 1.7 CategoricalArrays / PooledArrays integration +# tier: 1 +# status: deprioritized +#= +**What it is.** When the input column is already a `CategoricalVector` or a `PooledArray`, the dense level mapping is already computed and stored in the column's `.refs` field. Use it directly instead of rebuilding via `Dict`. + +**Why it matters.** Most real-world DataFrames use `CategoricalArrays.jl` for factor columns. Skipping the rebuild eliminates allocation entirely for the common case and gets us "for free" interop with the standard categorical-data ecosystem. + +**Implementation.** Two design choices: +- **Hard dep**: add `CategoricalArrays` to vimpl.jl's deps, dispatch on `CategoricalVector`, read `levelcode.(col)` and `levels(col)` directly. +- **Duck-typed**: sniff for the `.refs` field and `levels` method without importing the package, falling back to the generic Dict path. + +Recommend hard dep — it's the standard for tabular Julia code, and the duck-type path is more code with no real win. Same for `PooledArrays`. + +The actual integration is small once the design is picked: a method specialization in `_gc_idx` and in the categorical-predictor path. Composes with the caching TODO above. + +**Verification.** Preset (or test) that builds a DataFrame with a `CategoricalVector` column and uses it as a grouping factor / categorical predictor. Confirm the gradient sanity check stays green and the per-VBRMI allocation count drops. + +=# diff --git a/web-macro/todos/2.1-interactions-a-b-a-b.jl b/web-macro/todos/2.1-interactions-a-b-a-b.jl new file mode 100644 index 0000000..bc026cc --- /dev/null +++ b/web-macro/todos/2.1-interactions-a-b-a-b.jl @@ -0,0 +1,23 @@ +# label: 2.1 interactions a:b, a*b +# tier: 2 +# status: open +#= +**What it is.** brms's `a:b` is the elementwise interaction term (a single coefficient multiplying `a[i] * b[i]`). `a*b` is the "main effects + interaction" shorthand: it desugars to `a + b + a:b`. + +**Why it matters.** Interactions are the most commonly missed feature in regression DSLs. Without them, every model that needs `a:b` has to manually create the interaction column in the input DataFrame. + +**Implementation.** +1. **Parser side.** Add a `:` case to `_x` so that `a:b` becomes `ExprColumn(:, NamedColumn(:a), NamedColumn(:b))` instead of falling through to a Symbol/Range parse error. +2. **Materialization side.** Add `vmeta_sampling_rhs(meta, x::ExprColumn{typeof(:)}; group)` that elementwise-multiplies the operands and dispatches to the float-vector path. For continuous × continuous it's a single coefficient on `a .* b`; for categorical × continuous it's `(k-1)` coefficients (one per non-reference level of the categorical, multiplied by the continuous); for categorical × categorical it's `(k₁-1)*(k₂-1)` coefficients via a 2D `_cat_lookup`. +3. **`a*b` desugaring.** At parse time in `_x`, rewrite `*` between formula terms as `+(a, b, :(a:b))`. This needs care because `*` also means multiplication elsewhere (e.g. `Normal(0, 2*sigma)`); the rewrite should only apply at formula-RHS top-level. + +**Verification.** Presets exercising each interaction type: +- continuous×continuous: `loc ~ 1 + a + b + a:b; y1 ~ Normal(loc, 1)` → dim 4 +- continuous×categorical: `loc ~ 1 + a + c1 + a:c1; y1 ~ Normal(loc, 1)` → dim 6 (1 + 1 + 2 + 2) +- categorical×categorical: `loc ~ 1 + c1 + c2 + c1:c2; y1 ~ Normal(loc, 1)` → dim 5 (1 + 2 + 1 + 2) +- shorthand: `loc ~ 1 + a*b; y1 ~ Normal(loc, 1)` should match `loc ~ 1 + a + b + a:b` exactly. + +=# + +loc ~ 1 + a + b + a:b +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/2.2-configurable-categorical-reference-level.jl b/web-macro/todos/2.2-configurable-categorical-reference-level.jl new file mode 100644 index 0000000..d9ff19b --- /dev/null +++ b/web-macro/todos/2.2-configurable-categorical-reference-level.jl @@ -0,0 +1,19 @@ +# label: 2.2 configurable categorical reference level +# tier: 2 +# status: deprioritized +#= +**What it is.** Currently the reference level for treatment-coded categoricals is `sort(unique(x))[1]`. brms / lme4 let you override this via `factor(x, ref="some_level")` or by reordering the factor's levels. + +**Why it matters.** The reference level changes the interpretation of the intercept (it becomes "the mean for the reference level") and of the coefficients (each becomes "the difference from reference"). For some analyses, changing the reference is the only way to make the coefficients directly answer the research question. + +**Implementation.** +1. Add `function factor end` to `macro.jl`. +2. Either store the override at parse time (rewrite `factor(x, ref=:level3)` into a wrapper that the materializer recognizes) or at materialization time via a `meta.factor_ref` NamedTuple keyed by column name. +3. The categorical-predictor path's `levels = sort(unique(x))` becomes `levels = sort(unique(x), by=l -> l == ref ? -Inf : l)` so the chosen reference always sorts first. + +**Verification.** Preset: `loc ~ 1 + factor(c1, ref=2); y1 ~ Normal(loc, 1)`. Compare the fitted coefficients against the default-reference version — they should differ by the level-2-vs-level-1 mean shift but produce the same logdensity. + +=# + +loc ~ 1 + factor(c1, ref=2) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/2.3-per-parameter-prior-scales.jl b/web-macro/todos/2.3-per-parameter-prior-scales.jl new file mode 100644 index 0000000..6fa369f --- /dev/null +++ b/web-macro/todos/2.3-per-parameter-prior-scales.jl @@ -0,0 +1,28 @@ +# label: 2.3 per-parameter prior scales +# tier: 2 +# status: deprioritized +#= +**What it is.** Currently every parameter is `Normal(0, 1)` in `lprior!`. brms / Stan-style models routinely set custom priors per coefficient: `b ~ Normal(0, 0.5)` for tight priors on slopes, `b ~ Cauchy(0, 1)` for heavy-tailed priors, etc. + +**Why it matters.** Default `Normal(0, 1)` is fine after standardization but poor on raw scales. Allowing per-parameter prior scales is the prerequisite for spike-and-slab, Horseshoe, and most prior sensitivity analyses. Without it, users have no way to express domain knowledge about parameter magnitudes. + +**Implementation.** This is the largest design decision in Tier 2 because it has knock-on effects for every other prior-related TODO. + +Two storage candidates: +- **Per-block scales**: extend `meta.block_data[group]` with a per-column scale vector. `lprior!` multiplies the standard-normal draw by the scale before storing. Simple but only handles Normal-with-scale priors. +- **Per-block prior distributions**: store a vector of `Distribution` objects per block. `lprior!` calls `logpdf(prior_i, xi)` for each parameter. More general; handles Cauchy, StudentT, Horseshoe, etc. + +Recommend the second — it's strictly more powerful and the runtime cost is identical (one `logpdf` call per parameter). Default value is `Normal(0, 1)` for backward compatibility. + +**Implementation sketch.** +1. Extend `meta.block_data` with a `priors` field per block. +2. The macro syntax `b ~ Normal(0, 0.5)` parses as a sampling statement with a Distribution-typed RHS. Currently this is reserved for likelihood declarations; it would need a new "is this a prior or a likelihood?" branch in `vmeta_sampling`. Likelihood: LHS is a data column. Prior: LHS is a maybelocal (parameter). +3. `lprior!` reads the per-column prior and calls `logpdf(prior, value)` instead of the hard-coded `logpdf(Normal(), value)`. + +**Verification.** Preset: `loc ~ 1 + a; b ~ Normal(0, 0.1); y1 ~ Normal(loc, 1)` — confirm the gradient is dampened on `b` compared to the default-prior version, and that the dead-param check still passes. + +=# + +coef_a ~ Normal(0, 0.1) +loc ~ coef_a * a +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/2.4-centered-non-centered-parameterization-toggle.jl b/web-macro/todos/2.4-centered-non-centered-parameterization-toggle.jl new file mode 100644 index 0000000..be83585 --- /dev/null +++ b/web-macro/todos/2.4-centered-non-centered-parameterization-toggle.jl @@ -0,0 +1,18 @@ +# label: 2.4 centered / non-centered parameterization toggle +# tier: 2 +# status: deprioritized +#= +**What it is.** Currently every random-effect block uses non-centered parameterization (we sample standard normals and apply `mul!(vi, C.L, xi)`). brms / Stan let you choose centered (sample directly from `Normal(0, σ)` per group) on a per-factor basis. + +**Why it matters.** Non-centered is the default for "weak data per group" cases (Neal's funnel pathology), but for "strong data per group" cases centered samples better. Letting users choose is a meaningful sampling speedup for the latter regime. + +**Implementation.** Small change to `lprior!` and `growblock!!`. Add a `centered::Bool` flag to `meta.block_data[group]`. In `lprior!`'s non-population branch, if the block is centered, sample directly from `Normal(0, exp(log_scale))` instead of `Normal(0, 1)` then multiplying by `L`. The Cholesky machinery for off-diagonal correlations still applies in the centered case — just on the column before the variance scaling rather than after. + +Once (2.3) is in place, the centered/non-centered choice could be encoded as `(1 | g) ~ Normal(0, σ)` (centered) vs the implicit non-centered default — but for now a per-block kwarg or a wrapper function (e.g. `centered((1 + a | g))`) is simpler. + +**Verification.** Same model with both parameterizations should produce the same logdensity at the same parameter values (after the appropriate change of variables). Sampling efficiency on a known-funnel dataset should differ. + +=# + +loc ~ 1 + (1 + a | centered(g1)) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/2.5-grouped-random-effects-per-factor-variance.jl b/web-macro/todos/2.5-grouped-random-effects-per-factor-variance.jl new file mode 100644 index 0000000..df65ed8 --- /dev/null +++ b/web-macro/todos/2.5-grouped-random-effects-per-factor-variance.jl @@ -0,0 +1,23 @@ +# label: 2.5 grouped random effects (per-factor variance) +# tier: 2 +# status: open +#= +**What it is.** Peter's "different variance by diagnosis" pattern: `(1 | subject) gr(diagnosis)` says "the random intercept by subject has a different variance per diagnosis level". In brms this is a custom group structure where the variance hyperparameter itself depends on a second factor. + +**Why it matters.** Common in clinical data where treatment groups have intrinsically different between-subject variability. Without this, you have to fit separate models per diagnosis or accept a single pooled variance. + +**Implementation.** Bigger than it looks because the variance is no longer a single scalar but a length-`n_levels(diagnosis)` vector that needs its own prior and its own gradient. + +Proposed shape: +- A new `gr(group_factor)` wrapper recognized in the `~` RHS via a `function gr end` stub (already exists in `macro.jl`). +- The wrapped block stores `n_levels(group_factor)` log-scale parameters instead of one. `lprior!` walks them, multiplying each subject's random intercept by the diagnosis-specific scale. +- Requires the gc_idx for the inner factor (subject) AND for the outer factor (diagnosis) — both vectors of length N. + +This composes naturally with (1.6 caching) and (2.3 per-parameter priors). + +**Verification.** Preset against synthetic data with two grouping factors, one nested inside the other, with intentionally different per-outer-level variance. Confirm the fitted scales recover the synthetic values. + +=# + +loc ~ 1 + (1 | gr(g1, g2)) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/2.6-multi-membership-random-effects-mm.jl b/web-macro/todos/2.6-multi-membership-random-effects-mm.jl new file mode 100644 index 0000000..8f0fe11 --- /dev/null +++ b/web-macro/todos/2.6-multi-membership-random-effects-mm.jl @@ -0,0 +1,20 @@ +# label: 2.6 multi-membership random effects mm() +# tier: 2 +# status: deprioritized +#= +**What it is.** brms's `mm(g1, g2, ...)` lets one observation belong to **multiple** levels of the same random factor simultaneously, with weights summing to 1. Standard use: a student belongs to multiple schools across the year, and we want their random effect to be a weighted average of the per-school effects. + +**Why it matters.** Standard random effects assume each observation belongs to exactly one group. Multi-membership is the only clean way to handle observations that span groups (mobile students, patients seen by multiple clinicians, etc.). + +**Implementation.** `_gc_idx` would have to return a row-of-vectors instead of a single Int per row. Two paths: +- **Sparse design matrix**: replace the `gc_idx` lookup with a sparse `(N × n_levels)` matrix where each row's nonzero entries are the membership weights. The materialized random effect becomes `sparse_membership * random_effects_vector`. +- **Per-row lookup loop**: keep the row-major view but make `_re_lookup` iterate the membership list per row, summing weighted contributions. + +The sparse matrix approach is more memory-efficient and SIMD-friendly. Needs a new wrapper in the formula syntax: `(1 | mm(g1, g2; weights=...))`. + +**Verification.** Preset against synthetic data where each observation has 2 random group memberships with weights summing to 1. Compare against the equivalent "fully observed in primary group only" model. + +=# + +loc ~ 1 + (1 | mm(g1, g2)) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl b/web-macro/todos/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl new file mode 100644 index 0000000..8232f06 --- /dev/null +++ b/web-macro/todos/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl @@ -0,0 +1,39 @@ +# label: 2.7 se() / weights() — already works (just a different distribution) +# tier: 2 +# status: done +#= +**Status: already works without any new code.** Same insight as 1.2 (offset) and 1.3 (`I()`): brms needs sidecar syntax (`y | se(sigma_y) ~ ...`, `y | weights(w) ~ ...`) because R's formula DSL has no other way to attach extra info to the LHS. Our DSL has no such constraint — the likelihood is a free-form Julia expression, so anything brms expresses via sidecar syntax we can express via constructor arguments or a wrapper distribution. + +**Meta-analysis (per-observation known SE).** Just pass the SE column as the scale argument to `Normal`: + +```julia +loc ~ 1 + a +y1 ~ Normal(loc, exposure) +``` + +`exposure` is a length-N data column; the broadcasted `Normal(loc, exposure)` constructs a per-row Normal at materialization time, and `logpdf(Normal(loc[i], exposure[i]), y1[i])` is what `llikelihood!` ends up summing. No new code in `vimpl.jl`. + +**Weighted regression (per-observation likelihood weights).** Define a thin wrapper distribution that multiplies the underlying logpdf by a weight, and use the existing `FBroadcasted{<:Type{<:Distribution}}` pass-through: + +```julia +struct WeightedLikelihood{D, W} <: Distributions.ContinuousUnivariateDistribution + base::D + weight::W +end +Distributions.logpdf(w::WeightedLikelihood, y) = w.weight * logpdf(w.base, y) + +# Then in the formula: +loc ~ 1 + a +y1 ~ WeightedLikelihood(Normal(loc, 1), exposure) +``` + +The wrapper is ~5 lines of plain Julia in the user's own scope, no changes to the DSL or to `vimpl.jl`. The pass-through walks the broadcasted `WeightedLikelihood` constructor, materializes `loc`/`exposure` per row, and `llikelihood!` calls the user-defined `logpdf` per row. + +**Conclusion.** Neither feature needs special parser or vimpl support. `se()` is a constructor argument; `weights()` is a one-page user-defined distribution. The general principle: anything brms expresses as an LHS sidecar, our DSL expresses as a positional argument or a user-supplied distribution. + +**Verification.** The form below loads the meta-analysis case directly. The gradient sanity check should be all-active and the per-row Normal scales should pick up the `exposure` column. + +=# + +loc ~ 1 + a +y1 ~ Normal(loc, exposure) diff --git a/web-macro/todos/3.1-multivariate-outcomes-cbind-y1-y2.jl b/web-macro/todos/3.1-multivariate-outcomes-cbind-y1-y2.jl new file mode 100644 index 0000000..63e6cc0 --- /dev/null +++ b/web-macro/todos/3.1-multivariate-outcomes-cbind-y1-y2.jl @@ -0,0 +1,39 @@ +# label: 3.1 multivariate outcomes — uncorrelated free, correlated needs MvNormal +# tier: 3 +# status: open +#= +**What it is.** brms's `mvbind(y1, y2) ~ x + (1 | g)` declares two outcomes sharing the same linear predictor structure with **correlated residuals** (joint MvNormal with a covariance matrix to estimate). brms also supports combining multiple per-outcome formulas via `bf(y1 ~ x) + bf(y2 ~ x)`, but **`+` on the LHS of `~` is not a thing in brms's formula DSL** — both their multivariate syntaxes are sidecars (`mvbind`) or formula-object combinators (`bf + bf`). + +**Half of this is already free in our DSL.** Multiple `~` lines on different data columns produce independent likelihoods that `llikelihood!` sums. If they reference the same intermediate variable, they share that linear predictor automatically — no parser changes needed: + +```julia +loc ~ 1 + a + (1 | g1) +y1 ~ Normal(loc, 1) +y2 ~ Normal(loc, 1) +``` + +That's "multiple uncorrelated outcomes with the same regression formula", already supported. Mixed-family is also free: + +```julia +loc ~ 1 + a + (1 | g1) +y1 ~ Normal(loc, 1) +k1 ~ Poisson(exp(loc)) +``` + +**What's actually missing.** Only the **correlated residuals** case — when the off-diagonal entries of the residual covariance matrix are themselves parameters to estimate (because you want to learn how much shared latent noise the outcomes have). For that you need: + +1. A new likelihood path that takes a *tuple* of data columns and computes a single joint `logpdf(MvNormal(loc_vec, Σ), [y1[i], y2[i], ...])` per row instead of summing marginal logpdfs. +2. A new parameter block holding the Cholesky factor of `Σ` over the outcomes (separate from the random-effects Cholesky on the random-effects block). +3. A way to spell this in our DSL. Two candidate syntaxes: + - `(y1, y2) ~ MvNormal(loc, Sigma)` — Julia parses `(y1, y2)` as a tuple expression, the `~` parser would need a new branch for tuple LHS. + - `mvbind(y1, y2) ~ MvNormal(loc, Sigma)` — explicit wrapper, parser-friendly because it's just a function call (and `mvbind` could be a stub like `gr`/`doublepipe`). + +**Use cases for correlated:** `mcelreath::waffle_divorce_multivariate` is the canonical example. Joint pharmacology/efficacy modeling. Mediation analysis where you model both the mediator and the outcome and want to know how much shared subject-level variance they have. + +**Verification.** The form below loads the **uncorrelated** case (already works). Two `Normal` likelihoods on `y1` and `y2` sharing one `loc` formula. The gradient sanity check should be all-active. For the correlated case, `mvbind(y1, y2) ~ MvNormal(loc, Sigma)` would currently error — that's the half this todo still tracks. + +=# + +loc ~ 1 + a + (1 | g1) +y1 ~ Normal(loc, 1) +y2 ~ Normal(loc, 1) diff --git a/web-macro/todos/3.10-dirichlet-process-non-parametric-models.jl b/web-macro/todos/3.10-dirichlet-process-non-parametric-models.jl new file mode 100644 index 0000000..52e1bf1 --- /dev/null +++ b/web-macro/todos/3.10-dirichlet-process-non-parametric-models.jl @@ -0,0 +1,18 @@ +# label: 3.10 Dirichlet process / non-parametric models +# tier: 3 +# status: deprioritized +#= +**What it is.** Models where the number of components / clusters / random-effect levels is itself inferred during sampling, via a Dirichlet process or stick-breaking prior. + +**Why it matters.** When you don't know how many clusters are in your data, fixing K is itself a strong assumption. DP priors let the model decide. + +**Implementation.** The heaviest item on the list. Needs sampling-time level inference, a stick-breaking parameter block, and a different `growblock!!` that grows during sampling rather than at `VBRMI` build time. Probably requires a different `lprior!` interface entirely. + +Defer until everything else is solid. + +**Verification.** Preset against synthetic data with an unknown number of latent clusters. Compare recovered K against the truth. + +=# + +loc ~ 1 + (1 | dp(g1)) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/3.11-zero-inflated-hurdle-likelihoods.jl b/web-macro/todos/3.11-zero-inflated-hurdle-likelihoods.jl new file mode 100644 index 0000000..ee5af2f --- /dev/null +++ b/web-macro/todos/3.11-zero-inflated-hurdle-likelihoods.jl @@ -0,0 +1,17 @@ +# label: 3.11 zero-inflated / hurdle likelihoods +# tier: 3 +# status: done +#= +**What it is.** ZI Poisson, ZI Negative Binomial, hurdle Poisson, hurdle Gamma, … — likelihoods that mix a point mass at zero (or a separate "is zero" Bernoulli) with a continuous/count distribution for the nonzero values. + +**Why it matters.** Count data with excess zeros (insurance claims, species abundance, healthcare utilization) is everywhere. Standard Poisson / NegBin underfits the zero count. + +**Implementation.** Should mostly work through the existing `Distribution` pass-through once we use `Distributions.jl`'s ZI distributions (or write small wrappers). The mixing weight needs its own linear predictor, which is just another distributional-regression-style `~` line. + +**Verification.** Preset against synthetic ZI Poisson data. Confirm the recovered zero-inflation probability matches the synthetic generator. + +=# + +log_rate ~ 1 + a +zi_logit ~ 1 +k1 ~ ZeroInflatedPoisson(exp(log_rate), logistic(zi_logit)) diff --git a/web-macro/todos/3.2-inferred-predictors-measurement-error-me.jl b/web-macro/todos/3.2-inferred-predictors-measurement-error-me.jl new file mode 100644 index 0000000..da974ea --- /dev/null +++ b/web-macro/todos/3.2-inferred-predictors-measurement-error-me.jl @@ -0,0 +1,20 @@ +# label: 3.2 inferred predictors / measurement error me() +# tier: 3 +# status: open +#= +**What it is.** brms's `me(x_obs, sd_x)` says "the predictor `x_obs` is itself measured with error of size `sd_x`; sample the latent true value during inference". The model sees both the observed value and the latent. + +**Why it matters.** Standard regression treats predictors as fixed/known. When predictors are themselves estimates (e.g. from a previous study or a noisy sensor), ignoring measurement error biases the slope estimates toward zero. `me()` is the principled fix. + +**Implementation.** Bigger architectural change: predictor columns become latent variables sampled during inference, not data columns evaluated once. Needs: +- A new column type analogous to `MissingColumn` but with an observation-driven prior `Normal(x_obs, sd_x)`. +- The latent column gets a slot in the population block (one parameter per row). +- `lprior!` adds the per-row Normal prior contribution. +- `vbroadcasted` resolves the column to the latent values, not the observed ones. + +**Verification.** Preset against synthetic data where the true `x` is known but only a noisy observed version is in the dataframe. Compare slope estimates with and without `me()`. + +=# + +loc ~ 1 + me(a, 0.1) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/3.3-ordinal-predictors-mo-monotonic-effects.jl b/web-macro/todos/3.3-ordinal-predictors-mo-monotonic-effects.jl new file mode 100644 index 0000000..f9891cd --- /dev/null +++ b/web-macro/todos/3.3-ordinal-predictors-mo-monotonic-effects.jl @@ -0,0 +1,18 @@ +# label: 3.3 ordinal predictors mo() (monotonic effects) +# tier: 3 +# status: open +#= +**What it is.** brms's `mo(x)` for an ordinal predictor with K levels: instead of `K-1` independent treatment-coded coefficients, fit a single "total effect" β plus a `K-1`-dim simplex of inter-level shape. Forces the effect to be monotonic in the ordering of `x`'s levels. + +**Why it matters.** Likert-scale predictors and ordered categorical inputs (e.g. age groups) have a natural ordering that treatment coding ignores. `mo()` enforces the monotonicity prior, dramatically reducing the parameter count and tightening posterior inference. + +**Implementation.** New block layout: one β coefficient + one Dirichlet-distributed simplex of length `K-1`. `_cat_lookup`-style materialization but the per-level contribution is `β * cumulative_sum(simplex)[level]` instead of `coefficients[level]`. + +Needs a new prior block type (Dirichlet) in `lprior!`, plus new parser support for `mo(x)` and a new `vmeta_sampling_rhs` overload. + +**Verification.** Preset against synthetic data where the true effect is monotonic but the levels are unordered in the data. Compare fitted shape parameters against the synthetic monotonic curve. + +=# + +loc ~ 1 + mo(c3) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/3.4-ordinal-outcomes-proportional-odds.jl b/web-macro/todos/3.4-ordinal-outcomes-proportional-odds.jl new file mode 100644 index 0000000..310d614 --- /dev/null +++ b/web-macro/todos/3.4-ordinal-outcomes-proportional-odds.jl @@ -0,0 +1,15 @@ +# label: 3.4 ordinal outcomes (proportional odds) +# tier: 3 +# status: open +#= +**What it is.** When `y` is itself ordered categorical (Likert response, severity grades, …), use a cumulative-link model: `Pr(y ≤ k) = logistic(α_k - η)` where `α_k` are K-1 cutpoints and η is the linear predictor. The likelihood is the difference of consecutive CDFs. + +**Why it matters.** Ordinal outcomes are common in survey data, clinical scoring, and any "rating" task. Treating them as continuous is statistically wrong; treating them as nominal categorical loses the ordering information. + +**Implementation.** New likelihood family with a vector of cutpoints as additional parameters. `Distributions.jl` has `OrderedLogistic` already — the pass-through path should mostly handle it once the parser knows to extract cutpoints from a `cumulative` wrapper. + +The cutpoints need an ordered prior (e.g. ordered transform of unconstrained reals), which means a new prior block type — similar to (3.3)'s simplex. + +**Verification.** Preset against synthetic Likert data. Compare cutpoints against `polr` from R's MASS package. + +=# diff --git a/web-macro/todos/3.5-mixture-models.jl b/web-macro/todos/3.5-mixture-models.jl new file mode 100644 index 0000000..f96b39f --- /dev/null +++ b/web-macro/todos/3.5-mixture-models.jl @@ -0,0 +1,18 @@ +# label: 3.5 mixture models +# tier: 3 +# status: open +#= +**What it is.** brms's `mixture(Normal, Normal)` lets the likelihood be a weighted mixture of K component distributions, with mixing weights estimated as parameters. + +**Why it matters.** Heterogeneous populations, latent class analysis, robust regression (Normal + heavy-tailed component), zero-inflated outcomes, … + +**Implementation.** Extends the existing `Distribution` pass-through. Need a new `MixtureModel` wrapper that holds component distributions plus a weights parameter block. `llikelihood!` uses `logsumexp(log_weights .+ logpdf.(components, y))` per row. + +Composes with (3.4 ordered priors) for the mixing weights' Dirichlet-like prior. + +**Verification.** Preset against synthetic two-component-Normal data. Confirm the recovered mixture weights and component parameters. + +=# + +loc ~ 1 + a +y1 ~ MixtureModel(Normal[Normal(loc, 0.5), Normal(0, 5)], [0.9, 0.1]) diff --git a/web-macro/todos/3.6-splines-gp-submodels-s-bs-gp-t2.jl b/web-macro/todos/3.6-splines-gp-submodels-s-bs-gp-t2.jl new file mode 100644 index 0000000..c85b8f0 --- /dev/null +++ b/web-macro/todos/3.6-splines-gp-submodels-s-bs-gp-t2.jl @@ -0,0 +1,18 @@ +# label: 3.6 splines / GP submodels s(), bs(), gp(), t2() +# tier: 3 +# status: open +#= +**What it is.** Smoothers in the linear predictor: `s(x)` for a generic spline, `bs(x, knots=...)` for a B-spline basis, `gp(x)` for a Gaussian process, `t2(x, y)` for a tensor-product spline. brms / mgcv use these heavily. + +**Why it matters.** Nonlinear effects without committing to a specific functional form. The de facto way to model dose-response curves, time effects, growth curves, spatial trends, … + +**Implementation.** Each smoother is a basis-matrix builder that grows a population block by `n_basis` columns and stores the basis matrix as part of `meta`. Function stubs (`s`, `bs`, `t2`, `gp`) already exist in `scripts/parsing.jl` so the parser side is partly done. + +For each smoother type, the materialization is `basis_matrix * coefficients` (length-N output). The smoothness prior is a structured prior on the coefficients (typically a Gaussian prior with a banded or 2D-difference penalty matrix), which requires (2.3 per-parameter priors) as a prerequisite. + +**Verification.** Preset against synthetic curve data. Compare fitted smoother against `mgcv::gam`. + +=# + +loc ~ 1 + s(a) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/3.7-autoregressive-submodels-ar-ar1.jl b/web-macro/todos/3.7-autoregressive-submodels-ar-ar1.jl new file mode 100644 index 0000000..8e32a37 --- /dev/null +++ b/web-macro/todos/3.7-autoregressive-submodels-ar-ar1.jl @@ -0,0 +1,18 @@ +# label: 3.7 autoregressive submodels ar(), ar1() +# tier: 3 +# status: open +#= +**What it is.** Add an AR(p) structure to the residuals: `y_t = η_t + φ * (y_{t-1} - η_{t-1}) + ε_t`. brms's `ar(time, p=1)` specifies the order and the time variable. + +**Why it matters.** Time-series and repeated-measures data routinely have autocorrelated errors. Ignoring AR structure inflates the effective sample size and gives overconfident posteriors. + +**Implementation.** Bigger than it looks because the likelihood is no longer per-row independent — it's a chain. Need a new `LikelihoodColumn`-like type that holds the time index and walks the data in time order, accumulating the AR contribution row by row. + +Composes with (2.5 grouped random effects) for per-subject AR structure. + +**Verification.** Preset against synthetic AR(1) data. Recover φ. + +=# + +loc ~ 1 + a + ar(g1, p=1) +y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/3.8-decompositions-qr-orthogonal-polar.jl b/web-macro/todos/3.8-decompositions-qr-orthogonal-polar.jl new file mode 100644 index 0000000..89bff2f --- /dev/null +++ b/web-macro/todos/3.8-decompositions-qr-orthogonal-polar.jl @@ -0,0 +1,15 @@ +# label: 3.8 decompositions (QR, orthogonal polar) +# tier: 3 +# status: open +#= +**What it is.** Numerical-stability transformations of the population design matrix. brms uses QR decomposition on the design matrix internally so the sampler sees an orthogonal-columns version, then transforms back at the end. Stan does the same. + +**Why it matters.** When population covariates are correlated (which is the norm), the unrotated design matrix gives a poorly-conditioned posterior that NUTS struggles with. QR fixes this with no statistical change. + +**Implementation.** Mostly orthogonal to the formula DSL — happens at `VBRMI` build time. Add a per-block transform: store both the original design matrix and the QR factor, run sampling on the rotated parameter space, transform back when extracting coefficients. + +Could be implemented today without any parser changes, as a `qr_transform=true` flag on `VBRMI`. Lift to a default-on once verified. + +**Verification.** Same model with and without QR should produce identical logdensity values, but the gradient should be better-conditioned (smaller condition number on the Hessian). + +=# diff --git a/web-macro/todos/3.9-spike-and-slab-horseshoe-priors.jl b/web-macro/todos/3.9-spike-and-slab-horseshoe-priors.jl new file mode 100644 index 0000000..5657033 --- /dev/null +++ b/web-macro/todos/3.9-spike-and-slab-horseshoe-priors.jl @@ -0,0 +1,17 @@ +# label: 3.9 spike-and-slab / Horseshoe priors +# tier: 3 +# status: open +#= +**What it is.** Sparsity-inducing priors for high-dimensional regression. Spike-and-slab puts a delta-spike at zero plus a wide slab; Horseshoe uses a half-Cauchy hyperprior on a per-coefficient scale, producing a heavy-tailed shrinkage prior. + +**Why it matters.** Without sparsity priors, high-dimensional regressions overfit. These are the standard solution in Bayesian variable selection and high-dim genomics / finance. + +**Implementation.** Depends entirely on (2.3 per-parameter priors). Once that's in place, spike-and-slab is `prior = Mixture(Normal(0, ε), Normal(0, slab_scale))` per coefficient, and Horseshoe is `Normal(0, λ_i * τ)` with `λ_i ~ HalfCauchy(0, 1)` and `τ ~ HalfCauchy(0, 1)` — both expressible in the existing Distribution pass-through once per-coefficient priors are wired up. + +**Verification.** Preset on a sparse synthetic regression (mostly-zero true coefficients with a few large ones). Confirm the Horseshoe-fitted coefficients shrink the noise toward zero and preserve the signal. + +=# + +coef_a ~ Horseshoe() +loc ~ coef_a * a +y1 ~ Normal(loc, 1) From e44321355d2e938d7acf3ee8a61ea11fa5f173c4 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Fri, 10 Apr 2026 10:34:25 +0200 Subject: [PATCH 09/23] vimpl: support `0` (no-intercept) in formula terms MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `slope ~ 0 + ztime` now skips the intercept — `vmeta_sampling_rhs(0; group)` returns `(meta, 0)` without calling `growblock!!`, so no parameter is allocated. The literal `0` broadcasts as a scalar in the parent `+`. Needed by the QT pipeline's two-formula pattern where the slope formula uses `0 + ztime` (no intercept, just the time covariate). Co-Authored-By: Claude Opus 4.6 (1M context) --- web-macro/src/vimpl.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/web-macro/src/vimpl.jl b/web-macro/src/vimpl.jl index 009e08a..2acc737 100644 --- a/web-macro/src/vimpl.jl +++ b/web-macro/src/vimpl.jl @@ -60,7 +60,8 @@ vmeta_sampling_rhs(meta, x::ExprColumn{typeof(|)}; kwargs...) = begin lhs, rhs = getargs(x, 2) vmeta_sampling_rhs(meta, lhs; group=rhs) end -vmeta_sampling_rhs(meta, ::Int; group) = begin +vmeta_sampling_rhs(meta, n::Int; group) = begin + n == 0 && return meta, 0 # `0` suppresses the intercept (no parameter allocated) meta, p = growblock!!(meta, group, 1) meta, _re_lookup(p, _gc_idx(group)) end From fc5a75fb6ceca2de4fd899821ac55235fcb57f10 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Fri, 10 Apr 2026 11:03:48 +0200 Subject: [PATCH 10/23] gitignore: exclude confidential bruno-*.jl todos from tracking web-macro/todos/bruno-*.jl files contain client-project model references that must stay local. The pattern is broad enough to cover any future Bruno-related TODO files. Co-Authored-By: Claude Opus 4.6 (1M context) --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 25770c7..0a35f31 100644 --- a/.gitignore +++ b/.gitignore @@ -27,3 +27,6 @@ Manifest*.toml # File generated by the Preferences package to store local preferences LocalPreferences.toml JuliaLocalPreferences.toml + +# Confidential client-project TODOs (local-only, never commit) +web-macro/todos/bruno-*.jl From b6b611c0aa7d63453d52e5323ccd6114556480d6 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Mon, 20 Apr 2026 13:49:16 +0200 Subject: [PATCH 11/23] vimpl: Part-based refactor, add mo1/mo monotonic terms MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - drop `Block` struct; `meta.blocks` is now NamedTuple[Symbol → Tuple of `Part{F,D}`], with each Part carrying its marker-fn dispatch tag and a NamedTuple of buffers. - split LKJ-Cholesky into `chol` + `grouped_normal` Parts sharing one L; `finalize` threads the expansion per-part. - add `simplex` primitive Part (log-scale stick-breaking, `alpha=1.0` kwarg on `lprior!`), and user-template block at the bottom defining `mo1(c)` and `mo(c)` = β·mo1(c), reusing the shared `_mo_contrast!` + `_scale_by_beta` helpers. - extract helpers: `push_parts!!`, `_scale_by_beta`, `_level_index` (with a `CategoricalArrays` fast path + Dict-fallback warning). - rename `page` → `__page__` in AppContext to match HTMXObjects naming. --- web-macro/Project.toml | 2 + web-macro/src/BRMMacroWeb.jl | 53 +++-- web-macro/src/vimpl.jl | 398 ++++++++++++++++++++++++++++------- 3 files changed, 353 insertions(+), 100 deletions(-) diff --git a/web-macro/Project.toml b/web-macro/Project.toml index 7f26610..5a0e9c4 100644 --- a/web-macro/Project.toml +++ b/web-macro/Project.toml @@ -3,6 +3,7 @@ uuid = "cd55decf-6be8-4ea5-907a-b2dc35e4cc14" version = "0.1.0" [deps] +CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0" @@ -20,6 +21,7 @@ LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" TestModules = "63c02187-99fd-4e5c-aaf0-4d6bfebc181c" Treebars = "e1e568c4-3a56-40a4-95fa-9b9c6c16fccb" Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" diff --git a/web-macro/src/BRMMacroWeb.jl b/web-macro/src/BRMMacroWeb.jl index 97f46f7..9adaec7 100644 --- a/web-macro/src/BRMMacroWeb.jl +++ b/web-macro/src/BRMMacroWeb.jl @@ -419,15 +419,15 @@ function vbrmi_card(vbrmi::VBRMI) end end for (key, value) in pairs(meta.materialized)] - # Blocks: show dimensions with styled block key + # Blocks: each block is a tuple of parts; print the block key + one line per part. blocks_rows = [begin - m, n = size(value) - role = key == :__population__ ? :derived : :data + role = key === :__population__ ? :derived : :data h.div( _styled_name(key, role), - h.span(": (n_levels=$m, n_params=$n)"; style="color:#999"), + [h.div(; style="color:#999;margin-left:1.5rem")(sprint(show, part)) + for part in parts]..., ) - end for (key, value) in pairs(meta.blocks)] + end for (key, parts) in pairs(meta.blocks)] h.article(; style="margin:0.5rem 0")( h.header(h.strong("VBRMI"), @@ -619,13 +619,13 @@ _index_body(formula::String) = h.div( ) @htmx struct AppContext - req = nothing + - # HTMXObjects auto-uses `page` to wrap any route's return value into a full - # page on direct browser navigation, while returning just the fragment for - # HTMX requests (see `_resolve_response` in HTMXObjects.jl). The sidebar's - # `hx-get` swaps target `#content` directly. - page(content) = htmx( + # HTMXObjects auto-uses `__page__` to wrap any route's return value into a + # full page on direct browser navigation, while returning just the fragment + # for HTMX requests (see `_resolve_response` in HTMXObjects.jl). The + # sidebar's `hx-get` swaps target `#content` directly. + __page__(content) = htmx( h.div(; style="display:flex;gap:1rem;align-items:flex-start")( nav_sidebar([ "Pipeline" => "/", @@ -670,16 +670,11 @@ _index_body(formula::String) = h.div( end @get todo = begin - todos = _load_todos() + todos = _load_todos(refresh=true) h.div( h.h1("TODO — what's missing for full BRM coverage"), - h.p("Items grouped by tier. Each item has a sketch of what it is, why it matters, how to implement, and how to verify. Sourced from .jl files under ", h.code("web-macro/todos/"), "; status edits and edited formulas are written back to disk."), - h.h2("Tier 1 — cheap wins"), - [_todo_card(t) for t in todos if t.tier == 1]..., - h.h2("Tier 2 — moderate (one design decision each)"), - [_todo_card(t) for t in todos if t.tier == 2]..., - h.h2("Tier 3 — bigger features (real new infrastructure)"), - [_todo_card(t) for t in todos if t.tier == 3]..., + h.p("Sorted by last modified. Each item has a sketch of what it is, why it matters, how to implement, and how to verify. Sourced from .jl files under ", h.code("web-macro/todos/"), "; status edits and edited formulas are written back to disk."), + [_todo_card(t) for t in todos]..., ) end end @@ -721,7 +716,7 @@ function _load_todos(; refresh::Bool=false) if refresh || !isassigned(_todos_cache) dir = _todos_dir() isdir(dir) || _migrate_todos!() - files = sort(filter(endswith(".jl"), readdir(dir; join=true))) + files = sort(filter(endswith(".jl"), readdir(dir; join=true)); by=mtime, rev=true) _todos_cache[] = TodoEntry[_parse_todo_file(f) for f in files] end _todos_cache[] @@ -828,6 +823,20 @@ end # ── Rendering: one Pico CSS article per TODO with status-colored border ──── +const _TIER_LABELS = ( + "T1", # tier 1 + "T2", # tier 2 + "T3", # tier 3 +) +const _TIER_COLORS = ("#4a7c59", "#5a6a8c", "#8c5a5a") + +_tier_pill(tier::Int) = h.span( + get(_TIER_LABELS, tier, "T$tier"); + style="font-size:0.7em;padding:0.1rem 0.4rem;border-radius:1rem;" * + "color:white;background:$(get(_TIER_COLORS, tier, "#888"));" * + "vertical-align:middle;font-weight:normal", +) + const _STATUS_COLORS = ( open = "#888", done = "#2e7d32", @@ -856,7 +865,9 @@ function _todo_card(todo::TodoEntry) )( h.details(; open=todo.status == :open)( h.summary(; style="cursor:pointer;list-style-position:outside")( - h.strong(todo.label), " ", _status_pills(todo.label, todo.status), + _tier_pill(todo.tier), " ", + h.strong(todo.label), " ", + _status_pills(todo.label, todo.status), ), h.div(; style="margin-top:0.5rem")(body_children...), ), diff --git a/web-macro/src/vimpl.jl b/web-macro/src/vimpl.jl index 2acc737..50c46f3 100644 --- a/web-macro/src/vimpl.jl +++ b/web-macro/src/vimpl.jl @@ -1,14 +1,53 @@ -using LogExpFunctions, InverseFunctions, Distributions, ElasticArrays, LogDensityProblems, LinearAlgebra +using LogExpFunctions, InverseFunctions, Distributions, ElasticArrays, LogDensityProblems, LinearAlgebra, SpecialFunctions +import CategoricalArrays as CA struct VBRMI{P<:BRMI,M<:NamedTuple} parent::P meta::M end VBRMI(p::BRMI) = VBRMI(p, finalize(foldl(vmeta, p.operations; init=(;materialized=(;), blocks=(;))))) -finalize(x) = merge(x, (;block_data=map(x.blocks) do values - m, n = size(values) - (;L=zeros(n, n)) -end)) + +""" + Part{F<:Function,D<:NamedTuple} + +Bundles a marker function `func` (the sole dispatch tag for per-part behavior — +`nparams`, `lprior!`, `Base.show`) with a NamedTuple `data` of per-kind state. +Each part owns its buffer and knows how to consume a slice of the unconstrained +vector, constrain it into the buffer, and return its log-prior contribution. + +`meta.blocks` is a NamedTuple keyed by coalescing Symbol (`:__population__`, or +a grouping-factor name like `:g1`); each value is a tuple of parts. +""" +struct Part{F<:Function,D<:NamedTuple} + func::F + data::D +end + +"""IID-normal population slots.""" +function normal end +"""Grouped column set, correlated via a shared L.""" +function grouped_normal end +"""LKJ-Cholesky owning the shared L for a grouped block.""" +function chol end +"""Unit simplex owning `values`; `length(values)` entries sum to 1, parameterized by `length(values)-1` unconstrained reals via stick-breaking.""" +function simplex end + +""" + finalize(meta) / finalize(parts) / finalize(part, acc) + +Threaded finalize: fold over each block's parts, each part emitting zero or +more replacement parts. `acc` is visible to later parts so a part can observe +what already landed (e.g. a `grouped_normal` could skip re-allocating L if a +`chol` sibling is already present — useful for the future multi-grouped_normal +case). Default per-part method is a pass-through; specializations emit their +own expansion (e.g. `grouped_normal` prepends a `chol` sharing the same L). +""" +finalize(x::NamedTuple) = merge(x, (; blocks = map(finalize, x.blocks))) +finalize(parts::Tuple) = foldl((acc, p) -> (acc..., finalize(p, acc)...), parts; init=()) +finalize(p::Part, _) = (p,) +finalize(p::Part{typeof(grouped_normal)}, _) = let n = size(p.data.values, 2), L = zeros(n, n) + (Part(chol, (; L)), Part(grouped_normal, merge(p.data, (; L)))) +end rmerge(x::NamedTuple, y::NamedTuple) = begin xykeys = (intersect(keys(x), keys(y))...,) merge(x, y, map(rmerge, NamedTuple{xykeys}(x), NamedTuple{xykeys}(y))) @@ -67,35 +106,15 @@ vmeta_sampling_rhs(meta, n::Int; group) = begin end vmeta_sampling_rhs(meta, x::NamedColumn; kwargs...) = vmeta_sampling_rhs(meta, meta.materialized[name(x)]; kwargs...) vmeta_sampling_rhs(meta, x::DataColumn; kwargs...) = vmeta_sampling_rhs(meta, parent(x); kwargs...) -vmeta_sampling_rhs(meta, x::AbstractVector{<:AbstractFloat}; group) = begin - meta, p = growblock!!(meta, group, 1) - meta, Base.broadcasted(*, x, _re_lookup(p, _gc_idx(group))) -end +vmeta_sampling_rhs(meta, x::AbstractVector{<:AbstractFloat}; group) = _scale_by_beta(meta, x; group) FBroadcasted{F,Style<:Union{Nothing, Base.Broadcast.BroadcastStyle},Axes} = Base.Broadcast.Broadcasted{Style,Axes,F} -vmeta_sampling_rhs(meta, x::FBroadcasted; group) = begin - meta, p = growblock!!(meta, group, 1) - meta, Base.broadcasted(*, x, _re_lookup(p, _gc_idx(group))) -end +vmeta_sampling_rhs(meta, x::FBroadcasted; group) = _scale_by_beta(meta, x; group) vmeta_sampling_rhs(meta, x::AbstractVector{<:Integer}; group) = begin # Treatment-coded categorical predictor: level 1 is the reference (drops out), # remaining k-1 levels each get their own coefficient slot in the block. - # TODO: figure out where to cache `levels`/`level_map`/`dense`/`gc_idx` so we - # don't rebuild them on every VBRMI construction. Candidates: a new - # `meta.factor` NamedTuple keyed by column name, or attach it to the - # materialized entry for the source column — _gc_idx() does the same - # dense-mapping work for the grouping-factor side. - # TODO: investigate hooking into an existing "categorical values vector" - # abstraction instead of building the dense map by hand. CategoricalArrays.jl - # (used by DataFrames) already exposes `levels`, `levelcode`, and a `.refs` - # field that *is* the dense Int8/16 code vector — if the input column is - # already a CategoricalVector we'd avoid the Dict round-trip entirely. Same - # idea applies to PooledArrays.jl. Need to decide whether vimpl.jl should - # take a hard dep on CategoricalArrays or sniff for the duck-typed interface. - levels = sort(unique(x)) - level_map = Dict(l => i for (i, l) in enumerate(levels)) - dense = [level_map[l] for l in x] - meta, p = growblock!!(meta, group, length(levels) - 1) - meta, _cat_broadcast(p, _gc_idx(group), dense) + K, c_idx = _level_index(x) + meta, p = growblock!!(meta, group, K - 1) + meta, _cat_broadcast(p, _gc_idx(group), c_idx) end # Population case (gc_idx === nothing): p is (1, k-1), look up the (level-1)-th # column via linear indexing. @@ -111,11 +130,44 @@ _cat_re_lookup(p, gc, level) = level == 1 ? zero(eltype(p)) : p[gc, level - 1] # `values`. Returns `nothing` for the special :__population__ marker so that the # population-level path can pass through `_re_lookup` unchanged. _gc_idx(::Symbol) = nothing -function _gc_idx(group::NamedColumn) - raw = parent(parent(group)) - levels = sort(unique(raw)) - level_map = Dict(l => i for (i, l) in enumerate(levels)) - [level_map[l] for l in raw] +_gc_idx(group::NamedColumn) = _level_index(parent(parent(group)))[2] + +""" + _level_index(raw) -> (K, c_idx) + +Return the number of distinct levels `K` and a dense `c_idx::Vector{Int}` +mapping each row to its 1-based level index. `CategoricalArrays.CategoricalVector` +is handled for free via `levels` + `levelcode`; plain vectors fall back to +building a `Dict` on the fly and warn once per call site. +""" +_level_index(raw::CA.CategoricalVector) = length(CA.levels(raw)), CA.levelcode.(raw) +_level_index(raw::AbstractVector) = begin + @warn "Building categorical level index from a plain vector; wrap in `categorical(…)` to avoid the per-call Dict." maxlog=1 _id=:vimpl_level_index + lvls = sort(unique(raw)) + lm = Dict(l => i for (i, l) in enumerate(lvls)) + length(lvls), [lm[l] for l in raw] +end + +""" + push_parts!!(meta, group::Symbol, parts::Part...) -> meta + +Append `parts` to `meta.blocks[group]`, creating the block if it didn't exist. +""" +push_parts!!(meta, group::Symbol, new::Part...) = + rmerge(meta, (; blocks = (; group => (get(meta.blocks, group, ())..., new...)))) + +""" + _scale_by_beta(meta, bcast; group) -> (meta, scaled_bcast) + +Allocate one normal slot via `growblock!!(meta, group, 1)` and return a +broadcasted expression that multiplies `bcast` by that β. Works at both the +population (scalar β) and group (one β per level, looked up via `_re_lookup`) +levels. Shared backbone of `AbstractVector{<:AbstractFloat}`, `FBroadcasted`, +and `mo(c)` parsers. +""" +_scale_by_beta(meta, bcast; group) = begin + meta, p = growblock!!(meta, group, 1) + meta, Base.broadcasted(*, bcast, _re_lookup(p, _gc_idx(group))) end # Wrap a (m, 1) growblock view as a length-N broadcasted lookup keyed by @@ -138,66 +190,112 @@ getbroadcast(x::MaterializedColumn) = getfield(x, :broadcast) Base.broadcastable(x::MaterializedColumn) = Base.broadcastable(parent(x)) n_levels(group::NamedColumn) = length(unique(parent(parent(group)))) -growblock!!(meta, group::Symbol, n) = growblock!!(meta, group, 1, n) -growblock!!(meta, group::Symbol, m, n) = begin - g = get(meta.blocks, group) do - ElasticMatrix(zeros(m, 0)) - end - idxs = (size(g, 2)+1):(size(g, 2)+n) - append!(g, zeros(m, n)) - rmerge(meta, (;blocks=(;group=>g))), view(g, :, idxs) + +_append!(values, m, n) = begin + append!(values, zeros(m, n)) + view(values, :, size(values, 2)-n+1:size(values, 2)) +end + +# Append a fresh part with an (m, n) buffer; always valid, no coalescing. +_push_part(parts::Tuple, f, m, n) = let values = ElasticMatrix(zeros(m, n)) + (parts..., Part(f, (; values))), view(values, :, 1:n) +end + +# Coalesce-if-same-kind: extend trailing part's ElasticMatrix by `n` cols when +# its func type matches, else append fresh. Dispatches on `(last_part, f)`. +_grow_or_push(parts::Tuple{}, f, m, n) = _push_part(parts, f, m, n) +_grow_or_push(parts::Tuple, f, m, n) = _grow_or_push_tail(parts, last(parts), f, m, n) +_grow_or_push_tail(parts::Tuple, tail::Part{F}, ::F, m, n) where F = + (parts, _append!(tail.data.values, m, n)) +_grow_or_push_tail(parts::Tuple, _, f, m, n) = _push_part(parts, f, m, n) + +growblock!!(meta, ::Symbol, n) = begin + parts = get(meta.blocks, :__population__, ()) + parts, p = _grow_or_push(parts, normal, 1, n) + rmerge(meta, (; blocks = (; __population__ = parts))), p +end +growblock!!(meta, group::NamedColumn, n) = begin + key, m = name(group), n_levels(group) + parts = get(meta.blocks, key, ()) + parts, p = _grow_or_push(parts, grouped_normal, m, n) + rmerge(meta, (; blocks = (; key => parts))), p end -growblock!!(meta, group::NamedColumn, n) = growblock!!(meta, name(group), n_levels(group), n) Base.show(io::IO, (;parent, broadcast)::MaterializedColumn) = print(io, eltype(parent), "[...] .= ", broadcast) Base.show(io::IO, (;parent, rhs)::LikelihoodColumn) = print(io, eltype(parent), "[...] .~ ", rhs) -Base.show(io::IO, vbrm::VBRMI) = begin - (;parent, meta) = vbrm + +Base.show(io::IO, p::Part{typeof(normal)}) = print(io, "Part{normal}", size(p.data.values)) +Base.show(io::IO, p::Part{typeof(grouped_normal)}) = print(io, "Part{grouped_normal}", size(p.data.values)) +Base.show(io::IO, p::Part{typeof(chol)}) = print(io, "Part{chol}(", size(p.data.L, 1), ")") +Base.show(io::IO, p::Part{typeof(simplex)}) = print(io, "Part{simplex}(", length(p.data.values), ")") +Base.show(io::IO, p::Part) = print(io, "Part{", p.func, "}") + +Base.show(io::IO, vbrm::VBRMI) = begin + (; parent, meta) = vbrm print(io, parent) print(io, "dim: ", LogDensityProblems.dimension(vbrm), "\n") print(io, "materialized:\n") for (key, value) in pairs(meta.materialized) print(io, " ", key, ": ", value, "\n") end - print(io, "blocks (n_levels, n_params):\n") - for (key, value) in pairs(meta.blocks) - print(io, " ", key, ": ", size(value), "\n") + print(io, "blocks:\n") + for (key, parts) in pairs(meta.blocks) + print(io, " ", key, ":\n") + for part in parts + print(io, " ", part, "\n") + end end end -LogDensityProblems.dimension(vbrm::VBRMI) = hyperdim(vbrm) + directdim(vbrm) -hyperdim(vbrm::VBRMI) = sum(pairs(vbrm.meta.blocks)) do (k, v) - n = size(v, 2) - k == :__population__ ? 0 : n * (n+1) ÷ 2 -end -directdim(vbrm::VBRMI) = sum(length, vbrm.meta.blocks) + +LogDensityProblems.dimension(vbrm::VBRMI) = nparams(vbrm.meta.blocks) + +nparams(x::Union{Tuple,NamedTuple}) = sum(nparams, x; init=0) +nparams(p::Part{typeof(normal)}) = size(p.data.values, 2) +nparams(p::Part{typeof(grouped_normal)}) = let (m, n) = size(p.data.values); m * n end +nparams(p::Part{typeof(chol)}) = let n = size(p.data.L, 1); n * (n + 1) ÷ 2 end +nparams(p::Part{typeof(simplex)}) = length(p.data.values) - 1 + advance!!(x, pos) = x[pos+1], pos+1 advance!!(x, pos, n) = view(x, pos+1:pos+n), pos+n -lprior!((;meta)::VBRMI, x::AbstractVector; init=(0., 0)) = foldl(pairs(meta.blocks); init) do (lprior, pos), (key, values) - m, n = size(values) - if key == :__population__ - xi, pos = advance!!(x, pos, n) - values[1, :] .= xi - lprior += sum(Base.Fix1(logpdf, Normal()), xi) - else - C = LinearAlgebra.Cholesky(meta.block_data[key].L, :L, 0) - lprior, pos = lprior!(C, x; init=(lprior, pos)) - for vi in eachrow(values) - xi, pos = advance!!(x, pos, n) - mul!(vi, C.L, xi) - lprior += sum(Base.Fix1(logpdf, Normal()), xi) - end - end - lprior, pos -end |> first -log_abs_tanh(x) = begin + +""" + lprior!(container, x) -> lp + +Recursive shell: each child (block-container, `Part`, …) is handed a view of +exactly `nparams(child)` unconstrained reals; bottoms out on `Part` methods, +each of which writes its constrained buffer and returns its log-prior + +Jacobian contribution. +""" +lprior!(vbrmi::VBRMI, x::AbstractVector) = lprior!(vbrmi.meta.blocks, x) +lprior!(xs::Union{Tuple,NamedTuple}, x::AbstractVector) = + foldl(xs; init=(0.0, 0)) do (total, pos), child + xi, pos = advance!!(x, pos, nparams(child)) + total + lprior!(child, xi), pos + end |> first + +lprior!(p::Part{typeof(normal)}, x) = begin + p.data.values[1, :] .= x + sum(Base.Fix1(logpdf, Normal()), x) +end + +log_abs_tanh(x) = begin z = -2*abs(x) (log1mexp(z) - log1pexp(z)) end log_square_tanh(x) = 2 * log_abs_tanh(x) -"Either wrong or better LKJCholesky unconstraining + prior" -lprior!((;L)::Cholesky, x; init, eta=1.) = begin - lprior, pos = init + +""" + lprior!(p::Part{typeof(chol)}, x; eta=1.0) -> lp + +Either wrong or better LKJCholesky unconstraining + prior. Writes `p.data.L` +in place from `x` (length n(n+1)/2) and returns the log-prior + Jacobian +contribution. `eta` is the LKJ shape parameter. +""" +lprior!(p::Part{typeof(chol)}, x; eta=1.0) = begin + L = p.data.L n = LinearAlgebra.checksquare(L) + pos = 0 + lprior = 0.0 log_scale, pos = advance!!(x, pos) lprior += logpdf(Normal(), log_scale) L[1, 1] = exp(log_scale) @@ -222,8 +320,52 @@ lprior!((;L)::Cholesky, x; init, eta=1.) = begin L[i, i] = exp(log_scale + .5 * log1mexp(log_sos)) lprior += (n - i + 2*eta-2) * .5 * log1mexp(log_sos) end - lprior, pos + lprior end + +lprior!(p::Part{typeof(grouped_normal)}, x) = begin + (; values, L) = p.data + _, n = size(values) + lprior = 0.0 + pos = 0 + for vi in eachrow(values) + xi, pos = advance!!(x, pos, n) + mul!(vi, LowerTriangular(L), xi) + lprior += sum(Base.Fix1(logpdf, Normal()), xi) + end + lprior +end + +""" + lprior!(p::Part{typeof(simplex)}, x; alpha=1.0) -> lp + +Log-scale stick-breaking with logistic slices: constrain `length(values)-1` +unconstrained reals `x` into the simplex `p.data.values` in place and return +`log|J| + logpdf(Dirichlet(K, α), values)` in one pass. Fully log-scale — +never forms `1 - Σ values_ (meta, broadcasted) + +Shared core of `mo1` and `mo`: for a categorical column `raw` with K levels, +emits `(Part(simplex, …), Part(mo1, …))` into the population block and returns +a broadcasted `getindex(contrast, c_idx)` expression. +""" +_mo_contrast!(meta, raw) = begin + K, c_idx = _level_index(raw) + K >= 2 || error("mo/mo1: categorical must have ≥ 2 levels (got $K)") + values = zeros(K - 1) + contrast = zeros(K) + meta = push_parts!!(meta, :__population__, + Part(simplex, (; values)), + Part(mo1, (; values, contrast, c_idx)), + ) + meta, Base.broadcasted(getindex, Ref(contrast), c_idx) +end + +""" + vmeta_sampling_rhs(meta, x::ExprColumn{typeof(mo1)}; group) -> (meta, broadcasted) + +Parse a `mo1(c)` term: brms-style monotonic effect with β fixed to 1. +Population-level only for now. +""" +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(mo1)}; group) = begin + group === :__population__ || error("mo1: only population-level supported for now (got group=$group)") + _mo_contrast!(meta, vbroadcasted(getargs(x, 1)[1]; meta)) +end + +""" + vmeta_sampling_rhs(meta, x::ExprColumn{typeof(mo)}; group) -> (meta, broadcasted) + +Parse a `mo(c)` term: brms-style monotonic effect with free β. Delegates to +`_mo_contrast!` for the monotonic contrast and multiplies by a fresh β slot +allocated via `growblock!!`. Population-level only for now. +""" +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(mo)}; group) = begin + group === :__population__ || error("mo: only population-level supported for now (got group=$group)") + meta, contrast_bcast = _mo_contrast!(meta, vbroadcasted(getargs(x, 1)[1]; meta)) + _scale_by_beta(meta, contrast_bcast; group) +end + +nparams(p::Part{typeof(mo1)}) = 0 + +""" + lprior!(p::Part{typeof(mo1)}, _) -> 0.0 + +Zero-parameter follow-up to a `simplex` sibling: refresh +`contrast = vcat(0, cumsum(values))` in place so downstream broadcasts see +up-to-date contrast values. Contributes nothing to the log-prior. +""" +lprior!(p::Part{typeof(mo1)}, _) = begin + (; values, contrast) = p.data + contrast[1] = 0.0 + for k in 2:length(contrast) + contrast[k] = contrast[k-1] + values[k-1] + end + 0.0 +end + +Base.show(io::IO, p::Part{typeof(mo1)}) = print(io, "Part{mo1}(K=", length(p.data.contrast), ")") + +# ============================================================================== +# End of `mo1` / `mo` template +# ============================================================================== \ No newline at end of file From 94e876ab7b0480bebc0437542a3f406347610f6f Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Mon, 20 Apr 2026 22:40:47 +0200 Subject: [PATCH 12/23] web-macro: AppContext + polling_fetchindex + @inline hot path - Wire AppData via struct-body `__appdata__ = APPDATA` singleton. - Rewrite `@get stage` with Treebars `polling_fetchindex` so pipeline runs off the request thread and the UI polls for progress. - Add dataset_namespace/dataset_container/dataset_extras hooks so per-TODO extras (e.g. bruno-ext) can be spliced into the data NamedTuple without touching macro.jl. - Split `render_output` -> `render_pipeline`, stream per-step Chairmarks `@be` benches (logdensity, full lprior!, each Part, each materialized column) so allocations can be localized. - Add formula safety whitelist before `Meta.parse`/`eval`. - TODO page with file-backed entries and mark/formula persistence. - `@inline` lprior!/llikelihood! leaves + container folds so the heterogeneous-tuple `foldl` specializes to zero allocations. - .gitignore: also exclude `web-macro/src/bruno-*.jl`. Co-Authored-By: Claude Opus 4.7 --- .gitignore | 1 + web-macro/src/BRMMacroWeb.jl | 385 ++++++++++++++++++++--------------- web-macro/src/vimpl.jl | 30 +-- 3 files changed, 240 insertions(+), 176 deletions(-) diff --git a/.gitignore b/.gitignore index 0a35f31..b076ed1 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,4 @@ JuliaLocalPreferences.toml # Confidential client-project TODOs (local-only, never commit) web-macro/todos/bruno-*.jl +web-macro/src/bruno-*.jl diff --git a/web-macro/src/BRMMacroWeb.jl b/web-macro/src/BRMMacroWeb.jl index 9adaec7..57db88c 100644 --- a/web-macro/src/BRMMacroWeb.jl +++ b/web-macro/src/BRMMacroWeb.jl @@ -1,6 +1,7 @@ module BRMMacroWeb using HTMXObjects +using Treebars: polling_fetchindex, initialize_progress! using Random using Chairmarks using DataFrames @@ -74,6 +75,29 @@ function synthetic_df(; n=64, seed=1) y1, y2, k1, k2, bin_n, bin_succ, bin_y) end +# Extension hook: extensions (e.g. the gitignored `bruno-ext.jl`) that need +# to contribute auxiliary data which doesn't fit as per-row DataFrame +# columns -- for example `dose_times::Vector{<:AbstractVector}` indexed by +# subject id -- add a method `dataset_extras(::Val{:ns}, df)` returning a +# NamedTuple of extras. The namespace is derived from the TODO label/slug +# (first dash/space-separated segment), so `bruno-qt-*` TODOs dispatch to +# `::Val{:bruno}`. Default is no extras. +dataset_extras(::Val, df) = (;) + +# Namespace extracted from a TODO label or slug; empty -> :default. +dataset_namespace(label::AbstractString) = isempty(strip(label)) ? :default : + Symbol(lowercase(first(split(strip(label), r"[\s:\-]+")))) + +# Build the container passed as `df` to `_brm`: a NamedTuple merging the +# synthetic-df columns with extras dispatched by namespace. `@brm` only +# needs `hasproperty`/`getproperty` on its data argument, so a NamedTuple +# works as a drop-in for the DataFrame and lets extensions splice in +# extras without touching macro.jl. +function dataset_container(df, ns::Symbol=:default) + cols = (; (Symbol(c) => df[!, c] for c in names(df))...) + merge(cols, dataset_extras(Val(ns), df)) +end + # ── Formula safety whitelist ──────────────────────────────────────────────── # # The formula textarea accepts arbitrary text that gets `Meta.parse`d then @@ -83,7 +107,7 @@ end # generous for formula-writing (math, distributions, data-column references) # but blocks I/O, shell, eval, include, ccall, macros, etc. -const _ALLOWED_CALLS = Set{Symbol}([ +_ALLOWED_CALLS = Set{Symbol}([ # DSL operators :~, :(+), :(-), :(*), :(/), :(^), :(|), :(||), # Comparison (may appear in ifelse-style expressions) @@ -104,7 +128,7 @@ const _ALLOWED_CALLS = Set{Symbol}([ :OrderedLogistic, :Categorical, # TODO stubs (not yet implemented but syntactically valid) :scale, :center, :standardize, :factor, :offset, - :s, :bs, :t2, :gp, :ar, :ar1, :mo, + :s, :bs, :t2, :gp, :ar, :ar1, :mo, :mo1, :cbind, :mvbind, :mm, :gr, :dp, :me, :centered, :Horseshoe, :ZeroInflatedPoisson, :weighted, # Data helpers @@ -112,7 +136,7 @@ const _ALLOWED_CALLS = Set{Symbol}([ ]) # Expression heads that are safe in a formula AST (literals, blocks, calls, …) -const _SAFE_HEADS = Set{Symbol}([ +_SAFE_HEADS = Set{Symbol}([ :block, :call, :., :(=), :(||), :tuple, :vect, :ref, :kw, :parameters, :(...), # Comparison chains @@ -173,7 +197,7 @@ end # intermediate step (parsing, transforming, wrapping, eval'ing, materializing, # benchmarking) and inspect the result without paying for the later stages. -const STAGES = ( +STAGES = ( :parse, # Meta.parse(formula) :transform, # parse!(...) — rewrites = and ~ into @n/@x macro calls :wrap, # _brm(formula; df) — full let-block ready to eval @@ -186,10 +210,11 @@ stage_index(s::Symbol) = something(findfirst(==(s), STAGES), length(STAGES)) # Run the pipeline up to (and including) `stage`. Returns a NamedTuple # carrying every intermediate value computed so far. -function pipeline(formula::AbstractString, stage::Symbol) +function pipeline(formula::AbstractString, stage::Symbol; + namespace::Symbol=:default) s = stage_index(stage) df = synthetic_df() - out = (; df) + out = (; df, namespace) s >= 1 || return out raw = Meta.parse("begin\n$formula\nend") @@ -205,7 +230,8 @@ function pipeline(formula::AbstractString, stage::Symbol) out = merge(out, (; transformed, alllocals)) s >= 3 || return out - wrapped = _brm(formula; df) + container = dataset_container(df, namespace) + wrapped = _brm(formula; df=container) out = merge(out, (; wrapped)) s >= 4 || return out @@ -216,29 +242,44 @@ function pipeline(formula::AbstractString, stage::Symbol) vbrmi = VBRMI(brmi) dim = LogDensityProblems.dimension(vbrmi) x0 = randn(Xoshiro(0), dim) - ldp = try - string(LogDensityProblems.logdensity(vbrmi, x0)) - catch e - "error: " * sprint(showerror, e) - end - grad = try - FiniteDifferences.grad( - central_fdm(5, 1), - Base.Fix1(LogDensityProblems.logdensity, vbrmi), - x0, - )[1] - catch e - e - end + ldp = string(LogDensityProblems.logdensity(vbrmi, x0)) + grad = FiniteDifferences.grad( + central_fdm(5, 1), + Base.Fix1(LogDensityProblems.logdensity, vbrmi), + x0, + )[1] out = merge(out, (; vbrmi, dim, ldp, x0, grad)) s >= 6 || return out - bench = try - @be randn(dim) LogDensityProblems.logdensity($vbrmi, _) - catch e - e + x_rand = randn(dim) + benches = Pair{String,Any}[] + push!(benches, "logdensity (total)" => + @be randn(dim) LogDensityProblems.logdensity($vbrmi, _)) + push!(benches, "lprior!" => + @be randn(dim) lprior!($vbrmi, _)) + # Per-Part lprior! split: the foldl in lprior!(blocks, x) hands each Part + # a view of exactly nparams(part) reals. Reconstruct those slices here so + # each Part's contribution can be benched in isolation. + let pos = 0 + for (group_key, parts) in pairs(vbrmi.meta.blocks) + for (i, part) in enumerate(parts) + n = nparams(part) + xi = view(x_rand, pos+1:pos+n) + push!(benches, " lprior!($group_key[$i] $(part))" => + @be lprior!($part, $xi)) + pos += n + end + end end - merge(out, (; bench)) + # llikelihood! splits: each materialized column (either a linear-predictor + # MaterializedColumn or a LikelihoodColumn). Bench each in isolation so the + # allocation / time cost of each step is visible. + _ = lprior!(vbrmi, x_rand) # pre-fill buffers so bench measures the per-step cost + for (key, m) in pairs(vbrmi.meta.materialized) + push!(benches, "llikelihood!($key)" => + @be llikelihood!($m)) + end + merge(out, (; benches)) end # ── Rendering helpers ─────────────────────────────────────────────────────── @@ -442,103 +483,82 @@ function vbrmi_card(vbrmi::VBRMI) end -function render_output(formula::AbstractString; stage::Symbol=:vbrmi) +function render_pipeline(out::NamedTuple) sections = Vector{Any}[] # one entry per stage; rendered most-recent-first - try - out = pipeline(formula, stage) - - # Synthetic data always pinned at the top, collapsed by default so the - # macro pipeline output stays the focus. - data_section = Any[ - h.details( - h.summary("Synthetic data ($(nrow(out.df)) rows × $(ncol(out.df)) cols: " * - join(string.(names(out.df)), ", ") * ") — click to expand"), - render_table(out.df; sortable=false), - ), - ] - if haskey(out, :raw) - push!(sections, Any[_section("1. Meta.parse — raw Julia AST", - sprint(show, out.raw))...]) - end - if haskey(out, :transformed) - push!(sections, Any[ - _section("2. parse! — rewritten AST (= → @n/@x assign, ~ → @n/@x ~)", - sprint(show, out.transformed))..., - h.h3(" locals classified by parse!"), - h.pre(sprint(show, out.alllocals)), - ]) - end - if haskey(out, :wrapped) - push!(sections, Any[_section("3. _brm — full let-block (df spliced as a literal)", - sprint(show, out.wrapped))...]) - end - if haskey(out, :brmi) - push!(sections, Any[ - h.h3("4. eval — BRMI value (parsed model)"), - brmi_card(out.brmi), - ]) - end - if haskey(out, :vbrmi) - # Build the FD check summary for the
toggle line - fd_summary = if out.grad isa Exception - h.span("logdensity + FD check: error"; style="color:crimson") - else - tol = 1e-8 - n_dead = count(<=(tol) ∘ abs, out.grad) - if n_dead == 0 - h.span("logdensity + FD check: $(out.dim)/$(out.dim) active ✓"; - style="color:green") - else - h.span("logdensity + FD check: $(n_dead) dead param(s)"; - style="color:crimson") - end - end - - fd_body = if out.grad isa Exception - h.pre("gradient error: " * sprint(showerror, out.grad)) - else - tol = 1e-8 - dead = findall(<=(tol) ∘ abs, out.grad) - h.div( - h.p("dim = ", string(out.dim), ", logdensity = ", out.ldp), - isempty(dead) ? "" : - h.p(; style="color:crimson")( - "dead param indices: ", string(dead)), - h.pre(sprint(show, MIME"text/plain"(), out.grad)), - ) - end - - push!(sections, Any[ - h.h3("5. VBRMI — materialized action (blocks, dim, columns)"), - vbrmi_card(out.vbrmi), - h.details(h.summary(fd_summary), fd_body), - ]) - end - if haskey(out, :bench) - bench_body = out.bench isa Exception ? - h.pre("benchmark error: " * sprint(showerror, out.bench)) : - h.pre(sprint(show, MIME"text/plain"(), out.bench)) - push!(sections, Any[h.h3("6. Chairmarks @be — primal logdensity"), bench_body]) - end + # Synthetic data always pinned at the top, collapsed by default so the + # macro pipeline output stays the focus. + data_section = Any[ + h.details( + h.summary("Synthetic data ($(nrow(out.df)) rows × $(ncol(out.df)) cols: " * + join(string.(names(out.df)), ", ") * ") — click to expand"), + render_table(out.df; sortable=false), + ), + ] - # Stages render most-recent-first; synthetic data sits at the very top. - children = reduce(vcat, reverse(sections); init=Any[]) - prepend!(children, data_section) - return h.div(; id="brm-macro-output")(children...) - catch e - return h.div(; id="brm-macro-output")( - h.h3("Error"), - h.pre(sprint(showerror, e, catch_backtrace())), + if haskey(out, :raw) + push!(sections, Any[_section("1. Meta.parse — raw Julia AST", + sprint(show, out.raw))...]) + end + if haskey(out, :transformed) + push!(sections, Any[ + _section("2. parse! — rewritten AST (= → @n/@x assign, ~ → @n/@x ~)", + sprint(show, out.transformed))..., + h.h3(" locals classified by parse!"), + h.pre(sprint(show, out.alllocals)), + ]) + end + if haskey(out, :wrapped) + push!(sections, Any[_section("3. _brm — full let-block (df spliced as a literal)", + sprint(show, out.wrapped))...]) + end + if haskey(out, :brmi) + push!(sections, Any[ + h.h3("4. eval — BRMI value (parsed model)"), + brmi_card(out.brmi), + ]) + end + if haskey(out, :vbrmi) + tol = 1e-8 + n_dead = count(<=(tol) ∘ abs, out.grad) + fd_summary = n_dead == 0 ? + h.span("logdensity + FD check: $(out.dim)/$(out.dim) active ✓"; + style="color:green") : + h.span("logdensity + FD check: $(n_dead) dead param(s)"; + style="color:crimson") + dead = findall(<=(tol) ∘ abs, out.grad) + fd_body = h.div( + h.p("dim = ", string(out.dim), ", logdensity = ", out.ldp), + isempty(dead) ? "" : + h.p(; style="color:crimson")( + "dead param indices: ", string(dead)), + h.pre(sprint(show, MIME"text/plain"(), out.grad)), ) + push!(sections, Any[ + h.h3("5. VBRMI — materialized action (blocks, dim, columns)"), + vbrmi_card(out.vbrmi), + h.details(h.summary(fd_summary), fd_body), + ]) end + if haskey(out, :benches) + bench_rows = [h.div( + h.strong(label), h.br(), + h.pre(sprint(show, MIME"text/plain"(), b)) + ) for (label, b) in out.benches] + push!(sections, Any[h.h3("6. Chairmarks @be — per-step"), bench_rows...]) + end + + # Stages render most-recent-first; synthetic data sits at the very top. + children = reduce(vcat, reverse(sections); init=Any[]) + prepend!(children, data_section) + h.div(; id="brm-macro-output")(children...) end # ── Routes ────────────────────────────────────────────────────────────────── -_stage_button(label, stage) = h.button(label; type="button", +_stage_button(self, label, stage) = h.button(label; type="button", id="stage-$stage", - hx_get="/stage/$stage", + hx_get=string(query_url(self/"stage/$stage"; force=true)), hx_include="#brm-macro-form", hx_target="#brm-macro-output", hx_swap="outerHTML") @@ -582,12 +602,12 @@ _preset_button(label, formula) = h.button(label; onclick="document.querySelector('textarea[name=formula]').value = this.dataset.formula; document.getElementById('stage-vbrmi').click()", style="font-size:0.8em;padding:0.2rem 0.5rem;margin:0") -_index_body(formula::String) = h.div( +_index_body(self, formula::String) = h.div( h.h1("BRM macro pipeline"), h.p( "Enter a ", h.code("@brm"), " formula and step through the macro pipeline: ", - h.code("Meta.parse"), " → ", h.code("parse!"), " → ", h.code("_brm"), - " let-block → ", h.code("eval"), " → ", h.code("VBRMI"), " action → ", + h.code("Meta.parse"), " -> ", h.code("parse!"), " -> ", h.code("_brm"), + " let-block -> ", h.code("eval"), " -> ", h.code("VBRMI"), " action -> ", h.code("Chairmarks"), " benchmark.", ), h.details( @@ -607,19 +627,19 @@ _index_body(formula::String) = h.div( style="width:100%;font-family:monospace"), ), h.fieldset(; class="grid")( - _stage_button("1. Parse", :parse), - _stage_button("2. Transform", :transform), - _stage_button("3. Wrap", :wrap), - _stage_button("4. BRMI", :brmi), - _stage_button("5. VBRMI", :vbrmi), - _stage_button("6. Benchmark", :bench), + _stage_button(self, "1. Parse", :parse), + _stage_button(self, "2. Transform", :transform), + _stage_button(self, "3. Wrap", :wrap), + _stage_button(self, "4. BRMI", :brmi), + _stage_button(self, "5. VBRMI", :vbrmi), + _stage_button(self, "6. Benchmark", :bench), ), ), - render_output(formula), + lazy(string(query_url(self/"stage/bench"; formula)); id="brm-macro-output"), ) @htmx struct AppContext - + __appdata__ = APPDATA # HTMXObjects auto-uses `__page__` to wrap any route's return value into a # full page on direct browser navigation, while returning just the fragment @@ -646,35 +666,54 @@ _index_body(formula::String) = h.div( # If a TODO form posted us a (label, formula) pair, persist the edited # formula to that TODO's .jl file so the next visit to the TODO page # shows the user's edits instead of the seed default. - isempty(label) || _save_todo!(label; new_formula=formula) - _index_body(formula) + isempty(label) || _save_todo!(__appdata__, label; new_formula=formula) + _index_body(__self__, formula) end @get mark(; label::String="", state::String="") = begin isempty(label) && return "" target = Symbol(state) - entry = _find_todo(label) + entry = _find_todo(__appdata__, label) entry === nothing && return "" next_status = entry.status == target ? :open : target - updated = _save_todo!(label; new_status=next_status) + updated = _save_todo!(__appdata__, label; new_status=next_status) # Re-render the whole card so the border + collapse state update # together with the pill text. - _todo_card(updated) + _todo_card(__self__, updated) end - @get stage(name::AbstractString; formula::String=default_formula(), label::String="") = begin + @get stage(name::AbstractString; formula::String=default_formula(), + label::String="", force::Bool=false) = begin # When called from a TODO card's form, persist the (possibly edited) # formula back to the todo's .jl file before rendering. - isempty(label) || _save_todo!(label; new_formula=formula) - render_output(formula; stage=Symbol(name)) + isempty(label) || _save_todo!(__appdata__, label; new_formula=formula) + ns = dataset_namespace(label) + polling_fetchindex(__appdata__.pipeline_result, + formula, Symbol(name), ns; + poll_url=string(query_url(__self__/"stage/$name"; formula, label)), + label="BRM pipeline - $name", + force) do out + render_pipeline(out) + end end - @get todo = begin - todos = _load_todos(refresh=true) + @get todo(; slug::String="") = begin + if !isempty(slug) + entry = _find_todo_by_slug(__appdata__, slug) + entry === nothing && return h.div( + h.p("No TODO with slug ", h.code(slug), "."), + h.a("<- Back to TODO list"; href="/todo"), + ) + return h.div( + h.p(h.a("<- Back to TODO list"; href="/todo")), + _todo_card(__self__, entry), + ) + end + todos = _load_todos(__appdata__) h.div( - h.h1("TODO — what's missing for full BRM coverage"), + h.h1("TODO - what's missing for full BRM coverage"), h.p("Sorted by last modified. Each item has a sketch of what it is, why it matters, how to implement, and how to verify. Sourced from .jl files under ", h.code("web-macro/todos/"), "; status edits and edited formulas are written back to disk."), - [_todo_card(t) for t in todos]..., + [_todo_card(__self__, t) for t in todos]..., ) end end @@ -710,25 +749,35 @@ end _todos_dir() = joinpath(dirname(@__DIR__), "todos") _slug(label::AbstractString) = lowercase(strip(replace(label, r"[^\w.]+" => "-"), '-')) -const _todos_cache = Ref{Vector{TodoEntry}}() +@dynamicstruct struct AppData + __status__ = initialize_progress!(:state; description="BRM pipeline") + @cached pipeline_result(formula, stage, namespace) = + pipeline(formula, stage; namespace) +end -function _load_todos(; refresh::Bool=false) - if refresh || !isassigned(_todos_cache) - dir = _todos_dir() - isdir(dir) || _migrate_todos!() - files = sort(filter(endswith(".jl"), readdir(dir; join=true)); by=mtime, rev=true) - _todos_cache[] = TodoEntry[_parse_todo_file(f) for f in files] - end - _todos_cache[] +function _load_todos(::AppData) + dir = _todos_dir() + isdir(dir) || _migrate_todos!() + files = sort(filter(endswith(".jl"), readdir(dir; join=true)); by=mtime, rev=true) + TodoEntry[_parse_todo_file(f) for f in files] end -function _find_todo(label::AbstractString) - for t in _load_todos() +function _find_todo(appdata::AppData, label::AbstractString) + for t in _load_todos(appdata) t.label == label && return t end nothing end +_todo_slug(t::TodoEntry) = replace(basename(t.path), r"\.jl$" => "") + +function _find_todo_by_slug(appdata::AppData, slug::AbstractString) + for t in _load_todos(appdata) + _todo_slug(t) == slug && return t + end + nothing +end + function _parse_todo_file(path::String) lines = readlines(path) header = Dict{String,String}() @@ -785,10 +834,10 @@ function _write_todo_file(todo::TodoEntry) write(todo.path, take!(io)) end -function _save_todo!(label::AbstractString; +function _save_todo!(appdata::AppData, label::AbstractString; new_status::Union{Symbol,Nothing}=nothing, new_formula::Union{String,Nothing}=nothing) - todo = _find_todo(label) + todo = _find_todo(appdata, label) todo === nothing && return nothing updated = TodoEntry( todo.path, todo.label, todo.tier, @@ -797,7 +846,6 @@ function _save_todo!(label::AbstractString; new_formula === nothing ? todo.formula : new_formula, ) _write_todo_file(updated) - _load_todos(refresh=true) updated end @@ -823,12 +871,12 @@ end # ── Rendering: one Pico CSS article per TODO with status-colored border ──── -const _TIER_LABELS = ( +_TIER_LABELS = ( "T1", # tier 1 "T2", # tier 2 "T3", # tier 3 ) -const _TIER_COLORS = ("#4a7c59", "#5a6a8c", "#8c5a5a") +_TIER_COLORS = ("#4a7c59", "#5a6a8c", "#8c5a5a") _tier_pill(tier::Int) = h.span( get(_TIER_LABELS, tier, "T$tier"); @@ -837,20 +885,20 @@ _tier_pill(tier::Int) = h.span( "vertical-align:middle;font-weight:normal", ) -const _STATUS_COLORS = ( +_STATUS_COLORS = ( open = "#888", done = "#2e7d32", deprioritized = "#a05a2c", ) _status_color(s::Symbol) = get(_STATUS_COLORS, s, "#888") -function _todo_card(todo::TodoEntry) +function _todo_card(self, todo::TodoEntry) border_color = _status_color(todo.status) body_children = Any[HTMXObjects.md_to_node(todo.body)] if todo.formula !== nothing - push!(body_children, _formula_form(todo.label, todo.formula)) + push!(body_children, _formula_form(self, todo.label, todo.formula)) # Inline pipeline-result target — the form's hx_get fills this div - # with `render_output(formula; stage=:vbrmi)` so the user sees the + # with `render_pipeline(out)` so the user sees the # VBRMI/finite-difference output right inside the card. push!(body_children, h.div(; id="todo-result-$(hash(todo.label))", style="margin-top:0.5rem")) @@ -867,16 +915,24 @@ function _todo_card(todo::TodoEntry) h.summary(; style="cursor:pointer;list-style-position:outside")( _tier_pill(todo.tier), " ", h.strong(todo.label), " ", - _status_pills(todo.label, todo.status), + _status_pills(todo.label, todo.status), " ", + _permalink(todo), ), h.div(; style="margin-top:0.5rem")(body_children...), ), ) end -function _formula_form(label::String, formula::String) +_permalink(todo::TodoEntry) = h.a("🔗"; + href="/todo?slug=$(HTTP.URIs.escapeuri(_todo_slug(todo)))", + title="Standalone URL", + onclick="event.stopPropagation()", + style="text-decoration:none;font-size:0.8em;margin-left:0.3rem;vertical-align:middle", +) + +function _formula_form(self, label::String, formula::String) h.form(; - hx_get="/stage/vbrmi", + hx_get=string(query_url(self/"stage/bench"; force=true)), hx_target="#todo-result-$(hash(label))", hx_swap="innerHTML", style="margin:0.5rem 0", @@ -902,10 +958,6 @@ function _status_pills(label::AbstractString, state::Symbol) ) end -# Convenience overload that fetches the current status from disk. -_status_pills(label::AbstractString) = _status_pills(label, - something(_find_todo(label), (status=:open,)).status) - # Pill text and color reflect the *current* state. When the pill's target_state # is currently active, it shows the active label (e.g. "✓ done") in the active # color and clicking it toggles back to :open. Otherwise it shows the inactive @@ -1274,8 +1326,15 @@ Defer until everything else is solid. """, ] +const APPDATA = AppData(; cache_type=:parallel) + function __init__() route!(AppContext()) end +# Bruno-specific extensions (gitignored); load if present. +let path = joinpath(@__DIR__, "bruno-ext.jl") + isfile(path) && include(path) +end + end # module diff --git a/web-macro/src/vimpl.jl b/web-macro/src/vimpl.jl index 50c46f3..ab6874e 100644 --- a/web-macro/src/vimpl.jl +++ b/web-macro/src/vimpl.jl @@ -188,6 +188,10 @@ end Base.parent(x::MaterializedColumn) = getfield(x, :parent) getbroadcast(x::MaterializedColumn) = getfield(x, :broadcast) Base.broadcastable(x::MaterializedColumn) = Base.broadcastable(parent(x)) +# A name bound via `~` (e.g. `loc1 ~ 1 + a`) lands in meta.materialized as a +# MaterializedColumn whose parent is the live buffer refreshed each draw. +# When it reappears on a later `~` RHS, scale the buffer by a fresh beta. +vmeta_sampling_rhs(meta, x::MaterializedColumn; group) = _scale_by_beta(meta, parent(x); group) n_levels(group::NamedColumn) = length(unique(parent(parent(group)))) @@ -255,8 +259,8 @@ nparams(p::Part{typeof(grouped_normal)}) = let (m, n) = size(p.data.values); m * nparams(p::Part{typeof(chol)}) = let n = size(p.data.L, 1); n * (n + 1) ÷ 2 end nparams(p::Part{typeof(simplex)}) = length(p.data.values) - 1 -advance!!(x, pos) = x[pos+1], pos+1 -advance!!(x, pos, n) = view(x, pos+1:pos+n), pos+n +@inline advance!!(x, pos) = x[pos+1], pos+1 +@inline advance!!(x, pos, n) = view(x, pos+1:pos+n), pos+n """ lprior!(container, x) -> lp @@ -266,14 +270,14 @@ exactly `nparams(child)` unconstrained reals; bottoms out on `Part` methods, each of which writes its constrained buffer and returns its log-prior + Jacobian contribution. """ -lprior!(vbrmi::VBRMI, x::AbstractVector) = lprior!(vbrmi.meta.blocks, x) -lprior!(xs::Union{Tuple,NamedTuple}, x::AbstractVector) = +@inline lprior!(vbrmi::VBRMI, x::AbstractVector) = lprior!(vbrmi.meta.blocks, x) +@inline lprior!(xs::Union{Tuple,NamedTuple}, x::AbstractVector) = foldl(xs; init=(0.0, 0)) do (total, pos), child xi, pos = advance!!(x, pos, nparams(child)) total + lprior!(child, xi), pos end |> first -lprior!(p::Part{typeof(normal)}, x) = begin +@inline lprior!(p::Part{typeof(normal)}, x) = begin p.data.values[1, :] .= x sum(Base.Fix1(logpdf, Normal()), x) end @@ -291,7 +295,7 @@ Either wrong or better LKJCholesky unconstraining + prior. Writes `p.data.L` in place from `x` (length n(n+1)/2) and returns the log-prior + Jacobian contribution. `eta` is the LKJ shape parameter. """ -lprior!(p::Part{typeof(chol)}, x; eta=1.0) = begin +@inline lprior!(p::Part{typeof(chol)}, x; eta=1.0) = begin L = p.data.L n = LinearAlgebra.checksquare(L) pos = 0 @@ -323,7 +327,7 @@ lprior!(p::Part{typeof(chol)}, x; eta=1.0) = begin lprior end -lprior!(p::Part{typeof(grouped_normal)}, x) = begin +@inline lprior!(p::Part{typeof(grouped_normal)}, x) = begin (; values, L) = p.data _, n = size(values) lprior = 0.0 @@ -348,7 +352,7 @@ transform (plus the constant `loggamma(K)`). Ported from `blog/posts/simplex/transforms/stickbreakingLogistic.stan`. """ -lprior!(p::Part{typeof(simplex)}, x; alpha=1.0) = begin +@inline lprior!(p::Part{typeof(simplex)}, x; alpha=1.0) = begin values = p.data.values K = length(values) lp = 0.0 @@ -366,14 +370,14 @@ lprior!(p::Part{typeof(simplex)}, x; alpha=1.0) = begin lp end -llikelihood!((;meta)::VBRMI) = foldl(meta.materialized; init=0.) do llikelihood, m +@inline llikelihood!((;meta)::VBRMI) = foldl(meta.materialized; init=0.) do llikelihood, m llikelihood + llikelihood!(m) end -llikelihood!(::DataColumn) = 0. -llikelihood!(x::MaterializedColumn) = (Base.materialize!(parent(x), getbroadcast(x)); 0.) +@inline llikelihood!(::DataColumn) = 0. +@inline llikelihood!(x::MaterializedColumn) = (Base.materialize!(parent(x), getbroadcast(x)); 0.) # llikelihood!(x::LikelihoodColumn) = sum(Base.broadcasted(logpdf, rhs(x), parent(x)); init=0.) # The below is faster for some reason? -llikelihood!(x::LikelihoodColumn) = ssum(Base.broadcasted(logpdf, rhs(x), parent(x)); init=0.) +@inline llikelihood!(x::LikelihoodColumn) = ssum(Base.broadcasted(logpdf, rhs(x), parent(x)); init=0.) llikelihood!(x) = error(typeof(x)) ssum(args...; kwargs...) = sum(args...; kwargs...) @@ -471,7 +475,7 @@ Zero-parameter follow-up to a `simplex` sibling: refresh `contrast = vcat(0, cumsum(values))` in place so downstream broadcasts see up-to-date contrast values. Contributes nothing to the log-prior. """ -lprior!(p::Part{typeof(mo1)}, _) = begin +@inline lprior!(p::Part{typeof(mo1)}, _) = begin (; values, contrast) = p.data contrast[1] = 0.0 for k in 2:length(contrast) From adae6e84fb51598ef6d0be0099662306debf0c27 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Mon, 20 Apr 2026 22:41:00 +0200 Subject: [PATCH 13/23] todos: status/formula edits written back via web UI Status toggles and textarea edits persisted from the /todo page (1.1 Bernoulli/Binomial formula switches c1 -> c2; several items reopened). Co-Authored-By: Claude Opus 4.7 --- web-macro/todos/1.1-verify-bernoulli-binomial-done.jl | 4 ++-- ...2-offset-fixed-exposure-already-works-without-a-wrapper.jl | 2 +- web-macro/todos/1.3-i-expr-likely-already-works.jl | 2 +- web-macro/todos/1.4-scale-x-standardize-x.jl | 2 +- web-macro/todos/1.6-cache-levels-level_map-dense-gc_idx.jl | 2 +- .../todos/1.7-categoricalarrays-pooledarrays-integration.jl | 2 +- .../todos/2.2-configurable-categorical-reference-level.jl | 2 +- web-macro/todos/2.3-per-parameter-prior-scales.jl | 2 +- .../2.4-centered-non-centered-parameterization-toggle.jl | 2 +- web-macro/todos/2.6-multi-membership-random-effects-mm.jl | 2 +- ....7-se-weights-for-meta-analysis-and-weighted-regression.jl | 2 +- .../todos/3.10-dirichlet-process-non-parametric-models.jl | 2 +- web-macro/todos/3.11-zero-inflated-hurdle-likelihoods.jl | 2 +- 13 files changed, 14 insertions(+), 14 deletions(-) diff --git a/web-macro/todos/1.1-verify-bernoulli-binomial-done.jl b/web-macro/todos/1.1-verify-bernoulli-binomial-done.jl index 6aa9a26..5065ea8 100644 --- a/web-macro/todos/1.1-verify-bernoulli-binomial-done.jl +++ b/web-macro/todos/1.1-verify-bernoulli-binomial-done.jl @@ -1,6 +1,6 @@ # label: 1.1 verify Bernoulli/Binomial # tier: 1 -# status: done +# status: open #= **Status: done.** ✓ Confirmed that the existing `FBroadcasted{<:Type{<:Distribution}}` pass-through in `vimpl.jl` handles both Bernoulli and Binomial cleanly. @@ -8,7 +8,7 @@ =# -log_odds_bin ~ 1 + c1 + (1 | g1) +log_odds_bin ~ 1 + c2 + (1 | g1) bin_succ ~ Binomial(bin_n, logistic(log_odds_bin)) log_odds_b ~ 1 + (1 | g1) diff --git a/web-macro/todos/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl b/web-macro/todos/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl index 9ed9e36..5b1d1f4 100644 --- a/web-macro/todos/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl +++ b/web-macro/todos/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl @@ -1,6 +1,6 @@ # label: 1.2 offset / fixed exposure — already works without a wrapper # tier: 1 -# status: done +# status: open #= **Status: already works without any new code.** brms needs `offset(z)` because R's formula syntax has no other way to put a "no-coefficient term" into the linear predictor — the only thing on the RHS of `~` is the formula DSL. In our DSL the linear predictor and the likelihood are *separate* `~` lines, and the second one (the likelihood) takes a free-form Julia expression. Anything inside that expression gets evaluated as plain code at materialization time via `vbroadcasted` — function calls dispatch to whatever Julia function the symbol resolves to, and data column references are pulled from the dataframe. diff --git a/web-macro/todos/1.3-i-expr-likely-already-works.jl b/web-macro/todos/1.3-i-expr-likely-already-works.jl index a1aac08..fcf615a 100644 --- a/web-macro/todos/1.3-i-expr-likely-already-works.jl +++ b/web-macro/todos/1.3-i-expr-likely-already-works.jl @@ -1,6 +1,6 @@ # label: 1.3 I(expr) — likely already works # tier: 1 -# status: done +# status: open #= **What it is.** brms's `I()` is a literal-escape: `I(x^2)` says "compute `x^2` from the data and treat it as a single column". brms needs it because `+`, `*`, `:`, `|`, … all have special meaning inside an R formula. diff --git a/web-macro/todos/1.4-scale-x-standardize-x.jl b/web-macro/todos/1.4-scale-x-standardize-x.jl index 1a7218f..6757944 100644 --- a/web-macro/todos/1.4-scale-x-standardize-x.jl +++ b/web-macro/todos/1.4-scale-x-standardize-x.jl @@ -1,6 +1,6 @@ # label: 1.4 scale(x) / standardize(x) # tier: 1 -# status: deprioritized +# status: open #= **What it is.** brms's `scale(x)` z-transforms a column at parse time: `scale(x) = (x - mean(x)) / std(x)`. The model sees the standardized column. Crucial for default priors (which are scale-invariant only after standardization) and sampler stability (well-conditioned linear predictors). diff --git a/web-macro/todos/1.6-cache-levels-level_map-dense-gc_idx.jl b/web-macro/todos/1.6-cache-levels-level_map-dense-gc_idx.jl index 15447b8..a7d7016 100644 --- a/web-macro/todos/1.6-cache-levels-level_map-dense-gc_idx.jl +++ b/web-macro/todos/1.6-cache-levels-level_map-dense-gc_idx.jl @@ -1,6 +1,6 @@ # label: 1.6 cache levels / level_map / dense / gc_idx # tier: 1 -# status: deprioritized +# status: open #= **What it is.** Stop rebuilding the dense level mapping (`Dict(level => row_index)`) and the gc_idx vector on every `VBRMI(brmi)` call. Cache them once per source data column. diff --git a/web-macro/todos/1.7-categoricalarrays-pooledarrays-integration.jl b/web-macro/todos/1.7-categoricalarrays-pooledarrays-integration.jl index 0240eca..cba1895 100644 --- a/web-macro/todos/1.7-categoricalarrays-pooledarrays-integration.jl +++ b/web-macro/todos/1.7-categoricalarrays-pooledarrays-integration.jl @@ -1,6 +1,6 @@ # label: 1.7 CategoricalArrays / PooledArrays integration # tier: 1 -# status: deprioritized +# status: open #= **What it is.** When the input column is already a `CategoricalVector` or a `PooledArray`, the dense level mapping is already computed and stored in the column's `.refs` field. Use it directly instead of rebuilding via `Dict`. diff --git a/web-macro/todos/2.2-configurable-categorical-reference-level.jl b/web-macro/todos/2.2-configurable-categorical-reference-level.jl index d9ff19b..8b3824b 100644 --- a/web-macro/todos/2.2-configurable-categorical-reference-level.jl +++ b/web-macro/todos/2.2-configurable-categorical-reference-level.jl @@ -1,6 +1,6 @@ # label: 2.2 configurable categorical reference level # tier: 2 -# status: deprioritized +# status: open #= **What it is.** Currently the reference level for treatment-coded categoricals is `sort(unique(x))[1]`. brms / lme4 let you override this via `factor(x, ref="some_level")` or by reordering the factor's levels. diff --git a/web-macro/todos/2.3-per-parameter-prior-scales.jl b/web-macro/todos/2.3-per-parameter-prior-scales.jl index 6fa369f..d6cf595 100644 --- a/web-macro/todos/2.3-per-parameter-prior-scales.jl +++ b/web-macro/todos/2.3-per-parameter-prior-scales.jl @@ -1,6 +1,6 @@ # label: 2.3 per-parameter prior scales # tier: 2 -# status: deprioritized +# status: open #= **What it is.** Currently every parameter is `Normal(0, 1)` in `lprior!`. brms / Stan-style models routinely set custom priors per coefficient: `b ~ Normal(0, 0.5)` for tight priors on slopes, `b ~ Cauchy(0, 1)` for heavy-tailed priors, etc. diff --git a/web-macro/todos/2.4-centered-non-centered-parameterization-toggle.jl b/web-macro/todos/2.4-centered-non-centered-parameterization-toggle.jl index be83585..0cceffd 100644 --- a/web-macro/todos/2.4-centered-non-centered-parameterization-toggle.jl +++ b/web-macro/todos/2.4-centered-non-centered-parameterization-toggle.jl @@ -1,6 +1,6 @@ # label: 2.4 centered / non-centered parameterization toggle # tier: 2 -# status: deprioritized +# status: open #= **What it is.** Currently every random-effect block uses non-centered parameterization (we sample standard normals and apply `mul!(vi, C.L, xi)`). brms / Stan let you choose centered (sample directly from `Normal(0, σ)` per group) on a per-factor basis. diff --git a/web-macro/todos/2.6-multi-membership-random-effects-mm.jl b/web-macro/todos/2.6-multi-membership-random-effects-mm.jl index 8f0fe11..c9f2013 100644 --- a/web-macro/todos/2.6-multi-membership-random-effects-mm.jl +++ b/web-macro/todos/2.6-multi-membership-random-effects-mm.jl @@ -1,6 +1,6 @@ # label: 2.6 multi-membership random effects mm() # tier: 2 -# status: deprioritized +# status: open #= **What it is.** brms's `mm(g1, g2, ...)` lets one observation belong to **multiple** levels of the same random factor simultaneously, with weights summing to 1. Standard use: a student belongs to multiple schools across the year, and we want their random effect to be a weighted average of the per-school effects. diff --git a/web-macro/todos/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl b/web-macro/todos/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl index 8232f06..a6179d0 100644 --- a/web-macro/todos/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl +++ b/web-macro/todos/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl @@ -1,6 +1,6 @@ # label: 2.7 se() / weights() — already works (just a different distribution) # tier: 2 -# status: done +# status: open #= **Status: already works without any new code.** Same insight as 1.2 (offset) and 1.3 (`I()`): brms needs sidecar syntax (`y | se(sigma_y) ~ ...`, `y | weights(w) ~ ...`) because R's formula DSL has no other way to attach extra info to the LHS. Our DSL has no such constraint — the likelihood is a free-form Julia expression, so anything brms expresses via sidecar syntax we can express via constructor arguments or a wrapper distribution. diff --git a/web-macro/todos/3.10-dirichlet-process-non-parametric-models.jl b/web-macro/todos/3.10-dirichlet-process-non-parametric-models.jl index 52e1bf1..785a6bf 100644 --- a/web-macro/todos/3.10-dirichlet-process-non-parametric-models.jl +++ b/web-macro/todos/3.10-dirichlet-process-non-parametric-models.jl @@ -1,6 +1,6 @@ # label: 3.10 Dirichlet process / non-parametric models # tier: 3 -# status: deprioritized +# status: open #= **What it is.** Models where the number of components / clusters / random-effect levels is itself inferred during sampling, via a Dirichlet process or stick-breaking prior. diff --git a/web-macro/todos/3.11-zero-inflated-hurdle-likelihoods.jl b/web-macro/todos/3.11-zero-inflated-hurdle-likelihoods.jl index ee5af2f..d5cdb29 100644 --- a/web-macro/todos/3.11-zero-inflated-hurdle-likelihoods.jl +++ b/web-macro/todos/3.11-zero-inflated-hurdle-likelihoods.jl @@ -1,6 +1,6 @@ # label: 3.11 zero-inflated / hurdle likelihoods # tier: 3 -# status: done +# status: open #= **What it is.** ZI Poisson, ZI Negative Binomial, hurdle Poisson, hurdle Gamma, … — likelihoods that mix a point mass at zero (or a separate "is zero" Bernoulli) with a continuous/count distribution for the nonzero values. From ba7fb5d9318b00704291bac282d9aa3d6907d58c Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Tue, 21 Apr 2026 04:17:11 +0200 Subject: [PATCH 14/23] web-macro: add StanBlocks backend (sbimpl) + :stan_code pipeline stage sbimpl walks a BRMI directly and emits a @slic model that transpiles to Stan. Covers population fixed effects, `mo(c)` / `mo1(c)` monotonic terms, treatment-coded categorical predictors (K-1 betas), LKJ-correlated random effects `(1 + x + ... | g)` with per-group merging, and a Normal likelihood. Pipeline branches after :brmi -- :bench stays on VBRMI, :stan_code shortcircuits into SBBRMI without paying for materialization. TODO cards sb.1..sb.6 scaffold the roadmap for remaining features (distributional, ranefs done, categorical done, non-Normal likelihoods, submodels). Co-Authored-By: Claude Opus 4.7 --- web-macro/Project.toml | 2 + web-macro/src/BRMMacroWeb.jl | 56 +- web-macro/src/sbimpl.jl | 496 ++++++++++++++++++ web-macro/todos/sb.1-linear-regression.jl | 19 + web-macro/todos/sb.2-distributional.jl | 15 + web-macro/todos/sb.3-random-effects.jl | 20 + .../todos/sb.4-categorical-predictors.jl | 20 + .../todos/sb.5-non-normal-likelihoods.jl | 26 + web-macro/todos/sb.6-submodels-mo-mm.jl | 27 + 9 files changed, 670 insertions(+), 11 deletions(-) create mode 100644 web-macro/src/sbimpl.jl create mode 100644 web-macro/todos/sb.1-linear-regression.jl create mode 100644 web-macro/todos/sb.2-distributional.jl create mode 100644 web-macro/todos/sb.3-random-effects.jl create mode 100644 web-macro/todos/sb.4-categorical-predictors.jl create mode 100644 web-macro/todos/sb.5-non-normal-likelihoods.jl create mode 100644 web-macro/todos/sb.6-submodels-mo-mm.jl diff --git a/web-macro/Project.toml b/web-macro/Project.toml index 5a0e9c4..71df3e7 100644 --- a/web-macro/Project.toml +++ b/web-macro/Project.toml @@ -22,6 +22,8 @@ OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" +StanBlocks = "2e771a56-c23a-4e0b-9282-20c2e37157e9" +StanLogDensityProblems = "a545de4d-8dba-46db-9d34-4e41d3f07807" TestModules = "63c02187-99fd-4e5c-aaf0-4d6bfebc181c" Treebars = "e1e568c4-3a56-40a4-95fa-9b9c6c16fccb" Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" diff --git a/web-macro/src/BRMMacroWeb.jl b/web-macro/src/BRMMacroWeb.jl index 57db88c..416c2d3 100644 --- a/web-macro/src/BRMMacroWeb.jl +++ b/web-macro/src/BRMMacroWeb.jl @@ -12,6 +12,7 @@ using FiniteDifferences: FiniteDifferences, central_fdm # StanBlocksImpl) include them via relative paths into here. include("macro.jl") include("vimpl.jl") +include("sbimpl.jl") # ── Default formula + synthetic data ──────────────────────────────────────── @@ -120,7 +121,7 @@ _ALLOWED_CALLS = Set{Symbol}([ :logistic, :logit, :softmax, :logsumexp, :log_abs_tanh, :log_square_tanh, # Distributions (Type constructors — the pass-through handles these) - :Normal, :Poisson, :Binomial, :Bernoulli, :Beta, :Gamma, + :Normal, :Poisson, :Binomial, :Bernoulli, :BernoulliLogit, :Beta, :Gamma, :Exponential, :Cauchy, :StudentT, :LogNormal, :Weibull, :NegativeBinomial, :Geometric, :Laplace, :Uniform, :MvNormal, :MixtureModel, :Dirichlet, @@ -203,7 +204,9 @@ STAGES = ( :wrap, # _brm(formula; df) — full let-block ready to eval :brmi, # eval(...) — BRMI value :vbrmi, # VBRMI(brmi) — materialized action with blocks/dim + # ── branches after vbrmi (pick one) ── :bench, # Chairmarks @be primal logdensity + :stan_code, # SBBRMI(brmi): emit @slic body (a) + transpile to Stan (b) ) stage_index(s::Symbol) = something(findfirst(==(s), STAGES), length(STAGES)) @@ -238,6 +241,16 @@ function pipeline(formula::AbstractString, stage::Symbol; brmi = eval(wrapped) out = merge(out, (; brmi)) + # ── divergent branches after brmi ───────────────────────────────────── + # :stan_code targets the StanBlocks backend, which does NOT need VBRMI — + # SBBRMI walks the BRMI directly. Short-circuit here so the SB branch + # doesn't pay for VBRMI materialization. + if stage === :stan_code + sbbrmi = SBBRMI(brmi) + stan_src = stan_code(sbbrmi) + return merge(out, (; sbbrmi, stan_src)) + end + s >= 5 || return out vbrmi = VBRMI(brmi) dim = LogDensityProblems.dimension(vbrmi) @@ -250,7 +263,7 @@ function pipeline(formula::AbstractString, stage::Symbol; )[1] out = merge(out, (; vbrmi, dim, ldp, x0, grad)) - s >= 6 || return out + stage === :bench || return out x_rand = randn(dim) benches = Pair{String,Any}[] push!(benches, "logdensity (total)" => @@ -547,6 +560,18 @@ function render_pipeline(out::NamedTuple) ) for (label, b) in out.benches] push!(sections, Any[h.h3("6. Chairmarks @be — per-step"), bench_rows...]) end + if haskey(out, :sbbrmi) + push!(sections, Any[ + h.h3("5. SBBRMI — emitted @slic body"), + h.pre(sprint(show, out.sbbrmi.model.model)), + h.p("data keys: ", + h.code(string(sort(collect(keys(out.sbbrmi.data)))))), + ]) + push!(sections, Any[ + h.h3("6. transpiled Stan source"), + h.pre(out.stan_src), + ]) + end # Stages render most-recent-first; synthetic data sits at the very top. children = reduce(vcat, reverse(sections); init=Any[]) @@ -632,7 +657,11 @@ _index_body(self, formula::String) = h.div( _stage_button(self, "3. Wrap", :wrap), _stage_button(self, "4. BRMI", :brmi), _stage_button(self, "5. VBRMI", :vbrmi), - _stage_button(self, "6. Benchmark", :bench), + ), + h.small("Pick a branch:"), + h.fieldset(; class="grid")( + _stage_button(self, "6. Benchmark", :bench), + _stage_button(self, "6. Stan code", :stan_code), ), ), lazy(string(query_url(self/"stage/bench"; formula)); id="brm-macro-output"), @@ -931,20 +960,25 @@ _permalink(todo::TodoEntry) = h.a("🔗"; ) function _formula_form(self, label::String, formula::String) - h.form(; - hx_get=string(query_url(self/"stage/bench"; force=true)), - hx_target="#todo-result-$(hash(label))", + # Two terminal branches — cimpl (julia/VBRMI benchmark) vs sbimpl (Stan + # source). Each button is `type=button` (not submit) and carries its own + # `hx_get`; `hx_include="closest form"` pulls the textarea + hidden label. + target = "#todo-result-$(hash(label))" + _branch_button(text, stage) = h.button(text; + type="button", + hx_get=string(query_url(self/"stage/$stage"; force=true)), + hx_include="closest form", + hx_target=target, hx_swap="innerHTML", - style="margin:0.5rem 0", - )( + style="font-size:0.85em;padding:0.3rem 0.8rem;margin:0.3rem 0.3rem 0 0") + h.form(; style="margin:0.5rem 0")( h.input(; type="hidden", name="label", value=label), h.textarea(formula; name="formula", rows=max(3, count('\n', formula) + 1), style="width:100%;font-family:monospace;font-size:0.85em"), - h.button("Try in pipeline ▶"; - type="submit", - style="font-size:0.85em;padding:0.3rem 0.8rem;margin:0.3rem 0 0 0"), + _branch_button("cimpl (bench) ▶", :bench), + _branch_button("sbimpl (Stan) ▶", :stan_code), ) end diff --git a/web-macro/src/sbimpl.jl b/web-macro/src/sbimpl.jl new file mode 100644 index 0000000..93d9b6c --- /dev/null +++ b/web-macro/src/sbimpl.jl @@ -0,0 +1,496 @@ +using StanBlocks + + +# ============================================================================== +# SlicModel helpers (ported verbatim from /home/niko/github/nsiccha/bruno/src/qt.jl, +# `popefs`/`ranefs`/`popranefs`/`cdirichlet` family, lines 303-344). Kept here +# as module-local bindings so the walker can emit calls to them by name without +# depending on bruno. Duplication is intentional for now. +# ============================================================================== + +const popefs = StanBlocks.@slic begin + n_covariates = dims(X)[2] + beta_pop ~ std_normal(; n=n_covariates) + return X * beta_pop +end + +const cdirichlet = StanBlocks.@slic begin + increments ~ dirichlet(alpha) + return cumulative_sum(increments) +end + +const c0dirichlet = StanBlocks.@slic begin + increments ~ dirichlet(alpha) + return cumulative_sum(increments) - increments[1] +end + +const c01dirichlet = StanBlocks.@slic begin + increments ~ dirichlet(alpha) + return append_row(0., cumulative_sum(increments)) +end + +# Monotonic effect contrast (Buerkner & Charpentier 2018). Returns the per-obs +# contrast vector; the walker hcat's it as one column of X_pop so popefs +# supplies the free beta (matches vimpl's free-beta `mo` variant, not `mo1`). +# Named `_sb_mo` to avoid clashing with vimpl's marker function `mo`. +const _sb_mo = StanBlocks.@slic begin + n_levels = maximum(x) + simplex_incr ~ dirichlet(rep_vector(1., n_levels - 1)) + return cumulative_sum(append_row(0., simplex_incr))[x] +end + +# (1 | g) random intercept. Mirrors vimpl's scalar grouped_normal + chol(n=1) +# collapse: Part{chol}(1x1) -> log_scale ~ N(0,1), L[1,1] = exp(log_scale); +# Part{grouped_normal}(n_groups, 1) -> xi ~ N(0,1), values = L[1,1] * xi; +# per-obs contribution is values[group_idx]. No LKJ needed at n=1. +const ranef_intercept = StanBlocks.@slic begin + log_scale ~ std_normal() + xi ~ std_normal(; n=n_groups) + return exp(log_scale) * xi[group_idx] +end + +# Correlated random effects for K terms x G groups. brms-style (1 + x + y | g). +# Non-centered parameterization: +# L ~ lkj_corr_cholesky(1, K) # K x K Cholesky factor +# tau ~ half-std_normal(; n=K) # per-term marginal scales +# z ~ std_normal(; n=K, m=n_groups) # K x n_groups std normal +# b = (diag_pre_multiply(tau, L) * z)' # n_groups x K correlated draws +# Per-row contribution = Z[i, :] . b[group_idx[i], :], returned as a length-n +# vector via rows_dot_product. Note: `(1 | g) + (0 + x | g)` and `(1 + x | g)` +# are equivalent -- the walker merges everything sharing a group symbol into +# one correlated block. +const ranef_correlated = StanBlocks.@slic begin + L ~ lkj_corr_cholesky(1.; n=n_terms) + tau ~ std_normal(; n=n_terms, lower=0.) + z_flat ~ std_normal(; n=n_terms * n_groups) + z = reshape(z_flat, n_terms, n_groups) + b = (diag_pre_multiply(tau, L) * z)' # n_groups x n_terms + return rows_dot_product(Z, b[group_idx, :]) +end + +# Treatment-coded categorical predictor. Allocates K-1 free betas; reference +# level 1 contributes 0. Mirrors vimpl's `AbstractVector{<:Integer}` dispatch. +# `x` is the per-row 1-based level index, `n_levels = K`. +const _sb_cat = StanBlocks.@slic begin + beta ~ std_normal(; n=n_levels - 1) + return append_row(0., beta)[x] +end + +# Categorical -> (n_levels::Int, per-row level index::Vector{Int}). Mirrors +# vimpl._level_index so the integer indices the walker stashes in `data` +# agree with what the cimpl-side uses. +_sb_level_index(raw::CA.CategoricalVector) = length(CA.levels(raw)), Int.(CA.levelcode.(raw)) +_sb_level_index(raw::AbstractVector) = begin + lvls = sort(unique(raw)) + lm = Dict(l => i for (i, l) in enumerate(lvls)) + length(lvls), [lm[l] for l in raw] +end + + +# ============================================================================== +# SBBRMI: walk a BRMI, emit a SlicModel whose body references raw data columns +# by name. Scope for phase 1: population fixed effects + Normal likelihood. +# Ranefs / categorical / non-Normal likelihoods error out clearly. +# +# Design contract: data columns are referenced in the emitted expression by +# their formula names. `hcat(rep_vector(1., n), a, c1)` builds the design +# matrix inside Stan at runtime, so changing the data size (= length of the +# vectors) does not require recompiling the Stan model. StanBlocks' activity +# analysis routes each statement to the right Stan block (data / transformed +# data / parameters / model). +# ============================================================================== + +struct SBBRMI{P<:BRMI, M, D<:AbstractDict} + parent::P + model::M + data::D +end + +SBBRMI(brmi::BRMI) = begin + stmts = Any[] + data = Dict{Symbol,Any}() + # Prepass: stash every data-backed NamedColumn so later intercept-only + # predictors have a length probe to hang `rep_vector(1., num_elements(...))` + # off, regardless of iteration order. + for (_, op) in pairs(brmi.operations) + _sb_collect_data!(data, op) + end + for (key, op) in pairs(brmi.operations) + op isa NamedColumn || error("sbimpl: top-level op `$key` is not a NamedColumn") + _sb_emit!(stmts, data, key, parent(op)) + end + body = Expr(:block, stmts...) + model = StanBlocks.SlicModel(body, data, @__MODULE__) + SBBRMI(brmi, model, data) +end + +_sb_collect_data!(data, x) = nothing +_sb_collect_data!(data, x::NamedColumn) = begin + d = parent(x) + d isa DataColumn && (data[name(x)] = parent(d)) + _sb_collect_data!(data, d) +end +_sb_collect_data!(data, x::ExprColumn) = foreach(a -> _sb_collect_data!(data, a), getargs(x)) + +stan_code(sb::SBBRMI) = StanBlocks.stan_code(sb.model) + +Base.show(io::IO, sb::SBBRMI) = begin + print(io, "SBBRMI with data keys = ", sort(collect(keys(sb.data))), "\n") + print(io, "emitted @slic body:\n") + print(io, sb.model.model) +end + + +# ---- top-level op dispatch --------------------------------------------------- + +_sb_emit!(stmts, data, key, op::ExprColumn) = _sb_emit_expr!(stmts, data, key, getf(op), op) +# Raw data / missing columns appear as top-level ops when the formula mentions +# them as bare references (e.g. `c2` in `loc ~ 1 + c2`). Nothing to emit — the +# prepass already stashed data columns in `data`. +_sb_emit!(_, _, _, ::DataColumn) = nothing +_sb_emit!(_, _, _, ::MissingColumn) = nothing +_sb_emit!(_, _, key, op) = error("sbimpl: top-level op for `$key` not an ExprColumn (got $(typeof(op)))") + +_sb_emit_expr!(stmts, data, key, ::typeof(~), op) = begin + lhs, rhs = getargs(op, 2) + _sb_sampling!(stmts, data, key, lhs, rhs) +end +_sb_emit_expr!(stmts, data, key, ::typeof(assign), op) = begin + _, rhs = getargs(op, 2) + target_expr = _sb_scalar_expr(rhs, data) + push!(stmts, :($key = $target_expr)) +end +_sb_emit_expr!(_, _, key, f, _) = error("sbimpl: unsupported top-level op `$f` for `$key`") + + +# ---- sampling: likelihood vs linear-predictor split -------------------------- + +# Extension hook. Overridden (e.g. in bruno-ext.jl) to route `target ~ f(...)` +# where `f` is a known submodel family (`logistic_dr`, `gamma_time`, etc.) +# straight to `target ~ (; kwargs)`, bypassing the pop-linear-predictor +# wrap (which would wrongly multiply the submodel output by a beta). Return +# anything non-`nothing` to claim the binding; return `nothing` to fall +# through. Default: no hook registered. +_sb_submodel_rhs!(stmts, data, target, f, rhs) = nothing + +# LHS backed by real data => this is a likelihood. Record the observed values +# under the formula name in `data` and emit `key ~ dist(args...)`. +_sb_sampling!(stmts, data, key, lhs::NamedColumn, rhs) = begin + backing = parent(lhs) + if backing isa DataColumn + data[key] = parent(backing) + push!(stmts, _sb_likelihood(key, rhs, data)) + elseif backing isa MissingColumn + if rhs isa ExprColumn && + _sb_submodel_rhs!(stmts, data, key, getf(rhs), rhs) !== nothing + return + end + _sb_linear_predictor!(stmts, data, key, rhs) + else + error("sbimpl: unsupported LHS backing for `$key` ($(typeof(backing)))") + end +end + +# Link-transformed LHS: `log(err) ~ 1 + d` etc. +_sb_sampling!(stmts, data, key, lhs::ExprColumn, rhs) = begin + f = getf(lhs) + f === log || error("sbimpl: only `log(x)` links supported for now (got `$f`)") + inner = getargs(lhs, 1)[1] + inner isa NamedColumn || error("sbimpl: expected NamedColumn inside link, got $(typeof(inner))") + inner_name = name(inner) + log_name = Symbol(:log_, inner_name) + _sb_linear_predictor!(stmts, data, log_name, rhs) + push!(stmts, :($inner_name = exp($log_name))) +end + + +# ---- linear predictor: emit `X_ = hcat(...); ~ popefs(; X=X_)` -- + +function _sb_linear_predictor!(stmts, data, target::Symbol, rhs) + terms = _sb_terms(rhs) + pop_terms = Any[] + ran_terms = Any[] # `(expr | group)` -> collected per-group below + direct_terms = Any[] # e.g. `mo1(c)` -> direct summand, no popefs beta + for t in terms + if t isa ExprColumn && getf(t) === (|) + push!(ran_terms, t) + elseif t isa ExprColumn && getf(t) === mo1 + push!(direct_terms, t) + elseif _sb_is_categorical(t) + push!(direct_terms, t) + else + push!(pop_terms, t) + end + end + isempty(pop_terms) && isempty(ran_terms) && isempty(direct_terms) && + error("sbimpl: empty RHS for `$target` — no predictor terms") + + summands = Symbol[] + + if !isempty(pop_terms) + col_exprs = Any[_sb_predictor_col(t, data, stmts) for t in pop_terms] + X_name = Symbol(:X_, target) + pop_name = Symbol(:pop_, target) + # StanBlocks `hcat` promotes a lone vector to matrix[n,1] and folds to + # append_col for two-or-more columns, so we can always just emit hcat. + push!(stmts, :($X_name = $(Expr(:call, :hcat, col_exprs...)))) + push!(stmts, :($pop_name ~ popefs(; X=$X_name))) + push!(summands, pop_name) + end + + for dt in direct_terms + _sb_emit_direct!(stmts, data, target, dt, summands) + end + + _sb_emit_ranefs!(stmts, data, target, ran_terms, summands) + + if length(summands) == 1 + push!(stmts, :($target = $(only(summands)))) + else + push!(stmts, :($target = $(Expr(:call, :+, summands...)))) + end +end + +# Classify a predictor term as "direct" (allocates its own parameters, no +# popefs multiplication). Matches vimpl: integer-backed NamedColumns are +# treated as treatment-coded categoricals; floats go through popefs. +_sb_is_categorical(t::NamedColumn) = begin + d = parent(t) + d isa DataColumn || return false + v = parent(d) + v isa AbstractVector{<:Integer} || v isa CA.CategoricalVector +end +_sb_is_categorical(_) = false + +# Free-summand terms (no popefs beta): `mo1(c)`, categorical NamedColumns. +# Categoricals emit `cat_ ~ _sb_cat(; x=_idx, n_levels=_n_levels)`. +# `mo1(c)` reuses the `_sb_mo` submodel already used for free-beta `mo(c)`. +function _sb_emit_direct!(stmts, data, target::Symbol, t, summands) + if t isa NamedColumn + _sb_emit_cat!(stmts, data, t, summands) + return + end + f = getf(t) + f === mo1 || error("sbimpl: unsupported direct-summand term `$f`") + inner = only(getargs(t)) + inner isa NamedColumn || error("sbimpl: `mo1(…)` expects a NamedColumn, got $(typeof(inner))") + backing = parent(inner) + backing isa DataColumn || error("sbimpl: `mo1($(name(inner)))` expects a raw data column, got $(typeof(backing))") + n_levels, idx = _sb_level_index(parent(backing)) + n_levels >= 2 || error("sbimpl: `mo1($(name(inner)))` needs >= 2 levels (got $n_levels)") + idx_name = Symbol(name(inner), :_idx) + col_name = Symbol(:mo1_, name(inner)) + data[idx_name] = idx + push!(stmts, :($col_name ~ _sb_mo(; x=$idx_name))) + push!(summands, col_name) +end + +# Categorical population-level predictor. Allocates K-1 betas via `_sb_cat` +# and pushes the per-row contribution column into `summands`. +function _sb_emit_cat!(stmts, data, t::NamedColumn, summands) + backing = parent(t) + n_levels, idx = _sb_level_index(parent(backing)) + n_levels >= 2 || error("sbimpl: categorical `$(name(t))` needs >= 2 levels (got $n_levels)") + idx_name = Symbol(name(t), :_idx) + n_name = Symbol(name(t), :_n_levels) + col_name = Symbol(:cat_, name(t)) + data[idx_name] = idx + data[n_name] = n_levels + push!(stmts, :($col_name ~ _sb_cat(; x=$idx_name, n_levels=$n_name))) + push!(summands, col_name) +end + +# Expand a ranef LHS term into one-or-more design-matrix column references. +# Continuous/intercept terms produce a single column; categorical NamedColumns +# expand to K-1 treatment-coded dummy columns (level 1 is reference), matching +# the design-matrix that brms / lme4 build for `(1 + c | g)`. +function _sb_ranef_cols!(cols, data, stmts, t) + if _sb_is_categorical(t) + backing = parent(t) + n_levels, idx = _sb_level_index(parent(backing)) + n_levels >= 2 || error("sbimpl: categorical ranef term `$(name(t))` needs >= 2 levels (got $n_levels)") + for lvl in 2:n_levels + col_name = Symbol(name(t), :_dummy_, lvl) + data[col_name] = Float64[l == lvl ? 1.0 : 0.0 for l in idx] + push!(cols, col_name) + end + else + push!(cols, _sb_predictor_col(t, data, stmts)) + end +end + +# Collected ranef handling. Terms that share a grouping symbol are merged into +# a single correlated block (matches vimpl / brms: `(1 | g) + (x | g)` is the +# same as `(1 + x | g)`; `(1 | g) + (0 + x | g)` also collapses -- the two +# terms' LKJ and tau parameters are shared by design). Per-group we build: +# Z__ = hcat(, , ...) # n x K +# and emit one `~ ranef_correlated(...)` (K >= 2) or `~ ranef_intercept(...)` +# (K == 1, intercept-only) call. Scalar-slope-only blocks (K == 1 but not an +# intercept) also go through ranef_correlated -- the math degenerates gracefully. +function _sb_emit_ranefs!(stmts, data, target::Symbol, ran_terms, summands) + isempty(ran_terms) && return + # Group `(expr | g)` terms by the group symbol, preserving first-seen order. + groups = Symbol[] + by_group = Dict{Symbol, Vector{Any}}() + for rt in ran_terms + lhs, group = getargs(rt, 2) + group isa NamedColumn || error("sbimpl: expected NamedColumn on RHS of `|`, got $(typeof(group))") + g_backing = parent(group) + g_backing isa DataColumn || error("sbimpl: group `$(name(group))` must be a raw data column") + g = name(group) + haskey(by_group, g) || (push!(groups, g); by_group[g] = Any[]) + # Flatten `lhs` on `+` (same walker as pop terms; reuses `0` drop). + append!(by_group[g], _sb_terms(lhs)) + end + for g in groups + gterms = by_group[g] + isempty(gterms) && error("sbimpl: ranef `(… | $g)` has no terms after dropping `0`") + # Resolve group index/count from the first occurrence (same backing data). + example = first(rt for rt in ran_terms if name(getargs(rt, 2)[2]) === g) + group_col = getargs(example, 2)[2] + n_levels, g_idx = _sb_level_index(parent(parent(group_col))) + idx_name = Symbol(g, :_idx) + n_name = Symbol(:n_, g) + data[idx_name] = g_idx + data[n_name] = n_levels + r_name = Symbol(:r_, target, :_, g) + if length(gterms) == 1 && gterms[1] === 1 + # (1 | g) fast/equivalent path -- stays on ranef_intercept so the + # emitted Stan matches existing sb.3 smoke tests bit for bit. + push!(stmts, :($r_name ~ ranef_intercept(; group_idx=$idx_name, n_groups=$n_name))) + else + col_exprs = Any[] + for t in gterms + _sb_ranef_cols!(col_exprs, data, stmts, t) + end + Z_name = Symbol(:Z_, target, :_, g) + k_name = Symbol(:n_terms_, target, :_, g) + data[k_name] = length(col_exprs) + push!(stmts, :($Z_name = $(Expr(:call, :hcat, col_exprs...)))) + push!(stmts, :($r_name ~ ranef_correlated(; + Z=$Z_name, group_idx=$idx_name, + n_groups=$n_name, n_terms=$k_name))) + end + push!(summands, r_name) + end +end + +# Extract pop terms from `1 + a + c1 [+ (...|g)]`. `0` is the standard +# formula-language drop-intercept marker (e.g. `loc ~ 0 + ftime`) and +# contributes no predictor column, so it's filtered out here. +_sb_terms(x) = (acc = Any[]; _sb_collect_terms!(acc, x); acc) +_sb_collect_terms!(acc, x::ExprColumn) = _sb_collect_terms_expr!(acc, getf(x), x) +_sb_collect_terms!(acc, x::Int) = x == 0 ? nothing : push!(acc, x) +_sb_collect_terms!(acc, x) = push!(acc, x) +_sb_collect_terms_expr!(acc, ::typeof(+), x) = foreach(a -> _sb_collect_terms!(acc, a), getargs(x)) +# `(expr | group)` is kept as-is; `_sb_linear_predictor!` splits it off into +# the ranef side of the additive linear predictor. +_sb_collect_terms_expr!(acc, ::typeof(|), x) = push!(acc, x) +_sb_collect_terms_expr!(acc, _, x) = push!(acc, x) + +# Predictor column emitter. `stmts` is threaded in so terms that need their own +# `~` statement (e.g. `mo(c)`) can push before returning their column symbol. +# Integer `1` -> intercept, NamedColumn -> reference by name, ExprColumn(mo, c) +# -> submodel-sampled contrast column. +_sb_predictor_col(t::Int, data, _stmts) = begin + t == 1 || error("sbimpl: integer term must be `1` for intercept, got `$t`") + # Use any data column to get n + probe = _sb_any_data_symbol(data) + :(rep_vector(1., num_elements($probe))) +end +_sb_predictor_col(t::NamedColumn, data, _stmts) = begin + d = parent(t) + # If the named column was already bound earlier in the walker (e.g. + # `ftime ~ gamma_time(...)` emitted a `ftime ~ _sb_gamma_time(...)` stmt), + # its parent is the sampling ExprColumn rather than a raw data column -- + # just reference the Stan variable by name. + d isa ExprColumn && return name(t) + d isa DataColumn || error("sbimpl: expected data-backed NamedColumn for `$(name(t))`, got $(typeof(d))") + v = parent(d) + v isa AbstractVector{<:Real} || + error("sbimpl: non-numeric predictor `$(name(t))` not supported yet (wrap in `categorical(…)` once we add it)") + data[name(t)] = collect(Float64, v) + name(t) +end +_sb_predictor_col(t::ExprColumn, data, stmts) = _sb_predictor_term!(stmts, data, getf(t), t) +_sb_predictor_col(t, _, _) = error("sbimpl: unsupported predictor term $(typeof(t)): $t") + +# Monotonic-effect predictor: emit `mo_ ~ _sb_mo(; x=_idx)` and return +# `mo_` as the column. Scope: single NamedColumn inner arg backed by raw +# data; `mo1`, `mm`, `s` etc. fall through the default error. +_sb_predictor_term!(stmts, data, ::typeof(mo), t) = begin + inner = only(getargs(t)) + inner isa NamedColumn || error("sbimpl: `mo(…)` expects a NamedColumn, got $(typeof(inner))") + backing = parent(inner) + backing isa DataColumn || error("sbimpl: `mo($(name(inner)))` expects a raw data column, got $(typeof(backing))") + n_levels, idx = _sb_level_index(parent(backing)) + n_levels >= 2 || error("sbimpl: `mo($(name(inner)))` needs >= 2 levels (got $n_levels)") + idx_name = Symbol(name(inner), :_idx) + col_name = Symbol(:mo_, name(inner)) + data[idx_name] = idx + push!(stmts, :($col_name ~ _sb_mo(; x=$idx_name))) + col_name +end +_sb_predictor_term!(_, _, f, _) = + error("sbimpl: unsupported predictor-term function `$f` (supported: `mo`, `mo1`)") + +_sb_n_obs_probe(terms) = begin + for t in terms + t isa NamedColumn && parent(t) isa DataColumn && return name(t) + end + nothing +end +_sb_any_data_symbol(data) = begin + isempty(data) && error("sbimpl: can't emit `rep_vector(1., n)` — no data column seen yet. Make sure an observed `~` comes before the intercept-only predictor, or add a concrete covariate.") + # Prefer any vector-valued entry + for (k, v) in data + v isa AbstractVector && return k + end + first(keys(data)) +end + + +# ---- likelihood emitters: `y ~ Normal(loc, sigma)` etc. ---------------------- + +function _sb_likelihood(target::Symbol, rhs::ExprColumn, data) + f = getf(rhs) + _sb_lik_family(target, f, getargs(rhs), data) +end +_sb_likelihood(target, rhs, _) = + error("sbimpl: likelihood RHS for `$target` must be an ExprColumn (got $(typeof(rhs)))") + +# One dispatch per likelihood family. Each method states the Stan name and +# implicitly the arity (by destructuring `args`). Caveat on NegativeBinomial: +# Distributions.jl parameterizes by (r, p); Stan's neg_binomial is (alpha, beta) +# — the emitted model is NOT posterior-identical to the Julia side. Convert +# upstream if that matters. +_sb_lik_stan(target, name::Symbol, args, data) = + Expr(:call, :~, target, + Expr(:call, name, (_sb_scalar_expr(a, data) for a in args)...)) + +_sb_lik_family(target, ::Type{<:Normal}, args::Tuple{Any,Any}, data) = _sb_lik_stan(target, :normal, args, data) +_sb_lik_family(target, ::Type{<:Bernoulli}, args::Tuple{Any}, data) = _sb_lik_stan(target, :bernoulli, args, data) +_sb_lik_family(target, ::Type{<:BernoulliLogit}, args::Tuple{Any}, data) = _sb_lik_stan(target, :bernoulli_logit, args, data) +_sb_lik_family(target, ::Type{<:Binomial}, args::Tuple{Any,Any}, data) = _sb_lik_stan(target, :binomial, args, data) +_sb_lik_family(target, ::Type{<:Poisson}, args::Tuple{Any}, data) = _sb_lik_stan(target, :poisson, args, data) +_sb_lik_family(target, ::Type{<:Gamma}, args::Tuple{Any,Any}, data) = _sb_lik_stan(target, :gamma, args, data) +_sb_lik_family(target, ::Type{<:NegativeBinomial}, args::Tuple{Any,Any}, data) = _sb_lik_stan(target, :neg_binomial, args, data) +_sb_lik_family(target, ::Type{<:Beta}, args::Tuple{Any,Any}, data) = _sb_lik_stan(target, :beta, args, data) + +_sb_lik_family(target, fam, args, _) = + error("sbimpl: likelihood family `$fam` (arity $(length(args))) not supported yet") + + +# ---- scalar-expression reducer (unwraps NamedColumn references etc.) -------- + +_sb_scalar_expr(x::Symbol, _) = x +_sb_scalar_expr(x::Real, _) = x +_sb_scalar_expr(x::NamedColumn, data) = begin + d = parent(x) + if d isa DataColumn + data[name(x)] = parent(d) + end + name(x) +end +_sb_scalar_expr(x::ExprColumn, data) = Expr(:call, getf(x), (_sb_scalar_expr(a, data) for a in getargs(x))...) +_sb_scalar_expr(x, _) = error("sbimpl: cannot lift to Stan expression: $(typeof(x)): $x") diff --git a/web-macro/todos/sb.1-linear-regression.jl b/web-macro/todos/sb.1-linear-regression.jl new file mode 100644 index 0000000..d75d37e --- /dev/null +++ b/web-macro/todos/sb.1-linear-regression.jl @@ -0,0 +1,19 @@ +# label: sb.1 linear regression (verified working) +# tier: 1 +# status: open +#= +**Status: should work.** Baseline sbimpl coverage: pop fixed effects (intercept + numeric columns), `log(x) ~ ...` link, `Normal(loc, scale)` likelihood. + +**What sbimpl emits.** +- Prepass stashes all data-backed NamedColumns into the SlicModel's data dict (so `num_elements()` resolves even when an intercept-only predictor precedes the first data reference). +- `loc ~ 1 + a + b` becomes `X_loc = hcat(rep_vector(1., num_elements(a)), a, b); loc ~ popefs(; X=X_loc)`. +- `log(err) ~ 1` becomes `X_log_err = reshape(rep_vector(1., num_elements(a)), :, 1); log_err ~ popefs(; X=X_log_err); err = exp(log_err)`. +- `y1 ~ Normal(loc, err)` becomes `y1 ~ normal(loc, err)` with `data[:y1] = df.y1`. + +**Verification.** Click "sbimpl (Stan) ▶" — should print a compilable Stan program. "cimpl (bench) ▶" should also finish with a healthy FD check. + +=# + +loc ~ 1 + a + b +log(err) ~ 1 +y1 ~ Normal(loc, err) diff --git a/web-macro/todos/sb.2-distributional.jl b/web-macro/todos/sb.2-distributional.jl new file mode 100644 index 0000000..21bce94 --- /dev/null +++ b/web-macro/todos/sb.2-distributional.jl @@ -0,0 +1,15 @@ +# label: sb.2 distributional (loc + scale predictors) +# tier: 1 +# status: open +#= +**Status: should work.** Two separate linear predictors — one for the location, one (via `log` link) for the scale — both pop-only. + +**Why it matters.** Distributional regression is the simplest place where the "activity analysis routes statements to the right Stan block" benefit shows up: `X_loc` / `X_log_err` are built in `transformed data`, `beta_pop`s land in `parameters`, the `~ popefs(...)` calls land in `model`. + +**Verification.** Stan source should show two `popefs` calls with independent `X_...` design matrices. + +=# + +loc ~ 1 + a + b +log(err) ~ 1 + c + d +y1 ~ Normal(loc, err) diff --git a/web-macro/todos/sb.3-random-effects.jl b/web-macro/todos/sb.3-random-effects.jl new file mode 100644 index 0000000..7d763c9 --- /dev/null +++ b/web-macro/todos/sb.3-random-effects.jl @@ -0,0 +1,20 @@ +# label: sb.3 random effects (1 | g) — needs impl +# tier: 2 +# status: open +#= +**Status: not implemented.** sbimpl currently errors on `ExprColumn{typeof(|)}` in `_sb_collect_terms_expr!`. + +**What to implement.** Emit `ranefs` / `popranefs` SB submodels that sample a non-centered grouped-normal block and add it to the popefs output. + +**Walker work.** When `_sb_collect_terms_expr!` sees `(expr | group)`: +1. Split terms into `pop_terms` and `ran_terms`, collect `Z` as the one-hot / index matrix for `group`. +2. Swap the emitted `loc ~ popefs(; X)` for `loc ~ popranefs(; X_pop, X_ran, Z)`. +3. Data: add `Z_` (or group indices) to the data dict. Design matrix widening still happens in Stan so data size doesn't bake in. + +**Verification.** Reuse the `bernoulli/binomial` smoke test once `sb.5` lands (non-Normal likelihoods), but a Normal version works as a first pass. + +=# + +loc ~ 1 + a + (1 + a | g1) +log(err) ~ 1 +y1 ~ Normal(loc, err) diff --git a/web-macro/todos/sb.4-categorical-predictors.jl b/web-macro/todos/sb.4-categorical-predictors.jl new file mode 100644 index 0000000..c87066e --- /dev/null +++ b/web-macro/todos/sb.4-categorical-predictors.jl @@ -0,0 +1,20 @@ +# label: sb.4 categorical predictors (c1, c2, ...) — needs impl +# tier: 2 +# status: open +#= +**Status: not implemented.** `_sb_predictor_col(t::NamedColumn, data)` errors when the backing vector isn't `AbstractVector{<:Real}`. + +**What to implement.** For a categorical column `c`: +1. At sbimpl time, one-hot encode the levels (drop first for reference coding, matching brms default). +2. Stash the one-hot matrix under `data[:_dummies]` (or just the integer level vector + lookup in Stan). +3. Emit additional columns into the `hcat(...)` design matrix. + +**Alternative.** Pass integer levels and let Stan construct the dummy matrix via a `@deffun` helper — keeps the Julia side thin, makes recompilation-on-data-size avoidable. + +**Verification.** `loc ~ 1 + c1` should produce `K-1` extra columns (where `K = length(levels(c1))`). Posterior dim should match the `cimpl` result. + +=# + +loc ~ 1 + a + c1 + c2 +log(err) ~ 1 +y1 ~ Normal(loc, err) diff --git a/web-macro/todos/sb.5-non-normal-likelihoods.jl b/web-macro/todos/sb.5-non-normal-likelihoods.jl new file mode 100644 index 0000000..a15d05f --- /dev/null +++ b/web-macro/todos/sb.5-non-normal-likelihoods.jl @@ -0,0 +1,26 @@ +# label: sb.5 non-Normal likelihoods — needs impl +# tier: 2 +# status: open +#= +**Status: not implemented.** `_sb_lik_family` dispatches only on `::Type{<:Normal}`. + +**What to implement.** Add a table of supported `Distribution` types -> Stan `_lpdf`/`_lpmf` names: + +| Julia | Stan | +|------------------------------|-----------------------------| +| `Normal(loc, scale)` | `normal(loc, scale)` | +| `Bernoulli(p)` | `bernoulli(p)` | +| `Binomial(n, p)` | `binomial(n, p)` | +| `Poisson(lambda)` | `poisson(lambda)` | +| `Gamma(alpha, beta)` | `gamma(alpha, beta)` | +| `NegativeBinomial(r, p)` | `neg_binomial(r, p)` | +| `Beta(a, b)` | `beta(a, b)` | + +Most map 1:1 — the Bernoulli/Binomial/Poisson logistic/exp links already canonicalize on the Julia side (`bin_succ ~ Binomial(bin_n, logistic(eta))`), so sbimpl just has to pass the arguments through `_sb_scalar_expr`. + +**Verification.** Reuse `1.1` formula (`bin_succ ~ Binomial(bin_n, logistic(log_odds_bin))` + hierarchical Bernoulli). Needs `sb.3` (ranefs) for the `(1 | g1)` piece. + +=# + +log_odds_b ~ 1 + a +bin_y ~ BernoulliLogit(log_odds_b) diff --git a/web-macro/todos/sb.6-submodels-mo-mm.jl b/web-macro/todos/sb.6-submodels-mo-mm.jl new file mode 100644 index 0000000..07148ee --- /dev/null +++ b/web-macro/todos/sb.6-submodels-mo-mm.jl @@ -0,0 +1,27 @@ +# label: sb.6 submodels (mo, mm, spline) as SlicModels — needs impl +# tier: 3 +# status: open +#= +**Status: not implemented.** The main architectural bet behind sbimpl: terms that introduce their own parameters (`mo(x)` monotonic effects, `mm(x)` multi-membership, `s(x)` splines, etc.) should map to their own `@slic` submodels with their own `~`-sampled hyperparameters, composed into the overall linear predictor via `popranefs`-style addition. + +**Why it matters.** Each submodel lives as a standalone `SlicModel`, so StanBlocks' activity analysis routes its parameters / priors / transforms to the right Stan block automatically. No per-term case analysis in sbimpl. + +**Prototype term: `mo(x)`** (Bürkner & Charpentier 2018). Emit: + +```julia +const mo = StanBlocks.@slic begin + n_levels = maximum(x) + simplex_incr ~ dirichlet(rep_vector(1., n_levels - 1)) + return cumulative_sum(append_row(0., simplex_incr))[x] +end +``` + +Then `loc ~ 1 + mo(c1)` walks to `mo_c1 = mo(; x=c1); loc ~ popefs(; X=hcat(..., mo_c1))`. + +**Verification.** `cimpl` has `Mo1Part` already (todo `3.3-ordinal-predictors-mo-monotonic-effects`); Stan dim should match. + +=# + +loc ~ 1 + a + mo(c1) +log(err) ~ 1 +y1 ~ Normal(loc, err) From 726461dd1b6a7773957dde09544362d3f4389d4b Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Tue, 21 Apr 2026 13:58:44 +0200 Subject: [PATCH 15/23] web-macro: stratified ranefs gr(g, by=b) + generalized link-transformed LHS Co-Authored-By: Claude Opus 4.7 --- web-macro/src/macro.jl | 6 + web-macro/src/sbimpl.jl | 198 ++++++++--- web-macro/src/vimpl.jl | 311 +++++++++++++++++- ...uped-random-effects-per-factor-variance.jl | 44 ++- 4 files changed, 497 insertions(+), 62 deletions(-) diff --git a/web-macro/src/macro.jl b/web-macro/src/macro.jl index 6f30dc6..3f6d2d0 100644 --- a/web-macro/src/macro.jl +++ b/web-macro/src/macro.jl @@ -45,6 +45,12 @@ end function assign end function doublepipe end function gr end +function gp end +function offset end +function scale end +function center end +function standardize end +function I end _brm(x::AbstractString; kwargs...) = _brm(Meta.parse(""" begin $x diff --git a/web-macro/src/sbimpl.jl b/web-macro/src/sbimpl.jl index 93d9b6c..8bf1f3e 100644 --- a/web-macro/src/sbimpl.jl +++ b/web-macro/src/sbimpl.jl @@ -68,6 +68,28 @@ const ranef_correlated = StanBlocks.@slic begin return rows_dot_product(Z, b[group_idx, :]) end +# `(expr | gr(g, by=b))` stratified random effects: independent LKJ-Cholesky + +# tau per level of `b`, so each stratum has its own full covariance structure. +# `stratum_idx[g]` maps each group-level to its stratum (walker pre-computes it +# and errors if any group straddles strata). +# L :: array[n_strata] cholesky_factor_corr[n_terms] +# tau :: array[n_strata] vector[n_terms] +# z :: array[n_groups] vector[n_terms] +# Per-group contribution: b[g, :] = (diag_pre_multiply(tau[s], L[s]) * z[g])' +# where s = stratum_idx[g]. Uses SB's `lkj_corr_cholesky_lpdfs` array-broadcast +# (one-liner added to StanBlocks.jl builtins) so L's sampling statement +# vectorizes across strata. +const ranef_correlated_by = StanBlocks.@slic begin + L ~ lkj_corr_cholesky(1.; n=n_terms, m=n_strata) + tau ~ std_normal(; n=n_terms, m=n_strata, lower=0., type=vector) + z ~ std_normal(; n=n_terms, m=n_groups, type=vector) + b = rep_matrix(0., n_groups, n_terms) + for g in 1:n_groups + b[g, :] = (diag_pre_multiply(tau[stratum_idx[g]], L[stratum_idx[g]]) * z[g])' + end + return rows_dot_product(Z, b[group_idx, :]) +end + # Treatment-coded categorical predictor. Allocates K-1 free betas; reference # level 1 contributes 0. Mirrors vimpl's `AbstractVector{<:Integer}` dispatch. # `x` is the per-row 1-based level index, `n_levels = K`. @@ -191,16 +213,20 @@ _sb_sampling!(stmts, data, key, lhs::NamedColumn, rhs) = begin end end -# Link-transformed LHS: `log(err) ~ 1 + d` etc. +# Link-transformed LHS: `log(err) ~ 1 + d`, `logit(p) ~ 1 + x`, etc. +# Sample the linear predictor on the linked scale, then invert to recover the +# response. Mirrors vimpl's `inverse(getf(lhs))` path — any link whose Julia +# `inverse` is a function with a Stan-known name (log/exp/logit/logistic/ +# sqrt/square, ...) works; unknown links error at transpile time. _sb_sampling!(stmts, data, key, lhs::ExprColumn, rhs) = begin f = getf(lhs) - f === log || error("sbimpl: only `log(x)` links supported for now (got `$f`)") - inner = getargs(lhs, 1)[1] + inner = only(getargs(lhs)) inner isa NamedColumn || error("sbimpl: expected NamedColumn inside link, got $(typeof(inner))") + inv_f = InverseFunctions.inverse(f) inner_name = name(inner) - log_name = Symbol(:log_, inner_name) - _sb_linear_predictor!(stmts, data, log_name, rhs) - push!(stmts, :($inner_name = exp($log_name))) + pre_name = Symbol(nameof(f), :_, inner_name) + _sb_linear_predictor!(stmts, data, pre_name, rhs) + push!(stmts, :($inner_name = $(Symbol(nameof(inv_f)))($pre_name))) end @@ -327,52 +353,134 @@ end # and emit one `~ ranef_correlated(...)` (K >= 2) or `~ ranef_intercept(...)` # (K == 1, intercept-only) call. Scalar-slope-only blocks (K == 1 but not an # intercept) also go through ranef_correlated -- the math degenerates gracefully. +# `(... | rhs)` -> walker-side group descriptor. Bare NamedColumn and `gr(g)` +# (no kwargs) both collapse to the inner NamedColumn (plain correlated block); +# `gr(g; by=b)` returns `(group, by)` so the emitter allocates the stratified +# `ranef_correlated_by` block. Mirrors vimpl's `_normalize_group`. +_sb_normalize_group(g::NamedColumn) = g +_sb_normalize_group(g::ExprColumn) = begin + getf(g) === gr || error("sbimpl: expected NamedColumn or `gr(...)` on RHS of `|`, got `$(getf(g))`") + args = getargs(g); kw = getkwargs(g) + length(args) == 1 || error("sbimpl: `gr(...)` expects exactly one positional group, got $(length(args))") + group = args[1] + group isa NamedColumn || error("sbimpl: `gr(...)` expects a NamedColumn group, got $(typeof(group))") + by = get(kw, :by, nothing) + by === nothing && return group + by isa NamedColumn || error("sbimpl: `gr(...; by=...)` expects a NamedColumn for `by`, got $(typeof(by))") + (group, by) +end +_sb_normalize_group(g) = error("sbimpl: expected NamedColumn or `gr(...)` on RHS of `|`, got $(typeof(g))") + +# Walker-side key used to coalesce ran_terms. Plain group -> `Symbol`; stratified +# `gr(g, by=b)` -> `(Symbol, Symbol)` so the two don't accidentally merge. +_sb_group_key(g::NamedColumn) = name(g) +_sb_group_key(g::Tuple{NamedColumn,NamedColumn}) = (name(g[1]), name(g[2])) + +# Per-group level of `g` -> stratum level of `by`. Errors if any group level +# straddles multiple strata. Ported from vimpl's `_stratum_idx`. +_sb_stratum_idx(g_idx::AbstractVector{Int}, b_idx::AbstractVector{Int}, gname, bname) = begin + m_groups = maximum(g_idx) + mapping = zeros(Int, m_groups) + for (gi, bi) in zip(g_idx, b_idx) + if mapping[gi] == 0 + mapping[gi] = bi + elseif mapping[gi] != bi + error("sbimpl: gr($gname, by=$bname): group level $gi straddles multiple strata ($(mapping[gi]) vs $bi)") + end + end + mapping +end + function _sb_emit_ranefs!(stmts, data, target::Symbol, ran_terms, summands) isempty(ran_terms) && return - # Group `(expr | g)` terms by the group symbol, preserving first-seen order. - groups = Symbol[] - by_group = Dict{Symbol, Vector{Any}}() + # Group `(expr | g)` / `(expr | gr(g, by=b))` terms by normalized group key, + # preserving first-seen order. Bare-NamedColumn and gr-by groups key + # differently so they never coalesce. + keys_seen = Any[] + by_group = Dict{Any, Vector{Any}}() + descs = Dict{Any, Any}() # key -> normalized group descriptor for rt in ran_terms - lhs, group = getargs(rt, 2) - group isa NamedColumn || error("sbimpl: expected NamedColumn on RHS of `|`, got $(typeof(group))") - g_backing = parent(group) - g_backing isa DataColumn || error("sbimpl: group `$(name(group))` must be a raw data column") - g = name(group) - haskey(by_group, g) || (push!(groups, g); by_group[g] = Any[]) + lhs, raw_group = getargs(rt, 2) + g = _sb_normalize_group(raw_group) + k = _sb_group_key(g) + haskey(by_group, k) || (push!(keys_seen, k); by_group[k] = Any[]; descs[k] = g) # Flatten `lhs` on `+` (same walker as pop terms; reuses `0` drop). - append!(by_group[g], _sb_terms(lhs)) + append!(by_group[k], _sb_terms(lhs)) end - for g in groups - gterms = by_group[g] - isempty(gterms) && error("sbimpl: ranef `(… | $g)` has no terms after dropping `0`") - # Resolve group index/count from the first occurrence (same backing data). - example = first(rt for rt in ran_terms if name(getargs(rt, 2)[2]) === g) - group_col = getargs(example, 2)[2] - n_levels, g_idx = _sb_level_index(parent(parent(group_col))) - idx_name = Symbol(g, :_idx) - n_name = Symbol(:n_, g) - data[idx_name] = g_idx - data[n_name] = n_levels - r_name = Symbol(:r_, target, :_, g) - if length(gterms) == 1 && gterms[1] === 1 - # (1 | g) fast/equivalent path -- stays on ranef_intercept so the - # emitted Stan matches existing sb.3 smoke tests bit for bit. - push!(stmts, :($r_name ~ ranef_intercept(; group_idx=$idx_name, n_groups=$n_name))) - else - col_exprs = Any[] - for t in gterms - _sb_ranef_cols!(col_exprs, data, stmts, t) - end - Z_name = Symbol(:Z_, target, :_, g) - k_name = Symbol(:n_terms_, target, :_, g) - data[k_name] = length(col_exprs) - push!(stmts, :($Z_name = $(Expr(:call, :hcat, col_exprs...)))) - push!(stmts, :($r_name ~ ranef_correlated(; - Z=$Z_name, group_idx=$idx_name, - n_groups=$n_name, n_terms=$k_name))) + for k in keys_seen + gterms = by_group[k] + desc = descs[k] + isempty(gterms) && error("sbimpl: ranef `(… | $k)` has no terms after dropping `0`") + _sb_emit_ranef_block!(stmts, data, target, desc, gterms, summands) + end +end + +# Emit a single ranef block for one normalized group descriptor. +function _sb_emit_ranef_block!(stmts, data, target::Symbol, group::NamedColumn, gterms, summands) + g_backing = parent(group) + g_backing isa DataColumn || error("sbimpl: group `$(name(group))` must be a raw data column") + g = name(group) + n_levels, g_idx = _sb_level_index(parent(g_backing)) + idx_name = Symbol(g, :_idx) + n_name = Symbol(:n_, g) + data[idx_name] = g_idx + data[n_name] = n_levels + r_name = Symbol(:r_, target, :_, g) + if length(gterms) == 1 && gterms[1] === 1 + # (1 | g) fast/equivalent path -- stays on ranef_intercept so the + # emitted Stan matches existing sb.3 smoke tests bit for bit. + push!(stmts, :($r_name ~ ranef_intercept(; group_idx=$idx_name, n_groups=$n_name))) + else + col_exprs = Any[] + for t in gterms + _sb_ranef_cols!(col_exprs, data, stmts, t) end - push!(summands, r_name) + Z_name = Symbol(:Z_, target, :_, g) + k_name = Symbol(:n_terms_, target, :_, g) + data[k_name] = length(col_exprs) + push!(stmts, :($Z_name = $(Expr(:call, :hcat, col_exprs...)))) + push!(stmts, :($r_name ~ ranef_correlated(; + Z=$Z_name, group_idx=$idx_name, + n_groups=$n_name, n_terms=$k_name))) + end + push!(summands, r_name) +end + +function _sb_emit_ranef_block!(stmts, data, target::Symbol, group::Tuple{NamedColumn,NamedColumn}, gterms, summands) + gcol, bcol = group + g_backing = parent(gcol) + b_backing = parent(bcol) + g_backing isa DataColumn || error("sbimpl: group `$(name(gcol))` must be a raw data column") + b_backing isa DataColumn || error("sbimpl: `by=$(name(bcol))` must be a raw data column") + g, b = name(gcol), name(bcol) + n_groups, g_idx = _sb_level_index(parent(g_backing)) + n_strata, b_idx = _sb_level_index(parent(b_backing)) + stratum_idx = _sb_stratum_idx(g_idx, b_idx, g, b) + # Block-local names include both `g` and `b` so this block never clashes + # with a plain `(… | g)` block against the same group column. + suffix = Symbol(g, :__by__, b) + idx_name = Symbol(suffix, :_idx) + n_name = Symbol(:n_, suffix) + s_idx_name = Symbol(suffix, :_stratum_idx) + n_strata_nm = Symbol(:n_strata_, suffix) + data[idx_name] = g_idx + data[n_name] = n_groups + data[s_idx_name] = stratum_idx + data[n_strata_nm] = n_strata + r_name = Symbol(:r_, target, :_, suffix) + col_exprs = Any[] + for t in gterms + _sb_ranef_cols!(col_exprs, data, stmts, t) end + Z_name = Symbol(:Z_, target, :_, suffix) + k_name = Symbol(:n_terms_, target, :_, suffix) + data[k_name] = length(col_exprs) + push!(stmts, :($Z_name = $(Expr(:call, :hcat, col_exprs...)))) + push!(stmts, :($r_name ~ ranef_correlated_by(; + Z=$Z_name, group_idx=$idx_name, + n_groups=$n_name, n_terms=$k_name, + stratum_idx=$s_idx_name, n_strata=$n_strata_nm))) + push!(summands, r_name) end # Extract pop terms from `1 + a + c1 [+ (...|g)]`. `0` is the standard diff --git a/web-macro/src/vimpl.jl b/web-macro/src/vimpl.jl index ab6874e..5561760 100644 --- a/web-macro/src/vimpl.jl +++ b/web-macro/src/vimpl.jl @@ -29,9 +29,25 @@ function normal end function grouped_normal end """LKJ-Cholesky owning the shared L for a grouped block.""" function chol end +"""Stratified grouped-normal: per-row L picked from `Ls[:,:,stratum_idx[i]]`. Implements brms `gr(g, by=b)` — each stratum gets its own covariance structure.""" +function grouped_normal_by end +"""Array-of-LKJ-Cholesky owning per-stratum Ls for a `grouped_normal_by` block.""" +function chol_by end """Unit simplex owning `values`; `length(values)` entries sum to 1, parameterized by `length(values)-1` unconstrained reals via stick-breaking.""" function simplex end +""" + GrGroup(group::NamedColumn, by::NamedColumn) + +Walker-side grouping tag for `gr(group, by=strata)`. `growblock!!` / +`_gc_idx` dispatch on it to allocate a `grouped_normal_by` block whose +per-row covariance is keyed off `by`. +""" +struct GrGroup + group + by +end + """ finalize(meta) / finalize(parts) / finalize(part, acc) @@ -48,6 +64,11 @@ finalize(p::Part, _) = (p,) finalize(p::Part{typeof(grouped_normal)}, _) = let n = size(p.data.values, 2), L = zeros(n, n) (Part(chol, (; L)), Part(grouped_normal, merge(p.data, (; L)))) end +finalize(p::Part{typeof(grouped_normal_by)}, _) = let + n = size(p.data.values, 2) + Ls = zeros(n, n, p.data.n_strata) + (Part(chol_by, (; Ls)), Part(grouped_normal_by, merge(p.data, (; Ls)))) +end rmerge(x::NamedTuple, y::NamedTuple) = begin xykeys = (intersect(keys(x), keys(y))...,) merge(x, y, map(rmerge, NamedTuple{xykeys}(x), NamedTuple{xykeys}(y))) @@ -67,6 +88,27 @@ vbroadcasted(;kwargs...) = (args...)->vbroadcasted(args...; kwargs...) vbroadcasted(x::NamedColumn{<:Any,<:DataColumn}; meta) = parent(meta.materialized[name(x)]) vbroadcasted(x::NamedColumn; meta) = meta.materialized[name(x)] vbroadcasted(x::ExprColumn; meta) = Base.broadcasted(getf(x), map(vbroadcasted(;meta), getargs(x))...) +# `I(expr)` is brms's literal-escape. In our DSL every call is already a first- +# class `ExprColumn` node, so `I` just needs to unwrap to its inner argument. +vbroadcasted(x::ExprColumn{typeof(I)}; meta) = vbroadcasted(only(getargs(x)); meta) +# `scale(x)` / `center(x)` / `standardize(x)` z-transform the inner column at +# VBRMI-materialization time: they materialize the inner broadcast once, apply +# the transform, and pass the resulting plain vector back up to the predictor +# pipeline. Because this fires inside `vbroadcasted`, they compose with every +# downstream consumer (pop predictor, ranefs, link transforms, ...). +vbroadcasted(x::ExprColumn{typeof(center)}; meta) = let raw = Base.materialize(vbroadcasted(only(getargs(x)); meta)) + raw .- _mean(raw) +end +vbroadcasted(x::ExprColumn{typeof(scale)}; meta) = let raw = Base.materialize(vbroadcasted(only(getargs(x)); meta)) + mu = _mean(raw); sd = _std(raw, mu) + sd > 0 || error("scale: zero variance in `$(only(getargs(x)))`") + (raw .- mu) ./ sd +end +vbroadcasted(x::ExprColumn{typeof(standardize)}; meta) = vbroadcasted(ExprColumn(scale, getargs(x)...); meta) + +_mean(xs) = sum(xs) / length(xs) +_std(xs, mu=_mean(xs)) = sqrt(sum(abs2, xs .- mu) / (length(xs) - 1)) + # Literals (Int, Float, etc.) inside formula expressions like `a^2` pass # through unchanged — they get broadcasted as scalars by Base.broadcasted. vbroadcasted(x::Number; meta) = x @@ -97,7 +139,56 @@ vmeta_sampling_rhs(meta, x::ExprColumn{typeof(*)}; kwargs...) = error("NOT IMPLE vmeta_sampling_rhs(meta, x::ExprColumn{typeof(&)}; kwargs...) = error("NOT IMPLEMENTED") vmeta_sampling_rhs(meta, x::ExprColumn{typeof(|)}; kwargs...) = begin lhs, rhs = getargs(x, 2) - vmeta_sampling_rhs(meta, lhs; group=rhs) + vmeta_sampling_rhs(meta, lhs; group=_normalize_group(rhs)) +end + +# `(a + b + ... || g)` -- zero-correlation random effects. Each term on the LHS +# gets its own 1-column block keyed by a synthetic `__nocor__N` suffix, so no +# `chol` sibling is ever allocated between them. Same total direct-parameter +# count as the correlated variant, minus the Cholesky factor parameters. +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(doublepipe)}; kwargs...) = begin + lhs, rhs = getargs(x, 2) + rhs isa NamedColumn || error("zerocorr `||`: RHS must be a NamedColumn group, got $(typeof(rhs))") + terms = lhs isa ExprColumn{typeof(+)} ? getargs(lhs) : (lhs,) + meta, args = foldl(enumerate(terms); init=(meta, ())) do (mm, aa), (i, term) + nocor_key = NamedColumn(Symbol(name(rhs), :__nocor__, i), parent(rhs)) + mm, arg = vmeta_sampling_rhs(mm, term; group=nocor_key) + mm, (aa..., arg) + end + meta, length(args) == 1 ? only(args) : Base.broadcasted(+, args...) +end + +# `a:b` -- continuous x continuous interaction. Single slot on the elementwise +# product. Categorical operands (integer/CategoricalVector) would need K-1 +# dummy expansion; errors out for now with a pointer to the 2.1 TODO. +vmeta_sampling_rhs(meta, x::ExprColumn{Colon}; group) = begin + length(getargs(x)) == 2 || error("interaction `:` expects exactly two operands, got $(length(getargs(x)))") + lhs, rhs = getargs(x, 2) + _check_cont_interaction(lhs); _check_cont_interaction(rhs) + lbc = vbroadcasted(lhs; meta); rbc = vbroadcasted(rhs; meta) + _scale_by_beta(meta, Base.broadcasted(*, lbc, rbc); group) +end +_check_cont_interaction(t) = nothing +_check_cont_interaction(t::NamedColumn) = let raw = parent(parent(t)) + raw isa AbstractVector{<:Real} && !(raw isa AbstractVector{<:Integer}) || + error("interaction `:`: only continuous x continuous is supported for now (got `$(name(t))`); see 2.1 TODO for cat expansions") +end + +# `(... | rhs)` -> walker-side group tag. Bare NamedColumn stays as-is; `gr(g)` +# with no kwargs normalizes to the inner NamedColumn (so it coalesces with +# plain `(... | g)` blocks); `gr(g; by=b)` becomes a `GrGroup` so growblock!! +# and _gc_idx can allocate the stratified structure. +_normalize_group(g) = g +_normalize_group(g::NamedColumn) = g +_normalize_group(g::ExprColumn{typeof(gr)}) = begin + args = getargs(g); kw = getkwargs(g) + length(args) == 1 || error("gr(...) expects exactly one positional group, got $(length(args))") + group = args[1] + group isa NamedColumn || error("gr(...) expects a NamedColumn group, got $(typeof(group))") + by = get(kw, :by, nothing) + by === nothing && return group + by isa NamedColumn || error("gr(...; by=...) expects a NamedColumn for `by`, got $(typeof(by))") + GrGroup(group, by) end vmeta_sampling_rhs(meta, n::Int; group) = begin n == 0 && return meta, 0 # `0` suppresses the intercept (no parameter allocated) @@ -213,6 +304,26 @@ _grow_or_push_tail(parts::Tuple, tail::Part{F}, ::F, m, n) where F = (parts, _append!(tail.data.values, m, n)) _grow_or_push_tail(parts::Tuple, _, f, m, n) = _push_part(parts, f, m, n) +# Stratified variant: Part.data carries `(values, stratum_idx, n_strata)`. +# Coalescing is valid only when the trailing part uses the same stratification +# (same stratum_idx + n_strata) — i.e. repeated `(… | gr(g, by=b))` with +# identical (g, b). Ls is allocated later by finalize once the final column +# count is known. +_push_part_by(parts::Tuple, f, m, n, stratum_idx, n_strata) = let + values = ElasticMatrix(zeros(m, n)) + (parts..., Part(f, (; values, stratum_idx, n_strata))), view(values, :, 1:n) +end +_grow_or_push_by(parts::Tuple{}, f, m, n, n_strata, stratum_idx) = + _push_part_by(parts, f, m, n, stratum_idx, n_strata) +_grow_or_push_by(parts::Tuple, f, m, n, n_strata, stratum_idx) = + _grow_or_push_tail_by(parts, last(parts), f, m, n, n_strata, stratum_idx) +_grow_or_push_tail_by(parts::Tuple, tail::Part{F}, ::F, m, n, n_strata, stratum_idx) where F = + (tail.data.n_strata == n_strata && tail.data.stratum_idx == stratum_idx) ? + (parts, _append!(tail.data.values, m, n)) : + _push_part_by(parts, f, m, n, stratum_idx, n_strata) +_grow_or_push_tail_by(parts::Tuple, _, f, m, n, n_strata, stratum_idx) = + _push_part_by(parts, f, m, n, stratum_idx, n_strata) + growblock!!(meta, ::Symbol, n) = begin parts = get(meta.blocks, :__population__, ()) parts, p = _grow_or_push(parts, normal, 1, n) @@ -224,13 +335,50 @@ growblock!!(meta, group::NamedColumn, n) = begin parts, p = _grow_or_push(parts, grouped_normal, m, n) rmerge(meta, (; blocks = (; key => parts))), p end +growblock!!(meta, group::GrGroup, n) = begin + # Stratified (gr(g, by=b)) block. Keyed distinctly from plain (_|g) so + # the two don't coalesce, even if the same `g` shows up in both forms. + m = n_levels(group.group) + n_strata = n_levels(group.by) + stratum_idx = _stratum_idx(group.group, group.by) + key = Symbol(name(group.group), :__by__, name(group.by)) + parts = get(meta.blocks, key, ()) + parts, p = _grow_or_push_by(parts, grouped_normal_by, m, n, n_strata, stratum_idx) + rmerge(meta, (; blocks = (; key => parts))), p +end + +# For each level of `group`, find the level of `by` it belongs to. Each group +# level must map to exactly one stratum (same stratum on every row of that +# group); otherwise the covariance structure is ill-defined. +_stratum_idx(group::NamedColumn, by::NamedColumn) = begin + _, g_idx = _level_index(parent(parent(group))) + _, b_idx = _level_index(parent(parent(by))) + m_groups = maximum(g_idx) + mapping = zeros(Int, m_groups) + for (gi, bi) in zip(g_idx, b_idx) + if mapping[gi] == 0 + mapping[gi] = bi + elseif mapping[gi] != bi + error("gr($(name(group)), by=$(name(by))): group level $gi straddles multiple strata ($(mapping[gi]) vs $bi)") + end + end + mapping +end + +# Extend `_re_lookup` / `_gc_idx` to the stratified group. The per-row lookup +# into `values` is still keyed by the primary group column — stratum handling +# lives inside `lprior!(::Part{grouped_normal_by})`, not in the lookup path. +n_levels(g::GrGroup) = n_levels(g.group) +_gc_idx(g::GrGroup) = _level_index(parent(parent(g.group)))[2] Base.show(io::IO, (;parent, broadcast)::MaterializedColumn) = print(io, eltype(parent), "[...] .= ", broadcast) Base.show(io::IO, (;parent, rhs)::LikelihoodColumn) = print(io, eltype(parent), "[...] .~ ", rhs) Base.show(io::IO, p::Part{typeof(normal)}) = print(io, "Part{normal}", size(p.data.values)) Base.show(io::IO, p::Part{typeof(grouped_normal)}) = print(io, "Part{grouped_normal}", size(p.data.values)) +Base.show(io::IO, p::Part{typeof(grouped_normal_by)}) = print(io, "Part{grouped_normal_by}", size(p.data.values), " x ", p.data.n_strata, " strata") Base.show(io::IO, p::Part{typeof(chol)}) = print(io, "Part{chol}(", size(p.data.L, 1), ")") +Base.show(io::IO, p::Part{typeof(chol_by)}) = print(io, "Part{chol_by}(", size(p.data.Ls, 1), " x ", size(p.data.Ls, 3), " strata)") Base.show(io::IO, p::Part{typeof(simplex)}) = print(io, "Part{simplex}(", length(p.data.values), ")") Base.show(io::IO, p::Part) = print(io, "Part{", p.func, "}") @@ -256,7 +404,9 @@ LogDensityProblems.dimension(vbrm::VBRMI) = nparams(vbrm.meta.blocks) nparams(x::Union{Tuple,NamedTuple}) = sum(nparams, x; init=0) nparams(p::Part{typeof(normal)}) = size(p.data.values, 2) nparams(p::Part{typeof(grouped_normal)}) = let (m, n) = size(p.data.values); m * n end +nparams(p::Part{typeof(grouped_normal_by)}) = let (m, n) = size(p.data.values); m * n end nparams(p::Part{typeof(chol)}) = let n = size(p.data.L, 1); n * (n + 1) ÷ 2 end +nparams(p::Part{typeof(chol_by)}) = let n = size(p.data.Ls, 1), S = size(p.data.Ls, 3); S * n * (n + 1) ÷ 2 end nparams(p::Part{typeof(simplex)}) = length(p.data.values) - 1 @inline advance!!(x, pos) = x[pos+1], pos+1 @@ -295,8 +445,12 @@ Either wrong or better LKJCholesky unconstraining + prior. Writes `p.data.L` in place from `x` (length n(n+1)/2) and returns the log-prior + Jacobian contribution. `eta` is the LKJ shape parameter. """ -@inline lprior!(p::Part{typeof(chol)}, x; eta=1.0) = begin - L = p.data.L +@inline lprior!(p::Part{typeof(chol)}, x; eta=1.0) = _chol_lprior!(p.data.L, x; eta) + +# Shared LKJ-Cholesky unconstraining body. Writes `L` in place from the first +# n(n+1)/2 entries of `x` and returns the log-prior + Jacobian contribution. +# Used by both `chol` (one shared L) and `chol_by` (per-stratum slice). +@inline _chol_lprior!(L, x; eta=1.0) = begin n = LinearAlgebra.checksquare(L) pos = 0 lprior = 0.0 @@ -327,6 +481,23 @@ contribution. `eta` is the LKJ shape parameter. lprior end +# Array-of-Cholesky: one L per stratum. Unconstrains n_strata independent LKJ +# factors, writing each slice `Ls[:, :, s]` in place. Same eta for every +# stratum (brms `gr(g, by=b)` shares the prior hyperparameters). +@inline lprior!(p::Part{typeof(chol_by)}, x; eta=1.0) = begin + (; Ls) = p.data + n = size(Ls, 1) + n_strata = size(Ls, 3) + n_per = n * (n + 1) ÷ 2 + lprior = 0.0 + pos = 0 + for s in 1:n_strata + xi, pos = advance!!(x, pos, n_per) + lprior += _chol_lprior!(view(Ls, :, :, s), xi; eta) + end + lprior +end + @inline lprior!(p::Part{typeof(grouped_normal)}, x) = begin (; values, L) = p.data _, n = size(values) @@ -340,6 +511,21 @@ end lprior end +# Stratified grouped-normal: per-row pick the stratum's Cholesky slice. +@inline lprior!(p::Part{typeof(grouped_normal_by)}, x) = begin + (; values, Ls, stratum_idx) = p.data + _, n = size(values) + lprior = 0.0 + pos = 0 + for (i, vi) in enumerate(eachrow(values)) + xi, pos = advance!!(x, pos, n) + s = stratum_idx[i] + mul!(vi, LowerTriangular(view(Ls, :, :, s)), xi) + lprior += sum(Base.Fix1(logpdf, Normal()), xi) + end + lprior +end + """ lprior!(p::Part{typeof(simplex)}, x; alpha=1.0) -> lp @@ -488,4 +674,121 @@ Base.show(io::IO, p::Part{typeof(mo1)}) = print(io, "Part{mo1}(K=", length(p.dat # ============================================================================== # End of `mo1` / `mo` template -# ============================================================================== \ No newline at end of file +# ============================================================================== + +# ============================================================================== +# User-term: `gp(x; k=K, c=C)` -- 1D Hilbert-space approx GP (squared-exp kernel) +# ------------------------------------------------------------------------------ +# Standard Riutort-Mayol et al. (2020) parameterization. Julia-side precompute: +# x_c = x .- mean(x) (center) +# L = c * maximum(abs, x_c) +# lambda[k] = (k * pi / (2L))^2 for k = 1..K (eigenvalues) +# PHI[i,k] = (1/sqrt(L)) * sin(sqrt(lambda[k]) * (x_c[i] + L)) +# +# Per-draw state (inside `lprior!`): +# log_rho, log_sigma, beta_1..beta_K <- K + 2 unconstrained params +# rho, sigma = exp(log_rho), exp(log_sigma) +# sqrt_spd[k] = sigma * (2*pi)^(1/4) * sqrt(rho) * exp(-0.25 * rho^2 * lambda[k]) +# contrib = PHI * (sqrt_spd .* beta) # length-n vector +# +# Priors: log_rho, log_sigma ~ Normal(0, 1) directly on the unconstrained scale +# (=> rho, sigma ~ LogNormal(0, 1), no Jacobian needed since we parameterize on +# log scale); beta_k ~ Normal(0, 1). Scope: population-level, data-x only. +# ============================================================================== + +""" + vmeta_sampling_rhs(meta, x::ExprColumn{typeof(offset)}; group) -> (meta, broadcasted) + +Parse an `offset(x)` term: fixed-slope (no beta) contribution to the linear +predictor. brms-compatible — adds `x` directly to the predictor without +allocating any parameters. Population-level only. +""" +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(offset)}; group) = begin + group === :__population__ || error("offset: only population-level supported (got group=$group)") + length(getargs(x)) == 1 || error("offset: expects exactly one positional argument, got $(length(getargs(x)))") + meta, vbroadcasted(getargs(x, 1)[1]; meta) +end + +"""Hilbert-space approx GP marker; population-level, 1D squared-exp kernel.""" +function hsgp end + +""" + _hsgp_basis(raw, K, c) -> (PHI, lambda) + +Precompute the Riutort-Mayol eigen-basis from raw x values. `PHI` is (n, K), +`lambda` is length-K. Center-then-scale-by-boundary-factor-c is standard. +""" +_hsgp_basis(raw::AbstractVector{<:Real}, K::Integer, c::Real) = begin + K >= 1 || error("gp: k must be >= 1 (got $K)") + c > 1 || error("gp: c must be > 1 (got $c)") + x_c = raw .- (sum(raw) / length(raw)) + L = c * maximum(abs, x_c) + L > 0 || error("gp: degenerate input (all x equal)") + lambda = [(k * pi / (2 * L))^2 for k in 1:K] + PHI = zeros(length(raw), K) + inv_sqrt_L = 1 / sqrt(L) + for k in 1:K, i in eachindex(x_c) + PHI[i, k] = inv_sqrt_L * sin(sqrt(lambda[k]) * (x_c[i] + L)) + end + PHI, lambda +end + +""" + vmeta_sampling_rhs(meta, x::ExprColumn{typeof(gp)}; group) -> (meta, broadcasted) + +Parse a `gp(x; k=K, c=C)` term. Precomputes `PHI` + `lambda` from raw data, +pushes a `Part(hsgp, …)` into the population block, and returns a broadcast +wrapper over the per-draw `contrib` buffer (refreshed in `lprior!`). +""" +vmeta_sampling_rhs(meta, x::ExprColumn{typeof(gp)}; group) = begin + group === :__population__ || error("gp: only population-level supported for now (got group=$group)") + args = getargs(x); kw = getkwargs(x) + length(args) == 1 || error("gp: expects exactly one positional argument, got $(length(args))") + inner = args[1] + inner isa NamedColumn || error("gp: expects a NamedColumn argument, got $(typeof(inner))") + raw = parent(parent(inner)) + raw isa AbstractVector{<:Real} || error("gp: `$(name(inner))` must be a real-valued vector (got $(typeof(raw)))") + K = get(kw, :k, 20) + c = get(kw, :c, 1.5) + PHI, lambda = _hsgp_basis(raw, K, c) + sqrt_spd = zeros(K) + beta = zeros(K) + contrib = zeros(length(raw)) + meta = push_parts!!(meta, :__population__, + Part(hsgp, (; PHI, lambda, sqrt_spd, beta, contrib)), + ) + meta, Base.broadcasted(identity, contrib) +end + +nparams(p::Part{typeof(hsgp)}) = length(p.data.lambda) + 2 + +""" + lprior!(p::Part{typeof(hsgp)}, x) -> lp + +Refresh `sqrt_spd` + `contrib` in place from the `K+2` unconstrained reals +(ordering: log_rho, log_sigma, beta_1..beta_K). Returns Normal(0,1) log-priors +on all three component groups (no Jacobian because we parameterize on the +log-scale directly). +""" +@inline lprior!(p::Part{typeof(hsgp)}, x) = begin + (; PHI, lambda, sqrt_spd, beta, contrib) = p.data + K = length(lambda) + pos = 0 + log_rho, pos = advance!!(x, pos) + log_sigma, pos = advance!!(x, pos) + beta_x, pos = advance!!(x, pos, K) + beta .= beta_x + rho = exp(log_rho) + sigma = exp(log_sigma) + # sqrt(S(sqrt(lambda[k]); rho, sigma)) for the 1D squared-exp spectral density. + c_base = sigma * (2 * pi)^(1/4) * sqrt(rho) + for k in 1:K + sqrt_spd[k] = c_base * exp(-0.25 * rho^2 * lambda[k]) + end + mul!(contrib, PHI, sqrt_spd .* beta) + lp = logpdf(Normal(), log_rho) + logpdf(Normal(), log_sigma) + lp += sum(Base.Fix1(logpdf, Normal()), beta_x) + lp +end + +Base.show(io::IO, p::Part{typeof(hsgp)}) = print(io, "Part{hsgp}(K=", length(p.data.lambda), ", n=", length(p.data.contrib), ")") \ No newline at end of file diff --git a/web-macro/todos/2.5-grouped-random-effects-per-factor-variance.jl b/web-macro/todos/2.5-grouped-random-effects-per-factor-variance.jl index df65ed8..0a03fa0 100644 --- a/web-macro/todos/2.5-grouped-random-effects-per-factor-variance.jl +++ b/web-macro/todos/2.5-grouped-random-effects-per-factor-variance.jl @@ -1,23 +1,41 @@ -# label: 2.5 grouped random effects (per-factor variance) +# label: 2.5 grouped random effects (gr(g, by=b)) — per-stratum covariance # tier: 2 -# status: open +# status: done (vimpl); sb backend open #= -**What it is.** Peter's "different variance by diagnosis" pattern: `(1 | subject) gr(diagnosis)` says "the random intercept by subject has a different variance per diagnosis level". In brms this is a custom group structure where the variance hyperparameter itself depends on a second factor. +**Status: done in vimpl.** `gr(g, by=b)` on the RHS of `|` allocates an +independent LKJ-Cholesky + grouped-normal block per level of `b`, so each +stratum gets its own full covariance structure (not just its own variance). -**Why it matters.** Common in clinical data where treatment groups have intrinsically different between-subject variability. Without this, you have to fit separate models per diagnosis or accept a single pooled variance. +**Semantics.** `(1 + x | gr(g, by=b))` says: random intercept + slope on `x` +across `g`-levels, but with a separate covariance matrix per `b`-level. If +`b` has 3 levels you get 3 independent (K, K) Cholesky factors (and 3 +independent tau vectors) — brms's canonical "different correlations per +diagnosis" pattern. `gr(g)` with no kwargs collapses to the plain `(... | g)` +path (same block, fully correlated shared prior). -**Implementation.** Bigger than it looks because the variance is no longer a single scalar but a length-`n_levels(diagnosis)` vector that needs its own prior and its own gradient. +**Constraint.** Every level of `g` must belong to exactly one level of `b`. +The walker (`_stratum_idx`) checks this at VBRMI time and errors if any +group-level straddles strata. -Proposed shape: -- A new `gr(group_factor)` wrapper recognized in the `~` RHS via a `function gr end` stub (already exists in `macro.jl`). -- The wrapped block stores `n_levels(group_factor)` log-scale parameters instead of one. `lprior!` walks them, multiplying each subject's random intercept by the diagnosis-specific scale. -- Requires the gc_idx for the inner factor (subject) AND for the outer factor (diagnosis) — both vectors of length N. +**Implementation sketch** (vimpl.jl): +- `GrGroup(group, by)` walker-side tag, produced by `_normalize_group`. +- New Part kinds: `grouped_normal_by` (values + stratum_idx + n_strata), + `chol_by` (Ls::Array{Float64, 3}). +- `finalize(::Part{grouped_normal_by})` allocates the (n, n, n_strata) Ls + and prepends a `chol_by` sibling sharing the tensor. +- `lprior!(::Part{chol_by})` loops strata, reuses `_chol_lprior!` per slice. +- `lprior!(::Part{grouped_normal_by})` loops rows, picks + `Ls[:, :, stratum_idx[i]]` for the per-row mul. -This composes naturally with (1.6 caching) and (2.3 per-parameter priors). +**sb backend.** See sibling TODO (TBD) — needs an SB loop over groups and +array-of-cholesky allocation, then the same `rows_dot_product` form. -**Verification.** Preset against synthetic data with two grouping factors, one nested inside the other, with intentionally different per-outer-level variance. Confirm the fitted scales recover the synthetic values. +**Verification.** Synthetic data with two grouping factors (one nested) and +intentionally different per-outer-level correlation structure. Confirm the +fitted scales + correlations recover the synthetic values. =# -loc ~ 1 + (1 | gr(g1, g2)) -y1 ~ Normal(loc, 1) +loc ~ 1 + a + (1 + a | gr(g1, by=g2)) +log(err) ~ 1 +y1 ~ Normal(loc, err) From 4a279632785b7ef9efbda49e789f560f79e6b36c Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Tue, 21 Apr 2026 14:09:12 +0200 Subject: [PATCH 16/23] todos: document new vimpl terms (I/scale/center/standardize/||/a:b/offset/gp) Updates existing TODOs to mark I(), scale/center/standardize, `||` zerocorr, and `a:b` cont x cont interactions as done, and adds new TODO examples for `offset(x)` and HSGP `gp(x; k, c)`. Co-Authored-By: Claude Opus 4.7 --- .../todos/1.3-i-expr-likely-already-works.jl | 24 ++++++---- web-macro/todos/1.4-scale-x-standardize-x.jl | 45 +++++++++++-------- ...1.5-zerocorr-independent-random-effects.jl | 41 +++++++++-------- web-macro/todos/1.8-offset-wrapper.jl | 28 ++++++++++++ web-macro/todos/2.1-interactions-a-b-a-b.jl | 38 ++++++++++------ web-macro/todos/2.8-hsgp-gp-term.jl | 41 +++++++++++++++++ 6 files changed, 158 insertions(+), 59 deletions(-) create mode 100644 web-macro/todos/1.8-offset-wrapper.jl create mode 100644 web-macro/todos/2.8-hsgp-gp-term.jl diff --git a/web-macro/todos/1.3-i-expr-likely-already-works.jl b/web-macro/todos/1.3-i-expr-likely-already-works.jl index fcf615a..95cdf31 100644 --- a/web-macro/todos/1.3-i-expr-likely-already-works.jl +++ b/web-macro/todos/1.3-i-expr-likely-already-works.jl @@ -1,16 +1,24 @@ -# label: 1.3 I(expr) — likely already works +# label: 1.3 I(expr) — brms literal-escape # tier: 1 -# status: open +# status: done (vimpl) #= -**What it is.** brms's `I()` is a literal-escape: `I(x^2)` says "compute `x^2` from the data and treat it as a single column". brms needs it because `+`, `*`, `:`, `|`, … all have special meaning inside an R formula. +**Status: done in vimpl.** Our DSL already supports arbitrary Julia function +calls on the RHS (e.g. `a^2`, `sqrt(abs(b))`, `log(exposure)`), so `I()` is +functionally unnecessary. But brms users who have internalized the `I()` +convention can now write it literally — we added a passthrough overload +(`vbroadcasted(::ExprColumn{typeof(I)}) = inner`) so `I(expr)` behaves the same +as `expr` alone, no parameter penalty. -**Why we probably don't need it.** Our DSL is parsed by Julia first, then walked by `_x`. `_x` recursively wraps every `Expr(:call, f, args...)` in an `ExprColumn`, regardless of whether `f` is special. So `loc ~ a + x^2` becomes `+(a, ^(x, 2))` → `ExprColumn(+, NamedColumn(:a), ExprColumn(^, NamedColumn(:x), 2))`. The `^` is just another function call, no special handling needed. +**Why brms needs I().** In R, `a:b`, `a*b`, `|` etc. have DSL meaning inside a +formula — `I()` escapes them so the inner expression is interpreted as plain +arithmetic. Our formula DSL is parsed by Julia first, so `a^2` is already a +plain function call and `I()` is a no-op. -The only operators that have DSL meaning in our system are `~` (sampling), `=` (assignment), and `|` / `||` inside random-effects specs. Everything else (`^`, `/`, `sqrt`, `log`, `exp`, `mod`, `min`, `max`, …) is a regular function call resolved at materialization time via `vbroadcasted`. - -**Verification.** The form below loads a model with three nonlinear terms (`a^2`, `sqrt(abs(b))`, `log(exposure)`) directly as population-level covariates. The VBRMI dim should match the number of distinct terms; the gradient sanity check should be all-active. If it works, that confirms `I()` is unnecessary because Julia function calls are first-class on the formula RHS. +**Verification.** The form below mixes the "naked" spelling (`a^2`, +`sqrt(abs(b))`, `log(exposure)`) with I()-wrapped forms — the VBRMI dim, +gradient check, and materialized predictor values must be identical. =# -loc ~ 1 + a + a^2 + sqrt(abs(b)) + log(exposure) +loc ~ 1 + a + I(a^2) + I(sqrt(abs(b))) + log(exposure) y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/1.4-scale-x-standardize-x.jl b/web-macro/todos/1.4-scale-x-standardize-x.jl index 6757944..2f6f9bc 100644 --- a/web-macro/todos/1.4-scale-x-standardize-x.jl +++ b/web-macro/todos/1.4-scale-x-standardize-x.jl @@ -1,28 +1,35 @@ -# label: 1.4 scale(x) / standardize(x) +# label: 1.4 scale(x) / center(x) / standardize(x) # tier: 1 -# status: open +# status: done (vimpl) #= -**What it is.** brms's `scale(x)` z-transforms a column at parse time: `scale(x) = (x - mean(x)) / std(x)`. The model sees the standardized column. Crucial for default priors (which are scale-invariant only after standardization) and sampler stability (well-conditioned linear predictors). +**Status: done in vimpl.** Three data-transform wrappers that z-transform (or +just center) a column at VBRMI-materialization time. brms does `scale(x)` +automatically in most vignettes for sampler stability + prior scale-invariance; +making it available as a formula-level wrapper avoids forcing users to +pre-transform their DataFrame columns. -**Why it matters.** Most brms vignettes do `scale(x)` automatically as a convenience. Without it, every formula has to either manually z-transform the data or accept poorly-scaled coefficients. +**Semantics.** +- `center(x)` -> `x - mean(x)` +- `scale(x)` -> `(x - mean(x)) / std(x)` +- `standardize(x)` -> alias for `scale(x)` -**Implementation.** -1. Add `function scale end` (and `function center end`, `function standardize end`) to `macro.jl`. -2. Add a `vmeta_sampling_rhs` overload in `vimpl.jl`: -```julia -vmeta_sampling_rhs(meta, x::ExprColumn{typeof(scale)}; group) = begin - inner = vbroadcasted(only(getargs(x)); meta) - materialized = Base.materialize(inner) - z = (materialized .- Statistics.mean(materialized)) ./ Statistics.std(materialized) - vmeta_sampling_rhs(meta, z; group) -end -``` -The standardization happens once when the BRMI is materialized into a VBRMI. Composes with the existing dense-map caching TODO. -3. Add `Statistics` to `vimpl.jl`'s using-list (or vendor `mean`/`std` inline). +Each fires inside `vbroadcasted`, so transforms compose with every downstream +consumer: pop predictor, ranef slope, link functions, interactions, etc. +`scale(a):scale(b)` works as expected (z-transformed both operands before the +elementwise product). -**Verification.** Preset: `loc ~ 1 + scale(a) + scale(b); y1 ~ Normal(loc, 1)`. Compare against the unscaled version: same dim, different posterior geometry. The fitted coefficients should be ≈ the unscaled coefficients × std(x). +**Implementation.** Three-line `vbroadcasted` overloads in vimpl.jl using +inlined `_mean`/`_std` helpers (no new dependency). The transform evaluates +once at VBRMI construction and caches the resulting vector. + +**Verification.** Same model spelled two ways: +- `loc ~ 1 + scale(a) + scale(b)` -- standardized at formula level +- `loc ~ 1 + a_z + b_z` (where a_z, b_z were pre-standardized in the DataFrame) + +VBRMI dim + log-density should be identical; fitted coefficients ≈ the +unscaled coefficients × `std(x)`. =# -loc ~ 1 + scale(a) + scale(b) +loc ~ 1 + scale(a) + center(b) y1 ~ Normal(loc, 1) diff --git a/web-macro/todos/1.5-zerocorr-independent-random-effects.jl b/web-macro/todos/1.5-zerocorr-independent-random-effects.jl index 938c869..92f493f 100644 --- a/web-macro/todos/1.5-zerocorr-independent-random-effects.jl +++ b/web-macro/todos/1.5-zerocorr-independent-random-effects.jl @@ -1,28 +1,31 @@ -# label: 1.5 zerocorr — independent random effects +# label: 1.5 zerocorr — `(terms || group)` for independent random effects # tier: 1 -# status: open +# status: done (vimpl) #= -**What it is.** brms (via lme4 syntax) lets you opt out of the LKJ correlation between multiple random terms in the same group. `(1 + x || group)` (double bar) says "estimate the random intercept and the random slope independently — don't fit a 2×2 Cholesky factor between them". Useful when there isn't enough data to estimate the correlations, or when you have prior reason to believe the terms are uncorrelated. +**Status: done in vimpl.** brms (via lme4 syntax) lets you opt out of the LKJ +correlation between multiple random terms on the same grouping factor: +`(1 + x || g)` says "fit a random intercept and a random slope, but treat them +as independent — don't estimate a 2x2 Cholesky between them". Useful when +there's not enough data to estimate the correlation, or when you have prior +reason to believe the terms are uncorrelated. -**Why it matters.** Multi-term random specs are common, and the LKJ correlation often dominates the prior cost without much identifiability. Letting users skip it is a meaningful sampling speedup and prior simplification. +**Semantics.** `(1 + a + b || g)` splits into three independent 1-column +grouped_normal blocks, each keyed distinctly from plain `(... | g)` so the +uncorrelated and correlated variants never accidentally coalesce. Per-term +1x1 Cholesky reduces to a single log_scale parameter (same `chol` machinery +used for the scalar `(1 | g)` case). -**Implementation.** Our `_x` walker already wraps `||` as `ExprColumn{typeof(doublepipe)}`. Add a `vmeta_sampling_rhs` overload that splits each term inside the `||` LHS into its own block (with a synthetic per-term key like `Symbol(group_name, :__nocor__, term_index)`): +**Implementation.** Our `_x` walker wraps `||` as `ExprColumn{typeof(doublepipe)}`. +We added a `vmeta_sampling_rhs` overload that flattens the LHS on `+`, then +foldl's over the terms with a synthetic per-term group key +(`Symbol(name(rhs), :__nocor__, i)`) so each term lands in its own block. -```julia -vmeta_sampling_rhs(meta, x::ExprColumn{typeof(doublepipe)}; kwargs...) = begin - lhs, rhs = getargs(x, 2) - terms = lhs isa ExprColumn{typeof(+)} ? getargs(lhs) : (lhs,) - foldl(enumerate(terms); init=(meta, ())) do (m, args), (i, term) - nocor_key = NamedColumn(Symbol(name(rhs), :__nocor__, i), parent(rhs)) - m, arg = vmeta_sampling_rhs(m, term; group=nocor_key) - m, (args..., arg) - end |> ((m, args),) -> (m, Base.broadcasted(+, args...)) -end -``` +**Verification.** Preset below vs. correlated counterpart +`loc ~ 1 + (1 + a | g1)`: +- correlated variant: 1 + `chol(2)` (3 params) + `grouped_normal(8, 2)` (16 params) = 20 params +- uncorrelated variant: 1 + 2 * (`chol(1)` (1 param) + `grouped_normal(8, 1)` (8 params)) = 19 params -Each per-term block ends up as 1×1 with one `log_scale` Cholesky parameter — `lprior!`'s existing single-column path handles this with no changes. - -**Verification.** Preset: `loc ~ 1 + (1 + a || g1); y1 ~ Normal(loc, 1)`. Compare its dim against the correlated `(1 + a | g1)` version: the correlated version has 3 Cholesky params (1+2/2 for a 2×2), the uncorrelated version has 2 (one log_scale per term). Same direct-parameter count (2 cols × 8 levels = 16) either way. +So the only difference is one fewer Cholesky parameter (the off-diagonal). =# diff --git a/web-macro/todos/1.8-offset-wrapper.jl b/web-macro/todos/1.8-offset-wrapper.jl new file mode 100644 index 0000000..bb20d7d --- /dev/null +++ b/web-macro/todos/1.8-offset-wrapper.jl @@ -0,0 +1,28 @@ +# label: 1.8 offset(x) wrapper — brms-compatible fixed-slope term +# tier: 1 +# status: done (vimpl) +#= +**Status: done in vimpl.** Sibling 1.2 showed the functional pattern (put +`log(exposure)` directly inside the likelihood expression). This adds the +brms-style `offset(x)` wrapper on the RHS of a linear predictor, for users +who prefer the familiar formula-language spelling. + +**Semantics.** `loc ~ 1 + a + offset(z)` allocates one population intercept +and one slope on `a`, but NO parameter for `z` — the raw `z` vector is added +directly to the predictor (fixed slope = 1). Equivalent to writing +`loc ~ 1 + a` and then manually `loc .+ z` before the likelihood, just inline. + +**Implementation sketch.** Single parser method in vimpl.jl: +`vmeta_sampling_rhs(meta, x::ExprColumn{typeof(offset)}; group)` returns +`vbroadcasted(inner; meta)` without going through `_scale_by_beta`. No Part, +no parameter slot. Pop-level only (brms restriction); group-level is unusual +and would need a walker tweak. + +**Verification.** Equivalent Poisson-with-exposure model spelled two ways: +the VBRMI dim, gradient check, and materialized log_rate values should all +match the sibling 1.2 version bit-for-bit. + +=# + +log_rate ~ 1 + a + offset(log(exposure)) +k1 ~ Poisson(exp(log_rate)) diff --git a/web-macro/todos/2.1-interactions-a-b-a-b.jl b/web-macro/todos/2.1-interactions-a-b-a-b.jl index bc026cc..a338e68 100644 --- a/web-macro/todos/2.1-interactions-a-b-a-b.jl +++ b/web-macro/todos/2.1-interactions-a-b-a-b.jl @@ -1,21 +1,33 @@ -# label: 2.1 interactions a:b, a*b +# label: 2.1 interactions a:b (cont x cont done; cat cases open) # tier: 2 -# status: open +# status: partial (cont x cont done; cont x cat / cat x cat open) #= -**What it is.** brms's `a:b` is the elementwise interaction term (a single coefficient multiplying `a[i] * b[i]`). `a*b` is the "main effects + interaction" shorthand: it desugars to `a + b + a:b`. +**Status: partial in vimpl.** Continuous-by-continuous interactions via `a:b` +are implemented -- one free beta times the elementwise product. Cont x cat and +cat x cat still error out (with a pointer to this TODO). -**Why it matters.** Interactions are the most commonly missed feature in regression DSLs. Without them, every model that needs `a:b` has to manually create the interaction column in the input DataFrame. +**Julia syntax quirk.** In Julia, `a:b` parses as `Expr(:call, :(:), :a, :b)` +-- i.e. the `Colon()` function applied to `a` and `b`. At BRMI-walk time it +lands as `ExprColumn{Colon}(a, b)`. Precedence is tighter than `+` but looser +than `*`, so `1 + a + a:b` behaves correctly. `a*b` / `*`-interaction sugar is +left out for now because `*` already means scalar/array multiplication in most +Julia contexts -- keeping its surface meaning avoids DSL ambiguity. -**Implementation.** -1. **Parser side.** Add a `:` case to `_x` so that `a:b` becomes `ExprColumn(:, NamedColumn(:a), NamedColumn(:b))` instead of falling through to a Symbol/Range parse error. -2. **Materialization side.** Add `vmeta_sampling_rhs(meta, x::ExprColumn{typeof(:)}; group)` that elementwise-multiplies the operands and dispatches to the float-vector path. For continuous × continuous it's a single coefficient on `a .* b`; for categorical × continuous it's `(k-1)` coefficients (one per non-reference level of the categorical, multiplied by the continuous); for categorical × categorical it's `(k₁-1)*(k₂-1)` coefficients via a 2D `_cat_lookup`. -3. **`a*b` desugaring.** At parse time in `_x`, rewrite `*` between formula terms as `+(a, b, :(a:b))`. This needs care because `*` also means multiplication elsewhere (e.g. `Normal(0, 2*sigma)`); the rewrite should only apply at formula-RHS top-level. +**Implementation.** `vmeta_sampling_rhs(meta, x::ExprColumn{Colon}; group)` +broadcasts the two operands, multiplies them elementwise, and hands off to +`_scale_by_beta` like any other continuous predictor. A small +`_check_cont_interaction` guard errors out if either operand is an integer / +categorical column. -**Verification.** Presets exercising each interaction type: -- continuous×continuous: `loc ~ 1 + a + b + a:b; y1 ~ Normal(loc, 1)` → dim 4 -- continuous×categorical: `loc ~ 1 + a + c1 + a:c1; y1 ~ Normal(loc, 1)` → dim 6 (1 + 1 + 2 + 2) -- categorical×categorical: `loc ~ 1 + c1 + c2 + c1:c2; y1 ~ Normal(loc, 1)` → dim 5 (1 + 2 + 1 + 2) -- shorthand: `loc ~ 1 + a*b; y1 ~ Normal(loc, 1)` should match `loc ~ 1 + a + b + a:b` exactly. +**Open (cat expansions).** For cont x cat, allocate `(K-1)` slots multiplied +by the treatment dummies. For cat x cat, allocate `(K1-1)*(K2-1)` slots via a +2D level-index lookup. Both reuse the existing `_cat_broadcast` / `_cat_lookup` +machinery but need a slightly different indexing path. + +**Verification presets.** +- cont x cont (working): `loc ~ 1 + a + b + a:b; y1 ~ Normal(loc, 1)` -> dim 4 +- cont x cat (open): `loc ~ 1 + a + c1 + a:c1; y1 ~ Normal(loc, 1)` -> dim 6 +- cat x cat (open): `loc ~ 1 + c1 + c2 + c1:c2; y1 ~ Normal(loc, 1)` -> dim 5 =# diff --git a/web-macro/todos/2.8-hsgp-gp-term.jl b/web-macro/todos/2.8-hsgp-gp-term.jl new file mode 100644 index 0000000..40c3ac1 --- /dev/null +++ b/web-macro/todos/2.8-hsgp-gp-term.jl @@ -0,0 +1,41 @@ +# label: 2.8 HSGP gp(x; k, c) — Hilbert-space approx Gaussian process +# tier: 2 +# status: done (vimpl); sb backend open +#= +**Status: done in vimpl.** `gp(x; k=K, c=C)` adds a 1D Hilbert-space approx GP +term to the linear predictor (Riutort-Mayol et al. 2020). Squared-exponential +kernel, population-level, data-x only. + +**Semantics.** `loc ~ 1 + a + gp(b; k=20, c=1.5)` adds a smooth function of `t` +to the predictor without committing to a functional form. Parameters (per gp() +call): log_rho (lengthscale), log_sigma (marginal SD), and K latent betas — +total K+2 unconstrained reals. Priors are Normal(0,1) on all three groups +(=> rho, sigma ~ LogNormal(0,1); beta_k ~ Normal(0,1)). + +**Implementation sketch** (vimpl.jl): +- `_hsgp_basis(raw, K, c)` precomputes `PHI::(n,K)` and `lambda::K` from + centered x (L = c * maximum(abs, x_centered); lambda[k] = (k*pi/(2L))^2; + PHI[i,k] = (1/sqrt(L)) * sin(sqrt(lambda[k]) * (x_centered[i] + L))). +- New Part kind `hsgp` with `(; PHI, lambda, sqrt_spd, beta, contrib)`. +- `lprior!(::Part{hsgp}, x)` reads K+2 params, refreshes + `sqrt_spd[k] = sigma * (2pi)^(1/4) * sqrt(rho) * exp(-0.25 * rho^2 * lambda[k])` + and `contrib = PHI * (sqrt_spd .* beta)`, returns Normal(0,1) log-priors. +- Parser returns `Base.broadcasted(identity, contrib)` so the linear predictor + just adds the per-draw `contrib` vector in. + +**Parameter-x extension.** If `x` itself is a Stan parameter (measurement-error +or latent predictor models), PHI must be rebuilt per draw — that lands as a +separate code path, not this TODO. Data-x covers ~all real uses. + +**sb backend.** Emit the same `_hsgp_basis` precompute into SB data, plus an +`_sb_hsgp` submodel that does the sqrt_spd matvec inside Stan. Straightforward +once array-indexed basis expansions land. + +**Verification.** Synthetic sinusoidal f(t) + noise. Confirm the fitted GP +recovers the curve at reasonable K (>= 15) and that the log-prior / gradient +stay finite across the parameter range. + +=# + +loc ~ 1 + a + gp(b; k=20, c=1.5) +y1 ~ Normal(loc, 1) From ea982af52c716921aa22c3c0895f2e25c3b7bc07 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Wed, 22 Apr 2026 11:17:43 +0200 Subject: [PATCH 17/23] web-macro: baseline before BRMMacroWeb.jl walk-through refactor Staged current state of the web module so diffs against the upcoming review pass are easy to read: - rename todos/ -> examples/ (status/formula edits already written back) - add brm-macro.css + html_expr.jl (styled BRMI/VBRMI cards) - Project.toml: BridgeStan + WarmupHMC direct deps - BRMMacroWeb.jl: AppContext + polling_fetchindex + sbimpl stages Co-Authored-By: Claude Opus 4.7 --- .gitignore | 4 +- web-macro/Project.toml | 2 + .../1.1-verify-bernoulli-binomial-done.jl | 0 ...xposure-already-works-without-a-wrapper.jl | 0 .../1.3-i-expr-likely-already-works.jl | 0 .../1.4-scale-x-standardize-x.jl | 0 ...1.5-zerocorr-independent-random-effects.jl | 0 ...1.6-cache-levels-level_map-dense-gc_idx.jl | 0 ...egoricalarrays-pooledarrays-integration.jl | 0 .../{todos => examples}/1.8-offset-wrapper.jl | 0 .../2.1-interactions-a-b-a-b.jl | 0 ...onfigurable-categorical-reference-level.jl | 0 .../2.3-per-parameter-prior-scales.jl | 0 ...ed-non-centered-parameterization-toggle.jl | 0 ...uped-random-effects-per-factor-variance.jl | 0 .../2.6-multi-membership-random-effects-mm.jl | 0 ...r-meta-analysis-and-weighted-regression.jl | 0 .../{todos => examples}/2.8-hsgp-gp-term.jl | 0 .../3.1-multivariate-outcomes-cbind-y1-y2.jl | 0 ...dirichlet-process-non-parametric-models.jl | 0 .../3.11-zero-inflated-hurdle-likelihoods.jl | 0 ...nferred-predictors-measurement-error-me.jl | 0 ...ordinal-predictors-mo-monotonic-effects.jl | 0 .../3.4-ordinal-outcomes-proportional-odds.jl | 0 .../{todos => examples}/3.5-mixture-models.jl | 0 .../3.6-splines-gp-submodels-s-bs-gp-t2.jl | 0 .../3.7-autoregressive-submodels-ar-ar1.jl | 0 .../3.8-decompositions-qr-orthogonal-polar.jl | 0 .../3.9-spike-and-slab-horseshoe-priors.jl | 0 .../sb.1-linear-regression.jl | 0 .../sb.2-distributional.jl | 0 .../sb.3-random-effects.jl | 0 .../sb.4-categorical-predictors.jl | 0 .../sb.5-non-normal-likelihoods.jl | 0 .../sb.6-submodels-mo-mm.jl | 0 web-macro/src/BRMMacroWeb.jl | 1844 ++++++----------- web-macro/src/brm-macro.css | 56 + web-macro/src/html_expr.jl | 176 ++ 38 files changed, 864 insertions(+), 1218 deletions(-) rename web-macro/{todos => examples}/1.1-verify-bernoulli-binomial-done.jl (100%) rename web-macro/{todos => examples}/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl (100%) rename web-macro/{todos => examples}/1.3-i-expr-likely-already-works.jl (100%) rename web-macro/{todos => examples}/1.4-scale-x-standardize-x.jl (100%) rename web-macro/{todos => examples}/1.5-zerocorr-independent-random-effects.jl (100%) rename web-macro/{todos => examples}/1.6-cache-levels-level_map-dense-gc_idx.jl (100%) rename web-macro/{todos => examples}/1.7-categoricalarrays-pooledarrays-integration.jl (100%) rename web-macro/{todos => examples}/1.8-offset-wrapper.jl (100%) rename web-macro/{todos => examples}/2.1-interactions-a-b-a-b.jl (100%) rename web-macro/{todos => examples}/2.2-configurable-categorical-reference-level.jl (100%) rename web-macro/{todos => examples}/2.3-per-parameter-prior-scales.jl (100%) rename web-macro/{todos => examples}/2.4-centered-non-centered-parameterization-toggle.jl (100%) rename web-macro/{todos => examples}/2.5-grouped-random-effects-per-factor-variance.jl (100%) rename web-macro/{todos => examples}/2.6-multi-membership-random-effects-mm.jl (100%) rename web-macro/{todos => examples}/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl (100%) rename web-macro/{todos => examples}/2.8-hsgp-gp-term.jl (100%) rename web-macro/{todos => examples}/3.1-multivariate-outcomes-cbind-y1-y2.jl (100%) rename web-macro/{todos => examples}/3.10-dirichlet-process-non-parametric-models.jl (100%) rename web-macro/{todos => examples}/3.11-zero-inflated-hurdle-likelihoods.jl (100%) rename web-macro/{todos => examples}/3.2-inferred-predictors-measurement-error-me.jl (100%) rename web-macro/{todos => examples}/3.3-ordinal-predictors-mo-monotonic-effects.jl (100%) rename web-macro/{todos => examples}/3.4-ordinal-outcomes-proportional-odds.jl (100%) rename web-macro/{todos => examples}/3.5-mixture-models.jl (100%) rename web-macro/{todos => examples}/3.6-splines-gp-submodels-s-bs-gp-t2.jl (100%) rename web-macro/{todos => examples}/3.7-autoregressive-submodels-ar-ar1.jl (100%) rename web-macro/{todos => examples}/3.8-decompositions-qr-orthogonal-polar.jl (100%) rename web-macro/{todos => examples}/3.9-spike-and-slab-horseshoe-priors.jl (100%) rename web-macro/{todos => examples}/sb.1-linear-regression.jl (100%) rename web-macro/{todos => examples}/sb.2-distributional.jl (100%) rename web-macro/{todos => examples}/sb.3-random-effects.jl (100%) rename web-macro/{todos => examples}/sb.4-categorical-predictors.jl (100%) rename web-macro/{todos => examples}/sb.5-non-normal-likelihoods.jl (100%) rename web-macro/{todos => examples}/sb.6-submodels-mo-mm.jl (100%) create mode 100644 web-macro/src/brm-macro.css create mode 100644 web-macro/src/html_expr.jl diff --git a/.gitignore b/.gitignore index b076ed1..789a4f9 100644 --- a/.gitignore +++ b/.gitignore @@ -28,6 +28,6 @@ Manifest*.toml LocalPreferences.toml JuliaLocalPreferences.toml -# Confidential client-project TODOs (local-only, never commit) -web-macro/todos/bruno-*.jl +# Confidential client-project examples (local-only, never commit) +web-macro/examples/bruno-*.jl web-macro/src/bruno-*.jl diff --git a/web-macro/Project.toml b/web-macro/Project.toml index 71df3e7..185615d 100644 --- a/web-macro/Project.toml +++ b/web-macro/Project.toml @@ -3,6 +3,7 @@ uuid = "cd55decf-6be8-4ea5-907a-b2dc35e4cc14" version = "0.1.0" [deps] +BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a" CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" @@ -27,3 +28,4 @@ StanLogDensityProblems = "a545de4d-8dba-46db-9d34-4e41d3f07807" TestModules = "63c02187-99fd-4e5c-aaf0-4d6bfebc181c" Treebars = "e1e568c4-3a56-40a4-95fa-9b9c6c16fccb" Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" +WarmupHMC = "60658175-6863-4866-a322-ab51a11c0cfe" diff --git a/web-macro/todos/1.1-verify-bernoulli-binomial-done.jl b/web-macro/examples/1.1-verify-bernoulli-binomial-done.jl similarity index 100% rename from web-macro/todos/1.1-verify-bernoulli-binomial-done.jl rename to web-macro/examples/1.1-verify-bernoulli-binomial-done.jl diff --git a/web-macro/todos/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl b/web-macro/examples/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl similarity index 100% rename from web-macro/todos/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl rename to web-macro/examples/1.2-offset-fixed-exposure-already-works-without-a-wrapper.jl diff --git a/web-macro/todos/1.3-i-expr-likely-already-works.jl b/web-macro/examples/1.3-i-expr-likely-already-works.jl similarity index 100% rename from web-macro/todos/1.3-i-expr-likely-already-works.jl rename to web-macro/examples/1.3-i-expr-likely-already-works.jl diff --git a/web-macro/todos/1.4-scale-x-standardize-x.jl b/web-macro/examples/1.4-scale-x-standardize-x.jl similarity index 100% rename from web-macro/todos/1.4-scale-x-standardize-x.jl rename to web-macro/examples/1.4-scale-x-standardize-x.jl diff --git a/web-macro/todos/1.5-zerocorr-independent-random-effects.jl b/web-macro/examples/1.5-zerocorr-independent-random-effects.jl similarity index 100% rename from web-macro/todos/1.5-zerocorr-independent-random-effects.jl rename to web-macro/examples/1.5-zerocorr-independent-random-effects.jl diff --git a/web-macro/todos/1.6-cache-levels-level_map-dense-gc_idx.jl b/web-macro/examples/1.6-cache-levels-level_map-dense-gc_idx.jl similarity index 100% rename from web-macro/todos/1.6-cache-levels-level_map-dense-gc_idx.jl rename to web-macro/examples/1.6-cache-levels-level_map-dense-gc_idx.jl diff --git a/web-macro/todos/1.7-categoricalarrays-pooledarrays-integration.jl b/web-macro/examples/1.7-categoricalarrays-pooledarrays-integration.jl similarity index 100% rename from web-macro/todos/1.7-categoricalarrays-pooledarrays-integration.jl rename to web-macro/examples/1.7-categoricalarrays-pooledarrays-integration.jl diff --git a/web-macro/todos/1.8-offset-wrapper.jl b/web-macro/examples/1.8-offset-wrapper.jl similarity index 100% rename from web-macro/todos/1.8-offset-wrapper.jl rename to web-macro/examples/1.8-offset-wrapper.jl diff --git a/web-macro/todos/2.1-interactions-a-b-a-b.jl b/web-macro/examples/2.1-interactions-a-b-a-b.jl similarity index 100% rename from web-macro/todos/2.1-interactions-a-b-a-b.jl rename to web-macro/examples/2.1-interactions-a-b-a-b.jl diff --git a/web-macro/todos/2.2-configurable-categorical-reference-level.jl b/web-macro/examples/2.2-configurable-categorical-reference-level.jl similarity index 100% rename from web-macro/todos/2.2-configurable-categorical-reference-level.jl rename to web-macro/examples/2.2-configurable-categorical-reference-level.jl diff --git a/web-macro/todos/2.3-per-parameter-prior-scales.jl b/web-macro/examples/2.3-per-parameter-prior-scales.jl similarity index 100% rename from web-macro/todos/2.3-per-parameter-prior-scales.jl rename to web-macro/examples/2.3-per-parameter-prior-scales.jl diff --git a/web-macro/todos/2.4-centered-non-centered-parameterization-toggle.jl b/web-macro/examples/2.4-centered-non-centered-parameterization-toggle.jl similarity index 100% rename from web-macro/todos/2.4-centered-non-centered-parameterization-toggle.jl rename to web-macro/examples/2.4-centered-non-centered-parameterization-toggle.jl diff --git a/web-macro/todos/2.5-grouped-random-effects-per-factor-variance.jl b/web-macro/examples/2.5-grouped-random-effects-per-factor-variance.jl similarity index 100% rename from web-macro/todos/2.5-grouped-random-effects-per-factor-variance.jl rename to web-macro/examples/2.5-grouped-random-effects-per-factor-variance.jl diff --git a/web-macro/todos/2.6-multi-membership-random-effects-mm.jl b/web-macro/examples/2.6-multi-membership-random-effects-mm.jl similarity index 100% rename from web-macro/todos/2.6-multi-membership-random-effects-mm.jl rename to web-macro/examples/2.6-multi-membership-random-effects-mm.jl diff --git a/web-macro/todos/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl b/web-macro/examples/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl similarity index 100% rename from web-macro/todos/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl rename to web-macro/examples/2.7-se-weights-for-meta-analysis-and-weighted-regression.jl diff --git a/web-macro/todos/2.8-hsgp-gp-term.jl b/web-macro/examples/2.8-hsgp-gp-term.jl similarity index 100% rename from web-macro/todos/2.8-hsgp-gp-term.jl rename to web-macro/examples/2.8-hsgp-gp-term.jl diff --git a/web-macro/todos/3.1-multivariate-outcomes-cbind-y1-y2.jl b/web-macro/examples/3.1-multivariate-outcomes-cbind-y1-y2.jl similarity index 100% rename from web-macro/todos/3.1-multivariate-outcomes-cbind-y1-y2.jl rename to web-macro/examples/3.1-multivariate-outcomes-cbind-y1-y2.jl diff --git a/web-macro/todos/3.10-dirichlet-process-non-parametric-models.jl b/web-macro/examples/3.10-dirichlet-process-non-parametric-models.jl similarity index 100% rename from web-macro/todos/3.10-dirichlet-process-non-parametric-models.jl rename to web-macro/examples/3.10-dirichlet-process-non-parametric-models.jl diff --git a/web-macro/todos/3.11-zero-inflated-hurdle-likelihoods.jl b/web-macro/examples/3.11-zero-inflated-hurdle-likelihoods.jl similarity index 100% rename from web-macro/todos/3.11-zero-inflated-hurdle-likelihoods.jl rename to web-macro/examples/3.11-zero-inflated-hurdle-likelihoods.jl diff --git a/web-macro/todos/3.2-inferred-predictors-measurement-error-me.jl b/web-macro/examples/3.2-inferred-predictors-measurement-error-me.jl similarity index 100% rename from web-macro/todos/3.2-inferred-predictors-measurement-error-me.jl rename to web-macro/examples/3.2-inferred-predictors-measurement-error-me.jl diff --git a/web-macro/todos/3.3-ordinal-predictors-mo-monotonic-effects.jl b/web-macro/examples/3.3-ordinal-predictors-mo-monotonic-effects.jl similarity index 100% rename from web-macro/todos/3.3-ordinal-predictors-mo-monotonic-effects.jl rename to web-macro/examples/3.3-ordinal-predictors-mo-monotonic-effects.jl diff --git a/web-macro/todos/3.4-ordinal-outcomes-proportional-odds.jl b/web-macro/examples/3.4-ordinal-outcomes-proportional-odds.jl similarity index 100% rename from web-macro/todos/3.4-ordinal-outcomes-proportional-odds.jl rename to web-macro/examples/3.4-ordinal-outcomes-proportional-odds.jl diff --git a/web-macro/todos/3.5-mixture-models.jl b/web-macro/examples/3.5-mixture-models.jl similarity index 100% rename from web-macro/todos/3.5-mixture-models.jl rename to web-macro/examples/3.5-mixture-models.jl diff --git a/web-macro/todos/3.6-splines-gp-submodels-s-bs-gp-t2.jl b/web-macro/examples/3.6-splines-gp-submodels-s-bs-gp-t2.jl similarity index 100% rename from web-macro/todos/3.6-splines-gp-submodels-s-bs-gp-t2.jl rename to web-macro/examples/3.6-splines-gp-submodels-s-bs-gp-t2.jl diff --git a/web-macro/todos/3.7-autoregressive-submodels-ar-ar1.jl b/web-macro/examples/3.7-autoregressive-submodels-ar-ar1.jl similarity index 100% rename from web-macro/todos/3.7-autoregressive-submodels-ar-ar1.jl rename to web-macro/examples/3.7-autoregressive-submodels-ar-ar1.jl diff --git a/web-macro/todos/3.8-decompositions-qr-orthogonal-polar.jl b/web-macro/examples/3.8-decompositions-qr-orthogonal-polar.jl similarity index 100% rename from web-macro/todos/3.8-decompositions-qr-orthogonal-polar.jl rename to web-macro/examples/3.8-decompositions-qr-orthogonal-polar.jl diff --git a/web-macro/todos/3.9-spike-and-slab-horseshoe-priors.jl b/web-macro/examples/3.9-spike-and-slab-horseshoe-priors.jl similarity index 100% rename from web-macro/todos/3.9-spike-and-slab-horseshoe-priors.jl rename to web-macro/examples/3.9-spike-and-slab-horseshoe-priors.jl diff --git a/web-macro/todos/sb.1-linear-regression.jl b/web-macro/examples/sb.1-linear-regression.jl similarity index 100% rename from web-macro/todos/sb.1-linear-regression.jl rename to web-macro/examples/sb.1-linear-regression.jl diff --git a/web-macro/todos/sb.2-distributional.jl b/web-macro/examples/sb.2-distributional.jl similarity index 100% rename from web-macro/todos/sb.2-distributional.jl rename to web-macro/examples/sb.2-distributional.jl diff --git a/web-macro/todos/sb.3-random-effects.jl b/web-macro/examples/sb.3-random-effects.jl similarity index 100% rename from web-macro/todos/sb.3-random-effects.jl rename to web-macro/examples/sb.3-random-effects.jl diff --git a/web-macro/todos/sb.4-categorical-predictors.jl b/web-macro/examples/sb.4-categorical-predictors.jl similarity index 100% rename from web-macro/todos/sb.4-categorical-predictors.jl rename to web-macro/examples/sb.4-categorical-predictors.jl diff --git a/web-macro/todos/sb.5-non-normal-likelihoods.jl b/web-macro/examples/sb.5-non-normal-likelihoods.jl similarity index 100% rename from web-macro/todos/sb.5-non-normal-likelihoods.jl rename to web-macro/examples/sb.5-non-normal-likelihoods.jl diff --git a/web-macro/todos/sb.6-submodels-mo-mm.jl b/web-macro/examples/sb.6-submodels-mo-mm.jl similarity index 100% rename from web-macro/todos/sb.6-submodels-mo-mm.jl rename to web-macro/examples/sb.6-submodels-mo-mm.jl diff --git a/web-macro/src/BRMMacroWeb.jl b/web-macro/src/BRMMacroWeb.jl index 416c2d3..93c2a68 100644 --- a/web-macro/src/BRMMacroWeb.jl +++ b/web-macro/src/BRMMacroWeb.jl @@ -6,594 +6,534 @@ using Random using Chairmarks using DataFrames using FiniteDifferences: FiniteDifferences, central_fdm +using BridgeStan: BridgeStan -# The @brm macro and the VBRMI implementation live alongside this module so -# Revise tracks them. The scripts/ entry points (parsing.jl, Benchmarking, -# StanBlocksImpl) include them via relative paths into here. +# The @brm macro and the VBRMI / SBBRMI implementations live alongside this +# module so Revise tracks them. The scripts/ entry points (parsing.jl, +# Benchmarking, StanBlocksImpl) include them via relative paths into here. include("macro.jl") include("vimpl.jl") include("sbimpl.jl") - -# ── Default formula + synthetic data ──────────────────────────────────────── - -default_formula() = """loc1 ~ 1 + a + c1 + (1 + b + c1 | g1) + (1 | g2) -log(err1) ~ 1 + d -y1 ~ Normal(loc1, err1) - -log_rate ~ 1 + a + (1 | g3) -k1 ~ Poisson(exp(log_rate)) - -log_odds_bin ~ 1 + c2 + (1 | g2) -bin_succ ~ Binomial(bin_n, logistic(log_odds_bin)) - -log_odds_b ~ 1 + b -bin_y ~ Bernoulli(logistic(log_odds_b)) -""" - -# Synthetic data shaped roughly like the benchmarking example: continuous -# covariates plus a categorical `group` column. vimpl.jl currently materializes -# population-level terms only, so the default formula stays scalar — but the -# `group` column is kept around so users can experiment with `(1 | group)` once -# vimpl.jl supports it. -function synthetic_df(; n=64, seed=1) - rng = Xoshiro(seed) - # Continuous covariates - a = randn(rng, n) - b = randn(rng, n) - c = randn(rng, n) - d = randn(rng, n) - # Grouping factors with different numbers of levels — use these on the - # right-hand side of `(... | gN)` to test multiple random-effects blocks. - g1 = repeat(1:8, inner=cld(n, 8))[1:n] # 8 levels - g2 = repeat(1:4, inner=cld(n, 4))[1:n] # 4 levels - g3 = rand(rng, 1:6, n) # 6 levels, unordered - # Categorical predictors (treatment-coded) with small level counts. - c1 = rand(rng, 1:3, n) - c2 = rand(rng, 1:2, n) - c3 = rand(rng, 1:4, n) - # Positive exposure column — for Poisson-with-offset patterns where the - # log-rate gets a row-specific shift `+ log(exposure)` inside the - # likelihood expression directly (no `offset(...)` wrapper needed). - exposure = 0.5 .+ rand(rng, n) - # Continuous outcomes - eta1 = 0.5 .+ 1.2 .* a .- 0.7 .* b .+ 0.3 .* c .+ 0.1 .* d - y1 = eta1 .+ 0.3 .* randn(rng, n) - y2 = -0.2 .+ 0.6 .* a .+ 0.4 .* b .+ 0.2 .* randn(rng, n) - # Non-negative integer (count) outcomes — Poisson likelihoods - k1 = rand.(rng, Distributions.Poisson.(exp.(0.5 .* eta1))) - k2 = rand.(rng, Distributions.Poisson.(exp.(0.3 .+ 0.4 .* a))) - # Binomial-likelihood pair: variable trial counts plus successes - # (the `trials(size)` brms sidecar is unnecessary in our DSL — `size` - # is just another positional argument to `Binomial`). - bin_n = rand(rng, 5:30, n) - bin_p_true = @. 1 / (1 + exp(-(0.2 + 0.5 * a))) - bin_succ = [rand(rng, Distributions.Binomial(n_i, p_i)) - for (n_i, p_i) in zip(bin_n, bin_p_true)] - # Bernoulli-likelihood column (0/1) — for hierarchical-Bernoulli models - # like Kruschke's `therapeutic_touch`. - bin_y = [rand(rng, Distributions.Bernoulli(p_i)) ? 1 : 0 for p_i in bin_p_true] - DataFrame(; a, b, c, d, g1, g2, g3, c1, c2, c3, exposure, - y1, y2, k1, k2, bin_n, bin_succ, bin_y) -end +# Styled HTML rendering for BRMI / VBRMI cards. Lives in its own file +# because it will eventually move to an ext of the main package. +include("html_expr.jl") # Extension hook: extensions (e.g. the gitignored `bruno-ext.jl`) that need # to contribute auxiliary data which doesn't fit as per-row DataFrame # columns -- for example `dose_times::Vector{<:AbstractVector}` indexed by # subject id -- add a method `dataset_extras(::Val{:ns}, df)` returning a -# NamedTuple of extras. The namespace is derived from the TODO label/slug -# (first dash/space-separated segment), so `bruno-qt-*` TODOs dispatch to -# `::Val{:bruno}`. Default is no extras. +# NamedTuple of extras. The namespace is derived from the example +# label/slug (first dash/space-separated segment), so `bruno-qt-*` examples +# dispatch to `::Val{:bruno}`. Default is no extras. dataset_extras(::Val, df) = (;) -# Namespace extracted from a TODO label or slug; empty -> :default. -dataset_namespace(label::AbstractString) = isempty(strip(label)) ? :default : - Symbol(lowercase(first(split(strip(label), r"[\s:\-]+")))) - -# Build the container passed as `df` to `_brm`: a NamedTuple merging the -# synthetic-df columns with extras dispatched by namespace. `@brm` only -# needs `hasproperty`/`getproperty` on its data argument, so a NamedTuple -# works as a drop-in for the DataFrame and lets extensions splice in -# extras without touching macro.jl. -function dataset_container(df, ns::Symbol=:default) - cols = (; (Symbol(c) => df[!, c] for c in names(df))...) - merge(cols, dataset_extras(Val(ns), df)) +# DOs in dependency order. Every feature is a focused @dynamicstruct: the +# pipeline stages live as derived properties on `BRMRun`, example file I/O +# on `ExampleStore`, synthetic + namespace-merged data on `Dataset`, and +# AST parse/safety/transform on `Formula`. `AppData` is just a thin holder +# of sub-DOs plus the `pipeline_run` polling entry point; `AppContext` is +# the HTMX routes and rendering layer. +struct FormulaSecurityError <: Exception + msg::String end - -# ── Formula safety whitelist ──────────────────────────────────────────────── -# -# The formula textarea accepts arbitrary text that gets `Meta.parse`d then -# `eval`'d. To prevent arbitrary code execution when the web app is shared, -# we walk the parsed AST *before* any eval and reject any function call or -# expression type that isn't on the allowlist. The allowlist is deliberately -# generous for formula-writing (math, distributions, data-column references) -# but blocks I/O, shell, eval, include, ccall, macros, etc. +Base.showerror(io::IO, e::FormulaSecurityError) = print(io, "FormulaSecurityError: ", e.msg) _ALLOWED_CALLS = Set{Symbol}([ - # DSL operators :~, :(+), :(-), :(*), :(/), :(^), :(|), :(||), - # Comparison (may appear in ifelse-style expressions) :(==), :(!=), :(<), :(>), :(<=), :(>=), - # Math :log, :log2, :log10, :log1p, :exp, :exp2, :expm1, :sqrt, :cbrt, :abs, :abs2, :sign, :floor, :ceil, :round, :sin, :cos, :tan, :asin, :acos, :atan, :min, :max, :clamp, :mod, :rem, :div, :logistic, :logit, :softmax, :logsumexp, :log_abs_tanh, :log_square_tanh, - # Distributions (Type constructors — the pass-through handles these) :Normal, :Poisson, :Binomial, :Bernoulli, :BernoulliLogit, :Beta, :Gamma, :Exponential, :Cauchy, :StudentT, :LogNormal, :Weibull, :NegativeBinomial, :Geometric, :Laplace, :Uniform, :MvNormal, :MixtureModel, :Dirichlet, :InverseGamma, :InverseGaussian, :VonMises, :Pareto, :OrderedLogistic, :Categorical, - # TODO stubs (not yet implemented but syntactically valid) :scale, :center, :standardize, :factor, :offset, :s, :bs, :t2, :gp, :ar, :ar1, :mo, :mo1, :cbind, :mvbind, :mm, :gr, :dp, :me, :centered, :Horseshoe, :ZeroInflatedPoisson, :weighted, - # Data helpers :length, :unique, :sort, :size, :eltype, :nrow, :ncol, ]) -# Expression heads that are safe in a formula AST (literals, blocks, calls, …) _SAFE_HEADS = Set{Symbol}([ :block, :call, :., :(=), :(||), :tuple, :vect, :ref, :kw, :parameters, :(...), - # Comparison chains :comparison, :&&, ]) - -struct FormulaSecurityError <: Exception - msg::String -end -Base.showerror(io::IO, e::FormulaSecurityError) = print(io, "FormulaSecurityError: ", e.msg) - -function _check_formula_safety!(x) - # Literals, symbols, line numbers — always safe - x isa Union{Number, AbstractString, Symbol, LineNumberNode, Nothing, Bool, QuoteNode} && return - x isa Expr || return - - # Reject dangerous expression types outright - if x.head == :macrocall - throw(FormulaSecurityError("macro calls are not allowed in formulas (got $(x.args[1]))")) - elseif x.head in (:cmd, :string) - throw(FormulaSecurityError("`\$(x.head)` expressions are not allowed in formulas")) - elseif x.head == :quote || x.head == :$ - throw(FormulaSecurityError("quote/interpolation expressions are not allowed in formulas")) - end - - # For :call expressions, check the function name is on the allowlist +walk(x) = begin + x isa Union{Number, AbstractString, Symbol, LineNumberNode, + Nothing, Bool, QuoteNode} && return nothing + x isa Expr || return nothing + x.head == :macrocall && + return FormulaSecurityError( + "macro calls are not allowed in formulas (got $(x.args[1]))") + x.head in (:cmd, :string) && + return FormulaSecurityError( + "`$(x.head)` expressions are not allowed in formulas") + (x.head == :quote || x.head == :$) && + return FormulaSecurityError( + "quote/interpolation expressions are not allowed in formulas") if x.head == :call fname = x.args[1] if fname isa Symbol && fname ∉ _ALLOWED_CALLS - throw(FormulaSecurityError( + return FormulaSecurityError( "function `$fname` is not in the formula allowlist. " * - "Allowed: arithmetic, math, distributions, DSL operators. " * - "See _ALLOWED_CALLS in BRMMacroWeb.jl for the full list.")) + "Allowed: arithmetic, math, distributions, DSL operators.") end - # Also allow Type{...} constructors if the type name is allowed if fname isa Expr && fname.head == :curly tname = fname.args[1] tname isa Symbol && tname ∉ _ALLOWED_CALLS && - throw(FormulaSecurityError("type constructor `$tname` is not in the formula allowlist")) + return FormulaSecurityError( + "type constructor `$tname` is not in the formula allowlist") end end + x.head ∉ _SAFE_HEADS && + return FormulaSecurityError( + "expression type `:$(x.head)` is not allowed in formulas") + for arg in x.args + v = walk(arg); v === nothing || return v + end + nothing +end - # Check the expression head is expected - if x.head ∉ _SAFE_HEADS - throw(FormulaSecurityError( - "expression type `:$(x.head)` is not allowed in formulas")) +@dynamicstruct struct Formula + text::String + + raw = Meta.parse("begin\n$text\nend") + + violation = walk(raw) + is_safe = violation === nothing + + _t = begin + alllocals = OrderedDict{Symbol,Symbol}() + (; ex=parse!(deepcopy(raw); info=(;alllocals)), alllocals) + end + transformed = _t.ex + alllocals = _t.alllocals +end +@dynamicstruct struct Dataset + n::Int = 64 + seed::Int = 1 + + df = begin + rng = Xoshiro(seed) + a = randn(rng, n) + b = randn(rng, n) + c = randn(rng, n) + d = randn(rng, n) + # Grouping factors with different numbers of levels -- use these on + # the right-hand side of `(... | gN)` to test multiple random-effects + # blocks. + g1 = repeat(1:8, inner=cld(n, 8))[1:n] + g2 = repeat(1:4, inner=cld(n, 4))[1:n] + g3 = rand(rng, 1:6, n) + c1 = rand(rng, 1:3, n) + c2 = rand(rng, 1:2, n) + c3 = rand(rng, 1:4, n) + exposure = 0.5 .+ rand(rng, n) + eta1 = 0.5 .+ 1.2 .* a .- 0.7 .* b .+ 0.3 .* c .+ 0.1 .* d + y1 = eta1 .+ 0.3 .* randn(rng, n) + y2 = -0.2 .+ 0.6 .* a .+ 0.4 .* b .+ 0.2 .* randn(rng, n) + k1 = rand.(rng, Distributions.Poisson.(exp.(0.5 .* eta1))) + k2 = rand.(rng, Distributions.Poisson.(exp.(0.3 .+ 0.4 .* a))) + bin_n = rand(rng, 5:30, n) + bin_p_true = @. 1 / (1 + exp(-(0.2 + 0.5 * a))) + bin_succ = [rand(rng, Distributions.Binomial(n_i, p_i)) + for (n_i, p_i) in zip(bin_n, bin_p_true)] + bin_y = [rand(rng, Distributions.Bernoulli(p_i)) ? 1 : 0 + for p_i in bin_p_true] + DataFrame(; a, b, c, d, g1, g2, g3, c1, c2, c3, exposure, + y1, y2, k1, k2, bin_n, bin_succ, bin_y) end - # Recurse into children - for arg in x.args - _check_formula_safety!(arg) + # NamedTuple view of `df`, merged with namespace-dispatched extras so + # extensions (e.g. bruno-ext.jl) can splice in `dose_times` etc. without + # touching macro.jl. `@brm` only needs `hasproperty`/`getproperty` on + # its data argument, so the NamedTuple stands in for the DataFrame. + container(namespace=:default) = begin + cols = (; (Symbol(c) => df[!, c] for c in names(df))...) + merge(cols, dataset_extras(Val(namespace), df)) end end -# ── Pipeline stages ───────────────────────────────────────────────────────── +# Each example is a `.jl` file under `web-macro/examples/`. File format: # -# Each stage is computed lazily on demand so the user can stop at any -# intermediate step (parsing, transforming, wrapping, eval'ing, materializing, -# benchmarking) and inspect the result without paying for the later stages. - -STAGES = ( - :parse, # Meta.parse(formula) - :transform, # parse!(...) — rewrites = and ~ into @n/@x macro calls - :wrap, # _brm(formula; df) — full let-block ready to eval - :brmi, # eval(...) — BRMI value - :vbrmi, # VBRMI(brmi) — materialized action with blocks/dim - # ── branches after vbrmi (pick one) ── - :bench, # Chairmarks @be primal logdensity - :stan_code, # SBBRMI(brmi): emit @slic body (a) + transpile to Stan (b) -) - -stage_index(s::Symbol) = something(findfirst(==(s), STAGES), length(STAGES)) - -# Run the pipeline up to (and including) `stage`. Returns a NamedTuple -# carrying every intermediate value computed so far. -function pipeline(formula::AbstractString, stage::Symbol; - namespace::Symbol=:default) - s = stage_index(stage) - df = synthetic_df() - out = (; df, namespace) - - s >= 1 || return out - raw = Meta.parse("begin\n$formula\nend") - # Safety check: reject any AST node that isn't in the formula whitelist - # before handing the expression to eval. This blocks arbitrary code - # execution from the formula textarea. - _check_formula_safety!(raw) - out = merge(out, (; raw)) - - s >= 2 || return out - alllocals = OrderedDict{Symbol,Symbol}() - transformed = parse!(deepcopy(raw); info=(;alllocals)) - out = merge(out, (; transformed, alllocals)) - - s >= 3 || return out - container = dataset_container(df, namespace) - wrapped = _brm(formula; df=container) - out = merge(out, (; wrapped)) - - s >= 4 || return out - brmi = eval(wrapped) - out = merge(out, (; brmi)) - - # ── divergent branches after brmi ───────────────────────────────────── - # :stan_code targets the StanBlocks backend, which does NOT need VBRMI — - # SBBRMI walks the BRMI directly. Short-circuit here so the SB branch - # doesn't pay for VBRMI materialization. - if stage === :stan_code - sbbrmi = SBBRMI(brmi) - stan_src = stan_code(sbbrmi) - return merge(out, (; sbbrmi, stan_src)) - end +# # label: 1.1 verify Bernoulli/Binomial — done +# # tier: 1 +# # status: open +# #= +# **Markdown body** with whatever explanation text you want. +# =# +# +# +# +# Header lines (`# key: value`) carry metadata. The `#= ... =#` block is the +# markdown body. Everything after the body block is the formula. The web app +# loads + parses these files on every render of the Examples page, and writes +# them back when the user toggles status or submits an edited formula. +# Reopening the server picks up exactly where the user left off — no in-memory +# state. - s >= 5 || return out - vbrmi = VBRMI(brmi) - dim = LogDensityProblems.dimension(vbrmi) - x0 = randn(Xoshiro(0), dim) - ldp = string(LogDensityProblems.logdensity(vbrmi, x0)) - grad = FiniteDifferences.grad( - central_fdm(5, 1), - Base.Fix1(LogDensityProblems.logdensity, vbrmi), - x0, - )[1] - out = merge(out, (; vbrmi, dim, ldp, x0, grad)) - - stage === :bench || return out - x_rand = randn(dim) - benches = Pair{String,Any}[] - push!(benches, "logdensity (total)" => - @be randn(dim) LogDensityProblems.logdensity($vbrmi, _)) - push!(benches, "lprior!" => - @be randn(dim) lprior!($vbrmi, _)) - # Per-Part lprior! split: the foldl in lprior!(blocks, x) hands each Part - # a view of exactly nparams(part) reals. Reconstruct those slices here so - # each Part's contribution can be benched in isolation. - let pos = 0 - for (group_key, parts) in pairs(vbrmi.meta.blocks) - for (i, part) in enumerate(parts) - n = nparams(part) - xi = view(x_rand, pos+1:pos+n) - push!(benches, " lprior!($group_key[$i] $(part))" => - @be lprior!($part, $xi)) - pos += n +@dynamicstruct struct ExampleEntry + path::String + + _STATUS_COLORS = (open="#888", done="#2e7d32", deprioritized="#a05a2c") + _TIER_LABELS = ("T1", "T2", "T3") + _TIER_COLORS = ("#4a7c59", "#5a6a8c", "#8c5a5a") + + _parsed = begin + lines = readlines(path) + header = Dict{String,String}() + i = 1 + while i <= length(lines) + m = match(r"^# (\w+):\s*(.*)$", lines[i]) + m === nothing && break + header[m[1]] = m[2] + i += 1 + end + body_lines = String[] + if i <= length(lines) && strip(lines[i]) == "#=" + i += 1 + while i <= length(lines) && strip(lines[i]) != "=#" + push!(body_lines, lines[i]) + i += 1 end + i <= length(lines) && (i += 1) # consume `=#` end + formula_text = strip(join(lines[i:end], '\n')) + (; header, + body=join(body_lines, '\n'), + formula=isempty(formula_text) ? nothing : String(formula_text)) end - # llikelihood! splits: each materialized column (either a linear-predictor - # MaterializedColumn or a LikelihoodColumn). Bench each in isolation so the - # allocation / time cost of each step is visible. - _ = lprior!(vbrmi, x_rand) # pre-fill buffers so bench measures the per-step cost - for (key, m) in pairs(vbrmi.meta.materialized) - push!(benches, "llikelihood!($key)" => - @be llikelihood!($m)) + label = get(_parsed.header, "label", basename(path)) + tier = parse(Int, get(_parsed.header, "tier", "1")) + status = Symbol(get(_parsed.header, "status", "open")) + body = _parsed.body + formula = _parsed.formula + slug = replace(basename(path), r"\.jl$" => "") + + border_color = get(_STATUS_COLORS, status, "#888") + tier_label = get(_TIER_LABELS, tier, "T$tier") + tier_color = get(_TIER_COLORS, tier, "#888") + + # DOM ids derived once — HTMX targets reference these (hx_target= / id=). + # Hashing the label keeps ids stable across requests without needing to + # URL-escape the label. + _label_hash = hash(label) + card_id = "example-card-$_label_hash" + result_id = "example-result-$_label_hash" + status_id = "status-$_label_hash" + + tier_pill = h.span(tier_label; + class="brm-tier-pill", + style="background:$tier_color") + + permalink = h.a("🔗"; + href="/examples/$(HTTP.URIs.escapeuri(slug))", + title="Standalone URL", + onclick="event.stopPropagation()", + class="brm-permalink") + + state_pill(target_state, active_text, inactive_text) = begin + is_active = status == target_state + bg = is_active ? get(_STATUS_COLORS, target_state, "#888") : "#888" + h.button(is_active ? active_text : inactive_text; + type="button", + class="brm-state-pill", + hx_get="/mark?label=$(HTTP.URIs.escapeuri(label))&state=$target_state", + hx_target="#$card_id", + hx_swap="outerHTML", + onclick="event.stopPropagation()", + style="background:$bg", + ) end - merge(out, (; benches)) -end -# ── Rendering helpers ─────────────────────────────────────────────────────── - -_section(title, body) = (h.h3(title), h.pre(body)) - -# ── Styled HTML rendering for BRMI / VBRMI cards ─────────────────────────── -# -# Each symbol gets a deterministic color, data columns are bold, parameters -# are italic, and likelihood statements are underlined. The tree walker -# (_html_expr) converts an ExprColumn AST into a nest of s. -# -# TODO: HTMX.jl should grow a generic `htmx_node(x)::Node` extension point -# so downstream packages can overload once and have every consumer dispatch -# automatically. For now these card functions are wired in by hand. - -# Deterministic HSL color per symbol (golden-ratio spread for visual variety). -_symbol_color(name::Symbol) = "hsl($(mod(hash(name) * 137, 360)), 60%, 40%)" - -# A colored with role-based font styling. -# Data columns are normal weight; parameters (latent/sampled) are bold. -_styled_name(name::Symbol, role::Symbol) = begin - s = "color:$(_symbol_color(name));" - role == :parameter && (s *= "font-weight:bold;") - h.span(string(name); style=s) -end + status_pills = h.span(; + id=status_id, + class="brm-status-pills")( + state_pill(:done, "✓ done", "mark done"), + state_pill(:deprioritized, "✓ deprioritized", "deprioritize"), + ) -# ── _html_expr: recursive ExprColumn → styled HTML ───────────────────────── - -_html_expr(x::NamedColumn{<:Any, <:DataColumn}) = _styled_name(name(x), :data) -_html_expr(x::NamedColumn{<:Any, MissingColumn}) = _styled_name(name(x), :parameter) -_html_expr(x::NamedColumn) = _styled_name(name(x), :derived) -_html_expr(x::Int) = h.span(string(x); style="color:#666") -_html_expr(x::Float64) = h.span(string(x); style="color:#666") -_html_expr(x::Number) = h.span(string(x); style="color:#666") -_html_expr(x::DataColumn) = h.span("data($(eltype(parent(x))))"; style="color:#999") -_html_expr(x::MaterializedColumn) = _html_expr(getbroadcast(x)) -_html_expr(x::LikelihoodColumn) = h.span( - _html_expr(parent(x)), h.span(" .~ "; style="color:#333"), _html_expr(rhs(x))) - -# Infix operators: always parenthesized so inner expressions like (1 + b | g1) -# keep their grouping. Top-level callers (_html_brmi_row) use _html_infix -# directly to skip the outermost parens. -_html_expr(x::ExprColumn{<:Union{typeof.((~,*,+,|,doublepipe,assign))...}}) = begin - h.span("(", _html_infix(x), ")") -end + formula_form(routes) = h.form(; class="brm-example-form")( + h.input(; type="hidden", name="label", value=label), + h.textarea(formula; + name="formula", + rows=max(3, count('\n', formula) + 1), + class="brm-example-textarea"), + h.button("cimpl (bench) ▶"; + type="button", + class="brm-branch-btn", + hx_get=string(query_url(routes/"stage/bench"; force=true)), + hx_include="closest form", + hx_target="#$result_id", + hx_swap="innerHTML"), + h.button("sbimpl (compile) ▶"; + type="button", + class="brm-branch-btn", + hx_get=string(query_url(routes/"stage/stan_compile"; force=true)), + hx_include="closest form", + hx_target="#$result_id", + hx_swap="innerHTML"), + ) -_html_infix(x::ExprColumn) = begin - op_str = " $(getop(x)) " - args = getargs(x) - parts = Any[] - for (i, arg) in enumerate(args) - i > 1 && push!(parts, h.span(op_str; style="color:#555")) - push!(parts, _html_expr(arg)) + card(routes) = begin + children = Any[HTMXObjects.md_to_node(body)] + if formula !== nothing + push!(children, formula_form(routes)) + # Inline pipeline-result target — the form's hx_get fills this div + # with `render_pipeline(out)` so the user sees the + # VBRMI/finite-difference output right inside the card. + push!(children, h.div(; + id=result_id, + class="brm-example-result")) + end + # `:open` status → expanded; `:done`/`:deprioritized` → collapsed by + # default. Pills sit inside the so they're always reachable, + # but their onclick stops propagation so clicking a pill doesn't also + # toggle the disclosure. + h.article(; + id=card_id, + class="brm-example-card", + style="border-left-color:$border_color", + )( + h.details(; open=(status == :open))( + h.summary(; class="brm-example-summary")( + tier_pill, " ", + h.strong(label), " ", + status_pills, " ", + permalink, + ), + h.div(; class="brm-example-body")(children...), + ), + ) end - h.span(parts...) -end -# Function-call style: fname(args...; kwargs...) -_html_expr(x::ExprColumn) = begin - fname = getf(x) isa Function ? nameof(getf(x)) : - getf(x) isa Type ? nameof(getf(x)) : string(getf(x)) - args = getargs(x) - kw = getkwargs(x) - parts = Any[h.span(string(fname); style="color:#777"), "("] - for (i, arg) in enumerate(args) - i > 1 && push!(parts, ", ") - push!(parts, _html_expr(arg)) - end - if length(kw) > 0 - push!(parts, "; ") - for (i, (k, v)) in enumerate(pairs(kw)) - i > 1 && push!(parts, ", ") - push!(parts, "$k=", _html_expr(v)) + write_with!(; new_status=status, new_formula=formula) = begin + io = IOBuffer() + println(io, "# label: ", label) + println(io, "# tier: ", tier) + println(io, "# status: ", new_status) + if !isempty(body) + println(io, "#=") + println(io, body) + println(io, "=#") end + if new_formula !== nothing && !isempty(new_formula) + println(io) + print(io, new_formula) + endswith(new_formula, "\n") || println(io) + end + write(path, take!(io)) + ExampleEntry(; path) end - push!(parts, ")") - h.span(parts...) end +@dynamicstruct struct ExampleStore + dir::String + + # `entries` is a method, not a cached field, because `save!` writes to + # disk and we want subsequent reads to see the new file mtime/content. + entries() = begin + isdir(dir) || mkpath(dir) + files = sort(filter(endswith(".jl"), + readdir(dir; join=true)); + by=mtime, rev=true) + ExampleEntry[ExampleEntry(; path=f) for f in files] + end -# Broadcasted objects (from VBRMI materialization): walk their inner structure -_html_expr(x::Base.Broadcast.Broadcasted) = begin - fname = x.f isa Function ? nameof(x.f) : - x.f isa Type ? nameof(x.f) : string(x.f) - args = x.args - # Infix for common operators - if x.f in (+, -, *, /) - parts = Any[] - for (i, arg) in enumerate(args) - i > 1 && push!(parts, h.span(" $(x.f) "; style="color:#555")) - push!(parts, _html_expr(arg)) + find(label) = begin + for e in entries() + e.label == label && return e end - return h.span(parts...) + nothing end - parts = Any[h.span(string(fname); style="color:#777"), "("] - for (i, arg) in enumerate(args) - i > 1 && push!(parts, ", ") - push!(parts, _html_expr(arg)) + find_by_slug(slug) = begin + for e in entries() + e.slug == slug && return e + end + nothing end - push!(parts, ")") - h.span(parts...) -end -# Arrays / views from block parameter slots: show as a compact shape description -_html_expr(x::SubArray) = h.span("param[$(join(size(x), "×"))]"; - style="font-style:italic;color:#888") -_html_expr(x::AbstractVector{<:Number}) = h.span("vec[$(length(x))]"; - style="font-weight:bold;color:#888") -_html_expr(x::Base.RefValue) = _html_expr(x[]) -_html_expr(x::AbstractMatrix) = h.span("mat[$(join(size(x), "×"))]"; - style="font-style:italic;color:#888") - -# Catch-all -_html_expr(x) = h.span(sprint(show, x; context=:compact=>true); style="color:#999") - -# ── BRMI card ─────────────────────────────────────────────────────────────── - -function brmi_card(brmi::BRMI) - rows = [_html_brmi_row(key, parent(value)) for (key, value) in pairs(brmi.operations)] - h.article(; style="margin:0.5rem 0")( - h.header(h.strong("BRMI"), - h.small(" — $(length(brmi.operations)) operations")), - h.div(; style="font-family:monospace;font-size:1.15em;line-height:1.8;padding:0.3rem 0")( - rows...), - ) + save!(label; new_status=nothing, new_formula=nothing) = begin + e = find(label) + e === nothing && return nothing + e.write_with!(; + new_status = new_status === nothing ? e.status : new_status, + new_formula = new_formula === nothing ? e.formula : new_formula, + ) + end end +# One run of the @brm pipeline for a given (text, namespace) pair. Every +# stage is a derived property -- accessing `run.brmi` triggers parse + eval, +# accessing `run.benches` triggers vbrmi + finite-difference check + the +# benchmark loop, etc. Unused branches don't compute. Safety is enforced +# at the one point where it matters: `wrapped` refuses to produce Julia +# code for an unsafe formula. + +@dynamicstruct struct BRMRun + text::String + namespace::Symbol = :default + (; dataset) = __parent__ + + formula = Formula(; text) + df = dataset.df + container = dataset.container(namespace) + + wrapped = begin + formula.is_safe || throw(formula.violation) + _brm(text; df=container) + end + brmi = eval(wrapped) -_leaf_column(x::NamedColumn) = x -_leaf_column(x::ExprColumn) = _leaf_column(getargs(x)[1]) -_leaf_column(x) = x - -function _html_brmi_row(key, op::ExprColumn{typeof(~)}) - lhs_leaf = _leaf_column(getargs(op)[1]) - is_likelihood = lhs_leaf isa NamedColumn && parent(lhs_leaf) isa DataColumn - content = _html_infix(op) # no outer parens at top level - style = is_likelihood ? - "text-decoration:underline;text-decoration-color:#aaa;text-underline-offset:3px;" : "" - h.div(content; style) -end -_html_brmi_row(key, op::ExprColumn{typeof(assign)}) = h.div(_html_infix(op)) -_html_brmi_row(key, op) = h.div( - _styled_name(key, :data), - h.span(": "; style="color:#666"), - h.span(sprint(show, op); style="color:#999"), -) - -# ── VBRMI card ────────────────────────────────────────────────────────────── - -function vbrmi_card(vbrmi::VBRMI) - brmi = getfield(vbrmi, :parent) - (; meta) = vbrmi - n_dim = LogDensityProblems.dimension(vbrmi) - n_mat = length(meta.materialized) - n_blocks = length(meta.blocks) - - # Materialized columns: show the BRMI-level symbolic expression (styled) - # plus a compact description of the materialized value type. - mat_rows = [begin - nc = get(brmi.operations, key, nothing) - inner = nc !== nothing ? parent(nc) : nothing - if inner isa DataColumn - # Data column: styled name + eltype - h.div(_styled_name(key, :data), - h.span(": data($(eltype(parent(inner))))"; style="color:#999")) - elseif value isa LikelihoodColumn - # Likelihood: full expression, underlined - expr = inner !== nothing ? _html_infix(inner) : _styled_name(key, :data) - h.div(expr; style="text-decoration:underline;text-decoration-color:#aaa;text-underline-offset:3px") - elseif inner !== nothing - # Materialized (sampled/derived): symbolic expression + shape - expr = inner isa ExprColumn{<:Union{typeof(~),typeof(assign)}} ? - _html_infix(inner) : _html_expr(inner) - shape = "$(eltype(parent(value)))[$(length(parent(value)))]" - h.div(expr, h.span(" → $shape"; style="color:#999;font-size:0.85em")) - else - h.div(_styled_name(key, :derived), - h.span(": $(sprint(show, value))"; style="color:#999")) + # ── cimpl branch ── + vbrmi = VBRMI(brmi) + dim = LogDensityProblems.dimension(vbrmi) + x0 = randn(Xoshiro(0), dim) + ldp = string(LogDensityProblems.logdensity(vbrmi, x0)) + grad = FiniteDifferences.grad( + central_fdm(5, 1), + Base.Fix1(LogDensityProblems.logdensity, vbrmi), + x0, + )[1] + + benches = begin + x_rand = randn(dim) + bs = Pair{String,Any}[] + push!(bs, "logdensity (total)" => + @be randn(dim) LogDensityProblems.logdensity($vbrmi, _)) + push!(bs, "lprior!" => + @be randn(dim) lprior!($vbrmi, _)) + # Per-Part lprior! split: the foldl in lprior!(blocks, x) hands + # each Part a view of exactly nparams(part) reals. Reconstruct + # those slices here so each Part's contribution can be benched + # in isolation. + let pos = 0 + for (group_key, parts) in pairs(vbrmi.meta.blocks) + for (i, part) in enumerate(parts) + n = nparams(part) + xi = view(x_rand, pos+1:pos+n) + push!(bs, " lprior!($group_key[$i] $(part))" => + @be lprior!($part, $xi)) + pos += n + end + end + end + # llikelihood! splits: each materialized column (either a + # linear-predictor MaterializedColumn or a LikelihoodColumn). + _ = lprior!(vbrmi, x_rand) + for (key, m) in pairs(vbrmi.meta.materialized) + push!(bs, "llikelihood!($key)" => + @be llikelihood!($m)) end - end for (key, value) in pairs(meta.materialized)] + bs + end - # Blocks: each block is a tuple of parts; print the block key + one line per part. - blocks_rows = [begin - role = key === :__population__ ? :derived : :data - h.div( - _styled_name(key, role), - [h.div(; style="color:#999;margin-left:1.5rem")(sprint(show, part)) - for part in parts]..., - ) - end for (key, parts) in pairs(meta.blocks)] - - h.article(; style="margin:0.5rem 0")( - h.header(h.strong("VBRMI"), - h.small(" — dim $n_dim, $n_mat materialized, $n_blocks blocks")), - h.h6(; style="margin-bottom:0.2rem")("materialized"), - h.div(; style="font-family:monospace;font-size:1.15em;line-height:1.8;margin-left:1rem")( - mat_rows...), - h.h6(; style="margin-bottom:0.2rem")("blocks"), - h.div(; style="font-family:monospace;font-size:1.15em;line-height:1.8;margin-left:1rem")( - blocks_rows...), - ) + # ── stan branch ── + sbbrmi = SBBRMI(brmi) + stan_src = stan_code(sbbrmi) + stan_file = begin + f = tempname() * ".stan" + write(f, stan_src) + f + end + stan_lib = BridgeStan.compile_model(stan_file) end +@dynamicstruct struct AppData + __status__ = initialize_progress!(:state; description="BRM pipeline") + examples_dir = joinpath(dirname(@__DIR__), "examples") + default_formula = """loc1 ~ 1 + a + c1 + (1 + b + c1 | g1) + (1 | g2) +log(err1) ~ 1 + d +y1 ~ Normal(loc1, err1) -function render_pipeline(out::NamedTuple) - sections = Vector{Any}[] # one entry per stage; rendered most-recent-first +log_rate ~ 1 + a + (1 | g3) +k1 ~ Poisson(exp(log_rate)) - # Synthetic data always pinned at the top, collapsed by default so the - # macro pipeline output stays the focus. - data_section = Any[ - h.details( - h.summary("Synthetic data ($(nrow(out.df)) rows × $(ncol(out.df)) cols: " * - join(string.(names(out.df)), ", ") * ") — click to expand"), - render_table(out.df; sortable=false), - ), - ] +log_odds_bin ~ 1 + c2 + (1 | g2) +bin_succ ~ Binomial(bin_n, logistic(log_odds_bin)) - if haskey(out, :raw) - push!(sections, Any[_section("1. Meta.parse — raw Julia AST", - sprint(show, out.raw))...]) - end - if haskey(out, :transformed) - push!(sections, Any[ - _section("2. parse! — rewritten AST (= → @n/@x assign, ~ → @n/@x ~)", - sprint(show, out.transformed))..., - h.h3(" locals classified by parse!"), - h.pre(sprint(show, out.alllocals)), - ]) - end - if haskey(out, :wrapped) - push!(sections, Any[_section("3. _brm — full let-block (df spliced as a literal)", - sprint(show, out.wrapped))...]) - end - if haskey(out, :brmi) - push!(sections, Any[ - h.h3("4. eval — BRMI value (parsed model)"), - brmi_card(out.brmi), - ]) - end - if haskey(out, :vbrmi) - tol = 1e-8 - n_dead = count(<=(tol) ∘ abs, out.grad) - fd_summary = n_dead == 0 ? - h.span("logdensity + FD check: $(out.dim)/$(out.dim) active ✓"; - style="color:green") : - h.span("logdensity + FD check: $(n_dead) dead param(s)"; - style="color:crimson") - dead = findall(<=(tol) ∘ abs, out.grad) - fd_body = h.div( - h.p("dim = ", string(out.dim), ", logdensity = ", out.ldp), - isempty(dead) ? "" : - h.p(; style="color:crimson")( - "dead param indices: ", string(dead)), - h.pre(sprint(show, MIME"text/plain"(), out.grad)), - ) - push!(sections, Any[ - h.h3("5. VBRMI — materialized action (blocks, dim, columns)"), - vbrmi_card(out.vbrmi), - h.details(h.summary(fd_summary), fd_body), - ]) - end - if haskey(out, :benches) - bench_rows = [h.div( - h.strong(label), h.br(), - h.pre(sprint(show, MIME"text/plain"(), b)) - ) for (label, b) in out.benches] - push!(sections, Any[h.h3("6. Chairmarks @be — per-step"), bench_rows...]) - end - if haskey(out, :sbbrmi) - push!(sections, Any[ - h.h3("5. SBBRMI — emitted @slic body"), - h.pre(sprint(show, out.sbbrmi.model.model)), - h.p("data keys: ", - h.code(string(sort(collect(keys(out.sbbrmi.data)))))), - ]) - push!(sections, Any[ - h.h3("6. transpiled Stan source"), - h.pre(out.stan_src), - ]) - end +log_odds_b ~ 1 + b +bin_y ~ Bernoulli(logistic(log_odds_b)) +""" - # Stages render most-recent-first; synthetic data sits at the very top. - children = reduce(vcat, reverse(sections); init=Any[]) - prepend!(children, data_section) - h.div(; id="brm-macro-output")(children...) + dataset = Dataset() + example_store = ExampleStore(; dir=examples_dir) + + namespace_from(label) = isempty(strip(label)) ? :default : + Symbol(lowercase(first(split(strip(label), r"[\s:\-]+")))) + + # Ordered pipeline stages. Index gates which BRMRun properties + # `pipeline_run` touches (and which sections `render_pipeline` shows). + stages = (:parse, :transform, :wrap, :brmi, + :vbrmi, :bench, + :slic_model, :stan_code, :stan_compile) + stage_index(s) = something(findfirst(==(s), stages), length(stages)) + + # Indexable property: `appdata.run[text, ns]` is cached in-memory per key, + # `appdata.run(text, ns)` is fresh each call. + run(text, namespace=:default) = + BRMRun(; __parent__=__self__, text, namespace) + + # Indexable fetch for `polling_fetchindex` (accessed via brackets by the + # caller). Touches BRMRun properties up through `stage` so the heavy work + # lands inside the polled task rather than the HTTP response callback. + pipeline_run(text, stage::Symbol, namespace=:default) = begin + r = run[text, namespace] + s = stage_index(stage) + s >= 1 && r.formula.raw + s >= 2 && r.formula.transformed + s >= 3 && r.wrapped + s >= 4 && r.brmi + stage === :vbrmi && r.grad + stage === :bench && r.benches + stage === :slic_model && r.sbbrmi + stage === :stan_code && r.stan_src + stage === :stan_compile && r.stan_lib + r + end end -# ── Routes ────────────────────────────────────────────────────────────────── - -_stage_button(self, label, stage) = h.button(label; type="button", - id="stage-$stage", - hx_get=string(query_url(self/"stage/$stage"; force=true)), - hx_include="#brm-macro-form", - hx_target="#brm-macro-output", - hx_swap="outerHTML") - -# Pre-canned formulas. The ones above the divider exercise individual features -# in isolation; the last one stacks everything into a single multi-likelihood -# model. Click loads the formula into the textarea — user still hits "Render" -# (or any stage button) to advance the pipeline. -function presets() - [ +APPDATA = AppData(; cache_type=:parallel) +@htmx struct AppContext + __appdata__ = APPDATA + (; default_formula, dataset, example_store, namespace_from, + stages, stage_index, run, pipeline_run) = __appdata__ + + # Page-level stylesheet read once at construction. Classes are consumed by + # ExampleEntry.card / html_expr.jl; per-symbol / per-status colors that + # are data-derived stay inline on the element. + css = read(joinpath(@__DIR__, "brm-macro.css"), String) + + # HTMXObjects auto-uses `__page__` to wrap any route's return value into a + # full page on direct browser navigation, while returning just the fragment + # for HTMX requests (see `_resolve_response` in HTMXObjects.jl). The + # sidebar's `hx-get` swaps target `#content` directly. + __page__(content) = htmx( + h.div(; class="brm-layout")( + nav_sidebar([ + "Pipeline" => "/", + "Examples" => "/examples", + ]), + h.main(; class="container brm-main")( + h.div(; id="content")(content), + ), + ); + pico_version="2", + extra_head=( + h.title("BRM macro action"), + h.style(__self__.css), + ), + ) + + # Pre-canned formulas. The ones above the divider exercise individual + # features in isolation; the last one stacks everything into a single + # multi-likelihood model. + presets = [ "min" => "loc ~ 1\nlog(err) ~ 1\ny1 ~ Normal(loc, err)\n", "linear" => "loc ~ 1 + a\nlog(err) ~ 1\ny1 ~ Normal(loc, err)\n", "multi-lin" => "loc ~ 1 + a + b + c + d\nlog(err) ~ 1\ny1 ~ Normal(loc, err)\n", @@ -606,767 +546,239 @@ function presets() "Poisson" => "log_rate ~ 1 + a + (1 | g1)\nk1 ~ Poisson(exp(log_rate))\n", "Binomial" => "log_odds ~ 1 + a + (1 | g1)\nbin_succ ~ Binomial(bin_n, logistic(log_odds))\n", "Bernoulli" => "log_odds ~ 1 + a + (1 | g1)\nbin_y ~ Bernoulli(logistic(log_odds))\n", - # Joint smoke test for Peter's verified non-Normal examples: - # brms::cbpp_binomial → categorical + random intercept + Binomial with - # per-row trial counts; kruschke::therapeutic_touch → hierarchical - # Bernoulli. Both share the same grouping factor here so they hit the - # cross-likelihood block-sharing path too. "cbpp + therapeutic touch" => """log_odds_bin ~ 1 + c1 + (1 | g1) bin_succ ~ Binomial(bin_n, logistic(log_odds_bin)) log_odds_b ~ 1 + (1 | g1) bin_y ~ Bernoulli(logistic(log_odds_b)) """, - "everything" => default_formula(), + "everything" => default_formula, ] -end -_preset_button(label, formula) = h.button(label; - type="button", - data_formula=formula, - onclick="document.querySelector('textarea[name=formula]').value = this.dataset.formula; document.getElementById('stage-vbrmi').click()", - style="font-size:0.8em;padding:0.2rem 0.5rem;margin:0") - -_index_body(self, formula::String) = h.div( - h.h1("BRM macro pipeline"), - h.p( - "Enter a ", h.code("@brm"), " formula and step through the macro pipeline: ", - h.code("Meta.parse"), " -> ", h.code("parse!"), " -> ", h.code("_brm"), - " let-block -> ", h.code("eval"), " -> ", h.code("VBRMI"), " action -> ", - h.code("Chairmarks"), " benchmark.", - ), - h.details( - h.summary(h.small("Allowed functions in formulas")), - h.p(h.small( - join(sort(collect(string.(s) for s in _ALLOWED_CALLS)), ", "), - )), - ), - h.form(; id="brm-macro-form")( - h.label("Load preset"), - h.div(; style="display:flex;flex-wrap:wrap;gap:0.3rem;margin-bottom:0.6rem")( - [_preset_button(label, body) for (label, body) in presets()]..., - ), - h.label("Formula")( - h.textarea(formula; - name="formula", rows=8, - style="width:100%;font-family:monospace"), - ), - h.fieldset(; class="grid")( - _stage_button(self, "1. Parse", :parse), - _stage_button(self, "2. Transform", :transform), - _stage_button(self, "3. Wrap", :wrap), - _stage_button(self, "4. BRMI", :brmi), - _stage_button(self, "5. VBRMI", :vbrmi), - ), - h.small("Pick a branch:"), - h.fieldset(; class="grid")( - _stage_button(self, "6. Benchmark", :bench), - _stage_button(self, "6. Stan code", :stan_code), - ), - ), - lazy(string(query_url(self/"stage/bench"; formula)); id="brm-macro-output"), -) + preset_button(label, formula) = h.button(label; + type="button", + class="brm-preset-btn", + data_formula=formula, + onclick="document.querySelector('textarea[name=formula]').value = this.dataset.formula; document.getElementById('stage-vbrmi').click()") -@htmx struct AppContext - __appdata__ = APPDATA + stage_button(label, stage) = h.button(label; + type="button", + id="stage-$stage", + hx_get=string(query_url(__self__/"stage/$stage"; force=true)), + hx_include="#brm-macro-form", + hx_target="#brm-macro-output", + hx_swap="outerHTML") + + render_pipeline(r, stage) = begin + s = stage_index(stage) + sections = Vector{Any}[] + + # Synthetic data pinned at top, collapsed by default so the macro + # pipeline output stays the focus. + data_section = Any[ + h.details( + h.summary("Synthetic data ($(nrow(r.df)) rows × $(ncol(r.df)) cols: " * + join(string.(names(r.df)), ", ") * ") — click to expand"), + render_table(r.df; sortable=false), + ), + ] - # HTMXObjects auto-uses `__page__` to wrap any route's return value into a - # full page on direct browser navigation, while returning just the fragment - # for HTMX requests (see `_resolve_response` in HTMXObjects.jl). The - # sidebar's `hx-get` swaps target `#content` directly. - __page__(content) = htmx( - h.div(; style="display:flex;gap:1rem;align-items:flex-start")( - nav_sidebar([ - "Pipeline" => "/", - "TODO list" => "/todo", - ]), - h.main(; class="container", style="flex:1;min-width:0")( - h.div(; id="content")(content), + s >= 1 && push!(sections, Any[ + h.h3("1. Meta.parse — raw Julia AST"), + h.pre(sprint(show, r.formula.raw)), + ]) + s >= 2 && push!(sections, Any[ + h.h3("2. parse! — rewritten AST (= → @n/@x assign, ~ → @n/@x ~)"), + h.pre(sprint(show, r.formula.transformed)), + h.h3(" locals classified by parse!"), + h.pre(sprint(show, r.formula.alllocals)), + ]) + s >= 3 && push!(sections, Any[ + h.h3("3. _brm — full let-block (df spliced as a literal)"), + h.pre(sprint(show, r.wrapped)), + ]) + s >= 4 && push!(sections, Any[ + h.h3("4. eval — BRMI value (parsed model)"), + brmi_card(r.brmi), + ]) + + if stage in (:vbrmi, :bench) + tol = 1e-8 + n_dead = count(<=(tol) ∘ abs, r.grad) + fd_summary = n_dead == 0 ? + h.span("logdensity + FD check: $(r.dim)/$(r.dim) active ✓"; + class="brm-status-ok") : + h.span("logdensity + FD check: $(n_dead) dead param(s)"; + class="brm-status-err") + dead = findall(<=(tol) ∘ abs, r.grad) + fd_body = h.div( + h.p("dim = ", string(r.dim), ", logdensity = ", r.ldp), + isempty(dead) ? "" : + h.p(; class="brm-status-err")( + "dead param indices: ", string(dead)), + h.pre(sprint(show, MIME"text/plain"(), r.grad)), + ) + push!(sections, Any[ + h.h3("5. VBRMI — materialized action (blocks, dim, columns)"), + vbrmi_card(r.vbrmi), + h.details(h.summary(fd_summary), fd_body), + ]) + end + if stage === :bench + bench_rows = [h.div( + h.strong(lbl), h.br(), + h.pre(sprint(show, MIME"text/plain"(), b)) + ) for (lbl, b) in r.benches] + push!(sections, Any[h.h3("6. Chairmarks @be — per-step"), bench_rows...]) + end + if stage in (:slic_model, :stan_code, :stan_compile) + push!(sections, Any[ + h.h3("5a. SlicModel — SBBRMI @slic body"), + h.pre(sprint(show, r.sbbrmi.model.model)), + h.p("data keys: ", + h.code(string(sort(collect(keys(r.sbbrmi.data)))))), + ]) + end + if stage in (:stan_code, :stan_compile) + push!(sections, Any[ + h.h3("5b. StanCode — transpiled Stan source"), + h.pre(r.stan_src), + ]) + end + if stage === :stan_compile + push!(sections, Any[ + h.h3("5c. StanCompile — BridgeStan shared library"), + h.p("stan file: ", h.code(r.stan_file)), + h.p("compiled .so: ", h.code(r.stan_lib)), + ]) + end + + # Stages render most-recent-first; synthetic data sits at the very top. + children = reduce(vcat, reverse(sections); init=Any[]) + prepend!(children, data_section) + h.div(; id="brm-macro-output")(children...) + end + + index_body(formula) = h.div( + h.h1("BRM macro pipeline"), + h.p( + "Enter a ", h.code("@brm"), " formula and step through the macro pipeline: ", + h.code("Meta.parse"), " -> ", h.code("parse!"), " -> ", h.code("_brm"), + " let-block -> ", h.code("eval"), " -> ", h.code("VBRMI"), " action -> ", + h.code("Chairmarks"), " benchmark.", + ), + h.details( + h.summary(h.small("Allowed functions in formulas")), + h.p(h.small( + join(sort(collect(string.(s) for s in _ALLOWED_CALLS)), ", "), + )), + ), + h.form(; id="brm-macro-form")( + h.label("Load preset"), + h.div(; class="brm-preset-row")( + [preset_button(lbl, body) for (lbl, body) in presets]..., + ), + h.label("Formula")( + h.textarea(formula; + name="formula", rows=8, + class="brm-formula-textarea"), + ), + h.fieldset(; class="grid")( + stage_button("1. Parse", :parse), + stage_button("2. Transform", :transform), + stage_button("3. Wrap", :wrap), + stage_button("4. BRMI", :brmi), + ), + h.small("Pick a branch:"), + h.fieldset(; class="grid")( + stage_button("5. VBRMI", :vbrmi), + stage_button("6. Benchmark", :bench), + ), + h.fieldset(; class="grid")( + stage_button("5a. SlicModel", :slic_model), + stage_button("5b. StanCode", :stan_code), + stage_button("5c. StanCompile",:stan_compile), ), - ); - pico_version="2", - extra_head=( - h.title("BRM macro action"), - h.style(":root { font-size: 87.5%; }"), ), + lazy(string(query_url(__self__/"stage/bench"; formula)); id="brm-macro-output"), ) - @get index(; formula::String=default_formula(), label::String="") = begin - # If a TODO form posted us a (label, formula) pair, persist the edited - # formula to that TODO's .jl file so the next visit to the TODO page - # shows the user's edits instead of the seed default. - isempty(label) || _save_todo!(__appdata__, label; new_formula=formula) - _index_body(__self__, formula) + @get index(; formula::String=default_formula, label::String="") = begin + # If an example form posted us a (label, formula) pair, persist the + # edited formula to that example's .jl file so the next visit to the + # Examples page shows the user's edits instead of the seed default. + isempty(label) || example_store.save!(label; new_formula=formula) + index_body(formula) end @get mark(; label::String="", state::String="") = begin isempty(label) && return "" target = Symbol(state) - entry = _find_todo(__appdata__, label) + entry = example_store.find(label) entry === nothing && return "" next_status = entry.status == target ? :open : target - updated = _save_todo!(__appdata__, label; new_status=next_status) + updated = example_store.save!(label; new_status=next_status) # Re-render the whole card so the border + collapse state update # together with the pill text. - _todo_card(__self__, updated) + updated === nothing ? "" : updated.card(__self__) end - @get stage(name::AbstractString; formula::String=default_formula(), + @get stage(name::AbstractString; formula::String=default_formula, label::String="", force::Bool=false) = begin - # When called from a TODO card's form, persist the (possibly edited) - # formula back to the todo's .jl file before rendering. - isempty(label) || _save_todo!(__appdata__, label; new_formula=formula) - ns = dataset_namespace(label) - polling_fetchindex(__appdata__.pipeline_result, - formula, Symbol(name), ns; + # When called from an example card's form, persist the (possibly + # edited) formula back to the example's .jl file before rendering. + isempty(label) || example_store.save!(label; new_formula=formula) + ns = namespace_from(label) + stage_sym = Symbol(name) + polling_fetchindex(pipeline_run, + formula, stage_sym, ns; poll_url=string(query_url(__self__/"stage/$name"; formula, label)), label="BRM pipeline - $name", - force) do out - render_pipeline(out) + force) do r + render_pipeline(r, stage_sym) end end - @get todo(; slug::String="") = begin + # Focused per-model views of the sbimpl intermediate artifacts. Each + # route runs the pipeline just far enough and returns the relevant + # source in `h.pre` (plus markdown_only serves the bare source via + # `?plain` / `Accept: text/plain`, for piping into agents or curl). + @get slic(; formula::String=default_formula, label::String="") = begin + isempty(label) || example_store.save!(label; new_formula=formula) + h.pre(sprint(show, run(formula, namespace_from(label)).sbbrmi.model.model)) + end + + @get stan(; formula::String=default_formula, label::String="") = begin + isempty(label) || example_store.save!(label; new_formula=formula) + h.pre(run(formula, namespace_from(label)).stan_src) + end + + @get examples(slug::String="") = begin if !isempty(slug) - entry = _find_todo_by_slug(__appdata__, slug) + entry = example_store.find_by_slug(slug) entry === nothing && return h.div( - h.p("No TODO with slug ", h.code(slug), "."), - h.a("<- Back to TODO list"; href="/todo"), + h.p("No example with slug ", h.code(slug), "."), + h.a("<- Back to Examples"; href="/examples"), ) return h.div( - h.p(h.a("<- Back to TODO list"; href="/todo")), - _todo_card(__self__, entry), + h.p(h.a("<- Back to Examples"; href="/examples")), + entry.card(__self__), ) end - todos = _load_todos(__appdata__) h.div( - h.h1("TODO - what's missing for full BRM coverage"), - h.p("Sorted by last modified. Each item has a sketch of what it is, why it matters, how to implement, and how to verify. Sourced from .jl files under ", h.code("web-macro/todos/"), "; status edits and edited formulas are written back to disk."), - [_todo_card(__self__, t) for t in todos]..., + h.h1("Examples - coverage gaps and demos for BRM"), + h.p("Sorted by last modified. Each item has a sketch of what it is, why it matters, how to implement, and how to verify. Sourced from .jl files under ", h.code("web-macro/examples/"), "; status edits and edited formulas are written back to disk."), + [e.card(__self__) for e in example_store.entries()]..., ) end end -# ── TODO content (file-backed) ────────────────────────────────────────────── -# -# Each TODO is a `.jl` file under `web-macro/todos/`. File format: -# -# # label: 1.1 verify Bernoulli/Binomial — done -# # tier: 1 -# # status: open -# #= -# **Markdown body** with whatever explanation text you want. -# =# -# -# -# -# Header lines (`# key: value`) carry metadata. The `#= ... =#` block is the -# markdown body. Everything after the body block is the formula. The web app -# loads + parses these files on every render of the TODO page, and writes them -# back when the user toggles status or submits an edited formula. Reopening the -# server picks up exactly where the user left off — no in-memory state. - -struct TodoEntry - path::String - label::String - tier::Int - status::Symbol # :open | :done | :deprioritized - body::String # markdown - formula::Union{String,Nothing} -end - -_todos_dir() = joinpath(dirname(@__DIR__), "todos") -_slug(label::AbstractString) = lowercase(strip(replace(label, r"[^\w.]+" => "-"), '-')) - -@dynamicstruct struct AppData - __status__ = initialize_progress!(:state; description="BRM pipeline") - @cached pipeline_result(formula, stage, namespace) = - pipeline(formula, stage; namespace) -end - -function _load_todos(::AppData) - dir = _todos_dir() - isdir(dir) || _migrate_todos!() - files = sort(filter(endswith(".jl"), readdir(dir; join=true)); by=mtime, rev=true) - TodoEntry[_parse_todo_file(f) for f in files] -end - -function _find_todo(appdata::AppData, label::AbstractString) - for t in _load_todos(appdata) - t.label == label && return t - end - nothing -end - -_todo_slug(t::TodoEntry) = replace(basename(t.path), r"\.jl$" => "") - -function _find_todo_by_slug(appdata::AppData, slug::AbstractString) - for t in _load_todos(appdata) - _todo_slug(t) == slug && return t - end - nothing -end - -function _parse_todo_file(path::String) - lines = readlines(path) - header = Dict{String,String}() - i = 1 - - # Header: leading lines matching `# key: value`. Stop at the first - # non-matching line. Line-based avoids any UTF-8 byte-index footguns - # (labels routinely contain multi-byte characters like `✓`). - while i <= length(lines) - m = match(r"^# (\w+):\s*(.*)$", lines[i]) - m === nothing && break - header[m[1]] = m[2] - i += 1 - end - - # Body: optional `#= ... =#` block. Both delimiters live on their own - # lines (the writer guarantees this), so a simple line scan suffices. - body_lines = String[] - if i <= length(lines) && strip(lines[i]) == "#=" - i += 1 - while i <= length(lines) && strip(lines[i]) != "=#" - push!(body_lines, lines[i]) - i += 1 - end - i <= length(lines) && (i += 1) # consume `=#` - end - body = join(body_lines, '\n') - - # Formula: everything that's left, stripped. - formula_text = strip(join(lines[i:end], '\n')) - formula = isempty(formula_text) ? nothing : String(formula_text) - - label = get(header, "label", basename(path)) - tier = parse(Int, get(header, "tier", "1")) - status = Symbol(get(header, "status", "open")) - TodoEntry(path, label, tier, status, body, formula) -end - -function _write_todo_file(todo::TodoEntry) - io = IOBuffer() - println(io, "# label: ", todo.label) - println(io, "# tier: ", todo.tier) - println(io, "# status: ", todo.status) - if !isempty(todo.body) - println(io, "#=") - println(io, todo.body) - println(io, "=#") - end - if todo.formula !== nothing && !isempty(todo.formula) - println(io) - print(io, todo.formula) - endswith(todo.formula, "\n") || println(io) - end - write(todo.path, take!(io)) -end - -function _save_todo!(appdata::AppData, label::AbstractString; - new_status::Union{Symbol,Nothing}=nothing, - new_formula::Union{String,Nothing}=nothing) - todo = _find_todo(appdata, label) - todo === nothing && return nothing - updated = TodoEntry( - todo.path, todo.label, todo.tier, - something(new_status, todo.status), - todo.body, - new_formula === nothing ? todo.formula : new_formula, - ) - _write_todo_file(updated) - updated -end - -# One-shot migration: dumps the in-source `_tier1()`/`_tier2()`/`_tier3()` -# seed lists to files on first run. Skips files that already exist, so user -# edits to existing files survive. Once everything is migrated, the in-source -# seed functions are dead code that can eventually be deleted. -function _migrate_todos!() - dir = _todos_dir() - isdir(dir) || mkpath(dir) - for (tier, items) in [(1, _tier1()), (2, _tier2()), (3, _tier3())] - for item in items - label, body, formula = item isa Pair ? - (first(item), last(item), nothing) : - (item.label, item.body, item.formula) - slug = _slug(label) - path = joinpath(dir, slug * ".jl") - isfile(path) && continue - _write_todo_file(TodoEntry(path, label, tier, :open, body, formula)) - end - end -end - -# ── Rendering: one Pico CSS article per TODO with status-colored border ──── - -_TIER_LABELS = ( - "T1", # tier 1 - "T2", # tier 2 - "T3", # tier 3 -) -_TIER_COLORS = ("#4a7c59", "#5a6a8c", "#8c5a5a") - -_tier_pill(tier::Int) = h.span( - get(_TIER_LABELS, tier, "T$tier"); - style="font-size:0.7em;padding:0.1rem 0.4rem;border-radius:1rem;" * - "color:white;background:$(get(_TIER_COLORS, tier, "#888"));" * - "vertical-align:middle;font-weight:normal", -) - -_STATUS_COLORS = ( - open = "#888", - done = "#2e7d32", - deprioritized = "#a05a2c", -) -_status_color(s::Symbol) = get(_STATUS_COLORS, s, "#888") - -function _todo_card(self, todo::TodoEntry) - border_color = _status_color(todo.status) - body_children = Any[HTMXObjects.md_to_node(todo.body)] - if todo.formula !== nothing - push!(body_children, _formula_form(self, todo.label, todo.formula)) - # Inline pipeline-result target — the form's hx_get fills this div - # with `render_pipeline(out)` so the user sees the - # VBRMI/finite-difference output right inside the card. - push!(body_children, h.div(; id="todo-result-$(hash(todo.label))", - style="margin-top:0.5rem")) - end - # `:open` status → expanded; `:done`/`:deprioritized` → collapsed by default. - # Pills sit inside the so they're always reachable, but their - # onclick stops propagation so clicking a pill doesn't also toggle the - # disclosure. - h.article(; - id="todo-card-$(hash(todo.label))", - style="border-left:6px solid $border_color;margin:0.8rem 0;padding:0.5rem 1rem", - )( - h.details(; open=todo.status == :open)( - h.summary(; style="cursor:pointer;list-style-position:outside")( - _tier_pill(todo.tier), " ", - h.strong(todo.label), " ", - _status_pills(todo.label, todo.status), " ", - _permalink(todo), - ), - h.div(; style="margin-top:0.5rem")(body_children...), - ), - ) -end - -_permalink(todo::TodoEntry) = h.a("🔗"; - href="/todo?slug=$(HTTP.URIs.escapeuri(_todo_slug(todo)))", - title="Standalone URL", - onclick="event.stopPropagation()", - style="text-decoration:none;font-size:0.8em;margin-left:0.3rem;vertical-align:middle", -) - -function _formula_form(self, label::String, formula::String) - # Two terminal branches — cimpl (julia/VBRMI benchmark) vs sbimpl (Stan - # source). Each button is `type=button` (not submit) and carries its own - # `hx_get`; `hx_include="closest form"` pulls the textarea + hidden label. - target = "#todo-result-$(hash(label))" - _branch_button(text, stage) = h.button(text; - type="button", - hx_get=string(query_url(self/"stage/$stage"; force=true)), - hx_include="closest form", - hx_target=target, - hx_swap="innerHTML", - style="font-size:0.85em;padding:0.3rem 0.8rem;margin:0.3rem 0.3rem 0 0") - h.form(; style="margin:0.5rem 0")( - h.input(; type="hidden", name="label", value=label), - h.textarea(formula; - name="formula", - rows=max(3, count('\n', formula) + 1), - style="width:100%;font-family:monospace;font-size:0.85em"), - _branch_button("cimpl (bench) ▶", :bench), - _branch_button("sbimpl (Stan) ▶", :stan_code), - ) -end - -# Mutually-exclusive status pills wrapped in a single span so one toggle -# replaces both at once via `outerHTML` targeting `#status-{hash(label)}`. -function _status_pills(label::AbstractString, state::Symbol) - h.span(; id="status-$(hash(label))", - style="margin-left:0.5rem;display:inline-flex;gap:0.3rem;vertical-align:middle")( - _state_pill(label, :done, state, "✓ done", "mark done"), - _state_pill(label, :deprioritized, state, "✓ deprioritized", "deprioritize"), - ) -end - -# Pill text and color reflect the *current* state. When the pill's target_state -# is currently active, it shows the active label (e.g. "✓ done") in the active -# color and clicking it toggles back to :open. Otherwise it shows the inactive -# action label (e.g. "mark done") in gray and clicking it sets the target state. -function _state_pill(label, target_state, current_state, active_text, inactive_text) - is_active = current_state == target_state - text = is_active ? active_text : inactive_text - bg = is_active ? _status_color(target_state) : "#888" - h.button(text; - type="button", - hx_get="/mark?label=$(HTTP.URIs.escapeuri(label))&state=$(target_state)", - hx_target="#todo-card-$(hash(label))", - hx_swap="outerHTML", - onclick="event.stopPropagation()", - style="font-size:0.7em;padding:0.15rem 0.6rem;border:none;border-radius:1rem;" * - "color:white;background:$bg;cursor:pointer", - ) -end - -_tier1() = [ - (label="1.1 verify Bernoulli/Binomial", - body=raw""" -**Status: done.** ✓ Confirmed that the existing `FBroadcasted{<:Type{<:Distribution}}` pass-through in `vimpl.jl` handles both Bernoulli and Binomial cleanly. - -**Verification.** The form below loads the **cbpp + therapeutic touch** model — a faithful translation of `brms::cbpp_binomial` (categorical predictor + random intercept + Binomial with per-row trial counts) and `kruschke::therapeutic_touch` (hierarchical Bernoulli) into one multi-likelihood model. brms's `incidence | trials(size) ~ ...` sidecar collapses to a plain positional argument: `bin_succ ~ Binomial(bin_n, logistic(η))`. If the gradient sanity check stays green every other `Distribution` family (Beta, Gamma, NegBinomial, …) should be free as well. -""", - formula="""log_odds_bin ~ 1 + c1 + (1 | g1) -bin_succ ~ Binomial(bin_n, logistic(log_odds_bin)) - -log_odds_b ~ 1 + (1 | g1) -bin_y ~ Bernoulli(logistic(log_odds_b)) -"""), - - (label="1.2 offset / fixed exposure — already works without a wrapper", - body=raw""" -**Status: already works without any new code.** brms needs `offset(z)` because R's formula syntax has no other way to put a "no-coefficient term" into the linear predictor — the only thing on the RHS of `~` is the formula DSL. In our DSL the linear predictor and the likelihood are *separate* `~` lines, and the second one (the likelihood) takes a free-form Julia expression. Anything inside that expression gets evaluated as plain code at materialization time via `vbroadcasted` — function calls dispatch to whatever Julia function the symbol resolves to, and data column references are pulled from the dataframe. - -So instead of `count ~ x + offset(log(exposure))`, you write the offset directly inside the likelihood: - -```julia -loc ~ 1 + a -k1 ~ Poisson(exp(loc + log(exposure))) -``` - -The `log(exposure)` here is just `Base.log` applied to the `exposure` data column, broadcasted across rows and added to `loc` (which is the materialized linear predictor). No parameter is allocated for it because `growblock!!` is never called for that branch — there's no `~` on the data side, just an argument to `Poisson(...)`. - -**Verification.** Form below loads exactly that model. The VBRMI dim should match the offset-free version (only the population intercept + slope on `a`); the gradient sanity check should stay green; and the materialized `k1` likelihood should incorporate the row-specific exposure shift. -""", - formula="""loc ~ 1 + a -k1 ~ Poisson(exp(loc + log(exposure))) -"""), - - (label="1.3 I(expr) — likely already works", - body=raw""" -**What it is.** brms's `I()` is a literal-escape: `I(x^2)` says "compute `x^2` from the data and treat it as a single column". brms needs it because `+`, `*`, `:`, `|`, … all have special meaning inside an R formula. - -**Why we probably don't need it.** Our DSL is parsed by Julia first, then walked by `_x`. `_x` recursively wraps every `Expr(:call, f, args...)` in an `ExprColumn`, regardless of whether `f` is special. So `loc ~ a + x^2` becomes `+(a, ^(x, 2))` → `ExprColumn(+, NamedColumn(:a), ExprColumn(^, NamedColumn(:x), 2))`. The `^` is just another function call, no special handling needed. - -The only operators that have DSL meaning in our system are `~` (sampling), `=` (assignment), and `|` / `||` inside random-effects specs. Everything else (`^`, `/`, `sqrt`, `log`, `exp`, `mod`, `min`, `max`, …) is a regular function call resolved at materialization time via `vbroadcasted`. - -**Verification.** The form below loads a model with three nonlinear terms (`a^2`, `sqrt(abs(b))`, `log(exposure)`) directly as population-level covariates. The VBRMI dim should match the number of distinct terms; the gradient sanity check should be all-active. If it works, that confirms `I()` is unnecessary because Julia function calls are first-class on the formula RHS. -""", - formula="""loc ~ 1 + a + a^2 + sqrt(abs(b)) + log(exposure) -y1 ~ Normal(loc, 1) -"""), - - "1.4 scale(x) / standardize(x)" => raw""" -**What it is.** brms's `scale(x)` z-transforms a column at parse time: `scale(x) = (x - mean(x)) / std(x)`. The model sees the standardized column. Crucial for default priors (which are scale-invariant only after standardization) and sampler stability (well-conditioned linear predictors). - -**Why it matters.** Most brms vignettes do `scale(x)` automatically as a convenience. Without it, every formula has to either manually z-transform the data or accept poorly-scaled coefficients. - -**Implementation.** -1. Add `function scale end` (and `function center end`, `function standardize end`) to `macro.jl`. -2. Add a `vmeta_sampling_rhs` overload in `vimpl.jl`: -```julia -vmeta_sampling_rhs(meta, x::ExprColumn{typeof(scale)}; group) = begin - inner = vbroadcasted(only(getargs(x)); meta) - materialized = Base.materialize(inner) - z = (materialized .- Statistics.mean(materialized)) ./ Statistics.std(materialized) - vmeta_sampling_rhs(meta, z; group) -end -``` -The standardization happens once when the BRMI is materialized into a VBRMI. Composes with the existing dense-map caching TODO. -3. Add `Statistics` to `vimpl.jl`'s using-list (or vendor `mean`/`std` inline). - -**Verification.** Preset: `loc ~ 1 + scale(a) + scale(b); y1 ~ Normal(loc, 1)`. Compare against the unscaled version: same dim, different posterior geometry. The fitted coefficients should be ≈ the unscaled coefficients × std(x). -""", - "1.5 zerocorr — independent random effects" => raw""" -**What it is.** brms (via lme4 syntax) lets you opt out of the LKJ correlation between multiple random terms in the same group. `(1 + x || group)` (double bar) says "estimate the random intercept and the random slope independently — don't fit a 2×2 Cholesky factor between them". Useful when there isn't enough data to estimate the correlations, or when you have prior reason to believe the terms are uncorrelated. - -**Why it matters.** Multi-term random specs are common, and the LKJ correlation often dominates the prior cost without much identifiability. Letting users skip it is a meaningful sampling speedup and prior simplification. - -**Implementation.** Our `_x` walker already wraps `||` as `ExprColumn{typeof(doublepipe)}`. Add a `vmeta_sampling_rhs` overload that splits each term inside the `||` LHS into its own block (with a synthetic per-term key like `Symbol(group_name, :__nocor__, term_index)`): - -```julia -vmeta_sampling_rhs(meta, x::ExprColumn{typeof(doublepipe)}; kwargs...) = begin - lhs, rhs = getargs(x, 2) - terms = lhs isa ExprColumn{typeof(+)} ? getargs(lhs) : (lhs,) - foldl(enumerate(terms); init=(meta, ())) do (m, args), (i, term) - nocor_key = NamedColumn(Symbol(name(rhs), :__nocor__, i), parent(rhs)) - m, arg = vmeta_sampling_rhs(m, term; group=nocor_key) - m, (args..., arg) - end |> ((m, args),) -> (m, Base.broadcasted(+, args...)) -end -``` - -Each per-term block ends up as 1×1 with one `log_scale` Cholesky parameter — `lprior!`'s existing single-column path handles this with no changes. - -**Verification.** Preset: `loc ~ 1 + (1 + a || g1); y1 ~ Normal(loc, 1)`. Compare its dim against the correlated `(1 + a | g1)` version: the correlated version has 3 Cholesky params (1+2/2 for a 2×2), the uncorrelated version has 2 (one log_scale per term). Same direct-parameter count (2 cols × 8 levels = 16) either way. -""", - "1.6 cache levels / level_map / dense / gc_idx" => raw""" -**What it is.** Stop rebuilding the dense level mapping (`Dict(level => row_index)`) and the gc_idx vector on every `VBRMI(brmi)` call. Cache them once per source data column. - -**Why it matters.** Today every `VBRMI` build re-traces the categorical / grouping columns, sorts unique values, builds a Dict, and walks the column to dense-encode it. For models with many categorical columns or many `VBRMI` rebuilds (e.g. during AD), this adds up. - -**Implementation.** Pick a storage layout for the per-column metadata. Two candidates: -- A new `meta.factor` NamedTuple keyed by source column name, holding `(; levels, level_map, dense, gc_idx)` per column. Built lazily on first reference, indexed via `name(column)`. -- Attach the metadata to `meta.materialized[column_name]` directly. More tightly coupled but avoids a parallel NamedTuple. - -The TODO already lives at `vimpl.jl:78–86`. Once a layout is picked, refactor `_gc_idx` and the inline dense map in the categorical path to read from the cache, falling back to a build-on-miss helper. - -**Verification.** No behavioral change — the gradient sanity check should stay green. Benchmark `VBRMI(brmi)` with Chairmarks before/after and confirm a measurable speedup on a model with multiple categorical/grouping columns. -""", - "1.7 CategoricalArrays / PooledArrays integration" => raw""" -**What it is.** When the input column is already a `CategoricalVector` or a `PooledArray`, the dense level mapping is already computed and stored in the column's `.refs` field. Use it directly instead of rebuilding via `Dict`. - -**Why it matters.** Most real-world DataFrames use `CategoricalArrays.jl` for factor columns. Skipping the rebuild eliminates allocation entirely for the common case and gets us "for free" interop with the standard categorical-data ecosystem. - -**Implementation.** Two design choices: -- **Hard dep**: add `CategoricalArrays` to vimpl.jl's deps, dispatch on `CategoricalVector`, read `levelcode.(col)` and `levels(col)` directly. -- **Duck-typed**: sniff for the `.refs` field and `levels` method without importing the package, falling back to the generic Dict path. - -Recommend hard dep — it's the standard for tabular Julia code, and the duck-type path is more code with no real win. Same for `PooledArrays`. - -The actual integration is small once the design is picked: a method specialization in `_gc_idx` and in the categorical-predictor path. Composes with the caching TODO above. - -**Verification.** Preset (or test) that builds a DataFrame with a `CategoricalVector` column and uses it as a grouping factor / categorical predictor. Confirm the gradient sanity check stays green and the per-VBRMI allocation count drops. -""", -] - -_tier2() = [ - "2.1 interactions a:b, a*b" => raw""" -**What it is.** brms's `a:b` is the elementwise interaction term (a single coefficient multiplying `a[i] * b[i]`). `a*b` is the "main effects + interaction" shorthand: it desugars to `a + b + a:b`. - -**Why it matters.** Interactions are the most commonly missed feature in regression DSLs. Without them, every model that needs `a:b` has to manually create the interaction column in the input DataFrame. - -**Implementation.** -1. **Parser side.** Add a `:` case to `_x` so that `a:b` becomes `ExprColumn(:, NamedColumn(:a), NamedColumn(:b))` instead of falling through to a Symbol/Range parse error. -2. **Materialization side.** Add `vmeta_sampling_rhs(meta, x::ExprColumn{typeof(:)}; group)` that elementwise-multiplies the operands and dispatches to the float-vector path. For continuous × continuous it's a single coefficient on `a .* b`; for categorical × continuous it's `(k-1)` coefficients (one per non-reference level of the categorical, multiplied by the continuous); for categorical × categorical it's `(k₁-1)*(k₂-1)` coefficients via a 2D `_cat_lookup`. -3. **`a*b` desugaring.** At parse time in `_x`, rewrite `*` between formula terms as `+(a, b, :(a:b))`. This needs care because `*` also means multiplication elsewhere (e.g. `Normal(0, 2*sigma)`); the rewrite should only apply at formula-RHS top-level. - -**Verification.** Presets exercising each interaction type: -- continuous×continuous: `loc ~ 1 + a + b + a:b; y1 ~ Normal(loc, 1)` → dim 4 -- continuous×categorical: `loc ~ 1 + a + c1 + a:c1; y1 ~ Normal(loc, 1)` → dim 6 (1 + 1 + 2 + 2) -- categorical×categorical: `loc ~ 1 + c1 + c2 + c1:c2; y1 ~ Normal(loc, 1)` → dim 5 (1 + 2 + 1 + 2) -- shorthand: `loc ~ 1 + a*b; y1 ~ Normal(loc, 1)` should match `loc ~ 1 + a + b + a:b` exactly. -""", - "2.2 configurable categorical reference level" => raw""" -**What it is.** Currently the reference level for treatment-coded categoricals is `sort(unique(x))[1]`. brms / lme4 let you override this via `factor(x, ref="some_level")` or by reordering the factor's levels. - -**Why it matters.** The reference level changes the interpretation of the intercept (it becomes "the mean for the reference level") and of the coefficients (each becomes "the difference from reference"). For some analyses, changing the reference is the only way to make the coefficients directly answer the research question. - -**Implementation.** -1. Add `function factor end` to `macro.jl`. -2. Either store the override at parse time (rewrite `factor(x, ref=:level3)` into a wrapper that the materializer recognizes) or at materialization time via a `meta.factor_ref` NamedTuple keyed by column name. -3. The categorical-predictor path's `levels = sort(unique(x))` becomes `levels = sort(unique(x), by=l -> l == ref ? -Inf : l)` so the chosen reference always sorts first. - -**Verification.** Preset: `loc ~ 1 + factor(c1, ref=2); y1 ~ Normal(loc, 1)`. Compare the fitted coefficients against the default-reference version — they should differ by the level-2-vs-level-1 mean shift but produce the same logdensity. -""", - "2.3 per-parameter prior scales" => raw""" -**What it is.** Currently every parameter is `Normal(0, 1)` in `lprior!`. brms / Stan-style models routinely set custom priors per coefficient: `b ~ Normal(0, 0.5)` for tight priors on slopes, `b ~ Cauchy(0, 1)` for heavy-tailed priors, etc. - -**Why it matters.** Default `Normal(0, 1)` is fine after standardization but poor on raw scales. Allowing per-parameter prior scales is the prerequisite for spike-and-slab, Horseshoe, and most prior sensitivity analyses. Without it, users have no way to express domain knowledge about parameter magnitudes. - -**Implementation.** This is the largest design decision in Tier 2 because it has knock-on effects for every other prior-related TODO. - -Two storage candidates: -- **Per-block scales**: extend `meta.block_data[group]` with a per-column scale vector. `lprior!` multiplies the standard-normal draw by the scale before storing. Simple but only handles Normal-with-scale priors. -- **Per-block prior distributions**: store a vector of `Distribution` objects per block. `lprior!` calls `logpdf(prior_i, xi)` for each parameter. More general; handles Cauchy, StudentT, Horseshoe, etc. - -Recommend the second — it's strictly more powerful and the runtime cost is identical (one `logpdf` call per parameter). Default value is `Normal(0, 1)` for backward compatibility. - -**Implementation sketch.** -1. Extend `meta.block_data` with a `priors` field per block. -2. The macro syntax `b ~ Normal(0, 0.5)` parses as a sampling statement with a Distribution-typed RHS. Currently this is reserved for likelihood declarations; it would need a new "is this a prior or a likelihood?" branch in `vmeta_sampling`. Likelihood: LHS is a data column. Prior: LHS is a maybelocal (parameter). -3. `lprior!` reads the per-column prior and calls `logpdf(prior, value)` instead of the hard-coded `logpdf(Normal(), value)`. - -**Verification.** Preset: `loc ~ 1 + a; b ~ Normal(0, 0.1); y1 ~ Normal(loc, 1)` — confirm the gradient is dampened on `b` compared to the default-prior version, and that the dead-param check still passes. -""", - "2.4 centered / non-centered parameterization toggle" => raw""" -**What it is.** Currently every random-effect block uses non-centered parameterization (we sample standard normals and apply `mul!(vi, C.L, xi)`). brms / Stan let you choose centered (sample directly from `Normal(0, σ)` per group) on a per-factor basis. - -**Why it matters.** Non-centered is the default for "weak data per group" cases (Neal's funnel pathology), but for "strong data per group" cases centered samples better. Letting users choose is a meaningful sampling speedup for the latter regime. - -**Implementation.** Small change to `lprior!` and `growblock!!`. Add a `centered::Bool` flag to `meta.block_data[group]`. In `lprior!`'s non-population branch, if the block is centered, sample directly from `Normal(0, exp(log_scale))` instead of `Normal(0, 1)` then multiplying by `L`. The Cholesky machinery for off-diagonal correlations still applies in the centered case — just on the column before the variance scaling rather than after. - -Once (2.3) is in place, the centered/non-centered choice could be encoded as `(1 | g) ~ Normal(0, σ)` (centered) vs the implicit non-centered default — but for now a per-block kwarg or a wrapper function (e.g. `centered((1 + a | g))`) is simpler. - -**Verification.** Same model with both parameterizations should produce the same logdensity at the same parameter values (after the appropriate change of variables). Sampling efficiency on a known-funnel dataset should differ. -""", - "2.5 grouped random effects (per-factor variance)" => raw""" -**What it is.** Peter's "different variance by diagnosis" pattern: `(1 | subject) gr(diagnosis)` says "the random intercept by subject has a different variance per diagnosis level". In brms this is a custom group structure where the variance hyperparameter itself depends on a second factor. - -**Why it matters.** Common in clinical data where treatment groups have intrinsically different between-subject variability. Without this, you have to fit separate models per diagnosis or accept a single pooled variance. - -**Implementation.** Bigger than it looks because the variance is no longer a single scalar but a length-`n_levels(diagnosis)` vector that needs its own prior and its own gradient. - -Proposed shape: -- A new `gr(group_factor)` wrapper recognized in the `~` RHS via a `function gr end` stub (already exists in `macro.jl`). -- The wrapped block stores `n_levels(group_factor)` log-scale parameters instead of one. `lprior!` walks them, multiplying each subject's random intercept by the diagnosis-specific scale. -- Requires the gc_idx for the inner factor (subject) AND for the outer factor (diagnosis) — both vectors of length N. - -This composes naturally with (1.6 caching) and (2.3 per-parameter priors). - -**Verification.** Preset against synthetic data with two grouping factors, one nested inside the other, with intentionally different per-outer-level variance. Confirm the fitted scales recover the synthetic values. -""", - "2.6 multi-membership random effects mm()" => raw""" -**What it is.** brms's `mm(g1, g2, ...)` lets one observation belong to **multiple** levels of the same random factor simultaneously, with weights summing to 1. Standard use: a student belongs to multiple schools across the year, and we want their random effect to be a weighted average of the per-school effects. - -**Why it matters.** Standard random effects assume each observation belongs to exactly one group. Multi-membership is the only clean way to handle observations that span groups (mobile students, patients seen by multiple clinicians, etc.). - -**Implementation.** `_gc_idx` would have to return a row-of-vectors instead of a single Int per row. Two paths: -- **Sparse design matrix**: replace the `gc_idx` lookup with a sparse `(N × n_levels)` matrix where each row's nonzero entries are the membership weights. The materialized random effect becomes `sparse_membership * random_effects_vector`. -- **Per-row lookup loop**: keep the row-major view but make `_re_lookup` iterate the membership list per row, summing weighted contributions. - -The sparse matrix approach is more memory-efficient and SIMD-friendly. Needs a new wrapper in the formula syntax: `(1 | mm(g1, g2; weights=...))`. - -**Verification.** Preset against synthetic data where each observation has 2 random group memberships with weights summing to 1. Compare against the equivalent "fully observed in primary group only" model. -""", - "2.7 se() / weights() for meta-analysis and weighted regression" => raw""" -**What it is.** brms's `y | se(sigma_y) ~ ...` lets each observation have its own known standard error (typical for meta-analysis where each `y` is itself a summary estimate). `y | weights(w) ~ ...` is observation-level weighting (typical for survey data or sample-size correction). - -**Why it matters.** Both are extremely common in applied work. Without them, meta-analysis can't be expressed at all in this DSL, and weighted regression has to be hacked via likelihood multiplication. - -**Implementation.** Both are sidecar modifiers on the LHS of `~`, so they need parser support similar to brms's `|` syntax. Or, more naturally for our DSL: pass them as positional arguments to the distribution itself. -- `y ~ Normal(mu, se_y)` for the meta-analysis case (already works! `se_y` is just another data column). -- For weights, define a `weighted` likelihood wrapper: `y ~ weighted(Normal(mu, sigma), w)` where `weights` multiplies the per-row logpdf by `w[i]`. Needs a new `vmeta_sampling_rhs` overload and a new `LikelihoodColumn`-like type with a per-row weight. - -Meta-analysis is essentially free (already works). Weights need ~15 lines. - -**Verification.** Preset for meta-analysis: `y ~ Normal(mu, se_y)` with `se_y` from synthetic data. Preset for weights: `y ~ weighted(Normal(mu, 1), weight); loc ~ 1 + a` confirming the gradient is rescaled per-row by `weight`. -""", -] - -_tier3() = [ - "3.1 multivariate outcomes cbind(y1, y2)" => raw""" -**What it is.** brms's `cbind(y1, y2) ~ x + (1 | g)` declares that `y1` and `y2` share the same linear predictor structure but have correlated residuals. The likelihood becomes multivariate normal (or multivariate-t) over `(y1, y2)` with a covariance matrix to estimate. - -**Why it matters.** Unblocks `mcelreath::waffle_divorce_multivariate` and any model where multiple outcomes share latent structure (joint pharmacology/efficacy, paired outcomes, mediation analysis). - -**Implementation.** Real new infrastructure: -1. New parser support for `cbind(...)` on the LHS of `~`. -2. New `LikelihoodColumn`-like type that holds a tuple of data columns and a multivariate distribution. -3. `llikelihood!` calls `logpdf(MvNormal(loc_vec, Σ), [y1[i], y2[i]])` per row. -4. `Σ` is a new parameter block: a Cholesky factor over the outcomes (separate from the random-effects Cholesky). - -**Verification.** Translate `mcelreath::waffle_divorce_multivariate` directly. Compare fitted parameters against the published reference. -""", - "3.2 inferred predictors / measurement error me()" => raw""" -**What it is.** brms's `me(x_obs, sd_x)` says "the predictor `x_obs` is itself measured with error of size `sd_x`; sample the latent true value during inference". The model sees both the observed value and the latent. - -**Why it matters.** Standard regression treats predictors as fixed/known. When predictors are themselves estimates (e.g. from a previous study or a noisy sensor), ignoring measurement error biases the slope estimates toward zero. `me()` is the principled fix. - -**Implementation.** Bigger architectural change: predictor columns become latent variables sampled during inference, not data columns evaluated once. Needs: -- A new column type analogous to `MissingColumn` but with an observation-driven prior `Normal(x_obs, sd_x)`. -- The latent column gets a slot in the population block (one parameter per row). -- `lprior!` adds the per-row Normal prior contribution. -- `vbroadcasted` resolves the column to the latent values, not the observed ones. - -**Verification.** Preset against synthetic data where the true `x` is known but only a noisy observed version is in the dataframe. Compare slope estimates with and without `me()`. -""", - "3.3 ordinal predictors mo() (monotonic effects)" => raw""" -**What it is.** brms's `mo(x)` for an ordinal predictor with K levels: instead of `K-1` independent treatment-coded coefficients, fit a single "total effect" β plus a `K-1`-dim simplex of inter-level shape. Forces the effect to be monotonic in the ordering of `x`'s levels. - -**Why it matters.** Likert-scale predictors and ordered categorical inputs (e.g. age groups) have a natural ordering that treatment coding ignores. `mo()` enforces the monotonicity prior, dramatically reducing the parameter count and tightening posterior inference. - -**Implementation.** New block layout: one β coefficient + one Dirichlet-distributed simplex of length `K-1`. `_cat_lookup`-style materialization but the per-level contribution is `β * cumulative_sum(simplex)[level]` instead of `coefficients[level]`. - -Needs a new prior block type (Dirichlet) in `lprior!`, plus new parser support for `mo(x)` and a new `vmeta_sampling_rhs` overload. - -**Verification.** Preset against synthetic data where the true effect is monotonic but the levels are unordered in the data. Compare fitted shape parameters against the synthetic monotonic curve. -""", - "3.4 ordinal outcomes (proportional odds)" => raw""" -**What it is.** When `y` is itself ordered categorical (Likert response, severity grades, …), use a cumulative-link model: `Pr(y ≤ k) = logistic(α_k - η)` where `α_k` are K-1 cutpoints and η is the linear predictor. The likelihood is the difference of consecutive CDFs. - -**Why it matters.** Ordinal outcomes are common in survey data, clinical scoring, and any "rating" task. Treating them as continuous is statistically wrong; treating them as nominal categorical loses the ordering information. - -**Implementation.** New likelihood family with a vector of cutpoints as additional parameters. `Distributions.jl` has `OrderedLogistic` already — the pass-through path should mostly handle it once the parser knows to extract cutpoints from a `cumulative` wrapper. - -The cutpoints need an ordered prior (e.g. ordered transform of unconstrained reals), which means a new prior block type — similar to (3.3)'s simplex. - -**Verification.** Preset against synthetic Likert data. Compare cutpoints against `polr` from R's MASS package. -""", - "3.5 mixture models" => raw""" -**What it is.** brms's `mixture(Normal, Normal)` lets the likelihood be a weighted mixture of K component distributions, with mixing weights estimated as parameters. - -**Why it matters.** Heterogeneous populations, latent class analysis, robust regression (Normal + heavy-tailed component), zero-inflated outcomes, … - -**Implementation.** Extends the existing `Distribution` pass-through. Need a new `MixtureModel` wrapper that holds component distributions plus a weights parameter block. `llikelihood!` uses `logsumexp(log_weights .+ logpdf.(components, y))` per row. - -Composes with (3.4 ordered priors) for the mixing weights' Dirichlet-like prior. - -**Verification.** Preset against synthetic two-component-Normal data. Confirm the recovered mixture weights and component parameters. -""", - "3.6 splines / GP submodels s(), bs(), gp(), t2()" => raw""" -**What it is.** Smoothers in the linear predictor: `s(x)` for a generic spline, `bs(x, knots=...)` for a B-spline basis, `gp(x)` for a Gaussian process, `t2(x, y)` for a tensor-product spline. brms / mgcv use these heavily. - -**Why it matters.** Nonlinear effects without committing to a specific functional form. The de facto way to model dose-response curves, time effects, growth curves, spatial trends, … - -**Implementation.** Each smoother is a basis-matrix builder that grows a population block by `n_basis` columns and stores the basis matrix as part of `meta`. Function stubs (`s`, `bs`, `t2`, `gp`) already exist in `scripts/parsing.jl` so the parser side is partly done. - -For each smoother type, the materialization is `basis_matrix * coefficients` (length-N output). The smoothness prior is a structured prior on the coefficients (typically a Gaussian prior with a banded or 2D-difference penalty matrix), which requires (2.3 per-parameter priors) as a prerequisite. - -**Verification.** Preset against synthetic curve data. Compare fitted smoother against `mgcv::gam`. -""", - "3.7 autoregressive submodels ar(), ar1()" => raw""" -**What it is.** Add an AR(p) structure to the residuals: `y_t = η_t + φ * (y_{t-1} - η_{t-1}) + ε_t`. brms's `ar(time, p=1)` specifies the order and the time variable. - -**Why it matters.** Time-series and repeated-measures data routinely have autocorrelated errors. Ignoring AR structure inflates the effective sample size and gives overconfident posteriors. - -**Implementation.** Bigger than it looks because the likelihood is no longer per-row independent — it's a chain. Need a new `LikelihoodColumn`-like type that holds the time index and walks the data in time order, accumulating the AR contribution row by row. - -Composes with (2.5 grouped random effects) for per-subject AR structure. - -**Verification.** Preset against synthetic AR(1) data. Recover φ. -""", - "3.8 decompositions (QR, orthogonal polar)" => raw""" -**What it is.** Numerical-stability transformations of the population design matrix. brms uses QR decomposition on the design matrix internally so the sampler sees an orthogonal-columns version, then transforms back at the end. Stan does the same. - -**Why it matters.** When population covariates are correlated (which is the norm), the unrotated design matrix gives a poorly-conditioned posterior that NUTS struggles with. QR fixes this with no statistical change. - -**Implementation.** Mostly orthogonal to the formula DSL — happens at `VBRMI` build time. Add a per-block transform: store both the original design matrix and the QR factor, run sampling on the rotated parameter space, transform back when extracting coefficients. - -Could be implemented today without any parser changes, as a `qr_transform=true` flag on `VBRMI`. Lift to a default-on once verified. - -**Verification.** Same model with and without QR should produce identical logdensity values, but the gradient should be better-conditioned (smaller condition number on the Hessian). -""", - "3.9 spike-and-slab / Horseshoe priors" => raw""" -**What it is.** Sparsity-inducing priors for high-dimensional regression. Spike-and-slab puts a delta-spike at zero plus a wide slab; Horseshoe uses a half-Cauchy hyperprior on a per-coefficient scale, producing a heavy-tailed shrinkage prior. - -**Why it matters.** Without sparsity priors, high-dimensional regressions overfit. These are the standard solution in Bayesian variable selection and high-dim genomics / finance. - -**Implementation.** Depends entirely on (2.3 per-parameter priors). Once that's in place, spike-and-slab is `prior = Mixture(Normal(0, ε), Normal(0, slab_scale))` per coefficient, and Horseshoe is `Normal(0, λ_i * τ)` with `λ_i ~ HalfCauchy(0, 1)` and `τ ~ HalfCauchy(0, 1)` — both expressible in the existing Distribution pass-through once per-coefficient priors are wired up. - -**Verification.** Preset on a sparse synthetic regression (mostly-zero true coefficients with a few large ones). Confirm the Horseshoe-fitted coefficients shrink the noise toward zero and preserve the signal. -""", - "3.10 Dirichlet process / non-parametric models" => raw""" -**What it is.** Models where the number of components / clusters / random-effect levels is itself inferred during sampling, via a Dirichlet process or stick-breaking prior. - -**Why it matters.** When you don't know how many clusters are in your data, fixing K is itself a strong assumption. DP priors let the model decide. - -**Implementation.** The heaviest item on the list. Needs sampling-time level inference, a stick-breaking parameter block, and a different `growblock!!` that grows during sampling rather than at `VBRMI` build time. Probably requires a different `lprior!` interface entirely. - -Defer until everything else is solid. - -**Verification.** Preset against synthetic data with an unknown number of latent clusters. Compare recovered K against the truth. -""", - "3.11 zero-inflated / hurdle likelihoods" => raw""" -**What it is.** ZI Poisson, ZI Negative Binomial, hurdle Poisson, hurdle Gamma, … — likelihoods that mix a point mass at zero (or a separate "is zero" Bernoulli) with a continuous/count distribution for the nonzero values. - -**Why it matters.** Count data with excess zeros (insurance claims, species abundance, healthcare utilization) is everywhere. Standard Poisson / NegBin underfits the zero count. - -**Implementation.** Should mostly work through the existing `Distribution` pass-through once we use `Distributions.jl`'s ZI distributions (or write small wrappers). The mixing weight needs its own linear predictor, which is just another distributional-regression-style `~` line. - -**Verification.** Preset against synthetic ZI Poisson data. Confirm the recovered zero-inflation probability matches the synthetic generator. -""", -] - -const APPDATA = AppData(; cache_type=:parallel) - function __init__() route!(AppContext()) end -# Bruno-specific extensions (gitignored); load if present. +# Bruno-specific extensions (gitignored); load if present. Adds a +# `dataset_extras(::Val{:bruno}, df)` method and optionally more. let path = joinpath(@__DIR__, "bruno-ext.jl") isfile(path) && include(path) end diff --git a/web-macro/src/brm-macro.css b/web-macro/src/brm-macro.css new file mode 100644 index 0000000..de6a4fb --- /dev/null +++ b/web-macro/src/brm-macro.css @@ -0,0 +1,56 @@ +:root { font-size: 87.5%; } + +/* Layout */ +.brm-layout { display: flex; gap: 1rem; align-items: flex-start; } +.brm-main { flex: 1; min-width: 0; } + +/* Pipeline form */ +.brm-preset-row { display: flex; flex-wrap: wrap; gap: 0.3rem; margin-bottom: 0.6rem; } +.brm-preset-btn { font-size: 0.8em; padding: 0.2rem 0.5rem; margin: 0; } +.brm-formula-textarea { width: 100%; font-family: monospace; } + +/* Pipeline status line */ +.brm-status-ok { color: green; } +.brm-status-err { color: crimson; } + +/* BRMI / VBRMI cards */ +.brm-card { margin: 0.5rem 0; } +.brm-expr-list { font-family: monospace; font-size: 1.15em; line-height: 1.8; padding: 0.3rem 0; } +.brm-indent { margin-left: 1rem; } +.brm-subhead { margin-bottom: 0.2rem; } +.brm-block-part { color: #999; margin-left: 1.5rem; } + +/* Styled AST rendering (per-symbol color is inline) */ +.brm-sym-param { font-weight: bold; } +.brm-num { color: #666; } +.brm-muted { color: #999; } +.brm-op { color: #555; } +.brm-op-dark { color: #333; } +.brm-fname { color: #777; } +.brm-shape { color: #999; font-size: 0.85em; } +.brm-arr { font-style: italic; color: #888; } +.brm-vec { font-weight: bold; color: #888; } +.brm-likelihood { text-decoration: underline; text-decoration-color: #aaa; text-underline-offset: 3px; } + +/* Example cards (border-left color, tier/status pill bg are inline) */ +.brm-example-card { border-left: 6px solid #888; margin: 0.8rem 0; padding: 0.5rem 1rem; } +.brm-example-summary { cursor: pointer; list-style-position: outside; } +.brm-example-body { margin-top: 0.5rem; } +.brm-example-result { margin-top: 0.5rem; } +.brm-example-form { margin: 0.5rem 0; } +.brm-example-textarea { width: 100%; font-family: monospace; font-size: 0.85em; } +.brm-branch-btn { font-size: 0.85em; padding: 0.3rem 0.8rem; margin: 0.3rem 0.3rem 0 0; } +.brm-permalink { text-decoration: none; font-size: 0.8em; margin-left: 0.3rem; vertical-align: middle; } + +/* Pills */ +.brm-tier-pill { + font-size: 0.7em; padding: 0.1rem 0.4rem; border-radius: 1rem; + color: white; vertical-align: middle; font-weight: normal; +} +.brm-status-pills { + margin-left: 0.5rem; display: inline-flex; gap: 0.3rem; vertical-align: middle; +} +.brm-state-pill { + font-size: 0.7em; padding: 0.15rem 0.6rem; border: none; border-radius: 1rem; + color: white; cursor: pointer; +} diff --git a/web-macro/src/html_expr.jl b/web-macro/src/html_expr.jl new file mode 100644 index 0000000..e1a1ddf --- /dev/null +++ b/web-macro/src/html_expr.jl @@ -0,0 +1,176 @@ +# Styled HTML rendering for BRMI / VBRMI cards. +# +# Each symbol gets a deterministic color, data columns are bold, parameters +# are italic, and likelihood statements are underlined. The tree walker +# (_html_expr) converts an ExprColumn AST into a nest of s. +# +# Eventually this belongs in an ext of the main package (rendering depends on +# HTMX.jl which is a web-only concern); keeping it in its own file makes that +# future move a single-file relocation. +# +# TODO: HTMX.jl should grow a generic `htmx_node(x)::Node` extension point +# so downstream packages can overload once and have every consumer dispatch +# automatically. For now these card functions are wired in by hand. + +# Deterministic HSL color per symbol (golden-ratio spread for visual variety). +_symbol_color(name::Symbol) = "hsl($(mod(hash(name) * 137, 360)), 60%, 40%)" + +# A colored with role-based font styling. Data columns are normal +# weight; parameters (latent/sampled) are bold. Per-symbol color is the only +# inline style -- it is data-derived, not presentational. +_styled_name(name::Symbol, role::Symbol) = h.span(string(name); + class=role == :parameter ? "brm-sym brm-sym-param" : "brm-sym", + style="color:$(_symbol_color(name))") + +_html_expr(x::NamedColumn{<:Any, <:DataColumn}) = _styled_name(name(x), :data) +_html_expr(x::NamedColumn{<:Any, MissingColumn}) = _styled_name(name(x), :parameter) +_html_expr(x::NamedColumn) = _styled_name(name(x), :derived) +_html_expr(x::Int) = h.span(string(x); class="brm-num") +_html_expr(x::Float64) = h.span(string(x); class="brm-num") +_html_expr(x::Number) = h.span(string(x); class="brm-num") +_html_expr(x::DataColumn) = h.span("data($(eltype(parent(x))))"; class="brm-muted") +_html_expr(x::MaterializedColumn) = _html_expr(getbroadcast(x)) +_html_expr(x::LikelihoodColumn) = h.span( + _html_expr(parent(x)), h.span(" .~ "; class="brm-op-dark"), _html_expr(rhs(x))) + +# Infix operators: always parenthesized so inner expressions like (1 + b | g1) +# keep their grouping. Top-level callers (_html_brmi_row) use _html_infix +# directly to skip the outermost parens. +_html_expr(x::ExprColumn{<:Union{typeof.((~,*,+,|,doublepipe,assign))...}}) = begin + h.span("(", _html_infix(x), ")") +end + +_html_infix(x::ExprColumn) = begin + op_str = " $(getop(x)) " + args = getargs(x) + parts = Any[] + for (i, arg) in enumerate(args) + i > 1 && push!(parts, h.span(op_str; class="brm-op")) + push!(parts, _html_expr(arg)) + end + h.span(parts...) +end + +# Function-call style: fname(args...; kwargs...) +_html_expr(x::ExprColumn) = begin + fname = getf(x) isa Function ? nameof(getf(x)) : + getf(x) isa Type ? nameof(getf(x)) : string(getf(x)) + args = getargs(x) + kw = getkwargs(x) + parts = Any[h.span(string(fname); class="brm-fname"), "("] + for (i, arg) in enumerate(args) + i > 1 && push!(parts, ", ") + push!(parts, _html_expr(arg)) + end + if length(kw) > 0 + push!(parts, "; ") + for (i, (k, v)) in enumerate(pairs(kw)) + i > 1 && push!(parts, ", ") + push!(parts, "$k=", _html_expr(v)) + end + end + push!(parts, ")") + h.span(parts...) +end + +# Broadcasted objects (from VBRMI materialization): walk their inner structure +_html_expr(x::Base.Broadcast.Broadcasted) = begin + fname = x.f isa Function ? nameof(x.f) : + x.f isa Type ? nameof(x.f) : string(x.f) + args = x.args + if x.f in (+, -, *, /) + parts = Any[] + for (i, arg) in enumerate(args) + i > 1 && push!(parts, h.span(" $(x.f) "; class="brm-op")) + push!(parts, _html_expr(arg)) + end + return h.span(parts...) + end + parts = Any[h.span(string(fname); class="brm-fname"), "("] + for (i, arg) in enumerate(args) + i > 1 && push!(parts, ", ") + push!(parts, _html_expr(arg)) + end + push!(parts, ")") + h.span(parts...) +end + +# Arrays / views from block parameter slots: show as a compact shape description +_html_expr(x::SubArray) = h.span("param[$(join(size(x), "×"))]"; class="brm-arr") +_html_expr(x::AbstractVector{<:Number}) = h.span("vec[$(length(x))]"; class="brm-vec") +_html_expr(x::Base.RefValue) = _html_expr(x[]) +_html_expr(x::AbstractMatrix) = h.span("mat[$(join(size(x), "×"))]"; class="brm-arr") + +_html_expr(x) = h.span(sprint(show, x; context=:compact=>true); class="brm-muted") + +function brmi_card(brmi::BRMI) + rows = [_html_brmi_row(key, parent(value)) for (key, value) in pairs(brmi.operations)] + h.article(; class="brm-card")( + h.header(h.strong("BRMI"), + h.small(" -- $(length(brmi.operations)) operations")), + h.div(; class="brm-expr-list")(rows...), + ) +end + +_leaf_column(x::NamedColumn) = x +_leaf_column(x::ExprColumn) = _leaf_column(getargs(x)[1]) +_leaf_column(x) = x + +function _html_brmi_row(key, op::ExprColumn{typeof(~)}) + lhs_leaf = _leaf_column(getargs(op)[1]) + is_likelihood = lhs_leaf isa NamedColumn && parent(lhs_leaf) isa DataColumn + content = _html_infix(op) + h.div(content; class=is_likelihood ? "brm-likelihood" : "") +end +_html_brmi_row(key, op::ExprColumn{typeof(assign)}) = h.div(_html_infix(op)) +_html_brmi_row(key, op) = h.div( + _styled_name(key, :data), + h.span(": "; class="brm-num"), + h.span(sprint(show, op); class="brm-muted"), +) + +function vbrmi_card(vbrmi::VBRMI) + brmi = getfield(vbrmi, :parent) + (; meta) = vbrmi + n_dim = LogDensityProblems.dimension(vbrmi) + n_mat = length(meta.materialized) + n_blocks = length(meta.blocks) + + mat_rows = [begin + nc = get(brmi.operations, key, nothing) + inner = nc !== nothing ? parent(nc) : nothing + if inner isa DataColumn + h.div(_styled_name(key, :data), + h.span(": data($(eltype(parent(inner))))"; class="brm-muted")) + elseif value isa LikelihoodColumn + expr = inner !== nothing ? _html_infix(inner) : _styled_name(key, :data) + h.div(expr; class="brm-likelihood") + elseif inner !== nothing + expr = inner isa ExprColumn{<:Union{typeof(~),typeof(assign)}} ? + _html_infix(inner) : _html_expr(inner) + shape = "$(eltype(parent(value)))[$(length(parent(value)))]" + h.div(expr, h.span(" -> $shape"; class="brm-shape")) + else + h.div(_styled_name(key, :derived), + h.span(": $(sprint(show, value))"; class="brm-muted")) + end + end for (key, value) in pairs(meta.materialized)] + + blocks_rows = [begin + role = key === :__population__ ? :derived : :data + h.div( + _styled_name(key, role), + [h.div(; class="brm-block-part")(sprint(show, part)) + for part in parts]..., + ) + end for (key, parts) in pairs(meta.blocks)] + + h.article(; class="brm-card")( + h.header(h.strong("VBRMI"), + h.small(" -- dim $n_dim, $n_mat materialized, $n_blocks blocks")), + h.h6(; class="brm-subhead")("materialized"), + h.div(; class="brm-expr-list brm-indent")(mat_rows...), + h.h6(; class="brm-subhead")("blocks"), + h.div(; class="brm-expr-list brm-indent")(blocks_rows...), + ) +end From 7194a58f41730b2668f400684bb873d5cd10ec06 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Thu, 23 Apr 2026 00:53:58 +0200 Subject: [PATCH 18/23] web-macro: rename DSL terms `I` -> `protect`, `scale` -> `zscale` Avoid collisions with `LinearAlgebra.I` and `Distributions.scale` imports in vimpl. `protect` follows GLM.jl naming; `zscale` clarifies the z-transform semantics that `standardize` already aliases. Co-Authored-By: Claude Opus 4.7 --- .../1.3-i-expr-likely-already-works.jl | 18 +++++++++--------- .../examples/1.4-scale-x-standardize-x.jl | 14 +++++++------- web-macro/src/macro.jl | 4 ++-- web-macro/src/vimpl.jl | 15 ++++++++------- 4 files changed, 26 insertions(+), 25 deletions(-) diff --git a/web-macro/examples/1.3-i-expr-likely-already-works.jl b/web-macro/examples/1.3-i-expr-likely-already-works.jl index 95cdf31..136e519 100644 --- a/web-macro/examples/1.3-i-expr-likely-already-works.jl +++ b/web-macro/examples/1.3-i-expr-likely-already-works.jl @@ -1,24 +1,24 @@ -# label: 1.3 I(expr) — brms literal-escape +# label: 1.3 protect(expr) — brms literal-escape # tier: 1 # status: done (vimpl) #= **Status: done in vimpl.** Our DSL already supports arbitrary Julia function -calls on the RHS (e.g. `a^2`, `sqrt(abs(b))`, `log(exposure)`), so `I()` is -functionally unnecessary. But brms users who have internalized the `I()` +calls on the RHS (e.g. `a^2`, `sqrt(abs(b))`, `log(exposure)`), so `protect()` is +functionally unnecessary. But brms users who have internalized the `protect()` convention can now write it literally — we added a passthrough overload -(`vbroadcasted(::ExprColumn{typeof(I)}) = inner`) so `I(expr)` behaves the same +(`vbroadcasted(::ExprColumn{typeof(protect)}) = inner`) so `protect(expr)` behaves the same as `expr` alone, no parameter penalty. -**Why brms needs I().** In R, `a:b`, `a*b`, `|` etc. have DSL meaning inside a -formula — `I()` escapes them so the inner expression is interpreted as plain +**Why brms needs protect().** In R, `a:b`, `a*b`, `|` etc. have DSL meaning inside a +formula — `protect()` escapes them so the inner expression is interpreted as plain arithmetic. Our formula DSL is parsed by Julia first, so `a^2` is already a -plain function call and `I()` is a no-op. +plain function call and `protect()` is a no-op. **Verification.** The form below mixes the "naked" spelling (`a^2`, -`sqrt(abs(b))`, `log(exposure)`) with I()-wrapped forms — the VBRMI dim, +`sqrt(abs(b))`, `log(exposure)`) with protect()-wrapped forms — the VBRMI dim, gradient check, and materialized predictor values must be identical. =# -loc ~ 1 + a + I(a^2) + I(sqrt(abs(b))) + log(exposure) +loc ~ 1 + a + protect(a^2) + protect(sqrt(abs(b))) + log(exposure) y1 ~ Normal(loc, 1) diff --git a/web-macro/examples/1.4-scale-x-standardize-x.jl b/web-macro/examples/1.4-scale-x-standardize-x.jl index 2f6f9bc..3e22008 100644 --- a/web-macro/examples/1.4-scale-x-standardize-x.jl +++ b/web-macro/examples/1.4-scale-x-standardize-x.jl @@ -1,21 +1,21 @@ -# label: 1.4 scale(x) / center(x) / standardize(x) +# label: 1.4 zscale(x) / center(x) / standardize(x) # tier: 1 # status: done (vimpl) #= **Status: done in vimpl.** Three data-transform wrappers that z-transform (or -just center) a column at VBRMI-materialization time. brms does `scale(x)` +just center) a column at VBRMI-materialization time. brms does `zscale(x)` automatically in most vignettes for sampler stability + prior scale-invariance; making it available as a formula-level wrapper avoids forcing users to pre-transform their DataFrame columns. **Semantics.** - `center(x)` -> `x - mean(x)` -- `scale(x)` -> `(x - mean(x)) / std(x)` -- `standardize(x)` -> alias for `scale(x)` +- `zscale(x)` -> `(x - mean(x)) / std(x)` +- `standardize(x)` -> alias for `zscale(x)` Each fires inside `vbroadcasted`, so transforms compose with every downstream consumer: pop predictor, ranef slope, link functions, interactions, etc. -`scale(a):scale(b)` works as expected (z-transformed both operands before the +`zscale(a):zscale(b)` works as expected (z-transformed both operands before the elementwise product). **Implementation.** Three-line `vbroadcasted` overloads in vimpl.jl using @@ -23,7 +23,7 @@ inlined `_mean`/`_std` helpers (no new dependency). The transform evaluates once at VBRMI construction and caches the resulting vector. **Verification.** Same model spelled two ways: -- `loc ~ 1 + scale(a) + scale(b)` -- standardized at formula level +- `loc ~ 1 + zscale(a) + zscale(b)` -- standardized at formula level - `loc ~ 1 + a_z + b_z` (where a_z, b_z were pre-standardized in the DataFrame) VBRMI dim + log-density should be identical; fitted coefficients ≈ the @@ -31,5 +31,5 @@ unscaled coefficients × `std(x)`. =# -loc ~ 1 + scale(a) + center(b) +loc ~ 1 + zscale(a) + center(b) y1 ~ Normal(loc, 1) diff --git a/web-macro/src/macro.jl b/web-macro/src/macro.jl index 3f6d2d0..151c86d 100644 --- a/web-macro/src/macro.jl +++ b/web-macro/src/macro.jl @@ -47,10 +47,10 @@ function doublepipe end function gr end function gp end function offset end -function scale end +function zscale end function center end function standardize end -function I end +function protect end _brm(x::AbstractString; kwargs...) = _brm(Meta.parse(""" begin $x diff --git a/web-macro/src/vimpl.jl b/web-macro/src/vimpl.jl index 5561760..eb5205f 100644 --- a/web-macro/src/vimpl.jl +++ b/web-macro/src/vimpl.jl @@ -88,10 +88,11 @@ vbroadcasted(;kwargs...) = (args...)->vbroadcasted(args...; kwargs...) vbroadcasted(x::NamedColumn{<:Any,<:DataColumn}; meta) = parent(meta.materialized[name(x)]) vbroadcasted(x::NamedColumn; meta) = meta.materialized[name(x)] vbroadcasted(x::ExprColumn; meta) = Base.broadcasted(getf(x), map(vbroadcasted(;meta), getargs(x))...) -# `I(expr)` is brms's literal-escape. In our DSL every call is already a first- -# class `ExprColumn` node, so `I` just needs to unwrap to its inner argument. -vbroadcasted(x::ExprColumn{typeof(I)}; meta) = vbroadcasted(only(getargs(x)); meta) -# `scale(x)` / `center(x)` / `standardize(x)` z-transform the inner column at +# `protect(expr)` is brms's `I()` literal-escape. In our DSL every call is +# already a first-class `ExprColumn` node, so `protect` just needs to unwrap to +# its inner argument. +vbroadcasted(x::ExprColumn{typeof(protect)}; meta) = vbroadcasted(only(getargs(x)); meta) +# `zscale(x)` / `center(x)` / `standardize(x)` z-transform the inner column at # VBRMI-materialization time: they materialize the inner broadcast once, apply # the transform, and pass the resulting plain vector back up to the predictor # pipeline. Because this fires inside `vbroadcasted`, they compose with every @@ -99,12 +100,12 @@ vbroadcasted(x::ExprColumn{typeof(I)}; meta) = vbroadcasted(only(getargs(x)); me vbroadcasted(x::ExprColumn{typeof(center)}; meta) = let raw = Base.materialize(vbroadcasted(only(getargs(x)); meta)) raw .- _mean(raw) end -vbroadcasted(x::ExprColumn{typeof(scale)}; meta) = let raw = Base.materialize(vbroadcasted(only(getargs(x)); meta)) +vbroadcasted(x::ExprColumn{typeof(zscale)}; meta) = let raw = Base.materialize(vbroadcasted(only(getargs(x)); meta)) mu = _mean(raw); sd = _std(raw, mu) - sd > 0 || error("scale: zero variance in `$(only(getargs(x)))`") + sd > 0 || error("zscale: zero variance in `$(only(getargs(x)))`") (raw .- mu) ./ sd end -vbroadcasted(x::ExprColumn{typeof(standardize)}; meta) = vbroadcasted(ExprColumn(scale, getargs(x)...); meta) +vbroadcasted(x::ExprColumn{typeof(standardize)}; meta) = vbroadcasted(ExprColumn(zscale, getargs(x)...); meta) _mean(xs) = sum(xs) / length(xs) _std(xs, mu=_mean(xs)) = sqrt(sum(abs2, xs .- mu) / (length(xs) - 1)) From a34d8bbc36e1bb5781a9ed35fd6fa6360d743264 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Thu, 23 Apr 2026 00:54:38 +0200 Subject: [PATCH 19/23] web-macro: Stan fit pipeline (SBC) + shared plot tabsets + truth overlays - Stan path: param_constrain(generated_unc) -> (long, wide, summary) DF triple, shared with Pathfinder/warmup fits via `constrain_draws` + `dfs_from_constrained` helpers (deduplicates ~60 lines of inline DataFrame plumbing across three callsites). - SBC setup: take one draw from the prior-predictive `generated` matrix, fold `*_gen` entries back into the Stan data dict as `*` (observed), and fit `fit_instance = StanModel(lib, bridgestan_data( fit_data_dict))`. Pathfinder / full warmup then sample `p(theta | y_sim)`; ground-truth `fit_truth_unc` stays available. - Treebars progress nesting: `pathfinder` / `posterior_warmup` are IPs that pass `progress=__status__` to WarmupHMC; `compute_steps` calls `fetchindex!(progress, ip, args...)` so the IP's substatus attaches under the step's phase. Kwargs (rng, maxiters, n_draws) supported natively on IP signatures. - StanProblem wrap: BridgeStan.StanModel lacks `logdensity_and_gradient`; IPs wrap with `StanLogDensityProblems.StanProblem(instance)` at call. - Shared `posterior_plots(long, wide, summary; id_prefix, kind, truth)` builds the 8-tab PI / LR / ECDF / Hist tabset. Used by `stan_generate`, `stan_fit_pathfinder`, `stan_fit_warmup`. `plot_fit` removed (redundant with the two fit steps). - Pre-aggregated summary path: `bands=[:q025=>:q975, :q10=>:q90, :q25=>:q75]` feeds `pointinterval(bands=..., orientation=:vertical)` and `lineribbon(bands=...)`; `Statistics.quantile/median` compute once per (param, index) group. - Truth overlays: `truth_df` (fit_draw_idx column constrained, one row per (param, index)) layers black `Scatter` over PI/LR and index-colored `VLines` over ECDF/Hist. Same overlay on prior-predictive and fit plots so user can see which draw was picked and how well the fit recovers it. - PI/LR share identical mapping `(:index, :median, row=:param)` with `indep_y`; ECDF/Hist share `(:value; row=:param, color=:index)` with `indep_x`. Picker tabs wrap each spec via `with_plot_caption`; PI/LR pickers list only `:param` in dims to avoid the pinned catch-all combo. - Project deps: +AlgebraOfGraphics, +AlgebraOfVega, +JSON, +Statistics (web-macro); +AlgebraOfVega, +WarmupHMC (web-macro/app). Co-Authored-By: Claude Opus 4.7 --- web-macro/Project.toml | 4 + web-macro/app/Project.toml | 2 + web-macro/src/BRMMacroWeb.jl | 1179 +++++++++++++++++++++++----------- 3 files changed, 798 insertions(+), 387 deletions(-) diff --git a/web-macro/Project.toml b/web-macro/Project.toml index 185615d..871fd7f 100644 --- a/web-macro/Project.toml +++ b/web-macro/Project.toml @@ -3,6 +3,8 @@ uuid = "cd55decf-6be8-4ea5-907a-b2dc35e4cc14" version = "0.1.0" [deps] +AlgebraOfGraphics = "cbdf2221-f076-402e-a563-3d30da359d67" +AlgebraOfVega = "a2420894-485a-4f42-a4fc-c4323ec6fb2b" BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a" CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" @@ -16,6 +18,7 @@ FlexiChains = "4a37a8b9-6e57-4b92-8664-298d46e639f7" HTMX = "27f3e1ef-6ef8-44dc-9e9a-2fb23ed44e83" HTMXObjects = "b12ef442-5798-4353-80f3-9562b03a0cb6" InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" @@ -25,6 +28,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" StanBlocks = "2e771a56-c23a-4e0b-9282-20c2e37157e9" StanLogDensityProblems = "a545de4d-8dba-46db-9d34-4e41d3f07807" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" TestModules = "63c02187-99fd-4e5c-aaf0-4d6bfebc181c" Treebars = "e1e568c4-3a56-40a4-95fa-9b9c6c16fccb" Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" diff --git a/web-macro/app/Project.toml b/web-macro/app/Project.toml index bced045..8a19296 100644 --- a/web-macro/app/Project.toml +++ b/web-macro/app/Project.toml @@ -1,4 +1,5 @@ [deps] +AlgebraOfVega = "a2420894-485a-4f42-a4fc-c4323ec6fb2b" BRMMacroWeb = "cd55decf-6be8-4ea5-907a-b2dc35e4cc14" DynamicObjects = "23d02862-63fe-4c6e-8fdb-1d52cbbd39d5" HTMX = "27f3e1ef-6ef8-44dc-9e9a-2fb23ed44e83" @@ -6,3 +7,4 @@ HTMXObjects = "b12ef442-5798-4353-80f3-9562b03a0cb6" Revise = "295af30f-e4ad-537b-8983-00126c2a3abe" TestModules = "63c02187-99fd-4e5c-aaf0-4d6bfebc181c" Treebars = "e1e568c4-3a56-40a4-95fa-9b9c6c16fccb" +WarmupHMC = "60658175-6863-4866-a322-ab51a11c0cfe" diff --git a/web-macro/src/BRMMacroWeb.jl b/web-macro/src/BRMMacroWeb.jl index 93c2a68..99ac67d 100644 --- a/web-macro/src/BRMMacroWeb.jl +++ b/web-macro/src/BRMMacroWeb.jl @@ -1,12 +1,21 @@ module BRMMacroWeb using HTMXObjects -using Treebars: polling_fetchindex, initialize_progress! +using DynamicObjects: fetchindex! +using Treebars: polling_fetchindex, initialize_progress!, + prepare_progress!, with_prepared_progress, + htmx_treebar_styles using Random using Chairmarks using DataFrames +using Statistics: quantile, median using FiniteDifferences: FiniteDifferences, central_fdm using BridgeStan: BridgeStan +using StanLogDensityProblems: StanProblem +using WarmupHMC: initialize_mcmc, adaptive_warmup_mcmc +using JSON +using AlgebraOfVega: vega_head, auto_remap_node, with_plot_caption, config, pointinterval, lineribbon, to_node, ECDFPlot, VLines +import AlgebraOfGraphics as AoG # The @brm macro and the VBRMI / SBBRMI implementations live alongside this # module so Revise tracks them. The scripts/ entry points (parsing.jl, @@ -27,12 +36,15 @@ include("html_expr.jl") # dispatch to `::Val{:bruno}`. Default is no extras. dataset_extras(::Val, df) = (;) -# DOs in dependency order. Every feature is a focused @dynamicstruct: the -# pipeline stages live as derived properties on `BRMRun`, example file I/O -# on `ExampleStore`, synthetic + namespace-merged data on `Dataset`, and -# AST parse/safety/transform on `Formula`. `AppData` is just a thin holder -# of sub-DOs plus the `pipeline_run` polling entry point; `AppContext` is -# the HTMX routes and rendering layer. +# DOs in dependency order. Every feature is a focused @dynamicstruct: +# - Backend/data lives on `AppData`: `dataset`, `run(text, ns)` with its pipeline +# data, `step_chain` / `compute_steps` for polling_fetchindex, `context` for +# per-request namespace/run bundles. +# - UI/HTML lives on the routes structs: `PipelineRoutes` owns the formula +# editor page, per-step render dispatch, and `context!`; the `@include +# examples` sub-struct owns the examples list/detail/mark routes and the +# inline `example_store` that constructs ExampleEntry instances with the +# right `__parent__` for URL construction. struct FormulaSecurityError <: Exception msg::String end @@ -53,7 +65,7 @@ _ALLOWED_CALLS = Set{Symbol}([ :MvNormal, :MixtureModel, :Dirichlet, :InverseGamma, :InverseGaussian, :VonMises, :Pareto, :OrderedLogistic, :Categorical, - :scale, :center, :standardize, :factor, :offset, + :zscale, :center, :standardize, :factor, :offset, :protect, :s, :bs, :t2, :gp, :ar, :ar1, :mo, :mo1, :cbind, :mvbind, :mm, :gr, :dp, :me, :centered, :Horseshoe, :ZeroInflatedPoisson, :weighted, @@ -110,14 +122,14 @@ end is_safe = violation === nothing _t = begin - alllocals = OrderedDict{Symbol,Symbol}() + local alllocals = OrderedDict{Symbol,Symbol}() (; ex=parse!(deepcopy(raw); info=(;alllocals)), alllocals) end transformed = _t.ex alllocals = _t.alllocals end @dynamicstruct struct Dataset - n::Int = 64 + n::Int = 16 seed::Int = 1 df = begin @@ -233,8 +245,12 @@ end class="brm-tier-pill", style="background:$tier_color") + # `__parent__` is the `@include examples` sub-struct (owns /examples/* routes). + # `__parent__.__parent__.pipeline` is the sibling PipelineRoutes sub-struct + # (owns /pipeline/* routes). URLs built via `query_url` so request @params + # auto-propagate and values auto-encode. permalink = h.a("🔗"; - href="/examples/$(HTTP.URIs.escapeuri(slug))", + href=string(__parent__/slug), title="Standalone URL", onclick="event.stopPropagation()", class="brm-permalink") @@ -245,7 +261,7 @@ end h.button(is_active ? active_text : inactive_text; type="button", class="brm-state-pill", - hx_get="/mark?label=$(HTTP.URIs.escapeuri(label))&state=$target_state", + hx_get=string(query_url(__parent__/"mark"; label, state=target_state)), hx_target="#$card_id", hx_swap="outerHTML", onclick="event.stopPropagation()", @@ -260,7 +276,7 @@ end state_pill(:deprioritized, "✓ deprioritized", "deprioritize"), ) - formula_form(routes) = h.form(; class="brm-example-form")( + formula_form = h.form(; class="brm-example-form")( h.input(; type="hidden", name="label", value=label), h.textarea(formula; name="formula", @@ -269,26 +285,26 @@ end h.button("cimpl (bench) ▶"; type="button", class="brm-branch-btn", - hx_get=string(query_url(routes/"stage/bench"; force=true)), + hx_get=string(query_url(__parent__.__parent__.pipeline/"stage/bench"; force=true)), hx_include="closest form", hx_target="#$result_id", hx_swap="innerHTML"), h.button("sbimpl (compile) ▶"; type="button", class="brm-branch-btn", - hx_get=string(query_url(routes/"stage/stan_compile"; force=true)), + hx_get=string(query_url(__parent__.__parent__.pipeline/"stage/stan_compile"; force=true)), hx_include="closest form", hx_target="#$result_id", hx_swap="innerHTML"), ) - card(routes) = begin + card = begin children = Any[HTMXObjects.md_to_node(body)] if formula !== nothing - push!(children, formula_form(routes)) + push!(children, formula_form) # Inline pipeline-result target — the form's hx_get fills this div - # with `render_pipeline(out)` so the user sees the - # VBRMI/finite-difference output right inside the card. + # with the pipeline output so the user sees the VBRMI/FD/stan-compile + # output right inside the card. push!(children, h.div(; id=result_id, class="brm-example-result")) @@ -314,7 +330,7 @@ end ) end - write_with!(; new_status=status, new_formula=formula) = begin + save!(; new_status=status, new_formula=formula) = begin io = IOBuffer() println(io, "# label: ", label) println(io, "# tier: ", tier) @@ -330,205 +346,569 @@ end endswith(new_formula, "\n") || println(io) end write(path, take!(io)) - ExampleEntry(; path) - end -end -@dynamicstruct struct ExampleStore - dir::String - - # `entries` is a method, not a cached field, because `save!` writes to - # disk and we want subsequent reads to see the new file mtime/content. - entries() = begin - isdir(dir) || mkpath(dir) - files = sort(filter(endswith(".jl"), - readdir(dir; join=true)); - by=mtime, rev=true) - ExampleEntry[ExampleEntry(; path=f) for f in files] + # Preserve __parent__ so the reloaded entry can still build route URLs. + ExampleEntry(; __parent__, path) end - find(label) = begin - for e in entries() - e.label == label && return e - end - nothing - end - find_by_slug(slug) = begin - for e in entries() - e.slug == slug && return e - end - nothing - end + toggle_status!(target) = + save!(; new_status = status == target ? :open : target) +end +# Element-returning counterpart to `Base.findfirst(pred, coll)` (which returns an +# index or `nothing`). Does the index-then-lookup dance once so callers don't. +findfirstelement(pred, coll) = begin + i = findfirst(pred, coll) + isnothing(i) ? nothing : coll[i] +end - save!(label; new_status=nothing, new_formula=nothing) = begin - e = find(label) - e === nothing && return nothing - e.write_with!(; - new_status = new_status === nothing ? e.status : new_status, - new_formula = new_formula === nothing ? e.formula : new_formula, +# Stan draws → DataFrames plumbing, shared between prior-predictive generation +# and posterior fits (pathfinder / warmup). Keep these as plain module-level +# helpers so callsites inside `@struct stan = …` don't accidentally become IPs. + +# Param-constrain each column of `unc_draws` (dim × n) with `include_tp=true, +# include_gq=true`, returning an (m × n) matrix where `m = length(param_names(instance; include_tp=true, include_gq=true))`. +constrain_draws(unc_draws, instance; rng_seed) = begin + rng = BridgeStan.StanRNG(instance, rng_seed) + m = length(BridgeStan.param_names(instance; include_tp=true, include_gq=true)) + n = size(unc_draws, 2) + mat = Matrix{Float64}(undef, m, n) + for i in 1:n + mat[:, i] = BridgeStan.param_constrain( + instance, collect(view(unc_draws, :, i)); + include_tp=true, include_gq=true, rng=rng, ) end + mat end -# One run of the @brm pipeline for a given (text, namespace) pair. Every -# stage is a derived property -- accessing `run.brmi` triggers parse + eval, -# accessing `run.benches` triggers vbrmi + finite-difference check + the -# benchmark loop, etc. Unused branches don't compute. Safety is enforced -# at the one point where it matters: `wrapped` refuses to produce Julia -# code for an unsafe formula. - -@dynamicstruct struct BRMRun - text::String - namespace::Symbol = :default - (; dataset) = __parent__ - - formula = Formula(; text) - df = dataset.df - container = dataset.container(namespace) - wrapped = begin - formula.is_safe || throw(formula.violation) - _brm(text; df=container) - end - brmi = eval(wrapped) - - # ── cimpl branch ── - vbrmi = VBRMI(brmi) - dim = LogDensityProblems.dimension(vbrmi) - x0 = randn(Xoshiro(0), dim) - ldp = string(LogDensityProblems.logdensity(vbrmi, x0)) - grad = FiniteDifferences.grad( - central_fdm(5, 1), - Base.Fix1(LogDensityProblems.logdensity, vbrmi), - x0, - )[1] - - benches = begin - x_rand = randn(dim) - bs = Pair{String,Any}[] - push!(bs, "logdensity (total)" => - @be randn(dim) LogDensityProblems.logdensity($vbrmi, _)) - push!(bs, "lprior!" => - @be randn(dim) lprior!($vbrmi, _)) - # Per-Part lprior! split: the foldl in lprior!(blocks, x) hands - # each Part a view of exactly nparams(part) reals. Reconstruct - # those slices here so each Part's contribution can be benched - # in isolation. - let pos = 0 - for (group_key, parts) in pairs(vbrmi.meta.blocks) - for (i, part) in enumerate(parts) - n = nparams(part) - xi = view(x_rand, pos+1:pos+n) - push!(bs, " lprior!($group_key[$i] $(part))" => - @be lprior!($part, $xi)) - pos += n - end - end - end - # llikelihood! splits: each materialized column (either a - # linear-predictor MaterializedColumn or a LikelihoodColumn). - _ = lprior!(vbrmi, x_rand) - for (key, m) in pairs(vbrmi.meta.materialized) - push!(bs, "llikelihood!($key)" => - @be llikelihood!($m)) - end - bs - end - - # ── stan branch ── - sbbrmi = SBBRMI(brmi) - stan_src = stan_code(sbbrmi) - stan_file = begin - f = tempname() * ".stan" - write(f, stan_src) - f - end - stan_lib = BridgeStan.compile_model(stan_file) +# Build the (long, wide, summary) DataFrame triple from a constrained draws +# matrix `constrained` (m × n) and its matching parameter `names` (length m). +# Splits indexed names on the first `.` into (:param, :index) with :index as Int +# (0 for scalars); summary groups by (:param, :index) with the bands columns +# expected by `pointinterval(bands=…)` / `lineribbon(bands=…)`. +dfs_from_constrained(constrained, names) = begin + n = size(constrained, 2) + splits = [split(nm, '.', limit=2) for nm in names] + base_names = [String(first(s)) for s in splits] + parse_idx(s) = (v = tryparse(Int, s); isnothing(v) ? 0 : v) + indices = [length(s) > 1 ? parse_idx(String(s[2])) : 0 for s in splits] + long = DataFrame( + param = repeat(base_names, inner=n), + index = repeat(indices, inner=n), + draw = repeat(1:n, outer=length(base_names)), + value = vec(constrained'), + ) + wide = DataFrame( + [Symbol(names[i]) => constrained[i, :] for i in eachindex(names)] + ) + summary = combine( + groupby(long, [:param, :index]), + :value => (v -> quantile(v, 0.025)) => :q025, + :value => (v -> quantile(v, 0.10)) => :q10, + :value => (v -> quantile(v, 0.25)) => :q25, + :value => median => :median, + :value => (v -> quantile(v, 0.75)) => :q75, + :value => (v -> quantile(v, 0.90)) => :q90, + :value => (v -> quantile(v, 0.975)) => :q975, + ) + (; long, wide, summary) end + @dynamicstruct struct AppData __status__ = initialize_progress!(:state; description="BRM pipeline") examples_dir = joinpath(dirname(@__DIR__), "examples") - default_formula = """loc1 ~ 1 + a + c1 + (1 + b + c1 | g1) + (1 | g2) -log(err1) ~ 1 + d -y1 ~ Normal(loc1, err1) + dataset = Dataset() -log_rate ~ 1 + a + (1 | g3) -k1 ~ Poisson(exp(log_rate)) + namespace_from(label) = isempty(strip(label)) ? :default : + Symbol(lowercase(first(split(strip(label), r"[\s:\-]+")))) -log_odds_bin ~ 1 + c2 + (1 | g2) -bin_succ ~ Binomial(bin_n, logistic(log_odds_bin)) + # One run of the @brm pipeline for a given (text, namespace) pair. Every + # stage is a derived property -- accessing `.brmi` triggers parse + eval, + # `.benches` triggers the benchmark loop, etc. Unused branches don't + # compute. Safety is enforced at `wrapped`, which refuses to produce Julia + # code for an unsafe formula. + # TODO(DO): investigate whether indexed inline structs (@struct) should + # accept default values for index params. Today `@struct run(text, namespace=:default)` + # errors with "index param must be a Symbol" because the default becomes + # a `:kw` Expr. Unclear if there's a sensible meaning for such defaults at + # all (cache keying, forwarding, etc.) or if we should just continue + # requiring bare Symbol params. + @struct run(text, namespace) = begin + formula = Formula(text) + df = dataset.df + container = dataset.container(namespace) + + wrapped = begin + formula.is_safe || throw(formula.violation) + _brm(text; df=container) + end + brmi = eval(wrapped) + + # ── cimpl branch ── + vbrmi = VBRMI(brmi) + dim = LogDensityProblems.dimension(vbrmi) + x0 = randn(Xoshiro(0), dim) + ldp = string(LogDensityProblems.logdensity(vbrmi, x0)) + grad = FiniteDifferences.grad( + central_fdm(5, 1), + Base.Fix1(LogDensityProblems.logdensity, vbrmi), + x0, + )[1] + + tol = 1e-8 + n_dead = count(<=(tol) ∘ abs, grad) + dead = findall(<=(tol) ∘ abs, grad) + + benches = vcat( + [ + "logdensity (total)" => @be(randn(dim), LogDensityProblems.logdensity($vbrmi, _)), + "lprior!" => @be(randn(dim), lprior!($vbrmi, _)), + ], + [ + " lprior!($group_key[$i] $(part))" => @be(randn(nparams(part)), lprior!($part, _)) + for (group_key, parts) in pairs(vbrmi.meta.blocks) + for (i, part) in enumerate(parts) + ], + [ + "llikelihood!($key)" => @be(llikelihood!($m)) + for (key, m) in pairs(vbrmi.meta.materialized) + ], + ) -log_odds_b ~ 1 + b -bin_y ~ Bernoulli(logistic(log_odds_b)) -""" + # ── stan branch ── + sbbrmi = SBBRMI(brmi) + + # Everything Stan-related bundled under `r.stan.*`. Nested access via + # step_chain's tuple-path specs (e.g. `(:stan, :src)`). + @struct stan = begin + src = stan_code(sbbrmi) + # Hash-keyed cache path so identical Stan source reuses the same + # .stan (and co-located .so) across requests. `BridgeStan.compile_model` + # invokes make, which skips when the .so is newer than the .stan — + # so a cache hit resolves in milliseconds instead of re-running + # the C++ build. + file = begin + p = joinpath(tempdir(), "brm_stan", string(hash(src)) * ".stan") + mkpath(dirname(p)) + isfile(p) || write(p, src) + p + end + lib = BridgeStan.compile_model(file) + # SB's `stan_data` walks the SlicModel → StanModel tracing which + # auto-declares `_n` / `_m` sizes for every vector / matrix, then + # `bridgestan_data` JSON-serializes with Stan's column-major + # matrix convention. + data = StanBlocks.stan.bridgestan_data(StanBlocks.stan_data(sbbrmi.model)) + instance = BridgeStan.StanModel(lib, data) + dim = BridgeStan.param_unc_num(instance) + # Fixed-rng narrow-normal init — deterministic, cache-friendly. + init = 0.1 .* randn(Xoshiro(42), dim) + + # Smoke-test evaluation: log density at the init params. Forces the + # model-loaded + data-bound path without running Pathfinder. + log_density = BridgeStan.log_density(instance, init) + + # Synthetic-data generation: sample N unconstrained parameter draws + # from a narrow zero-mean normal, then `param_constrain` each with + # `include_tp` + `include_gq` so the output matrix also carries + # transformed parameters and generated-quantities (the synthetic + # outcomes `y_sim` live in the GQ block when the model defines + # one). + generated_n = 50 + generated_unc = 0.1 .* randn(Xoshiro(44), dim, generated_n) + generated_names = BridgeStan.param_names(instance; include_tp=true, include_gq=true) + generated = constrain_draws(generated_unc, instance; rng_seed=45) + generated_dfs = dfs_from_constrained(generated, generated_names) + generated_df = generated_dfs.long + generated_wide_df = generated_dfs.wide + generated_summary_df = generated_dfs.summary + # Ground-truth overlay: the `fit_draw_idx`-th column of the + # constrained generated matrix — one value per indexed parameter. + # Same (param, index) layout as `generated_df` so plots can join on it. + truth_df = begin + truth_col = view(generated, :, fit_draw_idx) + splits = [split(nm, '.', limit=2) for nm in generated_names] + base_names = [String(first(s)) for s in splits] + parse_idx(s) = (v = tryparse(Int, s); isnothing(v) ? 0 : v) + indices = [length(s) > 1 ? parse_idx(String(s[2])) : 0 for s in splits] + DataFrame( + param = base_names, + index = indices, + truth = collect(truth_col), + ) + end + # Simulation-based calibration setup: pick one draw from the prior + # predictive `generated` matrix, extract every `*_gen[.i.j]` entry, + # and fold them back into the Stan data dict as their `*` (observed) + # counterparts. The resulting `fit_instance` shares the compiled + # .so with `instance` but is bound to this synthetic observed data, + # so Pathfinder / full warmup samples `p(theta | y_sim)` and should + # recover the ground-truth `generated_unc[:, fit_draw_idx]`. + fit_draw_idx = 1 + fit_truth_unc = collect(view(generated_unc, :, fit_draw_idx)) + fit_data_dict = begin + base = StanBlocks.stan_data(sbbrmi.model) + col = view(generated, :, fit_draw_idx) + groups = Dict{Symbol, Vector{Tuple{Vector{Int}, Float64}}}() + for (i, name) in enumerate(generated_names) + m = match(r"^(.+)_gen(?:\.(.+))?$", name) + isnothing(m) && continue + base_name = Symbol(m.captures[1]) + idx_str = m.captures[2] + idxs = isnothing(idx_str) ? Int[] : + [Base.parse(Int, s) for s in split(idx_str, ".")] + push!(get!(Vector{Tuple{Vector{Int}, Float64}}, groups, base_name), + (idxs, col[i])) + end + overrides = Dict{Symbol, Any}() + for (name, entries) in groups + if length(entries) == 1 && isempty(entries[1][1]) + overrides[name] = entries[1][2] + else + orig = base[name] + out = similar(orig, Float64) + for (idxs, v) in entries + out[idxs...] = v + end + overrides[name] = out + end + end + merge(base, overrides) + end + fit_instance = BridgeStan.StanModel(lib, + StanBlocks.stan.bridgestan_data(fit_data_dict)) + + # IP: Pathfinder init (fast, no MCMC). The `progress=__status__` + # hook lets Treebars nest the 100 maxiters subtree under whatever + # node called `fetchindex!(status, …, pathfinder, instance, init)`. + pathfinder(instance, init; rng=Xoshiro(42), maxiters=100) = + initialize_mcmc(StanProblem(instance), init; rng, progress=__status__, maxiters) + # IP: full Stan + WarmupHMC fit. Same progress-hooking pattern. + # Returns a rich NamedTuple with `.posterior_position`, `.ess`, + # `.n_divergent_samples`, etc. + posterior_warmup(instance, init; rng=Xoshiro(42), n_draws=200) = + adaptive_warmup_mcmc(rng, StanProblem(instance); init, n_draws, progress=__status__) + # Gaussian-approximation draws from Pathfinder. Reads the IP via + # `@memo` so if `compute_steps` already computed it with progress + # nesting, we get the cached value for free. + posterior_pathfinder = begin + pf = @memo pathfinder(fit_instance, init) + pf.position .+ pf.scale * randn(Xoshiro(43), dim, 200) + end + # Default. Switch this alias to the warmup draws (`@memo + # posterior_warmup(instance, init).posterior_position`) to promote + # the full fit, or expose a toggle via a param later. + posterior = posterior_pathfinder + # Constrained posterior draws (+TP+GQ), plus (long, wide, summary). + posterior_constrained = constrain_draws(posterior, fit_instance; rng_seed=46) + posterior_dfs = dfs_from_constrained(posterior_constrained, generated_names) + posterior_long_df = posterior_dfs.long + posterior_wide_df = posterior_dfs.wide + posterior_summary_df = posterior_dfs.summary + + # Parallel set for the warmup+MCMC path. The warmup IP cache is + # warmed by `fetchindex!` in `compute_steps`; `@memo` hits it here. + posterior_warmup_draws = (@memo posterior_warmup(fit_instance, init)).posterior_position + posterior_warmup_constrained = constrain_draws(posterior_warmup_draws, fit_instance; rng_seed=47) + posterior_warmup_dfs = dfs_from_constrained(posterior_warmup_constrained, generated_names) + posterior_warmup_long_df = posterior_warmup_dfs.long + posterior_warmup_wide_df = posterior_warmup_dfs.wide + posterior_warmup_summary_df = posterior_warmup_dfs.summary + posterior_warmup_diagnostics = begin + w = @memo posterior_warmup(fit_instance, init) + (; w.n_divergent_samples, ess=w.ess) + end + end - dataset = Dataset() - example_store = ExampleStore(; dir=examples_dir) + # Stage-named aliases. `step_chain` / `compute_steps` extract step + # outputs by looking up these names via `getproperty`. + parse = formula.raw + transform = (; formula.transformed, formula.alllocals) + wrap = wrapped + end - namespace_from(label) = isempty(strip(label)) ? :default : - Symbol(lowercase(first(split(strip(label), r"[\s:\-]+")))) + # Per-request bundle for a (label, formula) pair. Pure construction; the + # routes side handles persistence before invoking this. + @struct context(label, formula) = begin + namespace = namespace_from(label) + run = __parent__.run(formula, namespace) + end - # Ordered pipeline stages. Index gates which BRMRun properties - # `pipeline_run` touches (and which sections `render_pipeline` shows). - stages = (:parse, :transform, :wrap, :brmi, - :vbrmi, :bench, - :slic_model, :stan_code, :stan_compile) - stage_index(s) = something(findfirst(==(s), stages), length(stages)) - - # Indexable property: `appdata.run[text, ns]` is cached in-memory per key, - # `appdata.run(text, ns)` is fresh each call. - run(text, namespace=:default) = - BRMRun(; __parent__=__self__, text, namespace) - - # Indexable fetch for `polling_fetchindex` (accessed via brackets by the - # caller). Touches BRMRun properties up through `stage` so the heavy work - # lands inside the polled task rather than the HTTP response callback. - pipeline_run(text, stage::Symbol, namespace=:default) = begin - r = run[text, namespace] - s = stage_index(stage) - s >= 1 && r.formula.raw - s >= 2 && r.formula.transformed - s >= 3 && r.wrapped - s >= 4 && r.brmi - stage === :vbrmi && r.grad - stage === :bench && r.benches - stage === :slic_model && r.sbbrmi - stage === :stan_code && r.stan_src - stage === :stan_compile && r.stan_lib - r + # DAG of step chains — one NamedTuple per stage target, keyed by step names + # in dependency order. Built incrementally via `merge`: each stage's chain + # is its parent's chain plus the one step it adds. The DAG branches are + # visible in the `merge` calls (e.g. `slic_model = merge(brmi, …)` forks + # off of `brmi`, parallel to `vbrmi`). Inner values are one of: + # - Symbol → `r.` (top-level property) + # - Tuple{Vararg{Symbol}} → nested access, e.g. `(:stan, :src)` → `r.stan.src` + # - NamedTuple of (Symbol|Tuple) → bundle, keys preserved in the result. + step_chain(name::Symbol) = begin + parse = (; parse=:parse) + transform = merge(parse, (; transform=:transform)) + wrap = merge(transform, (; wrap=:wrap)) + brmi = merge(wrap, (; brmi=:brmi)) + vbrmi = merge(brmi, (; vbrmi=(; vbrmi=:vbrmi, dim=:dim, ldp=:ldp, grad=:grad, n_dead=:n_dead, dead=:dead))) + bench = merge(vbrmi, (; bench=:benches)) + slic_model = merge(brmi, (; slic_model=:sbbrmi)) + stan_code = merge(slic_model, (; stan_code=(:stan, :src))) + stan_compile = merge(stan_code, (; stan_compile=(; file=(:stan, :file), lib=(:stan, :lib)))) + stan_instantiate = merge(stan_compile, (; stan_instantiate=(; instance=(:stan, :instance), dim=(:stan, :dim), init=(:stan, :init)))) + stan_eval = merge(stan_instantiate, (; stan_eval=(:stan, :log_density))) + stan_generate = merge(stan_eval, (; stan_generate=(; long=(:stan, :generated_df), wide=(:stan, :generated_wide_df), summary=(:stan, :generated_summary_df), truth=(:stan, :truth_df)))) + # Pathfinder / full warmup are computed via `fetchindex!` in + # `compute_steps` (special-cased below by step name) so the IP's + # progress subtree attaches to the step's phase. Chain-level specs + # read the resulting cached values back out as plain properties. + stan_fit_pathfinder = merge(stan_generate, (; stan_fit_pathfinder=(; long=(:stan, :posterior_long_df), wide=(:stan, :posterior_wide_df), summary=(:stan, :posterior_summary_df), truth=(:stan, :truth_df)))) + stan_fit_warmup = merge(stan_fit_pathfinder, (; stan_fit_warmup=(; long=(:stan, :posterior_warmup_long_df), wide=(:stan, :posterior_warmup_wide_df), summary=(:stan, :posterior_warmup_summary_df), diagnostics=(:stan, :posterior_warmup_diagnostics), truth=(:stan, :truth_df)))) + (; parse, transform, wrap, brmi, vbrmi, bench, slic_model, stan_code, stan_compile, + stan_instantiate, stan_eval, stan_generate, stan_fit_pathfinder, stan_fit_warmup)[name] + end + + # Fetch target for `polling_fetchindex`. Pre-enumerates each step in the + # requested chain as a pending progress child (so the whole pipeline is + # visible up front as dim "pending" nodes), then runs each step under its + # phase — `with_prepared_progress` handles start/finalize/fail around the + # property access. Heavy work runs in polling_fetchindex's background task + # via DO's lazy property cascading. Returns a NamedTuple keyed by step + # names plus `:data` (the synthetic-data frame for the pipeline top pin). + "Pipeline($name)" + compute_steps(text, namespace, name::Symbol) = begin + r = run(text, namespace) + chain = step_chain(name) + phases = [prepare_progress!(__status__; description=string(k)) for k in keys(chain)] + # Resolve each spec against `r`: + # Symbol → `r.` + # Tuple of Symbol → nested path `r.....` + # NamedTuple → bundle, recurse per value. + resolve(s::Symbol) = getproperty(r, s) + resolve(p::Tuple{Vararg{Symbol}}) = foldl(getproperty, p; init=r) + resolve(b::NamedTuple) = map(resolve, b) + vals = map(pairs(chain), phases) do (step_name, spec), phase + with_prepared_progress(phase) do progress + if step_name === :stan_fit_pathfinder + # Warm the pathfinder IP cache under this phase's progress, + # then resolve the (long, wide, summary) bundle (which + # reads back the cached value via `@memo`). + fetchindex!(progress, r.stan.pathfinder, r.stan.fit_instance, r.stan.init) + resolve(spec) + elseif step_name === :stan_fit_warmup + # Warm the warmup IP cache under this phase's progress, + # then resolve the (long, wide, summary, diagnostics) + # bundle (which reads back the cached value via `@memo`). + fetchindex!(progress, r.stan.posterior_warmup, r.stan.fit_instance, r.stan.init) + resolve(spec) + else + resolve(spec) + end + end + end + # Stages render most-recent-first; synthetic data pinned at the top. + merge((; data=r.df), + NamedTuple{reverse(keys(chain))}(Tuple(reverse(vals)))) end end APPDATA = AppData(; cache_type=:parallel) -@htmx struct AppContext - __appdata__ = APPDATA - (; default_formula, dataset, example_store, namespace_from, - stages, stage_index, run, pipeline_run) = __appdata__ - # Page-level stylesheet read once at construction. Classes are consumed by - # ExampleEntry.card / html_expr.jl; per-symbol / per-status colors that - # are data-derived stay inline on the element. - css = read(joinpath(@__DIR__, "brm-macro.css"), String) +# Pipeline-page routes mounted at /pipeline. The formula editor, stage polling, +# and sbimpl source views all live here. The top-level AppContext just includes +# this struct plus the Examples section and the page chrome. +@htmx struct PipelineRoutes + (; context, compute_steps) = __appdata__ + (; default_formula) = __parent__ + @param (; formula, label) = __parent__ + + # Persist + context. Reaches into the sibling Examples include for the + # examples store (UI concern: writing the edited formula back to the .jl + # file corresponding to `label`), then returns the pure run context. + context!() = begin + isempty(label) || __parent__.examples.example_store.persist!(label, formula) + context(label, formula) + end - # HTMXObjects auto-uses `__page__` to wrap any route's return value into a - # full page on direct browser navigation, while returning just the fragment - # for HTMX requests (see `_resolve_response` in HTMXObjects.jl). The - # sidebar's `hx-get` swaps target `#content` directly. - __page__(content) = htmx( - h.div(; class="brm-layout")( - nav_sidebar([ - "Pipeline" => "/", - "Examples" => "/examples", - ]), - h.main(; class="container brm-main")( - h.div(; id="content")(content), - ), - ); - pico_version="2", - extra_head=( - h.title("BRM macro action"), - h.style(__self__.css), - ), - ) + # Per-step HTML rendering. `getproperty(render, step_key)(value)` emits the + # section for that step; `compute_steps` produces the NamedTuple whose keys + # drive dispatch here. + @struct render = begin + data(df) = h.details( + h.summary("Synthetic data ($(nrow(df)) rows × $(ncol(df)) cols: " * + join(names(df), ", ") * ") — click to expand"), + render_table(df; sortable=false), + ) + parse(x) = h.section( + h.h3("1. Meta.parse — raw Julia AST"), + h.pre(x), + ) + transform(x) = h.section( + h.h3("2. parse! — rewritten AST (= → @n/@x assign, ~ → @n/@x ~)"), + h.pre(x.transformed), + h.h3(" locals classified by parse!"), + h.pre(x.alllocals), + ) + wrap(x) = h.section( + h.h3("3. _brm — full let-block (df spliced as a literal)"), + h.pre(x), + ) + brmi(x) = h.section( + h.h3("4. eval — BRMI value (parsed model)"), + brmi_card(x), + ) + vbrmi(x) = begin + fd_summary = x.n_dead == 0 ? + h.span("logdensity + FD check: $(x.dim)/$(x.dim) active ✓"; class="brm-status-ok") : + h.span("logdensity + FD check: $(x.n_dead) dead param(s)"; class="brm-status-err") + fd_body = h.div( + h.p("dim = ", x.dim, ", logdensity = ", x.ldp), + isempty(x.dead) ? "" : + h.p(; class="brm-status-err")("dead param indices: ", x.dead), + h.pre(x.grad), + ) + h.section( + h.h3("5. VBRMI — materialized action (blocks, dim, columns)"), + vbrmi_card(x.vbrmi), + h.details(h.summary(fd_summary), fd_body), + ) + end + bench(x) = h.section( + h.h3("6. Chairmarks @be — per-step"), + [h.article(h.header(lbl), h.pre(b)) for (lbl, b) in x]..., + ) + slic_model(x) = h.section( + h.h3("5a. SlicModel — SBBRMI @slic body"), + h.pre(x.model.model), + h.p("data keys: ", h.code(sort(collect(keys(x.data))))), + ) + stan_code(x) = h.section( + h.h3("5b. StanCode — transpiled Stan source"), + h.pre(x), + ) + stan_compile(x) = h.section( + h.h3("5c. StanCompile — BridgeStan shared library"), + h.p("stan file: ", h.code(x.file)), + h.p("compiled .so: ", h.code(x.lib)), + ) + stan_instantiate(x) = h.section( + h.h3("6a. StanInstantiate — model bound to data"), + h.p("param_unc_num = ", x.dim), + h.p("init (narrow normal, rng=Xoshiro(42)):"), + h.pre(x.init), + ) + stan_eval(x) = h.section( + h.h3("6b. StanEval — log density at init"), + h.p("log_density = ", x), + ) + # Shared plot-tabset builder used by stan_generate and the fit stages. + # `kind` goes into tab titles ("prior predictive" / "posterior") and + # plot ids. Returns the tabset + wide-table details block. + posterior_plots(long, wide, summary; id_prefix, kind, truth=nothing) = begin + bands = [:q025 => :q975, :q10 => :q90, :q25 => :q75] + pi_title = "$kind (N=$(nrow(wide)) draws)" + ecdf_title = "$kind — ECDF" + lr_title = "$kind — line + ribbon" + den_title = "$kind — histogram" + indep_x = config(facet=(; linkxaxes=:none)) + indep_y = config(facet=(; linkyaxes=:none)) + # Overlay layers: for (x=:index, y=:value) plots, plot truth as + # filled black dots at (:index, :truth); for (x=:value) plots, + # overlay vertical rules at truth values, colored by :index to + # match the base layer's coloring. + overlay_xy = isnothing(truth) ? nothing : + AoG.data(truth) * AoG.mapping(:index, :truth, row=:param) * + AoG.visual(AoG.Scatter; color=:black) + overlay_vrule = isnothing(truth) ? nothing : + AoG.data(truth) * AoG.mapping(:truth; row=:param, color=:index) * + AoG.visual(VLines) + add(spec, overlay) = isnothing(overlay) ? spec : spec + overlay + spec_pi = add(AoG.data(summary) * + AoG.mapping(:index, :median, row=:param) * + pointinterval(; bands, orientation=:vertical), + overlay_xy) * + config(title=pi_title) * indep_y + spec_lr = add(AoG.data(summary) * + AoG.mapping(:index, :median, row=:param) * + lineribbon(; bands), + overlay_xy) * + config(title=lr_title) * indep_y + spec_hist = add(AoG.data(long) * + AoG.mapping(:value; row=:param, color=:index) * + AoG.visual(ECDFPlot), + overlay_vrule) * + config(title=ecdf_title) * indep_x + spec_den = add(AoG.data(long) * + AoG.mapping(:value; row=:param, color=:index) * + AoG.histogram(), + overlay_vrule) * + config(title=den_title) * indep_x + tabs = tabset( + "Point + Interval" => to_node(spec_pi; id="$id_prefix-pi"), + "Line + Ribbon" => to_node(spec_lr; id="$id_prefix-lr"), + "ECDF" => to_node(spec_hist; id="$id_prefix-ecdf"), + "Histogram" => to_node(spec_den; id="$id_prefix-hist"), + "Point + Interval (picker)" => with_plot_caption(spec_pi; + auto_remap=(; dims=["param" => "Parameter / TP / GQ"]), + title=pi_title, plot_id="$id_prefix-pi-pick"), + "Line + Ribbon (picker)" => with_plot_caption(spec_lr; + auto_remap=(; dims=["param" => "Parameter / TP / GQ"]), + title=lr_title, plot_id="$id_prefix-lr-pick"), + "ECDF (picker)" => with_plot_caption(spec_hist; + auto_remap=(; dims=["param" => "Parameter / TP / GQ", + "index" => "Index (vector/matrix position)"]), + title=ecdf_title, plot_id="$id_prefix-ecdf-pick"), + "Histogram (picker)" => with_plot_caption(spec_den; + auto_remap=(; dims=["param" => "Parameter / TP / GQ", + "index" => "Index (vector/matrix position)"]), + title=den_title, plot_id="$id_prefix-hist-pick"); + id="$id_prefix-tabs", + ) + wide_details = h.details( + h.summary("Wide-format table (one row per draw, one column per indexed parameter)"), + render_table(wide; sortable=true), + ) + (; tabs, wide_details) + end + stan_generate(x) = begin + (; long, wide, summary, truth) = x + p = posterior_plots(long, wide, summary; + id_prefix="brm-plot-generated", + kind="Generated data (prior predictive)", + truth) + h.section( + h.h3("6c. StanGenerate — synthetic data from narrow-normal prior + param_constrain"), + h.p("long format: ", nrow(long), " rows · ", ncol(long), " cols · ", + "wide format: ", nrow(wide), " rows · ", ncol(wide), " cols"), + p.tabs, p.wide_details, + ) + end + stan_fit_pathfinder(x) = begin + (; long, wide, summary, truth) = x + p = posterior_plots(long, wide, summary; + id_prefix="brm-plot-pf", + kind="Pathfinder posterior", + truth) + h.section( + h.h3("6d. StanFit (Pathfinder) — variational approximation draws"), + h.p("long format: ", nrow(long), " rows · ", ncol(long), " cols · ", + "wide format: ", nrow(wide), " rows · ", ncol(wide), " cols"), + p.tabs, p.wide_details, + ) + end + stan_fit_warmup(x) = begin + (; long, wide, summary, diagnostics, truth) = x + p = posterior_plots(long, wide, summary; + id_prefix="brm-plot-warmup", + kind="Warmup+MCMC posterior", + truth) + h.section( + h.h3("6d'. StanFit (Warmup+MCMC) — full Stan fit"), + h.p("n_divergent_samples: ", diagnostics.n_divergent_samples, + " · min ESS: ", minimum(diagnostics.ess)), + h.p("long format: ", nrow(long), " rows · ", ncol(long), " cols · ", + "wide format: ", nrow(wide), " rows · ", ncol(wide), " cols"), + p.tabs, p.wide_details, + ) + end + end # Pre-canned formulas. The ones above the divider exercise individual # features in isolation; the last one stacks everything into a single @@ -555,221 +935,246 @@ bin_y ~ Bernoulli(logistic(log_odds_b)) "everything" => default_formula, ] - preset_button(label, formula) = h.button(label; - type="button", - class="brm-preset-btn", - data_formula=formula, - onclick="document.querySelector('textarea[name=formula]').value = this.dataset.formula; document.getElementById('stage-vbrmi').click()") - - stage_button(label, stage) = h.button(label; - type="button", - id="stage-$stage", - hx_get=string(query_url(__self__/"stage/$stage"; force=true)), - hx_include="#brm-macro-form", - hx_target="#brm-macro-output", - hx_swap="outerHTML") - - render_pipeline(r, stage) = begin - s = stage_index(stage) - sections = Vector{Any}[] - - # Synthetic data pinned at top, collapsed by default so the macro - # pipeline output stays the focus. - data_section = Any[ - h.details( - h.summary("Synthetic data ($(nrow(r.df)) rows × $(ncol(r.df)) cols: " * - join(string.(names(r.df)), ", ") * ") — click to expand"), - render_table(r.df; sortable=false), - ), - ] - - s >= 1 && push!(sections, Any[ - h.h3("1. Meta.parse — raw Julia AST"), - h.pre(sprint(show, r.formula.raw)), - ]) - s >= 2 && push!(sections, Any[ - h.h3("2. parse! — rewritten AST (= → @n/@x assign, ~ → @n/@x ~)"), - h.pre(sprint(show, r.formula.transformed)), - h.h3(" locals classified by parse!"), - h.pre(sprint(show, r.formula.alllocals)), - ]) - s >= 3 && push!(sections, Any[ - h.h3("3. _brm — full let-block (df spliced as a literal)"), - h.pre(sprint(show, r.wrapped)), - ]) - s >= 4 && push!(sections, Any[ - h.h3("4. eval — BRMI value (parsed model)"), - brmi_card(r.brmi), - ]) - - if stage in (:vbrmi, :bench) - tol = 1e-8 - n_dead = count(<=(tol) ∘ abs, r.grad) - fd_summary = n_dead == 0 ? - h.span("logdensity + FD check: $(r.dim)/$(r.dim) active ✓"; - class="brm-status-ok") : - h.span("logdensity + FD check: $(n_dead) dead param(s)"; - class="brm-status-err") - dead = findall(<=(tol) ∘ abs, r.grad) - fd_body = h.div( - h.p("dim = ", string(r.dim), ", logdensity = ", r.ldp), - isempty(dead) ? "" : - h.p(; class="brm-status-err")( - "dead param indices: ", string(dead)), - h.pre(sprint(show, MIME"text/plain"(), r.grad)), - ) - push!(sections, Any[ - h.h3("5. VBRMI — materialized action (blocks, dim, columns)"), - vbrmi_card(r.vbrmi), - h.details(h.summary(fd_summary), fd_body), - ]) - end - if stage === :bench - bench_rows = [h.div( - h.strong(lbl), h.br(), - h.pre(sprint(show, MIME"text/plain"(), b)) - ) for (lbl, b) in r.benches] - push!(sections, Any[h.h3("6. Chairmarks @be — per-step"), bench_rows...]) - end - if stage in (:slic_model, :stan_code, :stan_compile) - push!(sections, Any[ - h.h3("5a. SlicModel — SBBRMI @slic body"), - h.pre(sprint(show, r.sbbrmi.model.model)), - h.p("data keys: ", - h.code(string(sort(collect(keys(r.sbbrmi.data)))))), - ]) - end - if stage in (:stan_code, :stan_compile) - push!(sections, Any[ - h.h3("5b. StanCode — transpiled Stan source"), - h.pre(r.stan_src), - ]) - end - if stage === :stan_compile - push!(sections, Any[ - h.h3("5c. StanCompile — BridgeStan shared library"), - h.p("stan file: ", h.code(r.stan_file)), - h.p("compiled .so: ", h.code(r.stan_lib)), - ]) - end + @struct preset(label, formula) = begin + button = h.button(label; + type="button", + class="brm-preset-btn", + data_formula=formula, + onclick="document.querySelector('textarea[name=formula]').value = this.dataset.formula; document.getElementById('stage-vbrmi').click()" + ) + end - # Stages render most-recent-first; synthetic data sits at the very top. - children = reduce(vcat, reverse(sections); init=Any[]) - prepend!(children, data_section) - h.div(; id="brm-macro-output")(children...) + @struct stage(label, id) = begin + button = h.button(label; + type="button", + id="stage-$id", + hx_get=string(query_url(__self__/"stage/$id"; force=true)), + hx_include="#brm-macro-form", + hx_target="#brm-macro-output", + # `innerHTML` keeps the `#brm-macro-output` wrapper in the DOM + # across swaps — including when polling_fetchindex throws and the + # response is a bare error article with no matching id. Without + # this, buttons target a gone id after the first failure. + hx_swap="innerHTML") end - index_body(formula) = h.div( - h.h1("BRM macro pipeline"), - h.p( - "Enter a ", h.code("@brm"), " formula and step through the macro pipeline: ", - h.code("Meta.parse"), " -> ", h.code("parse!"), " -> ", h.code("_brm"), - " let-block -> ", h.code("eval"), " -> ", h.code("VBRMI"), " action -> ", - h.code("Chairmarks"), " benchmark.", - ), - h.details( - h.summary(h.small("Allowed functions in formulas")), - h.p(h.small( - join(sort(collect(string.(s) for s in _ALLOWED_CALLS)), ", "), - )), - ), - h.form(; id="brm-macro-form")( - h.label("Load preset"), - h.div(; class="brm-preset-row")( - [preset_button(lbl, body) for (lbl, body) in presets]..., - ), - h.label("Formula")( - h.textarea(formula; - name="formula", rows=8, - class="brm-formula-textarea"), + @get index = begin + # If an example form posted us a (label, formula) pair, persist the + # edited formula to that example's .jl file so the next visit to the + # Examples page shows the user's edits instead of the seed default. + context!() + h.div( + h.h1("BRM macro pipeline"), + h.p( + "Enter a ", h.code("@brm"), " formula and step through the macro pipeline: ", + h.code("Meta.parse"), " -> ", h.code("parse!"), " -> ", h.code("_brm"), + " let-block -> ", h.code("eval"), " -> ", h.code("VBRMI"), " action -> ", + h.code("Chairmarks"), " benchmark.", ), - h.fieldset(; class="grid")( - stage_button("1. Parse", :parse), - stage_button("2. Transform", :transform), - stage_button("3. Wrap", :wrap), - stage_button("4. BRMI", :brmi), + h.details( + h.summary(h.small("Allowed functions in formulas")), + h.p(h.small( + join(sort(collect(string.(s) for s in _ALLOWED_CALLS)), ", "), + )), ), - h.small("Pick a branch:"), - h.fieldset(; class="grid")( - stage_button("5. VBRMI", :vbrmi), - stage_button("6. Benchmark", :bench), + h.form(; id="brm-macro-form")( + h.label("Load preset"), + h.div(; class="brm-preset-row")( + [preset(lbl, body).button for (lbl, body) in presets]..., + ), + h.label("Formula")( + h.textarea(formula; + name="formula", rows=8, + class="brm-formula-textarea"), + ), + h.fieldset(; class="grid")( + stage("1. Parse", :parse).button, + stage("2. Transform", :transform).button, + stage("3. Wrap", :wrap).button, + stage("4. BRMI", :brmi).button, + ), + h.small("Pick a branch:"), + h.fieldset(; class="grid")( + stage("5. VBRMI", :vbrmi).button, + stage("6. Benchmark", :bench).button, + ), + h.fieldset(; class="grid")( + stage("5a. SlicModel", :slic_model).button, + stage("5b. StanCode", :stan_code).button, + stage("5c. StanCompile", :stan_compile).button, + ), + h.fieldset(; class="grid")( + stage("6a. StanInstantiate", :stan_instantiate).button, + stage("6b. StanEval", :stan_eval).button, + stage("6c. StanGenerate", :stan_generate).button, + stage("6d. StanFit (PF)", :stan_fit_pathfinder).button, + stage("6d'. StanFit (Warmup)", :stan_fit_warmup).button, + ), + h.small("Bug-report helper:"), + h.button("SB repro (current formula)"; + type="submit", + formaction=string(__self__/"sb_repro"), + class="secondary"), ), - h.fieldset(; class="grid")( - stage_button("5a. SlicModel", :slic_model), - stage_button("5b. StanCode", :stan_code), - stage_button("5c. StanCompile",:stan_compile), + # Persistent wrapper — buttons swap `innerHTML` into here so the + # id survives polling/error responses. + h.div(; id="brm-macro-output")( + lazy(query_url(__self__/"stage/bench"; formula)), ), - ), - lazy(string(query_url(__self__/"stage/bench"; formula)); id="brm-macro-output"), - ) - - @get index(; formula::String=default_formula, label::String="") = begin - # If an example form posted us a (label, formula) pair, persist the - # edited formula to that example's .jl file so the next visit to the - # Examples page shows the user's edits instead of the seed default. - isempty(label) || example_store.save!(label; new_formula=formula) - index_body(formula) - end - - @get mark(; label::String="", state::String="") = begin - isempty(label) && return "" - target = Symbol(state) - entry = example_store.find(label) - entry === nothing && return "" - next_status = entry.status == target ? :open : target - updated = example_store.save!(label; new_status=next_status) - # Re-render the whole card so the border + collapse state update - # together with the pill text. - updated === nothing ? "" : updated.card(__self__) + ) end - @get stage(name::AbstractString; formula::String=default_formula, - label::String="", force::Bool=false) = begin - # When called from an example card's form, persist the (possibly - # edited) formula back to the example's .jl file before rendering. - isempty(label) || example_store.save!(label; new_formula=formula) - ns = namespace_from(label) - stage_sym = Symbol(name) - polling_fetchindex(pipeline_run, - formula, stage_sym, ns; - poll_url=string(query_url(__self__/"stage/$name"; formula, label)), - label="BRM pipeline - $name", - force) do r - render_pipeline(r, stage_sym) - end + @get stage(name::Symbol; force::Bool=false) = polling_fetchindex( + compute_steps, formula, context!().namespace, name; + poll_url=query_url(__self__/"stage/$name"; formula, label), + label="BRM pipeline - $name", + force, + ) do result + # No id on this wrapper — the outer `#brm-macro-output` div in the + # form is the persistent target (see buttons' `hx_swap="innerHTML"`); + # putting the id here too would duplicate ids after a button swap. + h.div( + (getproperty(render, k)(v) for (k, v) in pairs(result))..., + ) end # Focused per-model views of the sbimpl intermediate artifacts. Each # route runs the pipeline just far enough and returns the relevant # source in `h.pre` (plus markdown_only serves the bare source via # `?plain` / `Accept: text/plain`, for piping into agents or curl). - @get slic(; formula::String=default_formula, label::String="") = begin - isempty(label) || example_store.save!(label; new_formula=formula) - h.pre(sprint(show, run(formula, namespace_from(label)).sbbrmi.model.model)) + @get slic = h.pre(context!().run.sbbrmi.model.model) + + @get stan = h.pre(context!().run.stan.src) + + # One-stop bug-report page for the StanBlocks agent. Renders the SlicModel + # body, generated Stan source, and BridgeStan compile output (success msg + # or full error) for the current formula. HTMXO's `_resolve_response` + # auto-converts to markdown when `Accept: text/plain` is requested, so the + # same URL works for humans (browser) and agents (curl). + # curl -H 'Accept: text/plain' 'http://localhost:/pipeline/sb_repro?formula=' + @get sb_repro = begin + r = context!().run + compile_out = try + r.stan.lib + "(compile succeeded — lib at `$(r.stan.lib)`)" + catch e + sprint(showerror, e) + end + h.div( + h.h1("StanBlocks bug report"), + h.h2("Formula"), + h.pre(formula), + h.h2("SlicModel body"), + h.p(h.code("r.sbbrmi.model.model")), + h.pre(r.sbbrmi.model.model), + h.h2("Generated Stan source"), + h.p(h.code("r.stan.src")), + h.pre(r.stan.src), + h.h2("BridgeStan compile output"), + h.p(h.code("r.stan.lib")), + h.pre(compile_out), + ) end +end + +@htmx struct AppContext + __appdata__ = APPDATA + + default_formula = """loc ~ 1 +log(err) ~ 1 +y1 ~ Normal(loc, err) +""" + + # Page-level stylesheet read once at construction. Classes are consumed by + # ExampleEntry.card / html_expr.jl; per-symbol / per-status colors that + # are data-derived stay inline on the element. + css = read(joinpath(@__DIR__, "brm-macro.css"), String) + + # HTMXObjects auto-uses `__page__` to wrap any route's return value into a + # full page on direct browser navigation, while returning just the fragment + # for HTMX requests (see `_resolve_response` in HTMXObjects.jl). The + # sidebar's `hx-get` swaps target `#content` directly. + __page__(content) = htmx( + h.div(; class="brm-layout")( + nav_sidebar([ + "Pipeline" => "/pipeline", + "Examples" => "/examples", + ]), + h.main(; class="container brm-main")( + h.div(; id="content")(content), + ), + ); + pico_version="2", + extra_head=( + h.title("BRM macro action"), + h.style(css), + htmx_treebar_styles(), + vega_head()..., + ), + ) - @get stan(; formula::String=default_formula, label::String="") = begin - isempty(label) || example_store.save!(label; new_formula=formula) - h.pre(run(formula, namespace_from(label)).stan_src) + @param begin + formula::String = default_formula + label::String = "" end - @get examples(slug::String="") = begin - if !isempty(slug) - entry = example_store.find_by_slug(slug) - entry === nothing && return h.div( - h.p("No example with slug ", h.code(slug), "."), - h.a("<- Back to Examples"; href="/examples"), - ) - return h.div( - h.p(h.a("<- Back to Examples"; href="/examples")), - entry.card(__self__), - ) + # `/` mirrors the pipeline landing page. + @get index = __self__.pipeline.index + + @include pipeline = PipelineRoutes() + + # The Examples section mounts at /examples. Both the list view and per-slug + # detail view share `@get index(slug)` (HTMXO registers `/examples` AND + # `/examples/{slug}` thanks to slug's default). `@get mark` lives here too + # since it operates exclusively on ExampleEntry; pill URLs hit /examples/mark. + @include examples = begin + (; examples_dir) = __appdata__ + # `label` is auto-forwarded from AppContext's `@param`; no explicit + # `@param (; label) = __parent__` needed (and declaring it explicitly + # collides with the auto-forward → "method overwritten" error). + + # Inline examples store. Constructs ExampleEntry with + # `__parent__=__parent__` (the Examples sub-struct) so each entry's + # rendering methods can build URLs via its parent chain. + @struct example_store = begin + entries() = begin + isdir(examples_dir) || mkpath(examples_dir) + files = sort(filter(endswith(".jl"), + readdir(examples_dir; join=true)); + by=mtime, rev=true) + ExampleEntry[ExampleEntry(; __parent__=__parent__, path=f) for f in files] + end + find(label) = findfirstelement(e -> e.label == label, entries()) + find_by_slug(slug) = findfirstelement(e -> e.slug == slug, entries()) + persist!(label, formula) = begin + e = find(label) + e === nothing || e.save!(; new_formula=formula) + end end - h.div( + + @get mark(; state::Symbol=Symbol("")) = + example_store.find(label).toggle_status!(state).card + + # List view vs detail view as separate derived-property methods; DO + # supports multi-method dispatch on a single property name (the route + # layer doesn't — see TODO below). + _index() = h.div( h.h1("Examples - coverage gaps and demos for BRM"), h.p("Sorted by last modified. Each item has a sketch of what it is, why it matters, how to implement, and how to verify. Sourced from .jl files under ", h.code("web-macro/examples/"), "; status edits and edited formulas are written back to disk."), - [e.card(__self__) for e in example_store.entries()]..., + [e.card for e in example_store.entries()]..., ) + + _index(slug::AbstractString) = h.div( + h.p(h.a("<- Back to Examples"; href=__prefix__)), + example_store.find_by_slug(slug).card, + ) + + # TODO(HTMXO): allow two `@get name` methods with distinct arities + # (e.g. `@get index()` + `@get index(slug::String)`) to be registered + # as separate paths `/examples` and `/examples/{slug}`. Today DO's + # meta dict rejects duplicate route property names, so we delegate to + # the multi-methoded `_index` helper above. + @get index(slug::String="") = isempty(slug) ? _index() : _index(slug) end end From d6dfd7ec92a3a008d5cf871252c432fa3d872019 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Thu, 23 Apr 2026 01:52:14 +0200 Subject: [PATCH 20/23] web-macro: StanShapes step, truth overlays, histogram binning, misc fixes - New pipeline step 6b' StanShapes between StanEval and StanGenerate: `param_shapes_df` (one row per base param, :n_indices count), rendered as a sortable table. - Truth overlay support in `posterior_plots`: filled black Scatter on PI/LR, VLines on ECDF/Hist. Pass-through via optional `truth=...` kwarg; threaded through stan_generate, stan_fit_pathfinder, stan_fit_warmup bundles. - Histogram: `bins=30, datalimits=extrema` for per-facet local binning (workaround for `linkxaxes=:none` + faceted histogram sharing x-axis). - Scatter fill: `filled=true` kwarg so dots render filled instead of Vega's hollow `point` default. - `nonnumeric` on :Int :index to keep ECDF/hist color discrete. - ExampleEntry: added explicit `__parent__ = nothing` field so positional `ExampleEntry(path; __parent__)` works after DO constructor tightening. - Reverted earlier explicit `examples_dir = __appdata__.examples_dir` back to `(; examples_dir) = __appdata__` destructure (DO b8cb6c8 fix makes the tuple-destructure LHS now registers as a parent property for inline-child auto-forwarding). Co-Authored-By: Claude Opus 4.7 --- web-macro/src/BRMMacroWeb.jl | 55 +++++++++++++++++++++++++----------- 1 file changed, 38 insertions(+), 17 deletions(-) diff --git a/web-macro/src/BRMMacroWeb.jl b/web-macro/src/BRMMacroWeb.jl index 99ac67d..daee28c 100644 --- a/web-macro/src/BRMMacroWeb.jl +++ b/web-macro/src/BRMMacroWeb.jl @@ -14,7 +14,7 @@ using BridgeStan: BridgeStan using StanLogDensityProblems: StanProblem using WarmupHMC: initialize_mcmc, adaptive_warmup_mcmc using JSON -using AlgebraOfVega: vega_head, auto_remap_node, with_plot_caption, config, pointinterval, lineribbon, to_node, ECDFPlot, VLines +using AlgebraOfVega: vega_head, auto_remap_node, with_plot_caption, config, pointinterval, lineribbon, to_node, ECDFPlot, VLines, nonnumeric import AlgebraOfGraphics as AoG # The @brm macro and the VBRMI / SBBRMI implementations live alongside this @@ -193,6 +193,7 @@ end @dynamicstruct struct ExampleEntry path::String + __parent__ = nothing _STATUS_COLORS = (open="#888", done="#2e7d32", deprioritized="#a05a2c") _TIER_LABELS = ("T1", "T2", "T3") @@ -347,7 +348,7 @@ end end write(path, take!(io)) # Preserve __parent__ so the reloaded entry can still build route URLs. - ExampleEntry(; __parent__, path) + ExampleEntry(path; __parent__) end toggle_status!(target) = @@ -517,6 +518,17 @@ end generated_n = 50 generated_unc = 0.1 .* randn(Xoshiro(44), dim, generated_n) generated_names = BridgeStan.param_names(instance; include_tp=true, include_gq=true) + # One row per base parameter name (before any `.`), with the + # number of indexed entries. Scalar params have n_indices=1. + param_shapes_df = begin + splits = [split(nm, '.', limit=2) for nm in generated_names] + base_names = [String(first(s)) for s in splits] + uniq = unique(base_names) + DataFrame( + param = uniq, + n_indices = [count(==(p), base_names) for p in uniq], + ) + end generated = constrain_draws(generated_unc, instance; rng_seed=45) generated_dfs = dfs_from_constrained(generated, generated_names) generated_df = generated_dfs.long @@ -654,7 +666,8 @@ end stan_compile = merge(stan_code, (; stan_compile=(; file=(:stan, :file), lib=(:stan, :lib)))) stan_instantiate = merge(stan_compile, (; stan_instantiate=(; instance=(:stan, :instance), dim=(:stan, :dim), init=(:stan, :init)))) stan_eval = merge(stan_instantiate, (; stan_eval=(:stan, :log_density))) - stan_generate = merge(stan_eval, (; stan_generate=(; long=(:stan, :generated_df), wide=(:stan, :generated_wide_df), summary=(:stan, :generated_summary_df), truth=(:stan, :truth_df)))) + stan_shapes = merge(stan_eval, (; stan_shapes=(:stan, :param_shapes_df))) + stan_generate = merge(stan_shapes, (; stan_generate=(; long=(:stan, :generated_df), wide=(:stan, :generated_wide_df), summary=(:stan, :generated_summary_df), truth=(:stan, :truth_df)))) # Pathfinder / full warmup are computed via `fetchindex!` in # `compute_steps` (special-cased below by step name) so the IP's # progress subtree attaches to the step's phase. Chain-level specs @@ -662,7 +675,7 @@ end stan_fit_pathfinder = merge(stan_generate, (; stan_fit_pathfinder=(; long=(:stan, :posterior_long_df), wide=(:stan, :posterior_wide_df), summary=(:stan, :posterior_summary_df), truth=(:stan, :truth_df)))) stan_fit_warmup = merge(stan_fit_pathfinder, (; stan_fit_warmup=(; long=(:stan, :posterior_warmup_long_df), wide=(:stan, :posterior_warmup_wide_df), summary=(:stan, :posterior_warmup_summary_df), diagnostics=(:stan, :posterior_warmup_diagnostics), truth=(:stan, :truth_df)))) (; parse, transform, wrap, brmi, vbrmi, bench, slic_model, stan_code, stan_compile, - stan_instantiate, stan_eval, stan_generate, stan_fit_pathfinder, stan_fit_warmup)[name] + stan_instantiate, stan_eval, stan_shapes, stan_generate, stan_fit_pathfinder, stan_fit_warmup)[name] end # Fetch target for `polling_fetchindex`. Pre-enumerates each step in the @@ -798,6 +811,12 @@ APPDATA = AppData(; cache_type=:parallel) h.h3("6b. StanEval — log density at init"), h.p("log_density = ", x), ) + stan_shapes(df) = h.section( + h.h3("6b'. StanShapes — index count per base parameter (p + tp + gq)"), + h.p("total indexed entries: ", sum(df.n_indices), + " across ", nrow(df), " base params"), + render_table(df; sortable=true), + ) # Shared plot-tabset builder used by stan_generate and the fit stages. # `kind` goes into tab titles ("prior predictive" / "posterior") and # plot ids. Returns the tabset + wide-table details block. @@ -807,39 +826,40 @@ APPDATA = AppData(; cache_type=:parallel) ecdf_title = "$kind — ECDF" lr_title = "$kind — line + ribbon" den_title = "$kind — histogram" - indep_x = config(facet=(; linkxaxes=:none)) - indep_y = config(facet=(; linkyaxes=:none)) # Overlay layers: for (x=:index, y=:value) plots, plot truth as # filled black dots at (:index, :truth); for (x=:value) plots, # overlay vertical rules at truth values, colored by :index to - # match the base layer's coloring. + # match the base layer's (nominal-sorted) coloring. overlay_xy = isnothing(truth) ? nothing : AoG.data(truth) * AoG.mapping(:index, :truth, row=:param) * - AoG.visual(AoG.Scatter; color=:black) + AoG.visual(AoG.Scatter; color=:black, filled=true) overlay_vrule = isnothing(truth) ? nothing : - AoG.data(truth) * AoG.mapping(:truth; row=:param, color=:index) * + AoG.data(truth) * AoG.mapping(:truth; row=:param, + color=:index => nonnumeric) * AoG.visual(VLines) add(spec, overlay) = isnothing(overlay) ? spec : spec + overlay spec_pi = add(AoG.data(summary) * AoG.mapping(:index, :median, row=:param) * pointinterval(; bands, orientation=:vertical), overlay_xy) * - config(title=pi_title) * indep_y + config(title=pi_title, facet=(; linkyaxes=:none)) spec_lr = add(AoG.data(summary) * AoG.mapping(:index, :median, row=:param) * lineribbon(; bands), overlay_xy) * - config(title=lr_title) * indep_y + config(title=lr_title, facet=(; linkyaxes=:none)) spec_hist = add(AoG.data(long) * - AoG.mapping(:value; row=:param, color=:index) * + AoG.mapping(:value; row=:param, + color=:index => nonnumeric) * AoG.visual(ECDFPlot), overlay_vrule) * - config(title=ecdf_title) * indep_x + config(title=ecdf_title, facet=(; linkxaxes=:none)) spec_den = add(AoG.data(long) * - AoG.mapping(:value; row=:param, color=:index) * - AoG.histogram(), + AoG.mapping(:value; row=:param, + color=:index => nonnumeric) * + AoG.histogram(; bins=30, datalimits=extrema), overlay_vrule) * - config(title=den_title) * indep_x + config(title=den_title, facet=(; linkxaxes=:none)) tabs = tabset( "Point + Interval" => to_node(spec_pi; id="$id_prefix-pi"), "Line + Ribbon" => to_node(spec_lr; id="$id_prefix-lr"), @@ -1006,6 +1026,7 @@ bin_y ~ Bernoulli(logistic(log_odds_b)) h.fieldset(; class="grid")( stage("6a. StanInstantiate", :stan_instantiate).button, stage("6b. StanEval", :stan_eval).button, + stage("6b'. StanShapes", :stan_shapes).button, stage("6c. StanGenerate", :stan_generate).button, stage("6d. StanFit (PF)", :stan_fit_pathfinder).button, stage("6d'. StanFit (Warmup)", :stan_fit_warmup).button, @@ -1142,7 +1163,7 @@ y1 ~ Normal(loc, err) files = sort(filter(endswith(".jl"), readdir(examples_dir; join=true)); by=mtime, rev=true) - ExampleEntry[ExampleEntry(; __parent__=__parent__, path=f) for f in files] + ExampleEntry[ExampleEntry(f; __parent__) for f in files] end find(label) = findfirstelement(e -> e.label == label, entries()) find_by_slug(slug) = findfirstelement(e -> e.slug == slug, entries()) From e5be55d6ce95babb30e7ad4116ab49757c3cd126 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Thu, 23 Apr 2026 02:13:52 +0200 Subject: [PATCH 21/23] web-macro: example cards get full stage button set + SB repro; permalink hardcoded MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Example card's formula form now renders the full 15-button pipeline stage set (1. Parse through 6d'. StanFit Warmup), each GETing /pipeline/stage/:id with the card's formula+label included so bruno- prefixed examples pick up `dataset_extras(::Val{:bruno}, df)` extras. - SB repro button added to card form too; shareable URL now carries label, so the sb_repro page for a bruno example fires the same namespace dispatch as the main pipeline page. - Permalink hardcoded to `/examples/$slug` (the `__parent__/slug` form was yielding `/slug` without the `/examples` prefix; reason pending HTMXO agent investigation). - `@get index(slug::AbstractString="")` — widened slug type so URL- parsed SubString{String} matches (was String-only → MethodError on /examples/{slug}). - Removed `filled=true` kwarg on scatter overlays — AoV now defaults Scatter marks to filled. Co-Authored-By: Claude Opus 4.7 --- web-macro/src/BRMMacroWeb.jl | 54 ++++++++++++++++++++++++------------ 1 file changed, 37 insertions(+), 17 deletions(-) diff --git a/web-macro/src/BRMMacroWeb.jl b/web-macro/src/BRMMacroWeb.jl index daee28c..d59efd0 100644 --- a/web-macro/src/BRMMacroWeb.jl +++ b/web-macro/src/BRMMacroWeb.jl @@ -251,7 +251,7 @@ end # (owns /pipeline/* routes). URLs built via `query_url` so request @params # auto-propagate and values auto-encode. permalink = h.a("🔗"; - href=string(__parent__/slug), + href="/examples/$slug", title="Standalone URL", onclick="event.stopPropagation()", class="brm-permalink") @@ -277,26 +277,46 @@ end state_pill(:deprioritized, "✓ deprioritized", "deprioritize"), ) + # Full pipeline stage list, mirroring the buttons on the main pipeline + # page. Each entry renders a button that GETs the corresponding /stage/:id + # and drops the response into this card's result div. + _stage_buttons = let + stages = [ + "1. Parse" => :parse, + "2. Transform" => :transform, + "3. Wrap" => :wrap, + "4. BRMI" => :brmi, + "5. VBRMI" => :vbrmi, + "6. Benchmark" => :bench, + "5a. SlicModel" => :slic_model, + "5b. StanCode" => :stan_code, + "5c. StanCompile" => :stan_compile, + "6a. StanInstantiate" => :stan_instantiate, + "6b. StanEval" => :stan_eval, + "6b'. StanShapes" => :stan_shapes, + "6c. StanGenerate" => :stan_generate, + "6d. StanFit (PF)" => :stan_fit_pathfinder, + "6d'. StanFit (Warmup)" => :stan_fit_warmup, + ] + [h.button(label; + type="button", + class="brm-branch-btn", + hx_get=string(query_url(__parent__.__parent__.pipeline/"stage/$id"; force=true)), + hx_include="closest form", + hx_target="#$result_id", + hx_swap="innerHTML") for (label, id) in stages] + end formula_form = h.form(; class="brm-example-form")( h.input(; type="hidden", name="label", value=label), h.textarea(formula; name="formula", rows=max(3, count('\n', formula) + 1), class="brm-example-textarea"), - h.button("cimpl (bench) ▶"; - type="button", - class="brm-branch-btn", - hx_get=string(query_url(__parent__.__parent__.pipeline/"stage/bench"; force=true)), - hx_include="closest form", - hx_target="#$result_id", - hx_swap="innerHTML"), - h.button("sbimpl (compile) ▶"; - type="button", - class="brm-branch-btn", - hx_get=string(query_url(__parent__.__parent__.pipeline/"stage/stan_compile"; force=true)), - hx_include="closest form", - hx_target="#$result_id", - hx_swap="innerHTML"), + _stage_buttons..., + h.button("SB repro ▶"; + type="submit", + formaction=string(__parent__.__parent__.pipeline/"sb_repro"), + class="secondary"), ) card = begin @@ -832,7 +852,7 @@ APPDATA = AppData(; cache_type=:parallel) # match the base layer's (nominal-sorted) coloring. overlay_xy = isnothing(truth) ? nothing : AoG.data(truth) * AoG.mapping(:index, :truth, row=:param) * - AoG.visual(AoG.Scatter; color=:black, filled=true) + AoG.visual(AoG.Scatter; color=:black) overlay_vrule = isnothing(truth) ? nothing : AoG.data(truth) * AoG.mapping(:truth; row=:param, color=:index => nonnumeric) * @@ -1195,7 +1215,7 @@ y1 ~ Normal(loc, err) # as separate paths `/examples` and `/examples/{slug}`. Today DO's # meta dict rejects duplicate route property names, so we delegate to # the multi-methoded `_index` helper above. - @get index(slug::String="") = isempty(slug) ? _index() : _index(slug) + @get index(slug::AbstractString="") = isempty(slug) ? _index() : _index(slug) end end From cb98bc7f6dfcf3f69cb26b4585b94e1717cb6032 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Thu, 23 Apr 2026 02:22:21 +0200 Subject: [PATCH 22/23] web-macro: revert permalink hardcode, drop histogram datalimits workaround MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Permalink `href=string(__parent__/slug)` now that HTMXO bc32014 + 6ae7866 land (per-include __prefix__ threading + :index-always-collapses path). - Histogram workaround `datalimits=extrema` removed — AoV now handles per-facet bin extents natively when `facet=(; linkxaxes=:none)`. `bins=30` kept (pending AoV kwarg-forwarding fix). Co-Authored-By: Claude Opus 4.7 --- web-macro/src/BRMMacroWeb.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/web-macro/src/BRMMacroWeb.jl b/web-macro/src/BRMMacroWeb.jl index d59efd0..c5a2ad8 100644 --- a/web-macro/src/BRMMacroWeb.jl +++ b/web-macro/src/BRMMacroWeb.jl @@ -251,7 +251,7 @@ end # (owns /pipeline/* routes). URLs built via `query_url` so request @params # auto-propagate and values auto-encode. permalink = h.a("🔗"; - href="/examples/$slug", + href=string(__parent__/slug), title="Standalone URL", onclick="event.stopPropagation()", class="brm-permalink") @@ -877,7 +877,7 @@ APPDATA = AppData(; cache_type=:parallel) spec_den = add(AoG.data(long) * AoG.mapping(:value; row=:param, color=:index => nonnumeric) * - AoG.histogram(; bins=30, datalimits=extrema), + AoG.histogram(; bins=30), overlay_vrule) * config(title=den_title, facet=(; linkxaxes=:none)) tabs = tabset( From e435cc2d792acaac22a8458ddf4a75c0e7efdab9 Mon Sep 17 00:00:00 2001 From: Nikolas Siccha Date: Thu, 23 Apr 2026 02:33:08 +0200 Subject: [PATCH 23/23] web-macro: sb_repro_example route + Bruno SB bug example - Factor sb_repro's render into `_sb_repro_html(run, formula_str)` helper so both verb paths share it. - New `@get sb_repro_example(; name::AbstractString="")` on PipelineRoutes: looks up an ExampleEntry by slug via `__parent__.examples.example_store. find_by_slug(name)`, reads that entry's label/formula, and emits the same bug-report HTML/markdown. No persist side effect (reads from disk only). - New example `examples/sb-bug-popefs-tp-size.jl` (label "Bruno SB bug: popefs TP matrix size scope", tier 2, open) with the minimal formula that trips the `pop_loc_loc_n_covariates` TP-size scope bug, for external agents to reproduce via the above route. Acceptance: curl -s -H 'Accept: text/plain' \\ 'http://localhost:.../pipeline/sb_repro_example?name=sb-bug-popefs-tp-size' \\ | head -40 Co-Authored-By: Claude Opus 4.7 --- web-macro/examples/sb-bug-popefs-tp-size.jl | 27 +++++++++++++++++++++ web-macro/src/BRMMacroWeb.jl | 23 +++++++++++++++--- 2 files changed, 47 insertions(+), 3 deletions(-) create mode 100644 web-macro/examples/sb-bug-popefs-tp-size.jl diff --git a/web-macro/examples/sb-bug-popefs-tp-size.jl b/web-macro/examples/sb-bug-popefs-tp-size.jl new file mode 100644 index 0000000..b1ce21d --- /dev/null +++ b/web-macro/examples/sb-bug-popefs-tp-size.jl @@ -0,0 +1,27 @@ +# label: Bruno SB bug: popefs TP matrix size scope +# tier: 2 +# status: open +#= +**SB bug repro (pending SB agent fix).** Pinned to the `pop_loc_loc_n_covariates` +scope issue in StanBlocks: when a `pop_*` reduction lifts the covariate matrix +`X_loc` into `transformed parameters`, the size parameter name it emits +(`pop_loc_loc_n_covariates`) is referenced outside its declaring scope and +Stan compilation fails. + +**How to reproduce.** GET +`/pipeline/sb_repro_example?name=sb-bug-popefs-tp-size` — the rendered bug +report's "BridgeStan compile output" section shows the error. + +**Minimal formula.** `ftime` is supplied by `bruno-ext.jl` as a parameter- +derived column; referencing it forces the intercept + slope linear predictor +into `transformed parameters`, which triggers the scope bug. `y_scale` on the +assay-indexed log scale matches bruno's production residual structure. + +**What's expected once fixed.** The Stan source should compile; the generated +`transformed parameters` block should declare `int pop_loc_loc_n_covariates` +in the same scope where it's consumed (or inline the size expression). +=# + +loc ~ 1 + ftime +log(y_scale) ~ 0 + assay_idx +y ~ Normal(loc, y_scale) diff --git a/web-macro/src/BRMMacroWeb.jl b/web-macro/src/BRMMacroWeb.jl index c5a2ad8..a5e2aa2 100644 --- a/web-macro/src/BRMMacroWeb.jl +++ b/web-macro/src/BRMMacroWeb.jl @@ -1093,8 +1093,10 @@ bin_y ~ Bernoulli(logistic(log_odds_b)) # auto-converts to markdown when `Accept: text/plain` is requested, so the # same URL works for humans (browser) and agents (curl). # curl -H 'Accept: text/plain' 'http://localhost:/pipeline/sb_repro?formula=' - @get sb_repro = begin - r = context!().run + # Shared renderer: takes a ready `run` context + the raw formula string and + # emits the bug-report HTML. Same output whether invoked via POST with an + # edited formula (`sb_repro`) or via GET by example slug (`sb_repro_example`). + _sb_repro_html(r, formula_str) = begin compile_out = try r.stan.lib "(compile succeeded — lib at `$(r.stan.lib)`)" @@ -1104,7 +1106,7 @@ bin_y ~ Bernoulli(logistic(log_odds_b)) h.div( h.h1("StanBlocks bug report"), h.h2("Formula"), - h.pre(formula), + h.pre(formula_str), h.h2("SlicModel body"), h.p(h.code("r.sbbrmi.model.model")), h.pre(r.sbbrmi.model.model), @@ -1116,6 +1118,21 @@ bin_y ~ Bernoulli(logistic(log_odds_b)) h.pre(compile_out), ) end + + @get sb_repro = _sb_repro_html(context!().run, formula) + + # Saved-example entry point for external agents: GET by slug so curl/agents + # can reproduce a StanBlocks bug without POSTing a formula. The slug is the + # URL-safe name used by /examples/. + # curl -H 'Accept: text/plain' 'http://.../pipeline/sb_repro_example?name=' + @get sb_repro_example(; name::AbstractString="") = begin + entry = __parent__.examples.example_store.find_by_slug(name) + isnothing(entry) && return h.div( + h.h1("Example not found"), + h.p("No example with slug ", h.code(name)), + ) + _sb_repro_html(context(entry.label, entry.formula).run, entry.formula) + end end @htmx struct AppContext