AmbientOcclusion.cs

Go to the documentation of this file.

using UnityEngine;
using System.Collections;
using System.Collections.Generic;
using Molecule.Model;
//using UI;

public class AmbientOcclusion {
    private static GameObject pdb2den;
    private static PDBtoDEN genDensity;
    private static float[,,] density;
    private static float[,,] occlusion;
    private static BallUpdate[] balls;
    private static Vector3 delta;
    private static Vector3 origin = MoleculeModel.MinValue;
    private static Vector3 fudgeFactor;
    
    private static Vector3 inverseDelta ;
    private static int xDim, yDim, zDim;
    private static List<Color> originalColors;
    
    private static float minDensity = float.MaxValue; // Unused ?
    private static float maxDensity = -float.MaxValue;
    
    public static bool isEnabled = false;
    public static bool reset = false;
    
    private const float faceCoef = 1f;
    private const float diagCoef = 0.707106781186548f; // 1/sqrt(2)
    private const float cornerCoef = 0.25f;
    private const float totalWeightFactor = 1f/(6f*faceCoef + 12f*diagCoef + 8f*cornerCoef);
    
    private static bool weighted = true;
    private static float compensationFactor = 1.30f;
    //private static float compensationTerm = 0.12f;
    //private static Color compensationColor = new Color(compensationTerm, compensationTerm, compensationTerm);
    
    
    private static float AverageOfSurroundingCubes(int i, int j, int k) {
        // j-1                  // i-1                      // i                        // i+1
        float lowerNine     =   density[i-1, j-1, k-1]  +   density[i, j-1, k-1]    +   density[i+1, j-1, k-1]  + // k-1
                                density[i-1, j-1, k]    +   density[i, j-1, k]      +   density[i+1, j-1, k]    + // k
                                density[i-1, j-1, k+1]  +   density[i, j-1, k+1]    +   density[i+1, j-1, k+1]  ; // k+1
        
        // j                    // i-1                      // i                        // i+1
        float middleEight   =   density[i-1, j, k-1]    +   density[i, j, k-1]      +   density[i+1, j, k-1]    + // k-1
                                density[i-1, j, k]      +   0f                      +   density[i+1, j, k]      + // k
                                density[i-1, j, k+1]    +   density[i, j, k+1]      +   density[i+1, j, k+1]    ; // k+1
        
        // j+1                  // i-1                      // i                        // i+1
        float upperNine     =   density[i-1, j+1, k-1]  +   density[i, j+1, k-1]    +   density[i+1, j+1, k-1]  + // k-1
                                density[i-1, j+1, k]    +   density[i, j+1, k]      +   density[i+1, j+1, k]    + // k
                                density[i-1, j+1, k+1]  +   density[i, j+1, k+1]    +   density[i+1, j+1, k+1]  ; // k+1
        
        float result = lowerNine + middleEight + upperNine;
        result /= 26;
        return result;
    }
    
    
    private static float WeightedAverageOfSurroundingCubes(int i, int j, int k) {
        // j-1                  // i-1                      // i                        // i+1
        float lowerNine     =   cornerCoef  * density[i-1, j-1, k-1]    +   diagCoef    * density[i, j-1, k-1]  +   cornerCoef  * density[i+1, j-1, k-1]    + // k-1
                                diagCoef    * density[i-1, j-1, k]      +   faceCoef    * density[i, j-1, k]    +   diagCoef    * density[i+1, j-1, k]      + // k
                                cornerCoef  * density[i-1, j-1, k+1]    +   diagCoef    * density[i, j-1, k+1]  +   cornerCoef  * density[i+1, j-1, k+1]    ; // k+1
        
        // j                    // i-1                      // i                        // i+1
        float middleEight   =   diagCoef    * density[i-1, j, k-1]      +   faceCoef    * density[i, j, k-1]    +   diagCoef    * density[i+1, j, k-1]      + // k-1
                                faceCoef    * density[i-1, j, k]        +   0f                                  +   faceCoef    * density[i+1, j, k]        + // k
                                diagCoef    * density[i-1, j, k+1]      +   faceCoef    * density[i, j, k+1]    +   diagCoef    * density[i+1, j, k+1]      ; // k+1
        
        // j+1                  // i-1                      // i                        // i+1
        float upperNine     =   cornerCoef  * density[i-1, j+1, k-1]    +   diagCoef    * density[i, j+1, k-1]  +   cornerCoef  * density[i+1, j+1, k-1]    + // k-1
                                diagCoef    * density[i-1, j+1, k]      +   faceCoef    * density[i, j+1, k]    +   diagCoef    * density[i+1, j+1, k]      + // k
                                cornerCoef  * density[i-1, j+1, k+1]    +   diagCoef    * density[i, j+1, k+1]  +   cornerCoef  * density[i+1, j+1, k+1]    ; // k+1
        
        float result = lowerNine + middleEight + upperNine;
        result *= totalWeightFactor;
        return result;
    }
    
    private static void BuildColorList() {
        originalColors = new List<Color>(MoleculeModel.atomsColorList);
    }
    
    private static void FillFacesWithZeros() {
        for(int i=0; i<xDim; i++) {
            for(int j=0; j<yDim; j++) {
                occlusion[i,        j,          0       ] = 0f; // front face
                occlusion[i,        j,          zDim-1  ] = 0f; // back face
            }
            for(int k=0; k<zDim; k++) {
                occlusion[i,        0,          k       ] = 0f; // bottom face
                occlusion[i,        yDim-1,     k       ] = 0f; // top face
            }
        }
        
        for(int j=0; j<yDim; j++) {
            for(int k=0; k<zDim; k++) {
                occlusion[0,        j,          k       ] = 0f; // left face
                occlusion[xDim-1,   j,          k       ] = 0f; // right face
            }
        }
    }
    
    private static void SetOcclusionGridBounds() {
        minDensity = float.MaxValue;
        maxDensity = -float.MaxValue;
        foreach(float f in density) {
            if(f < minDensity)
                minDensity = f;
            if(f > maxDensity)
                maxDensity = f;
        }
        
        /*
        float threshold = maxDensity * 0.3f;
        foreach(float f in occlusion)
            if(f > threshold)
                Debug.Log(f.ToString());
        */
        
        Debug.Log("Density bounds, min and max:");
        Debug.Log(minDensity.ToString());
        Debug.Log(maxDensity.ToString());
    }
    
    public AmbientOcclusion () {
        // First we use PDBtoDen to create a density grid.
        pdb2den = new GameObject();
        pdb2den.AddComponent<PDBtoDEN>();
        genDensity = pdb2den.GetComponent<PDBtoDEN>();
        
        
        float delta_scalar;
        int moleculeSize = MoleculeModel.atomsLocationlist.Count;
        if(moleculeSize < 2000)
            delta_scalar = 0.2f;
        else if (moleculeSize < 5000)
            delta_scalar = 0.4f;
        else if (moleculeSize < 10000)
            delta_scalar = 0.2f;
        else
            delta_scalar = 0.16f;
        //fudgeFactor = PDBtoDEN.fudgeFactor /delta_scalar; ???
        delta = new Vector3(delta_scalar, delta_scalar, delta_scalar);
        genDensity.TranPDBtoDEN(delta_scalar, false);
        density = PDBtoDEN.GridS;
        fudgeFactor = PDBtoDEN.fudgeFactor;
        
        // We need to refresh the molecule's origin when it's not the first molecule for which we generate a volumetric density.
        origin = MoleculeModel.MinValue;
        
        xDim = density.GetLength(0);
        yDim = density.GetLength(1);
        zDim = density.GetLength(2);
        
        Debug.Log("Density grid size :");
        Debug.Log(xDim.ToString());
        Debug.Log(yDim.ToString());
        Debug.Log(zDim.ToString());
        
        occlusion = new float[xDim, yDim, zDim];
        
        for(int i=1; i<(xDim-1); i++)
            for(int j=1;  j<(yDim-1); j++)
                for(int k=1; k<(zDim-1); k++)
                    if(weighted)
                        occlusion[i,j,k] = WeightedAverageOfSurroundingCubes(i,j,k);
                    else
                        occlusion[i,j,k] = AverageOfSurroundingCubes(i,j,k);
        
        FillFacesWithZeros();
        SetOcclusionGridBounds();
        balls = GameObject.FindObjectsOfType(typeof(BallUpdate)) as BallUpdate[];
        BuildColorList();
        reset = false;
    }
    
    public void Occlude() {
        int x, y, z;
        Vector3 position = Vector3.zero;
        Vector3 indices = Vector3.zero;
        float colorFactor;
        Debug.Log(maxDensity.ToString());
        //float threshold = 0.3f * maxDensity;
        
        /*
        foreach(float f in occlusion) {
            if(f > threshold)
                Debug.Log(f.ToString());
        }
        */
        Color color;
        foreach(BallUpdate bu in balls) {
            position = bu.transform.position;
            //indices = Vector3.Scale(delta, (position + fudgeFactor - origin));
            indices = Vector3.Scale(delta, (position - origin)) + fudgeFactor;
            /*
            x = (int) indices.x;
            y = (int) indices.y;
            z = (int) indices.z;
            */
            x = Mathf.RoundToInt(indices.x);
            y = Mathf.RoundToInt(indices.y);
            z = Mathf.RoundToInt(indices.z);
            colorFactor = (maxDensity - occlusion[x,y,z]) / maxDensity;
            
            /*
            if(occlusion[x,y,z] > threshold) {
                Debug.Log(occlusion[x,y,z].ToString());
                Debug.Log(colorFactor.ToString());
            }

            
            //if(colorFactor > 0.1f)
            //Debug.Log(colorFactor + " - " + maxDensity);
            */
            //bu.atomcolor = bu.atomcolor * colorFactor;
            color = MoleculeModel.atomsColorList[(int)bu.number];
            color = (color * colorFactor * compensationFactor);// + compensationColor;
            MoleculeModel.atomsColorList[(int)bu.number] = color;
            /*
            Color test = bu.atomcolor;
            test.a = test.a * (1 - colorFactor);
            bu.atomcolor = test;
            */
        }
        BallUpdate.resetColors = true;
        isEnabled = true;
    }
    
    public void Revert() {
        Debug.Log("AmbientOcclusion: reverting to original colors.");
        MoleculeModel.atomsColorList = originalColors;
        BallUpdate.resetColors = true;
        isEnabled = false;
    }
}