enthought / mayavi

3D visualization of scientific data in Python
http://docs.enthought.com/mayavi/mayavi/
Other
1.28k stars 282 forks source link

how-to: volume plot on irregular grid? #1268

Open sthinius87 opened 11 months ago

sthinius87 commented 11 months ago

Friendly "hello", Inspired by the chemistry examples on the homepage I started to make volume (electron density) plots of atomic systems. This works quiet well for orthogonal cells, but fails for systems with tilted cells. image Is there an easy way to transform the volume object matching the cell? this is my code so far:

import numpy as np
from ase.data import covalent_radii
from ase.io.cube import read_cube_data
from ase.data.colors import cpk_colors
from ase.units import Bohr
from tvtk.util.ctf import ColorTransferFunction
from tvtk.util.ctf import PiecewiseFunction
from tvtk.util import ctf
from matplotlib.pyplot import cm
from mayavi import mlab
import matplotlib.pyplot as plt

def plot_vol(atoms, data):

    mlab.figure(1, bgcolor=(1, 1, 1) ,size=(1200, 800))  # make a white figure

    ## Plot the atoms as spheres:
    for pos, Z in zip(atoms.positions, atoms.numbers):
        mlab.points3d(*pos,
                      scale_factor=covalent_radii[Z]*0.5,
                      resolution=20,
                      color=tuple(cpk_colors[Z]))

    ## Draw the unit cell:
    A = atoms.cell
    for i1, a in enumerate(A):
        i2 = (i1 + 1) % 3
        i3 = (i1 + 2) % 3
        for b in [np.zeros(3), A[i2]]:
            for c in [np.zeros(3), A[i3]]:
                p1 = b + c
                p2 = p1 + a
                mlab.plot3d([p1[0], p2[0]],
                            [p1[1], p2[1]],
                            [p1[2], p2[2]],
                            tube_radius=0.1)

    min = data.min()
    max = data.max()

    ##try to adjust x,y,z manually to the cell
    nx,ny,nz=np.shape(data)
    start=np.array([0.0,0.0,0.0])
    pts=[]
    xstep,ystep,zstep=atoms.cell[0]/nx,atoms.cell[1]/ny,atoms.cell[2]/nz
    for ix in range(1,nx+1):
        #outer
        tmp_outer=ix*xstep
        for iy in range(1,ny+1):
            #middle
            tmp_middle=iy*ystep+tmp_outer
            for iz in range(1,nz+1):
                #inner
                tmp_inner=iz*zstep+tmp_middle
                pts.append(tmp_inner)
    pts=np.asarray(pts).copy()
    x=np.reshape(pts[:,0],np.shape(data))
    y=np.reshape(pts[:,1],np.shape(data))
    z=np.reshape(pts[:,2],np.shape(data))

    source = mlab.pipeline.scalar_field(x,y,z,data) #x,y,z, ,colormap='GnBu'

    ### scale density to atoms
    #engine = mlab.get_engine()
    #ss = engine.scenes[0].children[-1] # the last child is the density plot
    #S=np.shape(data)
    #C=atoms.cell.array
    #ss.spacing = np.array([C[0,0]/S[0], C[1,1]/S[1], C[2,2]/S[2]])
    #ss.origin  = np.array([0.0,0.0,0.0])

    vol = mlab.pipeline.volume(source, vmin=0.0, vmax=0.005)

    ## Changing the ctf:
    ctf = ColorTransferFunction()

    cmap=cm.get_cmap('viridis').colors[::2]
    lcmap=len(cmap)
    for num,cstep in enumerate(cmap):
        ctf.add_rgb_point((lcmap-num)/lcmap,cstep[0],cstep[1],cstep[2])

    vol._volume_property.set_color(ctf)
    vol._ctf = ctf
    vol.update_ctf = True
    # Changing the otf:
    otf = PiecewiseFunction()
    for r,o in zip(np.arange(0.0,0.2,0.02),np.arange(0.0,0.8,0.08)):
        otf.add_point(r, o)
    vol._otf = otf
    vol._volume_property.set_scalar_opacity(otf)

    mlab.view(azimuth=155, elevation=70, distance='auto')
    # Show the 3d plot:
    mlab.show()

dens = 'density.cube'
data, atoms = read_cube_data(dens)
plot_vol(atoms,data)

Any approaches to solutions are very appreciated.

sthinius87 commented 11 months ago

I already found a solution that worked for contour3d plots, but not for volume plots. See, ase_mlab