text-metrics-0.3.2: Data/Text/Metrics.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MultiWayIf #-}
-- |
-- Module : Data.Text.Metrics
-- Copyright : © 2016–present Mark Karpov
-- License : BSD 3 clause
--
-- Maintainer : Mark Karpov <markkarpov92@gmail.com>
-- Stability : experimental
-- Portability : portable
--
-- The module provides efficient implementations of various strings metric
-- algorithms. It works with strict 'Text' values.
--
-- __Note__: before version /0.3.0/ the package used C implementations of
-- the algorithms under the hood. Beginning from version /0.3.0/, the
-- implementations are written in Haskell while staying almost as fast, see:
--
-- <https://markkarpov.com/post/migrating-text-metrics.html>
module Data.Text.Metrics
( -- * Levenshtein variants
levenshtein,
levenshteinNorm,
damerauLevenshtein,
damerauLevenshteinNorm,
-- * Treating inputs like sets
overlap,
jaccard,
-- * Other
hamming,
jaro,
jaroWinkler,
)
where
import Control.Monad
import Control.Monad.ST
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import Data.Ratio
import Data.Text
import qualified Data.Text as T
import qualified Data.Text.Internal as T
import qualified Data.Text.Unsafe as TU
import qualified Data.Vector.Unboxed.Mutable as VUM
import GHC.Exts (inline)
----------------------------------------------------------------------------
-- Levenshtein variants
-- | Return the Levenshtein distance between two 'Text' values. The
-- Levenshtein distance between two strings is the minimal number of
-- operations necessary to transform one string into another. For the
-- Levenshtein distance allowed operations are: deletion, insertion, and
-- substitution.
--
-- See also: <https://en.wikipedia.org/wiki/Levenshtein_distance>.
--
-- __Heads up__, before version /0.3.0/ this function returned
-- 'Data.Numeric.Natural'.
levenshtein :: Text -> Text -> Int
levenshtein a b = fst (levenshtein_ a b)
-- | Return the normalized Levenshtein distance between two 'Text' values.
-- Result is a non-negative rational number (represented as @'Ratio'
-- 'Data.Numeric.Natural'@), where 0 signifies no similarity between the
-- strings, while 1 means exact match.
--
-- See also: <https://en.wikipedia.org/wiki/Levenshtein_distance>.
--
-- __Heads up__, before version /0.3.0/ this function returned @'Ratio'
-- 'Data.Numeric.Natural'@.
levenshteinNorm :: Text -> Text -> Ratio Int
levenshteinNorm = norm levenshtein_
-- | An internal helper, returns the Levenshtein distance as the first
-- element of the tuple and max length of the two inputs as the second
-- element of the tuple.
levenshtein_ :: Text -> Text -> (Int, Int)
levenshtein_ a b
| T.null a = (lenb, lenm)
| T.null b = (lena, lenm)
| otherwise = runST $ do
let v_len = lenb + 1
v <- VUM.unsafeNew (v_len * 2)
let gov !i =
when (i < v_len) $ do
VUM.unsafeWrite v i i
gov (i + 1)
goi !i !na !v0 !v1 = do
let !(TU.Iter ai da) = TU.iter a na
goj !j !nb =
when (j < lenb) $ do
let !(TU.Iter bj db) = TU.iter b nb
cost = if ai == bj then 0 else 1
x <- (+ 1) <$> VUM.unsafeRead v (v1 + j)
y <- (+ 1) <$> VUM.unsafeRead v (v0 + j + 1)
z <- (+ cost) <$> VUM.unsafeRead v (v0 + j)
VUM.unsafeWrite v (v1 + j + 1) (min x (min y z))
goj (j + 1) (nb + db)
when (i < lena) $ do
VUM.unsafeWrite v v1 (i + 1)
goj 0 0
goi (i + 1) (na + da) v1 v0
gov 0
goi 0 0 0 v_len
ld <- VUM.unsafeRead v (lenb + if even lena then 0 else v_len)
return (ld, lenm)
where
lena = T.length a
lenb = T.length b
lenm = max lena lenb
{-# INLINE levenshtein_ #-}
-- | Return the Damerau-Levenshtein distance between two 'Text' values. The
-- function works like 'levenshtein', but the collection of allowed
-- operations also includes transposition of two /adjacent/ characters.
--
-- See also: <https://en.wikipedia.org/wiki/Damerau%E2%80%93Levenshtein_distance>.
--
-- __Heads up__, before version /0.3.0/ this function returned
-- 'Data.Numeric.Natural'.
damerauLevenshtein :: Text -> Text -> Int
damerauLevenshtein a b = fst (damerauLevenshtein_ a b)
-- | Return the normalized Damerau-Levenshtein distance between two 'Text'
-- values. 0 signifies no similarity between the strings, while 1 means
-- exact match.
--
-- See also: <https://en.wikipedia.org/wiki/Damerau%E2%80%93Levenshtein_distance>.
--
-- __Heads up__, before version /0.3.0/ this function returned @'Ratio'
-- 'Data.Numeric.Natural'@.
damerauLevenshteinNorm :: Text -> Text -> Ratio Int
damerauLevenshteinNorm = norm damerauLevenshtein_
-- | An internal helper, returns the Damerau-Levenshtein distance as the
-- first element of the tuple and max length of the two inputs as the second
-- element of the tuple.
damerauLevenshtein_ :: Text -> Text -> (Int, Int)
damerauLevenshtein_ a b
| T.null a = (lenb, lenm)
| T.null b = (lena, lenm)
| otherwise = runST $ do
let v_len = lenb + 1
v <- VUM.unsafeNew (v_len * 3)
let gov !i =
when (i < v_len) $ do
VUM.unsafeWrite v i i
gov (i + 1)
goi !i !na !ai_1 !v0 !v1 !v2 = do
let !(TU.Iter ai da) = TU.iter a na
goj !j !nb !bj_1 =
when (j < lenb) $ do
let !(TU.Iter bj db) = TU.iter b nb
cost = if ai == bj then 0 else 1
x <- (+ 1) <$> VUM.unsafeRead v (v1 + j)
y <- (+ 1) <$> VUM.unsafeRead v (v0 + j + 1)
z <- (+ cost) <$> VUM.unsafeRead v (v0 + j)
let g = min x (min y z)
val <- (+ cost) <$> VUM.unsafeRead v (v2 + j - 1)
VUM.unsafeWrite v (v1 + j + 1) $
if i > 0 && j > 0 && ai == bj_1 && ai_1 == bj && val < g
then val
else g
goj (j + 1) (nb + db) bj
when (i < lena) $ do
VUM.unsafeWrite v v1 (i + 1)
goj 0 0 'a'
goi (i + 1) (na + da) ai v1 v2 v0
gov 0
goi 0 0 'a' 0 v_len (v_len * 2)
ld <- VUM.unsafeRead v (lenb + (lena `mod` 3) * v_len)
return (ld, lenm)
where
lena = T.length a
lenb = T.length b
lenm = max lena lenb
{-# INLINE damerauLevenshtein_ #-}
----------------------------------------------------------------------------
-- Treating inputs like sets
-- | Return the overlap coefficient for two 'Text' values. Returned value is
-- in the range from 0 (no similarity) to 1 (exact match). Return 1 if both
-- 'Text' values are empty.
--
-- See also: <https://en.wikipedia.org/wiki/Overlap_coefficient>.
--
-- @since 0.3.0
overlap :: Text -> Text -> Ratio Int
overlap a b =
if d == 0
then 1 % 1
else intersectionSize (mkTextMap a) (mkTextMap b) % d
where
d = min (T.length a) (T.length b)
-- | Return the Jaccard similarity coefficient for two 'Text' values.
-- Returned value is in the range from 0 (no similarity) to 1 (exact match).
-- Return 1 if both
--
-- See also: <https://en.wikipedia.org/wiki/Jaccard_index>
--
-- @since 0.3.0
jaccard :: Text -> Text -> Ratio Int
jaccard a b =
if d == 0
then 1 % 1
else intersectionSize ma mb % d
where
ma = mkTextMap a
mb = mkTextMap b
d = unionSize ma mb
-- | Make a map from 'Char' to 'Int' representing how many times the 'Char'
-- appears in the input 'Text'.
mkTextMap :: Text -> Map Char Int
mkTextMap = T.foldl' f M.empty
where
f m ch = M.insertWith (+) ch 1 m
{-# INLINE mkTextMap #-}
-- | Return intersection size between two 'Text'-maps.
intersectionSize :: Map Char Int -> Map Char Int -> Int
intersectionSize a b = M.foldl' (+) 0 (M.intersectionWith min a b)
{-# INLINE intersectionSize #-}
-- | Return union size between two 'Text'-maps.
unionSize :: Map Char Int -> Map Char Int -> Int
unionSize a b = M.foldl' (+) 0 (M.unionWith max a b)
{-# INLINE unionSize #-}
----------------------------------------------------------------------------
-- Other
-- | /O(n)/ Return the Hamming distance between two 'Text' values. Hamming
-- distance is defined as the number of positions at which the corresponding
-- symbols are different. The input 'Text' values should be of equal length
-- or 'Nothing' will be returned.
--
-- See also: <https://en.wikipedia.org/wiki/Hamming_distance>.
--
-- __Heads up__, before version /0.3.0/ this function returned @'Maybe'
-- 'Data.Numeric.Natural'@.
hamming :: Text -> Text -> Maybe Int
hamming a@(T.Text _ _ len) b =
if T.length a == T.length b
then Just (go 0 0 0)
else Nothing
where
go !na !nb !r =
let !(TU.Iter cha da) = TU.iter a na
!(TU.Iter chb db) = TU.iter b nb
in if
| na == len -> r
| cha /= chb -> go (na + da) (nb + db) (r + 1)
| otherwise -> go (na + da) (nb + db) r
-- | Return the Jaro distance between two 'Text' values. Returned value is
-- in the range from 0 (no similarity) to 1 (exact match).
--
-- While the algorithm is pretty clear for artificial examples (like those
-- from the linked Wikipedia article), for /arbitrary/ strings, it may be
-- hard to decide which of two strings should be considered as one having
-- “reference” order of characters (order of matching characters in an
-- essential part of the definition of the algorithm). This makes us
-- consider the first string the “reference” string (with correct order of
-- characters). Thus generally,
--
-- > jaro a b ≠ jaro b a
--
-- This asymmetry can be found in all implementations of the algorithm on
-- the internet, AFAIK.
--
-- See also: <https://en.wikipedia.org/wiki/Jaro%E2%80%93Winkler_distance>
--
-- @since 0.2.0
--
-- __Heads up__, before version /0.3.0/ this function returned @'Ratio'
-- 'Data.Numeric.Natural'@.
jaro :: Text -> Text -> Ratio Int
jaro a b =
if T.null a || T.null b
then 0 % 1
else runST $ do
let lena = T.length a
lenb = T.length b
d =
if lena >= 2 && lenb >= 2
then max lena lenb `quot` 2 - 1
else 0
v <- VUM.replicate lenb (0 :: Int)
r <- VUM.replicate 3 (0 :: Int) -- tj, m, t
let goi !i !na !fromb = do
let !(TU.Iter ai da) = TU.iter a na
(from, fromb') =
if i >= d
then (i - d, fromb + TU.iter_ b fromb)
else (0, 0)
to = min (i + d + 1) lenb
goj !j !nb =
when (j < to) $ do
let !(TU.Iter bj db) = TU.iter b nb
used <- (== 1) <$> VUM.unsafeRead v j
if not used && ai == bj
then do
tj <- VUM.unsafeRead r 0
if j < tj
then VUM.unsafeModify r (+ 1) 2
else VUM.unsafeWrite r 0 j
VUM.unsafeWrite v j 1
VUM.unsafeModify r (+ 1) 1
else goj (j + 1) (nb + db)
when (i < lena) $ do
goj from fromb
goi (i + 1) (na + da) fromb'
goi 0 0 0
m <- VUM.unsafeRead r 1
t <- VUM.unsafeRead r 2
return $
if m == 0
then 0 % 1
else
( (m % lena)
+ (m % lenb)
+ ((m - t) % m)
)
/ 3
-- | Return the Jaro-Winkler distance between two 'Text' values. Returned
-- value is in range from 0 (no similarity) to 1 (exact match).
--
-- See also: <https://en.wikipedia.org/wiki/Jaro%E2%80%93Winkler_distance>
--
-- @since 0.2.0
--
-- __Heads up__, before version /0.3.0/ this function returned @'Ratio'
-- 'Data.Numeric.Natural'@.
jaroWinkler :: Text -> Text -> Ratio Int
jaroWinkler a b = dj + (1 % 10) * l * (1 - dj)
where
dj = inline (jaro a b)
l = fromIntegral (min 4 (commonPrefix a b))
-- | Return the length of the common prefix two 'Text' values have.
commonPrefix :: Text -> Text -> Int
commonPrefix a b = case T.commonPrefixes a b of
Nothing -> 0
Just (pref, _, _) -> T.length pref
{-# INLINE commonPrefix #-}
----------------------------------------------------------------------------
-- Helpers
norm :: (Text -> Text -> (Int, Int)) -> Text -> Text -> Ratio Int
norm f a b =
let (r, l) = f a b
in if r == 0
then 1 % 1
else 1 % 1 - r % l
{-# INLINE norm #-}