From ed58400c4b651e8b4bf22547ff0fbe41cd80a3d4 Mon Sep 17 00:00:00 2001 From: James Krieger Date: Mon, 18 Dec 2023 19:00:17 +0000 Subject: [PATCH 1/5] various fixes for handling ensembles with selections as ref --- prody/ensemble/ensemble.py | 41 ++++++++++++++++++++++++++++++----- prody/ensemble/pdbensemble.py | 5 +++++ 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/prody/ensemble/ensemble.py b/prody/ensemble/ensemble.py index 2e0c20361..d4337f5e0 100644 --- a/prody/ensemble/ensemble.py +++ b/prody/ensemble/ensemble.py @@ -238,7 +238,7 @@ def setAtoms(self, atoms): n_atoms = self._n_atoms if n_atoms: - if atoms.numAtoms() > n_atoms: + if atoms.numAtoms() > n_atoms and atoms.ca.numAtoms() > n_atoms: raise ValueError('atoms must be same size or smaller than ' 'the ensemble') @@ -261,7 +261,15 @@ def setAtoms(self, atoms): ag = atoms.getAtomGroup() except AttributeError: ag = atoms - if ag.numAtoms() != n_atoms: + try: + self_ag = self._atoms.getAtomGroup() + except AttributeError: + self_ag = self._atoms + try: + self_ag_n_atoms = self_ag.numAtoms() + except AttributeError: + self_ag_n_atoms = 0 + if ag.numAtoms() != n_atoms and ag.numAtoms() != self_ag_n_atoms and self_ag_n_atoms != 0: raise ValueError('size mismatch between this ensemble ({0} atoms) and atoms ({1} atoms)' .format(n_atoms, ag.numAtoms())) self._atoms = ag @@ -294,12 +302,26 @@ def getCoords(self, selected=True): return None if self._indices is None or not selected: return self._coords.copy() - return self._coords[self._indices].copy() + + selids = self._indices + if self.hasSelectionIssue(): + selids = self.getIndices(calphas=True) + return self._coords[selids].copy() - def getIndices(self): + def getIndices(self, calphas=False): """Returns a copy of indices of selected columns""" + if calphas: + return array([list(self._atoms.ca.getIndices()).index(idx) + for idx in self._indices]) return copy(self._indices) + + def hasSelectionIssue(self): + selids = self._indices + if (selids.max() > self._coords.shape[0] + and set(selids).issubset(set(self._atoms.ca.getIndices()))): + return True + return False def setIndices(self, value): if not isListLike(value): @@ -359,10 +381,17 @@ def getWeights(self, selected=True): return None if self._indices is None or not selected: return self._weights.copy() + + + if self.hasSelectionIssue(): + selids = self.getIndices(calphas=True) + else: + selids = self._indices + if self._weights.ndim == 2: - return self._weights[self._indices].copy() + return self._weights[selids].copy() else: - return self._weights[:, self._indices].copy() + return self._weights[:, selids].copy() def _getWeights(self, selected=True): diff --git a/prody/ensemble/pdbensemble.py b/prody/ensemble/pdbensemble.py index 12edcebf2..49540cd00 100644 --- a/prody/ensemble/pdbensemble.py +++ b/prody/ensemble/pdbensemble.py @@ -459,6 +459,9 @@ def getMSA(self, indices=None, selected=True): atom_indices = self._indices if selected else slice(None, None, None) indices = indices if indices is not None else slice(None, None, None) + + if self.hasSelectionIssue(): + atom_indices = self.getIndices(calphas=True) return self._msa[indices, atom_indices] @@ -491,6 +494,8 @@ def getCoordsets(self, indices=None, selected=True): confs[i, which] = coords[which] else: selids = self._indices + if self.hasSelectionIssue(): + selids = self.getIndices(calphas=True) coords = coords[selids] confs = self._confs[indices, selids].copy() for i, w in enumerate(self._weights[indices]): From 89ba879822a3b7a7a1474859b978e36e9f87bb6c Mon Sep 17 00:00:00 2001 From: James Krieger Date: Mon, 18 Dec 2023 19:10:43 +0000 Subject: [PATCH 2/5] sliceAtoms allow same --- prody/atomic/functions.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/prody/atomic/functions.py b/prody/atomic/functions.py index c2cab1b6c..433250712 100644 --- a/prody/atomic/functions.py +++ b/prody/atomic/functions.py @@ -285,7 +285,7 @@ def sortAtoms(atoms, label, reverse=False): return AtomMap(ag, sort, acsi) -def sliceAtoms(atoms, select): +def sliceAtoms(atoms, select, allowSame=False): """Slice *atoms* using the selection defined by *select*. :arg atoms: atoms to be selected from @@ -297,6 +297,8 @@ def sliceAtoms(atoms, select): """ if atoms == select: + if allowSame: + return atoms._getSubset('all'), atoms.all raise ValueError('atoms and select arguments are the same') try: From 11f2ee7d5ce6d6835ddd06feabdf958f2562943d Mon Sep 17 00:00:00 2001 From: James Krieger Date: Mon, 18 Dec 2023 19:11:24 +0000 Subject: [PATCH 3/5] use sliceAtoms allow same to fix loadEnsemble --- prody/ensemble/ensemble.py | 4 ++-- prody/ensemble/functions.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/prody/ensemble/ensemble.py b/prody/ensemble/ensemble.py index d4337f5e0..10985d7ac 100644 --- a/prody/ensemble/ensemble.py +++ b/prody/ensemble/ensemble.py @@ -217,7 +217,7 @@ def getAtoms(self, selected=True): return self._atoms return self._atoms[self._indices] - def setAtoms(self, atoms): + def setAtoms(self, atoms, allowSame=False): """Set *atoms* or specify a selection of atoms to be considered in calculations and coordinate requests. When a selection is set, corresponding subset of coordinates will be considered in, for @@ -273,7 +273,7 @@ def setAtoms(self, atoms): raise ValueError('size mismatch between this ensemble ({0} atoms) and atoms ({1} atoms)' .format(n_atoms, ag.numAtoms())) self._atoms = ag - self._indices, _ = sliceAtoms(self._atoms, atoms) + self._indices, _ = sliceAtoms(self._atoms, atoms, allowSame=allowSame) else: # if assigning atoms to a new ensemble self._n_atoms = atoms.numAtoms() diff --git a/prody/ensemble/functions.py b/prody/ensemble/functions.py index ac0b992c4..fc48c743f 100644 --- a/prody/ensemble/functions.py +++ b/prody/ensemble/functions.py @@ -184,7 +184,7 @@ def loadEnsemble(filename, **kwargs): data[key] = arr else: atoms = None - ensemble.setAtoms(atoms) + ensemble.setAtoms(atoms, allowSame=True) if '_indices' in attr_dict: indices = attr_dict['_indices'] From 3f16d3b409cdab7d669369e809c2fa016d327f7c Mon Sep 17 00:00:00 2001 From: James Krieger Date: Mon, 18 Dec 2023 20:04:11 +0000 Subject: [PATCH 4/5] fix hasSelectionIssue for no atoms --- prody/ensemble/ensemble.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/prody/ensemble/ensemble.py b/prody/ensemble/ensemble.py index 10985d7ac..bf3f3df96 100644 --- a/prody/ensemble/ensemble.py +++ b/prody/ensemble/ensemble.py @@ -317,6 +317,9 @@ def getIndices(self, calphas=False): return copy(self._indices) def hasSelectionIssue(self): + if self._atoms is None or self._atoms.ca is None: + return False + selids = self._indices if (selids.max() > self._coords.shape[0] and set(selids).issubset(set(self._atoms.ca.getIndices()))): From ef29c9725ce4a39ca93c01f81cdc3ab17269ec31 Mon Sep 17 00:00:00 2001 From: James Krieger Date: Mon, 18 Dec 2023 20:04:57 +0000 Subject: [PATCH 5/5] fix hasSelectionIssue for no indices --- prody/ensemble/ensemble.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/prody/ensemble/ensemble.py b/prody/ensemble/ensemble.py index bf3f3df96..9556ee1ba 100644 --- a/prody/ensemble/ensemble.py +++ b/prody/ensemble/ensemble.py @@ -321,6 +321,9 @@ def hasSelectionIssue(self): return False selids = self._indices + if selids is None: + return False + if (selids.max() > self._coords.shape[0] and set(selids).issubset(set(self._atoms.ca.getIndices()))): return True