Paste: GP

USING: kernel accessors math combinators sequences continuations random
arrays math.order sorting quotations grouping fry prettyprint
io locals ;
IN: ga
! this is a tool-box vocab to create simple genetic-programs
! you have to:
!   - declare a new class derivated from ga
!   - declare at least compare-output-to-awaited-result method
!   - adapt other methods when needed
! useful words are after :
!   - init-population
!   - evaluate-population to give each genome a fitness
!   - ga-step to do a loop
!   - ga-step. to do a loop and display best result
!   - n-ga-step. to do n loops max and display best result
!                 loop stops when best fitness is 0
!   - ga-execute to use best genome

! ------------------------------------
! GA class
! ------------------------------------

: nop ( -- ) ;
: test-quot-on-stack ( stack quot -- stack )
   '[ _ _ with-datastack ] [ drop f ] recover ;
: rand-in-%? ( n -- ? )  100 random > ;

TUPLE: genome 
    { sequence initial: f } 
    { fitness initial: f }  ;
! a genome has a false fitness when it doesn't compile

: <genome> (  -- obj )  \ genome new ;

: seq>genome ( seq -- obj ) <genome> swap >>sequence ;

TUPLE: ga population-size possible-genes 
    crossover-rate mutation-rate genome-length 
    training-input-stack training-output-stack
    { population  initial: f } 
    { #generation initial: 0 } 
    { evaluated?  initial: f } ;

! GA words that can be defined for a new type of GA
GENERIC: rand-gene> ( ga -- gene )
GENERIC: genome-decoder ( genome ga -- obj )
GENERIC: calculate-genome-fitness ( genome ga -- )
GENERIC: generate-random-genome-sequence ( ga -- seq )
GENERIC: renew-population ( ga -- ga )
GENERIC: compare-output-to-awaited-result ( stack ga -- x )
GENERIC: genome-stack-formalizer ( quot ga -- quot )
! GENERIC: wedding ( ga -- ga )
! M: ga wedding 
!  [ dup 2 <clumps> 
!  [ [ crossover ] with-datastack ] map concat append ] 
! change-population 

<PRIVATE
! ------------------------------------------
! mutations
! ------------------------------------------

: sequence-mutate ( seq ga -- seq )
    [ mutation-rate>> ] [ ] bi
    '[  _ rand-in-%? [ drop _ rand-gene> ] when ] map ;

: genome-mutate ( genome ga -- )
    [ sequence-mutate ] curry change-sequence drop ;

: mutate-population ( ga -- ga )
    ! 1 tail to exclude best genome
    [ ] [ population>> 1 tail ] [ ] tri 
    [ genome-mutate ] curry each ;

! ------------------------------------------
! cross-over
! ------------------------------------------
! this random is not very nice
: sequence-crossover ( seq seq -- seq seq )
    [ dup length 1+ random cut ] bi@ 
    spin [ append ] 2bi@ ;

! [ sequence-crossover ] 2bi@-change-sequence ? :-)
: genome-crossover ( genome genome -- genome genome )
    2dup [ sequence>> ]  bi@
    sequence-crossover
    swapd
    [ >>sequence ] 2bi@ ;

! =================================

: sort-genome-seq ( seq -- seq )
    [ [ fitness>> ] compare ] sort ;

: order-population ( ga -- ga ) 
    ! remove all genomes whose fitness is false
    [   [ fitness>> ] filter 
        sort-genome-seq 
    ] change-population ;

: calculate-fitness ( ga -- ga )
    [ ] 
    [ population>> ]
    [ '[ _ calculate-genome-fitness ] ] 
    tri each  ;

: evaluate-population ( ga -- ga )
    calculate-fitness
    order-population
    t >>evaluated? ;

: n-random-genomes ( n ga -- seq )
    '[ _ generate-random-genome-sequence seq>genome ] 
    replicate ;

: crossover-population ( ga -- ga ) 
    dup crossover-rate>>
    '[   unclip clone tuck clone 
        [   _ rand-in-%? 
            [ clone genome-crossover drop ] 
            [ drop ] if        
        ] curry map 
        swap prefix    
    ] change-population ;

: complete-population ( ga -- ga ) 
   dup [ population-size>> ] [ population>> length ]  bi -
   over n-random-genomes
   [ append ] curry change-population  ;

: execute-genome ( genome ga -- )  genome-decoder call ; inline

PRIVATE>

: init-population ( ga -- ga )
    dup
    [ population-size>> ] keep
    n-random-genomes  >>population 
    0 >>#generation
    f >>evaluated? ;

: new-generation ( ga -- ga )
    dup   population>>  length 0 = 
    [ init-population ] [ renew-population ]  if
    f >>evaluated?
    [ 1+ ] change-#generation ;
 
: ga-step ( ga -- ga ) 
    new-generation 
    evaluate-population ;

: best-genome> ( ga -- genome ) population>> first ;
: best-fitness> ( ga -- x ) best-genome> fitness>> ;
: ga-execute ( ga -- ) [ best-genome> ] keep execute-genome ; 

: ga-situation ( ga -- ga ) 
    [ ] [ "generation:   " write #generation>> . ] [ population>> ] tri
    [ first 
    [ "best genome:  "  write sequence>> . ] 
    [ "best fitness: " write fitness>>  . ] bi 
    ] [ 2drop "nothing good to display" . ] recover ;

:: (n-ga-step) ( ga n -- ga )
    n
    [ 1- dup  1 < ga best-fitness> 0 = or ] 
    [ ga ga-step drop ] 
    [ drop ga ] 
    until ;

: n-ga-step ( ga n -- ga ) 
    [ ga-step ] dip (n-ga-step) ;
   
: ga-step.   ( ga -- ga )      ga-step ga-situation ;
: n-ga-step. ( ga n -- ga )    n-ga-step  ga-situation ;

! ---------------------------------------------------------------------------
! default ga methods
! ---------------------------------------------------------------------------

M: ga rand-gene>                possible-genes>> random ;
M: ga genome-decoder            drop sequence>> >quotation ;
M: ga genome-stack-formalizer   drop ;
M: ga calculate-genome-fitness
    over sequence>> length 0 = [ drop f ] [ ! remove empty sequence
        [ ]
        [ nip training-input-stack>> ] ! ( genome ga -- stack )
        [ [ genome-decoder ] keep genome-stack-formalizer ] 
                    ! ( genome ga -- quot )
        2tri 
        ! we could use stack effect here to select genome
        test-quot-on-stack ! ( stack quot -- stack )
        [ swap compare-output-to-awaited-result ] [ drop f ] if*  
        ! ( genome ga stack -- genome x )
    ] if
    >>fitness drop ! ( genome x -- )  
    ;

M: ga renew-population 
    mutate-population
    complete-population
    crossover-population ;

M: ga generate-random-genome-sequence
   [ genome-length>> ] keep
   [ rand-gene> ] curry replicate ;