Create emission profile

import numpy as np
import matplotlib.pyplot as plt

# Core and external imports
from raysect.core import Vector3D, Point3D
from raysect.core.ray import Ray as CoreRay
from raysect.primitive import Cylinder
from raysect.primitive.mesh import Mesh, import_stl
from raysect.optical import Ray, Spectrum
from raysect.optical.material.absorber import AbsorbingSurface
from raysect.optical.observer import FibreOptic
from raysect.optical.observer.pipeline.spectral import SpectralPowerPipeline0D
from raysect.optical import World, translate, rotate, rotate_basis

# Cherab and raysect imports
from cherab.core import Species, Maxwellian, Plasma
from cherab.core.atomic import deuterium, carbon, Line
from cherab.core.model import ExcitationLine, RecombinationLine, StarkBroadenedLine
from cherab.openadas import OpenADAS
from cherab.iter.observer import FoV
from cherab.imas import EdgeProfilesIDS


shot = 122264
run = 1
time = 0
view = "DIM_UP1"

print("loading edge solution")
edge_ids = EdgeProfilesIDS(shot, run=run, user="public", machine="iter")
r_range, z_range = edge_ids.get_geometry_range(time)
rmin, rmax = r_range
zmin, zmax = z_range

fov = FoV(view)
pupils = fov.view["PUPILS"]
fov_center = fov.cs[2]

# === Create Plasma object from EdgeProfileIDS ================================

world = World()

# create plasma
plasma = edge_ids.create_plasma(time, parent=world)
plasma.atomic_data = OpenADAS(permit_extrapolation=True)
plasma.geometry = Cylinder(rmax, zmax - zmin, transform=translate(0, 0, -(zmax - zmin) / 2))
plasma.geometry_transform = translate(0, 0, -(zmax - zmin) / 2)

# Setup elements.deuterium lines
d_alpha_excit = ExcitationLine(Line(deuterium, 0, (3, 2)), lineshape=StarkBroadenedLine)
d_beta_excit = ExcitationLine(Line(deuterium, 0, (4, 2)), lineshape=StarkBroadenedLine)
d_gamma_excit = ExcitationLine(Line(deuterium, 0, (5, 2)), lineshape=StarkBroadenedLine)
d_delta_excit = ExcitationLine(Line(deuterium, 0, (6, 2)), lineshape=StarkBroadenedLine)
d_epsilon_excit = ExcitationLine(Line(deuterium, 0, (7, 2)), lineshape=StarkBroadenedLine)
d_alpha_recom = RecombinationLine(Line(deuterium, 0, (3, 2)), lineshape=StarkBroadenedLine)
d_beta_recom = RecombinationLine(Line(deuterium, 0, (4, 2)), lineshape=StarkBroadenedLine)
d_gamma_recom = RecombinationLine(Line(deuterium, 0, (5, 2)), lineshape=StarkBroadenedLine)
d_delta_recom = RecombinationLine(Line(deuterium, 0, (6, 2)), lineshape=StarkBroadenedLine)
d_epsilon_recom = RecombinationLine(Line(deuterium, 0, (7, 2)), lineshape=StarkBroadenedLine)

plasma.models = [
    d_alpha_excit,
    d_beta_excit,
    d_gamma_excit,
    d_delta_excit,
    d_epsilon_excit,
    d_alpha_recom,
    d_beta_recom,
    d_gamma_recom,
    d_delta_recom,
    d_epsilon_recom,
]

log = True
emission_samples = np.zeros((500, 500))

xrange = np.linspace(rmin, rmax, 500)
yrange = np.linspace(zmin, zmax, 500)

for i, x in enumerate(xrange):
    for j, y in enumerate(yrange):
        for emitter in plasma.models:
            emission_samples[j, i] += emitter.emission(
                Point3D(x, 0.0, y), Vector3D(0, 0, 0), Spectrum(400, 657, 800)
            ).total()

if log:
    emission_samples = np.log(emission_samples)
plt.figure()
plt.imshow(emission_samples, extent=[rmin, rmax, zmin, zmax], origin="lower")
plt.colorbar()
plt.xlim(rmin, rmax)
plt.ylim(zmin, zmax)
plt.show()
print(emission_samples)

# Normalise the emission arrays
# dalpha /= dalpha.sum()
# dgamma /= dgamma.sum()
# dbeta /= dbeta.sum()
# ddelta /= ddelta.sum()
# depsilon /= depsilon.sum()

# Plot the trajectory parameters
# sim.plot_pec_emission_lines([d_alpha_excit, d_alpha_recom], title='D_alpha')
# plt.plot(ray_r_points, ray_z_points, 'k')
# plt.plot(ray_r_points[0], ray_z_points[0], 'b.')
# plt.plot(ray_r_points[-1], ray_z_points[-1], 'r.')

# plt.figure()
# plt.plot(distance, te)
# plt.xlabel("Ray distance (m)")
# plt.ylabel("Electron temperature (eV)")
# plt.title("Electron temperature (eV) along ray path")

# plt.figure()
# plt.plot(distance, ne)
# plt.xlabel("Ray distance (m)")
# plt.ylabel("Electron density (m^-3)")
# plt.title("Electron density (m^-3) along ray path")

# plt.figure()
# plt.plot(distance, dalpha, label='Dalpha')
# plt.plot(distance, dgamma, label='Dgamma')
# plt.plot(distance, dbeta, label='Dbeta')
# plt.plot(distance, ddelta, label='Ddelta')
# plt.plot(distance, depsilon, label='Depsilon')
# plt.xlabel("Ray distance (m)")
# plt.ylabel("Normalised emission")
# plt.title("Normalised emission along ray path")
# plt.legend()


# plt.ion()
# r = Ray(origin=Point3D(*pupils), direction=Vector3D(*fov_center),
#         min_wavelength=200, max_wavelength=1000, bins=1000)
# s = r.trace(world)
# plt.plot(s.wavelengths, s.samples)
# plt.yscale('log')
# plt.ylim(10**-4, None)
# plt.xlabel('Wavelength (nm)')
# plt.ylabel('Radiance (W/m^2/str/nm)')
# plt.title('Observed Spectrum')
# plt.ioff()
# plt.show()

# Fibre position and observation direction
# start_point = Point3D(*pupils)
# forward_vector = Vector3D(*fov_center).normalise()
# up_vector = Vector3D(0, 0, 1.0)

# spectra = SpectralPowerPipeline0D()
# fibre = FibreOptic([spectra], acceptance_angle=1, radius=0.001, spectral_bins=8000, spectral_rays=1,
#                    pixel_samples=5, transform=translate(*start_point)*rotate_basis(forward_vector, up_vector), parent=world)

# fibre.min_wavelength = 350.0
# fibre.max_wavelength = 700.0

# fibre.observe()