root / ase / visualize / primiplotter.py @ 4
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
|