1 #include "extensions.h"
2 #include <stdlib.h>
3 #include <numpy/arrayobject.h>
4
5
_pw_insert(int nG,int nQ,double complex * c_G,npy_int32 * Q_G,double scale,double complex * tmp_Q)6 void _pw_insert(int nG,
7 int nQ,
8 double complex* c_G,
9 npy_int32* Q_G,
10 double scale,
11 double complex* tmp_Q)
12 // Does the same as these two lines of Python:
13 //
14 // tmp_Q[:] = 0.0
15 // tmp_Q.ravel()[Q_G] = c_G * scale
16 {
17 int Q1 = 0;
18 for (int G = 0; G < nG; G++) {
19 int Q2 = Q_G[G];
20 for (; Q1 < Q2; Q1++)
21 tmp_Q[Q1] = 0.0;
22 tmp_Q[Q1++] = c_G[G] * scale;
23 }
24 for (; Q1 < nQ; Q1++)
25 tmp_Q[Q1] = 0.0;
26 }
27
28
pw_insert(PyObject * self,PyObject * args)29 PyObject *pw_insert(PyObject *self, PyObject *args)
30 // Python wrapper
31 {
32 PyArrayObject *c_G_obj, *Q_G_obj, *tmp_Q_obj;
33 double scale;
34 if (!PyArg_ParseTuple(args, "OOdO",
35 &c_G_obj, &Q_G_obj, &scale, &tmp_Q_obj))
36 return NULL;
37 double complex *c_G = PyArray_DATA(c_G_obj);
38 npy_int32 *Q_G = PyArray_DATA(Q_G_obj);
39 double complex *tmp_Q = PyArray_DATA(tmp_Q_obj);
40 int nG = PyArray_SIZE(c_G_obj);
41 int nQ = PyArray_SIZE(tmp_Q_obj);
42 _pw_insert(nG, nQ, c_G, Q_G, scale, tmp_Q);
43 Py_RETURN_NONE;
44 }
45
46
pw_precond(PyObject * self,PyObject * args)47 PyObject *pw_precond(PyObject *self, PyObject *args)
48 {
49 PyArrayObject *G2_G_obj;
50 PyArrayObject *R_G_obj;
51 double ekin;
52 PyArrayObject *out_G_obj;
53
54 if (!PyArg_ParseTuple(args, "OOdO",
55 &G2_G_obj, &R_G_obj, &ekin, &out_G_obj))
56 return NULL;
57
58 double *G2_G = PyArray_DATA(G2_G_obj);
59 double complex *R_G = PyArray_DATA(R_G_obj);
60 double complex *out_G = PyArray_DATA(out_G_obj);
61 int nG = PyArray_SIZE(G2_G_obj);
62
63 for (int G = 0; G < nG; G++) {
64 double x = 1 / ekin / 3 * G2_G[G];
65 double a = 27.0 + x * (18.0 + x * (12.0 + x * 8.0));
66 double xx = x * x;
67 out_G[G] = -4.0 / 3 / ekin * a / (a + 16.0 * xx * xx) * R_G[G];
68 }
69 Py_RETURN_NONE;
70 }
71
72
pwlfc_expand(PyObject * self,PyObject * args)73 PyObject *pwlfc_expand(PyObject *self, PyObject *args)
74 {
75 PyArrayObject *f_Gs_obj;
76 PyArrayObject *emiGR_Ga_obj;
77 PyArrayObject *Y_GL_obj;
78 PyArrayObject *l_s_obj;
79 PyArrayObject *a_J_obj;
80 PyArrayObject *s_J_obj;
81 int cc;
82 PyArrayObject *f_GI_obj;
83
84 if (!PyArg_ParseTuple(args, "OOOOOOiO",
85 &f_Gs_obj, &emiGR_Ga_obj, &Y_GL_obj,
86 &l_s_obj, &a_J_obj, &s_J_obj,
87 &cc, &f_GI_obj))
88 return NULL;
89
90 double *f_Gs = PyArray_DATA(f_Gs_obj);
91 double complex *emiGR_Ga = PyArray_DATA(emiGR_Ga_obj);
92 double *Y_GL = PyArray_DATA(Y_GL_obj);
93 npy_int32 *l_s = PyArray_DATA(l_s_obj);
94 npy_int32 *a_J = PyArray_DATA(a_J_obj);
95 npy_int32 *s_J = PyArray_DATA(s_J_obj);
96 double *f_GI = PyArray_DATA(f_GI_obj);
97
98 int nG = PyArray_DIM(emiGR_Ga_obj, 0);
99 int nJ = PyArray_DIM(a_J_obj, 0);
100 int nL = PyArray_DIM(Y_GL_obj, 1);
101 int natoms = PyArray_DIM(emiGR_Ga_obj, 1);
102 int nsplines = PyArray_DIM(f_Gs_obj, 1);
103
104 double complex imag_powers[4] = {1.0, -I, -1.0, I};
105
106 if (PyArray_ITEMSIZE(f_GI_obj) == 16)
107 for(int G = 0; G < nG; G++) {
108 for (int J = 0; J < nJ; J++) {
109 int s = s_J[J];
110 int l = l_s[s];
111 double complex f1 = (emiGR_Ga[a_J[J]] *
112 f_Gs[s] *
113 imag_powers[l % 4]);
114 for (int m = 0; m < 2 * l + 1; m++) {
115 double complex f = f1 * Y_GL[l * l + m];
116 *f_GI++ = creal(f);
117 *f_GI++ = cc ? -cimag(f) : cimag(f);
118 }
119 }
120 f_Gs += nsplines;
121 emiGR_Ga += natoms;
122 Y_GL += nL;
123 }
124 else {
125 int nI = PyArray_DIM(f_GI_obj, 1);
126 for(int G = 0; G < nG; G++) {
127 for (int J = 0; J < nJ; J++) {
128 int s = s_J[J];
129 int l = l_s[s];
130 double complex f1 = (emiGR_Ga[a_J[J]] *
131 f_Gs[s] *
132 imag_powers[l % 4]);
133 for (int m = 0; m < 2 * l + 1; m++) {
134 double complex f = f1 * Y_GL[l * l + m];
135 f_GI[0] = creal(f);
136 f_GI[nI] = cc ? -cimag(f) : cimag(f);
137 f_GI++;
138 }
139 }
140 f_Gs += nsplines;
141 emiGR_Ga += natoms;
142 Y_GL += nL;
143 f_GI += nI;
144 }
145 }
146
147 Py_RETURN_NONE;
148 }
149
150
plane_wave_grid(PyObject * self,PyObject * args)151 PyObject *plane_wave_grid(PyObject *self, PyObject *args)
152 {
153 PyArrayObject* beg_c;
154 PyArrayObject* end_c;
155 PyArrayObject* h_c;
156 PyArrayObject* k_c;
157 PyArrayObject* r0_c;
158 PyArrayObject* pw_g;
159 if (!PyArg_ParseTuple(args, "OOOOOO", &beg_c, &end_c, &h_c,
160 &k_c, &r0_c, &pw_g))
161 return NULL;
162
163 long *beg = LONGP(beg_c);
164 long *end = LONGP(end_c);
165 double *h = DOUBLEP(h_c);
166 double *vk = DOUBLEP(k_c);
167 double *vr0 = DOUBLEP(r0_c);
168 double_complex *pw = COMPLEXP(pw_g);
169
170 double kr[3], kr0[3];
171 int n[3], ij;
172 for (int c = 0; c < 3; c++) {
173 n[c] = end[c] - beg[c];
174 kr0[c] = vk[c] * vr0[c];
175 }
176 for (int i = 0; i < n[0]; i++) {
177 kr[0] = vk[0] * h[0] * (beg[0] + i) - kr0[0];
178 for (int j = 0; j < n[1]; j++) {
179 kr[1] = kr[0] + vk[1] * h[1] * (beg[1] + j) - kr0[1];
180 ij = (i*n[1] + j)*n[2];
181 for (int k = 0; k < n[2]; k++) {
182 kr[2] = kr[1] + vk[2] * h[2] * (beg[2] + k) - kr0[2];
183 pw[ij + k] = cos(kr[2]) + I * sin(kr[2]);
184 }
185 }
186 }
187 Py_RETURN_NONE;
188 }
189