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 parser prettyprint.backend prettyprint.custom
   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 XVECTOR{       DEFINES ${T}vector{
 142
 143 WHERE
 144
 145 TUPLE: VECTOR < blas-vector-base ;
 146 : <VECTOR> ( underlying length inc -- vector ) VECTOR boa ; inline
 147
 148 : >VECTOR ( seq -- v )
 149     [ >ARRAY underlying>> ] [ length ] bi 1 <VECTOR> ;
 150
 151 M: VECTOR clone
 152     TYPE heap-size (prepare-copy)
 153     [ XCOPY ] 3dip <VECTOR> ;
 154
 155 M: VECTOR element-type
 156     drop TYPE ;
 157 M: VECTOR Vswap
 158     (prepare-swap) [ XSWAP ] 2dip ;
 159 M: VECTOR Viamax
 160     (prepare-nrm2) IXAMAX ;
 161
 162 M: VECTOR (blas-vector-like)
 163     drop <VECTOR> ;
 164
 165 M: VECTOR (blas-direct-array)
 166     [ underlying>> ]
 167     [ [ length>> ] [ inc>> ] bi * ] bi
 168     <DIRECT-ARRAY> ;
 169
 170 : XVECTOR{ \ } [ >VECTOR ] parse-literal ; parsing
 171
 172 M: VECTOR pprint-delims
 173     drop \ XVECTOR{ \ } ;
 174
 175 ;FUNCTOR
 176
 177
 178 FUNCTOR: (define-real-blas-vector) ( TYPE T -- )
 179
 180 VECTOR         IS ${TYPE}-blas-vector
 181 XDOT           IS cblas_${T}dot
 182 XNRM2          IS cblas_${T}nrm2
 183 XASUM          IS cblas_${T}asum
 184 XAXPY          IS cblas_${T}axpy
 185 XSCAL          IS cblas_${T}scal
 186
 187 WHERE
 188
 189 M: VECTOR V.
 190     (prepare-dot) XDOT ;
 191 M: VECTOR V.conj
 192     (prepare-dot) XDOT ;
 193 M: VECTOR Vnorm
 194     (prepare-nrm2) XNRM2 ;
 195 M: VECTOR Vasum
 196     (prepare-nrm2) XASUM ;
 197 M: VECTOR n*V+V!
 198     (prepare-axpy) [ XAXPY ] dip ;
 199 M: VECTOR n*V!
 200     (prepare-scal) [ XSCAL ] dip ;
 201
 202 ;FUNCTOR
 203
 204
 205 FUNCTOR: (define-complex-helpers) ( TYPE -- )
 206
 207 <DIRECT-COMPLEX-ARRAY> DEFINES <direct-${TYPE}-complex-array>
 208 >COMPLEX-ARRAY         DEFINES >${TYPE}-complex-array
 209 ARG>COMPLEX            DEFINES arg>${TYPE}-complex
 210 COMPLEX>ARG            DEFINES ${TYPE}-complex>arg
 211 <DIRECT-ARRAY>         IS      <direct-${TYPE}-array>
 212 >ARRAY                 IS      >${TYPE}-array
 213
 214 WHERE
 215
 216 : <DIRECT-COMPLEX-ARRAY> ( alien len -- sequence )
 217     1 shift <DIRECT-ARRAY> <complex-sequence> ;
 218 : >COMPLEX-ARRAY ( sequence -- sequence )
 219     <complex-components> >ARRAY ;
 220 : COMPLEX>ARG ( complex -- alien )
 221     >rect 2array >ARRAY underlying>> ;
 222 : ARG>COMPLEX ( alien -- complex )
 223     2 <DIRECT-ARRAY> first2 rect> ;
 224
 225 ;FUNCTOR
 226
 227
 228 FUNCTOR: (define-complex-blas-vector) ( TYPE C S -- )
 229
 230 VECTOR         IS ${TYPE}-blas-vector
 231 XDOTU_SUB      IS cblas_${C}dotu_sub
 232 XDOTC_SUB      IS cblas_${C}dotc_sub
 233 XXNRM2         IS cblas_${S}${C}nrm2
 234 XXASUM         IS cblas_${S}${C}asum
 235 XAXPY          IS cblas_${C}axpy
 236 XSCAL          IS cblas_${C}scal
 237 TYPE>ARG       IS ${TYPE}>arg
 238 ARG>TYPE       IS arg>${TYPE}
 239
 240 WHERE
 241
 242 M: VECTOR V.
 243     (prepare-dot) TYPE <c-object>
 244     [ XDOTU_SUB ] keep
 245     ARG>TYPE ;
 246 M: VECTOR V.conj
 247     (prepare-dot) TYPE <c-object>
 248     [ XDOTC_SUB ] keep
 249     ARG>TYPE ;
 250 M: VECTOR Vnorm
 251     (prepare-nrm2) XXNRM2 ;
 252 M: VECTOR Vasum
 253     (prepare-nrm2) XXASUM ;
 254 M: VECTOR n*V+V!
 255     [ TYPE>ARG ] 2dip
 256     (prepare-axpy) [ XAXPY ] dip ;
 257 M: VECTOR n*V!
 258     [ TYPE>ARG ] dip
 259     (prepare-scal) [ XSCAL ] dip ;
 260
 261 ;FUNCTOR
 262
 263
 264 : define-real-blas-vector ( TYPE T -- )
 265     [ (define-blas-vector) ]
 266     [ (define-real-blas-vector) ] 2bi ;
 267 :: define-complex-blas-vector ( TYPE C S -- )
 268     TYPE (define-complex-helpers)
 269     TYPE "-complex" append
 270     [ C (define-blas-vector) ]
 271     [ C S (define-complex-blas-vector) ] bi ;
 272
 273 "float"  "s" define-real-blas-vector
 274 "double" "d" define-real-blas-vector
 275 "float"  "c" "s" define-complex-blas-vector
 276 "double" "z" "d" define-complex-blas-vector
 277
 278 >>
 279
 280 M: blas-vector-base >pprint-sequence ;
 281 M: blas-vector-base pprint* pprint-object ;