Source code for kvazar.record

# -*- coding: utf-8 -*-
import os
import copy
import yaml
import numpy
import shutil
import logging
from zipfile import ZipFile
from zipfile import ZIP_STORED

from kvazar.fileman import FileManager

ZIP_PATH   = {"coords" : "coords/crd_%d.npy", 
              "prev_coords" : "coords/pcrd_%d.npy",
              "velocity" : "vels/vel_%d.npy",
              "half_velocity" : "vels/hvel_%d.npy",
              "accel" : "accels/acc_%d.npy", 
              "params" : "params/param_%d.txt",
              "T" : "npy_params/T.npy",
              "U" : "npy_params/U.npy",
              "K" : "npy_params/K.npy",
              "J" : "npy_params/J.npy"}

class RecordException(Exception):
    def __init__(self, message):
        super(RecordException, self).__init__(message)

class Record(object):
    """
    It's a class for saving intermediate data of calculation and archive it 
    in one output file. 
    """
    def __init__(self, rec_file, rec_type = 'w'):
        self.rec_file = rec_file
        rf_name, rf_ext = os.path.splitext(rec_file)
        self.state_count = 0        
        self.work_dir = rf_name + "_work_dir"
        self.fm = FileManager()
        
        index_dir = 1        
        while os.path.exists(self.work_dir):
            self.work_dir = rf_name + "_work_dir_%d" % index_dir
            index_dir += 1
        os.mkdir(self.work_dir)
        if rec_type == 'a' and os.path.isfile(rec_file):
            logging.info("Opening %s" % rec_file)
            self.zip_file = ZipFile(rec_file, 'a')
            self.zip_file.extractall(self.work_dir)
            mask = ZIP_PATH["coords"]
            while os.path.exists(self.work_dir + "/" + (mask % self.state_count)):
                self.state_count += 1
            acd_files = self.get_path_to_files_with_ext("ACD")
            if len(acd_files) > 0:
                self.struct = self.fm.load(acd_files[0])
            else:
                raise RecordException("Bad record file. Don't find struct.ACD")
            if len(acd_files) > 1:
                logging.warning("Too many .ACD files in %s." % rec_file)
        elif rec_type == 'w':
            logging.info("Creating %s" % rec_file)
            self.zip_file = ZipFile(rec_file, 'w')
        else:
            raise RecordException('Unknown modifier %s' %rec_type + "\nUse 'r' or 'a'!")
                
    
    def get_path_to_files_with_ext(self, ext):
        import re
        tree = os.walk(self.work_dir)
        p_ext = re.compile('.*[.]%s' % ext)
        
        path_to_files = []
        
        for d in tree:
            ext_files = filter(lambda x : p_ext.match(x), d[2])                        
            for cur_file in ext_files:
                path_to_files.append(d[0] + "/" + cur_file)
        
        return path_to_files
    
    
    def initialize(self, struct, config = None):
        necessary_dirs = ["coords", "vels", "accels", "params"]
        self.set_file(struct)
        if config:
            self.set_config(config)
        map(lambda x : self.fm.ensure_dir(self.work_dir + "/" + x), necessary_dirs)
        logging.info("Initializing of record file was succesfully completed.")
        
    def set_file(self, struct):
        self.struct = copy.deepcopy(struct)
        fm = FileManager()
        fm.set_struct(self.struct)
        pstruct = self.work_dir + "/struct.ACD"
        fm.save(pstruct)
        self.zip_file.write(pstruct, "struct.ACD", ZIP_STORED)
    
    def set_config(self, path_to_config):
        pconf = self.work_dir + "/config.cfg"
        shutil.copy(path_to_config, pconf)
        self.zip_file.write(pconf, "struct.ACD", ZIP_STORED)
    
    def append_state(self, state):
        npy = ["coords", "velocity", "accel", "T", "J", "U", "K", "half_velocity", "prev_coords"]
        yml = ["params"]
        
        def make_params_file(path_txt, monitor_data):
            cur_details = open(path_txt, "w")
            for d in monitor_data.keys():
                data = "%s : %s\n\n" % (d, monitor_data[d])
                cur_details.write(data)            
            cur_details.close()
        
        for parm in state.keys():
            if parm in npy:
                path = ZIP_PATH[parm] % self.state_count
                pparm = self.work_dir + "/" + path
                numpy.save(pparm, state[parm])
                self.zip_file.write(pparm, path, ZIP_STORED)
            elif parm in yml:
                data = state[parm]
                path = ZIP_PATH[parm] % self.state_count
                pdata = self.work_dir + "/" +  path
                make_params_file(pdata, data)
                self.zip_file.write(pdata, path, ZIP_STORED)
                
        self.state_count += 1
    
    def close(self):
        self.zip_file.close()
        shutil.rmtree(self.work_dir)
        
    def update_state(self, number):
        struct_parms = ["coords", "velocity", "accel", "half_velocity", "prev_coords"]
        
        for prm in struct_parms:
            path = self.work_dir + "/" + ZIP_PATH[prm] % number
            if os.path.exists(path):
                setattr(self.struct, prm, numpy.load(path))

        self.params = yaml.load(open((self.work_dir + "/" + ZIP_PATH["params"]) % number))

            
    
    

def main():
    logging.basicConfig(filename='record.log',level=logging.DEBUG)
    
    fm = FileManager()
    struct = fm.load("/home/savegor/projects/kvazar/python/problems/DPPC.ACD")
    a = Record("/home/savegor/projects/kvazar/python/output/lol.nmd")
    a.initialize(struct)
    a.close()
    
    print "Completed"
    
    

    


if __name__=="__main__":
    main()