basis/math/blas/vectors/vectors.factor

   1 USING: accessors alien alien.c-types arrays byte-arrays combinators
   2 combinators.short-circuit fry kernel math math.blas.cblas
   3 math.complex math.functions math.order sequences.complex
   4 sequences.complex-components sequences sequences.private
   5 functors words locals
   6 specialized-arrays.float specialized-arrays.double
   7 specialized-arrays.direct.float specialized-arrays.direct.double ;
   8 IN: math.blas.vectors
   9
  10 TUPLE: blas-vector-base underlying length inc ;
  11
  12 INSTANCE: blas-vector-base virtual-sequence
  13
  14 GENERIC: element-type ( v -- type )
  15
  16 GENERIC: n*V+V! ( alpha x y -- y=alpha*x+y )
  17 GENERIC: n*V!   ( alpha x -- x=alpha*x )
  18 GENERIC: V. ( x y -- x.y )
  19 GENERIC: V.conj ( x y -- xconj.y )
  20 GENERIC: Vnorm ( x -- norm )
  21 GENERIC: Vasum ( x -- sum )
  22 GENERIC: Vswap ( x y -- x=y y=x )
  23 GENERIC: Viamax ( x -- max-i )
  24
  25 <PRIVATE
  26
  27 GENERIC: (blas-vector-like) ( data length inc exemplar -- vector )
  28
  29 GENERIC: (blas-direct-array) ( blas-vector -- direct-array )
  30
  31 : shorter-length ( v1 v2 -- length )
  32     [ length>> ] bi@ min ; inline
  33 : data-and-inc ( v -- data inc )
  34     [ underlying>> ] [ inc>> ] bi ; inline
  35 : datas-and-incs ( v1 v2 -- v1-data v1-inc v2-data v2-inc )
  36     [ data-and-inc ] bi@ ; inline
  37
  38 :: (prepare-copy)
  39     ( v element-size -- length v-data v-inc v-dest-data v-dest-inc
  40                         copy-data copy-length copy-inc )
  41     v [ length>> ] [ data-and-inc ] bi
  42     v length>> element-size * <byte-array>
  43     1
  44     over v length>> 1 ;
  45
  46 : (prepare-swap)
  47     ( v1 v2 -- length v1-data v1-inc v2-data v2-inc
  48                v1 v2 )
  49     [ shorter-length ] [ datas-and-incs ] [ ] 2tri ;
  50
  51 :: (prepare-axpy)
  52     ( n v1 v2 -- length n v1-data v1-inc v2-data v2-inc
  53                  v2 )
  54     v1 v2 shorter-length
  55     n
  56     v1 v2 datas-and-incs
  57     v2 ;
  58
  59 :: (prepare-scal)
  60     ( n v -- length n v-data v-inc
  61              v )
  62     v length>>
  63     n
  64     v data-and-inc
  65     v ;
  66
  67 : (prepare-dot) ( v1 v2 -- length v1-data v1-inc v2-data v2-inc )
  68     [ shorter-length ] [ datas-and-incs ] 2bi ;
  69
  70 : (prepare-nrm2) ( v -- length data inc )
  71     [ length>> ] [ data-and-inc ] bi ;
  72
  73 PRIVATE>
  74
  75 : n*V+V ( alpha x y -- alpha*x+y ) clone n*V+V! ; inline
  76 : n*V ( alpha x -- alpha*x ) clone n*V! ; inline
  77
  78 : V+ ( x y -- x+y )
  79     1.0 -rot n*V+V ; inline
  80 : V- ( x y -- x-y )
  81     -1.0 spin n*V+V ; inline
  82
  83 : Vneg ( x -- -x )
  84     -1.0 swap n*V ; inline
  85
  86 : V*n ( x alpha -- x*alpha )
  87     swap n*V ; inline
  88 : V/n ( x alpha -- x/alpha )
  89     recip swap n*V ; inline
  90
  91 : Vamax ( x -- max )
  92     [ Viamax ] keep nth ; inline
  93
  94 :: Vsub ( v start length -- sub )
  95     v inc>> start * v element-type heap-size *
  96     v underlying>> <displaced-alien>
  97     length v inc>> v (blas-vector-like) ;
  98
  99 : <zero-vector> ( exemplar -- zero )
 100     [ element-type <c-object> ]
 101     [ length>> 0 ]
 102     [ (blas-vector-like) ] tri ;
 103
 104 : <empty-vector> ( length exemplar -- vector )
 105     [ element-type <c-array> ]
 106     [ 1 swap ] 2bi
 107     (blas-vector-like) ;
 108
 109 M: blas-vector-base equal?
 110     {
 111         [ [ length ] bi@ = ]
 112         [ [ = ] 2all? ]
 113     } 2&& ;
 114
 115 M: blas-vector-base length
 116     length>> ;
 117 M: blas-vector-base virtual-seq
 118     (blas-direct-array) ;
 119 M: blas-vector-base virtual@
 120     [ inc>> * ] [ nip (blas-direct-array) ] 2bi ;
 121
 122 : float>arg ( f -- f ) ; inline
 123 : double>arg ( f -- f ) ; inline
 124 : arg>float ( f -- f ) ; inline
 125 : arg>double ( f -- f ) ; inline
 126
 127 <<
 128
 129 FUNCTOR: (define-blas-vector) ( TYPE T -- )
 130
 131 <DIRECT-ARRAY> IS <direct-${TYPE}-array>
 132 >ARRAY         IS >${TYPE}-array
 133 XCOPY          IS cblas_${T}copy
 134 XSWAP          IS cblas_${T}swap
 135 IXAMAX         IS cblas_i${T}amax
 136
 137 VECTOR         DEFINES ${TYPE}-blas-vector
 138 <VECTOR>       DEFINES <${TYPE}-blas-vector>
 139 >VECTOR        DEFINES >${TYPE}-blas-vector
 140
 141 WHERE
 142
 143 TUPLE: VECTOR < blas-vector-base ;
 144 : <VECTOR> ( underlying length inc -- vector ) VECTOR boa ; inline
 145
 146 : >VECTOR ( seq -- v )
 147     [ >ARRAY execute underlying>> ] [ length ] bi 1 <VECTOR> execute ;
 148
 149 M: VECTOR clone
 150     TYPE heap-size (prepare-copy)
 151     [ XCOPY execute ] 3dip <VECTOR> execute ;
 152
 153 M: VECTOR element-type
 154     drop TYPE ;
 155 M: VECTOR Vswap
 156     (prepare-swap) [ XSWAP execute ] 2dip ;
 157 M: VECTOR Viamax
 158     (prepare-nrm2) IXAMAX execute ;
 159
 160 M: VECTOR (blas-vector-like)
 161     drop <VECTOR> execute ;
 162
 163 M: VECTOR (blas-direct-array)
 164     [ underlying>> ]
 165     [ [ length>> ] [ inc>> ] bi * ] bi
 166     <DIRECT-ARRAY> execute ;
 167
 168 ;FUNCTOR
 169
 170
 171 FUNCTOR: (define-real-blas-vector) ( TYPE T -- )
 172
 173 VECTOR         IS ${TYPE}-blas-vector
 174 XDOT           IS cblas_${T}dot
 175 XNRM2          IS cblas_${T}nrm2
 176 XASUM          IS cblas_${T}asum
 177 XAXPY          IS cblas_${T}axpy
 178 XSCAL          IS cblas_${T}scal
 179
 180 WHERE
 181
 182 M: VECTOR V.
 183     (prepare-dot) XDOT execute ;
 184 M: VECTOR V.conj
 185     (prepare-dot) XDOT execute ;
 186 M: VECTOR Vnorm
 187     (prepare-nrm2) XNRM2 execute ;
 188 M: VECTOR Vasum
 189     (prepare-nrm2) XASUM execute ;
 190 M: VECTOR n*V+V!
 191     (prepare-axpy) [ XAXPY execute ] dip ;
 192 M: VECTOR n*V!
 193     (prepare-scal) [ XSCAL execute ] dip ;
 194
 195 ;FUNCTOR
 196
 197
 198 FUNCTOR: (define-complex-helpers) ( TYPE -- )
 199
 200 <DIRECT-COMPLEX-ARRAY> DEFINES <direct-${TYPE}-complex-array>
 201 >COMPLEX-ARRAY         DEFINES >${TYPE}-complex-array
 202 ARG>COMPLEX            DEFINES arg>${TYPE}-complex
 203 COMPLEX>ARG            DEFINES ${TYPE}-complex>arg
 204 <DIRECT-ARRAY>         IS      <direct-${TYPE}-array>
 205 >ARRAY                 IS      >${TYPE}-array
 206
 207 WHERE
 208
 209 : <DIRECT-COMPLEX-ARRAY> ( alien len -- sequence )
 210     1 shift <DIRECT-ARRAY> execute <complex-sequence> ;
 211 : >COMPLEX-ARRAY ( sequence -- sequence )
 212     <complex-components> >ARRAY execute ;
 213 : COMPLEX>ARG ( complex -- alien )
 214     >rect 2array >ARRAY execute underlying>> ;
 215 : ARG>COMPLEX ( alien -- complex )
 216     2 <DIRECT-ARRAY> execute first2 rect> ;
 217
 218 ;FUNCTOR
 219
 220
 221 FUNCTOR: (define-complex-blas-vector) ( TYPE C S -- )
 222
 223 VECTOR         IS ${TYPE}-blas-vector
 224 XDOTU_SUB      IS cblas_${C}dotu_sub
 225 XDOTC_SUB      IS cblas_${C}dotc_sub
 226 XXNRM2         IS cblas_${S}${C}nrm2
 227 XXASUM         IS cblas_${S}${C}asum
 228 XAXPY          IS cblas_${C}axpy
 229 XSCAL          IS cblas_${C}scal
 230 TYPE>ARG       IS ${TYPE}>arg
 231 ARG>TYPE       IS arg>${TYPE}
 232
 233 WHERE
 234
 235 M: VECTOR V.
 236     (prepare-dot) TYPE <c-object>
 237     [ XDOTU_SUB execute ] keep
 238     ARG>TYPE execute ;
 239 M: VECTOR V.conj
 240     (prepare-dot) TYPE <c-object>
 241     [ XDOTC_SUB execute ] keep
 242     ARG>TYPE execute ;
 243 M: VECTOR Vnorm
 244     (prepare-nrm2) XXNRM2 execute ;
 245 M: VECTOR Vasum
 246     (prepare-nrm2) XXASUM execute ;
 247 M: VECTOR n*V+V!
 248     [ TYPE>ARG execute ] 2dip
 249     (prepare-axpy) [ XAXPY execute ] dip ;
 250 M: VECTOR n*V!
 251     [ TYPE>ARG execute ] dip
 252     (prepare-scal) [ XSCAL execute ] dip ;
 253
 254 ;FUNCTOR
 255
 256
 257 : define-real-blas-vector ( TYPE T -- )
 258     [ (define-blas-vector) ]
 259     [ (define-real-blas-vector) ] 2bi ;
 260 :: define-complex-blas-vector ( TYPE C S -- )
 261     TYPE (define-complex-helpers)
 262     TYPE "-complex" append
 263     [ C (define-blas-vector) ]
 264     [ C S (define-complex-blas-vector) ] bi ;
 265
 266 "float"  "s" define-real-blas-vector
 267 "double" "d" define-real-blas-vector
 268 "float"  "c" "s" define-complex-blas-vector
 269 "double" "z" "d" define-complex-blas-vector
 270
 271 >>
 272