Implement Complexes

atomium/structures/__init__.py

@@ -1,4 +1,4 @@
-from .models import Model
+from .models import Model, Complex
 from .chains import Chain
 from .molecules import Residue, Molecule
 from .atoms import Atom

atomium/structures/atoms.py

@@ -79,7 +79,7 @@ def __init__(self, element, x=0, y=0, z=0, id=0, name=None, charge=0,
         self._bfactor = bfactor
         self._bonds = set()
         self._residue, self._chain, self._molecule = None, None, None
-        self._model = None
+        self._model, self._complex = None, None
 
 
     def __repr__(self):
@@ -390,6 +390,16 @@ def model(self):
         return self._model
 
 
+    @property
+    def complex(self):
+        """Returns the :py:class:`.Complex` the atom is part of, or ``None`` if
+        it is not part of one.
+
+        :rtype: ``Complex``"""
+
+        return self._complex
+
+
     @property
     def bonds(self):
         """The atomic :py:class:`.Bond` objects that the atom is associated

atomium/structures/models.py

@@ -5,9 +5,7 @@
 from .chains import Chain
 
 class Model(AtomicStructure):
-    """Base classes: :py:class:`.AtomicStructure` and
-    :py:class:`.ResidueStructure` and :py:class:`.ChainStructure` and
-    :py:class:`.MoleculeStructure`
+    """Base class: :py:class:`.AtomicStructure`
 
     Represents molecular systems. These are essentially the isolated universes
     in which the other structures live.
@@ -20,3 +18,50 @@ def __init__(self, *atoms):
         AtomicStructure.__init__(self, *atoms)
         for atom in self._atoms:
             atom._model = self
+
+
+
+class Complex(AtomicStructure):
+    """Base class: :py:class:`.AtomicStructure`.
+
+    Clusters of chains which form a single functional unit.
+
+    :param \*atoms: The atoms that make up the model. These can also be\
+    :py:class:`.AtomicStructure` objects, in which case the atoms of that\
+    structure will be used in its place."""
+
+    def __init__(self, *atoms, id=None, name=None):
+        AtomicStructure.__init__(self, *atoms)
+        if id is not None and not isinstance(id, str):
+            raise TypeError("ID {} is not a string".format(id))
+        if name is not None and not isinstance(name, str):
+            raise TypeError("Complex name {} is not a string".format(name))
+        self._id = id
+        self._name = name
+        for atom in self._atoms:
+            atom._complex = self
+
+
+    @property
+    def id(self):
+        """The complex's unique string ID.
+
+        :rtype: ``str``"""
+
+        return self._id
+
+
+    @property
+    def name(self):
+        """The complex's name.
+
+        :raises TypeError: if the name given is not str."""
+
+        return self._name
+
+
+    @name.setter
+    def name(self, name):
+        if not isinstance(name, str):
+            raise TypeError("Complex name '{}' is not str".format(name))
+        self._name = name

atomium/structures/molecules.py

@@ -266,6 +266,39 @@ def chain(self, *args, **kwargs):
         for chain in chains: return chain
 
 
+    def complexes(self, id=None, name=None):
+        """Returns all the :py:class:`.Complex` objects in the structure which
+        match the given criteria.
+
+        :param str id: Filter by complex ID.
+        :param str name: Filter by name.
+        :rtype: ``Complex``"""
+
+        complexes = set()
+        for atom in self._atoms:
+            complexes.add(atom.complex)
+        try:
+            complexes.remove(None)
+        except KeyError: pass
+        if id:
+            complexes = set(filter(lambda r: r.id == id, complexes))
+        if name:
+            complexes = set(filter(lambda r: r.name == name, complexes))
+        return complexes
+
+
+    def complex(self, *args, **kwargs):
+        """Returns the first :py:class:`.Complex` object in the structure which
+        matches the given criteria.
+
+        :param str id: Filter by complex ID.
+        :param str name: Filter by name.
+        :rtype: ``Complex``"""
+
+        complexes = self.complexes(*args, **kwargs)
+        for complex in complexes: return complex
+
+
     def trim(self, places):
         """Rounds the coordinate values to a given number of decimal places.
         Useful for removing floating point rounding errors after transformation.

tests/integration/test_structures.py

@@ -1,6 +1,6 @@
 import math
 from tests.integration.base import IntegratedTest
-from atomium.structures import Model, Atom, Residue, Chain, Molecule
+from atomium.structures import Model, Atom, Residue, Chain, Molecule, Complex
 import atomium
 
 class CreationTests(IntegratedTest):
@@ -322,6 +322,12 @@ def test_atoms_in_chains(self):
         self.assertIs(self.atoms[-1].chain, chainb)
         self.assertIs(res6.chain, chainb)
 
+        # Complexes
+        complex = Complex(chaina, chainb, id="1", name="HEAVY")
+        self.assertEqual(complex.chains(), {chaina, chainb})
+        for atom in self.atoms[:9] + self.atoms[18:]:
+            self.assertIs(atom.complex, complex)
+
 
     def test_atoms_in_models(self):
         """Full model processing"""
@@ -339,7 +345,8 @@ def test_atoms_in_models(self):
         mol1 = Molecule(*self.atoms[18:21], id="A1000", name="XMP")
         mol2 = Molecule(*self.atoms[21:24], id="A1001", name="BIS")
         mol3 = Molecule(*self.atoms[24:27], id="A1002", name="BIS")
-        model = Model(chaina, chainb, mol1, mol2, mol3)
+        complex = Complex(chaina, chainb, id="1", name="HEAVY")
+        model = Model(complex, mol1, mol2, mol3)
 
         # Atoms in model
         for atom in self.atoms:
@@ -383,6 +390,10 @@ def test_atoms_in_models(self):
         self.assertEqual(model.chains(), {chaina, chainb})
         self.assertIs(model.chain("B"), chainb)
 
+        # Complexes in model
+        self.assertEqual(model.complexes(name="HEAVY"), {complex})
+        self.assertEqual(model.complex("1"), complex)
+
         # Model grid
         self.assertEqual(list(model.grid()), [(x, y, z)
          for x in range(3) for y in range(3) for z in range(3)])

tests/unit/structure_tests/test_atomic_structures.py

@@ -468,6 +468,77 @@ def test_chain_can_return_none(self, mock_chains):
 
 
 
+class StructureComplexesTests(AtomicStructureTest):
+
+    def setUp(self):
+        AtomicStructureTest.setUp(self)
+        self.structure = AtomicStructure(self.atom1, self.atom2, self.atom3)
+        self.complex1, self.complex2, self.complex3 = Mock(), Mock(), Mock()
+        self.atom1.complex, self.atom2.complex, self.atom3.complex = (
+         self.complex1, self.complex2, self.complex3
+        )
+        self.complex1.id, self.complex2.id, self.complex3.id = "1", "2", "3"
+        self.complex1.name, self.complex2.name, self.complex3.name = "AA", "CC", "CC"
+
+
+    def test_can_get_complexes(self):
+        self.assertEqual(
+         self.structure.complexes(),
+         set([self.complex1, self.complex2, self.complex3])
+        )
+
+
+    def test_can_filter_none_from_complexes(self):
+        self.atom3.complex = None
+        self.assertEqual(self.structure.complexes(), {self.complex1, self.complex2})
+
+
+    def test_can_get_complexes_by_id(self):
+        self.assertEqual(
+         self.structure.complexes(id="1"), {self.complex1}
+        )
+        self.assertEqual(
+         self.structure.complexes(id="2"), {self.complex2}
+        )
+        self.assertEqual(self.structure.complexes(id="4"), set())
+
+
+    def test_can_get_complexes_by_name(self):
+        self.assertEqual(
+         self.structure.complexes(name="AA"), set([self.complex1])
+        )
+        self.assertEqual(
+         self.structure.complexes(name="CC"), set([self.complex2, self.complex3])
+        )
+        self.assertEqual(self.structure.complexes(name="DD"), set())
+
+
+
+class StructureComplexTest(AtomicStructureTest):
+
+    def setUp(self):
+        AtomicStructureTest.setUp(self)
+        self.structure = AtomicStructure(self.atom1, self.atom2, self.atom3)
+        self.complex1, self.complex2, self.complex3 = Mock(), Mock(), Mock()
+        self.complex1.id, self.complex2.id, self.complex3.id = "1", "2", "3"
+        self.complex1.name, self.complex2.name, self.complex3.name = "AA", "BB", "CC"
+
+
+    @patch("atomium.structures.molecules.AtomicStructure.complexes")
+    def test_complex_calls_complexes(self, mock_complexes):
+        mock_complexes.return_value = set([self.complex3])
+        complex = self.structure.complex(name="1")
+        mock_complexes.assert_called_with(name="1")
+        self.assertIs(complex, self.complex3)
+
+
+    @patch("atomium.structures.molecules.AtomicStructure.complexes")
+    def test_complex_can_return_none(self, mock_complexes):
+        mock_complexes.return_value = set()
+        self.assertIs(self.structure.complex(name="AA"), None)
+
+
+
 class AtomicStructureTrimmingTests(AtomicStructureTest):
 
     @patch("atomium.structures.molecules.AtomicStructure.atoms")

tests/unit/structure_tests/test_atoms.py

@@ -20,6 +20,7 @@ def test_can_create_atom(self):
         self.assertEqual(atom._chain, None)
         self.assertEqual(atom._molecule, None)
         self.assertEqual(atom._model, None)
+        self.assertEqual(atom._complex, None)
 
 
     def test_atom_element_must_be_str(self):
@@ -540,6 +541,16 @@ def test_model_property(self):
         self.assertIs(atom.model, model)
 
 
+
+class AtomComplexTests(TestCase):
+
+    def test_complex_property(self):
+        complex = Mock()
+        atom = Atom("C", 2, 3, 5)
+        atom._complex = complex
+        self.assertIs(atom.complex, complex)
+
+
 class AtomBondsTests(TestCase):
 
     def test_bonds_property(self):

tests/unit/structure_tests/test_complexes.py

@@ -0,0 +1,92 @@
+from unittest import TestCase
+from unittest.mock import patch, Mock
+from atomium.structures.models import Complex
+from atomium.structures.molecules import AtomicStructure
+from atomium.structures.atoms import Atom
+
+class ComplexTest(TestCase):
+
+    def setUp(self):
+        self.atom1, self.atom2, self.atom3 = Mock(Atom), Mock(Atom), Mock(Atom)
+        self.atoms = [self.atom1, self.atom2, self.atom3]
+        def mock_init(obj, *args, **kwargs):
+            obj._atoms = set(args)
+        self.patch1 = patch("atomium.structures.molecules.AtomicStructure.__init__")
+        self.mock_init = self.patch1.start()
+        self.mock_init.side_effect = mock_init
+
+
+
+class ComplexCreationTests(ComplexTest):
+
+    @patch("atomium.structures.molecules.AtomicStructure.__init__")
+    def test_model_is_atomic_structure(self, mock_init):
+        mock_init.side_effect = self.mock_init
+        cmplx = Complex(*self.atoms)
+        self.assertIsInstance(cmplx, AtomicStructure)
+        self.assertTrue(mock_init.called)
+        self.assertIsNone(cmplx._id)
+        self.assertIsNone(cmplx._name)
+
+
+    def test_can_create_complex_with_id(self):
+        complx = Complex(self.atom1, self.atom2, self.atom3, id="1")
+        self.assertEqual(complx._id, "1")
+
+
+    def test_complex_id_must_be_str(self):
+        with self.assertRaises(TypeError):
+            Complex(self.atom1, self.atom2, self.atom3, id=1000)
+
+
+    def test_can_create_complex_with_name(self):
+        complx = Complex(self.atom1, self.atom2, self.atom3, name="HEAVY")
+        self.assertEqual(complx._name, "HEAVY")
+
+
+    def test_complex_name_must_be_str(self):
+        with self.assertRaises(TypeError):
+            Complex(self.atom1, self.atom2, self.atom3, name=1000)
+
+
+    def test_atoms_are_linked_to_complex(self):
+        cmplx = Complex(*self.atoms)
+        self.assertIs(self.atom1._complex, cmplx)
+        self.assertIs(self.atom2._complex, cmplx)
+        self.assertIs(self.atom3._complex, cmplx)
+
+
+
+class ComplexReprTests(ComplexTest):
+
+    def test_complex_repr(self):
+        cmplx = Complex(*self.atoms)
+        self.assertEqual(str(cmplx), "<Complex (3 atoms)>")
+
+
+
+class ComplexIdTests(ComplexTest):
+
+    def test_complex_id_property(self):
+        complex = Complex(self.atom1, self.atom2, self.atom3, id="B10C")
+        self.assertIs(complex._id, complex.id)
+
+
+
+class ComplexNameTests(ComplexTest):
+
+    def test_complex_name_property(self):
+        complex = Complex(self.atom1, self.atom2, self.atom3, name="VAL")
+        self.assertIs(complex._name, complex.name)
+
+
+    def test_can_update_complex_name(self):
+        complex = Complex(self.atom1, self.atom2, self.atom3, name="VAL")
+        complex.name = "HIS"
+        self.assertEqual(complex._name, "HIS")
+
+
+    def test_complex_name_must_be_str(self):
+        complex = Complex(self.atom1, self.atom2, self.atom3, name="VAL")
+        with self.assertRaises(TypeError):
+            complex.name = 10