Index: kernel/src/restraints/ConnectivityRestraint.cpp
===================================================================
--- kernel/src/restraints/ConnectivityRestraint.cpp (revision 420)
+++ kernel/src/restraints/ConnectivityRestraint.cpp (working copy)
@@ -13,9 +13,11 @@
#include "IMP/restraints/ConnectivityRestraint.h"
#include "IMP/PairScore.h"
-#include <cmath>
-#include <set>
+#include <boost/graph/kruskal_min_spanning_tree.hpp>
+#include <boost/graph/adjacency_list.hpp>
+#include <limits>
+
namespace IMP
{
@@ -67,68 +69,54 @@
{
IMP_CHECK_OBJECT(ps_.get());
unsigned int num_sets= set_offsets_.size()-1;
- typedef std::vector<std::pair<Pair, float> > Edges;
- typedef std::map<Pair, Edges> EMap;
- EMap edges;
+
+ typedef boost::adjacency_list<boost::vecS, boost::vecS,
+ boost::undirectedS, boost::no_property,
+ boost::property<boost::edge_weight_t, float> > Graph;
+ typedef boost::graph_traits<Graph>::edge_descriptor Edge;
+ typedef boost::graph_traits<Graph>::vertex_descriptor Vertex;
+ Graph g(Restraint::number_of_particles());
+
+ const float tag_weight= -std::numeric_limits<float>::max();
for (unsigned int i=0; i< num_sets; ++i) {
+ for (unsigned int ii= set_offsets_[i]+1; ii < set_offsets_[i+1]; ++ii) {
+ Edge e = boost::add_edge(ii-1, ii, g).first;
+ boost::put(boost::edge_weight_t(), g, e, tag_weight);
+ }
for (unsigned int j=0; j< i; ++j) {
- Pair cp(i,j);
for (unsigned int ii= set_offsets_[i]; ii != set_offsets_[i+1]; ++ii) {
for (unsigned int ij= set_offsets_[j]; ij != set_offsets_[j+1]; ++ij) {
- Pair pp(ii, ij);
float d= ps_->evaluate(get_particle(ii), get_particle(ij), NULL);
- edges[cp].push_back(std::make_pair(pp,d));
+ IMP_check(d != tag_weight, "ConnectivityRestraint uses the distance "
+ << tag_weight << " as a marker, so the distance function "
+ << " should not return it. This can be fixed if you "
+ << " complain.",
+ ValueException("Tag weight matches distance in"\
+ " ConnectivityRestraint"));
+ IMP_LOG(VERBOSE, "ConnectivityRestraint edge between "
+ << get_particle(ii)->get_index() << " and "
+ << get_particle(ij)->get_index() << " with weight "
+ << d << std::endl);
+ Edge e = boost::add_edge(ii, ij, g).first;
+ boost::put(boost::edge_weight_t(), g, e, d);
}
}
}
}
- std::set<int> in;
- in.insert(0);
- Edges picked_edges;
- while (in.size() != num_sets) {
- float best=std::numeric_limits<float>::max();
- Pair best_edge, best_sets;
- IMP_assert(!edges.empty(), "No more edges :-(");
- for (EMap::iterator it = edges.begin(); it != edges.end(); ++it) {
- Pair pit= it->first;
- /*IMP_LOG(VERBOSE, "Looking at edges between " << pit[0] << " "
- << pit[1] << std::endl);*/
- IMP_assert(in.find(pit[0]) == in.end()
- || in.find(pit[1]) == in.end(),
- "Edge only involves already visited nodes");
- if (in.find(pit[0]) == in.end() && in.find(pit[1]) == in.end()) {
- // want to maintain a connected tree
- continue;
- }
- for (unsigned int i=0; i< it->second.size(); ++i) {
- if (it->second[i].second < best) {
- best= it->second[i].second;
- best_edge= it->second[i].first;
- IMP_assert(best_edge != Pair(), "Bad pair");
- best_sets= it->first;
- }
- }
- }
- int newi=best_sets[0];
- if (in.find(newi) != in.end()) newi= best_sets[1];
+ std::vector<Edge> mst;
+ boost::kruskal_minimum_spanning_tree(g, std::back_inserter(mst));
- IMP_LOG(VERBOSE, "Adding set " << newi << " with score "
- << best << std::endl);
- picked_edges.push_back(std::make_pair(best_edge, best));
- in.insert(newi);
- IMP_assert(edges.find(best_sets) != edges.end(),
- "Where did they go?");
- /*IMP_LOG(VERBOSE, "Removing edges " << best_sets[0] << " "
- << best_sets[1] << std::endl);*/
- edges.erase(best_sets);
- }
-
float sum=0;
// could be more clever if accum is NULL
- for (unsigned int i=0; i< picked_edges.size(); ++i) {
- sum+= ps_->evaluate(get_particle(picked_edges[i].first[0]),
- get_particle(picked_edges[i].first[1]),
+ for (unsigned int index=0; index< mst.size(); ++index) {
+ if (boost::get(boost::edge_weight_t(), g, mst[index]) == tag_weight) continue;
+ int i= boost::target(mst[index], g);
+ int j= boost::source(mst[index], g);
+ IMP_LOG(VERBOSE, "ConnectivityRestraint edge between " << get_particle(i)->get_index()
+ << " and " << get_particle(j)->get_index() << std::endl);
+ sum+= ps_->evaluate(get_particle(i),
+ get_particle(j),
accum);
}
return sum;
Index: kernel/test/connectivity/test_connectivity.py
===================================================================
--- kernel/test/connectivity/test_connectivity.py (revision 420)
+++ kernel/test/connectivity/test_connectivity.py (working copy)
@@ -31,7 +31,8 @@
p0=IMP.Particles()
p1=IMP.Particles()
p2=IMP.Particles()
- for i in range(12):
+ p3=IMP.Particles()
+ for i in range(13):
p= IMP.Particle()
m.add_particle(p)
d= IMP.XYZDecorator.create(p)
@@ -45,6 +46,12 @@
else :
p2.append(p)
p.add_attribute(rk, 3);
+ p= IMP.Particle()
+ m.add_particle(p)
+ d= IMP.XYZDecorator.create(p)
+ d.set_coordinates_are_optimized(True)
+ p3.append(p)
+ p.add_attribute(rk, 1);
o = IMP.ConjugateGradients()
o.set_threshold(1e-4)
@@ -60,19 +67,45 @@
r.add_set(p0)
r.add_set(p1)
r.add_set(p2)
+ r.add_set(p3)
o.optimize(1000)
d01= self.min_distance(p0, p1)
d02= self.min_distance(p0, p2)
d12= self.min_distance(p1, p2)
+ d03= self.min_distance(p0, p3)
+ d13= self.min_distance(p1, p3)
+ d23= self.min_distance(p2, p3)
ok01= (d01 < 3.5)
ok02= (d02 < 4.5)
ok12= (d12 < 5.5)
+ ok03= (d03 < 2.5)
+ ok13= (d13 < 3.5)
+ ok23= (d23 < 4.5)
+ print ok01
+ print ok02
+ print ok12
+ print ok03
+ print ok13
+ print ok23
+ sum = ok01 + ok02 + ok12 + ok03 + ok13 + ok23
+ print sum
print d01
print d02
print d12
+ print d03
+ print d13
+ print d23
score= m.evaluate(False)
- self.assert_(ok01 and ok02 or ok01 and ok12 or ok12 and ok02,
- "Particles are too far")
+ self.assert_(sum >= 3,
+ "Wrong number of close pairs")
+ self.assert_(ok01 or ok02 or ok03,
+ "Point 0 is not connected")
+ self.assert_(ok01 or ok12 or ok13,
+ "Point 1 is not connected")
+ self.assert_(ok02 or ok12 or ok23,
+ "Point 2 is not connected")
+ self.assert_(ok03 or ok13 or ok23,
+ "Point 3 is not connected")
self.assert_(score < 10, "Score too high");