Search code examples
pythonvtkenthoughtmayavi

Change color depending on height in Mayavi iso_surface


Is is possible to change the colour of an iso-surface depending on height of the points (in python / mayavi) ? I can create an iso-surface visualization with my script, but I don't know how to make the iso_surface change colour with z axis so that it will be let's say black at the bottom and white at the top of the plot. I need this in order to make sense of the visualization when it is viewed from directly above the graph. If you know any other way to achieve this, please let me know as well. I only want to show one iso_surface plot.


Solution

  • I managed to do this by combining some code from examples http://docs.enthought.com/mayavi/mayavi/auto/example_atomic_orbital.html#example-atomic-orbital and http://docs.enthought.com/mayavi/mayavi/auto/example_custom_colormap.html . Basically you must create a surface as in atomic-orbital example and then make it change colour depending on x. You must create an array of values for x. My code is (the relevant part) :

      #src.image_data.point_data.add_array(np.indices(list(self.data.shape)[self.nx,self.ny,self.nz])[2].T.ravel())
      src.image_data.point_data.add_array(np.indices(list(self.data.shape))[0].T.ravel())
      src.image_data.point_data.get_array(1).name = 'z'
      # Make sure that the dataset is up to date with the different arrays:
      src.image_data.point_data.update()
      # We select the 'scalar' attribute, ie the norm of Phi
      src2 = mlab.pipeline.set_active_attribute(src, point_scalars='scalar')
      # Cut isosurfaces of the norm
      contour = mlab.pipeline.contour(src2)
      # contour.filter.contours=[plotIsoSurfaceContours]
      # contour.filter.contours=[plotIsoSurfaceContours[0]]
      min_c = min(contour.filter._data_min * 1.05,contour.filter._data_max)
      max_c = max(contour.filter._data_max * 0.95,contour.filter._data_min)
      plotIsoSurfaceContours = [ max(min(max_c,x),min_c) for x in plotIsoSurfaceContours ]
      contour.filter.contours= plotIsoSurfaceContours
    
      # Now we select the 'angle' attribute, ie the phase of Phi
      contour2 = mlab.pipeline.set_active_attribute(contour, point_scalars='z')
      # And we display the surface. The colormap is the current attribute: the phase.
      # mlab.pipeline.surface(contour2, colormap='hsv')
      xxx = mlab.pipeline.surface(contour2, colormap='gist_ncar')
      colorbar = xxx.module_manager.scalar_lut_manager
      colorbar.reverse_lut = True
      lut = xxx.module_manager.scalar_lut_manager.lut.table.to_array()
      lut[:,-1] = int(plotIsoSurfaceOpacity * 254)
      xxx.module_manager.scalar_lut_manager.lut.table = lut
      # mlab.colorbar(title='Phase', orientation='vertical', nb_labels=3)
    

    self.data is my data, and for unknown reasons if you want to set opacity of your surface you must reverse the lut first and then set the opacity. Multiplication by 254 instead of 255 is done to avoid a possible bug in mayavi. I hope this helps someone.