t_element_1d = Interpolate1DCubic(psin_1d, t_element_profile_1d)
n_element_1d = Interpolate1DCubic(psin_1d, n_element_profile_1d)
n_element2_1d = Interpolate1DCubic(psin_1d, n_element2_profile_1d)
n_tcx_donor_1d = Interpolate1DCubic(psin_1d, n_tcx_donor_profile_1d)
psin_2d = np.zeros(equilibrium.psi_data.shape)
for index0, value0 in enumerate(equilibrium.r_data):
for index1, value1 in enumerate(equilibrium.z_data):
psin_2d[index0, index1] = equilibrium.psi_normalised(value0, value1)
t_e_profile_2d = np.zeros_like(psin_2d)
for index in np.ndindex(*t_e_profile_2d.shape):
t_e_profile_2d[index] = t_e_1d(psin_2d[index])
t_e_2d = Interpolate2DCubic(equilibrium.r_data, equilibrium.z_data, t_e_profile_2d)
n_e_profile_2d = np.zeros_like(psin_2d)
for index in np.ndindex(*n_e_profile_2d.shape):
n_e_profile_2d[index] = n_e_1d(psin_2d[index])
n_e_2d = Interpolate2DCubic(equilibrium.r_data, equilibrium.z_data, n_e_profile_2d)
t_element_profile_2d = np.zeros_like(psin_2d)
for index in np.ndindex(*t_element_profile_2d.shape):
t_element_profile_2d[index] = t_element_1d(psin_2d[index])
t_element_2d = Interpolate2DCubic(equilibrium.r_data, equilibrium.z_data, t_element_profile_2d)
n_element_profile_2d = np.zeros_like(psin_2d)
for index in np.ndindex(*n_element_profile_2d.shape):
def map_psi_omp_to_divertor(x_axis_omp, divertor_coords, fiesta, location='lfs'):
"""
Function mapping the normalised psi from the specified coordinates at the
OMP to the specified coordinates at the divertor. Currently the divertor is
assumed to be represented by a 1D polynomial function, y = ax + b.
:param np.ndarray x_axis_omp: Numpy array with the radial coordinates we wish to map at the OMP
:param Fiesta fiesta: A Fiesta object with the 2D equilibrium we wish to map
:param np.ndarray divertor_coords: A 2-x-2 numpy array containg the corner
points of the divertor in the 2D projection
:param string location: a string with the location to evaluate, either 'hfs'
or 'lfs'. Default is 'lfs'
:rtype: dict
:return: A dictionary containing:
"R_div" : an n-x-1 array
with the R-coordinates at the divertor tile
corresponding to the same psi_n as at the OMP
"Z_div" : an n-x-1 array
with the Z-coordinates at the divertor tile
corresponding to the same psi_n as at the OMP
"Angles" : an n-x-1 array
with the angles between the field lines and the divertor tile
corresponding to the same psi_n as at the OMP
"Flux_expansion" : an n-x-1 array
with the flux expasion at the divertor tile
corresponding to the same psi_n as at the OMP
"""
# Define the 1D polynomial that represents the divertor
divertor_x = divertor_coords[0]
divertor_y = divertor_coords[1]
divertor_func = np.polyfit(divertor_x, divertor_y, 1)
divertor_polyfit = np.poly1d(divertor_func)
# Define a 1D polynomial for a surface just above the divertor (to be used
# for evaluating the angle of incidence)
divertor_func_above = [divertor_func[0], divertor_func[1] + 0.001]
divertor_polyfit_above = np.poly1d(divertor_func_above)
# Define a 1D polynomial for a surface just below the divertor (to be used
# for evaluating the angle of incidence)
divertor_func_below = [divertor_func[0], divertor_func[1] - 0.001]
divertor_polyfit_below = np.poly1d(divertor_func_below)
# Interpolate psi_n
psi_n_interp = Interpolate2DCubic(fiesta.r_vec, fiesta.z_vec, fiesta.psi_n.T)
# Interpolate b_pol (to be used when evaluating the flux expansion)
b_pol = fiesta.b_theta.T #np.sqrt(fiesta.b_r**2 + fiesta.b_theta**2 + fiesta.b_z**2).T
b_pol_interp = Interpolate2DCubic(fiesta.r_vec, fiesta.z_vec, b_pol)
r_div = []
r_div_above = []
r_div_below = []
flux_expansion = []
for point in x_axis_omp:
# Evaluate which flux surface the point at the OMP is on
psi_n = psi_n_interp(point, 0)
# Evaluate the corresponding point at the divertor
r_mid = _shooting_algorithm(divertor_x, psi_n_interp,
divertor_polyfit, psi_n, location=location)
# Evaluate the corresponding point just above the divertor (for
# calculating the angle of incidence)
r_above = _shooting_algorithm(divertor_x, psi_n_interp,
divertor_polyfit_above, psi_n, location=location)
# Evaluate the corresponding point just below the divertor (for
# calculating the angle of incidence)
r_below = _shooting_algorithm(divertor_x, psi_n_interp,
divertor_polyfit_below, psi_n, location=location)
# Calculate the flux expansion
f_x = _flux_expansion(b_pol_interp, [point, 0],
[r_mid, divertor_polyfit(r_mid)])
r_div.append(r_mid)
r_div_above.append(r_above)
r_div_below.append(r_below)
flux_expansion.append(f_x)
r_div = np.array(r_div)
r_div_above = np.array(r_div_above)
r_div_below = np.array(r_div_below)
flux_expansion = np.array(flux_expansion)
# Determine the vectors at each of the corresponding points at the divertor
v_1_x = r_div_above - r_div_below
v_1_y = divertor_polyfit_above(r_div_above) - divertor_polyfit_below(r_div_below)
v_1_vecs = np.array([v_1_x, v_1_y]).T
# Determine the divertor vector
v_2 = np.array([divertor_x[1] - divertor_x[0], divertor_y[1] - divertor_y[0]])
# Calculate the angle between the vectors and thus the angle of incidence
# defined as the positive angle with respect to the surface normal of the
# divertor
angles = _calculate_angles(v_1_vecs, v_2)
divertor_map = {}
for i in range(len(r_div)):
temp_dict = {}
temp_dict["R_pos"] = r_div[i]
temp_dict["Z_pos"] = divertor_polyfit(r_div[i])
temp_dict["alpha"] = angles[i]
temp_dict["f_x"] = flux_expansion[i]
divertor_map[x_axis_omp[i]] = temp_dict
return divertor_map
import matplotlib.pyplot as plt
from vita.utility import get_resource
from vita.modules.projection.projection2D.field_line.field_line import FieldLine
from vita.modules.projection.projection2D.field_line.field_line_projection import project_field_lines
from vita.modules.equilibrium.fiesta.fiesta_interface import Fiesta
from vita.modules.sol_heat_flux.eich.eich import Eich
from cherab.core.math import Interpolate2DCubic
from vita.modules.projection.projection2D.psi_map_projection import map_psi_omp_to_divertor
from vita.modules.projection.projection2D.project_heat_flux import project_heat_flux
if __name__ == "__main__":
FILEPATH = get_resource("ST40", "equilibrium", "eq002")
FIESTA = Fiesta(FILEPATH)
B_POL = np.sqrt(FIESTA.b_r**2 + FIESTA.b_theta**2 + FIESTA.b_z**2).T
B_POL_INTERP = Interpolate2DCubic(FIESTA.r_vec, FIESTA.z_vec, B_POL)
FIELD_LINE = FieldLine(FILEPATH)
MID_PLANE_LOC = FIESTA.get_midplane_lcfs()[1]
FOOTPRINT = Eich(0.0025, 0.0005,
r0_lfs=MID_PLANE_LOC) # lambda_q=2.5, S=0.5
X_OMP = np.linspace(0, 10, 100) * 1e-3
FOOTPRINT.set_coordinates(X_OMP)
FOOTPRINT.s_disconnected_dn_max = 2.1
FOOTPRINT.fx_in_out = 5.
FOOTPRINT.calculate_heat_flux_density("lfs")
Q_PARALLEL = FOOTPRINT._q
X_AFTER_LCFS = FOOTPRINT.get_global_coordinates()
def project_field_lines(x_axis_omp, surface_coords, fiesta):
'''
Function mapping the field-lines from the specified coordinates at the
OMP to the specified coordinates at a given surface. Currently the surface
is assumed to be represented by a 1D polynomial function, y = ax + b.
Parameters
----------
x_axis_omp : n-x-1 np.array
Numpy array with the radial coordinates we wish to map at the OMP
fiesta : Fiesta
A Fiesta object with the 2D equilibrium we wish to map
divertor_coords : 2-x-2 np.array
A 2-x-2 numpy array containg the corner points of the divertor in the
2D projection
Returns
-------
divertor_map : dictionary
A dictionary containing:
"R_div" : an n-x-1 array
with the R-coordinates at the divertor tile
corresponding to the same psi_n as at the OMP
"Z_div" : an n-x-1 array
with the Z-coordinates at the divertor tile
corresponding to the same psi_n as at the OMP
"Angles" : an n-x-1 array
with the angles between the field lines and the divertor tile
corresponding to the same psi_n as at the OMP
"Flux_expansion" : an n-x-1 array
with the flux expasion at the divertor tile
corresponding to the same psi_n as at the OMP
'''
# Interpolate b_pol (to be used when evaluating the flux expansion)
b_pol = fiesta.b_theta.T
b_pol_interp = Interpolate2DCubic(fiesta.r_vec, fiesta.z_vec, b_pol)
field_lines = {}
if os.path.exists("Random_filename"):
field_lines = load_pickle("Random_filename")
else:
field_line = FieldLine(fiesta)
for i in x_axis_omp:
p_0 = [i, 0, 0]
field_line_dict = field_line.follow_field_in_plane(p_0=p_0,
max_length=15.0)
field_lines[i] = field_line_dict
_v_2 = np.array([
surface_coords[0, 1] - surface_coords[0, 0],
surface_coords[1, 1] - surface_coords[1, 0]
])
divertor_map = {}
for i in field_lines:
func1 = np.array((field_lines[i]['R'], field_lines[i]['Z']))
(i_func1, _), (x_at_surface,
y_at_surface) = intersection(func1, (surface_coords))
if np.isnan(x_at_surface):
break
_x_component = field_lines[i]['R'][int(i_func1)] - field_lines[i]['R'][
int(i_func1) + 1]
_y_component = field_lines[i]['Z'][int(i_func1)] - field_lines[i]['Z'][
int(i_func1) + 1]
_v_1 = [_x_component, _y_component]
temp_dict = {}
temp_dict["R_pos"] = x_at_surface[0]
temp_dict["Z_pos"] = y_at_surface[0]
temp_dict["alpha"] = calculate_angle(_v_1, _v_2)
temp_dict["f_x"] = _flux_expansion(b_pol_interp, (i, 0),
(x_at_surface[0], y_at_surface[0]))
divertor_map[i] = temp_dict
return divertor_map
def import_equ_psi(equ_file):
"""
Imports a 2D Psi grid from an equ mesh equilibrium.
:param str equ_file: The .EQU mesh generator equilibrium file.
:return: A 2D Psi(r,z) function
:rtype: Interpolate2DCubic
"""
fh = open(equ_file, 'r')
file_lines = fh.readlines()
fh.close()
# Load r array
line_i = 0
while True:
line = file_lines[line_i]
match = re.match('^\s*r\(1:jm\);', line)
if not match:
line_i += 1
continue
line_i += 1
r_values = []
line = file_lines[line_i]
while line and not line.isspace():
print(line)
values = line.split()
print(values)
for v in values:
r_values.append(float(v))
line_i += 1
line = file_lines[line_i]
jm = len(r_values)
break
r_values = np.array(r_values)
# Load z array
while True:
line = file_lines[line_i]
match = re.match('^\s*z\(1:km\);', line)
if not match:
line_i += 1
continue
line_i += 1
z_values = []
line = file_lines[line_i]
while not re.match('^\s*$', line):
values = line.split()
for v in values:
z_values.append(float(v))
line_i += 1
line = file_lines[line_i]
km = len(z_values)
break
z_values = np.array(z_values)
# Load (r, z) array
while True:
line = file_lines[line_i]
match = re.match('^\s*\(\(psi\(j,k\)-psib,j=1,jm\),k=1,km\)', line)
if not match:
line_i += 1
continue
line_i += 1
psi_values = []
line = file_lines[line_i]
while not re.match('^\s*$', line):
values = line.split()
for v in values:
psi_values.append(float(v))
line_i += 1
try:
line = file_lines[line_i]
except IndexError:
break
if len(psi_values) != km * jm:
raise ValueError(
"Number of values in r, z array not equal to (km, jm).")
break
psi_raw = np.array(psi_values)
psi = psi_raw.reshape((km, jm))
psi = np.swapaxes(psi, 0, 1)
return Interpolate2DCubic(r_values, z_values, psi)