Use GaussianRandomFields to calculate optimal proposal covariance

Project.toml

@@ -9,6 +9,7 @@ EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FieldMetadata = "bf96fef3-21d2-5d20-8afa-0e7d4c32a885"
 Future = "9fa8497b-333b-5362-9e8d-4d0656e87820"
+GaussianRandomFields = "e4b2fa32-6e09-5554-b718-106ed5adafe9"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
@@ -23,6 +24,7 @@ Distributions = "0.22, 0.23, 0.24, 0.25"
 EllipsisNotation = "1.1.0"
 FFTW = "1"
 FieldMetadata = "0.3.1"
+GaussianRandomFields = "2.1.1"
 HDF5 = "0.14, 0.15, 0.16"
 MPI = "0.19"
 TimerOutputs = "0.5"

src/OptimalFilter.jl

@@ -3,6 +3,7 @@ using LinearAlgebra
 using SparseArrays
 using Random
 using Distributions
+using GaussianRandomFields
 
 struct OfflineMatrices{T<:AbstractArray, S<:AbstractArray, U<:AbstractArray}
 
@@ -45,15 +46,13 @@ struct Grid{T}
     y_length::T
 end
 
-# Covariance function r(x,y), equation 1 in Dietrich and Newsam 96
-function covariance(x::T, y::T, noise_params::NamedTuple) where T
-
-    return noise_params.sigma^2 * exp(-(abs(x) + abs(y))/(2 * noise_params.lambda))
-
+# Covariance function r(x,y), equation 1 in Dietrich & Newsam 1996
+function covariance(x::T, y::T, covariance_structure::IsotropicCovarianceStructure{T}) where T
+    return covariance_structure.σ^2 * apply(covariance_structure, [x, y])
 end
 
-# Extended covariance function /bar r(x,y), equation 8 of Dietrich and Newsam 96
-function extended_covariance(x::T, y::T, grid::Grid, noise_params::NamedTuple) where T
+# Extended covariance function ̄r(x,y), equation 8 of Dietrich & Newsam 1996
+function extended_covariance(x::T, y::T, grid::Grid, covariance_structure::IsotropicCovarianceStructure{T}) where T
 
     if 0 <= x <= grid.x_length
         x_ext = x
@@ -71,12 +70,18 @@ function extended_covariance(x::T, y::T, grid::Grid, noise_params::NamedTuple) w
         @error "value of y is out of bounds"
     end
 
-    return covariance(x_ext, y_ext, noise_params)
+    return covariance(x_ext, y_ext, covariance_structure)
 
 end
 
-# Covariance between observations and the extended grid /bar R_21, from equation 11 in Dietrich & Newsam 96
-function covariance_stations_extended_grid!(cov::AbstractMatrix{T}, grid::Grid, grid_ext::Grid, stations::NamedTuple, noise_params::NamedTuple) where T
+# Covariance between observations and the extended grid R̅₂₁, from equation 11 in Dietrich & Newsam 1996
+function covariance_stations_extended_grid!(
+    cov::AbstractMatrix{T},
+    grid::Grid,
+    grid_ext::Grid,
+    stations::NamedTuple,
+    covariance_structure::IsotropicCovarianceStructure{T}
+) where T
 
     xid = 1:grid_ext.nx
     yid = 1:grid_ext.ny
@@ -88,13 +93,13 @@ function covariance_stations_extended_grid!(cov::AbstractMatrix{T}, grid::Grid,
     for (i,ist,jst) in zip(1:stations.nst, stations.ist, stations.jst)
         cov[i,:] .= extended_covariance.(abs.(getindex.(c, 1) .- ist) .* grid_ext.dx,
                                          abs.(getindex.(c, 2) .- jst) .* grid_ext.dy,
-                                         (grid,), (noise_params,))
+                                         (grid,), (covariance_structure,))
     end
 
 end
 
-# Covariance between stations and the original grid R_21, from equations 3-4 in Dietrich & Newsam 96
-function covariance_grid_stations!(cov::AbstractMatrix{T}, grid::Grid, stations::NamedTuple, noise_params::NamedTuple) where T
+# Covariance between stations and the original grid R₂₁, from equations 3 and 4 in Dietrich & Newsam 1996
+function covariance_grid_stations!(cov::AbstractMatrix{T}, grid::Grid, stations::NamedTuple, covariance_structure::IsotropicCovarianceStructure{T}) where T
 
     xid = 1:grid.nx
     yid = 1:grid.ny
@@ -106,24 +111,25 @@ function covariance_grid_stations!(cov::AbstractMatrix{T}, grid::Grid, stations:
     for (i,ist,jst) in zip(1:stations.nst, stations.ist, stations.jst)
         cov[:,i] .= extended_covariance.(abs.(getindex.(c, 1) .- ist) .* grid.dx,
                                          abs.(getindex.(c, 2) .- jst) .* grid.dy,
-                                         (grid,), (noise_params,))
+                                         (grid,), (covariance_structure,))
     end
 
 end
 
-# Covariance between observations R_22, from equationd 3-4 in in Dietrich & Newsam 96
-# TODO: Ask Alex why we add sigma^2 on the diagonal
-function covariance_stations!(cov::AbstractMatrix{T}, grid::Grid, stations::NamedTuple, noise_params::NamedTuple, std::T) where T
+# Covariance between observations R̅₂₂ = R₂₂ + Σ, where R₂₂ is as defined in equations 3 and 4 in
+# Dietrich & Newsam 96 and Σ is the observation noise covariance (see note at end of
+# Section 2 in Dietrich & Newsam 1996)
+function covariance_stations!(cov::AbstractMatrix{T}, grid::Grid, stations::NamedTuple, covariance_structure::IsotropicCovarianceStructure{T}, std::T) where T
 
     cov .= covariance.(abs.(stations.ist .- stations.ist') .* grid.dx,
                        abs.(stations.jst' .- stations.jst) .* grid.dy,
-                       (noise_params,)) .+ I(stations.nst) .* std.^2
+                       (covariance_structure,)) .+ I(stations.nst) .* std.^2
 
 end
 
-# First column vector rho_bar of covariance matrix among pairs of points of the extended grid R11_bar,
-# from Dietrich & Newsam 96 described in text between equations 11 and 12
-function first_column_covariance_grid!(rho::AbstractVector{T}, grid::Grid, grid_ext::Grid, noise_params::NamedTuple) where T
+# First column vector rho_bar of covariance matrix among pairs of points of the extended grid R̅₁₁,
+# from Dietrich & Newsam 1996 described in text between equations 11 and 12
+function first_column_covariance_grid!(rho::AbstractVector{T}, grid::Grid, grid_ext::Grid, covariance_structure::IsotropicCovarianceStructure{T}) where T
 
     xid = 1:grid_ext.nx
     yid = 1:grid_ext.ny
@@ -132,12 +138,12 @@ function first_column_covariance_grid!(rho::AbstractVector{T}, grid::Grid, grid_
 
     rho .= extended_covariance.((getindex.(c, 1) .- 1) .* grid_ext.dx,
                                 (getindex.(c, 2) .- 1) .* grid_ext.dy,
-                                (grid,), (noise_params,))
+                                (grid,), (covariance_structure,))
 
 
 end
 
-# Normalized two-dimension discrete Fourier transofrm normalized by sqrt(n1_bar).
+# Normalized two-dimension discrete Fourier transofrm normalized by sqrt(̄n₁).
 # Operates on the 2d data stored as a matrix.
 # The argument `f` lets you switch between forward transform (`f=identity`) and
 # inverse transform (`f=inv`).
@@ -154,7 +160,7 @@ function normalized_2d_fft!(transformed_array::AbstractMatrix{<:Complex}, array:
 
 end
 
-# Normalized two-dimension discrete Fourier transofrm normalized by sqrt(n1_bar).
+# Normalized two-dimension discrete Fourier transofrm normalized by sqrt(̄n₁).
 # Operates on the 2d data stored as a vector.
 # The argument `f` lets you switch between forward transform (`f=identity`) and
 # inverse transform (`f=inv`).
@@ -168,7 +174,7 @@ function normalized_2d_fft!(transformed_vector::AbstractVector{<:Complex}, vecto
 
 end
 
-# Decomposition of R11, equation 12 of Deitrich and Newsam
+# Decomposition of R̅₁₁, equation 12 of Deitrich & Newsam 1996
 function WΛWH_decomposition!(transformed_array::AbstractMatrix{T},
                              array::AbstractMatrix{T},
                              offline_matrices::OfflineMatrices,
@@ -211,7 +217,7 @@ end
 function init_offline_matrices(grid::Grid,
                                grid_ext::Grid,
                                stations::NamedTuple,
-                               noise_params::NamedTuple,
+                               covariance_structure::IsotropicCovarianceStructure{T},
                                obs_noise_std::T,
                                fft_plan::FFTW.FFTWPlan,
                                fft_plan!::FFTW.FFTWPlan,
@@ -238,10 +244,10 @@ function init_offline_matrices(grid::Grid,
                                Matrix{F}(undef, grid.nx, grid.ny)  #buf2
                                )
 
-    first_column_covariance_grid!(matrices.rho_bar, grid, grid_ext, noise_params)
-    covariance_grid_stations!(matrices.R12, grid, stations, noise_params)
-    covariance_stations_extended_grid!(matrices.R21_bar, grid, grid_ext, stations, noise_params)
-    covariance_stations!(matrices.R22, grid, stations, noise_params, obs_noise_std)
+    first_column_covariance_grid!(matrices.rho_bar, grid, grid_ext, covariance_structure)
+    covariance_grid_stations!(matrices.R12, grid, stations, covariance_structure)
+    covariance_stations_extended_grid!(matrices.R21_bar, grid, grid_ext, stations, covariance_structure)
+    covariance_stations!(matrices.R22, grid, stations, covariance_structure, obs_noise_std)
     matrices.R22_inv .= inv(matrices.R22)
     matrices.R12_invR22 .= matrices.R12 * matrices.R22_inv
 

src/ParticleDA.jl

@@ -50,9 +50,10 @@ Return standard deviation of observation noise. Required for optimal filter only
 function get_obs_noise_std end
 
 """
-    ParticleDa.get_model_noise(model_data) -> NamedTuple({:sigma, :lambda, :nu},Tuple{Float, Float, Float})
+    ParticleDa.get_model_noise(model_data) -> GaussianRandomFields.IsotropicCovarianceStructure
 
-Return parameters of the noise covariance function of the first state variable (height). Required for optimal filter only.
+Return structure defining parameters of the covariance function of the state noise in
+the observed state variable. Required for optimal filter only.
 """
 function get_model_noise_params end
 

test/Project.toml

@@ -8,6 +8,7 @@ HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/model/model.jl

@@ -439,9 +439,9 @@ ParticleDA.get_stations(d::ModelData) = (nst = d.model_params.nobs,
                                          ist = d.stations.ist,
                                          jst = d.stations.jst)
 ParticleDA.get_obs_noise_std(d::ModelData) = d.model_params.obs_noise_std
-ParticleDA.get_model_noise_params(d::ModelData) = (sigma = d.model_params.sigma[1],
-                                                   lambda = d.model_params.lambda[1],
-                                                   nu = d.model_params.nu[1])
+ParticleDA.get_model_noise_params(d::ModelData) = Matern(d.model_params.lambda[1],
+                                                         d.model_params.nu[1],
+                                                         σ=d.model_params.sigma[1])
 
 function ParticleDA.set_particles!(d::ModelData, particles::AbstractArray{T}) where T
 

test/optimal_filter_validation.jl

@@ -9,6 +9,7 @@ using Distributions
 using LinearAlgebra
 using FFTW
 using Random
+using SpecialFunctions
 
 ####################################################################
 ####################################################################
@@ -17,8 +18,9 @@ Random.seed!(123)
 
 #### Parameters for Covariance Function
 struct f_th
-    s::Float64
-    l::Float64
+    σ::Float64
+    λ::Float64
+    ν::Float64
 end
 
 #### Parameters for Grid
@@ -40,7 +42,7 @@ end
 ####################################################################
 #### Example #######################################################
 
-th = f_th(1.0, 1.0)
+th = f_th(1.0, 1.0, 0.5)
 
 gr = f_gr(20, 20, 40.0, 40.0)
 
@@ -95,12 +97,15 @@ Sobs = 0.01
 ###################################################
 ###################################################
 
-#### Example Specification of Covariance Function
-function R(x::Float64, y::Float64, th)
+#### Example Specification of (Matern) Covariance Function
+function R(x::Float64, y::Float64, th::f_th)
 
-  Cov = (th.s^2)*exp(-(abs(x)+abs(y))/(2.0*th.l))
-
-  return(Cov)
+  ρ = hypot(x, y)
+  if iszero(ρ)
+    float(one(x))
+  else
+    th.σ^2 * 2^(1 - th.ν) / gamma(th.ν) * (ρ / th.λ)^th.ν * besselk(th.ν, ρ / th.λ)
+  end
 
 end
 

test/runtests.jl

@@ -3,6 +3,7 @@ using LinearAlgebra, Test, HDF5, Random, YAML
 using MPI
 using StableRNGs
 using FFTW
+using GaussianRandomFields: Matern
 
 using ParticleDA: FilterParameters
 
@@ -278,8 +279,10 @@ end
     seed = 123
     Random.seed!(seed)
 
-    # Use default parameters
-    model_params = ModelParameters()
+    # Use default parameters other than smoothness parameter ν for state noise Matern 
+    # covariance function  which is fixed to 0.5 so that covariance function reduces to 
+    # r(x, y) = σ^2 exp(-norm([x, y]) / λ)
+    model_params = ModelParameters(nu=[0.5, 0.5, 0.5])
     filter_params = FilterParameters()
     grid = ParticleDA.Grid(model_params.nx,
                            model_params.ny,
@@ -293,14 +296,14 @@ end
                                grid.dy,
                                (grid.x_length-grid.dx)*2,
                                (grid.y_length-grid.dy)*2)
-    noise_params = (sigma = model_params.sigma[1], lambda = model_params.lambda[1], nu = model_params.nu[1])
+    noise_params =  Matern(model_params.lambda[1], model_params.nu[1], σ = model_params.sigma[1])
 
     # Set station coordinates
     ist = rand(1:model_params.nx, model_params.nobs)
     jst = rand(1:model_params.ny, model_params.nobs)
     stations = (nst = model_params.nobs, ist = ist, jst = jst)
     cov_ext = ParticleDA.extended_covariance(0.0, 0.5 * grid.y_length, grid, noise_params)
-    @test cov_ext ≈ exp(-0.5 * model_params.y_length / (2 * noise_params.lambda))
+    @test cov_ext ≈ noise_params.σ^2 * exp(-norm([0.0, 0.5 * model_params.y_length]) / noise_params.λ)
     @test cov_ext ≈ ParticleDA.extended_covariance(2.0 * grid.x_length, 0.5 * grid.y_length, grid, noise_params)
     @test cov_ext ≈ ParticleDA.extended_covariance(0.0, 1.5 * grid.y_length, grid, noise_params)
     arr = rand(ComplexF64,10,10)
@@ -363,7 +366,7 @@ end
                                (grid.x_length-grid.dx)*2,
                                (grid.y_length-grid.dy)*2)
 
-    noise_params = (sigma = model_params.sigma[1], lambda = model_params.lambda[1], nu = model_params.nu[1])
+    noise_params =  Matern(model_params.lambda[1], model_params.nu[1], σ = model_params.sigma[1])
 
     stations = (nst = model_params.nobs, ist = st.st_ij[:,1].+1, jst = st.st_ij[:,2].+1)
 
@@ -390,6 +393,7 @@ end
     Yobs_t = copy(obs)
     FH_t = copy(reshape(permutedims(height, [3 1 2]), filter_params.nprt, (model_params.nx)*(model_params.ny)))
 
+    th = f_th(noise_params.σ[1], noise_params.λ[1], noise_params.ν[1])
     Mean_height = Calculate_Mean(FH_t, th, st, Yobs_t, Sobs, gr)
     Samples = Sample_Height_Proposal(FH_t, th, st, Yobs_t, Sobs, gr)