From 671068f7f6bb1ca4dffca0482c6f40fc0a830c92 Mon Sep 17 00:00:00 2001 From: Daniel Thwaites Date: Mon, 2 May 2022 01:12:27 +0100 Subject: [PATCH] Add internal docs for palette generator :bulb: So that I don't forget how it works :) --- flake.nix | 15 ++++- palette-generator/Ai/Evolutionary.hs | 96 +++++++++++++++++++++++----- palette-generator/Data/Colour.hs | 4 ++ palette-generator/Stylix/Main.hs | 10 ++- palette-generator/Stylix/Output.hs | 8 +++ palette-generator/Stylix/Palette.hs | 70 +++++++++++++++----- 6 files changed, 169 insertions(+), 34 deletions(-) diff --git a/flake.nix b/flake.nix index 040ffd01..58a695db 100644 --- a/flake.nix +++ b/flake.nix @@ -30,13 +30,26 @@ installPhase = "install -D Stylix/Main $out/bin/palette-generator"; }; + # Internal documentation + palette-generator-haddock = pkgs.stdenvNoCC.mkDerivation { + name = "palette-generator-haddock"; + src = ./palette-generator; + buildInputs = [ ghc ]; + buildPhase = + "haddock $src/**/*.hs --html --ignore-all-exports --odir $out"; + dontInstall = true; + dontFixup = true; + }; + palette-generator-app = utils.lib.mkApp { drv = palette-generator; name = "palette-generator"; }; in { - packages.palette-generator = palette-generator; + packages = { + inherit palette-generator palette-generator-haddock; + }; apps.palette-generator = palette-generator-app; })) // { nixosModules.stylix = { pkgs, ... }@args: { diff --git a/palette-generator/Ai/Evolutionary.hs b/palette-generator/Ai/Evolutionary.hs index f7bcd715..7cf02144 100644 --- a/palette-generator/Ai/Evolutionary.hs +++ b/palette-generator/Ai/Evolutionary.hs @@ -8,20 +8,35 @@ import Data.List ( mapAccumR, sortBy ) import Data.Ord ( Down(Down), comparing ) import System.Random ( RandomGen, randomR ) +{- | +Find every possible combination of two values, with the first value +coming from one list and the second value coming from a different list. +-} cartesianProduct :: [a] -> [b] -> [(a, b)] cartesianProduct = liftA2 (,) +{- | +Find every possible combination of two values, with both values coming +from the same list. Values are allowed to be paired with themself. +-} cartesianSquare :: [a] -> [(a, a)] cartesianSquare as = as `cartesianProduct` as +-- | Chain a function a set number of times. repeatCall :: Int -> (a -> a) -> a -> a repeatCall n f = (!! n) . iterate f +-- | Pick a random element from a list using a random generator. randomFromList :: (RandomGen r) => r -> [a] -> (a, r) randomFromList generator list = let (index, generator') = randomR (0, length list - 1) generator in (list !! index, generator') +{- | +Map over a list, passing a random generator into the mapped +function each time it is called. A random generator is returned +along with the new list. +-} mapWithGen :: (r -> a -> (r, b)) -> (r, [a]) -> (r, [b]) mapWithGen = uncurry . mapAccumR @@ -32,28 +47,63 @@ unfoldWithGen f size generator = (generator'', a) = f generator' in (generator'', a:as) +{- | +A genotype is a value which is generated by the genetic algorithm. + +The environment is used to specify the problem for which +we are trying to find the optimal genotype. +-} class Species environment genotype where + -- | Generate a new genotype at random. generate :: (RandomGen r) => environment -> r -> (r, genotype) + + -- | Randomly combine two genotypes. crossover :: (RandomGen r) => environment -> r -> genotype -> genotype -> (r, genotype) + + -- | Randomly mutate a genotype using the given environment. mutate :: (RandomGen r) => environment -> r -> genotype -> (r, genotype) + + -- | Score a genotype. Higher numbers are better. fitness :: environment -> genotype -> Double -data EvolutionConfig = EvolutionConfig { populationSize :: Int - , survivors :: Int - , mutationProbability :: Double - , generations :: Int - } +-- | Parameters for the genetic algorithm. +data EvolutionConfig = EvolutionConfig + { -- | The number of genotypes processed on each pass. + populationSize :: Int, + -- | How many genotypes make it through to the next pass. + survivors :: Int, + -- | The chance of a genotype being randomly changed + -- before crossover. Between 0 and 1. + mutationProbability :: Double, + -- | Number of passes of the algorithm. + generations :: Int + } -randomMutation :: - (RandomGen r, Species e g) => - e -> EvolutionConfig -> r -> g -> (r, g) +{- | +Randomly mutate the given genotype, if the mutation probability +from the 'EvolutionConfig' says yes. +-} +randomMutation :: (RandomGen r, Species e g) + => e -- ^ Environment + -> EvolutionConfig + -> r -- ^ Random generator + -> g -- ^ Genotype to mutate + -> (r, g) randomMutation environment config generator chromosome = let (r, generator') = randomR (0.0, 1.0) generator in if r <= mutationProbability config then mutate environment generator' chromosome else (generator', chromosome) -naturalSelection :: (Species e g) => e -> EvolutionConfig -> [g] -> [g] +{- | +Select the fittest survivors from a population, +to be moved to the next pass of the algorithm. +-} +naturalSelection :: (Species e g) + => e -- ^ Environment + -> EvolutionConfig + -> [g] -- ^ Original population + -> [g] -- ^ Survivors naturalSelection environment config = map snd . take (survivors config) @@ -62,9 +112,12 @@ naturalSelection environment config -- Down reverses the sort order so that the best fitness comes first . map (\genotype -> (Down $ fitness environment genotype, genotype)) -evolveGeneration :: - (RandomGen r, Species e g) => - e -> EvolutionConfig -> (r, [g]) -> (r, [g]) +-- | Run one pass of the genetic algorithm over a given population. +evolveGeneration :: (RandomGen r, Species e g) + => e -- ^ Environment + -> EvolutionConfig + -> (r, [g]) -- ^ Random generator, original population + -> (r, [g]) -- ^ New random generator, new population evolveGeneration environment config (generator, population) = second (naturalSelection environment config) $ mapWithGen (randomMutation environment config) @@ -73,13 +126,24 @@ evolveGeneration environment config (generator, population) randomCrossover gen = let (pair, gen') = randomFromList gen pairs in (uncurry $ crossover environment gen') pair -initialGeneration :: - (RandomGen r, Species e g) => - e -> EvolutionConfig -> r -> (r, [g]) +{- | +Create the initial population, to be fed into the first +pass of the genetic algorithm. +-} +initialGeneration :: (RandomGen r, Species e g) + => e -- ^ Environment + -> EvolutionConfig + -> r -- ^ Random generator + -> (r, [g]) -- ^ New random generator, population initialGeneration environment config = unfoldWithGen (generate environment) (survivors config) -evolve :: (RandomGen r, Species e g) => e -> EvolutionConfig -> r -> (r, g) +-- | Run the full genetic algorithm. +evolve :: (RandomGen r, Species e g) + => e -- ^ Environment + -> EvolutionConfig + -> r -- ^ Random generator + -> (r, g) -- ^ New random generator, optimal genotype evolve environment config generator = second head $ repeatCall (generations config) (evolveGeneration environment config) diff --git a/palette-generator/Data/Colour.hs b/palette-generator/Data/Colour.hs index c894802f..6641e9b5 100644 --- a/palette-generator/Data/Colour.hs +++ b/palette-generator/Data/Colour.hs @@ -1,10 +1,12 @@ module Data.Colour ( LAB(..), RGB(..), deltaE, lab2rgb, rgb2lab ) where +-- | Lightness A-B data LAB a = LAB { lightness :: a , channelA :: a , channelB :: a } +-- | Red, Green, Blue data RGB a = RGB { red :: a , green :: a , blue :: a @@ -29,6 +31,7 @@ deltaE (LAB l1 a1 b1) (LAB l2 a2 b2) = i = deltaL^2 + deltaCkcsc^2 + deltaHkhsh^2 in if i < 0 then 0 else sqrt i +-- | Convert a 'LAB' colour to a 'RGB' colour lab2rgb :: (Floating a, Ord a) => LAB a -> RGB a lab2rgb (LAB l a bx) = let y = (l + 16) / 116 @@ -48,6 +51,7 @@ lab2rgb (LAB l a bx) = , blue = max 0 (min 1 b') * 255 } +-- | Convert a 'RGB' colour to a 'LAB' colour rgb2lab :: (Floating a, Ord a) => RGB a -> LAB a rgb2lab (RGB r g b) = let r' = r / 255 diff --git a/palette-generator/Stylix/Main.hs b/palette-generator/Stylix/Main.hs index f7c978db..1448b6dc 100644 --- a/palette-generator/Stylix/Main.hs +++ b/palette-generator/Stylix/Main.hs @@ -9,10 +9,14 @@ import System.Exit ( die ) import System.Random ( mkStdGen ) import Text.JSON ( encode ) -selectColours :: (Floating a, Real a) => V.Vector (LAB a) -> V.Vector (LAB a) +-- | Run the genetic algorithm to generate a palette from the given image. +selectColours :: (Floating a, Real a) + => V.Vector (LAB a) -- ^ Colours of the source image + -> V.Vector (LAB a) -- ^ Generated palette selectColours image = snd $ evolve image (EvolutionConfig 1000 100 0.5 150) (mkStdGen 0) +-- | Convert a 'DynamicImage' to a simple 'V.Vector' of colours. unpackImage :: (Num a) => DynamicImage -> V.Vector (RGB a) unpackImage image = do let image' = convertRGB8 image @@ -21,7 +25,9 @@ unpackImage image = do let (PixelRGB8 r g b) = pixelAt image' x y return $ RGB (fromIntegral r) (fromIntegral g) (fromIntegral b) -loadImage :: String -> IO DynamicImage +-- | Load an image file. +loadImage :: String -- ^ Path to the file + -> IO DynamicImage loadImage input = either error id <$> readImage input mainProcess :: (String, String) -> IO () diff --git a/palette-generator/Stylix/Output.hs b/palette-generator/Stylix/Output.hs index 15be0983..45d53af9 100644 --- a/palette-generator/Stylix/Output.hs +++ b/palette-generator/Stylix/Output.hs @@ -6,12 +6,20 @@ import Data.Word ( Word8 ) import Text.JSON ( JSObject, toJSObject ) import Text.Printf ( printf ) +-- | Convert any 'RGB' colour to store integers between 0 and 255. toWord8 :: (RealFrac a) => RGB a -> RGB Word8 toWord8 (RGB r g b) = RGB (truncate r) (truncate g) (truncate b) +{- | +Convert a colour to a hexdecimal string. + +>>> toHex (RGB 255 255 255) +"#ffffff" +-} toHex :: RGB Word8 -> String toHex (RGB r g b) = printf "%02x%02x%02x" r g b +-- | Convert a palette to the JSON format expected by Stylix's NixOS modules. makeOutputTable :: (RealFrac a) => V.Vector (RGB a) -> JSObject String makeOutputTable = toJSObject diff --git a/palette-generator/Stylix/Palette.hs b/palette-generator/Stylix/Palette.hs index 7fb65ee8..64da6a1e 100644 --- a/palette-generator/Stylix/Palette.hs +++ b/palette-generator/Stylix/Palette.hs @@ -9,21 +9,36 @@ import Data.Vector ( (!), (//) ) import qualified Data.Vector as V import System.Random ( RandomGen, randomR ) +-- | Extract the primary scale from a pallete. primary :: V.Vector a -> V.Vector a primary = V.take 8 +-- | Extract the accent colours from a palette. accent :: V.Vector a -> V.Vector a accent = V.drop 8 +{- | +Combine two palettes by taking a colour from the left, +then the right, then the left, and so on until we have +taken enough colours for a new palette. +-} alternatingZip :: V.Vector a -> V.Vector a -> V.Vector a alternatingZip = V.izipWith (\i a b -> if even i then a else b) -randomFromVector :: (RandomGen r) => r -> V.Vector a -> (a, r) +-- | Select a random item from a vector. +randomFromVector :: (RandomGen r) + => r -- ^ Random generator + -> V.Vector a + -> (a, r) -- ^ Chosen item, new random generator randomFromVector generator vector = let (index, generator') = randomR (0, V.length vector - 1) generator in (vector ! index, generator') instance (Floating a, Real a) => Species (V.Vector (LAB a)) (V.Vector (LAB a)) where + {- | + Palettes in the initial population are created by randomly + sampling 16 colours from the source image. + -} generate image = generateColour 16 where generateColour 0 generator = (generator, V.empty) generateColour n generator @@ -32,23 +47,48 @@ instance (Floating a, Real a) => Species (V.Vector (LAB a)) (V.Vector (LAB a)) w crossover _ generator a b = (generator, alternatingZip a b) + {- | + Mutation is done by replacing a random slot in the palette with + a new colour, which is randomly sampled from the source image. + -} mutate image generator palette = let (index, generator') = randomR (0, 15) generator (colour, generator'') = randomFromVector generator' image in (generator'', palette // [(index, colour)]) fitness _ palette - = realToFrac $ accentDifference - min lightScheme darkScheme - where accentDifference = minimum $ do - a <- accent palette - b <- accent palette - return $ deltaE a b - lightnesses = V.map lightness palette - difference a b = abs $ a - b - lightnessError primaryScale accentValue - = sum (V.zipWith difference primaryScale $ primary lightnesses) - + sum (V.map (difference accentValue) $ accent lightnesses) - lightScheme - = lightnessError (V.fromList [90, 70, 55, 35, 25, 10, 5, 5]) 40 - darkScheme - = lightnessError (V.fromList [10, 30, 45, 65, 75, 90, 95, 95]) 60 + = realToFrac $ + accentDifference - + -- Either light schemes or dark themes are allowed. + -- We try to converge on whichever theme we are closer to. + min lightScheme darkScheme + where + -- The accent colours should be as different as possible. + accentDifference = minimum $ do + a <- accent palette + b <- accent palette + return $ deltaE a b + + -- Helpers for the function below. + lightnesses = V.map lightness palette + difference a b = abs $ a - b + + lightnessError primaryScale accentValue + -- The primary scale's lightnesses should match the given pattern. + = sum (V.zipWith difference primaryScale $ primary lightnesses) + -- The accent colours should all have the given lightness. + + sum (V.map (difference accentValue) $ accent lightnesses) + + {- + For light themes, the background is bright and the text is dark. + The accent colours are slightly darker. + -} + lightScheme + = lightnessError (V.fromList [90, 70, 55, 35, 25, 10, 5, 5]) 40 + + {- + For dark themes, the background is dark and the text is bright. + The accent colours are slightly brighter. + -} + darkScheme + = lightnessError (V.fromList [10, 30, 45, 65, 75, 90, 95, 95]) 60