Merge pull request #26 from ggebbie/23-source-water-inversion-with-mu…

…ltiple-state-variables-bh

23 source water inversion with multiple state variables bh

src/BLUEs.jl

@@ -2,6 +2,7 @@ module BLUEs
 
 using LinearAlgebra, Statistics, Unitful, UnitfulLinearAlgebra, Measurements
 using DimensionalData
+using Revise
 
 export Estimate, DimEstimate, OverdeterminedProblem, UnderdeterminedProblem
 export solve, show, cost, datacost, controlcost
@@ -101,6 +102,19 @@ function show(io::IO, mime::MIME{Symbol("text/plain")}, x::Union{DimEstimate,Est
     # show(io, mime, F.V)
 end
 
+function show(io::IO, mime::MIME{Symbol("text/plain")}, x::DimArray{Quantity{Float64}, 3})
+    summary(io, x); println(io)
+    statevars = x.dims[3]
+    for (i, s) in enumerate(statevars)
+        if i != 1
+            println()
+        end
+
+        println(io, "State Variable " * string(i) * ": " * string(s))
+        show(io, mime, x[:,:,At(s)])
+    end
+end
+
 """
     function getproperty(x::Estimate, d::Symbol)
 
@@ -473,12 +487,12 @@ function impulseresponse(funk::Function,x,args...)
     y₀ = funk(x,args...)
     #Eunits = expectedunits(y₀,x)
     #Eu = Quantity.(zeros(length(y₀),length(x)),Eunits)
-
-    Ep = zeros(length(y₀),length(x))
+    Ep = zeros(length(y₀), length(x))
+
     for rr in eachindex(x)
         #u = zeros(length(x)).*unit.(x)[:]
         if length(x) > 1
-            u = Quantity.(zeros(length(x)),unit.(vec(x)))
+            u = Quantity.(zeros(length(x)), unit.(vec(x)))
         else
             u = Quantity(0.0,unit(x))
         end
@@ -489,7 +503,9 @@ function impulseresponse(funk::Function,x,args...)
         y = funk(x₁,args...)
         #println(y)
         if y isa AbstractDimArray
-            Ep[:,rr] .= vec(parent((y - y₀)/Δu))
+            #Ep[:,rr] .= vec(parent((y - y₀)/Δu))
+            Ep[:, rr] .= ustrip.(vec(parent((y - y₀)/Δu)))
+
         else
             tmp = ustrip.((y - y₀)/Δu)
             # getting ugly around here
@@ -500,7 +516,8 @@ function impulseresponse(funk::Function,x,args...)
             end
         end
     end
-
+
+
     # This function could use vcat to be cleaner (but maybe slower)
     # Note: Need to add ability to return sparse matrix 
     #return E = UnitfulMatrix(Ep,unit.(vec(y₀)),unit.(vec(x)))
@@ -559,10 +576,44 @@ function convolve(x::AbstractDimArray,E::AbstractDimArray)
     lags = first(dims(E))
     return sum([E[ii,:] ⋅ x[Near(tnow-ll),:] for (ii,ll) in enumerate(lags)])
 end
+"""
+function convolve(x::AbstractDimArray, E::AbstractDimArray, coeffs::UnitfulMatrix}
+    the `coeffs` argument signifies that x is a 3D array (i.e. >1 state variables)
+
+    this function both convolves, and linearly combines the propagated state variables
+"""
+function convolve(x::AbstractDimArray,E::AbstractDimArray, coeffs::UnitfulMatrix)
+    statevars = x.dims[3]
+    mat = UnitfulMatrix(transpose([convolve(x[:,:,At(s)], E) for s in statevars])) * coeffs
+    return getindexqty(mat, 1,1)
+end
+
 function convolve(x::AbstractDimArray,E::AbstractDimArray,t::Number)
     lags = first(dims(E))
     return sum([E[ii,:] ⋅ x[Near(t-ll),:] for (ii,ll) in enumerate(lags)])
 end
+
+#coeffs signifies that x is 3D 
+function convolve(x::AbstractDimArray, E::AbstractDimArray, t::Number, coeffs::UnitfulMatrix)
+    statevars = x.dims[3]
+    mat = UnitfulMatrix(transpose([convolve(x[:,:,At(s)], E, t) for s in statevars]))*coeffs
+    return getindexqty(mat, 1,1) 
+end
+
+#don't handle the ndims(M) == 3 case here but I'll get back to it
+function convolve(x::AbstractDimArray, M::AbstractDimArray, Tx::Union{Ti, Vector}, coeffs::UnitfulMatrix)
+    if ndims(M) == 2
+        return DimArray([convolve(x,M,Tx[tt], coeffs) for (tt, yy) in enumerate(Tx)], Tx)
+    elseif ndims(M) == 3
+        error("some code should go here")
+    else
+        error("M has wrong number of dims") 
+    end
+
+
+end
+
+
 function convolve(x::AbstractDimArray,M::AbstractDimArray,Tx::Union{Ti,Vector})
     if ndims(M) == 2 
         return DimArray([convolve(x,M,Tx[tt]) for (tt,yy) in enumerate(Tx)],Tx)
@@ -626,5 +677,4 @@ function addcontrol!(x::AbstractDimArray,u)
     end
     return x
 end
-
 end # module

test/runtests.jl

@@ -260,14 +260,15 @@ include("test_functions.jl")
         y = predictobs(convolve,x,M)
 
         ## invert for y for x̃
-        # Given, M and y. Make first guess for x.
+        # Given, M and y. Make first guess for x.x
         # add adjustment
         # DimArray is good enough. This is an array, not necessarily a matrix.
         x₀ = DimArray(zeros(size(x))K,(Ti(years),last(dims(M))))
 
         # probe to get E matrix. Use function convolve
         E = impulseresponse(convolve,x₀,M)
-
+
+        @test (E*UnitfulMatrix(vec(x₀)))[1] .== ustrip(convolve(x₀, M))
         # Does E matrix work properly?
         ỹ = E*UnitfulMatrix(vec(x))
         @test isapprox(y,getindexqty(ỹ,1))
@@ -282,9 +283,9 @@ include("test_functions.jl")
         σₙ = 0.01
         σₓ = 100.0
 
-        Cnn = UnitfulMatrix(Diagonal(fill(σₙ,length(y))),fill(unit.(y).^1,length(y)),fill(unit.(y).^-1,length(y)),exact=false)
+        Cnn = UnitfulMatrix(Diagonal(fill(σₙ^2,length(y))),fill(unit.(y).^1,length(y)),fill(unit.(y).^-1,length(y)),exact=false)
 
-        Cxx = UnitfulMatrix(Diagonal(fill(σₓ,length(x₀))),unit.(x₀)[:],unit.(x₀)[:].^-1,exact=false)
+        Cxx = UnitfulMatrix(Diagonal(fill(σₓ^2,length(x₀))),unit.(x₀)[:],unit.(x₀)[:].^-1,exact=false)
 
         problem = UnderdeterminedProblem(UnitfulMatrix([y]),E,Cnn,Cxx,x₀)
         x̃ = solve(problem)
@@ -293,6 +294,108 @@ include("test_functions.jl")
         @test cost(x̃,problem) < 5e-2 # no noise in ob
     end
 
+    @testset "source water inversion: obs at one time, many surface regions, with circulation lag, TWO STATE VARIABLES" begin 
+
+        @dim YearCE "years Common Era"
+        @dim SurfaceRegion "surface location"
+        @dim InteriorLocation "interior location"
+        @dim StateVariable "state variable" 
+
+        yr = u"yr"
+        nτ = 5 # how much of a lag is possible?
+        lags = (0:(nτ-1))yr
+        surfaceregions = [:NATL,:ANT,:SUBANT]
+        years = (1990:2000)yr
+        n = length(surfaceregions)
+
+        M = source_water_matrix_with_lag(surfaceregions,lags,years)
+
+        statevariables = [:θ, :d18O] 
+        x = source_water_solution(surfaceregions,years, statevariables)
+
+        coeffs = UnitfulMatrix(reshape(rand(2), (2,1)), [unit(1.0), K], [K*permil^-1])
+        y = convolve(x, M, coeffs)
+
+        x₀ = copy(x).*0
+        E = impulseresponse(convolve,x₀,M,coeffs)
+
+        #E matrix tests from prior example
+        ỹ = E*UnitfulMatrix(vec(x))
+        @test isapprox(y,getindexqty(ỹ,1))
+
+        #had to get a little crafty to make this work... 
+        x̂ = E\UnitfulMatrix(reshape([ustrip.(y)], (1,1)), [permil], [unit(1.0)])
+        @test isapprox(y,getindexqty(E*x̂,1,1))
+
+        σₙ = 0.01
+        σₓ = 100.0
+
+        Cnn = UnitfulMatrix(Diagonal(fill(σₙ.^2,length(y))),fill(unit(y).^1,length(y)),fill(unit(y).^-1,length(y)),exact=false)
+
+        Cxx = UnitfulMatrix(Diagonal(fill(σₓ.^2,length(x₀))),unit.(x₀)[:],unit.(x₀)[:].^-1,exact=false)
+        problem = UnderdeterminedProblem(UnitfulMatrix([y]), E, Cnn, Cxx, x₀)
+        x̃ = solve(problem)
+        @test within(y, getindexqty(E*x̃.v, 1), 3σₙ)
+        @test cost(x̃,problem) < 5e-2
+        @test cost(x̃, problem) == datacost(x̃, problem) + controlcost(x̃, problem)
+    end
+
+    @testset "source water inversion: obs TIMESERIES, many surface regions, with circulation lag, TWO STATE VARIABLES" begin
+
+        @dim YearCE "years Common Era"
+        @dim SurfaceRegion "surface location"
+        @dim InteriorLocation "interior location"
+        @dim StateVariable "state variable" 
+
+        yr = u"yr"
+        nτ = 5 # how much of a lag is possible?
+        lags = (0:(nτ-1))yr
+        surfaceregions = [:NATL,:ANT,:SUBANT]
+        years = (1990:2000)yr
+        n = length(surfaceregions)
+
+        M = source_water_matrix_with_lag(surfaceregions,lags,years)
+
+        statevariables = [:θ, :d18O] 
+        x = source_water_solution(surfaceregions,years, statevariables)
+
+        Tx = first(dims(x))
+        x₀ = copy(x).*0
+
+        #choose random linear coefficients 
+        coeffs = UnitfulMatrix(reshape(rand(2), (2,1)), [unit(1.0), K], [K*permil^-1])
+
+        #synthetic timeseries 
+        y = convolve(x, M, Tx, coeffs)
+        E = impulseresponse(convolve, x₀, M, Tx,coeffs)
+
+        # Does E matrix work properly?
+        ỹ = E*UnitfulMatrix(vec(x))
+        for jj in eachindex(vec(y))
+            @test isapprox.(vec(y)[jj],getindexqty(ỹ,jj))
+        end
+        x̂ = E\UnitfulMatrix(vec(y))
+        for jj in eachindex(vec(y))
+            @test isapprox.(vec(y)[jj],getindexqty(E*x̂,jj))
+        end
+
+        #generate Cnn and Cxx matrices
+        σₙ = 0.01
+        σₓ = 100.0
+        Cnn = UnitfulMatrix(Diagonal(fill(σₙ.^2,length(y))),vec(unit.(y)).^1,vec(unit.(y)).^-1,exact=true)
+        Cxx = UnitfulMatrix(Diagonal(fill(σₓ.^2,length(x₀))),vec(unit.(x₀)),vec(unit.(x₀)).^-1,exact=true)
+
+        #create problem and solve
+        problem = UnderdeterminedProblem(UnitfulMatrix(vec(y)),E,Cnn,Cxx,x₀)
+        x̃ = solve(problem)
+        for jj in eachindex(vec(y))
+            @test within(y[jj],getindexqty(E*x̃.v,jj),3σₙ) # within 3-sigma
+        end
+
+        @test cost(x̃,problem) < 0.5 # ad-hoc choice
+
+    end
+
     @testset "source water inversion: obs TIMESERIES, many surface regions, with circulation lag" begin
 
         @dim YearCE "years Common Era"
@@ -346,6 +449,9 @@ include("test_functions.jl")
         # no noise in obs but some control penalty
         @test cost(x̃,problem) < 0.5 # ad-hoc choice
 
+        @test cost(x̃, problem) == datacost(x̃, problem) + controlcost(x̃, problem)
+
+
     end
 
     @testset "source water inversion: MANY OBS TIMESERIES, many surface regions, with circulation lag" begin

test/test_functions.jl

@@ -106,6 +106,17 @@ function source_water_solution(surfaceregions,years)
     return x
 end
 
+function source_water_solution(surfaceregions, years, statevar)
+    yr = u"yr"
+    K = u"K"
+    permil = u"permille"
+    m = length(years)
+    n = length(surfaceregions)
+    mat = cat(randn(m, n, 1)K, randn(m, n, 1)permil; dims = 3)
+    x = DimArray(mat, (Ti(years), SurfaceRegion(surfaceregions), StateVariable(statevar)))
+    return x
+end
+
 """
     function source_water_matrix_vector_pair_with_lag(M)