| 447 | {{{ |
| 448 | redBlack:: Array.Shape dim => Double -> Double -> DArray (dim :*: Int :*: Int :*: Int) Double -> |
| 449 | DArray (dim :*: Int :*: Int :*: Int) Double -> DArray (dim :*: Int :*: Int :*: Int) Double |
| 450 | redBlack factor hsq f arr@(DArray (d :*: l :*: n :*:m) fn) = |
| 451 | applyFactor $ insert odd arr' $ sumD $ getBlack $ stencil arr' |
| 452 | |
| 453 | where |
| 454 | arr' = applyFactor $ insert even arr $ sumD $ getRed $ stencil arr |
| 455 | |
| 456 | applyFactor = zipWith (\fi -> \si -> factor * (hsq * fi + si)) f |
| 457 | |
| 458 | sumD arr = fold (+) 0 arr |
| 459 | |
| 460 | getRed (DArray (sh :*: l :*: m :*: n :*: c) f ) = |
| 461 | DArray (sh :*: l-2 :*: m-2 :*: (n-1)`div` 2 :*: c) |
| 462 | (\(sh :*: h :*: i :*: j :*: c) -> f(sh :*: h+1 :*: i+1 :*: 2*j+1 :*: c)) |
| 463 | getBlack (DArray (sh :*: l :*: m :*: n :*: c) f) = |
| 464 | DArray (sh :*: l-2 :*: m-2 :*: (n-2) `div` 2:*: c) |
| 465 | (\(sh :*: h :*: i :*: j:*:c) -> f (sh :*: h+1 :*: i+1 :*: 2*j+2:*:c)) |
| 466 | |
| 467 | isBorder (d :*: h :*: i :*: j) = ((h * i * j) == 0) || |
| 468 | (h >= (l-1)) || (i >= (m-1)) || (j >= (n-1)) |
| 469 | |
| 470 | insert p (DArray sh f) (DArray sh' f') = |
| 471 | DArray sh (\d@(sh :*: h :*: i :*: j) -> if ((isBorder d) || p j) |
| 472 | then f d |
| 473 | else (f' (sh :*: h-1 :*: i-1 :*: (j-1)`div` 2))) |
| 474 | |
| 475 | stencil (DArray sh f) = |
| 476 | DArray (sh :*: 6) f' |
| 477 | where |
| 478 | f' (sh :*: n :*: m :*: 0) = f (sh :*: n :*: m+1) |
| 479 | f' (sh :*: n :*: m :*: 1) = f (sh :*: n :*: m-1) |
| 480 | f' (sh :*: n :*: m :*: 2) = f (sh :*: n+1 :*: m) |
| 481 | f' (sh :*: n :*: m :*: 3) = f (sh :*: n-1 :*: m) |
| 482 | f' (sh :*: k :*: n :*: m :*: 4) = f (sh :*: k+1 :*: n :*: m) |
| 483 | f' (sh :*: k :*: n :*: m :*: 5) = f (sh :*: k-1 :*: n :*: m) |
| 484 | |
| 485 | }}} |
| 486 | |
448 | | |
449 | | |
450 | | |
| 488 | Applied FFT to each vector in a three dimensional matrix, once along each of the three axes, iterate a given number of times. |
| 489 | |
| 490 | {{{ |
| 491 | fft3d:: Int -> DArray Array.DIM3 Complex -> DArray Array.DIM3 Complex -> DArray Array.DIM3 Complex |
| 492 | fft3d it rofu m | it < 1 = m |
| 493 | | otherwise = fft3d (it-1) rofu $ fftTrans $ fftTrans $ fftTrans m |
| 494 | where |
| 495 | fftTrans = forceDArray . (fft rofu) . transpose' |
| 496 | transpose' darr@(DArray (() :*: k :*: l :*: m) _) = |
| 497 | backpermute darr (() :*: m :*: k :*: l) |
| 498 | (\(() :*: m' :*: k' :*: l') -> (() :*: k' :*: l' :*: m')) |
| 499 | |
| 500 | fft:: Array.Subshape dim dim => |
| 501 | DArray (dim :*: Int) Complex -> DArray (dim :*: Int) Complex -> DArray (dim :*: Int) Complex |
| 502 | fft rofu@(DArray ( _ :*: s) _ ) v@(DArray sh@(_ :*: n) f) |
| 503 | | n > 2 = assert (2 * s == n) $ |
| 504 | append (fft_left + fft_right) (fft_left - fft_right) sh |
| 505 | | n == 2 = assert (2 * s == n) $ |
| 506 | DArray sh f' |
| 507 | where |
| 508 | f' (sh :*: 0) = f (sh :*: 0) + f (sh :*: 1) |
| 509 | f' (sh :*: 1) = f (sh :*: 0) - f (sh :*: 1) |
| 510 | f' (sh :*: x) = error ("error in fft - f:" ++ (show x) ++ "/" ++ (show sh)) |
| 511 | |
| 512 | rofu' = split rofu (\(sh :*: i) -> (sh :*: 2*i)) |
| 513 | fft_left = forceDArray $ rofu * (fft rofu' (split v (\(sh:*: i) -> (sh :*: 2*i)))) |
| 514 | fft_right = forceDArray $ fft rofu' (split v (\(sh:*: i) -> (sh :*: 2*i+1))) |
| 515 | |
| 516 | split:: Array.Shape dim => |
| 517 | DArray (dim :*: Int) Complex -> ((dim :*: Int) -> (dim :*: Int)) -> DArray (dim :*: Int) Complex |
| 518 | split arr@(DArray (sh :*: n) fn) sel = |
| 519 | (DArray (sh :*: (n `div` 2)) (fn . sel)) |
| 520 | |
| 521 | }}} |
| 522 | |