Search code examples
haskellforallautomatic-differentiationrank-n-types

How to properly match types when using Numeric.AD in Haskell?

I'm trying to implement Newton–Raphson root-finding algorithm using the ad package, but I can't properly match function types. I know there's a proper answer to a similar question, which was answered by the creator of ad himself, but the package have changed greatly since version 1.0.6 (current version is 4.3.4).

This first minimum example compiles and works when I iterate over it: 

import Numeric.AD
import Numeric.AD.Internal.Forward

g :: Fractional a => a -> a
g x = - x + 2

g' :: Fractional a => a -> a
g' x = diff g x

newtonG :: Fractional a => a -> a
newtonG x = x - (g x) / (g' x)

But if I try to abstract the function, like this: 

import Numeric.AD
import Numeric.AD.Internal.Forward

g :: Fractional a => a -> a
g x = - x + 2

newton :: Fractional a => (a -> a) -> a -> a
newton f x = x - (f x) / (diff f x)

GHC returns the following error: 

fsolve.hs:8:32: error:
    * Couldn't match type `a' with `AD s (Forward a)'
      `a' is a rigid type variable bound by
        the type signature for:
          newton :: forall a. Fractional a => (a -> a) -> a -> a
        at fsolve.hs:7:11
      Expected type: AD s (Forward a) -> AD s (Forward a)
        Actual type: a -> a
    * In the first argument of `diff', namely `f'
      In the second argument of `(/)', namely `(diff f x)'
      In the second argument of `(-)', namely `(f x) / (diff f x)'
    * Relevant bindings include
        x :: a (bound at fsolve.hs:8:10)
        f :: a -> a (bound at fsolve.hs:8:8)
        newton :: (a -> a) -> a -> a (bound at fsolve.hs:8:1)

If I use Numeric.AD.Rank1.Forward instead of Numeric.AD, the compiler says it cannot match a with Forward a, instead of a with AD s (Forward a). I also tried several ways of creating a dual number from x to pass it to f, e.g. snd . unbundle . f $ bundle x 1, but it only works if I create a new g' x using it, like in first case. Using this within newton doesn't work either.

In Numeric.AD, diff :: Num a => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> a. And in Numeric.AD.Rank1.Forward, it's diff :: Num a => (Forward a -> Forward a) -> a -> a. So why do they accept a function of type a -> a in the first case, but not in the second? Besides using polymorphic functions, should I take any special care when creating functions to use with Numeric.AD? Lastly, how should I change my code to make it work? I know the package already has a function to find roots, but I don't want to use yet (since I'm still learning Haskell), and looking at the documentation trying to sort this out, felt like running in circles.

Solution

  • Observe that your function:

    newton :: Fractional a => (a -> a) -> a -> a
newton f x = x - (f x) / (diff f x)

    uses the function argument f in two places. First, the subexpression f x uses f with the type:

    f :: Fractional a => a -> a

    Second, due to the use of diff the subexpression diff f x uses f with the type:

    f :: forall s a. Fractional a => AD s (Forward a) -> AD s (Forward a)

    The error message you got is the type system observing that these are different types that cannot be unified.

    The solution is to explicitly quantify the function argument of newton to work for all types satisfying the appropriate numeric type class constraint(s). This requires the RankNTypes language extension:

    {-# LANGUAGE RankNTypes #-}

newton
  :: Fractional a
  => (forall b. Fractional b => b -> b)
  -> a
  -> a
newton f x = x - (f x) / (diff f x)