Aquila-Consortium/libsharp2
0
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

First draft of SHTs in Python wrapper

Browse source
This commit is contained in:
Dag Sverre Seljebotn 2015-05-15 10:26:40 +02:00
parent c680bbb2e5
commit 48e213151a
7 changed files with 397 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

79
python/libsharp/libsharp.pxd Normal file
View file

@ -0,0 +1,79 @@
cdef extern from "sharp.h":
ctypedef long ptrdiff_t
void sharp_legendre_transform_s(float *bl, float *recfac, ptrdiff_t lmax, float *x,
float *out, ptrdiff_t nx)
void sharp_legendre_transform(double *bl, double *recfac, ptrdiff_t lmax, double *x,
double *out, ptrdiff_t nx)
void sharp_legendre_transform_recfac(double *r, ptrdiff_t lmax)
void sharp_legendre_transform_recfac_s(float *r, ptrdiff_t lmax)
void sharp_legendre_roots(int n, double *x, double *w)
# sharp_lowlevel.h
ctypedef struct sharp_alm_info:
pass
ctypedef struct sharp_geom_info:
pass
void sharp_make_alm_info (int lmax, int mmax, int stride,
ptrdiff_t *mvstart, sharp_alm_info **alm_info)
void sharp_make_geom_info (int nrings, int *nph, ptrdiff_t *ofs,
int *stride, double *phi0, double *theta,
double *wgt, sharp_geom_info **geom_info)
void sharp_destroy_alm_info(sharp_alm_info *info)
void sharp_destroy_geom_info(sharp_geom_info *info)
ptrdiff_t sharp_map_size(sharp_geom_info *info)
ptrdiff_t sharp_alm_count(sharp_alm_info *self);
ctypedef enum sharp_jobtype:
SHARP_YtW
SHARP_Yt
SHARP_WY
SHARP_Y
ctypedef enum:
SHARP_DP
SHARP_ADD
void sharp_execute(sharp_jobtype type_,
int spin,
void *alm,
void *map,
sharp_geom_info *geom_info,
sharp_alm_info *alm_info,
int ntrans,
int flags,
double *time,
unsigned long long *opcnt) nogil
ctypedef enum:
SHARP_ERROR_NO_MPI
int sharp_execute_mpi_maybe (void *pcomm, sharp_jobtype type, int spin,
void *alm, void *map, sharp_geom_info *geom_info,
sharp_alm_info *alm_info, int ntrans, int flags, double *time,
unsigned long long *opcnt) nogil
cdef extern from "sharp_geomhelpers.h":
void sharp_make_subset_healpix_geom_info(
int nside, int stride, int nrings,
int *rings, double *weight, sharp_geom_info **geom_info)
void sharp_make_gauss_geom_info(
int nrings, int nphi, double phi0,
int stride_lon, int stride_lat, sharp_geom_info **geom_info)
cdef extern from "sharp_almhelpers.h":
void sharp_make_triangular_alm_info (int lmax, int mmax, int stride,
sharp_alm_info **alm_info)
void sharp_make_rectangular_alm_info (int lmax, int mmax, int stride,
sharp_alm_info **alm_info)
void sharp_make_mmajor_real_packed_alm_info (int lmax, int stride,
int nm, const int *ms, sharp_alm_info **alm_info)

210
python/libsharp/libsharp.pyx
View file

@ -1,17 +1,8 @@
import numpy as np
__all__ = ['legendre_transform', 'legendre_roots']
cdef extern from "sharp.h":
ctypedef long ptrdiff_t
void sharp_legendre_transform_s(float *bl, float *recfac, ptrdiff_t lmax, float *x,
float *out, ptrdiff_t nx)
void sharp_legendre_transform(double *bl, double *recfac, ptrdiff_t lmax, double *x,
double *out, ptrdiff_t nx)
void sharp_legendre_transform_recfac(double *r, ptrdiff_t lmax)
void sharp_legendre_transform_recfac_s(float *r, ptrdiff_t lmax)
void sharp_legendre_roots(int n, double *x, double *w)
__all__ = ['legendre_transform', 'legendre_roots', 'sht', 'synthesis', 'adjoint_synthesis',
'analysis', 'adjoint_analysis', 'healpix_grid', 'triangular_order', 'rectangular_order',
'packed_real_order']
def legendre_transform(x, bl, out=None):
@ -54,3 +45,198 @@ def legendre_roots(n):
if n > 0:
sharp_legendre_roots(n, &x_buf[0], &w_buf[0])
return x, w
JOBTYPE_TO_CONST = {
'Y': SHARP_Y,
'Yt': SHARP_Yt,
'WY': SHARP_WY,
'YtW': SHARP_YtW
}
def sht(jobtype, geom_info ginfo, alm_info ainfo, double[::1] input,
int spin=0, comm=None, add=False):
cdef void *comm_ptr
cdef int flags = SHARP_DP | (SHARP_ADD if add else 0)
cdef double *palm
cdef double *pmap
cdef int r
cdef sharp_jobtype jobtype_i
cdef double[::1] output_buf
try:
jobtype_i = JOBTYPE_TO_CONST[jobtype]
except KeyError:
raise ValueError('jobtype must be one of: %s' % ', '.join(sorted(JOBTYPE_TO_CONST.keys())))
if jobtype_i == SHARP_Y or jobtype_i == SHARP_WY:
output = np.empty(ginfo.local_size(), dtype=np.float64)
output_buf = output
pmap = &output_buf[0]
palm = &input[0]
else:
output = np.empty(ainfo.local_size(), dtype=np.float64)
output_buf = output
pmap = &input[0]
palm = &output_buf[0]
if comm is None:
with nogil:
sharp_execute (
jobtype_i,
geom_info=ginfo.ginfo, alm_info=ainfo.ainfo,
spin=spin, alm=&palm, map=&pmap,
ntrans=1, flags=flags, time=NULL, opcnt=NULL)
else:
from mpi4py import MPI
if not isinstance(comm, MPI.Comm):
raise TypeError('comm must be an mpi4py communicator')
comm_ptr = <void*><size_t>MPI._addressof(comm)
with nogil:
r = sharp_execute_mpi_maybe (
comm_ptr, jobtype_i,
geom_info=ginfo.ginfo, alm_info=ainfo.ainfo,
spin=spin, alm=&palm, map=&pmap,
ntrans=1, flags=flags, time=NULL, opcnt=NULL)
if r == SHARP_ERROR_NO_MPI:
raise Exception('MPI requested, but not available')
return output
def synthesis(*args, **kw):
return sht('Y', *args, **kw)
def adjoint_synthesis(*args, **kw):
return sht('Yt', *args, **kw)
def analysis(*args, **kw):
return sht('YtW', *args, **kw)
def adjoint_analysis(*args, **kw):
return sht('WY', *args, **kw)
#
# geom_info
#
class NotInitializedError(Exception):
pass
cdef class geom_info:
cdef sharp_geom_info *ginfo
def __cinit__(self, *args, **kw):
self.ginfo = NULL
def local_size(self):
if self.ginfo == NULL:
raise NotInitializedError()
return sharp_map_size(self.ginfo)
def __dealloc__(self):
if self.ginfo != NULL:
sharp_destroy_geom_info(self.ginfo)
self.ginfo = NULL
cdef class healpix_grid(geom_info):
_weight_cache = {} # { (nside, 'T'/'Q'/'U') -> numpy array of ring weights cached from file }
def __init__(self, int nside, stride=1, int[::1] rings=None, double[::1] weights=None):
if weights is not None and weights.shape[0] != 2 * nside:
raise ValueError('weights must have length 2 * nside')
sharp_make_subset_healpix_geom_info(nside, stride,
nrings=4 * nside - 1 if rings is None else rings.shape[0],
rings=NULL if rings is None else &rings[0],
weight=NULL if weights is None else &weights[0],
geom_info=&self.ginfo)
@classmethod
def load_ring_weights(cls, nside, fields):
"""
Loads HEALPix ring weights from file. The environment variable
HEALPIX should be set, and this routine will look in the `data`
subdirectory.
Parameters
----------
nside: int
HEALPix nside parameter
fields: tuple of str
Which weights to extract; pass ('T',) to only get scalar
weights back, or ('T', 'Q', 'U') to get all the weights
Returns
-------
List of NumPy arrays, according to fields parameter.
"""
import os
from astropy.io import fits
data_path = os.path.join(os.environ['HEALPIX'], 'data')
fits_field_names = {
'T': 'TEMPERATURE WEIGHTS',
'Q': 'Q-POLARISATION WEIGHTS',
'U': 'U-POLARISATION WEIGHTS'}
must_load = [field for field in fields if (nside, field) not in cls._weight_cache]
if must_load:
hdulist = fits.open(os.path.join(data_path, 'weight_ring_n%05d.fits' % nside))
try:
for field in must_load:
w = hdulist[1].data.field(fits_field_names[field]).ravel().astype(np.double)
w += 1
cls._weight_cache[nside, field] = w
finally:
hdulist.close()
return [cls._weight_cache[(nside, field)].copy() for field in fields]
#
# alm_info
#
cdef class alm_info:
cdef sharp_alm_info *ainfo
def __cinit__(self, *args, **kw):
self.ainfo = NULL
def local_size(self):
if self.ainfo == NULL:
raise NotInitializedError()
return sharp_alm_count(self.ainfo)
def __dealloc__(self):
if self.ainfo != NULL:
sharp_destroy_alm_info(self.ainfo)
self.ainfo = NULL
cdef class triangular_order(alm_info):
def __init__(self, int lmax, mmax=None, stride=1):
mmax = mmax if mmax is not None else lmax
sharp_make_triangular_alm_info(lmax, mmax, stride, &self.ainfo)
cdef class rectangular_order(alm_info):
def __init__(self, int lmax, mmax=None, stride=1):
mmax = mmax if mmax is not None else lmax
sharp_make_rectangular_alm_info(lmax, mmax, stride, &self.ainfo)
cdef class packed_real_order(alm_info):
def __init__(self, int lmax, stride=1, int[::1] ms=None):
sharp_make_mmajor_real_packed_alm_info(lmax=lmax, stride=stride,
nm=lmax + 1 if ms is None else ms.shape[0],
ms=NULL if ms is None else &ms[0],
alm_info=&self.ainfo)

32
python/libsharp/libsharp_mpi.pyx Normal file
View file

@ -0,0 +1,32 @@
"""
We keep MPI support in a seperate module because we need to import mpi4py,
and that will fire up the linked MPI library, which is something we only
want to do if a comm was passed.
"""
from mpi4py.MPI cimport Comm
from mpi4py.mpi_c cimport MPI_Comm
cimport libsharp
cdef extern from "sharp_mpi.h":
void sharp_execute_mpi (MPI_Comm comm, sharp_jobtype type, int spin,
void *alm, void *map, sharp_geom_info *geom_info,
sharp_alm_info *alm_info, int ntrans, int flags, double *time,
unsigned long long *opcnt) nogil
def _sht_mpi(MPI_comm comm, int jobtype, int spin, int flags, int ntrans, grid_info ginfo, double[::1] map,
alm_info ainfo, double[::1] alm):
with nogil:
with nogil:
sharp_execute_mpi(comm=comm.ob_mpi, jobtype=jobtype, spin=spin,
alm=&alm[0], map=&map[0], geom_info=ginfo.ginfo,
alm_info=ainfo.ainfo, ntrans=ntrans, flags=flags,
time=NULL, opcnt=NULL)
def sht_mpi(object comm, int jobtype, grid g, double[::1] map, alm_order, double[::1] alm, spin=0, comm=None, add=False):
def mpi4py_Comm_to_handle(Comm mpi4py_comm, capsule_to_c_comm):
cdef MPI_Comm *out = <MPI_Comm*>PyCapsule_GetPointer(capsule_to_c_comm,
"_mpi4py_bridge_MPI_Comm")
out[0] = mpi4py_comm.ob_mpi

33
python/libsharp/tests/test_sht.py Normal file
View file

@ -0,0 +1,33 @@
import numpy as np
import healpy
from scipy.special import legendre
from scipy.special import p_roots
from numpy.testing import assert_allclose
import libsharp
from mpi4py import MPI
def test_basic():
lmax = 10
nside = 8
rank = MPI.COMM_WORLD.Get_rank()
ms = np.arange(rank, lmax + 1, MPI.COMM_WORLD.Get_size(), dtype=np.int32)
order = libsharp.packed_real_order(lmax, ms=ms)
grid = libsharp.healpix_grid(nside)
alm = np.zeros(order.local_size())
if rank == 0:
alm[0] = 1
elif rank == 1:
alm[0] = 1
map = libsharp.synthesis(grid, order, alm, comm=MPI.COMM_WORLD)
if rank == 0:
healpy.mollzoom(map)
from matplotlib.pyplot import show
show()