QUBO.hs

{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_HADDOCK show-extensions #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Converter.QUBO
-- Copyright   :  (c) Masahiro Sakai 2018
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable
--
-----------------------------------------------------------------------------
module ToySolver.Converter.QUBO
  ( qubo2pb
  , QUBO2PBInfo (..)

  , pb2qubo
  , PB2QUBOInfo

  , pbAsQUBO
  , PBAsQUBOInfo (..)

  , qubo2ising
  , QUBO2IsingInfo (..)

  , ising2qubo
  , Ising2QUBOInfo (..)
  ) where

import Control.Monad
import Control.Monad.State
import qualified Data.Aeson as J
import Data.Aeson ((.=), (.:))
import Data.Array.Unboxed
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.List
import Data.Maybe
import qualified Data.PseudoBoolean as PBFile
import Data.Ratio
import ToySolver.Converter.Base
import ToySolver.Converter.PB (pb2qubo', PB2QUBOInfo')
import ToySolver.Internal.JSON (withTypedObject)
import qualified ToySolver.QUBO as QUBO
import qualified ToySolver.SAT.Types as SAT

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

qubo2pb :: Real a => QUBO.Problem a -> (PBFile.Formula, QUBO2PBInfo a)
qubo2pb prob =
  ( PBFile.Formula
    { PBFile.pbObjectiveFunction = Just $
        [ (c, if x1==x2 then [x1+1] else [x1+1, x2+1])
        | (x1, row) <- IntMap.toList m2
        , (x2, c) <- IntMap.toList row
        ]
    , PBFile.pbConstraints = []
    , PBFile.pbNumVars = QUBO.quboNumVars prob
    , PBFile.pbNumConstraints = 0
    }
  , QUBO2PBInfo d
  )
  where
    m1 = fmap (fmap toRational) $ QUBO.quboMatrix prob
    d = foldl' lcm 1 [denominator c | row <- IntMap.elems m1, c <- IntMap.elems row, c /= 0]
    m2 = fmap (fmap (\c -> numerator c * (d ` div` denominator c))) m1

newtype QUBO2PBInfo a = QUBO2PBInfo Integer
  deriving (Eq, Show, Read)

instance (Eq a, Show a, Read a) => Transformer (QUBO2PBInfo a) where
  type Source (QUBO2PBInfo a) = QUBO.Solution
  type Target (QUBO2PBInfo a) = SAT.Model

instance (Eq a, Show a, Read a) => ForwardTransformer (QUBO2PBInfo a) where
  transformForward (QUBO2PBInfo _) sol = ixmap (lb+1,ub+1) (subtract 1) sol
    where
      (lb,ub) = bounds sol

instance (Eq a, Show a, Read a) => BackwardTransformer (QUBO2PBInfo a) where
  transformBackward (QUBO2PBInfo _) m = ixmap (lb-1,ub-1) (+1) m
    where
      (lb,ub) = bounds m

instance (Eq a, Show a, Read a) => ObjValueTransformer (QUBO2PBInfo a) where
  type SourceObjValue (QUBO2PBInfo a) = a
  type TargetObjValue (QUBO2PBInfo a) = Integer

instance (Eq a, Show a, Read a, Real a) => ObjValueForwardTransformer (QUBO2PBInfo a) where
  transformObjValueForward (QUBO2PBInfo d) obj = round (toRational obj) * d

instance (Eq a, Show a, Read a, Num a) => ObjValueBackwardTransformer (QUBO2PBInfo a) where
  transformObjValueBackward (QUBO2PBInfo d) obj = fromInteger $ (obj + d - 1) `div` d

instance J.ToJSON (QUBO2PBInfo a) where
  toJSON (QUBO2PBInfo d) =
    J.object
    [ "type" .= ("QUBO2PBInfo" :: J.Value)
    , "objective_function_scale_factor" .= d
    ]

instance J.FromJSON (QUBO2PBInfo a) where
  parseJSON =
    withTypedObject "QUBO2PBInfo" $ \obj ->
      QUBO2PBInfo <$> obj .: "objective_function_scale_factor"

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

pbAsQUBO :: forall a. Real a => PBFile.Formula -> Maybe (QUBO.Problem a, PBAsQUBOInfo a)
pbAsQUBO formula = do
  (prob, offset) <- runStateT body 0
  return $ (prob, PBAsQUBOInfo offset)
  where
    body :: StateT Integer Maybe (QUBO.Problem a)
    body = do
      guard $ null (PBFile.pbConstraints formula)
      let f :: PBFile.WeightedTerm -> StateT Integer Maybe [(Integer, Int, Int)]
          f (c,[]) = modify (subtract c) >> return []
          f (c,[x]) = return [(c,x,x)]
          f (c,[x1,x2]) = return [(c,x1,x2)]
          f _ = mzero
      xs <- mapM f $ SAT.removeNegationFromPBSum $ fromMaybe [] $ PBFile.pbObjectiveFunction formula
      return $
        QUBO.Problem
        { QUBO.quboNumVars = PBFile.pbNumVars formula
        , QUBO.quboMatrix = mkMat $
            [ (x1', x2', fromInteger c)
            | (c,x1,x2) <- concat xs, let x1' = min x1 x2 - 1, let x2' = max x1 x2 - 1
            ]
        }

newtype PBAsQUBOInfo a = PBAsQUBOInfo Integer
  deriving (Eq, Show, Read)

instance Transformer (PBAsQUBOInfo a) where
  type Source (PBAsQUBOInfo a) = SAT.Model
  type Target (PBAsQUBOInfo a) = QUBO.Solution

instance ForwardTransformer (PBAsQUBOInfo a) where
  transformForward (PBAsQUBOInfo _offset) m = ixmap (lb-1,ub-1) (+1) m
    where
      (lb,ub) = bounds m

instance BackwardTransformer (PBAsQUBOInfo a) where
  transformBackward (PBAsQUBOInfo _offset) sol = ixmap (lb+1,ub+1) (subtract 1) sol
    where
      (lb,ub) = bounds sol

instance ObjValueTransformer (PBAsQUBOInfo a) where
  type SourceObjValue (PBAsQUBOInfo a) = Integer
  type TargetObjValue (PBAsQUBOInfo a) = a

instance Num a => ObjValueForwardTransformer (PBAsQUBOInfo a) where
  transformObjValueForward (PBAsQUBOInfo offset) obj = fromInteger (obj + offset)

instance Real a => ObjValueBackwardTransformer (PBAsQUBOInfo a) where
  transformObjValueBackward (PBAsQUBOInfo offset) obj = round (toRational obj) - offset

instance J.ToJSON (PBAsQUBOInfo a) where
  toJSON (PBAsQUBOInfo offset) =
    J.object
    [ "type" .= ("PBAsQUBOInfo" :: J.Value)
    , "objective_function_offset" .= offset
    ]

instance J.FromJSON (PBAsQUBOInfo a) where
  parseJSON =
    withTypedObject "PBAsQUBOInfo" $ \obj -> do
      offset <- obj .: "objective_function_offset"
      pure (PBAsQUBOInfo offset)

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

pb2qubo :: Real a => PBFile.Formula -> ((QUBO.Problem a, a), PB2QUBOInfo a)
pb2qubo formula = ((qubo, transformObjValueForward info2 th), ComposedTransformer info1 info2)
  where
    ((qubo', th), info1) = pb2qubo' formula
    Just (qubo, info2) = pbAsQUBO qubo'

type PB2QUBOInfo a = ComposedTransformer PB2QUBOInfo' (PBAsQUBOInfo a)

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

qubo2ising :: (Eq a, Show a, Fractional a) => QUBO.Problem a -> (QUBO.IsingModel a, QUBO2IsingInfo a)
qubo2ising QUBO.Problem{ QUBO.quboNumVars = n, QUBO.quboMatrix = qq } =
  ( QUBO.IsingModel
    { QUBO.isingNumVars = n
    , QUBO.isingInteraction = normalizeMat $ jj'
    , QUBO.isingExternalMagneticField = normalizeVec h'
    }
  , QUBO2IsingInfo (- c')
  )
  where
    {-
       Let xi = (si + 1)/2.

       Then,
         Qij xi xj
       = Qij (si + 1)/2 (sj + 1)/2
       = 1/4 Qij (si sj + si + sj + 1).

       Also,
         Qii xi xi
       = Qii (si + 1)/2 (si + 1)/2
       = 1/4 Qii (si si + 2 si + 1)
       = 1/4 Qii (2 si + 2).
    -}
    (jj', h', c') = foldl' f (IntMap.empty, IntMap.empty, 0) $ do
      (i, row)  <- IntMap.toList qq
      (j, q_ij) <- IntMap.toList row
      if i /= j then
        return
          ( IntMap.singleton (min i j) $ IntMap.singleton (max i j) (q_ij / 4)
          , IntMap.fromList [(i, q_ij / 4), (j, q_ij / 4)]
          , q_ij / 4
          )
      else
        return
          ( IntMap.empty
          , IntMap.singleton i (q_ij / 2)
          , q_ij / 2
          )

    f (jj1, h1, c1) (jj2, h2, c2) =
      ( IntMap.unionWith (IntMap.unionWith (+)) jj1 jj2
      , IntMap.unionWith (+) h1 h2
      , c1+c2
      )

newtype QUBO2IsingInfo a = QUBO2IsingInfo a
  deriving (Eq, Show, Read)

instance (Eq a, Show a) => Transformer (QUBO2IsingInfo a) where
  type Source (QUBO2IsingInfo a) = QUBO.Solution
  type Target (QUBO2IsingInfo a) = QUBO.Solution

instance (Eq a, Show a) => ForwardTransformer (QUBO2IsingInfo a) where
  transformForward _ sol = sol

instance (Eq a, Show a) => BackwardTransformer (QUBO2IsingInfo a) where
  transformBackward _ sol = sol

instance ObjValueTransformer (QUBO2IsingInfo a) where
  type SourceObjValue (QUBO2IsingInfo a) = a
  type TargetObjValue (QUBO2IsingInfo a) = a

instance (Eq a, Show a, Num a) => ObjValueForwardTransformer (QUBO2IsingInfo a) where
  transformObjValueForward (QUBO2IsingInfo offset) obj = obj + offset

instance (Eq a, Show a, Num a) => ObjValueBackwardTransformer (QUBO2IsingInfo a) where
  transformObjValueBackward (QUBO2IsingInfo offset) obj = obj - offset

instance J.ToJSON a => J.ToJSON (QUBO2IsingInfo a) where
  toJSON (QUBO2IsingInfo offset) =
    J.object
    [ "type" .= J.String "QUBO2IsingInfo"
    , "objective_function_offset" .= offset
    ]

instance J.FromJSON a => J.FromJSON (QUBO2IsingInfo a) where
  parseJSON =
    withTypedObject "QUBO2IsingInfo" $ \obj -> do
      offset <- obj .: "objective_function_offset"
      pure (QUBO2IsingInfo offset)

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

ising2qubo :: (Eq a, Num a) => QUBO.IsingModel a -> (QUBO.Problem a, Ising2QUBOInfo a)
ising2qubo QUBO.IsingModel{ QUBO.isingNumVars = n, QUBO.isingInteraction = jj, QUBO.isingExternalMagneticField = h } =
  ( QUBO.Problem
    { QUBO.quboNumVars = n
    , QUBO.quboMatrix = mkMat m
    }
  , Ising2QUBOInfo (- offset)
  )
  where
    {-
       Let si = 2 xi - 1

       Then,
         Jij si sj
       = Jij (2 xi - 1) (2 xj - 1)
       = Jij (4 xi xj - 2 xi - 2 xj + 1)
       = 4 Jij xi xj - 2 Jij xi    - 2 Jij xj    + Jij
       = 4 Jij xi xj - 2 Jij xi xi - 2 Jij xj xj + Jij

         hi si
       = hi (2 xi - 1)
       = 2 hi xi - hi
       = 2 hi xi xi - hi
    -}
    m =
      concat
      [ [(i, j, 4 * jj_ij), (i, i,  - 2 * jj_ij), (j, j,  - 2 * jj_ij)]
      | (i, row) <- IntMap.toList jj, (j, jj_ij) <- IntMap.toList row
      ] ++
      [ (i, i,  2 * hi) | (i, hi) <- IntMap.toList h ]
    offset =
        sum [jj_ij | row <- IntMap.elems jj, jj_ij <- IntMap.elems row]
      - sum (IntMap.elems h)

newtype Ising2QUBOInfo a = Ising2QUBOInfo a
  deriving (Eq, Show, Read)

instance (Eq a, Show a) => Transformer (Ising2QUBOInfo a) where
  type Source (Ising2QUBOInfo a) = QUBO.Solution
  type Target (Ising2QUBOInfo a) = QUBO.Solution

instance (Eq a, Show a) => ForwardTransformer (Ising2QUBOInfo a) where
  transformForward _ sol = sol

instance (Eq a, Show a) => BackwardTransformer (Ising2QUBOInfo a) where
  transformBackward _ sol = sol

instance (Eq a, Show a) => ObjValueTransformer (Ising2QUBOInfo a) where
  type SourceObjValue (Ising2QUBOInfo a) = a
  type TargetObjValue (Ising2QUBOInfo a) = a

instance (Eq a, Show a, Num a) => ObjValueForwardTransformer (Ising2QUBOInfo a) where
  transformObjValueForward (Ising2QUBOInfo offset) obj = obj + offset

instance (Eq a, Show a, Num a) => ObjValueBackwardTransformer (Ising2QUBOInfo a) where
  transformObjValueBackward (Ising2QUBOInfo offset) obj = obj - offset

instance J.ToJSON a => J.ToJSON (Ising2QUBOInfo a) where
  toJSON (Ising2QUBOInfo offset) =
    J.object
    [ "type" .= J.String "Ising2QUBOInfo"
    , "objective_function_offset" .= offset
    ]

instance J.FromJSON a => J.FromJSON (Ising2QUBOInfo a) where
  parseJSON =
    withTypedObject "Ising2QUBOInfo" $ \obj -> do
      offset <- obj .: "objective_function_offset"
      pure (Ising2QUBOInfo offset)

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

mkMat :: (Eq a, Num a) => [(Int,Int,a)] -> IntMap (IntMap a)
mkMat m = normalizeMat $
  IntMap.unionsWith (IntMap.unionWith (+)) $
  [IntMap.singleton i (IntMap.singleton j qij) | (i,j,qij) <- m]

normalizeMat :: (Eq a, Num a) => IntMap (IntMap a) -> IntMap (IntMap a)
normalizeMat = IntMap.mapMaybe ((\m -> if IntMap.null m then Nothing else Just m) . normalizeVec)

normalizeVec :: (Eq a, Num a) => IntMap a -> IntMap a
normalizeVec = IntMap.filter (/=0)