<commit_before>from dolo.symbolic.model import *
from dolo.symbolic.symbolic import *
from dolo.compiler.compiler import Compiler
import numpy as np
import numpy.linalg

class MirFacCompiler(Compiler):
    
    def __init__(self,model):
        self.model = model
        # we assume model has already been checked

    def read_model(self):

        dmodel = self.model
        model = dmodel

        def_eqs = [eq for eq in dmodel.equations if 'def' == eq.tags['eq_type']]


        from dolo.misc.misc import map_function_to_expression
        from dolo.symbolic.symbolic import Variable
        def timeshift(v,n):
            if isinstance(v,Variable):
                return v(n)
            else:
                return v

        import sympy

        #### built substitution dict
        def_dict = {}
        for eq in def_eqs:
            v = eq.lhs
            rhs = sympy.sympify( eq.rhs )
            def_dict[v] = rhs
            def_dict[v(1)] = map_function_to_expression( lambda x: timeshift(x,1), rhs)

        new_equations = []
        tbr = []
        for i,eq in enumerate(dmodel.equations) :
            if not ('def' == eq.tags['eq_type']):
                lhs = sympy.sympify( eq.lhs ).subs(def_dict)
                rhs = sympy.sympify( eq.rhs ).subs(def_dict)
                neq = Equation(lhs,rhs).tag(**eq.tags)
                new_equations.append( neq )

        dmodel['equations'] = new_equations
        dmodel.check_consistency()



        f_eqs = [eq for eq in dmodel.equations if eq.tags['eq_type'] in ('f','arbitrage')]
        g_eqs = [eq for eq in dmodel.equations if eq.tags['eq_type'] in ('g','transition')]
        h_eqs = [eq for eq in dmodel.equations if eq.tags['eq_type'] in ('h','expectation')]


        
        states_vars = [eq.lhs for eq in g_eqs]
        exp_vars =  [eq.lhs for eq in h_eqs]
        controls = set(dmodel.variables) - set(states_vars + exp_vars)
        controls = list(controls)

        states_vars = [v for v in dmodel.variables if v in states_vars]
        exp_vars = [v for v in dmodel.variables if v in exp_vars]
        controls = [v for v in dmodel.variables if v in controls]


        # now we remove the left side of equations
        f_eqs = [eq.gap for eq in f_eqs]
        g_eqs = [eq.rhs for eq in g_eqs]
        h_eqs = [eq.rhs for eq in h_eqs]




        g_eqs = [map_function_to_expression(lambda x: timeshift(x,1),eq) for eq in g_eqs]
        h_eqs = [map_function_to_expression(lambda x: timeshift(x,-1),eq) for eq in h_eqs]


    #    sub_list[v] = v.name

        # read complementarity conditions
        compcond = {}
        of_eqs = [eq for eq in dmodel.equations if eq.tags['eq_type'] in ('f','arbitrage')]
        locals = {}
        import sympy
        locals['inf'] = sympy.Symbol('inf')
        for v in model.variables + model.parameters:
            locals[v.name] = v
        import re
        compregex = re.compile('(.*)<=(.*)<=(.*)')
        for eq in of_eqs:
            tg = eq.tags['complementarity']
            [lhs,mhs,rhs] = compregex.match(tg).groups()
            [lhs,mhs,rhs] = [eval(x,locals) for x in [lhs,mhs,rhs] ]
            compcond[mhs] = (lhs,rhs)
        
        complementarities = [compcond[v] for v in controls]

        inf_bounds = [c[0] for c in complementarities]
        sup_bounds = [c[1] for c in complementarities]
        
        data = {
            'f_eqs': f_eqs,
            'g_eqs': g_eqs,
            'h_eqs': h_eqs,
            'controls': controls,
            'states_vars': states_vars,
            'exp_vars': exp_vars,
            'inf_bounds': inf_bounds,
            'sup_bounds': sup_bounds
        }
        return data


    def perturbation_solution(self):
        """
        Returns perturbation solution around the steady-state
        Result is [X,Y,Z] representing :
        s_t = X s_{t-1} + Y e_t
        x_t = Z s_t
        """

        model = self.model
        data = self.read_model()
        
        controls_i = [model.variables.index(v) for v in data['controls'] ]
        states_i = [model.variables.index(v) for v in data['states_vars'] ]

        from dolo.numeric.perturbations import solve_decision_rule
        dr = solve_decision_rule(model)


        A = dr['g_a'][states_i,:]
        B = dr['g_a'][controls_i,:]

        [Z,err,rank,junk] = np.linalg.lstsq(A.T,B.T)
        Z = Z.T

        X = A[:,states_i] + np.dot( A[:,controls_i], Z )
        Y = dr.ghu[states_i,:]

        # We also compute bounds for distribution at a given probability interval

        M = np.linalg.solve( 1-X, Y)

        Sigma = np.array(model.covariances).astype(np.float64)
        [V,P] = np.linalg.eigh(Sigma)
        # we have Sigma == P * diag(V) * P.T
        # unconditional distribution is ( P*diag(V) ) * N
        # where N is the normal distribution associated
        H = np.dot(Y,np.dot(P,np.diag(V)))
        n_s = Sigma.shape[0]
        I = np.eye(n_s)
        points = np.concatenate([H,-H],axis=1)
        lam = 2 # this coefficient should be more cleverly defined
        p_infs = np.min(points,axis = 1) * lam
        p_max  = np.max(points,axis = 1) * lam

        # steady_state values
        s_ss = dr['ys'][states_i,]
        x_ss = dr['ys'][controls_i,]

        bounds = np.row_stack([
            s_ss + p_infs,
            s_ss + p_max
        ])<commit_msg>Changes in compiler_mirfac - bugfix : copy original model and cached read_model method - computation of bounds is optional<commit_after>from dolo.symbolic.model import *
from dolo.symbolic.symbolic import *
from dolo.compiler.compiler import Compiler
import numpy as np
import numpy.linalg

class MirFacCompiler(Compiler):
    
    def __init__(self,model):
        self.model = model
        self.__transformed_model__ = None
        # we assume model has already been checked

    def read_model(self):
    
        if self.__transformed_model__:
            return self.__transformed_model__

        dmodel = Model(**self.model) # copy the model

        def_eqs = [eq for eq in dmodel.equations if eq.tags['eq_type'] in ('def', 'definition')]

        from dolo.misc.misc import map_function_to_expression
        from dolo.symbolic.symbolic import Variable
        def timeshift(v,n):
            if isinstance(v,Variable):
                return v(n)
            else:
                return v

        import sympy

        #### build substitution dict
        def_dict = {}
        for eq in def_eqs:
            v = eq.lhs
            rhs = sympy.sympify( eq.rhs )
            def_dict[v] = rhs
            def_dict[v(1)] = map_function_to_expression( lambda x: timeshift(x,1), rhs)

        new_equations = []
        tbr = []
        for i,eq in enumerate(dmodel.equations) :
            if not ('def' == eq.tags['eq_type']):
                lhs = sympy.sympify( eq.lhs ).subs(def_dict)
                rhs = sympy.sympify( eq.rhs ).subs(def_dict)
                neq = Equation(lhs,rhs).tag(**eq.tags)
                new_equations.append( neq )

        dmodel['equations'] = new_equations
        dmodel.check_consistency()



        f_eqs = [eq for eq in dmodel.equations if eq.tags['eq_type'] in ('f','arbitrage')]
        g_eqs = [eq for eq in dmodel.equations if eq.tags['eq_type'] in ('g','transition')]
        h_eqs = [eq for eq in dmodel.equations if eq.tags['eq_type'] in ('h','expectation')]


        
        states_vars = [eq.lhs for eq in g_eqs]
        exp_vars =  [eq.lhs for eq in h_eqs]
        controls = set(dmodel.variables) - set(states_vars + exp_vars)
        controls = list(controls)

        states_vars = [v for v in dmodel.variables if v in states_vars]
        exp_vars = [v for v in dmodel.variables if v in exp_vars]
        controls = [v for v in dmodel.variables if v in controls]


        # now we remove the left side of equations
        f_eqs = [eq.gap for eq in f_eqs]
        g_eqs = [eq.rhs for eq in g_eqs]
        h_eqs = [eq.rhs for eq in h_eqs]

        g_eqs = [map_function_to_expression(lambda x: timeshift(x,1),eq) for eq in g_eqs]
        h_eqs = [map_function_to_expression(lambda x: timeshift(x,-1),eq) for eq in h_eqs]


    #    sub_list[v] = v.name

        # read complementarity conditions
        compcond = {}
        of_eqs = [eq for eq in dmodel.equations if eq.tags['eq_type'] in ('f','arbitrage')]
        locals = {}
        import sympy
        locals['inf'] = sympy.Symbol('inf')
        for v in dmodel.variables + dmodel.parameters:
            locals[v.name] = v
        import re
        compregex = re.compile('(.*)<=(.*)<=(.*)')
        for eq in of_eqs:
            tg = eq.tags['complementarity']
            [lhs,mhs,rhs] = compregex.match(tg).groups()
            [lhs,mhs,rhs] = [eval(x,locals) for x in [lhs,mhs,rhs] ]
            compcond[mhs] = (lhs,rhs)
        
        complementarities = [compcond[v] for v in controls]

        inf_bounds = [c[0] for c in complementarities]
        sup_bounds = [c[1] for c in complementarities]
        
        data = {
            'f_eqs': f_eqs,
            'g_eqs': g_eqs,
            'h_eqs': h_eqs,
            'controls': controls,
            'states_vars': states_vars,
            'exp_vars': exp_vars,
            'inf_bounds': inf_bounds,
            'sup_bounds': sup_bounds
        }
        
        self.__transformed_model__ = data # cache computation

        return data


    def perturbation_solution(self,with_bounds=False):
        """
        Returns perturbation solution around the steady-state
        Result is [X,Y,Z] representing :
        s_t = X s_{t-1} + Y e_t
        x_t = Z s_t
        """

        model = self.model
        data = self.read_model()
        
        controls_i = [model.variables.index(v) for v in data['controls'] ]
        states_i = [model.variables.index(v) for v in data['states_vars'] ]

        from dolo.numeric.perturbations import solve_decision_rule
        dr = solve_decision_rule(model)


        A = dr['g_a'][states_i,:]
        B = dr['g_a'][controls_i,:]

        [Z,err,rank,junk] = np.linalg.lstsq(A.T,B.T)
        Z = Z.T

        X = A[:,states_i] + np.dot( A[:,controls_i], Z )
        Y = dr.ghu[states_i,:]

        # steady_state values
        s_ss = dr['ys'][states_i,]
        x_ss = dr['ys'][controls_i,]

        # We also compute bounds for distribution at a given probability interval
        if with_bounds:
            M = np.linalg.solve( 1-X, Y)

            Sigma = np.array(model.covariances).astype(np.float64)
            [V,P] = np.linalg.eigh(Sigma)
            # we have Sigma == P * diag(V) * P.T
            # unconditional distribution is ( P*diag(V) ) * N
            # where N is the normal distribution associated
            H = np.dot(Y,np.dot(P,np.diag(V)))
            n_s = Sigma.shape[0]
            I = np.eye(n_s)
            points = np.concatenate([H,-H],axis=1)
            lam = 2 # this coefficient should be more cleverly defined
            p_infs = np.min(points,axis = 1) * lam
            p_max  = np.max(points,axis = 1) * lam

            bounds = np.row_stack([
                s_ss + p_infs,
                s_ss + p_max
            ])
        else:
            bounds = None
            
{'fraction_non_alphanumeric': 0.1151650312221231, 'fraction_numerical': 0.0026761819803746653, 'mean_word_length': 1.9371233953366518, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, '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': '7303511', 'n_tokens_mistral': 3814, 'n_tokens_neox': 3694, 'n_words': 923}