← Back to branch summary

~wavemol/wavemol/wavemol.fileaccess

« back to all changes in this revision

Viewing changes to lib/wavemol/fileaccess/general/XYZFile.py

  • Committer: Stefano Borini
  • Date: 2009-12-23 12:28:39 UTC
  • Revision ID: stefano.borini@gmail.com-20091223122839-yru9p52de35b7yns
added XYZfile, CSVFile and tests. Not ready, needs cleanup

Show diffs side-by-side

added added

removed removed

Lines of Context:
lib/wavemol/fileaccess/general/XYZFile.py
 
1
import re
 
2
from theochempy._theochempy import Measure
 
3
from theochempy._theochempy import Units
 
4
from theochempy._theochempy.Chemistry import PeriodicTable
 
5
 
 
6
__all__ = ['ParseError', 'XYZFile']
 
7
 
 
8
class ParseError(Exception): pass
 
9
 
 
10
class Molecule:
 
11
    def __init__(self):
 
12
        self.comment = None
 
13
        self.atom_list = []
 
14
 
 
15
class XYZFile:
 
16
 
 
17
    def __init__(self, filename = None):
 
18
        self._filename = filename
 
19
        self._molecules = []
 
20
 
 
21
        if filename is None:
 
22
            return
 
23
 
 
24
        f=file(self._filename,"r")
 
25
       
 
26
        STATE_START, STATE_NUM_ATOMS, STATE_COMMENT, STATE_ATOM_ROW = range(4)
 
27
 
 
28
        state = STATE_START 
 
29
 
 
30
        for line_num, line in enumerate(f):
 
31
            m = re.search("^\s*\d+\s*$", line)
 
32
            if m is not None:
 
33
                if state == STATE_START or state == STATE_ATOM_ROW:
 
34
                    current_molecule = Molecule()
 
35
                    self._molecules.append(current_molecule)
 
36
                    state = STATE_NUM_ATOMS
 
37
                    continue
 
38
                else:
 
39
                    raise ParseError("Found number of atoms at invalid position, line %d, file %s. " % (line_num, filename))
 
40
                    
 
41
 
 
42
            m = re.search("^\s*([A-Z]+[a-z]*[a-z]*)\s+(-?\d*\.\d+)\s*(-?\d*\.\d+)\s*(-?\d*\.\d+)\s+$", line)
 
43
            if m is not None:
 
44
                if state == STATE_COMMENT or state == STATE_ATOM_ROW:
 
45
                    element = PeriodicTable.getElementBySymbol(m.group(1))
 
46
                    if element is None:
 
47
                        raise ParseError("Invalid element symbol %s at line %d, file %s" % (m.group(1), line_num, filename))
 
48
                    
 
49
                    position = Measure.Measure( ( float(m.group(2)), float(m.group(3)), float(m.group(4))), Units.angstrom )
 
50
 
 
51
                    current_molecule.atom_list.append( (element, position) )
 
52
                    state = STATE_ATOM_ROW
 
53
                    continue
 
54
                else:
 
55
                    raise ParseError("Found atom line at invalid position, line %d, file %s. " % (line_num, filename))
 
56
                
 
57
            m = re.search("^\s*.*\s*$", line)
 
58
            if m is not None:
 
59
                if state == STATE_NUM_ATOMS:
 
60
                    current_molecule.comment = line.strip()
 
61
                    state = STATE_COMMENT
 
62
                    continue
 
63
                else:
 
64
                    raise ParseError("Found comment line at invalid position, line %d, file $s. " % (line_num, filename))
 
65
        
 
66
        f.close()
 
67
        
 
68
    def comment(self, molecule_index=-1):
 
69
        if len(self._molecules) == 0:
 
70
            return None
 
71
        return self._molecules[molecule_index].comment
 
72
 
 
73
    def numOfMolecules(self):
 
74
        return len(self._molecules)
 
75
 
 
76
    def numOfAtoms(self, molecule_index=-1):
 
77
        if len(self._molecules) == 0:
 
78
            return None
 
79
        return len(self._molecules[molecule_index].atom_list)
 
80
 
 
81
    def atom(self, *args):
 
82
        if len(args) == 1:
 
83
            atom_index = args[0]
 
84
            molecule_index = -1
 
85
        elif len(args) == 2:
 
86
            molecule_index = args[0]
 
87
            atom_index = args[1]
 
88
        else:
 
89
            raise TypeError("Invalid number of arguments %d" % len(args))
 
90
 
 
91
        if len(self._molecules) == 0:
 
92
            return None
 
93
        return self._molecules[molecule_index].atom_list[atom_index]
 
94
 
 
95
 
 
96
    def createMolecule(self, molecule_index=None):
 
97
        current_molecule = Molecule()
 
98
 
 
99
        if molecule_index is None:
 
100
            self._molecules.append(current_molecule)
 
101
            return len(self._molecules)-1
 
102
        else:
 
103
            self._molecules.insert(molecule_index, current_molecule)
 
104
            return molecule_index
 
105
 
 
106
     
 
107
    def setComment(self, *args):
 
108
        if len(args) == 1:
 
109
            comment = args[0]
 
110
            molecule_index = -1
 
111
        elif len(args) == 2:
 
112
            molecule_index = args[0]
 
113
            comment = args[1]
 
114
        else: 
 
115
            raise TypeError("Invalid number of arguments %d" % len(args))
 
116
 
 
117
        self._molecules[molecule_index].comment = comment
 
118
        
 
119
 
 
120
    def addAtom(self, *args):
 
121
        if len(args) == 1:
 
122
            atom = args[0]
 
123
            molecule_index = -1
 
124
        elif len(args) == 2:
 
125
            molecule_index = args[0]
 
126
            atom = args[1]
 
127
        else:
 
128
            raise TypeError("Invalid number of arguments %d" % len(args))
 
129
 
 
130
        self._molecules[molecule_index].atom_list.append(atom)
 
131
 
 
132
    def setAtom(self, *args):
 
133
        if len(args) == 2:
 
134
            atom_index = args[0]
 
135
            atom = args[1]
 
136
            molecule_index = -1
 
137
        elif len(args) == 2:
 
138
            molecule_index = args[0]
 
139
            atom_index = args[1]
 
140
            atom = args[2]
 
141
        else:
 
142
            raise TypeError("Invalid number of arguments %d" % len(args))
 
143
        self._molecules[molecule_index].atom_list[atom_index] = atom
 
144
 
 
145
    def saveTo(self, filename):
 
146
        f = file(filename, "w")
 
147
        for molecule in self._molecules:
 
148
            f.write(" %d\n" % len(molecule.atom_list))
 
149
            f.write("%s\n" % molecule.comment)
 
150
            for atom in molecule.atom_list:
 
151
                element, coordinate = atom
 
152
                coordinate_angstromg = coordinate.asUnit(Units.angstrom)
 
153
                f.write("%s  %20.10f %20.10f %20.10f\n" % (element.symbol(), coordinate_angstromg.value()[0], coordinate_angstromg.value()[1], coordinate_angstromg.value()[2]))
 
154
 
 
155
        f.close()