lots of changes that has been sitting locally for a month now

Author: HugoGranstrom

src/numericalnim/rbf.nim

@@ -27,19 +27,25 @@ type
 template km(point: Tensor[float], index: int, delta: float): int =
   int(ceil(point[0, index] / delta))
 
-iterator neighbours*[T](grid: RbfGrid[T], k: int): int =
+iterator neighbours*[T](grid: RbfGrid[T], k: int, searchLevels: int = 1): int =
   # TODO: Create product iterator that doesn't need to allocate 3^gridDim seqs
-  for dir in product(@[@[-1, 0, 1]].cycle(grid.gridDim)):
+  let directions = @[toSeq(-searchLevels .. searchLevels)].cycle(grid.gridDim)
+  for dir in product(directions):
     block loopBody:
       var kNeigh = k
       for i, x in dir:
         let step = grid.gridSize ^ (grid.gridDim - i - 1)
-        if (k div step) mod grid.gridSize == 0 and x == -1:
+        for level in 1 .. searchLevels:
+          if (k div step) mod grid.gridSize == level - 1 and x <= -level:
+            break loopBody
+          elif (k div step) mod grid.gridSize == grid.gridSize - level and x >= level:
+            break loopBody
+        #[ if ((k div step) mod grid.gridSize == 0 and x == -1) or ((k div step) mod grid.gridSize == 1 and x == -2):
           break loopBody
         elif (k div step) mod grid.gridSize == grid.gridSize - 1 and x == 1:
           break loopBody
-        else:
-          kNeigh += x * step
+        else: ]#
+        kNeigh += x * step
       if kNeigh >= 0 and kNeigh < grid.gridSize ^ grid.gridDim:
         yield kNeigh
 
@@ -67,11 +73,23 @@ template dist2(p1, p2: Tensor[float]): float =
 proc findAllWithin*[T](grid: RbfGrid[T], x: Tensor[float], rho: float): seq[int] =
   assert x.shape.len == 2 and x.shape[0] == 1
   let index = grid.findIndex(x)
-  for k in grid.neighbours(index):
+  let searchLevels = (rho / grid.gridDelta).ceil.int
+  for k in grid.neighbours(index, searchLevels):
     for i in grid.indices[k]:
       if dist2(x, grid.points[i, _]) <= rho*rho:
         result.add i
 
+proc findAllBetween*[T](grid: RbfGrid[T], x: Tensor[float], rho1, rho2: float): seq[int] =
+  assert x.shape.len == 2 and x.shape[0] == 1
+  assert rho2 > rho1
+  let index = grid.findIndex(x)
+  let searchLevels = (rho2 / grid.gridDelta).ceil.int
+  for k in grid.neighbours(index, searchLevels):
+    for i in grid.indices[k]:
+      let d = dist2(x, grid.points[i, _])
+      if rho1*rho1 <= d and d <= rho2*rho2:
+        result.add i
+
 proc newRbfGrid*[T](points: Tensor[float], values: Tensor[T], gridSize: int = 0): RbfGrid[T] =
   let nPoints = points.shape[0]
   let nDims = points.shape[1]
@@ -108,7 +126,10 @@ template compactRbfFuncScalar*(r: float, epsilon: float): float =
 
 proc compactRbfFunc*(r: Tensor[float], epsilon: float): Tensor[float] =
   result = map_inline(r):
-    (1 - x/epsilon) ^ 4 * (4*x/epsilon + 1) * float(x < epsilon)
+    let xeps = x / epsilon
+    let temp = (1 - xeps)
+    let temp2 = temp * temp
+    temp2*temp2 * (4*xeps + 1) * float(xeps < 1)
 
 proc newRbf*[T](points: Tensor[float], values: Tensor[T], rbfFunc: RbfFunc = compactRbfFunc, epsilon: float = 1): RbfType[T] =
   assert points.shape[0] == values.shape[0]
@@ -150,7 +171,8 @@ proc newRbfPu*[T](points: Tensor[float], values: Tensor[T], gridSize: int = 0, r
 
 proc eval*[T](rbf: RbfPUType[T], x: Tensor[float]): Tensor[T] =
   assert x.shape.len == 2
-  assert (not ((x <=. rbf.limits.upper) and (x >=. rbf.limits.lower))).astype(int).sum() == 0, "Some of your points are outside the allowed limits" 
+  assert (not ((x <=. rbf.limits.upper) and (x >=. rbf.limits.lower))).astype(int).sum() == 0, "Some of your points are outside the allowed limits"
+
   let nPoints = x.shape[0] 
   let x = x.scalePoint(rbf.limits)
   result = newTensor[T](nPoints, rbf.nValues)
@@ -170,6 +192,32 @@ proc eval*[T](rbf: RbfPUType[T], x: Tensor[float]): Tensor[T] =
     else:
       result[row, _] = T(Nan) # allow to pass default value to newRbfPU?
 
+proc evalAlt*[T](rbf: RbfPUType[T], x: Tensor[float]): Tensor[T] =
+  assert x.shape.len == 2
+  assert (not ((x <=. rbf.limits.upper) and (x >=. rbf.limits.lower))).astype(int).sum() == 0, "Some of your points are outside the allowed limits"
+
+  let nPoints = x.shape[0] 
+  let x = x.scalePoint(rbf.limits)
+  result = newTensor[T](nPoints, rbf.nValues)
+  var c = newTensor[float](nPoints, 1)
+  var isSet = newTensor[bool](nPoints, 1)
+  let nPatches = rbf.grid.points.shape[0]
+  let pointGrid = newRbfGrid(x, x, rbf.grid.gridSize)
+  for row in 0 ..< nPatches:
+    let center = rbf.grid.points[row, _]
+    let indices = pointGrid.findAllWithin(center, rbf.grid.gridDelta)
+    if indices.len > 0:
+      let vals = rbf.grid.values[row, 0].eval(x[indices, _])
+      for i, index in indices:
+        let r = sqrt(dist2(center, x[index, _]))
+        let ci = compactRbfFuncScalar(r, rbf.grid.gridDelta)
+        result[index, _] = result[index, _] + ci * vals[i, _]
+        c[index] += ci
+        isSet[index, 0] = true
+
+  result /.= c
+  result[not isSet, _] = T(NaN)
+
 when isMainModule:
   let x1 = @[@[0.01, 0.01, 0.01], @[1.0, 1.0, 0.0], @[1.0, 2.0, 4.0]].toTensor
   let x2 = @[@[0.0, 0.0, 1.0], @[1.0, 1.0, 2.0], @[1.0, 2.0, 3.0]].toTensor
@@ -179,18 +227,32 @@ when isMainModule:
 
   import benchy
 
-  #let pos = randomTensor(5000, 3, 1.0)
-  #let vals = randomTensor(5000, 1, 1.0)
-  # timeIt "Rbf":
-  #  keep newRbf(pos, vals)
+  var pos = randomTensor(100000, 3, 1.0)
+  pos[0, _] = [[1.0, 1.0, 1.0]]
+  pos[1, _] = [[0.0, 0.0, 0.0]]
+  let vals = randomTensor(100000, 1, 1.0)
+  let evalPos = randomTensor(100000, 3, 0.9)
+  #let rb = newRbf(pos, vals)
+  #let rbPU = newRbfPu(pos, vals)
+  #timeIt "RbfPU eval":
+  #  keep rbPU.eval(evalPos)
+  #timeIt "Rbf eval":
+  #  keep rb.eval(evalPos)
+  
   #let rbfPu = newRbfPu(x1, values, 3)
   #echo rbfPu.grid.values[1, 0]
   #echo rbfPu.eval(x1[[2, 1, 0], _])
 
   #echo rbfPu.eval(sqrt x1)
   echo "----------------"
-  #[ let xGrid = [[0.1, 0.1], [0.2, 0.3], [0.9, 0.9], [0.4, 0.4]].toTensor
+  for y in countdown(4, 0):
+    var row = ""
+    for x in 0 .. 4:
+      row &= $(x*5 + y) & "\t"
+    echo row
+
+  let xGrid = [[0.1, 0.1], [0.2, 0.3], [0.9, 0.9], [0.4, 0.4]].toTensor
   let valuesGrid = @[0, 1, 9, 5].toTensor.reshape(4, 1)
-  let grid = newRbfGrid(xGrid, valuesGrid, 2)
-  echo grid 
-  echo grid.neighbours(3).toSeq ]#
+  let grid = newRbfGrid(xGrid, valuesGrid, 5)
+  #echo grid 
+  echo grid.neighbours(7, 1).toSeq