IMP  2.0.0
The Integrative Modeling Platform
rmf_color.cpp
1 /**
2  * Copyright 2007-2013 IMP Inventors. All rights reserved.
3  */
4 #include <IMP/rmf/atom_io.h>
5 #include <RMF/FileHandle.h>
6 #include <IMP/rmf/geometry_io.h>
7 #include <IMP/display/Color.h>
8 #include <IMP/base/flags.h>
9 namespace {
10 int get_count(RMF::NodeHandle nh, int level) {
11  int ret=0;
12  if (level> 0) {
13  RMF::NodeHandles ch= nh.get_children();
14  for (unsigned int i=0; i< ch.size(); ++i) {
15  ret+= get_count(ch[i], level-1);
16  }
17  } else {
18  ret+= 1;
19  }
20  return ret;
21 }
22 
23 int recolor(RMF::NodeHandle nh, int level, int total,
24  int so_far, std::string method,
25  RMF::FloatKey red,
26  RMF::FloatKey green,
27  RMF::FloatKey blue) {
28  if (level> 0) {
29  RMF::NodeHandles ch= nh.get_children();
30  for (unsigned int i=0; i< ch.size(); ++i) {
31  so_far= recolor(ch[i], level-1, total, so_far, method,
32  red, green, blue);
33  }
34  } else {
35  double f= static_cast<double>(so_far)/total;
37  if (method =="jet") {
39  } else {
41  }
42  nh.set_value(red, c.get_red());
43  nh.set_value(green, c.get_green());
44  nh.set_value(blue, c.get_blue());
45  ++so_far;
46  }
47  return so_far;
48 }
49 }
50 
51 int main(int argc, char **argv) {
52  try {
53  boost::int64_t level=0;
54  std::string colormap="jet";
55  IMP::base::AddIntFlag lf("level",
56  "What level to apply the color to in the RMF (0 is root).",
57  &level);
58  IMP::base::AddStringFlag cf("colormap",
59  "The colormap to use: display, jet",
60  &colormap);
61  IMP::Strings files
62  = IMP::base::setup_from_argv(argc, argv, "Add color to an RMF file",
63  "file0.rmf ...", -1);
64  for (unsigned int i = 0; i < files.size(); ++i) {
65  RMF::FileHandle rh= RMF::open_rmf_file(files[i]);
66  int count= get_count(rh.get_root_node(), level);
67  RMF::Category cat= rh.get_category("shape");
68  RMF::FloatKey red
69  = rh.get_key<RMF::FloatTraits>(cat, "rgb color red");
70  RMF::FloatKey green
71  = rh.get_key<RMF::FloatTraits>(cat, "rgb color green");
72  RMF::FloatKey blue
73  = rh.get_key<RMF::FloatTraits>(cat, "rgb color blue");
74  recolor(rh.get_root_node(), level, count, 0, colormap, red, green, blue);
75  }
76  } catch (const IMP::base::Exception &e) {
77  std::cerr << "Error: " << e.what() << std::endl;
78  return 1;
79  } catch (const std::exception &e) {
80  std::cerr << "Error: " << e.what() << std::endl;
81  }
82  return 0;
83 }