Softmax Gaussian Field generator

def kernel(X1, X2, l=1.0, eta=1.0):
    '''
    Isotropic squared exponential kernel. Computes 
    a covariance matrix from points in X1 and X2.
    
    Args:
        X1: Array of m points (m x d).
        X2: Array of n points (n x d).

    Returns:
        Covariance matrix (m x n).
    '''
    sqdist = np.sum(X1**2, 1).reshape(-1, 1) + np.sum(X2**2, 1) - 2 * np.dot(X1, X2.T)
    return eta**2 * np.exp(-0.5 / l**2 * sqdist)
    
    
    np.random.seed(192)

n_cell_types = 5 #zones
n1, n2 = (14*3, 8*3) #saptial dimensions 
x1 = np.linspace(0, 140, n1)[:,None] #saptial dimensions 
x2 = np.linspace(0, 80, n2)[:,None] #saptial dimensions 

# make cartesian grid out of each dimension x1 and x2
X = pm.math.cartesian(x1[:,None], x2[:,None])

l1_true = [8, 8, 12, 15, 10] #bw parameter
l2_true = [8, 8, 12, 15, 10]
eta_true = 1 #variance, defnies overlapping

#cov1, cov2 = kernel(x1, x1, l=l1_true), kernel(x2, x2, l=l2_true)
K = [np.kron(kernel(x1, x1, l=l1_true[i], eta=eta_true), kernel(x2, x2, l=l2_true[i], eta=eta_true)) for i in range(n_cell_types)]

gaus_true = np.stack([np.random.multivariate_normal(np.zeros(X.shape[0]), 2*K[i]) for i in range(n_cell_types)]).T #samples from GP
N_true = np.exp(gaus_true)/np.exp(gaus_true).sum(axis=1)[:,None] #sofmtax transofrm 


plt.figure(figsize=(25,3))
for ct in range(n_cell_types):
    plt.subplot(1,n_cell_types+1,ct+1)
    plt.imshow(N_true[:,ct].reshape(n1,n2).T, cmap=plt.cm.get_cmap('Reds'))
    plt.colorbar()
    plt.title('subclone {}'.format(ct+1))
    
plt.subplot(1,n_cell_types+1,n_cell_types+1) 
plt.imshow(N_true.sum(axis=1).reshape(n1,n2).T, cmap=plt.cm.get_cmap('Greys'))
plt.colorbar()
plt.title('n_cells')