AI bioinformatics: porting algorithms from Python to C
Another example here of a good use case for AI in coding generall, and in scientific software in particular. One of the main tradeoffs in software development, since long before the AI era, is between high- and low-level languages.
High-level languages, like Python, let us express ideas more concisely and at a more abstract level, without worrying about the details of hardware and memory management, but tend to be less performant than low-level ones like C.
Since, for most programs, a small part of the code contributes most of the computational time, a common pattern is to write most of a program in a high-level language, and the small, computationally expensive part in a faster one. But this is a stretch for many scientific projects, where having an experienced programmer for two languages is unlikely.
We could just hand the whole project over to a coding agent, but consider a better option: write a Python implementation then use that as a reference for the C version.
Kmer counting example
By way of an example, let’s do a classic bioinformatics problem - kmer counting. Here’s our Python reference implementation in the simplest possible form:
def count_kmers_py(sequence, k):
counts = {}
for start in range(len(sequence) - k + 1):
kmer = sequence[start : start + k]
counts[kmer] = counts.get(kmer, 0) + 1
return countsWe start with a string and end up with a dictionary of kmer counts. Once we have this function written, we can do all sorts of useful things - try it out on some examples:
count_kmers_py('ATGCATCGATTAGCTAGCTATGC', 2){'AT': 4,
'TG': 2,
'GC': 4,
'CA': 1,
'TC': 1,
'CG': 1,
'GA': 1,
'TT': 1,
'TA': 3,
'AG': 2,
'CT': 2}and write some tests:
assert count_kmers_py("ACGTACGT", k=2) == {"AC": 2, "CG": 2, "GT": 2, "TA": 1}
assert count_kmers_py("AAAA", k=2) == {"AA": 3}
assert count_kmers_py("AAAA", k=3) == {"AAA": 2}
assert count_kmers_py("ACGT", k=1) == {"A": 1, "C": 1, "G": 1, "T": 1}
assert count_kmers_py("ACGT", k=4) == {"ACGT": 1}
assert count_kmers_py("ACGT", k=5) == {}We can run some benchmarks:
import random
def make_sequence(n, alphabet="ACGT"):
return "".join(random.choice(alphabet) for _ in range(n))
test_input = make_sequence(1000)
%timeit count_kmers_py(test_input, k=10)90 μs ± 931 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)But most importantly, we can study the code, show it to other people, and generally reason about its behaviour quite easily, because it’s written in such a high level language.
Want to see future experiments? sign up for the mailing list:
C implementation
Once we are happy with the behaviour of our function, and satisfied that it passes all the tests, we can hand over the function and tests together to a coding agent. For this experiment I used GLM5.2, but this should be a simple enough task for any modern model. I deliberately kept the prompt minimal:
in @file:kmers.py there is a python function. Write a faster implementation in C with Python bindings so that we can run the tests without any c hanges. Benchmark both versions to report the performance increase.
The resulting C code is, unsurprisingly, much more verbose and harder to read - I’ve put it here in collapsed form, click to expand:
#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <stdint.h>
#include <string.h>
/*
* Fast k-mer counter.
*
* Uses Python's dict directly via the C API, avoiding bytecode-interpreter
* overhead and eliminating the per-iteration method-dispatch cost of the
* pure-Python ``counts.get(kmer, 0) + 1`` pattern. Each k-mer window is
* materialised as a Python str with PyUnicode_FromStringAndSize (faster than
* Python slicing), looked up once with PyDict_GetItemWithError, and updated.
*
* For sequences with many duplicate k-mers (small k), an intermediate C hash
* table avoids creating redundant Python str objects during counting. The
* implementation picks the strategy based on the expected number of unique
* k-mers.
*/
/* ---- Strategy 1: custom open-addressing hash table (good for small k) ---- */
typedef struct {
const char *key; /* points into `sequence` buffer; not owned */
uint64_t hash; /* cached hash of key */
Py_ssize_t count;
} Entry;
static inline uint64_t fnv1a(const char *p, Py_ssize_t k) {
uint64_t h = 0xcbf29ce484222325ULL;
for (Py_ssize_t i = 0; i < k; i++) {
h ^= (uint8_t)p[i];
h *= 0x100000001b3ULL;
}
return h;
}
static PyObject *
count_via_htable(const char *seq, Py_ssize_t seqlen, long k)
{
Py_ssize_t n_windows = seqlen - k + 1;
Py_ssize_t cap = 16;
while (cap < n_windows * 2) {
cap <<= 1;
}
Entry *table = (Entry *)PyMem_Calloc(cap, sizeof(Entry));
if (!table) {
PyErr_NoMemory();
return NULL;
}
Py_ssize_t mask = cap - 1;
Py_ssize_t used = 0;
for (Py_ssize_t start = 0; start < n_windows; start++) {
const char *kmer = seq + start;
uint64_t h = fnv1a(kmer, k);
Py_ssize_t idx = (Py_ssize_t)(h & (uint64_t)mask);
while (table[idx].key != NULL) {
if (table[idx].hash == h &&
memcmp(table[idx].key, kmer, (size_t)k) == 0) {
table[idx].count++;
goto next;
}
idx = (idx + 1) & mask;
}
table[idx].key = kmer;
table[idx].hash = h;
table[idx].count = 1;
used++;
if (used * 10 > cap * 7) {
Py_ssize_t newcap = cap << 1;
Entry *newtable = (Entry *)PyMem_Calloc(newcap, sizeof(Entry));
if (!newtable) {
PyMem_Free(table);
PyErr_NoMemory();
return NULL;
}
Py_ssize_t newmask = newcap - 1;
for (Py_ssize_t i = 0; i < cap; i++) {
if (table[i].key) {
uint64_t h2 = table[i].hash;
Py_ssize_t j = (Py_ssize_t)(h2 & (uint64_t)newmask);
while (newtable[j].key) {
j = (j + 1) & newmask;
}
newtable[j] = table[i];
}
}
PyMem_Free(table);
table = newtable;
cap = newcap;
mask = newmask;
}
next:;
}
PyObject *result = PyDict_New();
if (!result) {
PyMem_Free(table);
return NULL;
}
for (Py_ssize_t i = 0; i < cap; i++) {
if (table[i].key) {
PyObject *key = PyUnicode_FromStringAndSize(table[i].key, k);
if (!key) {
Py_DECREF(result);
PyMem_Free(table);
return NULL;
}
PyObject *val = PyLong_FromSsize_t(table[i].count);
if (!val) {
Py_DECREF(key);
Py_DECREF(result);
PyMem_Free(table);
return NULL;
}
if (PyDict_SetItem(result, key, val) < 0) {
Py_DECREF(key);
Py_DECREF(val);
Py_DECREF(result);
PyMem_Free(table);
return NULL;
}
Py_DECREF(key);
Py_DECREF(val);
}
}
PyMem_Free(table);
return result;
}
/* ---- Strategy 2: Python dict directly via C API (good for large k) ---- */
static PyObject *
count_via_pydict(const char *seq, Py_ssize_t seqlen, long k)
{
Py_ssize_t n_windows = seqlen - k + 1;
PyObject *result = PyDict_New();
if (!result) return NULL;
/* Pre-allocate the integer 1; reused for first sightings. */
PyObject *one = PyLong_FromLong(1);
if (!one) {
Py_DECREF(result);
return NULL;
}
for (Py_ssize_t start = 0; start < n_windows; start++) {
PyObject *kmer = PyUnicode_FromStringAndSize(seq + start, k);
if (!kmer) {
Py_DECREF(one);
Py_DECREF(result);
return NULL;
}
/* Borrowed reference; does not raise on missing key. */
PyObject *existing = PyDict_GetItemWithError(result, kmer);
if (existing) {
long newval = PyLong_AsLong(existing) + 1;
PyObject *newobj = PyLong_FromLong(newval);
if (!newobj) {
Py_DECREF(kmer);
Py_DECREF(one);
Py_DECREF(result);
return NULL;
}
/* SetItem handles the case where key already exists (updates value). */
if (PyDict_SetItem(result, kmer, newobj) < 0) {
Py_DECREF(newobj);
Py_DECREF(kmer);
Py_DECREF(one);
Py_DECREF(result);
return NULL;
}
Py_DECREF(newobj);
} else if (PyErr_Occurred()) {
Py_DECREF(kmer);
Py_DECREF(one);
Py_DECREF(result);
return NULL;
} else {
/* Key not present; insert 1. */
if (PyDict_SetItem(result, kmer, one) < 0) {
Py_DECREF(kmer);
Py_DECREF(one);
Py_DECREF(result);
return NULL;
}
}
Py_DECREF(kmer);
}
Py_DECREF(one);
return result;
}
/* ---- Dispatch ---- */
/*
* Heuristic: when the number of possible k-mers (4^k for DNA) is much larger
* than the sequence length, almost every k-mer is unique and the custom hash
* table just adds overhead on top of building the Python dict. In that case,
* go straight to the Python-dict strategy. Otherwise, use the custom table to
* avoid creating redundant Python str objects for duplicates.
*
* The threshold is: if 4^k > n_windows, use pydict strategy.
*/
static PyObject *
count_kmers(PyObject *self, PyObject *args)
{
const char *seq;
Py_ssize_t seqlen;
long k;
if (!PyArg_ParseTuple(args, "s#l", &seq, &seqlen, &k)) {
return NULL;
}
if (k < 0) {
PyErr_SetString(PyExc_ValueError, "k must be non-negative");
return NULL;
}
if (k == 0 || k > seqlen) {
return PyDict_New();
}
Py_ssize_t n_windows = seqlen - k + 1;
/* Compute 4^k, capping at a large value to avoid overflow. */
double possible = 1.0;
for (long i = 0; i < k; i++) {
possible *= 4.0;
if (possible > 1e18) break;
}
if (possible > (double)n_windows * 4.0) {
return count_via_pydict(seq, seqlen, k);
} else {
return count_via_htable(seq, seqlen, k);
}
}
static PyMethodDef CkmersMethods[] = {
{"count_kmers", count_kmers, METH_VARARGS,
"Count k-mers in a sequence (C implementation)."},
{NULL, NULL, 0, NULL}
};
static struct PyModuleDef ckmersmodule = {
PyModuleDef_HEAD_INIT,
"ckmers",
"Fast C implementation of k-mer counting.",
-1,
CkmersMethods
};
PyMODINIT_FUNC
PyInit_ckmers(void)
{
return PyModule_Create(&ckmersmodule);
}But it passes the same tests, and is an easy drop-in replacement for the original Python function, while having better performance:
seq_len k python (s) C (s) speedup
10000 5 0.0009 0.0002 5.8x
100000 7 0.0104 0.0030 3.5x
1000000 10 0.2865 0.2054 1.4x
5000000 15 1.7603 1.5808 1.1x
5000000 8 0.6906 0.2488 2.8xBetter than vibe-coding the whole thing
What advantage does this approach have versus just asking an agent for the C version straight away? A few things jump out:
because we started with the Python version, we get a chance to explore the problem space in a high-level language (not particularly exciting in this well-known example, but a bigger deal for real research problems)
similarly, we get to figure out any ambiguities in the output we are looking for, and add tests specifically to handle them. For example, maybe we want to exclude ambiguous bases or add some error handling
because the coding agent gets access to the Python implementation in its context, it implicitly gets all the details about the exact behaviour we want
because the coding agent gets access to the tests, it can check and correct any errors in the C implementation as it is progressing
because we can trivially run the C implementation against the existing tests, we can be very confident at the end that we have a drop-in replacement for our original code