diff --git a/LICENSE b/LICENSE
new file mode 100644
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,27 @@
+Copyright (c) Henning Thielemann 2014
+
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions
+are met:
+1. Redistributions of source code must retain the above copyright
+   notice, this list of conditions and the following disclaimer.
+2. Redistributions in binary form must reproduce the above copyright
+   notice, this list of conditions and the following disclaimer in the
+   documentation and/or other materials provided with the distribution.
+3. Neither the name of the author nor the names of his contributors
+   may be used to endorse or promote products derived from this software
+   without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
+OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
+HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+SUCH DAMAGE.
diff --git a/Setup.lhs b/Setup.lhs
new file mode 100644
--- /dev/null
+++ b/Setup.lhs
@@ -0,0 +1,3 @@
+#! /usr/bin/env runhaskell
+> import Distribution.Simple
+> main = defaultMain
diff --git a/knead-arithmetic.cabal b/knead-arithmetic.cabal
new file mode 100644
--- /dev/null
+++ b/knead-arithmetic.cabal
@@ -0,0 +1,38 @@
+Name:             knead-arithmetic
+Version:          0.0
+License:          BSD3
+License-File:     LICENSE
+Author:           Henning Thielemann <haskell@henning-thielemann.de>
+Maintainer:       Henning Thielemann <haskell@henning-thielemann.de>
+Homepage:         http://hub.darcs.net/thielema/knead-arithmetic/
+Category:         Math
+Synopsis:         Linear algebra and interpolation using LLVM JIT
+Description:
+  Linear algebra and interpolation using LLVM JIT via the @knead@ package.
+Tested-With:      GHC==7.4.2, GHC==7.8.4, GHC==8.0.1
+Cabal-Version:    >=1.6
+Build-Type:       Simple
+
+Source-Repository this
+  Tag:         0.0
+  Type:        darcs
+  Location:    http://hub.darcs.net/thielema/knead-arithmetic/
+
+Source-Repository head
+  Type:        darcs
+  Location:    http://hub.darcs.net/thielema/knead-arithmetic/
+
+Library
+  Build-Depends:
+    knead >=0.2 && <0.3,
+    llvm-extra >=0.6 && <0.7,
+    llvm-tf >=3.0.3 && <3.0.4,
+    utility-ht >=0.0.8 && <0.1,
+    base >=4.5 && <5
+
+  GHC-Options:      -Wall
+  Hs-Source-Dirs:   src
+  Exposed-Modules:
+    Data.Array.Knead.Arithmetic.LinearAlgebra
+    Data.Array.Knead.Arithmetic.Sparse
+    Data.Array.Knead.Arithmetic.Interpolation
diff --git a/src/Data/Array/Knead/Arithmetic/Interpolation.hs b/src/Data/Array/Knead/Arithmetic/Interpolation.hs
new file mode 100644
--- /dev/null
+++ b/src/Data/Array/Knead/Arithmetic/Interpolation.hs
@@ -0,0 +1,226 @@
+{-# LANGUAGE TypeFamilies #-}
+module Data.Array.Knead.Arithmetic.Interpolation (
+   bisect,
+   lookupInterval,
+   Interpolator13, sampleBasisFunctions13,
+   ) where
+
+import qualified Data.Array.Knead.Arithmetic.LinearAlgebra as LinAlg
+import qualified Data.Array.Knead.Arithmetic.Sparse as Sparse
+import Data.Array.Knead.Arithmetic.LinearAlgebra
+          (Scalar, Vector, Matrix, IOScalar)
+
+import qualified Data.Array.Knead.Parameterized.Physical as Phys
+import qualified Data.Array.Knead.Parameterized.Symbolic as SymP
+import qualified Data.Array.Knead.Simple.Physical as SimPhys
+import qualified Data.Array.Knead.Simple.ShapeDependent as ShapeDep
+import qualified Data.Array.Knead.Simple.Symbolic as Sym
+import qualified Data.Array.Knead.Index.Nested.Shape as Shape
+import qualified Data.Array.Knead.Expression as Expr
+import Data.Array.Knead.Expression (Exp)
+
+import qualified LLVM.Extra.Multi.Value.Memory as MultiValueMemory
+import qualified LLVM.Extra.Multi.Value as MultiValue
+import LLVM.Extra.Multi.Value (atom)
+
+import qualified LLVM.Core as LLVM
+
+import Foreign.Storable (Storable)
+
+import Control.Arrow (arr)
+import Control.Monad.HT (chain)
+import Control.Applicative (pure)
+
+import qualified Data.List.Match as Match
+
+
+bisect ::
+   (Shape.C coll, Shape.C nodes, Shape.Index nodes ~ i,
+    MultiValue.IntegerConstant i, MultiValue.Integral i,
+    MultiValue.Select i,
+    MultiValue.Comparison a) =>
+   Vector p coll nodes a ->
+   Scalar p coll a ->
+   Scalar p coll (i, i) ->
+   Scalar p coll (i, i)
+bisect nodes xs bounds =
+   let centers =
+          Sym.map
+             (Expr.modify (atom, atom) $ \(lower, upper) ->
+                Expr.idiv (Expr.add lower upper) $ Expr.fromInteger' 2)
+             bounds
+   in  Sym.zipWith3
+          (Expr.liftM3 $ \center interval leftBranch ->
+              MultiValue.select leftBranch
+                 (MultiValue.mapSnd (const center) interval)
+                 (MultiValue.mapFst (const center) interval))
+          centers bounds $
+       Sym.zipWith (Expr.liftM2 $ MultiValue.cmp LLVM.CmpLT) xs $
+       Sym.gather (Sym.mapWithIndex Expr.zip centers) nodes
+
+nestLog2 ::
+   (Integral i, Monad m) =>
+   i -> (a -> m a) -> a -> m a
+nestLog2 i f =
+   chain $ Match.replicate (takeWhile (>1) $ iterate (flip div 2) i) f
+
+lookupInterval ::
+   (Shape.C coll, Shape.C nodes, Shape.Index nodes ~ i, nodes ~ i,
+    MultiValue.IntegerConstant i, MultiValue.Integral i,
+    MultiValue.Select i, Num i,
+    MultiValue.Comparison a,
+    MultiValueMemory.C nodes, Storable nodes,
+    MultiValueMemory.Struct nodes ~ nodesStruct, LLVM.IsSized nodesStruct,
+    MultiValueMemory.C i, Storable i,
+    MultiValueMemory.Struct i ~ iStruct, LLVM.IsSized iStruct,
+    MultiValueMemory.C coll, Storable coll,
+    MultiValueMemory.C a,
+    Storable a) =>
+   Vector p coll nodes a ->
+   Scalar p coll a ->
+   IOScalar p coll i
+lookupInterval nodes x = do
+   fill <-
+      Phys.render $
+      SymP.fill (arr fst) (fmap ((,) 0) $ arr snd)
+   bis <-
+      Phys.render $
+      bisect
+         (SymP.extendParameter fst nodes)
+         (SymP.extendParameter fst x)
+         (Phys.feed $ arr snd)
+   getFst <- Phys.render $ Sym.map Expr.fst $ Phys.feed $ arr id
+   getNodesShape <- Phys.renderShape nodes
+   getXShape <- Phys.renderShape x
+   return $ \p -> do
+      (_,numElems) <- getNodesShape p
+      (xShape,_) <- getXShape p
+      getFst =<<
+         nestLog2 numElems (curry bis p) =<<
+         fill (xShape, fromIntegral numElems)
+
+
+outerVector ::
+   (Shape.C coll, Shape.C dim) =>
+   (Exp a -> Exp b -> Exp c) ->
+   Scalar p coll a -> SymP.Array p dim b -> Vector p coll dim c
+outerVector =
+   ShapeDep.backpermute2 Expr.zip Expr.fst Expr.snd
+
+
+zipWithScalar ::
+   (Shape.C shape) =>
+   (Exp a -> Exp b -> Exp c) ->
+   SymP.Array p () a -> SymP.Array p shape b -> SymP.Array p shape c
+zipWithScalar =
+   ShapeDep.backpermute2
+      (flip const)
+      (const Expr.unit)
+      id
+
+
+{- |
+One node before index 0 and three nodes starting from index 0.
+-}
+type Interpolator13 a = (a,a) -> (a,a) -> (a,a) -> (a,a) -> a -> a
+
+sampleBasisFunctions13Aux ::
+   (Shape.C coll, Shape.C rows, Shape.C nodes,
+    Shape.C set, MultiValueMemory.C set, Storable set, Num set,
+    Shape.Index nodes ~ i,
+    MultiValue.Comparison i, MultiValue.PseudoRing i,
+    MultiValue.IntegerConstant i,
+    MultiValueMemory.C i, Storable i, Num i,
+    MultiValue.Select a, MultiValue.Real a,
+    MultiValue.Field a, MultiValue.RationalConstant a,
+    Num a, Storable a, MultiValueMemory.C a,
+    MultiValueMemory.Struct a ~ astruct, LLVM.IsSized astruct,
+    MultiValueMemory.Struct i ~ istruct, LLVM.IsSized istruct) =>
+   Interpolator13 (Exp a) ->
+   SymP.Array p () (i,i) ->
+   Vector p coll rows i ->
+   SymP.Array p nodes a ->
+   Vector p coll rows a ->
+   IO (Matrix p coll rows set (i, a))
+sampleBasisFunctions13Aux interpolate minMaxIx indices nodes zs = do
+   let limitIndices =
+          zipWithScalar
+             (\mm ->
+                case Expr.unzip mm of
+                   (minIx,maxIx) -> Expr.max minIx . Expr.min maxIx)
+             minMaxIx indices
+       gatherFromNodes d =
+          Sym.gather (Sym.map (d+) limitIndices) nodes
+   units <-
+      SimPhys.vectorFromList
+         [(-1, (1,0,0,0)), (0, (0,1,0,0)), (1, (0,0,1,0)), (2, (0,0,0,1))]
+   return $
+      ShapeDep.backpermute LinAlg.balanceRight LinAlg.balanceLeft $
+      outerVector
+         (Expr.liftM2 $
+          MultiValue.modifyF2
+             (atom, atom, atom, (atom, atom, atom, atom))
+             ((atom, atom), (atom, (atom, atom, atom, atom))) $
+          \(n, ln, z, (xm1,x0,x1,x2)) ((minIx, maxIx), (k, (ym1,y0,y1,y2))) -> do
+             lnk <- MultiValue.add ln k
+             tooSmall <- MultiValue.cmp LLVM.CmpLT n minIx
+             tooLarge <- MultiValue.cmp LLVM.CmpGT n maxIx
+             y <-
+                MultiValue.select tooSmall y0 =<<
+                MultiValue.select tooLarge y1 =<<
+                Expr.unExp
+                   (interpolate
+                       (Expr.lift0 xm1, Expr.lift0 ym1)
+                       (Expr.lift0 x0, Expr.lift0 y0)
+                       (Expr.lift0 x1, Expr.lift0 y1)
+                       (Expr.lift0 x2, Expr.lift0 y2)
+                       (Expr.lift0 z))
+             return (lnk, y))
+         (Sym.zip4 indices limitIndices zs
+            (Sym.zip4
+               (gatherFromNodes (-1))
+               (gatherFromNodes 0)
+               (gatherFromNodes 1)
+               (gatherFromNodes 2)))
+         (zipWithScalar Expr.zip minMaxIx $ Phys.feed $ pure units)
+
+
+sampleBasisFunctions13 ::
+   (Shape.Index nodes ~ nodes,
+    Shape.C coll, Shape.C rows, Shape.C nodes,
+    Shape.C set, MultiValueMemory.C set, Storable set, Num set,
+    MultiValue.Comparison nodes, MultiValue.PseudoRing nodes,
+    MultiValue.IntegerConstant nodes, MultiValue.Integral nodes,
+    MultiValue.Select nodes,
+    MultiValueMemory.C nodes, Storable nodes, Num nodes,
+    MultiValueMemory.C rows, Storable rows,
+    MultiValueMemory.C coll, Storable coll,
+    MultiValue.Select a, MultiValue.Comparison a,
+    MultiValue.Field a, MultiValue.RationalConstant a,
+    Num a, Storable a, MultiValueMemory.C a,
+    MultiValueMemory.Struct a ~ astruct, LLVM.IsSized astruct,
+    MultiValueMemory.Struct nodes ~ nodesstruct, LLVM.IsSized nodesstruct,
+    MultiValueMemory.Struct coll ~ collstruct, LLVM.IsSized collstruct,
+    MultiValueMemory.Struct rows ~ rowsstruct, LLVM.IsSized rowsstruct) =>
+   Interpolator13 (Exp a) ->
+   SymP.Array p nodes a ->
+   Vector p coll rows a ->
+   IO (p -> IO (Sparse.RowMatrix p coll rows set nodes a))
+sampleBasisFunctions13 interpolate nodes zs = do
+   indices <- lookupInterval (outerVector (flip const) zs nodes) zs
+   return $ \p -> do
+      indexArr <- indices p
+      let minMaxIx =
+             Sym.map (\numElems -> Expr.zip 1 (numElems - 3)) $
+             ShapeDep.shape nodes
+      basis <-
+         sampleBasisFunctions13Aux interpolate minMaxIx
+            (Phys.feed $ pure indexArr) nodes zs
+      return $ Sparse.RowMatrix $
+         ShapeDep.backpermuteExtra
+            (Expr.modify2 (atom, (atom,atom)) atom $
+             \(coll, (dim,set)) numElems ->
+                (coll, (dim, Sparse.Dim set numElems)))
+            id
+            basis
+            nodes
diff --git a/src/Data/Array/Knead/Arithmetic/LinearAlgebra.hs b/src/Data/Array/Knead/Arithmetic/LinearAlgebra.hs
new file mode 100644
--- /dev/null
+++ b/src/Data/Array/Knead/Arithmetic/LinearAlgebra.hs
@@ -0,0 +1,210 @@
+{-# LANGUAGE TypeFamilies #-}
+module Data.Array.Knead.Arithmetic.LinearAlgebra where
+
+import qualified Data.Array.Knead.Parameterized.Physical as Phys
+import qualified Data.Array.Knead.Parameterized.Symbolic as SymP
+import qualified Data.Array.Knead.Simple.Symbolic as Sym
+import qualified Data.Array.Knead.Simple.ShapeDependent as ShapeDep
+import qualified Data.Array.Knead.Index.Nested.Shape as Shape
+import qualified Data.Array.Knead.Expression as Expr
+import Data.Array.Knead.Expression (Exp)
+
+import qualified LLVM.Extra.Multi.Value.Memory as MultiValueMemory
+import qualified LLVM.Extra.Multi.Value as MultiValue
+import LLVM.Extra.Multi.Value (atom)
+
+import qualified LLVM.Core as LLVM
+
+import Foreign.Storable (Storable)
+
+import Control.Arrow (arr)
+import Control.Monad.HT (chain)
+
+import qualified Data.List as List
+
+
+type Scalar p coll a = SymP.Array p coll a
+type Vector p coll dim a = SymP.Array p (coll, dim) a
+type Matrix p coll rows cols a = SymP.Array p (coll, (rows, cols)) a
+
+type PhysScalar coll a = Phys.Array coll a
+type PhysVector coll dim a = Phys.Array (coll, dim) a
+type PhysMatrix coll rows cols a = Phys.Array (coll, (rows, cols)) a
+
+type IOScalar p coll a = IO (p -> IO (PhysScalar coll a))
+type IOVector p coll dim a = IO (p -> IO (PhysVector coll dim a))
+type IOMatrix p coll rows cols a = IO (p -> IO (PhysMatrix coll rows cols a))
+
+
+dotProduct ::
+   (Shape.C coll, Shape.C dim, MultiValue.PseudoRing a) =>
+   Vector p coll dim a ->
+   Vector p coll dim a ->
+   Scalar p coll a
+dotProduct a b =
+   Sym.fold1 Expr.add $
+   Sym.zipWith Expr.mul a b
+
+outer ::
+   (Shape.C coll, Shape.C rows, Shape.C cols, MultiValue.PseudoRing a) =>
+   Vector p coll rows a ->
+   Vector p coll cols a ->
+   Matrix p coll rows cols a
+outer =
+   ShapeDep.backpermute2
+      (Expr.modify2 (atom,atom) (atom,atom) $ \(colla,rows) (collb,cols) ->
+         (Shape.intersect colla collb, (rows, cols)))
+      (Expr.mapSnd Expr.fst)
+      (Expr.mapSnd Expr.snd)
+      Expr.mul
+
+multiplyMatrixVector ::
+   (Shape.C coll, Shape.C rows, Shape.C cols, MultiValue.PseudoRing a) =>
+   Matrix p coll rows cols a ->
+   Vector p coll cols a ->
+   Vector p coll rows a
+multiplyMatrixVector m v =
+   Sym.fold1 Expr.add $
+   ShapeDep.backpermute2
+      (Expr.modify2 (atom, (atom,atom)) (atom,atom) $
+       \(collM, (rowsM, colsM)) (collV, colsV) ->
+          ((Shape.intersect collM collV, rowsM),
+            Shape.intersect colsM colsV))
+      balanceRight
+      (Expr.mapFst Expr.fst)
+      Expr.mul m v
+
+multiplyMatrixMatrix ::
+   (Shape.C coll, Shape.C rows, Shape.C glue, Shape.C cols,
+    MultiValue.PseudoRing a) =>
+   Matrix p coll rows glue a ->
+   Matrix p coll glue cols a ->
+   Matrix p coll rows cols a
+multiplyMatrixMatrix a b =
+   Sym.fold1 Expr.add $
+   ShapeDep.backpermute2
+      (Expr.modify2 (atom, (atom,atom)) (atom, (atom,atom)) $
+         \(collA, (rows, glueA)) (collB, (glueB, cols)) ->
+            ((Shape.intersect collA collB, (rows, cols)),
+             Shape.intersect glueA glueB))
+      (Expr.modify ((atom, (atom,atom)), atom) $
+         \((coll, (rows, _cols)), glue) -> (coll, (rows, glue)))
+      (Expr.modify ((atom, (atom,atom)), atom) $
+         \((coll, (_rows, cols)), glue) -> (coll, (glue, cols)))
+      Expr.mul a b
+{-
+   transpose $
+   ShapeDep.backpermute balanceRight balanceLeft $
+   multiplyMatrixVector a $
+   ShapeDep.backpermute balanceLeft balanceRight $
+   transpose b
+-}
+
+{-
+For efficient computation of x*a*x
+we must cache (a*x) or (x*a).
+
+Is there an efficient joint multiplication?
+
+xa_i_k = sum_j x_i_j * a_j_k
+xax_i_l
+   = sum_k (sum_j x_i_j * a_j_k) * x_k_l
+   = sum_j sum_k x_i_j * a_j_k * x_k_l
+-}
+matrixInverseNewtonStepNaive ::
+   (Shape.C coll, Shape.C rows, Shape.C cols, MultiValue.PseudoRing a) =>
+   Matrix p coll rows cols a ->
+   Matrix p coll cols rows a ->
+   Matrix p coll cols rows a
+matrixInverseNewtonStepNaive a x =
+   Sym.zipWith Expr.sub (Sym.map double x) $
+   multiplyMatrixMatrix x $ multiplyMatrixMatrix a x
+
+matrixInverseNewtonStep ::
+   (Shape.C coll, Shape.C rows, Shape.C cols, MultiValue.PseudoRing a,
+    MultiValueMemory.C a, Storable a,
+    MultiValueMemory.C rows, Storable rows,
+    MultiValueMemory.C cols, Storable cols,
+    MultiValueMemory.C coll, Storable coll,
+    MultiValueMemory.Struct rows ~ rowsstruct, LLVM.IsSized rowsstruct,
+    MultiValueMemory.Struct cols ~ colsstruct, LLVM.IsSized colsstruct,
+    MultiValueMemory.Struct coll ~ collstruct, LLVM.IsSized collstruct) =>
+   Matrix p coll rows cols a ->
+   Matrix p coll cols rows a ->
+   IOMatrix p coll cols rows a
+matrixInverseNewtonStep a x = do
+   ax <- Phys.render $ multiplyMatrixMatrix a x
+   result <-
+      Phys.render $
+         let xe = SymP.extendParameter fst x
+         in  Sym.zipWith Expr.sub (Sym.map double xe) $
+             multiplyMatrixMatrix xe $ Phys.feed $ arr snd
+   return $ \p ->
+      curry result p =<< ax p
+
+nest ::
+   (Integral i, Monad m) =>
+   i -> (a -> m a) -> a -> m a
+nest i f = chain $ List.genericReplicate i f
+
+matrixInverseNewton ::
+   (Shape.C coll, Shape.C rows, Shape.C cols, MultiValue.PseudoRing a,
+    MultiValueMemory.C a, Storable a,
+    MultiValueMemory.C rows, Storable rows,
+    MultiValueMemory.C cols, Storable cols,
+    MultiValueMemory.C coll, Storable coll,
+    MultiValueMemory.Struct rows ~ rowsstruct, LLVM.IsSized rowsstruct,
+    MultiValueMemory.Struct cols ~ colsstruct, LLVM.IsSized colsstruct,
+    MultiValueMemory.Struct coll ~ collstruct, LLVM.IsSized collstruct) =>
+   Int ->
+   Matrix p coll rows cols a ->
+   Matrix p coll cols rows a ->
+   IOMatrix p coll cols rows a
+matrixInverseNewton n a x = do
+   physx <- Phys.render x
+   step <-
+      matrixInverseNewtonStep
+         (SymP.extendParameter fst a)
+         (Phys.feed $ arr snd)
+   return $ \p ->
+      nest n (curry step p) =<< physx p
+
+
+double :: (MultiValue.Additive a) => Exp a -> Exp a
+double = Expr.liftM $ \x -> MultiValue.add x x
+
+transpose ::
+   (Shape.C coll, Shape.C rows, Shape.C cols) =>
+   Matrix p coll rows cols a ->
+   Matrix p coll cols rows a
+transpose =
+   ShapeDep.backpermute
+      (Expr.mapSnd Expr.swap)
+      (Expr.mapSnd Expr.swap)
+
+
+scaleRows ::
+   (Shape.C coll, Shape.C rows, Shape.C cols, MultiValue.PseudoRing a) =>
+   Vector p coll rows a ->
+   Matrix p coll rows cols a ->
+   Matrix p coll rows cols a
+scaleRows =
+   ShapeDep.backpermute2
+      (Expr.modify2 (atom, atom) (atom, (atom,atom)) $
+       \(collV, rowsV) (collM, (rowsM, colsM)) ->
+           (Shape.intersect collV collM,
+            (Shape.intersect rowsV rowsM, colsM)))
+      (Expr.mapSnd Expr.fst)
+      id
+      Expr.mul
+
+
+balanceLeft :: (Expr.Value val) => val (a,(b,c)) -> val ((a,b),c)
+balanceLeft =
+   Expr.lift1 $
+   MultiValue.modify (atom,(atom,atom)) $ \(a,(b,c)) -> ((a,b),c)
+
+balanceRight :: (Expr.Value val) => val ((a,b),c) -> val (a,(b,c))
+balanceRight =
+   Expr.lift1 $
+   MultiValue.modify ((atom,atom),atom) $ \((a,b),c) -> (a,(b,c))
diff --git a/src/Data/Array/Knead/Arithmetic/Sparse.hs b/src/Data/Array/Knead/Arithmetic/Sparse.hs
new file mode 100644
--- /dev/null
+++ b/src/Data/Array/Knead/Arithmetic/Sparse.hs
@@ -0,0 +1,273 @@
+{-# LANGUAGE TypeFamilies #-}
+module Data.Array.Knead.Arithmetic.Sparse where
+
+import qualified Data.Array.Knead.Arithmetic.LinearAlgebra as LinAlg
+import Data.Array.Knead.Arithmetic.LinearAlgebra
+          (Vector, Matrix, IOVector, IOMatrix)
+
+import qualified Data.Array.Knead.Parameterized.Physical as Phys
+import qualified Data.Array.Knead.Parameterized.Symbolic as SymP
+import qualified Data.Array.Knead.Simple.ShapeDependent as ShapeDep
+import qualified Data.Array.Knead.Simple.Symbolic as Sym
+import qualified Data.Array.Knead.Index.Nested.Shape as Shape
+import qualified Data.Array.Knead.Expression as Expr
+import Data.Array.Knead.Expression (Exp)
+
+import qualified LLVM.Extra.Multi.Value.Memory as MultiValueMemory
+import qualified LLVM.Extra.Multi.Value as MultiValue
+import LLVM.Extra.Multi.Value (atom)
+
+import qualified LLVM.Core as LLVM
+
+import Foreign.Storable (Storable)
+
+
+data Dim set dim = Dim set dim
+
+sparseSet :: (Expr.Value val) => val (Dim set dim) -> val set
+sparseSet =
+   Expr.lift1 $
+   \(MultiValue.Cons (Dim set _dim)) -> MultiValue.Cons set
+
+sparseDim :: (Expr.Value val) => val (Dim set dim) -> val dim
+sparseDim =
+   Expr.lift1 $
+   \(MultiValue.Cons (Dim _set dim)) -> MultiValue.Cons dim
+
+pairFromDim ::
+   (Expr.Value val) => val (Dim set dim) -> val (set, dim)
+pairFromDim =
+   Expr.lift1 $
+   \(MultiValue.Cons (Dim set dim)) -> MultiValue.Cons (set, dim)
+
+dimFromPair ::
+   (Expr.Value val) => val (set, dim) -> val (Dim set dim)
+dimFromPair =
+   Expr.lift1 $
+   \(MultiValue.Cons (set, dim)) -> MultiValue.Cons (Dim set dim)
+
+
+instance
+   (MultiValue.C set, MultiValue.C dim) =>
+      MultiValue.C (Dim set dim) where
+   type Repr f (Dim set dim) =
+           Dim (MultiValue.Repr f set) (MultiValue.Repr f dim)
+   cons (Dim set dim) =
+      dimFromPair $ MultiValue.cons (set, dim)
+   zero = dimFromPair MultiValue.zero
+   undef = dimFromPair MultiValue.undef
+   phis bb = fmap dimFromPair . MultiValue.phis bb . pairFromDim
+   addPhis bb a b =
+      MultiValue.addPhis bb (pairFromDim a) (pairFromDim b)
+
+instance
+   (MultiValue.Compose set, MultiValue.Compose dim) =>
+      MultiValue.Compose (Dim set dim) where
+   type Composed (Dim set dim) =
+           Dim (MultiValue.Composed set) (MultiValue.Composed dim)
+   compose (Dim set dim) =
+      dimFromPair $ MultiValue.compose (set,dim)
+
+instance
+   (Expr.Compose set, Expr.Compose dim) =>
+      Expr.Compose (Dim set dim) where
+   type Composed (Dim set dim) =
+           Dim (Expr.Composed set) (Expr.Composed dim)
+   compose (Dim set dim) =
+      dimFromPair $ Expr.compose (set,dim)
+
+instance
+   (Shape.C set, Shape.C dim) =>
+      Shape.C (Dim set dim) where
+   type Index (Dim set dim) = Shape.Index set
+   {-
+   Not really useful.
+   Only intended for the case that all dimensions match.
+   -}
+   intersectCode
+         s0@(MultiValue.Cons (Dim _set0 _dim))
+         s1@(MultiValue.Cons (Dim _set1 dim)) = do
+      MultiValue.Cons set <-
+         Shape.intersectCode (sparseSet s0) (sparseSet s1)
+      return $ MultiValue.Cons $ Dim set dim
+   sizeCode sh = Shape.sizeCode $ sparseSet sh
+   size (Dim set _dim) = Shape.size set
+   flattenIndexRec sh ix =
+      Shape.flattenIndexRec (sparseSet sh) ix
+   loop f = Shape.loop f . sparseSet
+
+
+{- |
+Sparse matrix with a definite number of non-zero entries per row.
+-}
+newtype RowMatrix p coll rows set cols a =
+   RowMatrix
+      (Matrix p coll rows (Dim set cols)
+          (Shape.Index cols, a))
+
+multiplyRowMatrixVector ::
+   (Shape.C coll, Shape.C rows, Shape.C cols,
+    Shape.C set, MultiValue.PseudoRing a) =>
+   RowMatrix p coll rows set cols a ->
+   Vector p coll cols a ->
+   Vector p coll rows a
+multiplyRowMatrixVector (RowMatrix m) v =
+   Sym.fold1 Expr.add $
+   ShapeDep.backpermute LinAlg.balanceLeft LinAlg.balanceRight $
+   Sym.zipWith Expr.mul (Sym.map Expr.snd m) $
+   Sym.gather (sparseRealIndex m) v
+
+
+{- |
+Sparse matrix with a definite number of non-zero entries per column.
+-}
+newtype ColumnMatrix p coll set rows cols a =
+   ColumnMatrix
+      (Matrix p coll (Dim set rows) cols
+          (Shape.Index rows, a))
+
+
+multiplyColumnMatrixVector ::
+   (Shape.C coll, Shape.C set, Shape.C rows, Shape.C cols,
+    MultiValue.PseudoRing a,
+    MultiValueMemory.C a, Storable a,
+    MultiValueMemory.C rows, Storable rows,
+    MultiValueMemory.C coll, Storable coll,
+    MultiValueMemory.Struct rows ~ rowsStruct, LLVM.IsSized rowsStruct,
+    MultiValueMemory.Struct coll ~ collStruct, LLVM.IsSized collStruct) =>
+   ColumnMatrix p coll set rows cols a ->
+   Vector p coll cols a ->
+   IOVector p coll rows a
+multiplyColumnMatrixVector (ColumnMatrix m) v =
+   Phys.scatter Expr.add
+      (ShapeDep.fill
+         (Expr.modify (atom,(atom,atom)) $
+            \(coll, (row,_col)) -> (coll, sparseDim row))
+         Expr.zero m)
+   $
+   Sym.mapWithIndex
+      (Expr.modify2 (atom,atom) (atom,atom) $
+         \(coll,_ix) (i,a) -> ((coll,i),a))
+   $
+   ShapeDep.backpermute2
+      (\msh _ -> msh)
+      id
+      (Expr.modify (atom,(atom,atom)) $
+         \(coll, (_row,col)) -> (coll,col))
+      -- different from mulCell for non-commutative multiplications
+      (Expr.modify2 (atom,atom) atom $
+         \(i,a) b -> (i, Expr.mul a b))
+      m v
+
+multiplyMatrixMatrix ::
+   (Shape.C coll, Shape.C set0, Shape.C set1,
+    Shape.C rows, Shape.C cols, Shape.C glue,
+    MultiValue.PseudoRing a,
+    MultiValueMemory.C a, Storable a,
+    MultiValueMemory.C cols, Storable cols,
+    MultiValueMemory.C rows, Storable rows,
+    MultiValueMemory.C coll, Storable coll,
+    MultiValueMemory.Struct cols ~ colsStruct, LLVM.IsSized colsStruct,
+    MultiValueMemory.Struct rows ~ rowsStruct, LLVM.IsSized rowsStruct,
+    MultiValueMemory.Struct coll ~ collStruct, LLVM.IsSized collStruct) =>
+   ColumnMatrix p coll set0 rows glue a ->
+   RowMatrix p coll glue set1 cols a ->
+   IOMatrix p coll rows cols a
+multiplyMatrixMatrix sx@(ColumnMatrix x) sy@(RowMatrix y) =
+   Phys.scatter Expr.add (fillMatrixMatrix Expr.zero sx sy)
+   $
+   Sym.mapWithIndex
+      (Expr.modify2 (atom,atom) (atom,atom) $
+         \(coll,_ix) (i,a) -> ((coll,i),a))
+   $
+   ShapeDep.backpermute2
+      (Expr.modify2 (atom,(atom,atom)) (atom,(atom,atom)) $
+         \(coll, (rows,glues)) (_coll, (_glues,cols)) ->
+             (coll, (rows, glues, cols)))
+      (Expr.modify (atom,(atom,atom,atom)) $
+         \(coll, (row, glue, _col)) -> (coll, (row, glue)))
+      (Expr.modify (atom,(atom,atom,atom)) $
+         \(coll, (_row, glue, col)) -> (coll, (glue, col)))
+      (Expr.modify2 (atom,atom) (atom,atom) $
+         \(i,a) (j,b) -> ((i,j), Expr.mul a b))
+      x y
+
+fillMatrixMatrix ::
+   (Shape.C coll, Shape.C rows, Shape.C cols, MultiValue.C a) =>
+   Exp a ->
+   ColumnMatrix p coll set0 rows glue a ->
+   RowMatrix p coll glue set1 cols a ->
+   SymP.Array p (coll, (rows, cols)) a
+fillMatrixMatrix a (ColumnMatrix x) (RowMatrix y) =
+   ShapeDep.backpermute2
+      (Expr.modify2 (atom,atom) atom $
+         \(coll,row) col -> (coll, (row,col)))
+      (Expr.modify (atom,(atom,atom)) $
+         \(coll, (row,_col)) -> (coll, row))
+      (Expr.modify (atom,(atom,atom)) $
+         \(_coll, (_row,col)) -> col)
+      asTypeOf
+      (ShapeDep.fill
+         (Expr.modify (atom,(atom,atom)) $
+            \(coll, (row,_col)) -> (coll, sparseDim row))
+         a x)
+      (ShapeDep.fill
+         (Expr.modify (atom,(atom,atom)) $
+            \(_coll, (_row,col)) -> sparseDim col)
+         a y)
+
+
+transposeColumnMatrix ::
+   (Shape.C coll, Shape.C set, Shape.C rows, Shape.C cols) =>
+   ColumnMatrix p coll set rows cols a ->
+   RowMatrix p coll cols set rows a
+transposeColumnMatrix (ColumnMatrix x) =
+   RowMatrix $ LinAlg.transpose x
+
+
+sparseRealIndex ::
+   (Shape.C coll, Shape.C rows, Shape.C cols) =>
+   Matrix p coll rows cols (i, a) ->
+   Matrix p coll rows cols (Shape.Index coll, i)
+sparseRealIndex =
+   Sym.mapWithIndex
+      (Expr.modify2 (atom,atom) (atom,atom) $
+          \(coll,_ix) (i,_a) -> (coll, i))
+
+scaleRowRows ::
+   (Shape.C coll, Shape.C rows, Shape.C cols, Shape.C set,
+    MultiValue.PseudoRing a) =>
+   Vector p coll rows a ->
+   RowMatrix p coll rows set cols a ->
+   RowMatrix p coll rows set cols a
+scaleRowRows v (RowMatrix m) =
+   RowMatrix $
+   ShapeDep.backpermute2
+      (flip const)
+      (Expr.mapSnd Expr.fst)
+      id
+      mulCell
+      v m
+
+scaleRowColumns ::
+   (Shape.C coll, Shape.C rows, Shape.C cols, Shape.C set,
+    MultiValue.PseudoRing a) =>
+   Vector p coll cols a ->
+   RowMatrix p coll rows set cols a ->
+   RowMatrix p coll rows set cols a
+scaleRowColumns v (RowMatrix m) =
+   RowMatrix $
+   Sym.zipWith (flip mulCell) m $
+   Sym.gather
+      (Sym.mapWithIndex
+          (Expr.modify2 (atom,atom) (atom,atom) $
+           \(coll, _rowCol) (i,_y) -> (coll,i))
+          m)
+      v
+
+mulCell ::
+   (MultiValue.PseudoRing a) =>
+   Exp a -> Exp (i, a) -> Exp (i, a)
+mulCell =
+   Expr.modify2 atom (atom, atom) $
+   \x (i,y) -> (i, Expr.mul x y)