-
Benjamin Cumming authored5249a934
cell_tree.hpp 7.62 KiB
#pragma once
#include <algorithm>
#include <iomanip>
#include <iostream>
#include <fstream>
#include <memory>
#include <numeric>
#include <ostream>
#include <vector>
#include "vector/include/Vector.hpp"
#include "tree.hpp"
#include "util.hpp"
/// The tree data structure that describes the segments of a cell tree.
/// A cell is represented as a tree where each node may have any number of
/// children. Typically in a cell only the soma has more than two segments,
/// however this rule does not appear to be strictly followed in the PCP data
/// sets.
///
/// To optimize some computations it is important the the tree be balanced,
/// in the sense that the depth of the tree be minimized. This means that it is
/// necessary that any node in the tree may be used as the root. In the PCP data
/// sets it appears that the soma was always index 0, however we need more
/// flexibility in choosing the root.
class cell_tree {
public :
// use a signed 16-bit integer for storage of indexes, which is reasonable given
// that typical cells have at most 1000-2000 segments
using int_type = int16_t;
using index_type = memory::HostVector<int_type>;
using index_view = index_type::view_type;
/// construct from a parent index
cell_tree(std::vector<int> const& parent_index)
{
// handle the case of an empty parent list, which implies a single-compartment model
std::vector<int> segment_index;
if(parent_index.size()>0) {
segment_index = tree_.init_from_parent_index(parent_index);
}
else {
segment_index = tree_.init_from_parent_index(std::vector<int>({0}));
}
// if needed, calculate meta-data like length[] and end[] arrays for data
}
/// construct from a tree
// copy constructor
cell_tree(tree const& t)
: tree_(t)
{ }
// 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() {
// find the new root
auto new_root = find_minimum_root();
// change the root on the tree
auto p = tree_.change_root(new_root);
// keep track of the soma_
if(p.size()) {
soma_ = p[soma_];
}
return new_root;
}
/// memory used to store cell tree (in bytes)
size_t memory() const {
return tree_.memory() + sizeof(cell_tree) - sizeof(tree);
}
/// returns the number of segments in the cell
size_t num_segments() const {
return tree_.num_nodes();
}
/// returns the number of child segments of segment b
size_t num_children(size_t b) const {
return tree_.num_children(b);
}
/// returns a list of the children of segment b
const index_view children(size_t b) const {
return tree_.children(b);
}
/// returns the parent index of segment b
int_type parent(size_t b) const {
return tree_.parent(b);
}
/// writes a graphviz .dot file that visualizes cell segment structure
void to_graphviz(std::string const& fname) const {
std::ofstream fid(fname);
fid << "graph cell {" << std::endl;
for(auto b : range(0,num_segments())) {
if(children(b).size()) {
for(auto c : children(b)) {
fid << " " << b << " -- " << c << ";" << std::endl;
}
}
}
fid << "}" << std::endl;
}
index_type depth_from_leaf()
{
tree::index_type depth(num_segments());
depth_from_leaf(depth, 0);
return depth;
}
index_type depth_from_root()
{
tree::index_type depth(num_segments());
depth[0] = 0;
depth_from_root(depth, 1);
return depth;
}
private :
/// helper type for sub-tree computation
/// use in balance()
struct sub_tree {
sub_tree(int r, int diam, int dpth)
: root(r), diameter(diam), depth(dpth)
{}
void set(int r, int diam, int dpth)
{
root = r;
diameter = diam;
depth = dpth;
}
std::string to_string() const {
std::string s;
s += "[" + std::to_string(root) + ","
+ std::to_string(diameter) + "," + std::to_string(depth) + "]";
return s;
}
int root;
int diameter;
int depth;
};
/// returns the index of the segment that would minimise the depth of the
/// tree if used as the root segment
int_type find_minimum_root() {
if(num_segments()==1) {
return 0;
}
// calculate the depth of each segment from the root
// pre-order traversal of the tree
auto depth = depth_from_root();
auto max = std::max_element(depth.begin(), depth.end());
auto max_leaf = std::distance(depth.begin(), max);
auto original_depth = *max;
// Calculate the depth of each compartment as the maximum distance
// from a child leaf
// post-order traversal of the tree
depth = depth_from_leaf();
// Walk back from the deepest leaf towards the root.
// At each node test to find the deepest sub-tree that doesn't include
// node max_leaf. Then check if the total diameter of this sub-tree and
// the sub-tree and node max_leaf is the largest found so far. When the
// walk has been completed to the root node, the node that has been
// selected will be the root of the sub-tree with the largest diameter.
sub_tree max_sub_tree(0, 0, 0);
auto distance_from_max_leaf = 1;
auto pnt = max_leaf;
auto pos = parent(max_leaf);
while(pos != -1) {
for(auto c : children(pos)) {
if(c!=pnt) {
auto diameter = depth[c] + 1 + distance_from_max_leaf;
if(diameter>max_sub_tree.diameter) {
max_sub_tree.set(pos, diameter, distance_from_max_leaf);
}
}
}
pnt = pos;
pos = parent(pos);
++distance_from_max_leaf;
}
// calculate the depth of the balanced tree
auto new_depth = (max_sub_tree.diameter+1) / 2;
// nothing to do if the current root is also the root of the
// balanced tree
if(max_sub_tree.root==0 && max_sub_tree.depth==new_depth) {
return 0;
}
// perform another walk from max leaf towards max_sub_tree.root
auto count = new_depth;
auto new_root = max_leaf;
while(count) {
new_root = parent(new_root);
--count;
}
return new_root;
}
int_type depth_from_leaf(index_type& depth, int segment)
{
int_type max_depth = 0;
for(auto c : children(segment)) {
max_depth = std::max(max_depth, depth_from_leaf(depth, c));
}
depth[segment] = max_depth;
return max_depth+1;
}
void depth_from_root(index_type& depth, int segment)
{
auto d = depth[parent(segment)] + 1;
depth[segment] = d;
for(auto c : children(segment)) {
depth_from_root(depth, c);
}
}
// storage for the tree structure of cell segments
tree tree_;
int_type soma_ = 0;
};