PyMolRenderer.java

/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package pdbws.renderer;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import javax.imageio.ImageIO;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Wrapper class to render images using PyMOL.
 * This class should be used in any application, using ProteinRendererInterface instead
 * is recommended.
 * <b>Note:</b> Currently needs /dev/shm (*NIX ramdisk) to save the images to.
 * This may be faster than using temporary directory files. The ramdisk size
 * should be appropriate for the number of premature images
 * (about 10MB/100 images). This class does NOT delete any premature PNGs by itself.
 * 
 * <b>Important:</b>The bean property xUsed MUST be false if there is no running
 * X server (or, the standard graphical interface on non-*NIX-OSs) on the server,
 * else rendering will not work.
 * Rendering with xUsed == false is much slower than with xUsed == true
 * (In this mode PyMOL can use OpenGL to render the premature images)
 * <b>How this works:</b>
 * At first a variable number of images are rendered in low-quality
 * and low-resolution mode, each with a differend (these are referred as 'premature').
 * An external class loads each image and scores it using different algorithms.
 * The image having the best score gives the rotation parameters for the 'mature'
 * (full resolution, high quality) image.
 * PyMOL is called by writing a PML script and executing PyMOL with the filename
 * as argument
 * (There is no method calling all steps at once;
 * the appropriate methods must be called separately)
 *
 * @author uli
 */
public class PyMOLRenderer
{
    //TODO Randomize pml file name each time
    private static final String colorByBFactorsCommand = "spectrum b,rainbow\n";
    private static final String hideAllCommand = "hide all\n";
    private static final String tempDir = "/dev/shm/";
    private static final String pmlFile = tempDir + "prot.pml";
    RandomNumberProvider rand = RandomNumberProvider.inst;
    private static Logger logger = LoggerFactory.getLogger(
            PyMOLRenderer.class);
    private Vec2d resolution;
    private ImageScore imageScore = new ImageScore();
    //Static configuration
    private static String pymol = "pymol "; //Generate by some setters
    private static boolean xUsed = true;

    public boolean isXUsed()
    {
        return xUsed;
    }

    public void setXUsed(boolean xUsed)
    {
        PyMOLRenderer.xUsed = xUsed;
        assemblePyMOLCommandline();
    }

    private static void assemblePyMOLCommandline()
    {
        pymol = "pymol ";
        if (!xUsed)
        {
            pymol += ("-qxci "); //Command line mode, no GUI
        }
    }

    @Override
    protected void finalize() throws Throwable
    {
        System.out.println("Finalizing");
        //Delete the script
        File pml = new File(pmlFile);
        if (pml.exists())
        {
            pml.delete();
        }
    }

    /**
     * Constructs a new PyMOLRender with the default resolution (640x480)
     */
    public PyMOLRenderer()
    {
        //Resolution defaults to 640x480 here:
        resolution = new Vec2d(640, 480);
    }

    /**
     * Constructs a new PyMOLRenderer with a given resolution
     * @param xres The x resolution (in pixels)
     * @param yres The y resolution (in pixels)
     */
    public PyMOLRenderer(int xres, int yres)
    {
        resolution = new Vec2d(xres, yres);
    }

    /**
     * Sets the resolution
     * @param xres The x resolution (in pixels)
     * @param yres The y resolution (in pixels)
     */
    public void setResolution(int xres, int yres)
    {
        resolution = new Vec2d(xres, yres);
    }

    /**
     * Utility data structure representing a vector in R^3
     */
    private class Vec3d
    {

        public int x = 0;
        public int y = 0;
        public int z = 0;

        /**
         * Constructs a new null vector
         */
        public Vec3d()
        {
            //Values already are initialized with 0
            //so nothing is to do here
        }

        /**
         * Fills this vector with random data
         */
        public Vec3d randomize()
        {
            x = RandomNumberProvider.inst.nextInt();
            y = RandomNumberProvider.inst.nextInt();
            z = RandomNumberProvider.inst.nextInt();
            return this;
        }

        /**
         * Constructs a new Vector with given values
         * @param x The x value
         * @param y The y value
         * @param z The z value
         */
        public Vec3d(int x, int y, int z)
        {
            this.x = x;
            this.y = y;
            this.z = z;
        }
    }

    /**
     * Class representing a vector in R^2
     */
    private class Vec2d
    {

        public int x = 0;
        public int y = 0;

        /**
         * Constructs a new Vec2d with both values given
         * @param x The x value
         * @param y The y value
         */
        public Vec2d(int x, int y)
        {
            this.x = x;
            this.y = y;
        }
    }

    /**
     * Writes a matures rendering script (final rendering; uses raytracing)
     * @param id The PDB ID (The appropriate PDB file has to be present)
     * @param rot The rotation vector; relative to null
     * @param properties The rendering properties
     */
    private void writeMaturePML(String id, Vec3d rot,
                                RenderProperties properties) throws IOException
    {
        BufferedWriter w = null;
        try
        {
            new File(pmlFile).delete();
            w = new BufferedWriter(new FileWriter(pmlFile));
            //Show the cartoon only
            w.write(hideAllCommand);
            w.write(properties.getCommand());
            //Color the cartoon
            w.write(colorByBFactorsCommand);
            w.write("rotate x," + rot.x + "\n");
            w.write("rotate y," + rot.y + "\n");
            w.write("rotate z," + rot.z + "\n");
            //Raytrace
            w.write(String.format("ray %d,%d\n", resolution.x, resolution.y));
            String filename = id + ".png"; //PDB ID with ".png" extensions
            w.write("png " + filename + "\n");
            w.write("quit\n");
            w.close();
        }
        finally
        {
            try
            {
                w.close();
            }
            catch (IOException ex)
            {
                logger.error(null, ex);
            }
        }
    }

    /**
     * Writes a PML script for premature (test for scoring) rendering
     * @param id The PDB ID (The appropriate PDB file has to be present)
     * @param n The number of times to test-render
     * @param properties The rendering properties
     * @return A map mapping filenames to the rotation vectors
     * @throws IOException
     */
    private Map<String, Vec3d> writePrematurePML(String id, int n,
                                                 RenderProperties properties)
            throws
            IOException
    {
        try
        {
            //If we should create a mature script, set n to 1
            //Delete the PyMol script if it exists
            new File(pmlFile).delete();
            BufferedWriter w = new BufferedWriter(new FileWriter(pmlFile));
            //Show the cartoon only
            w.write(hideAllCommand);
            w.write(properties.getCommand());
            //Color the cartoon
            w.write(colorByBFactorsCommand);
            //Rotate by random degrees n times
            Map<String, Vec3d> fileProperties = new HashMap<String, Vec3d>(); //Maps the filename to the rotation
            Vec3d rot;
            for (int i = 0; i < n; i++)
            {
                //Apply a random rotation
                w.write("reset\n");
                rot = new Vec3d(rand.nextInt(360), rand.nextInt(360),
                                rand.nextInt(360));
                w.write("rotate x," + rot.x + "\n");
                w.write("rotate y," + rot.y + "\n");
                w.write("rotate z," + rot.z + "\n");
                //If we should create a mature script, raytrace
                String filename = tempDir + id + i + ".png";
                fileProperties.put(filename, rot);
                w.write("png " + filename + "\n");
            }
            w.write("quit\n");
            w.close();
            return fileProperties;
        }
        catch (IOException ex)
        {
            logger.error(null, ex);
        }
        return null;
    }

    /**
     * Modifies the image scoring weights
     */
    public void modifyValues(int tooEmpty, int tooLessContrast,
                             int notEnoughDetail)
    {
        imageScore.modifyValues(tooEmpty, tooLessContrast, notEnoughDetail);
    }

    /**
     * Fast rendering function. Doesn't use raytracing or scoring
     * @param input The PDB input file
     * @param properties The rendering properties 
     * @return
     * @throws IOException
     * @throws InterruptedException
     */
    public File renderTest(File input, RenderProperties properties) throws
            IOException, InterruptedException
    {
        String cmd = pymol + input.getAbsolutePath() + " " + pmlFile;
        String id = input.getName().split("\\.")[0]; //Guess the PDB id is the input filename minus the extension
        writeMaturePML(id, new Vec3d().randomize(), properties);
        Runtime.getRuntime().exec(cmd.toString(), null, input.getParentFile()).
                waitFor();
        return new File(tempDir + id + ".png");
    }

    /**
     * Renders a protein with using raytracing and scoring
     * @param id The PDB ID
     * @param input The PDB file
     * @param tries The number of times to test-render
     * @param properties The rendering properties
     * @return A fiel containing the finally-rendered PNG data
     * @throws IOException
     * @throws InterruptedException
     */
    public File renderRayByScore(String id, File input, int tries,
                                 RenderProperties properties) throws IOException,
                                                                     InterruptedException
    {
        String cmd = pymol + input.getAbsolutePath() + " " + pmlFile;
        double score = Double.MIN_VALUE; //There MUST be a score greater than this
        Vec3d highestScoredVec = null;
        for (int i = 0; i < tries; i += 300) //PyMOL can handle only ~300 PNGs in one script
        {
            //Score all images and pick the highest-scored one
            int numTries = 300; //Number of tries in one PyMol run, max. 300
            if ((tries - i) < 300)
            {
                numTries = tries;
            }
            Map<String, Vec3d> filenames =
                               writePrematurePML(id, numTries, properties);
            //Render the premature pictures
            Runtime.getRuntime().exec(cmd.toString()).waitFor();
            //Score all images and pick the best
            for (String filename : filenames.keySet())
            {
                //Score the image
                imageScore.setImg(ImageIO.read(new File(filename)));
                double newScore = imageScore.getScore();
                if (newScore > score)
                {
                    highestScoredVec = filenames.get(filename);
                    score = newScore;
                }
                //Delete all premature PNGs
                new File(filename).delete();
            }
        }
        //Write the mature script and render
        writeMaturePML(id, highestScoredVec, properties);
        Runtime.getRuntime().exec(cmd.toString(), null, input.getParentFile()).
                waitFor();
        //Delete the PML file
        new File(pmlFile).delete();
        //Return the highest-scoring file
        return new File(tempDir + id + ".png");
    }
}