[Date Prev][Date Next][Thread Prev][Thread Next][Date Index][Thread Index]

[IMP-dev] CMM log optimizer state



hey,

A CMMLogOptimizerState class is attached. It is similar to the VRML one.

Keren.
/**
 *  \file CMMLogOptimizerState.h   
 *  \brief A state which writes a series of CMM files.
 *
 *  Copyright 2007-8 Sali Lab. All rights reserved.
 *
 */

#ifndef __IMP_CMM_LOG_STATE_H
#define __IMP_CMM_LOG_STATE_H

#include <iostream>

#include "../../IMP_config.h"
#include "../../base_types.h"
#include "../../OptimizerState.h"

namespace IMP
{

//! A state that writes a series of CMM files.
/** The State writes a series of files generated from a 
    printf-style format string. These files contain spheres
    and points for a set of particles. The particles must have
    x,y,z coordinates (as defined by XYZDecorator) and,
    optionally have a radius.

    Documentation for the CMM file format can be found at
    http://www.cgl.ucsf.edu/chimera/docs/ContributedSoftware/volumepathtracer/volumepathtracer.html#markerfiles

    \ingroup log
 */
class IMPDLLEXPORT CMMLogOptimizerState : public OptimizerState
{
 public:
  CMMLogOptimizerState(std::string filename,
		       const Particles &pis=Particles());
  virtual ~CMMLogOptimizerState(){}

  virtual void update();

  virtual void show(std::ostream& out = std::cout) const;
  virtual std::string version() const {
    return "0.5";
  }
  virtual std::string last_modified_by() const {
    return "Keren Lasker";
  }
  //! Set the number of update calls to skip between writing files
  /** The first update call always writes a file.
   */
  void set_skip_steps(unsigned int i) {
    skip_steps_=i;
  }

  //! The float key to use for the radius
  /** Particles without such an attribute are drawn as fixed sized markers.
   */
  void set_radius(FloatKey k) {
    radius_=k;
  }
  //! The three color components
  /** Color values should be between 0 and 1. They will be snapped if needed.
   */
  void set_color(FloatKey r, FloatKey g, FloatKey b) {
    r_=r; g_=g; b_=b;
  }

  //! Set the particles to use.
  void set_particles(const Particles &pis) {
    pis_=pis;
  }
  void write(std::string name) const;

  //! A helper function to just write a list of particles to a file
  static void write(const Particles &pis,
		    const std::string &marker_set_name,
                    FloatKey radius_key,
                    FloatKey r_key,FloatKey g_key, FloatKey b_key,
                    std::ostream &out);

protected:
  Particles pis_;
  std::string filename_;
  int file_number_;
  int call_number_;
  int skip_steps_;
  FloatKey radius_;
  FloatKey r_, g_, b_;
};


IMP_OUTPUT_OPERATOR(CMMLogOptimizerState);

} // namespace IMP

#endif  /* __IMP_CMM_LOG_STATE_H */
/**
 *  \file CMMLogOptimizerState.h
 *  \brief A state which writes a series of CMM files.
 *
 *  Copyright 2007-8 Sali Lab. All rights reserved.
 *
 */

#include <fstream>
#include <sstream>

#include "IMP/optimizers/states/CMMLogOptimizerState.h"
#include "IMP/decorators/XYZDecorator.h"

namespace IMP
{

CMMLogOptimizerState::CMMLogOptimizerState(std::string filename,
                                             const Particles &pis) :
    pis_(pis), filename_(filename), file_number_(0), call_number_(0),
    skip_steps_(0)
{
}

void CMMLogOptimizerState::update()
{
  //  std::cout << " CMMLogOptimizerState::update skip_steps_: " << skip_steps_ << " call_number_: " << call_number_ << std::endl;
  if (skip_steps_ == 0 || (call_number_ % skip_steps_) == 0) {
    char buf[1000];
    sprintf(buf, filename_.c_str(), file_number_);
    ++file_number_;
    write(buf);
  }
  ++call_number_;
}

void CMMLogOptimizerState::write(std::string buf) const
{
  std::ofstream out(buf.c_str());
  if (!out) {
    IMP_WARN("Error opening CMM log file " << buf);
  } else {
    IMP_LOG(VERBOSE, "Writing " << pis_.size()
            << " particles to file " << buf << "..." << std::flush);
    std::stringstream filename;
    filename << "optimization_step_" << call_number_;
    write(pis_, filename.str() ,radius_, r_, g_, b_, out);
    //IMP_LOG(TERSE, "done" << std::endl);
  }
}

static Float snap(Float f)
{
  if (f < 0) return 0;
  if (f > 1) return 1;
  return f;
}

void CMMLogOptimizerState::write(const Particles &pis,
				 const std::string &marker_set_name,
				 FloatKey radius_key,
				 FloatKey r_key, FloatKey g_key, FloatKey b_key,
				 std::ostream &out)
{
  out << "<marker_set name=\"" <<marker_set_name << "\">"<<std::endl;
  for (unsigned int i = 0; i < pis.size(); ++i) {
    try {
      Particle *p = pis[i];
      XYZDecorator xyz = XYZDecorator::cast(p);
      float x = xyz.get_x();
      float y = xyz.get_y();
      float z = xyz.get_z();
      Float rv = 0., gv = 0., bv = 0.;
      if (r_key != FloatKey() && b_key != FloatKey() && g_key != FloatKey()
          && p->has_attribute(r_key) && p->has_attribute(g_key)
          && p->has_attribute(b_key)) {
        rv = snap(p->get_value(r_key));
        gv = snap(p->get_value(g_key));
        bv = snap(p->get_value(b_key));
      }
      Float radius = 5.;
      if (radius_key != FloatKey() && p->has_attribute(radius_key)) {
        radius = p->get_value(radius_key);
      }
      out <<"<marker id=\""<<i<<"\""
	" x=\"" <<x <<"\""<<
	" y=\"" <<y <<"\""<<
	" z=\"" <<z <<"\""<<
	" radius=\""<<radius<<"\""<<
	" r=\"" <<rv <<"\""<<
	" g=\"" <<gv <<"\""<<
	" b=\"" <<bv << "\"/>"<<
	std::endl;
    }

    catch (InvalidStateException &e) {
      IMP_WARN("Particle " << pis[i] << " does not have "
               << " cartesian coordinates");
    }
  }
  out << "</marker_set>" << std::endl;
}


void CMMLogOptimizerState::show(std::ostream &out) const
{
  out << "Writing CMM files " << filename_ << std::endl;
}

} // namespace IMP
import unittest
import IMP, IMP.test
import os.path

class TestBL(IMP.test.TestCase):
    def setUp(self):
        IMP.set_log_level(IMP.TERSE)

    def _testit(self, rk, r,g,b, pref):
        """Test logging to a CMM file"""
        m= IMP.Model()
        o= IMP.SteepestDescent()
        o.set_model(m)
        nm="/tmp/"+pref+"cmmtest%03d.cmm"
        p0= IMP.Particle()
        m.add_particle(p0)
        d0= IMP.XYZDecorator.create(p0)
        p0.add_attribute(rk, 1.5, False)
        d0.set_x(0)
        d0.set_y(0)
        d0.set_z(0)

        p1= IMP.Particle()
        m.add_particle(p1)
        d1= IMP.XYZDecorator.create(p1)
        p1.add_attribute(r, 1.0, False)
        p1.add_attribute(g, 0.0, False)
        p1.add_attribute(b, 0.0, False)
        d1.set_x(1)
        d1.set_y(1)
        d1.set_z(1)
        a= IMP.CMMLogOptimizerState(nm, IMP.Particles([p0,p1]))
        a.set_radius(rk)
        a.set_color(r, g, b)
        o.add_optimizer_state(a)
        a.update()

        self.assert_(os.path.isfile("/tmp/"+pref+"cmmtest000.cmm"))
        os.remove("/tmp/"+pref+"cmmtest000.cmm")
    def test_1(self):
        """Testing the CMM log"""
        self._testit(IMP.FloatKey("radius"),
                     IMP.FloatKey("red"),
                     IMP.FloatKey("green"),
                     IMP.FloatKey("blue"), "test1")

    def test_2(self):
        """Testing the CMM log with new attribute names"""
        self._testit(IMP.FloatKey("another_radius"),
                     IMP.FloatKey("red5"),
                     IMP.FloatKey("green5"),
                     IMP.FloatKey("blue5"),
                     "test1")
    def test_skip(self):
        """Test skipping steps in the CMM log"""
        m= IMP.Model()
        o= IMP.SteepestDescent()
        o.set_model(m)
        nm="/tmp/"+"skip"+"cmmtest%03d.cmm"
        p0= IMP.Particle()
        m.add_particle(p0)
        d0= IMP.XYZDecorator.create(p0)
        p0.add_attribute(IMP.FloatKey("radius"), 1.5, False)
        d0.set_x(0)
        d0.set_y(0)
        d0.set_z(0)
        p1= IMP.Particle()
        m.add_particle(p1)
        d1= IMP.XYZDecorator.create(p1)
        d1.set_x(1)
        d1.set_y(1)
        d1.set_z(1)
        a= IMP.CMMLogOptimizerState(nm, IMP.Particles([p0,p1]))
        a.set_skip_steps(20) 
        r= IMP.DistanceRestraint(p0, p1, IMP.Harmonic(0,10));
        m.add_restraint(r);
        o.add_optimizer_state(a)
        o.optimize(11)
        self.assert_(os.path.isfile("/tmp/"+"skip"+"cmmtest000.cmm"))
        self.assert_(os.path.isfile("/tmp/"+"skip"+"cmmtest001.cmm"))
        self.assert_(not os.path.isfile("/tmp/"+"skip"+"cmmtest002.cmm"))
        os.remove("/tmp/"+"skip"+"cmmtest000.cmm")
        os.remove("/tmp/"+"skip"+"cmmtest001.cmm")
        if (os.path.isfile("/tmp/"+"skip"+"cmmtest002.cmm")):
            os.remove("/tmp/"+"skip"+"cmmtest002.cmm")
        print "4"
if __name__ == '__main__':
    unittest.main()