Coverage for /builds/ase/ase/ase/io/xsd.py: 69.28%
166 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +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.utils import writer
12def read_xsd(fd):
13 tree = ET.parse(fd)
14 root = tree.getroot()
16 atomtreeroot = root.find('AtomisticTreeRoot')
17 # if periodic system
18 if atomtreeroot.find('SymmetrySystem') is not None:
19 symmetrysystem = atomtreeroot.find('SymmetrySystem')
20 mappingset = symmetrysystem.find('MappingSet')
21 mappingfamily = mappingset.find('MappingFamily')
22 system = mappingfamily.find('IdentityMapping')
24 coords = []
25 cell = []
26 formula = ''
28 for atom in system:
29 if atom.tag == 'Atom3d':
30 symbol = atom.get('Components')
31 formula += symbol
33 xyz = atom.get('XYZ')
34 if xyz:
35 coord = [float(coord) for coord in xyz.split(',')]
36 else:
37 coord = [0.0, 0.0, 0.0]
38 coords.append(coord)
39 elif atom.tag == 'SpaceGroup':
40 avec = [float(vec) for vec in atom.get('AVector').split(',')]
41 bvec = [float(vec) for vec in atom.get('BVector').split(',')]
42 cvec = [float(vec) for vec in atom.get('CVector').split(',')]
44 cell.append(avec)
45 cell.append(bvec)
46 cell.append(cvec)
48 atoms = Atoms(formula, cell=cell, pbc=True)
49 atoms.set_scaled_positions(coords)
50 return atoms
51 # if non-periodic system
52 elif atomtreeroot.find('Molecule') is not None:
53 system = atomtreeroot.find('Molecule')
55 coords = []
56 formula = ''
58 for atom in system:
59 if atom.tag == 'Atom3d':
60 symbol = atom.get('Components')
61 formula += symbol
63 xyz = atom.get('XYZ')
64 coord = [float(coord) for coord in xyz.split(',')]
65 coords.append(coord)
67 atoms = Atoms(formula, pbc=False)
68 atoms.set_scaled_positions(coords)
69 return atoms
72def CPK_or_BnS(element):
73 """Determine how atom is visualized"""
74 if element in ['C', 'H', 'O', 'S', 'N']:
75 visualization_choice = 'Ball and Stick'
76 else:
77 visualization_choice = 'CPK'
78 return visualization_choice
81def SetChild(parent, childname, props):
82 Child = ET.SubElement(parent, childname)
83 for key in props:
84 Child.set(key, props[key])
85 return Child
88def SetBasicChilds():
89 """
90 Basic property setup for Material Studio File
91 """
92 XSD = ET.Element('XSD')
93 XSD.set('Version', '6.0')
95 ATR = SetChild(XSD, 'AtomisticTreeRoot', dict(
96 ID='1',
97 NumProperties='40',
98 NumChildren='1',
99 ))
100 SetChild(ATR, 'Property', dict(
101 DefinedOn='ClassicalEnergyHolder',
102 Name='AngleEnergy',
103 Type='Double',
104 ))
105 SetChild(ATR, 'Property', dict(
106 DefinedOn='ClassicalEnergyHolder',
107 Name='BendBendEnergy',
108 Type='Double',
109 ))
110 SetChild(ATR, 'Property', dict(
111 DefinedOn='ClassicalEnergyHolder',
112 Name='BendTorsionBendEnergy',
113 Type='Double',
114 ))
115 SetChild(ATR, 'Property', dict(
116 DefinedOn='ClassicalEnergyHolder',
117 Name='BondEnergy',
118 Type='Double',
119 ))
120 SetChild(ATR, 'Property', dict(
121 DefinedOn='Atom',
122 Name='EFGAsymmetry',
123 Type='Double',
124 ))
125 SetChild(ATR, 'Property', dict(
126 DefinedOn='Atom',
127 Name='EFGQuadrupolarCoupling',
128 Type='Double',
129 ))
130 SetChild(ATR, 'Property', dict(
131 DefinedOn='ClassicalEnergyHolder',
132 Name='ElectrostaticEnergy',
133 Type='Double',
134 ))
135 SetChild(ATR, 'Property', dict(
136 DefinedOn='GrowthFace',
137 Name='FaceMillerIndex',
138 Type='MillerIndex',
139 ))
140 SetChild(ATR, 'Property', dict(
141 DefinedOn='GrowthFace',
142 Name='FacetTransparency',
143 Type='Float',
144 ))
145 SetChild(ATR, 'Property', dict(
146 DefinedOn='Bondable',
147 Name='Force',
148 Type='CoDirection',
149 ))
150 SetChild(ATR, 'Property', dict(
151 DefinedOn='ClassicalEnergyHolder',
152 Name='HydrogenBondEnergy',
153 Type='Double',
154 ))
155 SetChild(ATR, 'Property', dict(
156 DefinedOn='Bondable',
157 Name='ImportOrder',
158 Type='UnsignedInteger',
159 ))
160 SetChild(ATR, 'Property', dict(
161 DefinedOn='ClassicalEnergyHolder',
162 Name='InversionEnergy',
163 Type='Double',
164 ))
165 SetChild(ATR, 'Property', dict(
166 DefinedOn='Atom',
167 Name='IsBackboneAtom',
168 Type='Boolean',
169 ))
170 SetChild(ATR, 'Property', dict(
171 DefinedOn='Atom',
172 Name='IsChiralCenter',
173 Type='Boolean',
174 ))
175 SetChild(ATR, 'Property', dict(
176 DefinedOn='Atom',
177 Name='IsOutOfPlane',
178 Type='Boolean',
179 ))
180 SetChild(ATR, 'Property', dict(
181 DefinedOn='BestFitLineMonitor',
182 Name='LineExtentPadding',
183 Type='Double',
184 ))
185 SetChild(ATR, 'Property', dict(
186 DefinedOn='Linkage',
187 Name='LinkageGroupName',
188 Type='String',
189 ))
190 SetChild(ATR, 'Property', dict(
191 DefinedOn='PropertyList',
192 Name='ListIdentifier',
193 Type='String',
194 ))
195 SetChild(ATR, 'Property', dict(
196 DefinedOn='Atom',
197 Name='NMRShielding',
198 Type='Double',
199 ))
200 SetChild(ATR, 'Property', dict(
201 DefinedOn='ClassicalEnergyHolder',
202 Name='NonBondEnergy',
203 Type='Double',
204 ))
205 SetChild(ATR, 'Property', dict(
206 DefinedOn='Bondable',
207 Name='NormalMode',
208 Type='Direction',
209 ))
210 SetChild(ATR, 'Property', dict(
211 DefinedOn='Bondable',
212 Name='NormalModeFrequency',
213 Type='Double',
214 ))
215 SetChild(ATR, 'Property', dict(
216 DefinedOn='Bondable',
217 Name='OrbitalCutoffRadius',
218 Type='Double',
219 ))
220 SetChild(ATR, 'Property', dict(
221 DefinedOn='BestFitPlaneMonitor',
222 Name='PlaneExtentPadding',
223 Type='Double',
224 ))
225 SetChild(ATR, 'Property', dict(
226 DefinedOn='ClassicalEnergyHolder',
227 Name='PotentialEnergy',
228 Type='Double',
229 ))
230 SetChild(ATR, 'Property', dict(
231 DefinedOn='ScalarFieldBase',
232 Name='QuantizationValue',
233 Type='Double',
234 ))
235 SetChild(ATR, 'Property', dict(
236 DefinedOn='ClassicalEnergyHolder',
237 Name='RestraintEnergy',
238 Type='Double',
239 ))
240 SetChild(ATR, 'Property', dict(
241 DefinedOn='ClassicalEnergyHolder',
242 Name='SeparatedStretchStretchEnergy',
243 Type='Double',
244 ))
245 SetChild(ATR, 'Property', dict(
246 DefinedOn='Trajectory',
247 Name='SimulationStep',
248 Type='Integer',
249 ))
250 SetChild(ATR, 'Property', dict(
251 DefinedOn='ClassicalEnergyHolder',
252 Name='StretchBendStretchEnergy',
253 Type='Double',
254 ))
255 SetChild(ATR, 'Property', dict(
256 DefinedOn='ClassicalEnergyHolder',
257 Name='StretchStretchEnergy',
258 Type='Double',
259 ))
260 SetChild(ATR, 'Property', dict(
261 DefinedOn='ClassicalEnergyHolder',
262 Name='StretchTorsionStretchEnergy',
263 Type='Double',
264 ))
265 SetChild(ATR, 'Property', dict(
266 DefinedOn='ClassicalEnergyHolder',
267 Name='TorsionBendBendEnergy',
268 Type='Double',
269 ))
270 SetChild(ATR, 'Property', dict(
271 DefinedOn='ClassicalEnergyHolder',
272 Name='TorsionEnergy',
273 Type='Double',
274 ))
275 SetChild(ATR, 'Property', dict(
276 DefinedOn='ClassicalEnergyHolder',
277 Name='TorsionStretchEnergy',
278 Type='Double',
279 ))
280 SetChild(ATR, 'Property', dict(
281 DefinedOn='ClassicalEnergyHolder',
282 Name='ValenceCrossTermEnergy',
283 Type='Double',
284 ))
285 SetChild(ATR, 'Property', dict(
286 DefinedOn='ClassicalEnergyHolder',
287 Name='ValenceDiagonalEnergy',
288 Type='Double',
289 ))
290 SetChild(ATR, 'Property', dict(
291 DefinedOn='ClassicalEnergyHolder',
292 Name='VanDerWaalsEnergy',
293 Type='Double',
294 ))
295 SetChild(ATR, 'Property', dict(
296 DefinedOn='SymmetrySystem',
297 Name='_Stress',
298 Type='Matrix',
299 ))
300 return ATR, XSD
303def _write_xsd_html(images, connectivity=None):
304 ATR, XSD = SetBasicChilds()
305 natoms = len(images[0])
306 atom_element = images[0].get_chemical_symbols()
307 atom_cell = images[0].get_cell()
308 atom_positions = images[0].get_positions()
309 # Set up bonds
310 bonds = []
311 if connectivity is not None:
312 for i in range(connectivity.shape[0]):
313 for j in range(i + 1, connectivity.shape[0]):
314 if connectivity[i, j]:
315 bonds.append([i, j])
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)