fujiisoup / py3nj

Wigner's 3J, 6J, 9J symbols for python
https://py3nj.readthedocs.io/
Apache License 2.0
18 stars 5 forks source link

Cython + parallel #9

Closed albakalaja closed 4 years ago

albakalaja commented 4 years ago

I'm trying to use the wigner3j within a function in cython, and I also need to parallelize it with the the cython.parallel module. However, the latter requires to release the GIL. Is there anyway to do this?

Here is an example of how I would like to use wigner3j

%load_ext cython
%%cython

import numpy as np
cimport numpy as cnp

import py3nj
from py3nj import wigner

DTYPE = np.double
ITYPE = np.int
ctypedef cnp.double_t DTYPE_t
ctypedef cnp.int_t ITYPE_t

cimport cython
@cython.boundscheck(False) # turn off bounds-checking for entire function
@cython.wraparound(False)  # turn off negative index wrapping for entire function

cpdef double get_gaunt(int l1,# variable for the parallel loop
                       int lmin,
                       int lmax,
                         ):

    # fisher variables
    cdef double gaunt = 0.0 # initialize 
    cdef int l2, l3
    cdef int int_l2, int_l1, l3min_index
    cdef cnp.ndarray[ITYPE_t,ndim = 1] l3_values_w
    cdef cnp.ndarray[DTYPE_t,ndim = 1] wigner3j

    for l2 in range(lmin,l1+1): 

        # -----------------
        int_l2 = int(l2*2)
        int_l1 = int(l1*2)
        l3_values_w, wigner3j = py3nj.wigner._drc3jj(int_l2,int_l1,0,0) 
        l3min_index = max(abs(l1-l2), 0) # needed to index wigner3j
        # -----------------

        for l3 in range(l3min_index,l2+1): 

            if (l1+l2+l3)%2==0 and l3>=lmin: 

                gaunt+=(2.0*l1+1.0)*(2.0*l2+1.0)*(2.0*l3+1.0)/(4.0*np.pi)*wigner3j[l3-l3min_index]*\
                            wigner3j[l3-l3min_index]

    return gaunt
fujiisoup commented 4 years ago

Thanks @albakalaja for raising an issue. I am sorry but have no idea how to use py3nj in cython.

As py3nj runs a fortran code trhough f2py, it should be fast enough. And if it is compiled correctly, it should run in parallel (at least in the original fortran code).

So maybe it would be easier to run them outside cython but make it sure run in parallel.

albakalaja commented 4 years ago

Hi @fujiisoup , thank you for the quick reply! I thought that running it outside would require too much, but I can actually do it and the call it in the main function.