Skip to content

Commit b2e5ac3

Browse files
lpsingerlpsinger
committed
Support arrays for nside
Fixes astropy#66.
1 parent e330968 commit b2e5ac3

File tree

3 files changed

+78
-54
lines changed

3 files changed

+78
-54
lines changed

‎astropy_healpix/core.py

+24-15
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,7 @@ def healpix_to_lonlat(healpix_index, nside, dx=None, dy=None, order='ring'):
267267
raise ValueError('Either both or neither dx and dy must be specified')
268268

269269
healpix_index = np.asarray(healpix_index, dtype=np.int64)
270+
nside = np.asarray(nside, dtype=np.int64)
270271

271272
if dx is None and dy is not None:
272273
dx = 0.5
@@ -282,14 +283,16 @@ def healpix_to_lonlat(healpix_index, nside, dx=None, dy=None, order='ring'):
282283
dx = dx.ravel()
283284
dy = dy.ravel()
284285

285-
shape = healpix_index.shape
286-
healpix_index = healpix_index.ravel()
287-
nside = int(nside)
288-
289286
_validate_healpix_index('healpix_index', healpix_index, nside)
290287
_validate_nside(nside)
291288
order = _validate_order(order)
292289

290+
healpix_index, nside = np.broadcast_arrays(healpix_index, nside)
291+
292+
shape = healpix_index.shape
293+
healpix_index = healpix_index.ravel()
294+
nside = nside.ravel()
295+
293296
if dx is None:
294297
lon, lat = core_cython.healpix_to_lonlat(healpix_index, nside, order)
295298
else:
@@ -331,14 +334,14 @@ def lonlat_to_healpix(lon, lat, nside, return_offsets=False, order='ring'):
331334
lon = lon.to(u.rad).value
332335
lat = lat.to(u.rad).value
333336

334-
lon, lat = np.broadcast_arrays(lon, lat)
337+
lon, lat, nside = np.broadcast_arrays(lon, lat, nside)
335338

336339
shape = np.shape(lon)
337340

338341
lon = lon.astype(float).ravel()
339342
lat = lat.astype(float).ravel()
343+
nside = nside.astype(np.int64).ravel()
340344

341-
nside = int(nside)
342345
_validate_nside(nside)
343346
order = _validate_order(order)
344347

@@ -368,9 +371,11 @@ def nested_to_ring(nested_index, nside):
368371
"""
369372

370373
nested_index = np.asarray(nested_index, dtype=np.int64)
374+
nside = np.asarray(nside, dtype=np.int64)
375+
nested_index, nside = np.broadcast_arrays(nested_index, nside)
371376
shape = nested_index.shape
372377
nested_index = nested_index.ravel()
373-
nside = int(nside)
378+
nside = nside.ravel()
374379

375380
_validate_healpix_index('nested_index', nested_index, nside)
376381
_validate_nside(nside)
@@ -397,9 +402,11 @@ def ring_to_nested(ring_index, nside):
397402
"""
398403

399404
ring_index = np.asarray(ring_index, dtype=np.int64)
405+
nside = np.asarray(nside, dtype=np.int64)
406+
ring_index, nside = np.broadcast_arrays(ring_index, nside)
400407
shape = ring_index.shape
401408
ring_index = ring_index.ravel()
402-
nside = int(nside)
409+
nside = nside.ravel()
403410

404411
_validate_healpix_index('ring_index', ring_index, nside)
405412
_validate_nside(nside)
@@ -436,13 +443,13 @@ def bilinear_interpolation_weights(lon, lat, nside, order='ring'):
436443
lon = lon.to(u.rad).value
437444
lat = lat.to(u.rad).value
438445

439-
lon, lat = np.broadcast_arrays(lon, lat)
446+
lon, lat, nside = np.broadcast_arrays(lon, lat, nside)
440447

441448
shape = np.shape(lon)
442449

443450
lon = lon.astype(float).ravel()
444451
lat = lat.astype(float).ravel()
445-
nside = int(nside)
452+
nside = nside.astype(np.int64).ravel()
446453

447454
order = _validate_order(order)
448455
_validate_nside(nside)
@@ -509,16 +516,18 @@ def neighbours(healpix_index, nside, order='ring'):
509516
"""
510517

511518
healpix_index = np.asarray(healpix_index, dtype=np.int64)
512-
nside = int(nside)
513-
514-
shape = np.shape(healpix_index)
515-
516-
healpix_index = healpix_index.ravel()
519+
nside = np.asarray(nside, dtype=np.int64)
517520

518521
_validate_healpix_index('healpix_index', healpix_index, nside)
519522
_validate_nside(nside)
520523
order = _validate_order(order)
521524

525+
healpix_index, nside = np.broadcast_arrays(healpix_index, nside)
526+
shape = np.shape(healpix_index)
527+
528+
healpix_index = healpix_index.ravel()
529+
nside = nside.ravel()
530+
522531
neigh = core_cython.neighbours(healpix_index, nside, order)
523532
return _restore_shape(neigh, shape=(8,) + shape)
524533

‎astropy_healpix/core_cython.pyx

+42-35
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,8 @@ def _validate_order(str order):
3434

3535
@cython.boundscheck(False)
3636
def healpix_to_lonlat(np.ndarray[int64_t, ndim=1, mode="c"] healpix_index,
37-
int nside, str order):
37+
np.ndarray[int64_t, ndim=1, mode="c"] nside,
38+
str order):
3839
"""
3940
Convert HEALPix indices to longitudes/latitudes.
4041
@@ -46,7 +47,7 @@ def healpix_to_lonlat(np.ndarray[int64_t, ndim=1, mode="c"] healpix_index,
4647
----------
4748
healpix_index : `~numpy.ndarray`
4849
1-D array of HEALPix indices
49-
nside : int
50+
nside : `~numpy.ndarray`
5051
Number of pixels along the side of each of the 12 top-level HEALPix tiles
5152
order : { 'nested' | 'ring' }
5253
Order of HEALPix pixels
@@ -71,12 +72,12 @@ def healpix_to_lonlat(np.ndarray[int64_t, ndim=1, mode="c"] healpix_index,
7172

7273
if order == 'nested':
7374
for i in prange(n, nogil=True, schedule='static'):
74-
xy_index = healpixl_nested_to_xy(healpix_index[i], nside)
75-
healpixl_to_radec(xy_index, nside, dx, dy, &lon[i], &lat[i])
75+
xy_index = healpixl_nested_to_xy(healpix_index[i], nside[i])
76+
healpixl_to_radec(xy_index, nside[i], dx, dy, &lon[i], &lat[i])
7677
elif order == 'ring':
7778
for i in prange(n, nogil=True, schedule='static'):
78-
xy_index = healpixl_ring_to_xy(healpix_index[i], nside)
79-
healpixl_to_radec(xy_index, nside, dx, dy, &lon[i], &lat[i])
79+
xy_index = healpixl_ring_to_xy(healpix_index[i], nside[i])
80+
healpixl_to_radec(xy_index, nside[i], dx, dy, &lon[i], &lat[i])
8081

8182
return lon, lat
8283

@@ -85,7 +86,8 @@ def healpix_to_lonlat(np.ndarray[int64_t, ndim=1, mode="c"] healpix_index,
8586
def healpix_with_offset_to_lonlat(np.ndarray[int64_t, ndim=1, mode="c"] healpix_index,
8687
np.ndarray[double_t, ndim=1, mode="c"] dx,
8788
np.ndarray[double_t, ndim=1, mode="c"] dy,
88-
int nside, str order):
89+
np.ndarray[int64_t, ndim=1, mode="c"] nside,
90+
str order):
8991
"""
9092
Convert HEALPix indices to longitudes/latitudes
9193
@@ -122,20 +124,21 @@ def healpix_with_offset_to_lonlat(np.ndarray[int64_t, ndim=1, mode="c"] healpix_
122124

123125
if order == 'nested':
124126
for i in prange(n, nogil=True, schedule='static'):
125-
xy_index = healpixl_nested_to_xy(healpix_index[i], nside)
126-
healpixl_to_radec(xy_index, nside, dx[i], dy[i], &lon[i], &lat[i])
127+
xy_index = healpixl_nested_to_xy(healpix_index[i], nside[i])
128+
healpixl_to_radec(xy_index, nside[i], dx[i], dy[i], &lon[i], &lat[i])
127129
elif order == 'ring':
128130
for i in prange(n, nogil=True, schedule='static'):
129-
xy_index = healpixl_ring_to_xy(healpix_index[i], nside)
130-
healpixl_to_radec(xy_index, nside, dx[i], dy[i], &lon[i], &lat[i])
131+
xy_index = healpixl_ring_to_xy(healpix_index[i], nside[i])
132+
healpixl_to_radec(xy_index, nside[i], dx[i], dy[i], &lon[i], &lat[i])
131133

132134
return lon, lat
133135

134136

135137
@cython.boundscheck(False)
136138
def lonlat_to_healpix(np.ndarray[double_t, ndim=1, mode="c"] lon,
137139
np.ndarray[double_t, ndim=1, mode="c"] lat,
138-
int nside, str order):
140+
np.ndarray[int64_t, ndim=1, mode="c"] nside,
141+
str order):
139142
"""
140143
Convert longitudes/latitudes to HEALPix indices
141144
@@ -168,20 +171,21 @@ def lonlat_to_healpix(np.ndarray[double_t, ndim=1, mode="c"] lon,
168171

169172
if order == 'nested':
170173
for i in prange(n, nogil=True, schedule='static'):
171-
xy_index = radec_to_healpixlf(lon[i], lat[i], nside, &dx, &dy)
172-
healpix_index[i] = healpixl_xy_to_nested(xy_index, nside)
174+
xy_index = radec_to_healpixlf(lon[i], lat[i], nside[i], &dx, &dy)
175+
healpix_index[i] = healpixl_xy_to_nested(xy_index, nside[i])
173176
elif order == 'ring':
174177
for i in prange(n, nogil=True, schedule='static'):
175-
xy_index = radec_to_healpixlf(lon[i], lat[i], nside, &dx, &dy)
176-
healpix_index[i] = healpixl_xy_to_ring(xy_index, nside)
178+
xy_index = radec_to_healpixlf(lon[i], lat[i], nside[i], &dx, &dy)
179+
healpix_index[i] = healpixl_xy_to_ring(xy_index, nside[i])
177180

178181
return healpix_index
179182

180183

181184
@cython.boundscheck(False)
182185
def lonlat_to_healpix_with_offset(np.ndarray[double_t, ndim=1, mode="c"] lon,
183186
np.ndarray[double_t, ndim=1, mode="c"] lat,
184-
int nside, str order):
187+
np.ndarray[int64_t, ndim=1, mode="c"] nside,
188+
str order):
185189
"""
186190
Convert longitudes/latitudes to healpix indices
187191
@@ -217,18 +221,19 @@ def lonlat_to_healpix_with_offset(np.ndarray[double_t, ndim=1, mode="c"] lon,
217221

218222
if order == 'nested':
219223
for i in prange(n, nogil=True, schedule='static'):
220-
xy_index = radec_to_healpixlf(lon[i], lat[i], nside, &dx[i], &dy[i])
221-
healpix_index[i] = healpixl_xy_to_nested(xy_index, nside)
224+
xy_index = radec_to_healpixlf(lon[i], lat[i], nside[i], &dx[i], &dy[i])
225+
healpix_index[i] = healpixl_xy_to_nested(xy_index, nside[i])
222226
elif order == 'ring':
223227
for i in prange(n, nogil=True, schedule='static'):
224-
xy_index = radec_to_healpixlf(lon[i], lat[i], nside, &dx[i], &dy[i])
225-
healpix_index[i] = healpixl_xy_to_ring(xy_index, nside)
228+
xy_index = radec_to_healpixlf(lon[i], lat[i], nside[i], &dx[i], &dy[i])
229+
healpix_index[i] = healpixl_xy_to_ring(xy_index, nside[i])
226230

227231
return healpix_index, dx, dy
228232

229233

230234
@cython.boundscheck(False)
231-
def nested_to_ring(np.ndarray[int64_t, ndim=1, mode="c"] nested_index, int nside):
235+
def nested_to_ring(np.ndarray[int64_t, ndim=1, mode="c"] nested_index,
236+
np.ndarray[int64_t, ndim=1, mode="c"] nside):
232237
"""
233238
Convert a HEALPix 'nested' index to a HEALPix 'ring' index
234239
@@ -250,13 +255,14 @@ def nested_to_ring(np.ndarray[int64_t, ndim=1, mode="c"] nested_index, int nside
250255
cdef np.ndarray[int64_t, ndim=1, mode="c"] ring_index = np.zeros(n, dtype=npy_int64)
251256

252257
for i in prange(n, nogil=True, schedule='static'):
253-
ring_index[i] = healpixl_xy_to_ring(healpixl_nested_to_xy(nested_index[i], nside), nside)
258+
ring_index[i] = healpixl_xy_to_ring(healpixl_nested_to_xy(nested_index[i], nside[i]), nside[i])
254259

255260
return ring_index
256261

257262

258263
@cython.boundscheck(False)
259-
def ring_to_nested(np.ndarray[int64_t, ndim=1, mode="c"] ring_index, int nside):
264+
def ring_to_nested(np.ndarray[int64_t, ndim=1, mode="c"] ring_index,
265+
np.ndarray[int64_t, ndim=1, mode="c"] nside):
260266
"""
261267
Convert a HEALPix 'ring' index to a HEALPix 'nested' index
262268
@@ -278,7 +284,7 @@ def ring_to_nested(np.ndarray[int64_t, ndim=1, mode="c"] ring_index, int nside):
278284
cdef np.ndarray[int64_t, ndim=1, mode="c"] nested_index = np.zeros(n, dtype=npy_int64)
279285

280286
for i in prange(n, nogil=True, schedule='static'):
281-
nested_index[i] = healpixl_xy_to_nested(healpixl_ring_to_xy(ring_index[i], nside), nside)
287+
nested_index[i] = healpixl_xy_to_nested(healpixl_ring_to_xy(ring_index[i], nside[i]), nside[i])
282288

283289
return nested_index
284290

@@ -287,7 +293,8 @@ def ring_to_nested(np.ndarray[int64_t, ndim=1, mode="c"] ring_index, int nside):
287293
@cython.boundscheck(False)
288294
def bilinear_interpolation_weights(np.ndarray[double_t, ndim=1, mode="c"] lon,
289295
np.ndarray[double_t, ndim=1, mode="c"] lat,
290-
int nside, str order):
296+
np.ndarray[int64_t, ndim=1, mode="c"] nside,
297+
str order):
291298
"""
292299
Get the four neighbours for each (lon, lat) position and the weight
293300
associated with each one for bilinear interpolation.
@@ -345,10 +352,10 @@ def bilinear_interpolation_weights(np.ndarray[double_t, ndim=1, mode="c"] lon,
345352
abort()
346353

347354
for i in range(n):
348-
interpolate_weights(lon[i], lat[i], indices_indiv, weights_indiv, nside)
355+
interpolate_weights(lon[i], lat[i], indices_indiv, weights_indiv, nside[i])
349356
for j in range(4):
350357
if order_int == 0:
351-
indices[j, i] = healpixl_xy_to_nested(healpixl_ring_to_xy(indices_indiv[j], nside), nside)
358+
indices[j, i] = healpixl_xy_to_nested(healpixl_ring_to_xy(indices_indiv[j], nside[i]), nside[i])
352359
else:
353360
indices[j, i] = indices_indiv[j]
354361
weights[j, i] = weights_indiv[j]
@@ -358,7 +365,7 @@ def bilinear_interpolation_weights(np.ndarray[double_t, ndim=1, mode="c"] lon,
358365

359366
@cython.boundscheck(False)
360367
def neighbours(np.ndarray[int64_t, ndim=1, mode="c"] healpix_index,
361-
int nside, str order):
368+
np.ndarray[int64_t, ndim=1, mode="c"] nside, str order):
362369
"""
363370
Find all the HEALPix pixels that are the neighbours of a HEALPix pixel
364371
@@ -417,8 +424,8 @@ def neighbours(np.ndarray[int64_t, ndim=1, mode="c"] healpix_index,
417424

418425
for i in prange(n, schedule='static'):
419426

420-
xy_index = healpixl_nested_to_xy(healpix_index[i], nside)
421-
healpixl_get_neighbours(xy_index, neighbours_indiv, nside)
427+
xy_index = healpixl_nested_to_xy(healpix_index[i], nside[i])
428+
healpixl_get_neighbours(xy_index, neighbours_indiv, nside[i])
422429

423430
for j in range(8):
424431
k = 4 - j
@@ -427,15 +434,15 @@ def neighbours(np.ndarray[int64_t, ndim=1, mode="c"] healpix_index,
427434
if neighbours_indiv[k] < 0:
428435
neighbours[j, i] = -1
429436
else:
430-
neighbours[j, i] = healpixl_xy_to_nested(neighbours_indiv[k], nside)
437+
neighbours[j, i] = healpixl_xy_to_nested(neighbours_indiv[k], nside[i])
431438

432439
elif order_int == 1:
433440

434441
for i in prange(n, schedule='static'):
435442

436-
xy_index = healpixl_ring_to_xy(healpix_index[i], nside)
443+
xy_index = healpixl_ring_to_xy(healpix_index[i], nside[i])
437444

438-
healpixl_get_neighbours(xy_index, neighbours_indiv, nside)
445+
healpixl_get_neighbours(xy_index, neighbours_indiv, nside[i])
439446

440447
for j in range(8):
441448
k = 4 - j
@@ -444,7 +451,7 @@ def neighbours(np.ndarray[int64_t, ndim=1, mode="c"] healpix_index,
444451
if neighbours_indiv[k] < 0:
445452
neighbours[j, i] = -1
446453
else:
447-
neighbours[j, i] = healpixl_xy_to_ring(neighbours_indiv[k], nside)
454+
neighbours[j, i] = healpixl_xy_to_ring(neighbours_indiv[k], nside[i])
448455

449456
free(neighbours_indiv)
450457

0 commit comments

Comments
 (0)