diff --git a/arbor/tree.hpp b/arbor/tree.hpp
index af0fb4bbee783cde19114480247d6006cb9700eb..5adb6f6d92995ec8064ac975daf03120954713d1 100644
--- a/arbor/tree.hpp
+++ b/arbor/tree.hpp
@@ -18,15 +18,12 @@ public:
using int_type = cell_lid_type;
using size_type = cell_local_size_type;
- using iarray = memory::host_vector<int_type>;
- using view_type = typename iarray::view_type;
- using const_view_type = typename iarray::const_view_type;
+ using iarray = std::vector<int_type>;
static constexpr int_type no_parent = (int_type)-1;
tree() = default;
tree& operator=(tree&& other) {
- std::swap(data_, other.data_);
std::swap(child_index_, other.child_index_);
std::swap(children_, other.children_);
std::swap(parents_, other.parents_);
@@ -34,8 +31,9 @@ public:
}
tree& operator=(tree const& other) {
- data_ = other.data_;
- set_ranges(other.num_segments());
+ children_ = other.children_;
+ child_index_ = other.child_index_;
+ parents_ = other.child_index_;
return *this;
}
@@ -92,22 +90,24 @@ public:
}
/// return the child index
- const_view_type child_index() {
+ const iarray& child_index() {
return child_index_;
}
/// return the list of all children
- const_view_type children() const {
+ const iarray& children() const {
return children_;
}
- /// return the list of all children of branch b
- const_view_type children(size_type b) const {
- return children_(child_index_[b], child_index_[b+1]);
+ /// return the list of all children of branch i
+ auto children(size_type i) const {
+ const auto b = child_index_[i];
+ const auto e = child_index_[i+1];
+ return util::subrange_view(children_, b, e);
}
/// return the list of parents
- const_view_type parents() const {
+ const iarray& parents() const {
return parents_;
}
@@ -121,166 +121,29 @@ public:
/// memory used to store tree (in bytes)
std::size_t memory() const {
- return sizeof(int_type)*data_.size() + sizeof(tree);
- }
-
- iarray change_root(size_t b) {
- assert(b<num_segments());
-
- // no need to rebalance if the root node has been requested
- if(b==0) {
- return iarray();
- }
-
- // create new tree with memory allocated
- tree new_tree;
- new_tree.init(num_segments());
-
- // add the root node
- new_tree.parents_[0] = no_parent;
- new_tree.child_index_[0] = 0;
-
- // allocate space for the permutation vector that
- // will represent the permutation performed on the branches
- // during the rebalancing
- iarray p(num_segments(), -1);
-
- // recersively rebalance the tree
- new_tree.add_children(0, b, 0, p, *this);
-
- // renumber the child indexes
- std::transform(
- new_tree.children_.begin(), new_tree.children_.end(),
- new_tree.children_.begin(), [&p] (int i) {return p[i];}
- );
-
- // copy in new data with a move because the information in
- // new_tree is not kept
- std::swap(data_, new_tree.data_);
- set_ranges(new_tree.num_segments());
-
- return p;
+ return sizeof(int_type)*(children_.size()+child_index_.size()+parents_.size())
+ + sizeof(tree);
}
private:
void init(size_type nnode) {
- auto nchild = nnode - 1;
-
- data_ = iarray(nchild + (nnode + 1) + nnode);
- set_ranges(nnode);
- }
-
- void set_ranges(size_type nnode) {
if (nnode) {
auto nchild = nnode - 1;
- // data_ is partitioned as follows:
- // data_ = [children_[nchild], child_index_[nnode+1], parents_[nnode]]
- assert(data_.size() == unsigned(nchild + (nnode+1) + nnode));
- children_ = data_(0, nchild);
- child_index_ = data_(nchild, nchild+nnode+1);
- parents_ = data_(nchild+nnode+1, memory::end);
-
- // check that arrays have appropriate size
- // this should be moved into a unit test
- assert(children_.size() == unsigned(nchild));
- assert(child_index_.size() == unsigned(nnode+1));
- assert(parents_.size() == unsigned(nnode));
+ children_.resize(nchild);
+ child_index_.resize(nnode+1);
+ parents_.resize(nnode);
}
else {
- children_ = data_(0, 0);
- child_index_ = data_(0, 0);
- parents_ = data_(0, 0);
- }
- }
-
- /// Renumber the sub-tree with old_node as its root with new_node as
- /// the new index of old_node. This is a helper function for the
- /// renumbering required when a new root node is chosen to improve
- /// the balance of the tree.
- /// Optionally add the parent of old_node as a child of new_node, which
- /// will be applied recursively until the old root has been processed,
- /// which indicates that the renumbering is finished.
- ///
- /// precondition - the node new_node has already been placed in the tree
- /// precondition - all of new_node's children have been added to the tree
- /// new_node : the new index of the node whose children are to be added
- /// to the tree
- /// old_node : the index of new_node in the original tree
- /// parent_node : equals index of old_node's parent in the original tree
- /// should be a child of new_node
- /// : equals -1 if the old_node's parent is not a child of
- /// new_node
- /// p : permutation vector, p[i] is the new index of node i in the old
- /// tree
- int_type add_children(
- int_type new_node,
- int_type old_node,
- int_type parent_node,
- view_type p,
- tree const& old_tree
- )
- {
- // check for the sentinel that indicates that the old root has
- // been processed
- if (old_node==no_parent) {
- return new_node;
+ children_.resize(0);
+ child_index_.resize(0);
+ parents_.resize(0);
}
-
- p[old_node] = new_node;
-
- // the list of the children of the original node
- auto old_children = old_tree.children(old_node);
-
- auto this_node = new_node;
- auto pos = child_index_[this_node];
-
- auto add_parent_as_child = parent_node!=no_parent && old_node>0;
- //
- // STEP 1 : add the child indexes for this_node
- //
- // first add the children of the node
- for(auto b : old_children) {
- if(b != parent_node) {
- children_[pos++] = b;
- parents_[pos] = new_node;
- }
- }
- // then add the node's parent as a child if applicable
- if(add_parent_as_child) {
- children_[pos++] = old_tree.parent(old_node);
- parents_[pos] = new_node;
- }
- child_index_[this_node+1] = pos;
-
- //
- // STEP 2 : recursively add each child's children
- //
- new_node++;
- for (auto b : old_children) {
- if (b != parent_node) {
- new_node = add_children(new_node, b, no_parent, p, old_tree);
- }
- }
- if (add_parent_as_child) {
- new_node =
- add_children(
- new_node, old_tree.parent(old_node), old_node, p, old_tree
- );
- }
-
- return new_node;
}
- //////////////////////////////////////////////////
// state
- //////////////////////////////////////////////////
- iarray data_;
-
- // provide default parameters so that tree type can
- // be default constructed
- view_type children_ = data_(0, 0);
- view_type child_index_= data_(0, 0);
- view_type parents_ = data_(0, 0);
+ iarray children_;
+ iarray child_index_;
+ iarray parents_;
};
template <typename C>
diff --git a/test/unit/test_tree.cpp b/test/unit/test_tree.cpp
index 9fea9278f3a6d265560d7c24671a78379483ad85..e23cc8ed7d9e2b1c5240b3a685001414b93da0a1 100644
--- a/test/unit/test_tree.cpp
+++ b/test/unit/test_tree.cpp
@@ -164,6 +164,13 @@ TEST(tree, from_segment_index) {
EXPECT_EQ(tree.num_children(4), 2u);
EXPECT_EQ(tree.num_children(5), 0u);
EXPECT_EQ(tree.num_children(6), 0u);
+
+ EXPECT_EQ(tree.children(0)[0], 1u);
+ EXPECT_EQ(tree.children(0)[1], 2u);
+ EXPECT_EQ(tree.children(1)[0], 3u);
+ EXPECT_EQ(tree.children(1)[1], 4u);
+ EXPECT_EQ(tree.children(4)[0], 5u);
+ EXPECT_EQ(tree.children(4)[1], 6u);
}
}