diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,6 @@ +# Changelog for `exact-kantorovich` + + +## 0.1.0.0 - 2024-05-08 + +First release. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,26 @@ +Copyright 2024 Stéphane Laurent + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,84 @@ +# exact-kantorovich + + +[![Stack-lts](https://github.com/stla/exact-kantorovich/actions/workflows/Stack-lts.yml/badge.svg)](https://github.com/stla/exact-kantorovich/actions/workflows/Stack-lts.yml) +[![Stack-nightly](https://github.com/stla/exact-kantorovich/actions/workflows/Stack-nightly.yml/badge.svg)](https://github.com/stla/exact-kantorovich/actions/workflows/Stack-nightly.yml) + + +***Exact Kantorovich distance between finite probability measures.*** + +This small package allows to compute the exact Kantorovich distance between two +finite probability measures. This assumes that the probability masses are +rational numbers and that the distance function takes rational values only. + +___ + +Let's say you have the two probability measures with masses $(1/7, 2/7, 4/7)$ +and $(1/4, 1/4, 1/2)$, both distributed on the set $\{1, 2, 3\}$, and you want +to get the Kantorovich distance between them when this set is endowed with the +distance function $d(i, j) = |i - j|$. To get it with this package, you will +use the `kantorovich` function. It has four arguments. The first two ones are +the two random variables corresponding these probability measures, and a random +variable is defined as a map from its states set to the rational numbers; this +map assigns to each element its probability mass: + +```haskell +import Data.Map.Strict ( fromList ) +import Data.Ratio ( (%) ) +mu, nu :: RandomVariable Int +mu = fromList $ zip [1, 2, 3] [1%7, 2%7, 4%7] +nu = fromList $ zip [1, 2, 3] [1%4, 1%4, 1%2] +``` + +The third one is the distance function; it must return a *positive* rational +number: + +```haskell +dist :: (Int, Int) -> Rational +dist (i, j) = toRational $ abs (i - j) +``` + +And the fourth one is a Boolean value that you set to `True` if you want to +print to the `stdout` stream some details of the simplex algorithm performed +by the `kantorovich` function: + +```haskell +import Math.Optimization.Kantorovich +result <- kantorovich mu nu dist False +``` + +The output of the `kantorovich` function has type +`IO (Maybe (KantorovichResult a b))` where here `a = b = Int` and the type +`KantorovichResult a b` is the pair of types +`(Rational, RandomVariable (a, b))`. The first element of an object of a +`KantorovichResult` object represents the value of the Kantorovich distance +and the second element represents a solution of the +underlying linear programming problem, that is to say a joining of the two +probability measures that achieves the Kantorovich distance. + +Here is the value of the Kantorovich distance for our example: + +```haskell +import Data.Maybe ( fromJust ) +fst $ fromJust result +-- 5 % 28 +``` + +You can display the solution in the style of a matrix with the help of the +`prettyKantorovichSolution` function: + +```haskell +putStrLn $ prettyKantorovichSolution result +``` + +This prints: + +``` +┌ ┐ +│ 1 % 7 0 % 1 0 % 1 │ +│ 3 % 28 5 % 28 0 % 1 │ +│ 0 % 1 1 % 14 1 % 2 │ +└ ┘ +``` + +That's all. There is no other function exported by this package. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/exact-kantorovich.cabal b/exact-kantorovich.cabal new file mode 100644 --- /dev/null +++ b/exact-kantorovich.cabal @@ -0,0 +1,60 @@ +cabal-version: 2.2 + +name: exact-kantorovich +version: 0.1.0.0 +synopsis: Exact Kantorovich distance between finite probability measures. +description: This small package allows to compute the exact Kantorovich distance between two finite probability measures. This assumes that the probability masses are rational numbers and that the distance function takes rational values only. +homepage: https://github.com/stla/exact-kantorovich#readme +license: BSD-3-Clause +license-file: LICENSE +author: Stéphane Laurent +maintainer: laurent_step@outlook.fr +copyright: 2024 Stéphane Laurent +category: Math, Optimization +build-type: Simple +extra-source-files: README.md +extra-doc-files: CHANGELOG.md + +library + hs-source-dirs: src + exposed-modules: Math.Optimization.Kantorovich + build-depends: base >= 4.7 && < 5 + , simplex-method >= 0.2.0.0 && < 0.3 + , containers >= 0.6.5.1 && < 0.7 + , monad-logger >= 0.3.40 && < 0.4 + , matrix >= 0.3.6.0 && < 0.4 + , extra >= 1.7 && < 1.8 + default-language: Haskell2010 + ghc-options: -Wall + -Wcompat + -Widentities + -Wincomplete-record-updates + -Wincomplete-uni-patterns + -Wmissing-export-lists + -Wmissing-home-modules + -Wpartial-fields + -Wredundant-constraints + +test-suite unit-tests + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: tests/ + Build-Depends: base >= 4.7 && < 5 + , tasty >= 1.4 && < 1.5 + , tasty-hunit >= 0.10 && < 0.11 + , containers >= 0.6.5.1 && < 0.7 + , exact-kantorovich + Default-Language: Haskell2010 + ghc-options: -Wall + -Wcompat + -Widentities + -Wincomplete-record-updates + -Wincomplete-uni-patterns + -Wmissing-export-lists + -Wmissing-home-modules + -Wpartial-fields + -Wredundant-constraints + +source-repository head + type: git + location: https://github.com/stla/exact-kantorovich diff --git a/src/Math/Optimization/Kantorovich.hs b/src/Math/Optimization/Kantorovich.hs new file mode 100644 --- /dev/null +++ b/src/Math/Optimization/Kantorovich.hs @@ -0,0 +1,174 @@ +{-| +Module : Math.Optimization.Kantorovich +Description : Exact Kantorovich distance. +Copyright : (c) Stéphane Laurent, 2024 +License : BSD-3-Clause +Maintainer : laurent_step@outlook.fr + +Computes the exact Kantorovich distance between two finite probability measures. +This assumes that the probability masses are rational numbers and that the +distance function takes rational values only. +-} +module Math.Optimization.Kantorovich + ( + RandomVariable + , KantorovichValue + , KantorovichSolution + , KantorovichResult + , kantorovich + , prettyKantorovichSolution + ) where +import Prelude hiding ( EQ ) +import Control.Monad.Logger ( + runStdoutLoggingT + , filterLogger + ) +import Data.List.Extra ( + nubSort + ) +import Data.Map.Strict ( + fromList + , mapKeys + , singleton + , Map + ) +import qualified Data.Map.Strict as DM +import Data.Matrix ( + fromLists + , prettyMatrix + ) +import Data.Maybe ( + isJust + , fromJust + ) +import Data.Tuple.Extra ( + (***) + ) +import Linear.Simplex.Solver.TwoPhase ( + twoPhaseSimplex + ) +import Linear.Simplex.Types ( + Result ( .. ) + , PolyConstraint ( .. ) + , ObjectiveFunction ( .. ) + ) +import Linear.Simplex.Util ( + simplifySystem + ) + +-- | A random variable is defined as a map from a set to the set of rational +-- numbers. It maps an element to its probability mass. +type RandomVariable a = Map a Rational + +-- | Type for the value of the Kantorovich distance. +type KantorovichValue = Rational + +-- | Type for the solution of the underlying linear programming problem used +-- to compute the Kantorovich distance. The solution is a joining of the two +-- random variables. It is then a random variable on the Cartesian product of +-- the sets on which the two random variables are distributed. +type KantorovichSolution a b = RandomVariable (a, b) + +-- | Type for the result of the `kantorovich` function. +type KantorovichResult a b = (KantorovichValue, KantorovichSolution a b) + +stack :: Int -> (Int, Int) -> Int +stack ncol (i, j) = (i - 1) * ncol + j + +unstack :: Int -> Int -> (Int, Int) +unstack ncol k = quotRem (k-1) ncol + +-- | Prints the random variable representing the Kantorovich solution in the style +-- of a matrix. +prettyKantorovichSolution :: + (Ord a, Ord b) + => Maybe (KantorovichResult a b) -- ^ an output of the `kantorovich` function + -> String +prettyKantorovichSolution maybeKantorovichResult = + if isJust maybeKantorovichResult then prettyMatrix m else "" + where + kantorovichSolution = snd $ fromJust maybeKantorovichResult + pairs = DM.keys kantorovichSolution + (rows, cols) = (nubSort *** nubSort) (unzip pairs) + m = fromLists + [ + [ + kantorovichSolution DM.! (i, j) | j <- cols + ] + | i <- rows + ] + +-- | Kantorovich distance between two probability measures (random variables). +kantorovich :: + (Ord a, Ord b) + => RandomVariable a -- ^ first random variable + -> RandomVariable b -- ^ second random variable + -> ((a, b) -> Rational) -- ^ distance function taking /positive/ rational values + -> Bool -- ^ whether to print the details of the simplex algorithm to stdout + -> IO (Maybe (KantorovichResult a b)) +kantorovich rvA rvB dist info = do + maybeResult <- runStdoutLoggingT $ filterLogger (\_ _ -> info) $ + twoPhaseSimplex objFunc polyConstraints + return $ getObjectiveValueAndSolution maybeResult + where + ncol = DM.size rvB + as = DM.keys rvA + mu = DM.elems rvA + bs = DM.keys rvB + nu = DM.elems rvB + objFunc = kantorovichObjectiveFunction as bs dist + polyConstraints = kantorovichConstraints mu nu + getObjectiveValueAndSolution maybeResult = + case maybeResult of + Just (Result var varLitMap) -> + Just ( + varLitMap DM.! var + , mapKeys (\k -> (((!!) as) *** ((!!) bs)) (unstack ncol k)) + (DM.delete var varLitMap) + ) + Nothing -> Nothing + +kantorovichObjectiveFunction :: + [a] -> [b] -> ((a, b) -> Rational) -> ObjectiveFunction +kantorovichObjectiveFunction as bs dist = Min + { + objective = fromList + [ (stack n (i, j), dist (as!!(i-1), bs!!(j-1))) | i <- rows, j <- cols ] + } + where + n = length bs + rows = [ 1 .. length as ] + cols = [ 1 .. n ] + +kantorovichConstraints :: [Rational] -> [Rational] -> [PolyConstraint] +kantorovichConstraints mu nu = + simplifySystem $ + positivityConstraints ++ rowMarginsConstraints ++ colMarginsConstraints + where + m = length mu + n = length nu + rows = [ 1 .. m ] + cols = [ 1 .. n ] + positivityConstraints = -- not useless because simplex-method drops some zeros + [ + GEQ { + lhs = singleton (stack n (i, j)) 1, rhs = 0 + } + | i <- rows, j <- cols + ] + rowMarginsConstraints = + [ + EQ { + lhs = fromList [ (stack n (i, j), 1) | j <- cols ] + , rhs = mu !! (i-1) + } + | i <- rows + ] + colMarginsConstraints = + [ + EQ { + lhs = fromList [ (stack n (i, j), 1) | i <- rows ] + , rhs = nu !! (j-1) + } + | j <- cols + ] diff --git a/tests/Main.hs b/tests/Main.hs new file mode 100644 --- /dev/null +++ b/tests/Main.hs @@ -0,0 +1,46 @@ +module Main ( main ) where +import Data.Map.Strict ( + Map + , elems + , fromList + , mapWithKey + ) +import Data.Maybe ( + fromJust + ) +import Data.Ratio ( (%) ) +import Math.Optimization.Kantorovich ( + kantorovich + , KantorovichResult + ) +import Test.Tasty ( defaultMain, testGroup ) +import Test.Tasty.HUnit ( testCase, assertEqual ) + +checkSolution :: KantorovichResult Int Int -> Map (Int, Int) Rational -> Bool +checkSolution (value, solution) dists = integral == value + where + integral = sum (zipWith (*) (elems dists) (elems solution)) + +main :: IO () +main = defaultMain $ + testGroup "Tests" + [ + + testCase "README example" $ do + let + mu = fromList $ zip [1, 2, 3] [1%7, 2%7, 4%7] + nu = fromList $ zip [1, 2, 3] [1%4, 1%4, 1%2] + dist :: (Int, Int) -> Rational + dist (i, j) = toRational $ abs (i - j) + maybeResult <- kantorovich mu nu dist False + let + result@(value, solution) = fromJust maybeResult + dists = mapWithKey (\ij _ -> dist ij) solution + assertEqual "" + ( + value + , checkSolution result dists + ) + (5%28, True) + + ]