2

This is a question of speed. I'm trying to solve a diffusion equation which has three behavior states where:

  • Lambda == 0 equilibrium
  • Lambda > 0 max diffusion
  • Lambda < 0 min diffusion

Bottleneck is the else statement in the diffusion operator function

The equilibrium state has a simple T operator and Diffusion operator. It just gets rather complicated for the other two states. And so far I haven't had the patience to sit through the code run times. As far as I know the equations are correct, and the output of the equilibrium state appears correct, perhaps someone has some tips to increase the speed for the non-equilibrium states?

(Euler solution (FTCS) instead of Runge-Kutta would be quicker I imagine. Haven't tried this yet.)

You can import any black and white image to try the code out on.

import numpy as np
import sympy as sp
import scipy.ndimage.filters as flt
from PIL import Image

# import image
im = Image.open("/home/will/Downloads/zebra.png")
arr = np.array(im)
arr=arr/253.

def T(lamda,x): 
    """
    T Operator
    lambda is a "steering" constant between 3 behavior states
    -----------------------------
    0     -> linearity 
    +inf  -> max
    -inf  -> min
    -----------------------------
    """    
    if lamda == 0:  # linearity
        return x
    elif lamda > 0: #  Half-wave rectification
        return np.max(x,0)
    elif lamda < 0: # Inverse half-wave rectification
        return np.min(x,0)


def Diffusion_operator(lamda,f,t): 
    """
    2D Spatially Discrete Non-Linear Diffusion
    ------------------------------------------
    Special case where lambda == 0, operator becomes Laplacian        


    Parameters
    ----------
    D : int                      diffusion coefficient
    h : int                      step size
    t0 : int                     stimulus injection point
    stimulus : array-like        luminance distribution     

    Returns
    ----------
    f : array-like               output of diffusion equation
    -----------------------------
    0     -> linearity (T[0])
    +inf  -> positive K(lamda)
    -inf  -> negative K(lamda)
    -----------------------------
    """
    if lamda == 0:  # linearity
        return flt.laplace(f)
    else:           # non-linearity
        f_new = np.zeros(f.shape)
        for x in np.arange(0,f.shape[0]-1):
            for y in np.arange(0,f.shape[1]-1):
                f_new[x,y]=T(lamda,f[x+1,y]-f[x,y]) + T(lamda,f[x-1,y]-f[x,y]) + T(lamda,f[x,y+1]-f[x,y])
                + T(lamda,f[x,y-1]-f[x,y])
        return f_new


def Dirac_delta_test(tester):
    # Limit injection to unitary multiplication, not infinite
    if np.sum(sp.DiracDelta(tester)) == 0:
        return 0
    else:
        return 1

def Runge_Kutta(stimulus,lamda,t0,h,N,D,t_N):
    """
    4th order Runge-Kutta solution to:
    linear and spatially discrete diffusion equation (ignoring leakage currents)

    Adiabatic boundary conditions prevent flux exchange over domain boundaries
    Parameters
    ---------------
    stimulus : array_like   input stimuli [t,x,y]
    lamda : int             0 +/- inf
    t0 : int                point of stimulus "injection"
    h : int                 step size
    N : int                 array size (stimulus.shape[1])
    D : int                 diffusion coefficient [constant]

    Returns
    ----------------
    f : array_like          computed diffused array

    """
    f = np.zeros((t_N+1,N,N)) #[time, equal shape space dimension]
    t = np.zeros(t_N+1)

    if lamda ==0:
        """    Linearity    """
        for n in np.arange(0,t_N):
            k1 = D*flt.laplace(f[t[n],:,:]) +     stimulus*Dirac_delta_test(t[n]-t0)
            k1 = k1.astype(np.float64)
            k2 = D*flt.laplace(f[t[n],:,:]+(0.5*h*k1)) +     stimulus*Dirac_delta_test((t[n]+(0.5*h))- t0)
            k2 = k2.astype(np.float64)
            k3 = D*flt.laplace(f[t[n],:,:]+(0.5*h*k2)) + stimulus*Dirac_delta_test((t[n]+(0.5*h))-t0)
            k3 = k3.astype(np.float64)
            k4 = D*flt.laplace(f[t[n],:,:]+(h*k3)) + stimulus*Dirac_delta_test((t[n]+h)-t0)
            k4 = k4.astype(np.float64)
            f[n+1,:,:] = f[n,:,:] + (h/6.) * (k1 + 2.*k2 + 2.*k3 + k4)
            t[n+1] = t[n] + h
        return f

    else:
        """    Non-Linearity   THIS IS SLOW  """
        for n in np.arange(0,t_N):
            k1 = D*Diffusion_operator(lamda,f[t[n],:,:],t[n]) + stimulus*Dirac_delta_test(t[n]-t0)
            k1 = k1.astype(np.float64)
            k2 = D*Diffusion_operator(lamda,(f[t[n],:,:]+(0.5*h*k1)),t[n]) + stimulus*Dirac_delta_test((t[n]+(0.5*h))- t0)
            k2 = k2.astype(np.float64)
            k3 = D*Diffusion_operator(lamda,(f[t[n],:,:]+(0.5*h*k2)),t[n]) + stimulus*Dirac_delta_test((t[n]+(0.5*h))-t0)
            k3 = k3.astype(np.float64)
            k4 = D*Diffusion_operator(lamda,(f[t[n],:,:]+(h*k3)),t[n]) + stimulus*Dirac_delta_test((t[n]+h)-t0)
            k4 = k4.astype(np.float64)
            f[n+1,:,:] = f[n,:,:] + (h/6.) * (k1 + 2.*k2 + 2.*k3 + k4)
            t[n+1] = t[n] + h

        return f

# Code to run
N=arr.shape[1]  # Image size
stimulus=arr[0:N,0:N,1]
D = 0.3   # Diffusion coefficient [0>D>1]
h = 1     # Runge-Kutta step size [h > 0]
t0 = 0    # Injection time 
t_N = 100 # End time

f_out_equil = Runge_Kutta(stimulus,0,t0,h,N,D,t_N)
f_out_min = Runge_Kutta(stimulus,-1,t0,h,N,D,t_N)
f_out_max = Runge_Kutta(stimulus,1,t0,h,N,D,t_N)

In short, f_out_equil is relatively quick to calculate, whereas min and max cases are expensive and slow.

Here's a link to an image I have been using: http://4.bp.blogspot.com/_KbtOtXslVZE/SweZiZWllzI/AAAAAAAAAIg/i9wc-yfdW78/s200/Zebra_Black_and_White_by_Jenvanw.jpg

Tips on improving my coding appreciated, Many thanks,

Here's a quick plotting script for the output

import matplotlib.pyplot as plt
fig1, (ax1,ax2,ax3,ax4,ax5) = plt.subplots(ncols=5, figsize=(15,5))
ax1.imshow(f_out_equil[1,:,:],cmap='gray')
ax2.imshow(f_out_equil[t_N/10,:,:],cmap='gray')
ax3.imshow(f_out_equil[t_N/2,:,:],cmap='gray')
ax4.imshow(f_out_equil[t_N/1.5,:,:],cmap='gray')
ax5.imshow(f_out_equil[t_N,:,:],cmap='gray')
William Baker Morrison
  • 1,642
  • 4
  • 21
  • 33

1 Answers1

5

For loops in python tend to be slow; you can get a massive speedup by vectorizing as much as you can. (This is going to help you a lot with any numerical problems going forward). The new T operator works on the full array all at once, and the calls to np.roll in Diffusion_operator line the image array up properly for the finite difference calculations.

Whole thing ran in about 10 s on my computer.

def T(lamda,x): 
    """
    T Operator
    lambda is a "steering" constant between 3 behavior states
    -----------------------------
    0     -> linearity 
    +inf  -> max
    -inf  -> min
    -----------------------------
    """    
    if lamda == 0:  # linearity
        return x
    elif lamda > 0: #  Half-wave rectification
        maxval = np.zeros_like(x)
        return np.array([x, maxval]).max(axis=0)
    elif lamda < 0: # Inverse half-wave rectification
        minval = np.zeros_like(x)
        return np.array([x, minval]).min(axis=0)


def Diffusion_operator(lamda,f,t): 
    """
    2D Spatially Discrete Non-Linear Diffusion
    ------------------------------------------
    Special case where lambda == 0, operator becomes Laplacian        


    Parameters
    ----------
    D : int                      diffusion coefficient
    h : int                      step size
    t0 : int                     stimulus injection point
    stimulus : array-like        luminance distribution     

    Returns
    ----------
    f : array-like               output of diffusion equation
    -----------------------------
    0     -> linearity (T[0])
    +inf  -> positive K(lamda)
    -inf  -> negative K(lamda)
    -----------------------------
    """
    if lamda == 0:  # linearity
        return flt.laplace(f)
    else:           # non-linearity
        f_new = T(lamda,np.roll(f,1, axis=0) - f) \
                + T(lamda,np.roll(f,-1, axis=0) - f) \
                + T(lamda,np.roll(f, 1, axis=1) - f) \
                + T(lamda,np.roll(f,-1, axis=1) - f)
        return f_new
Elliot
  • 2,541
  • 17
  • 27
  • Your redefinition of T function with this change of f_new works for me `f_new=T(lamda,np.roll(f,1, axis=0)-f)+T(lamda,np.roll(f,-1, axis=0)-f)+T(lamda,np.roll(f, 1, axis=1)-f)+T(lamda,np.roll(f,-1, axis=1)-f)` Thanks – William Baker Morrison Mar 31 '16 at 10:09
  • Alright, the answer now includes that. As a side note, if you'd wanted the edges in the original program to not be zero, you would have needed `for x in np.arange(f.shape[0]):`. `xrange(n)`, `range(n)`, and `np.arange(n)` all generate values in [0, 1, ..., n-1]. – Elliot Mar 31 '16 at 11:43
  • Thanks for this comment. I'm not sure about the relevance of edges at this stage but I'm sure I'll come back to it. – William Baker Morrison Apr 01 '16 at 08:35
  • would you have any ideas about how to set up boundary conditions that reduce the diffusion operators at the edges to existing neighbors? These are my initial thoughts: `for i in np.arange(N): #Adiabatic Boundary Conditions f[0,i] = np.sum(f[1,i] +f[0,i+1]+f[0,i-1])/3 f[N,i] = np.sum(f[N-1,i]+f[N,i+1]+f[N,i-1])/3 f[i,0] = np.sum(f[i,1] +f[i+1,0]+f[i-1,0])/3 f[i,N] = np.sum(f[i,N-1]+f[i+1,N]+f[i-1,N])/3` I'm aware this will get dimension errors, haven't run it yet @Elliot – William Baker Morrison Apr 18 '16 at 08:23
  • If you've got the equations for the boundary conditions you can figure out what the ghost point has to be from the points you have. What you've posted looks more like a meaningless average than anything else sorry to say. – Elliot Apr 18 '16 at 11:59
  • Maybe if I explain it would make more sense: f[0,i] is x axis minimum edge, f[N,i] is x maximum edge, f[i,0] is y min edge, f[i,N] is y max edge. Then it averages the pixel neighbor hood into its own value, i.e. up,down,left,right ignoring the pixel outside of the image. So basically I'm making up boundary condition equations to match the description. – William Baker Morrison Apr 19 '16 at 08:19
  • (sorry this took so long) "So basically I'm making up boundary condition equations to match the description"; that's kind of a problem. If you want adiabatic boundaries, you need to actually take care to make adiabatic boundaries (i.e. dT/dx=0 at the boundary). This probably means that the point beyond the boundary equals the point in front of the boundary, but you'll have to confirm that and implement it yourself. Yes it's a pain and may screw up your algorithm, but that's boundary conditions for ya. – Elliot Apr 26 '16 at 14:32