Add maxwell boltzmann (#40)

Project.toml

@@ -1,6 +1,12 @@
 name = "QEDevents"
 uuid = "fc3ce04a-5be5-4f3a-acff-eceaab723759"
-authors = ["Uwe Hernandez Acosta <[email protected]>", "Simeon Ehrig", "Klaus Steiniger", "Tom Jungnickel", "Anton Reinhard"]
+authors = [
+    "Uwe Hernandez Acosta <[email protected]>",
+    "Simeon Ehrig",
+    "Klaus Steiniger",
+    "Tom Jungnickel",
+    "Anton Reinhard",
+]
 version = "0.1.0"
 
 [deps]
@@ -9,18 +15,21 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 QEDbase = "10e22c08-3ccb-4172-bfcf-7d7aa3d04d93"
 QEDcore = "35dc0263-cb5f-4c33-a114-1d7f54ab753e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
 Distributions = "0.25"
 DocStringExtensions = "0.9"
 QEDbase = "0.2.2"
 QEDcore = "0.1"
+StatsBase = "0.34"
 julia = "1.6"
 
 [extras]
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Random", "SafeTestsets", "Test"]
+test = ["LinearAlgebra", "Random", "SafeTestsets", "Test"]

src/QEDevents.jl

@@ -1,11 +1,14 @@
 module QEDevents
 
-export ParticleSampleable, weight
+export ParticleSampleable, weight, max_weight
 
 export SingleParticleDistribution
 export MultiParticleDistribution
 export ScatteringProcessDistribution
 
+# single particle distributions
+export MaxwellBoltzmannParticle, temperature
+
 import Random: AbstractRNG
 import Distributions: rand, rand!, _rand!
 using Distributions: Distributions
@@ -15,11 +18,16 @@ using QEDcore
 
 using DocStringExtensions
 
+# patch Distributions.jl
+include("patch_Distributions.jl")
+export MaxwellBoltzmann
+
 include("utils.jl")
 
 include("interfaces/particle_distribution.jl")
 include("interfaces/single_particle_distribution.jl")
 include("interfaces/multi_particle_distribution.jl")
 include("interfaces/process_distribution.jl")
 
+include("sampler/single_particle_dists/maxwell_boltzmann.jl")
 end

src/patch_Distributions.jl

@@ -0,0 +1,103 @@
+### Maxwell Boltzmann distribution
+# https://en.wikipedia.org/wiki/Maxwell–Boltzmann_distribution
+
+const SQRT_TWO_OVER_PI = sqrt(2 / pi)
+const _CHI3 = Distributions.Chi(3; check_args=false)
+
+"""
+    MaxwellBoltzmann(scale::Real)
+
+
+The *Maxwell-Boltzmann distribution* with scale parameter `a` has the probability density function
+
+```math
+f(x,a) = \\sqrt{\\frac{2}{\\pi}}\\frac{x^2}{a^3}\\exp\\left(\\frac{-x^2}{2a^2}\\right)
+```
+
+The Maxwell-Boltzmann distribution is related to the `Chi` distribution via the property
+``X\\sim \\operatorname{MaxwellBoltzmann}(a=1)``, then ``X\\sim\\chi(\\mathrm{dof}=3)``.
+
+
+External links
+
+* [Maxwell-Boltzmann distribution on Wikipedia](https://en.wikipedia.org/wiki/Maxwell–Boltzmann_distribution)
+* [Maxwell-Boltzmann distribution on Wolfram MathWorld](https://mathworld.wolfram.com/MaxwellDistribution.html)
+* [Maxwell-Boltzmann distribution implementation in Scipy](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.maxwell.html)
+"""
+struct MaxwellBoltzmann{T<:Real} <: Distributions.ContinuousUnivariateDistribution
+    a::T # scale
+    MaxwellBoltzmann{T}(a::T) where {T} = new{T}(a)
+end
+
+function MaxwellBoltzmann(a::T; check_args::Bool=true) where {T<:Real}
+    Distributions.@check_args MaxwellBoltzmann (a, a > zero(a))
+    return MaxwellBoltzmann{T}(a)
+end
+
+function MaxwellBoltzmann(a::Integer; check_args::Bool=true)
+    return MaxwellBoltzmann(float(a); check_args=check_args)
+end
+MaxwellBoltzmann() = MaxwellBoltzmann(1.0)
+
+Distributions.@distr_support MaxwellBoltzmann 0.0 Inf
+
+### Conversions
+convert(::Type{MaxwellBoltzmann{T}}, a::S) where {T<:Real,S<:Real} = MaxwellBoltzmann(T(a))
+function Base.convert(::Type{MaxwellBoltzmann{T}}, d::MaxwellBoltzmann) where {T<:Real}
+    return MaxwellBoltzmann(T(d.a))
+end
+Base.convert(::Type{MaxwellBoltzmann{T}}, d::MaxwellBoltzmann{T}) where {T<:Real} = d
+
+### Parameters
+Distributions.scale(d::MaxwellBoltzmann) = d.a
+Distributions.params(d::MaxwellBoltzmann) = (d.a,)
+Distributions.partype(d::MaxwellBoltzmann{T}) where {T} = T
+
+### Statistics
+Distributions.mean(d::MaxwellBoltzmann) = 2 * SQRT_TWO_OVER_PI * d.a
+Distributions.median(d::MaxwellBoltzmann{T}) where {T<:Real} = d.a * T(1.5381722544550522) # a * sqrt(2 Q^(-1)(3/2, 1/2))
+Distributions.mode(::MaxwellBoltzmann{T}) where {T<:Real} = sqrt(2) * d.a
+
+Distributions.var(d::MaxwellBoltzmann) = d.θ^2 * (3 * pi - 8) / pi
+function Distributions.skewness(::MaxwellBoltzmann{T}) where {T}
+    return T(2 * sqrt(2) * (16 - 5 * pi) / (3 * pi - 8)^(3 / 2))
+end
+function Distributions.kurtosis(::MaxwellBoltzmann{T}) where {T}
+    return T(4 * (-96 + 40 * pi - 3 * pi^2) / ((3 * pi - 8)^2))
+end
+
+function Distributions.entropy(d::MaxwellBoltzmann{T}) where {T}
+    # 0.5772156649 is the Euler-Machgeroni constant
+    return T(log(d.a * sqrt(2 * pi) + 0.5772156649 - 1 / 2))
+end
+
+#function kldivergence(p::MaxwellBoltzmann, q::MaxwellBoltzmann)
+#   # TODO:Implement this!
+#end
+
+### pdf, cdf, ...
+# derived from Chi(3)
+
+for func in (
+    :pdf,
+    :logpdf,
+    :cdf,
+    :ccdf,
+    :logcdf,
+    :logccdf,
+    :quantile,
+    :cquantile,
+    :invlogcdf,
+    :invlogccdf,
+)
+    eval(
+        quote
+            function ($Distributions.$func)(d::MaxwellBoltzmann, x::Real)
+                return (a = d.a; $Distributions.$func($_CHI3, x / a) / a)
+            end
+        end,
+    )
+end
+
+### sampling
+rand(rng::AbstractRNG, d::MaxwellBoltzmann) = rand(rng, _CHI3) * d.a

src/sampler/single_particle_dists/maxwell_boltzmann.jl

@@ -0,0 +1,96 @@
+"""
+MaxwellBoltzmannParticle(
+    dir::ParticleDirection,
+    part::AbstractParticleType,
+    temperature::Real
+)
+
+The Maxwell-Boltzmann particle distribution is a single-particle distribution, where the
+spatial components of the four-momentum are normally distributed with localion ``\\mu = 0``
+and variance ``\\sigma = m k_B T`` (``m`` is the particle's mass, ``k_B`` is the Boltzmann constant,
+and ``T`` is the temperature). The three-magnitude ``\\varrho = \\sqrt{p_x^2 + p_y^2 + p_z^2}`` of the generated
+particle-statefuls is [`MaxwellBoltzmann`](@ref) distributed.
+
+External links
+
+* [Maxwell-Boltzmann distributed four-momenta on Wikipedia](https://en.wikipedia.org/wiki/Maxwell–Boltzmann_distribution#Distribution_for_the_momentum_vector)
+
+"""
+struct MaxwellBoltzmannParticle{D,P,T,DIST} <: SingleParticleDistribution
+    dir::D
+    part::P
+    temperature::T
+    rho_dist::DIST
+
+    function MaxwellBoltzmannParticle(
+        dir::D, particle::P, temperature::T
+    ) where {D<:ParticleDirection,P<:AbstractParticleType,T<:Real}
+        a = sqrt(mass(particle) * temperature)
+        return new{D,P,T,MaxwellBoltzmann{T}}(
+            dir, particle, temperature, MaxwellBoltzmann(a)
+        )
+    end
+end
+
+_particle(d::MaxwellBoltzmannParticle) = d.part
+_particle_direction(d::MaxwellBoltzmannParticle) = d.dir
+temperature(d::MaxwellBoltzmannParticle) = d.temperature
+
+"""
+
+    _weight(
+        d::MaxwellBoltzmannParticle{D,P,T}, ps::ParticleStateful{D,P}
+    ) where {D,P,T}
+
+Unsafe weight-function for [`MaxwellBoltzmannParticle`](@ref), which is given by
+
+```math
+
+w(p) = \\begin{cases}
+\\frac{1}{4\\pi}\\sqrt{\\frac{2}{\\pi}}\\frac{\\varrho^2}{a^3}\\exp\\left(\\frac{-\\varrho^2}{2a^2}\\right)\\quad &\\text{,for } p^2=m^2\\\\
+0 &\\mathrm{elsewhere}.
+\\end{cases}
+
+```
+
+with ``\\varrho^2 = p_x^2 + p_y^2 + p_z^2`` and ``a = \\sqrt{m k_B T}`` (``m`` is the particle's mass, ``k_B`` is the Boltzmann constant,
+and ``T`` is the temperature).
+"""
+function _weight(
+    d::MaxwellBoltzmannParticle{D,P,T}, ps::ParticleStateful{D,P}
+) where {D,P,T}
+    mom = momentum(ps)
+
+    mag = getMag(mom)
+    E = getE(mom)
+    m = mass(_particle(d))
+
+    if abs(E^2 - m^2 - mag^2) <= 1e-5
+        return Distributions.pdf(d.rho_dist, mag) / (4 * pi)
+    else
+        return zero(T)
+    end
+end
+
+function max_weight(d::MaxwellBoltzmannParticle)
+    return Distributions.pdf(d.rho_dist, sqrt(2) * d.rho_dist.a) / (4 * pi)
+end
+
+function QEDevents._randmom(rng::AbstractRNG, d::MaxwellBoltzmannParticle)
+    rho = rand(rng, d.rho_dist)
+    cth = rand(rng) * 2 - 1
+    sth = sqrt(1 - cth^2)
+    phi = rand(rng) * 2 * pi
+    sphi, cphi = sincos(phi)
+
+    E = sqrt(rho^2 + mass(d.part)^2)
+    px = rho * sth * cphi
+    py = rho * sth * sphi
+    pz = rho * cth
+    return SFourMomentum(E, px, py, pz)
+end
+
+# consider writing this function for all single particle dists generically,
+# and implement _rand! accordingly. This could increase speed if a more
+# performant implementation for several samples exists.
+function _randmom(rng::AbstractRNG, d::MaxwellBoltzmannParticle, n::Dims) end

test/runtests.jl

@@ -12,4 +12,8 @@ begin
     @time @safetestset "scattering process distribution" begin
         include("interfaces/process_distribution.jl")
     end
+
+    @time @safetestset "Maxwell Boltzmann" begin
+        include("sampler/single_particle_dists/maxwell_boltzmann.jl")
+    end
 end

test/sampler/single_particle_dists/maxwell_boltzmann.jl

@@ -0,0 +1,69 @@
+using QEDbase
+using QEDcore
+using QEDevents
+using Random: Random
+import Distributions: Normal, Gamma
+
+RNG = Random.MersenneTwister(137137137)
+
+include("../../test_implementation/TestImpl.jl")
+include("../testutils.jl")
+
+test_particle = rand(RNG, TestImpl.PARTICLE_SET)
+test_direction = rand(RNG, (Incoming(), Outgoing(), UnknownDirection()))
+
+const N_SAMPLES = 1_000_000 # samples to be tested
+const TEMPERATURES = (1e-6, 1e-3, 1.0, 1e3, 1e6)
+
+@testset "$temp" for temp in TEMPERATURES
+    test_dist = MaxwellBoltzmannParticle(test_direction, test_particle, temp)
+
+    @testset "dist properties" begin
+        @test temperature(test_dist) == temp
+        @test QEDevents._particle(test_dist) == test_particle
+        @test QEDevents._particle_direction(test_dist) == test_direction
+    end
+
+    @testset "sample properties" begin
+        test_sample = rand(RNG, test_dist)
+        @test QEDcore.particle_species(test_sample) == test_particle
+        @test QEDcore.particle_direction(test_sample) == test_direction
+    end
+
+    @testset "sample distribution" begin
+        # maybe this comes in handy: https://github.com/JuliaStats/Distributions.jl/blob/47c040beef8c61bad3e1eefa4fc8194e3a62b55a/test/testutils.jl#L188C10-L188C22
+        test_sample = rand(RNG, test_dist, N_SAMPLES)
+
+        @testset "magnitude" begin
+            a = sqrt(mass(test_particle) * temp)
+            mag_projection = TestProjection(getMag, MaxwellBoltzmann(a))
+            @test test_univariate_samples(mag_projection, test_sample)
+        end
+
+        @testset "x component" begin
+            a = sqrt(mass(test_particle) * temp)
+            x_projection = TestProjection(getX, Normal(0.0, a))
+            @test test_univariate_samples(x_projection, test_sample)
+        end
+
+        @testset "y component" begin
+            a = sqrt(mass(test_particle) * temp)
+            y_projection = TestProjection(getY, Normal(0.0, a))
+            @test test_univariate_samples(y_projection, test_sample)
+        end
+
+        @testset "z component" begin
+            a = sqrt(mass(test_particle) * temp)
+            z_projection = TestProjection(getZ, Normal(0.0, a))
+            @test test_univariate_samples(z_projection, test_sample)
+        end
+    end
+
+    @testset "weight properties" begin
+        # TODO:
+        # * write a utils function for that (for general dists?)
+        test_samples = rand(RNG, test_dist, N_SAMPLES)
+        test_weights = weight.(test_dist, test_samples)
+        @test all(test_weights .<= max_weight(test_dist))
+    end
+end

test/sampler/testutils.jl

@@ -0,0 +1,41 @@
+using Distributions: Distributions
+import Distributions: pdf, quantile, Normal
+import StatsBase: fit, Histogram, middle
+using LinearAlgebra
+
+struct TestProjection{F,D<:Distributions.UnivariateDistribution}
+    proj::F
+    target_dist::D
+end
+
+# https://github.com/JuliaStats/Distributions.jl/blob/47c040beef8c61bad3e1eefa4fc8194e3a62b55a/test/testutils.jl#L188C10-L188C22
+"""
+    test_univariate_samples(p::TestProjection,samples)
+
+Tests if the given projection of the given samples follow the target_dist.
+"""
+function test_univariate_samples(
+    p::TestProjection, samples::AbstractVector{<:ParticleStateful}; nbins=50, q=1e-4
+)
+    moms = momentum.(samples)
+    samples_proj = p.proj.(moms)
+
+    h = normalize(fit(Histogram, samples_proj; nbins=nbins, closed=:right); mode=:pdf)
+    ww = h.weights
+    w = map(Base.Fix1(pdf, p.target_dist), h.edges[1])
+    n = length(h.edges[1]) - 1
+    max_w = maximum(w)
+
+    for i in 1:n
+        m = middle(w[i + 1], w[i])
+        bp = Normal(m, max_w * sqrt(q))
+        clb = max(quantile(bp, q / 2), 0.0)
+        cub = quantile(bp, 1 - q / 2)
+        @assert cub >= clb
+        if !(clb <= ww[i]) || !(ww[i] <= cub)
+            return false
+        end
+    end
+
+    return true
+end