Repository: dask/dask
File: dask/dataframe/core.py (Python)

Code before:

def cov_corr_chunk(df, corr=False):
    """Chunk part of a covariance or correlation computation
    """
    shape = (df.shape[1], df.shape[1])
    df = df.astype("float64", copy=False)
    sums = np.zeros_like(df.values, shape=shape)
    counts = np.zeros_like(df.values, shape=shape)
    for idx, col in enumerate(df):
        mask = df.iloc[:, idx].notnull()
        sums[idx] = df[mask].sum().values
        counts[idx] = df[mask].count().values
    cov = df.cov().values
    dtype = [("sum", sums.dtype), ("count", counts.dtype), ("cov", cov.dtype)]
    if corr:
        with warnings.catch_warnings(record=True):
            warnings.simplefilter("always")
            mu = (sums / counts).T
        m = np.zeros_like(df.values, shape=shape)
        mask = df.isnull().values
        for idx, x in enumerate(df):
            # Avoid using ufunc.outer, since it is not supported by cupy
            mu_discrepancy = np.zeros_like(df.values, shape=(len(df), len(mu)))
            for j in range(len(mu)):
                mu_discrepancy[:, j] = np.subtract(df.iloc[:, idx].values, mu[idx, j])
            mu_discrepancy = mu_discrepancy ** 2
            mu_discrepancy[mask] = np.nan
            m[idx] = np.nansum(mu_discrepancy, axis=0)
        m = m.T
        dtype.append(("m", m.dtype))

    out = {}
    out["sum"] = sums
    out["count"] = counts
    out["cov"] = cov * (counts - 1)
    if corr:
        out["m"] = m
    return out


def cov_corr_combine(data_in, corr=False):

    items = ["sum", "count", "cov"]
    if corr:
        items.append("m")

    data = {}
    for k in items:
        data[k] = []
        for d in data_in:
            data[k].append(d[k])

Review comment (suggestion):
I wonder if you might be able to do this instead with slicing with new dimensions, something like this ...

```python
x[:, None] - y[None, :]
```

Code after:

def cov_corr_chunk(df, corr=False):
    """Chunk part of a covariance or correlation computation
    """
    shape = (df.shape[1], df.shape[1])
    df = df.astype("float64", copy=False)
    sums = zeros_like_safe(df.values, shape=shape)
    counts = zeros_like_safe(df.values, shape=shape)
    for idx, col in enumerate(df):
        mask = df.iloc[:, idx].notnull()
        sums[idx] = df[mask].sum().values
        counts[idx] = df[mask].count().values
    cov = df.cov().values
    dtype = [("sum", sums.dtype), ("count", counts.dtype), ("cov", cov.dtype)]
    if corr:
        with warnings.catch_warnings(record=True):
            warnings.simplefilter("always")
            mu = (sums / counts).T
        m = zeros_like_safe(df.values, shape=shape)
        mask = df.isnull().values
        for idx, x in enumerate(df):
            # Avoid using ufunc.outer (not supported by cupy)
            mu_discrepancy = (
                np.subtract(df.iloc[:, idx].values[:, None], mu[idx][None, :]) ** 2
            )
            mu_discrepancy[mask] = np.nan
            m[idx] = np.nansum(mu_discrepancy, axis=0)
        m = m.T
        dtype.append(("m", m.dtype))

    out = {"sum": sums, "count": counts, "cov": cov * (counts - 1)}
    if corr:
        out["m"] = m
    return out


def cov_corr_combine(data_in, corr=False):

    data = {"sum": None, "count": None, "cov": None}
    if corr:
        data["m"] = None

    for k in data.keys():
        data[k] = [d[k] for d in data_in]
        data[k] = np.concatenate(data[k]).reshape((len(data[k]),) + data[k][0].shape)

    sums = np.nan_to_num(data["sum"])
    counts = data["count"]

    cum_sums = np.cumsum(sums, 0)
    cum_counts = np.cumsum(counts, 0)