/*****************************************************************************
*                                                                            *
*   CORRCOEF.SPL     Copyright (C) 2021 DSP Development Corporation          *
*                               All Rights Reserved                          *
*                                                                            *
*   Author:      Randy Race                                                  *
*                                                                            *
*   Synopsis:    Calculates the correlation coefficient                      *
*                                                                            *
*   Revisions:   30 Sep 2021  RRR  Creation                                  *
*                                                                            *
*****************************************************************************/

#if @HELP_CORRCOEF

    CORRCOEF

    Purpose: Calculates the correlation coefficient an array

    Syntax:  CORRCOEF(m1, m2)

              m1 - an array or series

              m2 - an optional second array or series

    Returns: An array

    Example:
             a = {{1, 5, 8},
                  {3, 4, 9},
                  {2, 6, 7}};

             b = corrcoef(a);
             c = diag(sqrt(b));
             d = colstdev(a)';

             b == {{ 1.0, -0.5,  0.5},
                   {-0.5,  1.0, -1.0},
                   { 0.5, -1.0,  1.0}}

             c == d

    Remarks:
             The mean is removed from each column before the
             correlation is computed.

             The standard deviations of each column can be
             calculated by:

                      diag(sqrt(corrcoef(m)))

    See Also:
             *^
             Colstdev
#endif


/* calculate correlation matrix */
corrcoef(m1, m2)
{
        local coef, d, j;

        if (argc < 2)
        {
                if (argc < 1) error(sprintf("%s - input series or array required", __FUNC__));
        }
        else
        {
                if (not(all(size(m1) == size(m2))))
                {
                        error(sprintf("%s - input sizes do not match", __FUNC__));
                }

                m1 = ravel(unravel(m1), unravel(m2));
        }

        if (numcols(m1) == 1 || numrows(m1) == 1)
        {
                if (any(isinf(m1) || isnan(m1)))
                {
                        coef = {nan};
                }
                else
                {
                        coef = {1};
                }
        }
        else
        {
                try
                {
                        /* VSL */
                        coef = correlation(m1);
                }
                catch
                {
                        coef = covm(m1) / std(m1);

                        d = diag(coef);
                        setvunits(d, "No Units");

                        coef = coef / d;
                }
        }

        if (outargc > 1)
        {
                if (iscomplex(coef))
                {
                        error(sprintf("%s - P value not valid for complex input", __FUNC__));
                }

                pval = corrcoef_pval(length(m1), coef);

                /* attributes */
                loop (j = 1..numcols(pval))
                {
                        setcomment(pval, getcomment(coef, 1, j), 1, j);
                }

                return(coef, pval);
        }
        else
        {
                return(coef);
        }
}


/* pval statistic */
corrcoef_pval(N, r)
{
        local df, rsqr, tsqr, pval;

        /* dof */
        df = N - 2;

        /* significance */
        rsqr = r * r;
        tsqr = rsqr * ((N - 2) / (1 - rsqr));

        pval = betainc(df / (df + tsqr), 0.5 * df, 0.5);

        if (numcols(pval) > 1)
        {
                pval += eye(pval);
        }

        pval[find(isnan(pval))] = 0;

        return(pval);
}