1from __future__ import division, absolute_import
2
3cimport cython
4cimport numpy as cnp
5
6import os
7import numpy as np
8
9# Parameters are linked to the direction numbers list.
10# See `initialize_direction_numbers` for more details.
11# Declared using DEF to be known at compilation time for ``poly`` et ``vinit``
12DEF MAXDIM = 21201  # max number of dimensions
13DEF MAXDEG = 18  # max polynomial degree
14DEF MAXBIT = 30  # max number of bits
15
16
17# Needed to be accessed with python
18cdef extern from *:
19    """
20    int MAXDIM_DEFINE = 21201;
21    int MAXDEG_DEFINE = 18;
22    int MAXBIT_DEFINE = 30;
23    """
24    int MAXDIM_DEFINE  # max number of dimensions
25    int MAXDEG_DEFINE  # max polynomial degree
26    int MAXBIT_DEFINE  # max number of bits
27
28_MAXDIM = MAXDIM_DEFINE
29_MAXDEG = MAXDEG_DEFINE
30_MAXBIT = MAXBIT_DEFINE
31
32cdef int poly[MAXDIM]
33cdef int vinit[MAXDIM][MAXDEG]
34
35cdef int LARGEST_NUMBER = 2 ** MAXBIT  # largest possible integer
36cdef float RECIPD = 1.0 / LARGEST_NUMBER  # normalization constant
37
38cdef bint is_initialized = False
39
40
41def initialize_direction_numbers():
42    """Load direction numbers.
43
44    Direction numbers obtained using the search criterion D(6)
45    up to the dimension 21201. This is the recommended choice by the authors.
46
47    Original data can be found at https://web.maths.unsw.edu.au/~fkuo/sobol/.
48    For additional details on the quantities involved, see [1].
49
50    [1] S. Joe and F. Y. Kuo. Remark on algorithm 659: Implementing sobol's
51        quasirandom sequence generator. ACM Trans. Math. Softw., 29(1):49-57,
52        Mar. 2003.
53
54    The C-code generated from putting the numbers in as literals is obscenely
55    large/inefficient. The data file was thus packaged and save as an .npz data
56    file for fast loading using the following code (this assumes that the file
57    https://web.maths.unsw.edu.au/~fkuo/sobol/new-joe-kuo-6.21201 is present in
58    the working directory):
59
60        import pandas as pd
61        import numpy as np
62
63        # read in file content
64        with open("./new-joe-kuo-6.21201", "r") as f:
65            lines = f.readlines()
66
67        rows = []
68
69        # parse data from file line by line
70        for l in lines[1:]:
71            nums = [int(n) for n in l.replace(" \n", "").split()]
72            d, s, a = nums[:3]
73            vs = {f"v{i}": int(v) for i,v in enumerate(nums[3:])}
74            rows.append({"d": d, "s": s, "a": a, **vs})
75
76
77        # read in as dataframe, explicitly use zero values
78        df = pd.DataFrame(rows).fillna(0).astype(int)
79
80        # peform conversion
81        df["poly"] = 2 * df["a"] + 2 ** df["s"] + 1
82
83        # ensure columns are properly ordered
84        vs = df[[f"v{i}" for i in range(18)]].values
85
86        # add the degenerate d=1 column (not included in the data file)
87        vs = np.vstack([vs[0][np.newaxis, :], vs])
88        poly = np.concatenate([[1], df["poly"].values])
89
90        # save as compressed .npz file to minimize size of distribution
91        np.savez_compressed("./_sobol_direction_numbers", vinit=vs, poly=poly)
92
93    """
94    cdef int[:] dns_poly
95    cdef int[:, :] dns_vinit
96
97    global is_initialized
98    if not is_initialized:
99        dns = np.load(os.path.join(os.path.dirname(__file__), "_sobol_direction_numbers.npz"))
100        dns_poly = dns["poly"].astype(np.intc)
101        dns_vinit = dns["vinit"].astype(np.intc)
102        for i in range(MAXDIM):
103            poly[i] = dns_poly[i]
104        for i in range(MAXDIM):
105            for j in range(MAXDEG):
106                vinit[i][j] = dns_vinit[i, j]
107        is_initialized = True
108
109
110@cython.boundscheck(False)
111@cython.wraparound(False)
112cdef int bit_length(const int n):
113    cdef int bits = 0
114    cdef int nloc = n
115    while nloc != 0:
116        nloc >>= 1
117        bits += 1
118    return bits
119
120
121@cython.boundscheck(False)
122@cython.wraparound(False)
123cdef int low_0_bit(const int x) nogil:
124    """Get the position of the right-most 0 bit for an integer.
125
126    Examples:
127    >>> low_0_bit(0)
128    1
129    >>> low_0_bit(1)
130    2
131    >>> low_0_bit(2)
132    1
133    >>> low_0_bit(5)
134    2
135    >>> low_0_bit(7)
136    4
137
138    Parameters
139    ----------
140    x: int
141        An integer.
142
143    Returns
144    -------
145    position: int
146        Position of the right-most 0 bit.
147
148    """
149    cdef int i = 0
150    while x & (1 << i) != 0:
151        i += 1
152    return i + 1
153
154
155@cython.boundscheck(False)
156@cython.wraparound(False)
157cdef int ibits(const int x, const int pos, const int length) nogil:
158    """Extract a sequence of bits from the bit representation of an integer.
159
160    Extract the sequence from position `pos` (inclusive) to ``pos + length``
161    (not inclusive), leftwise.
162
163    Examples:
164    >>> ibits(1, 0, 1)
165    1
166    >>> ibits(1, 1, 1)
167    0
168    >>> ibits(2, 0, 1)
169    0
170    >>> ibits(2, 0, 2)
171    2
172    >>> ibits(25, 1, 5)
173    12
174
175
176    Parameters
177    ----------
178    x: int
179        Integer to convert to bit representation.
180    pos: int
181        Starting position of sequence in bit representation of integer.
182    length: int
183        Length of sequence (number of bits).
184
185    Returns
186    -------
187    ibits: int
188        Integer value corresponding to bit sequence.
189
190    """
191    return (x >> pos) & ((1 << length) - 1)
192
193
194@cython.boundscheck(False)
195@cython.wraparound(False)
196cpdef void initialize_v(cnp.int_t[:, :] v, const int dim):
197    cdef int d, i, j, k, m, p, newv, pow2
198
199    if dim == 0:
200        return
201
202    # first row of v is all 1s
203    for i in range(MAXBIT):
204        v[0, i] = 1
205
206    # Remaining rows of v (row 2 through dim, indexed by [1:dim])
207    for d in range(1, dim):
208        p = poly[d]
209        m = bit_length(p) - 1
210
211        # First m elements of row d comes from vinit
212        for j in range(m):
213            v[d, j] = vinit[d][j]
214
215        # Fill in remaining elements of v as in Section 2 (top of pg. 90) of:
216        #
217        # P. Bratley and B. L. Fox. Algorithm 659: Implementing sobol's
218        # quasirandom sequence generator. ACM Trans.
219        # Math. Softw., 14(1):88-100, Mar. 1988.
220        #
221        for j in range(m, MAXBIT):
222            newv = v[d, j - m]
223            pow2 = 1
224            for k in range(m):
225                pow2 = pow2 << 1
226                if (p >> (m - 1 - k)) & 1:
227                    newv = newv ^ (pow2 * v[d, j - k - 1])
228            v[d, j] = newv
229
230    # Multiply each column of v by power of 2:
231    # v * [2^(maxbit-1), 2^(maxbit-2),..., 2, 1]
232    pow2 = 1
233    for d in range(MAXBIT):
234        for i in range(dim):
235            v[i, MAXBIT - 1 - d] *= pow2
236        pow2 = pow2 << 1
237
238
239@cython.boundscheck(False)
240@cython.wraparound(False)
241cpdef void _draw(const int n,
242                 const int num_gen,
243                 const int dim,
244                 cnp.int_t[:, :] sv,
245                 cnp.int_t[:] quasi,
246                 cnp.float_t[:, :] result) nogil:
247    cdef int i, j, l, qtmp
248    cdef int num_gen_loc = num_gen
249    for i in range(n):
250        l = low_0_bit(num_gen_loc)
251        for j in range(dim):
252            qtmp = quasi[j] ^ sv[j, l - 1]
253            quasi[j] = qtmp
254            result[i, j] = qtmp * RECIPD
255        num_gen_loc += 1
256
257
258@cython.boundscheck(False)
259@cython.wraparound(False)
260cpdef void _fast_forward(const int n,
261                         const int num_gen,
262                         const int dim,
263                         cnp.int_t[:, :] sv,
264                         cnp.int_t[:] quasi) nogil:
265    cdef int i, j, l
266    cdef int num_gen_loc = num_gen
267    for i in range(n):
268        l = low_0_bit(num_gen_loc)
269        for j in range(dim):
270            quasi[j] = quasi[j] ^ sv[j, l - 1]
271        num_gen_loc += 1
272
273
274@cython.boundscheck(False)
275@cython.wraparound(False)
276cdef int cdot_pow2(cnp.int_t[:] a) nogil:
277    cdef int i
278    cdef int size = a.shape[0]
279    cdef int z = 0
280    cdef int pow2 = 1
281    for i in range(size):
282        z += a[size - 1 - i] * pow2
283        pow2 *= 2
284    return z
285
286
287@cython.boundscheck(False)
288@cython.wraparound(False)
289cpdef void _cscramble(const int dim,
290                      cnp.int_t[:, :, :] ltm,
291                      cnp.int_t[:, :] sv) nogil:
292    cdef int d, i, j, k, l, lsm, lsmdp, p, t1, t2, vdj
293
294    # Set diagonals of maxbit x maxbit arrays to 1
295    for d in range(dim):
296        for i in range(MAXBIT):
297            ltm[d, i, i] = 1
298
299    for d in range(dim):
300        for j in range(MAXBIT):
301            vdj = sv[d, j]
302            l = 1
303            t2 = 0
304            for p in range(MAXBIT - 1, -1, -1):
305                lsmdp = cdot_pow2(ltm[d, p, :])
306                t1 = 0
307                for k in range(MAXBIT):
308                    t1 += ibits(lsmdp, k, 1) * ibits(vdj, k, 1)
309                t1 = t1 % 2
310                t2 = t2 + t1 * l
311                l = 2 * l
312            sv[d, j] = t2
313
314
315@cython.boundscheck(False)
316@cython.wraparound(False)
317cpdef void _fill_p_cumulative(cnp.float_t[:] p,
318                              cnp.float_t[:] p_cumulative) nogil:
319    cdef int i
320    cdef int len_p = p.shape[0]
321    cdef float tot = 0
322    cdef float t
323    for i in range(len_p):
324        t = tot + p[i]
325        p_cumulative[i] = t
326        tot = t
327
328
329@cython.boundscheck(False)
330@cython.wraparound(False)
331cpdef void _categorize(cnp.float_t[:] draws,
332                       cnp.float_t[:] p_cumulative,
333                       cnp.int_t[:] result) nogil:
334    cdef int i
335    cdef int n_p = p_cumulative.shape[0]
336    for i in range(draws.shape[0]):
337        j = _find_index(p_cumulative, n_p, draws[i])
338        result[j] = result[j] + 1
339
340
341@cython.boundscheck(False)
342@cython.wraparound(False)
343cdef int _find_index(cnp.float_t[:] p_cumulative,
344                     const int size,
345                     const float value) nogil:
346    cdef int l = 0
347    cdef int r = size - 1
348    cdef int m
349    while r > l:
350        m = (l + r) // 2
351        if value > p_cumulative[m]:
352            l = m + 1
353        else:
354            r = m
355    return r
356
357
358def _test_find_index(p_cumulative, size, value):
359    # type: (np.ndarray, int, float) -> int
360    """Wrapper for testing in python"""
361    return _find_index(p_cumulative, size, value)
362