Correlation Heatmap

// This Pine Script® code is subject to the terms of the Mozilla Public License 2.0 at https://mozilla.org/MPL/2.0/
// © TradingView

//@version=6
indicator("Correlation Heatmap")

// Correlation Heatmap
// v2, 2026.01.09

// This code's style is based on the recommendations from the Pine Script User Manual's Style guide:
//    https://www.tradingview.com/pine-script-docs/writing/style-guide/



import TradingView/ta/12 as TVta
import TradingView/ValueAtTime/2 as TVvt



//#region ———————————————————— Constants and inputs


// Tooltips
string TT_SYM = "A comma-separated list of symbols or ticker IDs with optional spaces."
string TT_COL = (
    "Defines the base colors for the correlation heatmap's color gradient:\n\n "
    + "- The first input sets the color for maximum anticorrelation (corr = -1).\n "
    + "- The second input sets the color for no correlation (corr = 0).\n "
    + "- The third input sets the color for maximum correlation (corr = 1)."
)
string TT_TF = "Timeframe of the sampled returns. Options are '1W' or '1M'."
string TT_MP = (
    "If enabled, the indicator limits the data used in the correlation calculation to the number of "
    + "periods specified in the input field. Otherwise, it uses the maximum number of periods allowed "
    + "by the chart's time span."
)
string TT_PR = "Specifies the number of fractional digits in the displayed correlation values."

// Symbol list, timeframe, and period control inputs
string symbolListInput = input.text_area("AAPL, GOLD, MSFT, BTCUSD, XOM, TLT, DBC, VNQ, SQQQ, VXX, GDX", "Symbol list", TT_SYM)
string tfInput         = input.timeframe("1M", "Returns timeframe", ["1W", "1M"], TT_TF)

// Data limit and precision inputs
bool limitInput      = input.bool(false, "Max periods",        tooltip = TT_MP, inline = "01")
int  maxPeriodsInput = input.int(60,     "         ",   2,     tooltip = TT_MP, inline = "01")
int  precisionInput  = input.int(2,      "Precision",   1, 16, tooltip = TT_PR)

// Gradient color inputs
string GRP1 = "Color gradient"
color negVeryStrongInput = input.color(#fb8c00, "", TT_COL, "01", group = GRP1)
color neutralInput       = input.color(#ffffff, "", TT_COL, "01", group = GRP1)
color posVeryStrongInput = input.color(#00acc1, "", TT_COL, "01", group = GRP1)
//#endregion



//#region ———————————————————— Functions


// @function            Calculates the arithmetic return for the instrument referenced by a specified symbol or ticker
//                      ID based on confirmed prices sampled at aligned periodic points. The function uses spread ticker
//                      IDs formed by the chart's symbol and the specified symbol in its requests to retrieve prices
//                      at times that align with periodic points on the current chart.
// @param symbol        (series string) The symbol of the instrument for which to request aligned price data. For
//                      consistency, the ticker IDs used by the requests include extended hours and dividend adjustment
//                      modifiers, and they inherit additional modifiers from the chart, such as "settlement-as-close"
//                      settings.
// @param timeframe     (series string) The timeframe of the time-aligned data request. The value should represent a
//                      timeframe that is higher than or equal to the chart's timeframe.
// @returns             (float) The arithmetic return of the instrument's prices retrieved at aligned sample points.
getAlignedReturns(series string symbol, series string timeframe) =>
    string chartTicker = ticker.new(syminfo.prefix, syminfo.tickerid)
    string standard1   = ticker.standard(chartTicker)
    string standard2   = ticker.standard(symbol)
    string sumSpread   = ticker.modify(
        ticker.inherit(chartTicker, standard1 + "+" + standard2), session.extended, adjustment.dividends
    )
    string diffSpread = ticker.modify(
        ticker.inherit(chartTicker, standard1 + "-" + standard2), session.extended, adjustment.dividends
    )
    float  currPrice   = close[1]
    float  prevPrice   = nz(close[2], open[1])
    [currSum, prevSum] = request.security(
        sumSpread, timeframe, [currPrice, prevPrice], barmerge.gaps_on, barmerge.lookahead_on
    )
    [currDiff, prevDiff] = request.security(
        diffSpread, timeframe, [currPrice, prevPrice], barmerge.gaps_on, barmerge.lookahead_on
    )
    TVta.changePercent(0.5 * (currSum - currDiff), 0.5 * (prevSum - prevDiff))


// @function            Calculates time-aligned periodic returns for the instruments referenced by an array of symbols,
//                      and stores the results in a matrix. Each row in the matrix corresponds to one of the
//                      specified instruments, and each column corresponds to a successive period where the data for
//                      each instrument is not `na`.
// @param symbols       (array<string>) References an array of symbols representing the instruments for
//                      which to request aligned price data. For consistency, each ticker ID used in the data requests
//                      includes extended hours and dividend adjustment modifiers, and it inherits additional modifiers
//                      from the chart, such as "settlement-as-close" settings.
// @param timeframe     (simple string) The timeframe of the time-aligned data request. The value must represent a
//                      timeframe that is greater than or equal to the chart's timeframe.
// @param maxPeriods    (simple int) Optional. The maximum number of periodic returns to store for each requested
//                      dataset. If the value is `na`, the maximum length of the data is the span of the current chart.
//                      The default is `na`.
// @returns             (matrix<float>) The ID of a matrix containing aligned periodic returns for each requested
//                      dataset.
collectReturns(array<string> symbols, simple string timeframe, simple int maxPeriods = na) =>
    var int           numSymbols   = symbols.size()
    var array<float>  returnsArray = array.new<float>(numSymbols)
    var matrix<float> result       = matrix.new<float>(numSymbols, 0)
    bool exclude = false
    for [i, symbol] in symbols
        float priceReturn = getAlignedReturns(symbol, timeframe)
        if na(priceReturn)
            exclude := true
            break
        returnsArray.set(i, priceReturn)
    int cols = result.columns()
    if not exclude
        result.add_col(cols, returnsArray)
        cols += 1
    if cols > maxPeriods
        result.remove_col(0)
    result


// @function            Calculates the Pearson correlation coefficient between the elements stored in two "float"
//                      arrays. The value represents the linear relationship between the arrays' elements. The
//                      coefficient is a measure of the covariance between the arrays, normalized by the product of
//                      their standard deviaitions.
// @param id1           (array<float>) References the first array to compare in the calculation. The array's size must
//                      match the size of the `id2` array.
// @param id2           (array<float>) References the second array to compare in the calculation. The array's size must
//                      match the size of the `id1` array.
// @returns             (float) The correlation coefficient between the `id1` and `id2` arrays' elements. The
//                      coefficient ranges from -1 to 1. A value of 1 indicates perfect positive correlation, -1
//                      indicates perfect negative correlation (anticorrelation), and 0 indicates no correlation.
arrayCorrelation(array<float> id1, array<float> id2) =>
    float covariance = array.covariance(id1, id2)
    float std1 = id1.stdev()
    float std2 = id2.stdev()
    float result = covariance / (std1 * std2)
//#endregion



//#region ———————————————————— Calculations and display


// @variable References an array of symbol strings representing the instruments for which to calculate correlation.
//           The script converts all letters in the `symbolListInput` to uppercase, splits the input value by its
//           commas, then trims all whitespaces from each item.
var array<string> symbolsArray = TVvt.getArrayFromString(str.upper(symbolListInput))

// @variable The number of symbol strings in the array referenced by `symbolsArray`.
var int numSymbols = symbolsArray.size()
// Calculate the height and width of each table cell, including headers, to evenly distribute cells within the pane.
var float cellWidth  = 95.0 / (numSymbols + 1)
var float cellHeight = 90.0 / (numSymbols + 1)

// @variable References a matrix containing monthly or weekly returns for each specified instrument.
//           Each row corresponds to an instrument (e.g., row 0 stores data for the instrument referenced by the first
//           symbol), and each column corresponds to a period (e.g., column 0 stores the oldest period's returns).
matrix<float> symbolReturns = collectReturns(symbolsArray, tfInput, limitInput ? maxPeriodsInput : na)

// @variable References a table with `numSymbols + 1` rows and columns for the correlation heatmap display.
var table heatmap = table.new(
    position.middle_center, numSymbols + 1, numSymbols + 1, border_color = #8f93a1, border_width = 1
)

// Logic to initialize the table's header cells on the first bar.
if barstate.isfirst
    // Initialize the top-left cell with blank text.
    heatmap.cell(0, 0, "", cellWidth, cellHeight, chart.fg_color, bgcolor = chart.bg_color)
    // Initialize the leftmost and top cell with symbol names.
    for [index, symbol] in symbolsArray
        heatmap.cell(0, index + 1, symbol, cellWidth, cellHeight, chart.fg_color, bgcolor = chart.bg_color)
        heatmap.cell(index + 1, 0, symbol, cellWidth, cellHeight, chart.fg_color, bgcolor = chart.bg_color)

// Logic to display correlation and length data on the last bar.
if barstate.islast
    // @variable The number of periods represented by the `symbolReturns` matrix.
    int numPeriods = symbolReturns.columns()
    // Update the table's first cell to show the number of analyzed periods.
    heatmap.cell_set_text(0, 0, str.format("Length:\n{0} {1}", numPeriods, tfInput == "1M" ? "months" : "weeks"))
    // Nested loops to calculate correlations and set the `heatmap` cells.
    for i = 0 to numSymbols - 1
        for j = 0 to i
            // Initialize the cell with a predefined text and background color if `j` and `i` are the same.
            if j == i
                string cellText = numPeriods < 2 ? "NaN" : "1"
                color  bgColor  = numPeriods < 2 ? color(na) : posVeryStrongInput
                heatmap.cell(i + 1, j + 1, cellText, cellWidth, cellHeight, #000000, bgcolor = bgColor)
                continue
            // @variable References an array of returns for the `i`-th listed instrument.
            array<float> returns1 = symbolReturns.row(i)
            // @variable References an array of returns for the `j`-th listed instrument.
            array<float> returns2 = symbolReturns.row(j)
            // @variable The correlation coefficient between the values in the `returns1` and `returns2` arrays.
            float correlation = arrayCorrelation(returns1, returns2)
            // @variable A string representing the correlation value, rounded to `precisionInput` fractional digits.
            string cellText = str.tostring(correlation, "0." + str.repeat("0", precisionInput))
            // @variable The cell's background color, which uses a polarized gradient based on the correlation's sign.
            color bgColor = switch math.sign(correlation)
                1 => color.from_gradient(correlation,  0, 1, neutralInput, posVeryStrongInput)
                =>   color.from_gradient(correlation, -1, 0, negVeryStrongInput, neutralInput)
            // Initialize the cells at `i + 1, j + 1`, and `j + 1, i + 1` using `cellText` and `bgColor`.
            heatmap.cell(i + 1, j + 1, cellText, cellWidth, cellHeight, #000000, bgcolor = bgColor)
            heatmap.cell(j + 1, i + 1, cellText, cellWidth, cellHeight, #000000, bgcolor = bgColor)

// Raise an error if the `tfInput` represents a timeframe lower than the chart's timeframe.
if timeframe.in_seconds(tfInput) < timeframe.in_seconds()
    runtime.error(
        "For accurate calculations, the requested timeframe must be higher than or equal to the chart's timeframe."
    )
//#endregion