python_cluster_reader/src/ClusterReader.c

#include "ClusterReader.h"
#include "arr_desc.h"
#include "cluster_reader.h"
#include "data_types.h"

//clang-format off
typedef struct {
    PyObject_HEAD FILE *fp;
    int n_left;
    Py_ssize_t chunk;
} ClusterFileReader;
//clang-format on

// Constructor: sets the fp to NULL then tries to open the file
// raises python exception if something goes wrong
// returned object should mean file is open and ready to read
static int ClusterFileReader_init(ClusterFileReader *self, PyObject *args,
                                  PyObject *kwds) {

    // Parse file name, accepts string or pathlike objects
    const char *fname = NULL;
    self->n_left = 0;
    self->chunk = 0;
    PyObject *fname_obj = NULL;
    PyObject *fname_bytes = NULL;
    Py_ssize_t len;

    static char *kwlist[] = {"fname", "chunk", NULL};

    // if (!PyArg_ParseTuple(args, "O&", PyUnicode_FSConverter, &buf))
    //     return -1;
    // PyBytes_AsStringAndSize(buf, &fname, &len);

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|n", kwlist, &fname_obj,
                                     &self->chunk)) {
        return -1;
    }

    if (fname_obj != Py_None)
        if (!PyUnicode_FSConverter(fname_obj, &fname_bytes))
            return -1;

    PyBytes_AsStringAndSize(fname_bytes, &fname, &len);
#ifdef CR_VERBOSE
    printf("Opening: %s\n chunk: %lu\n", fname, self->chunk);
#endif

    self->fp = fopen(fname, "rb");
    

    // Keep the return code to not return before releasing buffer
    int rc = 0;

    // Raise python exception using information from errno
    if (self->fp == NULL) {
        PyErr_SetFromErrnoWithFilename(PyExc_OSError, fname);
        rc = -1;
    }
    // Release buffer
    Py_DECREF(fname_bytes);

    // Success or fail
    return rc;
}

// Custom destructor to make sure we close the file
static void ClusterFileReader_dealloc(ClusterFileReader *self) {
    if (self->fp) {
#ifdef CR_VERBOSE
        printf("Closing file\n");
#endif
        fclose(self->fp);
        self->fp = NULL;
    }
    Py_TYPE(self)->tp_free((PyObject *)self);
}

// read method
static PyObject *ClusterFileReader_read(ClusterFileReader *self,
                                        PyObject *args) {

    const int ndim = 1;
    Py_ssize_t size = 0;
    PyObject *noise_obj = NULL;
    PyObject *noise_array = NULL;
    if (!PyArg_ParseTuple(args, "|nO", &size, &noise_obj)) {
        PyErr_SetString(PyExc_TypeError, "Could not parse args.");
        return NULL;
    }

    // Fall back on object default/config
    if (size == 0)
        size = self->chunk;

    npy_intp dims[] = {size};

    // If possible numpy will
    // use the underlying buffer, otherwise it will create a copy, for example
    // if data type is different or we pass in a list. The
    // NPY_ARRAY_C_CONTIGUOUS flag ensures that we have contiguous memory.

#ifdef CR_VERBOSE
    printf("Getting ready to read: %lu clusters. Noise map: %p\n", size,
           noise_obj);
#endif

    // If the user passed a noise map we fetch a pointer to that array as well
    int nx = 0, ny = 0;
    double *noise_map = NULL;
    if (noise_obj) {
        noise_array =
            PyArray_FROM_OTF(noise_obj, NPY_DOUBLE, NPY_ARRAY_C_CONTIGUOUS);

        int ndim_noise = PyArray_NDIM((PyArrayObject *)(noise_array));
        npy_intp *noise_shape = PyArray_SHAPE((PyArrayObject *)(noise_array));

        // For the C++ function call we need pointers (or another C++ type/data
        // structure)

        noise_map = (double *)(PyArray_DATA((PyArrayObject *)(noise_array)));

        /* for (int i=0; i< ndim_noise; i++) { */
        /*   printf("Dimension %d size %d pointer \n",i,noise_shape[i],
         * noise_map); */

        /* } */

        if (ndim_noise == 2) {

            nx = noise_shape[0];
            ny = noise_shape[1];

            //	printf("Noise map found size %d %d %d\n",nx,ny,noise_map);

        } else {
            nx = 0;
            if (ndim_noise == 1)
                nx = noise_shape[0];
            ny = 0;
            noise_map = NULL;
            //	printf("NO Noise map found %d %d %d
            //%d\n",ndim_noise,nx,ny,noise_map);
        }
    }

    // Create an uninitialized numpy array
    PyObject *clusters = PyArray_SimpleNewFromDescr(ndim, dims, cluster_dt());

    // Fill with zeros
    PyArray_FILLWBYTE((PyArrayObject *)clusters, 0);

    // Get a pointer to the array memory
    void *buf = PyArray_DATA((PyArrayObject *)clusters);

    // Call the standalone C code to read clusters from file
    // Here goes the looping, removing frame numbers etc.
    int n_read = 0;
    if (noise_map)
        n_read = read_clusters_with_cut(self->fp, size, buf, &self->n_left,
                                        noise_map, nx, ny);
    else
        n_read = read_clusters(self->fp, size, buf, &self->n_left);

    if (n_read != size) {
        // resize the array to match the number of read photons
        // this will reallocate memory

        // create a new_shape struct on the stack
        PyArray_Dims new_shape;

        // reuse dims for the shape
        dims[0] = n_read;
        new_shape.ptr = dims;
        new_shape.len = 1;

        // resize the array to match the number of clusters read
        PyArray_Resize((PyArrayObject *)clusters, &new_shape, 1, NPY_ANYORDER);
    }

    return clusters;
}

// List all methods in our ClusterFileReader class
static PyMethodDef ClusterFileReader_methods[] = {
    {"read", (PyCFunction)ClusterFileReader_read, METH_VARARGS,
     "Read clusters"},
    // {"clusterize", (PyCFunction)ClusterFileReader_clusterize, METH_VARARGS,
    //  "Analyze clusters"},
    {NULL, NULL, 0, NULL} /* Sentinel */
};

// Class defenition
static PyTypeObject ClusterFileReaderType = {
    PyVarObject_HEAD_INIT(NULL, 0).tp_name = "creader.ClusterFileReader",
    .tp_doc = PyDoc_STR("ClusterFileReader implemented in C"),
    .tp_basicsize = sizeof(ClusterFileReader),
    .tp_itemsize = 0,
    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
    .tp_new = PyType_GenericNew,
    .tp_dealloc = (destructor)ClusterFileReader_dealloc,
    .tp_init = (initproc)ClusterFileReader_init,
    .tp_methods = ClusterFileReader_methods,
};

PyObject *init_ClusterFileReader(PyObject *m) {
    import_array();
    if (PyType_Ready(&ClusterFileReaderType) < 0)
        return NULL;

    Py_INCREF(&ClusterFileReaderType);
    if (PyModule_AddObject(m, "ClusterFileReader",
                           (PyObject *)&ClusterFileReaderType) < 0) {
        Py_DECREF(&ClusterFileReaderType);
        Py_DECREF(m);
        return NULL;
    }
    return m;
}