diff --git a/.gitignore b/.gitignore index 196c9310..a11b8556 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,4 @@ chemfiles.egg-info/ .tox/ MANIFEST .coverage +*.pyc diff --git a/doc/tutorials.rst b/doc/tutorials.rst index 5f2d4c1e..b66b03b1 100644 --- a/doc/tutorials.rst +++ b/doc/tutorials.rst @@ -28,28 +28,20 @@ Then we open a Trajectory and read the first frame: :language: python :lines: 7-8 -We can now create a list to store the indices of the atoms with ``x < 5``, and -get the positions of the atoms in the frame with the :py:func:`Frame.positions` -function +Iterating through the atoms in the frame, we store the indices of the atoms with +``x < 5`` in a list. ``len(frame.atoms)`` gives the number of atoms in the +frame, which is also the size of the ``frame.positions`` array. This array is a +numpy array which shape is ``(len(frame), 3)``. .. literalinclude:: ../examples/indexes.py :language: python - :lines: 10-11 + :lines: 10-13 -Iterating through the atoms in the frame, we get the ones matching our -condition. ``len(frame)`` gives the number of atoms in the frame, which is also -the size of the ``positions`` array. This array is a numpy array which shape is -``(len(frame), 3)``. +We can then print our results .. literalinclude:: ../examples/indexes.py :language: python - :lines: 13-15 - -And finally we can print our results - -.. literalinclude:: ../examples/indexes.py - :language: python - :lines: 17-19 + :lines: 15-17 .. htmlhidden:: :toggle: Click here to see the whole program @@ -62,7 +54,7 @@ And finally we can print our results For more information about reading frame in a trajectory, see the following functions: -- :py:func:`Trajectory.nsteps` gives the number of frame in a Trajectory. +- :py:attr:`Trajectory.nsteps` is the number of frame in a Trajectory. - :py:func:`Trajectory.read_step` to directlty read a given step. - :py:func:`Trajectory.set_cell` and :py:func:`Trajectory.set_topology` to specify an unit cell or a topology for all frames in a trajectory. @@ -98,7 +90,7 @@ We can then set the atomic positions: .. literalinclude:: ../examples/generate.py :language: python - :lines: 20-23 + :lines: 20-22 Another possibility is to directly add atoms to the frame. Here we define a second molecule representing carbon dioxyde. :py:func:`Frame.add_atom` takes @@ -106,20 +98,20 @@ two arguments: the atom, and the position of the atom as a 3-element list .. literalinclude:: ../examples/generate.py :language: python - :lines: 25-29 + :lines: 24-28 Finally, we can set the :py:class:`UnitCell` associated with this frame. .. literalinclude:: ../examples/generate.py :language: python - :lines: 31 + :lines: 30 Now that our frame is constructed, it is time to write it to a file. For that, we open a trajectory in write (``'w'``) mode, and write to it. .. literalinclude:: ../examples/generate.py :language: python - :lines: 33-34 + :lines: 32-33 .. htmlhidden:: :toggle: Click here to see the whole program @@ -176,13 +168,20 @@ frame. :language: python :lines: 14-15 -Finally, we can write the cleaned frame to the output file, and start the next +We can then write the cleaned frame to the output file, and start the next iteration. .. literalinclude:: ../examples/select.py :language: python :lines: 16 +Finally, we close the input and output files to ensure that all the written data +is flushed to the disk. + +.. literalinclude:: ../examples/select.py + :language: python + :lines: 18-19 + .. htmlhidden:: :toggle: Click here to see the whole program :before-not-html: The whole program look like this: diff --git a/examples/generate.py b/examples/generate.py index a0646466..8aab0267 100644 --- a/examples/generate.py +++ b/examples/generate.py @@ -6,21 +6,20 @@ from chemfiles import Topology, Frame, Atom, UnitCell, Trajectory topology = Topology() -topology.add_atom(Atom("H")) -topology.add_atom(Atom("O")) -topology.add_atom(Atom("H")) +topology.atoms.append(Atom("H")) +topology.atoms.append(Atom("O")) +topology.atoms.append(Atom("H")) topology.add_bond(0, 1) topology.add_bond(2, 1) frame = Frame() frame.resize(3) -frame.set_topology(topology) +frame.topology = topology -positions = frame.positions() -positions[0, :] = np.array([1.0, 0.0, 0.0]) -positions[1, :] = np.array([0.0, 0.0, 0.0]) -positions[2, :] = np.array([0.0, 1.0, 0.0]) +frame.positions[0, :] = np.array([1.0, 0.0, 0.0]) +frame.positions[1, :] = np.array([0.0, 0.0, 0.0]) +frame.positions[2, :] = np.array([0.0, 1.0, 0.0]) frame.add_atom(Atom("O"), [5.0, 0.0, 0.0]) frame.add_atom(Atom("C"), [6.0, 0.0, 0.0]) @@ -28,7 +27,7 @@ frame.add_bond(3, 4) frame.add_bond(4, 5) -frame.set_cell(UnitCell(10, 10, 10)) +frame.cell = UnitCell(10, 10, 10) -trajectory = Trajectory("water-co2.pdb", 'w') -trajectory.write(frame) +with Trajectory("water-co2.pdb", 'w') as trajectory: + trajectory.write(frame) diff --git a/examples/indexes.py b/examples/indexes.py index 96af5135..b99c47a9 100644 --- a/examples/indexes.py +++ b/examples/indexes.py @@ -4,14 +4,12 @@ #!/usr/bin/env python from chemfiles import Trajectory -trajectory = Trajectory("filename.xyz") -frame = trajectory.read() +with Trajectory("filename.xyz") as trajectory: + frame = trajectory.read() less_than_five = [] -positions = frame.positions() - -for i in range(len(frame)): - if positions[i, 0] < 5: +for i in range(len(frame.atoms)): + if frame.positions[i, 0] < 5: less_than_five.append(i) print("Atoms with x < 5: ") diff --git a/examples/select.py b/examples/select.py index 5af9c723..a4f5c931 100644 --- a/examples/select.py +++ b/examples/select.py @@ -14,3 +14,6 @@ for i in reversed(sorted(to_remove)): frame.remove(i) output.write(frame) + +trajectory.close() +output.close() diff --git a/tests/examples.py b/tests/examples.py new file mode 100644 index 00000000..8e7d7c3a --- /dev/null +++ b/tests/examples.py @@ -0,0 +1,57 @@ +# -*- coding=utf-8 -*- +from __future__ import absolute_import, print_function, unicode_literals +import unittest +import tempfile +import shutil +import os + +from chemfiles import Trajectory, Frame, Atom + +ROOT = os.path.dirname(__file__) + + +class TestExamples(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls._cwd = os.getcwd() + cls._tmpdir = tempfile.mkdtemp() + os.chdir(cls._tmpdir) + + @classmethod + def tearDownClass(cls): + os.chdir(cls._cwd) + shutil.rmtree(cls._tmpdir) + + def test_generate(self): + path = os.path.join(ROOT, "..", "examples", "generate.py") + exec(open(path).read(), globals()) + + def test_indexes(self): + # Create an input file + frame = Frame() + for i in range(120): + frame.add_atom(Atom('X'), [i % 10, i + 1 % 10, i + 2 % 10]) + + with Trajectory("filename.xyz", "w") as file: + file.write(frame) + + path = os.path.join(ROOT, "..", "examples", "indexes.py") + # disable output + exec(open(path).read(), globals(), {'print': lambda _: None}) + + def test_select(self): + # Create an input file + frame = Frame() + NAMES = ["N", "Zn", "C", "O"] + for i in range(120): + frame.add_atom(Atom(NAMES[i % 4]), [i % 10, i + 1 % 10, i + 2 % 10]) + + with Trajectory("input.arc", "w")as file: + file.write(frame) + + path = os.path.join(ROOT, "..", "examples", "select.py") + exec(open(path).read(), globals()) + + +if __name__ == "__main__": + unittest.main()