From 4b5f216ca9e1890c40bbf8053cfed94e3961849f Mon Sep 17 00:00:00 2001 From: Michael Orlitzky Date: Fri, 7 Feb 2014 10:19:53 -0500 Subject: [PATCH] Add the 'ifoldl2' function to Linear.Matrix. --- src/Linear/Matrix.hs | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/src/Linear/Matrix.hs b/src/Linear/Matrix.hs index dd89a9f..6603302 100644 --- a/src/Linear/Matrix.hs +++ b/src/Linear/Matrix.hs @@ -37,6 +37,7 @@ import qualified Data.Vector.Fixed as V ( and, fromList, head, + ifoldl, length, map, maximum, @@ -823,3 +824,33 @@ matmap f (Mat rows) = Mat $ V.map g rows where g = V.map f + + +-- | Fold over the entire matrix passing the coordinates @i@ and @j@ +-- (of the row/column) to the accumulation function. +-- +-- Examples: +-- +-- >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int +-- >>> ifoldl2 (\i j cur _ -> cur + i + j) 0 m +-- 18 +-- +ifoldl2 :: forall a b m n. + (Int -> Int -> b -> a -> b) + -> b + -> Mat m n a + -> b +ifoldl2 f initial (Mat rows) = + V.ifoldl row_function initial rows + where + -- | The order that we need this in (so that @g idx@ makes sense) + -- is a little funny. So that we don't need to pass weird + -- functions into ifoldl2, we swap the second and third + -- arguments of @f@ calling the result @g@. + g :: Int -> b -> Int -> a -> b + g w x y = f w y x + + row_function :: b -> Int -> Vec n a -> b + row_function rowinit idx r = V.ifoldl (g idx) rowinit r + + -- 2.43.2