Discrete heat equation in Futhark

heat.fut

-- Apply the transformation field -> updated field for
-- a discrete temperature field f)i,j)
--
--    f'(i,j) = (1-beta)*f(i,j) + beta*average(i,j)
--
-- where average(i,j) is the average of the
-- temperatures of the cells to the North, South,
-- East. and West.

-- This example is adapted from blur.fut

-- Compute the new value of the temperature field at the given cell.  The
-- cell must not be located on the edges or an out-of-bounds access
-- will occur.
let newValue [rows][cols]
             (beta: f32) (field: [rows][cols]f32) (row: i32) (col: i32): f32 =
  -- The Futhark compiler cannot prove that these accesses are safe,
  -- and cannot perform dynamic bounds checks in parallel code.  We
  -- use the 'unsafe' keyword to elide the bounds checks.  If we did
  -- not do this, the code generator would fail with an error message.
  unsafe
  let sum =
      field[row-1,col] + field[row+1,col]
    + field[row,  col-1] + field[row,  col+1]
  in (1-beta) * field[row,col] + beta * (sum / 4f32)

-- Compute the new field: call newValue on every cell in the interior,
-- leaving the edges unchanged.
let newField [rows][cols]
                (beta: f32) (field: [rows][cols]f32): [rows][cols]f32 =
  unsafe
  map (\row ->
        map(\col ->
              if row > 0 && row < rows-1 && col > 0 && col < cols-1
              then newValue beta field row col
              else field[row,col])
            (0...cols-1))
      (0...rows-1)

-- Perform the specified number of updates om the given temperature field.
--   SIMPLEST TEST:
--   [8]> let data = [[0, 0, 0], [0, 1, 0], [0, 0, 0]]:[3][3]f32
--   [9]> main 1 0.5 data
--       [[0.0f32, 0.0f32, 0.0f32], [0.0f32, 0.5f32, 0.0f32], [0.0f32, 0.0f32, 0.0f32]]
let main [rows][cols]
         (iterations: i32) (beta: f32) (field: [rows][cols]f32): [rows][cols]f32 =
  let field = loop field for _i < iterations do newField beta field
  in field

Below is another try. Encouraging, but the types
are wrong. What should I be doing?

Python program

# File: heatDemo.py
import numpy as np
import heat

np.random.seed(0)                              # seed for reproducibility
x2 = np.random.uniform(0,0.1,25).reshape(5,5)  # 2D array of low-level noise
x2[2,2] = 1.0                                  # hot spot

heat.heat.main(heat.heat(),1, 0.5, x2)

Result

 $ python heatDemo.py
/Applications/anaconda/lib/python2.7/site-packages/pyopencl/__init__.py:235: CompilerWarning: Non-empty compiler output encountered. Set the environment variable PYOPENCL_COMPILER_OUTPUT=1 to see more.
  "to see more.", CompilerWarning)
Traceback (most recent call last):
  File "heatDemo.py", line 8, in <module>
    heat.heat.main(heat.heat(),1, 0.5, x2)
  File "/Users/carlson/dev/futhark/heat/heat.py", line 2354, in main
    field_mem_4750_ext))
TypeError: Argument #2 has invalid value
Futhark type: [][]f32
Argument has Python type <type 'numpy.ndarray'> and value: [[ 0.05488135  0.07151894  0.06027634  0.05448832  0.04236548]
 [ 0.06458941  0.04375872  0.0891773   0.09636628  0.03834415]
 [ 0.0791725   0.05288949  1.          0.09255966  0.00710361]
 [ 0.00871293  0.00202184  0.08326198  0.07781568  0.08700121]
 [ 0.09786183  0.07991586  0.04614794  0.07805292  0.01182744]]