QUBO.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
module Test.QUBO (quboTestGroup) where

import Control.Monad
import qualified Data.Aeson as J
import Data.Array.IArray
import Data.ByteString.Builder
import qualified Data.IntMap.Strict as IntMap
import Data.Maybe
import qualified Data.PseudoBoolean as PBFile
import Data.Scientific

import Test.Tasty
import Test.Tasty.QuickCheck
import Test.Tasty.HUnit
import Test.Tasty.TH
import qualified Test.QuickCheck.Monadic as QM

import ToySolver.Converter
import qualified ToySolver.FileFormat as FF
import qualified ToySolver.QUBO as QUBO
import ToySolver.Converter.QUBO
import qualified ToySolver.SAT.Types as SAT

import Test.SAT.Utils

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

instance (Arbitrary a, Eq a, Num a) => Arbitrary (QUBO.Problem a) where
  arbitrary = do
    nv <- choose (1,10)
    m <- choose (0,nv*nv)
    jj <- liftM (f . IntMap.unionsWith (IntMap.unionWith (+))) $ replicateM m $ do
      i <- choose (0,nv-1)
      j <- choose (i,nv-1)
      jj_ij <- arbitrary
      return $ IntMap.singleton i $ IntMap.singleton j jj_ij
    return $
      QUBO.Problem
      { QUBO.quboNumVars = nv
      , QUBO.quboMatrix = jj
      }
    where
      f = IntMap.mapMaybe (g . IntMap.filter (/= 0))
      g m = if IntMap.null m then Nothing else Just m

arbitrarySolution :: Int -> Gen QUBO.Solution
arbitrarySolution nv =
  liftM (array (0,nv-1) . zip [0..]) $ replicateM nv arbitrary

instance (Arbitrary a, Eq a, Num a) => Arbitrary (QUBO.IsingModel a) where
  arbitrary = do
    nv <- choose (1,10)

    m <- choose (0,nv*nv)
    qq <- liftM (f . IntMap.unionsWith (IntMap.unionWith (+))) $ replicateM m $ do
      i <- choose (0,nv-1)
      j <- choose (i,nv-1)
      qq_ij <- arbitrary
      return $ IntMap.singleton i $ IntMap.singleton j qq_ij

    h <- liftM (\h -> IntMap.fromList [(i,hi)| (i, Just hi) <- zip [0..] h]) $ replicateM nv arbitrary

    return $
      QUBO.IsingModel
      { QUBO.isingNumVars = nv
      , QUBO.isingInteraction = qq
      , QUBO.isingExternalMagneticField = h
      }
    where
      f = IntMap.mapMaybe (g . IntMap.filter (/= 0))
      g m = if IntMap.null m then Nothing else Just m

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

prop_QUBO_ReadWrite_Invariance :: Property
prop_QUBO_ReadWrite_Invariance = forAll g $ \qubo ->
  let s = toLazyByteString (FF.render qubo)
   in counterexample (show s) $ FF.parse s === Right qubo
  where
    g = do
      qubo <- arbitrary
      return $ fmap fromFloatDigits (qubo :: QUBO.Problem Double)

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

prop_qubo2pb :: Property
prop_qubo2pb = forAll arbitrary $ \(qubo :: QUBO.Problem Integer) ->
  let (pbo,_) = qubo2pb qubo
   in Just qubo === fmap fst (pbAsQUBO pbo)

prop_qubo2pb_forward :: Property
prop_qubo2pb_forward = forAll arbitrary $ \(qubo :: QUBO.Problem Integer) ->
  let (pbo, info) = qubo2pb qubo
   in counterexample (show pbo) $
        forAll (arbitrarySolution (QUBO.quboNumVars qubo)) $ \sol ->
          Just (transformObjValueForward info (QUBO.eval sol qubo)) === SAT.evalPBFormula (transformForward info sol) pbo

prop_qubo2pb_backward :: Property
prop_qubo2pb_backward = forAll arbitrary $ \(qubo :: QUBO.Problem Integer) ->
  let (pbo, info) = qubo2pb qubo
   in counterexample (show pbo) $
        forAll (arbitraryAssignment (PBFile.pbNumVars pbo)) $ \m ->
          Just (QUBO.eval (transformBackward info m) qubo) === fmap (transformObjValueBackward info) (SAT.evalPBFormula m pbo)

prop_qubo2pb_json :: Property
prop_qubo2pb_json = forAll arbitrary $ \(qubo :: QUBO.Problem Integer) ->
  let ret@(pbo, info) = qubo2pb qubo
      json = J.encode info
   in counterexample (show ret) $ counterexample (show json) $
        J.eitherDecode json === Right info

prop_pb2qubo :: Property
prop_pb2qubo = forAll arbitraryPBFormula $ \formula ->
  let ((qubo :: QUBO.Problem Integer, th), info) = pb2qubo formula
   in counterexample (show (qubo,th,info)) $
        conjoin
        [ forAll (arbitraryAssignment (PBFile.pbNumVars formula)) $ \m ->
            case SAT.evalPBFormula m formula of
              Nothing ->
                property (QUBO.eval (transformForward info m) qubo > th)
              Just o ->
                conjoin
                [ QUBO.eval (transformForward info m) qubo === transformObjValueForward info o
                , transformObjValueBackward info (transformObjValueForward info o) === o
                , property (transformObjValueForward info o <= th)
                ]
        , forAll (arbitrarySolution (QUBO.quboNumVars qubo)) $ \sol ->
            let o = QUBO.eval sol qubo
             in if (o <= th) then
                  (SAT.evalPBFormula (transformBackward info sol) formula === Just (transformObjValueBackward info o))
                  .&&.
                  transformObjValueForward info (transformObjValueBackward info o) === o
                else
                  property True
        ]

prop_pb2qubo_json :: Property
prop_pb2qubo_json = forAll arbitraryPBFormula $ \formula ->
  let ret@(_, info) = pb2qubo formula
      json = J.encode info
   in counterexample (show ret) $ counterexample (show json) $
        J.eitherDecode json === Right info

prop_qubo2ising :: Property
prop_qubo2ising = forAll arbitrary $ \(qubo :: QUBO.Problem Rational) ->
  let (ising, _) = qubo2ising qubo
   in qubo === fst (ising2qubo ising)

prop_qubo2ising_forward :: Property
prop_qubo2ising_forward = forAll arbitrary $ \(qubo :: QUBO.Problem Rational) ->
  let (ising, info) = qubo2ising qubo
   in counterexample (show ising) $
        forAll (arbitrarySolution (QUBO.quboNumVars qubo)) $ \sol ->
          transformObjValueForward info (QUBO.eval sol qubo) === QUBO.evalIsingModel (transformForward info sol) ising

prop_qubo2ising_backward :: Property
prop_qubo2ising_backward = forAll arbitrary $ \(qubo :: QUBO.Problem Rational) ->
  let (ising, info) = qubo2ising qubo
   in counterexample (show ising) $
        forAll (arbitrarySolution (QUBO.isingNumVars ising)) $ \sol ->
          QUBO.eval (transformBackward info sol) qubo === transformObjValueBackward info (QUBO.evalIsingModel sol ising)

prop_qubo2ising_json :: Property
prop_qubo2ising_json = forAll arbitrary $ \(qubo :: QUBO.Problem Rational) ->
  let ret@(_, info) = qubo2ising qubo
      json = J.encode info
   in counterexample (show ret) $ counterexample (show json) $
        J.eitherDecode json === Right info

prop_ising2qubo_forward :: Property
prop_ising2qubo_forward = forAll arbitrary $ \(ising :: QUBO.IsingModel Integer) ->
  let (qubo, info) = ising2qubo ising
   in counterexample (show qubo) $
        forAll (arbitrarySolution (QUBO.isingNumVars ising)) $ \sol ->
          transformObjValueForward info (QUBO.evalIsingModel sol ising) === QUBO.eval (transformForward info sol) qubo

prop_ising2qubo_backward :: Property
prop_ising2qubo_backward = forAll arbitrary $ \(ising :: QUBO.IsingModel Integer) ->
  let (qubo, info) = ising2qubo ising
   in counterexample (show qubo) $
        forAll (arbitrarySolution (QUBO.quboNumVars qubo)) $ \sol ->
          QUBO.evalIsingModel (transformBackward info sol) ising === transformObjValueBackward info (QUBO.eval sol qubo)

prop_ising2qubo_json :: Property
prop_ising2qubo_json = forAll arbitrary $ \(ising :: QUBO.IsingModel Integer) ->
  let ret@(_, info) = ising2qubo ising
      json = J.encode info
   in counterexample (show ret) $ counterexample (show json) $
        J.eitherDecode json === Right info

------------------------------------------------------------------------
-- Test harness

quboTestGroup :: TestTree
quboTestGroup = $(testGroupGenerator)