qpms.trans_calculator.get_AB working

(translations.c needs cleanup, remove duplicate code)


Former-commit-id: e99463b00456030384bf85fd3f782f47d2a44f5b

qpms/qpms_c.pyx

@@ -175,6 +175,11 @@ cdef extern from "translations.h":
     cdouble qpms_trans_calculator_get_B_ext(const qpms_trans_calculator* c,
             int m, int n, int mu, int nu, double kdlj_r, double kdlj_th, double kdlj_phi,
             int r_ge_d, int J)
+    int qpms_trans_calculator_get_AB_p_ext(const qpms_trans_calculator* c,
+            cdouble *Adest, cdouble *Bdest,
+            int m, int n, int mu, int nu, double kdlj_r, double kdlj_th, double kdlj_phi,
+            int r_ge_d, int J)
+
 
 
 
@@ -369,11 +374,10 @@ cdef void trans_calculator_loop_E_C_DD_iiiidddii_As_lllldddbl_DD(char **args, np
 cdef np.PyUFuncGenericFunction trans_calculator_get_X_loop_funcs[1]
 trans_calculator_get_X_loop_funcs[0] = trans_calculator_loop_D_Ciiiidddii_As_D_lllldddbl
 
-#TODO
+cdef np.PyUFuncGenericFunction trans_calculator_get_AB_loop_funcs[1]
-#cdef np.PyUFuncGenericFunction trans_calculator_get_AB_loop_funcs[1]
+trans_calculator_get_AB_loop_funcs[0] = trans_calculator_loop_E_C_DD_iiiidddii_As_lllldddbl_DD
-#trans_calculator_get_AB_loop_funcs[0] = trans_calculator_loop_E_C_DD_iiiidddii_As_lllldddbl_DD
+cdef void *trans_calculator_get_AB_elementwise_funcs[1]
-#cdef void *trans_calculator_get_AB_elementwise_funcs[1]
+trans_calculator_get_AB_elementwise_funcs[0] = <void *>qpms_trans_calculator_get_AB_p_ext
-#trans_calculator_get_AB_elementwise_funcs[0] = qpms_trans_calculator_get_AB_p_ext
 
 cdef class trans_calculator:
     cdef qpms_trans_calculator* c
@@ -386,7 +390,8 @@ cdef class trans_calculator:
     cdef trans_calculator_get_X_data_t get_AB_data[1]
     cdef trans_calculator_get_X_data_t* get_AB_data_p[1]
     cdef public: # TODO CHECK FOR CORRECT REFERENCE COUNTING AND LEAKS
-        object get_A, get_B, #get_AB
+        # have to be cdef public in order that __init__ can set these attributes
+        object get_A, get_B, get_AB
 
     def __cinit__(self, int lMax, int normalization = 1):
         self.c = qpms_trans_calculator_init(lMax, normalization)
@@ -426,9 +431,26 @@ cdef class trans_calculator:
                 """, # doc
                 0 # unused
                 )
-
+        self.get_AB_data[0].c = self.c
+        self.get_AB_data[0].cmethod = <void *>qpms_trans_calculator_get_AB_p_ext
+        self.get_AB_data_p[0] = &(self.get_AB_data[0])
+        self.get_AB = <object>np.PyUFunc_FromFuncAndData(# TODO CHECK FOR CORRECT REFERENCE COUNTING AND LEAKS
+                trans_calculator_get_AB_loop_funcs, # func
+                <void **>self.get_AB_data_p, #data
+                ufunc__get_both_coeff_types, #types
+                1, # ntypes: number of supported input types
+                9, # nin: number of input args
+                2, # nout: number of output args
+                0, # identity element, unused
+                "get_AB", #name
+                """
+                TODO doc
+                """, # doc
+                0 # unused
+                )
     def __dealloc__(self):
         qpms_trans_calculator_free(self.c)
+        # TODO Reference counts to get_A, get_B, get_AB?
 
 
     # TODO make possible to access the attributes (to show normalization etc)

qpms/translations.c

@@ -480,6 +480,56 @@ complex double qpms_trans_calculator_get_B_buf(const qpms_trans_calculator *c,
 	assert(0);
 }
 
+int qpms_trans_calculator_get_AB_buf_p(const qpms_trans_calculator *c,
+		complex double *Adest, complex double *Bdest,
+		int m, int n, int mu, int nu, sph_t kdlj,
+		bool r_ge_d, qpms_bessel_t J,
+		complex double *bessel_buf, double *legendre_buf) {
+	if (r_ge_d) J = QPMS_BESSEL_REGULAR;
+	switch(c->normalization) {
+		case QPMS_NORMALIZATION_TAYLOR:
+			{
+				double costheta = cos(kdlj.theta);
+				if (gsl_sf_legendre_array_e(GSL_SF_LEGENDRE_NONE,n+nu+1,
+																costheta,-1,legendre_buf)) abort();
+				if (qpms_sph_bessel_array(J, n+nu+2, kdlj.r, bessel_buf)) abort();
+				size_t i = qpms_trans_calculator_index_mnmunu(c, m, n, mu, nu);
+				size_t Qmax = c->B_multipliers[i+1] - c->B_multipliers[i] - 1;
+				assert(Qmax == gaunt_q_max(-m,n+1,mu,nu));
+				complex double Bsum = 0;
+				for(int q = 0; q <= Qmax; ++q) {
+					int p = n+nu-2*q;
+					double Pp_ = legendre_buf[gsl_sf_legendre_array_index(p+1, abs(mu-m))];
+					complex double eimf = cexp(I * kdlj.phi);
+					complex double zp_ = bessel_buf[p+1];
+					complex double multiplier = c->B_multipliers[i][q];
+					Bsum += Pp_ * zp_ * multiplier;
+				}
+
+				size_t qmax = c->A_multipliers[i+1] - c->A_multipliers[i] - 1;
+				assert(qmax == gaunt_q_max(-m,n,mu,nu));
+				complex double Asum = 0;
+				for(size_t q = 0; q <= qmax; ++q) {
+					int p = n+nu-2*q;
+					double Pp = legendre_buf[gsl_sf_legendre_array_index(p, abs(mu-m))];
+					complex double zp = bessel_buf[p];
+					complex double multiplier = c->A_multipliers[i][q];
+					Asum += Pp * zp *  multiplier;
+				}
+				complex double eimf =  cexp(I*(mu-m)*kdlj.phi);
+
+				*Adest = Asum * eimf;
+				*Bdest = Bsum * eimf;
+				return 0;
+			}
+			break;
+		default:
+			abort();
+	}
+	assert(0);
+}
+
+
 complex double qpms_trans_calculator_get_A(const qpms_trans_calculator *c,
                 int m, int n, int mu, int nu, sph_t kdlj,
                 bool r_ge_d, qpms_bessel_t J) {
@@ -498,6 +548,16 @@ complex double qpms_trans_calculator_get_B(const qpms_trans_calculator *c,
 			bes,leg);
 }
 
+int qpms_trans_calculator_get_AB_p(const qpms_trans_calculator *c,
+		complex double *Adest, complex double *Bdest,
+                int m, int n, int mu, int nu, sph_t kdlj,
+                bool r_ge_d, qpms_bessel_t J) {
+	double leg[gsl_sf_legendre_array_n(n+nu+1)];
+	complex double bes[n+nu+2];
+	return qpms_trans_calculator_get_AB_buf_p(c,Adest, Bdest,m,n,mu,nu,kdlj,r_ge_d,J,
+			bes,leg);
+}
+
 complex double qpms_trans_calculator_get_A_ext(const qpms_trans_calculator *c,
                 int m, int n, int mu, int nu,
 	       	double kdlj_r, double kdlj_theta, double kdlj_phi,
@@ -514,15 +574,12 @@ complex double qpms_trans_calculator_get_B_ext(const qpms_trans_calculator *c,
 	return qpms_trans_calculator_get_B(c,m,n,mu,nu,kdlj,r_ge_d,J);
 }
 
-#if 0
+int qpms_trans_calculator_get_AB_p_ext(const qpms_trans_calculator *c,
-int qpms_trans_calculator_get_AB_p(const qpms_trans_calculator *c,
+		complex double *Adest, complex double *Bdest,
-                complex double *Adest, complex double *Bdest,
+                int m, int n, int mu, int nu,
-                int m, int n, int mu, int nu, sph_t kdlj,
+	       	double kdlj_r, double kdlj_theta, double kdlj_phi,
-                bool r_ge_d, qpms_bessel_t J) {
+                int r_ge_d, int J) {
-	double leg[gsl_sf_legendre_array_n(n+nu+1)];
+	sph_t kdlj = {kdlj_r, kdlj_theta, kdlj_phi};
-	complex double bes[n+nu+2];
+	return qpms_trans_calculator_get_AB_p(c,Adest,Bdest,m,n,mu,nu,kdlj,r_ge_d,J);
-	return qpms_trans_calculator_get_AB_buf(c,Adest,Bdest,
-			m,n,mu,nu,kdlj,r_ge_d,J,
-			bes,leg);
 }
-#endif
+

qpms/translations.h

@@ -61,12 +61,14 @@ complex double qpms_trans_calculator_get_B_ext(const qpms_trans_calculator *c,
 		double kdlj_th, double kdlj_phi, int r_ge_d, int J);
 
 
-#if 0
 int qpms_trans_calculator_get_AB_p(const qpms_trans_calculator *c,
 		complex double *Adest, complex double *Bdest,
 		int m, int n, int mu, int nu, sph_t kdlj,
 		bool r_ge_d, qpms_bessel_t J);
-#endif
+int qpms_trans_calculator_get_AB_p_ext(const qpms_trans_calculator *c,
+		complex double *Adest, complex double *Bdest,
+		int m, int n, int mu, int nu, double kdlj_r,
+		double kdlj_th, double kdlj_phi, int r_ge_d, int J);
 
 
 #endif // QPMS_TRANSLATIONS_H