Make keywords functional and allow integer output from frequency()

Author: dpryan79

README.md

@@ -103,6 +103,11 @@ This returns a dictionary with bases as keys and their frequency as values. Note
 
 The start and end position are as with the `sequence()` method described above.
 
+If integer counts are preferred, then they can instead be returned.
+
+    >>> tb.frequency("chr1", 24, 74, True)
+    {'A': 6, 'C': 6, 'T': 6, 'G': 6}
+
 ## Close a file
 
 A `TwoBit` object can be closed with the `close()` method.

lib2bit/2bit.c

@@ -272,14 +272,20 @@ void increment(char base, uint32_t *A, uint32_t *C, uint32_t *T, uint32_t *G) {
     }
 }
 
-double *twobitFrequencyWorker(TwoBit *tb, uint32_t tid, uint32_t start, uint32_t end) {
-    double *out = malloc(4 * sizeof(double));
+void *twobitFrequencyWorker(TwoBit *tb, uint32_t tid, uint32_t start, uint32_t end, int fraction) {
+    void *out;
     uint32_t sz = end - start, pos = 0;
     uint32_t A = 0, C = 0, T = 0, G = 0, len = end - start;
     uint32_t maskIdx = -1, maskStart = -1, maskEnd = -1;
     uint32_t blockStart, offset;
     uint8_t byte, base;
     int rv = 0;
+
+    if(fraction) {
+        out = malloc(4 * sizeof(double));
+    } else {
+        out = malloc(4 * sizeof(uint32_t));
+    }
     if(!out) return NULL;
 
     getMask(tb, tid, start, end, &maskIdx, &maskStart, &maskEnd);
@@ -303,10 +309,17 @@ double *twobitFrequencyWorker(TwoBit *tb, uint32_t tid, uint32_t start, uint32_t
         if(rv != 1) offset = 0;
     }
 
-    out[0] = ((double) A)/((double) len);
-    out[1] = ((double) C)/((double) len);
-    out[2] = ((double) T)/((double) len);
-    out[3] = ((double) G)/((double) len);
+    if(fraction) {
+        ((double*) out)[0] = ((double) A)/((double) len);
+        ((double*) out)[1] = ((double) C)/((double) len);
+        ((double*) out)[2] = ((double) T)/((double) len);
+        ((double*) out)[3] = ((double) G)/((double) len);
+    } else {
+        ((uint32_t*) out)[0] = A;
+        ((uint32_t*) out)[1] = C;
+        ((uint32_t*) out)[2] = T;
+        ((uint32_t*) out)[3] = G;
+    }
 
     return out;
 
@@ -315,7 +328,7 @@ double *twobitFrequencyWorker(TwoBit *tb, uint32_t tid, uint32_t start, uint32_t
     return NULL;
 }
 
-double *twobitFrequency(TwoBit *tb, char *chrom, uint32_t start, uint32_t end) {
+void *twobitFrequency(TwoBit *tb, char *chrom, uint32_t start, uint32_t end, int fraction) {
     uint32_t tid = 0, i;
 
     //Get the chromosome ID
@@ -337,7 +350,7 @@ double *twobitFrequency(TwoBit *tb, char *chrom, uint32_t start, uint32_t end) {
     if(end > tb->idx->size[tid]) return NULL;
     if(start >= end) return NULL;
 
-    return twobitFrequencyWorker(tb, tid, start, end);
+    return twobitFrequencyWorker(tb, tid, start, end, fraction);
 }
 
 /*

lib2bit/2bit.h

@@ -47,8 +47,8 @@ uint32_t twobitChromLen(TwoBit *tb, char *chrom);
 //Return the sequence of the given range (or a whole chromosome if start=end=0
 char *twobitSequence(TwoBit *tb, char *chrom, uint32_t start, uint32_t end);
 
-//We can try to assume that N is always a T and see if that's correct...
-double *twobitFrequency(TwoBit *tb, char *chrom, uint32_t start, uint32_t end);
+//Return a pointer to either 4 doubles or 4 uint32_ts holding per-base frequencies or counts.
+void *twobitFrequency(TwoBit *tb, char *chrom, uint32_t start, uint32_t end, int fractional);
 
 #ifdef __cplusplus
 }

py2bit.c

@@ -2,16 +2,17 @@
 #include <inttypes.h>
 #include "py2bit.h"
 
-static PyObject *py2bitOpen(PyObject *self, PyObject *pyArgs) {
+static PyObject *py2bitOpen(PyObject *self, PyObject *args, PyObject *kwds) {
     char *fname = NULL;
-    PyObject *pystoreMasked = Py_False;
+    PyObject *storeMaskedO = Py_False;
     pyTwoBit_t *pytb;
     int storeMasked = 0;
     TwoBit *tb = NULL;
+    static char *kwd_list[] = {"fname", "storeMasked", NULL};
 
-    if(!PyArg_ParseTuple(pyArgs, "s|O", &fname, &pystoreMasked)) goto error;
+    if(!PyArg_ParseTupleAndKeywords(args, kwds, "s|O", kwd_list, &fname, &storeMaskedO)) goto error;
 
-    if(pystoreMasked == Py_True) storeMasked = 1;
+    if(storeMaskedO == Py_True) storeMasked = 1;
 
     //Open the file
     tb = twobitOpen(fname, storeMasked);
@@ -151,14 +152,15 @@ PyObject *PyString_FromString(char *seq) {
 }
 #endif
 
-static PyObject *py2bitSequence(pyTwoBit_t *self, PyObject *args) {
+static PyObject *py2bitSequence(pyTwoBit_t *self, PyObject *args, PyObject *kwds) {
     PyObject *ret = NULL;
     TwoBit *tb = self->tb;
     char *seq, *chrom;
     unsigned long startl = 0, endl = 0;
     uint32_t start, end, len;
+    static char *kwd_list[] = {"chrom", "start", "end", NULL};
 
-    if(!PyArg_ParseTuple(args, "s|kk", &chrom, &startl, &endl)) {
+    if(!PyArg_ParseTupleAndKeywords(args, kwds, "s|kk", kwd_list, &chrom, &startl, &endl)) {
         PyErr_SetString(PyExc_RuntimeError, "You must supply at least a chromosome!");
         return NULL;
     }
@@ -191,15 +193,18 @@ static PyObject *py2bitSequence(pyTwoBit_t *self, PyObject *args) {
     return ret;
 }
 
-static PyObject *py2bitFrequency(pyTwoBit_t *self, PyObject *args) {
+static PyObject *py2bitFrequency(pyTwoBit_t *self, PyObject *args, PyObject *kwds) {
     PyObject *ret = NULL, *val = NULL;
+    PyObject *fractionO = Py_True;
     TwoBit *tb = self->tb;
     char *chrom;
-    double *o = NULL;
+    void *o = NULL;
     unsigned long startl = 0, endl = 0;
     uint32_t start, end, len;
+    static char *kwd_list[] = {"chrom", "start", "end", "fraction", NULL};
+    int fraction = 1;
 
-    if(!PyArg_ParseTuple(args, "s|kk", &chrom, &startl, &endl)) {
+    if(!PyArg_ParseTupleAndKeywords(args, kwds, "s|kkO", kwd_list, &chrom, &startl, &endl, &fractionO)) {
         PyErr_SetString(PyExc_RuntimeError, "You must supply at least a chromosome!");
         return NULL;
     }
@@ -217,7 +222,9 @@ static PyObject *py2bitFrequency(pyTwoBit_t *self, PyObject *args) {
     }
     start = (uint32_t) startl;
 
-    o = twobitFrequency(tb, chrom, start, end);
+    if(fractionO == Py_False) fraction = 0;
+
+    o = twobitFrequency(tb, chrom, start, end, fraction);
     if(!o) {
         PyErr_SetString(PyExc_RuntimeError, "Received an error while determining the per-base frequency.");
         return NULL;
@@ -227,29 +234,35 @@ static PyObject *py2bitFrequency(pyTwoBit_t *self, PyObject *args) {
     if(!ret) goto error;
 
     //A
-    val = PyFloat_FromDouble(o[0]);
+    if(fraction) val = PyFloat_FromDouble(((double*)o)[0]);
+    else val = PyLong_FromUnsignedLong(((uint32_t*)o)[0]);
     if(!val) goto error;
     if(PyDict_SetItemString(ret, "A", val) == -1) goto error;
     Py_DECREF(val);
 
     //C
-    val = PyFloat_FromDouble(o[1]);
+    if(fraction) val = PyFloat_FromDouble(((double*)o)[1]);
+    else val = PyLong_FromUnsignedLong(((uint32_t*)o)[1]);
     if(!val) goto error;
     if(PyDict_SetItemString(ret, "C", val) == -1) goto error;
     Py_DECREF(val);
 
     //T
-    val = PyFloat_FromDouble(o[2]);
+    if(fraction) val = PyFloat_FromDouble(((double*)o)[2]);
+    else val = PyLong_FromUnsignedLong(((uint32_t*)o)[2]);
     if(!val) goto error;
     if(PyDict_SetItemString(ret, "T", val) == -1) goto error;
     Py_DECREF(val);
 
     //G
-    val = PyFloat_FromDouble(o[3]);
+    if(fraction) val = PyFloat_FromDouble(((double*)o)[3]);
+    else val = PyLong_FromUnsignedLong(((uint32_t*)o)[3]);
     if(!val) goto error;
     if(PyDict_SetItemString(ret, "G", val) == -1) goto error;
     Py_DECREF(val);
 
+    free(o);
+
     return ret;
 
 error:

py2bit.h

@@ -7,12 +7,12 @@ typedef struct {
     int storeMasked; //Whether storeMasked was set. 0 = False, 1 = True
 } pyTwoBit_t;
 
-static PyObject* py2bitOpen(PyObject *self, PyObject *args);
+static PyObject* py2bitOpen(PyObject *self, PyObject *args, PyObject *kwds);
 static PyObject *py2bitInfo(pyTwoBit_t *pybw, PyObject *args);
 static PyObject* py2bitClose(pyTwoBit_t *pybw, PyObject *args);
 static PyObject* py2bitChroms(pyTwoBit_t *pybw, PyObject *args);
-static PyObject *py2bitSequence(pyTwoBit_t *pybw, PyObject *args);
-static PyObject *py2bitFrequency(pyTwoBit_t *pybw, PyObject *args);
+static PyObject *py2bitSequence(pyTwoBit_t *pybw, PyObject *args, PyObject *kwds);
+static PyObject *py2bitFrequency(pyTwoBit_t *pybw, PyObject *args, PyObject *kwds);
 static void py2bitDealloc(pyTwoBit_t *pybw);
 
 static PyMethodDef tbMethods[] = {
@@ -117,22 +117,27 @@ Positional arguments:\n\
 Keyword arguments:\n\
     start: Starting position (0-based)\n\
     end:   Ending position (1-based)\n\
+    fraction: Whether to return fractional or integer values (default 'True',\n\
+              so fractional values are returned)\n\
 \n\
 Returns:\n\
-    A dictionary with nucleotide as the key and fraction as the value.\n\
+    A dictionary with nucleotide as the key and fraction (or count) as the\n\
+    value.\n\
 \n\
 If start and end aren't specified, the entire chromosome is returned. If the\n\
 end value is beyond the end of the chromosome then it is adjusted accordingly.\n\
 \n\
 Note that the fractions will sum to much less than 1 if there are hard-masked\n\
-bases.\n\
+bases. Counts may sum to less than the length of the region for the same reason.\n\
 \n\
 >>> import py2bit\n\
 >>> tb = py2bit.open(\"test/test.2bit\")\n\
 >>> tb.frequency(tb, \"chr1\")\n\
 {'A': 0.08, 'C': 0.08, 'T': 0.08666666666666667, 'G': 0.08666666666666667}\n\
 >>> tb.frequency(tb, \"chr1\", 24, 74)\n\
 {'A': 0.12, 'C': 0.12, 'T': 0.12, 'G': 0.12}\n\
+>>> tb.frequency(tb, \"chr1\", 24, 74, True)\n\
+{'A': 6, 'C': 6, 'T': 6, 'G': 6}\n\
 >>> tb.close()"},
     {NULL, NULL, 0, NULL}
 };

py2bitTest/test.py

@@ -39,4 +39,5 @@ def testFrequency(self):
         tb = py2bit.open(self.fname, True)
         assert(tb.frequency("chr1") == {'A': 0.08, 'C': 0.08, 'T': 0.08666666666666667, 'G': 0.08666666666666667})
         assert(tb.frequency("chr1", 24, 74) == {'A': 0.12, 'C': 0.12, 'T': 0.12, 'G': 0.12})
+        assert(tb.frequency("chr1", 24, 74, False) == {'A': 6, 'C': 6, 'T': 6, 'G': 6})
         tb.close()