math.matrices: Add a diagonal-matrix word that makes an nxn matrix from a sequence.

Optimize identity-matrix word -- 2-3x faster using diagonal-matrix.
Add an eye word that makes an mxn diagonal matrix with a diagonal offset by k.

diff --git a/basis/math/matrices/matrices-tests.factor b/basis/math/matrices/matrices-tests.factor
index 3996a475ba37110ca58f01e3017fbfae2fbdd45b..4e3cf066673ae376c071e87cf9219ce86c12f8c7 100644 (file)
--- a/basis/math/matrices/matrices-tests.factor
+++ b/basis/math/matrices/matrices-tests.factor
@@ -7,14 +7,84 @@ USING: math.matrices math.vectors tools.test math ;
     3 1 zero-matrix
 ] unit-test
 
-[
+{
     { { 1 0 0 }
        { 0 1 0 }
        { 0 0 1 } }
-] [
+} [
     3 identity-matrix
 ] unit-test
 
+{
+    { { 1 0 0 }
+       { 0 2 0 }
+       { 0 0 3 } }
+} [
+    { 1 2 3 } diagonal-matrix
+] unit-test
+
+{
+    {
+        { 1 0 0 }
+        { 0 1 0 }
+        { 0 0 1 }
+    }
+} [
+    3 3 0 eye
+] unit-test
+
+{
+    {
+        { 0 1 0 }
+        { 0 0 1 }
+        { 0 0 0 }
+    }
+} [
+    3 3 1 eye
+] unit-test
+
+{
+    {
+        { 0 0 0 }
+        { 1 0 0 }
+        { 0 1 0 }
+    }
+} [
+    3 3 -1 eye
+] unit-test
+
+{
+    {
+        { 1 0 0 0 }
+        { 0 1 0 0 }
+        { 0 0 1 0 }
+    }
+} [
+    3 4 0 eye
+] unit-test
+
+{
+    {
+        { 0 1 0 }
+        { 0 0 1 }
+        { 0 0 0 }
+        { 0 0 0 }
+    }
+} [
+    4 3 1 eye
+] unit-test
+
+{
+    {
+        { 0 0 0 }
+        { 1 0 0 }
+        { 0 1 0 }
+        { 0 0 1 }
+    }
+} [
+    4 3 -1 eye
+] unit-test
+
 [
     { { 1 0 4 }
        { 0 7 0 }

diff --git a/basis/math/matrices/matrices.factor b/basis/math/matrices/matrices.factor
index 9fc4f879e82c7fb3bc2857ef8af0a59936db85ff..763b133cf3547e7042aa622dd14a580abd6552a4 100644 (file)
--- a/basis/math/matrices/matrices.factor
+++ b/basis/math/matrices/matrices.factor
@@ -9,9 +9,15 @@ IN: math.matrices
 : zero-matrix ( m n -- matrix )
     '[ _ 0 <array> ] replicate ;
 
+: diagonal-matrix ( diagonal-seq -- matrix )
+    dup length dup zero-matrix
+    [ '[ dup _ nth set-nth ] each-index ] keep ; inline
+
 : identity-matrix ( n -- matrix )
-    #! Make a nxn identity matrix.
-    iota dup [ = 1 0 ? ] cartesian-map ;
+    1 <repetition> diagonal-matrix ; inline
+
+: eye ( m n k -- matrix )
+    [ [ iota ] bi@ ] dip neg '[ _ + = 1 0 ? ] cartesian-map ;
 
 :: rotation-matrix3 ( axis theta -- matrix )
     theta cos :> c