Gaussian Markov Random Field simulation using Fourier properties and base matrices for efficiency

gmrf_simulation.py

'''
Gaussian Markov Random Field simulation using Fourier properties and base matrices for efficiency
References from Gaussian Markov Random Fields: Theory and Applications, Havard Rue and Leonhard Held

We treat the 2D case, with 0 mean, stationary variance and exponential correlation function
'''

import matplotlib.pyplot as plt
import numpy as np
from scipy.fftpack import fft2, ifft2

def euclidean_dist_torus(x1, x2, y1, y2, lx, ly):
    '''
    Euclidean distance on the torus
    '''
    return np.sqrt(min(np.abs(x1 - x2), lx - np.abs(x1 - x2)) ** 2 +
                   min(np.abs(y1 - y2), ly - np.abs(y1 - y2)) ** 2)

def corr_exp(x1, x2, y1, y2, lx, ly, r):
    '''
    Exponential correlation function with range r
    '''
    try:
        return np.exp(-euclidean_dist_torus(x1, x2, y1, y2, lx, ly) / r)
    except FloatingPointError as e:
        return 0

def get_base_invert(b):
    '''
    If b is the base of a matrix B, returns the base of B^-1 using Fourier space
    Formula 2.50 from the book
    '''
    lx, ly = b.shape
    B = np.fft.fft2(b, norm='ortho')
    np.seterr(all='ignore')
    res = 1 / (lx * ly) * np.real(np.fft.ifft2(np.power(B, -1), norm='ortho'))
    np.seterr(all='raise')
    return res

def fourier_sampling_gaussian_field(r, lx, ly):
    '''
    Algorithm 2.10 from the book
    '''
    X = np.zeros((lx, ly), dtype=int)
    Z = np.random.normal(X, 1)
    Z = Z.astype(np.complex)
    Z += 1j * np.random.normal(X, 1)
    Z = Z.reshape((lx, ly))
    L = np.fft.fft2(r)
    Y = np.real(np.fft.fft2(np.multiply(
            np.power(L, -0.5), Z), norm="ortho"))

    return Y

if __name__ == "__main__":
    range1 = 20
    sigma = 1
    lx, ly = 512, 512
    # construct the base of the correlation matrix
    r = np.zeros((lx, ly))
    for x in range(lx):
        for y in range(ly):
            r[x, y] = sigma**2 * corr_exp(0, x, 0, y, lx, ly, range1)
    # get the base of the precision matrix
    q = get_base_invert(r)

    # check inversion formula
    rr = get_base_invert(q)
    assert r.all() == rr.all()
    
    # sample and print the result
    Y = fourier_sampling_gaussian_field(q, lx, ly)
    print("Estimated field standard deviation:", np.std(Y))
    print("Estimated field mean:", np.mean(Y))
    fig = plt.figure()
    ax = fig.add_subplot('111')
    ax.imshow(Y, cmap="gray")
    fig.show()
    plt.show()