Skip to content
This repository was archived by the owner on Dec 8, 2024. It is now read-only.

Commit 50cc6a7

Browse files
WardLTlarsbratholm
authored andcommitted
Add ability to create compounds from file-like objects (#103)
* Made Compound able to read file-like objects I added a six dependency to make life easier * Give Compound a name only if there is a filename
1 parent 4143b49 commit 50cc6a7

File tree

3 files changed

+22
-9
lines changed

3 files changed

+22
-9
lines changed

qml/data/compound.py

+11-7
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
from __future__ import print_function
2424

25+
from six import string_types
2526
import numpy as np
2627
import collections
2728

@@ -341,20 +342,23 @@ def generate_acsf(self, elements = [1,6,7,8,16], nRs2 = 3, nRs3 = 3, nTs = 3, et
341342
def read_xyz(self, filename):
342343
"""(Re-)initializes the Compound-object with data from an xyz-file.
343344
344-
:param filename: Input xyz-filename.
345-
:type filename: string
346-
"""
345+
:param filename: Input xyz-filename or file-like obejct
346+
:type filename: string or file-like object
347+
"""
347348

348-
f = open(filename, "r")
349-
lines = f.readlines()
350-
f.close()
349+
if isinstance(filename, string_types):
350+
with open(filename, "r") as f:
351+
lines = f.readlines()
352+
else:
353+
lines = filename.readlines()
351354

352355
self.natoms = int(lines[0])
353356
self.atomtypes = []
354357
self.nuclear_charges = np.empty(self.natoms, dtype=int)
355358
self.coordinates = np.empty((self.natoms, 3), dtype=float)
356359

357-
self.name = filename
360+
# Give the Compound a name if it is a string
361+
self.name = filename if isinstance(filename, string_types) else "Compound"
358362

359363
for i, line in enumerate(lines[2:self.natoms+2]):
360364
tokens = line.split()

requirements.txt

+1
Original file line numberDiff line numberDiff line change
@@ -2,3 +2,4 @@ numpy>=1.13
22
scipy
33
scikit-learn
44
ase
5+
six

test/test_compound.py

+10-2
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,16 @@ def test_compound():
5050
ref_atomtypes = ['C', 'Cl', 'Br', 'H', 'H']
5151
ref_charges = [ 6, 17, 35, 1 , 1]
5252

53-
assert compare_lists(ref_atomtypes, c.atomtypes), "Failed parsing atomtypes"
54-
assert compare_lists(ref_charges, c.nuclear_charges), "Failed parsing nuclear_charges"
53+
assert compare_lists(ref_atomtypes, c2.atomtypes), "Failed parsing atomtypes"
54+
assert compare_lists(ref_charges, c2.nuclear_charges), "Failed parsing nuclear_charges"
55+
56+
# Test extended xyz from a file pointer rather than a filename
57+
with open(test_dir + "/data/compound_test.exyz") as fp:
58+
c3 = Compound(xyz=fp)
59+
60+
assert compare_lists(ref_atomtypes, c3.atomtypes), "Failed parsing atomtypes"
61+
assert compare_lists(ref_charges, c3.nuclear_charges), "Failed parsing nuclear_charges"
62+
5563

5664
if __name__ == "__main__":
5765

0 commit comments

Comments
 (0)