diff --git a/debye_bec/bec_widgets/widgets/digital_twin/calc_positions.py b/debye_bec/bec_widgets/widgets/digital_twin/calc_positions.py
index 55a1cca..f463286 100644
--- a/debye_bec/bec_widgets/widgets/digital_twin/calc_positions.py
+++ b/debye_bec/bec_widgets/widgets/digital_twin/calc_positions.py
@@ -1,242 +1,259 @@
-import os
+"""
+Calculates the positions of axes based on a beamline config
+"""
+
 import numpy as np
 from bec_lib import bec_logger
 
-os.environ["USE_XRT"] = "False"
 import debye_bec.bec_widgets.widgets.x01da_parameters as bl
+from debye_bec.bec_widgets.widgets.digital_twin.types import ConfigDict
 
 logger = bec_logger.logger
 
-def calc_positions(cfg):
+
+def calc_positions(cfg: ConfigDict) -> dict[str, dict[str, float]]:
+    """
+    Calculates the positions of axes based on a beamline config.
+
+    Args:
+        cfg(ConfigDict): Dictionary with beamline config
+
+    Returns:
+        dict[str, dict[str, float]]: Dictionary mapping device names to dictionaries
+            containing a "value" key with the corresponding float value (position).
+    """
 
     pos = {}
 
     ## FE slits
-    trxr = -np.arctan(cfg['h_acc'])*bl.feSlits.center1[1]
-    trxw = (np.arctan(cfg['h_acc'])*bl.feSlits.center1[1])/bl.feSlits.center1[1]*bl.feSlits.center2[1]
-    tryb = -np.arctan(cfg['v_acc'])*bl.feSlits.center1[1]
-    tryt = (np.arctan(cfg['v_acc'])*bl.feSlits.center1[1])/bl.feSlits.center1[1]*bl.feSlits.center2[1]
+    trxr = -np.arctan(cfg["h_acc"]) * bl.feSlits.center1[1]
+    trxw = (
+        (np.arctan(cfg["h_acc"]) * bl.feSlits.center1[1])
+        / bl.feSlits.center1[1]
+        * bl.feSlits.center2[1]
+    )
+    tryb = -np.arctan(cfg["v_acc"]) * bl.feSlits.center1[1]
+    tryt = (
+        (np.arctan(cfg["v_acc"]) * bl.feSlits.center1[1])
+        / bl.feSlits.center1[1]
+        * bl.feSlits.center2[1]
+    )
 
-    # trxw_proj = trxw/bl.feSlits.center2[1]*bl.feSlits.center1[1]
-    # tryt_proj = tryt/bl.feSlits.center2[1]*bl.feSlits.center1[1]
-
-    # xcen = (trxr + trxw) / 2
-    # ycen = (tryb + tryt) / 2
     xgap = trxw - trxr
     ygap = tryt - tryb
 
-    pos['sldi_gapx'] = {'value': xgap}
-    pos['sldi_gapy'] = {'value': ygap}
+    pos["sldi_gapx"] = {"value": xgap}
+    pos["sldi_gapy"] = {"value": ygap}
 
     ## Collimating Mirror
-    obj_dist = bl.cm.center[1] # object distance
-    beam_vs = 2 * obj_dist * np.tan(cfg['v_acc']) # vertical size of beam after CM
+    obj_dist = bl.cm.center[1]  # object distance
+    beam_vs = 2 * obj_dist * np.tan(cfg["v_acc"])  # vertical size of beam after CM
 
     # TRX
-    try:
-        index = bl.cm.surface.index(cfg['cm_stripe'])
-    except:
+    if cfg["cm_stripe"] in bl.cm.surface:
+        index = bl.cm.surface.index(cfg["cm_stripe"])
+    else:
         raise ValueError(f"Requested stripe {cfg['cm_stripe']} not found in parameters!")
     cm_trx = -(bl.cm.limOptX[0][index] + bl.cm.limOptX[1][index]) / 2
-    pos['cm_trx'] = {'value': cm_trx}
+    pos["cm_trx"] = {"value": cm_trx}
 
     # TRY
-    height = obj_dist * np.tan(cfg['v_acc'])**2 * 1 / np.tan(cfg['cm_pitch'])
-    pos['cm_try'] = {'value': height}
+    height = obj_dist * np.tan(cfg["v_acc"]) ** 2 * 1 / np.tan(cfg["cm_pitch"])
+    pos["cm_try"] = {"value": height}
 
     # Pitch
-    pos['cm_rotx'] = {'value': -cfg["cm_pitch"]*1e3} # invert and convert to mrad (same as EGU of rotx axis)
+    pos["cm_rotx"] = {
+        "value": -cfg["cm_pitch"] * 1e3
+    }  # invert and convert to mrad (same as EGU of rotx axis)
 
     # Bending Radius
-    radius = 2. * obj_dist / np.sin(cfg['cm_pitch']) # Elements of modern X-ray Physics, page 108 ff.
-    pos['cm_bnd_radius'] = {'value': radius * 1e-6} # Convert to km
+    radius = (
+        2.0 * obj_dist / np.sin(cfg["cm_pitch"])
+    )  # Elements of modern X-ray Physics, page 108 ff.
+    pos["cm_bnd_radius"] = {"value": radius * 1e-6}  # Convert to km
 
     ## Monochromator
-    # Bragg Angle
-    # if cfg['mo1_mode'] == 'Monochromatic':
-    #     # Add 2x CM pitch to the bragg angle
-    #     bragg =  ((2 * cfg['cm_pitch']) + cfg['mo1_bragg']) / np.pi * 180
-    # elif cfg['mo1_mode'] == 'Pinkbeam':
-    #     # Align xtal surfaces parallel to beam
-    #     bragg =  (2 * cfg['cm_pitch']) / np.pi * 180
-    # else:
-    #     raise Exception('Monochromator mode not supported')
-    if cfg['mo1_mode'] == 'Monochromatic':
+    if cfg["mo1_mode"] == "Monochromatic":
         # Add 2x CM pitch to the bragg angle
-        bragg = cfg['mo1_bragg']
-    elif cfg['mo1_mode'] == 'Pinkbeam':
+        bragg = cfg["mo1_bragg"]
+    elif cfg["mo1_mode"] == "Pinkbeam":
         # Align xtal surfaces parallel to beam
-        bragg =  0
+        bragg = 0
     else:
-        raise Exception('Monochromator mode not supported')
-    pos['mo1_bragg_angle'] = {'value': bragg/np.pi*180} # Bragg angle in deg
+        raise ValueError("Monochromator mode not supported")
+    pos["mo1_bragg_angle"] = {"value": bragg / np.pi * 180}  # Bragg angle in deg
 
     # TRY, Height
-    l = bl.mo1.xtalGap[0]/np.sin(cfg['mo1_bragg'])
-    yhor = l*np.cos(2.*(cfg['mo1_bragg']+cfg['cm_pitch']))
-    yver = yhor*np.tan(2.*cfg['cm_pitch'])
+    l = bl.mo1.xtalGap[0] / np.sin(cfg["mo1_bragg"])
+    yhor = l * np.cos(2.0 * (cfg["mo1_bragg"] + cfg["cm_pitch"]))
+    yver = yhor * np.tan(2.0 * cfg["cm_pitch"])
 
-    if cfg['mo1_mode'] == 'Monochromatic':
-        beamOffsetCCM = l*np.sin(2.*(cfg['mo1_bragg']+cfg['cm_pitch']))-yver # Resultat ist korrekt!
-    elif cfg['mo1_mode'] == 'Pinkbeam':
-        beamOffsetCCM = 0
+    if cfg["mo1_mode"] == "Monochromatic":
+        beam_offset_mo1 = (
+            l * np.sin(2.0 * (cfg["mo1_bragg"] + cfg["cm_pitch"])) - yver
+        )  # Resultat ist korrekt!
+    elif cfg["mo1_mode"] == "Pinkbeam":
+        beam_offset_mo1 = 0
     else:
-        raise Exception('Monochromator mode not supported')
+        raise ValueError("Monochromator mode not supported")
 
     def csc(a):
-        return 1/np.sin(a)
+        return 1 / np.sin(a)
 
     def cot(a):
-        return 1/np.tan(a)
+        return 1 / np.tan(a)
 
     # calculate height of center of first crystal surface
-    f = bl.mo1.rotOffset # rotation offset, mm
-    # logger.info(f'f = {f}')
-    d = bl.mo1.heightOffset # xtal height offset, mm
-    # logger.info(f'd = {d}')
-    c = d*csc(cfg['mo1_bragg'])-f*cot(cfg['mo1_bragg'])
-    # logger.info(f'c = {c}')
+    f = bl.mo1.rotOffset  # rotation offset, mm
+    d = bl.mo1.heightOffset  # xtal height offset, mm
+    c = d * csc(cfg["mo1_bragg"]) - f * cot(cfg["mo1_bragg"])
 
     # Calculate height of center of rotation
-    b = np.sqrt(d**2*csc(cfg['mo1_bragg'])**2-2*d*f*cot(cfg['mo1_bragg'])*csc(cfg['mo1_bragg'])+f**2*cot(cfg['mo1_bragg'])**2+f**2)
-    # logger.info(f'b = {b}')
-    h = np.cos(np.pi/2-np.arctan(f/c)-cfg['mo1_bragg']-2*cfg['cm_pitch'])*b
-    # logger.info(f'h = {h}')
-    h2 = ((bl.mo1.center[1] - bl.cm.center[1])-np.sqrt(b**2-h**2))*np.tan(2*cfg['cm_pitch'])
-    # logger.info(f'mo1 = {bl.mo1.center[1]}')
-    # logger.info(f'cm = {bl.cm.center[1]}')
-    # logger.info(f'pitch = {cfg["cm_pitch"]}')
-    # logger.info(f'h2 = {h2}')
-    #TODO Mono height not exactly the same as in raytracing
-    heightCCM1real = h + h2 # per design, the height should not change if the pitch of the CM is not changed!
-    # heightCCM1real = heightCCM1real - 30 # Zero position of stage is at 1430 mm from ground.
-    if cfg['mo1_mode'] == 'Monochromatic':
+    b = np.sqrt(
+        d**2 * csc(cfg["mo1_bragg"]) ** 2
+        - 2 * d * f * cot(cfg["mo1_bragg"]) * csc(cfg["mo1_bragg"])
+        + f**2 * cot(cfg["mo1_bragg"]) ** 2
+        + f**2
+    )
+    h = np.cos(np.pi / 2 - np.arctan(f / c) - cfg["mo1_bragg"] - 2 * cfg["cm_pitch"]) * b
+    h2 = ((bl.mo1.center[1] - bl.cm.center[1]) - np.sqrt(b**2 - h**2)) * np.tan(2 * cfg["cm_pitch"])
+    height_mo1_real = (
+        h + h2
+    )  # per design, the height should not change if the pitch of the CM is not changed!
+    if cfg["mo1_mode"] == "Monochromatic":
         pass
-    elif cfg['mo1_mode'] == 'Pinkbeam':
-        heightCCM1real = heightCCM1real - 13 # Move down to let beam pass between both crystal without touching copper cooler
+    elif cfg["mo1_mode"] == "Pinkbeam":
+        height_mo1_real = (
+            height_mo1_real - 13
+        )  # Move down to let beam pass between both crystal without touching copper cooler
     else:
-        raise Exception('Monochromator mode not supported')
-    pos['mo1_try'] = {'value': heightCCM1real}
+        raise ValueError("Monochromator mode not supported")
+    pos["mo1_try"] = {"value": height_mo1_real}
 
     # TRX, Crystal selection
-    if cfg['mo1_mode'] == 'Monochromatic':
-        try:
-            xtal = cfg['mo1_xtal'].translate(str.maketrans('', '', '()')) # Remove brackets from xtal name to conform with parameters
+    if cfg["mo1_mode"] == "Monochromatic":
+        xtal = cfg["mo1_xtal"].translate(
+            str.maketrans("", "", "()")
+        )  # Remove brackets from xtal name to conform with parameters
+        if xtal in bl.mo1.xtal:
             index = bl.mo1.xtal.index(xtal)
-        except:
+        else:
             raise ValueError(f"Requested xtal {xtal} not found in parameters!")
-        pos['mo1_trx'] = {'value': bl.mo1.xtalOffsetX[index]}
+        pos["mo1_trx"] = {"value": bl.mo1.xtalOffsetX[index]}
     else:
-        pos['mo1_trx'] = {'value': 0}
+        pos["mo1_trx"] = {"value": 0}
 
-
-    #TODO move to mono, calc for beam Z-movement between crystal surfaces
-    diag = bl.mo1.xtalGap[0] / np.sin(cfg['mo1_bragg']) # Calculations for Mono
-    dz = diag * np.cos(2 * (cfg['cm_pitch'] + cfg['mo1_bragg']))
+    diag = bl.mo1.xtalGap[0] / np.sin(cfg["mo1_bragg"])  # Calculations for Mono
+    dz = diag * np.cos(2 * (cfg["cm_pitch"] + cfg["mo1_bragg"]))
 
     ## Slits 1
     d = bl.opSlits1.center[1] - bl.cm.center[1] - dz
-    sl1_beam_height = d * np.tan(2 * cfg['cm_pitch']) + beamOffsetCCM
-    pos['sl1_centery'] = {'value': sl1_beam_height}
-    pos['sl1_gapy'] = {'value': beam_vs + 1} # Add 0.5 mm space on both sides of the beam
+    sl1_beam_height = d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1
+    pos["sl1_centery"] = {"value": sl1_beam_height}
+    pos["sl1_gapy"] = {"value": beam_vs + 1}  # Add 0.5 mm space on both sides of the beam
 
     ## Beam Monitor 1
     d = bl.opBM1.center[1] - bl.cm.center[1] - dz
-    # logger.info(f'distance: {d}')
-    # logger.info(f'cm pitch: {cfg["cm_pitch"]}')
-    # logger.info(f'mono offset: {beamOffsetCCM}')
-    bm1_beam_height = d * np.tan(2 * cfg['cm_pitch']) + beamOffsetCCM
-    pos['bm1_try'] = {'value': bm1_beam_height}
+    bm1_beam_height = d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1
+    pos["bm1_try"] = {"value": bm1_beam_height}
 
     ## Focusing Mirror
     p = bl.fm.center[1]
-    q = cfg['smpl'] - bl.fm.center[1]
-    f = (p*q)/(p+q) # focal length
+    q = cfg["smpl"] - bl.fm.center[1]
+    f = (p * q) / (p + q)  # focal length
 
     # Bender radius
-    if cfg['fm_qy'] is None:
-        radius = 2 * q / np.sin(cfg['fm_rotx']) # ideal bending radius for focused beam
+    if cfg["fm_qy"] is None:
+        radius = 2 * q / np.sin(cfg["fm_rotx"])  # ideal bending radius for focused beam
     else:
-        radius = 2 * cfg['fm_qy'] / np.sin(cfg['fm_rotx']) # ideal bending radius for unfocused beam
-    pos['fm_bnd_radius'] = {'value': radius * 1e-6} # Convert to km
+        radius = (
+            2 * cfg["fm_qy"] / np.sin(cfg["fm_rotx"])
+        )  # ideal bending radius for unfocused beam
+    pos["fm_bnd_radius"] = {"value": radius * 1e-6}  # Convert to km
 
     # Pitch
     d = bl.fm.center[1] - bl.cm.center[1] - dz
-    fm_rotx = 2 * cfg['cm_pitch'] - cfg['fm_rotx'] # calculate pitch in absolute values (according to horizontal plane)
-    pos['fm_rotx'] = {'value': -fm_rotx * 1e3} # invert and convert to mrad (same as EGU of rotx axis)
+    fm_rotx = (
+        2 * cfg["cm_pitch"] - cfg["fm_rotx"]
+    )  # calculate pitch in absolute values (according to horizontal plane)
+    pos["fm_rotx"] = {
+        "value": -fm_rotx * 1e3
+    }  # invert and convert to mrad (same as EGU of rotx axis)
 
-    if cfg['fm_stripe'] in ('Rh (toroid)', 'Pt (toroid)'):
+    if cfg["fm_stripe"] in ("Rh (toroid)", "Pt (toroid)"):
 
         # TRY
-        if cfg['fm_stripe'] in 'Rh (toroid)':
+        if cfg["fm_stripe"] in "Rh (toroid)":
             r = bl.fm.r[0]
             h_cyl = bl.fm.hToroid[0]
-        else: # PT toroid
+        else:  # PT toroid
             r = bl.fm.r[1]
             h_cyl = bl.fm.hToroid[1]
-        widthBeam = 2 * bl.fm.center[1] * np.tan(cfg['h_acc'] * 1e-3)
-        alpha = np.arccos(1 - widthBeam**2 / (2 * r**2))
+        width_beam = 2 * bl.fm.center[1] * np.tan(cfg["h_acc"] * 1e-3)
+        alpha = np.arccos(1 - width_beam**2 / (2 * r**2))
         h = r - (r * np.cos(alpha / 2))
-        fm_beam_height = (d * np.tan(2 * cfg['cm_pitch']) + beamOffsetCCM) * cfg['fm_gain_height']
-        fm_height = (d * np.tan(2 * cfg['cm_pitch']) + beamOffsetCCM - h_cyl + h / 2) * cfg['fm_gain_height']
-        pos['fm_try'] = {'value': fm_height}
+        fm_beam_height = (d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1) * cfg["fm_gain_height"]
+        fm_height = (d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1 - h_cyl + h / 2) * cfg[
+            "fm_gain_height"
+        ]
+        pos["fm_try"] = {"value": fm_height}
 
         # TRX
-        if cfg['fm_stripe'] in 'Rh (toroid)':
-            x_cyl = - bl.fm.xToroid[0]
+        if cfg["fm_stripe"] in "Rh (toroid)":
+            x_cyl = -bl.fm.xToroid[0]
         else:
-            x_cyl = - bl.fm.xToroid[1]
-        pos['fm_trx'] = {'value': x_cyl}
+            x_cyl = -bl.fm.xToroid[1]
+        pos["fm_trx"] = {"value": x_cyl}
 
-    elif cfg['fm_stripe'] in ('Rh (flat)', 'Pt (flat)'):
+    elif cfg["fm_stripe"] in ("Rh (flat)", "Pt (flat)"):
 
         # TRY
-        fm_height = (d * np.tan(2 * cfg['cm_pitch']) + beamOffsetCCM) * cfg['fm_gain_height']
+        fm_height = (d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1) * cfg["fm_gain_height"]
         fm_beam_height = fm_height
-        pos['fm_try'] = {'value': fm_height}
+        pos["fm_try"] = {"value": fm_height}
 
         # TRX
-        if cfg['fm_stripe'] in 'Rh (flat)':
-            x_flat = - bl.fm.xFlat[0]
+        if cfg["fm_stripe"] in "Rh (flat)":
+            x_flat = -bl.fm.xFlat[0]
         else:
-            x_flat = - bl.fm.xFlat[1]
-        pos['fm_trx'] = {'value': x_flat}
+            x_flat = -bl.fm.xFlat[1]
+        pos["fm_trx"] = {"value": x_flat}
 
     else:
-        raise Exception('FM Stripe selection not valid')
+        raise ValueError("FM Stripe selection not valid")
 
-    pos['fm_roty'] = {'value': 0}
-    pos['fm_rotz'] = {'value': 0}
+    pos["fm_roty"] = {"value": 0}
+    pos["fm_rotz"] = {"value": 0}
 
     ## Slits 2
     d = bl.opSlits2.center[1] - bl.fm.center[1]
-    sl2_beam_height = fm_beam_height - d * np.tan(-(2 * cfg['cm_pitch'] - 2 * cfg['fm_rotx']))
-    pos['sl2_centery'] = {'value': sl2_beam_height}
-    pos['sl2_gapy'] = {'value': beam_vs + 1} # Add 0.5 mm space on both sides of the beam
+    sl2_beam_height = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"]))
+    pos["sl2_centery"] = {"value": sl2_beam_height}
+    pos["sl2_gapy"] = {"value": beam_vs + 1}  # Add 0.5 mm space on both sides of the beam
 
     ## Beam Monitor 2
     d = bl.opBM2.center[1] - bl.fm.center[1]
-    bm2_beam_height = fm_beam_height - d * np.tan(-(2 * cfg['cm_pitch'] - 2 * cfg['fm_rotx']))
-    pos['bm2_try'] = {'value': bm2_beam_height}
+    bm2_beam_height = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"]))
+    pos["bm2_try"] = {"value": bm2_beam_height}
 
     ## Optical Table
 
     # TRY
     d = bl.ehWindow.center[1] - bl.fm.center[1]
-    ot_height = fm_beam_height - d * np.tan(-(2 * cfg['cm_pitch'] - 2 * cfg['fm_rotx']))
-    # logger.info(fm_height)
-    # logger.info(d * np.tan((2 * cfg['cm_pitch'] - 2 * cfg['fm_rotx'])))
-    pos['ot_try'] = {'value': ot_height}
+    ot_height = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"]))
+    pos["ot_try"] = {"value": ot_height}
 
     # Pitch
-    ot_pitch = - (2 * cfg['cm_pitch'] - 2 * cfg['fm_rotx'])
-    pos['ot_rotx'] = {'value': ot_pitch * 1e3}
+    ot_pitch = -(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])
+    pos["ot_rotx"] = {"value": ot_pitch * 1e3}
 
     # TRZ ES1
-    ot_es1_trz = cfg['smpl']
-    pos['ot_es1_trz'] = {'value': ot_es1_trz}
+    ot_es1_trz = cfg["smpl"]
+    pos["ot_es1_trz"] = {"value": ot_es1_trz}
 
     # ES0 exit window
-    pos['es0wi_try'] = {'value': 5} # At 5mm, the middle of the window is 500 mm from the table (neutral position)
+    pos["es0wi_try"] = {
+        "value": 5
+    }  # At 5mm, the middle of the window is 500 mm from the table (neutral position)
 
     return pos
diff --git a/debye_bec/bec_widgets/widgets/digital_twin/calc_sideview.py b/debye_bec/bec_widgets/widgets/digital_twin/calc_sideview.py
index 55d09c5..092d5d4 100644
--- a/debye_bec/bec_widgets/widgets/digital_twin/calc_sideview.py
+++ b/debye_bec/bec_widgets/widgets/digital_twin/calc_sideview.py
@@ -1,15 +1,23 @@
+"""
+Calculates the sideview coordinates based on a beamline config.
+"""
+
 import numpy as np
 
 import debye_bec.bec_widgets.widgets.x01da_parameters as bl
-from debye_bec.bec_widgets.widgets.digital_twin.types import DataDict
+from debye_bec.bec_widgets.widgets.digital_twin.types import ConfigDict, DataDict
 
 
-def calc_sideview(cfg):
+def calc_sideview(cfg: ConfigDict) -> DataDict:
+    """
+    Calculates the sideview coordinates based on a beamline config.
 
-    # Calculate height of beam after CM
-    # height = 2 * bl.cm.center[1] * np.tan(cfg["v_acc"])
+    Args:
+        cfg(ConfigDict): Dictionary with beamline config
 
-    # beam height (Y=height, Z=along beam)
+    Returns:
+        DataDict: Sideview data
+    """
 
     beam: DataDict = {"x": [], "y": []}
 
@@ -59,11 +67,4 @@ def calc_sideview(cfg):
             + np.tan(2 * (cfg["cm_pitch"] - cfg["fm_rotx"])) * (cfg["smpl"] - bl.fm.center[1])
         )
 
-    dy_fm_ex = beam["y"][-1] - beam["y"][-2]
-    dz_fm_ex = beam["x"][-1] - beam["x"][-2]
-    dz_fm_win = bl.ehWindow.center[1] - beam["x"][-2]
-    h_at_win = beam["y"][-2] + dy_fm_ex / dz_fm_ex * dz_fm_win
-
-    # beam["heightWindow"] = h_at_win
-
     return beam
diff --git a/debye_bec/bec_widgets/widgets/digital_twin/calc_surfaces.py b/debye_bec/bec_widgets/widgets/digital_twin/calc_surfaces.py
index d8e80ce..168f977 100644
--- a/debye_bec/bec_widgets/widgets/digital_twin/calc_surfaces.py
+++ b/debye_bec/bec_widgets/widgets/digital_twin/calc_surfaces.py
@@ -1,17 +1,28 @@
-import os
+"""
+Calculates the surface coordinates based on a beamline config.
+"""
+
 import re
 
 import numpy as np
 from bec_lib import bec_logger
 
+import debye_bec.bec_widgets.widgets.x01da_parameters as bl
+from debye_bec.bec_widgets.widgets.digital_twin.types import ConfigDict, SurfaceDict
+
 logger = bec_logger.logger
 
-os.environ["USE_XRT"] = "False"
-import debye_bec.bec_widgets.widgets.x01da_parameters as bl
-from debye_bec.bec_widgets.widgets.digital_twin.types import SurfaceDict
 
+def calc_surfaces(cfg: ConfigDict) -> SurfaceDict:
+    """
+    Calculates the surface coordinates based on a beamline config.
 
-def calc_surfaces(cfg):
+    Args:
+        cfg(ConfigDict): Dictionary with beamline config
+
+    Returns:
+        SurfaceDict: Surface data
+    """
 
     out: SurfaceDict = {
         "cm": {"x": [], "y": []},
@@ -45,39 +56,39 @@ def calc_surfaces(cfg):
     )  # Remove brackets from xtal name to conform with parameters
     index = bl.mo1.xtal.index(xtal)
 
-    xtalPos = bl.mo1.xtalOffsetX[index]
-    xtalLength1 = bl.mo1.xtalLength1[index]
-    xtalLength2 = bl.mo1.xtalLength2[index]
+    xtal_pos = bl.mo1.xtalOffsetX[index]
+    xtal_length_1 = bl.mo1.xtalLength1[index]
+    xtal_length_2 = bl.mo1.xtalLength2[index]
 
-    widthBeam = 2 * bl.mo1.center[1] * np.tan(cfg["h_acc"])
+    width_beam = 2 * bl.mo1.center[1] * np.tan(cfg["h_acc"])
 
-    heightBeam = 2 * bl.cm.center[1] * np.tan(cfg["v_acc"])
-    w = heightBeam / np.sin(cfg["mo1_bragg"])
+    height_beam = 2 * bl.cm.center[1] * np.tan(cfg["v_acc"])
+    w = height_beam / np.sin(cfg["mo1_bragg"])
 
     if cfg["mo1_mode"] in "Monochromatic":
         out["mo1_1"]["x"] = [
-            xtalPos - widthBeam / 2,
-            xtalPos + widthBeam / 2,
-            xtalPos + widthBeam / 2,
-            xtalPos - widthBeam / 2,
+            xtal_pos - width_beam / 2,
+            xtal_pos + width_beam / 2,
+            xtal_pos + width_beam / 2,
+            xtal_pos - width_beam / 2,
         ]
         out["mo1_1"]["y"] = [
-            xtalLength1 / 2 - c - w / 2,
-            xtalLength1 / 2 - c - w / 2,
-            xtalLength1 / 2 - c + w / 2,
-            xtalLength1 / 2 - c + w / 2,
+            xtal_length_1 / 2 - c - w / 2,
+            xtal_length_1 / 2 - c - w / 2,
+            xtal_length_1 / 2 - c + w / 2,
+            xtal_length_1 / 2 - c + w / 2,
         ]
         out["mo1_2"]["x"] = [
-            xtalPos - widthBeam / 2,
-            xtalPos + widthBeam / 2,
-            xtalPos + widthBeam / 2,
-            xtalPos - widthBeam / 2,
+            xtal_pos - width_beam / 2,
+            xtal_pos + width_beam / 2,
+            xtal_pos + width_beam / 2,
+            xtal_pos - width_beam / 2,
         ]
         out["mo1_2"]["y"] = [
-            -xtalLength2 / 2 + e - w / 2,
-            -xtalLength2 / 2 + e - w / 2,
-            -xtalLength2 / 2 + e + w / 2,
-            -xtalLength2 / 2 + e + w / 2,
+            -xtal_length_2 / 2 + e - w / 2,
+            -xtal_length_2 / 2 + e - w / 2,
+            -xtal_length_2 / 2 + e + w / 2,
+            -xtal_length_2 / 2 + e + w / 2,
         ]
     else:  # Pinkbeam
         out["mo1_1"]["x"] = []
@@ -98,50 +109,44 @@ def calc_surfaces(cfg):
         r = bl.fm.r[index]
     off = -cfg["fm_trx"]
 
-    widthBeam = 2 * bl.fm.center[1] * np.tan(cfg["h_acc"])
+    width_beam = 2 * bl.fm.center[1] * np.tan(cfg["h_acc"])
 
     if cfg["fm_stripe"] in ("Rh (toroid)", "Pt (toroid)"):
 
-        l = heightBeam / np.sin(cfg["fm_rotx"])
-        alpha = np.arccos(1 - widthBeam**2 / (2 * r**2))
+        l = height_beam / np.sin(cfg["fm_rotx"])
+        alpha = np.arccos(1 - width_beam**2 / (2 * r**2))
         h = r - (r * np.cos(alpha / 2))
         z = h / np.tan(cfg["fm_rotx"])
 
-        x = [off - widthBeam / 2, off - widthBeam / 2]
+        x = [off - width_beam / 2, off - width_beam / 2]
         y = [l / 2 - z / 2, -l / 2 - z / 2]
 
-        # logger.info(f'stripe: {cfg["fm_stripe"]}')
-        # logger.info(f'fm_rotx: {cfg["fm_rotx"]}')
-        # logger.info(f'h: {h}')
-        # logger.info(f'z: {z}')
-        # logger.info(f'r: {r}')
-
         res = 20
-        xElipse = np.linspace(0, np.pi, res)
-        yElipse = np.linspace(0, np.pi, res)
-        xElipse = [-widthBeam / 2 * np.cos(i) + off for i in xElipse]
-        yElipse = [widthBeam * np.sin(i) * z / widthBeam - l / 2 - z / 2 for i in yElipse]
+        x_elipse = np.linspace(0, np.pi, res)
+        y_elipse = np.linspace(0, np.pi, res)
+        x_elipse = [-width_beam / 2 * np.cos(i) + off for i in x_elipse]
+        y_elipse = [width_beam * np.sin(i) * z / width_beam - l / 2 - z / 2 for i in y_elipse]
 
-        x.extend(xElipse)
-        y.extend(yElipse)
+        x.extend(x_elipse)
+        y.extend(y_elipse)
 
-        x.extend([off + widthBeam / 2, off + widthBeam / 2])
+        x.extend([off + width_beam / 2, off + width_beam / 2])
         y.extend([-l / 2 - z / 2, l / 2 - z / 2])
 
         res = 50
-        xElipse = np.linspace(np.pi, 0, res)
-        yElipse = np.linspace(np.pi, 0, res)
-        xElipse = [-widthBeam / 2 * np.cos(i) + off for i in xElipse]
-        yElipse = [widthBeam * np.sin(i) * z / widthBeam + l / 2 - z / 2 for i in yElipse]
+        x_elipse = np.linspace(np.pi, 0, res)
+        y_elipse = np.linspace(np.pi, 0, res)
+        x_elipse = [-width_beam / 2 * np.cos(i) + off for i in x_elipse]
+        y_elipse = [width_beam * np.sin(i) * z / width_beam + l / 2 - z / 2 for i in y_elipse]
 
-        x.extend(xElipse)
-        y.extend(yElipse)
+        x.extend(x_elipse)
+        y.extend(y_elipse)
 
         out["fm"]["x"] = x
         out["fm"]["y"] = y
 
     else:  # flat surface, no toroid
-        l = heightBeam / np.sin(cfg["fm_rotx"])
+        l = height_beam / np.sin(cfg["fm_rotx"])
 
         w1 = 2 * (bl.fm.center[1] - l / 2) * np.tan(cfg["h_acc"])
         w2 = 2 * (bl.fm.center[1] + l / 2) * np.tan(cfg["h_acc"])
diff --git a/debye_bec/bec_widgets/widgets/digital_twin/calc_varia.py b/debye_bec/bec_widgets/widgets/digital_twin/calc_varia.py
index 710212a..9ec7cc8 100644
--- a/debye_bec/bec_widgets/widgets/digital_twin/calc_varia.py
+++ b/debye_bec/bec_widgets/widgets/digital_twin/calc_varia.py
@@ -1,4 +1,9 @@
+"""
+Various calculations for the digital twin
+"""
+
 import re
+from typing import Literal, cast
 
 import numpy as np
 from bec_lib import bec_logger
@@ -9,19 +14,41 @@ import debye_bec.bec_widgets.widgets.x01da_parameters as bl
 
 logger = bec_logger.logger
 
+H = 6.62606957e-34
+E = 1.602176634e-19
+C = 299792458
+RE = 2.8179e-15
 
-def sldi_gap_to_acc(sldi_gapx, sldi_gapy):
+
+def sldi_gap_to_acc(sldi_gapx: float, sldi_gapy: float) -> tuple[float, float]:
+    """
+    Calculate the slits acceptance based on the gap values
+
+    Args:
+        sldi_gapx(float): GAPX value of the slits in mm
+        sldi_gapy(float): GAPY value of the slits in mm
+
+    Returns:
+        tuple[float, float]: Horizontal and vertical acceptance in rad
+    """
     d1 = bl.feSlits.center1[1]
     d2 = bl.feSlits.center2[1]
     h_acc = np.tan(sldi_gapx / (d2 + d1))
     v_acc = np.tan(sldi_gapy / (d2 + d1))
-
-    # h_acc = np.tan(sldi_gapx / (2 * d1))
-    # v_acc = np.tan(sldi_gapy / (2 * d1))
     return h_acc, v_acc
 
 
-def cm_trx_to_stripe(cm_trx):
+def cm_trx_to_stripe(cm_trx: float) -> str | None:
+    """
+    Based on the trx value of the collimating mirror, return
+    the correct stripe
+
+    Args:
+        cm_trx(float): Collimating mirror trx value
+
+    Returns
+        str | None: Stripe of the mirror, None if not found
+    """
     cm_stripe = None
     for name, low, high in zip(bl.cm.surface, bl.cm.limOptX[0], bl.cm.limOptX[1]):
         if low <= cm_trx <= high:
@@ -29,14 +56,34 @@ def cm_trx_to_stripe(cm_trx):
     return cm_stripe
 
 
-def cm_stripe_to_trx(cm_stripe):
+def cm_stripe_to_trx(cm_stripe: str) -> float | None:
+    """
+    Based on the stripe of the collimating mirror, return
+    the trx value
+
+    Args:
+        cm_stripe(str): Stripe of the collimating mirror
+
+    Returns:
+        float | None: TRX value of the stripe. None if not found
+    """
     for name, low, high in zip(bl.cm.surface, bl.cm.limOptX[0], bl.cm.limOptX[1]):
         if cm_stripe == name:
             return -(low + high) / 2
-    return 0
+    return None
 
 
-def fm_trx_to_stripe(fm_trx):
+def fm_trx_to_stripe(fm_trx: float) -> str | None:
+    """
+    Based on the trx value of the focusing mirror, return
+    the correct stripe
+
+    Args:
+        fm_trx(float): focusing mirror trx value
+
+    Returns
+        str | None: Stripe of the mirror, None if not found
+    """
     fm_stripe = None
     for name, low, high in zip(bl.fm.surfaceFlat, bl.fm.limOptXFlat[1], bl.fm.limOptXFlat[0]):
         if low <= fm_trx <= high:
@@ -47,17 +94,37 @@ def fm_trx_to_stripe(fm_trx):
     return fm_stripe
 
 
-def fm_stripe_to_trx(fm_stripe):
+def fm_stripe_to_trx(fm_stripe: str) -> float | None:
+    """
+    Based on the stripe of the focusing mirror, return
+    the trx value
+
+    Args:
+        fm_stripe(str): Stripe of the focusing mirror
+
+    Returns:
+        float | None: TRX value of the stripe. None if not found
+    """
     for name, low, high in zip(bl.fm.surfaceFlat, bl.fm.limOptXFlat[1], bl.fm.limOptXFlat[0]):
         if fm_stripe == name + " (flat)":
             return (low + high) / 2
     for name, low, high in zip(bl.fm.surfaceToroid, bl.fm.limOptXToroid[1], bl.fm.limOptXToroid[0]):
         if fm_stripe == name + " (toroid)":
             return -(low + high) / 2
-    return 0
+    return None
 
 
-def mo1_energy_resolution(xtal, energy):
+def mo1_energy_resolution(xtal: Literal["Si111", "Si311"], energy: float) -> float:
+    """
+    Calculate the energy resolution of the monochromator
+
+    Args:
+        xtal(str): Xtal name. "Si111" or "Si311"
+        energy(float): Energy in eV
+
+    Returns:
+        float: Energy resolution in eV
+    """
     index = bl.mo1.xtal.index(xtal)
     crystal = bl.mo1.material1[index]
 
@@ -69,29 +136,54 @@ def mo1_energy_resolution(xtal, energy):
 
     # FWHM of the DCM curve
     spline = UnivariateSpline(dtheta, refl2 - refl2.max() / 2, s=0)
-    r1, r2 = spline.roots()
+    roots = cast(np.ndarray, spline.roots())
+    r1, r2 = float(roots[0]), float(roots[1])
     fwhm_rad = (r2 - r1) * 1e-6  # µrad → rad
 
     # Energy resolution
-    theta_B = crystal.get_Bragg_angle(energy)
-    dE_over_E = fwhm_rad / np.tan(theta_B)
-    dE = dE_over_E * energy
+    theta_b = crystal.get_Bragg_angle(energy)
+    de_over_e = fwhm_rad / np.tan(theta_b)
+    de = de_over_e * energy
 
     # logger.info(f"DCM FWHM      : {r2-r1:.2f} µrad")
     # logger.info(f"ΔE/E          : {dE_over_E:.2e}")
     # logger.info(f"ΔE            : {dE:.3f} eV  at {E} eV")
 
-    return dE
+    return de
 
 
-def cm_reflectivity(cm_stripe, cm_pitch, energy):
+def cm_reflectivity(cm_stripe: str, cm_pitch: float, energy: float) -> float:
+    """
+    Calculate the reflectivity of the mirror stripe based
+    on the pitch and energy.
+
+    Args:
+        cm_stripe(str): Mirror stripe
+        cm_pitch(float): Pitch of the mirror (beam incidence angle)
+        energy(float): Energy of the beam in eV
+
+    Returns:
+        float: Reflectivity [0-1]
+    """
     index = bl.cm.surface.index(cm_stripe)
-    rs, rp = bl.cm.material[index].get_amplitude(energy, np.sin(cm_pitch))[0:2]
+    rs, _ = bl.cm.material[index].get_amplitude(energy, np.sin(cm_pitch))[0:2]
     refl = abs(rs) ** 2
     return refl
 
 
-def fm_reflectivity(fm_stripe, fm_pitch, energy):
+def fm_reflectivity(fm_stripe: str, fm_pitch: float, energy: float) -> float:
+    """
+    Calculate the reflectivity of the mirror stripe based
+    on the pitch and energy.
+
+    Args:
+        cm_stripe(str): Mirror stripe
+        cm_pitch(float): Pitch of the mirror (beam incidence angle)
+        energy(float): Energy of the beam in eV
+
+    Returns:
+        float: Reflectivity [0-1]
+    """
     if fm_stripe in ("Rh (toroid)", "Pt (toroid)"):
         surface = bl.fm.surfaceToroid
         material = bl.fm.materialToroid
@@ -102,35 +194,75 @@ def fm_reflectivity(fm_stripe, fm_pitch, energy):
         material = bl.fm.materialFlat
         stripe = re.sub(r"\s*\(.*?\)", "", fm_stripe).strip()
         index = surface.index(stripe)
-    rs, rp = material[index].get_amplitude(energy, np.sin(fm_pitch))[0:2]
+    rs, _ = material[index].get_amplitude(energy, np.sin(fm_pitch))[0:2]
     refl = abs(rs) ** 2
     return refl
 
 
-def mo1_bragg_angle(mo_mode, d_spacing, energy, cm_pitch):
-    H = 6.62606957e-34
-    E = 1.602176634e-19
-    C = 299792458
+def mo1_bragg_angle(
+    mo_mode: Literal["Monochromatic", "Pinkbeam"], d_spacing: float, energy: float, cm_pitch: float
+) -> tuple[float, float]:
+    """
+    Calculate the bragg angle of the monochromator.
+    Corrects for the collimating mirror pitch.
+
+    Args:
+        mo_mode(str): Monochromator mode. "Monochromatic" or "Pinkbeam"
+        d_spacing(float): D-spacing of the crystal in Angstrom
+        energy(float): Energy of the beam in eV
+        cm_pitch(float): Pitch of collimating mirror in rad
+
+    Returns:
+        tuple[float, float]: Bragg angle and corrected bragg angle
+    """
     wl = C * H / (E * energy)
     val = wl / (2 * d_spacing * 1e-10)
     bragg_angle = 0
     if val > -1 and val < 1:
         bragg_angle = np.asin(val)
-    if mo_mode in "Monochromatic":
+    if mo_mode == "Monochromatic":
         # Add 2x CM pitch to the bragg angle
         bragg_angle_cor = (2 * cm_pitch) + bragg_angle
-    elif mo_mode in "Pinkbeam":
+    else:
         # Align xtal surfaces parallel to beam
         bragg_angle_cor = 2 * cm_pitch
     return bragg_angle, bragg_angle_cor
 
 
 def fm_ideal_pitch(
-    fm_focus, fm_stripe, smpl, sldi_hacc=None, sldi_vacc=None, fm_focx=None, fm_focy=None
-):
+    fm_focus: Literal["Defocused", "Focused", "Manual"],
+    fm_stripe: str,
+    smpl: float,
+    sldi_hacc: float | None = None,
+    sldi_vacc: float | None = None,
+    fm_focx: float | None = None,
+    fm_focy: float | None = None,
+) -> tuple[float, float | None]:
+    """
+    Calculates the ideal pitch for the focusing mirror depending on the
+    focusing strategy.
+    If "Defocused" is chosed, sldi_hacc, sldi_vacc, fm_focx and fm_focy
+    must be provided.
+
+    Args:
+        fm_focus(str): Focus strategy. "Defocused", "Focused" or "Manual
+        fm_stripe(str): Mirror stripe
+        smpl(float): Sample position in mm from source
+        sldi_hacc(float): Horizontal acceptance of frontend slits. Defaults to None
+        sldi_vacc(float): Vertical acceptance of frontend slits. Defaults to None
+        fm_focx(float): Requested horizontal spot size in mm. Defaults to None
+        fm_focy(float): Requested vertical spot size in mm. Defaults to None
+
+    Returns:
+        tuple[float, float | None]: Pitch of mirror in rad, qy in mm
+    """
     p = bl.fm.center[1]  # posFM
     q = smpl - bl.fm.center[1]  # dist posFM to posEX
     if fm_focus in "Defocused":
+        assert sldi_hacc is not None, "sldi_hacc must be provided for Defocused mode"
+        assert sldi_vacc is not None, "sldi_vacc must be provided for Defocused mode"
+        assert fm_focx is not None, "fm_focx must be provided for Defocused mode"
+        assert fm_focy is not None, "fm_focy must be provided for Defocused mode"
         a = 2 * np.tan(sldi_hacc) * bl.fm.center[1]  # Beam width at focusing mirror
         b = (
             2 * np.tan(sldi_vacc) * bl.cm.center[1]
@@ -151,45 +283,77 @@ def fm_ideal_pitch(
     return pitch, qy
 
 
-def cm_critical_angle(cm_stripe, energy):
+def cm_critical_angle(cm_stripe: Literal["Si", "Pt", "Rh"], energy) -> float:
+    """
+    Calculate the critical angle of the mirror stripe
+
+    Args:
+        cm_stripe(str): Mirror stripe. "Si", "Pt" or "Rh"
+        energy(float): Energy in eV
+
+    Returns:
+        float: Critical angle in rad
+    """
     if cm_stripe in "Si":
         stripe = bl.stripeSi
     elif cm_stripe in "Pt":
         stripe = bl.stripePt
-    elif cm_stripe in "Rh":
-        stripe = bl.stripeRh
     else:
-        raise Exception(f"Stripe {stripe} not found in beamline parameters!")
+        stripe = bl.stripeRh
     w = CHeVcm / 100 / energy  # convert energy [eV] to wavelength [m]
-    # Calculate critical angle for mirror
     f1 = stripe.elements[0].Z + np.real(stripe.elements[0].get_f1f2(energy))
-    numberDensity = stripe.rho * 1e3 * AVOGADRO / (stripe.elements[0].mass / 1e3)
-    criticalAngle = np.sqrt(numberDensity * 2.8179e-15 * w**2 * f1 / np.pi)
-    return criticalAngle
+    number_density = stripe.rho * 1e3 * AVOGADRO / (stripe.elements[0].mass / 1e3)
+    critical_angle = np.sqrt(number_density * RE * w**2 * f1 / np.pi)
+    return critical_angle
 
 
-def mirror_surface_geometries(mirror):
+def mirror_surface_geometries(
+    mirror: Literal["cm", "fm_toroid", "fm_flat"],
+) -> dict[str, tuple[float, float, float, float]]:
+    """
+    Return the mirror stripe geometries
+
+    Args:
+        mirror(str): Mirror. "cm", "fm_toroid" or "fm_flat"
+
+    Returns:
+        dict[str, tuple[float, float, float, float]]: Dictionary mapping surface
+        names to tuples of (x, y, width, height).
+    """
     if mirror in "cm":
         surface = bl.cm.surface
-        limOptX = bl.cm.limOptX
-        limOptY = bl.cm.limOptY
+        lim_opt_x = bl.cm.limOptX
+        lim_opt_y = bl.cm.limOptY
     elif mirror in "fm_toroid":
         surface = bl.fm.surfaceToroid
-        limOptX = bl.fm.limOptXToroid
-        limOptY = bl.fm.limOptYToroid
+        lim_opt_x = bl.fm.limOptXToroid
+        lim_opt_y = bl.fm.limOptYToroid
     elif mirror in "fm_flat":
         surface = bl.fm.surfaceFlat
-        limOptX = bl.fm.limOptXFlat
-        limOptY = bl.fm.limOptYFlat
+        lim_opt_x = bl.fm.limOptXFlat
+        lim_opt_y = bl.fm.limOptYFlat
     else:
         raise ValueError(f"Requested mirror {mirror} not available!")
     geom = {}
-    for sf, lx, hx, ly, hy in zip(surface, limOptX[0], limOptX[1], limOptY[0], limOptY[1]):
+    for sf, lx, hx, ly, hy in zip(surface, lim_opt_x[0], lim_opt_x[1], lim_opt_y[0], lim_opt_y[1]):
         geom[sf] = (lx, ly, hx - lx, hy - ly)
     return geom
 
 
-def mo_surface_geometries(mo, plane):
+def mo_surface_geometries(
+    mo: Literal["mo1"], plane: Literal[0, 1]
+) -> dict[str, tuple[float, float, float, float]]:
+    """
+    Return the monochromator xtal geometries
+
+    Args:
+        mo(str): Monochromator. Only "mo1" implemented
+        plane(int): Surface of xtal. 0 and 1 (First and second)
+
+    Returns:
+        dict[str, tuple[float, float, float, float]]: Dictionary mapping surface
+        names to tuples of (x, y, width, height).
+    """
     if mo in "mo1":
         xtal = bl.mo1.xtal
         xtal_width = bl.mo1.xtalWidth
@@ -199,14 +363,20 @@ def mo_surface_geometries(mo, plane):
         else:
             xtal_length = bl.mo1.xtalLength2
     else:
-        raise ValueError(f"Requested mono {mo} not available!")
+        return {}
     geom = {}
     for sf, w, offx, length in zip(xtal, xtal_width, xtal_offset_x, xtal_length):
         geom[sf] = (offx - w / 2, -length / 2, w, length)
     return geom
 
 
-def wall_geometries():
+def wall_geometries() -> list[list[float]]:
+    """
+    Return the wall geometries
+
+    Returns:
+        list[list[float]]: List of [x, y, width, height] geometry values for each wall.
+    """
     geom = []
     for i, _ in enumerate(bl.walls.start):
         geom.append(
@@ -220,7 +390,15 @@ def wall_geometries():
     return geom
 
 
-def pipe_geometries():
+def pipe_geometries() -> list[dict[str, np.ndarray]]:
+    """
+    Return the wall geometries
+
+    Returns:
+        list[dict[str, np.ndarray]]: List of dictionaries with keys "x" and "y",
+            each containing a numpy array of two float values representing
+            the start and end coordinates of the pipe top and bottom edges.
+    """
     pipes = []
     for i, _ in enumerate(bl.vacuum_pipes.center):
         top = bl.vacuum_pipes.center[i] + bl.vacuum_pipes.diameter[i] / 2 + bl.sourceHeight
diff --git a/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py b/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py
index d98975d..985c8f5 100644
--- a/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py
+++ b/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py
@@ -4,7 +4,7 @@ Digital Twin: Custom BEC widget to support the beamline alignment.
 
 import sys
 from pathlib import Path
-from typing import Literal
+from typing import Literal, cast
 
 import numpy as np
 import yaml
@@ -48,6 +48,7 @@ from debye_bec.bec_widgets.widgets.digital_twin.input_panel import InputPanel
 from debye_bec.bec_widgets.widgets.digital_twin.mover_panel import MoverPanel
 from debye_bec.bec_widgets.widgets.digital_twin.plots import SideviewPlot, SurfacePlots
 from debye_bec.bec_widgets.widgets.digital_twin.settings_panel import SettingsPanel
+from debye_bec.bec_widgets.widgets.digital_twin.types import ConfigDict
 
 logger = bec_logger.logger
 
@@ -62,7 +63,7 @@ class DigitalTwin(BECWidget, QWidget):
     PLUGIN = True
     ICON_NAME = "lightbulb"
 
-    def __init__(self, parent=None, *arg, **kwargs):
+    def __init__(self, *arg, parent=None, **kwargs):
         super().__init__(parent=parent, theme_update=True, *arg, **kwargs)
         self.get_bec_shortcuts()
 
@@ -115,15 +116,15 @@ class DigitalTwin(BECWidget, QWidget):
         self.settings.reload_offsets.clicked_connect(self.load_offsets)
         self.settings.unload_offsets.clicked_connect(self.unload_offsets)
 
-        self.bragg_angle = 0
-        self.qy = 0
+        self.bragg_angle = 0.0
+        self.qy = 0.0
         self.offsets = {}
 
         # Initialize all values
         self.load_offsets(recalculate=False)
         self.calc_assistant(identifier="init")
 
-        # Timer: update plot every 1 second
+        # Timer: update plots every 1 second
         self._timer = QTimer(self)
         self._timer.setInterval(100)
         self._timer.timeout.connect(self.calc_reality)
@@ -142,6 +143,12 @@ class DigitalTwin(BECWidget, QWidget):
 
     @SafeSlot()
     def check_config(self, *args):
+        """
+        Checks the BEC config and opens a window if not all necessary
+        devices are loaded in the config. If called from a slot from
+        BEC dispatcher whenever there is a config update, stop the timer
+        that updates the plot in the background.
+        """
         reload = (args[0] if args else {}).get("action") == "reload"
         if reload:
             self._timer.stop()
@@ -221,14 +228,21 @@ class DigitalTwin(BECWidget, QWidget):
             dialog.show()
             info.setMinimumHeight(info.heightForWidth(info.width()))
             if dialog.exec_() == QDialog.DialogCode.Rejected:
-                app = QApplication.instance()
-                if app is not None:
-                    app.exit(0)
+                running_app = QApplication.instance()
+                if running_app is not None:
+                    running_app.exit(0)
         if reload:
             self._timer.start()
 
     @SafeSlot()
-    def calc_assistant(self, *args, **kwargs):
+    def calc_assistant(self, *_, **kwargs):
+        """
+        Calculates various values for the assistant.
+        If called from a qt slot, the identifier represents
+        the button pressed / value changed. Based on the identifier,
+        calculate different values.
+        Note: identifier=init calculates all values
+        """
         identifier = kwargs["identifier"]
         match identifier:
             case "init":
@@ -281,6 +295,13 @@ class DigitalTwin(BECWidget, QWidget):
         self.calc_assistant_surfaces()
 
     def get_assistant_config(self, apply_offset: bool = False):
+        """
+        Assembles the digital twin config from the assistants input.
+
+        Args:
+            apply_offset(bool): Applies the offset values to the config.
+                                Defaults to False
+        """
         fm_focus = self.input.fm_focus.currentText()
         if fm_focus in "Manual":
             fm_rotx = self.input.fm_rotx.value()
@@ -297,7 +318,10 @@ class DigitalTwin(BECWidget, QWidget):
         fm_stripe = self.input.fm_stripe.currentText()
         fm_trx = fm_stripe_to_trx(fm_stripe)
 
-        config = {
+        assert cm_trx is not None, f"No cm_trx found for given stripe {cm_stripe}!"
+        assert fm_trx is not None, f"No fm_trx found for given stripe {fm_stripe}!"
+
+        config: ConfigDict = {
             "energy": self.input.energy.value(),
             "h_acc": self.input.sldi_hacc.value(),
             "v_acc": self.input.sldi_vacc.value(),
@@ -320,16 +344,10 @@ class DigitalTwin(BECWidget, QWidget):
             for axis, _ in config.items():
                 if axis in self.offsets:
                     axis_offsets = self.offsets[axis]
-                    logger.info(f"Axis: {axis}")
                     if "modifier" in axis_offsets and "offset" in axis_offsets:
                         for idx, rng in enumerate(axis_offsets["modifier"]["range"]):
-                            logger.info(f"rng: {rng}")
-                            logger.info(f'value: {config[axis_offsets["modifier"]["axis"]]}')
                             if rng[0] < config[axis_offsets["modifier"]["axis"]] < rng[1]:
-                                logger.info(f'offset: {axis_offsets["offset"][idx]}')
-                                # logger.info(f"axis_data before: {axis_data}")
                                 config[axis] += axis_offsets["offset"][idx]
-                                # logger.info(f"axis_data after: {axis_data}")
                                 break
                     elif "offset" in axis_offsets:
                         config[axis] += axis_offsets["offset"]
@@ -344,6 +362,9 @@ class DigitalTwin(BECWidget, QWidget):
         return config
 
     def get_reality_config(self):
+        """
+        Assembles the digital twin config based on the real axis positions.
+        """
         mo1_trx = self.dev.mo1_trx.read(cached=True)["mo1_trx"]["value"]
         if abs(mo1_trx) > 5:
             mo1_mode = "Monochromatic"
@@ -361,7 +382,7 @@ class DigitalTwin(BECWidget, QWidget):
         fm_rotx = self.dev.fm_rotx.read(cached=True)["fm_rotx"]["value"]
         fm_rotx_real = 2 * cm_pitch - fm_rotx
         smpl = self.dev.ot_es1_trz.read(cached=True)["ot_es1_trz"]["value"]
-        config = {  # Config in SI units!
+        raw = {  # Config in SI units!
             "energy": mo1_bragg["mo1_bragg"]["value"],
             "h_acc": h_acc,
             "v_acc": v_acc,
@@ -374,9 +395,11 @@ class DigitalTwin(BECWidget, QWidget):
             "fm_rotx": -fm_rotx_real * 1e-3,
             "fm_stripe": fm_stripe,
             "fm_trx": fm_trx,
+            "fm_qy": None,
             "fm_gain_height": 1,
             "smpl": smpl,
         }
+        config = cast(ConfigDict, raw)
         # logger.info(f'Config created: {config}')
 
         abs_open = self.dev.abs.read(cached=True)["abs_status_string"]["value"] == "OPEN"
@@ -430,7 +453,12 @@ class DigitalTwin(BECWidget, QWidget):
         self.mover.abs.set_feedback(abs_open)
         return config
 
-    def adapt_reality(self, *args):
+    @SafeSlot()
+    def adapt_reality(self, *_):
+        """
+        Based on the real axis positions, adjust the assistant to reflect
+        the reality.
+        """
         pos = {}
         pos["sldi_gapx"] = self.dev.sldi_gapx.read(cached=True)["sldi_gapx"]["value"]
         pos["sldi_gapy"] = self.dev.sldi_gapy.read(cached=True)["sldi_gapy"]["value"]
@@ -474,7 +502,15 @@ class DigitalTwin(BECWidget, QWidget):
         self.input.smpl.set_number(pos["ot_es1_trz"])
         self.calc_assistant(identifier="init")
 
-    def load_offsets(self, recalculate=True, *args):
+    @SafeSlot()
+    def load_offsets(self, *_, recalculate: bool = True):
+        """
+        Loads the offsets from the file
+
+        Args:
+            recalculate(bool): Recalculates the assistant values.
+                               Defaults to True
+        """
         file = Path(OFFSET_FILE)
         if not file.exists():
             raise FileNotFoundError(f"Offset file not found: {OFFSET_FILE}")
@@ -490,11 +526,19 @@ class DigitalTwin(BECWidget, QWidget):
         if recalculate:
             self.calc_assistant(identifier="init")
 
-    def unload_offsets(self, *args):
+    @SafeSlot()
+    def unload_offsets(self, *_):
+        """
+        Removes the offsets and recalculates the assistant values.
+        """
         self.offsets = {}
         self.calc_assistant(identifier="init")
 
     def update_fm_mode(self):
+        """
+        Updates the focusing mirror input group based on the
+        selection of the focus strategy.
+        """
         fm_focus = self.input.fm_focus.currentText()
         if fm_focus in "Manual":
             self.input.fm_rotx.setVisible(True)
@@ -515,22 +559,32 @@ class DigitalTwin(BECWidget, QWidget):
             self.input.fm_focy.setVisible(True)
             self.input.fm_rotx_ideal.setLabel("Incidence Angle for defocused beam")
 
+    @SafeSlot()
     def calc_reality(self):
+        """
+        Updates the plots for the reality scene
+        """
         config = self.get_reality_config()
         data = calc_sideview(config)
         self.sideview_plot.update_curves("reality", data=data)
-        # logger.info('Calc reality surfaces')
         surfaces = calc_surfaces(config)
         self.surface_plots.update_surfaces(scene="reality", data=surfaces)
 
     def calc_mo1_energy_resolution(self):
+        """
+        Calculates the energy resolution of the monochromator
+        """
         xtal = self.input.mo1_xtal.currentText().translate(
             str.maketrans("", "", "()")
         )  # Remove brackets from xtal name to conform with parameters
+        xtal = cast(Literal["Si111", "Si311"], xtal)
         energy = self.input.energy.value()
         self.input.mo1_eres.setValue(mo1_energy_resolution(xtal, energy))
 
     def calc_cm_reflectivity(self):
+        """
+        Calculates the collimating mirror reflectivity
+        """
         cm_stripe = self.input.cm_stripe.currentText()
         cm_pitch = -self.input.cm_pitch.value() * 1e-3
         energy = self.input.energy.value()
@@ -540,6 +594,9 @@ class DigitalTwin(BECWidget, QWidget):
         self.input.cm_refl_harm.setLabel(f"Reflectivity at \n{3*energy:.0f} eV")
 
     def calc_fm_reflectivity(self):
+        """
+        Calculates the focusing mirror reflectivity
+        """
         fm_stripe = self.input.fm_stripe.currentText()
         fm_focus = self.input.fm_focus.currentText()
         if fm_focus in "Manual":
@@ -553,6 +610,9 @@ class DigitalTwin(BECWidget, QWidget):
         self.input.fm_refl_harm.setLabel(f"Reflectivity at \n{3*energy:.0f} eV")
 
     def calc_cm_fm_harm_suppr(self):
+        """
+        Calculates the combined harmonics suppression of both mirrors
+        """
         harm_suppr = (self.input.cm_refl.value() * self.input.fm_refl.value()) / (
             self.input.cm_refl_harm.value() * self.input.fm_refl_harm.value()
         )
@@ -562,16 +622,24 @@ class DigitalTwin(BECWidget, QWidget):
         )
 
     def calc_assistant_sideview(self):
+        """
+        Updates the sideview plot based on the assistant values
+        """
         config = self.get_assistant_config(apply_offset=True)
         data = calc_sideview(config)
         self.sideview_plot.update_curves("assistant", data)
 
     def calc_assistant_surfaces(self):
-        # logger.info('Calc assistant surfaces')
+        """
+        Updates the surface plot based on the assistant values
+        """
         surfaces = calc_surfaces(self.get_assistant_config())
         self.surface_plots.update_surfaces(scene="assistant", data=surfaces)
 
     def calc_positions(self):
+        """
+        Calculates the positions for the axes based on the assistant values
+        """
         out = calc_positions(self.get_assistant_config())
 
         # Apply offsets
@@ -628,13 +696,17 @@ class DigitalTwin(BECWidget, QWidget):
         else:
             raise ValueError(f"Invalid xtal selection: {xtal}")
         cm_pitch = -self.dev.cm_rotx.read(cached=True)["cm_rotx"]["value"] * 1e-3
-        mo1_mode = self.input.mo1_mode.currentText()
+        mo1_mode = cast(Literal["Monochromatic", "Pinkbeam"], self.input.mo1_mode.currentText())
         energy = self.input.energy.value()
         theta, _ = mo1_bragg_angle(mo1_mode, d_spacing, energy, cm_pitch)
         self.bragg_angle = theta
         self.input.mo1_bragg_angle.setValue(theta / np.pi * 180)
 
     def update_mo1_mode(self):
+        """
+        Updates the monochromator input group based on the
+        selection of the mode.
+        """
         if self.input.mo1_mode.currentText() in "Monochromatic":
             self.input.mo1_xtal.setVisible(True)
             self.input.mo1_bragg_angle.setVisible(True)
@@ -645,7 +717,12 @@ class DigitalTwin(BECWidget, QWidget):
             self.input.mo1_eres.setVisible(False)
 
     def calc_fm_ideal_pitch(self):
-        fm_focus = self.input.fm_focus.currentText()
+        """
+        Calculate the ideal pitch for the focusing mirror.
+        """
+        fm_focus = cast(
+            Literal["Defocused", "Focused", "Manual"], self.input.fm_focus.currentText()
+        )
         fm_stripe = self.input.fm_stripe.currentText()
         smpl = self.input.smpl.value()
         sldi_hacc = self.input.sldi_hacc.value() * 1e-3
@@ -659,7 +736,10 @@ class DigitalTwin(BECWidget, QWidget):
         self.input.fm_rotx_ideal.setValue(-fm_rotx * 1e3)
 
     def calc_cm_crit_pitch(self):
-        cm_stripe = self.input.cm_stripe.currentText()
+        """
+        Calculate the critical pitch for the collimating mirror
+        """
+        cm_stripe = cast(Literal["Si", "Pt", "Rh"], self.input.cm_stripe.currentText())
         energy = self.input.energy.value()
         self.input.cm_pitch_critical.setValue(-cm_critical_angle(cm_stripe, energy) * 1e3)
 
diff --git a/debye_bec/bec_widgets/widgets/digital_twin/types.py b/debye_bec/bec_widgets/widgets/digital_twin/types.py
index 737845b..674b59f 100644
--- a/debye_bec/bec_widgets/widgets/digital_twin/types.py
+++ b/debye_bec/bec_widgets/widgets/digital_twin/types.py
@@ -1,8 +1,47 @@
-"""Types used for plotting data"""
+"""Types used for the beamline config and for plotting data"""
 
 from typing import TypedDict
 
 
+class ConfigDict(TypedDict):
+    """
+    Typed dictionary representing the beamline configuration.
+
+    Attributes:
+        energy (float): Beam energy.
+        h_acc (float): Horizontal acceptance.
+        v_acc (float): Vertical acceptance.
+        cm_pitch (float): CM pitch angle.
+        cm_stripe (str): CM stripe name.
+        cm_trx (float): CM translation x.
+        mo1_mode (str): MO1 mode.
+        mo1_xtal (str): MO1 crystal.
+        mo1_bragg (float): MO1 Bragg angle.
+        fm_rotx (float): FM rotation x.
+        fm_stripe (str): FM stripe name.
+        fm_trx (float): FM translation x.
+        fm_qy (float): FM qy value.
+        fm_gain_height (int): FM gain height.
+        smpl (float): Sample value.
+    """
+
+    energy: float
+    h_acc: float
+    v_acc: float
+    cm_pitch: float
+    cm_stripe: str
+    cm_trx: float
+    mo1_mode: str
+    mo1_xtal: str
+    mo1_bragg: float
+    fm_rotx: float
+    fm_stripe: str
+    fm_trx: float
+    fm_qy: None | float
+    fm_gain_height: int
+    smpl: float
+
+
 class DataDict(TypedDict):
     """
     Typed dictionary representing plot data.