Source: molecules.js

var container, stats;

var camera, scene, depthScene, lightMapScene, lmTestScene, renderer, effectComposer;
var controls;

var material, depthMaterial, lightMapMaterial, depthRenderTarget, lightMapRenderTarget, lightMapRenderTarget2, lmTestMaterial, aoRenderTarget;
var depthPass, lightMapPass, lmTestPass, aoPass;
var depthScale = 1.0;
var postprocessing = { enabled : false, renderMode: 0};
var LIGHT_MAP_SIZE = 4096;

var isWebGL2;

var root, depthQuad, lightMapQuad, lmTestQuad, helper;
var atomCount = 0;

var dir = "assets/"
var MOLECULES = {
  Testosterone: 'testosterone.pdb',
  Test: 'test.pdb',
  Porin: "porin.pdb",
  CO2: "formicacid.pdb",
  "3QE5": '3QE5.pdb',
  "1AON": '1aon.pdb',
  // "!Nanostuff": 'nanostuff.pdb',
  // "!4RJW": '4RJW.pdb',
};

var guiParams = {
  Molecule: Object.keys(MOLECULES)[0],
  Saturation: 0.7,
  cutOff: 0.1,
  pushBack: 1.0,
  AO: true,
}

var directions = [];
var min, max, avg;

// stuff for loading shaders. Actual shaders to be used go into shaders variable.
var vertexShaders       = $('script[type="x-shader/x-vertex"]');
var fragmentShaders     = $('script[type="x-shader/x-fragment"]');
var shadersLoaderCount  = vertexShaders.length + fragmentShaders.length;

var shaderSourceCodes = {};

var loader = new THREE.PDBLoader();

/** Initializes the WebGL components and loads the first molecule.
*/
function init() {

  if ( !Detector.webgl ) {

    Detector.addGetWebGLMessage();
    return false;

  }

  for(var i = 0; i < 60; i++) {

    var d;
    if (i % 2 == 0) {
      var x = Math.random() * 2-1;
      var y = Math.random() * 2-1;
      var z = Math.random() * 2-1;

      d = new THREE.Vector3(x, y, z).normalize();

    } else {
      d = new THREE.Vector3().copy(directions[i-1]).multiplyScalar(-1);
    }
    directions.push(d);
  }

  renderer = new THREE.WebGLRenderer({antialias: false});
  renderer.setPixelRatio( window.devicePixelRatio );
  renderer.setSize( window.innerWidth, window.innerHeight );
  renderer.setClearColor( 0xaaaaaa );
  renderer.autoClear = true;

  if ( renderer.extensions.get( 'ANGLE_instanced_arrays' ) === false ||
  renderer.extensions.get('EXT_frag_depth') === false) {
    document.getElementById( "notSupported" ).style.display = "";
    return false;
  }

  container = document.createElement( 'div' );
  container.appendChild( renderer.domElement );
  document.body.appendChild( container );

  var gl;
  try {
    gl = document.createElement('canvas').getContext('webgl2');
  } catch (err) {
    console.error(err);
  }
  isWebGL2 = Boolean(gl);

  camera = new THREE.PerspectiveCamera( 50, window.innerWidth / window.innerHeight, 1, 5000 );
  // camera = new THREE.OrthographicCamera(window.innerWidth / - 2, window.innerWidth / 2, window.innerHeight / 2, window.innerHeight / - 2, 1, 1000)
  camera.position.z = 1000;

  scene = new THREE.Scene();
  scene.add(camera);

  controls = new THREE.TrackballControls(camera, renderer.domElement );

  controls.dynamicDampingFactor = 0.2;

  stats = new Stats();
  stats.domElement.style.position = 'absolute';
  stats.domElement.style.top = '0px';
  container.appendChild( stats.domElement );

  material = new THREE.RawShaderMaterial( {
    uniforms: {
      AOTexture: {type: 't', value: 0},
      saturation: {type: 'f', value: 1 - guiParams.Saturation},
      cutOff: {type: 'f', value: guiParams.cutOff},
      pushBack: {type: 'f', value: guiParams.pushBack},
      AO: {type: 'b', value: guiParams.AO },
    },
    vertexShader: shaderSourceCodes.basic.vertex,
    fragmentShader: shaderSourceCodes.basic.fragment,
    depthTest: true,
    depthWrite: true
  } );

  root = new THREE.Group();
  helper = new THREE.Group();
  scene.add(helper);

  initPreprocessing();

  loadMolecule( MOLECULES[Object.keys(MOLECULES)[0]] );

  //
  window.addEventListener( 'resize', onWindowResize, false );

  return true;

}

/** Event listener to resize the renderer, render target and camera accordingly.
  @param {UIEvent} event - the triggered event
*/
function onWindowResize( event ) {

  var width = window.innerWidth;
  var height = window.innerHeight;

  camera.aspect = width / height;
  camera.updateProjectionMatrix();

  renderer.setSize( width, height );

  var pixelRatio = renderer.getPixelRatio();
  var newWidth  = Math.floor( width * pixelRatio ) || 1;
  var newHeight = Math.floor( height * pixelRatio ) || 1;
  depthRenderTarget.setSize( newWidth, newHeight );

}

/**loads a molecule and calls functions to calculate the ambient occlusion map.
@param {string} url - the URL from which to request the PDB file to load.
*/
function loadMolecule( url ) {
  scene.remove(root);
  root = new THREE.Group();
  scene.add(root);

  loader.load(dir+url, function(json) {

    var geometry = new THREE.InstancedBufferGeometry();
    geometry.copy( new THREE.PlaneBufferGeometry( 1,1,1,1 ) );


    atomCount = json.atoms.length;

    var translateArray = new Float32Array( atomCount * 3 );
    var scaleArray = new Float32Array( atomCount );
    var colorsArray = new Float32Array( atomCount * 3 );
    var indexArray = new Float32Array( atomCount );

    max = new THREE.Vector3(-Infinity, -Infinity, -Infinity);
    min = new THREE.Vector3(Infinity, Infinity, Infinity);
    avg = new THREE.Vector3();

    for ( var i = 0, i3 = 0, l = atomCount; i < l; i ++, i3 += 3 ) {

      translateArray[ i3 + 0 ] = json.atoms[i][0];
      max.x = Math.max(max.x, json.atoms[i][0]);
      min.x = Math.min(min.x, json.atoms[i][0]);

      translateArray[ i3 + 1 ] = json.atoms[i][1];
      max.y = Math.max(max.y, json.atoms[i][1]);
      min.y = Math.min(min.y, json.atoms[i][1]);

      translateArray[ i3 + 2 ] = json.atoms[i][2];
      max.z = Math.max(max.z, json.atoms[i][2]);
      min.z = Math.min(min.z, json.atoms[i][2]);

      colorsArray[ i3 + 0 ] = json.atoms[i][3][0]/255;
      colorsArray[ i3 + 1 ] = json.atoms[i][3][1]/255;
      colorsArray[ i3 + 2 ] = json.atoms[i][3][2]/255;

      scaleArray[ i ] = json.atoms[i][4];

      indexArray[ i ] = i;
    }

    console.log('This many atoms baby: '+atomCount);
    avg.lerpVectors(max,min,0.5);
    camera.position.set(avg.x,avg.y,max.z*10);
    camera.lookAt(avg);
    camera.up.set(0,1,0);
    controls.target = avg;

    geometry.addAttribute( "translate", new THREE.InstancedBufferAttribute( translateArray, 3, 1 ) );
    geometry.addAttribute( "sphereRadius", new THREE.InstancedBufferAttribute( scaleArray, 1, 1 ));
    geometry.addAttribute( "color", new THREE.InstancedBufferAttribute( colorsArray, 3, 1 ) );
    geometry.addAttribute( "impostorIndex", new THREE.InstancedBufferAttribute( indexArray, 1, 1 ) );

    var mesh = new THREE.Mesh( geometry, material );
    //mesh.scale.set(100,100,100);
    mesh.frustumCulled = false;
    root.add( mesh );

    updatePreprocessing();
    preprocess();
    updateLightMapOutput();
  },
  function( xhr ){
    if (xhr.lengthComputable) {
      console.log(Math.round( 100*xhr.loaded/xhr.total, 2 )+ '% loaded');
    }
  },
  function( xhr ){
    console.log(Error(xhr),url);
  });
}

/** The rendering loop. */
function animate() {

  requestAnimationFrame( animate );

  controls.update();


  render();
  stats.update();
}

/** The rendering function. */
function render() {

  if ( postprocessing.enabled ) {

    renderer.render(lmTestScene, camera);

  } else {
    scene.overrideMaterial = null;
    material.uniforms['invViewMatrix'] = { value: camera.matrixWorld };
    renderer.render( scene, camera );
  }

}

/**binary search for getting patch edge size.
  Implemented after http://stackoverflow.com/questions/6463297/algorithm-to-fill-rectangle-with-small-squares.
  @return {number} the patch width.
*/
function computePatchWidth(){
  var hi, lo;
  hi = LIGHT_MAP_SIZE;
  lo = 0.0;
  while( Math.abs(hi-lo)>0.01) {
    var mid = (lo+hi)/2.0;
    midval = Math.floor(LIGHT_MAP_SIZE/mid) * Math.floor(LIGHT_MAP_SIZE/mid);
    if(midval >= atomCount){
      lo = mid;
    }
    else {
      hi = mid;
    }
  }
  return LIGHT_MAP_SIZE/Math.floor(LIGHT_MAP_SIZE/lo)
}

/**
  Updates the render targets containing the light map for Ambient Occlusion.
  */
function updateLightMapTarget(){
  var lightMapRenderTargetParams = {
    format: THREE.RGBFormat,
    minFilter: THREE.NearestFilter,
    magFilter: THREE.NearestFilter,
    generateMipmaps: false,
    stencilBuffer: false,
    depthBuffer: false
  };

  lightMapRenderTarget = new THREE.WebGLRenderTarget( LIGHT_MAP_SIZE, LIGHT_MAP_SIZE, lightMapRenderTargetParams);
  lightMapRenderTarget2 = new THREE.WebGLRenderTarget( LIGHT_MAP_SIZE, LIGHT_MAP_SIZE, lightMapRenderTargetParams);

  lightMapMaterial = new THREE.RawShaderMaterial({
    uniforms: {
      cameraNear: { type: 'f', value: 1.0 },
      cameraFar: { type: 'f', value: 5000.0 },
      size: { type: 'f', value: LIGHT_MAP_SIZE },
      patchSize: {type: 'f', value: computePatchWidth()},
      lightStepWidth: { type: 'f', value: (2.0/directions.length) },
      numberOfImpostors: { type: 'f', value: atomCount },
      tDepth : { type: 't', value: depthRenderTarget.depthTexture},
      tLightMap : { type: 't', value: lightMapRenderTarget2.texture}
    },
    fragmentShader: shaderSourceCodes.lightmap.fragment,
    vertexShader: shaderSourceCodes.lightmap.vertex,
    depthTest: false,
    depthWrite: false
  });

  material.uniforms['size'] = { type: 'f', value: LIGHT_MAP_SIZE };
  material.uniforms['patchSize'] = {type: 'f', value: computePatchWidth()};
  material.uniforms['numberOfImpostors'] = { type: 'f', value: atomCount };

  lightMapPlane = new THREE.PlaneBufferGeometry(2, 2);
  lightMapQuad = new THREE.Mesh(lightMapPlane, lightMapMaterial);
  lightMapQuad.frustumCulled = false;
  lightMapScene.add(lightMapQuad);
}

/**
  Updates the output of the light map (used for debugging).
*/
function updateLightMapOutput(){
  lmTestMaterial = new THREE.RawShaderMaterial({
    uniforms:{
      tLightMap : {type: 't', value: lightMapRenderTarget.texture}
    },
    fragmentShader: shaderSourceCodes.lmtest.fragment,
    vertexShader: shaderSourceCodes.lmtest.vertex
  });

  lmTestPlane = new THREE.PlaneBufferGeometry(2, 2);
  lmTestQuad = new THREE.Mesh(lmTestPlane, lmTestMaterial);
  lmTestQuad.frustumCulled = false;
  lmTestScene.add(lmTestQuad);

}

/**
  Sets up rendering components for Ambient Occlusion.
*/
function initPreprocessing() {
  // Setup depth rendering pass
  depthRenderTarget = new THREE.WebGLRenderTarget( window.innerWidth * window.devicePixelRatio, window.innerHeight * window.devicePixelRatio);
  depthRenderTarget.texture.format = THREE.RGBFormat;
  depthRenderTarget.texture.minFilter = THREE.NearestFilter;
  depthRenderTarget.texture.magFilter = THREE.NearestFilter;
  depthRenderTarget.texture.generateMipmaps = false;
  depthRenderTarget.stencilBuffer = false;
  depthRenderTarget.depthBuffer = true;
  depthRenderTarget.depthTexture = new THREE.DepthTexture();
  depthRenderTarget.depthTexture.type = isWebGL2 ? THREE.FloatType : THREE.UnsignedShortType;


  lightMapScene = new THREE.Scene();
  lmTestScene = new THREE.Scene();
  updateLightMapTarget();
}

/**
  Updates rendering components for Ambient Occlusion.
*/
function updatePreprocessing(){
  lightMapScene.remove(lightMapQuad);
  lmTestScene.remove(lmTestQuad);
  updateLightMapTarget();
}

/**
  Initializes the GUI and adds callbacks for change events.
*/
function initGui() {
  var gui = new dat.GUI();

  gui.add(guiParams, 'Molecule', Object.keys(MOLECULES)).onChange(function(v){loadMolecule(MOLECULES[v])});
  gui.add(guiParams, 'Saturation', 0.0, 1.0).onChange(function(v){material.uniforms.saturation.value = 1.0 - v});
  gui.add(guiParams, 'cutOff', 0.0, 1.0).onChange(function(v){material.uniforms.cutOff.value = v});
  gui.add(guiParams, 'pushBack', 0.0, 5.0).onChange(function(v){material.uniforms.pushBack.value = v});
  gui.add(guiParams, 'AO').onChange(function(v){material.uniforms.AO.value = v});

}

/**
  Loads a shader.
  @param {DOMElement} shader - the element describing the shader.
  @param {"fragment"|"vertex"} type - the type of shader.
*/
function loadShader(shader, type) {
  var $shader = $(shader);
  if(shaderSourceCodes[$shader.attr('title')]===undefined){
    shaderSourceCodes[$shader.attr('title')]={};
  }
  $.ajax({
    url: $shader.attr('src'),
    dataType: 'text',
    context: {
      name: $shader.attr('title'),
      type: type
    },
    complete: function( jqXHR, textStatus ) {
      shadersLoaderCount--;

      shaderSourceCodes[$shader.attr('title')][this.type] = jqXHR.responseText;
      if ( !shadersLoaderCount ) {
        shadersLoadComplete();
      }
    }
  });
}

/**
  Creates the light map for Ambient Occlusion.
*/
function preprocess(){
  directions.forEach(function(direction){
    orthoCamera = orthoView(direction);
    renderer.render(scene, orthoCamera, depthRenderTarget);

    //ping pong
    var tempTexture = lightMapRenderTarget2;
    lightMapRenderTarget2 = lightMapRenderTarget;
    lightMapRenderTarget = tempTexture;
    lightMapMaterial.uniforms.tLightMap.value = lightMapRenderTarget2.texture;

    scene.overrideMaterial = lightMapMaterial;
    renderer.render(scene, orthoCamera, lightMapRenderTarget);
    scene.overrideMaterial = null;
  });

  material.uniforms.AOTexture.value = lightMapRenderTarget.texture;
  console.log('Preprocessing finished.');
}

/**
  Creates a THREE.OrthographicCamera facing the molecule from the given direction.
  The molecule has to completely be in the camera's frustum.
  @param {THREE.Vector3} direction - the direction in which the camera should face.
  @param {boolean} show - creates and adds a THREE.CameraHelper to the scene.
  @return {THREE.OrthographicCamera} The camera looking at the molecule from the given direction.
*/
function orthoView(direction, show) {
  var ortho = new THREE.OrthographicCamera(-1,1,1,-1,0.1,100);

  var target = new THREE.Vector3();
  target.addVectors(avg,direction);

  ortho.position.copy(avg);
  ortho.lookAt(target);
  ortho.updateMatrix();
  ortho.updateMatrixWorld();

  var translateArray = root.children[0].geometry.getAttribute('translate').array;


  var max = new THREE.Vector3(-Infinity, -Infinity, -Infinity);
  var min = new THREE.Vector3(Infinity, Infinity, Infinity);

  for (var i = translateArray.length - 1; i >= 0; i-=3) {
    var z = translateArray[i];
    var y = translateArray[i-1];
    var x = translateArray[i-2];

    var camSpace = ortho.worldToLocal(new THREE.Vector3(x,y,z));

    max.x = Math.max(max.x, camSpace.x);
    max.y = Math.max(max.y, camSpace.y);
    max.z = Math.max(max.z, camSpace.z);

    min.x = Math.min(min.x, camSpace.x);
    min.y = Math.min(min.y, camSpace.y);
    min.z = Math.min(min.z, camSpace.z);

  }

  var width = (10 + max.x - min.x)/2;
  var height = (10 + max.y - min.y)/2;
  var depth = 10 + max.z - min.z;

  var a = new THREE.Vector3();
  a.lerpVectors(max,min,0.5);
  a = ortho.localToWorld(a);

  ortho.left = -width;
  ortho.right = width;
  ortho.top = height;
  ortho.bottom = -height;
  ortho.far = depth;
  ortho.near = 0.1;

  var eye = new THREE.Vector3();
  eye.copy(direction);
  eye.multiplyScalar(-(depth/2 + 0.1));

  ortho.position.addVectors(a, eye);
  ortho.lookAt(a);

  ortho.updateMatrix();
  ortho.updateMatrixWorld();
  ortho.updateProjectionMatrix();

  helper.remove(helper.children[0]);

  if(show) {
    console.log(ortho);
    var ch = new THREE.CameraHelper(ortho);
    helper.add(ch);
  } else {
    return ortho;
  }
}


/**
  Starts initialization of the actual application after the shaders have been loaded.
*/
function shadersLoadComplete() {
  console.log(Object.keys(shaderSourceCodes).length + " shaders loaded. Initializing..." );
  if(init()) {
    initGui();
    // initPreprocessing();
    // preprocess();
    console.log("Initialization succesfull. Start rendering...");
    animate();
  }
}

/**
  Entry point. First loads all the shaders, then init() etc. is called after load completion.
*/
function start() {

  for(var i =0; i<vertexShaders.length;i++){
    loadShader(vertexShaders[i], 'vertex');
  }
  for(var i =0; i<fragmentShaders.length;i++){
    loadShader(fragmentShaders[i], 'fragment');
  }
}

start();