1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
|
{-# LANGUAGE TypeOperators, ExplicitForAll, FlexibleContexts #-}
module Data.Array.Repa.Operators.Interleave
( interleave2
, interleave3
, interleave4)
where
import Data.Array.Repa.Shape
import Data.Array.Repa.Index
import Data.Array.Repa.Base
import Data.Array.Repa.Repr.Delayed
import Data.Array.Repa.Operators.Traversal
import Prelude hiding ((++))
-- Interleave -----------------------------------------------------------------
-- | Interleave the elements of two arrays.
-- All the input arrays must have the same extent, else `error`.
-- The lowest dimension of the result array is twice the size of the inputs.
--
-- @
-- interleave2 a1 a2 b1 b2 => a1 b1 a2 b2
-- a3 a4 b3 b4 a3 b3 a4 b4
-- @
--
interleave2
:: ( Shape sh
, Source r1 a, Source r2 a)
=> Array r1 (sh :. Int) a
-> Array r2 (sh :. Int) a
-> Array D (sh :. Int) a
{-# INLINE [2] interleave2 #-}
interleave2 arr1 arr2
= unsafeTraverse2 arr1 arr2 shapeFn elemFn
where
shapeFn dim1 dim2
| dim1 == dim2
, sh :. len <- dim1
= sh :. (len * 2)
| otherwise
= error "Data.Array.Repa.interleave2: arrays must have same extent"
elemFn get1 get2 (sh :. ix)
= case ix `mod` 2 of
0 -> get1 (sh :. ix `div` 2)
1 -> get2 (sh :. ix `div` 2)
_ -> error "Data.Array.Repa.interleave2: this never happens :-P"
-- | Interleave the elements of three arrays.
interleave3
:: ( Shape sh
, Source r1 a, Source r2 a, Source r3 a)
=> Array r1 (sh :. Int) a
-> Array r2 (sh :. Int) a
-> Array r3 (sh :. Int) a
-> Array D (sh :. Int) a
{-# INLINE [2] interleave3 #-}
interleave3 arr1 arr2 arr3
= unsafeTraverse3 arr1 arr2 arr3 shapeFn elemFn
where
shapeFn dim1 dim2 dim3
| dim1 == dim2
, dim1 == dim3
, sh :. len <- dim1
= sh :. (len * 3)
| otherwise
= error "Data.Array.Repa.interleave3: arrays must have same extent"
elemFn get1 get2 get3 (sh :. ix)
= case ix `mod` 3 of
0 -> get1 (sh :. ix `div` 3)
1 -> get2 (sh :. ix `div` 3)
2 -> get3 (sh :. ix `div` 3)
_ -> error "Data.Array.Repa.interleave3: this never happens :-P"
-- | Interleave the elements of four arrays.
interleave4
:: ( Shape sh
, Source r1 a, Source r2 a, Source r3 a, Source r4 a)
=> Array r1 (sh :. Int) a
-> Array r2 (sh :. Int) a
-> Array r3 (sh :. Int) a
-> Array r4 (sh :. Int) a
-> Array D (sh :. Int) a
{-# INLINE [2] interleave4 #-}
interleave4 arr1 arr2 arr3 arr4
= unsafeTraverse4 arr1 arr2 arr3 arr4 shapeFn elemFn
where
shapeFn dim1 dim2 dim3 dim4
| dim1 == dim2
, dim1 == dim3
, dim1 == dim4
, sh :. len <- dim1
= sh :. (len * 4)
| otherwise
= error "Data.Array.Repa.interleave4: arrays must have same extent"
elemFn get1 get2 get3 get4 (sh :. ix)
= case ix `mod` 4 of
0 -> get1 (sh :. ix `div` 4)
1 -> get2 (sh :. ix `div` 4)
2 -> get3 (sh :. ix `div` 4)
3 -> get4 (sh :. ix `div` 4)
_ -> error "Data.Array.Repa.interleave4: this never happens :-P"
|
