Coverage for ase / io / xsd.py: 68.90%
164 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +0000
1# fmt: off
3import xml.etree.ElementTree as ET
4from xml.dom import minidom
6import numpy as np
8from ase import Atoms
9from ase.io.utils import connectivity2bonds
10from ase.utils import writer
13def read_xsd(fd):
14 tree = ET.parse(fd)
15 root = tree.getroot()
17 atomtreeroot = root.find('AtomisticTreeRoot')
18 # if periodic system
19 if atomtreeroot.find('SymmetrySystem') is not None:
20 symmetrysystem = atomtreeroot.find('SymmetrySystem')
21 mappingset = symmetrysystem.find('MappingSet')
22 mappingfamily = mappingset.find('MappingFamily')
23 system = mappingfamily.find('IdentityMapping')
25 coords = []
26 cell = []
27 formula = ''
29 for atom in system:
30 if atom.tag == 'Atom3d':
31 symbol = atom.get('Components')
32 formula += symbol
34 xyz = atom.get('XYZ')
35 if xyz:
36 coord = [float(coord) for coord in xyz.split(',')]
37 else:
38 coord = [0.0, 0.0, 0.0]
39 coords.append(coord)
40 elif atom.tag == 'SpaceGroup':
41 avec = [float(vec) for vec in atom.get('AVector').split(',')]
42 bvec = [float(vec) for vec in atom.get('BVector').split(',')]
43 cvec = [float(vec) for vec in atom.get('CVector').split(',')]
45 cell.append(avec)
46 cell.append(bvec)
47 cell.append(cvec)
49 atoms = Atoms(formula, cell=cell, pbc=True)
50 atoms.set_scaled_positions(coords)
51 return atoms
52 # if non-periodic system
53 elif atomtreeroot.find('Molecule') is not None:
54 system = atomtreeroot.find('Molecule')
56 coords = []
57 formula = ''
59 for atom in system:
60 if atom.tag == 'Atom3d':
61 symbol = atom.get('Components')
62 formula += symbol
64 xyz = atom.get('XYZ')
65 coord = [float(coord) for coord in xyz.split(',')]
66 coords.append(coord)
68 atoms = Atoms(formula, pbc=False)
69 atoms.set_scaled_positions(coords)
70 return atoms
73def CPK_or_BnS(element):
74 """Determine how atom is visualized"""
75 if element in ['C', 'H', 'O', 'S', 'N']:
76 visualization_choice = 'Ball and Stick'
77 else:
78 visualization_choice = 'CPK'
79 return visualization_choice
82def SetChild(parent, childname, props):
83 Child = ET.SubElement(parent, childname)
84 for key in props:
85 Child.set(key, props[key])
86 return Child
89def SetBasicChilds():
90 """
91 Basic property setup for Material Studio File
92 """
93 XSD = ET.Element('XSD')
94 XSD.set('Version', '6.0')
96 ATR = SetChild(XSD, 'AtomisticTreeRoot', dict(
97 ID='1',
98 NumProperties='40',
99 NumChildren='1',
100 ))
101 SetChild(ATR, 'Property', dict(
102 DefinedOn='ClassicalEnergyHolder',
103 Name='AngleEnergy',
104 Type='Double',
105 ))
106 SetChild(ATR, 'Property', dict(
107 DefinedOn='ClassicalEnergyHolder',
108 Name='BendBendEnergy',
109 Type='Double',
110 ))
111 SetChild(ATR, 'Property', dict(
112 DefinedOn='ClassicalEnergyHolder',
113 Name='BendTorsionBendEnergy',
114 Type='Double',
115 ))
116 SetChild(ATR, 'Property', dict(
117 DefinedOn='ClassicalEnergyHolder',
118 Name='BondEnergy',
119 Type='Double',
120 ))
121 SetChild(ATR, 'Property', dict(
122 DefinedOn='Atom',
123 Name='EFGAsymmetry',
124 Type='Double',
125 ))
126 SetChild(ATR, 'Property', dict(
127 DefinedOn='Atom',
128 Name='EFGQuadrupolarCoupling',
129 Type='Double',
130 ))
131 SetChild(ATR, 'Property', dict(
132 DefinedOn='ClassicalEnergyHolder',
133 Name='ElectrostaticEnergy',
134 Type='Double',
135 ))
136 SetChild(ATR, 'Property', dict(
137 DefinedOn='GrowthFace',
138 Name='FaceMillerIndex',
139 Type='MillerIndex',
140 ))
141 SetChild(ATR, 'Property', dict(
142 DefinedOn='GrowthFace',
143 Name='FacetTransparency',
144 Type='Float',
145 ))
146 SetChild(ATR, 'Property', dict(
147 DefinedOn='Bondable',
148 Name='Force',
149 Type='CoDirection',
150 ))
151 SetChild(ATR, 'Property', dict(
152 DefinedOn='ClassicalEnergyHolder',
153 Name='HydrogenBondEnergy',
154 Type='Double',
155 ))
156 SetChild(ATR, 'Property', dict(
157 DefinedOn='Bondable',
158 Name='ImportOrder',
159 Type='UnsignedInteger',
160 ))
161 SetChild(ATR, 'Property', dict(
162 DefinedOn='ClassicalEnergyHolder',
163 Name='InversionEnergy',
164 Type='Double',
165 ))
166 SetChild(ATR, 'Property', dict(
167 DefinedOn='Atom',
168 Name='IsBackboneAtom',
169 Type='Boolean',
170 ))
171 SetChild(ATR, 'Property', dict(
172 DefinedOn='Atom',
173 Name='IsChiralCenter',
174 Type='Boolean',
175 ))
176 SetChild(ATR, 'Property', dict(
177 DefinedOn='Atom',
178 Name='IsOutOfPlane',
179 Type='Boolean',
180 ))
181 SetChild(ATR, 'Property', dict(
182 DefinedOn='BestFitLineMonitor',
183 Name='LineExtentPadding',
184 Type='Double',
185 ))
186 SetChild(ATR, 'Property', dict(
187 DefinedOn='Linkage',
188 Name='LinkageGroupName',
189 Type='String',
190 ))
191 SetChild(ATR, 'Property', dict(
192 DefinedOn='PropertyList',
193 Name='ListIdentifier',
194 Type='String',
195 ))
196 SetChild(ATR, 'Property', dict(
197 DefinedOn='Atom',
198 Name='NMRShielding',
199 Type='Double',
200 ))
201 SetChild(ATR, 'Property', dict(
202 DefinedOn='ClassicalEnergyHolder',
203 Name='NonBondEnergy',
204 Type='Double',
205 ))
206 SetChild(ATR, 'Property', dict(
207 DefinedOn='Bondable',
208 Name='NormalMode',
209 Type='Direction',
210 ))
211 SetChild(ATR, 'Property', dict(
212 DefinedOn='Bondable',
213 Name='NormalModeFrequency',
214 Type='Double',
215 ))
216 SetChild(ATR, 'Property', dict(
217 DefinedOn='Bondable',
218 Name='OrbitalCutoffRadius',
219 Type='Double',
220 ))
221 SetChild(ATR, 'Property', dict(
222 DefinedOn='BestFitPlaneMonitor',
223 Name='PlaneExtentPadding',
224 Type='Double',
225 ))
226 SetChild(ATR, 'Property', dict(
227 DefinedOn='ClassicalEnergyHolder',
228 Name='PotentialEnergy',
229 Type='Double',
230 ))
231 SetChild(ATR, 'Property', dict(
232 DefinedOn='ScalarFieldBase',
233 Name='QuantizationValue',
234 Type='Double',
235 ))
236 SetChild(ATR, 'Property', dict(
237 DefinedOn='ClassicalEnergyHolder',
238 Name='RestraintEnergy',
239 Type='Double',
240 ))
241 SetChild(ATR, 'Property', dict(
242 DefinedOn='ClassicalEnergyHolder',
243 Name='SeparatedStretchStretchEnergy',
244 Type='Double',
245 ))
246 SetChild(ATR, 'Property', dict(
247 DefinedOn='Trajectory',
248 Name='SimulationStep',
249 Type='Integer',
250 ))
251 SetChild(ATR, 'Property', dict(
252 DefinedOn='ClassicalEnergyHolder',
253 Name='StretchBendStretchEnergy',
254 Type='Double',
255 ))
256 SetChild(ATR, 'Property', dict(
257 DefinedOn='ClassicalEnergyHolder',
258 Name='StretchStretchEnergy',
259 Type='Double',
260 ))
261 SetChild(ATR, 'Property', dict(
262 DefinedOn='ClassicalEnergyHolder',
263 Name='StretchTorsionStretchEnergy',
264 Type='Double',
265 ))
266 SetChild(ATR, 'Property', dict(
267 DefinedOn='ClassicalEnergyHolder',
268 Name='TorsionBendBendEnergy',
269 Type='Double',
270 ))
271 SetChild(ATR, 'Property', dict(
272 DefinedOn='ClassicalEnergyHolder',
273 Name='TorsionEnergy',
274 Type='Double',
275 ))
276 SetChild(ATR, 'Property', dict(
277 DefinedOn='ClassicalEnergyHolder',
278 Name='TorsionStretchEnergy',
279 Type='Double',
280 ))
281 SetChild(ATR, 'Property', dict(
282 DefinedOn='ClassicalEnergyHolder',
283 Name='ValenceCrossTermEnergy',
284 Type='Double',
285 ))
286 SetChild(ATR, 'Property', dict(
287 DefinedOn='ClassicalEnergyHolder',
288 Name='ValenceDiagonalEnergy',
289 Type='Double',
290 ))
291 SetChild(ATR, 'Property', dict(
292 DefinedOn='ClassicalEnergyHolder',
293 Name='VanDerWaalsEnergy',
294 Type='Double',
295 ))
296 SetChild(ATR, 'Property', dict(
297 DefinedOn='SymmetrySystem',
298 Name='_Stress',
299 Type='Matrix',
300 ))
301 return ATR, XSD
304def _write_xsd_html(images, connectivity=None):
305 ATR, XSD = SetBasicChilds()
306 natoms = len(images[0])
307 atom_element = images[0].get_chemical_symbols()
308 atom_cell = images[0].get_cell()
309 atom_positions = images[0].get_positions()
310 # Set up bonds
311 if connectivity is not None:
312 bonds = connectivity2bonds(connectivity)
313 else:
314 bonds = []
316 nbonds = len(bonds)
318 # non-periodic system
319 if not images[0].pbc.all():
320 Molecule = SetChild(ATR, 'Molecule', dict(
321 ID='2',
322 NumChildren=str(natoms + nbonds),
323 Name='Lattice="1.0',
324 ))
325 # writing images[0]
326 for x in range(natoms):
327 Props = dict(
328 ID=str(x + 3),
329 Name=(atom_element[x] + str(x + 1)),
330 UserID=str(x + 1),
331 DisplayStyle=CPK_or_BnS(atom_element[x]),
332 XYZ=','.join('%1.16f' % xi for xi in atom_positions[x]),
333 Components=atom_element[x],
334 )
335 bondstr = []
336 for i, bond in enumerate(bonds):
337 if x in bond:
338 bondstr.append(str(i + 3 + natoms))
339 if bondstr:
340 Props['Connections'] = ','.join(bondstr)
341 SetChild(Molecule, 'Atom3d', Props)
342 for x in range(nbonds):
343 SetChild(Molecule, 'Bond', dict(
344 ID=str(x + 3 + natoms),
345 Connects='%i,%i' % (bonds[x][0] + 3, bonds[x][1] + 3),
346 ))
347 # periodic system
348 else:
349 atom_positions = np.dot(atom_positions, np.linalg.inv(atom_cell))
350 Props = dict(
351 ID='2',
352 Mapping='3',
353 Children=','.join(map(str, range(4, natoms + nbonds + 5))),
354 Normalized='1',
355 Name='SymmSys',
356 UserID=str(natoms + 18),
357 XYZ='0.00000000000000,0.00000000000000,0.000000000000000',
358 OverspecificationTolerance='0.05',
359 PeriodicDisplayType='Original',
360 )
361 SymmSys = SetChild(ATR, 'SymmetrySystem', Props)
363 Props = dict(
364 ID=str(natoms + nbonds + 5),
365 SymmetryDefinition=str(natoms + 4),
366 ActiveSystem='2',
367 NumFamilies='1',
368 OwnsTotalConstraintMapping='1',
369 TotalConstraintMapping='3',
370 )
371 MappngSet = SetChild(SymmSys, 'MappingSet', Props)
373 Props = dict(ID=str(natoms + nbonds + 6), NumImageMappings='0')
374 MappngFamily = SetChild(MappngSet, 'MappingFamily', Props)
376 Props = dict(
377 ID=str(natoms + len(bonds) + 7),
378 Element='1,0,0,0,0,1,0,0,0,0,1,0',
379 Constraint='1,0,0,0,0,1,0,0,0,0,1,0',
380 MappedObjects=','.join(map(str, range(4, natoms + len(bonds) + 4))),
381 DefectObjects='%i,%i' % (natoms + nbonds + 4, natoms + nbonds + 8),
382 NumImages=str(natoms + len(bonds)),
383 NumDefects='2',
384 )
385 IdentMappng = SetChild(MappngFamily, 'IdentityMapping', Props)
387 SetChild(MappngFamily, 'MappingRepairs', {'NumRepairs': '0'})
389 # writing atoms
390 for x in range(natoms):
391 Props = dict(
392 ID=str(x + 4),
393 Mapping=str(natoms + len(bonds) + 7),
394 Parent='2',
395 Name=(atom_element[x] + str(x + 1)),
396 UserID=str(x + 1),
397 DisplayStyle=CPK_or_BnS(atom_element[x]),
398 Components=atom_element[x],
399 XYZ=','.join(['%1.16f' % xi for xi in atom_positions[x]]),
400 )
401 bondstr = []
402 for i, bond in enumerate(bonds):
403 if x in bond:
404 bondstr.append(str(i + 4 * natoms + 1))
405 if bondstr:
406 Props['Connections'] = ','.join(bondstr)
407 SetChild(IdentMappng, 'Atom3d', Props)
409 for x in range(len(bonds)):
410 SetChild(IdentMappng, 'Bond', dict(
411 ID=str(x + 4 + natoms + 1),
412 Mapping=str(natoms + len(bonds) + 7),
413 Parent='2',
414 Connects='%i,%i' % (bonds[x][0] + 4, bonds[x][1] + 4),
415 ))
417 Props = dict(
418 ID=str(natoms + 4),
419 Parent='2',
420 Children=str(natoms + len(bonds) + 8),
421 DisplayStyle='Solid',
422 XYZ='0.00,0.00,0.00',
423 Color='0,0,0,0',
424 AVector=','.join(['%1.16f' % atom_cell[0, x] for x in range(3)]),
425 BVector=','.join(['%1.16f' % atom_cell[1, x] for x in range(3)]),
426 CVector=','.join(['%1.16f' % atom_cell[2, x] for x in range(3)]),
427 OrientationBase='C along Z, B in YZ plane',
428 Centering='3D Primitive-Centered',
429 Lattice='3D Triclinic',
430 GroupName='GroupName',
431 Operators='1,0,0,0,0,1,0,0,0,0,1,0',
432 DisplayRange='0,1,0,1,0,1',
433 LineThickness='2',
434 CylinderRadius='0.2',
435 LabelAxes='1',
436 ActiveSystem='2',
437 ITNumber='1',
438 LongName='P 1',
439 Qualifier='Origin-1',
440 SchoenfliesName='C1-1',
441 System='Triclinic',
442 Class='1',
443 )
444 SetChild(IdentMappng, 'SpaceGroup', Props)
446 SetChild(IdentMappng, 'ReciprocalLattice3D', dict(
447 ID=str(natoms + len(bonds) + 8),
448 Parent=str(natoms + 4),
449 ))
451 SetChild(MappngSet, 'InfiniteMapping', dict(
452 ID='3',
453 Element='1,0,0,0,0,1,0,0,0,0,1,0',
454 MappedObjects='2',
455 ))
457 return XSD, ATR
460@writer
461def write_xsd(fd, images, connectivity=None):
462 """Takes Atoms object, and write materials studio file
463 atoms: Atoms object
464 filename: path of the output file
465 connectivity: number of atoms by number of atoms matrix for connectivity
466 between atoms (0 not connected, 1 connected)
468 note: material studio file cannot use a partial periodic system. If partial
469 perodic system was inputted, full periodicity was assumed.
470 """
471 if hasattr(images, 'get_positions'):
472 images = [images]
474 XSD, _ATR = _write_xsd_html(images, connectivity)
476 # Return a pretty-printed XML string for the Element.
477 rough_string = ET.tostring(XSD, 'utf-8')
478 reparsed = minidom.parseString(rough_string)
479 Document = reparsed.toprettyxml(indent='\t')
481 fd.write(Document)