more rbf and test

Author: HugoGranstrom

src/numericalnim/interpolate.nim

@@ -3,7 +3,10 @@ import arraymancer, cdt/[dt, vectors, edges, types]
 import
   ./utils,
   ./common/commonTypes,
-  ./private/arraymancerOverloads
+  ./private/arraymancerOverloads,
+  ./rbf
+
+export rbf
 
 type
   InterpolatorType*[T] = ref object

src/numericalnim/rbf.nim

@@ -3,46 +3,56 @@ import arraymancer
 
 
 type
+  RbfFunc* = proc (r: Tensor[float], epsilon: float): Tensor[float]
   RbfType*[T] = object
     points*: Tensor[float] # (n_points, n_dim)
     values*: Tensor[T] # (n_points, n_values)
     coeffs*: Tensor[float] # (n_points, n_values)
     epsilon*: float
+    f*: RbfFunc
 
+# Idea: blocked distance matrix for better cache friendliness
 proc distanceMatrix(p1, p2: Tensor[float]): Tensor[float] =
   ## Returns distance matrix of shape (n_points, n_points)
-  let n_points = p1.shape[0]
+  let n_points1 = p1.shape[0]
+  let n_points2 = p2.shape[0]
   let n_dims = p1.shape[1]
-  result = newTensor[float](n_points, n_points)
-  for i in 0 ..< n_points:
-    for j in 0 ..< n_points:
+  result = newTensor[float](n_points2, n_points1)
+  for i in 0 ..< n_points2:
+    for j in 0 ..< n_points1:
       var r2 = 0.0
       for k in 0 ..< n_dims:
-        let diff = p1[i,k] - p2[j,k]
+        let diff = p2[i,k] - p1[j,k]
         r2 += diff * diff
       result[i, j] = sqrt(r2)
 
 proc compactRbfFunc*(r: Tensor[float], epsilon: float): Tensor[float] =
   result = map_inline(r):
     (1 - x/epsilon) ^ 4 * (4*x/epsilon + 1) * float(x < epsilon)
 
-proc newRbf*[T](points: Tensor[float], values: Tensor[T], rbfFunc: proc (r: Tensor[float], epsilon: float): Tensor[float] = compactRbfFunc, epsilon: float = 1): RbfType[T] =
+proc newRbf*[T](points: Tensor[float], values: Tensor[T], rbfFunc: RbfFunc = compactRbfFunc, epsilon: float = 1): RbfType[T] =
   assert points.shape[0] == values.shape[0]
   let dist = distanceMatrix(points, points)
   let A = rbfFunc(dist, epsilon)
   let coeffs = solve(A, values)
-  result = RbfType[T](points: points, values: values, coeffs: coeffs, epsilon: epsilon)
-
-let x1 = @[@[0.0, 0.0, 0.0], @[1.0, 1.0, 0.0], @[1.0, 2.0, 0.0]].toTensor
-let x2 = @[@[0.0, 0.0, 1.0], @[1.0, 1.0, 2.0], @[1.0, 2.0, 3.0]].toTensor
-echo distanceMatrix(x1, x1)
-let values = @[0.0, 1.0, 2.0].toTensor
-echo newRbf(x1, values, epsilon=10)
-
-import benchy
-
-let pos = randomTensor(5000, 3, 1.0)
-let vals = randomTensor(5000, 1, 1.0)
-timeIt "Rbf":
-  keep newRbf(pos, vals)
+  result = RbfType[T](points: points, values: values, coeffs: coeffs, epsilon: epsilon, f: rbfFunc)
+
+proc eval*[T](rbf: RbfType[T], x: Tensor[float]): Tensor[T] =
+  let dist = distanceMatrix(rbf.points, x)
+  let A = rbf.f(dist, rbf.epsilon)
+  result = A * rbf.coeffs
+
+when isMainModule:
+  let x1 = @[@[0.0, 0.0, 0.0], @[1.0, 1.0, 0.0], @[1.0, 2.0, 0.0]].toTensor
+  let x2 = @[@[0.0, 0.0, 1.0], @[1.0, 1.0, 2.0], @[1.0, 2.0, 3.0]].toTensor
+  echo distanceMatrix(x1, x1)
+  let values = @[0.0, 1.0, 2.0].toTensor
+  echo newRbf(x1, values, epsilon=10)
+
+  import benchy
+
+  let pos = randomTensor(5000, 3, 1.0)
+  let vals = randomTensor(5000, 1, 1.0)
+  timeIt "Rbf":
+    keep newRbf(pos, vals)
 

tests/test_interpolate.nim

@@ -1,8 +1,7 @@
-import unittest, math, sequtils
+import unittest, math, sequtils, random
 import numericalnim
 import arraymancer
 
-
 proc f(x: float): float = sin(x)
 proc df(x: float): float = cos(x)
 let t = linspace(0.0, 10.0, 100)
@@ -410,4 +409,18 @@ test "Trilinear f = x*y*z T: Tensor[float]":
            for k in z:
                check abs(spline.eval(i, j, k)[0] - i*j*k) < 1e-12
                check abs(spline.eval(i, j, k)[1] - i*j*k) < 1e-12
-               check abs(spline.eval(i, j, k)[2] - 1) < 1e-16
+               check abs(spline.eval(i, j, k)[2] - 1) < 1e-16
+
+test "rbf f=x*y*z":
+    randomize(1337)
+    let pos = randomTensor[float](100, 3, 1.0)
+    let vals = pos[_, 0] *. pos[_, 1] *. pos[_, 2]
+    let rbfObj = newRbf(pos, vals)
+
+    # We want test points in the interior to avoid the edges
+    let xTest = meshgrid(arraymancer.linspace(0.1, 0.9, 10), arraymancer.linspace(0.1, 0.9, 10), arraymancer.linspace(0.1, 0.9, 10))
+    let yTest = rbfObj.eval(xTest)
+    let yCorrect = xTest[_, 0] *. xTest[_, 1] *. xTest[_, 2]
+    for x in abs(yCorrect - yTest):
+        check x < 0.16
+    check mean_squared_error(yTest, yCorrect) < 2e-4