DyLab_3D_MOT/src/coil_class.py

# -*- coding: utf-8 -*-
"""
Created on Tue Aug 17 15:17:44 2021

@author: Joschka
"""
import numpy as np
import matplotlib.pyplot as plt
import logging as log
from scipy import special as sp

import time

from src import physical_constants as cs


# TODO: Docstrings for every function
# TODO: Implement conventions

class BCoil:
    def __init__(self, HH, distance, radius, layers, windings, wire_width, wire_height, insulation_thickness=0.
                 , is_round=False, winding_offset=False):
        """
        Creates object that represents a configuration of two coils, with symmetry axis = z
        :param HH: HH = 1 --> current in same direction, homogeneous field, HH = -1 --> current in opposite direction, quadrupole field
        :param distance: distance between center of coils
        :param radius: radius to center of coil
        :param layers: number of radial (horizontal) layers
        :param windings: number of axial (vertical) layers
        :param wire_width: width of conductive wire core
        :param wire_height: height of conductive wire core
        :param insulation_thickness: thickness of wire insulation
        :param is_round: True --> Round wire, False --> rectangular wire
        :param winding_offset: layer offset by half winding (especially for round wires)
        """
        if not is_round:
            if winding_offset:
                log.warning('Is there a reason you want to wind a not round wire like that?')
        if is_round:
            if wire_width != wire_height:
                log.error('Wire is round but width != height')

        self.HH = HH
        self.distance = distance * 1e-3
        self.radius = radius * 1e-3
        self.layers = layers
        self.windings = windings
        self.wire_width = wire_width * 1e-3
        self.wire_height = wire_height * 1e-3
        self.insulation_thickness = insulation_thickness * 1e-3
        self.is_round = is_round
        self.winding_offset = winding_offset

    def get_wire_area(self):
        """
        calculates wire area in m^2
        :return: float
        """
        if self.is_round:
            return np.pi * (self.wire_width / 2) ** 2
        return self.wire_width * self.wire_height

    def get_N(self):
        """
        Calculates number of windings
        """
        return self.layers * self.windings

    def get_wire_length(self):
        """
        :return: Approximate length of wire per coil (m)
        """
        return self.get_N() * 2 * self.radius * np.pi

    def get_tot_wire_height(self):
        """ returns wire height incl. insulation"""
        return self.wire_height + 2 * self.insulation_thickness

    def get_tot_wire_width(self):
        """ returns wire width incl. insulation"""
        return self.wire_width + 2 * self.insulation_thickness

    def get_coil_height(self):
        if self.winding_offset:
            return self.get_tot_wire_height() * (self.windings + 0.5)
        return self.get_tot_wire_height() * self.windings

    def get_coil_width(self):
        if self.winding_offset:
            if self.is_round:
                log.info(
                    "Be aware of the fact that this is an idealized situation of coil winding (slope of windings changes each layer)")
                return self.layers * self.get_tot_wire_width() - (self.layers - 1) * (
                            2 - np.sqrt(3)) * self.get_tot_wire_width() / 2  # width is reduced due to winding offset
        return self.get_tot_wire_width() * self.layers

    def get_zmin(self):
        return self.distance / 2 - self.get_coil_height() / 2

    def get_zmax(self):
        return self.distance / 2 + self.get_coil_height() / 2

    def get_radius(self):
        return self.radius

    def get_R_inner(self):
        return self.radius - self.get_coil_width() / 2

    def get_R_outer(self):
        return self.radius + self.get_coil_width() / 2

    def set_R_outer(self, R_outer):
        R_outer *= 1e-3
        self.radius = R_outer - self.get_coil_width() / 2

    def set_R_inner(self, R_inner):
        R_inner *= 1e-3
        self.radius = R_inner + self.get_coil_width() / 2

    def set_d_min(self, d_min):
        d_min *= 1e-3
        self.distance = d_min + self.get_coil_height()

    def set_d_max(self, d_max):
        d_max *= 1e-3
        self.distance = d_max - self.get_coil_height()

    def winding_raster(self):
        """
        generates raster of flowing currents
        :param raster_value: wire height/raster_value is distance between rastered points in one wire
        :return: 2 dim array [[z1,R1], [z2,R2], ...]
        """
        outer_raster = np.zeros((self.get_N(), 2))
        it = 0
        z_start = self.get_zmin() + self.get_tot_wire_height() / 2  # (distance_coils/2 - windings * wire_height/2 + wire_height/2)*1e-3
        R_start = self.get_R_inner() + self.get_tot_wire_width() / 2  # (R_inner + wire_width/2 )

        for xx in range(0, self.layers):
            for zz in range(0, self.windings):
                z_pos = z_start + zz * self.get_tot_wire_height()
                R_pos = R_start + xx * self.get_tot_wire_width()

                # correct for different winding techniques and round wire
                if self.winding_offset:
                    if xx % 2 == 1:
                        z_pos += self.get_tot_wire_height() / 2
                    if self.is_round:
                        R_pos -= xx * ((2 - np.sqrt(3)) * self.get_tot_wire_width() / 2)

                outer_raster[it] = [z_pos, R_pos]
                it += 1

        return outer_raster

    def inner_raster(self, raster_value):
        """
        Gives back inner raster for one wire around 0,0
        Args:
            raster_value (int): if N produces a N x N raster for rectangular and cut out of this for round
        Returns: array containing raster around 0,0 [[z_pos_in_1,r_pos_in_1],...]
        """

        inner_raster = np.zeros((raster_value ** 2, 2))
        it = 0
        for xx_in in range(0, raster_value):
            for zz_in in range(0, raster_value):
                z_pos_in = - self.wire_height / 2 + zz_in * self.wire_height / (raster_value - 1)
                r_pos_in = - self.wire_width / 2 + xx_in * self.wire_width / (raster_value - 1)

                inner_raster[it] = [z_pos_in, r_pos_in]
                it += 1

        # delete points out of round geometry
        if self.is_round:
            delete_list = []
            for i in range(0, len(inner_raster)):
                abs = np.sqrt(inner_raster[i, 0] ** 2 + inner_raster[i, 1] ** 2)
                if abs > self.wire_width / 2:
                    delete_list.append(i)
            inner_raster = np.delete(inner_raster, delete_list, axis=0)

        return inner_raster

    def full_raster(self, raster_value):
        """

        Args:
            raster_value: rastering value

        Returns: [ wire 1:[[z_in1,r_in2], [z_in2,r_in,2],...],wire 2:[ [,] , ]...]...]

        """
        outer_ras = self.winding_raster()
        inner_ras = self.inner_raster(raster_value)
        full_ras = np.zeros((len(outer_ras), len(inner_ras), 2))

        for i in range(0, len(full_ras)):
            full_ras[i] = outer_ras[i] + inner_ras

        return full_ras

    def plot_raster(self, raster_value):
        full_structure = self.full_raster(100) * 1e3
        if self.get_coil_width() > self.get_coil_height():
            extension = self.get_coil_width()
        else:
            extension = self.get_coil_height()
        extension *= 1e3
        plt.figure(77, figsize=(5, 5))
        plt.scatter(full_structure[:, :, 1], full_structure[:, :, 0], linewidths=0.01)
        plt.xlabel("radius [mm]")
        plt.ylabel("z position [mm]")
        plt.xlim(1e3 * self.get_R_inner() - 0.5, 1e3 * self.get_R_inner() + extension + 0.5)
        plt.ylim(1e3 * self.get_zmin() - 0.5, 1e3 * self.get_zmin() + extension + 0.5)

        plt.show()
        plt.close(77)

    def print_info(self):
        print(" ")
        # print(f"{self.get_zmin()}")
        # print(self.__name__)
        print(
            f"HH = {self.HH}, Distance = {self.distance * 1e3} mm, z_min = {self.get_zmin() * 1e3} mm, z_max = {self.get_zmax() * 1e3} mm")
        print(
            f"Radius = {self.radius * 1e3} mm, Radius_inner = {self.get_R_inner() * 1e3} mm, Radius_outer = {self.get_R_outer() * 1e3:.2f} mm")
        print(
            f"Coil width = {self.get_coil_width() * 1e3} mm, Coil height = {self.get_coil_height() * 1e3} mm"
        )
        print(
            f"layers = {self.layers}, windings = {self.windings}, wire_width = {self.wire_width * 1e3}, wire_height = {self.wire_height * 1e3} mm, round wire = {self.is_round} ")
        print(" ")

    def print_basic_info(self):
        print(
            f"HH = {self.HH}, Distance = {self.distance * 1e3} mm, Radius = {self.radius * 1e3} mm,layers = {self.layers}, windings = {self.windings}, round wire = {self.is_round}")

    def B_field_ideal_z(self, I_current, z_arg):
        """ 
        Calculate B-field for ideal point like wires with R_radius and d_distance
        HH = +1 --> Helmholtz configuration, HH = -1 --> Anti Helmholtz configuration
        z_arg in mm
        """
        z_SI = z_arg * 1e-3
        N_windings = self.layers * self.windings

        B = cs.mu_0 * N_windings * I_current / 2 * self.radius ** 2 * (
                1 / (self.radius ** 2 + (z_SI - self.distance / 2) ** 2) ** (3 / 2) + self.HH * 1 / (
                self.radius ** 2 + (z_arg + self.distance / 2) ** 2) ** (3 / 2))
        B *= 1e4  # conversion Gauss
        return B

    @staticmethod
    def k_sq(R_loop, z_loc, r, z):
        """ Argument for elliptic integrals"""
        return (4 * R_loop * r) / ((R_loop + r) ** 2 + (z - z_loc) ** 2)

    @staticmethod
    def B_z_loop(I_current, R_loop, z_loc, r, z):
        """calculate z-component of B-field at position r and z for each individual loop"""
        B_z = 2e-7 * I_current * 1 / ((R_loop + r) ** 2 + (z - z_loc) ** 2) ** (1 / 2) * (
                sp.ellipk(BCoil.k_sq(R_loop, z_loc, r, z)) + sp.ellipe(BCoil.k_sq(R_loop, z_loc, r, z)) * (
                R_loop ** 2 - r ** 2 - (z - z_loc) ** 2) / ((R_loop - r) ** 2 + (z - z_loc) ** 2))
        B_z *= 1e4  # conversion to gauss
        return B_z

    @staticmethod
    def B_r_loop(I_current, R_loop, z_loc, r, z):
        """calculate r-component of B-field at position r and z for each individual loop"""
        B_r = 2e-7 * I_current / r * (z - z_loc) / ((R_loop + r) ** 2 + (z - z_loc) ** 2) ** (1 / 2) * (
                -sp.ellipk(BCoil.k_sq(R_loop, z_loc, r, z)) + sp.ellipe(BCoil.k_sq(R_loop, z_loc, r, z)) * (
                R_loop ** 2 + r ** 2 + (z - z_loc) ** 2) / ((R_loop - r) ** 2 + (z - z_loc) ** 2))
        B_r *= 1e4  # conversion to gauss
        return B_r

    def B_field(self, I_current, x, z, raster=10):
        pass
        """
        Returns B_z along z-axis and B_x along x-axis,
        HH = +1 --> Helmholtz configuration, HH = -1 --> Anti Helmholtz configuration
        """

        x_SI = x * 1e-3
        z_SI = z * 1e-3

        if x[0] <= 0:  # divide array into positive and negative values
            x_neg = np.abs([el for el in x_SI if el < 0])
            x_pos = np.array([el for el in x_SI if el > 0])
            x_zero = np.array([el for el in x_SI if el == 0])

        B_z = np.zeros(len(z))
        B_x_pos = np.zeros(len(x_pos))
        B_x_neg = np.zeros(len(x_neg))
        B_x_zero = x_zero  # 0 in x-array --> field is zero at this position

        calc_raster = self.full_raster(raster)

        if self.get_N() != len(calc_raster):
            log.error("N is not equal length of raster")

        rastering_value = len(calc_raster[0])

        I_current /= rastering_value  # divide current into smaller currents for mapping the whole wire
        # start = time.time()
        for wire in range(0, self.get_N()):
            for ii in range(0, rastering_value):
                # extract position information out of raster
                z_pos = calc_raster[wire, ii, 0]
                r_pos = calc_raster[wire, ii, 1]

                B_z += BCoil.B_z_loop(I_current, r_pos, z_pos, 0, z_SI) + BCoil.B_z_loop(self.HH * I_current,
                                                                                         r_pos, -z_pos, 0, z_SI)
                B_x_pos += BCoil.B_r_loop(I_current, r_pos, z_pos, x_pos, 0) + BCoil.B_r_loop(
                    self.HH * I_current, r_pos, -z_pos, x_pos, 0)
                B_x_neg += BCoil.B_r_loop(I_current, r_pos, z_pos, x_neg, 0) + BCoil.B_r_loop(
                    self.HH * I_current, r_pos, -z_pos, x_neg, 0)

        B_x_neg *= -1  # B_x_neg is pointing in opposite x_direction
        B_x = np.concatenate((B_x_neg, B_x_zero, B_x_pos))
        # end = time.time()
        # print(f"time = {end - start} s")
        return B_z, B_x

    def B_tot_along_axis(self, I_current, x, z, raster=10):
        """
        return B_tot_z, B_tot_x
        Returns B_tot_z along z-axis and B_tot_x along x-axis,
        HH = +1 --> Helmholtz configuration, HH = -1 --> Anti Helmholtz configuration
        """

        # Convert axis to SI units
        x_SI = x * 1e-3
        z_SI = z * 1e-3

        if x[0] <= 0:
            x_neg = np.abs([el for el in x_SI if el < 0])
            x_pos = np.array([el for el in x_SI if el > 0])
            x_zero = np.array([el for el in x_SI if el == 0])

        B_tot_z = np.zeros(len(z))
        B_x_pos = np.zeros(len(x_pos))
        B_x_neg = np.zeros(len(x_neg))
        B_x_zero = x_zero  # If there is a zero in the x-array it is assured that the x component of the field at this point is zero

        B_z_x = np.zeros(len(x_SI))
        # dividing each wire into 10 x 10 raster
        calc_raster = self.full_raster(raster)

        if self.get_N() != len(calc_raster):
            log.error("N is not equal length of raster")

        rastering_value = len(calc_raster[0])

        I_current /= rastering_value  # divide current into smaller currents for mapping the whole wire
        # start = time.time()
        for wire in range(0, self.get_N()):
            for ii in range(0, rastering_value):
                # extract position information out of raster
                z_pos = calc_raster[wire, ii, 0]
                r_pos = calc_raster[wire, ii, 1]

                # z-field along z-axis (x-Field always zero)
                B_tot_z += BCoil.B_z_loop(I_current, r_pos, z_pos, 0, z_SI) + BCoil.B_z_loop(
                    self.HH * I_current, r_pos, -z_pos, 0, z_SI)

                # x-field along x-axis
                B_x_pos += BCoil.B_r_loop(I_current, r_pos, z_pos, x_pos, 0) + BCoil.B_r_loop(
                    self.HH * I_current, r_pos, -z_pos, x_pos, 0)
                B_x_neg += BCoil.B_r_loop(I_current, r_pos, z_pos, x_neg, 0) + BCoil.B_r_loop(
                    self.HH * I_current, r_pos, -z_pos, x_neg, 0)
                # B_z along x-axis:
                B_z_x += BCoil.B_z_loop(I_current, r_pos, z_pos, x_SI, 0) + BCoil.B_z_loop(self.HH * I_current,
                                                                                           r_pos, -z_pos, x_SI, 0)

        B_x_neg *= -1  # B_x_neg is pointing in opposite x_direction

        # B_x_x = np.zeros(len(x_SI))
        B_x_x = np.concatenate((B_x_neg, B_x_zero, B_x_pos))
        B_tot_x = np.sqrt(B_x_x ** 2 + B_z_x ** 2)
        B_tot_z = np.abs(B_tot_z)
        return B_tot_z, B_tot_x

    def B_multiple_3d(self, I_current, x, z, raster=4):
        z_start = self.get_zmin() + self.wire_height / 2
        R_start = self.get_R_inner() + self.wire_width / 2

        x_SI = x * 1e-3
        z_SI = z * 1e-3

        B = np.zeros([len(z_SI), len(x_SI), 2])

        calc_raster = self.full_raster(raster)

        if self.get_N() != len(calc_raster):
            log.error("N is not equal length of raster")

        rastering_value = len(calc_raster[0])
        # TODO: why division by zero?

        I_current /= rastering_value  # divide current into smaller currents for mapping the whole wire
        # start = time.time()
        for el in range(0,len(x_SI)):
            for wire in range(0, self.get_N()):
                for ii in range(0, rastering_value):
                    # extract position information out of raster
                    z_pos = calc_raster[wire, ii, 0]
                    r_pos = calc_raster[wire, ii, 1]

                    # compute z-value of field
                    B[:, el, 1] += BCoil.B_z_loop(I_current, r_pos, z_pos, np.abs(x_SI[el]),
                                                  z_SI) + BCoil.B_z_loop(self.HH * I_current, r_pos, -z_pos,
                                                                         np.abs(x_SI[el]), z_SI)

                    # compute x-value
                    if x_SI[el] < 0:
                        B[:, el, 0] -= BCoil.B_r_loop(I_current, r_pos, z_pos, -x_SI[el],
                                                      z_SI) + BCoil.B_r_loop(self.HH * I_current, r_pos, -z_pos,
                                                                             -x_SI[el], z_SI)

                    elif x_SI[el] > 0:
                        B[:, el, 0] += BCoil.B_r_loop(I_current, r_pos, z_pos, x_SI[el], z_SI) + BCoil.B_r_loop(
                            self.HH * I_current, r_pos, -z_pos, x_SI[el], z_SI)

                    elif x_SI[el] == 0:
                        B[:, el, 0] = 0

        return B

    @staticmethod
    def B_tot_3d(B):
        return np.sqrt(B[:, :, 0] ** 2 + B[:, :, 1] ** 2)

    def plot_3d(self, I_current, x_lim, z_lim):
        x = np.arange(-x_lim, x_lim, 5)
        z = np.arange(-z_lim, z_lim, 5)
        x_m, z_m = np.meshgrid(x, z)

        B = self.B_multiple_3d(I_current, x, z, raster=2)
        B_tot = BCoil.B_tot_3d(B)

        for xx in range(0, len(x)):
            for zz in range(0, len(z)):
                if B_tot[xx, zz] > 15:
                    B[xx, zz, :] /= B_tot[xx, zz] / 15

        plt.figure(33)
        plt.quiver(x_m, z_m, B[:, :, 0], B[:, :, 1])
        plt.xlabel("x-axis [mm]")
        plt.ylabel("z-axis [mm]")
        plt.show()

    def curv(B_field, z):
        return np.gradient(np.gradient(B_field, z), z) * 1e2

    def Bgrad(B_field, z):
        return np.gradient(B_field, z) * 1e1

    def Bz_plot_HH(self, I_current, x, z):
        B_z, B_x = self.B_field(I_current, x, z)

        B_z_curvature = np.gradient(np.gradient(B_z, z), z)

        plt.figure(100, figsize=(13, 10))

        # plt.rcParams.update({'font.size': 15})
        plt.suptitle("Helmholtz coil field B_z along z-axis")

        # Field plot
        ##########################
        plt.subplot(2, 1, 1)
        plt.plot(z, B_z, linestyle="solid", label=r"$B_z$: Result via elliptic integrals")

        # plt.xlim(-0.01,0.01)
        plt.title("B-field")

        plt.ylabel(r"$B_z$ [G]")
        plt.xlabel("z-axis [mm]")
        plt.legend()

        plt.subplot(2, 1, 2)
        plt.plot(z, B_z_curvature, linestyle="solid", label=r"$\nabla_z^2 B_z$: Result via elliptic integrals")

        plt.ylabel(r"$\nabla_z^2 B_z [G/cm^2]$")
        plt.xlabel("z-axis [mm]")
        plt.xlim(-10, 10)
        plt.title("Curvature of B-field")
        plt.legend(loc='lower right')

        plt.show()

    def Bz_plot_HH_comp(self, Coil2, I_current, x, z):
        B_z, B_x = self.B_field(I_current, x, z)
        B_z_2, B_x_2 = Coil2.B_field(I_current, x, z)

        B_z_curvature = np.gradient(np.gradient(B_z, z), z) * 1e2
        B_z_curvature_2 = np.gradient(np.gradient(B_z_2, z), z) * 1e2

        plt.figure(100, figsize=(13, 10))

        # plt.rcParams.update({'font.size': 15})
        plt.suptitle("Helmholtz coil field B_z along z-axis")

        # Field plot
        ##########################
        plt.subplot(2, 1, 1)
        plt.plot(z, B_z, linestyle="solid", label=r"$B_z$")
        plt.plot(z, B_z_2, linestyle="solid", label=r"$B_{z2}$")
        # plt.xlim(-0.01,0.01)
        plt.title("B-field")

        plt.ylabel(r"$B_z$ [G]")
        plt.xlabel("z-axis [mm]")
        plt.legend()

        plt.subplot(2, 1, 2)
        plt.plot(z, B_z_curvature, linestyle="solid", label=r"$\nabla_z^2 B_z$")
        plt.plot(z, B_z_curvature_2, linestyle="solid", label=r"$\nabla_z^2 B_{z2}$")

        plt.ylabel(r"$\nabla_z^2 B_z [G/cm^2]$")
        plt.xlabel("z-axis [mm]")
        plt.xlim(-10, 10)
        plt.title("Curvature of B-field")
        plt.legend(loc='lower right')

        plt.show()

    def cooling(self, I_current, T):
        """
        Print current density and power

        Parameters
        ----------
        I_current : current in A
        T : temperature in degree Celsius

        Returns
        -------
        None.

        """

        j_dens = I_current / self.get_wire_area() * 1e-6

        Power = self.resistance(T) * I_current ** 2

        print(f"current density = {j_dens} A/mm^2")
        print(f"Power = {Power} W")

    def power(self, I_current, T):
        P = self.resistance(T) * I_current ** 2
        return P

    def inductivity(self):
        return cs.mu_0 * (self.layers * self.windings) ** 2 * self.radius ** 2 * np.pi / self.get_coil_height()

    def induct_perry(self):
        """
        Calculates inductance of rectangular coil via empirical formular by perry

        """
        L = 4 * np.pi * (self.windings * self.layers) ** 2 * (1e2 * self.radius) ** 2 / (
                0.2317 * 100 * self.radius + 0.44 * 100 * self.get_coil_height() + 0.39 * 100 * self.get_coil_width())
        return L * 1e-9

    @staticmethod
    def resistivity_copper(T):
        R_ref = cs.rho_copper_20
        T_ref = 20  # degree celsius
        alpha = 0.00393
        R = R_ref * (1 + alpha * (T - T_ref))
        return R

    def resistance(self, T):
        """
        Calculates resistance of one coil of configuration

        Parameters
        ----------
        T : double
            temperature in degree Celsius

        Returns
        -------
        double
            resistance in Ohm
        """

        return BCoil.resistivity_copper(T) * self.get_wire_length() / self.get_wire_area()


def main():
    HH_Coil = BCoil(HH=1, distance=50, radius=40, layers=8, windings=8, wire_height=0.4, wire_width=0.4,
                    insulation_thickness=0.1, is_round=True, winding_offset=True)
    x = np.linspace(-50, 50, 300)
    z = np.linspace(-50, 50, 300)
    # Bz, Bx = HH_Coil.B_tot_along_axis(1.25, x, z)
    # Bz, Bx = HH_Coil.B_tot_along_axis(1.25, x, z)
    # #print(Bx)
    # #print(Bz)
    # plt.plot(x,Bx,label = "Bx")
    # plt.plot(z,Bz, label = "Bz")
    # plt.legend()
    # plt.show()
    #
    # HH_Coil.plot_raster(3)
    # ras = HH_Coil.full_raster(10)
    HH_Coil.plot_3d(1.25,50,50)

main()
# if __name__ == "__main__":
#     print("g")
#     main()