KIDS  ver-0.0.1
KIDS : Kernel Integrated Dynamics Simulator
Loading...
Searching...
No Matches
array_xtensor.h
Go to the documentation of this file.
1#ifndef Array_XTENSOR_H
2#define Array_XTENSOR_H
3
4#include <array>
5#include <xtensor-blas/xlinalg.hpp>
6#include <xtensor/xadapt.hpp>
7#include <xtensor/xarray.hpp>
8
9namespace ARRAY_XT {
10
11template <class T>
12void ARRAY_CLEAR(T* A, size_t N) {
13 memset(A, 0, N * sizeof(T));
14}
15
16template <class TA, class TB, class TC>
17void ARRAY_MATMUL(TA* A, TB* B, TC* C, size_t N1, size_t N2, size_t N3) {
18 std::array<size_t, 2> shapeA = {N1, N3};
19 std::array<size_t, 2> shapeB = {N1, N2};
20 std::array<size_t, 2> shapeC = {N2, N3};
21 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
22 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
23 auto c = xt::adapt(C, N2 * N3, xt::no_ownership(), shapeC);
24 a = xt::linalg::dot(b, c);
25}
26
27template <class TA, class TC>
28void ARRAY_MATMUL_TRANS1(TA* A, num_real* B, TC* C, size_t N1, size_t N2, size_t N3) {
29 std::array<size_t, 2> shapeA = {N1, N3};
30 std::array<size_t, 2> shapeB = {N2, N1};
31 std::array<size_t, 2> shapeC = {N2, N3};
32 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
33 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
34 auto c = xt::adapt(C, N2 * N3, xt::no_ownership(), shapeC);
35 a = xt::linalg::dot(xt::transpose(b), c);
36}
37
38template <class TA, class TC>
39void ARRAY_MATMUL_TRANS1(TA* A, num_complex* B, TC* C, size_t N1, size_t N2, size_t N3) {
40 std::array<size_t, 2> shapeA = {N1, N3};
41 std::array<size_t, 2> shapeB = {N2, N1};
42 std::array<size_t, 2> shapeC = {N2, N3};
43 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
44 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
45 auto c = xt::adapt(C, N2 * N3, xt::no_ownership(), shapeC);
46 a = xt::linalg::dot(xt::transpose(xt::conj(b)), c);
47}
48
49template <class TA, class TB>
50void ARRAY_MATMUL_TRANS2(TA* A, TB* B, num_real* C, size_t N1, size_t N2, size_t N3) {
51 std::array<size_t, 2> shapeA = {N1, N3};
52 std::array<size_t, 2> shapeB = {N1, N2};
53 std::array<size_t, 2> shapeC = {N3, N2};
54 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
55 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
56 auto c = xt::adapt(C, N2 * N3, xt::no_ownership(), shapeC);
57 a = xt::linalg::dot(b, xt::transpose(c));
58}
59
60template <class TA, class TB>
61void ARRAY_MATMUL_TRANS2(TA* A, TB* B, num_complex* C, size_t N1, size_t N2, size_t N3) {
62 std::array<size_t, 2> shapeA = {N1, N3};
63 std::array<size_t, 2> shapeB = {N1, N2};
64 std::array<size_t, 2> shapeC = {N3, N2};
65 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
66 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
67 auto c = xt::adapt(C, N2 * N3, xt::no_ownership(), shapeC);
68 a = xt::linalg::dot(b, xt::transpose(xt::conj(c)));
69}
70
71template <class T>
72void ARRAY_OUTER_CONJ2(T* A, T* B, T* C, size_t N1, size_t N2) {
73 ARRAY_MATMUL_TRANS2(A, B, C, N1, 1, N2);
74}
75
76template <class TA, class TC>
77void ARRAY_MATMUL3_TRANS1(TA* A, num_real* B, TC* C, num_real* D, size_t N1, size_t N2, size_t N0, size_t N3) {
78 if (N0 == 0) {
79 std::array<size_t, 2> shapeA = {N1, N3};
80 std::array<size_t, 2> shapeB = {N2, N1};
81 std::array<size_t, 2> shapeC = {N2, 1};
82 std::array<size_t, 2> shapeD = {N2, N3};
83 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
84 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
85 auto c = xt::adapt(C, N2 * 1, xt::no_ownership(), shapeC);
86 auto d = xt::adapt(D, N2 * N3, xt::no_ownership(), shapeD);
87
88 a = xt::linalg::dot(xt::transpose(b), xt::linalg::dot(xt::diag(c), d));
89
90 } else { // N0 == N2
91 std::array<size_t, 2> shapeA = {N1, N3};
92 std::array<size_t, 2> shapeB = {N2, N1};
93 std::array<size_t, 2> shapeC = {N2, N2};
94 std::array<size_t, 2> shapeD = {N2, N3};
95 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
96 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
97 auto c = xt::adapt(C, N2 * N2, xt::no_ownership(), shapeC);
98 auto d = xt::adapt(D, N2 * N3, xt::no_ownership(), shapeD);
99
100 a = xt::linalg::dot(xt::transpose(b), xt::linalg::dot(c, d));
101 }
102}
103
104template <class TA, class TC>
105void ARRAY_MATMUL3_TRANS1(TA* A, num_complex* B, TC* C, num_complex* D, size_t N1, size_t N2, size_t N0, size_t N3) {
106 if (N0 == 0) {
107 std::array<size_t, 2> shapeA = {N1, N3};
108 std::array<size_t, 2> shapeB = {N2, N1};
109 std::array<size_t, 2> shapeC = {N2, 1};
110 std::array<size_t, 2> shapeD = {N2, N3};
111 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
112 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
113 auto c = xt::adapt(C, N2 * 1, xt::no_ownership(), shapeC);
114 auto d = xt::adapt(D, N2 * N3, xt::no_ownership(), shapeD);
115
116 a = xt::linalg::dot(xt::transpose(xt::conj(b)), xt::linalg::dot(xt::diag(c), d));
117
118 } else { // N0 == N2
119 std::array<size_t, 2> shapeA = {N1, N3};
120 std::array<size_t, 2> shapeB = {N2, N1};
121 std::array<size_t, 2> shapeC = {N2, N2};
122 std::array<size_t, 2> shapeD = {N2, N3};
123 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
124 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
125 auto c = xt::adapt(C, N2 * N2, xt::no_ownership(), shapeC);
126 auto d = xt::adapt(D, N2 * N3, xt::no_ownership(), shapeD);
127
128 a = xt::linalg::dot(xt::transpose(xt::conj(b)), xt::linalg::dot(c, d));
129 }
130}
131
132template <class TA, class TC>
133void ARRAY_MATMUL3_TRANS2(TA* A, num_real* B, TC* C, num_real* D, size_t N1, size_t N2, size_t N0, size_t N3) {
134 if (N0 == 0) {
135 std::array<size_t, 2> shapeA = {N1, N3};
136 std::array<size_t, 2> shapeB = {N1, N2};
137 std::array<size_t, 2> shapeC = {N2, 1};
138 std::array<size_t, 2> shapeD = {N3, N2};
139 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
140 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
141 auto c = xt::adapt(C, N2 * 1, xt::no_ownership(), shapeC);
142 auto d = xt::adapt(D, N2 * N3, xt::no_ownership(), shapeD);
143
144 a = xt::linalg::dot(b, xt::linalg::dot(xt::diag(c), xt::transpose(d)));
145 } else { // N0 == N2
146 std::array<size_t, 2> shapeA = {N1, N3};
147 std::array<size_t, 2> shapeB = {N1, N2};
148 std::array<size_t, 2> shapeC = {N2, N2};
149 std::array<size_t, 2> shapeD = {N3, N2};
150 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
151 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
152 auto c = xt::adapt(C, N2 * N2, xt::no_ownership(), shapeC);
153 auto d = xt::adapt(D, N2 * N3, xt::no_ownership(), shapeD);
154
155 a = xt::linalg::dot(b, xt::linalg::dot(c, xt::transpose(d)));
156 }
157}
158
159template <class TA, class TC>
160void ARRAY_MATMUL3_TRANS2(TA* A, num_complex* B, TC* C, num_complex* D, size_t N1, size_t N2, size_t N0, size_t N3) {
161 if (N0 == 0) {
162 std::array<size_t, 2> shapeA = {N1, N3};
163 std::array<size_t, 2> shapeB = {N1, N2};
164 std::array<size_t, 2> shapeC = {N2, 1};
165 std::array<size_t, 2> shapeD = {N3, N2};
166 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
167 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
168 auto c = xt::adapt(C, N2 * 1, xt::no_ownership(), shapeC);
169 auto d = xt::adapt(D, N2 * N3, xt::no_ownership(), shapeD);
170
171 a = xt::linalg::dot(b, xt::linalg::dot(xt::diag(c), xt::transpose(xt::conj(d))));
172 } else { // N0 == N2
173 std::array<size_t, 2> shapeA = {N1, N3};
174 std::array<size_t, 2> shapeB = {N1, N2};
175 std::array<size_t, 2> shapeC = {N2, N2};
176 std::array<size_t, 2> shapeD = {N3, N2};
177 auto a = xt::adapt(A, N1 * N3, xt::no_ownership(), shapeA);
178 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
179 auto c = xt::adapt(C, N2 * N2, xt::no_ownership(), shapeC);
180 auto d = xt::adapt(D, N2 * N3, xt::no_ownership(), shapeD);
181
182 a = xt::linalg::dot(b, xt::linalg::dot(c, xt::transpose(xt::conj(d))));
183 }
184}
185
186template <class TB, class TC>
187TB ARRAY_TRACE2(TB* B, TC* C, size_t N1, size_t N2) {
188 std::array<size_t, 2> shapeB = {N1, N2};
189 std::array<size_t, 2> shapeC = {N2, N1};
190 auto b = xt::adapt(B, N1 * N2, xt::no_ownership(), shapeB);
191 auto c = xt::adapt(C, N2 * N1, xt::no_ownership(), shapeC);
192
193 TB res = xt::sum(b * xt::transpose(c))();
194 return res;
195}
196
197template <class T>
198void ARRAY_EYE(T* A, size_t n) {
199 std::array<size_t, 2> shapeA = {n, n};
200 auto a = xt::adapt(A, n * n, xt::no_ownership(), shapeA);
201 a = xt::eye<T>(n);
202}
203
204
205}; // namespace ARRAY_XT
206
207#endif // Array_XTENSOR_H
ARRAY_MATMUL_TRANS2
#define ARRAY_MATMUL_TRANS2(_A, _B, _C, _n1, _n2, _n3)
Definition array_macro.h:179
ARRAY_MATMUL3_TRANS1
#define ARRAY_MATMUL3_TRANS1(_A, _B, _C, _D, _n1, _n2, _n0, _n3)
Definition array_macro.h:199
ARRAY_MATMUL3_TRANS2
#define ARRAY_MATMUL3_TRANS2(_A, _B, _C, _D, _n1, _n2, _n0, _n3)
Definition array_macro.h:219
ARRAY_OUTER_CONJ2
#define ARRAY_OUTER_CONJ2(_A, _B, _C, _n1, _n2)
Definition array_macro.h:136
ARRAY_MATMUL
#define ARRAY_MATMUL(_A, _B, _C, _n1, _n2, _n3)
Definition array_macro.h:146
ARRAY_EYE
#define ARRAY_EYE(_A, _n)
Definition array_macro.h:75
ARRAY_MATMUL_TRANS1
#define ARRAY_MATMUL_TRANS1(_A, _B, _C, _n1, _n2, _n3)
Definition array_macro.h:162
ARRAY_TRACE2
#define ARRAY_TRACE2(_B, _C, _n1, _n2)
Definition array_macro.h:121
ARRAY_CLEAR
#define ARRAY_CLEAR(_A, _n)
Definition array_macro.h:36
ARRAY_XT
Definition array_xtensor.h:9