z_slices.py

import numpy as np


def sv_main(contour_res=10, num_slices=10):

    in_sockets = [
        ['s', 'contour_res', contour_res],
        ['s', 'num_slices', num_slices]]

    bbox = (-2.5, 2.5)
    xmin, xmax, ymin, ymax, zmin, zmax = bbox * 3
    A = np.linspace(xmin, xmax, contour_res)  # resolution of the contour
    B = np.linspace(xmin, xmax, num_slices)  # number of slices
    A1, A2 = np.meshgrid(A, A)  # grid on which the contour is plotted

    def goursat(x, y, z):
        a, b, c = 0.0, -5.0, 11.8
        return x ** 4 + y ** 4 + z ** 4 + a * (x ** 2 + y ** 2 + z ** 2) ** 2 + b * (x ** 2 + y ** 2 + z ** 2) + c

    verts = []
    fn = lambda x, y, z: goursat(x, y, z)

    def add_slices(X, Y, Z, flist, zdir):
        for xyz in zip(X, Y, Z):
            f = np.vstack(xyz).T
            for v in f:
                vec = v.tolist()
                idx = {'z': 2, 'y': 1, 'x': 0}[zdir]
                vec[idx] = flist[0]
                verts.append(vec)

    for z in B:  # plot contours in the XY plane
        X, Y = A1, A2
        Z = fn(X, Y, z)
        add_slices(X, Y, Z + z, [z], zdir='z')

    for y in B:  # plot contours in the XZ plane
        X, Z = A1, A2
        Y = fn(X, y, Z)
        add_slices(X, Y + y, Z, [y], zdir='y')

    for x in B:  # plot contours in the YZ plane
        Y, Z = A1, A2
        X = fn(x, Y, Z)
        add_slices(X + x, Y, Z, [x], zdir='x')

    # out boilerplate
    out_sockets = [
        ['v', 'Vecs', [verts]],
    ]

    return in_sockets, out_sockets