magnetic_field.py

from typing import Tuple
import numpy as np
from mayavi import mlab
from mayavi.modules.iso_surface import IsoSurface
from mayavi.modules.streamline import Streamline
from scipy import special


class MagneticField(object):
    def __init__(self, radius: float = 0.1, current: float = 1.0, spacing=50j) -> None:
        self.radius = radius
        self.current = current
        self.Lx = self.radius * 4
        self.Ly = self.radius * 4
        self.Lz = self.radius * 4
        self.sp = spacing  # grid spacing (complex number = inclusive bounds

    def sphere_el(self, radius: float, current: float) -> None:

        B = self.compute_all_coils(radius, current)

        objs = self.scene_setup(B[0], B[1], B[2], seedtype="sphere")
        # objs["streamlines"].seed.widget.phi_resolution = 10
        # objs["streamlines"].seed.widget.theta_resolution = 10
        objs["streamlines"].seed.widget.radius = radius

        MagneticField.scene_style(objs)

    def plane_el(self, radius: float, current: float) -> None:

        B = self.compute_all_coils(radius, current)

        objs = self.scene_setup(B[0], B[1], B[2], seedtype="plane")
        objs["streamlines"].stream_tracer.maximum_propagation = 40.0
        objs["streamlines"].seed.widget.resolution = 20
        objs["streamlines"].seed.widget.handle_size = 0.5
        objs["streamlines"].seed.widget.representation = "outline"

        MagneticField.scene_style(objs)

    def line_el(self, radius: float, current: float) -> None:

        B = self.compute_all_coils(radius, current)

        objs = self.scene_setup(B[0], B[1], B[2], seedtype="line")
        objs["streamlines"].stream_tracer.maximum_propagation = 150
        objs["streamlines"].seed.widget.resolution = 30
        # objs["streamlines"].seed.widget.point1 = [95, 100.5, 100]  # placing seed
        # objs["streamlines"].seed.widget.point2 = [105, 100.5, 100]

        MagneticField.scene_style(objs)

    def compute_all_coils(self, radius: float, current: float) -> np.ndarray:
        self.__init__(radius, current)
        x, y, z = self.get_grid("sparse")
        # Get the grid spacing. Normally defined as a complex number since
        # mgrid and ogrid need that for the interval to be inclusive on the upper
        # bound.
        sp = self.sp
        if int(sp.imag) != 0:
            sp = int(sp.imag)
        else:
            sp = int(sp.real)

        h = radius / 2.0
        fig = mlab.figure(1, size=(800, 600), bgcolor=(1, 1, 1), fgcolor=(0, 0, 0))

        B = np.zeros((3, sp, sp, sp))
        # TODO: add to a loop
        B += MagneticField.magnetic_field_single_coil(
            x,
            y,
            z,
            radius,
            current,
            centre=[0, 0, +h],
        )
        B += MagneticField.magnetic_field_single_coil(
            x,
            y,
            z,
            radius,
            current,
            centre=[0, 0, -h],
        )
        MagneticField.draw_coil(
            radius,
            name="Coil 1",
            color=(0, 0, 1),
            centre=[0, 0, +h],
        )
        MagneticField.draw_coil(
            radius,
            name="Coil 2",
            color=(0, 1, 1),
            centre=[0, 0, -h],
        )
        del x, y, z

        return B

    @staticmethod
    def draw_coil(
        radius: float,
        name="Coil",
        color=(0, 0, 1),
        centre: np.array = np.array([0, 0, 0]),
        scale: np.array = np.array([1, 1, 1]),
    ):
        l = 40
        theta = np.linspace(0, 2 * np.pi, l)
        y = ((radius * np.sin(theta)) + centre[0]) * scale[0]
        x = ((radius * np.cos(theta)) + centre[1]) * scale[1]
        z = (np.zeros(l) + centre[2]) * scale[2]
        return mlab.plot3d(
            x,
            y,
            z,
            name=name,
            tube_radius=radius * 0.1,
            color=color,
        )

    @staticmethod
    def magnetic_field_single_coil(
        x: np.ndarray,
        y: np.ndarray,
        z: np.ndarray,
        radius: float,
        current: float,
        centre: np.array = np.array([0.0, 0.0, 0.0]),
        scale: np.array = np.array([1.0, 1.0, 1.0]),
        normalise: bool = False,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Calculate the magnetic field of a single coil loop

        The equations in Cylindrical coordinates are given by:
        .. math::
        B_z = \frac{\mu_0I}{2\pi}\frac{1}{[(R + \rho)^2 + z^2]^{1/2}} \times
            \left[K(k^2)+\frac{R^2 - \rho^2 - z^2}{(R-\rho)^2 + z^2}E(k^2)\right]
        B_\rho = \frac{\mu_0I}{2\pi\rho} \frac{z}{[R+\rho)^2 +z^2]^{1/2}} \times
            \left[-K(k^2) + E(k^2)\frac{R^2+\rho^2+z^2}{(R-\rho)^2+z^2}\right]

        where K and E are the elliptic integrals

        Args:
            x (np.ndarray): generated by `np.mgrid`
            y (np.ndarray): generated by `np.mgrid`
            z (np.ndarray): generated by `np.mgrid`
            radius (float): radius of the coil
            current (float): current through the coil
            normalise (bool, optional): normalise field. Defaults to False.

        Returns:
            Tuple[np.ndarray, np.ndarray, np.ndarray]: Magnetic field in Cartesian coordinates
        """

        R = radius
        I = current
        if abs(radius) < 1e-10:
            R = 1e-10
        if abs(current) < 1e-10:
            I = 1e-10

        x = (x - centre[0]) * scale[0]
        y = (y - centre[1]) * scale[1]
        z = (z - centre[2]) * scale[2]

        rho = np.sqrt(x ** 2 + y ** 2)

        mu = 4 * np.pi * 10.0 ** (-7)  # μ0 constant
        Bz_norm_factor = 1
        Brho_norm_factor = 1
        if normalise:
            Bz_norm_factor = mu / (2 * np.pi)
            Brho_norm_factor = Bz_norm_factor / rho

        # Special ellipse E and K
        E = special.ellipe((4 * R * rho) / ((R + rho) ** 2 + z ** 2))
        K = special.ellipk((4 * R * rho) / ((R + rho) ** 2 + z ** 2))
        Bz = (
            I
            * Bz_norm_factor
            / (np.sqrt((R + rho) ** 2 + z ** 2))
            * (K + (R ** 2 - rho ** 2 - z ** 2) / ((R - rho) ** 2 + z ** 2) * E)
        )
        Brho = (
            I
            * Brho_norm_factor
            * z
            / (rho * np.sqrt((R + rho) ** 2 + z ** 2))
            * (-K + (R ** 2 + rho ** 2 + z ** 2) / ((R - rho) ** 2 + z ** 2) * E)
        )

        # At origin we get division by zero which yields NaN, physically the
        # field is zero in that location
        Brho[np.isnan(Brho)] = 0
        Brho[np.isinf(Brho)] = 0
        Bz[np.isnan(Bz)] = 0
        Bz[np.isinf(Bz)] = 0

        Bx, By = (x / rho) * Brho, (y / rho) * Brho
        B = np.array([Bx, By, Bz])
        del Brho, E, K

        return B

    def scene_setup(
        self,
        Bx: np.ndarray,
        By: np.ndarray,
        Bz: np.ndarray,
        seedtype: str,
    ) -> dict:

        # vector_field does not take sparse grids (as far as I know) so we have
        # to generate a dense xyz grid to get the vector field coordinates right
        x, y, z = self.get_grid("dense")
        Bnorm = np.sqrt(Bx ** 2 + By ** 2 + Bz ** 2)
        field = mlab.pipeline.vector_field(x, y, z, Bx, By, Bz, name="B field")
        del x, y, z
        magnitude = mlab.pipeline.extract_vector_norm(field)
        contours: IsoSurface = mlab.pipeline.iso_surface(
            magnitude,
            contours=4,
            transparent=True,
            opacity=0.6,
            colormap="YlGnBu",
            # vmin=0,
            # vmax=0.5,
        )

        streamlines: Streamline = mlab.pipeline.streamline(
            magnitude,
            seedtype=seedtype,
            integration_direction="both",
            transparent=True,
            opacity=0.2,
            colormap="jet",
            # vmin=0,
            # vmax=0.5,
        )

        contours.actor.property.frontface_culling = True
        contours.normals.filter.feature_angle = 90

        streamlines.stream_tracer.maximum_propagation = 100

        pipeline_obj = {
            "field": field,
            "mag": magnitude,
            "contours": contours,
            "streamlines": streamlines,
        }

        return pipeline_obj

    @staticmethod
    @mlab.show
    def scene_style(objs: dict, **kwargs) -> None:
        sc = mlab.scalarbar(
            objs["streamlines"],
            title="Field\nStrength [T]",
            orientation="vertical",
            nb_labels=4,
        )
        # horizontal and vertical position from left->right, bottom->top
        sc.scalar_bar_representation.position = np.array([0.85, 0.1])
        # width and height of colourbar
        sc.scalar_bar_representation.position2 = np.array([0.1, 0.8])
        # The title of colourbar does not scale so we need to manually set it
        sc.scalar_bar.unconstrained_font_size = True
        sc.title_text_property.font_size = 20
        sc.label_text_property.font_size = 20
        ax = mlab.axes()
        mlab.view(azimuth=42, elevation=73)

    def get_grid(
        self, mat_type: str = "sparse"
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Return either a dense or a sparse 3D grid

        Args:
            mat_type (str, optional): dense or sparse. Defaults to "sparse".

        Returns:
            Tuple[np.ndarray, np.ndarray, np.ndarray]: x, y, z
        """

        if mat_type == "sparse":
            x, y, z = np.ogrid[
                -self.Lx : self.Lx : self.sp,
                -self.Ly : self.Ly : self.sp,
                -self.Lz : self.Lz : self.sp,
            ]
            return x, y, z
        elif mat_type == "dense":
            x, y, z = np.mgrid[
                -self.Lx : self.Lx : self.sp,
                -self.Ly : self.Ly : self.sp,
                -self.Lz : self.Lz : self.sp,
            ]
            return x, y, z
        else:
            raise (f"Input grid format: {mat_type} not supported")


if __name__ == "__main__":
    plot = MagneticField()
    plot.sphere_el(0.1, 0.01)
    plot.plane_el(1.0, 0.1)
    plot.line_el(1.0, 0.1)