root / ase / visualize / primiplotter.py @ 3
Historique | Voir | Annoter | Télécharger (34,75 ko)
1 | 1 | tkerber | """An experimental package for making plots during a simulation.
|
---|---|---|---|
2 | 1 | tkerber |
|
3 | 1 | tkerber | A PrimiPlotter can plot a list of atoms on one or more output devices.
|
4 | 1 | tkerber | """
|
5 | 1 | tkerber | |
6 | 1 | tkerber | from numpy import * |
7 | 1 | tkerber | from ase.visualize.colortable import color_table |
8 | 1 | tkerber | import ase.data |
9 | 1 | tkerber | import sys, os, time, weakref |
10 | 1 | tkerber | |
11 | 1 | tkerber | class PrimiPlotterBase: |
12 | 1 | tkerber | "Base class for PrimiPlotter and Povrayplotter."
|
13 | 1 | tkerber | #def set_dimensions(self, dims):
|
14 | 1 | tkerber | # "Set the size of the canvas (a 2-tuple)."
|
15 | 1 | tkerber | # self.dims = dims
|
16 | 1 | tkerber | |
17 | 1 | tkerber | def set_rotation(self, rotation): |
18 | 1 | tkerber | "Set the rotation angles (in degrees)."
|
19 | 1 | tkerber | self.angles[:] = array(rotation) * (pi/180) |
20 | 1 | tkerber | |
21 | 1 | tkerber | def set_radii(self, radii): |
22 | 1 | tkerber | """Set the atomic radii. Give an array or a single number."""
|
23 | 1 | tkerber | self.radius = radii
|
24 | 1 | tkerber | |
25 | 1 | tkerber | def set_colors(self, colors): |
26 | 1 | tkerber | """Explicitly set the colors of the atoms."""
|
27 | 1 | tkerber | self.colors = colors
|
28 | 1 | tkerber | |
29 | 1 | tkerber | def set_color_function(self, colors): |
30 | 1 | tkerber | """Set a color function, to be used to color the atoms."""
|
31 | 1 | tkerber | if callable(colors): |
32 | 1 | tkerber | self.colorfunction = colors
|
33 | 1 | tkerber | else:
|
34 | 1 | tkerber | raise TypeError, "The color function is not callable." |
35 | 1 | tkerber | |
36 | 1 | tkerber | def set_invisible(self, inv): |
37 | 1 | tkerber | """Choose invisible atoms."""
|
38 | 1 | tkerber | self.invisible = inv
|
39 | 1 | tkerber | |
40 | 1 | tkerber | def set_invisibility_function(self, invfunc): |
41 | 1 | tkerber | """Set an invisibility function."""
|
42 | 1 | tkerber | if callable(invfunc): |
43 | 1 | tkerber | self.invisibilityfunction = invfunc
|
44 | 1 | tkerber | else:
|
45 | 1 | tkerber | raise TypeError, "The invisibility function is not callable." |
46 | 1 | tkerber | |
47 | 1 | tkerber | def set_cut(self, xmin=None, xmax=None, ymin=None, ymax=None, |
48 | 1 | tkerber | zmin=None, zmax=None): |
49 | 1 | tkerber | self.cut = {"xmin":xmin, "xmax":xmax, "ymin":ymin, "ymax":ymax, |
50 | 1 | tkerber | "zmin":zmin, "zmax":zmax} |
51 | 1 | tkerber | |
52 | 1 | tkerber | def update(self, newatoms = None): |
53 | 1 | tkerber | """Cause a plot (respecting the interval setting).
|
54 | 1 | tkerber |
|
55 | 1 | tkerber | update causes a plot to be made. If the interval variable was
|
56 | 1 | tkerber | specified when the plotter was create, it will only produce a
|
57 | 1 | tkerber | plot with that interval. update takes an optional argument,
|
58 | 1 | tkerber | newatoms, which can be used to replace the list of atoms with
|
59 | 1 | tkerber | a new one.
|
60 | 1 | tkerber | """
|
61 | 1 | tkerber | if newatoms is not None: |
62 | 1 | tkerber | self.atoms = newatoms
|
63 | 1 | tkerber | if self.skipnext <= 0: |
64 | 1 | tkerber | self.plot()
|
65 | 1 | tkerber | self.skipnext = self.interval |
66 | 1 | tkerber | self.skipnext -= 1 |
67 | 1 | tkerber | |
68 | 1 | tkerber | def set_log(self, log): |
69 | 1 | tkerber | """Sets a file for logging.
|
70 | 1 | tkerber |
|
71 | 1 | tkerber | log may be an open file or a filename.
|
72 | 1 | tkerber | """
|
73 | 1 | tkerber | if hasattr(log, "write"): |
74 | 1 | tkerber | self.logfile = log
|
75 | 1 | tkerber | self.ownlogfile = False |
76 | 1 | tkerber | else:
|
77 | 1 | tkerber | self.logfile = open(log, "w") |
78 | 1 | tkerber | self.ownlogfile = True |
79 | 1 | tkerber | |
80 | 1 | tkerber | def log(self, message): |
81 | 1 | tkerber | """logs a message to the file set by set_log."""
|
82 | 1 | tkerber | if self.logfile is not None: |
83 | 1 | tkerber | self.logfile.write(message+"\n") |
84 | 1 | tkerber | self.logfile.flush()
|
85 | 1 | tkerber | self._verb(message)
|
86 | 1 | tkerber | |
87 | 1 | tkerber | def _verb(self, txt): |
88 | 1 | tkerber | if self.verbose: |
89 | 1 | tkerber | sys.stderr.write(txt+"\n")
|
90 | 1 | tkerber | |
91 | 1 | tkerber | def _starttimer(self): |
92 | 1 | tkerber | self.starttime = time.time()
|
93 | 1 | tkerber | |
94 | 1 | tkerber | def _stoptimer(self): |
95 | 1 | tkerber | elapsedtime = time.time() - self.starttime
|
96 | 1 | tkerber | self.totaltime = self.totaltime + elapsedtime |
97 | 1 | tkerber | print "plotting time %s sec (total %s sec)" % (elapsedtime, |
98 | 1 | tkerber | self.totaltime)
|
99 | 1 | tkerber | |
100 | 1 | tkerber | def _getpositions(self): |
101 | 1 | tkerber | return self.atoms.get_positions() |
102 | 1 | tkerber | |
103 | 1 | tkerber | def _getradii(self): |
104 | 1 | tkerber | if self.radius is not None: |
105 | 1 | tkerber | if hasattr(self.radius, "shape"): |
106 | 1 | tkerber | return self.radius # User has specified an array |
107 | 1 | tkerber | else:
|
108 | 1 | tkerber | return self.radius * ones(len(self.atoms), float) |
109 | 1 | tkerber | # No radii specified. Try getting them from the atoms.
|
110 | 1 | tkerber | try:
|
111 | 1 | tkerber | return self.atoms.get_atomic_radii() |
112 | 1 | tkerber | except AttributeError: |
113 | 1 | tkerber | try:
|
114 | 1 | tkerber | z = self._getatomicnumbers()
|
115 | 1 | tkerber | except AttributeError: |
116 | 1 | tkerber | pass
|
117 | 1 | tkerber | else:
|
118 | 1 | tkerber | return ase.data.covalent_radii[z]
|
119 | 1 | tkerber | # No radius available. Defaulting to 1.0
|
120 | 1 | tkerber | return ones(len(self.atoms), float) |
121 | 1 | tkerber | |
122 | 1 | tkerber | def _getatomicnumbers(self): |
123 | 1 | tkerber | return self.atoms.get_atomic_numbers() |
124 | 1 | tkerber | |
125 | 1 | tkerber | def _getcolors(self): |
126 | 1 | tkerber | # Try any explicitly given colors
|
127 | 1 | tkerber | if self.colors is not None: |
128 | 1 | tkerber | if type(self.colors) == type({}): |
129 | 1 | tkerber | self.log("Explicit colors dictionary") |
130 | 1 | tkerber | return _colorsfromdict(self.colors, |
131 | 1 | tkerber | asarray(self.atoms.get_tags(),int)) |
132 | 1 | tkerber | else:
|
133 | 1 | tkerber | self.log("Explicit colors") |
134 | 1 | tkerber | return self.colors |
135 | 1 | tkerber | # Try the color function, if given
|
136 | 1 | tkerber | if self.colorfunction is not None: |
137 | 1 | tkerber | self.log("Calling color function.") |
138 | 1 | tkerber | return self.colorfunction(self.atoms) |
139 | 1 | tkerber | # Maybe the atoms know their own colors
|
140 | 1 | tkerber | try:
|
141 | 1 | tkerber | c = self.atoms.get_colors()
|
142 | 1 | tkerber | except AttributeError: |
143 | 1 | tkerber | c = None
|
144 | 1 | tkerber | if c is not None: |
145 | 1 | tkerber | if type(c) == type({}): |
146 | 1 | tkerber | self.log("Color dictionary from atoms.get_colors()") |
147 | 1 | tkerber | return _colorsfromdict(c, asarray(self.atoms.get_tags(),int)) |
148 | 1 | tkerber | else:
|
149 | 1 | tkerber | self.log("Colors from atoms.get_colors()") |
150 | 1 | tkerber | return c
|
151 | 1 | tkerber | # Default to white atoms
|
152 | 1 | tkerber | self.log("No colors: using white") |
153 | 1 | tkerber | return ones(len(self.atoms), float) |
154 | 1 | tkerber | |
155 | 1 | tkerber | def _getinvisible(self): |
156 | 1 | tkerber | if self.invisible is not None: |
157 | 1 | tkerber | inv = self.invisible
|
158 | 1 | tkerber | else:
|
159 | 1 | tkerber | inv = zeros(len(self.atoms)) |
160 | 1 | tkerber | if self.invisibilityfunction: |
161 | 1 | tkerber | inv = logical_or(inv, self.invisibilityfunction(self.atoms)) |
162 | 1 | tkerber | r = self._getpositions()
|
163 | 1 | tkerber | if len(r) > len(inv): |
164 | 1 | tkerber | # This will happen in parallel simulations due to ghost atoms.
|
165 | 1 | tkerber | # They are invisible. Hmm, this may cause trouble.
|
166 | 1 | tkerber | i2 = ones(len(r))
|
167 | 1 | tkerber | i2[:len(inv)] = inv
|
168 | 1 | tkerber | inv = i2 |
169 | 1 | tkerber | del i2
|
170 | 1 | tkerber | if self.cut["xmin"] is not None: |
171 | 1 | tkerber | inv = logical_or(inv, less(r[:,0], self.cut["xmin"])) |
172 | 1 | tkerber | if self.cut["xmax"] is not None: |
173 | 1 | tkerber | inv = logical_or(inv, greater(r[:,0], self.cut["xmax"])) |
174 | 1 | tkerber | if self.cut["ymin"] is not None: |
175 | 1 | tkerber | inv = logical_or(inv, less(r[:,1], self.cut["ymin"])) |
176 | 1 | tkerber | if self.cut["ymax"] is not None: |
177 | 1 | tkerber | inv = logical_or(inv, greater(r[:,1], self.cut["ymax"])) |
178 | 1 | tkerber | if self.cut["zmin"] is not None: |
179 | 1 | tkerber | inv = logical_or(inv, less(r[:,2], self.cut["zmin"])) |
180 | 1 | tkerber | if self.cut["zmax"] is not None: |
181 | 1 | tkerber | inv = logical_or(inv, greater(r[:,2], self.cut["zmax"])) |
182 | 1 | tkerber | return inv
|
183 | 1 | tkerber | |
184 | 1 | tkerber | def __del__(self): |
185 | 1 | tkerber | if self.ownlogfile: |
186 | 1 | tkerber | self.logfile.close()
|
187 | 1 | tkerber | |
188 | 1 | tkerber | class PrimiPlotter(PrimiPlotterBase): |
189 | 1 | tkerber | """Primitive PostScript-based plots during a simulation.
|
190 | 1 | tkerber |
|
191 | 1 | tkerber | The PrimiPlotter plots atoms during simulations, extracting the
|
192 | 1 | tkerber | relevant information from the list of atoms. It is created using
|
193 | 1 | tkerber | the list of atoms as an argument to the constructor. Then one or
|
194 | 1 | tkerber | more output devices must be attached using set_output(device). The
|
195 | 1 | tkerber | list of supported output devices is at the end.
|
196 | 1 | tkerber |
|
197 | 1 | tkerber | The atoms are plotted as circles. The system is first rotated
|
198 | 1 | tkerber | using the angles specified by set_rotation([vx, vy, vz]). The
|
199 | 1 | tkerber | rotation is vx degrees around the x axis (positive from the y
|
200 | 1 | tkerber | toward the z axis), then vy degrees around the y axis (from x
|
201 | 1 | tkerber | toward z), then vz degrees around the z axis (from x toward y).
|
202 | 1 | tkerber | The rotation matrix is the same as the one used by RasMol.
|
203 | 1 | tkerber |
|
204 | 1 | tkerber | Per default, the system is scaled so it fits within the canvas
|
205 | 1 | tkerber | (autoscale mode). Autoscale mode is enabled and disables using
|
206 | 1 | tkerber | autoscale("on") or autoscale("off"). A manual scale factor can be
|
207 | 1 | tkerber | set with set_scale(scale), this implies autoscale("off"). The
|
208 | 1 | tkerber | scale factor (from the last autoscale event or from set_scale) can
|
209 | 1 | tkerber | be obtained with get_scale(). Finally, an explicit autoscaling can
|
210 | 1 | tkerber | be triggered with autoscale("now"), this is mainly useful before
|
211 | 1 | tkerber | calling get_scale or before disabling further autoscaling.
|
212 | 1 | tkerber | Finally, a relative scaling factor can be set with
|
213 | 1 | tkerber | SetRelativeScaling(), it is multiplied to the usual scale factor
|
214 | 1 | tkerber | (from autoscale or from set_scale). This is probably only useful in
|
215 | 1 | tkerber | connection with autoscaling.
|
216 | 1 | tkerber |
|
217 | 1 | tkerber | The radii of the atoms are obtained from the first of the following
|
218 | 1 | tkerber | methods which work:
|
219 | 1 | tkerber |
|
220 | 1 | tkerber | 1. If the radii are specified using PrimiPlotter.set_radii(r),
|
221 | 1 | tkerber | they are used. Must be an array, or a single number.
|
222 | 1 | tkerber |
|
223 | 1 | tkerber | 2. If the atoms has a get_atomic_radii() method, it is used. This is
|
224 | 1 | tkerber | unlikely.
|
225 | 1 | tkerber |
|
226 | 1 | tkerber | 3. If the atoms has a get_atomic_numbers() method, the
|
227 | 1 | tkerber | corresponding covalent radii are extracted from the
|
228 | 1 | tkerber | ASE.ChemicalElements module.
|
229 | 1 | tkerber |
|
230 | 1 | tkerber | 4. If all else fails, the radius is set to 1.0 Angstrom.
|
231 | 1 | tkerber |
|
232 | 1 | tkerber | The atoms are colored using the first of the following methods
|
233 | 1 | tkerber | which work.
|
234 | 1 | tkerber |
|
235 | 1 | tkerber | 1. If colors are explicitly set using PrimiPlotter.set_colors(),
|
236 | 1 | tkerber | they are used.
|
237 | 1 | tkerber |
|
238 | 1 | tkerber | 2. If these colors are specified as a dictionary, the tags
|
239 | 1 | tkerber | (from atoms.get_tags()) are used as an index into the
|
240 | 1 | tkerber | dictionary to get the actual colors of the atoms.
|
241 | 1 | tkerber |
|
242 | 1 | tkerber | 3. If a color function has been set using
|
243 | 1 | tkerber | PrimiPlotter.set_color_function(), it is called with the atoms
|
244 | 1 | tkerber | as an argument, and is expected to return an array of colors.
|
245 | 1 | tkerber |
|
246 | 1 | tkerber | 4. If the atoms have a get_colors() method, it is used to get the
|
247 | 1 | tkerber | colors.
|
248 | 1 | tkerber |
|
249 | 1 | tkerber | 5. If these colors are specified as a dictionary, the tags
|
250 | 1 | tkerber | (from atoms.get_tags()) are used as an index into the
|
251 | 1 | tkerber | dictionary to get the actual colors of the atoms.
|
252 | 1 | tkerber |
|
253 | 1 | tkerber | 6. If all else fails, the atoms will be white.
|
254 | 1 | tkerber |
|
255 | 1 | tkerber | The colors are specified as an array of colors, one color per
|
256 | 1 | tkerber | atom. Each color is either a real number from 0.0 to 1.0,
|
257 | 1 | tkerber | specifying a grayscale (0.0 = black, 1.0 = white), or an array of
|
258 | 1 | tkerber | three numbers from 0.0 to 1.0, specifying RGB values. The colors
|
259 | 1 | tkerber | of all atoms are thus a Numerical Python N-vector or a 3xN matrix.
|
260 | 1 | tkerber |
|
261 | 1 | tkerber | In cases 1a and 3a above, the keys of the dictionary are integers,
|
262 | 1 | tkerber | and the values are either numbers (grayscales) or 3-vectors (RGB
|
263 | 1 | tkerber | values), or strings with X11 color names, which are then
|
264 | 1 | tkerber | translated to RGB values. Only in case 1a and 3a are strings
|
265 | 1 | tkerber | recognized as colors.
|
266 | 1 | tkerber |
|
267 | 1 | tkerber | Some atoms may be invisible, and thus left out of the plot.
|
268 | 1 | tkerber | Invisible atoms are determined from the following algorithm.
|
269 | 1 | tkerber | Unlike the radius or the coloring, all points below are tried and
|
270 | 1 | tkerber | if an atom is invisible by any criterion, it is left out of the plot.
|
271 | 1 | tkerber |
|
272 | 1 | tkerber | 1. All atoms are visible.
|
273 | 1 | tkerber |
|
274 | 1 | tkerber | 2. If PrimiPlotter.set_invisible() has be used to specify invisible
|
275 | 1 | tkerber | atoms, any atoms for which the value is non-zero becomes invisible.
|
276 | 1 | tkerber |
|
277 | 1 | tkerber | 3. If an invisiblility function has been set with
|
278 | 1 | tkerber | PrimiPlotter.set_invisibility_function(), it is called with the
|
279 | 1 | tkerber | atoms as argument. It is expected to return an integer per
|
280 | 1 | tkerber | atom, any non-zero value makes that atom invisible.
|
281 | 1 | tkerber |
|
282 | 1 | tkerber | 4. If a cut has been specified using set_cut, any atom outside the
|
283 | 1 | tkerber | cut is made invisible.
|
284 | 1 | tkerber |
|
285 | 1 | tkerber | Note that invisible atoms are still included in the algorithm for
|
286 | 1 | tkerber | positioning and scaling the plot.
|
287 | 1 | tkerber |
|
288 | 1 | tkerber |
|
289 | 1 | tkerber | The following output devices are implemented.
|
290 | 1 | tkerber |
|
291 | 1 | tkerber | PostScriptFile(prefix): Create PS files names prefix0000.ps etc.
|
292 | 1 | tkerber |
|
293 | 1 | tkerber | PnmFile(prefix): Similar, but makes PNM files.
|
294 | 1 | tkerber |
|
295 | 1 | tkerber | GifFile(prefix): Similar, but makes GIF files.
|
296 | 1 | tkerber |
|
297 | 1 | tkerber | JpegFile(prefix): Similar, but makes JPEG files.
|
298 | 1 | tkerber |
|
299 | 1 | tkerber | X11Window(): Show the plot in an X11 window using ghostscript.
|
300 | 1 | tkerber |
|
301 | 1 | tkerber | Output devices writing to files take an extra optional argument to
|
302 | 1 | tkerber | the constructor, compress, specifying if the output file should be
|
303 | 1 | tkerber | gzipped. This is not allowed for some (already compressed) file
|
304 | 1 | tkerber | formats.
|
305 | 1 | tkerber |
|
306 | 1 | tkerber | Instead of a filename prefix, a filename containing a % can be
|
307 | 1 | tkerber | used. In that case the filename is expected to expand to a real
|
308 | 1 | tkerber | filename when used with the Python string formatting operator (%)
|
309 | 1 | tkerber | with the frame number as argument. Avoid generating spaces in the
|
310 | 1 | tkerber | file names: use e.g. %03d instead of %3d.
|
311 | 1 | tkerber | """
|
312 | 1 | tkerber | def __init__(self, atoms, verbose=0, timing=0, interval=1, initframe=0): |
313 | 1 | tkerber | """
|
314 | 1 | tkerber |
|
315 | 1 | tkerber | Parameters to the constructor:
|
316 | 1 | tkerber |
|
317 | 1 | tkerber | atoms: The atoms to be plottet.
|
318 | 1 | tkerber |
|
319 | 1 | tkerber | verbose = 0: Write progress information to stderr.
|
320 | 1 | tkerber |
|
321 | 1 | tkerber | timing = 0: Collect timing information.
|
322 | 1 | tkerber |
|
323 | 1 | tkerber | interval = 1: If specified, a plot is only made every
|
324 | 1 | tkerber | interval'th time update() is called. Deprecated, normally you
|
325 | 1 | tkerber | should use the interval argument when attaching the plotter to
|
326 | 1 | tkerber | e.g. the dynamics.
|
327 | 1 | tkerber |
|
328 | 1 | tkerber | initframe = 0: Initial frame number, i.e. the number of the
|
329 | 1 | tkerber | first plot.
|
330 | 1 | tkerber |
|
331 | 1 | tkerber | """
|
332 | 1 | tkerber | self.atoms = atoms
|
333 | 1 | tkerber | self.outputdevice = []
|
334 | 1 | tkerber | self.angles = zeros(3, float) |
335 | 1 | tkerber | self.dims = (512, 512) |
336 | 1 | tkerber | self.verbose = verbose
|
337 | 1 | tkerber | self.timing = timing
|
338 | 1 | tkerber | self.totaltime = 0.0 |
339 | 1 | tkerber | self.radius = None |
340 | 1 | tkerber | self.colors = None |
341 | 1 | tkerber | self.colorfunction = None |
342 | 1 | tkerber | self.n = initframe
|
343 | 1 | tkerber | self.interval = interval
|
344 | 1 | tkerber | self.skipnext = 0 # Number of calls to update before anything happens. |
345 | 1 | tkerber | self.a_scale = 1 |
346 | 1 | tkerber | self.relativescale = 1.0 |
347 | 1 | tkerber | self.invisible = None |
348 | 1 | tkerber | self.invisibilityfunction = None |
349 | 1 | tkerber | self.set_cut() # No cut |
350 | 1 | tkerber | self.isparallel = 0 |
351 | 1 | tkerber | self.logfile = None |
352 | 1 | tkerber | self.ownlogfile = False |
353 | 1 | tkerber | |
354 | 1 | tkerber | def set_output(self, device): |
355 | 1 | tkerber | self.outputdevice.append(device)
|
356 | 1 | tkerber | device.set_dimensions(self.dims)
|
357 | 1 | tkerber | device.set_owner(weakref.proxy(self))
|
358 | 1 | tkerber | |
359 | 1 | tkerber | def set_dimensions(self, dims): |
360 | 1 | tkerber | "Set the size of the canvas (a 2-tuple)."
|
361 | 1 | tkerber | if self.outputdevice: |
362 | 1 | tkerber | raise RuntimeError("Cannot set dimensions after an output device has been specified.") |
363 | 1 | tkerber | self.dims = dims
|
364 | 1 | tkerber | |
365 | 1 | tkerber | def autoscale(self, mode): |
366 | 1 | tkerber | if mode == "on": |
367 | 1 | tkerber | self.a_scale = 1 |
368 | 1 | tkerber | elif mode == "off": |
369 | 1 | tkerber | self.a_scale = 0 |
370 | 1 | tkerber | elif mode == "now": |
371 | 1 | tkerber | coords = self._rotate(self.atoms.get_positions()) |
372 | 1 | tkerber | radii = self._getradii()
|
373 | 1 | tkerber | self._autoscale(coords, radii)
|
374 | 1 | tkerber | else:
|
375 | 1 | tkerber | raise ValueError, "Unknown autoscale mode: ",+str(mode) |
376 | 1 | tkerber | |
377 | 1 | tkerber | def set_scale(self, scale): |
378 | 1 | tkerber | self.autoscale("off") |
379 | 1 | tkerber | self.scale = scale
|
380 | 1 | tkerber | |
381 | 1 | tkerber | def get_scale(self): |
382 | 1 | tkerber | return self.scale |
383 | 1 | tkerber | |
384 | 1 | tkerber | def set_relative_scale(self, rscale = 1.0): |
385 | 1 | tkerber | self.relativescale = rscale
|
386 | 1 | tkerber | |
387 | 1 | tkerber | def plot(self): |
388 | 1 | tkerber | """Create a plot now. Does not respect the interval timer.
|
389 | 1 | tkerber |
|
390 | 1 | tkerber | This method makes a plot unconditionally. It does not look at
|
391 | 1 | tkerber | the interval variable, nor is this plot taken into account in
|
392 | 1 | tkerber | the counting done by the update() method if an interval
|
393 | 1 | tkerber | variable was specified.
|
394 | 1 | tkerber | """
|
395 | 1 | tkerber | if self.timing: |
396 | 1 | tkerber | self._starttimer()
|
397 | 1 | tkerber | self.log("PrimiPlotter: Starting plot at " |
398 | 1 | tkerber | + time.strftime("%a, %d %b %Y %H:%M:%S"))
|
399 | 1 | tkerber | colors = self._getcolors()
|
400 | 1 | tkerber | invisible = self._getinvisible()
|
401 | 1 | tkerber | coords = self._rotate(self._getpositions()) |
402 | 1 | tkerber | radii = self._getradii()
|
403 | 1 | tkerber | if self.a_scale: |
404 | 1 | tkerber | self._autoscale(coords,radii)
|
405 | 1 | tkerber | scale = self.scale * self.relativescale |
406 | 1 | tkerber | coords = scale * coords |
407 | 1 | tkerber | center = self._getcenter(coords)
|
408 | 1 | tkerber | offset = array(self.dims + (0.0,))/2.0 - center |
409 | 1 | tkerber | coords = coords + offset |
410 | 1 | tkerber | self.log("Scale is %f and size is (%d, %d)" |
411 | 1 | tkerber | % (scale, self.dims[0], self.dims[1])) |
412 | 1 | tkerber | self.log("Physical size of plot is %f Angstrom times %f Angstrom" |
413 | 1 | tkerber | % (self.dims[0] / scale, self.dims[1] / scale)) |
414 | 1 | tkerber | |
415 | 1 | tkerber | self._verb("Sorting.") |
416 | 1 | tkerber | order = argsort(coords[:,2])
|
417 | 1 | tkerber | coords = coords[order] ### take(coords, order)
|
418 | 1 | tkerber | radii = radii[order] ### take(radii, order)
|
419 | 1 | tkerber | colors = colors[order] ### take(colors, order)
|
420 | 1 | tkerber | invisible = invisible[order] ### take(invisible, order)
|
421 | 1 | tkerber | if self.isparallel: |
422 | 1 | tkerber | id = arange(len(coords))[order] ### take(arange(len(coords)), order) |
423 | 1 | tkerber | else:
|
424 | 1 | tkerber | id = None
|
425 | 1 | tkerber | |
426 | 1 | tkerber | radii = radii * scale |
427 | 1 | tkerber | selector = self._computevisibility(coords, radii, invisible, id) |
428 | 1 | tkerber | coords = compress(selector, coords, 0)
|
429 | 1 | tkerber | radii = compress(selector, radii) |
430 | 1 | tkerber | colors = compress(selector, colors, 0)
|
431 | 1 | tkerber | self._makeoutput(scale, coords, radii, colors)
|
432 | 1 | tkerber | self.log("PrimiPlotter: Finished plotting at " |
433 | 1 | tkerber | + time.strftime("%a, %d %b %Y %H:%M:%S"))
|
434 | 1 | tkerber | self.log("\n\n") |
435 | 1 | tkerber | if self.timing: |
436 | 1 | tkerber | self._stoptimer()
|
437 | 1 | tkerber | |
438 | 1 | tkerber | def _computevisibility(self, coords, rad, invisible, id, zoom = 1): |
439 | 1 | tkerber | xy = coords[:,:2]
|
440 | 1 | tkerber | typradius = sum(rad) / len(rad) |
441 | 1 | tkerber | if typradius < 4.0: |
442 | 1 | tkerber | self.log("Refining visibility check.") |
443 | 1 | tkerber | if zoom >= 16: |
444 | 1 | tkerber | raise RuntimeError, "Cannot check visibility - too deep recursion." |
445 | 1 | tkerber | return self._computevisibility(xy*2, rad*2, invisible, id, zoom*2) |
446 | 1 | tkerber | else:
|
447 | 1 | tkerber | self.log("Visibility(r_typ = %.1f pixels)" % (typradius,)) |
448 | 1 | tkerber | dims = array(self.dims) * zoom
|
449 | 1 | tkerber | maxr = int(ceil(max(rad))) + 2 |
450 | 1 | tkerber | canvas = zeros((dims[0] + 4*maxr, dims[1] + 4*maxr), int8) |
451 | 1 | tkerber | # Atoms are only invisible if they are within the canvas, or closer
|
452 | 1 | tkerber | # to its edge than their radius
|
453 | 1 | tkerber | visible = (greater(xy[:,0], -rad) * less(xy[:,0], dims[0]+rad) |
454 | 1 | tkerber | * greater(xy[:,1], -rad) * less(xy[:,1], dims[1]+rad) |
455 | 1 | tkerber | * logical_not(invisible)) |
456 | 1 | tkerber | # Atoms are visible if not hidden behind other atoms
|
457 | 1 | tkerber | xy = floor(xy + 2*maxr + 0.5).astype(int) |
458 | 1 | tkerber | masks = {} |
459 | 1 | tkerber | for i in xrange(len(rad)-1, -1, -1): |
460 | 1 | tkerber | if (i % 100000) == 0 and i: |
461 | 1 | tkerber | self._verb(str(i)) |
462 | 1 | tkerber | if not visible[i]: |
463 | 1 | tkerber | continue
|
464 | 1 | tkerber | x, y = xy[i] |
465 | 1 | tkerber | r = rad[i] |
466 | 1 | tkerber | try:
|
467 | 1 | tkerber | mask, invmask, rn = masks[r] |
468 | 1 | tkerber | except KeyError: |
469 | 1 | tkerber | rn = int(ceil(r))
|
470 | 1 | tkerber | nmask = 2*rn+1 |
471 | 1 | tkerber | mask = (arange(nmask) - rn)**2
|
472 | 1 | tkerber | mask = less(mask[:,newaxis]+mask[newaxis,:], r*r).astype(int8) |
473 | 1 | tkerber | invmask = equal(mask, 0).astype(int8)
|
474 | 1 | tkerber | masks[r] = (mask, invmask, rn) |
475 | 1 | tkerber | window = logical_or(canvas[x-rn:x+rn+1, y-rn:y+rn+1], invmask) |
476 | 1 | tkerber | hidden = alltrue(window.flat) |
477 | 1 | tkerber | if hidden:
|
478 | 1 | tkerber | visible[i] = 0
|
479 | 1 | tkerber | else:
|
480 | 1 | tkerber | canvas[x-rn:x+rn+1, y-rn:y+rn+1] = logical_or(canvas[x-rn:x+rn+1, y-rn:y+rn+1], mask) |
481 | 1 | tkerber | self.log("%d visible, %d hidden out of %d" % |
482 | 1 | tkerber | (sum(visible), len(visible) - sum(visible), len(visible))) |
483 | 1 | tkerber | return visible
|
484 | 1 | tkerber | |
485 | 1 | tkerber | def _rotate(self, positions): |
486 | 1 | tkerber | self.log("Rotation angles: %f %f %f" % tuple(self.angles)) |
487 | 1 | tkerber | mat = dot(dot(_rot(self.angles[2], 2), |
488 | 1 | tkerber | _rot(self.angles[1], 1)), |
489 | 1 | tkerber | _rot(self.angles[0]+pi, 0)) |
490 | 1 | tkerber | return dot(positions, mat)
|
491 | 1 | tkerber | |
492 | 1 | tkerber | def _getcenter(self, coords): |
493 | 1 | tkerber | return array((max(coords[:,0]) + min(coords[:,0]), |
494 | 1 | tkerber | max(coords[:,1]) + min(coords[:,1]), 0.0)) / 2.0 |
495 | 1 | tkerber | |
496 | 1 | tkerber | def _autoscale(self, coords, radii): |
497 | 1 | tkerber | x = coords[:,0]
|
498 | 1 | tkerber | y = coords[:,1]
|
499 | 1 | tkerber | maxradius = max(radii)
|
500 | 1 | tkerber | deltax = max(x) - min(x) + 2*maxradius |
501 | 1 | tkerber | deltay = max(y) - min(y) + 2*maxradius |
502 | 1 | tkerber | scalex = self.dims[0] / deltax |
503 | 1 | tkerber | scaley = self.dims[1] / deltay |
504 | 1 | tkerber | self.scale = 0.95 * min(scalex, scaley) |
505 | 1 | tkerber | self.log("Autoscale: %f" % self.scale) |
506 | 1 | tkerber | |
507 | 1 | tkerber | def _makeoutput(self, scale, coords, radii, colors): |
508 | 1 | tkerber | for device in self.outputdevice: |
509 | 1 | tkerber | device.inform_about_scale(scale) |
510 | 1 | tkerber | device.plot(self.n, coords, radii, colors)
|
511 | 1 | tkerber | self.n = self.n + 1 |
512 | 1 | tkerber | |
513 | 1 | tkerber | |
514 | 1 | tkerber | class ParallelPrimiPlotter(PrimiPlotter): |
515 | 1 | tkerber | """A version of PrimiPlotter for parallel ASAP simulations.
|
516 | 1 | tkerber |
|
517 | 1 | tkerber | Used like PrimiPlotter, but only the output devices on the master
|
518 | 1 | tkerber | node are used. Most of the processing is distributed on the
|
519 | 1 | tkerber | nodes, but the actual output is only done on the master. See the
|
520 | 1 | tkerber | PrimiPlotter docstring for details.
|
521 | 1 | tkerber | """
|
522 | 1 | tkerber | def __init__(self, *args, **kwargs): |
523 | 1 | tkerber | apply(PrimiPlotter.__init__, (self,)+args, kwargs) |
524 | 1 | tkerber | self.isparallel = 1 |
525 | 1 | tkerber | import Scientific.MPI |
526 | 1 | tkerber | self.MPI = Scientific.MPI
|
527 | 1 | tkerber | self.mpi = Scientific.MPI.world
|
528 | 1 | tkerber | if self.mpi is None: |
529 | 1 | tkerber | raise RuntimeError, "MPI is not available." |
530 | 1 | tkerber | self.master = self.mpi.rank == 0 |
531 | 1 | tkerber | self.mpitag = 42 # Reduce chance of collision with other modules. |
532 | 1 | tkerber | |
533 | 1 | tkerber | def set_output(self, device): |
534 | 1 | tkerber | if self.master: |
535 | 1 | tkerber | PrimiPlotter.set_output(self, device)
|
536 | 1 | tkerber | |
537 | 1 | tkerber | def set_log(self, log): |
538 | 1 | tkerber | if self.master: |
539 | 1 | tkerber | PrimiPlotter.set_log(self, log)
|
540 | 1 | tkerber | |
541 | 1 | tkerber | def _getpositions(self): |
542 | 1 | tkerber | realpos = self.atoms.get_positions()
|
543 | 1 | tkerber | ghostpos = self.atoms.GetGhostCartesianPositions()
|
544 | 1 | tkerber | self.numberofrealatoms = len(realpos) |
545 | 1 | tkerber | self.numberofghostatoms = len(ghostpos) |
546 | 1 | tkerber | return concatenate((realpos, ghostpos))
|
547 | 1 | tkerber | |
548 | 1 | tkerber | def _getatomicnumbers(self): |
549 | 1 | tkerber | realz = self.atoms.get_atomic_numbers()
|
550 | 1 | tkerber | ghostz = self.atoms.GetGhostAtomicNumbers()
|
551 | 1 | tkerber | return concatenate((realz, ghostz))
|
552 | 1 | tkerber | |
553 | 1 | tkerber | def _getradius(self): |
554 | 1 | tkerber | r = PrimiPlotter._getradius(self)
|
555 | 1 | tkerber | if len(r) == self.numberofrealatoms + self.numberofghostatoms: |
556 | 1 | tkerber | # Must have calculated radii from atomic numbers
|
557 | 1 | tkerber | return r
|
558 | 1 | tkerber | else:
|
559 | 1 | tkerber | assert len(r) == self.numberofrealatoms |
560 | 1 | tkerber | # Heuristic: use minimum r for the ghosts
|
561 | 1 | tkerber | ghostr = min(r) * ones(self.numberofghostatoms, float) |
562 | 1 | tkerber | return concatenate((r, ghostr))
|
563 | 1 | tkerber | |
564 | 1 | tkerber | def _getcenter(self, coords): |
565 | 1 | tkerber | # max(x) and min(x) only works for rank-1 arrays in Numeric version 17.
|
566 | 1 | tkerber | maximal = maximum.reduce(coords[:,0:2]) |
567 | 1 | tkerber | minimal = minimum.reduce(coords[:,0:2]) |
568 | 1 | tkerber | recvmax = zeros(2, maximal.typecode())
|
569 | 1 | tkerber | recvmin = zeros(2, minimal.typecode())
|
570 | 1 | tkerber | self.mpi.allreduce(maximal, recvmax, self.MPI.max) |
571 | 1 | tkerber | self.mpi.allreduce(minimal, recvmin, self.MPI.min) |
572 | 1 | tkerber | maxx, maxy = recvmax |
573 | 1 | tkerber | minx, miny = recvmin |
574 | 1 | tkerber | return array([maxx + minx, maxy + miny, 0.0]) / 2.0 |
575 | 1 | tkerber | |
576 | 1 | tkerber | def _computevisibility(self, xy, rad, invisible, id, zoom = 1): |
577 | 1 | tkerber | # Find visible atoms, allowing ghost atoms to hide real atoms.
|
578 | 1 | tkerber | v = PrimiPlotter._computevisibility(self, xy, rad, invisible, id, zoom) |
579 | 1 | tkerber | # Then remove ghost atoms
|
580 | 1 | tkerber | return v * less(id, self.numberofrealatoms) |
581 | 1 | tkerber | |
582 | 1 | tkerber | def _autoscale(self, coords, radii): |
583 | 1 | tkerber | self._verb("Autoscale") |
584 | 1 | tkerber | n = len(self.atoms) |
585 | 1 | tkerber | x = coords[:n,0]
|
586 | 1 | tkerber | y = coords[:n,1]
|
587 | 1 | tkerber | assert len(x) == len(self.atoms) |
588 | 1 | tkerber | maximal = array([max(x), max(y), max(radii[:n])]) |
589 | 1 | tkerber | minimal = array([min(x), min(y)]) |
590 | 1 | tkerber | recvmax = zeros(3, maximal.typecode())
|
591 | 1 | tkerber | recvmin = zeros(2, minimal.typecode())
|
592 | 1 | tkerber | self.mpi.allreduce(maximal, recvmax, self.MPI.max) |
593 | 1 | tkerber | self.mpi.allreduce(minimal, recvmin, self.MPI.min) |
594 | 1 | tkerber | maxx, maxy, maxradius = recvmax |
595 | 1 | tkerber | minx, miny = recvmin |
596 | 1 | tkerber | deltax = maxx - minx + 2*maxradius
|
597 | 1 | tkerber | deltay = maxy - miny + 2*maxradius
|
598 | 1 | tkerber | scalex = self.dims[0] / deltax |
599 | 1 | tkerber | scaley = self.dims[1] / deltay |
600 | 1 | tkerber | self.scale = 0.95 * min(scalex, scaley) |
601 | 1 | tkerber | self.log("Autoscale: %f" % self.scale) |
602 | 1 | tkerber | |
603 | 1 | tkerber | def _getcolors(self): |
604 | 1 | tkerber | col = PrimiPlotter._getcolors(self)
|
605 | 1 | tkerber | nghost = len(self.atoms.GetGhostCartesianPositions()) |
606 | 1 | tkerber | newcolshape = (nghost + col.shape[0],) + col.shape[1:] |
607 | 1 | tkerber | newcol = zeros(newcolshape, col.typecode()) |
608 | 1 | tkerber | newcol[:len(col)] = col
|
609 | 1 | tkerber | return newcol
|
610 | 1 | tkerber | |
611 | 1 | tkerber | def _makeoutput(self, scale, coords, radii, colors): |
612 | 1 | tkerber | if len(colors.shape) == 1: |
613 | 1 | tkerber | # Greyscales
|
614 | 1 | tkerber | ncol = 1
|
615 | 1 | tkerber | else:
|
616 | 1 | tkerber | ncol = colors.shape[1] # 1 or 3. |
617 | 1 | tkerber | assert ncol == 3 # RGB values |
618 | 1 | tkerber | # If one processor says RGB, all must convert
|
619 | 1 | tkerber | ncolthis = array([ncol]) |
620 | 1 | tkerber | ncolmax = zeros((1,), ncolthis.typecode())
|
621 | 1 | tkerber | self.mpi.allreduce(ncolthis, ncolmax, self.MPI.max) |
622 | 1 | tkerber | ncolmax = ncolmax[0]
|
623 | 1 | tkerber | if ncolmax > ncol:
|
624 | 1 | tkerber | assert ncol == 1 |
625 | 1 | tkerber | colors = colors[:,newaxis] + zeros(ncolmax)[newaxis,:] |
626 | 1 | tkerber | ncol = ncolmax |
627 | 1 | tkerber | assert colors.shape == (len(coords), ncol) |
628 | 1 | tkerber | # Now send data from slaves to master
|
629 | 1 | tkerber | data = zeros((len(coords)+1, 4+ncol), float) |
630 | 1 | tkerber | data[:-1,:3] = coords |
631 | 1 | tkerber | data[:-1,3] = radii |
632 | 1 | tkerber | data[-1,-1] = 4+ncol # Used to communicate shape |
633 | 1 | tkerber | if ncol == 1: |
634 | 1 | tkerber | data[:-1,4] = colors |
635 | 1 | tkerber | else:
|
636 | 1 | tkerber | data[:-1,4:] = colors |
637 | 1 | tkerber | if not self.master: |
638 | 1 | tkerber | self.mpi.send(data, 0, self.mpitag) |
639 | 1 | tkerber | else:
|
640 | 1 | tkerber | total = [data[:-1]] # Last row is the dimensions. |
641 | 1 | tkerber | n = len(coords)
|
642 | 1 | tkerber | colsmin = colsmax = 4+ncol
|
643 | 1 | tkerber | for proc in range(1, self.mpi.size): |
644 | 1 | tkerber | self._verb("Receiving from processor "+str(proc)) |
645 | 1 | tkerber | fdat = self.mpi.receive(float, proc, self.mpitag)[0] |
646 | 1 | tkerber | fdat.shape = (-1, fdat[-1]) |
647 | 1 | tkerber | fdat = fdat[:-1] # Last row is the dimensions. |
648 | 1 | tkerber | total.append(fdat) |
649 | 1 | tkerber | n = n + len(fdat)
|
650 | 1 | tkerber | if fdat.shape[1] < colsmin: |
651 | 1 | tkerber | colsmin = fdat.shape[1]
|
652 | 1 | tkerber | if fdat.shape[1] > colsmax: |
653 | 1 | tkerber | colsmax = fdat.shape[1]
|
654 | 1 | tkerber | self._verb("Merging data") |
655 | 1 | tkerber | # Some processors may have only greyscales whereas others
|
656 | 1 | tkerber | # may have RGB. That will cause difficulties.
|
657 | 1 | tkerber | trouble = colsmax != colsmin |
658 | 1 | tkerber | data = zeros((n, colsmax), float)
|
659 | 1 | tkerber | if trouble:
|
660 | 1 | tkerber | assert data.shape[1] == 7 |
661 | 1 | tkerber | else:
|
662 | 1 | tkerber | assert data.shape[1] == 7 or data.shape[1] == 5 |
663 | 1 | tkerber | i = 0
|
664 | 1 | tkerber | for d in total: |
665 | 1 | tkerber | if not trouble or d.shape[1] == 7: |
666 | 1 | tkerber | data[i:i+len(d)] = d
|
667 | 1 | tkerber | else:
|
668 | 1 | tkerber | assert d.shape[1] == 5 |
669 | 1 | tkerber | data[i:i+len(d), :5] = d |
670 | 1 | tkerber | data[i:i+len(d), 5] = d[4] |
671 | 1 | tkerber | data[i:i+len(d), 6] = d[4] |
672 | 1 | tkerber | i = i + len(d)
|
673 | 1 | tkerber | assert i == len(data) |
674 | 1 | tkerber | # Now all data is on the master
|
675 | 1 | tkerber | self._verb("Sorting merged data") |
676 | 1 | tkerber | order = argsort(data[:,2])
|
677 | 1 | tkerber | data = data[order] ### take(data, order)
|
678 | 1 | tkerber | coords = data[:,:3]
|
679 | 1 | tkerber | radii = data[:,3]
|
680 | 1 | tkerber | if data.shape[1] == 5: |
681 | 1 | tkerber | colors = data[:,4]
|
682 | 1 | tkerber | else:
|
683 | 1 | tkerber | colors = data[:,4:]
|
684 | 1 | tkerber | PrimiPlotter._makeoutput(self, scale, coords, radii, colors)
|
685 | 1 | tkerber | |
686 | 1 | tkerber | class _PostScriptDevice: |
687 | 1 | tkerber | """PostScript based output device."""
|
688 | 1 | tkerber | offset = (0,0) # Will be changed by some classes |
689 | 1 | tkerber | def __init__(self): |
690 | 1 | tkerber | self.scale = 1 |
691 | 1 | tkerber | self.linewidth = 1 |
692 | 1 | tkerber | self.outline = 1 |
693 | 1 | tkerber | |
694 | 1 | tkerber | def set_dimensions(self, dims): |
695 | 1 | tkerber | self.dims = dims
|
696 | 1 | tkerber | |
697 | 1 | tkerber | def set_owner(self, owner): |
698 | 1 | tkerber | self.owner = owner
|
699 | 1 | tkerber | |
700 | 1 | tkerber | def inform_about_scale(self, scale): |
701 | 1 | tkerber | self.linewidth = 0.1 * scale |
702 | 1 | tkerber | |
703 | 1 | tkerber | def set_outline(self, value): |
704 | 1 | tkerber | self.outline = value
|
705 | 1 | tkerber | return self # Can chain these calls in set_output() |
706 | 1 | tkerber | |
707 | 1 | tkerber | def plot(self, *args, **kargs): |
708 | 1 | tkerber | self.Doplot(self.PSplot, *args, **kargs) |
709 | 1 | tkerber | |
710 | 1 | tkerber | def plotArray(self, *args, **kargs): |
711 | 1 | tkerber | self.Doplot(self.PSplotArray, *args, **kargs) |
712 | 1 | tkerber | |
713 | 1 | tkerber | def PSplot(self, file, n, coords, r, colors, noshowpage=0): |
714 | 1 | tkerber | xy = coords[:,:2]
|
715 | 1 | tkerber | assert(len(xy) == len(r) and len(xy) == len(colors)) |
716 | 1 | tkerber | if len(colors.shape) == 1: |
717 | 1 | tkerber | gray = 1
|
718 | 1 | tkerber | else:
|
719 | 1 | tkerber | gray = 0
|
720 | 1 | tkerber | assert(colors.shape[1] == 3) |
721 | 1 | tkerber | file.write("%!PS-Adobe-2.0\n") |
722 | 1 | tkerber | file.write("%%Creator: Primiplot\n") |
723 | 1 | tkerber | file.write("%%Pages: 1\n") |
724 | 1 | tkerber | file.write("%%%%BoundingBox: %d %d %d %d\n" % |
725 | 1 | tkerber | (self.offset + (self.offset[0] + self.dims[0], |
726 | 1 | tkerber | self.offset[1] + self.dims[1]))) |
727 | 1 | tkerber | file.write("%%EndComments\n") |
728 | 1 | tkerber | file.write("\n") |
729 | 1 | tkerber | file.write("% Enforce BoundingBox\n") |
730 | 1 | tkerber | file.write("%d %d moveto %d 0 rlineto 0 %d rlineto -%d 0 rlineto\n" % |
731 | 1 | tkerber | ((self.offset + self.dims + (self.dims[0],)))) |
732 | 1 | tkerber | file.write("closepath clip newpath\n\n") |
733 | 1 | tkerber | file.write("%f %f scale\n" % (2*(1.0/self.scale,))) |
734 | 1 | tkerber | file.write("%d %d translate\n" % (self.scale * self.offset[0], |
735 | 1 | tkerber | self.scale * self.offset[1])) |
736 | 1 | tkerber | file.write("\n") |
737 | 1 | tkerber | if gray:
|
738 | 1 | tkerber | if self.outline: |
739 | 1 | tkerber | file.write("/circ { 0 360 arc gsave setgray fill grestore stroke } def\n") |
740 | 1 | tkerber | else:
|
741 | 1 | tkerber | file.write("/circ { 0 360 arc setgray fill } def\n") |
742 | 1 | tkerber | else:
|
743 | 1 | tkerber | if self.outline: |
744 | 1 | tkerber | file.write("/circ { 0 360 arc gsave setrgbcolor fill grestore stroke } def\n") |
745 | 1 | tkerber | else:
|
746 | 1 | tkerber | file.write("/circ { 0 360 arc setrgbcolor fill } def\n") |
747 | 1 | tkerber | file.write("%f setlinewidth 0.0 setgray\n" % |
748 | 1 | tkerber | (self.linewidth * self.scale,)) |
749 | 1 | tkerber | |
750 | 1 | tkerber | if gray:
|
751 | 1 | tkerber | data = zeros((len(xy), 4), float) |
752 | 1 | tkerber | data[:,0] = colors
|
753 | 1 | tkerber | data[:,1:3] = (self.scale * xy) |
754 | 1 | tkerber | data[:,3] = (self.scale * r) |
755 | 1 | tkerber | for point in data: |
756 | 1 | tkerber | file.write("%.3f %.2f %.2f %.2f circ\n" % tuple(point)) |
757 | 1 | tkerber | else:
|
758 | 1 | tkerber | data = zeros((len(xy), 6), float) |
759 | 1 | tkerber | data[:,0:3] = colors |
760 | 1 | tkerber | data[:,3:5] = (self.scale * xy) |
761 | 1 | tkerber | data[:,5] = (self.scale * r) |
762 | 1 | tkerber | for point in data: |
763 | 1 | tkerber | file.write("%.3f %.3f %.3f %.2f %.2f %.2f circ\n" % tuple(point)) |
764 | 1 | tkerber | if not noshowpage: |
765 | 1 | tkerber | file.write("showpage\n") |
766 | 1 | tkerber | |
767 | 1 | tkerber | def PSplotArray(self, file, n, data, noshowpage=0): |
768 | 1 | tkerber | assert(len(data.shape) == 3) |
769 | 1 | tkerber | assert(data.shape[0] == self.dims[1] and data.shape[1] == self.dims[0]) |
770 | 1 | tkerber | data = clip((256*data).astype(int), 0, 255) |
771 | 1 | tkerber | file.write("%!PS-Adobe-2.0\n") |
772 | 1 | tkerber | file.write("%%Creator: Fieldplotter\n") |
773 | 1 | tkerber | file.write("%%Pages: 1\n") |
774 | 1 | tkerber | file.write("%%%%BoundingBox: %d %d %d %d\n" % |
775 | 1 | tkerber | (self.offset + (self.offset[0] + self.dims[0], |
776 | 1 | tkerber | self.offset[1] + self.dims[1]))) |
777 | 1 | tkerber | file.write("%%EndComments\n") |
778 | 1 | tkerber | file.write("\n") |
779 | 1 | tkerber | file.write("%d %d translate\n" % self.offset) |
780 | 1 | tkerber | file.write("%f %f scale\n" % self.dims) |
781 | 1 | tkerber | file.write("\n") |
782 | 1 | tkerber | file.write("% String holding a single line\n") |
783 | 1 | tkerber | file.write("/pictline %d string def\n" %(data.shape[1]*data.shape[2],)) |
784 | 1 | tkerber | file.write("\n") |
785 | 1 | tkerber | file.write("%d %d 8\n" % self.dims) |
786 | 1 | tkerber | file.write("[%d 0 0 %d 0 0]\n" % self.dims) |
787 | 1 | tkerber | file.write("{currentfile pictline readhexstring pop}\n") |
788 | 1 | tkerber | file.write("false %d colorimage\n" % (data.shape[2],)) |
789 | 1 | tkerber | file.write("\n") |
790 | 1 | tkerber | s = ""
|
791 | 1 | tkerber | for d in data.flat: |
792 | 1 | tkerber | s += ("%02X" % d)
|
793 | 1 | tkerber | if len(s) >= 72: |
794 | 1 | tkerber | file.write(s+"\n") |
795 | 1 | tkerber | s = ""
|
796 | 1 | tkerber | file.write(s+"\n") |
797 | 1 | tkerber | file.write("\n") |
798 | 1 | tkerber | if not noshowpage: |
799 | 1 | tkerber | file.write("showpage\n") |
800 | 1 | tkerber | |
801 | 1 | tkerber | class _PostScriptToFile(_PostScriptDevice): |
802 | 1 | tkerber | """Output device for PS files."""
|
803 | 1 | tkerber | compr_suffix = None
|
804 | 1 | tkerber | def __init__(self, prefix, compress = 0): |
805 | 1 | tkerber | self.compress = compress
|
806 | 1 | tkerber | if "'" in prefix: |
807 | 1 | tkerber | raise ValueError, "Filename may not contain a quote ('): "+prefix |
808 | 1 | tkerber | if "%" in prefix: |
809 | 1 | tkerber | # Assume the user knows what (s)he is doing
|
810 | 1 | tkerber | self.filenames = prefix
|
811 | 1 | tkerber | else:
|
812 | 1 | tkerber | self.filenames = prefix + "%04d" + self.suffix |
813 | 1 | tkerber | if compress:
|
814 | 1 | tkerber | if self.compr_suffix is None: |
815 | 1 | tkerber | raise RuntimeError, "Compression not supported." |
816 | 1 | tkerber | self.filenames = self.filenames + self.compr_suffix |
817 | 1 | tkerber | _PostScriptDevice.__init__(self)
|
818 | 1 | tkerber | |
819 | 1 | tkerber | class PostScriptFile(_PostScriptToFile): |
820 | 1 | tkerber | suffix = ".ps"
|
821 | 1 | tkerber | compr_suffix = ".gz"
|
822 | 1 | tkerber | offset = (50,50) |
823 | 1 | tkerber | # Inherits __init__
|
824 | 1 | tkerber | |
825 | 1 | tkerber | def Doplot(self, plotmethod, n, *args, **kargs): |
826 | 1 | tkerber | filename = self.filenames % (n,)
|
827 | 1 | tkerber | self.owner.log("Output to PostScript file "+filename) |
828 | 1 | tkerber | if self.compress: |
829 | 1 | tkerber | file = os.popen("gzip > '"+filename+"'", "w") |
830 | 1 | tkerber | else:
|
831 | 1 | tkerber | file = open(filename, "w") |
832 | 1 | tkerber | apply(plotmethod, (file, n)+args, kargs) |
833 | 1 | tkerber | file.close()
|
834 | 1 | tkerber | |
835 | 1 | tkerber | class _PS_via_PnmFile(_PostScriptToFile): |
836 | 1 | tkerber | gscmd = "gs -q -sDEVICE=pnmraw -sOutputFile=- -dDEVICEWIDTH=%d -dDEVICEHEIGHT=%d - "
|
837 | 1 | tkerber | # Inherits __init__
|
838 | 1 | tkerber | |
839 | 1 | tkerber | def Doplot(self, plotmethod, n, *args, **kargs): |
840 | 1 | tkerber | filename = self.filenames % (n,)
|
841 | 1 | tkerber | self.owner.log("Output to bitmapped file " + filename) |
842 | 1 | tkerber | cmd = self.gscmd + self.converter |
843 | 1 | tkerber | if self.compress: |
844 | 1 | tkerber | cmd = cmd + "| gzip "
|
845 | 1 | tkerber | |
846 | 1 | tkerber | cmd = (cmd+" > '%s'") % (self.dims[0], self.dims[1], filename) |
847 | 1 | tkerber | file = os.popen(cmd, "w")
|
848 | 1 | tkerber | apply(plotmethod, (file, n)+args, kargs) |
849 | 1 | tkerber | file.close()
|
850 | 1 | tkerber | |
851 | 1 | tkerber | class PnmFile(_PS_via_PnmFile): |
852 | 1 | tkerber | suffix = ".pnm"
|
853 | 1 | tkerber | compr_suffix = ".gz"
|
854 | 1 | tkerber | converter = ""
|
855 | 1 | tkerber | |
856 | 1 | tkerber | class GifFile(_PS_via_PnmFile): |
857 | 1 | tkerber | suffix = ".gif"
|
858 | 1 | tkerber | converter = "| ppmquant -floyd 256 2>/dev/null | ppmtogif 2>/dev/null"
|
859 | 1 | tkerber | |
860 | 1 | tkerber | class JpegFile(_PS_via_PnmFile): |
861 | 1 | tkerber | suffix = ".jpeg"
|
862 | 1 | tkerber | converter = "| ppmtojpeg --smooth=5"
|
863 | 1 | tkerber | |
864 | 1 | tkerber | class X11Window(_PostScriptDevice): |
865 | 1 | tkerber | """Shows the plot in an X11 window."""
|
866 | 1 | tkerber | #Inherits __init__
|
867 | 1 | tkerber | gscmd = "gs -q -sDEVICE=x11 -dDEVICEWIDTH=%d -dDEVICEHEIGHT=%d -r72x72 -"
|
868 | 1 | tkerber | def Doplot(self, plotmethod, n, *args, **kargs): |
869 | 1 | tkerber | self.owner.log("Output to X11 window") |
870 | 1 | tkerber | try:
|
871 | 1 | tkerber | file = self.pipe
|
872 | 1 | tkerber | self.pipe.write("showpage\n") |
873 | 1 | tkerber | except AttributeError: |
874 | 1 | tkerber | filename = self.gscmd % tuple(self.dims) |
875 | 1 | tkerber | file = os.popen(filename, "w")
|
876 | 1 | tkerber | self.pipe = file |
877 | 1 | tkerber | kargs["noshowpage"] = 1 |
878 | 1 | tkerber | apply(plotmethod, (file, n)+args, kargs) |
879 | 1 | tkerber | file.write("flushpage\n") |
880 | 1 | tkerber | file.flush()
|
881 | 1 | tkerber | |
882 | 1 | tkerber | # Helper functions
|
883 | 1 | tkerber | def _rot(v, axis): |
884 | 1 | tkerber | ax1, ax2 = ((1, 2), (0, 2), (0, 1))[axis] |
885 | 1 | tkerber | c, s = cos(v), sin(v) |
886 | 1 | tkerber | m = zeros((3,3), float) |
887 | 1 | tkerber | m[axis,axis] = 1.0
|
888 | 1 | tkerber | m[ax1,ax1] = c |
889 | 1 | tkerber | m[ax2,ax2] = c |
890 | 1 | tkerber | m[ax1,ax2] = s |
891 | 1 | tkerber | m[ax2,ax1] = -s |
892 | 1 | tkerber | return m
|
893 | 1 | tkerber | |
894 | 1 | tkerber | def _colorsfromdict(dict, cls): |
895 | 1 | tkerber | """Extract colors from dictionary using cls as key."""
|
896 | 1 | tkerber | assert(type(dict) == type({})) |
897 | 1 | tkerber | # Allow local modifications, to replace strings with rgb values.
|
898 | 1 | tkerber | dict = dict.copy()
|
899 | 1 | tkerber | isgray, isrgb = 0, 0 |
900 | 1 | tkerber | for k in dict.keys(): |
901 | 1 | tkerber | v = dict[k]
|
902 | 1 | tkerber | if type(v) == type("string"): |
903 | 1 | tkerber | v = color_table[v] |
904 | 1 | tkerber | dict[k] = v
|
905 | 1 | tkerber | try:
|
906 | 1 | tkerber | if len(v) == 3: |
907 | 1 | tkerber | isrgb = 1 # Assume it is an RGB value |
908 | 1 | tkerber | if not hasattr(v, "shape"): |
909 | 1 | tkerber | dict[k] = array(v) # Convert to array |
910 | 1 | tkerber | else:
|
911 | 1 | tkerber | raise RuntimeError, "Unrecognized color object "+repr(v) |
912 | 1 | tkerber | except TypeError: |
913 | 1 | tkerber | isgray = 1 # Assume it is a number |
914 | 1 | tkerber | if isgray and isrgb: |
915 | 1 | tkerber | # Convert all to RGB
|
916 | 1 | tkerber | for k in dict.keys(): |
917 | 1 | tkerber | v = dict[k]
|
918 | 1 | tkerber | if not hasattr(v, "shape"): |
919 | 1 | tkerber | dict[k] = v * ones(3, float) |
920 | 1 | tkerber | # Now the dictionary is ready
|
921 | 1 | tkerber | if isrgb:
|
922 | 1 | tkerber | colors = zeros((len(cls),3), float) |
923 | 1 | tkerber | else:
|
924 | 1 | tkerber | colors = zeros((len(cls),), float) |
925 | 1 | tkerber | for i in xrange(len(cls)): |
926 | 1 | tkerber | colors[i] = dict[cls[i]]
|
927 | 1 | tkerber | return colors
|