prescribed_roots_pyx.spyx

#clang c
#cinclude $SAGE_LOCAL/include/flint/
#clib gomp
#cargs -fopenmp
#cfile all_roots_in_interval.c
#cfile power_sums.c

from cpython cimport array
import array
from cython.parallel import prange
from libc.stdlib cimport malloc, free, rand
cimport cython

cdef extern from "power_sums.h":
    ctypedef struct ps_static_data_t:
        pass
    ctypedef struct ps_dynamic_data_t:
        long count

    ps_static_data_t *ps_static_init(int d, int lead, int sign, int q,
    		     		     int cofactor, 
                                     int *modlist,
                                     int verbosity, long node_count)
    ps_dynamic_data_t *ps_dynamic_init(int d, int *Q0)
    ps_dynamic_data_t *ps_dynamic_clone(ps_dynamic_data_t *dy_data)
    ps_dynamic_data_t *ps_dynamic_split(ps_dynamic_data_t *dy_data)
    void extract_pol(int *Q, ps_dynamic_data_t *dy_data)
    void extract_symmetrized_pol(int *Q, ps_dynamic_data_t *dy_data)
    void ps_static_clear(ps_static_data_t *st_data)
    void ps_dynamic_clear(ps_dynamic_data_t *dy_data)
    int next_pol(ps_static_data_t *st_data, ps_dynamic_data_t *dy_data) nogil

cdef class process_queue:
    cdef int d, verbosity
    cdef long node_count
    cdef public long count
    cdef public int k
    cdef public array.array Q0_array
    cdef int[:] Q0
    cdef public array.array Qsym_array
    cdef int[:] Qsym
    cdef int[:] modlist
    cdef public array.array modlist_array
    cdef int sign
    cdef int cofactor
    cdef ps_static_data_t *ps_st_data
    cdef ps_dynamic_data_t *ps_dy_data

    def __init__(self, int d, int n, int lead, int sign, int q, int cofactor,
                 modlist, node_count, verbosity, Q):
        cdef int i
        self.d = d
        self.k = d
        self.sign = sign
        self.cofactor = cofactor
        self.Q0_array = array.array('i', [0,] * (d+1))
        self.Q0 = self.Q0_array
        self.Qsym_array = array.array('i', [0,] * (2*d+3))
        self.Qsym = self.Qsym_array
        self.modlist_array = array.array('i', [0,] * (d+1))
        self.modlist = self.modlist_array
        for i in range(d+1):
            self.modlist[i] = modlist[d-i]
            self.Q0[i] = Q[i]
        if verbosity == None:
            self.verbosity = -1
        else:
            self.verbosity = verbosity
        if node_count == None:
            self.node_count = -1
        else:
            self.node_count = node_count
        self.count = 0
        self.ps_st_data = ps_static_init(d, lead, sign, q, cofactor,
                                    self.modlist_array.data.as_ints,
                                         self.verbosity, self.node_count)
        self.ps_dy_data = ps_dynamic_init(d, self.Q0_array.data.as_ints)

    def clear(self):
        ps_static_clear(self.ps_st_data)
        ps_dynamic_clear(self.ps_dy_data)

    cpdef int exhaust_next_answer(self):
        cdef int t
        t = next_pol(self.ps_st_data, self.ps_dy_data)
        extract_symmetrized_pol(self.Qsym_array.data.as_ints, self.ps_dy_data)
        self.count = self.ps_dy_data.count
        return(t)

    cpdef object parallel_exhaust(process_queue self, int num_processes, f=None):
        cdef ps_dynamic_data_t **dy_data_buf
        cdef int i, j, k, m, d = self.d, t=1, np = num_processes
        ans = []
        dy_data_buf = <ps_dynamic_data_t **>malloc(np*cython.sizeof(cython.pointer(ps_dynamic_data_t)))
        dy_data_buf[0] = ps_dynamic_clone(self.ps_dy_data)
        cdef int *res = <int *>malloc(np*sizeof(int))

        for i in range(1, np):
            dy_data_buf[i] = NULL
        k=0
        while (t>0):
            t = 0
            k = k%(np-1) + 1
            with nogil: # Drop GIL for this parallel loop
                for i in prange(np, schedule='dynamic', num_threads=np):
                    if dy_data_buf[i] != NULL:
                        res[i] = next_pol(self.ps_st_data, dy_data_buf[i])
            for i in range(np):
                if dy_data_buf[i] != NULL:
                    if res[i] > 0:
                        t += 1
                        extract_symmetrized_pol(self.Qsym_array.data.as_ints,
                                                dy_data_buf[i])
                        if (f != None):
                            f.write(str(list(self.Qsym_array)))
                            f.write("\n")
                        else: ans.append(list(self.Qsym_array))
                    else:
                        self.count += dy_data_buf[i].count
                        ps_dynamic_clear(dy_data_buf[i])
                        dy_data_buf[i] = NULL
                if dy_data_buf[i] == NULL:
                    j = (i-k) % np
                    dy_data_buf[i] = ps_dynamic_split(dy_data_buf[j])
        free(dy_data_buf)
        free(res)
        if (f != None): return None
        else: return(ans)