process_pdb.py 文件源码

def atoms_to_density_map(atoms, voxelSZ):
    (x, y, z) = atoms[:,1:4].T.copy()
    (x_min, x_max) = (x.min(), x.max())
    (y_min, y_max) = (y.min(), y.max())
    (z_min, z_max) = (z.min(), z.max())

    grid_len = max([x_max - x_min, y_max - y_min, z_max - z_min])
    R = np.int(np.ceil(grid_len / voxelSZ))
    if R % 2 == 0:
        R += 1
    msg = "Length of particle (voxels), %d"%(R)
    logging.info(msg)
    elec_den = atoms[:,0].copy()

    #x = (x-x_min)/voxelSZ
    #y = (y-y_min)/voxelSZ
    #z = (z-z_min)/voxelSZ
    x = (x-0.5*(x_max+x_min-grid_len))/voxelSZ
    y = (y-0.5*(y_max+y_min-grid_len))/voxelSZ
    z = (z-0.5*(z_max+z_min-grid_len))/voxelSZ

    bins = np.arange(R+1)
    all_bins = np.vstack((bins,bins,bins))
    coords = np.asarray([x,y,z]).T
    #(h, h_edges) = np.histogramdd(coords, bins=all_bins, weights=elec_den)
    #return h
    #return griddata(coords, elec_den, np.mgrid[0:R,0:R,0:R].T, method='linear', fill_value=0.).T
    integ = np.floor(coords)
    frac = coords - integ
    ix = integ[:,0]; iy = integ[:,1]; iz = integ[:,2]
    fx = frac[:,0]; fy = frac[:,1]; fz = frac[:,2]
    cx = 1. - fx; cy = 1. - fy; cz = 1. - fz
    h_total = np.histogramdd(np.asarray([ix,iy,iz]).T, weights=elec_den*cx*cy*cz, bins=all_bins)[0]
    h_total += np.histogramdd(np.asarray([ix,iy,iz+1]).T, weights=elec_den*cx*cy*fz, bins=all_bins)[0]
    h_total += np.histogramdd(np.asarray([ix,iy+1,iz]).T, weights=elec_den*cx*fy*cz, bins=all_bins)[0]
    h_total += np.histogramdd(np.asarray([ix,iy+1,iz+1]).T, weights=elec_den*cx*fy*fz, bins=all_bins)[0]
    h_total += np.histogramdd(np.asarray([ix+1,iy,iz]).T, weights=elec_den*fx*cy*cz, bins=all_bins)[0]
    h_total += np.histogramdd(np.asarray([ix+1,iy,iz+1]).T, weights=elec_den*fx*cy*fz, bins=all_bins)[0]
    h_total += np.histogramdd(np.asarray([ix+1,iy+1,iz]).T, weights=elec_den*fx*fy*cz, bins=all_bins)[0]
    h_total += np.histogramdd(np.asarray([ix+1,iy+1,iz+1]).T, weights=elec_den*fx*fy*fz, bins=all_bins)[0]
    return h_total