python - Recombine numpy arrays via original boolean index - Stack Overflow

admin管理员组

文章数量:1345325

I have been working with numpy arrays that require a specific carve, due to a rare area in scipy that doesn't work as planned.

Therefore, I need to use mpmath for that area, and then I want to recombine the array for further operations that depend upon the original array position.

import mpmath as mp
from scipy.special import hyp2f1 as sc_hyp2f1


def carve(z):
    """Scipy hyp2f1 cannot deal with a specific region."""
    return (np.abs(z) > 0.9) & (np.abs(z) < 1.1) & (np.abs(1 - z) >= 0.9) & (z.real >= 0)

def hf2(z, a, b, c):
    """Wrapper to handle hyp2f1"""
    mp_hyp2f1 = np.frompyfunc(mp.hyp2f1, 4, 1)
    f_idx = carve(z)
    zm = z[f_idx]   # values for mpmath (slow!)
    zs = z[~f_idx]  # values for scipy (fast!)
    ma = mp_hyp2f1(a, b, c, zm).astype(npplex128)
    sa = sc_hyp2f1(a, b, c, zs)
    return np.hstack([sa, ma])  # Eep! All the values have been moved!

I have left out the (minimal) logic that avoids attempting to evaluate empty areas.

Hmm I think I have the answer to this..

    ....
    ans = np.zeros_like(z)
    ans[~f_idx] = sa
    ans[f_idx] = ma
    return ans

本文标签： pythonRecombine numpy arrays via original boolean indexStack Overflow