algos_take_helper.pxi.in 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222
  1. """
  2. Template for each `dtype` helper function for take
  3. WARNING: DO NOT edit .pxi FILE directly, .pxi is generated from .pxi.in
  4. """
  5. # ----------------------------------------------------------------------
  6. # take_1d, take_2d
  7. # ----------------------------------------------------------------------
  8. {{py:
  9. # c_type_in, c_type_out
  10. dtypes = [
  11. ('uint8_t', 'uint8_t'),
  12. ('uint8_t', 'object'),
  13. ('int8_t', 'int8_t'),
  14. ('int8_t', 'int32_t'),
  15. ('int8_t', 'int64_t'),
  16. ('int8_t', 'float64_t'),
  17. ('int16_t', 'int16_t'),
  18. ('int16_t', 'int32_t'),
  19. ('int16_t', 'int64_t'),
  20. ('int16_t', 'float64_t'),
  21. ('int32_t', 'int32_t'),
  22. ('int32_t', 'int64_t'),
  23. ('int32_t', 'float64_t'),
  24. ('int64_t', 'int64_t'),
  25. ('int64_t', 'float64_t'),
  26. ('float32_t', 'float32_t'),
  27. ('float32_t', 'float64_t'),
  28. ('float64_t', 'float64_t'),
  29. ('object', 'object'),
  30. ]
  31. def get_dispatch(dtypes):
  32. for (c_type_in, c_type_out) in dtypes:
  33. def get_name(dtype_name):
  34. if dtype_name == "object":
  35. return "object"
  36. if dtype_name == "uint8_t":
  37. return "bool"
  38. return dtype_name[:-2]
  39. name = get_name(c_type_in)
  40. dest = get_name(c_type_out)
  41. args = dict(name=name, dest=dest, c_type_in=c_type_in,
  42. c_type_out=c_type_out)
  43. yield (name, dest, c_type_in, c_type_out)
  44. }}
  45. {{for name, dest, c_type_in, c_type_out in get_dispatch(dtypes)}}
  46. @cython.wraparound(False)
  47. @cython.boundscheck(False)
  48. {{if c_type_in != "object"}}
  49. def take_1d_{{name}}_{{dest}}(const {{c_type_in}}[:] values,
  50. {{else}}
  51. def take_1d_{{name}}_{{dest}}(ndarray[{{c_type_in}}, ndim=1] values,
  52. {{endif}}
  53. const intp_t[:] indexer,
  54. {{c_type_out}}[:] out,
  55. fill_value=np.nan):
  56. cdef:
  57. Py_ssize_t i, n, idx
  58. {{c_type_out}} fv
  59. n = indexer.shape[0]
  60. fv = fill_value
  61. {{if c_type_out != "object"}}
  62. with nogil:
  63. {{else}}
  64. if True:
  65. {{endif}}
  66. for i in range(n):
  67. idx = indexer[i]
  68. if idx == -1:
  69. out[i] = fv
  70. else:
  71. {{if c_type_in == "uint8_t" and c_type_out == "object"}}
  72. out[i] = True if values[idx] > 0 else False
  73. {{else}}
  74. out[i] = values[idx]
  75. {{endif}}
  76. @cython.wraparound(False)
  77. @cython.boundscheck(False)
  78. {{if c_type_in != "object"}}
  79. def take_2d_axis0_{{name}}_{{dest}}(const {{c_type_in}}[:, :] values,
  80. {{else}}
  81. def take_2d_axis0_{{name}}_{{dest}}(ndarray[{{c_type_in}}, ndim=2] values,
  82. {{endif}}
  83. ndarray[intp_t, ndim=1] indexer,
  84. {{c_type_out}}[:, :] out,
  85. fill_value=np.nan):
  86. cdef:
  87. Py_ssize_t i, j, k, n, idx
  88. {{c_type_out}} fv
  89. {{if c_type_in == c_type_out != "object"}}
  90. const {{c_type_out}} *v
  91. {{c_type_out}} *o
  92. {{endif}}
  93. n = len(indexer)
  94. k = values.shape[1]
  95. fv = fill_value
  96. {{if c_type_in == c_type_out != "object"}}
  97. # GH#3130
  98. if (values.strides[1] == out.strides[1] and
  99. values.strides[1] == sizeof({{c_type_out}}) and
  100. sizeof({{c_type_out}}) * n >= 256):
  101. for i in range(n):
  102. idx = indexer[i]
  103. if idx == -1:
  104. for j in range(k):
  105. out[i, j] = fv
  106. else:
  107. v = &values[idx, 0]
  108. o = &out[i, 0]
  109. memmove(o, v, <size_t>(sizeof({{c_type_out}}) * k))
  110. return
  111. {{endif}}
  112. for i in range(n):
  113. idx = indexer[i]
  114. if idx == -1:
  115. for j in range(k):
  116. out[i, j] = fv
  117. else:
  118. for j in range(k):
  119. {{if c_type_in == "uint8_t" and c_type_out == "object"}}
  120. out[i, j] = True if values[idx, j] > 0 else False
  121. {{else}}
  122. out[i, j] = values[idx, j]
  123. {{endif}}
  124. @cython.wraparound(False)
  125. @cython.boundscheck(False)
  126. {{if c_type_in != "object"}}
  127. def take_2d_axis1_{{name}}_{{dest}}(const {{c_type_in}}[:, :] values,
  128. {{else}}
  129. def take_2d_axis1_{{name}}_{{dest}}(ndarray[{{c_type_in}}, ndim=2] values,
  130. {{endif}}
  131. ndarray[intp_t, ndim=1] indexer,
  132. {{c_type_out}}[:, :] out,
  133. fill_value=np.nan):
  134. cdef:
  135. Py_ssize_t i, j, k, n, idx
  136. {{c_type_out}} fv
  137. n = len(values)
  138. k = len(indexer)
  139. if n == 0 or k == 0:
  140. return
  141. fv = fill_value
  142. for i in range(n):
  143. for j in range(k):
  144. idx = indexer[j]
  145. if idx == -1:
  146. out[i, j] = fv
  147. else:
  148. {{if c_type_in == "uint8_t" and c_type_out == "object"}}
  149. out[i, j] = True if values[i, idx] > 0 else False
  150. {{else}}
  151. out[i, j] = values[i, idx]
  152. {{endif}}
  153. @cython.wraparound(False)
  154. @cython.boundscheck(False)
  155. def take_2d_multi_{{name}}_{{dest}}(ndarray[{{c_type_in}}, ndim=2] values,
  156. indexer,
  157. ndarray[{{c_type_out}}, ndim=2] out,
  158. fill_value=np.nan):
  159. cdef:
  160. Py_ssize_t i, j, k, n, idx
  161. ndarray[intp_t, ndim=1] idx0 = indexer[0]
  162. ndarray[intp_t, ndim=1] idx1 = indexer[1]
  163. {{c_type_out}} fv
  164. n = len(idx0)
  165. k = len(idx1)
  166. fv = fill_value
  167. for i in range(n):
  168. idx = idx0[i]
  169. if idx == -1:
  170. for j in range(k):
  171. out[i, j] = fv
  172. else:
  173. for j in range(k):
  174. if idx1[j] == -1:
  175. out[i, j] = fv
  176. else:
  177. {{if c_type_in == "uint8_t" and c_type_out == "object"}}
  178. out[i, j] = True if values[idx, idx1[j]] > 0 else False
  179. {{else}}
  180. out[i, j] = values[idx, idx1[j]]
  181. {{endif}}
  182. {{endfor}}