feectools 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (98) hide show
  1. feectools/__init__.py +0 -0
  2. feectools/accelerate/__init__.py +0 -0
  3. feectools/accelerate/accelerate.py +220 -0
  4. feectools/accelerate/compile_psydac.mk +52 -0
  5. feectools/api/__init__.py +0 -0
  6. feectools/api/essential_bc.py +122 -0
  7. feectools/api/fem_bilinear_form.py +2226 -0
  8. feectools/api/fem_common.py +286 -0
  9. feectools/api/fem_sum_form.py +123 -0
  10. feectools/api/settings.py +82 -0
  11. feectools/core/__init__.py +11 -0
  12. feectools/core/bsplines.py +1107 -0
  13. feectools/core/bsplines_kernels.py +1349 -0
  14. feectools/core/field_evaluation_kernels.py +5015 -0
  15. feectools/core/tests/__init__.py +0 -0
  16. feectools/core/tests/test_bsplines.py +263 -0
  17. feectools/core/tests/test_bsplines_kernel.py +40 -0
  18. feectools/core/tests/test_bsplines_pyccel.py +752 -0
  19. feectools/ddm/__init__.py +3 -0
  20. feectools/ddm/basic.py +78 -0
  21. feectools/ddm/blocking_data_exchanger.py +348 -0
  22. feectools/ddm/cart.py +1835 -0
  23. feectools/ddm/interface_data_exchanger.py +122 -0
  24. feectools/ddm/mpi.py +109 -0
  25. feectools/ddm/nonblocking_data_exchanger.py +331 -0
  26. feectools/ddm/partition.py +207 -0
  27. feectools/ddm/petsc.py +112 -0
  28. feectools/ddm/tests/__init__.py +0 -0
  29. feectools/ddm/tests/test_cart_1d.py +138 -0
  30. feectools/ddm/tests/test_cart_2d.py +164 -0
  31. feectools/ddm/tests/test_cart_3d.py +158 -0
  32. feectools/ddm/tests/test_multicart_2d.py +173 -0
  33. feectools/ddm/tests/test_partition.py +124 -0
  34. feectools/ddm/utilities.py +24 -0
  35. feectools/feec/__init__.py +0 -0
  36. feectools/feec/derivatives.py +780 -0
  37. feectools/feec/dof_kernels.py +210 -0
  38. feectools/feec/global_geometric_projectors.py +1073 -0
  39. feectools/feec/hodge.py +148 -0
  40. feectools/fem/__init__.py +0 -0
  41. feectools/fem/basic.py +465 -0
  42. feectools/fem/grid.py +181 -0
  43. feectools/fem/partitioning.py +344 -0
  44. feectools/fem/projectors.py +160 -0
  45. feectools/fem/splines.py +559 -0
  46. feectools/fem/tensor.py +1393 -0
  47. feectools/fem/tests/__init__.py +0 -0
  48. feectools/fem/tests/analytical_profiles_1d.py +100 -0
  49. feectools/fem/tests/analytical_profiles_base.py +34 -0
  50. feectools/fem/tests/splines_error_bounds.py +155 -0
  51. feectools/fem/tests/test_spline_histopolation.py +120 -0
  52. feectools/fem/tests/test_spline_interpolation.py +182 -0
  53. feectools/fem/tests/test_splines.py +184 -0
  54. feectools/fem/tests/test_splines_par.py +46 -0
  55. feectools/fem/tests/test_vector_spaces.py +150 -0
  56. feectools/fem/tests/utilities.py +47 -0
  57. feectools/fem/vector.py +729 -0
  58. feectools/linalg/__init__.py +0 -0
  59. feectools/linalg/basic.py +1386 -0
  60. feectools/linalg/block.py +1451 -0
  61. feectools/linalg/direct_solvers.py +201 -0
  62. feectools/linalg/fft.py +258 -0
  63. feectools/linalg/kernels/__init__.py +0 -0
  64. feectools/linalg/kernels/axpy_kernels.py +57 -0
  65. feectools/linalg/kernels/inner_kernels.py +100 -0
  66. feectools/linalg/kernels/matvec_kernels.py +206 -0
  67. feectools/linalg/kernels/stencil2IJV_kernels.py +227 -0
  68. feectools/linalg/kernels/stencil2coo_kernels.py +179 -0
  69. feectools/linalg/kernels/transpose_kernels.py +263 -0
  70. feectools/linalg/kron.py +911 -0
  71. feectools/linalg/solvers.py +1914 -0
  72. feectools/linalg/sparse.py +114 -0
  73. feectools/linalg/stencil.py +2923 -0
  74. feectools/linalg/stencil_dot_kernels.py +317 -0
  75. feectools/linalg/stencil_transpose_kernels.py +372 -0
  76. feectools/linalg/tests/__init__.py +0 -0
  77. feectools/linalg/tests/test_block.py +1588 -0
  78. feectools/linalg/tests/test_fft.py +106 -0
  79. feectools/linalg/tests/test_kron_stencil_matrix.py +114 -0
  80. feectools/linalg/tests/test_linalg.py +1065 -0
  81. feectools/linalg/tests/test_matrix_free.py +128 -0
  82. feectools/linalg/tests/test_solvers.py +213 -0
  83. feectools/linalg/tests/test_stencil_interface_matrix.py +379 -0
  84. feectools/linalg/tests/test_stencil_vector.py +1036 -0
  85. feectools/linalg/tests/test_stencil_vector_space.py +440 -0
  86. feectools/linalg/topetsc.py +522 -0
  87. feectools/linalg/utilities.py +200 -0
  88. feectools/utilities/__init__.py +0 -0
  89. feectools/utilities/quadratures.py +113 -0
  90. feectools/utilities/utils.py +166 -0
  91. feectools/version.py +1 -0
  92. feectools-0.1.0.dist-info/METADATA +66 -0
  93. feectools-0.1.0.dist-info/RECORD +98 -0
  94. feectools-0.1.0.dist-info/WHEEL +5 -0
  95. feectools-0.1.0.dist-info/entry_points.txt +3 -0
  96. feectools-0.1.0.dist-info/licenses/AUTHORS +22 -0
  97. feectools-0.1.0.dist-info/licenses/LICENSE +21 -0
  98. feectools-0.1.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,317 @@
1
+ def matvec_1d_kernel(mat: 'float[:, :]',
2
+ x: 'float[:]',
3
+ out: 'float[:]',
4
+ s_in: int,
5
+ p_in: int,
6
+ add: int,
7
+ s_out: int,
8
+ e_out: int,
9
+ p_out: int):
10
+
11
+ for i1 in range(s_out, e_out): # global row index
12
+ i1_loc = i1 - s_out # local row index
13
+ val = 0.
14
+ for d1 in range(2*p_in + 1):
15
+ val += mat[p_out + i1_loc, d1] * x[i1 + d1 - s_in]
16
+
17
+ out[p_out + i1_loc] = val
18
+
19
+ # last row treated separately
20
+ i1 = e_out
21
+ i1_loc = i1 - s_out # local row index
22
+ val = 0.
23
+ for d1 in range(2*p_in + add):
24
+ val += mat[p_out + i1_loc, d1] * x[i1 + d1 - s_in]
25
+
26
+ out[p_out + i1_loc] = val
27
+
28
+
29
+ def matvec_2d_kernel(mat: 'float[:, :, :, :]',
30
+ x: 'float[:, :]',
31
+ out: 'float[:, :]',
32
+ s_in: 'int[:]',
33
+ p_in: 'int[:]',
34
+ add: 'int[:]',
35
+ s_out: 'int[:]',
36
+ e_out: 'int[:]',
37
+ p_out: 'int[:]'):
38
+
39
+ #####################################
40
+ #####################################
41
+ # without last row in 1st direction #
42
+ #####################################
43
+ #####################################
44
+ for i1 in range(s_out[0], e_out[0]):
45
+ i1_loc = i1 - s_out[0]
46
+
47
+ #####################################
48
+ # without last row in 2nd direction #
49
+ #####################################
50
+ for i2 in range(s_out[1], e_out[1]):
51
+ i2_loc = i2 - s_out[1]
52
+
53
+ val = 0.
54
+ for d1 in range(2 * p_in[0] + 1):
55
+ for d2 in range(2 * p_in[1] + 1):
56
+ val += mat[p_out[0] + i1_loc,
57
+ p_out[1] + i2_loc,
58
+ d1, d2] * x[i1 + d1 - s_in[0],
59
+ i2 + d2 - s_in[1]]
60
+ out[p_out[0] + i1_loc,
61
+ p_out[1] + i2_loc] = val
62
+
63
+ ##############################################
64
+ # treat last row in 2nd direction separately #
65
+ ##############################################
66
+ i2 = e_out[1]
67
+ i2_loc = i2 - s_out[1]
68
+
69
+ val = 0.
70
+ for d1 in range(2 * p_in[0] + 1):
71
+ for d2 in range(2 * p_in[1] + add[1]):
72
+ val += mat[p_out[0] + i1_loc,
73
+ p_out[1] + i2_loc,
74
+ d1, d2] * x[i1 + d1 - s_in[0],
75
+ i2 + d2 - s_in[1]]
76
+ out[p_out[0] + i1_loc,
77
+ p_out[1] + i2_loc] = val
78
+
79
+ ##############################################
80
+ ##############################################
81
+ # treat last row in 1st direction separately #
82
+ ##############################################
83
+ ##############################################
84
+ i1 = e_out[0]
85
+ i1_loc = i1 - s_out[0]
86
+
87
+ #####################################
88
+ # without last row in 2nd direction #
89
+ #####################################
90
+ for i2 in range(s_out[1], e_out[1]):
91
+ i2_loc = i2 - s_out[1]
92
+
93
+ val = 0.
94
+ for d1 in range(2 * p_in[0] + add[0]):
95
+ for d2 in range(2 * p_in[1] + 1):
96
+ val += mat[p_out[0] + i1_loc,
97
+ p_out[1] + i2_loc,
98
+ d1, d2] * x[i1 + d1 - s_in[0],
99
+ i2 + d2 - s_in[1]]
100
+ out[p_out[0] + i1_loc,
101
+ p_out[1] + i2_loc] = val
102
+
103
+ ##############################################
104
+ # treat last row in 2nd direction separately #
105
+ ##############################################
106
+ i2 = e_out[1]
107
+ i2_loc = i2 - s_out[1]
108
+
109
+ val = 0.
110
+ for d1 in range(2 * p_in[0] + add[0]):
111
+ for d2 in range(2 * p_in[1] + add[1]):
112
+ val += mat[p_out[0] + i1_loc,
113
+ p_out[1] + i2_loc,
114
+ d1, d2] * x[i1 + d1 - s_in[0],
115
+ i2 + d2 - s_in[1]]
116
+
117
+ out[p_out[0] + i1_loc,
118
+ p_out[1] + i2_loc] = val
119
+
120
+
121
+ def matvec_3d_kernel(mat: 'float[:, :, :, :, :, :]',
122
+ x: 'float[:, :, :]',
123
+ out: 'float[:, :, :]',
124
+ s_in: 'int[:]',
125
+ p_in: 'int[:]',
126
+ add: 'int[:]',
127
+ s_out: 'int[:]',
128
+ e_out: 'int[:]',
129
+ p_out: 'int[:]'):
130
+
131
+ #####################################
132
+ #####################################
133
+ # without last row in 1st direction #
134
+ #####################################
135
+ #####################################
136
+ for i1 in range(s_out[0], e_out[0]):
137
+ i1_loc = i1 - s_out[0]
138
+
139
+ #####################################
140
+ # without last row in 2nd direction #
141
+ #####################################
142
+ for i2 in range(s_out[1], e_out[1]):
143
+ i2_loc = i2 - s_out[1]
144
+
145
+ # without last row in 3rd direction
146
+ for i3 in range(s_out[2], e_out[2]):
147
+ i3_loc = i3 - s_out[2]
148
+
149
+ val = 0.
150
+ for d1 in range(2*p_in[0] + 1):
151
+ for d2 in range(2*p_in[1] + 1):
152
+ for d3 in range(2*p_in[2] + 1):
153
+
154
+ val += mat[p_out[0] + i1_loc,
155
+ p_out[1] + i2_loc,
156
+ p_out[2] + i3_loc,
157
+ d1, d2, d3] * x[i1 + d1 - s_in[0],
158
+ i2 + d2 - s_in[1],
159
+ i3 + d3 - s_in[2]]
160
+ out[p_out[0] + i1_loc,
161
+ p_out[1] + i2_loc,
162
+ p_out[2] + i3_loc] = val
163
+
164
+ # treat last row in 3rd direction separately
165
+ i3 = e_out[2]
166
+ i3_loc = i3 - s_out[2]
167
+ val = 0.
168
+ for d1 in range(2*p_in[0] + 1):
169
+ for d2 in range(2*p_in[1] + 1):
170
+ for d3 in range(2*p_in[2] + add[2]):
171
+
172
+ val += mat[p_out[0] + i1_loc,
173
+ p_out[1] + i2_loc,
174
+ p_out[2] + i3_loc,
175
+ d1, d2, d3] * x[i1 + d1 - s_in[0],
176
+ i2 + d2 - s_in[1],
177
+ i3 + d3 - s_in[2]]
178
+ out[p_out[0] + i1_loc,
179
+ p_out[1] + i2_loc,
180
+ p_out[2] + i3_loc] = val
181
+
182
+ ##############################################
183
+ # treat last row in 2nd direction separately #
184
+ ##############################################
185
+ i2 = e_out[1]
186
+ i2_loc = i2 - s_out[1]
187
+
188
+ # without last row in 3rd direction
189
+ for i3 in range(s_out[2], e_out[2]):
190
+ i3_loc = i3 - s_out[2]
191
+
192
+ val = 0.
193
+ for d1 in range(2*p_in[0] + 1):
194
+ for d2 in range(2*p_in[1] + add[1]):
195
+ for d3 in range(2*p_in[2] + 1):
196
+
197
+ val += mat[p_out[0] + i1_loc,
198
+ p_out[1] + i2_loc,
199
+ p_out[2] + i3_loc,
200
+ d1, d2, d3] * x[i1 + d1 - s_in[0],
201
+ i2 + d2 - s_in[1],
202
+ i3 + d3 - s_in[2]]
203
+ out[p_out[0] + i1_loc,
204
+ p_out[1] + i2_loc,
205
+ p_out[2] + i3_loc] = val
206
+
207
+ # treat last row in 3rd direction separately
208
+ i3 = e_out[2]
209
+ i3_loc = i3 - s_out[2]
210
+ val = 0.
211
+ for d1 in range(2*p_in[0] + 1):
212
+ for d2 in range(2*p_in[1] + add[1]):
213
+ for d3 in range(2*p_in[2] + add[2]):
214
+
215
+ val += mat[p_out[0] + i1_loc,
216
+ p_out[1] + i2_loc,
217
+ p_out[2] + i3_loc,
218
+ d1, d2, d3] * x[i1 + d1 - s_in[0],
219
+ i2 + d2 - s_in[1],
220
+ i3 + d3 - s_in[2]]
221
+ out[p_out[0] + i1_loc,
222
+ p_out[1] + i2_loc,
223
+ p_out[2] + i3_loc] = val
224
+
225
+ ##############################################
226
+ ##############################################
227
+ # treat last row in 1st direction separately #
228
+ ##############################################
229
+ ##############################################
230
+ i1 = e_out[0]
231
+ i1_loc = i1 - s_out[0]
232
+
233
+ #####################################
234
+ # without last row in 2nd direction #
235
+ #####################################
236
+ for i2 in range(s_out[1], e_out[1]):
237
+ i2_loc = i2 - s_out[1]
238
+
239
+ # without last row in 3rd direction
240
+ for i3 in range(s_out[2], e_out[2]):
241
+ i3_loc = i3 - s_out[2]
242
+
243
+ val = 0.
244
+ for d1 in range(2*p_in[0] + add[0]):
245
+ for d2 in range(2*p_in[1] + 1):
246
+ for d3 in range(2*p_in[2] + 1):
247
+
248
+ val += mat[p_out[0] + i1_loc,
249
+ p_out[1] + i2_loc,
250
+ p_out[2] + i3_loc,
251
+ d1, d2, d3] * x[i1 + d1 - s_in[0],
252
+ i2 + d2 - s_in[1],
253
+ i3 + d3 - s_in[2]]
254
+ out[p_out[0] + i1_loc,
255
+ p_out[1] + i2_loc,
256
+ p_out[2] + i3_loc] = val
257
+
258
+ # treat last row in 3rd direction separately
259
+ i3 = e_out[2]
260
+ i3_loc = i3 - s_out[2]
261
+ val = 0.
262
+ for d1 in range(2*p_in[0] + add[0]):
263
+ for d2 in range(2*p_in[1] + 1):
264
+ for d3 in range(2*p_in[2] + add[2]):
265
+
266
+ val += mat[p_out[0] + i1_loc,
267
+ p_out[1] + i2_loc,
268
+ p_out[2] + i3_loc,
269
+ d1, d2, d3] * x[i1 + d1 - s_in[0],
270
+ i2 + d2 - s_in[1],
271
+ i3 + d3 - s_in[2]]
272
+ out[p_out[0] + i1_loc,
273
+ p_out[1] + i2_loc,
274
+ p_out[2] + i3_loc] = val
275
+
276
+ ##############################################
277
+ # treat last row in 2nd direction separately #
278
+ ##############################################
279
+ i2 = e_out[1]
280
+ i2_loc = i2 - s_out[1]
281
+
282
+ # without last row in 3rd direction
283
+ for i3 in range(s_out[2], e_out[2]):
284
+ i3_loc = i3 - s_out[2]
285
+
286
+ val = 0.
287
+ for d1 in range(2*p_in[0] + add[0]):
288
+ for d2 in range(2*p_in[1] + add[1]):
289
+ for d3 in range(2*p_in[2] + 1):
290
+
291
+ val += mat[p_out[0] + i1_loc,
292
+ p_out[1] + i2_loc,
293
+ p_out[2] + i3_loc,
294
+ d1, d2, d3] * x[i1 + d1 - s_in[0],
295
+ i2 + d2 - s_in[1],
296
+ i3 + d3 - s_in[2]]
297
+ out[p_out[0] + i1_loc,
298
+ p_out[1] + i2_loc,
299
+ p_out[2] + i3_loc] = val
300
+
301
+ # treat last row in 3rd direction separately
302
+ i3 = e_out[2]
303
+ i3_loc = i3 - s_out[2]
304
+ val = 0.
305
+ for d1 in range(2*p_in[0] + add[0]):
306
+ for d2 in range(2*p_in[1] + add[1]):
307
+ for d3 in range(2*p_in[2] + add[2]):
308
+
309
+ val += mat[p_out[0] + i1_loc,
310
+ p_out[1] + i2_loc,
311
+ p_out[2] + i3_loc,
312
+ d1, d2, d3] * x[i1 + d1 - s_in[0],
313
+ i2 + d2 - s_in[1],
314
+ i3 + d3 - s_in[2]]
315
+ out[p_out[0] + i1_loc,
316
+ p_out[1] + i2_loc,
317
+ p_out[2] + i3_loc] = val
@@ -0,0 +1,372 @@
1
+ def transpose_1d_kernel(mat: 'float[:, :]',
2
+ matT: 'float[:, :]',
3
+ s_in: int, # refers to matT
4
+ p_in: int,
5
+ add: int,
6
+ s_out: int,
7
+ e_out: int,
8
+ p_out: int):
9
+
10
+ for i1 in range(s_out, e_out): # global row index of matT = global column index of mat
11
+ i1_loc = i1 - s_out # local row index of matT
12
+ for d1 in range(2*p_in + 1):
13
+ j1 = i1 - p_in + d1 # global column index of matT
14
+ j1_loc = j1 - s_in # local column index of matT = local row index of mat
15
+
16
+ matT[p_out + i1_loc, d1] = mat[p_in + j1_loc, p_out + i1 - j1]
17
+
18
+ # last row treated separately
19
+ i1 = e_out
20
+ i1_loc = i1 - s_out # local row index of matT
21
+ for d1 in range(2*p_in + add):
22
+ j1 = i1 - p_in + d1 # global column index of matT
23
+ j1_loc = j1 - s_in # local column index of matT = local row index of mat
24
+
25
+ matT[p_out + i1_loc, d1] = mat[p_in + j1_loc, p_out + i1 - j1]
26
+
27
+ def transpose_2d_kernel(mat: 'float[:, :, :, :]',
28
+ matT: 'float[:, :, :, :]',
29
+ s_in: 'int[:]', # refers to matT
30
+ p_in: 'int[:]',
31
+ add: 'int[:]',
32
+ s_out: 'int[:]',
33
+ e_out: 'int[:]',
34
+ p_out: 'int[:]'):
35
+
36
+ #####################################
37
+ #####################################
38
+ # without last row in 1st direction #
39
+ #####################################
40
+ #####################################
41
+ for i1 in range(s_out[0], e_out[0]): # global row index of matT = global column index of mat
42
+ i1_loc = i1 - s_out[0] # local row index of matT
43
+
44
+ #####################################
45
+ # without last row in 2nd direction #
46
+ #####################################
47
+ for i2 in range(s_out[1], e_out[1]): # global row index of matT = global column index of mat
48
+ i2_loc = i2 - s_out[1] # local row index of matT
49
+
50
+ for d1 in range(2 * p_in[0] + 1):
51
+ j1 = i1 - p_in[0] + d1 # global column index of matT
52
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
53
+ for d2 in range(2 * p_in[1] + 1):
54
+ j2 = i2 - p_in[1] + d2 # global column index of matT
55
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
56
+
57
+ matT[p_out[0] + i1_loc,
58
+ p_out[1] + i2_loc,
59
+ d1, d2] = mat[p_in[0] + j1_loc,
60
+ p_in[1] + j2_loc,
61
+ p_out[0] + i1 - j1,
62
+ p_out[1] + i2 - j2]
63
+
64
+ ##############################################
65
+ # treat last row in 2nd direction separately #
66
+ ##############################################
67
+ i2 = e_out[1]
68
+ i2_loc = i2 - s_out[1] # local row index of matT
69
+
70
+ for d1 in range(2 * p_in[0] + 1):
71
+ j1 = i1 - p_in[0] + d1 # global column index of matT
72
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
73
+ for d2 in range(2 * p_in[1] + add[1]):
74
+ j2 = i2 - p_in[1] + d2 # global column index of matT
75
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
76
+
77
+ matT[p_out[0] + i1_loc,
78
+ p_out[1] + i2_loc,
79
+ d1, d2] = mat[p_in[0] + j1_loc,
80
+ p_in[1] + j2_loc,
81
+ p_out[0] + i1 - j1,
82
+ p_out[1] + i2 - j2]
83
+
84
+ ##############################################
85
+ ##############################################
86
+ # treat last row in 1st direction separately #
87
+ ##############################################
88
+ ##############################################
89
+ i1 = e_out[0]
90
+ i1_loc = i1 - s_out[0] # local row index of matT
91
+
92
+ #####################################
93
+ # without last row in 2nd direction #
94
+ #####################################
95
+ for i2 in range(s_out[1], e_out[1]):
96
+ i2_loc = i2 - s_out[1] # local row index of matT
97
+
98
+ for d1 in range(2 * p_in[0] + add[0]):
99
+ j1 = i1 - p_in[0] + d1 # global column index of matT
100
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
101
+ for d2 in range(2 * p_in[1] + 1):
102
+ j2 = i2 - p_in[1] + d2 # global column index of matT
103
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
104
+
105
+ matT[p_out[0] + i1_loc,
106
+ p_out[1] + i2_loc,
107
+ d1, d2] = mat[p_in[0] + j1_loc,
108
+ p_in[1] + j2_loc,
109
+ p_out[0] + i1 - j1,
110
+ p_out[1] + i2 - j2]
111
+
112
+ ##############################################
113
+ # treat last row in 2nd direction separately #
114
+ ##############################################
115
+ i2 = e_out[1]
116
+ i2_loc = i2 - s_out[1] # local row index of matT
117
+
118
+ for d1 in range(2 * p_in[0] + add[0]):
119
+ j1 = i1 - p_in[0] + d1 # global column index of matT
120
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
121
+ for d2 in range(2 * p_in[1] + add[1]):
122
+ j2 = i2 - p_in[1] + d2 # global column index of matT
123
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
124
+
125
+ matT[p_out[0] + i1_loc,
126
+ p_out[1] + i2_loc,
127
+ d1, d2] = mat[p_in[0] + j1_loc,
128
+ p_in[1] + j2_loc,
129
+ p_out[0] + i1 - j1,
130
+ p_out[1] + i2 - j2]
131
+
132
+
133
+ def transpose_3d_kernel(mat: 'float[:, :, :, :, :, :]',
134
+ matT: 'float[:, :, :, :, :, :]',
135
+ s_in: 'int[:]', # refers to matT
136
+ p_in: 'int[:]',
137
+ add: 'int[:]',
138
+ s_out: 'int[:]',
139
+ e_out: 'int[:]',
140
+ p_out: 'int[:]'):
141
+
142
+ #####################################
143
+ #####################################
144
+ # without last row in 1st direction #
145
+ #####################################
146
+ #####################################
147
+ for i1 in range(s_out[0], e_out[0]): # global row index of matT = global column index of mat
148
+ i1_loc = i1 - s_out[0] # local row index of matT
149
+
150
+ #####################################
151
+ # without last row in 2nd direction #
152
+ #####################################
153
+ for i2 in range(s_out[1], e_out[1]): # global row index of matT = global column index of mat
154
+ i2_loc = i2 - s_out[1] # local row index of matT
155
+
156
+ # without last row in 3rd direction
157
+ for i3 in range(s_out[2], e_out[2]): # global row index of matT = global column index of mat
158
+ i3_loc = i3 - s_out[2] # local row index of matT
159
+
160
+ for d1 in range(2*p_in[0] + 1):
161
+ j1 = i1 - p_in[0] + d1 # global column index of matT
162
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
163
+ for d2 in range(2*p_in[1] + 1):
164
+ j2 = i2 - p_in[1] + d2 # global column index of matT
165
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
166
+ for d3 in range(2*p_in[2] + 1):
167
+ j3 = i3 - p_in[2] + d3 # global column index of matT
168
+ j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat
169
+
170
+ matT[p_out[0] + i1_loc,
171
+ p_out[1] + i2_loc,
172
+ p_out[2] + i3_loc,
173
+ d1, d2, d3] = mat[p_in[0] + j1_loc,
174
+ p_in[1] + j2_loc,
175
+ p_in[2] + j3_loc,
176
+ p_out[0] + i1 - j1,
177
+ p_out[1] + i2 - j2,
178
+ p_out[2] + i3 - j3]
179
+
180
+ # treat last row in 3rd direction separately
181
+ i3 = e_out[2]
182
+ i3_loc = i3 - s_out[2] # local row index of matT
183
+
184
+ for d1 in range(2*p_in[0] + 1):
185
+ j1 = i1 - p_in[0] + d1 # global column index of matT
186
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
187
+ for d2 in range(2*p_in[1] + 1):
188
+ j2 = i2 - p_in[1] + d2 # global column index of matT
189
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
190
+ for d3 in range(2*p_in[2] + add[2]):
191
+ j3 = i3 - p_in[2] + d3 # global column index of matT
192
+ j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat
193
+
194
+ matT[p_out[0] + i1_loc,
195
+ p_out[1] + i2_loc,
196
+ p_out[2] + i3_loc,
197
+ d1, d2, d3] = mat[p_in[0] + j1_loc,
198
+ p_in[1] + j2_loc,
199
+ p_in[2] + j3_loc,
200
+ p_out[0] + i1 - j1,
201
+ p_out[1] + i2 - j2,
202
+ p_out[2] + i3 - j3]
203
+
204
+ ##############################################
205
+ # treat last row in 2nd direction separately #
206
+ ##############################################
207
+ i2 = e_out[1]
208
+ i2_loc = i2 - s_out[1] # local row index of matT
209
+
210
+ # without last row in 3rd direction
211
+ for i3 in range(s_out[2], e_out[2]): # global row index of matT = global column index of mat
212
+ i3_loc = i3 - s_out[2] # local row index of matT
213
+
214
+ for d1 in range(2*p_in[0] + 1):
215
+ j1 = i1 - p_in[0] + d1 # global column index of matT
216
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
217
+ for d2 in range(2*p_in[1] + add[1]):
218
+ j2 = i2 - p_in[1] + d2 # global column index of matT
219
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
220
+ for d3 in range(2*p_in[2] + 1):
221
+ j3 = i3 - p_in[2] + d3 # global column index of matT
222
+ j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat
223
+
224
+ matT[p_out[0] + i1_loc,
225
+ p_out[1] + i2_loc,
226
+ p_out[2] + i3_loc,
227
+ d1, d2, d3] = mat[p_in[0] + j1_loc,
228
+ p_in[1] + j2_loc,
229
+ p_in[2] + j3_loc,
230
+ p_out[0] + i1 - j1,
231
+ p_out[1] + i2 - j2,
232
+ p_out[2] + i3 - j3]
233
+
234
+ # treat last row in 3rd direction separately
235
+ i3 = e_out[2]
236
+ i3_loc = i3 - s_out[2] # local row index of matT
237
+
238
+ for d1 in range(2*p_in[0] + 1):
239
+ j1 = i1 - p_in[0] + d1 # global column index of matT
240
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
241
+ for d2 in range(2*p_in[1] + add[1]):
242
+ j2 = i2 - p_in[1] + d2 # global column index of matT
243
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
244
+ for d3 in range(2*p_in[2] + add[2]):
245
+ j3 = i3 - p_in[2] + d3 # global column index of matT
246
+ j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat
247
+
248
+ matT[p_out[0] + i1_loc,
249
+ p_out[1] + i2_loc,
250
+ p_out[2] + i3_loc,
251
+ d1, d2, d3] = mat[p_in[0] + j1_loc,
252
+ p_in[1] + j2_loc,
253
+ p_in[2] + j3_loc,
254
+ p_out[0] + i1 - j1,
255
+ p_out[1] + i2 - j2,
256
+ p_out[2] + i3 - j3]
257
+
258
+ ##############################################
259
+ ##############################################
260
+ # treat last row in 1st direction separately #
261
+ ##############################################
262
+ ##############################################
263
+ i1 = e_out[0]
264
+ i1_loc = i1 - s_out[0] # local row index of matT
265
+
266
+ #####################################
267
+ # without last row in 2nd direction #
268
+ #####################################
269
+ for i2 in range(s_out[1], e_out[1]): # global row index of matT = global column index of mat
270
+ i2_loc = i2 - s_out[1] # local row index of matT
271
+
272
+ # without last row in 3rd direction
273
+ for i3 in range(s_out[2], e_out[2]): # global row index of matT = global column index of mat
274
+ i3_loc = i3 - s_out[2] # local row index of matT
275
+
276
+ for d1 in range(2*p_in[0] + add[0]):
277
+ j1 = i1 - p_in[0] + d1 # global column index of matT
278
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
279
+ for d2 in range(2*p_in[1] + 1):
280
+ j2 = i2 - p_in[1] + d2 # global column index of matT
281
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
282
+ for d3 in range(2*p_in[2] + 1):
283
+ j3 = i3 - p_in[2] + d3 # global column index of matT
284
+ j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat
285
+
286
+ matT[p_out[0] + i1_loc,
287
+ p_out[1] + i2_loc,
288
+ p_out[2] + i3_loc,
289
+ d1, d2, d3] = mat[p_in[0] + j1_loc,
290
+ p_in[1] + j2_loc,
291
+ p_in[2] + j3_loc,
292
+ p_out[0] + i1 - j1,
293
+ p_out[1] + i2 - j2,
294
+ p_out[2] + i3 - j3]
295
+
296
+ # treat last row in 3rd direction separately
297
+ i3 = e_out[2]
298
+ i3_loc = i3 - s_out[2] # local row index of matT
299
+
300
+ for d1 in range(2*p_in[0] + add[0]):
301
+ j1 = i1 - p_in[0] + d1 # global column index of matT
302
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
303
+ for d2 in range(2*p_in[1] + 1):
304
+ j2 = i2 - p_in[1] + d2 # global column index of matT
305
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
306
+ for d3 in range(2*p_in[2] + add[2]):
307
+ j3 = i3 - p_in[2] + d3 # global column index of matT
308
+ j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat
309
+
310
+ matT[p_out[0] + i1_loc,
311
+ p_out[1] + i2_loc,
312
+ p_out[2] + i3_loc,
313
+ d1, d2, d3] = mat[p_in[0] + j1_loc,
314
+ p_in[1] + j2_loc,
315
+ p_in[2] + j3_loc,
316
+ p_out[0] + i1 - j1,
317
+ p_out[1] + i2 - j2,
318
+ p_out[2] + i3 - j3]
319
+
320
+ ##############################################
321
+ # treat last row in 2nd direction separately #
322
+ ##############################################
323
+ i2 = e_out[1]
324
+ i2_loc = i2 - s_out[1] # local row index of matT
325
+
326
+ # without last row in 3rd direction
327
+ for i3 in range(s_out[2], e_out[2]): # global row index of matT = global column index of mat
328
+ i3_loc = i3 - s_out[2] # local row index of matT
329
+
330
+ for d1 in range(2*p_in[0] + add[0]):
331
+ j1 = i1 - p_in[0] + d1 # global column index of matT
332
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
333
+ for d2 in range(2*p_in[1] + add[1]):
334
+ j2 = i2 - p_in[1] + d2 # global column index of matT
335
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
336
+ for d3 in range(2*p_in[2] + 1):
337
+ j3 = i3 - p_in[2] + d3 # global column index of matT
338
+ j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat
339
+
340
+ matT[p_out[0] + i1_loc,
341
+ p_out[1] + i2_loc,
342
+ p_out[2] + i3_loc,
343
+ d1, d2, d3] = mat[p_in[0] + j1_loc,
344
+ p_in[1] + j2_loc,
345
+ p_in[2] + j3_loc,
346
+ p_out[0] + i1 - j1,
347
+ p_out[1] + i2 - j2,
348
+ p_out[2] + i3 - j3]
349
+
350
+ # treat last row in 3rd direction separately
351
+ i3 = e_out[2]
352
+ i3_loc = i3 - s_out[2] # local row index of matT
353
+
354
+ for d1 in range(2*p_in[0] + add[0]):
355
+ j1 = i1 - p_in[0] + d1 # global column index of matT
356
+ j1_loc = j1 - s_in[0] # local column index of matT = local row index of mat
357
+ for d2 in range(2*p_in[1] + add[1]):
358
+ j2 = i2 - p_in[1] + d2 # global column index of matT
359
+ j2_loc = j2 - s_in[1] # local column index of matT = local row index of mat
360
+ for d3 in range(2*p_in[2] + add[2]):
361
+ j3 = i3 - p_in[2] + d3 # global column index of matT
362
+ j3_loc = j3 - s_in[2] # local column index of matT = local row index of mat
363
+
364
+ matT[p_out[0] + i1_loc,
365
+ p_out[1] + i2_loc,
366
+ p_out[2] + i3_loc,
367
+ d1, d2, d3] = mat[p_in[0] + j1_loc,
368
+ p_in[1] + j2_loc,
369
+ p_in[2] + j3_loc,
370
+ p_out[0] + i1 - j1,
371
+ p_out[1] + i2 - j2,
372
+ p_out[2] + i3 - j3]
File without changes