Heat equation in Bohrium/Numpy and Futhark

bh.py

#!/usr/bin/env python

import bohrium
import bohrium as np
import pyopencl as cl
import time

w = 1000
h = 1000
epsilon=(42.0)*(np.sqrt(w*h)/10)

def heat2d(height, width, epsilon):
  G = np.zeros((height+2,width+2),dtype=np.float64)
  G[:,0]  = -273.15
  G[:,-1] = -273.15
  G[-1,:] = -273.15
  G[0,:]  = 40.0
  center = G[1:-1,1:-1]
  north  = G[:-2,1:-1]
  south  = G[2:,1:-1]
  east   = G[1:-1,:-2]
  west   = G[1:-1,2:]
  delta  = epsilon+1
  while delta > epsilon:
    tmp = 0.2*(center+north+south+east+west)
    delta = np.sum(np.abs(tmp-center))
    center[:] = tmp
  return center

bohrium_res = heat2d(w, h, epsilon)

import heat2d

ctx = bohrium.interop_pyopencl.get_context()
c = heat2d.heat2d(command_queue=cl.CommandQueue(ctx))
futhark_res = c.heat2d(w, h, epsilon).get()

print (np.all(bohrium_res == futhark_res))

heat2d.fut

-- Compile with 'futhark-pyopencl --library heat2d.fut'

entry heat2d (height: i32) (width: i32) (epsilon: f64) =
  let G = replicate (height+2) (replicate (width+2) 0.0)
  let G[:,0]        = replicate (height+2) (-273.15)
  let G[:,width+1]  = replicate (height+2) (-273.15)
  let G[height+1,:] = replicate (width+2)  (-273.15)
  let G[0,:]        = replicate (width+2)  40.0
  let delta = epsilon + 1.0
  let (G', _) = loop (G, delta) while delta > epsilon do
     let center = G[1:height+1,1:width+1]
     let north  = G[:height,1:width+1]
     let south  = G[2:,     1:width+1]
     let east   = G[1:height+1,:width]
     let west   = G[1:height+1,2:]
     let tmp = map (\r -> map (\(c,n,s,e,w) -> 0.2 * (c + n + s + e + w)) r)
                   (zip@1 center north south east west)
     let delta = reduce (+) 0.0 (map f64.abs (map (-) (flatten tmp) (flatten center)))
     let G[1:height+1,1:width+1] = tmp
     in (G, delta)
  in G'[1:height+1,1:width+1]