1# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
2# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
3#
4# MDAnalysis --- https://www.mdanalysis.org
5# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
6# (see the file AUTHORS for the full list of names)
7#
8# Released under the GNU Public Licence, v2 or any higher version
9#
10# Please cite your use of MDAnalysis in published work:
11#
12# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
13# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
14# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
15# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
16# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
17#
18# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
19# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
20# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
21#
22from __future__ import absolute_import
23
24import pytest
25from six.moves import range
26
27import os
28from numpy.testing import (
29    assert_equal, assert_array_equal,
30    assert_array_almost_equal, TestCase,
31)
32
33from MDAnalysisTests.datafiles import (
34    mol2_molecules, mol2_molecule, mol2_broken_molecule,
35    mol2_zinc,
36)
37from MDAnalysis import Universe
38import MDAnalysis as mda
39from MDAnalysisTests import tempdir, make_Universe
40
41
42class TestMol2(TestCase):
43    def setUp(self):
44        self.tempdir = tempdir.TempDir()
45        self.outfile = os.path.join(self.tempdir.name, 'test.mol2')
46
47    def tearDown(self):
48        del self.tempdir
49
50    def test_read(self):
51        u = Universe(mol2_molecules)
52        assert_equal(len(u.atoms), 49)
53        assert_equal(u.trajectory.n_frames, 200)
54
55        u.trajectory[199]
56        assert_array_almost_equal(u.atoms.positions[0], [1.7240, 11.2730, 14.1200])
57
58    def test_write(self):
59        ref = Universe(mol2_molecules)
60        ref.atoms.write(self.outfile)
61        u = Universe(self.outfile)
62        assert_equal(len(u.atoms), len(ref.atoms))
63        assert_equal(len(u.trajectory), 1)
64        assert_array_equal(u.atoms.positions, ref.atoms.positions)
65
66    def test_write_selection(self):
67        ref = Universe(mol2_molecule)
68        gr0 = ref.select_atoms("name C*")
69        gr0.write(self.outfile)
70        u = Universe(self.outfile)
71        gr1 = u.select_atoms("name C*")
72        assert_equal(len(gr0), len(gr1))
73
74    def test_write_in_loop(self):
75        ref = Universe(mol2_molecules)
76
77        with mda.Writer(self.outfile) as W:
78            for ts in ref.trajectory:
79                W.write(ref.atoms)
80        u = Universe(self.outfile)
81        assert_equal(len(u.atoms), len(ref.atoms))
82        assert_equal(len(u.trajectory), len(ref.trajectory))
83        assert_array_equal(u.atoms.positions, ref.atoms.positions)
84        u.trajectory[199]
85        ref.trajectory[199]
86        assert_array_equal(u.atoms.positions, ref.atoms.positions)
87
88    def test_broken_molecule(self):
89        with pytest.raises(ValueError):
90            Universe(mol2_broken_molecule)
91
92        # This doesn't work with 2.6
93        # Checks the text of the error message, so it low priority
94        #with self.assertRaises(Exception) as context:
95        #    u = Universe(mol2_broken_molecule)
96        #self.assertEqual("The mol2 block (BrokenMolecule.mol2:0) has no atoms" in context.exception.message,
97        # True)
98
99
100class TestMol2_traj(TestCase):
101    def setUp(self):
102        self.universe = Universe(mol2_molecules)
103        self.traj = self.universe.trajectory
104        self.ts = self.universe.coord
105
106    def tearDown(self):
107        del self.universe
108        del self.traj
109        del self.ts
110
111    def test_rewind_traj(self):
112        self.traj.rewind()
113        assert_equal(self.ts.frame, 0, "rewinding to frame 0")
114
115    def test_next_traj(self):
116        self.traj.rewind()
117        self.traj.next()
118        assert_equal(self.ts.frame, 1, "loading frame 1")
119
120    def test_jump_traj(self):
121        self.traj[15]  # index is 0-based and frames are 0-based
122        assert_equal(self.ts.frame, 15, "jumping to frame 15")
123
124    def test_jump_lastframe_traj(self):
125        self.traj[-1]
126        assert_equal(self.ts.frame, 199, "indexing last frame with traj[-1]")
127
128    def test_slice_traj(self):
129        frames = [ts.frame for ts in self.traj[5:17:3]]
130        assert_equal(frames, [5, 8, 11, 14], "slicing traj [5:17:3]")
131
132    def test_reverse_traj(self):
133        frames = [ts.frame for ts in self.traj[20:5:-1]]
134        assert_equal(frames, list(range(20, 5, -1)),
135                     "reversing traj [20:5:-1]")
136
137    def test_n_frames(self):
138        assert_equal(self.universe.trajectory.n_frames, 200, "wrong number of frames in traj")
139
140
141class TestMOL2NoSubstructure(object):
142    """MOL2 file without substructure
143
144    """
145    n_atoms = 45
146
147    def test_load(self):
148        r = mda.coordinates.MOL2.MOL2Reader(mol2_zinc, n_atoms=self.n_atoms)
149        assert r.n_atoms == 45
150
151    def test_universe(self):
152        u = mda.Universe(mol2_zinc)
153        assert len(u.atoms) == self.n_atoms
154
155    def test_write_nostructure(self):
156        mytempdir = tempdir.TempDir()
157        outfile = os.path.join(mytempdir.name, 'test.mol2')
158
159        u = mda.Universe(mol2_zinc)
160        with mda.Writer(outfile) as W:
161            W.write(u.atoms)
162
163        u2 = mda.Universe(outfile)
164
165        assert len(u.atoms) == len(u2.atoms)
166
167
168def test_mol2_write_NIE():
169    mytempdir = tempdir.TempDir()
170    outfile = os.path.join(mytempdir.name, 'test.mol2')
171    u = make_Universe(trajectory=True)
172    with pytest.raises(NotImplementedError):
173        u.atoms.write(outfile)
174