Add function partitionNatural.

lib/numeric/cardano-numeric.cabal

@@ -34,6 +34,7 @@ library
   exposed-modules:
       Cardano.Numeric.PositiveNatural
       Cardano.Numeric.Rounding
+      Cardano.Numeric.Util
 
 test-suite unit
   default-language:
@@ -61,3 +62,4 @@ test-suite unit
       Main.hs
   other-modules:
       Cardano.Numeric.PositiveNatural.Gen
+      Cardano.Numeric.UtilSpec

lib/numeric/src/Cardano/Numeric/Util.hs

@@ -0,0 +1,121 @@
+{-# LANGUAGE TypeApplications #-}
+
+module Cardano.Numeric.Util
+    ( partitionNatural
+    , partitionNaturalMaybe
+    ) where
+
+import Prelude hiding
+    ( round )
+
+import Cardano.Numeric.PositiveNatural
+    ( PositiveNatural )
+import Cardano.Numeric.Rounding
+    ( RoundingDirection (..), round )
+import Control.Arrow
+    ( (&&&) )
+import Data.Function
+    ( (&) )
+import Data.List.NonEmpty
+    ( NonEmpty (..) )
+import Data.Maybe
+    ( fromMaybe )
+import Data.Ord
+    ( Down (..), comparing )
+import Data.Ratio
+    ( (%) )
+import Numeric.Natural
+    ( Natural )
+
+import qualified Cardano.Numeric.PositiveNatural as PN
+import qualified Data.Foldable as F
+import qualified Data.List.NonEmpty as NE
+
+--------------------------------------------------------------------------------
+-- Public functions
+--------------------------------------------------------------------------------
+
+--------------------------------------------------------------------------------
+-- Partitioning natural numbers
+--------------------------------------------------------------------------------
+
+partitionNatural :: Natural -> NonEmpty PositiveNatural -> NonEmpty Natural
+partitionNatural target weights =
+    fromMaybe impossible $
+        partitionNaturalMaybe target (PN.toNatural <$> weights)
+  where
+    -- This error condition should only happen with a non-zero sum of weights.
+    -- But this should be impossible, as we pass a non-empty list of positive
+    -- natural values.
+    impossible = error "partitionNatural"
+
+partitionNaturalMaybe :: Natural -> NonEmpty Natural -> Maybe (NonEmpty Natural)
+partitionNaturalMaybe target weights
+    | totalWeight == 0 = Nothing
+    | otherwise = Just portionsRounded
+  where
+    portionsRounded :: NonEmpty Natural
+    portionsRounded
+        -- 1. Start with the list of unrounded portions:
+        = portionsUnrounded
+        -- 2. Attach an index to each portion, so that we can remember the
+        --    original order:
+        & NE.zip indices
+        -- 3. Sort the portions into descending order of their fractional
+        --    parts, and then sort each subsequence with equal fractional
+        --    parts into descending order of their integral parts:
+        & NE.sortBy (comparing (Down . (fractionalPart &&& integralPart) . snd))
+        -- 4. Apply pre-computed roundings to each portion:
+        & NE.zipWith (fmap . round) roundings
+        -- 5. Restore the original order:
+        & NE.sortBy (comparing fst)
+        -- 6. Strip away the indices:
+        & fmap snd
+      where
+        indices :: NonEmpty Int
+        indices = 0 :| [1 ..]
+
+    portionsUnrounded :: NonEmpty Rational
+    portionsUnrounded = computeIdealPortion <$> weights
+      where
+        computeIdealPortion c
+            = fromIntegral target
+            * fromIntegral c
+            % fromIntegral totalWeight
+
+    roundings :: NonEmpty RoundingDirection
+    roundings =
+        applyN shortfall (NE.cons RoundUp) (NE.repeat RoundDown)
+      where
+        shortfall
+            = fromIntegral target
+            - fromIntegral @Integer
+                (F.sum $ round RoundDown <$> portionsUnrounded)
+
+    totalWeight :: Natural
+    totalWeight = F.sum weights
+
+--------------------------------------------------------------------------------
+-- Internal functions
+--------------------------------------------------------------------------------
+
+-- Apply the same function multiple times to a value.
+--
+applyN :: Int -> (a -> a) -> a -> a
+applyN n f = F.foldr (.) id (replicate n f)
+
+-- Extract the fractional part of a rational number.
+--
+-- Examples:
+--
+-- >>> fractionalPart (3 % 2)
+-- 1 % 2
+--
+-- >>> fractionalPart (11 % 10)
+-- 1 % 10
+--
+fractionalPart :: Rational -> Rational
+fractionalPart = snd . properFraction @_ @Integer
+
+integralPart :: Rational -> Integer
+integralPart = floor

lib/numeric/test/unit/Cardano/Numeric/UtilSpec.hs

@@ -0,0 +1,114 @@
+{-# OPTIONS_GHC -fno-warn-orphans #-}
+
+module Cardano.Numeric.UtilSpec where
+
+import Prelude
+
+import Cardano.Numeric.PositiveNatural
+    ( PositiveNatural )
+import Cardano.Numeric.PositiveNatural.Gen
+    ( genPositiveNaturalAny, shrinkPositiveNaturalAny )
+import Cardano.Numeric.Util
+    ( partitionNaturalMaybe )
+import Data.List.NonEmpty
+    ( NonEmpty (..) )
+import Data.Maybe
+    ( catMaybes )
+import Data.Ratio
+    ( (%) )
+import Numeric.Natural
+    ( Natural )
+import Test.Hspec
+    ( Spec, describe, it )
+import Test.QuickCheck
+    ( Arbitrary (..)
+    , Property
+    , arbitrarySizedNatural
+    , checkCoverage
+    , property
+    , shrink
+    , shrinkIntegral
+    , withMaxSuccess
+    , (.&&.)
+    , (===)
+    )
+
+import qualified Data.Foldable as F
+import qualified Data.List.NonEmpty as NE
+
+spec :: Spec
+spec = do
+
+    describe "partitionNaturalMaybe" $ do
+
+        it "prop_partitionNaturalMaybe_length" $
+            property prop_partitionNaturalMaybe_length
+        it "prop_partitionNaturalMaybe_sum" $
+            property prop_partitionNaturalMaybe_sum
+        it "prop_partitionNaturalMaybe_fair" $
+            withMaxSuccess 1000 $ checkCoverage prop_partitionNaturalMaybe_fair
+
+--------------------------------------------------------------------------------
+-- Partitioning natural numbers
+--------------------------------------------------------------------------------
+
+prop_partitionNaturalMaybe_length
+    :: Natural
+    -> NonEmpty Natural
+    -> Property
+prop_partitionNaturalMaybe_length target weights =
+    case partitionNaturalMaybe target weights of
+        Nothing -> F.sum weights === 0
+        Just ps -> F.length ps === F.length weights
+
+prop_partitionNaturalMaybe_sum
+    :: Natural
+    -> NonEmpty Natural
+    -> Property
+prop_partitionNaturalMaybe_sum target weights =
+    case partitionNaturalMaybe target weights of
+        Nothing -> F.sum weights === 0
+        Just ps -> F.sum ps === target
+
+-- | Check that portions are all within unity of ideal unrounded portions.
+--
+prop_partitionNaturalMaybe_fair
+    :: Natural
+    -> NonEmpty Natural
+    -> Property
+prop_partitionNaturalMaybe_fair target weights =
+    case partitionNaturalMaybe target weights of
+        Nothing -> F.sum weights === 0
+        Just ps -> prop ps
+  where
+    prop portions = (.&&.)
+        (F.all (uncurry (<=)) (NE.zip portions portionUpperBounds))
+        (F.all (uncurry (>=)) (NE.zip portions portionLowerBounds))
+      where
+        portionUpperBounds = ceiling . computeIdealPortion <$> weights
+        portionLowerBounds = floor   . computeIdealPortion <$> weights
+
+        computeIdealPortion :: Natural -> Rational
+        computeIdealPortion c
+            = fromIntegral target
+            * fromIntegral c
+            % fromIntegral totalWeight
+
+        totalWeight :: Natural
+        totalWeight = F.sum weights
+
+--------------------------------------------------------------------------------
+-- Arbitrary instances
+--------------------------------------------------------------------------------
+
+instance Arbitrary a => Arbitrary (NE.NonEmpty a) where
+    arbitrary = (:|) <$> arbitrary <*> arbitrary
+    shrink xs = catMaybes $ NE.nonEmpty <$> shrink (NE.toList xs)
+
+instance Arbitrary Natural where
+    arbitrary = arbitrarySizedNatural
+    shrink = shrinkIntegral
+
+instance Arbitrary PositiveNatural where
+    arbitrary = genPositiveNaturalAny
+    shrink = shrinkPositiveNaturalAny