<commit_before>from __future__ import division, print_function

'''Enhanced steady state sampling (Bhatt, <NAME>, 2010) for WE,
originally implemented by <NAME> (2010), re-implemented for this version
of WEMD by <NAME> (February 2011) and updated by <NAME> (May 2011).'''

import os, sys

import logging
log = logging.getLogger('w_steady_state')
import argparse

import wemd, wemdtools
from wemd import Segment
import numpy, operator, itertools
from itertools import izip
from math import ceil, floor, log10

class steady_state(object):
    """Calculate the Transition Probability Matrix for lag time [1]  
    """
    def __init__(self, sim_manager, args):
        self._sim_manager = sim_manager
        self._data_manager = sim_manager.data_manager

        self._system_driver = sim_manager.system
        self.miniter = args.start_iter
        self.maxiter = args.stop_iter
        self.count_adjust = args.count_adjust
        self.debug = None
        if args.ss_h5 is not None:
            import h5py
            self.debug = h5py.File(args.ss_h5,'w')
            
        self.quiet_mode = args.quiet_mode
        self.symmetrize = args.symmetrize

    def get_new_weights(self):                    
        Cij = self.calcCountMatrix()
        if self.debug is not None:        
            self.debug.create_dataset('Cij',data=Cij)
        
        eps = numpy.finfo(type(Cij[0,0])).eps

        #Reduce Cij to remove columns/row which are empty
        #oldindex[Ri] == Ci (ie it maps the reduced i,j to the original i,j which correspond to the original bins)        
        Rij, oldindex = self.reduceCountMatrix(Cij)
        #prevents a problem due to a low # of transitions
        Rij += self.count_adjust

        if self.debug is not None:          
            self.debug.create_dataset('Rij',data=Rij)
            
        Q = self.convCount2ProbMatrix(Rij)
        
        if self.debug is not None:             
            self.debug.create_dataset('Q',data=Q)
                
        region_set = self._system_driver.region_set
    
        #Cij -> i or j is the index into bins
        bins = region_set.get_all_bins()
        
        target_regions = self._system_driver.target_states
        
        flat_target_regions = []
        for target_region in target_regions:
            for ibin in xrange(0,len(bins)):
                if target_region is bins[ibin]:
                    break

            if ibin in oldindex: #it is possible that the target region was removed (ie if no recycling has occurred)
                new_ibin = oldindex.index(ibin)  #this is in terms of the Rij
                flat_target_regions.append(new_ibin)
            
        n_targets = len(flat_target_regions)
        
        bindim = Rij.shape[0]
        
        SSP, ActiveBins = self.steady_state_approximation(Q, bindim, flat_target_regions)
        if self.debug is not None:     
            self.debug.create_dataset('SSP',data=SSP)
        
        #now remap new_weights onto the bin indices
        MappedActiveBins = []
        for ibin in ActiveBins:
            MappedActiveBins.append(oldindex[ibin])

        mapped_new_weights = numpy.zeros(len(bins))
        for i in xrange(0,len(SSP)):
            mapped_new_weights[MappedActiveBins[i]] = SSP[i]
            #else it will be 0
  
        binnorm = 0.0
        for ibin in MappedActiveBins:
            binnorm += bins[ibin].weight
            
        if self.debug is not None:     
            self.debug.create_dataset('binnorm',data=binnorm)
            if MappedActiveBins:
                self.debug.create_dataset('MappedActiveBins',data=MappedActiveBins)
            self.debug.create_dataset('mapped_new_weights_not_norm',data=mapped_new_weights)

        #scale weights so sum is 1.0 for active bins (and 0 for inactive bins)
        for ibin in xrange(0,len(bins)):
            if ibin in MappedActiveBins:
                mapped_new_weights[ibin] *= 1.0/binnorm
            else:
                mapped_new_weights[ibin] = 0.0

        if self.debug is not None:     
            self.debug.create_dataset('final_weights',data=mapped_new_weights)
            self.debug.close()
                
        return mapped_new_weights

    def calcCountMatrix(self):
        """ Calculate the count matrix 
            Returns
            Cij - the count matrix    
        """
        if sys.stdout.isatty() and not self.quiet_mode:
            sys.stdout.write('Calculating Count Matrix\n')
        
        <commit_msg>Fixed regression in w_steady_state.<commit_after>from __future__ import division, print_function

'''Enhanced steady state sampling (Bhatt, Zhang, Zuckerman, 2010) for WE,
originally implemented by <NAME> (2010), re-implemented for this version
of WEMD by <NAME> (February 2011) and updated by <NAME> (May 2011).'''

import os, sys

import logging
log = logging.getLogger('w_steady_state')
import argparse

import wemd, wemdtools
from wemd import Segment
import numpy, operator, itertools
from itertools import izip
from math import ceil, floor, log10

class steady_state(object):
    """Calculate the Transition Probability Matrix for lag time [1]  
    """
    def __init__(self, sim_manager, args):
        self._sim_manager = sim_manager
        self._data_manager = sim_manager.data_manager

        self._system_driver = sim_manager.system
        self.miniter = args.start_iter
        self.maxiter = args.stop_iter
        self.count_adjust = args.count_adjust
        self.debug = None
        if args.ss_h5 is not None:
            import h5py
            self.debug = h5py.File(args.ss_h5,'w')
            
        self.quiet_mode = args.quiet_mode
        self.symmetrize = args.symmetrize

    def get_new_weights(self):                    
        Cij = self.calcCountMatrix()
        if self.debug is not None:        
            self.debug.create_dataset('Cij',data=Cij)
        
        eps = numpy.finfo(type(Cij[0,0])).eps

        #Reduce Cij to remove columns/row which are empty
        #oldindex[Ri] == Ci (ie it maps the reduced i,j to the original i,j which correspond to the original bins)        
        Rij, oldindex = self.reduceCountMatrix(Cij)
        #prevents a problem due to a low # of transitions
        Rij += self.count_adjust

        if self.debug is not None:          
            self.debug.create_dataset('Rij',data=Rij)
            
        Q = self.convCount2ProbMatrix(Rij)
        
        if self.debug is not None:             
            self.debug.create_dataset('Q',data=Q)
                
        region_set = self._system_driver.region_set
    
        #Cij -> i or j is the index into bins
        bins = region_set.get_all_bins()
        
        target_regions = self._system_driver.target_states
        
        flat_target_regions = []
        for target_region in target_regions:
            for ibin in xrange(0,len(bins)):
                if target_region.bin is bins[ibin]:
                    break

            if ibin in oldindex: #it is possible that the target region was removed (ie if no recycling has occurred)
                new_ibin = oldindex.index(ibin)  #this is in terms of the Rij
                flat_target_regions.append(new_ibin)
            
        n_targets = len(flat_target_regions)
        
        bindim = Rij.shape[0]
        
        SSP, ActiveBins = self.steady_state_approximation(Q, bindim, flat_target_regions)
        if self.debug is not None:     
            self.debug.create_dataset('SSP',data=SSP)
        
        #now remap new_weights onto the bin indices
        MappedActiveBins = []
        for ibin in ActiveBins:
            MappedActiveBins.append(oldindex[ibin])

        mapped_new_weights = numpy.zeros(len(bins))
        for i in xrange(0,len(SSP)):
            mapped_new_weights[MappedActiveBins[i]] = SSP[i]
            #else it will be 0
  
        binnorm = 0.0
        for ibin in MappedActiveBins:
            binnorm += bins[ibin].weight
            
        if self.debug is not None:     
            self.debug.create_dataset('binnorm',data=binnorm)
            if MappedActiveBins:
                self.debug.create_dataset('MappedActiveBins',data=MappedActiveBins)
            self.debug.create_dataset('mapped_new_weights_not_norm',data=mapped_new_weights)

        #scale weights so sum is 1.0 for active bins (and 0 for inactive bins)
        for ibin in xrange(0,len(bins)):
            if ibin in MappedActiveBins:
                mapped_new_weights[ibin] *= 1.0/binnorm
            else:
                mapped_new_weights[ibin] = 0.0

        if self.debug is not None:     
            self.debug.create_dataset('final_weights',data=mapped_new_weights)
            self.debug.close()
            
        #assert not (mapped_new_weights < 0).any()
                
        return mapped_new_weights

    def calcCountMatrix(self):
        """ Calculate the count matrix 
            Returns
            Cij - the count matrix    
        """
        if sys.stdout.isatty() and not self.quiet_mode:
            sys.stdout.write('Calculating Count Matrix\n')
        
        
{'fraction_non_alphanumeric': 0.07833698030634573, 'fraction_numerical': 0.008862144420131291, 'mean_word_length': 2.0088874259381173, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1053492', 'n_tokens_mistral': 2808, 'n_tokens_neox': 2672, 'n_words': 698}