diff --git a/python/ergodic_control_SMC_2D.py b/python/ergodic_control_SMC_2D.py
index 68a896e4936046e4cdc6ead63608f7a60094c484..5555297fce17881966887ea31d53af7ef494dc89 100644
--- a/python/ergodic_control_SMC_2D.py
+++ b/python/ergodic_control_SMC_2D.py
@@ -13,6 +13,7 @@ License: GPL-3.0-only
 import numpy as np
 #from math import exp
 import matplotlib.pyplot as plt
+from scipy.interpolate import interpn
 
 
 # Helper functions
@@ -96,12 +97,34 @@ Sigma[:, :, 1] = (
 )
 Alpha = np.ones(nbGaussian) / nbGaussian
 
-
-# Compute Fourier series coefficients w_hat of desired spatial distribution
+# Fourier basis functions (for a discretized map)
 # ===============================
 rg = np.arange(0, nbFct, dtype=float)
 KX = np.zeros((nbVar, nbFct, nbFct))
 KX[0, :, :], KX[1, :, :] = np.meshgrid(rg, rg)
+Lambda = np.array(KX[0, :].flatten() ** 2 + KX[1, :].flatten() ** 2 + 1).T ** (-sp)
+
+xm1d = np.linspace(xlim[0], xlim[1], nbRes)  # Spatial range for 1D
+xm = np.zeros((nbGaussian, nbRes, nbRes))  # Spatial range
+xm[0, :, :], xm[1, :, :] = np.meshgrid(xm1d, xm1d)
+# Mind the flatten() !!!
+arg1 = (
+    KX[0, :, :].flatten().T[:, np.newaxis] @ xm[0, :, :].flatten()[:, np.newaxis].T * om
+)
+arg2 = (
+    KX[1, :, :].flatten().T[:, np.newaxis] @ xm[1, :, :].flatten()[:, np.newaxis].T * om
+)
+phim = np.cos(arg1) * np.cos(arg2) * 2 ** (nbVar)  # Fourier basis functions
+
+# Some weird +1, -1 due to 0 index!!!
+xx, yy = np.meshgrid(np.arange(1, nbFct + 1), np.arange(1, nbFct + 1))
+hk = np.concatenate(([1], 2 * np.ones(nbFct)))
+HK = hk[xx.flatten() - 1] * hk[yy.flatten() - 1]
+phim = phim * np.tile(HK, (nbRes**nbVar, 1)).T
+
+
+# Compute Fourier series coefficients w_hat of desired spatial Gaussian distribution
+# ===============================
 # Mind the flatten() !!!
 # Weighting vector (Eq.(16))
 Lambda = np.array(KX[0, :].flatten() ** 2 + KX[1, :].flatten() ** 2 + 1).T ** (-sp)
@@ -122,27 +145,32 @@ for j in range(nbGaussian):
         w_hat = w_hat + Alpha[j] * cos_term * exp_term
 w_hat = w_hat / (L**nbVar) / (op.shape[1])
 
-# Fourier basis functions (for a discretized map)
-# ===============================
-xm1d = np.linspace(xlim[0], xlim[1], nbRes)  # Spatial range for 1D
-xm = np.zeros((nbGaussian, nbRes, nbRes))  # Spatial range
-xm[0, :, :], xm[1, :, :] = np.meshgrid(xm1d, xm1d)
-# Mind the flatten() !!!
-arg1 = (
-    KX[0, :, :].flatten().T[:, np.newaxis] @ xm[0, :, :].flatten()[:, np.newaxis].T * om
-)
-arg2 = (
-    KX[1, :, :].flatten().T[:, np.newaxis] @ xm[1, :, :].flatten()[:, np.newaxis].T * om
-)
-phim = np.cos(arg1) * np.cos(arg2) * 2 ** (nbVar)  # Fourrier basis functions
-
-# Some weird +1, -1 due to 0 index!!!
-xx, yy = np.meshgrid(np.arange(1, nbFct + 1), np.arange(1, nbFct + 1))
-hk = np.concatenate(([1], 2 * np.ones(nbFct)))
-HK = hk[xx.flatten() - 1] * hk[yy.flatten() - 1]
-phim = phim * np.tile(HK, (nbRes**nbVar, 1)).T
 
-# Desired spatial distribution
+# # Uncomment if a non-Gaussian spatial distribution should be used
+# # Here, a 2D array representing the SDF of a joined rectangle and circle
+# # is used as the distribution to target.
+# # ===============================
+# data = np.load('../data/sdf01.npy', allow_pickle=True).item()
+#
+# # Load data, inverse value and threshold to keep only the postive value (inside of SDF)
+# g = (-1.0 * data['y']).reshape((40, 40)) # sdf01 is a 40x40 array
+# g[g<0.0] = 0.0
+#
+# # Rescale distribution to match the defined resolution
+# x = np.linspace(0, 1, g.shape[0])
+# xi = np.linspace(0, 1, nbRes)
+# grid_xx, grid_yy = np.meshgrid(xi, xi, indexing='ij')
+# new_grid = np.stack((grid_xx, grid_yy), axis=-1)
+# g = interpn((x, x), g, new_grid)
+# g = g.reshape((nbRes * nbRes,))
+# g = g * nbRes**nbVar / np.sum(g)
+#
+# # Compute Fourier coefficients
+# phi_inv = np.cos(arg1) * np.cos(arg2) / L**nbVar / nbRes**nbVar
+# w_hat = phi_inv @ g
+
+
+# Compute desired spatial distribution from w_hat
 g = w_hat.T @ phim
 
 # Ergodic control