|
| 1 | +import numpy as np |
| 2 | + |
| 3 | +cdef extern from "sharp.h": |
| 4 | + ctypedef long ptrdiff_t |
| 5 | + |
| 6 | + void sharp_legendre_transform_s(float *bl, float *recfac, ptrdiff_t lmax, float *x, |
| 7 | + float *out, ptrdiff_t nx) |
| 8 | + void sharp_legendre_transform(double *bl, double *recfac, ptrdiff_t lmax, double *x, |
| 9 | + double *out, ptrdiff_t nx) |
| 10 | + void sharp_legendre_transform_recfac(double *r, ptrdiff_t lmax) |
| 11 | + void sharp_legendre_transform_recfac_s(float *r, ptrdiff_t lmax) |
| 12 | + |
| 13 | + |
| 14 | +def legendre_transform(x, bl, out=None): |
| 15 | + if out is None: |
| 16 | + out = np.empty_like(x) |
| 17 | + if x.dtype == np.float64: |
| 18 | + if bl.dtype != np.float64: |
| 19 | + bl = bl.astype(np.float64) |
| 20 | + return _legendre_transform(x, bl, out=out) |
| 21 | + elif x.dtype == np.float32: |
| 22 | + if bl.dtype != np.float32: |
| 23 | + bl = bl.astype(np.float32) |
| 24 | + return _legendre_transform_s(x, bl, out=out) |
| 25 | + else: |
| 26 | + raise ValueError("unsupported dtype") |
| 27 | + |
| 28 | + |
| 29 | +def _legendre_transform(double[::1] x, double[::1] bl, double[::1] out): |
| 30 | + if out.shape[0] != x.shape[0]: |
| 31 | + raise ValueError('x and out must have same shape') |
| 32 | + sharp_legendre_transform(&bl[0], NULL, bl.shape[0] - 1, &x[0], &out[0], x.shape[0]) |
| 33 | + return np.asarray(out) |
| 34 | + |
| 35 | + |
| 36 | +def _legendre_transform_s(float[::1] x, float[::1] bl, float[::1] out): |
| 37 | + if out.shape[0] != x.shape[0]: |
| 38 | + raise ValueError('x and out must have same shape') |
| 39 | + sharp_legendre_transform_s(&bl[0], NULL, bl.shape[0] - 1, &x[0], &out[0], x.shape[0]) |
| 40 | + return np.asarray(out) |
| 41 | + |
0 commit comments