<commit_before>"""
Functions for working with the discrete-continuous EGM (DCEGM) algorithm as
described in "The endogenous grid method for discrete-continuous dynamic
choice models with (or without) taste shocks" by Iskhakov et al. (2016)
[https://doi.org/10.3982/QE643 and ijrsDCEGM2017 in our Zotero]

Example can be found in https://github.com/econ-ark/DemARK/blob/master/notebooks/DCEGM-Upper-Envelope.ipynb
"""
import numpy as np
from HARK.interpolation import LinearInterp

def calcCrossPoints(mGrid, condVs, optIdx):
    '''
    Given a grid of m values, a matrix of the conditional values of different
    actions at every grid point, and a vector indicating the optimal action
    at each grid point, this function computes the coordinates of the
    crossing points that happen when the optimal action changes
    
    Parameters
    ----------
    mGrid : np.array
        Market resources grid.
    condVs : np.array must have as many rows as possible discrete actions, and
             as many columns as m gridpoints there are.
        Conditional value functions
        
    optIdx : np.array of indices
        Optimal decision at each grid point
    Returns
    -------
    xing_points: [tuple]
        List of crossing points, each as an (m,v) tuple.

    segments: np.array with two columns and as many rows as xing points.
        Each row represents a crossing point. The first column is the index
        of the optimal action to the left, and the second, to the right.
    '''
    # Compute differences in the optimal index,
    # to find positions of choice-changes
    diff_max = np.insert(np.diff(optIdx),len(optIdx)-1,0)
    idx_change = np.where(diff_max != 0)[0]
    
    # If no crossings, return an empty list
    if len(idx_change) == 0:
        return []
    else:
        
        # To find the crossing points we need the extremes of the intervals in
        # which they happen, and the two candidate segments evaluated in both
        # extremes. switchMs[0] has the left points and switchMs[1] the right
        # points of these intervals.
        switchMs = np.stack([mGrid[idx_change], mGrid[idx_change + 1]],
                            axis = 1)
            
        # Store the indices of the two segments involved in the changes, by
        # looking at the argmax in the switching possitions.
        # Columns are [0]: left extreme, [1]: right extreme,
        # Rows are individual crossing points.
        segments = np.stack([optIdx[idx_change], optIdx[idx_change + 1]],
                            axis = 1)
        
        # Values of both segments on the left point
        left_v  = np.stack([condVs[idx_change, segments[:,0]],
                            condVs[idx_change, segments[:,1]]], axis = 1)
        # and the right point
        right_v = np.stack([condVs[idx_change+1, segments[:,0]],
                            condVs[idx_change+1, segments[:,1]]], axis = 1)
            
        # A valid crossing must have both switching segments well defined at the
        # encompassing gridpoints. Filter those that do not.
        valid = np.logical_and(~np.isnan(left_v).any(axis = 1),
                               ~np.isnan(right_v).any(axis = 1))
            
        segments = segments[valid,:]
        switchMs = switchMs[valid,:]
        left_v   = left_v[valid,:]
        right_v  = right_v[valid,:]
            
        # Find crossing points. Returns a list (m,v) tuples. Keep m's only
        xing_points = [calcLinearCrossing(switchMs[i,:],
                                          left_v[i,:], right_v[i,:])
                       for i in range(segments.shape[0])]
        
        return xing_points, segments

def calcPrimKink(mGrid, vTGrids, Choices):
    '''
    Parameters
    ----------
    mGrid : np.array
        Common m grid
    vTGrids : [np.array], length  = # choices, each element has length = len(mGrid)
        value functions evaluated on the common m grid.
    Choices : [np.array], length  = # choices, each element has length = len(mGrid)
        Optimal choices. In the form of choice probability vectors that must
        be degenerate

    Returns
    -------
    None
    '''

    # Construct a vector with the optimal choice at each m point
    optChoice = np.empty(len(mGrid), dtype = int)
    optChoice[:] = np.nan
    for i in range(len(vTGrids)):
        idx = Choices[i] == 1
        optChoice[idx] = i
        
    return calcCrossPoints(mGrid, vTGrids.T, optChoice)
        

def calcSegments(x, v):
    """<commit_msg>Primary kink comment<commit_after>"""
Functions for working with the discrete-continuous EGM (DCEGM) algorithm as
described in "The endogenous grid method for discrete-continuous dynamic
choice models with (or without) taste shocks" by Iskhakov et al. (2016)
[https://doi.org/10.3982/QE643 and ijrsDCEGM2017 in our Zotero]

Example can be found in https://github.com/econ-ark/DemARK/blob/master/notebooks/DCEGM-Upper-Envelope.ipynb
"""
import numpy as np
from HARK.interpolation import LinearInterp

def calcCrossPoints(mGrid, condVs, optIdx):
    '''
    Given a grid of m values, a matrix of the conditional values of different
    actions at every grid point, and a vector indicating the optimal action
    at each grid point, this function computes the coordinates of the
    crossing points that happen when the optimal action changes
    
    Parameters
    ----------
    mGrid : np.array
        Market resources grid.
    condVs : np.array must have as many rows as possible discrete actions, and
             as many columns as m gridpoints there are.
        Conditional value functions
        
    optIdx : np.array of indices
        Optimal decision at each grid point
    Returns
    -------
    xing_points: [tuple]
        List of crossing points, each as an (m,v) tuple.

    segments: np.array with two columns and as many rows as xing points.
        Each row represents a crossing point. The first column is the index
        of the optimal action to the left, and the second, to the right.
    '''
    # Compute differences in the optimal index,
    # to find positions of choice-changes
    diff_max = np.insert(np.diff(optIdx),len(optIdx)-1,0)
    idx_change = np.where(diff_max != 0)[0]
    
    # If no crossings, return an empty list
    if len(idx_change) == 0:
        return []
    else:
        
        # To find the crossing points we need the extremes of the intervals in
        # which they happen, and the two candidate segments evaluated in both
        # extremes. switchMs[0] has the left points and switchMs[1] the right
        # points of these intervals.
        switchMs = np.stack([mGrid[idx_change], mGrid[idx_change + 1]],
                            axis = 1)
            
        # Store the indices of the two segments involved in the changes, by
        # looking at the argmax in the switching possitions.
        # Columns are [0]: left extreme, [1]: right extreme,
        # Rows are individual crossing points.
        segments = np.stack([optIdx[idx_change], optIdx[idx_change + 1]],
                            axis = 1)
        
        # Values of both segments on the left point
        left_v  = np.stack([condVs[idx_change, segments[:,0]],
                            condVs[idx_change, segments[:,1]]], axis = 1)
        # and the right point
        right_v = np.stack([condVs[idx_change+1, segments[:,0]],
                            condVs[idx_change+1, segments[:,1]]], axis = 1)
            
        # A valid crossing must have both switching segments well defined at the
        # encompassing gridpoints. Filter those that do not.
        valid = np.logical_and(~np.isnan(left_v).any(axis = 1),
                               ~np.isnan(right_v).any(axis = 1))
            
        segments = segments[valid,:]
        switchMs = switchMs[valid,:]
        left_v   = left_v[valid,:]
        right_v  = right_v[valid,:]
            
        # Find crossing points. Returns a list (m,v) tuples. Keep m's only
        xing_points = [calcLinearCrossing(switchMs[i,:],
                                          left_v[i,:], right_v[i,:])
                       for i in range(segments.shape[0])]
        
        return xing_points, segments

def calcPrimKink(mGrid, vTGrids, Choices):
    '''
    Parameters
    ----------
    mGrid : np.array
        Common m grid
    vTGrids : [np.array], length  = # choices, each element has length = len(mGrid)
        value functions evaluated on the common m grid.
    Choices : [np.array], length  = # choices, each element has length = len(mGrid)
        Optimal choices. In the form of choice probability vectors that must
        be degenerate

    Returns
    -------
    kinks: [(mCoord, vTCoor)]
        list of kink points
    segments: [(left, right)]
        List of the same length as kinks, where each element is a tuple
        indicating which segments are optimal on each side of the kink.
    '''

    # Construct a vector with the optimal choice at each m point
    optChoice = np.empty(len(mGrid), dtype = int)
    optChoice[:] = np.nan
    for i in range(len(vTGrids)):
        idx = Choices[i] == 1
        optChoice[idx] = i
        
    return calcCrossPoints(mGrid, vTGrids.T, optChoice)
        

def calcSegments(x, v):
    """
{'fraction_non_alphanumeric': 0.0891089108910891, 'fraction_numerical': 0.009040034438226431, 'mean_word_length': 2.1814447107155086, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '51491', 'n_tokens_mistral': 2669, 'n_tokens_neox': 2571, 'n_words': 965}