How to plot implicit functions of 3 variables

Question

I want to write a program that takes in strings representing functions, such as "x^2+y^2+z^2=30" and plot it in matplotlib. I want my program to be able to handle general cases. Previously for plotting general 2d implicit functions in matplotlib I was using spb.plot, which handled implicit functions very nicely. Unfortunately I can't find anything that would work for handling a wide variety of 3d functions.

import matplotlib.pyplot as plt
import spb
import numpy as np


def queryMatplotlib3dEq(qVal):

    fig = plt.figure()
    lhs, rhs = qVal.split("=")
    fig = plt.figure()
    ax = plt.axes(projection="3d")

    def f(x, y):
        if(lhs == "z"):
            return eval(rhs)
        else:
            return eval(lhs)

    x = np.linspace(-6, 6, 30)
    y = np.linspace(-6, 6, 30)
    X, Y = np.meshgrid(x, y)
    if(lhs =='z' or rhs == 'z'):
        Z = f(X, Y)
 
    ax.plot_surface(X, Y, Z, cmap='viridis')
    plt.show()


#Can do this
qVal = "z=x**2+y**2"

#Cannot do this
qVal = "x^2+y^2+z^2=30"

queryMatplotlib3dEq(qVal)

Samy

Answer 1

• Use sympy to solve the equation for z and then plot both halves.
• There's a discontinuity at z == 0, so the halves of the sphere don't quite meet.
  • More samples, num, in x = np.linspace(-6, 6, num=1000) gets the halves closer.
  • The color map needs to be reversed for the negative half of the equation, so the colors sync on the plot.
• <string>:1: RuntimeWarning: invalid value encountered in sqrt, which I ignore.
• Note Why is using 'eval' a bad practice?
  • Evaluating a mathematical expression in a string may offer alternate solutions.
• Tested in python 3.11.2, matplotlib 3.7.1, sympy 1.11.1

import sympy as sym
import numpy as np
import matplotlib.pyplot as plt


def f(x, y):
    return eval(qVal)


def queryMatplotlib3dEq(qVal, ax):
    x = np.linspace(-6, 6, num=1000)
    y = np.linspace(-6, 6, num=1000)
    X, Y = np.meshgrid(x, y)
    Z = f(X, Y)
    
    cmap = plt.colormaps.get_cmap('viridis')
    
    if qVal[0] == '-':
        cmap = cmap.reversed()
        
    ax.plot_surface(X, Y, Z, cmap=cmap)


# solve the equation with sympy
x, y, z = sym.symbols('x y z')
eq = sym.Eq(30, x**2+y**2+z**2)

# replace 'sqrt` with 'np.sqrt'
eqs = [v.replace('sqrt', 'np.sqrt') for v in map(str, sym.solve(eq, z))]

# create the figure and axes outside the function, so axes will be reused for each loop of the equation
fig = plt.figure()
ax = plt.axes(projection='3d')

# setting the limits and aspect ensures the sphere doesn't look like an ellipse
ax.set_zlim(-6, 6)
ax.set_xlim(-6, 6)
ax.set_ylim(-6, 6)
ax.set_aspect('equal')

for qVal in eqs:
    print(qVal)
    queryMatplotlib3dEq(qVal, ax)

plt.show()

num=1000

num=40000 (this will eat your system memory)

Gatecrasher