Newton Excel Bach, not (just) an Excel Blog

import numpy as np

import matplotlib as mpl
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as axes3d
from matplotlib import colors

import pyxll
from pyxll import xl_func, xl_arg, xl_return, plot

@xl_func
@xl_arg('rtn', 'int')
@xl_return("numpy_array")
def PlotWireFrame(val = 0.05, rtn = 2):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.set_title('Wireframe Plot')

    #getting test data
    X, Y, Z = axes3d.get_test_data(val)

    #drawing wireframe plot
    cb = ax.plot_wireframe(X, Y, Z, cstride=10,
                            rstride=10, color="green")

    plot(cb)
    dat = X, Y, Z
    return dat[rtn] 

@xl_func
@xl_arg('Nodes', 'numpy_array')
@xl_arg('Beams', 'numpy_array<int>')
@xl_arg('CenXYZ', 'numpy_array', ndim=1)
@xl_arg('ViewAng', 'numpy_array', ndim=1)
@xl_arg('DisplayAx', 'bool')
@xl_arg('Xrange', 'float')
@xl_arg('LineWidth', 'float')
def Plot3D(Nodes, Beams, ViewAng = None, CenXYZ = None, Xrange = 0., DisplayAx = False, LineWidth = 1):
    fig = mpl.figure.Figure()
    ax = axes3d.Axes3D(fig)
    
    ax.set_box_aspect([1, 1, 1])
    ax.set_proj_type('persp')

    # Set viewing angle if specified, or create axes with defaults
    try:
        ax = axes3d.Axes3D(fig, azim = ViewAng[0], elev = ViewAng[1]) # roll = ViewAng[2] to be added in next release
    except:
        ax = axes3d.Axes3D(fig)
    
    # Set axis limits
    if Xrange == 0:
        Xrange = (np.max(Nodes[:,0]) - np.min(Nodes[:,0]))/2
    rng = Xrange/2
    
    try:
        X, Y, Z = CenXYZ[0:3]
    except:
        X = (np.max(Nodes[:,0]) + np.min(Nodes[:,0]))/2
        Y = (np.max(Nodes[:,1]) + np.min(Nodes[:,1]))/2 
        Z = (np.max(Nodes[:,2]) + np.min(Nodes[:,2]))/2    
    
    ax.set_xlim3d(X-rng,X+rng)
    ax.set_ylim3d(Y-rng,Y+rng)
    ax.set_zlim3d(Z-rng, Z+rng)
    
    # Read beams and coordinates
    rows = Beams.shape[0]
    mats = np.max(Beams[:,0])
    nummata = np.zeros(mats, dtype = int)
    for i in range(0, mats):
        nummata[i] = np.count_nonzero(Beams[:,0] == i+1)

    colors =['tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'tab:purple', 'tab:brown', 'tab:pink', 'tab:gray', 'tab:olive', 'tab:cyan']
   
    if DisplayAx == False:
        ax.set_axis_off()

    # Create lines for each material
    for mat in range(0, mats):
        rowsm = nummata[mat]    
        rows2 = rowsm*3
        x_line = np.empty(rows2)
        y_line = np.empty(rows2)
        z_line = np.empty(rows2)
        matcol = np.mod(mat, 10)
        col = colors[matcol]
        j = 0
        for i in range(0, rows):
            if Beams[i, 0] == mat+1:
                n1 = Beams[i,1]-1
                n2 = Beams[i, 2]-1
                x_line[j] = Nodes[n1,0]
                x_line[j+1] = Nodes[n2,0]
                x_line[j+2] = None
                y_line[j] = Nodes[n1,1]
                y_line[j+1] = Nodes[n2,1]
                y_line[j+2] = None
                z_line[j] = Nodes[n1,2]
                z_line[j+1] = Nodes[n2,2]
                z_line[j+2] = None
                j = j+3
        
        ax.plot3D(x_line, y_line, z_line, col, linewidth=LineWidth)
    plot(ax)