extra/math/dual/dual.factor

   1 ! Copyright (C) 2009 Jason W. Merrill.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: kernel math math.functions math.derivatives accessors
   4     macros generic compiler.units words effects vocabs
   5     sequences arrays assocs generalizations fry make
   6     combinators.smart help help.markup ;
   7
   8 IN: math.dual
   9
  10 TUPLE: dual ordinary-part epsilon-part ;
  11
  12 C: <dual> dual
  13
  14 ! Ordinary numbers implement the dual protocol by returning
  15 ! themselves as the ordinary part, and 0 as the epsilon part.
  16 M: number ordinary-part>> ;
  17
  18 M: number epsilon-part>> drop 0 ;
  19
  20 : unpack-dual ( dual -- ordinary-part epsilon-part )
  21     [ ordinary-part>> ] [ epsilon-part>> ] bi ;
  22
  23 <PRIVATE
  24
  25 : input-length ( word -- n ) stack-effect in>> length ;
  26
  27 MACRO: ordinary-op ( word -- o )
  28     [ input-length ] keep
  29     '[ [ ordinary-part>> ] _ napply _ execute ] ;
  30
  31 ! Takes N dual numbers <o1,e1> <o2,e2> ... <oN,eN> and weaves
  32 ! their ordinary and epsilon parts to produce
  33 ! e1 o1 o2 ... oN e2 o1 o2 ... oN ... eN o1 o2 ... oN
  34 ! This allows a set of partial derivatives each to be evaluated
  35 ! at the same point.
  36 MACRO: duals>nweave ( n -- )
  37    dup dup dup
  38    '[
  39        [ [ epsilon-part>> ] _ napply ]
  40        _ nkeep
  41        [ ordinary-part>> ] _ napply
  42        _ nweave
  43     ] ;
  44
  45 MACRO: chain-rule ( word -- e )
  46     [ input-length '[ _ duals>nweave ] ]
  47     [ "derivative" word-prop ]
  48     [ input-length 1 + '[ _ nspread ] ]
  49     tri
  50     '[ [ @ _ @ ] sum-outputs ] ;
  51
  52 : set-dual-help ( word dword -- )
  53     [ swap
  54         [ stack-effect [ in>> ] [ out>> ] bi append
  55             [ dual ] { } map>assoc { $values } prepend
  56         ]
  57         [ [ { $description } % "Version of " ,
  58                    { $link } swap suffix ,
  59                    " extended to work on dual numbers." , ]
  60             { } make
  61         ]
  62         bi* 2array
  63     ] keep set-word-help ;
  64
  65 PRIVATE>
  66
  67 MACRO: dual-op ( word -- )
  68     [ '[ _ ordinary-op ] ]
  69     [ input-length '[ _ nkeep ] ]
  70     [ '[ _ chain-rule ] ]
  71     tri
  72     '[ _ @ @ <dual> ] ;
  73
  74 : define-dual ( word -- )
  75     dup name>> "d" prepend "math.dual" create
  76     [ [ stack-effect ] dip set-stack-effect ]
  77     [ set-dual-help ]
  78     [ swap '[ _ dual-op ] define ]
  79     2tri ;
  80
  81 ! Specialize math functions to operate on dual numbers.
  82 [ all-words [ "derivative" word-prop ] filter
  83     [ define-dual ] each ] with-compilation-unit