[Chimera-users] Exporting VRML of MSMS surfaces

[Thomas Goddard]

The VRMLoutput/__init__.py I posted in my previous two emails worked
for the newly added MSMS surface output, but was broken for volume contour
surfaces.  Here is the corrected code.

   Tom


----
VRMLoutput/__init__.py follows:

# -----------------------------------------------------------------------------
# Output displayed _surface models as VRML.
#

# -----------------------------------------------------------------------------
#
def surfaces_as_vrml(surface_models):

    surf_vrml = ['#VRML V2.0 utf8\n',
                 'Transform {',
                 '  children [']
    import _surface
    import chimera
    for m in surface_models:
    	if m.display:
            if isinstance(m, _surface.Surface_Model):
                glist = filter(lambda g: g.display, m.surface_groups())
                for g in glist:
                    vrml = surface_group_as_vrml(g)
                    surf_vrml.append(vrml)
            elif isinstance(m, chimera.MSMSModel):
                vrml = msms_surface_as_vrml(m)
                surf_vrml.append(vrml)
                
    surf_vrml.extend(['  ]', '}'])
    vrml = '\n'.join(surf_vrml)

    return vrml

# -----------------------------------------------------------------------------
#
def surface_group_as_vrml(g):

    vertex_rgba = g.vertex_colors()
    solid_rgba = g.color()
    m = g.model()
    mat = m.material
    shine = mat.shininess / 128.0
    srgb = mat.specular
    if g.display_style() == g.Solid:
        v, vi = g.geometry()
        n = g.normals()
        ts = g.two_sided_lighting()
        vrml = surface_vrml(v, vi, n, vertex_rgba, solid_rgba,
                            ts, shine, srgb)
    else:
        v, vi = g.edges()
        # Mesh lighting not supported in VRML97.
        vrml = mesh_vrml(v, vi, vertex_rgba, solid_rgba)
    return vrml

# -----------------------------------------------------------------------------
#
def msms_surface_as_vrml(m):

    vertex_rgba, solid_rgba, mat, shine, srgb = msms_colors(m)

    # TODO: Use m.drawMode (Filled, Mesh, Dot) to control vrml style
    vfloat, vint, tri = m.triangleData()
    v = vfloat[:,:3]
    n = vfloat[:,3:6]
    vi = tri[:,:3]
    ts = True
    vrml = surface_vrml(v, vi, n, vertex_rgba, solid_rgba,
                        ts, shine, srgb)
    return vrml

# -----------------------------------------------------------------------------
#
def msms_colors(m):

    if m.colorMode == m.Custom:
        vertex_rgba = colors_to_rgba(m.customColors)
    elif m.colorMode == m.ByAtom:
        vfloat, vint, tri = m.triangleData()
        alist = m.atomMap()
        vatoms = map(lambda i: alist[i-1], vint[:,1])
        vertex_rgba = colors_to_rgba(map(atom_color, vatoms))
    else:
        vertex_rgba = None

    mol = m.molecule
    if mol:
        c = mol.surfaceColor
        if c == None:
            c = mol.color
        solid_rgba = c.rgba()
        mat = c.material
    else:
        solid_rgba = (1,1,1,1)
        import chimera
        mat = chimera.MaterialColor(0,0,0).material
    shine = mat.shininess / 128.0
    srgb = mat.specular

    return vertex_rgba, solid_rgba, mat, shine, srgb

# -----------------------------------------------------------------------------
#
def colors_to_rgba(color_list):

    rgba = map(lambda c: c.rgba(), color_list)
    import Numeric
    rgba_f32 = Numeric.array(rgba, Numeric.Float32)
    return rgba_f32

# -----------------------------------------------------------------------------
#
def atom_color(atom):

    c = atom.surfaceColor
    if c == None:
        c = atom.color
        if c == None:
            c = atom.molecule.surfaceColor
            if c == None:
                c = atom.molecule.color
    return c

# -----------------------------------------------------------------------------
#
def surface_vrml(vertices, vertex_indices, vertex_normals, vertex_rgba,
                 solid_rgba, two_sided_lighting, shininess, specular_rgb):

    vrml_template = \
'''Shape { 
    appearance Appearance {
      material Material {
        diffuseColor %s
        transparency %.3f
        shininess %.3f
        specularColor %s
      }
    }
    geometry IndexedFaceSet {
        coord Coordinate { 
            point [
%s
            ]
        }
        normal Normal { 
            vector [
%s
            ]
        }
	coordIndex [
%s
        ]
%s
	solid %s
    }
}
'''

    pts = tuple_string(vertices, '%.6g %.6g %.6g')
    ns = tuple_string(vertex_normals, '%.6g %.6g %.6g')

    m = polygon_size(vertex_indices)
    vi_format = ' '.join(['%d'] * m + ['-1'])
    ci = tuple_string(vertex_indices, vi_format)

    if vertex_rgba:
        clrs = tuple_string(vertex_rgba[:,:3], '%.3f %.3f %.3f')
        cspec = 'color Color {\n  color [\n%s]\n            }\n' % clrs
    else:
        cspec = ''

    srgb, trans = material_colors(vertex_rgba, solid_rgba)
    specular = '%.3f %.3f %.3f' % specular_rgb
    
    if two_sided_lighting:
        solid = 'FALSE'
    else:
        solid = 'TRUE'

    vrml = vrml_template % (srgb, trans, shininess, specular,
                            pts, ns, ci, cspec, solid)
    return vrml

# -----------------------------------------------------------------------------
# VRML does not handle normals for lines.
#
def mesh_vrml(vertices, vertex_indices, vertex_rgba, solid_rgba):

    vrml_template = \
'''Shape { 
    appearance Appearance {
      material Material {
        emissiveColor %s
        transparency %.3f
      }
    }
    geometry IndexedLineSet {
        coord Coordinate { 
            point [
%s
            ]
        }
	coordIndex [
%s
        ]
%s
    }
}
'''

    pts = tuple_string(vertices, '%.6g %.6g %.6g')

    m = polygon_size(vertex_indices)
    vi_format = ' '.join(['%d'] * m + ['-1'])
    ci = tuple_string(vertex_indices, vi_format)

    if vertex_rgba:
        clrs = tuple_string(vertex_rgba[:,:3], '%.3f %.3f %.3f')
        cspec = 'color Color {\n  color [\n%s]\n            }\n' % clrs
    else:
        cspec = ''

    srgb, trans = material_colors(vertex_rgba, solid_rgba)
    
    vrml = vrml_template % (srgb, trans, pts, ci, cspec)
    return vrml

# -----------------------------------------------------------------------------
#
def material_colors(vertex_rgba, solid_rgba):
    
    if vertex_rgba:
        srgb = '0 0 0'
        trans = 1 - max(vertex_rgba[:,3])
    else:
        srgb = '%.3f %.3f %.3f' % solid_rgba[:3]
        trans = 1 - solid_rgba[3]

    return srgb, trans

# -----------------------------------------------------------------------------
#
def polygon_size(vertex_indices):

    if vertex_indices:
        return len(vertex_indices[0])
    return 0

# -----------------------------------------------------------------------------
#
def tuple_string(tuples, format):

    tuple_strings = []
    for t in tuples:
        tuple_strings.append(format % tuple(t))
    tstring = '\n'.join(tuple_strings)
    return tstring

# -----------------------------------------------------------------------------
#
def surface_models():

    import chimera
    import _surface
    mlist = chimera.openModels.list(modelTypes = [_surface.Surface_Model,
                                                  chimera.MSMSModel])
    return mlist

# -----------------------------------------------------------------------------
#
def write_surfaces_dialog():

    import os.path
    dir = os.path.dirname(__file__)
    help_path = os.path.join(dir, 'helpdir', 'vrml_output.html')
    help_url = 'file://' + help_path
    filters = (('VRML', ('*.wrl', '*.vrml'), '.vrml'),)

    import OpenSave
    od = OpenSave.SaveModeless(title = 'Save Surfaces as VRML',
                               filters = filters,
                               defaultFilter = 0,
                               command = write_surfaces_cb,
                               multiple = 0,
			       help = help_url)
    od.enter()

# -----------------------------------------------------------------------------
#
def write_surfaces_cb(okayed, dialog):

    if okayed:
        paths = dialog.getPaths()
        if len(paths) == 1:
            write_surfaces_as_vrml(paths[0])

# -----------------------------------------------------------------------------
#
def write_surfaces_as_vrml(path):

    mlist = surface_models()
    vrml = surfaces_as_vrml(mlist)
    f = open(path, 'w')
    f.write(vrml)
    f.close()