betere distancemap

git-svn-id: svn+ssh://gitlab/srv/svn-repos/pdb-redo/trunk@347 a1961a4f-ab94-4bcc-80e8-33b5a54de466

include/cif++/DistanceMap.h

@@ -16,7 +16,8 @@ namespace mmcif
 class DistanceMap
 {
   public:
-	DistanceMap(const Structure& p, const clipper::Spacegroup& spacegroup, const clipper::Cell& cell);
+	DistanceMap(const Structure& p, const clipper::Spacegroup& spacegroup, const clipper::Cell& cell,
+		float maxDistance);
 
 	// simplified version for subsets of atoms (used in refining e.g.)	
 //	DistanceMap(const Structure& p, const std::vector<Atom>& atoms);
@@ -24,29 +25,25 @@ class DistanceMap
 	DistanceMap(const DistanceMap&) = delete;
 	DistanceMap& operator=(const DistanceMap&) = delete;
 
-	float operator()(const Atom& a, const Atom& b) const;
+//	float operator()(const Atom& a, const Atom& b) const;
 
 	std::vector<Atom> near(const Atom& a, float maxDistance = 3.5f) const;
 	std::vector<Atom> near(const Point& p, float maxDistance = 3.5f) const;
 
   private:
 
-	const Structure& structure;
-	size_t dim;
-	std::unordered_map<std::string,size_t> index;
-	std::map<size_t,std::string> rIndex;
+	typedef std::map<std::tuple<size_t,size_t>,std::tuple<float,int32>> DistMap;
+
+	void AddDistancesForAtoms(const Residue& a, const Residue& b, DistMap& dm, int32 rtix);
+
+	const Structure&						structure;
+	size_t									dim;
+	std::unordered_map<std::string,size_t>	index;
+	std::map<size_t,std::string>			rIndex;
 	
-	struct key_hash
-	{
-		size_t operator()(std::tuple<size_t,size_t> const& k) const noexcept
-		{
-			size_t i[2];
-			std::tie(i[0], i[1]) = k;
-			return boost::hash_range(i, i + 1);
-		}
-	};
+	float									mMaxDistance, mMaxDistanceSQ;
 	
-	std::vector<std::tuple<float,size_t>>	mA;
+	std::vector<std::tuple<float,int32>>	mA;
 	std::vector<size_t>						mIA, mJA;
 	Point									mD;			// needed to move atoms to center
 	std::vector<clipper::RTop_orth>			mRtOrth;

include/cif++/Structure.h

@@ -2,6 +2,8 @@
 
 #pragma once
 
+#include <numeric>
+
 #include <boost/filesystem/operations.hpp>
 #include <boost/math/quaternion.hpp>
 
@@ -165,10 +167,12 @@ typedef std::vector<Atom> AtomView;
 class Residue
 {
   public:
+	// constructors should be private, but that's not possible for now (needed in emplace)
+	// constructor for waters
+	Residue(const Structure& structure, const std::string& asymID, Atom water, int ndbSeqID);
+
 	Residue(const Structure& structure, const std::string& compoundID,
-		const std::string& asymID, int seqID = 0)
-		: mStructure(&structure), mCompoundID(compoundID)
-		, mAsymID(asymID), mSeqID(seqID) {}
+		const std::string& asymID, int seqID = 0);
 
 	Residue(const Residue& rhs) = delete;
 	Residue& operator=(const Residue& rhs) = delete;
@@ -204,6 +208,9 @@ class Residue
 	const Structure& structure() const		{ return *mStructure; }
 
 	bool empty() const						{ return mStructure == nullptr; }
+	
+	// some routines for 3d work
+	std::tuple<Point,float> centerAndRadius() const;
 
   protected:
 
@@ -211,9 +218,11 @@ class Residue
 
 	friend class Polymer;
 
+	void calculateCenterAndRadius();
+
 	const Structure* mStructure = nullptr;
 	std::string	mCompoundID, mAsymID;
-	int mSeqID = 0;
+	int mSeqID = 0, mNDBSeqID = 0;
 	AtomView mAtoms;
 };
 
@@ -338,12 +347,11 @@ class Structure
 
 	File& getFile() const;
 
-	const AtomView& atoms() const;
+	const AtomView& atoms() const							{ return mAtoms; }
 	AtomView waters() const;
 	
-	const std::list<Polymer>& polymers() const;
-
-	std::vector<Residue> nonPolymers() const;
+	const std::list<Polymer>& polymers() const				{ return mPolymers; }
+	const std::vector<Residue>& nonPolymers() const			{ return mNonPolymers; }
 
 	Atom getAtomById(std::string id) const;
 	Atom getAtomByLocation(Point pt, float maxDistance) const;
@@ -381,22 +389,80 @@ class Structure
 	void changeResidue(Residue& res, const std::string& newCompound,
 		const std::vector<std::tuple<std::string,std::string>>& remappedAtoms);
 	
+	// iterator for all residues
+	
+	class residue_iterator : public std::iterator<std::forward_iterator_tag, const Residue>
+	{
+	  public:
+		typedef std::iterator<std::forward_iterator_tag, const Residue>	baseType;
+		typedef typename baseType::pointer								pointer;
+		typedef typename baseType::reference							reference;
+		
+		typedef std::list<Polymer>::const_iterator						poly_iterator;
+		
+		residue_iterator(const Structure* s, poly_iterator polyIter, size_t polyResIndex, size_t nonPolyIndex);
+		
+		reference operator*();
+		pointer operator->();
+		
+		residue_iterator& operator++();
+		residue_iterator operator++(int);
+		
+		bool operator==(const residue_iterator& rhs) const;
+		bool operator!=(const residue_iterator& rhs) const;
+		
+	  private:
+		const Structure&	mStructure;
+		poly_iterator		mPolyIter;
+		size_t				mPolyResIndex;
+		size_t				mNonPolyIndex;
+	};
+	
+	class residue_view
+	{
+	  public:
+		residue_view(const Structure* s) : mStructure(s) {}
+		residue_view(const residue_view& rhs) : mStructure(rhs.mStructure) {}
+		residue_view& operator=(residue_view& rhs)
+		{
+			mStructure = rhs.mStructure;
+			return *this;
+		}
+		
+		residue_iterator begin() const		{ return residue_iterator(mStructure, mStructure->mPolymers.begin(), 0, 0); }
+		residue_iterator end() const		{ return residue_iterator(mStructure, mStructure->mPolymers.end(), 0, mStructure->mNonPolymers.size()); }
+		size_t size() const
+		{
+			size_t ps = std::accumulate(mStructure->mPolymers.begin(), mStructure->mPolymers.end(), 0UL, [](size_t s, auto& p) { return s + p.size(); });
+			return ps + mStructure->mNonPolymers.size();
+		}
+
+	  private:
+		const Structure* mStructure;
+	};
+	
+	residue_view residues() const			{ return residue_view(this); }
+	
   private:
 	friend Polymer;
 	friend Residue;
+	friend residue_view;
+	friend residue_iterator;
 
 	cif::Category& category(const char* name) const;
 	cif::Datablock& datablock() const;
 
 	void insertCompound(const std::string& compoundID, bool isEntity);
 	
+	void loadData();
 	void updateAtomIndex();
-
+	
 	File&					mFile;
 	uint32					mModelNr;
 	AtomView				mAtoms;
 	std::vector<size_t>		mAtomIndex;
 	std::list<Polymer>		mPolymers;
+	std::vector<Residue>	mNonPolymers;
 };
 
 }

src/DistanceMap.cpp

@@ -33,6 +33,9 @@ vector<clipper::RTop_orth> AlternativeSites(const clipper::Spacegroup& spacegrou
 {
 	vector<clipper::RTop_orth> result;
 	
+	// to make index 0 equal to identity, no 
+	result.push_back(clipper::RTop_orth::identity());
+	
 	for (int i = 0; i < spacegroup.num_symops(); ++i)
 	{
 		const auto& symop = spacegroup.symop(i);
@@ -53,11 +56,10 @@ vector<clipper::RTop_orth> AlternativeSites(const clipper::Spacegroup& spacegrou
 
 // --------------------------------------------------------------------
 
-DistanceMap::DistanceMap(const Structure& p, const clipper::Spacegroup& spacegroup, const clipper::Cell& cell)
-	: structure(p), dim(0)
+DistanceMap::DistanceMap(const Structure& p, const clipper::Spacegroup& spacegroup, const clipper::Cell& cell,
+		float maxDistance)
+	: structure(p), dim(0), mMaxDistance(maxDistance), mMaxDistanceSQ(maxDistance * maxDistance)
 {
-	const float kMaxDistance = 5, kMaxDistanceSQ = kMaxDistance * kMaxDistance;
-	
 	auto& atoms = p.atoms();
 	dim = atoms.size();
 	
@@ -138,72 +140,140 @@ DistanceMap::DistanceMap(const Structure& p, const clipper::Spacegroup& spacegro
 		for_each(locations.begin(), locations.end(), [&](auto& p) { p += mD; });
 	}
 	
-	pMin -= kMaxDistance;	// extend bounding box
-	pMax += kMaxDistance;
+	pMin -= mMaxDistance;	// extend bounding box
+	pMax += mMaxDistance;
 
 	mRtOrth = AlternativeSites(spacegroup, cell);
 	
-	cif::Progress progress(locations.size() - 1, "Creating distance map");
+	DistMap dist;
 	
-	boost::thread_group t;
-	size_t N = boost::thread::hardware_concurrency();
-	atomic<size_t> next(0);
-	mutex m;
+	vector<const Residue*> residues;
+	residues.reserve(p.residues().size());
+	for (auto& r: p.residues())
+	{
+		residues.emplace_back(&r);
+
+		// Add distances for atoms in this residue		
+		AddDistancesForAtoms(r, r, dist, 0);
+	}
 	
-	map<tuple<size_t,size_t>,tuple<float,size_t>> dist;
+	cif::Progress progress(residues.size(), "Creating distance map");
 	
-	for (size_t i = 0; i < N; ++i)
-		t.create_thread([&]()
+	for (size_t i = 0; i + 1 < residues.size(); ++i)
+	{
+		progress.consumed(1);
+		
+		auto& ri = *residues[i];
+		
+		Point centerI;
+		float radiusI;
+		tie(centerI, radiusI) = ri.centerAndRadius();
+		
+		for (size_t j = i + 1; j < residues.size(); ++j)
 		{
-			for (;;)
+			auto& rj = *residues[j];
+			
+			// first case, no symmetry operations
+			
+			Point centerJ;
+			float radiusJ;
+			tie(centerJ, radiusJ) = rj.centerAndRadius();
+			
+			auto d = Distance(centerI, centerJ) - radiusI - radiusJ;
+			if (d < mMaxDistance)
 			{
-				size_t i = next++;
-				
-				if (i >= locations.size())
-					break;
+//				cout << ri.labelID() << " en " << rj.labelID() << " liggen dicht bij elkaar: " << (d - radiusI - radiusJ) << endl;
+
+				AddDistancesForAtoms(ri, rj, dist, 0);
+				continue;
+			}
+			
+			// now try all symmetry operations to see if we can move rj close to ri
+			
+			clipper::Coord_orth cI = centerI;
+			clipper::Coord_orth cJ = centerJ;
+			
+			auto minR2 = d;
+			
+			int32 kbest = 0;
+			for (int32 k = 1; k < static_cast<int32>(mRtOrth.size()); ++k)
+			{
+				auto& rt = mRtOrth[k];
 				
-				clipper::Coord_orth pi = locations[i];
+				auto pJ = cJ.transform(rt);
+				double r2 = sqrt((cI - pJ).lengthsq()) - radiusI - radiusJ;
 
-				for (size_t j = 0; j < locations.size(); ++j)
+				if (minR2 > r2)
 				{
-					if (j == i)
-						continue;
-					
-					// find nearest location based on spacegroup/cell
-					double minR2 = numeric_limits<double>::max();
-					
-					size_t kb = 0;
-					for (size_t k = 0; k < mRtOrth.size(); ++k)
-					{
-						auto& rt = mRtOrth[k];
-						
-						auto pj = locations[j];
-						
-						pj = pj.transform(rt);
-						double r2 = (pi - pj).lengthsq();
-
-						if (minR2 > r2)
-						{
-							minR2 = r2;
-							kb = k;
-						}
-					}
-					
-					if (minR2 < kMaxDistanceSQ)
-					{
-						float d = sqrt(minR2);
-						auto key = make_tuple(i, j);
-						
-						lock_guard<mutex> lock(m);
-						dist[key] = make_tuple(d, kb);
-					}
+					minR2 = r2;
+					kbest = k;
 				}
-
-				progress.consumed(1);
 			}
-		});
-	
-	t.join_all();
+			
+			if (minR2 < mMaxDistance)
+			{
+//cout << ri.labelID() << " en " << rj.labelID() << " liggen dicht bij elkaar na symmetrie operatie: " << kbest << endl;
+				AddDistancesForAtoms(ri, rj, dist, kbest);
+			}
+		}
+	}
+
+//	cif::Progress progress(locations.size() - 1, "Creating distance map");
+//	
+//	boost::thread_group t;
+//	size_t N = boost::thread::hardware_concurrency();
+//	atomic<size_t> next(0);
+//	mutex m;
+//	
+//	for (size_t i = 0; i < N; ++i)
+//		t.create_thread([&]()
+//		{
+//			for (;;)
+//			{
+//				size_t i = next++;
+//				
+//				if (i >= locations.size())
+//					break;
+//				
+//				clipper::Coord_orth pi = locations[i];
+//
+//				for (size_t j = i + 1; j < locations.size(); ++j)
+//				{
+//					// find nearest location based on spacegroup/cell
+//					double minR2 = numeric_limits<double>::max();
+//					
+//					size_t kb = 0;
+//					for (size_t k = 0; k < mRtOrth.size(); ++k)
+//					{
+//						auto& rt = mRtOrth[k];
+//						
+//						auto pj = locations[j];
+//						
+//						pj = pj.transform(rt);
+//						double r2 = (pi - pj).lengthsq();
+//
+//						if (minR2 > r2)
+//						{
+//							minR2 = r2;
+//							kb = k;
+//						}
+//					}
+//					
+//					if (minR2 < mMaxDistanceSQ)
+//					{
+//						float d = sqrt(minR2);
+//						auto key = make_tuple(i, j);
+//						
+//						lock_guard<mutex> lock(m);
+//						dist[key] = make_tuple(d, kb);
+//					}
+//				}
+//
+//				progress.consumed(1);
+//			}
+//		});
+//	
+//	t.join_all();
 	
 	// Store as a sparse CSR compressed matrix
 	
@@ -306,52 +376,89 @@ DistanceMap::DistanceMap(const Structure& p, const clipper::Spacegroup& spacegro
 //		}
 }
 
-float DistanceMap::operator()(const Atom& a, const Atom& b) const
+// --------------------------------------------------------------------
+
+void DistanceMap::AddDistancesForAtoms(const Residue& a, const Residue& b, DistMap& dm, int32 rtix)
 {
-	size_t ixa, ixb;
-	
-	try
+	for (auto& aa: a.atoms())
 	{
-		ixa = index.at(a.id());
-	}
-	catch (const out_of_range& ex)
-	{
-		throw runtime_error("atom " + a.id() + " not found in distance map");
-	}
+		clipper::Coord_orth pa = aa.location();
+		size_t ixa = index[aa.id()];
 		
-	try
-	{
-		ixb = index.at(b.id());
-	}
-	catch (const out_of_range& ex)
-	{
-		throw runtime_error("atom " + b.id() + " not found in distance map");
-	}
-	
-//	if (ixb < ixa)
-//		std::swap(ixa, ixb);
-
-	size_t L = mIA[ixa];
-	size_t R = mIA[ixa + 1] - 1;
-	
-	while (L <= R)
-	{
-		size_t i = (L + R) / 2;
-
-		if (mJA[i] == ixb)
-			return get<0>(mA[i]);
-
-		if (mJA[i] < ixb)
-			L = i + 1;
-		else
-			R = i - 1;
+		for (auto& bb: b.atoms())
+		{
+			if (aa.id() == bb.id())
+				continue;
+			
+			clipper::Coord_orth pb = bb.location();
+			
+			if (rtix)
+				pb = pb.transform(mRtOrth[rtix]);
+			
+			auto d = (pa - pb).lengthsq();
+			if (d > mMaxDistanceSQ)
+				continue;
+
+			d = sqrt(d);
+
+			size_t ixb = index[bb.id()];
+
+			dm[make_tuple(ixa, ixb)] = make_tuple(d, rtix);
+			dm[make_tuple(ixb, ixa)] = make_tuple(d, -rtix);
+		}
 	}
-	
-	return 100.f;
 }
 
+//float DistanceMap::operator()(const Atom& a, const Atom& b) const
+//{
+//	size_t ixa, ixb;
+//	
+//	try
+//	{
+//		ixa = index.at(a.id());
+//	}
+//	catch (const out_of_range& ex)
+//	{
+//		throw runtime_error("atom " + a.id() + " not found in distance map");
+//	}
+//		
+//	try
+//	{
+//		ixb = index.at(b.id());
+//	}
+//	catch (const out_of_range& ex)
+//	{
+//		throw runtime_error("atom " + b.id() + " not found in distance map");
+//	}
+//	
+////	if (ixb < ixa)
+////		std::swap(ixa, ixb);
+//
+//	size_t L = mIA[ixa];
+//	size_t R = mIA[ixa + 1] - 1;
+//	
+//	while (L <= R)
+//	{
+//		size_t i = (L + R) / 2;
+//
+//		if (mJA[i] == ixb)
+//			return get<0>(mA[i]);
+//
+//		if (mJA[i] < ixb)
+//			L = i + 1;
+//		else
+//			R = i - 1;
+//	}
+//	
+//	return 100.f;
+//}
+
 vector<Atom> DistanceMap::near(const Atom& a, float maxDistance) const
 {
+	assert(maxDistance <= mMaxDistance);
+	if (maxDistance > mMaxDistance)
+		throw runtime_error("Invalid max distance in DistanceMap::near");
+	
 	size_t ixa;
 	try
 	{
@@ -364,17 +471,41 @@ vector<Atom> DistanceMap::near(const Atom& a, float maxDistance) const
 
 	vector<Atom> result;
 	
-	for (size_t bi = mIA[ixa]; bi < mIA[ixa + 1]; ++bi)
+	for (size_t i = mIA[ixa]; i < mIA[ixa + 1]; ++i)
 	{
 		float d;
-		size_t rti;
-		tie(d, rti) = mA[bi];
+		int32 rti;
+		tie(d, rti) = mA[i];
 
 		if (d > maxDistance)
 			continue;
 
-		Atom a = structure.getAtomById(rIndex.at(mJA[bi]));
-		result.emplace_back(a.symmetryCopy(mD, mRtOrth.at(rti)));
+		size_t ixb = mJA[i];
+		Atom b = structure.getAtomById(rIndex.at(ixb));
+		
+		if (rti > 0)
+			result.emplace_back(b.symmetryCopy(mD, mRtOrth.at(rti)));
+		else if (rti < 0)
+			result.emplace_back(b.symmetryCopy(mD, mRtOrth.at(-rti).inverse()));
+		else
+			result.emplace_back(b);
+
+#if 0 //DEBUG
+		if (rti != 0)
+			cerr << "symmetrie contact " << a.labelID() << " en " << result.back().labelID()
+				 << " d: " << d
+				 << " rti: " << rti
+				 << endl;
+		
+		auto d2 = Distance(a, result.back());
+		if (abs(d2 - d) > 0.01)
+			cerr << "Afstand " << a.labelID() << " en " << result.back().labelID()
+				 << " is niet gelijk aan verwachtte waarde:"
+				 << "d: " << d
+				 << " d2: " << d2
+				 << " rti: " << rti
+				 << endl;
+#endif
 	}
 	
 	return result;

src/Structure.cpp

@@ -589,9 +589,33 @@ float Atom::radius() const
 // --------------------------------------------------------------------
 // residue
 
+Residue::Residue(const Structure& structure, const std::string& compoundID,
+	const std::string& asymID, int seqID)
+	: mStructure(&structure), mCompoundID(compoundID)
+	, mAsymID(asymID), mSeqID(seqID)
+{
+	for (auto& a: mStructure->atoms())
+	{
+		if (mSeqID > 0 and a.labelSeqId() != mSeqID)
+			continue;
+		
+		if (a.labelAsymId() != mAsymID or
+			a.labelCompId() != mCompoundID)
+				continue;
+		
+		mAtoms.push_back(a);
+	}
+}
+
+Residue::Residue(const Structure& structure, const std::string& asymID, Atom water, int ndbSeqID)
+	: mStructure(&structure), mCompoundID("HOH")
+	, mAsymID(asymID), mNDBSeqID(ndbSeqID), mAtoms({water})
+{
+}
+
 Residue::Residue(Residue&& rhs)
 	: mStructure(rhs.mStructure), mCompoundID(move(rhs.mCompoundID)), mAsymID(move(rhs.mAsymID))
-	, mSeqID(rhs.mSeqID), mAtoms(move(rhs.mAtoms))
+	, mSeqID(rhs.mSeqID), mNDBSeqID(rhs.mNDBSeqID), mAtoms(move(rhs.mAtoms))
 {
 //cerr << "move constructor residue" << endl;
 	rhs.mStructure = nullptr;
@@ -604,6 +628,7 @@ Residue& Residue::operator=(Residue&& rhs)
 	mCompoundID = move(rhs.mCompoundID);
 	mAsymID = move(rhs.mAsymID);
 	mSeqID = rhs.mSeqID;
+	mNDBSeqID = rhs.mNDBSeqID;
 	mAtoms = move(rhs.mAtoms);
 
 	return *this;
@@ -663,65 +688,11 @@ const AtomView& Residue::atoms() const
 	if (mStructure == nullptr)
 		throw runtime_error("Invalid Residue object");
 
-	if (mAtoms.empty())
-	{
-		
-
-//#if 0
-		for (auto& a: mStructure->atoms())
-		{
-			if (mSeqID > 0 and a.labelSeqId() != mSeqID)
-				continue;
-			
-			if (a.labelAsymId() != mAsymID or
-				a.labelCompId() != mCompoundID)
-					continue;
-			
-//			if (not mAltID.empty() and a.property<string>("label_alt_id") != mAltID)
-//				continue;
-//			
-			const_cast<AtomView&>(mAtoms).push_back(a);
-		}
-
-if (not mAtoms.empty())
-	cerr << "loaded atoms for " << labelID() << endl;
-
-//#else
-//		auto& atomSites = mStructure->category("atom_site");
-//		
-//		auto query = cif::Key("label_asym_id") == mAsymID and cif::Key("label_comp_id") == mCompoundID;
-//		
-//		if (mSeqID != 0)
-//			query = move(query) and cif::Key("label_seq_id") == mSeqID;
-//		
-//		if (not mAltID.empty())
-//			query = move(query) and (cif::Key("label_alt_id") == cif::Empty() or cif::Key("label_alt_id") == mAltID);
-//	
-//		auto cifAtoms = atomSites.find(move(query));
-//
-//		for (auto cifAtom: cifAtoms)
-//			mAtoms.push_back(mStructure->getAtomById(cifAtom["id"].as<string>()));
-//#endif
-	}
-	
 	return mAtoms;
-	
-//	auto& atomSites = mStructure->category("atom_site");
-//	
-//	auto query = cif::Key("label_asym_id") == mAsymID and cif::Key("label_comp_id") == mCompoundID;
-//	
-//	if (not mSeqID.empty())
-//		query = move(query) and cif::Key("label_seq_id") == mSeqID;
-//	
-//	if (not mAltID.empty())
-//		query = move(query) and (cif::Key("label_alt_id") == cif::Empty() or cif::Key("label_alt_id") == mAltID);
-//
-//	return atomSites.find(move(query));
 }
 
 Atom Residue::atomByID(const string& atomID) const
 {
-//	for (auto& a: atoms())
 	for (auto& a: mAtoms)
 	{
 		if (a.labelAtomId() == atomID)
@@ -769,7 +740,29 @@ string Residue::authID() const
 
 string Residue::labelID() const
 {
-	return mAsymID + to_string(mSeqID);
+	if (mCompoundID == "HOH")
+		return mAsymID + to_string(mNDBSeqID);
+	else
+		return mAsymID + to_string(mSeqID);
+}
+
+tuple<Point,float> Residue::centerAndRadius() const
+{
+	vector<Point> pts;
+	for (auto& a: mAtoms)
+		pts.push_back(a.location());
+	
+	auto center = Centroid(pts);
+	float radius = 0;
+	
+	for (auto& pt: pts)
+	{
+		auto d = Distance(pt, center);
+		if (radius < d)
+			radius = d;
+	}
+	
+	return make_tuple(center, radius);
 }
 
 // --------------------------------------------------------------------
@@ -1080,16 +1073,6 @@ Polymer::Polymer(const Structure& s, const string& entityID, const string& asymI
 		ix[seqID] = index;
 		emplace_back(*this, index, seqID, compoundID);
 	}
-	
-	for (auto atom: s.atoms())
-	{
-		if (atom.labelAsymId() != mAsymID)
-			continue;
-		
-		uint32 index = ix.at(atom.labelSeqId());
-		
-		at(index).mAtoms.push_back(atom);
-	}
 }
 
 string Polymer::chainID() const
@@ -1259,22 +1242,7 @@ Structure::Structure(File& f, uint32 modelNr)
 			mAtoms.emplace_back(new AtomImpl(f, a["id"].as<string>(), a));
 	}
 	
-//	sort(mAtoms.begin(), mAtoms.end(), [](auto& a, auto& b) { return a.id() < b.id(); });
-
-	updateAtomIndex();
-
-	// polymers
-	auto& polySeqScheme = category("pdbx_poly_seq_scheme");
-	
-	for (auto& r: polySeqScheme)
-	{
-		string asymID, entityID, seqID, monID;
-		cif::tie(asymID, entityID, seqID, monID) =
-			r.get("asym_id", "entity_id", "seq_id", "mon_id");
-		
-		if (mPolymers.empty() or mPolymers.back().asymID() != asymID or mPolymers.back().entityID() != entityID)
-			mPolymers.emplace_back(*this, entityID, asymID);
-	}
+	loadData();
 }
 
 Structure::Structure(const Structure& s)
@@ -1283,7 +1251,16 @@ Structure::Structure(const Structure& s)
 	mAtoms.reserve(s.mAtoms.size());
 	for (auto& atom: s.mAtoms)
 		mAtoms.emplace_back(atom.clone());
+	
+	loadData();
+}
 
+Structure::~Structure()
+{
+}
+
+void Structure::loadData()
+{
 	updateAtomIndex();
 
 	auto& polySeqScheme = category("pdbx_poly_seq_scheme");
@@ -1297,10 +1274,40 @@ Structure::Structure(const Structure& s)
 		if (mPolymers.empty() or mPolymers.back().asymID() != asymID or mPolymers.back().entityID() != entityID)
 			mPolymers.emplace_back(*this, entityID, asymID);
 	}
-}
 
-Structure::~Structure()
-{
+	auto& nonPolyScheme = category("pdbx_nonpoly_scheme");
+	
+	vector<string> waterIDs;
+	for (auto& a: mAtoms)
+	{
+		if (a.isWater())
+			waterIDs.push_back(a.id());
+	}
+	auto waterIDIter = waterIDs.begin();
+	
+	for (auto& r: nonPolyScheme)
+	{
+		string asymID, monID;
+		int ndbSeqNum;
+		cif::tie(asymID, monID, ndbSeqNum) =
+			r.get("asym_id", "mon_id", "ndb_seq_num");
+		
+		if (monID == "HOH")
+		{
+			if (waterIDIter == waterIDs.end())
+				throw runtime_error("pdbx_nonpoly_scheme contains more waters than water atoms are available in atom_site");
+			
+			Atom water = getAtomById(*waterIDIter);
+			++waterIDIter;
+			
+			mNonPolymers.emplace_back(*this, asymID, water, ndbSeqNum);
+		}
+		else if (mNonPolymers.empty() or mNonPolymers.back().asymID() != asymID)
+			mNonPolymers.emplace_back(*this, monID, asymID);
+	}
+	
+	if (waterIDIter != waterIDs.end())
+		throw runtime_error("there are more waters in atom_site than referenced in pdbx_nonpoly_scheme");
 }
 
 void Structure::updateAtomIndex()
@@ -1312,11 +1319,6 @@ void Structure::updateAtomIndex()
 	sort(mAtomIndex.begin(), mAtomIndex.end(), [this](size_t a, size_t b) { return mAtoms[a].id() < mAtoms[b].id(); });
 }
 
-const AtomView& Structure::atoms() const
-{
-	return mAtoms;
-}
-
 AtomView Structure::waters() const
 {
 	AtomView result;
@@ -1346,30 +1348,6 @@ AtomView Structure::waters() const
 	return result;
 }
 
-const list<Polymer>& Structure::polymers() const
-{
-	return mPolymers;
-}
-
-vector<Residue> Structure::nonPolymers() const
-{
-	vector<Residue> result;
-
-	auto& nonPolyScheme = category("pdbx_nonpoly_scheme");
-	
-	for (auto& r: nonPolyScheme)
-	{
-		string asymID, monID;
-		cif::tie(asymID, monID) =
-			r.get("asym_id", "mon_id");
-		
-		if (result.empty() or result.back().asymID() != asymID)
-			result.emplace_back(*this, monID, asymID);
-	}
-	
-	return result;
-}
-
 Atom Structure::getAtomById(string id) const
 {
 	auto i = lower_bound(mAtomIndex.begin(), mAtomIndex.end(),
@@ -1576,6 +1554,67 @@ void Structure::insertCompound(const string& compoundID, bool isEntity)
 	}
 }
 
+// --------------------------------------------------------------------
+
+Structure::residue_iterator::residue_iterator(const Structure* s, poly_iterator polyIter, size_t polyResIndex, size_t nonPolyIndex)
+	: mStructure(*s), mPolyIter(polyIter), mPolyResIndex(polyResIndex), mNonPolyIndex(nonPolyIndex)
+{
+	while (mPolyIter != mStructure.mPolymers.end() and mPolyResIndex == mPolyIter->size())
+		++mPolyIter;
+}
+
+auto Structure::residue_iterator::operator*() -> reference
+{
+	if (mPolyIter != mStructure.mPolymers.end())
+		return (*mPolyIter)[mPolyResIndex];
+	else
+		return mStructure.mNonPolymers[mNonPolyIndex];
+}
+
+auto Structure::residue_iterator::operator->() -> pointer
+{
+	if (mPolyIter != mStructure.mPolymers.end())
+		return &(*mPolyIter)[mPolyResIndex];
+	else
+		return &mStructure.mNonPolymers[mNonPolyIndex];
+}
+		
+Structure::residue_iterator& Structure::residue_iterator::operator++()
+{
+	if (mPolyIter != mStructure.mPolymers.end())
+	{
+		++mPolyResIndex;
+		if (mPolyResIndex >= mPolyIter->size())
+		{
+			++mPolyIter;
+			mPolyResIndex = 0;
+		}
+	}
+	else
+		++mNonPolyIndex;
+	
+	return *this;
+}
+
+Structure::residue_iterator Structure::residue_iterator::operator++(int)
+{
+	auto result = *this;
+	operator++();
+	return result;
+}
+		
+bool Structure::residue_iterator::operator==(const Structure::residue_iterator& rhs) const
+{
+	return mPolyIter == rhs.mPolyIter and mPolyResIndex == rhs.mPolyResIndex and mNonPolyIndex == rhs.mNonPolyIndex;
+}
+
+bool Structure::residue_iterator::operator!=(const Structure::residue_iterator& rhs) const
+{
+	return mPolyIter != rhs.mPolyIter or mPolyResIndex != rhs.mPolyResIndex or mNonPolyIndex != rhs.mNonPolyIndex;
+}
+
+// --------------------------------------------------------------------
+
 void Structure::removeAtom(Atom& a)
 {
 	cif::Datablock& db = *mFile.impl().mDb;