Ticket #937: color.py

File color.py, 22.3 KB (added by Tristan Croll, 8 years ago)
Line 
1# vim: set expandtab shiftwidth=4 softtabstop=4:
2
3# === UCSF ChimeraX Copyright ===
4# Copyright 2016 Regents of the University of California.
5# All rights reserved. This software provided pursuant to a
6# license agreement containing restrictions on its disclosure,
7# duplication and use. For details see:
8# http://www.rbvi.ucsf.edu/chimerax/docs/licensing.html
9# This notice must be embedded in or attached to all copies,
10# including partial copies, of the software or any revisions
11# or derivations thereof.
12# === UCSF ChimeraX Copyright ===
13
14_SpecialColors = ["byatom", "byelement", "byhetero", "bychain", "bypolymer", "bymodel",
15 "byproperty", "fromatoms", "random"]
16
17_SequentialLevels = ["residues", "chains", "polymers", "structures"]
18# More possible sequential levels: "helix", "helices", "strands", "SSEs", "volmodels", "allmodels"
19
20def color(session, objects, color=None, byproperty=None, what=None,
21 target=None, transparency=None,
22 sequential=None, palette=None, halfbond=None,
23 map=None, range=None, offset=0, zone=None, distance=2,
24 undo_name="color"):
25 """Color atoms, ribbons, surfaces, ....
26
27 Parameters
28 ----------
29 objects : Objects
30 Which objects to color.
31 color : Color
32 Color can be a standard color name or "byatom", "byelement", "byhetero",
33 "bychain", "bypolymer", "bymodel"
34 byproperty : one of "bfactors", "occupancies"
35 what : 'atoms', 'cartoons', 'ribbons', 'surfaces', 'bonds', 'pseudobonds' or None
36 What to color. Everything is colored if option is not specified.
37 target : string
38 Alternative to the "what" option for specifying what to color.
39 Characters indicating what to color, a = atoms, c = cartoon, r = cartoon, s = surfaces,
40 l = labels, b = bonds, p = pseudobonds, d = distances.
41 Everything is colored if no target is specified.
42 transparency : float
43 Percent transparency to use. If not specified current transparency is preserved.
44 sequential : "residues", "chains", "polymers", "structures"
45 Assigns each object a color from a color map.
46 palette : :class:`.Colormap`
47 Color map to use with sequential coloring.
48 halfbond : bool
49 Whether to color each half of a bond to match the connected atoms.
50 If halfbond is false the bond is given the single color assigned to the bond.
51 map : Volume
52 Color specified surfaces by sampling from this density map using palette, range, and offset options.
53 range : 2 comma-separated floats or "full"
54 Specifies the range of map values used for sampling from a palette.
55 offset : float
56 Displacement distance along surface normals for sampling map when using map option. Default 0.
57 zone : Atoms
58 Color surfaces to match closest atom within specified zone distance.
59 distance : float
60 Zone distance used with zone option.
61 """
62 if objects is None:
63 from . import all_objects
64 objects = all_objects(session)
65 atoms = objects.atoms
66 if byproperty is not None and color is not None:
67 from ..errors import UserError
68 raise UserError('Cannot set a single color and color by property. '
69 +'If you want to set a custom color scale, use '
70 +'"palette" instead.')
71
72 if color == "byhetero":
73 atoms = atoms.filter(atoms.element_numbers != 6)
74
75 default_target = (target is None and what is None)
76 if default_target:
77 target = 'acslbpd'
78 if target and 'r' in target:
79 target += 'c'
80
81 if what is not None:
82 what_target = {'atoms':'a', 'cartoons':'c', 'ribbons':'c',
83 'surfaces':'s', 'bonds':'b', 'pseudobonds':'p'}
84 if target is None:
85 target = ''
86 for w in what:
87 target += what_target[w]
88
89 from ..undo import UndoState
90 undo_state = UndoState(undo_name)
91
92 # Decide whether to set or preserve transparency
93 opacity = None
94 if transparency is not None:
95 opacity = min(255, max(0, int(2.56 * (100 - transparency))))
96 if getattr(color, 'explicit_transparency', False):
97 opacity = color.uint8x4()[3]
98
99 if halfbond is not None:
100 bonds = objects.bonds
101 if len(bonds) > 0:
102 undo_state.add(bonds, "halfbonds", bonds.halfbonds, halfbond)
103 bonds.halfbonds = halfbond
104 if 'p' in target:
105 pbonds = objects.pseudobonds
106 if len(pbonds) > 0:
107 undo_state.add(pbonds, "halfbonds", pbonds.halfbonds, halfbond)
108 pbonds.halfbonds = halfbond
109
110 if sequential is not None:
111 try:
112 f = _SequentialColor[sequential]
113 except KeyError:
114 from ..errors import UserError
115 raise UserError("sequential \"%s\" not implemented yet"
116 % sequential)
117 else:
118 f(session, objects, palette, opacity, target, undo_state)
119 session.undo.register(undo_state)
120 return
121
122 if zone is not None:
123 from ..atomic import MolecularSurface, Structure
124 slist = [m for m in objects.models
125 if not m.empty_drawing() and not isinstance(m, (Structure, MolecularSurface))]
126 for m in objects.models:
127 if hasattr(m, 'surface_drawings_for_vertex_coloring'):
128 slist.extend(m.surface_drawings_for_vertex_coloring())
129 bonds = None
130 auto_update = False
131 from ..surface.colorzone import points_and_colors, color_zone
132 for s in slist:
133 points, colors = points_and_colors(zone, bonds)
134 # TODO: save undo data
135 s.scene_position.inverse().move(points) # Transform points to surface coordinates
136 color_zone(s, points, colors, distance, auto_update)
137
138 what = []
139
140 bgcolor = session.main_view.background_color
141
142 if byproperty is not None:
143 if atoms is not None:
144 from ..colors import BuiltinColormaps
145 if palette is None:
146 palette = 'blue-white-red'
147 cmap = BuiltinColormaps[palette]
148 if range is None or range == 'full':
149 prop_vals = getattr(atoms, byproperty)
150 cm = cmap.linear_range(prop_vals.min(), prop_vals.max())
151 else:
152 cm = cmap.linear_range(*range)
153 property_color_map = cm
154
155
156 if 'a' in target:
157 # atoms/bonds
158 if atoms is not None:
159 if color is not None:
160 _set_atom_colors(atoms, color, opacity, bgcolor, undo_state)
161 elif byproperty is not None:
162 _set_atom_colors_by_property(atoms, byproperty, opacity, property_color_map, undo_state)
163 what.append('%d atoms' % len(atoms))
164
165 if 'l' in target:
166 if not default_target:
167 session.logger.warning('Label colors not supported yet')
168
169 if 's' in target and (color is not None or map is not None):
170 # TODO: save undo data
171 from ..atomic import MolecularSurface, concatenate, Structure, PseudobondGroup
172 msatoms = [m.atoms for m in objects.models
173 if isinstance(m, MolecularSurface) and not m.atoms.intersects(atoms)]
174 satoms = concatenate(msatoms + [atoms]) if msatoms else atoms
175 if color == "byhetero":
176 satoms = satoms.filter(satoms.element_numbers != 6)
177 ns = _set_surface_colors(session, satoms, color, opacity, bgcolor,
178 map, palette, range, offset, undo_state=undo_state)
179 # Handle non-molecular surfaces like density maps
180 if color not in _SpecialColors:
181 mlist = [m for m in objects.models if not isinstance(m, (Structure, MolecularSurface, PseudobondGroup))]
182 for m in mlist:
183 _set_model_colors(session, m, color, map, opacity, palette, range, offset)
184 ns += len(mlist)
185 what.append('%d surfaces' % ns)
186
187 if 'c' in target:
188 residues = atoms.unique_residues
189 if color is not None:
190 _set_ribbon_colors(residues, color, opacity, bgcolor, undo_state)
191 elif byproperty is not None:
192 _set_ribbon_colors_by_property(residues, byproperty, opacity, property_color_map, undo_state)
193 what.append('%d residues' % len(residues))
194
195 if 'b' in target:
196 if color not in _SpecialColors and color is not None:
197 bonds = objects.bonds
198 if len(bonds) > 0:
199 if color not in _SpecialColors:
200 color_array = color.uint8x4()
201 undo_state.add(bonds, "colors", bonds.colors, color_array)
202 bonds.colors = color_array
203 what.append('%d bonds' % len(bonds))
204
205 if 'p' in target:
206 if color not in _SpecialColors and color is not None:
207 pbonds = objects.pseudobonds
208 if len(pbonds) > 0:
209 color_array = color.uint8x4()
210 undo_state.add(pbonds, "colors", pbonds.colors, color_array)
211 pbonds.colors = color_array
212 what.append('%d pseudobonds' % len(pbonds))
213
214 if 'd' in target:
215 if not default_target:
216 session.logger.warning('Distances colors not supported yet')
217
218 if not what:
219 what.append('nothing')
220
221 from . import cli
222 session.logger.status('Colored %s' % cli.commas(what, ' and'))
223 session.undo.register(undo_state)
224
225
226def _computed_atom_colors(atoms, color, opacity, bgcolor):
227 if color in ("byatom", "byelement", "byhetero"):
228 c = _element_colors(atoms, opacity)
229 elif color == "bychain":
230 from ..atomic.colors import chain_colors
231 c = chain_colors(atoms.residues.chain_ids)
232 c[:, 3] = atoms.colors[:, 3] if opacity is None else opacity
233 elif color == "bypolymer":
234 from ..atomic.colors import polymer_colors
235 c = atoms.colors.copy()
236 sc,amask = polymer_colors(atoms.residues)
237 c[amask,:] = sc[amask,:]
238 c[amask, 3] = atoms.colors[amask, 3] if opacity is None else opacity
239 elif color == "bymodel":
240 c = atoms.colors.copy()
241 for m, matoms in atoms.by_structure:
242 color = m.initial_color(bgcolor).uint8x4()
243 mi = atoms.mask(matoms)
244 c[mi, :3] = color[:3]
245 if opacity is not None:
246 c[mi, 3] = opacity
247 elif color == "random":
248 from numpy import random, uint8
249 c = random.randint(0, 255, (len(atoms), 4)).astype(uint8)
250 c[:, 3] = 255 # Opaque
251 else:
252 # Other "colors" do not apply to atoms
253 c = None
254 return c
255
256
257def _element_colors(atoms, opacity=None):
258 from ..atomic.colors import element_colors
259 c = element_colors(atoms.element_numbers)
260 c[:, 3] = atoms.colors[:, 3] if opacity is None else opacity
261 return c
262
263
264def _set_atom_colors(atoms, color, opacity, bgcolor, undo_state):
265 if color in _SpecialColors:
266 c = _computed_atom_colors(atoms, color, opacity, bgcolor)
267 if c is not None:
268 undo_state.add(atoms, "colors", atoms.colors, c)
269 atoms.colors = c
270 else:
271 c = atoms.colors
272 c[:, :3] = color.uint8x4()[:3] # Preserve transparency
273 if opacity is not None:
274 c[:, 3] = opacity
275 undo_state.add(atoms, "colors", atoms.colors, c)
276 atoms.colors = c
277
278def _set_atom_colors_by_property(atoms, prop, opacity, cmap, undo_state):
279 prop_vals = getattr(atoms, prop)
280 c = atoms.colors
281 c[:,:3] = cmap.interpolated_rgba8(prop_vals)[:,:3]
282 if opacity is not None:
283 c[:,3] = opacity
284 undo_state.add(atoms, "colors", atoms.colors, c)
285 atoms.colors = c
286
287
288
289
290def _set_ribbon_colors(residues, color, opacity, bgcolor, undo_state):
291 if color not in _SpecialColors:
292 c = residues.ribbon_colors
293 c[:, :3] = color.uint8x4()[:3] # Preserve transparency
294 if opacity is not None:
295 c[:, 3] = opacity
296 undo_state.add(residues, "ribbon_colors", residues.ribbon_colors, c)
297 residues.ribbon_colors = c
298 elif color == 'bychain':
299 from ..atomic.colors import chain_colors
300 c = chain_colors(residues.chain_ids)
301 c[:, 3] = residues.ribbon_colors[:, 3] if opacity is None else opacity
302 undo_state.add(residues, "ribbon_colors", residues.ribbon_colors, c)
303 residues.ribbon_colors = c
304 elif color == "bypolymer":
305 from ..atomic.colors import polymer_colors
306 c,rmask = polymer_colors(residues)
307 c[rmask, 3] = residues.ribbon_colors[rmask, 3] if opacity is None else opacity
308 masked_residues = residues.filter(rmask)
309 undo_state.add(masked_residues, "ribbon_colors", masked_residues.ribbon_colors, c[rmask,:])
310 masked_residues.ribbon_colors = c[rmask,:]
311 elif color == 'bymodel':
312 for m, res in residues.by_structure:
313 c = res.ribbon_colors
314 c[:, :3] = m.initial_color(bgcolor).uint8x4()[:3]
315 if opacity is not None:
316 c[:, 3] = opacity
317 undo_state.add(res, "ribbon_colors", res.ribbon_colors, c)
318 res.ribbon_colors = c
319 elif color == 'random':
320 from numpy import random, uint8
321 c = random.randint(0, 255, (len(residues), 4)).astype(uint8)
322 c[:, 3] = 255 # No transparency
323 undo_state.add(residues, "ribbon_colors", residues.ribbon_colors, c)
324 residues.ribbon_colors = c
325
326def _set_ribbon_colors_by_property(residues, prop, opacity, cmap, undo_state):
327 c = residues.ribbon_colors
328 import numpy
329 vals = numpy.array([getattr(r.atoms, prop).mean() for r in residues])
330 c[:,:3] = cmap.interpolated_rgba8(vals)[:,:3]
331 if opacity is not None:
332 c[:,3] = opacity
333 undo_state.add(residues, "ribbon_colors", residues.ribbon_colors, c)
334 residues.ribbon_colors = c
335
336
337def _set_surface_colors(session, atoms, color, opacity, bgcolor=None,
338 map=None, palette=None, range=None, offset=0, undo_state=None):
339 # TODO: save undo data
340 from .scolor import color_surfaces_at_atoms, color_surfaces_by_map_value
341 if color in _SpecialColors:
342 if color == 'fromatoms':
343 ns = color_surfaces_at_atoms(atoms, opacity=opacity)
344 else:
345 # Surface colored different from atoms
346 c = _computed_atom_colors(atoms, color, opacity, bgcolor)
347 ns = color_surfaces_at_atoms(atoms, opacity=opacity, per_atom_colors=c)
348
349 elif map:
350 ns = color_surfaces_by_map_value(atoms, opacity=opacity, map=map, palette=palette,
351 range=range, offset=offset)
352 else:
353 ns = color_surfaces_at_atoms(atoms, color, opacity=opacity)
354 return ns
355
356def _set_model_colors(session, m, color, map, opacity, palette, range, offset):
357 if map is None:
358 c = color.uint8x4()
359 if not opacity is None:
360 c[3] = opacity
361 elif not m.single_color is None:
362 c[3] = m.single_color[3]
363 m.single_color = c
364 else:
365 if hasattr(m, 'surface_drawings_for_vertex_coloring'):
366 surfs = m.surface_drawings_for_vertex_coloring()
367 elif not m.empty_drawing():
368 surfs = [m]
369 else:
370 surfs = []
371 from .scolor import color_surface_by_map_value
372 for s in surfs:
373 color_surface_by_map_value(s, map, palette=palette, range=range,
374 offset=offset, opacity=opacity)
375
376# -----------------------------------------------------------------------------
377# Chain ids in each structure are colored from color map ordered alphabetically.
378#
379def _set_sequential_chain(session, selected, cmap, opacity, target, undo_state):
380 # Organize selected atoms by structure and then chain
381 uc = selected.atoms.residues.chains.unique()
382 chain_atoms = {}
383 for c in uc:
384 chain_atoms.setdefault(c.structure, []).append((c.chain_id, c.existing_residues.atoms))
385 # Make sure there is a colormap
386 if cmap is None:
387 from .. import colors
388 cmap = colors.BuiltinColormaps["rainbow"]
389 # Each structure is colored separately with cmap applied by chain
390 import numpy
391 from ..colors import Color
392 for sl in chain_atoms.values():
393 sl.sort(key = lambda ca: ca[0]) # Sort by chain id
394 colors = cmap.interpolated_rgba(numpy.linspace(0.0, 1.0, len(sl)))
395 for color, (chain_id, atoms) in zip(colors, sl):
396 c = Color(color)
397 if target is None or 'a' in target:
398 _set_atom_colors(atoms, c, opacity, None, undo_state)
399 if target is None or 'c' in target:
400 res = atoms.unique_residues
401 _set_ribbon_colors(res, c, opacity, None, undo_state)
402 if target is None or 's' in target:
403 _set_surface_colors(session, atoms, c, opacity, undo_state=undo_state)
404
405# ----------------------------------------------------------------------------------
406# Polymers (unique sequences) in each structure are colored from color map ordered
407# by polymer length.
408#
409def _set_sequential_polymer(session, objects, cmap, opacity, target, undo_state):
410 # Organize atoms by structure and then polymer sequence
411 uc = objects.atoms.residues.chains.unique()
412 seq_atoms = {}
413 for c in uc:
414 seq_atoms.setdefault(c.structure, {}).setdefault(c.characters, []).append(c.existing_residues.atoms)
415 # Make sure there is a colormap
416 if cmap is None:
417 from .. import colors
418 cmap = colors.BuiltinColormaps["rainbow"]
419 # Each structure is colored separately with cmap applied by chain
420 import numpy
421 from ..colors import Color
422 for sl in seq_atoms.values():
423 sseq = list(sl.items())
424 sseq.sort(key = lambda sa: len(sa[0])) # Sort by sequence length
425 colors = cmap.interpolated_rgba(numpy.linspace(0.0, 1.0, len(sseq)))
426 for color, (seq, alist) in zip(colors, sseq):
427 c = Color(color)
428 for atoms in alist:
429 if target is None or 'a' in target:
430 _set_atom_colors(atoms, c, opacity, None, undo_state)
431 if target is None or 'c' in target:
432 res = atoms.unique_residues
433 _set_ribbon_colors(res, c, opacity, None, undo_state)
434 if target is None or 's' in target:
435 _set_surface_colors(session, atoms, c, opacity, undo_state=undo_state)
436
437# -----------------------------------------------------------------------------
438#
439def _set_sequential_residue(session, selected, cmap, opacity, target, undo_state):
440 # Make sure there is a colormap
441 if cmap is None:
442 from .. import colors
443 cmap = colors.BuiltinColormaps["rainbow"]
444 # Get chains and atoms in chains with "by_chain"
445 # Each chain is colored separately with cmap applied by residue
446 import numpy
447 from ..colors import Color
448 structure_chain_ids = {}
449 for structure, chain_id, atoms in selected.atoms.by_chain:
450 try:
451 cids = structure_chain_ids[structure]
452 except KeyError:
453 structure_chain_ids[structure] = cids = set()
454 cids.add(chain_id)
455 for structure, cids in structure_chain_ids.items():
456 for chain in structure.chains:
457 if chain.chain_id not in cids:
458 continue
459 residues = chain.existing_residues
460 colors = cmap.interpolated_rgba(numpy.linspace(0.0, 1.0, len(residues)))
461 for color, r in zip(colors, residues):
462 c = Color(color)
463 if target is None or 'a' in target:
464 _set_atom_colors(r.atoms, c, opacity, None, undo_state)
465 if target is None or 'c' in target:
466 rgba = c.uint8x4()
467 if opacity is not None:
468 rgba[3] = opacity
469 undo_state.add(r, "ribbon_color", r.ribbon_color, rgba)
470 r.ribbon_color = rgba
471
472# -----------------------------------------------------------------------------
473#
474def _set_sequential_structures(session, selected, cmap, opacity, target, undo_state):
475 # Make sure there is a colormap
476 if cmap is None:
477 from .. import colors
478 cmap = colors.BuiltinColormaps["rainbow"]
479
480 from ..atomic import Structure
481 models = list(m for m in selected.models if isinstance(m, Structure))
482 models.sort(key = lambda m: m.id)
483 if len(models) == 0:
484 return
485
486 # Each structure is colored separately with cmap applied by chain
487 import numpy
488 from ..colors import Color
489 colors = cmap.interpolated_rgba(numpy.linspace(0.0, 1.0, len(models)))
490 for color, m in zip(colors, models):
491 c = Color(color)
492 if 'a' in target:
493 _set_atom_colors(m.atoms, c, opacity, None, undo_state)
494 if 'c' in target:
495 _set_ribbon_colors(m.residues, c, opacity, None, undo_state)
496 if 's' in target:
497 # TODO: save surface undo data
498 from .scolor import color_surfaces_at_atoms
499 color_surfaces_at_atoms(m.atoms, c)
500
501# -----------------------------------------------------------------------------
502#
503_SequentialColor = {
504 "polymers": _set_sequential_polymer,
505 "chains": _set_sequential_chain,
506 "residues": _set_sequential_residue,
507 "structures": _set_sequential_structures,
508}
509
510# -----------------------------------------------------------------------------
511#
512def register_command(session):
513 from . import register, CmdDesc, ColorArg, ColormapArg, ColormapRangeArg, ObjectsArg, create_alias
514 from . import EmptyArg, Or, EnumOf, StringArg, ListOf, FloatArg, BoolArg, AtomsArg
515 from ..map import MapArg
516 what_arg = ListOf(EnumOf(('atoms', 'cartoons', 'ribbons', 'surfaces', 'bonds', 'pseudobonds')))
517 property_arg = EnumOf(('bfactors', 'occupancies'))
518 desc = CmdDesc(required=[('objects', Or(ObjectsArg, EmptyArg))],
519 optional=[('color', Or(ColorArg, EnumOf(_SpecialColors))),
520 ('what', what_arg)],
521 keyword=[('target', StringArg),
522 ('transparency', FloatArg),
523 ('sequential', EnumOf(_SequentialLevels)),
524 ('halfbond', BoolArg),
525 ('map', MapArg),
526 ('palette', ColormapArg),
527 ('range', ColormapRangeArg),
528 ('offset', FloatArg),
529 ('zone', AtomsArg),
530 ('distance', FloatArg),
531 ('byproperty', property_arg),
532 ],
533 synopsis="color objects")
534 register('color', desc, color, logger=session.logger)
535 create_alias('colour', 'color $*', logger=session.logger)