Statistiques
| Révision :

root / ase / visualize / primiplotter.py @ 1

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