<commit_before>{-# LANGUAGE
  GADTs
 #-}

module Primitive where

import Numeric.SpecFunctions
import Data.Number.LogFloat (LogFloat, logFloat, logToLogFloat)
import Data.Typeable

-- | Primitive distributions for which we can compute density.
-- Here the weights of Categorical must be normalized.
data Primitive a where
    Categorical :: (Eq a, Typeable a) => [(a,LogFloat)] -> Primitive a
    Normal :: Double -> Double -> Primitive Double
    Gamma :: Double -> Double -> Primitive Double
    Beta :: Double -> Double -> Primitive Double

-- | The probability density function.
pdf :: Primitive a -> a -> LogFloat
pdf (Categorical d) = \x -> case lookup x d of
                              Just p -> p
                              Nothing -> logFloat 0
pdf (Normal m s) = normalPdf m s
pdf (Gamma  a b) = gammaPdf a b
pdf (Beta   a b) = betaPdf a b


-- | PDF of normal distribution parameterized by mean and stddev.
normalPdf :: (Real a) => a -> a -> a -> LogFloat
normalPdf mu sigma x =
  logToLogFloat $ (-0.5 * log (2 * pi * sigma2)) +
  ((-((realToFrac x) - (realToFrac mu))^2) / (2 * sigma2))
  where
    sigma2 = realToFrac sigma^2<commit_msg>make (Primitive a) comparable for equality<commit_after>{-# LANGUAGE
  GADTs
 #-}

module Primitive where

import Numeric.SpecFunctions
import Data.Number.LogFloat (LogFloat, logFloat, logToLogFloat)
import Data.Typeable

-- | Primitive distributions for which we can compute density.
-- Here the weights of Categorical must be normalized.
data Primitive a where
    Categorical :: (Eq a, Typeable a) => [(a,LogFloat)] -> Primitive a
    Normal :: Double -> Double -> Primitive Double
    Gamma :: Double -> Double -> Primitive Double
    Beta :: Double -> Double -> Primitive Double

instance Eq (Primitive a) where
  Categorical xs == Categorical ys = xs == ys
  Normal m s     == Normal m' s'   = m == m' && s == s'
  Gamma  a b     == Gamma  a' b'   = a == a' && b == b'
  Beta   a b     == Beta   a' b'   = a == a' && b == b'
  _              == _              = False

-- | The probability density function.
pdf :: Primitive a -> a -> LogFloat
pdf (Categorical d) = \x -> case lookup x d of
                              Just p -> p
                              Nothing -> logFloat 0
pdf (Normal m s) = normalPdf m s
pdf (Gamma  a b) = gammaPdf a b
pdf (Beta   a b) = betaPdf a b


-- | PDF of normal distribution parameterized by mean and stddev.
normalPdf :: (Real a) => a -> a -> a -> LogFloat
normalPdf mu sigma x =
  logToLogFloat $ (-0.5 * log (2 * pi * sigma2)) +
  ((-((realToFrac x) - (realToFrac mu))^2) / (2 * sigma2))
  where
    sigma2 = realToFrac sigma^2
{'fraction_non_alphanumeric': 0.11963539688568173, 'fraction_numerical': 0.007595898214963919, 'mean_word_length': 2.4163424124513617, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 37, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5853700', 'n_tokens_mistral': 858, 'n_tokens_neox': 830, 'n_words': 324}