From 46ac9063679d4d2ff8664a8dc390d05ec704eb35 Mon Sep 17 00:00:00 2001
From: monty <velimir.vesselinov@gmail.com>
Date: Fri, 21 Jul 2023 12:43:27 -0600
Subject: [PATCH] save jld2

---
 Project.toml               |  4 +++-
 src/AffineInvariantMCMC.jl | 40 ++++++++++----------------------------
 2 files changed, 13 insertions(+), 31 deletions(-)

diff --git a/Project.toml b/Project.toml
index 8e03c0a..4b0288a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,14 +1,16 @@
 name = "AffineInvariantMCMC"
 uuid = "a0f608ac-68d0-5108-ad9a-0cdca73324f3"
-version = "1.2.0"
+version = "1.2.1"
 
 [deps]
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RobustPmap = "27aeedcb-f738-516b-a0b8-3211cf1146e5"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
+JLD2 = "0.4, 0.5, 1"
 RobustPmap = "1"
 Statistics = "1"
 julia = "1.9"
diff --git a/src/AffineInvariantMCMC.jl b/src/AffineInvariantMCMC.jl
index 322024c..6554dc6 100644
--- a/src/AffineInvariantMCMC.jl
+++ b/src/AffineInvariantMCMC.jl
@@ -1,33 +1,8 @@
 __precompile__()
 
 """
-MADS: Model Analysis & Decision Support in Julia (Mads.jl v1.0) 2016
-module AffineInvariantMCMC
-
-https://mads.lanl.gov
-https://github.com/madsjulia
-
-Licensing: GPLv3: http://www.gnu.org/licenses/gpl-3.0.html
-
-Copyright 2016.  Los Alamos National Security, LLC.  All rights reserved.
-
-This material was produced under U.S. Government contract DE-AC52-06NA25396 for
-Los Alamos National Laboratory, which is operated by Los Alamos National Security, LLC for
-the U.S. Department of Energy. The Government is granted for itself and others acting on its
-behalf a paid-up, nonexclusive, irrevocable worldwide license in this material to reproduce,
-prepare derivative works, and perform publicly and display publicly. Beginning five (5) years after
---------------- November 17, 2015, ----------------------------------------------------------------
-subject to additional five-year worldwide renewals, the Government is granted for itself and
-others acting on its behalf a paid-up, nonexclusive, irrevocable worldwide license in this
-material to reproduce, prepare derivative works, distribute copies to the public, perform
-publicly and display publicly, and to permit others to do so.
-
-NEITHER THE UNITED STATES NOR THE UNITED STATES DEPARTMENT OF ENERGY, NOR LOS ALAMOS NATIONAL SECURITY, LLC,
-NOR ANY OF THEIR EMPLOYEES, MAKES ANY WARRANTY, EXPRESS OR IMPLIED, OR ASSUMES ANY LEGAL LIABILITY OR
-RESPONSIBILITY FOR THE ACCURACY, COMPLETENESS, OR USEFULNESS OF ANY INFORMATION, APPARATUS, PRODUCT, OR
-PROCESS DISCLOSED, OR REPRESENTS THAT ITS USE WOULD NOT INFRINGE PRIVATELY OWNED RIGHTS.
-
-LA-CC-15-080; Copyright Number Assigned: C16008
+MADS: Model Analysis & Decision Support (Mads.jl v1.0) (c) 2023
+Module: AffineInvariantMCMC
 """
 module AffineInvariantMCMC
 
@@ -65,8 +40,10 @@ Reference:
 
 Goodman & Weare, "Ensemble samplers with affine invariance", Communications in Applied Mathematics and Computational Science, DOI: 10.2140/camcos.2010.5.65, 2010.
 """
-function sample(llhood::Function, numwalkers::Int, x0::Array, numsamples_perwalker::Integer, thinning::Integer, a::Number=2.; rng::Random.AbstractRNG=Random.GLOBAL_RNG)
-	@assert length(size(x0)) == 2
+function sample(llhood::Function, numwalkers::Integer, x0::AbstractMatrix, numsamples_perwalker::Integer, thinning::Integer, a::Number=2.; filename::AbstractString, load::Bool=true, save::Bool=true, rng::Random.AbstractRNG=Random.GLOBAL_RNG)
+	if isfile(filename) && load
+		chain, llhoodvals = JLD2.load(filename, "chain", "llhoods")
+	end
 	x = copy(x0)
 	chain = Array{Float64}(undef, size(x0, 1), numwalkers, div(numsamples_perwalker, thinning))
 	lastllhoodvals = RobustPmap.rpmap(llhood, map(i->x[:, i], 1:size(x, 2)))
@@ -98,11 +75,14 @@ function sample(llhood::Function, numwalkers::Int, x0::Array, numsamples_perwalk
 			end
 		end
 	end
+	if save &&  filename != ""
+		JLD2.save(filename, "chain", chain, "llhoods", llhoodvals)
+	end
 	return chain, llhoodvals
 end
 
 "Flatten MCMC arrays"
-function flattenmcmcarray(chain::Array, llhoodvals::Array)
+function flattenmcmcarray(chain::AbstractArray, llhoodvals::AbstractArray)
 	numdims, numwalkers, numsteps = size(chain)
 	newchain = Array{Float64}(undef, numdims, numwalkers * numsteps)
 	for j = 1:numsteps