assembly_graph_inl.h 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668
  1. /*
  2. * nvbio
  3. * Copyright (c) 2011-2014, NVIDIA CORPORATION. All rights reserved.
  4. *
  5. * Redistribution and use in source and binary forms, with or without
  6. * modification, are permitted provided that the following conditions are met:
  7. * * Redistributions of source code must retain the above copyright
  8. * notice, this list of conditions and the following disclaimer.
  9. * * Redistributions in binary form must reproduce the above copyright
  10. * notice, this list of conditions and the following disclaimer in the
  11. * documentation and/or other materials provided with the distribution.
  12. * * Neither the name of the NVIDIA CORPORATION nor the
  13. * names of its contributors may be used to endorse or promote products
  14. * derived from this software without specific prior written permission.
  15. *
  16. * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
  17. * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  18. * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  19. * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
  20. * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  21. * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  22. * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
  23. * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  24. * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  25. * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  26. */
  27. #include <stdio.h>
  28. #include <fstream>
  29. #include "kmers.h"
  30. #define TOPO_SORT_LOCAL_NODE_SET_SIZE 200
  31. #define MAX_IN_DEGREE 20
  32. struct graph_functor
  33. {
  34. debruijn_graph::view graph;
  35. NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
  36. graph_functor(debruijn_graph::view _graph) : graph(_graph) { }
  37. };
  38. struct topological_sort_functor : public graph_functor
  39. {
  40. NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
  41. topological_sort_functor(debruijn_graph::view _graph) : graph_functor(_graph) { }
  42. NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
  43. void operator() (const uint32 graph_id)
  44. {
  45. const uint32 source = graph.subgraph_source_ids[graph_id];
  46. const uint32 n_nodes = graph.subgraph_n_nodes[graph_id];
  47. uint32 L_offset = graph.topo_sorted_offsets[graph_id]; // ordered UID output offset
  48. uint32 n_sorted = 0; // number of nodes sorted so far
  49. uint32 S[TOPO_SORT_LOCAL_NODE_SET_SIZE]; // set of nodes with in-degree 0
  50. S[0] = source;
  51. uint32 s_first = 0;
  52. uint32 s_last = 1;
  53. while(s_last - s_first != 0) {
  54. const uint32 n = S[s_first];
  55. // remove from S
  56. s_first++;
  57. // add to ordered list
  58. graph.topo_sorted_uids[L_offset + n_sorted] = n;
  59. n_sorted++;
  60. // traverse the out neighbors
  61. uint32 out_degree = graph.node_out_degrees[n];
  62. uint32 nbr_offset = graph.node_adj_offests[n];
  63. for(uint32 i = 0; i < out_degree; i++) {
  64. uint32 m = graph.node_adjancency_map[nbr_offset + i];
  65. graph.node_in_degrees[m]--;
  66. if(graph.node_in_degrees[m] == 0) {
  67. // add to S
  68. if(s_first > 0) {
  69. s_first--;
  70. S[s_first] = m;
  71. } else {
  72. if(s_last >= TOPO_SORT_LOCAL_NODE_SET_SIZE) {
  73. printf("ERROR: Topological sort exceeded max locally allocated set size!\n");
  74. graph.cycle_presence[graph_id] = true; //TODO: handle this case
  75. }
  76. S[s_last] = m;
  77. s_last++;
  78. }
  79. }
  80. }
  81. }
  82. if(n_sorted != n_nodes) {
  83. printf("Found cycle! %u %u \n", n_sorted, n_nodes);
  84. graph.cycle_presence[graph_id] = true; // found a cycle!
  85. }
  86. graph.cycle_presence[graph_id] = false;
  87. }
  88. };
  89. struct find_k_best_paths_functor : public graph_functor
  90. {
  91. const uint32 k;
  92. uint32* topk_scores;
  93. NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
  94. find_k_best_paths_functor(debruijn_graph::view _graph, const uint32 _k, uint32* _topk_scores) :
  95. graph_functor(_graph), k(_k), topk_scores(_topk_scores) { }
  96. NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
  97. void operator() (const uint32 graph_id)
  98. {
  99. const uint32 n_nodes = graph.subgraph_n_nodes[graph_id];
  100. const uint32 order_offset = graph.topo_sorted_offsets[graph_id];
  101. for(uint32 i = 0; i < n_nodes; i++) { // for each node in topological order
  102. const uint32 n = graph.topo_sorted_uids[order_offset + i];
  103. const uint32 in_degree = graph.node_in_degrees[n];
  104. const uint32 in_nbr_offset = graph.node_in_adj_offests[n];
  105. // compute the top k paths to n
  106. uint8 nbr_top_offsets[MAX_IN_DEGREE]; // keep track of how many paths consumed per in-neighbor
  107. for(uint32 j = 0; j < k; j++) {
  108. uint32 max_score = 0;
  109. bool set_score = false;
  110. for(uint32 v = 0; v < in_degree; v++) {
  111. const uint32 m = graph.node_in_adjancency_map[in_nbr_offset + v];
  112. uint32 j_score_offset = nbr_top_offsets[v]*n_nodes + m;
  113. if(topk_scores[j_score_offset] == 0) continue;
  114. set_score = true;
  115. // find edge weight (TODO: pre-compute this)
  116. uint32 edge_weight = 0;
  117. for(uint32 l = 0; l < graph.node_out_degrees[m]; l++) {
  118. if(graph.node_adjancency_map[graph.node_adj_offests[m + l]] == m) {
  119. edge_weight = graph.edge_weights[graph.node_adj_offests[m + l]];
  120. break;
  121. }
  122. }
  123. uint32 new_score = topk_scores[j_score_offset] + edge_weight;
  124. if(new_score > max_score) {
  125. max_score = new_score;
  126. nbr_top_offsets[v]++;
  127. }
  128. }
  129. if(set_score) {
  130. topk_scores[j*n_nodes + n] = max_score;
  131. // TODO: keep pointers
  132. } else {
  133. break; // consumed all the paths to this node
  134. }
  135. }
  136. }
  137. }
  138. };
  139. struct edge_weights_functor : public graph_functor
  140. {
  141. NVBIO_HOST_DEVICE
  142. edge_weights_functor(debruijn_graph::view _graph) : graph_functor(_graph) { }
  143. NVBIO_DEVICE
  144. void operator() (const uint32 node_uid)
  145. {
  146. const uint32 nbr_offset = graph.node_adjancency_map[node_uid];
  147. const uint32 out_degree = graph.node_in_degrees[node_uid];
  148. uint32 total_support = 0; // TODO: compute as a reduction
  149. for(uint32 i = 0; i < out_degree; i++) {
  150. total_support += graph.edge_counts[nbr_offset + i];
  151. }
  152. for(uint32 i = 0; i < out_degree; i++) {
  153. graph.edge_weights[nbr_offset + i] = log10((float)graph.edge_counts[nbr_offset + i]) - log10((float)total_support);
  154. }
  155. }
  156. };
  157. struct flip_edges : public graph_functor
  158. {
  159. NVBIO_HOST_DEVICE
  160. flip_edges(debruijn_graph::view _graph) : graph_functor(_graph) { }
  161. NVBIO_DEVICE // TODO: add atomicCAS to atomics.h for both host and device
  162. void operator() (const uint32 node_uid)
  163. {
  164. const uint32 nbr_offset = graph.node_adjancency_map[node_uid];
  165. const uint32 n_nbrs = graph.node_in_degrees[node_uid];
  166. for(uint32 i = 0; i < n_nbrs; i++) {
  167. const uint32 nbr_uid = graph.node_adjancency_map[nbr_offset + i];
  168. const uint32 offset = graph.node_in_adj_offests[nbr_uid];
  169. // atomically get an index
  170. uint32 idx = 0;
  171. uint32 val = atomicCAS(&graph.node_in_adjancency_map[idx], (uint32) -1, node_uid);
  172. while(val != (uint32) -1) {
  173. idx++;
  174. val = atomicCAS(&graph.node_in_adjancency_map[idx], (uint32) -1, node_uid);
  175. }
  176. }
  177. }
  178. };
  179. // converts a coordinate to the corresponding sequence -- used for debugging
  180. template <typename string_set_type>
  181. struct coord_to_seq
  182. {
  183. typedef typename string_set_type::string_type sequence;
  184. const string_set_type string_set;
  185. const uint32 kmer_size;
  186. const SequenceSetKmerCoord* nodes;
  187. uint8* sequences;
  188. NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
  189. coord_to_seq(const uint32 _kmer_size, const string_set_type _string_set, const SequenceSetKmerCoord* _nodes, uint8* _sequences) :
  190. kmer_size(_kmer_size), string_set(_string_set), nodes(_nodes),
  191. sequences(_sequences) {}
  192. NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
  193. void operator() (const uint64 node_idx) const
  194. {
  195. const SequenceSetKmerCoord node_coord = nodes[node_idx];
  196. const sequence seq = string_set[node_coord.x];
  197. const uint32 seq_pos = node_coord.y;
  198. const uint64 offset = kmer_size * node_idx;
  199. for (uint32 i = 0; i < kmer_size; i++) {
  200. uint8 c = dna_to_char(seq[seq_pos + i]);
  201. sequences[offset + i] = c;
  202. }
  203. }
  204. };
  205. template <typename string_set_type>
  206. struct extract_active_region_source_sink
  207. {
  208. const string_set_type string_set;
  209. const SequenceSetKmerCoord* coords;
  210. const uint32* seq_active_region_ids;
  211. const uint32* global_to_UID_map;
  212. const uint32* global_to_sorted_map;
  213. uint32* source_ids;
  214. uint32* sink_ids;
  215. uint32 kmer_size;
  216. NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
  217. extract_active_region_source_sink(const string_set_type _string_set, const uint32 _kmer_size,
  218. const SequenceSetKmerCoord* _coords,
  219. const uint32* _seq_active_region_ids, const uint32* _global_to_UID_map, const uint32* _global_to_sorted_map,
  220. uint32* _source_ids, uint32* _sink_ids):
  221. string_set(_string_set), kmer_size(_kmer_size),
  222. coords(_coords), seq_active_region_ids(_seq_active_region_ids),
  223. global_to_UID_map(_global_to_UID_map), global_to_sorted_map(_global_to_sorted_map),
  224. source_ids(_source_ids), sink_ids(_sink_ids) { }
  225. NVBIO_FORCEINLINE NVBIO_HOST_DEVICE
  226. void operator() (const uint32 global_coord_idx) const
  227. {
  228. const SequenceSetKmerCoord source_coord = coords[global_to_sorted_map[global_coord_idx]];
  229. // determine if the coord is the source coord of a reference sequence (first seq in each region)
  230. if(global_coord_idx != 0 && coords[global_to_sorted_map[global_coord_idx - 1]].w == source_coord.w) return;
  231. // set the source
  232. source_ids[source_coord.w] = global_to_UID_map[source_coord.z];
  233. // find and set the sink
  234. const uint32 seq_len = string_set[source_coord.x].length();
  235. const uint32 n_kmers = seq_len - kmer_size + 1;
  236. const SequenceSetKmerCoord sink_coord = coords[global_to_sorted_map[global_coord_idx + n_kmers - 1]]; // last coord of the ref
  237. sink_ids[sink_coord.w] = global_to_UID_map[sink_coord.z];
  238. }
  239. };
  240. // construct a debruijn assembly graph with the given kmer length
  241. bool debruijn_graph::construct_graph(const D_SequenceSet& d_sequence_set, const D_VectorU32& d_active_region_ids,
  242. const uint32 _kmer_size, bool allow_low_complexity)
  243. {
  244. kmer_size = _kmer_size;
  245. D_KmerSet<D_SequenceSet> graph_kmer_set(d_sequence_set, d_active_region_ids);
  246. // --- NODE extraction ---
  247. graph_kmer_set.kmer_size = kmer_size;
  248. graph_kmer_set.gen_kmer_coords();
  249. graph_kmer_set.init_alloc_temp_space();
  250. graph_kmer_set.gen_kmer_64b_keys();
  251. graph_kmer_set.sort_kmers_by_64b_keys();
  252. graph_kmer_set.partition_kmers_by_uniqueness();
  253. graph_kmer_set.gen_global_to_sorted_id_map();
  254. graph_kmer_set.mark_unique_kmers();
  255. graph_kmer_set.extract_super_kmers(); // handle kmers that are repeats
  256. segmented_sort_super_kmers_lexicographic(
  257. graph_kmer_set.string_set,
  258. graph_kmer_set.n_super_coords,
  259. graph_kmer_set.super_coords,
  260. graph_kmer_set.super_prefix_uids,
  261. graph_kmer_set.super_suffix_uids);
  262. // extract repeat kmer nodes (as well as U-R, R-R, and R-U repeat edges)
  263. D_VectorU32 repeat_edges_from_uids;
  264. D_VectorU32 repeat_edges_to_uids;
  265. D_VectorU32 repeat_edge_counts;
  266. graph_kmer_set.collapse_and_extract_non_overlaps(nodes, repeat_edges_from_uids, repeat_edges_to_uids, repeat_edge_counts);
  267. // extract unique kmer nodes
  268. thrust::copy_n(
  269. thrust::make_permutation_iterator(
  270. thrust::make_permutation_iterator( // unique coordinates
  271. graph_kmer_set.coords.begin(),
  272. graph_kmer_set.kmers_64b_unique_idxs.begin()),
  273. thrust::make_permutation_iterator(
  274. graph_kmer_set.global_to_UID_map.begin(), // global id -> UID
  275. thrust::make_transform_iterator( // coord -> global ids
  276. thrust::make_permutation_iterator( // unique coordinates
  277. graph_kmer_set.coords.begin(),
  278. graph_kmer_set.kmers_64b_unique_idxs.begin()),
  279. get_global_id()))),
  280. graph_kmer_set.n_unique,
  281. nodes.begin());
  282. // extract the source and sink subgraph nodes (for each active region)
  283. /*n_subgraphs = d_active_region_ids[d_active_region_ids.size()-1] + 1;
  284. subgraph_source_ids.resize(n_subgraphs);
  285. subgraph_sink_ids.resize(n_subgraphs);
  286. thrust::for_each(
  287. thrust::make_counting_iterator(0u),
  288. thrust::make_counting_iterator(0u) + graph_kmer_set.n_kmers,
  289. extract_active_region_source_sink<D_SequenceSet>(d_sequence_set, kmer_size,
  290. plain_view(graph_kmer_set.coords), plain_view(d_active_region_ids),
  291. plain_view(graph_kmer_set.global_to_UID_map), plain_view(graph_kmer_set.global_to_sorted_id_map),
  292. plain_view(subgraph_source_ids), plain_view(subgraph_sink_ids)));*/
  293. printf("Number of kmers %u \n", graph_kmer_set.n_kmers);
  294. printf("Number distinct kmers %u \n", graph_kmer_set.n_distinct);
  295. printf("Number unique kmers %u \n", graph_kmer_set.n_unique);
  296. // --- ADJACENCY MAP construction ---
  297. uint32 n_unique_nodes = graph_kmer_set.n_unique;
  298. n_nodes = graph_kmer_set.n_unique + graph_kmer_set.n_repeat;
  299. node_adj_offests.resize(n_nodes + 1);
  300. node_out_degrees.resize(n_nodes); // number of outgoing edges per nodes
  301. // generate edge (k+1)mers
  302. graph_kmer_set.kmer_size = kmer_size + 1;
  303. graph_kmer_set.gen_kmer_coords();
  304. graph_kmer_set.filter_coords_by_prefix_uniqueness(graph_kmer_set.global_unique_flag_map); // keep only the edges starting with a unique kmer
  305. graph_kmer_set.gen_kmer_64b_keys();
  306. graph_kmer_set.segmented_sort_kmers_by_64b_keys(); // to guarantee same prefix edges are consecutive by region
  307. graph_kmer_set.count_kmers(); // U-U and U-R edge counts
  308. // get the prefix node UID of each U-U and U-R edge
  309. D_VectorU32 unique_prefix_node_uids(graph_kmer_set.n_unique);
  310. thrust::transform(
  311. thrust::make_transform_iterator(
  312. graph_kmer_set.kmers_64b_unique_idxs.begin(),
  313. get_prefix_global_id_by_idx(plain_view(graph_kmer_set.coords))),
  314. thrust::make_transform_iterator(
  315. graph_kmer_set.kmers_64b_unique_idxs.begin() + graph_kmer_set.n_unique,
  316. get_prefix_global_id_by_idx(plain_view(graph_kmer_set.coords))),
  317. unique_prefix_node_uids.begin(),
  318. global_to_uid(plain_view(graph_kmer_set.global_to_UID_map)));
  319. D_VectorU32 unique_suffix_node_uids(graph_kmer_set.n_unique); // TODO: fuse with above
  320. thrust::transform(
  321. thrust::make_transform_iterator(
  322. graph_kmer_set.kmers_64b_unique_idxs.begin(),
  323. get_suffix_global_id_by_idx(plain_view(graph_kmer_set.coords))),
  324. thrust::make_transform_iterator(
  325. graph_kmer_set.kmers_64b_unique_idxs.begin() + graph_kmer_set.n_unique,
  326. get_suffix_global_id_by_idx(plain_view(graph_kmer_set.coords))),
  327. unique_suffix_node_uids.begin(),
  328. global_to_uid(plain_view(graph_kmer_set.global_to_UID_map)));
  329. // count the number of outgoing edges per unique kmer (i.e. of type U-U and U-R)
  330. // reduce by UID (edges starting with the same unique kmer will be consecutive)
  331. D_VectorU32 unique_prefix_uids_idx(n_unique_nodes);
  332. D_VectorU32 unique_prefix_counts(n_unique_nodes);
  333. uint32 n_unique_prefix_uids = thrust::reduce_by_key(
  334. thrust::make_counting_iterator(0),
  335. thrust::make_counting_iterator(0) + graph_kmer_set.n_unique,
  336. thrust::constant_iterator<uint32>(1),
  337. unique_prefix_uids_idx.begin(),
  338. unique_prefix_counts.begin(),
  339. kmer_uid_eq(plain_view(unique_prefix_node_uids))).first - unique_prefix_uids_idx.begin();
  340. // scatter counts based on UIDs
  341. thrust::scatter(
  342. unique_prefix_counts.begin(),
  343. unique_prefix_counts.begin() + n_unique_prefix_uids,
  344. thrust::make_permutation_iterator(
  345. unique_prefix_node_uids.begin(),
  346. unique_prefix_uids_idx.begin()),
  347. node_out_degrees.begin());
  348. // count the number of outgoing edges per repeat kmer (i.e. of type U-R, R-R and R-U)
  349. thrust::sort_by_key(
  350. repeat_edges_from_uids.begin(),
  351. repeat_edges_from_uids.end(),
  352. thrust::make_zip_iterator(thrust::make_tuple(repeat_edges_to_uids.begin(), repeat_edge_counts.begin())));
  353. D_VectorU32 repeat_prefix_uids_idx(repeat_edges_from_uids.size());
  354. D_VectorU32 repeat_prefix_counts(repeat_edges_from_uids.size());
  355. uint32 n_repeat_prefix_uids = thrust::reduce_by_key(
  356. thrust::make_counting_iterator(0),
  357. thrust::make_counting_iterator(0) + repeat_edges_from_uids.size(),
  358. thrust::constant_iterator<uint32>(1),
  359. repeat_prefix_uids_idx.begin(),
  360. repeat_prefix_counts.begin(),
  361. kmer_uid_eq (plain_view(repeat_edges_from_uids))).first - repeat_prefix_uids_idx.begin();
  362. // we need to discard the beginning of this vector corresponding to U-R edges
  363. // by counting the number of unique prefixes
  364. thrust::device_vector<uint8> temp_storage;
  365. uint32 n_unique_prefixes = cuda::reduce(
  366. n_repeat_prefix_uids,
  367. thrust::make_transform_iterator(
  368. thrust::make_permutation_iterator(
  369. repeat_edges_from_uids.begin(),
  370. repeat_prefix_uids_idx.begin()),
  371. is_unique_uid(n_unique_nodes)),
  372. thrust::plus<uint32>(),
  373. temp_storage);
  374. // scatter counts based on UIDs
  375. thrust::scatter(
  376. repeat_prefix_counts.begin() + n_unique_prefixes,
  377. repeat_prefix_counts.begin() + n_repeat_prefix_uids,
  378. thrust::make_permutation_iterator(
  379. repeat_edges_from_uids.begin(),
  380. repeat_prefix_uids_idx.begin() + n_unique_prefixes),
  381. node_out_degrees.begin());
  382. // prefix sum to get the offsets
  383. cuda::inclusive_scan(
  384. n_nodes,
  385. node_out_degrees.begin(),
  386. node_adj_offests.begin() + 1,
  387. thrust::plus<uint32>(),
  388. temp_storage);
  389. // fill out the adjacency map: U-R, R-R, R-U edges
  390. n_edges = node_adj_offests[n_nodes];
  391. node_adjancency_map.resize(n_edges);
  392. edge_counts.resize(n_edges);
  393. thrust::for_each(
  394. repeat_prefix_uids_idx.begin(),
  395. repeat_prefix_uids_idx.begin() + n_repeat_prefix_uids,
  396. extract_repeat_adjacencies(
  397. plain_view(node_adj_offests),
  398. plain_view(repeat_edges_from_uids),
  399. plain_view(repeat_edges_to_uids),
  400. plain_view(repeat_edge_counts),
  401. plain_view(node_adjancency_map),
  402. plain_view(edge_counts)));
  403. // fill out U-U edges
  404. thrust::for_each(
  405. unique_prefix_uids_idx.begin(),
  406. unique_prefix_uids_idx.begin() + n_unique_prefix_uids,
  407. extract_unique_adjacencies(
  408. plain_view(node_adj_offests),
  409. plain_view(unique_prefix_node_uids),
  410. plain_view(unique_suffix_node_uids),
  411. plain_view(graph_kmer_set.kmers_counts),
  412. plain_view(node_adjancency_map),
  413. plain_view(edge_counts)));
  414. // // TODO: prune low weight chains
  415. // // TODO: merge chains
  416. return true;
  417. }
  418. // --- note: the functionality below has not yet been fully tested
  419. void debruijn_graph::compute_edge_weights()
  420. {
  421. edge_weights.resize(n_edges);
  422. thrust::for_each(
  423. thrust::make_counting_iterator(0u),
  424. thrust::make_counting_iterator(0u) + n_nodes,
  425. edge_weights_functor(*this));
  426. }
  427. // compute the in-degree of each node in the graph
  428. // the number of times a node UID occurs in the adjacency map
  429. // is the node's in-degree
  430. void debruijn_graph::compute_in_degrees()
  431. {
  432. // sort the adjacency map
  433. D_VectorU32 adj_map(node_adjancency_map);
  434. thrust::sort(adj_map.begin(), adj_map.end());
  435. D_VectorU32 distinct_node_UIDs(n_nodes);
  436. D_VectorU32 counts(n_nodes);
  437. uint32 n_distinct = thrust::reduce_by_key(
  438. adj_map.begin(),
  439. adj_map.end(),
  440. thrust::make_constant_iterator<uint32>(1),
  441. distinct_node_UIDs.begin(),
  442. counts.begin()).first - distinct_node_UIDs.begin();
  443. node_in_degrees.resize(n_nodes);
  444. thrust::scatter(
  445. counts.begin(),
  446. counts.begin() + n_distinct,
  447. distinct_node_UIDs.begin(),
  448. node_in_degrees.begin());
  449. // prefix sum to get the offsets
  450. thrust::device_vector<uint8> temp_storage;
  451. node_in_adj_offests.resize(n_nodes + 1);
  452. cuda::inclusive_scan(
  453. n_nodes,
  454. node_in_degrees.begin(),
  455. node_in_adj_offests.begin() + 1,
  456. thrust::plus<uint32>(),
  457. temp_storage);
  458. }
  459. void debruijn_graph::compute_in_adjacencies()
  460. {
  461. // sort the adjacency map
  462. node_in_adjancency_map.resize(n_edges);
  463. thrust::fill(node_in_adjancency_map.begin(), node_in_adjancency_map.end(), (uint32) -1);
  464. thrust::for_each(
  465. thrust::make_counting_iterator(0u),
  466. thrust::make_counting_iterator(0u) + n_nodes,
  467. flip_edges(*this));
  468. }
  469. void debruijn_graph::compute_subgraph_n_nodes()
  470. {
  471. // sort the adjacency map
  472. D_VectorU32 active_region_ids(n_nodes);
  473. thrust::transform(
  474. nodes.begin(),
  475. nodes.begin() + n_nodes,
  476. active_region_ids.begin(),
  477. get_kmer_reg_id());
  478. thrust::sort(active_region_ids.begin(), active_region_ids.end());
  479. D_VectorU32 active_regions(n_subgraphs);
  480. D_VectorU32 counts(n_subgraphs);
  481. uint32 n_distinct = thrust::reduce_by_key(
  482. active_region_ids.begin(),
  483. active_region_ids.end(),
  484. thrust::make_constant_iterator<uint32>(1),
  485. active_regions.begin(),
  486. counts.begin()).first - active_regions.begin();
  487. subgraph_n_nodes.resize(n_subgraphs);
  488. thrust::scatter(
  489. counts.begin(),
  490. counts.begin() + n_distinct,
  491. active_regions.begin(),
  492. subgraph_n_nodes.begin());
  493. // prefix sum to get the offsets
  494. thrust::device_vector<uint8> temp_storage;
  495. topo_sorted_offsets.resize(n_subgraphs + 1);
  496. cuda::inclusive_scan(
  497. n_subgraphs,
  498. subgraph_n_nodes.begin(),
  499. topo_sorted_offsets.begin() + 1,
  500. thrust::plus<uint32>(),
  501. temp_storage);
  502. }
  503. void debruijn_graph::compute_complexity()
  504. {
  505. //returns true if n_repeats * 4 > n_unique
  506. }
  507. // performs a topological sort of the nodes in each subgraph
  508. // for each subgraph: all nodes are on a path from source to sink
  509. void debruijn_graph::topological_sort()
  510. {
  511. compute_in_degrees();
  512. compute_subgraph_n_nodes();
  513. topo_sorted_uids.resize(n_nodes);
  514. cycle_presence.resize(n_subgraphs);
  515. thrust::for_each(
  516. thrust::make_counting_iterator(0u),
  517. thrust::make_counting_iterator(0u) + n_subgraphs,
  518. topological_sort_functor(*this));
  519. }
  520. void debruijn_graph::find_k_best_paths(const uint32 k)
  521. {
  522. compute_in_degrees();
  523. compute_in_adjacencies();
  524. // allocate top k space for each node
  525. D_VectorU32 topk_scores(k*n_nodes);
  526. thrust::for_each(
  527. thrust::make_counting_iterator(0u),
  528. thrust::make_counting_iterator(0u) + n_subgraphs,
  529. find_k_best_paths_functor(*this, k, plain_view(topk_scores)));
  530. }
  531. // writes a DOT graph representation
  532. void debruijn_graph::print_dot_graph(const D_SequenceSet& string_set)
  533. {
  534. D_VectorU8 d_labels(kmer_size*n_nodes);
  535. thrust::for_each(
  536. thrust::make_counting_iterator<uint64>(0u),
  537. thrust::make_counting_iterator<uint64>(0u) + n_nodes,
  538. coord_to_seq<D_SequenceSet>(kmer_size, string_set, plain_view(nodes), plain_view(d_labels)));
  539. H_VectorU8 labels = d_labels;
  540. std::ofstream dot_file;
  541. dot_file.open("cuda_test_graph.txt");
  542. for(uint64 i = 0; i < n_nodes; i++) {
  543. uint32 n_out_edges = node_adj_offests[i+1] - node_adj_offests[i];
  544. uint64 offset = node_adj_offests[i];
  545. uint64 label_offset = i*kmer_size;
  546. std::string node_seq(labels.begin() + label_offset, labels.begin() + label_offset + kmer_size);
  547. for(uint32 j = 0; j < n_out_edges; j++) {
  548. uint64 nbr_uid = node_adjancency_map[offset + j];
  549. //uint32 count = edge_counts[offset + j];
  550. label_offset = nbr_uid*kmer_size;
  551. std::string nbr_seq(labels.begin() + label_offset, labels.begin() + label_offset + kmer_size);
  552. dot_file << node_seq << "_" << i << "\t" << nbr_seq << "_" << nbr_uid << "\n";
  553. }
  554. }
  555. dot_file.close();
  556. dot_file.open("cuda_test_graph.dot");
  557. dot_file << "digraph assemblyGraphs { \n";
  558. for(uint64 i = 0; i < n_nodes; i++) {
  559. uint32 n_out_edges = node_adj_offests[i+1] - node_adj_offests[i];
  560. uint64 offset = node_adj_offests[i];
  561. uint64 label_offset = i*kmer_size;
  562. std::string node_seq(labels.begin() + label_offset, labels.begin() + label_offset + kmer_size);
  563. for(uint32 j = 0; j < n_out_edges; j++) {
  564. uint64 nbr_uid = node_adjancency_map[offset + j];
  565. uint32 count = edge_counts[offset + j];
  566. label_offset = nbr_uid*kmer_size;
  567. std::string nbr_seq(labels.begin() + label_offset, labels.begin() + label_offset + kmer_size);
  568. dot_file << node_seq << "_" << i << " -> " << nbr_seq << "_" << nbr_uid << " [label=\"" << count<< "\"];\n";
  569. }
  570. }
  571. for(uint64 i = 0; i < n_nodes; i++) {
  572. uint64 offset = i*kmer_size;
  573. std::string node_seq(labels.begin() + offset, labels.begin() + offset + kmer_size);
  574. dot_file << node_seq << "_" << i << " [label=\"" << node_seq << "\",shape=box];\n";
  575. }
  576. dot_file << "}";
  577. dot_file.close();
  578. }