P464 - MHD Project¶

Proper Implementation (v1.0)¶

  • last update: 14-03-2024
    • currently using 'relative anti-symmetric' ghost-zone type with 3 cells, 6th order derivatives, and boundary conditions (whatever got from initial function).
In [ ]:
# importing the required libraries
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

fig_path = "figures/"
fig_save = True
In [ ]:
# defining spatial deritivative functions
def spatial_derivative(B, dx, order = 6, ghost_zone_type = 'relative anti-symmetric'):
    '''
    This function calculates the spatial derivatives of the magnetic field B.
    
    Parameters:
    B: 2D array
        The magnetic field B.
    dx: float
        The spatial step size.
    order: int
        The order of the finite difference scheme. It can be 2, 4, 6, 8 or 10.
        Default is 6.
    ghost_zone_type: str
        The type of the ghost zone. It can be 'anti-symmetric', 'symmetric', 'relative anti-symmetric' or None.
        Default is 'relative anti-symmetric'.
        
    Returns:
    dBdx: 2D array
        The spatial-derivative of the magnetic field B.
    ddBddx: 2D array
        The second z-derivative of the magnetic field B.
    '''

    dBdx = np.zeros(B.shape)
    ddBddx = np.zeros(B.shape)

    # print(f"The shape of B is {B.shape}")
    # print(f"The shape of dBdx is {dBdx.shape}")
    # print(f"The shape of ddBddx is {ddBddx.shape}")
    
    B = np.copy(B) # to avoid modifying the original array
    
    if order not in ['np',2, 4, 6, 8, 10]:
        raise ValueError("Order should be 'np',2, 4, 6, 8 or 10")
    
    if ghost_zone_type not in ['anti-symmetric', 'symmetric', 'relative anti-symmetric', None]:
        raise ValueError('ghost_zone_type should be anti-symmetric, symmetric, relative anti-symmetric or None')

    if order == 'np':
        ghost_zone = 3 # 3 cells on each side
    else:
        ghost_zone = int(order/2)
    
    if ghost_zone_type == 'relative anti-symmetric':
        B = np.pad(B, ((0, 0), (ghost_zone, ghost_zone)), mode='constant')
        for i in range(ghost_zone):
            B[:,i] = (2 * B[:,ghost_zone]) - B[:,ghost_zone + ghost_zone - i] # correcting the start of the array
            B[:, -1 - i] = 2 * B[:,-ghost_zone - 1] - B[:, (- ghost_zone - 1) - ghost_zone + i] # correcting the end of the array
    elif ghost_zone_type == 'anti-symmetric':
        B = np.pad(B, ((0, 0), (ghost_zone, ghost_zone)), mode='reflect') # makes symmetric padding
        for i in range(ghost_zone): 
            B[:,i] = - B[:, i] # making anti-symmetric at the start of the array
            B[:, -1 - i] = - B[:, -1 - i] # making anti-symmetric at the end of the array
    elif ghost_zone_type == 'symmetric':
        B = np.pad(B, ((0, 0), (ghost_zone, ghost_zone)), mode='symmetric')
    else:
        B = np.pad(B, ((0, 0), (ghost_zone, ghost_zone)), mode='constant') # zero padding

    # print(f"The shape of B after padding is {B.shape}")
    # print(f"The shape of dBdx after padding is {dBdx.shape}")
    # print(f"The shape of ddBddx after padding is {ddBddx.shape}")
    
    if order == 6:
        for i in range(ghost_zone, B.shape[1] - ghost_zone):
            dBdx[:,i-ghost_zone] = ((- B[:,i - 3]) + (9 * B[:,i - 2]) - (45 * B[:,i - 1]) + (45 * B[:,i + 1]) - (9 * B[:,i + 2]) + (B[:,i + 3])) / (60 * dx)
            ddBddx[:,i-ghost_zone] = ((2 * B[:,i - 3]) - (27 * B[:,i - 2]) + (270 * B[:,i - 1]) - (490 * B[:, i]) + (270 * B[:,i + 1]) - (27 * B[:,i + 2]) + (2 * B[:,i + 3])) / (180 * (dx ** 2))
    elif order == 2:
        for i in range(ghost_zone, B.shape[1] - ghost_zone):
            dBdx[:,i-ghost_zone] = ((- B[:,i - 1]) + (B[:,i + 1])) / (2 * dx)
            ddBddx[:,i-ghost_zone] = (B[:,i - 1] - (2 * B[:, i]) + B[:,i + 1]) / (dx ** 2)
    elif order == 4:
        for i in range(ghost_zone, B.shape[1] - ghost_zone):
            dBdx[:,i-ghost_zone] = ((B[:,i - 2]) - (8 * B[:,i - 1]) + (8 * B[:,i + 1]) - (B[:,i + 2])) / (12 * dx)
            ddBddx[:,i-ghost_zone] = ((- B[:,i - 2]) + (16 * B[:,i - 1]) - (30 * B[:,i]) + (16 * B[:,i + 1]) - (B[:,i + 2])) / (12 * (dx ** 2))
    elif order == 8:
        for i in range(ghost_zone, B.shape[1] - ghost_zone):
            dBdx[:,i-ghost_zone] = ((3*B[:,i - 4]) - (32 * B[:,i - 3]) + (168 * B[:,i - 2]) - (672 * B[:,i - 1]) + (672 * B[:,i + 1]) - (168 * B[:,i + 2]) + (32 * B[:,i + 3]) - (3 * B[:,i + 4])) / (840 * dx)
            ddBddx[:,i-ghost_zone] = ((-9 * B[:,i - 4]) + (128 * B[:,i - 3]) - (1008 * B[:,i - 2]) + (8064 * B[:,i - 1]) - (14350 * B[:,i]) + (8064 * B[:,i + 1]) - (1008 * B[:,i + 2]) + (128 * B[:,i + 3]) - (9 * B[:,i + 4])) / (5040 * (dx ** 2))
    elif order == 10:
        for i in range(ghost_zone, B.shape[1] - ghost_zone):
            dBdx[:,i-ghost_zone] = ((-2 * B[:,i - 5]) + (25 * B[:,i - 4]) - (150 * B[:,i - 3]) + (600 * B[:,i - 2]) - (2100 * B[:,i - 1]) + (2100 * B[:,i + 1]) - (600 * B[:,i + 2]) + (150 * B[:,i + 3]) - (25 * B[:,i + 4]) + (2 * B[:,i + 5])) / (2520 * dx)
            ddBddx[:,i-ghost_zone] = ((8 * B[:,i - 5]) - (125 * B[:,i - 4]) + (1000 * B[:,i - 3]) - (6000 * B[:,i - 2]) + (42000 * B[:,i - 1]) - (73766 * B[:, i]) + (42000 * B[:,i + 1]) - (6000 * B[:,i + 2]) + (1000 * B[:,i + 3]) - (125 * B[:,i + 4]) + (8 * B[:,i + 5])) / (25200 * (dx ** 2))
    else:
        dBdx = np.gradient(B, dx, axis=1)
        ddBddx = np.gradient(dBdx, dx, axis=1)
        
        # removing the ghost zones
        dBdx = dBdx[:,ghost_zone:-ghost_zone]
        ddBddx = ddBddx[:,ghost_zone:-ghost_zone]
            
    # update with other orders
    
    return dBdx, ddBddx

z implementation¶

In [ ]:
def simulate_B_z_field(z_i, z_f, Nz, T, Nt, B_r0, B_phi0, dBdt, order = 6, ghost_zone_type = 'relative anti-symmetric', iterations_to_save = None):
    '''    
    This function simulates the time evolution of the magnetic field B.
    
    Parameters:
    z_i: float
        The start of the spatial domain.
    z_f: float
        The end of the spatial domain.
    Nz: int
        The number of spatial grid points.
    T: float
        The total simulation time.
    Nt: int
        The number of time steps.
    B_r0: function
        The initial condition for the magnetic field B_r.
    B_phi0: function
        The initial condition for the magnetic field B_phi.
    dBdt: function
        The time evolution of the magnetic field B.
    order: int
        The order of the finite difference scheme. It can be 'np', 2, 4, 6, 8 or 10.
        Default is 6.
    ghost_zone_type: str
        The type of the ghost zone. It can be 'anti-symmetric', 'symmetric', 'relative anti-symmetric' or None.
        Default is 'relative anti-symmetric'.
    iterations_to_save: list
        A list of the time steps to save the solution.
        If None, it will save the solution at the start and end of the simulation.
        If 'all', it will save the solution at every time step.
        
    Returns:
    B_list: 3-D array
        A list of the magnetic field B at the required time steps from iterations_to_save.
        size: (len(iterations_to_save), 2, Nz)
    '''
    dz = (z_f - z_i) / (Nz - 1)  # Spatial step size
    dt = T / Nt  # Temporal step size
    
    # Create arrays to store the solution in time
    B_list = []

    # creating initial conditions from z, B_r0 and B_phi0
    z = np.linspace(z_i, z_f, Nz)
    B = np.zeros((2, Nz))
    B[0, :] = B_r0(0, 0, z) # first row is B_r
    B[1, :] = B_phi0(0, 0, z) # second row is B_phi
    
    if iterations_to_save is None:
        iterations_to_save = [0, Nt]
    elif iterations_to_save == 'all':
        iterations_to_save = list(range(0, Nt + 1))

    if 0 in iterations_to_save:
        B_list.append(np.copy(B))
    
    #--------- RK4 time-stepping block ---------#
    for n in tqdm(range(1, Nt + 1)):
        # Compute spatial derivatives
        dBdz, ddBddz = spatial_derivative(B, dz, order, ghost_zone_type)


        k1 = dBdt(B, dBdz, ddBddz, n * dt)
        k2 = dBdt(B + (0.5 * dt * k1), dBdz, ddBddz, (n * dt) + (0.5 * dt))
        k3 = dBdt(B + (0.5 * dt * k2), dBdz, ddBddz, (n * dt) + (0.5 * dt))
        k4 = dBdt(B + (dt * k3), dBdz, ddBddz, (n * dt) + dt)
        
        B = B + (dt / 6.0) * (k1 + (2 * k2) + (2 * k3) + k4)
        
        # print(B)
        
        if n in iterations_to_save:
            B_list.append(np.copy(B))
    #------------------------------------------#
    
    B_list = np.array(B_list)
    
    return B_list
In [ ]:
# defining constants and parameters
eta = 1

delta = 0.01 # safety net to avoid division by zero
z_i = -1.0 + delta  # start of spatial domain
z_f = 1.0  + delta # end of spatial domain
T = 1.0  # Total simulation time
Nz = 80  # Number of spatial grid points
Nt_z = 5000  # Number of time steps
h = 1

# order = 'np'
order = 6
ghost_zone_type = 'relative anti-symmetric'

iterations_to_save = [int(i * Nt_z/10) for i in range(10 + 1)] # which time steps to save

# defining initial conditions
def B_phi0_z(r, phi, z):
    return (3 * np.sin(((z - z_i)/(z_f - z_i))*np.pi)) + (1 * np.sin(((z - z_i)/(z_f - z_i))*2*np.pi)) + (2 * np.sin(((z - z_i)/(z_f - z_i))*5*np.pi))
def B_r0_z(r, phi, z): # temporary; since we are solving the same equation in r and phi; not using now
    return 0 

# definiing the time evolution of the magnetic field
def dBdt_z(B, dBdz, ddBdz, t):
    return eta * ddBdz
In [ ]:
# plotting initial conditions
z = np.linspace(z_i, z_f, Nz)
B = np.zeros((2, Nz))
B[0, :] = B_r0_z(0, 0, z) # first row is B_r
B[1, :] = B_phi0_z(0, 0, z) # second row is B_phi

# plt.plot(z, B[0, :], label = 'B_r')
plt.plot(z, B[1, :], label = r'$B_{\phi}$')
plt.xlabel(r'z')
plt.ylabel(r'B')
plt.title('Initial condition of B')
plt.grid()
plt.legend()
if fig_save:
    plt.savefig(fig_path + "initial_condition_B.png")
plt.show()

if order == 'np':
    ghost_zone = 3
else:
    ghost_zone = int(order/2)

# testing the spatial_derivative function
dBdx = np.zeros(B.shape)
ddBddx = np.zeros(B.shape)
B = np.pad(B, ((0, 0), (ghost_zone, ghost_zone)), mode='constant')
print(B.shape)
for i in range(ghost_zone):
    B[:,i] = 2 * B[:,ghost_zone] - B[:,ghost_zone + ghost_zone - i] # correcting the start of the array
    B[:, -1 - i] = 2 * B[:,-ghost_zone - 1] - B[:, (- ghost_zone - 1) - ghost_zone + i] # correcting the end of the array
dx = (z_f - z_i) / (Nz - 1)  # Spatial step size
if order == 6:
    for i in range(ghost_zone, B.shape[1] - ghost_zone):
        dBdx[:, i-ghost_zone] = ((- B[:,i - 3]) + (9 * B[:,i - 2]) - (45 * B[:,i - 1]) + (45 * B[:,i + 1]) - (9 * B[:,i + 2]) + (B[:,i + 3])) / (60 * dx)
        ddBddx[:,i-ghost_zone] = ((2 * B[:,i - 3]) - (27 * B[:,i - 2]) + (270 * B[:,i - 1]) - (490 * B[:, i]) + (270 * B[:,i + 1]) - (27 * B[:,i + 2]) + (2 * B[:,i + 3])) / (180 * (dx ** 2))
else:
    dBdx = np.gradient(B, dx, axis=1)
    ddBddx = np.gradient(dBdx, dx, axis=1)

    # removing the ghost zones
    dBdx = dBdx[:,ghost_zone:-ghost_zone]
    ddBddx = ddBddx[:,ghost_zone:-ghost_zone]

# plt.plot(z, dBdx[0,:], label = 'dB_r/dz')
plt.plot(z, dBdx[1,:], label = r'$\frac{\partial B_z}{\partial z}$')
plt.xlabel('z')
plt.ylabel(r'$\frac{\partial B}{\partial z}$')
plt.title('first spatial derivative')
plt.grid()
plt.legend()
if fig_save:
    plt.savefig(fig_path + "first_spatial_derivative.png")
plt.show()

# plt.plot(z, ddBddx[0,:], label = 'd^2B_r/dz^2')
plt.plot(z, ddBddx[1,:], label = r'$\frac{\partial^2B_z}{\partial^2z}$')
plt.xlabel('z')
plt.ylabel(r'$\frac{\partial^2B}{\partial^2z}$')
plt.title('second spatial derivative')
plt.grid()
plt.legend()
if fig_save:
    plt.savefig(fig_path + "second_spatial_derivative.png")
plt.show()
No description has been provided for this image
(2, 86)
No description has been provided for this image
No description has been provided for this image
In [ ]:
# running the simulation
B_list = simulate_B_z_field(z_i, z_f, Nz, T, Nt_z, B_r0_z, B_phi0_z, dBdt_z, order, ghost_zone_type, iterations_to_save)
  0%|          | 0/5000 [00:00<?, ?it/s]100%|██████████| 5000/5000 [00:14<00:00, 338.64it/s]
In [ ]:
# #plotting the results
# plt.figure(figsize=(10, 5))
# plt.title('Magnetic field B_r')
# plt.xlabel('z')
# plt.ylabel('B_r')
# plt.grid(True)
# for i in range(len(iterations_to_save)):
#     plt.plot(np.linspace(z_i, z_f, Nz), B_list[i, 0, :], label=f't = {iterations_to_save[i] * T / Nt:.4f}')
# plt.legend()
# plt.show()

plt.figure(figsize=(10, 5))
plt.title(r'Evolution of Magnetic field $B_{\phi}$')
plt.xlabel(r'z')
plt.ylabel(r'$B_{\phi}$')
plt.grid(True)
for i in range(len(iterations_to_save)):
    plt.plot(np.linspace(z_i, z_f, Nz), B_list[i, 1, :], label=f't = {iterations_to_save[i] * T / Nt_z:.4f}')
plt.legend()
if fig_save:
    plt.savefig(fig_path + "evolution_B_phi_z.png")
plt.show()
No description has been provided for this image

$r$ implementation¶

In [ ]:
def simulate_B_r_field(r_i, r_f, Nr, T, Nt, B_r0, B_phi0, dBdt_r, order = 6, ghost_zone_type = 'relative anti-symmetric', iterations_to_save = None):
    '''    
    This function simulates the time evolution of the magnetic field B.
    
    Parameters:
    z_i: float
        The start of the spatial domain.
    z_f: float
        The end of the spatial domain.
    Nr: int
        The number of spatial grid points.
    T: float
        The total simulation time.
    Nt: int
        The number of time steps.
    B_r0: function
        The initial condition for the magnetic field B_r.
    B_phi0: function
        The initial condition for the magnetic field B_phi.
    dBdt_r: function
        The time evolution of the magnetic field B.
    order: int
        The order of the finite difference scheme. It can be 'np', 2, 4, 6, 8 or 10.
        Default is 6.
    ghost_zone_type: str
        The type of the ghost zone. It can be 'anti-symmetric', 'symmetric', 'relative anti-symmetric' or None.
        Default is 'relative anti-symmetric'.
    iterations_to_save: list
        A list of the time steps to save the solution. 
        If None, it will save the solution at the start and end of the simulation.
        if 'all', it will save the solution at every time step.
        
    Returns:
    B_list: 3-D array
        A list of the magnetic field B at the required time steps from iterations_to_save.
        size: (len(iterations_to_save), 2, Nz)
    '''
    dr = (r_f - r_i) / (Nr - 1)  # Spatial step size
    dt = T / Nt  # Temporal step size
    
    # Create arrays to store the solution in time
    B_list = []

    # creating initial conditions from z, B_r0 and B_phi0
    r = np.linspace(r_i, r_f, Nr)
    r_ = np.array([r,r]) # Under the assumption that the equation is same for both B_r and B_phi
                         # We need this r_ for the equation
    
    B = np.zeros((2, Nr))
    B[0, :] = B_r0(r, 0, 0) # first row is B_r
    B[1, :] = B_phi0(r, 0, 0) # second row is B_phi
    
    if iterations_to_save is None:
        iterations_to_save = [0, Nt]
    elif iterations_to_save == 'all':
        iterations_to_save = list(range(Nt + 1))

    if 0 in iterations_to_save:
        B_list.append(np.copy(B))
    
    #--------- RK4 time-stepping block ---------#
    for n in tqdm(range(1, Nt + 1)):
        # Compute spatial derivatives
        dBdr, ddBddr = spatial_derivative(B, dr, order, ghost_zone_type)


        k1 = dBdt_r(B, r_, dBdr, ddBddr, n * dt)
        k2 = dBdt_r(B + (0.5 * dt * k1), r_, dBdr, ddBddr, (n * dt) + (0.5 * dt))
        k3 = dBdt_r(B + (0.5 * dt * k2), r_, dBdr, ddBddr, (n * dt) + (0.5 * dt))
        k4 = dBdt_r(B + (dt * k3), r_, dBdr, ddBddr, (n * dt) + dt)
        
        B = B + (dt / 6.0) * (k1 + (2 * k2) + (2 * k3) + k4)
        
        # print(B)
        
        if n in iterations_to_save:
            B_list.append(np.copy(B))
    #------------------------------------------#
    
    B_list = np.array(B_list)
    
    return B_list
In [ ]:
# defining constants and parameters
eta = 1
r_i = 0.01  # start of spatial domain
r_f = 8.01 # end of spatial domain
T = 1.0  # Total simulation time
Nr = 100  # Number of spatial grid points
Nt_r = 4000  # Number of time steps
h = 1
delta = 0.00001 # for avoiding division by zero


r_i = r_i + delta
r_f = r_f + delta

r = np.linspace(r_i, r_f, Nr)

order = 'np'
order = 6
ghost_zone_type = 'relative anti-symmetric'

iterations_to_save2 = [int(i * Nt_r/8) for i in range(8 + 1)] # which time steps to save

# defining initial conditions
def B_phi0_r(r, phi, z):
    return (3 * np.sin(((r - r_i)/(r_f - r_i))*np.pi)) + (1 * np.sin(((r - r_i)/(r_f - r_i))*2*np.pi)) + (2 * np.sin(((r - r_i)/(r_f - r_i))*5*np.pi))

def B_r0_r(r, phi, z):
    return 0

# definiing the time evolution of the magnetic field
def dBdt_r(B, r_, dBdr, ddBdr, t):
    return eta * (ddBdr + (dBdr / r_) - (B/(r_ ** 2)) - (((np.pi ** 2) / (4 * (h**2))) * B))
In [ ]:
# plotting initial magnetic field
r = np.linspace(r_i, r_f, Nr)
B = np.zeros((2, Nr))
B[0, :] = B_r0_r(r, 0, 0) # first row is B_r
B[1, :] = B_phi0_r(r, 0, 0) # second row is B_phi

# plt.plot(r, B[0, :], label = 'B_r')
plt.plot(r, B[1, :], label = 'B_phi')
plt.xlabel('r')
plt.ylabel('B')
plt.grid(True)
plt.title('Initial conditions')
plt.legend()
plt.show()
No description has been provided for this image
In [ ]:
# simulating the magnetic field
B_list2 = simulate_B_r_field(r_i, r_f, Nr, T, Nt_r, B_r0_r, B_phi0_r, dBdt_r, order, ghost_zone_type, iterations_to_save2)
100%|██████████| 4000/4000 [00:14<00:00, 273.69it/s]
In [ ]:
# plotting the results
# plt.figure(figsize=(10, 5))
# plt.title('Magnetic field B_r')
# plt.xlabel('r')
# plt.ylabel('B_r')
# plt.grid(True)
# for i in range(len(iterations_to_save)):
#     plt.plot(np.linspace(r_i, r_f, Nr), B_list2[i, 0, :], label=f't = {iterations_to_save[i] * T / Nt:.4f}')
# plt.legend()
# plt.show()

plt.figure(figsize=(10, 5))
plt.title(r'Time evolution of Magnetic field $B_{\phi}$')
plt.xlabel(r'r')
plt.ylabel(r'$B_{\phi}$')
plt.grid(True)
for i in range(len(iterations_to_save2)):
    plt.plot(np.linspace(r_i, r_f, Nr), B_list2[i, 1, :], label=f't = {iterations_to_save2[i] * T / Nt_r:.4f}')
plt.legend()
if fig_save:
    plt.savefig(fig_path + "evolution_B_phi_r.png")
plt.show()
No description has been provided for this image

Evolution of Magnetic field strength over time at the center of the grid + some other points¶

In [ ]:
# running the simulation to get every time step (to be used even further)
B_list = simulate_B_z_field(z_i, z_f, Nz, T, Nt_z, B_r0_z, B_phi0_z, dBdt_z, order, ghost_zone_type, iterations_to_save='all')
B_list2 = simulate_B_r_field(r_i, r_f, Nr, T, Nt_r, B_r0_r, B_phi0_r, dBdt_r, order, ghost_zone_type, iterations_to_save='all')
100%|██████████| 5000/5000 [00:14<00:00, 341.03it/s]
100%|██████████| 4000/4000 [00:14<00:00, 273.51it/s]
In [ ]:
# taking 7 equally placed points within the domain
z = np.linspace(z_i, z_f, 9) # 9 points to exclude the boundaries
z = z[1:-1]

spacesteps_to_save = np.linspace(0, Nz - 1, 9, dtype = int) # 9 points to exclude the boundaries
spacesteps_to_save = spacesteps_to_save[1:-1]

# plotting the magnetic field at different time steps
plt.figure(figsize=(10, 5))
plt.title(r'Magnetic field $|B_{\phi}|$ at different time steps for z-implementation')
plt.xlabel(r't')
plt.ylabel(r'$B_{\phi}$')
for i in range(len(spacesteps_to_save)):
    plt.plot(np.linspace(0, T, Nt_z + 1), abs(B_list[:, 1, spacesteps_to_save[i]]), label=f'z = {z[i]:.4f}')
plt.grid(True)
plt.legend()
if fig_save:
    plt.savefig(fig_path + "B_phi_z_strength__timesteps.png")
plt.show()

# # plotting in log scale
# plt.figure(figsize=(10, 5))
# plt.title(r'Magnetic field $|B_{\phi}|$ at different time steps for z-implementation in log scale')
# plt.xlabel(r't')
# plt.ylabel(r'$log(B_{\phi})$')
# for i in range(len(spacesteps_to_save)):
#     plt.plot(np.linspace(0, T, Nt_z + 1), abs(B_list[:, 1, spacesteps_to_save[i]]), label=f'z = {z[i]:.4f}')
# plt.yscale('log')
# plt.legend()
# plt.grid(True)
# if fig_save:
#     plt.savefig(fig_path + "B_phi_z_strength__timesteps_log.png")
# plt.show()
No description has been provided for this image
In [ ]:
# calculating the decay rate for the middle point
midz = int(Nz/2)
B_phi_midz = abs(B_list[:, 1, midz])
t = np.linspace(0, T, Nt_z + 1)
log_B_phi_midz = np.log(B_phi_midz)
slope, intercept = np.polyfit(t, log_B_phi_midz, 1)
print(f"The decay rate for the middle point is {slope:.4f}")

# plotting the magnetic field decay and fitting
plt.figure(figsize=(10, 5))
plt.title(r'Magnetic field $|B_{\phi}|$ decay for the middle point for z-implementation')
plt.xlabel(r't')
plt.ylabel(r'$log(B_{\phi})$')
plt.plot(t, log_B_phi_midz, label = 'log(B_phi)')
plt.plot(t, slope*t + intercept,'k--', label = f'fit: ({slope:.4f})t + {intercept:.4f}')
plt.grid(True)
plt.legend()
if fig_save:
    plt.savefig(fig_path + "B_phi_z_decay.png")
plt.show()
The decay rate for the middle point is -2.5170
No description has been provided for this image
In [ ]:
r = np.linspace(r_i, r_f, 9) # 9 points to exclude the boundaries
r = r[1:-1]

spacesteps_to_save2 = np.linspace(0, Nr - 1, 9, dtype = int) # 9 points to exclude the boundaries
spacesteps_to_save2 = spacesteps_to_save2[1:-1]

# plotting the magnetic field at different time steps
plt.figure(figsize=(10, 5))
plt.title(r'Magnetic field $|B_{\phi}|$ at different time steps for r-implementation')
plt.xlabel(r't')
plt.ylabel(r'$B_{\phi}$')
for i in range(len(spacesteps_to_save2)):
    plt.plot(np.linspace(0, T, Nt_r + 1), abs(B_list2[:, 1, spacesteps_to_save2[i]]), label=f'r = {r[i]:.4f}')
plt.grid(True)
plt.legend()
if fig_save:
    plt.savefig(fig_path + "B_phi_r_strength__timesteps.png")
plt.show()

# # plotting in log scale
# plt.figure(figsize=(10, 5))
# plt.title(r'Magnetic field $|B_{\phi}|$ at different time steps for r-implementation in log scale')
# plt.xlabel(r't')
# plt.ylabel(r'$log(B_{\phi})$')
# for i in range(len(spacesteps_to_save2)):
#     plt.plot(np.linspace(0, T, Nt_r + 1), abs(B_list2[:, 1, spacesteps_to_save2[i]]), label=f'r = {r[i]:.4f}')
# plt.yscale('log')
# if fig_save:
#     plt.savefig(fig_path + "B_phi_r_strength__timesteps_log.png")
# plt.legend()
# plt.grid(True)
# plt.show()
No description has been provided for this image
In [ ]:
# calculating the decay rate for the middle point
midr = int(Nr/2)
B_phi_midr = abs(B_list2[:, 1, midr])
t = np.linspace(0, T, Nt_r + 1)
log_B_phi_midr = np.log(B_phi_midr)
slope, intercept = np.polyfit(t, log_B_phi_midr, 1)
print(f"The decay rate for the middle point is {slope:.4f}")

# plotting the magnetic field decay and fitting
plt.figure(figsize=(10, 5))
plt.title(r'Magnetic field $|B_{\phi}|$ decay for the middle point for r-implementation')
plt.xlabel(r't')
plt.ylabel(r'$log(B_{\phi})$')
plt.plot(t, log_B_phi_midr, label = 'log(B_phi)')
plt.plot(t, slope*t + intercept,'k--', label = f'fit: ({slope:.4f})t + {intercept:.4f}')
plt.grid(True)
plt.legend()
if fig_save:
    plt.savefig(fig_path + "B_phi_r_decay.png")
plt.show()
The decay rate for the middle point is -3.1719
No description has been provided for this image

Evolution of both B_r and B_phi over time at (the center of the grid + some other points), and that of pitch angle¶

In [ ]:
# need to run the simulation for every time step (again)

# defining both the initial conditions for the magnetic field
def B_phi0_r(r, phi, z):
    return (3 * np.sin(((r - r_i)/(r_f - r_i))*np.pi)) + (1 * np.sin(((r - r_i)/(r_f - r_i))*2*np.pi)) + (2 * np.sin(((r - r_i)/(r_f - r_i))*5*np.pi))
def B_r0_r(r, phi, z):
    return (3 * np.sin(((r - r_i)/(r_f - r_i))*np.pi)) + (2.5 * np.sin(((r - r_i)/(r_f - r_i))*2*np.pi)) + (1 * np.sin(((r - r_i)/(r_f - r_i))*5*np.pi))
def B_phi0_z(r, phi, z):
    return (3 * np.sin(((z - z_i)/(z_f - z_i))*np.pi)) + (1 * np.sin(((z - z_i)/(z_f - z_i))*2*np.pi)) + (2 * np.sin(((z - z_i)/(z_f - z_i))*5*np.pi))
def B_r0_z(r, phi, z):
    return (3 * np.sin(((z - z_i)/(z_f - z_i))*np.pi)) + (2.5 * np.sin(((z - z_i)/(z_f - z_i))*2*np.pi)) + (1 * np.sin(((z - z_i)/(z_f - z_i))*5*np.pi))

B_list = simulate_B_z_field(z_i, z_f, Nz, T, Nt_z, B_r0_z, B_phi0_z, dBdt_z, order, ghost_zone_type, iterations_to_save='all')
B_list2 = simulate_B_r_field(r_i, r_f, Nr, T, Nt_r, B_r0_r, B_phi0_r, dBdt_r, order, ghost_zone_type, iterations_to_save='all')
  0%|          | 0/5000 [00:00<?, ?it/s]100%|██████████| 5000/5000 [00:14<00:00, 342.05it/s]
100%|██████████| 4000/4000 [00:14<00:00, 273.34it/s]
In [ ]:
print(B_list.shape)
print(B_list2.shape)

print(iterations_to_save)
print(iterations_to_save2)
(5001, 2, 80)
(4001, 2, 100)
[0, 500, 1000, 1500, 2000, 2500, 3000, 3500, 4000, 4500, 5000]
[0, 500, 1000, 1500, 2000, 2500, 3000, 3500, 4000]
In [ ]:
# plot 4 figures for the magnetic field at different time steps in spatial domain, side by side
# each plot has spatial domain in x-axis and magnetic field in y-axis
plt.suptitle('Magnetic field $|B_{\phi}|$ at different time steps for both z and r implementations')
fig, axs = plt.subplots(2, 2, figsize=(20, 10))

# z-implementation B-phi
axs[0, 0].set_title(r'Magnetic field $|B_{\phi}|$ at different time steps for z-implementation')
axs[0, 0].set_xlabel(r'z')
axs[0, 0].set_ylabel(r'$B_{\phi}$')
for i in range(len(iterations_to_save)):
    axs[0, 0].plot(np.linspace(z_i, z_f, Nz), B_list[iterations_to_save[i], 1, :], label=f't = {iterations_to_save[i] * T / Nt_z:.4f}')
axs[0, 0].grid(True)
axs[0, 0].legend()

# r-implementation B-phi
axs[0, 1].set_title(r'Magnetic field $|B_{\phi}|$ at different time steps for r-implementation')
axs[0, 1].set_xlabel(r'r')
axs[0, 1].set_ylabel(r'$B_{\phi}$')
for i in range(len(iterations_to_save2)):
    axs[0, 1].plot(np.linspace(r_i, r_f, Nr), B_list2[iterations_to_save2[i], 1, :], label=f't = {iterations_to_save2[i] * T / Nt_r:.4f}')
axs[0, 1].grid(True)
axs[0, 1].legend()

# z-implementation B_r
axs[1, 0].set_title(r'Magnetic field $|B_{r}|$ at different time steps for z-implementation')
axs[1, 0].set_xlabel(r'z')
axs[1, 0].set_ylabel(r'$B_{r}$')
for i in range(len(iterations_to_save)):
    axs[1, 0].plot(np.linspace(z_i, z_f, Nz), B_list[iterations_to_save[i], 0, :], label=f't = {iterations_to_save[i] * T / Nt_z:.4f}')
axs[1, 0].grid(True)
axs[1, 0].legend()

# r-implementation B_r
axs[1, 1].set_title(r'Magnetic field $|B_{r}|$ at different time steps for r-implementation')
axs[1, 1].set_xlabel(r'r')
axs[1, 1].set_ylabel(r'$B_{r}$')
for i in range(len(iterations_to_save2)):
    axs[1, 1].plot(np.linspace(r_i, r_f, Nr), B_list2[iterations_to_save2[i], 0, :], label=f't = {iterations_to_save2[i] * T / Nt_r:.4f}')
axs[1, 1].grid(True)
axs[1, 1].legend()
if fig_save:
    plt.savefig(fig_path + "Bphi_Br_z_r_evolution.png")
plt.show()
<Figure size 640x480 with 0 Axes>
No description has been provided for this image
In [ ]:
# computing pitch angle at middle point for both z and r implementations

midz = int(Nz/2)
midr = int(Nr/2)

# z-implementation
B_phi_midz = B_list[:, 1, midz]
B_r_midz = B_list[:, 0, midz]
pitch_angle_z = np.arctan(B_phi_midz/B_r_midz)
t_z = np.linspace(0, T, Nt_z + 1)

# r-implementation
B_phi_midr = B_list2[:, 1, midr]
B_r_midr = B_list2[:, 0, midr]
pitch_angle_r = np.arctan(B_phi_midr/B_r_midr)
t_r = np.linspace(0, T, Nt_r + 1)

# converting to degrees
pitch_angle_z = pitch_angle_z * 180 / np.pi
pitch_angle_r = pitch_angle_r * 180 / np.pi

# plotting the pitch angle for both z and r implementations
plt.figure(figsize=(10, 5))
plt.title(r'Pitch angle at the middle point for both z and r implementations')
plt.xlabel(r't')
plt.ylabel(r'Pitch angle')
plt.plot(t_z, pitch_angle_z, label = 'z-implementation')
plt.plot(t_r, pitch_angle_r, label = 'r-implementation')
plt.grid(True)
plt.legend()
if fig_save:
    plt.savefig(fig_path + "pitch_angle_z_r_mid.png")
plt.show()
No description has been provided for this image
In [ ]:
# pitch angle evolution for both z and r implementations at 5 different points within the spatial domain

space_steps_to_save = np.linspace(0, Nz - 1, 7, dtype = int) # 7 points to exclude the boundaries
space_steps_to_save = space_steps_to_save[1:-1]

space_steps_to_save2 = np.linspace(0, Nr - 1, 7, dtype = int) # 8 points to exclude the boundaries
space_steps_to_save2 = space_steps_to_save2[1:-1]

# z-implementation
pitch_angle_z = np.zeros((len(space_steps_to_save), Nt_z + 1))
for i in range(0,len(space_steps_to_save)):
    B_phi = B_list[:, 1, space_steps_to_save[i]]
    B_r = B_list[:, 0, space_steps_to_save[i]]
    pitch_angle_z[i, :] = np.arctan(B_phi/B_r)
    
# r-implementation
pitch_angle_r = np.zeros((len(space_steps_to_save2), Nt_r + 1))
for i in range(len(space_steps_to_save2)):
    B_phi = B_list2[:, 1, space_steps_to_save2[i]]
    B_r = B_list2[:, 0, space_steps_to_save2[i]]
    pitch_angle_r[i, :] = np.arctan(B_phi/B_r)
    
# converting the pitch angle to degrees
pitch_angle_z = pitch_angle_z * 180 / np.pi
pitch_angle_r = pitch_angle_r * 180 / np.pi
    
# plotting the pitch angle evolution for both z and r implementations
plt.figure(figsize=(10, 5))
plt.title(r'Pitch angle evolution for z implementation at 6 different points')
plt.xlabel(r't')
plt.ylabel(r'Pitch angle')
for i in range(0,len(space_steps_to_save)):
    plt.plot(t_z, pitch_angle_z[i, :], label = f'z = {z_i + (space_steps_to_save[i] * (z_f - z_i) / (Nz - 1)):.4f}')
plt.grid(True)
plt.legend()
if fig_save:
    plt.savefig(fig_path + "pitch_angle_z_timeevolution_at_diff_spatialsteps.png")
plt.show()

plt.figure(figsize=(10, 5))
plt.title(r'Pitch angle evolution for r implementation at 7 different points')
plt.xlabel(r't')
plt.ylabel(r'Pitch angle')
for i in range(len(space_steps_to_save2)):
    plt.plot(t_r, pitch_angle_r[i, :], label = f'r = {r_i + (space_steps_to_save2[i] * (r_f - r_i) / (Nr - 1)):.4f}')
plt.grid(True)
plt.legend()
if fig_save:
    plt.savefig(fig_path + "pitch_angle_r_timeevolution_at_diff_spatialsteps.png")
plt.show()
No description has been provided for this image
No description has been provided for this image
In [ ]:
# Pitch angle across the spatial domain at different time steps

iterations_to_save_ = [int(i * Nt_z/4) for i in range(4 + 1)] # which time steps to save
iterations_to_save_2 = [int(i * Nt_r/4) for i in range(4 + 1)] # which time steps to save

# z-implementation
pitch_angle_z = np.zeros((len(iterations_to_save), Nz))
for i in range(len(iterations_to_save_)):
    B_phi = B_list[iterations_to_save_[i], 1, :]
    B_r = B_list[iterations_to_save_[i], 0, :]
    pitch_angle_z[i, :] = np.arctan(B_phi/B_r)

# r-implementation
pitch_angle_r = np.zeros((len(iterations_to_save_2), Nr))
for i in range(len(iterations_to_save_2)):
    B_phi = B_list2[iterations_to_save_2[i], 1, :]
    B_r = B_list2[iterations_to_save_2[i], 0, :]
    pitch_angle_r[i, :] = np.arctan(B_phi/B_r)
    
# converting the pitch angle to degrees
pitch_angle_z = pitch_angle_z * 180 / np.pi
pitch_angle_r = pitch_angle_r * 180 / np.pi

# plotting the pitch angle across the spatial domain at different time steps for both z and r implementations
plt.figure(figsize=(10, 5))
plt.title(r'Pitch angle across the spatial domain at different time steps for z implementation')
plt.xlabel(r'z')
plt.ylabel(r'Pitch angle')
for i in range(len(iterations_to_save_)):
    plt.plot(np.linspace(z_i, z_f, Nz), pitch_angle_z[i, :], label = f't = {iterations_to_save_[i] * T / Nt_z:.4f}')
plt.grid(True)
plt.legend()
if fig_save:
    plt.savefig(fig_path + "pitch_angle_z_spatialevolution_at_diff_timesteps.png")
plt.show()

plt.figure(figsize=(10, 5))
plt.title(r'Pitch angle across the spatial domain at different time steps for r implementation')
plt.xlabel(r'r')
plt.ylabel(r'Pitch angle')
for i in range(len(iterations_to_save_2)):
    plt.plot(np.linspace(r_i, r_f, Nr), pitch_angle_r[i, :], label = f't = {iterations_to_save_2[i] * T / Nt_r:.4f}')
plt.grid(True)
plt.legend()
if fig_save:
    plt.savefig(fig_path + "pitch_angle_r_spatialevolution_at_diff_timesteps.png")
plt.show()
C:\Users\adhil\AppData\Local\Temp\ipykernel_20556\3662263396.py:11: RuntimeWarning: invalid value encountered in divide
  pitch_angle_z[i, :] = np.arctan(B_phi/B_r)
C:\Users\adhil\AppData\Local\Temp\ipykernel_20556\3662263396.py:18: RuntimeWarning: invalid value encountered in divide
  pitch_angle_r[i, :] = np.arctan(B_phi/B_r)
No description has been provided for this image
No description has been provided for this image

Exploring different magnetic fields and boundary conditions evolve¶

In [ ]:
# trial 1 (sin wave and zero boundary condition)

def B_phi0_z(r, phi, z):
    return (3 * np.sin(((z - z_i)/(z_f - z_i))*np.pi)) + (1 * np.sin(((z - z_i)/(z_f - z_i))*2*np.pi)) + (2 * np.sin(((z - z_i)/(z_f - z_i))*5*np.pi))
def B_r0_z(r, phi, z):
    return (3 * np.sin(((z - z_i)/(z_f - z_i))*np.pi)) + (2.5 * np.sin(((z - z_i)/(z_f - z_i))*2*np.pi)) + (1 * np.sin(((z - z_i)/(z_f - z_i))*5*np.pi))

def B_phi0_r(r, phi, z):
    return (3 * np.sin(((r - r_i)/(r_f - r_i))*np.pi)) + (1 * np.sin(((r - r_i)/(r_f - r_i))*2*np.pi)) + (2 * np.sin(((r - r_i)/(r_f - r_i))*5*np.pi))
def B_r0_r(r, phi, z):
    return (3 * np.sin(((r - r_i)/(r_f - r_i))*np.pi)) + (2.5 * np.sin(((r - r_i)/(r_f - r_i))*2*np.pi)) + (1 * np.sin(((r - r_i)/(r_f - r_i))*5*np.pi))

B_list = simulate_B_z_field(z_i, z_f, Nz, T, Nt_z, B_r0_z, B_phi0_z, dBdt_z, order, ghost_zone_type, iterations_to_save=iterations_to_save)
B_list2 = simulate_B_r_field(r_i, r_f, Nr, T, Nt_r, B_r0_r, B_phi0_r, dBdt_r, order, ghost_zone_type, iterations_to_save=iterations_to_save2)
  0%|          | 16/5000 [00:00<00:32, 152.71it/s]100%|██████████| 5000/5000 [00:29<00:00, 171.81it/s]
100%|██████████| 4000/4000 [00:22<00:00, 176.10it/s]
In [ ]:
# plotting the magnetic fields at different time steps (figures side by side)
plt.suptitle('Magnetic field $|B_{\phi}|$ and $|B_{r}|$ at different time steps for both z and r implementations')
fig, axs = plt.subplots(2, 2, figsize=(20, 10))

# z-implementation B-phi
axs[0, 0].set_title(r'Magnetic field $|B_{\phi}|$ at different time steps for z-implementation')
axs[0, 0].set_xlabel(r'z')
axs[0, 0].set_ylabel(r'$B_{\phi}$')
for i in range(len(iterations_to_save)):
    axs[0, 0].plot(np.linspace(z_i, z_f, Nz), B_list[i, 1, :], label=f't = {iterations_to_save[i] * T / Nt_z:.4f}')
axs[0, 0].grid(True)
axs[0, 0].legend()

# r-implementation B-phi
axs[0, 1].set_title(r'Magnetic field $|B_{\phi}|$ at different time steps for r-implementation')
axs[0, 1].set_xlabel(r'r')
axs[0, 1].set_ylabel(r'$B_{\phi}$')
for i in range(len(iterations_to_save2)):
    axs[0, 1].plot(np.linspace(r_i, r_f, Nr), B_list2[i, 1, :], label=f't = {iterations_to_save2[i] * T / Nt_r:.4f}')
axs[0, 1].grid(True)
axs[0, 1].legend()

# z-implementation B_r
axs[1, 0].set_title(r'Magnetic field $|B_{r}|$ at different time steps for z-implementation')
axs[1, 0].set_xlabel(r'z')
axs[1, 0].set_ylabel(r'$B_{r}$')
for i in range(len(iterations_to_save)):
    axs[1, 0].plot(np.linspace(z_i, z_f, Nz), B_list[i, 0, :], label=f't = {iterations_to_save[i] * T / Nt_z:.4f}')
axs[1, 0].grid(True)
axs[1, 0].legend()

# r-implementation B_r
axs[1, 1].set_title(r'Magnetic field $|B_{r}|$ at different time steps for r-implementation')
axs[1, 1].set_xlabel(r'r')
axs[1, 1].set_ylabel(r'$B_{r}$')
for i in range(len(iterations_to_save2)):
    axs[1, 1].plot(np.linspace(r_i, r_f, Nr), B_list2[i, 0, :], label=f't = {iterations_to_save2[i] * T / Nt_r:.4f}')
axs[1, 1].grid(True)
axs[1, 1].legend()
if fig_save:
    plt.savefig(fig_path + "Bphi_Br_z_r_evolution_trial1.png")
plt.show()
<Figure size 640x480 with 0 Axes>
No description has been provided for this image
In [ ]:
# trial 2 (cos wave and non-zero boundary condition)

def B_phi0_z(r, phi, z):
    return (3 * np.cos(((z - z_i)/(z_f - z_i))*np.pi)) + (1 * np.cos(((z - z_i)/(z_f - z_i))*2*np.pi)) + (2 * np.cos(((z - z_i)/(z_f - z_i))*5*np.pi))
def B_r0_z(r, phi, z):
    return (3 * np.cos(((z - z_i)/(z_f - z_i))*np.pi)) + (2.5 * np.cos(((z - z_i)/(z_f - z_i))*2*np.pi)) + (1 * np.cos(((z - z_i)/(z_f - z_i))*5*np.pi))

def B_phi0_r(r, phi, z):
    return (3 * np.cos(((r - r_i)/(r_f - r_i))*np.pi)) + (1 * np.cos(((r - r_i)/(r_f - r_i))*2*np.pi)) + (2 * np.cos(((r - r_i)/(r_f - r_i))*5*np.pi))
def B_r0_r(r, phi, z):
    return (3 * np.cos(((r - r_i)/(r_f - r_i))*np.pi)) + (2.5 * np.cos(((r - r_i)/(r_f - r_i))*2*np.pi)) + (1 * np.cos(((r - r_i)/(r_f - r_i))*5*np.pi))

B_list = simulate_B_z_field(z_i, z_f, Nz, T, Nt_z, B_r0_z, B_phi0_z, dBdt_z, order, ghost_zone_type, iterations_to_save=iterations_to_save)
B_list2 = simulate_B_r_field(r_i, r_f, Nr, T, Nt_r, B_r0_r, B_phi0_r, dBdt_r, order, ghost_zone_type, iterations_to_save=iterations_to_save2)
  0%|          | 0/5000 [00:00<?, ?it/s]100%|██████████| 5000/5000 [00:27<00:00, 184.53it/s]
100%|██████████| 4000/4000 [00:24<00:00, 164.86it/s]
In [ ]:
# plotting the magnetic fields at different time steps (figures side by side)

plt.suptitle('Magnetic field $|B_{\phi}|$ and $|B_{r}|$ at different time steps for both z and r implementations')
fig, axs = plt.subplots(2, 2, figsize=(20, 10))

# z-implementation B-phi
axs[0, 0].set_title(r'Magnetic field $|B_{\phi}|$ at different time steps for z-implementation')
axs[0, 0].set_xlabel(r'z')
axs[0, 0].set_ylabel(r'$B_{\phi}$')
for i in range(len(iterations_to_save)):
    axs[0, 0].plot(np.linspace(z_i, z_f, Nz), B_list[i, 1, :], label=f't = {iterations_to_save[i] * T / Nt_z:.4f}')
axs[0, 0].grid(True)
axs[0, 0].legend()

# r-implementation B-phi
axs[0, 1].set_title(r'Magnetic field $|B_{\phi}|$ at different time steps for r-implementation')
axs[0, 1].set_xlabel(r'r')
axs[0, 1].set_ylabel(r'$B_{\phi}$')
for i in range(len(iterations_to_save2)):
    axs[0, 1].plot(np.linspace(r_i, r_f, Nr), B_list2[i, 1, :], label=f't = {iterations_to_save2[i] * T / Nt_r:.4f}')
axs[0, 1].grid(True)
axs[0, 1].legend()

# z-implementation B_r
axs[1, 0].set_title(r'Magnetic field $|B_{r}|$ at different time steps for z-implementation')
axs[1, 0].set_xlabel(r'z')
axs[1, 0].set_ylabel(r'$B_{r}$')
for i in range(len(iterations_to_save)):
    axs[1, 0].plot(np.linspace(z_i, z_f, Nz), B_list[i, 0, :], label=f't = {iterations_to_save[i] * T / Nt_z:.4f}')
axs[1, 0].grid(True)
axs[1, 0].legend()

# r-implementation B_r
axs[1, 1].set_title(r'Magnetic field $|B_{r}|$ at different time steps for r-implementation')
axs[1, 1].set_xlabel(r'r')
axs[1, 1].set_ylabel(r'$B_{r}$')
for i in range(len(iterations_to_save2)):
    axs[1, 1].plot(np.linspace(r_i, r_f, Nr), B_list2[i, 0, :], label=f't = {iterations_to_save2[i] * T / Nt_r:.4f}')
axs[1, 1].grid(True)
axs[1, 1].legend()
if fig_save:
    plt.savefig(fig_path + "Bphi_Br_z_r_evolution_trial2.png")
plt.show()
<Figure size 640x480 with 0 Axes>
No description has been provided for this image

Attempt at extreme generalization of sim_B (to be updated)¶

In [ ]:
# defining spatial deritivative functions
def spatial_derivative(B, dx, order = 6, ghost_zone_type = 'relative anti-symmetric'):
    '''
    This function calculates the spatial derivatives of the magnetic field B.
    
    Parameters:
    B: 2D array
        The magnetic field B.
    dx: float
        The spatial step size.
    order: int
        The order of the finite difference scheme. It can be 2, 4, 6, 8 or 10.
        Default is 6.
    ghost_zone_type: str
        The type of the ghost zone. It can be 'anti-symmetric', 'symmetric', 'relative anti-symmetric' or None.
        Default is 'relative anti-symmetric'.
        
    Returns:
    dBdx: 2D array
        The spatial-derivative of the magnetic field B.
    ddBddx: 2D array
        The second z-derivative of the magnetic field B.
    '''

    dBdx = np.zeros(B.shape)
    ddBddx = np.zeros(B.shape)

    # print(f"The shape of B is {B.shape}")
    # print(f"The shape of dBdx is {dBdx.shape}")
    # print(f"The shape of ddBddx is {ddBddx.shape}")
    
    B = np.copy(B) # to avoid modifying the original array
    
    if order not in ['np',2, 4, 6, 8, 10]:
        raise ValueError("Order should be 'np',2, 4, 6, 8 or 10")
    
    if ghost_zone_type not in ['anti-symmetric', 'symmetric', 'relative anti-symmetric', None]:
        raise ValueError('ghost_zone_type should be anti-symmetric, symmetric, relative anti-symmetric or None')

    if order == 'np':
        ghost_zone = 3 # 3 cells on each side
    else:
        ghost_zone = int(order/2)
    
    if ghost_zone_type == 'relative anti-symmetric':
        B = np.pad(B, ((0, 0), (ghost_zone, ghost_zone)), mode='constant')
        for i in range(ghost_zone):
            B[:,i] = (2 * B[:,ghost_zone]) - B[:,ghost_zone + ghost_zone - i] # correcting the start of the array
            B[:, -1 - i] = 2 * B[:,-ghost_zone - 1] - B[:, (- ghost_zone - 1) - ghost_zone + i] # correcting the end of the array
    elif ghost_zone_type == 'None':
        B = np.pad(B, ((0, 0), (ghost_zone, ghost_zone)), mode='constant')
                
        # update with other types of ghost zones

    # print(f"The shape of B after padding is {B.shape}")
    # print(f"The shape of dBdx after padding is {dBdx.shape}")
    # print(f"The shape of ddBddx after padding is {ddBddx.shape}")
    
    if order == 6:
        for i in range(ghost_zone, B.shape[1] - ghost_zone):
            dBdx[:,i-ghost_zone] = ((- B[:,i - 3]) + (9 * B[:,i - 2]) - (45 * B[:,i - 1]) + (45 * B[:,i + 1]) - (9 * B[:,i + 2]) + (B[:,i + 3])) / (60 * dx)
            ddBddx[:,i-ghost_zone] = ((2 * B[:,i - 3]) - (27 * B[:,i - 2]) + (270 * B[:,i - 1]) - (490 * B[:, i]) + (270 * B[:,i + 1]) - (27 * B[:,i + 2]) + (2 * B[:,i + 3])) / (180 * (dx ** 2))
    elif order == 'np':
        dBdx = np.gradient(B, dx, axis=1)
        ddBddx = np.gradient(dBdx, dx, axis=1)
        
        # removing the ghost zones
        dBdx = dBdx[:,ghost_zone:-ghost_zone]
        ddBddx = ddBddx[:,ghost_zone:-ghost_zone]
            
    # update with other orders
    
    return dBdx, ddBddx

def simulate_B_field(r_i, r_f, Nr, T, Nt, B_r0, B_phi0, dBdt_r, order = 6, ghost_zone_type = 'relative anti-symmetric', iterations_to_save = None):
    '''    
    This function simulates the time evolution of the magnetic field B.
    
    Parameters:
    z_i: float
        The start of the spatial domain.
    z_f: float
        The end of the spatial domain.
    Nr: int
        The number of spatial grid points.
    T: float
        The total simulation time.
    Nt: int
        The number of time steps.
    B_r0: function
        The initial condition for the magnetic field B_r.
    B_phi0: function
        The initial condition for the magnetic field B_phi.
    dBdt_r: function
        The time evolution of the magnetic field B.
    order: int
        The order of the finite difference scheme. It can be 'np', 2, 4, 6, 8 or 10.
        Default is 6.
    ghost_zone_type: str
        The type of the ghost zone. It can be 'anti-symmetric', 'symmetric', 'relative anti-symmetric' or None.
        Default is 'relative anti-symmetric'.
    iterations_to_save: list
        A list of the time steps to save the solution. If None, it will save the solution at the start and end of the simulation.
        
    Returns:
    B_list: 3-D array
        A list of the magnetic field B at the required time steps from iterations_to_save.
        size: (len(iterations_to_save), 2, Nz)
    '''
    dr = (r_f - r_i) / (Nr - 1)  # Spatial step size
    dt = T / Nt  # Temporal step size
    
    # Create arrays to store the solution in time
    B_list = []

    # creating initial conditions from z, B_r0 and B_phi0
    r = np.linspace(r_i, r_f, Nr)
    r_ = np.array([r,r]) # Under the assumption that the equation is same for both B_r and B_phi
                         # We need this r_ for the equation
    
    B = np.zeros((2, Nr))
    B[0, :] = B_r0(r, 0, 0) # first row is B_r
    B[1, :] = B_phi0(r, 0, 0) # second row is B_phi
    
    if iterations_to_save is None:
        iterations_to_save = [0, Nt]

    if 0 in iterations_to_save:
        B_list.append(np.copy(B))
    
    #--------- RK4 time-stepping block ---------#
    for n in tqdm(range(1, Nt + 1)):
        # Compute spatial derivatives
        dBdr, ddBddr = spatial_derivative(B, dr, order, ghost_zone_type)


        k1 = dBdt_r(B, r_, dBdr, ddBddr, n * dt)
        k2 = dBdt_r(B + (0.5 * dt * k1), r_, dBdr, ddBddr, (n * dt) + (0.5 * dt))
        k3 = dBdt_r(B + (0.5 * dt * k2), r_, dBdr, ddBddr, (n * dt) + (0.5 * dt))
        k4 = dBdt_r(B + (dt * k3), r_, dBdr, ddBddr, (n * dt) + dt)
        
        B = B + (dt / 6.0) * (k1 + (2 * k2) + (2 * k3) + k4)
        
        # print(B)
        
        if n in iterations_to_save:
            B_list.append(np.copy(B))
    #------------------------------------------#
    
    B_list = np.array(B_list)
    
    return B_list