Simpl calculation of RTM

import numpy as np
from raysect.optical import World, translate, rotate
from raysect.optical.material import AbsorbingSurface
from raysect.optical.observer import FullFrameSampler2D
from raysect.primitive import Cylinder, Subtract
from cherab.iter.machine import import_iter_mesh
from cherab.iter.observer import FovCamera, FoV
from cherab.iter.emitter import RtmRegularRPhiZEmitter, RtmRegularRPhiZIntegrator
from cherab.iter.observer import RtmPipeline2D


# === configuration ===========================================================
view = "DIM_EP1"
dtheta = 0.1  # [deg]
dphi = 0.6  # [deg]
cutoff = 0.1
cutoff_frac = 0.01
pixel_samples = 2000
wlmin, wlmax = 655.6, 656.6
step = 0.001  # [m]
izstep = 95
eps_rz = 1.0e-4
nr = 460
nz = 950
rmin, rmax = 3.95, 8.55  # [m]
zmin, zmax = -4.65, 4.8  # [m]
r, dr = np.linspace(rmin, rmax, nr + 1, retstep=True)
z, dz = np.linspace(zmin, zmax, nz + 1, retstep=True)

world = World()

# === ITER PFC ================================================================
import_iter_mesh(world)
central_column = Cylinder(rmin, zmax - zmin, material=AbsorbingSurface(), parent=world, transform=translate(0, 0, zmin))

# === Camera ==================================================================
fov = FoV(view)
pupils = fov.view["PUPILS"]
_, theta = fov.view["FOV"]
phi = dphi
rotate_angles = fov.euler()

npix_x = 1  # just one pixel (0.6deg-wide)
npix_y = int(round(theta / dtheta))
pixels = (npix_x, npix_y)

pipeline = RtmPipeline2D()
sampler = FullFrameSampler2D()
camera = FovCamera(pixels, fov=(phi, theta), parent=world, pipelines=[pipeline], frame_sampler=sampler)
camera.transform = translate(*pupils) * rotate(*rotate_angles)
camera.min_wavelength = 656.0
camera.max_wavelength = 656.2
camera.spectral_rays = 1
camera.pixel_samples = pixel_samples
camera.ray_extinction_prob = 0
camera.ray_importance_sampling = False

# === RTM Calculation =========================================================
rtm = np.zeros((npix_x, npix_y, nr, nz), dtype=np.float32)

for iz in range(0, nz, izstep):
    iz_end = min(iz + izstep, nz)

    print("#-------------- ", iz, "-", iz_end, "/", nz, " --------------#")

    emitt_material = RtmRegularRPhiZEmitter(nr, iz_end - iz, dr, dz, rmin, integrator=RtmRegularRPhiZIntegrator(step))
    emitter = Subtract(
        Cylinder(rmax - eps_rz, z[iz_end] - z[iz] - eps_rz),
        Cylinder(rmin + eps_rz, z[iz_end] - z[iz] - eps_rz),
        parent=world,
    )
    emitter.material = emitt_material
    emitter.transform = translate(0, 0, z[iz])
    camera.spectral_bins = nr * (iz_end - iz)

    camera.observe()

    partial_rtm = pipeline.rtm
    rtm[:, :, :, iz:iz_end] = partial_rtm.astype(np.float32).reshape((npix_x, npix_y, nr, iz_end - iz))

np.save("RTM_%s.npy" % view, rtm)