diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,48 @@ +# Changelog + +All notable changes to `moonlight-linalg` are documented here. + +## 0.1.0.0 - 2026-08-21 + +- Initial public release. +- Typed dense row validation and matrix/vector combinators, with dense rows + documented as authoring/projection rather than hot dense storage. +- Dynamic dense matrices, affine transforms, vector primitives, exterior helpers, + and finite-field `GF2` support. +- Sparse COO/CSR/CSC/Packed carriers and sparse solver helpers. +- Row-reduction, PLU, RREF, Smith-normal-form, symmetric eigen, and domain-level + matrix operations. +- Arnoldi, Lanczos, restarted Lanczos, block Lanczos, projected subspaces, + tridiagonal and block-tridiagonal projected operators, and spectral request/result + surfaces above Krylov. +- Selected symmetric-tridiagonal/path-Laplacian spectral fast path documented as + the package's flagship hot kernel. +- Finalization campaign: perf storm follow-through, laws sublibrary split, thick + restart hardening, selected certification split, IC(0) sparse preconditioning, + GF2 sparse reducer, closed-form geometry eigen path, `-O2` shared properties, + and IEEE constant `encodeFloat` repair. +- Graded sublibrary decomposition: the private core dissolved into ten public + slices (`carrier` through `native`) with per-slice source directories, the + dependency DAG cabal-enforced, native LAPACK/Accelerate linkage confined to + `moonlight-linalg-native`, and `ArchitectureSpec` slice-discipline guards. + Architecture decisions recorded in `docs/ARCHITECTURE.md`. +- Multimodular Smith engine: `smithDiagonalForm` at `Integer` dispatches (via a + `NOINLINE`-pinned rewrite rule) to a CRT/Iliopoulos engine — word-prime + determinant/rank sweep on unboxed carriers, exact determinant by CRT under the + Hadamard certificate, Smith elimination mod `2·|det|` on tiered flat carriers, + unimodular Hermite compression for rectangular/rank-deficient inputs. Beats + the classical diagonal route ×1.8–3.6 on dense random n=8–32 and removes the + coefficient-explosion scaling wall; verified against the classical route and + FLINT with exact agreement gates. Classical Smith diagonals (both routes) now + canonicalize invariant factors to nonnegative canonical associates, with unit + flips absorbed into the left witnesses. +- Witnessed Smith engine: `smithNormalForm` at `Integer` dispatches square + certified-nonsingular inputs of dimension ≥ 25 to a mod-determinant fast path — + Domich–Kannan–Trotter `[A; R·I]` stack Hermite alternation mod `R = 2·|det|`, + per-phase transform recovery by early-terminated CRT over word primes with + deterministic exact verification, triangular back-substitution solves, and + inverse witnesses by exact diagonal division. Kills the 111k-bit intermediate + explosion (witnesses stay ~100–190 bits at n=32) and runs n=32 dense random + with torsion in 11 ms (was 69 ms); smaller and rectangular/rank-deficient + inputs keep the alternating arena unchanged. +- Hackage-facing metadata, documentation, and explicit package bounds. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Blue Rose + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,102 @@ +# moonlight-linalg + +> Part of **Moonlight**, the sheaf-theoretic computation layer beneath +> [Melusine](https://bluerose.blue) and Pale Meridian. + +Typed dense, sparse, finite-field, and Krylov linear algebra for Pale Meridian's +foundation packages. + +`moonlight-linalg` is Moonlight's numerical linear-algebra tier. Building on +[`moonlight-core`](../moonlight-core) and +[`moonlight-algebra`](../moonlight-algebra), it provides the matrix, vector, +GF(2), sparse-storage, Smith-normal-form, eigen, and Krylov machinery used by +homology, analysis, sheaf, geometry, and solver packages. + +The front door is the umbrella module `Moonlight.LinAlg`, whose header carries +the type-indexed `MoonlightError` contract, the role map across the public +surface, and the quick-start recipe. This page maps the public modules and the +benchmark tooling. + +## Public modules + +| Module | Surface | +| --- | --- | +| `Moonlight.LinAlg` | Broad public surface for dense, sparse, operator, spectral, domain, geometry, statics, and immutable Krylov modules. | +| `Moonlight.LinAlg.Dense` | Dense vectors/matrices, validated dense-row authoring, GF(2), exterior algebra, basic operations, decompositions, field operations, direct solvers, and primitives. | +| `Moonlight.LinAlg.Sparse` | Sparse matrix carriers, packed sparse operators, sealed preconditioner families, and sparse iterative solvers. | +| `Moonlight.LinAlg.Operator` | Abstract affine-normalized linear operators with explicit self-adjoint construction boundaries. | +| `Moonlight.LinAlg.Spectral` | Eigenvalue/eigenpair requests and contiguous result views, dispatched by demand and operator structure above Krylov. | +| `Moonlight.LinAlg.Krylov` | Public Arnoldi/Lanczos decomposition, projected tridiagonal/block-tridiagonal carriers, and block Lanczos surface. | +| `Moonlight.LinAlg.Native` | Effectful native LAPACK backend boundary. On macOS it links Accelerate; elsewhere it expects BLAS/LAPACK libraries. | +| `Moonlight.LinAlg.Domain` | Domain-level algebraic operations, including Smith normal form. | +| `Moonlight.LinAlg.Geometry` | `Vec2`, `Vec3`, AABB/AABB2, frames, affine transforms, and compact symmetric 2D/3D carriers. | +| `Moonlight.LinAlg.Statics` | Statics types, assembly, equilibrium compilation, and support checking. | + +The `Moonlight.LinAlg.Pure.*`, `Moonlight.LinAlg.Internal.*`, and +`Moonlight.LinAlg.Effect.*` leaves live in graded implementation sublibraries +(`carrier`, `structured`, `eigen`, `geometry`, `dense`, `domain`, `sparse`, +`statics`, `spectral`, `native`), with the dependency DAG cabal-enforced and +native linkage confined to `moonlight-linalg-native`. Public callers use +the public modules above; the slice modules define implementation ownership behind that +public vocabulary. + +## Benchmark artifacts + +Repository tooling generates benchmark artifacts: + +```sh +scripts/tooling/generate_moonlight_linalg_bench_artifacts.py +``` + +The generator runs `cabal test moonlight-linalg-test -j1`, then runs the +short default `moonlight-linalg-bench` target with a CSV `tasty-bench` report and +renders SVG artifacts under `/tmp` by default. Use `--output-dir` to choose a +destination. + +The default bench is decomposed across dense-row validation, dense decompositions/solvers, +sparse storage, sparse iterative solvers, domain algebra, GF(2), exterior powers, +geometry/statics, spectral demand dispatch, sparse Krylov, native LAPACK, and +structured projected block eigensolve. Heavier strata stay opt-in: + +- `--broad-medium` / `MOONLIGHT_LINALG_BENCH_ENABLE_BROAD_MEDIUM=1` +- `--broad-large` / `MOONLIGHT_LINALG_BENCH_ENABLE_BROAD_LARGE=1` +- `--sparse-large` / `MOONLIGHT_LINALG_BENCH_ENABLE_SPARSE_LARGE=1` +- `--include-100k` / `MOONLIGHT_LINALG_BENCH_ENABLE_100K=1` +- `--projected-medium` / `MOONLIGHT_LINALG_BENCH_ENABLE_PROJECTED_MEDIUM=1` +- `--large-projected` / `MOONLIGHT_LINALG_BENCH_ENABLE_PROJECTED_LARGE=1` +- `--native-large` / `MOONLIGHT_LINALG_BENCH_ENABLE_NATIVE_LARGE=1` + +Use `--diagnostic-sweep` for the medium broad rows, 50k sparse row, and +144-dimensional projected rows. + +The default native LAPACK group keeps small DSYEV rows and a small DSTEMR +selected-tridiagonal row. The 10k DSTEMR path-Laplacian row is opt-in because it +is a native-boundary stress case for deeper runs. + +For a fast local sanity sweep, skip the default calibrated sampling ceremony: + +```sh +cabal bench moonlight-linalg:moonlight-linalg-bench -j1 --benchmark-options='--once' +``` + +This executes every default benchmark row once through the same workload owners +and reports the slowest rows. Use the calibrated default only when the numbers +are going into evidence. + +Benchmark outputs are generated explicitly for each measurement run. + +## Relationship to external linear-algebra packages + +General-purpose Haskell linear algebra packages are better choices for ordinary +numerical applications. `moonlight-linalg` exists because Pale Meridian needs compact +compiler-local carriers, GF(2) and integer-domain hooks, exact shape/domain failures, +and structured Krylov/projected-operator types that compose with the rest of the +Moonlight foundation stack. Its strongest hot path is selected structured spectra, +especially path-Laplacian/tridiagonal modes. Dense nested rows serve validated +authoring. The native LAPACK boundary is deliberately effectful and isolated from +pure APIs. + +## License + +MIT; see [`LICENSE`](./LICENSE). Third-party attribution is recorded in +[`THIRD_PARTY_NOTICES.md`](./THIRD_PARTY_NOTICES.md). diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 --- /dev/null +++ b/THIRD_PARTY_NOTICES.md @@ -0,0 +1,15 @@ +# Third-party notices + +No third-party source code is vendored or adapted into `moonlight-linalg`. + +The native LAPACK backend links the platform BLAS/LAPACK implementation: + +- macOS: Apple Accelerate framework. +- non-macOS Cabal builds: system `lapack` and `blas` libraries. + +Those libraries' license and copyright information is governed by the platform or +library distribution that supplies them. `moonlight-linalg` does not vendor LAPACK, +ARPACK, PRIMME, SciPy, or Netlib source code. + +The package also depends on Haskell libraries through Cabal. Their license and +copyright information is governed by those packages' own distributions. diff --git a/bench/Main.hs b/bench/Main.hs new file mode 100644 --- /dev/null +++ b/bench/Main.hs @@ -0,0 +1,166 @@ +module Main + ( main, + ) +where + +import Data.Foldable (traverse_) +import Data.List (sortOn) +import Data.Maybe (mapMaybe) +import Data.Ord (Down (..)) +import Env + ( BenchmarkSelection, + benchmarkNotice, + readBenchmarkSelection, + ) +import DenseCore (denseCoreBenchmarks, denseCoreOnceBenchmarks) +import DenseDecomposition (denseDecompositionBenchmarks, denseDecompositionOnceBenchmarks) +import DomainAlgebra (domainAlgebraBenchmarks, domainAlgebraOnceBenchmarks) +import GeometryStatics (geometryStaticsBenchmarks, geometryStaticsOnceBenchmarks) +import NativeLapack (nativeLapackBenchmarks, nativeLapackOnceBenchmarks) +import ProjectedBlock (projectedBlockBenchmarks, projectedBlockOnceBenchmarks) +import SparseKrylov (sparseKrylovBenchmarks, sparseKrylovOnceBenchmarks) +import SparseSolvers (sparseSolverBenchmarks, sparseSolverOnceBenchmarks) +import SpectralDispatch (spectralDispatchBenchmarks, spectralDispatchOnceBenchmarks) +import SparseStorage (sparseStorageBenchmarks, sparseStorageOnceBenchmarks) +import Types + ( OnceBenchmark, + OnceBenchmarkResult (..), + OnceBenchmarkStats (..), + runOnceBenchmark, + ) +import System.Environment (getArgs) +import Test.Tasty.Bench (defaultMain) +import Text.Printf (printf) +import Prelude + +main :: IO () +main = do + benchmarkSelection <- readBenchmarkSelection + putStrLn (benchmarkNotice benchmarkSelection) + args <- getArgs + case args of + ["--once"] -> + runOnceBenchmarks (linalgOnceBenchmarks benchmarkSelection) + _ -> + defaultMain + [ denseCoreBenchmarks benchmarkSelection, + denseDecompositionBenchmarks, + sparseStorageBenchmarks benchmarkSelection, + sparseSolverBenchmarks benchmarkSelection, + spectralDispatchBenchmarks benchmarkSelection, + domainAlgebraBenchmarks, + geometryStaticsBenchmarks benchmarkSelection, + sparseKrylovBenchmarks benchmarkSelection, + nativeLapackBenchmarks benchmarkSelection, + projectedBlockBenchmarks benchmarkSelection + ] + +linalgOnceBenchmarks :: BenchmarkSelection -> [OnceBenchmark] +linalgOnceBenchmarks benchmarkSelection = + denseCoreOnceBenchmarks benchmarkSelection + <> denseDecompositionOnceBenchmarks + <> sparseStorageOnceBenchmarks benchmarkSelection + <> sparseSolverOnceBenchmarks benchmarkSelection + <> spectralDispatchOnceBenchmarks benchmarkSelection + <> domainAlgebraOnceBenchmarks + <> geometryStaticsOnceBenchmarks benchmarkSelection + <> sparseKrylovOnceBenchmarks benchmarkSelection + <> nativeLapackOnceBenchmarks benchmarkSelection + <> projectedBlockOnceBenchmarks benchmarkSelection + +runOnceBenchmarks :: [OnceBenchmark] -> IO () +runOnceBenchmarks benchmarks = do + results <- traverse runOnceBenchmark benchmarks + case sequence results of + Left failureText -> + ioError (userError ("moonlight-linalg once benchmark failed: " <> failureText)) + Right successfulResults -> + reportOnceResults successfulResults + +reportOnceResults :: [OnceBenchmarkResult] -> IO () +reportOnceResults results = do + putStrLn + ( "All " + <> show (length results) + <> " benchmark rows executed once (measured CPU " + <> secondsText totalSeconds + <> "; checksum " + <> printf "%.6f" checksumTotal + <> onceStatsSummaryText results + <> ")." + ) + putStrLn "Slowest once rows:" + traverse_ (putStrLn . renderOnceResult) (take 10 (sortOn (Down . onceResultElapsedSeconds) results)) + case onceResultsWithStats results of + [] -> putStrLn "RTS per-row allocation stats disabled; rerun with +RTS -T." + statsRows -> do + putStrLn "Highest allocation once rows:" + traverse_ (putStrLn . renderAllocatedOnceResult) (take 10 (sortOn (Down . onceAllocatedBytes . snd) statsRows)) + where + totalSeconds = + sum (onceResultElapsedSeconds <$> results) + checksumTotal = + sum (onceResultChecksum <$> results) + +renderOnceResult :: OnceBenchmarkResult -> String +renderOnceResult result = + " " + <> secondsText (onceResultElapsedSeconds result) + <> maybe "" renderInlineOnceStats (onceResultStats result) + <> " " + <> onceResultLabel result + +renderAllocatedOnceResult :: (OnceBenchmarkResult, OnceBenchmarkStats) -> String +renderAllocatedOnceResult (result, statsValue) = + " " + <> bytesText (onceAllocatedBytes statsValue) + <> " allocated; " + <> bytesText (onceLiveBytesAfterMajorGC statsValue) + <> " live after major GC; " + <> bytesText (onceProcessMaximumLiveBytes statsValue) + <> " process maximum residency " + <> onceResultLabel result + +onceStatsSummaryText :: [OnceBenchmarkResult] -> String +onceStatsSummaryText results = + case onceResultsWithStats results of + [] -> "; RTS allocation stats disabled" + statsRows -> + "; allocated " + <> bytesText (sum (onceAllocatedBytes . snd <$> statsRows)) + <> "; max retained live after row GC " + <> bytesText (maximum (0 : (onceLiveBytesAfterMajorGC . snd <$> statsRows))) + <> "; process maximum residency " + <> bytesText (maximum (0 : (onceProcessMaximumLiveBytes . snd <$> statsRows))) + +onceResultsWithStats :: [OnceBenchmarkResult] -> [(OnceBenchmarkResult, OnceBenchmarkStats)] +onceResultsWithStats = + mapMaybe withStats + where + withStats result = + case onceResultStats result of + Just statsValue -> Just (result, statsValue) + Nothing -> Nothing + +renderInlineOnceStats :: OnceBenchmarkStats -> String +renderInlineOnceStats statsValue = + " / alloc " + <> bytesText (onceAllocatedBytes statsValue) + <> " / live " + <> bytesText (onceLiveBytesAfterMajorGC statsValue) + <> " / max " + <> bytesText (onceProcessMaximumLiveBytes statsValue) + +secondsText :: Double -> String +secondsText secondsValue + | secondsValue < 1.0e-6 = printf "%.3f ns" (secondsValue * 1.0e9) + | secondsValue < 1.0e-3 = printf "%.3f us" (secondsValue * 1.0e6) + | secondsValue < 1.0 = printf "%.3f ms" (secondsValue * 1.0e3) + | otherwise = printf "%.3f s" secondsValue + +bytesText :: Integer -> String +bytesText byteCount + | byteCount < 1024 = show byteCount <> " B" + | byteCount < 1024 * 1024 = printf "%.3f KiB" (fromIntegral byteCount / 1024 :: Double) + | byteCount < 1024 * 1024 * 1024 = printf "%.3f MiB" (fromIntegral byteCount / (1024 * 1024) :: Double) + | otherwise = printf "%.3f GiB" (fromIntegral byteCount / (1024 * 1024 * 1024) :: Double) diff --git a/bench/dense/DenseCore.hs b/bench/dense/DenseCore.hs new file mode 100644 --- /dev/null +++ b/bench/dense/DenseCore.hs @@ -0,0 +1,200 @@ +module DenseCore + ( denseCoreBenchmarks, + denseCoreOnceBenchmarks, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Bifunctor (first) +import Data.Vector.Storable qualified as S +import Env (BenchmarkSelection (..)) +import Fixtures + ( denseBenchmarkRows, + denseBenchmarkVector, + ) +import Types + ( BenchmarkWeight (..), + OnceBenchmark (..), + benchmarkWeightEither, + eitherBenchmarkWeight, + ) +import Moonlight.LinAlg.Dense + ( DenseDoubleMatrix, + denseDoubleMatrixToRowMajorVector, + denseDoubleMatrixVectorProduct, + mkDenseDoubleMatrixRowMajor, + ) +import Moonlight.LinAlg.Dense.Primitives + ( matrixVectorProduct, + ) +import Moonlight.LinAlg.Dense.Rows + ( hcatRowsExact, + matrixProductRowsWith, + transposeRowsExact, + vcatRowsExact, + ) +import Moonlight.LinAlg.Native + ( denseDoubleMatrixProductBlas, + ) +import Test.Tasty.Bench (Benchmark, bench, bgroup, nf, nfIO) +import Prelude + +data DenseCoreCase = DenseCoreCase + { denseCoreLabel :: !String, + denseCoreRows :: ![[Double]], + denseCoreVector :: ![Double], + denseCoreFlatVector :: !(S.Vector Double), + denseCoreFlatMatrix :: !(Either String DenseDoubleMatrix) + } + +instance NFData DenseCoreCase where + rnf benchmarkCase = + rnf (denseCoreLabel benchmarkCase) + `seq` rnf (denseCoreRows benchmarkCase) + `seq` rnf (denseCoreVector benchmarkCase) + `seq` rnf (denseCoreFlatVector benchmarkCase) + `seq` forceFlatMatrix (denseCoreFlatMatrix benchmarkCase) + `seq` () + +denseCoreBenchmarks :: BenchmarkSelection -> Benchmark +denseCoreBenchmarks benchmarkSelection = + bgroup + "dense row validation surface" + (denseCoreCaseBenchmarks =<< denseCoreCases benchmarkSelection) + +denseCoreOnceBenchmarks :: BenchmarkSelection -> [OnceBenchmark] +denseCoreOnceBenchmarks benchmarkSelection = + denseCoreCaseOnceBenchmarks =<< denseCoreCases benchmarkSelection + +denseCoreCases :: BenchmarkSelection -> [DenseCoreCase] +denseCoreCases benchmarkSelection = + fmap + denseCoreCase + ( [32] + <> [96 | includeBroadMedium benchmarkSelection || includeBroadLarge benchmarkSelection] + <> [192 | includeBroadLarge benchmarkSelection] + ) + +denseCoreCase :: Int -> DenseCoreCase +denseCoreCase dimension = + let rowValues = denseBenchmarkRows dimension + vectorValues = denseBenchmarkVector dimension + flatPayload = S.fromList (concat rowValues) + in DenseCoreCase + { denseCoreLabel = "n=" <> show dimension, + denseCoreRows = rowValues, + denseCoreVector = vectorValues, + denseCoreFlatVector = S.fromList vectorValues, + denseCoreFlatMatrix = + first show + ( mkDenseDoubleMatrixRowMajor + dimension + dimension + flatPayload + ) + } + +denseCoreCaseBenchmarks :: DenseCoreCase -> [Benchmark] +denseCoreCaseBenchmarks benchmarkCase = + [ bench (denseCoreLabel benchmarkCase <> " matrix-vector reference rows") (nf denseMatvecWeight benchmarkCase), + bench (denseCoreLabel benchmarkCase <> " flat matrix-vector") (nf denseFlatMatvecWeight benchmarkCase), + bench (denseCoreLabel benchmarkCase <> " matrix-product reference rows") (nf denseMatmulWeight benchmarkCase), + bench (denseCoreLabel benchmarkCase <> " native BLAS matrix-product") (nfIO (denseNativeMatmulWeight benchmarkCase)), + bench (denseCoreLabel benchmarkCase <> " transpose") (nf denseTransposeWeight benchmarkCase), + bench (denseCoreLabel benchmarkCase <> " hcat/vcat") (nf denseConcatWeight benchmarkCase) + ] + +denseCoreCaseOnceBenchmarks :: DenseCoreCase -> [OnceBenchmark] +denseCoreCaseOnceBenchmarks benchmarkCase = + [ denseCoreOnceBenchmark benchmarkCase "matrix-vector reference rows" denseMatvecWeight, + denseCoreOnceBenchmark benchmarkCase "flat matrix-vector" denseFlatMatvecWeight, + denseCoreOnceBenchmark benchmarkCase "matrix-product reference rows" denseMatmulWeight, + denseCoreOnceBenchmarkIO benchmarkCase "native BLAS matrix-product" denseNativeMatmulWeight, + denseCoreOnceBenchmark benchmarkCase "transpose" denseTransposeWeight, + denseCoreOnceBenchmark benchmarkCase "hcat/vcat" denseConcatWeight + ] + +denseCoreOnceBenchmark :: DenseCoreCase -> String -> (DenseCoreCase -> BenchmarkWeight) -> OnceBenchmark +denseCoreOnceBenchmark benchmarkCase label measure = + OnceBenchmark + { onceBenchmarkLabel = "dense row validation surface." <> denseCoreLabel benchmarkCase <> " " <> label, + onceBenchmarkAction = pure (benchmarkWeightEither (measure benchmarkCase)) + } + +denseCoreOnceBenchmarkIO :: DenseCoreCase -> String -> (DenseCoreCase -> IO BenchmarkWeight) -> OnceBenchmark +denseCoreOnceBenchmarkIO benchmarkCase label measure = + OnceBenchmark + { onceBenchmarkLabel = "dense row validation surface." <> denseCoreLabel benchmarkCase <> " " <> label, + onceBenchmarkAction = benchmarkWeightEither <$> measure benchmarkCase + } + +denseMatvecWeight :: DenseCoreCase -> BenchmarkWeight +denseMatvecWeight benchmarkCase = + eitherBenchmarkWeight + (denseCoreLabel benchmarkCase <> " matrix-vector reference rows") + vectorChecksum + (matrixVectorProduct (denseCoreRows benchmarkCase) (denseCoreVector benchmarkCase)) + +denseFlatMatvecWeight :: DenseCoreCase -> BenchmarkWeight +denseFlatMatvecWeight benchmarkCase = + eitherBenchmarkWeight + (denseCoreLabel benchmarkCase <> " flat matrix-vector") + storableVectorChecksum + ( do + matrixValue <- denseCoreFlatMatrix benchmarkCase + first show (denseDoubleMatrixVectorProduct matrixValue (denseCoreFlatVector benchmarkCase)) + ) + +forceFlatMatrix :: Either String DenseDoubleMatrix -> () +forceFlatMatrix matrixResult = + case matrixResult of + Left failureText -> rnf failureText + Right matrixValue -> rnf (denseDoubleMatrixToRowMajorVector matrixValue) + +denseMatmulWeight :: DenseCoreCase -> BenchmarkWeight +denseMatmulWeight benchmarkCase = + eitherBenchmarkWeight + (denseCoreLabel benchmarkCase <> " matrix-product reference rows") + matrixChecksum + (matrixProductRowsWith (*) (+) 0.0 (denseCoreRows benchmarkCase) (denseCoreRows benchmarkCase)) + +denseNativeMatmulWeight :: DenseCoreCase -> IO BenchmarkWeight +denseNativeMatmulWeight benchmarkCase = + case denseCoreFlatMatrix benchmarkCase of + Left failureText -> + pure (BenchmarkMeasurementFailure (denseCoreLabel benchmarkCase <> " native BLAS matrix-product: " <> failureText)) + Right matrixValue -> + eitherBenchmarkWeight + (denseCoreLabel benchmarkCase <> " native BLAS matrix-product") + (storableVectorChecksum . denseDoubleMatrixToRowMajorVector) + <$> denseDoubleMatrixProductBlas matrixValue matrixValue + +denseTransposeWeight :: DenseCoreCase -> BenchmarkWeight +denseTransposeWeight benchmarkCase = + eitherBenchmarkWeight + (denseCoreLabel benchmarkCase <> " transpose") + matrixChecksum + (transposeRowsExact (denseCoreRows benchmarkCase)) + +denseConcatWeight :: DenseCoreCase -> BenchmarkWeight +denseConcatWeight benchmarkCase = + eitherBenchmarkWeight + (denseCoreLabel benchmarkCase <> " hcat/vcat") + matrixChecksum + ( do + horizontal <- hcatRowsExact [denseCoreRows benchmarkCase, denseCoreRows benchmarkCase] + vertical <- vcatRowsExact [denseCoreRows benchmarkCase, denseCoreRows benchmarkCase] + pure (horizontal <> vertical) + ) + +vectorChecksum :: [Double] -> Double +vectorChecksum values = + sum (abs <$> values) + +storableVectorChecksum :: S.Vector Double -> Double +storableVectorChecksum values = + S.foldl' (\accumulator value -> accumulator + abs value) 0.0 values + +matrixChecksum :: [[Double]] -> Double +matrixChecksum rows = + sum (vectorChecksum <$> rows) diff --git a/bench/dense/DenseDecomposition.hs b/bench/dense/DenseDecomposition.hs new file mode 100644 --- /dev/null +++ b/bench/dense/DenseDecomposition.hs @@ -0,0 +1,194 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} + +module DenseDecomposition + ( denseDecompositionBenchmarks, + denseDecompositionOnceBenchmarks, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Bifunctor (first) +import Data.Vector.Storable qualified as S +import Fixtures + ( denseBenchmarkRows, + denseBenchmarkVector, + denseSpdRows, + ) +import Types + ( BenchmarkSetup (..), + BenchmarkWeight, + OnceBenchmark, + PreparedBenchmarkRow (..), + eigenpairsChecksum, + eitherBenchmarkWeight, + renderPreparedBenchmark, + renderPreparedOnceBenchmark, + ) +import Moonlight.LinAlg.Dense + ( DenseDoubleMatrix, + Matrix, + Vector, + denseDoubleMatrixToRowMajorVector, + fromListMatrix, + fromListVector, + mkDenseDoubleMatrixRowMajor, + toListMatrix, + toListVector, + ) +import Moonlight.LinAlg.Dense.Decomposition + ( choleskyDecomp, + qrDecompFullColumnRank, + symmetricEigen, + thinSvdFullColumnRank, + ) +import Moonlight.LinAlg.Dense.Solver + ( solveCG, + solveDirect, + solveGMRES, + ) +import Moonlight.LinAlg.Native + ( denseDoubleLinearSolveLapack, + denseDoubleSymmetricEigenpairsLapack, + ) +import Test.Tasty.Bench (Benchmark, bgroup) +import Prelude + +data DenseDecompositionCase = DenseDecompositionCase + { denseDecompositionGeneral :: !(Matrix 12 12 Double), + denseDecompositionSpd :: !(Matrix 12 12 Double), + denseDecompositionRhs :: !(Vector 12 Double), + denseDecompositionFlatGeneral :: !DenseDoubleMatrix, + denseDecompositionFlatSpd :: !DenseDoubleMatrix, + denseDecompositionFlatRhs :: !(S.Vector Double) + } + +instance NFData DenseDecompositionCase where + rnf benchmarkCase = + rnf (toListMatrix (denseDecompositionGeneral benchmarkCase)) + `seq` rnf (toListMatrix (denseDecompositionSpd benchmarkCase)) + `seq` rnf (toListVector (denseDecompositionRhs benchmarkCase)) + `seq` rnf (denseDoubleMatrixToRowMajorVector (denseDecompositionFlatGeneral benchmarkCase)) + `seq` rnf (denseDoubleMatrixToRowMajorVector (denseDecompositionFlatSpd benchmarkCase)) + `seq` rnf (denseDecompositionFlatRhs benchmarkCase) + +denseDecompositionBenchmarks :: Benchmark +denseDecompositionBenchmarks = + bgroup + "dense decomposition and solvers" + (renderPreparedBenchmark prepareDenseDecompositionCase <$> denseDecompositionRows) + +denseDecompositionOnceBenchmarks :: [OnceBenchmark] +denseDecompositionOnceBenchmarks = + renderPreparedOnceBenchmark "dense decomposition and solvers." prepareDenseDecompositionCase <$> denseDecompositionRows + +denseDecompositionRows :: [PreparedBenchmarkRow DenseDecompositionCase] +denseDecompositionRows = + [ PurePreparedBenchmarkRow "qr 12x12" qrWeight, + PurePreparedBenchmarkRow "cholesky 12x12" choleskyWeight, + PurePreparedBenchmarkRow "symmetric eigen pure 12x12" symmetricEigenWeight, + EffectfulPreparedBenchmarkRow "native LAPACK symmetric eigen certified 12x12" nativeSymmetricEigenWeight, + PurePreparedBenchmarkRow "svd 12x12" svdWeight, + PurePreparedBenchmarkRow "direct solve 12x12" directSolveWeight, + EffectfulPreparedBenchmarkRow "native LAPACK direct solve 12x12" nativeDirectSolveWeight, + PurePreparedBenchmarkRow "dense CG 12x12" denseCgWeight, + PurePreparedBenchmarkRow "dense GMRES 12x12" denseGmresWeight + ] + +prepareDenseDecompositionCase :: BenchmarkSetup DenseDecompositionCase +prepareDenseDecompositionCase = + BenchmarkSetup $ do + let generalRows = denseBenchmarkRows 12 + spdRows = denseSpdRows 12 + rhsValues = denseBenchmarkVector 12 + generalMatrix <- first show (fromListMatrix @12 @12 @Double (concat generalRows)) + spdMatrix <- first show (fromListMatrix @12 @12 @Double (concat spdRows)) + rhsVector <- first show (fromListVector @12 @Double rhsValues) + flatGeneral <- first show (mkDenseDoubleMatrixRowMajor 12 12 (S.fromList (concat generalRows))) + flatSpd <- first show (mkDenseDoubleMatrixRowMajor 12 12 (S.fromList (concat spdRows))) + pure + DenseDecompositionCase + { denseDecompositionGeneral = generalMatrix, + denseDecompositionSpd = spdMatrix, + denseDecompositionRhs = rhsVector, + denseDecompositionFlatGeneral = flatGeneral, + denseDecompositionFlatSpd = flatSpd, + denseDecompositionFlatRhs = S.fromList rhsValues + } + +qrWeight :: DenseDecompositionCase -> BenchmarkWeight +qrWeight benchmarkCase = + eitherBenchmarkWeight + "qr 12x12" + (\(qMatrix, rMatrix) -> matrixChecksum (toListMatrix qMatrix) + matrixChecksum (toListMatrix rMatrix)) + (qrDecompFullColumnRank (denseDecompositionGeneral benchmarkCase)) + +choleskyWeight :: DenseDecompositionCase -> BenchmarkWeight +choleskyWeight benchmarkCase = + eitherBenchmarkWeight + "cholesky 12x12" + (matrixChecksum . toListMatrix) + (choleskyDecomp (denseDecompositionSpd benchmarkCase)) + +symmetricEigenWeight :: DenseDecompositionCase -> BenchmarkWeight +symmetricEigenWeight benchmarkCase = + eitherBenchmarkWeight + "symmetric eigen pure 12x12" + (\(values, vectors) -> vectorChecksum (toListVector values) + matrixChecksum (toListMatrix vectors)) + (symmetricEigen (denseDecompositionSpd benchmarkCase)) + +nativeSymmetricEigenWeight :: DenseDecompositionCase -> IO BenchmarkWeight +nativeSymmetricEigenWeight benchmarkCase = + eitherBenchmarkWeight + "native LAPACK symmetric eigen certified 12x12" + eigenpairsChecksum + <$> denseDoubleSymmetricEigenpairsLapack (denseDecompositionFlatSpd benchmarkCase) + +svdWeight :: DenseDecompositionCase -> BenchmarkWeight +svdWeight benchmarkCase = + eitherBenchmarkWeight + "svd 12x12" + (\(uMatrix, sMatrix, vtMatrix) -> matrixChecksum (toListMatrix uMatrix) + matrixChecksum (toListMatrix sMatrix) + matrixChecksum (toListMatrix vtMatrix)) + (thinSvdFullColumnRank (denseDecompositionGeneral benchmarkCase)) + +directSolveWeight :: DenseDecompositionCase -> BenchmarkWeight +directSolveWeight benchmarkCase = + eitherBenchmarkWeight + "direct solve 12x12" + (vectorChecksum . toListVector) + (solveDirect (denseDecompositionGeneral benchmarkCase) (denseDecompositionRhs benchmarkCase)) + +nativeDirectSolveWeight :: DenseDecompositionCase -> IO BenchmarkWeight +nativeDirectSolveWeight benchmarkCase = + eitherBenchmarkWeight + "native LAPACK direct solve 12x12" + storableVectorChecksum + <$> denseDoubleLinearSolveLapack + (denseDecompositionFlatGeneral benchmarkCase) + (denseDecompositionFlatRhs benchmarkCase) + +denseCgWeight :: DenseDecompositionCase -> BenchmarkWeight +denseCgWeight benchmarkCase = + eitherBenchmarkWeight + "dense CG 12x12" + (vectorChecksum . toListVector) + (solveCG (denseDecompositionSpd benchmarkCase) (denseDecompositionRhs benchmarkCase)) + +denseGmresWeight :: DenseDecompositionCase -> BenchmarkWeight +denseGmresWeight benchmarkCase = + eitherBenchmarkWeight + "dense GMRES 12x12" + (vectorChecksum . toListVector) + (solveGMRES (denseDecompositionGeneral benchmarkCase) (denseDecompositionRhs benchmarkCase)) + +vectorChecksum :: [Double] -> Double +vectorChecksum values = + sum (abs <$> values) + +storableVectorChecksum :: S.Vector Double -> Double +storableVectorChecksum values = + S.foldl' (\accumulator value -> accumulator + abs value) 0.0 values + +matrixChecksum :: [Double] -> Double +matrixChecksum values = + sum (abs <$> values) diff --git a/bench/domain/DomainAlgebra.hs b/bench/domain/DomainAlgebra.hs new file mode 100644 --- /dev/null +++ b/bench/domain/DomainAlgebra.hs @@ -0,0 +1,231 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} + +module DomainAlgebra + ( domainAlgebraBenchmarks, + domainAlgebraOnceBenchmarks, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Bifunctor (first) +import Data.Vector qualified as Boxed +import Data.Vector.Unboxed qualified as Unboxed +import Fixtures + ( gf2BenchmarkValues, + ) +import Types + ( BenchmarkSetup (..), + BenchmarkWeight (..), + OnceBenchmark, + PreparedBenchmarkRow (..), + eitherBenchmarkWeight, + renderPreparedBenchmark, + renderPreparedOnceBenchmark, + ) +import Moonlight.LinAlg.Dense + ( Matrix, + fromListMatrix, + toListMatrix, + ) +import Moonlight.LinAlg.Dense.Exterior + ( exteriorPowerMatrix, + ) +import Moonlight.LinAlg.Dense.Field + ( rank, + ) +import Moonlight.LinAlg.Dense.GF2 + ( GF2 (..), + GF2PackedMatrix, + GF2SparseColumn, + defaultGF2SparseReducerConfig, + gf2SparseColumnRows, + gf2PackedWords, + mkGF2SparseColumn, + mkGF2PackedMatrixFromRowMajor, + rankGF2SparseColumns, + rankGF2PackedMatrix, + ) +import Moonlight.LinAlg.Domain + ( SmithDiagonalForm (..), + SmithNormalForm (..), + smithDiagonalForm, + smithNormalForm, + ) +import Test.Tasty.Bench (Benchmark, bgroup) +import Prelude + +data DomainAlgebraCase = DomainAlgebraCase + { domainSmithMatrix :: !(Matrix 4 4 Integer), + domainRankMatrix :: !(Matrix 8 8 Double), + domainExteriorRows :: ![[Integer]], + domainGF2Matrix :: !GF2PackedMatrix, + domainGF2SparseColumns :: !(Boxed.Vector GF2SparseColumn) + } + +instance NFData DomainAlgebraCase where + rnf benchmarkCase = + domainSmithMatrix benchmarkCase + `seq` domainRankMatrix benchmarkCase + `seq` rnf (domainExteriorRows benchmarkCase) + `seq` rnf (Unboxed.toList (gf2PackedWords (domainGF2Matrix benchmarkCase))) + `seq` rnf (gf2SparseColumnRows <$> Boxed.toList (domainGF2SparseColumns benchmarkCase)) + +domainAlgebraBenchmarks :: Benchmark +domainAlgebraBenchmarks = + bgroup + "domain algebra, exterior powers, GF2" + (renderPreparedBenchmark prepareDomainAlgebraCase <$> domainAlgebraRows) + +domainAlgebraOnceBenchmarks :: [OnceBenchmark] +domainAlgebraOnceBenchmarks = + renderPreparedOnceBenchmark "domain algebra, exterior powers, GF2." prepareDomainAlgebraCase <$> domainAlgebraRows + +domainAlgebraRows :: [PreparedBenchmarkRow DomainAlgebraCase] +domainAlgebraRows = + [ PurePreparedBenchmarkRow "Smith normal form 4x4 Integer" smithWeight, + PurePreparedBenchmarkRow "Smith diagonal form 4x4 Integer" smithDiagonalWeight, + PurePreparedBenchmarkRow "field rank 8x8 Double" fieldRankWeight, + PurePreparedBenchmarkRow "exterior power k=2 n=8 Integer" exteriorPowerWeight, + PurePreparedBenchmarkRow "GF2 packed rank 128x192" gf2RankWeight, + PurePreparedBenchmarkRow "GF2 sparse-column rank 128x192" gf2SparseRankWeight + ] + +prepareDomainAlgebraCase :: BenchmarkSetup DomainAlgebraCase +prepareDomainAlgebraCase = + BenchmarkSetup $ do + smithMatrix <- first show (fromListMatrix @4 @4 @Integer smithEntries) + rankMatrix <- first show (fromListMatrix @8 @8 @Double rankEntries) + let gf2Values = gf2BenchmarkValues 128 192 + gf2Matrix <- first show (mkGF2PackedMatrixFromRowMajor 128 192 gf2Values) + gf2SparseColumns <- gf2SparseBenchmarkColumns 128 192 gf2Values + pure + DomainAlgebraCase + { domainSmithMatrix = smithMatrix, + domainRankMatrix = rankMatrix, + domainExteriorRows = exteriorRows, + domainGF2Matrix = gf2Matrix, + domainGF2SparseColumns = gf2SparseColumns + } + +smithEntries :: [Integer] +smithEntries = + [ 6, 10, 14, 22, + 9, 15, 21, 33, + 4, 8, 12, 16, + 5, 11, 17, 23 + ] + +rankEntries :: [Double] +rankEntries = + [ rankEntry rowIndex columnIndex + | rowIndex <- [0 .. 7], + columnIndex <- [0 .. 7] + ] + +rankEntry :: Int -> Int -> Double +rankEntry rowIndex columnIndex = + let diagonalContribution = if rowIndex == columnIndex then 3 else 0 + smoothContribution = fromIntegral (((rowIndex + 1) * (columnIndex + 2)) `mod` 7) / 13.0 + in diagonalContribution + smoothContribution + +exteriorRows :: [[Integer]] +exteriorRows = + [ [ exteriorEntry rowIndex columnIndex | columnIndex <- [0 .. 7] ] + | rowIndex <- [0 .. 7] + ] + +exteriorEntry :: Int -> Int -> Integer +exteriorEntry rowIndex columnIndex + | rowIndex == columnIndex = 2 + fromIntegral rowIndex + | abs (rowIndex - columnIndex) == 1 = 1 + | otherwise = 0 + +smithWeight :: DomainAlgebraCase -> BenchmarkWeight +smithWeight benchmarkCase = + eitherBenchmarkWeight + "Smith normal form 4x4 Integer" + smithChecksum + (smithNormalForm (domainSmithMatrix benchmarkCase)) + +smithDiagonalWeight :: DomainAlgebraCase -> BenchmarkWeight +smithDiagonalWeight benchmarkCase = + eitherBenchmarkWeight + "Smith diagonal form 4x4 Integer" + smithDiagonalChecksum + (smithDiagonalForm (domainSmithMatrix benchmarkCase)) + +fieldRankWeight :: DomainAlgebraCase -> BenchmarkWeight +fieldRankWeight benchmarkCase = + eitherBenchmarkWeight + "field rank 8x8 Double" + fromIntegral + (rank (domainRankMatrix benchmarkCase)) + +exteriorPowerWeight :: DomainAlgebraCase -> BenchmarkWeight +exteriorPowerWeight benchmarkCase = + eitherBenchmarkWeight + "exterior power k=2 n=8 Integer" + integerRowsChecksum + (exteriorPowerMatrix 2 (domainExteriorRows benchmarkCase)) + +gf2RankWeight :: DomainAlgebraCase -> BenchmarkWeight +gf2RankWeight benchmarkCase = + BenchmarkWeight (fromIntegral (rankGF2PackedMatrix (domainGF2Matrix benchmarkCase))) + +gf2SparseRankWeight :: DomainAlgebraCase -> BenchmarkWeight +gf2SparseRankWeight benchmarkCase = + eitherBenchmarkWeight + "GF2 sparse-column rank 128x192" + fromIntegral + (rankGF2SparseColumns defaultGF2SparseReducerConfig 128 192 (domainGF2SparseColumns benchmarkCase)) + +smithChecksum :: SmithNormalForm 4 4 Integer -> Double +smithChecksum smithValue = + integerVectorChecksum (toListMatrix (smithDiagonal smithValue)) + + integerVectorChecksum (toListMatrix (smithLeft smithValue)) + + integerVectorChecksum (toListMatrix (smithRight smithValue)) + +smithDiagonalChecksum :: SmithDiagonalForm 4 4 Integer -> Double +smithDiagonalChecksum smithValue = + integerVectorChecksum (toListMatrix (smithDiagonalMatrix smithValue)) + +integerRowsChecksum :: [[Integer]] -> Double +integerRowsChecksum rows = + integerVectorChecksum (concat rows) + +integerVectorChecksum :: [Integer] -> Double +integerVectorChecksum values = + fromIntegral (sum (abs <$> values)) + +gf2SparseBenchmarkColumns :: Int -> Int -> [GF2] -> Either String (Boxed.Vector GF2SparseColumn) +gf2SparseBenchmarkColumns rowCount columnCount values + | length values /= rowCount * columnCount = + Left "GF2 sparse-column benchmark fixture has malformed row-major length" + | otherwise = + first show + ( Boxed.fromList + <$> traverse + ( \columnIndex -> + mkGF2SparseColumn + ("GF2 sparse-column benchmark column " <> show columnIndex) + rowCount + columnIndex + (gf2SparseBenchmarkSupport columnIndex) + ) + [0 .. columnCount - 1] + ) + where + indexedValues = + zip + [ (rowIndex, columnIndex) + | rowIndex <- [0 .. rowCount - 1], + columnIndex <- [0 .. columnCount - 1] + ] + values + + gf2SparseBenchmarkSupport columnIndex = + [ rowIndex + | ((rowIndex, valueColumn), GF2One) <- indexedValues, + valueColumn == columnIndex + ] diff --git a/bench/native/NativeLapack.hs b/bench/native/NativeLapack.hs new file mode 100644 --- /dev/null +++ b/bench/native/NativeLapack.hs @@ -0,0 +1,350 @@ +module NativeLapack + ( nativeLapackBenchmarks, + nativeLapackOnceBenchmarks, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Bifunctor (first) +import Env (BenchmarkSelection (..)) +import Fixtures + ( benchmarkSeedBlock, + denseOperator, + genericBenchmarkTridiagonal, + pathLaplacianTridiagonal, + projectedBenchmarkDimension, + projectedBenchmarkRows, + projectedBlockBenchmarkCases, + ) +import Types + ( BenchmarkSetup (..), + BenchmarkWeight (..), + OnceBenchmark (..), + PreparedBenchmarkRow (..), + ProjectedBlockBenchmarkCase (..), + eigenpairsChecksum, + benchmarkWeightEither, + renderPreparedBenchmark, + renderPreparedOnceBenchmark, + ) +import Data.Vector.Unboxed qualified as U +import Moonlight.LinAlg.Dense (DynMatrix, mkDynMatrix) +import Moonlight.LinAlg.Krylov + ( SpectrumEnd (SmallestEigenvalues), + blockLanczosProjectedBlockTridiagonal, + blockLanczosSymmetric, + defaultBlockLanczosConfig, + mkPositiveCount, + withBlockLanczosBlockSize, + withBlockLanczosIterations, + ) +import Moonlight.LinAlg.Native + ( selectedSymmetricBlockTridiagonalEigenRequestLapack, + selectedSymmetricTridiagonalEigenRequestLapack, + symmetricEigenRequestLapack, + ) +import Moonlight.LinAlg.Pure.Structured.BlockTridiagonal (SymmetricBlockTridiagonal) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal (SymmetricTridiagonal) +import Moonlight.LinAlg.Spectral + ( Eigenpairs, + EigenRequest (..), + ) +import Test.Tasty.Bench (Benchmark, bench, bgroup, nfIO) +import Prelude + +data NativeTridiagonalLapackCase = NativeTridiagonalLapackCase + { nativeTridiagonalLapackKind :: !NativeTridiagonalLapackKind, + nativeTridiagonalLapackDimension :: !Int, + nativeTridiagonalLapackModes :: !Int + } + +data NativeTridiagonalLapackKind + = NativePathLaplacianTridiagonal + | NativeGenericTridiagonal + deriving stock (Eq, Show) + +data NativeProjectedBandPreparedCase = NativeProjectedBandPreparedCase + { nativeProjectedBandPreparedCase :: !ProjectedBlockBenchmarkCase, + nativeProjectedBandOperator :: !SymmetricBlockTridiagonal + } + +instance NFData NativeProjectedBandPreparedCase where + rnf preparedCase = + nativeProjectedBandPreparedCase preparedCase + `seq` nativeProjectedBandOperator preparedCase + `seq` () + +nativeLapackBenchmarks :: BenchmarkSelection -> Benchmark +nativeLapackBenchmarks benchmarkSelection = + bgroup + "native LAPACK symmetric eigensolve" + ( (renderNativeLapackBenchmark . nativeLapackMeasuredRow <$> projectedBlockBenchmarkCases benchmarkSelection) + <> concatMap (fmap renderNativeLapackBenchmark . nativeDenseSelectedRows) (projectedBlockBenchmarkCases benchmarkSelection) + <> concatMap + (\benchmarkCase -> renderPreparedBenchmark (prepareNativeProjectedBandCase benchmarkCase) <$> nativeProjectedBandRows benchmarkCase) + (projectedBlockBenchmarkCases benchmarkSelection) + <> concatMap (fmap renderNativeLapackBenchmark . nativeTridiagonalLapackRows) (nativeTridiagonalLapackCases benchmarkSelection) + ) + +nativeLapackOnceBenchmarks :: BenchmarkSelection -> [OnceBenchmark] +nativeLapackOnceBenchmarks benchmarkSelection = + (renderNativeLapackOnceBenchmark . nativeLapackMeasuredRow <$> projectedBlockBenchmarkCases benchmarkSelection) + <> concatMap (fmap renderNativeLapackOnceBenchmark . nativeDenseSelectedRows) (projectedBlockBenchmarkCases benchmarkSelection) + <> concatMap + (\benchmarkCase -> renderPreparedOnceBenchmark "native LAPACK symmetric eigensolve." (prepareNativeProjectedBandCase benchmarkCase) <$> nativeProjectedBandRows benchmarkCase) + (projectedBlockBenchmarkCases benchmarkSelection) + <> concatMap (fmap renderNativeLapackOnceBenchmark . nativeTridiagonalLapackRows) (nativeTridiagonalLapackCases benchmarkSelection) + +nativeLapackMeasuredRow :: ProjectedBlockBenchmarkCase -> (String, IO BenchmarkWeight) +nativeLapackMeasuredRow benchmarkCase = + (nativeLapackBenchmarkLabel benchmarkCase, nativeLapackWeight benchmarkCase) + +renderNativeLapackBenchmark :: (String, IO BenchmarkWeight) -> Benchmark +renderNativeLapackBenchmark (rowLabel, rowAction) = + bench rowLabel (nfIO rowAction) + +renderNativeLapackOnceBenchmark :: (String, IO BenchmarkWeight) -> OnceBenchmark +renderNativeLapackOnceBenchmark (rowLabel, rowAction) = + OnceBenchmark + { onceBenchmarkLabel = "native LAPACK symmetric eigensolve." <> rowLabel, + onceBenchmarkAction = benchmarkWeightEither <$> rowAction + } + +nativeLapackBenchmarkLabel :: ProjectedBlockBenchmarkCase -> String +nativeLapackBenchmarkLabel benchmarkCase = + projectedBenchmarkLabel benchmarkCase + <> " profile=" + <> show (projectedBenchmarkSpectrumProfile benchmarkCase) + <> " n=" + <> show (projectedBenchmarkDimension benchmarkCase) + +nativeLapackWeight :: ProjectedBlockBenchmarkCase -> IO BenchmarkWeight +nativeLapackWeight benchmarkCase = + case mkDynMatrix + (projectedBenchmarkDimension benchmarkCase) + (projectedBenchmarkDimension benchmarkCase) + (concat (projectedBenchmarkRows benchmarkCase)) of + Left err -> pure (BenchmarkMeasurementFailure (projectedBenchmarkLabel benchmarkCase <> ": " <> show err)) + Right matrixValue -> + case mkPositiveCount (projectedBenchmarkDimension benchmarkCase) of + Left err -> pure (BenchmarkMeasurementFailure (projectedBenchmarkLabel benchmarkCase <> ": " <> show err)) + Right requestedCount -> + symmetricEigenRequestLapack (EigenpairsRequest SmallestEigenvalues requestedCount) matrixValue + >>= nativeEigenpairsWeight (projectedBenchmarkLabel benchmarkCase) + +nativeDenseSelectedRows :: ProjectedBlockBenchmarkCase -> [(String, IO BenchmarkWeight)] +nativeDenseSelectedRows benchmarkCase = + [ (nativeDenseSelectedValuesLabel benchmarkCase, nativeDenseSelectedValuesWeight benchmarkCase), + (nativeDenseSelectedPairsLabel benchmarkCase, nativeDenseSelectedPairsWeight benchmarkCase) + ] + +nativeDenseSelectedValuesLabel :: ProjectedBlockBenchmarkCase -> String +nativeDenseSelectedValuesLabel = + nativeDenseSelectedLabel "DSYEVX dense values" + +nativeDenseSelectedPairsLabel :: ProjectedBlockBenchmarkCase -> String +nativeDenseSelectedPairsLabel = + nativeDenseSelectedLabel "DSYEVX dense pairs" + +nativeDenseSelectedLabel :: String -> ProjectedBlockBenchmarkCase -> String +nativeDenseSelectedLabel requestLabel benchmarkCase = + projectedBenchmarkLabel benchmarkCase + <> " " + <> requestLabel + <> " modes=" + <> show (projectedBenchmarkRequestedModes benchmarkCase) + <> " profile=" + <> show (projectedBenchmarkSpectrumProfile benchmarkCase) + <> " n=" + <> show (projectedBenchmarkDimension benchmarkCase) + +nativeDenseSelectedValuesWeight :: ProjectedBlockBenchmarkCase -> IO BenchmarkWeight +nativeDenseSelectedValuesWeight benchmarkCase = + case prepareNativeDenseMatrix benchmarkCase of + Left err -> pure (BenchmarkMeasurementFailure (nativeDenseSelectedValuesLabel benchmarkCase <> ": " <> err)) + Right matrixValue -> + case mkPositiveCount (projectedBenchmarkRequestedModes benchmarkCase) of + Left err -> pure (BenchmarkMeasurementFailure (nativeDenseSelectedValuesLabel benchmarkCase <> ": " <> show err)) + Right requestedCount -> + symmetricEigenRequestLapack (EigenvaluesRequest SmallestEigenvalues requestedCount) matrixValue + >>= nativeEigenvaluesWeight (nativeDenseSelectedValuesLabel benchmarkCase) + +nativeDenseSelectedPairsWeight :: ProjectedBlockBenchmarkCase -> IO BenchmarkWeight +nativeDenseSelectedPairsWeight benchmarkCase = + case prepareNativeDenseMatrix benchmarkCase of + Left err -> pure (BenchmarkMeasurementFailure (nativeDenseSelectedPairsLabel benchmarkCase <> ": " <> err)) + Right matrixValue -> + case mkPositiveCount (projectedBenchmarkRequestedModes benchmarkCase) of + Left err -> pure (BenchmarkMeasurementFailure (nativeDenseSelectedPairsLabel benchmarkCase <> ": " <> show err)) + Right requestedCount -> + symmetricEigenRequestLapack (EigenpairsRequest SmallestEigenvalues requestedCount) matrixValue + >>= nativeEigenpairsWeight (nativeDenseSelectedPairsLabel benchmarkCase) + +prepareNativeDenseMatrix :: ProjectedBlockBenchmarkCase -> Either String (DynMatrix Double) +prepareNativeDenseMatrix benchmarkCase = + first + show + ( mkDynMatrix + (projectedBenchmarkDimension benchmarkCase) + (projectedBenchmarkDimension benchmarkCase) + (concat (projectedBenchmarkRows benchmarkCase)) + ) + +nativeProjectedBandRows :: ProjectedBlockBenchmarkCase -> [PreparedBenchmarkRow NativeProjectedBandPreparedCase] +nativeProjectedBandRows benchmarkCase = + [ EffectfulPreparedBenchmarkRow (nativeProjectedBandValuesLabel benchmarkCase) nativeProjectedBandValuesWeight, + EffectfulPreparedBenchmarkRow (nativeProjectedBandPairsLabel benchmarkCase) nativeProjectedBandPairsWeight + ] + +prepareNativeProjectedBandCase :: ProjectedBlockBenchmarkCase -> BenchmarkSetup NativeProjectedBandPreparedCase +prepareNativeProjectedBandCase benchmarkCase = + BenchmarkSetup $ do + iterationCount <- + first + (\err -> "invalid native projected-band iteration count for " <> projectedBenchmarkLabel benchmarkCase <> ": " <> show err) + (mkPositiveCount (projectedBenchmarkIterations benchmarkCase)) + blockSize <- + first + (\err -> "invalid native projected-band block size for " <> projectedBenchmarkLabel benchmarkCase <> ": " <> show err) + (mkPositiveCount (projectedBenchmarkBlockSize benchmarkCase)) + operatorValue <- + first + (\err -> "native projected-band operator construction failed for " <> projectedBenchmarkLabel benchmarkCase <> ": " <> err) + (denseOperator (projectedBenchmarkRows benchmarkCase)) + let operatorDimension = projectedBenchmarkDimension benchmarkCase + seedBlock = benchmarkSeedBlock operatorDimension (projectedBenchmarkBlockSize benchmarkCase) + blockConfig = + withBlockLanczosBlockSize + blockSize + (withBlockLanczosIterations iterationCount defaultBlockLanczosConfig) + decomposition <- + first + (\err -> "native projected-band decomposition failed for " <> projectedBenchmarkLabel benchmarkCase <> ": " <> show err) + (blockLanczosSymmetric blockConfig operatorValue seedBlock) + pure + NativeProjectedBandPreparedCase + { nativeProjectedBandPreparedCase = benchmarkCase, + nativeProjectedBandOperator = blockLanczosProjectedBlockTridiagonal decomposition + } + +nativeProjectedBandValuesLabel :: ProjectedBlockBenchmarkCase -> String +nativeProjectedBandValuesLabel benchmarkCase = + nativeProjectedBandLabel "DSBEVX projected block values" benchmarkCase + +nativeProjectedBandPairsLabel :: ProjectedBlockBenchmarkCase -> String +nativeProjectedBandPairsLabel benchmarkCase = + nativeProjectedBandLabel "DSBEVX projected block pairs" benchmarkCase + +nativeProjectedBandLabel :: String -> ProjectedBlockBenchmarkCase -> String +nativeProjectedBandLabel requestLabel benchmarkCase = + projectedBenchmarkLabel benchmarkCase + <> " " + <> requestLabel + <> " modes=" + <> show (projectedBenchmarkRequestedModes benchmarkCase) + <> " profile=" + <> show (projectedBenchmarkSpectrumProfile benchmarkCase) + <> " n=" + <> show (projectedBenchmarkDimension benchmarkCase) + +nativeProjectedBandValuesWeight :: NativeProjectedBandPreparedCase -> IO BenchmarkWeight +nativeProjectedBandValuesWeight preparedCase = + case mkPositiveCount (projectedBenchmarkRequestedModes (nativeProjectedBandPreparedCase preparedCase)) of + Left err -> pure (BenchmarkMeasurementFailure (projectedBenchmarkLabel (nativeProjectedBandPreparedCase preparedCase) <> " DSBEVX values: " <> show err)) + Right requestedCount -> + selectedSymmetricBlockTridiagonalEigenRequestLapack + (EigenvaluesRequest SmallestEigenvalues requestedCount) + (nativeProjectedBandOperator preparedCase) + >>= nativeEigenvaluesWeight (nativeProjectedBandValuesLabel (nativeProjectedBandPreparedCase preparedCase)) + +nativeProjectedBandPairsWeight :: NativeProjectedBandPreparedCase -> IO BenchmarkWeight +nativeProjectedBandPairsWeight preparedCase = + case mkPositiveCount (projectedBenchmarkRequestedModes (nativeProjectedBandPreparedCase preparedCase)) of + Left err -> pure (BenchmarkMeasurementFailure (projectedBenchmarkLabel (nativeProjectedBandPreparedCase preparedCase) <> " DSBEVX pairs: " <> show err)) + Right requestedCount -> + selectedSymmetricBlockTridiagonalEigenRequestLapack + (EigenpairsRequest SmallestEigenvalues requestedCount) + (nativeProjectedBandOperator preparedCase) + >>= nativeEigenpairsWeight (nativeProjectedBandPairsLabel (nativeProjectedBandPreparedCase preparedCase)) + +nativeTridiagonalLapackCases :: BenchmarkSelection -> [NativeTridiagonalLapackCase] +nativeTridiagonalLapackCases benchmarkSelection = + [ NativeTridiagonalLapackCase NativePathLaplacianTridiagonal 512 4, + NativeTridiagonalLapackCase NativeGenericTridiagonal 512 4 + ] + <> [NativeTridiagonalLapackCase NativePathLaplacianTridiagonal 10000 4 | includeNativeLarge benchmarkSelection] + +nativeTridiagonalLapackRows :: NativeTridiagonalLapackCase -> [(String, IO BenchmarkWeight)] +nativeTridiagonalLapackRows benchmarkCase = + [ (nativeTridiagonalLapackValuesLabel benchmarkCase, nativeTridiagonalLapackValuesWeight benchmarkCase), + (nativeTridiagonalLapackPairsLabel benchmarkCase, nativeTridiagonalLapackPairsWeight benchmarkCase) + ] + +nativeTridiagonalLapackValuesLabel :: NativeTridiagonalLapackCase -> String +nativeTridiagonalLapackValuesLabel benchmarkCase = + nativeTridiagonalLapackLabelPrefix benchmarkCase + <> " DSTEMR selected tridiagonal values modes=" + <> show (nativeTridiagonalLapackModes benchmarkCase) + +nativeTridiagonalLapackPairsLabel :: NativeTridiagonalLapackCase -> String +nativeTridiagonalLapackPairsLabel benchmarkCase = + nativeTridiagonalLapackLabelPrefix benchmarkCase + <> " DSTEMR selected tridiagonal pairs modes=" + <> show (nativeTridiagonalLapackModes benchmarkCase) + +nativeTridiagonalLapackLabelPrefix :: NativeTridiagonalLapackCase -> String +nativeTridiagonalLapackLabelPrefix benchmarkCase = + nativeTridiagonalLapackKindLabel (nativeTridiagonalLapackKind benchmarkCase) + <> show (nativeTridiagonalLapackDimension benchmarkCase) + +nativeTridiagonalLapackKindLabel :: NativeTridiagonalLapackKind -> String +nativeTridiagonalLapackKindLabel benchmarkKind = + case benchmarkKind of + NativePathLaplacianTridiagonal -> "path-laplacian-" + NativeGenericTridiagonal -> "generic-tridiagonal-" + +nativeTridiagonalLapackValuesWeight :: NativeTridiagonalLapackCase -> IO BenchmarkWeight +nativeTridiagonalLapackValuesWeight benchmarkCase = + case mkPositiveCount (nativeTridiagonalLapackModes benchmarkCase) of + Left err -> pure (BenchmarkMeasurementFailure (nativeTridiagonalLapackValuesLabel benchmarkCase <> ": " <> show err)) + Right requestedCount -> + case nativeTridiagonalLapackOperator benchmarkCase of + Left err -> pure (BenchmarkMeasurementFailure (nativeTridiagonalLapackValuesLabel benchmarkCase <> ": " <> err)) + Right tridiagonalValue -> + selectedSymmetricTridiagonalEigenRequestLapack (EigenvaluesRequest SmallestEigenvalues requestedCount) tridiagonalValue + >>= nativeEigenvaluesWeight (nativeTridiagonalLapackValuesLabel benchmarkCase) + +nativeTridiagonalLapackPairsWeight :: NativeTridiagonalLapackCase -> IO BenchmarkWeight +nativeTridiagonalLapackPairsWeight benchmarkCase = + case mkPositiveCount (nativeTridiagonalLapackModes benchmarkCase) of + Left err -> pure (BenchmarkMeasurementFailure (nativeTridiagonalLapackPairsLabel benchmarkCase <> ": " <> show err)) + Right requestedCount -> + case nativeTridiagonalLapackOperator benchmarkCase of + Left err -> pure (BenchmarkMeasurementFailure (nativeTridiagonalLapackPairsLabel benchmarkCase <> ": " <> err)) + Right tridiagonalValue -> + selectedSymmetricTridiagonalEigenRequestLapack (EigenpairsRequest SmallestEigenvalues requestedCount) tridiagonalValue + >>= nativeEigenpairsWeight (nativeTridiagonalLapackPairsLabel benchmarkCase) + +nativeTridiagonalLapackOperator :: NativeTridiagonalLapackCase -> Either String SymmetricTridiagonal +nativeTridiagonalLapackOperator benchmarkCase = + case nativeTridiagonalLapackKind benchmarkCase of + NativePathLaplacianTridiagonal -> + pathLaplacianTridiagonal (nativeTridiagonalLapackDimension benchmarkCase) + NativeGenericTridiagonal -> + genericBenchmarkTridiagonal (nativeTridiagonalLapackDimension benchmarkCase) + +nativeEigenvaluesWeight :: Show err => String -> Either err (U.Vector Double) -> IO BenchmarkWeight +nativeEigenvaluesWeight label eigenResult = + pure + ( case eigenResult of + Left err -> BenchmarkMeasurementFailure (label <> ": " <> show err) + Right values -> BenchmarkWeight (U.sum values) + ) + +nativeEigenpairsWeight :: Show err => String -> Either err Eigenpairs -> IO BenchmarkWeight +nativeEigenpairsWeight label eigenResult = + pure + ( case eigenResult of + Left err -> BenchmarkMeasurementFailure (label <> ": " <> show err) + Right pairs -> BenchmarkWeight (eigenpairsChecksum pairs) + ) diff --git a/bench/sparse/SparseSolvers.hs b/bench/sparse/SparseSolvers.hs new file mode 100644 --- /dev/null +++ b/bench/sparse/SparseSolvers.hs @@ -0,0 +1,232 @@ +module SparseSolvers + ( sparseSolverBenchmarks, + sparseSolverOnceBenchmarks, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Bifunctor (first) +import qualified Data.Vector.Unboxed as U +import Env (BenchmarkSelection (..)) +import Fixtures + ( bandedSpdCSR, + denseBenchmarkVector, + diagonalBenchmarkValues, + ) +import Types + ( BenchmarkSetup (..), + BenchmarkWeight, + OnceBenchmark, + PreparedBenchmarkRow (..), + eitherBenchmarkWeight, + renderPreparedBenchmark, + renderPreparedOnceBenchmark, + ) +import Moonlight.LinAlg.Sparse + ( IC0Config (..), + SparseCSR, + SparseConjugateGradientConfig (..), + SparseGMRESConfig (..), + SparseIterativeResult (..), + SparsePreconditionerFamily (..), + SparseStationaryIterationConfig (..), + csrColumnIndicesVector, + csrRowOffsetsVector, + csrValuesVector, + diagonalCSR, + solveSparseCG, + solveSparseGMRES, + solveSparseJacobi, + solveSparseRichardson, + ) +import Test.Tasty.Bench (Benchmark, bgroup) +import Prelude + +data SparseSolverCase = SparseSolverCase + { sparseSolverLabel :: !String, + sparseSolverMatrix :: !(SparseCSR Double), + sparseSolverDiagonalMatrix :: !(SparseCSR Double), + sparseSolverRhs :: !(U.Vector Double), + sparseSolverInitialGuess :: !(U.Vector Double), + sparseSolverCGConfig :: !SparseConjugateGradientConfig, + sparseSolverGMRESConfig :: !SparseGMRESConfig, + sparseSolverStationaryConfig :: !SparseStationaryIterationConfig + } + +instance NFData SparseSolverCase where + rnf benchmarkCase = + rnf (sparseSolverLabel benchmarkCase) + `seq` rnf (csrRowOffsetsVector (sparseSolverMatrix benchmarkCase)) + `seq` rnf (csrColumnIndicesVector (sparseSolverMatrix benchmarkCase)) + `seq` rnf (csrValuesVector (sparseSolverMatrix benchmarkCase)) + `seq` rnf (csrRowOffsetsVector (sparseSolverDiagonalMatrix benchmarkCase)) + `seq` rnf (csrColumnIndicesVector (sparseSolverDiagonalMatrix benchmarkCase)) + `seq` rnf (csrValuesVector (sparseSolverDiagonalMatrix benchmarkCase)) + `seq` rnf (sparseSolverRhs benchmarkCase) + `seq` rnf (sparseSolverInitialGuess benchmarkCase) + `seq` sparseSolverCGConfig benchmarkCase + `seq` sparseSolverGMRESConfig benchmarkCase + `seq` sparseSolverStationaryConfig benchmarkCase + `seq` () + +sparseSolverBenchmarks :: BenchmarkSelection -> Benchmark +sparseSolverBenchmarks benchmarkSelection = + bgroup + "sparse iterative solvers" + (sparseSolverBenchmark <$> sparseSolverDimensions benchmarkSelection) + +sparseSolverOnceBenchmarks :: BenchmarkSelection -> [OnceBenchmark] +sparseSolverOnceBenchmarks benchmarkSelection = + sparseSolverOnceBenchmark =<< sparseSolverDimensions benchmarkSelection + +sparseSolverDimensions :: BenchmarkSelection -> [Int] +sparseSolverDimensions benchmarkSelection = + [64] + <> [128 | includeBroadMedium benchmarkSelection || includeBroadLarge benchmarkSelection] + <> [256 | includeBroadLarge benchmarkSelection] + +sparseSolverBenchmark :: Int -> Benchmark +sparseSolverBenchmark dimension = + bgroup + ("n=" <> show dimension) + (renderPreparedBenchmark (prepareSparseSolverCase dimension) <$> sparseSolverRows) + +sparseSolverOnceBenchmark :: Int -> [OnceBenchmark] +sparseSolverOnceBenchmark dimension = + renderPreparedOnceBenchmark ("sparse iterative solvers.n=" <> show dimension <> ".") (prepareSparseSolverCase dimension) + <$> sparseSolverRows + +sparseSolverRows :: [PreparedBenchmarkRow SparseSolverCase] +sparseSolverRows = + [ PurePreparedBenchmarkRow "CG" sparseCgWeight, + PurePreparedBenchmarkRow "PCG diagonal" sparsePcgDiagonalWeight, + PurePreparedBenchmarkRow "PCG SSOR" sparsePcgSsorWeight, + PurePreparedBenchmarkRow "PCG IC0" sparsePcgIC0Weight, + PurePreparedBenchmarkRow "GMRES" sparseGmresWeight, + PurePreparedBenchmarkRow "Jacobi diagonal" sparseJacobiWeight, + PurePreparedBenchmarkRow "Richardson diagonal" sparseRichardsonWeight + ] + +prepareSparseSolverCase :: Int -> BenchmarkSetup SparseSolverCase +prepareSparseSolverCase dimension = + BenchmarkSetup $ do + sparseMatrix <- first (("sparse solver matrix fixture failed: " <>)) (bandedSpdCSR dimension) + diagonalMatrix <- first (("sparse solver diagonal fixture failed: " <>) . show) (diagonalCSR (diagonalBenchmarkValues dimension)) + let rhsValues = U.fromList (denseBenchmarkVector dimension) + pure + SparseSolverCase + { sparseSolverLabel = "n=" <> show dimension, + sparseSolverMatrix = sparseMatrix, + sparseSolverDiagonalMatrix = diagonalMatrix, + sparseSolverRhs = rhsValues, + sparseSolverInitialGuess = U.replicate dimension 0.0, + sparseSolverCGConfig = + SparseConjugateGradientConfig + { scgcTolerance = 1.0e-8, + scgcIterationLimit = max 64 (dimension * 4), + scgcPreconditionerFamily = IdentitySparsePreconditionerFamily + }, + sparseSolverGMRESConfig = + SparseGMRESConfig + { sgcTolerance = 1.0e-8, + sgcIterationLimit = max 64 (dimension * 2), + sgcRestartDimension = min 24 (max 4 dimension), + sgcPreconditionerFamily = IdentitySparsePreconditionerFamily + }, + sparseSolverStationaryConfig = + SparseStationaryIterationConfig + { ssicTolerance = 1.0e-8, + ssicIterationLimit = max 64 (dimension * 4), + ssicDamping = 0.9 + } + } + +sparseCgWeight :: SparseSolverCase -> BenchmarkWeight +sparseCgWeight benchmarkCase = + eitherBenchmarkWeight + (sparseSolverLabel benchmarkCase <> " CG") + sparseResultChecksum + ( solveSparseCG + (sparseSolverCGConfig benchmarkCase) + (sparseSolverMatrix benchmarkCase) + (sparseSolverRhs benchmarkCase) + (sparseSolverInitialGuess benchmarkCase) + ) + +sparsePcgDiagonalWeight :: SparseSolverCase -> BenchmarkWeight +sparsePcgDiagonalWeight benchmarkCase = + eitherBenchmarkWeight + (sparseSolverLabel benchmarkCase <> " PCG diagonal") + sparseResultChecksum + ( solveSparseCG + ((sparseSolverCGConfig benchmarkCase) {scgcPreconditionerFamily = DiagonalJacobiSparsePreconditionerFamily}) + (sparseSolverMatrix benchmarkCase) + (sparseSolverRhs benchmarkCase) + (sparseSolverInitialGuess benchmarkCase) + ) + +sparsePcgSsorWeight :: SparseSolverCase -> BenchmarkWeight +sparsePcgSsorWeight benchmarkCase = + eitherBenchmarkWeight + (sparseSolverLabel benchmarkCase <> " PCG SSOR") + sparseResultChecksum + ( solveSparseCG + ((sparseSolverCGConfig benchmarkCase) {scgcPreconditionerFamily = SsorSparsePreconditionerFamily 1.0}) + (sparseSolverMatrix benchmarkCase) + (sparseSolverRhs benchmarkCase) + (sparseSolverInitialGuess benchmarkCase) + ) + +sparsePcgIC0Weight :: SparseSolverCase -> BenchmarkWeight +sparsePcgIC0Weight benchmarkCase = + eitherBenchmarkWeight + (sparseSolverLabel benchmarkCase <> " PCG IC0") + sparseResultChecksum + ( solveSparseCG + ((sparseSolverCGConfig benchmarkCase) {scgcPreconditionerFamily = IncompleteCholesky0SparsePreconditionerFamily (IC0Config Nothing)}) + (sparseSolverMatrix benchmarkCase) + (sparseSolverRhs benchmarkCase) + (sparseSolverInitialGuess benchmarkCase) + ) + +sparseGmresWeight :: SparseSolverCase -> BenchmarkWeight +sparseGmresWeight benchmarkCase = + eitherBenchmarkWeight + (sparseSolverLabel benchmarkCase <> " GMRES") + sparseResultChecksum + ( solveSparseGMRES + (sparseSolverGMRESConfig benchmarkCase) + (sparseSolverMatrix benchmarkCase) + (sparseSolverRhs benchmarkCase) + (sparseSolverInitialGuess benchmarkCase) + ) + +sparseJacobiWeight :: SparseSolverCase -> BenchmarkWeight +sparseJacobiWeight benchmarkCase = + eitherBenchmarkWeight + (sparseSolverLabel benchmarkCase <> " Jacobi diagonal") + sparseResultChecksum + ( solveSparseJacobi + (sparseSolverStationaryConfig benchmarkCase) + (sparseSolverDiagonalMatrix benchmarkCase) + (sparseSolverRhs benchmarkCase) + (sparseSolverInitialGuess benchmarkCase) + ) + +sparseRichardsonWeight :: SparseSolverCase -> BenchmarkWeight +sparseRichardsonWeight benchmarkCase = + eitherBenchmarkWeight + (sparseSolverLabel benchmarkCase <> " Richardson diagonal") + sparseResultChecksum + ( solveSparseRichardson + (sparseSolverStationaryConfig benchmarkCase) + (sparseSolverDiagonalMatrix benchmarkCase) + (sparseSolverRhs benchmarkCase) + (sparseSolverInitialGuess benchmarkCase) + ) + +sparseResultChecksum :: SparseIterativeResult -> Double +sparseResultChecksum resultValue = + fromIntegral (sparseIterations resultValue) + + sparseResidualNorm resultValue + + U.sum (U.map abs (sparseSolution resultValue)) diff --git a/bench/sparse/SparseStorage.hs b/bench/sparse/SparseStorage.hs new file mode 100644 --- /dev/null +++ b/bench/sparse/SparseStorage.hs @@ -0,0 +1,228 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} + +module SparseStorage + ( sparseStorageBenchmarks, + sparseStorageOnceBenchmarks, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Bifunctor (first) +import Data.Vector.Unboxed qualified as Unboxed +import Env (BenchmarkSelection (..)) +import Fixtures + ( bandedDenseRows, + bandedSpdCSR, + denseBenchmarkVector, + packedSparseBenchmarkOperator, + ) +import Types + ( BenchmarkSetup (..), + BenchmarkWeight (..), + OnceBenchmark, + PreparedBenchmarkRow (..), + eitherBenchmarkWeight, + renderPreparedBenchmark, + renderPreparedOnceBenchmark, + ) +import Moonlight.LinAlg.Dense + ( Matrix, + fromListMatrix, + ) +import Moonlight.LinAlg.Sparse + ( GraphEdge (..), + PackedSparseOperator, + SparseCOO, + SparseCSC, + SparseCSR, + applyPackedSparseOperatorDense, + cooEntries, + csrColumnIndicesVector, + csrMatVecVector, + csrRowOffsetsVector, + csrToCOO, + csrToCSC, + csrValuesVector, + cscColumnOffsetsVector, + cscRowIndicesVector, + cscValuesVector, + denseToCOO, + denseToCSC, + denseToCSR, + graphLaplacianCSR, + ) +import Test.Tasty.Bench (Benchmark, bgroup) +import Prelude + +data SparseStorageCase = SparseStorageCase + { sparseStorageLabel :: !String, + sparseStorageDense32 :: !(Matrix 32 32 Double), + sparseStorageCSR :: !(SparseCSR Double), + sparseStoragePackedOperator :: !(PackedSparseOperator Double), + sparseStoragePackedVector :: !(Unboxed.Vector Double), + sparseStorageGraphVertices :: ![Int], + sparseStorageGraphEdges :: ![GraphEdge Int] + } + +instance NFData SparseStorageCase where + rnf benchmarkCase = + rnf (sparseStorageLabel benchmarkCase) + `seq` sparseStorageDense32 benchmarkCase + `seq` rnf (csrRowOffsetsVector (sparseStorageCSR benchmarkCase)) + `seq` rnf (csrColumnIndicesVector (sparseStorageCSR benchmarkCase)) + `seq` rnf (csrValuesVector (sparseStorageCSR benchmarkCase)) + `seq` sparseStoragePackedOperator benchmarkCase + `seq` rnf (sparseStoragePackedVector benchmarkCase) + `seq` rnf (sparseStorageGraphVertices benchmarkCase) + `seq` rnf + ( fmap + ( \edgeValue -> + ( graphEdgeLeft edgeValue, + graphEdgeRight edgeValue, + graphEdgeWeight edgeValue + ) + ) + (sparseStorageGraphEdges benchmarkCase) + ) + `seq` () + +sparseStorageBenchmarks :: BenchmarkSelection -> Benchmark +sparseStorageBenchmarks benchmarkSelection = + bgroup + "sparse storage and packed kernels" + (sparseStorageBenchmark <$> sparseStorageDimensions benchmarkSelection) + +sparseStorageOnceBenchmarks :: BenchmarkSelection -> [OnceBenchmark] +sparseStorageOnceBenchmarks benchmarkSelection = + sparseStorageOnceBenchmark =<< sparseStorageDimensions benchmarkSelection + +sparseStorageDimensions :: BenchmarkSelection -> [Int] +sparseStorageDimensions benchmarkSelection = + [512] + <> [2048 | includeBroadMedium benchmarkSelection || includeBroadLarge benchmarkSelection] + <> [8192 | includeBroadLarge benchmarkSelection] + +sparseStorageBenchmark :: Int -> Benchmark +sparseStorageBenchmark dimension = + bgroup + ("n=" <> show dimension) + (renderPreparedBenchmark (prepareSparseStorageCase dimension) <$> sparseStorageRows) + +sparseStorageOnceBenchmark :: Int -> [OnceBenchmark] +sparseStorageOnceBenchmark dimension = + renderPreparedOnceBenchmark ("sparse storage and packed kernels.n=" <> show dimension <> ".") (prepareSparseStorageCase dimension) + <$> sparseStorageRows + +sparseStorageRows :: [PreparedBenchmarkRow SparseStorageCase] +sparseStorageRows = + [ PurePreparedBenchmarkRow "dense 32x32 -> COO" denseToCooWeight, + PurePreparedBenchmarkRow "dense 32x32 -> CSR" denseToCsrWeight, + PurePreparedBenchmarkRow "dense 32x32 -> CSC" denseToCscWeight, + PurePreparedBenchmarkRow "CSR -> COO" csrToCooWeight, + PurePreparedBenchmarkRow "CSR -> CSC" csrToCscWeight, + PurePreparedBenchmarkRow "CSR matvec" csrMatvecWeight, + PurePreparedBenchmarkRow "packed sparse apply" packedSparseWeight, + PurePreparedBenchmarkRow "graph Laplacian construction" graphLaplacianWeight + ] + +prepareSparseStorageCase :: Int -> BenchmarkSetup SparseStorageCase +prepareSparseStorageCase dimension = + BenchmarkSetup $ do + dense32 <- first show (fromListMatrix @32 @32 @Double (concat (bandedDenseRows 32))) + csrValue <- first (("banded sparse fixture failed: " <>)) (bandedSpdCSR dimension) + packedOperator <- first (("packed sparse fixture failed: " <>)) (packedSparseBenchmarkOperator dimension) + let vectorValues = denseBenchmarkVector dimension + pure + SparseStorageCase + { sparseStorageLabel = "n=" <> show dimension, + sparseStorageDense32 = dense32, + sparseStorageCSR = csrValue, + sparseStoragePackedOperator = packedOperator, + sparseStoragePackedVector = Unboxed.fromList vectorValues, + sparseStorageGraphVertices = [0 .. dimension - 1], + sparseStorageGraphEdges = duplicatePathEdges dimension + } + +denseToCooWeight :: SparseStorageCase -> BenchmarkWeight +denseToCooWeight benchmarkCase = + BenchmarkWeight (cooChecksum (denseToCOO (sparseStorageDense32 benchmarkCase))) + +denseToCsrWeight :: SparseStorageCase -> BenchmarkWeight +denseToCsrWeight benchmarkCase = + BenchmarkWeight (csrChecksum (denseToCSR (sparseStorageDense32 benchmarkCase))) + +denseToCscWeight :: SparseStorageCase -> BenchmarkWeight +denseToCscWeight benchmarkCase = + BenchmarkWeight (cscChecksum (denseToCSC (sparseStorageDense32 benchmarkCase))) + +csrToCooWeight :: SparseStorageCase -> BenchmarkWeight +csrToCooWeight benchmarkCase = + eitherBenchmarkWeight + (sparseStorageLabel benchmarkCase <> " CSR -> COO") + cooChecksum + (csrToCOO (sparseStorageCSR benchmarkCase)) + +csrToCscWeight :: SparseStorageCase -> BenchmarkWeight +csrToCscWeight benchmarkCase = + eitherBenchmarkWeight + (sparseStorageLabel benchmarkCase <> " CSR -> CSC") + cscChecksum + (csrToCSC (sparseStorageCSR benchmarkCase)) + +csrMatvecWeight :: SparseStorageCase -> BenchmarkWeight +csrMatvecWeight benchmarkCase = + eitherBenchmarkWeight + (sparseStorageLabel benchmarkCase <> " CSR matvec") + unboxedVectorChecksum + (csrMatVecVector (sparseStorageCSR benchmarkCase) (sparseStoragePackedVector benchmarkCase)) + +packedSparseWeight :: SparseStorageCase -> BenchmarkWeight +packedSparseWeight benchmarkCase = + eitherBenchmarkWeight + (sparseStorageLabel benchmarkCase <> " packed sparse apply") + unboxedVectorChecksum + (applyPackedSparseOperatorDense (sparseStoragePackedOperator benchmarkCase) (sparseStoragePackedVector benchmarkCase)) + +graphLaplacianWeight :: SparseStorageCase -> BenchmarkWeight +graphLaplacianWeight benchmarkCase = + eitherBenchmarkWeight + (sparseStorageLabel benchmarkCase <> " graph Laplacian construction") + csrChecksum + ( graphLaplacianCSR + (sparseStorageGraphVertices benchmarkCase) + (sparseStorageGraphEdges benchmarkCase) + ) + +duplicatePathEdges :: Int -> [GraphEdge Int] +duplicatePathEdges dimension = + concatMap + ( \vertexIndex -> + [ GraphEdge vertexIndex (vertexIndex + 1) 0.75, + GraphEdge (vertexIndex + 1) vertexIndex 0.25 + ] + ) + [0 .. dimension - 2] + +cooChecksum :: SparseCOO Double -> Double +cooChecksum cooValue = + fromIntegral (length (cooEntries cooValue)) + + sum ((\(rowIndex, columnIndex, entryValue) -> fromIntegral rowIndex + fromIntegral columnIndex + abs entryValue) <$> cooEntries cooValue) + +csrChecksum :: SparseCSR Double -> Double +csrChecksum csrValue = + fromIntegral (Unboxed.length (csrValuesVector csrValue)) + + fromIntegral (Unboxed.sum (csrRowOffsetsVector csrValue)) + + fromIntegral (Unboxed.sum (csrColumnIndicesVector csrValue)) + + unboxedVectorChecksum (csrValuesVector csrValue) + +cscChecksum :: SparseCSC Double -> Double +cscChecksum cscValue = + fromIntegral (Unboxed.length (cscValuesVector cscValue)) + + fromIntegral (Unboxed.sum (cscColumnOffsetsVector cscValue)) + + fromIntegral (Unboxed.sum (cscRowIndicesVector cscValue)) + + unboxedVectorChecksum (cscValuesVector cscValue) + +unboxedVectorChecksum :: Unboxed.Vector Double -> Double +unboxedVectorChecksum values = + Unboxed.foldl' (\acc value -> acc + abs value) 0.0 values diff --git a/bench/spectral/ProjectedBlock.hs b/bench/spectral/ProjectedBlock.hs new file mode 100644 --- /dev/null +++ b/bench/spectral/ProjectedBlock.hs @@ -0,0 +1,337 @@ +{-# LANGUAGE DataKinds #-} + +module ProjectedBlock + ( projectedBlockBenchmarks, + projectedBlockOnceBenchmarks, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Bifunctor (first) +import qualified Data.Vector.Unboxed as U +import Env (BenchmarkSelection) +import Fixtures + ( benchmarkSeedBlock, + denseOperator, + projectedBenchmarkDimension, + projectedBenchmarkRows, + projectedBlockBenchmarkCases, + ) +import Types + ( BenchmarkSetup (..), + BenchmarkWeight (..), + OnceBenchmark, + PreparedBenchmarkRow (..), + ProjectedBlockBenchmarkCase (..), + ProjectedBlockPreparedCase (..), + eitherBenchmarkWeight, + eigenpairsChecksum, + renderPreparedBenchmark, + renderPreparedOnceBenchmark, + ) +import Moonlight.LinAlg.Krylov + ( SpectrumEnd (..), + blockLanczosSymmetric, + defaultBlockLanczosConfig, + defaultLanczosConfig, + lanczosSymmetric, + mkPositiveCount, + withBlockLanczosBlockSize, + withBlockLanczosIterations, + withLanczosIterations, + ) +import Moonlight.LinAlg.Operator + ( LinearOperator, + OperatorSymmetry (..), + pathLaplacianLinearOperator, + ) +import Moonlight.LinAlg.Pure.Dense.Decomposition (symmetricEigenPairs) +import Moonlight.LinAlg.Pure.Krylov.Projected + ( ProjectedSubspace, + SymmetricProjectedOperator (..), + applySymmetricProjectedOperatorU, + projectedEigenpairs, + projectedEigenvalues, + projectedSubspaceDimension, + projectedSubspaceFromBlockLanczos, + projectedSubspaceFromLanczos, + projectedSubspaceOperator, + symmetricProjectedOperatorDimension, + ) +import Moonlight.LinAlg.Pure.Krylov.Selection (sortRawPairsForSpectrum) +import Moonlight.LinAlg.Native (selectedSymmetricBlockTridiagonalEigenRequestLapack) +import Moonlight.LinAlg.Spectral (EigenRequest (..)) +import Moonlight.LinAlg.Pure.Structured.BlockTridiagonal (SymmetricBlockTridiagonal) +import Test.Tasty.Bench (Benchmark, bgroup) +import Prelude + +data ProjectedTridiagonalBenchmarkCase = ProjectedTridiagonalBenchmarkCase + { projectedTridiagonalLabel :: !String, + projectedTridiagonalDimension :: !Int, + projectedTridiagonalIterations :: !Int, + projectedTridiagonalRequestedModes :: !Int + } + +data ProjectedTridiagonalPreparedCase = ProjectedTridiagonalPreparedCase + { projectedTridiagonalPreparedCase :: !ProjectedTridiagonalBenchmarkCase, + projectedTridiagonalPreparedOperator :: !(LinearOperator 'SelfAdjointOperator), + projectedTridiagonalPreparedSubspace :: !ProjectedSubspace + } + +instance NFData ProjectedTridiagonalPreparedCase where + rnf preparedCase = + projectedTridiagonalPreparedCase preparedCase + `seq` projectedTridiagonalPreparedOperator preparedCase + `seq` projectedTridiagonalPreparedSubspace preparedCase + `seq` () + +projectedBlockBenchmarks :: BenchmarkSelection -> Benchmark +projectedBlockBenchmarks benchmarkSelection = + bgroup + "projected structured eigensolve" + ( concatMap + (\benchmarkCase -> renderPreparedBenchmark (prepareProjectedTridiagonalCase benchmarkCase) <$> projectedTridiagonalRows benchmarkCase) + projectedTridiagonalBenchmarkCases + <> concatMap + (\benchmarkCase -> renderPreparedBenchmark (prepareProjectedBlockCase benchmarkCase) <$> projectedBlockRows benchmarkCase) + (projectedBlockBenchmarkCases benchmarkSelection) + ) + +projectedBlockOnceBenchmarks :: BenchmarkSelection -> [OnceBenchmark] +projectedBlockOnceBenchmarks benchmarkSelection = + concatMap + (\benchmarkCase -> renderPreparedOnceBenchmark "projected structured eigensolve." (prepareProjectedTridiagonalCase benchmarkCase) <$> projectedTridiagonalRows benchmarkCase) + projectedTridiagonalBenchmarkCases + <> concatMap + (\benchmarkCase -> renderPreparedOnceBenchmark "projected structured eigensolve." (prepareProjectedBlockCase benchmarkCase) <$> projectedBlockRows benchmarkCase) + (projectedBlockBenchmarkCases benchmarkSelection) + +projectedTridiagonalBenchmarkCases :: [ProjectedTridiagonalBenchmarkCase] +projectedTridiagonalBenchmarkCases = + [ProjectedTridiagonalBenchmarkCase "tridiagonal-path-512" 512 16 4] + +projectedTridiagonalRows :: ProjectedTridiagonalBenchmarkCase -> [PreparedBenchmarkRow ProjectedTridiagonalPreparedCase] +projectedTridiagonalRows benchmarkCase = + [ PurePreparedBenchmarkRow (projectedTridiagonalValuesBenchmarkLabel benchmarkCase) projectedTridiagonalValuesWeight, + PurePreparedBenchmarkRow (projectedTridiagonalPairsBenchmarkLabel benchmarkCase) projectedTridiagonalPairsWeight + ] + +projectedTridiagonalValuesBenchmarkLabel :: ProjectedTridiagonalBenchmarkCase -> String +projectedTridiagonalValuesBenchmarkLabel = + projectedTridiagonalBenchmarkLabel "values" + +projectedTridiagonalPairsBenchmarkLabel :: ProjectedTridiagonalBenchmarkCase -> String +projectedTridiagonalPairsBenchmarkLabel = + projectedTridiagonalBenchmarkLabel "pairs" + +projectedTridiagonalBenchmarkLabel :: String -> ProjectedTridiagonalBenchmarkCase -> String +projectedTridiagonalBenchmarkLabel requestLabel benchmarkCase = + projectedTridiagonalLabel benchmarkCase + <> " " + <> requestLabel + <> " n=" + <> show (projectedTridiagonalDimension benchmarkCase) + <> " m=" + <> show (projectedTridiagonalIterations benchmarkCase) + +prepareProjectedTridiagonalCase :: ProjectedTridiagonalBenchmarkCase -> BenchmarkSetup ProjectedTridiagonalPreparedCase +prepareProjectedTridiagonalCase benchmarkCase = + BenchmarkSetup $ do + iterationCount <- + first + (\err -> "invalid projected tridiagonal iteration count for " <> projectedTridiagonalLabel benchmarkCase <> ": " <> show err) + (mkPositiveCount (projectedTridiagonalIterations benchmarkCase)) + operatorValue <- + first + (\err -> "projected tridiagonal operator construction failed for " <> projectedTridiagonalLabel benchmarkCase <> ": " <> show err) + (pathLaplacianLinearOperator (projectedTridiagonalDimension benchmarkCase)) + subspace <- + first + (\err -> "projected tridiagonal decomposition failed for " <> projectedTridiagonalLabel benchmarkCase <> ": " <> show err) + ( projectedSubspaceFromLanczos + <$> lanczosSymmetric + (withLanczosIterations iterationCount defaultLanczosConfig) + operatorValue + (projectedSeedVector (projectedTridiagonalDimension benchmarkCase)) + ) + pure + ProjectedTridiagonalPreparedCase + { projectedTridiagonalPreparedCase = benchmarkCase, + projectedTridiagonalPreparedOperator = operatorValue, + projectedTridiagonalPreparedSubspace = subspace + } + +projectedTridiagonalValuesWeight :: ProjectedTridiagonalPreparedCase -> BenchmarkWeight +projectedTridiagonalValuesWeight preparedCase = + case + projectedEigenvalues + SmallestEigenvalues + (projectedTridiagonalRequestedModes (projectedTridiagonalPreparedCase preparedCase)) + (projectedTridiagonalPreparedOperator preparedCase) + (projectedTridiagonalPreparedSubspace preparedCase) of + Left err -> BenchmarkMeasurementFailure (projectedTridiagonalLabel (projectedTridiagonalPreparedCase preparedCase) <> " values: " <> show err) + Right values -> BenchmarkWeight (U.sum values) + +projectedTridiagonalPairsWeight :: ProjectedTridiagonalPreparedCase -> BenchmarkWeight +projectedTridiagonalPairsWeight preparedCase = + case + projectedEigenpairs + SmallestEigenvalues + (projectedTridiagonalRequestedModes (projectedTridiagonalPreparedCase preparedCase)) + (projectedTridiagonalPreparedOperator preparedCase) + (projectedTridiagonalPreparedSubspace preparedCase) of + Left err -> BenchmarkMeasurementFailure (projectedTridiagonalLabel (projectedTridiagonalPreparedCase preparedCase) <> " pairs: " <> show err) + Right pairs -> BenchmarkWeight (eigenpairsChecksum pairs) + +projectedBlockRows :: ProjectedBlockBenchmarkCase -> [PreparedBenchmarkRow ProjectedBlockPreparedCase] +projectedBlockRows benchmarkCase = + [ EffectfulPreparedBenchmarkRow (projectedBlockValuesBenchmarkLabel benchmarkCase) projectedBlockValuesWeight, + EffectfulPreparedBenchmarkRow (projectedBlockPairsBenchmarkLabel benchmarkCase) projectedBlockPairsWeight, + PurePreparedBenchmarkRow (projectedDenseOracleBenchmarkLabel benchmarkCase) projectedDenseOracleWeight + ] + +projectedBlockValuesBenchmarkLabel :: ProjectedBlockBenchmarkCase -> String +projectedBlockValuesBenchmarkLabel benchmarkCase = + projectedBlockBenchmarkLabel "values" benchmarkCase + +projectedBlockPairsBenchmarkLabel :: ProjectedBlockBenchmarkCase -> String +projectedBlockPairsBenchmarkLabel benchmarkCase = + projectedBlockBenchmarkLabel "pairs" benchmarkCase + +projectedDenseOracleBenchmarkLabel :: ProjectedBlockBenchmarkCase -> String +projectedDenseOracleBenchmarkLabel benchmarkCase = + projectedBlockBenchmarkLabel "generic dense oracle" benchmarkCase + +projectedBlockBenchmarkLabel :: String -> ProjectedBlockBenchmarkCase -> String +projectedBlockBenchmarkLabel requestLabel benchmarkCase = + projectedBenchmarkLabel benchmarkCase + <> " " + <> requestLabel + <> " profile=" + <> show (projectedBenchmarkSpectrumProfile benchmarkCase) + <> " n=" + <> show (projectedBenchmarkDimension benchmarkCase) + +prepareProjectedBlockCase :: ProjectedBlockBenchmarkCase -> BenchmarkSetup ProjectedBlockPreparedCase +prepareProjectedBlockCase benchmarkCase = + BenchmarkSetup $ do + iterationCount <- + first + (\err -> "invalid projected benchmark iteration count for " <> projectedBenchmarkLabel benchmarkCase <> ": " <> show err) + (mkPositiveCount (projectedBenchmarkIterations benchmarkCase)) + blockSize <- + first + (\err -> "invalid projected benchmark block size for " <> projectedBenchmarkLabel benchmarkCase <> ": " <> show err) + (mkPositiveCount (projectedBenchmarkBlockSize benchmarkCase)) + operatorValue <- + first + (\err -> "projected benchmark operator construction failed for " <> projectedBenchmarkLabel benchmarkCase <> ": " <> err) + (denseOperator (projectedBenchmarkRows benchmarkCase)) + let operatorDimension = projectedBenchmarkDimension benchmarkCase + seedBlock = benchmarkSeedBlock operatorDimension (projectedBenchmarkBlockSize benchmarkCase) + blockConfig = + withBlockLanczosBlockSize + blockSize + (withBlockLanczosIterations iterationCount defaultBlockLanczosConfig) + subspace <- + first + (\err -> "projected benchmark decomposition failed for " <> projectedBenchmarkLabel benchmarkCase <> ": " <> show err) + (projectedSubspaceFromBlockLanczos <$> blockLanczosSymmetric blockConfig operatorValue seedBlock) + pure + ProjectedBlockPreparedCase + { projectedPreparedCase = benchmarkCase, + projectedPreparedOperator = operatorValue, + projectedPreparedSubspace = subspace, + projectedPreparedDimension = projectedSubspaceDimension subspace + } + +projectedBlockValuesWeight :: ProjectedBlockPreparedCase -> IO BenchmarkWeight +projectedBlockValuesWeight preparedCase = + case nativeProjectedBlockOperator preparedCase of + Left err -> pure (BenchmarkMeasurementFailure (projectedBenchmarkLabel (projectedPreparedCase preparedCase) <> " values: " <> err)) + Right blockValue -> + case mkPositiveCount (projectedBenchmarkRequestedModes (projectedPreparedCase preparedCase)) of + Left err -> pure (BenchmarkMeasurementFailure (projectedBenchmarkLabel (projectedPreparedCase preparedCase) <> " values: " <> show err)) + Right countValue -> + eitherBenchmarkWeight + (projectedBenchmarkLabel (projectedPreparedCase preparedCase) <> " values") + U.sum + <$> selectedSymmetricBlockTridiagonalEigenRequestLapack + (EigenvaluesRequest SmallestEigenvalues countValue) + blockValue + +projectedBlockPairsWeight :: ProjectedBlockPreparedCase -> IO BenchmarkWeight +projectedBlockPairsWeight preparedCase = + case nativeProjectedBlockOperator preparedCase of + Left err -> pure (BenchmarkMeasurementFailure (projectedBenchmarkLabel (projectedPreparedCase preparedCase) <> " pairs: " <> err)) + Right blockValue -> + case mkPositiveCount (projectedBenchmarkRequestedModes (projectedPreparedCase preparedCase)) of + Left err -> pure (BenchmarkMeasurementFailure (projectedBenchmarkLabel (projectedPreparedCase preparedCase) <> " pairs: " <> show err)) + Right countValue -> + eitherBenchmarkWeight + (projectedBenchmarkLabel (projectedPreparedCase preparedCase) <> " pairs") + eigenpairsChecksum + <$> selectedSymmetricBlockTridiagonalEigenRequestLapack + (EigenpairsRequest SmallestEigenvalues countValue) + blockValue + +nativeProjectedBlockOperator :: ProjectedBlockPreparedCase -> Either String SymmetricBlockTridiagonal +nativeProjectedBlockOperator preparedCase = + case projectedSubspaceOperator (projectedPreparedSubspace preparedCase) of + BlockTridiagonalProjectedOperator blockValue -> Right blockValue + TridiagonalProjectedOperator _ -> Left "expected block-tridiagonal projected operator" + +projectedDenseOracleWeight :: ProjectedBlockPreparedCase -> BenchmarkWeight +projectedDenseOracleWeight preparedCase = + case projectedDenseOraclePairs SmallestEigenvalues (projectedBenchmarkRequestedModes (projectedPreparedCase preparedCase)) (projectedSubspaceOperator (projectedPreparedSubspace preparedCase)) of + Left err -> BenchmarkMeasurementFailure (projectedBenchmarkLabel (projectedPreparedCase preparedCase) <> " dense oracle: " <> show err) + Right pairs -> BenchmarkWeight (projectedDenseOracleChecksum pairs) + +projectedDenseOraclePairs :: + SpectrumEnd -> + Int -> + SymmetricProjectedOperator -> + Either String [(Double, [Double])] +projectedDenseOraclePairs spectrumEnd requestedModes projectedOperator = + let projectedDimension = symmetricProjectedOperatorDimension projectedOperator + in if requestedModes <= 0 + then Left "projected dense oracle requested count must be positive" + else + if requestedModes > projectedDimension + then Left "projected dense oracle requested count exceeds projected dimension" + else do + projectedRows <- projectedOperatorDenseRows projectedOperator + rawPairs <- first show (symmetricEigenPairs projectedDimension projectedRows) + Right (take requestedModes (sortRawPairsForSpectrum spectrumEnd rawPairs)) + +projectedOperatorDenseRows :: SymmetricProjectedOperator -> Either String [[Double]] +projectedOperatorDenseRows projectedOperator = + let projectedDimension = symmetricProjectedOperatorDimension projectedOperator + in do + imageColumns <- + traverse + (\coordinateIndex -> first show (applySymmetricProjectedOperatorU projectedOperator (coordinateBasisVector projectedDimension coordinateIndex))) + [0 .. projectedDimension - 1] + traverse (projectedDenseRow imageColumns) [0 .. projectedDimension - 1] + +projectedDenseRow :: [U.Vector Double] -> Int -> Either String [Double] +projectedDenseRow imageColumns rowIndex = + traverse (projectedColumnEntry rowIndex) imageColumns + +projectedColumnEntry :: Int -> U.Vector Double -> Either String Double +projectedColumnEntry rowIndex columnValue = + case columnValue U.!? rowIndex of + Just entryValue -> Right entryValue + Nothing -> Left "projected dense oracle column dimension mismatch" + +coordinateBasisVector :: Int -> Int -> U.Vector Double +coordinateBasisVector dimension columnIndex = + U.generate dimension (\rowIndex -> if rowIndex == columnIndex then 1.0 else 0.0) + +projectedSeedVector :: Int -> U.Vector Double +projectedSeedVector dimension = + U.generate dimension (\indexValue -> if indexValue == 0 then 1.0 else 1.0 / fromIntegral (indexValue + 1)) + +projectedDenseOracleChecksum :: [(Double, [Double])] -> Double +projectedDenseOracleChecksum pairs = + sum ((\(eigenvalue, eigenvector) -> eigenvalue + sum (abs <$> eigenvector)) <$> pairs) diff --git a/bench/spectral/SparseKrylov.hs b/bench/spectral/SparseKrylov.hs new file mode 100644 --- /dev/null +++ b/bench/spectral/SparseKrylov.hs @@ -0,0 +1,87 @@ +module SparseKrylov + ( sparseKrylovBenchmarks, + sparseKrylovOnceBenchmarks, + ) +where + +import Data.Bifunctor (first) +import qualified Data.Vector.Unboxed as U +import Fixtures + ( pathLaplacianTridiagonal, + sparseKrylovBenchmarkCases, + ) +import Env (BenchmarkSelection) +import Types + ( BenchmarkSetup (..), + OnceBenchmark (..), + SparseKrylovBenchmarkCase (..), + SparseKrylovPreparedCase (..), + BenchmarkWeight (..), + prepareBenchmarkSetup, + benchmarkWeightEither, + ) +import Moonlight.LinAlg.Krylov + ( mkPositiveCount, + SpectrumEnd (..), + ) +import Moonlight.LinAlg.Operator (symmetricTridiagonalLinearOperator) +import Moonlight.LinAlg.Spectral + ( defaultEigenSolveConfig, + EigenRequest (..), + solveEigenRequest, + ) +import Test.Tasty.Bench (Benchmark, bench, bgroup, env, nf) +import Prelude + +sparseKrylovBenchmarks :: BenchmarkSelection -> Benchmark +sparseKrylovBenchmarks benchmarkSelection = + bgroup + "selected tridiagonal eigenvalue solve" + (sparseKrylovBenchmark <$> sparseKrylovBenchmarkCases benchmarkSelection) + +sparseKrylovOnceBenchmarks :: BenchmarkSelection -> [OnceBenchmark] +sparseKrylovOnceBenchmarks benchmarkSelection = + sparseKrylovOnceBenchmark <$> sparseKrylovBenchmarkCases benchmarkSelection + +sparseKrylovBenchmark :: SparseKrylovBenchmarkCase -> Benchmark +sparseKrylovBenchmark benchmarkCase = + env (prepareBenchmarkSetup (prepareSparseKrylovCase benchmarkCase)) $ \preparedCase -> + bench (sparseBenchmarkLabel benchmarkCase) (nf sparseKrylovWeight preparedCase) + +sparseKrylovOnceBenchmark :: SparseKrylovBenchmarkCase -> OnceBenchmark +sparseKrylovOnceBenchmark benchmarkCase = + OnceBenchmark + { onceBenchmarkLabel = "selected tridiagonal eigenvalue solve." <> sparseBenchmarkLabel benchmarkCase, + onceBenchmarkAction = + pure + (runBenchmarkSetup (prepareSparseKrylovCase benchmarkCase) >>= benchmarkWeightEither . sparseKrylovWeight) + } + +prepareSparseKrylovCase :: SparseKrylovBenchmarkCase -> BenchmarkSetup SparseKrylovPreparedCase +prepareSparseKrylovCase benchmarkCase = + BenchmarkSetup $ do + tridiagonalOperator <- + first + (\err -> "benchmark tridiagonal fixture failed for " <> sparseBenchmarkLabel benchmarkCase <> ": " <> err) + (pathLaplacianTridiagonal (sparseBenchmarkDimension benchmarkCase)) + pure + SparseKrylovPreparedCase + { sparsePreparedLabel = sparseBenchmarkLabel benchmarkCase, + sparsePreparedRequestedModes = sparseBenchmarkRequestedModes benchmarkCase, + sparsePreparedTridiagonal = tridiagonalOperator + } + +sparseKrylovWeight :: SparseKrylovPreparedCase -> BenchmarkWeight +sparseKrylovWeight preparedCase = + case + do + requestedCount <- first show (mkPositiveCount (sparsePreparedRequestedModes preparedCase)) + first + show + ( solveEigenRequest + defaultEigenSolveConfig + (symmetricTridiagonalLinearOperator (sparsePreparedTridiagonal preparedCase)) + (EigenvaluesRequest SmallestEigenvalues requestedCount) + ) of + Left err -> BenchmarkMeasurementFailure (sparsePreparedLabel preparedCase <> ": " <> err) + Right eigenvalues -> BenchmarkWeight (sum (U.toList eigenvalues)) diff --git a/bench/spectral/SpectralDispatch.hs b/bench/spectral/SpectralDispatch.hs new file mode 100644 --- /dev/null +++ b/bench/spectral/SpectralDispatch.hs @@ -0,0 +1,281 @@ +{-# LANGUAGE DataKinds #-} + +module SpectralDispatch + ( spectralDispatchBenchmarks, + spectralDispatchOnceBenchmarks, + ) +where + +import Data.Bifunctor (first) +import Control.DeepSeq (NFData (..)) +import qualified Data.Vector.Unboxed as U +import Env (BenchmarkSelection (..)) +import Fixtures + ( bandedSpdCSR, + diagonalBenchmarkValues, + genericBenchmarkTridiagonal, + reducibleBenchmarkTridiagonal, + ) +import Types + ( BenchmarkSetup (..), + BenchmarkWeight, + OnceBenchmark (..), + benchmarkWeightEither, + eigenpairsChecksum, + eigenpairsResidualValidationChecksum, + eitherBenchmarkWeight, + prepareBenchmarkSetup, + ) +import Moonlight.LinAlg.Krylov + ( SpectrumEnd (..), + defaultLanczosConfig, + mkPositiveCount, + withLanczosIterations, + ) +import Moonlight.LinAlg.Operator + ( LinearOperator, + OperatorSymmetry (..), + diagonalLinearOperator, + operatorDimension, + pathLaplacianLinearOperator, + runOperatorU, + selfAdjointCSRLinearOperator, + symmetricTridiagonalLinearOperator, + ) +import Moonlight.LinAlg.Pure.Krylov.SelectedTridiagonal (symmetricTridiagonalFromCSR) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal (symmetricTridiagonalDimension) +import Moonlight.LinAlg.Spectral + ( Eigenpairs, + EigenRequest (..), + EigenSolveConfig, + defaultEigenSolveConfig, + eigenpairCount, + solveEigenRequest, + withEigenFallbackLanczosConfig, + withEigenFallbackInitialVector, + ) +import Test.Tasty.Bench (Benchmark, bench, bgroup, env, nf) +import Prelude + +data SpectralDispatchCase = SpectralDispatchCase + { spectralCaseLabel :: !String, + spectralCaseDimension :: !Int, + spectralCaseRequestedModes :: !Int, + spectralCaseKind :: !SpectralDispatchKind + } + +data SpectralDispatchKind + = PathDispatch + | DiagonalDispatch + | GenericTridiagonalDispatch + | ReducibleTridiagonalDispatch + | GenericCSRDispatch + deriving stock (Eq, Show) + +data SpectralPreparedCase = SpectralPreparedCase + { spectralPreparedLabel :: !String, + spectralPreparedRequestedModes :: !Int, + spectralPreparedOperator :: !(LinearOperator 'SelfAdjointOperator), + spectralPreparedConfig :: !EigenSolveConfig + } + +data SpectralResidualPreparedCase = SpectralResidualPreparedCase + { spectralResidualPreparedLabel :: !String, + spectralResidualPreparedOperator :: !(LinearOperator 'SelfAdjointOperator), + spectralResidualPreparedPairs :: !Eigenpairs + } + +instance NFData SpectralPreparedCase where + rnf preparedCase = + spectralPreparedLabel preparedCase + `seq` spectralPreparedRequestedModes preparedCase + `seq` spectralPreparedOperator preparedCase + `seq` spectralPreparedConfig preparedCase + `seq` () + +instance NFData SpectralResidualPreparedCase where + rnf preparedCase = + spectralResidualPreparedLabel preparedCase + `seq` spectralResidualPreparedOperator preparedCase + `seq` spectralResidualPreparedPairs preparedCase + `seq` () + +spectralDispatchBenchmarks :: BenchmarkSelection -> Benchmark +spectralDispatchBenchmarks benchmarkSelection = + bgroup + "spectral demand dispatch" + (spectralDispatchBenchmark <$> spectralDispatchCases benchmarkSelection) + +spectralDispatchOnceBenchmarks :: BenchmarkSelection -> [OnceBenchmark] +spectralDispatchOnceBenchmarks benchmarkSelection = + spectralDispatchOnceBenchmark =<< spectralDispatchCases benchmarkSelection + +spectralDispatchCases :: BenchmarkSelection -> [SpectralDispatchCase] +spectralDispatchCases benchmarkSelection = + [ SpectralDispatchCase "path-values-pairs-1024" 1024 4 PathDispatch, + SpectralDispatchCase "diagonal-values-pairs-4096" 4096 4 DiagonalDispatch, + SpectralDispatchCase "generic-tridiagonal-values-pairs-512" 512 4 GenericTridiagonalDispatch, + SpectralDispatchCase "reducible-tridiagonal-values-pairs-512" 512 4 ReducibleTridiagonalDispatch, + SpectralDispatchCase "generic-csr-fallback-values-pairs-96" 96 4 GenericCSRDispatch + ] + <> [SpectralDispatchCase "generic-csr-fallback-values-pairs-192" 192 6 GenericCSRDispatch | includeBroadMedium benchmarkSelection || includeBroadLarge benchmarkSelection] + <> ( if includeBroadLarge benchmarkSelection + then + [ SpectralDispatchCase "generic-csr-dense-fallback-values-pairs-384" 384 8 GenericCSRDispatch, + SpectralDispatchCase "generic-csr-dense-fallback-values-pairs-512" 512 8 GenericCSRDispatch, + SpectralDispatchCase "generic-csr-high-demand-dense-values-pairs-513" 513 513 GenericCSRDispatch, + SpectralDispatchCase "generic-csr-high-demand-dense-values-pairs-1024" 1024 1024 GenericCSRDispatch + ] + else [] + ) + +spectralDispatchBenchmark :: SpectralDispatchCase -> Benchmark +spectralDispatchBenchmark benchmarkCase = + env (prepareBenchmarkSetup (prepareSpectralDispatchCase benchmarkCase)) $ \preparedCase -> + bgroup + (spectralCaseLabel benchmarkCase) + [ bench "construction/classification" (nf spectralConstructionClassificationWeight benchmarkCase), + bench "values" (nf spectralValuesWeight preparedCase), + bench "pairs" (nf spectralPairsWeight preparedCase), + env (prepareBenchmarkSetup (prepareSpectralResidualCase benchmarkCase)) $ + \residualCase -> + bench "residual validation" (nf spectralResidualValidationWeight residualCase) + ] + +spectralDispatchOnceBenchmark :: SpectralDispatchCase -> [OnceBenchmark] +spectralDispatchOnceBenchmark benchmarkCase = + [ spectralDispatchOnceBenchmarkRow benchmarkCase "values" spectralValuesWeight, + spectralDispatchOnceBenchmarkRow benchmarkCase "pairs" spectralPairsWeight + ] + +spectralDispatchOnceBenchmarkRow :: SpectralDispatchCase -> String -> (SpectralPreparedCase -> BenchmarkWeight) -> OnceBenchmark +spectralDispatchOnceBenchmarkRow benchmarkCase rowLabel measure = + OnceBenchmark + { onceBenchmarkLabel = "spectral demand dispatch." <> spectralCaseLabel benchmarkCase <> "." <> rowLabel, + onceBenchmarkAction = + pure + (runBenchmarkSetup (prepareSpectralDispatchCase benchmarkCase) >>= benchmarkWeightEither . measure) + } + +prepareSpectralDispatchCase :: SpectralDispatchCase -> BenchmarkSetup SpectralPreparedCase +prepareSpectralDispatchCase benchmarkCase = + BenchmarkSetup $ do + operatorValue <- spectralOperator benchmarkCase + pure + SpectralPreparedCase + { spectralPreparedLabel = spectralCaseLabel benchmarkCase, + spectralPreparedRequestedModes = spectralCaseRequestedModes benchmarkCase, + spectralPreparedOperator = operatorValue, + spectralPreparedConfig = spectralConfig benchmarkCase + } + +spectralOperator :: SpectralDispatchCase -> Either String (LinearOperator 'SelfAdjointOperator) +spectralOperator benchmarkCase = + case spectralCaseKind benchmarkCase of + PathDispatch -> + first show (pathLaplacianLinearOperator (spectralCaseDimension benchmarkCase)) + DiagonalDispatch -> + first show (diagonalLinearOperator (U.fromList (diagonalBenchmarkValues (spectralCaseDimension benchmarkCase)))) + GenericTridiagonalDispatch -> + symmetricTridiagonalLinearOperator <$> genericBenchmarkTridiagonal (spectralCaseDimension benchmarkCase) + ReducibleTridiagonalDispatch -> + symmetricTridiagonalLinearOperator <$> reducibleBenchmarkTridiagonal (spectralCaseDimension benchmarkCase) + GenericCSRDispatch -> + bandedSpdCSR (spectralCaseDimension benchmarkCase) >>= first show . selfAdjointCSRLinearOperator + +spectralConfig :: SpectralDispatchCase -> EigenSolveConfig +spectralConfig benchmarkCase = + let fallbackIterations = max 8 (min (spectralCaseDimension benchmarkCase) 32) + in case mkPositiveCount fallbackIterations of + Left _ -> defaultEigenSolveConfig + Right iterationCount -> + withEigenFallbackInitialVector (seedVector (spectralCaseDimension benchmarkCase)) + ( withEigenFallbackLanczosConfig + (withLanczosIterations iterationCount defaultLanczosConfig) + defaultEigenSolveConfig + ) + +spectralValuesWeight :: SpectralPreparedCase -> BenchmarkWeight +spectralValuesWeight preparedCase = + eitherBenchmarkWeight + (spectralPreparedLabel preparedCase <> " values") + U.sum + ( do + requestedCount <- first show (mkPositiveCount (spectralPreparedRequestedModes preparedCase)) + first + show + ( solveEigenRequest + (spectralPreparedConfig preparedCase) + (spectralPreparedOperator preparedCase) + (EigenvaluesRequest SmallestEigenvalues requestedCount) + ) + ) + +spectralPairsWeight :: SpectralPreparedCase -> BenchmarkWeight +spectralPairsWeight preparedCase = + eitherBenchmarkWeight + (spectralPreparedLabel preparedCase <> " pairs") + eigenpairsChecksum + (spectralPairsResult preparedCase) + +spectralPairsResult :: SpectralPreparedCase -> Either String Eigenpairs +spectralPairsResult preparedCase = do + requestedCount <- first show (mkPositiveCount (spectralPreparedRequestedModes preparedCase)) + first + show + ( solveEigenRequest + (spectralPreparedConfig preparedCase) + (spectralPreparedOperator preparedCase) + (EigenpairsRequest SmallestEigenvalues requestedCount) + ) + +prepareSpectralResidualCase :: SpectralDispatchCase -> BenchmarkSetup SpectralResidualPreparedCase +prepareSpectralResidualCase benchmarkCase = + BenchmarkSetup $ do + preparedCase <- runBenchmarkSetup (prepareSpectralDispatchCase benchmarkCase) + pairs <- spectralPairsResult preparedCase + eigenpairsChecksum pairs `seq` + pure + SpectralResidualPreparedCase + { spectralResidualPreparedLabel = spectralPreparedLabel preparedCase, + spectralResidualPreparedOperator = spectralPreparedOperator preparedCase, + spectralResidualPreparedPairs = pairs + } + +spectralConstructionClassificationWeight :: SpectralDispatchCase -> BenchmarkWeight +spectralConstructionClassificationWeight benchmarkCase = + eitherBenchmarkWeight + (spectralCaseLabel benchmarkCase <> " construction/classification") + id + ( do + operatorValue <- spectralOperator benchmarkCase + classificationChecksum <- spectralClassificationChecksum benchmarkCase + pure (fromIntegral (operatorDimension operatorValue) + classificationChecksum) + ) + +spectralClassificationChecksum :: SpectralDispatchCase -> Either String Double +spectralClassificationChecksum benchmarkCase = + case spectralCaseKind benchmarkCase of + GenericCSRDispatch -> do + csrValue <- bandedSpdCSR (spectralCaseDimension benchmarkCase) + case symmetricTridiagonalFromCSR csrValue of + Left err -> Left (show err) + Right (Left _) -> Right 0.0 + Right (Right tridiagonalValue) -> Right (fromIntegral (symmetricTridiagonalDimension tridiagonalValue)) + _ -> Right 0.0 + +spectralResidualValidationWeight :: SpectralResidualPreparedCase -> BenchmarkWeight +spectralResidualValidationWeight residualCase = + eitherBenchmarkWeight + (spectralResidualPreparedLabel residualCase <> " residual validation") + id + ( ( + fromIntegral (eigenpairCount (spectralResidualPreparedPairs residualCase)) + ) + <$> eigenpairsResidualValidationChecksum + (runOperatorU (spectralResidualPreparedOperator residualCase)) + (spectralResidualPreparedPairs residualCase) + ) + +seedVector :: Int -> U.Vector Double +seedVector dimension = + U.generate dimension (\indexValue -> if indexValue == 0 then 1.0 else 1.0 / fromIntegral (indexValue + 1)) diff --git a/bench/statics/GeometryStatics.hs b/bench/statics/GeometryStatics.hs new file mode 100644 --- /dev/null +++ b/bench/statics/GeometryStatics.hs @@ -0,0 +1,256 @@ +{-# LANGUAGE DataKinds #-} + +module GeometryStatics + ( geometryStaticsBenchmarks, + geometryStaticsOnceBenchmarks, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Bifunctor (first) +import Data.List.NonEmpty qualified as NonEmpty +import Data.Map.Strict qualified as Map +import Env (BenchmarkSelection (..)) +import Fixtures + ( staticsBenchmarkNetwork, + ) +import Types + ( BenchmarkSetup (..), + BenchmarkWeight (..), + OnceBenchmark, + PreparedBenchmarkRow (..), + eitherBenchmarkWeight, + renderPreparedBenchmark, + renderPreparedOnceBenchmark, + ) +import Moonlight.LinAlg.Dense + ( Matrix, + Vector, + dynMatrixToList, + dynVectorToList, + toListMatrix, + toListVector, + ) +import Moonlight.LinAlg.Geometry + ( AABB, + Symmetric3 (..), + Vec3 (..), + aabbRadius, + crossVec3, + dotVec3, + eigendecomposeSymmetric3, + expandAabb, + magnitudeVec3, + normalizeVec3, + symmetricAabb, + symmetric3Entries, + translateAabb, + unionAabb, + vec3ToList, + ) +import Moonlight.LinAlg.Statics + ( CompiledEquilibrium, + EquilibriumResult (..), + EquilibriumSolution (..), + EquilibriumViolation (..), + ForceNetwork, + assembleEquilibriumEquations, + checkEquilibrium, + compiledCoefficientMatrix, + compiledNodeOrder, + compiledRightHandSide, + compiledUnknownOrder, + ) +import Test.Tasty.Bench (Benchmark, bgroup) +import Prelude + +data GeometryStaticsCase = GeometryStaticsCase + { geometryStaticsLabel :: !String, + geometryVectors :: ![Vec3], + geometryBoxes :: ![AABB], + geometrySymmetricTensors :: ![Symmetric3 Double], + geometryStaticsNetwork :: !ForceNetwork + } + +instance NFData GeometryStaticsCase where + rnf benchmarkCase = + rnf (geometryStaticsLabel benchmarkCase) + `seq` rnf (vec3ToList =<< geometryVectors benchmarkCase) + `seq` rnf (aabbRadius <$> geometryBoxes benchmarkCase) + `seq` rnf (symmetric3Entries =<< geometrySymmetricTensors benchmarkCase) + `seq` geometryStaticsNetwork benchmarkCase + `seq` () + +geometryStaticsBenchmarks :: BenchmarkSelection -> Benchmark +geometryStaticsBenchmarks benchmarkSelection = + bgroup + "geometry and statics" + (geometryStaticsBenchmark <$> geometryStaticsSpans benchmarkSelection) + +geometryStaticsOnceBenchmarks :: BenchmarkSelection -> [OnceBenchmark] +geometryStaticsOnceBenchmarks benchmarkSelection = + geometryStaticsOnceBenchmark =<< geometryStaticsSpans benchmarkSelection + +geometryStaticsSpans :: BenchmarkSelection -> [Int] +geometryStaticsSpans benchmarkSelection = + [8] + <> [24 | includeBroadMedium benchmarkSelection || includeBroadLarge benchmarkSelection] + <> [64 | includeBroadLarge benchmarkSelection] + +geometryStaticsBenchmark :: Int -> Benchmark +geometryStaticsBenchmark spanCount = + bgroup + ("spans=" <> show spanCount) + (renderPreparedBenchmark (prepareGeometryStaticsCase spanCount) <$> geometryStaticsRows) + +geometryStaticsOnceBenchmark :: Int -> [OnceBenchmark] +geometryStaticsOnceBenchmark spanCount = + renderPreparedOnceBenchmark ("geometry and statics.spans=" <> show spanCount <> ".") (prepareGeometryStaticsCase spanCount) + <$> geometryStaticsRows + +geometryStaticsRows :: [PreparedBenchmarkRow GeometryStaticsCase] +geometryStaticsRows = + [ PurePreparedBenchmarkRow "Vec3 normalize/cross batch" vec3BatchWeight, + PurePreparedBenchmarkRow "AABB union/translate batch" aabbBatchWeight, + PurePreparedBenchmarkRow "Symmetric3 eigendecompose batch" symmetric3EigenBatchWeight, + PurePreparedBenchmarkRow "assemble equilibrium" staticsAssembleWeight, + PurePreparedBenchmarkRow "check equilibrium" staticsSolveWeight + ] + +prepareGeometryStaticsCase :: Int -> BenchmarkSetup GeometryStaticsCase +prepareGeometryStaticsCase spanCount = + BenchmarkSetup $ do + networkValue <- first (("statics fixture failed: " <>)) (staticsBenchmarkNetwork spanCount) + boxes <- boxFixture (spanCount * 16) + pure + GeometryStaticsCase + { geometryStaticsLabel = "spans=" <> show spanCount, + geometryVectors = vectorFixture (spanCount * 32), + geometryBoxes = boxes, + geometrySymmetricTensors = symmetric3Fixture (spanCount * 16), + geometryStaticsNetwork = networkValue + } + +vectorFixture :: Int -> [Vec3] +vectorFixture count = + [ Vec3 + (1.0 + fromIntegral (indexValue `mod` 13)) + (2.0 + fromIntegral (indexValue `mod` 17) / 2.0) + (3.0 + fromIntegral (indexValue `mod` 19) / 3.0) + | indexValue <- [0 .. count - 1] + ] + +boxFixture :: Int -> Either String [AABB] +boxFixture count = + traverse boxAt [0 .. count - 1] + where + boxAt :: Int -> Either String AABB + boxAt indexValue = + let translationVector = + Vec3 + (fromIntegral indexValue * 0.01) + (fromIntegral (indexValue `mod` 11) * 0.02) + 0.0 + halfY = 1.0 + fromIntegral (indexValue `mod` 5) * 0.1 + in maybe + (Left "AABB fixture half-extents violated constructor contract") + (Right . translateAabb translationVector) + (symmetricAabb 1.0 halfY 1.5) + +symmetric3Fixture :: Int -> [Symmetric3 Double] +symmetric3Fixture count = + tensorAt <$> [0 .. count - 1] + where + tensorAt :: Int -> Symmetric3 Double + tensorAt indexValue = + let xValue = 1.0 + fromIntegral (indexValue `mod` 7) * 0.125 + yValue = 2.0 + fromIntegral (indexValue `mod` 11) * 0.0625 + zValue = 3.0 + fromIntegral (indexValue `mod` 13) * 0.03125 + couplingScale = fromIntegral (indexValue `mod` 5) * 0.01 + in Symmetric3 + { sym3XX = xValue, + sym3XY = couplingScale, + sym3XZ = -0.5 * couplingScale, + sym3YY = yValue, + sym3YZ = 0.25 * couplingScale, + sym3ZZ = zValue + } + +vec3BatchWeight :: GeometryStaticsCase -> BenchmarkWeight +vec3BatchWeight benchmarkCase = + eitherBenchmarkWeight + (geometryStaticsLabel benchmarkCase <> " Vec3 normalize/cross batch") + vectorBatchChecksum + ( do + normalized <- traverse normalizeVec3 (geometryVectors benchmarkCase) + pure (zipWith crossVec3 normalized (drop 1 normalized), normalized) + ) + +aabbBatchWeight :: GeometryStaticsCase -> BenchmarkWeight +aabbBatchWeight benchmarkCase = + eitherBenchmarkWeight + (geometryStaticsLabel benchmarkCase <> " AABB union/translate batch") + aabbRadius + ( do + seedBox <- maybeToEither "seed AABB half-extents violated constructor contract" (symmetricAabb 1.0 1.0 1.0) + expandedBoxes <- traverse (maybeToEither "AABB expansion inverted a box" . expandAabb 0.05) (geometryBoxes benchmarkCase) + pure (foldr unionAabb seedBox expandedBoxes) + ) + +symmetric3EigenBatchWeight :: GeometryStaticsCase -> BenchmarkWeight +symmetric3EigenBatchWeight benchmarkCase = + eitherBenchmarkWeight + (geometryStaticsLabel benchmarkCase <> " Symmetric3 eigendecompose batch") + symmetric3EigenBatchChecksum + (traverse eigendecomposeSymmetric3 (geometrySymmetricTensors benchmarkCase)) + +staticsAssembleWeight :: GeometryStaticsCase -> BenchmarkWeight +staticsAssembleWeight benchmarkCase = + eitherBenchmarkWeight + (geometryStaticsLabel benchmarkCase <> " assemble equilibrium") + compiledEquilibriumChecksum + (assembleEquilibriumEquations (geometryStaticsNetwork benchmarkCase)) + +staticsSolveWeight :: GeometryStaticsCase -> BenchmarkWeight +staticsSolveWeight benchmarkCase = + eitherBenchmarkWeight + (geometryStaticsLabel benchmarkCase <> " check equilibrium") + equilibriumResultChecksum + (checkEquilibrium (geometryStaticsNetwork benchmarkCase)) + +vectorBatchChecksum :: ([Vec3], [Vec3]) -> Double +vectorBatchChecksum (crossVectors, normalizedVectors) = + sum (magnitudeVec3 <$> crossVectors) + + sum (zipWith dotVec3 normalizedVectors (drop 1 normalizedVectors)) + +symmetric3EigenBatchChecksum :: [(Vector 3 Double, Matrix 3 3 Double)] -> Double +symmetric3EigenBatchChecksum = + sum . fmap symmetric3EigenChecksum + +symmetric3EigenChecksum :: (Vector 3 Double, Matrix 3 3 Double) -> Double +symmetric3EigenChecksum (eigenvalues, eigenvectors) = + sum (abs <$> toListVector eigenvalues) + sum (abs <$> toListMatrix eigenvectors) + +compiledEquilibriumChecksum :: CompiledEquilibrium -> Double +compiledEquilibriumChecksum compiledValue = + fromIntegral (length (compiledNodeOrder compiledValue)) + + fromIntegral (length (compiledUnknownOrder compiledValue)) + + sum (abs <$> dynMatrixToList (compiledCoefficientMatrix compiledValue)) + + sum (abs <$> dynVectorToList (compiledRightHandSide compiledValue)) + +equilibriumResultChecksum :: EquilibriumResult -> Double +equilibriumResultChecksum resultValue = + case resultValue of + InEquilibrium solutionValue -> equilibriumSolutionChecksum solutionValue + Disequilibrium violations -> sum (violationResidualMagnitude <$> NonEmpty.toList violations) + +equilibriumSolutionChecksum :: EquilibriumSolution -> Double +equilibriumSolutionChecksum solutionValue = + fromIntegral (Map.size (equilibriumMemberForces solutionValue)) + + fromIntegral (Map.size (equilibriumReactionForces solutionValue)) + + fromIntegral (Map.size (equilibriumResidualForces solutionValue)) + + sum (abs <$> Map.elems (equilibriumMemberForces solutionValue)) + +maybeToEither :: err -> Maybe value -> Either err value +maybeToEither failureValue = + maybe (Left failureValue) Right diff --git a/bench/support/Env.hs b/bench/support/Env.hs new file mode 100644 --- /dev/null +++ b/bench/support/Env.hs @@ -0,0 +1,90 @@ +module Env + ( BenchmarkSelection (..), + benchmarkNotice, + readBenchmarkSelection, + ) +where + +import System.Environment (lookupEnv) +import Prelude + +data BenchmarkSelection = BenchmarkSelection + { includeSparseLarge :: !Bool, + includeSparse100k :: !Bool, + includeBroadMedium :: !Bool, + includeBroadLarge :: !Bool, + includeProjectedMedium :: !Bool, + includeProjectedLarge :: !Bool, + includeNativeLarge :: !Bool + } + +readBenchmarkSelection :: IO BenchmarkSelection +readBenchmarkSelection = + BenchmarkSelection + <$> gateEnabled SparseLargeGate + <*> gateEnabled Sparse100kGate + <*> gateEnabled BroadMediumGate + <*> gateEnabled BroadLargeGate + <*> gateEnabled ProjectedMediumGate + <*> gateEnabled ProjectedLargeGate + <*> gateEnabled NativeLargeGate + +benchmarkNotice :: BenchmarkSelection -> String +benchmarkNotice selection = + mconcat + [ "sparse-large benchmark ", + gateNotice SparseLargeGate (includeSparseLarge selection || includeSparse100k selection), + "; 100k sparse Krylov benchmark ", + gateNotice Sparse100kGate (includeSparse100k selection), + "; medium broad linalg benchmarks ", + gateNotice BroadMediumGate (includeBroadMedium selection || includeBroadLarge selection), + "; large broad linalg benchmarks ", + gateNotice BroadLargeGate (includeBroadLarge selection), + "; medium projected structured-block benchmarks ", + gateNotice ProjectedMediumGate (includeProjectedMedium selection || includeProjectedLarge selection), + "; large projected structured-block benchmarks ", + gateNotice ProjectedLargeGate (includeProjectedLarge selection), + "; large native LAPACK benchmarks ", + gateNotice NativeLargeGate (includeNativeLarge selection), + "." + ] + +data BenchmarkGate + = SparseLargeGate + | Sparse100kGate + | BroadMediumGate + | BroadLargeGate + | ProjectedMediumGate + | ProjectedLargeGate + | NativeLargeGate + +benchmarkGateEnvName :: BenchmarkGate -> String +benchmarkGateEnvName benchmarkGate = + case benchmarkGate of + SparseLargeGate -> "MOONLIGHT_LINALG_BENCH_ENABLE_SPARSE_LARGE" + Sparse100kGate -> "MOONLIGHT_LINALG_BENCH_ENABLE_100K" + BroadMediumGate -> "MOONLIGHT_LINALG_BENCH_ENABLE_BROAD_MEDIUM" + BroadLargeGate -> "MOONLIGHT_LINALG_BENCH_ENABLE_BROAD_LARGE" + ProjectedMediumGate -> "MOONLIGHT_LINALG_BENCH_ENABLE_PROJECTED_MEDIUM" + ProjectedLargeGate -> "MOONLIGHT_LINALG_BENCH_ENABLE_PROJECTED_LARGE" + NativeLargeGate -> "MOONLIGHT_LINALG_BENCH_ENABLE_NATIVE_LARGE" + +gateEnabled :: BenchmarkGate -> IO Bool +gateEnabled = + fmap parseTruthy . lookupEnv . benchmarkGateEnvName + +parseTruthy :: Maybe String -> Bool +parseTruthy maybeValue = + case maybeValue of + Just "1" -> True + Just "true" -> True + Just "TRUE" -> True + Just "yes" -> True + Just "YES" -> True + _ -> False + +gateNotice :: BenchmarkGate -> Bool -> String +gateNotice benchmarkGate enabled = + if enabled + then "enabled via " <> benchmarkGateEnvName benchmarkGate + else "skipped by default; set " <> benchmarkGateEnvName benchmarkGate <> "=1 to opt in" diff --git a/bench/support/Fixtures.hs b/bench/support/Fixtures.hs new file mode 100644 --- /dev/null +++ b/bench/support/Fixtures.hs @@ -0,0 +1,355 @@ +{-# LANGUAGE DataKinds #-} + +module Fixtures + ( benchmarkSeedBlock, + bandedDenseRows, + bandedSpdCSR, + denseBenchmarkRows, + denseBenchmarkVector, + denseOperator, + denseSpdRows, + diagonalBenchmarkValues, + genericBenchmarkTridiagonal, + gf2BenchmarkValues, + packedSparseBenchmarkOperator, + pathLaplacianTridiagonal, + projectedBenchmarkDimension, + projectedBenchmarkRows, + projectedBlockBenchmarkCases, + reducibleBenchmarkTridiagonal, + sparseKrylovBenchmarkCases, + staticsBenchmarkNetwork, + ) +where + +import Types + ( ProjectedBlockBenchmarkCase (..), + SparseKrylovBenchmarkCase (..), + SpectrumProfile (..), + ) +import Env (BenchmarkSelection (..)) +import Data.Bifunctor (first) +import Data.Foldable (traverse_) +import qualified Data.Vector as Box +import qualified Data.Vector.Unboxed as U +import Moonlight.LinAlg.Dense.GF2 (GF2 (..)) +import Moonlight.LinAlg.Operator + ( LinearOperator, + OperatorSymmetry (..), + declaredSelfAdjointVectorLinearOperator, + ) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal + ( SymmetricTridiagonal, + mkSymmetricTridiagonal, + pathLaplacianBands, + ) +import Moonlight.LinAlg.Sparse + ( PackedSparseOperator, + PackedSparseEntry, + SparseCSR, + canonicalCSRFromEntries, + mkPackedSparseOperator, + packedSparseEntry, + ) +import Moonlight.LinAlg.Statics + ( ForceNetwork, + NetworkDeclaration, + Vec3 (..), + load, + member, + network, + support, + ) +import Prelude + +sparseKrylovBenchmarkCases :: BenchmarkSelection -> [SparseKrylovBenchmarkCase] +sparseKrylovBenchmarkCases benchmarkSelection = + [SparseKrylovBenchmarkCase "path-laplacian-10k" 10000 4] + <> [SparseKrylovBenchmarkCase "path-laplacian-50k" 50000 4 | includeSparseLarge benchmarkSelection || includeSparse100k benchmarkSelection] + <> [SparseKrylovBenchmarkCase "path-laplacian-100k" 100000 4 | includeSparse100k benchmarkSelection] + +projectedBlockBenchmarkCases :: BenchmarkSelection -> [ProjectedBlockBenchmarkCase] +projectedBlockBenchmarkCases benchmarkSelection = + [ ProjectedBlockBenchmarkCase "block-clustered-24" 4 6 6 4 ClusteredSpectrum, + ProjectedBlockBenchmarkCase "block-separated-24" 4 6 6 4 SeparatedSpectrum + ] + <> if includeProjectedMedium benchmarkSelection || includeProjectedLarge benchmarkSelection + then + [ ProjectedBlockBenchmarkCase "block-clustered-144" 24 6 12 6 ClusteredSpectrum, + ProjectedBlockBenchmarkCase "block-separated-144" 24 6 12 6 SeparatedSpectrum + ] + else [] + <> if includeProjectedLarge benchmarkSelection + then + [ ProjectedBlockBenchmarkCase "block-clustered-256" 32 8 16 8 ClusteredSpectrum, + ProjectedBlockBenchmarkCase "block-separated-256" 32 8 16 8 SeparatedSpectrum + ] + else [] + +projectedBenchmarkDimension :: ProjectedBlockBenchmarkCase -> Int +projectedBenchmarkDimension benchmarkCase = + projectedBenchmarkBlockCount benchmarkCase * projectedBenchmarkBlockSize benchmarkCase + +benchmarkSeedBlock :: Int -> Int -> Box.Vector (U.Vector Double) +benchmarkSeedBlock dimension blockSize = + Box.fromList + [ U.fromList [seedEntry rowIndex columnIndex | rowIndex <- [0 .. dimension - 1]] + | columnIndex <- [0 .. blockSize - 1] + ] + +seedEntry :: Int -> Int -> Double +seedEntry rowIndex columnIndex = + let rowOffset = fromIntegral (rowIndex + 1) + columnOffset = fromIntegral (columnIndex + 1) + diagonalContribution = if rowIndex == columnIndex then 1.0 else 0.0 + smoothContribution = 1.0 / (rowOffset + columnOffset) + in diagonalContribution + smoothContribution + +denseOperator :: [[Double]] -> Either String (LinearOperator 'SelfAdjointOperator) +denseOperator rows = + let rowVectors = Box.fromList (U.fromList <$> rows) + rowCount = Box.length rowVectors + in validateDenseOperatorRows rowCount rowVectors + *> first show (declaredSelfAdjointVectorLinearOperator rowCount (denseOperatorApply rowVectors)) + +validateDenseOperatorRows :: Int -> Box.Vector (U.Vector Double) -> Either String () +validateDenseOperatorRows dimension rowVectors = + traverse_ validateRow rowVectors + where + validateRow rowVector + | U.length rowVector /= dimension = + Left "benchmark dense operator rows must form a square matrix" + | U.any (not . isFiniteDouble) rowVector = + Left "benchmark dense operator rows must contain finite entries" + | otherwise = Right () + +denseOperatorApply :: Box.Vector (U.Vector Double) -> U.Vector Double -> Either errorValue (U.Vector Double) +denseOperatorApply rowVectors inputVector = + Right + ( U.generate + (Box.length rowVectors) + (\rowIndex -> dotDenseRow inputVector (rowVectors `Box.unsafeIndex` rowIndex)) + ) + +dotDenseRow :: U.Vector Double -> U.Vector Double -> Double +dotDenseRow inputVector rowVector = + U.ifoldl' + (\accumulator columnIndex rowEntry -> accumulator + rowEntry * (inputVector `U.unsafeIndex` columnIndex)) + 0.0 + rowVector + +isFiniteDouble :: Double -> Bool +isFiniteDouble value = + not (isNaN value || isInfinite value) + +projectedBenchmarkRows :: ProjectedBlockBenchmarkCase -> [[Double]] +projectedBenchmarkRows benchmarkCase = + let dimension = projectedBenchmarkDimension benchmarkCase + blockSize = projectedBenchmarkBlockSize benchmarkCase + spectrumProfile = projectedBenchmarkSpectrumProfile benchmarkCase + in [ [ projectedBenchmarkEntry spectrumProfile blockSize rowIndex columnIndex + | columnIndex <- [0 .. dimension - 1] + ] + | rowIndex <- [0 .. dimension - 1] + ] + +projectedBenchmarkEntry :: SpectrumProfile -> Int -> Int -> Int -> Double +projectedBenchmarkEntry spectrumProfile blockSize rowIndex columnIndex = + let (rowBlock, rowWithinBlock) = rowIndex `divMod` blockSize + (columnBlock, columnWithinBlock) = columnIndex `divMod` blockSize + in case compare rowBlock columnBlock of + EQ -> diagonalBlockEntry spectrumProfile rowBlock rowWithinBlock columnWithinBlock + LT -> + if rowBlock + 1 == columnBlock + then offDiagonalBlockEntry spectrumProfile rowBlock rowWithinBlock columnWithinBlock + else 0.0 + GT -> + if columnBlock + 1 == rowBlock + then offDiagonalBlockEntry spectrumProfile columnBlock columnWithinBlock rowWithinBlock + else 0.0 + +diagonalBlockEntry :: SpectrumProfile -> Int -> Int -> Int -> Double +diagonalBlockEntry spectrumProfile blockIndex rowWithinBlock columnWithinBlock = + let separationBase = + case spectrumProfile of + ClusteredSpectrum -> 12.0 + 0.005 * fromIntegral blockIndex + SeparatedSpectrum -> 2.0 + 1.5 * fromIntegral blockIndex + localOffset = 0.02 * fromIntegral (rowWithinBlock + columnWithinBlock) + entryWeight = + if rowWithinBlock == columnWithinBlock + then separationBase + 1.0 + localOffset + else 0.04 / fromIntegral (1 + abs (rowWithinBlock - columnWithinBlock)) + in entryWeight + +offDiagonalBlockEntry :: SpectrumProfile -> Int -> Int -> Int -> Double +offDiagonalBlockEntry spectrumProfile blockIndex rowWithinBlock columnWithinBlock = + let couplingBase = + case spectrumProfile of + ClusteredSpectrum -> 0.06 + 0.002 * fromIntegral (blockIndex `mod` 3) + SeparatedSpectrum -> 0.03 + 0.001 * fromIntegral (blockIndex `mod` 3) + in couplingBase / fromIntegral (1 + abs (rowWithinBlock - columnWithinBlock)) + +pathLaplacianTridiagonal :: Int -> Either String SymmetricTridiagonal +pathLaplacianTridiagonal dimension = + first show (pathLaplacianBands dimension >>= uncurry mkSymmetricTridiagonal) + +genericBenchmarkTridiagonal :: Int -> Either String SymmetricTridiagonal +genericBenchmarkTridiagonal dimension = + first show + ( mkSymmetricTridiagonal + (genericTridiagonalDiagonalEntry <$> [0 .. dimension - 1]) + (genericTridiagonalOffDiagonalEntry <$> [0 .. dimension - 2]) + ) + +genericTridiagonalDiagonalEntry :: Int -> Double +genericTridiagonalDiagonalEntry indexValue = + 2.0 + fromIntegral (indexValue `mod` 17) / 17.0 + +genericTridiagonalOffDiagonalEntry :: Int -> Double +genericTridiagonalOffDiagonalEntry indexValue = + -0.35 - 0.01 * fromIntegral (indexValue `mod` 5) + +reducibleBenchmarkTridiagonal :: Int -> Either String SymmetricTridiagonal +reducibleBenchmarkTridiagonal dimension = + first show + ( mkSymmetricTridiagonal + (genericTridiagonalDiagonalEntry <$> [0 .. dimension - 1]) + (reducibleTridiagonalOffDiagonalEntry <$> [0 .. dimension - 2]) + ) + +reducibleTridiagonalOffDiagonalEntry :: Int -> Double +reducibleTridiagonalOffDiagonalEntry indexValue = + if (indexValue + 1) `mod` 32 == 0 + then 0.0 + else genericTridiagonalOffDiagonalEntry indexValue + +denseBenchmarkRows :: Int -> [[Double]] +denseBenchmarkRows dimension = + [ [denseBenchmarkEntry rowIndex columnIndex | columnIndex <- [0 .. dimension - 1]] + | rowIndex <- [0 .. dimension - 1] + ] + +denseBenchmarkEntry :: Int -> Int -> Double +denseBenchmarkEntry rowIndex columnIndex = + let rowWeight = fromIntegral (rowIndex + 1) + columnWeight = fromIntegral (columnIndex + 1) + diagonalContribution = + if rowIndex == columnIndex + then 2.0 + 0.01 * rowWeight + else 0.0 + smoothContribution = 1.0 / (rowWeight + 2.0 * columnWeight + 3.0) + in diagonalContribution + smoothContribution + +denseSpdRows :: Int -> [[Double]] +denseSpdRows dimension = + [ [denseSpdEntry dimension rowIndex columnIndex | columnIndex <- [0 .. dimension - 1]] + | rowIndex <- [0 .. dimension - 1] + ] + +denseSpdEntry :: Int -> Int -> Int -> Double +denseSpdEntry dimension rowIndex columnIndex = + if rowIndex == columnIndex + then fromIntegral dimension + 2.0 + 0.05 * fromIntegral rowIndex + else 1.0 / fromIntegral (2 + abs (rowIndex - columnIndex)) + +denseBenchmarkVector :: Int -> [Double] +denseBenchmarkVector dimension = + fmap (\indexValue -> 1.0 + fromIntegral (indexValue `mod` 7) / 7.0) [0 .. dimension - 1] + +bandedDenseRows :: Int -> [[Double]] +bandedDenseRows dimension = + [ [bandedDenseEntry dimension rowIndex columnIndex | columnIndex <- [0 .. dimension - 1]] + | rowIndex <- [0 .. dimension - 1] + ] + +bandedDenseEntry :: Int -> Int -> Int -> Double +bandedDenseEntry dimension rowIndex columnIndex + | rowIndex == columnIndex = 4.0 + 0.001 * fromIntegral dimension + | abs (rowIndex - columnIndex) == 1 = -1.0 + | abs (rowIndex - columnIndex) == 2 = 0.25 + | otherwise = 0.0 + +bandedSpdCSR :: Int -> Either String (SparseCSR Double) +bandedSpdCSR dimension = + case + canonicalCSRFromEntries + dimension + dimension + (bandedSpdEntries dimension) + of + Left err -> Left (show err) + Right csrValue -> Right csrValue + +bandedSpdEntries :: Int -> [(Int, Int, Double)] +bandedSpdEntries dimension = + concatMap + ( \rowIndex -> + (\(columnIndex, value) -> (rowIndex, columnIndex, value)) + <$> bandedSpdRowEntries dimension rowIndex + ) + [0 .. dimension - 1] + +bandedSpdRowEntries :: Int -> Int -> [(Int, Double)] +bandedSpdRowEntries dimension rowIndex = + filter + (\(columnIndex, _) -> columnIndex >= 0 && columnIndex < dimension) + [ (rowIndex - 2, 0.25), + (rowIndex - 1, -1.0), + (rowIndex, 4.0), + (rowIndex + 1, -1.0), + (rowIndex + 2, 0.25) + ] + +diagonalBenchmarkValues :: Int -> [Double] +diagonalBenchmarkValues dimension = + let positiveDimension = max 1 dimension + in fmap (\indexValue -> 2.0 + fromIntegral (indexValue `mod` positiveDimension) / fromIntegral positiveDimension) [0 .. dimension - 1] + +packedSparseBenchmarkOperator :: Int -> Either String (PackedSparseOperator Double) +packedSparseBenchmarkOperator dimension = + case mkPackedSparseOperator (fromIntegral dimension) (fromIntegral dimension) (packedSparseEntries dimension) of + Left err -> Left (show err) + Right operatorValue -> Right operatorValue + +packedSparseEntries :: Int -> [PackedSparseEntry Double] +packedSparseEntries dimension = + [ packedSparseEntry sourceOffset targetOffset (packedSparseCoefficient sourceOffset targetOffset) + | targetOffset <- [0 .. dimension - 1], + sourceOffset <- [targetOffset - 1, targetOffset, targetOffset + 1], + sourceOffset >= 0, + sourceOffset < dimension + ] + +packedSparseCoefficient :: Int -> Int -> Double +packedSparseCoefficient sourceOffset targetOffset = + if sourceOffset == targetOffset + then 2.0 + else (-0.5) + +gf2BenchmarkValues :: Int -> Int -> [GF2] +gf2BenchmarkValues rowCount columnCount = + [ if gf2BenchmarkBit rowIndex columnIndex then GF2One else GF2Zero + | rowIndex <- [0 .. rowCount - 1], + columnIndex <- [0 .. columnCount - 1] + ] + +gf2BenchmarkBit :: Int -> Int -> Bool +gf2BenchmarkBit rowIndex columnIndex = + rowIndex == columnIndex + || ((rowIndex * 17 + columnIndex * 31 + rowIndex * columnIndex) `mod` 23 == 0) + +staticsBenchmarkNetwork :: Int -> Either String ForceNetwork +staticsBenchmarkNetwork spanCount = + case network (concatMap spanDeclarations [0 .. spanCount - 1]) of + Left err -> Left (show err) + Right networkValue -> Right networkValue + +spanDeclarations :: Int -> [NetworkDeclaration] +spanDeclarations spanIndex = + let supportLabel = "support-" <> show spanIndex + loadLabel = "load-" <> show spanIndex + coordinate = fromIntegral spanIndex + in [ support supportLabel (Vec3 coordinate 0.0 0.0), + load loadLabel (Vec3 coordinate 1.0 0.0) (Vec3 0.0 (-10.0 - coordinate) 0.0), + member supportLabel loadLabel + ] diff --git a/bench/support/Types.hs b/bench/support/Types.hs new file mode 100644 --- /dev/null +++ b/bench/support/Types.hs @@ -0,0 +1,289 @@ +{-# LANGUAGE DataKinds #-} + +module Types + ( BenchmarkSetup (..), + PreparedBenchmarkRow (..), + OnceBenchmark (..), + OnceBenchmarkResult (..), + OnceBenchmarkStats (..), + BenchmarkWeight (..), + ProjectedBlockBenchmarkCase (..), + ProjectedBlockPreparedCase (..), + SparseKrylovBenchmarkCase (..), + SparseKrylovPreparedCase (..), + SpectrumProfile (..), + eigenpairsChecksum, + eigenpairsResidualValidationChecksum, + benchmarkWeightEither, + eitherBenchmarkWeight, + prepareBenchmarkSetup, + renderPreparedBenchmark, + renderPreparedOnceBenchmark, + runOnceBenchmark, + ) +where + +import Control.DeepSeq (NFData (..), force) +import Control.Exception (evaluate) +import Control.Monad (foldM) +import Data.Bifunctor (first) +import qualified Data.Vector.Unboxed as U +import GHC.Stats + ( RTSStats, + allocated_bytes, + gc, + gcdetails_live_bytes, + getRTSStats, + getRTSStatsEnabled, + max_live_bytes, + ) +import Moonlight.LinAlg.Pure.Krylov.Projected + ( ProjectedSubspace, + ) +import Moonlight.LinAlg.Operator + ( LinearOperator, + OperatorSymmetry (..), + ) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal (SymmetricTridiagonal) +import Moonlight.LinAlg.Spectral + ( Eigenpairs, + eigenpairCount, + eigenpairDimension, + eigenpairResidualNorms, + eigenpairVectorAt, + eigenpairValues, + eigenpairVectorsColumnMajor, + ) +import System.CPUTime (getCPUTime) +import System.Mem (performMajorGC) +import Test.Tasty.Bench (Benchmark, bench, env, nf, nfIO) +import Prelude + +newtype BenchmarkSetup value = BenchmarkSetup + { runBenchmarkSetup :: Either String value + } + +data OnceBenchmark = OnceBenchmark + { onceBenchmarkLabel :: !String, + onceBenchmarkAction :: IO (Either String Double) + } + +data OnceBenchmarkResult = OnceBenchmarkResult + { onceResultLabel :: !String, + onceResultElapsedSeconds :: !Double, + onceResultChecksum :: !Double, + onceResultStats :: !(Maybe OnceBenchmarkStats) + } + +data OnceBenchmarkStats = OnceBenchmarkStats + { onceAllocatedBytes :: !Integer, + onceLiveBytesAfterMajorGC :: !Integer, + onceProcessMaximumLiveBytes :: !Integer + } + +data BenchmarkWeight + = BenchmarkWeight !Double + | BenchmarkMeasurementFailure !String + deriving stock (Show) + +data PreparedBenchmarkRow preparedCase + = PurePreparedBenchmarkRow + !String + (preparedCase -> BenchmarkWeight) + | EffectfulPreparedBenchmarkRow + !String + (preparedCase -> IO BenchmarkWeight) + +instance NFData BenchmarkWeight where + rnf weightValue = + case weightValue of + BenchmarkWeight checksumValue -> rnf checksumValue + BenchmarkMeasurementFailure failureText -> failBenchmarkMeasurement failureText + +data SparseKrylovBenchmarkCase = SparseKrylovBenchmarkCase + { sparseBenchmarkLabel :: !String, + sparseBenchmarkDimension :: !Int, + sparseBenchmarkRequestedModes :: !Int + } + +data SparseKrylovPreparedCase = SparseKrylovPreparedCase + { sparsePreparedLabel :: !String, + sparsePreparedRequestedModes :: !Int, + sparsePreparedTridiagonal :: !SymmetricTridiagonal + } + +instance NFData SparseKrylovPreparedCase where + rnf preparedCase = + sparsePreparedLabel preparedCase + `seq` sparsePreparedRequestedModes preparedCase + `seq` sparsePreparedTridiagonal preparedCase + `seq` () + +data ProjectedBlockBenchmarkCase = ProjectedBlockBenchmarkCase + { projectedBenchmarkLabel :: !String, + projectedBenchmarkBlockCount :: !Int, + projectedBenchmarkBlockSize :: !Int, + projectedBenchmarkIterations :: !Int, + projectedBenchmarkRequestedModes :: !Int, + projectedBenchmarkSpectrumProfile :: !SpectrumProfile + } + +data SpectrumProfile = ClusteredSpectrum | SeparatedSpectrum + deriving stock (Eq, Show) + +data ProjectedBlockPreparedCase = ProjectedBlockPreparedCase + { projectedPreparedCase :: !ProjectedBlockBenchmarkCase, + projectedPreparedOperator :: !(LinearOperator 'SelfAdjointOperator), + projectedPreparedSubspace :: !ProjectedSubspace, + projectedPreparedDimension :: !Int + } + +instance NFData ProjectedBlockPreparedCase where + rnf preparedCase = + projectedPreparedCase preparedCase + `seq` projectedPreparedOperator preparedCase + `seq` projectedPreparedSubspace preparedCase + `seq` projectedPreparedDimension preparedCase + `seq` () + +prepareBenchmarkSetup :: BenchmarkSetup value -> IO value +prepareBenchmarkSetup = + either failBenchmarkSetup pure . runBenchmarkSetup + +renderPreparedBenchmark :: + NFData preparedCase => + BenchmarkSetup preparedCase -> + PreparedBenchmarkRow preparedCase -> + Benchmark +renderPreparedBenchmark benchmarkSetup benchmarkRow = + case benchmarkRow of + PurePreparedBenchmarkRow benchmarkLabel benchmarkWeight -> + env (prepareBenchmarkSetup benchmarkSetup) $ \preparedCase -> + bench benchmarkLabel (nf benchmarkWeight preparedCase) + EffectfulPreparedBenchmarkRow benchmarkLabel benchmarkWeight -> + env (prepareBenchmarkSetup benchmarkSetup) $ \preparedCase -> + bench benchmarkLabel (nfIO (benchmarkWeight preparedCase)) + +renderPreparedOnceBenchmark :: + String -> + BenchmarkSetup preparedCase -> + PreparedBenchmarkRow preparedCase -> + OnceBenchmark +renderPreparedOnceBenchmark benchmarkGroup benchmarkSetup benchmarkRow = + case benchmarkRow of + PurePreparedBenchmarkRow benchmarkLabel benchmarkWeight -> + OnceBenchmark + { onceBenchmarkLabel = benchmarkGroup <> benchmarkLabel, + onceBenchmarkAction = + pure + (runBenchmarkSetup benchmarkSetup >>= benchmarkWeightEither . benchmarkWeight) + } + EffectfulPreparedBenchmarkRow benchmarkLabel benchmarkWeight -> + OnceBenchmark + { onceBenchmarkLabel = benchmarkGroup <> benchmarkLabel, + onceBenchmarkAction = + case runBenchmarkSetup benchmarkSetup of + Left err -> pure (Left err) + Right preparedCase -> benchmarkWeightEither <$> benchmarkWeight preparedCase + } + +runOnceBenchmark :: OnceBenchmark -> IO (Either String OnceBenchmarkResult) +runOnceBenchmark benchmarkValue = do + statsEnabled <- getRTSStatsEnabled + beforeStats <- beforeOnceStats statsEnabled + startTime <- getCPUTime + resultValue <- onceBenchmarkAction benchmarkValue >>= evaluate . force + endTime <- getCPUTime + afterStats <- afterOnceStats statsEnabled + pure + ( fmap + ( \checksumValue -> + OnceBenchmarkResult + { onceResultLabel = onceBenchmarkLabel benchmarkValue, + onceResultElapsedSeconds = fromIntegral (endTime - startTime) / 1.0e12, + onceResultChecksum = checksumValue, + onceResultStats = onceStatsDelta <$> beforeStats <*> afterStats + } + ) + resultValue + ) + +beforeOnceStats :: Bool -> IO (Maybe RTSStats) +beforeOnceStats statsEnabled = + if statsEnabled + then performMajorGC *> (Just <$> getRTSStats) + else pure Nothing + +afterOnceStats :: Bool -> IO (Maybe RTSStats) +afterOnceStats statsEnabled = + if statsEnabled + then performMajorGC *> (Just <$> getRTSStats) + else pure Nothing + +onceStatsDelta :: RTSStats -> RTSStats -> OnceBenchmarkStats +onceStatsDelta beforeStats afterStats = + OnceBenchmarkStats + { onceAllocatedBytes = + toInteger (allocated_bytes afterStats - allocated_bytes beforeStats), + onceLiveBytesAfterMajorGC = + toInteger (gcdetails_live_bytes (gc afterStats)), + onceProcessMaximumLiveBytes = + toInteger (max_live_bytes afterStats) + } + +failBenchmarkSetup :: String -> IO value +failBenchmarkSetup failureText = do + putStrLn ("moonlight-linalg benchmark setup failed: " <> failureText) + ioError (userError failureText) + +failBenchmarkMeasurement :: String -> value +failBenchmarkMeasurement failureText = + error ("moonlight-linalg benchmark measurement failed: " <> failureText) + +eigenpairsChecksum :: Eigenpairs -> Double +eigenpairsChecksum pairs = + U.sum (eigenpairValues pairs) + + U.sum (eigenpairResidualNorms pairs) + + U.sum (U.map abs (eigenpairVectorsColumnMajor pairs)) + +eigenpairsResidualValidationChecksum :: Show err => (U.Vector Double -> Either err (U.Vector Double)) -> Eigenpairs -> Either String Double +eigenpairsResidualValidationChecksum applyVector pairs = + foldM accumulateResidualChecksum 0.0 [0 .. eigenpairCount pairs - 1] + where + accumulateResidualChecksum accumulatedChecksum columnIndex = do + eigenvalue <- eigenvalueAt columnIndex + eigenvector <- first show (eigenpairVectorAt columnIndex pairs) + imageVector <- first show (applyVector eigenvector) + if U.length imageVector == eigenpairDimension pairs + then + pure + ( accumulatedChecksum + + normU + ( U.zipWith + (-) + imageVector + (U.map (* eigenvalue) eigenvector) + ) + ) + else Left "eigenpair residual validation apply dimension mismatch" + + eigenvalueAt columnIndex = + case eigenpairValues pairs U.!? columnIndex of + Just eigenvalue -> Right eigenvalue + Nothing -> Left "eigenpair residual validation value index out of bounds" + +normU :: U.Vector Double -> Double +normU vectorValue = + sqrt (U.sum (U.map (\entryValue -> entryValue * entryValue) vectorValue)) + +eitherBenchmarkWeight :: Show err => String -> (value -> Double) -> Either err value -> BenchmarkWeight +eitherBenchmarkWeight label checksum = + either + (\err -> BenchmarkMeasurementFailure (label <> ": " <> show err)) + (BenchmarkWeight . checksum) + +benchmarkWeightEither :: BenchmarkWeight -> Either String Double +benchmarkWeightEither weightValue = + case weightValue of + BenchmarkWeight checksumValue -> Right checksumValue + BenchmarkMeasurementFailure failureText -> Left failureText diff --git a/cbits/moonlight_linalg_native.c b/cbits/moonlight_linalg_native.c new file mode 100644 --- /dev/null +++ b/cbits/moonlight_linalg_native.c @@ -0,0 +1,37 @@ +#include + +extern void dgemm_(char *transa, char *transb, int *m, int *n, int *k, + double *alpha, const double *a, int *lda, + const double *b, int *ldb, double *beta, + double *c, int *ldc); + +void moonlight_dgemm_row_major(int left_rows, + int right_columns, + int inner_dimension, + const double *left_row_major, + const double *right_row_major, + double *output_row_major) { + char no_transpose = 'N'; + int m = right_columns; + int n = left_rows; + int k = inner_dimension; + int left_as_column_major_leading_dimension = right_columns; + int right_as_column_major_leading_dimension = inner_dimension; + int output_leading_dimension = right_columns; + double alpha = 1.0; + double beta = 0.0; + + dgemm_(&no_transpose, + &no_transpose, + &m, + &n, + &k, + &alpha, + right_row_major, + &left_as_column_major_leading_dimension, + left_row_major, + &right_as_column_major_leading_dimension, + &beta, + output_row_major, + &output_leading_dimension); +} diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md new file mode 100644 --- /dev/null +++ b/docs/ARCHITECTURE.md @@ -0,0 +1,234 @@ +# moonlight-linalg architecture + +`moonlight-linalg` supplies Moonlight's numerical linear-algebra tier: typed dense, +sparse, finite-field, geometric, and Krylov/spectral machinery beneath homology, +analysis, sheaf, and solver packages. This document records the sublibrary slice +map, the boundary verdicts behind it, and the numeric-floor rulings the campaign +established; `README.md` owns usage. + +## Sublibrary slices + +The implementation is split into graded public sublibraries, one source directory +each, with the dependency DAG enforced by cabal and guarded by +`ArchitectureSpec` slice-discipline source checks: + +| slice | contents | in-package deps | +|---|---|---| +| `carrier` | `Internal.{Primitives, DenseList, Storage, VectorOps, Discrete, GF2.*}`, `Pure.Dense.{Types, Rows, Flat}` | — | +| `structured` | `Pure.Structured.{Tridiagonal, BlockTridiagonal}` | — | +| `eigen` | `Internal.Eigen.*` | carrier | +| `geometry` | `Pure.Geometry.*` | carrier | +| `dense` | `Pure.Dense.*` (rest), `Internal.Backend.*`, `Internal.Dense.*` | carrier, eigen | +| `domain` | `Pure.Domain.{Bareiss, Smith, Smith.Multimodular, Smith.Witnessed}` | carrier, dense | +| `sparse` | `Pure.Sparse.*` | carrier, structured | +| `statics` | `Pure.Statics.*` | carrier, dense, geometry | +| `spectral` | `Pure.Operator(+.Internal)`, `Pure.Krylov.*`, `Pure.Spectral.*` | carrier, eigen, sparse, structured | +| `native` | `Effect.Native.{Dispatch, LAPACK}` + cbits | carrier, dense, eigen, spectral, structured | +| `laws` | `src-laws/**` harnesses and registry | public library | +| public | `src-public/*` facades, unchanged surface | all slices | + +`moonlight-linalg-native` is the sole owner of the C bits and the +Accelerate (Darwin) / LAPACK+BLAS (elsewhere) linkage; the quarantine is +cabal-enforced, not conventional. + +Two edges the campaign plan predicted turned out dead in live imports and were +omitted rather than transcribed: `sparse` does not depend on `dense` (its solvers +speak carrier and structured vocabulary only), and `spectral` does not depend on +`dense` (its dense fallback speaks `Internal.Eigen.Symmetric` and the carrier +`Pure.Dense.Flat` directly). `native` depends on `eigen`, which the plan missed. + +## Facades + +The public library re-exports through `Moonlight.LinAlg` and the focused facades +(`.Dense`, `.Sparse`, `.Operator`, `.Spectral`, `.Krylov`, `.Native`, `.Domain`, +`.Geometry`, `.Statics`). `Pure.*`, `Internal.*`, and `Effect.*` module names are +slice ownership, not a second public vocabulary; downstream packages import +facades only, and the boundary is verified clean (zero `Internal.*`/`Pure.*` +imports outside the package). + +`Moonlight.LinAlg.Native` is a true re-export facade: its implementation body +lives in `Effect.Native.Dispatch` inside the native slice, so the public module +carries no logic above the quarantine. + +## Boundary verdicts + +- **Carrier redraw (V1).** Three alleged DAG violations dissolved by reclassifying + `Pure.Dense.{Types, Rows, Flat}` and `Internal.Storage` as carriers: they import + only carrier siblings and core, so the bottom slice owns them and the + `Internal.Storage → Rows` inversion never existed. `Pure.Structured.*` imports + nothing in-package and owns its own low slice, dissolving the alleged + spectral↔sparse cycle. `Dense.Decomposition → Internal.Eigen.Symmetric` is + lawful grading (dense sits above eigen), not a wound. +- **Geometry gate (V2) — PASSED.** 2×2/3×3 symmetric eigendecomposition is + closed-form (analytic 2×2; stable hybrid trigonometric 3×3) behind the + preserved `eigendecomposeSymmetric{2,3}With` injection seams. The + `GeometrySymmetricEigen{Reconstructs,Orthonormal}` laws, including generated + near-degenerate spectra, adjudicated the gate; geometry consequently depends + only on carrier. +- **Gram-SVD deleted (V4).** `thinSvdFullColumnRank` is one-sided Jacobi with a + typed failure ADT; the condition-squaring Gram path is gone with no fallback. + The reconstruction law is the reference. +- **Certification is a result distinction, not a parallel API (V5).** + `symmetricEigenPairs` returns fast unchecked results; + `certifySymmetricEigenResult` is the explicit certification morphism. + `SymmetricEigenUncheckedPassesCertification` prevents drift. +- **Orphan ruling (V7).** The `-Wno-orphans` pragma was vestigial — all + `Internal.Discrete` instances are for the locally-defined `GF2`. Deleted; the + rebuild surfaced no genuine orphan. + +## Numeric-floor rulings + +- **Extreme-exponent literals are not constant-folded.** GHC 9.14 does not fold + `fromRational` for decimal scientific literals near the `Double` range limits; + their exact rationals survive to runtime as bignum arithmetic, and an `INLINE` + pragma on such a constant pastes GMP division into every consumer's hot loop. + `epsDouble`, `safeMinimumDouble`, and `maxFiniteDouble` are therefore + `encodeFloat` forms under `NOINLINE` (`Internal.Eigen.Kernels`), and new + numeric constants at extreme exponents must follow that shape. +- **Libraries build at `-O2`.** cabal defaults libraries to `-O1`; SpecConstr and + LiberateCase are worth ×3–4 on the Krylov and tridiagonal inner loops + (verified 89.3ms → 23.6ms on identical source). `-O2` lives in the cabal + `shared-properties` block. +- **Delivery floors are absorbed, never re-tightened.** A gate above a + certified-approximate engine must use the engine's certified tolerance tier: + Ritz lock gates lift candidates and re-test them against + `ritzLockToleranceBound` with demotion back to the unlocked pool, and + agreement laws between the selected (inverse-iteration, floor ≈ 1e7·ε·n·scale) + and dense routes assert at the residual tier, not the approx tier. A tighter + tolerance falsifies honestly-agreeing routes. +- **Demand-aware dense fallback dispatch.** The generic CSR spectral fallback + densifies and dense-solves at or below + `denseSpectralFallbackDimensionThreshold = 512`. Through dimension 1,024 it + also keeps requests for at least one quarter of the spectrum on the bounded + dense route; smaller requests descend to thick-restart Lanczos. Structure is + still authoritative, so diagonal, path-Laplacian, and tridiagonal sources + bypass this generic policy. Restart is convergence-bound on clustered + spectra; the crossover and demand provenance rows live in + `BENCHMARKS-m4-pro.md`. +- **Certified graph-Laplacian descent.** `graphLaplacianLinearOperator` is the + sole constructor that can retain `GraphLaplacianCSRSource` provenance after + checked sparse assembly; an arbitrary symmetric CSR is never promoted by + inspection. Smallest-mode requests of at most eight columns at dimension + 4,096 and above use private heavy-edge cascadic descent: aggregate masses and + edge weights form the coarse overlap data, mass-normalized prolongation glues + coarse sections back to the fine cover, and block Rayleigh--Ritz/Jacobi + refinement returns only columns whose residuals were recomputed against the + authoritative fine CSR. Stalled coarsening is the one typed inapplicability + obstruction that descends to generic Lanczos; rank loss, invalid requests, + and exhausted refinement remain failures. Values are derived from the pair + owner rather than solved through a second route. + +## Law surface + +`moonlight-linalg-laws` exposes the effectful harnesses and a closed `LawName` +ADT of 67 laws across dense algebra, decompositions, field/GF2, domain +(Smith/Bareiss), sparse, operator, preconditioner, Krylov/spectral, geometry, +and statics owners, with a registry-totality manifest test. Tolerance tiers are +stated once in `Effect.Harness.Core` (`approxTolerance` 1e-8, +`residualTolerance` 1e-5, `orthonormalTolerance` 1e-6) — never per-law magic +epsilons. Two agreement obligations live as deterministic anchors in the test +suite rather than the registry: GF2 sparse-column rank versus packed rank +(`GF2Spec`) and native-versus-pure selected eigensolves (`KrylovSpec`), both +fixture-carrier comparisons rather than generated laws. + +## Construction and validation + +Boundary constructors validate and return `Either MoonlightError`: +`fromListMatrix` and friends for typed dense shapes, `mkSparseCOO`/`cooToCSR` +for sparse carriers, `mkDenseDoubleMatrixRowMajor` for flat dense storage, +`mkSymmetricTridiagonal` for structured operators, and explicit self-adjoint +operator construction in `Pure.Operator`. Failure-prone kernels carry typed +failure ADTs (SVD non-finite/rank/sweep-budget, IC(0) pivot/nullspace +breakdowns, Lanczos state sections) instead of partial numerics. + +The exact `[[a]]` elimination tower survives deliberately for `Field` +polymorphism; flat unboxed `Double` kernels own the hot paths. Dense rows are +validated authoring, not hot storage. + +## Multimodular Smith engine + +The modular/CRT Smith deferral was adjudicated by measured demand: the external +referent harness (`potentialimprovements/linalg-external-referents/`) showed +FLINT `fmpz_mat_snf` beating the classical diagonal route ×11–35 on dense +random n=8–32, a coefficient-explosion gap that widens with n. +`Pure.Domain.Smith.Multimodular` now owns the `Integer` diagonal-only engine: +an unboxed `Word64` word-prime sweep (fixed deterministic 31-bit ladder) gives +per-prime determinant and rank; CRT with symmetric lift recovers the exact +determinant, with the Hadamard minor bound serving only as the +prime-consumption certificate; square full-rank inputs then run Iliopoulos +Smith elimination with all arithmetic mod `R = 2·|det|`, and rectangular or +rank-deficient inputs reach that core through unimodular Hermite-style +compression. The mod-`R` phase runs on flat mutable carriers tiered by modulus +size (`Word64` below 2^32, `Word64` with 128-bit primop intermediates below +2^62, flat boxed `Integer` above). Dispatch is a `NOINLINE`-pinned rewrite rule +on `smithDiagonalForm` at `Integer` — a perf-only substitution, semantically +identical; unoptimized builds fall back to the classical route. Referees: four +fixture agreements with +by-name assertions in `DomainSpec`, the generated +`SmithDiagonalOnlyAgreesWithFull` law, and the external FLINT agreement gate +(exact, including signs). Post-surgery the engine beats the prior classical +diagonal route ×1.8–3.6 across n=8–32 and sits ×2.8–12.3 from FLINT (was +×11–35), with the asymptotic wall removed. + +Ruling absorbed on the way: classical Smith diagonals carried algorithm-artifact +unit signs (e.g. `-1, 1, -493`). Both classical routes now canonicalize +invariant factors to nonnegative via `gcdDomain _ zero`, folding the unit flips +into the left witness and its inverse (unimodularity and reconstruction laws +preserved). Smith diagonals are canonical associates everywhere; the FLINT +agreement is exact rather than up-to-units. + +## Witnessed Smith engine + +Witness-carrying `smithNormalForm` at `Integer` suffered the same +coefficient-explosion wall as the classical diagonal route, but worse: the +alternating Hermite arena reached 111,000-bit work entries at n=32 (69 ms) +where the final diagonal needs 100 bits. Measured mechanism: integer echelon +stays Hadamard-bounded exactly while pivots are ±1, and the first gcd +pair-transform at a non-unit pivot destroys the Schur minor structure — no +pivot policy avoids it. `Pure.Domain.Smith.Witnessed` now dispatches square +inputs of dimension ≥ 25 with certified full rank (word-prime sweep) to a +mod-determinant engine; everything else keeps the alternating arena, which is +faster below the floor and produces pristine witnesses. + +The fast path alternates row/column Hermite reduction with all work-matrix +arithmetic bounded mod `R = 2·|det|`, on the Domich–Kannan–Trotter stack +`[A; R·I]`: virtual rows `R·e_j` are carried exactly in the pool, preserving +the invariant `L(rows ∪ R·Zⁿ) = L(A)`. Centering entries mod `R` without the +stack silently shrinks the lattice — the forward transform still verifies while +the inverse fails; the stack is the soundness boundary, not an optimization. +No transform is tracked inside the HNF. Each phase transform is recovered +afterward by per-prime linear solves (upper-triangular back-substitution when +the denominator matrix is triangular — always true after round one — augmented +word Gauss–Jordan otherwise) CRT-combined with symmetric lift, terminated +early when the lift stabilizes across consecutive primes, and then verified +deterministically: the candidate satisfies its congruence mod the accumulated +CRT modulus by construction, so verification consumes fresh primes only until +the combined modulus exceeds the product entry bound, at which point equality +is exact, not probabilistic. The Cramer/Hadamard cap (power-of-two square-root +bound; an upper bound is all a cap needs) is evaluated only if stabilization +never fires, and prime-ladder exhaustion is a typed failure. Phase outputs +equal to their inputs short-circuit to identity transforms. The finale takes +`L·A` from state (maintained compositionally — round one's `L·A` is the row +HNF itself), computes both inverse witnesses by exact diagonal division +(`L⁻¹ = A·R·D⁻¹`, `R⁻¹ = D⁻¹·L·A`) with typed inexact-quotient failures, and +seeds the existing arena for divisibility-chain repair and unit normalization. +`L·A·R = D` needs no final re-verification: it follows by associativity from +the per-phase verified equations. + +At n=32 dense random with torsion the engine runs 11.0 ms (was 69 ms), n=24 +stays on the alternating arena at 2.2 ms, and witness entries stay ≤ 189 bits. +Referees: the ≥ 25 nonsingular fast-path fixture in `DomainSpec` (strict +diagonal dominance guarantees the dispatch engages), the witnessed +reconstruction and unimodularity assertions, the generated Smith laws, and the +external referent agreement gate. + +Sealed 2026-07-06 by Fable. + +## Deferred with cause + +MRRR (complexity out of proportion to the selected-pairs need), general-purpose +AMG preconditioning beyond the certified graph-spectral route, +support-graph/KMP preconditioners +(paper-only), M4RI (until measured demand), SIMD kernels (GHC native codegen +rejects `DoubleX2#` on aarch64; LLVM-only), fast approximate normalize +(correctness policy). diff --git a/docs/BENCHMARKS-m4-pro.md b/docs/BENCHMARKS-m4-pro.md new file mode 100644 --- /dev/null +++ b/docs/BENCHMARKS-m4-pro.md @@ -0,0 +1,261 @@ +# moonlight-linalg Benchmarks — Apple M4 Pro + +Measured on Apple M4 Pro via `tasty-bench` on macOS. + +## Environment + +| Field | Value | +|---|---| +| Machine | MacBook Pro | +| Model identifier | Mac16,7 | +| Chip | Apple M4 Pro | +| CPU cores | 14 total: 10 performance, 4 efficiency | +| Memory | 48 GB unified memory | +| Architecture | aarch64 / arm64 macOS | +| macOS | 26.5.1 (25F80) | +| GHC | 9.14.1 | +| Cabal | 3.16.1.0 | +| Benchmark runner | `tasty-bench` | +| Cabal target | `moonlight-linalg:moonlight-linalg-bench` | +| Cabal parallelism | `-j1` | +| Benchmark timeout | `--timeout=1s` | +| Pre-sat2 CSV | `/tmp/moonlight-linalg-pre-sat2-closure.csv` | +| Baseline `HEAD` CSV | `/tmp/moonlight-linalg-baseline-closure.csv` | +| Current overlay CSV | `/tmp/moonlight-linalg-current-closure.csv` | + +The pre-sat2 baseline was measured from a clean worktree at commit `a7ad9c832`. The `HEAD` baseline and current overlay were measured from clean detached worktrees so unrelated dirty files in the main checkout could not pollute the result. The current overlay is `HEAD` plus the linalg cut. + +## Focused graph-spectral P2 campaign (2026-07-21) + +The P1 restart/allocation cut left one honest algorithmic collapse: the real +10,000-vertex homology graph-Laplacian row still timed out after 30 seconds and +allocated 72,968,137,200 bytes. P2 retains graph-Laplacian provenance through +the checked operator constructor and sends only large, low-cardinality smallest +requests through heavy-edge cascadic descent. Arbitrary self-adjoint CSR remains +on the generic dense/restarted-Lanczos policy; symmetry is not counterfeit proof +of graph structure. + +The retained hierarchy follows the graph-specific cascadic method rather than +pretending an RHS-dependent approximate linear solve is an exact shift-invert +operator. Heavy-edge aggregates descend locally, mass-normalized prolongation +glues the coarse sections, and the authoritative fine CSR supplies the returned +residual evidence. This is consonant with the algebraic multilevel Fiedler +method of [Urschel and Hu](https://arxiv.org/abs/1412.0565) and the multigrid- +preconditioned block eigensolver evidence in +[BLOPEX](https://arxiv.org/abs/0705.2626). + +| Downstream successor carrier | P1 | P2 | P2 allocation / peak | +|---|---:|---:|---:| +| n=1,024, 3 smallest modes | 134.9 ms ± 28.7 ms | 161 ms ± 32 ms | 332 MB / 14 MB | +| n=10,000, 3 smallest modes | >30 s timeout | 48.0 ms ± 14 ms | 127 MB / 19 MB | +| n=50,000, 3 smallest modes | not retained | 265 ms ± 46 ms | 619 MB / 51 MB | +| n=100,000, 3 smallest modes | not retained | 496 ms ± 60 ms | 1.2 GB / 81 MB | + +The 10k row is therefore more than 625x faster than the bounded P1 observation +and allocates about 575x less. The 50k and 100k rows preserve the intended +near-linear scaling on this carrier family. The 1k row deliberately remains on +P1's restarted route: its timing intervals overlap and its allocation fell from +354,183,382 bytes to 332 MB, so there is no material small-carrier regression. + +```sh +cabal run moonlight-homology:bench:moonlight-homology-bench -j1 -- \ + --hide-progress --stdev 100 --timeout=10s \ + --pattern 'successor-carrier-1k' +RTS -T -s + +MOONLIGHT_HOMOLOGY_SPARSE_SPECTRAL_BENCH_ENABLE_LARGE=1 cabal run \ + moonlight-homology:bench:moonlight-homology-bench -j1 -- \ + --hide-progress --stdev 20 --timeout=30s \ + --pattern 'successor-carrier-10k' +RTS -T -s +``` + +## Focused generic-spectral P1 campaign (2026-07-21) + +The current live checkout at `32246311d7` exposed one concentrated failure: +the generic CSR n=384 values request crossed the old n=256 boundary into thick +restart, taking 2.564 s and allocating 10.090 GiB for one evaluation. The +retained cut widens the measured dense crossover to n=512, dispatches bounded +high-cardinality requests by demand through n=1,024, uses a 24-vector default +restart window, fuses scaled subtraction, and gives both projected owners one +generated-output basis-linear-combination kernel. + +```sh +MOONLIGHT_LINALG_BENCH_ENABLE_BROAD_LARGE=1 cabal run \ + moonlight-linalg:bench:moonlight-linalg-bench -j1 -- \ + --hide-progress --stdev 20 --timeout=20s \ + --pattern 'generic-csr-dense-fallback-values-pairs-384.values' + +MOONLIGHT_LINALG_BENCH_ENABLE_BROAD_LARGE=1 cabal run \ + moonlight-linalg:bench:moonlight-linalg-bench -j1 -- --once +RTS -T +``` + +| Evidence | Before | After | Result | +|---|---:|---:|---:| +| generic CSR n=384 values, calibrated | 2.564 s ± 97.5 ms | 134.8 ms ± 29.1 ms | 19.0x faster | +| generic CSR n=384 values, once allocation | 10.090 GiB | 9.813 MiB | about 1,053x lower | +| generic CSR n=384 pairs, once | 2.56 s-scale route | 144.3 ms / 10.147 MiB | bounded dense route | +| generic CSR base-threshold boundary | adjacent n=513 restart probe failed its residual gate after 4.66 s | n=512 dense: 304 ms ± 29 ms | measured boundary, not a fabricated restart timing | +| generic CSR n=513, all values | restart is the wrong selected-mode route | 344.5 ms ± 86.2 ms | demand-aware dense route | +| generic CSR n=1,024, all values | restart is the wrong selected-mode route | 2.575 s ± 141 ms / 61.399 MiB once | measured high-demand ceiling | +| default once sweep | 41.517 ms / 36.840 MiB | 36.758 ms / 34.741 MiB | no default regression | + +The n=513 low-demand banded fixture was also probed through restart. It failed +honestly at the selected-tridiagonal inverse-iteration residual gate rather than +being laundered into a timing row. The benchmark cover therefore retains the +two authoritative dense rows, while the downstream homology carrier remains the +real convergent restarted-Lanczos stress case. The 2026-07-06 sections below are +historical provenance for the old threshold, not current policy. + +```sh +cd /tmp/pale-meridian-linalg-pre-sat2-closure/compiler +cabal run moonlight-linalg:bench:moonlight-linalg-bench -- \ + --hide-progress --stdev 20 --timeout=1s \ + --csv /tmp/moonlight-linalg-pre-sat2-closure.csv + +cd /tmp/pale-meridian-linalg-baseline-closure/compiler +cabal run moonlight-linalg:bench:moonlight-linalg-bench -- \ + --hide-progress --stdev 20 --timeout=1s \ + --csv /tmp/moonlight-linalg-baseline-closure.csv + +cd /tmp/pale-meridian-linalg-current-closure/compiler +cabal run moonlight-linalg:bench:moonlight-linalg-bench -- \ + --hide-progress --stdev 20 --timeout=1s \ + --baseline /tmp/moonlight-linalg-baseline-closure.csv \ + --csv /tmp/moonlight-linalg-current-closure.csv +``` + +## Closure comparison + +The current overlay preserves `HEAD` performance: 76 shared sampled rows, 0 significant regressions, and `tasty-bench` reported every row as baseline-equivalent within the sampled window. Using `mean delta > max(2*stdev)` as the conservative CSV gate gives 76 noise-equivalent rows, 0 significant improvements, and 0 significant regressions against `HEAD`. + +Against the pre-sat2 cut, 36 rows are directly name-matched. The same conservative CSV gate finds 9 significant improvements and 0 significant regressions. These are the real hot-path wins; the rest are noise-equivalent. + +| Shared row | Pre-sat2 | Current | Delta | +|---|---:|---:|---:| +| sparse iterative solvers / n=64 / PCG SSOR | 1.869 ms ± 438.619 us | 32.189 us ± 8.142 us | 58.1x faster / 98.3% lower | +| sparse iterative solvers / n=64 / Jacobi diagonal | 75.490 us ± 15.566 us | 13.026 us ± 3.290 us | 5.8x faster / 82.7% lower | +| sparse iterative solvers / n=64 / PCG diagonal | 243.283 us ± 74.419 us | 51.258 us ± 15.002 us | 4.7x faster / 78.9% lower | +| sparse iterative solvers / n=64 / Richardson diagonal | 216.427 us ± 53.428 us | 51.554 us ± 13.876 us | 4.2x faster / 76.2% lower | +| sparse iterative solvers / n=64 / CG | 148.997 us ± 54.511 us | 46.743 us ± 13.892 us | 3.2x faster / 68.6% lower | +| sparse iterative solvers / n=64 / GMRES | 281.931 us ± 105.343 us | 95.055 us ± 37.649 us | 3.0x faster / 66.3% lower | +| sparse storage and packed kernels / n=512 / CSR matvec | 21.499 us ± 7.920 us | 8.631 us ± 3.307 us | 2.5x faster / 59.9% lower | +| sparse storage and packed kernels / n=512 / dense 32x32 -> CSR | 19.000 us ± 6.912 us | 11.761 us ± 3.299 us | 1.6x faster / 38.1% lower | +| selected tridiagonal eigenvalue solve / path-laplacian-10k | 23.030 us ± 7.104 us | 15.379 us ± 3.302 us | 1.5x faster / 33.2% lower | + +The projected-block policy rows are not direct name matches because the old fake policy suite was deleted. The replacement rows expose demand and structure explicitly: values rows are values-only, pair rows return `Eigenpairs`, and dense oracle rows are reference evidence rather than production fallback. + +| Case | Old reuse-first policy | Old dense-fallback policy | Current structured values | Current structured pairs | Current dense oracle | +|---|---:|---:|---:|---:|---:| +| block-clustered-24 | 981.278 us | 990.648 us | 28.217 us | 90.107 us | 968.766 us | +| block-separated-24 | 953.539 us | 935.845 us | 29.654 us | 87.952 us | 897.297 us | + +For the same default fixtures, the current structured projected-block values rows are 32.2x-34.8x faster than the old reuse-first policy rows, and the current pair rows are 10.6x-11.0x faster than the old dense-fallback policy rows. That is not a compatibility story; the old policy surface is gone, and the measured demand-specific replacement is cheaper. + +## Current evidence rows + +The native LAPACK rows call `Moonlight.LinAlg.Native` with `EigenRequest` and consume the same `Eigenpairs` owner as the pure spectral surface. Benchmark callers do not import raw `FortranIndexRange` or driver functions. DSTEMR, DSYEVX, and DSBEVX vocabulary stays behind the native boundary. + +| Row | Mean | 2*Stdev | Purpose | +|---|---:|---:|---| +| path-laplacian-512 DSTEMR selected tridiagonal values modes=4 | 2.92 ms | 993 us | native tridiagonal values-only request | +| path-laplacian-512 DSTEMR selected tridiagonal pairs modes=4 | 3.14 ms | 955 us | native tridiagonal pair request through `Eigenpairs` | +| generic-tridiagonal-512 DSTEMR selected tridiagonal values modes=4 | 582 us | 210 us | native generic-tridiagonal values-only request | +| generic-tridiagonal-512 DSTEMR selected tridiagonal pairs modes=4 | 1.01 ms | 234 us | native generic-tridiagonal pair request through `Eigenpairs` | +| block-clustered-24 DSYEVX dense values | 72.9 us | 28.2 us | dense selected values oracle through public `EigenRequest` | +| block-clustered-24 DSYEVX dense pairs | 111 us | 33.2 us | dense selected pairs oracle through `Eigenpairs` | +| block-separated-24 DSYEVX dense values | 73.9 us | 29.5 us | dense selected values oracle through public `EigenRequest` | +| block-separated-24 DSYEVX dense pairs | 109 us | 26.9 us | dense selected pairs oracle through `Eigenpairs` | +| block-clustered-24 DSBEVX projected block values | 29.1 us | 8.95 us | native symmetric-band selected values through public `EigenRequest` | +| block-clustered-24 DSBEVX projected block pairs | 88.3 us | 26.6 us | native symmetric-band selected pairs through `Eigenpairs` | +| block-separated-24 DSBEVX projected block values | 31.3 us | 11.1 us | native symmetric-band selected values through public `EigenRequest` | +| block-separated-24 DSBEVX projected block pairs | 90.5 us | 28.5 us | native symmetric-band selected pairs through `Eigenpairs` | +| projected tridiagonal-path-512 values | 221 us | 61.2 us | projected tridiagonal values-only path | +| projected tridiagonal-path-512 pairs | 547 us | 145 us | projected tridiagonal pair path with ambient lift and residuals | +| projected block-clustered-24 values | 28.2 us | 7.26 us | projected block values-only path | +| projected block-clustered-24 pairs | 90.1 us | 32.9 us | projected block pair path through the native symmetric-band executor | +| projected block-separated-24 values | 29.7 us | 7.42 us | projected block values-only path | +| projected block-separated-24 pairs | 88.0 us | 26.4 us | projected block pair path through the native symmetric-band executor | + +## Current default group totals + +| Group | Rows | Sum of row means | +|---|---:|---:| +| dense row validation surface | 4 | 809 us | +| dense decomposition and solvers | 7 | 793 us | +| sparse storage and packed kernels | 8 | 1.44 ms | +| sparse iterative solvers | 6 | 290 us | +| spectral demand dispatch | 20 | 345.19 ms | +| domain algebra, exterior powers, GF2 | 4 | 290 us | +| geometry and statics | 4 | 479 us | +| selected tridiagonal eigenvalue solve | 1 | 15.4 us | +| native LAPACK symmetric eigensolve | 14 | 8.74 ms | +| projected structured eigensolve | 8 | 2.87 ms | +| **Total** | **76** | **360.91 ms** | + +## Allocation and retained-live evidence + +The once runner reports per-row allocation when RTS stats are enabled. Each row forces a major GC before and after the measured action, records the allocation delta, reports live bytes after the post-row major GC, and reports the RTS process maximum residency observed after the row. Row retained-live is row-local; RTS maximum residency is process cumulative and is therefore labeled as such instead of being smuggled in as per-row truth. + +```sh +cd /Users/bluerose/Developer/pale-meridian/compiler +cabal run moonlight-linalg:bench:moonlight-linalg-bench -- --once +RTS -T +``` + +Current result: + +| Rows | CPU once time | Checksum | Heap allocated | Max retained live after row GC | Process maximum residency | +|---:|---:|---:|---:|---:|---:| +| 66 | 360.970 ms | 3436947.625310 | 746.235 MiB | 1.502 MiB | 2.186 MiB | + +Highest allocation rows: + +| Row | Allocated | Live after major GC | Process maximum residency | +|---|---:|---:|---:| +| spectral demand dispatch / generic-tridiagonal-values-pairs-512 / pairs | 498.052 MiB | 173.648 KiB | 2.186 MiB | +| spectral demand dispatch / reducible-tridiagonal-values-pairs-512 / pairs | 49.684 MiB | 178.445 KiB | 2.186 MiB | +| spectral demand dispatch / reducible-tridiagonal-values-pairs-512 / values | 37.229 MiB | 176.125 KiB | 2.186 MiB | +| spectral demand dispatch / generic-csr-fallback-values-pairs-96 / pairs | 19.916 MiB | 183.242 KiB | 2.186 MiB | +| spectral demand dispatch / generic-csr-fallback-values-pairs-96 / values | 18.298 MiB | 180.922 KiB | 2.186 MiB | +| spectral demand dispatch / generic-tridiagonal-values-pairs-512 / values | 11.277 MiB | 171.328 KiB | 171.328 KiB | +| sparse storage and packed kernels / n=512 / CSR -> CSC | 7.254 MiB | 161.859 KiB | 161.859 KiB | +| sparse storage and packed kernels / n=512 / graph Laplacian construction | 6.057 MiB | 168.820 KiB | 168.820 KiB | +| projected structured eigensolve / block-clustered-24 generic dense oracle | 5.733 MiB | 248.938 KiB | 2.186 MiB | +| projected structured eigensolve / block-separated-24 generic dense oracle | 5.648 MiB | 255.992 KiB | 2.186 MiB | + +The detached pre-sat2 runner did not have per-row RTS stats. Its old whole-suite `+RTS -s` aggregate remains useful only as a process-level reference, not a row-normalized comparison: 44 rows, 26.847 ms CPU once time, 133,145,352 bytes allocated, 309,632 bytes maximum residency. The row-matched speed gate remains the CSV comparison above. + +## Reading the cut + +The benchmark cover now measures the requested spectral demand split instead of smuggling solver-selection toggles through the suite. The slow generic-tridiagonal pair row is intentionally visible; values are cheap, full pairs are not. + +Sparse CG, PCG, GMRES, Jacobi, and Richardson now run through a sealed `Double` `ST` workspace. The public preconditioner is an abstract ADT, not a closure-shaped backdoor. The old persistent sparse-solver row walkers were deleted; the mutable arena is sealed behind immutable `U.Vector Double` results. + +The native tridiagonal evidence has separate values-only and pair rows. Projected evidence has separate tridiagonal values/pairs, block values/pairs, DSYEVX dense oracle rows, dense projected-space oracle rows, and DSBEVX symmetric-band rows. Comparable current rows carry the no-regression claim, and values-only rows prove demand is no longer silently flattened. + +## Wave-2 finalization rows (2026-07-06) + +Two package-wide causes were found and repaired during Wave-2 reconciliation, and every row below reflects both: + +- The IEEE boundary constants (`epsDouble`, `safeMinimumDouble`, `maxFiniteDouble`) were decimal scientific literals whose exact rationals carry hundreds of digits. GHC 9.14 does not constant-fold `fromRational` at those exponents, and the `INLINE` pragmas pasted the runtime conversion — GMP integer division and gcd — into every consumer's hot loop, including the implicit-QL negligibility test and the Sturm bisection tolerances. They are now `encodeFloat` forms under `NOINLINE`, evaluated once. +- The libraries built at cabal's default `-O1` while every probe and target assumed `-O2`. SpecConstr and LiberateCase are worth ×3–4 on the Krylov and QL inner loops. `-O2` now lives in `shared-properties`. + +| Row | Mean | Prior | Purpose | +|---|---:|---:|---| +| symmetric eigen pure 12x12 | 17.6 us | 208 us | dense unchecked eigensolve, certification split off the hot path | +| selected tridiagonal path-laplacian-10k | 12.8 us | 305 ms / 498 MiB | bisection + shifted inverse iteration replacing the QL global solve | +| reducible-tridiagonal-values-pairs-512 pairs | 356 us | 7.43 ms | reducible split + selected solve per block | +| generic-csr-fallback-values-pairs-96 values | 22.8 ms | timeout-scale | bordered Wu–Simon thick restart, capacity 32, 4 modes | +| generic-csr-fallback-values-pairs-192 values | 1.50 s | 6.08 s | same path, 6 modes; convergence-bound, see below | +| generic-csr-thick-restart-values-pairs-384 values | 2.09 s | never completed | gated broad-large row | +| svd 12x12 | 32.4 us | — | one-sided Jacobi, Gram path deleted | + +The generic CSR fallback rows are convergence-bound, not overhead-bound: the banded SPD fixture's smallest eigenvalues cluster at gaps near 1e-5, and the restart loop needs ~50 cycles at capacity 32 to lock four of them, with the cycle interior running at flop cost. Locking is honest — candidates pass a scalar projected-residual gate, are lifted, and must then survive the ambient-residual predicate; failures demote back to the unlocked pool. Dense eigensolve of the densified operator beats restart at every benched dimension (~2.5 ms at n=96), so the fallback dispatch gains a dimension threshold in Wave 3; the restart path remains the only route where densification is prohibitive. + +## Wave-3 dispatch rows (2026-07-06) + +The generic CSR fallback now dispatches on dimension: at or below `denseSpectralFallbackDimensionThreshold = 256` the operator is densified once (one apply per basis vector) and solved by the flat unchecked dense eigensolver, with pair residuals computed against the densified matrix rather than per-pair operator re-applies; above the threshold the thick-restart path stands. + +| Row | Mean | Prior (restart-only) | Route | +|---|---:|---:|---| +| generic-csr-fallback-values-pairs-96 values | 2.40 ms | 22.8 ms | dense | +| generic-csr-fallback-values-pairs-192 values | 17.6 ms | 1.50 s | dense | +| generic-csr-thick-restart-values-pairs-384 values | 2.09 s | 2.09 s | restart (above threshold, by design) | diff --git a/docs/CONSTRUCTION.md b/docs/CONSTRUCTION.md new file mode 100644 --- /dev/null +++ b/docs/CONSTRUCTION.md @@ -0,0 +1,108 @@ +# Construction + +Moonlight constructors validate once and return the owning carriers directly. +There is no deferred linear-algebra expression tree. + +## Dense matrices + +```haskell +matrixValue <- + matrixRows @2 @2 + [ [1.0, 2.0], + [3.0, 4.0] + ] +``` + +`matrixRows` is row-major. Type-level dimensions are checked against the nested +rows, including zero-row matrices whose column count cannot be inferred from the +value alone. Treat `DenseRows` and nested-list construction as validated +authoring/projection surfaces; they preserve shape failures, not dense hot-path +performance. + +For dynamic dimensions: + +```haskell +matrixValue <- + dynMatrixFromRows + [ [1.0, 2.0], + [3.0, 4.0] + ] +``` + +## Structured sparse matrices + +```haskell +sparseOperator <- + tridiagonalCSR + [2.0, 2.0, 2.0] + [-1.0, -1.0] +``` + +```haskell +pathOperator <- + pathLaplacianCSR 8 +``` + +```haskell +graphOperator <- + graphLaplacianCSR + ["a", "b", "c"] + [ GraphEdge "a" "b" 1.0, + GraphEdge "b" "c" 2.0 + ] +``` + +The graph vertex list defines matrix row and column order. Edges are undirected; +parallel and reversed edges are combined. Weights must be finite and +non-negative. + +## Matrix-free eigen requests + +```haskell +count <- mkPositiveCount 4 +operator <- declaredSelfAdjointVectorLinearOperator dimension applyA +let config = + withEigenFallbackInitialVector seed defaultEigenSolveConfig +eigenvalues <- solveEigenRequest config operator (EigenvaluesRequest SmallestEigenvalues count) +``` + +```haskell +count <- mkPositiveCount 4 +operator <- selfAdjointCSRLinearOperator csr +eigenpairs <- solveEigenRequest defaultEigenSolveConfig operator (EigenpairsRequest SmallestEigenvalues count) +``` + +The request lives above Krylov: values-only requests can avoid ambient vector +lifting, while eigenpair requests return a contiguous `Eigenpairs` payload. + +## Selected structured spectra + +```haskell +count <- mkPositiveCount 4 +operator <- pathLaplacianLinearOperator 10000 +eigenvalues <- + solveEigenRequest + defaultEigenSolveConfig + operator + (EigenvaluesRequest SmallestEigenvalues count) +``` + +This is the package's flagship hot path: Moonlight recognizes certified operator structure and computes only the requested spectral data. Use this route when downstream code needs a small spectral slice; do not build dense rows just to throw most of the spectrum away. + +## Force networks + +```haskell +forceNetwork <- + network + [ support "a" (Vec3 0.0 0.0 0.0), + load + "b" + (Vec3 0.0 1.0 0.0) + (Vec3 0.0 (-10.0) 0.0), + member "a" "b" + ] +``` + +Declarations may occur in any order. Repeated loads add, repeated members are +idempotent, and repeated node declarations must agree on position. `network` +returns a validated `ForceNetwork` directly. diff --git a/moonlight-linalg.cabal b/moonlight-linalg.cabal new file mode 100644 --- /dev/null +++ b/moonlight-linalg.cabal @@ -0,0 +1,427 @@ +cabal-version: 3.0 +name: moonlight-linalg +version: 0.1.0.0 +homepage: https://github.com/PaleRoses/moonlight +bug-reports: https://github.com/PaleRoses/moonlight/issues +synopsis: Dense tensor and algebraic matrix core for Pale Meridian. +description: Typed dense and sparse matrices, GF(2) and Smith-normal-form backends, symmetry-indexed operators, and restarted Krylov and Lanczos spectral solvers over a typed shape-and-domain failure vocabulary. +license: MIT +license-file: LICENSE +copyright: (c) 2026 Blue Rose +author: Blue Rose +maintainer: rosaliafialkova@gmail.com +category: Math +build-type: Simple +tested-with: GHC == 9.14.1 +extra-doc-files: + README.md + CHANGELOG.md + THIRD_PARTY_NOTICES.md + docs/ARCHITECTURE.md + docs/CONSTRUCTION.md + docs/BENCHMARKS-m4-pro.md + +common shared-properties + default-language: GHC2024 + ghc-options: + -Wall + -Wcompat + -Wincomplete-record-updates + -Wincomplete-uni-patterns + -Wredundant-constraints + -Wpartial-fields + -Wno-missing-import-lists + -O2 + default-extensions: + TypeFamilies + UndecidableInstances + FunctionalDependencies + +library moonlight-linalg-carrier + import: shared-properties + visibility: public + hs-source-dirs: src-carrier + exposed-modules: + Moonlight.LinAlg.Internal.DenseList + Moonlight.LinAlg.Internal.Discrete + Moonlight.LinAlg.Internal.GF2.SparseColumn + Moonlight.LinAlg.Internal.GF2.Xor + Moonlight.LinAlg.Internal.Primitives + Moonlight.LinAlg.Internal.Storage + Moonlight.LinAlg.Internal.VectorOps + Moonlight.LinAlg.Pure.Dense.Flat + Moonlight.LinAlg.Pure.Dense.Rows + Moonlight.LinAlg.Pure.Dense.Types + build-depends: + base >= 4.22 && < 5 + , containers >= 0.6 && < 0.9 + , moonlight-algebra:abstract >= 0.1 && < 0.2 + , moonlight-core >= 0.1 && < 0.2 + , primitive >= 0.8 && < 0.10 + , vector >= 0.13 && < 0.14 + +library moonlight-linalg-structured + import: shared-properties + visibility: public + hs-source-dirs: src-structured + exposed-modules: + Moonlight.LinAlg.Pure.Structured.BlockTridiagonal + Moonlight.LinAlg.Pure.Structured.Tridiagonal + build-depends: + base >= 4.22 && < 5 + , moonlight-core >= 0.1 && < 0.2 + , primitive >= 0.8 && < 0.10 + , vector >= 0.13 && < 0.14 + +library moonlight-linalg-eigen + import: shared-properties + visibility: public + hs-source-dirs: src-eigen + exposed-modules: + Moonlight.LinAlg.Internal.Eigen.DenseWork + Moonlight.LinAlg.Internal.Eigen.Householder + Moonlight.LinAlg.Internal.Eigen.Input + Moonlight.LinAlg.Internal.Eigen.Kernels + Moonlight.LinAlg.Internal.Eigen.Residual + Moonlight.LinAlg.Internal.Eigen.Symmetric + Moonlight.LinAlg.Internal.Eigen.Tridiagonal + build-depends: + base >= 4.22 && < 5 + , moonlight-core >= 0.1 && < 0.2 + , moonlight-linalg:moonlight-linalg-carrier + , primitive >= 0.8 && < 0.10 + , vector >= 0.13 && < 0.14 + +library moonlight-linalg-geometry + import: shared-properties + visibility: public + hs-source-dirs: src-geometry + exposed-modules: + Moonlight.LinAlg.Pure.Geometry.AABB + Moonlight.LinAlg.Pure.Geometry.AABB2 + Moonlight.LinAlg.Pure.Geometry.Frame + Moonlight.LinAlg.Pure.Geometry.Symmetric + Moonlight.LinAlg.Pure.Geometry.Transform.Affine + Moonlight.LinAlg.Pure.Geometry.Vec2 + Moonlight.LinAlg.Pure.Geometry.Vec3 + build-depends: + base >= 4.22 && < 5 + , deepseq >= 1.4 && < 1.6 + , moonlight-algebra:abstract >= 0.1 && < 0.2 + , moonlight-core >= 0.1 && < 0.2 + , moonlight-linalg:moonlight-linalg-carrier + , vector >= 0.13 && < 0.14 + +library moonlight-linalg-dense + import: shared-properties + visibility: public + hs-source-dirs: src-dense + exposed-modules: + Moonlight.LinAlg.Internal.Backend.Core + Moonlight.LinAlg.Internal.Backend.Elimination + Moonlight.LinAlg.Internal.Backend.PLU + Moonlight.LinAlg.Internal.Backend.RREF + Moonlight.LinAlg.Internal.Backend.RowOps + Moonlight.LinAlg.Internal.Backend.RowStore + Moonlight.LinAlg.Internal.Backend.Smith + Moonlight.LinAlg.Internal.Dense.DoubleFactorization + Moonlight.LinAlg.Internal.Dense.OneSidedJacobiSVD + Moonlight.LinAlg.Pure.Dense.Basic + Moonlight.LinAlg.Pure.Dense.Block + Moonlight.LinAlg.Pure.Dense.Decomposition + Moonlight.LinAlg.Pure.Dense.Dynamic + Moonlight.LinAlg.Pure.Dense.Exterior + Moonlight.LinAlg.Pure.Dense.Field + Moonlight.LinAlg.Pure.Dense.GF2 + Moonlight.LinAlg.Pure.Dense.Solver + build-depends: + base >= 4.22 && < 5 + , moonlight-algebra:abstract >= 0.1 && < 0.2 + , moonlight-core >= 0.1 && < 0.2 + , moonlight-linalg:moonlight-linalg-carrier + , moonlight-linalg:moonlight-linalg-eigen + , primitive >= 0.8 && < 0.10 + , vector >= 0.13 && < 0.14 + +library moonlight-linalg-domain + import: shared-properties + visibility: public + hs-source-dirs: src-domain + exposed-modules: + Moonlight.LinAlg.Pure.Domain.Bareiss + Moonlight.LinAlg.Pure.Domain.Smith + Moonlight.LinAlg.Pure.Domain.Smith.Multimodular + Moonlight.LinAlg.Pure.Domain.Smith.Witnessed + build-depends: + base >= 4.22 && < 5 + , moonlight-algebra:abstract >= 0.1 && < 0.2 + , moonlight-core >= 0.1 && < 0.2 + , moonlight-linalg:moonlight-linalg-carrier + , moonlight-linalg:moonlight-linalg-dense + , vector >= 0.13 && < 0.14 + +library moonlight-linalg-sparse + import: shared-properties + visibility: public + hs-source-dirs: src-sparse + exposed-modules: + Moonlight.LinAlg.Pure.Sparse.Assembly + Moonlight.LinAlg.Pure.Sparse.Packed + Moonlight.LinAlg.Pure.Sparse.Solver.CG + Moonlight.LinAlg.Pure.Sparse.Solver.Common + Moonlight.LinAlg.Pure.Sparse.Solver.GMRES + Moonlight.LinAlg.Pure.Sparse.Solver.IncompleteCholesky0 + Moonlight.LinAlg.Pure.Sparse.Solver.Mutable + Moonlight.LinAlg.Pure.Sparse.Solver.Preconditioner + Moonlight.LinAlg.Pure.Sparse.Solver.Stationary + Moonlight.LinAlg.Pure.Sparse.Solver.Types + Moonlight.LinAlg.Pure.Sparse.Structured + Moonlight.LinAlg.Pure.Sparse.Types + build-depends: + base >= 4.22 && < 5 + , containers >= 0.6 && < 0.9 + , moonlight-core >= 0.1 && < 0.2 + , moonlight-linalg:moonlight-linalg-carrier + , moonlight-linalg:moonlight-linalg-structured + , primitive >= 0.8 && < 0.10 + , transformers >= 0.6 && < 1 + , vector >= 0.13 && < 0.14 + +library moonlight-linalg-statics + import: shared-properties + visibility: public + hs-source-dirs: src-statics + exposed-modules: + Moonlight.LinAlg.Pure.Statics.Algebra + Moonlight.LinAlg.Pure.Statics.Compile + Moonlight.LinAlg.Pure.Statics.Core + Moonlight.LinAlg.Pure.Statics.Network + Moonlight.LinAlg.Pure.Statics.Types + build-depends: + base >= 4.22 && < 5 + , containers >= 0.6 && < 0.9 + , moonlight-core >= 0.1 && < 0.2 + , moonlight-linalg:moonlight-linalg-carrier + , moonlight-linalg:moonlight-linalg-dense + , moonlight-linalg:moonlight-linalg-geometry + +library moonlight-linalg-spectral + import: shared-properties + visibility: public + hs-source-dirs: src-spectral + exposed-modules: + Moonlight.LinAlg.Pure.Operator + Moonlight.LinAlg.Pure.Krylov.Arnoldi + Moonlight.LinAlg.Pure.Krylov.Block + Moonlight.LinAlg.Pure.Krylov.Config + Moonlight.LinAlg.Pure.Krylov.Decomposition + Moonlight.LinAlg.Pure.Krylov.Internal + Moonlight.LinAlg.Pure.Krylov.Lanczos + Moonlight.LinAlg.Pure.Krylov.Projected + Moonlight.LinAlg.Pure.Krylov.Selection + Moonlight.LinAlg.Pure.Krylov.SelectedTridiagonal + Moonlight.LinAlg.Pure.Spectral.Request + Moonlight.LinAlg.Pure.Spectral.Result + Moonlight.LinAlg.Pure.Spectral.Solve + other-modules: + Moonlight.LinAlg.Pure.Operator.Internal + Moonlight.LinAlg.Pure.Krylov.CascadicGraph + build-depends: + base >= 4.22 && < 5 + , containers >= 0.6 && < 0.9 + , moonlight-core >= 0.1 && < 0.2 + , moonlight-linalg:moonlight-linalg-carrier + , moonlight-linalg:moonlight-linalg-eigen + , moonlight-linalg:moonlight-linalg-sparse + , moonlight-linalg:moonlight-linalg-structured + , primitive >= 0.8 && < 0.10 + , vector >= 0.13 && < 0.14 + +library moonlight-linalg-native + import: shared-properties + visibility: public + hs-source-dirs: src-native + exposed-modules: + Moonlight.LinAlg.Effect.Native.Dispatch + Moonlight.LinAlg.Effect.Native.LAPACK + Moonlight.LinAlg.Native + c-sources: + cbits/moonlight_linalg_native.c + build-depends: + base >= 4.22 && < 5 + , moonlight-core >= 0.1 && < 0.2 + , moonlight-linalg:moonlight-linalg-carrier + , moonlight-linalg:moonlight-linalg-dense + , moonlight-linalg:moonlight-linalg-eigen + , moonlight-linalg:moonlight-linalg-spectral + , moonlight-linalg:moonlight-linalg-structured + , vector >= 0.13 && < 0.14 + if os(darwin) + frameworks: + Accelerate + else + extra-libraries: + lapack + blas + +library + import: shared-properties + hs-source-dirs: src-public + exposed-modules: + Moonlight.LinAlg + Moonlight.LinAlg.Dense + Moonlight.LinAlg.Dense.Block + Moonlight.LinAlg.Dense.Decomposition + Moonlight.LinAlg.Dense.Exterior + Moonlight.LinAlg.Dense.Field + Moonlight.LinAlg.Dense.GF2 + Moonlight.LinAlg.Dense.Primitives + Moonlight.LinAlg.Dense.Rows + Moonlight.LinAlg.Dense.Solver + Moonlight.LinAlg.Domain + Moonlight.LinAlg.Geometry + Moonlight.LinAlg.Krylov + Moonlight.LinAlg.Operator + Moonlight.LinAlg.Sparse + Moonlight.LinAlg.Spectral + Moonlight.LinAlg.Statics + build-depends: + base >= 4.22 && < 5 + , moonlight-linalg:moonlight-linalg-carrier + , moonlight-linalg:moonlight-linalg-dense + , moonlight-linalg:moonlight-linalg-domain + , moonlight-linalg:moonlight-linalg-geometry + , moonlight-linalg:moonlight-linalg-sparse + , moonlight-linalg:moonlight-linalg-spectral + , moonlight-linalg:moonlight-linalg-statics + +library moonlight-linalg-laws + import: shared-properties + visibility: public + hs-source-dirs: src-laws + -- This test-law vocabulary has no production or benchmark consumers. + ghc-options: -O0 + exposed-modules: + Moonlight.LinAlg.Effect.Harness + Moonlight.LinAlg.Effect.Harness.Core + Moonlight.LinAlg.Effect.Harness.Dense + Moonlight.LinAlg.Effect.Harness.Decomposition + Moonlight.LinAlg.Effect.Harness.Domain + Moonlight.LinAlg.Effect.Harness.Field + Moonlight.LinAlg.Effect.Harness.Geometry + Moonlight.LinAlg.Effect.Harness.KrylovSpectral + Moonlight.LinAlg.Effect.Harness.Operator + Moonlight.LinAlg.Effect.Harness.Preconditioner + Moonlight.LinAlg.Effect.Harness.Sparse + Moonlight.LinAlg.Effect.Harness.Statics + Moonlight.LinAlg.Effect.LawNames + Moonlight.LinAlg.Effect.Laws + build-depends: + base >= 4.22 && < 5 + , containers >= 0.6 && < 0.9 + , moonlight-core >= 0.1 && < 0.2 + , moonlight-linalg + , moonlight-linalg:moonlight-linalg-carrier + , moonlight-linalg:moonlight-linalg-eigen + , moonlight-pale:test-laws >= 0.1 && < 0.2 + , tasty >= 1.4 && < 1.6 + , tasty-hunit >= 0.10 && < 0.11 + , tasty-quickcheck >= 0.10 && < 0.12 + , vector >= 0.13 && < 0.14 + +test-suite moonlight-linalg-laws-test + import: shared-properties + type: exitcode-stdio-1.0 + hs-source-dirs: test-laws + main-is: Main.hs + ghc-options: -Wall -Wcompat -O0 + build-depends: + base >= 4.22 && < 5 + , moonlight-linalg:moonlight-linalg-laws + , tasty >= 1.4 + +test-suite moonlight-linalg-test + type: exitcode-stdio-1.0 + default-language: GHC2024 + hs-source-dirs: test test/dense test/domain test/geometry test/sparse test/statics test/spectral test/architecture test/support + main-is: Main.hs + ghc-options: -Wall -Wcompat -O0 + default-extensions: + TypeFamilies + other-modules: + AdvancedSpec + ArchitectureSpec + BasicSpec + BlockSpec + DenseFlatSpec + DenseRowsSpec + KrylovSpec + FieldSpec + DomainSpec + DynamicSpec + ExteriorSpec + GF2Spec + GeometryStorageSpec + SymmetricSpec + StaticsSpec + SparseSolverSpec + SparsePackedSpec + Helpers + build-depends: + base >= 4.22 && < 5 + , moonlight-linalg + , moonlight-linalg:moonlight-linalg-carrier + , moonlight-linalg:moonlight-linalg-dense + , moonlight-linalg:moonlight-linalg-domain + , moonlight-linalg:moonlight-linalg-eigen + , moonlight-linalg:moonlight-linalg-native + , moonlight-linalg:moonlight-linalg-sparse + , moonlight-linalg:moonlight-linalg-spectral + , moonlight-linalg:moonlight-linalg-structured + , moonlight-core >= 0.1 && < 0.2 + , moonlight-algebra:abstract >= 0.1 && < 0.2 + , containers >= 0.6 && < 0.9 + , directory >= 1.3 && < 1.4 + , filepath >= 1.4 + , tasty >= 1.4 + , tasty-hunit >= 0.10 + , tasty-quickcheck >= 0.10 + , vector >= 0.13 && < 0.14 + +benchmark moonlight-linalg-bench + import: shared-properties + type: exitcode-stdio-1.0 + hs-source-dirs: bench bench/dense bench/domain bench/statics bench/native bench/sparse bench/spectral bench/support + main-is: Main.hs + other-modules: + DenseCore + DenseDecomposition + DomainAlgebra + Env + Fixtures + GeometryStatics + NativeLapack + ProjectedBlock + SparseKrylov + SparseSolvers + SpectralDispatch + SparseStorage + Types + ghc-options: -O2 -rtsopts + ghc-prof-options: -fprof-auto-top + build-depends: + base >= 4.22 && < 5 + , containers >= 0.6 && < 0.9 + , deepseq >= 1.4 && < 1.6 + , moonlight-linalg + , moonlight-linalg:moonlight-linalg-dense + , moonlight-linalg:moonlight-linalg-native + , moonlight-linalg:moonlight-linalg-spectral + , moonlight-linalg:moonlight-linalg-structured + , tasty-bench + , vector >= 0.13 && < 0.14 + +source-repository head + type: git + location: https://github.com/PaleRoses/moonlight.git + subdir: moonlight-linalg diff --git a/src-carrier/Moonlight/LinAlg/Internal/DenseList.hs b/src-carrier/Moonlight/LinAlg/Internal/DenseList.hs new file mode 100644 --- /dev/null +++ b/src-carrier/Moonlight/LinAlg/Internal/DenseList.hs @@ -0,0 +1,42 @@ +module Moonlight.LinAlg.Internal.DenseList + ( dotProductWith, + matrixVectorProductWith, + zipMatrixWith, + scaleMatrixWith, + outerProductWith, + ) +where + +import Data.Function ((&)) +import Prelude + +dotProductWith :: (left -> right -> product) -> (product -> accumulator -> accumulator) -> accumulator -> [left] -> [right] -> Either String accumulator +dotProductWith multiply append zeroValue left right = + if leftLength == rightLength + then Right (foldr append zeroValue (zipWith multiply left right)) + else Left ("length mismatch (left=" <> show leftLength <> ", right=" <> show rightLength <> ")") + where + leftLength = length left + rightLength = length right + +matrixVectorProductWith :: (entry -> value -> product) -> (product -> accumulator -> accumulator) -> accumulator -> [[entry]] -> [value] -> Either String [accumulator] +matrixVectorProductWith multiply append zeroValue matrixRows vectorValue = + traverse + (\(rowIndex, rowValues) -> + dotProductWith multiply append zeroValue rowValues vectorValue + & either + (Left . (\message -> "row " <> show rowIndex <> ": " <> message)) + Right + ) + (zip [0 :: Int ..] matrixRows) + +zipMatrixWith :: (left -> right -> result) -> [[left]] -> [[right]] -> [[result]] +zipMatrixWith combine = zipWith (zipWith combine) + +scaleMatrixWith :: (scalar -> value -> result) -> scalar -> [[value]] -> [[result]] +scaleMatrixWith multiply scalarValue = + map (map (multiply scalarValue)) + +outerProductWith :: (left -> right -> result) -> [left] -> [right] -> [[result]] +outerProductWith multiply left right = + map (\leftEntry -> map (multiply leftEntry) right) left diff --git a/src-carrier/Moonlight/LinAlg/Internal/Discrete.hs b/src-carrier/Moonlight/LinAlg/Internal/Discrete.hs new file mode 100644 --- /dev/null +++ b/src-carrier/Moonlight/LinAlg/Internal/Discrete.hs @@ -0,0 +1,190 @@ +module Moonlight.LinAlg.Internal.Discrete + ( GF2 (..), + gf2Zero, + gf2One, + gf2ToBool, + gf2FromBool, + rankPackedRows, + PackedBitMatrix (..), + matrixRowWords, + packedBitMatrixFromRowMajor, + packedBitMatrixFromXorEntries, + ) +where + +import Data.Bits (Bits (bit, xor, (.|.)), zeroBits) +import qualified Data.IntMap.Strict as IntMap +import Data.Kind (Type) +import qualified Data.Vector.Unboxed as U +import Data.Word (Word64) +import Moonlight.Algebra.Pure.Ring + ( EuclideanDomain (..), + GCDDomain (..), + IntegralDomain (..), + ) +import Moonlight.Core + ( AdditiveGroup (..), + AdditiveMonoid (..), + CommutativeRing, + Field (..), + MultiplicativeMonoid (..), + Ring, + Semiring, + ) +import qualified Moonlight.LinAlg.Internal.GF2.Xor as GF2Xor +import Prelude + +type GF2 :: Type +data GF2 = GF2Zero | GF2One + deriving stock (Eq, Ord, Show) + +instance Num GF2 where + (+) = add + (*) = mul + negate = id + abs = id + signum value = + case value of + GF2Zero -> GF2Zero + GF2One -> GF2One + fromInteger integerValue = + if odd integerValue + then GF2One + else GF2Zero + +gf2Zero :: GF2 +gf2Zero = GF2Zero + +gf2One :: GF2 +gf2One = GF2One + +gf2FromBool :: Bool -> GF2 +gf2FromBool flag = if flag then GF2One else GF2Zero + +gf2ToBool :: GF2 -> Bool +gf2ToBool value = case value of + GF2Zero -> False + GF2One -> True + +type PackedBitMatrix :: Type +data PackedBitMatrix = PackedBitMatrix + { packedRows :: Int, + packedCols :: Int, + packedWordsPerRow :: Int, + packedWords :: U.Vector Word64 + } + deriving stock (Eq, Show) + +wordWidth :: Int -> Int +wordWidth count + | count <= 0 = 0 + | otherwise = + count `Prelude.div` 64 + + if count `mod` 64 == 0 then 0 else 1 + +packRowMajorBits :: Int -> Int -> [GF2] -> U.Vector Word64 +packRowMajorBits rowCount columnCount values = + let wordsPerRowValue = wordWidth columnCount + packedWordCount = rowCount * wordsPerRowValue + updateWord mapValue (linearIndex, bitValue) = + if bitValue == GF2Zero || columnCount <= 0 + then mapValue + else + let rowIndex = linearIndex `Prelude.div` columnCount + columnIndex = linearIndex `mod` columnCount + wordIndex = (rowIndex * wordsPerRowValue) + (columnIndex `Prelude.div` 64) + bitIndex = columnIndex `mod` 64 + bitMask = bit bitIndex :: Word64 + in IntMap.insertWith (.|.) wordIndex bitMask mapValue + packedMap = foldl' updateWord IntMap.empty (zip [0 :: Int ..] values) + in U.generate packedWordCount (\wordIndex -> IntMap.findWithDefault zeroBits wordIndex packedMap) + +packedBitMatrixFromRowMajor :: Int -> Int -> [GF2] -> PackedBitMatrix +packedBitMatrixFromRowMajor rowCount columnCount values = + PackedBitMatrix + { packedRows = rowCount, + packedCols = columnCount, + packedWordsPerRow = wordWidth columnCount, + packedWords = packRowMajorBits rowCount columnCount values + } + +packedBitMatrixFromXorEntries :: Int -> Int -> [(Int, Int)] -> PackedBitMatrix +packedBitMatrixFromXorEntries rowCount columnCount entries = + let wordsPerRowValue = wordWidth columnCount + packedWordCount = rowCount * wordsPerRowValue + updateWord mapValue (rowIndex, columnIndex) = + let wordIndex = (rowIndex * wordsPerRowValue) + (columnIndex `Prelude.div` 64) + bitIndex = columnIndex `mod` 64 + bitMask = bit bitIndex :: Word64 + in IntMap.insertWith xor wordIndex bitMask mapValue + packedMap = foldl' updateWord IntMap.empty entries + in PackedBitMatrix + { packedRows = rowCount, + packedCols = columnCount, + packedWordsPerRow = wordsPerRowValue, + packedWords = U.generate packedWordCount (\wordIndex -> IntMap.findWithDefault zeroBits wordIndex packedMap) + } + +matrixRowWords :: PackedBitMatrix -> Int -> U.Vector Word64 +matrixRowWords matrixValue rowIndex = + let wordsPerRowValue = packedWordsPerRow matrixValue + startIndex = rowIndex * wordsPerRowValue + in U.take wordsPerRowValue (U.drop startIndex (packedWords matrixValue)) + +rankPackedRows :: Int -> [U.Vector Word64] -> Int +rankPackedRows = + GF2Xor.rankPackedRowsByReduction + +instance AdditiveMonoid GF2 where + zero = GF2Zero + add left right = + case (left, right) of + (GF2Zero, value) -> value + (value, GF2Zero) -> value + (GF2One, GF2One) -> GF2Zero + +instance AdditiveGroup GF2 where + neg = id + +instance MultiplicativeMonoid GF2 where + one = GF2One + mul left right = + case (left, right) of + (GF2One, GF2One) -> GF2One + _ -> GF2Zero + +instance Ring GF2 + +instance Field GF2 where + tryInv value = case value of + GF2Zero -> Nothing + GF2One -> Just GF2One + +instance Semiring GF2 + +instance CommutativeRing GF2 + +instance IntegralDomain GF2 where + isZero value = value == GF2Zero + unitInverse value = case value of + GF2Zero -> Nothing + GF2One -> Just GF2One + +instance GCDDomain GF2 where + gcdDomain left right + | left == GF2Zero && right == GF2Zero = GF2Zero + | otherwise = GF2One + extendedGcdDomain left right + | left == GF2Zero && right == GF2Zero = (GF2Zero, GF2Zero, GF2Zero) + | left == GF2One = (GF2One, GF2One, GF2Zero) + | otherwise = (GF2One, GF2Zero, GF2One) + +instance EuclideanDomain GF2 where + type Degree GF2 = Int + + divideWithRemainder numerator _ = (numerator, GF2Zero) + + degree value = + case value of + GF2Zero -> 0 + GF2One -> 0 diff --git a/src-carrier/Moonlight/LinAlg/Internal/GF2/SparseColumn.hs b/src-carrier/Moonlight/LinAlg/Internal/GF2/SparseColumn.hs new file mode 100644 --- /dev/null +++ b/src-carrier/Moonlight/LinAlg/Internal/GF2/SparseColumn.hs @@ -0,0 +1,298 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE StandaloneKindSignatures #-} + +module Moonlight.LinAlg.Internal.GF2.SparseColumn + ( GF2SparseColumn + , gf2SparseColumnIndex + , gf2SparseColumnRows + , mkGF2SparseColumn + , GF2SparseReducerConfig + , gf2SparseDensifyThreshold + , mkGF2SparseReducerConfig + , defaultGF2SparseReducerConfig + , GF2SparseColumnReduction (..) + , reduceGF2SparseColumns + , rankGF2SparseColumns + , independentGF2SparseColumns + , kernelBasisGF2SparseColumns + ) where + +import Control.Monad (foldM, unless) +import Data.Foldable (traverse_) +import Data.IntMap.Strict (IntMap) +import Data.IntMap.Strict qualified as IntMap +import Data.Kind (Type) +import Data.List (sortOn) +import Data.Vector (Vector) +import Data.Vector qualified as V +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Internal.GF2.Xor + ( PackedRow + , packedRowFromIndices + , packedRowIndices + , packedRowIsZero + , packedRowNonZeroCount + , packedRowXor + , unitPackedRow + ) + +-- | Sparse low-pivot GF2 column reduction with optional packed fallback. +type GF2SparseColumn :: Type +data GF2SparseColumn = GF2SparseColumn + { gf2SparseColumnIndex :: !Int + , gf2SparseColumnRows :: ![Int] + } + deriving stock (Eq, Show) + +type GF2SparseReducerConfig :: Type +data GF2SparseReducerConfig = GF2SparseReducerConfig + { gf2SparseDensifyThreshold :: !Int + } + deriving stock (Eq, Show) + +type GF2SparseColumnReduction :: Type +data GF2SparseColumnReduction = GF2SparseColumnReduction + { gf2SparseReductionRank :: !Int + , gf2SparseIndependentColumns :: !(Vector Int) + , gf2SparseKernelBasis :: !(Vector PackedRow) + } + deriving stock (Eq, Show) + +type SparseColumnBody :: Type +data SparseColumnBody + = SparseRows ![Int] + | PackedRows !PackedRow + deriving stock (Eq, Show) + +type TrackedSparseBasisColumn :: Type +data TrackedSparseBasisColumn = TrackedSparseBasisColumn + { tsbcData :: !SparseColumnBody + , tsbcWitness :: !PackedRow + } + deriving stock (Eq, Show) + +mkGF2SparseColumn :: String -> Int -> Int -> [Int] -> Either MoonlightError GF2SparseColumn +mkGF2SparseColumn context rowCount columnIndex rowsValue = do + unless (rowCount >= 0) + (Left (InvariantViolation (context <> ": negative sparse GF2 row count " <> show rowCount))) + unless (columnIndex >= 0) + (Left (InvariantViolation (context <> ": negative sparse GF2 column index " <> show columnIndex))) + traverse_ validateRow rowsValue + Right + GF2SparseColumn + { gf2SparseColumnIndex = columnIndex + , gf2SparseColumnRows = canonicalGF2Support rowsValue + } + where + validateRow rowIndex + | rowIndex < 0 || rowIndex >= rowCount = + Left + ( InvariantViolation + ( context + <> ": sparse GF2 row " + <> show rowIndex + <> " is outside row count " + <> show rowCount + ) + ) + | otherwise = Right () + +mkGF2SparseReducerConfig :: String -> Int -> Either MoonlightError GF2SparseReducerConfig +mkGF2SparseReducerConfig context thresholdValue + | thresholdValue < 0 = + Left (InvariantViolation (context <> ": negative sparse GF2 densify threshold " <> show thresholdValue)) + | otherwise = + Right GF2SparseReducerConfig {gf2SparseDensifyThreshold = thresholdValue} + +defaultGF2SparseReducerConfig :: GF2SparseReducerConfig +defaultGF2SparseReducerConfig = + GF2SparseReducerConfig {gf2SparseDensifyThreshold = 64} + +reduceGF2SparseColumns :: + GF2SparseReducerConfig -> + Int -> + Int -> + Vector GF2SparseColumn -> + Either MoonlightError GF2SparseColumnReduction +reduceGF2SparseColumns configValue rowCount columnCount columnsValue = do + unless (rowCount >= 0 && columnCount >= 0) + (Left (InvariantViolation ("reduceGF2SparseColumns: negative sparse GF2 shape " <> show (rowCount, columnCount)))) + orderedColumns <- validateColumnCover columnCount columnsValue + (_, independentReversed, kernelReversed) <- + foldM + (reduceColumn configValue rowCount columnCount) + (IntMap.empty, [], []) + (V.toList orderedColumns) + let independentColumns = V.fromList (reverse independentReversed) + Right + GF2SparseColumnReduction + { gf2SparseReductionRank = V.length independentColumns + , gf2SparseIndependentColumns = independentColumns + , gf2SparseKernelBasis = V.fromList (reverse kernelReversed) + } + +rankGF2SparseColumns :: + GF2SparseReducerConfig -> + Int -> + Int -> + Vector GF2SparseColumn -> + Either MoonlightError Int +rankGF2SparseColumns configValue rowCount columnCount columnsValue = + gf2SparseReductionRank <$> reduceGF2SparseColumns configValue rowCount columnCount columnsValue + +independentGF2SparseColumns :: + GF2SparseReducerConfig -> + Int -> + Int -> + Vector GF2SparseColumn -> + Either MoonlightError (Vector Int) +independentGF2SparseColumns configValue rowCount columnCount columnsValue = + gf2SparseIndependentColumns <$> reduceGF2SparseColumns configValue rowCount columnCount columnsValue + +kernelBasisGF2SparseColumns :: + GF2SparseReducerConfig -> + Int -> + Int -> + Vector GF2SparseColumn -> + Either MoonlightError (Vector PackedRow) +kernelBasisGF2SparseColumns configValue rowCount columnCount columnsValue = + gf2SparseKernelBasis <$> reduceGF2SparseColumns configValue rowCount columnCount columnsValue + +reduceColumn :: + GF2SparseReducerConfig -> + Int -> + Int -> + (IntMap TrackedSparseBasisColumn, [Int], [PackedRow]) -> + GF2SparseColumn -> + Either MoonlightError (IntMap TrackedSparseBasisColumn, [Int], [PackedRow]) +reduceColumn configValue rowCount columnCount (basisColumns, independentReversed, kernelReversed) columnValue = do + witnessValue <- unitPackedRow "reduceGF2SparseColumns: witness" columnCount (gf2SparseColumnIndex columnValue) + initialBody <- normalizeRows configValue rowCount (gf2SparseColumnRows columnValue) + (reducedData, reducedWitness) <- reduceSparseTracked configValue rowCount basisColumns initialBody witnessValue + case sparseBodyLowPivot reducedData of + Nothing -> Right (basisColumns, independentReversed, reducedWitness : kernelReversed) + Just pivotIndex -> + Right + ( IntMap.insert + pivotIndex + TrackedSparseBasisColumn + { tsbcData = reducedData + , tsbcWitness = reducedWitness + } + basisColumns + , gf2SparseColumnIndex columnValue : independentReversed + , kernelReversed + ) + +reduceSparseTracked :: + GF2SparseReducerConfig -> + Int -> + IntMap TrackedSparseBasisColumn -> + SparseColumnBody -> + PackedRow -> + Either MoonlightError (SparseColumnBody, PackedRow) +reduceSparseTracked configValue rowCount basisColumns dataValue witnessValue = + case sparseBodyLowPivot dataValue of + Nothing -> Right (dataValue, witnessValue) + Just pivotIndex -> + case IntMap.lookup pivotIndex basisColumns of + Nothing -> Right (dataValue, witnessValue) + Just TrackedSparseBasisColumn {tsbcData, tsbcWitness} -> do + reducedData <- sparseBodyXor configValue rowCount dataValue tsbcData + reducedWitness <- packedRowXor "reduceGF2SparseColumns: witness xor" witnessValue tsbcWitness + reduceSparseTracked configValue rowCount basisColumns reducedData reducedWitness + +sparseBodyXor :: + GF2SparseReducerConfig -> + Int -> + SparseColumnBody -> + SparseColumnBody -> + Either MoonlightError SparseColumnBody +sparseBodyXor configValue rowCount leftBody rightBody = + case (leftBody, rightBody) of + (PackedRows leftPacked, PackedRows rightPacked) -> + packedRowXor "reduceGF2SparseColumns: packed sparse body xor" leftPacked rightPacked + >>= normalizePacked configValue + _ -> + normalizeRows + configValue + rowCount + (xorSortedSupports (sparseBodyRows leftBody) (sparseBodyRows rightBody)) + +normalizeRows :: GF2SparseReducerConfig -> Int -> [Int] -> Either MoonlightError SparseColumnBody +normalizeRows configValue rowCount rowsValue + | supportPastThreshold configValue rowsValue = + PackedRows <$> packedRowFromIndices "reduceGF2SparseColumns: densified sparse body" rowCount rowsValue + | otherwise = Right (SparseRows rowsValue) + +normalizePacked :: GF2SparseReducerConfig -> PackedRow -> Either MoonlightError SparseColumnBody +normalizePacked configValue packedValue + | packedRowIsZero packedValue = Right (SparseRows []) + | packedRowNonZeroCount packedValue >= gf2SparseDensifyThreshold configValue = + Right (PackedRows packedValue) + | otherwise = Right (SparseRows (packedRowIndices packedValue)) + +supportPastThreshold :: GF2SparseReducerConfig -> [Int] -> Bool +supportPastThreshold configValue rowsValue = + not (null rowsValue) && length rowsValue >= gf2SparseDensifyThreshold configValue + +sparseBodyLowPivot :: SparseColumnBody -> Maybe Int +sparseBodyLowPivot bodyValue = + foldl' (\_ rowIndex -> Just rowIndex) Nothing (sparseBodyRows bodyValue) + +sparseBodyRows :: SparseColumnBody -> [Int] +sparseBodyRows bodyValue = + case bodyValue of + SparseRows rowsValue -> rowsValue + PackedRows packedValue -> packedRowIndices packedValue + +validateColumnCover :: Int -> Vector GF2SparseColumn -> Either MoonlightError (Vector GF2SparseColumn) +validateColumnCover columnCount columnsValue = do + unless (V.length columnsValue == columnCount) + ( Left + ( InvariantViolation + ( "reduceGF2SparseColumns: received " + <> show (V.length columnsValue) + <> " sparse GF2 columns for column count " + <> show columnCount + ) + ) + ) + traverse_ validateIndexedColumn (zip [0 .. columnCount - 1] orderedColumns) + Right (V.fromList orderedColumns) + where + orderedColumns = + sortOn gf2SparseColumnIndex (V.toList columnsValue) + + validateIndexedColumn (expectedIndex, columnValue) + | gf2SparseColumnIndex columnValue == expectedIndex = Right () + | otherwise = + Left + ( InvariantViolation + ( "reduceGF2SparseColumns: sparse GF2 column cover expected index " + <> show expectedIndex + <> " but found " + <> show (gf2SparseColumnIndex columnValue) + ) + ) + +canonicalGF2Support :: [Int] -> [Int] +canonicalGF2Support = + IntMap.keys . foldl' toggleRow IntMap.empty + where + toggleRow supportMap rowIndex = + case IntMap.lookup rowIndex supportMap of + Nothing -> IntMap.insert rowIndex () supportMap + Just () -> IntMap.delete rowIndex supportMap + +xorSortedSupports :: [Int] -> [Int] -> [Int] +xorSortedSupports leftRows rightRows = + case (leftRows, rightRows) of + ([], _) -> rightRows + (_, []) -> leftRows + (leftRow : remainingLeft, rightRow : remainingRight) -> + case compare leftRow rightRow of + LT -> leftRow : xorSortedSupports remainingLeft rightRows + EQ -> xorSortedSupports remainingLeft remainingRight + GT -> rightRow : xorSortedSupports leftRows remainingRight diff --git a/src-carrier/Moonlight/LinAlg/Internal/GF2/Xor.hs b/src-carrier/Moonlight/LinAlg/Internal/GF2/Xor.hs new file mode 100644 --- /dev/null +++ b/src-carrier/Moonlight/LinAlg/Internal/GF2/Xor.hs @@ -0,0 +1,820 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE StandaloneKindSignatures #-} + +module Moonlight.LinAlg.Internal.GF2.Xor + ( PackedRow + , packedRowWidth + , packedRowNonZeroCount + , emptyPackedRow + , unitPackedRow + , packedRowFromIndices + , packedRowIndices + , packedRowMember + , packedRowIsZero + , packedRowXor + , packedRowRemap + , PackedLinearMap + , packedLinearMapDomain + , packedLinearMapCodomain + , packedLinearMapColumns + , packedLinearMapFromColumns + , packedLinearMapFromEntries + , zeroPackedLinearMap + , identityPackedLinearMap + , applyPackedLinearMap + , composePackedLinearMaps + , addPackedLinearMaps + , packedLinearMapIsZero + , PackedSpan + , emptyPackedSpan + , packedSpanFromRows + , reducePackedRow + , admitPackedRow + , ColumnReduction (..) + , reducePackedColumns + , PackedCoordinateSolver + , packedCoordinateSolver + , coordinatesInPackedBasis + , inverseFromPackedBasisColumns + , rankPackedRowsByReduction + ) where + +import Control.Monad (foldM, unless) +import Control.Monad.ST (ST, runST) +import Data.Bits + ( bit + , clearBit + , complement + , countTrailingZeros + , popCount + , testBit + , xor + , (.&.) + ) +import Data.IntMap.Strict (IntMap) +import Data.IntMap.Strict qualified as IntMap +import Data.Kind (Type) +import Data.Vector (Vector) +import Data.Vector qualified as V +import Data.Vector.Unboxed qualified as U +import Data.Vector.Unboxed.Mutable qualified as UM +import Data.Word (Word64) +import Moonlight.Core (MoonlightError (..)) + +wordBits :: Int +wordBits = 64 + +wordCountForWidth :: Int -> Int +wordCountForWidth widthValue + | widthValue <= 0 = 0 + | otherwise = (widthValue + wordBits - 1) `div` wordBits + +lastWordMask :: Int -> Word64 +lastWordMask widthValue + | widthValue <= 0 = 0 + | remainderValue == 0 = complement 0 + | otherwise = bit remainderValue - 1 + where + remainderValue = widthValue `mod` wordBits + +type PackedRow :: Type +data PackedRow = PackedRow + { prWidth :: !Int + , prWords :: !(U.Vector Word64) + , prNonZeroCount :: !Int + } + deriving stock (Eq, Show) + +packedRowWidth :: PackedRow -> Int +packedRowWidth = prWidth + +packedRowNonZeroCount :: PackedRow -> Int +packedRowNonZeroCount = prNonZeroCount + +packedRowFromWords :: Int -> U.Vector Word64 -> PackedRow +packedRowFromWords widthValue rawWords = + let expectedWords = wordCountForWidth widthValue + paddedWords = + U.generate + expectedWords + (\wordIndex -> maybe 0 id (rawWords U.!? wordIndex)) + maskedWords + | U.null paddedWords = U.empty + | otherwise = + U.imap + (\wordIndex wordValue -> + if wordIndex == U.length paddedWords - 1 + then wordValue .&. lastWordMask widthValue + else wordValue + ) + paddedWords + in PackedRow + { prWidth = widthValue + , prWords = maskedWords + , prNonZeroCount = U.foldl' (\countValue wordValue -> countValue + popCount wordValue) 0 maskedWords + } + +emptyPackedRow :: Int -> Either MoonlightError PackedRow +emptyPackedRow widthValue + | widthValue < 0 = + Left (InvariantViolation ("emptyPackedRow: negative width " <> show widthValue)) + | otherwise = + Right (packedRowFromWords widthValue U.empty) + +unitPackedRow :: String -> Int -> Int -> Either MoonlightError PackedRow +unitPackedRow context widthValue indexValue = + packedRowFromIndices context widthValue [indexValue] + +packedRowFromIndices :: String -> Int -> [Int] -> Either MoonlightError PackedRow +packedRowFromIndices context widthValue indicesValue = do + unless (widthValue >= 0) + (Left (InvariantViolation (context <> ": negative packed-row width " <> show widthValue))) + traverse_ validateIndex indicesValue + let wordValues = + U.create $ do + mutableWords <- UM.replicate (wordCountForWidth widthValue) 0 + traverse_ (toggleIndex mutableWords) indicesValue + pure mutableWords + Right (packedRowFromWords widthValue wordValues) + where + validateIndex indexValue + | indexValue < 0 || indexValue >= widthValue = + Left + ( InvariantViolation + ( context + <> ": packed coordinate " + <> show indexValue + <> " is outside width " + <> show widthValue + ) + ) + | otherwise = Right () + + toggleIndex :: UM.MVector state Word64 -> Int -> ST state () + toggleIndex mutableWords indexValue = do + let wordIndex = indexValue `div` wordBits + bitIndex = indexValue `mod` wordBits + oldWord <- UM.read mutableWords wordIndex + UM.write mutableWords wordIndex (oldWord `xor` bit bitIndex) + +packedRowIndices :: PackedRow -> [Int] +packedRowIndices PackedRow {prWidth, prWords} = + reverse (U.ifoldl' collectWord [] prWords) + where + collectWord accumulated wordIndex wordValue = + collectBits accumulated (wordIndex * wordBits) wordValue + + collectBits accumulated baseIndex remainingWord + | remainingWord == 0 = accumulated + | otherwise = + let bitIndex = countTrailingZeros remainingWord + coordinateValue = baseIndex + bitIndex + nextWord = clearBit remainingWord bitIndex + in collectBits + ( if coordinateValue < prWidth + then coordinateValue : accumulated + else accumulated + ) + baseIndex + nextWord + +packedRowMember :: Int -> PackedRow -> Bool +packedRowMember indexValue PackedRow {prWidth, prWords} + | indexValue < 0 || indexValue >= prWidth = False + | otherwise = + maybe + False + (`testBit` (indexValue `mod` wordBits)) + (prWords U.!? (indexValue `div` wordBits)) + +packedRowIsZero :: PackedRow -> Bool +packedRowIsZero = (== 0) . prNonZeroCount + +xorSameWidth :: PackedRow -> PackedRow -> PackedRow +xorSameWidth leftRow rightRow = + packedRowFromWords + (prWidth leftRow) + (U.zipWith xor (prWords leftRow) (prWords rightRow)) + +packedRowXor :: String -> PackedRow -> PackedRow -> Either MoonlightError PackedRow +packedRowXor context leftRow rightRow + | prWidth leftRow /= prWidth rightRow = + Left + ( InvariantViolation + ( context + <> ": packed-row width mismatch " + <> show (prWidth leftRow, prWidth rightRow) + ) + ) + | otherwise = Right (xorSameWidth leftRow rightRow) + +packedRowRemap :: + String -> + Int -> + (Int -> Maybe Int) -> + PackedRow -> + Either MoonlightError PackedRow +packedRowRemap context targetWidth remapIndex sourceRow = do + remappedIndices <- traverse remapOne (packedRowIndices sourceRow) + packedRowFromIndices context targetWidth remappedIndices + where + remapOne sourceIndex = + case remapIndex sourceIndex of + Nothing -> + Left + ( InvariantViolation + ( context + <> ": no target coordinate for source coordinate " + <> show sourceIndex + ) + ) + Just targetIndex -> Right targetIndex + +type PackedLinearMap :: Type +data PackedLinearMap = PackedLinearMap + { plmDomain :: !Int + , plmCodomain :: !Int + , plmColumns :: !(Vector PackedRow) + } + deriving stock (Eq, Show) + +packedLinearMapDomain :: PackedLinearMap -> Int +packedLinearMapDomain = plmDomain + +packedLinearMapCodomain :: PackedLinearMap -> Int +packedLinearMapCodomain = plmCodomain + +packedLinearMapColumns :: PackedLinearMap -> Vector PackedRow +packedLinearMapColumns = plmColumns + +packedLinearMapFromColumns :: + String -> + Int -> + Int -> + Vector PackedRow -> + Either MoonlightError PackedLinearMap +packedLinearMapFromColumns context domainValue codomainValue columnValues + | domainValue < 0 || codomainValue < 0 = + Left + ( InvariantViolation + ( context + <> ": negative linear-map shape " + <> show (codomainValue, domainValue) + ) + ) + | V.length columnValues /= domainValue = + Left + ( InvariantViolation + ( context + <> ": received " + <> show (V.length columnValues) + <> " columns for domain dimension " + <> show domainValue + ) + ) + | otherwise = do + traverse_ validateColumn (V.toList (V.indexed columnValues)) + Right + PackedLinearMap + { plmDomain = domainValue + , plmCodomain = codomainValue + , plmColumns = columnValues + } + where + validateColumn (columnIndex, columnValue) + | prWidth columnValue == codomainValue = Right () + | otherwise = + Left + ( InvariantViolation + ( context + <> ": column " + <> show columnIndex + <> " has width " + <> show (prWidth columnValue) + <> ", expected " + <> show codomainValue + ) + ) + +packedLinearMapFromEntries :: + String -> + Int -> + Int -> + [(Int, Int)] -> + Either MoonlightError PackedLinearMap +packedLinearMapFromEntries context domainValue codomainValue entriesValue = do + unless (domainValue >= 0 && codomainValue >= 0) + (Left (InvariantViolation (context <> ": negative linear-map shape " <> show (codomainValue, domainValue)))) + traverse_ validateEntry entriesValue + columnsValue <- + traverse + (\columnIndex -> + packedRowFromIndices + (context <> ": column " <> show columnIndex) + codomainValue + (IntMap.findWithDefault [] columnIndex entriesByColumn) + ) + [0 .. domainValue - 1] + packedLinearMapFromColumns context domainValue codomainValue (V.fromList columnsValue) + where + entriesByColumn = + foldl' + (\entryMap (rowIndex, columnIndex) -> + IntMap.insertWith (flip (<>)) columnIndex [rowIndex] entryMap + ) + IntMap.empty + entriesValue + + validateEntry (rowIndex, columnIndex) + | rowIndex < 0 || rowIndex >= codomainValue = + Left + ( InvariantViolation + ( context + <> ": row index " + <> show rowIndex + <> " is outside codomain dimension " + <> show codomainValue + ) + ) + | columnIndex < 0 || columnIndex >= domainValue = + Left + ( InvariantViolation + ( context + <> ": column index " + <> show columnIndex + <> " is outside domain dimension " + <> show domainValue + ) + ) + | otherwise = Right () + +zeroPackedLinearMap :: String -> Int -> Int -> Either MoonlightError PackedLinearMap +zeroPackedLinearMap context domainValue codomainValue + | domainValue < 0 || codomainValue < 0 = + Left + ( InvariantViolation + ( context + <> ": negative linear-map shape " + <> show (codomainValue, domainValue) + ) + ) + | otherwise = do + zeroColumn <- emptyPackedRow codomainValue + packedLinearMapFromColumns context domainValue codomainValue (V.replicate domainValue zeroColumn) + +identityPackedLinearMap :: String -> Int -> Either MoonlightError PackedLinearMap +identityPackedLinearMap context dimensionValue = do + columnsValue <- traverse (unitPackedRow context dimensionValue) [0 .. dimensionValue - 1] + packedLinearMapFromColumns context dimensionValue dimensionValue (V.fromList columnsValue) + +xorIntoMutable :: UM.MVector state Word64 -> U.Vector Word64 -> ST state () +xorIntoMutable mutableTarget sourceWords = + traverse_ xorWordAt [0 .. UM.length mutableTarget - 1] + where + xorWordAt wordIndex = do + oldWord <- UM.read mutableTarget wordIndex + let sourceWord = maybe 0 id (sourceWords U.!? wordIndex) + UM.write mutableTarget wordIndex (oldWord `xor` sourceWord) + +applyPackedLinearMap :: String -> PackedLinearMap -> PackedRow -> Either MoonlightError PackedRow +applyPackedLinearMap context linearMap sourceRow + | prWidth sourceRow /= plmDomain linearMap = + Left + ( InvariantViolation + ( context + <> ": vector width " + <> show (prWidth sourceRow) + <> " does not match map domain " + <> show (plmDomain linearMap) + ) + ) + | otherwise = do + selectedColumns <- + traverse + (\columnIndex -> lookupVector (context <> ": source column") columnIndex (plmColumns linearMap)) + (packedRowIndices sourceRow) + let resultWords = + runST $ do + mutableResult <- UM.replicate (wordCountForWidth (plmCodomain linearMap)) 0 + traverse_ (xorIntoMutable mutableResult . prWords) selectedColumns + U.freeze mutableResult + Right (packedRowFromWords (plmCodomain linearMap) resultWords) + +composePackedLinearMaps :: + String -> + PackedLinearMap -> + PackedLinearMap -> + Either MoonlightError PackedLinearMap +composePackedLinearMaps context leftMap rightMap + | plmDomain leftMap /= plmCodomain rightMap = + Left + ( InvariantViolation + ( context + <> ": incompatible map shapes " + <> show (plmCodomain leftMap, plmDomain leftMap) + <> " and " + <> show (plmCodomain rightMap, plmDomain rightMap) + ) + ) + | otherwise = do + productColumns <- + traverse + (applyPackedLinearMap (context <> ": product column") leftMap) + (plmColumns rightMap) + packedLinearMapFromColumns context (plmDomain rightMap) (plmCodomain leftMap) productColumns + +addPackedLinearMaps :: + String -> + PackedLinearMap -> + PackedLinearMap -> + Either MoonlightError PackedLinearMap +addPackedLinearMaps context leftMap rightMap + | mapShape leftMap /= mapShape rightMap = + Left + ( InvariantViolation + ( context + <> ": linear-map shape mismatch " + <> show (mapShape leftMap, mapShape rightMap) + ) + ) + | otherwise = + packedLinearMapFromColumns + context + (plmDomain leftMap) + (plmCodomain leftMap) + (V.zipWith xorSameWidth (plmColumns leftMap) (plmColumns rightMap)) + +packedLinearMapIsZero :: PackedLinearMap -> Bool +packedLinearMapIsZero = V.all packedRowIsZero . plmColumns + +mapShape :: PackedLinearMap -> (Int, Int) +mapShape mapValue = (plmCodomain mapValue, plmDomain mapValue) + +type PackedSpan :: Type +data PackedSpan = PackedSpan + { psWidth :: !Int + , psBasis :: !(IntMap PackedRow) + } + deriving stock (Eq, Show) + +emptyPackedSpan :: Int -> Either MoonlightError PackedSpan +emptyPackedSpan widthValue + | widthValue < 0 = Left (InvariantViolation ("emptyPackedSpan: negative width " <> show widthValue)) + | otherwise = Right PackedSpan {psWidth = widthValue, psBasis = IntMap.empty} + +packedSpanFromRows :: String -> Int -> [PackedRow] -> Either MoonlightError PackedSpan +packedSpanFromRows context widthValue rowsValue = do + initialSpan <- emptyPackedSpan widthValue + foldM (\spanValue rowValue -> snd <$> admitPackedRow context rowValue spanValue) initialSpan rowsValue + +mutablePivot :: Int -> UM.MVector state Word64 -> ST state (Maybe Int) +mutablePivot widthValue mutableWords = do + candidates <- traverse pivotAtWord [0 .. UM.length mutableWords - 1] + pure (foldr firstJust Nothing candidates) + where + pivotAtWord wordIndex = do + wordValue <- UM.read mutableWords wordIndex + pure + ( if wordValue == 0 + then Nothing + else + let pivotIndex = wordIndex * wordBits + countTrailingZeros wordValue + in if pivotIndex < widthValue then Just pivotIndex else Nothing + ) + + firstJust :: Maybe Int -> Maybe Int -> Maybe Int + firstJust left right = + case left of + Nothing -> right + Just _ -> left + +reduceMutableWords :: Int -> IntMap PackedRow -> UM.MVector state Word64 -> ST state () +reduceMutableWords widthValue basisRows mutableWords = do + maybePivot <- mutablePivot widthValue mutableWords + case maybePivot of + Nothing -> pure () + Just pivotIndex -> + case IntMap.lookup pivotIndex basisRows of + Nothing -> pure () + Just basisRow -> do + xorIntoMutable mutableWords (prWords basisRow) + reduceMutableWords widthValue basisRows mutableWords + +reducePackedRowUnchecked :: PackedSpan -> PackedRow -> PackedRow +reducePackedRowUnchecked PackedSpan {psWidth, psBasis} rowValue = + let reducedWords = runST $ do + mutableWords <- U.thaw (prWords rowValue) + reduceMutableWords psWidth psBasis mutableWords + U.freeze mutableWords + in packedRowFromWords psWidth reducedWords + +reducePackedRow :: String -> PackedSpan -> PackedRow -> Either MoonlightError PackedRow +reducePackedRow context spanValue@PackedSpan {psWidth} rowValue + | prWidth rowValue /= psWidth = + Left + ( InvariantViolation + ( context + <> ": row width " + <> show (prWidth rowValue) + <> " does not match span width " + <> show psWidth + ) + ) + | otherwise = Right (reducePackedRowUnchecked spanValue rowValue) + +packedRowPivot :: PackedRow -> Maybe Int +packedRowPivot PackedRow {prWidth, prWords} = + U.ifoldl' firstPivot Nothing prWords + where + firstPivot (Just pivotIndex) _ _ = + Just pivotIndex + firstPivot Nothing wordIndex wordValue + | wordValue == 0 = Nothing + | otherwise = + let pivotIndex = wordIndex * wordBits + countTrailingZeros wordValue + in if pivotIndex < prWidth then Just pivotIndex else Nothing + +admitPackedRow :: + String -> + PackedRow -> + PackedSpan -> + Either MoonlightError (Maybe PackedRow, PackedSpan) +admitPackedRow context candidateRow spanValue@PackedSpan {psWidth, psBasis} = do + reducedRow <- reducePackedRow context spanValue candidateRow + case packedRowPivot reducedRow of + Nothing -> Right (Nothing, spanValue) + Just pivotIndex -> + Right + ( Just reducedRow + , PackedSpan + { psWidth = psWidth + , psBasis = IntMap.insert pivotIndex reducedRow psBasis + } + ) + +type TrackedBasisRow :: Type +data TrackedBasisRow = TrackedBasisRow + { tbrData :: !PackedRow + , tbrWitness :: !PackedRow + } + deriving stock (Eq, Show) + +type ColumnReduction :: Type +data ColumnReduction = ColumnReduction + { crIndependentIndices :: !(Vector Int) + , crKernelBasis :: !(Vector PackedRow) + } + deriving stock (Eq, Show) + +reduceMutableTracked :: + Int -> + IntMap TrackedBasisRow -> + UM.MVector state Word64 -> + UM.MVector state Word64 -> + ST state () +reduceMutableTracked dataWidth basisRows mutableData mutableWitness = do + maybePivot <- mutablePivot dataWidth mutableData + case maybePivot of + Nothing -> pure () + Just pivotIndex -> + case IntMap.lookup pivotIndex basisRows of + Nothing -> pure () + Just TrackedBasisRow {tbrData, tbrWitness} -> do + xorIntoMutable mutableData (prWords tbrData) + xorIntoMutable mutableWitness (prWords tbrWitness) + reduceMutableTracked dataWidth basisRows mutableData mutableWitness + +reduceTrackedRows :: + String -> + IntMap TrackedBasisRow -> + PackedRow -> + PackedRow -> + Either MoonlightError (PackedRow, PackedRow) +reduceTrackedRows context basisRows dataRow witnessRow = do + traverse_ validateBasis (IntMap.toList basisRows) + let (dataWords, witnessWords) = + runST $ do + mutableData <- U.thaw (prWords dataRow) + mutableWitness <- U.thaw (prWords witnessRow) + reduceMutableTracked (prWidth dataRow) basisRows mutableData mutableWitness + frozenData <- U.freeze mutableData + frozenWitness <- U.freeze mutableWitness + pure (frozenData, frozenWitness) + Right + ( packedRowFromWords (prWidth dataRow) dataWords + , packedRowFromWords (prWidth witnessRow) witnessWords + ) + where + validateBasis (pivotIndex, TrackedBasisRow {tbrData, tbrWitness}) + | pivotIndex < 0 || pivotIndex >= prWidth dataRow = + Left (InvariantViolation (context <> ": tracked pivot outside data width: " <> show pivotIndex)) + | prWidth tbrData /= prWidth dataRow = + Left (InvariantViolation (context <> ": tracked data width mismatch")) + | prWidth tbrWitness /= prWidth witnessRow = + Left (InvariantViolation (context <> ": tracked witness width mismatch")) + | otherwise = Right () + +reducePackedColumns :: + String -> + Int -> + Vector PackedRow -> + Either MoonlightError ColumnReduction +reducePackedColumns context codomainWidth columnsValue = do + unless (codomainWidth >= 0) + (Left (InvariantViolation (context <> ": negative codomain width " <> show codomainWidth))) + traverse_ validateColumn (V.toList (V.indexed columnsValue)) + let domainWidth = V.length columnsValue + (_, independentReversed, kernelReversed) <- + foldM + (reduceColumn domainWidth) + (IntMap.empty, [], []) + (V.toList (V.indexed columnsValue)) + Right + ColumnReduction + { crIndependentIndices = V.fromList (reverse independentReversed) + , crKernelBasis = V.fromList (reverse kernelReversed) + } + where + validateColumn (columnIndex, columnValue) + | prWidth columnValue == codomainWidth = Right () + | otherwise = + Left + ( InvariantViolation + ( context + <> ": column " + <> show columnIndex + <> " has width " + <> show (prWidth columnValue) + <> ", expected " + <> show codomainWidth + ) + ) + + reduceColumn domainWidth (basisRows, independentReversed, kernelReversed) (columnIndex, columnValue) = do + witnessValue <- unitPackedRow (context <> ": witness") domainWidth columnIndex + (reducedData, reducedWitness) <- reduceTrackedRows context basisRows columnValue witnessValue + case packedRowPivot reducedData of + Nothing -> Right (basisRows, independentReversed, reducedWitness : kernelReversed) + Just pivotIndex -> + Right + ( IntMap.insert + pivotIndex + TrackedBasisRow {tbrData = reducedData, tbrWitness = reducedWitness} + basisRows + , columnIndex : independentReversed + , kernelReversed + ) + +type PackedCoordinateSolver :: Type +data PackedCoordinateSolver = PackedCoordinateSolver + { pcsAmbientWidth :: !Int + , pcsBasisCardinality :: !Int + , pcsBasisRows :: !(IntMap TrackedBasisRow) + } + deriving stock (Eq, Show) + +packedCoordinateSolver :: + String -> + Int -> + Vector PackedRow -> + Either MoonlightError PackedCoordinateSolver +packedCoordinateSolver context ambientWidth basisColumns = do + let basisCardinality = V.length basisColumns + unless (ambientWidth >= 0) + (Left (InvariantViolation (context <> ": negative ambient width " <> show ambientWidth))) + unless (basisCardinality <= ambientWidth) + ( Left + ( InvariantViolation + ( context + <> ": basis cardinality " + <> show basisCardinality + <> " exceeds ambient width " + <> show ambientWidth + ) + ) + ) + basisRows <- + foldM insertBasisColumn IntMap.empty (V.toList (V.indexed basisColumns)) + Right + PackedCoordinateSolver + { pcsAmbientWidth = ambientWidth + , pcsBasisCardinality = basisCardinality + , pcsBasisRows = basisRows + } + where + insertBasisColumn basisRows (basisIndex, columnValue) + | prWidth columnValue /= ambientWidth = + Left + ( InvariantViolation + ( context + <> ": basis column " + <> show basisIndex + <> " has width " + <> show (prWidth columnValue) + <> ", expected " + <> show ambientWidth + ) + ) + | otherwise = do + witnessValue <- unitPackedRow (context <> ": basis witness") (V.length basisColumns) basisIndex + (reducedData, reducedWitness) <- + reduceTrackedRows (context <> ": basis reduction") basisRows columnValue witnessValue + case packedRowPivot reducedData of + Nothing -> + Left + ( InvariantViolation + ( context + <> ": supplied basis columns are linearly dependent at column " + <> show basisIndex + ) + ) + Just pivotIndex -> + Right + ( IntMap.insert + pivotIndex + TrackedBasisRow {tbrData = reducedData, tbrWitness = reducedWitness} + basisRows + ) + +coordinatesInPackedBasis :: + String -> + PackedCoordinateSolver -> + PackedRow -> + Either MoonlightError (Maybe PackedRow) +coordinatesInPackedBasis context PackedCoordinateSolver {pcsAmbientWidth, pcsBasisCardinality, pcsBasisRows} vectorValue + | prWidth vectorValue /= pcsAmbientWidth = + Left + ( InvariantViolation + ( context + <> ": vector width " + <> show (prWidth vectorValue) + <> " does not match ambient width " + <> show pcsAmbientWidth + ) + ) + | otherwise = do + zeroWitness <- emptyPackedRow pcsBasisCardinality + (reducedData, reducedWitness) <- + reduceTrackedRows context pcsBasisRows vectorValue zeroWitness + Right (if packedRowIsZero reducedData then Just reducedWitness else Nothing) + +inverseFromPackedBasisColumns :: + String -> + Vector PackedRow -> + Either MoonlightError PackedLinearMap +inverseFromPackedBasisColumns context basisColumns = do + let dimensionValue = V.length basisColumns + solver <- packedCoordinateSolver (context <> ": coordinate solver") dimensionValue basisColumns + inverseColumns <- traverse (inverseColumn solver) [0 .. dimensionValue - 1] + packedLinearMapFromColumns context dimensionValue dimensionValue (V.fromList inverseColumns) + where + inverseColumn solver columnIndex = do + unitVector <- unitPackedRow (context <> ": inverse unit vector") (V.length basisColumns) columnIndex + maybeCoordinates <- coordinatesInPackedBasis (context <> ": inverse coordinates") solver unitVector + case maybeCoordinates of + Nothing -> Left (InvariantViolation (context <> ": basis columns do not span the ambient space")) + Just coordinatesValue -> Right coordinatesValue + +rankPackedRowsByReduction :: Int -> [U.Vector Word64] -> Int +rankPackedRowsByReduction widthValue rowWords + | widthValue <= 0 = 0 + | otherwise = + IntMap.size + ( psBasis + ( foldl' + admitUnchecked + PackedSpan {psWidth = widthValue, psBasis = IntMap.empty} + (packedRowFromWords widthValue <$> rowWords) + ) + ) + where + admitUnchecked spanValue@PackedSpan {psWidth, psBasis} rowValue = + let reducedRow = reducePackedRowUnchecked spanValue rowValue + in case packedRowPivot reducedRow of + Nothing -> spanValue + Just pivotIndex -> + PackedSpan + { psWidth = psWidth + , psBasis = IntMap.insert pivotIndex reducedRow psBasis + } + +lookupVector :: String -> Int -> Vector value -> Either MoonlightError value +lookupVector context indexValue vectorValue = + case vectorValue V.!? indexValue of + Nothing -> + Left + ( InvariantViolation + ( context + <> ": index " + <> show indexValue + <> " is outside vector length " + <> show (V.length vectorValue) + ) + ) + Just value -> Right value + +traverse_ :: Applicative f => (a -> f b) -> [a] -> f () +traverse_ actionValue = + foldr (\value accumulated -> actionValue value *> accumulated) (pure ()) diff --git a/src-carrier/Moonlight/LinAlg/Internal/Primitives.hs b/src-carrier/Moonlight/LinAlg/Internal/Primitives.hs new file mode 100644 --- /dev/null +++ b/src-carrier/Moonlight/LinAlg/Internal/Primitives.hs @@ -0,0 +1,261 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} + +module Moonlight.LinAlg.Internal.Primitives + ( MatrixIndex, + RowIndex, + ColumnIndex, + mkRowIndex, + mkColumnIndex, + rowIndexInt, + columnIndexInt, + rowIndices, + columnIndices, + natInt, + epsilon, + selectAt, + selectAtIndex, + requireAt, + requireAtIndex, + requireRow, + requireColumnEntry, + requireMatrixEntry, + requireMatrixEntryAt, + updateAt, + replaceAt, + replaceAtIndexChecked, + replaceAtChecked, + replaceRowChecked, + replaceColumnEntryChecked, + swapAtIndexChecked, + swapAtChecked, + swapRowsChecked, + swapColumnsChecked, + dotProduct, + vectorNorm, + scaleVector, + addVector, + subVector, + matrixVectorProduct, + matrixSubtract, + scaleMatrix, + outerProduct, + basisVector, + linearCombination, + ) +where + +import Control.Monad (foldM) +import Data.Bifunctor (first) +import Data.Kind (Type) +import Data.Proxy (Proxy (..)) +import GHC.TypeNats (KnownNat, natVal) +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Internal.DenseList (matrixVectorProductWith, outerProductWith, scaleMatrixWith) +import Prelude + +type RowAxis :: Type +data RowAxis +type ColumnAxis :: Type +data ColumnAxis + +type MatrixIndex :: Type -> Type +newtype MatrixIndex axis = MatrixIndex Int + deriving stock (Eq, Ord) + +instance Show (MatrixIndex axis) where + show = show . matrixIndexInt + +type RowIndex :: Type +type RowIndex = MatrixIndex RowAxis +type ColumnIndex :: Type +type ColumnIndex = MatrixIndex ColumnAxis + +natInt :: forall n. KnownNat n => Int +natInt = fromIntegral (natVal (Proxy @n)) + +epsilon :: Double +epsilon = 1.0e-12 + +mkIndex :: MoonlightError -> Int -> Int -> Either MoonlightError (MatrixIndex axis) +mkIndex indexError upperBound candidateIndex + | candidateIndex < 0 = Left indexError + | candidateIndex >= upperBound = Left indexError + | otherwise = Right (MatrixIndex candidateIndex) + +mkRowIndex :: MoonlightError -> Int -> Int -> Either MoonlightError RowIndex +mkRowIndex = mkIndex + +mkColumnIndex :: MoonlightError -> Int -> Int -> Either MoonlightError ColumnIndex +mkColumnIndex = mkIndex + +matrixIndexInt :: MatrixIndex axis -> Int +matrixIndexInt (MatrixIndex indexValue) = indexValue + +rowIndexInt :: RowIndex -> Int +rowIndexInt = matrixIndexInt + +columnIndexInt :: ColumnIndex -> Int +columnIndexInt = matrixIndexInt + +rowIndices :: Int -> [RowIndex] +rowIndices rowCount = map MatrixIndex [0 .. rowCount - 1] + +columnIndices :: Int -> [ColumnIndex] +columnIndices columnCount = map MatrixIndex [0 .. columnCount - 1] + +selectAt :: Int -> [a] -> Maybe a +selectAt targetIndex values + | targetIndex < 0 = Nothing + | otherwise = + case drop targetIndex values of + value : _ -> Just value + [] -> Nothing + +selectAtIndex :: MatrixIndex axis -> [a] -> Maybe a +selectAtIndex targetIndex = selectAt (matrixIndexInt targetIndex) + +requireAt :: MoonlightError -> Int -> [a] -> Either MoonlightError a +requireAt lookupError targetIndex values = + maybe + (Left lookupError) + Right + (selectAt targetIndex values) + +requireAtIndex :: MoonlightError -> MatrixIndex axis -> [a] -> Either MoonlightError a +requireAtIndex lookupError targetIndex values = + maybe + (Left lookupError) + Right + (selectAtIndex targetIndex values) + +requireRow :: MoonlightError -> RowIndex -> [[a]] -> Either MoonlightError [a] +requireRow = requireAtIndex + +requireColumnEntry :: MoonlightError -> ColumnIndex -> [a] -> Either MoonlightError a +requireColumnEntry = requireAtIndex + +requireMatrixEntry :: MoonlightError -> Int -> Int -> [[a]] -> Either MoonlightError a +requireMatrixEntry lookupError rowIndex columnIndex matrixRows = + requireAt lookupError rowIndex matrixRows + >>= requireAt lookupError columnIndex + +requireMatrixEntryAt :: MoonlightError -> RowIndex -> ColumnIndex -> [[a]] -> Either MoonlightError a +requireMatrixEntryAt lookupError rowIndex columnIndex matrixRows = + requireRow lookupError rowIndex matrixRows + >>= requireColumnEntry lookupError columnIndex + +updateAt :: Int -> (a -> a) -> [a] -> [a] +updateAt targetIndex fn = + map + (\(indexValue, value) -> if indexValue == targetIndex then fn value else value) + . zip [0 :: Int ..] + +replaceAt :: Int -> a -> [a] -> [a] +replaceAt targetIndex replacement = updateAt targetIndex (const replacement) + +replaceAtIndexChecked :: MoonlightError -> MatrixIndex axis -> a -> [a] -> Either MoonlightError [a] +replaceAtIndexChecked updateError targetIndex replacement values = + requireAtIndex updateError targetIndex values + >>= const (Right (replaceAt (matrixIndexInt targetIndex) replacement values)) + +replaceAtChecked :: MoonlightError -> Int -> a -> [a] -> Either MoonlightError [a] +replaceAtChecked updateError targetIndex replacement values = + requireAt updateError targetIndex values + >>= const (Right (replaceAt targetIndex replacement values)) + +replaceRowChecked :: MoonlightError -> RowIndex -> [a] -> [[a]] -> Either MoonlightError [[a]] +replaceRowChecked = replaceAtIndexChecked + +replaceColumnEntryChecked :: MoonlightError -> ColumnIndex -> a -> [a] -> Either MoonlightError [a] +replaceColumnEntryChecked = replaceAtIndexChecked + +swapAtIndexChecked :: MoonlightError -> MatrixIndex axis -> MatrixIndex axis -> [a] -> Either MoonlightError [a] +swapAtIndexChecked updateError leftIndex rightIndex values = do + leftValue <- requireAtIndex updateError leftIndex values + rightValue <- requireAtIndex updateError rightIndex values + replaceAtIndexChecked updateError leftIndex rightValue values + >>= replaceAtIndexChecked updateError rightIndex leftValue + +swapAtChecked :: MoonlightError -> Int -> Int -> [a] -> Either MoonlightError [a] +swapAtChecked updateError leftIndex rightIndex values = do + leftValue <- requireAt updateError leftIndex values + rightValue <- requireAt updateError rightIndex values + replaceAtChecked updateError leftIndex rightValue values + >>= replaceAtChecked updateError rightIndex leftValue + +swapRowsChecked :: MoonlightError -> RowIndex -> RowIndex -> [[a]] -> Either MoonlightError [[a]] +swapRowsChecked = swapAtIndexChecked + +swapColumnsChecked :: MoonlightError -> ColumnIndex -> ColumnIndex -> [a] -> Either MoonlightError [a] +swapColumnsChecked = swapAtIndexChecked + +dotProduct :: [Double] -> [Double] -> Either MoonlightError Double +dotProduct left right = + go 0.0 left right + where + go !accumulator leftValues rightValues = + case (leftValues, rightValues) of + ([], []) -> Right accumulator + (leftValue : leftRest, rightValue : rightRest) -> + go + (accumulator + leftValue * rightValue) + leftRest + rightRest + _ -> + Left + ( InvariantViolation + ( "dotProduct: length mismatch (left=" + <> show (length left) + <> ", right=" + <> show (length right) + <> ")" + ) + ) +{-# INLINE dotProduct #-} + +vectorNorm :: [Double] -> Either MoonlightError Double +vectorNorm v = fmap sqrt (dotProduct v v) + +scaleVector :: Double -> [Double] -> [Double] +scaleVector scalarValue = map (\value -> scalarValue * value) + +addVector :: [Double] -> [Double] -> Either MoonlightError [Double] +addVector left right + | length left /= length right = + Left (InvariantViolation ("addVector: length mismatch (left=" <> show (length left) <> ", right=" <> show (length right) <> ")")) + | otherwise = Right (zipWith (+) left right) + +subVector :: [Double] -> [Double] -> Either MoonlightError [Double] +subVector left right + | length left /= length right = + Left (InvariantViolation ("subVector: length mismatch (left=" <> show (length left) <> ", right=" <> show (length right) <> ")")) + | otherwise = Right (zipWith (-) left right) + +matrixVectorProduct :: [[Double]] -> [Double] -> Either MoonlightError [Double] +matrixVectorProduct matrixRows vectorValue = + first (\msg -> InvariantViolation ("matrixVectorProduct: " <> msg)) (matrixVectorProductWith (*) (+) 0.0 matrixRows vectorValue) + +matrixSubtract :: [[Double]] -> [[Double]] -> Either MoonlightError [[Double]] +matrixSubtract left right + | length left /= length right = + Left (InvariantViolation ("matrixSubtract: row count mismatch (left=" <> show (length left) <> ", right=" <> show (length right) <> ")")) + | otherwise = traverse (\(l, r) -> subVector l r) (zip left right) + +scaleMatrix :: Double -> [[Double]] -> [[Double]] +scaleMatrix = scaleMatrixWith (*) + +outerProduct :: [Double] -> [Double] -> [[Double]] +outerProduct = outerProductWith (*) + +basisVector :: Int -> Int -> [Double] +basisVector size indexValue = + map (\position -> if position == indexValue then 1.0 else 0.0) [0 .. size - 1] + +linearCombination :: [(Double, [Double])] -> Either MoonlightError [Double] +linearCombination [] = Right [] +linearCombination ((firstCoefficient, firstVector) : rest) = + foldM + (\accumulator (coefficient, vectorValue) -> addVector accumulator (scaleVector coefficient vectorValue)) + (scaleVector firstCoefficient firstVector) + rest diff --git a/src-carrier/Moonlight/LinAlg/Internal/Storage.hs b/src-carrier/Moonlight/LinAlg/Internal/Storage.hs new file mode 100644 --- /dev/null +++ b/src-carrier/Moonlight/LinAlg/Internal/Storage.hs @@ -0,0 +1,130 @@ +module Moonlight.LinAlg.Internal.Storage + ( checkFlatLength, + chunkRows, + matrixMultiplyList, + matrixTransposeList, + matrixZipList, + matrixMapList, + unchunkRows, + ) +where + +import Data.Bifunctor (first) +import Moonlight.LinAlg.Pure.Dense.Rows (transposeRowsExact) +import Moonlight.LinAlg.Internal.DenseList (dotProductWith) +import Moonlight.Core + ( AdditiveMonoid (..), + MoonlightError (..), + MultiplicativeMonoid (..), + Semiring, + checkedNonNegativeProduct, + ) +import Prelude + +checkFlatLength :: Int -> Int -> [a] -> Either MoonlightError () +checkFlatLength rowCount columnCount values + | rowCount < 0 || columnCount < 0 = + Left (InvariantViolation "matrix dimensions must be non-negative") + | otherwise = do + expectedLength <- + first + (const (InvariantViolation "matrix dimensions exceed Int cardinality")) + (checkedNonNegativeProduct rowCount columnCount) + if expectedLength /= length values + then + Left + ( InvariantViolation + ( "flat payload length mismatch: expected " + <> show expectedLength + <> " values but received " + <> show (length values) + ) + ) + else Right () + +chunkRows :: Int -> [a] -> Either MoonlightError [[a]] +chunkRows columnCount values + | columnCount <= 0 && not (null values) = Left (InvariantViolation "column count must be positive when payload is non-empty") + | columnCount <= 0 = Right [] + | otherwise = Right (go values) + where + go [] = [] + go rest = + let (rowValues, nextValues) = splitAt columnCount rest + in rowValues : go nextValues + +unchunkRows :: [[a]] -> [a] +unchunkRows = concat + +matrixMapList :: Int -> Int -> (a -> b) -> [a] -> Either MoonlightError [b] +matrixMapList rowCount columnCount fn values = + checkFlatLength rowCount columnCount values *> pure (map fn values) + +matrixZipList :: + Int -> + Int -> + Int -> + Int -> + (a -> b -> c) -> + [a] -> + [b] -> + Either MoonlightError [c] +matrixZipList leftRows leftCols rightRows rightCols fn leftValues rightValues + | leftRows /= rightRows || leftCols /= rightCols = + Left + ( InvariantViolation + ( "matrix shape mismatch: left " + <> show (leftRows, leftCols) + <> " right " + <> show (rightRows, rightCols) + ) + ) + | otherwise = + checkFlatLength leftRows leftCols leftValues + *> checkFlatLength rightRows rightCols rightValues + *> pure (zipWith fn leftValues rightValues) + +matrixTransposeList :: Int -> Int -> [a] -> Either MoonlightError [a] +matrixTransposeList rowCount columnCount values = do + checkFlatLength rowCount columnCount values + rows <- chunkRows columnCount values + unchunkRows <$> transposeRowsExact rows + +matrixMultiplyList :: + Semiring a => + Int -> + Int -> + Int -> + Int -> + [a] -> + [a] -> + Either MoonlightError [a] +matrixMultiplyList leftRows leftCols rightRows rightCols leftValues rightValues + | leftCols /= rightRows = + Left + ( InvariantViolation + ( "matrix multiplication shape mismatch: left " + <> show (leftRows, leftCols) + <> " right " + <> show (rightRows, rightCols) + ) + ) + | otherwise = do + checkFlatLength leftRows leftCols leftValues + checkFlatLength rightRows rightCols rightValues + leftRowValues <- chunkRows leftCols leftValues + rightRowValues <- chunkRows rightCols rightValues + rightColumns <- transposeRowsExact rightRowValues + productRows <- + traverse + ( \rowValues -> + traverse + ( \columnValues -> + case dotProductWith mul add zero rowValues columnValues of + Left err -> Left (InvariantViolation err) + Right dotProductValue -> Right dotProductValue + ) + rightColumns + ) + leftRowValues + pure (unchunkRows productRows) diff --git a/src-carrier/Moonlight/LinAlg/Internal/VectorOps.hs b/src-carrier/Moonlight/LinAlg/Internal/VectorOps.hs new file mode 100644 --- /dev/null +++ b/src-carrier/Moonlight/LinAlg/Internal/VectorOps.hs @@ -0,0 +1,521 @@ +{-# LANGUAGE BangPatterns #-} + +module Moonlight.LinAlg.Internal.VectorOps + ( dotU, + normU, + scaleU, + subU, + subScaledU, + csrMatVecU, + csrMatVecValidatedU, + csrContiguousBandMatVecValidatedU, + csrMatVecBoxedDouble, + csrMatVecBoxedDoubleValidated, + ) +where + +import qualified Data.Vector as Box +import Control.Monad.ST (runST) +import Data.Primitive (sizeOf) +import Data.Primitive.ByteArray + ( indexByteArray, + newByteArray, + unsafeFreezeByteArray, + writeByteArray, + ) +import qualified Data.Vector.Primitive as P +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as MU +import qualified Data.Vector.Unboxed.Base as UB +import Moonlight.Core (MoonlightError (..)) +import Prelude + +dotU :: U.Vector Double -> U.Vector Double -> Either MoonlightError Double +dotU left right = + if U.length left == U.length right + then Right (U.sum (U.zipWith (*) left right)) + else + Left + ( InvariantViolation + ( "unboxed vector dot length mismatch: left " + <> show (U.length left) + <> " right " + <> show (U.length right) + ) + ) + +normU :: U.Vector Double -> Double +normU vectorValue = + sqrt (U.sum (U.map (\entry -> entry * entry) vectorValue)) + +scaleU :: Double -> U.Vector Double -> U.Vector Double +scaleU factor = U.map (factor *) + +subU :: U.Vector Double -> U.Vector Double -> Either MoonlightError (U.Vector Double) +subU left = subScaledU left 1.0 + +subScaledU :: U.Vector Double -> Double -> U.Vector Double -> Either MoonlightError (U.Vector Double) +subScaledU left factor right = + if U.length left == U.length right + then Right (U.zipWith (\leftEntry rightEntry -> leftEntry - factor * rightEntry) left right) + else + Left + ( InvariantViolation + ( "unboxed vector subtraction length mismatch: left " + <> show (U.length left) + <> " right " + <> show (U.length right) + ) + ) + +csrMatVecU :: + Int -> + U.Vector Int -> + U.Vector Int -> + U.Vector Double -> + U.Vector Double -> + Either MoonlightError (U.Vector Double) +csrMatVecU rowCount rowOffsets columnIndices values inputVector = + validateCSRKernelVectors rowCount rowOffsets columnIndices values inputVector + *> Right (csrMatVecValidatedU rowCount rowOffsets columnIndices values inputVector) + +csrMatVecValidatedU :: + Int -> + U.Vector Int -> + U.Vector Int -> + U.Vector Double -> + U.Vector Double -> + U.Vector Double +csrMatVecValidatedU + rowCount + (UB.V_Int (P.Vector rowOffsetBase _ rowOffsetArray)) + (UB.V_Int (P.Vector columnBase _ columnArray)) + (UB.V_Double (P.Vector coefficientBase _ coefficientArray)) + (UB.V_Double (P.Vector inputBase _ inputArray)) = + UB.V_Double + ( P.Vector + 0 + rowCount + ( runST $ do + targetArray <- + newByteArray + (rowCount * sizeOf (0.0 :: Double)) + + let writeRows !rowIndex + | rowIndex >= rowCount = + unsafeFreezeByteArray targetArray + | otherwise = do + let !startIndex = + indexByteArray + rowOffsetArray + (rowOffsetBase + rowIndex) + !stopIndex = + indexByteArray + rowOffsetArray + (rowOffsetBase + rowIndex + 1) + !rowValue = + accumulateRow startIndex stopIndex (0.0 :: Double) + writeByteArray targetArray rowIndex rowValue + writeRows (rowIndex + 1) + + accumulateRow :: Int -> Int -> Double -> Double + accumulateRow !entryIndex !stopIndex !accumulator + | entryIndex >= stopIndex = accumulator + | otherwise = + let !columnIndex = + indexByteArray + columnArray + (columnBase + entryIndex) + !coefficient = + ( indexByteArray + coefficientArray + (coefficientBase + entryIndex) + :: Double + ) + !inputValue = + ( indexByteArray + inputArray + (inputBase + columnIndex) + :: Double + ) + in accumulateRow + (entryIndex + 1) + stopIndex + (accumulator + coefficient * inputValue) + + writeRows 0 + ) + ) +{-# INLINE csrMatVecValidatedU #-} + +csrContiguousBandMatVecValidatedU :: + Int -> + Int -> + Int -> + U.Vector Double -> + U.Vector Double -> + U.Vector Double +csrContiguousBandMatVecValidatedU + rowCount + lowerBandwidth + upperBandwidth + coefficients + inputVector + | lowerBandwidth == 2 + && upperBandwidth == 2 + && rowCount >= 5 = + csrPentadiagonalMatVecValidatedU + rowCount + coefficients + inputVector + | lowerBandwidth == 1 + && upperBandwidth == 1 + && rowCount >= 3 = + csrTridiagonalMatVecValidatedU + rowCount + coefficients + inputVector + | otherwise = + csrContiguousBandMatVecGeneralValidatedU + rowCount + lowerBandwidth + upperBandwidth + coefficients + inputVector +{-# INLINE csrContiguousBandMatVecValidatedU #-} + +csrPentadiagonalMatVecValidatedU :: + Int -> + U.Vector Double -> + U.Vector Double -> + U.Vector Double +csrPentadiagonalMatVecValidatedU + rowCount + (UB.V_Double (P.Vector coefficientBase _ coefficientArray)) + (UB.V_Double (P.Vector inputBase _ inputArray)) = + UB.V_Double + ( P.Vector + 0 + rowCount + ( runST $ do + targetArray <- + newByteArray + (rowCount * sizeOf (0.0 :: Double)) + + let coefficientAt !indexValue = + ( indexByteArray + coefficientArray + (coefficientBase + indexValue) + :: Double + ) + inputAt !indexValue = + ( indexByteArray + inputArray + (inputBase + indexValue) + :: Double + ) + !row0 = + coefficientAt 0 * inputAt 0 + + coefficientAt 1 * inputAt 1 + + coefficientAt 2 * inputAt 2 + !row1 = + coefficientAt 3 * inputAt 0 + + coefficientAt 4 * inputAt 1 + + coefficientAt 5 * inputAt 2 + + coefficientAt 6 * inputAt 3 + + writeByteArray targetArray 0 row0 + writeByteArray targetArray 1 row1 + + let writeInterior !rowIndex + | rowIndex + 2 >= rowCount = pure () + | otherwise = do + let !entryIndex = 5 * rowIndex - 3 + !rowValue = + coefficientAt entryIndex + * inputAt (rowIndex - 2) + + coefficientAt (entryIndex + 1) + * inputAt (rowIndex - 1) + + coefficientAt (entryIndex + 2) + * inputAt rowIndex + + coefficientAt (entryIndex + 3) + * inputAt (rowIndex + 1) + + coefficientAt (entryIndex + 4) + * inputAt (rowIndex + 2) + writeByteArray targetArray rowIndex rowValue + writeInterior (rowIndex + 1) + + writeInterior 2 + + let !penultimateRow = rowCount - 2 + !penultimateEntry = 5 * rowCount - 13 + !penultimateValue = + coefficientAt penultimateEntry + * inputAt (rowCount - 4) + + coefficientAt (penultimateEntry + 1) + * inputAt (rowCount - 3) + + coefficientAt (penultimateEntry + 2) + * inputAt (rowCount - 2) + + coefficientAt (penultimateEntry + 3) + * inputAt (rowCount - 1) + !lastRow = rowCount - 1 + !lastEntry = 5 * rowCount - 9 + !lastValue = + coefficientAt lastEntry + * inputAt (rowCount - 3) + + coefficientAt (lastEntry + 1) + * inputAt (rowCount - 2) + + coefficientAt (lastEntry + 2) + * inputAt (rowCount - 1) + + writeByteArray targetArray penultimateRow penultimateValue + writeByteArray targetArray lastRow lastValue + unsafeFreezeByteArray targetArray + ) + ) +{-# INLINE csrPentadiagonalMatVecValidatedU #-} + +csrTridiagonalMatVecValidatedU :: + Int -> + U.Vector Double -> + U.Vector Double -> + U.Vector Double +csrTridiagonalMatVecValidatedU + rowCount + (UB.V_Double (P.Vector coefficientBase _ coefficientArray)) + (UB.V_Double (P.Vector inputBase _ inputArray)) = + UB.V_Double + ( P.Vector + 0 + rowCount + ( runST $ do + targetArray <- + newByteArray + (rowCount * sizeOf (0.0 :: Double)) + + let coefficientAt !indexValue = + ( indexByteArray + coefficientArray + (coefficientBase + indexValue) + :: Double + ) + inputAt !indexValue = + ( indexByteArray + inputArray + (inputBase + indexValue) + :: Double + ) + !firstValue = + coefficientAt 0 * inputAt 0 + + coefficientAt 1 * inputAt 1 + + writeByteArray targetArray 0 firstValue + + let writeInterior !rowIndex + | rowIndex + 1 >= rowCount = pure () + | otherwise = do + let !entryIndex = 3 * rowIndex - 1 + !rowValue = + coefficientAt entryIndex + * inputAt (rowIndex - 1) + + coefficientAt (entryIndex + 1) + * inputAt rowIndex + + coefficientAt (entryIndex + 2) + * inputAt (rowIndex + 1) + writeByteArray targetArray rowIndex rowValue + writeInterior (rowIndex + 1) + + writeInterior 1 + + let !lastRow = rowCount - 1 + !lastEntry = 3 * rowCount - 4 + !lastValue = + coefficientAt lastEntry + * inputAt (rowCount - 2) + + coefficientAt (lastEntry + 1) + * inputAt (rowCount - 1) + writeByteArray targetArray lastRow lastValue + unsafeFreezeByteArray targetArray + ) + ) +{-# INLINE csrTridiagonalMatVecValidatedU #-} + +csrContiguousBandMatVecGeneralValidatedU :: + Int -> + Int -> + Int -> + U.Vector Double -> + U.Vector Double -> + U.Vector Double +csrContiguousBandMatVecGeneralValidatedU + rowCount + lowerBandwidth + upperBandwidth + (UB.V_Double (P.Vector coefficientBase _ coefficientArray)) + (UB.V_Double (P.Vector inputBase _ inputArray)) = + UB.V_Double + ( P.Vector + 0 + rowCount + ( runST $ do + targetArray <- + newByteArray + (rowCount * sizeOf (0.0 :: Double)) + + let writeRows !rowIndex !entryIndex + | rowIndex >= rowCount = + unsafeFreezeByteArray targetArray + | otherwise = do + let !firstColumn = + max 0 (rowIndex - lowerBandwidth) + !lastColumn = + min + (rowCount - 1) + (rowIndex + upperBandwidth) + !entryCount = + lastColumn - firstColumn + 1 + !rowValue = + accumulateBand + entryIndex + firstColumn + entryCount + 0.0 + writeByteArray targetArray rowIndex rowValue + writeRows + (rowIndex + 1) + (entryIndex + entryCount) + + accumulateBand :: Int -> Int -> Int -> Double -> Double + accumulateBand + !entryIndex + !columnIndex + !remaining + !accumulator + | remaining <= 0 = accumulator + | otherwise = + let !coefficient = + ( indexByteArray + coefficientArray + (coefficientBase + entryIndex) + :: Double + ) + !inputValue = + ( indexByteArray + inputArray + (inputBase + columnIndex) + :: Double + ) + in accumulateBand + (entryIndex + 1) + (columnIndex + 1) + (remaining - 1) + (accumulator + coefficient * inputValue) + + writeRows 0 0 + ) + ) +{-# INLINE csrContiguousBandMatVecGeneralValidatedU #-} + +csrMatVecBoxedDouble :: + Int -> + U.Vector Int -> + U.Vector Int -> + Box.Vector Double -> + U.Vector Double -> + Either MoonlightError (U.Vector Double) +csrMatVecBoxedDouble rowCount rowOffsets columnIndices values inputVector = + validateCSRBoxedDoubleKernelVectors rowCount rowOffsets columnIndices values inputVector + *> Right (csrMatVecBoxedDoubleValidated rowCount rowOffsets columnIndices values inputVector) + +csrMatVecBoxedDoubleValidated :: + Int -> + U.Vector Int -> + U.Vector Int -> + Box.Vector Double -> + U.Vector Double -> + U.Vector Double +csrMatVecBoxedDoubleValidated rowCount rowOffsets columnIndices values inputVector = + U.create $ do + targetVector <- MU.unsafeNew rowCount + let writeRows !rowIndex + | rowIndex >= rowCount = pure targetVector + | otherwise = do + let !startIndex = rowOffsets `U.unsafeIndex` rowIndex + !stopIndex = rowOffsets `U.unsafeIndex` (rowIndex + 1) + !rowValue = accumulateRow startIndex stopIndex 0.0 + MU.unsafeWrite targetVector rowIndex rowValue + writeRows (rowIndex + 1) + + accumulateRow !entryIndex !stopIndex !accumulator + | entryIndex >= stopIndex = accumulator + | otherwise = + let !columnIndex = columnIndices `U.unsafeIndex` entryIndex + !coefficient = values `Box.unsafeIndex` entryIndex + !inputValue = inputVector `U.unsafeIndex` columnIndex + in accumulateRow + (entryIndex + 1) + stopIndex + (accumulator + coefficient * inputValue) + + writeRows 0 +{-# INLINE csrMatVecBoxedDoubleValidated #-} + +validateCSRKernelVectors :: + Int -> + U.Vector Int -> + U.Vector Int -> + U.Vector Double -> + U.Vector Double -> + Either MoonlightError () +validateCSRKernelVectors rowCount rowOffsets columnIndices values inputVector + = validateCSRKernelShape rowCount rowOffsets columnIndices (U.length values) inputVector + +validateCSRBoxedDoubleKernelVectors :: + Int -> + U.Vector Int -> + U.Vector Int -> + Box.Vector Double -> + U.Vector Double -> + Either MoonlightError () +validateCSRBoxedDoubleKernelVectors rowCount rowOffsets columnIndices values inputVector = + validateCSRKernelShape rowCount rowOffsets columnIndices (Box.length values) inputVector + +validateCSRKernelShape :: + Int -> + U.Vector Int -> + U.Vector Int -> + Int -> + U.Vector Double -> + Either MoonlightError () +validateCSRKernelShape rowCount rowOffsets columnIndices entryCount inputVector + | rowCount < 0 = Left (InvariantViolation "CSR matvec row count must be non-negative") + | U.length rowOffsets /= rowCount + 1 = + Left + ( InvariantViolation + ( "CSR row offset length mismatch: expected " + <> show (rowCount + 1) + <> " but received " + <> show (U.length rowOffsets) + ) + ) + | U.length columnIndices /= entryCount = + Left + ( InvariantViolation + ( "CSR column/value length mismatch: " + <> show (U.length columnIndices) + <> " columns but " + <> show entryCount + <> " values" + ) + ) + | not offsetsValid = Left (InvariantViolation "CSR row offsets are not a valid nondecreasing range") + | not columnsValid = Left (InvariantViolation "CSR column index out of input-vector bounds") + | otherwise = Right () + where + offsetsValid = + maybe False (== 0) (rowOffsets U.!? 0) + && maybe False (== entryCount) (rowOffsets U.!? rowCount) + && U.and (U.zipWith (<=) rowOffsets (U.drop 1 rowOffsets)) + && U.all (\offsetValue -> offsetValue >= 0 && offsetValue <= entryCount) rowOffsets + inputLength = U.length inputVector + columnsValid = U.all (\columnIndex -> columnIndex >= 0 && columnIndex < inputLength) columnIndices diff --git a/src-carrier/Moonlight/LinAlg/Pure/Dense/Flat.hs b/src-carrier/Moonlight/LinAlg/Pure/Dense/Flat.hs new file mode 100644 --- /dev/null +++ b/src-carrier/Moonlight/LinAlg/Pure/Dense/Flat.hs @@ -0,0 +1,138 @@ +{-# LANGUAGE StrictData #-} + +-- | Flat row-major dense storage for hot Double kernels. +-- +-- Nested lists remain the validation and authoring surface; this module owns +-- contiguous row-major execution. +module Moonlight.LinAlg.Pure.Dense.Flat + ( DenseDoubleMatrix, + mkDenseDoubleMatrixRowMajor, + mkDenseDoubleMatrixRows, + trustedDenseDoubleMatrixRowMajor, + denseDoubleMatrixShape, + denseDoubleMatrixToRowMajorVector, + denseDoubleMatrixToRows, + denseDoubleMatrixVectorProduct, + ) +where + +import Data.Bifunctor (first) +import Data.Kind (Type) +import Data.Vector.Storable qualified as S +import Moonlight.Core + ( MoonlightError (..), + checkedNonNegativeProduct, + fieldValueValid, + ) +import Moonlight.LinAlg.Pure.Dense.Rows + ( denseRowsShape, + denseRowsToLists, + mkDenseRows, + ) +import Prelude + +type DenseDoubleMatrix :: Type +data DenseDoubleMatrix = DenseDoubleMatrix + { denseDoubleMatrixRowCount :: !Int, + denseDoubleMatrixColumnCount :: !Int, + denseDoubleMatrixPayload :: !(S.Vector Double) + } + deriving stock (Eq, Show) + +mkDenseDoubleMatrixRowMajor :: Int -> Int -> S.Vector Double -> Either MoonlightError DenseDoubleMatrix +mkDenseDoubleMatrixRowMajor rowCount columnCount rowMajorValues + | rowCount < 0 || columnCount < 0 = + Left (InvariantViolation "dense Double matrix dimensions must be non-negative") + | otherwise = do + expectedLength <- + first + (const (InvariantViolation "dense Double matrix dimensions exceed Int cardinality")) + (checkedNonNegativeProduct rowCount columnCount) + if S.length rowMajorValues /= expectedLength + then + Left + ( InvariantViolation + ( "dense Double row-major payload length mismatch: expected " + <> show expectedLength + <> " values but received " + <> show (S.length rowMajorValues) + ) + ) + else + if S.any (not . fieldValueValid) rowMajorValues + then Left (InvariantViolation "dense Double row-major payload requires finite entries") + else + Right + ( trustedDenseDoubleMatrixRowMajor + rowCount + columnCount + rowMajorValues + ) + +trustedDenseDoubleMatrixRowMajor :: Int -> Int -> S.Vector Double -> DenseDoubleMatrix +trustedDenseDoubleMatrixRowMajor rowCount columnCount rowMajorValues = + DenseDoubleMatrix + { denseDoubleMatrixRowCount = rowCount, + denseDoubleMatrixColumnCount = columnCount, + denseDoubleMatrixPayload = rowMajorValues + } + +mkDenseDoubleMatrixRows :: [[Double]] -> Either MoonlightError DenseDoubleMatrix +mkDenseDoubleMatrixRows rowValues = do + denseRowsValue <- mkDenseRows rowValues + let (rowCount, columnCount) = denseRowsShape denseRowsValue + mkDenseDoubleMatrixRowMajor + rowCount + columnCount + (S.fromList (concat (denseRowsToLists denseRowsValue))) + +denseDoubleMatrixShape :: DenseDoubleMatrix -> (Int, Int) +denseDoubleMatrixShape matrixValue = + (denseDoubleMatrixRowCount matrixValue, denseDoubleMatrixColumnCount matrixValue) + +denseDoubleMatrixToRowMajorVector :: DenseDoubleMatrix -> S.Vector Double +denseDoubleMatrixToRowMajorVector = denseDoubleMatrixPayload + +denseDoubleMatrixToRows :: DenseDoubleMatrix -> [[Double]] +denseDoubleMatrixToRows matrixValue = + fmap rowValues [0 .. denseDoubleMatrixRowCount matrixValue - 1] + where + columnCount = denseDoubleMatrixColumnCount matrixValue + payload = denseDoubleMatrixPayload matrixValue + rowValues rowIndex = + S.toList (S.slice (rowIndex * columnCount) columnCount payload) + +denseDoubleMatrixVectorProduct :: DenseDoubleMatrix -> S.Vector Double -> Either MoonlightError (S.Vector Double) +denseDoubleMatrixVectorProduct matrixValue vectorValue = + if S.length vectorValue /= denseDoubleMatrixColumnCount matrixValue + then + Left + ( InvariantViolation + ( "dense Double matrix/vector shape mismatch (matrix=" + <> show (denseDoubleMatrixShape matrixValue) + <> ", vector=" + <> show (S.length vectorValue) + <> ")" + ) + ) + else + Right + ( S.generate + (denseDoubleMatrixRowCount matrixValue) + (denseDoubleMatrixRowDot matrixValue vectorValue) + ) +{-# INLINE denseDoubleMatrixVectorProduct #-} + +denseDoubleMatrixRowDot :: DenseDoubleMatrix -> S.Vector Double -> Int -> Double +denseDoubleMatrixRowDot matrixValue vectorValue rowIndex = + S.ifoldl' accumulateEntry 0.0 vectorValue + where + columnCount = denseDoubleMatrixColumnCount matrixValue + rowOffset = rowIndex * columnCount + payload = denseDoubleMatrixPayload matrixValue + + accumulateEntry accumulator columnIndex vectorEntry = + accumulator + + (payload `S.unsafeIndex` (rowOffset + columnIndex)) + * vectorEntry +{-# INLINE denseDoubleMatrixRowDot #-} diff --git a/src-carrier/Moonlight/LinAlg/Pure/Dense/Rows.hs b/src-carrier/Moonlight/LinAlg/Pure/Dense/Rows.hs new file mode 100644 --- /dev/null +++ b/src-carrier/Moonlight/LinAlg/Pure/Dense/Rows.hs @@ -0,0 +1,295 @@ +{-# LANGUAGE DerivingStrategies #-} + +-- | Validated rectangular row authoring surface. +-- +-- `DenseRows` exists to seal rectangular nested-list input and return precise +-- shape errors. It is deliberately not the hot dense-storage owner; use vector, +-- sparse, tridiagonal, or native kernels for benchmark-sensitive work. +module Moonlight.LinAlg.Pure.Dense.Rows + ( DenseRows, + mkDenseRows, + mkDenseRowsWithShape, + mkDenseRowsFromFlat, + denseRowsShape, + denseRowsToLists, + transposeRowsExact, + zipRowsExactWith, + matrixVectorProductRowsWith, + matrixProductRowsWith, + hcatRowsExact, + vcatRowsExact, + ) +where + +import Data.Bifunctor (first) +import Data.Kind (Type) +import Moonlight.Core + ( MoonlightError (..), + checkedNonNegativeProduct, + ) +import Prelude + +type DenseRows :: Type -> Type +data DenseRows a = DenseRows + { denseRowCount :: !Int, + denseColumnCount :: !Int, + denseRowsData :: ![[a]] + } + deriving stock (Eq, Show) + +mkDenseRows :: [[a]] -> Either MoonlightError (DenseRows a) +mkDenseRows rowValues = + case rowValues of + [] -> + Right (DenseRows 0 0 []) + firstRow : _ -> + mkDenseRowsWithShape + (length rowValues) + (length firstRow) + rowValues + +mkDenseRowsWithShape :: Int -> Int -> [[a]] -> Either MoonlightError (DenseRows a) +mkDenseRowsWithShape expectedRowCount expectedColumnCount rowValues + | expectedRowCount < 0 = + Left + ( InvariantViolation + ( "dense row matrix row count must be non-negative, received " + <> show expectedRowCount + ) + ) + | expectedColumnCount < 0 = + Left + ( InvariantViolation + ( "dense row matrix column count must be non-negative, received " + <> show expectedColumnCount + ) + ) + | actualRowCount /= expectedRowCount = + Left + ( InvariantViolation + ( "dense row matrix row count mismatch: expected " + <> show expectedRowCount + <> " rows but received " + <> show actualRowCount + ) + ) + | otherwise = + case firstMismatchedRowWidth expectedColumnCount rowValues of + Nothing -> + Right + DenseRows + { denseRowCount = expectedRowCount, + denseColumnCount = expectedColumnCount, + denseRowsData = rowValues + } + Just (rowIndex, actualColumnCount) -> + Left + ( InvariantViolation + ( "dense row matrix is ragged at row " + <> show rowIndex + <> " (expected " + <> show expectedColumnCount + <> " columns, got " + <> show actualColumnCount + <> ")" + ) + ) + where + actualRowCount = length rowValues + +firstMismatchedRowWidth :: Int -> [[a]] -> Maybe (Int, Int) +firstMismatchedRowWidth expectedColumnCount = + foldr firstMismatch Nothing . zip [0 :: Int ..] + where + firstMismatch (rowIndex, rowValues) remainingMismatch = + let actualColumnCount = length rowValues + in if actualColumnCount == expectedColumnCount + then remainingMismatch + else Just (rowIndex, actualColumnCount) + +mkDenseRowsFromFlat :: Int -> Int -> [a] -> Either MoonlightError (DenseRows a) +mkDenseRowsFromFlat rowCount columnCount values + | rowCount < 0 || columnCount < 0 = + Left (InvariantViolation "dense row matrix dimensions must be non-negative") + | otherwise = do + expectedLength <- + first + (const (InvariantViolation "dense row matrix dimensions exceed Int cardinality")) + (checkedNonNegativeProduct rowCount columnCount) + if expectedLength /= length values + then + Left + ( InvariantViolation + ( "dense row flat payload length mismatch: expected " + <> show expectedLength + <> " values but received " + <> show (length values) + ) + ) + else + if columnCount == 0 + then mkDenseRowsWithShape rowCount columnCount (replicate rowCount []) + else mkDenseRowsWithShape rowCount columnCount (flatRows rowCount columnCount values) + +flatRows :: Int -> Int -> [a] -> [[a]] +flatRows remainingRows columnCount values + | remainingRows <= 0 = [] + | otherwise = + let (rowValues, restValues) = splitAt columnCount values + in rowValues : flatRows (remainingRows - 1) columnCount restValues + +denseRowsShape :: DenseRows a -> (Int, Int) +denseRowsShape denseRowsValue = + (denseRowCount denseRowsValue, denseColumnCount denseRowsValue) + +denseRowsToLists :: DenseRows a -> [[a]] +denseRowsToLists = denseRowsData + +transposeRowsExact :: [[a]] -> Either MoonlightError [[a]] +transposeRowsExact = + fmap (denseRowsToLists . transposeDenseRows) . mkDenseRows + +zipRowsExactWith :: (left -> right -> result) -> [[left]] -> [[right]] -> Either MoonlightError [[result]] +zipRowsExactWith combine leftRows rightRows = do + leftDenseRows <- mkDenseRows leftRows + rightDenseRows <- mkDenseRows rightRows + denseRowsToLists <$> zipDenseRowsWith combine leftDenseRows rightDenseRows + +matrixVectorProductRowsWith :: + (entry -> value -> product) -> + (product -> accumulator -> accumulator) -> + accumulator -> + [[entry]] -> + [value] -> + Either MoonlightError [accumulator] +matrixVectorProductRowsWith multiply append zeroValue rowValues vectorValue = + do + denseRowsValue <- mkDenseRows rowValues + if length vectorValue /= denseColumnCount denseRowsValue + then + Left + ( InvariantViolation + ( "dense row matrix/vector shape mismatch (matrix=" + <> show (denseRowsShape denseRowsValue) + <> ", vector=" + <> show (length vectorValue) + <> ")" + ) + ) + else + Right + ( (\rowValue -> trustedDotProductWith multiply append zeroValue rowValue vectorValue) + <$> denseRowsData denseRowsValue + ) + +matrixProductRowsWith :: + (left -> right -> product) -> + (product -> accumulator -> accumulator) -> + accumulator -> + [[left]] -> + [[right]] -> + Either MoonlightError [[accumulator]] +matrixProductRowsWith multiply append zeroValue leftRows rightRows = do + leftDenseRows <- mkDenseRows leftRows + rightDenseRows <- mkDenseRows rightRows + if denseColumnCount leftDenseRows /= denseRowCount rightDenseRows + then + Left + ( InvariantViolation + ( "dense row matrix product shape mismatch (left=" + <> show (denseRowsShape leftDenseRows) + <> ", right=" + <> show (denseRowsShape rightDenseRows) + <> ")" + ) + ) + else + let rightColumns = denseRowsData (transposeDenseRows rightDenseRows) + in Right + ( (\leftRow -> trustedDotProductWith multiply append zeroValue leftRow <$> rightColumns) + <$> denseRowsData leftDenseRows + ) + +trustedDotProductWith :: + (left -> right -> product) -> + (product -> accumulator -> accumulator) -> + accumulator -> + [left] -> + [right] -> + accumulator +trustedDotProductWith multiply append zeroValue left right = + foldr append zeroValue (zipWith multiply left right) +{-# INLINE trustedDotProductWith #-} + +hcatRowsExact :: [[[a]]] -> Either MoonlightError [[a]] +hcatRowsExact rowMatrices = do + denseRowMatrices <- traverse mkDenseRows rowMatrices + case denseRowMatrices of + [] -> + Right [] + firstDenseRows : remainingDenseRows -> + if all (\denseRowsValue -> denseRowCount denseRowsValue == denseRowCount firstDenseRows) remainingDenseRows + then + pure + ( foldr + (zipWith (++)) + (replicate (denseRowCount firstDenseRows) []) + (map denseRowsData denseRowMatrices) + ) + else + Left + ( InvariantViolation + ( "dense horizontal concatenation requires equal row counts, got " + <> show (map denseRowsShape denseRowMatrices) + ) + ) + +vcatRowsExact :: [[[a]]] -> Either MoonlightError [[a]] +vcatRowsExact rowMatrices = do + denseRowMatrices <- traverse mkDenseRows rowMatrices + case denseRowMatrices of + [] -> + Right [] + firstDenseRows : remainingDenseRows -> + if all (\denseRowsValue -> denseColumnCount denseRowsValue == denseColumnCount firstDenseRows) remainingDenseRows + then pure (denseRowMatrices >>= denseRowsData) + else + Left + ( InvariantViolation + ( "dense vertical concatenation requires equal column counts, got " + <> show (map denseRowsShape denseRowMatrices) + ) + ) + +transposeDenseRows :: DenseRows a -> DenseRows a +transposeDenseRows denseRowsValue = + DenseRows + { denseRowCount = denseColumnCount denseRowsValue, + denseColumnCount = denseRowCount denseRowsValue, + denseRowsData = foldr (zipWith (:)) (replicate (denseColumnCount denseRowsValue) []) (denseRowsData denseRowsValue) + } + +zipDenseRowsWith :: (left -> right -> result) -> DenseRows left -> DenseRows right -> Either MoonlightError (DenseRows result) +zipDenseRowsWith combine leftDenseRows rightDenseRows + | denseRowsShape leftDenseRows /= denseRowsShape rightDenseRows = + Left + ( InvariantViolation + ( "dense row matrix zip shape mismatch (left=" + <> show (denseRowsShape leftDenseRows) + <> ", right=" + <> show (denseRowsShape rightDenseRows) + <> ")" + ) + ) + | otherwise = + Right + ( DenseRows + { denseRowCount = denseRowCount leftDenseRows, + denseColumnCount = denseColumnCount leftDenseRows, + denseRowsData = + zipWith + (zipWith combine) + (denseRowsData leftDenseRows) + (denseRowsData rightDenseRows) + } + ) diff --git a/src-carrier/Moonlight/LinAlg/Pure/Dense/Types.hs b/src-carrier/Moonlight/LinAlg/Pure/Dense/Types.hs new file mode 100644 --- /dev/null +++ b/src-carrier/Moonlight/LinAlg/Pure/Dense/Types.hs @@ -0,0 +1,114 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Moonlight.LinAlg.Pure.Dense.Types + ( Vector, + Matrix, + fromListVector, + fromListMatrix, + matrixRows, + toListVector, + toListMatrix, + vectorLength, + matrixShape, + matrixDenseRows, + matrixToRows, + ) +where + +import Data.Kind (Type) +import Data.Proxy (Proxy (..)) +import GHC.TypeNats (KnownNat, Nat, natVal) +import Moonlight.Core + ( MoonlightError (..), + checkedNaturalToInt, + ) +import Moonlight.LinAlg.Internal.Discrete () +import Moonlight.LinAlg.Internal.Storage (checkFlatLength) +import Moonlight.LinAlg.Pure.Dense.Rows + ( DenseRows, + denseRowsToLists, + mkDenseRowsFromFlat, + mkDenseRowsWithShape, + ) +import Prelude + +type Vector :: Nat -> Type -> Type +data Vector (n :: Nat) a = Vector + { vectorDimension :: !Int, + vectorPayload :: [a] + } + +type Matrix :: Nat -> Nat -> Type -> Type +data Matrix (r :: Nat) (c :: Nat) a = Matrix + { matrixRowCount :: !Int, + matrixColumnCount :: !Int, + matrixPayload :: [a] + } + +checkedTypeLevelDimension :: forall n. KnownNat n => Either MoonlightError Int +checkedTypeLevelDimension = + either + (const (Left (InvariantViolation "type-level dimension exceeds Int cardinality"))) + Right + (checkedNaturalToInt (natVal (Proxy @n))) + +fromListVector :: forall n a. KnownNat n => [a] -> Either MoonlightError (Vector n a) +fromListVector values = do + expectedLength <- checkedTypeLevelDimension @n + if length values /= expectedLength + then + Left + ( InvariantViolation + ( "vector length mismatch: expected " + <> show expectedLength + <> " values but received " + <> show (length values) + ) + ) + else Right (Vector expectedLength values) + +fromListMatrix :: forall r c a. (KnownNat r, KnownNat c) => [a] -> Either MoonlightError (Matrix r c a) +fromListMatrix values = do + rowCount <- checkedTypeLevelDimension @r + columnCount <- checkedTypeLevelDimension @c + checkFlatLength rowCount columnCount values + Right (Matrix rowCount columnCount values) + +matrixRows :: forall r c a. (KnownNat r, KnownNat c) => [[a]] -> Either MoonlightError (Matrix r c a) +matrixRows rowValues = do + rowCount <- checkedTypeLevelDimension @r + columnCount <- checkedTypeLevelDimension @c + denseRowsValue <- + mkDenseRowsWithShape + rowCount + columnCount + rowValues + Right (Matrix rowCount columnCount (concat (denseRowsToLists denseRowsValue))) + +toListVector :: Vector n a -> [a] +toListVector = vectorPayload + +toListMatrix :: Matrix r c a -> [a] +toListMatrix = matrixPayload + +vectorLength :: forall n a. KnownNat n => Vector n a -> Int +vectorLength vectorValue = + natVal (Proxy @n) `seq` vectorDimension vectorValue + +matrixShape :: forall r c a. (KnownNat r, KnownNat c) => Matrix r c a -> (Int, Int) +matrixShape matrixValue = + natVal (Proxy @r) + `seq` natVal (Proxy @c) + `seq` (matrixRowCount matrixValue, matrixColumnCount matrixValue) + +matrixDenseRows :: forall r c a. (KnownNat r, KnownNat c) => Matrix r c a -> Either MoonlightError (DenseRows a) +matrixDenseRows matrixValue = + let (rowCount, columnCount) = matrixShape matrixValue + in mkDenseRowsFromFlat + rowCount + columnCount + (toListMatrix matrixValue) + +matrixToRows :: forall r c a. (KnownNat r, KnownNat c) => Matrix r c a -> Either MoonlightError [[a]] +matrixToRows = + fmap denseRowsToLists . matrixDenseRows diff --git a/src-dense/Moonlight/LinAlg/Internal/Backend/Core.hs b/src-dense/Moonlight/LinAlg/Internal/Backend/Core.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Internal/Backend/Core.hs @@ -0,0 +1,90 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DefaultSignatures #-} + +module Moonlight.LinAlg.Internal.Backend.Core + ( DenseRankBackend (..), + runPluDecomp, + runKernel, + runSmithNormalForm, + runSmithDiagonalForm, + ) +where + +import Data.Kind (Constraint, Type) +import GHC.TypeNats (KnownNat) +import Moonlight.Algebra.Pure.Ring (EuclideanDomain) +import Moonlight.Core + ( Field, + MoonlightError (..), + ) +import Moonlight.LinAlg.Internal.Backend.PLU (PLU, pluDecompPure) +import Moonlight.LinAlg.Internal.Backend.RREF (KernelBasis, kernelPure, rankPure) +import Moonlight.LinAlg.Internal.Backend.Smith (SmithDiagonalForm, SmithNormalForm, smithDiagonalFormPure, smithNormalFormPure) +import Moonlight.LinAlg.Pure.Dense.GF2 + ( GF2 + , mkGF2PackedMatrixFromRowMajor + , rankGF2PackedMatrix + ) +import Moonlight.LinAlg.Pure.Dense.Types (Matrix, matrixShape, toListMatrix) +import Prelude + +type DenseRankBackend :: Type -> Constraint +class DenseRankBackend a where + runRank :: + forall r c. + (KnownNat r, KnownNat c, Eq a, Field a) => + Matrix r c a -> + Either MoonlightError Int + default runRank :: + forall r c. + (KnownNat r, KnownNat c, Eq a, Field a) => + Matrix r c a -> + Either MoonlightError Int + runRank = rankPure + +instance DenseRankBackend Double + +instance DenseRankBackend GF2 where + runRank matrixValue = + let (rowCount, columnCount) = + matrixShape matrixValue + in either + (Left . InvariantViolation . ("GF2 DenseRank: " <>) . show) + (Right . rankGF2PackedMatrix) + ( mkGF2PackedMatrixFromRowMajor + (fromIntegral rowCount) + (fromIntegral columnCount) + (toListMatrix matrixValue) + ) + +instance DenseRankBackend Integer + +instance DenseRankBackend Rational + +runPluDecomp :: + forall r c a. + (KnownNat r, KnownNat c, Field a) => + Matrix r c a -> + Either MoonlightError (PLU r c a) +runPluDecomp = pluDecompPure + +runKernel :: + forall r c a. + (KnownNat r, KnownNat c, Eq a, Field a) => + Matrix r c a -> + Either MoonlightError (KernelBasis c a) +runKernel = kernelPure + +runSmithNormalForm :: + forall r c a. + (KnownNat r, KnownNat c, EuclideanDomain a) => + Matrix r c a -> + Either MoonlightError (SmithNormalForm r c a) +runSmithNormalForm = smithNormalFormPure + +runSmithDiagonalForm :: + forall r c a. + (KnownNat r, KnownNat c, EuclideanDomain a) => + Matrix r c a -> + Either MoonlightError (SmithDiagonalForm r c a) +runSmithDiagonalForm = smithDiagonalFormPure diff --git a/src-dense/Moonlight/LinAlg/Internal/Backend/Elimination.hs b/src-dense/Moonlight/LinAlg/Internal/Backend/Elimination.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Internal/Backend/Elimination.hs @@ -0,0 +1,119 @@ +module Moonlight.LinAlg.Internal.Backend.Elimination + ( EliminationScope (..), + PivotResult (..), + EliminationState (..), + EliminationConfig (..), + runElimination, + ) +where + +import Control.Monad (foldM) +import Data.Kind (Type) +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Internal.Backend.RowOps (swapAt) +import Moonlight.LinAlg.Internal.Primitives + ( ColumnIndex, + RowIndex, + mkRowIndex, + replaceRowChecked, + requireRow, + rowIndexInt, + rowIndices, + ) +import Prelude + +type EliminationScope :: Type +data EliminationScope + = ForwardOnly + | FullReduction + +type PivotResult :: Type +data PivotResult + = PivotFound RowIndex + | NoPivotSkip + | NoPivotFail + +type EliminationState :: Type -> Type -> Type +data EliminationState a s = EliminationState + { elimRows :: [[a]], + elimSide :: s, + elimPivots :: [ColumnIndex] + } + +type EliminationConfig :: Type -> Type -> Type +data EliminationConfig a s = EliminationConfig + { elimSelectPivot :: Int -> ColumnIndex -> [[a]] -> Either MoonlightError PivotResult, + elimCandidateColumns :: Int -> [ColumnIndex] -> [ColumnIndex], + elimNormalizePivot :: RowIndex -> ColumnIndex -> [[a]] -> Either MoonlightError [[a]], + elimScope :: EliminationScope, + elimReduceRow :: [a] -> [a] -> ColumnIndex -> RowIndex -> RowIndex -> s -> Either MoonlightError ([a], s), + elimOnSwap :: RowIndex -> RowIndex -> s -> Either MoonlightError s, + elimMaxSteps :: Int + } + +runElimination :: + EliminationConfig a s -> + [[a]] -> + s -> + [ColumnIndex] -> + Either MoonlightError (EliminationState a s) +runElimination config initialRows initialSide initialColumns = + go 0 initialRows initialSide initialColumns [] + where + go step rows side remainingCols pivotsSoFar + | step >= elimMaxSteps config = + Right (EliminationState rows side (reverse pivotsSoFar)) + | otherwise = + tryColumns step rows side remainingCols pivotsSoFar (elimCandidateColumns config step remainingCols) + + tryColumns _step rows side _remainingCols pivotsSoFar [] = + Right (EliminationState rows side (reverse pivotsSoFar)) + tryColumns step rows side remainingCols pivotsSoFar (col : moreCols) = do + pivotResult <- elimSelectPivot config step col rows + case pivotResult of + NoPivotFail -> + Left (InvariantViolation ("elimination failed: no pivot at step " <> show step)) + NoPivotSkip -> + tryColumns step rows side remainingCols pivotsSoFar moreCols + PivotFound sourceRow -> do + let targetRowInt = step + rowCount = length rows + targetRow <- targetRowIndex rowCount targetRowInt + swappedRows <- swapAt targetRow sourceRow rows + swappedSide <- elimOnSwap config targetRow sourceRow side + normalizedRows <- elimNormalizePivot config targetRow col swappedRows + pivotRowValues <- + requireRow + (InvariantViolation ("elimination pivot row missing at index " <> show targetRow)) + targetRow + normalizedRows + let targetIndices = case elimScope config of + ForwardOnly -> drop (step + 1) (rowIndices rowCount) + FullReduction -> filter (\ri -> rowIndexInt ri /= targetRowInt) (rowIndices rowCount) + (eliminatedRows, finalSide) <- + foldM + ( \(currentRows, currentSide) ri -> do + targetRowValues <- + requireRow + (InvariantViolation ("elimination target row missing at index " <> show ri)) + ri + currentRows + (reducedRow, nextSide) <- elimReduceRow config pivotRowValues targetRowValues col targetRow ri currentSide + nextRows <- + replaceRowChecked + (InvariantViolation ("elimination row replacement failed at index " <> show ri)) + ri + reducedRow + currentRows + Right (nextRows, nextSide) + ) + (normalizedRows, swappedSide) + targetIndices + let nextCols = dropWhile (<= col) remainingCols + go (step + 1) eliminatedRows finalSide nextCols (col : pivotsSoFar) + + targetRowIndex rowCount idx = + mkRowIndex + (InvariantViolation ("elimination target row index out of bounds: " <> show idx)) + rowCount + idx diff --git a/src-dense/Moonlight/LinAlg/Internal/Backend/PLU.hs b/src-dense/Moonlight/LinAlg/Internal/Backend/PLU.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Internal/Backend/PLU.hs @@ -0,0 +1,205 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Moonlight.LinAlg.Internal.Backend.PLU + ( PLU (..), + pluDecompPure, + ) +where + +import Data.Kind (Type) +import GHC.TypeNats (KnownNat, Nat) +import Moonlight.Core + ( AdditiveGroup (..), + Field (..), + MoonlightError (..), + MultiplicativeMonoid (..), + requireInvertible, + ) +import Moonlight.Core (note, safeIndex) +import Moonlight.LinAlg.Internal.Backend.Elimination + ( EliminationConfig (..), + EliminationScope (..), + EliminationState (..), + PivotResult (..), + runElimination, + ) +import Moonlight.LinAlg.Internal.Backend.RowOps + ( findPivotRow, + identityRows, + permutationRows, + swapLowerPrefix, + ) +import Moonlight.LinAlg.Internal.Primitives + ( ColumnIndex, + RowIndex, + columnIndices, + mkRowIndex, + replaceColumnEntryChecked, + replaceRowChecked, + requireColumnEntry, + requireRow, + rowIndexInt, + ) +import Moonlight.LinAlg.Pure.Dense.Types + ( Matrix, + fromListMatrix, + ) +import qualified Moonlight.LinAlg.Pure.Dense.Types as DenseTypes +import Prelude + +type PLU :: Nat -> Nat -> Type -> Type +data PLU r c a = PLU + { pluPermutation :: Matrix r r a, + pluLower :: Matrix r r a, + pluUpper :: Matrix r c a + } + +type PLUSideState :: Type -> Type +data PLUSideState a = PLUSideState + { pluSidePermutation :: [Int], + pluSideLower :: [[a]], + pluSideStep :: Int + } + +pluConfig :: + Field a => + Int -> + Int -> + EliminationConfig a (PLUSideState a) +pluConfig rowCount columnCount = + EliminationConfig + { elimSelectPivot = pluSelectPivot rowCount, + elimCandidateColumns = \_ cols -> take 1 cols, + elimNormalizePivot = \_ _ rows -> Right rows, + elimScope = ForwardOnly, + elimReduceRow = pluReduceRow, + elimOnSwap = pluOnSwap, + elimMaxSteps = min rowCount columnCount + } + +pluSelectPivot :: + Field a => + Int -> + Int -> + ColumnIndex -> + [[a]] -> + Either MoonlightError PivotResult +pluSelectPivot rowCount step col rows = do + pivotRowIndex <- + mkRowIndex + (InvariantViolation ("PLU pivot row out of bounds at index " <> show step)) + rowCount + step + findPivotRow pivotRowIndex col rows >>= \maybePivot -> + case maybePivot of + Nothing -> + Right NoPivotFail + Just pivotIndex -> + Right (PivotFound pivotIndex) + +pluOnSwap :: + RowIndex -> + RowIndex -> + PLUSideState a -> + Either MoonlightError (PLUSideState a) +pluOnSwap targetRow sourceRow sideState = do + let step = pluSideStep sideState + swappedPermutation <- + swapPermutationAt targetRow sourceRow (pluSidePermutation sideState) + swappedLower <- swapLowerPrefix step targetRow sourceRow (pluSideLower sideState) + Right (sideState {pluSidePermutation = swappedPermutation, pluSideLower = swappedLower, pluSideStep = step + 1}) + +pluReduceRow :: + (Field a) => + [a] -> + [a] -> + ColumnIndex -> + RowIndex -> + RowIndex -> + PLUSideState a -> + Either MoonlightError ([a], PLUSideState a) +pluReduceRow pivotRowValues targetRowValues pivotColumn _pivotRow targetRow sideState = do + pivotValue <- + requireColumnEntry + (InvariantViolation ("PLU pivot entry missing at column " <> show pivotColumn)) + pivotColumn + pivotRowValues + pivotInverse <- + requireInvertible + (InvariantViolation ("PLU decomposition failed: pivot is not invertible at column " <> show pivotColumn)) + pivotValue + factorEntry <- + requireColumnEntry + (InvariantViolation ("PLU factor entry missing at pivot column " <> show pivotColumn)) + pivotColumn + targetRowValues + let factor = factorEntry `mul` pivotInverse + updatedRow = zipWith (\entry pivotEntry -> entry `sub` (factor `mul` pivotEntry)) targetRowValues pivotRowValues + currentLRow <- + requireRow + (InvariantViolation ("PLU lower row missing at index " <> show targetRow)) + targetRow + (pluSideLower sideState) + updatedLRow <- + replaceColumnEntryChecked + (InvariantViolation ("PLU lower factor placement failed at row " <> show targetRow <> ", column " <> show pivotColumn)) + pivotColumn + factor + currentLRow + updatedLRows <- + replaceRowChecked + (InvariantViolation ("PLU lower row replacement failed at index " <> show targetRow)) + targetRow + updatedLRow + (pluSideLower sideState) + Right (updatedRow, sideState {pluSideLower = updatedLRows}) + +swapPermutationAt :: RowIndex -> RowIndex -> [Int] -> Either MoonlightError [Int] +swapPermutationAt targetRow sourceRow permIndices = do + let targetIdx = rowIndexInt targetRow + sourceIdx = rowIndexInt sourceRow + if targetIdx == sourceIdx + then Right permIndices + else do + targetVal <- + note (InvariantViolation ("PLU permutation swap out of bounds at index " <> show targetIdx)) + (safeIndex targetIdx permIndices) + sourceVal <- + note (InvariantViolation ("PLU permutation swap out of bounds at index " <> show sourceIdx)) + (safeIndex sourceIdx permIndices) + Right (replaceAtPure sourceIdx targetVal (replaceAtPure targetIdx sourceVal permIndices)) + +replaceAtPure :: Int -> a -> [a] -> [a] +replaceAtPure idx val xs = + zipWith (\i x -> if i == idx then val else x) [0 :: Int ..] xs + +pluDecompPure :: + forall r c a. + (KnownNat r, KnownNat c, Field a) => + Matrix r c a -> + Either MoonlightError (PLU r c a) +pluDecompPure matrixValue = do + initialRows <- DenseTypes.matrixToRows matrixValue + let (rowCount, columnCount) = DenseTypes.matrixShape matrixValue + initialSide = + PLUSideState + { pluSidePermutation = [0 .. rowCount - 1], + pluSideLower = identityRows rowCount, + pluSideStep = 0 + } + eliminationResult <- + runElimination + (pluConfig rowCount columnCount) + initialRows + initialSide + (columnIndices columnCount) + let finalSide = elimSide eliminationResult + pMatrix <- fromListMatrix @r @r (concat (permutationRows (pluSidePermutation finalSide))) + lMatrix <- fromListMatrix @r @r (concat (pluSideLower finalSide)) + uMatrix <- fromListMatrix @r @c (concat (elimRows eliminationResult)) + pure + PLU + { pluPermutation = pMatrix, + pluLower = lMatrix, + pluUpper = uMatrix + } diff --git a/src-dense/Moonlight/LinAlg/Internal/Backend/RREF.hs b/src-dense/Moonlight/LinAlg/Internal/Backend/RREF.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Internal/Backend/RREF.hs @@ -0,0 +1,214 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Moonlight.LinAlg.Internal.Backend.RREF + ( RREF (..), + KernelBasis (..), + rrefRowsFrom, + rrefFromMatrix, + rankPure, + kernelPure, + ) +where + +import Data.Kind (Type) +import GHC.TypeNats (KnownNat, Nat) +import Moonlight.Core + ( AdditiveGroup (..), + AdditiveMonoid (..), + Field (..), + MoonlightError (..), + MultiplicativeMonoid (..), + requireInvertible, + ) +import Moonlight.LinAlg.Internal.Backend.Elimination + ( EliminationConfig (..), + EliminationScope (..), + EliminationState (..), + PivotResult (..), + runElimination, + ) +import Moonlight.LinAlg.Internal.Backend.RowOps + ( findPivotRow, + ) +import Moonlight.LinAlg.Internal.Primitives + ( ColumnIndex, + RowIndex, + columnIndexInt, + columnIndices, + mkRowIndex, + replaceAt, + requireColumnEntry, + requireRow, + rowIndexInt, + ) +import Moonlight.LinAlg.Pure.Dense.Types + ( Matrix, + Vector, + fromListVector, + ) +import qualified Moonlight.LinAlg.Pure.Dense.Types as DenseTypes +import Prelude + +type RREF :: Type -> Type -> Type +data RREF pivot a = RREF + { rrefPivotColumns :: [pivot], + rrefRows :: [[a]] + } + deriving stock (Eq, Show) + +type KernelBasis :: Nat -> Type -> Type +newtype KernelBasis c a = KernelBasis + { kernelBasisVectors :: [Vector c a] + } + +rrefConfigAt :: + (Field a, Eq a) => + Int -> + Int -> + EliminationConfig a () +rrefConfigAt rowCount pivotRow = + EliminationConfig + { elimSelectPivot = \step col rs -> do + pivotRowIndex <- + mkRowIndex + (InvariantViolation ("RREF pivot row out of bounds at index " <> show (pivotRow + step))) + rowCount + (pivotRow + step) + findPivotRow pivotRowIndex col rs >>= \maybePivot -> + case maybePivot of + Nothing -> Right NoPivotSkip + Just pivotIndex -> Right (PivotFound pivotIndex), + elimCandidateColumns = \_ cols -> cols, + elimNormalizePivot = rrefNormalizePivot, + elimScope = FullReduction, + elimReduceRow = rrefReduceRow, + elimOnSwap = \_ _ s -> Right s, + elimMaxSteps = rowCount - pivotRow + } + +rrefNormalizePivot :: + (Field a) => + RowIndex -> + ColumnIndex -> + [[a]] -> + Either MoonlightError [[a]] +rrefNormalizePivot pivotRow pivotColumn rows = do + pivotRowValues <- + requireRow + (InvariantViolation ("RREF pivot row missing after swap at row " <> show pivotRow)) + pivotRow + rows + pivotValue <- + requireColumnEntry + (InvariantViolation ("RREF pivot column missing at column " <> show pivotColumn)) + pivotColumn + pivotRowValues + pivotInverse <- + requireInvertible + (InvariantViolation ("RREF pivot at column " <> show pivotColumn <> " is not invertible")) + pivotValue + let normalizedRow = map (\entry -> entry `mul` pivotInverse) pivotRowValues + Right (replaceAt (rowIndexInt pivotRow) normalizedRow rows) + +rrefReduceRow :: + (Field a, Eq a) => + [a] -> + [a] -> + ColumnIndex -> + RowIndex -> + RowIndex -> + () -> + Either MoonlightError ([a], ()) +rrefReduceRow pivotRowValues targetRowValues pivotColumn _ _ sideState = do + factorEntry <- + requireColumnEntry + (InvariantViolation ("RREF elimination missing pivot column " <> show pivotColumn)) + pivotColumn + targetRowValues + if factorEntry == zero + then Right (targetRowValues, sideState) + else Right (zipWith (\entry pivotEntry -> entry `sub` (factorEntry `mul` pivotEntry)) targetRowValues pivotRowValues, sideState) + +rrefRowsFrom :: (Field a, Eq a) => Int -> Int -> Int -> [[a]] -> Either MoonlightError (RREF Int a) +rrefRowsFrom rowCount columnCount pivotRow rows + | pivotRow < 0 = Left (InvariantViolation ("RREF pivot row cannot be negative: " <> show pivotRow)) + | otherwise = do + result <- + runElimination + (rrefConfigAt rowCount pivotRow) + rows + () + (columnIndices columnCount) + Right + RREF + { rrefPivotColumns = map columnIndexInt (elimPivots result), + rrefRows = elimRows result + } + +rrefFromMatrixIndexed :: + forall r c a. + (KnownNat r, KnownNat c, Eq a, Field a) => + Matrix r c a -> + Either MoonlightError (RREF ColumnIndex a) +rrefFromMatrixIndexed matrixValue = do + rows <- DenseTypes.matrixToRows matrixValue + let (rowCount, columnCount) = DenseTypes.matrixShape matrixValue + result <- + runElimination + (rrefConfigAt rowCount 0) + rows + () + (columnIndices columnCount) + Right + RREF + { rrefPivotColumns = elimPivots result, + rrefRows = elimRows result + } + +rrefFromMatrix :: + forall r c a. + (KnownNat r, KnownNat c, Eq a, Field a) => + Matrix r c a -> + Either MoonlightError (RREF Int a) +rrefFromMatrix matrixValue = + fmap + ( \rrefValue -> + RREF + { rrefPivotColumns = map columnIndexInt (rrefPivotColumns rrefValue), + rrefRows = rrefRows rrefValue + } + ) + (rrefFromMatrixIndexed matrixValue) + +rankPure :: + forall r c a. + (KnownNat r, KnownNat c, Eq a, Field a) => + Matrix r c a -> + Either MoonlightError Int +rankPure matrixValue = + fmap (length . rrefPivotColumns) (rrefFromMatrixIndexed matrixValue) + +kernelPure :: + forall r c a. + (KnownNat r, KnownNat c, Eq a, Field a) => + Matrix r c a -> + Either MoonlightError (KernelBasis c a) +kernelPure matrixValue = do + rrefValue <- rrefFromMatrixIndexed matrixValue + let (_, columnCount) = DenseTypes.matrixShape matrixValue + allColumns = columnIndices columnCount + pivotColumns = rrefPivotColumns rrefValue + reducedRows = rrefRows rrefValue + pivotRows = zip pivotColumns reducedRows + freeColumns = filter (\columnIndex -> columnIndex `notElem` pivotColumns) allColumns + basisVector freeColumn = + let entryAt columnIndex = + if columnIndex == freeColumn + then one + else + case lookup columnIndex pivotRows of + Nothing -> zero + Just pivotRowValues -> + maybe zero neg (lookup freeColumn (zip allColumns pivotRowValues)) + in fromListVector (map entryAt allColumns) + KernelBasis <$> traverse basisVector freeColumns diff --git a/src-dense/Moonlight/LinAlg/Internal/Backend/RowOps.hs b/src-dense/Moonlight/LinAlg/Internal/Backend/RowOps.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Internal/Backend/RowOps.hs @@ -0,0 +1,105 @@ +module Moonlight.LinAlg.Internal.Backend.RowOps + ( swapAt, + identityRows, + permutationRows, + swapLowerPrefix, + rowEliminate, + findPivotRow, + ) +where + +import Data.List (findIndex) +import Moonlight.Core + ( AdditiveGroup (..), + AdditiveMonoid (..), + Field (..), + MoonlightError (..), + MultiplicativeMonoid (..), + ) +import Moonlight.LinAlg.Internal.Primitives + ( ColumnIndex, + MatrixIndex, + RowIndex, + replaceRowChecked, + requireColumnEntry, + requireRow, + rowIndices, + rowIndexInt, + selectAt, + selectAtIndex, + swapAtIndexChecked, + ) +import Prelude + +swapAt :: MatrixIndex axis -> MatrixIndex axis -> [a] -> Either MoonlightError [a] +swapAt leftIndex rightIndex values = + swapAtIndexChecked + (InvariantViolation ("row operation swap out of bounds at indices " <> show (leftIndex, rightIndex))) + leftIndex + rightIndex + values + +identityRows :: (AdditiveGroup a, MultiplicativeMonoid a) => Int -> [[a]] +identityRows size = + map + (\rowIndex -> map (\columnIndex -> if rowIndex == columnIndex then one else zero) [0 .. size - 1]) + [0 .. size - 1] + +permutationRows :: (AdditiveGroup a, MultiplicativeMonoid a) => [Int] -> [[a]] +permutationRows permutationIndices = + map + (\sourceRowIndex -> map (\columnIndex -> if columnIndex == sourceRowIndex then one else zero) [0 .. length permutationIndices - 1]) + permutationIndices + +swapLowerPrefix :: Int -> RowIndex -> RowIndex -> [[a]] -> Either MoonlightError [[a]] +swapLowerPrefix prefixLength leftIndex rightIndex rows = + requireRow + (InvariantViolation ("row-prefix swap left row missing at index " <> show leftIndex)) + leftIndex + rows + >>= \leftRow -> + requireRow + (InvariantViolation ("row-prefix swap right row missing at index " <> show rightIndex)) + rightIndex + rows + >>= \rightRow -> + let swappedLeft = take prefixLength rightRow <> drop prefixLength leftRow + swappedRight = take prefixLength leftRow <> drop prefixLength rightRow + in replaceRowChecked + (InvariantViolation ("row-prefix swap could not replace left row at index " <> show leftIndex)) + leftIndex + swappedLeft + rows + >>= replaceRowChecked + (InvariantViolation ("row-prefix swap could not replace right row at index " <> show rightIndex)) + rightIndex + swappedRight + +rowEliminate :: (Field a, Eq a) => [a] -> [a] -> ColumnIndex -> Either MoonlightError [a] +rowEliminate pivotRow rowValues pivotColumn = + requireColumnEntry + (InvariantViolation ("row elimination missing pivot column " <> show pivotColumn)) + pivotColumn + rowValues + >>= \factor -> + if factor == zero + then Right rowValues + else Right (zipWith (\entry pivotEntry -> entry `sub` (factor `mul` pivotEntry)) rowValues pivotRow) + +findPivotRow :: Field a => RowIndex -> ColumnIndex -> [[a]] -> Either MoonlightError (Maybe RowIndex) +findPivotRow pivotRow pivotColumn rows = + fmap + (\candidateFlags -> + findIndex id candidateFlags + >>= \offsetIndex -> + selectAt (rowIndexInt pivotRow + offsetIndex) (rowIndices (length rows)) + ) + (traverse candidateFlag (drop (rowIndexInt pivotRow) rows)) + where + candidateFlag rowValues = + case selectAtIndex pivotColumn rowValues of + Nothing -> Left (InvariantViolation ("pivot search missing column " <> show pivotColumn)) + Just entryValue + | not (fieldValueValid entryValue) -> + Left (InvariantViolation ("pivot search encountered an invalid field value at column " <> show pivotColumn)) + | otherwise -> Right (canInvert entryValue) diff --git a/src-dense/Moonlight/LinAlg/Internal/Backend/RowStore.hs b/src-dense/Moonlight/LinAlg/Internal/Backend/RowStore.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Internal/Backend/RowStore.hs @@ -0,0 +1,137 @@ +module Moonlight.LinAlg.Internal.Backend.RowStore + ( RowStore, + rowStoreFromRows, + rowStoreToRows, + rowStoreFlatten, + rowStoreShape, + rowStoreRowAt, + rowStoreRowAtInt, + rowStoreValueAt, + rowStoreValueAtInt, + replaceRowStore, + replaceRowStoreAtInt, + swapRowsStore, + swapRowsStoreAtInt, + swapColumnsStore, + columnStore, + replaceColumnStore, + traverseRowStoreWithIndex, + ) +where + +import Data.Kind (Type) +import Data.Vector qualified as Box +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Internal.Primitives + ( ColumnIndex, + RowIndex, + columnIndexInt, + rowIndexInt, + ) +import Prelude + +type RowStore :: Type -> Type +newtype RowStore a = RowStore (Box.Vector (Box.Vector a)) + deriving stock (Eq, Show) + +rowStoreFromRows :: [[a]] -> RowStore a +rowStoreFromRows = + RowStore . Box.fromList . fmap Box.fromList + +rowStoreToRows :: RowStore a -> [[a]] +rowStoreToRows (RowStore rows) = + Box.toList (fmap Box.toList rows) + +rowStoreFlatten :: RowStore a -> [a] +rowStoreFlatten = + concat . rowStoreToRows + +rowStoreShape :: RowStore a -> (Int, Int) +rowStoreShape (RowStore rows) = + ( Box.length rows, + maybe 0 Box.length (rows Box.!? 0) + ) + +rowStoreRowAt :: MoonlightError -> RowIndex -> RowStore a -> Either MoonlightError (Box.Vector a) +rowStoreRowAt failure rowIndex = + rowStoreRowAtInt failure (rowIndexInt rowIndex) + +rowStoreRowAtInt :: MoonlightError -> Int -> RowStore a -> Either MoonlightError (Box.Vector a) +rowStoreRowAtInt failure rowIndex (RowStore rows) = + maybe (Left failure) Right (rows Box.!? rowIndex) + +rowStoreValueAt :: MoonlightError -> RowIndex -> ColumnIndex -> RowStore a -> Either MoonlightError a +rowStoreValueAt failure rowIndex columnIndex = + rowStoreValueAtInt failure (rowIndexInt rowIndex) (columnIndexInt columnIndex) + +rowStoreValueAtInt :: MoonlightError -> Int -> Int -> RowStore a -> Either MoonlightError a +rowStoreValueAtInt failure rowIndex columnIndex store = + rowStoreRowAtInt failure rowIndex store + >>= \rowValues -> maybe (Left failure) Right (rowValues Box.!? columnIndex) + +replaceRowStore :: MoonlightError -> RowIndex -> Box.Vector a -> RowStore a -> Either MoonlightError (RowStore a) +replaceRowStore failure rowIndex = + replaceRowStoreAtInt failure (rowIndexInt rowIndex) + +replaceRowStoreAtInt :: MoonlightError -> Int -> Box.Vector a -> RowStore a -> Either MoonlightError (RowStore a) +replaceRowStoreAtInt failure rowIndex replacement (RowStore rows) = + case rows Box.!? rowIndex of + Nothing -> Left failure + Just _ -> Right (RowStore (rows Box.// [(rowIndex, replacement)])) + +swapRowsStore :: MoonlightError -> RowIndex -> RowIndex -> RowStore a -> Either MoonlightError (RowStore a) +swapRowsStore failure leftIndex rightIndex = + swapRowsStoreAtInt failure (rowIndexInt leftIndex) (rowIndexInt rightIndex) + +swapRowsStoreAtInt :: MoonlightError -> Int -> Int -> RowStore a -> Either MoonlightError (RowStore a) +swapRowsStoreAtInt failure leftIndex rightIndex (RowStore rows) = + case (rows Box.!? leftIndex, rows Box.!? rightIndex) of + (Just leftRow, Just rightRow) -> + Right (RowStore (rows Box.// [(leftIndex, rightRow), (rightIndex, leftRow)])) + _ -> Left failure + +swapColumnsStore :: MoonlightError -> ColumnIndex -> ColumnIndex -> RowStore a -> Either MoonlightError (RowStore a) +swapColumnsStore failure leftIndex rightIndex = + traverseRowStoreWithIndex + ( \_ rowValues -> + swapVectorAt failure (columnIndexInt leftIndex) (columnIndexInt rightIndex) rowValues + ) + +columnStore :: MoonlightError -> ColumnIndex -> RowStore a -> Either MoonlightError (Box.Vector a) +columnStore failure columnIndex (RowStore rows) = + traverse + (\rowValues -> maybe (Left failure) Right (rowValues Box.!? columnIndexInt columnIndex)) + rows + +replaceColumnStore :: MoonlightError -> ColumnIndex -> Box.Vector a -> RowStore a -> Either MoonlightError (RowStore a) +replaceColumnStore failure columnIndex columnValues store@(RowStore rows) + | Box.length columnValues /= fst (rowStoreShape store) = Left failure + | otherwise = + traverseRowStoreWithIndex + ( \rowIndex rowValues -> + case columnValues Box.!? rowIndex of + Nothing -> Left failure + Just columnValue -> + replaceVectorAt failure (columnIndexInt columnIndex) columnValue rowValues + ) + (RowStore rows) + +traverseRowStoreWithIndex :: + (Int -> Box.Vector a -> Either MoonlightError (Box.Vector b)) -> + RowStore a -> + Either MoonlightError (RowStore b) +traverseRowStoreWithIndex transform (RowStore rows) = + RowStore <$> Box.imapM transform rows + +swapVectorAt :: MoonlightError -> Int -> Int -> Box.Vector a -> Either MoonlightError (Box.Vector a) +swapVectorAt failure leftIndex rightIndex values = + case (values Box.!? leftIndex, values Box.!? rightIndex) of + (Just leftValue, Just rightValue) -> + Right (values Box.// [(leftIndex, rightValue), (rightIndex, leftValue)]) + _ -> Left failure + +replaceVectorAt :: MoonlightError -> Int -> a -> Box.Vector a -> Either MoonlightError (Box.Vector a) +replaceVectorAt failure indexValue replacement values = + case values Box.!? indexValue of + Nothing -> Left failure + Just _ -> Right (values Box.// [(indexValue, replacement)]) diff --git a/src-dense/Moonlight/LinAlg/Internal/Backend/Smith.hs b/src-dense/Moonlight/LinAlg/Internal/Backend/Smith.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Internal/Backend/Smith.hs @@ -0,0 +1,952 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Moonlight.LinAlg.Internal.Backend.Smith + ( SmithNormalForm (..), + SmithDiagonalForm (..), + smithNormalFormPure, + smithDiagonalFormPure, + ) +where + +import Control.Monad (foldM) +import Data.Bifunctor (first) +import Data.Kind (Type) +import Data.List (minimumBy) +import Data.Maybe (mapMaybe) +import Data.Ord (comparing) +import Data.Vector qualified as Box +import GHC.TypeNats (KnownNat, Nat) +import Moonlight.Algebra.Pure.Ring + ( EuclideanDomain (..), + GCDDomain (..), + IntegralDomain (..), + mkNonZeroDivisor, + ) +import Moonlight.Core + ( AdditiveGroup (..), + AdditiveMonoid (..), + MoonlightError (..), + MultiplicativeMonoid (..), + checkedNonNegativeProduct, + ) +import Moonlight.LinAlg.Internal.Backend.RowOps (identityRows) +import Moonlight.LinAlg.Internal.Backend.RowStore + ( RowStore, + columnStore, + replaceColumnStore, + replaceRowStore, + rowStoreFlatten, + rowStoreFromRows, + rowStoreRowAt, + rowStoreShape, + rowStoreValueAt, + swapColumnsStore, + swapRowsStore, + ) +import Moonlight.LinAlg.Internal.Primitives + ( ColumnIndex, + RowIndex, + columnIndexInt, + columnIndices, + mkColumnIndex, + mkRowIndex, + rowIndexInt, + rowIndices, + ) +import Moonlight.LinAlg.Pure.Dense.Types + ( Matrix, + fromListMatrix, + ) +import qualified Moonlight.LinAlg.Pure.Dense.Types as DenseTypes +import Prelude + +type SmithNormalForm :: Nat -> Nat -> Type -> Type +data SmithNormalForm r c a = SmithNormalForm + { smithLeft :: Matrix r r a, + smithDiagonal :: Matrix r c a, + smithRight :: Matrix c c a, + smithLeftInverse :: Matrix r r a, + smithRightInverse :: Matrix c c a + } + +type SmithDiagonalForm :: Nat -> Nat -> Type -> Type +newtype SmithDiagonalForm r c a = SmithDiagonalForm + { smithDiagonalMatrix :: Matrix r c a + } + +type SmithWitnessState :: Type -> Type +data SmithWitnessState a = SmithWitnessState + { smithWitnessLeft :: RowStore a, + smithWitnessRight :: RowStore a, + smithWitnessLeftInverse :: RowStore a, + smithWitnessRightInverse :: RowStore a + } + deriving stock (Eq) + +type SmithState :: Type -> Type +data SmithState a = SmithState + { smithStateMatrix :: RowStore a, + smithStateWitness :: Maybe (SmithWitnessState a) + } + deriving stock (Eq) + +matrixValueAt :: RowIndex -> ColumnIndex -> RowStore a -> Either MoonlightError a +matrixValueAt rowIndex columnIndex = + rowStoreValueAt + (InvariantViolation ("Smith normal form entry lookup failed at " <> show (rowIndex, columnIndex))) + rowIndex + columnIndex + +rowIndexAsColumn :: Int -> RowIndex -> Either MoonlightError ColumnIndex +rowIndexAsColumn columnCount rowIndex = + mkColumnIndex + (InvariantViolation ("Smith inverse witness row/column conversion failed at row " <> show rowIndex)) + columnCount + (rowIndexInt rowIndex) + +columnIndexAsRow :: Int -> ColumnIndex -> Either MoonlightError RowIndex +columnIndexAsRow rowCount columnIndex = + mkRowIndex + (InvariantViolation ("Smith inverse witness column/row conversion failed at column " <> show columnIndex)) + rowCount + (columnIndexInt columnIndex) + +swapRowsState :: RowIndex -> RowIndex -> SmithState a -> Either MoonlightError (SmithState a) +swapRowsState leftIndex rightIndex stateValue = do + swappedMatrix <- + swapRowsStore + (InvariantViolation ("Smith row swap failed at " <> show (leftIndex, rightIndex))) + leftIndex + rightIndex + (smithStateMatrix stateValue) + swappedWitness <- traverse swapWitness (smithStateWitness stateValue) + Right + stateValue + { smithStateMatrix = swappedMatrix, + smithStateWitness = swappedWitness + } + where + swapWitness witnessValue = do + swappedLeft <- + swapRowsStore + (InvariantViolation ("Smith left witness row swap failed at " <> show (leftIndex, rightIndex))) + leftIndex + rightIndex + (smithWitnessLeft witnessValue) + leftColumn <- rowIndexAsColumn (fst (rowStoreShape (smithWitnessLeftInverse witnessValue))) leftIndex + rightColumn <- rowIndexAsColumn (fst (rowStoreShape (smithWitnessLeftInverse witnessValue))) rightIndex + swappedLeftInverse <- + swapColumnsStore + (InvariantViolation ("Smith left inverse column swap failed at " <> show (leftColumn, rightColumn))) + leftColumn + rightColumn + (smithWitnessLeftInverse witnessValue) + Right + witnessValue + { smithWitnessLeft = swappedLeft, + smithWitnessLeftInverse = swappedLeftInverse + } + +swapColsState :: ColumnIndex -> ColumnIndex -> SmithState a -> Either MoonlightError (SmithState a) +swapColsState leftIndex rightIndex stateValue = do + swappedMatrix <- + swapColumnsStore + (InvariantViolation ("Smith column swap failed at " <> show (leftIndex, rightIndex))) + leftIndex + rightIndex + (smithStateMatrix stateValue) + swappedWitness <- traverse swapWitness (smithStateWitness stateValue) + Right + stateValue + { smithStateMatrix = swappedMatrix, + smithStateWitness = swappedWitness + } + where + swapWitness witnessValue = do + swappedRight <- + swapColumnsStore + (InvariantViolation ("Smith right witness column swap failed at " <> show (leftIndex, rightIndex))) + leftIndex + rightIndex + (smithWitnessRight witnessValue) + leftRow <- columnIndexAsRow (fst (rowStoreShape (smithWitnessRightInverse witnessValue))) leftIndex + rightRow <- columnIndexAsRow (fst (rowStoreShape (smithWitnessRightInverse witnessValue))) rightIndex + swappedRightInverse <- + swapRowsStore + (InvariantViolation ("Smith right inverse row swap failed at " <> show (leftRow, rightRow))) + leftRow + rightRow + (smithWitnessRightInverse witnessValue) + Right + witnessValue + { smithWitnessRight = swappedRight, + smithWitnessRightInverse = swappedRightInverse + } + +rowLinearCombination :: + (AdditiveGroup a, MultiplicativeMonoid a) => + a -> + Box.Vector a -> + a -> + Box.Vector a -> + Either MoonlightError (Box.Vector a) +rowLinearCombination leftCoefficient leftRow rightCoefficient rightRow = + if Box.length leftRow == Box.length rightRow + then + Right + ( Box.zipWith + (\leftEntry rightEntry -> (leftCoefficient `mul` leftEntry) `add` (rightCoefficient `mul` rightEntry)) + leftRow + rightRow + ) + else Left (InvariantViolation "Smith row combination length mismatch") + +rowCombine :: + (AdditiveGroup a, MultiplicativeMonoid a) => + Box.Vector a -> + Box.Vector a -> + a -> + Either MoonlightError (Box.Vector a) +rowCombine targetRow sourceRow coefficient = + rowLinearCombination one targetRow (neg coefficient) sourceRow + +rowAddScaled :: + (AdditiveGroup a, MultiplicativeMonoid a) => + Box.Vector a -> + Box.Vector a -> + a -> + Either MoonlightError (Box.Vector a) +rowAddScaled targetRow sourceRow coefficient = + rowLinearCombination one targetRow coefficient sourceRow + +replaceRowPair :: + MoonlightError -> + RowIndex -> + Box.Vector a -> + RowIndex -> + Box.Vector a -> + RowStore a -> + Either MoonlightError (RowStore a) +replaceRowPair failure leftIndex leftRow rightIndex rightRow rows = + replaceRowStore failure leftIndex leftRow rows + >>= replaceRowStore failure rightIndex rightRow + +replaceColumnPair :: + MoonlightError -> + ColumnIndex -> + Box.Vector a -> + ColumnIndex -> + Box.Vector a -> + RowStore a -> + Either MoonlightError (RowStore a) +replaceColumnPair failure leftIndex leftColumn rightIndex rightColumn rows = + replaceColumnStore failure leftIndex leftColumn rows + >>= replaceColumnStore failure rightIndex rightColumn + +rowPairTransform :: + (AdditiveGroup a, MultiplicativeMonoid a) => + MoonlightError -> + RowIndex -> + RowIndex -> + a -> + a -> + a -> + a -> + RowStore a -> + Either MoonlightError (RowStore a) +rowPairTransform failure leftIndex rightIndex aa ab ba bb rows = do + leftRow <- rowStoreRowAt failure leftIndex rows + rightRow <- rowStoreRowAt failure rightIndex rows + transformedLeft <- rowLinearCombination aa leftRow ab rightRow + transformedRight <- rowLinearCombination ba leftRow bb rightRow + replaceRowPair failure leftIndex transformedLeft rightIndex transformedRight rows + +columnPairTransform :: + (AdditiveGroup a, MultiplicativeMonoid a) => + MoonlightError -> + ColumnIndex -> + ColumnIndex -> + a -> + a -> + a -> + a -> + RowStore a -> + Either MoonlightError (RowStore a) +columnPairTransform failure leftIndex rightIndex aa ab ba bb rows = do + leftColumn <- columnStore failure leftIndex rows + rightColumn <- columnStore failure rightIndex rows + transformedLeft <- rowLinearCombination aa leftColumn ab rightColumn + transformedRight <- rowLinearCombination ba leftColumn bb rightColumn + replaceColumnPair failure leftIndex transformedLeft rightIndex transformedRight rows + +columnAddScaledInRows :: + (AdditiveGroup a, MultiplicativeMonoid a) => + ColumnIndex -> + ColumnIndex -> + a -> + RowStore a -> + Either MoonlightError (RowStore a) +columnAddScaledInRows targetIndex sourceIndex coefficient rows = do + targetColumn <- columnStore (InvariantViolation ("Smith column lookup failed at column " <> show targetIndex)) targetIndex rows + sourceColumn <- columnStore (InvariantViolation ("Smith column lookup failed at column " <> show sourceIndex)) sourceIndex rows + updatedColumn <- rowAddScaled targetColumn sourceColumn coefficient + replaceColumnStore + (InvariantViolation ("Smith column replacement failed at column " <> show targetIndex)) + targetIndex + updatedColumn + rows + +rowAddScaledInRows :: + (AdditiveGroup a, MultiplicativeMonoid a) => + RowIndex -> + RowIndex -> + a -> + RowStore a -> + Either MoonlightError (RowStore a) +rowAddScaledInRows targetIndex sourceIndex coefficient rows = do + targetRow <- + rowStoreRowAt + (InvariantViolation ("Smith inverse witness target row missing at index " <> show targetIndex)) + targetIndex + rows + sourceRow <- + rowStoreRowAt + (InvariantViolation ("Smith inverse witness source row missing at index " <> show sourceIndex)) + sourceIndex + rows + updatedRow <- rowAddScaled targetRow sourceRow coefficient + replaceRowStore + (InvariantViolation ("Smith inverse witness row replacement failed at index " <> show targetIndex)) + targetIndex + updatedRow + rows + +rowCombineState :: + (AdditiveGroup a, MultiplicativeMonoid a) => + RowIndex -> + RowIndex -> + a -> + SmithState a -> + Either MoonlightError (SmithState a) +rowCombineState targetIndex sourceIndex coefficient stateValue = do + sourceMatrixRow <- + rowStoreRowAt + (InvariantViolation ("Smith row-combine source matrix row missing at index " <> show sourceIndex)) + sourceIndex + (smithStateMatrix stateValue) + targetMatrixRow <- + rowStoreRowAt + (InvariantViolation ("Smith row-combine target matrix row missing at index " <> show targetIndex)) + targetIndex + (smithStateMatrix stateValue) + updatedMatrixRow <- rowCombine targetMatrixRow sourceMatrixRow coefficient + updatedMatrixRows <- + replaceRowStore + (InvariantViolation ("Smith row-combine matrix replacement failed at index " <> show targetIndex)) + targetIndex + updatedMatrixRow + (smithStateMatrix stateValue) + updatedWitness <- traverse updateWitness (smithStateWitness stateValue) + Right + stateValue + { smithStateMatrix = updatedMatrixRows, + smithStateWitness = updatedWitness + } + where + updateWitness witnessValue = do + sourceLeftRow <- + rowStoreRowAt + (InvariantViolation ("Smith row-combine source witness row missing at index " <> show sourceIndex)) + sourceIndex + (smithWitnessLeft witnessValue) + targetLeftRow <- + rowStoreRowAt + (InvariantViolation ("Smith row-combine target witness row missing at index " <> show targetIndex)) + targetIndex + (smithWitnessLeft witnessValue) + updatedLeftRow <- rowCombine targetLeftRow sourceLeftRow coefficient + updatedLeftRows <- + replaceRowStore + (InvariantViolation ("Smith row-combine witness replacement failed at index " <> show targetIndex)) + targetIndex + updatedLeftRow + (smithWitnessLeft witnessValue) + sourceColumn <- rowIndexAsColumn (fst (rowStoreShape (smithWitnessLeftInverse witnessValue))) sourceIndex + targetColumn <- rowIndexAsColumn (fst (rowStoreShape (smithWitnessLeftInverse witnessValue))) targetIndex + updatedLeftInverseRows <- + columnAddScaledInRows + sourceColumn + targetColumn + coefficient + (smithWitnessLeftInverse witnessValue) + Right + witnessValue + { smithWitnessLeft = updatedLeftRows, + smithWitnessLeftInverse = updatedLeftInverseRows + } + +colCombineState :: + (AdditiveGroup a, MultiplicativeMonoid a) => + ColumnIndex -> + ColumnIndex -> + a -> + SmithState a -> + Either MoonlightError (SmithState a) +colCombineState targetIndex sourceIndex coefficient stateValue = do + matrixTarget <- columnStore (InvariantViolation ("Smith column lookup failed at column " <> show targetIndex)) targetIndex (smithStateMatrix stateValue) + matrixSource <- columnStore (InvariantViolation ("Smith column lookup failed at column " <> show sourceIndex)) sourceIndex (smithStateMatrix stateValue) + updatedMatrixColumn <- rowCombine matrixTarget matrixSource coefficient + updatedMatrix <- + replaceColumnStore + (InvariantViolation ("Smith column replacement failed at column " <> show targetIndex)) + targetIndex + updatedMatrixColumn + (smithStateMatrix stateValue) + updatedWitness <- traverse updateWitness (smithStateWitness stateValue) + Right + stateValue + { smithStateMatrix = updatedMatrix, + smithStateWitness = updatedWitness + } + where + updateWitness witnessValue = do + rightTarget <- columnStore (InvariantViolation ("Smith right witness column lookup failed at column " <> show targetIndex)) targetIndex (smithWitnessRight witnessValue) + rightSource <- columnStore (InvariantViolation ("Smith right witness column lookup failed at column " <> show sourceIndex)) sourceIndex (smithWitnessRight witnessValue) + updatedRightColumn <- rowCombine rightTarget rightSource coefficient + updatedRight <- + replaceColumnStore + (InvariantViolation ("Smith right witness column replacement failed at column " <> show targetIndex)) + targetIndex + updatedRightColumn + (smithWitnessRight witnessValue) + sourceRow <- columnIndexAsRow (fst (rowStoreShape (smithWitnessRightInverse witnessValue))) sourceIndex + targetRow <- columnIndexAsRow (fst (rowStoreShape (smithWitnessRightInverse witnessValue))) targetIndex + updatedRightInverseRows <- + rowAddScaledInRows + sourceRow + targetRow + coefficient + (smithWitnessRightInverse witnessValue) + Right + witnessValue + { smithWitnessRight = updatedRight, + smithWitnessRightInverse = updatedRightInverseRows + } + +exactQuotient :: EuclideanDomain a => String -> a -> a -> Either MoonlightError a +exactQuotient context numerator denominator = do + (quotientValue, remainderValue) <- divideWithRemainderChecked context numerator denominator + if isZero remainderValue + then Right quotientValue + else Left (InvariantViolation ("Smith exact quotient had nonzero remainder during " <> context)) + +divideWithRemainderChecked :: EuclideanDomain a => String -> a -> a -> Either MoonlightError (a, a) +divideWithRemainderChecked context numerator denominator = + case mkNonZeroDivisor denominator of + Nothing -> Left (InvariantViolation ("Smith division received a zero divisor during " <> context)) + Just divisor -> Right (divideWithRemainder numerator divisor) + +dividesNonZero :: EuclideanDomain a => a -> a -> Bool +dividesNonZero denominator numerator = + case mkNonZeroDivisor denominator of + Nothing -> False + Just divisor -> isZero (snd (divideWithRemainder numerator divisor)) + +gcdCombineRowsState :: + EuclideanDomain a => + RowIndex -> + RowIndex -> + ColumnIndex -> + SmithState a -> + Either MoonlightError (SmithState a) +gcdCombineRowsState pivotRow candidateRow pivotColumn stateValue = do + pivotValue <- matrixValueAt pivotRow pivotColumn (smithStateMatrix stateValue) + entryValue <- matrixValueAt candidateRow pivotColumn (smithStateMatrix stateValue) + let (gcdValue, pivotCoefficient, entryCoefficient) = extendedGcdDomain pivotValue entryValue + pivotQuotient <- exactQuotient "row gcd pivot quotient" pivotValue gcdValue + entryQuotient <- exactQuotient "row gcd entry quotient" entryValue gcdValue + updatedMatrix <- + rowPairTransform + (InvariantViolation ("Smith row gcd transform failed at " <> show (pivotRow, candidateRow))) + pivotRow + candidateRow + pivotCoefficient + entryCoefficient + (neg entryQuotient) + pivotQuotient + (smithStateMatrix stateValue) + updatedWitness <- traverse (updateWitness pivotCoefficient entryCoefficient pivotQuotient entryQuotient) (smithStateWitness stateValue) + Right + stateValue + { smithStateMatrix = updatedMatrix, + smithStateWitness = updatedWitness + } + where + updateWitness pivotCoefficient entryCoefficient pivotQuotient entryQuotient witnessValue = do + updatedLeft <- + rowPairTransform + (InvariantViolation ("Smith left witness row gcd transform failed at " <> show (pivotRow, candidateRow))) + pivotRow + candidateRow + pivotCoefficient + entryCoefficient + (neg entryQuotient) + pivotQuotient + (smithWitnessLeft witnessValue) + pivotColumnWitness <- rowIndexAsColumn (fst (rowStoreShape (smithWitnessLeftInverse witnessValue))) pivotRow + candidateColumnWitness <- rowIndexAsColumn (fst (rowStoreShape (smithWitnessLeftInverse witnessValue))) candidateRow + updatedLeftInverse <- + columnPairTransform + (InvariantViolation ("Smith left inverse row gcd transform failed at " <> show (pivotColumnWitness, candidateColumnWitness))) + pivotColumnWitness + candidateColumnWitness + pivotQuotient + entryQuotient + (neg entryCoefficient) + pivotCoefficient + (smithWitnessLeftInverse witnessValue) + Right + witnessValue + { smithWitnessLeft = updatedLeft, + smithWitnessLeftInverse = updatedLeftInverse + } + +gcdCombineColsState :: + EuclideanDomain a => + RowIndex -> + ColumnIndex -> + ColumnIndex -> + SmithState a -> + Either MoonlightError (SmithState a) +gcdCombineColsState pivotRow pivotColumn candidateColumn stateValue = do + pivotValue <- matrixValueAt pivotRow pivotColumn (smithStateMatrix stateValue) + entryValue <- matrixValueAt pivotRow candidateColumn (smithStateMatrix stateValue) + let (gcdValue, pivotCoefficient, entryCoefficient) = extendedGcdDomain pivotValue entryValue + pivotQuotient <- exactQuotient "column gcd pivot quotient" pivotValue gcdValue + entryQuotient <- exactQuotient "column gcd entry quotient" entryValue gcdValue + updatedMatrix <- + columnPairTransform + (InvariantViolation ("Smith column gcd transform failed at " <> show (pivotColumn, candidateColumn))) + pivotColumn + candidateColumn + pivotCoefficient + entryCoefficient + (neg entryQuotient) + pivotQuotient + (smithStateMatrix stateValue) + updatedWitness <- traverse (updateWitness pivotCoefficient entryCoefficient pivotQuotient entryQuotient) (smithStateWitness stateValue) + Right + stateValue + { smithStateMatrix = updatedMatrix, + smithStateWitness = updatedWitness + } + where + updateWitness pivotCoefficient entryCoefficient pivotQuotient entryQuotient witnessValue = do + updatedRight <- + columnPairTransform + (InvariantViolation ("Smith right witness column gcd transform failed at " <> show (pivotColumn, candidateColumn))) + pivotColumn + candidateColumn + pivotCoefficient + entryCoefficient + (neg entryQuotient) + pivotQuotient + (smithWitnessRight witnessValue) + pivotRowWitness <- columnIndexAsRow (fst (rowStoreShape (smithWitnessRightInverse witnessValue))) pivotColumn + candidateRowWitness <- columnIndexAsRow (fst (rowStoreShape (smithWitnessRightInverse witnessValue))) candidateColumn + updatedRightInverse <- + rowPairTransform + (InvariantViolation ("Smith right inverse column gcd transform failed at " <> show (pivotRowWitness, candidateRowWitness))) + pivotRowWitness + candidateRowWitness + pivotQuotient + entryQuotient + (neg entryCoefficient) + pivotCoefficient + (smithWitnessRightInverse witnessValue) + Right + witnessValue + { smithWitnessRight = updatedRight, + smithWitnessRightInverse = updatedRightInverse + } + +findPivot :: EuclideanDomain a => RowIndex -> ColumnIndex -> RowStore a -> Either MoonlightError (Maybe (RowIndex, ColumnIndex, a)) +findPivot startRow startCol rows = + let (rowCount, columnCount) = rowStoreShape rows + in traverse + ( \rowIndex -> + traverse + ( \columnIndex -> + fmap + (\value -> if isZero value then Nothing else Just (rowIndex, columnIndex, value)) + (matrixValueAt rowIndex columnIndex rows) + ) + (dropWhile (< startCol) (columnIndices columnCount)) + ) + (dropWhile (< startRow) (rowIndices rowCount)) + >>= \candidateRows -> + let candidateTriples = mapMaybe id (concat candidateRows) + in Right + ( if null candidateTriples + then Nothing + else Just (minimumBy (comparing (degree . (\(_, _, value) -> value))) candidateTriples) + ) + +columnCleared :: IntegralDomain a => RowIndex -> ColumnIndex -> RowStore a -> Either MoonlightError Bool +columnCleared pivotRow pivotColumn rows = + fmap + and + ( traverse + ( \rowIndex -> + if rowIndex == pivotRow + then Right True + else fmap isZero (matrixValueAt rowIndex pivotColumn rows) + ) + (rowIndices (fst (rowStoreShape rows))) + ) + +rowCleared :: IntegralDomain a => RowIndex -> ColumnIndex -> RowStore a -> Either MoonlightError Bool +rowCleared pivotRow pivotColumn rows = + rowStoreRowAt + (InvariantViolation ("Smith row-clear pivot row missing at index " <> show pivotRow)) + pivotRow + rows + >>= \pivotRowValues -> + Right + ( all + (\(columnIndex, value) -> columnIndex == columnIndexInt pivotColumn || isZero value) + (zip [0 :: Int ..] (Box.toList pivotRowValues)) + ) + +clearColumn :: + forall a. + EuclideanDomain a => + RowIndex -> + ColumnIndex -> + SmithState a -> + Either MoonlightError (SmithState a) +clearColumn pivotRow pivotColumn stateValue = + let rows = smithStateMatrix stateValue + (rowCount, _) = rowStoreShape rows + in traverse + ( \rowIndex -> + fmap + (\entryValue -> if rowIndex /= pivotRow && not (isZero entryValue) then Just rowIndex else Nothing) + (matrixValueAt rowIndex pivotColumn rows) + ) + (rowIndices rowCount) + >>= \candidateMarks -> + case mapMaybe id candidateMarks of + [] -> Right stateValue + candidateRow : _ -> do + pivotValue <- matrixValueAt pivotRow pivotColumn rows + entryValue <- matrixValueAt candidateRow pivotColumn rows + if isZero pivotValue + then Left (InvariantViolation "Smith normal form pivot became zero during column reduction") + else do + (quotientValue, remainderValue) <- divideWithRemainderChecked "column reduction" entryValue pivotValue + reducedState <- + if isZero remainderValue + then rowCombineState candidateRow pivotRow quotientValue stateValue + else gcdCombineRowsState pivotRow candidateRow pivotColumn stateValue + clearColumn pivotRow pivotColumn reducedState + +clearRow :: + forall a. + EuclideanDomain a => + RowIndex -> + ColumnIndex -> + SmithState a -> + Either MoonlightError (SmithState a) +clearRow pivotRow pivotColumn stateValue = + let rows = smithStateMatrix stateValue + in rowStoreRowAt + (InvariantViolation ("Smith row reduction pivot row missing at index " <> show pivotRow)) + pivotRow + rows + >>= \pivotRowValues -> + case map fst (filter (\(columnIndex, value) -> columnIndex /= pivotColumn && not (isZero value)) (zip (columnIndices (Box.length pivotRowValues)) (Box.toList pivotRowValues))) of + [] -> Right stateValue + candidateCol : _ -> do + pivotValue <- matrixValueAt pivotRow pivotColumn rows + entryValue <- matrixValueAt pivotRow candidateCol rows + if isZero pivotValue + then Left (InvariantViolation "Smith normal form pivot became zero during row reduction") + else do + (quotientValue, remainderValue) <- divideWithRemainderChecked "row reduction" entryValue pivotValue + reducedState <- + if isZero remainderValue + then colCombineState candidateCol pivotColumn quotientValue stateValue + else gcdCombineColsState pivotRow pivotColumn candidateCol stateValue + clearRow pivotRow pivotColumn reducedState + +normalizePivot :: + forall a. + EuclideanDomain a => + RowIndex -> + ColumnIndex -> + Int -> + SmithState a -> + Either MoonlightError (SmithState a) +normalizePivot pivotRow pivotColumn remainingBudget stateValue + | remainingBudget <= 0 = Left (InvariantViolation "Smith normal form normalization exhausted iteration budget") + | otherwise = do + clearedColumn <- columnCleared pivotRow pivotColumn (smithStateMatrix stateValue) + clearedRow <- rowCleared pivotRow pivotColumn (smithStateMatrix stateValue) + if clearedColumn && clearedRow + then Right stateValue + else do + columnReduced <- clearColumn pivotRow pivotColumn stateValue + rowReduced <- clearRow pivotRow pivotColumn columnReduced + if smithStateMatrix rowReduced == smithStateMatrix stateValue + then Left (InvariantViolation "Smith normal form normalization stalled before reaching diagonal form") + else normalizePivot pivotRow pivotColumn (remainingBudget - 1) rowReduced + +smithStep :: + forall a. + EuclideanDomain a => + Int -> + Int -> + Int -> + Int -> + SmithState a -> + Either MoonlightError (SmithState a) +smithStep pivotIndex rowCount columnCount normalizationBudget stateValue + | pivotIndex >= min rowCount columnCount = Right stateValue + | otherwise = do + pivotRowIndex <- + mkRowIndex + (InvariantViolation ("Smith normal form pivot row out of bounds at index " <> show pivotIndex)) + rowCount + pivotIndex + pivotColumnIndex <- + mkColumnIndex + (InvariantViolation ("Smith normal form pivot column out of bounds at index " <> show pivotIndex)) + columnCount + pivotIndex + pivotCandidate <- findPivot pivotRowIndex pivotColumnIndex (smithStateMatrix stateValue) + case pivotCandidate of + Nothing -> Right stateValue + Just (pivotRow, pivotCol, _) -> do + pivotMoved <- swapRowsState pivotRowIndex pivotRow stateValue >>= swapColsState pivotColumnIndex pivotCol + normalized <- normalizePivot pivotRowIndex pivotColumnIndex normalizationBudget pivotMoved + smithStep (pivotIndex + 1) rowCount columnCount normalizationBudget normalized + +enforceDivisibilityChain :: + forall a. + EuclideanDomain a => + Int -> + Int -> + Int -> + Int -> + SmithState a -> + Either MoonlightError (SmithState a) +enforceDivisibilityChain rowCount columnCount normalizationBudget divisibilityBudget stateValue = + go divisibilityBudget stateValue + where + diagSize = min rowCount columnCount + + go remainingBudget currentState + | remainingBudget <= 0 = + case findViolation 0 (smithStateMatrix currentState) of + Nothing -> Right currentState + Just _ -> Left (InvariantViolation "Smith normal form divisibility chain exhausted iteration budget") + | otherwise = + case findViolation 0 (smithStateMatrix currentState) of + Nothing -> Right currentState + Just violationIndex -> do + rowI <- + mkRowIndex + (InvariantViolation ("divisibility chain row index out of bounds at " <> show violationIndex)) + rowCount + violationIndex + rowJ <- + mkRowIndex + (InvariantViolation ("divisibility chain row index out of bounds at " <> show (violationIndex + 1))) + rowCount + (violationIndex + 1) + colI <- + mkColumnIndex + (InvariantViolation ("divisibility chain column index out of bounds at " <> show violationIndex)) + columnCount + violationIndex + colJ <- + mkColumnIndex + (InvariantViolation ("divisibility chain column index out of bounds at " <> show (violationIndex + 1))) + columnCount + (violationIndex + 1) + combined <- rowCombineState rowI rowJ (neg one) currentState + normalizedI <- normalizePivot rowI colI normalizationBudget combined + normalizedJ <- normalizePivot rowJ colJ normalizationBudget normalizedI + go (remainingBudget - 1) normalizedJ + + findViolation idx matrixRows + | idx + 1 >= diagSize = Nothing + | otherwise = + case diagonalPair idx matrixRows of + Left _ -> Nothing + Right (dI, dJ) + | isZero dI -> findViolation (idx + 1) matrixRows + | isZero dJ -> findViolation (idx + 1) matrixRows + | dividesNonZero dI dJ -> findViolation (idx + 1) matrixRows + | otherwise -> Just idx + + diagonalPair idx matrixRows = do + rowI <- mkRowIndex (InvariantViolation "divisibility diagonal lookup") rowCount idx + colI <- mkColumnIndex (InvariantViolation "divisibility diagonal lookup") columnCount idx + rowJ <- mkRowIndex (InvariantViolation "divisibility diagonal lookup") rowCount (idx + 1) + colJ <- mkColumnIndex (InvariantViolation "divisibility diagonal lookup") columnCount (idx + 1) + dI <- matrixValueAt rowI colI matrixRows + dJ <- matrixValueAt rowJ colJ matrixRows + Right (dI, dJ) + +smithStateFromRows :: + [[a]] -> + Maybe (SmithWitnessState a) -> + SmithState a +smithStateFromRows rows witnessValue = + SmithState + { smithStateMatrix = rowStoreFromRows rows, + smithStateWitness = witnessValue + } + +fullWitnessState :: + (AdditiveGroup a, MultiplicativeMonoid a) => + Int -> + Int -> + SmithWitnessState a +fullWitnessState rowCount columnCount = + SmithWitnessState + { smithWitnessLeft = rowStoreFromRows (identityRows rowCount), + smithWitnessRight = rowStoreFromRows (identityRows columnCount), + smithWitnessLeftInverse = rowStoreFromRows (identityRows rowCount), + smithWitnessRightInverse = rowStoreFromRows (identityRows columnCount) + } + +runSmithState :: + forall r c a. + (KnownNat r, KnownNat c, EuclideanDomain a) => + Maybe (SmithWitnessState a) -> + Matrix r c a -> + Either MoonlightError (SmithState a) +runSmithState witnessValue matrixValue = do + let (rowCount, columnCount) = DenseTypes.matrixShape matrixValue + diagonalSize = min rowCount columnCount + matrixCardinality <- checkedSmithProduct "matrix normalization budget" rowCount columnCount + normalizationCardinality <- checkedSmithProduct "matrix normalization budget" matrixCardinality 2 + divisibilityCardinality <- checkedSmithProduct "divisibility-chain budget" diagonalSize diagonalSize + let normalizationBudget = max 1 normalizationCardinality + rows <- DenseTypes.matrixToRows matrixValue + let initialState = smithStateFromRows rows witnessValue + steppedState <- smithStep 0 rowCount columnCount normalizationBudget initialState + repairedState <- enforceDivisibilityChain rowCount columnCount normalizationBudget divisibilityCardinality steppedState + normalizeDiagonalUnits rowCount columnCount repairedState + +checkedSmithProduct :: String -> Int -> Int -> Either MoonlightError Int +checkedSmithProduct context leftFactor rightFactor = + first + (const (InvariantViolation ("Smith " <> context <> " exceeds Int range"))) + (checkedNonNegativeProduct leftFactor rightFactor) + +normalizeDiagonalUnits :: + forall a. + EuclideanDomain a => + Int -> + Int -> + SmithState a -> + Either MoonlightError (SmithState a) +normalizeDiagonalUnits rowCount columnCount stateValue = + foldM normalizeAt stateValue [0 .. min rowCount columnCount - 1] + where + normalizeAt currentState indexValue = do + rowIndex <- mkRowIndex (InvariantViolation ("Smith unit normalization row index failed at " <> show indexValue)) rowCount indexValue + columnIndex <- mkColumnIndex (InvariantViolation ("Smith unit normalization column index failed at " <> show indexValue)) columnCount indexValue + entryValue <- matrixValueAt rowIndex columnIndex (smithStateMatrix currentState) + let canonicalValue = gcdDomain entryValue zero + if canonicalValue == entryValue + then Right currentState + else do + (unitValue, remainderValue) <- divideWithRemainderChecked "unit normalization" entryValue canonicalValue + if remainderValue == zero + then do + inverseUnit <- + case unitInverse unitValue of + Just value -> Right value + Nothing -> Left (InvariantViolation ("Smith unit normalization met a nonunit cofactor at " <> show indexValue)) + scaledMatrix <- scaleRowStore rowIndex inverseUnit (smithStateMatrix currentState) + scaledWitness <- traverse (scaleWitness rowIndex unitValue inverseUnit) (smithStateWitness currentState) + Right + currentState + { smithStateMatrix = scaledMatrix, + smithStateWitness = scaledWitness + } + else Left (InvariantViolation ("Smith unit normalization division was inexact at " <> show indexValue)) + + scaleWitness :: RowIndex -> a -> a -> SmithWitnessState a -> Either MoonlightError (SmithWitnessState a) + scaleWitness rowIndex unitValue inverseUnit witnessValue = do + scaledLeft <- scaleRowStore rowIndex inverseUnit (smithWitnessLeft witnessValue) + witnessColumn <- rowIndexAsColumn (fst (rowStoreShape (smithWitnessLeftInverse witnessValue))) rowIndex + scaledLeftInverse <- scaleColumnStore witnessColumn unitValue (smithWitnessLeftInverse witnessValue) + Right + witnessValue + { smithWitnessLeft = scaledLeft, + smithWitnessLeftInverse = scaledLeftInverse + } + +scaleRowStore :: + MultiplicativeMonoid a => + RowIndex -> + a -> + RowStore a -> + Either MoonlightError (RowStore a) +scaleRowStore rowIndex factor store = do + let failure = InvariantViolation ("Smith unit normalization row scale failed at " <> show rowIndex) + rowValue <- rowStoreRowAt failure rowIndex store + replaceRowStore failure rowIndex (Box.map (mul factor) rowValue) store + +scaleColumnStore :: + MultiplicativeMonoid a => + ColumnIndex -> + a -> + RowStore a -> + Either MoonlightError (RowStore a) +scaleColumnStore columnIndex factor store = do + let failure = InvariantViolation ("Smith unit normalization column scale failed at " <> show columnIndex) + columnValue <- columnStore failure columnIndex store + replaceColumnStore failure columnIndex (Box.map (mul factor) columnValue) store + +smithNormalFormPure :: + forall r c a. + (KnownNat r, KnownNat c, EuclideanDomain a) => + Matrix r c a -> + Either MoonlightError (SmithNormalForm r c a) +smithNormalFormPure matrixValue = do + let (rowCount, columnCount) = DenseTypes.matrixShape matrixValue + finalState <- runSmithState (Just (fullWitnessState rowCount columnCount)) matrixValue + witnessValue <- + case smithStateWitness finalState of + Just value -> Right value + Nothing -> Left (InvariantViolation "Smith full decomposition lost witness state") + leftMatrix <- fromListMatrix @r @r (rowStoreFlatten (smithWitnessLeft witnessValue)) + diagonalMatrix <- fromListMatrix @r @c (rowStoreFlatten (smithStateMatrix finalState)) + rightMatrix <- fromListMatrix @c @c (rowStoreFlatten (smithWitnessRight witnessValue)) + leftInverseMatrix <- fromListMatrix @r @r (rowStoreFlatten (smithWitnessLeftInverse witnessValue)) + rightInverseMatrix <- fromListMatrix @c @c (rowStoreFlatten (smithWitnessRightInverse witnessValue)) + pure + SmithNormalForm + { smithLeft = leftMatrix, + smithDiagonal = diagonalMatrix, + smithRight = rightMatrix, + smithLeftInverse = leftInverseMatrix, + smithRightInverse = rightInverseMatrix + } + +smithDiagonalFormPure :: + forall r c a. + (KnownNat r, KnownNat c, EuclideanDomain a) => + Matrix r c a -> + Either MoonlightError (SmithDiagonalForm r c a) +smithDiagonalFormPure matrixValue = do + finalState <- runSmithState Nothing matrixValue + diagonalMatrix <- fromListMatrix @r @c (rowStoreFlatten (smithStateMatrix finalState)) + pure (SmithDiagonalForm diagonalMatrix) diff --git a/src-dense/Moonlight/LinAlg/Internal/Dense/DoubleFactorization.hs b/src-dense/Moonlight/LinAlg/Internal/Dense/DoubleFactorization.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Internal/Dense/DoubleFactorization.hs @@ -0,0 +1,622 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Internal.Dense.DoubleFactorization + ( choleskyLower, + qrFullColumnRank, + solveSquareLinearSystem, + ) +where + +import Control.Monad.ST (ST, runST) +import Data.Bifunctor (first) +import Data.Primitive.PrimArray qualified as PrimArray +import Moonlight.Core + ( MoonlightError (..), + checkedNonNegativeProduct, + fieldValueValid, + ) +import Moonlight.LinAlg.Internal.Eigen.Kernels (hypotStable) +import Moonlight.LinAlg.Internal.Primitives (epsilon) +import Prelude + +solveSquareLinearSystem :: Int -> [Double] -> [Double] -> Either MoonlightError [Double] +solveSquareLinearSystem !matrixSize matrixValues rightHandSideValues = do + matrixEntryCount <- requireSquarePayload "direct solve" matrixSize matrixValues + requireVectorPayload "direct solve" matrixSize rightHandSideValues + runST $ do + matrixWork <- newDoubleArray matrixEntryCount + rhsWork <- newDoubleArray matrixSize + pivots <- newIntArray matrixSize + copyDoubleList matrixWork matrixValues + copyDoubleList rhsWork rightHandSideValues + factorResult <- factorPLU matrixSize matrixWork pivots + case factorResult of + Left err -> pure (Left err) + Right () -> do + applyPivotVector matrixSize pivots rhsWork + forwardSolveUnitLower matrixSize matrixWork rhsWork + backResult <- backwardSolveUpper matrixSize matrixWork rhsWork + case backResult of + Left err -> pure (Left err) + Right () -> Right <$> freezeDoubleList rhsWork + +qrFullColumnRank :: Int -> Int -> [Double] -> Either MoonlightError ([Double], [Double]) +qrFullColumnRank !rowCount !columnCount matrixValues + | rowCount < columnCount = + Left (InvariantViolation "full-column-rank QR requires row count greater than or equal to column count") + | otherwise = do + matrixEntryCount <- requireMatrixPayload "QR decomposition" rowCount columnCount matrixValues + upperEntryCount <- + first + (const (InvariantViolation "QR upper-factor cardinality exceeds non-negative Int range")) + (checkedNonNegativeProduct columnCount columnCount) + runST $ do + matrixWork <- newDoubleArray matrixEntryCount + reflectorScalars <- newDoubleArray columnCount + copyDoubleList matrixWork matrixValues + setDoubleArray reflectorScalars columnCount 0.0 + factorResult <- factorQR rowCount columnCount matrixWork reflectorScalars + case factorResult of + Left err -> pure (Left err) + Right () -> do + qValues <- formThinQ rowCount columnCount matrixEntryCount matrixWork reflectorScalars + rValues <- extractUpperR rowCount columnCount upperEntryCount matrixWork + pure (Right (qValues, rValues)) + +choleskyLower :: Int -> [Double] -> Either MoonlightError [Double] +choleskyLower !matrixSize matrixValues = do + matrixEntryCount <- requireSquarePayload "Cholesky decomposition" matrixSize matrixValues + runST $ do + matrixWork <- newDoubleArray matrixEntryCount + copyDoubleList matrixWork matrixValues + symmetryResult <- checkSymmetricMatrix matrixSize matrixWork + case symmetryResult of + Left err -> pure (Left err) + Right () -> do + factorResult <- factorCholesky matrixSize matrixWork + case factorResult of + Left err -> pure (Left err) + Right () -> do + zeroStrictUpper matrixSize matrixWork + Right <$> freezeDoubleList matrixWork + +newDoubleArray :: Int -> ST s (PrimArray.MutablePrimArray s Double) +newDoubleArray = PrimArray.newPrimArray +{-# INLINE newDoubleArray #-} + +newIntArray :: Int -> ST s (PrimArray.MutablePrimArray s Int) +newIntArray = PrimArray.newPrimArray +{-# INLINE newIntArray #-} + +setDoubleArray :: PrimArray.MutablePrimArray s Double -> Int -> Double -> ST s () +setDoubleArray !target !entryCount !entryValue = + PrimArray.setPrimArray target 0 entryCount entryValue +{-# INLINE setDoubleArray #-} + +copyDoubleList :: PrimArray.MutablePrimArray s Double -> [Double] -> ST s () +copyDoubleList !target = go 0 + where + go !_ [] = pure () + go !entryIndex (entryValue : restValues) = do + PrimArray.writePrimArray target entryIndex entryValue + go (entryIndex + 1) restValues +{-# INLINE copyDoubleList #-} + +freezeDoubleList :: PrimArray.MutablePrimArray s Double -> ST s [Double] +freezeDoubleList !values = + PrimArray.primArrayToList <$> PrimArray.unsafeFreezePrimArray values +{-# INLINE freezeDoubleList #-} + +rowMajorIndex :: Int -> Int -> Int -> Int +rowMajorIndex !columnCount !rowIndex !columnIndex = + rowIndex * columnCount + columnIndex +{-# INLINE rowMajorIndex #-} + +readMatrix :: Int -> PrimArray.MutablePrimArray s Double -> Int -> Int -> ST s Double +readMatrix !columnCount !matrixValues !rowIndex !columnIndex = + PrimArray.readPrimArray matrixValues (rowMajorIndex columnCount rowIndex columnIndex) +{-# INLINE readMatrix #-} + +writeMatrix :: Int -> PrimArray.MutablePrimArray s Double -> Int -> Int -> Double -> ST s () +writeMatrix !columnCount !matrixValues !rowIndex !columnIndex !entryValue = + PrimArray.writePrimArray matrixValues (rowMajorIndex columnCount rowIndex columnIndex) entryValue +{-# INLINE writeMatrix #-} + +readMatrixWithColumnCount :: Int -> PrimArray.MutablePrimArray s Double -> Int -> Int -> ST s Double +readMatrixWithColumnCount = readMatrix +{-# INLINE readMatrixWithColumnCount #-} + +writeMatrixWithColumnCount :: Int -> PrimArray.MutablePrimArray s Double -> Int -> Int -> Double -> ST s () +writeMatrixWithColumnCount = writeMatrix +{-# INLINE writeMatrixWithColumnCount #-} + +requireSquarePayload :: String -> Int -> [Double] -> Either MoonlightError Int +requireSquarePayload label matrixSize matrixValues = do + expectedLength <- + first + (const (InvariantViolation (label <> " square cardinality exceeds non-negative Int range"))) + (checkedNonNegativeProduct matrixSize matrixSize) + if length matrixValues /= expectedLength + then + Left + ( InvariantViolation + ( label + <> " square payload length mismatch: expected " + <> show expectedLength + <> " values but received " + <> show (length matrixValues) + ) + ) + else requireFiniteEntries label matrixValues *> Right expectedLength + +requireMatrixPayload :: String -> Int -> Int -> [Double] -> Either MoonlightError Int +requireMatrixPayload label rowCount columnCount matrixValues = do + expectedLength <- + first + (const (InvariantViolation (label <> " dense cardinality exceeds non-negative Int range"))) + (checkedNonNegativeProduct rowCount columnCount) + if length matrixValues /= expectedLength + then + Left + ( InvariantViolation + ( label + <> " dense payload length mismatch: expected " + <> show expectedLength + <> " values but received " + <> show (length matrixValues) + ) + ) + else requireFiniteEntries label matrixValues *> Right expectedLength + +requireVectorPayload :: String -> Int -> [Double] -> Either MoonlightError () +requireVectorPayload label vectorSize vectorValues + | vectorSize < 0 = + Left (InvariantViolation (label <> " requires a non-negative vector size")) + | length vectorValues /= vectorSize = + Left + ( InvariantViolation + ( label + <> " vector payload length mismatch: expected " + <> show vectorSize + <> " values but received " + <> show (length vectorValues) + ) + ) + | otherwise = + requireFiniteEntries label vectorValues + +requireFiniteEntries :: String -> [Double] -> Either MoonlightError () +requireFiniteEntries label values = + if all fieldValueValid values + then Right () + else Left (InvariantViolation (label <> " requires finite entries")) + +checkSymmetricMatrix :: Int -> PrimArray.MutablePrimArray s Double -> ST s (Either MoonlightError ()) +checkSymmetricMatrix !matrixSize !matrixWork = goRow 0 + where + tolerance = sqrt epsilon + + goRow !rowIndex + | rowIndex >= matrixSize = pure (Right ()) + | otherwise = goColumn rowIndex (rowIndex + 1) + + goColumn !rowIndex !columnIndex + | columnIndex >= matrixSize = goRow (rowIndex + 1) + | otherwise = do + leftValue <- readMatrix matrixSize matrixWork rowIndex columnIndex + rightValue <- readMatrix matrixSize matrixWork columnIndex rowIndex + if abs (leftValue - rightValue) <= tolerance + then goColumn rowIndex (columnIndex + 1) + else pure (Left (InvariantViolation "Cholesky decomposition requires a symmetric matrix")) +{-# INLINE checkSymmetricMatrix #-} + +factorPLU :: + Int -> + PrimArray.MutablePrimArray s Double -> + PrimArray.MutablePrimArray s Int -> + ST s (Either MoonlightError ()) +factorPLU !matrixSize !matrixWork !pivotRows = go 0 + where + go !pivotIndex + | pivotIndex >= matrixSize = pure (Right ()) + | otherwise = do + (selectedRow, selectedMagnitude) <- findPivotRow matrixSize matrixWork pivotIndex + if selectedMagnitude <= epsilon + then pure (Left (InvariantViolation ("direct solve failed during PLU factorization: non-invertible pivot at column " <> show pivotIndex))) + else do + PrimArray.writePrimArray pivotRows pivotIndex selectedRow + swapMatrixRows matrixSize matrixWork pivotIndex selectedRow + pivotValue <- readMatrix matrixSize matrixWork pivotIndex pivotIndex + eliminatePLUColumn matrixSize matrixWork pivotIndex pivotValue + go (pivotIndex + 1) + +findPivotRow :: Int -> PrimArray.MutablePrimArray s Double -> Int -> ST s (Int, Double) +findPivotRow !matrixSize !matrixWork !columnIndex = do + firstValue <- readMatrix matrixSize matrixWork columnIndex columnIndex + go (columnIndex + 1) columnIndex (abs firstValue) + where + go !rowIndex !bestRow !bestMagnitude + | rowIndex >= matrixSize = pure (bestRow, bestMagnitude) + | otherwise = do + candidateValue <- readMatrix matrixSize matrixWork rowIndex columnIndex + let !candidateMagnitude = abs candidateValue + if candidateMagnitude > bestMagnitude + then go (rowIndex + 1) rowIndex candidateMagnitude + else go (rowIndex + 1) bestRow bestMagnitude +{-# INLINE findPivotRow #-} + +swapMatrixRows :: Int -> PrimArray.MutablePrimArray s Double -> Int -> Int -> ST s () +swapMatrixRows !columnCount !matrixWork !leftRow !rightRow + | leftRow == rightRow = pure () + | otherwise = go 0 + where + go !columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + leftValue <- readMatrix columnCount matrixWork leftRow columnIndex + rightValue <- readMatrix columnCount matrixWork rightRow columnIndex + writeMatrix columnCount matrixWork leftRow columnIndex rightValue + writeMatrix columnCount matrixWork rightRow columnIndex leftValue + go (columnIndex + 1) +{-# INLINE swapMatrixRows #-} + +eliminatePLUColumn :: Int -> PrimArray.MutablePrimArray s Double -> Int -> Double -> ST s () +eliminatePLUColumn !matrixSize !matrixWork !pivotIndex !pivotValue = + goRow (pivotIndex + 1) + where + goRow !rowIndex + | rowIndex >= matrixSize = pure () + | otherwise = do + factorEntry <- readMatrix matrixSize matrixWork rowIndex pivotIndex + let !multiplier = factorEntry / pivotValue + writeMatrix matrixSize matrixWork rowIndex pivotIndex multiplier + updateTrailingRow rowIndex multiplier (pivotIndex + 1) + goRow (rowIndex + 1) + + updateTrailingRow !rowIndex !multiplier !columnIndex + | columnIndex >= matrixSize = pure () + | otherwise = do + currentValue <- readMatrix matrixSize matrixWork rowIndex columnIndex + pivotRowValue <- readMatrix matrixSize matrixWork pivotIndex columnIndex + writeMatrix matrixSize matrixWork rowIndex columnIndex (currentValue - multiplier * pivotRowValue) + updateTrailingRow rowIndex multiplier (columnIndex + 1) +{-# INLINE eliminatePLUColumn #-} + +applyPivotVector :: Int -> PrimArray.MutablePrimArray s Int -> PrimArray.MutablePrimArray s Double -> ST s () +applyPivotVector !matrixSize !pivotRows !rhsWork = go 0 + where + go !pivotIndex + | pivotIndex >= matrixSize = pure () + | otherwise = do + selectedRow <- PrimArray.readPrimArray pivotRows pivotIndex + swapRhsEntries rhsWork pivotIndex selectedRow + go (pivotIndex + 1) +{-# INLINE applyPivotVector #-} + +swapRhsEntries :: PrimArray.MutablePrimArray s Double -> Int -> Int -> ST s () +swapRhsEntries !rhsWork !leftIndex !rightIndex + | leftIndex == rightIndex = pure () + | otherwise = do + leftValue <- PrimArray.readPrimArray rhsWork leftIndex + rightValue <- PrimArray.readPrimArray rhsWork rightIndex + PrimArray.writePrimArray rhsWork leftIndex rightValue + PrimArray.writePrimArray rhsWork rightIndex leftValue +{-# INLINE swapRhsEntries #-} + +forwardSolveUnitLower :: Int -> PrimArray.MutablePrimArray s Double -> PrimArray.MutablePrimArray s Double -> ST s () +forwardSolveUnitLower !matrixSize !matrixWork !rhsWork = goRow 0 + where + goRow !rowIndex + | rowIndex >= matrixSize = pure () + | otherwise = do + contribution <- lowerDot rowIndex 0 0.0 + rhsValue <- PrimArray.readPrimArray rhsWork rowIndex + PrimArray.writePrimArray rhsWork rowIndex (rhsValue - contribution) + goRow (rowIndex + 1) + + lowerDot !rowIndex !columnIndex !accumulator + | columnIndex >= rowIndex = pure accumulator + | otherwise = do + lowerValue <- readMatrix matrixSize matrixWork rowIndex columnIndex + solvedValue <- PrimArray.readPrimArray rhsWork columnIndex + lowerDot rowIndex (columnIndex + 1) (accumulator + lowerValue * solvedValue) +{-# INLINE forwardSolveUnitLower #-} + +backwardSolveUpper :: Int -> PrimArray.MutablePrimArray s Double -> PrimArray.MutablePrimArray s Double -> ST s (Either MoonlightError ()) +backwardSolveUpper !matrixSize !matrixWork !rhsWork = goRow (matrixSize - 1) + where + goRow !rowIndex + | rowIndex < 0 = pure (Right ()) + | otherwise = do + contribution <- upperDot rowIndex (rowIndex + 1) 0.0 + diagonalValue <- readMatrix matrixSize matrixWork rowIndex rowIndex + rhsValue <- PrimArray.readPrimArray rhsWork rowIndex + if abs diagonalValue <= epsilon + then pure (Left (InvariantViolation "direct solve failed during backward substitution: zero diagonal pivot")) + else do + PrimArray.writePrimArray rhsWork rowIndex ((rhsValue - contribution) / diagonalValue) + goRow (rowIndex - 1) + + upperDot !rowIndex !columnIndex !accumulator + | columnIndex >= matrixSize = pure accumulator + | otherwise = do + upperValue <- readMatrix matrixSize matrixWork rowIndex columnIndex + solvedValue <- PrimArray.readPrimArray rhsWork columnIndex + upperDot rowIndex (columnIndex + 1) (accumulator + upperValue * solvedValue) +{-# INLINE backwardSolveUpper #-} + +factorQR :: + Int -> + Int -> + PrimArray.MutablePrimArray s Double -> + PrimArray.MutablePrimArray s Double -> + ST s (Either MoonlightError ()) +factorQR !rowCount !columnCount !matrixWork !reflectorScalars = go 0 + where + go !columnIndex + | columnIndex >= columnCount = pure (Right ()) + | otherwise = do + reflectorResult <- makeHouseholderReflector rowCount columnCount matrixWork reflectorScalars columnIndex + case reflectorResult of + Left err -> pure (Left err) + Right tauValue -> do + applyQRReflectorToRemainder rowCount columnCount matrixWork columnIndex tauValue + go (columnIndex + 1) + +makeHouseholderReflector :: + Int -> + Int -> + PrimArray.MutablePrimArray s Double -> + PrimArray.MutablePrimArray s Double -> + Int -> + ST s (Either MoonlightError Double) +makeHouseholderReflector !rowCount !columnCount !matrixWork !reflectorScalars !columnIndex = do + alphaValue <- readMatrix columnCount matrixWork columnIndex columnIndex + tailNorm <- columnTailNorm rowCount columnCount matrixWork columnIndex + if tailNorm == 0.0 + then + if abs alphaValue <= epsilon + then pure (Left (InvariantViolation "QR decomposition failed: dependent or zero column encountered")) + else do + PrimArray.writePrimArray reflectorScalars columnIndex 0.0 + pure (Right 0.0) + else do + let !normValue = hypotStable alphaValue tailNorm + !betaValue = + if alphaValue < 0.0 || isNegativeZero alphaValue + then normValue + else negate normValue + if abs betaValue <= epsilon + then pure (Left (InvariantViolation "QR decomposition failed: dependent or zero column encountered")) + else do + let !tauValue = (betaValue - alphaValue) / betaValue + !scaleValue = 1.0 / (alphaValue - betaValue) + scaleReflectorTail rowCount columnCount matrixWork columnIndex scaleValue + writeMatrix columnCount matrixWork columnIndex columnIndex betaValue + PrimArray.writePrimArray reflectorScalars columnIndex tauValue + pure (Right tauValue) +{-# INLINE makeHouseholderReflector #-} + +columnTailNorm :: Int -> Int -> PrimArray.MutablePrimArray s Double -> Int -> ST s Double +columnTailNorm !rowCount !columnCount !matrixWork !columnIndex = + go (columnIndex + 1) 0.0 1.0 + where + go !rowIndex !scaleValue !sumSquares + | rowIndex >= rowCount = pure (scaleValue * sqrt sumSquares) + | otherwise = do + entryValue <- readMatrix columnCount matrixWork rowIndex columnIndex + let !entryMagnitude = abs entryValue + if entryMagnitude == 0.0 + then go (rowIndex + 1) scaleValue sumSquares + else + if scaleValue < entryMagnitude + then + let !scaled = scaleValue / entryMagnitude + in go (rowIndex + 1) entryMagnitude (1.0 + sumSquares * scaled * scaled) + else + let !scaled = entryMagnitude / scaleValue + in go (rowIndex + 1) scaleValue (sumSquares + scaled * scaled) +{-# INLINE columnTailNorm #-} + +scaleReflectorTail :: Int -> Int -> PrimArray.MutablePrimArray s Double -> Int -> Double -> ST s () +scaleReflectorTail !rowCount !columnCount !matrixWork !columnIndex !scaleValue = + go (columnIndex + 1) + where + go !rowIndex + | rowIndex >= rowCount = pure () + | otherwise = do + entryValue <- readMatrix columnCount matrixWork rowIndex columnIndex + writeMatrix columnCount matrixWork rowIndex columnIndex (entryValue * scaleValue) + go (rowIndex + 1) +{-# INLINE scaleReflectorTail #-} + +applyQRReflectorToRemainder :: Int -> Int -> PrimArray.MutablePrimArray s Double -> Int -> Double -> ST s () +applyQRReflectorToRemainder !rowCount !columnCount !matrixWork !reflectorIndex !tauValue + | tauValue == 0.0 = pure () + | otherwise = goColumn (reflectorIndex + 1) + where + goColumn !targetColumn + | targetColumn >= columnCount = pure () + | otherwise = do + dotValue <- matrixReflectorDot rowCount columnCount matrixWork reflectorIndex targetColumn + let !scaledDot = tauValue * dotValue + pivotValue <- readMatrix columnCount matrixWork reflectorIndex targetColumn + writeMatrix columnCount matrixWork reflectorIndex targetColumn (pivotValue - scaledDot) + updateTail targetColumn scaledDot (reflectorIndex + 1) + goColumn (targetColumn + 1) + + updateTail !targetColumn !scaledDot !rowIndex + | rowIndex >= rowCount = pure () + | otherwise = do + reflectorEntry <- readMatrix columnCount matrixWork rowIndex reflectorIndex + targetEntry <- readMatrix columnCount matrixWork rowIndex targetColumn + writeMatrix columnCount matrixWork rowIndex targetColumn (targetEntry - reflectorEntry * scaledDot) + updateTail targetColumn scaledDot (rowIndex + 1) +{-# INLINE applyQRReflectorToRemainder #-} + +matrixReflectorDot :: Int -> Int -> PrimArray.MutablePrimArray s Double -> Int -> Int -> ST s Double +matrixReflectorDot !rowCount !columnCount !matrixWork !reflectorIndex !targetColumn = do + pivotValue <- readMatrix columnCount matrixWork reflectorIndex targetColumn + go (reflectorIndex + 1) pivotValue + where + go !rowIndex !accumulator + | rowIndex >= rowCount = pure accumulator + | otherwise = do + reflectorEntry <- readMatrix columnCount matrixWork rowIndex reflectorIndex + targetEntry <- readMatrix columnCount matrixWork rowIndex targetColumn + go (rowIndex + 1) (accumulator + reflectorEntry * targetEntry) +{-# INLINE matrixReflectorDot #-} + +formThinQ :: + Int -> + Int -> + Int -> + PrimArray.MutablePrimArray s Double -> + PrimArray.MutablePrimArray s Double -> + ST s [Double] +formThinQ !rowCount !columnCount !matrixEntryCount !matrixWork !reflectorScalars = do + qWork <- newDoubleArray matrixEntryCount + setDoubleArray qWork matrixEntryCount 0.0 + setThinIdentity rowCount columnCount qWork + applyReflectors (columnCount - 1) qWork + freezeDoubleList qWork + where + applyReflectors !reflectorIndex !qWork + | reflectorIndex < 0 = pure () + | otherwise = do + tauValue <- PrimArray.readPrimArray reflectorScalars reflectorIndex + applyQRReflectorToQ rowCount columnCount matrixWork qWork reflectorIndex tauValue + applyReflectors (reflectorIndex - 1) qWork +{-# INLINE formThinQ #-} + +setThinIdentity :: Int -> Int -> PrimArray.MutablePrimArray s Double -> ST s () +setThinIdentity !rowCount !columnCount !qWork = + go 0 + where + diagonalCount = min rowCount columnCount + + go !diagonalIndex + | diagonalIndex >= diagonalCount = pure () + | otherwise = do + writeMatrixWithColumnCount columnCount qWork diagonalIndex diagonalIndex 1.0 + go (diagonalIndex + 1) +{-# INLINE setThinIdentity #-} + +applyQRReflectorToQ :: + Int -> + Int -> + PrimArray.MutablePrimArray s Double -> + PrimArray.MutablePrimArray s Double -> + Int -> + Double -> + ST s () +applyQRReflectorToQ !rowCount !columnCount !matrixWork !qWork !reflectorIndex !tauValue + | tauValue == 0.0 = pure () + | otherwise = goColumn reflectorIndex + where + goColumn !targetColumn + | targetColumn >= columnCount = pure () + | otherwise = do + dotValue <- qReflectorDot targetColumn reflectorIndex 0.0 + let !scaledDot = tauValue * dotValue + pivotValue <- readMatrixWithColumnCount columnCount qWork reflectorIndex targetColumn + writeMatrixWithColumnCount columnCount qWork reflectorIndex targetColumn (pivotValue - scaledDot) + updateTail targetColumn scaledDot (reflectorIndex + 1) + goColumn (targetColumn + 1) + + qReflectorDot !targetColumn !rowIndex !accumulator + | rowIndex >= rowCount = pure accumulator + | rowIndex == reflectorIndex = do + qEntry <- readMatrixWithColumnCount columnCount qWork rowIndex targetColumn + qReflectorDot targetColumn (rowIndex + 1) (accumulator + qEntry) + | otherwise = do + reflectorEntry <- readMatrix columnCount matrixWork rowIndex reflectorIndex + qEntry <- readMatrixWithColumnCount columnCount qWork rowIndex targetColumn + qReflectorDot targetColumn (rowIndex + 1) (accumulator + reflectorEntry * qEntry) + + updateTail !targetColumn !scaledDot !rowIndex + | rowIndex >= rowCount = pure () + | otherwise = do + reflectorEntry <- readMatrix columnCount matrixWork rowIndex reflectorIndex + qEntry <- readMatrixWithColumnCount columnCount qWork rowIndex targetColumn + writeMatrixWithColumnCount columnCount qWork rowIndex targetColumn (qEntry - reflectorEntry * scaledDot) + updateTail targetColumn scaledDot (rowIndex + 1) +{-# INLINE applyQRReflectorToQ #-} + +extractUpperR :: Int -> Int -> Int -> PrimArray.MutablePrimArray s Double -> ST s [Double] +extractUpperR !rowCount !columnCount !upperEntryCount !matrixWork = do + rWork <- newDoubleArray upperEntryCount + setDoubleArray rWork upperEntryCount 0.0 + goRow 0 rWork + freezeDoubleList rWork + where + goRow !rowIndex !rWork + | rowIndex >= columnCount = pure () + | otherwise = do + goColumn rowIndex rowIndex rWork + goRow (rowIndex + 1) rWork + + goColumn !rowIndex !columnIndex !rWork + | columnIndex >= columnCount = pure () + | rowIndex >= rowCount = pure () + | otherwise = do + entryValue <- readMatrix columnCount matrixWork rowIndex columnIndex + writeMatrix columnCount rWork rowIndex columnIndex entryValue + goColumn rowIndex (columnIndex + 1) rWork +{-# INLINE extractUpperR #-} + +factorCholesky :: Int -> PrimArray.MutablePrimArray s Double -> ST s (Either MoonlightError ()) +factorCholesky !matrixSize !matrixWork = goColumn 0 + where + goColumn !columnIndex + | columnIndex >= matrixSize = pure (Right ()) + | otherwise = do + diagonalContribution <- lowerSelfDot columnIndex 0 0.0 + diagonalInput <- readMatrix matrixSize matrixWork columnIndex columnIndex + let !diagonalResidual = diagonalInput - diagonalContribution + if diagonalResidual <= 0.0 || not (fieldValueValid diagonalResidual) + then pure (Left (InvariantViolation "Cholesky decomposition failed: matrix is not positive-definite")) + else do + let !diagonalValue = sqrt diagonalResidual + writeMatrix matrixSize matrixWork columnIndex columnIndex diagonalValue + updateColumnTail columnIndex diagonalValue (columnIndex + 1) + goColumn (columnIndex + 1) + + lowerSelfDot !rowIndex !columnIndex !accumulator + | columnIndex >= rowIndex = pure accumulator + | otherwise = do + lowerValue <- readMatrix matrixSize matrixWork rowIndex columnIndex + lowerSelfDot rowIndex (columnIndex + 1) (accumulator + lowerValue * lowerValue) + + lowerCrossDot !leftRow !rightRow !columnIndex !accumulator + | columnIndex >= rightRow = pure accumulator + | otherwise = do + leftValue <- readMatrix matrixSize matrixWork leftRow columnIndex + rightValue <- readMatrix matrixSize matrixWork rightRow columnIndex + lowerCrossDot leftRow rightRow (columnIndex + 1) (accumulator + leftValue * rightValue) + + updateColumnTail !columnIndex !diagonalValue !rowIndex + | rowIndex >= matrixSize = pure () + | otherwise = do + crossContribution <- lowerCrossDot rowIndex columnIndex 0 0.0 + inputValue <- readMatrix matrixSize matrixWork rowIndex columnIndex + writeMatrix matrixSize matrixWork rowIndex columnIndex ((inputValue - crossContribution) / diagonalValue) + updateColumnTail columnIndex diagonalValue (rowIndex + 1) +{-# INLINE factorCholesky #-} + +zeroStrictUpper :: Int -> PrimArray.MutablePrimArray s Double -> ST s () +zeroStrictUpper !matrixSize !matrixWork = goRow 0 + where + goRow !rowIndex + | rowIndex >= matrixSize = pure () + | otherwise = do + goColumn rowIndex (rowIndex + 1) + goRow (rowIndex + 1) + + goColumn !rowIndex !columnIndex + | columnIndex >= matrixSize = pure () + | otherwise = do + writeMatrix matrixSize matrixWork rowIndex columnIndex 0.0 + goColumn rowIndex (columnIndex + 1) +{-# INLINE zeroStrictUpper #-} diff --git a/src-dense/Moonlight/LinAlg/Internal/Dense/OneSidedJacobiSVD.hs b/src-dense/Moonlight/LinAlg/Internal/Dense/OneSidedJacobiSVD.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Internal/Dense/OneSidedJacobiSVD.hs @@ -0,0 +1,287 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE StrictData #-} + +-- | One-sided Jacobi SVD over sealed flat Double workspaces. +module Moonlight.LinAlg.Internal.Dense.OneSidedJacobiSVD + ( ThinSvdFailure (..), + ThinSvdResult (..), + thinSvdFullColumnRank, + ) +where + +import Control.Monad.ST (ST, runST) +import Data.Bifunctor (first) +import Data.List (sortBy) +import Data.Ord (comparing) +import Data.Vector.Storable qualified as S +import Data.Vector.Storable.Mutable qualified as SM +import Moonlight.Core (checkedNonNegativeProduct) +import Moonlight.LinAlg.Internal.Eigen.DenseWork + ( MutableDenseWork (..), + dotDenseColumns, + newDenseWork, + readDenseWork, + setIdentityDenseWork, + writeDenseWork, + ) +import Moonlight.LinAlg.Internal.Eigen.Kernels + ( epsDouble, + finiteDouble, + forIndex, + hypotStable, + ) +import Moonlight.LinAlg.Pure.Dense.Flat + ( DenseDoubleMatrix, + denseDoubleMatrixShape, + denseDoubleMatrixToRowMajorVector, + trustedDenseDoubleMatrixRowMajor, + ) +import Prelude + +data ThinSvdFailure + = ThinSvdNonFiniteInput + | ThinSvdDimensionViolation !String + | ThinSvdRankDeficient !Int !Double + | ThinSvdSweepBudgetNonConvergence !Int !Double + deriving stock (Eq, Show) + +data ThinSvdResult = ThinSvdResult + { thinSvdLeftSingularVectors :: !DenseDoubleMatrix, + thinSvdSingularValues :: !(S.Vector Double), + thinSvdRightSingularVectorsTransposed :: !DenseDoubleMatrix + } + deriving stock (Eq, Show) + +thinSvdFullColumnRank :: DenseDoubleMatrix -> Either ThinSvdFailure ThinSvdResult +thinSvdFullColumnRank matrixValue + | rowCount < columnCount = + Left (ThinSvdDimensionViolation "thin Jacobi SVD requires row count greater than or equal to column count") + | columnCount > 64 = + Left (ThinSvdDimensionViolation "thin Jacobi SVD supports at most 64 columns") + | S.any (not . finiteDouble) payload = + Left ThinSvdNonFiniteInput + | otherwise = do + matrixEntryCount <- checkedThinSvdCardinality "thin Jacobi SVD matrix" rowCount columnCount + rightEntryCount <- checkedThinSvdCardinality "thin Jacobi SVD right singular vectors" columnCount columnCount + if S.length payload /= matrixEntryCount + then Left (ThinSvdDimensionViolation "thin Jacobi SVD payload does not match its declared matrix shape") + else runST (thinSvdFullColumnRankST rowCount columnCount matrixEntryCount rightEntryCount matrixValue) + where + !(rowCount, columnCount) = denseDoubleMatrixShape matrixValue + payload = denseDoubleMatrixToRowMajorVector matrixValue + +checkedThinSvdCardinality :: String -> Int -> Int -> Either ThinSvdFailure Int +checkedThinSvdCardinality context leftCount rightCount = + first + (const (ThinSvdDimensionViolation (context <> " cardinality exceeds non-negative Int range"))) + (checkedNonNegativeProduct leftCount rightCount) + +thinSvdFullColumnRankST :: Int -> Int -> Int -> Int -> DenseDoubleMatrix -> ST s (Either ThinSvdFailure ThinSvdResult) +thinSvdFullColumnRankST !rowCount !columnCount !matrixEntryCount !rightEntryCount matrixValue = do + leftColumns <- newDenseWork rowCount columnCount + rightVectors <- newDenseWork columnCount columnCount + setIdentityDenseWork rightVectors + copyRowMajorToColumns rowCount columnCount matrixValue leftColumns + sweepResult <- runJacobiSweeps rowCount columnCount leftColumns rightVectors + case sweepResult of + Left err -> pure (Left err) + Right () -> do + singularValues <- singularValuesFromColumns columnCount leftColumns + let !maximumSingular = maximumSingularValue singularValues + !rankTolerance = fromIntegral (max 1 rowCount) * epsDouble * max 1.0 maximumSingular + case firstRankDeficiency rankTolerance singularValues of + Just (columnIndex, singularValue) -> pure (Left (ThinSvdRankDeficient columnIndex singularValue)) + Nothing -> Right <$> projectThinSvdResult rowCount columnCount matrixEntryCount rightEntryCount leftColumns rightVectors singularValues + +copyRowMajorToColumns :: Int -> Int -> DenseDoubleMatrix -> MutableDenseWork s -> ST s () +copyRowMajorToColumns !rowCount !columnCount matrixValue columns = + forIndex 0 rowCount $ \rowIndex -> + forIndex 0 columnCount $ \columnIndex -> + writeDenseWork columns rowIndex columnIndex (payload `S.unsafeIndex` (rowIndex * columnCount + columnIndex)) + where + payload = denseDoubleMatrixToRowMajorVector matrixValue +{-# INLINE copyRowMajorToColumns #-} + +runJacobiSweeps :: + Int -> + Int -> + MutableDenseWork s -> + MutableDenseWork s -> + ST s (Either ThinSvdFailure ()) +runJacobiSweeps !rowCount !columnCount leftColumns rightVectors = sweepAt 0 + where + !sweepBudget = max 8 (12 * max 1 columnCount) + !pairTolerance = 64.0 * epsDouble + + sweepAt !sweepIndex + | columnCount <= 1 = pure (Right ()) + | sweepIndex >= sweepBudget = do + finalCross <- maximumNormalizedCross columnCount leftColumns + pure (Left (ThinSvdSweepBudgetNonConvergence sweepBudget finalCross)) + | otherwise = do + summary <- sweepColumnPairs rowCount columnCount pairTolerance leftColumns rightVectors + if sweepMaximumCross summary <= pairTolerance + then pure (Right ()) + else sweepAt (sweepIndex + 1) + +data SweepSummary = SweepSummary + { sweepMaximumCross :: !Double + } + +sweepColumnPairs :: + Int -> + Int -> + Double -> + MutableDenseWork s -> + MutableDenseWork s -> + ST s SweepSummary +sweepColumnPairs !rowCount !columnCount !pairTolerance leftColumns rightVectors = + goLeft 0 0.0 + where + goLeft !leftColumn !maximumCross + | leftColumn >= columnCount - 1 = pure (SweepSummary maximumCross) + | otherwise = do + nextMaximum <- goRight leftColumn (leftColumn + 1) maximumCross + goLeft (leftColumn + 1) nextMaximum + + goRight !leftColumn !rightColumn !maximumCross + | rightColumn >= columnCount = pure maximumCross + | otherwise = do + alpha <- dotDenseColumns leftColumns leftColumn leftColumn + beta <- dotDenseColumns leftColumns rightColumn rightColumn + gamma <- dotDenseColumns leftColumns leftColumn rightColumn + let !crossValue = normalizedCross alpha beta gamma + if crossValue > pairTolerance + then do + rotateJacobiColumns rowCount leftColumns rightVectors leftColumn rightColumn alpha beta gamma + goRight leftColumn (rightColumn + 1) (max maximumCross crossValue) + else goRight leftColumn (rightColumn + 1) (max maximumCross crossValue) + +maximumNormalizedCross :: Int -> MutableDenseWork s -> ST s Double +maximumNormalizedCross !columnCount leftColumns = goLeft 0 0.0 + where + goLeft !leftColumn !maximumCross + | leftColumn >= columnCount - 1 = pure maximumCross + | otherwise = do + nextMaximum <- goRight leftColumn (leftColumn + 1) maximumCross + goLeft (leftColumn + 1) nextMaximum + + goRight !leftColumn !rightColumn !maximumCross + | rightColumn >= columnCount = pure maximumCross + | otherwise = do + alpha <- dotDenseColumns leftColumns leftColumn leftColumn + beta <- dotDenseColumns leftColumns rightColumn rightColumn + gamma <- dotDenseColumns leftColumns leftColumn rightColumn + goRight leftColumn (rightColumn + 1) (max maximumCross (normalizedCross alpha beta gamma)) + +normalizedCross :: Double -> Double -> Double -> Double +normalizedCross !alpha !beta !gamma = + let !denominator = sqrt (max 0.0 alpha * max 0.0 beta) + in if denominator <= 0.0 + then 0.0 + else abs gamma / denominator +{-# INLINE normalizedCross #-} + +rotateJacobiColumns :: + Int -> + MutableDenseWork s -> + MutableDenseWork s -> + Int -> + Int -> + Double -> + Double -> + Double -> + ST s () +rotateJacobiColumns !rowCount leftColumns rightVectors !leftColumn !rightColumn !alpha !beta !gamma = do + let !(cosineValue, sineValue) = jacobiRotation alpha beta gamma + rotateColumnPair rowCount leftColumns leftColumn rightColumn cosineValue sineValue + rotateColumnPair columnCount rightVectors leftColumn rightColumn cosineValue sineValue + where + MutableDenseWork columnCount _ _ = rightVectors + +jacobiRotation :: Double -> Double -> Double -> (Double, Double) +jacobiRotation !alpha !beta !gamma = + let !tauValue = (beta - alpha) / (2.0 * gamma) + !tangentValue = + if tauValue < 0.0 + then (-1.0) / ((-tauValue) + hypotStable tauValue 1.0) + else 1.0 / (tauValue + hypotStable tauValue 1.0) + !cosineValue = 1.0 / hypotStable 1.0 tangentValue + !sineValue = tangentValue * cosineValue + in (cosineValue, sineValue) +{-# INLINE jacobiRotation #-} + +rotateColumnPair :: Int -> MutableDenseWork s -> Int -> Int -> Double -> Double -> ST s () +rotateColumnPair !rowCount work !leftColumn !rightColumn !cosineValue !sineValue = + forIndex 0 rowCount $ \rowIndex -> do + leftValue <- readDenseWork work rowIndex leftColumn + rightValue <- readDenseWork work rowIndex rightColumn + writeDenseWork work rowIndex leftColumn (cosineValue * leftValue - sineValue * rightValue) + writeDenseWork work rowIndex rightColumn (sineValue * leftValue + cosineValue * rightValue) +{-# INLINE rotateColumnPair #-} + +singularValuesFromColumns :: Int -> MutableDenseWork s -> ST s (S.Vector Double) +singularValuesFromColumns !columnCount leftColumns = do + singularValueBuffer <- SM.new columnCount + forIndex 0 columnCount $ \columnIndex -> do + normSquared <- dotDenseColumns leftColumns columnIndex columnIndex + SM.write singularValueBuffer columnIndex (sqrt (max 0.0 normSquared)) + S.unsafeFreeze singularValueBuffer + +maximumSingularValue :: S.Vector Double -> Double +maximumSingularValue singularValues = + S.foldl' max 0.0 singularValues + +firstRankDeficiency :: Double -> S.Vector Double -> Maybe (Int, Double) +firstRankDeficiency !rankTolerance singularValues = go 0 + where + go !columnIndex + | columnIndex >= S.length singularValues = Nothing + | otherwise = + let !singularValue = singularValues `S.unsafeIndex` columnIndex + in if singularValue <= rankTolerance + then Just (columnIndex, singularValue) + else go (columnIndex + 1) + +projectThinSvdResult :: + Int -> + Int -> + Int -> + Int -> + MutableDenseWork s -> + MutableDenseWork s -> + S.Vector Double -> + ST s ThinSvdResult +projectThinSvdResult !rowCount !columnCount !matrixEntryCount !rightEntryCount leftColumns rightVectors singularValues = do + uBuffer <- SM.new matrixEntryCount + sigmaBuffer <- SM.new columnCount + vtBuffer <- SM.new rightEntryCount + let orderedColumns = + fmap fst + . sortBy (flip (comparing snd)) + $ [(columnIndex, singularValues `S.unsafeIndex` columnIndex) | columnIndex <- [0 .. columnCount - 1]] + writeOrderedColumns orderedColumns 0 uBuffer sigmaBuffer vtBuffer + uValues <- S.unsafeFreeze uBuffer + sigmaValues <- S.unsafeFreeze sigmaBuffer + vtValues <- S.unsafeFreeze vtBuffer + pure + ThinSvdResult + { thinSvdLeftSingularVectors = trustedDenseDoubleMatrixRowMajor rowCount columnCount uValues, + thinSvdSingularValues = sigmaValues, + thinSvdRightSingularVectorsTransposed = trustedDenseDoubleMatrixRowMajor columnCount columnCount vtValues + } + where + writeOrderedColumns orderedColumns !targetColumn uBuffer sigmaBuffer vtBuffer = + case orderedColumns of + [] -> pure () + sourceColumn : remainingColumns -> do + let !singularValue = singularValues `S.unsafeIndex` sourceColumn + !inverseSingular = 1.0 / singularValue + SM.write sigmaBuffer targetColumn singularValue + forIndex 0 rowCount $ \rowIndex -> do + leftEntry <- readDenseWork leftColumns rowIndex sourceColumn + SM.write uBuffer (rowIndex * columnCount + targetColumn) (leftEntry * inverseSingular) + forIndex 0 columnCount $ \columnIndex -> do + rightEntry <- readDenseWork rightVectors columnIndex sourceColumn + SM.write vtBuffer (targetColumn * columnCount + columnIndex) rightEntry + writeOrderedColumns remainingColumns (targetColumn + 1) uBuffer sigmaBuffer vtBuffer diff --git a/src-dense/Moonlight/LinAlg/Pure/Dense/Basic.hs b/src-dense/Moonlight/LinAlg/Pure/Dense/Basic.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Pure/Dense/Basic.hs @@ -0,0 +1,67 @@ + +module Moonlight.LinAlg.Pure.Dense.Basic + ( mapMatrix, + add, + mult, + transpose, + ) +where + +import GHC.TypeNats (KnownNat) +import Moonlight.Core (AdditiveGroup, MoonlightError, Semiring) +import qualified Moonlight.Core as Core +import Moonlight.LinAlg.Internal.Storage + ( matrixMultiplyList, + matrixTransposeList, + matrixZipList, + ) +import Moonlight.LinAlg.Pure.Dense.Types + ( Matrix, + fromListMatrix, + matrixShape, + toListMatrix, + ) +import Prelude + +mapMatrix :: + forall r c a b. + (KnownNat r, KnownNat c) => + (a -> b) -> + Matrix r c a -> + Either MoonlightError (Matrix r c b) +mapMatrix fn matrixValue = + fromListMatrix @r @c (map fn (toListMatrix matrixValue)) + +add :: + forall r c a. + (KnownNat r, KnownNat c, AdditiveGroup a) => + Matrix r c a -> + Matrix r c a -> + Either MoonlightError (Matrix r c a) +add left right = do + let (rowCount, columnCount) = matrixShape left + (rightRows, rightCols) = matrixShape right + values <- matrixZipList rowCount columnCount rightRows rightCols Core.add (toListMatrix left) (toListMatrix right) + fromListMatrix @r @c values + +mult :: + forall r m c a. + (KnownNat r, KnownNat m, KnownNat c, Semiring a) => + Matrix r m a -> + Matrix m c a -> + Either MoonlightError (Matrix r c a) +mult left right = do + let (leftRows, leftCols) = matrixShape left + (rightRows, rightCols) = matrixShape right + values <- matrixMultiplyList leftRows leftCols rightRows rightCols (toListMatrix left) (toListMatrix right) + fromListMatrix @r @c values + +transpose :: + forall r c a. + (KnownNat r, KnownNat c) => + Matrix r c a -> + Either MoonlightError (Matrix c r a) +transpose matrixValue = do + let (rowCount, columnCount) = matrixShape matrixValue + values <- matrixTransposeList rowCount columnCount (toListMatrix matrixValue) + fromListMatrix @c @r values diff --git a/src-dense/Moonlight/LinAlg/Pure/Dense/Block.hs b/src-dense/Moonlight/LinAlg/Pure/Dense/Block.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Pure/Dense/Block.hs @@ -0,0 +1,285 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE StandaloneKindSignatures #-} + +module Moonlight.LinAlg.Pure.Dense.Block + ( BlockMatrixFailure (..), + invertRationalBlock, + invertGF2Block, + invertUnimodularIntegerBlock, + ) +where + +import Control.Monad (foldM) +import Data.Kind (Type) +import Data.List (transpose) +import Data.Ratio (denominator, numerator) +import Moonlight.LinAlg.Internal.Discrete (GF2 (..)) + +-- | Failures for finite dense block inversion. These are obstruction values, +-- not runtime accidents: Schur contraction is only valid when the pivot block is +-- square and invertible in the requested coefficient domain. +type BlockMatrixFailure :: Type +data BlockMatrixFailure + = BlockMatrixNotSquare !Int ![Int] + | BlockMatrixSingular !Int + | BlockMatrixNonUnimodular ![[Rational]] + | BlockMatrixInverseLawFailed !Int + deriving stock (Eq, Show) + +type FieldBlockOps :: Type -> Type +data FieldBlockOps coefficient = FieldBlockOps + { fboZero :: !coefficient, + fboOne :: !coefficient, + fboAdd :: coefficient -> coefficient -> coefficient, + fboNegate :: coefficient -> coefficient, + fboMultiply :: coefficient -> coefficient -> coefficient, + fboInverse :: coefficient -> Maybe coefficient + } + +invertRationalBlock :: [[Rational]] -> Either BlockMatrixFailure [[Rational]] +invertRationalBlock = invertFieldBlock rationalBlockOps +{-# INLINEABLE invertRationalBlock #-} + +invertGF2Block :: [[GF2]] -> Either BlockMatrixFailure [[GF2]] +invertGF2Block = invertFieldBlock gf2BlockOps +{-# INLINEABLE invertGF2Block #-} + +invertUnimodularIntegerBlock :: [[Integer]] -> Either BlockMatrixFailure [[Integer]] +invertUnimodularIntegerBlock matrix = do + dimension <- squareDimension matrix + rationalInverse <- invertRationalBlock (fmap (fmap toRational) matrix) + integerInverse <- + maybe + (Left (BlockMatrixNonUnimodular rationalInverse)) + Right + (traverse (traverse rationalToIntegerExact) rationalInverse) + if matrixProductInteger matrix integerInverse == identityMatrix dimension + && matrixProductInteger integerInverse matrix == identityMatrix dimension + then Right integerInverse + else Left (BlockMatrixInverseLawFailed dimension) +{-# INLINEABLE invertUnimodularIntegerBlock #-} + +invertFieldBlock :: Eq coefficient => FieldBlockOps coefficient -> [[coefficient]] -> Either BlockMatrixFailure [[coefficient]] +invertFieldBlock ops matrix = do + dimension <- squareDimension matrix + let augmented = zipWith (<>) matrix (identityMatrixWith ops dimension) + reduced <- foldM (rrefPivotStep ops dimension) augmented [0 .. dimension - 1] + let leftBlock = fmap (take dimension) reduced + rightBlock = fmap (drop dimension) reduced + identity = identityMatrixWith ops dimension + if leftBlock == identity + && matrixProductWith ops matrix rightBlock == identity + && matrixProductWith ops rightBlock matrix == identity + then Right rightBlock + else Left (BlockMatrixInverseLawFailed dimension) +{-# INLINEABLE invertFieldBlock #-} + +rationalBlockOps :: FieldBlockOps Rational +rationalBlockOps = + FieldBlockOps + { fboZero = 0, + fboOne = 1, + fboAdd = (+), + fboNegate = negate, + fboMultiply = (*), + fboInverse = \value -> if value == 0 then Nothing else Just (recip value) + } + +gf2BlockOps :: FieldBlockOps GF2 +gf2BlockOps = + FieldBlockOps + { fboZero = GF2Zero, + fboOne = GF2One, + fboAdd = (+), + fboNegate = id, + fboMultiply = (*), + fboInverse = \value -> case value of + GF2Zero -> Nothing + GF2One -> Just GF2One + } + +squareDimension :: [[coefficient]] -> Either BlockMatrixFailure Int +squareDimension matrix = + let rowCount = length matrix + widths = fmap length matrix + in if all (== rowCount) widths + then Right rowCount + else Left (BlockMatrixNotSquare rowCount widths) +{-# INLINEABLE squareDimension #-} + +rrefPivotStep :: Eq coefficient => FieldBlockOps coefficient -> Int -> [[coefficient]] -> Int -> Either BlockMatrixFailure [[coefficient]] +rrefPivotStep ops dimension rows pivotIndex = do + pivotRowIndex <- + requireBlockValue + (BlockMatrixSingular dimension) + (findPivotRow ops pivotIndex rows) + swappedRows <- swapRows pivotIndex pivotRowIndex rows + pivotRow <- + requireBlockValue + (BlockMatrixSingular dimension) + (rowAt pivotIndex swappedRows) + pivotValue <- + requireBlockValue + (BlockMatrixSingular dimension) + (entryAt pivotIndex pivotRow) + pivotInverse <- + requireBlockValue + (BlockMatrixSingular dimension) + (fboInverse ops pivotValue) + let normalizedPivot = fmap (fboMultiply ops pivotInverse) pivotRow + normalizedRows = replaceRow pivotIndex normalizedPivot swappedRows + pure (eliminatePivotColumn ops pivotIndex normalizedPivot normalizedRows) +{-# INLINEABLE rrefPivotStep #-} + +findPivotRow :: Eq coefficient => FieldBlockOps coefficient -> Int -> [[coefficient]] -> Maybe Int +findPivotRow ops pivotIndex = + fmap fst + . findFirst + ( \(rowIndex, rowValues) -> + rowIndex >= pivotIndex + && maybe False (/= fboZero ops) (entryAt pivotIndex rowValues) + ) + . zip [0 ..] +{-# INLINEABLE findPivotRow #-} + +eliminatePivotColumn :: Eq coefficient => FieldBlockOps coefficient -> Int -> [coefficient] -> [[coefficient]] -> [[coefficient]] +eliminatePivotColumn ops pivotIndex normalizedPivot = + fmap + ( \(rowIndex, rowValues) -> + if rowIndex == pivotIndex + then normalizedPivot + else + case entryAt pivotIndex rowValues of + Nothing -> rowValues + Just factor + | factor == fboZero ops -> rowValues + | otherwise -> subtractMultiple ops factor normalizedPivot rowValues + ) + . zip [0 ..] +{-# INLINEABLE eliminatePivotColumn #-} + +subtractMultiple :: FieldBlockOps coefficient -> coefficient -> [coefficient] -> [coefficient] -> [coefficient] +subtractMultiple ops factor pivotRow targetRow = + zipWith + (\targetEntry pivotEntry -> fboAdd ops targetEntry (fboNegate ops (fboMultiply ops factor pivotEntry))) + targetRow + pivotRow +{-# INLINEABLE subtractMultiple #-} + +swapRows :: Int -> Int -> [[coefficient]] -> Either BlockMatrixFailure [[coefficient]] +swapRows leftIndex rightIndex rows = do + leftRow <- requireBlockValue (BlockMatrixSingular (length rows)) (rowAt leftIndex rows) + rightRow <- requireBlockValue (BlockMatrixSingular (length rows)) (rowAt rightIndex rows) + pure + ( fmap + ( \(rowIndex, rowValues) -> + if rowIndex == leftIndex + then rightRow + else + if rowIndex == rightIndex + then leftRow + else rowValues + ) + (zip [0 ..] rows) + ) +{-# INLINEABLE swapRows #-} + +replaceRow :: Int -> [coefficient] -> [[coefficient]] -> [[coefficient]] +replaceRow targetIndex replacement = + fmap + (\(rowIndex, rowValues) -> if rowIndex == targetIndex then replacement else rowValues) + . zip [0 ..] +{-# INLINEABLE replaceRow #-} + +requireBlockValue :: BlockMatrixFailure -> Maybe value -> Either BlockMatrixFailure value +requireBlockValue failureValue = + maybe (Left failureValue) Right +{-# INLINEABLE requireBlockValue #-} + +rowAt :: Int -> [row] -> Maybe row +rowAt indexValue rows + | indexValue < 0 = Nothing + | otherwise = + case drop indexValue rows of + rowValue : _ -> Just rowValue + [] -> Nothing +{-# INLINE rowAt #-} + +entryAt :: Int -> [entry] -> Maybe entry +entryAt indexValue entries + | indexValue < 0 = Nothing + | otherwise = + case drop indexValue entries of + entryValue : _ -> Just entryValue + [] -> Nothing +{-# INLINE entryAt #-} + +identityMatrix :: Num coefficient => Int -> [[coefficient]] +identityMatrix dimension = + [ [ if rowIndex == columnIndex then 1 else 0 + | columnIndex <- [0 .. dimension - 1] + ] + | rowIndex <- [0 .. dimension - 1] + ] +{-# INLINEABLE identityMatrix #-} + +identityMatrixWith :: FieldBlockOps coefficient -> Int -> [[coefficient]] +identityMatrixWith ops dimension = + [ [ if rowIndex == columnIndex then fboOne ops else fboZero ops + | columnIndex <- [0 .. dimension - 1] + ] + | rowIndex <- [0 .. dimension - 1] + ] +{-# INLINEABLE identityMatrixWith #-} + +matrixProductInteger :: [[Integer]] -> [[Integer]] -> [[Integer]] +matrixProductInteger = matrixProductWith integerBlockOps +{-# INLINEABLE matrixProductInteger #-} + +integerBlockOps :: FieldBlockOps Integer +integerBlockOps = + FieldBlockOps + { fboZero = 0, + fboOne = 1, + fboAdd = (+), + fboNegate = negate, + fboMultiply = (*), + fboInverse = \value -> case value of + 1 -> Just 1 + -1 -> Just (-1) + _ -> Nothing + } + +matrixProductWith :: FieldBlockOps coefficient -> [[coefficient]] -> [[coefficient]] -> [[coefficient]] +matrixProductWith ops left right = + let rightColumns = transpose right + in fmap + ( \leftRow -> + fmap + (dotWith ops leftRow) + rightColumns + ) + left +{-# INLINEABLE matrixProductWith #-} + +dotWith :: FieldBlockOps coefficient -> [coefficient] -> [coefficient] -> coefficient +dotWith ops left right = + foldl' + (fboAdd ops) + (fboZero ops) + (zipWith (fboMultiply ops) left right) +{-# INLINEABLE dotWith #-} + +rationalToIntegerExact :: Rational -> Maybe Integer +rationalToIntegerExact value = + if denominator value == 1 + then Just (numerator value) + else Nothing +{-# INLINEABLE rationalToIntegerExact #-} + +findFirst :: (a -> Bool) -> [a] -> Maybe a +findFirst predicate = + foldr + (\value rest -> if predicate value then Just value else rest) + Nothing +{-# INLINE findFirst #-} diff --git a/src-dense/Moonlight/LinAlg/Pure/Dense/Decomposition.hs b/src-dense/Moonlight/LinAlg/Pure/Dense/Decomposition.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Pure/Dense/Decomposition.hs @@ -0,0 +1,152 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE RecordWildCards #-} + +module Moonlight.LinAlg.Pure.Dense.Decomposition + ( qrDecompFullColumnRank, + choleskyDecomp, + symmetricEigen, + symmetricEigenPairs, + thinSvdFullColumnRank, + ) +where + +import Data.List (sortBy) +import Data.Ord (comparing) +import Data.Vector.Storable qualified as S +import GHC.TypeNats (KnownNat) +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Internal.Dense.DoubleFactorization + ( choleskyLower, + qrFullColumnRank, + ) +import Moonlight.LinAlg.Internal.Dense.OneSidedJacobiSVD qualified as JacobiSVD +import Moonlight.LinAlg.Pure.Dense.Rows (transposeRowsExact) +import Moonlight.LinAlg.Internal.Eigen.Input + ( validateSymmetricEigenInput, + ) +import Moonlight.LinAlg.Internal.Eigen.Symmetric + ( SymmetricEigenResult (..), + symmetricEigenPairsDenseUnchecked, + ) +import Moonlight.LinAlg.Pure.Dense.Flat + ( denseDoubleMatrixShape, + denseDoubleMatrixToRowMajorVector, + denseDoubleMatrixToRows, + mkDenseDoubleMatrixRows, + ) +import Moonlight.LinAlg.Pure.Dense.Types (Matrix, Vector, fromListMatrix, fromListVector, toListMatrix) +import qualified Moonlight.LinAlg.Pure.Dense.Types as DenseTypes +import Prelude + +qrDecompFullColumnRank :: + forall r c. + (KnownNat r, KnownNat c) => + Matrix r c Double -> + Either MoonlightError (Matrix r c Double, Matrix c c Double) +qrDecompFullColumnRank matrixValue = do + let (rowCount, columnCount) = DenseTypes.matrixShape matrixValue + (qValues, rValues) <- qrFullColumnRank rowCount columnCount (toListMatrix matrixValue) + qMatrix <- fromListMatrix @r @c qValues + rMatrix <- fromListMatrix @c @c rValues + pure (qMatrix, rMatrix) + +choleskyDecomp :: + forall n. + KnownNat n => + Matrix n n Double -> + Either MoonlightError (Matrix n n Double) +choleskyDecomp matrixValue = do + let (matrixSize, _) = DenseTypes.matrixShape matrixValue + lowerValues <- choleskyLower matrixSize (toListMatrix matrixValue) + fromListMatrix @n @n lowerValues + +symmetricEigenPairs :: Int -> [[Double]] -> Either MoonlightError [(Double, [Double])] +symmetricEigenPairs matrixSize matrixRows = do + validateSymmetricEigenInput "symmetric eigen decomposition" matrixSize matrixRows + matrixValue <- mkDenseDoubleMatrixRows matrixRows + eigenResultToPairs <$> symmetricEigenPairsDenseUnchecked matrixSize matrixValue + +symmetricEigen :: + forall n. + KnownNat n => + Matrix n n Double -> + Either MoonlightError (Vector n Double, Matrix n n Double) +symmetricEigen matrixValue = do + let (matrixSize, _) = DenseTypes.matrixShape matrixValue + matrixRows <- DenseTypes.matrixToRows matrixValue + validateSymmetricEigenInput "symmetric eigen decomposition" matrixSize matrixRows + denseMatrix <- mkDenseDoubleMatrixRows matrixRows + eigenResult <- symmetricEigenPairsDenseUnchecked matrixSize denseMatrix + let orderedPairs = sortBy (flip (comparing fst)) (eigenResultToPairs eigenResult) + eigenvalues = map fst orderedPairs + eigenvectors = map snd orderedPairs + eigenvalueVector <- fromListVector @n eigenvalues + eigenvectorRows <- transposeRowsExact eigenvectors + eigenvectorMatrix <- fromListMatrix @n @n (concat eigenvectorRows) + pure (eigenvalueVector, eigenvectorMatrix) + +diagonalRows :: [Double] -> [[Double]] +diagonalRows diagonalValues = + let indexedDiagonalValues = zip [0 :: Int ..] diagonalValues + size = length diagonalValues + in map + (\(rowIndex, diagonalValue) -> map (\columnIndex -> if rowIndex == columnIndex then diagonalValue else 0.0) [0 .. size - 1]) + indexedDiagonalValues + +thinSvdFullColumnRank :: + forall r c. + (KnownNat r, KnownNat c) => + Matrix r c Double -> + Either MoonlightError (Matrix r c Double, Matrix c c Double, Matrix c c Double) +thinSvdFullColumnRank matrixValue = do + rows <- DenseTypes.matrixToRows matrixValue + denseMatrix <- mkDenseDoubleMatrixRows rows + JacobiSVD.ThinSvdResult {..} <- + firstMoonlightSvdFailure (JacobiSVD.thinSvdFullColumnRank denseMatrix) + let singularValues = S.toList thinSvdSingularValues + sRows = diagonalRows singularValues + uRows = denseDoubleMatrixToRows thinSvdLeftSingularVectors + vTRows = denseDoubleMatrixToRows thinSvdRightSingularVectorsTransposed + uMatrix <- fromListMatrix @r @c (concat uRows) + sMatrix <- fromListMatrix @c @c (concat sRows) + vTMatrix <- fromListMatrix @c @c (concat vTRows) + pure (uMatrix, sMatrix, vTMatrix) + +eigenResultToPairs :: SymmetricEigenResult -> [(Double, [Double])] +eigenResultToPairs SymmetricEigenResult {..} = + fmap eigenPairAt [0 .. matrixSize - 1] + where + !(matrixSize, _) = denseDoubleMatrixShape symmetricEigenResultVectors + eigenvectorPayload = denseDoubleMatrixToRowMajorVector symmetricEigenResultVectors + + eigenPairAt !columnIndex = + ( symmetricEigenResultValues `S.unsafeIndex` columnIndex, + fmap + (\rowIndex -> eigenvectorPayload `S.unsafeIndex` (rowIndex * matrixSize + columnIndex)) + [0 .. matrixSize - 1] + ) + +firstMoonlightSvdFailure :: Either JacobiSVD.ThinSvdFailure value -> Either MoonlightError value +firstMoonlightSvdFailure resultValue = + case resultValue of + Right value -> Right value + Left failureValue -> Left (InvariantViolation (thinSvdFailureMessage failureValue)) + +thinSvdFailureMessage :: JacobiSVD.ThinSvdFailure -> String +thinSvdFailureMessage failureValue = + case failureValue of + JacobiSVD.ThinSvdNonFiniteInput -> + "thin Jacobi SVD requires finite entries" + JacobiSVD.ThinSvdDimensionViolation message -> + message + JacobiSVD.ThinSvdRankDeficient columnIndex singularValue -> + "thin Jacobi SVD requires full column rank; column " + <> show columnIndex + <> " singular value " + <> show singularValue + <> " is below rank tolerance" + JacobiSVD.ThinSvdSweepBudgetNonConvergence sweepBudget maximumCross -> + "thin Jacobi SVD exhausted " + <> show sweepBudget + <> " sweeps; maximum normalized column cross=" + <> show maximumCross diff --git a/src-dense/Moonlight/LinAlg/Pure/Dense/Dynamic.hs b/src-dense/Moonlight/LinAlg/Pure/Dense/Dynamic.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Pure/Dense/Dynamic.hs @@ -0,0 +1,193 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Moonlight.LinAlg.Pure.Dense.Dynamic + ( DynVector, + DynMatrix, + mkDynVector, + mkDynMatrix, + dynMatrixFromRows, + dynMatrixToRows, + toDynVector, + toDynMatrix, + fromDynVector, + fromDynMatrix, + withDynVector, + withDynMatrix, + dynVectorLength, + dynMatrixShape, + dynMatrixDenseRows, + dynVectorToList, + dynMatrixToList, + ) +where + +import Data.Kind (Type) +import Data.Proxy (Proxy (..)) +import GHC.TypeNats (KnownNat, SomeNat (..), natVal, someNatVal) +import Moonlight.Core + ( MoonlightError (..), + checkedNaturalToInt, + ) +import Moonlight.LinAlg.Internal.Storage (checkFlatLength, chunkRows) +import Moonlight.LinAlg.Pure.Dense.Rows + ( DenseRows, + denseRowsShape, + denseRowsToLists, + mkDenseRows, + mkDenseRowsFromFlat, + ) +import Moonlight.LinAlg.Pure.Dense.Types + ( Matrix, + Vector, + fromListMatrix, + fromListVector, + matrixShape, + toListVector, + toListMatrix, + vectorLength, + ) +import Prelude + +type DynVector :: Type -> Type +data DynVector a = DynVector + { dynVectorLength :: Int, + dynVectorValues :: [a] + } + +type DynMatrix :: Type -> Type +data DynMatrix a = DynMatrix + { dynRows :: Int, + dynCols :: Int, + dynMatrixValues :: [a] + } + +mkDynVector :: Int -> [a] -> Either MoonlightError (DynVector a) +mkDynVector expected values + | expected < 0 = Left (InvariantViolation "dynamic vector length must be non-negative") + | expected /= length values = Left (InvariantViolation "dynamic vector payload length mismatch") + | otherwise = Right (DynVector expected values) + +mkDynMatrix :: Int -> Int -> [a] -> Either MoonlightError (DynMatrix a) +mkDynMatrix rowCount columnCount values = do + checkFlatLength rowCount columnCount values + Right (DynMatrix rowCount columnCount values) + +dynMatrixFromRows :: [[a]] -> Either MoonlightError (DynMatrix a) +dynMatrixFromRows rowValues = do + denseRowsValue <- mkDenseRows rowValues + let (rowCount, columnCount) = denseRowsShape denseRowsValue + Right + DynMatrix + { dynRows = rowCount, + dynCols = columnCount, + dynMatrixValues = concat (denseRowsToLists denseRowsValue) + } + +dynMatrixToRows :: DynMatrix a -> Either MoonlightError [[a]] +dynMatrixToRows dynValue = do + let rowCount = dynRows dynValue + columnCount = dynCols dynValue + values = dynMatrixValues dynValue + checkFlatLength rowCount columnCount values + if columnCount == 0 + then Right (replicate rowCount []) + else chunkRows columnCount values + +toDynVector :: KnownNat n => Vector n a -> DynVector a +toDynVector vectorValue = + DynVector + { dynVectorLength = vectorLength vectorValue, + dynVectorValues = toListVector vectorValue + } + +toDynMatrix :: (KnownNat r, KnownNat c) => Matrix r c a -> DynMatrix a +toDynMatrix matrixValue = + let (rowCount, columnCount) = matrixShape matrixValue + in DynMatrix + { dynRows = rowCount, + dynCols = columnCount, + dynMatrixValues = toListMatrix matrixValue + } + +fromDynVector :: forall n a. KnownNat n => DynVector a -> Either MoonlightError (Vector n a) +fromDynVector dynValue = do + expected <- checkedStaticDimension @n + let actual = dynVectorLength dynValue + if actual /= expected + then + Left + ( InvariantViolation + ( "dynamic vector shape does not match static dimension: expected " + <> show expected + <> " but received " + <> show actual + ) + ) + else fromListVector @n (dynVectorValues dynValue) + +fromDynMatrix :: forall r c a. (KnownNat r, KnownNat c) => DynMatrix a -> Either MoonlightError (Matrix r c a) +fromDynMatrix dynValue = do + expectedRows <- checkedStaticDimension @r + expectedColumns <- checkedStaticDimension @c + let expected = (expectedRows, expectedColumns) + actual = dynMatrixShape dynValue + if actual /= expected + then + Left + ( InvariantViolation + ( "dynamic matrix shape does not match static dimensions: expected " + <> show expected + <> " but received " + <> show actual + ) + ) + else fromListMatrix @r @c (dynMatrixValues dynValue) + +checkedStaticDimension :: forall n. KnownNat n => Either MoonlightError Int +checkedStaticDimension = + either + (const (Left (InvariantViolation "static dimension exceeds Int cardinality"))) + Right + (checkedNaturalToInt (natVal (Proxy @n))) + +withDynVector :: + forall a b. + DynVector a -> + (forall n. KnownNat n => Vector n a -> b) -> + Either MoonlightError b +withDynVector dynValue callback + | dynVectorLength dynValue < 0 = Left (InvariantViolation "dynamic vector length must be non-negative") + | otherwise = + case someNatVal (fromIntegral (dynVectorLength dynValue)) of + SomeNat (_proxyN :: Proxy n) -> + callback <$> (fromListVector (dynVectorValues dynValue) :: Either MoonlightError (Vector n a)) + +withDynMatrix :: + forall a b. + DynMatrix a -> + (forall r c. (KnownNat r, KnownNat c) => Matrix r c a -> b) -> + Either MoonlightError b +withDynMatrix dynValue callback + | dynRows dynValue < 0 || dynCols dynValue < 0 = Left (InvariantViolation "dynamic matrix dimensions must be non-negative") + | otherwise = + case someNatVal (fromIntegral (dynRows dynValue)) of + SomeNat (_proxyR :: Proxy r) -> + case someNatVal (fromIntegral (dynCols dynValue)) of + SomeNat (_proxyC :: Proxy c) -> + callback <$> (fromListMatrix (dynMatrixValues dynValue) :: Either MoonlightError (Matrix r c a)) + +dynMatrixShape :: DynMatrix a -> (Int, Int) +dynMatrixShape dynValue = (dynRows dynValue, dynCols dynValue) + +dynMatrixDenseRows :: DynMatrix a -> Either MoonlightError (DenseRows a) +dynMatrixDenseRows dynValue = + mkDenseRowsFromFlat + (dynRows dynValue) + (dynCols dynValue) + (dynMatrixValues dynValue) + +dynVectorToList :: DynVector a -> [a] +dynVectorToList = dynVectorValues + +dynMatrixToList :: DynMatrix a -> [a] +dynMatrixToList = dynMatrixValues diff --git a/src-dense/Moonlight/LinAlg/Pure/Dense/Exterior.hs b/src-dense/Moonlight/LinAlg/Pure/Dense/Exterior.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Pure/Dense/Exterior.hs @@ -0,0 +1,234 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE StandaloneKindSignatures #-} + +module Moonlight.LinAlg.Pure.Dense.Exterior + ( ExteriorBasis (..), + ExteriorPowerFailure (..), + choose, + exteriorBasis, + exteriorBasisCardinality, + exteriorPowerMatrix, + exteriorPowerMatrixWithShape, + ) +where + +import Data.Kind (Type) +import Data.Vector qualified as Box + +-- | Ordered basis of Λ^p(k^n), represented by increasing source-coordinate +-- subsets. The order is lexicographic and therefore stable across coefficient +-- rings. +type ExteriorBasis :: Type +data ExteriorBasis = ExteriorBasis + { ebDegree :: !Int, + ebRank :: !Int, + ebBasisVectors :: ![[Int]] + } + deriving stock (Eq, Ord, Show) + +type ExteriorPowerFailure :: Type +data ExteriorPowerFailure + = ExteriorNegativeDegree !Int + | ExteriorNegativeSourceRank !Int + | ExteriorNegativeTargetRank !Int + | ExteriorMatrixShapeMismatch + !Int + !Int + !Int + ![Int] + deriving stock (Eq, Ord, Show) + +choose :: Int -> Int -> Integer +choose n k + | n < 0 || k < 0 || k > n = 0 + | otherwise = choosePositive n (min k (n - k)) + where + choosePositive :: Int -> Int -> Integer + choosePositive _ 0 = 1 + choosePositive nValue kValue = + product [fromIntegral (nValue - kValue + 1) .. fromIntegral nValue] + `div` product [1 .. fromIntegral kValue] +{-# INLINEABLE choose #-} + +exteriorBasis :: Int -> Int -> Either ExteriorPowerFailure ExteriorBasis +exteriorBasis degree rankValue + | degree < 0 = Left (ExteriorNegativeDegree degree) + | rankValue < 0 = Left (ExteriorNegativeSourceRank rankValue) + | otherwise = + Right + ExteriorBasis + { ebDegree = degree, + ebRank = rankValue, + ebBasisVectors = combinationsOf degree [0 .. rankValue - 1] + } +{-# INLINEABLE exteriorBasis #-} + +exteriorBasisCardinality :: ExteriorBasis -> Int +exteriorBasisCardinality = + length . ebBasisVectors +{-# INLINE exteriorBasisCardinality #-} + +-- | Infer source/target ranks from a rectangular dense matrix and compute the +-- induced Λ^p matrix. Use 'exteriorPowerMatrixWithShape' when a zero-row matrix +-- must still remember its source rank. +exteriorPowerMatrix :: Num coefficient => Int -> [[coefficient]] -> Either ExteriorPowerFailure [[coefficient]] +exteriorPowerMatrix degree matrix = + case inferredShape matrix of + Left failure -> Left failure + Right (targetRank, sourceRank) -> + exteriorPowerMatrixWithShape degree targetRank sourceRank matrix +{-# INLINEABLE exteriorPowerMatrix #-} + +-- | Compute the induced matrix Λ^p(f) for an explicitly shaped matrix +-- f : k^sourceRank -> k^targetRank. Rows are target coordinates; columns are +-- source coordinates. Entry (I,J) is the determinant of the I×J minor. +exteriorPowerMatrixWithShape :: + Num coefficient => + Int -> + Int -> + Int -> + [[coefficient]] -> + Either ExteriorPowerFailure [[coefficient]] +exteriorPowerMatrixWithShape degree targetRank sourceRank matrix + | degree < 0 = Left (ExteriorNegativeDegree degree) + | sourceRank < 0 = Left (ExteriorNegativeSourceRank sourceRank) + | targetRank < 0 = Left (ExteriorNegativeTargetRank targetRank) + | actualShapeRows matrix /= targetRank || any (/= sourceRank) (actualShapeColumns matrix) = + Left + ( ExteriorMatrixShapeMismatch + targetRank + sourceRank + (actualShapeRows matrix) + (actualShapeColumns matrix) + ) + | otherwise = do + targetBasis <- exteriorBasisWithRole targetRank + sourceBasis <- exteriorBasisWithRole sourceRank + let entries = denseEntries matrix + shapeFailure = + ExteriorMatrixShapeMismatch + targetRank + sourceRank + (actualShapeRows matrix) + (actualShapeColumns matrix) + traverse + (\targetVector -> traverse (minorDeterminant shapeFailure sourceRank entries targetVector) (ebBasisVectors sourceBasis)) + (ebBasisVectors targetBasis) + where + exteriorBasisWithRole rankValue = + case exteriorBasis degree rankValue of + Left (ExteriorNegativeSourceRank badRank) -> Left (ExteriorNegativeTargetRank badRank) + Left failure -> Left failure + Right basis -> Right basis +{-# INLINEABLE exteriorPowerMatrixWithShape #-} + +inferredShape :: [[coefficient]] -> Either ExteriorPowerFailure (Int, Int) +inferredShape matrix = + case actualShapeColumns matrix of + [] -> Right (0, 0) + firstWidth : widths -> + if all (== firstWidth) widths + then Right (length matrix, firstWidth) + else Left (ExteriorMatrixShapeMismatch (length matrix) firstWidth (length matrix) (firstWidth : widths)) + +actualShapeRows :: [[coefficient]] -> Int +actualShapeRows = + length +{-# INLINE actualShapeRows #-} + +actualShapeColumns :: [[coefficient]] -> [Int] +actualShapeColumns = + fmap length +{-# INLINE actualShapeColumns #-} + +combinationsOf :: Int -> [a] -> [[a]] +combinationsOf degree values + | degree < 0 = [] + | otherwise = + case (degree, values) of + (0, _) -> [[]] + (_, []) -> [] + (remaining, value : rest) -> + fmap (value :) (combinationsOf (remaining - 1) rest) + <> combinationsOf remaining rest +{-# INLINEABLE combinationsOf #-} + +denseEntries :: [[coefficient]] -> Box.Vector coefficient +denseEntries matrix = + Box.fromList (concat matrix) +{-# INLINE denseEntries #-} + +entryAt :: + ExteriorPowerFailure -> + Int -> + Box.Vector coefficient -> + Int -> + Int -> + Either ExteriorPowerFailure coefficient +entryAt failure sourceRank entries targetIndex sourceIndex = + maybe (Left failure) Right (entries Box.!? (targetIndex * sourceRank + sourceIndex)) +{-# INLINE entryAt #-} + +minorDeterminant :: + Num coefficient => + ExteriorPowerFailure -> + Int -> + Box.Vector coefficient -> + [Int] -> + [Int] -> + Either ExteriorPowerFailure coefficient +minorDeterminant failure sourceRank entries targetVector sourceVector = + case (targetVector, sourceVector) of + ([], []) -> Right 1 + ([row0], [column0]) -> + entryAt failure sourceRank entries row0 column0 + ([row0, row1], [column0, column1]) -> do + a <- entryAt failure sourceRank entries row0 column0 + b <- entryAt failure sourceRank entries row0 column1 + c <- entryAt failure sourceRank entries row1 column0 + d <- entryAt failure sourceRank entries row1 column1 + Right ((a * d) - (b * c)) + ([row0, row1, row2], [column0, column1, column2]) -> do + a <- entryAt failure sourceRank entries row0 column0 + b <- entryAt failure sourceRank entries row0 column1 + c <- entryAt failure sourceRank entries row0 column2 + d <- entryAt failure sourceRank entries row1 column0 + e <- entryAt failure sourceRank entries row1 column1 + f <- entryAt failure sourceRank entries row1 column2 + g <- entryAt failure sourceRank entries row2 column0 + h <- entryAt failure sourceRank entries row2 column1 + i <- entryAt failure sourceRank entries row2 column2 + Right ((a * e * i) + (b * f * g) + (c * d * h) - (c * e * g) - (b * d * i) - (a * f * h)) + _ -> + fmap determinant + ( traverse + (\targetIndex -> traverse (entryAt failure sourceRank entries targetIndex) sourceVector) + targetVector + ) +{-# INLINE minorDeterminant #-} + +determinant :: Num coefficient => [[coefficient]] -> coefficient +determinant matrix = + case matrix of + [] -> 1 + [singleRow] -> + case singleRow of + [value] -> value + _ -> 0 + firstRow : remainingRows -> + sum + ( fmap + (\(columnIndex, value) -> signFor columnIndex * value * determinant (removeColumn columnIndex remainingRows)) + (zip [0 ..] firstRow) + ) +{-# INLINEABLE determinant #-} + +removeColumn :: Int -> [[coefficient]] -> [[coefficient]] +removeColumn columnIndex = + fmap (fmap snd . filter ((/= columnIndex) . fst) . zip [0 :: Int ..]) +{-# INLINE removeColumn #-} + +signFor :: Num coefficient => Int -> coefficient +signFor columnIndex = + if even columnIndex then 1 else (-1) +{-# INLINE signFor #-} diff --git a/src-dense/Moonlight/LinAlg/Pure/Dense/Field.hs b/src-dense/Moonlight/LinAlg/Pure/Dense/Field.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Pure/Dense/Field.hs @@ -0,0 +1,37 @@ +module Moonlight.LinAlg.Pure.Dense.Field + ( DenseRankBackend, + PLU (..), + KernelBasis (..), + pluDecompFullRank, + rank, + kernel, + ) +where + +import GHC.TypeNats (KnownNat) +import Moonlight.Core (Field, MoonlightError) +import Moonlight.LinAlg.Internal.Backend.Core (DenseRankBackend, runKernel, runPluDecomp, runRank) +import Moonlight.LinAlg.Internal.Backend.PLU (PLU (..)) +import Moonlight.LinAlg.Internal.Backend.RREF (KernelBasis (..)) +import Moonlight.LinAlg.Pure.Dense.Types (Matrix) + +pluDecompFullRank :: + forall r c a. + (KnownNat r, KnownNat c, Field a) => + Matrix r c a -> + Either MoonlightError (PLU r c a) +pluDecompFullRank = runPluDecomp + +rank :: + forall r c a. + (KnownNat r, KnownNat c, Eq a, Field a, DenseRankBackend a) => + Matrix r c a -> + Either MoonlightError Int +rank = runRank + +kernel :: + forall r c a. + (KnownNat r, KnownNat c, Eq a, Field a) => + Matrix r c a -> + Either MoonlightError (KernelBasis c a) +kernel = runKernel diff --git a/src-dense/Moonlight/LinAlg/Pure/Dense/GF2.hs b/src-dense/Moonlight/LinAlg/Pure/Dense/GF2.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Pure/Dense/GF2.hs @@ -0,0 +1,330 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE StandaloneKindSignatures #-} + +module Moonlight.LinAlg.Pure.Dense.GF2 + ( GF2 (..), + gf2Zero, + gf2One, + gf2FromBool, + gf2ToBool, + PackedRow, + packedRowWidth, + packedRowNonZeroCount, + emptyPackedRow, + unitPackedRow, + packedRowFromIndices, + packedRowIndices, + packedRowMember, + packedRowIsZero, + packedRowXor, + packedRowRemap, + PackedLinearMap, + packedLinearMapDomain, + packedLinearMapCodomain, + packedLinearMapColumns, + packedLinearMapFromColumns, + packedLinearMapFromEntries, + zeroPackedLinearMap, + identityPackedLinearMap, + applyPackedLinearMap, + composePackedLinearMaps, + addPackedLinearMaps, + packedLinearMapIsZero, + PackedSpan, + emptyPackedSpan, + packedSpanFromRows, + reducePackedRow, + admitPackedRow, + ColumnReduction (..), + reducePackedColumns, + PackedCoordinateSolver, + packedCoordinateSolver, + coordinatesInPackedBasis, + inverseFromPackedBasisColumns, + GF2MatrixEntry (..), + GF2PackedMatrix, + gf2PackedRows, + gf2PackedColumns, + gf2PackedWordsPerRow, + gf2PackedWords, + GF2PackedMatrixFailure (..), + mkGF2PackedMatrix, + mkGF2PackedMatrixFromRowMajor, + rankGF2PackedMatrix, + gf2PackedMatrixLinearMap, + inverseGF2PackedMatrix, + GF2SparseColumn, + gf2SparseColumnIndex, + gf2SparseColumnRows, + mkGF2SparseColumn, + GF2SparseReducerConfig, + gf2SparseDensifyThreshold, + mkGF2SparseReducerConfig, + defaultGF2SparseReducerConfig, + GF2SparseColumnReduction (..), + reduceGF2SparseColumns, + rankGF2SparseColumns, + independentGF2SparseColumns, + kernelBasisGF2SparseColumns, + ) +where + +import Data.Bifunctor (first) +import Data.Bits + ( testBit, + ) +import Data.Kind + ( Type, + ) +import Data.Maybe + ( listToMaybe, + ) +import Data.Vector + ( Vector, + ) +import Data.Vector qualified as V +import Data.Vector.Unboxed qualified as U +import Data.Word + ( Word64, + ) +import Moonlight.Core + ( MoonlightError, + checkedNaturalToInt, + checkedNonNegativeProduct, + ) +import Moonlight.LinAlg.Internal.GF2.SparseColumn + ( GF2SparseColumn, + GF2SparseColumnReduction (..), + GF2SparseReducerConfig, + defaultGF2SparseReducerConfig, + gf2SparseColumnIndex, + gf2SparseColumnRows, + gf2SparseDensifyThreshold, + independentGF2SparseColumns, + kernelBasisGF2SparseColumns, + mkGF2SparseColumn, + mkGF2SparseReducerConfig, + rankGF2SparseColumns, + reduceGF2SparseColumns, + ) +import Moonlight.LinAlg.Internal.GF2.Xor + ( ColumnReduction (..), + PackedCoordinateSolver, + PackedLinearMap, + PackedRow, + PackedSpan, + addPackedLinearMaps, + admitPackedRow, + applyPackedLinearMap, + composePackedLinearMaps, + coordinatesInPackedBasis, + emptyPackedRow, + emptyPackedSpan, + identityPackedLinearMap, + inverseFromPackedBasisColumns, + packedCoordinateSolver, + packedLinearMapCodomain, + packedLinearMapColumns, + packedLinearMapDomain, + packedLinearMapFromColumns, + packedLinearMapFromEntries, + packedLinearMapIsZero, + packedRowFromIndices, + packedRowIndices, + packedRowIsZero, + packedRowMember, + packedRowNonZeroCount, + packedRowRemap, + packedRowWidth, + packedRowXor, + packedSpanFromRows, + reducePackedColumns, + reducePackedRow, + unitPackedRow, + zeroPackedLinearMap, + ) +import Moonlight.LinAlg.Internal.Discrete + ( GF2 (..), + PackedBitMatrix (..), + gf2FromBool, + gf2One, + gf2ToBool, + gf2Zero, + matrixRowWords, + packedBitMatrixFromRowMajor, + packedBitMatrixFromXorEntries, + rankPackedRows, + ) +import Numeric.Natural + ( Natural, + ) + +type GF2MatrixEntry :: Type +data GF2MatrixEntry = GF2MatrixEntry + { gf2EntryRow :: !Int, + gf2EntryColumn :: !Int + } + deriving stock (Eq, Ord, Show) + +type GF2PackedMatrix :: Type +type GF2PackedMatrix = PackedBitMatrix + +gf2PackedRows :: GF2PackedMatrix -> Int +gf2PackedRows = + packedRows + +gf2PackedColumns :: GF2PackedMatrix -> Int +gf2PackedColumns = + packedCols + +gf2PackedWordsPerRow :: GF2PackedMatrix -> Int +gf2PackedWordsPerRow = + packedWordsPerRow + +gf2PackedWords :: GF2PackedMatrix -> U.Vector Word64 +gf2PackedWords = + packedWords + +type GF2PackedMatrixFailure :: Type +data GF2PackedMatrixFailure + = GF2PackedMatrixEntryOutOfBounds !Int !Int !Int !Int + | GF2PackedMatrixFlatLengthMismatch !Int !Int + | GF2PackedMatrixCardinalityOutOfBounds !Natural !Natural + deriving stock (Eq, Show) + +mkGF2PackedMatrix :: + Natural -> + Natural -> + [GF2MatrixEntry] -> + Either GF2PackedMatrixFailure GF2PackedMatrix +mkGF2PackedMatrix rowCountValue columnCountValue entries = do + (rowCount, columnCount) <- checkedGF2PackedDimensions rowCountValue columnCountValue + case firstOutOfBoundsEntry rowCount columnCount entries of + Just entryValue -> + Left + ( GF2PackedMatrixEntryOutOfBounds + (gf2EntryRow entryValue) + (gf2EntryColumn entryValue) + rowCount + columnCount + ) + Nothing -> + Right + ( packedBitMatrixFromXorEntries + rowCount + columnCount + (entryCoordinates <$> entries) + ) + +mkGF2PackedMatrixFromRowMajor :: + Natural -> + Natural -> + [GF2] -> + Either GF2PackedMatrixFailure GF2PackedMatrix +mkGF2PackedMatrixFromRowMajor rowCountValue columnCountValue values = do + (rowCount, columnCount) <- checkedGF2PackedDimensions rowCountValue columnCountValue + expectedEntryCount <- + mapGF2CardinalityFailure rowCountValue columnCountValue + (checkedNonNegativeProduct rowCount columnCount) + let actualEntryCount = length values + if actualEntryCount /= expectedEntryCount + then Left (GF2PackedMatrixFlatLengthMismatch expectedEntryCount actualEntryCount) + else + Right + ( packedBitMatrixFromRowMajor + rowCount + columnCount + values + ) + +checkedGF2PackedDimensions :: + Natural -> + Natural -> + Either GF2PackedMatrixFailure (Int, Int) +checkedGF2PackedDimensions rowCountValue columnCountValue = do + rowCount <- + mapGF2CardinalityFailure rowCountValue columnCountValue + (checkedNaturalToInt rowCountValue) + columnCount <- + mapGF2CardinalityFailure rowCountValue columnCountValue + (checkedNaturalToInt columnCountValue) + let wordsPerRow = + columnCount `quot` 64 + + if columnCount `rem` 64 == 0 then 0 else 1 + _ <- + mapGF2CardinalityFailure rowCountValue columnCountValue + (checkedNonNegativeProduct rowCount wordsPerRow) + Right (rowCount, columnCount) + +mapGF2CardinalityFailure :: + Natural -> + Natural -> + Either cardinalityFailure value -> + Either GF2PackedMatrixFailure value +mapGF2CardinalityFailure rowCountValue columnCountValue = + first + (const (GF2PackedMatrixCardinalityOutOfBounds rowCountValue columnCountValue)) + +rankGF2PackedMatrix :: GF2PackedMatrix -> Int +rankGF2PackedMatrix matrixValue = + rankPackedRows + (packedCols matrixValue) + (matrixRowWords matrixValue <$> [0 .. packedRows matrixValue - 1]) + +gf2PackedMatrixLinearMap :: GF2PackedMatrix -> Either MoonlightError PackedLinearMap +gf2PackedMatrixLinearMap matrixValue = do + columnRows <- gf2PackedMatrixColumnRows "gf2PackedMatrixLinearMap" matrixValue + packedLinearMapFromColumns + "gf2PackedMatrixLinearMap" + (packedCols matrixValue) + (packedRows matrixValue) + columnRows + +inverseGF2PackedMatrix :: GF2PackedMatrix -> Either MoonlightError (Maybe PackedLinearMap) +inverseGF2PackedMatrix matrixValue + | packedRows matrixValue /= packedCols matrixValue = + Right Nothing + | otherwise = do + columnRows <- gf2PackedMatrixColumnRows "inverseGF2PackedMatrix" matrixValue + reductionValue <- reducePackedColumns "inverseGF2PackedMatrix" (packedRows matrixValue) columnRows + if V.length (crIndependentIndices reductionValue) == packedCols matrixValue + then Just <$> inverseFromPackedBasisColumns "inverseGF2PackedMatrix" columnRows + else Right Nothing + +gf2PackedMatrixColumnRows :: String -> GF2PackedMatrix -> Either MoonlightError (Vector PackedRow) +gf2PackedMatrixColumnRows context matrixValue = + V.fromList + <$> traverse + ( \columnIndex -> + packedRowFromIndices + (context <> ": column " <> show columnIndex) + (packedRows matrixValue) + (columnSupport columnIndex) + ) + [0 .. packedCols matrixValue - 1] + where + columnSupport columnIndex = + filter + (\rowIndex -> matrixEntryPresent rowIndex columnIndex) + [0 .. packedRows matrixValue - 1] + + matrixEntryPresent rowIndex columnIndex = + let rowWords = matrixRowWords matrixValue rowIndex + wordIndex = columnIndex `div` 64 + bitIndex = columnIndex `mod` 64 + in maybe False (`testBit` bitIndex) (rowWords U.!? wordIndex) + +firstOutOfBoundsEntry :: Int -> Int -> [GF2MatrixEntry] -> Maybe GF2MatrixEntry +firstOutOfBoundsEntry rowCount columnCount = + listToMaybe . filter (not . entryWithinBounds rowCount columnCount) + +entryWithinBounds :: Int -> Int -> GF2MatrixEntry -> Bool +entryWithinBounds rowCount columnCount entry = + gf2EntryRow entry >= 0 + && gf2EntryRow entry < rowCount + && gf2EntryColumn entry >= 0 + && gf2EntryColumn entry < columnCount + +entryCoordinates :: GF2MatrixEntry -> (Int, Int) +entryCoordinates entry = + (gf2EntryRow entry, gf2EntryColumn entry) diff --git a/src-dense/Moonlight/LinAlg/Pure/Dense/Solver.hs b/src-dense/Moonlight/LinAlg/Pure/Dense/Solver.hs new file mode 100644 --- /dev/null +++ b/src-dense/Moonlight/LinAlg/Pure/Dense/Solver.hs @@ -0,0 +1,250 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Moonlight.LinAlg.Pure.Dense.Solver + ( solveDirect, + solveCG, + solveGMRES, + ) +where + +import Control.Monad (foldM) +import Data.Bifunctor (first) +import qualified Data.Vector as Box +import qualified Data.Vector.Unboxed as U +import GHC.TypeNats (KnownNat) +import Moonlight.Core + ( MoonlightError (..), + checkedNonNegativeProduct, + checkedNonNegativeSum, + ) +import Moonlight.LinAlg.Internal.Dense.DoubleFactorization (solveSquareLinearSystem) +import Moonlight.LinAlg.Internal.Primitives + ( ColumnIndex, + RowIndex, + addVector, + dotProduct, + epsilon, + linearCombination, + matrixVectorProduct, + mkColumnIndex, + mkRowIndex, + replaceColumnEntryChecked, + replaceRowChecked, + requireRow, + rowIndices, + scaleVector, + subVector, + vectorNorm, + ) +import Moonlight.LinAlg.Pure.Dense.Types (Matrix, Vector, fromListVector, toListMatrix, toListVector) +import qualified Moonlight.LinAlg.Pure.Dense.Types as DenseTypes +import Prelude + +solveDirect :: + forall n. + KnownNat n => + Matrix n n Double -> + Vector n Double -> + Either MoonlightError (Vector n Double) +solveDirect matrixValue rightHandSide = do + let (dimension, _) = DenseTypes.matrixShape matrixValue + solutionValues <- + solveSquareLinearSystem + dimension + (toListMatrix matrixValue) + (toListVector rightHandSide) + fromListVector @n solutionValues + +solveCG :: + forall n. + KnownNat n => + Matrix n n Double -> + Vector n Double -> + Either MoonlightError (Vector n Double) +solveCG matrixValue rightHandSide = do + matrixRows <- DenseTypes.matrixToRows matrixValue + let (dimension, _) = DenseTypes.matrixShape matrixValue + iterationLimit <- checkedSolverProduct "CG iteration budget" dimension 20 + let rhsValues = toListVector rightHandSide + initialGuess = replicate dimension 0.0 + cgTolerance = 1.0e-10 + boundedIterationLimit = max 1 iterationLimit + initialResidual <- subVector rhsValues =<< matrixVectorProduct matrixRows initialGuess + initialResidualNormSquared <- dotProduct initialResidual initialResidual + let iterateCg iterationIndex guessVector residualVector directionVector residualNormSquared + | sqrt residualNormSquared <= cgTolerance = Right guessVector + | iterationIndex >= boundedIterationLimit = Left (InvariantViolation "CG solver exhausted iteration budget") + | otherwise = do + imageDirection <- matrixVectorProduct matrixRows directionVector + denominator <- dotProduct directionVector imageDirection + directionNorm <- vectorNorm directionVector + imageDirectionNorm <- vectorNorm imageDirection + if nearZeroConjugateGradientDenominator denominator directionNorm imageDirectionNorm + then Left (InvariantViolation "CG solver encountered near-zero denominator") + else do + let alphaValue = residualNormSquared / denominator + nextGuess <- addVector guessVector (scaleVector alphaValue directionVector) + nextResidual <- subVector residualVector (scaleVector alphaValue imageDirection) + nextResidualNormSquared <- dotProduct nextResidual nextResidual + if sqrt nextResidualNormSquared <= cgTolerance + then Right nextGuess + else do + let betaValue = nextResidualNormSquared / residualNormSquared + nextDirection <- addVector nextResidual (scaleVector betaValue directionVector) + iterateCg (iterationIndex + 1) nextGuess nextResidual nextDirection nextResidualNormSquared + solutionValues <- iterateCg 0 initialGuess initialResidual initialResidual initialResidualNormSquared + fromListVector @n solutionValues + +nearZeroConjugateGradientDenominator :: Double -> Double -> Double -> Bool +nearZeroConjugateGradientDenominator denominator directionNorm imageDirectionNorm = + abs denominator <= epsilon * directionNorm * imageDirectionNorm + +updateEntry :: RowIndex -> ColumnIndex -> Double -> [[Double]] -> Either MoonlightError [[Double]] +updateEntry rowIndex columnIndex value matrixRows = do + rowValues <- + requireRow + (InvariantViolation ("GMRES update entry row missing at index " <> show rowIndex)) + rowIndex + matrixRows + updatedRow <- + replaceColumnEntryChecked + (InvariantViolation ("GMRES update entry column missing at index " <> show columnIndex)) + columnIndex + value + rowValues + replaceRowChecked + (InvariantViolation ("GMRES row replacement failed at index " <> show rowIndex)) + rowIndex + updatedRow + matrixRows + +solveHessenbergLeastSquares :: [[Double]] -> [Double] -> Either MoonlightError [Double] +solveHessenbergLeastSquares hRows rhsValues = + let rows = Box.fromList (fmap U.fromList hRows) + rhs = U.fromList rhsValues + colCount = if Box.null rows then 0 else U.length (rows Box.! 0) + givensRotation :: Double -> Double -> (Double, Double) + givensRotation aVal bVal + | abs bVal <= 1.0e-15 = (1.0, 0.0) + | abs aVal <= 1.0e-15 = (0.0, if bVal >= 0 then 1.0 else -1.0) + | otherwise = + let rVal = sqrt (aVal * aVal + bVal * bVal) + in (aVal / rVal, bVal / rVal) + applyRotation :: + Double -> + Double -> + Int -> + Box.Vector (U.Vector Double) -> + U.Vector Double -> + (Box.Vector (U.Vector Double), U.Vector Double) + applyRotation cs sn pivotCol rs b = + let rowI = rs Box.! pivotCol + rowJ = rs Box.! (pivotCol + 1) + rotatedI = U.zipWith (\ri rj -> cs * ri + sn * rj) rowI rowJ + rotatedJ = U.zipWith (\ri rj -> negate sn * ri + cs * rj) rowI rowJ + updatedRs = rs Box.// [(pivotCol, rotatedI), (pivotCol + 1, rotatedJ)] + bI = b U.! pivotCol + bJ = b U.! (pivotCol + 1) + updatedB = b U.// [(pivotCol, cs * bI + sn * bJ), (pivotCol + 1, negate sn * bI + cs * bJ)] + in (updatedRs, updatedB) + eliminateSubdiagonal col (rs, b) + | col >= colCount = (rs, b) + | col + 1 >= Box.length rs = (rs, b) + | otherwise = + let aVal = (rs Box.! col) U.! col + bVal = (rs Box.! (col + 1)) U.! col + (cs, sn) = givensRotation aVal bVal + (nextRs, nextB) = applyRotation cs sn col rs b + in eliminateSubdiagonal (col + 1) (nextRs, nextB) + (triangularRows, transformedRhs) = eliminateSubdiagonal 0 (rows, rhs) + backSolve col solution + | col < 0 = solution + | otherwise = + let rowVec = triangularRows Box.! col + diagVal = rowVec U.! col + trailingLen = colCount - col - 1 + trailingSum = + if trailingLen <= 0 + then 0.0 + else U.sum (U.zipWith (*) (U.slice (col + 1) trailingLen rowVec) (U.slice (col + 1) trailingLen solution)) + rhsVal = transformedRhs U.! col + solvedVal = if abs diagVal <= 1.0e-15 then 0.0 else (rhsVal - trailingSum) / diagVal + in backSolve (col - 1) (solution U.// [(col, solvedVal)]) + in Right (U.toList (backSolve (colCount - 1) (U.replicate colCount 0.0))) + +solveGMRES :: + forall n. + KnownNat n => + Matrix n n Double -> + Vector n Double -> + Either MoonlightError (Vector n Double) +solveGMRES matrixValue rightHandSide = do + matrixRows <- DenseTypes.matrixToRows matrixValue + let (dimension, _) = DenseTypes.matrixShape matrixValue + maxIterations = max 1 dimension + krylovRowCount <- checkedSolverSum "GMRES Krylov row count" maxIterations 1 + _ <- checkedSolverProduct "GMRES Hessenberg workspace" krylovRowCount maxIterations + let rhsValues = toListVector rightHandSide + gmresTolerance = 1.0e-10 + betaValue <- vectorNorm rhsValues + if betaValue <= gmresTolerance + then fromListVector @n (replicate dimension 0.0) + else do + let firstBasis = scaleVector (1.0 / betaValue) rhsValues + initialHessenberg = replicate krylovRowCount (replicate maxIterations 0.0) + arnoldiStep iterationIndex basisVectors hessenbergRows + | iterationIndex >= maxIterations = Right (basisVectors, hessenbergRows, maxIterations) + | otherwise = do + iterationBasisIndex <- + mkRowIndex + (InvariantViolation ("GMRES basis index out of bounds at iteration " <> show iterationIndex)) + (length basisVectors) + iterationIndex + iterationColumnIndex <- + mkColumnIndex + (InvariantViolation ("GMRES Hessenberg column out of bounds at iteration " <> show iterationIndex)) + maxIterations + iterationIndex + currentBasis <- + requireRow + (InvariantViolation ("GMRES basis lookup failed at iteration " <> show iterationBasisIndex)) + iterationBasisIndex + basisVectors + initialVector <- matrixVectorProduct matrixRows currentBasis + let orthogonalize (workingVector, workingHessenberg) (basisIndex, basisVectorValue) = do + coefficient <- dotProduct basisVectorValue workingVector + nextVector <- subVector workingVector (scaleVector coefficient basisVectorValue) + nextHessenberg <- updateEntry basisIndex iterationColumnIndex coefficient workingHessenberg + Right (nextVector, nextHessenberg) + (reducedVector, filledHessenberg) <- + foldM orthogonalize (initialVector, hessenbergRows) (zip (take (iterationIndex + 1) (rowIndices (length basisVectors))) basisVectors) + nextNorm <- vectorNorm reducedVector + nextRowIndex <- + mkRowIndex + (InvariantViolation ("GMRES Hessenberg row out of bounds at iteration " <> show (iterationIndex + 1))) + (length filledHessenberg) + (iterationIndex + 1) + completedHessenberg <- updateEntry nextRowIndex iterationColumnIndex nextNorm filledHessenberg + if nextNorm <= gmresTolerance + then Right (basisVectors, completedHessenberg, iterationIndex + 1) + else + let nextBasis = basisVectors <> [scaleVector (1.0 / nextNorm) reducedVector] + in arnoldiStep (iterationIndex + 1) nextBasis completedHessenberg + (basisVectors, hessenbergRows, iterationCount) <- arnoldiStep 0 [firstBasis] initialHessenberg + let reducedHessenberg = map (take iterationCount) (take (iterationCount + 1) hessenbergRows) + leastSquaresRhs = betaValue : replicate iterationCount 0.0 + coefficients <- solveHessenbergLeastSquares reducedHessenberg leastSquaresRhs + solutionValues <- linearCombination (zip coefficients (take iterationCount basisVectors)) + fromListVector @n solutionValues + +checkedSolverProduct :: String -> Int -> Int -> Either MoonlightError Int +checkedSolverProduct context leftFactor rightFactor = + first + (const (InvariantViolation (context <> " exceeds Int range"))) + (checkedNonNegativeProduct leftFactor rightFactor) + +checkedSolverSum :: String -> Int -> Int -> Either MoonlightError Int +checkedSolverSum context leftTerm rightTerm = + first + (const (InvariantViolation (context <> " exceeds Int range"))) + (checkedNonNegativeSum leftTerm rightTerm) diff --git a/src-domain/Moonlight/LinAlg/Pure/Domain/Bareiss.hs b/src-domain/Moonlight/LinAlg/Pure/Domain/Bareiss.hs new file mode 100644 --- /dev/null +++ b/src-domain/Moonlight/LinAlg/Pure/Domain/Bareiss.hs @@ -0,0 +1,354 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} + +module Moonlight.LinAlg.Pure.Domain.Bareiss + ( BareissExactDivision (..), + BareissExactDivisionObligation (..), + BareissElimination (..), + euclideanBareissExactDivision, + bareissEliminationWith, + bareissElimination, + bareissRankWith, + bareissRank, + bareissDeterminantWith, + bareissDeterminant, + bareissEchelonWith, + bareissEchelon, + ) +where + +import Control.Monad (foldM) +import Data.Kind (Type) +import Data.Vector qualified as Box +import GHC.TypeNats (KnownNat, Nat) +import Moonlight.Algebra.Pure.Ring + ( EuclideanDomain (..), + IntegralDomain (..), + mkNonZeroDivisor, + ) +import Moonlight.Core + ( AdditiveGroup (..), + AdditiveMonoid (..), + MoonlightError (..), + MultiplicativeMonoid (..), + ) +import Moonlight.LinAlg.Internal.Backend.RowStore + ( RowStore, + rowStoreFlatten, + rowStoreFromRows, + rowStoreRowAtInt, + rowStoreShape, + rowStoreValueAtInt, + swapRowsStoreAtInt, + traverseRowStoreWithIndex, + ) +import Moonlight.LinAlg.Pure.Dense.Types + ( Matrix, + fromListMatrix, + ) +import qualified Moonlight.LinAlg.Pure.Dense.Types as DenseTypes +import Prelude + +type BareissExactDivisionObligation :: Type -> Type +data BareissExactDivisionObligation a = BareissExactDivisionObligation + { bareissExactDivisionStep :: !Int, + bareissExactDivisionRow :: !Int, + bareissExactDivisionColumn :: !Int, + bareissExactDividend :: a, + bareissExactDivisor :: a + } + +type BareissExactDivision :: Type -> Type +newtype BareissExactDivision a = BareissExactDivision + { runBareissExactDivision :: BareissExactDivisionObligation a -> Either MoonlightError a + } + +type BareissElimination :: Nat -> Nat -> Type -> Type +data BareissElimination r c a = BareissElimination + { bareissResultRank :: !Int, + bareissResultDeterminant :: !(Maybe a), + bareissResultEchelon :: Matrix r c a + } + +type BareissState :: Type -> Type +data BareissState a = BareissState + { bareissStateRows :: RowStore a, + bareissStateRank :: !Int, + bareissStatePreviousPivot :: a, + bareissStateDetSign :: a + } + +euclideanBareissExactDivision :: EuclideanDomain a => BareissExactDivision a +euclideanBareissExactDivision = + BareissExactDivision $ \obligation -> + case mkNonZeroDivisor (bareissExactDivisor obligation) of + Nothing -> Left (InvariantViolation "Bareiss exact division received a zero divisor") + Just divisor -> + let (quotientValue, remainderValue) = + divideWithRemainder + (bareissExactDividend obligation) + divisor + in if isZero remainderValue + then Right quotientValue + else + Left + ( InvariantViolation + ( "Bareiss exact division obligation failed at " + <> show + ( bareissExactDivisionStep obligation, + bareissExactDivisionRow obligation, + bareissExactDivisionColumn obligation + ) + ) + ) + +bareissEliminationWith :: + forall r c a. + (KnownNat r, KnownNat c, IntegralDomain a) => + BareissExactDivision a -> + Matrix r c a -> + Either MoonlightError (BareissElimination r c a) +bareissEliminationWith division matrixValue = do + initialRows <- DenseTypes.matrixToRows matrixValue + let (rowCount, columnCount) = DenseTypes.matrixShape matrixValue + initialState = + BareissState + { bareissStateRows = rowStoreFromRows initialRows, + bareissStateRank = 0, + bareissStatePreviousPivot = one, + bareissStateDetSign = one + } + finalState <- foldlBareiss division columnCount initialState + echelonMatrix <- fromListMatrix @r @c (rowStoreFlatten (bareissStateRows finalState)) + pure + BareissElimination + { bareissResultRank = bareissStateRank finalState, + bareissResultDeterminant = determinantValue rowCount columnCount finalState, + bareissResultEchelon = echelonMatrix + } + +bareissElimination :: + forall r c a. + (KnownNat r, KnownNat c, EuclideanDomain a) => + Matrix r c a -> + Either MoonlightError (BareissElimination r c a) +bareissElimination = + bareissEliminationWith euclideanBareissExactDivision + +bareissRankWith :: + forall r c a. + (KnownNat r, KnownNat c, IntegralDomain a) => + BareissExactDivision a -> + Matrix r c a -> + Either MoonlightError Int +bareissRankWith division = + fmap bareissResultRank . bareissEliminationWith division + +bareissRank :: + forall r c a. + (KnownNat r, KnownNat c, EuclideanDomain a) => + Matrix r c a -> + Either MoonlightError Int +bareissRank = + bareissRankWith euclideanBareissExactDivision + +bareissDeterminantWith :: + forall n a. + (KnownNat n, IntegralDomain a) => + BareissExactDivision a -> + Matrix n n a -> + Either MoonlightError a +bareissDeterminantWith division matrixValue = + bareissEliminationWith division matrixValue + >>= \result -> + case bareissResultDeterminant result of + Just determinant -> Right determinant + Nothing -> Left (InvariantViolation "Bareiss square determinant produced no determinant") + +bareissDeterminant :: + forall n a. + (KnownNat n, EuclideanDomain a) => + Matrix n n a -> + Either MoonlightError a +bareissDeterminant = + bareissDeterminantWith euclideanBareissExactDivision + +bareissEchelonWith :: + forall r c a. + (KnownNat r, KnownNat c, IntegralDomain a) => + BareissExactDivision a -> + Matrix r c a -> + Either MoonlightError (Matrix r c a) +bareissEchelonWith division = + fmap bareissResultEchelon . bareissEliminationWith division + +bareissEchelon :: + forall r c a. + (KnownNat r, KnownNat c, EuclideanDomain a) => + Matrix r c a -> + Either MoonlightError (Matrix r c a) +bareissEchelon = + bareissEchelonWith euclideanBareissExactDivision + +foldlBareiss :: + IntegralDomain a => + BareissExactDivision a -> + Int -> + BareissState a -> + Either MoonlightError (BareissState a) +foldlBareiss division columnCount initialState = + foldM + (\stateValue columnIndex -> bareissStep division columnIndex stateValue) + initialState + [0 .. columnCount - 1] + +bareissStep :: + IntegralDomain a => + BareissExactDivision a -> + Int -> + BareissState a -> + Either MoonlightError (BareissState a) +bareissStep division pivotColumn stateValue + | bareissStateRank stateValue >= rowCount = Right stateValue + | otherwise = do + pivotCandidate <- findPivotRowAt pivotRow pivotColumn rows + case pivotCandidate of + Nothing -> Right stateValue + Just sourceRow -> do + swappedRows <- + if sourceRow == pivotRow + then Right rows + else + swapRowsStoreAtInt + (InvariantViolation ("Bareiss row pivot swap failed at " <> show (pivotRow, sourceRow))) + pivotRow + sourceRow + rows + pivotValue <- + rowStoreValueAtInt + (InvariantViolation ("Bareiss pivot lookup failed at " <> show (pivotRow, pivotColumn))) + pivotRow + pivotColumn + swappedRows + pivotRowValues <- + rowStoreRowAtInt + (InvariantViolation ("Bareiss pivot row lookup failed at " <> show pivotRow)) + pivotRow + swappedRows + eliminatedRows <- + eliminateBareissColumn + division + pivotRow + pivotColumn + pivotValue + (bareissStatePreviousPivot stateValue) + pivotRowValues + swappedRows + Right + stateValue + { bareissStateRows = eliminatedRows, + bareissStateRank = pivotRow + 1, + bareissStatePreviousPivot = pivotValue, + bareissStateDetSign = + if sourceRow == pivotRow + then bareissStateDetSign stateValue + else neg (bareissStateDetSign stateValue) + } + where + rows = bareissStateRows stateValue + (rowCount, _) = rowStoreShape rows + pivotRow = bareissStateRank stateValue + +findPivotRowAt :: IntegralDomain a => Int -> Int -> RowStore a -> Either MoonlightError (Maybe Int) +findPivotRowAt startRow columnIndex rows = + fmap + (fmap fst . firstNonzero) + ( traverse + ( \rowIndex -> + fmap + (\value -> (rowIndex, value)) + (rowStoreValueAtInt (InvariantViolation ("Bareiss pivot search failed at " <> show (rowIndex, columnIndex))) rowIndex columnIndex rows) + ) + [startRow .. fst (rowStoreShape rows) - 1] + ) + where + firstNonzero :: IntegralDomain a => [(Int, a)] -> Maybe (Int, a) + firstNonzero = + foldr + ( \candidate rest -> + if isZero (snd candidate) + then rest + else Just candidate + ) + Nothing + +eliminateBareissColumn :: + IntegralDomain a => + BareissExactDivision a -> + Int -> + Int -> + a -> + a -> + Box.Vector a -> + RowStore a -> + Either MoonlightError (RowStore a) +eliminateBareissColumn division pivotRow pivotColumn pivotValue previousPivot pivotRowValues rows = + traverseRowStoreWithIndex transformRow rows + where + transformRow rowIndex rowValues + | rowIndex <= pivotRow = Right rowValues + | otherwise = do + targetPivotEntry <- + maybe + (Left (InvariantViolation ("Bareiss target pivot lookup failed at " <> show (rowIndex, pivotColumn)))) + Right + (rowValues Box.!? pivotColumn) + updatedRow <- + Box.imapM + ( \columnIndex entryValue -> + if columnIndex < pivotColumn + then Right entryValue + else + if columnIndex == pivotColumn + then Right zero + else do + pivotRowEntry <- + maybe + (Left (InvariantViolation ("Bareiss pivot row lookup failed at column " <> show columnIndex))) + Right + (pivotRowValues Box.!? columnIndex) + divideBareissEntry + division + pivotRow + rowIndex + columnIndex + ((pivotValue `mul` entryValue) `sub` (targetPivotEntry `mul` pivotRowEntry)) + previousPivot + ) + rowValues + Right updatedRow + +divideBareissEntry :: + BareissExactDivision a -> + Int -> + Int -> + Int -> + a -> + a -> + Either MoonlightError a +divideBareissEntry division step rowIndex columnIndex dividend divisor = + runBareissExactDivision + division + BareissExactDivisionObligation + { bareissExactDivisionStep = step, + bareissExactDivisionRow = rowIndex, + bareissExactDivisionColumn = columnIndex, + bareissExactDividend = dividend, + bareissExactDivisor = divisor + } + +determinantValue :: IntegralDomain a => Int -> Int -> BareissState a -> Maybe a +determinantValue rowCount columnCount stateValue + | rowCount /= columnCount = Nothing + | rowCount == 0 = Just one + | bareissStateRank stateValue < rowCount = Just zero + | otherwise = Just (bareissStateDetSign stateValue `mul` bareissStatePreviousPivot stateValue) diff --git a/src-domain/Moonlight/LinAlg/Pure/Domain/Smith.hs b/src-domain/Moonlight/LinAlg/Pure/Domain/Smith.hs new file mode 100644 --- /dev/null +++ b/src-domain/Moonlight/LinAlg/Pure/Domain/Smith.hs @@ -0,0 +1,35 @@ +module Moonlight.LinAlg.Pure.Domain.Smith + ( SmithNormalForm (..), + SmithDiagonalForm (..), + smithNormalForm, + smithDiagonalForm, + ) +where + +import GHC.TypeNats (KnownNat) +import Moonlight.Algebra.Pure.Ring (EuclideanDomain) +import Moonlight.Core (MoonlightError) +import Moonlight.LinAlg.Internal.Backend.Core (runSmithDiagonalForm, runSmithNormalForm) +import Moonlight.LinAlg.Internal.Backend.Smith (SmithDiagonalForm (..), SmithNormalForm (..)) +import Moonlight.LinAlg.Pure.Domain.Smith.Multimodular (smithDiagonalFormMultimodular) +import Moonlight.LinAlg.Pure.Domain.Smith.Witnessed (smithNormalFormWitnessed) +import Moonlight.LinAlg.Pure.Dense.Types (Matrix) + +smithNormalForm :: + forall r c a. + (KnownNat r, KnownNat c, EuclideanDomain a) => + Matrix r c a -> + Either MoonlightError (SmithNormalForm r c a) +smithNormalForm = runSmithNormalForm + +smithDiagonalForm :: + forall r c a. + (KnownNat r, KnownNat c, EuclideanDomain a) => + Matrix r c a -> + Either MoonlightError (SmithDiagonalForm r c a) +smithDiagonalForm = runSmithDiagonalForm + +{-# NOINLINE smithNormalForm #-} +{-# NOINLINE smithDiagonalForm #-} +{-# RULES "smithNormalForm/Integer" forall matrixValue. smithNormalForm matrixValue = smithNormalFormWitnessed matrixValue #-} +{-# RULES "smithDiagonalForm/Integer" forall matrixValue. smithDiagonalForm matrixValue = smithDiagonalFormMultimodular matrixValue #-} diff --git a/src-domain/Moonlight/LinAlg/Pure/Domain/Smith/Multimodular.hs b/src-domain/Moonlight/LinAlg/Pure/Domain/Smith/Multimodular.hs new file mode 100644 --- /dev/null +++ b/src-domain/Moonlight/LinAlg/Pure/Domain/Smith/Multimodular.hs @@ -0,0 +1,1136 @@ +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} + +module Moonlight.LinAlg.Pure.Domain.Smith.Multimodular + ( PrimeSweep (..), + certifiedPrimeSweep, + integerResidueWord, + modInverseWord, + modMul, + smithDiagonalFormMultimodular, + wordPrimeLadder, + ) +where + +import Control.Monad.ST (ST, runST) +import Data.Bifunctor (first) +import Data.List (sortBy) +import Data.Vector qualified as V +import Data.Vector.Mutable qualified as MV +import Data.Vector.Unboxed qualified as U +import Data.Vector.Unboxed.Mutable qualified as MU +import Data.Word (Word64) +import GHC.Exts (quotRemWord2#, timesWord2#, word64ToWord#, wordToWord64#) +import GHC.TypeNats (KnownNat) +import GHC.Word (Word64 (W64#)) +import Moonlight.Algebra.Pure.Ring (EuclideanDomain (..), GCDDomain (..), mkNonZeroDivisor) +import Moonlight.Core + ( AdditiveGroup (..), + AdditiveMonoid (..), + MoonlightError (..), + MultiplicativeMonoid (..), + checkedNonNegativeProduct, + ) +import Moonlight.LinAlg.Internal.Backend.Smith (SmithDiagonalForm (..)) +import Moonlight.LinAlg.Pure.Dense.Types + ( Matrix, + fromListMatrix, + ) +import qualified Moonlight.LinAlg.Pure.Dense.Types as DenseTypes +import Prelude + +data PrimeSweep = PrimeSweep + { primeSweepRank :: !Int, + primeSweepDeterminant :: !(Maybe Integer) + } + deriving stock (Eq, Show) + +data CrtState = CrtState + { crtResidue :: !Integer, + crtModulus :: !Integer, + crtRank :: !Int + } + deriving stock (Eq, Show) + +data PrimeElimination = PrimeElimination + { primeEliminationRank :: !Int, + primeEliminationDeterminant :: !Word64 + } + deriving stock (Eq, Show) + +data SmithTier + = SmithTierWord32 + | SmithTierWord62 + | SmithTierInteger + deriving stock (Eq, Show) + +data SmithCarrier s + = SmithWord32Carrier !Word64 !(MU.MVector s Word64) + | SmithWord62Carrier !Word64 !(MU.MVector s Word64) + | SmithIntegerCarrier !(MV.MVector s Integer) + +data MutableSmithState s = MutableSmithState + { mutableSmithRowCount :: !Int, + mutableSmithColumnCount :: !Int, + mutableSmithModulus :: !Integer, + mutableSmithCarrier :: !(SmithCarrier s) + } + +data SmithPivot = SmithPivot + { pivotRowIndex :: !Int, + pivotColumnIndex :: !Int + } + deriving stock (Eq, Show) + +data SmithPhaseFailure + = SmithPhaseBudgetExhausted !String + | SmithPhaseNormalizationStalled + | SmithPhasePivotBecameZero + | SmithPhaseInexactDivision !String + deriving stock (Eq, Show) + +data SmithPhaseResult + = SmithPhaseDone ![Integer] + | SmithPhaseFailed !SmithPhaseFailure + +smithDiagonalFormMultimodular :: + forall r c. + (KnownNat r, KnownNat c) => + Matrix r c Integer -> + Either MoonlightError (SmithDiagonalForm r c Integer) +smithDiagonalFormMultimodular matrixValue = do + rows <- DenseTypes.matrixToRows matrixValue + let (rowCount, columnCount) = DenseTypes.matrixShape matrixValue + diagonalSize = min rowCount columnCount + primeSweep <- certifiedPrimeSweep rowCount columnCount rows + invariantFactors <- + certifiedInvariantFactors + rowCount + columnCount + diagonalSize + rows + primeSweep + diagonalMatrix <- fromListMatrix @r @c (diagonalFlatEntries rowCount columnCount invariantFactors) + pure (SmithDiagonalForm diagonalMatrix) + +certifiedInvariantFactors :: + Int -> + Int -> + Int -> + [[Integer]] -> + PrimeSweep -> + Either MoonlightError [Integer] +certifiedInvariantFactors rowCount columnCount diagonalSize rows primeSweep + | diagonalSize == 0 = Right [] + | primeSweepRank primeSweep == 0 = Right [] + | rowCount == columnCount && primeSweepRank primeSweep == rowCount = + case primeSweepDeterminant primeSweep of + Just determinantValue -> smithNonsingularInvariantFactors determinantValue rows + Nothing -> Left (InvariantViolation "Smith multimodular square prime sweep did not return a determinant") + | otherwise = smithCompressedInvariantFactors rowCount columnCount (primeSweepRank primeSweep) rows + +certifiedPrimeSweep :: Int -> Int -> [[Integer]] -> Either MoonlightError PrimeSweep +certifiedPrimeSweep rowCount columnCount rows = + let determinantBound = hadamardMinorBound rowCount columnCount rows + target = max 2 (2 * determinantBound + 1) + in finishPrimeSweep rowCount columnCount target rows initialCrtState wordPrimeLadder + where + initialCrtState = + CrtState + { crtResidue = 0, + crtModulus = 1, + crtRank = 0 + } + +finishPrimeSweep :: Int -> Int -> Integer -> [[Integer]] -> CrtState -> [Word64] -> Either MoonlightError PrimeSweep +finishPrimeSweep rowCount columnCount target rows stateValue primes + | crtModulus stateValue > target = + Right + PrimeSweep + { primeSweepRank = crtRank stateValue, + primeSweepDeterminant = + if rowCount == columnCount + then Just (symmetricLift (crtModulus stateValue) (crtResidue stateValue)) + else Nothing + } + | otherwise = + case primes of + [] -> Left (InvariantViolation "Smith multimodular prime ladder exhausted before Hadamard certification") + primeValue : remainingPrimes -> do + let primeMatrix = residueVectorForPrime primeValue rows + primeResult = primeElimination rowCount columnCount primeValue primeMatrix + updatedState <- extendCrt rowCount columnCount stateValue primeValue primeResult + finishPrimeSweep rowCount columnCount target rows updatedState remainingPrimes + +extendCrt :: Int -> Int -> CrtState -> Word64 -> PrimeElimination -> Either MoonlightError CrtState +extendCrt rowCount columnCount stateValue primeValue primeResult = do + let primeInteger = toInteger primeValue + nextRank = max (crtRank stateValue) (primeEliminationRank primeResult) + nextResidue <- + if rowCount == columnCount + then combineCrt (crtResidue stateValue) (crtModulus stateValue) primeInteger (toInteger (primeEliminationDeterminant primeResult)) + else Right (crtResidue stateValue) + Right + CrtState + { crtResidue = nextResidue, + crtModulus = crtModulus stateValue * primeInteger, + crtRank = nextRank + } + +combineCrt :: Integer -> Integer -> Integer -> Integer -> Either MoonlightError Integer +combineCrt residueValue modulusValue primeValue primeResidue = do + inverseValue <- modularInverseInteger (modulusValue `mod` primeValue) primeValue + let deltaValue = (primeResidue - residueValue) `mod` primeValue + correction = (deltaValue * inverseValue) `mod` primeValue + nextModulus = modulusValue * primeValue + Right ((residueValue + modulusValue * correction) `mod` nextModulus) + +modularInverseInteger :: Integer -> Integer -> Either MoonlightError Integer +modularInverseInteger value modulusValue = + let (gcdValue, coefficient, _) = extendedGcdDomain value modulusValue + in if gcdValue == one + then Right (coefficient `mod` modulusValue) + else Left (InvariantViolation "Smith multimodular CRT encountered a noninvertible modulus section") + +symmetricLift :: Integer -> Integer -> Integer +symmetricLift modulusValue residueValue + | 2 * residueValue > modulusValue = residueValue - modulusValue + | otherwise = residueValue + +hadamardMinorBound :: Int -> Int -> [[Integer]] -> Integer +hadamardMinorBound rowCount columnCount rows = + let minorDimension = min rowCount columnCount + squaredNorms = fmap rowSquaredNorm rows + in integerCeilingSquareRoot (product (takeLargest minorDimension squaredNorms)) + +rowSquaredNorm :: [Integer] -> Integer +rowSquaredNorm = + foldl' (\total entry -> total + entry * entry) 0 + +takeLargest :: Int -> [Integer] -> [Integer] +takeLargest count = + take count . sortBy (flip compare) + +integerCeilingSquareRoot :: Integer -> Integer +integerCeilingSquareRoot value + | value <= 0 = 0 + | otherwise = + let rootValue = integerSquareRoot value + in if rootValue * rootValue == value + then rootValue + else rootValue + 1 + +integerSquareRoot :: Integer -> Integer +integerSquareRoot value = + go 0 (value + 1) + where + go low high + | high - low <= 1 = low + | midpoint * midpoint <= value = go midpoint high + | otherwise = go low midpoint + where + midpoint = (low + high) `quot` 2 + +wordPrimeLadder :: [Word64] +wordPrimeLadder = + [ 2147483647, + 2147483629, + 2147483587, + 2147483579, + 2147483563, + 2147483549, + 2147483543, + 2147483497, + 2147483489, + 2147483477, + 2147483423, + 2147483399, + 2147483353, + 2147483323, + 2147483269, + 2147483249, + 2147483237, + 2147483179, + 2147483171, + 2147483137, + 2147483123, + 2147483077, + 2147483069, + 2147483059, + 2147483053, + 2147483033, + 2147483029, + 2147482951, + 2147482949, + 2147482943, + 2147482937, + 2147482921, + 2147482877, + 2147482873, + 2147482819, + 2147482817, + 2147482811, + 2147482801, + 2147482763, + 2147482739, + 2147482697, + 2147482693, + 2147482681, + 2147482663, + 2147482661, + 2147482621, + 2147482591, + 2147482589, + 2147482577, + 2147482507, + 2147482501, + 2147482481, + 2147482417, + 2147482409, + 2147482367, + 2147482361, + 2147482349, + 2147482343, + 2147482327, + 2147482297, + 2147482291, + 2147482273, + 2147482237, + 2147482231 + ] + +residueVectorForPrime :: Word64 -> [[Integer]] -> U.Vector Word64 +residueVectorForPrime primeValue rows = + U.fromList (concatMap (fmap (integerResidueWord primeValue)) rows) + +integerResidueWord :: Word64 -> Integer -> Word64 +integerResidueWord primeValue entryValue = + fromInteger (entryValue `mod` toInteger primeValue) + +primeElimination :: Int -> Int -> Word64 -> U.Vector Word64 -> PrimeElimination +primeElimination rowCount columnCount primeValue entries = + runST $ do + work <- U.thaw entries + let readEntry rowIndex columnIndex = + MU.read work (flatIndex columnCount rowIndex columnIndex) + writeEntry rowIndex columnIndex entryValue = + MU.write work (flatIndex columnCount rowIndex columnIndex) entryValue + swapRows leftRow rightRow = + swapRowEntries readEntry writeEntry columnCount leftRow rightRow 0 + eliminateRows pivotRow pivotColumn inversePivot rowIndex + | rowIndex >= rowCount = pure () + | otherwise = do + entryValue <- readEntry rowIndex pivotColumn + if entryValue == 0 + then eliminateRows pivotRow pivotColumn inversePivot (rowIndex + 1) + else do + let factor = modMul primeValue entryValue inversePivot + eliminateRowEntries readEntry writeEntry primeValue pivotRow rowIndex pivotColumn factor columnCount + eliminateRows pivotRow pivotColumn inversePivot (rowIndex + 1) + step rankValue columnIndex determinantProduct signNegative + | columnIndex >= columnCount || rankValue >= rowCount = + pure (rankValue, determinantProduct, signNegative) + | otherwise = do + pivotCandidate <- findWordPivot readEntry rowCount rankValue columnIndex + case pivotCandidate of + Nothing -> step rankValue (columnIndex + 1) determinantProduct signNegative + Just pivotRow -> do + swapRows rankValue pivotRow + pivotValue <- readEntry rankValue columnIndex + let nextSignNegative = if pivotRow == rankValue then signNegative else not signNegative + nextDeterminant = modMul primeValue determinantProduct pivotValue + inversePivot = modInverseWord primeValue pivotValue + eliminateRows rankValue columnIndex inversePivot (rankValue + 1) + step (rankValue + 1) (columnIndex + 1) nextDeterminant nextSignNegative + (rankValue, determinantProduct, signNegative) <- step 0 0 1 False + let determinantValue = + if rowCount == columnCount && rankValue == rowCount + then if signNegative then modNeg primeValue determinantProduct else determinantProduct + else 0 + pure + PrimeElimination + { primeEliminationRank = rankValue, + primeEliminationDeterminant = determinantValue + } + +flatIndex :: Int -> Int -> Int -> Int +flatIndex columnCount rowIndex columnIndex = + rowIndex * columnCount + columnIndex + +findWordPivot :: (Int -> Int -> ST s Word64) -> Int -> Int -> Int -> ST s (Maybe Int) +findWordPivot readEntry rowCount startRow columnIndex = + go startRow + where + go rowIndex + | rowIndex >= rowCount = pure Nothing + | otherwise = do + entryValue <- readEntry rowIndex columnIndex + if entryValue == 0 + then go (rowIndex + 1) + else pure (Just rowIndex) + +swapRowEntries :: (Int -> Int -> ST s Word64) -> (Int -> Int -> Word64 -> ST s ()) -> Int -> Int -> Int -> Int -> ST s () +swapRowEntries readEntry writeEntry columnCount leftRow rightRow columnIndex + | leftRow == rightRow = pure () + | columnIndex >= columnCount = pure () + | otherwise = do + leftValue <- readEntry leftRow columnIndex + rightValue <- readEntry rightRow columnIndex + writeEntry leftRow columnIndex rightValue + writeEntry rightRow columnIndex leftValue + swapRowEntries readEntry writeEntry columnCount leftRow rightRow (columnIndex + 1) + +eliminateRowEntries :: (Int -> Int -> ST s Word64) -> (Int -> Int -> Word64 -> ST s ()) -> Word64 -> Int -> Int -> Int -> Word64 -> Int -> ST s () +eliminateRowEntries readEntry writeEntry primeValue pivotRow targetRow columnIndex factor columnCount + | columnIndex >= columnCount = pure () + | otherwise = do + targetValue <- readEntry targetRow columnIndex + pivotValue <- readEntry pivotRow columnIndex + let updatedValue = modSub primeValue targetValue (modMul primeValue factor pivotValue) + writeEntry targetRow columnIndex updatedValue + eliminateRowEntries readEntry writeEntry primeValue pivotRow targetRow (columnIndex + 1) factor columnCount + +modSub :: Word64 -> Word64 -> Word64 -> Word64 +modSub primeValue leftValue rightValue + | leftValue >= rightValue = leftValue - rightValue + | otherwise = primeValue - (rightValue - leftValue) + +modMul :: Word64 -> Word64 -> Word64 -> Word64 +modMul primeValue leftValue rightValue = + (leftValue * rightValue) `rem` primeValue + +modNeg :: Word64 -> Word64 -> Word64 +modNeg primeValue value + | value == 0 = 0 + | otherwise = primeValue - value + +modInverseWord :: Word64 -> Word64 -> Word64 +modInverseWord primeValue value = + modPow primeValue value (primeValue - 2) + +modPow :: Word64 -> Word64 -> Word64 -> Word64 +modPow primeValue baseValue exponentValue = + go baseValue exponentValue 1 + where + go currentBase currentExponent accumulator + | currentExponent == 0 = accumulator + | odd currentExponent = go (modMul primeValue currentBase currentBase) (currentExponent `quot` 2) (modMul primeValue accumulator currentBase) + | otherwise = go (modMul primeValue currentBase currentBase) (currentExponent `quot` 2) accumulator + +smithNonsingularInvariantFactors :: Integer -> [[Integer]] -> Either MoonlightError [Integer] +smithNonsingularInvariantFactors determinantValue rows + | determinantValue == zero = Left (InvariantViolation "Smith multimodular nonsingular phase received a zero determinant") + | otherwise = do + let rowCount = length rows + columnCount = firstRowLength rows + modulusValue = 2 * abs determinantValue + validateSmithPhaseCardinalities rowCount columnCount + case runSmithPhase rowCount columnCount modulusValue rows of + SmithPhaseFailed failureValue -> Left (smithPhaseFailureError failureValue) + SmithPhaseDone diagonalValues -> do + let invariantFactors = fmap (smithFactorFromResidue modulusValue) diagonalValues + certifyNonsingularFactors determinantValue invariantFactors + Right invariantFactors + +smithFactorFromResidue :: Integer -> Integer -> Integer +smithFactorFromResidue modulusValue residueValue = + abs (gcd residueValue modulusValue) + +certifyNonsingularFactors :: Integer -> [Integer] -> Either MoonlightError () +certifyNonsingularFactors determinantValue invariantFactors + | product invariantFactors /= abs determinantValue = + Left (InvariantViolation "Smith multimodular determinant-modulus factors failed determinant product certification") + | otherwise = certifyDivisibilityFactors invariantFactors + +certifyDivisibilityFactors :: [Integer] -> Either MoonlightError () +certifyDivisibilityFactors values = + case values of + [] -> Right () + [_] -> Right () + leftValue : rightValue : restValues -> + if leftValue == zero || maybe False ((== zero) . snd) (divideIntegerMaybe rightValue leftValue) + then certifyDivisibilityFactors (rightValue : restValues) + else Left (InvariantViolation "Smith multimodular invariant factors violate the divisibility chain") + +smithCompressedInvariantFactors :: Int -> Int -> Int -> [[Integer]] -> Either MoonlightError [Integer] +smithCompressedInvariantFactors rowCount columnCount certifiedRank rows = do + validateSmithPhaseCardinalities rowCount columnCount + case runSmithPhase rowCount columnCount 0 rows of + SmithPhaseFailed failureValue -> Left (smithPhaseFailureError failureValue) + SmithPhaseDone diagonalValues -> do + let compressedFactors = filter (/= zero) (fmap abs diagonalValues) + if length compressedFactors /= certifiedRank + then Left (InvariantViolation "Smith multimodular rank certificate disagreed with Hermite compression") + else + case compressedFactors of + [] -> Right [] + _ -> smithNonsingularInvariantFactors (product compressedFactors) (diagonalCoreRows compressedFactors) + +validateSmithPhaseCardinalities :: Int -> Int -> Either MoonlightError () +validateSmithPhaseCardinalities rowCount columnCount = do + matrixEntryCount <- checkedSmithPhaseProduct "matrix entries" rowCount columnCount + _ <- checkedSmithPhaseProduct "normalization budget" matrixEntryCount 2 + _ <- checkedSmithPhaseProduct "divisibility-chain budget" (min rowCount columnCount) (min rowCount columnCount) + Right () + +checkedSmithPhaseProduct :: String -> Int -> Int -> Either MoonlightError Int +checkedSmithPhaseProduct context leftFactor rightFactor = + first + (const (InvariantViolation ("Smith multimodular " <> context <> " exceed Int cardinality"))) + (checkedNonNegativeProduct leftFactor rightFactor) + +smithPhaseFailureError :: SmithPhaseFailure -> MoonlightError +smithPhaseFailureError failureValue = + case failureValue of + SmithPhaseBudgetExhausted context -> + InvariantViolation ("Smith multimodular " <> context <> " exhausted iteration budget") + SmithPhaseNormalizationStalled -> + InvariantViolation "Smith multimodular normalization stalled before reaching diagonal form" + SmithPhasePivotBecameZero -> + InvariantViolation "Smith multimodular pivot became zero during reduction" + SmithPhaseInexactDivision context -> + InvariantViolation ("Smith multimodular exact quotient had nonzero remainder during " <> context) + +runSmithPhase :: Int -> Int -> Integer -> [[Integer]] -> SmithPhaseResult +runSmithPhase rowCount columnCount modulusValue rows = + runST $ do + stateValue <- newMutableSmithState rowCount columnCount modulusValue rows + smithFailure <- smithStepMutable 0 stateValue + case smithFailure of + Just failureValue -> pure (SmithPhaseFailed failureValue) + Nothing -> do + chainFailure <- enforceDivisibilityChainMutable stateValue + case chainFailure of + Just failureValue -> pure (SmithPhaseFailed failureValue) + Nothing -> SmithPhaseDone <$> readDiagonalMutable (min rowCount columnCount) stateValue + +newMutableSmithState :: Int -> Int -> Integer -> [[Integer]] -> ST s (MutableSmithState s) +newMutableSmithState rowCount columnCount modulusValue rows = + case smithTierForModulus modulusValue of + SmithTierWord32 -> do + let modulusWord = fromInteger modulusValue + entries <- U.thaw (U.fromList (flattenRows (integerResidueWord modulusWord) rows)) + pure + MutableSmithState + { mutableSmithRowCount = rowCount, + mutableSmithColumnCount = columnCount, + mutableSmithModulus = modulusValue, + mutableSmithCarrier = SmithWord32Carrier modulusWord entries + } + SmithTierWord62 -> do + let modulusWord = fromInteger modulusValue + entries <- U.thaw (U.fromList (flattenRows (integerResidueWord modulusWord) rows)) + pure + MutableSmithState + { mutableSmithRowCount = rowCount, + mutableSmithColumnCount = columnCount, + mutableSmithModulus = modulusValue, + mutableSmithCarrier = SmithWord62Carrier modulusWord entries + } + SmithTierInteger -> do + entries <- V.thaw (V.fromList (flattenRows (centerResidue modulusValue) rows)) + pure + MutableSmithState + { mutableSmithRowCount = rowCount, + mutableSmithColumnCount = columnCount, + mutableSmithModulus = modulusValue, + mutableSmithCarrier = SmithIntegerCarrier entries + } + +smithTierForModulus :: Integer -> SmithTier +smithTierForModulus modulusValue + | modulusValue > 0 && modulusValue < word32ModulusLimit = SmithTierWord32 + | modulusValue > 0 && modulusValue < word62ModulusLimit = SmithTierWord62 + | otherwise = SmithTierInteger + +word32ModulusLimit :: Integer +word32ModulusLimit = + 2 ^ (32 :: Int) + +word62ModulusLimit :: Integer +word62ModulusLimit = + 2 ^ (62 :: Int) + +flattenRows :: (Integer -> a) -> [[Integer]] -> [a] +flattenRows transformEntry = + concatMap (fmap transformEntry) + +firstRowLength :: [[a]] -> Int +firstRowLength rows = + case rows of + [] -> 0 + rowValue : _ -> length rowValue + +readDiagonalMutable :: Int -> MutableSmithState s -> ST s [Integer] +readDiagonalMutable diagonalSize stateValue = + go 0 [] + where + go diagonalIndex diagonalValues + | diagonalIndex >= diagonalSize = pure (reverse diagonalValues) + | otherwise = do + diagonalValue <- readSmithEntry stateValue diagonalIndex diagonalIndex + go (diagonalIndex + 1) (diagonalValue : diagonalValues) + +smithStepMutable :: Int -> MutableSmithState s -> ST s (Maybe SmithPhaseFailure) +smithStepMutable pivotIndex stateValue + | pivotIndex >= min (mutableSmithRowCount stateValue) (mutableSmithColumnCount stateValue) = pure Nothing + | otherwise = do + pivotCandidate <- findPivotMutable pivotIndex pivotIndex stateValue + case pivotCandidate of + Nothing -> pure Nothing + Just pivotValue -> do + swapRowsMutable pivotIndex (pivotRowIndex pivotValue) stateValue + swapColumnsMutable pivotIndex (pivotColumnIndex pivotValue) stateValue + normalizationFailure <- + normalizePivotMutable + pivotIndex + pivotIndex + (max 1 (mutableSmithRowCount stateValue * mutableSmithColumnCount stateValue * 2)) + stateValue + case normalizationFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> smithStepMutable (pivotIndex + 1) stateValue + +enforceDivisibilityChainMutable :: MutableSmithState s -> ST s (Maybe SmithPhaseFailure) +enforceDivisibilityChainMutable stateValue = + go (diagonalSize * diagonalSize) + where + diagonalSize = min (mutableSmithRowCount stateValue) (mutableSmithColumnCount stateValue) + + go remainingBudget + | remainingBudget <= 0 = do + violationValue <- findDivisibilityViolationMutable diagonalSize stateValue + case violationValue of + Nothing -> pure Nothing + Just _ -> pure (Just (SmithPhaseBudgetExhausted "divisibility chain")) + | otherwise = do + violationValue <- findDivisibilityViolationMutable diagonalSize stateValue + case violationValue of + Nothing -> pure Nothing + Just violationIndex -> do + rowCombineMutable violationIndex (violationIndex + 1) (neg one) stateValue + leftFailure <- + normalizePivotMutable + violationIndex + violationIndex + (max 1 (mutableSmithRowCount stateValue * mutableSmithColumnCount stateValue * 2)) + stateValue + case leftFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + rightFailure <- + normalizePivotMutable + (violationIndex + 1) + (violationIndex + 1) + (max 1 (mutableSmithRowCount stateValue * mutableSmithColumnCount stateValue * 2)) + stateValue + case rightFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> go (remainingBudget - 1) + +findDivisibilityViolationMutable :: Int -> MutableSmithState s -> ST s (Maybe Int) +findDivisibilityViolationMutable diagonalSize stateValue = + go 0 + where + go diagonalIndex + | diagonalIndex >= diagonalSize - 1 = pure Nothing + | otherwise = do + leftDiagonal <- readSmithEntry stateValue diagonalIndex diagonalIndex + rightDiagonal <- readSmithEntry stateValue (diagonalIndex + 1) (diagonalIndex + 1) + if leftDiagonal == zero + || rightDiagonal == zero + || maybe False ((== zero) . snd) (divideIntegerMaybe rightDiagonal leftDiagonal) + then go (diagonalIndex + 1) + else pure (Just diagonalIndex) + +findPivotMutable :: Int -> Int -> MutableSmithState s -> ST s (Maybe SmithPivot) +findPivotMutable startRow startColumn stateValue = + fmap fst <$> goRows startRow Nothing + where + goRows rowIndex bestValue + | rowIndex >= mutableSmithRowCount stateValue = pure bestValue + | otherwise = do + rowBest <- goColumns rowIndex startColumn bestValue + goRows (rowIndex + 1) rowBest + + goColumns rowIndex columnIndex bestValue + | columnIndex >= mutableSmithColumnCount stateValue = pure bestValue + | otherwise = do + entryMagnitude <- entryMagnitudeMaybeMutable stateValue rowIndex columnIndex + nextBest <- + case entryMagnitude of + Nothing -> pure bestValue + Just magnitudeValue -> pure (betterPivot bestValue (SmithPivot rowIndex columnIndex, magnitudeValue)) + goColumns rowIndex (columnIndex + 1) nextBest + +betterPivot :: Maybe (SmithPivot, Integer) -> (SmithPivot, Integer) -> Maybe (SmithPivot, Integer) +betterPivot bestValue candidateValue = + case bestValue of + Nothing -> Just candidateValue + Just currentValue -> + if pivotOrderingKey candidateValue < pivotOrderingKey currentValue + then Just candidateValue + else bestValue + +pivotOrderingKey :: (SmithPivot, Integer) -> (Integer, Int, Int) +pivotOrderingKey (pivotValue, magnitudeValue) = + (magnitudeValue, pivotRowIndex pivotValue, pivotColumnIndex pivotValue) + +normalizePivotMutable :: Int -> Int -> Int -> MutableSmithState s -> ST s (Maybe SmithPhaseFailure) +normalizePivotMutable pivotRow pivotColumn remainingBudget stateValue + | remainingBudget <= 0 = pure (Just (SmithPhaseBudgetExhausted "normalization")) + | otherwise = do + (columnFailure, columnChanged) <- clearColumnMutable pivotRow pivotColumn stateValue + case columnFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + (rowFailure, rowChanged) <- clearRowMutable pivotRow pivotColumn stateValue + case rowFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + clearedColumn <- columnClearedMutable pivotRow pivotColumn stateValue + clearedRow <- rowClearedMutable pivotRow pivotColumn stateValue + if clearedColumn && clearedRow + then pure Nothing + else + if columnChanged || rowChanged + then normalizePivotMutable pivotRow pivotColumn (remainingBudget - 1) stateValue + else pure (Just SmithPhaseNormalizationStalled) + +clearColumnMutable :: Int -> Int -> MutableSmithState s -> ST s (Maybe SmithPhaseFailure, Bool) +clearColumnMutable pivotRow pivotColumn stateValue = do + candidateRow <- firstColumnEntryMutable pivotRow pivotColumn stateValue + case candidateRow of + Nothing -> pure (Nothing, False) + Just rowIndex -> do + pivotValue <- readSmithEntry stateValue pivotRow pivotColumn + entryValue <- readSmithEntry stateValue rowIndex pivotColumn + case divideIntegerMaybe entryValue pivotValue of + Nothing -> pure (Just SmithPhasePivotBecameZero, False) + Just (quotientValue, remainderValue) -> do + reductionFailure <- + if remainderValue == zero + then rowCombineMutable rowIndex pivotRow quotientValue stateValue *> pure Nothing + else gcdCombineRowsMutable pivotRow rowIndex pivotColumn stateValue + case reductionFailure of + Just failureValue -> pure (Just failureValue, True) + Nothing -> do + (nextFailure, _) <- clearColumnMutable pivotRow pivotColumn stateValue + pure (nextFailure, True) + +clearRowMutable :: Int -> Int -> MutableSmithState s -> ST s (Maybe SmithPhaseFailure, Bool) +clearRowMutable pivotRow pivotColumn stateValue = do + candidateColumn <- firstRowEntryMutable pivotRow pivotColumn stateValue + case candidateColumn of + Nothing -> pure (Nothing, False) + Just columnIndex -> do + pivotValue <- readSmithEntry stateValue pivotRow pivotColumn + entryValue <- readSmithEntry stateValue pivotRow columnIndex + case divideIntegerMaybe entryValue pivotValue of + Nothing -> pure (Just SmithPhasePivotBecameZero, False) + Just (quotientValue, remainderValue) -> do + reductionFailure <- + if remainderValue == zero + then columnCombineMutable columnIndex pivotColumn quotientValue stateValue *> pure Nothing + else gcdCombineColumnsMutable pivotRow pivotColumn columnIndex stateValue + case reductionFailure of + Just failureValue -> pure (Just failureValue, True) + Nothing -> do + (nextFailure, _) <- clearRowMutable pivotRow pivotColumn stateValue + pure (nextFailure, True) + +gcdCombineRowsMutable :: Int -> Int -> Int -> MutableSmithState s -> ST s (Maybe SmithPhaseFailure) +gcdCombineRowsMutable pivotRow candidateRow pivotColumn stateValue = do + pivotValue <- readSmithEntry stateValue pivotRow pivotColumn + entryValue <- readSmithEntry stateValue candidateRow pivotColumn + let (gcdValue, pivotCoefficient, entryCoefficient) = extendedGcdDomain pivotValue entryValue + case (exactQuotientMaybe "row gcd pivot quotient" pivotValue gcdValue, exactQuotientMaybe "row gcd entry quotient" entryValue gcdValue) of + (Right pivotQuotient, Right entryQuotient) -> do + rowPairTransformMutable pivotRow candidateRow pivotCoefficient entryCoefficient (neg entryQuotient) pivotQuotient stateValue + pure Nothing + (Left failureValue, _) -> pure (Just failureValue) + (_, Left failureValue) -> pure (Just failureValue) + +gcdCombineColumnsMutable :: Int -> Int -> Int -> MutableSmithState s -> ST s (Maybe SmithPhaseFailure) +gcdCombineColumnsMutable pivotRow pivotColumn candidateColumn stateValue = do + pivotValue <- readSmithEntry stateValue pivotRow pivotColumn + entryValue <- readSmithEntry stateValue pivotRow candidateColumn + let (gcdValue, pivotCoefficient, entryCoefficient) = extendedGcdDomain pivotValue entryValue + case (exactQuotientMaybe "column gcd pivot quotient" pivotValue gcdValue, exactQuotientMaybe "column gcd entry quotient" entryValue gcdValue) of + (Right pivotQuotient, Right entryQuotient) -> do + columnPairTransformMutable pivotColumn candidateColumn pivotCoefficient entryCoefficient (neg entryQuotient) pivotQuotient stateValue + pure Nothing + (Left failureValue, _) -> pure (Just failureValue) + (_, Left failureValue) -> pure (Just failureValue) + +exactQuotientMaybe :: String -> Integer -> Integer -> Either SmithPhaseFailure Integer +exactQuotientMaybe context numerator denominator = + case divideIntegerMaybe numerator denominator of + Nothing -> Left (SmithPhaseInexactDivision context) + Just (quotientValue, remainderValue) + | remainderValue == zero -> Right quotientValue + | otherwise -> Left (SmithPhaseInexactDivision context) + +divideIntegerMaybe :: Integer -> Integer -> Maybe (Integer, Integer) +divideIntegerMaybe numerator denominator = + divideWithRemainder numerator <$> mkNonZeroDivisor denominator + +columnClearedMutable :: Int -> Int -> MutableSmithState s -> ST s Bool +columnClearedMutable pivotRow pivotColumn stateValue = + go 0 + where + go rowIndex + | rowIndex >= mutableSmithRowCount stateValue = pure True + | rowIndex == pivotRow = go (rowIndex + 1) + | otherwise = do + isZeroEntry <- entryIsZeroMutable stateValue rowIndex pivotColumn + if isZeroEntry + then go (rowIndex + 1) + else pure False + +rowClearedMutable :: Int -> Int -> MutableSmithState s -> ST s Bool +rowClearedMutable pivotRow pivotColumn stateValue = + go 0 + where + go columnIndex + | columnIndex >= mutableSmithColumnCount stateValue = pure True + | columnIndex == pivotColumn = go (columnIndex + 1) + | otherwise = do + isZeroEntry <- entryIsZeroMutable stateValue pivotRow columnIndex + if isZeroEntry + then go (columnIndex + 1) + else pure False + +firstColumnEntryMutable :: Int -> Int -> MutableSmithState s -> ST s (Maybe Int) +firstColumnEntryMutable pivotRow pivotColumn stateValue = + go 0 + where + go rowIndex + | rowIndex >= mutableSmithRowCount stateValue = pure Nothing + | rowIndex == pivotRow = go (rowIndex + 1) + | otherwise = do + isZeroEntry <- entryIsZeroMutable stateValue rowIndex pivotColumn + if isZeroEntry + then go (rowIndex + 1) + else pure (Just rowIndex) + +firstRowEntryMutable :: Int -> Int -> MutableSmithState s -> ST s (Maybe Int) +firstRowEntryMutable pivotRow pivotColumn stateValue = + go 0 + where + go columnIndex + | columnIndex >= mutableSmithColumnCount stateValue = pure Nothing + | columnIndex == pivotColumn = go (columnIndex + 1) + | otherwise = do + isZeroEntry <- entryIsZeroMutable stateValue pivotRow columnIndex + if isZeroEntry + then go (columnIndex + 1) + else pure (Just columnIndex) + +readSmithEntry :: MutableSmithState s -> Int -> Int -> ST s Integer +readSmithEntry stateValue rowIndex columnIndex = + case mutableSmithCarrier stateValue of + SmithWord32Carrier modulusWord entries -> do + entryValue <- MU.read entries (flatIndex (mutableSmithColumnCount stateValue) rowIndex columnIndex) + pure (symmetricLift (toInteger modulusWord) (toInteger entryValue)) + SmithWord62Carrier modulusWord entries -> do + entryValue <- MU.read entries (flatIndex (mutableSmithColumnCount stateValue) rowIndex columnIndex) + pure (symmetricLift (toInteger modulusWord) (toInteger entryValue)) + SmithIntegerCarrier entries -> + MV.read entries (flatIndex (mutableSmithColumnCount stateValue) rowIndex columnIndex) + +entryIsZeroMutable :: MutableSmithState s -> Int -> Int -> ST s Bool +entryIsZeroMutable stateValue rowIndex columnIndex = + case mutableSmithCarrier stateValue of + SmithWord32Carrier _ entries -> (== 0) <$> MU.read entries (flatIndex (mutableSmithColumnCount stateValue) rowIndex columnIndex) + SmithWord62Carrier _ entries -> (== 0) <$> MU.read entries (flatIndex (mutableSmithColumnCount stateValue) rowIndex columnIndex) + SmithIntegerCarrier entries -> (== zero) <$> MV.read entries (flatIndex (mutableSmithColumnCount stateValue) rowIndex columnIndex) + +entryMagnitudeMaybeMutable :: MutableSmithState s -> Int -> Int -> ST s (Maybe Integer) +entryMagnitudeMaybeMutable stateValue rowIndex columnIndex = + case mutableSmithCarrier stateValue of + SmithWord32Carrier modulusWord entries -> do + entryValue <- MU.read entries (flatIndex (mutableSmithColumnCount stateValue) rowIndex columnIndex) + pure + ( if entryValue == 0 + then Nothing + else Just (toInteger (wordSymmetricMagnitude modulusWord entryValue)) + ) + SmithWord62Carrier modulusWord entries -> do + entryValue <- MU.read entries (flatIndex (mutableSmithColumnCount stateValue) rowIndex columnIndex) + pure + ( if entryValue == 0 + then Nothing + else Just (toInteger (wordSymmetricMagnitude modulusWord entryValue)) + ) + SmithIntegerCarrier entries -> do + entryValue <- MV.read entries (flatIndex (mutableSmithColumnCount stateValue) rowIndex columnIndex) + pure + ( if entryValue == zero + then Nothing + else Just (abs entryValue) + ) + +wordSymmetricMagnitude :: Word64 -> Word64 -> Word64 +wordSymmetricMagnitude modulusWord entryValue = + let complementValue = modulusWord - entryValue + in if entryValue <= complementValue + then entryValue + else complementValue + +swapRowsMutable :: Int -> Int -> MutableSmithState s -> ST s () +swapRowsMutable leftRow rightRow stateValue = + case mutableSmithCarrier stateValue of + SmithWord32Carrier _ entries -> swapRowsWord (mutableSmithColumnCount stateValue) entries leftRow rightRow 0 + SmithWord62Carrier _ entries -> swapRowsWord (mutableSmithColumnCount stateValue) entries leftRow rightRow 0 + SmithIntegerCarrier entries -> swapRowsInteger (mutableSmithColumnCount stateValue) entries leftRow rightRow 0 + +swapColumnsMutable :: Int -> Int -> MutableSmithState s -> ST s () +swapColumnsMutable leftColumn rightColumn stateValue = + case mutableSmithCarrier stateValue of + SmithWord32Carrier _ entries -> swapColumnsWord (mutableSmithColumnCount stateValue) entries leftColumn rightColumn 0 (mutableSmithRowCount stateValue) + SmithWord62Carrier _ entries -> swapColumnsWord (mutableSmithColumnCount stateValue) entries leftColumn rightColumn 0 (mutableSmithRowCount stateValue) + SmithIntegerCarrier entries -> swapColumnsInteger (mutableSmithColumnCount stateValue) entries leftColumn rightColumn 0 (mutableSmithRowCount stateValue) + +swapRowsWord :: Int -> MU.MVector s Word64 -> Int -> Int -> Int -> ST s () +swapRowsWord columnCount entries leftRow rightRow columnIndex + | leftRow == rightRow = pure () + | columnIndex >= columnCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount leftRow columnIndex + rightIndex = flatIndex columnCount rightRow columnIndex + leftValue <- MU.read entries leftIndex + rightValue <- MU.read entries rightIndex + MU.write entries leftIndex rightValue + MU.write entries rightIndex leftValue + swapRowsWord columnCount entries leftRow rightRow (columnIndex + 1) + +swapRowsInteger :: Int -> MV.MVector s Integer -> Int -> Int -> Int -> ST s () +swapRowsInteger columnCount entries leftRow rightRow columnIndex + | leftRow == rightRow = pure () + | columnIndex >= columnCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount leftRow columnIndex + rightIndex = flatIndex columnCount rightRow columnIndex + leftValue <- MV.read entries leftIndex + rightValue <- MV.read entries rightIndex + MV.write entries leftIndex rightValue + MV.write entries rightIndex leftValue + swapRowsInteger columnCount entries leftRow rightRow (columnIndex + 1) + +swapColumnsWord :: Int -> MU.MVector s Word64 -> Int -> Int -> Int -> Int -> ST s () +swapColumnsWord columnCount entries leftColumn rightColumn rowIndex rowCount + | leftColumn == rightColumn = pure () + | rowIndex >= rowCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount rowIndex leftColumn + rightIndex = flatIndex columnCount rowIndex rightColumn + leftValue <- MU.read entries leftIndex + rightValue <- MU.read entries rightIndex + MU.write entries leftIndex rightValue + MU.write entries rightIndex leftValue + swapColumnsWord columnCount entries leftColumn rightColumn (rowIndex + 1) rowCount + +swapColumnsInteger :: Int -> MV.MVector s Integer -> Int -> Int -> Int -> Int -> ST s () +swapColumnsInteger columnCount entries leftColumn rightColumn rowIndex rowCount + | leftColumn == rightColumn = pure () + | rowIndex >= rowCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount rowIndex leftColumn + rightIndex = flatIndex columnCount rowIndex rightColumn + leftValue <- MV.read entries leftIndex + rightValue <- MV.read entries rightIndex + MV.write entries leftIndex rightValue + MV.write entries rightIndex leftValue + swapColumnsInteger columnCount entries leftColumn rightColumn (rowIndex + 1) rowCount + +rowCombineMutable :: Int -> Int -> Integer -> MutableSmithState s -> ST s () +rowCombineMutable targetRow sourceRow coefficient stateValue = + case mutableSmithCarrier stateValue of + SmithWord32Carrier modulusWord entries -> + rowCombineWord wordMul32 (mutableSmithColumnCount stateValue) modulusWord entries targetRow sourceRow (integerResidueWord modulusWord (neg coefficient)) 0 + SmithWord62Carrier modulusWord entries -> + rowCombineWord wordMul62 (mutableSmithColumnCount stateValue) modulusWord entries targetRow sourceRow (integerResidueWord modulusWord (neg coefficient)) 0 + SmithIntegerCarrier entries -> + rowCombineInteger (mutableSmithColumnCount stateValue) (mutableSmithModulus stateValue) entries targetRow sourceRow coefficient 0 + +columnCombineMutable :: Int -> Int -> Integer -> MutableSmithState s -> ST s () +columnCombineMutable targetColumn sourceColumn coefficient stateValue = + case mutableSmithCarrier stateValue of + SmithWord32Carrier modulusWord entries -> + columnCombineWord wordMul32 (mutableSmithColumnCount stateValue) modulusWord entries targetColumn sourceColumn (integerResidueWord modulusWord (neg coefficient)) 0 (mutableSmithRowCount stateValue) + SmithWord62Carrier modulusWord entries -> + columnCombineWord wordMul62 (mutableSmithColumnCount stateValue) modulusWord entries targetColumn sourceColumn (integerResidueWord modulusWord (neg coefficient)) 0 (mutableSmithRowCount stateValue) + SmithIntegerCarrier entries -> + columnCombineInteger (mutableSmithColumnCount stateValue) (mutableSmithModulus stateValue) entries targetColumn sourceColumn coefficient 0 (mutableSmithRowCount stateValue) + +rowPairTransformMutable :: Int -> Int -> Integer -> Integer -> Integer -> Integer -> MutableSmithState s -> ST s () +rowPairTransformMutable leftRow rightRow aa ab ba bb stateValue = + case mutableSmithCarrier stateValue of + SmithWord32Carrier modulusWord entries -> + rowPairTransformWord wordMul32 (mutableSmithColumnCount stateValue) modulusWord entries leftRow rightRow (integerResidueWord modulusWord aa) (integerResidueWord modulusWord ab) (integerResidueWord modulusWord ba) (integerResidueWord modulusWord bb) 0 + SmithWord62Carrier modulusWord entries -> + rowPairTransformWord wordMul62 (mutableSmithColumnCount stateValue) modulusWord entries leftRow rightRow (integerResidueWord modulusWord aa) (integerResidueWord modulusWord ab) (integerResidueWord modulusWord ba) (integerResidueWord modulusWord bb) 0 + SmithIntegerCarrier entries -> + rowPairTransformInteger (mutableSmithColumnCount stateValue) (mutableSmithModulus stateValue) entries leftRow rightRow aa ab ba bb 0 + +columnPairTransformMutable :: Int -> Int -> Integer -> Integer -> Integer -> Integer -> MutableSmithState s -> ST s () +columnPairTransformMutable leftColumn rightColumn aa ab ba bb stateValue = + case mutableSmithCarrier stateValue of + SmithWord32Carrier modulusWord entries -> + columnPairTransformWord wordMul32 (mutableSmithColumnCount stateValue) modulusWord entries leftColumn rightColumn (integerResidueWord modulusWord aa) (integerResidueWord modulusWord ab) (integerResidueWord modulusWord ba) (integerResidueWord modulusWord bb) 0 (mutableSmithRowCount stateValue) + SmithWord62Carrier modulusWord entries -> + columnPairTransformWord wordMul62 (mutableSmithColumnCount stateValue) modulusWord entries leftColumn rightColumn (integerResidueWord modulusWord aa) (integerResidueWord modulusWord ab) (integerResidueWord modulusWord ba) (integerResidueWord modulusWord bb) 0 (mutableSmithRowCount stateValue) + SmithIntegerCarrier entries -> + columnPairTransformInteger (mutableSmithColumnCount stateValue) (mutableSmithModulus stateValue) entries leftColumn rightColumn aa ab ba bb 0 (mutableSmithRowCount stateValue) + +rowCombineWord :: (Word64 -> Word64 -> Word64 -> Word64) -> Int -> Word64 -> MU.MVector s Word64 -> Int -> Int -> Word64 -> Int -> ST s () +rowCombineWord multiplyMod columnCount modulusWord entries targetRow sourceRow coefficientWord columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + let targetIndex = flatIndex columnCount targetRow columnIndex + sourceIndex = flatIndex columnCount sourceRow columnIndex + targetValue <- MU.read entries targetIndex + sourceValue <- MU.read entries sourceIndex + MU.write entries targetIndex (wordAddMod modulusWord targetValue (multiplyMod modulusWord coefficientWord sourceValue)) + rowCombineWord multiplyMod columnCount modulusWord entries targetRow sourceRow coefficientWord (columnIndex + 1) + +columnCombineWord :: (Word64 -> Word64 -> Word64 -> Word64) -> Int -> Word64 -> MU.MVector s Word64 -> Int -> Int -> Word64 -> Int -> Int -> ST s () +columnCombineWord multiplyMod columnCount modulusWord entries targetColumn sourceColumn coefficientWord rowIndex rowCount + | rowIndex >= rowCount = pure () + | otherwise = do + let targetIndex = flatIndex columnCount rowIndex targetColumn + sourceIndex = flatIndex columnCount rowIndex sourceColumn + targetValue <- MU.read entries targetIndex + sourceValue <- MU.read entries sourceIndex + MU.write entries targetIndex (wordAddMod modulusWord targetValue (multiplyMod modulusWord coefficientWord sourceValue)) + columnCombineWord multiplyMod columnCount modulusWord entries targetColumn sourceColumn coefficientWord (rowIndex + 1) rowCount + +rowPairTransformWord :: (Word64 -> Word64 -> Word64 -> Word64) -> Int -> Word64 -> MU.MVector s Word64 -> Int -> Int -> Word64 -> Word64 -> Word64 -> Word64 -> Int -> ST s () +rowPairTransformWord multiplyMod columnCount modulusWord entries leftRow rightRow aa ab ba bb columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount leftRow columnIndex + rightIndex = flatIndex columnCount rightRow columnIndex + leftValue <- MU.read entries leftIndex + rightValue <- MU.read entries rightIndex + MU.write entries leftIndex (wordLinearCombination multiplyMod modulusWord aa leftValue ab rightValue) + MU.write entries rightIndex (wordLinearCombination multiplyMod modulusWord ba leftValue bb rightValue) + rowPairTransformWord multiplyMod columnCount modulusWord entries leftRow rightRow aa ab ba bb (columnIndex + 1) + +columnPairTransformWord :: (Word64 -> Word64 -> Word64 -> Word64) -> Int -> Word64 -> MU.MVector s Word64 -> Int -> Int -> Word64 -> Word64 -> Word64 -> Word64 -> Int -> Int -> ST s () +columnPairTransformWord multiplyMod columnCount modulusWord entries leftColumn rightColumn aa ab ba bb rowIndex rowCount + | rowIndex >= rowCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount rowIndex leftColumn + rightIndex = flatIndex columnCount rowIndex rightColumn + leftValue <- MU.read entries leftIndex + rightValue <- MU.read entries rightIndex + MU.write entries leftIndex (wordLinearCombination multiplyMod modulusWord aa leftValue ab rightValue) + MU.write entries rightIndex (wordLinearCombination multiplyMod modulusWord ba leftValue bb rightValue) + columnPairTransformWord multiplyMod columnCount modulusWord entries leftColumn rightColumn aa ab ba bb (rowIndex + 1) rowCount + +rowCombineInteger :: Int -> Integer -> MV.MVector s Integer -> Int -> Int -> Integer -> Int -> ST s () +rowCombineInteger columnCount modulusValue entries targetRow sourceRow coefficient columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + let targetIndex = flatIndex columnCount targetRow columnIndex + sourceIndex = flatIndex columnCount sourceRow columnIndex + targetValue <- MV.read entries targetIndex + sourceValue <- MV.read entries sourceIndex + MV.write entries targetIndex (centerResidue modulusValue (targetValue - coefficient * sourceValue)) + rowCombineInteger columnCount modulusValue entries targetRow sourceRow coefficient (columnIndex + 1) + +columnCombineInteger :: Int -> Integer -> MV.MVector s Integer -> Int -> Int -> Integer -> Int -> Int -> ST s () +columnCombineInteger columnCount modulusValue entries targetColumn sourceColumn coefficient rowIndex rowCount + | rowIndex >= rowCount = pure () + | otherwise = do + let targetIndex = flatIndex columnCount rowIndex targetColumn + sourceIndex = flatIndex columnCount rowIndex sourceColumn + targetValue <- MV.read entries targetIndex + sourceValue <- MV.read entries sourceIndex + MV.write entries targetIndex (centerResidue modulusValue (targetValue - coefficient * sourceValue)) + columnCombineInteger columnCount modulusValue entries targetColumn sourceColumn coefficient (rowIndex + 1) rowCount + +rowPairTransformInteger :: Int -> Integer -> MV.MVector s Integer -> Int -> Int -> Integer -> Integer -> Integer -> Integer -> Int -> ST s () +rowPairTransformInteger columnCount modulusValue entries leftRow rightRow aa ab ba bb columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount leftRow columnIndex + rightIndex = flatIndex columnCount rightRow columnIndex + leftValue <- MV.read entries leftIndex + rightValue <- MV.read entries rightIndex + MV.write entries leftIndex (centerResidue modulusValue (aa * leftValue + ab * rightValue)) + MV.write entries rightIndex (centerResidue modulusValue (ba * leftValue + bb * rightValue)) + rowPairTransformInteger columnCount modulusValue entries leftRow rightRow aa ab ba bb (columnIndex + 1) + +columnPairTransformInteger :: Int -> Integer -> MV.MVector s Integer -> Int -> Int -> Integer -> Integer -> Integer -> Integer -> Int -> Int -> ST s () +columnPairTransformInteger columnCount modulusValue entries leftColumn rightColumn aa ab ba bb rowIndex rowCount + | rowIndex >= rowCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount rowIndex leftColumn + rightIndex = flatIndex columnCount rowIndex rightColumn + leftValue <- MV.read entries leftIndex + rightValue <- MV.read entries rightIndex + MV.write entries leftIndex (centerResidue modulusValue (aa * leftValue + ab * rightValue)) + MV.write entries rightIndex (centerResidue modulusValue (ba * leftValue + bb * rightValue)) + columnPairTransformInteger columnCount modulusValue entries leftColumn rightColumn aa ab ba bb (rowIndex + 1) rowCount + +wordLinearCombination :: (Word64 -> Word64 -> Word64 -> Word64) -> Word64 -> Word64 -> Word64 -> Word64 -> Word64 -> Word64 +wordLinearCombination multiplyMod modulusWord leftCoefficient leftValue rightCoefficient rightValue = + wordAddMod modulusWord (multiplyMod modulusWord leftCoefficient leftValue) (multiplyMod modulusWord rightCoefficient rightValue) + +wordAddMod :: Word64 -> Word64 -> Word64 -> Word64 +wordAddMod modulusWord leftValue rightValue = + let sumValue = leftValue + rightValue + in if sumValue >= modulusWord + then sumValue - modulusWord + else sumValue + +wordMul32 :: Word64 -> Word64 -> Word64 -> Word64 +wordMul32 modulusWord leftValue rightValue = + (leftValue * rightValue) `rem` modulusWord + +wordMul62 :: Word64 -> Word64 -> Word64 -> Word64 +wordMul62 (W64# modulusWord#) (W64# leftValue#) (W64# rightValue#) = + case timesWord2# (word64ToWord# leftValue#) (word64ToWord# rightValue#) of + (# highWord#, lowWord# #) -> + case quotRemWord2# highWord# lowWord# (word64ToWord# modulusWord#) of + (# _, remainderWord# #) -> W64# (wordToWord64# remainderWord#) + +diagonalCoreRows :: [Integer] -> [[Integer]] +diagonalCoreRows factors = + let coreSize = length factors + in [ [ if rowIndex == columnIndex then diagonalValueAt factors rowIndex else zero + | columnIndex <- [0 .. coreSize - 1] + ] + | rowIndex <- [0 .. coreSize - 1] + ] + +diagonalFlatEntries :: Int -> Int -> [Integer] -> [Integer] +diagonalFlatEntries rowCount columnCount invariantFactors = + [ if rowIndex == columnIndex then diagonalValueAt invariantFactors rowIndex else zero + | rowIndex <- [0 .. rowCount - 1], + columnIndex <- [0 .. columnCount - 1] + ] + +diagonalValueAt :: [Integer] -> Int -> Integer +diagonalValueAt values indexValue = + maybe zero id (values !? indexValue) + +centerResidue :: Integer -> Integer -> Integer +centerResidue modulusValue value + | modulusValue <= 1 = value + | doubled > modulusValue = residueValue - modulusValue + | otherwise = residueValue + where + residueValue = value `mod` modulusValue + doubled = 2 * residueValue + +(!?) :: [a] -> Int -> Maybe a +values !? targetIndex + | targetIndex < 0 = Nothing + | otherwise = + case drop targetIndex values of + [] -> Nothing + value : _ -> Just value diff --git a/src-domain/Moonlight/LinAlg/Pure/Domain/Smith/Witnessed.hs b/src-domain/Moonlight/LinAlg/Pure/Domain/Smith/Witnessed.hs new file mode 100644 --- /dev/null +++ b/src-domain/Moonlight/LinAlg/Pure/Domain/Smith/Witnessed.hs @@ -0,0 +1,1821 @@ +module Moonlight.LinAlg.Pure.Domain.Smith.Witnessed + ( smithNormalFormWitnessed, + ) +where + +import Control.Monad.ST (ST, runST) +import Data.Bifunctor (first) +import Data.Bits (shiftL) +import Data.List ((!?)) +import Data.Vector qualified as V +import Data.Vector.Mutable qualified as MV +import Data.Vector.Unboxed qualified as U +import Data.Vector.Unboxed.Mutable qualified as MU +import Data.Word (Word64) +import GHC.TypeNats (KnownNat) +import Moonlight.Algebra.Pure.Ring (GCDDomain (..)) +import Moonlight.Core + ( MoonlightError (..), + checkedNonNegativeProduct, + checkedNonNegativeSum, + ) +import Moonlight.LinAlg.Internal.Backend.Smith (SmithNormalForm (..)) +import Moonlight.LinAlg.Pure.Dense.Types + ( Matrix, + fromListMatrix, + ) +import qualified Moonlight.LinAlg.Pure.Dense.Types as DenseTypes +import Moonlight.LinAlg.Pure.Domain.Smith.Multimodular + ( PrimeSweep (..), + certifiedPrimeSweep, + integerResidueWord, + modInverseWord, + modMul, + wordPrimeLadder, + ) +import Prelude + +data SmithWitnessArena s = SmithWitnessArena + { smithWitnessRowCount :: !Int, + smithWitnessColumnCount :: !Int, + smithWitnessWork :: !(MV.MVector s Integer), + smithWitnessLeftRows :: !(MV.MVector s Integer), + smithWitnessRightRows :: !(MV.MVector s Integer), + smithWitnessLeftInverseRows :: !(MV.MVector s Integer), + smithWitnessRightInverseRows :: !(MV.MVector s Integer) + } + +data SmithWitnessFailure + = SmithWitnessBudgetExhausted !String + | SmithWitnessNormalizationStalled + | SmithWitnessPivotBecameZero + | SmithWitnessInexactDivision !String + | SmithWitnessTransformRecoveryFailed !String + | SmithWitnessVerificationFailed !String + deriving stock (Eq, Show) + +data SmithWitnessResult + = SmithWitnessResult ![Integer] ![Integer] ![Integer] ![Integer] ![Integer] + | SmithWitnessFailed !SmithWitnessFailure + +data SmithExactQuotient + = SmithExactQuotient !Integer + | SmithInexactQuotient !SmithWitnessFailure + +data SmithPivot = SmithPivot + { smithPivotRowIndex :: !Int, + smithPivotColumnIndex :: !Int + } + deriving stock (Eq, Show) + +smithNormalFormWitnessed :: + forall r c. + (KnownNat r, KnownNat c) => + Matrix r c Integer -> + Either MoonlightError (SmithNormalForm r c Integer) +smithNormalFormWitnessed matrixValue = do + rows <- DenseTypes.matrixToRows matrixValue + let (rowCount, columnCount) = DenseTypes.matrixShape matrixValue + witnessResult <- runWitnessedSmith rowCount columnCount rows + case witnessResult of + SmithWitnessFailed failureValue -> + Left (InvariantViolation ("Smith witnessed Integer normal form failed: " <> show failureValue)) + SmithWitnessResult leftEntries diagonalEntries rightEntries leftInverseEntries rightInverseEntries -> do + leftMatrix <- fromListMatrix @r @r leftEntries + diagonalMatrix <- fromListMatrix @r @c diagonalEntries + rightMatrix <- fromListMatrix @c @c rightEntries + leftInverseMatrix <- fromListMatrix @r @r leftInverseEntries + rightInverseMatrix <- fromListMatrix @c @c rightInverseEntries + pure + SmithNormalForm + { smithLeft = leftMatrix, + smithDiagonal = diagonalMatrix, + smithRight = rightMatrix, + smithLeftInverse = leftInverseMatrix, + smithRightInverse = rightInverseMatrix + } + +runWitnessedSmith :: Int -> Int -> [[Integer]] -> Either MoonlightError SmithWitnessResult +runWitnessedSmith rowCount columnCount rows = do + validateSmithWitnessCardinalities rowCount columnCount + if rowCount == columnCount && rowCount >= fastWitnessSizeFloor + then do + primeSweep <- certifiedPrimeSweep rowCount columnCount rows + case primeSweepDeterminant primeSweep of + Just determinantValue + | primeSweepRank primeSweep == rowCount && determinantValue /= 0 -> + pure (runFastNonsingularWitnessedSmith rowCount determinantValue rows) + _ -> pure (runAlternatingWitnessedSmith rowCount columnCount (concat rows)) + else pure (runAlternatingWitnessedSmith rowCount columnCount (concat rows)) + +validateSmithWitnessCardinalities :: Int -> Int -> Either MoonlightError () +validateSmithWitnessCardinalities rowCount columnCount = do + matrixEntryCount <- checkedWitnessProduct "matrix entries" rowCount columnCount + rowWitnessEntryCount <- checkedWitnessProduct "left witness entries" rowCount rowCount + columnWitnessEntryCount <- checkedWitnessProduct "right witness entries" columnCount columnCount + _ <- checkedWitnessProduct "normalization budget" matrixEntryCount 16 + _ <- checkedWitnessProduct "divisibility-chain budget" (min rowCount columnCount) (min rowCount columnCount) + dimensionSum <- checkedWitnessSum "alternation dimension sum" rowCount columnCount + doubledDimensionSum <- checkedWitnessProduct "alternation budget" 2 dimensionSum + _ <- checkedWitnessSum "alternation budget" 64 doubledDimensionSum + doubledMaximumDimension <- checkedWitnessProduct "augmented witness width" 2 (max rowCount columnCount) + _ <- checkedWitnessSum "fast witness budget" 8 doubledMaximumDimension + _ <- checkedWitnessProduct "augmented left witness entries" rowWitnessEntryCount 2 + _ <- checkedWitnessProduct "augmented right witness entries" columnWitnessEntryCount 2 + Right () + +checkedWitnessProduct :: String -> Int -> Int -> Either MoonlightError Int +checkedWitnessProduct context leftFactor rightFactor = + first + (const (InvariantViolation ("Smith witnessed " <> context <> " exceed Int cardinality"))) + (checkedNonNegativeProduct leftFactor rightFactor) + +checkedWitnessSum :: String -> Int -> Int -> Either MoonlightError Int +checkedWitnessSum context leftTerm rightTerm = + first + (const (InvariantViolation ("Smith witnessed " <> context <> " exceed Int cardinality"))) + (checkedNonNegativeSum leftTerm rightTerm) + +fastWitnessSizeFloor :: Int +fastWitnessSizeFloor = 25 + +runAlternatingWitnessedSmith :: Int -> Int -> [Integer] -> SmithWitnessResult +runAlternatingWitnessedSmith rowCount columnCount entries = + runST $ do + arenaValue <- newSmithWitnessArena rowCount columnCount entries + stepFailure <- alternatingHermiteMutable (alternationBudget arenaValue) arenaValue + case stepFailure of + Just failureValue -> pure (SmithWitnessFailed failureValue) + Nothing -> do + chainFailure <- enforceDivisibilityChainMutable arenaValue + case chainFailure of + Just failureValue -> pure (SmithWitnessFailed failureValue) + Nothing -> do + normalizeFailure <- normalizeDiagonalUnitsMutable arenaValue + case normalizeFailure of + Just failureValue -> pure (SmithWitnessFailed failureValue) + Nothing -> + SmithWitnessResult + <$> readFlatVector (smithWitnessLeftRows arenaValue) + <*> readFlatVector (smithWitnessWork arenaValue) + <*> readFlatVector (smithWitnessRightRows arenaValue) + <*> readFlatVector (smithWitnessLeftInverseRows arenaValue) + <*> readFlatVector (smithWitnessRightInverseRows arenaValue) + +data FastWitnessState = FastWitnessState + { fastWitnessWork :: ![Integer], + fastWitnessLeft :: ![Integer], + fastWitnessLeftTimesOriginal :: ![Integer], + fastWitnessRight :: ![Integer], + fastWitnessLeftInverse :: ![Integer], + fastWitnessRightInverse :: ![Integer] + } + deriving stock (Eq, Show) + +data FastTransformOrientation + = FastLeftTimesInverseRight + | FastInverseLeftTimesRight + deriving stock (Eq, Show) + +runFastNonsingularWitnessedSmith :: Int -> Integer -> [[Integer]] -> SmithWitnessResult +runFastNonsingularWitnessedSmith matrixSize determinantValue rows = + case fastNonsingularWitness matrixSize determinantValue rows of + Left failureValue -> SmithWitnessFailed failureValue + Right stateValue -> finalizeFastWitness matrixSize stateValue + + +fastNonsingularWitness :: Int -> Integer -> [[Integer]] -> Either SmithWitnessFailure FastWitnessState +fastNonsingularWitness matrixSize determinantValue rows = + fastWitnessStep matrixSize modulusValue originalEntries (fastWitnessBudget matrixSize) initialState + where + modulusValue :: Integer + modulusValue = 2 * abs determinantValue + + originalEntries :: [Integer] + originalEntries = concat rows + + initialState :: FastWitnessState + initialState = + FastWitnessState + { fastWitnessWork = originalEntries, + fastWitnessLeft = identityList matrixSize, + fastWitnessLeftTimesOriginal = originalEntries, + fastWitnessRight = identityList matrixSize, + fastWitnessLeftInverse = identityList matrixSize, + fastWitnessRightInverse = identityList matrixSize + } + +fastWitnessBudget :: Int -> Int +fastWitnessBudget matrixSize = + 8 + 2 * matrixSize + +fastWitnessStep :: Int -> Integer -> [Integer] -> Int -> FastWitnessState -> Either SmithWitnessFailure FastWitnessState +fastWitnessStep matrixSize modulusValue originalEntries remainingBudget stateValue + | remainingBudget <= 0 = Left (SmithWitnessBudgetExhausted "mod-det witnessed alternation") + | matrixIsDiagonal matrixSize (fastWitnessWork stateValue) = completeFastWitness matrixSize originalEntries stateValue + | otherwise = do + rowState <- fastWitnessRowHermite matrixSize modulusValue stateValue + columnState <- fastWitnessColumnHermite matrixSize modulusValue rowState + if fastWitnessWork columnState == fastWitnessWork stateValue + then + if matrixIsDiagonal matrixSize (fastWitnessWork columnState) + then completeFastWitness matrixSize originalEntries columnState + else Left SmithWitnessNormalizationStalled + else fastWitnessStep matrixSize modulusValue originalEntries (remainingBudget - 1) columnState + + + + +fastWitnessRowHermite :: Int -> Integer -> FastWitnessState -> Either SmithWitnessFailure FastWitnessState +fastWitnessRowHermite matrixSize modulusValue stateValue = do + rowHermiteEntries <- rowHermiteModulo matrixSize modulusValue (fastWitnessWork stateValue) + if rowHermiteEntries == fastWitnessWork stateValue + then Right stateValue + else do + rowTransform <- + recoverTransform + matrixSize + FastLeftTimesInverseRight + rowHermiteEntries + (fastWitnessWork stateValue) + (transformRecoveryBound matrixSize modulusValue rowHermiteEntries (fastWitnessWork stateValue)) + "row HNF transform" + Right + stateValue + { fastWitnessWork = rowHermiteEntries, + fastWitnessLeft = composeFastFactor matrixSize rowTransform (fastWitnessLeft stateValue), + fastWitnessLeftTimesOriginal = + if fastWitnessLeftTimesOriginal stateValue == fastWitnessWork stateValue + then rowHermiteEntries + else composeFastFactor matrixSize rowTransform (fastWitnessLeftTimesOriginal stateValue) + } + +fastWitnessColumnHermite :: Int -> Integer -> FastWitnessState -> Either SmithWitnessFailure FastWitnessState +fastWitnessColumnHermite matrixSize modulusValue stateValue = do + columnHermiteEntries <- columnHermiteModulo matrixSize modulusValue (fastWitnessWork stateValue) + if columnHermiteEntries == fastWitnessWork stateValue + then Right stateValue + else do + columnTransform <- + recoverTransform + matrixSize + FastInverseLeftTimesRight + (fastWitnessWork stateValue) + columnHermiteEntries + (transformRecoveryBound matrixSize modulusValue columnHermiteEntries (fastWitnessWork stateValue)) + "column HNF transform" + Right + stateValue + { fastWitnessWork = columnHermiteEntries, + fastWitnessRight = composeFastFactor matrixSize (fastWitnessRight stateValue) columnTransform + } + +composeFastFactor :: Int -> [Integer] -> [Integer] -> [Integer] +composeFastFactor matrixSize leftEntries rightEntries + | leftEntries == identityList matrixSize = rightEntries + | rightEntries == identityList matrixSize = leftEntries + | otherwise = matrixProduct matrixSize matrixSize matrixSize leftEntries rightEntries + +completeFastWitness :: Int -> [Integer] -> FastWitnessState -> Either SmithWitnessFailure FastWitnessState +completeFastWitness matrixSize originalEntries stateValue = do + let leftTimesOriginal = fastWitnessLeftTimesOriginal stateValue + originalTimesRight = matrixProduct matrixSize matrixSize matrixSize originalEntries (fastWitnessRight stateValue) + diagonalValues = [valueAt (fastWitnessWork stateValue) (flatIndex matrixSize axisIndex axisIndex) | axisIndex <- [0 .. matrixSize - 1]] + leftInverseEntries <- divideColumnsByDiagonal matrixSize "left inverse diagonal division" diagonalValues originalTimesRight + rightInverseEntries <- divideRowsByDiagonal matrixSize "right inverse diagonal division" diagonalValues leftTimesOriginal + Right + stateValue + { fastWitnessLeftInverse = leftInverseEntries, + fastWitnessRightInverse = rightInverseEntries + } + +divideColumnsByDiagonal :: Int -> String -> [Integer] -> [Integer] -> Either SmithWitnessFailure [Integer] +divideColumnsByDiagonal matrixSize context diagonalValues entries = + traverse divideEntry (zip [0 ..] entries) + where + divideEntry :: (Int, Integer) -> Either SmithWitnessFailure Integer + divideEntry (entryIndex, entryValue) = + case exactQuotientMutable context entryValue (valueAt diagonalValues (entryIndex `rem` matrixSize)) of + SmithExactQuotient quotientValue -> Right quotientValue + SmithInexactQuotient failureValue -> Left failureValue + +divideRowsByDiagonal :: Int -> String -> [Integer] -> [Integer] -> Either SmithWitnessFailure [Integer] +divideRowsByDiagonal matrixSize context diagonalValues entries = + traverse divideEntry (zip [0 ..] entries) + where + divideEntry :: (Int, Integer) -> Either SmithWitnessFailure Integer + divideEntry (entryIndex, entryValue) = + case exactQuotientMutable context entryValue (valueAt diagonalValues (entryIndex `quot` matrixSize)) of + SmithExactQuotient quotientValue -> Right quotientValue + SmithInexactQuotient failureValue -> Left failureValue + +finalizeFastWitness :: Int -> FastWitnessState -> SmithWitnessResult +finalizeFastWitness matrixSize stateValue = + runST $ do + arenaValue <- + newSmithWitnessArenaFromWitnesses + matrixSize + (fastWitnessWork stateValue) + (fastWitnessLeft stateValue) + (fastWitnessRight stateValue) + (fastWitnessLeftInverse stateValue) + (fastWitnessRightInverse stateValue) + chainFailure <- enforceDivisibilityChainMutable arenaValue + case chainFailure of + Just failureValue -> pure (SmithWitnessFailed failureValue) + Nothing -> do + normalizeFailure <- normalizeDiagonalUnitsMutable arenaValue + case normalizeFailure of + Just failureValue -> pure (SmithWitnessFailed failureValue) + Nothing -> + SmithWitnessResult + <$> readFlatVector (smithWitnessLeftRows arenaValue) + <*> readFlatVector (smithWitnessWork arenaValue) + <*> readFlatVector (smithWitnessRightRows arenaValue) + <*> readFlatVector (smithWitnessLeftInverseRows arenaValue) + <*> readFlatVector (smithWitnessRightInverseRows arenaValue) + +newSmithWitnessArenaFromWitnesses :: Int -> [Integer] -> [Integer] -> [Integer] -> [Integer] -> [Integer] -> ST s (SmithWitnessArena s) +newSmithWitnessArenaFromWitnesses matrixSize workEntries leftEntries rightEntries leftInverseEntries rightInverseEntries = do + work <- V.thaw (V.fromList workEntries) + leftRows <- V.thaw (V.fromList leftEntries) + rightRows <- V.thaw (V.fromList rightEntries) + leftInverseRows <- V.thaw (V.fromList leftInverseEntries) + rightInverseRows <- V.thaw (V.fromList rightInverseEntries) + pure + SmithWitnessArena + { smithWitnessRowCount = matrixSize, + smithWitnessColumnCount = matrixSize, + smithWitnessWork = work, + smithWitnessLeftRows = leftRows, + smithWitnessRightRows = rightRows, + smithWitnessLeftInverseRows = leftInverseRows, + smithWitnessRightInverseRows = rightInverseRows + } + +rowHermiteModulo :: Int -> Integer -> [Integer] -> Either SmithWitnessFailure [Integer] +rowHermiteModulo matrixSize modulusValue entries = + runST $ do + work <- V.thaw (V.fromList (fmap (centerResidue modulusValue) entries <> fmap (modulusValue *) (identityList matrixSize))) + failureValue <- rowHermiteModuloAt matrixSize modulusValue work 0 + case failureValue of + Just hermiteFailure -> pure (Left hermiteFailure) + Nothing -> Right . take (matrixSize * matrixSize) <$> readFlatVector work + +columnHermiteModulo :: Int -> Integer -> [Integer] -> Either SmithWitnessFailure [Integer] +columnHermiteModulo matrixSize modulusValue entries = + runST $ do + work <- V.thaw (V.fromList (augmentedColumnPool matrixSize modulusValue entries)) + failureValue <- columnHermiteModuloAt matrixSize modulusValue work 0 + case failureValue of + Just hermiteFailure -> pure (Left hermiteFailure) + Nothing -> Right . extractColumnPool matrixSize <$> V.freeze work + +augmentedColumnPool :: Int -> Integer -> [Integer] -> [Integer] +augmentedColumnPool matrixSize modulusValue entries = + concat + [ fmap (centerResidue modulusValue) (take matrixSize (drop (rowIndex * matrixSize) entries)) + <> [if columnIndex == rowIndex then modulusValue else 0 | columnIndex <- [0 .. matrixSize - 1]] + | rowIndex <- [0 .. matrixSize - 1] + ] + +extractColumnPool :: Int -> V.Vector Integer -> [Integer] +extractColumnPool matrixSize pool = + [ vectorValueAt pool (flatIndex (2 * matrixSize) rowIndex columnIndex) + | rowIndex <- [0 .. matrixSize - 1], + columnIndex <- [0 .. matrixSize - 1] + ] + +rowHermiteModuloAt :: forall s. Int -> Integer -> MV.MVector s Integer -> Int -> ST s (Maybe SmithWitnessFailure) +rowHermiteModuloAt matrixSize modulusValue work pivotIndex + | pivotIndex >= matrixSize = pure Nothing + | otherwise = do + pivotCandidate <- findColumnNonZeroModulo (2 * matrixSize) matrixSize work pivotIndex pivotIndex + case pivotCandidate of + Nothing -> pure (Just SmithWitnessPivotBecameZero) + Just pivotRow -> do + swapRowsVector matrixSize work pivotIndex pivotRow pivotIndex + signFailure <- normalizeModuloPivotRow matrixSize modulusValue work pivotIndex + case signFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + clearFailure <- clearColumnBelowModulo (2 * matrixSize) matrixSize modulusValue work pivotIndex + case clearFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + finalSignFailure <- normalizeModuloPivotRow matrixSize modulusValue work pivotIndex + case finalSignFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + reduceColumnAboveModulo matrixSize modulusValue work pivotIndex + rowHermiteModuloAt matrixSize modulusValue work (pivotIndex + 1) + +columnHermiteModuloAt :: forall s. Int -> Integer -> MV.MVector s Integer -> Int -> ST s (Maybe SmithWitnessFailure) +columnHermiteModuloAt matrixSize modulusValue work pivotIndex + | pivotIndex >= matrixSize = pure Nothing + | otherwise = do + pivotCandidate <- findRowNonZeroModulo (2 * matrixSize) work pivotIndex pivotIndex + case pivotCandidate of + Nothing -> pure (Just SmithWitnessPivotBecameZero) + Just pivotColumn -> do + swapColumnsVector (2 * matrixSize) work pivotIndex pivotColumn pivotIndex matrixSize + signFailure <- normalizeModuloPivotColumn (2 * matrixSize) matrixSize modulusValue work pivotIndex + case signFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + clearFailure <- clearRowRightModulo (2 * matrixSize) matrixSize modulusValue work pivotIndex + case clearFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + finalSignFailure <- normalizeModuloPivotColumn (2 * matrixSize) matrixSize modulusValue work pivotIndex + case finalSignFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + reduceRowLeftModulo (2 * matrixSize) matrixSize modulusValue work pivotIndex + columnHermiteModuloAt matrixSize modulusValue work (pivotIndex + 1) + +findColumnNonZeroModulo :: forall s. Int -> Int -> MV.MVector s Integer -> Int -> Int -> ST s (Maybe Int) +findColumnNonZeroModulo rowCount columnCount work columnIndex rowIndex + | rowIndex >= rowCount = pure Nothing + | otherwise = do + entryValue <- MV.unsafeRead work (flatIndex columnCount rowIndex columnIndex) + if entryValue == 0 + then findColumnNonZeroModulo rowCount columnCount work columnIndex (rowIndex + 1) + else pure (Just rowIndex) + +findRowNonZeroModulo :: forall s. Int -> MV.MVector s Integer -> Int -> Int -> ST s (Maybe Int) +findRowNonZeroModulo poolColumnCount work rowIndex columnIndex + | columnIndex >= poolColumnCount = pure Nothing + | otherwise = do + entryValue <- MV.unsafeRead work (flatIndex poolColumnCount rowIndex columnIndex) + if entryValue == 0 + then findRowNonZeroModulo poolColumnCount work rowIndex (columnIndex + 1) + else pure (Just columnIndex) + +normalizeModuloPivotRow :: forall s. Int -> Integer -> MV.MVector s Integer -> Int -> ST s (Maybe SmithWitnessFailure) +normalizeModuloPivotRow matrixSize modulusValue work pivotIndex = do + pivotValue <- MV.unsafeRead work (flatIndex matrixSize pivotIndex pivotIndex) + if pivotValue < 0 + then scaleRowModuloVector matrixSize modulusValue work pivotIndex (-1) pivotIndex *> pure Nothing + else pure Nothing + +normalizeModuloPivotColumn :: forall s. Int -> Int -> Integer -> MV.MVector s Integer -> Int -> ST s (Maybe SmithWitnessFailure) +normalizeModuloPivotColumn poolColumnCount rowCount modulusValue work pivotIndex = do + pivotValue <- MV.unsafeRead work (flatIndex poolColumnCount pivotIndex pivotIndex) + if pivotValue < 0 + then scaleColumnModuloVector poolColumnCount modulusValue work pivotIndex (-1) pivotIndex rowCount *> pure Nothing + else pure Nothing + +clearColumnBelowModulo :: forall s. Int -> Int -> Integer -> MV.MVector s Integer -> Int -> ST s (Maybe SmithWitnessFailure) +clearColumnBelowModulo rowCount columnCount modulusValue work pivotIndex = + scanRows (pivotIndex + 1) + where + scanRows :: Int -> ST s (Maybe SmithWitnessFailure) + scanRows rowIndex + | rowIndex >= rowCount = pure Nothing + | otherwise = do + entryValue <- MV.unsafeRead work (flatIndex columnCount rowIndex pivotIndex) + if entryValue == 0 + then scanRows (rowIndex + 1) + else do + pivotValue <- MV.unsafeRead work (flatIndex columnCount pivotIndex pivotIndex) + if pivotValue == 0 + then pure (Just SmithWitnessPivotBecameZero) + else do + let (quotientValue, _) = balancedDivMod entryValue pivotValue + if quotientValue /= 0 + then rowCombineModuloVector columnCount modulusValue work rowIndex pivotIndex quotientValue pivotIndex + else pure () + reducedEntry <- MV.unsafeRead work (flatIndex columnCount rowIndex pivotIndex) + if reducedEntry == 0 + then scanRows (rowIndex + 1) + else do + gcdFailure <- gcdCombineRowsModulo columnCount modulusValue work pivotIndex rowIndex pivotIndex + case gcdFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> scanRows (rowIndex + 1) + +clearRowRightModulo :: forall s. Int -> Int -> Integer -> MV.MVector s Integer -> Int -> ST s (Maybe SmithWitnessFailure) +clearRowRightModulo poolColumnCount rowCount modulusValue work pivotIndex = + scanColumns (pivotIndex + 1) + where + scanColumns :: Int -> ST s (Maybe SmithWitnessFailure) + scanColumns columnIndex + | columnIndex >= poolColumnCount = pure Nothing + | otherwise = do + entryValue <- MV.unsafeRead work (flatIndex poolColumnCount pivotIndex columnIndex) + if entryValue == 0 + then scanColumns (columnIndex + 1) + else do + pivotValue <- MV.unsafeRead work (flatIndex poolColumnCount pivotIndex pivotIndex) + if pivotValue == 0 + then pure (Just SmithWitnessPivotBecameZero) + else do + let (quotientValue, _) = balancedDivMod entryValue pivotValue + if quotientValue /= 0 + then columnCombineModuloVector poolColumnCount modulusValue work columnIndex pivotIndex quotientValue pivotIndex rowCount + else pure () + reducedEntry <- MV.unsafeRead work (flatIndex poolColumnCount pivotIndex columnIndex) + if reducedEntry == 0 + then scanColumns (columnIndex + 1) + else do + gcdFailure <- gcdCombineColumnsModulo poolColumnCount rowCount modulusValue work pivotIndex pivotIndex columnIndex + case gcdFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> scanColumns (columnIndex + 1) + +gcdCombineRowsModulo :: forall s. Int -> Integer -> MV.MVector s Integer -> Int -> Int -> Int -> ST s (Maybe SmithWitnessFailure) +gcdCombineRowsModulo matrixSize modulusValue work pivotRow candidateRow pivotColumn = do + pivotValue <- MV.unsafeRead work (flatIndex matrixSize pivotRow pivotColumn) + entryValue <- MV.unsafeRead work (flatIndex matrixSize candidateRow pivotColumn) + let (gcdValue, pivotCoefficient, entryCoefficient) = extendedGcdDomain pivotValue entryValue + case (exactQuotientMutable "row mod-det gcd pivot quotient" pivotValue gcdValue, exactQuotientMutable "row mod-det gcd entry quotient" entryValue gcdValue) of + (SmithExactQuotient pivotQuotient, SmithExactQuotient entryQuotient) -> do + rowPairTransformModuloVector matrixSize modulusValue work pivotRow candidateRow pivotCoefficient entryCoefficient (negate entryQuotient) pivotQuotient pivotColumn + pure Nothing + (SmithInexactQuotient failureValue, _) -> pure (Just failureValue) + (_, SmithInexactQuotient failureValue) -> pure (Just failureValue) + +gcdCombineColumnsModulo :: forall s. Int -> Int -> Integer -> MV.MVector s Integer -> Int -> Int -> Int -> ST s (Maybe SmithWitnessFailure) +gcdCombineColumnsModulo poolColumnCount rowCount modulusValue work pivotRow pivotColumn candidateColumn = do + pivotValue <- MV.unsafeRead work (flatIndex poolColumnCount pivotRow pivotColumn) + entryValue <- MV.unsafeRead work (flatIndex poolColumnCount pivotRow candidateColumn) + let (gcdValue, pivotCoefficient, entryCoefficient) = extendedGcdDomain pivotValue entryValue + case (exactQuotientMutable "column mod-det gcd pivot quotient" pivotValue gcdValue, exactQuotientMutable "column mod-det gcd entry quotient" entryValue gcdValue) of + (SmithExactQuotient pivotQuotient, SmithExactQuotient entryQuotient) -> do + columnPairTransformModuloVector poolColumnCount modulusValue work pivotColumn candidateColumn pivotCoefficient entryCoefficient (negate entryQuotient) pivotQuotient pivotRow rowCount + pure Nothing + (SmithInexactQuotient failureValue, _) -> pure (Just failureValue) + (_, SmithInexactQuotient failureValue) -> pure (Just failureValue) + +reduceColumnAboveModulo :: forall s. Int -> Integer -> MV.MVector s Integer -> Int -> ST s () +reduceColumnAboveModulo matrixSize modulusValue work pivotIndex = + scanRows 0 + where + scanRows :: Int -> ST s () + scanRows rowIndex + | rowIndex >= pivotIndex = pure () + | otherwise = do + pivotValue <- MV.unsafeRead work (flatIndex matrixSize pivotIndex pivotIndex) + entryValue <- MV.unsafeRead work (flatIndex matrixSize rowIndex pivotIndex) + if pivotValue == 0 + then scanRows (rowIndex + 1) + else do + let quotientValue = entryValue `div` pivotValue + if quotientValue /= 0 + then rowCombineModuloVector matrixSize modulusValue work rowIndex pivotIndex quotientValue pivotIndex + else pure () + scanRows (rowIndex + 1) + +reduceRowLeftModulo :: forall s. Int -> Int -> Integer -> MV.MVector s Integer -> Int -> ST s () +reduceRowLeftModulo poolColumnCount rowCount modulusValue work pivotIndex = + scanColumns 0 + where + scanColumns :: Int -> ST s () + scanColumns columnIndex + | columnIndex >= pivotIndex = pure () + | otherwise = do + pivotValue <- MV.unsafeRead work (flatIndex poolColumnCount pivotIndex pivotIndex) + entryValue <- MV.unsafeRead work (flatIndex poolColumnCount pivotIndex columnIndex) + if pivotValue == 0 + then scanColumns (columnIndex + 1) + else do + let quotientValue = entryValue `div` pivotValue + if quotientValue /= 0 + then columnCombineModuloVector poolColumnCount modulusValue work columnIndex pivotIndex quotientValue pivotIndex rowCount + else pure () + scanColumns (columnIndex + 1) + +rowCombineModuloVector :: Int -> Integer -> MV.MVector s Integer -> Int -> Int -> Integer -> Int -> ST s () +rowCombineModuloVector columnCount modulusValue entries targetRow sourceRow coefficient columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + let targetIndex = flatIndex columnCount targetRow columnIndex + sourceIndex = flatIndex columnCount sourceRow columnIndex + sourceValue <- MV.unsafeRead entries sourceIndex + if sourceValue == 0 + then pure () + else do + targetValue <- MV.unsafeRead entries targetIndex + MV.unsafeWrite entries targetIndex (centerResidue modulusValue (targetValue - coefficient * sourceValue)) + rowCombineModuloVector columnCount modulusValue entries targetRow sourceRow coefficient (columnIndex + 1) + +columnCombineModuloVector :: Int -> Integer -> MV.MVector s Integer -> Int -> Int -> Integer -> Int -> Int -> ST s () +columnCombineModuloVector columnCount modulusValue entries targetColumn sourceColumn coefficient rowIndex rowCount + | rowIndex >= rowCount = pure () + | otherwise = do + let targetIndex = flatIndex columnCount rowIndex targetColumn + sourceIndex = flatIndex columnCount rowIndex sourceColumn + sourceValue <- MV.unsafeRead entries sourceIndex + if sourceValue == 0 + then pure () + else do + targetValue <- MV.unsafeRead entries targetIndex + MV.unsafeWrite entries targetIndex (centerResidue modulusValue (targetValue - coefficient * sourceValue)) + columnCombineModuloVector columnCount modulusValue entries targetColumn sourceColumn coefficient (rowIndex + 1) rowCount + +rowPairTransformModuloVector :: Int -> Integer -> MV.MVector s Integer -> Int -> Int -> Integer -> Integer -> Integer -> Integer -> Int -> ST s () +rowPairTransformModuloVector columnCount modulusValue entries leftRow rightRow aa ab ba bb columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount leftRow columnIndex + rightIndex = flatIndex columnCount rightRow columnIndex + leftValue <- MV.unsafeRead entries leftIndex + rightValue <- MV.unsafeRead entries rightIndex + if leftValue == 0 && rightValue == 0 + then pure () + else do + MV.unsafeWrite entries leftIndex (centerResidue modulusValue (aa * leftValue + ab * rightValue)) + MV.unsafeWrite entries rightIndex (centerResidue modulusValue (ba * leftValue + bb * rightValue)) + rowPairTransformModuloVector columnCount modulusValue entries leftRow rightRow aa ab ba bb (columnIndex + 1) + +columnPairTransformModuloVector :: Int -> Integer -> MV.MVector s Integer -> Int -> Int -> Integer -> Integer -> Integer -> Integer -> Int -> Int -> ST s () +columnPairTransformModuloVector columnCount modulusValue entries leftColumn rightColumn aa ab ba bb rowIndex rowCount + | rowIndex >= rowCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount rowIndex leftColumn + rightIndex = flatIndex columnCount rowIndex rightColumn + leftValue <- MV.unsafeRead entries leftIndex + rightValue <- MV.unsafeRead entries rightIndex + if leftValue == 0 && rightValue == 0 + then pure () + else do + MV.unsafeWrite entries leftIndex (centerResidue modulusValue (aa * leftValue + ab * rightValue)) + MV.unsafeWrite entries rightIndex (centerResidue modulusValue (ba * leftValue + bb * rightValue)) + columnPairTransformModuloVector columnCount modulusValue entries leftColumn rightColumn aa ab ba bb (rowIndex + 1) rowCount + +scaleRowModuloVector :: Int -> Integer -> MV.MVector s Integer -> Int -> Integer -> Int -> ST s () +scaleRowModuloVector columnCount modulusValue entries rowIndex factor columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + let entryIndex = flatIndex columnCount rowIndex columnIndex + entryValue <- MV.unsafeRead entries entryIndex + MV.unsafeWrite entries entryIndex (centerResidue modulusValue (factor * entryValue)) + scaleRowModuloVector columnCount modulusValue entries rowIndex factor (columnIndex + 1) + +scaleColumnModuloVector :: Int -> Integer -> MV.MVector s Integer -> Int -> Integer -> Int -> Int -> ST s () +scaleColumnModuloVector columnCount modulusValue entries columnIndex factor rowIndex rowCount + | rowIndex >= rowCount = pure () + | otherwise = do + let entryIndex = flatIndex columnCount rowIndex columnIndex + entryValue <- MV.unsafeRead entries entryIndex + MV.unsafeWrite entries entryIndex (centerResidue modulusValue (factor * entryValue)) + scaleColumnModuloVector columnCount modulusValue entries columnIndex factor (rowIndex + 1) rowCount + +recoverTransform :: Int -> FastTransformOrientation -> [Integer] -> [Integer] -> Integer -> String -> Either SmithWitnessFailure [Integer] +recoverTransform matrixSize orientation leftEntries rightEntries coefficientBound context = + searchPrimes initialResidues 1 Nothing wordPrimeLadder + where + leftVector :: V.Vector Integer + leftVector = V.fromList leftEntries + + rightVector :: V.Vector Integer + rightVector = V.fromList rightEntries + + target :: Integer + target = max 2 (2 * coefficientBound + 1) + + initialResidues :: V.Vector Integer + initialResidues = V.replicate (matrixSize * matrixSize) 0 + + leftIsUpperTriangular :: Bool + leftIsUpperTriangular = matrixIsUpperTriangularVector matrixSize leftVector + + rightIsLowerTriangular :: Bool + rightIsLowerTriangular = matrixIsLowerTriangularVector matrixSize rightVector + + solveTransformPrime :: Word64 -> Maybe (U.Vector Word64) + solveTransformPrime primeValue = + case orientation of + FastInverseLeftTimesRight + | leftIsUpperTriangular -> solveUpperTriangularModuloPrime matrixSize primeValue leftVector rightVector + | otherwise -> do + leftInverse <- invertMatrixModuloPrime matrixSize primeValue leftVector + Just (matrixProductModuloPrime matrixSize primeValue leftInverse (residueVector primeValue rightVector)) + FastLeftTimesInverseRight + | rightIsLowerTriangular -> + transposeResidueMatrix matrixSize <$> solveUpperTriangularModuloPrime matrixSize primeValue transposedRight transposedLeft + | otherwise -> + transposeResidueMatrix matrixSize <$> solveRightQuotientTransposedModuloPrime matrixSize primeValue rightVector leftVector + + transposedLeft :: V.Vector Integer + transposedLeft = transposeIntegerMatrix matrixSize leftVector + + transposedRight :: V.Vector Integer + transposedRight = transposeIntegerMatrix matrixSize rightVector + + verifyCandidate :: Integer -> [Word64] -> [Integer] -> Either SmithWitnessFailure () + verifyCandidate knownModulus freshPrimes candidate = + case orientation of + FastLeftTimesInverseRight -> verifyProductModuloPrimesFrom knownModulus freshPrimes matrixSize context candidate rightEntries leftEntries + FastInverseLeftTimesRight -> verifyProductModuloPrimesFrom knownModulus freshPrimes matrixSize context leftEntries candidate rightEntries + + searchPrimes :: V.Vector Integer -> Integer -> Maybe (V.Vector Integer) -> [Word64] -> Either SmithWitnessFailure [Integer] + searchPrimes residues modulusValue previousLift primes = + case primes of + [] -> Left (SmithWitnessTransformRecoveryFailed (context <> ": prime ladder exhausted")) + primeValue : remainingPrimes -> + case solveTransformPrime primeValue of + Nothing -> searchPrimes residues modulusValue previousLift remainingPrimes + Just primeResidues -> do + nextResidues <- combineCrtVector residues modulusValue primeValue primeResidues + let nextModulus = modulusValue * toInteger primeValue + candidateLift = V.map (symmetricLiftInteger nextModulus) nextResidues + if Just candidateLift == previousLift + then case verifyCandidate nextModulus remainingPrimes (V.toList candidateLift) of + Right () -> Right (V.toList candidateLift) + Left _ -> continueSearch nextResidues nextModulus candidateLift remainingPrimes + else continueSearch nextResidues nextModulus candidateLift remainingPrimes + + continueSearch :: V.Vector Integer -> Integer -> V.Vector Integer -> [Word64] -> Either SmithWitnessFailure [Integer] + continueSearch nextResidues nextModulus candidateLift remainingPrimes + | nextModulus > target = + let liftEntries = V.toList candidateLift + in verifyCandidate nextModulus remainingPrimes liftEntries *> Right liftEntries + | otherwise = searchPrimes nextResidues nextModulus (Just candidateLift) remainingPrimes + +matrixIsUpperTriangularVector :: Int -> V.Vector Integer -> Bool +matrixIsUpperTriangularVector matrixSize entries = + and + [ V.unsafeIndex entries (flatIndex matrixSize rowIndex columnIndex) == 0 + | rowIndex <- [1 .. matrixSize - 1], + columnIndex <- [0 .. rowIndex - 1] + ] + +matrixIsLowerTriangularVector :: Int -> V.Vector Integer -> Bool +matrixIsLowerTriangularVector matrixSize entries = + and + [ V.unsafeIndex entries (flatIndex matrixSize rowIndex columnIndex) == 0 + | rowIndex <- [0 .. matrixSize - 2], + columnIndex <- [rowIndex + 1 .. matrixSize - 1] + ] + +transposeIntegerMatrix :: Int -> V.Vector Integer -> V.Vector Integer +transposeIntegerMatrix matrixSize entries = + V.generate + (matrixSize * matrixSize) + ( \entryIndex -> + let (rowIndex, columnIndex) = entryIndex `quotRem` matrixSize + in V.unsafeIndex entries (flatIndex matrixSize columnIndex rowIndex) + ) + +transposeResidueMatrix :: Int -> U.Vector Word64 -> U.Vector Word64 +transposeResidueMatrix matrixSize entries = + U.generate + (matrixSize * matrixSize) + ( \entryIndex -> + let (rowIndex, columnIndex) = entryIndex `quotRem` matrixSize + in U.unsafeIndex entries (flatIndex matrixSize columnIndex rowIndex) + ) + +solveUpperTriangularModuloPrime :: Int -> Word64 -> V.Vector Integer -> V.Vector Integer -> Maybe (U.Vector Word64) +solveUpperTriangularModuloPrime matrixSize primeValue leftEntries rightEntries = + if U.any (== 0) pivotResidues + then Nothing + else Just solvedEntries + where + leftResidues :: U.Vector Word64 + leftResidues = residueVector primeValue leftEntries + + rightResidues :: U.Vector Word64 + rightResidues = residueVector primeValue rightEntries + + pivotResidues :: U.Vector Word64 + pivotResidues = U.generate matrixSize (\axisIndex -> U.unsafeIndex leftResidues (flatIndex matrixSize axisIndex axisIndex)) + + solvedEntries :: U.Vector Word64 + solvedEntries = runST $ do + work <- MU.replicate (matrixSize * matrixSize) 0 + solveRowsBottomUp matrixSize primeValue leftResidues rightResidues pivotResidues work (matrixSize - 1) + U.freeze work + +solveRowsBottomUp :: forall s. Int -> Word64 -> U.Vector Word64 -> U.Vector Word64 -> U.Vector Word64 -> MU.MVector s Word64 -> Int -> ST s () +solveRowsBottomUp matrixSize primeValue leftResidues rightResidues pivotResidues work rowIndex + | rowIndex < 0 = pure () + | otherwise = do + let pivotInverse = modInverseWord primeValue (U.unsafeIndex pivotResidues rowIndex) + solveRowColumns matrixSize primeValue leftResidues rightResidues work rowIndex pivotInverse 0 + solveRowsBottomUp matrixSize primeValue leftResidues rightResidues pivotResidues work (rowIndex - 1) + +solveRowColumns :: forall s. Int -> Word64 -> U.Vector Word64 -> U.Vector Word64 -> MU.MVector s Word64 -> Int -> Word64 -> Int -> ST s () +solveRowColumns matrixSize primeValue leftResidues rightResidues work rowIndex pivotInverse columnIndex + | columnIndex >= matrixSize = pure () + | otherwise = do + accumulated <- accumulateSolvedTail matrixSize primeValue leftResidues work rowIndex columnIndex (rowIndex + 1) 0 + let rhsValue = U.unsafeIndex rightResidues (flatIndex matrixSize rowIndex columnIndex) + MU.unsafeWrite work (flatIndex matrixSize rowIndex columnIndex) (modMul primeValue pivotInverse (modSubWord primeValue rhsValue accumulated)) + solveRowColumns matrixSize primeValue leftResidues rightResidues work rowIndex pivotInverse (columnIndex + 1) + +accumulateSolvedTail :: forall s. Int -> Word64 -> U.Vector Word64 -> MU.MVector s Word64 -> Int -> Int -> Int -> Word64 -> ST s Word64 +accumulateSolvedTail matrixSize primeValue leftResidues work rowIndex columnIndex sharedIndex accumulator + | sharedIndex >= matrixSize = pure accumulator + | otherwise = do + solvedValue <- MU.unsafeRead work (flatIndex matrixSize sharedIndex columnIndex) + accumulateSolvedTail matrixSize primeValue leftResidues work rowIndex columnIndex (sharedIndex + 1) (modAddWord primeValue accumulator (modMul primeValue (U.unsafeIndex leftResidues (flatIndex matrixSize rowIndex sharedIndex)) solvedValue)) + +residueVector :: Word64 -> V.Vector Integer -> U.Vector Word64 +residueVector primeValue entries = + U.generate (V.length entries) (integerResidueWord primeValue . V.unsafeIndex entries) + +verifyProductModuloPrimesFrom :: Integer -> [Word64] -> Int -> String -> [Integer] -> [Integer] -> [Integer] -> Either SmithWitnessFailure () +verifyProductModuloPrimesFrom priorModulus freshPrimes matrixSize context leftEntries rightEntries expectedEntries = + checkPrimes priorModulus freshPrimes + where + leftVector :: V.Vector Integer + leftVector = V.fromList leftEntries + + rightVector :: V.Vector Integer + rightVector = V.fromList rightEntries + + expectedVector :: V.Vector Integer + expectedVector = V.fromList expectedEntries + + entryBound :: Integer + entryBound = + toInteger matrixSize * maxAbsEntry leftEntries * maxAbsEntry rightEntries + maxAbsEntry expectedEntries + 1 + + checkPrimes :: Integer -> [Word64] -> Either SmithWitnessFailure () + checkPrimes modulusValue primes + | modulusValue > entryBound = Right () + | otherwise = + case primes of + [] -> Left (SmithWitnessVerificationFailed (context <> ": verification prime ladder exhausted")) + primeValue : remainingPrimes -> + let productResidues = matrixProductModuloPrime matrixSize primeValue (residueVector primeValue leftVector) (residueVector primeValue rightVector) + expectedResidues = residueVector primeValue expectedVector + in if productResidues == expectedResidues + then checkPrimes (modulusValue * toInteger primeValue) remainingPrimes + else Left (SmithWitnessVerificationFailed context) + +invertMatrixModuloPrime :: Int -> Word64 -> V.Vector Integer -> Maybe (U.Vector Word64) +invertMatrixModuloPrime matrixSize primeValue entries = + runST $ do + work <- MU.replicate (matrixSize * matrixSize * 2) 0 + writeAugmentedModuloMatrix matrixSize primeValue entries work 0 + invertFailure <- invertModuloAt matrixSize primeValue work 0 + case invertFailure of + Just () -> pure Nothing + Nothing -> Just <$> readInverseModuloMatrix matrixSize work + +solveRightQuotientTransposedModuloPrime :: Int -> Word64 -> V.Vector Integer -> V.Vector Integer -> Maybe (U.Vector Word64) +solveRightQuotientTransposedModuloPrime matrixSize primeValue denominatorEntries numeratorEntries = + runST $ do + work <- MU.replicate (matrixSize * matrixSize * 2) 0 + writeAugmentedTransposedPairModuloMatrix matrixSize primeValue denominatorEntries numeratorEntries work 0 + invertFailure <- invertModuloAt matrixSize primeValue work 0 + case invertFailure of + Just () -> pure Nothing + Nothing -> Just <$> readInverseModuloMatrix matrixSize work + +writeAugmentedTransposedPairModuloMatrix :: forall s. Int -> Word64 -> V.Vector Integer -> V.Vector Integer -> MU.MVector s Word64 -> Int -> ST s () +writeAugmentedTransposedPairModuloMatrix matrixSize primeValue denominatorEntries numeratorEntries work entryIndex + | entryIndex >= matrixSize * matrixSize = pure () + | otherwise = do + let (rowIndex, columnIndex) = entryIndex `quotRem` matrixSize + transposedIndex = flatIndex matrixSize columnIndex rowIndex + MU.unsafeWrite work (augmentedIndex matrixSize rowIndex columnIndex) (integerResidueWord primeValue (V.unsafeIndex denominatorEntries transposedIndex)) + MU.unsafeWrite work (augmentedIndex matrixSize rowIndex (columnIndex + matrixSize)) (integerResidueWord primeValue (V.unsafeIndex numeratorEntries transposedIndex)) + writeAugmentedTransposedPairModuloMatrix matrixSize primeValue denominatorEntries numeratorEntries work (entryIndex + 1) + +writeAugmentedModuloMatrix :: forall s. Int -> Word64 -> V.Vector Integer -> MU.MVector s Word64 -> Int -> ST s () +writeAugmentedModuloMatrix matrixSize primeValue entries work entryIndex + | entryIndex >= matrixSize * matrixSize = pure () + | otherwise = do + let (rowIndex, columnIndex) = entryIndex `quotRem` matrixSize + sourceValue = vectorValueAt entries entryIndex + MU.unsafeWrite work (augmentedIndex matrixSize rowIndex columnIndex) (integerResidueWord primeValue sourceValue) + MU.unsafeWrite work (augmentedIndex matrixSize rowIndex (columnIndex + matrixSize)) (if rowIndex == columnIndex then 1 else 0) + writeAugmentedModuloMatrix matrixSize primeValue entries work (entryIndex + 1) + +invertModuloAt :: forall s. Int -> Word64 -> MU.MVector s Word64 -> Int -> ST s (Maybe ()) +invertModuloAt matrixSize primeValue work pivotIndex + | pivotIndex >= matrixSize = pure Nothing + | otherwise = do + pivotCandidate <- findModuloPivot matrixSize work pivotIndex pivotIndex + case pivotCandidate of + Nothing -> pure (Just ()) + Just pivotRow -> do + swapAugmentedRows matrixSize work pivotIndex pivotRow 0 + pivotValue <- MU.unsafeRead work (augmentedIndex matrixSize pivotIndex pivotIndex) + let inversePivot = modInverseWord primeValue pivotValue + scaleAugmentedRow matrixSize primeValue work pivotIndex inversePivot 0 + eliminateModuloColumn matrixSize primeValue work pivotIndex 0 + invertModuloAt matrixSize primeValue work (pivotIndex + 1) + +findModuloPivot :: forall s. Int -> MU.MVector s Word64 -> Int -> Int -> ST s (Maybe Int) +findModuloPivot matrixSize work pivotColumn rowIndex + | rowIndex >= matrixSize = pure Nothing + | otherwise = do + entryValue <- MU.unsafeRead work (augmentedIndex matrixSize rowIndex pivotColumn) + if entryValue == 0 + then findModuloPivot matrixSize work pivotColumn (rowIndex + 1) + else pure (Just rowIndex) + +swapAugmentedRows :: forall s. Int -> MU.MVector s Word64 -> Int -> Int -> Int -> ST s () +swapAugmentedRows matrixSize work leftRow rightRow columnIndex + | leftRow == rightRow = pure () + | columnIndex >= 2 * matrixSize = pure () + | otherwise = do + let leftIndex = augmentedIndex matrixSize leftRow columnIndex + rightIndex = augmentedIndex matrixSize rightRow columnIndex + leftValue <- MU.unsafeRead work leftIndex + rightValue <- MU.unsafeRead work rightIndex + MU.unsafeWrite work leftIndex rightValue + MU.unsafeWrite work rightIndex leftValue + swapAugmentedRows matrixSize work leftRow rightRow (columnIndex + 1) + +scaleAugmentedRow :: forall s. Int -> Word64 -> MU.MVector s Word64 -> Int -> Word64 -> Int -> ST s () +scaleAugmentedRow matrixSize primeValue work rowIndex factor columnIndex + | columnIndex >= 2 * matrixSize = pure () + | otherwise = do + let entryIndex = augmentedIndex matrixSize rowIndex columnIndex + entryValue <- MU.unsafeRead work entryIndex + MU.unsafeWrite work entryIndex (modMul primeValue factor entryValue) + scaleAugmentedRow matrixSize primeValue work rowIndex factor (columnIndex + 1) + +eliminateModuloColumn :: forall s. Int -> Word64 -> MU.MVector s Word64 -> Int -> Int -> ST s () +eliminateModuloColumn matrixSize primeValue work pivotIndex rowIndex + | rowIndex >= matrixSize = pure () + | rowIndex == pivotIndex = eliminateModuloColumn matrixSize primeValue work pivotIndex (rowIndex + 1) + | otherwise = do + factor <- MU.unsafeRead work (augmentedIndex matrixSize rowIndex pivotIndex) + if factor == 0 + then eliminateModuloColumn matrixSize primeValue work pivotIndex (rowIndex + 1) + else eliminateModuloRow matrixSize primeValue work pivotIndex rowIndex factor 0 *> eliminateModuloColumn matrixSize primeValue work pivotIndex (rowIndex + 1) + +eliminateModuloRow :: forall s. Int -> Word64 -> MU.MVector s Word64 -> Int -> Int -> Word64 -> Int -> ST s () +eliminateModuloRow matrixSize primeValue work pivotRow targetRow factor columnIndex + | columnIndex >= 2 * matrixSize = pure () + | otherwise = do + let targetIndex = augmentedIndex matrixSize targetRow columnIndex + pivotEntryIndex = augmentedIndex matrixSize pivotRow columnIndex + targetValue <- MU.unsafeRead work targetIndex + pivotValue <- MU.unsafeRead work pivotEntryIndex + MU.unsafeWrite work targetIndex (modSubWord primeValue targetValue (modMul primeValue factor pivotValue)) + eliminateModuloRow matrixSize primeValue work pivotRow targetRow factor (columnIndex + 1) + +readInverseModuloMatrix :: forall s. Int -> MU.MVector s Word64 -> ST s (U.Vector Word64) +readInverseModuloMatrix matrixSize work = + U.generateM + (matrixSize * matrixSize) + ( \entryIndex -> do + let (rowIndex, columnIndex) = entryIndex `quotRem` matrixSize + MU.unsafeRead work (augmentedIndex matrixSize rowIndex (columnIndex + matrixSize)) + ) + +matrixProductModuloPrime :: Int -> Word64 -> U.Vector Word64 -> U.Vector Word64 -> U.Vector Word64 +matrixProductModuloPrime matrixSize primeValue leftEntries rightEntries = + U.generate + (matrixSize * matrixSize) + ( \entryIndex -> + let (rowIndex, columnIndex) = entryIndex `quotRem` matrixSize + in dotModuloPrime matrixSize primeValue leftEntries rightEntries rowIndex columnIndex 0 0 + ) + +dotModuloPrime :: Int -> Word64 -> U.Vector Word64 -> U.Vector Word64 -> Int -> Int -> Int -> Word64 -> Word64 +dotModuloPrime matrixSize primeValue leftEntries rightEntries rowIndex columnIndex sharedIndex accumulator + | sharedIndex >= matrixSize = accumulator + | otherwise = + let leftValue = U.unsafeIndex leftEntries (flatIndex matrixSize rowIndex sharedIndex) + nextAccumulator = + if leftValue == 0 + then accumulator + else modAddWord primeValue accumulator (modMul primeValue leftValue (U.unsafeIndex rightEntries (flatIndex matrixSize sharedIndex columnIndex))) + in dotModuloPrime matrixSize primeValue leftEntries rightEntries rowIndex columnIndex (sharedIndex + 1) nextAccumulator + +combineCrtVector :: V.Vector Integer -> Integer -> Word64 -> U.Vector Word64 -> Either SmithWitnessFailure (V.Vector Integer) +combineCrtVector residues modulusValue primeValue primeResidues + | V.length residues /= U.length primeResidues = Left (SmithWitnessTransformRecoveryFailed "CRT residue vector shape mismatch") + | modulusSection == 0 = Left (SmithWitnessTransformRecoveryFailed "CRT modulus section vanished") + | otherwise = Right (V.imap (\entryIndex residueValue -> combineEntry residueValue (U.unsafeIndex primeResidues entryIndex)) residues) + where + primeInteger :: Integer + primeInteger = toInteger primeValue + + modulusSection :: Word64 + modulusSection = fromInteger (modulusValue `mod` primeInteger) + + inverseValue :: Integer + inverseValue = toInteger (modInverseWord primeValue modulusSection) + + combineEntry :: Integer -> Word64 -> Integer + combineEntry residueValue primeResidue = + let deltaValue = (toInteger primeResidue - residueValue) `mod` primeInteger + correction = (deltaValue * inverseValue) `mod` primeInteger + in residueValue + modulusValue * correction + +transformRecoveryBound :: Int -> Integer -> [Integer] -> [Integer] -> Integer +transformRecoveryBound matrixSize modulusValue numeratorEntries denominatorEntries = + (toInteger matrixSize * maxAbsEntry numeratorEntries * hadamardMinorBoundDimension (matrixSize - 1) matrixSize denominatorEntries) `quot` max 1 (modulusValue `quot` 2) + 1 + +hadamardMinorBoundDimension :: Int -> Int -> [Integer] -> Integer +hadamardMinorBoundDimension minorDimension columnCount entries = + powerOfTwoSquareRootBound (product (takeLargestWitness minorDimension (rowSquaredNorms columnCount entries))) + +powerOfTwoSquareRootBound :: Integer -> Integer +powerOfTwoSquareRootBound value + | value <= 1 = max 0 value + | otherwise = narrow 0 (expand 1) + where + exceeds :: Int -> Bool + exceeds exponentValue = (1 :: Integer) `shiftL` (2 * exponentValue) > value + + expand :: Int -> Int + expand exponentValue + | exceeds exponentValue = exponentValue + | otherwise = expand (2 * exponentValue) + + narrow :: Int -> Int -> Integer + narrow low high + | high - low <= 1 = 1 `shiftL` high + | exceeds middle = narrow low middle + | otherwise = narrow middle high + where + middle :: Int + middle = (low + high) `quot` 2 + +rowSquaredNorms :: Int -> [Integer] -> [Integer] +rowSquaredNorms columnCount entries + | columnCount <= 0 = [] + | otherwise = rowNorms entries + where + rowNorms :: [Integer] -> [Integer] + rowNorms [] = [] + rowNorms remaining = + let (rowEntries, rest) = splitAt columnCount remaining + in foldl' (\accumulator entryValue -> accumulator + entryValue * entryValue) 0 rowEntries : rowNorms rest + +takeLargestWitness :: Int -> [Integer] -> [Integer] +takeLargestWitness count values = + take count (descendingInsertionSort values) + +descendingInsertionSort :: [Integer] -> [Integer] +descendingInsertionSort = + foldr insertDescending [] + +insertDescending :: Integer -> [Integer] -> [Integer] +insertDescending value values = + case values of + [] -> [value] + currentValue : remainingValues -> + if value >= currentValue + then value : values + else currentValue : insertDescending value remainingValues + +matrixProduct :: Int -> Int -> Int -> [Integer] -> [Integer] -> [Integer] +matrixProduct rowCount sharedCount columnCount leftEntries rightEntries = + matrixProductVector rowCount sharedCount columnCount (V.fromList leftEntries) (V.fromList rightEntries) + +matrixProductVector :: Int -> Int -> Int -> V.Vector Integer -> V.Vector Integer -> [Integer] +matrixProductVector rowCount sharedCount columnCount leftEntries rightEntries = + [ dotProductEntry sharedCount columnCount leftEntries rightEntries rowIndex columnIndex 0 0 + | rowIndex <- [0 .. rowCount - 1], + columnIndex <- [0 .. columnCount - 1] + ] + +dotProductEntry :: Int -> Int -> V.Vector Integer -> V.Vector Integer -> Int -> Int -> Int -> Integer -> Integer +dotProductEntry sharedCount columnCount leftEntries rightEntries rowIndex columnIndex sharedIndex accumulator + | sharedIndex >= sharedCount = accumulator + | otherwise = + let leftValue = vectorValueAt leftEntries (rowIndex * sharedCount + sharedIndex) + nextAccumulator = + if leftValue == 0 + then accumulator + else accumulator + leftValue * vectorValueAt rightEntries (sharedIndex * columnCount + columnIndex) + in dotProductEntry sharedCount columnCount leftEntries rightEntries rowIndex columnIndex (sharedIndex + 1) nextAccumulator + +matrixIsDiagonal :: Int -> [Integer] -> Bool +matrixIsDiagonal matrixSize entries = + and + [ rowIndex == columnIndex || valueAt entries (flatIndex matrixSize rowIndex columnIndex) == 0 + | rowIndex <- [0 .. matrixSize - 1], + columnIndex <- [0 .. matrixSize - 1] + ] + +identityList :: Int -> [Integer] +identityList matrixSize = + [ if rowIndex == columnIndex then 1 else 0 + | rowIndex <- [0 .. matrixSize - 1], + columnIndex <- [0 .. matrixSize - 1] + ] + +maxAbsEntry :: [Integer] -> Integer +maxAbsEntry = + foldl' (\current entryValue -> max current (abs entryValue)) 0 + +centerResidue :: Integer -> Integer -> Integer +centerResidue modulusValue value + | modulusValue <= 1 = value + | doubled > modulusValue = residueValue - modulusValue + | otherwise = residueValue + where + residueValue :: Integer + residueValue = value `mod` modulusValue + + doubled :: Integer + doubled = 2 * residueValue + +symmetricLiftInteger :: Integer -> Integer -> Integer +symmetricLiftInteger modulusValue residueValue + | 2 * residueValue > modulusValue = residueValue - modulusValue + | otherwise = residueValue + +modAddWord :: Word64 -> Word64 -> Word64 -> Word64 +modAddWord primeValue leftValue rightValue = + let sumValue = leftValue + rightValue + in if sumValue >= primeValue + then sumValue - primeValue + else sumValue + +modSubWord :: Word64 -> Word64 -> Word64 -> Word64 +modSubWord primeValue leftValue rightValue + | leftValue >= rightValue = leftValue - rightValue + | otherwise = primeValue - (rightValue - leftValue) + +augmentedIndex :: Int -> Int -> Int -> Int +augmentedIndex matrixSize rowIndex columnIndex = + rowIndex * (2 * matrixSize) + columnIndex + +valueAt :: [Integer] -> Int -> Integer +valueAt values indexValue = + maybe 0 id (values !? indexValue) + + +vectorValueAt :: V.Vector Integer -> Int -> Integer +vectorValueAt values indexValue = + maybe 0 id (values V.!? indexValue) + +newSmithWitnessArena :: Int -> Int -> [Integer] -> ST s (SmithWitnessArena s) +newSmithWitnessArena rowCount columnCount entries = do + work <- V.thaw (V.fromList entries) + leftRows <- V.thaw (identityVector rowCount) + rightRows <- V.thaw (identityVector columnCount) + leftInverseRows <- V.thaw (identityVector rowCount) + rightInverseRows <- V.thaw (identityVector columnCount) + pure + SmithWitnessArena + { smithWitnessRowCount = rowCount, + smithWitnessColumnCount = columnCount, + smithWitnessWork = work, + smithWitnessLeftRows = leftRows, + smithWitnessRightRows = rightRows, + smithWitnessLeftInverseRows = leftInverseRows, + smithWitnessRightInverseRows = rightInverseRows + } + +identityVector :: Int -> V.Vector Integer +identityVector sizeValue = + V.generate + (sizeValue * sizeValue) + ( \entryIndex -> + let (rowIndex, columnIndex) = entryIndex `quotRem` sizeValue + in if rowIndex == columnIndex then 1 else 0 + ) + +flatIndex :: Int -> Int -> Int -> Int +flatIndex columnCount rowIndex columnIndex = + rowIndex * columnCount + columnIndex + +readWorkEntry :: SmithWitnessArena s -> Int -> Int -> ST s Integer +readWorkEntry arenaValue rowIndex columnIndex = + MV.read (smithWitnessWork arenaValue) (flatIndex (smithWitnessColumnCount arenaValue) rowIndex columnIndex) + +entryIsZeroMutable :: SmithWitnessArena s -> Int -> Int -> ST s Bool +entryIsZeroMutable arenaValue rowIndex columnIndex = + (== 0) <$> readWorkEntry arenaValue rowIndex columnIndex + +readFlatVector :: forall s. MV.MVector s Integer -> ST s [Integer] +readFlatVector entries = + readFlatAt 0 [] + where + entryCount :: Int + entryCount = MV.length entries + + readFlatAt :: Int -> [Integer] -> ST s [Integer] + readFlatAt entryIndex values + | entryIndex >= entryCount = pure (reverse values) + | otherwise = do + entryValue <- MV.read entries entryIndex + readFlatAt (entryIndex + 1) (entryValue : values) + +alternationBudget :: SmithWitnessArena s -> Int +alternationBudget arenaValue = + 64 + 2 * (smithWitnessRowCount arenaValue + smithWitnessColumnCount arenaValue) + +normalizationBudget :: SmithWitnessArena s -> Int +normalizationBudget arenaValue = + max 1 (smithWitnessRowCount arenaValue * smithWitnessColumnCount arenaValue * 16) + +alternatingHermiteMutable :: forall s. Int -> SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure) +alternatingHermiteMutable remainingBudget arenaValue + | remainingBudget <= 0 = pure (Just (SmithWitnessBudgetExhausted "hermite alternation")) + | otherwise = do + rowFailure <- rowHermitePhaseMutable arenaValue + case rowFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + columnFailure <- columnHermitePhaseMutable arenaValue + case columnFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + cleared <- offDiagonalClearMutable arenaValue + if cleared + then pure Nothing + else alternatingHermiteMutable (remainingBudget - 1) arenaValue + +rowHermitePhaseMutable :: forall s. SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure) +rowHermitePhaseMutable arenaValue = + phaseStep 0 + where + diagonalSize :: Int + diagonalSize = min (smithWitnessRowCount arenaValue) (smithWitnessColumnCount arenaValue) + + phaseStep :: Int -> ST s (Maybe SmithWitnessFailure) + phaseStep pivotIndex + | pivotIndex >= diagonalSize = backwardReduceAboveMutable pivotIndex arenaValue *> pure Nothing + | otherwise = do + pivotCandidate <- findPivotMutable pivotIndex pivotIndex arenaValue + case pivotCandidate of + Nothing -> backwardReduceAboveMutable pivotIndex arenaValue *> pure Nothing + Just pivotValue -> do + swapRowsWitnessed pivotIndex (smithPivotRowIndex pivotValue) arenaValue + swapColumnsWitnessed pivotIndex (smithPivotColumnIndex pivotValue) arenaValue + signFailure <- normalizePivotSignMutable pivotIndex pivotIndex arenaValue + case signFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + clearFailure <- clearColumnBelowMutable pivotIndex arenaValue + case clearFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> phaseStep (pivotIndex + 1) + +columnHermitePhaseMutable :: forall s. SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure) +columnHermitePhaseMutable arenaValue = + phaseStep 0 + where + diagonalSize :: Int + diagonalSize = min (smithWitnessRowCount arenaValue) (smithWitnessColumnCount arenaValue) + + phaseStep :: Int -> ST s (Maybe SmithWitnessFailure) + phaseStep pivotIndex + | pivotIndex >= diagonalSize = backwardReduceLeftMutable pivotIndex arenaValue *> pure Nothing + | otherwise = do + pivotCandidate <- findPivotMutable pivotIndex pivotIndex arenaValue + case pivotCandidate of + Nothing -> backwardReduceLeftMutable pivotIndex arenaValue *> pure Nothing + Just pivotValue -> do + swapRowsWitnessed pivotIndex (smithPivotRowIndex pivotValue) arenaValue + swapColumnsWitnessed pivotIndex (smithPivotColumnIndex pivotValue) arenaValue + signFailure <- normalizePivotSignMutable pivotIndex pivotIndex arenaValue + case signFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + clearFailure <- clearRowRightMutable pivotIndex arenaValue + case clearFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> phaseStep (pivotIndex + 1) + +clearColumnBelowMutable :: forall s. Int -> SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure) +clearColumnBelowMutable pivotIndex arenaValue = + scanRows (pivotIndex + 1) + where + scanRows :: Int -> ST s (Maybe SmithWitnessFailure) + scanRows rowIndex + | rowIndex >= smithWitnessRowCount arenaValue = pure Nothing + | otherwise = do + entryValue <- readWorkEntry arenaValue rowIndex pivotIndex + if entryValue == 0 + then scanRows (rowIndex + 1) + else do + pivotValue <- readWorkEntry arenaValue pivotIndex pivotIndex + if pivotValue == 0 + then pure (Just SmithWitnessPivotBecameZero) + else do + let (quotientValue, remainderValue) = balancedDivMod entryValue pivotValue + if quotientValue /= 0 + then rowCombineWitnessed rowIndex pivotIndex quotientValue arenaValue + else pure () + if remainderValue == 0 + then scanRows (rowIndex + 1) + else do + gcdFailure <- gcdCombineRowsWitnessed pivotIndex rowIndex pivotIndex arenaValue + case gcdFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> scanRows (rowIndex + 1) + +clearRowRightMutable :: forall s. Int -> SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure) +clearRowRightMutable pivotIndex arenaValue = + scanColumns (pivotIndex + 1) + where + scanColumns :: Int -> ST s (Maybe SmithWitnessFailure) + scanColumns columnIndex + | columnIndex >= smithWitnessColumnCount arenaValue = pure Nothing + | otherwise = do + entryValue <- readWorkEntry arenaValue pivotIndex columnIndex + if entryValue == 0 + then scanColumns (columnIndex + 1) + else do + pivotValue <- readWorkEntry arenaValue pivotIndex pivotIndex + if pivotValue == 0 + then pure (Just SmithWitnessPivotBecameZero) + else do + let (quotientValue, remainderValue) = balancedDivMod entryValue pivotValue + if quotientValue /= 0 + then columnCombineWitnessed columnIndex pivotIndex quotientValue arenaValue + else pure () + if remainderValue == 0 + then scanColumns (columnIndex + 1) + else do + gcdFailure <- gcdCombineColumnsWitnessed pivotIndex pivotIndex columnIndex arenaValue + case gcdFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> scanColumns (columnIndex + 1) + +backwardReduceAboveMutable :: forall s. Int -> SmithWitnessArena s -> ST s () +backwardReduceAboveMutable settledCount arenaValue = + scanRows (settledCount - 1) + where + scanRows :: Int -> ST s () + scanRows rowIndex + | rowIndex < 0 = pure () + | otherwise = do + scanColumns rowIndex (rowIndex + 1) + scanRows (rowIndex - 1) + + scanColumns :: Int -> Int -> ST s () + scanColumns rowIndex columnIndex + | columnIndex >= settledCount = pure () + | otherwise = do + pivotValue <- readWorkEntry arenaValue columnIndex columnIndex + if pivotValue == 0 + then scanColumns rowIndex (columnIndex + 1) + else do + entryValue <- readWorkEntry arenaValue rowIndex columnIndex + let (quotientValue, _) = balancedDivMod entryValue pivotValue + if quotientValue /= 0 + then rowCombineWitnessed rowIndex columnIndex quotientValue arenaValue + else pure () + scanColumns rowIndex (columnIndex + 1) + +backwardReduceLeftMutable :: forall s. Int -> SmithWitnessArena s -> ST s () +backwardReduceLeftMutable settledCount arenaValue = + scanColumns (settledCount - 1) + where + scanColumns :: Int -> ST s () + scanColumns columnIndex + | columnIndex < 0 = pure () + | otherwise = do + scanRows columnIndex (columnIndex + 1) + scanColumns (columnIndex - 1) + + scanRows :: Int -> Int -> ST s () + scanRows columnIndex rowIndex + | rowIndex >= settledCount = pure () + | otherwise = do + pivotValue <- readWorkEntry arenaValue rowIndex rowIndex + if pivotValue == 0 + then scanRows columnIndex (rowIndex + 1) + else do + entryValue <- readWorkEntry arenaValue rowIndex columnIndex + let (quotientValue, _) = balancedDivMod entryValue pivotValue + if quotientValue /= 0 + then columnCombineWitnessed columnIndex rowIndex quotientValue arenaValue + else pure () + scanRows columnIndex (rowIndex + 1) + +offDiagonalClearMutable :: forall s. SmithWitnessArena s -> ST s Bool +offDiagonalClearMutable arenaValue = + scanRows 0 + where + scanRows :: Int -> ST s Bool + scanRows rowIndex + | rowIndex >= smithWitnessRowCount arenaValue = pure True + | otherwise = do + rowClear <- scanColumns rowIndex 0 + if rowClear + then scanRows (rowIndex + 1) + else pure False + + scanColumns :: Int -> Int -> ST s Bool + scanColumns rowIndex columnIndex + | columnIndex >= smithWitnessColumnCount arenaValue = pure True + | rowIndex == columnIndex = scanColumns rowIndex (columnIndex + 1) + | otherwise = do + isZeroEntry <- entryIsZeroMutable arenaValue rowIndex columnIndex + if isZeroEntry + then scanColumns rowIndex (columnIndex + 1) + else pure False + +findPivotMutable :: forall s. Int -> Int -> SmithWitnessArena s -> ST s (Maybe SmithPivot) +findPivotMutable startRow startColumn arenaValue = + fmap fst <$> scanRows startRow Nothing + where + scanRows :: Int -> Maybe (SmithPivot, Integer) -> ST s (Maybe (SmithPivot, Integer)) + scanRows rowIndex bestValue + | rowIndex >= smithWitnessRowCount arenaValue = pure bestValue + | otherwise = do + rowBest <- scanColumns rowIndex startColumn bestValue + scanRows (rowIndex + 1) rowBest + + scanColumns :: Int -> Int -> Maybe (SmithPivot, Integer) -> ST s (Maybe (SmithPivot, Integer)) + scanColumns rowIndex columnIndex bestValue + | columnIndex >= smithWitnessColumnCount arenaValue = pure bestValue + | otherwise = do + entryValue <- readWorkEntry arenaValue rowIndex columnIndex + let nextBest = + if entryValue == 0 + then bestValue + else betterPivot bestValue (SmithPivot rowIndex columnIndex, abs entryValue) + scanColumns rowIndex (columnIndex + 1) nextBest + +betterPivot :: Maybe (SmithPivot, Integer) -> (SmithPivot, Integer) -> Maybe (SmithPivot, Integer) +betterPivot bestValue candidateValue = + case bestValue of + Nothing -> Just candidateValue + Just currentValue -> + if pivotOrderingKey candidateValue < pivotOrderingKey currentValue + then Just candidateValue + else bestValue + +pivotOrderingKey :: (SmithPivot, Integer) -> (Integer, Int, Int) +pivotOrderingKey (pivotValue, magnitudeValue) = + (magnitudeValue, smithPivotRowIndex pivotValue, smithPivotColumnIndex pivotValue) + +normalizePivotMutable :: Int -> Int -> Int -> SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure) +normalizePivotMutable pivotRow pivotColumn remainingBudget arenaValue + | remainingBudget <= 0 = pure (Just (SmithWitnessBudgetExhausted "normalization")) + | otherwise = do + signFailure <- normalizePivotSignMutable pivotRow pivotColumn arenaValue + case signFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + (columnFailure, columnChanged) <- reduceDiagonalColumnMutable pivotRow arenaValue + case columnFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + (rowFailure, rowChanged) <- reduceDiagonalRowMutable pivotRow arenaValue + case rowFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + clearedColumn <- columnClearedMutable pivotRow pivotColumn arenaValue + clearedRow <- rowClearedMutable pivotRow pivotColumn arenaValue + if clearedColumn && clearedRow + then pure Nothing + else + if columnChanged || rowChanged + then normalizePivotMutable pivotRow pivotColumn (remainingBudget - 1) arenaValue + else pure (Just SmithWitnessNormalizationStalled) + +normalizePivotSignMutable :: Int -> Int -> SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure) +normalizePivotSignMutable pivotRow pivotColumn arenaValue = do + pivotValue <- readWorkEntry arenaValue pivotRow pivotColumn + if pivotValue < 0 + then scaleRowWitnessed pivotRow (-1) arenaValue *> pure Nothing + else pure Nothing + +reduceDiagonalColumnMutable :: forall s. Int -> SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure, Bool) +reduceDiagonalColumnMutable diagonalIndex arenaValue = + scanRows 0 False + where + scanRows :: Int -> Bool -> ST s (Maybe SmithWitnessFailure, Bool) + scanRows rowIndex changed + | rowIndex >= smithWitnessRowCount arenaValue = pure (Nothing, changed) + | rowIndex == diagonalIndex = scanRows (rowIndex + 1) changed + | otherwise = do + entryValue <- readWorkEntry arenaValue rowIndex diagonalIndex + if entryValue == 0 + then scanRows (rowIndex + 1) changed + else do + pivotValue <- readWorkEntry arenaValue diagonalIndex diagonalIndex + if pivotValue == 0 + then pure (Just SmithWitnessPivotBecameZero, changed) + else do + let (quotientValue, remainderValue) = balancedDivMod entryValue pivotValue + if quotientValue /= 0 + then rowCombineWitnessed rowIndex diagonalIndex quotientValue arenaValue + else pure () + reducedEntry <- readWorkEntry arenaValue rowIndex diagonalIndex + reductionFailure <- + if remainderValue == 0 && reducedEntry == 0 + then pure Nothing + else gcdCombineRowsWitnessed diagonalIndex rowIndex diagonalIndex arenaValue + case reductionFailure of + Just failureValue -> pure (Just failureValue, True) + Nothing -> scanRows 0 True + +reduceDiagonalRowMutable :: forall s. Int -> SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure, Bool) +reduceDiagonalRowMutable diagonalIndex arenaValue = + scanColumns 0 False + where + scanColumns :: Int -> Bool -> ST s (Maybe SmithWitnessFailure, Bool) + scanColumns columnIndex changed + | columnIndex >= smithWitnessColumnCount arenaValue = pure (Nothing, changed) + | columnIndex == diagonalIndex = scanColumns (columnIndex + 1) changed + | otherwise = do + entryValue <- readWorkEntry arenaValue diagonalIndex columnIndex + if entryValue == 0 + then scanColumns (columnIndex + 1) changed + else do + pivotValue <- readWorkEntry arenaValue diagonalIndex diagonalIndex + if pivotValue == 0 + then pure (Just SmithWitnessPivotBecameZero, changed) + else do + let (quotientValue, remainderValue) = balancedDivMod entryValue pivotValue + if quotientValue /= 0 + then columnCombineWitnessed columnIndex diagonalIndex quotientValue arenaValue + else pure () + reducedEntry <- readWorkEntry arenaValue diagonalIndex columnIndex + reductionFailure <- + if remainderValue == 0 && reducedEntry == 0 + then pure Nothing + else gcdCombineColumnsWitnessed diagonalIndex diagonalIndex columnIndex arenaValue + case reductionFailure of + Just failureValue -> pure (Just failureValue, True) + Nothing -> scanColumns 0 True + +balancedDivMod :: Integer -> Integer -> (Integer, Integer) +balancedDivMod numerator denominator = + let positiveDenominator = abs denominator + (floorQuotient, floorRemainder) = numerator `divMod` positiveDenominator + (quotientValue, remainderValue) = + if 2 * floorRemainder > positiveDenominator + then (floorQuotient + 1, floorRemainder - positiveDenominator) + else (floorQuotient, floorRemainder) + in if denominator < 0 + then (negate quotientValue, remainderValue) + else (quotientValue, remainderValue) + +columnClearedMutable :: forall s. Int -> Int -> SmithWitnessArena s -> ST s Bool +columnClearedMutable pivotRow pivotColumn arenaValue = + scanRows 0 + where + scanRows :: Int -> ST s Bool + scanRows rowIndex + | rowIndex >= smithWitnessRowCount arenaValue = pure True + | rowIndex == pivotRow = scanRows (rowIndex + 1) + | otherwise = do + isZeroEntry <- entryIsZeroMutable arenaValue rowIndex pivotColumn + if isZeroEntry + then scanRows (rowIndex + 1) + else pure False + +rowClearedMutable :: forall s. Int -> Int -> SmithWitnessArena s -> ST s Bool +rowClearedMutable pivotRow pivotColumn arenaValue = + scanColumns 0 + where + scanColumns :: Int -> ST s Bool + scanColumns columnIndex + | columnIndex >= smithWitnessColumnCount arenaValue = pure True + | columnIndex == pivotColumn = scanColumns (columnIndex + 1) + | otherwise = do + isZeroEntry <- entryIsZeroMutable arenaValue pivotRow columnIndex + if isZeroEntry + then scanColumns (columnIndex + 1) + else pure False + +gcdCombineRowsWitnessed :: Int -> Int -> Int -> SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure) +gcdCombineRowsWitnessed pivotRow candidateRow pivotColumn arenaValue = do + pivotValue <- readWorkEntry arenaValue pivotRow pivotColumn + entryValue <- readWorkEntry arenaValue candidateRow pivotColumn + let (gcdValue, pivotCoefficient, entryCoefficient) = extendedGcdDomain pivotValue entryValue + case (exactQuotientMutable "row gcd pivot quotient" pivotValue gcdValue, exactQuotientMutable "row gcd entry quotient" entryValue gcdValue) of + (SmithExactQuotient pivotQuotient, SmithExactQuotient entryQuotient) -> do + rowPairTransformWitnessed pivotRow candidateRow pivotCoefficient entryCoefficient (negate entryQuotient) pivotQuotient arenaValue + pure Nothing + (SmithInexactQuotient failureValue, _) -> pure (Just failureValue) + (_, SmithInexactQuotient failureValue) -> pure (Just failureValue) + +gcdCombineColumnsWitnessed :: Int -> Int -> Int -> SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure) +gcdCombineColumnsWitnessed pivotRow pivotColumn candidateColumn arenaValue = do + pivotValue <- readWorkEntry arenaValue pivotRow pivotColumn + entryValue <- readWorkEntry arenaValue pivotRow candidateColumn + let (gcdValue, pivotCoefficient, entryCoefficient) = extendedGcdDomain pivotValue entryValue + case (exactQuotientMutable "column gcd pivot quotient" pivotValue gcdValue, exactQuotientMutable "column gcd entry quotient" entryValue gcdValue) of + (SmithExactQuotient pivotQuotient, SmithExactQuotient entryQuotient) -> do + columnPairTransformWitnessed pivotColumn candidateColumn pivotCoefficient entryCoefficient (negate entryQuotient) pivotQuotient arenaValue + pure Nothing + (SmithInexactQuotient failureValue, _) -> pure (Just failureValue) + (_, SmithInexactQuotient failureValue) -> pure (Just failureValue) + +exactQuotientMutable :: String -> Integer -> Integer -> SmithExactQuotient +exactQuotientMutable context numerator denominator + | denominator == 0 = SmithInexactQuotient (SmithWitnessInexactDivision context) + | remainderValue == 0 = SmithExactQuotient quotientValue + | otherwise = SmithInexactQuotient (SmithWitnessInexactDivision context) + where + (quotientValue, remainderValue) = numerator `divMod` denominator + +enforceDivisibilityChainMutable :: forall s. SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure) +enforceDivisibilityChainMutable arenaValue = + repairAt (diagonalSize * diagonalSize) + where + diagonalSize :: Int + diagonalSize = min (smithWitnessRowCount arenaValue) (smithWitnessColumnCount arenaValue) + + repairAt :: Int -> ST s (Maybe SmithWitnessFailure) + repairAt remainingBudget + | remainingBudget <= 0 = do + violationValue <- findDivisibilityViolationMutable diagonalSize arenaValue + case violationValue of + Nothing -> pure Nothing + Just _ -> pure (Just (SmithWitnessBudgetExhausted "divisibility chain")) + | otherwise = do + violationValue <- findDivisibilityViolationMutable diagonalSize arenaValue + case violationValue of + Nothing -> pure Nothing + Just violationIndex -> do + rowCombineWitnessed violationIndex (violationIndex + 1) (-1) arenaValue + leftFailure <- normalizePivotMutable violationIndex violationIndex (normalizationBudget arenaValue) arenaValue + case leftFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> do + rightFailure <- normalizePivotMutable (violationIndex + 1) (violationIndex + 1) (normalizationBudget arenaValue) arenaValue + case rightFailure of + Just failureValue -> pure (Just failureValue) + Nothing -> repairAt (remainingBudget - 1) + +findDivisibilityViolationMutable :: forall s. Int -> SmithWitnessArena s -> ST s (Maybe Int) +findDivisibilityViolationMutable diagonalSize arenaValue = + scanDiagonal 0 + where + scanDiagonal :: Int -> ST s (Maybe Int) + scanDiagonal diagonalIndex + | diagonalIndex >= diagonalSize - 1 = pure Nothing + | otherwise = do + leftDiagonal <- readWorkEntry arenaValue diagonalIndex diagonalIndex + rightDiagonal <- readWorkEntry arenaValue (diagonalIndex + 1) (diagonalIndex + 1) + if leftDiagonal == 0 + || rightDiagonal == 0 + || rightDiagonal `mod` leftDiagonal == 0 + then scanDiagonal (diagonalIndex + 1) + else pure (Just diagonalIndex) + +normalizeDiagonalUnitsMutable :: forall s. SmithWitnessArena s -> ST s (Maybe SmithWitnessFailure) +normalizeDiagonalUnitsMutable arenaValue = + normalizeAt 0 + where + diagonalSize :: Int + diagonalSize = min (smithWitnessRowCount arenaValue) (smithWitnessColumnCount arenaValue) + + normalizeAt :: Int -> ST s (Maybe SmithWitnessFailure) + normalizeAt diagonalIndex + | diagonalIndex >= diagonalSize = pure Nothing + | otherwise = do + diagonalValue <- readWorkEntry arenaValue diagonalIndex diagonalIndex + if diagonalValue < 0 + then scaleRowWitnessed diagonalIndex (-1) arenaValue *> normalizeAt (diagonalIndex + 1) + else normalizeAt (diagonalIndex + 1) + +swapRowsWitnessed :: Int -> Int -> SmithWitnessArena s -> ST s () +swapRowsWitnessed leftRow rightRow arenaValue = do + swapRowsVector (smithWitnessColumnCount arenaValue) (smithWitnessWork arenaValue) leftRow rightRow 0 + swapRowsVector (smithWitnessRowCount arenaValue) (smithWitnessLeftRows arenaValue) leftRow rightRow 0 + swapColumnsVector (smithWitnessRowCount arenaValue) (smithWitnessLeftInverseRows arenaValue) leftRow rightRow 0 (smithWitnessRowCount arenaValue) + +swapColumnsWitnessed :: Int -> Int -> SmithWitnessArena s -> ST s () +swapColumnsWitnessed leftColumn rightColumn arenaValue = do + swapColumnsVector (smithWitnessColumnCount arenaValue) (smithWitnessWork arenaValue) leftColumn rightColumn 0 (smithWitnessRowCount arenaValue) + swapColumnsVector (smithWitnessColumnCount arenaValue) (smithWitnessRightRows arenaValue) leftColumn rightColumn 0 (smithWitnessColumnCount arenaValue) + swapRowsVector (smithWitnessColumnCount arenaValue) (smithWitnessRightInverseRows arenaValue) leftColumn rightColumn 0 + +rowCombineWitnessed :: Int -> Int -> Integer -> SmithWitnessArena s -> ST s () +rowCombineWitnessed targetRow sourceRow coefficient arenaValue = do + rowCombineVector (smithWitnessColumnCount arenaValue) (smithWitnessWork arenaValue) targetRow sourceRow coefficient 0 + rowCombineVector (smithWitnessRowCount arenaValue) (smithWitnessLeftRows arenaValue) targetRow sourceRow coefficient 0 + columnAddScaledVector (smithWitnessRowCount arenaValue) (smithWitnessLeftInverseRows arenaValue) sourceRow targetRow coefficient 0 (smithWitnessRowCount arenaValue) + +columnCombineWitnessed :: Int -> Int -> Integer -> SmithWitnessArena s -> ST s () +columnCombineWitnessed targetColumn sourceColumn coefficient arenaValue = do + columnCombineVector (smithWitnessColumnCount arenaValue) (smithWitnessWork arenaValue) targetColumn sourceColumn coefficient 0 (smithWitnessRowCount arenaValue) + columnCombineVector (smithWitnessColumnCount arenaValue) (smithWitnessRightRows arenaValue) targetColumn sourceColumn coefficient 0 (smithWitnessColumnCount arenaValue) + rowAddScaledVector (smithWitnessColumnCount arenaValue) (smithWitnessRightInverseRows arenaValue) sourceColumn targetColumn coefficient 0 + +rowPairTransformWitnessed :: Int -> Int -> Integer -> Integer -> Integer -> Integer -> SmithWitnessArena s -> ST s () +rowPairTransformWitnessed leftRow rightRow aa ab ba bb arenaValue = do + rowPairTransformVector (smithWitnessColumnCount arenaValue) (smithWitnessWork arenaValue) leftRow rightRow aa ab ba bb 0 + rowPairTransformVector (smithWitnessRowCount arenaValue) (smithWitnessLeftRows arenaValue) leftRow rightRow aa ab ba bb 0 + columnPairTransformVector (smithWitnessRowCount arenaValue) (smithWitnessLeftInverseRows arenaValue) leftRow rightRow bb (negate ba) (negate ab) aa 0 (smithWitnessRowCount arenaValue) + +columnPairTransformWitnessed :: Int -> Int -> Integer -> Integer -> Integer -> Integer -> SmithWitnessArena s -> ST s () +columnPairTransformWitnessed leftColumn rightColumn aa ab ba bb arenaValue = do + columnPairTransformVector (smithWitnessColumnCount arenaValue) (smithWitnessWork arenaValue) leftColumn rightColumn aa ab ba bb 0 (smithWitnessRowCount arenaValue) + columnPairTransformVector (smithWitnessColumnCount arenaValue) (smithWitnessRightRows arenaValue) leftColumn rightColumn aa ab ba bb 0 (smithWitnessColumnCount arenaValue) + rowPairTransformVector (smithWitnessColumnCount arenaValue) (smithWitnessRightInverseRows arenaValue) leftColumn rightColumn bb (negate ba) (negate ab) aa 0 + +scaleRowWitnessed :: Int -> Integer -> SmithWitnessArena s -> ST s () +scaleRowWitnessed rowIndex factor arenaValue = do + scaleRowVector (smithWitnessColumnCount arenaValue) (smithWitnessWork arenaValue) rowIndex factor 0 + scaleRowVector (smithWitnessRowCount arenaValue) (smithWitnessLeftRows arenaValue) rowIndex factor 0 + scaleColumnVector (smithWitnessRowCount arenaValue) (smithWitnessLeftInverseRows arenaValue) rowIndex factor 0 (smithWitnessRowCount arenaValue) + +swapRowsVector :: Int -> MV.MVector s Integer -> Int -> Int -> Int -> ST s () +swapRowsVector columnCount entries leftRow rightRow columnIndex + | leftRow == rightRow = pure () + | columnIndex >= columnCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount leftRow columnIndex + rightIndex = flatIndex columnCount rightRow columnIndex + leftValue <- MV.unsafeRead entries leftIndex + rightValue <- MV.unsafeRead entries rightIndex + MV.unsafeWrite entries leftIndex rightValue + MV.unsafeWrite entries rightIndex leftValue + swapRowsVector columnCount entries leftRow rightRow (columnIndex + 1) + +swapColumnsVector :: Int -> MV.MVector s Integer -> Int -> Int -> Int -> Int -> ST s () +swapColumnsVector columnCount entries leftColumn rightColumn rowIndex rowCount + | leftColumn == rightColumn = pure () + | rowIndex >= rowCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount rowIndex leftColumn + rightIndex = flatIndex columnCount rowIndex rightColumn + leftValue <- MV.unsafeRead entries leftIndex + rightValue <- MV.unsafeRead entries rightIndex + MV.unsafeWrite entries leftIndex rightValue + MV.unsafeWrite entries rightIndex leftValue + swapColumnsVector columnCount entries leftColumn rightColumn (rowIndex + 1) rowCount + +rowCombineVector :: Int -> MV.MVector s Integer -> Int -> Int -> Integer -> Int -> ST s () +rowCombineVector columnCount entries targetRow sourceRow coefficient columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + let targetIndex = flatIndex columnCount targetRow columnIndex + sourceIndex = flatIndex columnCount sourceRow columnIndex + targetValue <- MV.read entries targetIndex + sourceValue <- MV.read entries sourceIndex + MV.write entries targetIndex (targetValue - coefficient * sourceValue) + rowCombineVector columnCount entries targetRow sourceRow coefficient (columnIndex + 1) + +columnCombineVector :: Int -> MV.MVector s Integer -> Int -> Int -> Integer -> Int -> Int -> ST s () +columnCombineVector columnCount entries targetColumn sourceColumn coefficient rowIndex rowCount + | rowIndex >= rowCount = pure () + | otherwise = do + let targetIndex = flatIndex columnCount rowIndex targetColumn + sourceIndex = flatIndex columnCount rowIndex sourceColumn + targetValue <- MV.read entries targetIndex + sourceValue <- MV.read entries sourceIndex + MV.write entries targetIndex (targetValue - coefficient * sourceValue) + columnCombineVector columnCount entries targetColumn sourceColumn coefficient (rowIndex + 1) rowCount + +rowAddScaledVector :: Int -> MV.MVector s Integer -> Int -> Int -> Integer -> Int -> ST s () +rowAddScaledVector columnCount entries targetRow sourceRow coefficient columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + let targetIndex = flatIndex columnCount targetRow columnIndex + sourceIndex = flatIndex columnCount sourceRow columnIndex + targetValue <- MV.read entries targetIndex + sourceValue <- MV.read entries sourceIndex + MV.write entries targetIndex (targetValue + coefficient * sourceValue) + rowAddScaledVector columnCount entries targetRow sourceRow coefficient (columnIndex + 1) + +columnAddScaledVector :: Int -> MV.MVector s Integer -> Int -> Int -> Integer -> Int -> Int -> ST s () +columnAddScaledVector columnCount entries targetColumn sourceColumn coefficient rowIndex rowCount + | rowIndex >= rowCount = pure () + | otherwise = do + let targetIndex = flatIndex columnCount rowIndex targetColumn + sourceIndex = flatIndex columnCount rowIndex sourceColumn + targetValue <- MV.read entries targetIndex + sourceValue <- MV.read entries sourceIndex + MV.write entries targetIndex (targetValue + coefficient * sourceValue) + columnAddScaledVector columnCount entries targetColumn sourceColumn coefficient (rowIndex + 1) rowCount + +rowPairTransformVector :: Int -> MV.MVector s Integer -> Int -> Int -> Integer -> Integer -> Integer -> Integer -> Int -> ST s () +rowPairTransformVector columnCount entries leftRow rightRow aa ab ba bb columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount leftRow columnIndex + rightIndex = flatIndex columnCount rightRow columnIndex + leftValue <- MV.read entries leftIndex + rightValue <- MV.read entries rightIndex + MV.write entries leftIndex (aa * leftValue + ab * rightValue) + MV.write entries rightIndex (ba * leftValue + bb * rightValue) + rowPairTransformVector columnCount entries leftRow rightRow aa ab ba bb (columnIndex + 1) + +columnPairTransformVector :: Int -> MV.MVector s Integer -> Int -> Int -> Integer -> Integer -> Integer -> Integer -> Int -> Int -> ST s () +columnPairTransformVector columnCount entries leftColumn rightColumn aa ab ba bb rowIndex rowCount + | rowIndex >= rowCount = pure () + | otherwise = do + let leftIndex = flatIndex columnCount rowIndex leftColumn + rightIndex = flatIndex columnCount rowIndex rightColumn + leftValue <- MV.read entries leftIndex + rightValue <- MV.read entries rightIndex + MV.write entries leftIndex (aa * leftValue + ab * rightValue) + MV.write entries rightIndex (ba * leftValue + bb * rightValue) + columnPairTransformVector columnCount entries leftColumn rightColumn aa ab ba bb (rowIndex + 1) rowCount + +scaleRowVector :: Int -> MV.MVector s Integer -> Int -> Integer -> Int -> ST s () +scaleRowVector columnCount entries rowIndex factor columnIndex + | columnIndex >= columnCount = pure () + | otherwise = do + let entryIndex = flatIndex columnCount rowIndex columnIndex + entryValue <- MV.read entries entryIndex + MV.write entries entryIndex (factor * entryValue) + scaleRowVector columnCount entries rowIndex factor (columnIndex + 1) + +scaleColumnVector :: Int -> MV.MVector s Integer -> Int -> Integer -> Int -> Int -> ST s () +scaleColumnVector columnCount entries columnIndex factor rowIndex rowCount + | rowIndex >= rowCount = pure () + | otherwise = do + let entryIndex = flatIndex columnCount rowIndex columnIndex + entryValue <- MV.read entries entryIndex + MV.write entries entryIndex (factor * entryValue) + scaleColumnVector columnCount entries columnIndex factor (rowIndex + 1) rowCount diff --git a/src-eigen/Moonlight/LinAlg/Internal/Eigen/DenseWork.hs b/src-eigen/Moonlight/LinAlg/Internal/Eigen/DenseWork.hs new file mode 100644 --- /dev/null +++ b/src-eigen/Moonlight/LinAlg/Internal/Eigen/DenseWork.hs @@ -0,0 +1,95 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Internal.Eigen.DenseWork + ( MutableDenseWork (..), + denseWorkIndex, + dotDenseColumns, + newDenseWork, + readDenseWork, + scaleDenseColumn, + setIdentityDenseWork, + swapDenseColumns, + writeDenseWork, + ) +where + +import Control.Monad.ST (ST) +import Data.Primitive.PrimArray + ( MutablePrimArray, + newPrimArray, + readPrimArray, + setPrimArray, + writePrimArray, + ) +import Moonlight.LinAlg.Internal.Eigen.Kernels (forIndex) +import Prelude + +data MutableDenseWork s = MutableDenseWork + { denseWorkRows :: !Int, + denseWorkColumns :: !Int, + denseWorkPayload :: !(MutablePrimArray s Double) + } + +denseWorkIndex :: Int -> Int -> Int -> Int +denseWorkIndex !rowCount !rowIndex !columnIndex = rowIndex + (columnIndex * rowCount) +{-# INLINE denseWorkIndex #-} + +newDenseWork :: Int -> Int -> ST s (MutableDenseWork s) +newDenseWork !rowCount !columnCount = do + payload <- newPrimArray (rowCount * columnCount) + setPrimArray payload 0 (rowCount * columnCount) 0.0 + pure + MutableDenseWork + { denseWorkRows = rowCount, + denseWorkColumns = columnCount, + denseWorkPayload = payload + } +{-# INLINE newDenseWork #-} + +readDenseWork :: MutableDenseWork s -> Int -> Int -> ST s Double +readDenseWork (MutableDenseWork rowCount _ payload) !rowIndex !columnIndex = + readPrimArray payload (denseWorkIndex rowCount rowIndex columnIndex) +{-# INLINE readDenseWork #-} + +writeDenseWork :: MutableDenseWork s -> Int -> Int -> Double -> ST s () +writeDenseWork (MutableDenseWork rowCount _ payload) !rowIndex !columnIndex !entryValue = + writePrimArray payload (denseWorkIndex rowCount rowIndex columnIndex) entryValue +{-# INLINE writeDenseWork #-} + +setIdentityDenseWork :: MutableDenseWork s -> ST s () +setIdentityDenseWork work@(MutableDenseWork rowCount columnCount payload) = do + setPrimArray payload 0 (rowCount * columnCount) 0.0 + forIndex 0 (min rowCount columnCount) $ \indexValue -> + writeDenseWork work indexValue indexValue 1.0 +{-# INLINE setIdentityDenseWork #-} + +swapDenseColumns :: MutableDenseWork s -> Int -> Int -> ST s () +swapDenseColumns work@(MutableDenseWork rowCount _ _) !leftColumn !rightColumn = + if leftColumn == rightColumn + then pure () + else + forIndex 0 rowCount $ \rowIndex -> do + leftValue <- readDenseWork work rowIndex leftColumn + rightValue <- readDenseWork work rowIndex rightColumn + writeDenseWork work rowIndex leftColumn rightValue + writeDenseWork work rowIndex rightColumn leftValue +{-# INLINE swapDenseColumns #-} + +dotDenseColumns :: MutableDenseWork s -> Int -> Int -> ST s Double +dotDenseColumns work@(MutableDenseWork rowCount _ _) !leftColumn !rightColumn = go 0 0.0 + where + go !rowIndex !accumulator + | rowIndex >= rowCount = pure accumulator + | otherwise = do + leftValue <- readDenseWork work rowIndex leftColumn + rightValue <- readDenseWork work rowIndex rightColumn + go (rowIndex + 1) (accumulator + leftValue * rightValue) +{-# INLINE dotDenseColumns #-} + +scaleDenseColumn :: MutableDenseWork s -> Int -> Double -> ST s () +scaleDenseColumn work@(MutableDenseWork rowCount _ _) !columnIndex !scaleValue = + forIndex 0 rowCount $ \rowIndex -> do + entryValue <- readDenseWork work rowIndex columnIndex + writeDenseWork work rowIndex columnIndex (scaleValue * entryValue) +{-# INLINE scaleDenseColumn #-} diff --git a/src-eigen/Moonlight/LinAlg/Internal/Eigen/Householder.hs b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Householder.hs new file mode 100644 --- /dev/null +++ b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Householder.hs @@ -0,0 +1,217 @@ +{-# LANGUAGE BangPatterns #-} + +module Moonlight.LinAlg.Internal.Eigen.Householder + ( backtransformLower, + tridiagonalizeLower, + ) +where + +import Control.Monad (when) +import Control.Monad.ST (ST) +import Data.Primitive.PrimArray + ( MutablePrimArray, + newPrimArray, + readPrimArray, + setPrimArray, + writePrimArray, + ) +import Moonlight.LinAlg.Internal.Eigen.DenseWork + ( MutableDenseWork (..), + readDenseWork, + writeDenseWork, + ) +import Moonlight.LinAlg.Internal.Eigen.Kernels + ( copySignMagnitude, + forDescendingIndex, + forIndex, + hypotStable, + ) +import Prelude + +tridiagonalizeLower :: + MutableDenseWork s -> + ST s (MutablePrimArray s Double, MutablePrimArray s Double, MutablePrimArray s Double) +tridiagonalizeLower work@(MutableDenseWork matrixSize _ _) = do + diagonalValues <- newPrimArray matrixSize + offDiagonalValues <- newPrimArray matrixSize + reflectorScalars <- newPrimArray matrixSize + matrixTimesReflector <- newPrimArray matrixSize + setPrimArray offDiagonalValues 0 matrixSize 0.0 + setPrimArray reflectorScalars 0 matrixSize 0.0 + forIndex 0 matrixSize $ \pivotIndex -> do + diagonalEntry <- readDenseWork work pivotIndex pivotIndex + writePrimArray diagonalValues pivotIndex diagonalEntry + when (pivotIndex < matrixSize - 1) $ do + (reflectorScalar, reflectedSubDiagonal) <- makeHouseholderColumn pivotIndex work + writePrimArray reflectorScalars pivotIndex reflectorScalar + writePrimArray offDiagonalValues pivotIndex reflectedSubDiagonal + when (reflectorScalar /= 0.0) $ do + writeDenseWork work (pivotIndex + 1) pivotIndex 1.0 + symmetricLowerMatrixVector pivotIndex work matrixTimesReflector + scaleScratch (matrixSize - pivotIndex - 1) reflectorScalar matrixTimesReflector + reflectorDot <- dotImplicitReflector pivotIndex work matrixTimesReflector + let !rankTwoCorrection = (-0.5) * reflectorScalar * reflectorDot + addScaledImplicitReflector pivotIndex work matrixTimesReflector rankTwoCorrection + rankTwoUpdateLower pivotIndex work matrixTimesReflector + writeDenseWork work (pivotIndex + 1) pivotIndex reflectedSubDiagonal + pure (diagonalValues, offDiagonalValues, reflectorScalars) + +makeHouseholderColumn :: Int -> MutableDenseWork s -> ST s (Double, Double) +makeHouseholderColumn !pivotIndex work@(MutableDenseWork matrixSize _ _) = + let !firstRow = pivotIndex + 1 + !reflectorLength = matrixSize - firstRow + in if reflectorLength <= 0 + then pure (0.0, 0.0) + else do + firstEntry <- readDenseWork work firstRow pivotIndex + tailNorm <- columnTailNorm work pivotIndex (firstRow + 1) matrixSize + if tailNorm == 0.0 + then pure (0.0, firstEntry) + else do + let !sourceNorm = hypotStable firstEntry tailNorm + !reflectedHead = negate (copySignMagnitude sourceNorm firstEntry) + !reflectorScalar = (reflectedHead - firstEntry) / reflectedHead + !tailScale = 1.0 / (firstEntry - reflectedHead) + scaleColumnTail work pivotIndex (firstRow + 1) matrixSize tailScale + pure (reflectorScalar, reflectedHead) +{-# INLINE makeHouseholderColumn #-} + +columnTailNorm :: MutableDenseWork s -> Int -> Int -> Int -> ST s Double +columnTailNorm work !columnIndex !startRow !stopRow = go startRow 0.0 1.0 + where + go !rowIndex !scaleValue !scaledSum + | rowIndex >= stopRow = + if scaleValue == 0.0 + then pure 0.0 + else pure (scaleValue * sqrt scaledSum) + | otherwise = do + rawEntry <- readDenseWork work rowIndex columnIndex + let !entryAbs = abs rawEntry + if entryAbs == 0.0 + then go (rowIndex + 1) scaleValue scaledSum + else + if scaleValue < entryAbs + then + let !scaledRatio = scaleValue / entryAbs + in go (rowIndex + 1) entryAbs (1.0 + scaledSum * scaledRatio * scaledRatio) + else + let !scaledRatio = entryAbs / scaleValue + in go (rowIndex + 1) scaleValue (scaledSum + scaledRatio * scaledRatio) +{-# INLINE columnTailNorm #-} + +scaleColumnTail :: MutableDenseWork s -> Int -> Int -> Int -> Double -> ST s () +scaleColumnTail work !columnIndex !startRow !stopRow !scaleValue = + forIndex startRow stopRow $ \rowIndex -> do + entryValue <- readDenseWork work rowIndex columnIndex + writeDenseWork work rowIndex columnIndex (scaleValue * entryValue) +{-# INLINE scaleColumnTail #-} + +implicitReflectorEntry :: Int -> MutableDenseWork s -> Int -> ST s Double +implicitReflectorEntry !pivotIndex work !localIndex = + if localIndex == 0 + then pure 1.0 + else readDenseWork work (pivotIndex + 1 + localIndex) pivotIndex +{-# INLINE implicitReflectorEntry #-} + +symmetricLowerMatrixVector :: + Int -> + MutableDenseWork s -> + MutablePrimArray s Double -> + ST s () +symmetricLowerMatrixVector !pivotIndex work@(MutableDenseWork matrixSize _ _) scratchValues = do + let !startRow = pivotIndex + 1 + !dimension = matrixSize - startRow + setPrimArray scratchValues 0 dimension 0.0 + forIndex 0 dimension $ \columnLocalIndex -> do + reflectorColumnEntry <- implicitReflectorEntry pivotIndex work columnLocalIndex + diagonalEntry <- readDenseWork work (startRow + columnLocalIndex) (startRow + columnLocalIndex) + scratchColumnEntry <- readPrimArray scratchValues columnLocalIndex + writePrimArray scratchValues columnLocalIndex (scratchColumnEntry + diagonalEntry * reflectorColumnEntry) + forIndex (columnLocalIndex + 1) dimension $ \rowLocalIndex -> do + lowerEntry <- readDenseWork work (startRow + rowLocalIndex) (startRow + columnLocalIndex) + reflectorRowEntry <- implicitReflectorEntry pivotIndex work rowLocalIndex + rowAccumulator <- readPrimArray scratchValues rowLocalIndex + writePrimArray scratchValues rowLocalIndex (rowAccumulator + lowerEntry * reflectorColumnEntry) + columnAccumulator <- readPrimArray scratchValues columnLocalIndex + writePrimArray scratchValues columnLocalIndex (columnAccumulator + lowerEntry * reflectorRowEntry) +{-# INLINE symmetricLowerMatrixVector #-} + +scaleScratch :: Int -> Double -> MutablePrimArray s Double -> ST s () +scaleScratch !entryCount !scaleValue scratchValues = + forIndex 0 entryCount $ \entryIndex -> do + entryValue <- readPrimArray scratchValues entryIndex + writePrimArray scratchValues entryIndex (scaleValue * entryValue) +{-# INLINE scaleScratch #-} + +dotImplicitReflector :: Int -> MutableDenseWork s -> MutablePrimArray s Double -> ST s Double +dotImplicitReflector !pivotIndex work@(MutableDenseWork matrixSize _ _) scratchValues = + let !dimension = matrixSize - pivotIndex - 1 + in go 0 0.0 dimension + where + go !localIndex !accumulator !dimension + | localIndex >= dimension = pure accumulator + | otherwise = do + reflectorEntry <- implicitReflectorEntry pivotIndex work localIndex + scratchEntry <- readPrimArray scratchValues localIndex + go (localIndex + 1) (accumulator + reflectorEntry * scratchEntry) dimension +{-# INLINE dotImplicitReflector #-} + +addScaledImplicitReflector :: + Int -> + MutableDenseWork s -> + MutablePrimArray s Double -> + Double -> + ST s () +addScaledImplicitReflector !pivotIndex work@(MutableDenseWork matrixSize _ _) scratchValues !scaleValue = + let !dimension = matrixSize - pivotIndex - 1 + in forIndex 0 dimension $ \localIndex -> do + reflectorEntry <- implicitReflectorEntry pivotIndex work localIndex + scratchEntry <- readPrimArray scratchValues localIndex + writePrimArray scratchValues localIndex (scratchEntry + scaleValue * reflectorEntry) +{-# INLINE addScaledImplicitReflector #-} + +rankTwoUpdateLower :: Int -> MutableDenseWork s -> MutablePrimArray s Double -> ST s () +rankTwoUpdateLower !pivotIndex work@(MutableDenseWork matrixSize _ _) updateVector = do + let !startRow = pivotIndex + 1 + !dimension = matrixSize - startRow + forIndex 0 dimension $ \columnLocalIndex -> do + reflectorColumnEntry <- implicitReflectorEntry pivotIndex work columnLocalIndex + updateColumnEntry <- readPrimArray updateVector columnLocalIndex + forIndex columnLocalIndex dimension $ \rowLocalIndex -> do + reflectorRowEntry <- implicitReflectorEntry pivotIndex work rowLocalIndex + updateRowEntry <- readPrimArray updateVector rowLocalIndex + matrixEntry <- readDenseWork work (startRow + rowLocalIndex) (startRow + columnLocalIndex) + writeDenseWork + work + (startRow + rowLocalIndex) + (startRow + columnLocalIndex) + (matrixEntry - reflectorRowEntry * updateColumnEntry - updateRowEntry * reflectorColumnEntry) +{-# INLINE rankTwoUpdateLower #-} + +backtransformLower :: + MutableDenseWork s -> + MutablePrimArray s Double -> + MutableDenseWork s -> + ST s () +backtransformLower reflectors reflectorScalars eigenvectors@(MutableDenseWork matrixSize _ _) = + forDescendingIndex (matrixSize - 2) 0 $ \pivotIndex -> do + reflectorScalar <- readPrimArray reflectorScalars pivotIndex + when (reflectorScalar /= 0.0) $ do + let !startRow = pivotIndex + 1 + !dimension = matrixSize - startRow + forIndex 0 matrixSize $ \columnIndex -> do + firstComponent <- readDenseWork eigenvectors startRow columnIndex + reflectorProduct <- dotTail 1 dimension firstComponent pivotIndex columnIndex startRow + let !projectionScale = reflectorScalar * reflectorProduct + writeDenseWork eigenvectors startRow columnIndex (firstComponent - projectionScale) + forIndex 1 dimension $ \localIndex -> do + reflectorEntry <- readDenseWork reflectors (startRow + localIndex) pivotIndex + eigenvectorEntry <- readDenseWork eigenvectors (startRow + localIndex) columnIndex + writeDenseWork eigenvectors (startRow + localIndex) columnIndex (eigenvectorEntry - projectionScale * reflectorEntry) + where + dotTail !localIndex !dimension !accumulator !pivotIndex !columnIndex !startRow + | localIndex >= dimension = pure accumulator + | otherwise = do + reflectorEntry <- readDenseWork reflectors (startRow + localIndex) pivotIndex + eigenvectorEntry <- readDenseWork eigenvectors (startRow + localIndex) columnIndex + dotTail (localIndex + 1) dimension (accumulator + reflectorEntry * eigenvectorEntry) pivotIndex columnIndex startRow diff --git a/src-eigen/Moonlight/LinAlg/Internal/Eigen/Input.hs b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Input.hs new file mode 100644 --- /dev/null +++ b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Input.hs @@ -0,0 +1,46 @@ +module Moonlight.LinAlg.Internal.Eigen.Input + ( validateSymmetricEigenInput, + isSymmetricWithin, + ) +where + +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Internal.Eigen.Kernels (finiteDouble) +import Moonlight.LinAlg.Internal.Primitives (epsilon) +import Prelude + +validateSymmetricEigenInput :: String -> Int -> [[Double]] -> Either MoonlightError () +validateSymmetricEigenInput context matrixSize matrixRows + | matrixSize < 0 = + Left (InvariantViolation (context <> " requires a non-negative matrix size")) + | length matrixRows /= matrixSize = + Left (InvariantViolation (context <> " row count does not match declared matrix size")) + | not (all ((== matrixSize) . length) matrixRows) = + Left (InvariantViolation (context <> " requires a square dense matrix")) + | not (all (all finiteDouble) matrixRows) = + Left (InvariantViolation (context <> " requires finite matrix entries")) + | not (isSymmetricWithin (sqrt epsilon) matrixRows) = + Left (InvariantViolation (context <> " requires a symmetric matrix")) + | otherwise = Right () + +isSymmetricWithin :: Double -> [[Double]] -> Bool +isSymmetricWithin tolerance matrixRows = + and + [ abs (leftEntry - rightEntry) <= tolerance + | (rowIndex, rowValues) <- zip [0 :: Int ..] matrixRows, + (columnIndex, leftEntry) <- zip [0 :: Int ..] rowValues, + rowIndex < columnIndex, + Just rightEntry <- [matrixEntryMaybe columnIndex rowIndex matrixRows] + ] + +matrixEntryMaybe :: Int -> Int -> [[a]] -> Maybe a +matrixEntryMaybe rowIndex columnIndex matrixRows = + listEntryMaybe rowIndex matrixRows >>= listEntryMaybe columnIndex + +listEntryMaybe :: Int -> [a] -> Maybe a +listEntryMaybe indexValue values + | indexValue < 0 = Nothing + | otherwise = + case drop indexValue values of + [] -> Nothing + entryValue : _ -> Just entryValue diff --git a/src-eigen/Moonlight/LinAlg/Internal/Eigen/Kernels.hs b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Kernels.hs new file mode 100644 --- /dev/null +++ b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Kernels.hs @@ -0,0 +1,77 @@ +{-# LANGUAGE BangPatterns #-} + +module Moonlight.LinAlg.Internal.Eigen.Kernels + ( copySignMagnitude, + epsDouble, + finiteDouble, + forDescendingIndex, + forIndex, + hypotStable, + maxFiniteDouble, + safeMinimumDouble, + ) +where + +import Control.Monad.ST (ST) +import Moonlight.Core (fieldValueValid) +import Prelude + +-- | IEEE-754 binary64 machine epsilon. The eigensolver's convergence tests need +-- the floating-point unit scale, not Moonlight's broader domain tolerance. +epsDouble :: Double +epsDouble = encodeFloat 1 (-52) +{-# NOINLINE epsDouble #-} + +safeMinimumDouble :: Double +safeMinimumDouble = encodeFloat 1 (-1022) +{-# NOINLINE safeMinimumDouble #-} + +maxFiniteDouble :: Double +maxFiniteDouble = encodeFloat 0x1fffffffffffff (971) +{-# NOINLINE maxFiniteDouble #-} + +finiteDouble :: Double -> Bool +finiteDouble = fieldValueValid +{-# INLINE finiteDouble #-} + +copySignMagnitude :: Double -> Double -> Double +copySignMagnitude !magnitudeValue !signReference = + if signReference < 0.0 || isNegativeZero signReference + then negate (abs magnitudeValue) + else abs magnitudeValue +{-# INLINE copySignMagnitude #-} + +hypotStable :: Double -> Double -> Double +hypotStable !leftValue !rightValue = + let !leftAbs = abs leftValue + !rightAbs = abs rightValue + in if leftAbs < rightAbs + then + if rightAbs == 0.0 + then 0.0 + else + let !scaled = leftAbs / rightAbs + in rightAbs * sqrt (1.0 + scaled * scaled) + else + if leftAbs == 0.0 + then 0.0 + else + let !scaled = rightAbs / leftAbs + in leftAbs * sqrt (1.0 + scaled * scaled) +{-# INLINE hypotStable #-} + +forIndex :: Int -> Int -> (Int -> ST s ()) -> ST s () +forIndex !startIndex !stopIndex action = go startIndex + where + go !indexValue + | indexValue >= stopIndex = pure () + | otherwise = action indexValue >> go (indexValue + 1) +{-# INLINE forIndex #-} + +forDescendingIndex :: Int -> Int -> (Int -> ST s ()) -> ST s () +forDescendingIndex !startIndex !stopIndex action = go startIndex + where + go !indexValue + | indexValue < stopIndex = pure () + | otherwise = action indexValue >> go (indexValue - 1) +{-# INLINE forDescendingIndex #-} diff --git a/src-eigen/Moonlight/LinAlg/Internal/Eigen/Residual.hs b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Residual.hs new file mode 100644 --- /dev/null +++ b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Residual.hs @@ -0,0 +1,155 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Internal.Eigen.Residual + ( ResidualReport (..), + residualReportPassesSymmetricEigenLimits, + symmetricEigenResidual, + ) +where + +import Data.Vector.Storable qualified as S +import Moonlight.LinAlg.Internal.Eigen.Kernels (epsDouble) +import Moonlight.LinAlg.Pure.Dense.Flat + ( DenseDoubleMatrix, + denseDoubleMatrixShape, + denseDoubleMatrixToRowMajorVector, + ) +import Prelude + +data ResidualReport = ResidualReport + { residualMatrixNorm :: !Double, + residualFrobenius :: !Double, + residualOrthogonality :: !Double, + residualScaled :: !Double, + residualOrthogonalityScaled :: !Double + } + deriving stock (Eq, Show) + +residualReportPassesSymmetricEigenLimits :: ResidualReport -> Bool +residualReportPassesSymmetricEigenLimits report = + residualScaled report <= 1.0e7 + && residualOrthogonalityScaled report <= 1.0e7 + +symmetricEigenResidual :: DenseDoubleMatrix -> S.Vector Double -> DenseDoubleMatrix -> ResidualReport +symmetricEigenResidual matrixValue eigenvalues eigenvectors = + let !(matrixSize, _) = denseDoubleMatrixShape matrixValue + !matrixNorm = frobeniusNorm matrixValue + !residualNorm = residualFrobeniusNorm matrixValue eigenvalues eigenvectors + !orthogonalityNorm = orthogonalityFrobeniusNorm eigenvectors + !dimensionScale = max 1.0 (fromIntegral matrixSize) + !residualDenominator = max 1.0 matrixNorm * dimensionScale * epsDouble + !orthogonalityDenominator = dimensionScale * epsDouble + in ResidualReport + { residualMatrixNorm = matrixNorm, + residualFrobenius = residualNorm, + residualOrthogonality = orthogonalityNorm, + residualScaled = residualNorm / residualDenominator, + residualOrthogonalityScaled = orthogonalityNorm / orthogonalityDenominator + } + +frobeniusNorm :: DenseDoubleMatrix -> Double +frobeniusNorm matrixValue = + S.foldl' accumulateScaledNorm scaledNormZero (denseDoubleMatrixToRowMajorVector matrixValue) + |> scaledNormValue + +residualFrobeniusNorm :: DenseDoubleMatrix -> S.Vector Double -> DenseDoubleMatrix -> Double +residualFrobeniusNorm matrixValue eigenvalues eigenvectors = + residualColumn 0 scaledNormZero |> scaledNormValue + where + !(matrixSize, _) = denseDoubleMatrixShape matrixValue + matrixPayload = denseDoubleMatrixToRowMajorVector matrixValue + eigenvectorPayload = denseDoubleMatrixToRowMajorVector eigenvectors + + residualColumn !columnIndex !normState + | columnIndex >= matrixSize = normState + | otherwise = + residualRow columnIndex 0 normState + |> residualColumn (columnIndex + 1) + + residualRow !columnIndex !rowIndex !normState + | rowIndex >= matrixSize = normState + | otherwise = + let !lambdaValue = eigenvalues `S.unsafeIndex` columnIndex + !vectorEntry = eigenvectorAt rowIndex columnIndex + !residualEntry = matrixVectorEntry rowIndex columnIndex - lambdaValue * vectorEntry + in residualRow columnIndex (rowIndex + 1) (accumulateScaledNorm normState residualEntry) + + matrixVectorEntry !rowIndex !columnIndex = + dotAt 0 0.0 + where + !rowOffset = rowIndex * matrixSize + + dotAt !entryIndex !accumulator + | entryIndex >= matrixSize = accumulator + | otherwise = + let !matrixEntry = matrixPayload `S.unsafeIndex` (rowOffset + entryIndex) + !vectorEntry = eigenvectorAt entryIndex columnIndex + in dotAt (entryIndex + 1) (accumulator + matrixEntry * vectorEntry) + + eigenvectorAt !rowIndex !columnIndex = + eigenvectorPayload `S.unsafeIndex` (rowIndex * matrixSize + columnIndex) + +orthogonalityFrobeniusNorm :: DenseDoubleMatrix -> Double +orthogonalityFrobeniusNorm eigenvectors = + orthogonalityColumn 0 scaledNormZero |> scaledNormValue + where + !(matrixSize, _) = denseDoubleMatrixShape eigenvectors + eigenvectorPayload = denseDoubleMatrixToRowMajorVector eigenvectors + + orthogonalityColumn !leftColumn !normState + | leftColumn >= matrixSize = normState + | otherwise = + orthogonalityPair leftColumn leftColumn normState + |> orthogonalityColumn (leftColumn + 1) + + orthogonalityPair !leftColumn !rightColumn !normState + | rightColumn >= matrixSize = normState + | otherwise = + let !targetValue = if leftColumn == rightColumn then 1.0 else 0.0 + !weightValue = if leftColumn == rightColumn then 1.0 else sqrt 2.0 + !entryValue = weightValue * (columnDot leftColumn rightColumn - targetValue) + in orthogonalityPair leftColumn (rightColumn + 1) (accumulateScaledNorm normState entryValue) + + columnDot !leftColumn !rightColumn = go 0 0.0 + where + go !rowIndex !accumulator + | rowIndex >= matrixSize = accumulator + | otherwise = + let !leftValue = eigenvectorAt rowIndex leftColumn + !rightValue = eigenvectorAt rowIndex rightColumn + in go (rowIndex + 1) (accumulator + leftValue * rightValue) + + eigenvectorAt !rowIndex !columnIndex = + eigenvectorPayload `S.unsafeIndex` (rowIndex * matrixSize + columnIndex) + +data ScaledNorm = ScaledNorm !Double !Double + +scaledNormZero :: ScaledNorm +scaledNormZero = ScaledNorm 0.0 1.0 + +accumulateScaledNorm :: ScaledNorm -> Double -> ScaledNorm +accumulateScaledNorm (ScaledNorm !scaleValue !scaledSum) !entryValue = + let !entryAbs = abs entryValue + in if entryAbs == 0.0 + then ScaledNorm scaleValue scaledSum + else + if scaleValue < entryAbs + then + let !scaledRatio = scaleValue / entryAbs + in ScaledNorm entryAbs (1.0 + scaledSum * scaledRatio * scaledRatio) + else + let !scaledRatio = entryAbs / scaleValue + in ScaledNorm scaleValue (scaledSum + scaledRatio * scaledRatio) +{-# INLINE accumulateScaledNorm #-} + +scaledNormValue :: ScaledNorm -> Double +scaledNormValue (ScaledNorm !scaleValue !scaledSum) = + if scaleValue == 0.0 + then 0.0 + else scaleValue * sqrt scaledSum +{-# INLINE scaledNormValue #-} + +(|>) :: a -> (a -> b) -> b +(|>) value function = function value +{-# INLINE (|>) #-} diff --git a/src-eigen/Moonlight/LinAlg/Internal/Eigen/Symmetric.hs b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Symmetric.hs new file mode 100644 --- /dev/null +++ b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Symmetric.hs @@ -0,0 +1,185 @@ +{-# LANGUAGE BangPatterns #-} + +module Moonlight.LinAlg.Internal.Eigen.Symmetric + ( CertifiedSymmetricEigenResult (..), + SymmetricEigenCertificationFailure (..), + SymmetricEigenResult (..), + certifySymmetricEigenResult, + symmetricEigenPairsDenseUnchecked, + ) +where + +import Control.Monad.ST (ST, runST) +import Data.Bifunctor (first) +import Data.Primitive.PrimArray + ( MutablePrimArray, + readPrimArray, + writePrimArray, + ) +import Data.Vector.Storable qualified as S +import Data.Vector.Storable.Mutable qualified as SM +import Moonlight.Core + ( MoonlightError (..), + checkedNonNegativeProduct, + ) +import Moonlight.LinAlg.Internal.Eigen.DenseWork + ( MutableDenseWork (..), + newDenseWork, + ) +import Moonlight.LinAlg.Internal.Eigen.Householder + ( backtransformLower, + tridiagonalizeLower, + ) +import Moonlight.LinAlg.Internal.Eigen.Kernels (forIndex) +import Moonlight.LinAlg.Internal.Eigen.Residual + ( ResidualReport, + residualReportPassesSymmetricEigenLimits, + symmetricEigenResidual, + ) +import Moonlight.LinAlg.Internal.Eigen.Tridiagonal + ( canonicalizeEigenvectorSigns, + newIdentityEigenvectors, + orthonormalizeDegenerateClusters, + solveTridiagonalEigenvectors, + sortEigenpairsAscending, + ) +import Moonlight.LinAlg.Pure.Dense.Flat + ( DenseDoubleMatrix, + denseDoubleMatrixShape, + denseDoubleMatrixToRowMajorVector, + trustedDenseDoubleMatrixRowMajor, + ) +import Prelude + +data SymmetricEigenResult = SymmetricEigenResult + { symmetricEigenResultValues :: !(S.Vector Double), + symmetricEigenResultVectors :: !DenseDoubleMatrix + } + deriving stock (Eq, Show) + +data CertifiedSymmetricEigenResult = CertifiedSymmetricEigenResult + { certifiedSymmetricEigenResult :: !SymmetricEigenResult, + certifiedSymmetricEigenResidualReport :: !ResidualReport + } + deriving stock (Eq, Show) + +data SymmetricEigenCertificationFailure + = SymmetricEigenCertificationShapeMismatch !String + | SymmetricEigenCertificationResidualExceeded !ResidualReport + deriving stock (Eq, Show) + +symmetricEigenPairsDenseUnchecked :: Int -> DenseDoubleMatrix -> Either MoonlightError SymmetricEigenResult +symmetricEigenPairsDenseUnchecked !matrixSize matrixValue + | matrixSize /= rowCount || matrixSize /= columnCount = + Left + ( InvariantViolation + ( "symmetric eigen dense input shape mismatch: requested " + <> show matrixSize + <> " but received " + <> show (rowCount, columnCount) + ) + ) + | matrixSize < 0 = + Left (InvariantViolation "symmetric eigen dense input dimension must be non-negative") + | matrixSize == 0 = + Right + SymmetricEigenResult + { symmetricEigenResultValues = S.empty, + symmetricEigenResultVectors = trustedDenseDoubleMatrixRowMajor 0 0 S.empty + } + | matrixSize == 1 = + Right + SymmetricEigenResult + { symmetricEigenResultValues = S.singleton (denseDoubleMatrixToRowMajorVector matrixValue `S.unsafeIndex` 0), + symmetricEigenResultVectors = trustedDenseDoubleMatrixRowMajor 1 1 (S.singleton 1.0) + } + | otherwise = do + entryCount <- + first + (const (InvariantViolation "symmetric eigen workspace cardinality exceeds Int range")) + (checkedNonNegativeProduct matrixSize matrixSize) + runST (symmetricEigenPairsST matrixSize entryCount matrixValue) + where + !(rowCount, columnCount) = denseDoubleMatrixShape matrixValue + +certifySymmetricEigenResult :: + DenseDoubleMatrix -> + SymmetricEigenResult -> + Either SymmetricEigenCertificationFailure CertifiedSymmetricEigenResult +certifySymmetricEigenResult matrixValue resultValue@SymmetricEigenResult {symmetricEigenResultValues, symmetricEigenResultVectors} + | matrixRows /= matrixColumns = + Left (SymmetricEigenCertificationShapeMismatch ("symmetric eigen certification requires square source matrix, received " <> show (matrixRows, matrixColumns))) + | S.length symmetricEigenResultValues /= matrixRows = + Left (SymmetricEigenCertificationShapeMismatch ("symmetric eigen certification eigenvalue count mismatch: expected " <> show matrixRows <> " but received " <> show (S.length symmetricEigenResultValues))) + | vectorRows /= matrixRows || vectorColumns /= matrixRows = + Left (SymmetricEigenCertificationShapeMismatch ("symmetric eigen certification eigenvector shape mismatch: expected " <> show (matrixRows, matrixRows) <> " but received " <> show (vectorRows, vectorColumns))) + | residualReportPassesSymmetricEigenLimits reportValue = + Right + CertifiedSymmetricEigenResult + { certifiedSymmetricEigenResult = resultValue, + certifiedSymmetricEigenResidualReport = reportValue + } + | otherwise = + Left (SymmetricEigenCertificationResidualExceeded reportValue) + where + !(matrixRows, matrixColumns) = denseDoubleMatrixShape matrixValue + !(vectorRows, vectorColumns) = denseDoubleMatrixShape symmetricEigenResultVectors + !reportValue = symmetricEigenResidual matrixValue symmetricEigenResultValues symmetricEigenResultVectors + +symmetricEigenPairsST :: Int -> Int -> DenseDoubleMatrix -> ST s (Either MoonlightError SymmetricEigenResult) +symmetricEigenPairsST !matrixSize !entryCount matrixValue = do + work <- newDenseWork matrixSize matrixSize + copyLowerFlatToWork matrixSize matrixValue work + (diagonalValues, offDiagonalValues, reflectorScalars) <- tridiagonalizeLower work + eigenvectors <- newIdentityEigenvectors matrixSize + solveResult <- solveTridiagonalEigenvectors matrixSize diagonalValues offDiagonalValues eigenvectors + case solveResult of + Left err -> pure (Left err) + Right () -> do + sortEigenpairsAscending matrixSize diagonalValues eigenvectors + backtransformLower work reflectorScalars eigenvectors + clusterResult <- orthonormalizeDegenerateClusters matrixSize diagonalValues eigenvectors + case clusterResult of + Left err -> pure (Left err) + Right () -> do + canonicalizeEigenvectorSigns matrixSize eigenvectors + Right <$> eigenResultFromMutable matrixSize entryCount diagonalValues eigenvectors + +copyLowerFlatToWork :: Int -> DenseDoubleMatrix -> MutableDenseWork s -> ST s () +copyLowerFlatToWork !matrixSize matrixValue (MutableDenseWork rowCount _ workPayload) = + forIndex 0 matrixSize $ \rowIndex -> + forIndex 0 (rowIndex + 1) $ \columnIndex -> + writePrimArray + workPayload + (rowIndex + columnIndex * rowCount) + (payload `S.unsafeIndex` (rowIndex * matrixSize + columnIndex)) + where + payload = denseDoubleMatrixToRowMajorVector matrixValue +{-# INLINE copyLowerFlatToWork #-} + +eigenResultFromMutable :: + Int -> + Int -> + MutablePrimArray s Double -> + MutableDenseWork s -> + ST s SymmetricEigenResult +eigenResultFromMutable !matrixSize !entryCount diagonalValues (MutableDenseWork rowCount _ eigenvectorPayload) = do + eigenvalueBuffer <- SM.new matrixSize + eigenvectorBuffer <- SM.new entryCount + forIndex 0 matrixSize $ \columnIndex -> do + eigenvalue <- readPrimArray diagonalValues columnIndex + SM.unsafeWrite eigenvalueBuffer columnIndex eigenvalue + forIndex 0 matrixSize $ \rowIndex -> do + entryValue <- readPrimArray eigenvectorPayload (rowIndex + columnIndex * rowCount) + SM.unsafeWrite eigenvectorBuffer (rowIndex * matrixSize + columnIndex) entryValue + eigenvalues <- S.unsafeFreeze eigenvalueBuffer + eigenvectorValues <- S.unsafeFreeze eigenvectorBuffer + pure + SymmetricEigenResult + { symmetricEigenResultValues = eigenvalues, + symmetricEigenResultVectors = + trustedDenseDoubleMatrixRowMajor + matrixSize + matrixSize + eigenvectorValues + } diff --git a/src-eigen/Moonlight/LinAlg/Internal/Eigen/Tridiagonal.hs b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Tridiagonal.hs new file mode 100644 --- /dev/null +++ b/src-eigen/Moonlight/LinAlg/Internal/Eigen/Tridiagonal.hs @@ -0,0 +1,404 @@ +{-# LANGUAGE BangPatterns #-} + +module Moonlight.LinAlg.Internal.Eigen.Tridiagonal + ( canonicalizeEigenvectorSigns, + eigenpairsFromMutable, + newIdentityEigenvectors, + orthonormalizeDegenerateClusters, + sortEigenpairsAscending, + solveTridiagonalEigenvectors, + ) +where + +import Control.Monad (when) +import Control.Monad.ST (ST) +import Data.Primitive.PrimArray + ( MutablePrimArray, + readPrimArray, + writePrimArray, + ) +import Moonlight.Core (MoonlightError (..), fieldValueValid) +import Moonlight.LinAlg.Internal.Eigen.DenseWork + ( MutableDenseWork (..), + newDenseWork, + setIdentityDenseWork, + ) +import Moonlight.LinAlg.Internal.Eigen.Kernels + ( copySignMagnitude, + epsDouble, + forIndex, + hypotStable, + maxFiniteDouble, + safeMinimumDouble, + ) +import Prelude + +newIdentityEigenvectors :: Int -> ST s (MutableDenseWork s) +newIdentityEigenvectors !matrixSize = do + eigenvectors <- newDenseWork matrixSize matrixSize + setIdentityDenseWork eigenvectors + pure eigenvectors + +solveTridiagonalEigenvectors :: + Int -> + MutablePrimArray s Double -> + MutablePrimArray s Double -> + MutableDenseWork s -> + ST s (Either MoonlightError ()) +solveTridiagonalEigenvectors !matrixSize diagonalValues offDiagonalValues eigenvectors = do + when (matrixSize > 0) $ + writePrimArray offDiagonalValues (matrixSize - 1) 0.0 + scaleValue <- scaleTridiagonal matrixSize diagonalValues offDiagonalValues + solveResult <- solveAllIndices 0 + unscaleDiagonal matrixSize diagonalValues scaleValue + pure solveResult + where + !iterationLimit = max 64 (matrixSize * 128) + + solveAllIndices !splitIndex + | splitIndex >= matrixSize = pure (Right ()) + | otherwise = do + indexResult <- convergeIndex splitIndex 0 + case indexResult of + Left err -> pure (Left err) + Right () -> solveAllIndices (splitIndex + 1) + + convergeIndex !splitIndex !iterationCount + | iterationCount >= iterationLimit = do + maxOffDiagonal <- maximumOffDiagonalMagnitude splitIndex matrixSize offDiagonalValues + pure + ( Left + ( InvariantViolation + ( "tridiagonal eigensolve exhausted implicit-QL iteration budget at block " + <> show (splitIndex, matrixSize - 1) + <> " after " + <> show iterationCount + <> " iterations; max off-diagonal=" + <> show maxOffDiagonal + ) + ) + ) + | otherwise = do + activeIndex <- findActiveSplit splitIndex + if activeIndex == splitIndex + then pure (Right ()) + else do + implicitQLStep splitIndex activeIndex diagonalValues offDiagonalValues eigenvectors + convergeIndex splitIndex (iterationCount + 1) + + findActiveSplit !splitIndex = go splitIndex + where + go !candidateIndex + | candidateIndex >= matrixSize - 1 = pure (matrixSize - 1) + | otherwise = do + offDiagonal <- readPrimArray offDiagonalValues candidateIndex + leftDiagonal <- readPrimArray diagonalValues candidateIndex + rightDiagonal <- readPrimArray diagonalValues (candidateIndex + 1) + if negligibleOffDiagonal offDiagonal leftDiagonal rightDiagonal + then writePrimArray offDiagonalValues candidateIndex 0.0 >> pure candidateIndex + else go (candidateIndex + 1) + +implicitQLStep :: + Int -> + Int -> + MutablePrimArray s Double -> + MutablePrimArray s Double -> + MutableDenseWork s -> + ST s () +implicitQLStep !splitIndex !activeIndex diagonalValues offDiagonalValues (MutableDenseWork eigenvectorRowCount _ eigenvectorPayload) = do + leftDiagonal <- readPrimArray diagonalValues splitIndex + nextDiagonal <- readPrimArray diagonalValues (splitIndex + 1) + leftOffDiagonal <- readPrimArray offDiagonalValues splitIndex + activeDiagonal <- readPrimArray diagonalValues activeIndex + let !shiftRatio = (nextDiagonal - leftDiagonal) / (2.0 * leftOffDiagonal) + !shiftRadius = hypotStable shiftRatio 1.0 + !shiftDenominator = shiftRatio + copySignMagnitude shiftRadius shiftRatio + !initialShift = activeDiagonal - leftDiagonal + (leftOffDiagonal / shiftDenominator) + sweepDown (activeIndex - 1) 1.0 1.0 initialShift 0.0 + where + sweepDown !indexValue !previousCosine !previousSine !currentShift !currentCorrection + | indexValue < splitIndex = do + updatedLeftDiagonal <- readPrimArray diagonalValues splitIndex + writePrimArray diagonalValues splitIndex (updatedLeftDiagonal - currentCorrection) + writePrimArray offDiagonalValues splitIndex currentShift + writePrimArray offDiagonalValues activeIndex 0.0 + | otherwise = do + currentOffDiagonal <- readPrimArray offDiagonalValues indexValue + currentDiagonal <- readPrimArray diagonalValues indexValue + nextDiagonal <- readPrimArray diagonalValues (indexValue + 1) + let !fValue = previousSine * currentOffDiagonal + !bValue = previousCosine * currentOffDiagonal + rotateWithGivens indexValue bValue currentDiagonal nextDiagonal currentCorrection fValue currentShift + + rotateWithGivens !indexValue !bValue !currentDiagonal !nextDiagonal !currentCorrection !fValue !gValue + | fValue == 0.0 && gValue == 0.0 = + finishRotation indexValue bValue currentDiagonal nextDiagonal currentCorrection 1.0 0.0 0.0 + | abs fValue >= abs gValue = + let !normalizedCosine = gValue / fValue + !radius = hypotStable normalizedCosine 1.0 + !sineValue = 1.0 / radius + !cosineValue = normalizedCosine * sineValue + in finishRotation indexValue bValue currentDiagonal nextDiagonal currentCorrection cosineValue sineValue (fValue * radius) + | otherwise = + let !normalizedSine = fValue / gValue + !radius = hypotStable normalizedSine 1.0 + !cosineValue = 1.0 / radius + !sineValue = normalizedSine * cosineValue + in finishRotation indexValue bValue currentDiagonal nextDiagonal currentCorrection cosineValue sineValue (gValue * radius) + + finishRotation !indexValue !bValue !currentDiagonal !nextDiagonal !currentCorrection !nextCosine !nextSine !updatedOffDiagonal = do + let !nextDiagonalBase = nextDiagonal - currentCorrection + !rotationRadius = ((currentDiagonal - nextDiagonalBase) * nextSine) + (2.0 * nextCosine * bValue) + !nextCorrection = nextSine * rotationRadius + !updatedNextDiagonal = nextDiagonalBase + nextCorrection + !nextShift = (nextCosine * rotationRadius) - bValue + writePrimArray offDiagonalValues (indexValue + 1) updatedOffDiagonal + writePrimArray diagonalValues (indexValue + 1) updatedNextDiagonal + rotateEigenvectorColumnsAt indexValue (indexValue + 1) nextCosine nextSine + sweepDown (indexValue - 1) nextCosine nextSine nextShift nextCorrection + + rotateEigenvectorColumnsAt !leftColumn !rightColumn !cosineValue !sineValue = rotateRows 0 + where + !leftBase = leftColumn * eigenvectorRowCount + !rightBase = rightColumn * eigenvectorRowCount + rotateRows !rowIndex + | rowIndex >= eigenvectorRowCount = pure () + | otherwise = do + leftEntry <- readPrimArray eigenvectorPayload (leftBase + rowIndex) + rightEntry <- readPrimArray eigenvectorPayload (rightBase + rowIndex) + writePrimArray eigenvectorPayload (leftBase + rowIndex) (cosineValue * leftEntry - sineValue * rightEntry) + writePrimArray eigenvectorPayload (rightBase + rowIndex) (sineValue * leftEntry + cosineValue * rightEntry) + rotateRows (rowIndex + 1) + +scaleTridiagonal :: Int -> MutablePrimArray s Double -> MutablePrimArray s Double -> ST s Double +scaleTridiagonal !matrixSize diagonalValues offDiagonalValues = do + maximumMagnitude <- maximumTridiagonalMagnitude matrixSize diagonalValues offDiagonalValues + let !safeMaximum = sqrt maxFiniteDouble * 0.25 + !safeMinimum = sqrt safeMinimumDouble / epsDouble + !scaleValue + | maximumMagnitude == 0.0 = 1.0 + | maximumMagnitude > safeMaximum = safeMaximum / maximumMagnitude + | maximumMagnitude < safeMinimum = safeMinimum / maximumMagnitude + | otherwise = 1.0 + when (scaleValue /= 1.0) $ do + forIndex 0 matrixSize $ \indexValue -> do + diagonalEntry <- readPrimArray diagonalValues indexValue + writePrimArray diagonalValues indexValue (scaleValue * diagonalEntry) + forIndex 0 (max 0 (matrixSize - 1)) $ \indexValue -> do + offDiagonalEntry <- readPrimArray offDiagonalValues indexValue + writePrimArray offDiagonalValues indexValue (scaleValue * offDiagonalEntry) + pure scaleValue +{-# INLINE scaleTridiagonal #-} + +unscaleDiagonal :: Int -> MutablePrimArray s Double -> Double -> ST s () +unscaleDiagonal !matrixSize diagonalValues !scaleValue = + when (scaleValue /= 1.0) $ do + let !inverseScale = 1.0 / scaleValue + forIndex 0 matrixSize $ \indexValue -> do + diagonalEntry <- readPrimArray diagonalValues indexValue + writePrimArray diagonalValues indexValue (inverseScale * diagonalEntry) +{-# INLINE unscaleDiagonal #-} + +maximumTridiagonalMagnitude :: Int -> MutablePrimArray s Double -> MutablePrimArray s Double -> ST s Double +maximumTridiagonalMagnitude !matrixSize diagonalValues offDiagonalValues = do + diagonalMaximum <- maximumArrayMagnitude 0 matrixSize diagonalValues 0.0 + maximumArrayMagnitude 0 (max 0 (matrixSize - 1)) offDiagonalValues diagonalMaximum +{-# INLINE maximumTridiagonalMagnitude #-} + +maximumOffDiagonalMagnitude :: Int -> Int -> MutablePrimArray s Double -> ST s Double +maximumOffDiagonalMagnitude !startIndex !matrixSize offDiagonalValues = + maximumArrayMagnitude startIndex (max startIndex (matrixSize - 1)) offDiagonalValues 0.0 +{-# INLINE maximumOffDiagonalMagnitude #-} + +maximumArrayMagnitude :: Int -> Int -> MutablePrimArray s Double -> Double -> ST s Double +maximumArrayMagnitude !startIndex !stopIndex arrayValues !initialMaximum = go startIndex initialMaximum + where + go !indexValue !currentMaximum + | indexValue >= stopIndex = pure currentMaximum + | otherwise = do + entryValue <- readPrimArray arrayValues indexValue + go (indexValue + 1) (max currentMaximum (abs entryValue)) +{-# INLINE maximumArrayMagnitude #-} + +negligibleOffDiagonal :: Double -> Double -> Double -> Bool +negligibleOffDiagonal !offDiagonal !leftDiagonal !rightDiagonal = + abs offDiagonal <= (64.0 * epsDouble * (abs leftDiagonal + abs rightDiagonal)) + safeMinimumDouble +{-# INLINE negligibleOffDiagonal #-} + +sortEigenpairsAscending :: Int -> MutablePrimArray s Double -> MutableDenseWork s -> ST s () +sortEigenpairsAscending !matrixSize diagonalValues eigenvectors = + forIndex 0 matrixSize $ \targetIndex -> do + minimumIndex <- findMinimumIndex targetIndex (targetIndex + 1) + when (minimumIndex /= targetIndex) $ do + targetValue <- readPrimArray diagonalValues targetIndex + minimumValue <- readPrimArray diagonalValues minimumIndex + writePrimArray diagonalValues targetIndex minimumValue + writePrimArray diagonalValues minimumIndex targetValue + swapDenseColumnsTight eigenvectors targetIndex minimumIndex + where + findMinimumIndex !bestIndex !candidateIndex + | candidateIndex >= matrixSize = pure bestIndex + | otherwise = do + bestValue <- readPrimArray diagonalValues bestIndex + candidateValue <- readPrimArray diagonalValues candidateIndex + if candidateValue < bestValue + then findMinimumIndex candidateIndex (candidateIndex + 1) + else findMinimumIndex bestIndex (candidateIndex + 1) + +orthonormalizeDegenerateClusters :: Int -> MutablePrimArray s Double -> MutableDenseWork s -> ST s (Either MoonlightError ()) +orthonormalizeDegenerateClusters !matrixSize diagonalValues eigenvectors = processCluster 0 + where + processCluster !clusterStart + | clusterStart >= matrixSize = pure (Right ()) + | otherwise = do + clusterStop <- findClusterStop clusterStart (clusterStart + 1) + clusterResult <- orthonormalizeColumns clusterStart clusterStop + case clusterResult of + Left err -> pure (Left err) + Right () -> processCluster clusterStop + + findClusterStop !clusterStart !candidateIndex + | candidateIndex >= matrixSize = pure matrixSize + | otherwise = do + leftValue <- readPrimArray diagonalValues (candidateIndex - 1) + rightValue <- readPrimArray diagonalValues candidateIndex + if sameEigenCluster leftValue rightValue + then findClusterStop clusterStart (candidateIndex + 1) + else pure candidateIndex + + orthonormalizeColumns !clusterStart !clusterStop = normalizeColumnAt clusterStart + where + normalizeColumnAt !columnIndex + | columnIndex >= clusterStop = pure (Right ()) + | otherwise = do + subtractPriorColumns clusterStart columnIndex + normalized <- normalizeEigenvectorColumn eigenvectors columnIndex + if normalized + then normalizeColumnAt (columnIndex + 1) + else pure (Left (InvariantViolation ("symmetric eigen decomposition produced zero eigenvector at column " <> show columnIndex))) + + subtractPriorColumns !priorIndex !columnIndex + | priorIndex >= columnIndex = pure () + | otherwise = do + projection <- dotDenseColumnsTight eigenvectors priorIndex columnIndex + addScaledColumn eigenvectors priorIndex columnIndex (negate projection) + subtractPriorColumns (priorIndex + 1) columnIndex + +sameEigenCluster :: Double -> Double -> Bool +sameEigenCluster !leftValue !rightValue = + abs (leftValue - rightValue) <= 128.0 * epsDouble * max 1.0 (max (abs leftValue) (abs rightValue)) +{-# INLINE sameEigenCluster #-} + +addScaledColumn :: MutableDenseWork s -> Int -> Int -> Double -> ST s () +addScaledColumn (MutableDenseWork rowCount _ payload) !sourceColumn !targetColumn !scaleValue = + forIndex 0 rowCount $ \rowIndex -> do + sourceEntry <- readPrimArray payload (sourceBase + rowIndex) + targetEntry <- readPrimArray payload (targetBase + rowIndex) + writePrimArray payload (targetBase + rowIndex) (targetEntry + scaleValue * sourceEntry) + where + !sourceBase = sourceColumn * rowCount + !targetBase = targetColumn * rowCount +{-# INLINE addScaledColumn #-} + +normalizeEigenvectorColumn :: forall s. MutableDenseWork s -> Int -> ST s Bool +normalizeEigenvectorColumn eigenvectors !columnIndex = do + normValue <- columnNorm eigenvectors columnIndex + if normValue <= 0.0 || not (fieldValueValid normValue) + then pure False + else scaleDenseColumnTight eigenvectors columnIndex (1.0 / normValue) >> pure True + where + columnNorm :: MutableDenseWork s -> Int -> ST s Double + columnNorm (MutableDenseWork rowCount _ payload) !targetColumn = go 0 0.0 1.0 + where + !targetBase = targetColumn * rowCount + go !rowIndex !scaleValue !scaledSum + | rowIndex >= rowCount = + if scaleValue == 0.0 + then pure 0.0 + else pure (scaleValue * sqrt scaledSum) + | otherwise = do + entryValue <- readPrimArray payload (targetBase + rowIndex) + let !entryAbs = abs entryValue + if entryAbs == 0.0 + then go (rowIndex + 1) scaleValue scaledSum + else + if scaleValue < entryAbs + then + let !scaledRatio = scaleValue / entryAbs + in go (rowIndex + 1) entryAbs (1.0 + scaledSum * scaledRatio * scaledRatio) + else + let !scaledRatio = entryAbs / scaleValue + in go (rowIndex + 1) scaleValue (scaledSum + scaledRatio * scaledRatio) +{-# INLINE normalizeEigenvectorColumn #-} + +canonicalizeEigenvectorSigns :: Int -> MutableDenseWork s -> ST s () +canonicalizeEigenvectorSigns !matrixSize eigenvectors = + forIndex 0 matrixSize $ \columnIndex -> do + (_, maximumMagnitude, representativeValue) <- maximumMagnitudeInColumn eigenvectors columnIndex + when (maximumMagnitude > 0.0 && representativeValue < 0.0) $ + scaleDenseColumnTight eigenvectors columnIndex (-1.0) + +maximumMagnitudeInColumn :: MutableDenseWork s -> Int -> ST s (Int, Double, Double) +maximumMagnitudeInColumn (MutableDenseWork rowCount _ payload) !columnIndex = go 0 0 0.0 0.0 + where + !columnBase = columnIndex * rowCount + go !rowIndex !bestIndex !bestMagnitude !bestValue + | rowIndex >= rowCount = pure (bestIndex, bestMagnitude, bestValue) + | otherwise = do + entryValue <- readPrimArray payload (columnBase + rowIndex) + let !entryMagnitude = abs entryValue + if entryMagnitude > bestMagnitude + then go (rowIndex + 1) rowIndex entryMagnitude entryValue + else go (rowIndex + 1) bestIndex bestMagnitude bestValue +{-# INLINE maximumMagnitudeInColumn #-} + +swapDenseColumnsTight :: MutableDenseWork s -> Int -> Int -> ST s () +swapDenseColumnsTight (MutableDenseWork rowCount _ payload) !leftColumn !rightColumn = + when (leftColumn /= rightColumn) $ + forIndex 0 rowCount $ \rowIndex -> do + leftValue <- readPrimArray payload (leftBase + rowIndex) + rightValue <- readPrimArray payload (rightBase + rowIndex) + writePrimArray payload (leftBase + rowIndex) rightValue + writePrimArray payload (rightBase + rowIndex) leftValue + where + !leftBase = leftColumn * rowCount + !rightBase = rightColumn * rowCount +{-# INLINE swapDenseColumnsTight #-} + +dotDenseColumnsTight :: MutableDenseWork s -> Int -> Int -> ST s Double +dotDenseColumnsTight (MutableDenseWork rowCount _ payload) !leftColumn !rightColumn = go 0 0.0 + where + !leftBase = leftColumn * rowCount + !rightBase = rightColumn * rowCount + go !rowIndex !accumulator + | rowIndex >= rowCount = pure accumulator + | otherwise = do + leftValue <- readPrimArray payload (leftBase + rowIndex) + rightValue <- readPrimArray payload (rightBase + rowIndex) + go (rowIndex + 1) (accumulator + leftValue * rightValue) +{-# INLINE dotDenseColumnsTight #-} + +scaleDenseColumnTight :: MutableDenseWork s -> Int -> Double -> ST s () +scaleDenseColumnTight (MutableDenseWork rowCount _ payload) !columnIndex !scaleValue = + forIndex 0 rowCount $ \rowIndex -> do + entryValue <- readPrimArray payload (columnBase + rowIndex) + writePrimArray payload (columnBase + rowIndex) (scaleValue * entryValue) + where + !columnBase = columnIndex * rowCount +{-# INLINE scaleDenseColumnTight #-} + +eigenpairsFromMutable :: Int -> MutablePrimArray s Double -> MutableDenseWork s -> ST s [(Double, [Double])] +eigenpairsFromMutable !matrixSize diagonalValues (MutableDenseWork rowCount _ eigenvectorPayload) = collectColumns 0 [] + where + collectColumns !columnIndex !revPairs + | columnIndex >= matrixSize = pure (reverse revPairs) + | otherwise = do + eigenvalue <- readPrimArray diagonalValues columnIndex + eigenvector <- collectColumnEntries columnIndex 0 [] + collectColumns (columnIndex + 1) ((eigenvalue, eigenvector) : revPairs) + + collectColumnEntries !columnIndex !rowIndex !revEntries + | rowIndex >= matrixSize = pure (reverse revEntries) + | otherwise = do + entryValue <- readPrimArray eigenvectorPayload (rowIndex + columnIndex * rowCount) + collectColumnEntries columnIndex (rowIndex + 1) (entryValue : revEntries) diff --git a/src-geometry/Moonlight/LinAlg/Pure/Geometry/AABB.hs b/src-geometry/Moonlight/LinAlg/Pure/Geometry/AABB.hs new file mode 100644 --- /dev/null +++ b/src-geometry/Moonlight/LinAlg/Pure/Geometry/AABB.hs @@ -0,0 +1,135 @@ +module Moonlight.LinAlg.Pure.Geometry.AABB + ( AABB, + mkAabb, + aabbMin, + aabbMax, + aabbDimensions, + aabbCenter, + aabbHalfExtent, + aabbRadius, + symmetricAabb, + unionAabb, + expandAabb, + intersectAabbMaybe, + translateAabb, + transformAabb, + scaleAabb, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Kind (Type) +import Moonlight.Algebra.Pure.Lattice (JoinSemilattice (..)) +import Moonlight.LinAlg.Pure.Geometry.Transform.Affine + ( AffineTransform, + affineMaxScale, + atTranslation, + ) +import Moonlight.LinAlg.Pure.Geometry.Vec3 + ( Vec3 (..), + addVec3, + averageVec3, + mapVec3, + maxVec3, + minVec3, + mulVec3, + scaleVec3, + subVec3, + ) +import Prelude (Double, Eq, Maybe (..), Ord, Show, abs, seq, sqrt, (*), (+), (<=), (&&)) + +type AABB :: Type +data AABB = AABB + { aabbMin :: {-# UNPACK #-} !Vec3, + aabbMax :: {-# UNPACK #-} !Vec3 + } + deriving stock (Eq, Ord, Show) + +instance NFData AABB where + rnf aabbValue = aabbValue `seq` () + +instance JoinSemilattice AABB where + join leftAabb rightAabb = + AABB + { aabbMin = minVec3 (aabbMin leftAabb) (aabbMin rightAabb), + aabbMax = maxVec3 (aabbMax leftAabb) (aabbMax rightAabb) + } + +mkAabb :: Vec3 -> Vec3 -> Maybe AABB +mkAabb minimumCorner maximumCorner = + let Vec3 minX minY minZ = minimumCorner + Vec3 maxX maxY maxZ = maximumCorner + in if minX <= maxX && minY <= maxY && minZ <= maxZ + then Just (AABB minimumCorner maximumCorner) + else Nothing + +aabbDimensions :: AABB -> Vec3 +aabbDimensions aabbValue = + subVec3 (aabbMax aabbValue) (aabbMin aabbValue) + +aabbCenter :: AABB -> Vec3 +aabbCenter aabbValue = + averageVec3 (aabbMin aabbValue) (aabbMax aabbValue) + +aabbHalfExtent :: AABB -> Vec3 +aabbHalfExtent aabbValue = + scaleVec3 0.5 (aabbDimensions aabbValue) + +aabbRadius :: AABB -> Double +aabbRadius aabbValue = + let Vec3 halfX halfY halfZ = aabbHalfExtent aabbValue + in sqrt (halfX * halfX + halfY * halfY + halfZ * halfZ) + +symmetricAabb :: Double -> Double -> Double -> Maybe AABB +symmetricAabb halfX halfY halfZ = + mkAabb + (Vec3 (-halfX) (-halfY) (-halfZ)) + (Vec3 halfX halfY halfZ) + +unionAabb :: AABB -> AABB -> AABB +unionAabb = join + +expandAabb :: Double -> AABB -> Maybe AABB +expandAabb radius aabbValue = + if 0.0 <= radius + then + mkAabb + (shiftVec3 (-radius) (aabbMin aabbValue)) + (shiftVec3 radius (aabbMax aabbValue)) + else Nothing + +intersectAabbMaybe :: AABB -> AABB -> Maybe AABB +intersectAabbMaybe leftAabb rightAabb = + mkAabb + (maxVec3 (aabbMin leftAabb) (aabbMin rightAabb)) + (minVec3 (aabbMax leftAabb) (aabbMax rightAabb)) + +translateAabb :: Vec3 -> AABB -> AABB +translateAabb translationVector aabbValue = + AABB + { aabbMin = addVec3 translationVector (aabbMin aabbValue), + aabbMax = addVec3 translationVector (aabbMax aabbValue) + } + +transformAabb :: AffineTransform -> AABB -> AABB +transformAabb affineTransform aabbValue = + let radius = aabbRadius aabbValue * affineMaxScale affineTransform + translatedCenter = addVec3 (atTranslation affineTransform) (aabbCenter aabbValue) + in AABB + { aabbMin = shiftVec3 (-radius) translatedCenter, + aabbMax = shiftVec3 radius translatedCenter + } + +scaleAabb :: Vec3 -> AABB -> AABB +scaleAabb scaleVector aabbValue = + let center = aabbCenter aabbValue + halfExtent = aabbHalfExtent aabbValue + scaledCenter = mulVec3 scaleVector center + scaledHalfExtent = mulVec3 (mapVec3 abs scaleVector) halfExtent + in AABB + { aabbMin = subVec3 scaledCenter scaledHalfExtent, + aabbMax = addVec3 scaledCenter scaledHalfExtent + } + +shiftVec3 :: Double -> Vec3 -> Vec3 +shiftVec3 offset = mapVec3 (+ offset) diff --git a/src-geometry/Moonlight/LinAlg/Pure/Geometry/AABB2.hs b/src-geometry/Moonlight/LinAlg/Pure/Geometry/AABB2.hs new file mode 100644 --- /dev/null +++ b/src-geometry/Moonlight/LinAlg/Pure/Geometry/AABB2.hs @@ -0,0 +1,132 @@ +module Moonlight.LinAlg.Pure.Geometry.AABB2 + ( AABB2, + singletonAabb2, + mkAabb2, + aabb2Min, + aabb2Max, + aabb2Dimensions, + aabb2Center, + aabb2HalfExtent, + aabb2Radius, + symmetricAabb2, + unionAabb2, + expandAabb2, + intersectAabb2Maybe, + translateAabb2, + scaleAabb2, + containsVec2, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Kind (Type) +import Moonlight.Algebra.Pure.Lattice (JoinSemilattice (..)) +import Moonlight.LinAlg.Pure.Geometry.Vec2 + ( Vec2 (..), + addVec2, + averageVec2, + mapVec2, + maxVec2, + minVec2, + mulVec2, + scaleVec2, + subVec2, + ) +import Prelude (Bool, Double, Eq, Maybe (..), Ord, Show, abs, seq, sqrt, (*), (+), (<=), (>=), (&&)) + +type AABB2 :: Type +data AABB2 = AABB2 + { aabb2Min :: {-# UNPACK #-} !Vec2, + aabb2Max :: {-# UNPACK #-} !Vec2 + } + deriving stock (Eq, Ord, Show) + +instance NFData AABB2 where + rnf aabbValue = aabbValue `seq` () + +instance JoinSemilattice AABB2 where + join leftAabb rightAabb = + AABB2 + { aabb2Min = minVec2 (aabb2Min leftAabb) (aabb2Min rightAabb), + aabb2Max = maxVec2 (aabb2Max leftAabb) (aabb2Max rightAabb) + } + +singletonAabb2 :: Vec2 -> AABB2 +singletonAabb2 point = + AABB2 + { aabb2Min = point, + aabb2Max = point + } + +mkAabb2 :: Vec2 -> Vec2 -> Maybe AABB2 +mkAabb2 minimumCorner maximumCorner = + let Vec2 minX minY = minimumCorner + Vec2 maxX maxY = maximumCorner + in if minX <= maxX && minY <= maxY + then Just (AABB2 minimumCorner maximumCorner) + else Nothing + +aabb2Dimensions :: AABB2 -> Vec2 +aabb2Dimensions aabbValue = + subVec2 (aabb2Max aabbValue) (aabb2Min aabbValue) + +aabb2Center :: AABB2 -> Vec2 +aabb2Center aabbValue = + averageVec2 (aabb2Min aabbValue) (aabb2Max aabbValue) + +aabb2HalfExtent :: AABB2 -> Vec2 +aabb2HalfExtent aabbValue = + scaleVec2 0.5 (aabb2Dimensions aabbValue) + +aabb2Radius :: AABB2 -> Double +aabb2Radius aabbValue = + let Vec2 halfX halfY = aabb2HalfExtent aabbValue + in sqrt (halfX * halfX + halfY * halfY) + +symmetricAabb2 :: Double -> Double -> Maybe AABB2 +symmetricAabb2 halfX halfY = + mkAabb2 + (Vec2 (-halfX) (-halfY)) + (Vec2 halfX halfY) + +unionAabb2 :: AABB2 -> AABB2 -> AABB2 +unionAabb2 = join + +expandAabb2 :: Double -> AABB2 -> Maybe AABB2 +expandAabb2 radius aabbValue = + if 0.0 <= radius + then + mkAabb2 + (mapVec2 (+ (-radius)) (aabb2Min aabbValue)) + (mapVec2 (+ radius) (aabb2Max aabbValue)) + else Nothing + +intersectAabb2Maybe :: AABB2 -> AABB2 -> Maybe AABB2 +intersectAabb2Maybe leftAabb rightAabb = + mkAabb2 + (maxVec2 (aabb2Min leftAabb) (aabb2Min rightAabb)) + (minVec2 (aabb2Max leftAabb) (aabb2Max rightAabb)) + +translateAabb2 :: Vec2 -> AABB2 -> AABB2 +translateAabb2 translationVector aabbValue = + AABB2 + { aabb2Min = addVec2 translationVector (aabb2Min aabbValue), + aabb2Max = addVec2 translationVector (aabb2Max aabbValue) + } + +scaleAabb2 :: Vec2 -> AABB2 -> AABB2 +scaleAabb2 scaleVector aabbValue = + let center = aabb2Center aabbValue + halfExtent = aabb2HalfExtent aabbValue + scaledCenter = mulVec2 scaleVector center + scaledHalfExtent = mulVec2 (mapVec2 abs scaleVector) halfExtent + in AABB2 + { aabb2Min = subVec2 scaledCenter scaledHalfExtent, + aabb2Max = addVec2 scaledCenter scaledHalfExtent + } + +containsVec2 :: AABB2 -> Vec2 -> Bool +containsVec2 aabbValue (Vec2 px py) = + let Vec2 minX minY = aabb2Min aabbValue + Vec2 maxX maxY = aabb2Max aabbValue + in px <= maxX && px >= minX && py <= maxY && py >= minY diff --git a/src-geometry/Moonlight/LinAlg/Pure/Geometry/Frame.hs b/src-geometry/Moonlight/LinAlg/Pure/Geometry/Frame.hs new file mode 100644 --- /dev/null +++ b/src-geometry/Moonlight/LinAlg/Pure/Geometry/Frame.hs @@ -0,0 +1,89 @@ +{-# LANGUAGE LambdaCase #-} + +module Moonlight.LinAlg.Pure.Geometry.Frame + ( OrthonormalFrame, + identityOrthonormalFrame, + orthonormalFrameColumns, + orthonormalFrameFromColumns, + orthonormalFrameFromMatrixEntries, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Kind (Type) +import Moonlight.LinAlg.Pure.Geometry.Vec3 + ( Vec3 (..), + dotVec3, + magnitudeVec3, + ) +import Prelude (Bool, Double, Eq, Maybe (..), Read, Show, abs, all, seq, (-), (&&), (<=)) + +type OrthonormalFrame :: Type +data OrthonormalFrame + = OrthonormalFrame + {-# UNPACK #-} !Vec3 + {-# UNPACK #-} !Vec3 + {-# UNPACK #-} !Vec3 + deriving stock (Eq, Show, Read) + +instance NFData OrthonormalFrame where + rnf frameValue = frameValue `seq` () + +identityOrthonormalFrame :: OrthonormalFrame +identityOrthonormalFrame = + OrthonormalFrame + (Vec3 1.0 0.0 0.0) + (Vec3 0.0 1.0 0.0) + (Vec3 0.0 0.0 1.0) + +orthonormalFrameColumns :: OrthonormalFrame -> (Vec3, Vec3, Vec3) +orthonormalFrameColumns (OrthonormalFrame axis1 axis2 axis3) = (axis1, axis2, axis3) + +orthonormalFrameFromColumns :: [Vec3] -> Maybe OrthonormalFrame +orthonormalFrameFromColumns = \case + axis1 : axis2 : axis3 : [] + | isOrthonormalTriple axis1 axis2 axis3 -> + Just (OrthonormalFrame axis1 axis2 axis3) + _ -> + Nothing + +orthonormalFrameFromMatrixEntries :: [Double] -> Maybe OrthonormalFrame +orthonormalFrameFromMatrixEntries = \case + [ r1c1, + r1c2, + r1c3, + r2c1, + r2c2, + r2c3, + r3c1, + r3c2, + r3c3 + ] -> + orthonormalFrameFromColumns + [ Vec3 r1c1 r2c1 r3c1, + Vec3 r1c2 r2c2 r3c2, + Vec3 r1c3 r2c3 r3c3 + ] + _ -> + Nothing + +isOrthonormalTriple :: Vec3 -> Vec3 -> Vec3 -> Bool +isOrthonormalTriple axis1 axis2 axis3 = + all isUnitVector [axis1, axis2, axis3] + && all + orthogonalWithinTolerance + [ (axis1, axis2), + (axis1, axis3), + (axis2, axis3) + ] + +isUnitVector :: Vec3 -> Bool +isUnitVector axisValue = + abs (magnitudeVec3 axisValue - 1.0) <= frameTolerance + +orthogonalWithinTolerance :: (Vec3, Vec3) -> Bool +orthogonalWithinTolerance (leftAxis, rightAxis) = + abs (dotVec3 leftAxis rightAxis) <= frameTolerance + +frameTolerance :: Double +frameTolerance = 1.0e-10 diff --git a/src-geometry/Moonlight/LinAlg/Pure/Geometry/Symmetric.hs b/src-geometry/Moonlight/LinAlg/Pure/Geometry/Symmetric.hs new file mode 100644 --- /dev/null +++ b/src-geometry/Moonlight/LinAlg/Pure/Geometry/Symmetric.hs @@ -0,0 +1,987 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeApplications #-} + +module Moonlight.LinAlg.Pure.Geometry.Symmetric + ( DiagonalizedSymmetric2 (..), + mapDiagonalizedSymmetric2, + diagonalizedSymmetric2ToTensor, + diagonalizedSymmetric2ToVec2, + eigendecomposeSymmetric2With, + Symmetric2 (..), + mapSymmetric2, + zipSymmetric2With, + diagonalSymmetric2, + scaleSymmetric2, + outerSymmetric2, + traceSymmetric2, + applySymmetric2, + symmetric2Entries, + symmetric2ToMatrix, + eigendecomposeSymmetric2, + DiagonalizedSymmetric3 (..), + mapDiagonalizedSymmetric3, + diagonalizedSymmetric3ToTensor, + diagonalizedSymmetric3ToVec3, + eigendecomposeSymmetric3With, + eigendecomposeSymmetric3OrthonormalFrame, + Symmetric3 (..), + mapSymmetric3, + zipSymmetric3With, + diagonalSymmetric3, + scaleSymmetric3, + outerSymmetric3, + traceSymmetric3, + applySymmetric3, + symmetric3Entries, + symmetric3ToMatrix, + eigendecomposeSymmetric3, + ) +where + +import Data.Kind (Type) +import Data.Maybe (fromMaybe) +import Moonlight.Algebra.Pure.Module (BilinearSpace (..), Module (..), VectorSpace) +import Moonlight.Core (AdditiveGroup (..), AdditiveMonoid (..), Field, Metric (..), MoonlightError (..), MultiplicativeMonoid (..), Ring, fieldValueValid) +import Moonlight.LinAlg.Pure.Geometry.Frame (OrthonormalFrame, identityOrthonormalFrame, orthonormalFrameFromMatrixEntries) +import Moonlight.LinAlg.Pure.Dense.Types (Matrix, Vector, fromListMatrix, fromListVector, toListMatrix, toListVector) +import Moonlight.LinAlg.Pure.Geometry.Vec2 (Vec2 (..)) +import Moonlight.LinAlg.Pure.Geometry.Vec3 (Vec3 (..)) +import Prelude + ( Bool, + Double, + Either (..), + Eq, + Int, + Maybe (..), + Monoid (..), + Ord, + Read, + Semigroup (..), + Show, + abs, + acos, + all, + cos, + foldr, + max, + min, + not, + otherwise, + pi, + pure, + sqrt, + (&&), + (||), + (*), + (+), + (-), + (.), + (/), + (<), + (<=), + (>=), + (>), + (>>=), + ) + +type DiagonalizedSymmetric2 :: Type -> Type -> Type +data DiagonalizedSymmetric2 axes a = DiagonalizedSymmetric2 + { diag2XX :: !a, + diag2YY :: !a, + diag2Axes :: !axes + } + deriving stock (Eq, Ord, Show, Read) + +mapDiagonalizedSymmetric2 :: + (a -> b) -> + DiagonalizedSymmetric2 axes a -> + DiagonalizedSymmetric2 axes b +mapDiagonalizedSymmetric2 transform diagonalizedValue = + DiagonalizedSymmetric2 + { diag2XX = transform (diag2XX diagonalizedValue), + diag2YY = transform (diag2YY diagonalizedValue), + diag2Axes = diag2Axes diagonalizedValue + } + +diagonalizedSymmetric2ToTensor :: + AdditiveGroup a => + DiagonalizedSymmetric2 axes a -> + Symmetric2 a +diagonalizedSymmetric2ToTensor diagonalizedValue = + diagonalSymmetric2 + (diag2XX diagonalizedValue) + (diag2YY diagonalizedValue) + +diagonalizedSymmetric2ToVec2 :: + DiagonalizedSymmetric2 axes Double -> + Vec2 +diagonalizedSymmetric2ToVec2 diagonalizedValue = + Vec2 + (diag2XX diagonalizedValue) + (diag2YY diagonalizedValue) + +eigendecomposeSymmetric2With :: + ([Double] -> Maybe axes) -> + axes -> + Symmetric2 Double -> + Either MoonlightError (DiagonalizedSymmetric2 axes Double) +eigendecomposeSymmetric2With decodeAxes fallbackAxes tensorValue = do + (eigenvalues, eigenvectors) <- eigendecomposeSymmetric2 tensorValue + let resolvedAxes = fromMaybe fallbackAxes (decodeAxes (toListMatrix eigenvectors)) + in pure + ( case toListVector eigenvalues of + lambda1 : lambda2 : _ -> + DiagonalizedSymmetric2 + { diag2XX = lambda1, + diag2YY = lambda2, + diag2Axes = resolvedAxes + } + lambda1 : _ -> + DiagonalizedSymmetric2 + { diag2XX = lambda1, + diag2YY = 0.0, + diag2Axes = resolvedAxes + } + [] -> + DiagonalizedSymmetric2 + { diag2XX = 0.0, + diag2YY = 0.0, + diag2Axes = resolvedAxes + } + ) + +type Symmetric2 :: Type -> Type +data Symmetric2 a = Symmetric2 + { sym2XX :: !a, + sym2XY :: !a, + sym2YY :: !a + } + deriving stock (Eq, Ord, Show, Read) + +instance AdditiveMonoid a => Semigroup (Symmetric2 a) where + (<>) = zipSymmetric2With add + +instance AdditiveMonoid a => Monoid (Symmetric2 a) where + mempty = diagonalSymmetric2 zero zero + +instance AdditiveMonoid a => AdditiveMonoid (Symmetric2 a) where + zero = mempty + add = (<>) + +instance AdditiveGroup a => AdditiveGroup (Symmetric2 a) where + neg = mapSymmetric2 neg + +instance Ring a => Module a (Symmetric2 a) where + scale = scaleSymmetric2 + +instance Field a => VectorSpace a (Symmetric2 a) + +instance Field a => BilinearSpace a (Symmetric2 a) where + bilinearForm leftValue rightValue = + let doubledUnit = add one one + in foldr + add + zero + [ mul (sym2XX leftValue) (sym2XX rightValue), + mul (sym2YY leftValue) (sym2YY rightValue), + mul doubledUnit (mul (sym2XY leftValue) (sym2XY rightValue)) + ] + +instance Metric (Symmetric2 Double) where + type Magnitude (Symmetric2 Double) = Double + magnitude tensorValue = sqrt (bilinearForm tensorValue tensorValue) + +mapSymmetric2 :: (a -> b) -> Symmetric2 a -> Symmetric2 b +mapSymmetric2 transform tensorValue = + Symmetric2 + { sym2XX = transform (sym2XX tensorValue), + sym2XY = transform (sym2XY tensorValue), + sym2YY = transform (sym2YY tensorValue) + } + +zipSymmetric2With :: (a -> b -> c) -> Symmetric2 a -> Symmetric2 b -> Symmetric2 c +zipSymmetric2With combine leftValue rightValue = + Symmetric2 + { sym2XX = combine (sym2XX leftValue) (sym2XX rightValue), + sym2XY = combine (sym2XY leftValue) (sym2XY rightValue), + sym2YY = combine (sym2YY leftValue) (sym2YY rightValue) + } + +diagonalSymmetric2 :: AdditiveMonoid a => a -> a -> Symmetric2 a +diagonalSymmetric2 diagonalX diagonalY = + Symmetric2 + { sym2XX = diagonalX, + sym2XY = zero, + sym2YY = diagonalY + } + +scaleSymmetric2 :: MultiplicativeMonoid a => a -> Symmetric2 a -> Symmetric2 a +scaleSymmetric2 = mapSymmetric2 . mul + +outerSymmetric2 :: Double -> Vec2 -> Symmetric2 Double +outerSymmetric2 weightValue (Vec2 xValue yValue) = + Symmetric2 + { sym2XX = weightValue * xValue * xValue, + sym2XY = weightValue * xValue * yValue, + sym2YY = weightValue * yValue * yValue + } + +traceSymmetric2 :: AdditiveGroup a => Symmetric2 a -> a +traceSymmetric2 tensorValue = + add (sym2XX tensorValue) (sym2YY tensorValue) + +applySymmetric2 :: Symmetric2 Double -> Vec2 -> Vec2 +applySymmetric2 tensorValue (Vec2 xValue yValue) = + Vec2 + (sym2XX tensorValue * xValue + sym2XY tensorValue * yValue) + (sym2XY tensorValue * xValue + sym2YY tensorValue * yValue) + +symmetric2Entries :: Symmetric2 a -> [a] +symmetric2Entries tensorValue = + [ sym2XX tensorValue, + sym2XY tensorValue, + sym2XY tensorValue, + sym2YY tensorValue + ] + +symmetric2ToMatrix :: Symmetric2 a -> Either MoonlightError (Matrix 2 2 a) +symmetric2ToMatrix = fromListMatrix @2 @2 . symmetric2Entries + +eigendecomposeSymmetric2 :: + Symmetric2 Double -> + Either MoonlightError (Vector 2 Double, Matrix 2 2 Double) +eigendecomposeSymmetric2 tensorValue = + if not (all fieldValueValid (symmetric2Entries tensorValue)) + then Left (InvariantViolation "symmetric2 eigendecomposition requires finite entries") + else + let scaleValue = symmetric2Scale tensorValue + in if scaleValue <= 0.0 + then symmetric2Result 1.0 0.0 (Vec2 1.0 0.0) 0.0 (Vec2 0.0 1.0) + else + let scaledTensor = scaleSymmetric2 (1.0 / scaleValue) tensorValue + aValue = sym2XX scaledTensor + bValue = sym2XY scaledTensor + dValue = sym2YY scaledTensor + meanValue = (aValue + dValue) / 2.0 + halfDifference = (aValue - dValue) / 2.0 + radiusValue = sqrt (halfDifference * halfDifference + bValue * bValue) + firstEigenvalue = meanValue + radiusValue + firstVector = symmetric2Eigenvector scaledTensor firstEigenvalue + secondVector = canonicalizeVec2Sign (perpendicularVec2 firstVector) + refinedFirst = rayleighQuotient2 tensorValue firstVector + refinedSecond = rayleighQuotient2 tensorValue secondVector + in if refinedFirst >= refinedSecond + then symmetric2Result 1.0 refinedFirst firstVector refinedSecond secondVector + else symmetric2Result 1.0 refinedSecond secondVector refinedFirst firstVector + +perpendicularVec2 :: Vec2 -> Vec2 +perpendicularVec2 (Vec2 xValue yValue) = + Vec2 (0.0 - yValue) xValue + +rayleighQuotient2 :: Symmetric2 Double -> Vec2 -> Double +rayleighQuotient2 tensorValue vectorValue@(Vec2 xValue yValue) = + let Vec2 imageX imageY = applySymmetric2 tensorValue vectorValue + in (imageX * xValue + imageY * yValue) / (xValue * xValue + yValue * yValue) + +type DiagonalizedSymmetric3 :: Type -> Type -> Type +data DiagonalizedSymmetric3 axes a = DiagonalizedSymmetric3 + { diag3XX :: a, + diag3YY :: a, + diag3ZZ :: a, + diag3Axes :: axes + } + deriving stock (Eq, Ord, Show, Read) + +mapDiagonalizedSymmetric3 :: + (a -> b) -> + DiagonalizedSymmetric3 axes a -> + DiagonalizedSymmetric3 axes b +mapDiagonalizedSymmetric3 transform diagonalizedValue = + DiagonalizedSymmetric3 + { diag3XX = transform (diag3XX diagonalizedValue), + diag3YY = transform (diag3YY diagonalizedValue), + diag3ZZ = transform (diag3ZZ diagonalizedValue), + diag3Axes = diag3Axes diagonalizedValue + } + +diagonalizedSymmetric3ToTensor :: + AdditiveGroup a => + DiagonalizedSymmetric3 axes a -> + Symmetric3 a +diagonalizedSymmetric3ToTensor diagonalizedValue = + diagonalSymmetric3 + (diag3XX diagonalizedValue) + (diag3YY diagonalizedValue) + (diag3ZZ diagonalizedValue) + +diagonalizedSymmetric3ToVec3 :: + DiagonalizedSymmetric3 axes Double -> + Vec3 +diagonalizedSymmetric3ToVec3 diagonalizedValue = + Vec3 + (diag3XX diagonalizedValue) + (diag3YY diagonalizedValue) + (diag3ZZ diagonalizedValue) + +eigendecomposeSymmetric3With :: + ([Double] -> Maybe axes) -> + axes -> + Symmetric3 Double -> + Either MoonlightError (DiagonalizedSymmetric3 axes Double) +eigendecomposeSymmetric3With decodeAxes fallbackAxes tensorValue = do + (eigenvalues, eigenvectors) <- eigendecomposeSymmetric3 tensorValue + let resolvedAxes = fromMaybe fallbackAxes (decodeAxes (toListMatrix eigenvectors)) + in pure + ( case toListVector eigenvalues of + lambda1 : lambda2 : lambda3 : _ -> + DiagonalizedSymmetric3 + { diag3XX = lambda1, + diag3YY = lambda2, + diag3ZZ = lambda3, + diag3Axes = resolvedAxes + } + lambda1 : lambda2 : _ -> + DiagonalizedSymmetric3 + { diag3XX = lambda1, + diag3YY = lambda2, + diag3ZZ = 0.0, + diag3Axes = resolvedAxes + } + lambda1 : _ -> + DiagonalizedSymmetric3 + { diag3XX = lambda1, + diag3YY = 0.0, + diag3ZZ = 0.0, + diag3Axes = resolvedAxes + } + [] -> + DiagonalizedSymmetric3 + { diag3XX = 0.0, + diag3YY = 0.0, + diag3ZZ = 0.0, + diag3Axes = resolvedAxes + } + ) + +eigendecomposeSymmetric3OrthonormalFrame :: + Symmetric3 Double -> + Either MoonlightError (DiagonalizedSymmetric3 OrthonormalFrame Double) +eigendecomposeSymmetric3OrthonormalFrame = + eigendecomposeSymmetric3With + orthonormalFrameFromMatrixEntries + identityOrthonormalFrame + +type Symmetric3 :: Type -> Type +data Symmetric3 a = Symmetric3 + { sym3XX :: a, + sym3XY :: a, + sym3XZ :: a, + sym3YY :: a, + sym3YZ :: a, + sym3ZZ :: a + } + deriving stock (Eq, Ord, Show, Read) + +instance AdditiveMonoid a => Semigroup (Symmetric3 a) where + (<>) = zipSymmetric3With add + +instance AdditiveMonoid a => Monoid (Symmetric3 a) where + mempty = diagonalSymmetric3 zero zero zero + +instance AdditiveMonoid a => AdditiveMonoid (Symmetric3 a) where + zero = mempty + add = (<>) + +instance AdditiveGroup a => AdditiveGroup (Symmetric3 a) where + neg = mapSymmetric3 neg + +instance Ring a => Module a (Symmetric3 a) where + scale = scaleSymmetric3 + +instance Field a => VectorSpace a (Symmetric3 a) + +instance Field a => BilinearSpace a (Symmetric3 a) where + bilinearForm leftValue rightValue = + let doubledUnit = add one one + in foldr + add + zero + [ mul (sym3XX leftValue) (sym3XX rightValue), + mul (sym3YY leftValue) (sym3YY rightValue), + mul (sym3ZZ leftValue) (sym3ZZ rightValue), + mul + doubledUnit + ( foldr + add + zero + [ mul (sym3XY leftValue) (sym3XY rightValue), + mul (sym3XZ leftValue) (sym3XZ rightValue), + mul (sym3YZ leftValue) (sym3YZ rightValue) + ] + ) + ] + +instance Metric (Symmetric3 Double) where + type Magnitude (Symmetric3 Double) = Double + magnitude tensorValue = sqrt (bilinearForm tensorValue tensorValue) + +mapSymmetric3 :: (a -> b) -> Symmetric3 a -> Symmetric3 b +mapSymmetric3 transform tensorValue = + Symmetric3 + { sym3XX = transform (sym3XX tensorValue), + sym3XY = transform (sym3XY tensorValue), + sym3XZ = transform (sym3XZ tensorValue), + sym3YY = transform (sym3YY tensorValue), + sym3YZ = transform (sym3YZ tensorValue), + sym3ZZ = transform (sym3ZZ tensorValue) + } + +zipSymmetric3With :: (a -> b -> c) -> Symmetric3 a -> Symmetric3 b -> Symmetric3 c +zipSymmetric3With combine leftValue rightValue = + Symmetric3 + { sym3XX = combine (sym3XX leftValue) (sym3XX rightValue), + sym3XY = combine (sym3XY leftValue) (sym3XY rightValue), + sym3XZ = combine (sym3XZ leftValue) (sym3XZ rightValue), + sym3YY = combine (sym3YY leftValue) (sym3YY rightValue), + sym3YZ = combine (sym3YZ leftValue) (sym3YZ rightValue), + sym3ZZ = combine (sym3ZZ leftValue) (sym3ZZ rightValue) + } + +diagonalSymmetric3 :: AdditiveMonoid a => a -> a -> a -> Symmetric3 a +diagonalSymmetric3 diagonalX diagonalY diagonalZ = + Symmetric3 + { sym3XX = diagonalX, + sym3XY = zero, + sym3XZ = zero, + sym3YY = diagonalY, + sym3YZ = zero, + sym3ZZ = diagonalZ + } + +scaleSymmetric3 :: MultiplicativeMonoid a => a -> Symmetric3 a -> Symmetric3 a +scaleSymmetric3 = mapSymmetric3 . mul + +outerSymmetric3 :: Double -> Vec3 -> Symmetric3 Double +outerSymmetric3 weightValue (Vec3 xValue yValue zValue) = + Symmetric3 + { sym3XX = weightValue * xValue * xValue, + sym3XY = weightValue * xValue * yValue, + sym3XZ = weightValue * xValue * zValue, + sym3YY = weightValue * yValue * yValue, + sym3YZ = weightValue * yValue * zValue, + sym3ZZ = weightValue * zValue * zValue + } + +traceSymmetric3 :: AdditiveGroup a => Symmetric3 a -> a +traceSymmetric3 tensorValue = + foldr add zero [sym3XX tensorValue, sym3YY tensorValue, sym3ZZ tensorValue] + +applySymmetric3 :: Symmetric3 Double -> Vec3 -> Vec3 +applySymmetric3 tensorValue (Vec3 xValue yValue zValue) = + Vec3 + (sym3XX tensorValue * xValue + sym3XY tensorValue * yValue + sym3XZ tensorValue * zValue) + (sym3XY tensorValue * xValue + sym3YY tensorValue * yValue + sym3YZ tensorValue * zValue) + (sym3XZ tensorValue * xValue + sym3YZ tensorValue * yValue + sym3ZZ tensorValue * zValue) + +symmetric3Entries :: Symmetric3 a -> [a] +symmetric3Entries tensorValue = + [ sym3XX tensorValue, + sym3XY tensorValue, + sym3XZ tensorValue, + sym3XY tensorValue, + sym3YY tensorValue, + sym3YZ tensorValue, + sym3XZ tensorValue, + sym3YZ tensorValue, + sym3ZZ tensorValue + ] + +symmetric3ToMatrix :: Symmetric3 a -> Either MoonlightError (Matrix 3 3 a) +symmetric3ToMatrix = fromListMatrix @3 @3 . symmetric3Entries + +eigendecomposeSymmetric3 :: + Symmetric3 Double -> + Either MoonlightError (Vector 3 Double, Matrix 3 3 Double) +eigendecomposeSymmetric3 tensorValue = + if not (all fieldValueValid (symmetric3Entries tensorValue)) + then Left (InvariantViolation "symmetric3 eigendecomposition requires finite entries") + else + let scaleValue = symmetric3Scale tensorValue + in if scaleValue <= 0.0 + then + symmetric3Result + 1.0 + ( EigenColumn3 0.0 (Vec3 1.0 0.0 0.0), + EigenColumn3 0.0 (Vec3 0.0 1.0 0.0), + EigenColumn3 0.0 (Vec3 0.0 0.0 1.0) + ) + else + let scaledTensor = scaleSymmetric3 (1.0 / scaleValue) tensorValue + analyticEigenvalues = symmetric3AnalyticEigenvalues scaledTensor + analyticCandidate = symmetric3AnalyticCandidate scaledTensor analyticEigenvalues + (firstColumn, secondColumn, thirdColumn) = + case analyticCandidate of + Just eigenColumns + | symmetric3CandidateAcceptable scaledTensor eigenColumns -> + eigenColumns + _ -> symmetric3JacobiEigenColumns scaledTensor + in symmetric3Result + 1.0 + ( sortEigenColumns3 + (rayleighRefineColumn3 tensorValue firstColumn) + (rayleighRefineColumn3 tensorValue secondColumn) + (rayleighRefineColumn3 tensorValue thirdColumn) + ) + +symmetric2Scale :: Symmetric2 Double -> Double +symmetric2Scale tensorValue = + max (abs (sym2XX tensorValue)) (max (abs (sym2XY tensorValue)) (abs (sym2YY tensorValue))) + +symmetric3Scale :: Symmetric3 Double -> Double +symmetric3Scale tensorValue = + max + (abs (sym3XX tensorValue)) + ( max + (abs (sym3XY tensorValue)) + ( max + (abs (sym3XZ tensorValue)) + ( max + (abs (sym3YY tensorValue)) + (max (abs (sym3YZ tensorValue)) (abs (sym3ZZ tensorValue))) + ) + ) + ) + +symmetric2Result :: + Double -> + Double -> + Vec2 -> + Double -> + Vec2 -> + Either MoonlightError (Vector 2 Double, Matrix 2 2 Double) +symmetric2Result scaleValue firstEigenvalue firstVector secondEigenvalue secondVector = do + eigenvalues <- fromListVector @2 [scaleValue * firstEigenvalue, scaleValue * secondEigenvalue] + eigenvectors <- fromListMatrix @2 @2 (matrix2Columns firstVector secondVector) + pure (eigenvalues, eigenvectors) + +symmetric2Eigenvector :: Symmetric2 Double -> Double -> Vec2 +symmetric2Eigenvector tensorValue eigenvalue = + canonicalizeVec2Sign + ( if abs (sym2XY tensorValue) <= 1.0e-14 + then + if abs (sym2XX tensorValue - eigenvalue) <= abs (sym2YY tensorValue - eigenvalue) + then Vec2 1.0 0.0 + else Vec2 0.0 1.0 + else + normalizeVec2Or + (Vec2 1.0 0.0) + ( largerVec2 + (Vec2 (sym2XY tensorValue) (eigenvalue - sym2XX tensorValue)) + (Vec2 (eigenvalue - sym2YY tensorValue) (sym2XY tensorValue)) + ) + ) + +largerVec2 :: Vec2 -> Vec2 -> Vec2 +largerVec2 leftValue rightValue = + if vec2NormSquared leftValue >= vec2NormSquared rightValue + then leftValue + else rightValue + +matrix2Columns :: Vec2 -> Vec2 -> [Double] +matrix2Columns (Vec2 x1 y1) (Vec2 x2 y2) = + [x1, x2, y1, y2] + +normalizeVec2Or :: Vec2 -> Vec2 -> Vec2 +normalizeVec2Or fallbackValue vectorValue = + let normValue = sqrt (vec2NormSquared vectorValue) + in if normValue <= 1.0e-24 + then fallbackValue + else scaleVec2Local (1.0 / normValue) vectorValue + +vec2NormSquared :: Vec2 -> Double +vec2NormSquared (Vec2 xValue yValue) = + xValue * xValue + yValue * yValue + +scaleVec2Local :: Double -> Vec2 -> Vec2 +scaleVec2Local scaleValue (Vec2 xValue yValue) = + Vec2 (scaleValue * xValue) (scaleValue * yValue) + +canonicalizeVec2Sign :: Vec2 -> Vec2 +canonicalizeVec2Sign vectorValue@(Vec2 xValue yValue) + | abs xValue > 1.0e-14 = + if xValue < 0.0 + then scaleVec2Local (-1.0) vectorValue + else vectorValue + | yValue < 0.0 = scaleVec2Local (-1.0) vectorValue + | otherwise = vectorValue + +data EigenColumn3 = EigenColumn3 + { eigenColumn3Value :: !Double, + eigenColumn3Vector :: !Vec3 + } + +data Symmetric3Pivot + = PivotXY + | PivotXZ + | PivotYZ + +symmetric3AnalyticEigenvalues :: Symmetric3 Double -> (Double, Double, Double) +symmetric3AnalyticEigenvalues tensorValue = + let meanValue = (sym3XX tensorValue + sym3YY tensorValue + sym3ZZ tensorValue) / 3.0 + xxCentered = sym3XX tensorValue - meanValue + yyCentered = sym3YY tensorValue - meanValue + zzCentered = sym3ZZ tensorValue - meanValue + secondMoment = + xxCentered * xxCentered + + yyCentered * yyCentered + + zzCentered * zzCentered + + 2.0 * (sym3XY tensorValue * sym3XY tensorValue + sym3XZ tensorValue * sym3XZ tensorValue + sym3YZ tensorValue * sym3YZ tensorValue) + scaleMoment = sqrt (secondMoment / 6.0) + in if scaleMoment <= 1.0e-24 + then (meanValue, meanValue, meanValue) + else + let inverseScaleMoment = 1.0 / scaleMoment + normalized = + Symmetric3 + { sym3XX = xxCentered * inverseScaleMoment, + sym3XY = sym3XY tensorValue * inverseScaleMoment, + sym3XZ = sym3XZ tensorValue * inverseScaleMoment, + sym3YY = yyCentered * inverseScaleMoment, + sym3YZ = sym3YZ tensorValue * inverseScaleMoment, + sym3ZZ = zzCentered * inverseScaleMoment + } + determinantHalf = symmetric3Determinant normalized / 2.0 + angleValue = + if determinantHalf <= -1.0 + then pi / 3.0 + else + if determinantHalf >= 1.0 + then 0.0 + else acos determinantHalf / 3.0 + largestValue = meanValue + 2.0 * scaleMoment * cos angleValue + smallestValue = meanValue + 2.0 * scaleMoment * cos (angleValue + 2.0 * pi / 3.0) + middleValue = 3.0 * meanValue - largestValue - smallestValue + in sortEigenvalues3 largestValue middleValue smallestValue + +symmetric3AnalyticCandidate :: + Symmetric3 Double -> + (Double, Double, Double) -> + Maybe (EigenColumn3, EigenColumn3, EigenColumn3) +symmetric3AnalyticCandidate tensorValue eigenvalues@(firstEigenvalue, secondEigenvalue, thirdEigenvalue) + | symmetric3EigenvaluesDegenerate eigenvalues = Nothing + | otherwise = + case symmetric3Eigenvector tensorValue firstEigenvalue of + Nothing -> Nothing + Just firstVector -> + case symmetric3Eigenvector tensorValue secondEigenvalue >>= orthogonalizeVec3Against firstVector of + Nothing -> Nothing + Just secondVector -> + case normalizeVec3Maybe (crossVec3Local firstVector secondVector) of + Nothing -> Nothing + Just thirdVector -> + Just + ( EigenColumn3 firstEigenvalue (canonicalizeVec3Sign firstVector), + EigenColumn3 secondEigenvalue (canonicalizeVec3Sign secondVector), + EigenColumn3 thirdEigenvalue (canonicalizeVec3Sign thirdVector) + ) + +symmetric3EigenvaluesDegenerate :: (Double, Double, Double) -> Bool +symmetric3EigenvaluesDegenerate (firstEigenvalue, secondEigenvalue, thirdEigenvalue) = + min (abs (firstEigenvalue - secondEigenvalue)) (abs (secondEigenvalue - thirdEigenvalue)) <= 1.0e-10 + +symmetric3Eigenvector :: Symmetric3 Double -> Double -> Maybe Vec3 +symmetric3Eigenvector tensorValue eigenvalue = + let rowX = Vec3 (sym3XX tensorValue - eigenvalue) (sym3XY tensorValue) (sym3XZ tensorValue) + rowY = Vec3 (sym3XY tensorValue) (sym3YY tensorValue - eigenvalue) (sym3YZ tensorValue) + rowZ = Vec3 (sym3XZ tensorValue) (sym3YZ tensorValue) (sym3ZZ tensorValue - eigenvalue) + candidateValue = + largestVec3 + (crossVec3Local rowX rowY) + (largestVec3 (crossVec3Local rowX rowZ) (crossVec3Local rowY rowZ)) + in normalizeVec3Maybe candidateValue + +symmetric3CandidateAcceptable :: + Symmetric3 Double -> + (EigenColumn3, EigenColumn3, EigenColumn3) -> + Bool +symmetric3CandidateAcceptable tensorValue (firstColumn, secondColumn, thirdColumn) = + let firstVector = eigenColumn3Vector firstColumn + secondVector = eigenColumn3Vector secondColumn + thirdVector = eigenColumn3Vector thirdColumn + residualBound = + max + (symmetric3ResidualNorm tensorValue firstColumn) + (max (symmetric3ResidualNorm tensorValue secondColumn) (symmetric3ResidualNorm tensorValue thirdColumn)) + in residualBound <= 1.0e-8 + && abs (vec3Norm firstVector - 1.0) <= 1.0e-10 + && abs (vec3Norm secondVector - 1.0) <= 1.0e-10 + && abs (vec3Norm thirdVector - 1.0) <= 1.0e-10 + && abs (dotVec3Local firstVector secondVector) <= 1.0e-9 + && abs (dotVec3Local firstVector thirdVector) <= 1.0e-9 + && abs (dotVec3Local secondVector thirdVector) <= 1.0e-9 + +rayleighRefineColumn3 :: Symmetric3 Double -> EigenColumn3 -> EigenColumn3 +rayleighRefineColumn3 tensorValue eigenColumn = + let vectorValue = eigenColumn3Vector eigenColumn + refinedValue = + dotVec3Local (applySymmetric3 tensorValue vectorValue) vectorValue + / dotVec3Local vectorValue vectorValue + in EigenColumn3 refinedValue vectorValue + +symmetric3ResidualNorm :: Symmetric3 Double -> EigenColumn3 -> Double +symmetric3ResidualNorm tensorValue eigenColumn = + vec3Norm + ( subVec3Local + (applySymmetric3 tensorValue (eigenColumn3Vector eigenColumn)) + (scaleVec3Local (eigenColumn3Value eigenColumn) (eigenColumn3Vector eigenColumn)) + ) + +symmetric3JacobiEigenColumns :: + Symmetric3 Double -> + (EigenColumn3, EigenColumn3, EigenColumn3) +symmetric3JacobiEigenColumns tensorValue = + let (diagonalizedTensor, firstVector, secondVector, thirdVector) = + symmetric3Jacobi 72 tensorValue (Vec3 1.0 0.0 0.0) (Vec3 0.0 1.0 0.0) (Vec3 0.0 0.0 1.0) + in sortEigenColumns3 + (EigenColumn3 (sym3XX diagonalizedTensor) (canonicalizeVec3Sign firstVector)) + (EigenColumn3 (sym3YY diagonalizedTensor) (canonicalizeVec3Sign secondVector)) + (EigenColumn3 (sym3ZZ diagonalizedTensor) (canonicalizeVec3Sign thirdVector)) + +symmetric3Jacobi :: + Int -> + Symmetric3 Double -> + Vec3 -> + Vec3 -> + Vec3 -> + (Symmetric3 Double, Vec3, Vec3, Vec3) +symmetric3Jacobi remainingSteps tensorValue firstVector secondVector thirdVector = + if remainingSteps <= 0 || symmetric3OffDiagonalMax tensorValue <= 1.0e-15 + then (tensorValue, firstVector, secondVector, thirdVector) + else + let pivotValue = symmetric3LargestPivot tensorValue + (cosineValue, sineValue, tangentValue) = symmetric3JacobiRotation tensorValue pivotValue + (nextTensor, nextFirstVector, nextSecondVector, nextThirdVector) = + symmetric3ApplyJacobi pivotValue cosineValue sineValue tangentValue tensorValue firstVector secondVector thirdVector + in symmetric3Jacobi (remainingSteps - 1) nextTensor nextFirstVector nextSecondVector nextThirdVector + +symmetric3LargestPivot :: Symmetric3 Double -> Symmetric3Pivot +symmetric3LargestPivot tensorValue = + let xyMagnitude = abs (sym3XY tensorValue) + xzMagnitude = abs (sym3XZ tensorValue) + yzMagnitude = abs (sym3YZ tensorValue) + in if xyMagnitude >= xzMagnitude && xyMagnitude >= yzMagnitude + then PivotXY + else + if xzMagnitude >= yzMagnitude + then PivotXZ + else PivotYZ + +symmetric3JacobiRotation :: + Symmetric3 Double -> + Symmetric3Pivot -> + (Double, Double, Double) +symmetric3JacobiRotation tensorValue pivotValue = + let (leftDiagonal, rightDiagonal, offDiagonal) = + case pivotValue of + PivotXY -> (sym3XX tensorValue, sym3YY tensorValue, sym3XY tensorValue) + PivotXZ -> (sym3XX tensorValue, sym3ZZ tensorValue, sym3XZ tensorValue) + PivotYZ -> (sym3YY tensorValue, sym3ZZ tensorValue, sym3YZ tensorValue) + in if abs offDiagonal <= 1.0e-30 + then (1.0, 0.0, 0.0) + else + let tauValue = (rightDiagonal - leftDiagonal) / (2.0 * offDiagonal) + signValue = + if tauValue < 0.0 + then -1.0 + else 1.0 + tangentValue = signValue / (abs tauValue + sqrt (1.0 + tauValue * tauValue)) + cosineValue = 1.0 / sqrt (1.0 + tangentValue * tangentValue) + sineValue = tangentValue * cosineValue + in (cosineValue, sineValue, tangentValue) + +symmetric3ApplyJacobi :: + Symmetric3Pivot -> + Double -> + Double -> + Double -> + Symmetric3 Double -> + Vec3 -> + Vec3 -> + Vec3 -> + (Symmetric3 Double, Vec3, Vec3, Vec3) +symmetric3ApplyJacobi pivotValue cosineValue sineValue tangentValue tensorValue firstVector secondVector thirdVector = + case pivotValue of + PivotXY -> + ( Symmetric3 + { sym3XX = sym3XX tensorValue - tangentValue * sym3XY tensorValue, + sym3XY = 0.0, + sym3XZ = cosineValue * sym3XZ tensorValue - sineValue * sym3YZ tensorValue, + sym3YY = sym3YY tensorValue + tangentValue * sym3XY tensorValue, + sym3YZ = sineValue * sym3XZ tensorValue + cosineValue * sym3YZ tensorValue, + sym3ZZ = sym3ZZ tensorValue + }, + combineRotatedVec3 cosineValue (-sineValue) firstVector secondVector, + combineRotatedVec3 sineValue cosineValue firstVector secondVector, + thirdVector + ) + PivotXZ -> + ( Symmetric3 + { sym3XX = sym3XX tensorValue - tangentValue * sym3XZ tensorValue, + sym3XY = cosineValue * sym3XY tensorValue - sineValue * sym3YZ tensorValue, + sym3XZ = 0.0, + sym3YY = sym3YY tensorValue, + sym3YZ = sineValue * sym3XY tensorValue + cosineValue * sym3YZ tensorValue, + sym3ZZ = sym3ZZ tensorValue + tangentValue * sym3XZ tensorValue + }, + combineRotatedVec3 cosineValue (-sineValue) firstVector thirdVector, + secondVector, + combineRotatedVec3 sineValue cosineValue firstVector thirdVector + ) + PivotYZ -> + ( Symmetric3 + { sym3XX = sym3XX tensorValue, + sym3XY = cosineValue * sym3XY tensorValue - sineValue * sym3XZ tensorValue, + sym3XZ = sineValue * sym3XY tensorValue + cosineValue * sym3XZ tensorValue, + sym3YY = sym3YY tensorValue - tangentValue * sym3YZ tensorValue, + sym3YZ = 0.0, + sym3ZZ = sym3ZZ tensorValue + tangentValue * sym3YZ tensorValue + }, + firstVector, + combineRotatedVec3 cosineValue (-sineValue) secondVector thirdVector, + combineRotatedVec3 sineValue cosineValue secondVector thirdVector + ) + +combineRotatedVec3 :: Double -> Double -> Vec3 -> Vec3 -> Vec3 +combineRotatedVec3 leftScale rightScale leftVector rightVector = + addVec3Local (scaleVec3Local leftScale leftVector) (scaleVec3Local rightScale rightVector) + +symmetric3OffDiagonalMax :: Symmetric3 Double -> Double +symmetric3OffDiagonalMax tensorValue = + max (abs (sym3XY tensorValue)) (max (abs (sym3XZ tensorValue)) (abs (sym3YZ tensorValue))) + +symmetric3Result :: + Double -> + (EigenColumn3, EigenColumn3, EigenColumn3) -> + Either MoonlightError (Vector 3 Double, Matrix 3 3 Double) +symmetric3Result scaleValue (firstColumn, secondColumn, thirdColumn) = do + eigenvalues <- + fromListVector + @3 + [ scaleValue * eigenColumn3Value firstColumn, + scaleValue * eigenColumn3Value secondColumn, + scaleValue * eigenColumn3Value thirdColumn + ] + eigenvectors <- + fromListMatrix + @3 + @3 + ( matrix3Columns + (eigenColumn3Vector firstColumn) + (eigenColumn3Vector secondColumn) + (eigenColumn3Vector thirdColumn) + ) + pure (eigenvalues, eigenvectors) + +matrix3Columns :: Vec3 -> Vec3 -> Vec3 -> [Double] +matrix3Columns (Vec3 x1 y1 z1) (Vec3 x2 y2 z2) (Vec3 x3 y3 z3) = + [x1, x2, x3, y1, y2, y3, z1, z2, z3] + +sortEigenvalues3 :: Double -> Double -> Double -> (Double, Double, Double) +sortEigenvalues3 firstValue secondValue thirdValue = + let (largestValue, smallerValue) = orderEigenvaluePair firstValue secondValue + (middleCandidate, smallestValue) = orderEigenvaluePair smallerValue thirdValue + (finalLargest, finalMiddle) = orderEigenvaluePair largestValue middleCandidate + in (finalLargest, finalMiddle, smallestValue) + +orderEigenvaluePair :: Double -> Double -> (Double, Double) +orderEigenvaluePair leftValue rightValue = + if leftValue >= rightValue + then (leftValue, rightValue) + else (rightValue, leftValue) + +sortEigenColumns3 :: + EigenColumn3 -> + EigenColumn3 -> + EigenColumn3 -> + (EigenColumn3, EigenColumn3, EigenColumn3) +sortEigenColumns3 firstColumn secondColumn thirdColumn = + let (largestColumn, smallerColumn) = orderEigenColumnPair firstColumn secondColumn + (middleCandidate, smallestColumn) = orderEigenColumnPair smallerColumn thirdColumn + (finalLargest, finalMiddle) = orderEigenColumnPair largestColumn middleCandidate + in (finalLargest, finalMiddle, smallestColumn) + +orderEigenColumnPair :: EigenColumn3 -> EigenColumn3 -> (EigenColumn3, EigenColumn3) +orderEigenColumnPair leftColumn rightColumn = + if eigenColumn3Value leftColumn >= eigenColumn3Value rightColumn + then (leftColumn, rightColumn) + else (rightColumn, leftColumn) + +symmetric3Determinant :: Symmetric3 Double -> Double +symmetric3Determinant tensorValue = + sym3XX tensorValue * sym3YY tensorValue * sym3ZZ tensorValue + + 2.0 * sym3XY tensorValue * sym3XZ tensorValue * sym3YZ tensorValue + - sym3XX tensorValue * sym3YZ tensorValue * sym3YZ tensorValue + - sym3YY tensorValue * sym3XZ tensorValue * sym3XZ tensorValue + - sym3ZZ tensorValue * sym3XY tensorValue * sym3XY tensorValue + +orthogonalizeVec3Against :: Vec3 -> Vec3 -> Maybe Vec3 +orthogonalizeVec3Against axisValue vectorValue = + normalizeVec3Maybe (subVec3Local vectorValue (scaleVec3Local (dotVec3Local axisValue vectorValue) axisValue)) + +largestVec3 :: Vec3 -> Vec3 -> Vec3 +largestVec3 leftValue rightValue = + if vec3NormSquared leftValue >= vec3NormSquared rightValue + then leftValue + else rightValue + +normalizeVec3Maybe :: Vec3 -> Maybe Vec3 +normalizeVec3Maybe vectorValue = + let normValue = vec3Norm vectorValue + in if normValue <= 1.0e-12 + then Nothing + else Just (scaleVec3Local (1.0 / normValue) vectorValue) + +vec3Norm :: Vec3 -> Double +vec3Norm vectorValue = + sqrt (vec3NormSquared vectorValue) + +vec3NormSquared :: Vec3 -> Double +vec3NormSquared (Vec3 xValue yValue zValue) = + xValue * xValue + yValue * yValue + zValue * zValue + +dotVec3Local :: Vec3 -> Vec3 -> Double +dotVec3Local (Vec3 x1 y1 z1) (Vec3 x2 y2 z2) = + x1 * x2 + y1 * y2 + z1 * z2 + +crossVec3Local :: Vec3 -> Vec3 -> Vec3 +crossVec3Local (Vec3 x1 y1 z1) (Vec3 x2 y2 z2) = + Vec3 + (y1 * z2 - z1 * y2) + (z1 * x2 - x1 * z2) + (x1 * y2 - y1 * x2) + +addVec3Local :: Vec3 -> Vec3 -> Vec3 +addVec3Local (Vec3 x1 y1 z1) (Vec3 x2 y2 z2) = + Vec3 (x1 + x2) (y1 + y2) (z1 + z2) + +subVec3Local :: Vec3 -> Vec3 -> Vec3 +subVec3Local (Vec3 x1 y1 z1) (Vec3 x2 y2 z2) = + Vec3 (x1 - x2) (y1 - y2) (z1 - z2) + +scaleVec3Local :: Double -> Vec3 -> Vec3 +scaleVec3Local scaleValue (Vec3 xValue yValue zValue) = + Vec3 (scaleValue * xValue) (scaleValue * yValue) (scaleValue * zValue) + +canonicalizeVec3Sign :: Vec3 -> Vec3 +canonicalizeVec3Sign vectorValue@(Vec3 xValue yValue zValue) + | abs xValue > 1.0e-14 = + if xValue < 0.0 + then scaleVec3Local (-1.0) vectorValue + else vectorValue + | abs yValue > 1.0e-14 = + if yValue < 0.0 + then scaleVec3Local (-1.0) vectorValue + else vectorValue + | zValue < 0.0 = scaleVec3Local (-1.0) vectorValue + | otherwise = vectorValue diff --git a/src-geometry/Moonlight/LinAlg/Pure/Geometry/Transform/Affine.hs b/src-geometry/Moonlight/LinAlg/Pure/Geometry/Transform/Affine.hs new file mode 100644 --- /dev/null +++ b/src-geometry/Moonlight/LinAlg/Pure/Geometry/Transform/Affine.hs @@ -0,0 +1,70 @@ +module Moonlight.LinAlg.Pure.Geometry.Transform.Affine + ( AffineTransform (..), + TransformMetricEffect (..), + affineMaxScale, + isIdentityScale, + affineMetricEffect, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Kind (Type) +import Moonlight.LinAlg.Pure.Geometry.Frame (OrthonormalFrame, orthonormalFrameColumns) +import Moonlight.LinAlg.Pure.Geometry.Vec3 (Vec3 (..), maxAbsComponentVec3, vec3ToTuple) +import Prelude (Bool, Double, Eq, Ord, Read, Show, abs, compare, otherwise, seq, (.), (==), (>), (&&)) + +type AffineTransform :: Type +data AffineTransform = AffineTransform + { atTranslation :: {-# UNPACK #-} !Vec3, + atRotationFrame :: {-# UNPACK #-} !OrthonormalFrame, + atScale :: {-# UNPACK #-} !Vec3 + } + deriving stock (Eq, Show) + +instance NFData AffineTransform where + rnf affineTransform = affineTransform `seq` () + +type TransformMetricEffect :: Type +data TransformMetricEffect + = MetricIsometry + | UniformMetricScale !Double + | AnisotropicMetricDistortion + deriving stock (Eq, Ord, Show, Read) + +instance Ord AffineTransform where + compare leftTransform rightTransform = + compare + ( vec3ToTuple (atTranslation leftTransform), + frameTuple (atRotationFrame leftTransform), + vec3ToTuple (atScale leftTransform) + ) + ( vec3ToTuple (atTranslation rightTransform), + frameTuple (atRotationFrame rightTransform), + vec3ToTuple (atScale rightTransform) + ) + +affineMaxScale :: AffineTransform -> Double +affineMaxScale = maxAbsComponentVec3 . atScale + +isIdentityScale :: Vec3 -> Bool +isIdentityScale scaleVector = scaleVector == Vec3 1.0 1.0 1.0 + +affineMetricEffect :: AffineTransform -> TransformMetricEffect +affineMetricEffect affine = + case atScale affine of + Vec3 sx sy sz + | ax == 1.0 && ay == 1.0 && az == 1.0 -> MetricIsometry + | ax == ay && ay == az && ax > 0.0 -> UniformMetricScale ax + | otherwise -> AnisotropicMetricDistortion + where + ax = abs sx + ay = abs sy + az = abs sz + +frameTuple :: + OrthonormalFrame -> + ((Double, Double, Double), (Double, Double, Double), (Double, Double, Double)) +frameTuple frameValue = + case orthonormalFrameColumns frameValue of + (axis1, axis2, axis3) -> + (vec3ToTuple axis1, vec3ToTuple axis2, vec3ToTuple axis3) diff --git a/src-geometry/Moonlight/LinAlg/Pure/Geometry/Vec2.hs b/src-geometry/Moonlight/LinAlg/Pure/Geometry/Vec2.hs new file mode 100644 --- /dev/null +++ b/src-geometry/Moonlight/LinAlg/Pure/Geometry/Vec2.hs @@ -0,0 +1,289 @@ +module Moonlight.LinAlg.Pure.Geometry.Vec2 + ( Vec2 (..), + Axis2 (..), + addVec2, + subVec2, + scaleVec2, + negateVec2, + zipVec2, + mapVec2, + mulVec2, + minVec2, + maxVec2, + vec2Zero, + dotVec2, + magnitudeVec2, + normalizeVec2, + normalizeVec2Safe, + normalizeVec2Or, + vec2FromList, + vec2ToList, + vec2ToTuple, + maxAbsComponentVec2, + axis2Component, + axis2Vector, + distanceVec2, + averageVec2, + rejectVec2, + areaLikeVec2, + perpVec2, + crossZVec2, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Kind (Type) +import Data.Vector.Generic qualified as G +import Data.Vector.Generic.Mutable qualified as M +import Data.Vector.Unboxed qualified as U +import Foreign.Storable (Storable (..), peekByteOff, pokeByteOff) +import Moonlight.Algebra.Pure.Module (BilinearSpace (..), Module (..), VectorSpace) +import Moonlight.Core (AdditiveGroup (..), AdditiveMonoid (..), Metric (..), MoonlightError (..)) +import Prelude (Bounded, Double, Either (..), Enum, Eq, Int, Ord, Read, Show, abs, length, max, min, seq, show, sqrt, (*), (+), (-), (/), (<$>), (<*>), (*>), (<>), (<=)) + +type Axis2 :: Type +data Axis2 + = Axis2X + | Axis2Y + deriving stock (Eq, Ord, Show, Read, Enum, Bounded) + +type Vec2 :: Type +data Vec2 = Vec2 + { vec2X :: {-# UNPACK #-} !Double, + vec2Y :: {-# UNPACK #-} !Double + } + deriving stock (Eq, Ord, Show, Read) + +instance NFData Vec2 where + rnf vectorValue = vectorValue `seq` () + +instance Storable Vec2 where + sizeOf _ = 2 * doubleStorageSize + alignment _ = alignment (0.0 :: Double) + peek pointerValue = + Vec2 + <$> peekByteOff pointerValue 0 + <*> peekByteOff pointerValue doubleStorageSize + poke pointerValue (Vec2 xValue yValue) = + pokeByteOff pointerValue 0 xValue + *> pokeByteOff pointerValue doubleStorageSize yValue + +newtype instance U.MVector s Vec2 = MV_Vec2 (U.MVector s (Double, Double)) + +newtype instance U.Vector Vec2 = V_Vec2 (U.Vector (Double, Double)) + +instance U.Unbox Vec2 + +instance M.MVector U.MVector Vec2 where + {-# INLINE basicLength #-} + basicLength (MV_Vec2 vectorValue) = M.basicLength vectorValue + + {-# INLINE basicUnsafeSlice #-} + basicUnsafeSlice offset lengthValue (MV_Vec2 vectorValue) = + MV_Vec2 (M.basicUnsafeSlice offset lengthValue vectorValue) + + {-# INLINE basicOverlaps #-} + basicOverlaps (MV_Vec2 leftVector) (MV_Vec2 rightVector) = + M.basicOverlaps leftVector rightVector + + {-# INLINE basicUnsafeNew #-} + basicUnsafeNew lengthValue = + MV_Vec2 <$> M.basicUnsafeNew lengthValue + + {-# INLINE basicInitialize #-} + basicInitialize (MV_Vec2 vectorValue) = + M.basicInitialize vectorValue + + {-# INLINE basicUnsafeReplicate #-} + basicUnsafeReplicate lengthValue (Vec2 xValue yValue) = + MV_Vec2 <$> M.basicUnsafeReplicate lengthValue (xValue, yValue) + + {-# INLINE basicUnsafeRead #-} + basicUnsafeRead (MV_Vec2 vectorValue) indexValue = + tupleToVec2 <$> M.basicUnsafeRead vectorValue indexValue + + {-# INLINE basicUnsafeWrite #-} + basicUnsafeWrite (MV_Vec2 vectorValue) indexValue vectorPayload = + M.basicUnsafeWrite vectorValue indexValue (vec2ToTuple vectorPayload) + + {-# INLINE basicClear #-} + basicClear (MV_Vec2 vectorValue) = + M.basicClear vectorValue + + {-# INLINE basicSet #-} + basicSet (MV_Vec2 vectorValue) vectorPayload = + M.basicSet vectorValue (vec2ToTuple vectorPayload) + + {-# INLINE basicUnsafeCopy #-} + basicUnsafeCopy (MV_Vec2 targetVector) (MV_Vec2 sourceVector) = + M.basicUnsafeCopy targetVector sourceVector + + {-# INLINE basicUnsafeMove #-} + basicUnsafeMove (MV_Vec2 targetVector) (MV_Vec2 sourceVector) = + M.basicUnsafeMove targetVector sourceVector + + {-# INLINE basicUnsafeGrow #-} + basicUnsafeGrow (MV_Vec2 vectorValue) lengthValue = + MV_Vec2 <$> M.basicUnsafeGrow vectorValue lengthValue + +instance G.Vector U.Vector Vec2 where + {-# INLINE basicUnsafeFreeze #-} + basicUnsafeFreeze (MV_Vec2 vectorValue) = + V_Vec2 <$> G.basicUnsafeFreeze vectorValue + + {-# INLINE basicUnsafeThaw #-} + basicUnsafeThaw (V_Vec2 vectorValue) = + MV_Vec2 <$> G.basicUnsafeThaw vectorValue + + {-# INLINE basicLength #-} + basicLength (V_Vec2 vectorValue) = G.basicLength vectorValue + + {-# INLINE basicUnsafeSlice #-} + basicUnsafeSlice offset lengthValue (V_Vec2 vectorValue) = + V_Vec2 (G.basicUnsafeSlice offset lengthValue vectorValue) + + {-# INLINE basicUnsafeIndexM #-} + basicUnsafeIndexM (V_Vec2 vectorValue) indexValue = + tupleToVec2 <$> G.basicUnsafeIndexM vectorValue indexValue + + {-# INLINE basicUnsafeCopy #-} + basicUnsafeCopy (MV_Vec2 targetVector) (V_Vec2 sourceVector) = + G.basicUnsafeCopy targetVector sourceVector + + {-# INLINE elemseq #-} + elemseq _ (Vec2 xValue yValue) resultValue = + G.elemseq doubleVectorWitness xValue (G.elemseq doubleVectorWitness yValue resultValue) + +instance AdditiveMonoid Vec2 where + zero = Vec2 0.0 0.0 + add (Vec2 x1 y1) (Vec2 x2 y2) = Vec2 (x1 + x2) (y1 + y2) + +instance AdditiveGroup Vec2 where + neg (Vec2 xValue yValue) = Vec2 (-xValue) (-yValue) + +instance Module Double Vec2 where + scale scaleValue (Vec2 xValue yValue) = Vec2 (scaleValue * xValue) (scaleValue * yValue) + +instance VectorSpace Double Vec2 + +instance BilinearSpace Double Vec2 where + bilinearForm (Vec2 x1 y1) (Vec2 x2 y2) = x1 * x2 + y1 * y2 + +instance Metric Vec2 where + type Magnitude Vec2 = Double + magnitude vectorValue = sqrt (bilinearForm vectorValue vectorValue) + +vec2Zero :: Vec2 +vec2Zero = zero + +addVec2 :: Vec2 -> Vec2 -> Vec2 +addVec2 = add + +subVec2 :: Vec2 -> Vec2 -> Vec2 +subVec2 = sub + +negateVec2 :: Vec2 -> Vec2 +negateVec2 = neg + +scaleVec2 :: Double -> Vec2 -> Vec2 +scaleVec2 = scale + +zipVec2 :: (Double -> Double -> Double) -> Vec2 -> Vec2 -> Vec2 +zipVec2 combine (Vec2 x1 y1) (Vec2 x2 y2) = + Vec2 (combine x1 x2) (combine y1 y2) + +mapVec2 :: (Double -> Double) -> Vec2 -> Vec2 +mapVec2 transform (Vec2 xValue yValue) = + Vec2 (transform xValue) (transform yValue) + +mulVec2 :: Vec2 -> Vec2 -> Vec2 +mulVec2 = zipVec2 (*) + +minVec2 :: Vec2 -> Vec2 -> Vec2 +minVec2 = zipVec2 min + +maxVec2 :: Vec2 -> Vec2 -> Vec2 +maxVec2 = zipVec2 max + +dotVec2 :: Vec2 -> Vec2 -> Double +dotVec2 = bilinearForm + +magnitudeVec2 :: Vec2 -> Double +magnitudeVec2 = magnitude + +normalizeVec2 :: Vec2 -> Either MoonlightError Vec2 +normalizeVec2 vectorValue = + let vectorMagnitude = magnitudeVec2 vectorValue + in if vectorMagnitude <= 1.0e-12 + then Left (InvariantViolation "member direction cannot be normalized from a zero-length vector") + else Right (scaleVec2 (1.0 / vectorMagnitude) vectorValue) + +normalizeVec2Safe :: Vec2 -> Vec2 +normalizeVec2Safe vectorValue = + let vectorMagnitude = magnitudeVec2 vectorValue + in if vectorMagnitude <= 1.0e-12 then vec2Zero else scaleVec2 (1.0 / vectorMagnitude) vectorValue + +normalizeVec2Or :: Vec2 -> Vec2 -> Vec2 +normalizeVec2Or fallback vectorValue = + let vectorMagnitude = magnitudeVec2 vectorValue + in if vectorMagnitude <= 1.0e-12 then fallback else scaleVec2 (1.0 / vectorMagnitude) vectorValue + +vec2FromList :: [Double] -> Either MoonlightError Vec2 +vec2FromList values = + case values of + [xValue, yValue] -> Right (Vec2 xValue yValue) + _ -> + Left + ( InvariantViolation + ( "Vec2 requires exactly 2 entries, received " + <> show (length values) + ) + ) + +vec2ToList :: Vec2 -> [Double] +vec2ToList (Vec2 xValue yValue) = [xValue, yValue] + +vec2ToTuple :: Vec2 -> (Double, Double) +vec2ToTuple (Vec2 xValue yValue) = (xValue, yValue) + +tupleToVec2 :: (Double, Double) -> Vec2 +tupleToVec2 (xValue, yValue) = Vec2 xValue yValue + +doubleStorageSize :: Int +doubleStorageSize = sizeOf (0.0 :: Double) + +doubleVectorWitness :: U.Vector Double +doubleVectorWitness = U.empty + +maxAbsComponentVec2 :: Vec2 -> Double +maxAbsComponentVec2 (Vec2 xValue yValue) = + max (abs xValue) (abs yValue) + +axis2Component :: Axis2 -> Vec2 -> Double +axis2Component Axis2X (Vec2 xValue _) = xValue +axis2Component Axis2Y (Vec2 _ yValue) = yValue + +axis2Vector :: Axis2 -> Double -> Vec2 +axis2Vector Axis2X magnitudeValue = Vec2 magnitudeValue 0.0 +axis2Vector Axis2Y magnitudeValue = Vec2 0.0 magnitudeValue + +distanceVec2 :: Vec2 -> Vec2 -> Double +distanceVec2 leftPosition rightPosition = + magnitudeVec2 (subVec2 leftPosition rightPosition) + +averageVec2 :: Vec2 -> Vec2 -> Vec2 +averageVec2 leftValue rightValue = + scaleVec2 0.5 (addVec2 leftValue rightValue) + +rejectVec2 :: Vec2 -> Vec2 -> Vec2 +rejectVec2 onto vectorValue = + subVec2 vectorValue (scaleVec2 (dotVec2 onto vectorValue) onto) + +areaLikeVec2 :: Vec2 -> Double +areaLikeVec2 (Vec2 xValue yValue) = abs (xValue * yValue) + +perpVec2 :: Vec2 -> Vec2 +perpVec2 (Vec2 xValue yValue) = Vec2 (-yValue) xValue + +crossZVec2 :: Vec2 -> Vec2 -> Double +crossZVec2 (Vec2 x1 y1) (Vec2 x2 y2) = x1 * y2 - y1 * x2 diff --git a/src-geometry/Moonlight/LinAlg/Pure/Geometry/Vec3.hs b/src-geometry/Moonlight/LinAlg/Pure/Geometry/Vec3.hs new file mode 100644 --- /dev/null +++ b/src-geometry/Moonlight/LinAlg/Pure/Geometry/Vec3.hs @@ -0,0 +1,289 @@ +module Moonlight.LinAlg.Pure.Geometry.Vec3 + ( Vec3 (..), + Axis (..), + addVec3, + subVec3, + scaleVec3, + negateVec3, + zipVec3, + mapVec3, + mulVec3, + minVec3, + maxVec3, + vec3Zero, + dotVec3, + magnitudeVec3, + normalizeVec3, + normalizeVec3Safe, + crossVec3, + vec3FromList, + vec3ToList, + vec3ToTuple, + maxAbsComponentVec3, + axisComponent, + axisVector, + distanceVec3, + averageVec3, + rejectVec3, + volumeLikeVec3, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Kind (Type) +import Data.Vector.Generic qualified as G +import Data.Vector.Generic.Mutable qualified as M +import Data.Vector.Unboxed qualified as U +import Foreign.Storable (Storable (..), peekByteOff, pokeByteOff) +import Moonlight.Algebra.Pure.Module (BilinearSpace (..), Module (..), VectorSpace) +import Moonlight.Core (AdditiveGroup (..), AdditiveMonoid (..), Metric (..), MoonlightError (..)) +import Prelude (Bounded, Double, Either (..), Enum, Eq, Int, Ord, Read, Show, abs, length, max, min, seq, show, sqrt, (*), (+), (-), (/), (<$>), (<*>), (*>), (<>), (<=)) + +type Axis :: Type +data Axis + = AxisX + | AxisY + | AxisZ + deriving stock (Eq, Ord, Show, Read, Enum, Bounded) + +type Vec3 :: Type +data Vec3 = Vec3 + { vecX :: {-# UNPACK #-} !Double, + vecY :: {-# UNPACK #-} !Double, + vecZ :: {-# UNPACK #-} !Double + } + deriving stock (Eq, Ord, Show, Read) + +instance NFData Vec3 where + rnf vectorValue = vectorValue `seq` () + +instance Storable Vec3 where + sizeOf _ = 3 * doubleStorageSize + alignment _ = alignment (0.0 :: Double) + peek pointerValue = + Vec3 + <$> peekByteOff pointerValue 0 + <*> peekByteOff pointerValue doubleStorageSize + <*> peekByteOff pointerValue (2 * doubleStorageSize) + poke pointerValue (Vec3 xValue yValue zValue) = + pokeByteOff pointerValue 0 xValue + *> pokeByteOff pointerValue doubleStorageSize yValue + *> pokeByteOff pointerValue (2 * doubleStorageSize) zValue + +newtype instance U.MVector s Vec3 = MV_Vec3 (U.MVector s (Double, Double, Double)) + +newtype instance U.Vector Vec3 = V_Vec3 (U.Vector (Double, Double, Double)) + +instance U.Unbox Vec3 + +instance M.MVector U.MVector Vec3 where + {-# INLINE basicLength #-} + basicLength (MV_Vec3 vectorValue) = M.basicLength vectorValue + + {-# INLINE basicUnsafeSlice #-} + basicUnsafeSlice offset lengthValue (MV_Vec3 vectorValue) = + MV_Vec3 (M.basicUnsafeSlice offset lengthValue vectorValue) + + {-# INLINE basicOverlaps #-} + basicOverlaps (MV_Vec3 leftVector) (MV_Vec3 rightVector) = + M.basicOverlaps leftVector rightVector + + {-# INLINE basicUnsafeNew #-} + basicUnsafeNew lengthValue = + MV_Vec3 <$> M.basicUnsafeNew lengthValue + + {-# INLINE basicInitialize #-} + basicInitialize (MV_Vec3 vectorValue) = + M.basicInitialize vectorValue + + {-# INLINE basicUnsafeReplicate #-} + basicUnsafeReplicate lengthValue (Vec3 xValue yValue zValue) = + MV_Vec3 <$> M.basicUnsafeReplicate lengthValue (xValue, yValue, zValue) + + {-# INLINE basicUnsafeRead #-} + basicUnsafeRead (MV_Vec3 vectorValue) indexValue = + tupleToVec3 <$> M.basicUnsafeRead vectorValue indexValue + + {-# INLINE basicUnsafeWrite #-} + basicUnsafeWrite (MV_Vec3 vectorValue) indexValue vectorPayload = + M.basicUnsafeWrite vectorValue indexValue (vec3ToTuple vectorPayload) + + {-# INLINE basicClear #-} + basicClear (MV_Vec3 vectorValue) = + M.basicClear vectorValue + + {-# INLINE basicSet #-} + basicSet (MV_Vec3 vectorValue) vectorPayload = + M.basicSet vectorValue (vec3ToTuple vectorPayload) + + {-# INLINE basicUnsafeCopy #-} + basicUnsafeCopy (MV_Vec3 targetVector) (MV_Vec3 sourceVector) = + M.basicUnsafeCopy targetVector sourceVector + + {-# INLINE basicUnsafeMove #-} + basicUnsafeMove (MV_Vec3 targetVector) (MV_Vec3 sourceVector) = + M.basicUnsafeMove targetVector sourceVector + + {-# INLINE basicUnsafeGrow #-} + basicUnsafeGrow (MV_Vec3 vectorValue) lengthValue = + MV_Vec3 <$> M.basicUnsafeGrow vectorValue lengthValue + +instance G.Vector U.Vector Vec3 where + {-# INLINE basicUnsafeFreeze #-} + basicUnsafeFreeze (MV_Vec3 vectorValue) = + V_Vec3 <$> G.basicUnsafeFreeze vectorValue + + {-# INLINE basicUnsafeThaw #-} + basicUnsafeThaw (V_Vec3 vectorValue) = + MV_Vec3 <$> G.basicUnsafeThaw vectorValue + + {-# INLINE basicLength #-} + basicLength (V_Vec3 vectorValue) = G.basicLength vectorValue + + {-# INLINE basicUnsafeSlice #-} + basicUnsafeSlice offset lengthValue (V_Vec3 vectorValue) = + V_Vec3 (G.basicUnsafeSlice offset lengthValue vectorValue) + + {-# INLINE basicUnsafeIndexM #-} + basicUnsafeIndexM (V_Vec3 vectorValue) indexValue = + tupleToVec3 <$> G.basicUnsafeIndexM vectorValue indexValue + + {-# INLINE basicUnsafeCopy #-} + basicUnsafeCopy (MV_Vec3 targetVector) (V_Vec3 sourceVector) = + G.basicUnsafeCopy targetVector sourceVector + + {-# INLINE elemseq #-} + elemseq _ (Vec3 xValue yValue zValue) resultValue = + G.elemseq doubleVectorWitness xValue (G.elemseq doubleVectorWitness yValue (G.elemseq doubleVectorWitness zValue resultValue)) + +instance AdditiveMonoid Vec3 where + zero = Vec3 0.0 0.0 0.0 + add (Vec3 x1 y1 z1) (Vec3 x2 y2 z2) = Vec3 (x1 + x2) (y1 + y2) (z1 + z2) + +instance AdditiveGroup Vec3 where + neg (Vec3 x y z) = Vec3 (-x) (-y) (-z) + +instance Module Double Vec3 where + scale s (Vec3 x y z) = Vec3 (s * x) (s * y) (s * z) + +instance VectorSpace Double Vec3 + +instance BilinearSpace Double Vec3 where + bilinearForm (Vec3 x1 y1 z1) (Vec3 x2 y2 z2) = x1 * x2 + y1 * y2 + z1 * z2 + +instance Metric Vec3 where + type Magnitude Vec3 = Double + magnitude v = sqrt (bilinearForm v v) + +vec3Zero :: Vec3 +vec3Zero = zero + +addVec3 :: Vec3 -> Vec3 -> Vec3 +addVec3 = add + +subVec3 :: Vec3 -> Vec3 -> Vec3 +subVec3 = sub + +negateVec3 :: Vec3 -> Vec3 +negateVec3 = neg + +scaleVec3 :: Double -> Vec3 -> Vec3 +scaleVec3 = scale + +zipVec3 :: (Double -> Double -> Double) -> Vec3 -> Vec3 -> Vec3 +zipVec3 combine (Vec3 x1 y1 z1) (Vec3 x2 y2 z2) = + Vec3 (combine x1 x2) (combine y1 y2) (combine z1 z2) + +mapVec3 :: (Double -> Double) -> Vec3 -> Vec3 +mapVec3 transform (Vec3 xValue yValue zValue) = + Vec3 (transform xValue) (transform yValue) (transform zValue) + +mulVec3 :: Vec3 -> Vec3 -> Vec3 +mulVec3 = zipVec3 (*) + +minVec3 :: Vec3 -> Vec3 -> Vec3 +minVec3 = zipVec3 min + +maxVec3 :: Vec3 -> Vec3 -> Vec3 +maxVec3 = zipVec3 max + +dotVec3 :: Vec3 -> Vec3 -> Double +dotVec3 = bilinearForm + +magnitudeVec3 :: Vec3 -> Double +magnitudeVec3 = magnitude + +normalizeVec3 :: Vec3 -> Either MoonlightError Vec3 +normalizeVec3 v = + let m = magnitudeVec3 v + in if m <= 1.0e-12 + then Left (InvariantViolation "member direction cannot be normalized from a zero-length vector") + else Right (scaleVec3 (1.0 / m) v) + +normalizeVec3Safe :: Vec3 -> Vec3 +normalizeVec3Safe v = + let m = magnitudeVec3 v + in if m <= 1.0e-12 then vec3Zero else scaleVec3 (1.0 / m) v + +crossVec3 :: Vec3 -> Vec3 -> Vec3 +crossVec3 (Vec3 x1 y1 z1) (Vec3 x2 y2 z2) = + Vec3 + (y1 * z2 - z1 * y2) + (z1 * x2 - x1 * z2) + (x1 * y2 - y1 * x2) + +vec3FromList :: [Double] -> Either MoonlightError Vec3 +vec3FromList xs = + case xs of + [x, y, z] -> Right (Vec3 x y z) + _ -> + Left + ( InvariantViolation + ( "Vec3 requires exactly 3 entries, received " + <> show (length xs) + ) + ) + +vec3ToList :: Vec3 -> [Double] +vec3ToList (Vec3 xValue yValue zValue) = [xValue, yValue, zValue] + +vec3ToTuple :: Vec3 -> (Double, Double, Double) +vec3ToTuple (Vec3 xValue yValue zValue) = (xValue, yValue, zValue) + +tupleToVec3 :: (Double, Double, Double) -> Vec3 +tupleToVec3 (xValue, yValue, zValue) = Vec3 xValue yValue zValue + +doubleStorageSize :: Int +doubleStorageSize = sizeOf (0.0 :: Double) + +doubleVectorWitness :: U.Vector Double +doubleVectorWitness = U.empty + +maxAbsComponentVec3 :: Vec3 -> Double +maxAbsComponentVec3 (Vec3 xValue yValue zValue) = + max (abs xValue) (max (abs yValue) (abs zValue)) + +axisComponent :: Axis -> Vec3 -> Double +axisComponent AxisX (Vec3 x _ _) = x +axisComponent AxisY (Vec3 _ y _) = y +axisComponent AxisZ (Vec3 _ _ z) = z + +axisVector :: Axis -> Double -> Vec3 +axisVector AxisX m = Vec3 m 0.0 0.0 +axisVector AxisY m = Vec3 0.0 m 0.0 +axisVector AxisZ m = Vec3 0.0 0.0 m + +distanceVec3 :: Vec3 -> Vec3 -> Double +distanceVec3 leftPosition rightPosition = + magnitudeVec3 (subVec3 leftPosition rightPosition) + +averageVec3 :: Vec3 -> Vec3 -> Vec3 +averageVec3 leftValue rightValue = + scaleVec3 0.5 (addVec3 leftValue rightValue) + +rejectVec3 :: Vec3 -> Vec3 -> Vec3 +rejectVec3 onto v = + subVec3 v (scaleVec3 (dotVec3 onto v) onto) + +volumeLikeVec3 :: Vec3 -> Double +volumeLikeVec3 (Vec3 x y z) = abs (x * y * z) diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness.hs b/src-laws/Moonlight/LinAlg/Effect/Harness.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness.hs @@ -0,0 +1,24 @@ +module Moonlight.LinAlg.Effect.Harness + ( module Dense, + module Decomposition, + module Domain, + module Field, + module Geometry, + module KrylovSpectral, + module Operator, + module Preconditioner, + module Sparse, + module Statics, + ) +where + +import Moonlight.LinAlg.Effect.Harness.Dense as Dense +import Moonlight.LinAlg.Effect.Harness.Decomposition as Decomposition +import Moonlight.LinAlg.Effect.Harness.Domain as Domain +import Moonlight.LinAlg.Effect.Harness.Field as Field +import Moonlight.LinAlg.Effect.Harness.Geometry as Geometry +import Moonlight.LinAlg.Effect.Harness.KrylovSpectral as KrylovSpectral +import Moonlight.LinAlg.Effect.Harness.Operator as Operator +import Moonlight.LinAlg.Effect.Harness.Preconditioner as Preconditioner +import Moonlight.LinAlg.Effect.Harness.Sparse as Sparse +import Moonlight.LinAlg.Effect.Harness.Statics as Statics diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness/Core.hs b/src-laws/Moonlight/LinAlg/Effect/Harness/Core.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness/Core.hs @@ -0,0 +1,107 @@ +module Moonlight.LinAlg.Effect.Harness.Core + ( approxTolerance, + orthonormalTolerance, + residualTolerance, + matrix3Product, + assertApproxEqual, + assertApproxEqualWith, + assertApproxList, + assertApproxListWith, + assertRightBool, + assertRightProperty, + exactRightProperty, + matrixRows3, + matrix3VectorProduct, + maxAbsDifference, + vectorDot, + vectorNorm, + ) +where + +import Test.Tasty.QuickCheck qualified as QC + +approxTolerance :: Double +approxTolerance = + 1.0e-8 + +residualTolerance :: Double +residualTolerance = + 1.0e-5 + +orthonormalTolerance :: Double +orthonormalTolerance = + 1.0e-6 + +assertApproxEqual :: Double -> Double -> Bool +assertApproxEqual expected actual = + assertApproxEqualWith approxTolerance expected actual + +assertApproxEqualWith :: Double -> Double -> Double -> Bool +assertApproxEqualWith tolerance expected actual = + abs (expected - actual) <= tolerance + +assertApproxList :: [Double] -> [Double] -> Bool +assertApproxList expected actual = + assertApproxListWith approxTolerance expected actual + +assertApproxListWith :: Double -> [Double] -> [Double] -> Bool +assertApproxListWith tolerance expected actual = + length expected == length actual + && and (zipWith (assertApproxEqualWith tolerance) expected actual) + +assertRightBool :: Either failure Bool -> Bool +assertRightBool = + either (const False) id + +exactRightProperty :: (Eq value, Show failure, Show value) => Either failure value -> Either failure value -> QC.Property +exactRightProperty left right = + case (left, right) of + (Right leftValue, Right rightValue) -> + QC.counterexample (show (leftValue, rightValue)) (leftValue == rightValue) + (Left leftFailure, _) -> + QC.counterexample (show leftFailure) False + (_, Left rightFailure) -> + QC.counterexample (show rightFailure) False + +assertRightProperty :: Show failure => Either failure Bool -> QC.Property +assertRightProperty result = + case result of + Left failure -> + QC.counterexample (show failure) False + Right accepted -> + QC.property accepted + +matrixRows3 :: [a] -> [[a]] +matrixRows3 values = + case values of + [a00, a01, a02, a10, a11, a12, a20, a21, a22] -> + [[a00, a01, a02], [a10, a11, a12], [a20, a21, a22]] + _ -> [] + +matrix3VectorProduct :: [[Double]] -> [Double] -> [Double] +matrix3VectorProduct rows vectorValue = + fmap (`vectorDot` vectorValue) rows + +matrix3Product :: [[Double]] -> [[Double]] -> [[Double]] +matrix3Product leftRows rightRows = + let rightColumns = transpose3 rightRows + in fmap (\leftRow -> fmap (vectorDot leftRow) rightColumns) leftRows + +vectorDot :: [Double] -> [Double] -> Double +vectorDot left right = + sum (zipWith (*) left right) + +vectorNorm :: [Double] -> Double +vectorNorm values = + sqrt (sum ((\entryValue -> entryValue * entryValue) <$> values)) + +maxAbsDifference :: [Double] -> [Double] -> Double +maxAbsDifference expected actual = + maximum (0.0 : zipWith (\leftValue rightValue -> abs (leftValue - rightValue)) expected actual) + +transpose3 :: [[a]] -> [[a]] +transpose3 rows = + case rows of + [[a00, a01, a02], [a10, a11, a12], [a20, a21, a22]] -> + [[a00, a10, a20], [a01, a11, a21], [a02, a12, a22]] + _ -> [] diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness/Decomposition.hs b/src-laws/Moonlight/LinAlg/Effect/Harness/Decomposition.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness/Decomposition.hs @@ -0,0 +1,294 @@ +module Moonlight.LinAlg.Effect.Harness.Decomposition + ( qrReconstructsInputLaw, + qrOrthonormalColumnsLaw, + choleskyReconstructsSpdLaw, + symmetricEigenReconstructsLaw, + symmetricEigenOrthonormalLaw, + symmetricEigenUncheckedPassesCertificationLaw, + thinSvdReconstructsLaw, + thinSvdOrthonormalFactorsLaw, + thinSvdSingularValuesOrderedNonnegativeLaw, + ) +where + +import Data.Bifunctor (first) +import Data.Vector.Storable qualified as S +import Moonlight.LinAlg + ( choleskyDecomp, + fromListMatrix, + mult, + qrDecompFullColumnRank, + symmetricEigen, + thinSvdFullColumnRank, + toListMatrix, + toListVector, + transpose, + ) +import Moonlight.LinAlg.Effect.Harness.Core + ( approxTolerance, + assertApproxList, + assertRightProperty, + maxAbsDifference, + ) +import Moonlight.LinAlg.Internal.Eigen.Symmetric + ( certifySymmetricEigenResult, + symmetricEigenPairsDenseUnchecked, + ) +import Moonlight.LinAlg.Pure.Dense.Flat (mkDenseDoubleMatrixRowMajor) +import Test.Tasty.QuickCheck qualified as QC + +newtype FullRankMatrix43 = FullRankMatrix43 [Double] + deriving stock (Eq, Show) + +newtype SpdMatrix3 = SpdMatrix3 [Double] + deriving stock (Eq, Show) + +newtype SymmetricMatrix3 = SymmetricMatrix3 [Double] + deriving stock (Eq, Show) + +newtype FullRankMatrix32 = FullRankMatrix32 [Double] + deriving stock (Eq, Show) + +instance QC.Arbitrary FullRankMatrix43 where + arbitrary = + FullRankMatrix43 <$> anchoredOrGenerated fullRankMatrix43Anchors generateFullRankMatrix43 + +instance QC.Arbitrary SpdMatrix3 where + arbitrary = + SpdMatrix3 <$> anchoredOrGenerated spdMatrix3Anchors generateSpdMatrix3 + +instance QC.Arbitrary SymmetricMatrix3 where + arbitrary = + SymmetricMatrix3 <$> anchoredOrGenerated symmetricMatrix3Anchors generateSymmetricMatrix3 + +instance QC.Arbitrary FullRankMatrix32 where + arbitrary = + FullRankMatrix32 <$> anchoredOrGenerated fullRankMatrix32Anchors generateFullRankMatrix32 + +anchoredOrGenerated :: [[Double]] -> QC.Gen [Double] -> QC.Gen [Double] +anchoredOrGenerated anchors generatedValues = + QC.frequency + [ (1, QC.elements anchors), + (9, generatedValues) + ] + +generateFullRankMatrix43 :: QC.Gen [Double] +generateFullRankMatrix43 = + fullRankMatrix43Entries + <$> generatedTriple generatedNonZeroEntry + <*> generatedTriple generatedEntry + <*> generatedTriple generatedEntry + +generateSpdMatrix3 :: QC.Gen [Double] +generateSpdMatrix3 = + spdMatrix3Entries + <$> generatedTriple (QC.choose (1.0, 3.0)) + <*> generatedTriple (QC.choose (-1.0, 1.0)) + +generateSymmetricMatrix3 :: QC.Gen [Double] +generateSymmetricMatrix3 = + symmetricMatrix3Entries + <$> generatedTriple generatedEntry + <*> generatedTriple generatedEntry + +generateFullRankMatrix32 :: QC.Gen [Double] +generateFullRankMatrix32 = + fullRankMatrix32Entries + <$> ((,) <$> generatedNonZeroEntry <*> generatedNonZeroEntry) + <*> generatedTriple generatedEntry + +generatedTriple :: QC.Gen value -> QC.Gen (value, value, value) +generatedTriple generatedValue = + (,,) <$> generatedValue <*> generatedValue <*> generatedValue + +fullRankMatrix43Entries :: (Double, Double, Double) -> (Double, Double, Double) -> (Double, Double, Double) -> [Double] +fullRankMatrix43Entries (d0, d1, d2) (l10, l20, l21) (r0, r1, r2) = + [d0, 0.0, 0.0, l10, d1, 0.0, l20, l21, d2, r0, r1, r2] + +spdMatrix3Entries :: (Double, Double, Double) -> (Double, Double, Double) -> [Double] +spdMatrix3Entries (d0, d1, d2) (l10, l20, l21) = + symmetricMatrix3Entries + (d0 * d0, l10 * l10 + d1 * d1, l20 * l20 + l21 * l21 + d2 * d2) + (d0 * l10, d0 * l20, l10 * l20 + d1 * l21) + +symmetricMatrix3Entries :: (Double, Double, Double) -> (Double, Double, Double) -> [Double] +symmetricMatrix3Entries (d0, d1, d2) (o01, o02, o12) = + [d0, o01, o02, o01, d1, o12, o02, o12, d2] + +fullRankMatrix32Entries :: (Double, Double) -> (Double, Double, Double) -> [Double] +fullRankMatrix32Entries (d0, d1) (l10, l20, l21) = + [d0, 0.0, l10, d1, l20, l21] + +generatedEntry :: QC.Gen Double +generatedEntry = + QC.choose (-4.0, 4.0) + +generatedNonZeroEntry :: QC.Gen Double +generatedNonZeroEntry = + QC.elements [-4.0, -3.0, -2.0, -1.0, 1.0, 2.0, 3.0, 4.0] + +fullRankMatrix43Anchors :: [[Double]] +fullRankMatrix43Anchors = + [ [1.0, 0.0, 2.0, 0.0, 1.0, -1.0, 2.0, 1.0, 0.0, 1.0, -1.0, 1.0], + [2.0, 1.0, 0.0, 1.0, 3.0, 1.0, 0.0, -1.0, 2.0, 1.0, 0.0, 1.0], + [1.0, 2.0, 1.0, 2.0, 0.0, -1.0, 0.0, 1.0, 3.0, 1.0, -1.0, 0.0] + ] + +spdMatrix3Anchors :: [[Double]] +spdMatrix3Anchors = + [ [6.0, 2.0, 1.0, 2.0, 5.0, 0.5, 1.0, 0.5, 4.0], + [5.0, -1.0, 0.5, -1.0, 4.0, 1.0, 0.5, 1.0, 3.5], + [9.0, 1.5, -0.5, 1.5, 7.0, 2.0, -0.5, 2.0, 6.0] + ] + +symmetricMatrix3Anchors :: [[Double]] +symmetricMatrix3Anchors = + [ [4.0, 1.0, 2.0, 1.0, 3.0, 0.5, 2.0, 0.5, 5.0], + [2.0, 0.0, 0.0, 0.0, 3.0, 0.0, 0.0, 0.0, 7.0], + [1.0, 1.0e-6, 0.0, 1.0e-6, 1.0 + 1.0e-12, -1.0e-6, 0.0, -1.0e-6, 3.0] + ] + +fullRankMatrix32Anchors :: [[Double]] +fullRankMatrix32Anchors = + [ [3.0, 0.0, 0.0, 2.0, 1.0, 1.0], + [1.0, 2.0, 2.0, -1.0, 0.5, 3.0], + [4.0, 1.0, 1.0, 3.0, -1.0, 2.0] + ] + +qrReconstructsInputLaw :: QC.Property +qrReconstructsInputLaw = + QC.property qrReconstructsInputLawProperty + +qrOrthonormalColumnsLaw :: QC.Property +qrOrthonormalColumnsLaw = + QC.property qrOrthonormalColumnsLawProperty + +choleskyReconstructsSpdLaw :: QC.Property +choleskyReconstructsSpdLaw = + QC.property choleskyReconstructsSpdLawProperty + +symmetricEigenReconstructsLaw :: QC.Property +symmetricEigenReconstructsLaw = + QC.property symmetricEigenReconstructsLawProperty + +symmetricEigenOrthonormalLaw :: QC.Property +symmetricEigenOrthonormalLaw = + QC.property symmetricEigenOrthonormalLawProperty + +symmetricEigenUncheckedPassesCertificationLaw :: QC.Property +symmetricEigenUncheckedPassesCertificationLaw = + QC.property symmetricEigenUncheckedPassesCertificationLawProperty + +thinSvdReconstructsLaw :: QC.Property +thinSvdReconstructsLaw = + QC.property thinSvdReconstructsLawProperty + +thinSvdOrthonormalFactorsLaw :: QC.Property +thinSvdOrthonormalFactorsLaw = + QC.property thinSvdOrthonormalFactorsLawProperty + +thinSvdSingularValuesOrderedNonnegativeLaw :: QC.Property +thinSvdSingularValuesOrderedNonnegativeLaw = + QC.property thinSvdSingularValuesOrderedNonnegativeLawProperty + +qrReconstructsInputLawProperty :: FullRankMatrix43 -> QC.Property +qrReconstructsInputLawProperty (FullRankMatrix43 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @4 @3 entries + (qMatrix, rMatrix) <- qrDecompFullColumnRank matrixValue + reconstructed <- mult qMatrix rMatrix + pure (maxAbsDifference entries (toListMatrix reconstructed) <= approxTolerance) + +qrOrthonormalColumnsLawProperty :: FullRankMatrix43 -> QC.Property +qrOrthonormalColumnsLawProperty (FullRankMatrix43 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @4 @3 entries + (qMatrix, _) <- qrDecompFullColumnRank matrixValue + transposedQ <- transpose qMatrix + gramMatrix <- mult transposedQ qMatrix + pure (assertApproxList identity3 (toListMatrix gramMatrix)) + +choleskyReconstructsSpdLawProperty :: SpdMatrix3 -> QC.Property +choleskyReconstructsSpdLawProperty (SpdMatrix3 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @3 @3 entries + lowerMatrix <- choleskyDecomp matrixValue + transposedLower <- transpose lowerMatrix + reconstructed <- mult lowerMatrix transposedLower + pure (maxAbsDifference entries (toListMatrix reconstructed) <= approxTolerance) + +symmetricEigenReconstructsLawProperty :: SymmetricMatrix3 -> QC.Property +symmetricEigenReconstructsLawProperty (SymmetricMatrix3 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @3 @3 entries + (eigenvalues, eigenvectors) <- symmetricEigen matrixValue + diagonalMatrix <- fromListMatrix @3 @3 (diagonal3 (toListVector eigenvalues)) + weightedEigenvectors <- mult eigenvectors diagonalMatrix + transposedEigenvectors <- transpose eigenvectors + reconstructed <- mult weightedEigenvectors transposedEigenvectors + pure (maxAbsDifference entries (toListMatrix reconstructed) <= approxTolerance) + +symmetricEigenOrthonormalLawProperty :: SymmetricMatrix3 -> QC.Property +symmetricEigenOrthonormalLawProperty (SymmetricMatrix3 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @3 @3 entries + (_, eigenvectors) <- symmetricEigen matrixValue + transposedEigenvectors <- transpose eigenvectors + gramMatrix <- mult transposedEigenvectors eigenvectors + pure (assertApproxList identity3 (toListMatrix gramMatrix)) + +symmetricEigenUncheckedPassesCertificationLawProperty :: SymmetricMatrix3 -> QC.Property +symmetricEigenUncheckedPassesCertificationLawProperty (SymmetricMatrix3 entries) = + assertRightProperty $ do + matrixValue <- first show (mkDenseDoubleMatrixRowMajor 3 3 (S.fromList entries)) + uncheckedResult <- first show (symmetricEigenPairsDenseUnchecked 3 matrixValue) + _ <- first show (certifySymmetricEigenResult matrixValue uncheckedResult) + pure True + +thinSvdReconstructsLawProperty :: FullRankMatrix32 -> QC.Property +thinSvdReconstructsLawProperty (FullRankMatrix32 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @3 @2 entries + (uMatrix, sMatrix, vTMatrix) <- thinSvdFullColumnRank matrixValue + usMatrix <- mult uMatrix sMatrix + reconstructed <- mult usMatrix vTMatrix + pure (maxAbsDifference entries (toListMatrix reconstructed) <= approxTolerance) + +thinSvdOrthonormalFactorsLawProperty :: FullRankMatrix32 -> QC.Property +thinSvdOrthonormalFactorsLawProperty (FullRankMatrix32 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @3 @2 entries + (uMatrix, _, vTMatrix) <- thinSvdFullColumnRank matrixValue + transposedU <- transpose uMatrix + uGram <- mult transposedU uMatrix + vMatrix <- transpose vTMatrix + vGram <- mult vTMatrix vMatrix + pure (assertApproxList identity2 (toListMatrix uGram) && assertApproxList identity2 (toListMatrix vGram)) + +thinSvdSingularValuesOrderedNonnegativeLawProperty :: FullRankMatrix32 -> QC.Property +thinSvdSingularValuesOrderedNonnegativeLawProperty (FullRankMatrix32 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @3 @2 entries + (_, sMatrix, _) <- thinSvdFullColumnRank matrixValue + pure (orderedNonnegativeDiagonal2 (toListMatrix sMatrix)) + +identity2 :: [Double] +identity2 = + [1.0, 0.0, 0.0, 1.0] + +identity3 :: [Double] +identity3 = + [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0] + +diagonal3 :: [Double] -> [Double] +diagonal3 values = + case values of + [d0, d1, d2] -> [d0, 0.0, 0.0, 0.0, d1, 0.0, 0.0, 0.0, d2] + _ -> [] + +orderedNonnegativeDiagonal2 :: [Double] -> Bool +orderedNonnegativeDiagonal2 entries = + case entries of + [s0, z01, z10, s1] -> + s0 >= 0.0 && s1 >= 0.0 && s0 >= s1 && assertApproxList [0.0, 0.0] [z01, z10] + _ -> False diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness/Dense.hs b/src-laws/Moonlight/LinAlg/Effect/Harness/Dense.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness/Dense.hs @@ -0,0 +1,177 @@ +module Moonlight.LinAlg.Effect.Harness.Dense + ( denseAddAssociativeLaw, + denseAddCommutativeLaw, + denseMultiplyAssociativeLaw, + denseLeftDistributiveLaw, + denseRightDistributiveLaw, + denseTransposeInvolutionLaw, + denseTransposeProductReversalLaw, + denseMapCompositionLaw, + ) +where + +import Moonlight.LinAlg (add, fromListMatrix, mapMatrix, mult, toListMatrix, transpose) +import Moonlight.LinAlg.Effect.Harness.Core (exactRightProperty) +import Test.Tasty.QuickCheck qualified as QC + +newtype RationalMatrix2 = RationalMatrix2 [Rational] + deriving stock (Eq, Show) + +instance QC.Arbitrary RationalMatrix2 where + arbitrary = + RationalMatrix2 + <$> QC.vectorOf 4 (fromIntegral <$> QC.chooseInt (-8, 8)) + +denseAddAssociativeLaw :: QC.Property +denseAddAssociativeLaw = + QC.property denseAddAssociativeLawProperty + +denseAddCommutativeLaw :: QC.Property +denseAddCommutativeLaw = + QC.property denseAddCommutativeLawProperty + +denseMultiplyAssociativeLaw :: QC.Property +denseMultiplyAssociativeLaw = + QC.property denseMultiplyAssociativeLawProperty + +denseLeftDistributiveLaw :: QC.Property +denseLeftDistributiveLaw = + QC.property denseLeftDistributiveLawProperty + +denseRightDistributiveLaw :: QC.Property +denseRightDistributiveLaw = + QC.property denseRightDistributiveLawProperty + +denseTransposeInvolutionLaw :: QC.Property +denseTransposeInvolutionLaw = + QC.property denseTransposeInvolutionLawProperty + +denseTransposeProductReversalLaw :: QC.Property +denseTransposeProductReversalLaw = + QC.property denseTransposeProductReversalLawProperty + +denseMapCompositionLaw :: QC.Property +denseMapCompositionLaw = + QC.property denseMapCompositionLawProperty + +denseAddAssociativeLawProperty :: RationalMatrix2 -> RationalMatrix2 -> RationalMatrix2 -> QC.Property +denseAddAssociativeLawProperty (RationalMatrix2 leftEntries) (RationalMatrix2 middleEntries) (RationalMatrix2 rightEntries) = + exactRightProperty leftAssociated rightAssociated + where + leftAssociated = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + middleMatrix <- fromListMatrix @2 @2 middleEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + leftMiddle <- add leftMatrix middleMatrix + fmap toListMatrix (add leftMiddle rightMatrix) + rightAssociated = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + middleMatrix <- fromListMatrix @2 @2 middleEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + middleRight <- add middleMatrix rightMatrix + fmap toListMatrix (add leftMatrix middleRight) + +denseAddCommutativeLawProperty :: RationalMatrix2 -> RationalMatrix2 -> QC.Property +denseAddCommutativeLawProperty (RationalMatrix2 leftEntries) (RationalMatrix2 rightEntries) = + exactRightProperty leftRight rightLeft + where + leftRight = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + fmap toListMatrix (add leftMatrix rightMatrix) + rightLeft = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + fmap toListMatrix (add rightMatrix leftMatrix) + +denseMultiplyAssociativeLawProperty :: RationalMatrix2 -> RationalMatrix2 -> RationalMatrix2 -> QC.Property +denseMultiplyAssociativeLawProperty (RationalMatrix2 leftEntries) (RationalMatrix2 middleEntries) (RationalMatrix2 rightEntries) = + exactRightProperty leftAssociated rightAssociated + where + leftAssociated = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + middleMatrix <- fromListMatrix @2 @2 middleEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + leftMiddle <- mult leftMatrix middleMatrix + fmap toListMatrix (mult leftMiddle rightMatrix) + rightAssociated = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + middleMatrix <- fromListMatrix @2 @2 middleEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + middleRight <- mult middleMatrix rightMatrix + fmap toListMatrix (mult leftMatrix middleRight) + +denseLeftDistributiveLawProperty :: RationalMatrix2 -> RationalMatrix2 -> RationalMatrix2 -> QC.Property +denseLeftDistributiveLawProperty (RationalMatrix2 leftEntries) (RationalMatrix2 middleEntries) (RationalMatrix2 rightEntries) = + exactRightProperty distributed expanded + where + distributed = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + middleMatrix <- fromListMatrix @2 @2 middleEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + middleRight <- add middleMatrix rightMatrix + fmap toListMatrix (mult leftMatrix middleRight) + expanded = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + middleMatrix <- fromListMatrix @2 @2 middleEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + leftMiddle <- mult leftMatrix middleMatrix + leftRight <- mult leftMatrix rightMatrix + fmap toListMatrix (add leftMiddle leftRight) + +denseRightDistributiveLawProperty :: RationalMatrix2 -> RationalMatrix2 -> RationalMatrix2 -> QC.Property +denseRightDistributiveLawProperty (RationalMatrix2 leftEntries) (RationalMatrix2 middleEntries) (RationalMatrix2 rightEntries) = + exactRightProperty distributed expanded + where + distributed = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + middleMatrix <- fromListMatrix @2 @2 middleEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + leftMiddle <- add leftMatrix middleMatrix + fmap toListMatrix (mult leftMiddle rightMatrix) + expanded = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + middleMatrix <- fromListMatrix @2 @2 middleEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + leftRight <- mult leftMatrix rightMatrix + middleRight <- mult middleMatrix rightMatrix + fmap toListMatrix (add leftRight middleRight) + +denseTransposeInvolutionLawProperty :: RationalMatrix2 -> QC.Property +denseTransposeInvolutionLawProperty (RationalMatrix2 entries) = + exactRightProperty original transposedTwice + where + original = Right entries + transposedTwice = do + matrixValue <- fromListMatrix @2 @2 entries + once <- transpose matrixValue + twice <- transpose once + pure (toListMatrix twice) + +denseTransposeProductReversalLawProperty :: RationalMatrix2 -> RationalMatrix2 -> QC.Property +denseTransposeProductReversalLawProperty (RationalMatrix2 leftEntries) (RationalMatrix2 rightEntries) = + exactRightProperty transposedProduct reversedProduct + where + transposedProduct = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + productMatrix <- mult leftMatrix rightMatrix + fmap toListMatrix (transpose productMatrix) + reversedProduct = do + leftMatrix <- fromListMatrix @2 @2 leftEntries + rightMatrix <- fromListMatrix @2 @2 rightEntries + leftTranspose <- transpose leftMatrix + rightTranspose <- transpose rightMatrix + fmap toListMatrix (mult rightTranspose leftTranspose) + +denseMapCompositionLawProperty :: RationalMatrix2 -> QC.Property +denseMapCompositionLawProperty (RationalMatrix2 entries) = + exactRightProperty staged composed + where + staged = do + matrixValue <- fromListMatrix @2 @2 entries + incremented <- mapMatrix (+ 1) matrixValue + fmap toListMatrix (mapMatrix (* 3) incremented) + composed = do + matrixValue <- fromListMatrix @2 @2 entries + fmap toListMatrix (mapMatrix ((* 3) . (+ 1)) matrixValue) diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness/Domain.hs b/src-laws/Moonlight/LinAlg/Effect/Harness/Domain.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness/Domain.hs @@ -0,0 +1,133 @@ +module Moonlight.LinAlg.Effect.Harness.Domain + ( smithDiagonalReconstructsInputLaw, + smithDivisibilityChainLaw, + smithWitnessesUnimodularLaw, + smithDiagonalOnlyAgreesWithFullLaw, + bareissRankAgreesWithFieldRankLaw, + bareissDeterminantAgreesWithExteriorLaw, + ) +where + +import Data.Bifunctor (first) +import Moonlight.LinAlg + ( bareissDeterminant, + bareissRank, + exteriorPowerMatrix, + fromListMatrix, + mult, + rank, + smithDiagonal, + smithDiagonalForm, + smithDiagonalMatrix, + smithLeft, + smithLeftInverse, + smithNormalForm, + smithRight, + smithRightInverse, + toListMatrix, + ) +import Moonlight.LinAlg.Effect.Harness.Core (assertRightProperty, matrixRows3) +import Test.Tasty.QuickCheck qualified as QC + +newtype IntegerMatrix3 = IntegerMatrix3 [Integer] + deriving stock (Eq, Show) + +instance QC.Arbitrary IntegerMatrix3 where + arbitrary = + IntegerMatrix3 + <$> QC.vectorOf 9 (fromIntegral <$> QC.chooseInt (-8, 8)) + +smithDiagonalReconstructsInputLaw :: QC.Property +smithDiagonalReconstructsInputLaw = + QC.property smithDiagonalReconstructsInputLawProperty + +smithDivisibilityChainLaw :: QC.Property +smithDivisibilityChainLaw = + QC.property smithDivisibilityChainLawProperty + +smithWitnessesUnimodularLaw :: QC.Property +smithWitnessesUnimodularLaw = + QC.property smithWitnessesUnimodularLawProperty + +smithDiagonalOnlyAgreesWithFullLaw :: QC.Property +smithDiagonalOnlyAgreesWithFullLaw = + QC.property smithDiagonalOnlyAgreesWithFullLawProperty + +bareissRankAgreesWithFieldRankLaw :: QC.Property +bareissRankAgreesWithFieldRankLaw = + QC.property bareissRankAgreesWithFieldRankLawProperty + +bareissDeterminantAgreesWithExteriorLaw :: QC.Property +bareissDeterminantAgreesWithExteriorLaw = + QC.property bareissDeterminantAgreesWithExteriorLawProperty + +smithDiagonalReconstructsInputLawProperty :: IntegerMatrix3 -> QC.Property +smithDiagonalReconstructsInputLawProperty (IntegerMatrix3 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @3 @3 entries + smithValue <- smithNormalForm matrixValue + leftTimesInput <- mult (smithLeft smithValue) matrixValue + reconstructed <- mult leftTimesInput (smithRight smithValue) + pure (toListMatrix reconstructed == toListMatrix (smithDiagonal smithValue)) + +smithDivisibilityChainLawProperty :: IntegerMatrix3 -> QC.Property +smithDivisibilityChainLawProperty (IntegerMatrix3 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @3 @3 entries + smithValue <- smithNormalForm matrixValue + pure (diagonalDivisibility (toListMatrix (smithDiagonal smithValue))) + +smithWitnessesUnimodularLawProperty :: IntegerMatrix3 -> QC.Property +smithWitnessesUnimodularLawProperty (IntegerMatrix3 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @3 @3 entries + smithValue <- smithNormalForm matrixValue + leftInverseLeft <- mult (smithLeftInverse smithValue) (smithLeft smithValue) + leftLeftInverse <- mult (smithLeft smithValue) (smithLeftInverse smithValue) + rightInverseRight <- mult (smithRightInverse smithValue) (smithRight smithValue) + rightRightInverse <- mult (smithRight smithValue) (smithRightInverse smithValue) + let identityEntries = [1, 0, 0, 0, 1, 0, 0, 0, 1] + pure + ( toListMatrix leftInverseLeft == identityEntries + && toListMatrix leftLeftInverse == identityEntries + && toListMatrix rightInverseRight == identityEntries + && toListMatrix rightRightInverse == identityEntries + ) + +smithDiagonalOnlyAgreesWithFullLawProperty :: IntegerMatrix3 -> QC.Property +smithDiagonalOnlyAgreesWithFullLawProperty (IntegerMatrix3 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @3 @3 entries + fullValue <- smithNormalForm matrixValue + diagonalOnlyValue <- smithDiagonalForm matrixValue + pure (toListMatrix (smithDiagonal fullValue) == toListMatrix (smithDiagonalMatrix diagonalOnlyValue)) + +bareissRankAgreesWithFieldRankLawProperty :: IntegerMatrix3 -> QC.Property +bareissRankAgreesWithFieldRankLawProperty (IntegerMatrix3 entries) = + assertRightProperty $ do + integerMatrix <- fromListMatrix @3 @3 entries + rationalMatrix <- fromListMatrix @3 @3 (fromInteger <$> entries :: [Rational]) + integerRank <- bareissRank integerMatrix + rationalRank <- rank rationalMatrix + pure (integerRank == rationalRank) + +bareissDeterminantAgreesWithExteriorLawProperty :: IntegerMatrix3 -> QC.Property +bareissDeterminantAgreesWithExteriorLawProperty (IntegerMatrix3 entries) = + assertRightProperty $ do + integerMatrix <- first show (fromListMatrix @3 @3 entries) + integerDeterminant <- first show (bareissDeterminant integerMatrix) + exteriorDeterminant <- first show (exteriorPowerMatrix 3 (matrixRows3 (fromInteger <$> entries :: [Rational]))) + pure (exteriorDeterminant == [[fromInteger integerDeterminant]]) + +diagonalDivisibility :: [Integer] -> Bool +diagonalDivisibility entries = + case entries of + [d0, z01, z02, z10, d1, z12, z20, z21, d2] -> + all (== 0) [z01, z02, z10, z12, z20, z21] + && divides d0 d1 + && divides d1 d2 + _ -> False + +divides :: Integer -> Integer -> Bool +divides left right = + left == 0 || right == 0 || right `rem` left == 0 diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness/Field.hs b/src-laws/Moonlight/LinAlg/Effect/Harness/Field.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness/Field.hs @@ -0,0 +1,169 @@ +module Moonlight.LinAlg.Effect.Harness.Field + ( pluReconstructsInputLaw, + rankKernelNullityLaw, + kernelVectorsAnnihilatedLaw, + packedLinearMapIdentityLaw, + packedLinearMapCompositionLaw, + gf2PackedInverseTwoSidedLaw, + ) +where + +import Data.Bifunctor (first) +import Data.Vector qualified as V +import Moonlight.LinAlg + ( GF2 (..), + applyPackedLinearMap, + composePackedLinearMaps, + fromListMatrix, + gf2PackedMatrixLinearMap, + identityPackedLinearMap, + inverseGF2PackedMatrix, + kernel, + kernelBasisVectors, + mult, + packedLinearMapColumns, + packedLinearMapFromEntries, + packedRowFromIndices, + packedRowIndices, + pluDecompFullRank, + pluLower, + pluPermutation, + pluUpper, + rank, + toListMatrix, + toListVector, + mkGF2PackedMatrixFromRowMajor, + ) +import Moonlight.LinAlg.Effect.Harness.Core (assertRightProperty) +import Test.Tasty.QuickCheck qualified as QC + +newtype InvertibleRational2 = InvertibleRational2 [Rational] + deriving stock (Eq, Show) + +newtype RationalMatrix23 = RationalMatrix23 [Rational] + deriving stock (Eq, Show) + +newtype PackedRow3 = PackedRow3 [Int] + deriving stock (Eq, Show) + +instance QC.Arbitrary InvertibleRational2 where + arbitrary = + InvertibleRational2 + <$> QC.suchThat + (QC.vectorOf 4 (fromIntegral <$> QC.chooseInt (-5, 5))) + invertible2 + +instance QC.Arbitrary RationalMatrix23 where + arbitrary = + RationalMatrix23 + <$> QC.vectorOf 6 (fromIntegral <$> QC.chooseInt (-5, 5)) + +instance QC.Arbitrary PackedRow3 where + arbitrary = + PackedRow3 + <$> QC.sublistOf [0, 1, 2] + +pluReconstructsInputLaw :: QC.Property +pluReconstructsInputLaw = + QC.property pluReconstructsInputLawProperty + +rankKernelNullityLaw :: QC.Property +rankKernelNullityLaw = + QC.property rankKernelNullityLawProperty + +kernelVectorsAnnihilatedLaw :: QC.Property +kernelVectorsAnnihilatedLaw = + QC.property kernelVectorsAnnihilatedLawProperty + +packedLinearMapIdentityLaw :: QC.Property +packedLinearMapIdentityLaw = + QC.property packedLinearMapIdentityLawProperty + +packedLinearMapCompositionLaw :: QC.Property +packedLinearMapCompositionLaw = + QC.property packedLinearMapCompositionLawProperty + +invertible2 :: [Rational] -> Bool +invertible2 entries = + case entries of + [a, b, c, d] -> a * d - b * c /= 0 + _ -> False + +pluReconstructsInputLawProperty :: InvertibleRational2 -> QC.Property +pluReconstructsInputLawProperty (InvertibleRational2 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @2 @2 entries + pluValue <- pluDecompFullRank matrixValue + leftSide <- mult (pluPermutation pluValue) matrixValue + rightSide <- mult (pluLower pluValue) (pluUpper pluValue) + pure (toListMatrix leftSide == toListMatrix rightSide) + +rankKernelNullityLawProperty :: RationalMatrix23 -> QC.Property +rankKernelNullityLawProperty (RationalMatrix23 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @2 @3 entries + rankValue <- rank matrixValue + kernelValue <- kernel matrixValue + pure (rankValue + length (kernelBasisVectors kernelValue) == 3) + +kernelVectorsAnnihilatedLawProperty :: RationalMatrix23 -> QC.Property +kernelVectorsAnnihilatedLawProperty (RationalMatrix23 entries) = + assertRightProperty $ do + matrixValue <- fromListMatrix @2 @3 entries + kernelValue <- kernel matrixValue + annihilated <- + traverse + ( \basisVector -> do + columnMatrix <- fromListMatrix @3 @1 (toListVector basisVector) + resultMatrix <- mult matrixValue columnMatrix + pure (all (== 0) (toListMatrix resultMatrix)) + ) + (kernelBasisVectors kernelValue) + pure (and annihilated) + +packedLinearMapIdentityLawProperty :: PackedRow3 -> QC.Property +packedLinearMapIdentityLawProperty (PackedRow3 indices) = + assertRightProperty $ do + rowValue <- packedRowFromIndices "packed identity law row" 3 indices + identityMap <- identityPackedLinearMap "packed identity law map" 3 + image <- applyPackedLinearMap "packed identity law apply" identityMap rowValue + pure (packedRowIndices image == packedRowIndices rowValue) + +packedLinearMapCompositionLawProperty :: PackedRow3 -> QC.Property +packedLinearMapCompositionLawProperty (PackedRow3 indices) = + assertRightProperty $ do + rowValue <- packedRowFromIndices "packed composition law row" 3 indices + leftMap <- + packedLinearMapFromEntries + "packed composition law left" + 3 + 3 + [(0, 0), (1, 0), (1, 2), (2, 1)] + rightMap <- + packedLinearMapFromEntries + "packed composition law right" + 3 + 3 + [(0, 1), (1, 2), (2, 0), (2, 2)] + composedMap <- composePackedLinearMaps "packed composition law composed" leftMap rightMap + directImage <- applyPackedLinearMap "packed composition law direct" composedMap rowValue + stagedRight <- applyPackedLinearMap "packed composition law first" rightMap rowValue + stagedImage <- applyPackedLinearMap "packed composition law second" leftMap stagedRight + pure (packedRowIndices directImage == packedRowIndices stagedImage) + +gf2PackedInverseTwoSidedLaw :: QC.Property +gf2PackedInverseTwoSidedLaw = + assertRightProperty $ do + matrixValue <- first show (mkGF2PackedMatrixFromRowMajor 3 3 [GF2One, GF2One, GF2Zero, GF2Zero, GF2One, GF2One, GF2Zero, GF2Zero, GF2One]) + matrixMap <- first show (gf2PackedMatrixLinearMap matrixValue) + maybeInverse <- first show (inverseGF2PackedMatrix matrixValue) + case maybeInverse of + Nothing -> pure False + Just inverseMap -> do + identityMap <- first show (identityPackedLinearMap "packed inverse law identity" 3) + leftIdentity <- first show (composePackedLinearMaps "packed inverse law left" inverseMap matrixMap) + rightIdentity <- first show (composePackedLinearMaps "packed inverse law right" matrixMap inverseMap) + let identityColumns = packedRowIndices <$> V.toList (packedLinearMapColumns identityMap) + leftColumns = packedRowIndices <$> V.toList (packedLinearMapColumns leftIdentity) + rightColumns = packedRowIndices <$> V.toList (packedLinearMapColumns rightIdentity) + pure (leftColumns == identityColumns && rightColumns == identityColumns) diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness/Geometry.hs b/src-laws/Moonlight/LinAlg/Effect/Harness/Geometry.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness/Geometry.hs @@ -0,0 +1,257 @@ +module Moonlight.LinAlg.Effect.Harness.Geometry + ( vec3AddCommutativeAssociativeLaw, + vec3DotSymmetricLaw, + vec3NormalizeUnitLaw, + aabbUnionCommutativeAssociativeLaw, + aabbUnionContainsOperandsLaw, + aabbIntersectionCommutativeLaw, + symmetricOuterApplyAgreementLaw, + geometrySymmetricEigenReconstructsLaw, + geometrySymmetricEigenOrthonormalLaw, + ) +where + +import Moonlight.LinAlg + ( AABB, + Symmetric3 (..), + Vec3 (..), + aabbMax, + aabbMin, + addVec3, + applySymmetric3, + diagonalSymmetric3, + dotVec3, + eigendecomposeSymmetric3, + magnitudeVec3, + mkAabb, + normalizeVec3, + outerSymmetric3, + scaleVec3, + symmetric3ToMatrix, + toListMatrix, + toListVector, + unionAabb, + ) +import Moonlight.LinAlg.Effect.Harness.Core (assertApproxEqual, assertApproxList, assertRightProperty, matrix3Product) +import Test.Tasty.QuickCheck qualified as QC + +newtype GeneratedVec3 = GeneratedVec3 Vec3 + deriving stock (Eq, Show) + +newtype NonZeroVec3 = NonZeroVec3 Vec3 + deriving stock (Eq, Show) + +newtype GeneratedAABB = GeneratedAABB AABB + deriving stock (Eq, Show) + +newtype GeneratedSymmetric3 = GeneratedSymmetric3 (Symmetric3 Double) + deriving stock (Eq, Show) + +newtype GeneratedWeight = GeneratedWeight Double + deriving stock (Eq, Show) + +instance QC.Arbitrary GeneratedVec3 where + arbitrary = + GeneratedVec3 + <$> (Vec3 <$> component <*> component <*> component) + +instance QC.Arbitrary NonZeroVec3 where + arbitrary = + NonZeroVec3 + <$> QC.suchThat + (Vec3 <$> component <*> component <*> component) + (\value -> magnitudeVec3 value > 1.0e-6) + +instance QC.Arbitrary GeneratedAABB where + arbitrary = do + minX <- component + minY <- component + minZ <- component + sizeX <- nonNegativeComponent + sizeY <- nonNegativeComponent + sizeZ <- nonNegativeComponent + case mkAabb (Vec3 minX minY minZ) (Vec3 (minX + sizeX) (minY + sizeY) (minZ + sizeZ)) of + Nothing -> QC.discard + Just boxValue -> pure (GeneratedAABB boxValue) + +instance QC.Arbitrary GeneratedSymmetric3 where + arbitrary = + GeneratedSymmetric3 <$> QC.oneof [repeatedSpectrum, tinyGapSpectrum, badlyScaledSpectrum, coupledSpectrum] + +instance QC.Arbitrary GeneratedWeight where + arbitrary = + GeneratedWeight <$> component + +vec3AddCommutativeAssociativeLaw :: QC.Property +vec3AddCommutativeAssociativeLaw = + QC.property vec3AddCommutativeAssociativeLawProperty + +vec3DotSymmetricLaw :: QC.Property +vec3DotSymmetricLaw = + QC.property vec3DotSymmetricLawProperty + +vec3NormalizeUnitLaw :: QC.Property +vec3NormalizeUnitLaw = + QC.property vec3NormalizeUnitLawProperty + +aabbUnionCommutativeAssociativeLaw :: QC.Property +aabbUnionCommutativeAssociativeLaw = + QC.property aabbUnionCommutativeAssociativeLawProperty + +aabbUnionContainsOperandsLaw :: QC.Property +aabbUnionContainsOperandsLaw = + QC.property aabbUnionContainsOperandsLawProperty + +aabbIntersectionCommutativeLaw :: QC.Property +aabbIntersectionCommutativeLaw = + QC.property aabbIntersectionCommutativeLawProperty + +symmetricOuterApplyAgreementLaw :: QC.Property +symmetricOuterApplyAgreementLaw = + QC.property symmetricOuterApplyAgreementLawProperty + +geometrySymmetricEigenReconstructsLaw :: QC.Property +geometrySymmetricEigenReconstructsLaw = + QC.property geometrySymmetricEigenReconstructsLawProperty + +geometrySymmetricEigenOrthonormalLaw :: QC.Property +geometrySymmetricEigenOrthonormalLaw = + QC.property geometrySymmetricEigenOrthonormalLawProperty + +component :: QC.Gen Double +component = + fromIntegral <$> QC.chooseInt (-8, 8) + +nonNegativeComponent :: QC.Gen Double +nonNegativeComponent = + fromIntegral <$> QC.chooseInt (0, 8) + +vec3AddCommutativeAssociativeLawProperty :: GeneratedVec3 -> GeneratedVec3 -> GeneratedVec3 -> QC.Property +vec3AddCommutativeAssociativeLawProperty (GeneratedVec3 leftValue) (GeneratedVec3 middleValue) (GeneratedVec3 rightValue) = + QC.property + ( addVec3 leftValue middleValue == addVec3 middleValue leftValue + && addVec3 (addVec3 leftValue middleValue) rightValue == addVec3 leftValue (addVec3 middleValue rightValue) + ) + +vec3DotSymmetricLawProperty :: GeneratedVec3 -> GeneratedVec3 -> QC.Property +vec3DotSymmetricLawProperty (GeneratedVec3 leftValue) (GeneratedVec3 rightValue) = + QC.property (assertApproxEqual (dotVec3 leftValue rightValue) (dotVec3 rightValue leftValue)) + +vec3NormalizeUnitLawProperty :: NonZeroVec3 -> QC.Property +vec3NormalizeUnitLawProperty (NonZeroVec3 value) = + assertRightProperty $ do + normalized <- normalizeVec3 value + pure (assertApproxEqual 1.0 (magnitudeVec3 normalized)) + +aabbUnionCommutativeAssociativeLawProperty :: GeneratedAABB -> GeneratedAABB -> GeneratedAABB -> QC.Property +aabbUnionCommutativeAssociativeLawProperty (GeneratedAABB leftValue) (GeneratedAABB middleValue) (GeneratedAABB rightValue) = + QC.property + ( unionAabb leftValue middleValue == unionAabb middleValue leftValue + && unionAabb (unionAabb leftValue middleValue) rightValue == unionAabb leftValue (unionAabb middleValue rightValue) + ) + +aabbUnionContainsOperandsLawProperty :: GeneratedAABB -> GeneratedAABB -> QC.Property +aabbUnionContainsOperandsLawProperty (GeneratedAABB leftValue) (GeneratedAABB rightValue) = + let unionValue = unionAabb leftValue rightValue + in QC.property (aabbContains unionValue leftValue && aabbContains unionValue rightValue) + +aabbIntersectionCommutativeLawProperty :: GeneratedAABB -> GeneratedAABB -> QC.Property +aabbIntersectionCommutativeLawProperty (GeneratedAABB leftValue) (GeneratedAABB rightValue) = + QC.property (intersectAabb leftValue rightValue == intersectAabb rightValue leftValue) + +symmetricOuterApplyAgreementLawProperty :: GeneratedWeight -> GeneratedVec3 -> GeneratedVec3 -> QC.Property +symmetricOuterApplyAgreementLawProperty (GeneratedWeight weightValue) (GeneratedVec3 basisValue) (GeneratedVec3 inputValue) = + let tensorValue = outerSymmetric3 weightValue basisValue + expectedValue = scaleVec3 (weightValue * dotVec3 basisValue inputValue) basisValue + in QC.property (vec3Approx expectedValue (applySymmetric3 tensorValue inputValue)) + +geometrySymmetricEigenReconstructsLawProperty :: GeneratedSymmetric3 -> QC.Property +geometrySymmetricEigenReconstructsLawProperty (GeneratedSymmetric3 tensorValue) = + assertRightProperty $ do + matrixValue <- symmetric3ToMatrix tensorValue + (eigenvalues, eigenvectors) <- eigendecomposeSymmetric3 tensorValue + let originalRows = rows3 (toListMatrix matrixValue) + eigenvectorRows = rows3 (toListMatrix eigenvectors) + diagonalRows = diagonal3 (toListVector eigenvalues) + reconstructedRows = matrix3Product (matrix3Product eigenvectorRows diagonalRows) (transposeRows3 eigenvectorRows) + pure (and (zipWith assertApproxList originalRows reconstructedRows)) + +geometrySymmetricEigenOrthonormalLawProperty :: GeneratedSymmetric3 -> QC.Property +geometrySymmetricEigenOrthonormalLawProperty (GeneratedSymmetric3 tensorValue) = + assertRightProperty $ do + (_, eigenvectors) <- eigendecomposeSymmetric3 tensorValue + let eigenvectorRows = rows3 (toListMatrix eigenvectors) + gramRows = matrix3Product (transposeRows3 eigenvectorRows) eigenvectorRows + pure (and (zipWith assertApproxList identityRows3 gramRows)) + +repeatedSpectrum :: QC.Gen (Symmetric3 Double) +repeatedSpectrum = + pure (diagonalSymmetric3 2.0 2.0 5.0) + +tinyGapSpectrum :: QC.Gen (Symmetric3 Double) +tinyGapSpectrum = + pure (diagonalSymmetric3 1.0 (1.0 + 1.0e-12) 3.0) + +badlyScaledSpectrum :: QC.Gen (Symmetric3 Double) +badlyScaledSpectrum = + pure (diagonalSymmetric3 1.0e-12 1.0 1.0e12) + +coupledSpectrum :: QC.Gen (Symmetric3 Double) +coupledSpectrum = + pure + Symmetric3 + { sym3XX = 3.0, + sym3XY = 1.0e-6, + sym3XZ = 2.0e-6, + sym3YY = 3.0 + 1.0e-12, + sym3YZ = -1.0e-6, + sym3ZZ = 7.0 + } + +intersectAabb :: AABB -> AABB -> Maybe AABB +intersectAabb leftValue rightValue = + mkAabb + (Vec3 (max lx rx) (max ly ry) (max lz rz)) + (Vec3 (min lxx rxx) (min lyy ryy) (min lzz rzz)) + where + Vec3 lx ly lz = aabbMin leftValue + Vec3 lxx lyy lzz = aabbMax leftValue + Vec3 rx ry rz = aabbMin rightValue + Vec3 rxx ryy rzz = aabbMax rightValue + +aabbContains :: AABB -> AABB -> Bool +aabbContains outerValue innerValue = + let Vec3 ox oy oz = aabbMin outerValue + Vec3 oxx oyy ozz = aabbMax outerValue + Vec3 ix iy iz = aabbMin innerValue + Vec3 ixx iyy izz = aabbMax innerValue + in ox <= ix && oy <= iy && oz <= iz && ixx <= oxx && iyy <= oyy && izz <= ozz + +vec3Approx :: Vec3 -> Vec3 -> Bool +vec3Approx (Vec3 ex ey ez) (Vec3 ax ay az) = + assertApproxList [ex, ey, ez] [ax, ay, az] + +rows3 :: [a] -> [[a]] +rows3 values = + case values of + [a00, a01, a02, a10, a11, a12, a20, a21, a22] -> + [[a00, a01, a02], [a10, a11, a12], [a20, a21, a22]] + _ -> [] + +transposeRows3 :: [[a]] -> [[a]] +transposeRows3 rowsValue = + case rowsValue of + [[a00, a01, a02], [a10, a11, a12], [a20, a21, a22]] -> + [[a00, a10, a20], [a01, a11, a21], [a02, a12, a22]] + _ -> [] + +diagonal3 :: [Double] -> [[Double]] +diagonal3 values = + case values of + [xValue, yValue, zValue] -> + [[xValue, 0.0, 0.0], [0.0, yValue, 0.0], [0.0, 0.0, zValue]] + _ -> [] + +identityRows3 :: [[Double]] +identityRows3 = + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]] diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness/KrylovSpectral.hs b/src-laws/Moonlight/LinAlg/Effect/Harness/KrylovSpectral.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness/KrylovSpectral.hs @@ -0,0 +1,299 @@ +{-# LANGUAGE DataKinds #-} + +module Moonlight.LinAlg.Effect.Harness.KrylovSpectral + ( arnoldiRelationHoldsLaw, + arnoldiBasisOrthonormalLaw, + lanczosProjectionTridiagonalLaw, + lanczosBasisOrthonormalLaw, + thickRestartLockedPairsResidualBoundedLaw, + selectedPairsResidualBoundedLaw, + selectedPairsClusterOrthonormalLaw, + tridiagonalSelectedValuesAgreeWithAllPairsLaw, + diagonalSpectralValuesExactLaw, + pathLaplacianSpectralValuesClosedFormLaw, + eigenRequestRejectsOversubscriptionLaw, + ) +where + +import Data.Bifunctor (first) +import Data.Vector qualified as Box +import Data.Vector.Unboxed qualified as U +import Moonlight.Core (fieldValueValid) +import Moonlight.LinAlg + ( EigenRequest (..), + EigenSolveConfig, + Eigenpairs, + LinearOperator, + OperatorSymmetry (..), + SpectrumEnd (..), + arnoldi, + arnoldiBasisColumns, + arnoldiHessenbergRows, + defaultArnoldiConfig, + defaultEigenSolveConfig, + defaultLanczosConfig, + diagonalLinearOperator, + eigenpairCount, + eigenpairResidualNorms, + eigenpairValues, + eigenpairVectorAt, + lanczosAlphaDiagonal, + lanczosBasisColumns, + lanczosBetaOffDiagonal, + lanczosStepsCompleted, + lanczosSymmetric, + mkNonNegativeConfigTolerance, + mkPositiveCount, + mkSparseCOO, + cooToCSR, + pathLaplacianLinearOperator, + runOperatorU, + selfAdjointCSRLinearOperator, + solveEigenRequest, + withArnoldiIterations, + withEigenFallbackInitialVector, + withEigenFallbackLanczosConfig, + withLanczosIterations, + withLanczosTolerance, + ) +import Moonlight.LinAlg.Effect.Harness.Core + ( approxTolerance, + assertApproxList, + assertApproxListWith, + assertRightProperty, + orthonormalTolerance, + residualTolerance, + ) +import Test.Tasty.QuickCheck qualified as QC + +arnoldiRelationHoldsLaw :: QC.Property +arnoldiRelationHoldsLaw = + assertRightProperty $ do + iterationCount <- mapLeftShow (mkPositiveCount 2) + operatorValue <- mapLeftShow (diagonalLinearOperator (U.fromList [1.0, 3.0])) + decomposition <- mapLeftShow (arnoldi (withArnoldiIterations iterationCount defaultArnoldiConfig) operatorValue (U.fromList [1.0, 1.0])) + relationHolds operatorValue (arnoldiBasisColumns decomposition) (arnoldiHessenbergRows decomposition) + +arnoldiBasisOrthonormalLaw :: QC.Property +arnoldiBasisOrthonormalLaw = + assertRightProperty $ do + iterationCount <- mapLeftShow (mkPositiveCount 3) + operatorValue <- mapLeftShow (diagonalLinearOperator (U.fromList [1.0, 2.0, 4.0])) + decomposition <- mapLeftShow (arnoldi (withArnoldiIterations iterationCount defaultArnoldiConfig) operatorValue (U.fromList [1.0, 1.0, 1.0])) + pure (orthonormalColumns (arnoldiBasisColumns decomposition)) + +lanczosProjectionTridiagonalLaw :: QC.Property +lanczosProjectionTridiagonalLaw = + assertRightProperty $ do + iterationCount <- mapLeftShow (mkPositiveCount 3) + operatorValue <- mapLeftShow (pathLaplacianLinearOperator 4) + decomposition <- mapLeftShow (lanczosSymmetric (withLanczosIterations iterationCount defaultLanczosConfig) operatorValue (U.fromList [1.0, 0.0, 0.0, 0.0])) + let stepCount = lanczosStepsCompleted decomposition + pure + ( stepCount > 0 + && length (lanczosAlphaDiagonal decomposition) == stepCount + && length (lanczosBetaOffDiagonal decomposition) == max 0 (stepCount - 1) + && Box.length (lanczosBasisColumns decomposition) == stepCount + ) + +lanczosBasisOrthonormalLaw :: QC.Property +lanczosBasisOrthonormalLaw = + assertRightProperty $ do + iterationCount <- mapLeftShow (mkPositiveCount 4) + operatorValue <- mapLeftShow (pathLaplacianLinearOperator 5) + decomposition <- mapLeftShow (lanczosSymmetric (withLanczosIterations iterationCount defaultLanczosConfig) operatorValue (U.fromList [1.0, 0.5, 0.25, 0.125, 0.0625])) + pure (orthonormalColumns (lanczosBasisColumns decomposition)) + +thickRestartLockedPairsResidualBoundedLaw :: QC.Property +thickRestartLockedPairsResidualBoundedLaw = + assertRightProperty $ do + countValue <- mapLeftShow (mkPositiveCount 3) + operatorValue <- genericPentadiagonalOperator 18 + solveConfig <- restartedSolveConfig 5 approxTolerance (restartSeedVector 18) + pairs <- mapLeftShow (solveEigenRequest solveConfig operatorValue (EigenpairsRequest SmallestEigenvalues countValue)) + pure (eigenpairCount pairs == 3 && eigenpairResidualsBounded pairs) + +selectedPairsResidualBoundedLaw :: QC.Property +selectedPairsResidualBoundedLaw = + assertRightProperty $ do + countValue <- mapLeftShow (mkPositiveCount 2) + operatorValue <- tridiagonalOperator [2.0, 2.5, 3.0, 3.5, 4.0] [-0.31, -0.27, -0.23, -0.19] + pairs <- mapLeftShow (solveEigenRequest defaultEigenSolveConfig operatorValue (EigenpairsRequest SmallestEigenvalues countValue)) + pure (eigenpairCount pairs == 2 && eigenpairResidualsBounded pairs) + +selectedPairsClusterOrthonormalLaw :: QC.Property +selectedPairsClusterOrthonormalLaw = + assertRightProperty $ do + countValue <- mapLeftShow (mkPositiveCount 3) + operatorValue <- mapLeftShow (diagonalLinearOperator (U.fromList [2.0, 2.0, 2.0, 5.0])) + pairs <- mapLeftShow (solveEigenRequest defaultEigenSolveConfig operatorValue (EigenpairsRequest SmallestEigenvalues countValue)) + pure (eigenpairCount pairs == 3 && orthonormalEigenpairs pairs) + +tridiagonalSelectedValuesAgreeWithAllPairsLaw :: QC.Property +tridiagonalSelectedValuesAgreeWithAllPairsLaw = + assertRightProperty $ do + selectedCount <- mapLeftShow (mkPositiveCount 3) + fullCount <- mapLeftShow (mkPositiveCount 5) + operatorValue <- tridiagonalOperator [2.0, 2.5, 3.0, 3.5, 4.0] [-0.31, -0.27, -0.23, -0.19] + selectedValues <- mapLeftShow (solveEigenRequest defaultEigenSolveConfig operatorValue (EigenvaluesRequest SmallestEigenvalues selectedCount)) + allPairs <- mapLeftShow (solveEigenRequest defaultEigenSolveConfig operatorValue (EigenpairsRequest SmallestEigenvalues fullCount)) + pure (assertApproxListWith residualTolerance (U.toList selectedValues) (take 3 (U.toList (eigenpairValues allPairs)))) + +diagonalSpectralValuesExactLaw :: QC.Property +diagonalSpectralValuesExactLaw = + assertRightProperty $ do + countValue <- mapLeftShow (mkPositiveCount 2) + operatorValue <- mapLeftShow (diagonalLinearOperator (U.fromList [3.0, -2.0, 7.0, 1.0])) + values <- mapLeftShow (solveEigenRequest defaultEigenSolveConfig operatorValue (EigenvaluesRequest SmallestEigenvalues countValue)) + pure (U.toList values == [-2.0, 1.0]) + +pathLaplacianSpectralValuesClosedFormLaw :: QC.Property +pathLaplacianSpectralValuesClosedFormLaw = + assertRightProperty $ do + countValue <- mapLeftShow (mkPositiveCount 3) + operatorValue <- mapLeftShow (pathLaplacianLinearOperator 5) + values <- mapLeftShow (solveEigenRequest defaultEigenSolveConfig operatorValue (EigenvaluesRequest SmallestEigenvalues countValue)) + pure (assertApproxList (pathLaplacianValues 5 [0, 1, 2]) (U.toList values)) + +eigenRequestRejectsOversubscriptionLaw :: QC.Property +eigenRequestRejectsOversubscriptionLaw = + assertRightProperty $ do + countValue <- mapLeftShow (mkPositiveCount 4) + operatorValue <- mapLeftShow (diagonalLinearOperator (U.fromList [1.0, 2.0, 3.0])) + let resultValue = solveEigenRequest defaultEigenSolveConfig operatorValue (EigenvaluesRequest SmallestEigenvalues countValue) + pure + ( case resultValue of + Left _ -> True + Right _ -> False + ) + +relationHolds :: + LinearOperator symmetry -> + Box.Vector (U.Vector Double) -> + Box.Vector (U.Vector Double) -> + Either String Bool +relationHolds operatorValue basisColumns hessenbergRows = + let basisValues = Box.toList basisColumns + hessenbergValues = U.toList <$> Box.toList hessenbergRows + stepCount = Box.length hessenbergRows - 1 + relationAt columnIndex = do + basisVector <- maybeToEither ("missing Arnoldi basis column " <> show columnIndex) (entryAt columnIndex basisValues) + imageVector <- mapLeftShow (runOperatorU operatorValue basisVector) + coefficients <- traverse (maybeToEither ("missing Arnoldi coefficient at column " <> show columnIndex) . entryAt columnIndex) hessenbergValues + pure + ( assertApproxList + (U.toList imageVector) + (U.toList (linearCombinationU (take (length basisValues) coefficients) basisValues)) + && assertApproxList [0.0] (drop (length basisValues) coefficients) + ) + in fmap and (traverse relationAt [0 .. stepCount - 1]) + +orthonormalColumns :: Box.Vector (U.Vector Double) -> Bool +orthonormalColumns columns = + and + [ assertApproxListWith orthonormalTolerance [expectedValue] [vectorDotU leftColumn rightColumn] + | (leftIndex, leftColumn) <- zip [0 :: Int ..] (Box.toList columns), + (rightIndex, rightColumn) <- zip [0 :: Int ..] (Box.toList columns), + leftIndex <= rightIndex, + let expectedValue = if leftIndex == rightIndex then 1.0 else 0.0 + ] + +orthonormalEigenpairs :: Eigenpairs -> Bool +orthonormalEigenpairs pairs = + case traverse (`eigenpairVectorAt` pairs) [0 .. eigenpairCount pairs - 1] of + Left _ -> False + Right columns -> orthonormalColumns (Box.fromList columns) + +eigenpairResidualsBounded :: Eigenpairs -> Bool +eigenpairResidualsBounded pairs = + U.all (\residualNorm -> fieldValueValid residualNorm && residualNorm <= residualTolerance) (eigenpairResidualNorms pairs) + +restartedSolveConfig :: Int -> Double -> U.Vector Double -> Either String EigenSolveConfig +restartedSolveConfig iterationLimit toleranceValue seedVector = do + iterationCount <- mapLeftShow (mkPositiveCount iterationLimit) + toleranceBound <- mapLeftShow (mkNonNegativeConfigTolerance toleranceValue) + let lanczosConfig = withLanczosTolerance toleranceBound (withLanczosIterations iterationCount defaultLanczosConfig) + pure (withEigenFallbackInitialVector seedVector (withEigenFallbackLanczosConfig lanczosConfig defaultEigenSolveConfig)) + +tridiagonalOperator :: [Double] -> [Double] -> Either String (LinearOperator 'SelfAdjointOperator) +tridiagonalOperator diagonalEntries offDiagonalEntries = + mapLeftShow + ( selfAdjointCSRLinearOperator + =<< (mkSparseCOO (length diagonalEntries) (length diagonalEntries) (tridiagonalEntries diagonalEntries offDiagonalEntries) >>= cooToCSR) + ) + +tridiagonalEntries :: [Double] -> [Double] -> [(Int, Int, Double)] +tridiagonalEntries diagonalEntries offDiagonalEntries = + zipWith (\entryIndex entryValue -> (entryIndex, entryIndex, entryValue)) [0 ..] diagonalEntries + <> concat + ( zipWith + ( \entryIndex entryValue -> + [(entryIndex, entryIndex + 1, entryValue), (entryIndex + 1, entryIndex, entryValue)] + ) + [0 ..] + offDiagonalEntries + ) + +genericPentadiagonalOperator :: Int -> Either String (LinearOperator 'SelfAdjointOperator) +genericPentadiagonalOperator dimension = + mapLeftShow + ( selfAdjointCSRLinearOperator + =<< (mkSparseCOO dimension dimension (genericPentadiagonalEntries dimension) >>= cooToCSR) + ) + +genericPentadiagonalEntries :: Int -> [(Int, Int, Double)] +genericPentadiagonalEntries dimension = + diagonalEntries <> firstOffDiagonalEntries <> secondOffDiagonalEntries + where + diagonalEntries = + (\rowIndex -> (rowIndex, rowIndex, 4.0 + 0.03 * fromIntegral (rowIndex `mod` 7))) + <$> [0 .. dimension - 1] + firstOffDiagonalEntries = + symmetricBandEntries dimension 1 (\rowIndex -> -1.0 - 0.01 * fromIntegral (rowIndex `mod` 5)) + secondOffDiagonalEntries = + symmetricBandEntries dimension 2 (\rowIndex -> -0.2 - 0.005 * fromIntegral (rowIndex `mod` 3)) + +symmetricBandEntries :: Int -> Int -> (Int -> Double) -> [(Int, Int, Double)] +symmetricBandEntries dimension offset entryValueAt = + concat + ( ( \rowIndex -> + let columnIndex = rowIndex + offset + entryValue = entryValueAt rowIndex + in [(rowIndex, columnIndex, entryValue), (columnIndex, rowIndex, entryValue)] + ) + <$> [0 .. dimension - offset - 1] + ) + +restartSeedVector :: Int -> U.Vector Double +restartSeedVector dimension = + U.generate dimension (\indexValue -> 1.0 / fromIntegral (indexValue + 1)) + +pathLaplacianValues :: Int -> [Int] -> [Double] +pathLaplacianValues dimension = + fmap (\modeIndex -> 2.0 - 2.0 * cos (pi * fromIntegral modeIndex / fromIntegral dimension)) + +linearCombinationU :: [Double] -> [U.Vector Double] -> U.Vector Double +linearCombinationU coefficients basisVectors = + case basisVectors of + [] -> U.empty + firstVector : _ -> + foldr (U.zipWith (+)) (U.replicate (U.length firstVector) 0.0) (zipWith (\coefficient vectorValue -> U.map (* coefficient) vectorValue) coefficients basisVectors) + +vectorDotU :: U.Vector Double -> U.Vector Double -> Double +vectorDotU leftVector rightVector = + U.sum (U.zipWith (*) leftVector rightVector) + +entryAt :: Int -> [value] -> Maybe value +entryAt targetIndex values = + case drop targetIndex values of + entryValue : _ -> Just entryValue + [] -> Nothing + +maybeToEither :: failure -> Maybe value -> Either failure value +maybeToEither failureValue value = + case value of + Just presentValue -> Right presentValue + Nothing -> Left failureValue + +mapLeftShow :: Show failure => Either failure value -> Either String value +mapLeftShow = first show diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness/Operator.hs b/src-laws/Moonlight/LinAlg/Effect/Harness/Operator.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness/Operator.hs @@ -0,0 +1,64 @@ +module Moonlight.LinAlg.Effect.Harness.Operator + ( scaledOperatorActionLaw, + shiftedIdentityActionLaw, + ) +where + +import Data.Vector.Unboxed qualified as U +import Moonlight.LinAlg + ( addScaledIdentity, + diagonalLinearOperator, + runOperatorU, + scaleLinearOperator, + ) +import Moonlight.LinAlg.Effect.Harness.Core + ( assertApproxList, + assertRightProperty, + ) +import Test.Tasty.QuickCheck qualified as QC + +newtype OperatorVector3 = OperatorVector3 [Double] + deriving stock (Eq, Show) + +newtype OperatorScale = OperatorScale Double + deriving stock (Eq, Show) + +newtype OperatorShift = OperatorShift Double + deriving stock (Eq, Show) + +instance QC.Arbitrary OperatorVector3 where + arbitrary = + OperatorVector3 + <$> QC.vectorOf 3 (fromIntegral <$> QC.chooseInt (-8, 8)) + +instance QC.Arbitrary OperatorScale where + arbitrary = + OperatorScale . fromIntegral <$> QC.chooseInt (-4, 4) + +instance QC.Arbitrary OperatorShift where + arbitrary = + OperatorShift . fromIntegral <$> QC.chooseInt (-4, 4) + +scaledOperatorActionLaw :: QC.Property +scaledOperatorActionLaw = + QC.property scaledOperatorActionLawProperty + +shiftedIdentityActionLaw :: QC.Property +shiftedIdentityActionLaw = + QC.property shiftedIdentityActionLawProperty + +scaledOperatorActionLawProperty :: OperatorScale -> OperatorVector3 -> QC.Property +scaledOperatorActionLawProperty (OperatorScale scaleValue) (OperatorVector3 vectorEntries) = + assertRightProperty $ do + operatorValue <- diagonalLinearOperator (U.fromList [2.0, -3.0, 5.0]) + baseImage <- runOperatorU operatorValue (U.fromList vectorEntries) + scaledImage <- runOperatorU (scaleLinearOperator scaleValue operatorValue) (U.fromList vectorEntries) + pure (assertApproxList (fmap (* scaleValue) (U.toList baseImage)) (U.toList scaledImage)) + +shiftedIdentityActionLawProperty :: OperatorShift -> OperatorVector3 -> QC.Property +shiftedIdentityActionLawProperty (OperatorShift shiftValue) (OperatorVector3 vectorEntries) = + assertRightProperty $ do + operatorValue <- diagonalLinearOperator (U.fromList [2.0, -3.0, 5.0]) + baseImage <- runOperatorU operatorValue (U.fromList vectorEntries) + shiftedImage <- runOperatorU (addScaledIdentity shiftValue operatorValue) (U.fromList vectorEntries) + pure (assertApproxList (zipWith (\baseValue inputValue -> baseValue + shiftValue * inputValue) (U.toList baseImage) vectorEntries) (U.toList shiftedImage)) diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness/Preconditioner.hs b/src-laws/Moonlight/LinAlg/Effect/Harness/Preconditioner.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness/Preconditioner.hs @@ -0,0 +1,155 @@ +module Moonlight.LinAlg.Effect.Harness.Preconditioner + ( ic0FactorSolveRoundTripLaw, + ic0RejectsNonpositivePivotLaw, + preconditionedCgConvergesOnSpdLaw, + ) +where + +import Data.Vector.Unboxed qualified as U +import Moonlight.Core (MoonlightError) +import Moonlight.LinAlg + ( IC0Config (..), + SparseConjugateGradientConfig (..), + SparseCSR, + SparseIterativeFailure (..), + SparsePreconditionerFamily (..), + SparseIterativeResult (..), + cooToCSR, + csrMatVecVector, + mkSparseCOO, + solveSparseCG, + ) +import Moonlight.LinAlg.Effect.Harness.Core + ( approxTolerance, + assertApproxList, + assertRightProperty, + vectorNorm, + ) +import Test.Tasty.QuickCheck qualified as QC + +newtype PositiveDiagonal3 = PositiveDiagonal3 [Double] + deriving stock (Eq, Show) + +newtype RightHandSide3 = RightHandSide3 [Double] + deriving stock (Eq, Show) + +instance QC.Arbitrary PositiveDiagonal3 where + arbitrary = + PositiveDiagonal3 + <$> QC.vectorOf 3 (fromIntegral <$> QC.chooseInt (1, 9)) + +instance QC.Arbitrary RightHandSide3 where + arbitrary = + RightHandSide3 + <$> QC.vectorOf 3 (fromIntegral <$> QC.chooseInt (-8, 8)) + +ic0FactorSolveRoundTripLaw :: QC.Property +ic0FactorSolveRoundTripLaw = + QC.property ic0FactorSolveRoundTripLawProperty + +ic0RejectsNonpositivePivotLaw :: QC.Property +ic0RejectsNonpositivePivotLaw = + assertRightProperty $ do + matrixValue <- + mkSparseCOO + 2 + 2 + [(0, 0, 1.0), (0, 1, 2.0), (1, 0, 2.0), (1, 1, 1.0)] + >>= cooToCSR + let resultValue = + solveSparseCG + (cgConfigWith 8 (IncompleteCholesky0SparsePreconditionerFamily (IC0Config Nothing))) + matrixValue + (U.fromList [1.0, 1.0]) + (U.fromList [0.0, 0.0]) + pure + ( case resultValue of + Left (SparseNonpositivePivot _ _) -> True + _ -> False + ) + +preconditionedCgConvergesOnSpdLaw :: QC.Property +preconditionedCgConvergesOnSpdLaw = + assertRightProperty $ do + matrixValue <- anchoredPathLaplacian 16 + let rhsValues = anchoredPathRightHandSide 16 + resultValue = + solveSparseCG + (cgConfigWith 128 (IncompleteCholesky0SparsePreconditionerFamily (IC0Config Nothing))) + matrixValue + rhsValues + (U.replicate 16 0.0) + pure + ( case resultValue of + Right sparseResult -> + sparseResidualNorm sparseResult <= scgcTolerance (cgConfigWith 128 (IncompleteCholesky0SparsePreconditionerFamily (IC0Config Nothing))) + && trueResidualNorm matrixValue rhsValues (sparseSolution sparseResult) <= scgcTolerance (cgConfigWith 128 (IncompleteCholesky0SparsePreconditionerFamily (IC0Config Nothing))) + Left _ -> False + ) + +ic0FactorSolveRoundTripLawProperty :: PositiveDiagonal3 -> RightHandSide3 -> QC.Property +ic0FactorSolveRoundTripLawProperty (PositiveDiagonal3 diagonalEntries) (RightHandSide3 rhsEntries) = + assertRightProperty $ do + matrixValue <- diagonalMatrix3 diagonalEntries + let rhsValues = U.fromList rhsEntries + expectedSolution = zipWith (/) rhsEntries diagonalEntries + resultValue = + solveSparseCG + (cgConfigWith 8 (IncompleteCholesky0SparsePreconditionerFamily (IC0Config Nothing))) + matrixValue + rhsValues + (U.replicate 3 0.0) + pure + ( case resultValue of + Right sparseResult -> + assertApproxList expectedSolution (U.toList (sparseSolution sparseResult)) + && trueResidualNorm matrixValue rhsValues (sparseSolution sparseResult) <= scgcTolerance (cgConfigWith 8 (IncompleteCholesky0SparsePreconditionerFamily (IC0Config Nothing))) + Left _ -> False + ) + +cgConfigWith :: Int -> SparsePreconditionerFamily -> SparseConjugateGradientConfig +cgConfigWith iterationLimit preconditionerFamily = + SparseConjugateGradientConfig + { scgcTolerance = approxTolerance, + scgcIterationLimit = iterationLimit, + scgcPreconditionerFamily = preconditionerFamily + } + +diagonalMatrix3 :: [Double] -> Either MoonlightError (SparseCSR Double) +diagonalMatrix3 diagonalEntries = + mkSparseCOO + 3 + 3 + (zipWith (\entryIndex entryValue -> (entryIndex, entryIndex, entryValue)) [0 ..] diagonalEntries) + >>= cooToCSR + +anchoredPathLaplacian :: Int -> Either MoonlightError (SparseCSR Double) +anchoredPathLaplacian dimension = + mkSparseCOO dimension dimension ((0, 0, 1.0) : concatMap anchoredPathEdgeEntries [0 .. dimension - 2]) + >>= cooToCSR + +anchoredPathEdgeEntries :: Int -> [(Int, Int, Double)] +anchoredPathEdgeEntries leftIndex = + let rightIndex = leftIndex + 1 + in [ (leftIndex, leftIndex, 1.0), + (leftIndex, rightIndex, -1.0), + (rightIndex, leftIndex, -1.0), + (rightIndex, rightIndex, 1.0) + ] + +anchoredPathRightHandSide :: Int -> U.Vector Double +anchoredPathRightHandSide dimension = + U.generate + dimension + ( \indexValue -> + let entryPhase = fromIntegral (indexValue + 1) + entrySkew = fromIntegral ((indexValue * 7) `mod` 11) + in 1.0 + sin entryPhase + 0.125 * entrySkew + ) + +trueResidualNorm :: SparseCSR Double -> U.Vector Double -> U.Vector Double -> Double +trueResidualNorm matrixValue rhsValues solutionValues = + either + (const (1.0 / 0.0)) + (\productValues -> vectorNorm (U.toList (U.zipWith (-) productValues rhsValues))) + (csrMatVecVector matrixValue solutionValues) diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness/Sparse.hs b/src-laws/Moonlight/LinAlg/Effect/Harness/Sparse.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness/Sparse.hs @@ -0,0 +1,161 @@ +module Moonlight.LinAlg.Effect.Harness.Sparse + ( cooCsrRoundTripLaw, + cooCscRoundTripLaw, + csrCscTransposeAgreementLaw, + csrMatVecAgreesWithDenseLaw, + canonicalCsrCombinesDuplicatesLaw, + graphLaplacianSymmetricRowSumsZeroLaw, + selfAdjointCsrRejectsAsymmetryLaw, + ) +where + +import Data.Vector.Unboxed qualified as U +import Moonlight.LinAlg + ( GraphEdge (..), + canonicalCSRFromEntries, + cooToCSC, + cooToCSR, + cooToDense, + cscToCSR, + cscToDense, + csrMatVecVector, + csrToCSC, + csrToDense, + fromListMatrix, + graphLaplacianCSR, + mkSparseCOO, + selfAdjointCSRLinearOperator, + toListMatrix, + ) +import Moonlight.LinAlg.Effect.Harness.Core (assertApproxList, assertRightProperty, matrix3VectorProduct) +import Test.Tasty.QuickCheck qualified as QC + +newtype SparseEntries3 = SparseEntries3 [(Int, Int, Double)] + deriving stock (Eq, Show) + +newtype DenseVector3 = DenseVector3 [Double] + deriving stock (Eq, Show) + +instance QC.Arbitrary SparseEntries3 where + arbitrary = + SparseEntries3 + <$> QC.listOf + ( (,,) + <$> QC.chooseInt (0, 2) + <*> QC.chooseInt (0, 2) + <*> (fromIntegral <$> QC.chooseInt (-5, 5)) + ) + +instance QC.Arbitrary DenseVector3 where + arbitrary = + DenseVector3 + <$> QC.vectorOf 3 (fromIntegral <$> QC.chooseInt (-5, 5)) + +cooCsrRoundTripLaw :: QC.Property +cooCsrRoundTripLaw = + QC.property cooCsrRoundTripLawProperty + +cooCscRoundTripLaw :: QC.Property +cooCscRoundTripLaw = + QC.property cooCscRoundTripLawProperty + +csrCscTransposeAgreementLaw :: QC.Property +csrCscTransposeAgreementLaw = + QC.property csrCscTransposeAgreementLawProperty + +csrMatVecAgreesWithDenseLaw :: QC.Property +csrMatVecAgreesWithDenseLaw = + QC.property csrMatVecAgreesWithDenseLawProperty + +cooCsrRoundTripLawProperty :: SparseEntries3 -> QC.Property +cooCsrRoundTripLawProperty (SparseEntries3 entries) = + assertRightProperty $ do + cooValue <- mkSparseCOO 3 3 entries + csrValue <- cooToCSR cooValue + originalDense <- cooToDense @3 @3 cooValue + roundTripDense <- csrToDense @3 @3 csrValue + pure (toListMatrix originalDense == toListMatrix roundTripDense) + +cooCscRoundTripLawProperty :: SparseEntries3 -> QC.Property +cooCscRoundTripLawProperty (SparseEntries3 entries) = + assertRightProperty $ do + cooValue <- mkSparseCOO 3 3 entries + cscValue <- cooToCSC cooValue + originalDense <- cooToDense @3 @3 cooValue + roundTripDense <- cscToDense @3 @3 cscValue + pure (toListMatrix originalDense == toListMatrix roundTripDense) + +csrCscTransposeAgreementLawProperty :: SparseEntries3 -> QC.Property +csrCscTransposeAgreementLawProperty (SparseEntries3 entries) = + assertRightProperty $ do + cooValue <- mkSparseCOO 3 3 entries + csrValue <- cooToCSR cooValue + cscValue <- csrToCSC csrValue + csrRoundTrip <- cscToCSR cscValue + originalDense <- csrToDense @3 @3 csrValue + roundTripDense <- csrToDense @3 @3 csrRoundTrip + pure (toListMatrix originalDense == toListMatrix roundTripDense) + +csrMatVecAgreesWithDenseLawProperty :: SparseEntries3 -> DenseVector3 -> QC.Property +csrMatVecAgreesWithDenseLawProperty (SparseEntries3 entries) (DenseVector3 vectorEntries) = + assertRightProperty $ do + cooValue <- mkSparseCOO 3 3 entries + csrValue <- cooToCSR cooValue + denseMatrix <- csrToDense @3 @3 csrValue + csrProduct <- csrMatVecVector csrValue (U.fromList vectorEntries) + pure (assertApproxList (matrix3VectorProduct (rows3 (toListMatrix denseMatrix)) vectorEntries) (U.toList csrProduct)) + +canonicalCsrCombinesDuplicatesLaw :: QC.Property +canonicalCsrCombinesDuplicatesLaw = + assertRightProperty $ do + csrValue <- + canonicalCSRFromEntries + 2 + 3 + ([(0, 1, 2.0), (0, 1, 3.0), (0, 2, 0.0), (1, 0, 5.0), (1, 0, -5.0), (1, 2, 4.0)] :: [(Int, Int, Double)]) + denseMatrix <- csrToDense @2 @3 csrValue + pure (toListMatrix denseMatrix == [0.0, 5.0, 0.0, 0.0, 0.0, 4.0]) + +graphLaplacianSymmetricRowSumsZeroLaw :: QC.Property +graphLaplacianSymmetricRowSumsZeroLaw = + assertRightProperty $ do + csrValue <- + graphLaplacianCSR + ["a", "b", "c"] + [GraphEdge "a" "b" 1.0, GraphEdge "b" "c" 2.0, GraphEdge "a" "c" 3.0] + denseMatrix <- csrToDense @3 @3 csrValue + let rowsValue = rows3 (toListMatrix denseMatrix) + pure (symmetricRows rowsValue && all (\rowValue -> assertApproxList [0.0] [sum rowValue]) rowsValue) + +selfAdjointCsrRejectsAsymmetryLaw :: QC.Property +selfAdjointCsrRejectsAsymmetryLaw = + assertRightProperty $ do + matrixValue <- fromListMatrix @3 @3 ([0.0, 1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0] :: [Double]) + let csrValue = cooToCSR =<< mkSparseCOO 3 3 [(0, 1, 1.0), (1, 2, 2.0)] + directValue = case csrValue of + Left _ -> False + Right value -> case selfAdjointCSRLinearOperator value of + Left _ -> True + Right _ -> False + pure (directValue && toListMatrix matrixValue == [0.0, 1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0]) + +rows3 :: [a] -> [[a]] +rows3 values = + case values of + [a00, a01, a02, a10, a11, a12, a20, a21, a22] -> + [[a00, a01, a02], [a10, a11, a12], [a20, a21, a22]] + _ -> [] + +symmetricRows :: [[Double]] -> Bool +symmetricRows rowsValue = + case rowsValue of + [[a00, a01, a02], [a10, a11, a12], [a20, a21, a22]] -> + and + [ a00 == a00, + assertApproxList [a01] [a10], + assertApproxList [a02] [a20], + a11 == a11, + assertApproxList [a12] [a21], + a22 == a22 + ] + _ -> False diff --git a/src-laws/Moonlight/LinAlg/Effect/Harness/Statics.hs b/src-laws/Moonlight/LinAlg/Effect/Harness/Statics.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Harness/Statics.hs @@ -0,0 +1,173 @@ +module Moonlight.LinAlg.Effect.Harness.Statics + ( networkDeclarationOrderInvariantLaw, + repeatedLoadsAccumulateLaw, + equilibriumAssemblyCanonicalOrderingLaw, + equilibriumSolutionResidualBoundedLaw, + unsupportedLoadProducesResidualViolationLaw, + ) +where + +import Data.Bifunctor (first) +import Data.List (permutations) +import Data.List.NonEmpty qualified as NonEmpty +import Data.Map.Strict qualified as Map +import Moonlight.LinAlg + ( Axis (..), + EquilibriumResult (..), + EquilibriumSolution (..), + EquilibriumViolation (..), + Vec3 (..), + assembleEquilibriumEquations, + checkEquilibrium, + compiledFoundationOrder, + compiledMemberOrder, + compiledNodeOrder, + joint, + load, + member, + mkMemberRef, + network, + networkNodeMap, + nodeLoad, + nodeRef, + supportOn, + mkSupportAxes, + ) +import Moonlight.LinAlg.Effect.Harness.Core (approxTolerance, assertApproxEqual, assertRightProperty) +import Test.Tasty.QuickCheck qualified as QC + +newtype VerticalLoad = VerticalLoad Double + deriving stock (Eq, Show) + +newtype LoadPair = LoadPair (Double, Double) + deriving stock (Eq, Show) + +instance QC.Arbitrary VerticalLoad where + arbitrary = + VerticalLoad . fromIntegral <$> QC.chooseInt (1, 20) + +instance QC.Arbitrary LoadPair where + arbitrary = + LoadPair + <$> ((,) <$> (fromIntegral <$> QC.chooseInt (-20, 20)) <*> (fromIntegral <$> QC.chooseInt (-20, 20))) + +networkDeclarationOrderInvariantLaw :: QC.Property +networkDeclarationOrderInvariantLaw = + QC.property networkDeclarationOrderInvariantLawProperty + +repeatedLoadsAccumulateLaw :: QC.Property +repeatedLoadsAccumulateLaw = + QC.property repeatedLoadsAccumulateLawProperty + +equilibriumAssemblyCanonicalOrderingLaw :: QC.Property +equilibriumAssemblyCanonicalOrderingLaw = + QC.property equilibriumAssemblyCanonicalOrderingLawProperty + +equilibriumSolutionResidualBoundedLaw :: QC.Property +equilibriumSolutionResidualBoundedLaw = + QC.property equilibriumSolutionResidualBoundedLawProperty + +unsupportedLoadProducesResidualViolationLaw :: QC.Property +unsupportedLoadProducesResidualViolationLaw = + QC.property unsupportedLoadProducesResidualViolationLawProperty + +networkDeclarationOrderInvariantLawProperty :: VerticalLoad -> QC.Property +networkDeclarationOrderInvariantLawProperty (VerticalLoad loadMagnitude) = + let declarations = + [ supportOn "a" (Vec3 0.0 0.0 0.0) (mkSupportAxes [AxisY]), + load "b" (Vec3 0.0 1.0 0.0) (Vec3 0.0 (negate loadMagnitude) 0.0), + member "a" "b" + ] + results = network <$> permutations declarations + in QC.counterexample (show results) (allEqual results) + +repeatedLoadsAccumulateLawProperty :: LoadPair -> QC.Property +repeatedLoadsAccumulateLawProperty (LoadPair (firstLoad, secondLoad)) = + case (nodeRef "p", network declarations) of + (Right pointRef, Right networkValue) -> + case Map.lookup pointRef (networkNodeMap networkValue) of + Just nodeValue -> + QC.property (nodeLoad nodeValue == Vec3 (firstLoad + secondLoad) 0.0 0.0) + Nothing -> QC.counterexample "missing generated node p" False + other -> QC.counterexample (show other) False + where + declarations = + [ load "p" (Vec3 0.0 0.0 0.0) (Vec3 firstLoad 0.0 0.0), + load "p" (Vec3 0.0 0.0 0.0) (Vec3 secondLoad 0.0 0.0) + ] + +equilibriumAssemblyCanonicalOrderingLawProperty :: VerticalLoad -> QC.Property +equilibriumAssemblyCanonicalOrderingLawProperty (VerticalLoad loadMagnitude) = + assertRightProperty $ do + networkValue <- + mapLeftShow $ + network + [ member "c" "a", + joint "c" (Vec3 0.0 1.0 0.0), + supportOn "a" (Vec3 (-1.0) 0.0 0.0) (mkSupportAxes [AxisY]), + member "b" "c", + supportOn "b" (Vec3 1.0 0.0 0.0) (mkSupportAxes [AxisY]), + load "c" (Vec3 0.0 1.0 0.0) (Vec3 0.0 (negate loadMagnitude) 0.0) + ] + compiledValue <- mapLeftShow (assembleEquilibriumEquations networkValue) + nodeA <- mapLeftShow (nodeRef "a") + nodeB <- mapLeftShow (nodeRef "b") + nodeC <- mapLeftShow (nodeRef "c") + leftMember <- mapLeftShow (mkMemberRef nodeA nodeC) + rightMember <- mapLeftShow (mkMemberRef nodeB nodeC) + pure + ( compiledNodeOrder compiledValue == [nodeA, nodeB, nodeC] + && compiledFoundationOrder compiledValue == [nodeA, nodeB] + && compiledMemberOrder compiledValue == [leftMember, rightMember] + ) + +equilibriumSolutionResidualBoundedLawProperty :: VerticalLoad -> QC.Property +equilibriumSolutionResidualBoundedLawProperty (VerticalLoad loadMagnitude) = + assertRightProperty $ do + networkValue <- + mapLeftShow $ + network + [ load "load" (Vec3 0.0 1.0 0.0) (Vec3 0.0 (negate loadMagnitude) 0.0), + supportOn "foundation" (Vec3 0.0 0.0 0.0) (mkSupportAxes [AxisY]), + member "foundation" "load" + ] + equilibriumResult <- mapLeftShow (checkEquilibrium networkValue) + pure + ( case equilibriumResult of + InEquilibrium solutionValue -> + all residualBounded (Map.elems (equilibriumResidualForces solutionValue)) + Disequilibrium _ -> False + ) + +unsupportedLoadProducesResidualViolationLawProperty :: VerticalLoad -> QC.Property +unsupportedLoadProducesResidualViolationLawProperty (VerticalLoad loadMagnitude) = + assertRightProperty $ do + networkValue <- + mapLeftShow $ + network + [ supportOn "foundation" (Vec3 0.0 0.0 0.0) (mkSupportAxes [AxisY]), + load "load" (Vec3 1.0 0.0 0.0) (Vec3 0.0 (negate loadMagnitude) 0.0), + member "foundation" "load" + ] + equilibriumResult <- mapLeftShow (checkEquilibrium networkValue) + pure + ( case equilibriumResult of + InEquilibrium _ -> False + Disequilibrium violations -> + any ((> approxTolerance) . violationResidualMagnitude) (NonEmpty.toList violations) + ) + +residualBounded :: Vec3 -> Bool +residualBounded (Vec3 xValue yValue zValue) = + assertApproxEqual 0.0 xValue + && assertApproxEqual 0.0 yValue + && assertApproxEqual 0.0 zValue + +allEqual :: Eq value => [value] -> Bool +allEqual values = + case values of + [] -> True + firstValue : remainingValues -> all (== firstValue) remainingValues + +mapLeftShow :: Show failure => Either failure value -> Either String value +mapLeftShow = first show diff --git a/src-laws/Moonlight/LinAlg/Effect/LawNames.hs b/src-laws/Moonlight/LinAlg/Effect/LawNames.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/LawNames.hs @@ -0,0 +1,160 @@ +{-# LANGUAGE DerivingStrategies #-} + +module Moonlight.LinAlg.Effect.LawNames + ( LawName (..), + allLawNames, + lawName, + ) +where + +import Data.Kind (Type) +import Moonlight.Core (IsLawName (..), constructorLawNameWithOverrides) + +type LawName :: Type +data LawName + = DenseAddAssociative + | DenseAddCommutative + | DenseMultiplyAssociative + | DenseLeftDistributive + | DenseRightDistributive + | DenseTransposeInvolution + | DenseTransposeProductReversal + | DenseMapComposition + | QRReconstructsInput + | QROrthonormalColumns + | CholeskyReconstructsSPD + | SymmetricEigenReconstructs + | SymmetricEigenOrthonormal + | SymmetricEigenUncheckedPassesCertification + | ThinSVDReconstructs + | ThinSVDOrthonormalFactors + | ThinSVDSingularValuesOrderedNonnegative + | PLUReconstructsInput + | RankKernelNullity + | KernelVectorsAnnihilated + | PackedLinearMapIdentity + | PackedLinearMapComposition + | GF2PackedInverseTwoSided + | SmithDiagonalReconstructsInput + | SmithDivisibilityChain + | SmithWitnessesUnimodular + | SmithDiagonalOnlyAgreesWithFull + | BareissRankAgreesWithFieldRank + | BareissDeterminantAgreesWithExterior + | COOCSRRoundTrip + | COOCSCRoundTrip + | CSRCSCTransposeAgreement + | CSRMatVecAgreesWithDense + | CanonicalCSRCombinesDuplicates + | GraphLaplacianSymmetricRowSumsZero + | SelfAdjointCSRRejectsAsymmetry + | ScaledOperatorAction + | ShiftedIdentityAction + | IC0FactorSolveRoundTrip + | IC0RejectsNonpositivePivot + | PreconditionedCGConvergesOnSPD + | ArnoldiRelationHolds + | ArnoldiBasisOrthonormal + | LanczosProjectionTridiagonal + | LanczosBasisOrthonormal + | ThickRestartLockedPairsResidualBounded + | SelectedPairsResidualBounded + | SelectedPairsClusterOrthonormal + | TridiagonalSelectedValuesAgreeWithAllPairs + | DiagonalSpectralValuesExact + | PathLaplacianSpectralValuesClosedForm + | EigenRequestRejectsOversubscription + | Vec3AddCommutativeAssociative + | Vec3DotSymmetric + | Vec3NormalizeUnit + | AABBUnionCommutativeAssociative + | AABBUnionContainsOperands + | AABBIntersectionCommutative + | SymmetricOuterApplyAgreement + | GeometrySymmetricEigenReconstructs + | GeometrySymmetricEigenOrthonormal + | NetworkDeclarationOrderInvariant + | RepeatedLoadsAccumulate + | EquilibriumAssemblyCanonicalOrdering + | EquilibriumSolutionResidualBounded + | UnsupportedLoadProducesResidualViolation + deriving stock (Bounded, Enum, Eq, Ord, Show) + +instance IsLawName LawName where + lawNameText = lawName + +allLawNames :: [LawName] +allLawNames = + [minBound .. maxBound] + +lawName :: LawName -> String +lawName = + constructorLawNameWithOverrides + [ ("DenseAddAssociative", "linalg_dense_add_associative"), + ("DenseAddCommutative", "linalg_dense_add_commutative"), + ("DenseMultiplyAssociative", "linalg_dense_multiply_associative"), + ("DenseLeftDistributive", "linalg_dense_left_distributive"), + ("DenseRightDistributive", "linalg_dense_right_distributive"), + ("DenseTransposeInvolution", "linalg_dense_transpose_involution"), + ("DenseTransposeProductReversal", "linalg_dense_transpose_product_reversal"), + ("DenseMapComposition", "linalg_dense_map_composition"), + ("QRReconstructsInput", "linalg_decomposition_qr_reconstructs_input"), + ("QROrthonormalColumns", "linalg_decomposition_qr_orthonormal_columns"), + ("CholeskyReconstructsSPD", "linalg_decomposition_cholesky_reconstructs_spd"), + ("SymmetricEigenReconstructs", "linalg_decomposition_symmetric_eigen_reconstructs"), + ("SymmetricEigenOrthonormal", "linalg_decomposition_symmetric_eigen_orthonormal"), + ("SymmetricEigenUncheckedPassesCertification", "linalg_decomposition_symmetric_eigen_unchecked_passes_certification"), + ("ThinSVDReconstructs", "linalg_decomposition_thin_svd_reconstructs"), + ("ThinSVDOrthonormalFactors", "linalg_decomposition_thin_svd_orthonormal_factors"), + ("ThinSVDSingularValuesOrderedNonnegative", "linalg_decomposition_thin_svd_singular_values_ordered_nonnegative"), + ("PLUReconstructsInput", "linalg_field_plu_reconstructs_input"), + ("RankKernelNullity", "linalg_field_rank_kernel_nullity"), + ("KernelVectorsAnnihilated", "linalg_field_kernel_vectors_annihilated"), + ("PackedLinearMapIdentity", "linalg_gf2_packed_linear_map_identity"), + ("PackedLinearMapComposition", "linalg_gf2_packed_linear_map_composition"), + ("GF2PackedInverseTwoSided", "linalg_gf2_packed_inverse_two_sided"), + ("SmithDiagonalReconstructsInput", "linalg_domain_smith_diagonal_reconstructs_input"), + ("SmithDivisibilityChain", "linalg_domain_smith_divisibility_chain"), + ("SmithWitnessesUnimodular", "linalg_domain_smith_witnesses_unimodular"), + ("SmithDiagonalOnlyAgreesWithFull", "linalg_domain_smith_diagonal_only_agrees_with_full"), + ("BareissRankAgreesWithFieldRank", "linalg_domain_bareiss_rank_agrees_with_field_rank"), + ("BareissDeterminantAgreesWithExterior", "linalg_domain_bareiss_determinant_agrees_with_exterior"), + ("COOCSRRoundTrip", "linalg_sparse_coo_csr_round_trip"), + ("COOCSCRoundTrip", "linalg_sparse_coo_csc_round_trip"), + ("CSRCSCTransposeAgreement", "linalg_sparse_csr_csc_transpose_agreement"), + ("CSRMatVecAgreesWithDense", "linalg_sparse_csr_matvec_agrees_with_dense"), + ("CanonicalCSRCombinesDuplicates", "linalg_sparse_canonical_csr_combines_duplicates"), + ("GraphLaplacianSymmetricRowSumsZero", "linalg_sparse_graph_laplacian_symmetric_row_sums_zero"), + ("SelfAdjointCSRRejectsAsymmetry", "linalg_sparse_self_adjoint_csr_rejects_asymmetry"), + ("ScaledOperatorAction", "linalg_operator_scaled_operator_action"), + ("ShiftedIdentityAction", "linalg_operator_shifted_identity_action"), + ("IC0FactorSolveRoundTrip", "linalg_preconditioner_ic0_factor_solve_round_trip"), + ("IC0RejectsNonpositivePivot", "linalg_preconditioner_ic0_rejects_nonpositive_pivot"), + ("PreconditionedCGConvergesOnSPD", "linalg_preconditioner_preconditioned_cg_converges_on_spd"), + ("ArnoldiRelationHolds", "linalg_krylov_spectral_arnoldi_relation_holds"), + ("ArnoldiBasisOrthonormal", "linalg_krylov_spectral_arnoldi_basis_orthonormal"), + ("LanczosProjectionTridiagonal", "linalg_krylov_spectral_lanczos_projection_tridiagonal"), + ("LanczosBasisOrthonormal", "linalg_krylov_spectral_lanczos_basis_orthonormal"), + ("ThickRestartLockedPairsResidualBounded", "linalg_krylov_spectral_thick_restart_locked_pairs_residual_bounded"), + ("SelectedPairsResidualBounded", "linalg_krylov_spectral_selected_pairs_residual_bounded"), + ("SelectedPairsClusterOrthonormal", "linalg_krylov_spectral_selected_pairs_cluster_orthonormal"), + ("TridiagonalSelectedValuesAgreeWithAllPairs", "linalg_krylov_spectral_tridiagonal_selected_values_agree_with_all_pairs"), + ("DiagonalSpectralValuesExact", "linalg_krylov_spectral_diagonal_spectral_values_exact"), + ("PathLaplacianSpectralValuesClosedForm", "linalg_krylov_spectral_path_laplacian_spectral_values_closed_form"), + ("EigenRequestRejectsOversubscription", "linalg_krylov_spectral_eigen_request_rejects_oversubscription"), + ("Vec3AddCommutativeAssociative", "linalg_geometry_vec3_add_commutative_associative"), + ("Vec3DotSymmetric", "linalg_geometry_vec3_dot_symmetric"), + ("Vec3NormalizeUnit", "linalg_geometry_vec3_normalize_unit"), + ("AABBUnionCommutativeAssociative", "linalg_geometry_aabb_union_commutative_associative"), + ("AABBUnionContainsOperands", "linalg_geometry_aabb_union_contains_operands"), + ("AABBIntersectionCommutative", "linalg_geometry_aabb_intersection_commutative"), + ("SymmetricOuterApplyAgreement", "linalg_geometry_symmetric_outer_apply_agreement"), + ("GeometrySymmetricEigenReconstructs", "linalg_geometry_symmetric_eigen_reconstructs"), + ("GeometrySymmetricEigenOrthonormal", "linalg_geometry_symmetric_eigen_orthonormal"), + ("NetworkDeclarationOrderInvariant", "linalg_statics_network_declaration_order_invariant"), + ("RepeatedLoadsAccumulate", "linalg_statics_repeated_loads_accumulate"), + ("EquilibriumAssemblyCanonicalOrdering", "linalg_statics_equilibrium_assembly_canonical_ordering"), + ("EquilibriumSolutionResidualBounded", "linalg_statics_equilibrium_solution_residual_bounded"), + ("UnsupportedLoadProducesResidualViolation", "linalg_statics_unsupported_load_produces_residual_violation") + ] + . show diff --git a/src-laws/Moonlight/LinAlg/Effect/Laws.hs b/src-laws/Moonlight/LinAlg/Effect/Laws.hs new file mode 100644 --- /dev/null +++ b/src-laws/Moonlight/LinAlg/Effect/Laws.hs @@ -0,0 +1,161 @@ +module Moonlight.LinAlg.Effect.Laws + ( linalgLawSuites, + registeredLawNames, + tests, + ) +where + +import Data.List (nub, (\\)) +import Moonlight.LinAlg.Effect.Harness qualified as Harness +import Moonlight.LinAlg.Effect.LawNames (LawName (..), allLawNames) +import Moonlight.Pale.Test.Laws.Suite + ( LawSuite, + hUnitLaw, + lawGroup, + namedQuickCheckLaw, + renderLawSuite, + ) +import Test.Tasty (TestTree) +import Test.Tasty.HUnit (assertBool) +import Test.Tasty.QuickCheck qualified as QC + +tests :: TestTree +tests = + renderLawSuite (lawGroup "moonlight-linalg laws" linalgLawSuites) + +data RegisteredLaw = RegisteredLaw !LawName !QC.Property + +data LawSection = LawSection !String ![RegisteredLaw] + +linalgLawSuites :: [LawSuite] +linalgLawSuites = + lawGroup + "manifest" + [hUnitLaw "linalg_law_manifest_totality" manifestTotalityAssertion] + : (lawSectionSuite <$> registeredLawSections) + +registeredLawNames :: [LawName] +registeredLawNames = + concatMap + (\(LawSection _ registrations) -> registeredLawName <$> registrations) + registeredLawSections + +lawSectionSuite :: LawSection -> LawSuite +lawSectionSuite (LawSection sectionName registrations) = + lawGroup sectionName (registeredLawSuite <$> registrations) + +registeredLawSuite :: RegisteredLaw -> LawSuite +registeredLawSuite (RegisteredLaw lawName lawProperty) = + namedQuickCheckLaw lawName lawProperty + +registeredLawName :: RegisteredLaw -> LawName +registeredLawName (RegisteredLaw lawName _) = lawName + +registeredLawSections :: [LawSection] +registeredLawSections = + [ LawSection + "dense algebra" + [ RegisteredLaw DenseAddAssociative Harness.denseAddAssociativeLaw, + RegisteredLaw DenseAddCommutative Harness.denseAddCommutativeLaw, + RegisteredLaw DenseMultiplyAssociative Harness.denseMultiplyAssociativeLaw, + RegisteredLaw DenseLeftDistributive Harness.denseLeftDistributiveLaw, + RegisteredLaw DenseRightDistributive Harness.denseRightDistributiveLaw, + RegisteredLaw DenseTransposeInvolution Harness.denseTransposeInvolutionLaw, + RegisteredLaw DenseTransposeProductReversal Harness.denseTransposeProductReversalLaw, + RegisteredLaw DenseMapComposition Harness.denseMapCompositionLaw + ], + LawSection + "decompositions" + [ RegisteredLaw QRReconstructsInput Harness.qrReconstructsInputLaw, + RegisteredLaw QROrthonormalColumns Harness.qrOrthonormalColumnsLaw, + RegisteredLaw CholeskyReconstructsSPD Harness.choleskyReconstructsSpdLaw, + RegisteredLaw SymmetricEigenReconstructs Harness.symmetricEigenReconstructsLaw, + RegisteredLaw SymmetricEigenOrthonormal Harness.symmetricEigenOrthonormalLaw, + RegisteredLaw SymmetricEigenUncheckedPassesCertification Harness.symmetricEigenUncheckedPassesCertificationLaw, + RegisteredLaw ThinSVDReconstructs Harness.thinSvdReconstructsLaw, + RegisteredLaw ThinSVDOrthonormalFactors Harness.thinSvdOrthonormalFactorsLaw, + RegisteredLaw ThinSVDSingularValuesOrderedNonnegative Harness.thinSvdSingularValuesOrderedNonnegativeLaw + ], + LawSection + "field and gf2" + [ RegisteredLaw PLUReconstructsInput Harness.pluReconstructsInputLaw, + RegisteredLaw RankKernelNullity Harness.rankKernelNullityLaw, + RegisteredLaw KernelVectorsAnnihilated Harness.kernelVectorsAnnihilatedLaw, + RegisteredLaw PackedLinearMapIdentity Harness.packedLinearMapIdentityLaw, + RegisteredLaw PackedLinearMapComposition Harness.packedLinearMapCompositionLaw, + RegisteredLaw GF2PackedInverseTwoSided Harness.gf2PackedInverseTwoSidedLaw + ], + LawSection + "domain" + [ RegisteredLaw SmithDiagonalReconstructsInput Harness.smithDiagonalReconstructsInputLaw, + RegisteredLaw SmithDivisibilityChain Harness.smithDivisibilityChainLaw, + RegisteredLaw SmithWitnessesUnimodular Harness.smithWitnessesUnimodularLaw, + RegisteredLaw SmithDiagonalOnlyAgreesWithFull Harness.smithDiagonalOnlyAgreesWithFullLaw, + RegisteredLaw BareissRankAgreesWithFieldRank Harness.bareissRankAgreesWithFieldRankLaw, + RegisteredLaw BareissDeterminantAgreesWithExterior Harness.bareissDeterminantAgreesWithExteriorLaw + ], + LawSection + "sparse storage" + [ RegisteredLaw COOCSRRoundTrip Harness.cooCsrRoundTripLaw, + RegisteredLaw COOCSCRoundTrip Harness.cooCscRoundTripLaw, + RegisteredLaw CSRCSCTransposeAgreement Harness.csrCscTransposeAgreementLaw, + RegisteredLaw CSRMatVecAgreesWithDense Harness.csrMatVecAgreesWithDenseLaw, + RegisteredLaw CanonicalCSRCombinesDuplicates Harness.canonicalCsrCombinesDuplicatesLaw, + RegisteredLaw GraphLaplacianSymmetricRowSumsZero Harness.graphLaplacianSymmetricRowSumsZeroLaw, + RegisteredLaw SelfAdjointCSRRejectsAsymmetry Harness.selfAdjointCsrRejectsAsymmetryLaw + ], + LawSection + "operator" + [ RegisteredLaw ScaledOperatorAction Harness.scaledOperatorActionLaw, + RegisteredLaw ShiftedIdentityAction Harness.shiftedIdentityActionLaw + ], + LawSection + "preconditioner" + [ RegisteredLaw IC0FactorSolveRoundTrip Harness.ic0FactorSolveRoundTripLaw, + RegisteredLaw IC0RejectsNonpositivePivot Harness.ic0RejectsNonpositivePivotLaw, + RegisteredLaw PreconditionedCGConvergesOnSPD Harness.preconditionedCgConvergesOnSpdLaw + ], + LawSection + "krylov spectral" + [ RegisteredLaw ArnoldiRelationHolds Harness.arnoldiRelationHoldsLaw, + RegisteredLaw ArnoldiBasisOrthonormal Harness.arnoldiBasisOrthonormalLaw, + RegisteredLaw LanczosProjectionTridiagonal Harness.lanczosProjectionTridiagonalLaw, + RegisteredLaw LanczosBasisOrthonormal Harness.lanczosBasisOrthonormalLaw, + RegisteredLaw ThickRestartLockedPairsResidualBounded Harness.thickRestartLockedPairsResidualBoundedLaw, + RegisteredLaw SelectedPairsResidualBounded Harness.selectedPairsResidualBoundedLaw, + RegisteredLaw SelectedPairsClusterOrthonormal Harness.selectedPairsClusterOrthonormalLaw, + RegisteredLaw TridiagonalSelectedValuesAgreeWithAllPairs Harness.tridiagonalSelectedValuesAgreeWithAllPairsLaw, + RegisteredLaw DiagonalSpectralValuesExact Harness.diagonalSpectralValuesExactLaw, + RegisteredLaw PathLaplacianSpectralValuesClosedForm Harness.pathLaplacianSpectralValuesClosedFormLaw, + RegisteredLaw EigenRequestRejectsOversubscription Harness.eigenRequestRejectsOversubscriptionLaw + ], + LawSection + "geometry" + [ RegisteredLaw Vec3AddCommutativeAssociative Harness.vec3AddCommutativeAssociativeLaw, + RegisteredLaw Vec3DotSymmetric Harness.vec3DotSymmetricLaw, + RegisteredLaw Vec3NormalizeUnit Harness.vec3NormalizeUnitLaw, + RegisteredLaw AABBUnionCommutativeAssociative Harness.aabbUnionCommutativeAssociativeLaw, + RegisteredLaw AABBUnionContainsOperands Harness.aabbUnionContainsOperandsLaw, + RegisteredLaw AABBIntersectionCommutative Harness.aabbIntersectionCommutativeLaw, + RegisteredLaw SymmetricOuterApplyAgreement Harness.symmetricOuterApplyAgreementLaw, + RegisteredLaw GeometrySymmetricEigenReconstructs Harness.geometrySymmetricEigenReconstructsLaw, + RegisteredLaw GeometrySymmetricEigenOrthonormal Harness.geometrySymmetricEigenOrthonormalLaw + ], + LawSection + "statics" + [ RegisteredLaw NetworkDeclarationOrderInvariant Harness.networkDeclarationOrderInvariantLaw, + RegisteredLaw RepeatedLoadsAccumulate Harness.repeatedLoadsAccumulateLaw, + RegisteredLaw EquilibriumAssemblyCanonicalOrdering Harness.equilibriumAssemblyCanonicalOrderingLaw, + RegisteredLaw EquilibriumSolutionResidualBounded Harness.equilibriumSolutionResidualBoundedLaw, + RegisteredLaw UnsupportedLoadProducesResidualViolation Harness.unsupportedLoadProducesResidualViolationLaw + ] + ] + +manifestTotalityAssertion :: IO () +manifestTotalityAssertion = do + let missing = allLawNames \\ registeredLawNames + unrecognized = registeredLawNames \\ allLawNames + duplicateCount = length registeredLawNames - length (nub registeredLawNames) + assertBool ("missing law registrations: " <> show missing) (null missing) + assertBool ("unrecognized law registrations: " <> show unrecognized) (null unrecognized) + assertBool ("duplicate law registrations: " <> show duplicateCount) (duplicateCount == 0) diff --git a/src-native/Moonlight/LinAlg/Effect/Native/Dispatch.hs b/src-native/Moonlight/LinAlg/Effect/Native/Dispatch.hs new file mode 100644 --- /dev/null +++ b/src-native/Moonlight/LinAlg/Effect/Native/Dispatch.hs @@ -0,0 +1,541 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE GADTs #-} + +-- | Native spectral dispatch over LAPACK result sections. +module Moonlight.LinAlg.Effect.Native.Dispatch + ( denseDoubleLinearSolveLapack, + denseDoubleMatrixProductBlas, + denseDoubleSymmetricEigenpairsLapack, + leastSquaresLapack, + symmetricEigenRequestLapack, + selectedSymmetricTridiagonalEigenRequestLapack, + selectedSymmetricBlockTridiagonalEigenRequestLapack, + ) +where + +import Data.Bifunctor (first) +import Data.Vector.Storable qualified as S +import Data.Vector.Unboxed qualified as U +import Moonlight.Core + ( MoonlightError (..), + checkedNonNegativeProduct, + ) +import Moonlight.LinAlg.Effect.Native.LAPACK + ( denseDoubleLinearSolveLapack, + denseDoubleMatrixProductBlas, + denseDoubleSymmetricEigenpairsRawLapack, + leastSquaresLapack, + selectedSymmetricEigenPairsLapack, + selectedSymmetricEigenValuesLapack, + selectedSymmetricBlockTridiagonalEigenPairsLapack, + selectedSymmetricBlockTridiagonalEigenValuesLapack, + selectedSymmetricTridiagonalEigenPairsLapack, + selectedSymmetricTridiagonalEigenValuesLapack, + ) +import Moonlight.LinAlg.Internal.Eigen.Kernels (epsDouble, finiteDouble) +import Moonlight.LinAlg.Internal.Eigen.Symmetric + ( CertifiedSymmetricEigenResult (..), + SymmetricEigenCertificationFailure, + SymmetricEigenResult (..), + certifySymmetricEigenResult, + ) +import Moonlight.LinAlg.Internal.VectorOps (normU, scaleU, subU) +import Moonlight.LinAlg.Pure.Dense.Dynamic + ( DynMatrix, + dynMatrixShape, + dynMatrixToList, + ) +import Moonlight.LinAlg.Pure.Dense.Flat + ( DenseDoubleMatrix, + denseDoubleMatrixShape, + denseDoubleMatrixToRowMajorVector, + trustedDenseDoubleMatrixRowMajor, + ) +import Moonlight.LinAlg.Pure.Krylov.Config (positiveCountValue) +import Moonlight.LinAlg.Pure.Krylov.Selection + ( SpectrumEnd (..), + sortRawPairsForSpectrum, + ) +import Moonlight.LinAlg.Pure.Spectral.Request (EigenRequest (..)) +import Moonlight.LinAlg.Pure.Spectral.Result + ( Eigenpairs, + mkEigenpairs, + ) +import Moonlight.LinAlg.Pure.Structured.BlockTridiagonal + ( SymmetricBlockTridiagonal, + applySymmetricBlockTridiagonalU, + symmetricBlockTridiagonalDimension, + symmetricBlockTridiagonalFrobeniusNorm, + ) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal + ( SymmetricTridiagonal, + symmetricTridiagonalDiagonalEntries, + symmetricTridiagonalDimension, + symmetricTridiagonalOffDiagonalEntries, + ) +import Prelude + +symmetricEigenRequestLapack :: + EigenRequest result -> + DynMatrix Double -> + IO (Either MoonlightError result) +symmetricEigenRequestLapack requestValue matrixValue = + case requestValue of + EigenvaluesRequest spectrumEnd countValue -> + fmap (fmap (orderNativeValues spectrumEnd)) $ + selectedSymmetricEigenValuesLapack spectrumEnd (positiveCountValue countValue) matrixValue + EigenpairsRequest spectrumEnd countValue -> + fmap (>>= denseEigenpairsFromRawColumns spectrumEnd matrixValue) $ + selectedSymmetricEigenPairsLapack spectrumEnd (positiveCountValue countValue) matrixValue + +selectedSymmetricTridiagonalEigenRequestLapack :: + EigenRequest result -> + SymmetricTridiagonal -> + IO (Either MoonlightError result) +selectedSymmetricTridiagonalEigenRequestLapack requestValue tridiagonalValue = + case requestValue of + EigenvaluesRequest spectrumEnd countValue -> + fmap (fmap (orderNativeValues spectrumEnd)) $ + selectedSymmetricTridiagonalEigenValuesLapack spectrumEnd (positiveCountValue countValue) tridiagonalValue + EigenpairsRequest spectrumEnd countValue -> + fmap (>>= tridiagonalEigenpairsFromRawPairs spectrumEnd tridiagonalValue) $ + selectedSymmetricTridiagonalEigenPairsLapack spectrumEnd (positiveCountValue countValue) tridiagonalValue + +selectedSymmetricBlockTridiagonalEigenRequestLapack :: + EigenRequest result -> + SymmetricBlockTridiagonal -> + IO (Either MoonlightError result) +selectedSymmetricBlockTridiagonalEigenRequestLapack requestValue blockValue = + case requestValue of + EigenvaluesRequest spectrumEnd countValue -> + fmap (fmap (orderNativeValues spectrumEnd)) $ + selectedSymmetricBlockTridiagonalEigenValuesLapack spectrumEnd (positiveCountValue countValue) blockValue + EigenpairsRequest spectrumEnd countValue -> + fmap (>>= blockTridiagonalEigenpairsFromRawPairs spectrumEnd blockValue) $ + selectedSymmetricBlockTridiagonalEigenPairsLapack spectrumEnd (positiveCountValue countValue) blockValue + +denseDoubleSymmetricEigenpairsLapack :: DenseDoubleMatrix -> IO (Either MoonlightError Eigenpairs) +denseDoubleSymmetricEigenpairsLapack matrixValue = + fmap (>>= denseDoubleEigenpairsFromRawColumns matrixValue) $ + denseDoubleSymmetricEigenpairsRawLapack matrixValue + +denseDoubleEigenpairsFromRawColumns :: + DenseDoubleMatrix -> + (S.Vector Double, S.Vector Double) -> + Either MoonlightError Eigenpairs +denseDoubleEigenpairsFromRawColumns matrixValue rawColumns = do + certifiedResult <- denseCertifiedEigenResultFromRawColumns matrixValue rawColumns + let resultValue = certifiedSymmetricEigenResult certifiedResult + !dimension = fst (denseDoubleMatrixShape matrixValue) + !eigenvalues = + storableVectorToUnboxed + (symmetricEigenResultValues resultValue) + !eigenvectors = + denseEigenvectorsColumnMajorUnboxed + dimension + (symmetricEigenResultVectors resultValue) + !residuals = + denseCertifiedPairResiduals + dimension + matrixValue + eigenvalues + eigenvectors + mkEigenpairs + dimension + eigenvalues + eigenvectors + residuals + +denseCertifiedEigenResultFromRawColumns :: + DenseDoubleMatrix -> + (S.Vector Double, S.Vector Double) -> + Either MoonlightError CertifiedSymmetricEigenResult +denseCertifiedEigenResultFromRawColumns matrixValue (rawEigenvalues, rawEigenvectors) = do + let (rowCount, columnCount) = denseDoubleMatrixShape matrixValue + if rowCount /= columnCount + then Left (InvariantViolation "native dense Double eigenpairs require a square matrix") + else do + expectedVectorCount <- + checkedColumnPayloadLength + rowCount + (S.length rawEigenvalues) + if S.length rawEigenvectors /= expectedVectorCount + then + Left + ( InvariantViolation + ( "native dense Double eigenvector payload mismatch: expected " + <> show expectedVectorCount + <> " entries but received " + <> show (S.length rawEigenvectors) + ) + ) + else + case certifySymmetricEigenResult matrixValue (rawSymmetricEigenResult rowCount rawEigenvalues rawEigenvectors) of + Left failureValue -> + Left + ( nativeCertificationFailure + "native dense Double symmetric eigensolve" + failureValue + ) + Right certifiedResult -> + Right certifiedResult + +rawSymmetricEigenResult :: + Int -> + S.Vector Double -> + S.Vector Double -> + SymmetricEigenResult +rawSymmetricEigenResult dimension eigenvalues eigenvectors = + SymmetricEigenResult + { symmetricEigenResultValues = eigenvalues, + symmetricEigenResultVectors = + trustedDenseDoubleMatrixRowMajor + dimension + dimension + (lapackColumnMajorEigenvectorsToRowMajor dimension eigenvectors) + } + +lapackColumnMajorEigenvectorsToRowMajor :: + Int -> + S.Vector Double -> + S.Vector Double +lapackColumnMajorEigenvectorsToRowMajor dimension eigenvectors = + S.generate + (S.length eigenvectors) + ( \payloadIndex -> + let (!rowIndex, !columnIndex) = payloadIndex `quotRem` dimension + in eigenvectors `S.unsafeIndex` (columnIndex * dimension + rowIndex) + ) + +denseEigenvectorsColumnMajorUnboxed :: + Int -> + DenseDoubleMatrix -> + U.Vector Double +denseEigenvectorsColumnMajorUnboxed dimension eigenvectors = + U.generate + (S.length eigenvectorPayload) + ( \payloadIndex -> + let (!columnIndex, !rowIndex) = payloadIndex `quotRem` dimension + in eigenvectorPayload `S.unsafeIndex` (rowIndex * dimension + columnIndex) + ) + where + eigenvectorPayload = denseDoubleMatrixToRowMajorVector eigenvectors + +denseCertifiedPairResiduals :: + Int -> + DenseDoubleMatrix -> + U.Vector Double -> + U.Vector Double -> + U.Vector Double +denseCertifiedPairResiduals dimension matrixValue eigenvalues eigenvectors = + let !matrixPayload = + storableVectorToUnboxed + (denseDoubleMatrixToRowMajorVector matrixValue) + in U.generate + (U.length eigenvalues) + ( denseResidualNormAt + dimension + matrixPayload + eigenvalues + eigenvectors + ) + +nativeCertificationFailure :: + String -> + SymmetricEigenCertificationFailure -> + MoonlightError +nativeCertificationFailure context failureValue = + InvariantViolation + ( context + <> " certification failed: " + <> show failureValue + ) + +denseEigenpairsFromRawColumns :: + SpectrumEnd -> + DynMatrix Double -> + (U.Vector Double, U.Vector Double) -> + Either MoonlightError Eigenpairs +denseEigenpairsFromRawColumns + spectrumEnd + matrixValue + rawColumns = do + let !dimension = dynMatrixDimension matrixValue + !matrixPayload = U.fromList (dynMatrixToList matrixValue) + (eigenvalues, eigenvectors) <- + orderNativeColumns + spectrumEnd + dimension + rawColumns + let !pairCount = U.length eigenvalues + !residuals = + U.generate + pairCount + ( denseResidualNormAt + dimension + matrixPayload + eigenvalues + eigenvectors + ) + validateNativeResidualNorms + "native dense symmetric eigensolve" + dimension + (frobeniusNormU matrixPayload) + residuals + mkEigenpairs + dimension + eigenvalues + eigenvectors + residuals + +orderNativeColumns :: + SpectrumEnd -> + Int -> + (U.Vector Double, U.Vector Double) -> + Either + MoonlightError + (U.Vector Double, U.Vector Double) +orderNativeColumns spectrumEnd dimension (eigenvalues, eigenvectors) = do + expectedVectorCount <- + checkedColumnPayloadLength + dimension + (U.length eigenvalues) + if U.length eigenvectors /= expectedVectorCount + then + Left + ( InvariantViolation + ( "native eigenvector payload mismatch: expected " + <> show expectedVectorCount + <> " entries but received " + <> show (U.length eigenvectors) + ) + ) + else + case spectrumEnd of + SmallestEigenvalues -> Right (eigenvalues, eigenvectors) + LargestEigenvalues -> + Right + ( U.reverse eigenvalues, + reverseEigenvectorColumns + dimension + (U.length eigenvalues) + eigenvectors + ) + +checkedColumnPayloadLength :: + Int -> + Int -> + Either MoonlightError Int +checkedColumnPayloadLength dimension columnCount + | dimension <= 0 = + Left + ( InvariantViolation + "native eigenpairs require a positive dimension" + ) + | columnCount < 0 = + Left + ( InvariantViolation + "native eigenpair count cannot be negative" + ) + | otherwise = + first + (const (InvariantViolation "native eigenvector payload exceeds Int range")) + (checkedNonNegativeProduct dimension columnCount) + +reverseEigenvectorColumns :: + Int -> + Int -> + U.Vector Double -> + U.Vector Double +reverseEigenvectorColumns dimension columnCount eigenvectors = + U.generate + (U.length eigenvectors) + ( \payloadIndex -> + let (!targetColumn, !rowIndex) = + payloadIndex `quotRem` dimension + !sourceColumn = columnCount - targetColumn - 1 + in eigenvectors + `U.unsafeIndex` + (sourceColumn * dimension + rowIndex) + ) + +storableVectorToUnboxed :: S.Vector Double -> U.Vector Double +storableVectorToUnboxed values = + U.generate (S.length values) (values `S.unsafeIndex`) +{-# INLINE storableVectorToUnboxed #-} + +tridiagonalEigenpairsFromRawPairs :: + SpectrumEnd -> + SymmetricTridiagonal -> + [(Double, [Double])] -> + Either MoonlightError Eigenpairs +tridiagonalEigenpairsFromRawPairs spectrumEnd tridiagonalValue rawPairs = + let sortedPairs = sortRawPairsForSpectrum spectrumEnd rawPairs + dimension = symmetricTridiagonalDimension tridiagonalValue + eigenvalues = U.fromList (fst <$> sortedPairs) + eigenvectors = U.fromList (sortedPairs >>= snd) + diagonalValues = U.fromList (symmetricTridiagonalDiagonalEntries tridiagonalValue) + offDiagonalValues = U.fromList (symmetricTridiagonalOffDiagonalEntries tridiagonalValue) + residuals = + U.generate + (length sortedPairs) + (tridiagonalResidualNormAt dimension diagonalValues offDiagonalValues eigenvalues eigenvectors) + in validateNativeResidualNorms "native tridiagonal eigensolve" dimension (tridiagonalFrobeniusNorm diagonalValues offDiagonalValues) residuals + *> mkEigenpairs dimension eigenvalues eigenvectors residuals + +blockTridiagonalEigenpairsFromRawPairs :: + SpectrumEnd -> + SymmetricBlockTridiagonal -> + [(Double, [Double])] -> + Either MoonlightError Eigenpairs +blockTridiagonalEigenpairsFromRawPairs spectrumEnd blockValue rawPairs = do + let sortedPairs = sortRawPairsForSpectrum spectrumEnd rawPairs + dimension = symmetricBlockTridiagonalDimension blockValue + eigenvalues = U.fromList (fst <$> sortedPairs) + eigenvectors = U.fromList (sortedPairs >>= snd) + residuals <- U.fromList <$> traverse (blockResidualNorm blockValue) sortedPairs + let matrixNorm = symmetricBlockTridiagonalFrobeniusNorm blockValue + validateNativeResidualNorms "native symmetric-band eigensolve" dimension matrixNorm residuals + mkEigenpairs dimension eigenvalues eigenvectors residuals + +orderNativeValues :: SpectrumEnd -> U.Vector Double -> U.Vector Double +orderNativeValues spectrumEnd values = + case spectrumEnd of + SmallestEigenvalues -> values + LargestEigenvalues -> U.reverse values + +denseResidualNormAt :: + Int -> + U.Vector Double -> + U.Vector Double -> + U.Vector Double -> + Int -> + Double +denseResidualNormAt + dimension + matrixPayload + eigenvalues + eigenvectors + columnIndex = + sqrt (rowLoop 0 0.0) + where + !eigenvalue = eigenvalues `U.unsafeIndex` columnIndex + !vectorStart = columnIndex * dimension + + rowLoop !rowIndex !sumSquares + | rowIndex >= dimension = sumSquares + | otherwise = + let !imageEntry = matrixRowDot rowIndex 0 0.0 + !vectorEntry = + eigenvectors + `U.unsafeIndex` + (vectorStart + rowIndex) + !residualEntry = + imageEntry - eigenvalue * vectorEntry + in rowLoop + (rowIndex + 1) + (sumSquares + residualEntry * residualEntry) + + matrixRowDot !rowIndex !columnIndexValue !accumulator + | columnIndexValue >= dimension = accumulator + | otherwise = + let !matrixEntry = + matrixPayload + `U.unsafeIndex` + (rowIndex * dimension + columnIndexValue) + !vectorEntry = + eigenvectors + `U.unsafeIndex` + (vectorStart + columnIndexValue) + in matrixRowDot + rowIndex + (columnIndexValue + 1) + (accumulator + matrixEntry * vectorEntry) + +tridiagonalResidualNormAt :: + Int -> + U.Vector Double -> + U.Vector Double -> + U.Vector Double -> + U.Vector Double -> + Int -> + Double +tridiagonalResidualNormAt dimension diagonalValues offDiagonalValues eigenvalues eigenvectors columnIndex = + sqrt (U.sum (U.generate dimension entryResidualSquare)) + where + eigenvalue = eigenvalues `U.unsafeIndex` columnIndex + vectorStart = columnIndex * dimension + + vectorEntry rowIndex = + eigenvectors `U.unsafeIndex` (vectorStart + rowIndex) + + offDiagonalEntry rowIndex = + offDiagonalValues `U.unsafeIndex` rowIndex + + entryResidualSquare rowIndex = + let mainEntry = diagonalValues `U.unsafeIndex` rowIndex + currentValue = vectorEntry rowIndex + lowerContribution = + if rowIndex <= 0 + then 0.0 + else offDiagonalEntry (rowIndex - 1) * vectorEntry (rowIndex - 1) + upperContribution = + if rowIndex + 1 >= dimension + then 0.0 + else offDiagonalEntry rowIndex * vectorEntry (rowIndex + 1) + residualValue = + lowerContribution + + mainEntry * currentValue + + upperContribution + - eigenvalue * currentValue + in residualValue * residualValue + +blockResidualNorm :: SymmetricBlockTridiagonal -> (Double, [Double]) -> Either MoonlightError Double +blockResidualNorm blockValue (eigenvalue, eigenvector) = do + let vectorValue = U.fromList eigenvector + imageVector <- applySymmetricBlockTridiagonalU blockValue vectorValue + residualVector <- subU imageVector (scaleU eigenvalue vectorValue) + pure (normU residualVector) + +dynMatrixDimension :: DynMatrix Double -> Int +dynMatrixDimension matrixValue = + case dynMatrixShape matrixValue of + (rowCount, _) -> rowCount + +validateNativeResidualNorms :: String -> Int -> Double -> U.Vector Double -> Either MoonlightError () +validateNativeResidualNorms context dimension matrixNorm residuals = + let residualLimit = + 1.0e7 + * max 1.0 matrixNorm + * max 1.0 (fromIntegral dimension) + * epsDouble + accepted residualValue = + finiteDouble residualValue && residualValue <= residualLimit + in if U.all accepted residuals + then Right () + else + Left + ( InvariantViolation + ( context + <> " residual exceeded tolerance; limit=" + <> show residualLimit + <> ", residuals=" + <> show (U.toList residuals) + ) + ) + +frobeniusNormU :: U.Vector Double -> Double +frobeniusNormU values = + sqrt + ( U.foldl' + (\accumulator entryValue -> accumulator + entryValue * entryValue) + 0.0 + values + ) + +tridiagonalFrobeniusNorm :: U.Vector Double -> U.Vector Double -> Double +tridiagonalFrobeniusNorm diagonalValues offDiagonalValues = + sqrt + ( U.sum (U.map (\entryValue -> entryValue * entryValue) diagonalValues) + + 2.0 * U.sum (U.map (\entryValue -> entryValue * entryValue) offDiagonalValues) + ) diff --git a/src-native/Moonlight/LinAlg/Effect/Native/LAPACK.hs b/src-native/Moonlight/LinAlg/Effect/Native/LAPACK.hs new file mode 100644 --- /dev/null +++ b/src-native/Moonlight/LinAlg/Effect/Native/LAPACK.hs @@ -0,0 +1,1866 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ForeignFunctionInterface #-} +{-# LANGUAGE LambdaCase #-} + +module Moonlight.LinAlg.Effect.Native.LAPACK + ( denseDoubleLinearSolveLapack, + denseDoubleMatrixProductBlas, + denseDoubleSymmetricEigenpairsRawLapack, + leastSquaresLapack, + selectedSymmetricEigenValuesLapack, + selectedSymmetricEigenPairsLapack, + selectedSymmetricBlockTridiagonalEigenValuesLapack, + selectedSymmetricBlockTridiagonalEigenPairsLapack, + selectedSymmetricTridiagonalEigenValuesLapack, + selectedSymmetricTridiagonalEigenPairsLapack, + ) +where + +import Control.Monad.ST (ST, runST) +import Data.Bifunctor (first) +import Data.Kind (Type) +import Data.List (transpose) +import qualified Data.Vector.Storable as S +import qualified Data.Vector.Storable.Mutable as MS +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as MU +import Foreign + ( Ptr, + alloca, + allocaArray, + castPtr, + peek, + peekArray, + peekElemOff, + poke, + pokeElemOff, + with, + withArray, + ) +import Foreign.C.String (castCharToCChar) +import Foreign.C.Types (CChar, CDouble (..), CInt (..)) +import Foreign.ForeignPtr (mallocForeignPtrArray, withForeignPtr) +import Moonlight.Core + ( MoonlightError (..), + checkedNonNegativeProduct, + checkedNonNegativeSum, + fieldValueValid, + ) +import Moonlight.LinAlg.Internal.Storage (chunkRows) +import Moonlight.LinAlg.Pure.Dense.Dynamic + ( DynMatrix, + DynVector, + dynMatrixShape, + dynMatrixToList, + dynMatrixToRows, + dynVectorToList, + ) +import Moonlight.LinAlg.Pure.Dense.Flat + ( DenseDoubleMatrix, + denseDoubleMatrixShape, + denseDoubleMatrixToRowMajorVector, + trustedDenseDoubleMatrixRowMajor, + ) +import Moonlight.LinAlg.Pure.Krylov.Selection (SpectrumEnd (..)) +import Moonlight.LinAlg.Pure.Structured.BlockTridiagonal + ( SymmetricBlockTridiagonal, + blockOffsets, + couplingPayloadOffsets, + diagonalLowerPacked, + diagonalPayloadOffsets, + lowerCouplingPayload, + symmetricBlockTridiagonalBandwidth, + symmetricBlockTridiagonalBlockCount, + symmetricBlockTridiagonalDimension, + ) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal + ( SymmetricTridiagonal, + symmetricTridiagonalDiagonalEntries, + symmetricTridiagonalOffDiagonalEntries, + ) +import Prelude + +type FortranIndexRange :: Type +data FortranIndexRange = FortranIndexRange + { fortranIndexRangeLower :: !Int, + fortranIndexRangeUpper :: !Int + } + deriving stock (Eq, Show) + +mkFortranIndexRange :: Int -> Int -> Either MoonlightError FortranIndexRange +mkFortranIndexRange lowerIndex upperIndex + | lowerIndex < 1 = + Left (InvariantViolation "Fortran index range lower bound must be positive") + | upperIndex < lowerIndex = + Left (InvariantViolation "Fortran index range upper bound must be at least the lower bound") + | otherwise = + Right + FortranIndexRange + { fortranIndexRangeLower = lowerIndex, + fortranIndexRangeUpper = upperIndex + } + +denseSelectedEigenIndexRange :: SpectrumEnd -> Int -> Int -> Int -> Either MoonlightError FortranIndexRange +denseSelectedEigenIndexRange spectrumEnd requestedCount rowCount columnCount + | rowCount /= columnCount = + Left (InvariantViolation "native dense symmetric eigensolve requires a square matrix") + | otherwise = + selectedEigenIndexRange spectrumEnd requestedCount rowCount + +selectedEigenIndexRange :: SpectrumEnd -> Int -> Int -> Either MoonlightError FortranIndexRange +selectedEigenIndexRange spectrumEnd requestedCount dimension = + case selectedNativeIndexBounds spectrumEnd requestedCount dimension of + Left err -> Left err + Right (lowerIndex, upperIndex) -> mkFortranIndexRange lowerIndex upperIndex + +selectedNativeIndexBounds :: SpectrumEnd -> Int -> Int -> Either MoonlightError (Int, Int) +selectedNativeIndexBounds spectrumEnd requestedCount dimension + | requestedCount <= 0 = Left (InvariantViolation "native eigen request count must be positive") + | requestedCount > dimension = Left (InvariantViolation "native eigen request count exceeds matrix dimension") + | otherwise = + Right + ( case spectrumEnd of + SmallestEigenvalues -> (1, requestedCount) + LargestEigenvalues -> (dimension - requestedCount + 1, dimension) + ) + +foreign import ccall unsafe "dsyev_" + lapackDsyev :: + Ptr CChar -> + Ptr CChar -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + IO () + +foreign import ccall unsafe "moonlight_dgemm_row_major" + moonlightDgemmRowMajor :: + CInt -> + CInt -> + CInt -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CDouble -> + IO () + +foreign import ccall unsafe "dgesv_" + lapackDgesv :: + Ptr CInt -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + IO () + +foreign import ccall unsafe "dsyevx_" + lapackDsyevx :: + Ptr CChar -> + Ptr CChar -> + Ptr CChar -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + Ptr CInt -> + Ptr CInt -> + IO () + +foreign import ccall unsafe "dgels_" + lapackDgels :: + Ptr CChar -> + Ptr CInt -> + Ptr CInt -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + IO () + +foreign import ccall unsafe "dstemr_" + lapackDstemr :: + Ptr CChar -> + Ptr CChar -> + Ptr CInt -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + Ptr CInt -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + Ptr CInt -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + Ptr CInt -> + Ptr CInt -> + IO () + +foreign import ccall unsafe "dsbevx_" + lapackDsbevx :: + Ptr CChar -> + Ptr CChar -> + Ptr CChar -> + Ptr CInt -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + Ptr CInt -> + IO () + +denseDoubleMatrixProductBlas :: DenseDoubleMatrix -> DenseDoubleMatrix -> IO (Either MoonlightError DenseDoubleMatrix) +denseDoubleMatrixProductBlas leftMatrix rightMatrix = + let (leftRows, leftColumns) = denseDoubleMatrixShape leftMatrix + (rightRows, rightColumns) = denseDoubleMatrixShape rightMatrix + in case validateDenseProductInput leftRows leftColumns rightRows rightColumns of + Left err -> pure (Left err) + Right () + | otherwise -> + case checkedProduct "BLAS dense matrix product output entry count" leftRows rightColumns of + Left err -> pure (Left err) + Right outputLength + | leftRows == 0 || rightColumns == 0 -> + pure (Right (trustedDenseDoubleMatrixRowMajor leftRows rightColumns S.empty)) + | leftColumns == 0 -> + pure + ( Right + ( trustedDenseDoubleMatrixRowMajor + leftRows + rightColumns + (S.replicate outputLength 0.0) + ) + ) + | otherwise -> + case traverse matrixSizeAsLapackInt [leftRows, rightColumns, leftColumns] of + Left err -> pure (Left err) + Right [lapackLeftRows, lapackRightColumns, lapackInner] -> + solveDenseProductBlas lapackLeftRows lapackRightColumns lapackInner leftRows rightColumns outputLength leftPayload rightPayload + Right _ -> pure (Left (InvariantViolation "BLAS dense matrix product internal dimension arity mismatch")) + where + leftPayload = denseDoubleMatrixToRowMajorVector leftMatrix + rightPayload = denseDoubleMatrixToRowMajorVector rightMatrix + +denseDoubleLinearSolveLapack :: DenseDoubleMatrix -> S.Vector Double -> IO (Either MoonlightError (S.Vector Double)) +denseDoubleLinearSolveLapack matrixValue rightHandSide = + let (rowCount, columnCount) = denseDoubleMatrixShape matrixValue + matrixPayload = denseDoubleMatrixToRowMajorVector matrixValue + in case validateDenseLinearSolveInput rowCount columnCount rightHandSide of + Left err -> pure (Left err) + Right () -> + case matrixSizeAsLapackInt rowCount of + Left err -> pure (Left err) + Right lapackSize -> + solveDenseLinearSystemLapack lapackSize rowCount matrixPayload rightHandSide + +denseDoubleSymmetricEigenpairsRawLapack :: DenseDoubleMatrix -> IO (Either MoonlightError (S.Vector Double, S.Vector Double)) +denseDoubleSymmetricEigenpairsRawLapack matrixValue = + let (rowCount, columnCount) = denseDoubleMatrixShape matrixValue + matrixPayload = denseDoubleMatrixToRowMajorVector matrixValue + in case validateDenseSymmetricEigenInput rowCount columnCount of + Left err -> pure (Left err) + Right () -> + case matrixSizeAsLapackInt rowCount of + Left err -> pure (Left err) + Right lapackSize -> + solveDenseSymmetricEigenpairsRawLapack lapackSize rowCount matrixPayload + +validateDenseProductInput :: + Int -> + Int -> + Int -> + Int -> + Either MoonlightError () +validateDenseProductInput leftRows leftColumns rightRows rightColumns + | leftColumns /= rightRows = + Left + ( InvariantViolation + ( "BLAS dense matrix product shape mismatch: left " + <> show (leftRows, leftColumns) + <> " right " + <> show (rightRows, rightColumns) + ) + ) + | otherwise = Right () + +validateDenseLinearSolveInput :: + Int -> + Int -> + S.Vector Double -> + Either MoonlightError () +validateDenseLinearSolveInput rowCount columnCount rightHandSide + | rowCount /= columnCount = + Left (InvariantViolation "LAPACK dense linear solve requires a square matrix") + | rowCount <= 0 = + Left (InvariantViolation "LAPACK dense linear solve requires a positive dimension") + | S.length rightHandSide /= rowCount = + Left + ( InvariantViolation + ( "LAPACK dense linear solve right-hand side length mismatch: matrix dimension " + <> show rowCount + <> ", vector " + <> show (S.length rightHandSide) + ) + ) + | S.any (not . fieldValueValid) rightHandSide = + Left (InvariantViolation "LAPACK dense linear solve requires finite right-hand side entries") + | otherwise = Right () + +validateDenseSymmetricEigenInput :: + Int -> + Int -> + Either MoonlightError () +validateDenseSymmetricEigenInput rowCount columnCount + | rowCount /= columnCount = + Left (InvariantViolation "LAPACK dense symmetric eigensolve requires a square matrix") + | rowCount <= 0 = + Left (InvariantViolation "LAPACK dense symmetric eigensolve requires a positive dimension") + | otherwise = Right () + +solveDenseProductBlas :: + CInt -> + CInt -> + CInt -> + Int -> + Int -> + Int -> + S.Vector Double -> + S.Vector Double -> + IO (Either MoonlightError DenseDoubleMatrix) +solveDenseProductBlas !lapackLeftRows !lapackRightColumns !lapackInner !leftRows !rightColumns !outputLength leftPayload rightPayload = do + outputPayload <- MS.unsafeNew outputLength + S.unsafeWith leftPayload $ \leftPointer -> + S.unsafeWith rightPayload $ \rightPointer -> + MS.unsafeWith outputPayload $ \outputPointer -> + moonlightDgemmRowMajor + lapackLeftRows + lapackRightColumns + lapackInner + (castPtr leftPointer) + (castPtr rightPointer) + (castPtr outputPointer) + frozenOutput <- S.unsafeFreeze outputPayload + pure + ( if S.any (not . fieldValueValid) frozenOutput + then Left (InvariantViolation "BLAS dense matrix product produced non-finite entries") + else + Right + ( trustedDenseDoubleMatrixRowMajor + leftRows + rightColumns + frozenOutput + ) + ) + +solveDenseLinearSystemLapack :: + CInt -> + Int -> + S.Vector Double -> + S.Vector Double -> + IO (Either MoonlightError (S.Vector Double)) +solveDenseLinearSystemLapack !lapackSize !matrixSize matrixPayload rightHandSide = + with lapackSize $ \sizePointer -> + with (1 :: CInt) $ \rightHandSideCountPointer -> + with lapackSize $ \leadingDimensionPointer -> + with lapackSize $ \rightHandSideLeadingDimensionPointer -> + allocaArray matrixSize $ \pivotPointer -> + alloca $ \infoPointer -> do + matrixWork <- S.thaw (denseRowMajorToColumnMajorSquare matrixSize matrixPayload) + rightHandSideWork <- S.thaw rightHandSide + MS.unsafeWith matrixWork $ \matrixPointer -> + MS.unsafeWith rightHandSideWork $ \rightHandSidePointer -> do + poke infoPointer 0 + lapackDgesv + sizePointer + rightHandSideCountPointer + (castPtr matrixPointer) + leadingDimensionPointer + pivotPointer + (castPtr rightHandSidePointer) + rightHandSideLeadingDimensionPointer + infoPointer + infoValue <- peek infoPointer + if infoValue /= 0 + then pure (Left (lapackLinearSolveInfoError infoValue)) + else do + solution <- S.unsafeFreeze rightHandSideWork + pure + ( if S.any (not . fieldValueValid) solution + then Left (InvariantViolation "LAPACK dense linear solve produced non-finite entries") + else Right solution + ) + +solveDenseSymmetricEigenpairsRawLapack :: + CInt -> + Int -> + S.Vector Double -> + IO (Either MoonlightError (S.Vector Double, S.Vector Double)) +solveDenseSymmetricEigenpairsRawLapack !lapackSize !matrixSize matrixPayload = + withLapackChar 'V' $ \jobPointer -> + withLapackChar 'U' $ \uploPointer -> + with lapackSize $ \sizePointer -> + with lapackSize $ \leadingDimensionPointer -> + alloca $ \infoPointer -> do + matrixWork <- S.thaw matrixPayload + eigenvalueWork <- MS.replicate matrixSize 0.0 + MS.unsafeWith matrixWork $ \matrixPointer -> + MS.unsafeWith eigenvalueWork $ \eigenvaluePointer -> + validateSymmetricDenseBuffer matrixSize (castPtr matrixPointer) >>= \case + Left err -> pure (Left err) + Right () -> + queryWorkspace + jobPointer + uploPointer + sizePointer + (castPtr matrixPointer) + leadingDimensionPointer + (castPtr eigenvaluePointer) + infoPointer + >>= \case + Left err -> pure (Left err) + Right workspaceSize -> + allocaArray workspaceSize $ \workspacePointer -> do + poke infoPointer 0 + with (fromIntegral workspaceSize) $ \workspaceSizePointer -> + lapackDsyev + jobPointer + uploPointer + sizePointer + (castPtr matrixPointer) + leadingDimensionPointer + (castPtr eigenvaluePointer) + workspacePointer + workspaceSizePointer + infoPointer + decodeDenseSymmetricEigenpairsRaw matrixWork eigenvalueWork infoPointer + +denseRowMajorToColumnMajorSquare :: Int -> S.Vector Double -> S.Vector Double +denseRowMajorToColumnMajorSquare matrixSize matrixPayload = + S.generate + (S.length matrixPayload) + ( \payloadIndex -> + let (!columnIndex, !rowIndex) = payloadIndex `quotRem` matrixSize + in matrixPayload `S.unsafeIndex` (rowIndex * matrixSize + columnIndex) + ) +{-# INLINE denseRowMajorToColumnMajorSquare #-} + +decodeDenseSymmetricEigenpairsRaw :: + MS.IOVector Double -> + MS.IOVector Double -> + Ptr CInt -> + IO (Either MoonlightError (S.Vector Double, S.Vector Double)) +decodeDenseSymmetricEigenpairsRaw matrixWork eigenvalueWork infoPointer = do + infoValue <- peek infoPointer + if infoValue /= 0 + then pure (Left (lapackInfoError "LAPACK DSYEV" infoValue)) + else do + eigenvalues <- S.unsafeFreeze eigenvalueWork + eigenvectors <- S.unsafeFreeze matrixWork + pure + ( if S.any (not . fieldValueValid) eigenvalues || S.any (not . fieldValueValid) eigenvectors + then Left (InvariantViolation "LAPACK dense symmetric eigensolve produced non-finite entries") + else Right (eigenvalues, eigenvectors) + ) + +selectedSymmetricEigenPairsLapack :: + SpectrumEnd -> + Int -> + DynMatrix Double -> + IO + ( Either + MoonlightError + (U.Vector Double, U.Vector Double) + ) +selectedSymmetricEigenPairsLapack spectrumEnd requestedCount matrixValue = + let (rowCount, columnCount) = dynMatrixShape matrixValue + in case + denseSelectedEigenIndexRange + spectrumEnd + requestedCount + rowCount + columnCount of + Left err -> pure (Left err) + Right indexRange -> + withSelectedSymmetricDenseBuffer + matrixValue + ( \matrixSize lapackSize matrixPointer -> + solveSelectedSymmetricPairsLapack + lapackSize + matrixSize + (fortranIndexRangeLower indexRange) + (fortranIndexRangeUpper indexRange) + matrixPointer + ) + +selectedSymmetricEigenValuesLapack :: + SpectrumEnd -> + Int -> + DynMatrix Double -> + IO (Either MoonlightError (U.Vector Double)) +selectedSymmetricEigenValuesLapack spectrumEnd requestedCount matrixValue = + let (rowCount, columnCount) = dynMatrixShape matrixValue + in case + denseSelectedEigenIndexRange + spectrumEnd + requestedCount + rowCount + columnCount of + Left err -> pure (Left err) + Right indexRange -> + withSelectedSymmetricDenseBuffer + matrixValue + ( \matrixSize lapackSize matrixPointer -> + let !lowerIndex = fortranIndexRangeLower indexRange + !upperIndex = fortranIndexRangeUpper indexRange + in if matrixSize <= smallDenseValuesFullThreshold + then + solveSelectedSymmetricValuesDsyev + lapackSize + matrixSize + lowerIndex + upperIndex + matrixPointer + else + solveSelectedSymmetricValuesLapack + lapackSize + matrixSize + lowerIndex + upperIndex + matrixPointer + ) + +smallDenseValuesFullThreshold :: Int +smallDenseValuesFullThreshold = 32 +{-# INLINE smallDenseValuesFullThreshold #-} + +withSelectedSymmetricDenseBuffer :: + DynMatrix Double -> + (Int -> CInt -> Ptr CDouble -> IO (Either MoonlightError result)) -> + IO (Either MoonlightError result) +withSelectedSymmetricDenseBuffer matrixValue useBuffer = + case validateSelectedDenseStorage matrixValue of + Left err -> pure (Left err) + Right (matrixSize, lapackSize, entryCount) -> do + matrixForeignPointer <- mallocForeignPtrArray entryCount + withForeignPtr matrixForeignPointer $ \matrixPointer -> do + copiedPayload <- + copyFiniteDensePayload + entryCount + (dynMatrixToList matrixValue) + matrixPointer + case copiedPayload of + Left err -> pure (Left err) + Right () -> do + symmetryResult <- + validateSymmetricDenseBuffer + matrixSize + matrixPointer + case symmetryResult of + Left err -> pure (Left err) + Right () -> + useBuffer matrixSize lapackSize matrixPointer + +validateSelectedDenseStorage :: + DynMatrix Double -> + Either MoonlightError (Int, CInt, Int) +validateSelectedDenseStorage matrixValue = do + let (rowCount, columnCount) = dynMatrixShape matrixValue + if rowCount /= columnCount + then + Left + ( InvariantViolation + "LAPACK selected symmetric eigensolve requires a square matrix" + ) + else pure () + if rowCount <= 0 + then + Left + ( InvariantViolation + "LAPACK selected symmetric eigensolve requires a positive dimension" + ) + else pure () + entryCount <- + checkedProduct + "LAPACK selected symmetric matrix entry count" + rowCount + columnCount + lapackSize <- matrixSizeAsLapackInt rowCount + pure (rowCount, lapackSize, entryCount) + +copyFiniteDensePayload :: + Int -> + [Double] -> + Ptr CDouble -> + IO (Either MoonlightError ()) +copyFiniteDensePayload expectedCount values targetPointer = + go 0 values + where + go !entryIndex remainingValues + | entryIndex >= expectedCount = + case remainingValues of + [] -> pure (Right ()) + _ -> + pure + ( Left + ( InvariantViolation + "LAPACK selected symmetric matrix payload contains excess entries" + ) + ) + | otherwise = + case remainingValues of + [] -> + pure + ( Left + ( InvariantViolation + ( "LAPACK selected symmetric matrix payload ended at offset " + <> show entryIndex + ) + ) + ) + entryValue : rest + | not (fieldValueValid entryValue) -> + pure + ( Left + ( InvariantViolation + ( "LAPACK selected symmetric eigensolve requires finite entries; invalid offset " + <> show entryIndex + ) + ) + ) + | otherwise -> do + pokeElemOff targetPointer entryIndex (CDouble entryValue) + go (entryIndex + 1) rest + +validateSymmetricDenseBuffer :: + Int -> + Ptr CDouble -> + IO (Either MoonlightError ()) +validateSymmetricDenseBuffer matrixSize matrixPointer = + validateRow 0 + where + !tolerance = 1.0e-6 + + validateRow !rowIndex + | rowIndex >= matrixSize = pure (Right ()) + | otherwise = validateColumn rowIndex (rowIndex + 1) + + validateColumn !rowIndex !columnIndex + | columnIndex >= matrixSize = validateRow (rowIndex + 1) + | otherwise = do + CDouble upperValue <- + peekElemOff + matrixPointer + (rowIndex * matrixSize + columnIndex) + CDouble lowerValue <- + peekElemOff + matrixPointer + (columnIndex * matrixSize + rowIndex) + if abs (upperValue - lowerValue) <= tolerance + then validateColumn rowIndex (columnIndex + 1) + else + pure + ( Left + ( InvariantViolation + ( "LAPACK selected symmetric eigensolve requires a symmetric matrix; mismatch at " + <> show (rowIndex, columnIndex) + ) + ) + ) + +leastSquaresLapack :: DynMatrix Double -> DynVector Double -> IO (Either MoonlightError [Double]) +leastSquaresLapack matrixValue rightHandSideValue = + case dynMatrixToRows matrixValue of + Left err -> pure (Left err) + Right matrixToRows -> + let (rowCount, columnCount) = dynMatrixShape matrixValue + rightHandSide = dynVectorToList rightHandSideValue + in case validateLeastSquaresInput rowCount columnCount matrixToRows rightHandSide of + Left err -> pure (Left err) + Right () -> + case (matrixSizeAsLapackInt rowCount, matrixSizeAsLapackInt columnCount) of + (Right lapackRows, Right lapackColumns) -> + solveLeastSquaresLapack lapackRows lapackColumns rowCount columnCount matrixToRows rightHandSide + (Left err, _) -> pure (Left err) + (_, Left err) -> pure (Left err) + +selectedSymmetricTridiagonalEigenPairsLapack :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + IO (Either MoonlightError [(Double, [Double])]) +selectedSymmetricTridiagonalEigenPairsLapack spectrumEnd requestedCount tridiagonalValue = + let diagonalValues = symmetricTridiagonalDiagonalEntries tridiagonalValue + offDiagonalValues = symmetricTridiagonalOffDiagonalEntries tridiagonalValue + in case selectedEigenIndexRange spectrumEnd requestedCount (length diagonalValues) of + Left err -> pure (Left err) + Right indexRange -> + let lowerIndex = fortranIndexRangeLower indexRange + upperIndex = fortranIndexRangeUpper indexRange + in case validateSelectedTridiagonalInput lowerIndex upperIndex diagonalValues offDiagonalValues of + Left err -> pure (Left err) + Right matrixSize -> + case matrixSizeAsLapackInt matrixSize of + Left err -> pure (Left err) + Right lapackSize -> + solveSelectedTridiagonalLapack lapackSize matrixSize lowerIndex upperIndex diagonalValues offDiagonalValues + +selectedSymmetricTridiagonalEigenValuesLapack :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + IO (Either MoonlightError (U.Vector Double)) +selectedSymmetricTridiagonalEigenValuesLapack spectrumEnd requestedCount tridiagonalValue = + let diagonalValues = symmetricTridiagonalDiagonalEntries tridiagonalValue + offDiagonalValues = symmetricTridiagonalOffDiagonalEntries tridiagonalValue + in case selectedEigenIndexRange spectrumEnd requestedCount (length diagonalValues) of + Left err -> pure (Left err) + Right indexRange -> + let lowerIndex = fortranIndexRangeLower indexRange + upperIndex = fortranIndexRangeUpper indexRange + in case validateSelectedTridiagonalInput lowerIndex upperIndex diagonalValues offDiagonalValues of + Left err -> pure (Left err) + Right matrixSize -> + case matrixSizeAsLapackInt matrixSize of + Left err -> pure (Left err) + Right lapackSize -> + solveSelectedTridiagonalValuesLapack lapackSize matrixSize lowerIndex upperIndex diagonalValues offDiagonalValues + +selectedSymmetricBlockTridiagonalEigenPairsLapack :: + SpectrumEnd -> + Int -> + SymmetricBlockTridiagonal -> + IO (Either MoonlightError [(Double, [Double])]) +selectedSymmetricBlockTridiagonalEigenPairsLapack spectrumEnd requestedCount blockValue = + let matrixSize = symmetricBlockTridiagonalDimension blockValue + bandwidth = symmetricBlockTridiagonalBandwidth blockValue + in case selectedEigenIndexRange spectrumEnd requestedCount matrixSize of + Left err -> pure (Left err) + Right indexRange -> + let lowerIndex = fortranIndexRangeLower indexRange + upperIndex = fortranIndexRangeUpper indexRange + in case symmetricBlockTridiagonalLowerBandPayload blockValue of + Left err -> pure (Left err) + Right lowerBandPayload -> + case validateSelectedBandInput lowerIndex upperIndex matrixSize bandwidth lowerBandPayload of + Left err -> pure (Left err) + Right leadingDimensionValue -> + case (matrixSizeAsLapackInt matrixSize, matrixSizeAsLapackInt bandwidth, matrixSizeAsLapackInt leadingDimensionValue) of + (Right lapackSize, Right lapackBandwidth, Right leadingDimension) -> + solveSelectedBandPairsLapack lapackSize lapackBandwidth leadingDimension matrixSize lowerIndex upperIndex lowerBandPayload + (Left err, _, _) -> pure (Left err) + (_, Left err, _) -> pure (Left err) + (_, _, Left err) -> pure (Left err) + +selectedSymmetricBlockTridiagonalEigenValuesLapack :: + SpectrumEnd -> + Int -> + SymmetricBlockTridiagonal -> + IO (Either MoonlightError (U.Vector Double)) +selectedSymmetricBlockTridiagonalEigenValuesLapack spectrumEnd requestedCount blockValue = + let matrixSize = symmetricBlockTridiagonalDimension blockValue + bandwidth = symmetricBlockTridiagonalBandwidth blockValue + in case selectedEigenIndexRange spectrumEnd requestedCount matrixSize of + Left err -> pure (Left err) + Right indexRange -> + let lowerIndex = fortranIndexRangeLower indexRange + upperIndex = fortranIndexRangeUpper indexRange + in case symmetricBlockTridiagonalLowerBandPayload blockValue of + Left err -> pure (Left err) + Right lowerBandPayload -> + case validateSelectedBandInput lowerIndex upperIndex matrixSize bandwidth lowerBandPayload of + Left err -> pure (Left err) + Right leadingDimensionValue -> + case (matrixSizeAsLapackInt matrixSize, matrixSizeAsLapackInt bandwidth, matrixSizeAsLapackInt leadingDimensionValue) of + (Right lapackSize, Right lapackBandwidth, Right leadingDimension) -> + solveSelectedBandValuesLapack lapackSize lapackBandwidth leadingDimension matrixSize lowerIndex upperIndex lowerBandPayload + (Left err, _, _) -> pure (Left err) + (_, Left err, _) -> pure (Left err) + (_, _, Left err) -> pure (Left err) + +symmetricBlockTridiagonalLowerBandPayload :: SymmetricBlockTridiagonal -> Either MoonlightError (U.Vector Double) +symmetricBlockTridiagonalLowerBandPayload blockValue = do + let dimension = symmetricBlockTridiagonalDimension blockValue + leadingDimension <- + checkedSum + "LAPACK symmetric block-tridiagonal leading dimension" + (symmetricBlockTridiagonalBandwidth blockValue) + 1 + payloadLength <- + checkedProduct + "LAPACK symmetric block-tridiagonal band payload" + leadingDimension + dimension + Right $ runST $ do + bandPayload <- MU.replicate payloadLength 0.0 + U.foldM' + (writeDiagonalBandBlock blockValue leadingDimension bandPayload) + () + (U.enumFromN 0 (symmetricBlockTridiagonalBlockCount blockValue)) + U.foldM' + (writeCouplingBandBlock blockValue leadingDimension bandPayload) + () + (U.enumFromN 0 (max 0 (symmetricBlockTridiagonalBlockCount blockValue - 1))) + U.unsafeFreeze bandPayload + +writeDiagonalBandBlock :: + SymmetricBlockTridiagonal -> + Int -> + MU.MVector s Double -> + () -> + Int -> + ST s () +writeDiagonalBandBlock blockValue leadingDimension bandPayload () blockIndex = + U.foldM' + (writeDiagonalBandRow blockValue leadingDimension bandPayload blockIndex blockStart) + () + (U.enumFromN 0 (nativeBlockSizeAt blockValue blockIndex)) + where + blockStart = nativeIntAt (blockOffsets blockValue) blockIndex +{-# INLINE writeDiagonalBandBlock #-} + +writeDiagonalBandRow :: + SymmetricBlockTridiagonal -> + Int -> + MU.MVector s Double -> + Int -> + Int -> + () -> + Int -> + ST s () +writeDiagonalBandRow blockValue leadingDimension bandPayload blockIndex blockStart () localRow = + U.foldM' + (writeDiagonalBandEntry blockValue leadingDimension bandPayload blockIndex blockStart localRow) + () + (U.enumFromN 0 (localRow + 1)) +{-# INLINE writeDiagonalBandRow #-} + +writeDiagonalBandEntry :: + SymmetricBlockTridiagonal -> + Int -> + MU.MVector s Double -> + Int -> + Int -> + Int -> + () -> + Int -> + ST s () +writeDiagonalBandEntry blockValue leadingDimension bandPayload blockIndex blockStart localRow () localColumn = + writeLowerBandEntry + leadingDimension + bandPayload + (blockStart + localRow) + (blockStart + localColumn) + (nativeDiagonalEntry blockValue blockIndex localRow localColumn) +{-# INLINE writeDiagonalBandEntry #-} + +writeCouplingBandBlock :: + SymmetricBlockTridiagonal -> + Int -> + MU.MVector s Double -> + () -> + Int -> + ST s () +writeCouplingBandBlock blockValue leadingDimension bandPayload () couplingIndex = + U.foldM' + (writeCouplingBandRow blockValue leadingDimension bandPayload couplingIndex upperBlockStart lowerBlockStart) + () + (U.enumFromN 0 (nativeBlockSizeAt blockValue (couplingIndex + 1))) + where + upperBlockStart = nativeIntAt (blockOffsets blockValue) couplingIndex + lowerBlockStart = nativeIntAt (blockOffsets blockValue) (couplingIndex + 1) +{-# INLINE writeCouplingBandBlock #-} + +writeCouplingBandRow :: + SymmetricBlockTridiagonal -> + Int -> + MU.MVector s Double -> + Int -> + Int -> + Int -> + () -> + Int -> + ST s () +writeCouplingBandRow blockValue leadingDimension bandPayload couplingIndex upperBlockStart lowerBlockStart () localRow = + U.foldM' + (writeCouplingBandEntry blockValue leadingDimension bandPayload couplingIndex upperBlockStart lowerBlockStart localRow) + () + (U.enumFromN 0 (nativeBlockSizeAt blockValue couplingIndex)) +{-# INLINE writeCouplingBandRow #-} + +writeCouplingBandEntry :: + SymmetricBlockTridiagonal -> + Int -> + MU.MVector s Double -> + Int -> + Int -> + Int -> + Int -> + () -> + Int -> + ST s () +writeCouplingBandEntry blockValue leadingDimension bandPayload couplingIndex _ lowerBlockStart localRow () localColumn = + writeLowerBandEntry + leadingDimension + bandPayload + (lowerBlockStart + localRow) + (nativeIntAt (blockOffsets blockValue) couplingIndex + localColumn) + (nativeCouplingEntry blockValue couplingIndex localRow localColumn) +{-# INLINE writeCouplingBandEntry #-} + +writeLowerBandEntry :: Int -> MU.MVector s Double -> Int -> Int -> Double -> ST s () +writeLowerBandEntry leadingDimension bandPayload rowIndex columnIndex entryValue = + MU.unsafeWrite bandPayload (columnIndex * leadingDimension + rowIndex - columnIndex) entryValue +{-# INLINE writeLowerBandEntry #-} + +nativeBlockSizeAt :: SymmetricBlockTridiagonal -> Int -> Int +nativeBlockSizeAt blockValue blockIndex = + nativeIntAt (blockOffsets blockValue) (blockIndex + 1) + - nativeIntAt (blockOffsets blockValue) blockIndex +{-# INLINE nativeBlockSizeAt #-} + +nativeDiagonalEntry :: SymmetricBlockTridiagonal -> Int -> Int -> Int -> Double +nativeDiagonalEntry blockValue blockIndex localRow localColumn + | localColumn <= localRow = + nativeDoubleAt (diagonalLowerPacked blockValue) (nativeDiagonalPayloadStart blockValue blockIndex + nativePackedLowerIndex localRow localColumn) + | otherwise = + nativeDoubleAt (diagonalLowerPacked blockValue) (nativeDiagonalPayloadStart blockValue blockIndex + nativePackedLowerIndex localColumn localRow) +{-# INLINE nativeDiagonalEntry #-} + +nativeCouplingEntry :: SymmetricBlockTridiagonal -> Int -> Int -> Int -> Double +nativeCouplingEntry blockValue couplingIndex localRow localColumn = + let couplingStart = nativeIntAt (couplingPayloadOffsets blockValue) couplingIndex + couplingColumns = nativeBlockSizeAt blockValue couplingIndex + in nativeDoubleAt (lowerCouplingPayload blockValue) (couplingStart + localRow * couplingColumns + localColumn) +{-# INLINE nativeCouplingEntry #-} + +nativeDiagonalPayloadStart :: SymmetricBlockTridiagonal -> Int -> Int +nativeDiagonalPayloadStart blockValue blockIndex = + nativeIntAt (diagonalPayloadOffsets blockValue) blockIndex +{-# INLINE nativeDiagonalPayloadStart #-} + +nativePackedLowerIndex :: Int -> Int -> Int +nativePackedLowerIndex rowIndex columnIndex = + rowIndex * (rowIndex + 1) `quot` 2 + columnIndex +{-# INLINE nativePackedLowerIndex #-} + +nativeIntAt :: U.Vector Int -> Int -> Int +nativeIntAt values indexValue = + maybe 0 id (values U.!? indexValue) +{-# INLINE nativeIntAt #-} + +nativeDoubleAt :: U.Vector Double -> Int -> Double +nativeDoubleAt values indexValue = + maybe 0.0 id (values U.!? indexValue) +{-# INLINE nativeDoubleAt #-} + +solveSelectedSymmetricPairsLapack :: + CInt -> + Int -> + Int -> + Int -> + Ptr CDouble -> + IO + ( Either + MoonlightError + (U.Vector Double, U.Vector Double) + ) +solveSelectedSymmetricPairsLapack + !lapackSize + !matrixSize + !lowerIndex + !upperIndex + matrixPointer = + case + (,,) + <$> checkedProduct "LAPACK selected symmetric eigenvector workspace" matrixSize selectedCount + <*> checkedProduct "LAPACK DSYEVX floating workspace" 8 matrixSize + <*> checkedProduct "LAPACK DSYEVX integer workspace" 5 matrixSize + of + Left err -> pure (Left err) + Right (eigenvectorEntryCount, workspaceCount, integerWorkspaceCount) -> withLapackChar 'V' $ \jobPointer -> + withLapackChar 'I' $ \rangePointer -> + -- The row-major payload is the column-major payload of A^T. Reading + -- its upper triangle therefore preserves the original lower triangle. + withLapackChar 'U' $ \uploPointer -> + with lapackSize $ \sizePointer -> + with lapackSize $ \leadingDimensionPointer -> + with 0.0 $ \lowerValuePointer -> + with 0.0 $ \upperValuePointer -> + with (fromIntegral lowerIndex) $ \lowerIndexPointer -> + with (fromIntegral upperIndex) $ \upperIndexPointer -> + with 0.0 $ \absoluteTolerancePointer -> + alloca $ \foundCountPointer -> + allocaArray matrixSize $ \eigenvaluePointer -> + allocaArray eigenvectorEntryCount $ \eigenvectorPointer -> + with lapackSize $ \eigenvectorLeadingDimensionPointer -> + allocaArray workspaceCount $ \workspacePointer -> + with (fromIntegral workspaceCount) $ \workspaceSizePointer -> + allocaArray integerWorkspaceCount $ \integerWorkspacePointer -> + allocaArray matrixSize $ \failedVectorPointer -> + alloca $ \infoPointer -> do + poke foundCountPointer 0 + poke infoPointer 0 + lapackDsyevx + jobPointer + rangePointer + uploPointer + sizePointer + matrixPointer + leadingDimensionPointer + lowerValuePointer + upperValuePointer + lowerIndexPointer + upperIndexPointer + absoluteTolerancePointer + foundCountPointer + eigenvaluePointer + eigenvectorPointer + eigenvectorLeadingDimensionPointer + workspacePointer + workspaceSizePointer + integerWorkspacePointer + failedVectorPointer + infoPointer + decodeSelectedSymmetricColumns + selectedCount + eigenvectorEntryCount + eigenvaluePointer + eigenvectorPointer + foundCountPointer + infoPointer + where + !selectedCount = upperIndex - lowerIndex + 1 + +solveSelectedSymmetricValuesLapack :: + CInt -> + Int -> + Int -> + Int -> + Ptr CDouble -> + IO (Either MoonlightError (U.Vector Double)) +solveSelectedSymmetricValuesLapack + !lapackSize + !matrixSize + !lowerIndex + !upperIndex + matrixPointer = + case + (,) + <$> checkedProduct "LAPACK DSYEVX values floating workspace" 8 matrixSize + <*> checkedProduct "LAPACK DSYEVX values integer workspace" 5 matrixSize + of + Left err -> pure (Left err) + Right (workspaceCount, integerWorkspaceCount) -> withLapackChar 'N' $ \jobPointer -> + withLapackChar 'I' $ \rangePointer -> + withLapackChar 'U' $ \uploPointer -> + with lapackSize $ \sizePointer -> + with lapackSize $ \leadingDimensionPointer -> + with 0.0 $ \lowerValuePointer -> + with 0.0 $ \upperValuePointer -> + with (fromIntegral lowerIndex) $ \lowerIndexPointer -> + with (fromIntegral upperIndex) $ \upperIndexPointer -> + with 0.0 $ \absoluteTolerancePointer -> + alloca $ \foundCountPointer -> + allocaArray matrixSize $ \eigenvaluePointer -> + allocaArray 1 $ \eigenvectorPointer -> + with (1 :: CInt) $ \eigenvectorLeadingDimensionPointer -> + allocaArray workspaceCount $ \workspacePointer -> + with (fromIntegral workspaceCount) $ \workspaceSizePointer -> + allocaArray integerWorkspaceCount $ \integerWorkspacePointer -> + allocaArray 1 $ \failedVectorPointer -> + alloca $ \infoPointer -> do + poke foundCountPointer 0 + poke infoPointer 0 + lapackDsyevx + jobPointer + rangePointer + uploPointer + sizePointer + matrixPointer + leadingDimensionPointer + lowerValuePointer + upperValuePointer + lowerIndexPointer + upperIndexPointer + absoluteTolerancePointer + foundCountPointer + eigenvaluePointer + eigenvectorPointer + eigenvectorLeadingDimensionPointer + workspacePointer + workspaceSizePointer + integerWorkspacePointer + failedVectorPointer + infoPointer + decodeSelectedSymmetricValues + selectedCount + eigenvaluePointer + foundCountPointer + infoPointer + where + !selectedCount = upperIndex - lowerIndex + 1 + +solveSelectedSymmetricValuesDsyev :: + CInt -> + Int -> + Int -> + Int -> + Ptr CDouble -> + IO (Either MoonlightError (U.Vector Double)) +solveSelectedSymmetricValuesDsyev + !lapackSize + !matrixSize + !lowerIndex + !upperIndex + matrixPointer = + case checkedProduct "LAPACK DSYEV values workspace" 66 matrixSize of + Left err -> pure (Left err) + Right workspaceCount -> withLapackChar 'N' $ \jobPointer -> + withLapackChar 'U' $ \uploPointer -> + with lapackSize $ \sizePointer -> + with lapackSize $ \leadingDimensionPointer -> + allocaArray matrixSize $ \eigenvaluePointer -> + allocaArray workspaceCount $ \workspacePointer -> + with (fromIntegral workspaceCount) $ \workspaceSizePointer -> + alloca $ \infoPointer -> do + poke infoPointer 0 + lapackDsyev + jobPointer + uploPointer + sizePointer + matrixPointer + leadingDimensionPointer + eigenvaluePointer + workspacePointer + workspaceSizePointer + infoPointer + infoValue <- peek infoPointer + if infoValue /= 0 + then pure (Left (lapackInfoError "LAPACK DSYEV" infoValue)) + else + Right + <$> peekCDoubleVectorSlice + (lowerIndex - 1) + selectedCount + eigenvaluePointer + where + !selectedCount = upperIndex - lowerIndex + 1 + +queryWorkspace :: + Ptr CChar -> + Ptr CChar -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + IO (Either MoonlightError Int) +queryWorkspace jobPointer uploPointer sizePointer matrixPointer leadingDimensionPointer eigenvaluePointer infoPointer = + alloca $ \workspaceQueryPointer -> + with (-1) $ \workspaceSizePointer -> do + poke infoPointer 0 + lapackDsyev + jobPointer + uploPointer + sizePointer + matrixPointer + leadingDimensionPointer + eigenvaluePointer + workspaceQueryPointer + workspaceSizePointer + infoPointer + infoValue <- peek infoPointer + workspaceQuery <- peek workspaceQueryPointer + pure + ( if infoValue == 0 + then checkedLapackWorkspaceQuery "LAPACK DSYEV workspace query" (realToFrac workspaceQuery) + else Left (lapackInfoError "LAPACK DSYEV workspace query" infoValue) + ) + +solveLeastSquaresLapack :: CInt -> CInt -> Int -> Int -> [[Double]] -> [Double] -> IO (Either MoonlightError [Double]) +solveLeastSquaresLapack !lapackRows !lapackColumns !rowCount !columnCount matrixToRows rightHandSide = + withLapackChar 'N' $ \transPointer -> + with lapackRows $ \rowPointer -> + with lapackColumns $ \columnPointer -> + with (1 :: CInt) $ \rightHandSideCountPointer -> + with lapackRows $ \leadingDimensionPointer -> + with (max lapackRows lapackColumns) $ \rightHandSideLeadingDimensionPointer -> + withArray (toColumnMajor matrixToRows) $ \matrixPointer -> + withArray (leastSquaresRightHandSidePayload rowCount columnCount rightHandSide) $ \rightHandSidePointer -> + alloca $ \infoPointer -> + queryLeastSquaresWorkspace + transPointer + rowPointer + columnPointer + rightHandSideCountPointer + matrixPointer + leadingDimensionPointer + rightHandSidePointer + rightHandSideLeadingDimensionPointer + infoPointer + >>= \case + Left err -> pure (Left err) + Right workspaceSize -> + allocaArray workspaceSize $ \workspacePointer -> do + poke infoPointer 0 + with (fromIntegral workspaceSize) $ \workspaceSizePointer -> do + lapackDgels + transPointer + rowPointer + columnPointer + rightHandSideCountPointer + matrixPointer + leadingDimensionPointer + rightHandSidePointer + rightHandSideLeadingDimensionPointer + workspacePointer + workspaceSizePointer + infoPointer + decodeLeastSquares columnCount rightHandSidePointer infoPointer + +queryLeastSquaresWorkspace :: + Ptr CChar -> + Ptr CInt -> + Ptr CInt -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + IO (Either MoonlightError Int) +queryLeastSquaresWorkspace transPointer rowPointer columnPointer rightHandSideCountPointer matrixPointer leadingDimensionPointer rightHandSidePointer rightHandSideLeadingDimensionPointer infoPointer = + alloca $ \workspaceQueryPointer -> + with (-1) $ \workspaceSizePointer -> do + poke infoPointer 0 + lapackDgels + transPointer + rowPointer + columnPointer + rightHandSideCountPointer + matrixPointer + leadingDimensionPointer + rightHandSidePointer + rightHandSideLeadingDimensionPointer + workspaceQueryPointer + workspaceSizePointer + infoPointer + infoValue <- peek infoPointer + workspaceQuery <- peek workspaceQueryPointer + pure + ( if infoValue == 0 + then checkedLapackWorkspaceQuery "LAPACK DGELS workspace query" (realToFrac workspaceQuery) + else Left (lapackInfoError "LAPACK DGELS workspace query" infoValue) + ) + +checkedLapackWorkspaceQuery :: String -> Double -> Either MoonlightError Int +checkedLapackWorkspaceQuery context workspaceQuery + | not (fieldValueValid workspaceQuery) || workspaceQuery < 0.0 = + Left (InvariantViolation (context <> " returned an invalid workspace cardinality")) + | roundedWorkspace > toInteger (maxBound :: Int) = + Left (InvariantViolation (context <> " exceeds Int storage range")) + | otherwise = + Right (max 1 (fromInteger roundedWorkspace)) + where + roundedWorkspace = ceiling workspaceQuery :: Integer + +decodeLeastSquares :: Int -> Ptr CDouble -> Ptr CInt -> IO (Either MoonlightError [Double]) +decodeLeastSquares !columnCount rightHandSidePointer infoPointer = do + infoValue <- peek infoPointer + if infoValue /= 0 + then pure (Left (lapackLeastSquaresInfoError infoValue)) + else Right . take columnCount . fromCDoubleList <$> peekArray columnCount rightHandSidePointer + +solveSelectedTridiagonalLapack :: CInt -> Int -> Int -> Int -> [Double] -> [Double] -> IO (Either MoonlightError [(Double, [Double])]) +solveSelectedTridiagonalLapack !lapackSize !matrixSize !lowerIndex !upperIndex diagonalValues offDiagonalValues = + case + (,,,) + <$> checkedProduct "LAPACK DSTEMR eigenvector workspace" matrixSize selectedCount + <*> checkedProduct "LAPACK DSTEMR support workspace" 2 selectedCount + <*> checkedProduct "LAPACK DSTEMR floating workspace" 18 matrixSize + <*> checkedProduct "LAPACK DSTEMR integer workspace" 10 matrixSize + of + Left err -> pure (Left err) + Right (eigenvectorEntryCount, supportCount, workspaceCount, integerWorkspaceCount) -> withLapackChar 'V' $ \jobPointer -> + withLapackChar 'I' $ \rangePointer -> + with lapackSize $ \sizePointer -> + withArray (toCDoubleList diagonalValues) $ \diagonalPointer -> + withArray (toCDoubleList (offDiagonalValues <> [0.0])) $ \offDiagonalPointer -> + with 0.0 $ \lowerValuePointer -> + with 0.0 $ \upperValuePointer -> + with (fromIntegral lowerIndex) $ \lowerIndexPointer -> + with (fromIntegral upperIndex) $ \upperIndexPointer -> + alloca $ \foundCountPointer -> + allocaArray matrixSize $ \eigenvaluePointer -> + allocaArray eigenvectorEntryCount $ \eigenvectorPointer -> + with lapackSize $ \eigenvectorLeadingDimensionPointer -> + with (fromIntegral selectedCount) $ \eigenvectorColumnCountPointer -> + allocaArray supportCount $ \supportPointer -> + with (0 :: CInt) $ \tryRacPointer -> + allocaArray workspaceCount $ \workspacePointer -> + with (fromIntegral workspaceCount) $ \workspaceSizePointer -> + allocaArray integerWorkspaceCount $ \integerWorkspacePointer -> + with (fromIntegral integerWorkspaceCount) $ \integerWorkspaceSizePointer -> + alloca $ \infoPointer -> do + poke foundCountPointer 0 + poke infoPointer 0 + lapackDstemr + jobPointer + rangePointer + sizePointer + diagonalPointer + offDiagonalPointer + lowerValuePointer + upperValuePointer + lowerIndexPointer + upperIndexPointer + foundCountPointer + eigenvaluePointer + eigenvectorPointer + eigenvectorLeadingDimensionPointer + eigenvectorColumnCountPointer + supportPointer + tryRacPointer + workspacePointer + workspaceSizePointer + integerWorkspacePointer + integerWorkspaceSizePointer + infoPointer + decodeSelectedTridiagonal matrixSize selectedCount eigenvectorEntryCount eigenvaluePointer eigenvectorPointer foundCountPointer infoPointer + where + selectedCount = upperIndex - lowerIndex + 1 + +solveSelectedTridiagonalValuesLapack :: CInt -> Int -> Int -> Int -> [Double] -> [Double] -> IO (Either MoonlightError (U.Vector Double)) +solveSelectedTridiagonalValuesLapack !lapackSize !matrixSize !lowerIndex !upperIndex diagonalValues offDiagonalValues = + case + (,) + <$> checkedProduct "LAPACK DSTEMR values floating workspace" 18 matrixSize + <*> checkedProduct "LAPACK DSTEMR values integer workspace" 10 matrixSize + of + Left err -> pure (Left err) + Right (workspaceCount, integerWorkspaceCount) -> withLapackChar 'N' $ \jobPointer -> + withLapackChar 'I' $ \rangePointer -> + with lapackSize $ \sizePointer -> + withArray (toCDoubleList diagonalValues) $ \diagonalPointer -> + withArray (toCDoubleList (offDiagonalValues <> [0.0])) $ \offDiagonalPointer -> + with 0.0 $ \lowerValuePointer -> + with 0.0 $ \upperValuePointer -> + with (fromIntegral lowerIndex) $ \lowerIndexPointer -> + with (fromIntegral upperIndex) $ \upperIndexPointer -> + alloca $ \foundCountPointer -> + allocaArray matrixSize $ \eigenvaluePointer -> + allocaArray 1 $ \eigenvectorPointer -> + with (1 :: CInt) $ \eigenvectorLeadingDimensionPointer -> + with (1 :: CInt) $ \eigenvectorColumnCountPointer -> + allocaArray 1 $ \supportPointer -> + with (0 :: CInt) $ \tryRacPointer -> + allocaArray workspaceCount $ \workspacePointer -> + with (fromIntegral workspaceCount) $ \workspaceSizePointer -> + allocaArray integerWorkspaceCount $ \integerWorkspacePointer -> + with (fromIntegral integerWorkspaceCount) $ \integerWorkspaceSizePointer -> + alloca $ \infoPointer -> do + poke foundCountPointer 0 + poke infoPointer 0 + lapackDstemr + jobPointer + rangePointer + sizePointer + diagonalPointer + offDiagonalPointer + lowerValuePointer + upperValuePointer + lowerIndexPointer + upperIndexPointer + foundCountPointer + eigenvaluePointer + eigenvectorPointer + eigenvectorLeadingDimensionPointer + eigenvectorColumnCountPointer + supportPointer + tryRacPointer + workspacePointer + workspaceSizePointer + integerWorkspacePointer + integerWorkspaceSizePointer + infoPointer + decodeSelectedTridiagonalValues selectedCount eigenvaluePointer foundCountPointer infoPointer + where + selectedCount = upperIndex - lowerIndex + 1 + +solveSelectedBandPairsLapack :: + CInt -> + CInt -> + CInt -> + Int -> + Int -> + Int -> + U.Vector Double -> + IO (Either MoonlightError [(Double, [Double])]) +solveSelectedBandPairsLapack !lapackSize !lapackBandwidth !leadingDimension !matrixSize !lowerIndex !upperIndex lowerBandPayload = + case + (,,,) + <$> checkedProduct "LAPACK selected band orthogonal workspace" matrixSize matrixSize + <*> checkedProduct "LAPACK selected band eigenvector workspace" matrixSize selectedCount + <*> checkedProduct "LAPACK DSBEVX floating workspace" 7 matrixSize + <*> checkedProduct "LAPACK DSBEVX integer workspace" 5 matrixSize + of + Left err -> pure (Left err) + Right (orthogonalEntryCount, eigenvectorEntryCount, workspaceCount, integerWorkspaceCount) -> withLapackChar 'V' $ \jobPointer -> + withLapackChar 'I' $ \rangePointer -> + withLapackChar 'L' $ \uploPointer -> + with lapackSize $ \sizePointer -> + with lapackBandwidth $ \bandwidthPointer -> + withArray (toCDoubleList (U.toList lowerBandPayload)) $ \bandPointer -> + with leadingDimension $ \leadingDimensionPointer -> + allocaArray orthogonalEntryCount $ \orthogonalMatrixPointer -> + with lapackSize $ \orthogonalLeadingDimensionPointer -> + with 0.0 $ \lowerValuePointer -> + with 0.0 $ \upperValuePointer -> + with (fromIntegral lowerIndex) $ \lowerIndexPointer -> + with (fromIntegral upperIndex) $ \upperIndexPointer -> + with 0.0 $ \absoluteTolerancePointer -> + alloca $ \foundCountPointer -> + allocaArray matrixSize $ \eigenvaluePointer -> + allocaArray eigenvectorEntryCount $ \eigenvectorPointer -> + with lapackSize $ \eigenvectorLeadingDimensionPointer -> + allocaArray workspaceCount $ \workspacePointer -> + allocaArray integerWorkspaceCount $ \integerWorkspacePointer -> + allocaArray matrixSize $ \failedVectorPointer -> + alloca $ \infoPointer -> do + poke foundCountPointer 0 + poke infoPointer 0 + lapackDsbevx + jobPointer + rangePointer + uploPointer + sizePointer + bandwidthPointer + bandPointer + leadingDimensionPointer + orthogonalMatrixPointer + orthogonalLeadingDimensionPointer + lowerValuePointer + upperValuePointer + lowerIndexPointer + upperIndexPointer + absoluteTolerancePointer + foundCountPointer + eigenvaluePointer + eigenvectorPointer + eigenvectorLeadingDimensionPointer + workspacePointer + integerWorkspacePointer + failedVectorPointer + infoPointer + decodeSelectedBand matrixSize selectedCount eigenvectorEntryCount eigenvaluePointer eigenvectorPointer foundCountPointer infoPointer + where + selectedCount = upperIndex - lowerIndex + 1 + +solveSelectedBandValuesLapack :: + CInt -> + CInt -> + CInt -> + Int -> + Int -> + Int -> + U.Vector Double -> + IO (Either MoonlightError (U.Vector Double)) +solveSelectedBandValuesLapack !lapackSize !lapackBandwidth !leadingDimension !matrixSize !lowerIndex !upperIndex lowerBandPayload = + case + (,) + <$> checkedProduct "LAPACK DSBEVX values floating workspace" 7 matrixSize + <*> checkedProduct "LAPACK DSBEVX values integer workspace" 5 matrixSize + of + Left err -> pure (Left err) + Right (workspaceCount, integerWorkspaceCount) -> withLapackChar 'N' $ \jobPointer -> + withLapackChar 'I' $ \rangePointer -> + withLapackChar 'L' $ \uploPointer -> + with lapackSize $ \sizePointer -> + with lapackBandwidth $ \bandwidthPointer -> + withArray (toCDoubleList (U.toList lowerBandPayload)) $ \bandPointer -> + with leadingDimension $ \leadingDimensionPointer -> + allocaArray 1 $ \orthogonalMatrixPointer -> + with (1 :: CInt) $ \orthogonalLeadingDimensionPointer -> + with 0.0 $ \lowerValuePointer -> + with 0.0 $ \upperValuePointer -> + with (fromIntegral lowerIndex) $ \lowerIndexPointer -> + with (fromIntegral upperIndex) $ \upperIndexPointer -> + with 0.0 $ \absoluteTolerancePointer -> + alloca $ \foundCountPointer -> + allocaArray matrixSize $ \eigenvaluePointer -> + allocaArray 1 $ \eigenvectorPointer -> + with (1 :: CInt) $ \eigenvectorLeadingDimensionPointer -> + allocaArray workspaceCount $ \workspacePointer -> + allocaArray integerWorkspaceCount $ \integerWorkspacePointer -> + allocaArray 1 $ \failedVectorPointer -> + alloca $ \infoPointer -> do + poke foundCountPointer 0 + poke infoPointer 0 + lapackDsbevx + jobPointer + rangePointer + uploPointer + sizePointer + bandwidthPointer + bandPointer + leadingDimensionPointer + orthogonalMatrixPointer + orthogonalLeadingDimensionPointer + lowerValuePointer + upperValuePointer + lowerIndexPointer + upperIndexPointer + absoluteTolerancePointer + foundCountPointer + eigenvaluePointer + eigenvectorPointer + eigenvectorLeadingDimensionPointer + workspacePointer + integerWorkspacePointer + failedVectorPointer + infoPointer + decodeSelectedBandValues selectedCount eigenvaluePointer foundCountPointer infoPointer + where + selectedCount = upperIndex - lowerIndex + 1 + +decodeSelectedTridiagonal :: + Int -> + Int -> + Int -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + IO (Either MoonlightError [(Double, [Double])]) +decodeSelectedTridiagonal !matrixSize !selectedCount !eigenvectorEntryCount eigenvaluePointer eigenvectorPointer foundCountPointer infoPointer = do + infoValue <- peek infoPointer + foundCount <- fromIntegral <$> peek foundCountPointer + if infoValue /= 0 + then pure (Left (lapackInfoError "LAPACK DSTEMR" infoValue)) + else + if foundCount /= selectedCount + then pure (Left (InvariantViolation ("LAPACK DSTEMR returned " <> show foundCount <> " eigenpairs; expected " <> show selectedCount))) + else do + eigenvalues <- fromCDoubleList <$> peekArray selectedCount eigenvaluePointer + eigenvectorPayload <- fromCDoubleList <$> peekArray eigenvectorEntryCount eigenvectorPointer + pure + ( do + eigenvectors <- take selectedCount <$> chunkRows matrixSize eigenvectorPayload + Right (zip eigenvalues eigenvectors) + ) + +decodeSelectedTridiagonalValues :: + Int -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + IO (Either MoonlightError (U.Vector Double)) +decodeSelectedTridiagonalValues !selectedCount eigenvaluePointer foundCountPointer infoPointer = do + infoValue <- peek infoPointer + foundCount <- fromIntegral <$> peek foundCountPointer + if infoValue /= 0 + then pure (Left (lapackInfoError "LAPACK DSTEMR" infoValue)) + else + if foundCount /= selectedCount + then pure (Left (InvariantViolation ("LAPACK DSTEMR returned " <> show foundCount <> " eigenvalues; expected " <> show selectedCount))) + else Right . U.fromList . fromCDoubleList <$> peekArray selectedCount eigenvaluePointer + +decodeSelectedSymmetricColumns :: + Int -> + Int -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + IO + ( Either + MoonlightError + (U.Vector Double, U.Vector Double) + ) +decodeSelectedSymmetricColumns + !selectedCount + !eigenvectorEntryCount + eigenvaluePointer + eigenvectorPointer + foundCountPointer + infoPointer = do + infoValue <- peek infoPointer + foundCount <- fromIntegral <$> peek foundCountPointer + if infoValue /= 0 + then pure (Left (lapackInfoError "LAPACK DSYEVX" infoValue)) + else + if foundCount /= selectedCount + then + pure + ( Left + ( InvariantViolation + ( "LAPACK DSYEVX returned " + <> show foundCount + <> " eigenpairs; expected " + <> show selectedCount + ) + ) + ) + else do + eigenvalues <- + peekCDoubleVectorSlice + 0 + selectedCount + eigenvaluePointer + eigenvectors <- + peekCDoubleVectorSlice + 0 + eigenvectorEntryCount + eigenvectorPointer + pure (Right (eigenvalues, eigenvectors)) + +decodeSelectedSymmetricValues :: + Int -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + IO (Either MoonlightError (U.Vector Double)) +decodeSelectedSymmetricValues + !selectedCount + eigenvaluePointer + foundCountPointer + infoPointer = do + infoValue <- peek infoPointer + foundCount <- fromIntegral <$> peek foundCountPointer + if infoValue /= 0 + then pure (Left (lapackInfoError "LAPACK DSYEVX" infoValue)) + else + if foundCount /= selectedCount + then + pure + ( Left + ( InvariantViolation + ( "LAPACK DSYEVX returned " + <> show foundCount + <> " eigenvalues; expected " + <> show selectedCount + ) + ) + ) + else + Right + <$> peekCDoubleVectorSlice + 0 + selectedCount + eigenvaluePointer + +peekCDoubleVectorSlice :: + Int -> + Int -> + Ptr CDouble -> + IO (U.Vector Double) +peekCDoubleVectorSlice !sourceOffset !elementCount sourcePointer = + U.generateM elementCount $ \entryIndex -> do + CDouble entryValue <- + peekElemOff + sourcePointer + (sourceOffset + entryIndex) + pure entryValue +{-# INLINE peekCDoubleVectorSlice #-} + +checkedProduct :: + String -> + Int -> + Int -> + Either MoonlightError Int +checkedProduct context leftCount rightCount = + first + (const (InvariantViolation (context <> " exceeds non-negative Int storage range"))) + (checkedNonNegativeProduct leftCount rightCount) + +checkedSum :: String -> Int -> Int -> Either MoonlightError Int +checkedSum context leftCount rightCount = + first + (const (InvariantViolation (context <> " exceeds non-negative Int storage range"))) + (checkedNonNegativeSum leftCount rightCount) + +decodeSelectedBand :: + Int -> + Int -> + Int -> + Ptr CDouble -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + IO (Either MoonlightError [(Double, [Double])]) +decodeSelectedBand !matrixSize !selectedCount !eigenvectorEntryCount eigenvaluePointer eigenvectorPointer foundCountPointer infoPointer = do + infoValue <- peek infoPointer + foundCount <- fromIntegral <$> peek foundCountPointer + if infoValue /= 0 + then pure (Left (lapackInfoError "LAPACK DSBEVX" infoValue)) + else + if foundCount /= selectedCount + then pure (Left (InvariantViolation ("LAPACK DSBEVX returned " <> show foundCount <> " eigenpairs; expected " <> show selectedCount))) + else do + eigenvalues <- fromCDoubleList <$> peekArray selectedCount eigenvaluePointer + eigenvectorPayload <- fromCDoubleList <$> peekArray eigenvectorEntryCount eigenvectorPointer + pure + ( do + eigenvectors <- take selectedCount <$> chunkRows matrixSize eigenvectorPayload + Right (zip eigenvalues eigenvectors) + ) + +decodeSelectedBandValues :: + Int -> + Ptr CDouble -> + Ptr CInt -> + Ptr CInt -> + IO (Either MoonlightError (U.Vector Double)) +decodeSelectedBandValues !selectedCount eigenvaluePointer foundCountPointer infoPointer = do + infoValue <- peek infoPointer + foundCount <- fromIntegral <$> peek foundCountPointer + if infoValue /= 0 + then pure (Left (lapackInfoError "LAPACK DSBEVX" infoValue)) + else + if foundCount /= selectedCount + then pure (Left (InvariantViolation ("LAPACK DSBEVX returned " <> show foundCount <> " eigenvalues; expected " <> show selectedCount))) + else Right . U.fromList . fromCDoubleList <$> peekArray selectedCount eigenvaluePointer + +withLapackChar :: Char -> (Ptr CChar -> IO value) -> IO value +withLapackChar charValue onPointer = + with (castCharToCChar charValue) onPointer + +toColumnMajor :: [[Double]] -> [CDouble] +toColumnMajor = + toCDoubleList . concat . transpose + +toCDoubleList :: [Double] -> [CDouble] +toCDoubleList = + fmap CDouble + +fromCDoubleList :: [CDouble] -> [Double] +fromCDoubleList = + fmap realToFrac + +matrixSizeAsLapackInt :: Int -> Either MoonlightError CInt +matrixSizeAsLapackInt matrixSize + | matrixSize < 0 = + Left (InvariantViolation "LAPACK matrix size must be non-negative") + | matrixSize > fromIntegral (maxBound :: CInt) = + Left (InvariantViolation "LAPACK matrix size exceeds CInt range") + | otherwise = Right (fromIntegral matrixSize) + +lapackInfoError :: String -> CInt -> MoonlightError +lapackInfoError context infoValue + | infoValue < 0 = InvariantViolation (context <> " rejected argument " <> show (negate infoValue)) + | otherwise = InvariantViolation (context <> " failed to converge; info=" <> show infoValue) + +lapackLeastSquaresInfoError :: CInt -> MoonlightError +lapackLeastSquaresInfoError infoValue + | infoValue < 0 = InvariantViolation ("LAPACK DGELS rejected argument " <> show (negate infoValue)) + | otherwise = InvariantViolation ("LAPACK DGELS detected exact rank deficiency at triangular factor diagonal " <> show infoValue) + +lapackLinearSolveInfoError :: CInt -> MoonlightError +lapackLinearSolveInfoError infoValue + | infoValue < 0 = InvariantViolation ("LAPACK DGESV rejected argument " <> show (negate infoValue)) + | otherwise = InvariantViolation ("LAPACK DGESV detected exact singularity at U diagonal " <> show infoValue) + +leastSquaresRightHandSidePayload :: Int -> Int -> [Double] -> [CDouble] +leastSquaresRightHandSidePayload rowCount columnCount rightHandSide = + toCDoubleList (rightHandSide <> replicate (max 0 (columnCount - rowCount)) 0.0) + +validateLeastSquaresInput :: Int -> Int -> [[Double]] -> [Double] -> Either MoonlightError () +validateLeastSquaresInput rowCount columnCount matrixToRows rightHandSide + | rowCount < 0 || columnCount < 0 = + Left (InvariantViolation "LAPACK least-squares dimensions must be non-negative") + | rowCount == 0 || columnCount == 0 = + Left (InvariantViolation "LAPACK least-squares requires positive dimensions") + | length matrixToRows /= rowCount = + Left (InvariantViolation "LAPACK least-squares row count mismatch") + | any ((/= columnCount) . length) matrixToRows = + Left (InvariantViolation "LAPACK least-squares requires rectangular rows") + | length rightHandSide /= rowCount = + Left (InvariantViolation "LAPACK least-squares RHS length mismatch") + | not (all fieldValueValid (concat matrixToRows <> rightHandSide)) = + Left (InvariantViolation "LAPACK least-squares requires finite entries") + | otherwise = Right () + +validateSelectedTridiagonalInput :: Int -> Int -> [Double] -> [Double] -> Either MoonlightError Int +validateSelectedTridiagonalInput lowerIndex upperIndex diagonalValues offDiagonalValues + | matrixSize <= 0 = + Left (InvariantViolation "LAPACK selected tridiagonal eigensolve requires a positive dimension") + | length offDiagonalValues /= matrixSize - 1 = + Left (InvariantViolation "LAPACK selected tridiagonal eigensolve off-diagonal length mismatch") + | lowerIndex < 1 || upperIndex < lowerIndex || upperIndex > matrixSize = + Left (InvariantViolation "LAPACK selected tridiagonal eigensolve index range is out of bounds") + | not (all fieldValueValid (diagonalValues <> offDiagonalValues)) = + Left (InvariantViolation "LAPACK selected tridiagonal eigensolve requires finite entries") + | otherwise = Right matrixSize + where + matrixSize = length diagonalValues + +validateSelectedBandInput :: Int -> Int -> Int -> Int -> U.Vector Double -> Either MoonlightError Int +validateSelectedBandInput lowerIndex upperIndex matrixSize bandwidth lowerBandPayload = do + if matrixSize <= 0 + then Left (InvariantViolation "LAPACK selected symmetric-band eigensolve requires a positive dimension") + else Right () + if bandwidth < 0 + then Left (InvariantViolation "LAPACK selected symmetric-band eigensolve bandwidth must be non-negative") + else Right () + if bandwidth >= matrixSize + then Left (InvariantViolation "LAPACK selected symmetric-band eigensolve bandwidth must be smaller than dimension") + else Right () + leadingDimension <- checkedSum "LAPACK selected symmetric-band leading dimension" bandwidth 1 + expectedPayloadLength <- checkedProduct "LAPACK selected symmetric-band payload" leadingDimension matrixSize + if U.length lowerBandPayload /= expectedPayloadLength + then + Left + ( InvariantViolation + ( "LAPACK selected symmetric-band payload length mismatch: expected " + <> show expectedPayloadLength + <> " but received " + <> show (U.length lowerBandPayload) + ) + ) + else Right () + if lowerIndex < 1 || upperIndex < lowerIndex || upperIndex > matrixSize + then Left (InvariantViolation "LAPACK selected symmetric-band eigensolve index range is out of bounds") + else Right () + if U.any (not . fieldValueValid) lowerBandPayload + then Left (InvariantViolation "LAPACK selected symmetric-band eigensolve requires finite entries") + else Right leadingDimension diff --git a/src-native/Moonlight/LinAlg/Native.hs b/src-native/Moonlight/LinAlg/Native.hs new file mode 100644 --- /dev/null +++ b/src-native/Moonlight/LinAlg/Native.hs @@ -0,0 +1,21 @@ +-- | Effectful native LAPACK backend boundary: on macOS it links Accelerate, elsewhere it expects BLAS\/LAPACK. +module Moonlight.LinAlg.Native + ( denseDoubleLinearSolveLapack, + denseDoubleMatrixProductBlas, + denseDoubleSymmetricEigenpairsLapack, + leastSquaresLapack, + symmetricEigenRequestLapack, + selectedSymmetricTridiagonalEigenRequestLapack, + selectedSymmetricBlockTridiagonalEigenRequestLapack, + ) +where + +import Moonlight.LinAlg.Effect.Native.Dispatch + ( denseDoubleLinearSolveLapack, + denseDoubleMatrixProductBlas, + denseDoubleSymmetricEigenpairsLapack, + leastSquaresLapack, + selectedSymmetricBlockTridiagonalEigenRequestLapack, + selectedSymmetricTridiagonalEigenRequestLapack, + symmetricEigenRequestLapack, + ) diff --git a/src-public/Moonlight/LinAlg.hs b/src-public/Moonlight/LinAlg.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg.hs @@ -0,0 +1,80 @@ +{-| +Typed dense, sparse, finite-field, and Krylov linear algebra — the single public +surface of @moonlight-linalg@, layered over "Moonlight.Core" and +"Moonlight.Algebra". Dense matrices carry their shape and scalar as type indices, +so an @r@-by-@c@ matrix over @a@ is a @Matrix r c a@ and a product's shared inner +dimension is fixed at the type level. Construction is validated: malformed +matrices, incompatible dimensions, and invalid solver configuration return typed +@MoonlightError@ failures rather than partial indexing or runtime bottoms. + +The re-exported surface, by role: + +* Dense — "Moonlight.LinAlg.Dense": typed vectors and matrices, validated + rectangular row authoring, GF(2) and packed bit matrices, exterior algebra, + decompositions, field operations, direct solvers, and primitives. +* Geometry — "Moonlight.LinAlg.Geometry": @Vec2@, @Vec3@, AABB\/AABB2, frames, + affine transforms, and compact symmetric 2D\/3D carriers. +* Sparse — "Moonlight.LinAlg.Sparse": COO\/CSR\/CSC\/packed encodings, sealed + preconditioner families, and sparse iterative solvers for graph and mesh + operators. +* Operators and spectra — "Moonlight.LinAlg.Operator" affine-normalized linear + operators with explicit self-adjoint boundaries, "Moonlight.LinAlg.Spectral" + eigenvalue\/eigenpair demand dispatched by operator structure, and + "Moonlight.LinAlg.Krylov" Arnoldi\/Lanczos decompositions with projected + tridiagonal and block-tridiagonal carriers. +* Domain and statics — "Moonlight.LinAlg.Domain" domain-level algebra including + Smith normal form, and "Moonlight.LinAlg.Statics" assembly, equilibrium + compilation, and support checking. + +The effectful native LAPACK boundary lives apart in "Moonlight.LinAlg.Native" +(not re-exported here); the @Moonlight.LinAlg.Pure.*@ and +@Moonlight.LinAlg.Internal.*@ leaves live in graded implementation sublibraries +behind this vocabulary. + +/Quick start./ A whole computation composes in @Either MoonlightError@: + +> import Moonlight.LinAlg (Matrix, fromListMatrix, mult, toListMatrix) +> import Moonlight.Core (MoonlightError) +> +> product22 :: Either MoonlightError (Matrix 2 2 Double) +> product22 = do +> left <- fromListMatrix @2 @2 @Double [1.0, 2.0, 3.0, 4.0] +> right <- fromListMatrix @2 @2 @Double [2.0, 0.0, 1.0, 2.0] +> mult left right + +@mult :: Matrix r m a -> Matrix m c a -> Either MoonlightError (Matrix r c a)@ +forces the two matrices to meet on @m@; the result reads back as +@toListMatrix \<$\> product22 == Right [4.0, 4.0, 10.0, 8.0]@. +-} +module Moonlight.LinAlg + ( module Dense, + module DenseBlock, + module DenseDecomposition, + module DenseExterior, + module DenseField, + module DenseGF2, + module DenseSolver, + module Domain, + module Geometry, + module Krylov, + module Operator, + module Sparse, + module Spectral, + module Statics, + ) +where + +import Moonlight.LinAlg.Dense as Dense +import Moonlight.LinAlg.Dense.Block as DenseBlock +import Moonlight.LinAlg.Dense.Decomposition as DenseDecomposition +import Moonlight.LinAlg.Dense.Exterior as DenseExterior +import Moonlight.LinAlg.Dense.Field as DenseField +import Moonlight.LinAlg.Dense.GF2 as DenseGF2 +import Moonlight.LinAlg.Dense.Solver as DenseSolver +import Moonlight.LinAlg.Domain as Domain +import Moonlight.LinAlg.Geometry as Geometry +import Moonlight.LinAlg.Krylov as Krylov +import Moonlight.LinAlg.Operator as Operator +import Moonlight.LinAlg.Sparse as Sparse +import Moonlight.LinAlg.Spectral as Spectral +import Moonlight.LinAlg.Statics as Statics diff --git a/src-public/Moonlight/LinAlg/Dense.hs b/src-public/Moonlight/LinAlg/Dense.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Dense.hs @@ -0,0 +1,87 @@ +-- | Dense vectors and matrices: validated row authoring, GF(2), exterior algebra, decompositions, field operations, direct solvers, and primitives. +module Moonlight.LinAlg.Dense + ( Vector, + Matrix, + fromListVector, + fromListMatrix, + matrixRows, + toListVector, + toListMatrix, + vectorLength, + matrixShape, + matrixToRows, + DynVector, + DynMatrix, + DenseDoubleMatrix, + mkDynVector, + mkDynMatrix, + mkDenseDoubleMatrixRowMajor, + mkDenseDoubleMatrixRows, + denseDoubleMatrixShape, + denseDoubleMatrixToRowMajorVector, + denseDoubleMatrixToRows, + denseDoubleMatrixVectorProduct, + dynMatrixFromRows, + dynMatrixToRows, + toDynVector, + toDynMatrix, + fromDynVector, + fromDynMatrix, + withDynVector, + withDynMatrix, + dynVectorLength, + dynMatrixShape, + dynVectorToList, + dynMatrixToList, + mapMatrix, + add, + mult, + transpose, + ) +where + +import Moonlight.LinAlg.Pure.Dense.Basic + ( add, + mapMatrix, + mult, + transpose, + ) +import Moonlight.LinAlg.Pure.Dense.Dynamic + ( DynMatrix, + DynVector, + dynMatrixFromRows, + dynMatrixShape, + dynMatrixToList, + dynMatrixToRows, + dynVectorLength, + dynVectorToList, + fromDynMatrix, + fromDynVector, + mkDynMatrix, + mkDynVector, + toDynMatrix, + toDynVector, + withDynMatrix, + withDynVector, + ) +import Moonlight.LinAlg.Pure.Dense.Flat + ( DenseDoubleMatrix, + denseDoubleMatrixShape, + denseDoubleMatrixToRowMajorVector, + denseDoubleMatrixToRows, + denseDoubleMatrixVectorProduct, + mkDenseDoubleMatrixRowMajor, + mkDenseDoubleMatrixRows, + ) +import Moonlight.LinAlg.Pure.Dense.Types + ( Matrix, + Vector, + fromListMatrix, + fromListVector, + matrixRows, + matrixShape, + matrixToRows, + toListMatrix, + toListVector, + vectorLength, + ) diff --git a/src-public/Moonlight/LinAlg/Dense/Block.hs b/src-public/Moonlight/LinAlg/Dense/Block.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Dense/Block.hs @@ -0,0 +1,15 @@ +-- | Block-matrix inversion over the rationals, GF(2), and the unimodular integers. +module Moonlight.LinAlg.Dense.Block + ( BlockMatrixFailure (..), + invertRationalBlock, + invertGF2Block, + invertUnimodularIntegerBlock, + ) +where + +import Moonlight.LinAlg.Pure.Dense.Block + ( BlockMatrixFailure (..), + invertGF2Block, + invertRationalBlock, + invertUnimodularIntegerBlock, + ) diff --git a/src-public/Moonlight/LinAlg/Dense/Decomposition.hs b/src-public/Moonlight/LinAlg/Dense/Decomposition.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Dense/Decomposition.hs @@ -0,0 +1,17 @@ +-- | Dense matrix factorizations: QR, Cholesky, symmetric eigendecomposition, and thin SVD. +module Moonlight.LinAlg.Dense.Decomposition + ( qrDecompFullColumnRank, + choleskyDecomp, + symmetricEigen, + symmetricEigenPairs, + thinSvdFullColumnRank, + ) +where + +import Moonlight.LinAlg.Pure.Dense.Decomposition + ( choleskyDecomp, + qrDecompFullColumnRank, + symmetricEigen, + symmetricEigenPairs, + thinSvdFullColumnRank, + ) diff --git a/src-public/Moonlight/LinAlg/Dense/Exterior.hs b/src-public/Moonlight/LinAlg/Dense/Exterior.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Dense/Exterior.hs @@ -0,0 +1,21 @@ +-- | Exterior powers of a dense matrix and their basis bookkeeping. +module Moonlight.LinAlg.Dense.Exterior + ( ExteriorBasis (..), + ExteriorPowerFailure (..), + choose, + exteriorBasis, + exteriorBasisCardinality, + exteriorPowerMatrix, + exteriorPowerMatrixWithShape, + ) +where + +import Moonlight.LinAlg.Pure.Dense.Exterior + ( ExteriorBasis (..), + ExteriorPowerFailure (..), + choose, + exteriorBasis, + exteriorBasisCardinality, + exteriorPowerMatrix, + exteriorPowerMatrixWithShape, + ) diff --git a/src-public/Moonlight/LinAlg/Dense/Field.hs b/src-public/Moonlight/LinAlg/Dense/Field.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Dense/Field.hs @@ -0,0 +1,19 @@ +-- | Field-level dense operations: full-rank PLU, rank, and kernel bases. +module Moonlight.LinAlg.Dense.Field + ( DenseRankBackend, + PLU (..), + KernelBasis (..), + pluDecompFullRank, + rank, + kernel, + ) +where + +import Moonlight.LinAlg.Pure.Dense.Field + ( DenseRankBackend, + KernelBasis (..), + PLU (..), + kernel, + pluDecompFullRank, + rank, + ) diff --git a/src-public/Moonlight/LinAlg/Dense/GF2.hs b/src-public/Moonlight/LinAlg/Dense/GF2.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Dense/GF2.hs @@ -0,0 +1,135 @@ +-- | GF(2) scalars, packed bit rows, and packed GF(2) linear maps for boundary and incidence computations. +module Moonlight.LinAlg.Dense.GF2 + ( GF2 (..), + gf2Zero, + gf2One, + gf2FromBool, + gf2ToBool, + PackedRow, + packedRowWidth, + packedRowNonZeroCount, + emptyPackedRow, + unitPackedRow, + packedRowFromIndices, + packedRowIndices, + packedRowMember, + packedRowIsZero, + packedRowXor, + packedRowRemap, + PackedLinearMap, + packedLinearMapDomain, + packedLinearMapCodomain, + packedLinearMapColumns, + packedLinearMapFromColumns, + packedLinearMapFromEntries, + zeroPackedLinearMap, + identityPackedLinearMap, + applyPackedLinearMap, + composePackedLinearMaps, + addPackedLinearMaps, + packedLinearMapIsZero, + PackedSpan, + emptyPackedSpan, + packedSpanFromRows, + reducePackedRow, + admitPackedRow, + ColumnReduction (..), + reducePackedColumns, + PackedCoordinateSolver, + packedCoordinateSolver, + coordinatesInPackedBasis, + inverseFromPackedBasisColumns, + GF2MatrixEntry (..), + GF2PackedMatrix, + gf2PackedRows, + gf2PackedColumns, + gf2PackedWordsPerRow, + gf2PackedWords, + GF2PackedMatrixFailure (..), + mkGF2PackedMatrix, + mkGF2PackedMatrixFromRowMajor, + rankGF2PackedMatrix, + gf2PackedMatrixLinearMap, + inverseGF2PackedMatrix, + GF2SparseColumn, + gf2SparseColumnIndex, + gf2SparseColumnRows, + mkGF2SparseColumn, + GF2SparseReducerConfig, + gf2SparseDensifyThreshold, + mkGF2SparseReducerConfig, + defaultGF2SparseReducerConfig, + GF2SparseColumnReduction (..), + reduceGF2SparseColumns, + rankGF2SparseColumns, + independentGF2SparseColumns, + kernelBasisGF2SparseColumns, + ) +where + +import Moonlight.LinAlg.Pure.Dense.GF2 + ( GF2 (..), + GF2MatrixEntry (..), + GF2PackedMatrix, + GF2PackedMatrixFailure (..), + GF2SparseColumn, + GF2SparseColumnReduction (..), + GF2SparseReducerConfig, + ColumnReduction (..), + PackedCoordinateSolver, + PackedLinearMap, + PackedRow, + PackedSpan, + addPackedLinearMaps, + admitPackedRow, + applyPackedLinearMap, + composePackedLinearMaps, + coordinatesInPackedBasis, + defaultGF2SparseReducerConfig, + emptyPackedRow, + emptyPackedSpan, + gf2FromBool, + gf2One, + gf2PackedMatrixLinearMap, + gf2PackedColumns, + gf2PackedRows, + gf2PackedWords, + gf2PackedWordsPerRow, + gf2SparseColumnIndex, + gf2SparseColumnRows, + gf2SparseDensifyThreshold, + gf2ToBool, + gf2Zero, + identityPackedLinearMap, + independentGF2SparseColumns, + inverseFromPackedBasisColumns, + inverseGF2PackedMatrix, + kernelBasisGF2SparseColumns, + mkGF2SparseColumn, + mkGF2SparseReducerConfig, + mkGF2PackedMatrix, + mkGF2PackedMatrixFromRowMajor, + packedCoordinateSolver, + packedLinearMapCodomain, + packedLinearMapColumns, + packedLinearMapDomain, + packedLinearMapFromColumns, + packedLinearMapFromEntries, + packedLinearMapIsZero, + packedRowFromIndices, + packedRowIndices, + packedRowIsZero, + packedRowMember, + packedRowNonZeroCount, + packedRowRemap, + packedRowWidth, + packedRowXor, + packedSpanFromRows, + rankGF2SparseColumns, + rankGF2PackedMatrix, + reduceGF2SparseColumns, + reducePackedColumns, + reducePackedRow, + unitPackedRow, + zeroPackedLinearMap, + ) diff --git a/src-public/Moonlight/LinAlg/Dense/Primitives.hs b/src-public/Moonlight/LinAlg/Dense/Primitives.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Dense/Primitives.hs @@ -0,0 +1,29 @@ +-- | Low-level vector and matrix primitives: dot products, norms, scaling, outer products, and linear combinations. +module Moonlight.LinAlg.Dense.Primitives + ( dotProduct, + vectorNorm, + scaleVector, + addVector, + subVector, + matrixVectorProduct, + matrixSubtract, + scaleMatrix, + outerProduct, + basisVector, + linearCombination, + ) +where + +import Moonlight.LinAlg.Internal.Primitives + ( addVector, + basisVector, + dotProduct, + linearCombination, + matrixSubtract, + matrixVectorProduct, + outerProduct, + scaleMatrix, + scaleVector, + subVector, + vectorNorm, + ) diff --git a/src-public/Moonlight/LinAlg/Dense/Rows.hs b/src-public/Moonlight/LinAlg/Dense/Rows.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Dense/Rows.hs @@ -0,0 +1,36 @@ +-- | Validated rectangular row authoring surface. +-- +-- This facade preserves nested-row shape failures as typed errors. It is not +-- the hot dense-storage owner; benchmark-sensitive code should use sealed +-- vector, sparse, tridiagonal, or native owners instead of pretending lists are +-- a BLAS implementation. +module Moonlight.LinAlg.Dense.Rows + ( DenseRows, + mkDenseRows, + mkDenseRowsWithShape, + mkDenseRowsFromFlat, + denseRowsShape, + denseRowsToLists, + transposeRowsExact, + zipRowsExactWith, + matrixVectorProductRowsWith, + matrixProductRowsWith, + hcatRowsExact, + vcatRowsExact, + ) +where + +import Moonlight.LinAlg.Pure.Dense.Rows + ( DenseRows, + denseRowsShape, + denseRowsToLists, + hcatRowsExact, + matrixProductRowsWith, + matrixVectorProductRowsWith, + mkDenseRows, + mkDenseRowsFromFlat, + mkDenseRowsWithShape, + transposeRowsExact, + vcatRowsExact, + zipRowsExactWith, + ) diff --git a/src-public/Moonlight/LinAlg/Dense/Solver.hs b/src-public/Moonlight/LinAlg/Dense/Solver.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Dense/Solver.hs @@ -0,0 +1,13 @@ +-- | Dense linear-system solvers: direct, conjugate-gradient, and GMRES. +module Moonlight.LinAlg.Dense.Solver + ( solveDirect, + solveCG, + solveGMRES, + ) +where + +import Moonlight.LinAlg.Pure.Dense.Solver + ( solveCG, + solveDirect, + solveGMRES, + ) diff --git a/src-public/Moonlight/LinAlg/Domain.hs b/src-public/Moonlight/LinAlg/Domain.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Domain.hs @@ -0,0 +1,9 @@ +-- | Domain-level algebraic operations over integer and Euclidean domains, including Smith normal form. +module Moonlight.LinAlg.Domain + ( module Bareiss, + module Smith, + ) +where + +import Moonlight.LinAlg.Pure.Domain.Bareiss as Bareiss +import Moonlight.LinAlg.Pure.Domain.Smith as Smith diff --git a/src-public/Moonlight/LinAlg/Geometry.hs b/src-public/Moonlight/LinAlg/Geometry.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Geometry.hs @@ -0,0 +1,19 @@ +-- | @Vec2@, @Vec3@, AABB\/AABB2, frames, affine transforms, and compact symmetric 2D\/3D carriers. +module Moonlight.LinAlg.Geometry + ( module AABB, + module AABB2, + module Affine, + module Frame, + module Symmetric, + module Vec2, + module Vec3, + ) +where + +import Moonlight.LinAlg.Pure.Geometry.AABB as AABB +import Moonlight.LinAlg.Pure.Geometry.AABB2 as AABB2 +import Moonlight.LinAlg.Pure.Geometry.Frame as Frame +import Moonlight.LinAlg.Pure.Geometry.Symmetric as Symmetric +import Moonlight.LinAlg.Pure.Geometry.Transform.Affine as Affine +import Moonlight.LinAlg.Pure.Geometry.Vec2 as Vec2 +import Moonlight.LinAlg.Pure.Geometry.Vec3 as Vec3 diff --git a/src-public/Moonlight/LinAlg/Krylov.hs b/src-public/Moonlight/LinAlg/Krylov.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Krylov.hs @@ -0,0 +1,70 @@ +-- | Public Arnoldi\/Lanczos decompositions, projected tridiagonal and block-tridiagonal carriers, and the block Lanczos surface. +module Moonlight.LinAlg.Krylov + ( KrylovConfigError (..), + krylovConfigErrorMessage, + PositiveCount, + mkPositiveCount, + positiveCountValue, + NonNegativeConfigTolerance, + mkNonNegativeConfigTolerance, + nonNegativeConfigToleranceValue, + ArnoldiConfig, + mkArnoldiConfig, + arnoldiIterations, + arnoldiTolerance, + arnoldiReorthogonalize, + withArnoldiIterations, + withArnoldiTolerance, + withArnoldiReorthogonalize, + defaultArnoldiConfig, + ArnoldiDecomposition, + arnoldiBasisColumns, + arnoldiHessenbergRows, + arnoldiStepsCompleted, + arnoldi, + LanczosConfig, + mkLanczosConfig, + lanczosIterations, + lanczosTolerance, + withLanczosIterations, + withLanczosTolerance, + defaultLanczosConfig, + LanczosDecomposition, + lanczosBasisColumns, + lanczosProjectedTridiagonal, + lanczosAlphaDiagonal, + lanczosBetaOffDiagonal, + lanczosResidualNorm, + lanczosStepsCompleted, + lanczosSymmetric, + LanczosRestartProjection, + lanczosRestartProjectionBasisColumns, + lanczosRestartProjectionProjectedPairs, + lanczosRestartedProjection, + BlockLanczosConfig, + mkBlockLanczosConfig, + blockLanczosIterations, + blockLanczosTolerance, + blockLanczosBlockSize, + blockLanczosReorthogonalize, + withBlockLanczosIterations, + withBlockLanczosTolerance, + withBlockLanczosBlockSize, + withBlockLanczosReorthogonalize, + defaultBlockLanczosConfig, + BlockLanczosDecomposition, + blockLanczosBasisColumns, + blockLanczosProjectedBlockTridiagonal, + blockLanczosBasisCount, + blockLanczosBlockSteps, + blockLanczosSymmetric, + SpectrumEnd (..), + ) +where + +import Moonlight.LinAlg.Pure.Krylov.Arnoldi +import Moonlight.LinAlg.Pure.Krylov.Block +import Moonlight.LinAlg.Pure.Krylov.Config +import Moonlight.LinAlg.Pure.Krylov.Decomposition +import Moonlight.LinAlg.Pure.Krylov.Lanczos +import Moonlight.LinAlg.Pure.Krylov.Selection diff --git a/src-public/Moonlight/LinAlg/Operator.hs b/src-public/Moonlight/LinAlg/Operator.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Operator.hs @@ -0,0 +1,43 @@ +{-# LANGUAGE DataKinds #-} + +-- | Abstract affine-normalized linear operators with explicit self-adjoint construction boundaries. +module Moonlight.LinAlg.Operator + ( OperatorSymmetry (..), + LinearOperator, + operatorShape, + operatorDimension, + mkVectorLinearOperator, + declaredSelfAdjointVectorLinearOperator, + runOperatorU, + csrLinearOperator, + selfAdjointCSRLinearOperator, + graphLaplacianLinearOperator, + diagonalLinearOperator, + pathLaplacianLinearOperator, + symmetricTridiagonalLinearOperator, + packedSparseLinearOperator, + scaleLinearOperator, + addScaledIdentity, + sigmaIdentityMinus, + ) +where + +import Moonlight.LinAlg.Pure.Operator + ( LinearOperator, + OperatorSymmetry (..), + addScaledIdentity, + csrLinearOperator, + declaredSelfAdjointVectorLinearOperator, + diagonalLinearOperator, + graphLaplacianLinearOperator, + mkVectorLinearOperator, + operatorDimension, + operatorShape, + packedSparseLinearOperator, + pathLaplacianLinearOperator, + runOperatorU, + scaleLinearOperator, + selfAdjointCSRLinearOperator, + sigmaIdentityMinus, + symmetricTridiagonalLinearOperator, + ) diff --git a/src-public/Moonlight/LinAlg/Sparse.hs b/src-public/Moonlight/LinAlg/Sparse.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Sparse.hs @@ -0,0 +1,155 @@ +-- | Sparse matrix carriers, packed sparse operators, sealed preconditioner families, and sparse iterative solvers. +module Moonlight.LinAlg.Sparse + ( SparseCOO, + mkSparseCOO, + cooRows, + cooCols, + cooEntries, + SparseCSR, + mkSparseCSR, + csrRows, + csrCols, + csrRowOffsetsVector, + csrColumnIndicesVector, + csrValuesVector, + SparseCSC, + mkSparseCSC, + cscRows, + cscCols, + cscColumnOffsetsVector, + cscRowIndicesVector, + cscValuesVector, + denseToCOO, + denseToCSR, + denseToCSC, + cooToCSR, + cooToCSC, + csrToCOO, + cscToCOO, + cooToDense, + csrToDense, + cscToDense, + csrToCSC, + cscToCSR, + canonicalCSRFromEntries, + GraphEdge (..), + diagonalCSR, + tridiagonalCSR, + pathLaplacianCSR, + graphLaplacianCSR, + csrMatVecVector, + validateCOO, + validateCSR, + validateCSC, + PackedSparseEntry, + packedSparseEntry, + packedSparseEntrySourceOffset, + packedSparseEntryTargetOffset, + packedSparseEntryCoefficient, + PackedSparseOperator, + packedSparseOperatorSourceCardinality, + packedSparseOperatorTargetCardinality, + packedSparseOperatorEntryCount, + packedSparseOperatorEntries, + PackedSparseOperatorShapeError (..), + PackedSparseApplyError (..), + mkPackedSparseOperator, + applyPackedSparseOperatorDense, + SparseIterativeFailure (..), + SparseIterativeResult (..), + SparseStationaryIterationConfig (..), + SparseConjugateGradientConfig (..), + SparseGMRESConfig (..), + IC0Config (..), + SparsePreconditionerFamily (..), + defaultSparsePreconditionerFamily, + solveSparseJacobi, + solveSparseRichardson, + solveSparseCG, + solveSparseGMRES, + ) +where + +import Moonlight.LinAlg.Pure.Sparse.Packed + ( PackedSparseApplyError (..), + PackedSparseEntry, + PackedSparseOperator, + PackedSparseOperatorShapeError (..), + applyPackedSparseOperatorDense, + mkPackedSparseOperator, + packedSparseEntry, + packedSparseEntryCoefficient, + packedSparseEntrySourceOffset, + packedSparseEntryTargetOffset, + packedSparseOperatorEntries, + packedSparseOperatorEntryCount, + packedSparseOperatorSourceCardinality, + packedSparseOperatorTargetCardinality, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.CG + ( solveSparseCG, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.GMRES + ( solveSparseGMRES, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Stationary + ( solveSparseJacobi, + solveSparseRichardson, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Types + ( IC0Config (..), + SparseConjugateGradientConfig (..), + SparseGMRESConfig (..), + SparseIterativeFailure (..), + SparseIterativeResult (..), + SparsePreconditionerFamily (..), + SparseStationaryIterationConfig (..), + defaultSparsePreconditionerFamily, + ) +import Moonlight.LinAlg.Pure.Sparse.Assembly + ( canonicalCSRFromEntries, + ) +import Moonlight.LinAlg.Pure.Sparse.Structured + ( GraphEdge (..), + diagonalCSR, + graphLaplacianCSR, + pathLaplacianCSR, + tridiagonalCSR, + ) +import Moonlight.LinAlg.Pure.Sparse.Types + ( SparseCOO, + SparseCSC, + SparseCSR, + cooCols, + cooEntries, + cooRows, + cooToCSC, + cooToCSR, + cooToDense, + cscCols, + cscColumnOffsetsVector, + cscRowIndicesVector, + cscRows, + cscToCOO, + cscToCSR, + cscToDense, + cscValuesVector, + csrCols, + csrColumnIndicesVector, + csrMatVecVector, + csrRows, + csrRowOffsetsVector, + csrToCOO, + csrToCSC, + csrToDense, + csrValuesVector, + denseToCOO, + denseToCSC, + denseToCSR, + mkSparseCOO, + mkSparseCSC, + mkSparseCSR, + validateCOO, + validateCSC, + validateCSR, + ) diff --git a/src-public/Moonlight/LinAlg/Spectral.hs b/src-public/Moonlight/LinAlg/Spectral.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Spectral.hs @@ -0,0 +1,49 @@ +-- | Eigenvalue and eigenpair requests and contiguous result views, dispatched by demand and operator structure above Krylov. +module Moonlight.LinAlg.Spectral + ( EigenRequest (..), + Eigenpairs, + CertifiedSelectedEigenpairResult (..), + SelectedEigenpairCertificationFailure (..), + SelectedEigenpairOrthonormalityEvidence (..), + SelectedEigenpairRequestOrderingEvidence (..), + SelectedEigenpairResidualEvidence (..), + certifySelectedEigenpairResult, + eigenpairDimension, + eigenpairValues, + eigenpairVectorsColumnMajor, + eigenpairResidualNorms, + eigenpairCount, + eigenpairVectorAt, + EigenSolveConfig, + defaultEigenSolveConfig, + withEigenFallbackLanczosConfig, + withEigenFallbackInitialVector, + solveEigenRequest, + ) +where + +import Moonlight.LinAlg.Pure.Spectral.Request + ( EigenRequest (..), + ) +import Moonlight.LinAlg.Pure.Spectral.Result + ( Eigenpairs, + CertifiedSelectedEigenpairResult (..), + SelectedEigenpairCertificationFailure (..), + SelectedEigenpairOrthonormalityEvidence (..), + SelectedEigenpairRequestOrderingEvidence (..), + SelectedEigenpairResidualEvidence (..), + certifySelectedEigenpairResult, + eigenpairCount, + eigenpairDimension, + eigenpairResidualNorms, + eigenpairValues, + eigenpairVectorAt, + eigenpairVectorsColumnMajor, + ) +import Moonlight.LinAlg.Pure.Spectral.Solve + ( EigenSolveConfig, + defaultEigenSolveConfig, + withEigenFallbackLanczosConfig, + withEigenFallbackInitialVector, + solveEigenRequest, + ) diff --git a/src-public/Moonlight/LinAlg/Statics.hs b/src-public/Moonlight/LinAlg/Statics.hs new file mode 100644 --- /dev/null +++ b/src-public/Moonlight/LinAlg/Statics.hs @@ -0,0 +1,135 @@ +-- | Statics types, assembly, equilibrium compilation, and support checking. +module Moonlight.LinAlg.Statics + ( NodeRef, + nodeRef, + nodeRefLabel, + Axis (..), + Vec3 (..), + MemberRef, + mkMemberRef, + memberEndpoints, + memberTouchesNode, + SupportAxes, + mkSupportAxes, + supportAxesList, + freeSupportAxes, + fixedSupportAxes, + ForceNode, + ForceNetwork, + nodePosition, + nodeLoad, + nodeSupportAxes, + nodeReactionAxes, + networkNodeMap, + networkMemberSet, + NetworkDeclaration, + NetworkBuildError (..), + joint, + support, + supportOn, + load, + member, + network, + UnknownForce (..), + EquationRef (..), + CompiledEquilibrium, + compiledNodeOrder, + compiledFoundationOrder, + compiledMemberOrder, + compiledMemberDirections, + compiledUnknownOrder, + compiledEquationOrder, + compiledCoefficientMatrix, + compiledRightHandSide, + EquilibriumSolution (..), + ForceSign (..), + EquilibriumViolation (..), + EquilibriumResult (..), + allAxes, + addVec3, + subVec3, + negateVec3, + scaleVec3, + dotVec3, + magnitudeVec3, + normalizeVec3, + normalizeVec3Safe, + axisComponent, + axisVector, + vec3Zero, + assembleEquilibriumEquations, + checkEquilibrium, + solveGraphicStatics, + ) +where + +import Moonlight.LinAlg.Pure.Statics.Algebra + ( addVec3, + allAxes, + axisComponent, + axisVector, + dotVec3, + magnitudeVec3, + memberEndpoints, + memberTouchesNode, + mkMemberRef, + negateVec3, + normalizeVec3, + normalizeVec3Safe, + scaleVec3, + subVec3, + vec3Zero, + ) +import Moonlight.LinAlg.Pure.Statics.Compile + ( assembleEquilibriumEquations, + ) +import Moonlight.LinAlg.Pure.Statics.Core + ( checkEquilibrium, + solveGraphicStatics, + ) +import Moonlight.LinAlg.Pure.Statics.Network + ( NetworkBuildError (..), + NetworkDeclaration, + joint, + load, + member, + network, + networkMemberSet, + networkNodeMap, + nodeLoad, + nodePosition, + nodeReactionAxes, + nodeRef, + nodeRefLabel, + nodeSupportAxes, + support, + supportOn, + ) +import Moonlight.LinAlg.Pure.Statics.Types + ( Axis (..), + CompiledEquilibrium, + EquationRef (..), + EquilibriumResult (..), + EquilibriumSolution (..), + EquilibriumViolation (..), + ForceNetwork, + ForceNode, + ForceSign (..), + MemberRef, + NodeRef, + SupportAxes, + UnknownForce (..), + Vec3 (..), + compiledCoefficientMatrix, + compiledEquationOrder, + compiledFoundationOrder, + compiledMemberDirections, + compiledMemberOrder, + compiledNodeOrder, + compiledRightHandSide, + compiledUnknownOrder, + fixedSupportAxes, + freeSupportAxes, + mkSupportAxes, + supportAxesList, + ) diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Assembly.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Assembly.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Assembly.hs @@ -0,0 +1,131 @@ +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Sparse.Assembly + ( canonicalCSRFromEntries, + orderedCSRFromEntries, + ) +where + +import Control.Monad (foldM) +import Data.Kind (Type) +import qualified Data.Vector.Unboxed as U +import Moonlight.Core (AdditiveGroup (..), AdditiveMonoid (..), MoonlightError (..)) +import Moonlight.LinAlg.Pure.Sparse.Types + ( SparseCSR, + canonicalCSRFromValidEntriesUnchecked, + mkSparseCSR, + validateCOOEntries, + ) +import Prelude + +canonicalCSRFromEntries :: + (Eq a, AdditiveGroup a, U.Unbox a) => + Int -> + Int -> + [(Int, Int, a)] -> + Either MoonlightError (SparseCSR a) +canonicalCSRFromEntries rowCount columnCount entries = do + validateCOOEntries rowCount columnCount entries + pure (canonicalCSRFromValidEntriesUnchecked rowCount columnCount entries) + +orderedCSRFromEntries :: + (Eq a, AdditiveMonoid a, U.Unbox a) => + Int -> + Int -> + [(Int, Int, a)] -> + Either MoonlightError (SparseCSR a) +orderedCSRFromEntries rowCount columnCount entries + | rowCount < 0 || columnCount < 0 = + Left + ( InvariantViolation + ( "ordered CSR dimensions must be non-negative, received " + <> show (rowCount, columnCount) + ) + ) + | otherwise = do + builtState <- + foldM + (appendOrderedEntry rowCount columnCount) + initialOrderedCSRState + entries + let completedState = closeRows rowCount builtState + mkSparseCSR + rowCount + columnCount + (reverse (orderedOffsetsRev completedState)) + (reverse (orderedColumnsRev completedState)) + (reverse (orderedValuesRev completedState)) + +type OrderedCSRState :: Type -> Type +data OrderedCSRState a = OrderedCSRState + { orderedCurrentRow :: !Int, + orderedEntryCount :: !Int, + orderedOffsetsRev :: [Int], + orderedColumnsRev :: [Int], + orderedValuesRev :: [a], + orderedPreviousCoordinate :: Maybe (Int, Int) + } + +initialOrderedCSRState :: OrderedCSRState a +initialOrderedCSRState = + OrderedCSRState + { orderedCurrentRow = 0, + orderedEntryCount = 0, + orderedOffsetsRev = [0], + orderedColumnsRev = [], + orderedValuesRev = [], + orderedPreviousCoordinate = Nothing + } + +appendOrderedEntry :: + Int -> + Int -> + OrderedCSRState a -> + (Int, Int, a) -> + Either MoonlightError (OrderedCSRState a) +appendOrderedEntry rowCount columnCount stateValue (rowIndex, columnIndex, entryValue) + | rowIndex < 0 || rowIndex >= rowCount || columnIndex < 0 || columnIndex >= columnCount = + Left + ( InvariantViolation + ( "ordered CSR entry index out of bounds: " + <> show (rowIndex, columnIndex) + <> " for shape " + <> show (rowCount, columnCount) + ) + ) + | not (coordinateStrictlyAfter (orderedPreviousCoordinate stateValue) (rowIndex, columnIndex)) = + Left + ( InvariantViolation + ( "ordered CSR entries must be in strictly increasing row-major order; encountered " + <> show (rowIndex, columnIndex) + <> " after " + <> show (orderedPreviousCoordinate stateValue) + ) + ) + | otherwise = + let rowClosedState = closeRows rowIndex stateValue + in Right + rowClosedState + { orderedEntryCount = orderedEntryCount rowClosedState + 1, + orderedColumnsRev = columnIndex : orderedColumnsRev rowClosedState, + orderedValuesRev = entryValue : orderedValuesRev rowClosedState, + orderedPreviousCoordinate = Just (rowIndex, columnIndex) + } + +coordinateStrictlyAfter :: Maybe (Int, Int) -> (Int, Int) -> Bool +coordinateStrictlyAfter previousCoordinate currentCoordinate = + case previousCoordinate of + Nothing -> True + Just previousValue -> previousValue < currentCoordinate + +closeRows :: Int -> OrderedCSRState a -> OrderedCSRState a +closeRows targetRow stateValue + | orderedCurrentRow stateValue >= targetRow = stateValue + | otherwise = + let closedRowCount = targetRow - orderedCurrentRow stateValue + in stateValue + { orderedCurrentRow = targetRow, + orderedOffsetsRev = + replicate closedRowCount (orderedEntryCount stateValue) + <> orderedOffsetsRev stateValue + } diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Packed.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Packed.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Packed.hs @@ -0,0 +1,186 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE StandaloneKindSignatures #-} + +module Moonlight.LinAlg.Pure.Sparse.Packed + ( PackedSparseEntry, + packedSparseEntry, + packedSparseEntrySourceOffset, + packedSparseEntryTargetOffset, + packedSparseEntryCoefficient, + PackedSparseOperator, + packedSparseOperatorSourceCardinality, + packedSparseOperatorTargetCardinality, + packedSparseOperatorEntryCount, + packedSparseOperatorEntries, + PackedSparseOperatorShapeError (..), + PackedSparseApplyError (..), + mkPackedSparseOperator, + applyPackedSparseOperatorDense, + ) +where + +import Data.Bifunctor (first) +import Data.Kind (Type) +import Data.Maybe (listToMaybe) +import Data.Vector.Unboxed qualified as Unboxed +import Moonlight.Core (checkedNaturalToInt) +import Numeric.Natural (Natural) + +-- | Canonical packed COO-style entry for a sparse operator. +-- +-- The entry is source/target oriented because boundary and coboundary +-- operators name their domain and codomain that way. The operator constructor +-- validates every entry against the declared shape before sealing the unboxed +-- vectors. +type PackedSparseEntry :: Type -> Type +data PackedSparseEntry coefficient = PackedSparseEntry + { pseSourceOffset :: !Int, + pseTargetOffset :: !Int, + pseCoefficient :: !coefficient + } + deriving stock (Eq, Ord, Show) + +packedSparseEntry :: Int -> Int -> coefficient -> PackedSparseEntry coefficient +packedSparseEntry sourceOffsetValue targetOffsetValue coefficientValue = + PackedSparseEntry + { pseSourceOffset = sourceOffsetValue, + pseTargetOffset = targetOffsetValue, + pseCoefficient = coefficientValue + } + +packedSparseEntrySourceOffset :: PackedSparseEntry coefficient -> Int +packedSparseEntrySourceOffset = + pseSourceOffset + +packedSparseEntryTargetOffset :: PackedSparseEntry coefficient -> Int +packedSparseEntryTargetOffset = + pseTargetOffset + +packedSparseEntryCoefficient :: PackedSparseEntry coefficient -> coefficient +packedSparseEntryCoefficient = + pseCoefficient + +-- | Shape-validated packed sparse operator. +-- +-- The constructor is hidden: once built, offset vectors are known in-bounds and +-- zero coefficients have been pruned, so dense apply can stay a pure unboxed +-- vector kernel. +type PackedSparseOperator :: Type -> Type +data PackedSparseOperator coefficient = PackedSparseOperator + { psoSourceCardinality :: !Int, + psoTargetCardinality :: !Int, + psoSourceOffsets :: !(Unboxed.Vector Int), + psoTargetOffsets :: !(Unboxed.Vector Int), + psoCoefficients :: !(Unboxed.Vector coefficient) + } + deriving stock (Eq, Show) + +packedSparseOperatorSourceCardinality :: PackedSparseOperator coefficient -> Int +packedSparseOperatorSourceCardinality = + psoSourceCardinality + +packedSparseOperatorTargetCardinality :: PackedSparseOperator coefficient -> Int +packedSparseOperatorTargetCardinality = + psoTargetCardinality + +packedSparseOperatorEntryCount :: + Unboxed.Unbox coefficient => + PackedSparseOperator coefficient -> + Int +packedSparseOperatorEntryCount = + Unboxed.length . psoCoefficients + +packedSparseOperatorEntries :: + Unboxed.Unbox coefficient => + PackedSparseOperator coefficient -> + [PackedSparseEntry coefficient] +packedSparseOperatorEntries packedOperator = + zipWith3 + packedSparseEntry + (Unboxed.toList (psoSourceOffsets packedOperator)) + (Unboxed.toList (psoTargetOffsets packedOperator)) + (Unboxed.toList (psoCoefficients packedOperator)) + +type PackedSparseOperatorShapeError :: Type +data PackedSparseOperatorShapeError + = PackedSparseCardinalityOutOfBounds !Natural + | PackedSparseEntryOutOfBounds !Int !Int !Int !Int + deriving stock (Eq, Show) + +type PackedSparseApplyError :: Type +data PackedSparseApplyError + = PackedSparseInputLengthMismatch !Int !Int + deriving stock (Eq, Show) + +mkPackedSparseOperator :: + (Eq coefficient, Num coefficient, Unboxed.Unbox coefficient) => + Natural -> + Natural -> + [PackedSparseEntry coefficient] -> + Either PackedSparseOperatorShapeError (PackedSparseOperator coefficient) +mkPackedSparseOperator sourceCardinalityValue targetCardinalityValue entries = do + sourceDimension <- packedSparseCardinalityToInt sourceCardinalityValue + targetDimension <- packedSparseCardinalityToInt targetCardinalityValue + case firstOutOfBoundsEntry sourceDimension targetDimension entries of + Just entryValue -> Left (entryOutOfBoundsError sourceDimension targetDimension entryValue) + Nothing -> + let nonzeroEntries = filter ((/= 0) . pseCoefficient) entries + in Right + PackedSparseOperator + { psoSourceCardinality = sourceDimension, + psoTargetCardinality = targetDimension, + psoSourceOffsets = Unboxed.fromList (fmap pseSourceOffset nonzeroEntries), + psoTargetOffsets = Unboxed.fromList (fmap pseTargetOffset nonzeroEntries), + psoCoefficients = Unboxed.fromList (fmap pseCoefficient nonzeroEntries) + } + +packedSparseCardinalityToInt :: Natural -> Either PackedSparseOperatorShapeError Int +packedSparseCardinalityToInt cardinalityValue = + first + (const (PackedSparseCardinalityOutOfBounds cardinalityValue)) + (checkedNaturalToInt cardinalityValue) + +applyPackedSparseOperatorDense :: + (Num coefficient, Unboxed.Unbox coefficient) => + PackedSparseOperator coefficient -> + Unboxed.Vector coefficient -> + Either PackedSparseApplyError (Unboxed.Vector coefficient) +applyPackedSparseOperatorDense packedOperator sourceVector = + if Unboxed.length sourceVector == packedSparseOperatorSourceCardinality packedOperator + then + Right + ( Unboxed.accumulate_ + (+) + (Unboxed.replicate (packedSparseOperatorTargetCardinality packedOperator) 0) + (psoTargetOffsets packedOperator) + (Unboxed.zipWith (*) (psoCoefficients packedOperator) (Unboxed.backpermute sourceVector (psoSourceOffsets packedOperator))) + ) + else + Left + ( PackedSparseInputLengthMismatch + (packedSparseOperatorSourceCardinality packedOperator) + (Unboxed.length sourceVector) + ) + +firstOutOfBoundsEntry :: Int -> Int -> [PackedSparseEntry coefficient] -> Maybe (PackedSparseEntry coefficient) +firstOutOfBoundsEntry sourceDimension targetDimension = + listToMaybe . filter (not . entryWithinBounds sourceDimension targetDimension) + +entryWithinBounds :: Int -> Int -> PackedSparseEntry coefficient -> Bool +entryWithinBounds sourceDimension targetDimension entry = + pseSourceOffset entry >= 0 + && pseSourceOffset entry < sourceDimension + && pseTargetOffset entry >= 0 + && pseTargetOffset entry < targetDimension + +entryOutOfBoundsError :: + Int -> + Int -> + PackedSparseEntry coefficient -> + PackedSparseOperatorShapeError +entryOutOfBoundsError sourceDimension targetDimension entry = + PackedSparseEntryOutOfBounds + (pseSourceOffset entry) + (pseTargetOffset entry) + sourceDimension + targetDimension diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/CG.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/CG.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/CG.hs @@ -0,0 +1,721 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RecordWildCards #-} + +module Moonlight.LinAlg.Pure.Sparse.Solver.CG + ( solveSparseCG, + ) +where + +import Control.Monad.ST (ST, runST) +import Data.Vector.Unboxed qualified as U +import Data.Vector.Unboxed.Mutable qualified as MU +import Moonlight.Core (fieldValueValid) +import Moonlight.LinAlg.Pure.Sparse.Solver.Common + ( validateSparseSolverConfiguration, + validateSparseSystemInput, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Mutable + ( MutableDoubleVector, + copyImmutableSquaredNormIntoMutable, + copyMutableVector, + csrMatVecDotIntoMutable, + csrResidualSquaredIntoMutable, + divideByDiagonalDotAndCopyMutable, + freezeMutableDoubleVector, + initializeZeroJacobiMutable, + updateDirectionMutable, + updateSolutionAndResidualSquaredMutable, + updateSolutionResidualJacobiMutable, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Preconditioner + ( SparsePreconditioner (..), + applySparsePreconditionerAndDotMutable, + compileSparsePreconditioner, + preconditionerDimension, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Types + ( SparseConjugateGradientConfig (..), + SparseIterativeFailure (..), + SparseIterativeResult (..), + ) +import Moonlight.LinAlg.Pure.Sparse.Types + ( SparseCSR, + csrRows, + ) +import Prelude + +solveSparseCG :: + SparseConjugateGradientConfig -> + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + Either SparseIterativeFailure SparseIterativeResult +solveSparseCG + SparseConjugateGradientConfig {..} + sparseMatrix + rhsValues + initialGuess = do + validateSparseSystemInput + sparseMatrix + rhsValues + initialGuess + validateSparseSolverConfiguration + "CG" + scgcTolerance + scgcIterationLimit + preconditioner <- + compileSparsePreconditioner + scgcPreconditionerFamily + sparseMatrix + validateCgPreconditioner sparseMatrix preconditioner + let !zeroInitialGuess = U.all (== 0.0) initialGuess + case preconditioner of + IdentitySparsePreconditioner _ -> + runST + ( solveIdentityCgMutable + scgcTolerance + scgcIterationLimit + zeroInitialGuess + sparseMatrix + rhsValues + initialGuess + ) + DiagonalSparsePreconditioner diagonalValues + | uniformDiagonal diagonalValues -> + -- For M = dI with d > 0, all factors of d cancel exactly from + -- alpha, beta, and the represented search direction. + runST + ( solveIdentityCgMutable + scgcTolerance + scgcIterationLimit + zeroInitialGuess + sparseMatrix + rhsValues + initialGuess + ) + | otherwise -> + runST + ( solveJacobiCgMutable + scgcTolerance + scgcIterationLimit + zeroInitialGuess + diagonalValues + sparseMatrix + rhsValues + initialGuess + ) + SsorSparsePreconditioner {} -> + runST + ( solveGenericPcgMutable + scgcTolerance + scgcIterationLimit + zeroInitialGuess + preconditioner + sparseMatrix + rhsValues + initialGuess + ) + IncompleteCholesky0SparsePreconditioner {} -> + runST + ( solveGenericPcgMutable + scgcTolerance + scgcIterationLimit + zeroInitialGuess + preconditioner + sparseMatrix + rhsValues + initialGuess + ) + +solveIdentityCgMutable :: + Double -> + Int -> + Bool -> + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + ST s (Either SparseIterativeFailure SparseIterativeResult) +solveIdentityCgMutable + !toleranceValue + !iterationLimit + zeroInitialGuess + sparseMatrix + rhsValues + initialGuess = do + let !dimension = csrRows sparseMatrix + guessVector <- U.thaw initialGuess + residualVector <- MU.unsafeNew dimension + directionVector <- MU.unsafeNew dimension + imageDirectionVector <- MU.unsafeNew dimension + + initialResidualSquared <- + if zeroInitialGuess + then + copyImmutableSquaredNormIntoMutable + rhsValues + residualVector + else + csrResidualSquaredIntoMutable + sparseMatrix + rhsValues + guessVector + residualVector + + case residualNormFromSquared initialResidualSquared of + Left failureValue -> pure (Left failureValue) + Right initialResidualNorm + | initialResidualNorm <= toleranceValue -> + Right + <$> freezeCgResult + 0 + initialResidualNorm + guessVector + | otherwise -> do + copyMutableVector residualVector directionVector + let iterateCg !iterationCount !residualSquared + | iterationCount >= iterationLimit = + pure + ( Left + ( SparseIterationBudgetExceeded + iterationLimit + ) + ) + | otherwise = do + denominator <- + csrMatVecDotIntoMutable + sparseMatrix + directionVector + imageDirectionVector + if not (positiveFinite denominator) + then + pure + ( Left + ( SparseInvalidInput + "CG requires p^T A p > 0; matrix is not SPD or arithmetic broke down" + ) + ) + else do + let !alphaValue = + residualSquared / denominator + if not (fieldValueValid alphaValue) + then + pure + ( Left + ( SparseInvalidInput + "CG produced a non-finite alpha" + ) + ) + else do + nextResidualSquared <- + updateSolutionAndResidualSquaredMutable + alphaValue + directionVector + imageDirectionVector + guessVector + residualVector + case residualNormFromSquared nextResidualSquared of + Left failureValue -> + pure (Left failureValue) + Right nextResidualNorm -> + let !nextIteration = iterationCount + 1 + in if nextResidualNorm <= toleranceValue + then certifyOrRestart nextIteration + else do + let !betaValue = + nextResidualSquared + / residualSquared + if not (fieldValueValid betaValue) + then + pure + ( Left + ( SparseInvalidInput + "CG produced a non-finite beta" + ) + ) + else do + updateDirectionMutable + betaValue + residualVector + directionVector + iterateCg + nextIteration + nextResidualSquared + + certifyOrRestart !iterationCount = do + trueResidualSquared <- + csrResidualSquaredIntoMutable + sparseMatrix + rhsValues + guessVector + residualVector + case residualNormFromSquared trueResidualSquared of + Left failureValue -> pure (Left failureValue) + Right trueResidualNorm + | trueResidualNorm <= toleranceValue -> + Right + <$> freezeCgResult + iterationCount + trueResidualNorm + guessVector + | otherwise -> do + copyMutableVector + residualVector + directionVector + iterateCg + iterationCount + trueResidualSquared + + iterateCg 0 initialResidualSquared + +solveJacobiCgMutable :: + Double -> + Int -> + Bool -> + U.Vector Double -> + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + ST s (Either SparseIterativeFailure SparseIterativeResult) +solveJacobiCgMutable + !toleranceValue + !iterationLimit + zeroInitialGuess + diagonalValues + sparseMatrix + rhsValues + initialGuess = do + let !dimension = csrRows sparseMatrix + guessVector <- U.thaw initialGuess + residualVector <- MU.unsafeNew dimension + preconditionedResidualVector <- MU.unsafeNew dimension + directionVector <- MU.unsafeNew dimension + imageDirectionVector <- MU.unsafeNew dimension + + (initialResidualSquared, initialRho) <- + if zeroInitialGuess + then + initializeZeroJacobiMutable + rhsValues + diagonalValues + residualVector + preconditionedResidualVector + directionVector + else do + residualSquared <- + csrResidualSquaredIntoMutable + sparseMatrix + rhsValues + guessVector + residualVector + rhoValue <- + divideByDiagonalDotAndCopyMutable + diagonalValues + residualVector + preconditionedResidualVector + directionVector + pure (residualSquared, rhoValue) + + case residualNormFromSquared initialResidualSquared of + Left failureValue -> pure (Left failureValue) + Right initialResidualNorm + | initialResidualNorm <= toleranceValue -> + Right + <$> freezeCgResult + 0 + initialResidualNorm + guessVector + | not (positiveFinite initialRho) -> + pure + ( Left + ( SparseInvalidInput + "Jacobi PCG requires r^T M^-1 r > 0" + ) + ) + | otherwise -> do + let iteratePcg !iterationCount !rhoValue + | iterationCount >= iterationLimit = + pure + ( Left + ( SparseIterationBudgetExceeded + iterationLimit + ) + ) + | otherwise = do + denominator <- + csrMatVecDotIntoMutable + sparseMatrix + directionVector + imageDirectionVector + if not (positiveFinite denominator) + then + pure + ( Left + ( SparseInvalidInput + "Jacobi PCG requires p^T A p > 0; matrix is not SPD or arithmetic broke down" + ) + ) + else do + let !alphaValue = rhoValue / denominator + if not (fieldValueValid alphaValue) + then + pure + ( Left + ( SparseInvalidInput + "Jacobi PCG produced a non-finite alpha" + ) + ) + else do + (nextResidualSquared, nextRho) <- + updateSolutionResidualJacobiMutable + alphaValue + diagonalValues + directionVector + imageDirectionVector + guessVector + residualVector + preconditionedResidualVector + case residualNormFromSquared nextResidualSquared of + Left failureValue -> + pure (Left failureValue) + Right nextResidualNorm -> + let !nextIteration = iterationCount + 1 + in if nextResidualNorm <= toleranceValue + then certifyOrRestart nextIteration + else + if not (positiveFinite nextRho) + then + pure + ( Left + ( SparseInvalidInput + "Jacobi PCG encountered non-positive r^T M^-1 r" + ) + ) + else do + let !betaValue = + nextRho / rhoValue + if not (fieldValueValid betaValue) + then + pure + ( Left + ( SparseInvalidInput + "Jacobi PCG produced a non-finite beta" + ) + ) + else do + updateDirectionMutable + betaValue + preconditionedResidualVector + directionVector + iteratePcg + nextIteration + nextRho + + certifyOrRestart !iterationCount = do + trueResidualSquared <- + csrResidualSquaredIntoMutable + sparseMatrix + rhsValues + guessVector + residualVector + case residualNormFromSquared trueResidualSquared of + Left failureValue -> pure (Left failureValue) + Right trueResidualNorm + | trueResidualNorm <= toleranceValue -> + Right + <$> freezeCgResult + iterationCount + trueResidualNorm + guessVector + | otherwise -> do + restartedRho <- + divideByDiagonalDotAndCopyMutable + diagonalValues + residualVector + preconditionedResidualVector + directionVector + if not (positiveFinite restartedRho) + then + pure + ( Left + ( SparseInvalidInput + "Jacobi PCG residual replacement produced non-positive r^T M^-1 r" + ) + ) + else + iteratePcg + iterationCount + restartedRho + + iteratePcg 0 initialRho + +solveGenericPcgMutable :: + Double -> + Int -> + Bool -> + SparsePreconditioner -> + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + ST s (Either SparseIterativeFailure SparseIterativeResult) +solveGenericPcgMutable + !toleranceValue + !iterationLimit + zeroInitialGuess + preconditioner + sparseMatrix + rhsValues + initialGuess = do + let !dimension = csrRows sparseMatrix + guessVector <- U.thaw initialGuess + residualVector <- MU.unsafeNew dimension + preconditionedResidualVector <- MU.unsafeNew dimension + directionVector <- MU.unsafeNew dimension + imageDirectionVector <- MU.unsafeNew dimension + preconditionerScratchA <- MU.unsafeNew dimension + preconditionerScratchB <- MU.unsafeNew dimension + + initialResidualSquared <- + if zeroInitialGuess + then + copyImmutableSquaredNormIntoMutable + rhsValues + residualVector + else + csrResidualSquaredIntoMutable + sparseMatrix + rhsValues + guessVector + residualVector + + case residualNormFromSquared initialResidualSquared of + Left failureValue -> pure (Left failureValue) + Right initialResidualNorm + | initialResidualNorm <= toleranceValue -> + Right + <$> freezeCgResult + 0 + initialResidualNorm + guessVector + | otherwise -> do + initialRho <- + applySparsePreconditionerAndDotMutable + preconditioner + residualVector + preconditionerScratchA + preconditionerScratchB + preconditionedResidualVector + if not (positiveFinite initialRho) + then + pure + ( Left + ( SparseInvalidInput + "PCG requires r^T M^-1 r > 0" + ) + ) + else do + copyMutableVector + preconditionedResidualVector + directionVector + let iteratePcg !iterationCount !rhoValue + | iterationCount >= iterationLimit = + pure + ( Left + ( SparseIterationBudgetExceeded + iterationLimit + ) + ) + | otherwise = do + denominator <- + csrMatVecDotIntoMutable + sparseMatrix + directionVector + imageDirectionVector + if not (positiveFinite denominator) + then + pure + ( Left + ( SparseInvalidInput + "PCG requires p^T A p > 0; matrix is not SPD or arithmetic broke down" + ) + ) + else do + let !alphaValue = rhoValue / denominator + if not (fieldValueValid alphaValue) + then + pure + ( Left + ( SparseInvalidInput + "PCG produced a non-finite alpha" + ) + ) + else do + nextResidualSquared <- + updateSolutionAndResidualSquaredMutable + alphaValue + directionVector + imageDirectionVector + guessVector + residualVector + case residualNormFromSquared nextResidualSquared of + Left failureValue -> + pure (Left failureValue) + Right nextResidualNorm -> + let !nextIteration = iterationCount + 1 + in if nextResidualNorm <= toleranceValue + then certifyOrRestart nextIteration + else do + nextRho <- + applySparsePreconditionerAndDotMutable + preconditioner + residualVector + preconditionerScratchA + preconditionerScratchB + preconditionedResidualVector + if not (positiveFinite nextRho) + then + pure + ( Left + ( SparseInvalidInput + "PCG encountered non-positive r^T M^-1 r" + ) + ) + else do + let !betaValue = + nextRho / rhoValue + if not (fieldValueValid betaValue) + then + pure + ( Left + ( SparseInvalidInput + "PCG produced a non-finite beta" + ) + ) + else do + updateDirectionMutable + betaValue + preconditionedResidualVector + directionVector + iteratePcg + nextIteration + nextRho + + certifyOrRestart !iterationCount = do + trueResidualSquared <- + csrResidualSquaredIntoMutable + sparseMatrix + rhsValues + guessVector + residualVector + case residualNormFromSquared trueResidualSquared of + Left failureValue -> pure (Left failureValue) + Right trueResidualNorm + | trueResidualNorm <= toleranceValue -> + Right + <$> freezeCgResult + iterationCount + trueResidualNorm + guessVector + | otherwise -> do + restartedRho <- + applySparsePreconditionerAndDotMutable + preconditioner + residualVector + preconditionerScratchA + preconditionerScratchB + preconditionedResidualVector + if not (positiveFinite restartedRho) + then + pure + ( Left + ( SparseInvalidInput + "PCG residual replacement produced non-positive r^T M^-1 r" + ) + ) + else do + copyMutableVector + preconditionedResidualVector + directionVector + iteratePcg + iterationCount + restartedRho + + iteratePcg 0 initialRho + +freezeCgResult :: + Int -> + Double -> + MutableDoubleVector s -> + ST s SparseIterativeResult +freezeCgResult iterationCount residualNormValue guessVector = do + solutionVector <- freezeMutableDoubleVector guessVector + pure + SparseIterativeResult + { sparseSolution = solutionVector, + sparseIterations = iterationCount, + sparseResidualNorm = residualNormValue + } + +validateCgPreconditioner :: + SparseCSR Double -> + SparsePreconditioner -> + Either SparseIterativeFailure () +validateCgPreconditioner sparseMatrix preconditioner + | preconditionerDimension preconditioner /= csrRows sparseMatrix = + Left + ( SparseInvalidInput + "CG preconditioner dimension must equal matrix dimension" + ) + | otherwise = + case preconditioner of + IdentitySparsePreconditioner _ -> Right () + DiagonalSparsePreconditioner diagonalValues + | U.all positiveFinite diagonalValues -> Right () + | otherwise -> + Left + ( SparseInvalidInput + "CG requires a finite positive-definite diagonal preconditioner" + ) + SsorSparsePreconditioner + omegaValue + diagonalValues + scaledDiagonalValues + _ + | omegaValue > 0.0 + && omegaValue < 2.0 + && U.all positiveFinite diagonalValues + && U.all positiveFinite scaledDiagonalValues -> + Right () + | otherwise -> + Left + ( SparseInvalidInput + "CG requires an SPD SSOR preconditioner with omega in (0,2) and positive diagonal" + ) + IncompleteCholesky0SparsePreconditioner _ -> Right () + +residualNormFromSquared :: + Double -> + Either SparseIterativeFailure Double +residualNormFromSquared squaredNorm + | not (fieldValueValid squaredNorm) = + Left + ( SparseInvalidInput + "CG residual squared norm became non-finite" + ) + | squaredNorm < 0.0 = + Left + ( SparseInvalidInput + "CG residual squared norm became negative" + ) + | otherwise = Right (sqrt squaredNorm) + +uniformDiagonal :: U.Vector Double -> Bool +uniformDiagonal diagonalValues + | U.null diagonalValues = False + | otherwise = + let !firstValue = diagonalValues `U.unsafeIndex` 0 + in U.all (== firstValue) diagonalValues +{-# INLINE uniformDiagonal #-} + +positiveFinite :: Double -> Bool +positiveFinite value = value > 0.0 && fieldValueValid value +{-# INLINE positiveFinite #-} diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Common.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Common.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Common.hs @@ -0,0 +1,103 @@ +{-# LANGUAGE BangPatterns #-} + +module Moonlight.LinAlg.Pure.Sparse.Solver.Common + ( finiteDouble, + validateSparseSystemInput, + validateSparseSolverConfiguration, + sparseDiagonal, + solverEpsilon, + shiftedDiagonalValue, + ) +where + +import Data.Vector.Unboxed qualified as U +import Moonlight.Core (fieldValueValid) +import Moonlight.LinAlg.Pure.Sparse.Solver.Types (SparseIterativeFailure (..)) +import Moonlight.LinAlg.Pure.Sparse.Types + ( SparseCSR, + csrCols, + csrColumnIndicesVector, + csrRows, + csrRowOffsetsVector, + csrValuesVector, + ) +import Prelude + +finiteDouble :: Double -> Bool +finiteDouble = fieldValueValid +{-# INLINE finiteDouble #-} + +validateSparseSystemInput :: + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + Either SparseIterativeFailure () +validateSparseSystemInput sparseMatrix rhsValues initialGuess + | dimension /= csrCols sparseMatrix = + Left (SparseInvalidInput "sparse iterative solver expects a square system matrix") + | U.length rhsValues /= dimension = + Left (SparseInvalidInput "sparse iterative solver expects RHS length equal to matrix dimension") + | U.length initialGuess /= dimension = + Left (SparseInvalidInput "sparse iterative solver expects initial guess length equal to matrix dimension") + | U.any (not . fieldValueValid) (csrValuesVector sparseMatrix) = + Left (SparseInvalidInput "sparse iterative solver matrix entries must be finite") + | U.any (not . fieldValueValid) rhsValues = + Left (SparseInvalidInput "sparse iterative solver right-hand side must be finite") + | U.any (not . fieldValueValid) initialGuess = + Left (SparseInvalidInput "sparse iterative solver initial guess must be finite") + | otherwise = Right () + where + dimension = csrRows sparseMatrix + +validateSparseSolverConfiguration :: + String -> + Double -> + Int -> + Either SparseIterativeFailure () +validateSparseSolverConfiguration methodName toleranceValue iterationLimit + | not (fieldValueValid toleranceValue) = + Left (SparseInvalidInput (methodName <> " tolerance must be finite")) + | toleranceValue < 0.0 = + Left (SparseInvalidInput (methodName <> " tolerance must be non-negative")) + | iterationLimit < 0 = + Left (SparseInvalidInput (methodName <> " iteration limit must be non-negative")) + | otherwise = Right () + +sparseDiagonal :: + SparseCSR Double -> + Either SparseIterativeFailure (U.Vector Double) +sparseDiagonal sparseMatrix = + Right + ( U.generate + (csrRows sparseMatrix) + (diagonalAt sparseMatrix) + ) + +diagonalAt :: SparseCSR Double -> Int -> Double +diagonalAt sparseMatrix rowIndex = + findDiagonal startOffset + where + !rowOffsets = csrRowOffsetsVector sparseMatrix + !columnIndices = csrColumnIndicesVector sparseMatrix + !values = csrValuesVector sparseMatrix + !startOffset = rowOffsets `U.unsafeIndex` rowIndex + !endOffset = rowOffsets `U.unsafeIndex` (rowIndex + 1) + + findDiagonal !entryIndex + | entryIndex >= endOffset = 0.0 + | otherwise = + let !columnIndex = + columnIndices `U.unsafeIndex` entryIndex + in case compare columnIndex rowIndex of + LT -> findDiagonal (entryIndex + 1) + EQ -> values `U.unsafeIndex` entryIndex + GT -> 0.0 +{-# INLINE diagonalAt #-} + +solverEpsilon :: Double +solverEpsilon = 1.0e-12 + +shiftedDiagonalValue :: Double -> Double -> Double +shiftedDiagonalValue shiftValue diagonalValue + | diagonalValue >= 0.0 = diagonalValue + shiftValue + | otherwise = diagonalValue - shiftValue diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/GMRES.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/GMRES.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/GMRES.hs @@ -0,0 +1,554 @@ +{-# LANGUAGE BangPatterns #-} + +module Moonlight.LinAlg.Pure.Sparse.Solver.GMRES + ( solveSparseGMRES, + ) +where + +import Control.Monad (unless, when) +import Control.Monad.ST (ST, runST) +import Control.Monad.Trans.Class (lift) +import Control.Monad.Trans.Except (ExceptT, except, runExceptT, throwE) +import Data.Bifunctor (first) +import Data.Kind (Type) +import Data.Vector.Unboxed qualified as U +import Data.Vector.Unboxed.Mutable qualified as MU +import Moonlight.Core + ( checkedNonNegativeProduct, + checkedNonNegativeSum, + fieldValueValid, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Common + ( solverEpsilon, + validateSparseSolverConfiguration, + validateSparseSystemInput, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Mutable + ( MutableDoubleVector, + addScaledMutableVector, + copyMutableVector, + csrMatVecIntoMutable, + dotMutableVector, + freezeMutableDoubleVector, + newMutableDoubleVector, + residualIntoMutable, + scaleMutableVector, + thawMutableDoubleVector, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Preconditioner + ( SparsePreconditioner, + applySparsePreconditionerMutable, + compileSparsePreconditioner, + preconditionerDimension, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Types + ( SparseGMRESConfig (..), + SparseIterativeFailure (..), + SparseIterativeResult (..), + ) +import Moonlight.LinAlg.Pure.Sparse.Types (SparseCSR, csrRows) +import Prelude + +solveSparseGMRES :: + SparseGMRESConfig -> + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + Either SparseIterativeFailure SparseIterativeResult +solveSparseGMRES config sparseMatrix rhsValues initialGuess = do + validateSparseSystemInput sparseMatrix rhsValues initialGuess + validateSparseSolverConfiguration "GMRES" (sgcTolerance config) (sgcIterationLimit config) + if sgcRestartDimension config <= 0 + then Left (SparseInvalidInput "GMRES restart dimension must be positive") + else Right () + workspaceSizes <- + checkedGMRESWorkspaceSizes + (csrRows sparseMatrix) + (sgcRestartDimension config) + preconditioner <- compileSparsePreconditioner (sgcPreconditionerFamily config) sparseMatrix + validateGmresPreconditioner sparseMatrix preconditioner + runST (solveSparseGMRESMutable config preconditioner sparseMatrix rhsValues initialGuess workspaceSizes) + +solveSparseGMRESMutable :: + SparseGMRESConfig -> + SparsePreconditioner -> + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + GMRESWorkspaceSizes -> + ST s (Either SparseIterativeFailure SparseIterativeResult) +solveSparseGMRESMutable config preconditioner sparseMatrix rhsValues initialGuess workspaceSizes = + runExceptT $ do + let !restartDimension = sgcRestartDimension config + !restartCycles = restartCycleCount (sgcIterationLimit config) restartDimension + workspace <- lift (newGMRESWorkspace workspaceSizes) + currentGuess <- lift (thawMutableDoubleVector initialGuess) + finalState <- + U.foldM' + (gmresRestartCycle config preconditioner sparseMatrix rhsValues currentGuess workspace) + (GmresRunning 0) + (U.enumFromN 0 restartCycles) + gmresResultFromState (sgcIterationLimit config) currentGuess finalState + +type GMRESWorkspace :: Type -> Type +data GMRESWorkspace s = GMRESWorkspace + { gmresBasisPayload :: !(MutableDoubleVector s), + gmresPreconditionedPayload :: !(MutableDoubleVector s), + gmresHessenbergPayload :: !(MutableDoubleVector s), + gmresCosines :: !(MutableDoubleVector s), + gmresSines :: !(MutableDoubleVector s), + gmresProjectedResidual :: !(MutableDoubleVector s), + gmresYValues :: !(MutableDoubleVector s), + gmresResidualVector :: !(MutableDoubleVector s), + gmresImageVector :: !(MutableDoubleVector s), + gmresWorkVector :: !(MutableDoubleVector s), + gmresPreconditionerScratchA :: !(MutableDoubleVector s), + gmresPreconditionerScratchB :: !(MutableDoubleVector s), + gmresDimension :: !Int, + gmresRestartDimension :: !Int, + gmresKrylovColumnCount :: !Int + } + +type GMRESWorkspaceSizes :: Type +data GMRESWorkspaceSizes = GMRESWorkspaceSizes + { gmresWorkspaceDimension :: !Int, + gmresWorkspaceRestartDimension :: !Int, + gmresWorkspaceKrylovColumnCount :: !Int, + gmresWorkspaceBasisPayloadLength :: !Int, + gmresWorkspacePreconditionedPayloadLength :: !Int, + gmresWorkspaceHessenbergPayloadLength :: !Int + } + +checkedGMRESWorkspaceSizes :: + Int -> + Int -> + Either SparseIterativeFailure GMRESWorkspaceSizes +checkedGMRESWorkspaceSizes dimension restartDimension = do + krylovColumnCount <- + checkedWorkspaceCardinality "GMRES restart dimension plus one" + (checkedNonNegativeSum restartDimension 1) + basisPayloadLength <- + checkedWorkspaceCardinality "GMRES basis workspace" + (checkedNonNegativeProduct krylovColumnCount dimension) + preconditionedPayloadLength <- + checkedWorkspaceCardinality "GMRES preconditioned basis workspace" + (checkedNonNegativeProduct restartDimension dimension) + hessenbergPayloadLength <- + checkedWorkspaceCardinality "GMRES Hessenberg workspace" + (checkedNonNegativeProduct krylovColumnCount restartDimension) + Right + GMRESWorkspaceSizes + { gmresWorkspaceDimension = dimension, + gmresWorkspaceRestartDimension = restartDimension, + gmresWorkspaceKrylovColumnCount = krylovColumnCount, + gmresWorkspaceBasisPayloadLength = basisPayloadLength, + gmresWorkspacePreconditionedPayloadLength = preconditionedPayloadLength, + gmresWorkspaceHessenbergPayloadLength = hessenbergPayloadLength + } + +checkedWorkspaceCardinality :: + String -> + Either cardinalityFailure Int -> + Either SparseIterativeFailure Int +checkedWorkspaceCardinality workspaceName = + first + (const (SparseInvalidInput (workspaceName <> " exceeds non-negative Int cardinality"))) + +newGMRESWorkspace :: GMRESWorkspaceSizes -> ST s (GMRESWorkspace s) +newGMRESWorkspace workspaceSizes = do + let !dimension = gmresWorkspaceDimension workspaceSizes + !restartDimension = gmresWorkspaceRestartDimension workspaceSizes + !krylovColumnCount = gmresWorkspaceKrylovColumnCount workspaceSizes + basisPayload <- newMutableDoubleVector (gmresWorkspaceBasisPayloadLength workspaceSizes) + preconditionedPayload <- newMutableDoubleVector (gmresWorkspacePreconditionedPayloadLength workspaceSizes) + hessenbergPayload <- newMutableDoubleVector (gmresWorkspaceHessenbergPayloadLength workspaceSizes) + cosines <- newMutableDoubleVector restartDimension + sines <- newMutableDoubleVector restartDimension + projectedResidual <- newMutableDoubleVector krylovColumnCount + yValues <- newMutableDoubleVector restartDimension + residualVector <- newMutableDoubleVector dimension + imageVector <- newMutableDoubleVector dimension + workVector <- newMutableDoubleVector dimension + preconditionerScratchA <- newMutableDoubleVector dimension + preconditionerScratchB <- newMutableDoubleVector dimension + pure + GMRESWorkspace + { gmresBasisPayload = basisPayload, + gmresPreconditionedPayload = preconditionedPayload, + gmresHessenbergPayload = hessenbergPayload, + gmresCosines = cosines, + gmresSines = sines, + gmresProjectedResidual = projectedResidual, + gmresYValues = yValues, + gmresResidualVector = residualVector, + gmresImageVector = imageVector, + gmresWorkVector = workVector, + gmresPreconditionerScratchA = preconditionerScratchA, + gmresPreconditionerScratchB = preconditionerScratchB, + gmresDimension = dimension, + gmresRestartDimension = restartDimension, + gmresKrylovColumnCount = krylovColumnCount + } + +type GmresState :: Type +data GmresState + = GmresRunning !Int + | GmresConverged !Int !Double + +gmresRestartCycle :: + SparseGMRESConfig -> + SparsePreconditioner -> + SparseCSR Double -> + U.Vector Double -> + MutableDoubleVector s -> + GMRESWorkspace s -> + GmresState -> + Int -> + ExceptT SparseIterativeFailure (ST s) GmresState +gmresRestartCycle config preconditioner sparseMatrix rhsValues currentGuess workspace stateValue _ = + case stateValue of + GmresConverged _ _ -> pure stateValue + GmresRunning totalIterations -> + if totalIterations >= sgcIterationLimit config + then pure stateValue + else do + betaValue <- + lift $ do + residualIntoMutable sparseMatrix rhsValues currentGuess (gmresImageVector workspace) (gmresResidualVector workspace) + stableNormMutableVector (gmresResidualVector workspace) + unless (fieldValueValid betaValue) $ + throwE (SparseInvalidInput "GMRES residual norm is not representable as a finite Double") + if betaValue <= sgcTolerance config + then pure (GmresConverged totalIterations betaValue) + else do + lift (prepareRestartBasis betaValue workspace) + arnoldiState <- + U.foldM' + (gmresArnoldiStep config preconditioner sparseMatrix workspace) + (ArnoldiRunning 0 betaValue) + (U.enumFromN 0 (min (gmresRestartDimension workspace) (sgcIterationLimit config - totalIterations))) + applyArnoldiCorrection currentGuess workspace arnoldiState + trueResidualNorm <- lift (gmresTrueResidualNorm sparseMatrix rhsValues currentGuess workspace) + except (gmresStateAfterArnoldi config totalIterations arnoldiState trueResidualNorm) + +type ArnoldiState :: Type +data ArnoldiState + = ArnoldiRunning !Int !Double + | ArnoldiConverged !Int !Double + | ArnoldiHappyBreakdown !Int !Double + +gmresArnoldiStep :: + SparseGMRESConfig -> + SparsePreconditioner -> + SparseCSR Double -> + GMRESWorkspace s -> + ArnoldiState -> + Int -> + ExceptT SparseIterativeFailure (ST s) ArnoldiState +gmresArnoldiStep config preconditioner sparseMatrix workspace stateValue _ = + case stateValue of + ArnoldiConverged _ _ -> pure stateValue + ArnoldiHappyBreakdown _ _ -> pure stateValue + ArnoldiRunning completedSteps residualNormValue -> + if residualNormValue <= sgcTolerance config + then pure (ArnoldiConverged completedSteps residualNormValue) + else do + let basisVector = basisColumn workspace completedSteps + preconditionedVector = preconditionedColumn workspace completedSteps + nextBasisNorm <- + lift $ do + applySparsePreconditionerMutable + preconditioner + basisVector + (gmresPreconditionerScratchA workspace) + (gmresPreconditionerScratchB workspace) + preconditionedVector + csrMatVecIntoMutable sparseMatrix preconditionedVector (gmresWorkVector workspace) + orthogonalizeAgainstBasis workspace completedSteps + stableNormMutableVector (gmresWorkVector workspace) + unless (fieldValueValid nextBasisNorm) $ throwE (SparseInvalidInput "GMRES Arnoldi norm is not representable as a finite Double") + lift $ do + writeHessenbergEntry workspace (completedSteps + 1) completedSteps nextBasisNorm + writeNextBasisColumn workspace completedSteps nextBasisNorm + residualAfterRotation <- + applyPreviousRotations workspace completedSteps + *> applyNextRotation workspace completedSteps + let !nextCompletedSteps = completedSteps + 1 + if nextBasisNorm <= solverEpsilon + then pure (ArnoldiHappyBreakdown nextCompletedSteps residualAfterRotation) + else + if residualAfterRotation <= sgcTolerance config + then pure (ArnoldiConverged nextCompletedSteps residualAfterRotation) + else pure (ArnoldiRunning nextCompletedSteps residualAfterRotation) + +prepareRestartBasis :: Double -> GMRESWorkspace s -> ST s () +prepareRestartBasis !betaValue workspace = do + MU.set (gmresProjectedResidual workspace) 0.0 + MU.unsafeWrite (gmresProjectedResidual workspace) 0 betaValue + copyMutableVector (gmresResidualVector workspace) (basisColumn workspace 0) + scaleMutableVector (1.0 / betaValue) (basisColumn workspace 0) + +orthogonalizeAgainstBasis :: GMRESWorkspace s -> Int -> ST s () +orthogonalizeAgainstBasis workspace stepIndex = + U.foldM' orthogonalizeColumn () (U.enumFromN 0 (stepIndex + 1)) + where + orthogonalizeColumn () basisIndex = do + coefficientValue <- dotMutableVector (gmresWorkVector workspace) (basisColumn workspace basisIndex) + writeHessenbergEntry workspace basisIndex stepIndex coefficientValue + addScaledMutableVector (negate coefficientValue) (basisColumn workspace basisIndex) (gmresWorkVector workspace) + +writeNextBasisColumn :: GMRESWorkspace s -> Int -> Double -> ST s () +writeNextBasisColumn workspace stepIndex nextBasisNorm = + if nextBasisNorm <= solverEpsilon + then pure () + else do + copyMutableVector (gmresWorkVector workspace) (basisColumn workspace (stepIndex + 1)) + scaleMutableVector (1.0 / nextBasisNorm) (basisColumn workspace (stepIndex + 1)) + +applyPreviousRotations :: GMRESWorkspace s -> Int -> ExceptT SparseIterativeFailure (ST s) () +applyPreviousRotations workspace stepIndex = + U.mapM_ applyRotation (U.enumFromN 0 stepIndex) + where + applyRotation rotationIndex = do + (cosValue, sinValue, firstEntry, secondEntry) <- + lift $ do + cosValue <- MU.unsafeRead (gmresCosines workspace) rotationIndex + sinValue <- MU.unsafeRead (gmresSines workspace) rotationIndex + firstEntry <- readHessenbergEntry workspace rotationIndex stepIndex + secondEntry <- readHessenbergEntry workspace (rotationIndex + 1) stepIndex + pure (cosValue, sinValue, firstEntry, secondEntry) + let !firstRotated = cosValue * firstEntry + sinValue * secondEntry + !secondRotated = negate sinValue * firstEntry + cosValue * secondEntry + unless (fieldValueValid firstRotated && fieldValueValid secondRotated) $ + throwE (SparseInvalidInput "GMRES previous Givens rotation produced a non-finite Hessenberg entry") + lift $ do + writeHessenbergEntry workspace rotationIndex stepIndex firstRotated + writeHessenbergEntry workspace (rotationIndex + 1) stepIndex secondRotated + +applyNextRotation :: GMRESWorkspace s -> Int -> ExceptT SparseIterativeFailure (ST s) Double +applyNextRotation workspace stepIndex = do + (diagonalEntry, subdiagonalEntry) <- + lift $ do + diagonalEntry <- readHessenbergEntry workspace stepIndex stepIndex + subdiagonalEntry <- readHessenbergEntry workspace (stepIndex + 1) stepIndex + pure (diagonalEntry, subdiagonalEntry) + GivensRotation cosValue sinValue rValue <- + except (gmresGivensCoefficients diagonalEntry subdiagonalEntry) + (projectedEntry, projectedNext) <- + lift $ do + projectedEntry <- MU.unsafeRead (gmresProjectedResidual workspace) stepIndex + projectedNext <- MU.unsafeRead (gmresProjectedResidual workspace) (stepIndex + 1) + pure (projectedEntry, projectedNext) + let !rotatedEntry = cosValue * projectedEntry + sinValue * projectedNext + !rotatedNext = negate sinValue * projectedEntry + cosValue * projectedNext + unless (fieldValueValid rotatedEntry && fieldValueValid rotatedNext) $ + throwE (SparseInvalidInput "GMRES Givens rotation produced a non-finite projected residual") + lift $ do + MU.unsafeWrite (gmresCosines workspace) stepIndex cosValue + MU.unsafeWrite (gmresSines workspace) stepIndex sinValue + writeHessenbergEntry workspace stepIndex stepIndex rValue + writeHessenbergEntry workspace (stepIndex + 1) stepIndex 0.0 + MU.unsafeWrite (gmresProjectedResidual workspace) stepIndex rotatedEntry + MU.unsafeWrite (gmresProjectedResidual workspace) (stepIndex + 1) rotatedNext + pure (abs rotatedNext) + +applyArnoldiCorrection :: MutableDoubleVector s -> GMRESWorkspace s -> ArnoldiState -> ExceptT SparseIterativeFailure (ST s) () +applyArnoldiCorrection currentGuess workspace arnoldiState = + case arnoldiStepCount arnoldiState of + 0 -> pure () + stepCount -> do + solveProjectedUpperTriangular workspace stepCount + lift $ do + MU.set (gmresWorkVector workspace) 0.0 + U.foldM' addColumnContribution () (U.enumFromN 0 stepCount) + addScaledMutableVector 1.0 (gmresWorkVector workspace) currentGuess + where + addColumnContribution () columnIndex = do + coefficientValue <- MU.unsafeRead (gmresYValues workspace) columnIndex + addScaledMutableVector coefficientValue (preconditionedColumn workspace columnIndex) (gmresWorkVector workspace) + +gmresTrueResidualNorm :: + SparseCSR Double -> + U.Vector Double -> + MutableDoubleVector s -> + GMRESWorkspace s -> + ST s Double +gmresTrueResidualNorm sparseMatrix rhsValues currentGuess workspace = do + residualIntoMutable sparseMatrix rhsValues currentGuess (gmresImageVector workspace) (gmresResidualVector workspace) + stableNormMutableVector (gmresResidualVector workspace) + +solveProjectedUpperTriangular :: GMRESWorkspace s -> Int -> ExceptT SparseIterativeFailure (ST s) () +solveProjectedUpperTriangular workspace stepCount = do + projectedScale <- lift (projectedUpperScale workspace stepCount) + unless (fieldValueValid projectedScale && projectedScale > 0.0) $ + throwE (SparseInvalidInput "GMRES projected triangular solve has no finite non-zero scale") + U.mapM_ (solveRow projectedScale) (U.enumFromN 0 stepCount) + where + solveRow projectedScale reverseOffset = do + let !rowIndex = stepCount - reverseOffset - 1 + (laterProduct, rhsValue, diagonalEntry) <- + lift $ do + laterProduct <- projectedLaterProduct workspace stepCount rowIndex + rhsValue <- MU.unsafeRead (gmresProjectedResidual workspace) rowIndex + diagonalEntry <- readHessenbergEntry workspace rowIndex rowIndex + pure (laterProduct, rhsValue, diagonalEntry) + let !numerator = rhsValue - laterProduct + !diagonalThreshold = solverEpsilon * projectedScale + unless (fieldValueValid rhsValue && fieldValueValid laterProduct && fieldValueValid numerator && fieldValueValid diagonalEntry) $ + throwE (SparseInvalidInput "GMRES projected triangular solve encountered non-finite arithmetic") + when (abs diagonalEntry <= diagonalThreshold) $ + throwE (SparseInvalidInput "GMRES projected triangular solve encountered a zero or scale-negligible diagonal") + let !solutionValue = numerator / diagonalEntry + unless (fieldValueValid solutionValue) $ + throwE (SparseInvalidInput "GMRES projected triangular solve would write a non-finite correction") + lift (MU.unsafeWrite (gmresYValues workspace) rowIndex solutionValue) + +projectedLaterProduct :: GMRESWorkspace s -> Int -> Int -> ST s Double +projectedLaterProduct workspace stepCount rowIndex = + U.foldM' accumulateLater 0.0 (U.enumFromN (rowIndex + 1) (stepCount - rowIndex - 1)) + where + accumulateLater !accumulator columnIndex = do + hEntry <- readHessenbergEntry workspace rowIndex columnIndex + yValue <- MU.unsafeRead (gmresYValues workspace) columnIndex + pure (accumulator + hEntry * yValue) + +projectedUpperScale :: GMRESWorkspace s -> Int -> ST s Double +projectedUpperScale workspace stepCount = + U.foldM' + accumulateRowScale + 0.0 + (U.enumFromN 0 stepCount) + where + accumulateRowScale currentScale rowIndex = + U.foldM' + (\rowScale columnIndex -> max rowScale . abs <$> readHessenbergEntry workspace rowIndex columnIndex) + currentScale + (U.enumFromN rowIndex (stepCount - rowIndex)) + +gmresStateAfterArnoldi :: SparseGMRESConfig -> Int -> ArnoldiState -> Double -> Either SparseIterativeFailure GmresState +gmresStateAfterArnoldi config totalIterations arnoldiState trueResidualNorm + | not (fieldValueValid trueResidualNorm) = + Left (SparseInvalidInput "GMRES restart produced a non-finite true residual") + | ArnoldiHappyBreakdown _ projectedResidualNorm <- arnoldiState, + trueResidualNorm > sgcTolerance config = + Left + ( SparseInvalidInput + ( "GMRES happy breakdown did not certify the true residual; projected residual " + <> show projectedResidualNorm + <> ", true residual " + <> show trueResidualNorm + ) + ) + | trueResidualNorm <= sgcTolerance config = + Right (GmresConverged nextTotal trueResidualNorm) + | otherwise = Right (GmresRunning nextTotal) + where + !nextTotal = totalIterations + arnoldiStepCount arnoldiState + +gmresResultFromState :: Int -> MutableDoubleVector s -> GmresState -> ExceptT SparseIterativeFailure (ST s) SparseIterativeResult +gmresResultFromState iterationLimit currentGuess stateValue = + case stateValue of + GmresRunning _ -> throwE (SparseIterationBudgetExceeded iterationLimit) + GmresConverged iterationCount residualNormValue -> do + solutionVector <- lift (freezeMutableDoubleVector currentGuess) + pure + SparseIterativeResult + { sparseSolution = solutionVector, + sparseIterations = iterationCount, + sparseResidualNorm = residualNormValue + } + +arnoldiStepCount :: ArnoldiState -> Int +arnoldiStepCount stateValue = + case stateValue of + ArnoldiRunning stepCount _ -> stepCount + ArnoldiConverged stepCount _ -> stepCount + ArnoldiHappyBreakdown stepCount _ -> stepCount + +type GivensRotation :: Type +data GivensRotation = GivensRotation !Double !Double !Double + +gmresGivensCoefficients :: Double -> Double -> Either SparseIterativeFailure GivensRotation +gmresGivensCoefficients diagonalEntry subdiagonalEntry + | not (fieldValueValid diagonalEntry && fieldValueValid subdiagonalEntry) = + Left (SparseInvalidInput "GMRES Givens rotation requires finite Hessenberg entries") + | scaleValue == 0.0 = Right (GivensRotation 1.0 0.0 0.0) + | otherwise = + let scaledDiagonal = diagonalEntry / scaleValue + scaledSubdiagonal = subdiagonalEntry / scaleValue + radius = scaleValue * sqrt (scaledDiagonal * scaledDiagonal + scaledSubdiagonal * scaledSubdiagonal) + in if fieldValueValid radius && radius > 0.0 + then Right (GivensRotation (diagonalEntry / radius) (subdiagonalEntry / radius) radius) + else Left (SparseInvalidInput "GMRES Givens radius is not representable as a finite positive Double") + where + scaleValue = max (abs diagonalEntry) (abs subdiagonalEntry) + +basisColumn :: GMRESWorkspace s -> Int -> MutableDoubleVector s +basisColumn workspace columnIndex = + MU.unsafeSlice (columnIndex * gmresDimension workspace) (gmresDimension workspace) (gmresBasisPayload workspace) +{-# INLINE basisColumn #-} + +preconditionedColumn :: GMRESWorkspace s -> Int -> MutableDoubleVector s +preconditionedColumn workspace columnIndex = + MU.unsafeSlice (columnIndex * gmresDimension workspace) (gmresDimension workspace) (gmresPreconditionedPayload workspace) +{-# INLINE preconditionedColumn #-} + +readHessenbergEntry :: GMRESWorkspace s -> Int -> Int -> ST s Double +readHessenbergEntry workspace rowIndex columnIndex = + MU.unsafeRead (gmresHessenbergPayload workspace) (hessenbergOffset workspace rowIndex columnIndex) +{-# INLINE readHessenbergEntry #-} + +writeHessenbergEntry :: GMRESWorkspace s -> Int -> Int -> Double -> ST s () +writeHessenbergEntry workspace rowIndex columnIndex value = + MU.unsafeWrite (gmresHessenbergPayload workspace) (hessenbergOffset workspace rowIndex columnIndex) value +{-# INLINE writeHessenbergEntry #-} + +hessenbergOffset :: GMRESWorkspace s -> Int -> Int -> Int +hessenbergOffset workspace rowIndex columnIndex = + rowIndex + columnIndex * gmresKrylovColumnCount workspace +{-# INLINE hessenbergOffset #-} + +restartCycleCount :: Int -> Int -> Int +restartCycleCount iterationLimit restartDimension = + if iterationLimit <= 0 + then 0 + else + iterationLimit `quot` restartDimension + + if iterationLimit `rem` restartDimension == 0 then 0 else 1 + +validateGmresPreconditioner :: SparseCSR Double -> SparsePreconditioner -> Either SparseIterativeFailure () +validateGmresPreconditioner sparseMatrix preconditioner = + if preconditionerDimension preconditioner == csrRows sparseMatrix + then Right () + else Left (SparseInvalidInput "GMRES requires preconditioner dimension equal to matrix dimension") + +stableNormMutableVector :: MutableDoubleVector s -> ST s Double +stableNormMutableVector vectorValue = do + (scaleValue, scaledSumSquares) <- + U.foldM' + accumulateScaledSquare + (0.0, 1.0) + (U.enumFromN 0 (MU.length vectorValue)) + pure + ( if scaleValue == 0.0 + then 0.0 + else scaleValue * sqrt scaledSumSquares + ) + where + accumulateScaledSquare (!scaleValue, !scaledSumSquares) entryIndex = do + entryValue <- abs <$> MU.unsafeRead vectorValue entryIndex + pure (accumulateEntry scaleValue scaledSumSquares entryValue) + + accumulateEntry :: Double -> Double -> Double -> (Double, Double) + accumulateEntry scaleValue scaledSumSquares entryValue + | not (fieldValueValid entryValue) = (entryValue, entryValue) + | entryValue == 0.0 = (scaleValue, scaledSumSquares) + | scaleValue < entryValue = + ( entryValue, + 1.0 + scaledSumSquares * (scaleValue / entryValue) * (scaleValue / entryValue) + ) + | otherwise = + ( scaleValue, + scaledSumSquares + (entryValue / scaleValue) * (entryValue / scaleValue) + ) +{-# INLINE stableNormMutableVector #-} diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/IncompleteCholesky0.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/IncompleteCholesky0.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/IncompleteCholesky0.hs @@ -0,0 +1,511 @@ +{-# LANGUAGE BangPatterns #-} + +module Moonlight.LinAlg.Pure.Sparse.Solver.IncompleteCholesky0 + ( IC0Factor, + ic0FactorDimension, + incompleteCholesky0Factor, + applyIC0FactorMutable, + applyIC0FactorAndDotMutable, + ) +where + +import Control.Monad.ST (ST, runST) +import Data.Kind (Type) +import Data.Vector.Unboxed qualified as U +import Data.Vector.Unboxed.Mutable qualified as MU +import Moonlight.Core (fieldValueValid) +import Moonlight.LinAlg.Pure.Sparse.Solver.Common (solverEpsilon) +import Moonlight.LinAlg.Pure.Sparse.Solver.Mutable + ( MutableDoubleVector, + copyMutableVector, + dotMutableVector, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Types + ( IC0Config (..), + SparseIterativeFailure (..), + ) +import Moonlight.LinAlg.Pure.Sparse.Types + ( SparseCSR, + csrCols, + csrColumnIndicesVector, + csrRows, + csrRowOffsetsVector, + csrValuesVector, + ) +import Prelude + +type IC0Factor :: Type +data IC0Factor = IC0Factor + { ic0FactorDimension :: !Int, + ic0FactorRowOffsets :: !(U.Vector Int), + ic0FactorColumnIndices :: !(U.Vector Int), + ic0FactorValues :: !(U.Vector Double), + ic0FactorDiagonal :: !(U.Vector Double), + ic0FactorPivots :: !(U.Vector Double) + } + deriving stock (Eq, Show) + +type IC0SymbolicPattern :: Type +data IC0SymbolicPattern = IC0SymbolicPattern + { ic0PatternRowOffsets :: !(U.Vector Int), + ic0PatternColumnIndices :: !(U.Vector Int), + ic0PatternValues :: !(U.Vector Double), + ic0PatternDiagonalValues :: !(U.Vector Double), + ic0PatternSuspectedNullspace :: !Bool + } + +type IC0RowPayload :: Type +data IC0RowPayload = IC0RowPayload + { ic0RowColumns :: ![Int], + ic0RowValues :: ![Double], + ic0RowDiagonal :: !(Maybe Double), + ic0RowSum :: !Double + } + +incompleteCholesky0Factor :: + IC0Config -> + SparseCSR Double -> + Either SparseIterativeFailure IC0Factor +incompleteCholesky0Factor configValue sparseMatrix = do + shiftValue <- validateIC0Shift configValue + validateIC0Shape sparseMatrix + symbolicPattern <- ic0SymbolicPattern sparseMatrix + factorIC0SymbolicPattern + shiftValue + symbolicPattern + +applyIC0FactorMutable :: + IC0Factor -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +applyIC0FactorMutable factorValue sourceVector scratchVector targetVector = do + ic0ForwardSolveIntoMutable factorValue sourceVector scratchVector + ic0BackwardSolveIntoMutable factorValue scratchVector targetVector +{-# INLINE applyIC0FactorMutable #-} + +applyIC0FactorAndDotMutable :: + IC0Factor -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +applyIC0FactorAndDotMutable factorValue sourceVector scratchVector targetVector = do + applyIC0FactorMutable factorValue sourceVector scratchVector targetVector + dotMutableVector sourceVector targetVector +{-# INLINE applyIC0FactorAndDotMutable #-} + +validateIC0Shift :: IC0Config -> Either SparseIterativeFailure Double +validateIC0Shift configValue = + case ic0DiagonalShift configValue of + Nothing -> Right 0.0 + Just shiftValue + | fieldValueValid shiftValue && shiftValue >= 0.0 -> Right shiftValue + | otherwise -> Left (SparseInvalidDiagonalShift shiftValue) + +validateIC0Shape :: SparseCSR Double -> Either SparseIterativeFailure () +validateIC0Shape sparseMatrix + | csrRows sparseMatrix /= csrCols sparseMatrix = + Left (SparseNonSquareSparsePreconditioner (csrRows sparseMatrix) (csrCols sparseMatrix)) + | otherwise = Right () + +ic0SymbolicPattern :: + SparseCSR Double -> + Either SparseIterativeFailure IC0SymbolicPattern +ic0SymbolicPattern sparseMatrix = do + rows <- traverse (ic0RowPayload sparseMatrix) [0 .. dimension - 1] + let diagonalValues = traverse ic0DiagonalFromPayload (zip [0 ..] rows) + case diagonalValues of + Left failureValue -> Left failureValue + Right rowDiagonals -> + let !rowCounts = length . ic0RowColumns <$> rows + !rowOffsets = U.fromList (scanl (+) 0 rowCounts) + !columnIndices = U.fromList (ic0RowColumns =<< rows) + !lowerValues = U.fromList (ic0RowValues =<< rows) + !diagonalVector = U.fromList rowDiagonals + !nullspaceLike = + all + (\rowValue -> abs (ic0RowSum rowValue) <= solverEpsilon) + rows + in Right + IC0SymbolicPattern + { ic0PatternRowOffsets = rowOffsets, + ic0PatternColumnIndices = columnIndices, + ic0PatternValues = lowerValues, + ic0PatternDiagonalValues = diagonalVector, + ic0PatternSuspectedNullspace = nullspaceLike + } + where + !dimension = csrRows sparseMatrix + +ic0DiagonalFromPayload :: + (Int, IC0RowPayload) -> + Either SparseIterativeFailure Double +ic0DiagonalFromPayload (rowIndex, rowValue) = + case ic0RowDiagonal rowValue of + Nothing -> Left (SparseMissingDiagonal rowIndex) + Just diagonalValue -> Right diagonalValue + +ic0RowPayload :: + SparseCSR Double -> + Int -> + Either SparseIterativeFailure IC0RowPayload +ic0RowPayload sparseMatrix rowIndex = + collectEntries startOffset [] [] Nothing 0.0 + where + !rowOffsets = csrRowOffsetsVector sparseMatrix + !columnIndices = csrColumnIndicesVector sparseMatrix + !values = csrValuesVector sparseMatrix + !startOffset = rowOffsets `U.unsafeIndex` rowIndex + !endOffset = rowOffsets `U.unsafeIndex` (rowIndex + 1) + + collectEntries !entryIndex !columnsRev !valuesRev !diagonalValue !rowSum + | entryIndex >= endOffset = + Right + IC0RowPayload + { ic0RowColumns = reverse columnsRev, + ic0RowValues = reverse valuesRev, + ic0RowDiagonal = diagonalValue, + ic0RowSum = rowSum + } + | otherwise = + let !columnIndex = columnIndices `U.unsafeIndex` entryIndex + !entryValue = values `U.unsafeIndex` entryIndex + !nextRowSum = rowSum + entryValue + in if not (fieldValueValid entryValue) + then Left (SparseNonFiniteUpdate rowIndex columnIndex entryValue) + else + case compare columnIndex rowIndex of + LT -> + case findCSRValue sparseMatrix columnIndex rowIndex of + Nothing -> Left (SparseStructuralAsymmetry rowIndex columnIndex) + Just _ -> + collectEntries + (entryIndex + 1) + (columnIndex : columnsRev) + (entryValue : valuesRev) + diagonalValue + nextRowSum + EQ -> + collectEntries + (entryIndex + 1) + columnsRev + valuesRev + (Just entryValue) + nextRowSum + GT -> + case findCSRValue sparseMatrix columnIndex rowIndex of + Nothing -> Left (SparseStructuralAsymmetry rowIndex columnIndex) + Just _ -> + collectEntries + (entryIndex + 1) + columnsRev + valuesRev + diagonalValue + nextRowSum + +factorIC0SymbolicPattern :: + Double -> + IC0SymbolicPattern -> + Either SparseIterativeFailure IC0Factor +factorIC0SymbolicPattern !shiftValue symbolicPattern = + runST $ do + lowerValues <- U.thaw (ic0PatternValues symbolicPattern) + diagonalValues <- MU.unsafeNew dimension + pivotValues <- MU.unsafeNew dimension + resultValue <- + factorRows + lowerValues + diagonalValues + pivotValues + 0 + case resultValue of + Left failureValue -> pure (Left failureValue) + Right () -> do + frozenValues <- U.unsafeFreeze lowerValues + frozenDiagonals <- U.unsafeFreeze diagonalValues + frozenPivots <- U.unsafeFreeze pivotValues + pure + ( Right + IC0Factor + { ic0FactorDimension = dimension, + ic0FactorRowOffsets = ic0PatternRowOffsets symbolicPattern, + ic0FactorColumnIndices = ic0PatternColumnIndices symbolicPattern, + ic0FactorValues = frozenValues, + ic0FactorDiagonal = frozenDiagonals, + ic0FactorPivots = frozenPivots + } + ) + where + !rowOffsets = ic0PatternRowOffsets symbolicPattern + !columnIndices = ic0PatternColumnIndices symbolicPattern + !matrixValues = ic0PatternValues symbolicPattern + !matrixDiagonal = ic0PatternDiagonalValues symbolicPattern + !dimension = U.length matrixDiagonal + + factorRows :: + MU.MVector s Double -> + MU.MVector s Double -> + MU.MVector s Double -> + Int -> + ST s (Either SparseIterativeFailure ()) + factorRows lowerValues diagonalValues pivotValues !rowIndex + | rowIndex >= dimension = pure (Right ()) + | otherwise = do + let !rowStart = rowOffsets `U.unsafeIndex` rowIndex + !rowEnd = rowOffsets `U.unsafeIndex` (rowIndex + 1) + offDiagonalResult <- + factorStrictLowerRow + lowerValues + diagonalValues + rowIndex + rowStart + case offDiagonalResult of + Left failureValue -> pure (Left failureValue) + Right () -> do + correctionValue <- lowerRowSquared lowerValues rowStart rowEnd 0.0 + let !pivotValue = + (matrixDiagonal `U.unsafeIndex` rowIndex) + + shiftValue + - correctionValue + if not (fieldValueValid pivotValue) + then pure (Left (SparseNonFiniteUpdate rowIndex rowIndex pivotValue)) + else + if pivotValue <= solverEpsilon + then + pure + ( Left + ( if ic0PatternSuspectedNullspace symbolicPattern + then SparseSuspectedNullspaceUnanchoredLaplacian rowIndex pivotValue + else SparseNonpositivePivot rowIndex pivotValue + ) + ) + else do + MU.unsafeWrite pivotValues rowIndex pivotValue + MU.unsafeWrite diagonalValues rowIndex (sqrt pivotValue) + factorRows + lowerValues + diagonalValues + pivotValues + (rowIndex + 1) + + factorStrictLowerRow :: + MU.MVector s Double -> + MU.MVector s Double -> + Int -> + Int -> + ST s (Either SparseIterativeFailure ()) + factorStrictLowerRow lowerValues diagonalValues !rowIndex !entryIndex + | entryIndex >= rowEnd = pure (Right ()) + | otherwise = do + let !columnIndex = columnIndices `U.unsafeIndex` entryIndex + !matrixValue = matrixValues `U.unsafeIndex` entryIndex + correctionValue <- + lowerIntersectionProduct + lowerValues + rowIndex + columnIndex + rowStart + (rowOffsets `U.unsafeIndex` columnIndex) + 0.0 + pivotDiagonal <- MU.unsafeRead diagonalValues columnIndex + let !factorValue = (matrixValue - correctionValue) / pivotDiagonal + if not (fieldValueValid factorValue) + then pure (Left (SparseNonFiniteUpdate rowIndex columnIndex factorValue)) + else do + MU.unsafeWrite lowerValues entryIndex factorValue + factorStrictLowerRow + lowerValues + diagonalValues + rowIndex + (entryIndex + 1) + where + !rowStart = rowOffsets `U.unsafeIndex` rowIndex + !rowEnd = rowOffsets `U.unsafeIndex` (rowIndex + 1) + + lowerIntersectionProduct :: + MU.MVector s Double -> + Int -> + Int -> + Int -> + Int -> + Double -> + ST s Double + lowerIntersectionProduct lowerValues !rowIndex !columnIndex !leftEntry !rightEntry !accumulator = + let !leftEnd = rowOffsets `U.unsafeIndex` (rowIndex + 1) + !rightEnd = rowOffsets `U.unsafeIndex` (columnIndex + 1) + in if leftEntry >= leftEnd || rightEntry >= rightEnd + then pure accumulator + else + let !leftColumn = columnIndices `U.unsafeIndex` leftEntry + !rightColumn = columnIndices `U.unsafeIndex` rightEntry + in if leftColumn >= columnIndex || rightColumn >= columnIndex + then pure accumulator + else + case compare leftColumn rightColumn of + LT -> + lowerIntersectionProduct + lowerValues + rowIndex + columnIndex + (leftEntry + 1) + rightEntry + accumulator + EQ -> do + leftValue <- MU.unsafeRead lowerValues leftEntry + rightValue <- MU.unsafeRead lowerValues rightEntry + lowerIntersectionProduct + lowerValues + rowIndex + columnIndex + (leftEntry + 1) + (rightEntry + 1) + (accumulator + leftValue * rightValue) + GT -> + lowerIntersectionProduct + lowerValues + rowIndex + columnIndex + leftEntry + (rightEntry + 1) + accumulator + +lowerRowSquared :: + MU.MVector s Double -> + Int -> + Int -> + Double -> + ST s Double +lowerRowSquared lowerValues !entryIndex !endEntry !accumulator + | entryIndex >= endEntry = pure accumulator + | otherwise = do + factorValue <- MU.unsafeRead lowerValues entryIndex + lowerRowSquared + lowerValues + (entryIndex + 1) + endEntry + (accumulator + factorValue * factorValue) +{-# INLINE lowerRowSquared #-} + +ic0ForwardSolveIntoMutable :: + IC0Factor -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +ic0ForwardSolveIntoMutable factorValue sourceVector targetVector = + solveRows 0 + where + !dimension = ic0FactorDimension factorValue + !rowOffsets = ic0FactorRowOffsets factorValue + !columnIndices = ic0FactorColumnIndices factorValue + !factorValues = ic0FactorValues factorValue + !diagonalValues = ic0FactorDiagonal factorValue + + solveRows !rowIndex + | rowIndex >= dimension = pure () + | otherwise = do + rhsValue <- MU.unsafeRead sourceVector rowIndex + knownProduct <- + lowerKnownProduct + columnIndices + factorValues + targetVector + (rowOffsets `U.unsafeIndex` rowIndex) + (rowOffsets `U.unsafeIndex` (rowIndex + 1)) + 0.0 + MU.unsafeWrite + targetVector + rowIndex + ( (rhsValue - knownProduct) + / (diagonalValues `U.unsafeIndex` rowIndex) + ) + solveRows (rowIndex + 1) +{-# INLINE ic0ForwardSolveIntoMutable #-} + +ic0BackwardSolveIntoMutable :: + IC0Factor -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +ic0BackwardSolveIntoMutable factorValue sourceVector targetVector = do + copyMutableVector sourceVector targetVector + solveRows (dimension - 1) + where + !dimension = ic0FactorDimension factorValue + !rowOffsets = ic0FactorRowOffsets factorValue + !columnIndices = ic0FactorColumnIndices factorValue + !factorValues = ic0FactorValues factorValue + !diagonalValues = ic0FactorDiagonal factorValue + + solveRows !rowIndex + | rowIndex < 0 = pure () + | otherwise = do + rhsValue <- MU.unsafeRead targetVector rowIndex + let !solutionValue = rhsValue / (diagonalValues `U.unsafeIndex` rowIndex) + MU.unsafeWrite targetVector rowIndex solutionValue + scatterLowerTranspose + (rowOffsets `U.unsafeIndex` rowIndex) + (rowOffsets `U.unsafeIndex` (rowIndex + 1)) + solutionValue + solveRows (rowIndex - 1) + + scatterLowerTranspose !entryIndex !endEntry !solutionValue + | entryIndex >= endEntry = pure () + | otherwise = do + let !columnIndex = columnIndices `U.unsafeIndex` entryIndex + !factorEntry = factorValues `U.unsafeIndex` entryIndex + targetValue <- MU.unsafeRead targetVector columnIndex + MU.unsafeWrite + targetVector + columnIndex + (targetValue - factorEntry * solutionValue) + scatterLowerTranspose + (entryIndex + 1) + endEntry + solutionValue +{-# INLINE ic0BackwardSolveIntoMutable #-} + +lowerKnownProduct :: + U.Vector Int -> + U.Vector Double -> + MutableDoubleVector s -> + Int -> + Int -> + Double -> + ST s Double +lowerKnownProduct columnIndices factorValues targetVector !entryIndex !endEntry !accumulator + | entryIndex >= endEntry = pure accumulator + | otherwise = do + let !columnIndex = columnIndices `U.unsafeIndex` entryIndex + !factorValue = factorValues `U.unsafeIndex` entryIndex + targetValue <- MU.unsafeRead targetVector columnIndex + lowerKnownProduct + columnIndices + factorValues + targetVector + (entryIndex + 1) + endEntry + (accumulator + factorValue * targetValue) +{-# INLINE lowerKnownProduct #-} + +findCSRValue :: SparseCSR Double -> Int -> Int -> Maybe Double +findCSRValue sparseMatrix rowIndex columnIndex = + binarySearch startOffset endOffset + where + !rowOffsets = csrRowOffsetsVector sparseMatrix + !columnIndices = csrColumnIndicesVector sparseMatrix + !values = csrValuesVector sparseMatrix + !startOffset = rowOffsets `U.unsafeIndex` rowIndex + !endOffset = rowOffsets `U.unsafeIndex` (rowIndex + 1) + + binarySearch !lo !hi + | lo >= hi = Nothing + | otherwise = + let !mid = lo + ((hi - lo) `div` 2) + !midColumn = columnIndices `U.unsafeIndex` mid + in case compare midColumn columnIndex of + LT -> binarySearch (mid + 1) hi + EQ -> Just (values `U.unsafeIndex` mid) + GT -> binarySearch lo mid +{-# INLINE findCSRValue #-} diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Mutable.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Mutable.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Mutable.hs @@ -0,0 +1,1221 @@ +{-# LANGUAGE BangPatterns #-} + +module Moonlight.LinAlg.Pure.Sparse.Solver.Mutable + ( MutableDoubleVector, + newMutableDoubleVector, + thawMutableDoubleVector, + freezeMutableDoubleVector, + copyImmutableIntoMutable, + copyImmutableSquaredNormIntoMutable, + copyMutableVector, + dotMutableVector, + normMutableVector, + addScaledMutableVector, + scaleMutableVector, + scaledCopyMutableVector, + addScaledPairIntoMutable, + subtractMutableInto, + csrMatVecIntoMutable, + csrMatVecDotIntoMutable, + residualIntoMutable, + csrResidualSquaredIntoMutable, + updateSolutionAndResidualSquaredMutable, + updateDirectionMutable, + initializeZeroJacobiMutable, + divideByDiagonalDotAndCopyMutable, + updateSolutionResidualJacobiMutable, + divideByDiagonalAndDotIntoMutable, + divideByDiagonalIntoMutable, + multiplyByDiagonalIntoMutable, + lowerTriangularSolveIntoMutable, + upperTriangularSolveIntoMutable, + ) +where + +import Control.Monad.ST (ST) +import Data.Kind (Type) +import Data.Primitive.ByteArray + ( ByteArray, + MutableByteArray, + indexByteArray, + readByteArray, + writeByteArray, + ) +import Data.Vector.Primitive qualified as P +import Data.Vector.Primitive.Mutable qualified as PM +import Data.Vector.Unboxed.Base qualified as UB +import Data.Vector.Unboxed qualified as U +import Data.Vector.Unboxed.Mutable qualified as MU +import Moonlight.LinAlg.Pure.Sparse.Types + ( CSRExecutionPlan (..), + SparseCSR, + csrColumnIndicesVector, + csrRows, + csrExecutionPlan, + csrRowOffsetsVector, + csrValuesVector, + ) +import Prelude + +type MutableDoubleVector :: Type -> Type +type MutableDoubleVector s = MU.MVector s Double + +newMutableDoubleVector :: Int -> ST s (MutableDoubleVector s) +newMutableDoubleVector !dimension = + MU.replicate dimension 0.0 +{-# INLINE newMutableDoubleVector #-} + +thawMutableDoubleVector :: U.Vector Double -> ST s (MutableDoubleVector s) +thawMutableDoubleVector = U.thaw +{-# INLINE thawMutableDoubleVector #-} + +freezeMutableDoubleVector :: MutableDoubleVector s -> ST s (U.Vector Double) +freezeMutableDoubleVector = U.freeze +{-# INLINE freezeMutableDoubleVector #-} + +copyImmutableIntoMutable :: + U.Vector Double -> + MutableDoubleVector s -> + ST s () +copyImmutableIntoMutable sourceVector targetVector = + go 0 + where + !dimension = U.length sourceVector + + go !indexValue + | indexValue >= dimension = pure () + | otherwise = do + MU.unsafeWrite + targetVector + indexValue + (sourceVector `U.unsafeIndex` indexValue) + go (indexValue + 1) +{-# INLINE copyImmutableIntoMutable #-} + +copyImmutableSquaredNormIntoMutable :: + U.Vector Double -> + MutableDoubleVector s -> + ST s Double +copyImmutableSquaredNormIntoMutable sourceVector targetVector = + go 0 0.0 + where + !dimension = U.length sourceVector + + go !indexValue !sumSquares + | indexValue >= dimension = pure sumSquares + | otherwise = do + let !entryValue = sourceVector `U.unsafeIndex` indexValue + MU.unsafeWrite targetVector indexValue entryValue + go + (indexValue + 1) + (sumSquares + entryValue * entryValue) +{-# INLINE copyImmutableSquaredNormIntoMutable #-} + +copyMutableVector :: + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +copyMutableVector sourceVector targetVector = + MU.unsafeCopy targetVector sourceVector +{-# INLINE copyMutableVector #-} + +dotMutableVector :: + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +dotMutableVector leftVector rightVector = + go 0 0.0 + where + !dimension = MU.length leftVector + + go !indexValue !accumulator + | indexValue >= dimension = pure accumulator + | otherwise = do + leftValue <- MU.unsafeRead leftVector indexValue + rightValue <- MU.unsafeRead rightVector indexValue + go + (indexValue + 1) + (accumulator + leftValue * rightValue) +{-# INLINE dotMutableVector #-} + +normMutableVector :: MutableDoubleVector s -> ST s Double +normMutableVector vectorValue = + sqrt <$> dotMutableVector vectorValue vectorValue +{-# INLINE normMutableVector #-} + +addScaledMutableVector :: + Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +addScaledMutableVector !scaleValue sourceVector targetVector = + go 0 + where + !dimension = MU.length sourceVector + + go !indexValue + | indexValue >= dimension = pure () + | otherwise = do + sourceValue <- MU.unsafeRead sourceVector indexValue + targetValue <- MU.unsafeRead targetVector indexValue + MU.unsafeWrite + targetVector + indexValue + (targetValue + scaleValue * sourceValue) + go (indexValue + 1) +{-# INLINE addScaledMutableVector #-} + +scaleMutableVector :: + Double -> + MutableDoubleVector s -> + ST s () +scaleMutableVector !scaleValue targetVector = + go 0 + where + !dimension = MU.length targetVector + + go !indexValue + | indexValue >= dimension = pure () + | otherwise = do + targetValue <- MU.unsafeRead targetVector indexValue + MU.unsafeWrite + targetVector + indexValue + (scaleValue * targetValue) + go (indexValue + 1) +{-# INLINE scaleMutableVector #-} + +scaledCopyMutableVector :: + Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +scaledCopyMutableVector !scaleValue sourceVector targetVector = + go 0 + where + !dimension = MU.length sourceVector + + go !indexValue + | indexValue >= dimension = pure () + | otherwise = do + sourceValue <- MU.unsafeRead sourceVector indexValue + MU.unsafeWrite + targetVector + indexValue + (scaleValue * sourceValue) + go (indexValue + 1) +{-# INLINE scaledCopyMutableVector #-} + +addScaledPairIntoMutable :: + Double -> + MutableDoubleVector s -> + Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +addScaledPairIntoMutable + !leftScale + leftVector + !rightScale + rightVector + targetVector = + go 0 + where + !dimension = MU.length targetVector + + go !indexValue + | indexValue >= dimension = pure () + | otherwise = do + leftValue <- MU.unsafeRead leftVector indexValue + rightValue <- MU.unsafeRead rightVector indexValue + MU.unsafeWrite + targetVector + indexValue + (leftScale * leftValue + rightScale * rightValue) + go (indexValue + 1) +{-# INLINE addScaledPairIntoMutable #-} + +subtractMutableInto :: + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +subtractMutableInto leftVector rightVector targetVector = + addScaledPairIntoMutable + 1.0 + leftVector + (-1.0) + rightVector + targetVector +{-# INLINE subtractMutableInto #-} + +csrMatVecIntoMutable :: + SparseCSR Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +csrMatVecIntoMutable sparseMatrix inputVector targetVector = + case csrExecutionPlan sparseMatrix of + CSRGeneral -> + csrMatVecGeneralIntoMutable + sparseMatrix + inputVector + targetVector + CSRContiguousBand 2 2 + | csrRows sparseMatrix >= 5 -> + pentadiagonalMatVecIntoMutable + (csrRows sparseMatrix) + (csrValuesVector sparseMatrix) + inputVector + targetVector + CSRContiguousBand lowerBandwidth upperBandwidth -> + contiguousBandMatVecIntoMutable + (csrRows sparseMatrix) + lowerBandwidth + upperBandwidth + (csrValuesVector sparseMatrix) + inputVector + targetVector +{-# INLINE csrMatVecIntoMutable #-} + +csrMatVecGeneralIntoMutable :: + SparseCSR Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +csrMatVecGeneralIntoMutable sparseMatrix inputVector targetVector = + writeRows 0 + where + !rowCount = csrRows sparseMatrix + !rowOffsets = csrRowOffsetsVector sparseMatrix + !columnIndices = csrColumnIndicesVector sparseMatrix + !coefficients = csrValuesVector sparseMatrix + + writeRows !rowIndex + | rowIndex >= rowCount = pure () + | otherwise = do + rowValue <- + csrRowDotMutable + rowOffsets + columnIndices + coefficients + inputVector + rowIndex + MU.unsafeWrite targetVector rowIndex rowValue + writeRows (rowIndex + 1) +{-# INLINE csrMatVecGeneralIntoMutable #-} + +pentadiagonalMatVecIntoMutable :: + Int -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +pentadiagonalMatVecIntoMutable + rowCount + (UB.V_Double (P.Vector coefficientBase _ coefficientArray)) + (UB.MV_Double (PM.MVector inputBase _ inputArray)) + (UB.MV_Double (PM.MVector targetBase _ targetArray)) = + go 0 + where + go !rowIndex + | rowIndex >= rowCount = pure () + | otherwise = do + rowValue <- + pentadiagonalRowDotMutable + rowCount + coefficientBase + coefficientArray + inputBase + inputArray + rowIndex + writeByteArray + targetArray + (targetBase + rowIndex) + rowValue + go (rowIndex + 1) +{-# INLINE pentadiagonalMatVecIntoMutable #-} + +pentadiagonalMatVecDotIntoMutable :: + Int -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +pentadiagonalMatVecDotIntoMutable + rowCount + (UB.V_Double (P.Vector coefficientBase _ coefficientArray)) + (UB.MV_Double (PM.MVector inputBase _ inputArray)) + (UB.MV_Double (PM.MVector targetBase _ targetArray)) = + go 0 0.0 + where + go !rowIndex !dotAccumulator + | rowIndex >= rowCount = pure dotAccumulator + | otherwise = do + rowValue <- + pentadiagonalRowDotMutable + rowCount + coefficientBase + coefficientArray + inputBase + inputArray + rowIndex + inputValue <- + readByteArray inputArray (inputBase + rowIndex) + writeByteArray + targetArray + (targetBase + rowIndex) + rowValue + go + (rowIndex + 1) + (dotAccumulator + (inputValue :: Double) * rowValue) +{-# INLINE pentadiagonalMatVecDotIntoMutable #-} + +pentadiagonalResidualSquaredIntoMutable :: + Int -> + U.Vector Double -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +pentadiagonalResidualSquaredIntoMutable + rowCount + (UB.V_Double (P.Vector coefficientBase _ coefficientArray)) + (UB.V_Double (P.Vector rhsBase _ rhsArray)) + (UB.MV_Double (PM.MVector guessBase _ guessArray)) + (UB.MV_Double (PM.MVector residualBase _ residualArray)) = + go 0 0.0 + where + go !rowIndex !sumSquares + | rowIndex >= rowCount = pure sumSquares + | otherwise = do + imageValue <- + pentadiagonalRowDotMutable + rowCount + coefficientBase + coefficientArray + guessBase + guessArray + rowIndex + let !rhsValue = + ( indexByteArray + rhsArray + (rhsBase + rowIndex) + :: Double + ) + !residualValue = rhsValue - imageValue + writeByteArray + residualArray + (residualBase + rowIndex) + residualValue + go + (rowIndex + 1) + (sumSquares + residualValue * residualValue) +{-# INLINE pentadiagonalResidualSquaredIntoMutable #-} + +pentadiagonalRowDotMutable :: + Int -> + Int -> + ByteArray -> + Int -> + MutableByteArray s -> + Int -> + ST s Double +pentadiagonalRowDotMutable + rowCount + coefficientBase + coefficientArray + inputBase + inputArray + rowIndex + | rowIndex == 0 = do + input0 <- readInput 0 + input1 <- readInput 1 + input2 <- readInput 2 + pure + ( coefficientAt 0 * input0 + + coefficientAt 1 * input1 + + coefficientAt 2 * input2 + ) + | rowIndex == 1 = do + input0 <- readInput 0 + input1 <- readInput 1 + input2 <- readInput 2 + input3 <- readInput 3 + pure + ( coefficientAt 3 * input0 + + coefficientAt 4 * input1 + + coefficientAt 5 * input2 + + coefficientAt 6 * input3 + ) + | rowIndex + 2 < rowCount = do + let !entryIndex = 5 * rowIndex - 3 + input0 <- readInput (rowIndex - 2) + input1 <- readInput (rowIndex - 1) + input2 <- readInput rowIndex + input3 <- readInput (rowIndex + 1) + input4 <- readInput (rowIndex + 2) + pure + ( coefficientAt entryIndex * input0 + + coefficientAt (entryIndex + 1) * input1 + + coefficientAt (entryIndex + 2) * input2 + + coefficientAt (entryIndex + 3) * input3 + + coefficientAt (entryIndex + 4) * input4 + ) + | rowIndex + 1 < rowCount = do + let !entryIndex = 5 * rowCount - 13 + input0 <- readInput (rowCount - 4) + input1 <- readInput (rowCount - 3) + input2 <- readInput (rowCount - 2) + input3 <- readInput (rowCount - 1) + pure + ( coefficientAt entryIndex * input0 + + coefficientAt (entryIndex + 1) * input1 + + coefficientAt (entryIndex + 2) * input2 + + coefficientAt (entryIndex + 3) * input3 + ) + | otherwise = do + let !entryIndex = 5 * rowCount - 9 + input0 <- readInput (rowCount - 3) + input1 <- readInput (rowCount - 2) + input2 <- readInput (rowCount - 1) + pure + ( coefficientAt entryIndex * input0 + + coefficientAt (entryIndex + 1) * input1 + + coefficientAt (entryIndex + 2) * input2 + ) + where + coefficientAt !entryIndex = + ( indexByteArray + coefficientArray + (coefficientBase + entryIndex) + :: Double + ) + readInput !columnIndex = + readByteArray inputArray (inputBase + columnIndex) +{-# INLINE pentadiagonalRowDotMutable #-} + +contiguousBandMatVecIntoMutable :: + Int -> + Int -> + Int -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +contiguousBandMatVecIntoMutable + rowCount + lowerBandwidth + upperBandwidth + coefficients + inputVector + targetVector = + writeRows 0 0 + where + writeRows !rowIndex !entryIndex + | rowIndex >= rowCount = pure () + | otherwise = do + let !firstColumn = max 0 (rowIndex - lowerBandwidth) + !lastColumn = min (rowCount - 1) (rowIndex + upperBandwidth) + !entryCount = lastColumn - firstColumn + 1 + rowValue <- + accumulateBand + entryIndex + firstColumn + entryCount + 0.0 + MU.unsafeWrite targetVector rowIndex rowValue + writeRows (rowIndex + 1) (entryIndex + entryCount) + + accumulateBand !entryIndex !columnIndex !remaining !accumulator + | remaining <= 0 = pure accumulator + | otherwise = do + inputValue <- MU.unsafeRead inputVector columnIndex + let !coefficient = coefficients `U.unsafeIndex` entryIndex + accumulateBand + (entryIndex + 1) + (columnIndex + 1) + (remaining - 1) + (accumulator + coefficient * inputValue) +{-# INLINE contiguousBandMatVecIntoMutable #-} + +csrMatVecDotIntoMutable :: + SparseCSR Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +csrMatVecDotIntoMutable sparseMatrix inputVector targetVector = + case csrExecutionPlan sparseMatrix of + CSRGeneral -> + csrMatVecDotGeneralIntoMutable + sparseMatrix + inputVector + targetVector + CSRContiguousBand 2 2 + | csrRows sparseMatrix >= 5 -> + pentadiagonalMatVecDotIntoMutable + (csrRows sparseMatrix) + (csrValuesVector sparseMatrix) + inputVector + targetVector + CSRContiguousBand lowerBandwidth upperBandwidth -> + contiguousBandMatVecDotIntoMutable + (csrRows sparseMatrix) + lowerBandwidth + upperBandwidth + (csrValuesVector sparseMatrix) + inputVector + targetVector +{-# INLINE csrMatVecDotIntoMutable #-} + +csrMatVecDotGeneralIntoMutable :: + SparseCSR Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +csrMatVecDotGeneralIntoMutable sparseMatrix inputVector targetVector = + writeRows 0 0.0 + where + !rowCount = csrRows sparseMatrix + !rowOffsets = csrRowOffsetsVector sparseMatrix + !columnIndices = csrColumnIndicesVector sparseMatrix + !coefficients = csrValuesVector sparseMatrix + + writeRows !rowIndex !dotAccumulator + | rowIndex >= rowCount = pure dotAccumulator + | otherwise = do + rowValue <- + csrRowDotMutable + rowOffsets + columnIndices + coefficients + inputVector + rowIndex + inputValue <- MU.unsafeRead inputVector rowIndex + MU.unsafeWrite targetVector rowIndex rowValue + writeRows + (rowIndex + 1) + (dotAccumulator + inputValue * rowValue) +{-# INLINE csrMatVecDotGeneralIntoMutable #-} + +contiguousBandMatVecDotIntoMutable :: + Int -> + Int -> + Int -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +contiguousBandMatVecDotIntoMutable + rowCount + lowerBandwidth + upperBandwidth + coefficients + inputVector + targetVector = + writeRows 0 0 0.0 + where + writeRows !rowIndex !entryIndex !dotAccumulator + | rowIndex >= rowCount = pure dotAccumulator + | otherwise = do + let !firstColumn = max 0 (rowIndex - lowerBandwidth) + !lastColumn = min (rowCount - 1) (rowIndex + upperBandwidth) + !entryCount = lastColumn - firstColumn + 1 + rowValue <- + accumulateBand + entryIndex + firstColumn + entryCount + 0.0 + inputValue <- MU.unsafeRead inputVector rowIndex + MU.unsafeWrite targetVector rowIndex rowValue + writeRows + (rowIndex + 1) + (entryIndex + entryCount) + (dotAccumulator + inputValue * rowValue) + + accumulateBand !entryIndex !columnIndex !remaining !accumulator + | remaining <= 0 = pure accumulator + | otherwise = do + inputValue <- MU.unsafeRead inputVector columnIndex + let !coefficient = coefficients `U.unsafeIndex` entryIndex + accumulateBand + (entryIndex + 1) + (columnIndex + 1) + (remaining - 1) + (accumulator + coefficient * inputValue) +{-# INLINE contiguousBandMatVecDotIntoMutable #-} + +residualIntoMutable :: + SparseCSR Double -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +residualIntoMutable + sparseMatrix + rhsValues + guessVector + imageVector + residualVector = do + csrMatVecIntoMutable sparseMatrix guessVector imageVector + writeResidual 0 + where + !dimension = U.length rhsValues + + writeResidual !indexValue + | indexValue >= dimension = pure () + | otherwise = do + imageValue <- MU.unsafeRead imageVector indexValue + MU.unsafeWrite + residualVector + indexValue + (rhsValues `U.unsafeIndex` indexValue - imageValue) + writeResidual (indexValue + 1) +{-# INLINE residualIntoMutable #-} + +csrResidualSquaredIntoMutable :: + SparseCSR Double -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +csrResidualSquaredIntoMutable + sparseMatrix + rhsValues + guessVector + residualVector = + case csrExecutionPlan sparseMatrix of + CSRGeneral -> + csrResidualSquaredGeneralIntoMutable + sparseMatrix + rhsValues + guessVector + residualVector + CSRContiguousBand 2 2 + | csrRows sparseMatrix >= 5 -> + pentadiagonalResidualSquaredIntoMutable + (csrRows sparseMatrix) + (csrValuesVector sparseMatrix) + rhsValues + guessVector + residualVector + CSRContiguousBand lowerBandwidth upperBandwidth -> + contiguousBandResidualSquaredIntoMutable + (csrRows sparseMatrix) + lowerBandwidth + upperBandwidth + (csrValuesVector sparseMatrix) + rhsValues + guessVector + residualVector +{-# INLINE csrResidualSquaredIntoMutable #-} + +csrResidualSquaredGeneralIntoMutable :: + SparseCSR Double -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +csrResidualSquaredGeneralIntoMutable + sparseMatrix + rhsValues + guessVector + residualVector = + writeRows 0 0.0 + where + !rowCount = csrRows sparseMatrix + !rowOffsets = csrRowOffsetsVector sparseMatrix + !columnIndices = csrColumnIndicesVector sparseMatrix + !coefficients = csrValuesVector sparseMatrix + + writeRows !rowIndex !sumSquares + | rowIndex >= rowCount = pure sumSquares + | otherwise = do + imageValue <- + csrRowDotMutable + rowOffsets + columnIndices + coefficients + guessVector + rowIndex + let !residualValue = + rhsValues `U.unsafeIndex` rowIndex - imageValue + MU.unsafeWrite residualVector rowIndex residualValue + writeRows + (rowIndex + 1) + (sumSquares + residualValue * residualValue) +{-# INLINE csrResidualSquaredGeneralIntoMutable #-} + +contiguousBandResidualSquaredIntoMutable :: + Int -> + Int -> + Int -> + U.Vector Double -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +contiguousBandResidualSquaredIntoMutable + rowCount + lowerBandwidth + upperBandwidth + coefficients + rhsValues + guessVector + residualVector = + writeRows 0 0 0.0 + where + writeRows !rowIndex !entryIndex !sumSquares + | rowIndex >= rowCount = pure sumSquares + | otherwise = do + let !firstColumn = max 0 (rowIndex - lowerBandwidth) + !lastColumn = min (rowCount - 1) (rowIndex + upperBandwidth) + !entryCount = lastColumn - firstColumn + 1 + imageValue <- + accumulateBand + entryIndex + firstColumn + entryCount + 0.0 + let !residualValue = + rhsValues `U.unsafeIndex` rowIndex - imageValue + MU.unsafeWrite residualVector rowIndex residualValue + writeRows + (rowIndex + 1) + (entryIndex + entryCount) + (sumSquares + residualValue * residualValue) + + accumulateBand !entryIndex !columnIndex !remaining !accumulator + | remaining <= 0 = pure accumulator + | otherwise = do + guessValue <- MU.unsafeRead guessVector columnIndex + let !coefficient = coefficients `U.unsafeIndex` entryIndex + accumulateBand + (entryIndex + 1) + (columnIndex + 1) + (remaining - 1) + (accumulator + coefficient * guessValue) +{-# INLINE contiguousBandResidualSquaredIntoMutable #-} + +updateSolutionAndResidualSquaredMutable :: + Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +updateSolutionAndResidualSquaredMutable + !alphaValue + directionVector + imageDirectionVector + guessVector + residualVector = + go 0 0.0 + where + !dimension = MU.length guessVector + + go !indexValue !sumSquares + | indexValue >= dimension = pure sumSquares + | otherwise = do + directionValue <- MU.unsafeRead directionVector indexValue + imageDirectionValue <- + MU.unsafeRead imageDirectionVector indexValue + guessValue <- MU.unsafeRead guessVector indexValue + residualValue <- MU.unsafeRead residualVector indexValue + let !nextGuessValue = + guessValue + alphaValue * directionValue + !nextResidualValue = + residualValue - alphaValue * imageDirectionValue + MU.unsafeWrite guessVector indexValue nextGuessValue + MU.unsafeWrite residualVector indexValue nextResidualValue + go + (indexValue + 1) + (sumSquares + nextResidualValue * nextResidualValue) +{-# INLINE updateSolutionAndResidualSquaredMutable #-} + +updateDirectionMutable :: + Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +updateDirectionMutable !betaValue sourceVector directionVector = + go 0 + where + !dimension = MU.length directionVector + + go !indexValue + | indexValue >= dimension = pure () + | otherwise = do + sourceValue <- MU.unsafeRead sourceVector indexValue + directionValue <- MU.unsafeRead directionVector indexValue + MU.unsafeWrite + directionVector + indexValue + (sourceValue + betaValue * directionValue) + go (indexValue + 1) +{-# INLINE updateDirectionMutable #-} + +initializeZeroJacobiMutable :: + U.Vector Double -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s (Double, Double) +initializeZeroJacobiMutable + rhsValues + diagonalValues + residualVector + preconditionedResidualVector + directionVector = + go 0 0.0 0.0 + where + !dimension = U.length rhsValues + + go !indexValue !residualSquared !rhoValue + | indexValue >= dimension = + pure (residualSquared, rhoValue) + | otherwise = do + let !residualValue = rhsValues `U.unsafeIndex` indexValue + !preconditionedValue = + residualValue + / (diagonalValues `U.unsafeIndex` indexValue) + MU.unsafeWrite residualVector indexValue residualValue + MU.unsafeWrite + preconditionedResidualVector + indexValue + preconditionedValue + MU.unsafeWrite + directionVector + indexValue + preconditionedValue + go + (indexValue + 1) + (residualSquared + residualValue * residualValue) + (rhoValue + residualValue * preconditionedValue) +{-# INLINE initializeZeroJacobiMutable #-} + +divideByDiagonalDotAndCopyMutable :: + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +divideByDiagonalDotAndCopyMutable + diagonalValues + sourceVector + targetVector + directionVector = + go 0 0.0 + where + !dimension = U.length diagonalValues + + go !indexValue !dotAccumulator + | indexValue >= dimension = pure dotAccumulator + | otherwise = do + sourceValue <- MU.unsafeRead sourceVector indexValue + let !targetValue = + sourceValue + / (diagonalValues `U.unsafeIndex` indexValue) + MU.unsafeWrite targetVector indexValue targetValue + MU.unsafeWrite directionVector indexValue targetValue + go + (indexValue + 1) + (dotAccumulator + sourceValue * targetValue) +{-# INLINE divideByDiagonalDotAndCopyMutable #-} + +updateSolutionResidualJacobiMutable :: + Double -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s (Double, Double) +updateSolutionResidualJacobiMutable + !alphaValue + diagonalValues + directionVector + imageDirectionVector + guessVector + residualVector + preconditionedResidualVector = + go 0 0.0 0.0 + where + !dimension = MU.length guessVector + + go !indexValue !residualSquared !rhoValue + | indexValue >= dimension = + pure (residualSquared, rhoValue) + | otherwise = do + directionValue <- MU.unsafeRead directionVector indexValue + imageDirectionValue <- + MU.unsafeRead imageDirectionVector indexValue + guessValue <- MU.unsafeRead guessVector indexValue + residualValue <- MU.unsafeRead residualVector indexValue + let !nextGuessValue = + guessValue + alphaValue * directionValue + !nextResidualValue = + residualValue - alphaValue * imageDirectionValue + !preconditionedValue = + nextResidualValue + / (diagonalValues `U.unsafeIndex` indexValue) + MU.unsafeWrite guessVector indexValue nextGuessValue + MU.unsafeWrite residualVector indexValue nextResidualValue + MU.unsafeWrite + preconditionedResidualVector + indexValue + preconditionedValue + go + (indexValue + 1) + (residualSquared + nextResidualValue * nextResidualValue) + (rhoValue + nextResidualValue * preconditionedValue) +{-# INLINE updateSolutionResidualJacobiMutable #-} + +divideByDiagonalAndDotIntoMutable :: + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +divideByDiagonalAndDotIntoMutable + diagonalValues + sourceVector + targetVector = + go 0 0.0 + where + !dimension = U.length diagonalValues + + go !indexValue !dotAccumulator + | indexValue >= dimension = pure dotAccumulator + | otherwise = do + sourceValue <- MU.unsafeRead sourceVector indexValue + let !targetValue = + sourceValue + / (diagonalValues `U.unsafeIndex` indexValue) + MU.unsafeWrite targetVector indexValue targetValue + go + (indexValue + 1) + (dotAccumulator + sourceValue * targetValue) +{-# INLINE divideByDiagonalAndDotIntoMutable #-} + +divideByDiagonalIntoMutable :: + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +divideByDiagonalIntoMutable diagonalValues sourceVector targetVector = + go 0 + where + !dimension = U.length diagonalValues + + go !indexValue + | indexValue >= dimension = pure () + | otherwise = do + sourceValue <- MU.unsafeRead sourceVector indexValue + MU.unsafeWrite + targetVector + indexValue + (sourceValue / (diagonalValues `U.unsafeIndex` indexValue)) + go (indexValue + 1) +{-# INLINE divideByDiagonalIntoMutable #-} + +multiplyByDiagonalIntoMutable :: + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +multiplyByDiagonalIntoMutable diagonalValues sourceVector targetVector = + go 0 + where + !dimension = U.length diagonalValues + + go !indexValue + | indexValue >= dimension = pure () + | otherwise = do + sourceValue <- MU.unsafeRead sourceVector indexValue + MU.unsafeWrite + targetVector + indexValue + (sourceValue * (diagonalValues `U.unsafeIndex` indexValue)) + go (indexValue + 1) +{-# INLINE multiplyByDiagonalIntoMutable #-} + +lowerTriangularSolveIntoMutable :: + U.Vector Double -> + SparseCSR Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +lowerTriangularSolveIntoMutable + diagonalValues + sparseMatrix + rhsVector + targetVector = + solveRows 0 + where + !dimension = csrRows sparseMatrix + !rowOffsets = csrRowOffsetsVector sparseMatrix + !columnIndices = csrColumnIndicesVector sparseMatrix + !coefficients = csrValuesVector sparseMatrix + + solveRows !rowIndex + | rowIndex >= dimension = pure () + | otherwise = do + rhsValue <- MU.unsafeRead rhsVector rowIndex + knownProduct <- + lowerKnownProduct + rowOffsets + columnIndices + coefficients + targetVector + rowIndex + MU.unsafeWrite + targetVector + rowIndex + ( (rhsValue - knownProduct) + / (diagonalValues `U.unsafeIndex` rowIndex) + ) + solveRows (rowIndex + 1) +{-# INLINE lowerTriangularSolveIntoMutable #-} + +upperTriangularSolveIntoMutable :: + U.Vector Double -> + SparseCSR Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +upperTriangularSolveIntoMutable + diagonalValues + sparseMatrix + rhsVector + targetVector = + solveRows (dimension - 1) + where + !dimension = csrRows sparseMatrix + !rowOffsets = csrRowOffsetsVector sparseMatrix + !columnIndices = csrColumnIndicesVector sparseMatrix + !coefficients = csrValuesVector sparseMatrix + + solveRows !rowIndex + | rowIndex < 0 = pure () + | otherwise = do + rhsValue <- MU.unsafeRead rhsVector rowIndex + knownProduct <- + upperKnownProduct + rowOffsets + columnIndices + coefficients + targetVector + rowIndex + MU.unsafeWrite + targetVector + rowIndex + ( (rhsValue - knownProduct) + / (diagonalValues `U.unsafeIndex` rowIndex) + ) + solveRows (rowIndex - 1) +{-# INLINE upperTriangularSolveIntoMutable #-} + +csrRowDotMutable :: + U.Vector Int -> + U.Vector Int -> + U.Vector Double -> + MutableDoubleVector s -> + Int -> + ST s Double +csrRowDotMutable + rowOffsets + columnIndices + coefficients + inputVector + rowIndex = + go startIndex 0.0 + where + !startIndex = rowOffsets `U.unsafeIndex` rowIndex + !stopIndex = rowOffsets `U.unsafeIndex` (rowIndex + 1) + + go !entryIndex !accumulator + | entryIndex >= stopIndex = pure accumulator + | otherwise = do + let !columnIndex = + columnIndices `U.unsafeIndex` entryIndex + !coefficientValue = + coefficients `U.unsafeIndex` entryIndex + inputValue <- MU.unsafeRead inputVector columnIndex + go + (entryIndex + 1) + (accumulator + coefficientValue * inputValue) +{-# INLINE csrRowDotMutable #-} + +lowerKnownProduct :: + U.Vector Int -> + U.Vector Int -> + U.Vector Double -> + MutableDoubleVector s -> + Int -> + ST s Double +lowerKnownProduct + rowOffsets + columnIndices + coefficients + solutionVector + rowIndex = + go startIndex 0.0 + where + !startIndex = rowOffsets `U.unsafeIndex` rowIndex + !stopIndex = rowOffsets `U.unsafeIndex` (rowIndex + 1) + + go !entryIndex !accumulator + | entryIndex >= stopIndex = pure accumulator + | otherwise = + let !columnIndex = + columnIndices `U.unsafeIndex` entryIndex + !coefficientValue = + coefficients `U.unsafeIndex` entryIndex + in if columnIndex < rowIndex + then do + solutionValue <- + MU.unsafeRead solutionVector columnIndex + go + (entryIndex + 1) + (accumulator + coefficientValue * solutionValue) + else go (entryIndex + 1) accumulator +{-# INLINE lowerKnownProduct #-} + +upperKnownProduct :: + U.Vector Int -> + U.Vector Int -> + U.Vector Double -> + MutableDoubleVector s -> + Int -> + ST s Double +upperKnownProduct + rowOffsets + columnIndices + coefficients + solutionVector + rowIndex = + go startIndex 0.0 + where + !startIndex = rowOffsets `U.unsafeIndex` rowIndex + !stopIndex = rowOffsets `U.unsafeIndex` (rowIndex + 1) + + go !entryIndex !accumulator + | entryIndex >= stopIndex = pure accumulator + | otherwise = + let !columnIndex = + columnIndices `U.unsafeIndex` entryIndex + !coefficientValue = + coefficients `U.unsafeIndex` entryIndex + in if columnIndex > rowIndex + then do + solutionValue <- + MU.unsafeRead solutionVector columnIndex + go + (entryIndex + 1) + (accumulator + coefficientValue * solutionValue) + else go (entryIndex + 1) accumulator +{-# INLINE upperKnownProduct #-} diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Preconditioner.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Preconditioner.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Preconditioner.hs @@ -0,0 +1,179 @@ +module Moonlight.LinAlg.Pure.Sparse.Solver.Preconditioner + ( SparsePreconditioner (..), + applySparsePreconditionerMutable, + applySparsePreconditionerAndDotMutable, + preconditionerDimension, + compileSparsePreconditioner, + ) +where + +import Control.Monad.ST (ST) +import Data.Kind (Type) +import Data.Vector.Unboxed qualified as U +import Moonlight.Core (fieldValueValid) +import Moonlight.LinAlg.Pure.Sparse.Solver.Common + ( shiftedDiagonalValue, + solverEpsilon, + sparseDiagonal, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.IncompleteCholesky0 + ( IC0Factor, + applyIC0FactorAndDotMutable, + applyIC0FactorMutable, + ic0FactorDimension, + incompleteCholesky0Factor, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Mutable + ( MutableDoubleVector, + copyMutableVector, + divideByDiagonalAndDotIntoMutable, + divideByDiagonalIntoMutable, + dotMutableVector, + lowerTriangularSolveIntoMutable, + multiplyByDiagonalIntoMutable, + scaleMutableVector, + upperTriangularSolveIntoMutable, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Types + ( SparseIterativeFailure (..), + SparsePreconditionerFamily (..), + ) +import Moonlight.LinAlg.Pure.Sparse.Types (SparseCSR, csrRows) +import Prelude + +type SparsePreconditioner :: Type +data SparsePreconditioner + = IdentitySparsePreconditioner !Int + | DiagonalSparsePreconditioner !(U.Vector Double) + | SsorSparsePreconditioner !Double !(U.Vector Double) !(U.Vector Double) !(SparseCSR Double) + | IncompleteCholesky0SparsePreconditioner !IC0Factor + +preconditionerDimension :: SparsePreconditioner -> Int +preconditionerDimension preconditionerValue = + case preconditionerValue of + IdentitySparsePreconditioner dimension -> dimension + DiagonalSparsePreconditioner diagonalValues -> U.length diagonalValues + SsorSparsePreconditioner _ diagonalValues _ _ -> U.length diagonalValues + IncompleteCholesky0SparsePreconditioner factorValue -> ic0FactorDimension factorValue + +applySparsePreconditionerMutable :: + SparsePreconditioner -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +applySparsePreconditionerMutable preconditionerValue sourceVector scratchA scratchB targetVector = + case preconditionerValue of + IdentitySparsePreconditioner _ -> + copyMutableVector sourceVector targetVector + DiagonalSparsePreconditioner diagonalValues -> + divideByDiagonalIntoMutable diagonalValues sourceVector targetVector + SsorSparsePreconditioner omegaValue diagonalValues scaledDiagonalValues sparseMatrix -> do + lowerTriangularSolveIntoMutable scaledDiagonalValues sparseMatrix sourceVector scratchA + multiplyByDiagonalIntoMutable diagonalValues scratchA scratchB + upperTriangularSolveIntoMutable scaledDiagonalValues sparseMatrix scratchB targetVector + scaleMutableVector ((2.0 - omegaValue) / omegaValue) targetVector + IncompleteCholesky0SparsePreconditioner factorValue -> + applyIC0FactorMutable factorValue sourceVector scratchA targetVector + +applySparsePreconditionerAndDotMutable :: + SparsePreconditioner -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s Double +applySparsePreconditionerAndDotMutable + preconditionerValue + sourceVector + scratchA + scratchB + targetVector = + case preconditionerValue of + IdentitySparsePreconditioner _ -> do + copyMutableVector sourceVector targetVector + dotMutableVector sourceVector sourceVector + DiagonalSparsePreconditioner diagonalValues -> + divideByDiagonalAndDotIntoMutable + diagonalValues + sourceVector + targetVector + SsorSparsePreconditioner + omegaValue + diagonalValues + scaledDiagonalValues + sparseMatrix -> do + lowerTriangularSolveIntoMutable + scaledDiagonalValues + sparseMatrix + sourceVector + scratchA + multiplyByDiagonalIntoMutable + diagonalValues + scratchA + scratchB + upperTriangularSolveIntoMutable + scaledDiagonalValues + sparseMatrix + scratchB + targetVector + scaleMutableVector + ((2.0 - omegaValue) / omegaValue) + targetVector + dotMutableVector sourceVector targetVector + IncompleteCholesky0SparsePreconditioner factorValue -> + applyIC0FactorAndDotMutable + factorValue + sourceVector + scratchA + targetVector +{-# INLINE applySparsePreconditionerAndDotMutable #-} + +compileSparsePreconditioner :: SparsePreconditionerFamily -> SparseCSR Double -> Either SparseIterativeFailure SparsePreconditioner +compileSparsePreconditioner preconditionerFamily sparseMatrix = + case preconditionerFamily of + IdentitySparsePreconditionerFamily -> + Right (IdentitySparsePreconditioner (csrRows sparseMatrix)) + DiagonalJacobiSparsePreconditionerFamily -> + diagonalPreconditioner sparseMatrix + ShiftedDiagonalJacobiSparsePreconditionerFamily shiftValue -> + shiftedDiagonalPreconditioner shiftValue sparseMatrix + SsorSparsePreconditionerFamily omegaValue -> + ssorPreconditioner omegaValue sparseMatrix + IncompleteCholesky0SparsePreconditionerFamily configValue -> + IncompleteCholesky0SparsePreconditioner + <$> incompleteCholesky0Factor configValue sparseMatrix + +diagonalPreconditioner :: SparseCSR Double -> Either SparseIterativeFailure SparsePreconditioner +diagonalPreconditioner sparseMatrix = do + diagonalValues <- sparseDiagonal sparseMatrix + if U.all strictlyNonZero diagonalValues + then Right (DiagonalSparsePreconditioner diagonalValues) + else Left (SparseInvalidInput "diagonal preconditioner requires a strictly non-zero diagonal") + +shiftedDiagonalPreconditioner :: Double -> SparseCSR Double -> Either SparseIterativeFailure SparsePreconditioner +shiftedDiagonalPreconditioner shiftValue sparseMatrix + | shiftValue <= solverEpsilon = + Left (SparseInvalidInput "shifted diagonal preconditioner requires a strictly positive shift") + | otherwise = do + diagonalValues <- sparseDiagonal sparseMatrix + let shiftedDiagonalValues = U.map (shiftedDiagonalValue shiftValue) diagonalValues + if U.all strictlyNonZero shiftedDiagonalValues + then Right (DiagonalSparsePreconditioner shiftedDiagonalValues) + else Left (SparseInvalidInput "shifted diagonal preconditioner requires a non-degenerate shifted diagonal") + +ssorPreconditioner :: Double -> SparseCSR Double -> Either SparseIterativeFailure SparsePreconditioner +ssorPreconditioner omegaValue sparseMatrix + | omegaValue <= solverEpsilon || omegaValue >= 2.0 - solverEpsilon = + Left (SparseInvalidInput "SSOR preconditioner requires a relaxation parameter strictly between 0 and 2") + | otherwise = do + diagonalValues <- sparseDiagonal sparseMatrix + let scaledDiagonalValues = U.map (/ omegaValue) diagonalValues + if U.all strictlyNonZero scaledDiagonalValues + then Right (SsorSparsePreconditioner omegaValue diagonalValues scaledDiagonalValues sparseMatrix) + else Left (SparseInvalidInput "SSOR preconditioner requires a strictly non-zero diagonal") + +strictlyNonZero :: Double -> Bool +strictlyNonZero value = + fieldValueValid value && abs value > solverEpsilon diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Stationary.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Stationary.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Stationary.hs @@ -0,0 +1,233 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RecordWildCards #-} + +module Moonlight.LinAlg.Pure.Sparse.Solver.Stationary + ( solveSparseJacobi, + solveSparseRichardson, + ) +where + +import Control.Monad.ST (ST, runST) +import Data.Kind (Type) +import Data.Vector.Unboxed qualified as U +import Moonlight.Core (fieldValueValid) +import Moonlight.LinAlg.Pure.Sparse.Solver.Common + ( validateSparseSolverConfiguration, + validateSparseSystemInput, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Mutable + ( MutableDoubleVector, + addScaledMutableVector, + freezeMutableDoubleVector, + newMutableDoubleVector, + normMutableVector, + residualIntoMutable, + scaledCopyMutableVector, + thawMutableDoubleVector, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Preconditioner + ( SparsePreconditioner, + applySparsePreconditionerMutable, + compileSparsePreconditioner, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Types + ( SparseIterativeFailure (..), + SparseIterativeResult (..), + SparsePreconditionerFamily (..), + SparseStationaryIterationConfig (..), + ) +import Moonlight.LinAlg.Pure.Sparse.Types + ( SparseCSR, + csrRowOffsetsVector, + csrRows, + csrValuesVector, + ) +import Prelude + +solveSparseJacobi :: + SparseStationaryIterationConfig -> + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + Either SparseIterativeFailure SparseIterativeResult +solveSparseJacobi config sparseMatrix rhsValues initialGuess = do + validateSparseStationaryInput + "Jacobi" + validJacobiDamping + config + sparseMatrix + rhsValues + initialGuess + preconditioner <- compileSparsePreconditioner DiagonalJacobiSparsePreconditionerFamily sparseMatrix + runValidatedSparseStationary config sparseMatrix rhsValues initialGuess (JacobiStationaryStep (ssicDamping config) preconditioner) + +-- | Solve by damped Richardson iteration under the caller-owned precondition +-- that the operator is symmetric positive-definite. The step is derived from +-- the maximum absolute row sum, an upper bound on the spectral radius for a +-- symmetric operator. +solveSparseRichardson :: + SparseStationaryIterationConfig -> + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + Either SparseIterativeFailure SparseIterativeResult +solveSparseRichardson config sparseMatrix rhsValues initialGuess = do + validateSparseStationaryInput + "Richardson" + validRichardsonDamping + config + sparseMatrix + rhsValues + initialGuess + stepSize <- conservativeRichardsonStep sparseMatrix + runValidatedSparseStationary config sparseMatrix rhsValues initialGuess (RichardsonStationaryStep (ssicDamping config * stepSize)) + +runValidatedSparseStationary :: + SparseStationaryIterationConfig -> + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + StationaryStep -> + Either SparseIterativeFailure SparseIterativeResult +runValidatedSparseStationary SparseStationaryIterationConfig {..} sparseMatrix rhsValues initialGuess stepKind = + runST (solveSparseStationaryMutable ssicTolerance ssicIterationLimit sparseMatrix rhsValues initialGuess stepKind) + +validateSparseStationaryInput :: + String -> + (Double -> Bool) -> + SparseStationaryIterationConfig -> + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + Either SparseIterativeFailure () +validateSparseStationaryInput methodName validDamping config sparseMatrix rhsValues initialGuess = do + validateSparseSystemInput sparseMatrix rhsValues initialGuess + validateSparseSolverConfiguration methodName (ssicTolerance config) (ssicIterationLimit config) + if validDamping (ssicDamping config) + then Right () + else Left (SparseInvalidInput (methodName <> " damping is outside its finite admissible range")) + +validJacobiDamping :: Double -> Bool +validJacobiDamping dampingValue = + fieldValueValid dampingValue && dampingValue > 0.0 && dampingValue <= 1.0 + +validRichardsonDamping :: Double -> Bool +validRichardsonDamping dampingValue = + fieldValueValid dampingValue && dampingValue > 0.0 && dampingValue < 2.0 + +type StationaryStep :: Type +data StationaryStep + = JacobiStationaryStep !Double !SparsePreconditioner + | RichardsonStationaryStep !Double + +solveSparseStationaryMutable :: + Double -> + Int -> + SparseCSR Double -> + U.Vector Double -> + U.Vector Double -> + StationaryStep -> + ST s (Either SparseIterativeFailure SparseIterativeResult) +solveSparseStationaryMutable !toleranceValue !iterationLimit sparseMatrix rhsValues initialGuess stepKind = do + let !dimension = csrRows sparseMatrix + currentVector <- thawMutableDoubleVector initialGuess + residualVector <- newMutableDoubleVector dimension + stepVector <- newMutableDoubleVector dimension + imageVector <- newMutableDoubleVector dimension + preconditionerScratchA <- newMutableDoubleVector dimension + preconditionerScratchB <- newMutableDoubleVector dimension + residualIntoMutable sparseMatrix rhsValues currentVector imageVector residualVector + initialResidualNorm <- normMutableVector residualVector + if initialResidualNorm <= toleranceValue + then Right <$> freezeStationaryResult 0 initialResidualNorm currentVector + else do + finalState <- + U.foldM' + (stationaryIteration toleranceValue sparseMatrix rhsValues currentVector residualVector stepVector imageVector preconditionerScratchA preconditionerScratchB stepKind) + StationaryRunning + (U.enumFromN 0 iterationLimit) + stationaryResultFromState iterationLimit currentVector finalState + +type StationaryState :: Type +data StationaryState + = StationaryRunning + | StationaryConverged !Int !Double + +stationaryIteration :: + Double -> + SparseCSR Double -> + U.Vector Double -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + StationaryStep -> + StationaryState -> + Int -> + ST s StationaryState +stationaryIteration !toleranceValue sparseMatrix rhsValues currentVector residualVector stepVector imageVector preconditionerScratchA preconditionerScratchB stepKind stepState iterationIndex = + case stepState of + StationaryConverged _ _ -> pure stepState + StationaryRunning -> do + writeStationaryStep stepKind residualVector preconditionerScratchA preconditionerScratchB stepVector + addScaledMutableVector 1.0 stepVector currentVector + residualIntoMutable sparseMatrix rhsValues currentVector imageVector residualVector + nextResidualNorm <- normMutableVector residualVector + if nextResidualNorm <= toleranceValue + then pure (StationaryConverged (iterationIndex + 1) nextResidualNorm) + else pure StationaryRunning + +writeStationaryStep :: + StationaryStep -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + MutableDoubleVector s -> + ST s () +writeStationaryStep stepKind residualVector scratchA scratchB stepVector = + case stepKind of + JacobiStationaryStep dampingValue preconditioner -> do + applySparsePreconditionerMutable preconditioner residualVector scratchA scratchB stepVector + scaledCopyMutableVector dampingValue stepVector stepVector + RichardsonStationaryStep richardsonScale -> + scaledCopyMutableVector richardsonScale residualVector stepVector + +stationaryResultFromState :: Int -> MutableDoubleVector s -> StationaryState -> ST s (Either SparseIterativeFailure SparseIterativeResult) +stationaryResultFromState iterationLimit currentVector stepState = + case stepState of + StationaryRunning -> pure (Left (SparseIterationBudgetExceeded iterationLimit)) + StationaryConverged iterationCount residualNormValue -> Right <$> freezeStationaryResult iterationCount residualNormValue currentVector + +freezeStationaryResult :: Int -> Double -> MutableDoubleVector s -> ST s SparseIterativeResult +freezeStationaryResult iterationCount residualNormValue currentVector = do + solutionVector <- freezeMutableDoubleVector currentVector + pure + SparseIterativeResult + { sparseSolution = solutionVector, + sparseIterations = iterationCount, + sparseResidualNorm = residualNormValue + } + +-- For a symmetric positive-definite operator, the maximum absolute row sum +-- bounds the largest eigenvalue. Damping in (0,2) therefore yields a +-- conservative Richardson step without pretending the diagonal is a spectral +-- bound. +conservativeRichardsonStep :: SparseCSR Double -> Either SparseIterativeFailure Double +conservativeRichardsonStep sparseMatrix + | not (fieldValueValid operatorBound) = + Left (SparseInvalidInput "Richardson absolute row-sum bound is non-finite") + | operatorBound <= 0.0 = + Left (SparseInvalidInput "Richardson iteration requires a non-zero operator bound") + | otherwise = Right (1.0 / operatorBound) + where + rowOffsets = csrRowOffsetsVector sparseMatrix + matrixValues = csrValuesVector sparseMatrix + rowAbsoluteSum rowIndex = + let startOffset = rowOffsets `U.unsafeIndex` rowIndex + endOffset = rowOffsets `U.unsafeIndex` (rowIndex + 1) + in U.sum (U.map abs (U.slice startOffset (endOffset - startOffset) matrixValues)) + operatorBound = + U.maximum + (U.cons 0.0 (U.generate (csrRows sparseMatrix) rowAbsoluteSum)) diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Types.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Types.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Types.hs @@ -0,0 +1,81 @@ +module Moonlight.LinAlg.Pure.Sparse.Solver.Types + ( SparseIterativeFailure (..), + SparseIterativeResult (..), + IC0Config (..), + SparsePreconditionerFamily (..), + defaultSparsePreconditionerFamily, + SparseStationaryIterationConfig (..), + SparseConjugateGradientConfig (..), + SparseGMRESConfig (..), + ) +where + +import Data.Kind (Type) +import qualified Data.Vector.Unboxed as U +import Moonlight.Core (MoonlightError) +import Prelude + +type SparseIterativeFailure :: Type +data SparseIterativeFailure + = SparseIterationBudgetExceeded Int + | SparseInvalidInput String + | SparseMissingDiagonal Int + | SparseNonpositivePivot Int Double + | SparseNonFiniteUpdate Int Int Double + | SparseStructuralAsymmetry Int Int + | SparseSuspectedNullspaceUnanchoredLaplacian Int Double + | SparseInvalidDiagonalShift Double + | SparseNonSquareSparsePreconditioner Int Int + | SparseBackendFailure MoonlightError + deriving stock (Eq, Show) + +type SparseIterativeResult :: Type +data SparseIterativeResult = SparseIterativeResult + { sparseSolution :: !(U.Vector Double), + sparseIterations :: !Int, + sparseResidualNorm :: !Double + } + deriving stock (Eq, Show) + +type IC0Config :: Type +data IC0Config = IC0Config + { ic0DiagonalShift :: !(Maybe Double) + } + deriving stock (Eq, Ord, Show, Read) + +type SparsePreconditionerFamily :: Type +data SparsePreconditionerFamily + = IdentitySparsePreconditionerFamily + | DiagonalJacobiSparsePreconditionerFamily + | ShiftedDiagonalJacobiSparsePreconditionerFamily Double + | SsorSparsePreconditionerFamily Double + | IncompleteCholesky0SparsePreconditionerFamily IC0Config + deriving stock (Eq, Ord, Show, Read) + +defaultSparsePreconditionerFamily :: SparsePreconditionerFamily +defaultSparsePreconditionerFamily = DiagonalJacobiSparsePreconditionerFamily + +type SparseStationaryIterationConfig :: Type +data SparseStationaryIterationConfig = SparseStationaryIterationConfig + { ssicTolerance :: Double, + ssicIterationLimit :: Int, + ssicDamping :: Double + } + deriving stock (Eq, Show) + +type SparseConjugateGradientConfig :: Type +data SparseConjugateGradientConfig = SparseConjugateGradientConfig + { scgcTolerance :: Double, + scgcIterationLimit :: Int, + scgcPreconditionerFamily :: SparsePreconditionerFamily + } + deriving stock (Eq, Show) + +type SparseGMRESConfig :: Type +data SparseGMRESConfig = SparseGMRESConfig + { sgcTolerance :: Double, + sgcIterationLimit :: Int, + sgcRestartDimension :: Int, + sgcPreconditionerFamily :: SparsePreconditionerFamily + } + deriving stock (Eq, Show) diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Structured.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Structured.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Structured.hs @@ -0,0 +1,1260 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Sparse.Structured + ( GraphEdge (..), + TridiagonalRejection (..), + symmetricTridiagonalFromCSR, + diagonalCSR, + tridiagonalCSR, + pathLaplacianCSR, + graphLaplacianCSR, + ) +where + +import Control.Monad (foldM) +import Control.Monad.ST (ST, runST) +import Data.Kind (Type) +import Data.List (sortBy) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Maybe (catMaybes, mapMaybe) +import Data.Ord (comparing) +import Data.Vector qualified as Box +import Data.Vector.Unboxed qualified as U +import Data.Vector.Unboxed.Mutable qualified as MU +import Moonlight.Core (AdditiveGroup (..), AdditiveMonoid (..), MoonlightError (..), fieldValueValid) +import Moonlight.LinAlg.Pure.Sparse.Assembly + ( orderedCSRFromEntries, + ) +import Moonlight.LinAlg.Pure.Sparse.Types + ( SparseCSR, + CSRExecutionPlan (..), + csrFromCanonicalVectorsUnchecked, + csrFromCanonicalVectorsWithPlanUnchecked, + csrCols, + csrColumnIndicesVector, + csrRows, + csrRowOffsetsVector, + csrValuesVector, + validateCSR, + ) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal + ( SymmetricTridiagonal, + mkSymmetricTridiagonalVectors, + pathLaplacianBands, + ) +import Prelude + +type TridiagonalRejection :: Type +data TridiagonalRejection + = TridiagonalNonSquare !Int !Int + | TridiagonalOutOfBandEntry !Int !Int + | TridiagonalAsymmetricOffDiagonal + deriving stock (Eq, Show) + +symmetricTridiagonalFromCSR :: + SparseCSR Double -> + Either + MoonlightError + (Either TridiagonalRejection SymmetricTridiagonal) +symmetricTridiagonalFromCSR csrValue = do + validateCSR csrValue + if csrRows csrValue /= csrCols csrValue + then + Right + ( Left + ( TridiagonalNonSquare + (csrRows csrValue) + (csrCols csrValue) + ) + ) + else do + let !values = csrValuesVector csrValue + if U.any (not . fieldValueValid) values + then + Left + ( InvariantViolation + "symmetric tridiagonal classification requires finite CSR entries" + ) + else + case classifyTridiagonalStorage csrValue of + Left rejection -> Right (Left rejection) + Right (diagonalEntries, lowerEntries, upperEntries) -> + if symmetricOffDiagonalEntries lowerEntries upperEntries + then do + let !matrixSize = U.length diagonalEntries + !offDiagonalEntries = + U.generate + (max 0 (matrixSize - 1)) + (U.unsafeIndex upperEntries) + Right + <$> mkSymmetricTridiagonalVectors + diagonalEntries + offDiagonalEntries + else Right (Left TridiagonalAsymmetricOffDiagonal) + +classifyTridiagonalStorage :: + SparseCSR Double -> + Either + TridiagonalRejection + (U.Vector Double, U.Vector Double, U.Vector Double) +classifyTridiagonalStorage csrValue = + runST $ do + diagonalEntries <- MU.replicate matrixSize 0.0 + lowerEntries <- MU.replicate matrixSize 0.0 + upperEntries <- MU.replicate matrixSize 0.0 + classification <- + classifyRows + diagonalEntries + lowerEntries + upperEntries + 0 + case classification of + Left rejection -> pure (Left rejection) + Right () -> do + frozenDiagonal <- U.unsafeFreeze diagonalEntries + frozenLower <- U.unsafeFreeze lowerEntries + frozenUpper <- U.unsafeFreeze upperEntries + pure + ( Right + (frozenDiagonal, frozenLower, frozenUpper) + ) + where + !matrixSize = csrRows csrValue + !rowOffsets = csrRowOffsetsVector csrValue + !columnIndices = csrColumnIndicesVector csrValue + !values = csrValuesVector csrValue + + classifyRows :: + MU.MVector s Double -> + MU.MVector s Double -> + MU.MVector s Double -> + Int -> + ST s (Either TridiagonalRejection ()) + classifyRows diagonalEntries lowerEntries upperEntries !rowIndex + | rowIndex >= matrixSize = pure (Right ()) + | otherwise = do + let !startIndex = rowOffsets `U.unsafeIndex` rowIndex + !stopIndex = rowOffsets `U.unsafeIndex` (rowIndex + 1) + rowClassification <- + classifyRowEntries + diagonalEntries + lowerEntries + upperEntries + rowIndex + startIndex + stopIndex + case rowClassification of + Left rejection -> pure (Left rejection) + Right () -> + classifyRows + diagonalEntries + lowerEntries + upperEntries + (rowIndex + 1) + + classifyRowEntries :: + MU.MVector s Double -> + MU.MVector s Double -> + MU.MVector s Double -> + Int -> + Int -> + Int -> + ST s (Either TridiagonalRejection ()) + classifyRowEntries + diagonalEntries + lowerEntries + upperEntries + !rowIndex + !entryIndex + !stopIndex + | entryIndex >= stopIndex = pure (Right ()) + | otherwise = do + let !columnIndex = columnIndices `U.unsafeIndex` entryIndex + !entryValue = values `U.unsafeIndex` entryIndex + !offset = columnIndex - rowIndex + case offset of + -1 -> do + addMutableEntry lowerEntries rowIndex entryValue + continue + 0 -> do + addMutableEntry diagonalEntries rowIndex entryValue + continue + 1 -> do + addMutableEntry upperEntries rowIndex entryValue + continue + _ -> + pure + ( Left + ( TridiagonalOutOfBandEntry + rowIndex + columnIndex + ) + ) + where + continue = + classifyRowEntries + diagonalEntries + lowerEntries + upperEntries + rowIndex + (entryIndex + 1) + stopIndex + +addMutableEntry :: + MU.MVector s Double -> + Int -> + Double -> + ST s () +addMutableEntry targetVector !entryIndex !entryValue = do + currentValue <- MU.unsafeRead targetVector entryIndex + MU.unsafeWrite + targetVector + entryIndex + (currentValue + entryValue) +{-# INLINE addMutableEntry #-} + +symmetricOffDiagonalEntries :: + U.Vector Double -> + U.Vector Double -> + Bool +symmetricOffDiagonalEntries lowerEntries upperEntries = + go 0 + where + !entryCount = max 0 (U.length lowerEntries - 1) + + go !entryIndex + | entryIndex >= entryCount = True + | upperEntries `U.unsafeIndex` entryIndex + == lowerEntries `U.unsafeIndex` (entryIndex + 1) = + go (entryIndex + 1) + | otherwise = False +{-# INLINE symmetricOffDiagonalEntries #-} + +type GraphEdge :: Type -> Type +data GraphEdge vertex = GraphEdge + { graphEdgeLeft :: !vertex, + graphEdgeRight :: !vertex, + graphEdgeWeight :: !Double + } + deriving stock (Eq, Ord, Show) + +type IndexedGraphEdge :: Type +data IndexedGraphEdge = IndexedGraphEdge + { indexedGraphEdgeLeft :: !Int, + indexedGraphEdgeRight :: !Int, + indexedGraphEdgeWeight :: !Double + } + deriving stock (Eq, Show) + +type IndexedGraphEdgeRows :: Type +data IndexedGraphEdgeRows = IndexedGraphEdgeRows + { indexedGraphEdgeOffsets :: !(U.Vector Int), + indexedGraphEdgeRights :: !(U.Vector Int), + indexedGraphEdgeWeights :: !(U.Vector Double) + } + +type CollectedGraphEdges :: Type +data CollectedGraphEdges = CollectedGraphEdges + { collectedGraphEdgeCount :: !Int, + collectedGraphEdgesAreSorted :: !Bool + } + +type VertexIndex :: Type -> Type +data VertexIndex vertex + = AscendingVertexIndex !(Box.Vector vertex) + | MapVertexIndex !(Map vertex Int) + +diagonalCSR :: + (Eq a, AdditiveGroup a, U.Unbox a) => + [a] -> + Either MoonlightError (SparseCSR a) +diagonalCSR diagonalEntries = + let dimension = length diagonalEntries + in orderedCSRFromEntries + dimension + dimension + (nonZeroDiagonalEntries diagonalEntries) + +tridiagonalCSR :: + (Eq a, AdditiveGroup a, U.Unbox a) => + [a] -> + [a] -> + Either MoonlightError (SparseCSR a) +tridiagonalCSR diagonalEntries offDiagonalEntries + | actualOffDiagonalCount /= expectedOffDiagonalCount = + Left + ( InvariantViolation + ( "symmetric tridiagonal CSR off-diagonal length mismatch: expected " + <> show expectedOffDiagonalCount + <> " but received " + <> show actualOffDiagonalCount + ) + ) + | otherwise = + orderedCSRFromEntries + dimension + dimension + (tridiagonalEntries diagonalEntries offDiagonalEntries) + where + dimension = length diagonalEntries + expectedOffDiagonalCount = max 0 (dimension - 1) + actualOffDiagonalCount = length offDiagonalEntries + +pathLaplacianCSR :: Int -> Either MoonlightError (SparseCSR Double) +pathLaplacianCSR dimension = do + (diagonalEntries, offDiagonalEntries) <- pathLaplacianBands dimension + tridiagonalCSR diagonalEntries offDiagonalEntries + +graphLaplacianCSR :: + (Ord vertex, Show vertex) => + [vertex] -> + [GraphEdge vertex] -> + Either MoonlightError (SparseCSR Double) +graphLaplacianCSR vertexOrder graphEdges = do + vertexIndices <- buildVertexIndices vertexOrder + let dimension = length vertexOrder + matrixValue <- + case vertexIndices of + AscendingVertexIndex vertices -> + case pathGraphLaplacianCSRFromAscendingEdges vertices graphEdges of + Just pathValue -> pure pathValue + Nothing -> graphLaplacianCSRFromGenericEdges vertexIndices dimension graphEdges + MapVertexIndex _ -> + graphLaplacianCSRFromGenericEdges vertexIndices dimension graphEdges + if U.all fieldValueValid (csrValuesVector matrixValue) + then Right matrixValue + else Left (InvariantViolation "graph Laplacian accumulation overflowed to a non-finite matrix entry") +{-# INLINE graphLaplacianCSR #-} + +graphLaplacianCSRFromGenericEdges :: + (Ord vertex, Show vertex) => + VertexIndex vertex -> + Int -> + [GraphEdge vertex] -> + Either MoonlightError (SparseCSR Double) +graphLaplacianCSRFromGenericEdges vertexIndices dimension graphEdges = do + combinedEdges <- + indexedGraphEdgesByLeft + vertexIndices + dimension + graphEdges + pure + ( case pathGraphLaplacianCSRFromIndexedEdges dimension combinedEdges of + Just pathValue -> pathValue + Nothing -> graphLaplacianCSRFromIndexedEdges dimension combinedEdges + ) + +pathGraphLaplacianCSRFromAscendingEdges :: + forall vertex. + Ord vertex => + Box.Vector vertex -> + [GraphEdge vertex] -> + Maybe (SparseCSR Double) +pathGraphLaplacianCSRFromAscendingEdges vertices graphEdges + | dimension <= 1 = + if null graphEdges + then Just (emptySquareCSR dimension) + else Nothing + | otherwise = + runST $ do + edgeWeights <- MU.replicate (dimension - 1) 0.0 + collectionResult <- collectAscendingPathEdges edgeWeights 0 graphEdges + case collectionResult of + Nothing -> pure Nothing + Just () -> do + frozenWeights <- U.unsafeFreeze edgeWeights + pure + ( if U.any (== 0.0) frozenWeights + then Nothing + else Just (pathGraphLaplacianCSRFromWeights dimension frozenWeights) + ) + where + !dimension = Box.length vertices + + collectAscendingPathEdges :: + MU.MVector s Double -> + Int -> + [GraphEdge vertex] -> + ST s (Maybe ()) + collectAscendingPathEdges _ !pathIndex [] + | pathIndex <= dimension - 1 = pure (Just ()) + | otherwise = pure Nothing + collectAscendingPathEdges edgeWeights !pathIndex edges@(edgeValue : remainingEdges) + | pathIndex >= dimension - 1 = pure Nothing + | ascendingPathEdgeMatches pathIndex edgeValue = + let !weightValue = graphEdgeWeight edgeValue + in if not (fieldValueValid weightValue) || weightValue < 0.0 + then pure Nothing + else + if weightValue == 0.0 + then collectAscendingPathEdges edgeWeights pathIndex remainingEdges + else do + addMutableEntry edgeWeights pathIndex weightValue + collectAscendingPathEdges edgeWeights pathIndex remainingEdges + | otherwise = + collectAscendingPathEdges edgeWeights (pathIndex + 1) edges + + ascendingPathEdgeMatches :: Int -> GraphEdge vertex -> Bool + ascendingPathEdgeMatches !pathIndex edgeValue = + let !leftVertex = graphEdgeLeft edgeValue + !rightVertex = graphEdgeRight edgeValue + !expectedLeft = vertices `Box.unsafeIndex` pathIndex + !expectedRight = vertices `Box.unsafeIndex` (pathIndex + 1) + in (leftVertex == expectedLeft && rightVertex == expectedRight) + || (leftVertex == expectedRight && rightVertex == expectedLeft) + +nonZeroDiagonalEntries :: + (Eq a, AdditiveGroup a) => + [a] -> + [(Int, Int, a)] +nonZeroDiagonalEntries entries = + mapMaybe + ( \(entryIndex, entryValue) -> + nonZeroEntry + entryIndex + entryIndex + entryValue + ) + (zip [0 ..] entries) + +tridiagonalEntries :: + (Eq a, AdditiveGroup a) => + [a] -> + [a] -> + [(Int, Int, a)] +tridiagonalEntries diagonalEntries offDiagonalEntries = + concat + ( zipWith + tridiagonalRowEntries + [0 ..] + (zip3 lowerEntries diagonalEntries upperEntries) + ) + where + lowerEntries = Nothing : (Just <$> offDiagonalEntries) + upperEntries = (Just <$> offDiagonalEntries) <> [Nothing] + +tridiagonalRowEntries :: + (Eq a, AdditiveGroup a) => + Int -> + (Maybe a, a, Maybe a) -> + [(Int, Int, a)] +tridiagonalRowEntries rowIndex (lowerValue, diagonalValue, upperValue) = + catMaybes + [ lowerValue >>= nonZeroEntry rowIndex (rowIndex - 1), + nonZeroEntry rowIndex rowIndex diagonalValue, + upperValue >>= nonZeroEntry rowIndex (rowIndex + 1) + ] + +nonZeroEntry :: + (Eq a, AdditiveGroup a) => + Int -> + Int -> + a -> + Maybe (Int, Int, a) +nonZeroEntry rowIndex columnIndex entryValue = + if entryValue == zero + then Nothing + else Just (rowIndex, columnIndex, entryValue) + +buildVertexIndices :: + (Ord vertex, Show vertex) => + [vertex] -> + Either MoonlightError (VertexIndex vertex) +buildVertexIndices vertexOrder + | isStrictlyAscending vertexOrder = + Right (AscendingVertexIndex (Box.fromList vertexOrder)) + | otherwise = + MapVertexIndex <$> foldMIndexed insertVertex Map.empty vertexOrder + +isStrictlyAscending :: Ord vertex => [vertex] -> Bool +isStrictlyAscending [] = True +isStrictlyAscending (vertexValue : remainingVertices) = + go vertexValue remainingVertices + where + go :: Ord vertex => vertex -> [vertex] -> Bool + go _ [] = True + go previousVertex (currentVertex : rest) + | previousVertex < currentVertex = go currentVertex rest + | otherwise = False + +insertVertex :: + (Ord vertex, Show vertex) => + Int -> + Map vertex Int -> + vertex -> + Either MoonlightError (Map vertex Int) +insertVertex vertexIndex vertexIndices vertexValue = + case Map.lookup vertexValue vertexIndices of + Just originalIndex -> + Left + ( InvariantViolation + ( "graph Laplacian vertex order contains duplicate vertex " + <> show vertexValue + <> " at indices " + <> show originalIndex + <> " and " + <> show vertexIndex + ) + ) + Nothing -> + Right (Map.insert vertexValue vertexIndex vertexIndices) + +foldMIndexed :: + Monad monadValue => + (Int -> state -> item -> monadValue state) -> + state -> + [item] -> + monadValue state +foldMIndexed step initialState items = + snd + <$> foldM + ( \(itemIndex, stateValue) itemValue -> + (\nextState -> (itemIndex + 1, nextState)) + <$> step itemIndex stateValue itemValue + ) + (0, initialState) + items + +canonicalGraphEdge :: + (Ord vertex, Show vertex) => + VertexIndex vertex -> + GraphEdge vertex -> + Either MoonlightError (Maybe IndexedGraphEdge) +canonicalGraphEdge vertexIndices edgeValue + | not (fieldValueValid weightValue) = + Left + ( InvariantViolation + ( "graph Laplacian edge weight must be finite, received " + <> show weightValue + ) + ) + | weightValue < 0.0 = + Left + ( InvariantViolation + ( "graph Laplacian edge weight must be non-negative, received " + <> show weightValue + ) + ) + | leftVertex == rightVertex = + Left + ( InvariantViolation + ( "graph Laplacian does not admit self-loop at vertex " + <> show leftVertex + ) + ) + | otherwise = do + leftIndex <- requireVertexIndex "left" leftVertex vertexIndices + rightIndex <- requireVertexIndex "right" rightVertex vertexIndices + if weightValue == 0.0 + then Right Nothing + else + Right + ( Just + IndexedGraphEdge + { indexedGraphEdgeLeft = min leftIndex rightIndex, + indexedGraphEdgeRight = max leftIndex rightIndex, + indexedGraphEdgeWeight = weightValue + } + ) + where + leftVertex = graphEdgeLeft edgeValue + rightVertex = graphEdgeRight edgeValue + weightValue = graphEdgeWeight edgeValue + +requireVertexIndex :: + (Ord vertex, Show vertex) => + String -> + vertex -> + VertexIndex vertex -> + Either MoonlightError Int +requireVertexIndex endpointRole vertexValue vertexIndices = + case lookupVertexIndex vertexValue vertexIndices of + Nothing -> + Left + ( InvariantViolation + ( "graph Laplacian " + <> endpointRole + <> " endpoint is absent from the explicit vertex order: " + <> show vertexValue + ) + ) + Just vertexIndex -> Right vertexIndex + +lookupVertexIndex :: Ord vertex => vertex -> VertexIndex vertex -> Maybe Int +lookupVertexIndex vertexValue vertexIndex = + case vertexIndex of + AscendingVertexIndex vertices -> + lookupAscendingVertex vertices vertexValue + MapVertexIndex vertexIndices -> + Map.lookup vertexValue vertexIndices + +lookupAscendingVertex :: Ord vertex => Box.Vector vertex -> vertex -> Maybe Int +lookupAscendingVertex vertices vertexValue = + go 0 (Box.length vertices - 1) + where + go !lowerBound !upperBound + | lowerBound > upperBound = Nothing + | otherwise = + let !midpoint = lowerBound + ((upperBound - lowerBound) `div` 2) + !midpointVertex = vertices `Box.unsafeIndex` midpoint + in case compare vertexValue midpointVertex of + LT -> go lowerBound (midpoint - 1) + EQ -> Just midpoint + GT -> go (midpoint + 1) upperBound + +indexedGraphEdgesByLeft :: + (Ord vertex, Show vertex) => + VertexIndex vertex -> + Int -> + [GraphEdge vertex] -> + Either MoonlightError IndexedGraphEdgeRows +indexedGraphEdgesByLeft vertexIndices dimension graphEdges = + runST $ do + let !edgeCapacity = length graphEdges + leftCounts <- MU.replicate dimension 0 + collectedLefts <- MU.unsafeNew edgeCapacity + collectedRights <- MU.unsafeNew edgeCapacity + collectedWeights <- MU.unsafeNew edgeCapacity + collectionResult <- + collectIndexedGraphEdges + vertexIndices + leftCounts + collectedLefts + collectedRights + collectedWeights + 0 + (-1) + (-1) + True + graphEdges + case collectionResult of + Left err -> pure (Left err) + Right collectionValue + | collectedGraphEdgesAreSorted collectionValue -> + Right + <$> compactSortedCollectedGraphEdges + dimension + collectedLefts + collectedRights + collectedWeights + (collectedGraphEdgeCount collectionValue) + | otherwise -> do + let !collectedCount = collectedGraphEdgeCount collectionValue + leftOffsets <- MU.replicate (dimension + 1) 0 + prefixMutableIntCountsWithStarts dimension leftCounts leftOffsets + scatteredRights <- MU.unsafeNew collectedCount + scatteredWeights <- MU.unsafeNew collectedCount + scatterCollectedGraphEdges + leftCounts + collectedLefts + collectedRights + collectedWeights + scatteredRights + scatteredWeights + 0 + collectedCount + rawOffsets <- U.unsafeFreeze leftOffsets + frozenRawRights <- U.unsafeFreeze scatteredRights + frozenRawWeights <- U.unsafeFreeze scatteredWeights + Right <$> compactGraphEdgeRows dimension rawOffsets frozenRawRights frozenRawWeights collectedCount +{-# INLINE indexedGraphEdgesByLeft #-} + +collectIndexedGraphEdges :: + (Ord vertex, Show vertex) => + VertexIndex vertex -> + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + Int -> + Int -> + Bool -> + [GraphEdge vertex] -> + ST s (Either MoonlightError CollectedGraphEdges) +collectIndexedGraphEdges _ _ _ _ _ !collectedCount _ _ !edgesAreSorted [] = + pure + ( Right + CollectedGraphEdges + { collectedGraphEdgeCount = collectedCount, + collectedGraphEdgesAreSorted = edgesAreSorted + } + ) +collectIndexedGraphEdges vertexIndices leftCounts collectedLefts collectedRights collectedWeights !collectedCount !previousLeft !previousRight !edgesAreSorted (edgeValue : remainingEdges) = + case canonicalGraphEdge vertexIndices edgeValue of + Left err -> pure (Left err) + Right Nothing -> + collectIndexedGraphEdges + vertexIndices + leftCounts + collectedLefts + collectedRights + collectedWeights + collectedCount + previousLeft + previousRight + edgesAreSorted + remainingEdges + Right (Just indexedEdge) -> do + let !leftIndex = indexedGraphEdgeLeft indexedEdge + !rightIndex = indexedGraphEdgeRight indexedEdge + !nextEdgesAreSorted = + edgesAreSorted + && ( previousLeft < 0 + || previousLeft < leftIndex + || (previousLeft == leftIndex && previousRight <= rightIndex) + ) + incrementMutableInt leftCounts leftIndex + MU.unsafeWrite collectedLefts collectedCount leftIndex + MU.unsafeWrite collectedRights collectedCount rightIndex + MU.unsafeWrite collectedWeights collectedCount (indexedGraphEdgeWeight indexedEdge) + collectIndexedGraphEdges + vertexIndices + leftCounts + collectedLefts + collectedRights + collectedWeights + (collectedCount + 1) + leftIndex + rightIndex + nextEdgesAreSorted + remainingEdges + +compactSortedCollectedGraphEdges :: + forall s. + Int -> + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + ST s IndexedGraphEdgeRows +compactSortedCollectedGraphEdges dimension collectedLefts collectedRights collectedWeights collectedCount = do + compactCounts <- MU.replicate dimension 0 + compactRights <- MU.unsafeNew collectedCount + compactWeights <- MU.unsafeNew collectedCount + compactCount <- + combineSortedCollectedGraphEdges + compactCounts + compactRights + compactWeights + 0 + collectedCount + 0 + compactOffsets <- MU.replicate (dimension + 1) 0 + prefixMutableIntCountsWithStarts dimension compactCounts compactOffsets + frozenOffsets <- U.unsafeFreeze compactOffsets + frozenRights <- U.unsafeFreeze compactRights + frozenWeights <- U.unsafeFreeze compactWeights + pure + IndexedGraphEdgeRows + { indexedGraphEdgeOffsets = frozenOffsets, + indexedGraphEdgeRights = U.slice 0 compactCount frozenRights, + indexedGraphEdgeWeights = U.slice 0 compactCount frozenWeights + } + where + combineSortedCollectedGraphEdges :: + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + Int -> + Int -> + ST s Int + combineSortedCollectedGraphEdges compactCounts compactRights compactWeights !entryIndex !entryStop !compactIndex + | entryIndex >= entryStop = pure compactIndex + | otherwise = do + leftIndex <- MU.unsafeRead collectedLefts entryIndex + rightIndex <- MU.unsafeRead collectedRights entryIndex + weightValue <- MU.unsafeRead collectedWeights entryIndex + combineSortedCollectedGraphEdge + compactCounts + compactRights + compactWeights + leftIndex + rightIndex + weightValue + (entryIndex + 1) + entryStop + compactIndex + + combineSortedCollectedGraphEdge :: + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + Int -> + Double -> + Int -> + Int -> + Int -> + ST s Int + combineSortedCollectedGraphEdge compactCounts compactRights compactWeights !leftIndex !rightIndex !weightValue !entryIndex !entryStop !compactIndex + | entryIndex >= entryStop = + writeSortedCollectedGraphEdge compactCounts compactRights compactWeights leftIndex rightIndex weightValue entryIndex entryStop compactIndex + | otherwise = do + nextLeft <- MU.unsafeRead collectedLefts entryIndex + nextRight <- MU.unsafeRead collectedRights entryIndex + if nextLeft == leftIndex && nextRight == rightIndex + then do + nextWeight <- MU.unsafeRead collectedWeights entryIndex + combineSortedCollectedGraphEdge + compactCounts + compactRights + compactWeights + leftIndex + rightIndex + (nextWeight + weightValue) + (entryIndex + 1) + entryStop + compactIndex + else + writeSortedCollectedGraphEdge compactCounts compactRights compactWeights leftIndex rightIndex weightValue entryIndex entryStop compactIndex + + writeSortedCollectedGraphEdge :: + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + Int -> + Double -> + Int -> + Int -> + Int -> + ST s Int + writeSortedCollectedGraphEdge compactCounts compactRights compactWeights !leftIndex !rightIndex !weightValue !nextEntryIndex !entryStop !compactIndex + | weightValue == 0.0 = + combineSortedCollectedGraphEdges compactCounts compactRights compactWeights nextEntryIndex entryStop compactIndex + | otherwise = do + incrementMutableInt compactCounts leftIndex + MU.unsafeWrite compactRights compactIndex rightIndex + MU.unsafeWrite compactWeights compactIndex weightValue + combineSortedCollectedGraphEdges compactCounts compactRights compactWeights nextEntryIndex entryStop (compactIndex + 1) + +scatterCollectedGraphEdges :: + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + Int -> + ST s () +scatterCollectedGraphEdges nextOffsets collectedLefts collectedRights collectedWeights scatteredRights scatteredWeights !entryIndex !entryStop + | entryIndex >= entryStop = pure () + | otherwise = do + leftIndex <- MU.unsafeRead collectedLefts entryIndex + targetIndex <- MU.unsafeRead nextOffsets leftIndex + rightIndex <- MU.unsafeRead collectedRights entryIndex + weightValue <- MU.unsafeRead collectedWeights entryIndex + MU.unsafeWrite scatteredRights targetIndex rightIndex + MU.unsafeWrite scatteredWeights targetIndex weightValue + MU.unsafeWrite nextOffsets leftIndex (targetIndex + 1) + scatterCollectedGraphEdges + nextOffsets + collectedLefts + collectedRights + collectedWeights + scatteredRights + scatteredWeights + (entryIndex + 1) + entryStop + +compactGraphEdgeRows :: + Int -> + U.Vector Int -> + U.Vector Int -> + U.Vector Double -> + Int -> + ST s IndexedGraphEdgeRows +compactGraphEdgeRows dimension rawOffsets rawRights rawWeights edgeCount = do + compactOffsets <- MU.replicate (dimension + 1) 0 + compactRights <- MU.unsafeNew edgeCount + compactWeights <- MU.unsafeNew edgeCount + finalCount <- compactLeftRows compactOffsets compactRights compactWeights 0 0 + frozenOffsets <- U.unsafeFreeze compactOffsets + frozenRights <- U.unsafeFreeze compactRights + frozenWeights <- U.unsafeFreeze compactWeights + pure + IndexedGraphEdgeRows + { indexedGraphEdgeOffsets = frozenOffsets, + indexedGraphEdgeRights = U.slice 0 finalCount frozenRights, + indexedGraphEdgeWeights = U.slice 0 finalCount frozenWeights + } + where + compactLeftRows :: + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + Int -> + ST s Int + compactLeftRows compactOffsets compactRights compactWeights !leftIndex !compactCount + | leftIndex >= dimension = do + MU.unsafeWrite compactOffsets dimension compactCount + pure compactCount + | otherwise = do + MU.unsafeWrite compactOffsets leftIndex compactCount + let !entryStart = rawOffsets `U.unsafeIndex` leftIndex + !entryStop = rawOffsets `U.unsafeIndex` (leftIndex + 1) + !orderedPairs = + sortBy + (comparing (\(rightIndex, weightValue) -> (rightIndex, weightValue))) + (collectGraphEdgePairs entryStart entryStop []) + compactStop <- writeCombinedGraphPairs compactRights compactWeights compactCount orderedPairs + compactLeftRows compactOffsets compactRights compactWeights (leftIndex + 1) compactStop + + collectGraphEdgePairs :: Int -> Int -> [(Int, Double)] -> [(Int, Double)] + collectGraphEdgePairs !entryIndex !entryStop rowPairs + | entryIndex >= entryStop = rowPairs + | otherwise = + collectGraphEdgePairs + (entryIndex + 1) + entryStop + ( ( rawRights `U.unsafeIndex` entryIndex, + rawWeights `U.unsafeIndex` entryIndex + ) + : rowPairs + ) + + writeCombinedGraphPairs :: + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + [(Int, Double)] -> + ST s Int + writeCombinedGraphPairs _ _ !compactIndex [] = + pure compactIndex + writeCombinedGraphPairs compactRights compactWeights !compactIndex ((rightIndex, weightValue) : rowPairs) = + writeCombinedGraphPair compactRights compactWeights compactIndex rightIndex weightValue rowPairs + + writeCombinedGraphPair :: + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + Int -> + Double -> + [(Int, Double)] -> + ST s Int + writeCombinedGraphPair compactRights compactWeights !compactIndex !rightIndex !weightValue [] = + writeNonZeroGraphPair compactRights compactWeights compactIndex rightIndex weightValue [] + writeCombinedGraphPair compactRights compactWeights !compactIndex !rightIndex !weightValue ((nextRight, nextWeight) : rowPairs) + | nextRight == rightIndex = + writeCombinedGraphPair compactRights compactWeights compactIndex rightIndex (nextWeight + weightValue) rowPairs + | otherwise = + writeNonZeroGraphPair compactRights compactWeights compactIndex rightIndex weightValue ((nextRight, nextWeight) : rowPairs) + + writeNonZeroGraphPair :: + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + Int -> + Double -> + [(Int, Double)] -> + ST s Int + writeNonZeroGraphPair compactRights compactWeights !compactIndex !rightIndex !weightValue rowPairs + | weightValue == 0.0 = + writeCombinedGraphPairs compactRights compactWeights compactIndex rowPairs + | otherwise = do + MU.unsafeWrite compactRights compactIndex rightIndex + MU.unsafeWrite compactWeights compactIndex weightValue + writeCombinedGraphPairs compactRights compactWeights (compactIndex + 1) rowPairs + +graphLaplacianCSRFromIndexedEdges :: Int -> IndexedGraphEdgeRows -> SparseCSR Double +graphLaplacianCSRFromIndexedEdges dimension indexedEdges = + runST $ do + lowerCounts <- MU.replicate dimension 0 + upperCounts <- MU.replicate dimension 0 + degrees <- MU.replicate dimension 0.0 + accumulateGraphEdgeRows lowerCounts upperCounts degrees indexedEdges 0 + rowOffsets <- MU.replicate (dimension + 1) 0 + prefixGraphRowOffsets lowerCounts upperCounts degrees rowOffsets 0 0 + finalCount <- MU.unsafeRead rowOffsets dimension + columnIndices <- MU.unsafeNew finalCount + values <- MU.unsafeNew finalCount + lowerNext <- MU.unsafeNew dimension + upperNext <- MU.unsafeNew dimension + initializeGraphRows lowerCounts degrees rowOffsets lowerNext upperNext columnIndices values 0 + writeGraphEdgeRows lowerNext upperNext columnIndices values indexedEdges 0 + frozenOffsets <- U.unsafeFreeze rowOffsets + frozenColumns <- U.unsafeFreeze columnIndices + frozenValues <- U.unsafeFreeze values + pure + ( csrFromCanonicalVectorsUnchecked + dimension + dimension + frozenOffsets + frozenColumns + frozenValues + ) +{-# INLINE graphLaplacianCSRFromIndexedEdges #-} + +pathGraphLaplacianCSRFromIndexedEdges :: Int -> IndexedGraphEdgeRows -> Maybe (SparseCSR Double) +pathGraphLaplacianCSRFromIndexedEdges dimension edgeRows + | dimension <= 1 = + if U.null edgeWeights + then Just (emptySquareCSR dimension) + else Nothing + | U.length edgeWeights /= dimension - 1 = + Nothing + | pathEdgesMatch 0 = + Just (pathGraphLaplacianCSRFromWeights dimension edgeWeights) + | otherwise = + Nothing + where + !edgeOffsets = indexedGraphEdgeOffsets edgeRows + !edgeRights = indexedGraphEdgeRights edgeRows + !edgeWeights = indexedGraphEdgeWeights edgeRows + + pathEdgesMatch !leftIndex + | leftIndex >= dimension - 1 = + edgeOffsets `U.unsafeIndex` dimension == U.length edgeWeights + | otherwise = + let !entryStart = edgeOffsets `U.unsafeIndex` leftIndex + !entryStop = edgeOffsets `U.unsafeIndex` (leftIndex + 1) + in entryStop - entryStart == 1 + && edgeRights `U.unsafeIndex` entryStart == leftIndex + 1 + && pathEdgesMatch (leftIndex + 1) + +emptySquareCSR :: Int -> SparseCSR Double +emptySquareCSR dimension = + csrFromCanonicalVectorsWithPlanUnchecked + dimension + dimension + (U.replicate (dimension + 1) 0) + U.empty + U.empty + CSRGeneral + +pathGraphLaplacianCSRFromWeights :: Int -> U.Vector Double -> SparseCSR Double +pathGraphLaplacianCSRFromWeights dimension edgeWeights = + runST $ do + let !entryCount = 3 * dimension - 2 + rowOffsets <- MU.unsafeNew (dimension + 1) + columnIndices <- MU.unsafeNew entryCount + values <- MU.unsafeNew entryCount + writePathGraphRows rowOffsets columnIndices values 0 0 + frozenOffsets <- U.unsafeFreeze rowOffsets + frozenColumns <- U.unsafeFreeze columnIndices + frozenValues <- U.unsafeFreeze values + pure + ( csrFromCanonicalVectorsWithPlanUnchecked + dimension + dimension + frozenOffsets + frozenColumns + frozenValues + (CSRContiguousBand 1 1) + ) + where + writePathGraphRows :: + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + Int -> + ST s () + writePathGraphRows rowOffsets columnIndices values !rowIndex !entryIndex + | rowIndex >= dimension = do + MU.unsafeWrite rowOffsets dimension entryIndex + | rowIndex == 0 = do + let !rightWeight = edgeWeights `U.unsafeIndex` 0 + MU.unsafeWrite rowOffsets rowIndex entryIndex + MU.unsafeWrite columnIndices entryIndex rowIndex + MU.unsafeWrite values entryIndex rightWeight + MU.unsafeWrite columnIndices (entryIndex + 1) (rowIndex + 1) + MU.unsafeWrite values (entryIndex + 1) (negate rightWeight) + writePathGraphRows rowOffsets columnIndices values (rowIndex + 1) (entryIndex + 2) + | rowIndex == dimension - 1 = do + let !leftWeight = edgeWeights `U.unsafeIndex` (rowIndex - 1) + MU.unsafeWrite rowOffsets rowIndex entryIndex + MU.unsafeWrite columnIndices entryIndex (rowIndex - 1) + MU.unsafeWrite values entryIndex (negate leftWeight) + MU.unsafeWrite columnIndices (entryIndex + 1) rowIndex + MU.unsafeWrite values (entryIndex + 1) leftWeight + writePathGraphRows rowOffsets columnIndices values (rowIndex + 1) (entryIndex + 2) + | otherwise = do + let !leftWeight = edgeWeights `U.unsafeIndex` (rowIndex - 1) + !rightWeight = edgeWeights `U.unsafeIndex` rowIndex + MU.unsafeWrite rowOffsets rowIndex entryIndex + MU.unsafeWrite columnIndices entryIndex (rowIndex - 1) + MU.unsafeWrite values entryIndex (negate leftWeight) + MU.unsafeWrite columnIndices (entryIndex + 1) rowIndex + MU.unsafeWrite values (entryIndex + 1) (leftWeight + rightWeight) + MU.unsafeWrite columnIndices (entryIndex + 2) (rowIndex + 1) + MU.unsafeWrite values (entryIndex + 2) (negate rightWeight) + writePathGraphRows rowOffsets columnIndices values (rowIndex + 1) (entryIndex + 3) + +accumulateGraphEdgeRows :: + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + IndexedGraphEdgeRows -> + Int -> + ST s () +accumulateGraphEdgeRows lowerCounts upperCounts degrees edgeRows !leftIndex + | leftIndex >= U.length edgeOffsets - 1 = pure () + | otherwise = do + let !entryStart = edgeOffsets `U.unsafeIndex` leftIndex + !entryStop = edgeOffsets `U.unsafeIndex` (leftIndex + 1) + accumulateGraphEdgeSpan lowerCounts upperCounts degrees leftIndex entryStart entryStop + accumulateGraphEdgeRows lowerCounts upperCounts degrees edgeRows (leftIndex + 1) + where + !edgeOffsets = indexedGraphEdgeOffsets edgeRows + !edgeRights = indexedGraphEdgeRights edgeRows + !edgeWeights = indexedGraphEdgeWeights edgeRows + + accumulateGraphEdgeSpan :: + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + Int -> + Int -> + ST s () + accumulateGraphEdgeSpan lowerCountSlots upperCountSlots degreeSlots !rowIndex !entryIndex !entryStop + | entryIndex >= entryStop = pure () + | otherwise = do + let !rightIndex = edgeRights `U.unsafeIndex` entryIndex + !weightValue = edgeWeights `U.unsafeIndex` entryIndex + incrementMutableInt upperCountSlots rowIndex + incrementMutableInt lowerCountSlots rightIndex + addMutableEntry degreeSlots rowIndex weightValue + addMutableEntry degreeSlots rightIndex weightValue + accumulateGraphEdgeSpan lowerCountSlots upperCountSlots degreeSlots rowIndex (entryIndex + 1) entryStop + +incrementMutableInt :: MU.MVector s Int -> Int -> ST s () +incrementMutableInt values !entryIndex = do + currentValue <- MU.unsafeRead values entryIndex + MU.unsafeWrite values entryIndex (currentValue + 1) +{-# INLINE incrementMutableInt #-} + +prefixMutableIntCountsWithStarts :: + Int -> + MU.MVector s Int -> + MU.MVector s Int -> + ST s () +prefixMutableIntCountsWithStarts axisCount counts offsets = + go 0 0 + where + go !axisIndex !runningTotal + | axisIndex >= axisCount = + MU.unsafeWrite offsets axisCount runningTotal + | otherwise = do + axisCountValue <- MU.unsafeRead counts axisIndex + MU.unsafeWrite offsets axisIndex runningTotal + MU.unsafeWrite counts axisIndex runningTotal + go (axisIndex + 1) (runningTotal + axisCountValue) + +prefixGraphRowOffsets :: + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + MU.MVector s Int -> + Int -> + Int -> + ST s () +prefixGraphRowOffsets lowerCounts upperCounts degrees rowOffsets !rowIndex !runningTotal + | rowIndex >= MU.length degrees = + MU.unsafeWrite rowOffsets rowIndex runningTotal + | otherwise = do + lowerCount <- MU.unsafeRead lowerCounts rowIndex + upperCount <- MU.unsafeRead upperCounts rowIndex + degreeValue <- MU.unsafeRead degrees rowIndex + MU.unsafeWrite rowOffsets rowIndex runningTotal + let !diagonalCount = + if degreeValue == 0.0 + then 0 + else 1 + prefixGraphRowOffsets + lowerCounts + upperCounts + degrees + rowOffsets + (rowIndex + 1) + (runningTotal + lowerCount + diagonalCount + upperCount) + +initializeGraphRows :: + MU.MVector s Int -> + MU.MVector s Double -> + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + ST s () +initializeGraphRows lowerCounts degrees rowOffsets lowerNext upperNext columnIndices values !rowIndex + | rowIndex >= MU.length degrees = pure () + | otherwise = do + rowStart <- MU.unsafeRead rowOffsets rowIndex + lowerCount <- MU.unsafeRead lowerCounts rowIndex + degreeValue <- MU.unsafeRead degrees rowIndex + let !diagonalIndex = rowStart + lowerCount + !upperStart = + if degreeValue == 0.0 + then diagonalIndex + else diagonalIndex + 1 + MU.unsafeWrite lowerNext rowIndex rowStart + MU.unsafeWrite upperNext rowIndex upperStart + if degreeValue == 0.0 + then pure () + else do + MU.unsafeWrite columnIndices diagonalIndex rowIndex + MU.unsafeWrite values diagonalIndex degreeValue + initializeGraphRows + lowerCounts + degrees + rowOffsets + lowerNext + upperNext + columnIndices + values + (rowIndex + 1) + +writeGraphEdgeRows :: + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + IndexedGraphEdgeRows -> + Int -> + ST s () +writeGraphEdgeRows lowerNext upperNext columnIndices values edgeRows !leftIndex + | leftIndex >= U.length edgeOffsets - 1 = pure () + | otherwise = do + let !entryStart = edgeOffsets `U.unsafeIndex` leftIndex + !entryStop = edgeOffsets `U.unsafeIndex` (leftIndex + 1) + writeGraphEdgeSpan lowerNext upperNext columnIndices values leftIndex entryStart entryStop + writeGraphEdgeRows lowerNext upperNext columnIndices values edgeRows (leftIndex + 1) + where + !edgeOffsets = indexedGraphEdgeOffsets edgeRows + !edgeRights = indexedGraphEdgeRights edgeRows + !edgeWeights = indexedGraphEdgeWeights edgeRows + + writeGraphEdgeSpan :: + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s Double -> + Int -> + Int -> + Int -> + ST s () + writeGraphEdgeSpan lowerSlots upperSlots columnSlots valueSlots !rowIndex !entryIndex !entryStop + | entryIndex >= entryStop = pure () + | otherwise = do + let !rightIndex = edgeRights `U.unsafeIndex` entryIndex + !weightValue = edgeWeights `U.unsafeIndex` entryIndex + upperIndex <- MU.unsafeRead upperSlots rowIndex + MU.unsafeWrite columnSlots upperIndex rightIndex + MU.unsafeWrite valueSlots upperIndex (negate weightValue) + MU.unsafeWrite upperSlots rowIndex (upperIndex + 1) + lowerIndex <- MU.unsafeRead lowerSlots rightIndex + MU.unsafeWrite columnSlots lowerIndex rowIndex + MU.unsafeWrite valueSlots lowerIndex (negate weightValue) + MU.unsafeWrite lowerSlots rightIndex (lowerIndex + 1) + writeGraphEdgeSpan lowerSlots upperSlots columnSlots valueSlots rowIndex (entryIndex + 1) entryStop diff --git a/src-sparse/Moonlight/LinAlg/Pure/Sparse/Types.hs b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Types.hs new file mode 100644 --- /dev/null +++ b/src-sparse/Moonlight/LinAlg/Pure/Sparse/Types.hs @@ -0,0 +1,1080 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Moonlight.LinAlg.Pure.Sparse.Types + ( SparseCOO, + mkSparseCOO, + cooRows, + cooCols, + cooEntries, + SparseCSR, + mkSparseCSR, + csrRows, + csrCols, + csrRowOffsetsVector, + csrColumnIndicesVector, + csrValuesVector, + CSRExecutionPlan (..), + csrExecutionPlan, + SparseCSC, + mkSparseCSC, + cscRows, + cscCols, + cscColumnOffsetsVector, + cscRowIndicesVector, + cscValuesVector, + denseToCOO, + denseToCSR, + denseToCSC, + cooToCSR, + cooToCSC, + canonicalCSRFromValidEntriesUnchecked, + csrFromCanonicalVectorsUnchecked, + csrFromCanonicalVectorsWithPlanUnchecked, + csrToCOO, + cscToCOO, + cooToDense, + csrToDense, + cscToDense, + csrToCSC, + cscToCSR, + csrMatVecVector, + validateCOO, + validateCOOEntries, + validateCSR, + validateCSC, + ) +where + +import Control.Monad.ST (ST, runST) +import Data.Bifunctor (first) +import Data.Foldable (traverse_) +import Data.Kind (Type) +import Data.List (sortBy) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Maybe (mapMaybe) +import Data.Ord (comparing) +import Data.Proxy (Proxy (..)) +import Data.Vector.Unboxed qualified as U +import Data.Vector.Unboxed.Mutable qualified as MU +import GHC.TypeNats (KnownNat, natVal) +import Moonlight.Core + ( AdditiveGroup (..), + AdditiveMonoid (..), + MoonlightError (..), + checkedNaturalToInt, + checkedNonNegativeProduct, + ) +import Moonlight.LinAlg.Internal.VectorOps + ( csrContiguousBandMatVecValidatedU, + csrMatVecValidatedU, + ) +import Moonlight.LinAlg.Pure.Dense.Types (Matrix, fromListMatrix, toListMatrix) +import qualified Moonlight.LinAlg.Pure.Dense.Types as DenseTypes +import Prelude + +type SparseCOO :: Type -> Type +data SparseCOO a = SparseCOO + { cooRows :: Int, + cooCols :: Int, + cooEntries :: [(Int, Int, a)] + } + deriving stock (Eq, Show) + +type CSRExecutionPlan :: Type +data CSRExecutionPlan + = CSRGeneral + | CSRContiguousBand !Int !Int + deriving stock (Eq, Show) + +type SparseCSR :: Type -> Type +data SparseCSR a = SparseCSR + { csrRows :: Int, + csrCols :: Int, + csrRowOffsetsVector :: U.Vector Int, + csrColumnIndicesVector :: U.Vector Int, + csrValuesVector :: U.Vector a, + csrExecutionPlan :: !CSRExecutionPlan + } + +instance (Eq a, U.Unbox a) => Eq (SparseCSR a) where + left == right = + csrRows left == csrRows right + && csrCols left == csrCols right + && csrRowOffsetsVector left == csrRowOffsetsVector right + && csrColumnIndicesVector left == csrColumnIndicesVector right + && csrValuesVector left == csrValuesVector right + +instance (Show a, U.Unbox a) => Show (SparseCSR a) where + showsPrec precedence csrValue = + showParen (precedence > 10) $ + showString "SparseCSR {csrRows = " + . shows (csrRows csrValue) + . showString ", csrCols = " + . shows (csrCols csrValue) + . showString ", csrRowOffsets = " + . shows (U.toList (csrRowOffsetsVector csrValue)) + . showString ", csrColumnIndices = " + . shows (U.toList (csrColumnIndicesVector csrValue)) + . showString ", csrValues = " + . shows (U.toList (csrValuesVector csrValue)) + . showString "}" + +type SparseCSC :: Type -> Type +data SparseCSC a = SparseCSC + { cscRows :: Int, + cscCols :: Int, + cscColumnOffsetsVector :: U.Vector Int, + cscRowIndicesVector :: U.Vector Int, + cscValuesVector :: U.Vector a + } + +instance (Eq a, U.Unbox a) => Eq (SparseCSC a) where + left == right = + cscRows left == cscRows right + && cscCols left == cscCols right + && cscColumnOffsetsVector left == cscColumnOffsetsVector right + && cscRowIndicesVector left == cscRowIndicesVector right + && cscValuesVector left == cscValuesVector right + +instance (Show a, U.Unbox a) => Show (SparseCSC a) where + showsPrec precedence cscValue = + showParen (precedence > 10) $ + showString "SparseCSC {cscRows = " + . shows (cscRows cscValue) + . showString ", cscCols = " + . shows (cscCols cscValue) + . showString ", cscColumnOffsets = " + . shows (U.toList (cscColumnOffsetsVector cscValue)) + . showString ", cscRowIndices = " + . shows (U.toList (cscRowIndicesVector cscValue)) + . showString ", cscValues = " + . shows (U.toList (cscValuesVector cscValue)) + . showString "}" + +mkSparseCSR :: (Eq a, AdditiveMonoid a, U.Unbox a) => Int -> Int -> [Int] -> [Int] -> [a] -> Either MoonlightError (SparseCSR a) +mkSparseCSR rows cols offsets colIndices values = do + let !offsetVector = U.fromList offsets + !columnVector = U.fromList colIndices + !valueVector = U.fromList values + !unplanned = + SparseCSR + { csrRows = rows, + csrCols = cols, + csrRowOffsetsVector = offsetVector, + csrColumnIndicesVector = columnVector, + csrValuesVector = valueVector, + csrExecutionPlan = CSRGeneral + } + validateCSR unplanned + pure + unplanned + { csrExecutionPlan = + detectCSRExecutionPlan + rows + cols + offsetVector + columnVector + } + +mkSparseCOO :: Int -> Int -> [(Int, Int, a)] -> Either MoonlightError (SparseCOO a) +mkSparseCOO rows cols entries = do + validateCOOEntries rows cols entries + Right (SparseCOO rows cols entries) + +mkSparseCSC :: (Eq a, AdditiveMonoid a, U.Unbox a) => Int -> Int -> [Int] -> [Int] -> [a] -> Either MoonlightError (SparseCSC a) +mkSparseCSC rows cols offsets rowIndices values = do + let csc = + SparseCSC + rows + cols + (U.fromList offsets) + (U.fromList rowIndices) + (U.fromList values) + validateCSC csc + Right csc + +denseToCOO :: + forall r c a. + (KnownNat r, KnownNat c, Eq a, AdditiveGroup a) => + Matrix r c a -> + SparseCOO a +denseToCOO matrixValue = + let (rowCount, columnCount) = DenseTypes.matrixShape matrixValue + indexedValues = zip [0 ..] (toListMatrix matrixValue) + toEntry (flatIndex, value) = + if value == zero + then Nothing + else + let rowIndex = flatIndex `div` columnCount + columnIndex = flatIndex `mod` columnCount + in Just (rowIndex, columnIndex, value) + in SparseCOO rowCount columnCount (mapMaybe toEntry indexedValues) + +denseToCSR :: + forall r c a. + (KnownNat r, KnownNat c, Eq a, AdditiveGroup a, U.Unbox a) => + Matrix r c a -> + SparseCSR a +denseToCSR = cooToCSRSortedUniqueUnchecked . denseToCOO + +denseToCSC :: + forall r c a. + (KnownNat r, KnownNat c, Eq a, AdditiveGroup a, U.Unbox a) => + Matrix r c a -> + SparseCSC a +denseToCSC = cooToCSCSortedUniqueUnchecked . denseToCOO + +cooToCSR :: (Eq a, AdditiveGroup a, U.Unbox a) => SparseCOO a -> Either MoonlightError (SparseCSR a) +cooToCSR cooValue = + validateCOO cooValue *> pure (cooToCSRUnchecked cooValue) + +cooToCSRUnchecked :: (Eq a, AdditiveGroup a, U.Unbox a) => SparseCOO a -> SparseCSR a +cooToCSRUnchecked cooValue = + canonicalCSRFromValidEntriesUnchecked + (cooRows cooValue) + (cooCols cooValue) + (cooEntries cooValue) +{-# INLINE cooToCSRUnchecked #-} + +cooToCSRSortedUniqueUnchecked :: U.Unbox a => SparseCOO a -> SparseCSR a +cooToCSRSortedUniqueUnchecked cooValue = + csrFromSortedEntries (cooRows cooValue) (cooCols cooValue) (cooEntries cooValue) + +csrFromSortedEntries :: U.Unbox a => Int -> Int -> [(Int, Int, a)] -> SparseCSR a +csrFromSortedEntries rowCount columnCount orderedEntries = + let !offsetVector = + U.fromList + ( offsetsFromSortedAxes + rowCount + ((\(rowIndex, _, _) -> rowIndex) <$> orderedEntries) + ) + !columnVector = + U.fromList + ((\(_, columnIndex, _) -> columnIndex) <$> orderedEntries) + !valueVector = + U.fromList + ((\(_, _, value) -> value) <$> orderedEntries) + in csrFromCanonicalVectorsUnchecked rowCount columnCount offsetVector columnVector valueVector + +csrFromCanonicalVectorsUnchecked :: + Int -> + Int -> + U.Vector Int -> + U.Vector Int -> + U.Vector a -> + SparseCSR a +csrFromCanonicalVectorsUnchecked rowCount columnCount offsetVector columnVector valueVector = + csrFromCanonicalVectorsWithPlanUnchecked + rowCount + columnCount + offsetVector + columnVector + valueVector + ( detectCSRExecutionPlan + rowCount + columnCount + offsetVector + columnVector + ) + +csrFromCanonicalVectorsWithPlanUnchecked :: + Int -> + Int -> + U.Vector Int -> + U.Vector Int -> + U.Vector a -> + CSRExecutionPlan -> + SparseCSR a +csrFromCanonicalVectorsWithPlanUnchecked rowCount columnCount offsetVector columnVector valueVector executionPlan = + SparseCSR + { csrRows = rowCount, + csrCols = columnCount, + csrRowOffsetsVector = offsetVector, + csrColumnIndicesVector = columnVector, + csrValuesVector = valueVector, + csrExecutionPlan = executionPlan + } + +canonicalCSRFromValidEntriesUnchecked :: + (Eq a, AdditiveGroup a, U.Unbox a) => + Int -> + Int -> + [(Int, Int, a)] -> + SparseCSR a +canonicalCSRFromValidEntriesUnchecked rowCount columnCount entries = + let SparseEntryVectors offsetVector columnVector valueVector = + canonicalCSRVectorsFromValidEntries rowCount columnCount entries + in csrFromCanonicalVectorsUnchecked rowCount columnCount offsetVector columnVector valueVector +{-# INLINE canonicalCSRFromValidEntriesUnchecked #-} + +canonicalCSRVectorsFromValidEntries :: + (Eq a, AdditiveGroup a, U.Unbox a) => + Int -> + Int -> + [(Int, Int, a)] -> + SparseEntryVectors a +canonicalCSRVectorsFromValidEntries rowCount columnCount entries = + compactCompressedRows columnCount (compressEntriesByRow rowCount entries) + +compressEntriesByRow :: + U.Unbox a => + Int -> + [(Int, Int, a)] -> + SparseEntryVectors a +compressEntriesByRow rowCount entries = + runST $ do + let !entryCount = length entries + rowCounts <- MU.replicate rowCount 0 + traverse_ (\(rowIndex, _, _) -> incrementMutableInt rowCounts rowIndex) entries + rowOffsets <- MU.replicate (rowCount + 1) 0 + prefixMutableCountsWithStarts rowCount rowCounts rowOffsets + columnVector <- MU.unsafeNew entryCount + valueVector <- MU.unsafeNew entryCount + traverse_ (scatterCompressedEntry rowCounts columnVector valueVector) entries + frozenOffsets <- U.unsafeFreeze rowOffsets + frozenColumns <- U.unsafeFreeze columnVector + frozenValues <- U.unsafeFreeze valueVector + pure + SparseEntryVectors + { sparseEntryOffsets = frozenOffsets, + sparseEntryIndices = frozenColumns, + sparseEntryValues = frozenValues + } + +scatterCompressedEntry :: + U.Unbox a => + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s a -> + (Int, Int, a) -> + ST s () +scatterCompressedEntry nextOffsets columnVector valueVector (rowIndex, columnIndex, entryValue) = do + targetIndex <- MU.unsafeRead nextOffsets rowIndex + MU.unsafeWrite columnVector targetIndex columnIndex + MU.unsafeWrite valueVector targetIndex entryValue + MU.unsafeWrite nextOffsets rowIndex (targetIndex + 1) + +compactCompressedRows :: + forall a. + (Eq a, AdditiveGroup a, U.Unbox a) => + Int -> + SparseEntryVectors a -> + SparseEntryVectors a +compactCompressedRows columnCount compressedEntries = + runST $ do + markerSlots <- MU.replicate columnCount (-1) + uniqueColumns <- MU.unsafeNew entryCount + uniqueValues <- MU.unsafeNew entryCount + compactOffsets <- MU.replicate (rowCount + 1) 0 + compactColumns <- MU.unsafeNew entryCount + compactValues <- MU.unsafeNew entryCount + finalCount <- + compactRows + markerSlots + uniqueColumns + uniqueValues + compactOffsets + compactColumns + compactValues + 0 + 0 + 0 + frozenOffsets <- U.unsafeFreeze compactOffsets + frozenColumns <- U.unsafeFreeze compactColumns + frozenValues <- U.unsafeFreeze compactValues + pure + SparseEntryVectors + { sparseEntryOffsets = frozenOffsets, + sparseEntryIndices = U.slice 0 finalCount frozenColumns, + sparseEntryValues = U.slice 0 finalCount frozenValues + } + where + !rowOffsets = sparseEntryOffsets compressedEntries + !rawColumns = sparseEntryIndices compressedEntries + !rawValues = sparseEntryValues compressedEntries + !rowCount = U.length rowOffsets - 1 + !entryCount = U.length rawValues + + compactRows :: + forall s. + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s a -> + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s a -> + Int -> + Int -> + Int -> + ST s Int + compactRows + markerSlots + uniqueColumns + uniqueValues + compactOffsets + compactColumns + compactValues + !rowIndex + !uniqueCount + !compactCount + | rowIndex >= rowCount = do + MU.unsafeWrite compactOffsets rowCount compactCount + pure compactCount + | otherwise = do + MU.unsafeWrite compactOffsets rowIndex compactCount + let !entryStart = rowOffsets `U.unsafeIndex` rowIndex + !entryStop = rowOffsets `U.unsafeIndex` (rowIndex + 1) + uniqueStop <- + combineCompressedRow + markerSlots + uniqueColumns + uniqueValues + uniqueCount + uniqueCount + entryStart + entryStop + rowPairs <- collectNonZeroPairs uniqueColumns uniqueValues uniqueCount uniqueStop [] + compactStop <- + writeSortedPairs + compactColumns + compactValues + compactCount + (sortBy (comparing fst) rowPairs) + compactRows + markerSlots + uniqueColumns + uniqueValues + compactOffsets + compactColumns + compactValues + (rowIndex + 1) + uniqueStop + compactStop + + combineCompressedRow :: + forall s. + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s a -> + Int -> + Int -> + Int -> + Int -> + ST s Int + combineCompressedRow + markerSlots + uniqueColumns + uniqueValues + !uniqueStart + !uniqueNext + !entryIndex + !entryStop + | entryIndex >= entryStop = pure uniqueNext + | entryValue == zero = + combineCompressedRow + markerSlots + uniqueColumns + uniqueValues + uniqueStart + uniqueNext + (entryIndex + 1) + entryStop + | otherwise = do + markerSlot <- MU.unsafeRead markerSlots columnIndex + if markerSlot >= uniqueStart + then do + currentValue <- MU.unsafeRead uniqueValues markerSlot + MU.unsafeWrite uniqueValues markerSlot (add entryValue currentValue) + combineCompressedRow + markerSlots + uniqueColumns + uniqueValues + uniqueStart + uniqueNext + (entryIndex + 1) + entryStop + else do + MU.unsafeWrite markerSlots columnIndex uniqueNext + MU.unsafeWrite uniqueColumns uniqueNext columnIndex + MU.unsafeWrite uniqueValues uniqueNext entryValue + combineCompressedRow + markerSlots + uniqueColumns + uniqueValues + uniqueStart + (uniqueNext + 1) + (entryIndex + 1) + entryStop + where + !columnIndex = rawColumns `U.unsafeIndex` entryIndex + !entryValue = rawValues `U.unsafeIndex` entryIndex + + collectNonZeroPairs :: + forall s. + MU.MVector s Int -> + MU.MVector s a -> + Int -> + Int -> + [(Int, a)] -> + ST s [(Int, a)] + collectNonZeroPairs uniqueColumns uniqueValues !entryIndex !entryStop !rowPairs + | entryIndex >= entryStop = pure rowPairs + | otherwise = do + entryValue <- MU.unsafeRead uniqueValues entryIndex + if entryValue == zero + then collectNonZeroPairs uniqueColumns uniqueValues (entryIndex + 1) entryStop rowPairs + else do + columnIndex <- MU.unsafeRead uniqueColumns entryIndex + collectNonZeroPairs uniqueColumns uniqueValues (entryIndex + 1) entryStop ((columnIndex, entryValue) : rowPairs) + + writeSortedPairs :: + forall s. + MU.MVector s Int -> + MU.MVector s a -> + Int -> + [(Int, a)] -> + ST s Int + writeSortedPairs _ _ !entryIndex [] = + pure entryIndex + writeSortedPairs compactColumns compactValues !entryIndex ((columnIndex, entryValue) : rowPairs) = do + MU.unsafeWrite compactColumns entryIndex columnIndex + MU.unsafeWrite compactValues entryIndex entryValue + writeSortedPairs compactColumns compactValues (entryIndex + 1) rowPairs + +cooToCSC :: (Eq a, AdditiveGroup a, U.Unbox a) => SparseCOO a -> Either MoonlightError (SparseCSC a) +cooToCSC cooValue = + validateCOO cooValue *> pure (cooToCSCUnchecked cooValue) + +cooToCSCUnchecked :: (Eq a, AdditiveGroup a, U.Unbox a) => SparseCOO a -> SparseCSC a +cooToCSCUnchecked cooValue = + csrToCSCUnchecked (cooToCSRUnchecked cooValue) + +cooToCSCSortedUniqueUnchecked :: U.Unbox a => SparseCOO a -> SparseCSC a +cooToCSCSortedUniqueUnchecked cooValue = + csrToCSCUnchecked (cooToCSRSortedUniqueUnchecked cooValue) + +type AxisOffsetState :: Type +data AxisOffsetState = AxisOffsetState + { axisOffsetCurrent :: !Int, + axisOffsetEntryCount :: !Int, + axisOffsetsRev :: [Int] + } + +type SparseEntryVectors :: Type -> Type +data SparseEntryVectors a = SparseEntryVectors + { sparseEntryOffsets :: !(U.Vector Int), + sparseEntryIndices :: !(U.Vector Int), + sparseEntryValues :: !(U.Vector a) + } + +offsetsFromSortedAxes :: Int -> [Int] -> [Int] +offsetsFromSortedAxes axisCount = + axisOffsets + . closeOffsetAxes axisCount + . foldl' acceptAxisOffset initialAxisOffsetState + where + axisOffsets = + reverse . axisOffsetsRev + +initialAxisOffsetState :: AxisOffsetState +initialAxisOffsetState = + AxisOffsetState + { axisOffsetCurrent = 0, + axisOffsetEntryCount = 0, + axisOffsetsRev = [0] + } + +acceptAxisOffset :: AxisOffsetState -> Int -> AxisOffsetState +acceptAxisOffset stateValue axisIndex = + let closedState = closeOffsetAxes axisIndex stateValue + in closedState {axisOffsetEntryCount = axisOffsetEntryCount closedState + 1} + +closeOffsetAxes :: Int -> AxisOffsetState -> AxisOffsetState +closeOffsetAxes targetAxis stateValue + | axisOffsetCurrent stateValue >= targetAxis = stateValue + | otherwise = + let closedAxisCount = targetAxis - axisOffsetCurrent stateValue + in stateValue + { axisOffsetCurrent = targetAxis, + axisOffsetsRev = + replicate closedAxisCount (axisOffsetEntryCount stateValue) + <> axisOffsetsRev stateValue + } + +incrementMutableInt :: MU.MVector s Int -> Int -> ST s () +incrementMutableInt values !entryIndex = do + currentValue <- MU.unsafeRead values entryIndex + MU.unsafeWrite values entryIndex (currentValue + 1) +{-# INLINE incrementMutableInt #-} + +prefixMutableCountsWithStarts :: + Int -> + MU.MVector s Int -> + MU.MVector s Int -> + ST s () +prefixMutableCountsWithStarts axisCount counts offsets = + go 0 0 + where + go !axisIndex !runningTotal + | axisIndex >= axisCount = + MU.unsafeWrite offsets axisCount runningTotal + | otherwise = do + axisCountValue <- MU.unsafeRead counts axisIndex + MU.unsafeWrite offsets axisIndex runningTotal + MU.unsafeWrite counts axisIndex runningTotal + go (axisIndex + 1) (runningTotal + axisCountValue) + +isMonotonicVector :: U.Vector Int -> Bool +isMonotonicVector values = + U.and (U.zipWith (<=) values (U.drop 1 values)) + +vectorEndpoints :: U.Vector Int -> Maybe (Int, Int) +vectorEndpoints values = + (,) <$> values U.!? 0 <*> values U.!? (U.length values - 1) + +detectCSRExecutionPlan :: + Int -> + Int -> + U.Vector Int -> + U.Vector Int -> + CSRExecutionPlan +detectCSRExecutionPlan rowCount columnCount rowOffsets columnIndices + | rowCount /= columnCount = CSRGeneral + | rowCount <= 0 = CSRContiguousBand 0 0 + | U.null columnIndices = CSRGeneral + | otherwise = + case discoverBandwidth 0 0 0 of + Nothing -> CSRGeneral + Just (!lowerBandwidth, !upperBandwidth) + | verifyRows lowerBandwidth upperBandwidth 0 -> + CSRContiguousBand lowerBandwidth upperBandwidth + | otherwise -> CSRGeneral + where + discoverBandwidth !rowIndex !lowerBandwidth !upperBandwidth + | rowIndex >= rowCount = + Just (lowerBandwidth, upperBandwidth) + | otherwise = + let !startIndex = rowOffsets `U.unsafeIndex` rowIndex + !stopIndex = rowOffsets `U.unsafeIndex` (rowIndex + 1) + in if startIndex >= stopIndex + then Nothing + else + let !firstColumn = + columnIndices `U.unsafeIndex` startIndex + !lastColumn = + columnIndices `U.unsafeIndex` (stopIndex - 1) + in discoverBandwidth + (rowIndex + 1) + (max lowerBandwidth (rowIndex - firstColumn)) + (max upperBandwidth (lastColumn - rowIndex)) + + verifyRows !lowerBandwidth !upperBandwidth !rowIndex + | rowIndex >= rowCount = True + | otherwise = + let !startIndex = rowOffsets `U.unsafeIndex` rowIndex + !stopIndex = rowOffsets `U.unsafeIndex` (rowIndex + 1) + !firstExpected = max 0 (rowIndex - lowerBandwidth) + !lastExpected = min (columnCount - 1) (rowIndex + upperBandwidth) + !expectedCount = lastExpected - firstExpected + 1 + in stopIndex - startIndex == expectedCount + && verifyColumns startIndex stopIndex firstExpected + && verifyRows lowerBandwidth upperBandwidth (rowIndex + 1) + + verifyColumns !entryIndex !stopIndex !expectedColumn + | entryIndex >= stopIndex = True + | columnIndices `U.unsafeIndex` entryIndex /= expectedColumn = False + | otherwise = + verifyColumns + (entryIndex + 1) + stopIndex + (expectedColumn + 1) +{-# INLINE detectCSRExecutionPlan #-} + +validateCOO :: SparseCOO a -> Either MoonlightError () +validateCOO cooValue = + validateCOOEntries (cooRows cooValue) (cooCols cooValue) (cooEntries cooValue) + +validateCOOEntries :: Int -> Int -> [(Int, Int, a)] -> Either MoonlightError () +validateCOOEntries rowCount columnCount entries + | rowCount < 0 || columnCount < 0 = + Left (InvariantViolation "COO dimensions must be non-negative") + | any invalidEntry entries = + Left (InvariantViolation "COO entry index out of bounds") + | otherwise = + Right () + where + invalidEntry (rowIndex, columnIndex, _) = + rowIndex < 0 + || rowIndex >= rowCount + || columnIndex < 0 + || columnIndex >= columnCount + +validateCSR :: (Eq a, AdditiveMonoid a, U.Unbox a) => SparseCSR a -> Either MoonlightError () +validateCSR csrValue = + let offsets = csrRowOffsetsVector csrValue + columnIndices = csrColumnIndicesVector csrValue + values = csrValuesVector csrValue + valueCount = U.length values + in case vectorEndpoints offsets of + Nothing -> Left (InvariantViolation "CSR row offsets must be non-empty") + Just (firstOffset, lastOffset) -> + validateCSRPayload csrValue offsets columnIndices values valueCount firstOffset lastOffset + +validateCSRPayload :: (Eq a, AdditiveMonoid a, U.Unbox a) => SparseCSR a -> U.Vector Int -> U.Vector Int -> U.Vector a -> Int -> Int -> Int -> Either MoonlightError () +validateCSRPayload csrValue offsets columnIndices values valueCount firstOffset lastOffset + | csrRows csrValue < 0 || csrCols csrValue < 0 = + Left (InvariantViolation "CSR dimensions must be non-negative") + | U.length offsets /= csrRows csrValue + 1 = + Left (InvariantViolation "CSR row-offset length must equal row count + 1") + | U.length columnIndices /= valueCount = + Left (InvariantViolation "CSR column-index and value payload lengths must match") + | U.any (< 0) offsets = + Left (InvariantViolation "CSR row offsets must be non-negative") + | not (isMonotonicVector offsets) = + Left (InvariantViolation "CSR row offsets must be monotonically non-decreasing") + | firstOffset /= 0 = + Left (InvariantViolation "CSR row offsets must begin at 0") + | lastOffset /= valueCount = + Left (InvariantViolation "CSR terminal row offset must equal value count") + | U.any (\columnIndex -> columnIndex < 0 || columnIndex >= csrCols csrValue) columnIndices = + Left (InvariantViolation "CSR column index out of bounds") + | not (csrRowsStrictlyCanonical csrValue offsets columnIndices) = + Left (InvariantViolation "CSR column indices must be strictly increasing within each row") + | U.any (== zero) values = + Left (InvariantViolation "CSR values must not store exact zeros") + | otherwise = Right () + +csrRowsStrictlyCanonical :: SparseCSR a -> U.Vector Int -> U.Vector Int -> Bool +csrRowsStrictlyCanonical csrValue offsets columnIndices = + U.all csrRowStrictlyCanonical (U.enumFromN 0 (csrRows csrValue)) + where + csrRowStrictlyCanonical rowIndex = + let startOffset = offsets `U.unsafeIndex` rowIndex + stopOffset = offsets `U.unsafeIndex` (rowIndex + 1) + rowColumns = U.slice startOffset (stopOffset - startOffset) columnIndices + in U.and (U.zipWith (<) rowColumns (U.drop 1 rowColumns)) + +csrToCOO :: U.Unbox a => SparseCSR a -> Either MoonlightError (SparseCOO a) +csrToCOO = + Right . csrToCOOUnchecked + +csrToCOOUnchecked :: U.Unbox a => SparseCSR a -> SparseCOO a +csrToCOOUnchecked csrValue = + SparseCOO + { cooRows = csrRows csrValue, + cooCols = csrCols csrValue, + cooEntries = + U.toList + ( U.zip3 + (offsetAxisIndicesVector (csrRows csrValue) (csrRowOffsetsVector csrValue)) + (csrColumnIndicesVector csrValue) + (csrValuesVector csrValue) + ) + } + +validateCSC :: (Eq a, AdditiveMonoid a, U.Unbox a) => SparseCSC a -> Either MoonlightError () +validateCSC cscValue = + let offsets = cscColumnOffsetsVector cscValue + rowIndices = cscRowIndicesVector cscValue + values = cscValuesVector cscValue + valueCount = U.length values + in case vectorEndpoints offsets of + Nothing -> Left (InvariantViolation "CSC column offsets must be non-empty") + Just (firstOffset, lastOffset) -> + validateCSCPayload cscValue offsets rowIndices values valueCount firstOffset lastOffset + +validateCSCPayload :: (Eq a, AdditiveMonoid a, U.Unbox a) => SparseCSC a -> U.Vector Int -> U.Vector Int -> U.Vector a -> Int -> Int -> Int -> Either MoonlightError () +validateCSCPayload cscValue offsets rowIndices values valueCount firstOffset lastOffset + | cscRows cscValue < 0 || cscCols cscValue < 0 = + Left (InvariantViolation "CSC dimensions must be non-negative") + | U.length offsets /= cscCols cscValue + 1 = + Left (InvariantViolation "CSC column-offset length must equal column count + 1") + | U.length rowIndices /= valueCount = + Left (InvariantViolation "CSC row-index and value payload lengths must match") + | U.any (< 0) offsets = + Left (InvariantViolation "CSC column offsets must be non-negative") + | not (isMonotonicVector offsets) = + Left (InvariantViolation "CSC column offsets must be monotonically non-decreasing") + | firstOffset /= 0 = + Left (InvariantViolation "CSC column offsets must begin at 0") + | lastOffset /= valueCount = + Left (InvariantViolation "CSC terminal column offset must equal value count") + | U.any (\rowIndex -> rowIndex < 0 || rowIndex >= cscRows cscValue) rowIndices = + Left (InvariantViolation "CSC row index out of bounds") + | not (cscColumnsStrictlyCanonical cscValue offsets rowIndices) = + Left (InvariantViolation "CSC row indices must be strictly increasing within each column") + | U.any (== zero) values = + Left (InvariantViolation "CSC values must not store exact zeros") + | otherwise = Right () + +cscColumnsStrictlyCanonical :: SparseCSC a -> U.Vector Int -> U.Vector Int -> Bool +cscColumnsStrictlyCanonical cscValue offsets rowIndices = + U.all cscColumnStrictlyCanonical (U.enumFromN 0 (cscCols cscValue)) + where + cscColumnStrictlyCanonical columnIndex = + let startOffset = offsets `U.unsafeIndex` columnIndex + stopOffset = offsets `U.unsafeIndex` (columnIndex + 1) + columnRows = U.slice startOffset (stopOffset - startOffset) rowIndices + in U.and (U.zipWith (<) columnRows (U.drop 1 columnRows)) + +cscToCOO :: U.Unbox a => SparseCSC a -> Either MoonlightError (SparseCOO a) +cscToCOO = + Right . cscToCOOUnchecked + +cscToCOOUnchecked :: U.Unbox a => SparseCSC a -> SparseCOO a +cscToCOOUnchecked cscValue = + SparseCOO + { cooRows = cscRows cscValue, + cooCols = cscCols cscValue, + cooEntries = + U.toList + ( U.zip3 + (cscRowIndicesVector cscValue) + (offsetAxisIndicesVector (cscCols cscValue) (cscColumnOffsetsVector cscValue)) + (cscValuesVector cscValue) + ) + } + +offsetAxisIndicesVector :: Int -> U.Vector Int -> U.Vector Int +offsetAxisIndicesVector axisCount offsets = + runST $ do + let !entryCount = offsets `U.unsafeIndex` axisCount + axisVector <- MU.unsafeNew entryCount + fillOffsetAxisIndices axisVector 0 + U.unsafeFreeze axisVector + where + fillOffsetAxisIndices :: forall s. MU.MVector s Int -> Int -> ST s () + fillOffsetAxisIndices axisVector !axisIndex + | axisIndex >= axisCount = pure () + | otherwise = do + let !entryStart = offsets `U.unsafeIndex` axisIndex + !entryStop = offsets `U.unsafeIndex` (axisIndex + 1) + fillAxisSpan axisVector axisIndex entryStart entryStop + fillOffsetAxisIndices axisVector (axisIndex + 1) + + fillAxisSpan :: forall s. MU.MVector s Int -> Int -> Int -> Int -> ST s () + fillAxisSpan axisVector !axisIndex !entryIndex !entryStop + | entryIndex >= entryStop = pure () + | otherwise = do + MU.unsafeWrite axisVector entryIndex axisIndex + fillAxisSpan axisVector axisIndex (entryIndex + 1) entryStop + +combineEntries :: + AdditiveGroup a => + [(Int, Int, a)] -> + Map (Int, Int) a +combineEntries = + foldl' + (\entryMap (rowIndex, columnIndex, value) -> Map.insertWith add (rowIndex, columnIndex) value entryMap) + Map.empty + +cooToDense :: + forall r c a. + (KnownNat r, KnownNat c, AdditiveGroup a) => + SparseCOO a -> + Either MoonlightError (Matrix r c a) +cooToDense cooValue = do + rowCount <- checkedSparseStaticDimension @r + columnCount <- checkedSparseStaticDimension @c + _ <- + first + (const (InvariantViolation "static sparse/dense shape exceeds Int cardinality")) + (checkedNonNegativeProduct rowCount columnCount) + if cooRows cooValue /= rowCount || cooCols cooValue /= columnCount + then + Left + ( InvariantViolation + ( "COO shape does not match static matrix dimensions: expected " + <> show (rowCount, columnCount) + <> " but received " + <> show (cooRows cooValue, cooCols cooValue) + ) + ) + else + if any (sparseDenseEntryOutOfBounds rowCount columnCount) (cooEntries cooValue) + then Left (InvariantViolation "COO entry index out of bounds") + else + let entryMap = combineEntries (cooEntries cooValue) + flatValues = + concatMap + (\rowIndex -> map (\columnIndex -> Map.findWithDefault zero (rowIndex, columnIndex) entryMap) [0 .. columnCount - 1]) + [0 .. rowCount - 1] + in fromListMatrix @r @c flatValues + +checkedSparseStaticDimension :: forall n. KnownNat n => Either MoonlightError Int +checkedSparseStaticDimension = + first + (const (InvariantViolation "static sparse/dense dimension exceeds Int cardinality")) + (checkedNaturalToInt (natVal (Proxy @n))) + +sparseDenseEntryOutOfBounds :: Int -> Int -> (Int, Int, a) -> Bool +sparseDenseEntryOutOfBounds rowCount columnCount (rowIndex, columnIndex, _) = + rowIndex < 0 + || rowIndex >= rowCount + || columnIndex < 0 + || columnIndex >= columnCount + +csrToDense :: + forall r c a. + (KnownNat r, KnownNat c, AdditiveGroup a, U.Unbox a) => + SparseCSR a -> + Either MoonlightError (Matrix r c a) +csrToDense csrValue = csrToCOO csrValue >>= cooToDense + +cscToDense :: + forall r c a. + (KnownNat r, KnownNat c, AdditiveGroup a, U.Unbox a) => + SparseCSC a -> + Either MoonlightError (Matrix r c a) +cscToDense cscValue = cscToCOO cscValue >>= cooToDense + +csrToCSC :: U.Unbox a => SparseCSR a -> Either MoonlightError (SparseCSC a) +csrToCSC csrValue = + Right (csrToCSCUnchecked csrValue) +{-# INLINE csrToCSC #-} + +csrToCSCUnchecked :: U.Unbox a => SparseCSR a -> SparseCSC a +csrToCSCUnchecked csrValue = + let SparseEntryVectors offsetVector rowVector valueVector = + countingTransposeCompressed + (csrRows csrValue) + (csrCols csrValue) + (csrRowOffsetsVector csrValue) + (csrColumnIndicesVector csrValue) + (csrValuesVector csrValue) + in SparseCSC + { cscRows = csrRows csrValue, + cscCols = csrCols csrValue, + cscColumnOffsetsVector = offsetVector, + cscRowIndicesVector = rowVector, + cscValuesVector = valueVector + } +{-# INLINE csrToCSCUnchecked #-} + +cscToCSR :: U.Unbox a => SparseCSC a -> Either MoonlightError (SparseCSR a) +cscToCSR cscValue = + Right (cscToCSRUnchecked cscValue) + +cscToCSRUnchecked :: U.Unbox a => SparseCSC a -> SparseCSR a +cscToCSRUnchecked cscValue = + let SparseEntryVectors offsetVector columnVector valueVector = + countingTransposeCompressed + (cscCols cscValue) + (cscRows cscValue) + (cscColumnOffsetsVector cscValue) + (cscRowIndicesVector cscValue) + (cscValuesVector cscValue) + in csrFromCanonicalVectorsUnchecked + (cscRows cscValue) + (cscCols cscValue) + offsetVector + columnVector + valueVector + +countingTransposeCompressed :: + forall a. + U.Unbox a => + Int -> + Int -> + U.Vector Int -> + U.Vector Int -> + U.Vector a -> + SparseEntryVectors a +countingTransposeCompressed majorCount minorCount majorOffsets minorIndices values = + runST $ do + minorCounts <- MU.replicate minorCount 0 + countMinorOccurrences minorCounts 0 + minorOffsets <- MU.replicate (minorCount + 1) 0 + prefixMutableCountsWithStarts minorCount minorCounts minorOffsets + majorIndices <- MU.unsafeNew entryCount + transposedValues <- MU.unsafeNew entryCount + scatterMajorEntries minorCounts majorIndices transposedValues 0 + frozenOffsets <- U.unsafeFreeze minorOffsets + frozenIndices <- U.unsafeFreeze majorIndices + frozenValues <- U.unsafeFreeze transposedValues + pure + SparseEntryVectors + { sparseEntryOffsets = frozenOffsets, + sparseEntryIndices = frozenIndices, + sparseEntryValues = frozenValues + } + where + !entryCount = U.length values + + countMinorOccurrences :: forall s. MU.MVector s Int -> Int -> ST s () + countMinorOccurrences minorCounts !entryIndex + | entryIndex >= entryCount = pure () + | otherwise = do + let !minorIndex = minorIndices `U.unsafeIndex` entryIndex + incrementMutableInt minorCounts minorIndex + countMinorOccurrences minorCounts (entryIndex + 1) + + scatterMajorEntries :: + forall s. + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s a -> + Int -> + ST s () + scatterMajorEntries nextOffsets majorIndices transposedValues !majorIndex + | majorIndex >= majorCount = pure () + | otherwise = do + let !entryStart = majorOffsets `U.unsafeIndex` majorIndex + !entryStop = majorOffsets `U.unsafeIndex` (majorIndex + 1) + scatterMajorSpan nextOffsets majorIndices transposedValues majorIndex entryStart entryStop + scatterMajorEntries nextOffsets majorIndices transposedValues (majorIndex + 1) + + scatterMajorSpan :: + forall s. + MU.MVector s Int -> + MU.MVector s Int -> + MU.MVector s a -> + Int -> + Int -> + Int -> + ST s () + scatterMajorSpan nextOffsets majorIndices transposedValues !majorIndex !entryIndex !entryStop + | entryIndex >= entryStop = pure () + | otherwise = do + let !minorIndex = minorIndices `U.unsafeIndex` entryIndex + !entryValue = values `U.unsafeIndex` entryIndex + targetIndex <- MU.unsafeRead nextOffsets minorIndex + MU.unsafeWrite majorIndices targetIndex majorIndex + MU.unsafeWrite transposedValues targetIndex entryValue + MU.unsafeWrite nextOffsets minorIndex (targetIndex + 1) + scatterMajorSpan nextOffsets majorIndices transposedValues majorIndex (entryIndex + 1) entryStop +{-# INLINE countingTransposeCompressed #-} + +csrMatVecVector :: SparseCSR Double -> U.Vector Double -> Either MoonlightError (U.Vector Double) +csrMatVecVector csrValue vectorValue = + if U.length vectorValue /= csrCols csrValue + then + Left + ( InvariantViolation + ( "CSR matvec dimension mismatch: expected " + <> show (csrCols csrValue) + <> " but received " + <> show (U.length vectorValue) + ) + ) + else + Right + ( case csrExecutionPlan csrValue of + CSRGeneral -> + csrMatVecValidatedU + (csrRows csrValue) + (csrRowOffsetsVector csrValue) + (csrColumnIndicesVector csrValue) + (csrValuesVector csrValue) + vectorValue + CSRContiguousBand lowerBandwidth upperBandwidth -> + csrContiguousBandMatVecValidatedU + (csrRows csrValue) + lowerBandwidth + upperBandwidth + (csrValuesVector csrValue) + vectorValue + ) diff --git a/src-spectral/Moonlight/LinAlg/Pure/Krylov/Arnoldi.hs b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Arnoldi.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Arnoldi.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Krylov.Arnoldi + ( arnoldi, + ) +where + +import qualified Data.Vector as Box +import qualified Data.Vector.Unboxed as U +import Moonlight.Core (MoonlightError) +import Moonlight.LinAlg.Internal.VectorOps (normU, scaleU) +import Moonlight.LinAlg.Pure.Krylov.Config (ArnoldiConfig, arnoldiIterations, arnoldiReorthogonalize, arnoldiTolerance) +import Moonlight.LinAlg.Pure.Krylov.Decomposition (ArnoldiDecomposition, mkArnoldiDecomposition) +import Moonlight.LinAlg.Pure.Krylov.Internal + ( normalizeSeed, + orthogonalizeAgainst, + requireBasisVector, + sparseColumnsToDenseRowVectors, + validateIterationCount, + validateSquareOperator, + ) +import Moonlight.LinAlg.Pure.Operator (LinearOperator, operatorShape, runOperatorU) +import Prelude + +arnoldi :: ArnoldiConfig -> LinearOperator symmetry -> U.Vector Double -> Either MoonlightError ArnoldiDecomposition +arnoldi config op seedVector = do + validateSquareOperator "Arnoldi" op + let (_, cols) = operatorShape op + targetIterations <- validateIterationCount "Arnoldi" (arnoldiIterations config) + firstBasis <- normalizeSeed "Arnoldi" cols (arnoldiTolerance config) seedVector + let boundedIterations = min targetIterations cols + go basisVectors hessenbergColumns iterationIndex = do + currentBasis <- requireBasisVector iterationIndex basisVectors + imageVector <- runOperatorU op currentBasis + (reducedVector, coefficients) <- + orthogonalizeAgainst (arnoldiReorthogonalize config) basisVectors imageVector + let nextNorm = normU reducedVector + nextColumn = U.snoc coefficients nextNorm + nextHessenbergColumns = hessenbergColumns `Box.snoc` nextColumn + if nextNorm <= arnoldiTolerance config || iterationIndex + 1 >= boundedIterations + then finalize basisVectors nextHessenbergColumns + else + let nextBasis = scaleU (1.0 / nextNorm) reducedVector + in go (basisVectors `Box.snoc` nextBasis) nextHessenbergColumns (iterationIndex + 1) + finalize basisVectors hessenbergColumns = + let hessenbergRows = sparseColumnsToDenseRowVectors (stepCount + 1) stepCount hessenbergColumns + stepCount = Box.length hessenbergColumns + in mkArnoldiDecomposition basisVectors hessenbergRows + in go (Box.singleton firstBasis) Box.empty 0 diff --git a/src-spectral/Moonlight/LinAlg/Pure/Krylov/Block.hs b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Block.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Block.hs @@ -0,0 +1,148 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Krylov.Block + ( blockLanczosSymmetric + ) +where + +import Data.Kind (Type) +import qualified Data.Vector as Box +import qualified Data.Vector.Unboxed as U +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Pure.Krylov.Config +import Moonlight.LinAlg.Pure.Krylov.Decomposition +import Moonlight.LinAlg.Pure.Operator +import Moonlight.LinAlg.Pure.Krylov.Internal + ( blockInnerBlock, + multiplyBasisByBlock, + normalizeSeedBlock, + orthonormalizeBlock, + selfAdjointBlockInnerBlock, + subtractBlocks, + validateIterationCount, + validateSquareOperator, + ) +import Moonlight.LinAlg.Pure.Structured.BlockTridiagonal + ( RowMajorBlock, + mkSymmetricBlockTridiagonal, + transposeRowMajorBlock, + ) +import Prelude + +type PreviousBlockContext :: Type +data PreviousBlockContext + = InitialBlockContext + | PreviousBlockContext !(Box.Vector (U.Vector Double)) !RowMajorBlock + +type BlockStepResult :: Type +data BlockStepResult = BlockStepResult + { stepDiagonalBlock :: !RowMajorBlock, + stepNextBlock :: !(Box.Vector (U.Vector Double)), + stepNextCouplingBlock :: !(Maybe RowMajorBlock) + } + +blockLanczosSymmetric :: + BlockLanczosConfig -> + LinearOperator 'SelfAdjointOperator -> + Box.Vector (U.Vector Double) -> + Either MoonlightError BlockLanczosDecomposition +blockLanczosSymmetric config op seedBlock = do + validateSquareOperator "Block Lanczos" op + let (_, dimension) = operatorShape op + boundedBlockSize = min dimension (blockLanczosBlockSize config) + boundedIterations <- validateIterationCount "Block Lanczos" (blockLanczosIterations config) + initialBlock <- + normalizeSeedBlock + "Block Lanczos" + dimension + (blockLanczosTolerance config) + boundedBlockSize + seedBlock + iterateBlocks boundedIterations initialBlock initialBlock InitialBlockContext [] [] 1 + where + iterateBlocks boundedIterations accumulatedBasis currentBlock previousContext diagonalBlocksRev couplingBlocksRev blockStepCount = do + imageBlock <- Box.fromList <$> traverse (runOperatorU op) (Box.toList currentBlock) + stepResult <- + blockLanczosStep + config + accumulatedBasis + currentBlock + previousContext + imageBlock + let nextDiagonalBlocksRev = stepDiagonalBlock stepResult : diagonalBlocksRev + nextCouplingBlocksRev = + maybe couplingBlocksRev (: couplingBlocksRev) (stepNextCouplingBlock stepResult) + nextAccumulatedBasis = accumulatedBasis <> stepNextBlock stepResult + in if blockStepCount >= boundedIterations || Box.length accumulatedBasis >= snd (operatorShape op) || Box.null (stepNextBlock stepResult) + then + finalize + accumulatedBasis + nextDiagonalBlocksRev + couplingBlocksRev + blockStepCount + else + iterateBlocks + boundedIterations + nextAccumulatedBasis + (stepNextBlock stepResult) + (case stepNextCouplingBlock stepResult of + Nothing -> InitialBlockContext + Just couplingBlock -> PreviousBlockContext currentBlock couplingBlock) + nextDiagonalBlocksRev + nextCouplingBlocksRev + (blockStepCount + 1) + + finalize basisVectors diagonalBlocksRev couplingBlocksRev blockStepCount = + let diagonalBlocks = Box.fromList (reverse diagonalBlocksRev) + couplingBlocks = Box.fromList (reverse couplingBlocksRev) + in do + projectedBlockTridiagonal <- mkSymmetricBlockTridiagonal diagonalBlocks couplingBlocks + mkBlockLanczosDecomposition basisVectors projectedBlockTridiagonal blockStepCount + +blockLanczosStep :: + BlockLanczosConfig -> + Box.Vector (U.Vector Double) -> + Box.Vector (U.Vector Double) -> + PreviousBlockContext -> + Box.Vector (U.Vector Double) -> + Either MoonlightError BlockStepResult +blockLanczosStep config accumulatedBasis currentBlock previousContext imageBlock = do + (alphaRows, recurrenceResidual) <- threeTermResidual currentBlock previousContext imageBlock + let stabilizationBasis = + if blockLanczosReorthogonalize config + then accumulatedBasis + else Box.empty + nextBlock <- + orthonormalizeBlock + (blockLanczosReorthogonalize config) + (blockLanczosTolerance config) + stabilizationBasis + recurrenceResidual + nextCouplingBlock <- + if Box.null nextBlock + then Right Nothing + else Just <$> blockInnerBlock nextBlock recurrenceResidual + Right + BlockStepResult + { stepDiagonalBlock = alphaRows, + stepNextBlock = nextBlock, + stepNextCouplingBlock = nextCouplingBlock + } + +threeTermResidual :: + Box.Vector (U.Vector Double) -> + PreviousBlockContext -> + Box.Vector (U.Vector Double) -> + Either MoonlightError (RowMajorBlock, Box.Vector (U.Vector Double)) +threeTermResidual currentBlock previousContext imageBlock = do + alphaRows <- selfAdjointBlockInnerBlock currentBlock imageBlock + currentContribution <- multiplyBasisByBlock currentBlock alphaRows + residualAfterCurrent <- subtractBlocks imageBlock currentContribution + recurrenceResidual <- + case previousContext of + InitialBlockContext -> Right residualAfterCurrent + PreviousBlockContext previousBlock couplingRows -> do + previousContribution <- multiplyBasisByBlock previousBlock (transposeRowMajorBlock couplingRows) + subtractBlocks residualAfterCurrent previousContribution + Right (alphaRows, recurrenceResidual) diff --git a/src-spectral/Moonlight/LinAlg/Pure/Krylov/CascadicGraph.hs b/src-spectral/Moonlight/LinAlg/Pure/Krylov/CascadicGraph.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Krylov/CascadicGraph.hs @@ -0,0 +1,782 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Krylov.CascadicGraph + ( CascadicGraphObstruction (..), + cascadicGraphLaplacianEigenpairs, + ) +where + +import Data.Bifunctor (first) +import Data.Function ((&)) +import Data.IntMap.Strict (IntMap) +import Data.IntMap.Strict qualified as IntMap +import Data.Kind (Type) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Vector qualified as Box +import Data.Vector.Storable qualified as S +import Data.Vector.Unboxed qualified as U +import Moonlight.Core + ( MoonlightError (..), + checkedNonNegativeProduct, + fieldValueValid, + ) +import Moonlight.LinAlg.Internal.Eigen.Kernels (epsDouble) +import Moonlight.LinAlg.Internal.Eigen.Symmetric + ( SymmetricEigenResult (..), + symmetricEigenPairsDenseUnchecked, + ) +import Moonlight.LinAlg.Internal.VectorOps (normU, subScaledU) +import Moonlight.LinAlg.Pure.Dense.Flat + ( denseDoubleMatrixToRowMajorVector, + mkDenseDoubleMatrixRowMajor, + ) +import Moonlight.LinAlg.Pure.Krylov.Config + ( LanczosConfig, + lanczosIterations, + lanczosTolerance, + ) +import Moonlight.LinAlg.Pure.Krylov.Internal + ( linearCombinationColumnsU, + orthonormalizeBlock, + ) +import Moonlight.LinAlg.Pure.Sparse.Assembly (canonicalCSRFromEntries) +import Moonlight.LinAlg.Pure.Sparse.Solver.Common (sparseDiagonal) +import Moonlight.LinAlg.Pure.Sparse.Structured + ( GraphEdge (..), + graphLaplacianCSR, + ) +import Moonlight.LinAlg.Pure.Sparse.Types + ( SparseCSR, + cooEntries, + csrCols, + csrMatVecVector, + csrRows, + csrToCOO, + ) +import Moonlight.LinAlg.Pure.Spectral.Result + ( Eigenpairs, + eigenpairsFromColumns, + ) +import Prelude + +type CascadicGraphObstruction :: Type +data CascadicGraphObstruction + = CascadicGraphBackendFailure !MoonlightError + | CascadicGraphCoarseningStalled !Int + | CascadicGraphIncompleteAssignment !Int + | CascadicGraphInvalidRequest !Int !Int + | CascadicGraphRankLoss !Int !Int + | CascadicGraphRefinementBudgetExceeded !Int !Double !Double + deriving stock (Eq, Show) + +type GraphLaplacianLevel :: Type +data GraphLaplacianLevel = GraphLaplacianLevel + { graphLevelMasses :: !(U.Vector Double), + graphLevelEdges :: !(Box.Vector (GraphEdge Int)), + graphLevelMatrix :: !(SparseCSR Double) + } + +type GraphAggregation :: Type +data GraphAggregation = GraphAggregation + { graphAggregationFineToCoarse :: !(U.Vector Int), + graphAggregationCoarseMasses :: !(U.Vector Double), + graphAggregationCoarseEdges :: !(Box.Vector (GraphEdge Int)) + } + +type AggregationFold :: Type +data AggregationFold = AggregationFold + { aggregationAssignments :: !(IntMap Int), + aggregationNextIndex :: !Int + } + +type WeightedNeighbor :: Type +data WeightedNeighbor = WeightedNeighbor + { weightedNeighborVertex :: !Int, + weightedNeighborEdgeWeight :: !Double + } + +type RitzBlock :: Type +data RitzBlock = RitzBlock + { ritzBlockValues :: !(U.Vector Double), + ritzBlockColumns :: !(Box.Vector (U.Vector Double)), + ritzBlockResidualVectors :: !(Box.Vector (U.Vector Double)), + ritzBlockResidualNorms :: !(U.Vector Double) + } + +type RefinementState :: Type +data RefinementState = RefinementState + { refinementStepCount :: !Int, + refinementRitzBlock :: !RitzBlock + } + +cascadicGraphLaplacianEigenpairs :: + LanczosConfig -> + Int -> + SparseCSR Double -> + Either CascadicGraphObstruction Eigenpairs +cascadicGraphLaplacianEigenpairs config requestedCount fineMatrix = do + validateCascadicRequest requestedCount fineMatrix + fineLevel <- initialGraphLaplacianLevel fineMatrix + if Box.null (graphLevelEdges fineLevel) + then zeroGraphEigenpairs requestedCount (csrRows fineMatrix) + else do + refinementLimit <- cascadicRefinementLimit config + finalBlock <- + solveCascadicLevel + requestedCount + refinementLimit + (cascadicResidualTarget config (csrRows fineMatrix)) + True + fineLevel + ritzBlockToEigenpairs finalBlock + +validateCascadicRequest :: + Int -> + SparseCSR Double -> + Either CascadicGraphObstruction () +validateCascadicRequest requestedCount matrixValue + | csrRows matrixValue <= 0 || csrRows matrixValue /= csrCols matrixValue = + Left + ( CascadicGraphBackendFailure + (InvariantViolation "cascadic graph eigensolve requires a positive square matrix") + ) + | requestedCount <= 0 || requestedCount > csrRows matrixValue = + Left (CascadicGraphInvalidRequest requestedCount (csrRows matrixValue)) + | otherwise = Right () + +initialGraphLaplacianLevel :: + SparseCSR Double -> + Either CascadicGraphObstruction GraphLaplacianLevel +initialGraphLaplacianLevel matrixValue = do + coordinateValue <- first CascadicGraphBackendFailure (csrToCOO matrixValue) + let dimension = csrRows matrixValue + edges = + Box.fromList + [ GraphEdge rowIndex columnIndex (negate entryValue) + | (rowIndex, columnIndex, entryValue) <- cooEntries coordinateValue, + rowIndex < columnIndex, + entryValue < 0.0 + ] + pure + GraphLaplacianLevel + { graphLevelMasses = U.replicate dimension 1.0, + graphLevelEdges = edges, + graphLevelMatrix = matrixValue + } + +zeroGraphEigenpairs :: + Int -> + Int -> + Either CascadicGraphObstruction Eigenpairs +zeroGraphEigenpairs requestedCount dimension = + first CascadicGraphBackendFailure + ( eigenpairsFromColumns + dimension + [ (0.0, unitVector dimension columnIndex, 0.0) + | columnIndex <- [0 .. requestedCount - 1] + ] + ) + +solveCascadicLevel :: + Int -> + Int -> + Double -> + Bool -> + GraphLaplacianLevel -> + Either CascadicGraphObstruction RitzBlock +solveCascadicLevel requestedCount refinementLimit residualTarget isFinest levelValue + | graphLevelDimension levelValue <= cascadicCoarsestDimension = + coarsestRitzBlock requestedCount levelValue + | otherwise = do + aggregation <- heavyEdgeAggregation levelValue + coarseLevel <- graphLaplacianCoarseLevel aggregation + coarseBlock <- + solveCascadicLevel + requestedCount + refinementLimit + residualTarget + False + coarseLevel + prolongedColumns <- + prolongRitzColumns + levelValue + aggregation + (ritzBlockColumns coarseBlock) + initialBlock <- + rayleighRitzBlock + requestedCount + (graphLevelMatrix levelValue) + prolongedColumns + refineRitzBlock + requestedCount + refinementLimit + residualTarget + isFinest + levelValue + initialBlock + +cascadicCoarsestDimension :: Int +cascadicCoarsestDimension = 256 + +cascadicResidualTarget :: LanczosConfig -> Int -> Double +cascadicResidualTarget config dimension = + max + (sqrt (lanczosTolerance config)) + (128.0 * epsDouble * sqrt (fromIntegral (max 1 dimension) :: Double)) + +cascadicRefinementLimit :: + LanczosConfig -> + Either CascadicGraphObstruction Int +cascadicRefinementLimit config = + first + ( CascadicGraphBackendFailure + . const (InvariantViolation "cascadic graph refinement budget exceeds Int range") + ) + (checkedNonNegativeProduct 5 (lanczosIterations config)) + +heavyEdgeAggregation :: + GraphLaplacianLevel -> + Either CascadicGraphObstruction GraphAggregation +heavyEdgeAggregation levelValue = do + fineToCoarse <- + U.generateM + fineDimension + ( \vertexIndex -> + case IntMap.lookup vertexIndex (aggregationAssignments finalFold) of + Nothing -> Left (CascadicGraphIncompleteAssignment vertexIndex) + Just coarseIndex -> Right coarseIndex + ) + let coarseDimension = aggregationNextIndex finalFold + if coarseDimension >= fineDimension + then Left (CascadicGraphCoarseningStalled fineDimension) + else + let coarseMasses = + U.accumulate + (+) + (U.replicate coarseDimension 0.0) + (U.zip fineToCoarse (graphLevelMasses levelValue)) + coarseEdgeMap = + Box.foldl' + (collapseFineEdge fineToCoarse) + Map.empty + (graphLevelEdges levelValue) + coarseEdges = + Box.fromList + [ GraphEdge leftIndex rightIndex edgeWeight + | ((leftIndex, rightIndex), edgeWeight) <- Map.toAscList coarseEdgeMap + ] + in Right + GraphAggregation + { graphAggregationFineToCoarse = fineToCoarse, + graphAggregationCoarseMasses = coarseMasses, + graphAggregationCoarseEdges = coarseEdges + } + where + fineDimension = graphLevelDimension levelValue + adjacency = graphAdjacency (graphLevelEdges levelValue) + finalFold = + foldl' + (assignHeavyEdgeAggregate adjacency) + (AggregationFold IntMap.empty 0) + [0 .. fineDimension - 1] + +graphAdjacency :: + Box.Vector (GraphEdge Int) -> + IntMap [WeightedNeighbor] +graphAdjacency = + Box.foldr + ( \edgeValue -> + IntMap.insertWith + (<>) + (graphEdgeLeft edgeValue) + [ WeightedNeighbor + (graphEdgeRight edgeValue) + (graphEdgeWeight edgeValue) + ] + . IntMap.insertWith + (<>) + (graphEdgeRight edgeValue) + [ WeightedNeighbor + (graphEdgeLeft edgeValue) + (graphEdgeWeight edgeValue) + ] + ) + IntMap.empty + +assignHeavyEdgeAggregate :: + IntMap [WeightedNeighbor] -> + AggregationFold -> + Int -> + AggregationFold +assignHeavyEdgeAggregate adjacency foldValue vertexIndex = + case IntMap.lookup vertexIndex (aggregationAssignments foldValue) of + Just _ -> foldValue + Nothing -> + let selectedNeighbor = + foldl' + (selectHeavierUnassignedNeighbor (aggregationAssignments foldValue)) + Nothing + (IntMap.findWithDefault [] vertexIndex adjacency) + aggregateIndex = aggregationNextIndex foldValue + withVertex = + IntMap.insert + vertexIndex + aggregateIndex + (aggregationAssignments foldValue) + withNeighbor = + case selectedNeighbor of + Nothing -> withVertex + Just neighborValue -> + IntMap.insert + (weightedNeighborVertex neighborValue) + aggregateIndex + withVertex + in AggregationFold + { aggregationAssignments = withNeighbor, + aggregationNextIndex = aggregateIndex + 1 + } + +selectHeavierUnassignedNeighbor :: + IntMap Int -> + Maybe WeightedNeighbor -> + WeightedNeighbor -> + Maybe WeightedNeighbor +selectHeavierUnassignedNeighbor assignments selected candidate + | IntMap.member (weightedNeighborVertex candidate) assignments = selected + | otherwise = + case selected of + Nothing -> Just candidate + Just selectedValue + | weightedNeighborEdgeWeight candidate > weightedNeighborEdgeWeight selectedValue -> + Just candidate + | weightedNeighborEdgeWeight candidate == weightedNeighborEdgeWeight selectedValue + && weightedNeighborVertex candidate < weightedNeighborVertex selectedValue -> + Just candidate + | otherwise -> selected + +collapseFineEdge :: + U.Vector Int -> + Map (Int, Int) Double -> + GraphEdge Int -> + Map (Int, Int) Double +collapseFineEdge fineToCoarse edgeMap edgeValue = + let leftAggregate = fineToCoarse `U.unsafeIndex` graphEdgeLeft edgeValue + rightAggregate = fineToCoarse `U.unsafeIndex` graphEdgeRight edgeValue + in if leftAggregate == rightAggregate + then edgeMap + else + Map.insertWith + (+) + (min leftAggregate rightAggregate, max leftAggregate rightAggregate) + (graphEdgeWeight edgeValue) + edgeMap + +graphLaplacianCoarseLevel :: + GraphAggregation -> + Either CascadicGraphObstruction GraphLaplacianLevel +graphLaplacianCoarseLevel aggregation = do + let coarseMasses = graphAggregationCoarseMasses aggregation + coarseDimension = U.length coarseMasses + coarseEdges = graphAggregationCoarseEdges aggregation + laplacianMatrix <- + first CascadicGraphBackendFailure + ( graphLaplacianCSR + [0 .. coarseDimension - 1] + (Box.toList coarseEdges) + ) + coordinateValue <- first CascadicGraphBackendFailure (csrToCOO laplacianMatrix) + normalizedMatrix <- + first CascadicGraphBackendFailure + ( canonicalCSRFromEntries + coarseDimension + coarseDimension + [ ( rowIndex, + columnIndex, + entryValue + / sqrt + ( (coarseMasses `U.unsafeIndex` rowIndex) + * (coarseMasses `U.unsafeIndex` columnIndex) + ) + ) + | (rowIndex, columnIndex, entryValue) <- cooEntries coordinateValue + ] + ) + pure + GraphLaplacianLevel + { graphLevelMasses = coarseMasses, + graphLevelEdges = coarseEdges, + graphLevelMatrix = normalizedMatrix + } + +prolongRitzColumns :: + GraphLaplacianLevel -> + GraphAggregation -> + Box.Vector (U.Vector Double) -> + Either CascadicGraphObstruction (Box.Vector (U.Vector Double)) +prolongRitzColumns fineLevel aggregation coarseColumns = + Box.mapM prolongColumn coarseColumns + where + fineMasses = graphLevelMasses fineLevel + coarseMasses = graphAggregationCoarseMasses aggregation + fineToCoarse = graphAggregationFineToCoarse aggregation + fineDimension = graphLevelDimension fineLevel + prolongColumn coarseColumn + | U.length coarseColumn /= U.length coarseMasses = + Left + ( CascadicGraphBackendFailure + (InvariantViolation "cascadic graph coarse eigenvector dimension mismatch") + ) + | otherwise = + Right + ( U.generate + fineDimension + ( \fineIndex -> + let coarseIndex = fineToCoarse `U.unsafeIndex` fineIndex + fineMass = fineMasses `U.unsafeIndex` fineIndex + coarseMass = coarseMasses `U.unsafeIndex` coarseIndex + in sqrt (fineMass / coarseMass) + * (coarseColumn `U.unsafeIndex` coarseIndex) + ) + ) + +coarsestRitzBlock :: + Int -> + GraphLaplacianLevel -> + Either CascadicGraphObstruction RitzBlock +coarsestRitzBlock requestedCount levelValue = do + eigenResult <- denseSymmetricEigenResult (graphLevelMatrix levelValue) + let dimension = graphLevelDimension levelValue + vectorPayload = + denseDoubleMatrixToRowMajorVector + (symmetricEigenResultVectors eigenResult) + initialColumns = + Box.generate + requestedCount + ( \columnIndex -> + U.generate + dimension + ( \rowIndex -> + vectorPayload S.! (rowIndex * dimension + columnIndex) + ) + ) + rayleighRitzBlock + requestedCount + (graphLevelMatrix levelValue) + initialColumns + +denseSymmetricEigenResult :: + SparseCSR Double -> + Either CascadicGraphObstruction SymmetricEigenResult +denseSymmetricEigenResult matrixValue = do + let dimension = csrRows matrixValue + entryCount <- + first + ( CascadicGraphBackendFailure + . const (InvariantViolation "cascadic coarse dense cardinality exceeds Int range") + ) + (checkedNonNegativeProduct dimension dimension) + imageColumns <- + traverse + ( first CascadicGraphBackendFailure + . csrMatVecVector matrixValue + . unitVector dimension + ) + [0 .. dimension - 1] + let columnPayload = U.concat imageColumns + rowMajorPayload = + S.generate + entryCount + ( \flatIndex -> + let (rowIndex, columnIndex) = flatIndex `quotRem` dimension + in columnPayload `U.unsafeIndex` (columnIndex * dimension + rowIndex) + ) + denseMatrix <- + first CascadicGraphBackendFailure + (mkDenseDoubleMatrixRowMajor dimension dimension rowMajorPayload) + first CascadicGraphBackendFailure + (symmetricEigenPairsDenseUnchecked dimension denseMatrix) + +refineRitzBlock :: + Int -> + Int -> + Double -> + Bool -> + GraphLaplacianLevel -> + RitzBlock -> + Either CascadicGraphObstruction RitzBlock +refineRitzBlock requestedCount refinementLimit residualTarget requireTarget levelValue initialBlock = do + diagonalValues <- + first + (CascadicGraphBackendFailure . InvariantViolation . show) + (sparseDiagonal (graphLevelMatrix levelValue)) + refineState + diagonalValues + RefinementState + { refinementStepCount = 0, + refinementRitzBlock = initialBlock + } + where + refineState diagonalValues stateValue + | ritzBlockMaximumResidual (refinementRitzBlock stateValue) <= residualTarget = + Right (refinementRitzBlock stateValue) + | refinementStepCount stateValue >= refinementLimit = + if requireTarget + then + Left + ( CascadicGraphRefinementBudgetExceeded + (graphLevelDimension levelValue) + residualTarget + (ritzBlockMaximumResidual (refinementRitzBlock stateValue)) + ) + else Right (refinementRitzBlock stateValue) + | otherwise = do + candidateColumns <- + smoothedRitzColumns + diagonalValues + (refinementRitzBlock stateValue) + nextBlock <- + rayleighRitzBlock + requestedCount + (graphLevelMatrix levelValue) + candidateColumns + refineState + diagonalValues + RefinementState + { refinementStepCount = refinementStepCount stateValue + 1, + refinementRitzBlock = nextBlock + } + +smoothedRitzColumns :: + U.Vector Double -> + RitzBlock -> + Either CascadicGraphObstruction (Box.Vector (U.Vector Double)) +smoothedRitzColumns diagonalValues blockValue = do + candidateColumns <- + Box.zipWithM + (smoothRitzColumn diagonalValues) + (ritzBlockColumns blockValue) + (ritzBlockResidualVectors blockValue) + orthonormalColumns <- + first CascadicGraphBackendFailure + ( orthonormalizeBlock + True + cascadicRankTolerance + Box.empty + candidateColumns + ) + requireBlockRank (Box.length candidateColumns) orthonormalColumns + +smoothRitzColumn :: + U.Vector Double -> + U.Vector Double -> + U.Vector Double -> + Either CascadicGraphObstruction (U.Vector Double) +smoothRitzColumn diagonalValues columnValue residualVector + | U.length columnValue /= U.length diagonalValues + || U.length residualVector /= U.length diagonalValues = + Left + ( CascadicGraphBackendFailure + (InvariantViolation "cascadic graph smoother dimension mismatch") + ) + | otherwise = + Right + ( U.generate + (U.length columnValue) + ( \entryIndex -> + let diagonalValue = diagonalValues `U.unsafeIndex` entryIndex + columnEntry = columnValue `U.unsafeIndex` entryIndex + residualEntry = residualVector `U.unsafeIndex` entryIndex + in if abs diagonalValue <= cascadicRankTolerance + then columnEntry + else + columnEntry + - cascadicSmoothingWeight + * residualEntry + / diagonalValue + ) + ) + +cascadicSmoothingWeight :: Double +cascadicSmoothingWeight = 0.72 + +cascadicRankTolerance :: Double +cascadicRankTolerance = 1.0e-13 + +rayleighRitzBlock :: + Int -> + SparseCSR Double -> + Box.Vector (U.Vector Double) -> + Either CascadicGraphObstruction RitzBlock +rayleighRitzBlock requestedCount matrixValue candidateColumns = do + orthonormalColumns <- + first CascadicGraphBackendFailure + ( orthonormalizeBlock + True + cascadicRankTolerance + Box.empty + candidateColumns + ) + basisColumns <- requireBlockRank requestedCount orthonormalColumns + basisImages <- + Box.mapM + (first CascadicGraphBackendFailure . csrMatVecVector matrixValue) + basisColumns + projectedResult <- projectedSymmetricEigenResult basisColumns basisImages + let projectedDimension = Box.length basisColumns + projectedVectorPayload = + denseDoubleMatrixToRowMajorVector + (symmetricEigenResultVectors projectedResult) + selectedValues = + U.generate + requestedCount + (S.unsafeIndex (symmetricEigenResultValues projectedResult)) + projectedColumn columnIndex = + U.generate + projectedDimension + ( \rowIndex -> + projectedVectorPayload + S.! (rowIndex * projectedDimension + columnIndex) + ) + selectedColumns <- + Box.generateM + requestedCount + (linearCombinationColumnsU basisColumns . projectedColumn) + & first CascadicGraphBackendFailure + selectedImages <- + Box.generateM + requestedCount + (linearCombinationColumnsU basisImages . projectedColumn) + & first CascadicGraphBackendFailure + residualVectors <- + Box.generateM + requestedCount + ( \columnIndex -> + case + ( selectedImages Box.!? columnIndex, + selectedColumns Box.!? columnIndex, + selectedValues U.!? columnIndex + ) + of + (Just imageVector, Just columnValue, Just eigenvalue) -> + first CascadicGraphBackendFailure + (subScaledU imageVector eigenvalue columnValue) + _ -> + Left + ( CascadicGraphBackendFailure + (InvariantViolation "cascadic graph Ritz column extraction failed") + ) + ) + let residualNorms = U.generate requestedCount (normU . Box.unsafeIndex residualVectors) + if U.all fieldValueValid selectedValues && U.all fieldValueValid residualNorms + then + Right + RitzBlock + { ritzBlockValues = selectedValues, + ritzBlockColumns = selectedColumns, + ritzBlockResidualVectors = residualVectors, + ritzBlockResidualNorms = residualNorms + } + else + Left + ( CascadicGraphBackendFailure + (InvariantViolation "cascadic graph Ritz solve produced non-finite evidence") + ) + +projectedSymmetricEigenResult :: + Box.Vector (U.Vector Double) -> + Box.Vector (U.Vector Double) -> + Either CascadicGraphObstruction SymmetricEigenResult +projectedSymmetricEigenResult basisColumns basisImages = do + let projectedDimension = Box.length basisColumns + projectedEntryCount <- + first + ( CascadicGraphBackendFailure + . const (InvariantViolation "cascadic projected cardinality exceeds Int range") + ) + (checkedNonNegativeProduct projectedDimension projectedDimension) + projectedPayload <- + S.generateM + projectedEntryCount + ( \flatIndex -> + let (rowIndex, columnIndex) = flatIndex `quotRem` projectedDimension + in case + ( basisColumns Box.!? rowIndex, + basisImages Box.!? columnIndex + ) + of + (Just rowVector, Just imageVector) -> + first CascadicGraphBackendFailure + (vectorInnerProduct rowVector imageVector) + _ -> + Left + ( CascadicGraphBackendFailure + (InvariantViolation "cascadic projected basis lookup failed") + ) + ) + projectedMatrix <- + first CascadicGraphBackendFailure + ( mkDenseDoubleMatrixRowMajor + projectedDimension + projectedDimension + projectedPayload + ) + first CascadicGraphBackendFailure + ( symmetricEigenPairsDenseUnchecked + projectedDimension + projectedMatrix + ) + +vectorInnerProduct :: + U.Vector Double -> + U.Vector Double -> + Either MoonlightError Double +vectorInnerProduct leftVector rightVector + | U.length leftVector /= U.length rightVector = + Left (InvariantViolation "cascadic graph inner-product dimension mismatch") + | otherwise = Right (U.sum (U.zipWith (*) leftVector rightVector)) + +requireBlockRank :: + Int -> + Box.Vector (U.Vector Double) -> + Either CascadicGraphObstruction (Box.Vector (U.Vector Double)) +requireBlockRank requiredCount columns + | Box.length columns < requiredCount = + Left (CascadicGraphRankLoss requiredCount (Box.length columns)) + | otherwise = Right (Box.take requiredCount columns) + +ritzBlockMaximumResidual :: RitzBlock -> Double +ritzBlockMaximumResidual = + U.foldl' max 0.0 . ritzBlockResidualNorms + +ritzBlockToEigenpairs :: + RitzBlock -> + Either CascadicGraphObstruction Eigenpairs +ritzBlockToEigenpairs blockValue = + case ritzBlockColumns blockValue Box.!? 0 of + Nothing -> Left (CascadicGraphRankLoss 1 0) + Just firstColumn -> + first CascadicGraphBackendFailure + ( eigenpairsFromColumns + (U.length firstColumn) + [ ( eigenvalue, + columnValue, + residualNorm + ) + | (eigenvalue, columnValue, residualNorm) <- + zip3 + (U.toList (ritzBlockValues blockValue)) + (Box.toList (ritzBlockColumns blockValue)) + (U.toList (ritzBlockResidualNorms blockValue)) + ] + ) + +graphLevelDimension :: GraphLaplacianLevel -> Int +graphLevelDimension = csrRows . graphLevelMatrix + +unitVector :: Int -> Int -> U.Vector Double +unitVector dimension selectedIndex = + U.generate + dimension + (\entryIndex -> if entryIndex == selectedIndex then 1.0 else 0.0) diff --git a/src-spectral/Moonlight/LinAlg/Pure/Krylov/Config.hs b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Config.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Config.hs @@ -0,0 +1,280 @@ +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Krylov.Config + ( KrylovConfigError (..), + krylovConfigErrorMessage, + PositiveCount, + mkPositiveCount, + positiveCountValue, + NonNegativeConfigTolerance, + mkNonNegativeConfigTolerance, + nonNegativeConfigToleranceValue, + ArnoldiConfig, + mkArnoldiConfig, + arnoldiIterations, + arnoldiTolerance, + arnoldiReorthogonalize, + withArnoldiIterations, + withArnoldiTolerance, + withArnoldiReorthogonalize, + defaultArnoldiConfig, + LanczosConfig, + mkLanczosConfig, + lanczosIterations, + lanczosTolerance, + withLanczosIterations, + withLanczosTolerance, + defaultLanczosConfig, + BlockLanczosConfig, + mkBlockLanczosConfig, + blockLanczosIterations, + blockLanczosTolerance, + blockLanczosBlockSize, + blockLanczosReorthogonalize, + withBlockLanczosIterations, + withBlockLanczosTolerance, + withBlockLanczosBlockSize, + withBlockLanczosReorthogonalize, + defaultBlockLanczosConfig, + ) +where + +import Data.Kind (Type) +import Moonlight.Core + ( mkNonNegativeFiniteWith, + mkPositiveIntWith, + ) +import Prelude + +type PositiveCount :: Type +newtype PositiveCount = PositiveCount + { positiveCountValue :: Int + } + deriving stock (Eq, Show) + +type KrylovConfigError :: Type +data KrylovConfigError + = PositiveCountMustBePositive !Int + | KrylovConfigToleranceMustBeNonNegative !Double + deriving stock (Eq, Show) + +krylovConfigErrorMessage :: KrylovConfigError -> String +krylovConfigErrorMessage configError = + case configError of + PositiveCountMustBePositive value -> + "positive count must be positive, received " <> show value + KrylovConfigToleranceMustBeNonNegative value -> + "Krylov config tolerance must be finite and non-negative, received " <> show value + +mkPositiveCount :: Int -> Either KrylovConfigError PositiveCount +mkPositiveCount = + mkPositiveIntWith PositiveCountMustBePositive PositiveCount + +type NonNegativeConfigTolerance :: Type +newtype NonNegativeConfigTolerance = NonNegativeConfigTolerance + { nonNegativeConfigToleranceValue :: Double + } + deriving stock (Eq, Show) + +mkNonNegativeConfigTolerance :: Double -> Either KrylovConfigError NonNegativeConfigTolerance +mkNonNegativeConfigTolerance = + mkNonNegativeFiniteWith + KrylovConfigToleranceMustBeNonNegative + NonNegativeConfigTolerance + +type ArnoldiConfig :: Type +data ArnoldiConfig = ArnoldiConfig + { arnoldiIterationCount :: !PositiveCount, + arnoldiToleranceBound :: !NonNegativeConfigTolerance, + arnoldiReorthogonalize :: !Bool + } + deriving stock (Eq) + +instance Show ArnoldiConfig where + showsPrec precedence config = + showParen (precedence > 10) $ + showString "ArnoldiConfig {arnoldiIterations = " + . shows (arnoldiIterations config) + . showString ", arnoldiTolerance = " + . shows (arnoldiTolerance config) + . showString ", arnoldiReorthogonalize = " + . shows (arnoldiReorthogonalize config) + . showString "}" + +arnoldiIterations :: ArnoldiConfig -> Int +arnoldiIterations = + positiveCountValue . arnoldiIterationCount + +arnoldiTolerance :: ArnoldiConfig -> Double +arnoldiTolerance = + nonNegativeConfigToleranceValue . arnoldiToleranceBound + +mkArnoldiConfig :: PositiveCount -> NonNegativeConfigTolerance -> Bool -> ArnoldiConfig +mkArnoldiConfig iterations tolerance reorthogonalize = + ArnoldiConfig + { arnoldiIterationCount = iterations, + arnoldiToleranceBound = tolerance, + arnoldiReorthogonalize = reorthogonalize + } + +withArnoldiIterations :: PositiveCount -> ArnoldiConfig -> ArnoldiConfig +withArnoldiIterations iterations config = + mkArnoldiConfig + iterations + (arnoldiToleranceBound config) + (arnoldiReorthogonalize config) + +withArnoldiTolerance :: NonNegativeConfigTolerance -> ArnoldiConfig -> ArnoldiConfig +withArnoldiTolerance tolerance config = + mkArnoldiConfig + (arnoldiIterationCount config) + tolerance + (arnoldiReorthogonalize config) + +withArnoldiReorthogonalize :: Bool -> ArnoldiConfig -> ArnoldiConfig +withArnoldiReorthogonalize reorthogonalize config = + mkArnoldiConfig + (arnoldiIterationCount config) + (arnoldiToleranceBound config) + reorthogonalize + +defaultArnoldiConfig :: ArnoldiConfig +defaultArnoldiConfig = + mkArnoldiConfig (PositiveCount 64) (NonNegativeConfigTolerance 1.0e-10) True + +type LanczosConfig :: Type +data LanczosConfig = LanczosConfig + { lanczosIterationCount :: !PositiveCount, + lanczosToleranceBound :: !NonNegativeConfigTolerance + } + deriving stock (Eq) + +instance Show LanczosConfig where + showsPrec precedence config = + showParen (precedence > 10) $ + showString "LanczosConfig {lanczosIterations = " + . shows (lanczosIterations config) + . showString ", lanczosTolerance = " + . shows (lanczosTolerance config) + . showString "}" + +lanczosIterations :: LanczosConfig -> Int +lanczosIterations = + positiveCountValue . lanczosIterationCount + +lanczosTolerance :: LanczosConfig -> Double +lanczosTolerance = + nonNegativeConfigToleranceValue . lanczosToleranceBound + +mkLanczosConfig :: PositiveCount -> NonNegativeConfigTolerance -> LanczosConfig +mkLanczosConfig iterations tolerance = + LanczosConfig + { lanczosIterationCount = iterations, + lanczosToleranceBound = tolerance + } + +withLanczosIterations :: PositiveCount -> LanczosConfig -> LanczosConfig +withLanczosIterations iterations config = + mkLanczosConfig + iterations + (lanczosToleranceBound config) + +withLanczosTolerance :: NonNegativeConfigTolerance -> LanczosConfig -> LanczosConfig +withLanczosTolerance tolerance config = + mkLanczosConfig + (lanczosIterationCount config) + tolerance + +-- | The 24-vector window is the measured minimum-allocation basin for the +-- generic sparse fallback; callers retain an explicit override. +defaultLanczosConfig :: LanczosConfig +defaultLanczosConfig = + mkLanczosConfig (PositiveCount 24) (NonNegativeConfigTolerance 1.0e-10) + +type BlockLanczosConfig :: Type +data BlockLanczosConfig = BlockLanczosConfig + { blockLanczosIterationCount :: !PositiveCount, + blockLanczosToleranceBound :: !NonNegativeConfigTolerance, + blockLanczosConfiguredBlockSize :: !PositiveCount, + blockLanczosReorthogonalize :: !Bool + } + deriving stock (Eq) + +instance Show BlockLanczosConfig where + showsPrec precedence config = + showParen (precedence > 10) $ + showString "BlockLanczosConfig {blockLanczosIterations = " + . shows (blockLanczosIterations config) + . showString ", blockLanczosTolerance = " + . shows (blockLanczosTolerance config) + . showString ", blockLanczosBlockSize = " + . shows (blockLanczosBlockSize config) + . showString ", blockLanczosReorthogonalize = " + . shows (blockLanczosReorthogonalize config) + . showString "}" + +blockLanczosIterations :: BlockLanczosConfig -> Int +blockLanczosIterations = + positiveCountValue . blockLanczosIterationCount + +blockLanczosTolerance :: BlockLanczosConfig -> Double +blockLanczosTolerance = + nonNegativeConfigToleranceValue . blockLanczosToleranceBound + +blockLanczosBlockSize :: BlockLanczosConfig -> Int +blockLanczosBlockSize = + positiveCountValue . blockLanczosConfiguredBlockSize + +mkBlockLanczosConfig :: + PositiveCount -> + NonNegativeConfigTolerance -> + PositiveCount -> + Bool -> + BlockLanczosConfig +mkBlockLanczosConfig iterations tolerance blockSize reorthogonalize = + BlockLanczosConfig + { blockLanczosIterationCount = iterations, + blockLanczosToleranceBound = tolerance, + blockLanczosConfiguredBlockSize = blockSize, + blockLanczosReorthogonalize = reorthogonalize + } + +withBlockLanczosIterations :: PositiveCount -> BlockLanczosConfig -> BlockLanczosConfig +withBlockLanczosIterations iterations config = + mkBlockLanczosConfig + iterations + (blockLanczosToleranceBound config) + (blockLanczosConfiguredBlockSize config) + (blockLanczosReorthogonalize config) + +withBlockLanczosTolerance :: NonNegativeConfigTolerance -> BlockLanczosConfig -> BlockLanczosConfig +withBlockLanczosTolerance tolerance config = + mkBlockLanczosConfig + (blockLanczosIterationCount config) + tolerance + (blockLanczosConfiguredBlockSize config) + (blockLanczosReorthogonalize config) + +withBlockLanczosBlockSize :: PositiveCount -> BlockLanczosConfig -> BlockLanczosConfig +withBlockLanczosBlockSize blockSize config = + mkBlockLanczosConfig + (blockLanczosIterationCount config) + (blockLanczosToleranceBound config) + blockSize + (blockLanczosReorthogonalize config) + +withBlockLanczosReorthogonalize :: Bool -> BlockLanczosConfig -> BlockLanczosConfig +withBlockLanczosReorthogonalize reorthogonalize config = + mkBlockLanczosConfig + (blockLanczosIterationCount config) + (blockLanczosToleranceBound config) + (blockLanczosConfiguredBlockSize config) + reorthogonalize + +defaultBlockLanczosConfig :: BlockLanczosConfig +defaultBlockLanczosConfig = + mkBlockLanczosConfig + (PositiveCount 48) + (NonNegativeConfigTolerance 1.0e-10) + (PositiveCount 2) + True diff --git a/src-spectral/Moonlight/LinAlg/Pure/Krylov/Decomposition.hs b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Decomposition.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Decomposition.hs @@ -0,0 +1,156 @@ +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Krylov.Decomposition + ( ArnoldiDecomposition, + mkArnoldiDecomposition, + arnoldiBasisColumns, + arnoldiHessenbergRows, + arnoldiStepsCompleted, + LanczosDecomposition, + mkLanczosDecomposition, + lanczosBasisColumns, + lanczosProjectedTridiagonal, + lanczosAlphaDiagonal, + lanczosBetaOffDiagonal, + lanczosResidualNorm, + lanczosStepsCompleted, + BlockLanczosDecomposition, + mkBlockLanczosDecomposition, + blockLanczosBasisColumns, + blockLanczosProjectedBlockTridiagonal, + blockLanczosBasisCount, + blockLanczosBlockSteps, + ) +where + +import Data.Kind (Type) +import Data.Vector qualified as Box +import Data.Vector.Unboxed qualified as U +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Pure.Structured.BlockTridiagonal + ( SymmetricBlockTridiagonal, + symmetricBlockTridiagonalDimension, + ) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal + ( SymmetricTridiagonal, + symmetricTridiagonalDiagonalEntries, + symmetricTridiagonalDimension, + symmetricTridiagonalOffDiagonalEntries, + ) +import Prelude + +type ArnoldiDecomposition :: Type +data ArnoldiDecomposition = ArnoldiDecomposition + { arnoldiBasisColumns :: !(Box.Vector (U.Vector Double)), + arnoldiHessenbergRows :: !(Box.Vector (U.Vector Double)) + } + deriving stock (Eq, Show) + +mkArnoldiDecomposition :: Box.Vector (U.Vector Double) -> Box.Vector (U.Vector Double) -> Either MoonlightError ArnoldiDecomposition +mkArnoldiDecomposition basisColumns hessenbergRows = do + stepCount <- validateBasisColumns "Arnoldi" basisColumns + if Box.length hessenbergRows /= stepCount + 1 + then + Left + ( InvariantViolation + ( "Arnoldi Hessenberg row count mismatch: expected " + <> show (stepCount + 1) + <> " but received " + <> show (Box.length hessenbergRows) + ) + ) + else + if any ((/= stepCount) . U.length) (Box.toList hessenbergRows) + then Left (InvariantViolation "Arnoldi Hessenberg rows must have equal length matching the step count") + else Right (ArnoldiDecomposition basisColumns hessenbergRows) + +arnoldiStepsCompleted :: ArnoldiDecomposition -> Int +arnoldiStepsCompleted = Box.length . arnoldiBasisColumns + +type LanczosDecomposition :: Type +data LanczosDecomposition = LanczosDecomposition + { lanczosBasisColumns :: !(Box.Vector (U.Vector Double)), + lanczosProjectedTridiagonal :: !SymmetricTridiagonal, + lanczosResidualNorm :: !Double + } + deriving stock (Eq, Show) + +mkLanczosDecomposition :: + Box.Vector (U.Vector Double) -> + SymmetricTridiagonal -> + Double -> + Either MoonlightError LanczosDecomposition +mkLanczosDecomposition basisColumns projectedTridiagonal residualNorm = do + basisCount <- validateBasisColumns "Lanczos" basisColumns + let projectedDimension = symmetricTridiagonalDimension projectedTridiagonal + if projectedDimension /= basisCount + then + Left + ( InvariantViolation + ( "Lanczos projected tridiagonal dimension mismatch: expected " + <> show basisCount + <> " but received " + <> show projectedDimension + ) + ) + else Right (LanczosDecomposition basisColumns projectedTridiagonal residualNorm) + +lanczosAlphaDiagonal :: LanczosDecomposition -> [Double] +lanczosAlphaDiagonal = symmetricTridiagonalDiagonalEntries . lanczosProjectedTridiagonal + +lanczosBetaOffDiagonal :: LanczosDecomposition -> [Double] +lanczosBetaOffDiagonal = symmetricTridiagonalOffDiagonalEntries . lanczosProjectedTridiagonal + +lanczosStepsCompleted :: LanczosDecomposition -> Int +lanczosStepsCompleted = + symmetricTridiagonalDimension . lanczosProjectedTridiagonal + +type BlockLanczosDecomposition :: Type +data BlockLanczosDecomposition = BlockLanczosDecomposition + { blockLanczosBasisColumns :: !(Box.Vector (U.Vector Double)), + blockLanczosProjectedBlockTridiagonal :: !SymmetricBlockTridiagonal, + blockLanczosBlockSteps :: !Int + } + deriving stock (Eq, Show) + +mkBlockLanczosDecomposition :: + Box.Vector (U.Vector Double) -> + SymmetricBlockTridiagonal -> + Int -> + Either MoonlightError BlockLanczosDecomposition +mkBlockLanczosDecomposition basisColumns projectedBlockTridiagonal blockStepCount = do + basisCount <- validateBasisColumns "Block Lanczos" basisColumns + let projectedDimension = symmetricBlockTridiagonalDimension projectedBlockTridiagonal + if blockStepCount <= 0 + then Left (InvariantViolation "Block Lanczos step count must be positive") + else + if projectedDimension /= basisCount + then + Left + ( InvariantViolation + ( "Block Lanczos projected operator dimension mismatch: expected " + <> show basisCount + <> " but received " + <> show projectedDimension + ) + ) + else Right (BlockLanczosDecomposition basisColumns projectedBlockTridiagonal blockStepCount) + +blockLanczosBasisCount :: BlockLanczosDecomposition -> Int +blockLanczosBasisCount = + symmetricBlockTridiagonalDimension . blockLanczosProjectedBlockTridiagonal + +validateBasisColumns :: String -> Box.Vector (U.Vector Double) -> Either MoonlightError Int +validateBasisColumns algorithmName basisColumns = + let basisCount = Box.length basisColumns + basisDimensions = U.length <$> Box.toList basisColumns + basisDimension = + case basisDimensions of + [] -> 0 + firstDimension : _ -> firstDimension + in if basisCount <= 0 + then Left (InvariantViolation (algorithmName <> " basis must be non-empty")) + else + if any (/= basisDimension) basisDimensions + then Left (InvariantViolation (algorithmName <> " basis columns must have equal length")) + else Right basisCount diff --git a/src-spectral/Moonlight/LinAlg/Pure/Krylov/Internal.hs b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Internal.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Internal.hs @@ -0,0 +1,387 @@ +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Krylov.Internal + ( validateSquareOperator, + validateIterationCount, + normalizeSeed, + normalizeSeedBlock, + unitVector, + requireBasisVector, + orthogonalizeAgainst, + orthonormalizeBlock, + blockInnerBlock, + selfAdjointBlockInnerBlock, + linearCombinationColumnsU, + multiplyBasisByBlock, + subtractBlocks, + sparseColumnsToDenseRowVectors, + ) +where + +import Control.Monad.ST (ST, runST) +import qualified Data.Vector as Box +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as MU +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Internal.VectorOps (dotU, normU, scaleU, subScaledU, subU) +import Moonlight.LinAlg.Pure.Operator +import Moonlight.LinAlg.Pure.Structured.BlockTridiagonal + ( RowMajorBlock, + mkRowMajorBlock, + rowMajorBlockColumns, + rowMajorBlockEntry, + rowMajorBlockRows, + symmetrizeRowMajorBlockLower, + ) +import Prelude + +validateSquareOperator :: String -> LinearOperator symmetry -> Either MoonlightError () +validateSquareOperator algorithmName op + | rows <= 0 || cols <= 0 = + Left (InvariantViolation (algorithmName <> " requires a positive square operator")) + | rows /= cols = + Left + ( InvariantViolation + ( algorithmName + <> " requires a square operator, but received " + <> show (rows, cols) + ) + ) + | otherwise = Right () + where + (rows, cols) = operatorShape op + +validateIterationCount :: String -> Int -> Either MoonlightError Int +validateIterationCount algorithmName iterationCount + | iterationCount <= 0 = + Left (InvariantViolation (algorithmName <> " iteration count must be positive")) + | otherwise = Right iterationCount + +normalizeSeed :: String -> Int -> Double -> U.Vector Double -> Either MoonlightError (U.Vector Double) +normalizeSeed algorithmName dimension tolerance seedVector + | U.null seedVector = Left (InvariantViolation (algorithmName <> " requires a non-empty seed vector")) + | U.length seedVector /= dimension = + Left + ( InvariantViolation + ( algorithmName + <> " seed length mismatch: expected " + <> show dimension + <> " but received " + <> show (U.length seedVector) + ) + ) + | otherwise = + let seedNorm = normU seedVector + in if seedNorm <= tolerance + then Left (InvariantViolation (algorithmName <> " seed vector is near-zero; provide a non-trivial seed")) + else Right (scaleU (1.0 / seedNorm) seedVector) + +normalizeSeedBlock :: + String -> + Int -> + Double -> + Int -> + Box.Vector (U.Vector Double) -> + Either MoonlightError (Box.Vector (U.Vector Double)) +normalizeSeedBlock algorithmName dimension tolerance blockSize seedBlock + | blockSize <= 0 = + Left (InvariantViolation (algorithmName <> " block size must be positive")) + | otherwise = do + let candidateSeeds = Box.take blockSize (seedBlock <> canonicalSeeds dimension) + acceptedSeeds <- + U.foldM' (acceptSeed candidateSeeds) Box.empty (U.enumFromN 0 (Box.length candidateSeeds)) + if Box.null acceptedSeeds + then Left (InvariantViolation (algorithmName <> " could not construct a non-zero orthogonal seed block")) + else Right acceptedSeeds + where + acceptSeed candidateSeeds acceptedVectors seedIndex = + case candidateSeeds Box.!? seedIndex of + Nothing -> Left (InvariantViolation "block seed lookup failed") + Just seedVector -> do + normalizedSeed <- normalizeSeed algorithmName dimension tolerance seedVector + (reducedSeed, _) <- orthogonalizeAgainst True acceptedVectors normalizedSeed + let reducedNorm = normU reducedSeed + pure + ( if reducedNorm <= tolerance + then acceptedVectors + else acceptedVectors `Box.snoc` scaleU (1.0 / reducedNorm) reducedSeed + ) + +canonicalSeeds :: Int -> Box.Vector (U.Vector Double) +canonicalSeeds dimension = + Box.generate dimension (unitVector dimension) + +unitVector :: Int -> Int -> U.Vector Double +unitVector dimension selectedIndex = + U.generate dimension + (\indexValue -> if indexValue == selectedIndex then 1.0 else 0.0) + +requireBasisVector :: Int -> Box.Vector (U.Vector Double) -> Either MoonlightError (U.Vector Double) +requireBasisVector indexValue basisVectors = + maybe + (Left (InvariantViolation ("Krylov basis lookup failed at index " <> show indexValue))) + Right + (basisVectors Box.!? indexValue) + +orthogonalizeAgainst :: + Bool -> + Box.Vector (U.Vector Double) -> + U.Vector Double -> + Either MoonlightError (U.Vector Double, U.Vector Double) +orthogonalizeAgainst reorthogonalize basisVectors inputVector = + do + (reducedOnce, coefficientsOnce) <- projectOnce basisVectors inputVector + if reorthogonalize + then do + (reducedTwice, coefficientsTwice) <- projectOnce basisVectors reducedOnce + coefficients <- addCoefficientVectors coefficientsOnce coefficientsTwice + Right (reducedTwice, coefficients) + else Right (reducedOnce, coefficientsOnce) + +projectOnce :: + Box.Vector (U.Vector Double) -> + U.Vector Double -> + Either MoonlightError (U.Vector Double, U.Vector Double) +projectOnce basisVectors inputVector = + runST $ do + coefficientValues <- MU.replicate (Box.length basisVectors) 0.0 + projectedVector <- + U.foldM' + (projectBasisIndex basisVectors coefficientValues) + (Right inputVector) + (U.enumFromN 0 (Box.length basisVectors)) + case projectedVector of + Left err -> pure (Left err) + Right reducedVector -> do + frozenCoefficients <- U.freeze coefficientValues + pure (Right (reducedVector, frozenCoefficients)) + +projectBasisIndex :: + Box.Vector (U.Vector Double) -> + MU.MVector s Double -> + Either MoonlightError (U.Vector Double) -> + Int -> + ST s (Either MoonlightError (U.Vector Double)) +projectBasisIndex basisVectors coefficientValues projectedVector basisIndex = + case projectedVector of + Left err -> pure (Left err) + Right workingVector -> + case basisVectors Box.!? basisIndex of + Nothing -> pure (Left (InvariantViolation ("Krylov basis lookup failed at index " <> show basisIndex))) + Just basisVector -> + case dotU basisVector workingVector of + Left err -> pure (Left err) + Right coefficient -> + case subScaledU workingVector coefficient basisVector of + Left err -> pure (Left err) + Right nextVector -> do + MU.unsafeWrite coefficientValues basisIndex coefficient + pure (Right nextVector) + +addCoefficientVectors :: + U.Vector Double -> + U.Vector Double -> + Either MoonlightError (U.Vector Double) +addCoefficientVectors left right = + if U.length left == U.length right + then Right (U.zipWith (+) left right) + else + Left + ( InvariantViolation + ( "coefficient vector length mismatch: left " + <> show (U.length left) + <> " right " + <> show (U.length right) + ) + ) + +orthonormalizeBlock :: + Bool -> + Double -> + Box.Vector (U.Vector Double) -> + Box.Vector (U.Vector Double) -> + Either MoonlightError (Box.Vector (U.Vector Double)) +orthonormalizeBlock reorthogonalize tolerance existingBasis candidateVectors = + U.foldM' acceptCandidateIndex Box.empty (U.enumFromN 0 (Box.length candidateVectors)) + where + acceptCandidateIndex acceptedVectors candidateIndex = + case candidateVectors Box.!? candidateIndex of + Nothing -> Left (InvariantViolation "block candidate lookup failed") + Just candidateVector -> do + let combinedBasis = existingBasis <> acceptedVectors + (reducedVector, _) <- orthogonalizeAgainst reorthogonalize combinedBasis candidateVector + let reducedNorm = normU reducedVector + Right + ( if reducedNorm <= tolerance + then acceptedVectors + else acceptedVectors `Box.snoc` scaleU (1.0 / reducedNorm) reducedVector + ) + +blockInnerBlock :: + Box.Vector (U.Vector Double) -> + Box.Vector (U.Vector Double) -> + Either MoonlightError RowMajorBlock +blockInnerBlock leftBasis rightBasis = do + payload <- + U.generateM + (Box.length leftBasis * Box.length rightBasis) + ( \payloadIndex -> + let rightCount = Box.length rightBasis + leftIndex = payloadIndex `quot` rightCount + rightIndex = payloadIndex `rem` rightCount + in case (leftBasis Box.!? leftIndex, rightBasis Box.!? rightIndex) of + (Just leftVector, Just rightVector) -> dotU leftVector rightVector + _ -> Left (InvariantViolation "block inner-product index out of bounds") + ) + mkRowMajorBlock + (Box.length leftBasis) + (Box.length rightBasis) + payload + +selfAdjointBlockInnerBlock :: + Box.Vector (U.Vector Double) -> + Box.Vector (U.Vector Double) -> + Either MoonlightError RowMajorBlock +selfAdjointBlockInnerBlock leftBasis rightBasis = + blockInnerBlock leftBasis rightBasis >>= symmetrizeRowMajorBlockLower + +linearCombinationColumnsU :: + Box.Vector (U.Vector Double) -> + U.Vector Double -> + Either MoonlightError (U.Vector Double) +linearCombinationColumnsU basisColumns coefficients = + case basisColumns Box.!? 0 of + Nothing -> + Left (InvariantViolation "basis linear combination requires a non-empty basis") + Just firstColumn -> + let basisCount = Box.length basisColumns + coefficientCount = U.length coefficients + ambientDimension = U.length firstColumn + in if coefficientCount /= basisCount + then + Left + ( InvariantViolation + ( "basis linear combination coefficient count mismatch: expected " + <> show basisCount + <> " but received " + <> show coefficientCount + ) + ) + else do + _ <- Box.ifoldM' (validateLinearCombinationColumn ambientDimension) () basisColumns + Right + ( U.generate + ambientDimension + ( \entryIndex -> + Box.ifoldl' + (accumulateLinearCombinationEntry coefficients entryIndex) + 0.0 + basisColumns + ) + ) + +validateLinearCombinationColumn :: + Int -> + () -> + Int -> + U.Vector Double -> + Either MoonlightError () +validateLinearCombinationColumn ambientDimension () columnIndex columnValue = + let actualDimension = U.length columnValue + in if actualDimension == ambientDimension + then Right () + else + Left + ( InvariantViolation + ( "basis linear combination column " + <> show columnIndex + <> " has dimension " + <> show actualDimension + <> " but expected " + <> show ambientDimension + ) + ) + +accumulateLinearCombinationEntry :: + U.Vector Double -> + Int -> + Double -> + Int -> + U.Vector Double -> + Double +accumulateLinearCombinationEntry coefficients entryIndex accumulator columnIndex columnVector = + accumulator + + coefficients `U.unsafeIndex` columnIndex + * columnVector `U.unsafeIndex` entryIndex + +multiplyBasisByBlock :: + Box.Vector (U.Vector Double) -> + RowMajorBlock -> + Either MoonlightError (Box.Vector (U.Vector Double)) +multiplyBasisByBlock basisVectors coefficientBlock + | Box.length basisVectors /= rowMajorBlockRows coefficientBlock = + Left (InvariantViolation "basis/vector coefficient block row count mismatch") + | Box.null basisVectors = + Left (InvariantViolation "basis/vector coefficient block requires a non-empty basis") + | otherwise = + Right + ( Box.generate + (rowMajorBlockColumns coefficientBlock) + (basisCombinationColumn basisVectors coefficientBlock) + ) + +basisCombinationColumn :: Box.Vector (U.Vector Double) -> RowMajorBlock -> Int -> U.Vector Double +basisCombinationColumn basisVectors coefficientBlock outputColumn = + let ambientDimension = + maybe 0 U.length (basisVectors Box.!? 0) + in U.generate + ambientDimension + ( \entryIndex -> + Box.ifoldl' + ( \accumulator basisIndex basisVector -> + accumulator + + rowMajorBlockEntry coefficientBlock basisIndex outputColumn + * maybe 0.0 id (basisVector U.!? entryIndex) + ) + 0.0 + basisVectors + ) + +subtractBlocks :: + Box.Vector (U.Vector Double) -> + Box.Vector (U.Vector Double) -> + Either MoonlightError (Box.Vector (U.Vector Double)) +subtractBlocks leftBlocks rightBlocks + | Box.length leftBlocks /= Box.length rightBlocks = + Left + ( InvariantViolation + ( "block vector count mismatch: left " + <> show (Box.length leftBlocks) + <> " right " + <> show (Box.length rightBlocks) + ) + ) + | otherwise = + Box.generateM + (Box.length leftBlocks) + ( \blockIndex -> + case (leftBlocks Box.!? blockIndex, rightBlocks Box.!? blockIndex) of + (Just leftBlock, Just rightBlock) -> subU leftBlock rightBlock + _ -> Left (InvariantViolation "block vector lookup failed") + ) + +sparseColumnsToDenseRowVectors :: Int -> Int -> Box.Vector (U.Vector Double) -> Box.Vector (U.Vector Double) +sparseColumnsToDenseRowVectors rowCount columnCount columnValues = + Box.generate + rowCount + ( \rowIndex -> + U.generate + columnCount + (sparseEntryAt rowIndex) + ) + where + sparseEntryAt rowIndex columnIndex = + maybe + 0.0 + (\columnValue -> maybe 0.0 id (columnValue U.!? rowIndex)) + (columnValues Box.!? columnIndex) diff --git a/src-spectral/Moonlight/LinAlg/Pure/Krylov/Lanczos.hs b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Lanczos.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Lanczos.hs @@ -0,0 +1,1314 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Krylov.Lanczos + ( lanczosSymmetric, + LanczosRestartProjection, + lanczosRestartProjectionBasisColumns, + lanczosRestartProjectionProjectedPairs, + lanczosRestartedProjection, + ritzLockThreshold, + ) +where + +import Control.Monad (foldM) +import Control.Monad.ST (ST, runST) +import Data.Either (partitionEithers) +import Data.Foldable (traverse_) +import Data.Maybe (catMaybes, listToMaybe) +import Data.Primitive.PrimArray + ( MutablePrimArray, + newPrimArray, + readPrimArray, + setPrimArray, + writePrimArray, + ) +import qualified Data.Vector as Box +import qualified Data.Vector.Mutable as BoxM +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as UM +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Internal.Eigen.Kernels + ( epsDouble, + finiteDouble, + forDescendingIndex, + forIndex, + hypotStable, + ) +import Moonlight.LinAlg.Internal.VectorOps (dotU, normU, scaleU, subScaledU) +import Moonlight.LinAlg.Pure.Krylov.Config (LanczosConfig, lanczosIterations, lanczosTolerance) +import Moonlight.LinAlg.Pure.Krylov.Decomposition (LanczosDecomposition, mkLanczosDecomposition) +import Moonlight.LinAlg.Pure.Krylov.Internal + ( linearCombinationColumnsU, + normalizeSeed, + validateIterationCount, + validateSquareOperator, + ) +import Moonlight.LinAlg.Pure.Krylov.Selection (SpectrumEnd (..), sortForSpectrumBy) +import Moonlight.LinAlg.Pure.Krylov.SelectedTridiagonal (inverseIterationResidualToleranceBound, selectedSymmetricTridiagonalEigenpairColumnsDirect) +import Moonlight.LinAlg.Pure.Operator (LinearOperator, OperatorSymmetry (SelfAdjointOperator), operatorShape, runOperatorU) +import Moonlight.LinAlg.Pure.Spectral.Result (Eigenpairs, eigenpairCount, eigenpairResidualNorms, eigenpairValues, eigenpairVectorAt, eigenpairsFromColumns) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal (SymmetricTridiagonal, mkSymmetricTridiagonal, mkSymmetricTridiagonalVectors) +import Prelude + +newtype ActiveDimension = ActiveDimension + { activeDimensionValue :: Int + } + deriving stock (Eq, Show) + +data LanczosState + = LanczosRunning !ActiveDimension !Double !(U.Vector Double) !(U.Vector Double) + | LanczosConverged !ActiveDimension !Double + | LanczosBreakdown !ActiveDimension !Double + | LanczosRestarting !ActiveDimension !Double + deriving stock (Eq, Show) + +data LanczosArena s = LanczosArena + { lanczosBasisArena :: !(BoxM.MVector s (U.Vector Double)), + lanczosAlphaArena :: !(UM.MVector s Double), + lanczosBetaArena :: !(UM.MVector s Double) + } + +data LanczosRestartProjection = LanczosRestartProjection + { lanczosRestartProjectionBasisColumns :: !(Box.Vector (U.Vector Double)), + lanczosRestartProjectionProjectedPairs :: !Eigenpairs + } + deriving stock (Eq, Show) + +data RitzPair = RitzPair + { ritzPairValue :: !Double, + ritzPairVector :: !(U.Vector Double), + ritzPairResidualNorm :: !Double, + ritzPairProjectedResidualNorm :: !Double, + ritzPairBoundaryCoupling :: !Double + } + deriving stock (Eq, Show) + +data RitzCandidate = RitzCandidate + { ritzCandidateValue :: !Double, + ritzCandidateProjectedVector :: !(U.Vector Double), + ritzCandidateProjectedResidualNorm :: !Double, + ritzCandidateBoundaryCoupling :: !Double + } + deriving stock (Eq, Show) + +data RestartSeed = RestartSeed + { restartSeedBasisColumns :: !(Box.Vector (U.Vector Double)), + restartSeedRetainedValues :: !(U.Vector Double), + restartSeedSpikeCouplings :: !(U.Vector Double) + } + deriving stock (Eq, Show) + +data RestartContext = RestartContext + { restartLockedPairs :: ![RitzPair], + restartSeed :: !RestartSeed, + restartRetainedPairs :: ![RitzPair] + } + deriving stock (Eq, Show) + +data ExpandedWindow = ExpandedWindow + { expandedWindowBasisColumns :: !(Box.Vector (U.Vector Double)), + expandedWindowProjectedOperator :: !BorderedProjectedOperator, + expandedWindowBoundaryResidualNorm :: !Double, + expandedWindowBoundaryVector :: !(Maybe (U.Vector Double)), + expandedWindowState :: !LanczosState + } + deriving stock (Eq, Show) + +data BorderedProjectedOperator = BorderedProjectedOperator + { borderedRetainedValues :: !(U.Vector Double), + borderedSpikeCouplings :: !(U.Vector Double), + borderedKrylovDiagonal :: !(U.Vector Double), + borderedKrylovOffDiagonal :: !(U.Vector Double) + } + deriving stock (Eq, Show) + +data BorderedProjectionReduction = BorderedProjectionReduction + { borderedReductionBasisColumns :: !(Box.Vector (U.Vector Double)), + borderedReductionTridiagonal :: !SymmetricTridiagonal + } + deriving stock (Eq, Show) + +data BorderedReductionArena s = BorderedReductionArena + { borderedReductionArenaDimension :: !Int, + borderedReductionArenaPayload :: !(MutablePrimArray s Double), + borderedReductionArenaMatrixOffset :: !Int, + borderedReductionArenaBasisOffset :: !Int + } + +data RestartSeedResult = RestartSeedResult + { restartSeedResultSeed :: !RestartSeed, + restartSeedResultRetainedPairs :: ![RitzPair] + } + deriving stock (Eq, Show) + +lanczosSymmetric :: LanczosConfig -> LinearOperator 'SelfAdjointOperator -> U.Vector Double -> Either MoonlightError LanczosDecomposition +lanczosSymmetric config op seedVector = do + validateSquareOperator "Lanczos" op + let (_, cols) = operatorShape op + targetIterations <- validateIterationCount "Lanczos" (lanczosIterations config) + firstBasis <- normalizeSeed "Lanczos" cols (lanczosTolerance config) seedVector + let boundedIterations = min targetIterations cols + zeroVector = U.replicate cols 0.0 + tolerance = lanczosTolerance config + in runST $ do + arena <- newLanczosArena boundedIterations + BoxM.unsafeWrite (lanczosBasisArena arena) 0 firstBasis + finalStateResult <- + runLanczosState + op + tolerance + boundedIterations + arena + (LanczosRunning (ActiveDimension 1) 0.0 zeroVector firstBasis) + case finalStateResult of + Left err -> pure (Left err) + Right finalState -> freezeLanczosState arena finalState + +lanczosRestartedProjection :: + LanczosConfig -> + SpectrumEnd -> + Int -> + LinearOperator 'SelfAdjointOperator -> + U.Vector Double -> + Either MoonlightError LanczosRestartProjection +lanczosRestartedProjection config spectrumEnd requestedCount op seedVector + | requestedCount <= 0 = + Left (InvariantViolation "restarted Lanczos requires a positive requested count") + | otherwise = do + validateSquareOperator "restarted Lanczos" op + let (_, cols) = operatorShape op + tolerance = lanczosTolerance config + if requestedCount > cols + then Left (InvariantViolation "restarted Lanczos requested count exceeds operator dimension") + else do + targetIterations <- validateIterationCount "restarted Lanczos" (lanczosIterations config) + firstBasis <- normalizeSeed "restarted Lanczos" cols tolerance seedVector + let capacity = min targetIterations cols + restartLoop + op + spectrumEnd + requestedCount + tolerance + capacity + cols + (maxRestartCycles cols capacity) + (RestartContext [] (initialRestartSeed firstBasis) []) + +initialRestartSeed :: U.Vector Double -> RestartSeed +initialRestartSeed firstBasis = + RestartSeed + { restartSeedBasisColumns = Box.singleton firstBasis, + restartSeedRetainedValues = U.empty, + restartSeedSpikeCouplings = U.empty + } + +newLanczosArena :: Int -> ST s (LanczosArena s) +newLanczosArena capacity = do + basisArena <- BoxM.unsafeNew capacity + alphaArena <- UM.unsafeNew capacity + betaArena <- UM.unsafeNew (max 0 (capacity - 1)) + pure + LanczosArena + { lanczosBasisArena = basisArena, + lanczosAlphaArena = alphaArena, + lanczosBetaArena = betaArena + } + +runLanczosState :: + LinearOperator 'SelfAdjointOperator -> + Double -> + Int -> + LanczosArena s -> + LanczosState -> + ST s (Either MoonlightError LanczosState) +runLanczosState op tolerance capacity arena state = + case state of + LanczosRunning activeDimension previousBeta previousBasis currentBasis -> do + nextState <- stepLanczosState op tolerance capacity arena activeDimension previousBeta previousBasis currentBasis + case nextState of + Left err -> pure (Left err) + Right stateValue -> runLanczosState op tolerance capacity arena stateValue + LanczosConverged{} -> pure (Right state) + LanczosBreakdown{} -> pure (Right state) + LanczosRestarting{} -> pure (Right state) + +stepLanczosState :: + LinearOperator 'SelfAdjointOperator -> + Double -> + Int -> + LanczosArena s -> + ActiveDimension -> + Double -> + U.Vector Double -> + U.Vector Double -> + ST s (Either MoonlightError LanczosState) +stepLanczosState op tolerance capacity arena activeDimension previousBeta previousBasis currentBasis = + case runOperatorU op currentBasis of + Left err -> + pure (Left err) + Right imageVector -> + case removePreviousDirection imageVector of + Left err -> + pure (Left err) + Right withPreviousRemoved -> + case dotU currentBasis withPreviousRemoved of + Left err -> + pure (Left err) + Right alphaValue -> + case subScaledU withPreviousRemoved alphaValue currentBasis of + Left err -> + pure (Left err) + Right projectedCurrent -> do + residualResult <- orthogonalizeAgainstArena arena activeDimension projectedCurrent + case residualResult of + Left err -> + pure (Left err) + Right residualVector -> do + let activeCount = activeDimensionValue activeDimension + currentIndex = activeCount - 1 + betaValue = normU residualVector + UM.unsafeWrite (lanczosAlphaArena arena) currentIndex alphaValue + if betaValue <= tolerance + then pure (Right (LanczosConverged activeDimension betaValue)) + else + if activeCount >= capacity + then pure (Right (LanczosRestarting activeDimension betaValue)) + else do + let nextBasis = scaleU (1.0 / betaValue) residualVector + nextActiveDimension = ActiveDimension (activeCount + 1) + UM.unsafeWrite (lanczosBetaArena arena) currentIndex betaValue + BoxM.unsafeWrite (lanczosBasisArena arena) activeCount nextBasis + pure (Right (LanczosRunning nextActiveDimension betaValue currentBasis nextBasis)) + where + removePreviousDirection imageVector = + if activeDimensionValue activeDimension == 1 + then Right imageVector + else subScaledU imageVector previousBeta previousBasis + +orthogonalizeAgainstArena :: + LanczosArena s -> + ActiveDimension -> + U.Vector Double -> + ST s (Either MoonlightError (U.Vector Double)) +orthogonalizeAgainstArena arena activeDimension inputVector = do + reducedOnce <- projectAgainstArenaOnce arena activeDimension inputVector + case reducedOnce of + Left err -> pure (Left err) + Right reducedVector -> projectAgainstArenaOnce arena activeDimension reducedVector + +projectAgainstArenaOnce :: + LanczosArena s -> + ActiveDimension -> + U.Vector Double -> + ST s (Either MoonlightError (U.Vector Double)) +projectAgainstArenaOnce arena activeDimension inputVector = + projectBasisIndex 0 inputVector + where + activeCount = activeDimensionValue activeDimension + projectBasisIndex basisIndex workingVector + | basisIndex >= activeCount = pure (Right workingVector) + | otherwise = do + basisVector <- BoxM.unsafeRead (lanczosBasisArena arena) basisIndex + case dotU basisVector workingVector of + Left err -> pure (Left err) + Right coefficient -> + case subScaledU workingVector coefficient basisVector of + Left err -> pure (Left err) + Right nextVector -> projectBasisIndex (basisIndex + 1) nextVector + +expandRestartWindow :: + LinearOperator 'SelfAdjointOperator -> + Double -> + Int -> + Box.Vector (U.Vector Double) -> + RestartSeed -> + Either MoonlightError ExpandedWindow +expandRestartWindow op tolerance capacity lockedVectors seedValue + | capacity <= 0 = + Left (InvariantViolation "restarted Lanczos active capacity must be positive") + | Box.null (restartSeedBasisColumns seedValue) = + Left (InvariantViolation "restarted Lanczos requires a non-empty restart basis") + | Box.length (restartSeedBasisColumns seedValue) > capacity = + Left (InvariantViolation "restarted Lanczos restart seed exceeds the active capacity") + | U.length (restartSeedRetainedValues seedValue) /= U.length (restartSeedSpikeCouplings seedValue) = + Left (InvariantViolation "restarted Lanczos retained Ritz values must match spike couplings") + | Box.length (restartSeedBasisColumns seedValue) /= U.length (restartSeedRetainedValues seedValue) + 1 = + Left (InvariantViolation "restarted Lanczos seed basis must contain retained vectors plus one Krylov boundary vector") + | otherwise = + runST $ do + arena <- newLanczosArena capacity + let seedBasis = restartSeedBasisColumns seedValue + seedCount = Box.length seedBasis + traverse_ (writeSeedBasis arena) (zip [0 :: Int ..] (Box.toList seedBasis)) + case reverse (Box.toList seedBasis) of + [] -> pure (Left (InvariantViolation "restarted Lanczos requires a non-empty bounded restart basis")) + currentBasis : _ -> + expandRestartFirstKrylovState op tolerance capacity lockedVectors seedValue arena seedCount currentBasis + +writeSeedBasis :: LanczosArena s -> (Int, U.Vector Double) -> ST s () +writeSeedBasis arena (basisIndex, basisVector) = + BoxM.unsafeWrite (lanczosBasisArena arena) basisIndex basisVector + +expandRestartFirstKrylovState :: + LinearOperator 'SelfAdjointOperator -> + Double -> + Int -> + Box.Vector (U.Vector Double) -> + RestartSeed -> + LanczosArena s -> + Int -> + U.Vector Double -> + ST s (Either MoonlightError ExpandedWindow) +expandRestartFirstKrylovState op tolerance capacity lockedVectors seedValue arena activeCount currentBasis = + case runOperatorU op currentBasis of + Left err -> pure (Left err) + Right imageVector -> + case removeRetainedDirections imageVector of + Left err -> pure (Left err) + Right withRetainedRemoved -> + case dotU currentBasis withRetainedRemoved of + Left err -> pure (Left err) + Right alphaValue -> + case subScaledU withRetainedRemoved alphaValue currentBasis of + Left err -> pure (Left err) + Right projectedCurrent -> do + residualResult <- + orthogonalizeAgainstLockedAndArena + lockedVectors + arena + (ActiveDimension activeCount) + projectedCurrent + case residualResult of + Left err -> + pure (Left err) + Right residualVector -> do + let currentIndex = activeCount - 1 + betaValue = normU residualVector + UM.unsafeWrite (lanczosAlphaArena arena) currentIndex alphaValue + if betaValue <= tolerance + then + freezeExpandedWindow + seedValue + arena + (LanczosConverged (ActiveDimension activeCount) betaValue) + Nothing + else + let nextBasis = scaleU (1.0 / betaValue) residualVector + in if activeCount >= capacity + then + freezeExpandedWindow + seedValue + arena + (LanczosRestarting (ActiveDimension activeCount) betaValue) + (Just nextBasis) + else do + BoxM.unsafeWrite (lanczosBasisArena arena) activeCount nextBasis + UM.unsafeWrite (lanczosBetaArena arena) currentIndex betaValue + expandRestartState + op + tolerance + capacity + lockedVectors + seedValue + arena + (activeCount + 1) + betaValue + currentBasis + nextBasis + where + retainedBasis = Box.take (U.length (restartSeedRetainedValues seedValue)) (restartSeedBasisColumns seedValue) + retainedCouplings = U.toList (restartSeedSpikeCouplings seedValue) + removeRetainedDirections imageVector = + foldM + (\workingVector (basisVector, couplingValue) -> subScaledU workingVector couplingValue basisVector) + imageVector + (zip (Box.toList retainedBasis) retainedCouplings) + +expandRestartState :: + LinearOperator 'SelfAdjointOperator -> + Double -> + Int -> + Box.Vector (U.Vector Double) -> + RestartSeed -> + LanczosArena s -> + Int -> + Double -> + U.Vector Double -> + U.Vector Double -> + ST s (Either MoonlightError ExpandedWindow) +expandRestartState op tolerance capacity lockedVectors seedValue arena activeCount previousBeta previousBasis currentBasis + | activeCount > capacity = + freezeExpandedWindow seedValue arena (LanczosRestarting (ActiveDimension activeCount) previousBeta) Nothing + | otherwise = + case runOperatorU op currentBasis of + Left err -> pure (Left err) + Right imageVector -> + case subScaledU imageVector previousBeta previousBasis of + Left err -> pure (Left err) + Right withPreviousRemoved -> + case dotU currentBasis withPreviousRemoved of + Left err -> pure (Left err) + Right alphaValue -> + case subScaledU withPreviousRemoved alphaValue currentBasis of + Left err -> pure (Left err) + Right projectedCurrent -> do + residualResult <- + orthogonalizeAgainstLockedAndArena + lockedVectors + arena + (ActiveDimension activeCount) + projectedCurrent + case residualResult of + Left err -> pure (Left err) + Right residualVector -> do + let currentIndex = activeCount - 1 + betaValue = normU residualVector + UM.unsafeWrite (lanczosAlphaArena arena) currentIndex alphaValue + if betaValue <= tolerance + then + freezeExpandedWindow + seedValue + arena + (LanczosConverged (ActiveDimension activeCount) betaValue) + Nothing + else do + let nextBasis = scaleU (1.0 / betaValue) residualVector + if activeCount >= capacity + then + freezeExpandedWindow + seedValue + arena + (LanczosRestarting (ActiveDimension activeCount) betaValue) + (Just nextBasis) + else do + BoxM.unsafeWrite (lanczosBasisArena arena) activeCount nextBasis + UM.unsafeWrite (lanczosBetaArena arena) currentIndex betaValue + expandRestartState + op + tolerance + capacity + lockedVectors + seedValue + arena + (activeCount + 1) + betaValue + currentBasis + nextBasis + +orthogonalizeAgainstLockedAndArena :: + Box.Vector (U.Vector Double) -> + LanczosArena s -> + ActiveDimension -> + U.Vector Double -> + ST s (Either MoonlightError (U.Vector Double)) +orthogonalizeAgainstLockedAndArena lockedVectors arena activeDimension inputVector = + case projectAgainstVectorListTwice (Box.toList lockedVectors) inputVector of + Left err -> pure (Left err) + Right selectivelyReduced -> orthogonalizeAgainstArena arena activeDimension selectivelyReduced + +freezeExpandedWindow :: RestartSeed -> LanczosArena s -> LanczosState -> Maybe (U.Vector Double) -> ST s (Either MoonlightError ExpandedWindow) +freezeExpandedWindow seedValue arena state boundaryVector = + let activeCount = + case state of + LanczosRunning activeDimension _ _ _ -> activeDimensionValue activeDimension + LanczosConverged activeDimension _ -> activeDimensionValue activeDimension + LanczosBreakdown activeDimension _ -> activeDimensionValue activeDimension + LanczosRestarting activeDimension _ -> activeDimensionValue activeDimension + retainedCount = U.length (restartSeedRetainedValues seedValue) + krylovCount = activeCount - retainedCount + boundaryResidual = + case state of + LanczosRunning _ residual _ _ -> residual + LanczosConverged _ residual -> residual + LanczosBreakdown _ residual -> residual + LanczosRestarting _ residual -> residual + in do + basisVectors <- Box.freeze (BoxM.slice 0 activeCount (lanczosBasisArena arena)) + krylovDiagonal <- U.freeze (UM.slice retainedCount krylovCount (lanczosAlphaArena arena)) + krylovOffDiagonal <- U.freeze (UM.slice retainedCount (max 0 (krylovCount - 1)) (lanczosBetaArena arena)) + pure $ do + projectedOperator <- + mkBorderedProjectedOperator + (restartSeedRetainedValues seedValue) + (restartSeedSpikeCouplings seedValue) + krylovDiagonal + krylovOffDiagonal + Right + ExpandedWindow + { expandedWindowBasisColumns = basisVectors, + expandedWindowProjectedOperator = projectedOperator, + expandedWindowBoundaryResidualNorm = boundaryResidual, + expandedWindowBoundaryVector = boundaryVector, + expandedWindowState = state + } + +restartLoop :: + LinearOperator 'SelfAdjointOperator -> + SpectrumEnd -> + Int -> + Double -> + Int -> + Int -> + Int -> + RestartContext -> + Either MoonlightError LanczosRestartProjection +restartLoop op spectrumEnd requestedCount tolerance capacity ambientDimension remainingCycles context + | length (restartLockedPairs context) >= requestedCount = + finalizeRestartProjection spectrumEnd requestedCount (restartLockedPairs context) [] + | remainingCycles <= 0 = + finalizeRestartProjection spectrumEnd requestedCount (restartLockedPairs context) (restartRetainedPairs context) + | otherwise = do + let lockedVectors = Box.fromList (ritzPairVector <$> restartLockedPairs context) + expandedWindow <- expandRestartWindow op tolerance capacity lockedVectors (restartSeed context) + let activeBasis = expandedWindowBasisColumns expandedWindow + selectedCount = + min + (Box.length activeBasis) + ( max + 1 + (requestedCount - length (restartLockedPairs context) + restartGuardCount requestedCount capacity) + ) + cycleCandidates <- + selectedRitzCandidatesFromProjectedOperator + spectrumEnd + selectedCount + (expandedWindowProjectedOperator expandedWindow) + (expandedWindowBoundaryResidualNorm expandedWindow) + let candidateIsLocked = + ritzCandidateIsLocked tolerance ambientDimension (expandedWindowProjectedOperator expandedWindow) + cycleLockedCandidates = + take + (requestedCount - length (restartLockedPairs context)) + (filter candidateIsLocked cycleCandidates) + cycleLiftedPairs <- traverse (ritzPairFromCandidate activeBasis) cycleLockedCandidates + let pairIsLocked = + ritzPairIsLocked tolerance ambientDimension (expandedWindowProjectedOperator expandedWindow) + (cycleLockedPairs, demotedCandidates) = + partitionEithers + [ if pairIsLocked liftedPair then Left liftedPair else Right liftedCandidate + | (liftedCandidate, liftedPair) <- zip cycleLockedCandidates cycleLiftedPairs + ] + cycleUnlockedCandidates = + demotedCandidates <> filter (not . candidateIsLocked) cycleCandidates + nextLockedPairs = + take requestedCount $ + sortForSpectrumBy + spectrumEnd + ritzPairValue + (restartLockedPairs context <> cycleLockedPairs) + if length nextLockedPairs >= requestedCount + then finalizeRestartProjection spectrumEnd requestedCount nextLockedPairs [] + else + if lanczosStateTerminal (expandedWindowState expandedWindow) + then do + let terminalUnlockedCandidates = + take + (requestedCount - length nextLockedPairs) + (sortForSpectrumBy spectrumEnd ritzCandidateValue cycleUnlockedCandidates) + cycleUnlockedPairs <- traverse (ritzPairFromCandidate activeBasis) terminalUnlockedCandidates + finalizeRestartProjection spectrumEnd requestedCount nextLockedPairs cycleUnlockedPairs + else do + seedResult <- + restartSeedFromRitzCandidates + spectrumEnd + requestedCount + tolerance + capacity + ambientDimension + expandedWindow + nextLockedPairs + cycleUnlockedCandidates + restartLoop + op + spectrumEnd + requestedCount + tolerance + capacity + ambientDimension + (remainingCycles - 1) + (RestartContext nextLockedPairs (restartSeedResultSeed seedResult) (restartSeedResultRetainedPairs seedResult)) + +lanczosStateTerminal :: LanczosState -> Bool +lanczosStateTerminal state = + case state of + LanczosConverged{} -> True + LanczosBreakdown{} -> True + LanczosRunning{} -> False + LanczosRestarting{} -> False + +finalizeRestartProjection :: + SpectrumEnd -> + Int -> + [RitzPair] -> + [RitzPair] -> + Either MoonlightError LanczosRestartProjection +finalizeRestartProjection spectrumEnd requestedCount lockedPairs candidatePairs = + let finalPairs = + take requestedCount $ + sortForSpectrumBy + spectrumEnd + ritzPairValue + (lockedPairs <> candidatePairs) + finalBasis = Box.fromList (ritzPairVector <$> finalPairs) + in if Box.length finalBasis < requestedCount + then Left (InvariantViolation "restarted Lanczos final subspace is smaller than the requested eigenspace") + else do + projectedPairs <- finalProjectedPairsFromRitzPairs finalPairs + Right + LanczosRestartProjection + { lanczosRestartProjectionBasisColumns = finalBasis, + lanczosRestartProjectionProjectedPairs = projectedPairs + } + +selectedRitzCandidatesFromProjectedOperator :: + SpectrumEnd -> + Int -> + BorderedProjectedOperator -> + Double -> + Either MoonlightError [RitzCandidate] +selectedRitzCandidatesFromProjectedOperator spectrumEnd requestedCount projectedOperator boundaryResidualNorm = do + projectedPairs <- selectedProjectedPairsFromBorderedOperator spectrumEnd requestedCount projectedOperator boundaryResidualNorm + traverse + (ritzCandidateFromProjectedPair projectedOperator boundaryResidualNorm projectedPairs) + [0 .. eigenpairCount projectedPairs - 1] + +selectedProjectedPairsFromBorderedOperator :: + SpectrumEnd -> + Int -> + BorderedProjectedOperator -> + Double -> + Either MoonlightError Eigenpairs +selectedProjectedPairsFromBorderedOperator spectrumEnd requestedCount projectedOperator boundaryResidualNorm + | requestedCount <= 0 = + Left (InvariantViolation "projected restarted Lanczos eigensolve requires a positive requested count") + | requestedCount > borderedProjectedOperatorDimension projectedOperator = + Left (InvariantViolation "projected restarted Lanczos eigensolve requested count exceeds basis dimension") + | otherwise = do + selectedColumns <- selectedBorderedProjectedColumns spectrumEnd requestedCount projectedOperator + columnsWithResiduals <- traverse (projectedPairColumn projectedOperator boundaryResidualNorm) selectedColumns + eigenpairsFromColumns + (borderedProjectedOperatorDimension projectedOperator) + columnsWithResiduals + +selectedBorderedProjectedColumns :: + SpectrumEnd -> + Int -> + BorderedProjectedOperator -> + Either MoonlightError [(Double, U.Vector Double, Double)] +selectedBorderedProjectedColumns spectrumEnd requestedCount projectedOperator = + if U.null (borderedRetainedValues projectedOperator) + then do + tridiagonalValue <- + mkSymmetricTridiagonalVectors + (borderedKrylovDiagonal projectedOperator) + (borderedKrylovOffDiagonal projectedOperator) + selectedSymmetricTridiagonalEigenpairColumnsDirect spectrumEnd requestedCount tridiagonalValue + else do + reduction <- reduceBorderedProjectedOperator projectedOperator + reducedColumns <- + selectedSymmetricTridiagonalEigenpairColumnsDirect + spectrumEnd + requestedCount + (borderedReductionTridiagonal reduction) + traverse (liftReducedBorderedColumn reduction) reducedColumns + +liftReducedBorderedColumn :: + BorderedProjectionReduction -> + (Double, U.Vector Double, Double) -> + Either MoonlightError (Double, U.Vector Double, Double) +liftReducedBorderedColumn reduction (eigenvalue, reducedVector, reducedResidualNorm) = do + projectedVector <- normalizeProjectedCoefficientVector =<< linearCombinationColumnsU (borderedReductionBasisColumns reduction) reducedVector + Right (eigenvalue, projectedVector, reducedResidualNorm) + +projectedPairColumn :: + BorderedProjectedOperator -> + Double -> + (Double, U.Vector Double, Double) -> + Either MoonlightError (Double, U.Vector Double, Double) +projectedPairColumn projectedOperator boundaryResidualNorm (eigenvalue, eigenvector, selectedResidualNorm) = do + projectedVector <- normalizeProjectedCoefficientVector eigenvector + residualEvidence <- projectedResidualEvidence projectedOperator boundaryResidualNorm eigenvalue projectedVector + if finiteDouble selectedResidualNorm + then Right (eigenvalue, projectedVector, max residualEvidence selectedResidualNorm) + else Left (InvariantViolation "bordered projected eigensolve produced a non-finite selected residual") + +ritzCandidateFromProjectedPair :: + BorderedProjectedOperator -> + Double -> + Eigenpairs -> + Int -> + Either MoonlightError RitzCandidate +ritzCandidateFromProjectedPair projectedOperator boundaryResidualNorm projectedPairs columnIndex = do + eigenvalue <- + case eigenpairValues projectedPairs U.!? columnIndex of + Nothing -> Left (InvariantViolation "restarted Lanczos projected eigenvalue index out of bounds") + Just value -> Right value + projectedVector <- eigenpairVectorAt columnIndex projectedPairs + projectedResidualNorm <- + case eigenpairResidualNorms projectedPairs U.!? columnIndex of + Nothing -> Left (InvariantViolation "restarted Lanczos projected residual index out of bounds") + Just value -> Right value + boundaryCoupling <- projectedBoundaryCoupling projectedOperator boundaryResidualNorm projectedVector + if finiteDouble projectedResidualNorm && finiteDouble boundaryCoupling + then + Right + RitzCandidate + { ritzCandidateValue = eigenvalue, + ritzCandidateProjectedVector = projectedVector, + ritzCandidateProjectedResidualNorm = projectedResidualNorm, + ritzCandidateBoundaryCoupling = boundaryCoupling + } + else Left (InvariantViolation "restarted Lanczos produced a non-finite projected Ritz residual") + +ritzPairFromCandidate :: + Box.Vector (U.Vector Double) -> + RitzCandidate -> + Either MoonlightError RitzPair +ritzPairFromCandidate basisColumns candidate = do + liftedVector <- normalizeLiftedVector =<< linearCombinationColumnsU basisColumns projectedVector + let residualNorm = ritzCandidateProjectedResidualNorm candidate + if finiteDouble residualNorm && finiteDouble projectedResidualNorm + then + Right + RitzPair + { ritzPairValue = ritzCandidateValue candidate, + ritzPairVector = liftedVector, + ritzPairResidualNorm = residualNorm, + ritzPairProjectedResidualNorm = projectedResidualNorm, + ritzPairBoundaryCoupling = ritzCandidateBoundaryCoupling candidate + } + else Left (InvariantViolation "restarted Lanczos produced a non-finite Ritz residual") + where + projectedVector = ritzCandidateProjectedVector candidate + projectedResidualNorm = ritzCandidateProjectedResidualNorm candidate + +normalizeLiftedVector :: U.Vector Double -> Either MoonlightError (U.Vector Double) +normalizeLiftedVector vectorValue = + let vectorNorm = normU vectorValue + in if finiteDouble vectorNorm && vectorNorm > 0.0 + then Right (scaleU (1.0 / vectorNorm) vectorValue) + else Left (InvariantViolation "restarted Lanczos produced a degenerate lifted Ritz vector") + +restartSeedFromRitzCandidates :: + SpectrumEnd -> + Int -> + Double -> + Int -> + Int -> + ExpandedWindow -> + [RitzPair] -> + [RitzCandidate] -> + Either MoonlightError RestartSeedResult +restartSeedFromRitzCandidates spectrumEnd requestedCount tolerance capacity ambientDimension expandedWindow lockedPairs cycleCandidates = do + let lockedVectors = ritzPairVector <$> lockedPairs + remainingWanted = max 1 (requestedCount - length lockedPairs) + retainedCandidates = + take + (restartRetainedCount capacity remainingWanted (length cycleCandidates)) + (sortForSpectrumBy spectrumEnd ritzCandidateValue cycleCandidates) + case expandedWindowBoundaryVector expandedWindow of + Just boundaryVector -> + if U.length boundaryVector == ambientDimension + then do + retainedPairs <- traverse (ritzPairFromCandidate (expandedWindowBasisColumns expandedWindow)) retainedCandidates + Right + RestartSeedResult + { restartSeedResultSeed = + RestartSeed + { restartSeedBasisColumns = Box.fromList ((ritzPairVector <$> retainedPairs) <> [boundaryVector]), + restartSeedRetainedValues = U.fromList (ritzPairValue <$> retainedPairs), + restartSeedSpikeCouplings = U.fromList (ritzPairBoundaryCoupling <$> retainedPairs) + }, + restartSeedResultRetainedPairs = retainedPairs + } + else Left (InvariantViolation "restarted Lanczos boundary vector dimension mismatch") + Nothing -> + if null retainedCandidates + then do + seedBasis <- canonicalRestartBasis tolerance ambientDimension lockedVectors + Right + RestartSeedResult + { restartSeedResultSeed = + RestartSeed + { restartSeedBasisColumns = seedBasis, + restartSeedRetainedValues = U.empty, + restartSeedSpikeCouplings = U.empty + }, + restartSeedResultRetainedPairs = [] + } + else Left (InvariantViolation "restarted Lanczos cannot retain Ritz values without a boundary vector") + +finalProjectedPairsFromRitzPairs :: [RitzPair] -> Either MoonlightError Eigenpairs +finalProjectedPairsFromRitzPairs finalPairs = + let projectedDimension = length finalPairs + in eigenpairsFromColumns + projectedDimension + (zipWith finalProjectedPairColumn [0 ..] finalPairs) + where + finalProjectedPairColumn columnIndex ritzPair = + ( ritzPairValue ritzPair, + unitVector (length finalPairs) columnIndex, + ritzPairProjectedResidualNorm ritzPair + ) + +mkBorderedProjectedOperator :: + U.Vector Double -> + U.Vector Double -> + U.Vector Double -> + U.Vector Double -> + Either MoonlightError BorderedProjectedOperator +mkBorderedProjectedOperator retainedValues spikeCouplings krylovDiagonal krylovOffDiagonal + | U.length retainedValues /= U.length spikeCouplings = + Left (InvariantViolation "bordered projected operator retained value count must match spike count") + | U.null krylovDiagonal = + Left (InvariantViolation "bordered projected operator requires a non-empty Krylov block") + | U.length krylovOffDiagonal /= U.length krylovDiagonal - 1 = + Left (InvariantViolation "bordered projected operator Krylov off-diagonal length mismatch") + | U.any (not . finiteDouble) retainedValues + || U.any (not . finiteDouble) spikeCouplings + || U.any (not . finiteDouble) krylovDiagonal + || U.any (not . finiteDouble) krylovOffDiagonal = + Left (InvariantViolation "bordered projected operator entries must be finite") + | otherwise = + Right + BorderedProjectedOperator + { borderedRetainedValues = retainedValues, + borderedSpikeCouplings = spikeCouplings, + borderedKrylovDiagonal = krylovDiagonal, + borderedKrylovOffDiagonal = krylovOffDiagonal + } + +borderedProjectedOperatorDimension :: BorderedProjectedOperator -> Int +borderedProjectedOperatorDimension projectedOperator = + U.length (borderedRetainedValues projectedOperator) + U.length (borderedKrylovDiagonal projectedOperator) + +applyBorderedProjectedOperatorU :: + BorderedProjectedOperator -> + U.Vector Double -> + Either MoonlightError (U.Vector Double) +applyBorderedProjectedOperatorU projectedOperator inputVector = + let retainedCount = U.length (borderedRetainedValues projectedOperator) + krylovCount = U.length (borderedKrylovDiagonal projectedOperator) + projectedDimension = retainedCount + krylovCount + in if U.length inputVector /= projectedDimension + then Left (InvariantViolation "bordered projected operator input dimension mismatch") + else + Right + ( U.generate + projectedDimension + (borderedProjectedOperatorEntry projectedOperator inputVector retainedCount krylovCount) + ) + +borderedProjectedOperatorEntry :: + BorderedProjectedOperator -> + U.Vector Double -> + Int -> + Int -> + Int -> + Double +borderedProjectedOperatorEntry projectedOperator inputVector retainedCount krylovCount entryIndex = + if entryIndex < retainedCount + then + let retainedValue = borderedRetainedValues projectedOperator `U.unsafeIndex` entryIndex + spikeValue = borderedSpikeCouplings projectedOperator `U.unsafeIndex` entryIndex + retainedEntry = inputVector `U.unsafeIndex` entryIndex + firstKrylovEntry = inputVector `U.unsafeIndex` retainedCount + in retainedValue * retainedEntry + spikeValue * firstKrylovEntry + else + let krylovIndex = entryIndex - retainedCount + diagonalValue = borderedKrylovDiagonal projectedOperator `U.unsafeIndex` krylovIndex + centerEntry = inputVector `U.unsafeIndex` entryIndex + leftEntry = + if krylovIndex <= 0 + then U.sum (U.zipWith (*) (borderedSpikeCouplings projectedOperator) (U.take retainedCount inputVector)) + else (borderedKrylovOffDiagonal projectedOperator `U.unsafeIndex` (krylovIndex - 1)) * (inputVector `U.unsafeIndex` (entryIndex - 1)) + rightEntry = + if krylovIndex + 1 >= krylovCount + then 0.0 + else (borderedKrylovOffDiagonal projectedOperator `U.unsafeIndex` krylovIndex) * (inputVector `U.unsafeIndex` (entryIndex + 1)) + in leftEntry + diagonalValue * centerEntry + rightEntry + +reduceBorderedProjectedOperator :: + BorderedProjectedOperator -> + Either MoonlightError BorderedProjectionReduction +reduceBorderedProjectedOperator projectedOperator = + runST $ do + let projectedDimension = borderedProjectedOperatorDimension projectedOperator + reductionTolerance = borderedReductionTolerance projectedOperator + arena <- newBorderedReductionArena projectedDimension + initializeBorderedReductionArena projectedOperator arena + chaseBorderedReductionBulges arena reductionTolerance + diagonalValues <- freezeBorderedReductionDiagonal arena + offDiagonalValues <- freezeBorderedReductionOffDiagonal arena + basisColumns <- freezeBorderedReductionBasis arena + pure $ do + tridiagonalValue <- mkSymmetricTridiagonalVectors diagonalValues offDiagonalValues + Right + BorderedProjectionReduction + { borderedReductionBasisColumns = basisColumns, + borderedReductionTridiagonal = tridiagonalValue + } + +newBorderedReductionArena :: Int -> ST s (BorderedReductionArena s) +newBorderedReductionArena projectedDimension = do + payload <- newPrimArray payloadLength + setPrimArray payload 0 payloadLength 0.0 + pure + BorderedReductionArena + { borderedReductionArenaDimension = projectedDimension, + borderedReductionArenaPayload = payload, + borderedReductionArenaMatrixOffset = matrixOffset, + borderedReductionArenaBasisOffset = basisOffset + } + where + matrixOffset = 0 + matrixLength = projectedDimension * projectedDimension + basisOffset = matrixOffset + matrixLength + basisLength = projectedDimension * projectedDimension + payloadLength = basisOffset + basisLength + +initializeBorderedReductionArena :: BorderedProjectedOperator -> BorderedReductionArena s -> ST s () +initializeBorderedReductionArena projectedOperator arena = do + forIndex 0 retainedCount $ \retainedIndex -> + writeBorderedMatrixEntry arena retainedIndex retainedIndex (borderedRetainedValues projectedOperator `U.unsafeIndex` retainedIndex) + forIndex 0 krylovCount $ \krylovIndex -> + writeBorderedMatrixEntry arena (retainedCount + krylovIndex) (retainedCount + krylovIndex) (borderedKrylovDiagonal projectedOperator `U.unsafeIndex` krylovIndex) + forIndex 0 retainedCount $ \retainedIndex -> + writeSymmetricBorderedMatrixEntry arena retainedIndex retainedCount (borderedSpikeCouplings projectedOperator `U.unsafeIndex` retainedIndex) + forIndex 0 (max 0 (krylovCount - 1)) $ \krylovIndex -> + writeSymmetricBorderedMatrixEntry + arena + (retainedCount + krylovIndex) + (retainedCount + krylovIndex + 1) + (borderedKrylovOffDiagonal projectedOperator `U.unsafeIndex` krylovIndex) + forIndex 0 projectedDimension $ \basisIndex -> + writeBorderedBasisEntry arena basisIndex basisIndex 1.0 + where + retainedCount = U.length (borderedRetainedValues projectedOperator) + krylovCount = U.length (borderedKrylovDiagonal projectedOperator) + projectedDimension = retainedCount + krylovCount + +chaseBorderedReductionBulges :: BorderedReductionArena s -> Double -> ST s () +chaseBorderedReductionBulges arena reductionTolerance = + forIndex 0 (max 0 (projectedDimension - 2)) $ \columnIndex -> + forDescendingIndex (projectedDimension - 1) (columnIndex + 2) $ \rowIndex -> + annihilateBorderedReductionEntry arena reductionTolerance columnIndex (rowIndex - 1) rowIndex + where + projectedDimension = borderedReductionArenaDimension arena + +annihilateBorderedReductionEntry :: + BorderedReductionArena s -> + Double -> + Int -> + Int -> + Int -> + ST s () +annihilateBorderedReductionEntry arena reductionTolerance columnIndex leftIndex rightIndex = do + targetValue <- readBorderedMatrixEntry arena rightIndex columnIndex + if abs targetValue <= reductionTolerance + then do + writeBorderedMatrixEntry arena rightIndex columnIndex 0.0 + writeBorderedMatrixEntry arena columnIndex rightIndex 0.0 + else do + pivotValue <- readBorderedMatrixEntry arena leftIndex columnIndex + let radiusValue = hypotStable pivotValue targetValue + if radiusValue <= 0.0 + then do + writeBorderedMatrixEntry arena rightIndex columnIndex 0.0 + writeBorderedMatrixEntry arena columnIndex rightIndex 0.0 + else do + let cosineValue = pivotValue / radiusValue + sineValue = targetValue / radiusValue + applyBorderedReductionGivens arena leftIndex rightIndex cosineValue sineValue + writeBorderedMatrixEntry arena leftIndex columnIndex radiusValue + writeBorderedMatrixEntry arena columnIndex leftIndex radiusValue + writeBorderedMatrixEntry arena rightIndex columnIndex 0.0 + writeBorderedMatrixEntry arena columnIndex rightIndex 0.0 + +applyBorderedReductionGivens :: + BorderedReductionArena s -> + Int -> + Int -> + Double -> + Double -> + ST s () +applyBorderedReductionGivens arena leftIndex rightIndex cosineValue sineValue = do + forIndex 0 projectedDimension $ \columnIndex -> do + leftEntry <- readBorderedMatrixEntry arena leftIndex columnIndex + rightEntry <- readBorderedMatrixEntry arena rightIndex columnIndex + writeBorderedMatrixEntry arena leftIndex columnIndex (cosineValue * leftEntry + sineValue * rightEntry) + writeBorderedMatrixEntry arena rightIndex columnIndex ((negate sineValue) * leftEntry + cosineValue * rightEntry) + forIndex 0 projectedDimension $ \rowIndex -> do + leftEntry <- readBorderedMatrixEntry arena rowIndex leftIndex + rightEntry <- readBorderedMatrixEntry arena rowIndex rightIndex + writeBorderedMatrixEntry arena rowIndex leftIndex (cosineValue * leftEntry + sineValue * rightEntry) + writeBorderedMatrixEntry arena rowIndex rightIndex ((negate sineValue) * leftEntry + cosineValue * rightEntry) + rotateBorderedReductionBasisColumns arena leftIndex rightIndex cosineValue sineValue + where + projectedDimension = borderedReductionArenaDimension arena + +rotateBorderedReductionBasisColumns :: + BorderedReductionArena s -> + Int -> + Int -> + Double -> + Double -> + ST s () +rotateBorderedReductionBasisColumns arena leftIndex rightIndex cosineValue sineValue = + forIndex 0 projectedDimension $ \rowIndex -> do + leftEntry <- readBorderedBasisEntry arena rowIndex leftIndex + rightEntry <- readBorderedBasisEntry arena rowIndex rightIndex + writeBorderedBasisEntry arena rowIndex leftIndex (cosineValue * leftEntry + sineValue * rightEntry) + writeBorderedBasisEntry arena rowIndex rightIndex ((negate sineValue) * leftEntry + cosineValue * rightEntry) + where + projectedDimension = borderedReductionArenaDimension arena + +freezeBorderedReductionDiagonal :: BorderedReductionArena s -> ST s (U.Vector Double) +freezeBorderedReductionDiagonal arena = + U.generateM projectedDimension $ \entryIndex -> + readBorderedMatrixEntry arena entryIndex entryIndex + where + projectedDimension = borderedReductionArenaDimension arena + +freezeBorderedReductionOffDiagonal :: BorderedReductionArena s -> ST s (U.Vector Double) +freezeBorderedReductionOffDiagonal arena = + U.generateM (max 0 (projectedDimension - 1)) $ \entryIndex -> + readBorderedMatrixEntry arena entryIndex (entryIndex + 1) + where + projectedDimension = borderedReductionArenaDimension arena + +freezeBorderedReductionBasis :: BorderedReductionArena s -> ST s (Box.Vector (U.Vector Double)) +freezeBorderedReductionBasis arena = + Box.generateM projectedDimension $ \columnIndex -> + U.generateM projectedDimension $ \rowIndex -> + readBorderedBasisEntry arena rowIndex columnIndex + where + projectedDimension = borderedReductionArenaDimension arena + +writeSymmetricBorderedMatrixEntry :: BorderedReductionArena s -> Int -> Int -> Double -> ST s () +writeSymmetricBorderedMatrixEntry arena rowIndex columnIndex entryValue = do + writeBorderedMatrixEntry arena rowIndex columnIndex entryValue + writeBorderedMatrixEntry arena columnIndex rowIndex entryValue + +readBorderedMatrixEntry :: BorderedReductionArena s -> Int -> Int -> ST s Double +readBorderedMatrixEntry arena rowIndex columnIndex = + readPrimArray (borderedReductionArenaPayload arena) (borderedMatrixEntryOffset arena rowIndex columnIndex) + +writeBorderedMatrixEntry :: BorderedReductionArena s -> Int -> Int -> Double -> ST s () +writeBorderedMatrixEntry arena rowIndex columnIndex entryValue = + writePrimArray (borderedReductionArenaPayload arena) (borderedMatrixEntryOffset arena rowIndex columnIndex) entryValue + +readBorderedBasisEntry :: BorderedReductionArena s -> Int -> Int -> ST s Double +readBorderedBasisEntry arena rowIndex columnIndex = + readPrimArray (borderedReductionArenaPayload arena) (borderedBasisEntryOffset arena rowIndex columnIndex) + +writeBorderedBasisEntry :: BorderedReductionArena s -> Int -> Int -> Double -> ST s () +writeBorderedBasisEntry arena rowIndex columnIndex entryValue = + writePrimArray (borderedReductionArenaPayload arena) (borderedBasisEntryOffset arena rowIndex columnIndex) entryValue + +borderedMatrixEntryOffset :: BorderedReductionArena s -> Int -> Int -> Int +borderedMatrixEntryOffset arena rowIndex columnIndex = + borderedReductionArenaMatrixOffset arena + rowIndex * borderedReductionArenaDimension arena + columnIndex + +borderedBasisEntryOffset :: BorderedReductionArena s -> Int -> Int -> Int +borderedBasisEntryOffset arena rowIndex columnIndex = + borderedReductionArenaBasisOffset arena + columnIndex * borderedReductionArenaDimension arena + rowIndex + +borderedReductionTolerance :: BorderedProjectedOperator -> Double +borderedReductionTolerance projectedOperator = + 256.0 + * epsDouble + * sqrt (fromIntegral (max 1 (borderedProjectedOperatorDimension projectedOperator)) :: Double) + * max 1.0 (borderedProjectedOperatorInfinityBound projectedOperator) + +borderedProjectedOperatorInfinityBound :: BorderedProjectedOperator -> Double +borderedProjectedOperatorInfinityBound projectedOperator = + maximum [1.0, retainedBound, firstKrylovBound, tailKrylovBound] + where + retainedValues = borderedRetainedValues projectedOperator + spikeValues = borderedSpikeCouplings projectedOperator + krylovDiagonal = borderedKrylovDiagonal projectedOperator + krylovOffDiagonal = borderedKrylovOffDiagonal projectedOperator + offDiagonalAt :: Int -> Double + offDiagonalAt entryIndex = maybe 0.0 abs (krylovOffDiagonal U.!? entryIndex) + retainedBound = + if U.null retainedValues + then 0.0 + else U.maximum (U.zipWith (\value spike -> abs value + abs spike) retainedValues spikeValues) + firstKrylovBound = + case krylovDiagonal U.!? 0 of + Nothing -> 0.0 + Just firstDiagonal -> abs firstDiagonal + U.sum (U.map abs spikeValues) + offDiagonalAt 0 + tailKrylovBound = + if U.length krylovDiagonal <= 1 + then 0.0 + else + U.maximum + ( U.imap + (\entryIndex diagonalValue -> offDiagonalAt entryIndex + abs diagonalValue + offDiagonalAt (entryIndex + 1)) + (U.drop 1 krylovDiagonal) + ) + +projectedResidualEvidence :: + BorderedProjectedOperator -> + Double -> + Double -> + U.Vector Double -> + Either MoonlightError Double +projectedResidualEvidence projectedOperator boundaryResidualNorm eigenvalue projectedVector = do + projectedImage <- applyBorderedProjectedOperatorU projectedOperator projectedVector + projectedResidual <- subScaledU projectedImage eigenvalue projectedVector + boundaryCoupling <- projectedBoundaryCoupling projectedOperator boundaryResidualNorm projectedVector + let projectedNorm = normU projectedResidual + residualNorm = sqrt (projectedNorm * projectedNorm + boundaryCoupling * boundaryCoupling) + if finiteDouble residualNorm + then Right residualNorm + else Left (InvariantViolation "bordered projected eigensolve produced a non-finite residual") + +projectedBoundaryCoupling :: + BorderedProjectedOperator -> + Double -> + U.Vector Double -> + Either MoonlightError Double +projectedBoundaryCoupling projectedOperator boundaryResidualNorm projectedVector = + case projectedVector U.!? (borderedProjectedOperatorDimension projectedOperator - 1) of + Nothing -> Left (InvariantViolation "bordered projected eigenvector boundary coefficient index out of bounds") + Just coefficient -> Right (boundaryResidualNorm * coefficient) + +normalizeProjectedCoefficientVector :: U.Vector Double -> Either MoonlightError (U.Vector Double) +normalizeProjectedCoefficientVector vectorValue = + let vectorNorm = normU vectorValue + in if finiteDouble vectorNorm && vectorNorm > 0.0 + then Right (scaleU (1.0 / vectorNorm) vectorValue) + else Left (InvariantViolation "bordered projected eigensolve produced a degenerate coefficient vector") + +unitVector :: Int -> Int -> U.Vector Double +unitVector dimension activeIndex = + U.generate dimension (\entryIndex -> if entryIndex == activeIndex then 1.0 else 0.0) + +orthonormalizeCandidateVectors :: + Double -> + [U.Vector Double] -> + [U.Vector Double] -> + Either MoonlightError (Box.Vector (U.Vector Double)) +orthonormalizeCandidateVectors tolerance lockedVectors candidateVectors = + Box.fromList . reverse + <$> foldM appendCandidate [] candidateVectors + where + appendCandidate acceptedRev candidateVector = do + lockedReduced <- projectAgainstVectorListTwice lockedVectors candidateVector + activeReduced <- projectAgainstVectorListTwice acceptedRev lockedReduced + let candidateNorm = normU activeReduced + if finiteDouble candidateNorm && candidateNorm > tolerance + then Right (scaleU (1.0 / candidateNorm) activeReduced : acceptedRev) + else Right acceptedRev + +canonicalRestartBasis :: + Double -> + Int -> + [U.Vector Double] -> + Either MoonlightError (Box.Vector (U.Vector Double)) +canonicalRestartBasis tolerance ambientDimension lockedVectors = + case listToMaybe (filter (not . Box.null) candidateBases) of + Just basisValue -> Right basisValue + Nothing -> Left (InvariantViolation "restarted Lanczos could not construct a restart vector orthogonal to locked Ritz vectors") + where + coordinateVectors = + U.generate ambientDimension + <$> [ \rowIndex -> if rowIndex == coordinateIndex then 1.0 else 0.0 + | coordinateIndex <- [0 .. ambientDimension - 1] + ] + candidateBases = + catMaybes + ( either + (const Nothing) + Just + . orthonormalizeCandidateVectors tolerance lockedVectors + . pure + <$> coordinateVectors + ) + +projectAgainstVectorListTwice :: + [U.Vector Double] -> + U.Vector Double -> + Either MoonlightError (U.Vector Double) +projectAgainstVectorListTwice basisVectors inputVector = + projectAgainstVectorListOnce basisVectors inputVector >>= projectAgainstVectorListOnce basisVectors + +projectAgainstVectorListOnce :: + [U.Vector Double] -> + U.Vector Double -> + Either MoonlightError (U.Vector Double) +projectAgainstVectorListOnce basisVectors inputVector = + foldM projectOne inputVector basisVectors + where + projectOne workingVector basisVector = do + coefficient <- dotU basisVector workingVector + subScaledU workingVector coefficient basisVector + +ritzPairIsLocked :: Double -> Int -> BorderedProjectedOperator -> RitzPair -> Bool +ritzPairIsLocked tolerance ambientDimension projectedOperator ritzPair = + max (ritzPairResidualNorm ritzPair) (ritzPairProjectedResidualNorm ritzPair) + <= ritzLockToleranceBound tolerance ambientDimension projectedOperator (ritzPairValue ritzPair) + +ritzCandidateIsLocked :: Double -> Int -> BorderedProjectedOperator -> RitzCandidate -> Bool +ritzCandidateIsLocked tolerance ambientDimension projectedOperator candidate = + ritzCandidateProjectedResidualNorm candidate + <= ritzLockToleranceBound tolerance ambientDimension projectedOperator (ritzCandidateValue candidate) + +ritzLockToleranceBound :: Double -> Int -> BorderedProjectedOperator -> Double -> Double +ritzLockToleranceBound tolerance ambientDimension projectedOperator eigenvalue = + max + (ritzLockThreshold tolerance ambientDimension eigenvalue) + ( inverseIterationResidualToleranceBound + (borderedProjectedOperatorInfinityBound projectedOperator) + eigenvalue + (borderedProjectedOperatorDimension projectedOperator) + ) + +ritzLockThreshold :: Double -> Int -> Double -> Double +ritzLockThreshold tolerance ambientDimension eigenvalue = + max + tolerance + (128.0 * epsDouble * sqrt (fromIntegral (max 1 ambientDimension) :: Double) * max 1.0 (abs eigenvalue)) + +restartGuardCount :: Int -> Int -> Int +restartGuardCount requestedCount capacity = + max 1 (min requestedCount (max 1 (capacity `quot` 2))) + +restartRetainedCount :: Int -> Int -> Int -> Int +restartRetainedCount capacity remainingWanted candidateCount = + min candidateCount (max 1 (min retainedRoom (remainingWanted + restartGuardCount remainingWanted capacity))) + where + retainedRoom = + if capacity <= 1 + then 1 + else capacity - 1 + +maxRestartCycles :: Int -> Int -> Int +maxRestartCycles ambientDimension capacity = + max 1 (4 * max 1 ambientDimension * max 1 (ambientDimension `quot` max 1 capacity)) + +freezeLanczosState :: LanczosArena s -> LanczosState -> ST s (Either MoonlightError LanczosDecomposition) +freezeLanczosState arena state = + case state of + LanczosConverged activeDimension finalResidual -> + freezeLanczosDecomposition arena activeDimension finalResidual + LanczosBreakdown activeDimension finalResidual -> + freezeLanczosDecomposition arena activeDimension finalResidual + LanczosRestarting activeDimension finalResidual -> + freezeLanczosDecomposition arena activeDimension finalResidual + LanczosRunning{} -> + pure (Left (InvariantViolation "Lanczos reached an unfinished running state")) + +freezeLanczosDecomposition :: LanczosArena s -> ActiveDimension -> Double -> ST s (Either MoonlightError LanczosDecomposition) +freezeLanczosDecomposition arena activeDimension finalResidual = do + basisVectors <- Box.freeze (BoxM.slice 0 activeCount (lanczosBasisArena arena)) + alphaValues <- U.freeze (UM.slice 0 activeCount (lanczosAlphaArena arena)) + betaValues <- U.freeze (UM.slice 0 (max 0 (activeCount - 1)) (lanczosBetaArena arena)) + pure $ do + projectedTridiagonal <- mkSymmetricTridiagonal (U.toList alphaValues) (U.toList betaValues) + mkLanczosDecomposition basisVectors projectedTridiagonal finalResidual + where + activeCount = activeDimensionValue activeDimension diff --git a/src-spectral/Moonlight/LinAlg/Pure/Krylov/Projected.hs b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Projected.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Projected.hs @@ -0,0 +1,419 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Krylov.Projected + ( SpectrumEnd (..), + SymmetricProjectedOperator (..), + symmetricProjectedOperatorDimension, + applySymmetricProjectedOperatorU, + ProjectedSubspace, + projectedSubspaceDimension, + mkStructuredProjectedSubspace, + projectedSubspaceBasisColumns, + projectedSubspaceOperator, + projectedSubspaceFromLanczos, + projectedSubspaceFromBlockLanczos, + projectedEigenvalues, + projectedEigenpairs, + projectedEigenvaluesFromRestartedLanczos, + projectedEigenpairsFromRestartedLanczos, + ) +where + +import Data.Kind (Type) +import qualified Data.Vector as Box +import qualified Data.Vector.Unboxed as U +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Internal.Eigen.Kernels (epsDouble, finiteDouble) +import Moonlight.LinAlg.Internal.VectorOps (normU, scaleU, subScaledU) +import Moonlight.LinAlg.Pure.Krylov.Config (LanczosConfig, lanczosTolerance) +import Moonlight.LinAlg.Pure.Krylov.Decomposition +import Moonlight.LinAlg.Pure.Krylov.Lanczos + ( lanczosRestartProjectionBasisColumns, + lanczosRestartProjectionProjectedPairs, + lanczosRestartedProjection, + lanczosSymmetric, + ritzLockThreshold, + ) +import Moonlight.LinAlg.Pure.Krylov.Internal (linearCombinationColumnsU) +import Moonlight.LinAlg.Pure.Spectral.Result + ( Eigenpairs, + eigenpairCount, + eigenpairResidualNorms, + eigenpairValues, + eigenpairVectorAt, + mkEigenpairs, + ) +import Moonlight.LinAlg.Pure.Operator + ( LinearOperator, + OperatorSymmetry (SelfAdjointOperator), + operatorShape, + runOperatorU, + ) +import Moonlight.LinAlg.Pure.Krylov.Selection (SpectrumEnd (..)) +import Moonlight.LinAlg.Pure.Krylov.SelectedTridiagonal + ( selectedSymmetricTridiagonalEigenpairsDirect, + selectedSymmetricTridiagonalEigenvaluesDirect, + ) +import Moonlight.LinAlg.Pure.Structured.BlockTridiagonal + ( SymmetricBlockTridiagonal, + applySymmetricBlockTridiagonalU, + symmetricBlockTridiagonalDimension, + ) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal + ( SymmetricTridiagonal, + applySymmetricTridiagonalU, + symmetricTridiagonalDimension, + ) +import Prelude + +type SymmetricProjectedOperator :: Type +data SymmetricProjectedOperator + = TridiagonalProjectedOperator !SymmetricTridiagonal + | BlockTridiagonalProjectedOperator !SymmetricBlockTridiagonal + deriving stock (Eq, Show) + +symmetricProjectedOperatorDimension :: SymmetricProjectedOperator -> Int +symmetricProjectedOperatorDimension projectedOperator = + case projectedOperator of + TridiagonalProjectedOperator tridiagonalValue -> symmetricTridiagonalDimension tridiagonalValue + BlockTridiagonalProjectedOperator blockTridiagonalValue -> symmetricBlockTridiagonalDimension blockTridiagonalValue + +applySymmetricProjectedOperatorU :: + SymmetricProjectedOperator -> + U.Vector Double -> + Either MoonlightError (U.Vector Double) +applySymmetricProjectedOperatorU projectedOperator inputVector = + let dimension = symmetricProjectedOperatorDimension projectedOperator + in if U.length inputVector /= dimension + then + Left + ( InvariantViolation + ( "Projected operator input dimension mismatch: expected " + <> show dimension + <> " but received " + <> show (U.length inputVector) + ) + ) + else + case projectedOperator of + TridiagonalProjectedOperator tridiagonalValue -> + Right (applySymmetricTridiagonalU tridiagonalValue inputVector) + BlockTridiagonalProjectedOperator blockTridiagonalValue -> + applySymmetricBlockTridiagonalU blockTridiagonalValue inputVector + +type ProjectedSubspace :: Type +data ProjectedSubspace = ProjectedSubspace + { subspaceBasisColumns :: !(Box.Vector (U.Vector Double)), + subspaceOperatorValue :: !SymmetricProjectedOperator + } + deriving stock (Eq, Show) + +projectedSubspaceDimension :: ProjectedSubspace -> Int +projectedSubspaceDimension = symmetricProjectedOperatorDimension . projectedSubspaceOperator + +mkStructuredProjectedSubspace :: + Box.Vector (U.Vector Double) -> + SymmetricProjectedOperator -> + Either MoonlightError ProjectedSubspace +mkStructuredProjectedSubspace basisColumns projectedOperator = + let basisCount = symmetricProjectedOperatorDimension projectedOperator + basisDimensions = U.length <$> Box.toList basisColumns + basisDimension = + case basisDimensions of + [] -> 0 + firstDimension : _ -> firstDimension + in if Box.length basisColumns /= basisCount + then Left (InvariantViolation "Projected subspace basis column count must match the projected dimension witness") + else + if any (/= basisDimension) basisDimensions + then Left (InvariantViolation "Projected subspace basis columns must have equal length") + else Right (ProjectedSubspace basisColumns projectedOperator) + +projectedSubspaceBasisColumns :: ProjectedSubspace -> Box.Vector (U.Vector Double) +projectedSubspaceBasisColumns = subspaceBasisColumns + +projectedSubspaceOperator :: ProjectedSubspace -> SymmetricProjectedOperator +projectedSubspaceOperator = subspaceOperatorValue + +projectedSubspaceFromLanczos :: LanczosDecomposition -> ProjectedSubspace +projectedSubspaceFromLanczos decomposition = + let basisColumns = lanczosBasisColumns decomposition + projectedTridiagonal = lanczosProjectedTridiagonal decomposition + in ProjectedSubspace + { subspaceBasisColumns = basisColumns, + subspaceOperatorValue = TridiagonalProjectedOperator projectedTridiagonal + } + +projectedSubspaceFromBlockLanczos :: BlockLanczosDecomposition -> ProjectedSubspace +projectedSubspaceFromBlockLanczos decomposition = + let basisColumns = blockLanczosBasisColumns decomposition + projectedBlockTridiagonal = blockLanczosProjectedBlockTridiagonal decomposition + in ProjectedSubspace + { subspaceBasisColumns = basisColumns, + subspaceOperatorValue = BlockTridiagonalProjectedOperator projectedBlockTridiagonal + } + +projectedEigenvalues :: + SpectrumEnd -> + Int -> + LinearOperator 'SelfAdjointOperator -> + ProjectedSubspace -> + Either MoonlightError (U.Vector Double) +projectedEigenvalues spectrumEnd requestedCount op subspace + | requestedCount <= 0 = + Left (InvariantViolation "Projected eigenvalue count must be positive") + | otherwise = do + basisCount <- validateProjectedSubspace op subspace + validateProjectedRequestedCount "Projected eigenvalue" requestedCount basisCount + symmetricProjectedEigenvalues spectrumEnd requestedCount (projectedSubspaceOperator subspace) + +projectedEigenpairs :: + SpectrumEnd -> + Int -> + LinearOperator 'SelfAdjointOperator -> + ProjectedSubspace -> + Either MoonlightError Eigenpairs +projectedEigenpairs spectrumEnd requestedCount op subspace + | requestedCount <= 0 = + Left (InvariantViolation "Projected eigenpair count must be positive") + | otherwise = do + basisCount <- validateProjectedSubspace op subspace + validateProjectedRequestedCount "Projected eigenpair" requestedCount basisCount + projectedPairs <- symmetricProjectedEigenpairs spectrumEnd requestedCount (projectedSubspaceOperator subspace) + liftProjectedEigenpairs op (projectedSubspaceBasisColumns subspace) projectedPairs + +projectedEigenvaluesFromRestartedLanczos :: + LanczosConfig -> + SpectrumEnd -> + Int -> + LinearOperator 'SelfAdjointOperator -> + U.Vector Double -> + Either MoonlightError (U.Vector Double) +projectedEigenvaluesFromRestartedLanczos config spectrumEnd requestedCount op seedVector = + case singleCycleCertifiedEigenpairs config spectrumEnd requestedCount op seedVector of + Just certifiedPairs -> Right (eigenpairValues certifiedPairs) + Nothing -> do + restartProjection <- lanczosRestartedProjection config spectrumEnd requestedCount op seedVector + pure (eigenpairValues (lanczosRestartProjectionProjectedPairs restartProjection)) + +projectedEigenpairsFromRestartedLanczos :: + LanczosConfig -> + SpectrumEnd -> + Int -> + LinearOperator 'SelfAdjointOperator -> + U.Vector Double -> + Either MoonlightError Eigenpairs +projectedEigenpairsFromRestartedLanczos config spectrumEnd requestedCount op seedVector = + case singleCycleCertifiedEigenpairs config spectrumEnd requestedCount op seedVector of + Just certifiedPairs -> Right certifiedPairs + Nothing -> do + restartProjection <- lanczosRestartedProjection config spectrumEnd requestedCount op seedVector + liftProjectedEigenpairs + op + (lanczosRestartProjectionBasisColumns restartProjection) + (lanczosRestartProjectionProjectedPairs restartProjection) + +singleCycleCertifiedEigenpairs :: + LanczosConfig -> + SpectrumEnd -> + Int -> + LinearOperator 'SelfAdjointOperator -> + U.Vector Double -> + Maybe Eigenpairs +singleCycleCertifiedEigenpairs config spectrumEnd requestedCount op seedVector = + case lanczosSymmetric config op seedVector of + Left _ -> Nothing + Right decomposition -> + let subspace = projectedSubspaceFromLanczos decomposition + basisCount = Box.length (projectedSubspaceBasisColumns subspace) + in if requestedCount <= 0 || requestedCount > basisCount + then Nothing + else case projectedEigenpairs spectrumEnd requestedCount op subspace of + Left _ -> Nothing + Right liftedPairs -> + let (_, ambientDimension) = operatorShape op + tolerance = lanczosTolerance config + pairCertified eigenvalue residualNorm = + residualNorm <= ritzLockThreshold tolerance ambientDimension eigenvalue + allCertified = + U.and + ( U.zipWith + pairCertified + (eigenpairValues liftedPairs) + (eigenpairResidualNorms liftedPairs) + ) + in if allCertified then Just liftedPairs else Nothing + +symmetricProjectedEigenvalues :: + SpectrumEnd -> + Int -> + SymmetricProjectedOperator -> + Either MoonlightError (U.Vector Double) +symmetricProjectedEigenvalues spectrumEnd requestedCount projectedOperator = do + let operatorDimension = symmetricProjectedOperatorDimension projectedOperator + if requestedCount <= 0 + then Left (InvariantViolation "Projected eigensolve requires a positive requested count") + else if requestedCount > operatorDimension + then Left (InvariantViolation "Projected eigensolve requested count exceeds projected dimension") + else + case projectedOperator of + TridiagonalProjectedOperator tridiagonalValue -> + selectedSymmetricTridiagonalEigenvaluesDirect spectrumEnd requestedCount tridiagonalValue + BlockTridiagonalProjectedOperator _ -> Left blockProjectedSpectralObstruction + +symmetricProjectedEigenpairs :: + SpectrumEnd -> + Int -> + SymmetricProjectedOperator -> + Either MoonlightError Eigenpairs +symmetricProjectedEigenpairs spectrumEnd requestedCount projectedOperator = do + let operatorDimension = symmetricProjectedOperatorDimension projectedOperator + if requestedCount <= 0 + then Left (InvariantViolation "Projected eigensolve requires a positive requested count") + else if requestedCount > operatorDimension + then Left (InvariantViolation "Projected eigensolve requested count exceeds projected dimension") + else + case projectedOperator of + TridiagonalProjectedOperator tridiagonalValue -> + selectedSymmetricTridiagonalEigenpairsDirect spectrumEnd requestedCount tridiagonalValue + BlockTridiagonalProjectedOperator _ -> Left blockProjectedSpectralObstruction + +blockProjectedSpectralObstruction :: MoonlightError +blockProjectedSpectralObstruction = + InvariantViolation "pure block-projected eigensolve has no exact block backend; use the native symmetric-band EigenRequest executor" + +validateProjectedSubspace :: LinearOperator 'SelfAdjointOperator -> ProjectedSubspace -> Either MoonlightError Int +validateProjectedSubspace op subspace = + let basisColumns = projectedSubspaceBasisColumns subspace + basisDimension = + case Box.toList basisColumns of + [] -> 0 + firstBasisColumn : _ -> U.length firstBasisColumn + basisCount = projectedSubspaceDimension subspace + (rowCount, columnCount) = operatorShape op + in if rowCount /= columnCount || basisDimension /= columnCount + then + Left + ( InvariantViolation + ( "Projected eigensolve basis dimension mismatch: operator " + <> show (rowCount, columnCount) + <> " basis vectors of length " + <> show basisDimension + ) + ) + else Right basisCount + +validateProjectedRequestedCount :: String -> Int -> Int -> Either MoonlightError () +validateProjectedRequestedCount context requestedCount projectedDimension = + if requestedCount > projectedDimension + then + Left + ( InvariantViolation + ( context + <> " count exceeds projected dimension: requested " + <> show requestedCount + <> " from " + <> show projectedDimension + ) + ) + else Right () + +liftProjectedEigenpairs :: + LinearOperator 'SelfAdjointOperator -> + Box.Vector (U.Vector Double) -> + Eigenpairs -> + Either MoonlightError Eigenpairs +liftProjectedEigenpairs op basisColumns projectedPairs = do + let ambientDimension = snd (operatorShape op) + projectedValues = eigenpairValues projectedPairs + projectedCount = eigenpairCount projectedPairs + liftedColumns <- Box.generateM projectedCount (liftProjectedColumn op basisColumns projectedPairs) + liftedVectors <- flattenLiftedColumns ambientDimension projectedCount liftedColumns + liftedResiduals <- projectedResidualVector projectedCount liftedColumns + mkEigenpairs ambientDimension projectedValues liftedVectors liftedResiduals + +liftProjectedColumn :: + LinearOperator 'SelfAdjointOperator -> + Box.Vector (U.Vector Double) -> + Eigenpairs -> + Int -> + Either MoonlightError (U.Vector Double, Double) +liftProjectedColumn op basisColumns projectedPairs columnIndex = do + eigenvalue <- + case eigenpairValues projectedPairs U.!? columnIndex of + Nothing -> Left (InvariantViolation "projected eigenpair value index out of bounds") + Just value -> Right value + projectedResidualNorm <- + case eigenpairResidualNorms projectedPairs U.!? columnIndex of + Nothing -> Left (InvariantViolation "projected eigenpair residual index out of bounds") + Just value -> Right value + projectedVector <- eigenpairVectorAt columnIndex projectedPairs + liftProjectedMode op basisColumns eigenvalue projectedResidualNorm projectedVector + +liftProjectedMode :: + LinearOperator 'SelfAdjointOperator -> + Box.Vector (U.Vector Double) -> + Double -> + Double -> + U.Vector Double -> + Either MoonlightError (U.Vector Double, Double) +liftProjectedMode op basisColumns eigenvalue projectedResidualNorm projectedVector + | not (finiteDouble eigenvalue) = + Left (InvariantViolation "projected eigensolve produced a non-finite projected eigenvalue") + | not (finiteDouble projectedResidualNorm) = + Left (InvariantViolation "projected eigensolve produced a non-finite projected residual") + | otherwise = do + liftedVector <- linearCombinationColumnsU basisColumns projectedVector + let liftedNorm = normU liftedVector + breakdownThreshold = eigenvectorBreakdownThreshold (U.length liftedVector) + if not (finiteDouble liftedNorm) + then Left (InvariantViolation "projected eigensolve produced a non-finite lifted projected eigenvector norm") + else + if liftedNorm <= breakdownThreshold + then + Left + ( InvariantViolation + ( "projected eigensolve produced a numerically zero lifted projected eigenvector; norm=" + <> show liftedNorm + <> ", threshold=" + <> show breakdownThreshold + ) + ) + else do + let normalizedVector = scaleU (1.0 / liftedNorm) liftedVector + imageVector <- runOperatorU op normalizedVector + residualVector <- subScaledU imageVector eigenvalue normalizedVector + let residualNorm = max projectedResidualNorm (normU residualVector) + if finiteDouble residualNorm + then + pure (normalizedVector, residualNorm) + else Left (InvariantViolation "projected eigensolve produced a non-finite projected eigen residual") + +flattenLiftedColumns :: Int -> Int -> Box.Vector (U.Vector Double, Double) -> Either MoonlightError (U.Vector Double) +flattenLiftedColumns ambientDimension projectedCount liftedColumns = + U.generateM + (ambientDimension * projectedCount) + ( \offset -> + let (columnIndex, rowIndex) = offset `quotRem` ambientDimension + in case liftedColumns Box.!? columnIndex of + Nothing -> Left (InvariantViolation "lifted projected eigenpair column index out of bounds") + Just (liftedVector, _) -> + case liftedVector U.!? rowIndex of + Nothing -> Left (InvariantViolation "lifted projected eigenpair row index out of bounds") + Just entryValue -> Right entryValue + ) + +projectedResidualVector :: Int -> Box.Vector (U.Vector Double, Double) -> Either MoonlightError (U.Vector Double) +projectedResidualVector projectedCount liftedColumns = + U.generateM + projectedCount + ( \columnIndex -> + case liftedColumns Box.!? columnIndex of + Nothing -> Left (InvariantViolation "lifted projected eigenpair residual index out of bounds") + Just (_, residualNorm) -> Right residualNorm + ) + +eigenvectorBreakdownThreshold :: Int -> Double +eigenvectorBreakdownThreshold ambientDimension = + 128.0 * epsDouble * sqrt (fromIntegral (max 1 ambientDimension) :: Double) diff --git a/src-spectral/Moonlight/LinAlg/Pure/Krylov/SelectedTridiagonal.hs b/src-spectral/Moonlight/LinAlg/Pure/Krylov/SelectedTridiagonal.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Krylov/SelectedTridiagonal.hs @@ -0,0 +1,1061 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE StrictData #-} + +-- | Selected symmetric-tridiagonal spectral fast path. +module Moonlight.LinAlg.Pure.Krylov.SelectedTridiagonal + ( SelectedTridiagonalAttempt (..), + TridiagonalRejection (..), + selectedSymmetricTridiagonalEigenvalues, + selectedSymmetricTridiagonalEigenvaluesDirect, + selectedSymmetricTridiagonalEigenpairColumnsDirect, + selectedSymmetricTridiagonalEigenpairsDirect, + selectedSymmetricTridiagonalEigenpairsFromCSR, + symmetricTridiagonalFromCSR, + inverseIterationResidualToleranceBound, + ) +where + +import Data.Kind (Type) +import Data.Foldable (foldlM) +import Data.List (mapAccumL, sortBy) +import Data.Ord (comparing) +import qualified Data.Vector.Unboxed as U +import Moonlight.Core (MoonlightError (..), fieldValueValid) +import Moonlight.LinAlg.Internal.Eigen.Kernels (epsDouble, safeMinimumDouble) +import Moonlight.LinAlg.Internal.VectorOps (normU) +import Moonlight.LinAlg.Pure.Krylov.Selection + ( SpectrumEnd (..), + sortForSpectrumBy, + ) +import Moonlight.LinAlg.Pure.Sparse.Structured + ( TridiagonalRejection (..), + symmetricTridiagonalFromCSR, + ) +import Moonlight.LinAlg.Pure.Sparse.Types (SparseCSR) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal + ( SymmetricTridiagonal, + mkSymmetricTridiagonalVectors, + symmetricTridiagonalDiagonalVector, + symmetricTridiagonalDimension, + symmetricTridiagonalOffDiagonalVector, + isPathLaplacianTridiagonal, + ) +import Moonlight.LinAlg.Pure.Spectral.Result + ( Eigenpairs, + eigenpairsFromColumns, + ) +import Prelude + +type TridiagonalBlock :: Type +data TridiagonalBlock = TridiagonalBlock + { tridiagonalBlockStart :: !Int, + tridiagonalBlockDiagonal :: !(U.Vector Double), + tridiagonalBlockOffDiagonal :: !(U.Vector Double) + } + +type RankInterval :: Type +data RankInterval = RankInterval + { rankIntervalLowerBound :: !Double, + rankIntervalUpperBound :: !Double, + rankIntervalLowerCount :: !Int, + rankIntervalUpperCount :: !Int, + rankIntervalRanks :: ![Int], + rankIntervalIteration :: !Int + } + +type SelectedTridiagonalAttempt :: Type +data SelectedTridiagonalAttempt + = SelectedTridiagonalSolved !Eigenpairs + | SelectedTridiagonalNotApplicable !TridiagonalRejection + deriving stock (Eq, Show) + +type SelectedEigenvalue :: Type +data SelectedEigenvalue = SelectedEigenvalue + { selectedEigenvalueOrdinal :: !Int, + selectedEigenvalueValue :: !Double + } + deriving stock (Eq, Show) + +type ClusterBasis :: Type +data ClusterBasis = ClusterBasis + { clusterBasisVectors :: ![U.Vector Double], + clusterBasisColumns :: ![(Double, U.Vector Double, Double)] + } + deriving stock (Eq, Show) + +type InverseIterationState :: Type +data InverseIterationState + = InverseIterationSearching !(U.Vector Double) + | InverseIterationConverged !(U.Vector Double) !Double + deriving stock (Eq, Show) + +type SelectedTridiagonalPairObstruction :: Type +data SelectedTridiagonalPairObstruction + = SelectedTridiagonalInverseIterationNonConverged !Int !Double !Double + | SelectedTridiagonalSolveNonFinite !Int !Double + | SelectedTridiagonalVectorDegenerate !Int !Double + | SelectedTridiagonalClusterBasisUnstable !Int !Double + deriving stock (Eq, Show) + +selectedSymmetricTridiagonalEigenpairsFromCSR :: + SpectrumEnd -> + Int -> + SparseCSR Double -> + Either MoonlightError SelectedTridiagonalAttempt +selectedSymmetricTridiagonalEigenpairsFromCSR spectrumEnd requestedCount csrValue + | requestedCount <= 0 = Left (InvariantViolation "selected tridiagonal eigensolve requires a positive requested count") + | otherwise = do + selectedOperator <- symmetricTridiagonalFromCSR csrValue + case selectedOperator of + Left rejection -> Right (SelectedTridiagonalNotApplicable rejection) + Right tridiagonalValue -> + SelectedTridiagonalSolved + <$> selectedSymmetricTridiagonalEigenpairs + spectrumEnd + requestedCount + tridiagonalValue + +selectedSymmetricTridiagonalEigenvalues :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + Either MoonlightError (U.Vector Double) +selectedSymmetricTridiagonalEigenvalues spectrumEnd requestedCount tridiagonalValue + | requestedCount <= 0 = Left (InvariantViolation "selected tridiagonal eigenvalue solve requires a positive requested count") + | requestedCount > symmetricTridiagonalDimension tridiagonalValue = + Left (InvariantViolation "selected tridiagonal eigenvalue count exceeds operator dimension") + | otherwise = selectedSymmetricTridiagonalEigenvaluesChecked spectrumEnd requestedCount tridiagonalValue + +selectedSymmetricTridiagonalEigenvaluesDirect :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + Either MoonlightError (U.Vector Double) +selectedSymmetricTridiagonalEigenvaluesDirect spectrumEnd requestedCount tridiagonalValue + | requestedCount <= 0 = Left (InvariantViolation "selected tridiagonal eigenvalue solve requires a positive requested count") + | requestedCount > symmetricTridiagonalDimension tridiagonalValue = + Left (InvariantViolation "selected tridiagonal eigenvalue count exceeds operator dimension") + | otherwise = selectedSymmetricTridiagonalEigenvaluesChecked spectrumEnd requestedCount tridiagonalValue + +selectedSymmetricTridiagonalEigenpairsDirect :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + Either MoonlightError Eigenpairs +selectedSymmetricTridiagonalEigenpairsDirect spectrumEnd requestedCount tridiagonalValue + | requestedCount <= 0 = Left (InvariantViolation "selected tridiagonal eigenpair solve requires a positive requested count") + | otherwise = selectedSymmetricTridiagonalEigenpairs spectrumEnd requestedCount tridiagonalValue + +selectedSymmetricTridiagonalEigenpairColumnsDirect :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + Either MoonlightError [(Double, U.Vector Double, Double)] +selectedSymmetricTridiagonalEigenpairColumnsDirect spectrumEnd requestedCount tridiagonalValue + | requestedCount <= 0 = Left (InvariantViolation "selected tridiagonal eigenpair solve requires a positive requested count") + | requestedCount > symmetricTridiagonalDimension tridiagonalValue = + Left (InvariantViolation "selected tridiagonal eigenpair count exceeds operator dimension") + | otherwise = selectedSymmetricTridiagonalEigenpairColumns spectrumEnd requestedCount tridiagonalValue + +selectedSymmetricTridiagonalEigenpairs :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + Either MoonlightError Eigenpairs +selectedSymmetricTridiagonalEigenpairs spectrumEnd requestedCount tridiagonalValue = + let !matrixSize = U.length (symmetricTridiagonalDiagonalVector tridiagonalValue) + in if requestedCount > matrixSize + then Left (InvariantViolation "selected tridiagonal eigenpair count exceeds operator dimension") + else + eigenpairsFromColumns matrixSize + =<< selectedSymmetricTridiagonalEigenpairColumns spectrumEnd requestedCount tridiagonalValue + +selectedSymmetricTridiagonalEigenpairColumns :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + Either MoonlightError [(Double, U.Vector Double, Double)] +selectedSymmetricTridiagonalEigenpairColumns spectrumEnd requestedCount tridiagonalValue = + case pathLaplacianEigenpairs spectrumEnd requestedCount tridiagonalValue of + Just pathPairs -> Right pathPairs + Nothing -> + case diagonalOperatorEigenpairs spectrumEnd tridiagonalValue of + Just diagonalPairs -> Right (take requestedCount diagonalPairs) + Nothing -> + if U.any (== 0.0) (symmetricTridiagonalOffDiagonalVector tridiagonalValue) + then selectedReducibleTridiagonalEigenpairColumnsViaInverseIteration spectrumEnd requestedCount tridiagonalValue + else selectedTridiagonalEigenpairColumnsViaInverseIteration spectrumEnd requestedCount tridiagonalValue + +selectedSymmetricTridiagonalEigenvaluesChecked :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + Either MoonlightError (U.Vector Double) +selectedSymmetricTridiagonalEigenvaluesChecked spectrumEnd requestedCount tridiagonalValue = + case pathLaplacianEigenvalues spectrumEnd requestedCount tridiagonalValue of + Just pathValues -> Right pathValues + Nothing -> + case diagonalOperatorEigenvalues spectrumEnd tridiagonalValue of + Just diagonalValues -> Right (U.take requestedCount diagonalValues) + Nothing -> + if U.any (== 0.0) (symmetricTridiagonalOffDiagonalVector tridiagonalValue) + then selectedReducibleTridiagonalEigenvaluesViaSturm spectrumEnd requestedCount tridiagonalValue + else + Right + (selectedIrreducibleTridiagonalEigenvalues spectrumEnd requestedCount tridiagonalValue) + +pathLaplacianEigenvalues :: SpectrumEnd -> Int -> SymmetricTridiagonal -> Maybe (U.Vector Double) +pathLaplacianEigenvalues spectrumEnd boundedCount tridiagonalValue = + let !matrixSize = U.length (symmetricTridiagonalDiagonalVector tridiagonalValue) + in if isPathLaplacianTridiagonal tridiagonalValue + then + Just + ( U.generate + boundedCount + ( \entryIndex -> + pathLaplacianEigenvalueAt matrixSize $ + case spectrumEnd of + SmallestEigenvalues -> entryIndex + LargestEigenvalues -> matrixSize - entryIndex - 1 + ) + ) + else Nothing + +pathLaplacianEigenpairs :: SpectrumEnd -> Int -> SymmetricTridiagonal -> Maybe [(Double, U.Vector Double, Double)] +pathLaplacianEigenpairs spectrumEnd boundedCount tridiagonalValue = + let !matrixSize = U.length (symmetricTridiagonalDiagonalVector tridiagonalValue) + in if isPathLaplacianTridiagonal tridiagonalValue + then + Just + ( pathLaplacianEigenpairAt matrixSize + <$> case spectrumEnd of + SmallestEigenvalues -> [0 .. boundedCount - 1] + LargestEigenvalues -> [matrixSize - 1, matrixSize - 2 .. matrixSize - boundedCount] + ) + else Nothing + +pathLaplacianEigenpairAt :: Int -> Int -> (Double, U.Vector Double, Double) +pathLaplacianEigenpairAt !matrixSize !modeIndex = + let !theta = pi * fromIntegral modeIndex / fromIntegral (max 1 matrixSize) + !eigenvalue = pathLaplacianEigenvalueAt matrixSize modeIndex + !eigenvector = + if modeIndex == 0 + then U.replicate matrixSize (1.0 / sqrt (fromIntegral (max 1 matrixSize))) + else + let !normalizer = sqrt (2.0 / fromIntegral matrixSize) + in U.generate + matrixSize + (\rowIndex -> normalizer * cos (theta * (fromIntegral rowIndex + 0.5))) + !residualNorm = pathLaplacianResidualNorm matrixSize eigenvalue eigenvector + in (eigenvalue, eigenvector, residualNorm) + +pathLaplacianEigenvalueAt :: Int -> Int -> Double +pathLaplacianEigenvalueAt !matrixSize !modeIndex = + 2.0 - 2.0 * cos (pi * fromIntegral modeIndex / fromIntegral (max 1 matrixSize)) +{-# INLINE pathLaplacianEigenvalueAt #-} + +pathLaplacianResidualNorm :: Int -> Double -> U.Vector Double -> Double +pathLaplacianResidualNorm !matrixSize !eigenvalue eigenvector = + sqrt + ( U.ifoldl' + ( \ !squaredNorm !rowIndex _ -> + let !residualEntry = pathLaplacianResidualEntry matrixSize eigenvalue eigenvector rowIndex + in squaredNorm + residualEntry * residualEntry + ) + 0.0 + eigenvector + ) + +pathLaplacianResidualEntry :: Int -> Double -> U.Vector Double -> Int -> Double +pathLaplacianResidualEntry !matrixSize !eigenvalue eigenvector !rowIndex = + let !centerValue = eigenvector `U.unsafeIndex` rowIndex + !degree + | matrixSize == 1 = 0.0 + | rowIndex == 0 || rowIndex + 1 == matrixSize = 1.0 + | otherwise = 2.0 + !leftValue = + if rowIndex <= 0 + then 0.0 + else eigenvector `U.unsafeIndex` (rowIndex - 1) + !rightValue = + if rowIndex + 1 >= matrixSize + then 0.0 + else eigenvector `U.unsafeIndex` (rowIndex + 1) + !imageValue = degree * centerValue - leftValue - rightValue + in imageValue - eigenvalue * centerValue +{-# INLINE pathLaplacianResidualEntry #-} + +diagonalOperatorEigenvalues :: SpectrumEnd -> SymmetricTridiagonal -> Maybe (U.Vector Double) +diagonalOperatorEigenvalues spectrumEnd tridiagonalValue = + if U.all (== 0.0) (symmetricTridiagonalOffDiagonalVector tridiagonalValue) + then + Just + ( U.fromList + ( snd + <$> sortForSpectrum + spectrumEnd + (U.toList (U.indexed (symmetricTridiagonalDiagonalVector tridiagonalValue))) + ) + ) + else Nothing + +diagonalOperatorEigenpairs :: SpectrumEnd -> SymmetricTridiagonal -> Maybe [(Double, U.Vector Double, Double)] +diagonalOperatorEigenpairs spectrumEnd tridiagonalValue = + if U.all (== 0.0) (symmetricTridiagonalOffDiagonalVector tridiagonalValue) + then + Just + ( fmap + ( \(entryIndex, eigenvalue) -> + ( eigenvalue, + unitVector (U.length (symmetricTridiagonalDiagonalVector tridiagonalValue)) entryIndex, + 0.0 + ) + ) + ( sortForSpectrum + spectrumEnd + (U.toList (U.indexed (symmetricTridiagonalDiagonalVector tridiagonalValue))) + ) + ) + else Nothing + +selectedIrreducibleTridiagonalEigenvalues :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + U.Vector Double +selectedIrreducibleTridiagonalEigenvalues spectrumEnd boundedCount tridiagonalValue = + let !matrixSize = U.length (symmetricTridiagonalDiagonalVector tridiagonalValue) + selectedRanks = + case spectrumEnd of + SmallestEigenvalues -> [1 .. boundedCount] + LargestEigenvalues -> [matrixSize - boundedCount + 1 .. matrixSize] + selectedValues = U.fromList (batchedBisectEigenvaluesAtRanks tridiagonalValue selectedRanks) + in case spectrumEnd of + SmallestEigenvalues -> selectedValues + LargestEigenvalues -> U.reverse selectedValues + +batchedBisectEigenvaluesAtRanks :: SymmetricTridiagonal -> [Int] -> [Double] +batchedBisectEigenvaluesAtRanks tridiagonalValue selectedRanks = + let (!initialLower, !initialUpper) = gershgorinBounds tridiagonalValue + !matrixSize = U.length (symmetricTridiagonalDiagonalVector tridiagonalValue) + !matrixScale = tridiagonalInfinityNormBound tridiagonalValue + initialInterval = + RankInterval + { rankIntervalLowerBound = initialLower, + rankIntervalUpperBound = initialUpper, + rankIntervalLowerCount = 0, + rankIntervalUpperCount = matrixSize, + rankIntervalRanks = selectedRanks, + rankIntervalIteration = 0 + } + in snd + <$> sortBy + (comparing fst) + (refineRankInterval matrixScale tridiagonalValue initialInterval) + +refineRankInterval :: Double -> SymmetricTridiagonal -> RankInterval -> [(Int, Double)] +refineRankInterval !matrixScale tridiagonalValue interval + | null (rankIntervalRanks interval) = [] + | rankIntervalIteration interval >= tridiagonalBisectionIterationLimit = + finalizeRankInterval interval + | rankIntervalUpperBound interval - rankIntervalLowerBound interval + <= eigenTolerance matrixScale (rankIntervalLowerBound interval) (rankIntervalUpperBound interval) = + finalizeRankInterval interval + | rankIntervalUpperBound interval == rankIntervalLowerBound interval = + finalizeRankInterval interval + | otherwise = + concatMap + (refineRankInterval matrixScale tridiagonalValue) + (splitRankInterval matrixScale tridiagonalValue interval) + +splitRankInterval :: Double -> SymmetricTridiagonal -> RankInterval -> [RankInterval] +splitRankInterval !matrixScale tridiagonalValue interval = + maybeInterval + lowerRanks + (rankIntervalLowerBound interval) + middleValue + (rankIntervalLowerCount interval) + middleCount + <> maybeInterval + upperRanks + middleValue + (rankIntervalUpperBound interval) + middleCount + (rankIntervalUpperCount interval) + where + !middleValue = midpoint (rankIntervalLowerBound interval) (rankIntervalUpperBound interval) + !middleCount = + clamp + (rankIntervalLowerCount interval) + (rankIntervalUpperCount interval) + (sturmCountLessEqual matrixScale tridiagonalValue middleValue) + lowerRanks = filter (<= middleCount) (rankIntervalRanks interval) + upperRanks = filter (> middleCount) (rankIntervalRanks interval) + maybeInterval ranks lowerBound upperBound lowerCount upperCount = + if null ranks + then [] + else + [ RankInterval + { rankIntervalLowerBound = lowerBound, + rankIntervalUpperBound = upperBound, + rankIntervalLowerCount = lowerCount, + rankIntervalUpperCount = upperCount, + rankIntervalRanks = ranks, + rankIntervalIteration = rankIntervalIteration interval + 1 + } + ] + +finalizeRankInterval :: RankInterval -> [(Int, Double)] +finalizeRankInterval interval = + (\rankValue -> (rankValue, midpoint (rankIntervalLowerBound interval) (rankIntervalUpperBound interval))) + <$> rankIntervalRanks interval + +tridiagonalBisectionIterationLimit :: Int +tridiagonalBisectionIterationLimit = 80 +{-# INLINE tridiagonalBisectionIterationLimit #-} + +selectedTridiagonalEigenpairColumnsViaInverseIteration :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + Either MoonlightError [(Double, U.Vector Double, Double)] +selectedTridiagonalEigenpairColumnsViaInverseIteration spectrumEnd requestedCount tridiagonalValue = + selectedTridiagonalPairResultToEither + ( selectedEigenpairColumnsFromValues + tridiagonalValue + ( U.toList + (selectedIrreducibleTridiagonalEigenvalues spectrumEnd requestedCount tridiagonalValue) + ) + ) + +selectedReducibleTridiagonalEigenvaluesViaSturm :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + Either MoonlightError (U.Vector Double) +selectedReducibleTridiagonalEigenvaluesViaSturm spectrumEnd requestedCount tridiagonalValue = + fmap + (U.fromList . take requestedCount . sortForSpectrumBy spectrumEnd id . concat) + (traverse (blockEigenvaluesViaSturm spectrumEnd requestedCount) (tridiagonalBlocks tridiagonalValue)) + +selectedReducibleTridiagonalEigenpairColumnsViaInverseIteration :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + Either MoonlightError [(Double, U.Vector Double, Double)] +selectedReducibleTridiagonalEigenpairColumnsViaInverseIteration spectrumEnd requestedCount tridiagonalValue = + fmap + ( take requestedCount + . sortForSpectrumBy spectrumEnd (\(eigenvalue, _, _) -> eigenvalue) + . concat + ) + (traverse (blockEigenpairColumnsViaInverseIteration spectrumEnd requestedCount tridiagonalValue) (tridiagonalBlocks tridiagonalValue)) + +blockEigenvaluesViaSturm :: SpectrumEnd -> Int -> TridiagonalBlock -> Either MoonlightError [Double] +blockEigenvaluesViaSturm spectrumEnd requestedCount blockValue = do + blockTridiagonal <- + mkSymmetricTridiagonalVectors + (tridiagonalBlockDiagonal blockValue) + (tridiagonalBlockOffDiagonal blockValue) + let blockRequestedCount = min requestedCount (symmetricTridiagonalDimension blockTridiagonal) + Right (U.toList (selectedIrreducibleTridiagonalEigenvalues spectrumEnd blockRequestedCount blockTridiagonal)) + +blockEigenpairColumnsViaInverseIteration :: + SpectrumEnd -> + Int -> + SymmetricTridiagonal -> + TridiagonalBlock -> + Either MoonlightError [(Double, U.Vector Double, Double)] +blockEigenpairColumnsViaInverseIteration spectrumEnd requestedCount tridiagonalValue blockValue = do + blockTridiagonal <- + mkSymmetricTridiagonalVectors + (tridiagonalBlockDiagonal blockValue) + (tridiagonalBlockOffDiagonal blockValue) + let blockRequestedCount = min requestedCount (symmetricTridiagonalDimension blockTridiagonal) + selectedTridiagonalPairResultToEither + ( fmap + (fmap (embedBlockEigenpairColumn tridiagonalValue blockValue)) + ( selectedEigenpairColumnsFromValues + blockTridiagonal + ( U.toList + (selectedIrreducibleTridiagonalEigenvalues spectrumEnd blockRequestedCount blockTridiagonal) + ) + ) + ) + +embedBlockEigenpairColumn :: + SymmetricTridiagonal -> + TridiagonalBlock -> + (Double, U.Vector Double, Double) -> + (Double, U.Vector Double, Double) +embedBlockEigenpairColumn tridiagonalValue blockValue (eigenvalue, blockVector, _) = + let eigenvector = + embedBlockVector + (symmetricTridiagonalDimension tridiagonalValue) + (tridiagonalBlockStart blockValue) + blockVector + in (eigenvalue, eigenvector, tridiagonalResidualNorm tridiagonalValue eigenvalue eigenvector) + +embedBlockVector :: Int -> Int -> U.Vector Double -> U.Vector Double +embedBlockVector dimension startOffset blockVector = + U.generate + dimension + ( \entryIndex -> + if entryIndex >= startOffset && entryIndex < startOffset + U.length blockVector + then vectorEntryOrZero blockVector (entryIndex - startOffset) + else 0.0 + ) + +tridiagonalBlocks :: SymmetricTridiagonal -> [TridiagonalBlock] +tridiagonalBlocks tridiagonalValue = + makeBlock <$> blockRanges (symmetricTridiagonalOffDiagonalVector tridiagonalValue) (symmetricTridiagonalDimension tridiagonalValue) + where + diagonalEntries = symmetricTridiagonalDiagonalVector tridiagonalValue + offDiagonalEntries = symmetricTridiagonalOffDiagonalVector tridiagonalValue + makeBlock (startIndex, stopIndex) = + let blockSize = stopIndex - startIndex + in TridiagonalBlock + { tridiagonalBlockStart = startIndex, + tridiagonalBlockDiagonal = U.slice startIndex blockSize diagonalEntries, + tridiagonalBlockOffDiagonal = U.slice startIndex (max 0 (blockSize - 1)) offDiagonalEntries + } + +blockRanges :: U.Vector Double -> Int -> [(Int, Int)] +blockRanges offDiagonalEntries dimension = + filter + (\(startIndex, stopIndex) -> startIndex < stopIndex) + (zip splitStarts splitStops) + where + zeroIndices = zeroCouplingIndices offDiagonalEntries + splitStarts = 0 : fmap (+ 1) zeroIndices + splitStops = fmap (+ 1) zeroIndices <> [dimension] + +zeroCouplingIndices :: U.Vector Double -> [Int] +zeroCouplingIndices offDiagonalEntries = + fst <$> filter ((== 0.0) . snd) (U.toList (U.indexed offDiagonalEntries)) + +vectorEntryOrZero :: U.Vector Double -> Int -> Double +vectorEntryOrZero values indexValue = + maybe 0.0 id (values U.!? indexValue) +{-# INLINE vectorEntryOrZero #-} + +selectedEigenpairColumnsFromValues :: + SymmetricTridiagonal -> + [Double] -> + Either SelectedTridiagonalPairObstruction [(Double, U.Vector Double, Double)] +selectedEigenpairColumnsFromValues tridiagonalValue eigenvalues = + fmap + concat + ( traverse + (solveSelectedEigenvalueCluster tridiagonalValue) + (clusterSelectedEigenvalues (tridiagonalInfinityNormBound tridiagonalValue) (zipWith SelectedEigenvalue [0 ..] eigenvalues)) + ) + +solveSelectedEigenvalueCluster :: + SymmetricTridiagonal -> + [SelectedEigenvalue] -> + Either SelectedTridiagonalPairObstruction [(Double, U.Vector Double, Double)] +solveSelectedEigenvalueCluster tridiagonalValue eigenvalueCluster = + clusterBasisColumns + <$> foldlM + (appendSelectedEigenpairColumn tridiagonalValue) + ClusterBasis {clusterBasisVectors = [], clusterBasisColumns = []} + eigenvalueCluster + +appendSelectedEigenpairColumn :: + SymmetricTridiagonal -> + ClusterBasis -> + SelectedEigenvalue -> + Either SelectedTridiagonalPairObstruction ClusterBasis +appendSelectedEigenpairColumn tridiagonalValue basis selectedValue = do + column@(_, eigenvector, _) <- + solveSelectedEigenpairColumn + tridiagonalValue + (clusterBasisVectors basis) + selectedValue + Right + basis + { clusterBasisVectors = clusterBasisVectors basis <> [eigenvector], + clusterBasisColumns = clusterBasisColumns basis <> [column] + } + +solveSelectedEigenpairColumn :: + SymmetricTridiagonal -> + [U.Vector Double] -> + SelectedEigenvalue -> + Either SelectedTridiagonalPairObstruction (Double, U.Vector Double, Double) +solveSelectedEigenpairColumn tridiagonalValue clusterVectors selectedValue = + let !matrixScale = tridiagonalInfinityNormBound tridiagonalValue + !eigenvalue = selectedEigenvalueValue selectedValue + !ordinal = selectedEigenvalueOrdinal selectedValue + attempts = + inverseIterationAttempt + tridiagonalValue + matrixScale + clusterVectors + selectedValue + <$> inverseIterationShiftSchedule matrixScale eigenvalue ordinal + in firstSuccessfulAttempt + (SelectedTridiagonalInverseIterationNonConverged ordinal eigenvalue (inverseIterationResidualTolerance matrixScale eigenvalue tridiagonalValue)) + attempts + +inverseIterationAttempt :: + SymmetricTridiagonal -> + Double -> + [U.Vector Double] -> + SelectedEigenvalue -> + Double -> + Either SelectedTridiagonalPairObstruction (Double, U.Vector Double, Double) +inverseIterationAttempt tridiagonalValue !matrixScale clusterVectors selectedValue !shiftValue = do + let !eigenvalue = selectedEigenvalueValue selectedValue + !ordinal = selectedEigenvalueOrdinal selectedValue + !initialVector = + inverseIterationSeed + (symmetricTridiagonalDimension tridiagonalValue) + ordinal + !residualLimit = inverseIterationResidualTolerance matrixScale eigenvalue tridiagonalValue + finalState <- + foldlM + (inverseIterationStep tridiagonalValue matrixScale clusterVectors selectedValue shiftValue residualLimit) + (InverseIterationSearching initialVector) + [1 .. inverseIterationStepLimit] + case finalState of + InverseIterationConverged eigenvector residualNorm -> + Right (tridiagonalRayleighQuotient tridiagonalValue eigenvector, eigenvector, residualNorm) + InverseIterationSearching eigenvector -> + let !certifiedEigenvalue = tridiagonalRayleighQuotient tridiagonalValue eigenvector + !residualNorm = tridiagonalResidualNorm tridiagonalValue certifiedEigenvalue eigenvector + in Left (SelectedTridiagonalInverseIterationNonConverged ordinal eigenvalue residualNorm) + +inverseIterationStep :: + SymmetricTridiagonal -> + Double -> + [U.Vector Double] -> + SelectedEigenvalue -> + Double -> + Double -> + InverseIterationState -> + Int -> + Either SelectedTridiagonalPairObstruction InverseIterationState +inverseIterationStep _ _ _ _ _ _ converged@(InverseIterationConverged _ _) _ = + Right converged +inverseIterationStep tridiagonalValue !matrixScale clusterVectors selectedValue !shiftValue !residualLimit (InverseIterationSearching eigenvector) _ = do + solvedVector <- + solveShiftedTridiagonal + matrixScale + tridiagonalValue + selectedValue + shiftValue + eigenvector + normalizedVector <- + normalizeClusterVector + matrixScale + selectedValue + clusterVectors + solvedVector + let !residualNorm = + tridiagonalResidualNorm + tridiagonalValue + (tridiagonalRayleighQuotient tridiagonalValue normalizedVector) + normalizedVector + Right + ( if residualNorm <= residualLimit + then InverseIterationConverged normalizedVector residualNorm + else InverseIterationSearching normalizedVector + ) + +solveShiftedTridiagonal :: + Double -> + SymmetricTridiagonal -> + SelectedEigenvalue -> + Double -> + U.Vector Double -> + Either SelectedTridiagonalPairObstruction (U.Vector Double) +solveShiftedTridiagonal !matrixScale tridiagonalValue selectedValue !shiftValue rhsVector = + let diagonalEntries = symmetricTridiagonalDiagonalVector tridiagonalValue + offDiagonalEntries = symmetricTridiagonalOffDiagonalVector tridiagonalValue + !matrixSize = U.length diagonalEntries + forwardStep (!previousUpper, !previousRhs) !entryIndex = + let !diagonalPivot = (diagonalEntries `U.unsafeIndex` entryIndex) - shiftValue + !lowerEntry = + if entryIndex <= 0 + then 0.0 + else offDiagonalEntries `U.unsafeIndex` (entryIndex - 1) + !rawPivot = diagonalPivot - lowerEntry * previousUpper + !pivotValue = safeTridiagonalSolvePivot matrixScale rawPivot + !upperEntry = + if entryIndex + 1 >= matrixSize + then 0.0 + else offDiagonalEntries `U.unsafeIndex` entryIndex + !forwardUpper = upperEntry / pivotValue + !forwardRhs = ((rhsVector `U.unsafeIndex` entryIndex) - lowerEntry * previousRhs) / pivotValue + in ((forwardUpper, forwardRhs), (forwardUpper, forwardRhs)) + (_, forwardValues) = + mapAccumL + forwardStep + (0.0, 0.0) + [0 .. matrixSize - 1] + (solutionValues, _) = + foldr + ( \(!forwardUpper, !forwardRhs) (!accumulatedValues, !nextValue) -> + let !solutionValue = forwardRhs - forwardUpper * nextValue + in (solutionValue : accumulatedValues, solutionValue) + ) + ([], 0.0) + forwardValues + solutionVector = U.fromList solutionValues + in if U.all fieldValueValid solutionVector + then Right solutionVector + else Left (SelectedTridiagonalSolveNonFinite (selectedEigenvalueOrdinal selectedValue) (selectedEigenvalueValue selectedValue)) + +normalizeClusterVector :: + Double -> + SelectedEigenvalue -> + [U.Vector Double] -> + U.Vector Double -> + Either SelectedTridiagonalPairObstruction (U.Vector Double) +normalizeClusterVector !matrixScale selectedValue clusterVectors candidateVector = do + firstPass <- + normalizeSelectedVector + matrixScale + selectedValue + (orthogonalizeAgainst clusterVectors candidateVector) + secondPass <- + normalizeSelectedVector + matrixScale + selectedValue + (orthogonalizeAgainst clusterVectors firstPass) + let !largestOverlap = + maximum + (0.0 : (abs . vectorDot secondPass <$> clusterVectors)) + in if largestOverlap <= clusterOrthogonalityTolerance matrixScale (U.length secondPass) + then Right secondPass + else Left (SelectedTridiagonalClusterBasisUnstable (selectedEigenvalueOrdinal selectedValue) (selectedEigenvalueValue selectedValue)) + +normalizeSelectedVector :: + Double -> + SelectedEigenvalue -> + U.Vector Double -> + Either SelectedTridiagonalPairObstruction (U.Vector Double) +normalizeSelectedVector !matrixScale selectedValue vectorValue = + let !vectorNorm = normU vectorValue + in if fieldValueValid vectorNorm && vectorNorm > vectorNormTolerance matrixScale (U.length vectorValue) + then Right (U.map (/ vectorNorm) vectorValue) + else Left (SelectedTridiagonalVectorDegenerate (selectedEigenvalueOrdinal selectedValue) (selectedEigenvalueValue selectedValue)) + +orthogonalizeAgainst :: [U.Vector Double] -> U.Vector Double -> U.Vector Double +orthogonalizeAgainst basisVectors vectorValue = + foldl' + ( \candidateVector basisVector -> + let !projectionScale = vectorDot candidateVector basisVector + in U.zipWith + (\candidateEntry basisEntry -> candidateEntry - projectionScale * basisEntry) + candidateVector + basisVector + ) + vectorValue + basisVectors + +vectorDot :: U.Vector Double -> U.Vector Double -> Double +vectorDot leftVector rightVector = + U.sum (U.zipWith (*) leftVector rightVector) +{-# INLINE vectorDot #-} + +inverseIterationSeed :: Int -> Int -> U.Vector Double +inverseIterationSeed !matrixSize !ordinal = + U.generate + matrixSize + ( \entryIndex -> + let !phase = + fromIntegral ((entryIndex + 1) * (ordinal + 1)) + * pi + / fromIntegral (matrixSize + ordinal + 2) + in sin phase + 0.5 * cos (phase * 0.5) + ) + +clusterSelectedEigenvalues :: Double -> [SelectedEigenvalue] -> [[SelectedEigenvalue]] +clusterSelectedEigenvalues !matrixScale = + reverse + . fmap reverse + . foldl' appendEigenvalueCluster [] + where + appendEigenvalueCluster [] eigenvalue = [[eigenvalue]] + appendEigenvalueCluster (cluster@(previousEigenvalue : _) : restClusters) eigenvalue + | eigenvalueGapInCluster matrixScale previousEigenvalue eigenvalue = + (eigenvalue : cluster) : restClusters + | otherwise = [eigenvalue] : cluster : restClusters + appendEigenvalueCluster ([] : restClusters) eigenvalue = [eigenvalue] : restClusters + +eigenvalueGapInCluster :: Double -> SelectedEigenvalue -> SelectedEigenvalue -> Bool +eigenvalueGapInCluster !matrixScale leftValue rightValue = + abs (selectedEigenvalueValue leftValue - selectedEigenvalueValue rightValue) + <= eigenvalueClusterTolerance matrixScale (selectedEigenvalueValue leftValue) (selectedEigenvalueValue rightValue) + +inverseIterationShiftSchedule :: Double -> Double -> Int -> [Double] +inverseIterationShiftSchedule !matrixScale !eigenvalue !ordinal = + (eigenvalue +) + <$> fmap + (* inverseIterationShiftUnit matrixScale eigenvalue) + (0.0 : concatMap signedShift [1 .. inverseIterationShiftAttemptLimit]) + where + signedShift attemptIndex = + let !shiftMagnitude = fromIntegral attemptIndex + in if even (ordinal + attemptIndex) + then [shiftMagnitude, negate shiftMagnitude] + else [negate shiftMagnitude, shiftMagnitude] + +firstSuccessfulAttempt :: SelectedTridiagonalPairObstruction -> [Either SelectedTridiagonalPairObstruction value] -> Either SelectedTridiagonalPairObstruction value +firstSuccessfulAttempt fallbackObstruction = + foldr + ( \attemptValue remainingAttempts -> + case attemptValue of + Right resultValue -> Right resultValue + Left _ -> remainingAttempts + ) + (Left fallbackObstruction) + +selectedTridiagonalPairResultToEither :: Either SelectedTridiagonalPairObstruction value -> Either MoonlightError value +selectedTridiagonalPairResultToEither resultValue = + case resultValue of + Right value -> Right value + Left obstruction -> Left (InvariantViolation (renderSelectedTridiagonalPairObstruction obstruction)) + +renderSelectedTridiagonalPairObstruction :: SelectedTridiagonalPairObstruction -> String +renderSelectedTridiagonalPairObstruction obstruction = + case obstruction of + SelectedTridiagonalInverseIterationNonConverged ordinal eigenvalue residualNorm -> + "selected tridiagonal inverse iteration did not converge at ordinal " + <> show ordinal + <> " for eigenvalue " + <> show eigenvalue + <> " with residual " + <> show residualNorm + SelectedTridiagonalSolveNonFinite ordinal eigenvalue -> + "selected tridiagonal inverse iteration produced a non-finite solve at ordinal " + <> show ordinal + <> " for eigenvalue " + <> show eigenvalue + SelectedTridiagonalVectorDegenerate ordinal eigenvalue -> + "selected tridiagonal inverse iteration produced a degenerate vector at ordinal " + <> show ordinal + <> " for eigenvalue " + <> show eigenvalue + SelectedTridiagonalClusterBasisUnstable ordinal eigenvalue -> + "selected tridiagonal clustered basis could not be stabilized at ordinal " + <> show ordinal + <> " for eigenvalue " + <> show eigenvalue + +inverseIterationStepLimit :: Int +inverseIterationStepLimit = 16 +{-# INLINE inverseIterationStepLimit #-} + +inverseIterationShiftAttemptLimit :: Int +inverseIterationShiftAttemptLimit = 4 +{-# INLINE inverseIterationShiftAttemptLimit #-} + +inverseIterationShiftUnit :: Double -> Double -> Double +inverseIterationShiftUnit !matrixScale !eigenvalue = + 16.0 * epsDouble * max 1.0 (max matrixScale (abs eigenvalue)) +{-# INLINE inverseIterationShiftUnit #-} + +inverseIterationResidualTolerance :: Double -> Double -> SymmetricTridiagonal -> Double +inverseIterationResidualTolerance !matrixScale !eigenvalue tridiagonalValue = + inverseIterationResidualToleranceBound matrixScale eigenvalue (symmetricTridiagonalDimension tridiagonalValue) +{-# INLINE inverseIterationResidualTolerance #-} + +inverseIterationResidualToleranceBound :: Double -> Double -> Int -> Double +inverseIterationResidualToleranceBound !matrixScale !eigenvalue !dimension = + 1.0e7 + * epsDouble + * max 1.0 (fromIntegral dimension) + * max 1.0 (max matrixScale (abs eigenvalue)) +{-# INLINE inverseIterationResidualToleranceBound #-} + +eigenvalueClusterTolerance :: Double -> Double -> Double -> Double +eigenvalueClusterTolerance !matrixScale !leftValue !rightValue = + 64.0 * sqrt epsDouble * max 1.0 (maximum [matrixScale, abs leftValue, abs rightValue]) +{-# INLINE eigenvalueClusterTolerance #-} + +clusterOrthogonalityTolerance :: Double -> Int -> Double +clusterOrthogonalityTolerance _ !matrixSize = + 256.0 * sqrt epsDouble * max 1.0 (fromIntegral matrixSize) +{-# INLINE clusterOrthogonalityTolerance #-} + +vectorNormTolerance :: Double -> Int -> Double +vectorNormTolerance !matrixScale !matrixSize = + 64.0 * safeMinimumDouble * max 1.0 matrixScale * max 1.0 (fromIntegral matrixSize) +{-# INLINE vectorNormTolerance #-} + +safeTridiagonalSolvePivot :: Double -> Double -> Double +safeTridiagonalSolvePivot !matrixScale !pivotValue + | abs pivotValue > solvePivotTolerance matrixScale = pivotValue + | pivotValue < 0.0 = negate (solvePivotTolerance matrixScale) + | otherwise = solvePivotTolerance matrixScale +{-# INLINE safeTridiagonalSolvePivot #-} + +solvePivotTolerance :: Double -> Double +solvePivotTolerance !matrixScale = + (128.0 * epsDouble * max 1.0 matrixScale) + safeMinimumDouble +{-# INLINE solvePivotTolerance #-} + +sturmCountLessEqual :: Double -> SymmetricTridiagonal -> Double -> Int +sturmCountLessEqual !matrixScale tridiagonalValue !shiftValue = + let diagonalEntries = symmetricTridiagonalDiagonalVector tridiagonalValue + offDiagonalEntries = symmetricTridiagonalOffDiagonalVector tridiagonalValue + !matrixSize = U.length diagonalEntries + pivotAt !indexValue !previousPivot = + let !diagonalPivot = (diagonalEntries `U.unsafeIndex` indexValue) - shiftValue + in if indexValue == 0 + then diagonalPivot + else + let !offDiagonal = offDiagonalEntries `U.unsafeIndex` (indexValue - 1) + !safePreviousPivot = nonzeroSturmPivot matrixScale previousPivot + in diagonalPivot - (offDiagonal * offDiagonal / safePreviousPivot) + countAt !indexValue !previousPivot !negativeCount + | indexValue >= matrixSize = negativeCount + | otherwise = + let !pivotValue = pivotAt indexValue previousPivot + !nextCount = + if pivotValue <= 0.0 + then negativeCount + 1 + else negativeCount + in countAt (indexValue + 1) pivotValue nextCount + in countAt 0 1.0 0 + +tridiagonalResidualNorm :: SymmetricTridiagonal -> Double -> U.Vector Double -> Double +tridiagonalResidualNorm tridiagonalValue !eigenvalue eigenvector = + normU + ( U.generate + (U.length eigenvector) + (tridiagonalResidualEntry tridiagonalValue eigenvalue eigenvector) + ) + +tridiagonalRayleighQuotient :: SymmetricTridiagonal -> U.Vector Double -> Double +tridiagonalRayleighQuotient tridiagonalValue eigenvector = + vectorDot eigenvector (tridiagonalApply tridiagonalValue eigenvector) + +tridiagonalApply :: SymmetricTridiagonal -> U.Vector Double -> U.Vector Double +tridiagonalApply tridiagonalValue eigenvector = + U.generate + (U.length eigenvector) + (tridiagonalApplyEntry tridiagonalValue eigenvector) + +tridiagonalApplyEntry :: SymmetricTridiagonal -> U.Vector Double -> Int -> Double +tridiagonalApplyEntry tridiagonalValue eigenvector !entryIndex = + let diagonalEntries = symmetricTridiagonalDiagonalVector tridiagonalValue + offDiagonalEntries = symmetricTridiagonalOffDiagonalVector tridiagonalValue + !matrixSize = U.length diagonalEntries + !centerValue = eigenvector `U.unsafeIndex` entryIndex + !leftValue = + if entryIndex <= 0 + then 0.0 + else (offDiagonalEntries `U.unsafeIndex` (entryIndex - 1)) * (eigenvector `U.unsafeIndex` (entryIndex - 1)) + !rightValue = + if entryIndex + 1 >= matrixSize + then 0.0 + else (offDiagonalEntries `U.unsafeIndex` entryIndex) * (eigenvector `U.unsafeIndex` (entryIndex + 1)) + in leftValue + (diagonalEntries `U.unsafeIndex` entryIndex) * centerValue + rightValue +{-# INLINE tridiagonalApplyEntry #-} + +tridiagonalResidualEntry :: SymmetricTridiagonal -> Double -> U.Vector Double -> Int -> Double +tridiagonalResidualEntry tridiagonalValue !eigenvalue eigenvector !entryIndex = + let diagonalEntries = symmetricTridiagonalDiagonalVector tridiagonalValue + offDiagonalEntries = symmetricTridiagonalOffDiagonalVector tridiagonalValue + !matrixSize = U.length diagonalEntries + !centerValue = eigenvector `U.unsafeIndex` entryIndex + !leftValue = + if entryIndex <= 0 + then 0.0 + else (offDiagonalEntries `U.unsafeIndex` (entryIndex - 1)) * (eigenvector `U.unsafeIndex` (entryIndex - 1)) + !rightValue = + if entryIndex + 1 >= matrixSize + then 0.0 + else (offDiagonalEntries `U.unsafeIndex` entryIndex) * (eigenvector `U.unsafeIndex` (entryIndex + 1)) + !imageValue = leftValue + (diagonalEntries `U.unsafeIndex` entryIndex) * centerValue + rightValue + in imageValue - eigenvalue * centerValue +{-# INLINE tridiagonalResidualEntry #-} + +gershgorinBounds :: SymmetricTridiagonal -> (Double, Double) +gershgorinBounds tridiagonalValue = + let diagonalEntries = symmetricTridiagonalDiagonalVector tridiagonalValue + offDiagonalEntries = symmetricTridiagonalOffDiagonalVector tridiagonalValue + !matrixSize = U.length diagonalEntries + rowLowerBound !rowIndex = + let !radius = offDiagonalRadius offDiagonalEntries matrixSize rowIndex + in (diagonalEntries `U.unsafeIndex` rowIndex) - radius + rowUpperBound !rowIndex = + let !radius = offDiagonalRadius offDiagonalEntries matrixSize rowIndex + in (diagonalEntries `U.unsafeIndex` rowIndex) + radius + lowerBound = U.minimum (U.generate matrixSize rowLowerBound) + upperBound = U.maximum (U.generate matrixSize rowUpperBound) + margin = 16.0 * eigenTolerance (tridiagonalInfinityNormBound tridiagonalValue) lowerBound upperBound + in (lowerBound - margin, upperBound + margin) + +tridiagonalInfinityNormBound :: SymmetricTridiagonal -> Double +tridiagonalInfinityNormBound tridiagonalValue = + let diagonalEntries = symmetricTridiagonalDiagonalVector tridiagonalValue + offDiagonalEntries = symmetricTridiagonalOffDiagonalVector tridiagonalValue + !matrixSize = U.length diagonalEntries + in if matrixSize <= 0 + then 0.0 + else + U.maximum + ( U.generate + matrixSize + ( \rowIndex -> + abs (diagonalEntries `U.unsafeIndex` rowIndex) + + offDiagonalRadius offDiagonalEntries matrixSize rowIndex + ) + ) + +offDiagonalRadius :: U.Vector Double -> Int -> Int -> Double +offDiagonalRadius offDiagonalEntries !matrixSize !rowIndex = + ( if rowIndex <= 0 + then 0.0 + else abs (offDiagonalEntries `U.unsafeIndex` (rowIndex - 1)) + ) + + ( if rowIndex + 1 >= matrixSize + then 0.0 + else abs (offDiagonalEntries `U.unsafeIndex` rowIndex) + ) +{-# INLINE offDiagonalRadius #-} + +sortForSpectrum :: SpectrumEnd -> [(Int, Double)] -> [(Int, Double)] +sortForSpectrum spectrumEnd = + sortBy + ( case spectrumEnd of + SmallestEigenvalues -> comparing snd + LargestEigenvalues -> flip (comparing snd) + ) + +unitVector :: Int -> Int -> U.Vector Double +unitVector !matrixSize !selectedIndex = + U.generate matrixSize (\entryIndex -> if entryIndex == selectedIndex then 1.0 else 0.0) + +midpoint :: Double -> Double -> Double +midpoint !leftValue !rightValue = + leftValue + 0.5 * (rightValue - leftValue) +{-# INLINE midpoint #-} + +clamp :: Ord value => value -> value -> value -> value +clamp lowerValue upperValue value = + max lowerValue (min upperValue value) +{-# INLINE clamp #-} + +eigenTolerance :: Double -> Double -> Double -> Double +eigenTolerance !matrixScale !leftValue !rightValue = + sqrt epsDouble * max 1.0 (maximum [matrixScale, abs leftValue, abs rightValue]) +{-# INLINE eigenTolerance #-} + +nonzeroSturmPivot :: Double -> Double -> Double +nonzeroSturmPivot !matrixScale !pivotValue + | abs pivotValue > sturmPivotTolerance matrixScale = pivotValue + | pivotValue > 0.0 = sturmPivotTolerance matrixScale + | otherwise = negate (sturmPivotTolerance matrixScale) +{-# INLINE nonzeroSturmPivot #-} + +sturmPivotTolerance :: Double -> Double +sturmPivotTolerance !matrixScale = + (64.0 * epsDouble * max 1.0 matrixScale) + safeMinimumDouble +{-# INLINE sturmPivotTolerance #-} diff --git a/src-spectral/Moonlight/LinAlg/Pure/Krylov/Selection.hs b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Selection.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Krylov/Selection.hs @@ -0,0 +1,28 @@ +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Krylov.Selection + ( SpectrumEnd (..), + sortForSpectrumBy, + sortRawPairsForSpectrum, + ) +where + +import Data.Kind (Type) +import Data.List (sortBy) +import Data.Ord (comparing) +import Prelude + +type SpectrumEnd :: Type +data SpectrumEnd = SmallestEigenvalues | LargestEigenvalues + deriving stock (Eq, Show) + +sortForSpectrumBy :: Ord keyValue => SpectrumEnd -> (value -> keyValue) -> [value] -> [value] +sortForSpectrumBy spectrumEnd projectValue = + sortBy + ( case spectrumEnd of + SmallestEigenvalues -> comparing projectValue + LargestEigenvalues -> flip (comparing projectValue) + ) + +sortRawPairsForSpectrum :: SpectrumEnd -> [(Double, vector)] -> [(Double, vector)] +sortRawPairsForSpectrum spectrumEnd = sortForSpectrumBy spectrumEnd fst diff --git a/src-spectral/Moonlight/LinAlg/Pure/Operator.hs b/src-spectral/Moonlight/LinAlg/Pure/Operator.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Operator.hs @@ -0,0 +1,44 @@ +{-# LANGUAGE DataKinds #-} + +module Moonlight.LinAlg.Pure.Operator + ( OperatorSymmetry (..), + LinearOperator, + ApplyU, + operatorShape, + operatorDimension, + mkVectorLinearOperator, + declaredSelfAdjointVectorLinearOperator, + runOperatorU, + csrLinearOperator, + selfAdjointCSRLinearOperator, + graphLaplacianLinearOperator, + diagonalLinearOperator, + pathLaplacianLinearOperator, + symmetricTridiagonalLinearOperator, + packedSparseLinearOperator, + scaleLinearOperator, + addScaledIdentity, + sigmaIdentityMinus, + ) +where + +import Moonlight.LinAlg.Pure.Operator.Internal + ( ApplyU, + LinearOperator, + OperatorSymmetry (..), + addScaledIdentity, + csrLinearOperator, + declaredSelfAdjointVectorLinearOperator, + diagonalLinearOperator, + graphLaplacianLinearOperator, + mkVectorLinearOperator, + operatorDimension, + operatorShape, + packedSparseLinearOperator, + pathLaplacianLinearOperator, + runOperatorU, + scaleLinearOperator, + selfAdjointCSRLinearOperator, + sigmaIdentityMinus, + symmetricTridiagonalLinearOperator, + ) diff --git a/src-spectral/Moonlight/LinAlg/Pure/Operator/Internal.hs b/src-spectral/Moonlight/LinAlg/Pure/Operator/Internal.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Operator/Internal.hs @@ -0,0 +1,376 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Operator.Internal + ( OperatorSymmetry (..), + LinearOperator (..), + OperatorSource (..), + ApplyU, + operatorShape, + operatorDimension, + mkVectorLinearOperator, + declaredSelfAdjointVectorLinearOperator, + runOperatorU, + csrLinearOperator, + selfAdjointCSRLinearOperator, + graphLaplacianLinearOperator, + diagonalLinearOperator, + pathLaplacianLinearOperator, + symmetricTridiagonalLinearOperator, + packedSparseLinearOperator, + scaleLinearOperator, + addScaledIdentity, + sigmaIdentityMinus, + applyOperatorSource, + operatorSourceShape, + ) +where + +import Data.Kind (Type) +import Data.Map.Strict qualified as Map +import Data.Vector.Unboxed qualified as U +import Moonlight.Core (MoonlightError (..), fieldValueValid) +import Moonlight.LinAlg.Internal.VectorOps (csrMatVecValidatedU) +import Moonlight.LinAlg.Pure.Sparse.Packed + ( PackedSparseApplyError (..), + PackedSparseOperator, + applyPackedSparseOperatorDense, + packedSparseOperatorSourceCardinality, + packedSparseOperatorTargetCardinality, + ) +import Moonlight.LinAlg.Pure.Sparse.Structured + ( GraphEdge, + graphLaplacianCSR, + symmetricTridiagonalFromCSR, + ) +import Moonlight.LinAlg.Pure.Sparse.Types + ( SparseCSR, + cooEntries, + csrCols, + csrColumnIndicesVector, + csrRows, + csrRowOffsetsVector, + csrToCOO, + csrValuesVector, + ) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal + ( SymmetricTridiagonal, + applyPathLaplacianValidatedU, + applySymmetricTridiagonalValidatedU, + isPathLaplacianTridiagonal, + symmetricTridiagonalDimension, + ) +import Prelude + +data OperatorSymmetry + = GeneralOperator + | SelfAdjointOperator + +type ApplyU :: Type +type ApplyU = U.Vector Double -> Either MoonlightError (U.Vector Double) + +type OperatorSource :: OperatorSymmetry -> Type +data OperatorSource symmetry where + OpaqueGeneralSource :: + !Int -> + !Int -> + !ApplyU -> + OperatorSource 'GeneralOperator + DeclaredSelfAdjointSource :: + !Int -> + !ApplyU -> + OperatorSource 'SelfAdjointOperator + CSRSource :: + !(SparseCSR Double) -> + OperatorSource 'GeneralOperator + SelfAdjointCSRSource :: + !(SparseCSR Double) -> + OperatorSource 'SelfAdjointOperator + GraphLaplacianCSRSource :: + !(SparseCSR Double) -> + OperatorSource 'SelfAdjointOperator + DiagonalSource :: + !(U.Vector Double) -> + OperatorSource 'SelfAdjointOperator + PathLaplacianSource :: + !Int -> + OperatorSource 'SelfAdjointOperator + SymmetricTridiagonalSource :: + !SymmetricTridiagonal -> + OperatorSource 'SelfAdjointOperator + PackedSparseSource :: + !(PackedSparseOperator Double) -> + OperatorSource 'GeneralOperator + +type LinearOperator :: OperatorSymmetry -> Type +data LinearOperator symmetry = LinearOperator + { operatorSourceScale :: !Double, + operatorIdentityShift :: !Double, + operatorSource :: !(OperatorSource symmetry) + } + +operatorShape :: LinearOperator symmetry -> (Int, Int) +operatorShape = operatorSourceShape . operatorSource + +operatorDimension :: LinearOperator 'SelfAdjointOperator -> Int +operatorDimension operatorValue = + case operatorShape operatorValue of + (rowCount, _) -> rowCount + +operatorSourceShape :: OperatorSource symmetry -> (Int, Int) +operatorSourceShape sourceValue = + case sourceValue of + OpaqueGeneralSource rowCount columnCount _ -> (rowCount, columnCount) + DeclaredSelfAdjointSource dimension _ -> (dimension, dimension) + CSRSource csrValue -> (csrRows csrValue, csrCols csrValue) + SelfAdjointCSRSource csrValue -> (csrRows csrValue, csrCols csrValue) + GraphLaplacianCSRSource csrValue -> (csrRows csrValue, csrCols csrValue) + DiagonalSource diagonalEntries -> (U.length diagonalEntries, U.length diagonalEntries) + PathLaplacianSource dimension -> (dimension, dimension) + SymmetricTridiagonalSource tridiagonalValue -> + let dimension = symmetricTridiagonalDimension tridiagonalValue + in (dimension, dimension) + PackedSparseSource packedOperator -> + ( packedSparseOperatorTargetCardinality packedOperator, + packedSparseOperatorSourceCardinality packedOperator + ) + +mkVectorLinearOperator :: Int -> Int -> ApplyU -> Either MoonlightError (LinearOperator 'GeneralOperator) +mkVectorLinearOperator rowCount columnCount applyVector = + validateRectangularDimensions "linear operator" rowCount columnCount + *> pure + LinearOperator + { operatorSourceScale = 1.0, + operatorIdentityShift = 0.0, + operatorSource = OpaqueGeneralSource rowCount columnCount (checkedApply rowCount columnCount applyVector) + } + +declaredSelfAdjointVectorLinearOperator :: Int -> ApplyU -> Either MoonlightError (LinearOperator 'SelfAdjointOperator) +declaredSelfAdjointVectorLinearOperator dimension applyVector = + validateDimension "declared self-adjoint linear operator" dimension + *> pure + LinearOperator + { operatorSourceScale = 1.0, + operatorIdentityShift = 0.0, + operatorSource = DeclaredSelfAdjointSource dimension (checkedApply dimension dimension applyVector) + } + +csrLinearOperator :: SparseCSR Double -> LinearOperator 'GeneralOperator +csrLinearOperator csrValue = + LinearOperator + { operatorSourceScale = 1.0, + operatorIdentityShift = 0.0, + operatorSource = CSRSource csrValue + } + +selfAdjointCSRLinearOperator :: SparseCSR Double -> Either MoonlightError (LinearOperator 'SelfAdjointOperator) +selfAdjointCSRLinearOperator csrValue = do + validateSelfAdjointCSR csrValue + sourceValue <- + classifyStructuredSelfAdjointCSR + SelfAdjointCSRSource + csrValue + pure + LinearOperator + { operatorSourceScale = 1.0, + operatorIdentityShift = 0.0, + operatorSource = sourceValue + } + +graphLaplacianLinearOperator :: + (Ord vertex, Show vertex) => + [vertex] -> + [GraphEdge vertex] -> + Either MoonlightError (LinearOperator 'SelfAdjointOperator) +graphLaplacianLinearOperator vertexOrder graphEdges = do + csrValue <- graphLaplacianCSR vertexOrder graphEdges + sourceValue <- + classifyStructuredSelfAdjointCSR + GraphLaplacianCSRSource + csrValue + pure + LinearOperator + { operatorSourceScale = 1.0, + operatorIdentityShift = 0.0, + operatorSource = sourceValue + } + +classifyStructuredSelfAdjointCSR :: + (SparseCSR Double -> OperatorSource 'SelfAdjointOperator) -> + SparseCSR Double -> + Either MoonlightError (OperatorSource 'SelfAdjointOperator) +classifyStructuredSelfAdjointCSR fallbackSource csrValue = do + classifiedStructure <- symmetricTridiagonalFromCSR csrValue + pure + ( case classifiedStructure of + Right tridiagonalValue + | isPathLaplacianTridiagonal tridiagonalValue -> + PathLaplacianSource + (symmetricTridiagonalDimension tridiagonalValue) + | otherwise -> + SymmetricTridiagonalSource tridiagonalValue + Left _ -> fallbackSource csrValue + ) + +diagonalLinearOperator :: U.Vector Double -> Either MoonlightError (LinearOperator 'SelfAdjointOperator) +diagonalLinearOperator diagonalEntries = + if U.any (not . fieldValueValid) diagonalEntries + then Left (InvariantViolation "diagonal linear operator requires finite entries") + else + pure + LinearOperator + { operatorSourceScale = 1.0, + operatorIdentityShift = 0.0, + operatorSource = DiagonalSource diagonalEntries + } + +pathLaplacianLinearOperator :: Int -> Either MoonlightError (LinearOperator 'SelfAdjointOperator) +pathLaplacianLinearOperator dimension = + validateDimension "path Laplacian linear operator" dimension + *> pure + LinearOperator + { operatorSourceScale = 1.0, + operatorIdentityShift = 0.0, + operatorSource = PathLaplacianSource dimension + } + +symmetricTridiagonalLinearOperator :: SymmetricTridiagonal -> LinearOperator 'SelfAdjointOperator +symmetricTridiagonalLinearOperator tridiagonalValue = + LinearOperator + { operatorSourceScale = 1.0, + operatorIdentityShift = 0.0, + operatorSource = SymmetricTridiagonalSource tridiagonalValue + } + +packedSparseLinearOperator :: PackedSparseOperator Double -> LinearOperator 'GeneralOperator +packedSparseLinearOperator packedOperator = + LinearOperator + { operatorSourceScale = 1.0, + operatorIdentityShift = 0.0, + operatorSource = PackedSparseSource packedOperator + } + +scaleLinearOperator :: Double -> LinearOperator symmetry -> LinearOperator symmetry +scaleLinearOperator scaleValue operatorValue = + operatorValue {operatorSourceScale = scaleValue * operatorSourceScale operatorValue, operatorIdentityShift = scaleValue * operatorIdentityShift operatorValue} + +addScaledIdentity :: Double -> LinearOperator 'SelfAdjointOperator -> LinearOperator 'SelfAdjointOperator +addScaledIdentity shiftValue operatorValue = + operatorValue {operatorIdentityShift = shiftValue + operatorIdentityShift operatorValue} + +sigmaIdentityMinus :: Double -> LinearOperator 'SelfAdjointOperator -> LinearOperator 'SelfAdjointOperator +sigmaIdentityMinus sigma operatorValue = + operatorValue + { operatorSourceScale = negate (operatorSourceScale operatorValue), + operatorIdentityShift = sigma - operatorIdentityShift operatorValue + } + +runOperatorU :: LinearOperator symmetry -> U.Vector Double -> Either MoonlightError (U.Vector Double) +runOperatorU operatorValue inputVector = do + sourceImage <- applyOperatorSource (operatorSource operatorValue) inputVector + applyAffineImage (operatorSourceScale operatorValue) (operatorIdentityShift operatorValue) sourceImage inputVector + +applyOperatorSource :: OperatorSource symmetry -> U.Vector Double -> Either MoonlightError (U.Vector Double) +applyOperatorSource sourceValue inputVector = + case sourceValue of + OpaqueGeneralSource rowCount columnCount applyVector -> checkedApply rowCount columnCount applyVector inputVector + DeclaredSelfAdjointSource dimension applyVector -> checkedApply dimension dimension applyVector inputVector + CSRSource csrValue -> applyCSR csrValue inputVector + SelfAdjointCSRSource csrValue -> applyCSR csrValue inputVector + GraphLaplacianCSRSource csrValue -> applyCSR csrValue inputVector + DiagonalSource diagonalEntries -> applyDiagonal diagonalEntries inputVector + PathLaplacianSource dimension -> + if U.length inputVector == dimension + then Right (applyPathLaplacianValidatedU dimension inputVector) + else + Left + ( InvariantViolation + ( "path Laplacian input dimension mismatch: expected " + <> show dimension + <> " but received " + <> show (U.length inputVector) + ) + ) + SymmetricTridiagonalSource tridiagonalValue -> + let dimension = symmetricTridiagonalDimension tridiagonalValue + in if U.length inputVector == dimension + then + Right + ( applySymmetricTridiagonalValidatedU + tridiagonalValue + inputVector + ) + else + Left + ( InvariantViolation + ( "symmetric tridiagonal operator input dimension mismatch: expected " + <> show dimension + <> " but received " + <> show (U.length inputVector) + ) + ) + PackedSparseSource packedOperator -> + case applyPackedSparseOperatorDense packedOperator inputVector of + Right output -> Right output + Left applyError -> Left (packedSparseApplyErrorToMoonlightError applyError) + +applyAffineImage :: Double -> Double -> U.Vector Double -> U.Vector Double -> Either MoonlightError (U.Vector Double) +applyAffineImage scaleValue shiftValue sourceImage inputVector + | shiftValue == 0.0 && scaleValue == 1.0 = Right sourceImage + | shiftValue == 0.0 = Right (U.map (scaleValue *) sourceImage) + | U.length sourceImage == U.length inputVector = + Right (U.zipWith (\imageEntry inputEntry -> scaleValue * imageEntry + shiftValue * inputEntry) sourceImage inputVector) + | otherwise = Left (InvariantViolation "identity shift requires a square operator image") + +applyCSR :: SparseCSR Double -> U.Vector Double -> Either MoonlightError (U.Vector Double) +applyCSR csrValue inputVector = + if U.length inputVector /= csrCols csrValue + then Left (InvariantViolation ("CSR matvec dimension mismatch: expected " <> show (csrCols csrValue) <> " but received " <> show (U.length inputVector))) + else Right (csrMatVecValidatedU (csrRows csrValue) (csrRowOffsetsVector csrValue) (csrColumnIndicesVector csrValue) (csrValuesVector csrValue) inputVector) + +applyDiagonal :: U.Vector Double -> U.Vector Double -> Either MoonlightError (U.Vector Double) +applyDiagonal diagonalEntries inputVector = + if U.length inputVector /= U.length diagonalEntries + then Left (InvariantViolation ("diagonal operator input dimension mismatch: expected " <> show (U.length diagonalEntries) <> " but received " <> show (U.length inputVector))) + else Right (U.zipWith (*) diagonalEntries inputVector) + +checkedApply :: Int -> Int -> ApplyU -> ApplyU +checkedApply rowCount columnCount applyVector inputVector = + if U.length inputVector /= columnCount + then Left (InvariantViolation ("linear operator input dimension mismatch: expected " <> show columnCount <> " but received " <> show (U.length inputVector))) + else do + outputVector <- applyVector inputVector + if U.length outputVector == rowCount + then Right outputVector + else Left (InvariantViolation ("linear operator output dimension mismatch: expected " <> show rowCount <> " but received " <> show (U.length outputVector))) + +validateRectangularDimensions :: String -> Int -> Int -> Either MoonlightError () +validateRectangularDimensions label rowCount columnCount = + if rowCount < 0 || columnCount < 0 + then Left (InvariantViolation (label <> " dimensions must be non-negative")) + else Right () + +validateDimension :: String -> Int -> Either MoonlightError () +validateDimension label dimension = + if dimension <= 0 + then Left (InvariantViolation (label <> " dimension must be positive, received " <> show dimension)) + else Right () + +validateSelfAdjointCSR :: SparseCSR Double -> Either MoonlightError () +validateSelfAdjointCSR csrValue = do + if csrRows csrValue /= csrCols csrValue + then Left (InvariantViolation "self-adjoint CSR operator requires a square matrix") + else do + cooValue <- csrToCOO csrValue + let entryMap = Map.fromList (((\(rowIndex, columnIndex, value) -> ((rowIndex, columnIndex), value)) <$> cooEntries cooValue)) + symmetricEntry ((rowIndex, columnIndex), value) = Map.lookup (columnIndex, rowIndex) entryMap == Just value + if all symmetricEntry (Map.toList entryMap) + then Right () + else Left (InvariantViolation "self-adjoint CSR operator requires exact symmetric storage") + +packedSparseApplyErrorToMoonlightError :: PackedSparseApplyError -> MoonlightError +packedSparseApplyErrorToMoonlightError errorValue = + case errorValue of + PackedSparseInputLengthMismatch expectedLength actualLength -> + InvariantViolation ("packed sparse operator input dimension mismatch: expected " <> show expectedLength <> " but received " <> show actualLength) diff --git a/src-spectral/Moonlight/LinAlg/Pure/Spectral/Request.hs b/src-spectral/Moonlight/LinAlg/Pure/Spectral/Request.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Spectral/Request.hs @@ -0,0 +1,25 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Spectral.Request + ( EigenRequest (..), + ) +where + +import Data.Kind (Type) +import Data.Vector.Unboxed qualified as U +import Moonlight.LinAlg.Pure.Krylov.Config (PositiveCount) +import Moonlight.LinAlg.Pure.Krylov.Selection (SpectrumEnd) +import Moonlight.LinAlg.Pure.Spectral.Result (Eigenpairs) +import Prelude + +type EigenRequest :: Type -> Type +data EigenRequest result where + EigenvaluesRequest :: + !SpectrumEnd -> + !PositiveCount -> + EigenRequest (U.Vector Double) + EigenpairsRequest :: + !SpectrumEnd -> + !PositiveCount -> + EigenRequest Eigenpairs diff --git a/src-spectral/Moonlight/LinAlg/Pure/Spectral/Result.hs b/src-spectral/Moonlight/LinAlg/Pure/Spectral/Result.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Spectral/Result.hs @@ -0,0 +1,253 @@ +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Spectral.Result + ( Eigenpairs, + CertifiedSelectedEigenpairResult (..), + SelectedEigenpairCertificationFailure (..), + SelectedEigenpairOrthonormalityEvidence (..), + SelectedEigenpairRequestOrderingEvidence (..), + SelectedEigenpairResidualEvidence (..), + mkEigenpairs, + certifySelectedEigenpairResult, + eigenpairDimension, + eigenpairValues, + eigenpairVectorsColumnMajor, + eigenpairResidualNorms, + eigenpairCount, + eigenpairVectorAt, + eigenpairsFromColumns, + mapEigenpairValues, + ) +where + +import Data.Bifunctor (first) +import Data.Kind (Type) +import Data.Vector.Unboxed qualified as U +import Moonlight.Core + ( MoonlightError (..), + checkedNonNegativeProduct, + fieldValueValid, + ) +import Moonlight.LinAlg.Pure.Krylov.Selection (SpectrumEnd (..)) +import Prelude + +type Eigenpairs :: Type +data Eigenpairs = Eigenpairs + { eigenpairDimension :: !Int, + eigenpairValues :: !(U.Vector Double), + eigenpairVectorsColumnMajor :: !(U.Vector Double), + eigenpairResidualNorms :: !(U.Vector Double) + } + deriving stock (Eq, Show) + +type SelectedEigenpairResidualEvidence :: Type +data SelectedEigenpairResidualEvidence = SelectedEigenpairResidualEvidence + { selectedEigenpairResidualBound :: !Double, + selectedEigenpairMaxResidualNorm :: !Double + } + deriving stock (Eq, Show) + +type SelectedEigenpairOrthonormalityEvidence :: Type +data SelectedEigenpairOrthonormalityEvidence = SelectedEigenpairOrthonormalityEvidence + { selectedEigenpairOrthonormalityBound :: !Double, + selectedEigenpairMaxOrthonormalityDeviation :: !Double + } + deriving stock (Eq, Show) + +type SelectedEigenpairRequestOrderingEvidence :: Type +data SelectedEigenpairRequestOrderingEvidence = SelectedEigenpairRequestOrderingEvidence + { selectedEigenpairRequestedCount :: !Int, + selectedEigenpairCertifiedCount :: !Int, + selectedEigenpairCertifiedOrdering :: !SpectrumEnd + } + deriving stock (Eq, Show) + +type CertifiedSelectedEigenpairResult :: Type +data CertifiedSelectedEigenpairResult = CertifiedSelectedEigenpairResult + { certifiedSelectedEigenpairResult :: !Eigenpairs, + certifiedSelectedEigenpairResidualEvidence :: !SelectedEigenpairResidualEvidence, + certifiedSelectedEigenpairOrthonormalityEvidence :: !SelectedEigenpairOrthonormalityEvidence, + certifiedSelectedEigenpairRequestOrderingEvidence :: !SelectedEigenpairRequestOrderingEvidence + } + deriving stock (Eq, Show) + +type SelectedEigenpairCertificationFailure :: Type +data SelectedEigenpairCertificationFailure + = SelectedEigenpairCertificationInvalidRequest !String + | SelectedEigenpairCertificationRequestedCountMismatch !Int !Int + | SelectedEigenpairCertificationResidualExceeded !Int !Double !Double + | SelectedEigenpairCertificationOrthonormalityExceeded !Int !Int !Double !Double + | SelectedEigenpairCertificationOrderingViolation !SpectrumEnd !Int !Double !Double + | SelectedEigenpairCertificationShapeMismatch !String + deriving stock (Eq, Show) + +mkEigenpairs :: Int -> U.Vector Double -> U.Vector Double -> U.Vector Double -> Either MoonlightError Eigenpairs +mkEigenpairs dimension values vectors residuals + | dimension <= 0 = Left (InvariantViolation "Eigenpairs require a positive ambient dimension") + | U.length residuals /= U.length values = + Left (InvariantViolation "Eigenpair residual count must match eigenvalue count") + | otherwise = do + expectedVectorCount <- + first + (const (InvariantViolation "Eigenpair vector payload cardinality exceeds Int range")) + (checkedNonNegativeProduct dimension (U.length values)) + if U.length vectors /= expectedVectorCount + then Left (InvariantViolation "Eigenpair vector payload length must equal dimension * eigenvalue count") + else + Right + Eigenpairs + { eigenpairDimension = dimension, + eigenpairValues = values, + eigenpairVectorsColumnMajor = vectors, + eigenpairResidualNorms = residuals + } + +certifySelectedEigenpairResult :: + SpectrumEnd -> + Int -> + Double -> + Double -> + Eigenpairs -> + Either SelectedEigenpairCertificationFailure CertifiedSelectedEigenpairResult +certifySelectedEigenpairResult spectrumEnd requestedCount residualBound orthonormalityBound pairs + | requestedCount <= 0 = + Left (SelectedEigenpairCertificationInvalidRequest ("selected eigenpair certification requires a positive requested count, received " <> show requestedCount)) + | not (finiteNonNegative residualBound) = + Left (SelectedEigenpairCertificationInvalidRequest ("selected eigenpair residual bound must be finite and non-negative, received " <> show residualBound)) + | not (finiteNonNegative orthonormalityBound) = + Left (SelectedEigenpairCertificationInvalidRequest ("selected eigenpair orthonormality bound must be finite and non-negative, received " <> show orthonormalityBound)) + | eigenpairCount pairs /= requestedCount = + Left (SelectedEigenpairCertificationRequestedCountMismatch requestedCount (eigenpairCount pairs)) + | otherwise = do + residualEvidence <- certifySelectedEigenpairResiduals residualBound pairs + columns <- + case traverse (`eigenpairVectorAt` pairs) [0 .. eigenpairCount pairs - 1] of + Left err -> Left (SelectedEigenpairCertificationShapeMismatch (show err)) + Right columnValues -> Right columnValues + orthonormalityEvidence <- certifySelectedEigenpairOrthonormality orthonormalityBound columns + requestOrderingEvidence <- certifySelectedEigenpairOrdering spectrumEnd requestedCount pairs + Right + CertifiedSelectedEigenpairResult + { certifiedSelectedEigenpairResult = pairs, + certifiedSelectedEigenpairResidualEvidence = residualEvidence, + certifiedSelectedEigenpairOrthonormalityEvidence = orthonormalityEvidence, + certifiedSelectedEigenpairRequestOrderingEvidence = requestOrderingEvidence + } + +eigenpairCount :: Eigenpairs -> Int +eigenpairCount = U.length . eigenpairValues + +eigenpairVectorAt :: Int -> Eigenpairs -> Either MoonlightError (U.Vector Double) +eigenpairVectorAt columnIndex pairs + | columnIndex < 0 || columnIndex >= eigenpairCount pairs = + Left (InvariantViolation "eigenpair vector index out of bounds") + | otherwise = + Right + ( U.slice + (columnIndex * eigenpairDimension pairs) + (eigenpairDimension pairs) + (eigenpairVectorsColumnMajor pairs) + ) + +eigenpairsFromColumns :: Int -> [(Double, U.Vector Double, Double)] -> Either MoonlightError Eigenpairs +eigenpairsFromColumns dimension columns = + let values = U.fromList ((\(value, _, _) -> value) <$> columns) + vectors = U.concat ((\(_, vector, _) -> vector) <$> columns) + residuals = U.fromList ((\(_, _, residual) -> residual) <$> columns) + in mkEigenpairs dimension values vectors residuals + +mapEigenpairValues :: (Double -> Double) -> (Double -> Double) -> Eigenpairs -> Either MoonlightError Eigenpairs +mapEigenpairValues mapValue mapResidual pairs = + mkEigenpairs + (eigenpairDimension pairs) + (U.map mapValue (eigenpairValues pairs)) + (eigenpairVectorsColumnMajor pairs) + (U.map mapResidual (eigenpairResidualNorms pairs)) + +certifySelectedEigenpairResiduals :: + Double -> + Eigenpairs -> + Either SelectedEigenpairCertificationFailure SelectedEigenpairResidualEvidence +certifySelectedEigenpairResiduals residualBound pairs = + case filter (\(_, residualValue) -> not (fieldValueValid residualValue) || residualValue < 0.0 || residualValue > residualBound) indexedResiduals of + [] -> + Right + SelectedEigenpairResidualEvidence + { selectedEigenpairResidualBound = residualBound, + selectedEigenpairMaxResidualNorm = foldr max 0.0 (abs . snd <$> indexedResiduals) + } + (columnIndex, residualValue) : _ -> + Left (SelectedEigenpairCertificationResidualExceeded columnIndex residualValue residualBound) + where + indexedResiduals = zip [0 :: Int ..] (U.toList (eigenpairResidualNorms pairs)) + +certifySelectedEigenpairOrthonormality :: + Double -> + [U.Vector Double] -> + Either SelectedEigenpairCertificationFailure SelectedEigenpairOrthonormalityEvidence +certifySelectedEigenpairOrthonormality orthonormalityBound columns = + case filter (\(_, _, deviationValue) -> not (fieldValueValid deviationValue) || abs deviationValue > orthonormalityBound) deviations of + [] -> + Right + SelectedEigenpairOrthonormalityEvidence + { selectedEigenpairOrthonormalityBound = orthonormalityBound, + selectedEigenpairMaxOrthonormalityDeviation = foldr max 0.0 (abs . thirdEntry <$> deviations) + } + (leftIndex, rightIndex, deviationValue) : _ -> + Left (SelectedEigenpairCertificationOrthonormalityExceeded leftIndex rightIndex deviationValue orthonormalityBound) + where + deviations = + [ (leftIndex, rightIndex, vectorDotU leftColumn rightColumn - expectedInnerProduct leftIndex rightIndex) + | (leftIndex, leftColumn) <- zip [0 :: Int ..] columns, + (rightIndex, rightColumn) <- zip [0 :: Int ..] columns, + leftIndex <= rightIndex + ] + +certifySelectedEigenpairOrdering :: + SpectrumEnd -> + Int -> + Eigenpairs -> + Either SelectedEigenpairCertificationFailure SelectedEigenpairRequestOrderingEvidence +certifySelectedEigenpairOrdering spectrumEnd requestedCount pairs = + case filter (not . orderedAdjacent spectrumEnd) adjacentValues of + [] -> + Right + SelectedEigenpairRequestOrderingEvidence + { selectedEigenpairRequestedCount = requestedCount, + selectedEigenpairCertifiedCount = eigenpairCount pairs, + selectedEigenpairCertifiedOrdering = spectrumEnd + } + (leftIndex, leftValue, rightValue) : _ -> + Left (SelectedEigenpairCertificationOrderingViolation spectrumEnd leftIndex leftValue rightValue) + where + values = U.toList (eigenpairValues pairs) + adjacentValues = zipWith (\indexValue (leftValue, rightValue) -> (indexValue, leftValue, rightValue)) [0 :: Int ..] (zip values (drop 1 values)) + +orderedAdjacent :: SpectrumEnd -> (Int, Double, Double) -> Bool +orderedAdjacent spectrumEnd (_, leftValue, rightValue) = + fieldValueValid leftValue + && fieldValueValid rightValue + && spectrumValuesOrdered spectrumEnd leftValue rightValue + +spectrumValuesOrdered :: SpectrumEnd -> Double -> Double -> Bool +spectrumValuesOrdered spectrumEnd leftValue rightValue = + case spectrumEnd of + SmallestEigenvalues -> leftValue <= rightValue + LargestEigenvalues -> leftValue >= rightValue + +expectedInnerProduct :: Int -> Int -> Double +expectedInnerProduct leftIndex rightIndex = + if leftIndex == rightIndex + then 1.0 + else 0.0 + +vectorDotU :: U.Vector Double -> U.Vector Double -> Double +vectorDotU leftVector rightVector = + U.sum (U.zipWith (*) leftVector rightVector) + +thirdEntry :: (left, right, value) -> value +thirdEntry (_, _, value) = value + +finiteNonNegative :: Double -> Bool +finiteNonNegative value = + fieldValueValid value && value >= 0.0 diff --git a/src-spectral/Moonlight/LinAlg/Pure/Spectral/Solve.hs b/src-spectral/Moonlight/LinAlg/Pure/Spectral/Solve.hs new file mode 100644 --- /dev/null +++ b/src-spectral/Moonlight/LinAlg/Pure/Spectral/Solve.hs @@ -0,0 +1,650 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Spectral.Solve + ( EigenSolveConfig (..), + defaultEigenSolveConfig, + withEigenFallbackLanczosConfig, + withEigenFallbackInitialVector, + denseSpectralFallbackDimensionThreshold, + solveEigenRequest, + ) +where + +import Data.Bifunctor (first) +import Data.List (sortBy) +import Data.Ord (comparing) +import Data.Vector.Storable qualified as S +import Data.Vector.Unboxed qualified as U +import Moonlight.Core + ( MoonlightError (..), + checkedNonNegativeProduct, + ) +import Moonlight.LinAlg.Internal.Eigen.Symmetric + ( SymmetricEigenResult (..), + symmetricEigenPairsDenseUnchecked, + ) +import Moonlight.LinAlg.Pure.Dense.Flat + ( DenseDoubleMatrix, + denseDoubleMatrixToRowMajorVector, + denseDoubleMatrixVectorProduct, + mkDenseDoubleMatrixRowMajor, + ) +import Moonlight.LinAlg.Pure.Krylov.Config (LanczosConfig, defaultLanczosConfig, positiveCountValue) +import Moonlight.LinAlg.Pure.Krylov.CascadicGraph + ( CascadicGraphObstruction (..), + cascadicGraphLaplacianEigenpairs, + ) +import Moonlight.LinAlg.Pure.Krylov.Projected + ( projectedEigenpairsFromRestartedLanczos, + projectedEigenvaluesFromRestartedLanczos, + ) +import Moonlight.LinAlg.Pure.Krylov.SelectedTridiagonal + ( symmetricTridiagonalFromCSR, + selectedSymmetricTridiagonalEigenpairsDirect, + selectedSymmetricTridiagonalEigenvaluesDirect, + ) +import Moonlight.LinAlg.Pure.Krylov.Selection (SpectrumEnd (..)) +import Moonlight.LinAlg.Pure.Sparse.Types (SparseCSR) +import Moonlight.LinAlg.Pure.Operator.Internal + ( LinearOperator (..), + OperatorSource (..), + OperatorSymmetry (SelfAdjointOperator), + operatorDimension, + runOperatorU, + ) +import Moonlight.LinAlg.Pure.Spectral.Request (EigenRequest (..)) +import Moonlight.LinAlg.Pure.Spectral.Result + ( Eigenpairs, + eigenpairValues, + eigenpairsFromColumns, + mapEigenpairValues, + ) +import Prelude + +data DiagonalOrder + = DiagonalAscending + | DiagonalDescending + +data DiagonalOrderScan = DiagonalOrderScan + { diagonalScanPrevious :: !Double, + diagonalScanAscending :: !Bool, + diagonalScanDescending :: !Bool + } + +data EigenSolveConfig = EigenSolveConfig + { eigenFallbackLanczosConfig :: !LanczosConfig, + eigenFallbackInitialVector :: !(Maybe (U.Vector Double)) + } + deriving stock (Eq, Show) + +defaultEigenSolveConfig :: EigenSolveConfig +defaultEigenSolveConfig = + EigenSolveConfig + { eigenFallbackLanczosConfig = defaultLanczosConfig, + eigenFallbackInitialVector = Nothing + } + +withEigenFallbackLanczosConfig :: LanczosConfig -> EigenSolveConfig -> EigenSolveConfig +withEigenFallbackLanczosConfig lanczosConfig config = + config {eigenFallbackLanczosConfig = lanczosConfig} + +withEigenFallbackInitialVector :: U.Vector Double -> EigenSolveConfig -> EigenSolveConfig +withEigenFallbackInitialVector seedVector config = + config {eigenFallbackInitialVector = Just seedVector} + +solveEigenRequest :: + EigenSolveConfig -> + LinearOperator 'SelfAdjointOperator -> + EigenRequest result -> + Either MoonlightError result +solveEigenRequest config operatorValue requestValue = do + let dimension = operatorDimension operatorValue + requestedCount = eigenRequestCount requestValue + scaleValue = operatorSourceScale operatorValue + shiftValue = operatorIdentityShift operatorValue + validateSpectralCount requestedCount dimension + if scaleValue == 0.0 + then solveZeroScale shiftValue dimension requestValue + else solveAffineRequest config operatorValue scaleValue shiftValue requestValue + +solveAffineRequest :: + EigenSolveConfig -> + LinearOperator 'SelfAdjointOperator -> + Double -> + Double -> + EigenRequest result -> + Either MoonlightError result +solveAffineRequest config operatorValue scaleValue shiftValue requestValue = + case requestValue of + EigenvaluesRequest spectrumEnd count -> + transformValues scaleValue shiftValue + <$> solveSourceEigenvalues config operatorValue (transportSpectrumEnd scaleValue spectrumEnd) (positiveCountValue count) + EigenpairsRequest spectrumEnd count -> + transformPairs scaleValue shiftValue + =<< solveSourceEigenpairs config operatorValue (transportSpectrumEnd scaleValue spectrumEnd) (positiveCountValue count) + +solveSourceEigenvalues :: + EigenSolveConfig -> + LinearOperator 'SelfAdjointOperator -> + SpectrumEnd -> + Int -> + Either MoonlightError (U.Vector Double) +solveSourceEigenvalues config operatorValue spectrumEnd requestedCount = + case operatorSource operatorValue of + DiagonalSource diagonalEntries -> diagonalValues spectrumEnd requestedCount diagonalEntries + PathLaplacianSource dimension -> pathLaplacianValues spectrumEnd requestedCount dimension + SymmetricTridiagonalSource tridiagonalValue -> + selectedSymmetricTridiagonalEigenvaluesDirect spectrumEnd requestedCount tridiagonalValue + SelfAdjointCSRSource csrValue -> + symmetricTridiagonalFromCSR csrValue >>= \case + Right tridiagonalValue -> selectedSymmetricTridiagonalEigenvaluesDirect spectrumEnd requestedCount tridiagonalValue + Left _ -> genericFallbackValues config (sourceOperator operatorValue) spectrumEnd requestedCount + GraphLaplacianCSRSource csrValue -> + graphLaplacianFallbackValues + config + (sourceOperator operatorValue) + csrValue + spectrumEnd + requestedCount + DeclaredSelfAdjointSource _ _ -> genericFallbackValues config (sourceOperator operatorValue) spectrumEnd requestedCount + +solveSourceEigenpairs :: + EigenSolveConfig -> + LinearOperator 'SelfAdjointOperator -> + SpectrumEnd -> + Int -> + Either MoonlightError Eigenpairs +solveSourceEigenpairs config operatorValue spectrumEnd requestedCount = + case operatorSource operatorValue of + DiagonalSource diagonalEntries -> diagonalPairs spectrumEnd requestedCount diagonalEntries + PathLaplacianSource dimension -> pathLaplacianPairs spectrumEnd requestedCount dimension + SymmetricTridiagonalSource tridiagonalValue -> + selectedSymmetricTridiagonalEigenpairsDirect spectrumEnd requestedCount tridiagonalValue + SelfAdjointCSRSource csrValue -> + symmetricTridiagonalFromCSR csrValue >>= \case + Right tridiagonalValue -> + selectedSymmetricTridiagonalEigenpairsDirect spectrumEnd requestedCount tridiagonalValue + Left _ -> genericFallbackPairs config (sourceOperator operatorValue) spectrumEnd requestedCount + GraphLaplacianCSRSource csrValue -> + graphLaplacianFallbackPairs + config + (sourceOperator operatorValue) + csrValue + spectrumEnd + requestedCount + DeclaredSelfAdjointSource _ _ -> genericFallbackPairs config (sourceOperator operatorValue) spectrumEnd requestedCount + +graphLaplacianFallbackValues :: + EigenSolveConfig -> + LinearOperator 'SelfAdjointOperator -> + SparseCSR Double -> + SpectrumEnd -> + Int -> + Either MoonlightError (U.Vector Double) +graphLaplacianFallbackValues config operatorValue csrValue spectrumEnd requestedCount = + eigenpairValues + <$> graphLaplacianFallbackPairs + config + operatorValue + csrValue + spectrumEnd + requestedCount + +graphLaplacianFallbackPairs :: + EigenSolveConfig -> + LinearOperator 'SelfAdjointOperator -> + SparseCSR Double -> + SpectrumEnd -> + Int -> + Either MoonlightError Eigenpairs +graphLaplacianFallbackPairs config operatorValue csrValue spectrumEnd requestedCount + | shouldUseCascadicGraphFallback spectrumEnd requestedCount operatorValue = + case + cascadicGraphLaplacianEigenpairs + (eigenFallbackLanczosConfig config) + requestedCount + csrValue + of + Right pairs -> Right pairs + Left (CascadicGraphCoarseningStalled _) -> + genericFallbackPairs config operatorValue spectrumEnd requestedCount + Left obstruction -> Left (cascadicGraphObstructionError obstruction) + | otherwise = genericFallbackPairs config operatorValue spectrumEnd requestedCount + +shouldUseCascadicGraphFallback :: + SpectrumEnd -> + Int -> + LinearOperator 'SelfAdjointOperator -> + Bool +shouldUseCascadicGraphFallback spectrumEnd requestedCount operatorValue = + spectrumEnd == SmallestEigenvalues + && operatorDimension operatorValue >= cascadicGraphFallbackDimensionThreshold + && requestedCount <= cascadicGraphFallbackModeCountThreshold + +cascadicGraphFallbackDimensionThreshold :: Int +cascadicGraphFallbackDimensionThreshold = 4096 + +cascadicGraphFallbackModeCountThreshold :: Int +cascadicGraphFallbackModeCountThreshold = 8 + +cascadicGraphObstructionError :: CascadicGraphObstruction -> MoonlightError +cascadicGraphObstructionError obstruction = + case obstruction of + CascadicGraphBackendFailure errorValue -> errorValue + CascadicGraphCoarseningStalled dimension -> + InvariantViolation + ( "cascadic graph coarsening stalled at dimension " + <> show dimension + ) + CascadicGraphIncompleteAssignment vertexIndex -> + InvariantViolation + ( "cascadic graph coarsening omitted vertex " + <> show vertexIndex + ) + CascadicGraphInvalidRequest requestedCount dimension -> + InvariantViolation + ( "cascadic graph eigensolve requested " + <> show requestedCount + <> " modes from dimension " + <> show dimension + ) + CascadicGraphRankLoss requiredCount actualCount -> + InvariantViolation + ( "cascadic graph refinement lost block rank: required " + <> show requiredCount + <> " but retained " + <> show actualCount + ) + CascadicGraphRefinementBudgetExceeded dimension residualTarget actualResidual -> + InvariantViolation + ( "cascadic graph refinement exhausted its budget at dimension " + <> show dimension + <> ": residual target " + <> show residualTarget + <> ", actual " + <> show actualResidual + ) + +data GenericSpectralFallback + = DenseSpectralFallback + | RestartedLanczosSpectralFallback + +-- | Densify generic self-adjoint fallback through n=512; measured banded SPD +-- benches favor dense below this cutoff. +denseSpectralFallbackDimensionThreshold :: Int +denseSpectralFallbackDimensionThreshold = 512 + +-- | High-demand requests remain cheaper as one bounded dense solve than as +-- hundreds of selected-mode restart cycles. +denseSpectralHighDemandDimensionThreshold :: Int +denseSpectralHighDemandDimensionThreshold = 1024 + +genericFallbackDispatch :: Int -> LinearOperator 'SelfAdjointOperator -> GenericSpectralFallback +genericFallbackDispatch requestedCount operatorValue + | dimension <= denseSpectralFallbackDimensionThreshold = DenseSpectralFallback + | dimension <= denseSpectralHighDemandDimensionThreshold + && requestedCount >= denseRequestCardinalityFloor dimension = DenseSpectralFallback + | otherwise = RestartedLanczosSpectralFallback + where + dimension = operatorDimension operatorValue + +denseRequestCardinalityFloor :: Int -> Int +denseRequestCardinalityFloor dimension = + max 1 ((dimension + 3) `quot` 4) + +genericFallbackValues :: + EigenSolveConfig -> + LinearOperator 'SelfAdjointOperator -> + SpectrumEnd -> + Int -> + Either MoonlightError (U.Vector Double) +genericFallbackValues config operatorValue spectrumEnd requestedCount = + case genericFallbackDispatch requestedCount operatorValue of + DenseSpectralFallback -> denseFallbackValues operatorValue spectrumEnd requestedCount + RestartedLanczosSpectralFallback -> lanczosFallbackValues config operatorValue spectrumEnd requestedCount + +genericFallbackPairs :: + EigenSolveConfig -> + LinearOperator 'SelfAdjointOperator -> + SpectrumEnd -> + Int -> + Either MoonlightError Eigenpairs +genericFallbackPairs config operatorValue spectrumEnd requestedCount = + case genericFallbackDispatch requestedCount operatorValue of + DenseSpectralFallback -> denseFallbackPairs operatorValue spectrumEnd requestedCount + RestartedLanczosSpectralFallback -> lanczosFallbackPairs config operatorValue spectrumEnd requestedCount + +denseFallbackValues :: + LinearOperator 'SelfAdjointOperator -> + SpectrumEnd -> + Int -> + Either MoonlightError (U.Vector Double) +denseFallbackValues operatorValue spectrumEnd requestedCount = do + (_, eigenResult) <- denseFallbackEigenResult operatorValue + let ascendingValues = symmetricEigenResultValues eigenResult + pure + ( U.fromList + ( (ascendingValues S.!) + <$> selectedSpectrumIndices spectrumEnd requestedCount (operatorDimension operatorValue) + ) + ) + +denseFallbackPairs :: + LinearOperator 'SelfAdjointOperator -> + SpectrumEnd -> + Int -> + Either MoonlightError Eigenpairs +denseFallbackPairs operatorValue spectrumEnd requestedCount = do + (denseMatrix, eigenResult) <- denseFallbackEigenResult operatorValue + let dimension = operatorDimension operatorValue + columns <- + traverse + (denseFallbackPairColumn denseMatrix eigenResult) + (selectedSpectrumIndices spectrumEnd requestedCount dimension) + eigenpairsFromColumns dimension columns + +denseFallbackEigenResult :: + LinearOperator 'SelfAdjointOperator -> + Either MoonlightError (DenseDoubleMatrix, SymmetricEigenResult) +denseFallbackEigenResult operatorValue = do + let dimension = operatorDimension operatorValue + entryCount <- + first + (const (InvariantViolation "dense spectral fallback cardinality exceeds Int range")) + (checkedNonNegativeProduct dimension dimension) + imageColumns <- traverse (runOperatorU operatorValue . unitVector dimension) [0 .. dimension - 1] + let columnPayload = U.concat imageColumns + rowMajorPayload = + S.generate + entryCount + ( \flatIndex -> + let (rowIndex, columnIndex) = flatIndex `quotRem` dimension + in columnPayload U.! (columnIndex * dimension + rowIndex) + ) + denseMatrix <- mkDenseDoubleMatrixRowMajor dimension dimension rowMajorPayload + eigenResult <- symmetricEigenPairsDenseUnchecked dimension denseMatrix + pure (denseMatrix, eigenResult) + +selectedSpectrumIndices :: SpectrumEnd -> Int -> Int -> [Int] +selectedSpectrumIndices spectrumEnd requestedCount dimension = + case spectrumEnd of + SmallestEigenvalues -> [0 .. requestedCount - 1] + LargestEigenvalues -> [dimension - 1, dimension - 2 .. dimension - requestedCount] + +denseFallbackPairColumn :: + DenseDoubleMatrix -> + SymmetricEigenResult -> + Int -> + Either MoonlightError (Double, U.Vector Double, Double) +denseFallbackPairColumn denseMatrix eigenResult columnIndex = do + let eigenvalue = symmetricEigenResultValues eigenResult S.! columnIndex + vectorPayload = denseDoubleMatrixToRowMajorVector (symmetricEigenResultVectors eigenResult) + dimension = S.length (symmetricEigenResultValues eigenResult) + eigenvector = U.generate dimension (\rowIndex -> vectorPayload S.! (rowIndex * dimension + columnIndex)) + imageVector <- denseDoubleMatrixVectorProduct denseMatrix (S.convert eigenvector) + pure (eigenvalue, eigenvector, residualNorm eigenvalue eigenvector (S.convert imageVector)) + +residualNorm :: Double -> U.Vector Double -> U.Vector Double -> Double +residualNorm eigenvalue eigenvector imageVector = + sqrt + ( U.sum + ( U.map + (\entryValue -> entryValue * entryValue) + (U.zipWith (\imageEntry vectorEntry -> imageEntry - eigenvalue * vectorEntry) imageVector eigenvector) + ) + ) + +lanczosFallbackValues :: + EigenSolveConfig -> + LinearOperator 'SelfAdjointOperator -> + SpectrumEnd -> + Int -> + Either MoonlightError (U.Vector Double) +lanczosFallbackValues config operatorValue spectrumEnd requestedCount = + projectedEigenvaluesFromRestartedLanczos + (eigenFallbackLanczosConfig config) + spectrumEnd + requestedCount + operatorValue + (fallbackSeed config (operatorDimension operatorValue)) + +lanczosFallbackPairs :: + EigenSolveConfig -> + LinearOperator 'SelfAdjointOperator -> + SpectrumEnd -> + Int -> + Either MoonlightError Eigenpairs +lanczosFallbackPairs config operatorValue spectrumEnd requestedCount = + projectedEigenpairsFromRestartedLanczos + (eigenFallbackLanczosConfig config) + spectrumEnd + requestedCount + operatorValue + (fallbackSeed config (operatorDimension operatorValue)) + +fallbackSeed :: EigenSolveConfig -> Int -> U.Vector Double +fallbackSeed config dimension = + case eigenFallbackInitialVector config of + Just seedVector -> seedVector + Nothing -> U.generate dimension (\indexValue -> if indexValue == 0 then 1.0 else 0.0) + +sourceOperator :: LinearOperator 'SelfAdjointOperator -> LinearOperator 'SelfAdjointOperator +sourceOperator operatorValue = + operatorValue {operatorSourceScale = 1.0, operatorIdentityShift = 0.0} + +diagonalValues :: SpectrumEnd -> Int -> U.Vector Double -> Either MoonlightError (U.Vector Double) +diagonalValues spectrumEnd requestedCount diagonalEntries = + Right (diagonalSelectedValues spectrumEnd requestedCount diagonalEntries) + +diagonalPairs :: SpectrumEnd -> Int -> U.Vector Double -> Either MoonlightError Eigenpairs +diagonalPairs spectrumEnd requestedCount diagonalEntries = + eigenpairsFromColumns (U.length diagonalEntries) + ( fmap + (\(entryIndex, eigenvalue) -> (eigenvalue, unitVector (U.length diagonalEntries) entryIndex, 0.0)) + (diagonalSelectedEntries spectrumEnd requestedCount diagonalEntries) + ) + +diagonalSelectedValues :: SpectrumEnd -> Int -> U.Vector Double -> U.Vector Double +diagonalSelectedValues spectrumEnd requestedCount diagonalEntries = + case diagonalOrder diagonalEntries of + Just DiagonalAscending -> orderedAscendingValues spectrumEnd requestedCount diagonalEntries + Just DiagonalDescending -> orderedDescendingValues spectrumEnd requestedCount diagonalEntries + Nothing -> + U.fromList . fmap snd $ + diagonalSelectedEntriesBySort spectrumEnd requestedCount diagonalEntries + +diagonalSelectedEntries :: SpectrumEnd -> Int -> U.Vector Double -> [(Int, Double)] +diagonalSelectedEntries spectrumEnd requestedCount diagonalEntries = + case diagonalOrder diagonalEntries of + Just DiagonalAscending -> orderedAscendingEntries spectrumEnd requestedCount diagonalEntries + Just DiagonalDescending -> orderedDescendingEntries spectrumEnd requestedCount diagonalEntries + Nothing -> diagonalSelectedEntriesBySort spectrumEnd requestedCount diagonalEntries + +diagonalOrder :: U.Vector Double -> Maybe DiagonalOrder +diagonalOrder diagonalEntries + | U.length diagonalEntries <= 1 = Just DiagonalAscending + | otherwise = + orderFromScan + ( U.foldl' + scanDiagonalOrder + (DiagonalOrderScan (diagonalEntries `U.unsafeIndex` 0) True True) + (U.drop 1 diagonalEntries) + ) + +scanDiagonalOrder :: DiagonalOrderScan -> Double -> DiagonalOrderScan +scanDiagonalOrder scanValue entryValue = + DiagonalOrderScan + { diagonalScanPrevious = entryValue, + diagonalScanAscending = diagonalScanAscending scanValue && diagonalScanPrevious scanValue <= entryValue, + diagonalScanDescending = diagonalScanDescending scanValue && diagonalScanPrevious scanValue >= entryValue + } + +orderFromScan :: DiagonalOrderScan -> Maybe DiagonalOrder +orderFromScan scanValue + | diagonalScanAscending scanValue = Just DiagonalAscending + | diagonalScanDescending scanValue = Just DiagonalDescending + | otherwise = Nothing + +orderedAscendingValues :: SpectrumEnd -> Int -> U.Vector Double -> U.Vector Double +orderedAscendingValues spectrumEnd requestedCount diagonalEntries = + case spectrumEnd of + SmallestEigenvalues -> U.take requestedCount diagonalEntries + LargestEigenvalues -> U.reverse (U.drop (U.length diagonalEntries - requestedCount) diagonalEntries) + +orderedDescendingValues :: SpectrumEnd -> Int -> U.Vector Double -> U.Vector Double +orderedDescendingValues spectrumEnd requestedCount diagonalEntries = + case spectrumEnd of + SmallestEigenvalues -> U.reverse (U.drop (U.length diagonalEntries - requestedCount) diagonalEntries) + LargestEigenvalues -> U.take requestedCount diagonalEntries + +orderedAscendingEntries :: SpectrumEnd -> Int -> U.Vector Double -> [(Int, Double)] +orderedAscendingEntries spectrumEnd requestedCount diagonalEntries = + diagonalEntriesAt + diagonalEntries + ( case spectrumEnd of + SmallestEigenvalues -> [0 .. requestedCount - 1] + LargestEigenvalues -> [U.length diagonalEntries - 1, U.length diagonalEntries - 2 .. U.length diagonalEntries - requestedCount] + ) + +orderedDescendingEntries :: SpectrumEnd -> Int -> U.Vector Double -> [(Int, Double)] +orderedDescendingEntries spectrumEnd requestedCount diagonalEntries = + diagonalEntriesAt + diagonalEntries + ( case spectrumEnd of + SmallestEigenvalues -> [U.length diagonalEntries - 1, U.length diagonalEntries - 2 .. U.length diagonalEntries - requestedCount] + LargestEigenvalues -> [0 .. requestedCount - 1] + ) + +diagonalEntriesAt :: U.Vector Double -> [Int] -> [(Int, Double)] +diagonalEntriesAt diagonalEntries = + fmap (\entryIndex -> (entryIndex, diagonalEntries `U.unsafeIndex` entryIndex)) + +diagonalSelectedEntriesBySort :: SpectrumEnd -> Int -> U.Vector Double -> [(Int, Double)] +diagonalSelectedEntriesBySort spectrumEnd requestedCount diagonalEntries = + take requestedCount (sortIndexedValues spectrumEnd (U.toList (U.indexed diagonalEntries))) + +pathLaplacianValues :: SpectrumEnd -> Int -> Int -> Either MoonlightError (U.Vector Double) +pathLaplacianValues spectrumEnd requestedCount dimension = + Right + ( U.generate + requestedCount + ( \entryIndex -> + pathLaplacianEigenvalueAt dimension $ + case spectrumEnd of + SmallestEigenvalues -> entryIndex + LargestEigenvalues -> dimension - entryIndex - 1 + ) + ) + +pathLaplacianPairs :: SpectrumEnd -> Int -> Int -> Either MoonlightError Eigenpairs +pathLaplacianPairs spectrumEnd requestedCount dimension = + eigenpairsFromColumns dimension $ + pathLaplacianColumn dimension <$> selectedModeIndices spectrumEnd requestedCount dimension + +pathLaplacianColumn :: Int -> Int -> (Double, U.Vector Double, Double) +pathLaplacianColumn dimension modeIndex = + let eigenvalue = pathLaplacianEigenvalueAt dimension modeIndex + theta = pi * fromIntegral modeIndex / fromIntegral (max 1 dimension) + eigenvector = + if modeIndex == 0 + then U.replicate dimension (1.0 / sqrt (fromIntegral (max 1 dimension))) + else + let normalizer = sqrt (2.0 / fromIntegral dimension) + in U.generate dimension (\rowIndex -> normalizer * cos (theta * (fromIntegral rowIndex + 0.5))) + in (eigenvalue, eigenvector, pathLaplacianResidualNorm dimension eigenvalue eigenvector) + +pathLaplacianResidualNorm :: Int -> Double -> U.Vector Double -> Double +pathLaplacianResidualNorm dimension eigenvalue eigenvector = + sqrt + ( U.ifoldl' + ( \squaredNorm rowIndex _ -> + let residualEntry = pathLaplacianResidualEntry dimension eigenvalue eigenvector rowIndex + in squaredNorm + residualEntry * residualEntry + ) + 0.0 + eigenvector + ) + +pathLaplacianResidualEntry :: Int -> Double -> U.Vector Double -> Int -> Double +pathLaplacianResidualEntry dimension eigenvalue eigenvector rowIndex = + let centerValue = eigenvector `U.unsafeIndex` rowIndex + degree + | dimension == 1 = 0.0 + | rowIndex == 0 || rowIndex + 1 == dimension = 1.0 + | otherwise = 2.0 + leftValue = + if rowIndex <= 0 + then 0.0 + else eigenvector `U.unsafeIndex` (rowIndex - 1) + rightValue = + if rowIndex + 1 >= dimension + then 0.0 + else eigenvector `U.unsafeIndex` (rowIndex + 1) + imageValue = degree * centerValue - leftValue - rightValue + in imageValue - eigenvalue * centerValue +{-# INLINE pathLaplacianResidualEntry #-} + +selectedModeIndices :: SpectrumEnd -> Int -> Int -> [Int] +selectedModeIndices spectrumEnd requestedCount dimension = + case spectrumEnd of + SmallestEigenvalues -> [0 .. requestedCount - 1] + LargestEigenvalues -> [dimension - 1, dimension - 2 .. dimension - requestedCount] + +pathLaplacianEigenvalueAt :: Int -> Int -> Double +pathLaplacianEigenvalueAt matrixSize modeIndex = + 2.0 - 2.0 * cos (pi * fromIntegral modeIndex / fromIntegral (max 1 matrixSize)) + +sortIndexedValues :: SpectrumEnd -> [(Int, Double)] -> [(Int, Double)] +sortIndexedValues spectrumEnd = + sortBy + ( case spectrumEnd of + SmallestEigenvalues -> comparing snd + LargestEigenvalues -> flip (comparing snd) + ) + +transformValues :: Double -> Double -> U.Vector Double -> U.Vector Double +transformValues scaleValue shiftValue = + U.map (\eigenvalue -> scaleValue * eigenvalue + shiftValue) + +transformPairs :: Double -> Double -> Eigenpairs -> Either MoonlightError Eigenpairs +transformPairs scaleValue shiftValue = + mapEigenpairValues + (\eigenvalue -> scaleValue * eigenvalue + shiftValue) + (abs scaleValue *) + +transportSpectrumEnd :: Double -> SpectrumEnd -> SpectrumEnd +transportSpectrumEnd scaleValue spectrumEnd + | scaleValue < 0.0 = + case spectrumEnd of + SmallestEigenvalues -> LargestEigenvalues + LargestEigenvalues -> SmallestEigenvalues + | otherwise = spectrumEnd + +solveZeroScale :: Double -> Int -> EigenRequest result -> Either MoonlightError result +solveZeroScale eigenvalue dimension requestValue = + case requestValue of + EigenvaluesRequest _ count -> Right (U.replicate (positiveCountValue count) eigenvalue) + EigenpairsRequest _ count -> + eigenpairsFromColumns + dimension + ((\entryIndex -> (eigenvalue, unitVector dimension entryIndex, 0.0)) <$> [0 .. positiveCountValue count - 1]) + +eigenRequestCount :: EigenRequest result -> Int +eigenRequestCount requestValue = + case requestValue of + EigenvaluesRequest _ count -> positiveCountValue count + EigenpairsRequest _ count -> positiveCountValue count + +validateSpectralCount :: Int -> Int -> Either MoonlightError () +validateSpectralCount requestedCount dimension + | dimension <= 0 = Left (InvariantViolation "spectral solve requires a positive operator dimension") + | requestedCount <= 0 = Left (InvariantViolation "spectral request count must be positive") + | requestedCount > dimension = Left (InvariantViolation "spectral request count exceeds operator dimension") + | otherwise = Right () + +unitVector :: Int -> Int -> U.Vector Double +unitVector dimension selectedIndex = + U.generate dimension (\entryIndex -> if entryIndex == selectedIndex then 1.0 else 0.0) diff --git a/src-statics/Moonlight/LinAlg/Pure/Statics/Algebra.hs b/src-statics/Moonlight/LinAlg/Pure/Statics/Algebra.hs new file mode 100644 --- /dev/null +++ b/src-statics/Moonlight/LinAlg/Pure/Statics/Algebra.hs @@ -0,0 +1,41 @@ +module Moonlight.LinAlg.Pure.Statics.Algebra + ( allAxes, + mkMemberRef, + memberEndpoints, + memberTouchesNode, + addVec3, + subVec3, + negateVec3, + scaleVec3, + dotVec3, + magnitudeVec3, + normalizeVec3, + normalizeVec3Safe, + axisComponent, + axisVector, + vec3Zero, + ) +where + +import Moonlight.LinAlg.Pure.Statics.Types + ( memberEndpoints, + memberTouchesNode, + mkMemberRef, + ) +import Moonlight.LinAlg.Pure.Geometry.Vec3 + ( Axis (..), + addVec3, + axisComponent, + axisVector, + dotVec3, + magnitudeVec3, + negateVec3, + normalizeVec3, + normalizeVec3Safe, + scaleVec3, + subVec3, + vec3Zero, + ) + +allAxes :: [Axis] +allAxes = [AxisX, AxisY, AxisZ] diff --git a/src-statics/Moonlight/LinAlg/Pure/Statics/Compile.hs b/src-statics/Moonlight/LinAlg/Pure/Statics/Compile.hs new file mode 100644 --- /dev/null +++ b/src-statics/Moonlight/LinAlg/Pure/Statics/Compile.hs @@ -0,0 +1,159 @@ +module Moonlight.LinAlg.Pure.Statics.Compile + ( assembleEquilibriumEquations, + ) +where + +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Set (Set) +import qualified Data.Set as Set +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Pure.Dense.Dynamic (mkDynMatrix, mkDynVector) +import Moonlight.LinAlg.Pure.Statics.Algebra + ( allAxes, + axisComponent, + memberEndpoints, + mkMemberRef, + normalizeVec3, + scaleVec3, + subVec3, + ) +import Moonlight.LinAlg.Pure.Statics.Types + ( CompiledEquilibrium, + EquationRef (..), + ForceNetwork, + ForceNode, + MemberRef, + NodeRef, + UnknownForce (..), + Vec3, + forceMembers, + forceNodeLoad, + forceNodePosition, + forceNodeReactionAxes, + forceNodes, + mkCompiledEquilibrium, + ) +import Prelude + +assembleEquilibriumEquations :: ForceNetwork -> Either MoonlightError CompiledEquilibrium +assembleEquilibriumEquations networkValue = do + memberOrder <- canonicalMemberOrder (forceMembers networkValue) + let nodeOrder = Map.keys (forceNodes networkValue) + foundationEntries = supportedNodeEntries (forceNodes networkValue) + foundationOrder = fst <$> foundationEntries + reactionUnknowns = + concatMap + ( \(nodeRefValue, nodeValue) -> + fmap + (ReactionUnknown nodeRefValue) + (forceNodeReactionAxes nodeValue) + ) + foundationEntries + equationOrder = + concatMap + (\nodeRefValue -> fmap (EquationRef nodeRefValue) allAxes) + nodeOrder + unknownOrder = + fmap MemberUnknown memberOrder + <> reactionUnknowns + memberDirections <- Map.fromList <$> traverse (directionEntry networkValue) memberOrder + coefficientRows <- traverse (equationCoefficients networkValue memberDirections unknownOrder) equationOrder + rightHandSideValues <- traverse (equationRightHandSide networkValue) equationOrder + coefficientMatrix <- mkDynMatrix (length equationOrder) (length unknownOrder) (concat coefficientRows) + rightHandSideVector <- mkDynVector (length equationOrder) rightHandSideValues + mkCompiledEquilibrium + nodeOrder + foundationOrder + memberOrder + memberDirections + unknownOrder + equationOrder + coefficientMatrix + rightHandSideVector + +canonicalMemberOrder :: Set MemberRef -> Either MoonlightError [MemberRef] +canonicalMemberOrder memberRefs = + fmap Set.toAscList + ( Set.fromList + <$> traverse + ( \memberRefValue -> + case memberEndpoints memberRefValue of + (leftRef, rightRef) -> mkMemberRef leftRef rightRef + ) + (Set.toAscList memberRefs) + ) + +directionEntry :: ForceNetwork -> MemberRef -> Either MoonlightError (MemberRef, Vec3) +directionEntry networkValue memberRefValue = + fmap ((,) memberRefValue) (memberDirection networkValue memberRefValue) + +memberDirection :: ForceNetwork -> MemberRef -> Either MoonlightError Vec3 +memberDirection networkValue memberRefValue = + case memberEndpoints memberRefValue of + (leftRef, rightRef) -> do + leftNode <- lookupNode networkValue leftRef + rightNode <- lookupNode networkValue rightRef + normalizeVec3 + ( subVec3 + (forceNodePosition rightNode) + (forceNodePosition leftNode) + ) + +equationCoefficients :: + ForceNetwork -> + Map MemberRef Vec3 -> + [UnknownForce] -> + EquationRef -> + Either MoonlightError [Double] +equationCoefficients networkValue memberDirections unknownOrder equationRefValue = do + _ <- lookupNode networkValue (equationNodeRef equationRefValue) + pure + ( fmap + (unknownCoefficient memberDirections equationRefValue) + unknownOrder + ) + +equationRightHandSide :: ForceNetwork -> EquationRef -> Either MoonlightError Double +equationRightHandSide networkValue equationRefValue = do + nodeValue <- lookupNode networkValue (equationNodeRef equationRefValue) + pure + ( negate + (axisComponent (equationAxis equationRefValue) (forceNodeLoad nodeValue)) + ) + +unknownCoefficient :: Map MemberRef Vec3 -> EquationRef -> UnknownForce -> Double +unknownCoefficient memberDirections equationRefValue unknownValue = + case unknownValue of + MemberUnknown memberRefValue -> + maybe 0.0 + ( \directionValue -> + maybe + 0.0 + (axisComponent (equationAxis equationRefValue)) + (memberContribution (equationNodeRef equationRefValue) memberRefValue directionValue) + ) + (Map.lookup memberRefValue memberDirections) + ReactionUnknown reactionNode reactionAxisValue -> + if reactionNode == equationNodeRef equationRefValue && reactionAxisValue == equationAxis equationRefValue + then 1.0 + else 0.0 + +memberContribution :: NodeRef -> MemberRef -> Vec3 -> Maybe Vec3 +memberContribution nodeRefValue memberRefValue directionValue = + case memberEndpoints memberRefValue of + (leftRef, rightRef) + | nodeRefValue == leftRef -> Just (scaleVec3 (-1.0) directionValue) + | nodeRefValue == rightRef -> Just directionValue + | otherwise -> Nothing + +lookupNode :: ForceNetwork -> NodeRef -> Either MoonlightError ForceNode +lookupNode networkValue nodeRefValue = + case Map.lookup nodeRefValue (forceNodes networkValue) of + Nothing -> + Left (InvariantViolation ("force network member references unknown node " <> show nodeRefValue)) + Just nodeValue -> Right nodeValue + +supportedNodeEntries :: Map NodeRef ForceNode -> [(NodeRef, ForceNode)] +supportedNodeEntries = + filter (not . null . forceNodeReactionAxes . snd) . Map.toAscList diff --git a/src-statics/Moonlight/LinAlg/Pure/Statics/Core.hs b/src-statics/Moonlight/LinAlg/Pure/Statics/Core.hs new file mode 100644 --- /dev/null +++ b/src-statics/Moonlight/LinAlg/Pure/Statics/Core.hs @@ -0,0 +1,409 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module Moonlight.LinAlg.Pure.Statics.Core + ( checkEquilibrium, + solveGraphicStatics, + ) +where + +import Control.Monad (join) +import Data.Graph (SCC (..), stronglyConnComp) +import Data.List.NonEmpty (NonEmpty) +import qualified Data.List.NonEmpty as NonEmpty +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Set (Set) +import qualified Data.Set as Set +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Pure.Dense.Rows (transposeRowsExact) +import Moonlight.LinAlg.Pure.Dense.Dynamic + ( DynMatrix, + DynVector, + dynMatrixToRows, + dynMatrixShape, + dynVectorLength, + dynVectorToList, + fromDynMatrix, + fromDynVector, + mkDynMatrix, + mkDynVector, + toDynVector, + withDynMatrix, + withDynVector, + ) +import Moonlight.LinAlg.Internal.Primitives (matrixVectorProduct) +import Moonlight.LinAlg.Pure.Dense.Decomposition (qrDecompFullColumnRank) +import Moonlight.LinAlg.Pure.Dense.Solver (solveDirect) +import Moonlight.LinAlg.Pure.Statics.Algebra + ( addVec3, + axisVector, + magnitudeVec3, + memberEndpoints, + memberTouchesNode, + vec3Zero, + ) +import Moonlight.LinAlg.Pure.Statics.Compile (assembleEquilibriumEquations) +import Moonlight.LinAlg.Pure.Statics.Types + ( CompiledEquilibrium, + EquationRef (..), + EquilibriumResult (..), + EquilibriumSolution (..), + EquilibriumViolation (..), + ForceNetwork, + ForceSign (..), + MemberRef, + NodeRef, + UnknownForce (..), + Vec3, + compiledCoefficientMatrix, + compiledEquationOrder, + compiledFoundationOrder, + compiledMemberOrder, + compiledNodeOrder, + compiledRightHandSide, + compiledUnknownOrder, + ) +import Moonlight.LinAlg.Pure.Dense.Types + ( Matrix, + Vector, + fromListVector, + matrixToRows, + toListVector, + ) +import Prelude + +checkEquilibrium :: ForceNetwork -> Either MoonlightError EquilibriumResult +checkEquilibrium networkValue = + assembleEquilibriumEquations networkValue >>= solveGraphicStatics + +solveGraphicStatics :: CompiledEquilibrium -> Either MoonlightError EquilibriumResult +solveGraphicStatics compiledValue = do + solvedUnknowns <- solveUnknowns compiledValue + solutionValues <- pure (dynVectorToList solvedUnknowns) + residualForces <- solveResiduals compiledValue solutionValues + let solutionValue = interpretSolution compiledValue solutionValues residualForces + violations = collectViolations compiledValue solutionValue + pure + ( maybe + (InEquilibrium solutionValue) + Disequilibrium + violations + ) + +solveUnknowns :: CompiledEquilibrium -> Either MoonlightError (DynVector Double) +solveUnknowns compiledValue = + case componentNodeSets compiledValue of + [] -> solveUnknownsDense compiledValue + [_] -> solveUnknownsDense compiledValue + components -> solveUnknownsByComponents components compiledValue + +solveUnknownsDense :: CompiledEquilibrium -> Either MoonlightError (DynVector Double) +solveUnknownsDense compiledValue = + solveDenseSystemByShape + (compiledCoefficientMatrix compiledValue) + (compiledRightHandSide compiledValue) + +solveUnknownsByComponents :: [Set NodeRef] -> CompiledEquilibrium -> Either MoonlightError (DynVector Double) +solveUnknownsByComponents components compiledValue = do + coefficientRows <- dynMatrixToRows (compiledCoefficientMatrix compiledValue) + solvedEntries <- + fmap concat + ( traverse + (solveComponentUnknowns compiledValue coefficientRows (dynVectorToList (compiledRightHandSide compiledValue))) + components + ) + let solvedMap = Map.fromList solvedEntries + unknownCount = length (compiledUnknownOrder compiledValue) + solutionValues <- + traverse + ( \unknownIndex -> + maybe + (Left (InvariantViolation ("graphic statics component solve omitted unknown index " <> show unknownIndex))) + Right + (Map.lookup unknownIndex solvedMap) + ) + [0 .. unknownCount - 1] + mkDynVector unknownCount solutionValues + +solveComponentUnknowns :: + CompiledEquilibrium -> + [[Double]] -> + [Double] -> + Set NodeRef -> + Either MoonlightError [(Int, Double)] +solveComponentUnknowns compiledValue coefficientRows rightHandSideValues componentNodes = do + let equationEntries = + filter + ( \(_, equationRefValue) -> + Set.member (equationNodeRef equationRefValue) componentNodes + ) + (indexedValues (compiledEquationOrder compiledValue)) + unknownEntries = + filter + (unknownEntryInComponent componentNodes) + (indexedValues (compiledUnknownOrder compiledValue)) + componentRows <- + traverse + ( \(equationIndex, _) -> do + rowValues <- selectIndex "graphic statics component equation row" equationIndex coefficientRows + traverse + (\(unknownIndex, _) -> selectIndex "graphic statics component unknown column" unknownIndex rowValues) + unknownEntries + ) + equationEntries + componentRightHandSide <- + traverse + (\(equationIndex, _) -> selectIndex "graphic statics component RHS" equationIndex rightHandSideValues) + equationEntries + componentMatrix <- mkDynMatrix (length equationEntries) (length unknownEntries) (concat componentRows) + componentVector <- mkDynVector (length equationEntries) componentRightHandSide + componentSolution <- solveComponentDenseSystem componentMatrix componentVector + let solutionValues = dynVectorToList componentSolution + if length solutionValues /= length unknownEntries + then Left (InvariantViolation "graphic statics component solve returned wrong unknown count") + else Right (zip (fst <$> unknownEntries) solutionValues) + +solveComponentDenseSystem :: DynMatrix Double -> DynVector Double -> Either MoonlightError (DynVector Double) +solveComponentDenseSystem = + solveDenseSystemByShape + +solveDenseSystemByShape :: DynMatrix Double -> DynVector Double -> Either MoonlightError (DynVector Double) +solveDenseSystemByShape coefficientMatrix rightHandSide = + let (rowCount, columnCount) = dynMatrixShape coefficientMatrix + in if rowCount == columnCount && rowCount == dynVectorLength rightHandSide + then solveSquareSystem coefficientMatrix rightHandSide + else solveLeastSquares coefficientMatrix rightHandSide + +componentNodeSets :: CompiledEquilibrium -> [Set NodeRef] +componentNodeSets compiledValue = + Set.fromList . flattenSCC + <$> stronglyConnComp + ( (\nodeRefValue -> (nodeRefValue, nodeRefValue, Map.findWithDefault [] nodeRefValue adjacencyMap)) + <$> compiledNodeOrder compiledValue + ) + where + adjacencyMap = + Map.fromListWith + (<>) + (componentMemberAdjacency =<< compiledMemberOrder compiledValue) + +componentMemberAdjacency :: MemberRef -> [(NodeRef, [NodeRef])] +componentMemberAdjacency memberRefValue = + case memberEndpoints memberRefValue of + (leftRef, rightRef) -> + [ (leftRef, [rightRef]), + (rightRef, [leftRef]) + ] + +flattenSCC :: SCC node -> [node] +flattenSCC component = + case component of + AcyclicSCC nodeValue -> [nodeValue] + CyclicSCC nodeValues -> nodeValues + +unknownEntryInComponent :: Set NodeRef -> (Int, UnknownForce) -> Bool +unknownEntryInComponent componentNodes (_, unknownValue) = + case unknownValue of + MemberUnknown memberRefValue -> + case memberEndpoints memberRefValue of + (leftRef, rightRef) -> + Set.member leftRef componentNodes || Set.member rightRef componentNodes + ReactionUnknown nodeRefValue _ -> + Set.member nodeRefValue componentNodes + +indexedValues :: [value] -> [(Int, value)] +indexedValues = + zip [0 ..] + +selectIndex :: String -> Int -> [value] -> Either MoonlightError value +selectIndex context indexValue values + | indexValue < 0 = + Left (InvariantViolation (context <> " index must be non-negative: " <> show indexValue)) + | otherwise = + case drop indexValue values of + value : _ -> Right value + [] -> + Left + ( InvariantViolation + ( context + <> " index out of bounds: index=" + <> show indexValue + <> ", length=" + <> show (length values) + ) + ) + +solveSquareSystem :: DynMatrix Double -> DynVector Double -> Either MoonlightError (DynVector Double) +solveSquareSystem coefficientMatrix rightHandSide + | rowCount /= columnCount = + Left (InvariantViolation "graphic statics direct solve requires a square coefficient matrix") + | rowCount /= dynVectorLength rightHandSide = + Left (InvariantViolation "graphic statics direct solve RHS length mismatch") + | otherwise = + join + ( withDynVector rightHandSide + ( \(staticRightHandSide :: Vector n Double) -> do + staticMatrix <- (fromDynMatrix coefficientMatrix :: Either MoonlightError (Matrix n n Double)) + toDynVector <$> solveDirect staticMatrix staticRightHandSide + ) + ) + where + (rowCount, columnCount) = dynMatrixShape coefficientMatrix + +solveLeastSquares :: DynMatrix Double -> DynVector Double -> Either MoonlightError (DynVector Double) +solveLeastSquares coefficientMatrix rightHandSide + | rowCount /= dynVectorLength rightHandSide = + Left (InvariantViolation "graphic statics least-squares RHS length mismatch") + | rowCount < columnCount = + Left (InvariantViolation "graphic statics QR least-squares requires row count greater than or equal to unknown count") + | otherwise = + join + ( withDynMatrix coefficientMatrix + ( \(staticMatrix :: Matrix rows columns Double) -> do + staticRightHandSide <- (fromDynVector rightHandSide :: Either MoonlightError (Vector rows Double)) + (qMatrix, rMatrix) <- qrDecompFullColumnRank staticMatrix + qRows <- matrixToRows qMatrix + qTransposeRows <- transposeRowsExact qRows + projectedRightHandSideValues <- matrixVectorProduct qTransposeRows (toListVector staticRightHandSide) + projectedRightHandSide <- fromListVector @columns projectedRightHandSideValues + toDynVector <$> solveDirect rMatrix projectedRightHandSide + ) + ) + where + (rowCount, columnCount) = dynMatrixShape coefficientMatrix + +solveResiduals :: CompiledEquilibrium -> [Double] -> Either MoonlightError (Map NodeRef Vec3) +solveResiduals compiledValue solvedUnknowns = do + coefficientRows <- dynMatrixToRows (compiledCoefficientMatrix compiledValue) + let rightHandSideValues = dynVectorToList (compiledRightHandSide compiledValue) + predictedValues <- matrixVectorProduct coefficientRows solvedUnknowns + if length predictedValues /= length rightHandSideValues + then Left (InvariantViolation "graphic statics residual computation length mismatch") + else + pure + ( foldl' + accumulateResidual + Map.empty + ( zip + (compiledEquationOrder compiledValue) + (zipWith (-) predictedValues rightHandSideValues) + ) + ) + +interpretSolution :: CompiledEquilibrium -> [Double] -> Map NodeRef Vec3 -> EquilibriumSolution +interpretSolution compiledValue solvedUnknowns residualForces = + let solutionEntries = zip (compiledUnknownOrder compiledValue) solvedUnknowns + (memberForces, reactionForces) = + foldl' + accumulateUnknown + (Map.empty, Map.empty) + solutionEntries + in EquilibriumSolution + { equilibriumMemberForces = memberForces, + equilibriumReactionForces = + foldl' + (\reactionMap nodeRefValue -> Map.insertWith addVec3 nodeRefValue vec3Zero reactionMap) + reactionForces + (compiledFoundationOrder compiledValue), + equilibriumResidualForces = + foldl' + (\residualMap nodeRefValue -> Map.insertWith addVec3 nodeRefValue vec3Zero residualMap) + residualForces + (compiledNodeOrder compiledValue) + } + +collectViolations :: CompiledEquilibrium -> EquilibriumSolution -> Maybe (NonEmpty EquilibriumViolation) +collectViolations compiledValue solutionValue = + NonEmpty.nonEmpty + ( foldMap + (violationAtNode compiledValue solutionValue) + (compiledNodeOrder compiledValue) + ) + +violationAtNode :: CompiledEquilibrium -> EquilibriumSolution -> NodeRef -> [EquilibriumViolation] +violationAtNode compiledValue solutionValue nodeRefValue = + let residualForce = + Map.findWithDefault vec3Zero nodeRefValue (equilibriumResidualForces solutionValue) + residualMagnitude = magnitudeVec3 residualForce + memberDetails = incidentMembers compiledValue solutionValue nodeRefValue + worstMember = strongestMember memberDetails + tensionMember = strongestTension memberDetails + selectedMember = maybe worstMember Just tensionMember + selectedSign = + fmap + (\(_, forceValue) -> if forceValue < 0.0 then Tension else Compression) + selectedMember + in if residualMagnitude > equilibriumTolerance || tensionMember /= Nothing + then + [ EquilibriumViolation + { violationNode = nodeRefValue, + violationResidualForce = residualForce, + violationResidualMagnitude = residualMagnitude, + violationWorstMember = fmap fst selectedMember, + violationMemberForceSign = selectedSign + } + ] + else [] + +incidentMembers :: CompiledEquilibrium -> EquilibriumSolution -> NodeRef -> [(MemberRef, Double)] +incidentMembers compiledValue solutionValue nodeRefValue = + fmap + (\memberRefValue -> (memberRefValue, Map.findWithDefault 0.0 memberRefValue (equilibriumMemberForces solutionValue))) + ( filter + (memberTouchesNode nodeRefValue) + (compiledMemberOrder compiledValue) + ) + +strongestMember :: [(MemberRef, Double)] -> Maybe (MemberRef, Double) +strongestMember = + foldl' + ( \currentBest candidate -> + case currentBest of + Nothing -> Just candidate + Just bestCandidate -> + if abs (snd candidate) > abs (snd bestCandidate) + then Just candidate + else currentBest + ) + Nothing + +strongestTension :: [(MemberRef, Double)] -> Maybe (MemberRef, Double) +strongestTension = + foldl' + ( \currentBest candidate -> + if snd candidate < (-equilibriumTolerance) + then + case currentBest of + Nothing -> Just candidate + Just bestCandidate -> + if snd candidate < snd bestCandidate + then Just candidate + else currentBest + else currentBest + ) + Nothing + +accumulateUnknown :: + (Map MemberRef Double, Map NodeRef Vec3) -> + (UnknownForce, Double) -> + (Map MemberRef Double, Map NodeRef Vec3) +accumulateUnknown (memberForces, reactionForces) (unknownValue, magnitudeValue) = + case unknownValue of + MemberUnknown memberRefValue -> + (Map.insert memberRefValue magnitudeValue memberForces, reactionForces) + ReactionUnknown nodeRefValue axisValue -> + ( memberForces, + Map.insertWith addVec3 nodeRefValue (axisVector axisValue magnitudeValue) reactionForces + ) + +accumulateResidual :: Map NodeRef Vec3 -> (EquationRef, Double) -> Map NodeRef Vec3 +accumulateResidual residuals (equationRefValue, magnitudeValue) = + Map.insertWith + addVec3 + (equationNodeRef equationRefValue) + (axisVector (equationAxis equationRefValue) magnitudeValue) + residuals + +equilibriumTolerance :: Double +equilibriumTolerance = 1.0e-8 diff --git a/src-statics/Moonlight/LinAlg/Pure/Statics/Network.hs b/src-statics/Moonlight/LinAlg/Pure/Statics/Network.hs new file mode 100644 --- /dev/null +++ b/src-statics/Moonlight/LinAlg/Pure/Statics/Network.hs @@ -0,0 +1,376 @@ +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Statics.Network + ( NetworkDeclaration, + NetworkBuildError (..), + nodeRef, + nodeRefLabel, + joint, + support, + supportOn, + load, + member, + network, + nodePosition, + nodeLoad, + nodeSupportAxes, + nodeReactionAxes, + networkNodeMap, + networkMemberSet, + ) +where + +import Control.Monad (foldM) +import Data.Foldable (traverse_) +import Data.Kind (Type) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Set (Set) +import qualified Data.Set as Set +import Moonlight.Core (fieldValueValid) +import Moonlight.LinAlg.Pure.Geometry.Vec3 + ( Axis, + Vec3 (..), + magnitudeVec3, + normalizeVec3, + subVec3, + ) +import Moonlight.LinAlg.Pure.Statics.Types + ( ForceNetwork (..), + ForceNode, + MemberRef, + NodeRef (..), + SupportAxes, + fixedSupportAxes, + forceMembers, + forceNodeLoad, + forceNodePosition, + forceNodeReactionAxes, + forceNodeSupportAxes, + forceNodes, + memberEndpoints, + mkMemberRef, + mkSupportAxes, + supportedForceNode, + supportAxesList, + ) +import Prelude + +type NetworkDeclaration :: Type +data NetworkDeclaration + = JointDeclaration !String !Vec3 + | SupportDeclaration !String !Vec3 !SupportAxes + | LoadDeclaration !String !Vec3 !Vec3 + | MemberDeclaration !String !String + deriving stock (Eq, Show) + +type NetworkBuildError :: Type +data NetworkBuildError + = EmptyNodeLabel + | NonFiniteNodePosition !String !Vec3 + | NonFiniteNodeLoad !String !Vec3 + | MissingNodePosition !String + | ConflictingNodePosition !String !Vec3 !Vec3 + | NonFiniteAccumulatedLoad !String !Vec3 + | UnknownMemberEndpoint !String + | SelfMember !String + | NonFiniteMemberGeometry !String !String + | DegenerateMember !String !String + deriving stock (Eq, Show) + +nodeRef :: String -> Either NetworkBuildError NodeRef +nodeRef labelValue + | null labelValue = Left EmptyNodeLabel + | otherwise = Right (NodeRef labelValue) + +nodeRefLabel :: NodeRef -> String +nodeRefLabel (NodeRef labelValue) = + labelValue + +joint :: String -> Vec3 -> NetworkDeclaration +joint = + JointDeclaration + +support :: String -> Vec3 -> NetworkDeclaration +support labelValue positionValue = + SupportDeclaration labelValue positionValue fixedSupportAxes + +supportOn :: String -> Vec3 -> SupportAxes -> NetworkDeclaration +supportOn = + SupportDeclaration + +load :: String -> Vec3 -> Vec3 -> NetworkDeclaration +load = + LoadDeclaration + +member :: String -> String -> NetworkDeclaration +member = + MemberDeclaration + +network :: [NetworkDeclaration] -> Either NetworkBuildError ForceNetwork +network declarations = do + accumulatedNetwork <- + foldM + collectDeclaration + emptyNetworkAccumulator + declarations + finalizedNodes <- + Map.traverseWithKey + finalizePartialNode + (networkAccumulatorNodes accumulatedNetwork) + traverse_ + (validateMemberGeometry finalizedNodes) + (Set.toAscList (networkAccumulatorMembers accumulatedNetwork)) + Right + ForceNetwork + { forceNodes = finalizedNodes, + forceMembers = networkAccumulatorMembers accumulatedNetwork + } + +nodePosition :: ForceNode -> Vec3 +nodePosition = + forceNodePosition + +nodeLoad :: ForceNode -> Vec3 +nodeLoad = + forceNodeLoad + +nodeSupportAxes :: ForceNode -> SupportAxes +nodeSupportAxes = + forceNodeSupportAxes + +nodeReactionAxes :: ForceNode -> [Axis] +nodeReactionAxes = + forceNodeReactionAxes + +networkNodeMap :: ForceNetwork -> Map NodeRef ForceNode +networkNodeMap = + forceNodes + +networkMemberSet :: ForceNetwork -> Set MemberRef +networkMemberSet = + forceMembers + +type ExactVec3 :: Type +data ExactVec3 = ExactVec3 + { exactVec3X :: !Rational, + exactVec3Y :: !Rational, + exactVec3Z :: !Rational + } + deriving stock (Eq, Show) + +zeroExactVec3 :: ExactVec3 +zeroExactVec3 = + ExactVec3 0 0 0 + +exactVec3FromVec3 :: Vec3 -> ExactVec3 +exactVec3FromVec3 (Vec3 xValue yValue zValue) = + ExactVec3 + (toRational xValue) + (toRational yValue) + (toRational zValue) + +addExactVec3 :: ExactVec3 -> ExactVec3 -> ExactVec3 +addExactVec3 leftValue rightValue = + ExactVec3 + { exactVec3X = exactVec3X leftValue + exactVec3X rightValue, + exactVec3Y = exactVec3Y leftValue + exactVec3Y rightValue, + exactVec3Z = exactVec3Z leftValue + exactVec3Z rightValue + } + +type PartialNode :: Type +data PartialNode = PartialNode + { partialNodePositions :: !(Set Vec3), + partialNodeLoadExact :: !ExactVec3, + partialNodeSupportAxes :: !(Set Axis) + } + deriving stock (Eq, Show) + +emptyPartialNode :: PartialNode +emptyPartialNode = + PartialNode + { partialNodePositions = Set.empty, + partialNodeLoadExact = zeroExactVec3, + partialNodeSupportAxes = Set.empty + } + +type NetworkAccumulator :: Type +data NetworkAccumulator = NetworkAccumulator + { networkAccumulatorNodes :: !(Map NodeRef PartialNode), + networkAccumulatorMembers :: !(Set MemberRef) + } + deriving stock (Eq, Show) + +emptyNetworkAccumulator :: NetworkAccumulator +emptyNetworkAccumulator = + NetworkAccumulator + { networkAccumulatorNodes = Map.empty, + networkAccumulatorMembers = Set.empty + } + +collectDeclaration :: + NetworkAccumulator -> + NetworkDeclaration -> + Either NetworkBuildError NetworkAccumulator +collectDeclaration accumulator declarationValue = + case declarationValue of + JointDeclaration labelValue positionValue -> + collectNodeDeclaration + labelValue + positionValue + Nothing + Set.empty + accumulator + SupportDeclaration labelValue positionValue supportAxes -> + collectNodeDeclaration + labelValue + positionValue + Nothing + (Set.fromList (supportAxesList supportAxes)) + accumulator + LoadDeclaration labelValue positionValue loadValue -> + collectNodeDeclaration + labelValue + positionValue + (Just loadValue) + Set.empty + accumulator + MemberDeclaration leftLabel rightLabel -> + collectMemberDeclaration + leftLabel + rightLabel + accumulator + +collectNodeDeclaration :: + String -> + Vec3 -> + Maybe Vec3 -> + Set Axis -> + NetworkAccumulator -> + Either NetworkBuildError NetworkAccumulator +collectNodeDeclaration labelValue positionValue maybeLoadValue supportAxes accumulator = do + nodeReference <- nodeRef labelValue + if finiteVec3 positionValue + then Right () + else Left (NonFiniteNodePosition labelValue positionValue) + traverse_ validateLoad maybeLoadValue + let currentPartialNode = + Map.findWithDefault + emptyPartialNode + nodeReference + (networkAccumulatorNodes accumulator) + loadContribution = + maybe zeroExactVec3 exactVec3FromVec3 maybeLoadValue + updatedPartialNode = + currentPartialNode + { partialNodePositions = + Set.insert positionValue (partialNodePositions currentPartialNode), + partialNodeLoadExact = + addExactVec3 (partialNodeLoadExact currentPartialNode) loadContribution, + partialNodeSupportAxes = + Set.union (partialNodeSupportAxes currentPartialNode) supportAxes + } + Right + accumulator + { networkAccumulatorNodes = + Map.insert + nodeReference + updatedPartialNode + (networkAccumulatorNodes accumulator) + } + where + validateLoad loadValue + | finiteVec3 loadValue = Right () + | otherwise = Left (NonFiniteNodeLoad labelValue loadValue) + +collectMemberDeclaration :: + String -> + String -> + NetworkAccumulator -> + Either NetworkBuildError NetworkAccumulator +collectMemberDeclaration leftLabel rightLabel accumulator = do + leftReference <- nodeRef leftLabel + rightReference <- nodeRef rightLabel + if leftReference == rightReference + then Left (SelfMember leftLabel) + else + case mkMemberRef leftReference rightReference of + Left _ -> Left (SelfMember leftLabel) + Right memberReference -> + Right + accumulator + { networkAccumulatorMembers = + Set.insert + memberReference + (networkAccumulatorMembers accumulator) + } + +finalizePartialNode :: + NodeRef -> + PartialNode -> + Either NetworkBuildError ForceNode +finalizePartialNode nodeReference partialNode = do + positionValue <- + case Set.toAscList (partialNodePositions partialNode) of + [] -> Left (MissingNodePosition (nodeRefLabel nodeReference)) + [singlePosition] -> Right singlePosition + firstPosition : secondPosition : _ -> + Left + ( ConflictingNodePosition + (nodeRefLabel nodeReference) + firstPosition + secondPosition + ) + loadValue <- + finalizeExactLoad + (nodeRefLabel nodeReference) + (partialNodeLoadExact partialNode) + let supportAxes = mkSupportAxes (Set.toAscList (partialNodeSupportAxes partialNode)) + Right (supportedForceNode positionValue loadValue supportAxes) + +finalizeExactLoad :: String -> ExactVec3 -> Either NetworkBuildError Vec3 +finalizeExactLoad labelValue exactLoad = + let loadValue = + Vec3 + (fromRational (exactVec3X exactLoad)) + (fromRational (exactVec3Y exactLoad)) + (fromRational (exactVec3Z exactLoad)) + in if finiteVec3 loadValue + then Right loadValue + else Left (NonFiniteAccumulatedLoad labelValue loadValue) + +validateMemberGeometry :: + Map NodeRef ForceNode -> + MemberRef -> + Either NetworkBuildError () +validateMemberGeometry nodeValues memberReference = do + let (leftReference, rightReference) = memberEndpoints memberReference + leftLabel = nodeRefLabel leftReference + rightLabel = nodeRefLabel rightReference + leftNode <- requireNode leftLabel leftReference nodeValues + rightNode <- requireNode rightLabel rightReference nodeValues + let displacement = + subVec3 + (forceNodePosition rightNode) + (forceNodePosition leftNode) + displacementMagnitude = magnitudeVec3 displacement + if not (finiteVec3 displacement) || not (fieldValueValid displacementMagnitude) + then Left (NonFiniteMemberGeometry leftLabel rightLabel) + else + case normalizeVec3 displacement of + Left _ -> Left (DegenerateMember leftLabel rightLabel) + Right _ -> Right () + +requireNode :: + String -> + NodeRef -> + Map NodeRef ForceNode -> + Either NetworkBuildError ForceNode +requireNode labelValue nodeReference nodeValues = + case Map.lookup nodeReference nodeValues of + Nothing -> Left (UnknownMemberEndpoint labelValue) + Just nodeValue -> Right nodeValue + +finiteVec3 :: Vec3 -> Bool +finiteVec3 (Vec3 xValue yValue zValue) = + fieldValueValid xValue && fieldValueValid yValue && fieldValueValid zValue diff --git a/src-statics/Moonlight/LinAlg/Pure/Statics/Types.hs b/src-statics/Moonlight/LinAlg/Pure/Statics/Types.hs new file mode 100644 --- /dev/null +++ b/src-statics/Moonlight/LinAlg/Pure/Statics/Types.hs @@ -0,0 +1,251 @@ +module Moonlight.LinAlg.Pure.Statics.Types + ( NodeRef (..), + Axis (..), + Vec3 (..), + MemberRef, + mkMemberRef, + memberEndpoints, + memberTouchesNode, + SupportAxes, + mkSupportAxes, + supportAxesList, + freeSupportAxes, + fixedSupportAxes, + ForceNode, + freeForceNode, + supportedForceNode, + fixedForceNode, + forceNodePosition, + forceNodeLoad, + forceNodeSupportAxes, + forceNodeReactionAxes, + ForceNetwork (..), + UnknownForce (..), + EquationRef (..), + CompiledEquilibrium, + mkCompiledEquilibrium, + compiledNodeOrder, + compiledFoundationOrder, + compiledMemberOrder, + compiledMemberDirections, + compiledUnknownOrder, + compiledEquationOrder, + compiledCoefficientMatrix, + compiledRightHandSide, + EquilibriumSolution (..), + ForceSign (..), + EquilibriumViolation (..), + EquilibriumResult (..), + ) +where + +import Data.Kind (Type) +import Data.List.NonEmpty (NonEmpty) +import Data.Map.Strict (Map) +import Data.Set (Set) +import qualified Data.Set as Set +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Pure.Dense.Dynamic + ( DynMatrix, + DynVector, + dynMatrixShape, + dynVectorLength, + ) +import Moonlight.LinAlg.Pure.Geometry.Vec3 (Axis (..), Vec3 (..)) +import Prelude + +type NodeRef :: Type +newtype NodeRef = NodeRef + { unNodeRef :: String + } + deriving stock (Eq, Ord, Show) + +type MemberRef :: Type +data MemberRef = MemberRef NodeRef NodeRef + deriving stock (Eq, Ord, Show) + +mkMemberRef :: NodeRef -> NodeRef -> Either MoonlightError MemberRef +mkMemberRef leftRef rightRef + | leftRef == rightRef = Left (InvariantViolation "member endpoints must be distinct") + | leftRef < rightRef = Right (MemberRef leftRef rightRef) + | otherwise = Right (MemberRef rightRef leftRef) + +memberEndpoints :: MemberRef -> (NodeRef, NodeRef) +memberEndpoints (MemberRef leftRef rightRef) = + (leftRef, rightRef) + +memberTouchesNode :: NodeRef -> MemberRef -> Bool +memberTouchesNode nodeRefValue memberRefValue = + case memberEndpoints memberRefValue of + (leftRef, rightRef) -> nodeRefValue == leftRef || nodeRefValue == rightRef + +type SupportAxes :: Type +newtype SupportAxes = SupportAxes + { supportAxesSet :: Set Axis + } + deriving stock (Eq, Ord, Show) + +mkSupportAxes :: [Axis] -> SupportAxes +mkSupportAxes = + SupportAxes . Set.fromList + +supportAxesList :: SupportAxes -> [Axis] +supportAxesList = + Set.toAscList . supportAxesSet + +freeSupportAxes :: SupportAxes +freeSupportAxes = + SupportAxes Set.empty + +fixedSupportAxes :: SupportAxes +fixedSupportAxes = + mkSupportAxes [AxisX, AxisY, AxisZ] + +type ForceNode :: Type +data ForceNode = ForceNode + { forceNodePosition :: Vec3, + forceNodeLoad :: Vec3, + forceNodeSupportAxesValue :: SupportAxes + } + deriving stock (Eq, Show) + +freeForceNode :: Vec3 -> Vec3 -> ForceNode +freeForceNode position load = + ForceNode position load freeSupportAxes + +supportedForceNode :: Vec3 -> Vec3 -> SupportAxes -> ForceNode +supportedForceNode = + ForceNode + +fixedForceNode :: Vec3 -> Vec3 -> ForceNode +fixedForceNode position load = + ForceNode position load fixedSupportAxes + +forceNodeSupportAxes :: ForceNode -> SupportAxes +forceNodeSupportAxes = + forceNodeSupportAxesValue + +forceNodeReactionAxes :: ForceNode -> [Axis] +forceNodeReactionAxes = + supportAxesList . forceNodeSupportAxes + +type ForceNetwork :: Type +data ForceNetwork = ForceNetwork + { forceNodes :: Map NodeRef ForceNode, + forceMembers :: Set MemberRef + } + deriving stock (Eq, Show) + +type UnknownForce :: Type +data UnknownForce + = MemberUnknown MemberRef + | ReactionUnknown NodeRef Axis + deriving stock (Eq, Ord, Show) + +type EquationRef :: Type +data EquationRef = EquationRef + { equationNodeRef :: NodeRef, + equationAxis :: Axis + } + deriving stock (Eq, Ord, Show) + +type CompiledEquilibrium :: Type +data CompiledEquilibrium = CompiledEquilibrium + { compiledNodeOrder :: [NodeRef], + compiledFoundationOrder :: [NodeRef], + compiledMemberOrder :: [MemberRef], + compiledMemberDirections :: Map MemberRef Vec3, + compiledUnknownOrder :: [UnknownForce], + compiledEquationOrder :: [EquationRef], + compiledCoefficientMatrix :: DynMatrix Double, + compiledRightHandSide :: DynVector Double + } + +mkCompiledEquilibrium :: + [NodeRef] -> + [NodeRef] -> + [MemberRef] -> + Map MemberRef Vec3 -> + [UnknownForce] -> + [EquationRef] -> + DynMatrix Double -> + DynVector Double -> + Either MoonlightError CompiledEquilibrium +mkCompiledEquilibrium nodeOrder foundationOrder memberOrder memberDirections unknownOrder equationOrder coefficientMatrix rightHandSide = + let (matrixRowCount, matrixColumnCount) = dynMatrixShape coefficientMatrix + equationCount = length equationOrder + unknownCount = length unknownOrder + in if matrixRowCount /= equationCount + then + Left + ( InvariantViolation + ( "compiled equilibrium coefficient row count mismatch: expected " + <> show equationCount + <> " rows but received " + <> show matrixRowCount + ) + ) + else + if matrixColumnCount /= unknownCount + then + Left + ( InvariantViolation + ( "compiled equilibrium coefficient column count mismatch: expected " + <> show unknownCount + <> " columns but received " + <> show matrixColumnCount + ) + ) + else + if dynVectorLength rightHandSide /= equationCount + then + Left + ( InvariantViolation + ( "compiled equilibrium RHS length mismatch: expected " + <> show equationCount + <> " entries but received " + <> show (dynVectorLength rightHandSide) + ) + ) + else + Right + CompiledEquilibrium + { compiledNodeOrder = nodeOrder, + compiledFoundationOrder = foundationOrder, + compiledMemberOrder = memberOrder, + compiledMemberDirections = memberDirections, + compiledUnknownOrder = unknownOrder, + compiledEquationOrder = equationOrder, + compiledCoefficientMatrix = coefficientMatrix, + compiledRightHandSide = rightHandSide + } + +type EquilibriumSolution :: Type +data EquilibriumSolution = EquilibriumSolution + { equilibriumMemberForces :: Map MemberRef Double, + equilibriumReactionForces :: Map NodeRef Vec3, + equilibriumResidualForces :: Map NodeRef Vec3 + } + deriving stock (Eq, Show) + +type ForceSign :: Type +data ForceSign + = Compression + | Tension + deriving stock (Eq, Ord, Show, Read) + +type EquilibriumViolation :: Type +data EquilibriumViolation = EquilibriumViolation + { violationNode :: NodeRef, + violationResidualForce :: Vec3, + violationResidualMagnitude :: Double, + violationWorstMember :: Maybe MemberRef, + violationMemberForceSign :: Maybe ForceSign + } + deriving stock (Eq, Show) + +type EquilibriumResult :: Type +data EquilibriumResult + = InEquilibrium EquilibriumSolution + | Disequilibrium (NonEmpty EquilibriumViolation) + deriving stock (Eq, Show) diff --git a/src-structured/Moonlight/LinAlg/Pure/Structured/BlockTridiagonal.hs b/src-structured/Moonlight/LinAlg/Pure/Structured/BlockTridiagonal.hs new file mode 100644 --- /dev/null +++ b/src-structured/Moonlight/LinAlg/Pure/Structured/BlockTridiagonal.hs @@ -0,0 +1,518 @@ +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Structured.BlockTridiagonal + ( RowMajorBlock, + mkRowMajorBlock, + rowMajorBlockRows, + rowMajorBlockColumns, + rowMajorBlockPayload, + transposeRowMajorBlock, + symmetrizeRowMajorBlockLower, + rowMajorBlockEntry, + SymmetricBlockTridiagonal, + mkSymmetricBlockTridiagonal, + symmetricBlockTridiagonalDimension, + symmetricBlockTridiagonalBlockCount, + symmetricBlockTridiagonalBandwidth, + symmetricBlockTridiagonalEntry, + blockOffsets, + diagonalPayloadOffsets, + diagonalLowerPacked, + couplingPayloadOffsets, + lowerCouplingPayload, + applySymmetricBlockTridiagonalU, + symmetricBlockTridiagonalUpperBound, + symmetricBlockTridiagonalFrobeniusNorm, + ) +where + +import Control.Applicative ((<|>)) +import Control.Monad.ST (ST, runST) +import Data.Bifunctor (first) +import Data.Foldable (traverse_) +import Data.Kind (Type) +import Data.Vector qualified as Box +import Data.Vector.Unboxed qualified as U +import Data.Vector.Unboxed.Mutable qualified as MU +import Moonlight.Core + ( MoonlightError (..), + checkedNaturalToInt, + checkedNonNegativeProduct, + fieldValueValid, + ) +import Numeric.Natural (Natural) +import Prelude + +type RowMajorBlock :: Type +data RowMajorBlock = RowMajorBlock + { rowMajorBlockRows :: !Int, + rowMajorBlockColumns :: !Int, + rowMajorBlockPayload :: !(U.Vector Double) + } + deriving stock (Eq, Show) + +mkRowMajorBlock :: Int -> Int -> U.Vector Double -> Either MoonlightError RowMajorBlock +mkRowMajorBlock rowCount columnCount payload + | rowCount <= 0 || columnCount <= 0 = + Left (InvariantViolation "row-major block dimensions must be positive") + | otherwise = do + expectedLength <- + first + (const (InvariantViolation "row-major block dimensions exceed Int cardinality")) + (checkedNonNegativeProduct rowCount columnCount) + if U.length payload /= expectedLength + then + Left + ( InvariantViolation + ( "row-major block payload length mismatch: expected " + <> show expectedLength + <> " but received " + <> show (U.length payload) + ) + ) + else + if U.any (not . fieldValueValid) payload + then Left (InvariantViolation "row-major block entries must be finite") + else + Right + RowMajorBlock + { rowMajorBlockRows = rowCount, + rowMajorBlockColumns = columnCount, + rowMajorBlockPayload = payload + } + +rowMajorBlockEntry :: RowMajorBlock -> Int -> Int -> Double +rowMajorBlockEntry blockValue rowIndex columnIndex = + doubleAt (rowMajorBlockPayload blockValue) (rowIndex * rowMajorBlockColumns blockValue + columnIndex) +{-# INLINE rowMajorBlockEntry #-} + +transposeRowMajorBlock :: RowMajorBlock -> RowMajorBlock +transposeRowMajorBlock blockValue = + RowMajorBlock + { rowMajorBlockRows = rowMajorBlockColumns blockValue, + rowMajorBlockColumns = rowMajorBlockRows blockValue, + rowMajorBlockPayload = + U.generate + (U.length (rowMajorBlockPayload blockValue)) + transposeEntry + } + where + transposeEntry payloadIndex = + let targetColumnCount = rowMajorBlockRows blockValue + rowIndex = payloadIndex `quot` targetColumnCount + columnIndex = payloadIndex `rem` targetColumnCount + in rowMajorBlockEntry blockValue columnIndex rowIndex +{-# INLINE transposeRowMajorBlock #-} + +symmetrizeRowMajorBlockLower :: RowMajorBlock -> Either MoonlightError RowMajorBlock +symmetrizeRowMajorBlockLower blockValue + | rowMajorBlockRows blockValue /= rowMajorBlockColumns blockValue = + Left (InvariantViolation "lower-authoritative symmetrization requires a square block") + | otherwise = + mkRowMajorBlock + blockSize + blockSize + (U.generate (U.length (rowMajorBlockPayload blockValue)) mirroredLowerEntry) + where + blockSize = rowMajorBlockRows blockValue + mirroredLowerEntry payloadIndex = + let rowIndex = payloadIndex `quot` blockSize + columnIndex = payloadIndex `rem` blockSize + in if columnIndex <= rowIndex + then rowMajorBlockEntry blockValue rowIndex columnIndex + else rowMajorBlockEntry blockValue columnIndex rowIndex + +type SymmetricBlockTridiagonal :: Type +data SymmetricBlockTridiagonal = SymmetricBlockTridiagonal + { blockOffsets :: !(U.Vector Int), + diagonalPayloadOffsets :: !(U.Vector Int), + diagonalLowerPacked :: !(U.Vector Double), + couplingPayloadOffsets :: !(U.Vector Int), + lowerCouplingPayload :: !(U.Vector Double) + } + deriving stock (Eq, Show) + +mkSymmetricBlockTridiagonal :: + Box.Vector RowMajorBlock -> + Box.Vector RowMajorBlock -> + Either MoonlightError SymmetricBlockTridiagonal +mkSymmetricBlockTridiagonal diagonalBlocks lowerCouplingBlocks = do + if Box.null diagonalBlocks + then Left (InvariantViolation "symmetric block tridiagonal requires at least one diagonal block") + else Right () + traverse_ validateDiagonalBlock (Box.toList diagonalBlocks) + let blockSizes = rowMajorBlockRows <$> diagonalBlocks + expectedCouplingCount = max 0 (Box.length diagonalBlocks - 1) + if Box.length lowerCouplingBlocks /= expectedCouplingCount + then + Left + ( InvariantViolation + ( "symmetric block tridiagonal coupling count mismatch: expected " + <> show expectedCouplingCount + <> " but received " + <> show (Box.length lowerCouplingBlocks) + ) + ) + else Right () + traverse_ + (validateCouplingBlock blockSizes) + (zip [0 :: Int ..] (Box.toList lowerCouplingBlocks)) + let diagonalPayloads = packLowerBlock <$> diagonalBlocks + couplingPayloads = rowMajorBlockPayload <$> lowerCouplingBlocks + blockOffsetValues <- checkedOffsetsFromSizes "block offsets" (Box.toList blockSizes) + diagonalOffsetValues <- checkedOffsetsFromSizes "diagonal payload offsets" (U.length <$> Box.toList diagonalPayloads) + couplingOffsetValues <- checkedOffsetsFromSizes "coupling payload offsets" (U.length <$> Box.toList couplingPayloads) + Right + SymmetricBlockTridiagonal + { blockOffsets = blockOffsetValues, + diagonalPayloadOffsets = diagonalOffsetValues, + diagonalLowerPacked = U.concat (Box.toList diagonalPayloads), + couplingPayloadOffsets = couplingOffsetValues, + lowerCouplingPayload = U.concat (Box.toList couplingPayloads) + } + +validateDiagonalBlock :: RowMajorBlock -> Either MoonlightError () +validateDiagonalBlock blockValue + | rowMajorBlockRows blockValue /= rowMajorBlockColumns blockValue = + Left (InvariantViolation "symmetric block tridiagonal diagonal blocks must be square") + | otherwise = + if U.and (U.generate (U.length (rowMajorBlockPayload blockValue)) symmetricEntry) + then Right () + else Left (InvariantViolation "symmetric block tridiagonal diagonal block is not exactly symmetric") + where + symmetricEntry payloadIndex = + let blockSize = rowMajorBlockRows blockValue + rowIndex = payloadIndex `quot` blockSize + columnIndex = payloadIndex `rem` blockSize + in rowMajorBlockEntry blockValue rowIndex columnIndex == rowMajorBlockEntry blockValue columnIndex rowIndex + +validateCouplingBlock :: Box.Vector Int -> (Int, RowMajorBlock) -> Either MoonlightError () +validateCouplingBlock blockSizes (couplingIndex, blockValue) = + let expectedRows = intBoxAt blockSizes (couplingIndex + 1) + expectedColumns = intBoxAt blockSizes couplingIndex + in if rowMajorBlockRows blockValue /= expectedRows || rowMajorBlockColumns blockValue /= expectedColumns + then + Left + ( InvariantViolation + ( "symmetric block tridiagonal coupling block " + <> show couplingIndex + <> " shape mismatch: expected " + <> show (expectedRows, expectedColumns) + <> " but received " + <> show (rowMajorBlockRows blockValue, rowMajorBlockColumns blockValue) + ) + ) + else Right () + +packLowerBlock :: RowMajorBlock -> U.Vector Double +packLowerBlock blockValue = + U.concat + ( ( \rowIndex -> + U.generate + (rowIndex + 1) + (\columnIndex -> rowMajorBlockEntry blockValue rowIndex columnIndex) + ) + <$> [0 .. rowMajorBlockRows blockValue - 1] + ) + +checkedOffsetsFromSizes :: String -> [Int] -> Either MoonlightError (U.Vector Int) +checkedOffsetsFromSizes context sizes + | any (< 0) sizes = Left cardinalityFailure + | otherwise = + U.fromList + <$> traverse + (first (const cardinalityFailure) . checkedNaturalToInt) + (scanl (+) 0 (fromIntegral <$> sizes :: [Natural])) + where + cardinalityFailure = + InvariantViolation ("symmetric block tridiagonal " <> context <> " exceed Int cardinality") + +symmetricBlockTridiagonalDimension :: SymmetricBlockTridiagonal -> Int +symmetricBlockTridiagonalDimension blockValue = + intAt (blockOffsets blockValue) (U.length (blockOffsets blockValue) - 1) + +symmetricBlockTridiagonalBlockCount :: SymmetricBlockTridiagonal -> Int +symmetricBlockTridiagonalBlockCount blockValue = + max 0 (U.length (blockOffsets blockValue) - 1) + +symmetricBlockTridiagonalBandwidth :: SymmetricBlockTridiagonal -> Int +symmetricBlockTridiagonalBandwidth blockValue = + let blockCount = symmetricBlockTridiagonalBlockCount blockValue + diagonalWidths = + U.generate blockCount (\blockIndex -> blockSizeAt blockValue blockIndex - 1) + couplingWidths = + U.generate + (max 0 (blockCount - 1)) + (\couplingIndex -> blockSizeAt blockValue couplingIndex + blockSizeAt blockValue (couplingIndex + 1) - 1) + in U.maximum (U.concat [diagonalWidths, couplingWidths]) + +symmetricBlockTridiagonalEntry :: SymmetricBlockTridiagonal -> Int -> Int -> Either MoonlightError Double +symmetricBlockTridiagonalEntry blockValue rowIndex columnIndex + | rowIndex < 0 || rowIndex >= dimension || columnIndex < 0 || columnIndex >= dimension = + Left + ( InvariantViolation + ( "symmetric block tridiagonal entry index out of bounds: " + <> show (rowIndex, columnIndex) + <> " for dimension " + <> show dimension + ) + ) + | otherwise = + case (blockLocalIndex blockValue rowIndex, blockLocalIndex blockValue columnIndex) of + (Just (rowBlockIndex, rowLocalIndex), Just (columnBlockIndex, columnLocalIndex)) -> + Right (entryFromLocal rowBlockIndex rowLocalIndex columnBlockIndex columnLocalIndex) + _ -> + Left (InvariantViolation "symmetric block tridiagonal entry index missing from block map") + where + dimension = symmetricBlockTridiagonalDimension blockValue + + entryFromLocal rowBlockIndex rowLocalIndex columnBlockIndex columnLocalIndex = + case compare rowBlockIndex columnBlockIndex of + EQ -> diagonalEntry blockValue rowBlockIndex rowLocalIndex columnLocalIndex + GT -> + if rowBlockIndex == columnBlockIndex + 1 + then couplingEntry blockValue columnBlockIndex rowLocalIndex columnLocalIndex + else 0.0 + LT -> + if columnBlockIndex == rowBlockIndex + 1 + then couplingEntry blockValue rowBlockIndex columnLocalIndex rowLocalIndex + else 0.0 + +applySymmetricBlockTridiagonalU :: SymmetricBlockTridiagonal -> U.Vector Double -> Either MoonlightError (U.Vector Double) +applySymmetricBlockTridiagonalU blockValue inputVector + | U.length inputVector /= symmetricBlockTridiagonalDimension blockValue = + Left + ( InvariantViolation + ( "symmetric block tridiagonal input dimension mismatch: expected " + <> show (symmetricBlockTridiagonalDimension blockValue) + <> " but received " + <> show (U.length inputVector) + ) + ) + | otherwise = + Right + (runST (applySymmetricBlockTridiagonalST blockValue inputVector)) + +applySymmetricBlockTridiagonalST :: SymmetricBlockTridiagonal -> U.Vector Double -> ST s (U.Vector Double) +applySymmetricBlockTridiagonalST blockValue inputVector = do + let dimension = symmetricBlockTridiagonalDimension blockValue + outputVector <- MU.unsafeNew dimension + U.foldM' + (writeApplyBlock blockValue inputVector outputVector) + () + (U.enumFromN 0 (symmetricBlockTridiagonalBlockCount blockValue)) + U.unsafeFreeze outputVector + +writeApplyBlock :: + SymmetricBlockTridiagonal -> + U.Vector Double -> + MU.MVector s Double -> + () -> + Int -> + ST s () +writeApplyBlock blockValue inputVector outputVector () blockIndex = + U.foldM' + (writeApplyLocalRow blockValue inputVector outputVector blockIndex blockStart) + () + (U.enumFromN 0 (blockSizeAt blockValue blockIndex)) + where + blockStart = intAt (blockOffsets blockValue) blockIndex +{-# INLINE writeApplyBlock #-} + +writeApplyLocalRow :: + SymmetricBlockTridiagonal -> + U.Vector Double -> + MU.MVector s Double -> + Int -> + Int -> + () -> + Int -> + ST s () +writeApplyLocalRow blockValue inputVector outputVector blockIndex blockStart () localRow = + MU.unsafeWrite + outputVector + (blockStart + localRow) + (applyBlockEntry blockValue inputVector blockIndex localRow) +{-# INLINE writeApplyLocalRow #-} + +symmetricBlockTridiagonalUpperBound :: SymmetricBlockTridiagonal -> Double +symmetricBlockTridiagonalUpperBound blockValue = + let dimension = symmetricBlockTridiagonalDimension blockValue + in if dimension <= 0 + then 0.0 + else U.maximum (U.generate dimension (rowAbsSum blockValue)) + +symmetricBlockTridiagonalFrobeniusNorm :: SymmetricBlockTridiagonal -> Double +symmetricBlockTridiagonalFrobeniusNorm blockValue = + sqrt + ( diagonalPackedWeightedSumSquares blockValue + + 2.0 * U.foldl' (\accumulator entryValue -> accumulator + squared entryValue) 0.0 (lowerCouplingPayload blockValue) + ) + +diagonalPackedWeightedSumSquares :: SymmetricBlockTridiagonal -> Double +diagonalPackedWeightedSumSquares blockValue = + U.foldl' + (\accumulator blockIndex -> + accumulator + + sumIndexRange + (blockSizeAt blockValue blockIndex) + (\localRow -> + sumIndexRange + (localRow + 1) + (\localColumn -> + (if localRow == localColumn then 1.0 else 2.0) + * squared (diagonalEntry blockValue blockIndex localRow localColumn) + ) + ) + ) + 0.0 + (U.enumFromN 0 (symmetricBlockTridiagonalBlockCount blockValue)) + +squared :: Double -> Double +squared value = value * value +{-# INLINE squared #-} + +applyBlockEntry :: SymmetricBlockTridiagonal -> U.Vector Double -> Int -> Int -> Double +applyBlockEntry blockValue inputVector blockIndex localRow = + diagonalContribution blockValue inputVector blockIndex localRow + + lowerContribution blockValue inputVector blockIndex localRow + + upperContribution blockValue inputVector blockIndex localRow +{-# INLINE applyBlockEntry #-} + +rowAbsSum :: SymmetricBlockTridiagonal -> Int -> Double +rowAbsSum blockValue rowIndex = + case blockLocalIndex blockValue rowIndex of + Nothing -> 0.0 + Just (blockIndex, localRow) -> + diagonalAbsSum blockValue blockIndex localRow + + lowerAbsSum blockValue blockIndex localRow + + upperAbsSum blockValue blockIndex localRow + +blockLocalIndex :: SymmetricBlockTridiagonal -> Int -> Maybe (Int, Int) +blockLocalIndex blockValue rowIndex = + U.foldl' selectBlock Nothing (U.enumFromN 0 (symmetricBlockTridiagonalBlockCount blockValue)) + where + selectBlock selectedBlock blockIndex = + selectedBlock + <|> let startOffset = intAt (blockOffsets blockValue) blockIndex + stopOffset = intAt (blockOffsets blockValue) (blockIndex + 1) + in if rowIndex >= startOffset && rowIndex < stopOffset + then Just (blockIndex, rowIndex - startOffset) + else Nothing + +blockSizeAt :: SymmetricBlockTridiagonal -> Int -> Int +blockSizeAt blockValue blockIndex = + (intAt (blockOffsets blockValue) (blockIndex + 1)) + - (intAt (blockOffsets blockValue) blockIndex) +{-# INLINE blockSizeAt #-} + +diagonalContribution :: SymmetricBlockTridiagonal -> U.Vector Double -> Int -> Int -> Double +diagonalContribution blockValue inputVector blockIndex localRow = + let blockStart = intAt (blockOffsets blockValue) blockIndex + blockSize = blockSizeAt blockValue blockIndex + in sumIndexRange + blockSize + ( \localColumn -> + diagonalEntry blockValue blockIndex localRow localColumn + * doubleAt inputVector (blockStart + localColumn) + ) +{-# INLINE diagonalContribution #-} + +lowerContribution :: SymmetricBlockTridiagonal -> U.Vector Double -> Int -> Int -> Double +lowerContribution blockValue inputVector blockIndex localRow + | blockIndex <= 0 = 0.0 + | otherwise = + let couplingIndex = blockIndex - 1 + previousStart = intAt (blockOffsets blockValue) couplingIndex + previousSize = blockSizeAt blockValue couplingIndex + in sumIndexRange + previousSize + ( \localColumn -> + couplingEntry blockValue couplingIndex localRow localColumn + * doubleAt inputVector (previousStart + localColumn) + ) +{-# INLINE lowerContribution #-} + +upperContribution :: SymmetricBlockTridiagonal -> U.Vector Double -> Int -> Int -> Double +upperContribution blockValue inputVector blockIndex localRow + | blockIndex + 1 >= symmetricBlockTridiagonalBlockCount blockValue = 0.0 + | otherwise = + let nextStart = intAt (blockOffsets blockValue) (blockIndex + 1) + nextSize = blockSizeAt blockValue (blockIndex + 1) + in sumIndexRange + nextSize + ( \nextLocalRow -> + couplingEntry blockValue blockIndex nextLocalRow localRow + * doubleAt inputVector (nextStart + nextLocalRow) + ) +{-# INLINE upperContribution #-} + +diagonalAbsSum :: SymmetricBlockTridiagonal -> Int -> Int -> Double +diagonalAbsSum blockValue blockIndex localRow = + sumIndexRange + (blockSizeAt blockValue blockIndex) + (\localColumn -> abs (diagonalEntry blockValue blockIndex localRow localColumn)) +{-# INLINE diagonalAbsSum #-} + +lowerAbsSum :: SymmetricBlockTridiagonal -> Int -> Int -> Double +lowerAbsSum blockValue blockIndex localRow + | blockIndex <= 0 = 0.0 + | otherwise = + sumIndexRange + (blockSizeAt blockValue (blockIndex - 1)) + (\localColumn -> abs (couplingEntry blockValue (blockIndex - 1) localRow localColumn)) +{-# INLINE lowerAbsSum #-} + +upperAbsSum :: SymmetricBlockTridiagonal -> Int -> Int -> Double +upperAbsSum blockValue blockIndex localRow + | blockIndex + 1 >= symmetricBlockTridiagonalBlockCount blockValue = 0.0 + | otherwise = + sumIndexRange + (blockSizeAt blockValue (blockIndex + 1)) + (\nextLocalRow -> abs (couplingEntry blockValue blockIndex nextLocalRow localRow)) +{-# INLINE upperAbsSum #-} + +sumIndexRange :: Int -> (Int -> Double) -> Double +sumIndexRange count valueAt = + U.foldl' (\accumulator indexValue -> accumulator + valueAt indexValue) 0.0 (U.enumFromN 0 count) +{-# INLINE sumIndexRange #-} + +diagonalEntry :: SymmetricBlockTridiagonal -> Int -> Int -> Int -> Double +diagonalEntry blockValue blockIndex localRow localColumn + | localColumn <= localRow = + doubleAt (diagonalLowerPacked blockValue) (diagonalPayloadStart blockValue blockIndex + packedLowerIndex localRow localColumn) + | otherwise = + doubleAt (diagonalLowerPacked blockValue) (diagonalPayloadStart blockValue blockIndex + packedLowerIndex localColumn localRow) +{-# INLINE diagonalEntry #-} + +couplingEntry :: SymmetricBlockTridiagonal -> Int -> Int -> Int -> Double +couplingEntry blockValue couplingIndex localRow localColumn = + let couplingStart = intAt (couplingPayloadOffsets blockValue) couplingIndex + couplingColumns = blockSizeAt blockValue couplingIndex + in doubleAt (lowerCouplingPayload blockValue) (couplingStart + localRow * couplingColumns + localColumn) +{-# INLINE couplingEntry #-} + +diagonalPayloadStart :: SymmetricBlockTridiagonal -> Int -> Int +diagonalPayloadStart blockValue blockIndex = + intAt (diagonalPayloadOffsets blockValue) blockIndex +{-# INLINE diagonalPayloadStart #-} + +packedLowerIndex :: Int -> Int -> Int +packedLowerIndex rowIndex columnIndex = + rowIndex * (rowIndex + 1) `quot` 2 + columnIndex +{-# INLINE packedLowerIndex #-} + +intAt :: U.Vector Int -> Int -> Int +intAt values indexValue = + maybe 0 id (values U.!? indexValue) +{-# INLINE intAt #-} + +intBoxAt :: Box.Vector Int -> Int -> Int +intBoxAt values indexValue = + maybe 0 id (values Box.!? indexValue) +{-# INLINE intBoxAt #-} + +doubleAt :: U.Vector Double -> Int -> Double +doubleAt values indexValue = + maybe 0.0 id (values U.!? indexValue) +{-# INLINE doubleAt #-} diff --git a/src-structured/Moonlight/LinAlg/Pure/Structured/Tridiagonal.hs b/src-structured/Moonlight/LinAlg/Pure/Structured/Tridiagonal.hs new file mode 100644 --- /dev/null +++ b/src-structured/Moonlight/LinAlg/Pure/Structured/Tridiagonal.hs @@ -0,0 +1,304 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE StrictData #-} + +module Moonlight.LinAlg.Pure.Structured.Tridiagonal + ( SymmetricTridiagonal, + mkSymmetricTridiagonal, + mkSymmetricTridiagonalVectors, + pathLaplacianBands, + symmetricTridiagonalDimension, + symmetricTridiagonalDiagonalEntries, + symmetricTridiagonalOffDiagonalEntries, + symmetricTridiagonalDiagonalVector, + symmetricTridiagonalOffDiagonalVector, + applyPathLaplacianValidatedU, + applySymmetricTridiagonalU, + applySymmetricTridiagonalValidatedU, + isPathLaplacianTridiagonal, + symmetricTridiagonalUpperBound, + ) +where + +import Control.Monad.ST (runST) +import Data.Kind (Type) +import Data.Primitive (sizeOf) +import Data.Primitive.ByteArray + ( indexByteArray, + newByteArray, + unsafeFreezeByteArray, + writeByteArray, + ) +import qualified Data.Vector.Primitive as P +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as MU +import qualified Data.Vector.Unboxed.Base as UB +import Moonlight.Core (MoonlightError (..), fieldValueValid) +import Prelude + +type SymmetricTridiagonal :: Type +data SymmetricTridiagonal = SymmetricTridiagonal + { symmetricTridiagonalDiagonalVector :: !(U.Vector Double), + symmetricTridiagonalOffDiagonalVector :: !(U.Vector Double) + } + deriving stock (Eq, Show) + +mkSymmetricTridiagonal :: [Double] -> [Double] -> Either MoonlightError SymmetricTridiagonal +mkSymmetricTridiagonal diagonalValues offDiagonalValues = + mkSymmetricTridiagonalVectors + (U.fromList diagonalValues) + (U.fromList offDiagonalValues) + +mkSymmetricTridiagonalVectors :: + U.Vector Double -> + U.Vector Double -> + Either MoonlightError SymmetricTridiagonal +mkSymmetricTridiagonalVectors diagonalValues offDiagonalValues = + let matrixSize = U.length diagonalValues + expectedOffDiagonalCount = max 0 (matrixSize - 1) + in if U.length offDiagonalValues /= expectedOffDiagonalCount + then + Left + ( InvariantViolation + ( "Symmetric tridiagonal off-diagonal length mismatch: expected " + <> show expectedOffDiagonalCount + <> " but received " + <> show (U.length offDiagonalValues) + ) + ) + else + if U.any (not . fieldValueValid) diagonalValues || U.any (not . fieldValueValid) offDiagonalValues + then Left (InvariantViolation "Symmetric tridiagonal entries must be finite") + else + Right + SymmetricTridiagonal + { symmetricTridiagonalDiagonalVector = diagonalValues, + symmetricTridiagonalOffDiagonalVector = offDiagonalValues + } + +pathLaplacianBands :: Int -> Either MoonlightError ([Double], [Double]) +pathLaplacianBands dimension + | dimension < 0 = + Left + ( InvariantViolation + ( "path Laplacian dimension must be non-negative, received " + <> show dimension + ) + ) + | dimension == 0 = Right ([], []) + | dimension == 1 = Right ([0.0], []) + | otherwise = + Right + ( 1.0 : (replicate (dimension - 2) 2.0 <> [1.0]), + replicate (dimension - 1) (-1.0) + ) + +symmetricTridiagonalDimension :: SymmetricTridiagonal -> Int +symmetricTridiagonalDimension = + U.length . symmetricTridiagonalDiagonalVector + +symmetricTridiagonalDiagonalEntries :: SymmetricTridiagonal -> [Double] +symmetricTridiagonalDiagonalEntries = + U.toList . symmetricTridiagonalDiagonalVector + +symmetricTridiagonalOffDiagonalEntries :: SymmetricTridiagonal -> [Double] +symmetricTridiagonalOffDiagonalEntries = + U.toList . symmetricTridiagonalOffDiagonalVector + +applyPathLaplacianValidatedU :: + Int -> + U.Vector Double -> + U.Vector Double +applyPathLaplacianValidatedU dimension inputVector + | dimension <= 0 = U.empty + | dimension == 1 = U.singleton 0.0 + | otherwise = + U.create $ do + targetVector <- MU.unsafeNew dimension + let !firstValue = inputVector `U.unsafeIndex` 0 + !secondValue = inputVector `U.unsafeIndex` 1 + MU.unsafeWrite targetVector 0 (firstValue - secondValue) + + let writeInterior !rowIndex + | rowIndex + 1 >= dimension = pure () + | otherwise = do + let !leftValue = inputVector `U.unsafeIndex` (rowIndex - 1) + !centerValue = inputVector `U.unsafeIndex` rowIndex + !rightValue = inputVector `U.unsafeIndex` (rowIndex + 1) + MU.unsafeWrite + targetVector + rowIndex + (2.0 * centerValue - leftValue - rightValue) + writeInterior (rowIndex + 1) + + writeInterior 1 + + let !lastIndex = dimension - 1 + !lastValue = inputVector `U.unsafeIndex` lastIndex + !penultimateValue = inputVector `U.unsafeIndex` (lastIndex - 1) + MU.unsafeWrite targetVector lastIndex (lastValue - penultimateValue) + pure targetVector +{-# INLINE applyPathLaplacianValidatedU #-} + +applySymmetricTridiagonalU :: + SymmetricTridiagonal -> + U.Vector Double -> + U.Vector Double +applySymmetricTridiagonalU = applySymmetricTridiagonalValidatedU +{-# INLINE applySymmetricTridiagonalU #-} + +applySymmetricTridiagonalValidatedU :: + SymmetricTridiagonal -> + U.Vector Double -> + U.Vector Double +applySymmetricTridiagonalValidatedU + ( SymmetricTridiagonal + (UB.V_Double (P.Vector diagonalBase matrixSize diagonalArray)) + (UB.V_Double (P.Vector offDiagonalBase _ offDiagonalArray)) + ) + (UB.V_Double (P.Vector inputBase _ inputArray)) + | matrixSize <= 0 = U.empty + | matrixSize == 1 = + U.singleton + ( (indexByteArray diagonalArray diagonalBase :: Double) + * (indexByteArray inputArray inputBase :: Double) + ) + | otherwise = + UB.V_Double + ( P.Vector + 0 + matrixSize + ( runST $ do + targetArray <- + newByteArray + (matrixSize * sizeOf (0.0 :: Double)) + + let !firstInput = + ( indexByteArray inputArray inputBase + :: Double + ) + !secondInput = + ( indexByteArray inputArray (inputBase + 1) + :: Double + ) + !firstValue = + ( indexByteArray diagonalArray diagonalBase + :: Double + ) + * firstInput + + ( indexByteArray + offDiagonalArray + offDiagonalBase + :: Double + ) + * secondInput + writeByteArray targetArray 0 firstValue + + let writeRows !rowIndex !previousInput !currentInput + | rowIndex + 1 >= matrixSize = do + let !lastValue = + ( indexByteArray + offDiagonalArray + (offDiagonalBase + rowIndex - 1) + :: Double + ) + * previousInput + + ( indexByteArray + diagonalArray + (diagonalBase + rowIndex) + :: Double + ) + * currentInput + writeByteArray + targetArray + rowIndex + lastValue + unsafeFreezeByteArray targetArray + | otherwise = do + let !nextInput = + ( indexByteArray + inputArray + (inputBase + rowIndex + 1) + :: Double + ) + !rowValue = + ( indexByteArray + offDiagonalArray + (offDiagonalBase + rowIndex - 1) + :: Double + ) + * previousInput + + ( indexByteArray + diagonalArray + (diagonalBase + rowIndex) + :: Double + ) + * currentInput + + ( indexByteArray + offDiagonalArray + (offDiagonalBase + rowIndex) + :: Double + ) + * nextInput + writeByteArray + targetArray + rowIndex + rowValue + writeRows + (rowIndex + 1) + currentInput + nextInput + + writeRows 1 firstInput secondInput + ) + ) +{-# INLINE applySymmetricTridiagonalValidatedU #-} + +isPathLaplacianTridiagonal :: SymmetricTridiagonal -> Bool +isPathLaplacianTridiagonal tridiagonalValue = + diagonalLoop 0 && U.all (== (-1.0)) offDiagonalValues + where + !diagonalValues = symmetricTridiagonalDiagonalVector tridiagonalValue + !offDiagonalValues = symmetricTridiagonalOffDiagonalVector tridiagonalValue + !matrixSize = U.length diagonalValues + + expectedDiagonal !rowIndex + | matrixSize == 1 = 0.0 + | rowIndex == 0 || rowIndex + 1 == matrixSize = 1.0 + | otherwise = 2.0 + + diagonalLoop !rowIndex + | rowIndex >= matrixSize = True + | diagonalValues `U.unsafeIndex` rowIndex == expectedDiagonal rowIndex = + diagonalLoop (rowIndex + 1) + | otherwise = False +{-# INLINE isPathLaplacianTridiagonal #-} + +symmetricTridiagonalUpperBound :: SymmetricTridiagonal -> Double +symmetricTridiagonalUpperBound tridiagonalValue = + let matrixSize = symmetricTridiagonalDimension tridiagonalValue + in if matrixSize <= 0 + then 0.0 + else U.maximum (U.generate matrixSize (rowUpperBound tridiagonalValue)) + +rowUpperBound :: SymmetricTridiagonal -> Int -> Double +rowUpperBound tridiagonalValue rowIndex = + let diagonalValues = symmetricTridiagonalDiagonalVector tridiagonalValue + offDiagonalValues = symmetricTridiagonalOffDiagonalVector tridiagonalValue + matrixSize = U.length diagonalValues + leftRadius = + if rowIndex <= 0 + then 0.0 + else abs (vectorEntryOrZero offDiagonalValues (rowIndex - 1)) + rightRadius = + if rowIndex + 1 >= matrixSize + then 0.0 + else abs (vectorEntryOrZero offDiagonalValues rowIndex) + in vectorEntryOrZero diagonalValues rowIndex + leftRadius + rightRadius +{-# INLINE rowUpperBound #-} + +vectorEntryOrZero :: U.Vector Double -> Int -> Double +vectorEntryOrZero values indexValue = + case values U.!? indexValue of + Nothing -> 0.0 + Just value -> value +{-# INLINE vectorEntryOrZero #-} diff --git a/test-laws/Main.hs b/test-laws/Main.hs new file mode 100644 --- /dev/null +++ b/test-laws/Main.hs @@ -0,0 +1,8 @@ +module Main (main) where + +import Moonlight.LinAlg.Effect.Laws qualified +import Test.Tasty (defaultMain) + +main :: IO () +main = + defaultMain Moonlight.LinAlg.Effect.Laws.tests diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,45 @@ +module Main (main) where + +import qualified AdvancedSpec as AdvancedSpec +import qualified ArchitectureSpec as ArchitectureSpec +import qualified BasicSpec as BasicSpec +import qualified BlockSpec as BlockSpec +import qualified DenseFlatSpec as DenseFlatSpec +import qualified DenseRowsSpec as DenseRowsSpec +import qualified KrylovSpec as KrylovSpec +import qualified DomainSpec as DomainSpec +import qualified DynamicSpec as DynamicSpec +import qualified FieldSpec as FieldSpec +import qualified ExteriorSpec as ExteriorSpec +import qualified GF2Spec as GF2Spec +import qualified GeometryStorageSpec as GeometryStorageSpec +import qualified SymmetricSpec as SymmetricSpec +import qualified StaticsSpec as StaticsSpec +import qualified SparseSolverSpec as SparseSolverSpec +import qualified SparsePackedSpec as SparsePackedSpec +import Test.Tasty (defaultMain, testGroup) + +main :: IO () +main = + defaultMain + ( testGroup + "moonlight-linalg" + [ ArchitectureSpec.tests, + BasicSpec.tests, + BlockSpec.tests, + DenseFlatSpec.tests, + DenseRowsSpec.tests, + FieldSpec.tests, + DomainSpec.tests, + DynamicSpec.tests, + ExteriorSpec.tests, + GF2Spec.tests, + GeometryStorageSpec.tests, + SymmetricSpec.tests, + AdvancedSpec.tests, + KrylovSpec.tests, + SparsePackedSpec.tests, + SparseSolverSpec.tests, + StaticsSpec.tests + ] + ) diff --git a/test/architecture/ArchitectureSpec.hs b/test/architecture/ArchitectureSpec.hs new file mode 100644 --- /dev/null +++ b/test/architecture/ArchitectureSpec.hs @@ -0,0 +1,687 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE StrictData #-} + +module ArchitectureSpec + ( tests, + ) +where + +import Control.Applicative ((<|>)) +import Data.List (isInfixOf, isSuffixOf, sort) +import Data.Maybe (catMaybes) +import System.Directory (doesDirectoryExist, doesFileExist, listDirectory) +import System.FilePath ((), takeExtension) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase) +import Prelude + +data SourceCheck = SourceCheck + { sourcePath :: !FilePath, + forbiddenFragments :: ![(String, String)] + } + +data MissingSource = MissingSource + { missingPath :: !FilePath, + missingLabel :: !String + } + +data SliceSourceRoot = SliceSourceRoot + { sliceSourceRoot :: !FilePath, + sliceForbiddenFragments :: ![(String, String)] + } + +tests :: TestTree +tests = + testGroup + "architecture" + [ testCase "obsolete linalg owners stay deleted" assertObsoleteOwnersStayDeleted, + testCase "public Krylov surface stays algorithm-only" assertPublicKrylovSurface, + testCase "public operator surface keeps source constructors hidden" assertPublicOperatorSurface, + testCase "Krylov hot coefficients stay vector-native" assertKrylovHotCoefficients, + testCase "selected spectral values stay vector-native" assertSelectedSpectralValuesVectorNative, + testCase "flat dense Double storage stays storable-native" assertDenseDoubleStorageStorable, + testCase "spectral fallback seed stays vector-native" assertSpectralSeedBoundary, + testCase "structured projected/block routes avoid dense row owners" assertStructuredRoutesAvoidDenseRowOwners, + testCase "native LAPACK details stay behind the native effect boundary" assertNativeLapackBoundary, + testCase "linalg sublibrary slices stay downward" assertSublibrarySliceDiscipline, + testCase "canonical sparse storage exposes vector payloads only" assertCanonicalSparseStorageSurface, + testCase "canonical CSR matvec does not revalidate in hot apply paths" assertCSRHotApplyBoundary, + testCase "public sparse solver surface stays canonical" assertPublicSparseSolverSurface, + testCase "sparse solvers stay off list orchestration" assertSparseSolverHotSurface, + testCase "linalg docs do not advertise deleted spectral APIs" assertLinalgDocsAvoidDeletedSpectralAPIs + ] + +assertObsoleteOwnersStayDeleted :: Assertion +assertObsoleteOwnersStayDeleted = do + existingSources <- catMaybes <$> traverse existingMissingSource obsoleteSources + assertBool + ("obsolete linalg owner files still exist:\n" <> unlines existingSources) + (null existingSources) + +obsoleteSources :: [MissingSource] +obsoleteSources = + [ MissingSource "src-carrier/Moonlight/LinAlg/Internal/Continuous.hs" "continuous dense-list helper", + MissingSource "src-dense/Moonlight/LinAlg/Pure/Dense/Classes.hs" "duplicate dense class owner", + MissingSource "src-dense/Moonlight/LinAlg/Pure/Dense/Primitives.hs" "duplicate dense primitive owner", + MissingSource "src-dense/Moonlight/LinAlg/Pure/Dense/VectorOps.hs" "duplicate dense vector ops", + MissingSource "src-spectral/Moonlight/LinAlg/Pure/Krylov.hs" "old core Krylov barrel", + MissingSource "src-spectral/Moonlight/LinAlg/Pure/Krylov/Restart.hs" "obsolete restarted solve front door", + MissingSource "src-spectral/Moonlight/LinAlg/Pure/Krylov/Structure.hs" "duplicate projected structure owner", + MissingSource "src-spectral/Moonlight/LinAlg/Pure/Krylov/TridiagonalSolve.hs" "tridiagonal wrapper solve owner", + MissingSource "src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver.hs" "old list sparse solver facade" + ] + +assertPublicKrylovSurface :: Assertion +assertPublicKrylovSurface = + assertForbiddenFragments + [ SourceCheck + "src-public/Moonlight/LinAlg/Krylov.hs" + [ ("public operator source tag", "LinearOperatorSource"), + ("public operator source accessor", "linearOperatorSource"), + ("public Ritz pair result", "RitzPair"), + ("public Ritz vector conversion", "ritzVectorDyn"), + ("public projected solve policy", "ProjectedSolvePolicy"), + ("public projected solve backend", "ProjectedSolveBackend"), + ("public projected policy runner", "projectedEigenpairsWithPolicy"), + ("public projected policy default", "defaultProjectedSolvePolicy"), + ("public projected backend accessor", "projectedSolveBackend"), + ("public projected operator internals", "SymmetricProjectedOperator"), + ("public projected subspace internals", "ProjectedSubspace"), + ("obsolete restarted solve front door", "restartedEigenpairsSymmetric"), + ("obsolete block solve front door", "blockEigenpairsSymmetric") + ], + SourceCheck + "src-public/Moonlight/LinAlg/Spectral.hs" + [ ("list fallback setter", "withEigenFallbackInitialList"), + ("public list eigenvalue projection", "eigenvaluesToList"), + ("public list eigenpair projection", "eigenpairsToListColumns"), + ("obsolete selected eigenvalue alias", "selectedEigenvalues"), + ("obsolete selected eigenpair alias", "selectedEigenpairs") + ] + ] + +assertPublicOperatorSurface :: Assertion +assertPublicOperatorSurface = + assertForbiddenFragments + [ SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Operator.hs" + [ ("public linear operator constructor export", "LinearOperator (..)"), + ("public operator source type leak", "OperatorSource"), + ("public operator source field leak", "operatorSource"), + ("public source interpreter leak", "applyOperatorSource"), + ("public source shape leak", "operatorSourceShape"), + ("public list-valued operator constructor", "mkLinearOperator"), + ("public list-valued self-adjoint constructor", "declaredSelfAdjointLinearOperator"), + ("public list-valued operator application", "applyLinearOperator") + ], + SourceCheck + "src-public/Moonlight/LinAlg/Operator.hs" + [ ("public linear operator constructor export", "LinearOperator (..)"), + ("public operator source type leak", "OperatorSource"), + ("public operator source field leak", "operatorSource"), + ("public source interpreter leak", "applyOperatorSource"), + ("public source shape leak", "operatorSourceShape"), + ("public list-valued operator constructor", "mkLinearOperator"), + ("public list-valued self-adjoint constructor", "declaredSelfAdjointLinearOperator"), + ("public list-valued operator application", "applyLinearOperator") + ], + SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Operator/Internal.hs" + [ ("list-valued operator constructor", "mkLinearOperator"), + ("list-valued self-adjoint constructor", "declaredSelfAdjointLinearOperator"), + ("list-valued operator application", "applyLinearOperator") + ] + ] + +assertKrylovHotCoefficients :: Assertion +assertKrylovHotCoefficients = + assertForbiddenFragments + [ SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Krylov/Internal.hs" + [ ("list-valued orthogonalization coefficients", "Either MoonlightError (U.Vector Double, [Double])"), + ("seed list normalizer signature", "normalizeSeed :: String -> Int -> Double -> [Double]"), + ("seed list materialization", "U.fromList seedValues"), + ("basis list traversal", "Box.toList"), + ("coefficient list zipper", "zipWithExact"), + ("restart seed residue", "nextRestartSeed") + ], + SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Krylov/Arnoldi.hs" + [ ("Arnoldi list seed signature", "-> [Double] ->"), + ("Arnoldi coefficient list materialization", "U.fromList (coefficients"), + ("Arnoldi column list accumulator", "columnsRev") + ], + SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Krylov/Lanczos.hs" + [("Lanczos list seed signature", "-> [Double] ->")] + ] + +assertSelectedSpectralValuesVectorNative :: Assertion +assertSelectedSpectralValuesVectorNative = + assertForbiddenFragments + [ SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Krylov/SelectedTridiagonal.hs" + [ ("list-valued selected tridiagonal direct values", "selectedSymmetricTridiagonalEigenvaluesDirect ::\n SpectrumEnd ->\n Int ->\n SymmetricTridiagonal ->\n Either MoonlightError [Double]"), + ("list-valued selected tridiagonal checked values", "selectedSymmetricTridiagonalEigenvaluesChecked ::\n SpectrumEnd ->\n Int ->\n SymmetricTridiagonal ->\n Either MoonlightError [Double]"), + ("CSR row-offset list roundtrip", "U.fromList (csrRowOffsets csrValue)"), + ("CSR column-index list roundtrip", "U.fromList (csrColumnIndices csrValue)"), + ("CSR value list roundtrip", "U.fromList (csrValues csrValue)"), + ("path tridiagonal pair sort before dispatch", "fmap (sortForSpectrumBy spectrumEnd (\\(eigenvalue, _, _) -> eigenvalue) . take requestedCount) $\n case pathLaplacianEigenpairs"), + ("path tridiagonal largest pair ascending modes", "LargestEigenvalues -> [matrixSize - boundedCount .. matrixSize - 1]"), + ("path tridiagonal pair residual via generic tridiagonal apply", "!residualNorm = tridiagonalResidualNorm tridiagonalValue eigenvalue eigenvector"), + ("path tridiagonal pair residual vector allocation", "normU (U.generate matrixSize (pathLaplacianResidualEntry matrixSize eigenvalue eigenvector))"), + ("reducible values through QL pair kernel", "traverse blockEigenvaluesViaQL (tridiagonalBlocks tridiagonalValue)"), + ("values-only selected tridiagonal QL detour", "selectedTridiagonalEigenvaluesViaQL"), + ("small values-only QL threshold", "smallTridiagonalSelectedQLThreshold") + ], + SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Spectral/Solve.hs" + [ ("selected value call-site list materialization", "U.fromList <$> selectedSymmetricTridiagonalEigenvaluesDirect"), + ("path value list materialization", "U.fromList\n . sortForSpectrumBy spectrumEnd id\n . fmap (pathLaplacianEigenvalueAt dimension)"), + ("path pair operator reconstruction", "pathLaplacianLinearOperator dimension"), + ("path pair residual via full operator apply", "runOperatorU operatorValue eigenvector"), + ("path pair residual vector allocation", "normU (U.generate dimension (pathLaplacianResidualEntry dimension eigenvalue eigenvector))"), + ("path pair resorting after ordered mode selection", "sortBy\n (spectrumPairOrdering spectrumEnd)\n (pathLaplacianColumn dimension <$> selectedModeIndices spectrumEnd requestedCount dimension)"), + ("diagonal values unconditional full sort", "Right . U.fromList . take requestedCount . fmap snd . sortIndexedValues spectrumEnd . U.toList $ U.indexed diagonalEntries") + ], + SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Spectral/Result.hs" + [ ("list eigenvalue projection", "eigenvaluesToList"), + ("list eigenpair projection", "eigenpairsToListColumns"), + ("list eigenpair column projection", "eigenpairsToColumns") + ], + SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Krylov/Projected.hs" + [ ("projected value call-site list materialization", "U.fromList <$> symmetricProjectedEigenvalues"), + ("projected selected value call-site list materialization", "U.fromList <$> selectedSymmetricTridiagonalEigenvaluesDirect"), + ("projected list-valued pair column request", "selectedSymmetricTridiagonalEigenpairColumnsDirect"), + ("projected raw pair list intermediate", "rawPairs <-"), + ("projected lifted column list intermediate", "liftedColumns <- traverse") + ] + , + SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Spectral/Request.hs" + [ ("obsolete selected eigenvalue alias", "selectedEigenvalues"), + ("obsolete selected eigenpair alias", "selectedEigenpairs") + ] + ] + +assertDenseDoubleStorageStorable :: Assertion +assertDenseDoubleStorageStorable = + assertForbiddenFragments + [ SourceCheck + "src-carrier/Moonlight/LinAlg/Pure/Dense/Flat.hs" + [ ("flat dense Double unboxed storage import", "Data.Vector.Unboxed"), + ("flat dense Double unboxed payload", "U.Vector Double") + ] + ] + +assertSpectralSeedBoundary :: Assertion +assertSpectralSeedBoundary = + assertForbiddenFragments + [ SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Spectral/Solve.hs" + [ ("list fallback setter", "withEigenFallbackInitialList"), + ("fallback seed list conversion", "U.toList (fallbackSeed") + ], + SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Krylov/Projected.hs" + [("projected Lanczos seed list conversion", "U.toList seedVector")], + SourceCheck + "bench/spectral/ProjectedBlock.hs" + [("projected benchmark list seed", "projectedSeedVector :: Int -> [Double]")], + SourceCheck + "bench/spectral/SpectralDispatch.hs" + [ ("spectral benchmark list seed", "seedVector :: Int -> [Double]"), + ("spectral benchmark list fallback setter", "withEigenFallbackInitialList") + ] + ] + +assertStructuredRoutesAvoidDenseRowOwners :: Assertion +assertStructuredRoutesAvoidDenseRowOwners = + assertForbiddenFragments + [ SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Krylov/Projected.hs" + [ ("projected dense row owner", "symmetricProjectedOperatorRows"), + ("projected dense row payload", "[[Double]]"), + ("fake block-projected values Lanczos fallback", "blockProjectedEigenvaluesViaLanczos"), + ("fake block-projected pairs Lanczos fallback", "blockProjectedEigenpairsViaLanczos"), + ("projected block fallback operator wrapper", "declaredSelfAdjointVectorLinearOperator") + ], + SourceCheck + "src-structured/Moonlight/LinAlg/Pure/Structured/BlockTridiagonal.hs" + [ ("block tridiagonal dense rows", "blockTridiagonalRows"), + ("nested dense block row payload", "[[[Double]]]"), + ("derived row-to-block map owner", "blockRowIndices"), + ("derived row-to-block map builder", "rowsToBlockIndices"), + ("block apply generated output vector", "Right\n ( U.generate"), + ("block apply generated row sum", "U.sum\n ( U.generate"), + ("block apply generated absolute row sum", "U.sum\n ( U.generate") + ], + SourceCheck + "src-structured/Moonlight/LinAlg/Pure/Structured/Tridiagonal.hs" + [ ("tridiagonal dense row view", "symmetricTridiagonalRows"), + ("tridiagonal dense row payload", "[[Double]]") + ], + SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Krylov/Decomposition.hs" + [("stored projected dimension witness", "ProjectedDimension")], + SourceCheck + "src-native/Moonlight/LinAlg/Native.hs" + [("native-owned block band payload builder", "symmetricBlockTridiagonalLowerBandPayload ::")] + ] + +assertNativeLapackBoundary :: Assertion +assertNativeLapackBoundary = + assertForbiddenFragments + [ SourceCheck + "src-native/Moonlight/LinAlg/Native.hs" + [ ("public native Fortran range type", "FortranIndexRange"), + ("public native Fortran range constructor", "mkFortranIndexRange"), + ("public native raw symmetric-band value driver", "selectedSymmetricBandEigenValuesLapack"), + ("public native raw symmetric-band pair driver", "selectedSymmetricBandEigenPairsLapack"), + ("public native block band payload import", "symmetricBlockTridiagonalLowerBandPayload") + ], + SourceCheck + "src-structured/Moonlight/LinAlg/Pure/Structured/BlockTridiagonal.hs" + [ ("structured carrier native band payload", "symmetricBlockTridiagonalLowerBandPayload") ] + ] + +assertCSRHotApplyBoundary :: Assertion +assertCSRHotApplyBoundary = + assertForbiddenFragments + [ SourceCheck + "src-spectral/Moonlight/LinAlg/Pure/Operator/Internal.hs" + [ ("canonical CSR operator revalidating kernel import", "import Moonlight.LinAlg.Internal.VectorOps (csrMatVecU)"), + ("canonical CSR operator revalidating kernel call", "csrMatVecU ") + ], + SourceCheck + "src-sparse/Moonlight/LinAlg/Pure/Sparse/Types.hs" + [ ("canonical CSR public matvec revalidating kernel import", "import Moonlight.LinAlg.Internal.VectorOps (csrMatVecU)"), + ("canonical CSR public matvec revalidating kernel call", "csrMatVecU ") + ] + ] + +assertCanonicalSparseStorageSurface :: Assertion +assertCanonicalSparseStorageSurface = + assertForbiddenFragments + [ SourceCheck + "src-sparse/Moonlight/LinAlg/Pure/Sparse/Types.hs" + [ ("CSR row-offset list accessor", "csrRowOffsets ::"), + ("CSR column-index list accessor", "csrColumnIndices ::"), + ("CSR value list accessor", "csrValues ::"), + ("CSC column-offset list accessor", "cscColumnOffsets ::"), + ("CSC row-index list accessor", "cscRowIndices ::"), + ("CSC value list accessor", "cscValues ::"), + ("CSR list matvec wrapper", "csrMatVec ::") + ], + SourceCheck + "src-public/Moonlight/LinAlg/Sparse.hs" + [ ("public CSR row-offset list accessor", "csrRowOffsets,"), + ("public CSR column-index list accessor", "csrColumnIndices,"), + ("public CSR value list accessor", "csrValues,"), + ("public CSC column-offset list accessor", "cscColumnOffsets,"), + ("public CSC row-index list accessor", "cscRowIndices,"), + ("public CSC value list accessor", "cscValues,"), + ("public CSR list matvec wrapper", "csrMatVec,") + ] + ] + +assertSparseSolverHotSurface :: Assertion +assertSparseSolverHotSurface = + assertForbiddenFragments + ( sparseSolverCheck + <$> [ "src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/CG.hs", + "src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/GMRES.hs", + "src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Stationary.hs", + "src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Preconditioner.hs", + "src-sparse/Moonlight/LinAlg/Pure/Sparse/Solver/Mutable.hs" + ] + ) + +sparseSolverCheck :: FilePath -> SourceCheck +sparseSolverCheck pathValue = + SourceCheck + pathValue + [ ("list vector payload", "[Double]"), + ("list drop in solver hot path", "drop "), + ("list replacement helper", "replaceAt"), + ("list concatenation in solver hot path", "++"), + ("unboxed vector to-list conversion", "U.toList"), + ("unboxed vector from-list conversion", "U.fromList"), + ("boxed vector to-list conversion", "Box.toList") + ] + +assertPublicSparseSolverSurface :: Assertion +assertPublicSparseSolverSurface = + assertForbiddenFragments + [ SourceCheck + "src-public/Moonlight/LinAlg/Sparse.hs" + [ ("public compiled sparse preconditioner type", "SparsePreconditioner,"), + ("public sparse preconditioner applier", "applySparsePreconditioner"), + ("public sparse preconditioner compiler", "compileSparsePreconditioner"), + ("public direct diagonal preconditioner compiler", "diagonalPreconditioner"), + ("public direct SSOR preconditioner compiler", "ssorPreconditioner"), + ("public direct shifted preconditioner compiler", "shiftedDiagonalPreconditioner"), + ("public PCG entry point", "solveSparsePCG"), + ("public GMRES family wrapper", "solveSparseGMRESWithFamily") + ] + ] + +assertLinalgDocsAvoidDeletedSpectralAPIs :: Assertion +assertLinalgDocsAvoidDeletedSpectralAPIs = + assertForbiddenFragments + [ SourceCheck + "README.md" + deletedSpectralDocFragments, + SourceCheck + "docs/CONSTRUCTION.md" + deletedSpectralDocFragments, + SourceCheck + "CHANGELOG.md" + deletedSpectralDocFragments + ] + +deletedSpectralDocFragments :: [(String, String)] +deletedSpectralDocFragments = + [ ("deleted list-valued self-adjoint constructor", "declaredSelfAdjointLinearOperator"), + ("deleted list-valued fallback seed setter", "withEigenFallbackInitialList"), + ("deleted public operator source tag", "LinearOperatorSource"), + ("deleted Ritz pair result", "RitzPair"), + ("deleted projected solve policy", "ProjectedSolvePolicy"), + ("deleted path tridiagonal helper name", "pathLaplacianTridiagonal") + ] + +assertSublibrarySliceDiscipline :: Assertion +assertSublibrarySliceDiscipline = do + discoveredSlices <- traverse discoverSliceSources sliceSourceRoots + let discoveredSources = concatMap snd discoveredSlices + missingNestedDomainSources = + filter + (\requiredSuffix -> not (any (requiredSuffix `isSuffixOf`) discoveredSources)) + [ "src-domain/Moonlight/LinAlg/Pure/Domain/Smith/Multimodular.hs", + "src-domain/Moonlight/LinAlg/Pure/Domain/Smith/Witnessed.hs" + ] + sourceChecks = + concatMap + (\(sourceRoot, sourcePaths) -> sourceCheckFor (sliceForbiddenFragments sourceRoot) <$> sourcePaths) + discoveredSlices + assertBool + ("recursive slice discovery missed nested domain modules: " <> show missingNestedDomainSources) + (null missingNestedDomainSources) + assertForbiddenFragments sourceChecks + +discoverSliceSources :: SliceSourceRoot -> IO (SliceSourceRoot, [FilePath]) +discoverSliceSources sourceRoot = do + resolvedRoot <- resolveSourceDirectory (sliceSourceRoot sourceRoot) + sourcePaths <- + case resolvedRoot of + Nothing -> + assertFailure ("architecture source root is not reachable from test cwd: " <> sliceSourceRoot sourceRoot) + *> pure [] + Just rootPath -> discoverHaskellSources rootPath + assertBool + ("architecture source root contains no Haskell modules: " <> sliceSourceRoot sourceRoot) + (not (null sourcePaths)) + pure (sourceRoot, sourcePaths) + +discoverHaskellSources :: FilePath -> IO [FilePath] +discoverHaskellSources rootPath = do + childNames <- sort <$> listDirectory rootPath + concat + <$> traverse + (\childName -> + let childPath = rootPath childName + in doesDirectoryExist childPath >>= \case + True -> discoverHaskellSources childPath + False -> pure [childPath | takeExtension childPath == ".hs"] + ) + childNames + +sourceCheckFor :: [(String, String)] -> FilePath -> SourceCheck +sourceCheckFor fragments pathValue = + SourceCheck pathValue fragments + +noInPackageImports :: [(String, String)] +noInPackageImports = + forbiddenSliceImports + [ ("in-package", "Moonlight.LinAlg.") + ] + +carrierForbiddenImports :: [(String, String)] +carrierForbiddenImports = + forbiddenSliceImports + [ ("dense slice", "Moonlight.LinAlg.Internal.Backend"), + ("dense slice", "Moonlight.LinAlg.Internal.Dense."), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Basic"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Block"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Decomposition"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Dynamic"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Exterior"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Field"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.GF2"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Solver"), + ("domain slice", "Moonlight.LinAlg.Pure.Domain"), + ("eigen slice", "Moonlight.LinAlg.Internal.Eigen"), + ("geometry slice", "Moonlight.LinAlg.Pure.Geometry"), + ("native slice", "Moonlight.LinAlg.Effect.Native"), + ("sparse slice", "Moonlight.LinAlg.Pure.Sparse"), + ("spectral Krylov slice", "Moonlight.LinAlg.Pure.Krylov"), + ("spectral operator slice", "Moonlight.LinAlg.Pure.Operator"), + ("spectral solve slice", "Moonlight.LinAlg.Pure.Spectral"), + ("statics slice", "Moonlight.LinAlg.Pure.Statics"), + ("structured slice", "Moonlight.LinAlg.Pure.Structured") + ] + +eigenForbiddenImports :: [(String, String)] +eigenForbiddenImports = + forbiddenSliceImports + [ ("dense slice", "Moonlight.LinAlg.Internal.Backend"), + ("dense slice", "Moonlight.LinAlg.Internal.Dense."), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Basic"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Block"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Decomposition"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Dynamic"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Exterior"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Field"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.GF2"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Solver"), + ("domain slice", "Moonlight.LinAlg.Pure.Domain"), + ("geometry slice", "Moonlight.LinAlg.Pure.Geometry"), + ("native slice", "Moonlight.LinAlg.Effect.Native"), + ("sparse slice", "Moonlight.LinAlg.Pure.Sparse"), + ("spectral Krylov slice", "Moonlight.LinAlg.Pure.Krylov"), + ("spectral operator slice", "Moonlight.LinAlg.Pure.Operator"), + ("spectral solve slice", "Moonlight.LinAlg.Pure.Spectral"), + ("statics slice", "Moonlight.LinAlg.Pure.Statics"), + ("structured slice", "Moonlight.LinAlg.Pure.Structured") + ] + +geometryForbiddenImports :: [(String, String)] +geometryForbiddenImports = + forbiddenSliceImports + [ ("dense slice", "Moonlight.LinAlg.Internal.Backend"), + ("dense slice", "Moonlight.LinAlg.Internal.Dense."), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Basic"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Block"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Decomposition"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Dynamic"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Exterior"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Field"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.GF2"), + ("dense slice", "Moonlight.LinAlg.Pure.Dense.Solver"), + ("domain slice", "Moonlight.LinAlg.Pure.Domain"), + ("eigen slice", "Moonlight.LinAlg.Internal.Eigen"), + ("native slice", "Moonlight.LinAlg.Effect.Native"), + ("sparse slice", "Moonlight.LinAlg.Pure.Sparse"), + ("spectral Krylov slice", "Moonlight.LinAlg.Pure.Krylov"), + ("spectral operator slice", "Moonlight.LinAlg.Pure.Operator"), + ("spectral solve slice", "Moonlight.LinAlg.Pure.Spectral"), + ("statics slice", "Moonlight.LinAlg.Pure.Statics"), + ("structured slice", "Moonlight.LinAlg.Pure.Structured") + ] + +denseForbiddenImports :: [(String, String)] +denseForbiddenImports = + forbiddenSliceImports + [ ("domain slice", "Moonlight.LinAlg.Pure.Domain"), + ("geometry slice", "Moonlight.LinAlg.Pure.Geometry"), + ("native slice", "Moonlight.LinAlg.Effect.Native"), + ("sparse slice", "Moonlight.LinAlg.Pure.Sparse"), + ("spectral Krylov slice", "Moonlight.LinAlg.Pure.Krylov"), + ("spectral operator slice", "Moonlight.LinAlg.Pure.Operator"), + ("spectral solve slice", "Moonlight.LinAlg.Pure.Spectral"), + ("statics slice", "Moonlight.LinAlg.Pure.Statics"), + ("structured slice", "Moonlight.LinAlg.Pure.Structured") + ] + +domainForbiddenImports :: [(String, String)] +domainForbiddenImports = + forbiddenSliceImports + [ ("eigen slice", "Moonlight.LinAlg.Internal.Eigen"), + ("geometry slice", "Moonlight.LinAlg.Pure.Geometry"), + ("native slice", "Moonlight.LinAlg.Effect.Native"), + ("sparse slice", "Moonlight.LinAlg.Pure.Sparse"), + ("spectral Krylov slice", "Moonlight.LinAlg.Pure.Krylov"), + ("spectral operator slice", "Moonlight.LinAlg.Pure.Operator"), + ("spectral solve slice", "Moonlight.LinAlg.Pure.Spectral"), + ("statics slice", "Moonlight.LinAlg.Pure.Statics"), + ("structured slice", "Moonlight.LinAlg.Pure.Structured") + ] + +sparseForbiddenImports :: [(String, String)] +sparseForbiddenImports = + forbiddenSliceImports + [ ("domain slice", "Moonlight.LinAlg.Pure.Domain"), + ("eigen slice", "Moonlight.LinAlg.Internal.Eigen"), + ("geometry slice", "Moonlight.LinAlg.Pure.Geometry"), + ("native slice", "Moonlight.LinAlg.Effect.Native"), + ("spectral Krylov slice", "Moonlight.LinAlg.Pure.Krylov"), + ("spectral operator slice", "Moonlight.LinAlg.Pure.Operator"), + ("spectral solve slice", "Moonlight.LinAlg.Pure.Spectral"), + ("statics slice", "Moonlight.LinAlg.Pure.Statics") + ] + +staticsForbiddenImports :: [(String, String)] +staticsForbiddenImports = + forbiddenSliceImports + [ ("domain slice", "Moonlight.LinAlg.Pure.Domain"), + ("eigen slice", "Moonlight.LinAlg.Internal.Eigen"), + ("native slice", "Moonlight.LinAlg.Effect.Native"), + ("sparse slice", "Moonlight.LinAlg.Pure.Sparse"), + ("spectral Krylov slice", "Moonlight.LinAlg.Pure.Krylov"), + ("spectral operator slice", "Moonlight.LinAlg.Pure.Operator"), + ("spectral solve slice", "Moonlight.LinAlg.Pure.Spectral"), + ("structured slice", "Moonlight.LinAlg.Pure.Structured") + ] + +spectralForbiddenImports :: [(String, String)] +spectralForbiddenImports = + forbiddenSliceImports + [ ("domain slice", "Moonlight.LinAlg.Pure.Domain"), + ("geometry slice", "Moonlight.LinAlg.Pure.Geometry"), + ("native slice", "Moonlight.LinAlg.Effect.Native"), + ("statics slice", "Moonlight.LinAlg.Pure.Statics") + ] + +nativeForbiddenImports :: [(String, String)] +nativeForbiddenImports = + forbiddenSliceImports + [ ("domain slice", "Moonlight.LinAlg.Pure.Domain"), + ("geometry slice", "Moonlight.LinAlg.Pure.Geometry"), + ("sparse slice", "Moonlight.LinAlg.Pure.Sparse"), + ("statics slice", "Moonlight.LinAlg.Pure.Statics") + ] + +forbiddenSliceImports :: [(String, String)] -> [(String, String)] +forbiddenSliceImports slicePrefixes = + concatMap importPrefixFragments slicePrefixes + +importPrefixFragments :: (String, String) -> [(String, String)] +importPrefixFragments (labelValue, prefixValue) = + [ (labelValue <> " import", "import " <> prefixValue), + (labelValue <> " qualified import", "import qualified " <> prefixValue) + ] + +sliceSourceRoots :: [SliceSourceRoot] +sliceSourceRoots = + [ SliceSourceRoot "src-carrier" carrierForbiddenImports, + SliceSourceRoot "src-structured" noInPackageImports, + SliceSourceRoot "src-eigen" eigenForbiddenImports, + SliceSourceRoot "src-geometry" geometryForbiddenImports, + SliceSourceRoot "src-dense" denseForbiddenImports, + SliceSourceRoot "src-domain" domainForbiddenImports, + SliceSourceRoot "src-sparse" sparseForbiddenImports, + SliceSourceRoot "src-statics" staticsForbiddenImports, + SliceSourceRoot "src-spectral" spectralForbiddenImports, + SliceSourceRoot "src-native" nativeForbiddenImports + ] + +assertForbiddenFragments :: [SourceCheck] -> Assertion +assertForbiddenFragments sourceChecks = do + violations <- concat <$> traverse checkSource sourceChecks + assertBool + ("forbidden architecture fragments remain:\n" <> unlines violations) + (null violations) + +checkSource :: SourceCheck -> IO [String] +checkSource sourceCheck = do + contents <- readSourceFile (sourcePath sourceCheck) + pure + ( violationMessage (sourcePath sourceCheck) + <$> filter + (\(_, forbiddenFragment) -> forbiddenFragment `isInfixOf` contents) + (forbiddenFragments sourceCheck) + ) + +violationMessage :: FilePath -> (String, String) -> String +violationMessage pathValue (labelValue, fragmentValue) = + pathValue <> " contains " <> labelValue <> ": " <> show fragmentValue + +readSourceFile :: FilePath -> IO String +readSourceFile relativePath = + resolveSourcePath relativePath >>= \case + Just resolvedPath -> readFile resolvedPath + Nothing -> assertFailure ("architecture source file is not reachable from test cwd: " <> relativePath) *> pure "" + +existingMissingSource :: MissingSource -> IO (Maybe String) +existingMissingSource missingSource = do + exists <- sourceExists (missingPath missingSource) + pure + ( if exists + then Just (missingPath missingSource <> " (" <> missingLabel missingSource <> ")") + else Nothing + ) + +sourceExists :: FilePath -> IO Bool +sourceExists relativePath = + maybe False (const True) <$> resolveSourcePath relativePath + +resolveSourcePath :: FilePath -> IO (Maybe FilePath) +resolveSourcePath relativePath = + firstJust <$> traverse existingCandidate (sourcePathCandidates relativePath) + +resolveSourceDirectory :: FilePath -> IO (Maybe FilePath) +resolveSourceDirectory relativePath = + firstJust <$> traverse existingDirectoryCandidate (sourcePathCandidates relativePath) + +sourcePathCandidates :: FilePath -> [FilePath] +sourcePathCandidates relativePath = + [ relativePath, + "foundation/moonlight-linalg/" <> relativePath, + "compiler/foundation/moonlight-linalg/" <> relativePath + ] + +existingCandidate :: FilePath -> IO (Maybe FilePath) +existingCandidate candidatePath = do + exists <- doesFileExist candidatePath + pure (if exists then Just candidatePath else Nothing) + +existingDirectoryCandidate :: FilePath -> IO (Maybe FilePath) +existingDirectoryCandidate candidatePath = do + exists <- doesDirectoryExist candidatePath + pure (if exists then Just candidatePath else Nothing) + +firstJust :: [Maybe value] -> Maybe value +firstJust = + foldr (<|>) Nothing diff --git a/test/dense/AdvancedSpec.hs b/test/dense/AdvancedSpec.hs new file mode 100644 --- /dev/null +++ b/test/dense/AdvancedSpec.hs @@ -0,0 +1,877 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE TypeApplications #-} + +module AdvancedSpec + ( tests, + ) +where + +import Data.Foldable qualified as Foldable +import Data.List (isInfixOf, sortBy) +import Data.Map.Strict qualified as Map +import Data.Maybe (catMaybes) +import Data.Ord (comparing) +import qualified Data.Vector.Unboxed as U +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg + ( choleskyDecomp, + canonicalCSRFromEntries, + cooEntries, + cooToDense, + cooToCSR, + csrCols, + csrColumnIndicesVector, + csrMatVecVector, + csrRows, + csrRowOffsetsVector, + csrToCSC, + cscToCOO, + cscColumnOffsetsVector, + cscRowIndicesVector, + cscToDense, + cscValuesVector, + csrToCOO, + csrToDense, + csrValuesVector, + denseToCOO, + denseToCSC, + denseToCSR, + diagonalCSR, + fromListMatrix, + fromListVector, + GraphEdge (..), + graphLaplacianCSR, + mkSparseCSC, + mkSparseCSR, + mkSparseCOO, + SparseCSC, + SparseCSR, + mult, + pathLaplacianCSR, + qrDecompFullColumnRank, + solveCG, + solveDirect, + solveGMRES, + thinSvdFullColumnRank, + symmetricEigen, + toListMatrix, + toListVector, + transpose, + tridiagonalCSR, + ) +import Moonlight.LinAlg.Pure.Dense.Field (PLU (..), pluDecompFullRank) +import Helpers (extractRight) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit + ( Assertion, + assertBool, + assertEqual, + assertFailure, + testCase, + ) +import Test.Tasty.QuickCheck qualified as QC + +tests :: TestTree +tests = + testGroup + "Advanced" + [ testCase "sparse COO/CSR/CSC conversions round-trip dense matrices" testSparseConversions, + QC.testProperty "counting CSR to CSC transpose agrees with sort-based conversion" propCountingCSRToCSCAgreesWithSort, + testCase "COO constructor rejects out-of-bounds entries" testSparseCOORejectsOutOfBounds, + testCase "COO to static dense rejects overflowing type-level shape" testCOOToDenseRejectsTypeLevelCardinalityOverflow, + testCase "COO to static dense rejects type-level dimensions outside Int range" testCOOToDenseRejectsTypeLevelDimensionOverflow, + testCase "COO to CSR combines duplicates and prunes zero storage entries" testCOOToCSRCombinesDuplicateAndPrunesZeroStorageEntries, + testCase "canonical CSR assembly combines duplicates and prunes zeros" testCanonicalCSRFromEntriesCombinesDuplicatesAndPrunesZeros, + testCase "canonical CSR assembly rejects out-of-bounds entries before pruning" testCanonicalCSRFromEntriesRejectsOutOfBoundsBeforePruning, + testCase "structured sparse constructors produce canonical CSR layouts" testStructuredSparseConstructors, + testCase "symmetric tridiagonal CSR has exact canonical storage" testTridiagonalCSR, + testCase "one-vertex path Laplacian is the zero operator" testOneVertexPathLaplacian, + testCase "weighted graph Laplacian canonicalizes parallel undirected edges" testGraphLaplacian, + QC.testProperty "edge-level graph Laplacian agrees with coordinate expansion" propGraphLaplacianAgreesWithCoordinateExpansion, + testCase "graph Laplacian rejects malformed graph declarations" testGraphLaplacianFailures, + testCase "qrDecompFullColumnRank reconstructs dense input" testQrDecomp, + testCase "choleskyDecomp reconstructs SPD matrix" testCholeskyDecomp, + testCase "choleskyDecomp rejects non-symmetric matrix" testCholeskyRejectsNonSymmetric, + testCase "symmetricEigen diagonalizes symmetric matrices" testSymmetricEigen, + testCase "symmetricEigen reconstructs coupled symmetric matrices" testSymmetricEigenReconstructsCoupledMatrix, + testCase "symmetricEigen matches Dirichlet second-difference spectrum" testSymmetricEigenDirichletSecondDifferenceSpectrum, + testCase "symmetricEigen rejects non-finite entries" testSymmetricEigenRejectsNonFinite, + testCase "thinSvdFullColumnRank reconstructs dense input with orthonormal factors" testSvdDecomp, + testCase "solveDirect solves linear systems via PLU" testSolveDirect, + testCase "solveDirect matches exact PLU semantics on generated systems" testSolveDirectGeneratedExactSemantics, + testCase "qrDecompFullColumnRank reconstructs generated matrices" testQrGeneratedResiduals, + testCase "choleskyDecomp reconstructs generated SPD matrices" testCholeskyGeneratedResiduals, + testCase "solveCG converges on SPD systems" testSolveCg, + testCase "solveGMRES converges on non-symmetric systems" testSolveGmres + ] + +assertApproxList :: String -> [Double] -> [Double] -> Assertion +assertApproxList label expected actual = + let tolerance = 1.0e-6 + closeEnough left right = abs (left - right) <= tolerance + in assertBool label (length expected == length actual && and (zipWith closeEnough expected actual)) + +testCOOToDenseRejectsTypeLevelCardinalityOverflow :: Assertion +testCOOToDenseRejectsTypeLevelCardinalityOverflow = + case mkSparseCOO 0 0 ([] :: [(Int, Int, Double)]) >>= cooToDense @4294967296 @4294967296 of + Left failureValue -> + assertEqual + "overflowing static shape" + (InvariantViolation "static sparse/dense shape exceeds Int cardinality") + failureValue + Right _ -> + assertFailure "expected overflowing type-level sparse/dense shape to fail" + +testCOOToDenseRejectsTypeLevelDimensionOverflow :: Assertion +testCOOToDenseRejectsTypeLevelDimensionOverflow = + case mkSparseCOO 0 0 ([] :: [(Int, Int, Double)]) >>= cooToDense @9223372036854775808 @0 of + Left failureValue -> + assertEqual + "out-of-range static dimension" + (InvariantViolation "static sparse/dense dimension exceeds Int cardinality") + failureValue + Right _ -> + assertFailure "expected out-of-range type-level sparse/dense dimension to fail" + +data GeneratedCSRCase = GeneratedCSRCase + { generatedCSRRows :: !Int, + generatedCSRCols :: !Int, + generatedCSREntries :: ![(Int, Int, Double)] + } + deriving stock (Show) + +instance QC.Arbitrary GeneratedCSRCase where + arbitrary = do + rowCount <- QC.chooseInt (0, 8) + columnCount <- QC.chooseInt (0, 8) + if rowCount == 0 || columnCount == 0 + then pure (GeneratedCSRCase rowCount columnCount []) + else do + entryCount <- QC.chooseInt (0, 32) + randomEntries <- + QC.vectorOf + entryCount + (generatedCOOEntry rowCount columnCount) + let duplicateEntries = + [ (0, 0, 1.0), + (0, 0, 2.0) + ] + <> if columnCount > 1 + then + [ (0, 1, 4.0), + (0, 1, -4.0) + ] + else [] + pure (GeneratedCSRCase rowCount columnCount (duplicateEntries <> randomEntries)) + +generatedCOOEntry :: Int -> Int -> QC.Gen (Int, Int, Double) +generatedCOOEntry rowCount columnCount = do + rowIndex <- + QC.chooseInt + ( 0, + if rowCount > 1 + then rowCount - 2 + else 0 + ) + columnIndex <- QC.chooseInt (0, columnCount - 1) + entryValue <- QC.elements [-4.0, -2.0, -1.0, 0.0, 0.5, 1.0, 2.0, 4.0] + pure (rowIndex, columnIndex, entryValue) + +data GeneratedGraphCase = GeneratedGraphCase + { generatedGraphVertices :: ![Int], + generatedGraphEdges :: ![GraphEdge Int] + } + deriving stock (Show) + +instance QC.Arbitrary GeneratedGraphCase where + arbitrary = do + vertexCount <- QC.chooseInt (0, 8) + if vertexCount < 2 + then pure (GeneratedGraphCase [0 .. vertexCount - 1] []) + else do + edgeCount <- QC.chooseInt (0, 40) + randomEdges <- QC.vectorOf edgeCount (generatedGraphEdge vertexCount) + let parallelEdges = + [ GraphEdge 0 1 0.5, + GraphEdge 1 0 1.5 + ] + pure (GeneratedGraphCase [0 .. vertexCount - 1] (parallelEdges <> randomEdges)) + +generatedGraphEdge :: Int -> QC.Gen (GraphEdge Int) +generatedGraphEdge vertexCount = do + leftIndex <- QC.chooseInt (0, vertexCount - 1) + offset <- QC.chooseInt (1, vertexCount - 1) + reversedEdge <- QC.arbitrary + weightValue <- QC.elements [0.0, 0.25, 0.5, 1.0, 2.0, 4.0] + let rightIndex = (leftIndex + offset) `mod` vertexCount + pure + ( if reversedEdge + then GraphEdge rightIndex leftIndex weightValue + else GraphEdge leftIndex rightIndex weightValue + ) + +propCountingCSRToCSCAgreesWithSort :: GeneratedCSRCase -> QC.Property +propCountingCSRToCSCAgreesWithSort GeneratedCSRCase {..} = + case resultValue of + Left err -> + QC.counterexample ("unexpected sparse generation failure: " <> show err) False + Right (countingFingerprint, sortedFingerprint) -> + QC.counterexample + ( "counting transpose = " + <> show countingFingerprint + <> ", sort transpose = " + <> show sortedFingerprint + ) + (countingFingerprint == sortedFingerprint) + where + resultValue = do + cooValue <- mkSparseCOO generatedCSRRows generatedCSRCols generatedCSREntries + csrValue <- cooToCSR cooValue + countingCsc <- csrToCSC csrValue + sortedCsc <- sortBasedCSRToCSC csrValue + pure (cscFingerprint countingCsc, cscFingerprint sortedCsc) + +sortBasedCSRToCSC :: SparseCSR Double -> Either MoonlightError (SparseCSC Double) +sortBasedCSRToCSC csrValue = do + cooValue <- csrToCOO csrValue + let orderedEntries = + sortBy + (comparing (\(rowIndex, columnIndex, _) -> (columnIndex, rowIndex))) + (cooEntries cooValue) + columnOffsets = + offsetsFromSortedAxesForTest + (csrCols csrValue) + ((\(_, columnIndex, _) -> columnIndex) <$> orderedEntries) + rowIndices = (\(rowIndex, _, _) -> rowIndex) <$> orderedEntries + values = (\(_, _, entryValue) -> entryValue) <$> orderedEntries + mkSparseCSC + (csrRows csrValue) + (csrCols csrValue) + columnOffsets + rowIndices + values + +cscFingerprint :: SparseCSC Double -> (U.Vector Int, U.Vector Int, U.Vector Double) +cscFingerprint cscValue = + ( cscColumnOffsetsVector cscValue, + cscRowIndicesVector cscValue, + cscValuesVector cscValue + ) + +propGraphLaplacianAgreesWithCoordinateExpansion :: GeneratedGraphCase -> QC.Property +propGraphLaplacianAgreesWithCoordinateExpansion GeneratedGraphCase {..} = + case (graphLaplacianCSR generatedGraphVertices generatedGraphEdges, coordinateExpansionGraphLaplacian generatedGraphVertices generatedGraphEdges) of + (Right edgeLevelValue, Right coordinateValue) -> + QC.counterexample + ( "edge-level = " + <> show (csrFingerprint edgeLevelValue) + <> ", coordinate = " + <> show (csrFingerprint coordinateValue) + ) + (csrFingerprint edgeLevelValue == csrFingerprint coordinateValue) + (Left leftError, Left rightError) -> + QC.counterexample + ("both constructors rejected generated graph: " <> show (leftError, rightError)) + True + otherResult -> + QC.counterexample ("constructor disagreement: " <> show otherResult) False + +coordinateExpansionGraphLaplacian :: [Int] -> [GraphEdge Int] -> Either MoonlightError (SparseCSR Double) +coordinateExpansionGraphLaplacian vertexOrder graphEdges = do + indexedEdges <- + catMaybes + <$> traverse + (coordinateExpansionGraphEdge (Map.fromList (zip vertexOrder [0 ..]))) + graphEdges + let orderedEdges = + sortBy + (comparing (\(leftIndex, rightIndex, weightValue) -> (leftIndex, rightIndex, weightValue))) + indexedEdges + orderedEntries = + fmap + (\((rowIndex, columnIndex), entryValue) -> (rowIndex, columnIndex, entryValue)) + . filter ((/= 0.0) . snd) + . Map.toAscList + . foldl' + ( \entryMap (rowIndex, columnIndex, entryValue) -> + Map.insertWith (+) (rowIndex, columnIndex) entryValue entryMap + ) + Map.empty + . concatMap coordinateExpansionGraphEdgeContributions + $ orderedEdges + dimension = length vertexOrder + rowOffsets = + offsetsFromSortedAxesForTest + dimension + ((\(rowIndex, _, _) -> rowIndex) <$> orderedEntries) + columnIndices = (\(_, columnIndex, _) -> columnIndex) <$> orderedEntries + values = (\(_, _, entryValue) -> entryValue) <$> orderedEntries + mkSparseCSR dimension dimension rowOffsets columnIndices values + +coordinateExpansionGraphEdge :: + Map.Map Int Int -> + GraphEdge Int -> + Either MoonlightError (Maybe (Int, Int, Double)) +coordinateExpansionGraphEdge vertexIndices edgeValue + | graphEdgeWeight edgeValue == 0.0 = Right Nothing + | otherwise = do + leftIndex <- requireGeneratedVertex "left" (graphEdgeLeft edgeValue) vertexIndices + rightIndex <- requireGeneratedVertex "right" (graphEdgeRight edgeValue) vertexIndices + pure + ( Just + ( min leftIndex rightIndex, + max leftIndex rightIndex, + graphEdgeWeight edgeValue + ) + ) + +requireGeneratedVertex :: String -> Int -> Map.Map Int Int -> Either MoonlightError Int +requireGeneratedVertex endpointRole vertexValue vertexIndices = + case Map.lookup vertexValue vertexIndices of + Nothing -> + Left + ( InvariantViolation + ( "generated graph " + <> endpointRole + <> " endpoint absent: " + <> show vertexValue + ) + ) + Just vertexIndex -> Right vertexIndex + +coordinateExpansionGraphEdgeContributions :: (Int, Int, Double) -> [(Int, Int, Double)] +coordinateExpansionGraphEdgeContributions (leftIndex, rightIndex, weightValue) = + [ (leftIndex, leftIndex, weightValue), + (leftIndex, rightIndex, negate weightValue), + (rightIndex, leftIndex, negate weightValue), + (rightIndex, rightIndex, weightValue) + ] + +offsetsFromSortedAxesForTest :: Int -> [Int] -> [Int] +offsetsFromSortedAxesForTest axisCount sortedAxes = + scanl + (+) + 0 + ( (\axisIndex -> length (filter (== axisIndex) sortedAxes)) + <$> [0 .. axisCount - 1] + ) + +csrFingerprint :: SparseCSR Double -> (U.Vector Int, U.Vector Int, U.Vector Double) +csrFingerprint csrValue = + ( csrRowOffsetsVector csrValue, + csrColumnIndicesVector csrValue, + csrValuesVector csrValue + ) + +testSparseConversions :: Assertion +testSparseConversions = + let result = do + denseMatrix <- fromListMatrix @3 @3 @Double [1.0, 0.0, 0.0, 0.0, 2.0, 3.0, 0.0, 0.0, 4.0] + let cooMatrix = denseToCOO denseMatrix + csrMatrix = denseToCSR denseMatrix + cscMatrix = denseToCSC denseMatrix + cooFromCsr <- csrToCOO csrMatrix + cooFromCsc <- cscToCOO cscMatrix + denseFromCsr <- csrToDense @3 @3 csrMatrix + denseFromCsc <- cscToDense @3 @3 cscMatrix + pure + ( cooEntries cooMatrix, + cooEntries cooFromCsr, + cooEntries cooFromCsc, + toListMatrix denseFromCsr, + toListMatrix denseFromCsc + ) + in extractRight result (\(baseEntries, csrEntries, cscEntries, csrDense, cscDense) -> do + assertEqual "COO non-zero entries" baseEntries csrEntries + assertEqual "CSC -> COO preserves entries" baseEntries cscEntries + assertEqual "CSR round-trip dense payload" [1.0, 0.0, 0.0, 0.0, 2.0, 3.0, 0.0, 0.0, 4.0] csrDense + assertEqual "CSC round-trip dense payload" [1.0, 0.0, 0.0, 0.0, 2.0, 3.0, 0.0, 0.0, 4.0] cscDense) + +testSparseCOORejectsOutOfBounds :: Assertion +testSparseCOORejectsOutOfBounds = + case mkSparseCOO 2 2 [(2, 0, 1.0 :: Double)] of + Left (InvariantViolation message) -> + assertBool "shape error should mention bounds" ("out of bounds" `isInfixOf` message) + Left err -> + assertFailure ("expected COO shape error, got: " <> show err) + Right _ -> + assertFailure "expected COO constructor to reject out-of-bounds entry" + +testCOOToCSRCombinesDuplicateAndPrunesZeroStorageEntries :: Assertion +testCOOToCSRCombinesDuplicateAndPrunesZeroStorageEntries = + let result = do + cooValue <- + mkSparseCOO + 2 + 3 + [ (1, 2, 4.0 :: Double), + (0, 1, 2.0), + (0, 1, 3.0), + (0, 2, 0.0) + ] + csrValue <- cooToCSR cooValue + denseFromCsr <- csrToDense @2 @3 csrValue + matvecResult <- csrMatVecVector csrValue (U.fromList [10.0, 20.0, 30.0]) + pure + ( csrRowOffsetsVector csrValue, + csrColumnIndicesVector csrValue, + csrValuesVector csrValue, + toListMatrix denseFromCsr, + matvecResult + ) + in extractRight result $ \(rowOffsets, columnIndices, values, denseValues, matvecValues) -> do + assertEqual "COO -> CSR row offsets" (U.fromList [0, 1, 2]) rowOffsets + assertEqual "COO -> CSR column indices" (U.fromList [1, 2]) columnIndices + assertEqual "COO -> CSR values combine duplicates and prune explicit zero" (U.fromList [5.0, 4.0]) values + assertEqual "dense conversion sums duplicate coordinates" [0.0, 5.0, 0.0, 0.0, 0.0, 4.0] denseValues + assertEqual "matvec sums duplicate stored entries" (U.fromList [100.0, 120.0]) matvecValues + +testCanonicalCSRFromEntriesCombinesDuplicatesAndPrunesZeros :: Assertion +testCanonicalCSRFromEntriesCombinesDuplicatesAndPrunesZeros = + let result = do + csrValue <- + canonicalCSRFromEntries + 2 + 3 + [ (0, 1, 2.0 :: Double), + (0, 1, 3.0), + (0, 2, 0.0), + (1, 0, 5.0), + (1, 0, -5.0), + (1, 2, 4.0) + ] + denseFromCsr <- csrToDense @2 @3 csrValue + pure + ( csrRowOffsetsVector csrValue, + csrColumnIndicesVector csrValue, + csrValuesVector csrValue, + toListMatrix denseFromCsr + ) + in extractRight result $ \(rowOffsets, columnIndices, values, denseValues) -> do + assertEqual "canonical CSR row offsets" (U.fromList [0, 1, 2]) rowOffsets + assertEqual "canonical CSR column indices" (U.fromList [1, 2]) columnIndices + assertEqual "canonical CSR values" (U.fromList [5.0, 4.0]) values + assertEqual "canonical dense payload" [0.0, 5.0, 0.0, 0.0, 0.0, 4.0] denseValues + +testCanonicalCSRFromEntriesRejectsOutOfBoundsBeforePruning :: Assertion +testCanonicalCSRFromEntriesRejectsOutOfBoundsBeforePruning = + case canonicalCSRFromEntries 1 1 [(2, 0, 1.0 :: Double), (2, 0, -1.0)] of + Left (InvariantViolation message) -> + assertBool "error should mention out of bounds" ("out of bounds" `isInfixOf` message) + Left other -> + assertFailure ("expected InvariantViolation, got: " <> show other) + Right _ -> + assertFailure "canonical CSR must reject invalid entries even when duplicates sum to zero" + +testStructuredSparseConstructors :: Assertion +testStructuredSparseConstructors = + let result = do + diagonalMatrix <- diagonalCSR [0.0 :: Double, 2.0, 0.0, 4.0] + pathMatrix <- pathLaplacianCSR 4 + densePathMatrix <- csrToDense @4 @4 pathMatrix + pure + ( csrRowOffsetsVector diagonalMatrix, + csrColumnIndicesVector diagonalMatrix, + csrValuesVector diagonalMatrix, + toListMatrix densePathMatrix + ) + in extractRight result $ \(diagonalOffsets, diagonalColumns, diagonalValues, pathDenseValues) -> do + assertEqual "diagonal CSR prunes zero diagonal entries" (U.fromList [0, 0, 1, 1, 2]) diagonalOffsets + assertEqual "diagonal CSR column indices" (U.fromList [1, 3]) diagonalColumns + assertEqual "diagonal CSR values" (U.fromList [2.0, 4.0]) diagonalValues + assertEqual + "path graph Laplacian dense payload" + [ 1.0, -1.0, 0.0, 0.0, + -1.0, 2.0, -1.0, 0.0, + 0.0, -1.0, 2.0, -1.0, + 0.0, 0.0, -1.0, 1.0 + ] + pathDenseValues + +testTridiagonalCSR :: Assertion +testTridiagonalCSR = + let result = + tridiagonalCSR + [2.0 :: Double, 3.0, 4.0] + [-1.0, -2.0] + in extractRight result $ \matrixValue -> do + assertEqual "row offsets" (U.fromList [0, 2, 5, 7]) (csrRowOffsetsVector matrixValue) + assertEqual "column indices" (U.fromList [0, 1, 0, 1, 2, 1, 2]) (csrColumnIndicesVector matrixValue) + assertEqual "values" (U.fromList [2.0, -1.0, -1.0, 3.0, -2.0, -2.0, 4.0]) (csrValuesVector matrixValue) + +testOneVertexPathLaplacian :: Assertion +testOneVertexPathLaplacian = + let result = do + matrixValue <- pathLaplacianCSR 1 + denseValue <- csrToDense @1 @1 matrixValue + pure (toListMatrix denseValue) + in extractRight result $ + assertEqual "P1 Laplacian" [0.0] + +testGraphLaplacian :: Assertion +testGraphLaplacian = + let result = do + matrixValue <- + graphLaplacianCSR + ["b", "a", "c"] + [ GraphEdge "a" "b" 1.0, + GraphEdge "b" "a" 2.0, + GraphEdge "b" "c" 4.0 + ] + denseValue <- csrToDense @3 @3 matrixValue + pure + ( csrRowOffsetsVector matrixValue, + csrColumnIndicesVector matrixValue, + csrValuesVector matrixValue, + toListMatrix denseValue + ) + in extractRight result $ \(offsets, columns, values, denseEntries) -> do + assertEqual "offsets" (U.fromList [0, 3, 5, 7]) offsets + assertEqual "columns" (U.fromList [0, 1, 2, 0, 1, 0, 2]) columns + assertEqual "values" (U.fromList [7.0, -3.0, -4.0, -3.0, 3.0, -4.0, 4.0]) values + assertEqual + "dense Laplacian" + [ 7.0, -3.0, -4.0, + -3.0, 3.0, 0.0, + -4.0, 0.0, 4.0 + ] + denseEntries + +testGraphLaplacianFailures :: Assertion +testGraphLaplacianFailures = do + assertGraphFailure "duplicate vertices" (graphLaplacianCSR ["a", "a"] []) + assertGraphFailure "unknown endpoint" (graphLaplacianCSR ["a"] [GraphEdge "a" "b" 1.0]) + assertGraphFailure "self loop" (graphLaplacianCSR ["a"] [GraphEdge "a" "a" 1.0]) + assertGraphFailure "negative weight" (graphLaplacianCSR ["a", "b"] [GraphEdge "a" "b" (-1.0)]) + assertGraphFailure "non-finite weight" (graphLaplacianCSR ["a", "b"] [GraphEdge "a" "b" (0.0 / 0.0)]) + +assertGraphFailure :: String -> Either MoonlightError value -> Assertion +assertGraphFailure label resultValue = + case resultValue of + Left _ -> pure () + Right _ -> assertFailure (label <> ": expected graph construction failure") + +testQrDecomp :: Assertion +testQrDecomp = + let result = do + matrixValue <- fromListMatrix @3 @2 @Double [1.0, 1.0, 1.0, 0.0, 1.0, 2.0] + (qMatrix, rMatrix) <- qrDecompFullColumnRank matrixValue + reconstructed <- mult qMatrix rMatrix + pure (toListMatrix reconstructed) + in extractRight result (\values -> assertApproxList "QR reconstruction" [1.0, 1.0, 1.0, 0.0, 1.0, 2.0] values) + +testCholeskyDecomp :: Assertion +testCholeskyDecomp = + let result = do + matrixValue <- fromListMatrix @2 @2 @Double [4.0, 2.0, 2.0, 3.0] + lowerMatrix <- choleskyDecomp matrixValue + transposedLower <- transpose lowerMatrix + reconstructed <- mult lowerMatrix transposedLower + pure (toListMatrix reconstructed) + in extractRight result (\values -> assertApproxList "Cholesky reconstruction" [4.0, 2.0, 2.0, 3.0] values) + +testCholeskyRejectsNonSymmetric :: Assertion +testCholeskyRejectsNonSymmetric = + let result = do + matrixValue <- fromListMatrix @2 @2 @Double [4.0, 1.0, 3.0, 3.0] + choleskyDecomp matrixValue + in case result of + Left (InvariantViolation msg) -> assertBool "error should mention symmetric" ("symmetric" `isInfixOf` msg) + Left other -> assertFailure ("expected InvariantViolation about symmetry, got: " <> show other) + Right _ -> assertFailure "Cholesky should reject non-symmetric matrix" + +testSymmetricEigen :: Assertion +testSymmetricEigen = + let result = do + matrixValue <- fromListMatrix @2 @2 @Double [2.0, 0.0, 0.0, 3.0] + (eigenvalues, eigenvectors) <- symmetricEigen matrixValue + pure (toListVector eigenvalues, toListMatrix eigenvectors) + in extractRight result (\(values, vectors) -> do + assertApproxList "eigenvalues" [3.0, 2.0] values + assertApproxList "eigenvector matrix" [0.0, 1.0, 1.0, 0.0] vectors) + +testSymmetricEigenReconstructsCoupledMatrix :: Assertion +testSymmetricEigenReconstructsCoupledMatrix = + let sourceRows = + [ 4.0, 1.0, 2.0, + 1.0, 3.0, 0.5, + 2.0, 0.5, 5.0 + ] + result = do + matrixValue <- fromListMatrix @3 @3 @Double sourceRows + (eigenvalues, eigenvectors) <- symmetricEigen matrixValue + diagonalized <- fromListMatrix @3 @3 @Double (diagonalMatrixEntries (toListVector eigenvalues)) + weightedEigenvectors <- mult eigenvectors diagonalized + transposedEigenvectors <- transpose eigenvectors + reconstructed <- mult weightedEigenvectors transposedEigenvectors + pure (toListMatrix reconstructed) + in extractRight result (assertApproxList "symmetric eigen reconstruction" sourceRows) + +testSymmetricEigenDirichletSecondDifferenceSpectrum :: Assertion +testSymmetricEigenDirichletSecondDifferenceSpectrum = + let result = do + matrixValue <- + fromListMatrix @3 @3 @Double + [ 2.0, -1.0, 0.0, + -1.0, 2.0, -1.0, + 0.0, -1.0, 2.0 + ] + (eigenvalues, _) <- symmetricEigen matrixValue + pure (toListVector eigenvalues) + expected = + [ 2.0 + sqrt 2.0, + 2.0, + 2.0 - sqrt 2.0 + ] + in extractRight result (assertApproxList "Dirichlet second-difference spectrum" expected) + +testSymmetricEigenRejectsNonFinite :: Assertion +testSymmetricEigenRejectsNonFinite = + let result = do + matrixValue <- fromListMatrix @2 @2 @Double [1.0, 0.0, 0.0, 0.0 / 0.0] + symmetricEigen matrixValue + in case result of + Left (InvariantViolation msg) -> assertBool "error should mention finite" ("finite" `isInfixOf` msg) + Left other -> assertFailure ("expected InvariantViolation about finite entries, got: " <> show other) + Right _ -> assertFailure "symmetricEigen should reject NaN entries" + +testSvdDecomp :: Assertion +testSvdDecomp = + let result = do + matrixValue <- fromListMatrix @2 @2 @Double [3.0, 0.0, 0.0, 2.0] + (uMatrix, sMatrix, vTMatrix) <- thinSvdFullColumnRank matrixValue + usMatrix <- mult uMatrix sMatrix + reconstructed <- mult usMatrix vTMatrix + uTMatrix <- transpose uMatrix + uOrthogonality <- mult uTMatrix uMatrix + vMatrix <- transpose vTMatrix + vOrthogonality <- mult vTMatrix vMatrix + pure (toListMatrix reconstructed, toListMatrix uOrthogonality, toListMatrix vOrthogonality) + in extractRight result $ \(reconstructedValues, uOrthogonalityValues, vOrthogonalityValues) -> do + assertApproxList "SVD reconstruction" [3.0, 0.0, 0.0, 2.0] reconstructedValues + assertApproxList "SVD U orthonormality" (identityMatrixEntries 2) uOrthogonalityValues + assertApproxList "SVD V orthonormality" (identityMatrixEntries 2) vOrthogonalityValues + +testSolveDirect :: Assertion +testSolveDirect = + let result = do + matrixValue <- fromListMatrix @2 @2 @Double [3.0, 1.0, 1.0, 2.0] + vectorValue <- fromListVector @2 @Double [9.0, 8.0] + solution <- solveDirect matrixValue vectorValue + pure (toListVector solution) + in extractRight result (\values -> assertApproxList "direct solver solution" [2.0, 3.0] values) + +testSolveDirectGeneratedExactSemantics :: Assertion +testSolveDirectGeneratedExactSemantics = + Foldable.traverse_ assertSeed generatedSeeds + where + assertSeed seedValue = + let matrixRational = generatedSolveMatrix seedValue + solutionRational = generatedSolveSolution seedValue + rhsRational = multiplySquareRowsVector 3 matrixRational solutionRational + result = do + exactSolution <- exactPluSolve3 matrixRational rhsRational + matrixValue <- fromListMatrix @3 @3 @Double (fmap fromRational matrixRational) + rhsValue <- fromListVector @3 @Double (fmap fromRational rhsRational) + solutionValue <- solveDirect matrixValue rhsValue + pure (fmap fromRational exactSolution, toListVector solutionValue, fmap fromRational rhsRational, fmap fromRational matrixRational) + in extractRight result $ \(expected, actual, rhsValues, matrixValues) -> do + assertApproxList ("generated exact solve seed " <> show seedValue) expected actual + assertResidualBelow + ("generated solve residual seed " <> show seedValue) + 1.0e-8 + (matrixVectorResidual 3 matrixValues actual rhsValues) + +testQrGeneratedResiduals :: Assertion +testQrGeneratedResiduals = + Foldable.traverse_ assertSeed generatedSeeds + where + assertSeed seedValue = + let matrixEntries = generatedQrMatrix seedValue + result = do + matrixValue <- fromListMatrix @4 @3 @Double matrixEntries + (qMatrix, rMatrix) <- qrDecompFullColumnRank matrixValue + reconstructed <- mult qMatrix rMatrix + pure (toListMatrix reconstructed) + in extractRight result $ + assertResidualBelow + ("generated QR residual seed " <> show seedValue) + 1.0e-8 + . maxAbsDifference matrixEntries + +testCholeskyGeneratedResiduals :: Assertion +testCholeskyGeneratedResiduals = + Foldable.traverse_ assertSeed generatedSeeds + where + assertSeed seedValue = + let matrixEntries = generatedSpdMatrix seedValue + result = do + matrixValue <- fromListMatrix @3 @3 @Double matrixEntries + lowerMatrix <- choleskyDecomp matrixValue + transposedLower <- transpose lowerMatrix + reconstructed <- mult lowerMatrix transposedLower + pure (toListMatrix reconstructed) + in extractRight result $ + assertResidualBelow + ("generated Cholesky residual seed " <> show seedValue) + 1.0e-8 + . maxAbsDifference matrixEntries + +testSolveCg :: Assertion +testSolveCg = + let result = do + matrixValue <- fromListMatrix @2 @2 @Double [4.0, 1.0, 1.0, 3.0] + vectorValue <- fromListVector @2 @Double [1.0, 2.0] + solution <- solveCG matrixValue vectorValue + pure (toListVector solution) + in extractRight result (\values -> assertApproxList "CG solver solution" [1.0 / 11.0, 7.0 / 11.0] values) + +testSolveGmres :: Assertion +testSolveGmres = + let result = do + matrixValue <- fromListMatrix @2 @2 @Double [3.0, 2.0, 0.0, 1.0] + vectorValue <- fromListVector @2 @Double [2.0, 1.0] + solution <- solveGMRES matrixValue vectorValue + pure (toListVector solution) + in extractRight result (\values -> assertApproxList "GMRES solver solution" [0.0, 1.0] values) + +diagonalMatrixEntries :: [Double] -> [Double] +diagonalMatrixEntries diagonalValues = + let matrixSize = length diagonalValues + in [ diagonalEntry rowIndex columnIndex + | rowIndex <- [0 .. matrixSize - 1], + columnIndex <- [0 .. matrixSize - 1] + ] + where + diagonalEntry rowIndex columnIndex + | rowIndex == columnIndex = + case drop rowIndex diagonalValues of + diagonalValue : _ -> diagonalValue + [] -> 0.0 + | otherwise = 0.0 + +identityMatrixEntries :: Int -> [Double] +identityMatrixEntries matrixSize = + [ if rowIndex == columnIndex then 1.0 else 0.0 + | rowIndex <- [0 .. matrixSize - 1], + columnIndex <- [0 .. matrixSize - 1] + ] + +generatedSeeds :: [Int] +generatedSeeds = [1, 2, 3, 5, 8, 13] + +generatedSolveMatrix :: Int -> [Rational] +generatedSolveMatrix seedValue = + fmap fromIntegral [generatedSolveEntry seedValue rowIndex columnIndex | rowIndex <- [0 .. 2], columnIndex <- [0 .. 2]] + +generatedSolveEntry :: Int -> Int -> Int -> Int +generatedSolveEntry seedValue rowIndex columnIndex + | rowIndex == columnIndex = 12 + seedValue + rowIndex + | otherwise = ((seedValue + rowIndex * 3 + columnIndex * 5) `mod` 5) - 2 + +generatedSolveSolution :: Int -> [Rational] +generatedSolveSolution seedValue = + fmap fromIntegral [seedValue + 1, 3 - seedValue, seedValue * 2 - 5] + +generatedQrMatrix :: Int -> [Double] +generatedQrMatrix seedValue = + [ generatedQrEntry seedValue rowIndex columnIndex + | rowIndex <- [0 .. 3], + columnIndex <- [0 .. 2] + ] + +generatedQrEntry :: Int -> Int -> Int -> Double +generatedQrEntry seedValue rowIndex columnIndex + | rowIndex == columnIndex = fromIntegral (8 + seedValue + columnIndex) + | otherwise = fromIntegral (((seedValue + rowIndex * 2 + columnIndex * 3) `mod` 7) - 3) / 5.0 + +generatedSpdMatrix :: Int -> [Double] +generatedSpdMatrix seedValue = + [ sum [generatedLowerEntry seedValue rowIndex k * generatedLowerEntry seedValue columnIndex k | k <- [0 .. 2]] + | rowIndex <- [0 .. 2], + columnIndex <- [0 .. 2] + ] + +generatedLowerEntry :: Int -> Int -> Int -> Double +generatedLowerEntry seedValue rowIndex columnIndex + | columnIndex > rowIndex = 0.0 + | rowIndex == columnIndex = fromIntegral (4 + seedValue + rowIndex) + | otherwise = fromIntegral (((seedValue + rowIndex * 3 + columnIndex * 2) `mod` 5) - 2) / 4.0 + +exactPluSolve3 :: [Rational] -> [Rational] -> Either MoonlightError [Rational] +exactPluSolve3 matrixValues rhsValues = do + matrixValue <- fromListMatrix @3 @3 @Rational matrixValues + pluValue <- pluDecompFullRank matrixValue + let permutationRows = rowMajorRows 3 (toListMatrix (pluPermutation pluValue)) + lowerRows = rowMajorRows 3 (toListMatrix (pluLower pluValue)) + upperRows = rowMajorRows 3 (toListMatrix (pluUpper pluValue)) + permutedRhs = multiplyRowsVector permutationRows rhsValues + forwardValues <- forwardSubstituteRational lowerRows permutedRhs + backwardSubstituteRational upperRows forwardValues + +forwardSubstituteRational :: [[Rational]] -> [Rational] -> Either MoonlightError [Rational] +forwardSubstituteRational lowerRows rhsValues = go 0 [] lowerRows rhsValues + where + go :: Int -> [Rational] -> [[Rational]] -> [Rational] -> Either MoonlightError [Rational] + go !_ solvedValues [] [] = Right solvedValues + go !rowIndex solvedValues (rowValues : remainingRows) (rhsValue : remainingRhs) = do + diagonalValue <- requireTestEntry ("exact forward diagonal missing at row " <> show rowIndex) rowIndex rowValues + if diagonalValue == 0 + then Left (InvariantViolation "exact forward substitution encountered zero diagonal") + else + let knownContribution = sum (zipWith (*) (take rowIndex rowValues) solvedValues) + nextValue = (rhsValue - knownContribution) / diagonalValue + in go (rowIndex + 1) (solvedValues <> [nextValue]) remainingRows remainingRhs + go _ _ _ _ = Left (InvariantViolation "exact forward substitution shape mismatch") + +backwardSubstituteRational :: [[Rational]] -> [Rational] -> Either MoonlightError [Rational] +backwardSubstituteRational upperRows rhsValues = go (length upperRows - 1) [] + where + go !rowIndex solvedSuffix + | rowIndex < 0 = Right solvedSuffix + | otherwise = do + rowValues <- requireTestEntry ("exact backward row missing at row " <> show rowIndex) rowIndex upperRows + rhsValue <- requireTestEntry ("exact backward RHS missing at row " <> show rowIndex) rowIndex rhsValues + diagonalValue <- requireTestEntry ("exact backward diagonal missing at row " <> show rowIndex) rowIndex rowValues + if diagonalValue == 0 + then Left (InvariantViolation "exact backward substitution encountered zero diagonal") + else + let knownContribution = sum (zipWith (*) (drop (rowIndex + 1) rowValues) solvedSuffix) + nextValue = (rhsValue - knownContribution) / diagonalValue + in go (rowIndex - 1) (nextValue : solvedSuffix) + +rowMajorRows :: Int -> [a] -> [[a]] +rowMajorRows columnCount values + | columnCount <= 0 = [] + | otherwise = + case splitAt columnCount values of + ([], []) -> [] + (rowValues, remainingValues) -> rowValues : rowMajorRows columnCount remainingValues + +multiplyRowsVector :: Num a => [[a]] -> [a] -> [a] +multiplyRowsVector rows vectorValues = + fmap (\rowValues -> sum (zipWith (*) rowValues vectorValues)) rows + +multiplySquareRowsVector :: Num a => Int -> [a] -> [a] -> [a] +multiplySquareRowsVector columnCount matrixValues = + multiplyRowsVector (rowMajorRows columnCount matrixValues) + +matrixVectorResidual :: Int -> [Double] -> [Double] -> [Double] -> Double +matrixVectorResidual columnCount matrixValues vectorValues rhsValues = + maxAbsDifference rhsValues (multiplySquareRowsVector columnCount matrixValues vectorValues) + +maxAbsDifference :: [Double] -> [Double] -> Double +maxAbsDifference expected actual = + maximum (0.0 : zipWith (\leftValue rightValue -> abs (leftValue - rightValue)) expected actual) + +assertResidualBelow :: String -> Double -> Double -> Assertion +assertResidualBelow label tolerance residualValue = + assertBool (label <> ": residual " <> show residualValue <> " exceeded " <> show tolerance) (residualValue <= tolerance) + +requireTestEntry :: String -> Int -> [a] -> Either MoonlightError a +requireTestEntry label targetIndex values = + case drop targetIndex values of + entryValue : _ -> Right entryValue + [] -> Left (InvariantViolation label) diff --git a/test/dense/BasicSpec.hs b/test/dense/BasicSpec.hs new file mode 100644 --- /dev/null +++ b/test/dense/BasicSpec.hs @@ -0,0 +1,74 @@ + +module BasicSpec + ( tests, + ) +where + +import Moonlight.LinAlg + ( Matrix, + add, + fromListMatrix, + gf2One, + gf2Zero, + mapMatrix, + mult, + toListMatrix, + transpose, + ) +import Moonlight.Core (MoonlightError) +import Helpers (extractRight) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit + ( Assertion, + assertEqual, + testCase, + ) + +tests :: TestTree +tests = + testGroup + "Basic" + [ testCase "add on Double matrices" testAdd, + testCase "multiply on Double matrices" testMultiply, + testCase "transpose on Double matrices" testTranspose, + testCase "mapMatrix changes scalar type through a direct function" testMapMatrix, + testCase "multiply on GF2 matrices uses algebraic semantics" testGF2Multiply + ] + +testAdd :: Assertion +testAdd = + let result = do + left <- fromListMatrix @2 @2 @Double [1.0, 2.0, 3.0, 4.0] + right <- fromListMatrix @2 @2 @Double [4.0, 3.0, 2.0, 1.0] + add left right + in extractRight result (\value -> assertEqual "matrix sum" [5.0, 5.0, 5.0, 5.0] (toListMatrix value)) + +testMultiply :: Assertion +testMultiply = + let result = do + left <- fromListMatrix @2 @2 @Double [1.0, 2.0, 3.0, 4.0] + right <- fromListMatrix @2 @2 @Double [2.0, 0.0, 1.0, 2.0] + mult left right + in extractRight result (\value -> assertEqual "matrix product" [4.0, 4.0, 10.0, 8.0] (toListMatrix value)) + +testTranspose :: Assertion +testTranspose = + let result = do + matrixValue <- fromListMatrix @2 @3 @Double [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + transpose matrixValue + in extractRight result (\value -> assertEqual "matrix transpose" [1.0, 4.0, 2.0, 5.0, 3.0, 6.0] (toListMatrix value)) + +testMapMatrix :: Assertion +testMapMatrix = + let result = do + matrixValue <- fromListMatrix @2 @2 @Double [1.2, 2.8, 3.4, 4.9] + mapMatrix round matrixValue :: Either MoonlightError (Matrix 2 2 Integer) + in extractRight result (\value -> assertEqual "mapped matrix" [1, 3, 3, 5] (toListMatrix value)) + +testGF2Multiply :: Assertion +testGF2Multiply = + let result = do + left <- fromListMatrix @2 @2 [gf2One, gf2One, gf2Zero, gf2One] + right <- fromListMatrix @2 @2 [gf2One, gf2Zero, gf2One, gf2One] + mult left right + in extractRight result (\value -> assertEqual "GF2 product" [gf2Zero, gf2One, gf2One, gf2One] (toListMatrix value)) diff --git a/test/dense/BlockSpec.hs b/test/dense/BlockSpec.hs new file mode 100644 --- /dev/null +++ b/test/dense/BlockSpec.hs @@ -0,0 +1,54 @@ +module BlockSpec + ( tests, + ) +where + +import Data.Ratio ((%)) +import Moonlight.LinAlg + ( BlockMatrixFailure (..), + GF2 (..), + invertGF2Block, + invertRationalBlock, + invertUnimodularIntegerBlock, + ) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertEqual, testCase) + +tests :: TestTree +tests = + testGroup + "block inverse" + [ testCase "rational 2x2 inverse satisfies both sides" testRationalInverse, + testCase "GF2 invertible block succeeds" testGF2Inverse, + testCase "integer unimodular block succeeds" testIntegerUnimodularInverse, + testCase "integer non-unimodular block is rejected" testIntegerNonUnimodularRejected + ] + +testRationalInverse :: IO () +testRationalInverse = + assertEqual + "rational inverse" + (Right [[(-2) :: Rational, 1], [3 % 2, (-1) % 2]]) + (invertRationalBlock [[1, 2], [3, 4]]) + +testGF2Inverse :: IO () +testGF2Inverse = + assertEqual + "GF2 inverse" + (Right [[GF2One, GF2One], [GF2Zero, GF2One]]) + (invertGF2Block [[GF2One, GF2One], [GF2Zero, GF2One]]) + +testIntegerUnimodularInverse :: IO () +testIntegerUnimodularInverse = + assertEqual + "integer unimodular inverse" + (Right [[1 :: Integer, -1], [0, 1]]) + (invertUnimodularIntegerBlock [[1, 1], [0, 1]]) + +testIntegerNonUnimodularRejected :: IO () +testIntegerNonUnimodularRejected = + assertEqual + "integer non-unimodular rejection" + (Left (BlockMatrixNonUnimodular [[1 % 2]])) + (invertUnimodularIntegerBlock [[2 :: Integer]]) + diff --git a/test/dense/DenseFlatSpec.hs b/test/dense/DenseFlatSpec.hs new file mode 100644 --- /dev/null +++ b/test/dense/DenseFlatSpec.hs @@ -0,0 +1,179 @@ +{-# LANGUAGE RecordWildCards #-} + +module DenseFlatSpec (tests) where + +import Data.Vector.Storable qualified as S +import Data.Vector.Unboxed qualified as U +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Internal.Eigen.Residual (ResidualReport (..)) +import Moonlight.LinAlg.Internal.Eigen.Symmetric + ( CertifiedSymmetricEigenResult (..), + certifySymmetricEigenResult, + symmetricEigenPairsDenseUnchecked, + ) +import Moonlight.LinAlg.Pure.Dense.Flat (trustedDenseDoubleMatrixRowMajor) +import Moonlight.LinAlg.Dense + ( denseDoubleMatrixShape, + denseDoubleMatrixToRows, + denseDoubleMatrixVectorProduct, + mkDenseDoubleMatrixRowMajor, + mkDenseDoubleMatrixRows, + ) +import Moonlight.LinAlg.Native + ( denseDoubleLinearSolveLapack, + denseDoubleMatrixProductBlas, + denseDoubleSymmetricEigenpairsLapack, + ) +import Moonlight.LinAlg.Spectral + ( eigenpairCount, + eigenpairResidualNorms, + eigenpairValues, + ) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, testCase) +import Prelude + +tests :: TestTree +tests = + testGroup + "Dense flat Double matrix" + [ testCase "row-major constructor rejects invalid payload length" $ + assertEqual + "shape error" + (Left (InvariantViolation "dense Double row-major payload length mismatch: expected 6 values but received 5")) + (mkDenseDoubleMatrixRowMajor 2 3 (S.fromList [1.0 .. 5.0])), + testCase "row-major constructor rejects non-finite payloads" $ + assertEqual + "finite payload" + (Left (InvariantViolation "dense Double row-major payload requires finite entries")) + (mkDenseDoubleMatrixRowMajor 1 1 (S.fromList [0 / 0])), + testCase "row-major constructor rejects wrapped shape cardinality" $ + let wrappedDimension = 2 ^ (32 :: Int) + in assertEqual + "wrapped shape" + (Left (InvariantViolation "dense Double matrix dimensions exceed Int cardinality")) + (mkDenseDoubleMatrixRowMajor wrappedDimension wrappedDimension S.empty), + testCase "row constructor preserves rectangular shape and projection" $ + assertEqual + "rows" + (Right ((2, 3), [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])) + (fmap (\matrixValue -> (denseDoubleMatrixShape matrixValue, denseDoubleMatrixToRows matrixValue)) (mkDenseDoubleMatrixRows [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])), + testCase "matrix/vector product uses validated flat storage" $ + assertEqual + "matvec" + (Right (S.fromList [140.0, 320.0])) + ( do + matrixValue <- mkDenseDoubleMatrixRowMajor 2 3 (S.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + denseDoubleMatrixVectorProduct matrixValue (S.fromList [10.0, 20.0, 30.0]) + ), + testCase "matrix/vector product rejects vector shape mismatch" $ + assertEqual + "shape error" + (Left (InvariantViolation "dense Double matrix/vector shape mismatch (matrix=(2,3), vector=2)")) + ( do + matrixValue <- mkDenseDoubleMatrixRowMajor 2 3 (S.fromList [1.0 .. 6.0]) + denseDoubleMatrixVectorProduct matrixValue (S.fromList [10.0, 20.0]) + ), + testCase "BLAS matrix product matches reference rows" testDenseDoubleMatrixProductBlas, + testCase "LAPACK dense solve matches reference solution" testDenseDoubleLinearSolveLapack, + testCase "LAPACK dense solve rejects singular input" testDenseDoubleLinearSolveSingular, + testCase "LAPACK dense symmetric eigenpairs are residual certified" testDenseDoubleSymmetricEigenpairsLapack, + testCase "pure dense symmetric eigen certification is explicit" testPureDenseSymmetricEigenCertification + , testCase "pure dense symmetric eigen rejects oversized workspace cardinality" testPureDenseSymmetricEigenWorkspaceOverflow + ] + +testDenseDoubleMatrixProductBlas :: Assertion +testDenseDoubleMatrixProductBlas = do + productResult <- + case (mkDenseDoubleMatrixRowMajor 2 3 (S.fromList [1.0 .. 6.0]), mkDenseDoubleMatrixRowMajor 3 2 (S.fromList [7.0 .. 12.0])) of + (Right leftMatrix, Right rightMatrix) -> + denseDoubleMatrixProductBlas leftMatrix rightMatrix + (Left err, _) -> pure (Left err) + (_, Left err) -> pure (Left err) + assertEqual + "matrix product rows" + (Right [[58.0, 64.0], [139.0, 154.0]]) + (denseDoubleMatrixToRows <$> productResult) + +testDenseDoubleLinearSolveLapack :: Assertion +testDenseDoubleLinearSolveLapack = do + solveResult <- + case mkDenseDoubleMatrixRowMajor 2 2 (S.fromList [3.0, 1.0, 1.0, 2.0]) of + Left err -> pure (Left err) + Right matrixValue -> + denseDoubleLinearSolveLapack matrixValue (S.fromList [9.0, 8.0]) + assertApproxStorableVector "solution" 1.0e-10 (S.fromList [2.0, 3.0]) solveResult + +testDenseDoubleLinearSolveSingular :: Assertion +testDenseDoubleLinearSolveSingular = do + solveResult <- + case mkDenseDoubleMatrixRowMajor 2 2 (S.fromList [1.0, 2.0, 2.0, 4.0]) of + Left err -> pure (Left err) + Right matrixValue -> + denseDoubleLinearSolveLapack matrixValue (S.fromList [1.0, 2.0]) + assertEqual + "singular solve" + (Left (InvariantViolation "LAPACK DGESV detected exact singularity at U diagonal 2")) + solveResult + +testDenseDoubleSymmetricEigenpairsLapack :: Assertion +testDenseDoubleSymmetricEigenpairsLapack = do + eigenResult <- + case mkDenseDoubleMatrixRowMajor 2 2 (S.fromList [2.0, 0.0, 0.0, 3.0]) of + Left err -> pure (Left err) + Right matrixValue -> denseDoubleSymmetricEigenpairsLapack matrixValue + case eigenResult of + Left err -> assertEqual "eigen success" (Right ()) (Left err) + Right pairs -> do + assertEqual "eigenpair count" 2 (eigenpairCount pairs) + assertApproxUnboxedVector "eigenvalues" 1.0e-10 (U.fromList [2.0, 3.0]) (Right (eigenpairValues pairs)) + assertBool + "residuals stay certified" + (U.all (<= 1.0e-10) (eigenpairResidualNorms pairs)) + +testPureDenseSymmetricEigenCertification :: Assertion +testPureDenseSymmetricEigenCertification = do + let resultValue = do + matrixValue <- mkDenseDoubleMatrixRowMajor 2 2 (S.fromList [2.0, 0.0, 0.0, 3.0]) + eigenResult <- symmetricEigenPairsDenseUnchecked 2 matrixValue + case certifySymmetricEigenResult matrixValue eigenResult of + Left err -> Left (InvariantViolation ("unexpected eigen certification failure: " <> show err)) + Right certified -> Right certified + case resultValue of + Left err -> assertEqual "certification success" (Right ()) (Left err) + Right CertifiedSymmetricEigenResult {certifiedSymmetricEigenResidualReport = ResidualReport {..}} -> do + assertBool "residual scale stays certified" (residualScaled <= 1.0e7) + assertBool "orthogonality scale stays certified" (residualOrthogonalityScaled <= 1.0e7) + +testPureDenseSymmetricEigenWorkspaceOverflow :: Assertion +testPureDenseSymmetricEigenWorkspaceOverflow = + let wrappedDimension = 2 ^ (32 :: Int) + in assertEqual + "oversized eigensolver workspace" + (Left (InvariantViolation "symmetric eigen workspace cardinality exceeds Int range")) + ( symmetricEigenPairsDenseUnchecked + wrappedDimension + (trustedDenseDoubleMatrixRowMajor wrappedDimension wrappedDimension S.empty) + ) + +assertApproxStorableVector :: String -> Double -> S.Vector Double -> Either MoonlightError (S.Vector Double) -> Assertion +assertApproxStorableVector label tolerance expected actualResult = + case actualResult of + Left err -> assertEqual label (Right expected) (Left err) + Right actual -> + assertBool + label + ( S.length expected == S.length actual + && S.and (S.zipWith (\left right -> abs (left - right) <= tolerance) expected actual) + ) + +assertApproxUnboxedVector :: String -> Double -> U.Vector Double -> Either MoonlightError (U.Vector Double) -> Assertion +assertApproxUnboxedVector label tolerance expected actualResult = + case actualResult of + Left err -> assertEqual label (Right expected) (Left err) + Right actual -> + assertBool + label + ( U.length expected == U.length actual + && U.and (U.zipWith (\left right -> abs (left - right) <= tolerance) expected actual) + ) diff --git a/test/dense/DenseRowsSpec.hs b/test/dense/DenseRowsSpec.hs new file mode 100644 --- /dev/null +++ b/test/dense/DenseRowsSpec.hs @@ -0,0 +1,100 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} + +module DenseRowsSpec (tests) where + +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Dense + ( matrixRows, + matrixShape, + toListMatrix, + ) +import Moonlight.LinAlg.Dense.Rows + ( hcatRowsExact, + matrixProductRowsWith, + matrixVectorProductRowsWith, + mkDenseRows, + mkDenseRowsFromFlat, + mkDenseRowsWithShape, + transposeRowsExact, + vcatRowsExact, + ) +import Moonlight.LinAlg.Pure.Dense.Dynamic + ( dynMatrixFromRows, + dynMatrixToRows, + mkDynMatrix, + ) +import Moonlight.LinAlg.Pure.Dense.Types + ( fromListMatrix, + matrixToRows, + ) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "DenseRows" + [ testCase "mkDenseRows rejects ragged input" $ + expectLeft + (InvariantViolation "dense row matrix is ragged at row 1 (expected 2 columns, got 1)") + (mkDenseRows [[1 :: Int, 2], [3]]) + , testCase "mkDenseRowsWithShape validates row count" $ + expectLeft + (InvariantViolation "dense row matrix row count mismatch: expected 2 rows but received 1") + (mkDenseRowsWithShape 2 2 [[1 :: Int, 2]]) + , testCase "mkDenseRowsFromFlat rejects wrapped shape cardinality" $ + let wrappedDimension = 2 ^ (32 :: Int) + in expectLeft + (InvariantViolation "dense row matrix dimensions exceed Int cardinality") + (mkDenseRowsFromFlat wrappedDimension wrappedDimension ([] :: [Int])) + , testCase "matrixRows stores entries in row-major order" $ + fmap toListMatrix (matrixRows @2 @2 [[1 :: Int, 2], [3, 4]]) + @?= Right [1, 2, 3, 4] + , testCase "matrixRows rejects static column mismatch" $ + expectLeft + (InvariantViolation "dense row matrix is ragged at row 1 (expected 2 columns, got 1)") + (matrixRows @2 @2 [[1 :: Int, 2], [3]]) + , testCase "matrixRows retains columns for zero-row matrices" $ + fmap matrixShape (matrixRows @0 @3 ([] :: [[Int]])) + @?= Right (0, 3) + , testCase "matrixToRows lowers static flat matrix through DenseRows" $ + (matrixToRows =<< fromListMatrix @2 @3 @Int [1, 2, 3, 4, 5, 6]) + @?= Right [[1, 2, 3], [4, 5, 6]] + , testCase "matrixToRows preserves zero-column static row count" $ + (matrixToRows =<< fromListMatrix @2 @0 @Int []) + @?= Right [[], []] + , testCase "dynMatrixToRows lowers dynamic flat matrix through DenseRows" $ + (dynMatrixToRows =<< mkDynMatrix 2 3 [1 :: Int, 2, 3, 4, 5, 6]) + @?= Right [[1, 2, 3], [4, 5, 6]] + , testCase "dynMatrixFromRows preserves zero-column dynamic row count" $ + (dynMatrixToRows =<< dynMatrixFromRows [[], [], [] :: [Int]]) + @?= Right [[], [], []] + , testCase "transposeRowsExact preserves rectangular data" $ + transposeRowsExact [[1 :: Int, 2, 3], [4, 5, 6]] + @?= Right [[1, 4], [2, 5], [3, 6]] + , testCase "matrixVectorProductRowsWith rejects vector shape mismatch" $ + expectLeft + (InvariantViolation "dense row matrix/vector shape mismatch (matrix=(2,2), vector=1)") + (matrixVectorProductRowsWith (*) (+) (0 :: Int) [[1, 2], [3, 4]] [9]) + , testCase "matrixProductRowsWith rejects incompatible shapes" $ + expectLeft + (InvariantViolation "dense row matrix product shape mismatch (left=(1,2), right=(1,1))") + (matrixProductRowsWith (*) (+) (0 :: Int) [[1, 2]] [[3]]) + , testCase "hcatRowsExact rejects mismatched row counts" $ + expectLeft + (InvariantViolation "dense horizontal concatenation requires equal row counts, got [(1,1),(2,1)]") + (hcatRowsExact [[[1 :: Int]], [[2], [3]]]) + , testCase "vcatRowsExact rejects mismatched column counts" $ + expectLeft + (InvariantViolation "dense vertical concatenation requires equal column counts, got [(1,1),(1,2)]") + (vcatRowsExact [[[1 :: Int]], [[2, 3]]]) + ] + +expectLeft :: (Eq e, Show e) => e -> Either e a -> Assertion +expectLeft expected value = + case value of + Left err -> + err @?= expected + Right _ -> + assertFailure "expected Left, got Right" diff --git a/test/dense/DynamicSpec.hs b/test/dense/DynamicSpec.hs new file mode 100644 --- /dev/null +++ b/test/dense/DynamicSpec.hs @@ -0,0 +1,99 @@ + +module DynamicSpec + ( tests, + ) +where + +import Moonlight.LinAlg + ( DynMatrix, + dynMatrixFromRows, + dynMatrixShape, + dynMatrixToList, + dynMatrixToRows, + fromDynMatrix, + fromListMatrix, + mkDynMatrix, + toDynMatrix, + toListMatrix, + withDynMatrix, + ) +import Helpers (extractRight) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit + ( Assertion, + assertEqual, + assertFailure, + testCase, + ) + +tests :: TestTree +tests = + testGroup + "Dynamic" + [ testCase "toDynMatrix preserves shape and payload" testToDyn, + testCase "fromDynMatrix reifies static dimensions" testFromDyn, + testCase "fromDynMatrix rejects equal-cardinality shape changes" testFromDynShapeMismatch, + testCase "dynamic nested rows preserve row-major shape" testDynamicRows, + testCase "dynamic zero-column rows retain row count" testDynamicZeroColumns, + testCase "withDynMatrix introduces existential static dimensions" testWithDyn + ] + +testToDyn :: Assertion +testToDyn = + let result = do + matrixValue <- fromListMatrix @2 @3 @Double [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + pure (toDynMatrix matrixValue) + in extractRight result (\dynValue -> do + assertEqual "dynamic shape" (2, 3) (dynMatrixShape dynValue) + assertEqual "dynamic payload" [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] (dynMatrixToList dynValue) + ) + +testFromDyn :: Assertion +testFromDyn = + let result = do + dynValue <- mkDynMatrix 2 2 ([1, 2, 3, 4] :: [Integer]) + fromDynMatrix @2 @2 dynValue + in extractRight result (\matrixValue -> assertEqual "static payload" [1, 2, 3, 4] (toListMatrix matrixValue)) + +testFromDynShapeMismatch :: Assertion +testFromDynShapeMismatch = + case + do + dynValue <- mkDynMatrix 1 4 ([1, 2, 3, 4] :: [Integer]) + fromDynMatrix @2 @2 dynValue + of + Left err -> + assertEqual + "shape failure" + "InvariantViolation \"dynamic matrix shape does not match static dimensions: expected (2,2) but received (1,4)\"" + (show err) + Right _ -> + assertFailure "fromDynMatrix must not reinterpret a 1x4 matrix as 2x2" + +testDynamicRows :: Assertion +testDynamicRows = + let result = do + matrixValue <- dynMatrixFromRows [[1 :: Integer, 2], [3, 4]] + rows <- dynMatrixToRows matrixValue + pure (dynMatrixShape matrixValue, dynMatrixToList matrixValue, rows) + in extractRight result $ \(shapeValue, payload, rows) -> do + assertEqual "shape" (2, 2) shapeValue + assertEqual "payload" [1, 2, 3, 4] payload + assertEqual "rows" [[1, 2], [3, 4]] rows + +testDynamicZeroColumns :: Assertion +testDynamicZeroColumns = + let result = do + matrixValue <- dynMatrixFromRows [[], [], [] :: [Integer]] + rows <- dynMatrixToRows matrixValue + pure (dynMatrixShape matrixValue, rows) + in extractRight result $ \(shapeValue, rows) -> do + assertEqual "shape" (3, 0) shapeValue + assertEqual "rows" [[], [], []] rows + +testWithDyn :: Assertion +testWithDyn = + let result = do + dynValue :: DynMatrix Double <- mkDynMatrix 1 3 [7.0, 8.0, 9.0] + withDynMatrix dynValue (\matrixValue -> toListMatrix matrixValue) + in extractRight result (\values -> assertEqual "existential reification" [7.0, 8.0, 9.0] values) diff --git a/test/dense/ExteriorSpec.hs b/test/dense/ExteriorSpec.hs new file mode 100644 --- /dev/null +++ b/test/dense/ExteriorSpec.hs @@ -0,0 +1,155 @@ +module ExteriorSpec + ( tests, + ) +where + +import Data.Foldable (traverse_) +import Moonlight.LinAlg + ( ExteriorBasis (..), + choose, + exteriorBasis, + exteriorPowerMatrix, + ) +import Test.Tasty + ( TestTree, + testGroup, + ) +import Test.Tasty.HUnit + ( Assertion, + assertEqual, + assertFailure, + testCase, + ) + +tests :: TestTree +tests = + testGroup + "exterior algebra" + [ testCase "rank of exterior basis is choose n p" testExteriorBasisRank, + testCase "induced exterior maps compose on a diagonal fixture" testExteriorMapComposition, + testCase "integer exterior determinant keeps minor signs" testExteriorMinorSigns, + testCase "degree zero exterior map is rank-one constant" testDegreeZeroExteriorMap, + testCase "direct exterior kernels agree with recursive minors" testExteriorDirectKernels + ] + +testExteriorBasisRank :: Assertion +testExteriorBasisRank = do + basis <- expectRight (exteriorBasis 2 4) + assertEqual "Λ^2(Q^4) has choose 4 2 generators" (fromIntegral (choose 4 2)) (length (ebBasisVectors basis)) + +testExteriorMapComposition :: Assertion +testExteriorMapComposition = do + lambdaF <- expectRight (exteriorPowerMatrix 2 ([[2, 0], [0, 3]] :: [[Integer]])) + lambdaG <- expectRight (exteriorPowerMatrix 2 ([[5, 0], [0, 7]] :: [[Integer]])) + lambdaGF <- expectRight (exteriorPowerMatrix 2 ([[10, 0], [0, 21]] :: [[Integer]])) + assertEqual "Λ²(g ∘ f) multiplies the determinant scalars" (multiply1x1 lambdaG lambdaF) lambdaGF + +testExteriorMinorSigns :: Assertion +testExteriorMinorSigns = do + lambdaTwo <- expectRight (exteriorPowerMatrix 2 ([[1, 2], [3, 4]] :: [[Integer]])) + assertEqual "Λ² records the signed determinant" [[-2]] lambdaTwo + +testDegreeZeroExteriorMap :: Assertion +testDegreeZeroExteriorMap = do + lambdaZero <- expectRight (exteriorPowerMatrix 0 ([[2, 4], [6, 8], [10, 12]] :: [[Integer]])) + assertEqual "Λ⁰ is the constant rank-one map" [[1]] lambdaZero + +testExteriorDirectKernels :: Assertion +testExteriorDirectKernels = + traverse_ + assertDirectKernel + [ (2, 4, 11), + (2, 5, 23), + (3, 3, 37), + (3, 5, 41) + ] + where + assertDirectKernel (degree, rankValue, seedValue) = do + let rows = generatedIntegerRows seedValue rankValue rankValue + direct <- expectRight (exteriorPowerMatrix degree rows) + recursive <- expectRight (referenceExteriorPowerMatrix degree rows) + assertEqual ("Λ^" <> show degree <> " direct minors at rank " <> show rankValue) recursive direct + +multiply1x1 :: Num coefficient => [[coefficient]] -> [[coefficient]] -> [[coefficient]] +multiply1x1 left right = + case (left, right) of + ([[leftValue]], [[rightValue]]) -> [[leftValue * rightValue]] + _ -> [] + +expectRight :: Show failure => Either failure value -> IO value +expectRight result = + case result of + Right value -> pure value + Left failure -> assertFailure ("unexpected failure: " <> show failure) + +generatedIntegerRows :: Int -> Int -> Int -> [[Integer]] +generatedIntegerRows seedValue rowCount columnCount = + [ [ generatedIntegerEntry seedValue rowIndex columnIndex + | columnIndex <- [0 .. columnCount - 1] + ] + | rowIndex <- [0 .. rowCount - 1] + ] + +generatedIntegerEntry :: Int -> Int -> Int -> Integer +generatedIntegerEntry seedValue rowIndex columnIndex = + fromIntegral ((((seedValue + 17 * rowIndex + 31 * columnIndex + 7 * rowIndex * columnIndex) `mod` 19) - 9) :: Int) + +referenceExteriorPowerMatrix :: Int -> [[Integer]] -> Either String [[Integer]] +referenceExteriorPowerMatrix degree rows = + case rows of + [] -> Right [] + firstRow : _ -> do + targetBasis <- either (Left . show) Right (exteriorBasis degree (length rows)) + sourceBasis <- either (Left . show) Right (exteriorBasis degree (length firstRow)) + traverse + ( \targetVector -> + traverse + (\sourceVector -> referenceMinorDeterminant rows targetVector sourceVector) + (ebBasisVectors sourceBasis) + ) + (ebBasisVectors targetBasis) + +referenceMinorDeterminant :: [[Integer]] -> [Int] -> [Int] -> Either String Integer +referenceMinorDeterminant rows targetVector sourceVector = + fmap + referenceDeterminant + ( traverse + ( \targetIndex -> + traverse + ( \sourceIndex -> + safeIndex targetIndex rows + >>= safeIndex sourceIndex + ) + sourceVector + ) + targetVector + ) + +referenceDeterminant :: [[Integer]] -> Integer +referenceDeterminant matrix = + case matrix of + [] -> 1 + [singleRow] -> + case singleRow of + [value] -> value + _ -> 0 + firstRow : remainingRows -> + sum + ( fmap + (\(columnIndex, value) -> referenceSignFor columnIndex * value * referenceDeterminant (referenceRemoveColumn columnIndex remainingRows)) + (zip [0 ..] firstRow) + ) + +referenceRemoveColumn :: Int -> [[Integer]] -> [[Integer]] +referenceRemoveColumn columnIndex = + fmap (fmap snd . filter ((/= columnIndex) . fst) . zip [0 :: Int ..]) + +referenceSignFor :: Int -> Integer +referenceSignFor columnIndex = + if even columnIndex then 1 else (-1) + +safeIndex :: Int -> [a] -> Either String a +safeIndex targetIndex values = + case drop targetIndex values of + value : _ -> Right value + [] -> Left ("index out of bounds: " <> show targetIndex) diff --git a/test/dense/FieldSpec.hs b/test/dense/FieldSpec.hs new file mode 100644 --- /dev/null +++ b/test/dense/FieldSpec.hs @@ -0,0 +1,94 @@ + +module FieldSpec + ( tests, + ) +where + +import Moonlight.Core (canInvert) +import Moonlight.LinAlg + ( fromListMatrix, + gf2One, + gf2Zero, + kernel, + kernelBasisVectors, + mult, + pluDecompFullRank, + pluLower, + pluPermutation, + pluUpper, + rank, + toListMatrix, + toListVector, + ) +import Helpers (extractRight) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit + ( Assertion, + assertBool, + assertEqual, + testCase, + ) + +tests :: TestTree +tests = + testGroup + "Field" + [ testCase "pluDecompFullRank reconstructs matrix for non-singular input" testPluReconstruction, + testCase "pluDecompFullRank captures row permutations for zero-leading pivots" testPluPivoting, + testCase "rank over GF2 returns expected value" testRank, + testCase "kernel over GF2 returns basis vectors" testKernel, + testCase "GF2 rank skips all-zero column (zero is not invertible)" testGF2ZeroColumnRank, + testCase "GF2 canInvert rejects zero" testGF2CanInvertZero + ] + +testPluReconstruction :: Assertion +testPluReconstruction = + let result = do + matrixValue <- fromListMatrix @2 @2 @Double [4.0, 3.0, 6.0, 3.0] + pluValue <- pluDecompFullRank matrixValue + let permutation = pluPermutation pluValue + lower = pluLower pluValue + upper = pluUpper pluValue + lhs <- mult permutation matrixValue + reconstructed <- mult lower upper + pure (toListMatrix lhs, toListMatrix reconstructed) + in extractRight result (\(lhsValues, rhsValues) -> assertEqual "PLU reconstruction" lhsValues rhsValues) + +testPluPivoting :: Assertion +testPluPivoting = + let result = do + matrixValue <- fromListMatrix @2 @2 @Double [0.0, 1.0, 1.0, 1.0] + pluValue <- pluDecompFullRank matrixValue + let permutation = pluPermutation pluValue + lower = pluLower pluValue + upper = pluUpper pluValue + lhs <- mult permutation matrixValue + rhs <- mult lower upper + pure (toListMatrix lhs, toListMatrix rhs) + in extractRight result (\(lhsValues, rhsValues) -> assertEqual "PLU reconstruction with permutation" lhsValues rhsValues) + +testRank :: Assertion +testRank = + let result = do + matrixValue <- fromListMatrix @2 @2 [gf2One, gf2One, gf2One, gf2One] + rank matrixValue + in extractRight result (\value -> assertEqual "GF2 rank" 1 value) + +testKernel :: Assertion +testKernel = + let result = do + matrixValue <- fromListMatrix @2 @2 [gf2One, gf2One, gf2One, gf2One] + fmap kernelBasisVectors (kernel matrixValue) + in extractRight result (\basis -> assertEqual "GF2 kernel basis" [[gf2One, gf2One]] (map toListVector basis)) + +testGF2ZeroColumnRank :: Assertion +testGF2ZeroColumnRank = + let result = do + matrixValue <- fromListMatrix @2 @3 [gf2Zero, gf2One, gf2Zero, gf2Zero, gf2Zero, gf2One] + rank matrixValue + in extractRight result (\value -> assertEqual "rank of matrix with all-zero first column" 2 value) + +testGF2CanInvertZero :: Assertion +testGF2CanInvertZero = do + assertBool "GF2Zero must not be invertible" (not (canInvert gf2Zero)) + assertBool "GF2One must be invertible" (canInvert gf2One) diff --git a/test/dense/GF2Spec.hs b/test/dense/GF2Spec.hs new file mode 100644 --- /dev/null +++ b/test/dense/GF2Spec.hs @@ -0,0 +1,261 @@ +module GF2Spec + ( tests, + ) +where + +import Data.Foldable (traverse_) +import Data.Vector.Unboxed qualified as U +import Data.Vector qualified as V +import Moonlight.Core + ( MoonlightError, + ) +import Numeric.Natural (Natural) +import Moonlight.LinAlg + ( GF2 (..), + GF2MatrixEntry (..), + GF2PackedMatrixFailure (..), + GF2SparseColumn, + PackedRow, + defaultGF2SparseReducerConfig, + gf2SparseColumnRows, + gf2PackedWords, + kernelBasisGF2SparseColumns, + mkGF2SparseColumn, + mkGF2SparseReducerConfig, + mkGF2PackedMatrix, + mkGF2PackedMatrixFromRowMajor, + packedRowIndices, + rankGF2SparseColumns, + rankGF2PackedMatrix, + ) +import Test.Tasty + ( TestTree, + testGroup, + ) +import Test.Tasty.HUnit + ( Assertion, + assertEqual, + assertFailure, + testCase, + ) + +tests :: TestTree +tests = + testGroup + "GF2 packed matrix" + [ testCase "rejects out-of-bounds entries" testRejectsOutOfBounds, + testCase "rejects row-major length mismatch" testRejectsRowMajorLengthMismatch, + testCase "rejects Natural dimensions outside Int range" testRejectsOversizedNaturalDimension, + testCase "rejects wrapped row-major cardinality" testRejectsWrappedRowMajorCardinality, + testCase "duplicate entries cancel by XOR" testDuplicateEntriesCancel, + testCase "row-major rank matches known full-rank fixture" testRowMajorRank, + testCase "sparse rank agrees with packed rank on generated matrices" testSparseRankAgreement, + testCase "sparse kernel witnesses annihilate columns" testSparseKernelWitnessAnnihilation, + testCase "sparse densify threshold preserves reduction semantics" testSparseThresholdCrossing + ] + +testRejectsOutOfBounds :: Assertion +testRejectsOutOfBounds = + case mkGF2PackedMatrix 2 3 [GF2MatrixEntry 2 0] of + Left (GF2PackedMatrixEntryOutOfBounds row column rowCount columnCount) -> + assertEqual "out-of-bounds entry" (2, 0, 2, 3) (row, column, rowCount, columnCount) + Left failureValue -> + assertFailure ("unexpected packed matrix failure: " <> show failureValue) + Right _ -> + assertFailure "expected out-of-bounds packed matrix construction to fail" + +testRejectsRowMajorLengthMismatch :: Assertion +testRejectsRowMajorLengthMismatch = + case mkGF2PackedMatrixFromRowMajor 2 2 [GF2One] of + Left (GF2PackedMatrixFlatLengthMismatch expectedCount actualCount) -> + assertEqual "flat length mismatch" (4, 1) (expectedCount, actualCount) + Left failureValue -> + assertFailure ("unexpected packed matrix failure: " <> show failureValue) + Right _ -> + assertFailure "expected row-major packed matrix construction to reject malformed length" + +testRejectsOversizedNaturalDimension :: Assertion +testRejectsOversizedNaturalDimension = + let oversizedDimension = fromIntegral (maxBound :: Int) + 1 :: Natural + in case mkGF2PackedMatrix oversizedDimension 0 [] of + Left (GF2PackedMatrixCardinalityOutOfBounds rowCount columnCount) -> + assertEqual "out-of-range Natural shape" (oversizedDimension, 0) (rowCount, columnCount) + Left failureValue -> + assertFailure ("unexpected packed matrix failure: " <> show failureValue) + Right _ -> + assertFailure "expected oversized Natural dimension to fail" + +testRejectsWrappedRowMajorCardinality :: Assertion +testRejectsWrappedRowMajorCardinality = + let dimension = (2 :: Natural) ^ (32 :: Int) + in case mkGF2PackedMatrixFromRowMajor dimension dimension [] of + Left (GF2PackedMatrixCardinalityOutOfBounds rowCount columnCount) -> + assertEqual "wrapped row-major shape" (dimension, dimension) (rowCount, columnCount) + Left failureValue -> + assertFailure ("unexpected packed matrix failure: " <> show failureValue) + Right _ -> + assertFailure "expected wrapped row-major cardinality to fail" + +testDuplicateEntriesCancel :: Assertion +testDuplicateEntriesCancel = + case mkGF2PackedMatrix 1 1 [GF2MatrixEntry 0 0, GF2MatrixEntry 0 0] of + Left failureValue -> + assertFailure ("packed matrix construction failed: " <> show failureValue) + Right matrixValue -> do + assertEqual "duplicate entry rank" 0 (rankGF2PackedMatrix matrixValue) + assertEqual "duplicate entry storage" [0] (U.toList (gf2PackedWords matrixValue)) + +testRowMajorRank :: Assertion +testRowMajorRank = + case mkGF2PackedMatrixFromRowMajor 2 2 [GF2One, GF2Zero, GF2One, GF2One] of + Left failureValue -> + assertFailure ("packed matrix construction failed: " <> show failureValue) + Right matrixValue -> + assertEqual "row-major rank" 2 (rankGF2PackedMatrix matrixValue) + +testSparseRankAgreement :: Assertion +testSparseRankAgreement = + traverse_ + assertGeneratedRankAgreement + [ (0, 0, 1), + (1, 3, 2), + (4, 5, 3), + (8, 9, 5), + (17, 23, 7) + ] + +assertGeneratedRankAgreement :: (Int, Int, Int) -> Assertion +assertGeneratedRankAgreement (rowCount, columnCount, saltValue) = + case ( mkGF2PackedMatrix (fromIntegral rowCount) (fromIntegral columnCount) (generatedEntries rowCount columnCount saltValue), + generatedSparseColumns rowCount columnCount saltValue + ) of + (Left failureValue, _) -> + assertFailure ("packed matrix construction failed: " <> show failureValue) + (_, Left errorValue) -> + assertFailure ("sparse column construction failed: " <> show errorValue) + (Right packedMatrix, Right sparseColumns) -> + case rankGF2SparseColumns defaultGF2SparseReducerConfig rowCount columnCount sparseColumns of + Left errorValue -> + assertFailure ("sparse rank failed: " <> show errorValue) + Right sparseRank -> + assertEqual + ("generated sparse rank " <> show (rowCount, columnCount, saltValue)) + (rankGF2PackedMatrix packedMatrix) + sparseRank + +testSparseKernelWitnessAnnihilation :: Assertion +testSparseKernelWitnessAnnihilation = + case dependentSparseColumns of + Left errorValue -> + assertFailure ("dependent sparse columns failed: " <> show errorValue) + Right sparseColumns -> + case kernelBasisGF2SparseColumns defaultGF2SparseReducerConfig 3 3 sparseColumns of + Left errorValue -> + assertFailure ("sparse kernel basis failed: " <> show errorValue) + Right kernelBasis -> do + assertEqual "sparse kernel dependency" [[0, 1, 2]] (packedRowIndices <$> V.toList kernelBasis) + assertKernelBasisAnnihilates "dependent sparse kernel" sparseColumns kernelBasis + +testSparseThresholdCrossing :: Assertion +testSparseThresholdCrossing = + case thresholdSparseColumns of + Left errorValue -> + assertFailure ("threshold sparse columns failed: " <> show errorValue) + Right sparseColumns -> + case ( mkGF2SparseReducerConfig "threshold low fixture" 2, + mkGF2SparseReducerConfig "threshold high fixture" 99 + ) of + (Right lowConfig, Right highConfig) -> + case ( rankGF2SparseColumns lowConfig 8 4 sparseColumns, + rankGF2SparseColumns highConfig 8 4 sparseColumns, + kernelBasisGF2SparseColumns lowConfig 8 4 sparseColumns, + kernelBasisGF2SparseColumns highConfig 8 4 sparseColumns + ) of + (Right lowRank, Right highRank, Right lowKernel, Right highKernel) -> do + assertEqual "threshold rank" highRank lowRank + assertEqual "threshold kernel width" (length (V.toList highKernel)) (length (V.toList lowKernel)) + assertKernelBasisAnnihilates "low-threshold sparse kernel" sparseColumns lowKernel + resultValue -> + assertFailure ("threshold reduction failed: " <> show resultValue) + resultValue -> + assertFailure ("threshold config failed: " <> show resultValue) + +dependentSparseColumns :: Either MoonlightError (V.Vector GF2SparseColumn) +dependentSparseColumns = + V.fromList + <$> sequence + [ mkGF2SparseColumn "dependent column 0" 3 0 [0, 2], + mkGF2SparseColumn "dependent column 1" 3 1 [1], + mkGF2SparseColumn "dependent column 2" 3 2 [0, 1, 2] + ] + +thresholdSparseColumns :: Either MoonlightError (V.Vector GF2SparseColumn) +thresholdSparseColumns = + V.fromList + <$> sequence + [ mkGF2SparseColumn "threshold column 0" 8 0 [0, 1, 2, 3], + mkGF2SparseColumn "threshold column 1" 8 1 [2, 3, 4, 5], + mkGF2SparseColumn "threshold column 2" 8 2 [0, 1, 4, 5], + mkGF2SparseColumn "threshold column 3" 8 3 [6, 7] + ] + +generatedSparseColumns :: Int -> Int -> Int -> Either MoonlightError (V.Vector GF2SparseColumn) +generatedSparseColumns rowCount columnCount saltValue = + V.fromList + <$> traverse + ( \columnIndex -> + mkGF2SparseColumn + ("generated sparse column " <> show columnIndex) + rowCount + columnIndex + (generatedSupport rowCount columnIndex saltValue) + ) + [0 .. columnCount - 1] + +generatedEntries :: Int -> Int -> Int -> [GF2MatrixEntry] +generatedEntries rowCount columnCount saltValue = + [ GF2MatrixEntry rowIndex columnIndex + | columnIndex <- [0 .. columnCount - 1], + rowIndex <- generatedSupport rowCount columnIndex saltValue + ] + +generatedSupport :: Int -> Int -> Int -> [Int] +generatedSupport rowCount columnIndex saltValue = + [ rowIndex + | rowIndex <- [0 .. rowCount - 1], + generatedBit rowIndex columnIndex saltValue + ] + +generatedBit :: Int -> Int -> Int -> Bool +generatedBit rowIndex columnIndex saltValue = + rowIndex == columnIndex + || ((rowIndex * 17 + columnIndex * 31 + saltValue * 13 + rowIndex * columnIndex) `mod` 11 == 0) + +assertKernelBasisAnnihilates :: String -> V.Vector GF2SparseColumn -> V.Vector PackedRow -> Assertion +assertKernelBasisAnnihilates label sparseColumns kernelBasis = + traverse_ + (assertKernelWitnessAnnihilates label sparseColumns) + (packedRowIndices <$> V.toList kernelBasis) + +assertKernelWitnessAnnihilates :: String -> V.Vector GF2SparseColumn -> [Int] -> Assertion +assertKernelWitnessAnnihilates label sparseColumns witnessColumns = + case traverse (`lookupSparseColumnRows` sparseColumns) witnessColumns of + Nothing -> + assertFailure (label <> ": kernel witness referenced an absent column") + Just supports -> + assertEqual (label <> ": annihilated support") [] (foldl' xorSortedSupports [] supports) + +lookupSparseColumnRows :: Int -> V.Vector GF2SparseColumn -> Maybe [Int] +lookupSparseColumnRows columnIndex sparseColumns = + gf2SparseColumnRows <$> (sparseColumns V.!? columnIndex) + +xorSortedSupports :: [Int] -> [Int] -> [Int] +xorSortedSupports leftRows rightRows = + case (leftRows, rightRows) of + ([], _) -> rightRows + (_, []) -> leftRows + (leftRow : remainingLeft, rightRow : remainingRight) -> + case compare leftRow rightRow of + LT -> leftRow : xorSortedSupports remainingLeft rightRows + EQ -> xorSortedSupports remainingLeft remainingRight + GT -> rightRow : xorSortedSupports leftRows remainingRight diff --git a/test/dense/SymmetricSpec.hs b/test/dense/SymmetricSpec.hs new file mode 100644 --- /dev/null +++ b/test/dense/SymmetricSpec.hs @@ -0,0 +1,524 @@ +module SymmetricSpec + ( tests, + ) +where + +import Moonlight.Algebra.Pure.Module (BilinearSpace (..)) +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg + ( DiagonalizedSymmetric2 (..), + DiagonalizedSymmetric3 (..), + Symmetric2 (..), + Symmetric3 (..), + Vec2 (..), + Vec3 (..), + applySymmetric2, + applySymmetric3, + diagonalSymmetric2, + diagonalSymmetric3, + diagonalizedSymmetric2ToTensor, + diagonalizedSymmetric2ToVec2, + diagonalizedSymmetric3ToTensor, + diagonalizedSymmetric3ToVec3, + eigendecomposeSymmetric2, + eigendecomposeSymmetric2With, + eigendecomposeSymmetric3, + eigendecomposeSymmetric3With, + outerSymmetric2, + outerSymmetric3, + symmetric2Entries, + symmetric2ToMatrix, + symmetric3Entries, + symmetric3ToMatrix, + toListMatrix, + toListVector, + vec2FromList, + vec3FromList, + ) +import Helpers (extractRight) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertEqual, assertFailure, testCase) +import Test.Tasty.QuickCheck qualified as QC + +tests :: TestTree +tests = + testGroup + "Symmetric" + [ testCase "Symmetric2 accumulates compact entries componentwise" testSymmetric2Monoid, + testCase "symmetric2ToMatrix expands the upper triangle into a dense matrix" testSymmetric2DenseExpansion, + testCase "Symmetric2 bilinearForm uses the Frobenius form of the represented matrix" testSymmetric2FrobeniusInnerProduct, + testCase "eigendecomposeSymmetric2 diagonalizes a symmetric tensor" testSymmetric2EigenDecomposition, + QC.testProperty "eigendecomposeSymmetric2 reconstructs generated near-degenerate tensors" propSymmetric2GeneratedReconstruction, + testCase "eigendecomposeSymmetric2With lifts the spectrum into a reusable diagonalized form" testSymmetric2DiagonalizedDecomposition, + testCase "outerSymmetric2 and applySymmetric2 agree on the induced linear map" testSymmetric2Apply, + testCase "vec2FromList rejects non-exact payloads" testVec2FromListRejectsNonExactPayload, + testCase "Semigroup and Monoid accumulate compact entries componentwise" testMonoid, + testCase "symmetric3ToMatrix expands the upper triangle into a dense matrix" testDenseExpansion, + testCase "bilinearForm uses the Frobenius form of the represented matrix" testFrobeniusInnerProduct, + testCase "eigendecomposeSymmetric3 diagonalizes a symmetric tensor" testEigenDecomposition, + QC.testProperty "eigendecomposeSymmetric3 reconstructs generated near-degenerate tensors" propSymmetric3GeneratedReconstruction, + testCase "eigendecomposeSymmetric3With lifts the spectrum into a reusable diagonalized form" testDiagonalizedDecomposition, + testCase "outerSymmetric3 and applySymmetric3 agree on the induced linear map" testApply, + testCase "vec3FromList rejects non-exact payloads" testVec3FromListRejectsNonExactPayload + ] + +closeTo :: Double -> Double -> Double -> Bool +closeTo tolerance expected actual = abs (expected - actual) <= tolerance + +data GeneratedSymmetric2Case = GeneratedSymmetric2Case + { generatedSymmetric2Tensor :: !(Symmetric2 Double), + generatedSymmetric2Scale :: !Double + } + deriving stock (Show) + +instance QC.Arbitrary GeneratedSymmetric2Case where + arbitrary = do + scaleValue <- generatedScale + baseValue <- QC.choose (-3.0, 3.0) + gapValue <- generatedEigenvalueGap + angleValue <- QC.choose (-pi, pi) + let cosineValue = cos angleValue + sineValue = sin angleValue + firstAxis = Vec2 cosineValue sineValue + secondAxis = Vec2 (-sineValue) cosineValue + firstEigenvalue = scaleValue * (baseValue + gapValue) + secondEigenvalue = scaleValue * (baseValue - gapValue) + pure + GeneratedSymmetric2Case + { generatedSymmetric2Tensor = + symmetric2FromEigenFrame firstEigenvalue secondEigenvalue firstAxis secondAxis, + generatedSymmetric2Scale = scaleValue + } + +data GeneratedSymmetric3Case = GeneratedSymmetric3Case + { generatedSymmetric3Tensor :: !(Symmetric3 Double), + generatedSymmetric3Scale :: !Double + } + deriving stock (Show) + +instance QC.Arbitrary GeneratedSymmetric3Case where + arbitrary = do + scaleValue <- generatedScale + baseValue <- QC.choose (-3.0, 3.0) + firstGap <- generatedEigenvalueGap + secondGap <- generatedEigenvalueGap + spectrumShape <- QC.elements [0 :: Int, 1, 2, 3] + thetaValue <- QC.choose (-pi, pi) + phiValue <- QC.choose (-pi, pi) + let (firstAxis, secondAxis, thirdAxis) = generatedOrthonormalFrame thetaValue phiValue + (firstEigenvalue, secondEigenvalue, thirdEigenvalue) = + case spectrumShape of + 0 -> (baseValue, baseValue, baseValue) + 1 -> (baseValue + firstGap, baseValue, baseValue - firstGap) + 2 -> (baseValue + 1.0, baseValue + firstGap, baseValue) + _ -> (baseValue + 1.0, baseValue + firstGap, baseValue - secondGap) + pure + GeneratedSymmetric3Case + { generatedSymmetric3Tensor = + symmetric3FromEigenFrame + (scaleValue * firstEigenvalue) + (scaleValue * secondEigenvalue) + (scaleValue * thirdEigenvalue) + firstAxis + secondAxis + thirdAxis, + generatedSymmetric3Scale = scaleValue + } + +generatedScale :: QC.Gen Double +generatedScale = + QC.elements [1.0e-12, 1.0e-6, 1.0, 1.0e6, 1.0e12] + +generatedEigenvalueGap :: QC.Gen Double +generatedEigenvalueGap = + QC.elements [0.0, 1.0e-12, 1.0e-10, 1.0e-8, 1.0e-4, 1.0] + +generatedOrthonormalFrame :: Double -> Double -> (Vec3, Vec3, Vec3) +generatedOrthonormalFrame thetaValue phiValue = + let cosineTheta = cos thetaValue + sineTheta = sin thetaValue + cosinePhi = cos phiValue + sinePhi = sin phiValue + in ( Vec3 cosineTheta sineTheta 0.0, + Vec3 (-sineTheta * cosinePhi) (cosineTheta * cosinePhi) sinePhi, + Vec3 (sineTheta * sinePhi) (-cosineTheta * sinePhi) cosinePhi + ) + +symmetric2FromEigenFrame :: Double -> Double -> Vec2 -> Vec2 -> Symmetric2 Double +symmetric2FromEigenFrame firstEigenvalue secondEigenvalue firstAxis secondAxis = + outerSymmetric2 firstEigenvalue firstAxis <> outerSymmetric2 secondEigenvalue secondAxis + +symmetric3FromEigenFrame :: Double -> Double -> Double -> Vec3 -> Vec3 -> Vec3 -> Symmetric3 Double +symmetric3FromEigenFrame firstEigenvalue secondEigenvalue thirdEigenvalue firstAxis secondAxis thirdAxis = + outerSymmetric3 firstEigenvalue firstAxis + <> outerSymmetric3 secondEigenvalue secondAxis + <> outerSymmetric3 thirdEigenvalue thirdAxis + +propSymmetric2GeneratedReconstruction :: GeneratedSymmetric2Case -> QC.Property +propSymmetric2GeneratedReconstruction generatedCase = + case eigendecomposeSymmetric2 (generatedSymmetric2Tensor generatedCase) of + Left err -> + QC.counterexample ("unexpected symmetric2 decomposition failure: " <> show err) False + Right (eigenvalues, eigenvectors) -> + case reconstructSymmetric2 (toListVector eigenvalues) (toListMatrix eigenvectors) of + Nothing -> + QC.counterexample "symmetric2 decomposition returned malformed carriers" False + Just reconstructedTensor -> + QC.counterexample + ( "symmetric2 reconstruction=" + <> show reconstructedTensor + <> ", expected=" + <> show (generatedSymmetric2Tensor generatedCase) + <> ", scale=" + <> show (generatedSymmetric2Scale generatedCase) + ) + ( symmetric2Close (generatedSymmetric2Tensor generatedCase) reconstructedTensor + && orthonormal2 (toListMatrix eigenvectors) + ) + +propSymmetric3GeneratedReconstruction :: GeneratedSymmetric3Case -> QC.Property +propSymmetric3GeneratedReconstruction generatedCase = + case eigendecomposeSymmetric3 (generatedSymmetric3Tensor generatedCase) of + Left err -> + QC.counterexample ("unexpected symmetric3 decomposition failure: " <> show err) False + Right (eigenvalues, eigenvectors) -> + case reconstructSymmetric3 (toListVector eigenvalues) (toListMatrix eigenvectors) of + Nothing -> + QC.counterexample "symmetric3 decomposition returned malformed carriers" False + Just reconstructedTensor -> + QC.counterexample + ( "symmetric3 reconstruction=" + <> show reconstructedTensor + <> ", expected=" + <> show (generatedSymmetric3Tensor generatedCase) + <> ", scale=" + <> show (generatedSymmetric3Scale generatedCase) + ) + ( symmetric3Close (generatedSymmetric3Tensor generatedCase) reconstructedTensor + && orthonormal3 (toListMatrix eigenvectors) + ) + +reconstructSymmetric2 :: [Double] -> [Double] -> Maybe (Symmetric2 Double) +reconstructSymmetric2 eigenvalues eigenvectors = + case (eigenvalues, eigenvectors) of + ([firstEigenvalue, secondEigenvalue], [x1, x2, y1, y2]) -> + Just + ( symmetric2FromEigenFrame + firstEigenvalue + secondEigenvalue + (Vec2 x1 y1) + (Vec2 x2 y2) + ) + _ -> Nothing + +reconstructSymmetric3 :: [Double] -> [Double] -> Maybe (Symmetric3 Double) +reconstructSymmetric3 eigenvalues eigenvectors = + case (eigenvalues, eigenvectors) of + ([firstEigenvalue, secondEigenvalue, thirdEigenvalue], [x1, x2, x3, y1, y2, y3, z1, z2, z3]) -> + Just + ( symmetric3FromEigenFrame + firstEigenvalue + secondEigenvalue + thirdEigenvalue + (Vec3 x1 y1 z1) + (Vec3 x2 y2 z2) + (Vec3 x3 y3 z3) + ) + _ -> Nothing + +symmetric2Close :: Symmetric2 Double -> Symmetric2 Double -> Bool +symmetric2Close expectedTensor actualTensor = + symmetric2MaxAbsDifference expectedTensor actualTensor <= 1.0e-7 * max 1.0 (symmetric2MaxAbs expectedTensor) + +symmetric3Close :: Symmetric3 Double -> Symmetric3 Double -> Bool +symmetric3Close expectedTensor actualTensor = + symmetric3MaxAbsDifference expectedTensor actualTensor <= 1.0e-7 * max 1.0 (symmetric3MaxAbs expectedTensor) + +symmetric2MaxAbs :: Symmetric2 Double -> Double +symmetric2MaxAbs tensorValue = + maximum (abs <$> [sym2XX tensorValue, sym2XY tensorValue, sym2YY tensorValue]) + +symmetric3MaxAbs :: Symmetric3 Double -> Double +symmetric3MaxAbs tensorValue = + maximum (abs <$> [sym3XX tensorValue, sym3XY tensorValue, sym3XZ tensorValue, sym3YY tensorValue, sym3YZ tensorValue, sym3ZZ tensorValue]) + +symmetric2MaxAbsDifference :: Symmetric2 Double -> Symmetric2 Double -> Double +symmetric2MaxAbsDifference expectedTensor actualTensor = + maximum + ( abs + <$> [ sym2XX expectedTensor - sym2XX actualTensor, + sym2XY expectedTensor - sym2XY actualTensor, + sym2YY expectedTensor - sym2YY actualTensor + ] + ) + +symmetric3MaxAbsDifference :: Symmetric3 Double -> Symmetric3 Double -> Double +symmetric3MaxAbsDifference expectedTensor actualTensor = + maximum + ( abs + <$> [ sym3XX expectedTensor - sym3XX actualTensor, + sym3XY expectedTensor - sym3XY actualTensor, + sym3XZ expectedTensor - sym3XZ actualTensor, + sym3YY expectedTensor - sym3YY actualTensor, + sym3YZ expectedTensor - sym3YZ actualTensor, + sym3ZZ expectedTensor - sym3ZZ actualTensor + ] + ) + +orthonormal2 :: [Double] -> Bool +orthonormal2 eigenvectors = + case eigenvectors of + [x1, x2, y1, y2] -> + closeTo 1.0e-8 1.0 (x1 * x1 + y1 * y1) + && closeTo 1.0e-8 1.0 (x2 * x2 + y2 * y2) + && closeTo 1.0e-8 0.0 (x1 * x2 + y1 * y2) + _ -> False + +orthonormal3 :: [Double] -> Bool +orthonormal3 eigenvectors = + case eigenvectors of + [x1, x2, x3, y1, y2, y3, z1, z2, z3] -> + closeTo 1.0e-8 1.0 (x1 * x1 + y1 * y1 + z1 * z1) + && closeTo 1.0e-8 1.0 (x2 * x2 + y2 * y2 + z2 * z2) + && closeTo 1.0e-8 1.0 (x3 * x3 + y3 * y3 + z3 * z3) + && closeTo 1.0e-8 0.0 (x1 * x2 + y1 * y2 + z1 * z2) + && closeTo 1.0e-8 0.0 (x1 * x3 + y1 * y3 + z1 * z3) + && closeTo 1.0e-8 0.0 (x2 * x3 + y2 * y3 + z2 * z3) + _ -> False + +testSymmetric2Monoid :: IO () +testSymmetric2Monoid = + let leftValue = + ( Symmetric2 + { sym2XX = 1.0, + sym2XY = 2.0, + sym2YY = 3.0 + } :: + Symmetric2 Double + ) + rightValue = + ( Symmetric2 + { sym2XX = 0.5, + sym2XY = -1.5, + sym2YY = 0.25 + } :: + Symmetric2 Double + ) + in assertEqual + "compact storage should add componentwise" + [1.5, 0.5, 0.5, 3.25] + (symmetric2Entries (leftValue <> rightValue <> mempty)) + +testSymmetric2DenseExpansion :: IO () +testSymmetric2DenseExpansion = + let tensorValue = + ( Symmetric2 + { sym2XX = 1.0, + sym2XY = 2.0, + sym2YY = 4.0 + } :: + Symmetric2 Double + ) + in extractRight + (symmetric2ToMatrix tensorValue) + (\matrixValue -> + assertEqual + "dense expansion should mirror the upper triangle" + [1.0, 2.0, 2.0, 4.0] + (toListMatrix matrixValue) + ) + +testSymmetric2FrobeniusInnerProduct :: IO () +testSymmetric2FrobeniusInnerProduct = + let leftValue = + ( Symmetric2 + { sym2XX = 1.0, + sym2XY = 2.0, + sym2YY = 3.0 + } :: + Symmetric2 Double + ) + rightValue = + ( Symmetric2 + { sym2XX = 2.0, + sym2XY = 1.5, + sym2YY = 1.0 + } :: + Symmetric2 Double + ) + expectedValue = 2.0 + 3.0 + 2.0 * 3.0 + in assertBool + "off-diagonal components should count twice under the Frobenius inner product" + (closeTo 1.0e-9 expectedValue (bilinearForm leftValue rightValue)) + +testSymmetric2EigenDecomposition :: IO () +testSymmetric2EigenDecomposition = + extractRight + (eigendecomposeSymmetric2 (diagonalSymmetric2 2.0 5.0)) + (\(eigenvalues, eigenvectors) -> do + assertEqual "eigenvalues should be sorted descending" [5.0, 2.0] (toListVector eigenvalues) + assertEqual + "eigenvectors should be the canonical basis for a diagonal tensor" + [0.0, 1.0, 1.0, 0.0] + (toListMatrix eigenvectors) + ) + +testSymmetric2DiagonalizedDecomposition :: IO () +testSymmetric2DiagonalizedDecomposition = + extractRight + (eigendecomposeSymmetric2With (Just . length) 0 (diagonalSymmetric2 2.0 5.0)) + (\diagonalizedValue -> do + assertEqual + "the lifted decomposition should expose sorted eigenvalues" + (DiagonalizedSymmetric2 5.0 2.0 4) + diagonalizedValue + assertEqual + "the lifted decomposition should reconstruct the diagonal tensor" + (diagonalSymmetric2 5.0 2.0 :: Symmetric2 Double) + (diagonalizedSymmetric2ToTensor diagonalizedValue) + assertEqual + "the lifted decomposition should expose the diagonal as a Vec2" + (Vec2 5.0 2.0) + (diagonalizedSymmetric2ToVec2 diagonalizedValue) + ) + +testSymmetric2Apply :: IO () +testSymmetric2Apply = + let tensorValue = outerSymmetric2 2.0 (Vec2 1.0 (-1.0)) + actualValue = applySymmetric2 tensorValue (Vec2 3.0 1.0) + in assertBool + "outerSymmetric2 should produce the expected rank-one linear map" + (actualValue == Vec2 4.0 (-4.0)) + +testVec2FromListRejectsNonExactPayload :: IO () +testVec2FromListRejectsNonExactPayload = + case vec2FromList [1.0, 2.0, 3.0] of + Left (InvariantViolation _) -> pure () + Left err -> assertFailure ("expected Vec2 shape error, got " <> show err) + Right value -> assertFailure ("expected Vec2 constructor failure, got " <> show value) + +testMonoid :: IO () +testMonoid = + let leftValue = + ( Symmetric3 + { sym3XX = 1.0, + sym3XY = 2.0, + sym3XZ = 3.0, + sym3YY = 4.0, + sym3YZ = 5.0, + sym3ZZ = 6.0 + } :: + Symmetric3 Double + ) + rightValue = + ( Symmetric3 + { sym3XX = 0.5, + sym3XY = 1.5, + sym3XZ = -2.0, + sym3YY = 0.25, + sym3YZ = 0.75, + sym3ZZ = 1.25 + } :: + Symmetric3 Double + ) + in assertEqual + "compact storage should add componentwise" + [1.5, 3.5, 1.0, 3.5, 4.25, 5.75, 1.0, 5.75, 7.25] + (symmetric3Entries (leftValue <> rightValue <> mempty)) + +testDenseExpansion :: IO () +testDenseExpansion = + let tensorValue = + ( Symmetric3 + { sym3XX = 1.0, + sym3XY = 2.0, + sym3XZ = 3.0, + sym3YY = 4.0, + sym3YZ = 5.0, + sym3ZZ = 6.0 + } :: + Symmetric3 Double + ) + in extractRight + (symmetric3ToMatrix tensorValue) + (\matrixValue -> + assertEqual + "dense expansion should mirror the upper triangle" + [1.0, 2.0, 3.0, 2.0, 4.0, 5.0, 3.0, 5.0, 6.0] + (toListMatrix matrixValue) + ) + +testFrobeniusInnerProduct :: IO () +testFrobeniusInnerProduct = + let leftValue = + ( Symmetric3 + { sym3XX = 1.0, + sym3XY = 2.0, + sym3XZ = 0.0, + sym3YY = 3.0, + sym3YZ = 4.0, + sym3ZZ = 5.0 + } :: + Symmetric3 Double + ) + rightValue = + ( Symmetric3 + { sym3XX = 2.0, + sym3XY = 1.5, + sym3XZ = 0.0, + sym3YY = 1.0, + sym3YZ = 0.5, + sym3ZZ = 4.0 + } :: + Symmetric3 Double + ) + expectedValue = 2.0 + 3.0 + 20.0 + 2.0 * (3.0 + 0.0 + 2.0) + in assertBool + "off-diagonal components should count twice under the Frobenius inner product" + (closeTo 1.0e-9 expectedValue (bilinearForm leftValue rightValue)) + +testEigenDecomposition :: IO () +testEigenDecomposition = + extractRight + (eigendecomposeSymmetric3 (diagonalSymmetric3 2.0 3.0 5.0)) + (\(eigenvalues, eigenvectors) -> do + assertEqual "eigenvalues should be sorted descending" [5.0, 3.0, 2.0] (toListVector eigenvalues) + assertEqual + "eigenvectors should be the canonical basis for a diagonal tensor" + [0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0] + (toListMatrix eigenvectors) + ) + +testDiagonalizedDecomposition :: IO () +testDiagonalizedDecomposition = + extractRight + (eigendecomposeSymmetric3With (const Nothing) "fallback" (diagonalSymmetric3 2.0 3.0 5.0)) + (\diagonalizedValue -> do + assertEqual + "the lifted decomposition should expose sorted eigenvalues and fallback axes" + (DiagonalizedSymmetric3 5.0 3.0 2.0 "fallback") + diagonalizedValue + assertEqual + "the lifted decomposition should reconstruct the diagonal tensor" + (diagonalSymmetric3 5.0 3.0 2.0 :: Symmetric3 Double) + (diagonalizedSymmetric3ToTensor diagonalizedValue) + assertEqual + "the lifted decomposition should expose the diagonal as a Vec3" + (Vec3 5.0 3.0 2.0) + (diagonalizedSymmetric3ToVec3 diagonalizedValue) + ) + +testApply :: IO () +testApply = + let tensorValue = outerSymmetric3 2.0 (Vec3 1.0 2.0 (-1.0)) + actualValue = applySymmetric3 tensorValue (Vec3 3.0 0.0 1.0) + in assertBool + "outerSymmetric3 should produce the expected rank-one linear map" + (actualValue == Vec3 4.0 8.0 (-4.0)) + +testVec3FromListRejectsNonExactPayload :: IO () +testVec3FromListRejectsNonExactPayload = + case vec3FromList [1.0, 2.0] of + Left (InvariantViolation _) -> pure () + Left err -> assertFailure ("expected Vec3 shape error, got " <> show err) + Right value -> assertFailure ("expected Vec3 constructor failure, got " <> show value) diff --git a/test/domain/DomainSpec.hs b/test/domain/DomainSpec.hs new file mode 100644 --- /dev/null +++ b/test/domain/DomainSpec.hs @@ -0,0 +1,504 @@ + +module DomainSpec + ( tests, + ) +where + +import Control.Monad (foldM) +import Data.Foldable (traverse_) +import Data.List (mapAccumL) +import Data.Proxy (Proxy (..)) +import GHC.TypeNats (KnownNat, natVal) +import Moonlight.Core (MoonlightError) +import Moonlight.LinAlg + ( bareissDeterminant, + bareissRank, + exteriorPowerMatrix, + fromListMatrix, + mult, + rank, + smithDiagonal, + smithDiagonalForm, + smithDiagonalMatrix, + smithLeft, + smithLeftInverse, + smithNormalForm, + smithRight, + smithRightInverse, + toListMatrix, + ) +import Moonlight.LinAlg.Pure.Domain.Smith.Multimodular (smithDiagonalFormMultimodular) +import Moonlight.LinAlg.Pure.Domain.Smith.Witnessed (smithNormalFormWitnessed) +import Helpers (extractRight) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit + ( Assertion, + assertBool, + assertEqual, + assertFailure, + testCase, + ) + +tests :: TestTree +tests = + testGroup + "Domain" + [ testCase "smithNormalForm returns diagonal matrix for diagonal input" testSmithDiagonal, + testCase "smithNormalForm clears off-diagonal entries on simple integer matrix" testSmithClearsOffDiagonal, + testCase "smithNormalForm enforces divisibility chain" testSmithDivisibilityChain, + testCase "smithNormalForm witness reconstruction" testSmithWitness, + testCase "smithNormalForm inverse witnesses survive row and column reductions" testSmithWitnessInversesRowColumn, + testCase "smithNormalForm inverse witnesses survive divisibility repair" testSmithWitnessInversesDivisibilityRepair, + testCase "smithDiagonalForm agrees with full Smith invariant factors" testSmithDiagonalOnlyAgreesWithFull, + testCase "smithDiagonalForm rectangular multimodular fixture agrees with full Smith" testSmithDiagonalRectangularMultimodular, + testCase "smithDiagonalForm rank-deficient multimodular fixture agrees with full Smith" testSmithDiagonalRankDeficientMultimodular, + testCase "smithDiagonalForm torsion-rich multimodular fixture agrees with full Smith" testSmithDiagonalTorsionRichMultimodular, + testCase "smithDiagonalForm adversarial large-determinant fixture agrees with full Smith" testSmithDiagonalLargeDeterminantMultimodular, + testCase "smithNormalForm invariant factors match the determinantal divisors of the 3x3 minors" testSmithDeterminantalDivisorsThreeByThree, + testCase "smithNormalForm invariant factors match the determinantal divisors of the 4x4 minors" testSmithDeterminantalDivisorsFourByFour, + testCase "smithNormalForm degenerate shapes yield an empty invariant-factor list" testSmithZeroDimensionalShapes, + testCase "smithNormalFormWitnessed rectangular fixture reconstructs and agrees with multimodular diagonal" testSmithWitnessedRectangular, + testCase "smithNormalFormWitnessed rank-deficient fixture reconstructs and agrees with multimodular diagonal" testSmithWitnessedRankDeficient, + testCase "smithNormalFormWitnessed torsion-rich fixture reconstructs and agrees with multimodular diagonal" testSmithWitnessedTorsionRich, + testCase "smithNormalFormWitnessed adversarial large-entry fixture reconstructs and agrees with multimodular diagonal" testSmithWitnessedLargeEntry, + testCase "smithNormalFormWitnessed nonsingular fast-path fixture reconstructs and agrees with multimodular diagonal" testSmithWitnessedFastPath, + testCase "Bareiss rank agrees with Rational field rank" testBareissRankAgreesWithRationalRank, + testCase "Bareiss determinant agrees with Rational exterior determinant" testBareissDeterminantAgreesWithRationalDeterminant, + testCase "smithNormalForm identity matrix" testSmithIdentity, + testCase "smithNormalForm zero matrix" testSmithZero + ] + +testSmithDiagonal :: Assertion +testSmithDiagonal = + let result = do + matrixValue <- fromListMatrix @2 @2 [2 :: Integer, 0, 0, 4] + fmap (toListMatrix . smithDiagonal) (smithNormalForm matrixValue) + in extractRight result (\values -> assertEqual "smith diagonal" [2, 0, 0, 4] values) + +testSmithClearsOffDiagonal :: Assertion +testSmithClearsOffDiagonal = + let result = do + matrixValue <- fromListMatrix @2 @2 [2 :: Integer, 4, 0, 2] + fmap (toListMatrix . smithDiagonal) (smithNormalForm matrixValue) + in extractRight result assertTwoByTwoOffDiagonalZero + +testSmithDivisibilityChain :: Assertion +testSmithDivisibilityChain = + let result = do + matrixValue <- fromListMatrix @2 @2 [6 :: Integer, 0, 0, 4] + smithValue <- smithNormalForm matrixValue + pure (toListMatrix (smithDiagonal smithValue)) + in extractRight result assertTwoByTwoDiagonalDivisibility + +testSmithWitness :: Assertion +testSmithWitness = + let result = do + matrixValue <- fromListMatrix @2 @2 [2 :: Integer, 4, 0, 2] + smithValue <- smithNormalForm matrixValue + let leftMatrix = smithLeft smithValue + diagonal = smithDiagonal smithValue + rightMatrix = smithRight smithValue + la <- mult leftMatrix matrixValue + lar <- mult la rightMatrix + let diagValues = toListMatrix diagonal + pure (diagValues, toListMatrix lar) + in extractRight result (\(diagValues, reconstructed) -> + assertEqual "L * A * R must equal diagonal" diagValues reconstructed) + +testSmithWitnessInversesRowColumn :: Assertion +testSmithWitnessInversesRowColumn = + assertSmithWitnessInverses "row and column reductions" [2, 4, 6, 8] + +testSmithWitnessInversesDivisibilityRepair :: Assertion +testSmithWitnessInversesDivisibilityRepair = + assertSmithWitnessInverses "divisibility repair" [6, 0, 0, 4] + +testSmithDiagonalOnlyAgreesWithFull :: Assertion +testSmithDiagonalOnlyAgreesWithFull = + traverse_ + assertDiagonalAgreement + [ generatedIntegerEntries 3 3 11, + generatedIntegerEntries 3 3 29, + [2, 0, 0, 0, 6, 0, 0, 0, 0], + [2, 4, 6, 1, 2, 3, 0, 0, 0] + ] + where + assertDiagonalAgreement entries = + let result = do + matrixValue <- fromListMatrix @3 @3 @Integer entries + fullValue <- smithNormalForm matrixValue + diagonalOnly <- smithDiagonalForm matrixValue + pure (toListMatrix (smithDiagonal fullValue), toListMatrix (smithDiagonalMatrix diagonalOnly)) + in extractRight result $ + \(fullDiagonal, diagonalOnly) -> + assertEqual "diagonal-only Smith factors" fullDiagonal diagonalOnly + +testSmithDiagonalRectangularMultimodular :: Assertion +testSmithDiagonalRectangularMultimodular = + assertSmithDiagonalAgreement (Proxy @3) (Proxy @4) "rectangular multimodular Smith diagonal" [6, 10, 14, 22, 9, 15, 21, 33, 3, 5, 7, 11] + +testSmithDiagonalRankDeficientMultimodular :: Assertion +testSmithDiagonalRankDeficientMultimodular = + assertSmithDiagonalAgreement (Proxy @4) (Proxy @4) "rank-deficient multimodular Smith diagonal" [4, 8, 12, 16, 6, 12, 18, 24, 10, 20, 30, 40, 0, 0, 0, 0] + +testSmithDiagonalTorsionRichMultimodular :: Assertion +testSmithDiagonalTorsionRichMultimodular = + assertSmithDiagonalAgreement (Proxy @4) (Proxy @4) "torsion-rich multimodular Smith diagonal" [12, 18, 30, 42, 0, 36, 54, 78, 0, 0, 90, 126, 6, 0, 0, 210] + +testSmithDiagonalLargeDeterminantMultimodular :: Assertion +testSmithDiagonalLargeDeterminantMultimodular = + assertSmithDiagonalAgreement (Proxy @3) (Proxy @3) "large-determinant multimodular Smith diagonal" [4294967296, 0, 0, 0, 4294967296, 0, 0, 0, 4294967296] + +-- The determinantal divisors are an oracle independent of every Smith route: +-- Delta_k is the gcd of all k x k minors, and d_k = Delta_k / Delta_(k-1). +testSmithDeterminantalDivisorsThreeByThree :: Assertion +testSmithDeterminantalDivisorsThreeByThree = + traverse_ + assertThreeByThreeCertificate + [ ("3x3 generated seed 11", generatedIntegerEntries 3 3 11), + ("3x3 generated seed 29", generatedIntegerEntries 3 3 29), + ("3x3 torsion diagonal", [2, 0, 0, 0, 6, 0, 0, 0, 0]), + ("3x3 rank-deficient", [2, 4, 6, 1, 2, 3, 0, 0, 0]), + ("3x3 large determinant", [4294967296, 0, 0, 0, 4294967296, 0, 0, 0, 4294967296]) + ] + where + assertThreeByThreeCertificate (label, entries) = + let result = do + matrixValue <- fromListMatrix @3 @3 @Integer entries + fullValue <- smithNormalForm matrixValue + firstDivisor <- minorDeterminantGcd (Proxy @1) 3 3 entries + secondDivisor <- minorDeterminantGcd (Proxy @2) 3 3 entries + thirdDivisor <- minorDeterminantGcd (Proxy @3) 3 3 entries + pure + ( invariantFactorsFromDivisors [firstDivisor, secondDivisor, thirdDivisor], + diagonalEntriesOf 3 3 (toListMatrix (smithDiagonal fullValue)) + ) + in extractRight result $ + \(determinantalFactors, smithFactors) -> + assertEqual (label <> ": d_k = Delta_k / Delta_(k-1)") determinantalFactors smithFactors + +testSmithDeterminantalDivisorsFourByFour :: Assertion +testSmithDeterminantalDivisorsFourByFour = + traverse_ + assertFourByFourCertificate + [ ("4x4 rank-deficient", [4, 8, 12, 16, 6, 12, 18, 24, 10, 20, 30, 40, 0, 0, 0, 0]), + ("4x4 torsion-rich", [12, 18, 30, 42, 0, 36, 54, 78, 0, 0, 90, 126, 6, 0, 0, 210]), + ("4x4 generated seed 7", generatedIntegerEntries 4 4 7) + ] + where + assertFourByFourCertificate (label, entries) = + let result = do + matrixValue <- fromListMatrix @4 @4 @Integer entries + fullValue <- smithNormalForm matrixValue + firstDivisor <- minorDeterminantGcd (Proxy @1) 4 4 entries + secondDivisor <- minorDeterminantGcd (Proxy @2) 4 4 entries + thirdDivisor <- minorDeterminantGcd (Proxy @3) 4 4 entries + fourthDivisor <- minorDeterminantGcd (Proxy @4) 4 4 entries + pure + ( invariantFactorsFromDivisors [firstDivisor, secondDivisor, thirdDivisor, fourthDivisor], + diagonalEntriesOf 4 4 (toListMatrix (smithDiagonal fullValue)) + ) + in extractRight result $ + \(determinantalFactors, smithFactors) -> + assertEqual (label <> ": d_k = Delta_k / Delta_(k-1)") determinantalFactors smithFactors + +testSmithZeroDimensionalShapes :: Assertion +testSmithZeroDimensionalShapes = do + extractRight (fromListMatrix @0 @3 @Integer [] >>= smithNormalForm) $ + \value -> assertEqual "0x3 Smith diagonal" [] (toListMatrix (smithDiagonal value)) + extractRight (fromListMatrix @3 @0 @Integer [] >>= smithNormalForm) $ + \value -> assertEqual "3x0 Smith diagonal" [] (toListMatrix (smithDiagonal value)) + extractRight (fromListMatrix @0 @0 @Integer [] >>= smithNormalForm) $ + \value -> assertEqual "0x0 Smith diagonal" [] (toListMatrix (smithDiagonal value)) + +minorDeterminantGcd :: + forall k. + KnownNat k => + Proxy k -> + Int -> + Int -> + [Integer] -> + Either MoonlightError Integer +minorDeterminantGcd _ rowCount columnCount entries = + foldM accumulateDivisor 0 minorSelections + where + minorSize = matrixNat (Proxy @k) + rows = chunkRowsOf columnCount entries + minorSelections = + [ (rowSelection, columnSelection) + | rowSelection <- combinationsOf minorSize [0 .. rowCount - 1], + columnSelection <- combinationsOf minorSize [0 .. columnCount - 1] + ] + accumulateDivisor divisorSoFar (rowSelection, columnSelection) = do + minorMatrix <- + fromListMatrix @k @k @Integer + (concatMap (selectIndexed columnSelection) (selectIndexed rowSelection rows)) + minorDeterminant <- bareissDeterminant minorMatrix + pure (gcd divisorSoFar (abs minorDeterminant)) + +invariantFactorsFromDivisors :: [Integer] -> [Integer] +invariantFactorsFromDivisors = + snd . mapAccumL nextFactor 1 + where + nextFactor :: Integer -> Integer -> (Integer, Integer) + nextFactor previousDivisor divisor + | previousDivisor == 0 || divisor == 0 = (0, 0) + | otherwise = (divisor, divisor `div` previousDivisor) + +combinationsOf :: Int -> [a] -> [[a]] +combinationsOf size values + | size <= 0 = [[]] + | otherwise = + case values of + [] -> [] + value : rest -> + fmap (value :) (combinationsOf (size - 1) rest) <> combinationsOf size rest + +selectIndexed :: [Int] -> [a] -> [a] +selectIndexed selection = + fmap snd . filter (\(index, _) -> index `elem` selection) . zip [0 :: Int ..] + +diagonalEntriesOf :: Int -> Int -> [Integer] -> [Integer] +diagonalEntriesOf rowCount columnCount entries + | columnCount <= 0 = [] + | otherwise = + fmap snd (filter (onDiagonal . fst) (zip [0 :: Int ..] entries)) + where + onDiagonal index = + case index `divMod` columnCount of + (rowIndex, columnIndex) -> + rowIndex == columnIndex && rowIndex < min rowCount columnCount + +testSmithWitnessedRectangular :: Assertion +testSmithWitnessedRectangular = + assertSmithWitnessedFixture (Proxy @3) (Proxy @4) "rectangular witnessed Smith" [6, 10, 14, 22, 9, 15, 21, 33, 3, 5, 7, 11] + +testSmithWitnessedRankDeficient :: Assertion +testSmithWitnessedRankDeficient = + assertSmithWitnessedFixture (Proxy @4) (Proxy @4) "rank-deficient witnessed Smith" [4, 8, 12, 16, 6, 12, 18, 24, 10, 20, 30, 40, 0, 0, 0, 0] + +testSmithWitnessedTorsionRich :: Assertion +testSmithWitnessedTorsionRich = + assertSmithWitnessedFixture (Proxy @4) (Proxy @4) "torsion-rich witnessed Smith" [12, 18, 30, 42, 0, 36, 54, 78, 0, 0, 90, 126, 6, 0, 0, 210] + +testSmithWitnessedLargeEntry :: Assertion +testSmithWitnessedLargeEntry = + assertSmithWitnessedFixture (Proxy @3) (Proxy @3) "large-entry witnessed Smith" [4294967291, 4294967279, 4294967231, 4294967197, 4294967189, 4294967161, 4294967143, 4294967111, 4294967087] + +testSmithWitnessedFastPath :: Assertion +testSmithWitnessedFastPath = + assertSmithWitnessedFixture (Proxy @26) (Proxy @26) "nonsingular fast-path witnessed Smith" diagonallyDominantEntries + where + diagonallyDominantEntries :: [Integer] + diagonallyDominantEntries = + [ if rowIndex == columnIndex + then 26 + fromIntegral (rowIndex `mod` 9) + else fromIntegral ((rowIndex * 31 + columnIndex * 17) `mod` 3) - 1 + | rowIndex <- [0 .. 25 :: Int], + columnIndex <- [0 .. 25 :: Int] + ] + +assertSmithDiagonalAgreement :: + forall r c. + (KnownNat r, KnownNat c) => + Proxy r -> + Proxy c -> + String -> + [Integer] -> + Assertion +assertSmithDiagonalAgreement _ _ label entries = + let result = do + matrixValue <- fromListMatrix @r @c @Integer entries + fullValue <- smithNormalForm matrixValue + diagonalOnly <- smithDiagonalForm matrixValue + multimodular <- smithDiagonalFormMultimodular matrixValue + pure + ( toListMatrix (smithDiagonal fullValue), + toListMatrix (smithDiagonalMatrix diagonalOnly), + toListMatrix (smithDiagonalMatrix multimodular) + ) + in extractRight result $ + \(fullDiagonal, diagonalOnly, multimodular) -> do + assertEqual label fullDiagonal diagonalOnly + assertEqual (label <> " engine by name") fullDiagonal multimodular + +assertSmithWitnessedFixture :: + forall r c. + (KnownNat r, KnownNat c) => + Proxy r -> + Proxy c -> + String -> + [Integer] -> + Assertion +assertSmithWitnessedFixture _ _ label entries = + let rowCount = matrixNat (Proxy @r) + columnCount = matrixNat (Proxy @c) + result = do + matrixValue <- fromListMatrix @r @c @Integer entries + smithValue <- smithNormalFormWitnessed matrixValue + multimodular <- smithDiagonalFormMultimodular matrixValue + leftApplied <- mult (smithLeft smithValue) matrixValue + reconstructed <- mult leftApplied (smithRight smithValue) + leftInverseLeft <- mult (smithLeftInverse smithValue) (smithLeft smithValue) + leftLeftInverse <- mult (smithLeft smithValue) (smithLeftInverse smithValue) + rightInverseRight <- mult (smithRightInverse smithValue) (smithRight smithValue) + rightRightInverse <- mult (smithRight smithValue) (smithRightInverse smithValue) + pure + ( toListMatrix (smithDiagonal smithValue), + toListMatrix (smithDiagonalMatrix multimodular), + toListMatrix reconstructed, + toListMatrix leftInverseLeft, + toListMatrix leftLeftInverse, + toListMatrix rightInverseRight, + toListMatrix rightRightInverse + ) + in extractRight result $ + \(witnessDiagonal, multimodularDiagonal, reconstructed, leftInverseLeftEntries, leftLeftInverseEntries, rightInverseRightEntries, rightRightInverseEntries) -> do + assertEqual (label <> ": L * A * R") witnessDiagonal reconstructed + assertEqual (label <> ": multimodular diagonal") multimodularDiagonal witnessDiagonal + assertEqual (label <> ": L^-1 * L") (identityEntries rowCount) leftInverseLeftEntries + assertEqual (label <> ": L * L^-1") (identityEntries rowCount) leftLeftInverseEntries + assertEqual (label <> ": R^-1 * R") (identityEntries columnCount) rightInverseRightEntries + assertEqual (label <> ": R * R^-1") (identityEntries columnCount) rightRightInverseEntries + +matrixNat :: forall n. KnownNat n => Proxy n -> Int +matrixNat _ = + fromIntegral (natVal (Proxy @n)) + +testBareissRankAgreesWithRationalRank :: Assertion +testBareissRankAgreesWithRationalRank = + traverse_ + assertRankAgreement + [ generatedIntegerEntries 3 4 41, + generatedIntegerEntries 3 4 53, + [1, 2, 3, 4, 2, 4, 6, 8, 0, 0, 0, 0] + ] + where + assertRankAgreement entries = + let result = do + integerMatrix <- fromListMatrix @3 @4 @Integer entries + rationalMatrix <- fromListMatrix @3 @4 @Rational (fmap fromInteger entries) + integerRank <- bareissRank integerMatrix + rationalRank <- rank rationalMatrix + pure (integerRank, rationalRank) + in extractRight result $ + \(integerRank, rationalRank) -> + assertEqual "Bareiss rank must match Rational field rank" rationalRank integerRank + +testBareissDeterminantAgreesWithRationalDeterminant :: Assertion +testBareissDeterminantAgreesWithRationalDeterminant = + traverse_ + assertDeterminantAgreement + [ generatedIntegerEntries 4 4 67, + generatedIntegerEntries 4 4 79, + [1, 2, 3, 4, 2, 4, 6, 8, 3, 6, 9, 12, 0, 0, 0, 0] + ] + where + assertDeterminantAgreement entries = + let integerResult = do + integerMatrix <- fromListMatrix @4 @4 @Integer entries + bareissDeterminant integerMatrix + rationalRows = chunkRowsOf 4 (fmap fromInteger entries :: [Rational]) + rationalResult = exteriorPowerMatrix 4 rationalRows + in case (integerResult, rationalResult) of + (Right integerDeterminant, Right [[rationalDeterminant]]) -> + assertEqual "Bareiss determinant must match Rational determinant" rationalDeterminant (fromInteger integerDeterminant) + (Left failure, _) -> + assertFailure ("Bareiss determinant failed: " <> show failure) + (_, Left failure) -> + assertFailure ("Rational exterior determinant failed: " <> show failure) + (_, Right unexpected) -> + assertFailure ("Rational exterior determinant was not 1x1: " <> show unexpected) + +assertSmithWitnessInverses :: String -> [Integer] -> Assertion +assertSmithWitnessInverses label matrixEntries = + let result = do + matrixValue <- fromListMatrix @2 @2 matrixEntries + smithValue <- smithNormalForm matrixValue + let leftMatrix = smithLeft smithValue + diagonalMatrix = smithDiagonal smithValue + rightMatrix = smithRight smithValue + leftInverseMatrix = smithLeftInverse smithValue + rightInverseMatrix = smithRightInverse smithValue + la <- mult leftMatrix matrixValue + lar <- mult la rightMatrix + leftInverseLeft <- mult leftInverseMatrix leftMatrix + leftLeftInverse <- mult leftMatrix leftInverseMatrix + rightInverseRight <- mult rightInverseMatrix rightMatrix + rightRightInverse <- mult rightMatrix rightInverseMatrix + pure + ( toListMatrix diagonalMatrix, + toListMatrix lar, + toListMatrix leftInverseLeft, + toListMatrix leftLeftInverse, + toListMatrix rightInverseRight, + toListMatrix rightRightInverse + ) + in extractRight result $ + \(diagonalEntries, reconstructedEntries, leftInverseLeftEntries, leftLeftInverseEntries, rightInverseRightEntries, rightRightInverseEntries) -> do + assertEqual (label <> ": L * A * R") diagonalEntries reconstructedEntries + assertEqual (label <> ": L^-1 * L") identityEntries2 leftInverseLeftEntries + assertEqual (label <> ": L * L^-1") identityEntries2 leftLeftInverseEntries + assertEqual (label <> ": R^-1 * R") identityEntries2 rightInverseRightEntries + assertEqual (label <> ": R * R^-1") identityEntries2 rightRightInverseEntries + assertTwoByTwoDiagonalDivisibility diagonalEntries + +identityEntries2 :: [Integer] +identityEntries2 = [1, 0, 0, 1] + +identityEntries :: Int -> [Integer] +identityEntries sizeValue = + [ if rowIndex == columnIndex then 1 else 0 + | rowIndex <- [0 .. sizeValue - 1], + columnIndex <- [0 .. sizeValue - 1] + ] + +assertTwoByTwoOffDiagonalZero :: [Integer] -> Assertion +assertTwoByTwoOffDiagonalZero values = + case values of + [_, offDiagonal01, offDiagonal10, _] -> + assertBool "off-diagonal entries must be zero" (offDiagonal01 == 0 && offDiagonal10 == 0) + _ -> + assertFailure ("expected a 2x2 matrix payload, got " <> show values) + +assertTwoByTwoDiagonalDivisibility :: [Integer] -> Assertion +assertTwoByTwoDiagonalDivisibility values = + case values of + [d0, offDiagonal01, offDiagonal10, d1] -> do + assertBool "off-diagonal entries must be zero" (offDiagonal01 == 0 && offDiagonal10 == 0) + assertBool "d0 must divide d1" (d1 == 0 || d0 == 0 || d1 `mod` d0 == 0) + _ -> + assertFailure ("expected a 2x2 diagonal matrix payload, got " <> show values) + +testSmithIdentity :: Assertion +testSmithIdentity = + let result = do + matrixValue <- fromListMatrix @2 @2 [1 :: Integer, 0, 0, 1] + fmap (toListMatrix . smithDiagonal) (smithNormalForm matrixValue) + in extractRight result (\values -> assertEqual "smith identity" [1, 0, 0, 1] values) + +testSmithZero :: Assertion +testSmithZero = + let result = do + matrixValue <- fromListMatrix @2 @2 [0 :: Integer, 0, 0, 0] + fmap (toListMatrix . smithDiagonal) (smithNormalForm matrixValue) + in extractRight result (\values -> assertEqual "smith zero" [0, 0, 0, 0] values) + +generatedIntegerEntries :: Int -> Int -> Int -> [Integer] +generatedIntegerEntries rowCount columnCount seedValue = + [ generatedIntegerEntry seedValue rowIndex columnIndex + | rowIndex <- [0 .. rowCount - 1], + columnIndex <- [0 .. columnCount - 1] + ] + +generatedIntegerEntry :: Int -> Int -> Int -> Integer +generatedIntegerEntry seedValue rowIndex columnIndex = + fromIntegral ((((seedValue + 13 * rowIndex + 23 * columnIndex + 5 * rowIndex * columnIndex) `mod` 17) - 8) :: Int) + +chunkRowsOf :: Int -> [a] -> [[a]] +chunkRowsOf columnCount values = + case values of + [] -> [] + _ -> + let (rowValues, restValues) = splitAt columnCount values + in rowValues : chunkRowsOf columnCount restValues diff --git a/test/geometry/GeometryStorageSpec.hs b/test/geometry/GeometryStorageSpec.hs new file mode 100644 --- /dev/null +++ b/test/geometry/GeometryStorageSpec.hs @@ -0,0 +1,88 @@ +module GeometryStorageSpec + ( tests, + ) +where + +import Data.Vector.Unboxed qualified as U +import Foreign.Marshal.Alloc (alloca) +import Foreign.Storable (Storable (..)) +import Moonlight.LinAlg.Geometry + ( Vec2 (..), + Vec3 (..), + ) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertEqual, testCase) + +tests :: TestTree +tests = + testGroup + "Geometry storage" + [ testCase "Vec2 unboxed vectors round-trip through fromList/toList" testVec2UnboxRoundTrip, + testCase "Vec3 unboxed vectors round-trip through fromList/toList" testVec3UnboxRoundTrip, + testCase "Vec2 storable layout round-trips through peek/poke" testVec2StorableRoundTrip, + testCase "Vec3 storable layout round-trips through peek/poke" testVec3StorableRoundTrip, + testCase "Vec2 unboxed indexing agrees with list indexing on generated batches" testVec2UnboxIndexing, + testCase "Vec3 unboxed indexing agrees with list indexing on generated batches" testVec3UnboxIndexing + ] + +testVec2UnboxRoundTrip :: Assertion +testVec2UnboxRoundTrip = + assertEqual + "Vec2 U.fromList/U.toList identity" + vec2Batch + (U.toList (U.fromList vec2Batch :: U.Vector Vec2)) + +testVec3UnboxRoundTrip :: Assertion +testVec3UnboxRoundTrip = + assertEqual + "Vec3 U.fromList/U.toList identity" + vec3Batch + (U.toList (U.fromList vec3Batch :: U.Vector Vec3)) + +testVec2StorableRoundTrip :: Assertion +testVec2StorableRoundTrip = + assertStorableRoundTrip "Vec2 peek/poke identity" (Vec2 3.25 (-8.5)) + +testVec3StorableRoundTrip :: Assertion +testVec3StorableRoundTrip = + assertStorableRoundTrip "Vec3 peek/poke identity" (Vec3 3.25 (-8.5) 13.75) + +testVec2UnboxIndexing :: Assertion +testVec2UnboxIndexing = + assertEqual + "Vec2 indexed unboxed vector agrees with indexed source list" + (zip [0 :: Int ..] vec2Batch) + (U.toList (U.indexed (U.fromList vec2Batch :: U.Vector Vec2))) + +testVec3UnboxIndexing :: Assertion +testVec3UnboxIndexing = + assertEqual + "Vec3 indexed unboxed vector agrees with indexed source list" + (zip [0 :: Int ..] vec3Batch) + (U.toList (U.indexed (U.fromList vec3Batch :: U.Vector Vec3))) + +assertStorableRoundTrip :: (Eq value, Show value, Storable value) => String -> value -> Assertion +assertStorableRoundTrip label value = + alloca $ \pointerValue -> do + poke pointerValue value + actualValue <- peek pointerValue + assertEqual label value actualValue + +vec2Batch :: [Vec2] +vec2Batch = + (\indexValue -> Vec2 (coordinateValue 17 indexValue) (coordinateValue 29 indexValue)) + <$> [0 .. 127] + +vec3Batch :: [Vec3] +vec3Batch = + ( \indexValue -> + Vec3 + (coordinateValue 17 indexValue) + (coordinateValue 29 indexValue) + (coordinateValue 43 indexValue) + ) + <$> [0 .. 127] + +coordinateValue :: Int -> Int -> Double +coordinateValue saltValue indexValue = + (fromIntegral ((indexValue * 1103515245 + saltValue * 12345) `mod` 65521) / 257.0) - 127.0 diff --git a/test/sparse/SparsePackedSpec.hs b/test/sparse/SparsePackedSpec.hs new file mode 100644 --- /dev/null +++ b/test/sparse/SparsePackedSpec.hs @@ -0,0 +1,139 @@ +module SparsePackedSpec + ( tests, + ) +where + +import Data.Vector.Unboxed qualified as Unboxed +import Moonlight.LinAlg.Sparse + ( PackedSparseApplyError (..), + PackedSparseOperatorShapeError (..), + applyPackedSparseOperatorDense, + mkPackedSparseOperator, + packedSparseEntry, + packedSparseOperatorEntryCount, + ) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertEqual, assertFailure, testCase) + +tests :: TestTree +tests = + testGroup + "SparsePacked" + [ testCase "packed integral sparse operator applies to dense vectors" testPackedApply, + testCase "packed floating sparse operator applies to dense vectors" testPackedFloatingApply, + testCase "packed floating zero coefficients vanish from sealed storage" testFloatingZeroCoefficientsVanish, + testCase "packed floating operator rejects out-of-bounds entries" testFloatingRejectsOutOfBoundsEntry, + testCase "packed floating apply rejects wrong source vector length" testFloatingRejectsWrongSourceVectorLength, + testCase "packed operator rejects cardinalities beyond Int range before allocation" testRejectsCardinalityOutOfBounds, + testCase "zero coefficients vanish from sealed storage" testZeroCoefficientsVanish, + testCase "packed operator rejects out-of-bounds entries" testRejectsOutOfBoundsEntry, + testCase "packed apply rejects wrong source vector length" testRejectsWrongSourceVectorLength + ] + +testPackedApply :: Assertion +testPackedApply = + case mkPackedSparseOperator 3 2 entries of + Left failureValue -> + assertFailure ("packed operator construction failed: " <> show failureValue) + Right packedOperator -> + case applyPackedSparseOperatorDense packedOperator (Unboxed.fromList [2, 5, 7]) of + Left failureValue -> + assertFailure ("packed operator apply failed: " <> show failureValue) + Right resultValue -> + assertEqual "dense packed apply result" [23, -5] (Unboxed.toList resultValue) + where + entries = + [ packedSparseEntry 0 0 (1 :: Int), + packedSparseEntry 2 0 3, + packedSparseEntry 1 1 (-1) + ] + +testPackedFloatingApply :: Assertion +testPackedFloatingApply = + case mkPackedSparseOperator 3 2 entries of + Left failureValue -> + assertFailure ("packed floating operator construction failed: " <> show failureValue) + Right packedOperator -> + case applyPackedSparseOperatorDense packedOperator (Unboxed.fromList [2.0, 5.0, 7.0]) of + Left failureValue -> + assertFailure ("packed floating operator apply failed: " <> show failureValue) + Right resultValue -> + assertEqual "dense packed floating apply result" [2.75, -2.5] (Unboxed.toList resultValue) + where + entries = + [ packedSparseEntry 0 0 (0.5 :: Double), + packedSparseEntry 2 0 0.25, + packedSparseEntry 1 1 (-0.5) + ] + +testFloatingZeroCoefficientsVanish :: Assertion +testFloatingZeroCoefficientsVanish = + case mkPackedSparseOperator 2 2 [packedSparseEntry 0 0 (0.0 :: Double), packedSparseEntry 1 1 4.5] of + Left failureValue -> + assertFailure ("packed floating operator construction failed: " <> show failureValue) + Right packedOperator -> + assertEqual "nonzero packed floating entry count" 1 (packedSparseOperatorEntryCount packedOperator) + +testFloatingRejectsOutOfBoundsEntry :: Assertion +testFloatingRejectsOutOfBoundsEntry = + case mkPackedSparseOperator 1 1 [packedSparseEntry 1 0 (1.0 :: Double)] of + Left (PackedSparseEntryOutOfBounds sourceOffset targetOffset sourceDimension targetDimension) -> + assertEqual "out-of-bounds packed floating entry" (1, 0, 1, 1) (sourceOffset, targetOffset, sourceDimension, targetDimension) + Left failureValue -> + assertFailure ("expected out-of-bounds packed floating entry, received: " <> show failureValue) + Right _ -> + assertFailure "expected packed floating operator construction to reject out-of-bounds entry" + +testFloatingRejectsWrongSourceVectorLength :: Assertion +testFloatingRejectsWrongSourceVectorLength = + case mkPackedSparseOperator 2 1 [packedSparseEntry 1 0 (4.0 :: Double)] of + Left failureValue -> + assertFailure ("packed floating operator construction failed: " <> show failureValue) + Right packedOperator -> + case applyPackedSparseOperatorDense packedOperator (Unboxed.fromList [3.0]) of + Left (PackedSparseInputLengthMismatch expectedLength actualLength) -> + assertEqual "source vector length mismatch" (2, 1) (expectedLength, actualLength) + Right _ -> + assertFailure "expected packed floating operator apply to reject wrong source vector length" + +testRejectsCardinalityOutOfBounds :: Assertion +testRejectsCardinalityOutOfBounds = + case mkPackedSparseOperator oversizedCardinality 1 [packedSparseEntry 0 0 (1 :: Int)] of + Left (PackedSparseCardinalityOutOfBounds rejectedCardinality) -> + assertEqual "out-of-bounds cardinality" oversizedCardinality rejectedCardinality + Left failureValue -> + assertFailure ("expected cardinality failure, received: " <> show failureValue) + Right _ -> + assertFailure "expected packed operator construction to reject oversized cardinality" + where + oversizedCardinality = fromIntegral (maxBound :: Int) + 1 + +testZeroCoefficientsVanish :: Assertion +testZeroCoefficientsVanish = + case mkPackedSparseOperator 2 2 [packedSparseEntry 0 0 (0 :: Int), packedSparseEntry 1 1 4] of + Left failureValue -> + assertFailure ("packed operator construction failed: " <> show failureValue) + Right packedOperator -> + assertEqual "nonzero packed entry count" 1 (packedSparseOperatorEntryCount packedOperator) + +testRejectsOutOfBoundsEntry :: Assertion +testRejectsOutOfBoundsEntry = + case mkPackedSparseOperator 1 1 [packedSparseEntry 1 0 (1 :: Int)] of + Left (PackedSparseEntryOutOfBounds sourceOffset targetOffset sourceDimension targetDimension) -> + assertEqual "out-of-bounds packed entry" (1, 0, 1, 1) (sourceOffset, targetOffset, sourceDimension, targetDimension) + Left failureValue -> + assertFailure ("expected out-of-bounds packed entry, received: " <> show failureValue) + Right _ -> + assertFailure "expected packed operator construction to reject out-of-bounds entry" + +testRejectsWrongSourceVectorLength :: Assertion +testRejectsWrongSourceVectorLength = + case mkPackedSparseOperator 2 1 [packedSparseEntry 1 0 (4 :: Int)] of + Left failureValue -> + assertFailure ("packed operator construction failed: " <> show failureValue) + Right packedOperator -> + case applyPackedSparseOperatorDense packedOperator (Unboxed.fromList [3]) of + Left (PackedSparseInputLengthMismatch expectedLength actualLength) -> + assertEqual "source vector length mismatch" (2, 1) (expectedLength, actualLength) + Right _ -> + assertFailure "expected packed operator apply to reject wrong source vector length" diff --git a/test/sparse/SparseSolverSpec.hs b/test/sparse/SparseSolverSpec.hs new file mode 100644 --- /dev/null +++ b/test/sparse/SparseSolverSpec.hs @@ -0,0 +1,583 @@ +module SparseSolverSpec + ( tests, + ) +where + +import Moonlight.Core (MoonlightError) +import Moonlight.LinAlg.Sparse + ( IC0Config (..), + SparseConjugateGradientConfig (..), + SparseGMRESConfig (..), + SparsePreconditionerFamily (..), + SparseIterativeFailure (..), + SparseIterativeResult, + SparseStationaryIterationConfig (..), + sparseIterations, + sparseResidualNorm, + sparseSolution, + solveSparseCG, + solveSparseGMRES, + solveSparseJacobi, + solveSparseRichardson, + ) +import Moonlight.LinAlg.Sparse + ( SparseCSR, + cooToCSR, + mkSparseCOO, + ) +import Moonlight.LinAlg.Pure.Sparse.Solver.Preconditioner + ( compileSparsePreconditioner, + ) +import qualified Data.Vector.Unboxed as U +import Data.Foldable (traverse_) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase) + +tests :: TestTree +tests = + testGroup + "SparseSolver" + [ testCase "diagonal preconditioner rejects zero diagonal" testDiagonalPreconditionerRejectsZeroDiagonal, + testCase "identity preconditioner family compiles without matrix assumptions" testIdentityPreconditionerFamilyCompiles, + testCase "shifted diagonal preconditioner regularizes zero diagonals" testShiftedDiagonalPreconditionerRegularizesZeroDiagonal, + testCase "SSOR preconditioner rejects invalid relaxation" testSsorPreconditionerRejectsInvalidRelaxation, + testCase "IC(0) preconditioner rejects nonpositive pivot" testIC0PreconditionerRejectsNonpositivePivot, + testCase "sparse CG solves a small SPD system" testSparseCGSolvesSpd, + testCase "sparse CG with diagonal family solves a small SPD system" testSparseCGWithFamilySolvesSpd, + testCase "sparse CG with shifted diagonal family solves a small SPD system" testSparseCGWithShiftedFamilySolvesSpd, + testCase "sparse CG with SSOR family solves a small SPD system" testSparseCGWithSsorFamilySolvesSpd, + testCase "sparse CG with IC(0) family solves an anchored SPD Laplacian" testSparseCGWithIC0FamilySolvesAnchoredLaplacian, + testCase "restarted sparse GMRES accumulates Arnoldi work across cycles" testRestartedSparseGmresAccumulatesArnoldiWork, + testCase "sparse Jacobi solves a diagonal system" testSparseJacobiSolvesDiagonal, + testCase "sparse Richardson solves a diagonal system" testSparseRichardsonSolvesDiagonal, + testCase "Richardson converges on the 6x6 SPD equicorrelation counterexample" testRichardsonEquicorrelationRegression, + testCase "all iterative solvers reject non-finite matrix entries" testSolversRejectNonFiniteMatrix, + testCase "all iterative solvers reject non-finite right-hand sides" testSolversRejectNonFiniteRhs, + testCase "all iterative solvers reject non-finite initial guesses" testSolversRejectNonFiniteGuess, + testCase "all iterative solvers reject negative and NaN tolerances" testSolversRejectInvalidTolerance, + testCase "stationary solvers reject method-invalid damping" testStationarySolversRejectInvalidDamping, + testCase "GMRES reports zero-operator projected breakdown before correction" testGmresRejectsZeroProjectedDiagonal, + testCase "GMRES reports scale-negligible projected breakdown" testGmresRejectsNearProjectedBreakdown, + testCase "GMRES reports non-finite residual arithmetic" testGmresRejectsNonFiniteResidualArithmetic, + testCase "GMRES handles very large finite Hessenberg entries" testGmresHandlesLargeFiniteHessenberg, + testCase "GMRES accepts certified happy breakdown" testGmresHappyBreakdown, + testCase "GMRES rejects overflowing workspace cardinality before allocation" testGmresRejectsOverflowingWorkspace + ] + +testDiagonalPreconditionerRejectsZeroDiagonal :: Assertion +testDiagonalPreconditionerRejectsZeroDiagonal = + withSparseCSRFixture (csrFixture 2 2 [(0, 0, 1.0), (1, 0, 2.0)]) $ \matrixValue -> + case compileSparsePreconditioner DiagonalJacobiSparsePreconditionerFamily matrixValue of + Left (SparseInvalidInput _) -> + pure () + Left failureValue -> + assertFailure ("unexpected sparse preconditioner failure: " <> show failureValue) + Right _ -> + assertFailure "expected diagonal preconditioner to reject a zero diagonal" + +testIdentityPreconditionerFamilyCompiles :: Assertion +testIdentityPreconditionerFamilyCompiles = + withSparseCSRFixture (csrFixture 2 2 [(0, 1, 2.0)]) $ \matrixValue -> + case compileSparsePreconditioner IdentitySparsePreconditionerFamily matrixValue of + Left failureValue -> + assertFailure ("identity preconditioner family failed: " <> show failureValue) + Right _ -> + pure () + +testShiftedDiagonalPreconditionerRegularizesZeroDiagonal :: Assertion +testShiftedDiagonalPreconditionerRegularizesZeroDiagonal = + withSparseCSRFixture (csrFixture 2 2 [(0, 1, 1.0), (1, 1, 3.0)]) $ \matrixValue -> + case compileSparsePreconditioner (ShiftedDiagonalJacobiSparsePreconditionerFamily 0.5) matrixValue of + Left failureValue -> + assertFailure ("shifted diagonal preconditioner failed: " <> show failureValue) + Right _ -> + pure () + +testSsorPreconditionerRejectsInvalidRelaxation :: Assertion +testSsorPreconditionerRejectsInvalidRelaxation = + withSparseCSRFixture smallSpdMatrix $ \matrixValue -> + case compileSparsePreconditioner (SsorSparsePreconditionerFamily 2.0) matrixValue of + Left (SparseInvalidInput _) -> + pure () + Left failureValue -> + assertFailure ("unexpected SSOR preconditioner failure: " <> show failureValue) + Right _ -> + assertFailure "expected SSOR preconditioner to reject relaxation outside (0, 2)" + +testIC0PreconditionerRejectsNonpositivePivot :: Assertion +testIC0PreconditionerRejectsNonpositivePivot = + withSparseCSRFixture indefiniteSymmetricMatrix $ \matrixValue -> + case compileSparsePreconditioner (IncompleteCholesky0SparsePreconditionerFamily (IC0Config Nothing)) matrixValue of + Left (SparseNonpositivePivot 1 _) -> + pure () + Left failureValue -> + assertFailure ("unexpected IC(0) preconditioner failure: " <> show failureValue) + Right _ -> + assertFailure "expected IC(0) preconditioner to reject a nonpositive pivot" + +testSparseCGSolvesSpd :: Assertion +testSparseCGSolvesSpd = + withSparseCSRFixture smallSpdMatrix $ \matrixValue -> + case solveSparseCG (cgConfigWith IdentitySparsePreconditionerFamily) matrixValue (U.fromList [1.0, 2.0]) (U.fromList [0.0, 0.0]) of + Left failureValue -> + assertFailure ("sparse CG failed: " <> show failureValue) + Right resultValue -> + assertVectorApprox [1.0 / 11.0, 7.0 / 11.0] (sparseSolution resultValue) + +testSparseCGWithFamilySolvesSpd :: Assertion +testSparseCGWithFamilySolvesSpd = + withSparseCSRFixture smallSpdMatrix $ \matrixValue -> + case solveSparseCG (cgConfigWith DiagonalJacobiSparsePreconditionerFamily) matrixValue (U.fromList [1.0, 2.0]) (U.fromList [0.0, 0.0]) of + Left failureValue -> + assertFailure ("sparse CG with family failed: " <> show failureValue) + Right resultValue -> + assertVectorApprox [1.0 / 11.0, 7.0 / 11.0] (sparseSolution resultValue) + +testSparseCGWithShiftedFamilySolvesSpd :: Assertion +testSparseCGWithShiftedFamilySolvesSpd = + withSparseCSRFixture smallSpdMatrix $ \matrixValue -> + case solveSparseCG (cgConfigWith (ShiftedDiagonalJacobiSparsePreconditionerFamily 0.25)) matrixValue (U.fromList [1.0, 2.0]) (U.fromList [0.0, 0.0]) of + Left failureValue -> + assertFailure ("sparse CG with shifted family failed: " <> show failureValue) + Right resultValue -> + assertVectorApprox [1.0 / 11.0, 7.0 / 11.0] (sparseSolution resultValue) + +testSparseCGWithSsorFamilySolvesSpd :: Assertion +testSparseCGWithSsorFamilySolvesSpd = + withSparseCSRFixture smallSpdMatrix $ \matrixValue -> + case solveSparseCG (cgConfigWith (SsorSparsePreconditionerFamily 1.0)) matrixValue (U.fromList [1.0, 2.0]) (U.fromList [0.0, 0.0]) of + Left failureValue -> + assertFailure ("sparse CG with SSOR family failed: " <> show failureValue) + Right resultValue -> + assertVectorApprox [1.0 / 11.0, 7.0 / 11.0] (sparseSolution resultValue) + +testSparseCGWithIC0FamilySolvesAnchoredLaplacian :: Assertion +testSparseCGWithIC0FamilySolvesAnchoredLaplacian = + withSparseCSRFixture matrixValue $ \csrValue -> + case solveSparseCG ic0AnchoredLaplacianConfig csrValue rhsValues initialGuess of + Left failureValue -> + assertFailure ("sparse CG with IC(0) family failed: " <> show failureValue) + Right resultValue -> + assertBool + ("IC(0) PCG residual too large: " <> show (sparseResidualNorm resultValue)) + (sparseResidualNorm resultValue <= scgcTolerance ic0AnchoredLaplacianConfig) + where + dimension = 32 + matrixValue = anchoredPathLaplacian dimension + rhsValues = anchoredPathRightHandSide dimension + initialGuess = U.replicate dimension 0.0 + +testRestartedSparseGmresAccumulatesArnoldiWork :: Assertion +testRestartedSparseGmresAccumulatesArnoldiWork = + let dimension = 64 + restartDimension = 8 + matrixValue = anchoredPathLaplacian dimension + rhsValues = anchoredPathRightHandSide dimension + initialGuess = U.replicate dimension 0.0 + in withSparseCSRFixture matrixValue $ \csrValue -> + let cgResult = solveSparseCG cgRestartComparisonConfig csrValue rhsValues initialGuess + gmresResult = solveSparseGMRES (gmresConfigWith restartDimension) csrValue rhsValues initialGuess + in case (cgResult, gmresResult) of + (Left failureValue, _) -> + assertFailure ("sparse CG comparison failed: " <> show failureValue) + (_, Left failureValue) -> + assertFailure ("restarted sparse GMRES failed: " <> show failureValue) + (Right cgValue, Right gmresValue) -> do + assertBool + ("GMRES iterations did not cross a restart boundary: " <> show (sparseIterations gmresValue)) + (sparseIterations gmresValue > restartDimension) + assertBool + ("GMRES true residual too large: " <> show (sparseResidualNorm gmresValue)) + (sparseResidualNorm gmresValue <= 1.0e-8) + assertVectorApproxWith 1.0e-4 (sparseSolution cgValue) (sparseSolution gmresValue) + +testSparseJacobiSolvesDiagonal :: Assertion +testSparseJacobiSolvesDiagonal = + withSparseCSRFixture diagonalMatrix $ \matrixValue -> + case solveSparseJacobi stationaryConfig matrixValue (U.fromList [4.0, 9.0]) (U.fromList [0.0, 0.0]) of + Left failureValue -> + assertFailure ("sparse Jacobi failed: " <> show failureValue) + Right resultValue -> + assertVectorApprox [2.0, 3.0] (sparseSolution resultValue) + +testSparseRichardsonSolvesDiagonal :: Assertion +testSparseRichardsonSolvesDiagonal = + withSparseCSRFixture diagonalMatrix $ \matrixValue -> + case solveSparseRichardson stationaryConfig matrixValue (U.fromList [4.0, 9.0]) (U.fromList [0.0, 0.0]) of + Left failureValue -> + assertFailure ("sparse Richardson failed: " <> show failureValue) + Right resultValue -> + assertVectorApprox [2.0, 3.0] (sparseSolution resultValue) + +testRichardsonEquicorrelationRegression :: Assertion +testRichardsonEquicorrelationRegression = + withSparseCSRFixture equicorrelationMatrix $ \matrixValue -> + case solveSparseRichardson stationaryConfig matrixValue (U.replicate 6 5.5) (U.replicate 6 0.0) of + Left failureValue -> + assertFailure ("Richardson failed on SPD equicorrelation matrix: " <> show failureValue) + Right resultValue -> + assertVectorApprox (replicate 6 1.0) (sparseSolution resultValue) + +testSolversRejectNonFiniteMatrix :: Assertion +testSolversRejectNonFiniteMatrix = + traverse_ + (\invalidValue -> + withSparseCSRFixture (csrFixture 1 1 [(0, 0, invalidValue)]) $ \matrixValue -> + assertSparseInvalidInvocations + (solverInvocations matrixValue (U.singleton 1.0) (U.singleton 0.0)) + ) + nonFiniteValues + +testSolversRejectNonFiniteRhs :: Assertion +testSolversRejectNonFiniteRhs = + withSparseCSRFixture identityMatrix $ \matrixValue -> + traverse_ + (\invalidValue -> + assertSparseInvalidInvocations + (solverInvocations matrixValue (U.singleton invalidValue) (U.singleton 0.0)) + ) + nonFiniteValues + +testSolversRejectNonFiniteGuess :: Assertion +testSolversRejectNonFiniteGuess = + withSparseCSRFixture identityMatrix $ \matrixValue -> + traverse_ + (\invalidValue -> + assertSparseInvalidInvocations + (solverInvocations matrixValue (U.singleton 1.0) (U.singleton invalidValue)) + ) + nonFiniteValues + +testSolversRejectInvalidTolerance :: Assertion +testSolversRejectInvalidTolerance = + withSparseCSRFixture identityMatrix $ \matrixValue -> + traverse_ + (\invalidTolerance -> + assertSparseInvalidInvocations + (invalidToleranceInvocations invalidTolerance matrixValue) + ) + [-1.0, nanValue] + +testStationarySolversRejectInvalidDamping :: Assertion +testStationarySolversRejectInvalidDamping = + withSparseCSRFixture identityMatrix $ \matrixValue -> + traverse_ + assertSparseInvalidInvocation + [ ( "Jacobi damping above one", + solveSparseJacobi (stationaryConfig {ssicDamping = 1.01}) matrixValue (U.singleton 1.0) (U.singleton 0.0) + ), + ( "Jacobi NaN damping", + solveSparseJacobi (stationaryConfig {ssicDamping = nanValue}) matrixValue (U.singleton 1.0) (U.singleton 0.0) + ), + ( "Richardson damping at two", + solveSparseRichardson (stationaryConfig {ssicDamping = 2.0}) matrixValue (U.singleton 1.0) (U.singleton 0.0) + ), + ( "Richardson infinite damping", + solveSparseRichardson (stationaryConfig {ssicDamping = infinityValue}) matrixValue (U.singleton 1.0) (U.singleton 0.0) + ) + ] + +testGmresRejectsZeroProjectedDiagonal :: Assertion +testGmresRejectsZeroProjectedDiagonal = + withSparseCSRFixture (csrFixture 1 1 []) $ \matrixValue -> + assertSparseInvalidInvocation + ( "zero operator", + solveSparseGMRES (identityGmresConfig 1) matrixValue (U.singleton 1.0) (U.singleton 0.0) + ) + +testGmresRejectsNearProjectedBreakdown :: Assertion +testGmresRejectsNearProjectedBreakdown = + withSparseCSRFixture nearlySingularMatrix $ \matrixValue -> + assertSparseInvalidInvocation + ( "near projected breakdown", + solveSparseGMRES (identityGmresConfig 2) matrixValue (U.fromList [1.0, 0.0]) (U.fromList [0.0, 0.0]) + ) + +testGmresRejectsNonFiniteResidualArithmetic :: Assertion +testGmresRejectsNonFiniteResidualArithmetic = + withSparseCSRFixture overflowingResidualMatrix $ \matrixValue -> + assertSparseInvalidInvocation + ( "non-finite residual arithmetic", + solveSparseGMRES + (identityGmresConfig 2) + matrixValue + (U.fromList [0.0, 0.0]) + (U.fromList [1.0e308, 1.0e308]) + ) + +testGmresHandlesLargeFiniteHessenberg :: Assertion +testGmresHandlesLargeFiniteHessenberg = + withSparseCSRFixture (csrFixture 1 1 [(0, 0, 1.0e300)]) $ \matrixValue -> + case solveSparseGMRES (identityGmresConfig 1) matrixValue (U.singleton 1.0e300) (U.singleton 0.0) of + Left failureValue -> + assertFailure ("GMRES rejected representable scaled Givens arithmetic: " <> show failureValue) + Right resultValue -> + assertVectorApprox [1.0] (sparseSolution resultValue) + +testGmresHappyBreakdown :: Assertion +testGmresHappyBreakdown = + withSparseCSRFixture identityMatrix $ \matrixValue -> + case solveSparseGMRES (identityGmresConfig 1) matrixValue (U.singleton 3.0) (U.singleton 0.0) of + Left failureValue -> + assertFailure ("GMRES happy breakdown failed certification: " <> show failureValue) + Right resultValue -> do + assertBool "happy breakdown should complete one Arnoldi step" (sparseIterations resultValue == 1) + assertVectorApprox [3.0] (sparseSolution resultValue) + +testGmresRejectsOverflowingWorkspace :: Assertion +testGmresRejectsOverflowingWorkspace = + withSparseCSRFixture (csrFixture 0 0 []) $ \matrixValue -> + assertSparseInvalidInvocation + ( "overflowing workspace", + solveSparseGMRES + ( (identityGmresConfig maxBound) + { sgcIterationLimit = 0 + } + ) + matrixValue + U.empty + U.empty + ) + +smallSpdMatrix :: Either MoonlightError (SparseCSR Double) +smallSpdMatrix = + csrFixture + 2 + 2 + [ (0, 0, 4.0), + (0, 1, 1.0), + (1, 0, 1.0), + (1, 1, 3.0) + ] + +indefiniteSymmetricMatrix :: Either MoonlightError (SparseCSR Double) +indefiniteSymmetricMatrix = + csrFixture + 2 + 2 + [ (0, 0, 1.0), + (0, 1, 2.0), + (1, 0, 2.0), + (1, 1, 1.0) + ] + +diagonalMatrix :: Either MoonlightError (SparseCSR Double) +diagonalMatrix = + csrFixture + 2 + 2 + [ (0, 0, 2.0), + (1, 1, 3.0) + ] + +identityMatrix :: Either MoonlightError (SparseCSR Double) +identityMatrix = + csrFixture 1 1 [(0, 0, 1.0)] + +nearlySingularMatrix :: Either MoonlightError (SparseCSR Double) +nearlySingularMatrix = + csrFixture + 2 + 2 + [ (0, 0, 1.0), + (0, 1, 1.0), + (1, 0, 1.0), + (1, 1, 1.0 + 1.0e-14) + ] + +overflowingResidualMatrix :: Either MoonlightError (SparseCSR Double) +overflowingResidualMatrix = + csrFixture + 2 + 2 + [ (0, 0, 1.0e308), + (0, 1, -1.0e308) + ] + +equicorrelationMatrix :: Either MoonlightError (SparseCSR Double) +equicorrelationMatrix = + csrFixture + 6 + 6 + ( concatMap + (\rowIndex -> + (\columnIndex -> (rowIndex, columnIndex, if rowIndex == columnIndex then 1.0 else 0.9)) + <$> [0 .. 5] + ) + [0 .. 5] + ) + +anchoredPathLaplacian :: Int -> Either MoonlightError (SparseCSR Double) +anchoredPathLaplacian dimension = + csrFixture + dimension + dimension + ((0, 0, 1.0) : concatMap anchoredPathEdgeEntries [0 .. dimension - 2]) + +anchoredPathEdgeEntries :: Int -> [(Int, Int, Double)] +anchoredPathEdgeEntries leftIndex = + let rightIndex = leftIndex + 1 + in [ (leftIndex, leftIndex, 1.0), + (leftIndex, rightIndex, -1.0), + (rightIndex, leftIndex, -1.0), + (rightIndex, rightIndex, 1.0) + ] + +anchoredPathRightHandSide :: Int -> U.Vector Double +anchoredPathRightHandSide dimension = + U.generate + dimension + ( \indexValue -> + let entryPhase = fromIntegral (indexValue + 1) + entrySkew = fromIntegral ((indexValue * 7) `mod` 11) + in 1.0 + sin entryPhase + 0.125 * entrySkew + ) + +csrFixture :: Int -> Int -> [(Int, Int, Double)] -> Either MoonlightError (SparseCSR Double) +csrFixture rowCount columnCount entries = + mkSparseCOO rowCount columnCount entries >>= cooToCSR + +withSparseCSRFixture :: Either MoonlightError (SparseCSR Double) -> (SparseCSR Double -> Assertion) -> Assertion +withSparseCSRFixture fixtureValue onFixture = + case fixtureValue of + Left err -> assertFailure ("invalid sparse solver fixture: " <> show err) + Right csrValue -> onFixture csrValue + +cgConfigWith :: SparsePreconditionerFamily -> SparseConjugateGradientConfig +cgConfigWith preconditionerFamily = + SparseConjugateGradientConfig + { scgcTolerance = 1.0e-10, + scgcIterationLimit = 32, + scgcPreconditionerFamily = preconditionerFamily + } + +gmresConfigWith :: Int -> SparseGMRESConfig +gmresConfigWith restartDimension = + SparseGMRESConfig + { sgcTolerance = 1.0e-8, + sgcIterationLimit = 4096, + sgcRestartDimension = restartDimension, + sgcPreconditionerFamily = DiagonalJacobiSparsePreconditionerFamily + } + +identityGmresConfig :: Int -> SparseGMRESConfig +identityGmresConfig restartDimension = + (gmresConfigWith restartDimension) + { sgcPreconditionerFamily = IdentitySparsePreconditionerFamily + } + +cgRestartComparisonConfig :: SparseConjugateGradientConfig +cgRestartComparisonConfig = + SparseConjugateGradientConfig + { scgcTolerance = 1.0e-10, + scgcIterationLimit = 256, + scgcPreconditionerFamily = IdentitySparsePreconditionerFamily + } + +ic0AnchoredLaplacianConfig :: SparseConjugateGradientConfig +ic0AnchoredLaplacianConfig = + SparseConjugateGradientConfig + { scgcTolerance = 1.0e-8, + scgcIterationLimit = 128, + scgcPreconditionerFamily = IncompleteCholesky0SparsePreconditionerFamily (IC0Config Nothing) + } + +stationaryConfig :: SparseStationaryIterationConfig +stationaryConfig = + SparseStationaryIterationConfig + { ssicTolerance = 1.0e-8, + ssicIterationLimit = 128, + ssicDamping = 1.0 + } + +type SolverInvocation = (String, Either SparseIterativeFailure SparseIterativeResult) + +solverInvocations :: SparseCSR Double -> U.Vector Double -> U.Vector Double -> [SolverInvocation] +solverInvocations matrixValue rhsValues initialGuess = + [ ( "CG", + solveSparseCG (cgConfigWith IdentitySparsePreconditionerFamily) matrixValue rhsValues initialGuess + ), + ( "GMRES", + solveSparseGMRES (identityGmresConfig 1) matrixValue rhsValues initialGuess + ), + ( "Jacobi", + solveSparseJacobi stationaryConfig matrixValue rhsValues initialGuess + ), + ( "Richardson", + solveSparseRichardson stationaryConfig matrixValue rhsValues initialGuess + ) + ] + +invalidToleranceInvocations :: Double -> SparseCSR Double -> [SolverInvocation] +invalidToleranceInvocations invalidTolerance matrixValue = + [ ( "CG invalid tolerance", + solveSparseCG + ((cgConfigWith IdentitySparsePreconditionerFamily) {scgcTolerance = invalidTolerance}) + matrixValue + (U.singleton 1.0) + (U.singleton 0.0) + ), + ( "GMRES invalid tolerance", + solveSparseGMRES + ((identityGmresConfig 1) {sgcTolerance = invalidTolerance}) + matrixValue + (U.singleton 1.0) + (U.singleton 0.0) + ), + ( "Jacobi invalid tolerance", + solveSparseJacobi + (stationaryConfig {ssicTolerance = invalidTolerance}) + matrixValue + (U.singleton 1.0) + (U.singleton 0.0) + ), + ( "Richardson invalid tolerance", + solveSparseRichardson + (stationaryConfig {ssicTolerance = invalidTolerance}) + matrixValue + (U.singleton 1.0) + (U.singleton 0.0) + ) + ] + +assertSparseInvalidInvocations :: [SolverInvocation] -> Assertion +assertSparseInvalidInvocations = + traverse_ assertSparseInvalidInvocation + +assertSparseInvalidInvocation :: SolverInvocation -> Assertion +assertSparseInvalidInvocation (solverName, solverResult) = + case solverResult of + Left (SparseInvalidInput _) -> pure () + Left failureValue -> + assertFailure (solverName <> " returned the wrong typed obstruction: " <> show failureValue) + Right resultValue -> + assertFailure (solverName <> " fabricated success: " <> show resultValue) + +nonFiniteValues :: [Double] +nonFiniteValues = [nanValue, infinityValue, negate infinityValue] + +nanValue :: Double +nanValue = 0.0 / 0.0 + +infinityValue :: Double +infinityValue = 1.0 / 0.0 + +assertVectorApprox :: [Double] -> U.Vector Double -> Assertion +assertVectorApprox expectedValues actualValues = + assertBool + ("expected " <> show expectedValues <> " but received " <> show actualValues) + ( and + ( zipWith + (\expectedValue actualValue -> abs (expectedValue - actualValue) <= 1.0e-5) + expectedValues + (U.toList actualValues) + ) + ) + +assertVectorApproxWith :: Double -> U.Vector Double -> U.Vector Double -> Assertion +assertVectorApproxWith tolerance expectedValues actualValues = + assertBool + ("expected " <> show expectedValues <> " but received " <> show actualValues) + ( U.length expectedValues == U.length actualValues + && U.and (U.zipWith (\expectedValue actualValue -> abs (expectedValue - actualValue) <= tolerance) expectedValues actualValues) + ) diff --git a/test/spectral/KrylovSpec.hs b/test/spectral/KrylovSpec.hs new file mode 100644 --- /dev/null +++ b/test/spectral/KrylovSpec.hs @@ -0,0 +1,1116 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE RecordWildCards #-} + +module KrylovSpec + ( tests, + ) +where + +import qualified Data.Vector as Box +import qualified Data.Vector.Unboxed as U +import Data.List (sort) +import Moonlight.Core (MoonlightError (..)) +import Moonlight.LinAlg.Dense (mkDynMatrix) +import Moonlight.LinAlg.Internal.Eigen.Kernels (epsDouble) +import Moonlight.LinAlg.Krylov + ( SpectrumEnd (..), + PositiveCount, + NonNegativeConfigTolerance, + blockLanczosBasisCount, + blockLanczosProjectedBlockTridiagonal, + blockLanczosSymmetric, + defaultBlockLanczosConfig, + defaultLanczosConfig, + LanczosConfig, + lanczosBasisColumns, + lanczosProjectedTridiagonal, + lanczosRestartProjectionBasisColumns, + lanczosRestartProjectionProjectedPairs, + lanczosRestartedProjection, + lanczosStepsCompleted, + lanczosSymmetric, + mkNonNegativeConfigTolerance, + mkPositiveCount, + withBlockLanczosBlockSize, + withBlockLanczosIterations, + withLanczosIterations, + withLanczosTolerance, + ) +import Moonlight.LinAlg.Operator + ( addScaledIdentity, + csrLinearOperator, + declaredSelfAdjointVectorLinearOperator, + diagonalLinearOperator, + graphLaplacianLinearOperator, + LinearOperator, + OperatorSymmetry (SelfAdjointOperator), + operatorShape, + pathLaplacianLinearOperator, + runOperatorU, + scaleLinearOperator, + selfAdjointCSRLinearOperator, + sigmaIdentityMinus, + ) +import Moonlight.LinAlg.Pure.Krylov.Projected + ( SymmetricProjectedOperator (..), + projectedEigenpairs, + projectedEigenpairsFromRestartedLanczos, + projectedEigenvalues, + projectedSubspaceDimension, + projectedSubspaceFromBlockLanczos, + projectedSubspaceFromLanczos, + projectedSubspaceOperator, + symmetricProjectedOperatorDimension, + ) +import Moonlight.LinAlg.Pure.Krylov.SelectedTridiagonal + ( TridiagonalRejection (..), + selectedSymmetricTridiagonalEigenpairsDirect, + selectedSymmetricTridiagonalEigenvaluesDirect, + symmetricTridiagonalFromCSR, + ) +import Moonlight.LinAlg.Pure.Structured.BlockTridiagonal + ( applySymmetricBlockTridiagonalU, + mkRowMajorBlock, + mkSymmetricBlockTridiagonal, + SymmetricBlockTridiagonal, + symmetrizeRowMajorBlockLower, + symmetricBlockTridiagonalBandwidth, + symmetricBlockTridiagonalDimension, + symmetricBlockTridiagonalEntry, + symmetricBlockTridiagonalFrobeniusNorm, + ) +import Moonlight.LinAlg.Pure.Structured.Tridiagonal + ( SymmetricTridiagonal, + mkSymmetricTridiagonal, + symmetricTridiagonalDiagonalEntries, + symmetricTridiagonalDimension, + symmetricTridiagonalOffDiagonalEntries, + ) +import Moonlight.LinAlg.Sparse + ( GraphEdge (..), + SparseCSR, + cooToCSR, + mkSparseCOO, + pathLaplacianCSR, + tridiagonalCSR, + ) +import Moonlight.LinAlg.Spectral + ( CertifiedSelectedEigenpairResult (..), + Eigenpairs, + EigenRequest (..), + SelectedEigenpairCertificationFailure (..), + SelectedEigenpairOrthonormalityEvidence (..), + SelectedEigenpairRequestOrderingEvidence (..), + SelectedEigenpairResidualEvidence (..), + certifySelectedEigenpairResult, + defaultEigenSolveConfig, + eigenpairCount, + eigenpairDimension, + eigenpairResidualNorms, + eigenpairValues, + eigenpairVectorAt, + solveEigenRequest, + withEigenFallbackInitialVector, + withEigenFallbackLanczosConfig, + ) +import Moonlight.LinAlg.Pure.Spectral.Solve (denseSpectralFallbackDimensionThreshold) +import Moonlight.LinAlg.Native + ( selectedSymmetricBlockTridiagonalEigenRequestLapack, + selectedSymmetricTridiagonalEigenRequestLapack, + symmetricEigenRequestLapack, + ) +import Helpers (extractRight) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, assertFailure, testCase) +import Test.Tasty.QuickCheck qualified as QC +import Prelude + +tests :: TestTree +tests = + testGroup + "Krylov" + [ testGroup + "operator" + [ testCase "operator shape" testOperatorShape, + testCase "CSR operator matvec" testCSROperatorMatvec, + testCase "affine normalization applies sigma I minus A" testSigmaIdentityMinusApplies, + testCase "zero affine scale has deterministic coordinate eigenpairs" testZeroScaleEigenpairs, + testCase "negative affine scale reverses the requested spectrum end" testNegativeScaleEigenvalues + ], + testGroup + "spectral dispatch" + [ testCase "diagonal values and pairs are selected directly" testDiagonalSpectralDispatch, + testCase "path Laplacian values and pairs are closed-form" testPathSpectralDispatch, + testCase "checked graph Laplacian construction preserves spectral semantics" testGraphLaplacianSpectralDispatch, + testCase "large cycle graph descent selects the analytic smallest spectrum" testLargeCycleGraphCascadicSpectrum, + testCase "large graph Laplacian smallest modes descend with certified residuals" testLargeGraphLaplacianCascadicDispatch, + testCase "generic low-demand fallback densifies through threshold and restarts above it" testGenericFallbackThresholdDispatch, + testCase "generic high-demand fallback stays on the bounded dense route" testGenericHighDemandFallbackDispatch, + testCase "requested count above dimension is rejected" testEigenCountRejectsOversubscription + ], + testGroup + "exact tridiagonal structure" + [ testCase "zero-dimensional tridiagonal requests are rejected" testZeroDimensionTridiagonalRejected, + testCase "one-dimensional tridiagonal values and pairs are exact" testOneDimensionalTridiagonal, + testCase "repeated diagonal eigenvalues are checked through the projector" testRepeatedDiagonalProjector, + testCase "CSR tridiagonal classification preserves tiny non-zero couplings" testTinyCouplingIsStructural, + testCase "CSR tridiagonal classification rejects out-of-band entries exactly" testOutOfBandCSRRejected, + testCase "CSR tridiagonal classification rejects asymmetric off-diagonals exactly" testAsymmetricCSRRejected, + testCase "reducible tridiagonal values split and interleave blocks" testReducibleTridiagonalValues, + testCase "reducible tridiagonal pairs split and interleave blocks" testReducibleTridiagonalPairs, + testCase "generic tridiagonal pairs are selected inverse-iteration residual-checked columns" testGenericTridiagonalPairs, + testCase "selected tridiagonal pairs certify residuals orthonormality count and order" testSelectedTridiagonalPairCertification, + testCase "generic selected tridiagonal pairs agree with all-pairs on small n" testGenericSelectedTridiagonalPairsAgreeWithAllPairs, + testCase "clustered irreducible tridiagonal values agree with LAPACK" testClusteredIrreducibleTridiagonalValues, + testCase "generic irreducible largest values agree with LAPACK order" testGenericIrreducibleLargestValues, + testCase "extreme scaled irreducible tridiagonal values agree with LAPACK" testExtremeScaledIrreducibleTridiagonalValues, + testCase "perturbed path is not silently solved as the exact path" testPerturbedPathNotExactPath, + testCase "extreme affine scale transports eigenvalues without losing order" testExtremeAffineScale + ], + testGroup + "projected block" + [ testCase "packed symmetric block tridiagonal applies without row materialization" testPackedBlockTridiagonalApply, + testCase "row-major block rejects wrapped cardinality" testRowMajorBlockCardinalityOverflow, + testCase "packed block Frobenius norm preserves overflow evidence" testBlockFrobeniusNormPreservesOverflow, + QC.testProperty "packed block Frobenius norm equals dense reconstruction on mixed block sizes" propBlockFrobeniusNormMatchesDense, + testCase "block Lanczos builds a structured projected operator" testBlockLanczosStructuredProjection, + testCase "projected tridiagonal pairs use selected tridiagonal columns" testProjectedTridiagonalSelectedPairs, + testCase "projected eigensolve rejects oversubscribed requests" testProjectedCountRejectsOversubscription + ], + testGroup + "native selected eigensolve" + [ testCase "dense selected values agree with selected pairs" testNativeDenseSelectedValuesAgreeWithPairs, + testCase "generic tridiagonal selected values and pairs use native selected path" testNativeGenericTridiagonalSelected, + testCase "block-band selected values and pairs agree with dense reference" testNativeBlockBandSelectedAgreesWithDense + ], + testGroup + "Lanczos" + [ testCase "Lanczos decomposition owns a SymmetricTridiagonal projection" testLanczosProjection, + testCase "thick restart fallback converges across multiple Krylov cycles" testLanczosThickRestartMultiCycle, + testCase "thick restart locked Ritz pairs satisfy residual bounds" testLanczosThickRestartLockedResiduals, + testCase "thick restart projected pairs carry residual evidence" testLanczosThickRestartProjectedResidualEvidence, + testCase "thick restart basis stays orthogonal across locked and active vectors" testLanczosThickRestartBasisOrthogonality + ] + ] + +tolerance :: Double +tolerance = 1.0e-7 + +approxEqual :: Double -> Double -> Bool +approxEqual expected actual = + abs (expected - actual) <= tolerance * max 1.0 (max (abs expected) (abs actual)) + +assertApproxList :: [Double] -> [Double] -> Assertion +assertApproxList expected actual = + assertBool + ("expected " <> show expected <> " but received " <> show actual) + (length expected == length actual && and (zipWith approxEqual expected actual)) + +assertApproxVector :: [Double] -> U.Vector Double -> Assertion +assertApproxVector expected actual = + assertApproxList expected (U.toList actual) + +withPositiveCount :: Int -> (PositiveCount -> Assertion) -> Assertion +withPositiveCount value onCount = + extractRight (mkPositiveCount value) onCount + +assertScaledEigenpairResiduals :: String -> Double -> Eigenpairs -> Assertion +assertScaledEigenpairResiduals context matrixNorm pairs = do + _ <- traverse (assertScaledEigenpairResidualAt context matrixNorm pairs) [0 .. eigenpairCount pairs - 1] + pure () + +assertScaledEigenpairResidualAt :: String -> Double -> Eigenpairs -> Int -> Assertion +assertScaledEigenpairResidualAt context matrixNorm pairs columnIndex = do + eigenvector <- extractEither (eigenpairVectorAt columnIndex pairs) + let eigenvalue = eigenpairValues pairs `U.unsafeIndex` columnIndex + residualNorm = eigenpairResidualNorms pairs `U.unsafeIndex` columnIndex + residualRatio = scaledEigenpairResidualRatio matrixNorm eigenvalue eigenvector residualNorm + residualLimit = scaledEigenpairResidualLimit pairs + assertBool + ( context + <> " residual ratio exceeded scaled machine bound at column " + <> show columnIndex + <> ": ratio=" + <> show residualRatio + <> ", limit=" + <> show residualLimit + <> ", residual=" + <> show residualNorm + ) + (isFiniteDouble residualRatio && residualRatio <= residualLimit) + +scaledEigenpairResidualRatio :: Double -> Double -> U.Vector Double -> Double -> Double +scaledEigenpairResidualRatio matrixNorm eigenvalue eigenvector residualNorm = + residualNorm / max 1.0 ((matrixNorm + abs eigenvalue) * vectorNormU eigenvector) + +scaledEigenpairResidualLimit :: Eigenpairs -> Double +scaledEigenpairResidualLimit pairs = + 1.0e7 * max 1.0 (fromIntegral (eigenpairDimension pairs)) * epsDouble + +vectorNormU :: U.Vector Double -> Double +vectorNormU values = + sqrt (U.sum (U.map (\entryValue -> entryValue * entryValue) values)) + +isFiniteDouble :: Double -> Bool +isFiniteDouble value = + not (isNaN value || isInfinite value) + +denseFrobeniusNorm :: [Double] -> Double +denseFrobeniusNorm values = + sqrt (sum ((\entryValue -> entryValue * entryValue) <$> values)) + +tridiagonalFrobeniusNorm :: [Double] -> [Double] -> Double +tridiagonalFrobeniusNorm diagonalValues offDiagonalValues = + sqrt + ( sum ((\entryValue -> entryValue * entryValue) <$> diagonalValues) + + 2.0 * sum ((\entryValue -> entryValue * entryValue) <$> offDiagonalValues) + ) + +tridiagonalMatrixNorm :: SymmetricTridiagonal -> Double +tridiagonalMatrixNorm tridiagonalValue = + tridiagonalFrobeniusNorm + (symmetricTridiagonalDiagonalEntries tridiagonalValue) + (symmetricTridiagonalOffDiagonalEntries tridiagonalValue) + +pathLaplacianFrobeniusNorm :: Int -> Double +pathLaplacianFrobeniusNorm dimension + | dimension <= 0 = 0.0 + | dimension == 1 = 0.0 + | otherwise = + tridiagonalFrobeniusNorm + (1.0 : (replicate (dimension - 2) 2.0 <> [1.0])) + (replicate (dimension - 1) (-1.0)) + +eigenpairProjectorDiagonal :: Eigenpairs -> Either MoonlightError (U.Vector Double) +eigenpairProjectorDiagonal pairs = do + columns <- traverse (`eigenpairVectorAt` pairs) [0 .. eigenpairCount pairs - 1] + pure + ( U.generate + (eigenpairDimension pairs) + ( \rowIndex -> + sum + ( (\columnVector -> + let !entryValue = columnVector `U.unsafeIndex` rowIndex + in entryValue * entryValue + ) + <$> columns + ) + ) + ) + +assertEigenpairColumnsOrthonormal :: String -> Eigenpairs -> Assertion +assertEigenpairColumnsOrthonormal context pairs = + case traverse (`eigenpairVectorAt` pairs) [0 .. eigenpairCount pairs - 1] of + Left err -> assertFailure (context <> ": eigenpair column extraction failed: " <> show err) + Right columns -> do + _ <- + traverse + (assertEigenpairColumnInnerProduct context) + [ (leftIndex, rightIndex, leftColumn, rightColumn) + | (leftIndex, leftColumn) <- zip [0 ..] columns, + (rightIndex, rightColumn) <- zip [0 ..] columns, + leftIndex <= rightIndex + ] + pure () + +assertEigenpairColumnInnerProduct :: String -> (Int, Int, U.Vector Double, U.Vector Double) -> Assertion +assertEigenpairColumnInnerProduct context (leftIndex, rightIndex, leftColumn, rightColumn) = + let !actual = vectorDotU leftColumn rightColumn + !expected = + if leftIndex == rightIndex + then 1.0 + else 0.0 + !limit = 1.0e-6 + in assertBool + ( context + <> " columns are not orthonormal at (" + <> show leftIndex + <> ", " + <> show rightIndex + <> "): " + <> show actual + ) + (abs (actual - expected) <= limit) + +assertBasisColumnsOrthonormal :: String -> Box.Vector (U.Vector Double) -> Assertion +assertBasisColumnsOrthonormal context basisColumns = do + _ <- + traverse + (assertEigenpairColumnInnerProduct context) + [ (leftIndex, rightIndex, leftColumn, rightColumn) + | (leftIndex, leftColumn) <- zip [0 ..] (Box.toList basisColumns), + (rightIndex, rightColumn) <- zip [0 ..] (Box.toList basisColumns), + leftIndex <= rightIndex + ] + pure () + +assertEigenpairResidualsBelow :: String -> Double -> Eigenpairs -> Assertion +assertEigenpairResidualsBelow context residualLimit pairs = + assertBool + (context <> " residuals exceeded " <> show residualLimit <> ": " <> show (U.toList (eigenpairResidualNorms pairs))) + (U.all (\residualNorm -> isFiniteDouble residualNorm && residualNorm <= residualLimit) (eigenpairResidualNorms pairs)) + +vectorDotU :: U.Vector Double -> U.Vector Double -> Double +vectorDotU leftVector rightVector = + U.sum (U.zipWith (*) leftVector rightVector) + +testOperatorShape :: Assertion +testOperatorShape = + extractRight (pathLaplacianLinearOperator 4) $ \operatorValue -> + assertEqual "path operator shape" (4, 4) (operatorShape operatorValue) + +testCSROperatorMatvec :: Assertion +testCSROperatorMatvec = + extractRight (tridiagonalCSR [2.0, 3.0, 4.0] [-1.0, -2.0]) $ \csrValue -> + case runOperatorU (csrLinearOperator csrValue) (U.fromList [1.0, 2.0, 3.0]) of + Left err -> assertFailure ("CSR operator failed: " <> show err) + Right actual -> assertApproxVector [0.0, -1.0, 8.0] actual + +testSigmaIdentityMinusApplies :: Assertion +testSigmaIdentityMinusApplies = + extractRight (diagonalLinearOperator (U.fromList [2.0, 5.0])) $ \operatorValue -> + case runOperatorU (sigmaIdentityMinus 7.0 operatorValue) (U.fromList [3.0, 11.0]) of + Left err -> assertFailure ("sigma identity minus failed: " <> show err) + Right actual -> assertApproxVector [15.0, 22.0] actual + +testZeroScaleEigenpairs :: Assertion +testZeroScaleEigenpairs = + withPositiveCount 2 $ \countValue -> + extractRight (diagonalLinearOperator (U.fromList [2.0, 5.0, 9.0])) $ \operatorValue -> + case solveEigenRequest defaultEigenSolveConfig (addScaledIdentity 4.0 (scaleLinearOperator 0.0 operatorValue)) (EigenpairsRequest SmallestEigenvalues countValue) of + Left err -> assertFailure ("zero-scale eigensolve failed: " <> show err) + Right pairs -> do + assertApproxVector [4.0, 4.0] (eigenpairValues pairs) + assertEqual "ambient dimension" 3 (eigenpairDimension pairs) + assertEqual "pair count" 2 (eigenpairCount pairs) + assertApproxVector [0.0, 0.0] (eigenpairResidualNorms pairs) + +testNegativeScaleEigenvalues :: Assertion +testNegativeScaleEigenvalues = + withPositiveCount 2 $ \countValue -> + extractRight (diagonalLinearOperator (U.fromList [1.0, 3.0, 9.0])) $ \operatorValue -> + case solveEigenRequest defaultEigenSolveConfig (scaleLinearOperator (-2.0) operatorValue) (EigenvaluesRequest SmallestEigenvalues countValue) of + Left err -> assertFailure ("negative-scale eigensolve failed: " <> show err) + Right values -> assertApproxVector [-18.0, -6.0] values + +testDiagonalSpectralDispatch :: Assertion +testDiagonalSpectralDispatch = + withPositiveCount 2 $ \countValue -> + extractRight (diagonalLinearOperator (U.fromList [3.0, -2.0, 7.0, 1.0])) $ \operatorValue -> do + case solveEigenRequest defaultEigenSolveConfig operatorValue (EigenvaluesRequest SmallestEigenvalues countValue) of + Left err -> assertFailure ("diagonal values failed: " <> show err) + Right values -> assertApproxVector [-2.0, 1.0] values + case solveEigenRequest defaultEigenSolveConfig operatorValue (EigenpairsRequest LargestEigenvalues countValue) of + Left err -> assertFailure ("diagonal pairs failed: " <> show err) + Right pairs -> do + assertApproxVector [7.0, 3.0] (eigenpairValues pairs) + assertApproxVector [0.0, 0.0] (eigenpairResidualNorms pairs) + +testPathSpectralDispatch :: Assertion +testPathSpectralDispatch = + withPositiveCount 3 $ \countValue -> + extractRight (pathLaplacianLinearOperator 5) $ \operatorValue -> do + case solveEigenRequest defaultEigenSolveConfig operatorValue (EigenvaluesRequest SmallestEigenvalues countValue) of + Left err -> assertFailure ("path values failed: " <> show err) + Right values -> assertApproxVector (pathLaplacianValues 5 [0, 1, 2]) values + case solveEigenRequest defaultEigenSolveConfig operatorValue (EigenpairsRequest SmallestEigenvalues countValue) of + Left err -> assertFailure ("path pairs failed: " <> show err) + Right pairs -> do + assertEqual "path pair count" 3 (eigenpairCount pairs) + assertScaledEigenpairResiduals "path Laplacian pairs" (pathLaplacianFrobeniusNorm 5) pairs + +testGraphLaplacianSpectralDispatch :: Assertion +testGraphLaplacianSpectralDispatch = + withPositiveCount 3 $ \countValue -> + extractRight + ( graphLaplacianLinearOperator + [0 :: Int, 1, 2] + [ GraphEdge 0 1 1.0, + GraphEdge 1 2 1.0, + GraphEdge 0 2 1.0 + ] + ) + $ \operatorValue -> do + case runOperatorU operatorValue (U.fromList [1.0, 2.0, 4.0]) of + Left err -> assertFailure ("graph Laplacian operator failed: " <> show err) + Right actual -> assertApproxVector [-4.0, -1.0, 5.0] actual + case solveEigenRequest defaultEigenSolveConfig operatorValue (EigenpairsRequest SmallestEigenvalues countValue) of + Left err -> assertFailure ("graph Laplacian eigensolve failed: " <> show err) + Right pairs -> do + assertApproxVector [0.0, 3.0, 3.0] (eigenpairValues pairs) + assertEigenpairResidualsBelow "checked graph Laplacian" 1.0e-10 pairs + assertEigenpairColumnsOrthonormal "checked graph Laplacian" pairs + +testLargeGraphLaplacianCascadicDispatch :: Assertion +testLargeGraphLaplacianCascadicDispatch = + let dimension = 4096 + in withPositiveCount 3 $ \countValue -> + extractRight + ( graphLaplacianLinearOperator + [0 .. dimension - 1] + (successorLikeGraphEdges dimension) + ) + $ \operatorValue -> + case solveEigenRequest defaultEigenSolveConfig operatorValue (EigenpairsRequest SmallestEigenvalues countValue) of + Left err -> assertFailure ("large graph Laplacian eigensolve failed: " <> show err) + Right pairs -> do + let values = U.toList (eigenpairValues pairs) + assertEqual "large graph pair count" 3 (eigenpairCount pairs) + assertEqual "large graph ambient dimension" dimension (eigenpairDimension pairs) + assertEqual "large graph eigenvalue order" (sort values) values + assertEigenpairResidualsBelow "large graph cascadic pairs" 1.0e-5 pairs + assertEigenpairColumnsOrthonormal "large graph cascadic pairs" pairs + +testLargeCycleGraphCascadicSpectrum :: Assertion +testLargeCycleGraphCascadicSpectrum = + let dimension = 4096 :: Int + firstNonzeroEigenvalue = 2.0 - 2.0 * cos (2.0 * pi / fromIntegral dimension) + expectedValues = [0.0, firstNonzeroEigenvalue, firstNonzeroEigenvalue] + cycleEdges = + GraphEdge (dimension - 1) 0 1.0 + : fmap + (\sourceVertex -> GraphEdge sourceVertex (sourceVertex + 1) 1.0) + [0 .. dimension - 2] + in withPositiveCount 3 $ \countValue -> + extractRight + (graphLaplacianLinearOperator [0 .. dimension - 1] cycleEdges) + $ \operatorValue -> + case solveEigenRequest defaultEigenSolveConfig operatorValue (EigenpairsRequest SmallestEigenvalues countValue) of + Left err -> assertFailure ("large cycle graph eigensolve failed: " <> show err) + Right pairs -> do + assertApproxList expectedValues (U.toList (eigenpairValues pairs)) + assertEigenpairResidualsBelow "large cycle cascadic pairs" 1.0e-5 pairs + assertEigenpairColumnsOrthonormal "large cycle cascadic pairs" pairs + +successorLikeGraphEdges :: Int -> [GraphEdge Int] +successorLikeGraphEdges dimension = + localEdges + <> strideEdges 8 0.25 + <> strideEdges 31 0.125 + where + localEdges = + fmap + (\sourceVertex -> GraphEdge sourceVertex (sourceVertex + 1) 1.0) + [0 .. dimension - 2] + strideEdges strideValue edgeWeight = + fmap + (\sourceVertex -> GraphEdge sourceVertex (sourceVertex + strideValue) edgeWeight) + [0, strideValue .. dimension - strideValue - 1] + +testGenericFallbackThresholdDispatch :: Assertion +testGenericFallbackThresholdDispatch = + withPositiveCount 3 $ \countValue -> + withPositiveCount 8 $ \iterationCount -> + extractRight (mkNonNegativeConfigTolerance 1.0e-12) $ \toleranceValue -> do + _ <- + traverse + (assertGenericFallbackThresholdDimension countValue iterationCount toleranceValue) + [ denseSpectralFallbackDimensionThreshold, + denseSpectralFallbackDimensionThreshold + 1 + ] + pure () + +testGenericHighDemandFallbackDispatch :: Assertion +testGenericHighDemandFallbackDispatch = + let dimension = denseSpectralFallbackDimensionThreshold + 1 + in withPositiveCount dimension $ \countValue -> + withPositiveCount 8 $ \iterationCount -> + extractRight (genericThresholdOperator dimension) $ \operatorValue -> do + let solveConfig = + withEigenFallbackLanczosConfig + (withLanczosIterations iterationCount defaultLanczosConfig) + defaultEigenSolveConfig + case solveEigenRequest solveConfig operatorValue (EigenvaluesRequest SmallestEigenvalues countValue) of + Left err -> assertFailure ("generic high-demand values failed: " <> show err) + Right values -> do + assertEqual "generic high-demand value count" dimension (U.length values) + assertApproxVector [1.0, 1.5, 2.25] (U.take 3 values) + +assertGenericFallbackThresholdDimension :: PositiveCount -> PositiveCount -> NonNegativeConfigTolerance -> Int -> Assertion +assertGenericFallbackThresholdDimension countValue iterationCount toleranceValue dimension = + extractRight (genericThresholdOperator dimension) $ \operatorValue -> do + let solveConfig = + withEigenFallbackInitialVector (restartSeedVector dimension) + ( withEigenFallbackLanczosConfig + (withLanczosTolerance toleranceValue (withLanczosIterations iterationCount defaultLanczosConfig)) + defaultEigenSolveConfig + ) + expectedValues = [1.0, 1.5, 2.25] + dispatchContext = "generic fallback n=" <> show dimension + case solveEigenRequest solveConfig operatorValue (EigenvaluesRequest SmallestEigenvalues countValue) of + Left err -> assertFailure (dispatchContext <> " values failed: " <> show err) + Right values -> assertApproxVector expectedValues values + case solveEigenRequest solveConfig operatorValue (EigenpairsRequest SmallestEigenvalues countValue) of + Left err -> assertFailure (dispatchContext <> " pairs failed: " <> show err) + Right pairs -> do + assertApproxVector expectedValues (eigenpairValues pairs) + assertEigenpairResidualsBelow dispatchContext 1.0e-5 pairs + assertEigenpairColumnsOrthonormal dispatchContext pairs + +genericThresholdOperator :: Int -> Either MoonlightError (LinearOperator 'SelfAdjointOperator) +genericThresholdOperator dimension = + declaredSelfAdjointVectorLinearOperator + dimension + (Right . U.imap (\entryIndex entryValue -> genericThresholdEigenvalue entryIndex * entryValue)) + +genericThresholdEigenvalue :: Int -> Double +genericThresholdEigenvalue entryIndex + | entryIndex == 0 = 1.0 + | entryIndex == 1 = 1.5 + | entryIndex == 2 = 2.25 + | otherwise = 10.0 + +testEigenCountRejectsOversubscription :: Assertion +testEigenCountRejectsOversubscription = + withPositiveCount 4 $ \countValue -> + extractRight (diagonalLinearOperator (U.fromList [1.0, 2.0, 3.0])) $ \operatorValue -> + case solveEigenRequest defaultEigenSolveConfig operatorValue (EigenvaluesRequest SmallestEigenvalues countValue) of + Left (InvariantViolation _) -> pure () + Left err -> assertFailure ("expected InvariantViolation, got " <> show err) + Right values -> assertFailure ("expected oversubscription rejection, got " <> show values) + +testZeroDimensionTridiagonalRejected :: Assertion +testZeroDimensionTridiagonalRejected = + extractRight (mkSymmetricTridiagonal [] []) $ \tridiagonalValue -> do + case selectedSymmetricTridiagonalEigenvaluesDirect SmallestEigenvalues 1 tridiagonalValue of + Left (InvariantViolation _) -> pure () + Left err -> assertFailure ("expected eigenvalue InvariantViolation, got " <> show err) + Right values -> assertFailure ("expected zero-dimensional eigenvalue rejection, got " <> show values) + case selectedSymmetricTridiagonalEigenpairsDirect SmallestEigenvalues 1 tridiagonalValue of + Left (InvariantViolation _) -> pure () + Left err -> assertFailure ("expected eigenpair InvariantViolation, got " <> show err) + Right pairs -> assertFailure ("expected zero-dimensional eigenpair rejection, got " <> show pairs) + +testOneDimensionalTridiagonal :: Assertion +testOneDimensionalTridiagonal = + extractRight (mkSymmetricTridiagonal [-3.0] []) $ \tridiagonalValue -> do + case selectedSymmetricTridiagonalEigenvaluesDirect SmallestEigenvalues 1 tridiagonalValue of + Left err -> assertFailure ("one-dimensional values failed: " <> show err) + Right values -> assertApproxVector [-3.0] values + case selectedSymmetricTridiagonalEigenpairsDirect LargestEigenvalues 1 tridiagonalValue of + Left err -> assertFailure ("one-dimensional pairs failed: " <> show err) + Right pairs -> do + assertApproxVector [-3.0] (eigenpairValues pairs) + assertScaledEigenpairResiduals "one-dimensional tridiagonal" 3.0 pairs + +testRepeatedDiagonalProjector :: Assertion +testRepeatedDiagonalProjector = + extractRight (mkSymmetricTridiagonal [2.0, 2.0, 3.0] [0.0, 0.0]) $ \tridiagonalValue -> + case selectedSymmetricTridiagonalEigenpairsDirect SmallestEigenvalues 2 tridiagonalValue of + Left err -> assertFailure ("repeated diagonal pairs failed: " <> show err) + Right pairs -> do + assertApproxVector [2.0, 2.0] (eigenpairValues pairs) + assertScaledEigenpairResiduals "repeated diagonal tridiagonal" (tridiagonalMatrixNorm tridiagonalValue) pairs + projectorDiagonal <- extractEither (eigenpairProjectorDiagonal pairs) + assertApproxVector [1.0, 1.0, 0.0] projectorDiagonal + +testTinyCouplingIsStructural :: Assertion +testTinyCouplingIsStructural = + extractRight (tridiagonalCSR [1.0, 2.0] [1.0e-300]) $ \csrValue -> + case symmetricTridiagonalFromCSR csrValue of + Left err -> assertFailure ("classification failed: " <> show err) + Right (Left rejection) -> assertFailure ("expected accepted tridiagonal, got " <> show rejection) + Right (Right tridiagonalValue) -> + assertEqual "tiny coupling must not be collapsed to structural zero" [1.0e-300] (symmetricTridiagonalOffDiagonalEntries tridiagonalValue) + +testOutOfBandCSRRejected :: Assertion +testOutOfBandCSRRejected = + withCSRFixture 3 3 [(0, 0, 1.0), (0, 2, 1.0e-300), (1, 1, 2.0), (2, 0, 1.0e-300), (2, 2, 3.0)] $ \csrValue -> + case symmetricTridiagonalFromCSR csrValue of + Left err -> assertFailure ("classification failed before rejection: " <> show err) + Right (Left (TridiagonalOutOfBandEntry 0 2)) -> pure () + Right (Left rejection) -> assertFailure ("unexpected rejection: " <> show rejection) + Right (Right tridiagonalValue) -> assertFailure ("expected out-of-band rejection, got " <> show tridiagonalValue) + +testAsymmetricCSRRejected :: Assertion +testAsymmetricCSRRejected = + withCSRFixture 2 2 [(0, 0, 1.0), (0, 1, 1.0), (1, 0, 1.0 + 1.0e-12), (1, 1, 2.0)] $ \csrValue -> + case symmetricTridiagonalFromCSR csrValue of + Left err -> assertFailure ("classification failed before rejection: " <> show err) + Right (Left TridiagonalAsymmetricOffDiagonal) -> pure () + Right (Left rejection) -> assertFailure ("unexpected rejection: " <> show rejection) + Right (Right tridiagonalValue) -> assertFailure ("expected asymmetric rejection, got " <> show tridiagonalValue) + +testReducibleTridiagonalValues :: Assertion +testReducibleTridiagonalValues = + extractRight (mkSymmetricTridiagonal [1.0, 3.0, 2.0, 4.0] [0.5, 0.0, 0.25]) $ \tridiagonalValue -> + case selectedSymmetricTridiagonalEigenvaluesDirect SmallestEigenvalues 3 tridiagonalValue of + Left err -> assertFailure ("reducible tridiagonal solve failed: " <> show err) + Right values -> assertApproxVector (take 3 (sortedValues (twoByTwoSymmetricEigenvalues 1.0 0.5 3.0 <> twoByTwoSymmetricEigenvalues 2.0 0.25 4.0))) values + +testReducibleTridiagonalPairs :: Assertion +testReducibleTridiagonalPairs = + extractRight (mkSymmetricTridiagonal [1.0, 3.0, 2.0, 4.0] [0.5, 0.0, 0.25]) $ \tridiagonalValue -> + case selectedSymmetricTridiagonalEigenpairsDirect SmallestEigenvalues 3 tridiagonalValue of + Left err -> assertFailure ("reducible tridiagonal pairs failed: " <> show err) + Right pairs -> do + assertApproxVector + (take 3 (sortedValues (twoByTwoSymmetricEigenvalues 1.0 0.5 3.0 <> twoByTwoSymmetricEigenvalues 2.0 0.25 4.0))) + (eigenpairValues pairs) + assertScaledEigenpairResiduals "reducible tridiagonal pairs" (tridiagonalMatrixNorm tridiagonalValue) pairs + +testGenericTridiagonalPairs :: Assertion +testGenericTridiagonalPairs = + extractRight (mkSymmetricTridiagonal [2.0, 2.0, 2.0, 2.0] [-0.75, -0.5, -0.25]) $ \tridiagonalValue -> + case selectedSymmetricTridiagonalEigenpairsDirect SmallestEigenvalues 2 tridiagonalValue of + Left err -> assertFailure ("generic tridiagonal pairs failed: " <> show err) + Right pairs -> do + assertEqual "pair dimension" 4 (eigenpairDimension pairs) + assertEqual "pair count" 2 (eigenpairCount pairs) + assertScaledEigenpairResiduals "generic tridiagonal selected inverse-iteration pairs" (tridiagonalMatrixNorm tridiagonalValue) pairs + assertEigenpairColumnsOrthonormal "generic tridiagonal selected inverse-iteration pairs" pairs + +testSelectedTridiagonalPairCertification :: Assertion +testSelectedTridiagonalPairCertification = + extractRight (mkSymmetricTridiagonal [1.0, 2.0, 4.0] [0.0, 0.0]) $ \tridiagonalValue -> + case selectedSymmetricTridiagonalEigenpairsDirect SmallestEigenvalues 2 tridiagonalValue of + Left err -> assertFailure ("selected tridiagonal pairs failed: " <> show err) + Right pairs -> + case certifySelectedEigenpairResult SmallestEigenvalues 2 1.0e-12 1.0e-12 pairs of + Left failureValue -> + assertFailure ("selected tridiagonal pair certification failed: " <> show failureValue) + Right + CertifiedSelectedEigenpairResult + { certifiedSelectedEigenpairResidualEvidence = residualEvidence, + certifiedSelectedEigenpairOrthonormalityEvidence = orthonormalityEvidence, + certifiedSelectedEigenpairRequestOrderingEvidence = requestOrderingEvidence + } -> do + assertEqual + "residual bound evidence" + 1.0e-12 + (selectedEigenpairResidualBound residualEvidence) + assertEqual + "orthonormality bound evidence" + 1.0e-12 + (selectedEigenpairOrthonormalityBound orthonormalityEvidence) + assertEqual + "requested count evidence" + 2 + (selectedEigenpairRequestedCount requestOrderingEvidence) + assertEqual + "certified count evidence" + 2 + (selectedEigenpairCertifiedCount requestOrderingEvidence) + assertEqual + "ordering evidence" + SmallestEigenvalues + (selectedEigenpairCertifiedOrdering requestOrderingEvidence) + case certifySelectedEigenpairResult LargestEigenvalues 2 1.0e-12 1.0e-12 pairs of + Left (SelectedEigenpairCertificationOrderingViolation LargestEigenvalues 0 _ _) -> + pure () + Left failureValue -> + assertFailure ("expected ordering violation, received " <> show failureValue) + Right _ -> + assertFailure "expected largest-order selected certification to reject ascending pairs" + +testGenericSelectedTridiagonalPairsAgreeWithAllPairs :: Assertion +testGenericSelectedTridiagonalPairsAgreeWithAllPairs = + withPositiveCount 5 $ \fullCount -> + extractRight (mkSymmetricTridiagonal [2.0, 2.5, 3.0, 3.5, 4.0] [-0.31, -0.27, -0.23, -0.19]) $ \tridiagonalValue -> do + allPairs <- + selectedSymmetricTridiagonalEigenRequestLapack (EigenpairsRequest SmallestEigenvalues fullCount) tridiagonalValue + >>= extractEither + case selectedSymmetricTridiagonalEigenpairsDirect SmallestEigenvalues 3 tridiagonalValue of + Left err -> assertFailure ("generic selected tridiagonal pairs failed: " <> show err) + Right pairs -> do + assertApproxVector (take 3 (U.toList (eigenpairValues allPairs))) (eigenpairValues pairs) + assertScaledEigenpairResiduals "generic selected tridiagonal all-pairs agreement" (tridiagonalMatrixNorm tridiagonalValue) pairs + assertEigenpairColumnsOrthonormal "generic selected tridiagonal all-pairs agreement" pairs + +testClusteredIrreducibleTridiagonalValues :: Assertion +testClusteredIrreducibleTridiagonalValues = + withPositiveCount 3 $ \countValue -> + extractRight (mkSymmetricTridiagonal [1.0, 1.0 + 1.0e-12, 1.0 + 2.0e-12, 1.0 + 3.0e-12] [1.0e-8, 1.0e-8, 1.0e-8]) $ \tridiagonalValue -> do + nativeValues <- + selectedSymmetricTridiagonalEigenRequestLapack (EigenvaluesRequest SmallestEigenvalues countValue) tridiagonalValue + >>= extractEither + nativePairs <- + selectedSymmetricTridiagonalEigenRequestLapack (EigenpairsRequest SmallestEigenvalues countValue) tridiagonalValue + >>= extractEither + case selectedSymmetricTridiagonalEigenvaluesDirect SmallestEigenvalues 3 tridiagonalValue of + Left err -> assertFailure ("clustered tridiagonal values failed: " <> show err) + Right values -> assertApproxVector (U.toList nativeValues) values + case selectedSymmetricTridiagonalEigenpairsDirect SmallestEigenvalues 3 tridiagonalValue of + Left err -> assertFailure ("clustered tridiagonal pairs failed: " <> show err) + Right pairs -> do + assertApproxVector (U.toList nativeValues) (eigenpairValues pairs) + assertScaledEigenpairResiduals "clustered pure tridiagonal pairs" (tridiagonalMatrixNorm tridiagonalValue) pairs + assertEigenpairColumnsOrthonormal "clustered pure tridiagonal pairs" pairs + assertScaledEigenpairResiduals "clustered native tridiagonal pairs" (tridiagonalMatrixNorm tridiagonalValue) nativePairs + +testGenericIrreducibleLargestValues :: Assertion +testGenericIrreducibleLargestValues = + withPositiveCount 4 $ \countValue -> + extractRight (genericTestTridiagonal 16) $ \tridiagonalValue -> do + nativeValues <- + selectedSymmetricTridiagonalEigenRequestLapack (EigenvaluesRequest LargestEigenvalues countValue) tridiagonalValue + >>= extractEither + case selectedSymmetricTridiagonalEigenvaluesDirect LargestEigenvalues 4 tridiagonalValue of + Left err -> assertFailure ("generic largest tridiagonal values failed: " <> show err) + Right values -> assertApproxVector (U.toList nativeValues) values + +testExtremeScaledIrreducibleTridiagonalValues :: Assertion +testExtremeScaledIrreducibleTridiagonalValues = + withPositiveCount 2 $ \countValue -> + extractRight (mkSymmetricTridiagonal [1.0e100, 2.0e100, 4.0e100] [1.0e90, -1.0e90]) $ \tridiagonalValue -> do + nativeValues <- + selectedSymmetricTridiagonalEigenRequestLapack (EigenvaluesRequest SmallestEigenvalues countValue) tridiagonalValue + >>= extractEither + case selectedSymmetricTridiagonalEigenvaluesDirect SmallestEigenvalues 2 tridiagonalValue of + Left err -> assertFailure ("extreme scaled tridiagonal values failed: " <> show err) + Right values -> assertApproxVector (U.toList nativeValues) values + +testPerturbedPathNotExactPath :: Assertion +testPerturbedPathNotExactPath = + extractRight (mkSymmetricTridiagonal [1.0, 2.0, 1.0] [-0.95, -1.0]) $ \tridiagonalValue -> + case selectedSymmetricTridiagonalEigenvaluesDirect SmallestEigenvalues 3 tridiagonalValue of + Left err -> assertFailure ("perturbed path solve failed: " <> show err) + Right values -> + assertBool + "perturbed path must not reuse exact path spectrum" + (not (and (zipWith approxEqual [0.0, 1.0, 3.0] (U.toList values)))) + +testExtremeAffineScale :: Assertion +testExtremeAffineScale = + withPositiveCount 2 $ \countValue -> + extractRight (diagonalLinearOperator (U.fromList [-1.0e100, 2.0e100, 3.0e100])) $ \operatorValue -> + case solveEigenRequest defaultEigenSolveConfig (addScaledIdentity 5.0 (scaleLinearOperator (-0.5) operatorValue)) (EigenvaluesRequest LargestEigenvalues countValue) of + Left err -> assertFailure ("extreme affine solve failed: " <> show err) + Right values -> assertApproxVector [5.0e99 + 5.0, -1.0e100 + 5.0] values + +testPackedBlockTridiagonalApply :: Assertion +testPackedBlockTridiagonalApply = do + diagonal0 <- extractEither (mkRowMajorBlock 2 2 (U.fromList [2.0, 0.5, 0.5, 3.0])) + diagonal1 <- extractEither (mkRowMajorBlock 1 1 (U.singleton 4.0)) + coupling0 <- extractEither (mkRowMajorBlock 1 2 (U.fromList [1.0, -1.0])) + blockValue <- extractEither (mkSymmetricBlockTridiagonal (Box.fromList [diagonal0, diagonal1]) (Box.singleton coupling0)) + assertEqual "packed block dimension" 3 (symmetricBlockTridiagonalDimension blockValue) + assertEqual "packed block bandwidth" 2 (symmetricBlockTridiagonalBandwidth blockValue) + extractRight (symmetricBlockTridiagonalEntry blockValue 0 2) (assertApproxList [1.0] . pure) + extractRight (symmetricBlockTridiagonalEntry blockValue 1 2) (assertApproxList [-1.0] . pure) + case applySymmetricBlockTridiagonalU blockValue (U.fromList [1.0, 2.0, 3.0]) of + Left err -> assertFailure ("block tridiagonal apply failed: " <> show err) + Right actual -> assertApproxVector [6.0, 3.5, 11.0] actual + +testRowMajorBlockCardinalityOverflow :: Assertion +testRowMajorBlockCardinalityOverflow = + let wrappedDimension = 2 ^ (32 :: Int) + in assertEqual + "oversized block cardinality" + (Left (InvariantViolation "row-major block dimensions exceed Int cardinality")) + (mkRowMajorBlock wrappedDimension wrappedDimension U.empty) + +data GeneratedBlockTridiagonal = GeneratedBlockTridiagonal + { generatedBlockSizes :: [Int], + generatedDiagonalPayloads :: [[Double]], + generatedCouplingPayloads :: [[Double]] + } + deriving stock (Show) + +instance QC.Arbitrary GeneratedBlockTridiagonal where + arbitrary = do + blockCount <- QC.chooseInt (2, 4) + firstSize <- QC.chooseInt (1, 4) + secondSize <- QC.elements (filter (/= firstSize) [1 .. 4]) + remainingSizes <- QC.vectorOf (blockCount - 2) (QC.chooseInt (1, 4)) + let blockSizes = firstSize : secondSize : remainingSizes + diagonalPayloads <- traverse generatedPayload ((\blockSize -> blockSize * blockSize) <$> blockSizes) + couplingPayloads <- traverse generatedPayload (zipWith (*) (drop 1 blockSizes) blockSizes) + pure + GeneratedBlockTridiagonal + { generatedBlockSizes = blockSizes, + generatedDiagonalPayloads = diagonalPayloads, + generatedCouplingPayloads = couplingPayloads + } + where + generatedPayload :: Int -> QC.Gen [Double] + generatedPayload entryCount = + QC.vectorOf entryCount (QC.choose (-8.0, 8.0)) + +propBlockFrobeniusNormMatchesDense :: GeneratedBlockTridiagonal -> QC.Property +propBlockFrobeniusNormMatchesDense GeneratedBlockTridiagonal {..} = + case generatedBlockValue of + Left err -> + QC.counterexample ("generated block construction failed: " <> show err) False + Right blockValue -> + case traverse (uncurry (symmetricBlockTridiagonalEntry blockValue)) denseCoordinates of + Left err -> + QC.counterexample ("generated dense reconstruction failed: " <> show err) False + Right denseEntries -> + QC.counterexample + ("packed norm differs from dense reconstruction for block sizes " <> show generatedBlockSizes) + (approxEqual (denseFrobeniusNorm denseEntries) (symmetricBlockTridiagonalFrobeniusNorm blockValue)) + where + generatedBlockValue = do + diagonalBlocks <- + traverse + (\(blockSize, payload) -> mkRowMajorBlock blockSize blockSize (U.fromList payload) >>= symmetrizeRowMajorBlockLower) + (zip generatedBlockSizes generatedDiagonalPayloads) + couplingBlocks <- + traverse + (\((previousSize, nextSize), payload) -> mkRowMajorBlock nextSize previousSize (U.fromList payload)) + (zip (zip generatedBlockSizes (drop 1 generatedBlockSizes)) generatedCouplingPayloads) + mkSymmetricBlockTridiagonal (Box.fromList diagonalBlocks) (Box.fromList couplingBlocks) + matrixDimension = sum generatedBlockSizes + denseCoordinates = + [ (rowIndex, columnIndex) + | rowIndex <- [0 .. matrixDimension - 1], + columnIndex <- [0 .. matrixDimension - 1] + ] + +testBlockFrobeniusNormPreservesOverflow :: Assertion +testBlockFrobeniusNormPreservesOverflow = + extractRight + ( do + diagonalBlock <- mkRowMajorBlock 1 1 (U.singleton 1.0e308) + mkSymmetricBlockTridiagonal (Box.singleton diagonalBlock) Box.empty + ) + (\blockValue -> + assertBool + "finite entries with an unrepresentable squared norm must not collapse to zero" + (isInfinite (symmetricBlockTridiagonalFrobeniusNorm blockValue)) + ) + +testBlockLanczosStructuredProjection :: Assertion +testBlockLanczosStructuredProjection = + withPositiveCount 3 $ \iterationCount -> + withPositiveCount 2 $ \blockSize -> + extractRight (pathLaplacianCSR 4 >>= selfAdjointCSRLinearOperator) $ \operatorValue -> do + let config = withBlockLanczosBlockSize blockSize (withBlockLanczosIterations iterationCount defaultBlockLanczosConfig) + seedBlock = Box.fromList [U.fromList [1.0, 0.0, 0.0, 0.0], U.fromList [0.0, 1.0, 0.0, 0.0]] + case blockLanczosSymmetric config operatorValue seedBlock of + Left err -> assertFailure ("block Lanczos failed: " <> show err) + Right decomposition -> do + assertEqual "projected block dimension" (blockLanczosBasisCount decomposition) (symmetricBlockTridiagonalDimension (blockLanczosProjectedBlockTridiagonal decomposition)) + assertEqual "projected subspace dimension" (blockLanczosBasisCount decomposition) (symmetricProjectedOperatorDimension (projectedSubspaceOperator (projectedSubspaceFromBlockLanczos decomposition))) + +testProjectedTridiagonalSelectedPairs :: Assertion +testProjectedTridiagonalSelectedPairs = + withPositiveCount 4 $ \iterationCount -> + extractRight (pathLaplacianLinearOperator 8) $ \operatorValue -> + case lanczosSymmetric (withLanczosIterations iterationCount defaultLanczosConfig) operatorValue (U.fromList [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) of + Left err -> assertFailure ("Lanczos failed: " <> show err) + Right decomposition -> do + let subspace = projectedSubspaceFromLanczos decomposition + case projectedSubspaceOperator subspace of + TridiagonalProjectedOperator tridiagonalValue -> + case selectedSymmetricTridiagonalEigenpairsDirect SmallestEigenvalues 2 tridiagonalValue of + Left err -> assertFailure ("selected projected tridiagonal reference failed: " <> show err) + Right selectedPairs -> + case projectedEigenpairs SmallestEigenvalues 2 operatorValue subspace of + Left err -> assertFailure ("projected selected pairs failed: " <> show err) + Right projectedPairs -> + assertApproxVector (U.toList (eigenpairValues selectedPairs)) (eigenpairValues projectedPairs) + BlockTridiagonalProjectedOperator blockValue -> + assertFailure ("expected tridiagonal projection, got block projection " <> show blockValue) + +testProjectedCountRejectsOversubscription :: Assertion +testProjectedCountRejectsOversubscription = + withPositiveCount 1 $ \iterationCount -> + extractRight (pathLaplacianLinearOperator 4) $ \operatorValue -> + case lanczosSymmetric (withLanczosIterations iterationCount defaultLanczosConfig) operatorValue (U.fromList [1.0, 0.0, 0.0, 0.0]) of + Left err -> assertFailure ("Lanczos failed: " <> show err) + Right decomposition -> + let subspace = projectedSubspaceFromLanczos decomposition + oversubscribedCount = projectedSubspaceDimension subspace + 1 + in case projectedEigenvalues SmallestEigenvalues oversubscribedCount operatorValue subspace of + Left (InvariantViolation _) -> pure () + Left err -> assertFailure ("expected InvariantViolation, got " <> show err) + Right values -> assertFailure ("expected projected oversubscription rejection, got " <> show values) + +testLanczosProjection :: Assertion +testLanczosProjection = + withPositiveCount 3 $ \iterationCount -> + extractRight (pathLaplacianLinearOperator 4) $ \operatorValue -> + case lanczosSymmetric (withLanczosIterations iterationCount defaultLanczosConfig) operatorValue (U.fromList [1.0, 0.0, 0.0, 0.0]) of + Left err -> assertFailure ("Lanczos failed: " <> show err) + Right decomposition -> do + assertEqual "Lanczos projected dimension" (Box.length (lanczosBasisColumns decomposition)) (symmetricTridiagonalDimension (lanczosProjectedTridiagonal decomposition)) + assertBool "Lanczos completed at least one step" (lanczosStepsCompleted decomposition > 0) + +testLanczosThickRestartMultiCycle :: Assertion +testLanczosThickRestartMultiCycle = + withPositiveCount 3 $ \requestedCount -> + withRestartedLanczosFixture 18 5 $ \operatorValue lanczosConfig seedVector -> do + let solveConfig = + withEigenFallbackInitialVector seedVector + (withEigenFallbackLanczosConfig lanczosConfig defaultEigenSolveConfig) + case solveEigenRequest solveConfig operatorValue (EigenpairsRequest SmallestEigenvalues requestedCount) of + Left err -> assertFailure ("thick restart spectral fallback failed: " <> show err) + Right pairs -> do + assertEqual "multi-cycle pair count" 3 (eigenpairCount pairs) + assertEigenpairResidualsBelow "multi-cycle restarted Lanczos" 1.0e-5 pairs + assertEigenpairColumnsOrthonormal "multi-cycle restarted Lanczos" pairs + +testLanczosThickRestartLockedResiduals :: Assertion +testLanczosThickRestartLockedResiduals = + withRestartedLanczosFixture 20 6 $ \operatorValue lanczosConfig seedVector -> + case projectedEigenpairsFromRestartedLanczos lanczosConfig SmallestEigenvalues 4 operatorValue seedVector of + Left err -> assertFailure ("restarted projected eigenpairs failed: " <> show err) + Right pairs -> do + assertEqual "locked residual pair count" 4 (eigenpairCount pairs) + assertEigenpairResidualsBelow "locked restarted Lanczos" 1.0e-5 pairs + +testLanczosThickRestartProjectedResidualEvidence :: Assertion +testLanczosThickRestartProjectedResidualEvidence = + withRestartedLanczosFixture 18 5 $ \operatorValue lanczosConfig seedVector -> + case lanczosRestartedProjection lanczosConfig SmallestEigenvalues 3 operatorValue seedVector of + Left err -> assertFailure ("restarted Lanczos projection failed: " <> show err) + Right restartProjection -> do + let projectedResiduals = eigenpairResidualNorms (lanczosRestartProjectionProjectedPairs restartProjection) + assertEqual "projected residual evidence count" 3 (U.length projectedResiduals) + assertBool + ("projected residual evidence must be finite: " <> show (U.toList projectedResiduals)) + (U.all isFiniteDouble projectedResiduals) + assertBool + ("projected residual evidence must not be stamped as zero: " <> show (U.toList projectedResiduals)) + (U.any (> 0.0) projectedResiduals) + +testLanczosThickRestartBasisOrthogonality :: Assertion +testLanczosThickRestartBasisOrthogonality = + withRestartedLanczosFixture 16 5 $ \operatorValue lanczosConfig seedVector -> + case lanczosRestartedProjection lanczosConfig SmallestEigenvalues 3 operatorValue seedVector of + Left err -> assertFailure ("restarted Lanczos projection failed: " <> show err) + Right restartProjection -> + assertBasisColumnsOrthonormal + "restarted Lanczos locked/active basis" + (lanczosRestartProjectionBasisColumns restartProjection) + +testNativeDenseSelectedValuesAgreeWithPairs :: Assertion +testNativeDenseSelectedValuesAgreeWithPairs = + withPositiveCount 2 $ \countValue -> do + matrixValue <- extractEither (mkDynMatrix 3 3 [2.0, 1.0, 0.0, 1.0, 2.0, 0.0, 0.0, 0.0, 5.0]) + values <- + symmetricEigenRequestLapack (EigenvaluesRequest SmallestEigenvalues countValue) matrixValue + >>= extractEither + pairs <- + symmetricEigenRequestLapack (EigenpairsRequest SmallestEigenvalues countValue) matrixValue + >>= extractEither + assertApproxVector (U.toList (eigenpairValues pairs)) values + assertScaledEigenpairResiduals "native dense selected pairs" (denseFrobeniusNorm [2.0, 1.0, 0.0, 1.0, 2.0, 0.0, 0.0, 0.0, 5.0]) pairs + +testNativeGenericTridiagonalSelected :: Assertion +testNativeGenericTridiagonalSelected = + withPositiveCount 4 $ \countValue -> + extractRight (genericTestTridiagonal 16) $ \tridiagonalValue -> do + pureValues <- + extractEither + (selectedSymmetricTridiagonalEigenvaluesDirect SmallestEigenvalues 4 tridiagonalValue) + nativeValues <- + selectedSymmetricTridiagonalEigenRequestLapack (EigenvaluesRequest SmallestEigenvalues countValue) tridiagonalValue + >>= extractEither + nativePairs <- + selectedSymmetricTridiagonalEigenRequestLapack (EigenpairsRequest SmallestEigenvalues countValue) tridiagonalValue + >>= extractEither + assertApproxVector (U.toList pureValues) nativeValues + assertApproxVector (U.toList nativeValues) (eigenpairValues nativePairs) + assertScaledEigenpairResiduals "native generic tridiagonal selected pairs" (tridiagonalMatrixNorm tridiagonalValue) nativePairs + +testNativeBlockBandSelectedAgreesWithDense :: Assertion +testNativeBlockBandSelectedAgreesWithDense = + withPositiveCount 2 $ \countValue -> do + blockValue <- nativeBlockFixture + denseMatrix <- extractEither (mkDynMatrix 3 3 [2.0, 0.5, 1.0, 0.5, 3.0, -1.0, 1.0, -1.0, 4.0]) + denseValues <- + symmetricEigenRequestLapack (EigenvaluesRequest SmallestEigenvalues countValue) denseMatrix + >>= extractEither + blockValues <- + selectedSymmetricBlockTridiagonalEigenRequestLapack (EigenvaluesRequest SmallestEigenvalues countValue) blockValue + >>= extractEither + blockPairs <- + selectedSymmetricBlockTridiagonalEigenRequestLapack (EigenpairsRequest SmallestEigenvalues countValue) blockValue + >>= extractEither + assertApproxVector (U.toList denseValues) blockValues + assertApproxVector (U.toList denseValues) (eigenpairValues blockPairs) + assertScaledEigenpairResiduals "native block-band selected pairs" (denseFrobeniusNorm [2.0, 0.5, 1.0, 0.5, 3.0, -1.0, 1.0, -1.0, 4.0]) blockPairs + +pathLaplacianValues :: Int -> [Int] -> [Double] +pathLaplacianValues dimension = + fmap (\modeIndex -> 2.0 - 2.0 * cos (pi * fromIntegral modeIndex / fromIntegral dimension)) + +genericTestTridiagonal :: Int -> Either MoonlightError SymmetricTridiagonal +genericTestTridiagonal dimension = + mkSymmetricTridiagonal + (genericTestTridiagonalDiagonalEntry <$> [0 .. dimension - 1]) + (genericTestTridiagonalOffDiagonalEntry <$> [0 .. dimension - 2]) + +genericTestTridiagonalDiagonalEntry :: Int -> Double +genericTestTridiagonalDiagonalEntry indexValue = + 2.0 + fromIntegral (indexValue `mod` 17) / 17.0 + +genericTestTridiagonalOffDiagonalEntry :: Int -> Double +genericTestTridiagonalOffDiagonalEntry indexValue = + -0.35 - 0.01 * fromIntegral (indexValue `mod` 5) + +withRestartedLanczosFixture :: + Int -> + Int -> + (LinearOperator 'SelfAdjointOperator -> LanczosConfig -> U.Vector Double -> Assertion) -> + Assertion +withRestartedLanczosFixture dimension iterationLimit onFixture = + withPositiveCount iterationLimit $ \iterationCount -> + extractRight (mkNonNegativeConfigTolerance 1.0e-8) $ \toleranceValue -> + extractRight (genericPentadiagonalCSR dimension >>= selfAdjointCSRLinearOperator) $ \operatorValue -> + onFixture + operatorValue + (withLanczosTolerance toleranceValue (withLanczosIterations iterationCount defaultLanczosConfig)) + (restartSeedVector dimension) + +genericPentadiagonalCSR :: Int -> Either MoonlightError (SparseCSR Double) +genericPentadiagonalCSR dimension = + mkSparseCOO dimension dimension (genericPentadiagonalEntries dimension) >>= cooToCSR + +genericPentadiagonalEntries :: Int -> [(Int, Int, Double)] +genericPentadiagonalEntries dimension = + diagonalEntries <> firstOffDiagonalEntries <> secondOffDiagonalEntries + where + diagonalEntries = + (\rowIndex -> (rowIndex, rowIndex, 4.0 + 0.03 * fromIntegral (rowIndex `mod` 7))) + <$> [0 .. dimension - 1] + firstOffDiagonalEntries = + symmetricBandEntries dimension 1 (\rowIndex -> -1.0 - 0.01 * fromIntegral (rowIndex `mod` 5)) + secondOffDiagonalEntries = + symmetricBandEntries dimension 2 (\rowIndex -> -0.2 - 0.005 * fromIntegral (rowIndex `mod` 3)) + +symmetricBandEntries :: Int -> Int -> (Int -> Double) -> [(Int, Int, Double)] +symmetricBandEntries dimension offset entryAt = + concatMap + ( \rowIndex -> + let columnIndex = rowIndex + offset + entryValue = entryAt rowIndex + in [(rowIndex, columnIndex, entryValue), (columnIndex, rowIndex, entryValue)] + ) + [0 .. dimension - offset - 1] + +restartSeedVector :: Int -> U.Vector Double +restartSeedVector dimension = + U.generate dimension (\indexValue -> 1.0 / fromIntegral (indexValue + 1)) + +twoByTwoSymmetricEigenvalues :: Double -> Double -> Double -> [Double] +twoByTwoSymmetricEigenvalues a b d = + let traceHalf = 0.5 * (a + d) + radius = sqrt (((a - d) * 0.5) * ((a - d) * 0.5) + b * b) + in [traceHalf - radius, traceHalf + radius] + +sortedValues :: [Double] -> [Double] +sortedValues = sort + +withCSRFixture :: Int -> Int -> [(Int, Int, Double)] -> (SparseCSR Double -> Assertion) -> Assertion +withCSRFixture rowCount columnCount entries onCSR = + case mkSparseCOO rowCount columnCount entries >>= cooToCSR of + Left err -> assertFailure ("invalid sparse fixture: " <> show err) + Right csrValue -> onCSR csrValue + +extractEither :: Either MoonlightError value -> IO value +extractEither value = + case value of + Left err -> assertFailure ("expected Right, got " <> show err) + Right resultValue -> pure resultValue + +nativeBlockFixture :: IO SymmetricBlockTridiagonal +nativeBlockFixture = do + diagonal0 <- extractEither (mkRowMajorBlock 2 2 (U.fromList [2.0, 0.5, 0.5, 3.0])) + diagonal1 <- extractEither (mkRowMajorBlock 1 1 (U.singleton 4.0)) + coupling0 <- extractEither (mkRowMajorBlock 1 2 (U.fromList [1.0, -1.0])) + extractEither (mkSymmetricBlockTridiagonal (Box.fromList [diagonal0, diagonal1]) (Box.singleton coupling0)) diff --git a/test/statics/StaticsSpec.hs b/test/statics/StaticsSpec.hs new file mode 100644 --- /dev/null +++ b/test/statics/StaticsSpec.hs @@ -0,0 +1,282 @@ +module StaticsSpec + ( tests, + ) +where + +import qualified Data.List.NonEmpty as NonEmpty +import qualified Data.Map.Strict as Map +import Moonlight.LinAlg + ( Axis (..), + EquationRef (..), + EquilibriumResult (..), + EquilibriumSolution (..), + EquilibriumViolation (..), + ForceNetwork, + ForceSign (..), + MemberRef, + NetworkBuildError (..), + NetworkDeclaration, + NodeRef, + UnknownForce (..), + Vec3 (..), + assembleEquilibriumEquations, + checkEquilibrium, + compiledEquationOrder, + compiledFoundationOrder, + compiledMemberOrder, + compiledNodeOrder, + compiledUnknownOrder, + joint, + load, + member, + mkMemberRef, + mkSupportAxes, + network, + networkNodeMap, + nodeLoad, + nodeRef, + support, + supportOn, + ) +import Helpers (extractRight) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit + ( Assertion, + assertBool, + assertEqual, + assertFailure, + testCase, + ) + +tests :: TestTree +tests = + testGroup + "Statics" + [ testCase "assembleEquilibriumEquations uses canonical NodeRef ordering" testCanonicalAssembly, + testCase "assembleEquilibriumEquations emits reactions only for supported axes" testAxisSpecificSupportUnknowns, + testCase "network declarations are order-independent" testDeclarationOrderIndependence, + testCase "network repeated loads accumulate exactly before rounding" testRepeatedLoadsUseOneRoundedExactSum, + testCase "network rejects conflicting node positions" testConflictingPositionsFail, + testCase "network rejects unknown member endpoints" testUnknownEndpointFails, + testCase "network rejects degenerate members" testDegenerateMemberFails, + testCase "checkEquilibrium resolves a compressive vertical support" testVerticalEquilibrium, + testCase "checkEquilibrium rejects tension-only hanging support" testCompressionOnlyViolation, + testCase "checkEquilibrium reports residual force for unsupported load" testResidualViolation + ] + +testCanonicalAssembly :: Assertion +testCanonicalAssembly = do + leftMember <- expectMember "c" "a" + rightMember <- expectMember "b" "c" + nodeA <- expectNodeRef "a" + nodeB <- expectNodeRef "b" + nodeC <- expectNodeRef "c" + networkValue <- + expectNetwork + [ member "c" "a", + joint "c" (Vec3 0.0 1.0 0.0), + support "a" (Vec3 (-1.0) 0.0 0.0), + member "b" "c", + support "b" (Vec3 1.0 0.0 0.0) + ] + extractRight (assembleEquilibriumEquations networkValue) $ \compiledValue -> do + assertEqual "node order" [nodeA, nodeB, nodeC] (compiledNodeOrder compiledValue) + assertEqual "foundation order" [nodeA, nodeB] (compiledFoundationOrder compiledValue) + assertEqual "member order" [leftMember, rightMember] (compiledMemberOrder compiledValue) + assertEqual + "equation order begins with first node" + [ EquationRef nodeA AxisX, + EquationRef nodeA AxisY, + EquationRef nodeA AxisZ + ] + (take 3 (compiledEquationOrder compiledValue)) + +testAxisSpecificSupportUnknowns :: Assertion +testAxisSpecificSupportUnknowns = do + foundationRef <- expectNodeRef "foundation" + networkValue <- + expectNetwork + [ supportOn + "foundation" + (Vec3 0.0 0.0 0.0) + (mkSupportAxes [AxisY]) + ] + extractRight (assembleEquilibriumEquations networkValue) $ \compiledValue -> + assertEqual + "axis-specific support emits only its declared reaction unknown" + [ReactionUnknown foundationRef AxisY] + (compiledUnknownOrder compiledValue) + +testDeclarationOrderIndependence :: Assertion +testDeclarationOrderIndependence = do + let declarations = + [ support "a" (Vec3 0.0 0.0 0.0), + load "b" (Vec3 0.0 1.0 0.0) (Vec3 0.0 (-4.0) 0.0), + load "b" (Vec3 0.0 1.0 0.0) (Vec3 0.0 (-6.0) 0.0), + member "a" "b" + ] + forwardNetwork <- expectNetwork declarations + reverseNetwork <- expectNetwork (reverse declarations) + assertEqual "declaration order" forwardNetwork reverseNetwork + +testRepeatedLoadsUseOneRoundedExactSum :: Assertion +testRepeatedLoadsUseOneRoundedExactSum = do + pointReference <- expectNodeRef "p" + networkValue <- + expectNetwork + [ load "p" (Vec3 0.0 0.0 0.0) (Vec3 1.0e16 0.0 0.0), + load "p" (Vec3 0.0 0.0 0.0) (Vec3 1.0 0.0 0.0), + load "p" (Vec3 0.0 0.0 0.0) (Vec3 (-1.0e16) 0.0 0.0) + ] + case Map.lookup pointReference (networkNodeMap networkValue) of + Nothing -> assertFailure "expected node p" + Just nodeValue -> + assertEqual + "exactly accumulated load" + (Vec3 1.0 0.0 0.0) + (nodeLoad nodeValue) + +testConflictingPositionsFail :: Assertion +testConflictingPositionsFail = + case + network + [ joint "p" (Vec3 0.0 0.0 0.0), + load "p" (Vec3 1.0 0.0 0.0) (Vec3 0.0 1.0 0.0) + ] + of + Left (ConflictingNodePosition "p" (Vec3 0.0 0.0 0.0) (Vec3 1.0 0.0 0.0)) -> + pure () + Left other -> + assertFailure ("unexpected construction error: " <> show other) + Right _ -> + assertFailure "expected conflicting positions to fail" + +testUnknownEndpointFails :: Assertion +testUnknownEndpointFails = + case + network + [ joint "a" (Vec3 0.0 0.0 0.0), + member "a" "b" + ] + of + Left (UnknownMemberEndpoint "b") -> + pure () + Left other -> + assertFailure ("unexpected construction error: " <> show other) + Right _ -> + assertFailure "expected unknown member endpoint to fail" + +testDegenerateMemberFails :: Assertion +testDegenerateMemberFails = + case + network + [ joint "a" (Vec3 0.0 0.0 0.0), + joint "b" (Vec3 0.0 0.0 0.0), + member "a" "b" + ] + of + Left (DegenerateMember "a" "b") -> + pure () + Left other -> + assertFailure ("unexpected construction error: " <> show other) + Right _ -> + assertFailure "expected zero-length member to fail" + +testVerticalEquilibrium :: Assertion +testVerticalEquilibrium = do + memberRefValue <- expectMember "foundation" "load" + foundationRef <- expectNodeRef "foundation" + networkValue <- + expectNetwork + [ load "load" (Vec3 0.0 1.0 0.0) (Vec3 0.0 (-10.0) 0.0), + supportOn "foundation" (Vec3 0.0 0.0 0.0) (mkSupportAxes [AxisY]), + member "foundation" "load" + ] + extractRight (checkEquilibrium networkValue) $ \equilibriumResult -> + case equilibriumResult of + InEquilibrium solutionValue -> do + assertApprox "member force" 10.0 (Map.findWithDefault 0.0 memberRefValue (equilibriumMemberForces solutionValue)) + assertVec3Approx + "foundation reaction" + (Vec3 0.0 10.0 0.0) + (Map.findWithDefault (Vec3 0.0 0.0 0.0) foundationRef (equilibriumReactionForces solutionValue)) + Disequilibrium violations -> + assertBool ("expected equilibrium, got " <> show violations) False + +testCompressionOnlyViolation :: Assertion +testCompressionOnlyViolation = do + networkValue <- + expectNetwork + [ support "left" (Vec3 (-1.0) 1.0 0.0), + support "right" (Vec3 1.0 1.0 0.0), + load "load" (Vec3 0.0 0.0 0.0) (Vec3 0.0 (-10.0) 0.0), + member "left" "load", + member "right" "load" + ] + extractRight (checkEquilibrium networkValue) $ \equilibriumResult -> + case equilibriumResult of + InEquilibrium solutionValue -> + assertBool ("expected compression-only violation, got " <> show solutionValue) False + Disequilibrium violations -> + assertBool + "expected at least one tension violation" + ( any + ((== Just Tension) . violationMemberForceSign) + (NonEmpty.toList violations) + ) + +testResidualViolation :: Assertion +testResidualViolation = do + networkValue <- + expectNetwork + [ supportOn "foundation" (Vec3 0.0 0.0 0.0) (mkSupportAxes [AxisY]), + load "load" (Vec3 1.0 0.0 0.0) (Vec3 0.0 (-10.0) 0.0), + member "foundation" "load" + ] + extractRight (checkEquilibrium networkValue) $ \equilibriumResult -> + case equilibriumResult of + InEquilibrium solutionValue -> + assertBool ("expected residual violation, got " <> show solutionValue) False + Disequilibrium violations -> + assertBool + "expected non-zero residual" + ( any + ((> 1.0e-6) . violationResidualMagnitude) + (NonEmpty.toList violations) + ) + +expectNetwork :: [NetworkDeclaration] -> IO ForceNetwork +expectNetwork declarations = + case network declarations of + Left buildError -> + assertFailure ("expected valid network, got " <> show buildError) + Right networkValue -> + pure networkValue + +expectNodeRef :: String -> IO NodeRef +expectNodeRef labelValue = + case nodeRef labelValue of + Left buildError -> + assertFailure ("expected valid node reference, got " <> show buildError) + Right nodeReference -> + pure nodeReference + +expectMember :: String -> String -> IO MemberRef +expectMember leftLabel rightLabel = do + leftReference <- expectNodeRef leftLabel + rightReference <- expectNodeRef rightLabel + case mkMemberRef leftReference rightReference of + Left err -> assertFailure ("expected valid member, got " <> show err) + Right memberRefValue -> pure memberRefValue + +assertApprox :: String -> Double -> Double -> Assertion +assertApprox message expected actual = + assertBool + (message <> ": expected " <> show expected <> " but received " <> show actual) + (abs (expected - actual) <= 1.0e-6) + +assertVec3Approx :: String -> Vec3 -> Vec3 -> Assertion +assertVec3Approx message expected actual = do + assertApprox (message <> " x") (vecX expected) (vecX actual) + assertApprox (message <> " y") (vecY expected) (vecY actual) + assertApprox (message <> " z") (vecZ expected) (vecZ actual) diff --git a/test/support/Helpers.hs b/test/support/Helpers.hs new file mode 100644 --- /dev/null +++ b/test/support/Helpers.hs @@ -0,0 +1,12 @@ +module Helpers + ( extractRight, + ) +where + +import Test.Tasty.HUnit (Assertion, assertFailure) + +extractRight :: Show e => Either e a -> (a -> Assertion) -> Assertion +extractRight value onRight = + case value of + Left err -> assertFailure ("expected Right, got Left: " <> show err) + Right rightValue -> onRight rightValue