Merge pull request #23 from McMasterU/add_simplify

[ solver ] simplify constraints & obj, add tests

Author: dandoh

src/HashedExpression/Codegen/CSimple.hs

@@ -383,8 +383,8 @@ instance Codegen CSimpleConfig where
       params :: [String]
       params = map fst $ paramNodesWithId expressionMap
       -- value nodes
-      vs :: [(String, Int)]
-      vs = sortOn fst $ varNodesWithId expressionMap ++ paramNodesWithId expressionMap
+      varsAndParams :: [(String, Int)]
+      varsAndParams = sortOn fst $ varNodesWithId expressionMap ++ paramNodesWithId expressionMap
       -- get shape of a variable
       variableShape :: String -> Shape
       variableShape name =
@@ -405,7 +405,7 @@ instance Codegen CSimpleConfig where
           let isOk (var, nId)
                 | Just val <- Map.lookup var valMaps = compatible (retrieveShape nId expressionMap) val
                 | otherwise = True,
-          Just (var, shape) <- find (not . isOk) vs =
+          Just (var, shape) <- find (not . isOk) varsAndParams =
           Just $ "variable " ++ var ++ "is of shape " ++ show shape ++ " but the value provided is not"
         | otherwise = Nothing
       -------------------------------------------------------------------------------
@@ -428,8 +428,8 @@ instance Codegen CSimpleConfig where
           offset = cAddress nId
           shape = retrieveShape nId expressionMap
       -------------------------------------------------------------------------------
-      writeResultCodeEach :: (String, NodeID) -> Code
-      writeResultCodeEach (name, nId)
+      writeResultCodeEach :: Variable -> Code
+      writeResultCodeEach variable
         | output == OutputHDF5 =
           scoped
             [ [i|printf("Writing #{name} to #{name}_out.h5...\\n");|],
@@ -455,6 +455,8 @@ instance Codegen CSimpleConfig where
                  )
               ++ ["fclose(file);"]
         where
+          nId = nodeId variable
+          name = varName variable
           offset = cAddress nId
           shape = retrieveShape nId expressionMap
       -------------------------------------------------------------------------------
@@ -555,12 +557,12 @@ instance Codegen CSimpleConfig where
       readValsCodes =
         ["void read_values() {"]
           ++ ["  srand(time(NULL));"] --
-          ++ scoped (concatMap readValCodeEach vs)
+          ++ scoped (concatMap readValCodeEach varsAndParams)
           ++ ["}"] --
           -------------------------------------------------------------------------------
       writeResultCodes =
         ["void write_result()"]
-          ++ scoped (concatMap writeResultCodeEach vs)
+          ++ scoped (concatMap writeResultCodeEach variables)
       -------------------------------------------------------------------------------
       evaluatingCodes =
         ["void evaluate_partial_derivatives_and_objective()"]

src/HashedExpression/Internal/Simplify.hs

@@ -7,7 +7,7 @@
 -- Portability :  unportable
 --
 -- Simplifying expressions
-module HashedExpression.Internal.Simplify (simplify) where
+module HashedExpression.Internal.Simplify (simplify, simplifyUnwrapped) where
 
 import Control.Monad.State.Strict
 import Data.Eq.HT (equating)
@@ -42,8 +42,8 @@ import HashedExpression.Prettify
 import Prelude hiding ((^))
 import qualified Prelude
 
-simplify :: forall d et. (Dimension d, ElementType et) => Expression d et -> Expression d et
-simplify = wrap . removeUnreachable . apply . unwrap
+simplifyUnwrapped :: (ExpressionMap, NodeID) -> (ExpressionMap, NodeID)
+simplifyUnwrapped = removeUnreachable . apply
   where
     apply =
       multipleTimes 1000 . toRecursiveTransformation . chainModifications $
@@ -73,6 +73,9 @@ simplify = wrap . removeUnreachable . apply . unwrap
                 ]
             )
 
+simplify :: forall d et. (Dimension d, ElementType et) => Expression d et -> Expression d et
+simplify = wrap . simplifyUnwrapped . unwrap
+
 -- | Predefined holes used for pattern matching with 'Pattern'
 [p, q, r, s, t, u, v, w, x, y, z, condition] = map PHole [1 .. 12]
 
@@ -199,17 +202,15 @@ zeroOneSumProdRules :: Modification
 zeroOneSumProdRules exp@(mp, n) =
   case retrieveOp n mp of
     Sum ns
-      -- to make sure filter (not . isZero mp) ns is not empty
-      | all (isZero mp) ns -> just $ head ns
       -- if the sumP has any zero, remove them
       -- sum(x, y, z, 0, t, 0) = sum(x, y, z, t)
-      | any (isZero mp) ns -> sum_ . map just . filter (not . isZero mp) $ ns
+      | (x : _, []) <- partition (isZero mp) ns -> just x
+      | (_, nonZeros) <- partition (isZero mp) ns -> sum_ . map just $ nonZeros
     Mul ns
-      -- to make sure filter (not . isOne mp) ns is not empty
-      | all (isOne mp) ns -> just $ head ns
       -- if the product has any one, remove them
       -- product(x, y, z, 1, t, 1) = product(x, y, z, t)
-      | any (isOne mp) ns -> product_ . map just . filter (not . isOne mp) $ ns
+      | (x : _, []) <- partition (isOne mp) ns -> just x
+      | (_, nonOnes) <- partition (isOne mp) ns -> product_ . map just $ nonOnes
       -- if any is zero, collapse to zero
       -- product(x, y, z, 0, t, u, v) = 0
       | nId : _ <- filter (isZero mp) ns -> just nId

src/HashedExpression/Internal/Utils.hs

@@ -1,7 +1,8 @@
 module HashedExpression.Internal.Utils where
 
 import Data.Array
-import Data.Complex
+import qualified Data.Complex as Complex
+import Data.Complex (Complex(..))
 import qualified Data.IntMap.Strict as IM
 import Data.List (foldl')
 import Data.List.Split (splitOn)
@@ -183,3 +184,10 @@ zipMp3 mp1 mp2 mp3 = foldl' f Map.empty $ Map.keys mp1
 
 instance PowerOp Double Int where
   (^) x y = x Prelude.^ y
+
+instance ComplexRealOp Double (Complex Double) where 
+  (+:) = (:+)
+  xRe = Complex.realPart
+  xIm = Complex.imagPart
+  conjugate = Complex.conjugate
+  

src/HashedExpression/Interp.hs

@@ -16,6 +16,10 @@ module HashedExpression.Interp
     evaluate2DComplex,
     evaluate3DReal,
     evaluate3DComplex,
+    fourierTransform1D,
+    fourierTransform2D,
+    fourierTransform3D,
+    FTMode(..),
   )
 where
 
@@ -463,8 +467,8 @@ evaluate1DComplex valMap (mp, n)
               ]
       Rotate [amount] arg ->
         rotate1D size amount (evaluate1DComplex valMap $ (mp, arg))
-      FT arg -> fourierTransform1D False size $ evaluate1DComplex valMap (mp, arg)
-      IFT arg -> fourierTransform1D True size $ evaluate1DComplex valMap (mp, arg)
+      FT arg -> fourierTransform1D FT_FORWARD size $ evaluate1DComplex valMap (mp, arg)
+      IFT arg -> fourierTransform1D FT_BACKWARD size $ evaluate1DComplex valMap (mp, arg)
       _ -> error "expression structure One C is wrong"
   | otherwise = error "one C but shape is not [size] ??"
 
@@ -600,8 +604,8 @@ evaluate2DComplex valMap (mp, n)
           (size1, size2)
           (amount1, amount2)
           (evaluate2DComplex valMap $ (mp, arg))
-      FT arg -> fourierTransform2D False (size1, size2) $ evaluate2DComplex valMap (mp, arg)
-      IFT arg -> fourierTransform2D True (size1, size2) $ evaluate2DComplex valMap (mp, arg)
+      FT arg -> fourierTransform2D FT_FORWARD (size1, size2) $ evaluate2DComplex valMap (mp, arg)
+      IFT arg -> fourierTransform2D FT_BACKWARD (size1, size2) $ evaluate2DComplex valMap (mp, arg)
       _ -> error "expression structure Two C is wrong"
   | otherwise = error "Two C but shape is not [size1, size2] ??"
 
@@ -739,8 +743,8 @@ evaluate3DComplex valMap (mp, n)
           (size1, size2, size3)
           (amount1, amount2, amount3)
           (evaluate3DComplex valMap $ (mp, arg))
-      FT arg -> fourierTransform3D False (size1, size2, size3) $ evaluate3DComplex valMap (mp, arg)
-      IFT arg -> fourierTransform3D True (size1, size2, size3) $ evaluate3DComplex valMap (mp, arg)
+      FT arg -> fourierTransform3D FT_FORWARD (size1, size2, size3) $ evaluate3DComplex valMap (mp, arg)
+      IFT arg -> fourierTransform3D FT_BACKWARD (size1, size2, size3) $ evaluate3DComplex valMap (mp, arg)
       _ -> error "expression structure Three C is wrong"
   | otherwise = error "Three C but shape is not [size1, size2, size3] ??"
 
@@ -816,6 +820,8 @@ rotate3D (size1, size2, size3) (amount1, amount2, amount3) arr =
         k <- [0 .. size3 - 1]
     ]
 
+data FTMode = FT_FORWARD | FT_BACKWARD deriving (Eq, Ord)
+
 -- | Fourier Transform in 1D.
 --  Frequency is just in one dimension.
 --  Consider a real-valued function, S(x),
@@ -824,14 +830,14 @@ rotate3D (size1, size2, size3) (amount1, amount2, amount3) arr =
 --  length of cycle is P/n, and frequency is n/P.
 --  so for input i the frequency is (2*pi*i*n)/P
 fourierTransform1D ::
-  Bool -> Int -> Array Int (Complex Double) -> Array Int (Complex Double)
-fourierTransform1D inverse size arr =
+  FTMode -> Int -> Array Int (Complex Double) -> Array Int (Complex Double)
+fourierTransform1D mode size arr =
   listArray (0, size - 1) [computeX i | i <- [0 .. size - 1]]
   where
-    s = if inverse then fromIntegral size else 1
+    s = if mode == FT_BACKWARD then fromIntegral size else 1
     computeX i = (sum $ zipWithA (*) arr (fourierBasis i)) / s
     fourierBasis i =
-      let frequency n = (2 * pi * fromIntegral (i * n) / fromIntegral size) * (if inverse then -1 else 1)
+      let frequency n = (2 * pi * fromIntegral (i * n) / fromIntegral size) * (if mode == FT_BACKWARD then -1 else 1)
        in listArray
             (0, size - 1)
             [ cos (frequency n) :+ (- sin (frequency n))
@@ -847,23 +853,23 @@ fourierTransform1D inverse size arr =
 --  so for input i the frequency is (2*pi*i*n)/P
 --  the frequency should be calculated in both dimensions for i and j
 fourierTransform2D ::
-  Bool ->
+  FTMode ->
   (Int, Int) ->
   Array (Int, Int) (Complex Double) ->
   Array (Int, Int) (Complex Double)
-fourierTransform2D inverse (size1, size2) arr =
+fourierTransform2D mode (size1, size2) arr =
   listArray
     ((0, 0), (size1 - 1, size2 - 1))
     [computeX i j | i <- [0 .. size1 - 1], j <- [0 .. size2 - 1]]
   where
-    s = if inverse then fromIntegral (size1 * size2) else 1
+    s = if mode == FT_BACKWARD then fromIntegral (size1 * size2) else 1
     computeX i j = (sum $ zipWithA (*) arr (fourierBasis i j)) / s
     fourierBasis i j =
       let frequency m n =
             ( 2 * pi * fromIntegral (i * m) / fromIntegral size1
                 + 2 * pi * fromIntegral (j * n) / fromIntegral size2
             )
-              * (if inverse then -1 else 1)
+              * (if mode == FT_BACKWARD then -1 else 1)
        in listArray
             ((0, 0), (size1 - 1, size2 - 1))
             [ cos (frequency m n) :+ (- sin (frequency m n))
@@ -880,11 +886,11 @@ fourierTransform2D inverse (size1, size2) arr =
 --   so for input i the frequency is (2*pi*i*n)/P
 --   the frequency should be calculated for all dimensions, i , j , k
 fourierTransform3D ::
-  Bool ->
+  FTMode ->
   (Int, Int, Int) ->
   Array (Int, Int, Int) (Complex Double) ->
   Array (Int, Int, Int) (Complex Double)
-fourierTransform3D inverse (size1, size2, size3) arr =
+fourierTransform3D mode (size1, size2, size3) arr =
   listArray
     ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
     [ computeX i j k
@@ -893,15 +899,15 @@ fourierTransform3D inverse (size1, size2, size3) arr =
         k <- [0 .. size3 - 1]
     ]
   where
-    s = if inverse then fromIntegral (size1 * size2) else 1
+    s = if mode == FT_BACKWARD then fromIntegral (size1 * size2) else 1
     computeX i j k = (sum $ zipWithA (*) arr (fourierBasis i j k)) / s
     fourierBasis i j k =
       let frequency m n p =
             ( 2 * pi * fromIntegral (i * m) / fromIntegral size1
                 + 2 * pi * fromIntegral (j * n) / fromIntegral size2
                 + 2 * pi * fromIntegral (k * p) / fromIntegral size3
             )
-              * (if inverse then -1 else 1)
+              * (if mode == FT_BACKWARD then -1 else 1)
        in listArray
             ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1))
             [ cos (frequency m n p) :+ (- sin (frequency m n p))

src/HashedExpression/Problem.hs

@@ -30,6 +30,7 @@ import HashedExpression.Internal.Expression
 import HashedExpression.Internal.Node
 import HashedExpression.Internal.OperationSpec
 import HashedExpression.Internal.Rewrite
+import HashedExpression.Internal.Simplify
 import HashedExpression.Internal.Utils
 import HashedExpression.Prettify (debugPrint)
 import HashedExpression.Value
@@ -195,6 +196,13 @@ getExpressionCS cs =
     Upper exp _ -> exp
     Between exp _ -> exp
 
+mapExpressionCS :: ((ExpressionMap, NodeID) -> (ExpressionMap, NodeID)) -> ConstraintStatement -> ConstraintStatement
+mapExpressionCS f cs =
+  case cs of
+    Lower exp v -> Lower (f exp) v
+    Upper exp v -> Upper (f exp) v
+    Between exp v -> Between (f exp) v
+
 -- | Extract the value from the 'ConstraintStatement'
 getValCS :: ConstraintStatement -> [Val]
 getValCS cs =
@@ -264,13 +272,6 @@ mergeToMain (mp, nID) = do
   put mergedMp
   return mergedNID
 
-mergeToMainMany :: (ExpressionMap, [NodeID]) -> ProblemConstructingM [NodeID]
-mergeToMainMany (mp, nIDs) = do
-  curMp <- get
-  let (mergedMp, resIDs) = safeMergeManyRoots curMp (mp, nIDs)
-  put mergedMp
-  return resIDs
-
 varsWithShape :: (ExpressionMap, NodeID) -> [(String, Shape)]
 varsWithShape = mapMaybe collect . IM.toList . fst
   where
@@ -298,7 +299,8 @@ constructProblemHelper obj (Constraint constraints) = do
   let processF exp = do
         let (mp, name2ID) = partialDerivativesMap exp
         let (names, beforeMergeIDs) = unzip $ Map.toList name2ID
-        Map.fromList . zip names <$> mergeToMainMany (mp, beforeMergeIDs)
+        afterMergedIDs <- mapM (mergeToMain . simplifyUnwrapped . (mp, )) beforeMergeIDs 
+        return $ Map.fromList $ zip names afterMergedIDs
   let lookupDerivative :: (String, Shape) -> Map String NodeID -> ProblemConstructingM NodeID
       lookupDerivative (name, shape) dMap = case Map.lookup name dMap of
         Just dID -> return dID
@@ -377,7 +379,10 @@ constructProblemHelper obj (Constraint constraints) = do
 
 -- | Construct a Problem from given objective function and constraints
 constructProblem :: Expression Scalar R -> Constraint -> ProblemResult
-constructProblem objectiveFunction constraint =
-  case runStateT (constructProblemHelper objectiveFunction constraint) IM.empty of
+constructProblem objectiveFunction (Constraint cs) =
+  case runStateT (constructProblemHelper simplifiedObjective simplifiedConstraint) IM.empty of
     Left reason -> ProblemInvalid reason
     Right (problem, _) -> ProblemValid problem
+  where
+    simplifiedObjective = simplify objectiveFunction
+    simplifiedConstraint = Constraint $ map (mapExpressionCS simplifyUnwrapped) cs

symphony/Symphony/Symphony.hs

@@ -131,6 +131,8 @@ generateCode outputPath (ValidSymphony objectiveExp vars consts css solver) = do
       case generateProblemCode CSimple.CSimpleConfig {output = CSimple.OutputHDF5} heProblem valMap of
         Invalid reason -> throwError $ GeneralError reason
         Success res -> do
+          liftIO $ putStrLn "Problem detail:"
+          liftIO $ print $ heProblem
           liftIO $ res outputPath
           liftIO $ putStrLn "Download solver & adapter......"
           liftIO $ downloadSolver outputPath solver

test/Commons.hs

@@ -8,7 +8,7 @@ module Commons where
 import Control.Applicative (liftA2)
 import Control.Monad (foldM, forM)
 import Data.Array
-import Data.Complex
+import Data.Complex (Complex(..))
 import Data.Function.HT (nest)
 import qualified Data.IntMap.Strict as IM
 import Data.List (intercalate, sort)
@@ -547,3 +547,9 @@ sz = IM.size . exMap
 instance (Ix i, Num a) => Num (Array i a) where
   (+) arr1 arr2 = listArray (bounds arr1) $ zipWith (+) (elems arr1) (elems arr2)
   (*) arr1 arr2 = listArray (bounds arr1) $ zipWith (*) (elems arr1) (elems arr2)
+
+instance (Ix i) => ComplexRealOp (Array i Double) (Array i (Complex Double)) where
+  (+:) arr1 arr2 = listArray (bounds arr1) $ zipWith (+:) (elems arr1) (elems arr2)
+  xRe arr = listArray (bounds arr) $ map xRe (elems arr)
+  xIm arr = listArray (bounds arr) $ map xIm (elems arr)
+  conjugate arr = listArray (bounds arr) $ map conjugate (elems arr)

test/SolverSpec.hs

@@ -48,10 +48,10 @@ runCommandIn cwd cmd = do
   (_, _, _, ph) <-
     createProcess
       (shell cmd)
-        { cwd = Just cwd
---          std_in = NoStream,
---          std_out = NoStream,
---          std_err = NoStream
+        { cwd = Just cwd,
+          std_in = NoStream,
+          std_out = NoStream,
+          std_err = NoStream
         }
   waitForProcess ph
 
@@ -127,5 +127,21 @@ spec =
           let xGot = getValue2D "x" (5, 5) res
           let xExpect = listArray ((0, 0), (4, 4)) $ replicate 25 (1 / 2.7182818285)
           xGot `shouldApprox` xExpect
+    specify "Fourier transform" $ do
+      let x = variable1D @10 "x"
+      let y = variable1D @10 "y"
+      let a = param1D @10 "a"
+      let b = param1D @10 "b"
+      let obj = norm2square (ft (x +: y) - (a +: b))
+      case constructProblem obj (Constraint []) of
+        ProblemValid p -> do
+          valA <- listArray (0, 9) <$> generate (vectorOf 10 arbitrary)
+          valB <- listArray (0, 9) <$> generate (vectorOf 10 arbitrary)
+          res <- solveProblem p (Map.fromList [("a", V1D valA), ("b", V1D valB)])
+          let xGot = getValue1D "x" 10 res
+          let yGot = getValue1D "y" 10 res
+          let xExpect = xRe (fourierTransform1D FT_BACKWARD 10 (valA +: valB))
+          let yExpect = xIm (fourierTransform1D FT_BACKWARD 10 (valA +: valB))
+          xGot `shouldApprox` xExpect
     specify "Banana function" $ do
       property prop_Rosenbrock

test/Spec.hs

@@ -35,7 +35,7 @@ main = do
     describe "HashedInterpSpec" InterpSpec.spec
     describe "StructureSpec" StructureSpec.spec
     describe "ReverseDifferentiationSpec" ReverseDifferentiationSpec.spec
---  hspecWith defaultConfig {configQuickCheckMaxSuccess = Just 10} $ do
---    describe "SolverSpec" SolverSpec.spec
+  hspecWith defaultConfig {configQuickCheckMaxSuccess = Just 10} $ do
+    describe "SolverSpec" SolverSpec.spec
   hspecWith defaultConfig {configQuickCheckMaxSuccess = Just 20} $ do
     describe "CSimpleSpec" CSimpleSpec.spec