extra/math/dual/dual.factor

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