module Domain.Math.Data.SquareRoot
( SquareRoot
, imaginary, imaginaryUnit
, con, toList, scale, fromSquareRoot
, sqrt, sqrtRational, isqrt, eval
) where
import Control.Monad
import Data.Ratio
import Domain.Math.Safe
import Prelude hiding (sqrt)
import Test.QuickCheck hiding (scale)
import qualified Data.Map as M
import qualified Domain.Math.Data.PrimeFactors as P
import qualified Prelude
data SquareRoot a = S
{ imaginary :: Bool
, squareRootMap :: SqMap a
} deriving (Eq, Ord)
type SqMap a = M.Map P.PrimeFactors a
makeMap :: (Eq a,Num a) => SqMap a -> SqMap a
makeMap = M.filter (/=0) . M.foldrWithKey f M.empty
where
f k a m
| a == 0 = m
| otherwise = M.unionWith (+) (fmap (*a) (sqrtPF k)) m
plusSqMap :: (Eq a,Num a) => SqMap a -> SqMap a -> SqMap a
plusSqMap m1 m2 = M.filter (/=0) (M.unionWith (+) m1 m2)
minusSqMap :: (Eq a,Num a) => SqMap a -> SqMap a -> SqMap a
minusSqMap m1 m2 = m1 `plusSqMap` negateSqMap m2
negateSqMap :: Num a => SqMap a -> SqMap a
negateSqMap = fmap negate
timesSqMap :: (Eq a,Num a) => SqMap a -> SqMap a -> SqMap a
timesSqMap m1 m2 =
case (M.toList m1, M.toList m2) of
([], _) -> M.empty
(_, []) -> M.empty
([(n, a)], _) | n==1 -> if a==0 then M.empty else fmap (*a) m2
(_, [(n, a)]) | n==1 -> if a==0 then M.empty else fmap (*a) m1
_ ->
let op n a = M.unionWith (+) (f n (fmap (a *) m1))
f i = M.mapKeys (*i)
in makeMap (M.foldrWithKey op M.empty m2)
recipSqMap :: (Eq a,Fractional a) => SqMap a -> SqMap a
recipSqMap m =
case M.toList m of
[] -> error "SquareRoot: division by zero"
[(n, x)] -> M.singleton n (recip (x * fromIntegral n))
_ -> (a .-. b) .*. recipSqMap (makeMap ((a .*. a) .-. (b .*. b)))
where
(ys, zs) = splitAt (length xs `div` 2) xs
(a, b) = (M.fromList ys, M.fromList zs)
xs = M.toList m
(.*.) = timesSqMap
(.-.) = minusSqMap
sqrtPF :: Num a => P.PrimeFactors -> SqMap a
sqrtPF n
| n == 0 = M.empty
| otherwise = M.singleton b (fromIntegral a)
where
(a, b) = P.splitPower 2 n
instance (Show a,Eq a,Num a) => Show (SquareRoot a) where
show (S isNeg m) = g (map f (M.toList m)) ++ imPart
where
f (n, a) = ( signum a == -1
, times (guard (abs a /= 1) >> Just (show (abs a)))
(guard (n /= 1) >> Just ("sqrt(" ++ show (toInteger n) ++ ")"))
)
imPart = if isNeg then " (imaginary number)" else ""
g [] = "0"
g ((b,x):xs) = (if b then "-" else "") ++ x ++ concatMap h xs
h (b, x) = (if b then " - " else " + ") ++ x
times (Just a) (Just b) = a ++ "*" ++ b
times (Just a) Nothing = a
times Nothing (Just b) = b
times Nothing Nothing = "1"
instance Functor SquareRoot where
fmap f (S b m) = S b (M.map f m)
instance (Eq a,Num a) => Num (SquareRoot a) where
S b1 m1 + S b2 m2 = S (b1 || b2) (plusSqMap m1 m2)
S b1 m1 - S b2 m2 = S (b1 || b2) (minusSqMap m1 m2)
S b1 m1 * S b2 m2 = S (b1 || b2) (timesSqMap m1 m2)
negate (S b m) = S b (negateSqMap m)
fromInteger = con . fromInteger
abs = error "abs not defined for square roots"
signum = error "signum not defined for square roots"
instance (Eq a,Fractional a) => SafeDiv (SquareRoot a) where
safeDiv x y
| y == 0 = Nothing
| otherwise = Just (x/y)
instance (Eq a,Fractional a) => Fractional (SquareRoot a) where
recip (S b m) = S b (recipSqMap m)
fromRational = con . fromRational
instance (Eq a,Fractional a) => Arbitrary (SquareRoot a) where
arbitrary = sized $ \n -> do
let make r1 r2 = fromRational r1 * sqrtRational r2
i <- choose (0, 10)
xs <- vectorOf i (liftM2 make (rationalGen n) (rationalGen n))
return (sum xs)
rationalGen :: Int -> Gen Rational
rationalGen n = do
c <- choose (0, n)
b <- choose (0, c)
a <- choose (0, n)
return $ fromIntegral a + if c==0 then 0
else fromIntegral b / fromIntegral c
imaginaryUnit :: Num a => SquareRoot a
imaginaryUnit = S True (M.singleton (-1) 1)
toList :: SquareRoot a -> [(a, Integer)]
toList = map (\(k, r) -> (r, toInteger k)) . M.toList . squareRootMap
fromSquareRoot :: Num a => SquareRoot a -> Maybe a
fromSquareRoot a =
case toList a of
[(b, n)] | n==1 -> Just b
[] -> Just 0
_ -> Nothing
con :: (Eq a,Num a) => a -> SquareRoot a
con a = S False (if a==0 then M.empty else M.singleton 1 a)
sqrt :: Num a => Integer -> SquareRoot a
sqrt n
| n < 0 = S True (M.mapKeys negate m)
| otherwise = S False m
where
m = sqrtPF (fromIntegral (abs n))
scale :: (Eq a,Num a) => a -> SquareRoot a -> SquareRoot a
scale a sr = if a==0 then 0 else fmap (*a) sr
isqrt :: Integer -> Integer
isqrt = (floor :: Double -> Integer) . Prelude.sqrt . fromInteger
sqrtRational :: (Eq a,Fractional a) => Rational -> SquareRoot a
sqrtRational r = scale (1/fromIntegral b) (sqrt (a*b))
where
(a, b) = (numerator r, denominator r)
eval :: Floating a => SquareRoot a -> a
eval (S _ m) = M.foldrWithKey f 0 m
where f n a b = a * Prelude.sqrt (fromIntegral n) + b