From 7ee0ed1c36107a56afc5d7750ce537ad7e022b8a Mon Sep 17 00:00:00 2001
From: James Krieger <jamesmkrieger@gmail.com>
Date: Fri, 3 Nov 2023 17:08:27 +0100
Subject: [PATCH] propagate anisous to other objects

---
 prody/atomic/atom.py      | 21 +++++++++++++++++++++
 prody/atomic/atomgroup.py |  4 ++--
 prody/atomic/fields.py    |  2 --
 prody/atomic/select.py    |  7 +++++++
 prody/atomic/subset.py    | 16 ++++++++++++++++
 5 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/prody/atomic/atom.py b/prody/atomic/atom.py
index 6d3b63d10..b4ea110e4 100644
--- a/prody/atomic/atom.py
+++ b/prody/atomic/atom.py
@@ -94,6 +94,27 @@ def setCoords(self, coords):
         self._ag._coords[acsi, self._index] = coords
         self._ag._setTimeStamp(acsi)
 
+    def getAnisou(self):
+        """Returns a copy of anisou of the atom from the active coordinate
+        set."""
+
+        if self._ag._anisous is not None:
+            return self._ag._anisous[self.getACSIndex(), self._index].copy()
+
+    def _getAnisou(self):
+        """Returns a view of anisou of the atom from the active coordinate
+        set."""
+
+        if self._ag._anisous is not None:
+            return self._ag._anisous[self.getACSIndex(), self._index]
+
+    def setAnisou(self, anisou):
+        """Set anisou of the atom in the active coordinate set."""
+
+        acsi = self.getACSIndex()
+        self._ag._anisous[acsi, self._index] = anisou
+        self._ag._setTimeStamp(acsi)
+
     def getCoordsets(self, indices=None):
         """Returns a copy of coordinate set(s) at given *indices*."""
 
diff --git a/prody/atomic/atomgroup.py b/prody/atomic/atomgroup.py
index 3a9b18f0a..0d6357d83 100644
--- a/prody/atomic/atomgroup.py
+++ b/prody/atomic/atomgroup.py
@@ -615,8 +615,8 @@ def setAnisous(self, anisous, label=''):
         try:
             checkAnisous(anisous, csets=True, dtype=(float, np.float32))
         except TypeError:
-            raise TypeError('coords must be a numpy array or an '
-                            'object with `getCoords` method')
+            raise TypeError('anisous must be a numpy array or an '
+                            'object with `getAnisous` method')
 
         self._setAnisous(anisous, label=label)
 
diff --git a/prody/atomic/fields.py b/prody/atomic/fields.py
index 2dd3288c7..26b7fa106 100644
--- a/prody/atomic/fields.py
+++ b/prody/atomic/fields.py
@@ -123,8 +123,6 @@ def getDocstr(self, meth, plural=True, selex=True):
     'altloc':    Field('altloc', DTYPE + '1',
                        doc='alternate location indicator',
                        selstr=('altloc A B', 'altloc _'),),
-    'anisou':    Field('anisou', float, doc='anisotropic temperature factor',
-                       ndim=2),
     'chain':     Field('chain', DTYPE + '6',  doc='chain identifier',
                        meth='Chid', none=HVNONE, synonym='chid',
                        selstr=('chain A', 'chid A B C', 'chain _')),
diff --git a/prody/atomic/select.py b/prody/atomic/select.py
index a58ad749e..1749ef2c3 100644
--- a/prody/atomic/select.py
+++ b/prody/atomic/select.py
@@ -2424,3 +2424,10 @@ def _getCoords(self):
         if self._coords is None:
             self._coords = self._atoms._getCoords()
         return self._coords
+
+    def _getAnisous(self):
+        """Returns anisous of atoms."""
+
+        if self._anisous is None:
+            self._anisous = self._atoms._getAnisous()
+        return self._anisous
diff --git a/prody/atomic/subset.py b/prody/atomic/subset.py
index b6fc466df..9099b1789 100644
--- a/prody/atomic/subset.py
+++ b/prody/atomic/subset.py
@@ -63,6 +63,15 @@ def getCoords(self):
 
     _getCoords = getCoords
 
+    def getAnisous(self):
+        """Returns a copy of anisous from the active coordinate set."""
+
+        if self._ag._anisous is not None:
+            # Since this is not slicing, a view is not returned
+            return self._ag._anisous[self.getACSIndex(), self._indices]
+
+    _getAnisous = getAnisous
+
     def setCoords(self, coords):
         """Set coordinates in the active coordinate set."""
 
@@ -70,6 +79,13 @@ def setCoords(self, coords):
             self._ag._coords[self.getACSIndex(), self._indices] = coords
             self._ag._setTimeStamp(self.getACSIndex())
 
+    def setAnisous(self, anisous):
+        """Set anisous in the active coordinate set."""
+
+        if self._ag._anisous is not None:
+            self._ag._anisous[self.getACSIndex(), self._indices] = anisous
+            self._ag._setTimeStamp(self.getACSIndex())
+
     def getCoordsets(self, indices=None):
         """Returns coordinate set(s) at given *indices*, which may be an integer
         or a list/array of integers."""