fix bugs in balance and add soma information to cell_tree

cell_tree.hpp

@@ -49,14 +49,32 @@ class cell_tree {
     }
 
     /// construct from a tree
+    // copy constructor
     cell_tree(tree const& t)
     : tree_(t)
-    {}
+    { }
 
-    /// construct from a tree
+    // move constructor
     cell_tree(tree&& t)
     : tree_(std::move(t))
-    {}
+    { }
+
+    /// construct from a cell tree
+    // copy constructor
+    cell_tree(cell_tree const& other)
+    : tree_(other.tree_),
+      soma_(other.soma())
+    { }
+
+    // move constructor
+    cell_tree(cell_tree&& other)
+    : tree_(std::move(other.tree_)),
+      soma_(other.soma())
+    { }
+
+    int_type soma() const {
+        return soma_;
+    }
 
     /// Minimize the depth of the tree.
     int_type balance() {
@@ -64,7 +82,12 @@ class cell_tree {
         auto new_root = find_minimum_root();
 
         // change the root on the tree
-        tree_.change_root(new_root);
+        auto p = tree_.change_root(new_root);
+
+        // keep track of the soma_
+        if(p.size()) {
+            soma_ = p[soma_];
+        }
 
         return new_root;
     }
@@ -241,4 +264,6 @@ class cell_tree {
 
     // storage for the tree structure of cell segments
     tree tree_;
+
+    int_type soma_ = 0;
 };

main.cpp

@@ -138,13 +138,13 @@ TEST(tree, change_root) {
         std::vector<int> parent_index = {0,0,0};
         tree t;
         t.init_from_parent_index(parent_index);
-        auto new_tree = t.change_root(1);
+        t.change_root(1);
 
-        EXPECT_EQ(new_tree.num_nodes(), 3);
+        EXPECT_EQ(t.num_nodes(), 3);
 
-        EXPECT_EQ(new_tree.num_children(0), 1);
-        EXPECT_EQ(new_tree.num_children(1), 1);
-        EXPECT_EQ(new_tree.num_children(2), 0);
+        EXPECT_EQ(t.num_children(0), 1);
+        EXPECT_EQ(t.num_children(1), 1);
+        EXPECT_EQ(t.num_children(2), 0);
     }
     {
         // a cell with the following structure
@@ -157,15 +157,15 @@ TEST(tree, change_root) {
         std::vector<int> parent_index = {0,0,0,1,1};
         tree t;
         t.init_from_parent_index(parent_index);
-        auto new_tree = t.change_root(1);
+        t.change_root(1);
 
-        EXPECT_EQ(new_tree.num_nodes(), 5);
+        EXPECT_EQ(t.num_nodes(), 5);
 
-        EXPECT_EQ(new_tree.num_children(0), 3);
-        EXPECT_EQ(new_tree.num_children(1), 0);
-        EXPECT_EQ(new_tree.num_children(2), 0);
-        EXPECT_EQ(new_tree.num_children(3), 1);
-        EXPECT_EQ(new_tree.num_children(4), 0);
+        EXPECT_EQ(t.num_children(0), 3);
+        EXPECT_EQ(t.num_children(1), 0);
+        EXPECT_EQ(t.num_children(2), 0);
+        EXPECT_EQ(t.num_children(3), 1);
+        EXPECT_EQ(t.num_children(4), 0);
     }
     {
         // a cell with the following structure
@@ -183,19 +183,17 @@ TEST(tree, change_root) {
         tree t;
         t.init_from_parent_index(parent_index);
 
-        auto new_tree = t.change_root(1);
+        t.change_root(1);
 
-        EXPECT_EQ(new_tree.num_nodes(), 7);
+        EXPECT_EQ(t.num_nodes(), 7);
 
-        EXPECT_EQ(new_tree.num_children(0), 3);
-        EXPECT_EQ(new_tree.num_children(1), 0);
-        EXPECT_EQ(new_tree.num_children(2), 2);
-        EXPECT_EQ(new_tree.num_children(3), 0);
-        EXPECT_EQ(new_tree.num_children(4), 0);
-        EXPECT_EQ(new_tree.num_children(5), 1);
-        EXPECT_EQ(new_tree.num_children(6), 0);
-
-        auto T = cell_tree(new_tree);
+        EXPECT_EQ(t.num_children(0), 3);
+        EXPECT_EQ(t.num_children(1), 0);
+        EXPECT_EQ(t.num_children(2), 2);
+        EXPECT_EQ(t.num_children(3), 0);
+        EXPECT_EQ(t.num_children(4), 0);
+        EXPECT_EQ(t.num_children(5), 1);
+        EXPECT_EQ(t.num_children(6), 0);
     }
 }
 
@@ -215,7 +213,24 @@ TEST(cell_tree, balance) {
 
         t.balance();
 
+        // the soma (original root) has moved to 5 in the new tree
+        EXPECT_EQ(t.soma(), 5);
+
         EXPECT_EQ(t.num_segments(), 7);
+        EXPECT_EQ(t.num_children(0),3);
+        EXPECT_EQ(t.num_children(1),0);
+        EXPECT_EQ(t.num_children(2),2);
+        EXPECT_EQ(t.num_children(3),0);
+        EXPECT_EQ(t.num_children(4),0);
+        EXPECT_EQ(t.num_children(5),1);
+        EXPECT_EQ(t.num_children(6),0);
+        EXPECT_EQ(t.parent(0),-1);
+        EXPECT_EQ(t.parent(1), 0);
+        EXPECT_EQ(t.parent(2), 0);
+        EXPECT_EQ(t.parent(3), 0);
+        EXPECT_EQ(t.parent(4), 2);
+        EXPECT_EQ(t.parent(5), 2);
+        EXPECT_EQ(t.parent(6), 5);
 
         t.to_graphviz("cell.dot");
     }

tree.hpp

@@ -175,12 +175,12 @@ class tree {
         return sizeof(int_type)*data_.size() + sizeof(tree);
     }
 
-    tree change_root(int b) {
+    index_type change_root(int b) {
         assert(b<num_nodes());
 
         // no need to rebalance if the root node has been requested
         if(b==0) {
-            return tree(*this);
+            return index_type();
         }
 
         // create new tree with memory allocated
@@ -188,6 +188,7 @@ class tree {
         new_tree.init(num_nodes());
 
         // add the root node
+        new_tree.data_(memory::all) = -std::numeric_limits<int_type>::min();
         new_tree.parents_[0] = -1;
         new_tree.child_index_[0] = 0;
 
@@ -205,13 +206,12 @@ class tree {
             new_tree.children_.begin(), [&p] (int i) {return p[i];}
         );
 
-        // copy in new data
-        // this should be done with a swap, to avoid a malloc-free, however
-        // using std::swap gives a seg fault... todo
-        data_ = new_tree.data_;
-        //std::swap(data_, new_tree.data_);
+        // 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_nodes());
 
-        return new_tree;
+        return p;
     }
 
     private :
@@ -289,11 +289,13 @@ class tree {
         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;