Docs/70 text refresh

.JuliaFormatter.toml

@@ -4,4 +4,4 @@ normalize_line_endings = "unix"
 v2_stable_multiline_strings = true
 conditional_to_if = false
 pipe_to_function_call = false
-style="blue"
+style="blue"

.github/workflows/Test.yml

@@ -25,8 +25,8 @@ jobs:
           - "1"
         os:
           - ubuntu-latest
-          
-          
+
+
         arch:
           - x64
         allow_failure: [false]

.gitignore

@@ -29,4 +29,4 @@ CODEX.md
 
 # random scripts for debugging
 /scripts
-/dev
+/dev

README.md

@@ -23,28 +23,34 @@
 
 ## What this package does
 
-**ParallelMCMC.jl** implements *parallel-across-the-sequence* MCMC in Julia: instead of generating samples one at a time, an entire trajectory of $T$ correlated steps is solved *simultaneously*. This makes wall-clock time per sample sublinear in chain length on multi-core CPUs and GPUs, where conventional sequential MCMC scales linearly.
+ParallelMCMC.jl implements MCMC methods that parallelize *within* a chain. Its
+main sampler uses DEER to solve a block of MALA transitions together instead of
+waiting for each transition before starting the next one.
 
-The flagship algorithm is **DEER** (Lim et al. 2024; Gonzalez et al. 2024), which reformulates a chain of $T$ MALA steps as a fixed-point problem and solves it with Newton iterations. Each iteration linearizes the per-step transition around the current trajectory guess and resolves the resulting linear recursion in $O(\log T)$ parallel work via an associative prefix scan. With shared input randomness, DEER converges to the exact sequential MALA trace up to a numerical tolerance — typically in tens of iterations even for chains of tens of thousands of samples.
+This is parallel-across-the-sequence MCMC, not just several independent chains
+running at once. It is most useful when density and derivative evaluations are
+expensive enough to keep parallel hardware busy. Small CPU targets may still be
+faster with ordinary sequential MALA.
 
-The approach and its scaling tricks (stochastic Hutchinson Jacobian estimators, damping, sliding windows) are described in:
+The method is described in:
 
 > Zoltowski, D. M., Wu, S., Gonzalez, X., Kozachkov, L., & Linderman, S. W. (2025).
 > **Parallelizing MCMC Across the Sequence Length.** *NeurIPS 2025.*
 > [arXiv:2508.18413](https://arxiv.org/abs/2508.18413)
 
-### Samplers
+## Samplers
 
 | Sampler | Role |
 |---|---|
-| [`ParallelMALASampler`](src/interface.jl) | **Primary** — parallel-across-sequence MALA via DEER; $O(\log T)$ per solve |
-| [`MALASampler`](src/interface.jl) | Baseline — sequential MALA with a fixed step size |
-| [`AdaptiveMALASampler`](src/interface.jl) | Baseline — sequential MALA with dual-averaging step-size adaptation |
+| [`ParallelMALASampler`](docs/src/95-reference.md) | Parallel-across-sequence MALA via DEER |
+| [`AdaptiveMALASampler`](docs/src/95-reference.md) | Sequential MALA with step-size adaptation |
+| [`MALASampler`](docs/src/95-reference.md) | Sequential MALA with a fixed step size |
 
-All samplers implement the [AbstractMCMC](https://github.com/TuringLang/AbstractMCMC.jl) interface and return [`FlexiChains`](https://pysm.dev/FlexiChains.jl) objects, so they slot into existing AbstractMCMC workflows.
+All three implement the
+[AbstractMCMC](https://github.com/TuringLang/AbstractMCMC.jl) interface and
+return [FlexiChains](https://pysm.dev/FlexiChains.jl) objects.
 
-
-### Quick start
+## Quick start
 
 ParallelMCMC can be installed via Julia's package manager. In the Julia REPL, press `]` to enter pkg mode, then run:
 
@@ -53,22 +59,34 @@ pkg> add ParallelMCMC
 ```
 
 ```julia
-using ParallelMCMC, FlexiChains
-using ADTypes, Enzyme
-
-logp(x)      = -0.5 * sum(abs2, x)            # 2-D standard normal
-grad_logp(x) = -x
-
-model   = DensityModel(logp, grad_logp, 2; param_names=[:x1, :x2])
-sampler = ParallelMALASampler(0.1; T=64, jacobian=:stoch_diag,
-                              backend=AutoEnzyme())
-
-chain = sample(model, sampler, 500; chain_type=VNChain)
+using ParallelMCMC, FlexiChains, Random
+
+logdensity(x) = -sum(abs2, x) / 2
+gradient(x) = -x
+hvp(x, v) = -v
+
+model = DensityModel(
+    logdensity,
+    gradient,
+    2;
+    hvp=hvp,
+    param_names=[:x1, :x2],
+)
+
+sampler = ParallelMALASampler(0.1; T=64)
+chain = sample(
+    MersenneTwister(42), model, sampler, 1_000;
+    initial_params=zeros(2),
+    chain_type=VNChain,
+)
 ```
 
-See the [Getting Started guide](docs/src/10-getting-started.md) for worked examples, GPU usage, and step-size tuning.
+Start with the [documentation](https://rsenne.github.io/ParallelMCMC.jl/dev/)
+or the repository's [getting-started guide](docs/src/10-getting-started.md).
+The docs cover step-size tuning, model and AD setup, GPU execution,
+troubleshooting, and the DEER algorithm.
 
-## How to Cite
+## How to cite
 
 If you use ParallelMCMC.jl in your work, please cite using the reference given in [CITATION.cff](https://github.com/rsenne/ParallelMCMC.jl/blob/main/CITATION.cff).
 

docs/make.jl

@@ -4,60 +4,19 @@ using Documenter.Remotes: GitHub
 
 DocMeta.setdocmeta!(ParallelMCMC, :DocTestSetup, :(using ParallelMCMC); recursive=true)
 
-# Add titles of sections and overrides page titles
-const titles = Dict(
-    # "10-tutorials" => "Tutorials", # example folder title
-    "91-developer.md" => "Developer docs",
-)
-
-function recursively_list_pages(folder; path_prefix="")
-    pages_list = Any[]
-    for file in readdir(folder)
-        if file == "index.md"
-            # We add index.md separately to make sure it is the first in the list
-            continue
-        end
-        # this is the relative path according to our prefix, not @__DIR__, i.e., relative to `src`
-        relpath = joinpath(path_prefix, file)
-        # full path of the file
-        fullpath = joinpath(folder, relpath)
-
-        if isdir(fullpath)
-            # If this is a folder, enter the recursion case
-            subsection = recursively_list_pages(fullpath; path_prefix=relpath)
-
-            # Ignore empty folders
-            if length(subsection) > 0
-                title = if haskey(titles, relpath)
-                    titles[relpath]
-                else
-                    @error "Bad usage: '$relpath' does not have a title set. Fix in 'docs/make.jl'"
-                    relpath
-                end
-                push!(pages_list, title => subsection)
-            end
-
-            continue
-        end
-
-        if splitext(file)[2] != ".md" # non .md files are ignored
-            continue
-        elseif haskey(titles, relpath) # case 'title => path'
-            push!(pages_list, titles[relpath] => relpath)
-        else # case 'title'
-            push!(pages_list, relpath)
-        end
-    end
-
-    return pages_list
-end
-
-function list_pages()
-    root_dir = joinpath(@__DIR__, "src")
-    pages_list = recursively_list_pages(root_dir)
-
-    return ["index.md"; pages_list]
-end
+const pages = [
+    "Home" => "index.md",
+    "Getting started" => "10-getting-started.md",
+    "Defining models" => "12-models.md",
+    "GPU execution" => "15-gpu.md",
+    "Troubleshooting" => "16-troubleshooting.md",
+    "How the algorithm works" => "20-algorithms.md",
+    "API reference" => "95-reference.md",
+    "Contributing" => [
+        "Contributor guide" => "90-contributing.md",
+        "Developer guide" => "91-developer.md",
+    ],
+]
 
 makedocs(;
     modules=[ParallelMCMC],
@@ -70,7 +29,7 @@ makedocs(;
         repolink="https://github.com/rsenne/ParallelMCMC.jl",
         edit_link="main",
     ),
-    pages=list_pages(),
+    pages=pages,
 )
 
 deploydocs(; repo="github.com/rsenne/ParallelMCMC.jl", devbranch="main")