pmf_filter.py

import numpy as np
from scipy.ndimage import gaussian_filter, laplace
from scipy.ndimage import binary_erosion
from scipy.signal import convolve2d
import matplotlib.pyplot as plt
import read_data as rd

# Set the diffusion parameters
kappa = 25  # Diffusion coefficient
gamma = 0.1  # Gradient threshold
dt = 1  # Time step


def convolve_around_mask(x, kernel, mask):
    # Compute the convolution of the mask with the kernel
    mask_conv = convolve2d(~mask, kernel, boundary='wrap', mode='same')

    # Convolve the kernel across x, ignoring NaNs
    result = convolve2d(x, kernel, boundary='wrap', mode='same')
    # print(np.any(np.isnan(mask_conv)), np.any(np.isnan(result)))

    # Divide by the sum of the kernel values at each pixel location
    mask_conv[mask_conv == 0] = 1
    result = result / mask_conv
    result[mask] = 0
    # print(np.any(np.isnan(x)), np.any(np.isnan(result)))

    return result


# Compute the gradient magnitude
def grad_mag(f, mask):
    # Define kernel for x differences
    kx = np.array([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
    # Define kernel for y differences
    ky = np.array([[1, 2, 1], [0, 0, 0],  [-1, -2, -1]])
    # Calculate x gradients
    fx = convolve_around_mask(f, kx, mask)
    # Calculate y gradients
    fy = convolve_around_mask(f, ky, mask)
    return np.sqrt(fx**2 + fy**2)


# Compute the conductivity function
def c(x, mask):
    return 1 / (1 + (grad_mag(x, mask) / gamma)**2)


# Apply the nonlinear anisotropic diffusion filtering
def anisotropic_diffusion(mdt):
    mask = np.isnan(mdt)

    # Compute the latitude-dependent diffusion coefficient
    n, m = mdt.shape
    lon, lat = np.meshgrid(np.linspace(0, 360, m, endpoint=False), np.linspace(-90, 90, n))
    D = kappa / np.abs(np.cos(np.deg2rad(lat)))

    # Initialize the filtered image
    mdt_filtered = mdt.copy()
    mdt_filtered[mask] = 0

    # Perform the anisotropic diffusion filtering
    for i in range(100):

        # Compute the Laplacian of the filtered image
        lap_mdt = laplace(np.nan_to_num(mdt_filtered), mode='wrap')

        # Compute the flux term while ignoring land values
        print(D.shape)
        flux = np.where(~mask.astype('bool'), c(mdt_filtered, mask) * D * lap_mdt, 0)
        # plt.imshow(flux)
        # plt.show()

        # Update the filtered image
        mdt_filtered += dt * flux
        plt.imshow(mdt_filtered, vmin=-1.5, vmax=1.5)
        plt.show()

    # Replace the land values with the original values
    mdt_filtered = np.where(~mask, mdt_filtered, mdt)

    return mdt_filtered


mdt = np.rot90(rd.read_surface('dtu18_gtim5_do0280_rr0004.dat', '../a_mdt_data/computations/mdts'))

mdt_filtered = anisotropic_diffusion(mdt)

# Display the results
fig, axs = plt.subplots(1, 2, figsize=(10, 5))
axs[0].imshow(mdt, vmin=-1.4, vmax=1.4, cmap='turbo')
axs[0].set_title('Original MDT')
axs[1].imshow(mdt_filtered, vmin=-1.4, vmax=1.4, cmap='turbo')
axs[1].set_title('Filtered MDT')
plt.show()