add support for inplace integrand

Project.toml

@@ -1,7 +1,7 @@
 name = "MCIntegration"
 uuid = "ea1e2de9-7db7-4b42-91ee-0cd1bf6df167"
 authors = ["Kun Chen", "Xiansheng Cai", "Pengcheng Hou"]
-version = "0.3.0"
+version = "0.3.1"
 
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"

README.md

@@ -81,6 +81,14 @@ julia> integrate(var = Continuous(0.0, 1.0), dof = [[2,], [3,]]) do X, c
        end
 ```
 
+If there are too many components of integrands, it is better to preallocate the integrand weights. The function `integrate` provide an `inplace` key argument to achieve this goal. It is turned off by default, and only applies to the solver `:vegas` and `:vegasmc`. Once `inplace` is turned on, `integrate` will call the user-defined integrand function with a preallocated vector to store the user calculated weights. The following example demonstrates its usage,  
+```julia
+julia> integrate(var = Continuous(0.0, 1.0), dof = [[2,], [3,]], inplace=true) do X, f, c
+           f[1] = (X[1]^2 + X[2]^2 < 1.0) ? 1.0 : 0.0
+           f[2] = (X[1]^2 + X[2]^2 + X[3]^2 < 1.0) ? 1.0 : 0.0
+       end
+```
+
 ## Measure Histogram
 You may want to study how an integral changes with a tuning parameter. The following example is how to solve the histogram measurement problem.
 ```julia

docs/src/index.md

@@ -84,6 +84,14 @@ julia> integrate(var = Continuous(0.0, 1.0), dof = [[2,], [3,]]) do X, c
        end
 ```
 
+If there are too many components of integrands, it is better to preallocate the integrand weights. The function `integrate` provide an `inplace` key argument to achieve this goal. It is turned off by default, and only applies to the solver `:vegas` and `:vegasmc`. Once `inplace` is turned on, `integrate` will call the user-defined integrand function with a preallocated vector to store the user calculated weights. The following example demonstrates its usage,  
+```julia
+julia> integrate(var = Continuous(0.0, 1.0), dof = [[2,], [3,]], inplace=true) do X, f, c
+           f[1] = (X[1]^2 + X[2]^2 < 1.0) ? 1.0 : 0.0
+           f[2] = (X[1]^2 + X[2]^2 + X[3]^2 < 1.0) ? 1.0 : 0.0
+       end
+```
+
 ## Measure Histogram
 You may want to study how an integral changes with a tuning parameter. The following example is how to solve the histogram measurement problem.
 ```julia

src/distribution/variable.jl

@@ -115,7 +115,7 @@ J(y) = J_i = N \\Delta x_i
 
 The grid point ``x_i`` is trained after each iteration.
 """
-function Continuous(lower::Float64, upper::Float64, size=MaxOrder; offset=0, alpha=2.0, adapt=true, ninc=1000, grid=collect(LinRange(lower, upper, ninc))) where {G}
+function Continuous(lower::Float64, upper::Float64, size=MaxOrder; offset=0, alpha=2.0, adapt=true, ninc=1000, grid=collect(LinRange(lower, upper, ninc)))
     @assert offset + 1 < size
     size = size + 1 # need one more element as cache for the swap operation
     @assert upper > lower + 2 * eps(1.0)

src/main.jl

@@ -24,7 +24,9 @@
 
  # Arguments
 
-- `integrand`:Function call to evaluate the integrand. It should accept an argument of the type [`Configuration`](@ref), and return a weight. 
+- `integrand`:Function call to evaluate the integrand.  
+              If inpace = false, then the signature of the integrand should be `integrand(var, config)`, where `var` is a vector of random variables, and `config` is the [`Configuration`](@ref) struct. It should return one or more weights, corresponding to the value of each component of the integrand for the given `var`.
+              If inpace = true, then the signature of the integrand should be `integrand(var, weights, config)`, where the additional argument `weights` are the value of the components of the integrand for the given `var`.
               Internally, MC only samples the absolute value of the weight. Therefore, it is also important to define Main.abs for the weight if its type is user-defined. 
 - `solver` :  :vegas, :vegasmc, or :mcmc. See Readme for more details.
 - `config`:   [`Configuration`](@ref) object to perform the MC integration. If `nothing`, it attempts to create a new one with Configuration(; kwargs...).
@@ -40,14 +42,18 @@
 - `ignore`:   ignore the iteration until the `ignore` round. By default, the first iteration is igonred if adapt=true, and non is ignored if adapt=false.
 - `measure`:  measurement function, See [`Vegas.montecarlo`](@ref), [`VegasMC.montecarlo`](@ref) and [`MCMC.montecarlo`](@ref) for more details.
 - `measurefreq`: how often perform the measurement for ever `measurefreq` MC steps. If a measurement is expansive, you may want to make the measurement less frequent.
+- `inplace`:  whether to use the inplace version of the integrand. Default is `false`, which is more convenient for integrand with a few return values but may cause type instability. Only useful for the :vegas and :vegasmc solver.
 - `kwargs`:   Keyword arguments. The supported keywords include,
   * `measure` and `measurefreq`: measurement function and how frequent it is called. 
   * If `config` is `nothing`, you may need to provide arguments for the `Configuration` constructor, check [`Configuration`](@ref) docs for more details.
 
 # Examples
 ```julia-repl
-julia> integrate((x, c)->(x[1]^2+x[2]^2); var = Continuous(0.0, 1.0), dof = 2, print=-1)
-Integral 1 = 0.6668625385256122 ± 0.0009960142738129768   (chi2/dof = 1.07)
+integrate((x, c)->(x[1]^2+x[2]^2); var = Continuous(0.0, 1.0), dof = 2, print=-2, solver=:vegas)
+Integral 1 = 0.6663652080622751 ± 0.000490978424216832   (chi2/dof = 0.645)
+
+julia> integrate((x, f, c)-> (f[1] = x[1]^2+x[2]^2); var = Continuous(0.0, 1.0), dof = 2, print=-2, solver=:vegas, inplace=true)
+Integral 1 = 0.6672083165915914 ± 0.0004919147870306026   (chi2/dof = 2.54)
 ```
 """
 function integrate(integrand::Function;
@@ -64,6 +70,7 @@ function integrate(integrand::Function;
     ignore::Int=adapt ? 1 : 0, #ignore the first `ignore` iteractions in average
     measure::Union{Nothing,Function}=nothing,
     measurefreq::Int=1,
+    inplace::Bool=false, # whether to use the inplace version of the integrand
     kwargs...
 )
     if isnothing(config)
@@ -119,10 +126,10 @@ function integrate(integrand::Function;
 
             if solver == :vegasmc
                 config = VegasMC.montecarlo(config, integrand, nevalperblock, print, save, timer, debug;
-                    measure=measure, measurefreq=measurefreq)
+                    measure=measure, measurefreq=measurefreq, inplace=inplace)
             elseif solver == :vegas
                 config = Vegas.montecarlo(config, integrand, nevalperblock, print, save, timer, debug;
-                    measure=measure, measurefreq=measurefreq)
+                    measure=measure, measurefreq=measurefreq, inplace=inplace)
             elseif solver == :mcmc
                 config = MCMC.montecarlo(config, integrand, nevalperblock, print, save, timer, debug;
                     measure=measure, measurefreq=measurefreq)

src/vegas/montecarlo.jl

@@ -79,7 +79,7 @@ Integral 1 = 0.667203631824444 ± 0.0005046485925614018   (chi2/dof = 1.46)
 """
 function montecarlo(config::Configuration{Ni,V,P,O,T}, integrand::Function, neval,
     print=0, save=0, timer=[], debug=false;
-    measure::Union{Nothing,Function}=nothing, measurefreq::Int=1
+    measure::Union{Nothing,Function}=nothing, measurefreq::Int=1, inplace::Bool=false
 ) where {Ni,V,P,O,T}
 
     @assert measurefreq > 0
@@ -89,10 +89,18 @@ function montecarlo(config::Configuration{Ni,V,P,O,T}, integrand::Function, neva
 
     ################## test integrand type stability ######################
     if debug
-        if (length(config.var) == 1)
-            MCUtility.test_type_stability(integrand, (config.var[1], config))
+        if inplace
+            if (length(config.var) == 1)
+                MCUtility.test_type_stability(integrand, (config.var[1], weights, config))
+            else
+                MCUtility.test_type_stability(integrand, (config.var, weights, config))
+            end
         else
-            MCUtility.test_type_stability(integrand, (config.var, config))
+            if (length(config.var) == 1)
+                MCUtility.test_type_stability(integrand, (config.var[1], config))
+            else
+                MCUtility.test_type_stability(integrand, (config.var, config))
+            end
         end
     end
     #######################################################################
@@ -129,7 +137,11 @@ function montecarlo(config::Configuration{Ni,V,P,O,T}, integrand::Function, neva
         end
         # weights = @_expanded_integrand(config, integrand, 1) # very fast, but requires explicit N
         # weights = integrand_wrap(config, integrand) #make a lot of allocations
-        weights = (fieldcount(V) == 1) ? integrand(config.var[1], config) : integrand(config.var, config)
+        if inplace
+            (fieldcount(V) == 1) ? integrand(config.var[1], weights, config) : integrand(config.var, weights, config)
+        else
+            weights = (fieldcount(V) == 1) ? integrand(config.var[1], config) : integrand(config.var, config)
+        end
 
         if (ne % measurefreq == 0)
             if isnothing(measure)

src/vegas_mc/montecarlo.jl

@@ -62,25 +62,35 @@ Integral 1 = 0.6640840471808533 ± 0.000916060916265263   (chi2/dof = 0.945)
 """
 function montecarlo(config::Configuration{N,V,P,O,T}, integrand::Function, neval,
     print=0, save=0, timer=[], debug=false;
-    measure::Union{Nothing,Function}=nothing, measurefreq::Int=1
+    measure::Union{Nothing,Function}=nothing, measurefreq::Int=1, inplace::Bool=false
 ) where {N,V,P,O,T}
 
     @assert measurefreq > 0
 
+    if isnothing(measure)
+        @assert (config.observable isa AbstractVector) && (length(config.observable) == config.N) && (eltype(config.observable) == T) "the default measure can only handle observable as Vector{$T} with $(config.N) elements!"
+    end
+    weights = zeros(T, N)
+    _weights = zeros(T, N)
+
     ################## test integrand type stability ######################
     if debug
-        if (length(config.var) == 1)
-            MCUtility.test_type_stability(integrand, (config.var[1], config))
+        if inplace
+            if (length(config.var) == 1)
+                MCUtility.test_type_stability(integrand, (config.var[1], weights, config))
+            else
+                MCUtility.test_type_stability(integrand, (config.var, weights, config))
+            end
         else
-            MCUtility.test_type_stability(integrand, (config.var, config))
+            if (length(config.var) == 1)
+                MCUtility.test_type_stability(integrand, (config.var[1], config))
+            else
+                MCUtility.test_type_stability(integrand, (config.var, config))
+            end
         end
     end
     #######################################################################
 
-    if isnothing(measure)
-        @assert (config.observable isa AbstractVector) && (length(config.observable) == config.N) && (eltype(config.observable) == T) "the default measure can only handle observable as Vector{$T} with $(config.N) elements!"
-    end
-    weights = zeros(T, N)
     relativeWeights = zeros(T, N)
     # padding probability for user and normalization integrands
     # should be something like [f_1(x)*g_0(y), f_2(x, y), 1*f_0(x)*g_0(y)], where f_0(x) and g_0(y) are the ansatz from the Vegas map
@@ -93,7 +103,11 @@ function montecarlo(config::Configuration{N,V,P,O,T}, integrand::Function, neval
         Dist.initialize!(var, config)
     end
 
-    _weights = (length(config.var) == 1) ? integrand(config.var[1], config) : integrand(config.var, config)
+    if inplace
+        (length(config.var) == 1) ? integrand(config.var[1], _weights, config) : integrand(config.var, _weights, config)
+    else
+        _weights = (length(config.var) == 1) ? integrand(config.var[1], config) : integrand(config.var, config)
+    end
 
     padding_probability .= [Dist.padding_probability(config, i) for i in 1:N+1]
     probability = config.reweight[config.norm] * padding_probability[config.norm] #normalization integral
@@ -110,8 +124,8 @@ function montecarlo(config::Configuration{N,V,P,O,T}, integrand::Function, neval
     # end
     if debug
         for _update in updates
-            MCUtility.test_type_stability(_update, (config, integrand, probability,
-                weights, padding_probability, padding_probability_cache))
+            MCUtility.test_type_stability(_update, (config, integrand, inplace, probability,
+                weights, _weights, padding_probability, padding_probability_cache))
         end
     end
 
@@ -121,8 +135,8 @@ function montecarlo(config::Configuration{N,V,P,O,T}, integrand::Function, neval
     for ne = 1:neval
         # config.neval += 1
         _update = rand(config.rng, updates) # randomly select an update
-        probability = _update(config, integrand, probability,
-            weights, padding_probability, padding_probability_cache)
+        probability = _update(config, integrand, inplace, probability,
+            weights, _weights, padding_probability, padding_probability_cache)
         if debug && (isfinite(probability) == false)
             @warn("integrand probability = $(probability) is not finite at step $(neval)")
         end

src/vegas_mc/updates.jl

@@ -42,8 +42,8 @@
 #     return
 # end
 
-function changeVariable(config::Configuration{N,V,P,O,T}, integrand,
-    currProbability::Float64, weights,
+function changeVariable(config::Configuration{N,V,P,O,T}, integrand, inplace,
+    currProbability::Float64, weights, _weights,
     padding_probability, _padding_probability) where {N,V,P,O,T}
     # update to change the variables of the current diagrams
     maxdof = config.maxdof
@@ -64,10 +64,15 @@ function changeVariable(config::Configuration{N,V,P,O,T}, integrand,
         return currProbability
     end
 
-    # weights = integrand_wrap(config, integrand)
-    _weights = (fieldcount(V) == 1) ?
-               integrand(config.var[1], config) :
-               integrand(config.var, config)
+    if inplace
+        (fieldcount(V) == 1) ?
+        integrand(config.var[1], _weights, config) :
+        integrand(config.var, config)
+    else
+        _weights = (fieldcount(V) == 1) ?
+                   integrand(config.var[1], config) :
+                   integrand(config.var, config)
+    end
     # evaulation acutally happens before this step
     config.neval += 1
 

test/montecarlo.jl

@@ -114,6 +114,19 @@ function TestComplex2(totalstep, alg)
     return res
 end
 
+function TestComplex2_inplace(totalstep, alg)
+    function integrand(x, f, c) #return a tuple (real, complex) 
+        #the code should handle real -> complex conversion
+        f[1] = x[1]
+        f[2] = x[1]^2 * 1im
+    end
+    res = integrate(integrand; dof=[[1,], [1,]], neval=totalstep, print=-1, type=ComplexF64, solver=alg, inplace=true, debug=true)
+    config = res.config
+    w = zeros(ComplexF64, 2)
+    @inferred integrand(config.var[1], w, config) #make sure the type is inferred for the integrand function
+    return res
+end
+
 @testset "MCMC Sampler" begin
     neval = 1000_00
     println("MCMC tests")
@@ -173,6 +186,8 @@ end
     println("Complex2")
     check_complex(TestComplex2(neval, :vegas), [0.5, 1.0 / 3 * 1im])
 
+    println("inplace Complex2")
+    check_complex(TestComplex2_inplace(neval, :vegas), [0.5, 1.0 / 3 * 1im])
 end
 
 @testset "Markov-Chain Vegas" begin
@@ -208,4 +223,6 @@ end
     println("Complex2")
     check_complex(TestComplex2(neval, :vegasmc), [0.5, 1.0 / 3 * 1im])
 
+    println("inplace Complex2")
+    check_complex(TestComplex2_inplace(neval, :vegasmc), [0.5, 1.0 / 3 * 1im])
 end