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

Code before:
        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])
        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)

    s1 = cum_sums[:-1]

Review comment (suggestion):
Very slight style preference to create the full dictionary in one shot rather than iteratively build it.

```python
return {
    "sum": sums,
    "count": counts,
    "cov": cov * (counts - 1)
}
```

Code after:
        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)

    s1 = cum_sums[:-1]
    s2 = sums[1:]
    n1 = cum_counts[:-1]
    n2 = counts[1:]
    with np.errstate(invalid="ignore"):
        d = (s2 / n2) - (s1 / n1)
        C = np.nansum(
            (n1 * n2) / (n1 + n2) * (d * d.transpose((0, 2, 1))), 0