Merge pull request #2 from climate-machine/sb/solus

rename to Solus.jl, add NEKI

Author: simonbyrne

Manifest.toml

@@ -16,10 +16,10 @@ uuid = "9e28174c-4ba2-5203-b857-d8d62c4213ee"
 version = "0.8.10"
 
 [[BinaryProvider]]
-deps = ["Libdl", "Pkg", "SHA", "Test"]
-git-tree-sha1 = "055eb2690182ebc31087859c3dd8598371d3ef9e"
+deps = ["Libdl", "SHA"]
+git-tree-sha1 = "c7361ce8a2129f20b0e05a89f7070820cfed6648"
 uuid = "b99e7846-7c00-51b0-8f62-c81ae34c0232"
-version = "0.5.3"
+version = "0.5.4"
 
 [[Compat]]
 deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
@@ -51,6 +51,12 @@ git-tree-sha1 = "dec0ebacfbc3a2126c614ab5e903c9ef063688d0"
 uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
 version = "0.17.0"
 
+[[DocStringExtensions]]
+deps = ["LibGit2", "Markdown", "Pkg", "Test"]
+git-tree-sha1 = "4d30e889c9f106a51ffa4791a88ffd4765bf20c3"
+uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+version = "0.7.0"
+
 [[InteractiveUtils]]
 deps = ["Markdown"]
 uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
@@ -73,25 +79,25 @@ deps = ["Base64"]
 uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
 
 [[Missings]]
-deps = ["Dates", "InteractiveUtils", "SparseArrays", "Test"]
-git-tree-sha1 = "d1d2585677f2bd93a97cfeb8faa7a0de0f982042"
+deps = ["SparseArrays", "Test"]
+git-tree-sha1 = "f0719736664b4358aa9ec173077d4285775f8007"
 uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
-version = "0.4.0"
+version = "0.4.1"
 
 [[Mmap]]
 uuid = "a63ad114-7e13-5084-954f-fe012c677804"
 
 [[OrderedCollections]]
 deps = ["Random", "Serialization", "Test"]
-git-tree-sha1 = "85619a3f3e17bb4761fe1b1fd47f0e979f964d5b"
+git-tree-sha1 = "c4c13474d23c60d20a67b217f1d7f22a40edf8f1"
 uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
-version = "1.0.2"
+version = "1.1.0"
 
 [[PDMats]]
 deps = ["Arpack", "LinearAlgebra", "SparseArrays", "SuiteSparse", "Test"]
-git-tree-sha1 = "b6c91fc0ab970c0563cbbe69af18d741a49ce551"
+git-tree-sha1 = "8b68513175b2dc4023a564cb0e917ce90e74fd69"
 uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
-version = "0.9.6"
+version = "0.9.7"
 
 [[Pkg]]
 deps = ["Dates", "LibGit2", "Markdown", "Printf", "REPL", "Random", "SHA", "UUIDs"]
@@ -155,10 +161,10 @@ deps = ["LinearAlgebra", "SparseArrays"]
 uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [[StatsBase]]
-deps = ["DataStructures", "DelimitedFiles", "LinearAlgebra", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "Test"]
-git-tree-sha1 = "435707791dc85a67d98d671c1c3fcf1b20b00f94"
+deps = ["DataStructures", "LinearAlgebra", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics"]
+git-tree-sha1 = "8a0f4b09c7426478ab677245ab2b0b68552143c7"
 uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
-version = "0.29.0"
+version = "0.30.0"
 
 [[StatsFuns]]
 deps = ["Rmath", "SpecialFunctions", "Test"]

Project.toml

@@ -1,9 +1,11 @@
-name = "EnsembleKalmanInversion"
+name = "Solus"
 uuid = "16b9c28e-565c-11e9-13a9-53e2645fa95d"
 authors = ["Simon Byrne <[email protected]>"]
 version = "0.1.0"
 
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

README.md

@@ -1,4 +1,4 @@
-# EnsembleKalmanInversion.jl
+# Solus.jl
 
 [![Travis Build Status](https://travis-ci.org/climate-machine/EnsembleKalmanInversion.jl.svg?branch=master)](https://travis-ci.org/climate-machine/EnsembleKalmanInversion.jl)
 [![Appveyor Build Status](https://ci.appveyor.com/api/projects/status/s7l958ngcd3efv9t/branch/master?svg=true)](https://ci.appveyor.com/project/simonbyrne/ensemblekalmaninversion-jl/branch/master)

src/EnsembleKalmanInversion.jl

@@ -0,0 +1,99 @@
+module Solus
+
+using Distributions, Statistics, LinearAlgebra, DocStringExtensions
+
+include("spaces.jl")
+
+"""
+    SolusProblem
+
+An uncertainty quantification problem.
+
+# Fields
+
+$(DocStringExtensions.FIELDS)
+"""
+struct SolusProblem{P,M,O,S<:HilbertSpace}
+    """
+    The prior distribution for the parameters ``θ``. Currently only `MvNormal` objects are supported.
+    """
+    prior::P
+
+    """
+    A forward model maps the parameter ``θ`` to a predicted observation.
+    """
+    forwardmodel::M
+    
+    """
+    The observed data.
+    """
+    obs::O
+
+    """
+    The space on which the forward model output and observation exist. `DefaultSpace()` is the default.
+    """
+    space::S
+end
+SolusProblem(prior, forwardmodel, obs) = SolusProblem(prior, forwardmodel, obs, DefaultSpace())
+
+struct FlatPrior
+end
+
+Distributions.logpdf(::FlatPrior, x) = 0.0
+
+function neglogposteriordensity(s::SolusProblem, θ)
+    -logpdf(s.prior, θ) + norm(s.forwardmodel(θ) - s.obs, s.space)
+end
+
+
+"""
+    Ensemble
+
+An ensemble of `inputs` and their corresponding `outputs` from the forward model.
+
+# Fields
+
+$(DocStringExtensions.FIELDS)
+"""
+struct Ensemble{I,O}
+    """
+    A matrix of inputs: each column corresponds to a vector of ``θ``s
+    """
+    inputs::Vector{I}
+    """
+    Result of forward model for each column of `inputs`.
+    """
+    outputs::Vector{O}
+end
+
+"""
+    neki_iter(prob::SolusProblem, ens::Ensemble)
+
+Peform an iteration of NEKI, returning a new `Ensemble` object.
+"""
+function neki_iter(prob::SolusProblem, ens::Ensemble)
+    covθ = cov(ens.inputs)
+    covθ += (tr(covθ)*1e-15)I
+    
+    m = mean(ens.outputs)
+
+    CG = [dot(Gθk - m, Gθj - prob.obs, prob.space) for Gθj in ens.outputs, Gθk in ens.outputs]
+    
+    Δt = 0.1 / norm(CG) # use better constants  
+    implicit = lu( I + Δt .* (covθ / cov(prob.prior)) ) # todo: incorporate means
+    noise = MvNormal(covθ)
+    
+    inputs = map(enumerate(ens.inputs)) do (j, θj)
+        X = sum(enumerate(ens.inputs)) do (k, θk)
+            CG[k,j]*θk
+        end
+        rhs = θj .- Δt .* X
+        (implicit \ rhs) .+ sqrt(Δt)*rand(noise)
+    end
+    
+    outputs = map(prob.forwardmodel, inputs)
+    Ensemble(inputs, outputs)
+end
+    
+
+end # module

src/Solus.jl

@@ -0,0 +1,30 @@
+import LinearAlgebra: norm, dot
+
+"""
+    HilbertSpace
+
+A `Space <: HilbertSpace` defines methods for
+- `norm(x, ::Space)`
+- `dot(x, y, ::Space)`
+"""
+abstract type HilbertSpace
+end
+
+"""
+    DefaultSpace
+
+Uses the default `norm` and `dot` implementations (e.g. `Array`s and `Array`s of `Array`s).
+"""
+struct DefaultSpace <: HilbertSpace
+end
+norm(x, ::DefaultSpace) = norm(x)
+dot(x,y, ::DefaultSpace) = dot(x,y)
+
+
+struct CovarianceSpace{M} <: HilbertSpace
+    Γ::M
+end
+norm(x, c::CovarianceSpace) = x'*(c.Γ\x)
+dot(x,y, c::CovarianceSpace) = x'*(c.Γ\y)
+
+# TODO: CovarianceSpace/PDSpace with a positive definite matrix

src/spaces.jl

@@ -1,35 +1,2 @@
 using Test
-using EnsembleKalmanInversion, Distributions
-
-let
-    iter_max = 5
-    J = 100
-    n = 5
-    r = 0.1
-    # r = 0.0
-
-    A = rand(Normal(0, 2), (n,n))
-    u_t = rand(Normal(0, 3), n)
-    g_t = A * u_t
-
-    cov = rand(Normal(0, 2), (n,n))
-    cov = cov' * cov
-    cov *= r^2
-
-    u_init = rand(Normal(0, 2), (n,J))
-    g_ens = A * u_init
-
-    eki = EKI(g_t, cov, u_init)
-
-    iter = 0
-    while iter < iter_max
-        g_ens = A * eki.u[end]
-        EnsembleKalmanInversion.update!(eki, g_ens)
-        if EnsembleKalmanInversion.residual(eki) < 0.01
-            break
-        end
-        iter += 1
-    end
-
-    @test u_t ≈ vec(EnsembleKalmanInversion.get_u(eki)) atol=0.01
-end
+using Solus