diff --git a/empress/support_files/js/bp-tree.js b/empress/support_files/js/bp-tree.js index 7ca8a17fc..8c737eedb 100644 --- a/empress/support_files/js/bp-tree.js +++ b/empress/support_files/js/bp-tree.js @@ -65,7 +65,7 @@ define(["ByteArray", "underscore"], function (ByteArray, _) { /** * @type {Array} * @private - * stores the name of each node in preorder. If names are not provided + * stores the name of each node in postorder. If names are not provided * then the names will be set to null by default. * Note: if memory becomes an issue this could be converted into a * Uint16Array @@ -85,7 +85,7 @@ define(["ByteArray", "underscore"], function (ByteArray, _) { /** * @type {Array} * @private - * Stores the length of the nodes in preorder. If lengths are not + * Stores the length of the nodes in postorder. If lengths are not * provided then lengths will be set to null. */ this.lengths_ = lengths ? lengths : null; @@ -971,5 +971,71 @@ define(["ByteArray", "underscore"], function (ByteArray, _) { return this._nameToNodes[name]; }; + /** + * Returns a new BPTree object that contains just the tips (and ancestors) + * of the nodes in keepTips. + * + * This method was ported from iow. + * https://github.com/wasade/improved-octo-waddle/blob/0e9e75b77238acda6752f59d940620f89607ba6b/bp/_bp.pyx#L732 + * + * @param {Set} keepTips The set of tip names to keep. + * + * @return {BPTree} The new BPTree. + */ + BPTree.prototype.shear = function (keepTips) { + // closure + var scope = this; + + // create new names and lengths array + var names = [null]; + var lengths = [null]; + + // create new bit array + var mask = []; + + // function to that will set open/close bits for a node + var set_bits = (node) => { + mask[node] = 1; + mask[scope.close(node)] = 0; + }; + + // set root open/close bits + set_bits(this.root()); + + // iterate over bp tree in post order and add all tips that are in + // keepTips plus their ancestors + var i; + for (i = 1; i <= this.size; i++) { + var node = this.postorderselect(i); + var name = this.name(node); + if (this.isleaf(node) && keepTips.has(name)) { + // set open/close bits for tip + set_bits(node); + + // set open/close bits for tips ancestors + var parent = this.parent(node); + while (parent !== this.root() || mask[parent] !== 1) { + set_bits(parent); + parent = this.parent(parent); + } + } + } + + var newBitArray = []; + for (i = 0; i < mask.length; i++) { + if (mask[i] !== undefined) { + newBitArray.push(mask[i]); + } + + // get name and length of node + // Note: names and lengths of nodes are stored in postorder + if (mask[i] === 0) { + names.push(this.name(i)); + lengths.push(this.length(i)); + } + } + return new BPTree(newBitArray, names, lengths, null); + }; + return BPTree; }); diff --git a/tests/test-bp-tree.js b/tests/test-bp-tree.js index 0bfb8282c..78bff215d 100644 --- a/tests/test-bp-tree.js +++ b/tests/test-bp-tree.js @@ -877,5 +877,107 @@ require(["jquery", "ByteArray", "BPTree"], function ($, ByteArray, BPTree) { "Error thrown when no length info given" ); }); + + test("Test shear", function () { + // test modified from https://github.com/wasade/improved-octo-waddle/blob/master/bp/tests/test_bp.py#L228 + // newick represenation + // ((3,4,(6)5)2, 7, ((10, 11)9)8)r; + var preShearArr = [ + 1, + 1, + 1, + 0, + 1, + 0, + 1, + 1, + 0, + 0, + 0, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 0, + 0, + 0, + 0, + ]; + var preShearNames = [ + null, + "3", + "4", + "6", + "5", + "2", + "7", + "10", + "11", + "9", + "8", + "r", + ]; + var preShearLenghts = [null, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]; + var preShearBPTree = new BPTree( + preShearArr, + preShearNames, + preShearLenghts, + null + ); + + var keep = new Set(["4", "6", "7", "10", "11"]); + var result = preShearBPTree.shear(keep); + deepEqual(result.b_, [ + 1, + 1, + 1, + 0, + 1, + 1, + 0, + 0, + 0, + 1, + 0, + 1, + 1, + 1, + 0, + 1, + 0, + 0, + 0, + 0, + ]); + deepEqual(result.names_, [ + null, + "4", + "6", + "5", + "2", + "7", + "10", + "11", + "9", + "8", + "r", + ]); + deepEqual(result.lengths_, [null, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]); + + keep = new Set(["7", "10", "11"]); + result = preShearBPTree.shear(keep); + deepEqual(result.b_, [1, 1, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0]); + deepEqual(result.names_, [null, "7", "10", "11", "9", "8", "r"]); + deepEqual(result.lengths_, [null, 6, 7, 8, 9, 10, 11]); + + keep = new Set([]); + result = preShearBPTree.shear(keep); + deepEqual(result.b_, [1, 0]); + deepEqual(result.names_, [null, "r"]); + deepEqual(result.lengths_, [null, 11]); + }); }); });