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 -#include +#include +#include +#include + namespace IMP { @@ -67,68 +69,54 @@ { IMP_CHECK_OBJECT(ps_.get()); unsigned int num_sets= set_offsets_.size()-1; - typedef std::vector > Edges; - typedef std::map EMap; - EMap edges; + + typedef boost::adjacency_list > Graph; + typedef boost::graph_traits::edge_descriptor Edge; + typedef boost::graph_traits::vertex_descriptor Vertex; + Graph g(Restraint::number_of_particles()); + + const float tag_weight= -std::numeric_limits::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 in; - in.insert(0); - Edges picked_edges; - while (in.size() != num_sets) { - float best=std::numeric_limits::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 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");