First draft of SHTs in Python wrapper

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

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

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

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)
+
+

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 = <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)
+

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 = <MPI_Comm*>PyCapsule_GetPointer(capsule_to_c_comm,
+                                                         "_mpi4py_bridge_MPI_Comm")
+    out[0] = mpi4py_comm.ob_mpi

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()