_sky_map.c 13.8 KB
Newer Older
1
/*
2
 * Copyright (C) 2013-2017  Leo Singer
3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with with program; see the file COPYING. If not, write to the
 * Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston,
 * MA  02111-1307  USA
 */

20
#include "config.h"
21
#include <Python.h>
22 23 24
/* Ignore warnings in Numpy API itself */
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wcast-qual"
25
#include <numpy/arrayobject.h>
26
#include <numpy/ufuncobject.h>
27
#pragma GCC diagnostic pop
28 29
#include <chealpix.h>
#include <gsl/gsl_errno.h>
30
#include <gsl/gsl_nan.h>
31
#include <lal/bayestar_sky_map.h>
32
#include <stddef.h>
33
#include "six.h"
34 35


36
static void capsule_free(PyObject *self)
37
{
38
    free(PyCapsule_GetPointer(self, NULL));
39 40 41 42 43
}


#define INPUT_LIST_OF_ARRAYS(NAME, NPYTYPE, DEPTH, CHECK) \
{ \
44
    const Py_ssize_t n = PySequence_Length(NAME##_obj); \
45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87
    if (n < 0) \
        goto fail; \
    else if (n != nifos) { \
        PyErr_SetString(PyExc_ValueError, #NAME \
        " appears to be the wrong length for the number of detectors"); \
        goto fail; \
    } \
    for (unsigned int iifo = 0; iifo < nifos; iifo ++) \
    { \
        PyObject *obj = PySequence_ITEM(NAME##_obj, iifo); \
        if (!obj) goto fail; \
        PyArrayObject *npy = (PyArrayObject *) \
            PyArray_ContiguousFromAny(obj, NPYTYPE, DEPTH, DEPTH); \
        Py_XDECREF(obj); \
        if (!npy) goto fail; \
        NAME##_npy[iifo] = npy; \
        { CHECK } \
        NAME[iifo] = PyArray_DATA(npy); \
    } \
}

#define FREE_INPUT_LIST_OF_ARRAYS(NAME) \
{ \
    for (unsigned int iifo = 0; iifo < nifos; iifo ++) \
        Py_XDECREF(NAME##_npy[iifo]); \
}

#define INPUT_VECTOR_NIFOS(CTYPE, NAME, NPYTYPE) \
    NAME##_npy = (PyArrayObject *) \
        PyArray_ContiguousFromAny(NAME##_obj, NPYTYPE, 1, 1); \
    if (!NAME##_npy) goto fail; \
    if (PyArray_DIM(NAME##_npy, 0) != nifos) \
    { \
        PyErr_SetString(PyExc_ValueError, #NAME \
            " appears to be the wrong length for the number of detectors"); \
        goto fail; \
    } \
    const CTYPE *NAME = PyArray_DATA(NAME##_npy);

#define INPUT_VECTOR_DOUBLE_NIFOS(NAME) \
    INPUT_VECTOR_NIFOS(double, NAME, NPY_DOUBLE)


88 89 90 91 92 93 94
static PyArray_Descr *sky_map_descr;


static PyArray_Descr *sky_map_create_descr(void)
{
    PyArray_Descr *dtype = NULL;
    PyObject *dtype_dict = Py_BuildValue("{s(ssss)s(cccc)s(IIII)}",
95
        "names", "UNIQ", "PROBDENSITY", "DISTMEAN", "DISTSTD",
96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112
        "formats", NPY_ULONGLONGLTR, NPY_DOUBLELTR, NPY_DOUBLELTR, NPY_DOUBLELTR,
        "offsets",
        (unsigned int) offsetof(bayestar_pixel, uniq),
        (unsigned int) offsetof(bayestar_pixel, value[0]),
        (unsigned int) offsetof(bayestar_pixel, value[1]),
        (unsigned int) offsetof(bayestar_pixel, value[2]));

    if (dtype_dict)
    {
        PyArray_DescrConverter(dtype_dict, &dtype);
        Py_DECREF(dtype_dict);
    }

    return dtype;
}


113 114
static PyObject *sky_map_toa_phoa_snr(
    PyObject *NPY_UNUSED(module), PyObject *args, PyObject *kwargs)
115
{
116 117 118
    /* Input arguments */
    double min_distance;
    double max_distance;
119
    int prior_distance_power;
120
    int cosmology;
121 122 123 124
    double gmst;
    unsigned int nifos;
    unsigned long nsamples = 0;
    double sample_rate;
125 126
    PyObject *epochs_obj;
    PyObject *snrs_obj;
127 128 129
    PyObject *responses_obj;
    PyObject *locations_obj;
    PyObject *horizons_obj;
130 131

    /* Names of arguments */
132
    static const char *keywords[] = {"min_distance", "max_distance",
133 134
        "prior_distance_power", "cosmology", "gmst", "sample_rate", "epochs",
        "snrs", "responses", "locations", "horizons", NULL};
135 136

    /* Parse arguments */
Leo Pound Singer's avatar
Leo Pound Singer committed
137 138 139
    /* FIXME: PyArg_ParseTupleAndKeywords should expect keywords to be const */
    #pragma GCC diagnostic push
    #pragma GCC diagnostic ignored "-Wincompatible-pointer-types"
140 141 142 143
    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "ddiiddOOOOO",
        keywords, &min_distance, &max_distance, &prior_distance_power,
        &cosmology, &gmst, &sample_rate, &epochs_obj, &snrs_obj,
        &responses_obj, &locations_obj, &horizons_obj)) return NULL;
Leo Pound Singer's avatar
Leo Pound Singer committed
144
    #pragma GCC diagnostic pop
145

146
    /* Determine number of detectors */
147
    {
148
        Py_ssize_t n = PySequence_Length(epochs_obj);
149 150
        if (n < 0) return NULL;
        nifos = n;
151 152
    }

153 154
    /* Return value */
    PyObject *out = NULL;
155
    double log_bci, log_bsn;
156

157
    /* Numpy array objects */
158 159 160
    PyArrayObject *epochs_npy = NULL, *snrs_npy[nifos], *responses_npy[nifos],
        *locations_npy[nifos], *horizons_npy = NULL;
    memset(snrs_npy, 0, sizeof(snrs_npy));
161 162 163 164
    memset(responses_npy, 0, sizeof(responses_npy));
    memset(locations_npy, 0, sizeof(locations_npy));

    /* Arrays of pointers for inputs with multiple dimensions */
165
    const float complex *snrs[nifos];
166 167 168 169
    const float (*responses[nifos])[3];
    const double *locations[nifos];

    /* Gather C-aligned arrays from Numpy types */
170
    INPUT_VECTOR_DOUBLE_NIFOS(epochs)
171
    INPUT_LIST_OF_ARRAYS(snrs, NPY_CFLOAT, 1,
172 173 174 175
        npy_intp dim = PyArray_DIM(npy, 0);
        if (iifo == 0)
            nsamples = dim;
        else if ((unsigned long)dim != nsamples)
176
        {
177
            PyErr_SetString(PyExc_ValueError,
178
                "expected elements of snrs to be vectors of the same length");
179 180
            goto fail;
        }
181 182 183
    )
    INPUT_LIST_OF_ARRAYS(responses, NPY_FLOAT, 2,
        if (PyArray_DIM(npy, 0) != 3 || PyArray_DIM(npy, 1) != 3)
184
        {
185 186
            PyErr_SetString(PyExc_ValueError,
                "expected elements of responses to be 3x3 arrays");
187 188
            goto fail;
        }
189 190 191 192 193 194 195 196 197 198 199 200
    )
    INPUT_LIST_OF_ARRAYS(locations, NPY_DOUBLE, 1,
        if (PyArray_DIM(npy, 0) != 3)
        {
            PyErr_SetString(PyExc_ValueError,
                "expected elements of locations to be vectors of length 3");
            goto fail;
        }
    )
    INPUT_VECTOR_DOUBLE_NIFOS(horizons)

    /* Call function */
201
    gsl_error_handler_t *old_handler = gsl_set_error_handler_off();
202
    size_t len;
203 204
    bayestar_pixel *pixels;
    Py_BEGIN_ALLOW_THREADS
205
    pixels = bayestar_sky_map_toa_phoa_snr(&len, &log_bci, &log_bsn,
206 207 208
        min_distance, max_distance, prior_distance_power, cosmology, gmst,
        nifos, nsamples, sample_rate, epochs, snrs, responses, locations,
        horizons);
209
    Py_END_ALLOW_THREADS
210 211
    gsl_set_error_handler(old_handler);

212 213
    PyErr_CheckSignals();

214
    if (!pixels)
215 216
        goto fail;

217
    /* Prepare output object */
218
    PyObject *capsule = PyCapsule_New(pixels, NULL, capsule_free);
219 220 221
    if (!capsule)
        goto fail;

222 223 224 225
    npy_intp dims[] = {len};
    Py_INCREF(sky_map_descr);
    out = PyArray_NewFromDescr(&PyArray_Type,
        sky_map_descr, 1, dims, NULL, pixels, NPY_ARRAY_DEFAULT, NULL);
226 227 228 229 230 231
    if (!out)
    {
        Py_DECREF(capsule);
        goto fail;
    }

232
    if (PyArray_SetBaseObject((PyArrayObject *) out, capsule))
233
    {
234 235 236
        Py_DECREF(out);
        out = NULL;
        goto fail;
237
    }
238

239
fail: /* Cleanup */
240 241
    Py_XDECREF(epochs_npy);
    FREE_INPUT_LIST_OF_ARRAYS(snrs)
242 243 244
    FREE_INPUT_LIST_OF_ARRAYS(responses)
    FREE_INPUT_LIST_OF_ARRAYS(locations)
    Py_XDECREF(horizons_npy);
245 246 247
    if (out) {
        out = Py_BuildValue("Ndd", out, log_bci, log_bsn);
    }
248 249 250 251
    return out;
};


252 253
static PyObject *log_likelihood_toa_phoa_snr(
    PyObject *NPY_UNUSED(module), PyObject *args, PyObject *kwargs)
254
{
255 256 257 258 259 260 261 262 263 264 265
    /* Input arguments */
    double ra;
    double sin_dec;
    double distance;
    double u;
    double twopsi;
    double t;
    double gmst;
    unsigned int nifos;
    unsigned long nsamples = 0;
    double sample_rate;
266 267
    PyObject *epochs_obj;
    PyObject *snrs_obj;
268 269 270
    PyObject *responses_obj;
    PyObject *locations_obj;
    PyObject *horizons_obj;
271 272

    /* Names of arguments */
273 274
    static const char *keywords[] = {"params", "gmst", "sample_rate", "epochs",
        "snrs", "responses", "locations", "horizons", NULL};
275 276

    /* Parse arguments */
Leo Pound Singer's avatar
Leo Pound Singer committed
277 278 279
    /* FIXME: PyArg_ParseTupleAndKeywords should expect keywords to be const */
    #pragma GCC diagnostic push
    #pragma GCC diagnostic ignored "-Wincompatible-pointer-types"
280
    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "(dddddd)ddOOOOO",
281
        keywords, &ra, &sin_dec, &distance, &u, &twopsi, &t, &gmst,
282 283
        &sample_rate, &epochs_obj, &snrs_obj, &responses_obj, &locations_obj,
        &horizons_obj)) return NULL;
Leo Pound Singer's avatar
Leo Pound Singer committed
284
    #pragma GCC diagnostic pop
285

286
    /* Determine number of detectors */
287
    {
288
        Py_ssize_t n = PySequence_Length(epochs_obj);
289 290
        if (n < 0) return NULL;
        nifos = n;
291 292
    }

293 294
    /* Return value */
    PyObject *out = NULL;
295

296
    /* Numpy array objects */
297 298 299
    PyArrayObject *epochs_npy = NULL, *snrs_npy[nifos], *responses_npy[nifos],
        *locations_npy[nifos], *horizons_npy = NULL;
    memset(snrs_npy, 0, sizeof(snrs_npy));
300 301 302 303
    memset(responses_npy, 0, sizeof(responses_npy));
    memset(locations_npy, 0, sizeof(locations_npy));

    /* Arrays of pointers for inputs with multiple dimensions */
304
    const float complex *snrs[nifos];
305 306 307 308
    const float (*responses[nifos])[3];
    const double *locations[nifos];

    /* Gather C-aligned arrays from Numpy types */
309
    INPUT_VECTOR_DOUBLE_NIFOS(epochs)
310
    INPUT_LIST_OF_ARRAYS(snrs, NPY_CFLOAT, 1,
311 312 313 314
        npy_intp dim = PyArray_DIM(npy, 0);
        if (iifo == 0)
            nsamples = dim;
        else if ((unsigned long)dim != nsamples)
315
        {
316
            PyErr_SetString(PyExc_ValueError,
317
                "expected elements of snrs to be vectors of the same length");
318 319
            goto fail;
        }
320 321 322
    )
    INPUT_LIST_OF_ARRAYS(responses, NPY_FLOAT, 2,
        if (PyArray_DIM(npy, 0) != 3 || PyArray_DIM(npy, 1) != 3)
323
        {
324 325
            PyErr_SetString(PyExc_ValueError,
                "expected elements of responses to be 3x3 arrays");
326 327
            goto fail;
        }
328 329 330 331 332 333 334 335 336 337 338 339
    )
    INPUT_LIST_OF_ARRAYS(locations, NPY_DOUBLE, 1,
        if (PyArray_DIM(npy, 0) != 3)
        {
            PyErr_SetString(PyExc_ValueError,
                "expected elements of locations to be vectors of length 3");
            goto fail;
        }
    )
    INPUT_VECTOR_DOUBLE_NIFOS(horizons)

    /* Call function */
340
    gsl_error_handler_t *old_handler = gsl_set_error_handler_off();
341
    const double ret = bayestar_log_likelihood_toa_phoa_snr(ra, sin_dec,
342 343
        distance, u, twopsi, t, gmst, nifos, nsamples, sample_rate, epochs,
        snrs, responses, locations, horizons);
344
    gsl_set_error_handler(old_handler);
345 346 347 348 349

    /* Prepare output object */
    out = PyFloat_FromDouble(ret);

fail: /* Cleanup */
350 351
    Py_XDECREF(epochs_npy);
    FREE_INPUT_LIST_OF_ARRAYS(snrs)
352 353 354
    FREE_INPUT_LIST_OF_ARRAYS(responses)
    FREE_INPUT_LIST_OF_ARRAYS(locations)
    Py_XDECREF(horizons_npy);
355 356 357 358
    return out;
};


359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379
static void signal_amplitude_model_loop(
    char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(data))
{
    const npy_intp n = dimensions[0];

    for (npy_intp i = 0; i < n; i ++)
    {
        /* FIXME: args must be void ** to avoid alignment warnings */
        #pragma GCC diagnostic push
        #pragma GCC diagnostic ignored "-Wcast-align"
        *(double complex *) &args[4][i * steps[4]] =
            bayestar_signal_amplitude_model(
            *(double complex *) &args[0][i * steps[0]],
            *(double complex *) &args[1][i * steps[1]],
            *(double *)         &args[2][i * steps[2]],
            *(double *)         &args[3][i * steps[3]]);
        #pragma GCC diagnostic pop
    }
}


380 381 382
static PyObject *test(
    PyObject *NPY_UNUSED(module), PyObject *NPY_UNUSED(arg))
{
383
    int ret;
384
    gsl_error_handler_t *old_handler = gsl_set_error_handler_off();
385 386 387
    Py_BEGIN_ALLOW_THREADS
    ret = bayestar_test();
    Py_END_ALLOW_THREADS
388 389
    gsl_set_error_handler(old_handler);
    return PyLong_FromLong(ret);
390 391 392
}


393
static PyMethodDef methods[] = {
Leo Pound Singer's avatar
Leo Pound Singer committed
394 395
    {"toa_phoa_snr", (PyCFunction)sky_map_toa_phoa_snr,
        METH_VARARGS | METH_KEYWORDS, "fill me in"},
396
    {"log_likelihood_toa_phoa_snr", (PyCFunction)log_likelihood_toa_phoa_snr,
Leo Pound Singer's avatar
Leo Pound Singer committed
397
        METH_VARARGS | METH_KEYWORDS, "fill me in"},
398 399
    {"test", (PyCFunction)test,
        METH_NOARGS, "fill me in"},
400 401 402 403
    {NULL, NULL, 0, NULL}
};


404 405
static PyModuleDef moduledef = {
    PyModuleDef_HEAD_INIT,
406 407
    "_sky_map", NULL, -1, methods,
    NULL, NULL, NULL, NULL
408 409 410
};


411 412 413 414 415 416 417 418 419
static const PyUFuncGenericFunction
    signal_amplitude_model_loops[] = {signal_amplitude_model_loop};

static const char signal_amplitude_model_ufunc_types[] = {
    NPY_CDOUBLE, NPY_CDOUBLE, NPY_DOUBLE, NPY_DOUBLE, NPY_CDOUBLE};

static void *const no_ufunc_data[] = {NULL};


420 421
PyMODINIT_FUNC PyInit__sky_map(void); /* Silence -Wmissing-prototypes */
PyMODINIT_FUNC PyInit__sky_map(void)
422
{
423 424
    PyObject *module = NULL;

425
    gsl_set_error_handler_off();
426
    import_array();
427
    import_umath();
428

429 430 431 432
    sky_map_descr = sky_map_create_descr();
    if (!sky_map_descr)
        goto done;

433
    module = PyModule_Create(&moduledef);
434 435
    if (!module)
        goto done;
436

437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452
    /* Ignore warnings in Numpy API */
    #pragma GCC diagnostic push
    #ifdef __clang__
    #pragma GCC diagnostic ignored "-Wincompatible-pointer-types-discards-qualifiers"
    #else
    #pragma GCC diagnostic ignored "-Wdiscarded-qualifiers"
    #endif

    PyModule_AddObject(
        module, "signal_amplitude_model", PyUFunc_FromFuncAndData(
            signal_amplitude_model_loops, no_ufunc_data,
            signal_amplitude_model_ufunc_types, 1, 4, 1, PyUFunc_None,
            "signal_amplitude_model", NULL, 0));

    #pragma GCC diagnostic pop

453
done:
454
    return module;
455
}
456 457 458


SIX_COMPAT_MODULE(_sky_map)