From 48e213151a6c056c8bc0cc24ae059094ba9bbc1b Mon Sep 17 00:00:00 2001 From: Dag Sverre Seljebotn Date: Fri, 15 May 2015 10:26:40 +0200 Subject: [PATCH] First draft of SHTs in Python wrapper --- Makefile | 11 +- libsharp/sharp.c | 25 ++++ libsharp/sharp_lowlevel.h | 24 ++++ python/libsharp/libsharp.pxd | 79 +++++++++++ python/libsharp/libsharp.pyx | 210 ++++++++++++++++++++++++++++-- python/libsharp/libsharp_mpi.pyx | 32 +++++ python/libsharp/tests/test_sht.py | 33 +++++ 7 files changed, 397 insertions(+), 17 deletions(-) create mode 100644 python/libsharp/libsharp.pxd create mode 100644 python/libsharp/libsharp_mpi.pyx create mode 100644 python/libsharp/tests/test_sht.py diff --git a/Makefile b/Makefile index a621948..7644974 100644 --- a/Makefile +++ b/Makefile @@ -62,9 +62,10 @@ perftest: compile_all genclean: rm libsharp/sharp_legendre.c || exit 0 -pytest: - rm python/libsharp/libsharp.so || exit 0 - cd python && LIBSHARP_INCLUDE=$(INCDIR) LIBSHARP_LIB=$(LIBDIR) python setup.py build_ext --inplace - cd python && nosetests libsharp - +python/libsharp/libsharp.so: python/libsharp/libsharp.pyx $(LIB_libsharp) + cython $< + $(CC) -fPIC `python-config --cflags` -I$(INCDIR) -o python/libsharp/libsharp.o -c python/libsharp/libsharp.c + $(CL) -shared python/libsharp/libsharp.o -L$(LIBDIR) -lsharp -lfftpack -lc_utils `python-config --libs` -o $@ +pytest: python/libsharp/libsharp.so + cd python && nosetests --nocapture libsharp/tests/test_sht.py diff --git a/libsharp/sharp.c b/libsharp/sharp.c index e5a5634..1eb8857 100644 --- a/libsharp/sharp.c +++ b/libsharp/sharp.c @@ -992,4 +992,29 @@ int sharp_nv_oracle (sharp_jobtype type, int spin, int ntrans) #ifdef USE_MPI #include "sharp_mpi.c" + +int sharp_execute_mpi_maybe (void *pcomm, sharp_jobtype type, int spin, + void *alm, void *map, const sharp_geom_info *geom_info, + const sharp_alm_info *alm_info, int ntrans, int flags, double *time, + unsigned long long *opcnt) + { + MPI_Comm comm = *(MPI_Comm*)pcomm; + sharp_execute_mpi((MPI_Comm)comm, type, spin, alm, map, geom_info, alm_info, ntrans, + flags, time, opcnt); + return 0; + } + +#else + +int sharp_execute_mpi_maybe (void *pcomm, sharp_jobtype type, int spin, + void *alm, void *map, const sharp_geom_info *geom_info, + const sharp_alm_info *alm_info, int ntrans, int flags, double *time, + unsigned long long *opcnt) + { + /* Suppress unused warning: */ + (void)pcomm; (void)type; (void)spin; (void)alm; (void)map; (void)geom_info; + (void)alm_info; (void)ntrans; (void)flags; (void)time; (void)opcnt; + return SHARP_ERROR_NO_MPI; + } + #endif diff --git a/libsharp/sharp_lowlevel.h b/libsharp/sharp_lowlevel.h index 0260c37..d9aa01b 100644 --- a/libsharp/sharp_lowlevel.h +++ b/libsharp/sharp_lowlevel.h @@ -239,6 +239,30 @@ void sharp_execute (sharp_jobtype type, int spin, void *alm, void *map, void sharp_set_chunksize_min(int new_chunksize_min); void sharp_set_nchunks_max(int new_nchunks_max); + +typedef enum { SHARP_ERROR_NO_MPI = 1, + /*!< libsharp not compiled with MPI support */ + } sharp_errors; + +/*! Works like sharp_execute_mpi, but is always present whether or not libsharp + is compiled with USE_MPI. This is primarily useful for wrapper code etc. + + Note that \a pcomm has the type MPI_Comm*, except we declare void* to avoid + pulling in MPI headers. I.e., the comm argument of sharp_execute_mpi + is *(MPI_Comm*)pcomm. + + Other parameters are the same as sharp_execute_mpi. + + Returns 0 if successful, or SHARP_ERROR_NO_MPI if MPI is not available + (in which case nothing is done). + */ +int sharp_execute_mpi_maybe (void *pcomm, sharp_jobtype type, int spin, + void *alm, void *map, const sharp_geom_info *geom_info, + const sharp_alm_info *alm_info, int ntrans, int flags, double *time, + unsigned long long *opcnt); + + + /*! \} */ #ifdef __cplusplus diff --git a/python/libsharp/libsharp.pxd b/python/libsharp/libsharp.pxd new file mode 100644 index 0000000..fc27f19 --- /dev/null +++ b/python/libsharp/libsharp.pxd @@ -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) + + diff --git a/python/libsharp/libsharp.pyx b/python/libsharp/libsharp.pyx index 9e83d3d..3144e03 100644 --- a/python/libsharp/libsharp.pyx +++ b/python/libsharp/libsharp.pyx @@ -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 = 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) + diff --git a/python/libsharp/libsharp_mpi.pyx b/python/libsharp/libsharp_mpi.pyx new file mode 100644 index 0000000..e5c951d --- /dev/null +++ b/python/libsharp/libsharp_mpi.pyx @@ -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 = PyCapsule_GetPointer(capsule_to_c_comm, + "_mpi4py_bridge_MPI_Comm") + out[0] = mpi4py_comm.ob_mpi diff --git a/python/libsharp/tests/test_sht.py b/python/libsharp/tests/test_sht.py new file mode 100644 index 0000000..e97b90f --- /dev/null +++ b/python/libsharp/tests/test_sht.py @@ -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()