{-# LANGUAGE TypeOperators #-}
{- |
Do not import this module. It is only for demonstration purposes.
-}
module Numeric.LAPACK.Example.EconomicAllocation where

import qualified Numeric.LAPACK.Matrix.Square as Square
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix.Square ((|=|))
import Numeric.LAPACK.Matrix (ShapeInt, shapeInt, (#-#), (#*|), (#\|), (\\#))
import Numeric.LAPACK.Vector ((|-|))
import Numeric.LAPACK.Format ((##))

import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Shape ((:+:)((:+:)))


{- $setup
>>> import Numeric.LAPACK.Example.EconomicAllocation
>>> import Test.Utility (approxVector)
>>>
>>> import qualified Numeric.LAPACK.Vector as Vector
>>> import Numeric.LAPACK.Vector ((+++))
>>>
>>> import qualified Data.Array.Comfort.Storable as Array
-}


type ZeroInt2 = ShapeInt:+:ShapeInt
type Vector sh = Vector.Vector sh Double
type Matrix height width = Matrix.General height width Double
type SquareMatrix size = Square.Square size Double


balances0 :: Vector ZeroInt2
balances0 :: Vector ZeroInt2
balances0 =
   ZeroInt2 -> [Double] -> Vector ZeroInt2
forall sh a. (C sh, Storable a) => sh -> [a] -> Vector sh a
Vector.fromList (Int -> ShapeInt
shapeInt Int
2 ShapeInt -> ShapeInt -> ZeroInt2
forall sh0 sh1. sh0 -> sh1 -> sh0 :+: sh1
:+: Int -> ShapeInt
shapeInt Int
2)
      [Double
100000, Double
90000, -Double
50000, -Double
120000]

expenses0 :: Matrix ShapeInt ZeroInt2
expenses0 :: Matrix ShapeInt ZeroInt2
expenses0 =
   ShapeInt -> ZeroInt2 -> [Double] -> Matrix ShapeInt ZeroInt2
forall height width a.
(C height, C width, Storable a) =>
height -> width -> [a] -> General height width a
Matrix.fromList (Int -> ShapeInt
shapeInt Int
2) (Int -> ShapeInt
shapeInt Int
2 ShapeInt -> ShapeInt -> ZeroInt2
forall sh0 sh1. sh0 -> sh1 -> sh0 :+: sh1
:+: Int -> ShapeInt
shapeInt Int
2) ([Double] -> Matrix ShapeInt ZeroInt2)
-> [Double] -> Matrix ShapeInt ZeroInt2
forall a b. (a -> b) -> a -> b
$
   [Double
16000,  Double
4000,  Double
8000, Double
12000,
    Double
10000, Double
30000, Double
40000, Double
20000]

normalize ::
   (Eq height, Shape.C height, Shape.C width) =>
   Matrix height width -> Matrix height width
normalize :: Matrix height width -> Matrix height width
normalize Matrix height width
x = Matrix height width -> Vector height Double
forall vert horiz height width a.
(C vert, C horiz, C height, C width, Floating a) =>
Full vert horiz height width a -> Vector height a
Matrix.rowSums Matrix height width
x Vector height Double -> Matrix height width -> Matrix height width
forall vert horiz height width a.
(C vert, C horiz, C height, Eq height, C width, Floating a) =>
Vector height a
-> Full vert horiz height width a -> Full vert horiz height width a
\\# Matrix height width
x

normalizeSplit ::
   (Shape.C sh0, Shape.C sh1, Eq sh1) =>
   Matrix sh1 (sh0:+:sh1) -> (Matrix sh0 sh1, SquareMatrix sh1)
normalizeSplit :: Matrix sh1 (sh0 :+: sh1) -> (Matrix sh0 sh1, SquareMatrix sh1)
normalizeSplit Matrix sh1 (sh0 :+: sh1)
expenses =
   let a :: Full Big Big (sh0 :+: sh1) sh1 Double
a = Matrix sh1 (sh0 :+: sh1) -> Full Big Big (sh0 :+: sh1) sh1 Double
forall vert horiz height width a.
(C vert, C horiz) =>
Full vert horiz height width a -> Full horiz vert width height a
Matrix.transpose (Matrix sh1 (sh0 :+: sh1) -> Full Big Big (sh0 :+: sh1) sh1 Double)
-> Matrix sh1 (sh0 :+: sh1)
-> Full Big Big (sh0 :+: sh1) sh1 Double
forall a b. (a -> b) -> a -> b
$ Matrix sh1 (sh0 :+: sh1) -> Matrix sh1 (sh0 :+: sh1)
forall height width.
(Eq height, C height, C width) =>
Matrix height width -> Matrix height width
normalize Matrix sh1 (sh0 :+: sh1)
expenses
   in (Full Big Big (sh0 :+: sh1) sh1 Double -> Matrix sh0 sh1
forall vert height0 height1 width a.
(C vert, C height0, C height1, C width, Floating a) =>
Full vert Big (height0 :+: height1) width a
-> Full vert Big height0 width a
Matrix.takeTop Full Big Big (sh0 :+: sh1) sh1 Double
a, General sh1 sh1 Double -> SquareMatrix sh1
forall sh a. Eq sh => General sh sh a -> Square sh a
Square.fromGeneral (General sh1 sh1 Double -> SquareMatrix sh1)
-> General sh1 sh1 Double -> SquareMatrix sh1
forall a b. (a -> b) -> a -> b
$ Full Big Big (sh0 :+: sh1) sh1 Double -> General sh1 sh1 Double
forall vert height0 height1 width a.
(C vert, C height0, C height1, C width, Floating a) =>
Full vert Big (height0 :+: height1) width a
-> Full vert Big height1 width a
Matrix.takeBottom Full Big Big (sh0 :+: sh1) sh1 Double
a)


completeIdSquare ::
   (Shape.C sh0, Eq sh0, Shape.C sh1, Eq sh1) =>
   Matrix sh1 (sh0:+:sh1) -> SquareMatrix (sh0:+:sh1)
completeIdSquare :: Matrix sh1 (sh0 :+: sh1) -> SquareMatrix (sh0 :+: sh1)
completeIdSquare Matrix sh1 (sh0 :+: sh1)
x =
   let (Matrix sh0 sh1
p,SquareMatrix sh1
k) = Matrix sh1 (sh0 :+: sh1) -> (Matrix sh0 sh1, SquareMatrix sh1)
forall sh0 sh1.
(C sh0, C sh1, Eq sh1) =>
Matrix sh1 (sh0 :+: sh1) -> (Matrix sh0 sh1, SquareMatrix sh1)
normalizeSplit Matrix sh1 (sh0 :+: sh1)
x
   in (Matrix sh0 sh1 -> Square sh0 Double
forall height width a.
(C height, C width, Floating a) =>
General height width a -> Square height a
Square.identityFromHeight Matrix sh0 sh1
p, Matrix sh0 sh1
p)
      (Square sh0 Double, Matrix sh0 sh1)
-> (Full Big Big sh1 sh0 Double, SquareMatrix sh1)
-> SquareMatrix (sh0 :+: sh1)
forall vert horiz sizeA sizeB a.
(C vert, C horiz, C sizeA, Eq sizeA, C sizeB, Eq sizeB,
 Floating a) =>
(Square sizeA a, Full vert horiz sizeA sizeB a)
-> (Full horiz vert sizeB sizeA a, Square sizeB a)
-> Square (sizeA :+: sizeB) a
|=|
      (Full Big Big sh1 sh0 -> Full Big Big sh1 sh0 Double
forall shape a.
(Homogeneous shape, Floating a) =>
shape -> ArrayMatrix shape a
Matrix.zero (Full Big Big sh1 sh0 -> Full Big Big sh1 sh0 Double)
-> Full Big Big sh1 sh0 -> Full Big Big sh1 sh0 Double
forall a b. (a -> b) -> a -> b
$ Full Big Big sh1 sh0 Double -> Full Big Big sh1 sh0
forall sh a. ArrayMatrix sh a -> sh
ArrMatrix.shape (Full Big Big sh1 sh0 Double -> Full Big Big sh1 sh0)
-> Full Big Big sh1 sh0 Double -> Full Big Big sh1 sh0
forall a b. (a -> b) -> a -> b
$ Matrix sh0 sh1 -> Full Big Big sh1 sh0 Double
forall vert horiz height width a.
(C vert, C horiz) =>
Full vert horiz height width a -> Full horiz vert width height a
Matrix.transpose Matrix sh0 sh1
p, SquareMatrix sh1
k)

iterated ::
   (Shape.C sh0, Eq sh0, Shape.C sh1, Eq sh1) =>
   Matrix sh1 (sh0:+:sh1) -> Vector (sh0:+:sh1) -> Vector (sh0:+:sh1)
iterated :: Matrix sh1 (sh0 :+: sh1)
-> Vector (sh0 :+: sh1) -> Vector (sh0 :+: sh1)
iterated Matrix sh1 (sh0 :+: sh1)
expenses =
   -- 'Stream.head' would be total
   [Vector (sh0 :+: sh1)] -> Vector (sh0 :+: sh1)
forall a. [a] -> a
head ([Vector (sh0 :+: sh1)] -> Vector (sh0 :+: sh1))
-> (Vector (sh0 :+: sh1) -> [Vector (sh0 :+: sh1)])
-> Vector (sh0 :+: sh1)
-> Vector (sh0 :+: sh1)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector (sh0 :+: sh1) -> Bool)
-> [Vector (sh0 :+: sh1)] -> [Vector (sh0 :+: sh1)]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile ((Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>=Double
1e-5) (Double -> Bool)
-> (Vector (sh0 :+: sh1) -> Double) -> Vector (sh0 :+: sh1) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector sh1 Double -> Double
forall sh a. (C sh, Floating a) => Vector sh a -> RealOf a
Vector.normInf (Vector sh1 Double -> Double)
-> (Vector (sh0 :+: sh1) -> Vector sh1 Double)
-> Vector (sh0 :+: sh1)
-> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (sh0 :+: sh1) -> Vector sh1 Double
forall sh0 sh1 a.
(C sh0, C sh1, Storable a) =>
Array (sh0 :+: sh1) a -> Array sh1 a
Vector.takeRight) ([Vector (sh0 :+: sh1)] -> [Vector (sh0 :+: sh1)])
-> (Vector (sh0 :+: sh1) -> [Vector (sh0 :+: sh1)])
-> Vector (sh0 :+: sh1)
-> [Vector (sh0 :+: sh1)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   (Vector (sh0 :+: sh1) -> Vector (sh0 :+: sh1))
-> Vector (sh0 :+: sh1) -> [Vector (sh0 :+: sh1)]
forall a. (a -> a) -> a -> [a]
iterate (Matrix sh1 (sh0 :+: sh1) -> SquareMatrix (sh0 :+: sh1)
forall sh0 sh1.
(C sh0, Eq sh0, C sh1, Eq sh1) =>
Matrix sh1 (sh0 :+: sh1) -> SquareMatrix (sh0 :+: sh1)
completeIdSquare Matrix sh1 (sh0 :+: sh1)
expenses SquareMatrix (sh0 :+: sh1)
-> Vector (sh0 :+: sh1)
-> Vector (HeightOf (Array (Square (sh0 :+: sh1)))) Double
forall typ width a.
(MultiplyVector typ, WidthOf typ ~ width, Eq width, Floating a) =>
Matrix typ a -> Vector width a -> Vector (HeightOf typ) a
#*|)


compensated ::
   (Shape.C sh0, Eq sh0, Shape.C sh1, Eq sh1) =>
   Matrix sh1 (sh0:+:sh1) -> Vector (sh0:+:sh1) -> Vector sh0
compensated :: Matrix sh1 (sh0 :+: sh1) -> Vector (sh0 :+: sh1) -> Vector sh0
compensated Matrix sh1 (sh0 :+: sh1)
expenses Vector (sh0 :+: sh1)
balances =
   let (Matrix sh0 sh1
p,SquareMatrix sh1
k) = Matrix sh1 (sh0 :+: sh1) -> (Matrix sh0 sh1, SquareMatrix sh1)
forall sh0 sh1.
(C sh0, C sh1, Eq sh1) =>
Matrix sh1 (sh0 :+: sh1) -> (Matrix sh0 sh1, SquareMatrix sh1)
normalizeSplit Matrix sh1 (sh0 :+: sh1)
expenses
       x :: Vector sh0
x = Vector (sh0 :+: sh1) -> Vector sh0
forall sh0 sh1 a.
(C sh0, C sh1, Storable a) =>
Array (sh0 :+: sh1) a -> Array sh0 a
Vector.takeLeft Vector (sh0 :+: sh1)
balances
       y :: Array sh1 Double
y = Vector (sh0 :+: sh1) -> Array sh1 Double
forall sh0 sh1 a.
(C sh0, C sh1, Storable a) =>
Array (sh0 :+: sh1) a -> Array sh1 a
Vector.takeRight Vector (sh0 :+: sh1)
balances
   in Vector sh0
x Vector sh0 -> Vector sh0 -> Vector sh0
forall sh a.
(C sh, Eq sh, Floating a) =>
Vector sh a -> Vector sh a -> Vector sh a
|-| Matrix sh0 sh1
p Matrix sh0 sh1
-> Array sh1 Double
-> Vector (HeightOf (Array (General sh0 sh1))) Double
forall typ width a.
(MultiplyVector typ, WidthOf typ ~ width, Eq width, Floating a) =>
Matrix typ a -> Vector width a -> Vector (HeightOf typ) a
#*| (SquareMatrix sh1
k SquareMatrix sh1 -> SquareMatrix sh1 -> SquareMatrix sh1
forall shape a.
(Additive shape, Floating a) =>
ArrayMatrix shape a -> ArrayMatrix shape a -> ArrayMatrix shape a
#-# SquareMatrix sh1 -> SquareMatrix sh1
forall sh a. (C sh, Floating a) => Square sh a -> Square sh a
Square.identityFrom SquareMatrix sh1
k) SquareMatrix sh1 -> Array sh1 Double -> Array sh1 Double
forall typ height a.
(Solve typ, HeightOf typ ~ height, Eq height, Floating a) =>
Matrix typ a -> Vector height a -> Vector height a
#\| Array sh1 Double
y


{- |
prop> let result = iterated expenses0 balances0 in approxVector result $ compensated expenses0 balances0 +++ Vector.zero (Array.shape $ Vector.takeRight result)
-}
main :: IO ()
main :: IO ()
main = do
   Matrix ShapeInt ZeroInt2 -> Vector ZeroInt2 -> Vector ZeroInt2
forall sh0 sh1.
(C sh0, Eq sh0, C sh1, Eq sh1) =>
Matrix sh1 (sh0 :+: sh1)
-> Vector (sh0 :+: sh1) -> Vector (sh0 :+: sh1)
iterated Matrix ShapeInt ZeroInt2
expenses0 Vector ZeroInt2
balances0 Vector ZeroInt2 -> String -> IO ()
forall a. Format a => a -> String -> IO ()
## String
"%10.2f"
   Matrix ShapeInt ZeroInt2 -> Vector ZeroInt2 -> Vector ShapeInt
forall sh0 sh1.
(C sh0, Eq sh0, C sh1, Eq sh1) =>
Matrix sh1 (sh0 :+: sh1) -> Vector (sh0 :+: sh1) -> Vector sh0
compensated Matrix ShapeInt ZeroInt2
expenses0 Vector ZeroInt2
balances0 Vector ShapeInt -> String -> IO ()
forall a. Format a => a -> String -> IO ()
## String
"%10.2f"