added missing dep to test env

docs/make.jl

@@ -30,8 +30,7 @@ using Pkg
 if any(contains.(ARGS, "--project"))
     @assert sum(contains.(ARGS, "--project")) == 1 "Only one environment can be specified using the `--project` argument."
     _path =
-        ARGS[findall(contains.(ARGS, "--project"))][1] |>
-        x -> replace(x, "--project=" => "")
+        (x -> replace(x, "--project=" => ""))(ARGS[findall(contains.(ARGS, "--project"))][1])
     Pkg.activate(_path)
 else
     Pkg.activate(@__DIR__)
@@ -57,12 +56,7 @@ makedocs(;
         edit_link="main",
         assets=String[],
     ),
-    pages=[
-        "Home" => "index.md",
-    ],
+    pages=["Home" => "index.md"],
 )
 
-deploydocs(;
-    repo="github.com/JuliaTrustworthyAI/EnergySamplers.jl",
-    devbranch="main",
-)
+deploydocs(; repo="github.com/JuliaTrustworthyAI/EnergySamplers.jl", devbranch="main")

src/optimizers.jl

@@ -16,10 +16,9 @@ struct SGLD <: AbstractSamplingRule
     b::Float64
     gamma::Float64
 end
-SGLD(; a::Real = 10.0, b::Real = 1000.0, γ::Real = 0.9) = SGLD(a, b, γ)
+SGLD(; a::Real=10.0, b::Real=1000.0, γ::Real=0.9) = SGLD(a, b, γ)
 
 function Optimisers.apply!(o::SGLD, state, x::AbstractArray{T}, Δ) where {T}
-
     a, b, γ = T(o.a), T(o.b), T(o.gamma)
 
     εt = @.(a * (b + state)^-γ)
@@ -48,7 +47,7 @@ struct ImproperSGLD <: AbstractSamplingRule
     alpha::Float64
     sigma::Float64
 end
-ImproperSGLD(; α::Real = 2.0, σ::Real = 0.01) = ImproperSGLD(α, σ)
+ImproperSGLD(; α::Real=2.0, σ::Real=0.01) = ImproperSGLD(α, σ)
 
 function Optimisers.apply!(o::ImproperSGLD, state, x::AbstractArray{T}, Δ) where {T}
     α, σ = T(o.alpha), T(o.sigma)

src/utils.jl

@@ -1,22 +1,22 @@
 using Flux
 using StatsBase
 
-get_logits(f::Flux.Chain, x) = f[end] isa Function ? f[1:end-1](x) : f(x)
+get_logits(f::Flux.Chain, x) = f[end] isa Function ? f[1:(end - 1)](x) : f(x)
 
 @doc raw"""
     energy(f, x)
 
 Computes the energy for unconditional samples $x \sim p_{\theta}(x)$: $E(x)=-\text{LogSumExp}_y f_{\theta}(x)[y]$.
 """
-function _energy(f, x; agg = mean)
+function _energy(f, x; agg=mean)
     if f isa Flux.Chain
         ŷ = get_logits(f, x)
     else
         ŷ = f(x)
     end
     if ndims(ŷ) > 1
         E = 0.0
-        E = agg(map(y -> -logsumexp(y), eachslice(ŷ, dims = ndims(ŷ))))
+        E = agg(map(y -> -logsumexp(y), eachslice(ŷ; dims=ndims(ŷ))))
         return E
     else
         return -logsumexp(ŷ)
@@ -28,7 +28,7 @@ end
 
 Computes the energy for conditional samples $x \sim p_{\theta}(x|y)$: $E(x)=- f_{\theta}(x)[y]$.
 """
-function _energy(f, x, y::Int; agg = mean)
+function _energy(f, x, y::Int; agg=mean)
     if f isa Flux.Chain
         ŷ = get_logits(f, x)
     else
@@ -37,7 +37,7 @@ function _energy(f, x, y::Int; agg = mean)
     _E(y, idx) = length(y) > 1 ? -y[idx] : (idx == 2 ? -y[1] : -(1.0 - y[1]))
     if ndims(ŷ) > 1
         E = 0.0
-        E = agg(map(_y -> _E(_y, y), eachslice(ŷ, dims = ndims(ŷ))))
+        E = agg(map(_y -> _E(_y, y), eachslice(ŷ; dims=ndims(ŷ))))
         return E
     else
         return _E(_y, y)

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+MLDataUtils = "cc2ba9b6-d476-5e6d-8eaf-a92d5412d41d"
 MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/samplers.jl

@@ -12,14 +12,13 @@ using EnergySamplers:
     PMC
 
 @testset "Samplers" begin
-
     f(x) = @.(2x + 1)  # dummy model
     nn = Chain(Dense(2, 1, σ))
     rule = ImproperSGLD()
 
     # Data:
     nobs = 2000
-    X, y = make_circles(nobs, noise = 0.1, factor = 0.5)
+    X, y = make_circles(nobs; noise=0.1, factor=0.5)
     Xmat = Float32.(permutedims(matrix(X)))
     X = table(permutedims(Xmat))
     batch_size = Int(round(nobs / 10))
@@ -36,23 +35,23 @@ using EnergySamplers:
 
     for (name, Sampler) in all_samplers
         @testset "$name" begin
-            smpler =
-                Sampler(𝒟x, 𝒟y, input_size = size(Xmat)[1:end-1], batch_size = batch_size)
+            smpler = Sampler(
+                𝒟x, 𝒟y; input_size=size(Xmat)[1:(end - 1)], batch_size=batch_size
+            )
             @test smpler isa AbstractSampler
 
-            X̂ = smpler(f, rule; n_samples = 10)
+            X̂ = smpler(f, rule; n_samples=10)
             @test size(X̂, 2) == 10
 
-            X̂ = smpler(nn, rule; n_samples = 10)
+            X̂ = smpler(nn, rule; n_samples=10)
             @test size(X̂, 2) == 10
-
         end
     end
 
     @testset "Persistent Markov Chains (PMC)" begin
 
         # Train a simple neural network on the data (classification)
-        Xtrain = MLJBase.matrix(X) |> permutedims
+        Xtrain = permutedims(MLJBase.matrix(X))
         ytrain = Flux.onehotbatch(y, levels(y))
         train_set = zip(eachcol(Xtrain), eachcol(ytrain))
         inputdim = size(first(train_set)[1], 1)
@@ -61,7 +60,7 @@ using EnergySamplers:
         loss(yhat, y) = Flux.logitcrossentropy(yhat, y)
         opt_state = Flux.setup(Flux.Adam(), nn)
         epochs = 5
-        for epoch = 1:epochs
+        for epoch in 1:epochs
             Flux.train!(nn, train_set, opt_state) do m, x, y
                 loss(m(x), y)
             end
@@ -74,13 +73,15 @@ using EnergySamplers:
         ntrans = 100
         niter = 20
         # Conditionally sample from first class:
-        smpler =
-            ConditionalSampler(𝒟x, 𝒟y, input_size = size(Xmat)[1:end-1], batch_size = bs)
-        x1 = PMC(smpler, nn, ImproperSGLD(); ntransitions = ntrans, niter = niter, y = 1)
+        smpler = ConditionalSampler(
+            𝒟x, 𝒟y; input_size=size(Xmat)[1:(end - 1)], batch_size=bs
+        )
+        x1 = PMC(smpler, nn, ImproperSGLD(); ntransitions=ntrans, niter=niter, y=1)
         # Conditionally sample from second class:
-        smpler =
-            ConditionalSampler(𝒟x, 𝒟y, input_size = size(Xmat)[1:end-1], batch_size = bs)
-        x2 = PMC(smpler, nn, ImproperSGLD(); ntransitions = ntrans, niter = niter, y = 2)
+        smpler = ConditionalSampler(
+            𝒟x, 𝒟y; input_size=size(Xmat)[1:(end - 1)], batch_size=bs
+        )
+        x2 = PMC(smpler, nn, ImproperSGLD(); ntransitions=ntrans, niter=niter, y=2)
 
         # using Plots
         # plt = scatter(Xtrain[1, :], Xtrain[2, :], color=Int.(y.refs), group=Int.(y.refs), label=["X|y=0" "X|y=1"], alpha=0.1)
@@ -107,28 +108,28 @@ using EnergySamplers:
         target = :Default
         numerics = [:Balance, :Income]
         features = [:Student, :Balance, :Income]
-        train, test = shuffleobs(data) |> d -> stratifiedobs(first, d, p = 0.7)
+        train, test = (d -> stratifiedobs(first, d; p=0.7))(shuffleobs(data))
 
         for feature in numerics
-            μ, σ = rescale!(train[!, feature], obsdim = 1)
-            rescale!(test[!, feature], μ, σ, obsdim = 1)
+            μ, σ = rescale!(train[!, feature]; obsdim=1)
+            rescale!(test[!, feature], μ, σ; obsdim=1)
         end
 
-        prep_X(x) = Matrix(x)' |> gpu
-        prep_y(y) = reshape(y, 1, :) |> gpu
+        prep_X(x) = gpu(Matrix(x)')
+        prep_y(y) = gpu(reshape(y, 1, :))
         train_X, test_X = prep_X.((train[:, features], test[:, features]))
         train_y, test_y = prep_y.((train[:, target], test[:, target]))
-        train_set = Flux.DataLoader((train_X, train_y), batchsize = 100, shuffle = false)
+        train_set = Flux.DataLoader((train_X, train_y); batchsize=100, shuffle=false)
 
-        function train_logreg(; steps::Int = 1000, opt = Flux.Descent(2))
+        function train_logreg(; steps::Int=1000, opt=Flux.Descent(2))
             Random.seed!(1)
 
             paramvec(θ) = reduce(hcat, cpu(θ))
-            model = Dense(length(features), 1, sigmoid) |> gpu
+            model = gpu(Dense(length(features), 1, sigmoid))
             θ = Flux.params(model)
             θ₀ = paramvec(θ)
 
-            predict(x; thres = 0.5) = model(x) .> thres
+            predict(x; thres=0.5) = model(x) .> thres
             accuracy(x, y) = mean(cpu(predict(x)) .== cpu(y))
 
             loss(yhat, y) = Flux.binarycrossentropy(yhat, y)
@@ -142,10 +143,8 @@ using EnergySamplers:
 
             opt_state = Flux.setup(opt, model)
 
-            for t = 1:steps
-
+            for t in 1:steps
                 for data in train_set
-
                     input, label = data
 
                     # Calculate the gradient of the objective
@@ -159,27 +158,25 @@ using EnergySamplers:
                 end
 
                 # Bookkeeping
-                weights[t+1, :] = cpu(paramvec(θ))
-                trainlosses[t+1] = cpu(trainloss())
-                testlosses[t+1] = cpu(testloss())
+                weights[t + 1, :] = cpu(paramvec(θ))
+                trainlosses[t + 1] = cpu(trainloss())
+                testlosses[t + 1] = cpu(testloss())
             end
 
             println("Final parameters are $(paramvec(θ))")
             println("Test accuracy is $(accuracy(test_X, test_y))")
 
-            model, weights, trainlosses, testlosses
+            return model, weights, trainlosses, testlosses
         end
 
-        results = train_logreg(steps = 1000, opt = SGLD(10.0, 1000.0, 0.9))
+        results = train_logreg(; steps=1000, opt=SGLD(10.0, 1000.0, 0.9))
         model, weights, trainlosses, testlosses = results
-        plot(weights; label = ["Student" "Balance" "Income" "Intercept"])
+        plot(weights; label=["Student" "Balance" "Income" "Intercept"])
 
-        results = train_logreg(steps = 100, opt = ImproperSGLD(2.0, 0.01))
+        results = train_logreg(; steps=100, opt=ImproperSGLD(2.0, 0.01))
         model, weights, trainlosses, testlosses = results
-        plot(weights; label = ["Student" "Balance" "Income" "Intercept"])
+        plot(weights; label=["Student" "Balance" "Income" "Intercept"])
 
         @test true
-
     end
-
 end