extra/machine-learning/rebalancing/rebalancing.factor

   1 ! Copyright (C) 2012 Doug Coleman.
   2 ! See https://factorcode.org/license.txt for BSD license.
   3 USING: arrays assocs kernel math math.functions math.statistics
   4 random sequences sorting ;
   5 IN: machine-learning.rebalancing
   6
   7 ERROR: probability-sum-not-one seq ;
   8
   9 : check-probabilities ( seq -- seq )
  10     dup sum 1.0 .00000000001 ~ [ probability-sum-not-one ] unless ;
  11
  12 : equal-probabilities ( n -- array )
  13     dup recip <array> ; inline
  14
  15 MEMO: probabilities-seq ( seq -- seq' )
  16     check-probabilities [ >float ] map cum-sum ;
  17
  18 : probabilities-quot ( seq -- quot )
  19     probabilities-seq
  20     '[ _ random-unit '[ _ > ] find drop ] ; inline
  21
  22 : stratified-sample ( stratified-sequences probability-sequence -- elt )
  23     probabilities-quot call swap nth random ; inline
  24
  25 : stratified-samples ( stratified-sequences probability-sequence n -- elt )
  26     [ '[ _ _ stratified-sample ] ] dip swap replicate ;
  27
  28 : equal-stratified-sample ( stratified-sequences -- elt )
  29     random random ; inline
  30
  31 : collect-indices ( seq -- indices )
  32     H{ } clone [ '[ swap _ push-at ] each-index ] keep ;
  33
  34 : balance-labels ( X y n -- X' y' )
  35     [
  36         dup collect-indices
  37         values '[
  38             _ _ _ equal-stratified-sample
  39             '[ _ swap nth ] bi@ 2array
  40         ]
  41     ] dip swap replicate [ keys ] [ values ] bi ;
  42
  43 : skew-labels ( X y probs n -- X' y' )
  44     [
  45         [ dup collect-indices sort-keys values ] dip
  46         '[
  47             _ _ _ _ stratified-sample
  48             '[ _ swap nth ] bi@ 2array
  49         ]
  50     ] dip swap replicate [ keys ] [ values ] bi ;