Optimize LinearAlgebra.backSubstitute

Author: muehlhaus

src/FsMath/Algebra/LinearAlgebra.fs

@@ -84,37 +84,115 @@ type LinearAlgebra =
 
         Matrix(m, n, qData), r
 
-    /// <summary>Back substitute to solve R * x = y</summary>
-    /// <remarks>R is upper triangular</remarks>
+    /// <summary>Forward substitute to solve L * x = y</summary>
+    /// <remarks>L is lower triangular</remarks>
+    static member inline forwardSubstitute<'T when 'T :> Numerics.INumber<'T>
+                 and 'T : (new: unit -> 'T)
+                 and 'T : struct
+                 and 'T : comparison
+                 and 'T :> ValueType>
+        (L : Matrix<'T>)
+        (y : Vector<'T>) : Vector<'T> =
+
+        let n = L.NumRows
+
+        if L.NumCols <> n || y.Length <> n then
+            invalidArg "dimensions" "L must be square and match the length of y"
+
+        let x     = Array.zeroCreate<'T> n
+        let cols  = L.NumCols
+        let lData = L.Data
+
+        // Again, scalar version; easy to SIMD the inner sum later
+        for i = 0 to n - 1 do
+            let mutable s = y.[i]
+            let rowOffset = i * cols
+            for j = 0 to i - 1 do
+                s <- s - lData.[rowOffset + j] * x.[j]
+            let diag = lData.[rowOffset + i]
+            if diag = 'T.Zero then
+                invalidArg $"Matrix[{i},{i}]" "Diagonal element is zero. Cannot divide."
+            x.[i] <- s / diag
+
+        x
+
+
+
+    ///// <summary>Back substitute to solve R * x = y</summary>
+    ///// <remarks>R is upper triangular</remarks>
+    //static member inline backSubstitute<'T when 'T :> Numerics.INumber<'T>
+    //            and 'T : (new: unit -> 'T)
+    //            and 'T : struct
+    //            and 'T : comparison
+    //            and 'T :> ValueType>
+    //    (r: Matrix<'T>) 
+    //    (y: Vector<'T>) : Vector<'T> =
+
+    //    let n = r.NumRows
+
+    //    if r.NumCols <> n || y.Length <> n then
+    //        invalidArg "dimensions" "R must be square and match the length of y"
+
+    //    let x = Array.zeroCreate<'T> n
+
+    //    for i = n - 1 downto 0 do
+    //        let mutable sum = y.[i]
+    //        for j = i + 1 to n - 1 do
+    //            sum <- sum - r.[i, j] * x.[j]
+    //        let diag = r.[i, i]
+    //        if diag = 'T.Zero then
+    //            invalidArg $"Matrix{i},{i}]" "Diagonal element is zero. Cannot divide."            
+    //        x.[i] <- sum / diag
+
+    //    x
+
+
     static member inline backSubstitute<'T when 'T :> Numerics.INumber<'T>
                 and 'T : (new: unit -> 'T)
                 and 'T : struct
                 and 'T : comparison
                 and 'T :> ValueType>
-        (r: Matrix<'T>) 
+        (R: Matrix<'T>) 
         (y: Vector<'T>) : Vector<'T> =
 
-        let n = r.NumRows
+        let n = R.NumRows
 
-        if r.NumCols <> n || y.Length <> n then
+        if R.NumCols <> n || y.Length <> n then
             invalidArg "dimensions" "R must be square and match the length of y"
 
         let x = Array.zeroCreate<'T> n
+        let cols = R.NumCols
+        let rData = R.Data  // row-major underlying array
 
+        // Backward substitution
         for i = n - 1 downto 0 do
             let mutable sum = y.[i]
-            for j = i + 1 to n - 1 do
-                sum <- sum - r.[i, j] * x.[j]
-            let diag = r.[i, i]
+
+            let startJ = i + 1
+            let len    = n - startJ
+
+            if len > 0 then
+                // row slice: r[i, i+1 .. n-1]
+                let rowOffset   = i * cols + startJ
+                let rowTailSpan = ReadOnlySpan<'T>(rData, rowOffset, len)
+
+                // x slice: x[i+1 .. n-1]
+                let xTailSpan   = ReadOnlySpan<'T>(x, startJ, len)
+
+                // subtract SIMD dot product
+                let dot = SpanMath.dot(rowTailSpan, xTailSpan)
+                sum <- sum - dot
+
+            let diag = R.[i, i]
             if diag = 'T.Zero then
-                invalidArg $"r[{i},{i}]" "Diagonal element is zero. Cannot divide."            
+                invalidArg $"Matrix[{i},{i}]" "Diagonal element is zero. Cannot divide."
+
             x.[i] <- sum / diag
 
         x
 
 
 
-
     /// Solve A * x = b for x, where A is a square matrix (n×n) and b is a vector (length n).
     static member inline solveLinearQR<'T when 'T :> Numerics.INumber<'T>
                 and 'T : (new: unit -> 'T)
@@ -249,10 +327,56 @@ type LinearAlgebra =
 
 
 
-    /// Solve K * x = v (triangular system) in-place, returning a copy of x.
+    ///// Solve K * x = v (triangular system) in-place, returning a copy of x.
+    ///// K must be n×n, v must be length n. 
+    ///// isLower = true => forward substitution
+    ///// isLower = false => backward substitution
+    //static member inline solveTriangularLinearSystem
+    //    (K       : Matrix<'T>)
+    //    (v       : Vector<'T>)
+    //    (isLower : bool)
+    //    : Vector<'T> =
+
+    //    let nK, mK = K.NumRows, K.NumCols
+    //    let nV = v.Length
+    //    if nK <> mK || nV <> nK then
+    //        invalidArg (nameof K) "K must be square, and v must match its dimension."
+
+    //    let x = Array.copy v
+    //    let Kdata = K.Data  // row-major flattened
+
+    //    // Forward or backward substitution
+    //    if isLower then
+    //        // For i in [0..n-1]:
+    //        //   x[i] <- ( x[i] - sum_{j=0..i-1}(K[i,j] * x[j]) ) / K[i,i]
+    //        for i = 0 to nK - 1 do
+    //            let mutable s = x.[i]
+    //            let rowOffset = i * nK
+    //            for j = 0 to i - 1 do
+    //                s <- s - (Kdata.[rowOffset + j] * x.[j])
+    //            let diag = Kdata.[rowOffset + i]
+    //            if diag = 'T.Zero then
+    //                invalidArg $"K[{i},{i}]" "Diagonal element is zero. Cannot divide."
+    //            x.[i] <- s / diag
+    //    else
+    //        // For i in [n-1..downto..0]:
+    //        //   x[i] <- ( x[i] - sum_{j=i+1..n-1}(K[i,j] * x[j]) ) / K[i,i]
+    //        for i = nK - 1 downto 0 do
+    //            let mutable s = x.[i]
+    //            let rowOffset = i * nK
+    //            for j = i + 1 to nK - 1 do
+    //                s <- s - (Kdata.[rowOffset + j] * x.[j])
+    //            let diag = Kdata.[rowOffset + i]
+    //            if diag = 'T.Zero then
+    //                invalidArg $"K[{i},{i}]" "Diagonal element is zero. Cannot divide."
+    //            x.[i] <- s / diag
+
+    //    x
+
+    /// Solve K * x = v (triangular system), returning a new x.
     /// K must be n×n, v must be length n. 
-    /// isLower = true => forward substitution
-    /// isLower = false => backward substitution
+    /// isLower = true  => forward substitution (K lower triangular)
+    /// isLower = false => backward substitution (K upper triangular)
     static member inline solveTriangularLinearSystem
         (K       : Matrix<'T>)
         (v       : Vector<'T>)
@@ -264,36 +388,12 @@ type LinearAlgebra =
         if nK <> mK || nV <> nK then
             invalidArg (nameof K) "K must be square, and v must match its dimension."
 
-        let x = Array.copy v
-        let Kdata = K.Data  // row-major flattened
-
-        // Forward or backward substitution
         if isLower then
-            // For i in [0..n-1]:
-            //   x[i] <- ( x[i] - sum_{j=0..i-1}(K[i,j] * x[j]) ) / K[i,i]
-            for i = 0 to nK - 1 do
-                let mutable s = x.[i]
-                let rowOffset = i * nK
-                for j = 0 to i - 1 do
-                    s <- s - (Kdata.[rowOffset + j] * x.[j])
-                let diag = Kdata.[rowOffset + i]
-                if diag = 'T.Zero then
-                    invalidArg $"K[{i},{i}]" "Diagonal element is zero. Cannot divide."
-                x.[i] <- s / diag
+            // L * x = v
+            LinearAlgebra.forwardSubstitute K v
         else
-            // For i in [n-1..downto..0]:
-            //   x[i] <- ( x[i] - sum_{j=i+1..n-1}(K[i,j] * x[j]) ) / K[i,i]
-            for i = nK - 1 downto 0 do
-                let mutable s = x.[i]
-                let rowOffset = i * nK
-                for j = i + 1 to nK - 1 do
-                    s <- s - (Kdata.[rowOffset + j] * x.[j])
-                let diag = Kdata.[rowOffset + i]
-                if diag = 'T.Zero then
-                    invalidArg $"K[{i},{i}]" "Diagonal element is zero. Cannot divide."
-                x.[i] <- s / diag
-
-        x
+            // R * x = v
+            LinearAlgebra.backSubstitute K v
 
 
 

tests/FsMath.Tests/LinearAlgebraErrorTestsAdditional.fs

@@ -62,7 +62,7 @@ module SolveTriangularSystemZeroDiagonalTests =
         let ex = Assert.Throws<ArgumentException>(fun () ->
             LinearAlgebra.solveTriangularLinearSystem L v true |> ignore)
         Assert.Contains("Diagonal element is zero", ex.Message)
-        Assert.Contains("K[0,0]", ex.Message)
+        Assert.Contains("Matrix[0,0]", ex.Message)
 
     [<Fact>]
     let ``solveTriangularLinearSystem throws on zero diagonal - forward sub in middle`` () =
@@ -76,7 +76,7 @@ module SolveTriangularSystemZeroDiagonalTests =
         let ex = Assert.Throws<ArgumentException>(fun () ->
             LinearAlgebra.solveTriangularLinearSystem L v true |> ignore)
         Assert.Contains("Diagonal element is zero", ex.Message)
-        Assert.Contains("K[2,2]", ex.Message)
+        Assert.Contains("Matrix[2,2]", ex.Message)
 
     [<Fact>]
     let ``solveTriangularLinearSystem throws on zero diagonal - backward sub at position 0`` () =
@@ -89,7 +89,7 @@ module SolveTriangularSystemZeroDiagonalTests =
         let ex = Assert.Throws<ArgumentException>(fun () ->
             LinearAlgebra.solveTriangularLinearSystem U v false |> ignore)
         Assert.Contains("Diagonal element is zero", ex.Message)
-        Assert.Contains("K[0,0]", ex.Message)
+        Assert.Contains("Matrix[0,0]", ex.Message)
 
     [<Fact>]
     let ``solveTriangularLinearSystem throws on zero diagonal - backward sub in middle`` () =
@@ -103,7 +103,7 @@ module SolveTriangularSystemZeroDiagonalTests =
         let ex = Assert.Throws<ArgumentException>(fun () ->
             LinearAlgebra.solveTriangularLinearSystem U v false |> ignore)
         Assert.Contains("Diagonal element is zero", ex.Message)
-        Assert.Contains("K[1,1]", ex.Message)
+        Assert.Contains("Matrix[1,1]", ex.Message)
 
     [<Fact>]
     let ``solveTriangularLinearSystem throws on zero diagonal - backward sub at last`` () =
@@ -115,7 +115,7 @@ module SolveTriangularSystemZeroDiagonalTests =
         let ex = Assert.Throws<ArgumentException>(fun () ->
             LinearAlgebra.solveTriangularLinearSystem U v false |> ignore)
         Assert.Contains("Diagonal element is zero", ex.Message)
-        Assert.Contains("K[1,1]", ex.Message)
+        Assert.Contains("Matrix[1,1]", ex.Message)
 
 
 module CholeskyNonSquareTests =

tests/FsMath.Tests/VectorOpsCoverageTests.fs

@@ -427,37 +427,3 @@ module VectorOpsCoverageTests =
         let result2 = evalQ <@ (scalar .* v1) .+ (scalar .* v2) @>
         floatArrayClose result1 result2 1e-10
 
-    // ========================================
-    // Quotation tests for @ operator
-    // Note: @ operator currently calls Power.Invoke (not Dot.Invoke despite the comment)
-    // ========================================
-
-    [<Fact>]
-    let ``operator @_Q: applies power operation`` () =
-        let v = [| 2.0; 3.0; 4.0 |]
-        let result = evalQ <@ v @ 2.0 @>
-        floatArrayClose [| 4.0; 9.0; 16.0 |] result 1e-10
-
-    [<Fact>]
-    let ``operator @_Q: fractional power`` () =
-        let v = [| 4.0; 9.0; 16.0 |]
-        let result = evalQ <@ v @ 0.5 @>
-        floatArrayClose [| 2.0; 3.0; 4.0 |] result 1e-10
-
-    [<Fact>]
-    let ``operator @_Q: negative power`` () =
-        let v = [| 2.0; 4.0; 5.0 |]
-        let result = evalQ <@ v @ -1.0 @>
-        floatArrayClose [| 0.5; 0.25; 0.2 |] result 1e-10
-
-    [<Fact>]
-    let ``operator @_Q: power of zero`` () =
-        let v = [| 2.0; 3.0; 4.0 |]
-        let result = evalQ <@ v @ 0.0 @>
-        floatArrayClose [| 1.0; 1.0; 1.0 |] result 1e-10
-
-    [<Fact>]
-    let ``operator @_Q: integer power`` () =
-        let v = [| 2.0; 3.0; 4.0 |]
-        let result = evalQ <@ v @ 3.0 @>
-        floatArrayClose [| 8.0; 27.0; 64.0 |] result 1e-10

tests/FsMath.Tests/VectorOpsTests.fs

@@ -302,48 +302,3 @@ module VectorOpsTests =
         let result2 = (scalar .* v1) .+ (scalar .* v2)
         floatArrayClose result1 result2 1e-10
 
-    // =============================================
-    // @ Operator Tests (Power operator - NOTE: Comment in VectorOps.fs incorrectly says "Dot product")
-    // =============================================
-
-    [<Fact>]
-    let ``@ operator applies power operation (float)`` () =
-        // The @ operator currently calls Power.Invoke, not Dot.Invoke
-        // despite the comment saying "// Dot product ( @ )"
-        let v = [| 2.0; 3.0; 4.0 |]
-        let power = 2.0
-        let result = v @ power
-        let expected = [| 4.0; 9.0; 16.0 |]
-        floatArrayClose expected result 1e-10
-
-    [<Fact>]
-    let ``@ operator with fractional power (float)`` () =
-        let v = [| 4.0; 9.0; 16.0 |]
-        let power = 0.5
-        let result = v @ power
-        let expected = [| 2.0; 3.0; 4.0 |]
-        floatArrayClose expected result 1e-10
-
-    [<Fact>]
-    let ``@ operator with negative power (float)`` () =
-        let v = [| 2.0; 4.0; 5.0 |]
-        let power = -1.0
-        let result = v @ power
-        let expected = [| 0.5; 0.25; 0.2 |]
-        floatArrayClose expected result 1e-10
-
-    [<Fact>]
-    let ``@ operator with zero power returns ones (float)`` () =
-        let v = [| 2.0; 3.0; 4.0 |]
-        let power = 0.0
-        let result = v @ power
-        let expected = [| 1.0; 1.0; 1.0 |]
-        floatArrayClose expected result 1e-10
-
-    [<Fact>]
-    let ``@ operator with integer power (float)`` () =
-        let v = [| 2.0; 3.0; 4.0 |]
-        let power = 3.0
-        let result = v @ power
-        let expected = [| 8.0; 27.0; 64.0 |]
-        floatArrayClose expected result 1e-10