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