Search code examples
pythonnumpymathsympysymbolic-math

Error applying special sympy functions with numpy

Im trying to create a graph with numerical solutions of a magnetic field. The object is a solenoid.

import numpy as np
import sympy as smp
from scipy.integrate import quad_vec

t, x, y, z = smp.symbols("t x y z")

l = smp.Matrix([0.5 * smp.cos(t), 
                0.5 * smp.sin(t), 
                (t / (600 * smp.pi)) * smp.Heaviside(t - 1) * smp.Heaviside(t)])

r = smp.Matrix([x, y, z])
sep = r-l

integrand = smp.diff(l, t).cross(sep) / sep.norm()**3

dBxdt = smp.lambdify([t, x, y, z], integrand[0])
dBydt = smp.lambdify([t, x, y, z], integrand[1])
dBzdt = smp.lambdify([t, x, y, z], integrand[2])

def B(x, y, z):
    return np.array([quad_vec(dBxdt, 0, 2*np.pi, args=(x, y, z))[0],
                     quad_vec(dBydt, 0, 2*np.pi, args=(x, y, z))[0],
                     quad_vec(dBzdt, 0, 2*np.pi, args=(x, y, z))[0]])

x = np.linspace(-2, 2, 20)
xv, yv, zv = np.meshgrid(x, x, x)

B_field = B(xv, yv, zv)
Bx, By, Bz = B_field

File <lambdifygenerated-8>:2, in _lambdifygenerated(t, x, y, z)
      1 def _lambdifygenerated(t, x, y, z):
----> 2     return (-(y - 0.5*sin(t))*((1/600)*t*DiracDelta(t)*select([less(t, 1),equal(t, 1),True], [0,1/2,1], default=nan)/pi + (1/600)*t*DiracDelta(t - 1)*select([less(t, 0),equal(t, 0),True], [0,1/2,1], default=nan)/pi + (1/600)*select([less(t, 0),equal(t, 0),True], [0,1/2,1], default=nan)*select([less(t, 1),equal(t, 1),True], [0,1/2,1], default=nan)/pi) + 0.5*(-1/600*t*select([less(t, 0),equal(t, 0),True], [0,1/2,1], default=nan)*select([less(t, 1),equal(t, 1),True], [0,1/2,1], default=nan)/pi + z)*cos(t))/(abs(x - 0.5*cos(t))**2 + abs(y - 0.5*sin(t))**2 + abs((1/600)*t*select([less(t, 0),equal(t, 0),True], [0,1/2,1], default=nan)*select([less(t, 1),equal(t, 1),True], [0,1/2,1], default=nan)/pi - z)**2)**(3/2)

NameError: name 'DiracDelta' is not defined
Solution

  • I'm going to use SymPy Plotting Backends.

    The first thing we can do is to rewrite Heaviside with Piecewise. This will create visually longer expressions, but they will evaluate much faster with mpmath.

    I'm going to create a 3D streamline plot. Evaluation with NumPy is going to fail because it doesn't know what DiracDelta is. But evaluation with mpmath will succeed. However, mpmath doesn't support vectorized operations, so evaluation will be slow.

    from spb import *

# rewrite integrand's Heaviside with Piecewise
a = smp.Integral(integrand[0].rewrite(smp.Piecewise), (t, 0, 2*smp.pi))
b = smp.Integral(integrand[1].rewrite(smp.Piecewise), (t, 0, 2*smp.pi))
c = smp.Integral(integrand[2].rewrite(smp.Piecewise), (t, 0, 2*smp.pi))

graphics(
    vector_field_3d(
        a, b, c, (x, -2, 2), (y, -2, 2), (z, -2, 2),

        # `a, b, c` are complicated expressions, whose
        # latex representation would make the backend fail.
        # let's set an easier label
        label="a", 

        # number of discretization points on each direction.
        # Remember you are discretizing a volume, so you will have
        # n^3 evaluation points. The higher this number the slower
        # the evaluation.
        n=10,

        # evaluates the expressions with mpmath.
        # mpmath appears to be able to deal with `DiracDelta`, and
        # it is usually faster than SymPy at evaluating an expression.
        # however, it is much much slower than NumPy, hence n must be
        # kept low
        modules="mpmath",

        # visualize streamlines (True) or quivers (False)
        streamlines=True,
        stream_kw={
            # request random starting locations for the streamlines
            "starts": True
        }
    ),

    # use the plotting library K3D
    backend=KB, 
)

    enter image description here

    It looks like a mess, but with K3D you can quickly add clipping planes (click K3D-Panel -> Controls -> Clipping Planes -> Add New).

    enter image description here

    If you look closely, all streamlines appears to lay on some meridian plane. So, we can chose a 2D plane, for example the plane XZ (y=0):

    a = smp.Integral(integrand[0].rewrite(smp.Piecewise).subs(y, 0), (t, 0, 2*smp.pi))
b = smp.Integral(integrand[1].rewrite(smp.Piecewise).subs(y, 0), (t, 0, 2*smp.pi))
c = smp.Integral(integrand[2].rewrite(smp.Piecewise).subs(y, 0), (t, 0, 2*smp.pi))

from spb import *
graphics(
    vector_field_2d(
        a, c, (x, -2, 2), (z, -2, 2),
        label="Magnitude",
        modules="mpmath",
        
        # do not visualize a contour plot with magnitude of the
        # vector field, as it evaluate over 100^2 points, which
        # is too slow for this complicated expression
        scalar=False,
        
        # visualize streamlines (True) or quivers (False)
        streamlines=True,
        
        # streamlines benefit from higher number of discretization points
        n=30
    ),
    backend=MB, aspect="equal", grid=False,
    xlabel="x", ylabel="z", title=r"$\vec{B}(x, 0, z)$"
)

    enter image description here