Convert the zoom function to use Values3D.

[spline3.git] / src / Grid.hs

diff --git a/src/Grid.hs b/src/Grid.hs
index 1a01ab7eb60bba873e31e087b645eaf8d83a3689..3d636b53a780f465b515c6e45f8ee53b6e7f10ba 100644 (file)
--- a/src/Grid.hs
+++ b/src/Grid.hs
@@ -12,7 +12,9 @@ import Misc (flatten)
 import Point (Point)
 import Tetrahedron (polynomial)
 import ThreeDimensional (contains_point)
-import Values (Values3D, dims, empty3d)
+import Values (Values3D, dims, empty3d, zoom_shape)
+
+import qualified Data.Array.Repa as R
 
 -- | Our problem is defined on a Grid. The grid size is given by the
 --   positive number h. The function values are the values of the
@@ -81,28 +83,24 @@ find_containing_cubes g p =
       contains_our_point = flip contains_point p
 
 
-
-zoom :: Grid -> Int -> [[[Double]]]
+zoom :: Grid -> Int -> Values3D
 zoom g scale_factor
-    | xsize == 0 || ysize == 0 || zsize == 0 = []
+    | xsize == 0 || ysize == 0 || zsize == 0 = empty3d
     | otherwise =
-        [[[f p | i <- [0..scaled_zsize],
-                 let i' = scale_dimension i,
-                 let j' = scale_dimension j,
-                 let k' = scale_dimension k,
-                 let p = (i', j', k') :: Point,
-                 let c = (find_containing_cubes g p) !! 0,
-                 let t = (find_containing_tetrahedra c p) !! 0,
-                 let f = polynomial t]
-                 | j <- [0..scaled_ysize]]
-                   | k <- [0..scaled_xsize]]
-  where
-    scale_dimension :: Int -> Double
-    scale_dimension x = (fromIntegral x) / (fromIntegral scale_factor)
-
-    fvs = function_values g
-    (xsize, ysize, zsize) = dims fvs
-    scaled_xsize = xsize * scale_factor
-    scaled_ysize = ysize * scale_factor
-    scaled_zsize = zsize * scale_factor
-
+        R.traverse arr transExtent (\_ -> newlookup)
+          where
+            fvs = function_values g
+            (xsize, ysize, zsize) = dims fvs
+            arr = fvs
+            transExtent = zoom_shape scale_factor
+            newlookup :: R.DIM3 -> Double
+            newlookup (R.Z R.:. i R.:. j R.:. k) =
+                f p
+                  where
+                    i' = fromIntegral i
+                    j' = fromIntegral j
+                    k' = fromIntegral k
+                    p = (i', j', k') :: Point
+                    c = head (find_containing_cubes g p)
+                    t = head (find_containing_tetrahedra c p)
+                    f = polynomial t