sdf.py 文件源码

def sphere_filter(center, radius, domain_width):

    def myfilter(chunk, mask=None):
        pos = np.array([chunk['x'], chunk['y'], chunk['z']]).T
        left = center-radius

        # This hurts, but is useful for periodicity. Probably should check
        # first if it is even needed for a given left/right
        for i in range(3):
            pos[:, i] = np.mod(pos[:, i] - left[i], domain_width[i]) + left[i]

        # Now get all particles that are within the radius
        if mask is None:
            mask = ((pos-center)**2).sum(axis=1)**0.5 < radius
        else:
            np.multiply(mask, np.linalg.norm(pos - center, 2) < radius, mask)
        return mask

    return myfilter