Merge pull request #7 from dagss/pythonsht

Declaring this ready for merging into master now.

.gitignore

@@ -14,3 +14,4 @@
 /sharp_oracle.inc
 
 /python/libsharp/libsharp.c
+/python/libsharp/libsharp_mpi.c

Makefile

@@ -19,6 +19,8 @@ include libfftpack/planck.make
 include libsharp/planck.make
 include docsrc/planck.make
 
+CYTHON_MODULES=python/libsharp/libsharp.so $(if $(MPI_CFLAGS), python/libsharp/libsharp_mpi.so)
+
 $(all_lib): %: | $(LIBDIR)_mkdir
 	@echo "#  creating library $*"
 	$(ARCREATE) $@ $^
@@ -62,9 +64,15 @@ 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
+$(CYTHON_MODULES): %.so: %.pyx
+ifndef PIC_CFLAGS
+	$(error Python extension must be built using the --enable-pic configure option.)
+endif
+	cython $<
+	$(CC) $(DEBUG_CFLAGS) $(OPENMP_CFLAGS) $(PIC_CFLAGS) `python-config --cflags` -I$(INCDIR) -o $(<:.pyx=.o) -c $(<:.pyx=.c)
+	$(CL) -shared $(<:.pyx=.o) $(OPENMP_CFLAGS) $(CYTHON_OBJ) -L$(LIBDIR) -lsharp -lfftpack -lc_utils -L`python-config --prefix`/lib `python-config --ldflags` -o $@
 
+python: $(all_lib) hdrcopy $(CYTHON_MODULES)
 
+pytest: python
+	cd python && nosetests --nocapture libsharp/tests/test_sht.py

config/config.auto.in

@@ -5,5 +5,8 @@ CL=@CC@
 CCFLAGS_NO_C=@CCFLAGS_NO_C@
 CCFLAGS=$(CCFLAGS_NO_C) -c
 CLFLAGS=-L. -L$(LIBDIR) @LDCCFLAGS@ -lm
-
+DEBUG_CFLAGS=@DEBUG_CFLAGS@
+MPI_CFLAGS=@MPI_CFLAGS@
+OPENMP_CFLAGS=@OPENMP_CFLAGS@
+PIC_CFLAGS=@PIC_CFLAGS@
 ARCREATE=@ARCREATE@

config/rules.common

@@ -1,9 +1,10 @@
-BLDROOT = $(SRCROOT)/build.$(SHARP_TARGET)
-PREFIX  = $(SRCROOT)/$(SHARP_TARGET)
-BINDIR	= $(PREFIX)/bin
-INCDIR	= $(PREFIX)/include
-LIBDIR	= $(PREFIX)/lib
-DOCDIR	= $(SRCROOT)/doc
+BLDROOT   = $(SRCROOT)/build.$(SHARP_TARGET)
+PREFIX    = $(SRCROOT)/$(SHARP_TARGET)
+BINDIR    = $(PREFIX)/bin
+INCDIR    = $(PREFIX)/include
+LIBDIR    = $(PREFIX)/lib
+DOCDIR    = $(SRCROOT)/doc
+PYTHONDIR = $(SRCROOT)/python/libsharp
 
 # do not use any suffix rules
 .SUFFIXES:
@@ -26,6 +27,7 @@ $(BLDROOT)/%.o : $(SRCROOT)/%.cc | echo_config
 
 clean:
 	rm -rf $(BLDROOT) $(PREFIX) $(DOCDIR) autom4te.cache/ config.log config.status
+	rm -rf $(PYTHONDIR)/*.c $(PYTHONDIR)/*.o $(PYTHONDIR)/*.so
 
 distclean: clean
 	rm -f config/config.auto

configure.ac

@@ -82,20 +82,20 @@ case $system in
     ;;
 esac
 
-CCFLAGS="$CCFLAGS $OPENMP_CFLAGS"
-
 if test $ENABLE_DEBUG = yes; then
-  CCFLAGS="$CCFLAGS -g"
+  DEBUG_CFLAGS="-g"
 fi
 
 if test $ENABLE_PIC = yes; then
-  CCFLAGS="$CCFLAGS -fPIC"
+  PIC_CFLAGS="-fPIC"
 fi
 
 if test $ENABLE_MPI = yes; then
-  CCFLAGS="$CCFLAGS -DUSE_MPI"
+  MPI_CFLAGS="-DUSE_MPI"
 fi
 
+CCFLAGS="$CCFLAGS $DEBUG_CFLAGS $OPENMP_CFLAGS $PIC_CFLAGS $MPI_CFLAGS"
+
 CCFLAGS_NO_C="$CCFLAGS $CPPFLAGS"
 
 LDCCFLAGS="$LDFLAGS $CCFLAGS"
@@ -104,6 +104,10 @@ AC_SUBST(SILENT_RULE)
 AC_SUBST(CC)
 AC_SUBST(CCFLAGS_NO_C)
 AC_SUBST(LDCCFLAGS)
+AC_SUBST(DEBUG_CFLAGS)
+AC_SUBST(MPI_CFLAGS)
+AC_SUBST(OPENMP_CFLAGS)
+AC_SUBST(PIC_CFLAGS)
 AC_SUBST(ARCREATE)
 
 AC_OUTPUT(config/config.auto)

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,212 @@ 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 int r
+    cdef sharp_jobtype jobtype_i
+    cdef double[:, :, ::1] output_buf
+    cdef int ntrans = input.shape[0] * input.shape[1]
+    cdef int i, j
+
+    if spin == 0 and input.shape[1] != 1:
+        raise ValueError('For spin == 0, we need input.shape[1] == 1')
+    elif spin != 0 and input.shape[1] != 2:
+        raise ValueError('For spin != 0, we need input.shape[1] == 2')
+
+
+    cdef size_t[::1] ptrbuf = np.empty(2 * ntrans, dtype=np.uintp)
+    cdef double **alm_ptrs = <double**>&ptrbuf[0]
+    cdef double **map_ptrs = <double**>&ptrbuf[ntrans]
+
+    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((input.shape[0], input.shape[1], ginfo.local_size()), dtype=np.float64)
+        output_buf = output
+        for i in range(input.shape[0]):
+            for j in range(input.shape[1]):
+                alm_ptrs[i * input.shape[1] + j] = &input[i, j, 0]
+                map_ptrs[i * input.shape[1] + j] = &output_buf[i, j, 0]
+    else:
+        output = np.empty((input.shape[0], input.shape[1], ainfo.local_size()), dtype=np.float64)
+        output_buf = output
+        for i in range(input.shape[0]):
+            for j in range(input.shape[1]):
+                alm_ptrs[i * input.shape[1] + j] = &output_buf[i, j, 0]
+                map_ptrs[i * input.shape[1] + j] = &input[i, j, 0]
+
+    if comm is None:
+        with nogil:
+            sharp_execute (
+                jobtype_i,
+                geom_info=ginfo.ginfo, alm_info=ainfo.ainfo,
+                spin=spin, alm=alm_ptrs, map=map_ptrs,
+                ntrans=ntrans, 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')
+        from .libsharp_mpi import _addressof
+        comm_ptr = <void*><size_t>_addressof(comm)
+        with nogil:
+            r = sharp_execute_mpi_maybe (
+                comm_ptr, jobtype_i,
+                geom_info=ginfo.ginfo, alm_info=ainfo.ainfo,
+                spin=spin, alm=alm_ptrs, map=map_ptrs,
+                ntrans=ntrans, 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,17 @@
+cdef extern from "mpi.h":
+    ctypedef void *MPI_Comm
+
+cdef extern from "Python.h":
+    object PyLong_FromVoidPtr(void*)
+
+cdef extern:
+    ctypedef class mpi4py.MPI.Comm [object PyMPICommObject]:
+        cdef MPI_Comm ob_mpi
+        cdef unsigned flags
+
+# For compatibility with mpi4py <= 1.3.1
+# Newer versions could use the MPI._addressof function
+def _addressof(Comm comm):
+    cdef void *ptr = NULL
+    ptr = <void*>&comm.ob_mpi
+    return PyLong_FromVoidPtr(ptr)

python/libsharp/tests/test_sht.py

@@ -0,0 +1,34 @@
+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, np.repeat(alm[None, None, :], 3, 0), comm=MPI.COMM_WORLD)
+    assert np.all(map[2, :] == map[1, :]) and np.all(map[1, :] == map[0, :])
+    map = map[0, 0, :]
+    if rank == 0:
+        healpy.mollzoom(map)
+        from matplotlib.pyplot import show
+        show()