-
Notifications
You must be signed in to change notification settings - Fork 15
/
kc-c2.c
144 lines (132 loc) · 4.11 KB
/
kc-c2.c
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
#include <stdio.h>
#include <stdint.h>
#include <zlib.h>
#include "ketopt.h" // command-line argument parser
#include "kseq.h" // FASTA/Q parser
KSEQ_INIT(gzFile, gzread)
#include "khashl.h" // hash table
#define KC_BITS 10
#define kc_c2_eq(a, b) ((a)>>KC_BITS == (b)>>KC_BITS) // lower 8 bits for counts; higher bits for k-mer
#define kc_c2_hash(a) ((a)>>KC_BITS)
KHASHL_SET_INIT(, kc_c2_t, kc_c2, uint64_t, kc_c2_hash, kc_c2_eq)
#define CALLOC(ptr, len) ((ptr) = (__typeof__(ptr))calloc((len), sizeof(*(ptr))))
const unsigned char seq_nt4_table[256] = { // translate ACGT to 0123
0, 1, 2, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 0, 4, 1, 4, 4, 4, 2, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 0, 4, 1, 4, 4, 4, 2, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4,
4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4
};
static inline uint64_t hash64(uint64_t key, uint64_t mask) // invertible integer hash function
{
key = (~key + (key << 21)) & mask; // key = (key << 21) - key - 1;
key = key ^ key >> 24;
key = ((key + (key << 3)) + (key << 8)) & mask; // key * 265
key = key ^ key >> 14;
key = ((key + (key << 2)) + (key << 4)) & mask; // key * 21
key = key ^ key >> 28;
key = (key + (key << 31)) & mask;
return key;
}
typedef struct {
int p;
kc_c2_t **h;
} kc_c2x_t;
static kc_c2x_t *c2x_init(int p)
{
int i;
kc_c2x_t *h;
CALLOC(h, 1);
CALLOC(h->h, 1<<p);
h->p = p;
for (i = 0; i < 1<<p; ++i)
h->h[i] = kc_c2_init();
return h;
}
static inline void c2x_insert(kc_c2x_t *h, uint64_t y) // insert a k-mer $y to hash table $h
{
int absent, pre = y & ((1<<h->p) - 1);
kc_c2_t *g = h->h[pre];
khint_t k;
k = kc_c2_put(g, y>>h->p<<KC_BITS, &absent);
if ((kh_key(g, k)&0xff) < 255) ++kh_key(g, k); // count if not saturated
}
static void count_seq(kc_c2x_t *h, int k, int len, char *seq) // insert k-mers in $seq to hash table $h
{
int i, l;
uint64_t x[2], mask = (1ULL<<k*2) - 1, shift = (k - 1) * 2;
for (i = l = 0, x[0] = x[1] = 0; i < len; ++i) {
int c = seq_nt4_table[(uint8_t)seq[i]];
if (c < 4) { // not an "N" base
x[0] = (x[0] << 2 | c) & mask; // forward strand
x[1] = x[1] >> 2 | (uint64_t)(3 - c) << shift; // reverse strand
if (++l >= k) { // we find a k-mer
uint64_t y = x[0] < x[1]? x[0] : x[1];
c2x_insert(h, hash64(y, mask));
}
} else l = 0, x[0] = x[1] = 0; // if there is an "N", restart
}
}
static kc_c2x_t *count_file(const char *fn, int k, int p)
{
gzFile fp;
kseq_t *ks;
kc_c2x_t *h;
if ((fp = gzopen(fn, "r")) == 0) return 0;
ks = kseq_init(fp);
h = c2x_init(p);
while (kseq_read(ks) >= 0)
count_seq(h, k, ks->seq.l, ks->seq.s);
kseq_destroy(ks);
gzclose(fp);
return h;
}
static void print_hist(const kc_c2x_t *h)
{
khint_t k;
uint64_t cnt[256];
int i;
for (i = 0; i < 256; ++i) cnt[i] = 0;
for (i = 0; i < 1<<h->p; ++i) {
kc_c2_t *g = h->h[i];
for (k = 0; k < kh_end(g); ++k)
if (kh_exist(g, k))
++cnt[kh_key(g, k)&0xff];
}
for (i = 1; i < 256; ++i)
printf("%d\t%ld\n", i, (long)cnt[i]);
}
int main(int argc, char *argv[])
{
kc_c2x_t *h;
int i, c, k = 31, p = KC_BITS;
ketopt_t o = KETOPT_INIT;
while ((c = ketopt(&o, argc, argv, 1, "k:p:", 0)) >= 0)
if (c == 'k') k = atoi(o.arg);
else if (c == 'p') p = atoi(o.arg);
if (argc - o.ind < 1) {
fprintf(stderr, "Usage: kc-c2 [-k %d] [-p %d] <in.fa>\n", k, p);
return 1;
}
if (p < KC_BITS) {
fprintf(stderr, "ERROR: -p should be at least %d\n", KC_BITS);
return 1;
}
h = count_file(argv[o.ind], k, p);
print_hist(h);
for (i = 0; i < 1<<p; ++i)
kc_c2_destroy(h->h[i]);
free(h->h); free(h);
return 0;
}