File: CommonNeighborAnalysis.py

package info (click to toggle)
ovito 0.9.5-2
  • links: PTS
  • area: main
  • in suites: wheezy
  • size: 24,824 kB
  • sloc: cpp: 61,876; ansic: 10,002; xml: 940; python: 523; sh: 43; makefile: 15
file content (167 lines) | stat: -rwxr-xr-x 8,188 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
#!/usr/bin/env ovitos

#######################################################################
# This utility script loads an atoms file, performs the
# Common Neighbor Analysis, and writes the results
# to a new atoms file.
#
# Note: The script always assumes periodic boundary conditions.
#######################################################################

from sys import *
from atomviz.ChemicalElements import *
from atomviz.AtomImport import *
from atomviz.AtomExport import *
from AtomViz import *
from optparse import OptionParser

# Define command line options.
parser = OptionParser(usage="usage: %prog [options] inputfile outputfile", version="OVITO Common Neighbor Analysis\nVersion 1.0.0\nAuthor: Alexander Stukowski (stukowski@mm.tu-darmstadt.de)")	

parser.add_option("-c", "--cutoff", action="store", type="float", dest="cutoff", metavar="RADIUS", help="Sets the nearest neighbor cutoff to RADIUS.")
parser.add_option("-e", "--element", action="store", type="string", dest="element", metavar="ELEMENT", help="Selects the optimal CNA cutoff radius for the given chemical element.")
parser.add_option("-f", "--format", action="store", type="int", dest="format", metavar="FORMAT", help="Selects the input/output file format. The supported formats are documented in the script file.")
parser.add_option("-l", "--listformat", action="store_true", dest="listformat", default=False, help="Print list of file formats; then exit.")
parser.add_option("--exclude", action="append", type="string", dest="exclude_types", metavar="CNATYPE", help="Exclude this atom type from output. Valid types: FCC, HCP, BCC, OTHER")
parser.add_option("--coordination", action="store_true", dest="coordination", default=False, help="Calculate coordination numbers.")

# Parse command line options.
options, args = parser.parse_args()
if options.listformat:
	print "List of file formats that can be selected with the --format option:"
	print ""
	print " 1 - LAMMPS text dump file"
	print "     Input columns: (1) atom index (2) atom type (3,4,5) XYZ"
	print "     On output, the 2nd column is overwritten with the CNA type."
	print ""
	print " 2 - IMD file format"
	print "     Input and output file format is IMD."
	print "     The CNA type is written to the second last column."
	print "     The coordination number is written to the last column."
	print ""
	print " 3 - LAMMPS data format"
	print "     Output file format is LAMMPS text dump."
	print "     The following columns are written to the output file:"
	print "     (1) Atom index (2) Atom type (3,4,5) X,Y,Z (6) CNA type"
	exit(0)

if len(args) < 1:
	parser.print_help()
	exit("ERROR: No input file specified.")
if len(args) < 2:
	parser.print_help()
	exit("ERROR: No output file specified.")

# Choose a reasonable cutoff radius for the neighbor analysis.
if options.element != None:
	# Retrieve number of element.
	if options.element not in element_symbols: exit("ERROR: Could not find chemical element %s in element database." % options.element)
	elementIndex = element_symbols.index(options.element)
	options.cutoff = GetCNARadiusForElement(elementIndex)
	print "Using CNA cutoff radius for chemical element %s from database:   r=%f" % (options.element, options.cutoff)

# Check cutoff radius
if options.cutoff == None:
	parser.print_help()
	exit("ERROR: No CNA cutoff radius specified. Please use the -e or the -c options.")
if options.cutoff <= 0.0:
	exit("ERROR: Invalid CNA cutoff radius.")

# What kind of file are we dealing with?
inputFileFormat = None
outputFileFormat = None
if options.format == None:
	parser.print_help()
	exit("ERROR: You have to specify the column format of the input and output files using the -f option. See source of script file for supported file format specifiers.")

# Now define the columns in the input and in the output format.
inputColumnMapping = ColumnChannelMapping()
outputChannelMapping = ChannelColumnMapping()

if options.format == 1:
	inputFileFormat = "LAMMPS.dump.text"
	inputColumnMapping.DefineStandardColumn(0, DataChannelIdentifier.AtomIndexChannel, 0)
	inputColumnMapping.DefineStandardColumn(1, DataChannelIdentifier.AtomTypeChannel, 0)
	inputColumnMapping.DefineStandardColumn(2, DataChannelIdentifier.PositionChannel, 0)
	inputColumnMapping.DefineStandardColumn(3, DataChannelIdentifier.PositionChannel, 1)
	inputColumnMapping.DefineStandardColumn(4, DataChannelIdentifier.PositionChannel, 2)
	outputFileFormat = "LAMMPS.dump.text"
	outputChannelMapping.InsertColumn(0, DataChannelIdentifier.AtomIndexChannel, "Atom Index", 0)
	outputChannelMapping.InsertColumn(1, DataChannelIdentifier.CNATypeChannel, "CNA Atom Type", 0)
	outputChannelMapping.InsertColumn(2, DataChannelIdentifier.PositionChannel, "Position", 0)
	outputChannelMapping.InsertColumn(3, DataChannelIdentifier.PositionChannel, "Position", 1)
	outputChannelMapping.InsertColumn(4, DataChannelIdentifier.PositionChannel, "Position", 2)
elif options.format == 2:
	inputFileFormat = "IMD"
	outputFileFormat = "IMD"
elif options.format == 3:
	inputFileFormat = "LAMMPS.data"
	outputFileFormat = "LAMMPS.dump.text"
	outputChannelMapping.InsertColumn(0, DataChannelIdentifier.AtomIndexChannel, "Atom Index", 0)
	outputChannelMapping.InsertColumn(1, DataChannelIdentifier.AtomTypeChannel, "Atom Type", 0)
	outputChannelMapping.InsertColumn(2, DataChannelIdentifier.PositionChannel, "Position", 0)
	outputChannelMapping.InsertColumn(3, DataChannelIdentifier.PositionChannel, "Position", 1)
	outputChannelMapping.InsertColumn(4, DataChannelIdentifier.PositionChannel, "Position", 2)
	outputChannelMapping.InsertColumn(5, DataChannelIdentifier.CNATypeChannel, "CNA Atom Type", 0)
else:
	exit("ERROR: Invalid format number: %s" % str(options.format))

# Import input atoms.
if inputFileFormat == "LAMMPS.dump.text":
	node = ImportLAMMPSDumpFile(args[0], binary = False, columns = inputColumnMapping)
elif inputFileFormat == "LAMMPS.dump.binary":
	node = ImportLAMMPSDumpFile(args[0], binary = True, columns = inputColumnMapping)
elif inputFileFormat == "LAMMPS.data":
	node = ImportLAMMPSDataFile(args[0])
elif inputFileFormat == "XYZ":
	node = ImportXYZFile(args[0], columns = inputColumnMapping)
elif inputFileFormat == "IMD":
	node = ImportIMDAtomFile(args[0])
else:
	exit("ERROR: Unknown input file format: %s" % inputFileFormat)

# Get the scene node that contains the atoms.
atoms = node.SceneObject.Atoms

# Apply CNA modifier
cnaMod = CommonNeighborAnalysisModifier()
cnaMod.NearestNeighborList.CutoffRadius = options.cutoff
node.ApplyModifier(cnaMod)
cnaMod.PerformAnalysis(AnimManager.Instance.FrameToTime(0), False)

# Apply Coordination Number modifier
if options.coordination:
	coordinationMod = CoordinationNumberModifier()
	coordinationMod.NearestNeighborList.CutoffRadius = options.cutoff
	node.ApplyModifier(coordinationMod)
	coordinationMod.PerformAnalysis(AnimManager.Instance.FrameToTime(0), False)

# Delete atoms of a certain kind when requested
for cnaType in options.exclude_types:
	if cnaType == "FCC": cnaTypeInt = CNAAtomType.FCC
	elif cnaType == "HCP": cnaTypeInt = CNAAtomType.HCP
	elif cnaType == "BCC": cnaTypeInt = CNAAtomType.BCC
	elif cnaType == "OTHER": cnaTypeInt = CNAAtomType.Other
	else: exit("Invalid CNA atom type: \"%s\"" % cnaType)
	print "Deleting atoms of type %s" % cnaType
	selectAtomTypeMod = SelectAtomTypeModifier()
	channelRef = DataChannelReference()
	channelRef.Id = DataChannelIdentifier.CNATypeChannel
	selectAtomTypeMod.SourceDataChannel = channelRef
	selectAtomTypeMod.SetSelectedAtomType(cnaTypeInt)
	node.ApplyModifier(selectAtomTypeMod)
	node.ApplyModifier(DeleteAtomsModifier())
    
# Write results to a new output file.
if outputFileFormat == "LAMMPS.dump.text":
	ExportLAMMPSDumpFile(args[1], binary = False, columns = outputChannelMapping)
elif outputFileFormat == "LAMMPS.dump.binary":
	ExportLAMMPSDumpFile(args[1], binary = True, columns = outputChannelMapping)
elif outputFileFormat == "LAMMPS.data":
	ExportLAMMPSDataFile(args[1])
elif outputFileFormat == "XYZ":
	ExportXYZFile(args[1], columns = outputChannelMapping)
elif outputFileFormat == "IMD":
	ExportIMDAtomFile(args[1])
else:
	exit("ERROR: Unknown output file format: %s" % outputFileFormat)