1# distutils: sources = FIXED_INTERP
2# distutils: include_dirs = LIB_DIR
3# distutils: libraries = STD_LIBS
4# distutils: language = c++
5# distutils: extra_compile_args = CPP14_FLAG
6# distutils: extra_link_args = CPP14_FLAG
7"""
8Image sampler definitions
9
10
11
12"""
13
14
15import numpy as np
16
17cimport cython
18cimport numpy as np
19from libc.stdlib cimport abs, calloc, free, malloc
20
21from .fixed_interpolator cimport offset_interpolate
22
23
24cdef class PartitionedGrid:
25
26    @cython.boundscheck(False)
27    @cython.wraparound(False)
28    @cython.cdivision(True)
29    def __cinit__(self,
30                  int parent_grid_id, data,
31                  mask,
32                  np.ndarray[np.float64_t, ndim=1] left_edge,
33                  np.ndarray[np.float64_t, ndim=1] right_edge,
34                  np.ndarray[np.int64_t, ndim=1] dims,
35                  int n_fields = -1):
36        # The data is likely brought in via a slice, so we copy it
37        cdef np.ndarray[np.float64_t, ndim=3] tdata
38        cdef np.ndarray[np.uint8_t, ndim=3] mask_data
39        self.container = NULL
40        self.parent_grid_id = parent_grid_id
41        self.LeftEdge = left_edge
42        self.RightEdge = right_edge
43        self.container = <VolumeContainer *> \
44            malloc(sizeof(VolumeContainer))
45        cdef VolumeContainer *c = self.container # convenience
46        if n_fields == -1:
47            n_fields = len(data)
48        cdef int n_data = len(data)
49
50        c.n_fields = n_fields
51        for i in range(3):
52            c.left_edge[i] = left_edge[i]
53            c.right_edge[i] = right_edge[i]
54            c.dims[i] = dims[i]
55            c.dds[i] = (c.right_edge[i] - c.left_edge[i])/dims[i]
56            c.idds[i] = 1.0/c.dds[i]
57        self.my_data = data
58        self.source_mask = mask
59        mask_data = mask
60        c.data = <np.float64_t **> malloc(sizeof(np.float64_t*) * n_fields)
61        for i in range(n_data):
62            tdata = data[i]
63            c.data[i] = <np.float64_t *> tdata.data
64        c.mask = <np.uint8_t *> mask_data.data
65
66    def __dealloc__(self):
67        # The data fields are not owned by the container, they are owned by us!
68        # So we don't need to deallocate them.
69        if self.container == NULL: return
70        if self.container.data != NULL: free(self.container.data)
71        free(self.container)
72
73    @cython.boundscheck(False)
74    @cython.wraparound(False)
75    @cython.cdivision(True)
76    def integrate_streamline(self, pos, np.float64_t h, mag):
77        cdef np.float64_t cmag[1]
78        cdef np.float64_t k1[3]
79        cdef np.float64_t k2[3]
80        cdef np.float64_t k3[3]
81        cdef np.float64_t k4[3]
82        cdef np.float64_t newpos[3]
83        cdef np.float64_t oldpos[3]
84        for i in range(3):
85            newpos[i] = oldpos[i] = pos[i]
86        self.get_vector_field(newpos, k1, cmag)
87        for i in range(3):
88            newpos[i] = oldpos[i] + 0.5*k1[i]*h
89
90        if not (self.LeftEdge[0] < newpos[0] and newpos[0] < self.RightEdge[0] and \
91                self.LeftEdge[1] < newpos[1] and newpos[1] < self.RightEdge[1] and \
92                self.LeftEdge[2] < newpos[2] and newpos[2] < self.RightEdge[2]):
93            if mag is not None:
94                mag[0] = cmag[0]
95            for i in range(3):
96                pos[i] = newpos[i]
97            return
98
99        self.get_vector_field(newpos, k2, cmag)
100        for i in range(3):
101            newpos[i] = oldpos[i] + 0.5*k2[i]*h
102
103        if not (self.LeftEdge[0] <= newpos[0] and newpos[0] <= self.RightEdge[0] and \
104                self.LeftEdge[1] <= newpos[1] and newpos[1] <= self.RightEdge[1] and \
105                self.LeftEdge[2] <= newpos[2] and newpos[2] <= self.RightEdge[2]):
106            if mag is not None:
107                mag[0] = cmag[0]
108            for i in range(3):
109                pos[i] = newpos[i]
110            return
111
112        self.get_vector_field(newpos, k3, cmag)
113        for i in range(3):
114            newpos[i] = oldpos[i] + k3[i]*h
115
116        if not (self.LeftEdge[0] <= newpos[0] and newpos[0] <= self.RightEdge[0] and \
117                self.LeftEdge[1] <= newpos[1] and newpos[1] <= self.RightEdge[1] and \
118                self.LeftEdge[2] <= newpos[2] and newpos[2] <= self.RightEdge[2]):
119            if mag is not None:
120                mag[0] = cmag[0]
121            for i in range(3):
122                pos[i] = newpos[i]
123            return
124
125        self.get_vector_field(newpos, k4, cmag)
126
127        for i in range(3):
128            pos[i] = oldpos[i] + h*(k1[i]/6.0 + k2[i]/3.0 + k3[i]/3.0 + k4[i]/6.0)
129
130        if mag is not None:
131            for i in range(3):
132                newpos[i] = pos[i]
133            self.get_vector_field(newpos, k4, cmag)
134            mag[0] = cmag[0]
135
136    @cython.boundscheck(False)
137    @cython.wraparound(False)
138    @cython.cdivision(True)
139    cdef void get_vector_field(self, np.float64_t pos[3],
140                               np.float64_t *vel, np.float64_t *vel_mag):
141        cdef np.float64_t dp[3]
142        cdef int ci[3]
143        cdef VolumeContainer *c = self.container # convenience
144
145        for i in range(3):
146            ci[i] = (int)((pos[i]-self.LeftEdge[i])/c.dds[i])
147            dp[i] = (pos[i] - ci[i]*c.dds[i] - self.LeftEdge[i])/c.dds[i]
148
149        cdef int offset = ci[0] * (c.dims[1] + 1) * (c.dims[2] + 1) \
150                          + ci[1] * (c.dims[2] + 1) + ci[2]
151
152        vel_mag[0] = 0.0
153        for i in range(3):
154            vel[i] = offset_interpolate(c.dims, dp, c.data[i] + offset)
155            vel_mag[0] += vel[i]*vel[i]
156        vel_mag[0] = np.sqrt(vel_mag[0])
157        if vel_mag[0] != 0.0:
158            for i in range(3):
159                vel[i] /= vel_mag[0]
160