Plot multiple R^3 shapes defined as a function in Python/plotly-CodePudding

Let's say I have an association [2,1] -> [3,n], assuming [i,j] represents an element in i'th row and j'th column, such as

$f(\begin{bmatrix}x \ y\end{bmatrix}) = \begin{bmatrix}xy^2 & y^2 & ...\x-xy & 1-y & ...\x^2 2y^2 & -2x & ... \\end{bmatrix}$

I want to represent this as n shapes on a single XYZ plot with different colors. So far I managed to do this:

import numpy as np
import plotly.graph_objects as go

def f1(x,y):
    return x * y*y

def f2(x,y):
    return  x - x*y

def f3(x,y):
    return x*x   2*y*y


X=np.linspace(-10, 10, 100)
Y=np.linspace(-10, 10, 100)
Z1 = f1(X,Y)
Z2 = f2(X,Y)
Z3 = f3(X,Y)

fig = go.Figure(data=[
    go.Scatter3d(x = X, y = Y, z = Z1, mode='lines'),
    go.Scatter3d(x = X, y = Y, z = Z2, mode='lines'),
    go.Scatter3d(x = X, y = Y, z = Z3, mode='lines'),
])

fig.show()

This plots the first column, but it is not scalable – I need to define 3 functions for each column. How to do this in a more vectorized way, assuming that I can have many columns in the output array.

CodePudding user response：

I've found a solution. This can be done with a list comprehension:

import numpy as np
import plotly.graph_objects as go

def f(x,y):
    return [
        [ x * y*y, x - x*y, x*x   2*y*y ],
        [ y*y, 1 - y, -2 * x]
    ]


X=np.linspace(-10, 10, 100)
Y=np.linspace(-10, 10, 100)

plots = [
    [ go.Scatter3d(x = X, y = Y, z = Z, mode='lines') for Z in V ]
    for V in f(X, Y)
]

plots_list = np.reshape(plots, -1).tolist()

fig = go.Figure(data = plots_list)

fig.show()