Add Linear.Matrix.set_idx.

[numerical-analysis.git] / src / Linear / Matrix.hs

diff --git a/src/Linear/Matrix.hs b/src/Linear/Matrix.hs
index 3d4daab3dc1f7be81277238b52253c5df54ae61f..4d1dceb417c07e63992ec7afefd981a23d484c39 100644 (file)
--- a/src/Linear/Matrix.hs
+++ b/src/Linear/Matrix.hs
@@ -662,11 +662,16 @@ vec5d (v,w,x,y,z) = Mat (mk5 (mk1 v) (mk1 w) (mk1 x) (mk1 y) (mk1 z))
 scalar :: a -> Mat1 a
 scalar x = Mat (mk1 (mk1 x))
 
-dot :: (RealRing.C a, m ~ S t, Arity t)
-       => Col m a
-       -> Col m a
+-- Get the scalar value out of a 1x1 matrix.
+unscalar :: Mat1 a -> a
+unscalar (Mat rows) = V.head $ V.head rows
+
+
+dot :: (Ring.C a, Arity m)
+       => Col (S m) a
+       -> Col (S m) a
        -> a
-v1 `dot` v2 = ((transpose v1) * v2) !!! (0, 0)
+v1 `dot` v2 = unscalar $ ((transpose v1) * v2)
 
 
 -- | The angle between @v1@ and @v2@ in Euclidean space.
@@ -969,3 +974,25 @@ reverse2 :: (Arity m, Arity n) => Mat m n a -> Mat m n a
 reverse2 (Mat rows) = Mat $ V.reverse $ V.map V.reverse rows
 
 
+-- | Unsafely set the (i,j) element of the given matrix.
+--
+--   Examples:
+--
+--   >>> let m = fromList [[1,2,3],[4,5,6],[7,8,9]] :: Mat3 Int
+--   >>> set_idx m (1,1) 17
+--   ((1,2,3),(4,17,6),(7,8,9))
+--
+set_idx :: forall m n a.
+           (Arity m, Arity n)
+        => Mat m n a
+        -> (Int, Int)
+        -> a
+        -> Mat m n a
+set_idx matrix (i,j) newval =
+  imap2 updater matrix
+  where
+    updater :: Int -> Int -> a -> a
+    updater k l existing =
+      if k == i && l == j
+      then newval
+      else existing