basis/math/blas/vectors/vectors.factor

   1 USING: accessors alien alien.c-types arrays byte-arrays combinators
   2 combinators.short-circuit fry kernel macros math math.blas.cblas
   3 math.complex math.functions math.order multi-methods qualified
   4 sequences sequences.private generalizations
   5 specialized-arrays.float specialized-arrays.double
   6 specialized-arrays.direct.float specialized-arrays.direct.double ;
   7 QUALIFIED: syntax
   8 IN: math.blas.vectors
   9
  10 TUPLE: blas-vector-base data length inc ;
  11 TUPLE: float-blas-vector < blas-vector-base ;
  12 TUPLE: double-blas-vector < blas-vector-base ;
  13 TUPLE: float-complex-blas-vector < blas-vector-base ;
  14 TUPLE: double-complex-blas-vector < blas-vector-base ;
  15
  16 INSTANCE: float-blas-vector sequence
  17 INSTANCE: double-blas-vector sequence
  18 INSTANCE: float-complex-blas-vector sequence
  19 INSTANCE: double-complex-blas-vector sequence
  20
  21 C: <float-blas-vector> float-blas-vector
  22 C: <double-blas-vector> double-blas-vector
  23 C: <float-complex-blas-vector> float-complex-blas-vector
  24 C: <double-complex-blas-vector> double-complex-blas-vector
  25
  26 GENERIC: n*V+V! ( alpha x y -- y=alpha*x+y )
  27 GENERIC: n*V!   ( alpha x -- x=alpha*x )
  28
  29 GENERIC: V. ( x y -- x.y )
  30 GENERIC: V.conj ( x y -- xconj.y )
  31 GENERIC: Vnorm ( x -- norm )
  32 GENERIC: Vasum ( x -- sum )
  33 GENERIC: Vswap ( x y -- x=y y=x )
  34
  35 GENERIC: Viamax ( x -- max-i )
  36
  37 GENERIC: element-type ( v -- type )
  38
  39 METHOD: element-type { float-blas-vector }
  40     drop "float" ;
  41 METHOD: element-type { double-blas-vector }
  42     drop "double" ;
  43 METHOD: element-type { float-complex-blas-vector }
  44     drop "CBLAS_C" ;
  45 METHOD: element-type { double-complex-blas-vector }
  46     drop "CBLAS_Z" ;
  47
  48 <PRIVATE
  49
  50 GENERIC: (blas-vector-like) ( data length inc exemplar -- vector )
  51
  52 METHOD: (blas-vector-like) { object object object float-blas-vector }
  53     drop <float-blas-vector> ;
  54 METHOD: (blas-vector-like) { object object object double-blas-vector }
  55     drop <double-blas-vector> ;
  56 METHOD: (blas-vector-like) { object object object float-complex-blas-vector }
  57     drop <float-complex-blas-vector> ;
  58 METHOD: (blas-vector-like) { object object object double-complex-blas-vector }
  59     drop <double-complex-blas-vector> ;
  60
  61 : (prepare-copy) ( v element-size -- length v-data v-inc v-dest-data v-dest-inc )
  62     [ [ length>> ] [ data>> ] [ inc>> ] tri ] dip
  63     4 npick * <byte-array>
  64     1 ;
  65
  66 MACRO: (do-copy) ( copy make-vector -- )
  67     '[ over 6 npick _ 2dip 1 @ ] ;
  68
  69 : (prepare-swap) ( v1 v2 -- length v1-data v1-inc v2-data v2-inc v1 v2 )
  70     [
  71         [ [ length>> ] bi@ min ]
  72         [ [ [ data>> ] [ inc>> ] bi ] bi@ ] 2bi
  73     ] 2keep ;
  74
  75 : (prepare-axpy) ( n v1 v2 -- length n v1-data v1-inc v2-data v2-inc v2 )
  76     [
  77         [ [ length>> ] bi@ min swap ]
  78         [ [ [ data>> ] [ inc>> ] bi ] bi@ ] 2bi
  79     ] keep ;
  80
  81 : (prepare-scal) ( n v -- length n v-data v-inc v )
  82     [ [ length>> swap ] [ data>> ] [ inc>> ] tri ] keep ;
  83
  84 : (prepare-dot) ( v1 v2 -- length v1-data v1-inc v2-data v2-inc )
  85     [ [ length>> ] bi@ min ]
  86     [ [ [ data>> ] [ inc>> ] bi ] bi@ ] 2bi ;
  87
  88 : (prepare-nrm2) ( v -- length v1-data v1-inc )
  89     [ length>> ] [ data>> ] [ inc>> ] tri ;
  90
  91 : (flatten-complex-sequence) ( seq -- seq' )
  92     [ [ real-part ] [ imaginary-part ] bi 2array ] map concat ;
  93
  94 : (>c-complex) ( complex -- alien )
  95     [ real-part ] [ imaginary-part ] bi float-array{ } 2sequence underlying>> ;
  96 : (>z-complex) ( complex -- alien )
  97     [ real-part ] [ imaginary-part ] bi double-array{ } 2sequence underlying>> ;
  98
  99 : (c-complex>) ( alien -- complex )
 100     2 <direct-float-array> first2 rect> ;
 101 : (z-complex>) ( alien -- complex )
 102     2 <direct-double-array> first2 rect> ;
 103
 104 : (prepare-nth) ( n v -- n*inc v-data )
 105     [ inc>> ] [ data>> ] bi [ * ] dip ;
 106
 107 MACRO: (complex-nth) ( nth-quot -- )
 108     '[
 109         [ 2 * dup 1+ ] dip
 110         _ curry bi@ rect>
 111     ] ;
 112
 113 : (c-complex-nth) ( n alien -- complex )
 114     [ float-nth ] (complex-nth) ;
 115 : (z-complex-nth) ( n alien -- complex )
 116     [ double-nth ] (complex-nth) ;
 117
 118 MACRO: (set-complex-nth) ( set-nth-quot -- )
 119     '[
 120         [
 121             [ [ real-part ] [ imaginary-part ] bi ]
 122             [ 2 * dup 1+ ] bi*
 123             swapd
 124         ] dip
 125         _ curry 2bi@
 126     ] ;
 127
 128 : (set-c-complex-nth) ( complex n alien -- )
 129     [ set-float-nth ] (set-complex-nth) ;
 130 : (set-z-complex-nth) ( complex n alien -- )
 131     [ set-double-nth ] (set-complex-nth) ;
 132
 133 PRIVATE>
 134
 135 : <zero-vector> ( exemplar -- zero )
 136     [ element-type <c-object> ]
 137     [ length>> 0 ]
 138     [ (blas-vector-like) ] tri ;
 139
 140 : <empty-vector> ( length exemplar -- vector )
 141     [ element-type <c-array> ]
 142     [ 1 swap ] 2bi
 143     (blas-vector-like) ;
 144
 145 syntax:M: blas-vector-base length
 146     length>> ;
 147
 148 syntax:M: float-blas-vector nth-unsafe
 149     (prepare-nth) float-nth ;
 150 syntax:M: float-blas-vector set-nth-unsafe
 151     (prepare-nth) set-float-nth ;
 152
 153 syntax:M: double-blas-vector nth-unsafe
 154     (prepare-nth) double-nth ;
 155 syntax:M: double-blas-vector set-nth-unsafe
 156     (prepare-nth) set-double-nth ;
 157
 158 syntax:M: float-complex-blas-vector nth-unsafe
 159     (prepare-nth) (c-complex-nth) ;
 160 syntax:M: float-complex-blas-vector set-nth-unsafe
 161     (prepare-nth) (set-c-complex-nth) ;
 162
 163 syntax:M: double-complex-blas-vector nth-unsafe
 164     (prepare-nth) (z-complex-nth) ;
 165 syntax:M: double-complex-blas-vector set-nth-unsafe
 166     (prepare-nth) (set-z-complex-nth) ;
 167
 168 syntax:M: blas-vector-base equal?
 169     {
 170         [ [ length ] bi@ = ]
 171         [ [ = ] 2all? ]
 172     } 2&& ;
 173
 174 : >float-blas-vector ( seq -- v )
 175     [ >float-array underlying>> ] [ length ] bi 1 <float-blas-vector> ;
 176 : >double-blas-vector ( seq -- v )
 177     [ >double-array underlying>> ] [ length ] bi 1 <double-blas-vector> ;
 178 : >float-complex-blas-vector ( seq -- v )
 179     [ (flatten-complex-sequence) >float-array underlying>> ] [ length ] bi
 180     1 <float-complex-blas-vector> ;
 181 : >double-complex-blas-vector ( seq -- v )
 182     [ (flatten-complex-sequence) >double-array underlying>> ] [ length ] bi
 183     1 <double-complex-blas-vector> ;
 184
 185 syntax:M: float-blas-vector clone
 186     "float" heap-size (prepare-copy)
 187     [ cblas_scopy ] [ <float-blas-vector> ] (do-copy) ;
 188 syntax:M: double-blas-vector clone
 189     "double" heap-size (prepare-copy)
 190     [ cblas_dcopy ] [ <double-blas-vector> ] (do-copy) ;
 191 syntax:M: float-complex-blas-vector clone
 192     "CBLAS_C" heap-size (prepare-copy)
 193     [ cblas_ccopy ] [ <float-complex-blas-vector> ] (do-copy) ;
 194 syntax:M: double-complex-blas-vector clone
 195     "CBLAS_Z" heap-size (prepare-copy)
 196     [ cblas_zcopy ] [ <double-complex-blas-vector> ] (do-copy) ;
 197
 198 METHOD: Vswap { float-blas-vector float-blas-vector }
 199     (prepare-swap) [ cblas_sswap ] 2dip ;
 200 METHOD: Vswap { double-blas-vector double-blas-vector }
 201     (prepare-swap) [ cblas_dswap ] 2dip ;
 202 METHOD: Vswap { float-complex-blas-vector float-complex-blas-vector }
 203     (prepare-swap) [ cblas_cswap ] 2dip ;
 204 METHOD: Vswap { double-complex-blas-vector double-complex-blas-vector }
 205     (prepare-swap) [ cblas_zswap ] 2dip ;
 206
 207 METHOD: n*V+V! { real float-blas-vector float-blas-vector }
 208     (prepare-axpy) [ cblas_saxpy ] dip ;
 209 METHOD: n*V+V! { real double-blas-vector double-blas-vector }
 210     (prepare-axpy) [ cblas_daxpy ] dip ;
 211 METHOD: n*V+V! { number float-complex-blas-vector float-complex-blas-vector }
 212     [ (>c-complex) ] 2dip
 213     (prepare-axpy) [ cblas_caxpy ] dip ;
 214 METHOD: n*V+V! { number double-complex-blas-vector double-complex-blas-vector }
 215     [ (>z-complex) ] 2dip
 216     (prepare-axpy) [ cblas_zaxpy ] dip ;
 217
 218 METHOD: n*V! { real float-blas-vector }
 219     (prepare-scal) [ cblas_sscal ] dip ;
 220 METHOD: n*V! { real double-blas-vector }
 221     (prepare-scal) [ cblas_dscal ] dip ;
 222 METHOD: n*V! { number float-complex-blas-vector }
 223     [ (>c-complex) ] dip
 224     (prepare-scal) [ cblas_cscal ] dip ;
 225 METHOD: n*V! { number double-complex-blas-vector }
 226     [ (>z-complex) ] dip
 227     (prepare-scal) [ cblas_zscal ] dip ;
 228
 229 : n*V+V ( alpha x y -- alpha*x+y ) clone n*V+V! ; inline
 230 : n*V ( alpha x -- alpha*x ) clone n*V! ; inline
 231
 232 : V+ ( x y -- x+y )
 233     1.0 -rot n*V+V ; inline
 234 : V- ( x y -- x-y )
 235     -1.0 spin n*V+V ; inline
 236
 237 : Vneg ( x -- -x )
 238     -1.0 swap n*V ; inline
 239
 240 : V*n ( x alpha -- x*alpha )
 241     swap n*V ; inline
 242 : V/n ( x alpha -- x/alpha )
 243     recip swap n*V ; inline
 244
 245 METHOD: V. { float-blas-vector float-blas-vector }
 246     (prepare-dot) cblas_sdot ;
 247 METHOD: V. { double-blas-vector double-blas-vector }
 248     (prepare-dot) cblas_ddot ;
 249 METHOD: V. { float-complex-blas-vector float-complex-blas-vector }
 250     (prepare-dot)
 251     "CBLAS_C" <c-object> [ cblas_cdotu_sub ] keep (c-complex>) ;
 252 METHOD: V. { double-complex-blas-vector double-complex-blas-vector }
 253     (prepare-dot)
 254     "CBLAS_Z" <c-object> [ cblas_zdotu_sub ] keep (z-complex>) ;
 255
 256 METHOD: V.conj { float-blas-vector float-blas-vector }
 257     (prepare-dot) cblas_sdot ;
 258 METHOD: V.conj { double-blas-vector double-blas-vector }
 259     (prepare-dot) cblas_ddot ;
 260 METHOD: V.conj { float-complex-blas-vector float-complex-blas-vector }
 261     (prepare-dot)
 262     "CBLAS_C" <c-object> [ cblas_cdotc_sub ] keep (c-complex>) ;
 263 METHOD: V.conj { double-complex-blas-vector double-complex-blas-vector }
 264     (prepare-dot)
 265     "CBLAS_Z" <c-object> [ cblas_zdotc_sub ] keep (z-complex>) ;
 266
 267 METHOD: Vnorm { float-blas-vector }
 268     (prepare-nrm2) cblas_snrm2 ;
 269 METHOD: Vnorm { double-blas-vector }
 270     (prepare-nrm2) cblas_dnrm2 ;
 271 METHOD: Vnorm { float-complex-blas-vector }
 272     (prepare-nrm2) cblas_scnrm2 ;
 273 METHOD: Vnorm { double-complex-blas-vector }
 274     (prepare-nrm2) cblas_dznrm2 ;
 275
 276 METHOD: Vasum { float-blas-vector }
 277     (prepare-nrm2) cblas_sasum ;
 278 METHOD: Vasum { double-blas-vector }
 279     (prepare-nrm2) cblas_dasum ;
 280 METHOD: Vasum { float-complex-blas-vector }
 281     (prepare-nrm2) cblas_scasum ;
 282 METHOD: Vasum { double-complex-blas-vector }
 283     (prepare-nrm2) cblas_dzasum ;
 284
 285 METHOD: Viamax { float-blas-vector }
 286     (prepare-nrm2) cblas_isamax ;
 287 METHOD: Viamax { double-blas-vector }
 288     (prepare-nrm2) cblas_idamax ;
 289 METHOD: Viamax { float-complex-blas-vector }
 290     (prepare-nrm2) cblas_icamax ;
 291 METHOD: Viamax { double-complex-blas-vector }
 292     (prepare-nrm2) cblas_izamax ;
 293
 294 : Vamax ( x -- max )
 295     [ Viamax ] keep nth ; inline
 296
 297 : Vsub ( v start length -- sub )
 298     rot [
 299         [
 300             nip [ inc>> ] [ element-type heap-size ] [ data>> ] tri
 301             [ * * ] dip <displaced-alien>
 302         ] [ swap 2nip ] [ 2nip inc>> ] 3tri
 303     ] keep (blas-vector-like) ;