RBC_Mathematica.nb

ClearAll["Global`*"]

AbsoluteTiming[

    Module[
        {
            alpha, beta, vProductivity, mTransition, mTransitionTransposed,
            kss, yss, css, vGridCapital, nGridCapital, nGridProductivity, mOutput,
            tolerance, mValueFunction, mPolicyFunction
        },

        (* 1. Calibration *)
        alpha = 0.333333333333;
        beta = 0.95;

        (* Productivity values *)
        vProductivity = {0.9792, 0.9896, 1.0000, 1.0106, 1.0212};

        (* Transition matrix *)
        mTransition = {
            {0.9727, 0.0273, 0.0000, 0.0000, 0.0000},
            {0.0041, 0.9806, 0.0153, 0.0000, 0.0000},
            {0.0000, 0.0082, 0.9837, 0.0082, 0.0000},
            {0.0000, 0.0000, 0.0153, 0.9806, 0.0041},
            {0.0000, 0.0000, 0.0000, 0.0273, 0.9727}};
        mTransitionTransposed = Transpose[mTransition];

        (*2. Steady State *)
        kss = (alpha beta)^(1 / (1 - alpha));
        yss = kss^alpha;
        css = yss - kss;

        (* 4. We generate the grid of capital *)
        vGridCapital = Range[0.5 kss, 1.5 kss, 0.00001];
        nGridCapital = Length[vGridCapital];
        nGridProductivity = Length[vProductivity];

        (* 5. We pre-build output for each point in the grid *)
        mOutput = Transpose[{vGridCapital^alpha}].{vProductivity};


        (* 6. Compiling the Inner Loop *)
        With[
            {
            (* Using undocumented function GetElement for faster access to Array elements *)
                part = Compile`GetElement,
                beta = beta
            },
            innerLoop = Compile[
                {
                    {mOutput, _Real, 2}, {vGridCapital, _Real, 1},  {nGridCapital, _Integer},
                    {nGridProductivity, _Integer}, {expectedValueFunction, _Real, 2}
                },
                Module[
                    {
                    (* Initializations *)
                        tmpOutput = Table[0., {2}, {nGridCapital}, {nGridProductivity}],
                        valueProvisional
                    },
                    Do[
                        Module[
                            {
                                gridCapitalNextPeriod = 1
                            },
                            Do[
                                Module[
                                    {
                                        valueHighSoFar = -1000.,
                                        capitalChoice = part[vGridCapital, -1],
                                        y = part[mOutput, nCapital, nProductivity]
                                    },
                                    Do[
                                        valueProvisional = (1 - beta) *
                                            Log[Subtract[y, part[vGridCapital, nCapitalNextPeriod]]] +
                                            beta part[expectedValueFunction, nCapitalNextPeriod, nProductivity];
                                        If[valueHighSoFar < valueProvisional,
                                            (
                                                valueHighSoFar = valueProvisional;
                                                capitalChoice = part[vGridCapital, nCapitalNextPeriod];
                                                gridCapitalNextPeriod = nCapitalNextPeriod;
                                            ),
                                            Break[]
                                        ],
                                        {nCapitalNextPeriod, gridCapitalNextPeriod, nGridCapital}
                                    ];
                                    tmpOutput[[1, nCapital, nProductivity]] = valueHighSoFar;
                                    tmpOutput[[2, nCapital, nProductivity]] = capitalChoice;
                                ],
                                {nCapital, nGridCapital}
                            ]
                        ],
                        {nProductivity, nGridProductivity}
                    ];
                    tmpOutput
                ],
                CompilationTarget -> "C",
                "RuntimeOptions" -> "Speed"
            ]
        ];

        (* 7. Value Function Iteration *)
        tolerance = 0.0000001;
        {mValueFunction, mPolicyFunction} =
            FixedPoint[
                innerLoop[
                    mOutput, vGridCapital, nGridCapital, nGridProductivity,
                    Dot[#[[1]], mTransitionTransposed]
                ] &,
                Table[0., {2}, {nGridCapital}, {nGridProductivity}] (* Starting value *),
                SameTest -> Module[
                    {
                        iteration = 1
                    },
                    Module[{dis = Max[Abs[Subtract[#1[[1]], #2[[1]]]]]},
                        If[
                            Mod[iteration, 10] == 0 || iteration == 1,
                            Print["Iteration = ", iteration, " Sup Diff = ", dis]

                        ];
                        iteration++;
                        dis < tolerance
                    ] &
                ]
            ];

        Print["My check = ", mPolicyFunction[[1000, 3]]];
    ];

]