Abstract mcmc rwm

.gitignore

@@ -4,3 +4,6 @@
 /Manifest*.toml
 /docs/Manifest*.toml
 /docs/build/
+
+
+**/.vscode

Project.toml

@@ -1,9 +1,23 @@
 name = "Arianna"
 uuid = "17c2f236-ed37-45c0-bff7-21ff0c2ec05e"
-authors = ["Matt Graham", "Ross Ah-Weng", "Callum Lau", "Anees Hussain", "Jordan Simbananiye and contributors"]
 version = "1.0.0-DEV"
+authors = ["Matt Graham", "Ross Ah-Weng", "Callum Lau", "Anees Hussain", "Jordan Simbananiye and contributors"]
+
+[deps]
+AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [compat]
+AbstractMCMC = "5.10.0"
+Distributions = "0.25.122"
+LinearAlgebra = "1.12.0"
+LogDensityProblems = "2.2.0"
+Plots = "1.41.1"
+Random = "1.11.0"
 julia = "1.6.7"
 
 [extras]

src/Arianna.jl

@@ -1,5 +1,6 @@
 module Arianna
 
-# Write your package code here.
+include("RandomWalk.jl")
+using .RandomWalk
 
 end

src/RandomWalk.jl

@@ -0,0 +1,78 @@
+module RandomWalk
+
+using AbstractMCMC
+using Random
+using Distributions
+using LogDensityProblems
+
+# Step 1 — Model
+struct DistributionModel <: AbstractMCMC.AbstractModel
+    dist::Distribution
+end
+
+LogDensityProblems.logdensity(model::DistributionModel, x) = 
+    logpdf(model.dist, x)
+
+LogDensityProblems.dimension(model::DistributionModel) = length(mean(model.dist))
+
+
+# Step 2 - Sampler
+struct RWSampler <: AbstractMCMC.AbstractSampler 
+    position::Vector{Float64}
+    stepsize::Float64
+end
+
+# Step 3 - Random Walk Metropolis Step
+function AbstractMCMC.step(
+    rng::AbstractRNG,
+    model::DistributionModel,
+    sampler::RWSampler
+    )
+
+    proposal = sampler.position .+ sampler.stepsize .* randn(rng, length(sampler.position))
+
+    logp_current = LogDensityProblems.logdensity(model, sampler.position)
+    logp_proposal = LogDensityProblems.logdensity(model, proposal)
+
+    log_accept_ratio = logp_proposal - logp_current
+
+    if log(rand(rng)) < log_accept_ratio
+        new_position = proposal
+    else
+        new_position = sampler.position
+    end
+
+    return RWSampler(new_position, sampler.stepsize), logp_current
+end
+
+function AbstractMCMC.step(rng::AbstractRNG,
+                           model::DistributionModel,
+                           sampler::RWSampler,
+                           _weight::Float64)
+
+    # Multiple dispatch, ignore the weight and call the usual step
+    return AbstractMCMC.step(rng, model, sampler)
+end
+
+# Step 4 - Sampling Interface
+function AbstractMCMC.sample(
+    rng::AbstractRNG,
+    model::DistributionModel,
+    sampler::RWSampler,
+    n_samples::Integer,
+    kwargs...
+    )
+
+    d = length(sampler.position)
+    samples = Matrix{Float64}(undef, n_samples, d)
+
+    current_sampler = sampler 
+
+    for i in 1:n_samples
+        current_sampler, logp = AbstractMCMC.step(rng, model, current_sampler)
+        samples[i, :] = current_sampler.position
+    end
+
+    return samples
+end
+end

test/runtests.jl

@@ -1,6 +1,51 @@
-using Arianna
 using Test
+using Arianna
+import Arianna.RandomWalk
+using Distributions
+using LogDensityProblems
+using Random
+using AbstractMCMC
+using Plots
+using LinearAlgebra
+
+@testset "Multidimensional Random Walk" begin
+    rng = Random.default_rng()
+
+    # 2D Gaussian
+    dist = MvNormal([0.0, 0.0], I(2))
+    model = RandomWalk.DistributionModel(dist)
+
+    # Initial sampler
+    s0 = RandomWalk.RWSampler([0.0, 0.0], 0.1)
+    @test s0.position isa Vector{Float64}
+    @test length(s0.position) == 2
+
+    # One step
+    s1, logp = AbstractMCMC.step(rng, model, s0)
+    @test s1.position isa Vector{Float64}
+    @test length(s1.position) == 2
+    @test logp isa Float64
 
-@testset "Arianna.jl" begin
-    # Write your tests here.
+    # Multiple samples
+    samples = AbstractMCMC.sample(rng, model, s0, 100)
+    @test samples isa Matrix{Float64}
+    @test size(samples) == (100, 2)
+
+    # # Plot each dimension separately
+    # x = samples[:,1]
+    # y = samples[:,2]
+    # t = 1:size(samples,1)
+
+    # p = plot3d(x, y, t,
+    # xlabel = "x₁",
+    # ylabel = "x₂",
+    # zlabel = "Step",
+    # title = "3D Random Walk Trajectory",
+    # linealpha = 5,
+    # legend = false)
+
+    # savefig("trace_mv.png")
 end
+
+
+