toysolver-0.10.0: src/ToySolver/Arith/Simplex/Textbook.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_HADDOCK show-extensions #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Arith.Simplex.Textbook
-- Copyright   :  (c) Masahiro Sakai 2011
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable
--
-- Naïve implementation of Simplex method
--
-- Reference:
--
-- * <http://www.math.cuhk.edu.hk/~wei/lpch3.pdf>
--
-----------------------------------------------------------------------------
module ToySolver.Arith.Simplex.Textbook
  (
  -- * The @Tableau@ type
    Tableau
  , RowIndex
  , ColIndex
  , Row
  , emptyTableau
  , objRowIndex
  , pivot
  , lookupRow
  , addRow
  , setObjFun

  -- * Optimization direction
  , module Data.OptDir

  -- * Reading status
  , currentValue
  , currentObjValue
  , isFeasible
  , isOptimal

  -- * High-level solving functions
  , simplex
  , dualSimplex
  , phaseI
  , primalDualSimplex

  -- * For debugging
  , isValidTableau
  , toCSV
  ) where

import Data.Ord
import Data.List (foldl', intersperse, maximumBy, minimumBy)
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Data.OptDir
import Data.VectorSpace
import Control.Exception

import qualified ToySolver.Data.LA as LA
import ToySolver.Data.IntVar

-- ---------------------------------------------------------------------------

type Tableau r = VarMap (Row r)
{-
tbl ! v == (m, val)
==>
var v .+. m .==. constant val
-}

-- | Basic variables
type RowIndex = Int

-- | Non-basic variables
type ColIndex = Int

type Row r = (VarMap r, r)

data PivotResult r = PivotUnbounded | PivotFinished | PivotSuccess (Tableau r)
  deriving (Show, Eq, Ord)

emptyTableau :: Tableau r
emptyTableau = IM.empty

objRowIndex :: RowIndex
objRowIndex = -1

pivot :: (Fractional r, Eq r) => RowIndex -> ColIndex -> Tableau r -> Tableau r
{-# INLINE pivot #-}
{-# SPECIALIZE pivot :: RowIndex -> ColIndex -> Tableau Rational -> Tableau Rational #-}
{-# SPECIALIZE pivot :: RowIndex -> ColIndex -> Tableau Double -> Tableau Double #-}
pivot r s tbl =
    assert (isValidTableau tbl) $  -- precondition
    assert (isValidTableau tbl') $ -- postcondition
      tbl'
  where
    tbl' = IM.insert s row_s $ IM.map f $ IM.delete r $ tbl
    f orig@(row_i, row_i_val) =
      case IM.lookup s row_i of
        Nothing -> orig
        Just c ->
          ( IM.filter (0/=) $ IM.unionWith (+) (IM.delete s row_i) (IM.map (negate c *) row_r')
          , row_i_val - c*row_r_val'
          )
    (row_r, row_r_val) = lookupRow r tbl
    a_rs = row_r IM.! s
    row_r' = IM.map (/ a_rs) $ IM.insert r 1 $ IM.delete s row_r
    row_r_val' = row_r_val / a_rs
    row_s = (row_r', row_r_val')

-- | Lookup a row by basic variable
lookupRow :: RowIndex -> Tableau r -> Row r
lookupRow r m = m IM.! r

-- 行の基底変数の列が0になるように変形
normalizeRow :: (Num r, Eq r) => Tableau r -> Row r -> Row r
normalizeRow a (row0,val0) = obj'
  where
    obj' = g $ foldl' f (IM.empty, val0) $
           [ case IM.lookup j a of
               Nothing -> (IM.singleton j x, 0)
               Just (row,val) -> (IM.map ((-x)*) (IM.delete j row), -x*val)
           | (j,x) <- IM.toList row0 ]
    f (m1,v1) (m2,v2) = (IM.unionWith (+) m1 m2, v1+v2)
    g (m,v) = (IM.filter (0/=) m, v)

setRow :: (Num r, Eq r) => Tableau r -> RowIndex -> Row r -> Tableau r
setRow tbl i row = assert (isValidTableau tbl) $ assert (isValidTableau tbl') $ tbl'
  where
    tbl' = IM.insert i (normalizeRow tbl row) tbl

addRow :: (Num r, Eq r) => Tableau r -> RowIndex -> Row r -> Tableau r
addRow tbl i row = assert (i `IM.notMember` tbl) $ setRow tbl i row

setObjFun :: (Num r, Eq r) => Tableau r -> LA.Expr r -> Tableau r
setObjFun tbl e = addRow tbl objRowIndex row
  where
    row =
      case LA.extract LA.unitVar e of
        (c, e') -> (LA.coeffMap (negateV e'), c)

copyObjRow :: (Num r, Eq r) => Tableau r -> Tableau r -> Tableau r
copyObjRow from to =
  case IM.lookup objRowIndex from of
    Nothing -> IM.delete objRowIndex to
    Just row -> addRow to objRowIndex row

currentValue :: Num r => Tableau r -> Var -> r
currentValue tbl v =
  case IM.lookup v tbl of
    Nothing -> 0
    Just (_, val) -> val

currentObjValue :: Tableau r -> r
currentObjValue = snd . lookupRow objRowIndex

isValidTableau :: Tableau r -> Bool
isValidTableau tbl =
  and [v `IM.notMember` m | (v, (m,_)) <- IM.toList tbl, v /= objRowIndex] &&
  and [IS.null (IM.keysSet m `IS.intersection` vs) | (m,_) <- IM.elems tbl']
  where
    tbl' = IM.delete objRowIndex tbl
    vs = IM.keysSet tbl'

isFeasible :: Real r => Tableau r -> Bool
isFeasible tbl =
  and [val >= 0 | (v, (_,val)) <- IM.toList tbl, v /= objRowIndex]

isOptimal :: Real r => OptDir -> Tableau r -> Bool
isOptimal optdir tbl =
  and [not (cmp cj) | cj <- IM.elems (fst (lookupRow objRowIndex tbl))]
  where
    cmp = case optdir of
            OptMin -> (0<)
            OptMax -> (0>)

isImproving :: Real r => OptDir -> Tableau r -> Tableau r -> Bool
isImproving OptMin from to = currentObjValue to <= currentObjValue from
isImproving OptMax from to = currentObjValue to >= currentObjValue from

-- ---------------------------------------------------------------------------
-- primal simplex

simplex :: (Real r, Fractional r) => OptDir -> Tableau r -> (Bool, Tableau r)
{-# SPECIALIZE simplex :: OptDir -> Tableau Rational -> (Bool, Tableau Rational) #-}
{-# SPECIALIZE simplex :: OptDir -> Tableau Double -> (Bool, Tableau Double) #-}
simplex optdir = go
  where
    go tbl = assert (isFeasible tbl) $
      case primalPivot optdir tbl of
        PivotFinished  -> assert (isOptimal optdir tbl) (True, tbl)
        PivotUnbounded -> (False, tbl)
        PivotSuccess tbl' -> assert (isImproving optdir tbl tbl') $ go tbl'

primalPivot :: (Real r, Fractional r) => OptDir -> Tableau r -> PivotResult r
{-# INLINE primalPivot #-}
primalPivot optdir tbl
  | null cs   = PivotFinished
  | null rs   = PivotUnbounded
  | otherwise = PivotSuccess (pivot r s tbl)
  where
    cmp = case optdir of
            OptMin -> (0<)
            OptMax -> (0>)
    cs = [(j,cj) | (j,cj) <- IM.toList (fst (lookupRow objRowIndex tbl)), cmp cj]
    -- smallest subscript rule
    s = fst $ head cs
    -- classical rule
    --s = fst $ (if optdir==OptMin then maximumBy else minimumBy) (compare `on` snd) cs
    rs = [ (i, y_i0 / y_is)
         | (i, (row_i, y_i0)) <- IM.toList tbl, i /= objRowIndex
         , let y_is = IM.findWithDefault 0 s row_i, y_is > 0
         ]
    r = fst $ minimumBy (comparing snd) rs

-- ---------------------------------------------------------------------------
-- dual simplex

dualSimplex :: (Real r, Fractional r) => OptDir -> Tableau r -> (Bool, Tableau r)
{-# SPECIALIZE dualSimplex :: OptDir -> Tableau Rational -> (Bool, Tableau Rational) #-}
{-# SPECIALIZE dualSimplex :: OptDir -> Tableau Double -> (Bool, Tableau Double) #-}
dualSimplex optdir = go
  where
    go tbl = assert (isOptimal optdir tbl) $
      case dualPivot optdir tbl of
        PivotFinished  -> assert (isFeasible tbl) $ (True, tbl)
        PivotUnbounded -> (False, tbl)
        PivotSuccess tbl' -> assert (isImproving optdir tbl' tbl) $ go tbl'

dualPivot :: (Real r, Fractional r) => OptDir -> Tableau r -> PivotResult r
{-# INLINE dualPivot #-}
dualPivot optdir tbl
  | null rs   = PivotFinished
  | null cs   = PivotUnbounded
  | otherwise = PivotSuccess (pivot r s tbl)
  where
    rs = [(i, row_i) | (i, (row_i, y_i0)) <- IM.toList tbl, i /= objRowIndex, 0 > y_i0]
    (r, row_r) = head rs
    cs = [ (j, if optdir==OptMin then y_0j / y_rj else - y_0j / y_rj)
         | (j, y_rj) <- IM.toList row_r
         , y_rj < 0
         , let y_0j = IM.findWithDefault 0 j obj
         ]
    s = fst $ minimumBy (comparing snd) cs
    (obj,_) = lookupRow objRowIndex tbl

-- ---------------------------------------------------------------------------
-- phase I of the two-phased method

phaseI :: (Real r, Fractional r) => Tableau r -> VarSet -> (Bool, Tableau r)
{-# SPECIALIZE phaseI :: Tableau Rational -> VarSet -> (Bool, Tableau Rational) #-}
{-# SPECIALIZE phaseI :: Tableau Double -> VarSet -> (Bool, Tableau Double) #-}
phaseI tbl avs
  | currentObjValue tbl1' /= 0 = (False, tbl1')
  | otherwise = (True, copyObjRow tbl $ removeArtificialVariables avs $ tbl1')
  where
    optdir = OptMax
    tbl1 = setObjFun tbl $ negateV $ sumV [LA.var v | v <- IS.toList avs]
    tbl1' = go tbl1
    go tbl2
      | currentObjValue tbl2 == 0 = tbl2
      | otherwise =
        case primalPivot optdir tbl2 of
          PivotSuccess tbl2' -> assert (isImproving optdir tbl2 tbl2') $ go tbl2'
          PivotFinished -> assert (isOptimal optdir tbl2) tbl2
          PivotUnbounded -> error "phaseI: should not happen"

-- post-processing of phaseI
-- pivotを使ってartificial variablesを基底から除いて、削除
removeArtificialVariables :: (Real r, Fractional r) => VarSet -> Tableau r -> Tableau r
removeArtificialVariables avs tbl0 = tbl2
  where
    tbl1 = foldl' f (IM.delete objRowIndex tbl0) (IS.toList avs)
    tbl2 = IM.map (\(row,val) ->  (IM.filterWithKey (\j _ -> j `IS.notMember` avs) row, val)) tbl1
    f tbl i =
      case IM.lookup i tbl of
        Nothing -> tbl
        Just row ->
          case [j | (j,c) <- IM.toList (fst row), c /= 0, j `IS.notMember` avs] of
            [] -> IM.delete i tbl
            j:_ -> pivot i j tbl

-- ---------------------------------------------------------------------------
-- primal-dual simplex

data PDResult = PDUnsat | PDOptimal | PDUnbounded

primalDualSimplex :: (Real r, Fractional r) => OptDir -> Tableau r -> (Bool, Tableau r)
{-# SPECIALIZE primalDualSimplex :: OptDir -> Tableau Rational -> (Bool, Tableau Rational) #-}
{-# SPECIALIZE primalDualSimplex :: OptDir -> Tableau Double -> (Bool, Tableau Double) #-}
primalDualSimplex optdir = go
  where
    go tbl =
      case pdPivot optdir tbl of
        Left PDOptimal   -> assert (isFeasible tbl) $ assert (isOptimal optdir tbl) $ (True, tbl)
        Left PDUnsat     -> assert (not (isFeasible tbl)) $ (False, tbl)
        Left PDUnbounded -> assert (not (isOptimal optdir tbl)) $ (False, tbl)
        Right tbl' -> go tbl'

pdPivot :: (Real r, Fractional r) => OptDir -> Tableau r -> Either PDResult (Tableau r)
pdPivot optdir tbl
  | null ps && null qs = Left PDOptimal -- optimal
  | otherwise =
      case ret of
        Left p -> -- primal update
          let rs = [ (i, (bi - t) / y_ip)
                   | (i, (row_i, bi)) <- IM.toList tbl, i /= objRowIndex
                   , let y_ip = IM.findWithDefault 0 p row_i, y_ip > 0
                   ]
              q = fst $ minimumBy (comparing snd) rs
          in if null rs
             then Left PDUnsat
             else Right (pivot q p tbl)
        Right q -> -- dual update
          let (row_q, _bq) = (tbl IM.! q)
              cs = [ (j, (cj'-t) / (-y_qj))
                   | (j, y_qj) <- IM.toList row_q
                   , y_qj < 0
                   , let cj  = IM.findWithDefault 0 j obj
                   , let cj' = if optdir==OptMax then cj else -cj
                   ]
              p = fst $ maximumBy (comparing snd) cs
              (obj,_) = lookupRow objRowIndex tbl
          in if null cs
             then Left PDUnbounded -- dual infeasible
             else Right (pivot q p tbl)
  where
    qs = [ (Right i, bi) | (i, (_row_i, bi)) <- IM.toList tbl, i /= objRowIndex, 0 > bi ]
    ps = [ (Left j, cj')
         | (j,cj) <- IM.toList (fst (lookupRow objRowIndex tbl))
         , let cj' = if optdir==OptMax then cj else -cj
         , 0 > cj' ]
    (ret, t) = minimumBy (comparing snd) (qs ++ ps)

-- ---------------------------------------------------------------------------

toCSV :: (Num r) => (r -> String) -> Tableau r -> String
toCSV showCell tbl = unlines . map (concat . intersperse ",") $ header : body
  where
    header :: [String]
    header = "" : map colName cols ++ [""]

    body :: [[String]]
    body = [showRow i (lookupRow i tbl) | i <- rows]

    rows :: [RowIndex]
    rows = IM.keys (IM.delete objRowIndex tbl) ++ [objRowIndex]

    cols :: [ColIndex]
    cols = [0..colMax]
      where
        colMax = maximum (-1 : [c | (row, _) <- IM.elems tbl, c <- IM.keys row])

    rowName :: RowIndex -> String
    rowName i = if i==objRowIndex then "obj" else "x" ++ show i

    colName :: ColIndex -> String
    colName j = "x" ++ show j

    showRow i (row, row_val) = rowName i : [showCell (IM.findWithDefault 0 j row') | j <- cols] ++ [showCell row_val]
      where row' = IM.insert i 1 row

-- ---------------------------------------------------------------------------