import numpy as np
from .misc import xy2rp, map_matrix, toroidal_period, div0
[docs]class Dipole(object):
"""
magnetic dipole class
"""
def __init__(self, **kwargs):
"""
Initialize empty class
"""
self.num = 0 # total number of dipoles
self.nfp = 1 # toroidal_periodicity
self.momentq = 1 # q expotent
self.sp_switch = False # switch to indicate if using spherical coordinates
self.xyz_switch = False # switch to indicate if using spherical coordinates
self.old = False # old format or new
# self.symmetry = 2 # 0: no symmetry; 1: periodicity; 2: stellarator symmetry (+periodicity)
if "mm" in kwargs: # spherical coord
self.ox = kwargs["ox"]
self.oy = kwargs["oy"]
self.oz = kwargs["oz"]
self.mm = kwargs["mm"]
self.mt = kwargs["mt"]
self.mp = kwargs["mp"]
self.pho = kwargs["pho"]
self.momentq = kwargs["momentq"]
self.Ic = kwargs["Ic"]
self.Lc = kwargs["Lc"]
self.num = len(self.ox)
self.sp_switch = True
self.filename = kwargs.get("filename", "dipole")
self.name = kwargs.get(
"name", ["pm_{:010d}".format(i) for i in range(1, self.num + 1)]
)
self.symm = kwargs.get("symm", 2 * np.ones(self.num, dtype=int))
self.rho = self.pho ** self.momentq
elif "mx" in kwargs: # cartesian coord
self.ox = kwargs["ox"]
self.oy = kwargs["oy"]
self.oz = kwargs["oz"]
self.mx = kwargs["mx"]
self.my = kwargs["my"]
self.mz = kwargs["mz"]
self.mm = np.sqrt(self.mx ** 2 + self.my ** 2 + self.mz ** 2)
self.num = len(self.ox)
self.pho = np.ones(self.num)
self.Ic = np.zeros(self.num, dtype=int)
self.Lc = np.zeros(self.num, dtype=int)
self.xyz_switch = True
self.filename = kwargs.get("filename", "dipole")
self.name = kwargs.get(
"name", ["pm_{:010d}".format(i) for i in range(1, self.num + 1)]
)
self.symm = kwargs.get("symm", 2 * np.ones(self.num, dtype=int))
self.rho = np.ones(self.num)
return
[docs] @classmethod
def open(cls, filename, verbose=False, **kwargs):
"""Read FAMUS dipoles
Args:
filename (str): path to open the file.
verbose (bool, optional): Whether to print additional info. Defaults to False.
Returns:
Dipole: Dipole class
"""
import pandas as pd
with open(filename, "r") as coilfile:
coilfile.readline()
line = coilfile.readline()
line = line.replace(",", " ")
line = line.split()
num = int(line[0])
try:
momentq = int(line[1])
except:
print("Moment Q factor was not read. Default=1.")
momentq = 1
data = pd.read_csv(filename, skiprows=3, header=None)
symm = np.array(data[1], dtype=int)[0:num]
name = np.array(data[2])[0:num] # type string
ox = np.array(data[3], dtype=float)[0:num]
oy = np.array(data[4], dtype=float)[0:num]
oz = np.array(data[5], dtype=float)[0:num]
Ic = np.array(data[6], dtype=int)[0:num]
mm = np.array(data[7], dtype=float)[0:num]
pho = np.array(data[8], dtype=float)[0:num]
Lc = np.array(data[9], dtype=int)[0:num]
mp = np.array(data[10], dtype=float)[0:num]
mt = np.array(data[11], dtype=float)[0:num]
if verbose:
print("Read {:d} dipoles from {:}".format(len(ox), filename))
return cls(
symm=symm,
name=name,
ox=ox,
oy=oy,
oz=oz,
Ic=Ic,
mm=mm,
Lc=Lc,
mp=mp,
mt=mt,
pho=pho,
momentq=momentq,
filename=filename,
)
[docs] @classmethod
def read_dipole_old(cls, filename, zeta=0.0, zeta1=np.pi * 2, **kwargs):
"""read diploes from FOCUS format (old)
Args:
filename (str): path to open the file.
zeta (float, optional): Starting toroidal angle for clipping. Defaults to 0.0.
zeta1 (float, optional): [Ending toroidal angle for clipping]. Defaults to np.pi*2.
Returns:
Dipole: Dipole class
"""
symm = []
name = []
ox = []
oy = []
oz = []
mm = []
mp = []
mt = []
Ic = []
Lc = []
with open(filename, "r") as coilfile:
coilfile.readline()
Ncoils = int(coilfile.readline()[0])
for icoil in range(Ncoils):
coilfile.readline()
coilfile.readline()
linelist = coilfile.readline().split()
if int(linelist[0]) == 3: # central current and background Bz
coilfile.readline()
linelist = coilfile.readline().split()
elif int(linelist[0]) == 2: # dipoles
symm.append(int(linelist[1]))
name.append(linelist[2])
coilfile.readline()
linelist = coilfile.readline().split()
r, phi = xy2rp(float(linelist[1]), float(linelist[2]))
if phi >= zeta and phi <= zeta1:
Lc.append(int(linelist[0]))
ox.append(float(linelist[1]))
oy.append(float(linelist[2]))
oz.append(float(linelist[3]))
Ic.append(int(linelist[4]))
mm.append(float(linelist[5]))
mt.append(float(linelist[6]))
mp.append(float(linelist[7]))
elif int(linelist[0]) == 1: # Fourier coils
for i in range(11):
coilfile.readline()
else:
raise ValueError(
"Invalid coiltype = {:d}.".format(int(linelist[0]))
)
nc = len(ox)
if nc == 0:
print("Warning: no dipoles was read from " + filename)
return
print(
"Read {:d} dipoles from {:}. Please manually set self.old=True.".format(
nc, filename
)
)
symm = np.array(symm)
ox = np.array(ox)
oy = np.array(oy)
oz = np.array(oz)
mm = np.array(mm)
mt = np.array(mt)
mp = np.array(mp)
Lc = np.array(Lc)
Ic = np.array(Ic)
pho = np.ones_like(ox)
momentq = 1
return cls(
symm=symm,
name=name,
filename=filename,
ox=ox,
oy=oy,
oz=oz,
Ic=Ic,
mm=mm,
Lc=Lc,
mp=mp,
mt=mt,
pho=pho,
momentq=momentq,
)
[docs] @classmethod
def from_regcoil(
cls,
regcoilname,
winding,
ilambda=-1,
symmetry="full",
num_pol=128,
num_tor=128,
m0=None,
half_shift=True,
):
"""Initialize from REGCOIL current potential
Args:
regcoilname (str): REGCOIL netcdf output file.
winding (str, optional): NESCOIL input format winding surface.
ilambda (int, optional): Lambda index in REGCOIL output. Defaults to -1.
symmetry (str, optional): Stellarator symmetry option. Defaults to 'full'.
num_pol (int, optional): Number of poloidal dipoles. Defaults to 128.
num_tor (int, optional): Number of toroidal dipoles. Defaults to 128.
m0 ([type], optional): Magnetization limit per thickness. Defaults to None.
half_shift (bool, optional): Logical flag to determine if half-grid shifted. Defaults to True.
Returns:
None
"""
# read regcoil output
from scipy.io import netcdf
f = netcdf.netcdf_file(regcoilname, "r", mmap=False)
nfp = f.variables["nfp"][()]
xm_potential = f.variables["xm_potential"][()]
xn_potential = f.variables["xn_potential"][()]
phi_mn = f.variables["single_valued_current_potential_mn"][()][ilambda, :]
symmetry_option = f.variables["symmetry_option"][()]
f.close()
mn_max = len(xm_potential)
phi_sin = np.zeros(mn_max)
phi_cos = np.zeros(mn_max)
if symmetry_option == 1:
phi_sin = phi_mn[0:mn_max]
elif symmetry_option == 2:
phi_cos = phi_mn[0:mn_max]
elif symmetry_option == 3:
phi_sin = phi_mn[0:mn_max]
phi_cos = phi_mn[mn_max:]
else:
raise ValueError(
"Something wrong the symmetry_option: {:}".format(symmetry_option)
)
def func_new(theta, zeta):
"""get current potential of (theta, zeta)
Parameters:
theta -- float array_like, poloidal angle
zeta -- float array_like, toroidal angle value
Returns:
phi -- float array_like
"""
assert len(np.atleast_1d(theta)) == len(
np.atleast_1d(zeta)
), "theta, zeta should be equal size"
# mt - nz (in matrix)
_mtnz = np.matmul(
np.reshape(xm_potential, (-1, 1)), np.reshape(theta, (1, -1))
) - np.matmul(np.reshape(xn_potential, (-1, 1)), np.reshape(zeta, (1, -1)))
_sin = np.sin(_mtnz)
_cos = np.cos(_mtnz)
phi = np.matmul(np.reshape(phi_sin, (1, -1)), _sin) + np.matmul(
np.reshape(phi_cos, (1, -1)), _cos
)
return phi.ravel()
# initialize dipoles
num = num_pol * num_tor
if symmetry.lower() == "full":
zeta_end = 2 * np.pi
symm = 0 + np.zeros(num, dtype=np.int)
elif symmetry.lower() == "one":
zeta_end = 2 * np.pi / nfp
symm = 1 + np.zeros(num, dtype=np.int)
elif symmetry.lower() == "half":
zeta_end = 2 * np.pi / nfp / 2
symm = 2 + np.zeros(num, dtype=np.int)
else:
raise ValueError("No such symmetry option!")
ox = []
oy = []
oz = []
mt = []
mp = []
mm = []
# discretization
if half_shift: # toroidally shift half grid
half = zeta_end / (num_tor * 2)
else:
half = 0
theta_array = np.linspace(0, 2 * np.pi, num_pol, endpoint=False)
zeta_array = half + np.linspace(0, zeta_end, num_tor, endpoint=False)
dtheta = 2 * np.pi / num_pol
dzeta = zeta_end / num_tor
tv, zv = np.meshgrid(theta_array, zeta_array, indexing="ij")
from .surface import FourSurf
wind = FourSurf.read_winding_surfce(winding)
data = wind.xyz(tv, zv, normal=True)
phi = func_new(tv, zv)
ox = data[0]
oy = data[1]
oz = data[2]
n = data[3]
nn = np.linalg.norm(n, axis=1)
mt = np.arccos(n[:, 2] / nn)
mp = np.arctan2(n[:, 1], n[:, 0])
if m0 is None: # not specify the allowable magnetization
mm = -phi * nn * dtheta * dzeta
pho = np.ones_like(mm)
else:
mm = m0 * nn * dtheta * dzeta
pho = -phi / m0
return cls(
symm=symm,
ox=ox,
oy=oy,
oz=oz,
mp=mp,
mt=mt,
mm=mm,
pho=pho,
Ic=np.ones_like(ox, dtype=int),
Lc=np.ones_like(ox, dtype=int),
momentq=1,
)
[docs] def sp2xyz(self):
"""
spherical coordinates to cartesian coordinates
"""
assert self.sp_switch == True, "You are not using spherical coordinates"
if self.old:
self.mx = self.mm * np.sin(self.mt) * np.cos(self.mp)
self.my = self.mm * np.sin(self.mt) * np.sin(self.mp)
self.mz = self.mm * np.cos(self.mt)
else:
self.mx = (
self.mm * self.pho ** self.momentq * np.sin(self.mt) * np.cos(self.mp)
)
self.my = (
self.mm * self.pho ** self.momentq * np.sin(self.mt) * np.sin(self.mp)
)
self.mz = self.mm * self.pho ** self.momentq * np.cos(self.mt)
self.xyz_switch = True
return
[docs] def xyz2sp(self):
"""
cartesian coordinates to spherical coordinates
"""
assert self.xyz_switch is True, "You are not using cartesian coordinates"
if self.old:
self.mm = np.sqrt(self.mx * self.mx + self.my * self.my + self.mz * self.mz)
div = np.divide(
self.mz, self.mm, out=np.zeros_like(self.mz), where=self.mm != 0
)
self.mt = np.arccos(div0(self.mz, self.mm))
# self.pho = np.ones_like(self.mm)
# self.momentq = 1
else:
mmag = np.sqrt(self.mx * self.mx + self.my * self.my + self.mz * self.mz)
self.pho = np.power(div0(mmag, self.mm), 1.0 / self.momentq)
self.mt = np.arccos(div0(self.mz, (self.mm * self.pho ** self.momentq)))
self.mp = np.arctan2(self.my, self.mx)
self.sp_switch = True
return
[docs] def save(self, filename, unique=False, tol=0):
"""write diploes from FOCUS format
Args:
filename (str): FOCUS file name.
unique (bool, optional): Writing dipole every self.nfp term. Defaults to False.
tol (float, optional): tolerance to skip zeros. Defaults to 0.
"""
if not self.sp_switch:
self.xyz2sp()
cond = np.abs(self.rho) >= tol
with open(filename, "w") as wfile:
wfile.write(" # Total number of dipoles, momentq \n")
if unique:
wfile.write(
"{:6d}, {:4d}\n".format(
np.count_nonzero(cond) / self.nfp, self.momentq
)
)
else:
wfile.write(
"{:6d}, {:4d}\n".format(np.count_nonzero(cond), self.momentq)
)
if self.old:
for icoil in range(self.num):
if unique:
if np.mod(icoil, self.nfp) == 0:
continue
if not cond[icoil]:
continue
wfile.write(
"#-----------------{}---------------------------\n".format(
icoil + 1
)
)
wfile.write("#coil_type symm coil_name \n")
wfile.write(
" {:1d} {:1d} {:} \n".format(
2, self.symm[icoil], self.name[icoil]
)
)
wfile.write("# Lc ox oy oz Ic I mt mp \n")
wfile.write(
"{:6d} {:23.15E} {:23.15E} {:23.15E} {:6d} {:23.15E} {:23.15E} {:23.15E}\n".format(
self.Lc[icoil],
self.ox[icoil],
self.oy[icoil],
self.oz[icoil],
self.Ic[icoil],
self.mm[icoil],
self.mt[icoil],
self.mp[icoil],
)
)
else:
wfile.write(
"#coiltype, symmetry, coilname, ox, oy, oz, Ic, M_0, pho, Lc, mp, mt \n"
)
for i in range(self.num):
if unique:
if np.mod(i, self.nfp) == 0:
continue
if not cond[i]:
continue
wfile.write(
" 2, {:1d}, {:}, {:15.8E}, {:15.8E}, {:15.8E}, {:2d}, {:15.8E},"
"{:15.8E}, {:2d}, {:15.8E}, {:15.8E} \n".format(
self.symm[i],
self.name[i],
self.ox[i],
self.oy[i],
self.oz[i],
self.Ic[i],
self.mm[i],
self.pho[i],
self.Lc[i],
self.mp[i],
self.mt[i],
)
)
return
[docs] def toVTK(
self, vtkname, dim=(1), close=False, ntnz=False, toroidal=False, **kwargs
):
"""write dipole data into a VTK file
Args:
vtkname (str): VTK filename, will be appended with .vts or .vtu.
dim (tuple, optional): Dimension information if saved as structured grids. Defaults to (1).
close (bool, optional): Logical flag to manually close the gaps. Defaults to False.
ntnz (bool, optional): Logical flag of theta-zeta order. Defaults to False.
toroidal (bool, optional): Logical flag of filling in the toroidal gap. Defaults to False.
Returns:
None
"""
from pyevtk.hl import gridToVTK, pointsToVTK
if not self.xyz_switch:
self.sp2xyz()
if vtkname is None:
vtkname = self.filename
dim = np.atleast_1d(dim)
if len(dim) == 1: # save as points
print("write VTK as points")
data = {"m": (self.mx, self.my, self.mz)}
if not self.old:
data.update({"rho": self.pho ** self.momentq})
data.update(kwargs)
pointsToVTK(
vtkname, self.ox, self.oy, self.oz, data=data
) # .update(kwargs))
else: # save as surfaces
assert len(dim) == 3
print("write VTK as closed surface")
if close:
# manually close the gap
phi = 2 * np.pi / self.nfp
def map_toroidal(vec):
rotate = np.array(
[
[np.cos(phi), np.sin(phi), 0],
[-np.sin(phi), np.cos(phi), 0],
[0, 0, 1],
]
)
return np.matmul(vec, rotate)
data_array = {
"ox": self.ox,
"oy": self.oy,
"oz": self.oz,
"mx": self.mx,
"my": self.my,
"mz": self.mz,
"Ic": self.Ic,
"rho": self.pho ** self.momentq,
}
data_array.update(kwargs)
nr, nz, nt = dim
for key in list(data_array.keys()):
new_vec = np.zeros((nr, nz + 1, nt + 1))
for ir in range(nr):
new_vec[ir, :, :] = map_matrix(
np.reshape(data_array[key], dim)[ir, :, :]
)
if toroidal:
data_array[key] = new_vec
else:
if ntnz:
data_array[key] = np.ascontiguousarray(new_vec[:, :, :-1])
else:
data_array[key] = np.ascontiguousarray(new_vec[:, :-1, :])
ox = np.copy(data_array["ox"])
oy = np.copy(data_array["oy"])
oz = np.copy(data_array["oz"])
del data_array["ox"]
del data_array["oy"]
del data_array["oz"]
data_array["m"] = (data_array["mx"], data_array["my"], data_array["mz"])
if toroidal and self.nfp >= 1: # not quite sure if should include nfp=1
for ir in range(nr):
if ntnz:
xyz = map_toroidal(
np.transpose([ox[ir, :, 0], oy[ir, :, 0], oz[ir, :, 0]])
)
ox[ir, :, nz] = xyz[:, 0]
oy[ir, :, nz] = xyz[:, 1]
oz[ir, :, nz] = xyz[:, 2]
moment = map_toroidal(
np.transpose(
[
data_array["mx"][ir, :, 0],
data_array["my"][ir, :, 0],
data_array["mz"][ir, :, 0],
]
)
)
data_array["m"][0][ir, :, nz] = moment[:, 0]
data_array["m"][1][ir, :, nz] = moment[:, 1]
data_array["m"][2][ir, :, nz] = moment[:, 2]
else:
xyz = map_toroidal(
np.transpose([ox[ir, 0, :], oy[ir, 0, :], oz[ir, 0, :]])
)
ox[ir, nz, :] = xyz[:, 0]
oy[ir, nz, :] = xyz[:, 1]
oz[ir, nz, :] = xyz[:, 2]
moment = map_toroidal(
np.transpose(
[
data_array["mx"][ir, 0, :],
data_array["my"][ir, 0, :],
data_array["mz"][ir, 0, :],
]
)
)
data_array["m"][0][ir, nz, :] = moment[:, 0]
data_array["m"][1][ir, nz, :] = moment[:, 1]
data_array["m"][2][ir, nz, :] = moment[:, 2]
del data_array["mx"]
del data_array["my"]
del data_array["mz"]
gridToVTK(vtkname, ox, oy, oz, pointData=data_array)
return
else:
ox = np.reshape(self.ox[: self.num], dim)
oy = np.reshape(self.oy[: self.num], dim)
oz = np.reshape(self.oz[: self.num], dim)
mx = np.reshape(self.mx[: self.num], dim)
my = np.reshape(self.my[: self.num], dim)
mz = np.reshape(self.mz[: self.num], dim)
rho = np.reshape(self.pho[: self.num] ** self.momentq, dim)
Ic = np.reshape(self.Ic[: self.num], dim)
data = {"m": (mx, my, mz), "rho": rho, "Ic": Ic}
data.update(kwargs)
gridToVTK(vtkname, ox, oy, oz, pointData=data)
return
[docs] def full_period(self, nfp=1, dim=None):
"""
map from one period to full periods
"""
assert nfp >= 1
self.nfp = nfp
if not self.xyz_switch:
self.sp2xyz()
# change order
if dim is not None:
self.ox = np.ravel(
np.transpose(np.reshape(self.ox, dim)[::-1, ::-1, ::-1], (2, 0, 1))
)
self.oy = np.ravel(
np.transpose(np.reshape(self.oy, dim)[::-1, ::-1, ::-1], (2, 0, 1))
)
self.oz = np.ravel(
np.transpose(np.reshape(self.oz, dim)[::-1, ::-1, ::-1], (2, 0, 1))
)
self.mx = np.ravel(
np.transpose(np.reshape(self.mx, dim)[::-1, ::-1, ::-1], (2, 0, 1))
)
self.my = np.ravel(
np.transpose(np.reshape(self.my, dim)[::-1, ::-1, ::-1], (2, 0, 1))
)
self.mz = np.ravel(
np.transpose(np.reshape(self.mz, dim)[::-1, ::-1, ::-1], (2, 0, 1))
)
self.mm = np.ravel(
np.transpose(np.reshape(self.mm, dim)[::-1, ::-1, ::-1], (2, 0, 1))
)
self.Ic = np.ravel(
np.transpose(np.reshape(self.Ic, dim)[::-1, ::-1, ::-1], (2, 0, 1))
)
self.Lc = np.ravel(
np.transpose(np.reshape(self.Lc, dim)[::-1, ::-1, ::-1], (2, 0, 1))
)
self.pho = np.ravel(
np.transpose(np.reshape(self.pho, dim)[::-1, ::-1, ::-1], (2, 0, 1))
)
if np.mean(self.symm) == 2:
# get the stellarator symmetry part first
# Here, we assume no dipoles on the symmetry plane, or only half are listed.
self.num *= 2
self.name = np.concatenate((self.name, self.name[::-1]))
self.ox = np.concatenate((self.ox, self.ox[::-1]))
self.oy = np.concatenate((self.oy, self.oy[::-1] * (-1)))
self.oz = np.concatenate((self.oz, self.oz[::-1] * (-1)))
self.mx = np.concatenate((self.mx, self.mx[::-1] * (-1)))
self.my = np.concatenate((self.my, self.my[::-1]))
self.mz = np.concatenate((self.mz, self.mz[::-1]))
self.mm = np.concatenate((self.mm, self.mm[::-1]))
self.pho = np.concatenate((self.pho, self.pho[::-1]))
self.Ic = np.concatenate((self.Ic, self.Ic[::-1]))
self.Lc = np.concatenate((self.Lc, self.Lc[::-1]))
xyz = toroidal_period(np.transpose([self.ox, self.oy, self.oz]), self.nfp)
self.ox = xyz[:, 0].copy()
self.oy = xyz[:, 1].copy()
self.oz = xyz[:, 2].copy()
moment = toroidal_period(np.transpose([self.mx, self.my, self.mz]), self.nfp)
self.mx = moment[:, 0].copy()
self.my = moment[:, 1].copy()
self.mz = moment[:, 2].copy()
self.mm = np.tile(self.mm, self.nfp)
self.pho = np.tile(self.pho, self.nfp)
self.rho = self.pho ** self.momentq
self.Ic = np.tile(self.Ic, self.nfp)
self.Lc = np.tile(self.Lc, self.nfp)
self.name = np.tile(self.name, self.nfp)
self.num *= self.nfp
self.symm = np.zeros(self.num, dtype=int)
self.sp_switch = False
return
[docs] def inverse(self):
"""get the stellarator symmetric part?"""
if not self.xyz_switch:
self.sp2xyz()
self.ox = np.copy(self.ox[::-1])
self.oy = np.copy(self.oy[::-1] * (-1))
self.oz = np.copy(self.oz[::-1] * (-1))
self.mx = np.copy(self.mx[::-1] * (-1))
self.my = np.copy(self.my[::-1])
self.mz = np.copy(self.mz[::-1])
self.mm = np.copy(self.mm[::-1])
self.pho = np.copy(self.pho[::-1])
self.Ic = np.copy(self.Ic[::-1])
self.Lc = np.copy(self.Lc[::-1])
return
[docs] def change_momentq(self, newq):
"""
change the q factor for the normalized density
"""
assert newq > 0
pho = self.pho ** self.momentq
# get signs
self.sp2xyz()
sign = np.sign(pho)
if newq % 2 == 0: # even number, flip the orientation when negative
# self.mx *= sign
# self.my *= sign
# self.mz *= sign
self.momentq = newq
self.xyz2sp()
# convert to positive rho
# self.pho = np.power(np.abs(pho), 1.0/newq)
else: # odd exponetial index
self.momentq = newq
# convert to positive rho
self.pho = np.power(np.abs(pho), 1.0 / newq)
self.pho *= sign
return
[docs] def plot_rho_profile(self, lower=0, upper=1, nrange=10, nofigure=False, **kwargs):
import matplotlib.pyplot as plt
pho = np.abs(self.pho ** self.momentq)
zone = np.linspace(lower, upper, nrange + 1, endpoint=True)
count = []
for i in range(nrange - 1):
count.append(((pho >= zone[i]) & (pho < zone[i + 1])).sum())
count.append(((pho >= zone[nrange - 1]) & (pho <= zone[nrange])).sum())
count = np.array(count)
if not nofigure:
if plt.get_fignums():
fig = plt.gcf()
ax = plt.gca()
else:
fig, ax = plt.subplots()
plt.bar(
zone[:-1], 100 * count / float(self.num), width=0.9 / nrange, **kwargs
)
ax.set_xlabel(r"$|\rho|$", fontsize=15)
ax.set_ylabel("fraction [%]", fontsize=15)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
return count
[docs] def volume(self, magnitization=1.1e6, **kwargs):
self.total_moment = np.sum(np.abs(self.rho * self.mm))
return self.total_moment / magnitization
[docs] def orientation(self, unit=True, uniform=False):
oldq = self.momentq
self.change_momentq(2) # to a even number
self.change_momentq(oldq) # recover
if unit:
self.mm[:] = 1.0
if uniform:
self.pho[:] = 1.0
self.sp2xyz()
return
[docs] def round_angle(self):
"""round off the moment orientation to the closest preferred axis
Returns:
Dipole class: return a new dipole class with rounded orientations
"""
from copy import deepcopy
new = deepcopy(self)
phi1 = np.arctan2(new.oy, new.ox)
phi2 = phi1 + np.pi / 2
mxyz = np.transpose(
[
np.sin(new.mt) * np.cos(new.mp),
np.sin(new.mt) * np.sin(new.mp),
np.cos(new.mt),
]
)
xvec = np.transpose([np.cos(phi1), np.sin(phi1), np.zeros_like(phi1)])
yvec = np.transpose([np.cos(phi2), np.sin(phi2), np.zeros_like(phi1)])
zvec = np.reshape(np.tile([0, 0, 1], new.num), (-1, 3))
cosx = np.sum(mxyz * xvec, axis=1)
cosy = np.sum(mxyz * yvec, axis=1)
cosz = np.sum(mxyz * zvec, axis=1)
cos_arr = np.transpose([cosx[:], cosy[:], cosz[:]])
argmin = np.argmax(np.abs(cos_arr), axis=1)
theta_arr = np.reshape(np.tile([np.pi / 2, np.pi / 2, 0], new.num), (-1, 3))
phi_arr = np.transpose([phi1, phi2, np.zeros_like(phi1)])
new.mt = np.array([theta_arr[i, argmin[i]] for i in range(new.num)])
new.mp = np.array([phi_arr[i, argmin[i]] for i in range(new.num)])
sign_arr = np.sign([cos_arr[i, argmin[i]] for i in range(new.num)])
new.pho *= sign_arr
new.rho *= sign_arr
return new
[docs] def mimic(self, template, **kwargs):
"""Mimic rho, mp, mt following the "template" dipole set
Args:
template (Dipole class): The template Dipole class to be mirrored.
Returns:
None: None
"""
# find the cloest dipole in real space
def closest(i):
dis = (
(self.ox[i] - template.ox) ** 2
+ (self.oy[i] - template.oy) ** 2
+ (self.oz[i] - template.oz) ** 2
)
return np.argmin(dis)
# assign rho, mt, mp
for i in range(self.num):
ind = closest(i)
self.rho[i] = template.rho[ind]
self.mt[i] = template.mt[ind]
self.mp[i] = template.mp[ind]
self.pho = self.rho ** (1.0 / self.momentq)
return
[docs] def plot(self, engine="pyplot", start=0, end=None, **kwargs):
if end is None:
end = self.num
if engine == "pyplot":
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.scatter(
self.ox[start:end], self.oy[start:end], self.oz[start:end], **kwargs
)
elif engine == "plotly":
import plotly.express as px
fig = px.scatter_3d(
x=self.ox[start:end],
y=self.oy[start:end],
z=self.oz[start:end],
color=self.rho[start:end],
**kwargs
)
fig.show()
return
[docs] def bfield(self, pos):
"""Calculate the magnetic field at an arbitrary position.
No symmetry info considered for now.
Args:
pos (array_like): [x,y,z] Cartesian coordinates in space.
Returns:
numpy.array: The total magnetic field produced by all dipoles
"""
# calculate mx, my, mz if needed
if not self.xyz_switch:
self.sp2xyz()
# Biot-Savart law
pos = np.reshape(pos, (3, 1))
oxyz = np.asarray([self.ox, self.oy, self.oz])
rxyz = oxyz - pos
r = np.linalg.norm(rxyz, axis=0)
mxyz = np.asarray([self.mx, self.my, self.mz])
Bvec = 3 * np.sum(mxyz * rxyz, axis=0) / r ** 5 * rxyz - 1 / r ** 3 * mxyz
return 1e-7 * np.sum(Bvec, axis=1)
def __repr__(self):
return "FAMUS dipole class, num={:d}, symmetry={:}, filename={:}".format(
self.num, np.mean(self.symm), self.filename
)
def __add__(self, other):
"""Combine two dipole files together"""
if not self.sp_switch:
self.xyz2sp()
assert self.momentq == other.momentq, "Two classes should use the same momentq."
return Dipole(
symm=np.concatenate((self.symm, other.symm)),
name=np.concatenate((self.name, other.name)),
ox=np.concatenate((self.ox, other.ox)),
oy=np.concatenate((self.oy, other.oy)),
oz=np.concatenate((self.oz, other.oz)),
Ic=np.concatenate((self.Ic, other.Ic)),
mm=np.concatenate((self.mm, other.mm)),
Lc=np.concatenate((self.Lc, other.Lc)),
mp=np.concatenate((self.mp, other.mp)),
mt=np.concatenate((self.mt, other.mt)),
pho=np.concatenate((self.pho, other.pho)),
momentq=self.momentq,
filename=self.filename + "+" + other.filename,
)
[docs] def truncate(self, cond):
"""Truncate dipole set following a condition
Args:
cond (array-like, bool): Slicing conditions
Returns:
Dipole: Truncated dipole
"""
return Dipole(
symm=self.symm[cond],
name=self.name,
ox=self.ox[cond],
oy=self.oy[cond],
oz=self.oz[cond],
Ic=self.Ic[cond],
mm=self.mm[cond],
Lc=self.Lc[cond],
mp=self.mp[cond],
mt=self.mt[cond],
pho=self.pho[cond],
momentq=self.momentq,
filename=self.filename + "_clip",
)
[docs]class GAdipole(Dipole):
def __init__(self, **kwargs):
super().__init__(**kwargs)
if __name__ == "main":
pass