1cimport numpy as np
2cimport cython
3
4ctypedef fused DTYPE_floating_t:
5    float
6    float complex
7    double
8    double complex
9    long double
10    long double complex
11
12ctypedef fused DTYPE_t:
13    DTYPE_floating_t
14    object
15
16
17# Once Cython 3.0 is out, we can just do the following below:
18#
19#     with nogil(DTYPE_t is not object):
20#
21# But until then, we'll neeed two copies of the loops, one with
22# nogil and another with gil.
23
24@cython.cdivision(True)
25@cython.boundscheck(False)
26@cython.wraparound(False)
27cdef void _sosfilt_float(DTYPE_floating_t [:, ::1] sos,
28                         DTYPE_floating_t [:, ::1] x,
29                         DTYPE_floating_t [:, :, ::1] zi) nogil:
30    # Modifies x and zi in place
31    cdef Py_ssize_t n_signals = x.shape[0]
32    cdef Py_ssize_t n_samples = x.shape[1]
33    cdef Py_ssize_t n_sections = sos.shape[0]
34    cdef Py_ssize_t i, n, s
35    cdef DTYPE_floating_t x_new, x_cur
36    cdef DTYPE_floating_t[:, ::1] zi_slice
37
38    # jumping through a few memoryview hoops to reduce array lookups,
39    # the original version is still in the gil version below.
40    for i in xrange(n_signals):
41        zi_slice = zi[i, :, :]
42        for n in xrange(n_samples):
43
44            x_cur = x[i, n]
45
46            for s in xrange(n_sections):
47                x_new = sos[s, 0] * x_cur + zi_slice[s, 0]
48                zi_slice[s, 0] = (sos[s, 1] * x_cur - sos[s, 4] * x_new
49                                  + zi_slice[s, 1])
50                zi_slice[s, 1] = sos[s, 2] * x_cur - sos[s, 5] * x_new
51                x_cur = x_new
52
53            x[i, n] = x_cur
54
55
56@cython.cdivision(True)
57@cython.boundscheck(False)
58@cython.wraparound(False)
59def _sosfilt_object(object [:, ::1] sos,
60                    object [:, ::1] x,
61                    object [:, :, ::1] zi):
62    # Modifies x and zi in place
63    cdef Py_ssize_t n_signals = x.shape[0]
64    cdef Py_ssize_t n_samples = x.shape[1]
65    cdef Py_ssize_t n_sections = sos.shape[0]
66    cdef Py_ssize_t i, n, s
67    cdef object x_n
68    for i in xrange(n_signals):
69        for n in xrange(n_samples):
70            for s in xrange(n_sections):
71                x_n = x[i, n]  # make a temporary copy
72                # Use direct II transposed structure:
73                x[i, n] = sos[s, 0] * x_n + zi[i, s, 0]
74                zi[i, s, 0] = (sos[s, 1] * x_n - sos[s, 4] * x[i, n] +
75                               zi[i, s, 1])
76                zi[i, s, 1] = (sos[s, 2] * x_n - sos[s, 5] * x[i, n])
77
78
79def _sosfilt(DTYPE_t [:, ::1] sos,
80             DTYPE_t [:, ::1] x,
81             DTYPE_t [:, :, ::1] zi):
82    if DTYPE_t is object:
83        _sosfilt_object(sos, x, zi)
84    else:
85        with nogil:
86            _sosfilt_float(sos, x, zi)
87