diff --git a/artist/field/kinematics_rigid_body.py b/artist/field/kinematics_rigid_body.py
index d3ca3eb4f..72bf0f979 100644
--- a/artist/field/kinematics_rigid_body.py
+++ b/artist/field/kinematics_rigid_body.py
@@ -1,10 +1,16 @@
+import logging
+
 import torch
 
+from artist.field.actuators import Actuators
 from artist.field.kinematics import Kinematics
-from artist.geometry import rotations, transforms
+from artist.geometry import coordinates, rotations, transforms
 from artist.util import constants, indices, type_registry
 from artist.util.env import get_device
 
+log = logging.getLogger(__name__)
+"""A logger for the rigid body kinematics."""
+
 
 class RigidBody(Kinematics):
     """
@@ -43,11 +49,14 @@ class RigidBody(Kinematics):
     active_motor_positions : torch.Tensor
         The motor positions of active heliostats.
         Shape is ``[number_of_active_heliostats, 2]``.
-    artist_standard_orientation : torch.Tensor
-        The standard orientation of the kinematics.
-        Shape is ``[4]``.
     actuators : Actuators
         The actuators used in the kinematics.
+    kinematics_standard_orientation : torch.Tensor
+        Standard orientation of the kinematics system: south (0, -1, 0, 0) in homogeneous ENU.
+        Shape is ``[4]``
+    initial_orientation_offsets : torch.Tensor
+        Rotation matrix to account for the initial orientation offset between surface mesh and kinematics system.
+        Shape is ``[1, 4, 4]``.
 
     Methods
     -------
@@ -149,11 +158,7 @@ def __init__(
             self.motor_positions, device=device
         )
 
-        self.artist_standard_orientation = torch.tensor(
-            [0.0, -1.0, 0.0, 0.0], device=device
-        )
-
-        self.actuators = type_registry.actuator_type_mapping[
+        self.actuators: Actuators = type_registry.actuator_type_mapping[
             actuator_parameters_non_optimizable[
                 0, indices.actuator_type, indices.actuator_one_index
             ].item()
@@ -163,6 +168,23 @@ def __init__(
             device=device,
         )
 
+        # The surface points and normals are always sampled from a model (converted NURBS from deflectometry or ideal NURBS) that lays
+        # flat on the ground, i.e., the surface normals are pointing upwards [0.0, 0.0, 1.0]. Since the kinematics in ARTIST expects the
+        # points and normals to be initially oriented to the south, an extra rotation needs to be applied.
+        self.kinematics_standard_orientation = torch.tensor(
+            [0.0, -1.0, 0.0, 0.0], device=device
+        )
+        sampled_surface_orientation = torch.tensor([0.0, 0.0, 1.0, 0.0], device=device)
+        east_angles, north_angles, up_angles = rotations.decompose_rotations(
+            initial_vector=sampled_surface_orientation[None, :],
+            target_vector=self.kinematics_standard_orientation,
+        )
+        self.initial_orientation_offsets = (
+            transforms.rotate_e(e=east_angles, device=device)
+            @ transforms.rotate_n(n=north_angles, device=device)
+            @ transforms.rotate_u(u=up_angles, device=device)
+        )
+
     def _compute_orientations_from_motor_positions(
         self,
         motor_positions: torch.Tensor,
@@ -187,6 +209,8 @@ def _compute_orientations_from_motor_positions(
             The orientation matrices.
             Shape is ``[number_of_active_heliostats, 4, 4]``.
         """
+        device = get_device(device=device)
+
         joint_angles = self.actuators.motor_positions_to_angles(
             motor_positions=motor_positions,
             device=device,
@@ -291,17 +315,25 @@ def _compute_orientations_from_motor_positions(
 
         return orientations
 
-    def _apply_initial_orientation_offsets(
-        self, orientations: torch.Tensor, device: torch.device | None = None
+    def _compute_motor_positions_from_normal(
+        self,
+        normals: torch.Tensor,
+        device: torch.device | None = None,
     ) -> torch.Tensor:
         """
-        Apply the initial orientation offsets to the given orientation matrices.
+        Compute the motor positions from a normal vector.
+
+        This is the inverse kinematics. First the joint angles are computed from the desired normal vector.
+        Then the motor positions are computed from the joint angles.
+        The inverse kinematics produces two solutions, the valid solution is chosen according to resulting
+        motor positions which must lie within the minimum and maximum allowed motor positions defined within
+        the actuator parameters.
 
         Parameters
         ----------
-        orientations : torch.Tensor
-            The orientation matrices.
-            Shape is ``[number_of_active_heliostats, 4, 4]``.
+        normals : torch.Tensor
+            Concentrator normals.
+            Shape is ``[number_of_active_heliostats, 4]``.
         device : torch.device | None
             The device on which to perform computations or load tensors and models (default is None).
             If None, ``ARTIST`` will automatically select the most appropriate
@@ -310,35 +342,165 @@ def _apply_initial_orientation_offsets(
         Returns
         -------
         torch.Tensor
-            The orientation matrices with the initial orientation offset.
-            Shape is ``[number_of_active_heliostats, 4, 4]``.
+            The joint angles.
+            Shape is ``[number_of_active_heliostats, 2]``.
         """
-        # The surface points and normals are always sampled from a model (converted NURBS from deflectometry or ideal NURBS) that lays
-        # flat on the ground, i.e., the surface normals are pointing upwards [0.0, 0.0, 1.0]. Since ARTIST expects the points and normals
-        # to be initially oriented to the south, an extra rotation needs to be applied.
-        sampled_surface_model_orientation = torch.tensor(
-            [[0.0, 0.0, 1.0, 0.0]], device=device
-        ).expand(self.number_of_active_heliostats, 4)
-        east_angles, north_angles, up_angles = rotations.decompose_rotations(
-            initial_vector=sampled_surface_model_orientation,
-            target_vector=self.artist_standard_orientation,
+        device = get_device(device=device)
+
+        first_rotation_axis_deviations = transforms.rotate_n(
+            n=self.active_rotation_deviation_parameters[:, indices.first_joint_tilt_n],
+            device=device,
+        ) @ transforms.rotate_u(
+            u=self.active_rotation_deviation_parameters[:, indices.first_joint_tilt_u],
+            device=device,
         )
 
-        orientations_with_initial_orientation_offsets = (
-            orientations
-            @ transforms.rotate_e(e=east_angles, device=device)
-            @ transforms.rotate_n(n=north_angles, device=device)
-            @ transforms.rotate_u(u=up_angles, device=device)
+        second_rotation_axis_deviations = transforms.rotate_e(
+            e=self.active_rotation_deviation_parameters[:, indices.second_joint_tilt_e],
+            device=device,
+        ) @ transforms.rotate_n(
+            n=self.active_rotation_deviation_parameters[:, indices.second_joint_tilt_n],
+            device=device,
+        )
+
+        normal_after_first_deviation = (
+            first_rotation_axis_deviations.transpose(-1, -2) @ normals[:, :, None]
+        )
+
+        # Determine second joint angles.
+        second_axis_deviation_00 = second_rotation_axis_deviations[:, 0, 0]
+        second_axis_deviation_01 = second_rotation_axis_deviations[:, 0, 1]
+
+        denominator = torch.sqrt(
+            second_axis_deviation_00**2 + second_axis_deviation_01**2
+        )[:, None]
+        phi = torch.atan2(-second_axis_deviation_01, second_axis_deviation_00)[:, None]
+
+        ratio = torch.clamp(
+            (normal_after_first_deviation[:, 0] / (denominator + 1e-12)),
+            -1.0 + 1e-7,
+            1.0 - 1e-7,
+        )
+        second_joint_angle_1 = torch.arcsin(ratio) - phi
+        second_joint_angle_2 = torch.pi - torch.arcsin(ratio) - phi
+
+        second_joint_angle_1 = torch.atan2(
+            torch.sin(second_joint_angle_1), torch.cos(second_joint_angle_1)
+        )
+        second_joint_angle_2 = torch.atan2(
+            torch.sin(second_joint_angle_2), torch.cos(second_joint_angle_2)
+        )
+
+        # Determine first joint angles. There are two valid solutions.
+        v_1 = (
+            second_rotation_axis_deviations
+            @ transforms.rotate_u(second_joint_angle_1[:, 0], device=device)
+            @ self.kinematics_standard_orientation
+        )
+        first_joint_angle_1 = torch.atan2(
+            v_1[:, 1] * normal_after_first_deviation[:, 2, 0]
+            - v_1[:, 2] * normal_after_first_deviation[:, 1, 0],
+            v_1[:, 1] * normal_after_first_deviation[:, 1, 0]
+            + v_1[:, 2] * normal_after_first_deviation[:, 2, 0],
+        )
+        v_2 = (
+            second_rotation_axis_deviations
+            @ transforms.rotate_u(second_joint_angle_2[:, 0], device=device)
+            @ self.kinematics_standard_orientation
+        )
+        first_joint_angle_2 = torch.atan2(
+            v_2[:, 1] * normal_after_first_deviation[:, 2, 0]
+            - v_2[:, 2] * normal_after_first_deviation[:, 1, 0],
+            v_2[:, 1] * normal_after_first_deviation[:, 1, 0]
+            + v_2[:, 2] * normal_after_first_deviation[:, 2, 0],
+        )
+
+        first_joint_angle_1 = torch.atan2(
+            torch.sin(first_joint_angle_1), torch.cos(first_joint_angle_1)
+        )
+        first_joint_angle_2 = torch.atan2(
+            torch.sin(first_joint_angle_2), torch.cos(first_joint_angle_2)
+        )
+
+        # Determine possible motor positions.
+        motor_positions_1 = self.actuators.angles_to_motor_positions(
+            angles=torch.stack(
+                [first_joint_angle_1, second_joint_angle_1[:, 0]], dim=-1
+            ),
+            device=device,
+        )
+        motor_positions_2 = self.actuators.angles_to_motor_positions(
+            angles=torch.stack(
+                [first_joint_angle_2, second_joint_angle_2[:, 0]], dim=-1
+            ),
+            device=device,
+        )
+
+        # Determine valid solution. Prefer motor_positions_1 when valid, otherwise use motor_positions_2.
+        min_pos = self.actuators.active_non_optimizable_parameters[
+            :, indices.actuator_min_motor_position
+        ]
+        max_pos = self.actuators.active_non_optimizable_parameters[
+            :, indices.actuator_max_motor_position
+        ]
+        valid_1 = (motor_positions_1 >= min_pos) & (motor_positions_1 <= max_pos)
+        valid_2 = (motor_positions_2 >= min_pos) & (motor_positions_2 <= max_pos)
+
+        solution_1_valid = valid_1.all(dim=1)
+        solution_2_valid = valid_2.all(dim=1)
+
+        # Detect active heliostats where neither solution works.
+        invalid_rows = torch.nonzero(~(solution_1_valid | solution_2_valid))[:, 0]
+
+        if invalid_rows.numel() > 0:
+            log.warning(
+                f"No valid motor position combination for active heliostat number(s): {invalid_rows.tolist()}."
+            )
+
+        motor_positions = torch.where(
+            solution_1_valid[:, None],
+            motor_positions_1,
+            motor_positions_2,
+        )
+
+        return motor_positions
+
+    def motor_positions_to_orientations(
+        self, motor_positions: torch.Tensor, device: torch.device | None = None
+    ) -> torch.Tensor:
+        """
+        Compute orientation matrices given the motor positions.
+
+        Parameters
+        ----------
+        motor_positions : torch.Tensor
+            The motor positions.
+            Shape is ``[number_of_active_heliostats, 2]``.
+        device : torch.device | None
+            The device on which to perform computations or load tensors and models (default is None).
+            If None, ``ARTIST`` will automatically select the most appropriate
+            device (CUDA or CPU) based on availability and OS.
+
+        Returns
+        -------
+        torch.Tensor
+            The orientation matrices.
+            Shape is ``[number_of_active_heliostats, 4, 4]``.
+        """
+        device = get_device(device=device)
+
+        orientations = self._compute_orientations_from_motor_positions(
+            motor_positions=motor_positions, device=device
         )
 
-        return orientations_with_initial_orientation_offsets
+        return orientations @ self.initial_orientation_offsets
 
     def incident_ray_directions_to_orientations(
         self,
         incident_ray_directions: torch.Tensor,
         aim_points: torch.Tensor,
         device: torch.device | None = None,
-        max_num_iterations: int = 2,
+        max_num_iterations: int = 4,
         min_eps: float = 0.0001,
     ) -> torch.Tensor:
         """
@@ -352,14 +514,14 @@ def incident_ray_directions_to_orientations(
         aim_points : torch.Tensor
             The aim points for the active heliostats.
             Shape is ``[number_of_active_heliostats, 4]``.
-        max_num_iterations : int
-            Maximum number of iterations (default is 2).
-        min_eps : float
-            Convergence criterion (default is 0.0001).
         device : torch.device | None
             The device on which to perform computations or load tensors and models (default is None).
             If None, ``ARTIST`` will automatically select the most appropriate
             device (CUDA or CPU) based on availability and OS.
+        max_num_iterations : int
+            Maximum number of iterations (default is 2).
+        min_eps : float
+            Convergence criterion (default is 0.0001).
 
         Returns
         -------
@@ -380,27 +542,42 @@ def incident_ray_directions_to_orientations(
 
         for _ in range(max_num_iterations):
             orientations = self._compute_orientations_from_motor_positions(
-                motor_positions=motor_positions, device=device
+                motor_positions=motor_positions,
+                device=device,
             )
 
             concentrator_normals = orientations @ torch.tensor(
-                [0, -1, 0, 0], dtype=torch.float32, device=device
+                [0.0, -1.0, 0.0, 0.0],
+                device=device,
             )
+
             concentrator_origins = orientations @ torch.tensor(
-                [0, 0, 0, 1], dtype=torch.float32, device=device
+                [0.0, 0.0, 0.0, 1.0],
+                device=device,
             )
 
-            # Compute desired normals.
             desired_reflection_directions = torch.nn.functional.normalize(
-                aim_points - concentrator_origins, p=2, dim=1
+                aim_points[:, :3] - concentrator_origins[:, :3],
+                p=2,
+                dim=1,
+                eps=1e-8,
             )
+
             desired_concentrator_normals = torch.nn.functional.normalize(
-                -incident_ray_directions + desired_reflection_directions, p=2, dim=1
+                -incident_ray_directions[:, :3] + desired_reflection_directions,
+                p=2,
+                dim=1,
+                eps=1e-8,
+            )
+
+            desired_concentrator_normals = (
+                coordinates.convert_3d_directions_to_4d_format(
+                    desired_concentrator_normals, device=device
+                )
             )
 
-            # Compute loss.
             loss = torch.abs(desired_concentrator_normals - concentrator_normals).mean(
-                dim=0
+                dim=-1
             )
 
             # Stop if converged.
@@ -410,111 +587,10 @@ def incident_ray_directions_to_orientations(
                     break
             last_iteration_loss = loss
 
-            # Analytical solution for joint angles.
-            joint_angles_1 = -torch.arcsin(
-                torch.clamp(
-                    -desired_concentrator_normals[:, indices.e]
-                    / torch.cos(
-                        self.active_translation_deviation_parameters[
-                            :, indices.second_joint_translation_n
-                        ]
-                    ),
-                    min=-1,
-                    max=1,
-                )
-            )
-
-            a = -torch.cos(
-                self.active_translation_deviation_parameters[
-                    :, indices.second_joint_translation_e
-                ]
-            ) * torch.cos(joint_angles_1) + torch.sin(
-                self.active_translation_deviation_parameters[
-                    :, indices.second_joint_translation_e
-                ]
-            ) * torch.sin(
-                self.active_translation_deviation_parameters[
-                    :, indices.second_joint_translation_n
-                ]
-            ) * torch.sin(joint_angles_1)
-            b = -torch.sin(
-                self.active_translation_deviation_parameters[
-                    :, indices.second_joint_translation_e
-                ]
-            ) * torch.cos(joint_angles_1) - torch.cos(
-                self.active_translation_deviation_parameters[
-                    :, indices.second_joint_translation_e
-                ]
-            ) * torch.sin(
-                self.active_translation_deviation_parameters[
-                    :, indices.second_joint_translation_n
-                ]
-            ) * torch.sin(joint_angles_1)
-
-            joint_angles_0 = (
-                torch.arctan2(
-                    a * -desired_concentrator_normals[:, indices.u]
-                    - b * -desired_concentrator_normals[:, indices.n],
-                    a * -desired_concentrator_normals[:, indices.n]
-                    + b * -desired_concentrator_normals[:, indices.u],
-                )
-                - torch.pi
-            )
-
-            joint_angles = torch.stack(
-                [
-                    joint_angles_0,
-                    joint_angles_1,
-                ],
-                dim=1,
-            )
-
-            motor_positions = self.actuators.angles_to_motor_positions(
-                angles=joint_angles, device=device
+            motor_positions = self._compute_motor_positions_from_normal(
+                normals=desired_concentrator_normals, device=device
             )
 
-        orientations_with_initial_orientation_offsets = (
-            self._apply_initial_orientation_offsets(
-                orientations=orientations, device=device
-            )
-        )
-
         self.active_motor_positions = motor_positions
 
-        return orientations_with_initial_orientation_offsets
-
-    def motor_positions_to_orientations(
-        self, motor_positions: torch.Tensor, device: torch.device | None = None
-    ) -> torch.Tensor:
-        """
-        Compute orientation matrices given the motor positions.
-
-        Parameters
-        ----------
-        motor_positions : torch.Tensor
-            The motor positions.
-            Shape is ``[number_of_active_heliostats, 2]``.
-        device : torch.device | None
-            The device on which to perform computations or load tensors and models (default is None).
-            If None, ``ARTIST`` will automatically select the most appropriate
-            device (CUDA or CPU) based on availability and OS.
-
-        Returns
-        -------
-        torch.Tensor
-            The orientation matrices.
-            Shape is ``[number_of_active_heliostats, 4, 4]``.
-        """
-        device = get_device(device=device)
-
-        orientations = self._compute_orientations_from_motor_positions(
-            motor_positions=motor_positions, device=device
-        )
-
-        orientations_with_initial_orientation_offsets = (
-            self._apply_initial_orientation_offsets(
-                orientations=orientations, device=device
-            )
-        )
-
-        return orientations_with_initial_orientation_offsets
+        return orientations @ self.initial_orientation_offsets
diff --git a/artist/optim/__init__.py b/artist/optim/__init__.py
index 2b42fefd5..3c3825187 100644
--- a/artist/optim/__init__.py
+++ b/artist/optim/__init__.py
@@ -1,3 +1,4 @@
+from .aim_point_optimizer import AimPointOptimizer
 from .kinematics_reconstructor import KinematicsReconstructor
 from .loss import (
     AngleLoss,
@@ -6,22 +7,22 @@
     Loss,
     PixelLoss,
     VectorLoss,
-    mean_loss_per_heliostat,
+    reduce_loss_per_sample,
 )
-from .motor_position_optimizer import MotorPositionsOptimizer
 from .regularizers import IdealSurfaceRegularizer, SmoothnessRegularizer
 from .surface_reconstructor import SurfaceReconstructor
 from .training import EarlyStopping, cyclic, exponential, reduce_on_plateau
 
 __all__ = [
     "KinematicsReconstructor",
-    "MotorPositionsOptimizer",
+    "AimPointOptimizer",
     "SurfaceReconstructor",
     "SmoothnessRegularizer",
     "IdealSurfaceRegularizer",
     "EarlyStopping",
     "Loss",
     "VectorLoss",
+    "reduce_loss_per_sample",
     "FocalSpotLoss",
     "PixelLoss",
     "KLDivergenceLoss",
diff --git a/artist/optim/motor_position_optimizer.py b/artist/optim/aim_point_optimizer.py
similarity index 93%
rename from artist/optim/motor_position_optimizer.py
rename to artist/optim/aim_point_optimizer.py
index b366c37c8..28b570d61 100644
--- a/artist/optim/motor_position_optimizer.py
+++ b/artist/optim/aim_point_optimizer.py
@@ -13,12 +13,12 @@
 from artist.util.env import DdpSetup, get_device
 
 log = logging.getLogger(__name__)
-"""A logger for the motor positions optimizer."""
+"""A logger for the aim point optimizer."""
 
 
-class MotorPositionsOptimizer:
+class AimPointOptimizer:
     """
-    An optimizer used to find optimal motor positions for the heliostats.
+    An optimizer used to find optimal aim points via individual motor positions for the heliostats.
 
     The optimization loss is defined as the loss between the combined predicted and target
     flux densities. Additionally, there is one constraint that maximizes the flux integral,
@@ -74,7 +74,7 @@ def __init__(
         device: torch.device | None = None,
     ) -> None:
         """
-        Initialize the motor-positions optimizer.
+        Initialize the aim points optimizer.
 
         Parameters
         ----------
@@ -110,7 +110,7 @@ def __init__(
         rank = ddp_setup["rank"]
 
         if rank == 0:
-            log.info("Create a motor positions optimizer.")
+            log.info("Create an aim points optimizer.")
 
         self.ddp_setup = ddp_setup
         self.scenario = scenario
@@ -130,7 +130,7 @@ def optimize(
         device: torch.device | None = None,
     ) -> tuple[torch.Tensor, dict[str, list], torch.Tensor, torch.Tensor, torch.Tensor]:
         r"""
-        Optimize the motor positions.
+        Optimize the motor positions for optimal aim points.
 
         The motor positions are optimized through a reparameterization to ensure stable training
         across different heliostats with widely varying initial motor positions and ranges. Motor
@@ -174,7 +174,7 @@ def optimize(
         Returns
         -------
         torch.Tensor
-            Final loss of the motor position optimization.
+            Final loss of the aim point optimization.
         dict[str, list]
             Loss history over epochs, with keys ``"total_loss"``, ``"flux_loss"``,
             ``"local_flux_constraint"``, ``"intercept_constraint"``, ``"flux_integral_constraint"``,
@@ -190,7 +190,7 @@ def optimize(
         rank = self.ddp_setup["rank"]
 
         if rank == 0:
-            log.info("Start the motor positions optimization.")
+            log.info("Start the aim point optimization.")
 
         optimizable_parameters_all_groups = []
         scales_all_groups = []
@@ -210,11 +210,10 @@ def optimize(
             ]
         )
 
-        # Per-group pre-alignment + reparameterization prep
+        # Per-group pre-alignment and parameterization
         for group_index, group in enumerate(
             self.scenario.heliostat_field.heliostat_groups
         ):
-            # All heliostats are initially active.
             active_heliostats_masks_all_groups.append(
                 torch.ones(
                     group.number_of_heliostats,
@@ -222,7 +221,6 @@ def optimize(
                     device=device,
                 )
             )
-            # All target indices are initially set to selected target area.
             target_area_indices_all_groups.append(
                 torch.full(
                     (group.number_of_heliostats,),
@@ -231,20 +229,15 @@ def optimize(
                     device=device,
                 )
             )
-            # Repeat the same incident ray direction per heliostat.
             incident_ray_directions_all_groups.append(
                 self.incident_ray_direction.repeat(group.number_of_heliostats, 1)
             )
 
             # Align all heliostats once, to the given incident ray direction and target, to set initial motor positions.
-            # The motor positions are set automatically in the ``align_surfaces_with_incident_ray_directions()`` method.
-            # Activate heliostats.
             group.activate_heliostats(
                 active_heliostats_mask=active_heliostats_masks_all_groups[group_index],
                 device=device,
             )
-
-            # Align heliostats.
             group.align_surfaces_with_incident_ray_directions(
                 aim_points=self.scenario.solar_tower.get_centers_of_target_areas(
                     target_area_indices=target_area_indices_all_groups[group_index],
@@ -260,7 +253,6 @@ def optimize(
                 group.kinematics.active_motor_positions.detach().clone()
             )
             initial_motor_positions_all_groups.append(initial_motor_positions)
-            # Read actuator min/max positions.
             motor_positions_minimum = (
                 group.kinematics.actuators.non_optimizable_parameters[
                     :, indices.actuator_min_motor_position
@@ -277,7 +269,6 @@ def optimize(
                 torch.minimum(lower_margin, upper_margin).clamp(min=1.0)
             )
 
-            # Create trainable unconstrained tensor.
             optimizable_parameters_all_groups.append(
                 torch.nn.Parameter(
                     torch.zeros_like(initial_motor_positions, device=device)
@@ -307,7 +298,6 @@ def optimize(
             relative=True,
         )
 
-        # Determine target-plane dimensions.
         target_plane_dimensions = torch.empty(2, device=device)
         target_areas, index = self.scenario.solar_tower.index_to_target_area[
             self.target_area_index
@@ -317,9 +307,9 @@ def optimize(
             < self.scenario.solar_tower.number_of_target_areas_per_type[
                 indices.planar_target_areas
             ]
-        ):  # Planar target
+        ):
             target_plane_dimensions = target_areas.dimensions[self.target_area_index]  # type: ignore[attr-defined]
-        else:  # Cylinder target
+        else:
             cylinder_indices = (
                 self.target_area_index
                 - self.scenario.solar_tower.number_of_target_areas_per_type[
@@ -383,7 +373,7 @@ def optimize(
             on_target_factors = torch.zeros_like(intercept_factors, device=device)
             blocking_factors = torch.zeros_like(intercept_factors, device=device)
 
-            # First per-group loop: Reconstruct and apply motor positions.
+            # First per-group loop: Align all heliostats such that blocking can be computed correctly.
             for heliostat_group_index in self.ddp_setup["groups_to_ranks_mapping"][
                 rank
             ]:
@@ -405,7 +395,6 @@ def optimize(
                     * scales_all_groups[heliostat_group_index]
                 )
 
-                # Activate heliostats.
                 heliostat_alignment_group.activate_heliostats(
                     active_heliostats_mask=active_heliostats_masks_all_groups[
                         heliostat_group_index
@@ -422,7 +411,7 @@ def optimize(
                     device=device,
                 )
 
-            # Second per-group loop: Trace rays and accumulate outputs.
+            # Second per-group loop: Trace rays and accumulate fluxes.
             for heliostat_group_index in self.ddp_setup["groups_to_ranks_mapping"][
                 rank
             ]:
@@ -464,7 +453,6 @@ def optimize(
                     device=device,
                 )
                 sample_indices_for_local_rank = ray_tracer.get_sampler_indices()
-                # Retrieve flux on selected target.
                 flux_distribution_on_target = ray_tracer.get_bitmaps_per_target(
                     bitmaps_per_heliostat=flux_distributions,
                     target_area_indices=target_area_indices_all_groups[
@@ -472,7 +460,6 @@ def optimize(
                     ][sample_indices_for_local_rank],
                     device=device,
                 )[self.target_area_index]
-                # Add to global flux.
                 total_flux = total_flux + flux_distribution_on_target
 
                 global_indices = (
@@ -482,7 +469,6 @@ def optimize(
                 on_target_factors[global_indices] = on_target_factor
                 blocking_factors[global_indices] = blocking_factor
 
-            # Reduce total flux in the distributed case.
             if self.ddp_setup["is_distributed"]:
                 total_flux = torch.distributed.nn.functional.all_reduce(
                     total_flux,
@@ -503,12 +489,10 @@ def optimize(
                 device=device,
             )
 
-            if isinstance(loss_definition, FocalSpotLoss):  # Focal-spot loss
+            if isinstance(loss_definition, FocalSpotLoss):
                 loss = flux_loss
 
-            if isinstance(
-                loss_definition, KLDivergenceLoss
-            ):  # Kullback-Leibler divergence loss
+            if isinstance(loss_definition, KLDivergenceLoss):
                 # Store references at epoch 0, i.e., baseline flux integral and intercept factors.
                 if epoch == 0:
                     flux_integral_reference = total_flux.sum().detach()
@@ -590,7 +574,6 @@ def optimize(
                             param.grad /= self.ddp_setup[constants.world_size]  # type: ignore
 
             optimizer.step()
-            # Scheduler update.
             if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                 scheduler.step(loss.detach())
             else:
@@ -641,7 +624,7 @@ def optimize(
             "flux_integral_constraint": flux_integral_constraint_history,
             "flux_integral": flux_integral,
         }
-        log.info(f"Rank: {rank}, motor positions optimized.")
+        log.info(f"Rank: {rank}, aim points optimized.")
 
         # Broadcast final motor positions for each heliostat group from source rank to others.
         if self.ddp_setup["is_distributed"]:
@@ -654,7 +637,7 @@ def optimize(
                     src=source[indices.first_rank_from_group],
                 )
 
-            log.info(f"Rank: {rank}, synchronized after motor positions optimization.")
+            log.info(f"Rank: {rank}, synchronized after aim point optimization.")
 
         return (
             loss.detach().cpu(),
diff --git a/artist/optim/kinematics_reconstructor.py b/artist/optim/kinematics_reconstructor.py
index 0b105663a..fef9ed992 100644
--- a/artist/optim/kinematics_reconstructor.py
+++ b/artist/optim/kinematics_reconstructor.py
@@ -1,8 +1,10 @@
 import logging
 import pathlib
-from typing import Any, cast
+from functools import partial
+from typing import Any, Callable, cast
 
 import torch
+from matplotlib import pyplot as plt
 from torch.optim.lr_scheduler import LRScheduler
 
 from artist.field.heliostat_group import HeliostatGroup
@@ -10,7 +12,13 @@
 from artist.geometry import coordinates
 from artist.io.calibration_parser import CalibrationDataParser
 from artist.optim import training
-from artist.optim.loss import Loss, mean_loss_per_heliostat
+from artist.optim.loss import (
+    FocalSpotLoss,
+    KLDivergenceLoss,
+    Loss,
+    PixelLoss,
+    reduce_loss_per_sample,
+)
 from artist.raytracing.heliostat_ray_tracer import HeliostatRayTracer
 from artist.scenario.scenario import Scenario
 from artist.util import constants, indices
@@ -46,6 +54,14 @@ class KinematicsReconstructor:
         Direct normal irradiance in W/m^2.
     reconstruction_method : str
         The reconstruction method. Currently, only reconstruction via ray tracing is implemented.
+    validation_loss_focal_spot : FocalSpotLoss
+        Flux loss used for validation.
+    validation_loss_pixel : PixelLoss
+        Pixel loss used for validation.
+    validation_loss_kl_div : KLDivergenceLoss
+        Kullback-Leibler divergence loss used for validation.
+    plot_results : bool
+        Create flux plots in the last epoch of training for each heliostat and all its samples (slow!).
 
     Note
     ----
@@ -70,6 +86,7 @@ def __init__(
         dni: float | None = None,
         reconstruction_method: str = constants.kinematics_reconstruction_raytracing,
         bitmap_resolution: torch.Tensor = torch.tensor([256, 256]),
+        plot_results: bool = False,
     ) -> None:
         """
         Initialize the kinematics optimizer.
@@ -92,6 +109,8 @@ def __init__(
         bitmap_resolution : torch.Tensor
             The resolution of all bitmaps during reconstruction (default is ``torch.tensor([256, 256])``).
             Shape is ``[2]``.
+        plot_results : bool
+            Create flux plots in the last epoch of training for each heliostat and all its samples (slow!, default is ``False``).
         """
         rank = ddp_setup["rank"]
         if rank == 0:
@@ -105,8 +124,13 @@ def __init__(
         self.dni = dni
         self.bitmap_resolution = bitmap_resolution
 
+        self.validation_loss_focal_spot = FocalSpotLoss(scenario=self.scenario)
+        self.validation_loss_pixel = PixelLoss()
+        self.validation_loss_kl_div = KLDivergenceLoss()
+
         if reconstruction_method in [
             constants.kinematics_reconstruction_raytracing,
+            constants.kinematics_reconstruction_alignment,
         ]:
             self.reconstruction_method = reconstruction_method
         else:
@@ -115,6 +139,8 @@ def __init__(
                 f"Please select another reconstruction method and try again!"
             )
 
+        self.plot_results = plot_results
+
     def reconstruct_kinematics(
         self,
         loss_definition: Loss,
@@ -154,6 +180,13 @@ def reconstruct_kinematics(
                     device=device,
                 )
             )
+        elif (
+            self.reconstruction_method == constants.kinematics_reconstruction_alignment
+        ):
+            loss, loss_history = self._reconstruct_kinematics_parameters_wih_alignment(
+                loss_definition=loss_definition,
+                device=device,
+            )
 
         else:
             raise ValueError(
@@ -163,16 +196,230 @@ def reconstruct_kinematics(
 
         return loss, loss_history
 
-    def _reconstruct_kinematics_parameters_with_raytracing(
+    def _validate(
+        self,
+        heliostat_group: HeliostatGroup,
+        data_split: training.TrainTestSplit,
+        reduction: Callable[..., Any],
+        device: torch.device | None = None,
+    ) -> torch.Tensor:
+        """
+        Validate the kinematic reconstruction for a specified heliostat group on the test data.
+
+        Parameters
+        ----------
+        heliostat_group : HeliostatGroup
+            Heliostat group to validate.
+        data_split : training.TrainTestSplit
+            Train/test split containing all test tensors and metadata.
+        reduction : Callable[..., Any]
+            Reduction function applied across the sample dimension for each heliostat.
+        device : torch.device | None
+            The device on which to perform computations or load tensors and models (default is None).
+            If None, ARTIST will automatically select the most appropriate
+            device (CUDA or CPU) based on availability and OS.
+
+        Returns
+        -------
+        torch.Tensor
+            Predicted flux distributions for the local validation samples.
+            Shape is ``[number_of_local_test_samples, height, width]``.
+        """
+        device = get_device(device=device)
+
+        heliostat_group.activate_heliostats(
+            active_heliostats_mask=data_split.active_heliostats_mask_test,
+            device=device,
+        )
+
+        heliostat_group.align_surfaces_with_motor_positions(
+            motor_positions=data_split.motor_positions_test,
+            active_heliostats_mask=data_split.active_heliostats_mask_test,
+            device=device,
+        )
+
+        ray_tracer = HeliostatRayTracer(
+            scenario=self.scenario,
+            heliostat_group=heliostat_group,
+            blocking_active=False,
+            world_size=self.ddp_setup["heliostat_group_world_size"],
+            rank=self.ddp_setup["heliostat_group_rank"],
+            batch_size=self.optimizer_dict[constants.batch_size],
+            random_seed=self.ddp_setup["heliostat_group_rank"],
+            dni=self.dni,
+            bitmap_resolution=self.bitmap_resolution,
+        )
+
+        flux_prediction, _, _, _ = ray_tracer.trace_rays(
+            incident_ray_directions=data_split.incident_ray_directions_test,
+            active_heliostats_mask=data_split.active_heliostats_mask_test,
+            target_area_indices=data_split.target_area_indices_test,
+            device=device,
+        )
+
+        indices_for_local_rank = ray_tracer.get_sampler_indices()
+
+        loss_focal_spot_per_sample = self.validation_loss_focal_spot(
+            prediction=flux_prediction,
+            ground_truth=data_split.flux_measured_test[indices_for_local_rank],
+            target_area_indices=data_split.target_area_indices_test[
+                indices_for_local_rank
+            ],
+            device=device,
+        )
+        loss_pixel_per_sample = self.validation_loss_pixel(
+            prediction=flux_prediction,
+            ground_truth=data_split.flux_measured_test[indices_for_local_rank],
+            reduction_dimensions=(
+                1,
+                2,
+            ),
+        )
+        loss_kl_div_per_sample = self.validation_loss_kl_div(
+            prediction=flux_prediction,
+            ground_truth=data_split.flux_measured_test[indices_for_local_rank],
+            reduction_dimensions=(
+                1,
+                2,
+            ),
+        )
+
+        test_loss_focal_spot = reduce_loss_per_sample(
+            loss_per_sample=loss_focal_spot_per_sample,
+            number_of_samples_per_heliostat=data_split.number_of_test_samples,
+            reduction=reduction,
+        )
+        test_loss_pixel = reduce_loss_per_sample(
+            loss_per_sample=loss_pixel_per_sample,
+            number_of_samples_per_heliostat=data_split.number_of_test_samples,
+            reduction=reduction,
+        )
+        test_loss_kl_div = reduce_loss_per_sample(
+            loss_per_sample=loss_kl_div_per_sample,
+            number_of_samples_per_heliostat=data_split.number_of_test_samples,
+            reduction=reduction,
+        )
+
+        log.info(
+            "test loss focal spot mean: %.5f, pixel mean: %.5f, kl div mean: %.5f",
+            torch.mean(test_loss_focal_spot).item(),
+            torch.mean(test_loss_pixel).item(),
+            torch.mean(test_loss_kl_div).item(),
+        )
+
+        return flux_prediction
+
+    def _plot_fluxes(
+        self,
+        flux_measured: torch.Tensor,
+        flux_prediction_train: torch.Tensor,
+        flux_prediction_test: torch.Tensor,
+        data_split: training.TrainTestSplit,
+        plot_name: str,
+    ) -> None:
+        """
+        Plot predicted and measured flux maps for each heliostat sample.
+
+        Each row in the generated figure corresponds to one sample of a
+        heliostat, where the left column contains the predicted flux and
+        the right column contains the measured flux.
+        The subplot borders are color-coded to indicate whether a sample
+        belongs to the training samples (green) or testing samples (red)
+        One figure is generated per heliostat.
+
+        Parameters
+        ----------
+        flux_measured : torch.Tensor
+            Ground-truth measured flux maps.
+        flux_prediction_train : torch.Tensor
+            Predicted flux maps of the training samples.
+        flux_prediction_test : torch.Tensor
+            Predicted flux maps of the testing samples.
+        data_split : training.TrainTestSplit
+            Information on the train/test split.
+        plot_name : str
+            Name suffix used when saving the generated plot files.
+        """
+        device = torch.device("cpu")
+
+        flux_predicted = torch.zeros_like(flux_measured, device=device)
+        flux_predicted[data_split.train_indices] = flux_prediction_train
+        flux_predicted[data_split.test_indices] = flux_prediction_test
+        samples_per_heliostat = data_split.number_of_samples_per_heliostat
+        total_samples = flux_measured.shape[0]
+        train_indices = set(data_split.train_indices.tolist())
+        test_indices = set(data_split.test_indices.tolist())
+
+        centers_of_mass_predicted = bitmap.get_center_of_mass(
+            flux_predicted, device=device
+        )
+        centers_of_mass_measured = bitmap.get_center_of_mass(
+            flux_measured, device=device
+        )
+
+        for heliostat_start_index in range(0, total_samples, samples_per_heliostat):
+            fig, axes = plt.subplots(
+                samples_per_heliostat,
+                2,
+                figsize=(8, samples_per_heliostat * 4),
+            )
+            heliostat_index = heliostat_start_index // samples_per_heliostat
+
+            for sample_offset in range(samples_per_heliostat):
+                sample_index = heliostat_start_index + sample_offset
+
+                if sample_index in train_indices:
+                    border_color = "green"
+                    split_name = "TRAIN"
+                elif sample_index in test_indices:
+                    border_color = "red"
+                    split_name = "TEST"
+
+                axes[sample_offset, 0].imshow(flux_predicted[sample_index])
+                axes[sample_offset, 0].set_title(
+                    f"Predicted Flux - Heliostat {heliostat_index} ({split_name})"
+                )
+
+                axes[sample_offset, 1].imshow(flux_measured[sample_index])
+                axes[sample_offset, 1].set_title(
+                    f"Measured Flux - Heliostat {heliostat_index} ({split_name})"
+                )
+
+                for ax in axes[sample_offset, :]:
+                    ax.scatter(
+                        centers_of_mass_measured[sample_index, 0],
+                        centers_of_mass_measured[sample_index, 1],
+                        c="black",
+                        s=30,
+                        marker="o",
+                    )
+                    ax.scatter(
+                        centers_of_mass_predicted[sample_index, 0],
+                        centers_of_mass_predicted[sample_index, 1],
+                        c="red",
+                        s=30,
+                        marker="x",
+                    )
+
+                for spine in axes[sample_offset, :].flat:
+                    for spines in spine.spines.values():
+                        spines.set_edgecolor(border_color)
+                        spines.set_linewidth(4)
+
+            plt.tight_layout()
+            plt.savefig(f"heliostat_{heliostat_index}_{plot_name}")
+            plt.close(fig)
+
+    def _reconstruct_kinematics_parameters_wih_alignment(
         self,
         loss_definition: Loss,
         device: torch.device | None = None,
     ) -> tuple[torch.Tensor, list[list[dict[str, list[float]]]]]:
         """
-        Reconstruct the kinematics parameters using ray tracing.
+        Reconstruct the kinematics parameters using alignment and geometry data.
 
-        This reconstruction method optimizes the kinematics parameters by extracting the focal points
-        of calibration images and using heliostat-tracing.
+        This reconstruction method optimizes the kinematics parameters by iteratively
+        aligning heliostats to reach a defined flux focal spot.
 
         Parameters
         ----------
@@ -221,11 +468,10 @@ def _reconstruct_kinematics_parameters_with_raytracing(
 
         # Iterate heliostat groups assigned to this rank.
         for heliostat_group_index in self.ddp_setup["groups_to_ranks_mapping"][rank]:
-            # Parse calibration inputs for current group to obtain measured flux, incident ray directions, mask of
-            # active heliostats, and target area indices.
             heliostat_group: HeliostatGroup = (
                 self.scenario.heliostat_field.heliostat_groups[heliostat_group_index]
             )
+            # Load data parser and input file mapping, then parse the calibration data.
             parser = cast(CalibrationDataParser, self.data[constants.data_parser])
             heliostat_mapping = cast(
                 list[tuple[str, list[pathlib.Path], list[pathlib.Path]]],
@@ -233,9 +479,9 @@ def _reconstruct_kinematics_parameters_with_raytracing(
             )
             (
                 flux_measured,
-                _,
+                focal_spots_measured,
                 incident_ray_directions,
-                _,
+                motor_positions,
                 active_heliostats_mask,
                 target_area_indices,
             ) = parser.parse_data_for_reconstruction(
@@ -246,68 +492,357 @@ def _reconstruct_kinematics_parameters_with_raytracing(
                 device=device,
             )
 
+            # Skip groups with no active heliostats.
             if active_heliostats_mask.sum() > 0:
-                # Calculate focal spot from measured flux.
-                focal_spots_bitmap_coordinates = bitmap.get_center_of_mass(
-                    bitmaps=flux_measured, device=device
+                data_split: training.TrainTestSplit = training.train_test_split(
+                    active_heliostats_mask=active_heliostats_mask,
+                    flux_measured=flux_measured,
+                    focal_spots_measured=focal_spots_measured,
+                    incident_ray_directions=incident_ray_directions,
+                    motor_positions=motor_positions,
+                    target_area_indices=target_area_indices,
+                    test_fraction=0.25,
+                    device=device,
                 )
-                focal_spots_measured = (
-                    coordinates.bitmap_coordinates_to_target_coordinates(
-                        bitmap_coordinates=focal_spots_bitmap_coordinates,
-                        bitmap_resolution=self.bitmap_resolution,
-                        solar_tower=self.scenario.solar_tower,
-                        target_area_indices=target_area_indices,
-                        device=device,
+                # Calculate focal spot from measured flux.
+                preferred_reflection_directions_measured = (
+                    torch.nn.functional.normalize(
+                        (
+                            focal_spots_measured[:, :3]
+                            - heliostat_group.positions.repeat_interleave(
+                                active_heliostats_mask, dim=0
+                            )[:, :3]
+                        ),
+                        p=2,
+                        dim=1,
                     )
                 )
+                normals_measured = coordinates.convert_3d_directions_to_4d_format(
+                    torch.nn.functional.normalize(
+                        preferred_reflection_directions_measured
+                        - incident_ray_directions[:, :3],
+                        dim=-1,
+                    ),
+                    device=device,
+                )
+
+                # Set up optimizer, scheduler, and early stopping.
+                optimizer_params = [
+                    {
+                        "params": heliostat_group.kinematics.rotation_deviation_parameters.requires_grad_(),
+                        "lr": self.optimizer_dict[
+                            constants.initial_learning_rate_rotation_deviation
+                        ],
+                    }
+                ]
+
+                optimizer = torch.optim.Adam(optimizer_params)
+
+                # Create a learning rate scheduler.
+                scheduler_fn = getattr(
+                    training,
+                    self.scheduler_dict[constants.scheduler_type],
+                )
+                scheduler: LRScheduler = scheduler_fn(
+                    optimizer=optimizer, parameters=self.scheduler_dict
+                )
+
+                # Set up early stopping.
+                early_stopper = training.EarlyStopping(
+                    window_size=self.optimizer_dict[constants.early_stopping_window],
+                    patience=self.optimizer_dict[constants.early_stopping_patience],
+                    min_improvement=self.optimizer_dict[constants.early_stopping_delta],
+                    relative=True,
+                )
 
-                # Reparametrize optimizable actuator parameters.
-                initial_actuator_params = (
-                    heliostat_group.kinematics.actuators.optimizable_parameters.detach()
+                loss_history_list = []
+
+                # Start the optimization.
+                loss = torch.inf
+                epoch = 0
+                log_step = (
+                    self.optimizer_dict[constants.max_epoch]
+                    if self.optimizer_dict[constants.log_step] == 0
+                    else self.optimizer_dict[constants.log_step]
                 )
-                if initial_actuator_params is not None:
-                    angle_mean = initial_actuator_params[
-                        :, indices.actuator_params_initial_angle
-                    ].mean()
-                    angle_std = (
-                        initial_actuator_params[
-                            :, indices.actuator_params_initial_angle
-                        ]
-                        .std()
-                        .clamp(min=1e-3)
+                while (
+                    loss > float(self.optimizer_dict[constants.tolerance])
+                    and epoch <= self.optimizer_dict[constants.max_epoch]
+                ):
+                    optimizer.zero_grad()
+
+                    # Activate heliostats.
+                    heliostat_group.activate_heliostats(
+                        active_heliostats_mask=data_split.active_heliostats_mask_train,
+                        device=device,
+                    )
+
+                    orientations = (
+                        heliostat_group.kinematics.motor_positions_to_orientations(
+                            motor_positions=data_split.motor_positions_train,
+                            device=device,
+                        )
                     )
 
-                    stroke_mean = initial_actuator_params[
-                        :, indices.actuator_params_initial_stroke_length
-                    ].mean()
-                    stroke_std = (
-                        initial_actuator_params[
-                            :, indices.actuator_params_initial_stroke_length
-                        ]
-                        .std()
-                        .clamp(min=1e-3)
+                    normals_predicted = orientations @ torch.tensor(
+                        [0.0, 0.0, 1.0, 0.0], device=device
                     )
 
-                    angle_normalized = (
-                        initial_actuator_params[
-                            :, indices.actuator_params_initial_angle
-                        ]
-                        - angle_mean
-                    ) / angle_std
-                    stroke_length_normalized = (
-                        initial_actuator_params[
-                            :, indices.actuator_params_initial_stroke_length
-                        ]
-                        - stroke_mean
-                    ) / stroke_std
-
-                    delta_angle = torch.zeros_like(angle_normalized, requires_grad=True)
-                    delta_stroke = torch.zeros_like(
-                        stroke_length_normalized, requires_grad=True
+                    # Compute loss from prediction vs. measured normals.
+                    loss_per_sample = loss_definition(
+                        prediction=normals_predicted,
+                        ground_truth=normals_measured[data_split.train_indices],
+                    )
+
+                    loss_per_heliostat = reduce_loss_per_sample(
+                        loss_per_sample=loss_per_sample,
+                        number_of_samples_per_heliostat=data_split.number_of_train_samples,
+                        reduction=partial(torch.mean),
+                    )
+
+                    loss = torch.mean(loss_per_heliostat)
+
+                    loss.backward()
+
+                    # Severely misaligned heliostat samples produce nan gradients. These are set to zero and the
+                    # sample has no influence on the training epoch.
+                    optimizer.param_groups[0]["params"][0].grad.nan_to_num_(
+                        nan=0.0, posinf=0.0, neginf=0.0
+                    )
+
+                    if self.ddp_setup["is_nested"]:
+                        # Reduce gradients within each heliostat group.
+                        for param_group in optimizer.param_groups:
+                            for param in param_group["params"]:
+                                if param.grad is not None:
+                                    param.grad = (
+                                        torch.distributed.nn.functional.all_reduce(
+                                            param.grad,
+                                            op=torch.distributed.ReduceOp.SUM,
+                                            group=self.ddp_setup["process_subgroup"],
+                                        )
+                                    )
+                                    param.grad /= self.ddp_setup[
+                                        "heliostat_group_world_size"
+                                    ]
+
+                    optimizer.step()
+                    if isinstance(
+                        scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau
+                    ):
+                        scheduler.step(loss.detach())
+                    else:
+                        scheduler.step()
+
+                    is_last_epoch = (
+                        epoch == self.optimizer_dict[constants.max_epoch] - 1
                     )
-                else:
-                    delta_angle = None
-                    delta_stroke = None
+                    stop = early_stopper.step(loss.item())
+
+                    if epoch % log_step == 0 or is_last_epoch or stop:
+                        log.info(
+                            f"Rank: {rank}, Epoch: {epoch}, Loss: {loss}",
+                        )
+
+                        with torch.no_grad():
+                            heliostat_group.activate_heliostats(
+                                active_heliostats_mask=data_split.active_heliostats_mask_train,
+                                device=device,
+                            )
+                            heliostat_group.align_surfaces_with_motor_positions(
+                                motor_positions=data_split.motor_positions_train,
+                                active_heliostats_mask=data_split.active_heliostats_mask_train,
+                                device=device,
+                            )
+                            ray_tracer = HeliostatRayTracer(
+                                scenario=self.scenario,
+                                heliostat_group=heliostat_group,
+                                blocking_active=False,
+                                world_size=self.ddp_setup["heliostat_group_world_size"],
+                                rank=self.ddp_setup["heliostat_group_rank"],
+                                batch_size=self.optimizer_dict[constants.batch_size],
+                                random_seed=self.ddp_setup["heliostat_group_rank"],
+                                dni=self.dni,
+                                bitmap_resolution=self.bitmap_resolution,
+                            )
+                            flux_prediction_train, _, _, _ = ray_tracer.trace_rays(
+                                incident_ray_directions=data_split.incident_ray_directions_train,
+                                active_heliostats_mask=data_split.active_heliostats_mask_train,
+                                target_area_indices=data_split.target_area_indices_train,
+                                device=device,
+                            )
+                            flux_prediction_test = self._validate(
+                                heliostat_group=heliostat_group,
+                                data_split=data_split,
+                                reduction=partial(torch.mean),
+                                device=device,
+                            )
+
+                    if self.plot_results and (is_last_epoch or stop):
+                        self._plot_fluxes(
+                            flux_measured=flux_measured.cpu().detach(),
+                            flux_prediction_train=flux_prediction_train.cpu().detach(),
+                            flux_prediction_test=flux_prediction_test.cpu().detach(),
+                            data_split=data_split,
+                            plot_name=f"alignment_{epoch}",
+                        )
+
+                    # Early stopping when loss did not improve for a predefined number of epochs.
+                    if stop:
+                        log.info(f"Early stopping at epoch {epoch}.")
+                        break
+
+                    loss_history_list.append(loss.detach().cpu().item())
+                    epoch += 1
+
+                loss_history.append(
+                    {
+                        "total_loss": loss_history_list,
+                    }
+                )
+
+                active_indices_group = torch.nonzero(
+                    active_heliostats_mask != 0, as_tuple=True
+                )[0]
+
+                final_indices = (
+                    active_indices_group
+                    + final_loss_start_indices[heliostat_group_index]
+                )
+
+                final_loss_per_heliostat[final_indices] = loss_per_heliostat
+
+                log.info(f"Rank: {rank}, Kinematics reconstructed.")
+
+        if self.ddp_setup["is_distributed"]:
+            for index, heliostat_group in enumerate(
+                self.scenario.heliostat_field.heliostat_groups
+            ):
+                source = self.ddp_setup["ranks_to_groups_mapping"][index]
+                torch.distributed.broadcast(
+                    heliostat_group.kinematics.rotation_deviation_parameters,
+                    src=source[indices.first_rank_from_group],
+                )
+                torch.distributed.broadcast(
+                    heliostat_group.kinematics.actuators.optimizable_parameters,
+                    src=source[indices.first_rank_from_group],
+                )
+            torch.distributed.all_reduce(
+                final_loss_per_heliostat, op=torch.distributed.ReduceOp.MIN
+            )
+
+            final_loss_history_all_groups: list[list[dict[str, list[float]]]] = [
+                [] for _ in range(self.ddp_setup["world_size"])
+            ]
+            torch.distributed.all_gather_object(
+                final_loss_history_all_groups, loss_history
+            )
+
+            log.info(f"Rank: {rank}, synchronized after kinematics reconstruction.")
+
+        else:
+            final_loss_history_all_groups = [loss_history]
+
+        for heliostat_group in self.scenario.heliostat_field.heliostat_groups:
+            heliostat_group.kinematics.rotation_deviation_parameters = (
+                heliostat_group.kinematics.rotation_deviation_parameters.detach()
+            )
+
+        return final_loss_per_heliostat.detach().cpu(), final_loss_history_all_groups
+
+    def _reconstruct_kinematics_parameters_with_raytracing(
+        self,
+        loss_definition: Loss,
+        device: torch.device | None = None,
+    ) -> tuple[torch.Tensor, list[list[dict[str, list[float]]]]]:
+        """
+        Reconstruct the kinematics parameters using ray tracing.
+
+        This reconstruction method optimizes the kinematics parameters by extracting the focal points
+        of calibration images and using heliostat-tracing.
+
+        Parameters
+        ----------
+        loss_definition : Loss
+            Definition of the loss function and pre-processing of the prediction.
+        device : torch.device | None
+            The device on which to perform computations or load tensors and models (default is None).
+            If None, ARTIST will automatically select the most appropriate
+            device (CUDA or CPU) based on availability and OS.
+
+        Returns
+        -------
+        torch.Tensor
+            The final loss of the kinematics reconstruction for each heliostat in each group.
+            Shape is ``[total_number_of_heliostats_in_scenario]``.
+        list[list[dict[str, list[float]]]]
+            Loss histories over epochs grouped by rank.
+            Outer list: one entry per rank.
+            Inner list: one entry per heliostat group processed on that rank.
+            Each group entry is a dict with key ``"total_loss"`` mapping to a list
+            of per-epoch scalar loss values.
+            In non-distributed mode, this is a single-rank container: ``[local_group_histories]``.
+        """
+        device = get_device(device=device)
+        rank = self.ddp_setup["rank"]
+
+        if rank == 0:
+            log.info("Beginning kinematics reconstruction with ray tracing.")
+
+        # Initialize final loss per heliostat, group offset table into global heliostat index space, and
+        # per-group loss curves for this rank.
+        final_loss_per_heliostat = torch.full(
+            (self.scenario.heliostat_field.number_of_heliostats_per_group.sum(),),
+            torch.inf,
+            device=device,
+        )
+        final_loss_start_indices = torch.cat(
+            [
+                torch.tensor([0], device=device),
+                self.scenario.heliostat_field.number_of_heliostats_per_group.cumsum(
+                    indices.heliostat_dimension
+                ),
+            ]
+        )
+        loss_history: list[dict[str, list[float]]] = []
+
+        # Iterate heliostat groups assigned to this rank.
+        for heliostat_group_index in self.ddp_setup["groups_to_ranks_mapping"][rank]:
+            heliostat_group: HeliostatGroup = (
+                self.scenario.heliostat_field.heliostat_groups[heliostat_group_index]
+            )
+            # Load data parser and input file mapping, then parse the calibration data.
+            parser = cast(CalibrationDataParser, self.data[constants.data_parser])
+            heliostat_mapping = cast(
+                list[tuple[str, list[pathlib.Path], list[pathlib.Path]]],
+                self.data[constants.heliostat_data_mapping],
+            )
+            (
+                flux_measured,
+                focal_spots_measured,
+                incident_ray_directions,
+                motor_positions,
+                active_heliostats_mask,
+                target_area_indices,
+            ) = parser.parse_data_for_reconstruction(
+                heliostat_data_mapping=heliostat_mapping,
+                heliostat_group=heliostat_group,
+                scenario=self.scenario,
+                bitmap_resolution=self.bitmap_resolution,
+                device=device,
+            )
+
+            if active_heliostats_mask.sum() > 0:
+                data_split: training.TrainTestSplit = training.train_test_split(
+                    active_heliostats_mask=active_heliostats_mask,
+                    flux_measured=flux_measured,
+                    focal_spots_measured=focal_spots_measured,
+                    incident_ray_directions=incident_ray_directions,
+                    motor_positions=motor_positions,
+                    target_area_indices=target_area_indices,
+                    test_fraction=0.25,
+                    device=device,
+                )
 
                 # Set up optimizer, scheduler, and early stopping.
                 optimizer_params = [
@@ -319,24 +854,6 @@ def _reconstruct_kinematics_parameters_with_raytracing(
                     }
                 ]
 
-                if initial_actuator_params is not None:
-                    optimizer_params.extend(
-                        [
-                            {
-                                "params": delta_angle,
-                                "lr": self.optimizer_dict[
-                                    constants.initial_learning_rate_initial_angles
-                                ],
-                            },
-                            {
-                                "params": delta_stroke,
-                                "lr": self.optimizer_dict[
-                                    constants.initial_learning_rate_initial_stroke_length
-                                ],
-                            },
-                        ]
-                    )
-
                 optimizer = torch.optim.Adam(optimizer_params)
 
                 # Create a learning rate scheduler.
@@ -372,37 +889,16 @@ def _reconstruct_kinematics_parameters_with_raytracing(
                 ):
                     optimizer.zero_grad()
 
-                    # Get actuator parameters from reparametrized version.
-                    actuator_params = torch.cat(
-                        [
-                            ((angle_normalized + delta_angle) * angle_std + angle_mean)[
-                                :, None, :
-                            ],
-                            (
-                                (stroke_length_normalized + delta_stroke) * stroke_std
-                                + stroke_mean
-                            )[:, None, :],
-                        ],
-                        dim=-1,
-                    ).view_as(
-                        heliostat_group.kinematics.actuators.optimizable_parameters
-                    )
-                    heliostat_group.kinematics.actuators.optimizable_parameters = (
-                        actuator_params
-                    )
-
                     # Activate heliostats.
                     heliostat_group.activate_heliostats(
-                        active_heliostats_mask=active_heliostats_mask, device=device
+                        active_heliostats_mask=data_split.active_heliostats_mask_train,
+                        device=device,
                     )
 
                     # Align heliostats.
-                    heliostat_group.align_surfaces_with_incident_ray_directions(
-                        aim_points=self.scenario.solar_tower.get_centers_of_target_areas(
-                            target_area_indices=target_area_indices, device=device
-                        ),
-                        incident_ray_directions=incident_ray_directions,
-                        active_heliostats_mask=active_heliostats_mask,
+                    heliostat_group.align_surfaces_with_motor_positions(
+                        motor_positions=data_split.motor_positions_train,
+                        active_heliostats_mask=data_split.active_heliostats_mask_train,
                         device=device,
                     )
 
@@ -420,42 +916,42 @@ def _reconstruct_kinematics_parameters_with_raytracing(
                     )
 
                     # Perform heliostat-based ray tracing.
-                    flux_distributions, _, _, _ = ray_tracer.trace_rays(
-                        incident_ray_directions=incident_ray_directions,
-                        active_heliostats_mask=active_heliostats_mask,
-                        target_area_indices=target_area_indices,
+                    flux_prediction_train, _, _, _ = ray_tracer.trace_rays(
+                        incident_ray_directions=data_split.incident_ray_directions_train,
+                        active_heliostats_mask=data_split.active_heliostats_mask_train,
+                        target_area_indices=data_split.target_area_indices_train,
                         device=device,
                     )
 
                     sample_indices_for_local_rank = ray_tracer.get_sampler_indices()
 
-                    # Compute loss from prediction vs. measured focal spots.
+                    # Compute loss from prediction vs. measured flux.
                     loss_per_sample = loss_definition(
-                        prediction=flux_distributions,
-                        ground_truth=focal_spots_measured[
+                        prediction=flux_prediction_train,
+                        ground_truth=data_split.flux_measured_train[
                             sample_indices_for_local_rank
                         ],
-                        target_area_indices=target_area_indices[
+                        target_area_indices=data_split.target_area_indices_train[
                             sample_indices_for_local_rank
                         ],
-                        reduction_dimensions=(indices.focal_spots,),
+                        reduction_dimensions=(
+                            indices.batched_bitmap_e,
+                            indices.batched_bitmap_u,
+                        ),
                         device=device,
                     )
 
-                    number_of_samples_per_heliostat = int(
-                        heliostat_group.active_heliostats_mask.sum()
-                        / (heliostat_group.active_heliostats_mask > 0).sum()
-                    )
-
-                    loss_per_heliostat = mean_loss_per_heliostat(
+                    loss_per_heliostat = reduce_loss_per_sample(
                         loss_per_sample=loss_per_sample,
-                        number_of_samples_per_heliostat=number_of_samples_per_heliostat,
+                        number_of_samples_per_heliostat=data_split.number_of_train_samples,
+                        reduction=partial(torch.median, dim=1),
                     )
 
-                    loss = loss_per_heliostat.mean()
+                    loss = torch.mean(loss_per_heliostat)
 
                     loss.backward()
 
+                    # Nested-DDP gradient synchronization within heliostat-group subgroup.
                     if self.ddp_setup["is_nested"]:
                         # Reduce gradients within each heliostat group.
                         for param_group in optimizer.param_groups:
@@ -480,20 +976,39 @@ def _reconstruct_kinematics_parameters_with_raytracing(
                     else:
                         scheduler.step()
 
-                    if epoch % log_step == 0:
+                    is_last_epoch = (
+                        epoch == self.optimizer_dict[constants.max_epoch] - 1
+                    )
+                    stop = early_stopper.step(loss.item())
+
+                    if epoch % log_step == 0 or is_last_epoch or stop:
                         log.info(
-                            f"Rank: {rank}, Epoch: {epoch}, Loss: {loss},",
+                            f"Rank: {rank}, Epoch: {epoch}, Loss: {loss}",
                         )
 
-                    loss_history_list.append(loss.detach().cpu().item())
+                        with torch.no_grad():
+                            flux_prediction_test = self._validate(
+                                heliostat_group=heliostat_group,
+                                data_split=data_split,
+                                reduction=partial(torch.median, dim=1),
+                                device=device,
+                            )
+
+                    if self.plot_results and (is_last_epoch or stop):
+                        self._plot_fluxes(
+                            flux_measured=flux_measured.cpu().detach(),
+                            flux_prediction_train=flux_prediction_train.cpu().detach(),
+                            flux_prediction_test=flux_prediction_test.cpu().detach(),
+                            data_split=data_split,
+                            plot_name=f"raytracing_{epoch}",
+                        )
 
                     # Early stopping when loss did not improve for a predefined number of epochs.
-                    stop = early_stopper.step(loss.item())
-
                     if stop:
                         log.info(f"Early stopping at epoch {epoch}.")
                         break
 
+                    loss_history_list.append(loss.detach().cpu().item())
                     epoch += 1
 
                 loss_history.append(
@@ -503,8 +1018,8 @@ def _reconstruct_kinematics_parameters_with_raytracing(
                 )
 
                 local_indices = (
-                    sample_indices_for_local_rank[::number_of_samples_per_heliostat]
-                    // number_of_samples_per_heliostat
+                    sample_indices_for_local_rank[:: data_split.number_of_train_samples]
+                    // data_split.number_of_train_samples
                 )
 
                 global_active_indices = torch.nonzero(
@@ -551,4 +1066,9 @@ def _reconstruct_kinematics_parameters_with_raytracing(
         else:
             final_loss_history_all_groups = [loss_history]
 
+        for heliostat_group in self.scenario.heliostat_field.heliostat_groups:
+            heliostat_group.kinematics.rotation_deviation_parameters = (
+                heliostat_group.kinematics.rotation_deviation_parameters.detach()
+            )
+
         return final_loss_per_heliostat.detach().cpu(), final_loss_history_all_groups
diff --git a/artist/optim/loss.py b/artist/optim/loss.py
index de1e67ecd..f176d8bb7 100644
--- a/artist/optim/loss.py
+++ b/artist/optim/loss.py
@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Callable
 
 import torch
 
@@ -122,7 +122,7 @@ def __call__(
 
 class FocalSpotLoss(Loss):
     """
-    A loss defined as Euclidean distance between the predicted focal spot coordinate and the ground-truth coordinate.
+    A loss defined as Euclidean distance between the predicted focal spot coordinate and the ground-truth focal spot coordinate.
 
     Attributes
     ----------
@@ -157,7 +157,7 @@ def __call__(
         r"""
         Compute the focal spot loss.
 
-        First the focal spots of the prediction are computed, then the loss is computed and reduced
+        First the focal spots of the prediction and ground truth flux maps are computed, then the loss is computed and reduced
         along the specified dimensions.
 
         Parameters
@@ -203,21 +203,48 @@ def __call__(
             device=device,
         )
 
-        focal_spot_coordinates = coordinates.bitmap_coordinates_to_target_coordinates(
-            bitmap_coordinates=focal_spots_bitmap,
-            bitmap_resolution=torch.tensor(
-                [
-                    prediction.shape[indices.batched_bitmap_u],
-                    prediction.shape[indices.batched_bitmap_e],
-                ],
+        focal_spot_coordinates_prediction = (
+            coordinates.bitmap_coordinates_to_target_coordinates(
+                bitmap_coordinates=focal_spots_bitmap,
+                bitmap_resolution=torch.tensor(
+                    [
+                        prediction.shape[indices.batched_bitmap_u],
+                        prediction.shape[indices.batched_bitmap_e],
+                    ],
+                    device=device,
+                ),
+                solar_tower=self.scenario.solar_tower,
+                target_area_indices=target_area_indices,
                 device=device,
-            ),
-            solar_tower=self.scenario.solar_tower,
-            target_area_indices=target_area_indices,
+            )
+        )
+
+        focal_spots_ground_truth = bitmap.get_center_of_mass(
+            bitmaps=ground_truth,
             device=device,
         )
 
-        return torch.norm(focal_spot_coordinates[:, :3] - ground_truth[:, :3], dim=1)
+        focal_spot_coordinates_ground_truth = (
+            coordinates.bitmap_coordinates_to_target_coordinates(
+                bitmap_coordinates=focal_spots_ground_truth,
+                bitmap_resolution=torch.tensor(
+                    [
+                        prediction.shape[indices.batched_bitmap_u],
+                        prediction.shape[indices.batched_bitmap_e],
+                    ],
+                    device=device,
+                ),
+                solar_tower=self.scenario.solar_tower,
+                target_area_indices=target_area_indices,
+                device=device,
+            )
+        )
+
+        return torch.norm(
+            focal_spot_coordinates_prediction[:, :3]
+            - focal_spot_coordinates_ground_truth[:, :3],
+            dim=1,
+        )
 
 
 class PixelLoss(Loss):
@@ -236,17 +263,9 @@ class PixelLoss(Loss):
     :class:`Loss` : Reference to the parent class.
     """
 
-    def __init__(self, scenario: Scenario) -> None:
-        """
-        Initialize the pixel loss.
-
-        Parameters
-        ----------
-        scenario : Scenario
-            The scenario.
-        """
+    def __init__(self) -> None:
+        """Initialize the pixel loss."""
         super().__init__(loss_function=torch.nn.MSELoss(reduction="none"))
-        self.scenario = scenario
 
     def __call__(
         self,
@@ -255,10 +274,10 @@ def __call__(
         **kwargs: Any,
     ) -> torch.Tensor:
         r"""
-        Compute the pixel loss.
+        Compute the normalized pixel-wise loss.
 
-        First the predicted bitmaps and the ground truth are normalized, then the loss is
-        computed and reduced along the specified dimensions.
+        To make the loss invariant to the overall magnitude of the ground truth flux, the summed loss is divided
+        by the total ground truth intensity per sample.
 
         Parameters
         ----------
@@ -270,7 +289,7 @@ def __call__(
             Shape is ``[number_of_samples, bitmap_resolution_e, bitmap_resolution_u]``.
         \*\*kwargs : Any
             Keyword arguments.
-            ``reduction_dimensions``, ``target_area_indices``, and ``device`` are expected keyword arguments for the pixel loss.
+            ``reduction_dimensions`` is an expected keyword arguments for the pixel loss.
 
         Raises
         ------
@@ -283,7 +302,7 @@ def __call__(
             The summed MSE pixel loss reduced along the specified dimensions.
             Shape is ``[number_of_samples]``.
         """
-        expected_kwargs = ["reduction_dimensions", "device", "target_area_indices"]
+        expected_kwargs = ["reduction_dimensions"]
         errors = []
         for key in expected_kwargs:
             if key not in kwargs:
@@ -294,9 +313,9 @@ def __call__(
                 + " ".join(errors)
             )
 
-        loss = self.loss_function(prediction, ground_truth)
-
-        return loss.sum(dim=kwargs["reduction_dimensions"])
+        return self.loss_function(prediction, ground_truth).sum(
+            dim=kwargs["reduction_dimensions"]
+        ) / ground_truth.sum(dim=(1, 2))
 
 
 class KLDivergenceLoss(Loss):
@@ -403,6 +422,55 @@ class AngleLoss(Loss):
     :class:`Loss` : Reference to the parent class.
     """
 
+    def __init__(self) -> None:
+        """Initialize the angle loss."""
+        super().__init__(loss_function=None)
+
+    def __call__(
+        self,
+        prediction: torch.Tensor,
+        ground_truth: torch.Tensor,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        r"""
+        Compute the angular distance between prediction and ground truth.
+
+        Parameters
+        ----------
+        prediction : torch.Tensor
+            The predicted values.
+            Shape is ``[number_of_samples, 4]``.
+        ground_truth : torch.Tensor
+            The ground truth.
+            Shape is ``[number_of_samples, 4]``.
+        \*\*kwargs : Any
+            Keyword arguments.
+
+        Returns
+        -------
+        torch.Tensor
+            The summed loss reduced along the specified dimensions.
+            Shape is ``[number_of_samples]``.
+        """
+        prediction = torch.nn.functional.normalize(prediction[:, :3])
+        ground_truth = torch.nn.functional.normalize(ground_truth[:, :3])
+        return torch.acos((prediction * ground_truth).sum(dim=-1).clamp(-1.0, 1.0))
+
+
+class CosineSimilarityLoss(Loss):
+    """
+    A loss defined as the cosine similarity between prediction and ground truth.
+
+    Attributes
+    ----------
+    loss_function : torch.nn.Module
+        A torch module implementing a loss.
+
+    See Also
+    --------
+    :class:`Loss` : Reference to the parent class.
+    """
+
     def __init__(self) -> None:
         """Initialize the angle loss."""
         super().__init__(loss_function=torch.nn.CosineSimilarity(dim=-1))
@@ -436,26 +504,31 @@ def __call__(
         return 1.0 - self.loss_function(prediction, ground_truth)
 
 
-def mean_loss_per_heliostat(
+def reduce_loss_per_sample(
     loss_per_sample: torch.Tensor,
     number_of_samples_per_heliostat: int,
+    reduction: Callable[..., Any],
 ) -> torch.Tensor:
     """
-    Calculate the mean loss per heliostat from a loss per sample.
+    Calculate the loss per heliostat from a loss per sample.
+
+    The reduction operation is chosen via the `reduction` parameter.
 
     Parameters
     ----------
     loss_per_sample : torch.Tensor
         Loss per sample.
-        Shape is ``[number_of_samples]``.
+        Tensor of shape [number_of_samples].
     number_of_samples_per_heliostat : int
         Number of samples per heliostat.
+    reduction : Callable[..., Any]
+        Reduction function applied across the sample dimension for each heliostat.
 
     Returns
     -------
     torch.Tensor
         Loss per heliostat.
-        Shape is ``[number_of_heliostats]``.
+        Tensor of shape [number_of_heliostats].
     """
     number_of_heliostats = int(
         loss_per_sample.numel() // number_of_samples_per_heliostat
@@ -464,8 +537,11 @@ def mean_loss_per_heliostat(
         : number_of_heliostats * number_of_samples_per_heliostat
     ]
 
-    loss_per_heliostat = loss_per_sample.view(
-        number_of_heliostats, number_of_samples_per_heliostat
-    ).mean(dim=1)
+    reduced = reduction(
+        loss_per_sample.view(number_of_heliostats, number_of_samples_per_heliostat)
+    )
+
+    if not isinstance(reduced, torch.Tensor):
+        reduced = reduced.values
 
-    return loss_per_heliostat
+    return reduced
diff --git a/artist/optim/surface_reconstructor.py b/artist/optim/surface_reconstructor.py
index be25f3a5a..f40c7f551 100644
--- a/artist/optim/surface_reconstructor.py
+++ b/artist/optim/surface_reconstructor.py
@@ -1,8 +1,10 @@
 import logging
 import pathlib
+from functools import partial
 from typing import Any, cast
 
 import torch
+from matplotlib import pyplot as plt
 from torch.optim.lr_scheduler import LRScheduler
 
 from artist.field.heliostat_group import HeliostatGroup
@@ -11,7 +13,7 @@
 from artist.nurbs.surfaces import NURBSSurfaces
 from artist.nurbs.utils import create_nurbs_evaluation_grid
 from artist.optim import training
-from artist.optim.loss import Loss, mean_loss_per_heliostat
+from artist.optim.loss import KLDivergenceLoss, Loss, PixelLoss, reduce_loss_per_sample
 from artist.optim.regularizers import IdealSurfaceRegularizer, SmoothnessRegularizer
 from artist.raytracing.heliostat_ray_tracer import HeliostatRayTracer
 from artist.scenario.scenario import Scenario
@@ -63,6 +65,12 @@ class SurfaceReconstructor:
         Shape is ``[2]``.
     epsilon : float | None
         Small numerical offset used to avoid division by zero in the energy constraint.
+    plot_results : bool
+        Create flux plots in the last epoch of training for each heliostat and all its samples (slow!).
+    validation_loss_pixel : PixelLoss
+        Pixel loss used for validation.
+    validation_loss_kl_div : KLDivergenceLoss
+        Kullback-Leibler divergence loss used for validation.
 
     Note
     ----
@@ -90,6 +98,7 @@ def __init__(
         number_of_surface_points: torch.Tensor = torch.tensor([50, 50]),
         bitmap_resolution: torch.Tensor = torch.tensor([256, 256]),
         epsilon: float | None = 1e-12,
+        plot_results: bool = False,
         device: torch.device | None = None,
     ) -> None:
         """
@@ -117,6 +126,8 @@ def __init__(
             Shape is ``[2]``.
         epsilon : float | None
             Small numerical offset used to avoid division by zero in the energy constraint (default is 1e-12).
+        plot_results : bool
+            Create flux plots in the last epoch of training for each heliostat and all its samples (slow!, default is ``False``).
         device : torch.device | None
             The device on which to perform computations or load tensors and models (default is None).
             If None, ``ARTIST`` will automatically select the most appropriate
@@ -139,6 +150,224 @@ def __init__(
         self.dni = dni
         self.bitmap_resolution = bitmap_resolution.to(device)
         self.epsilon = epsilon
+        self.plot_results = plot_results
+
+        self.validation_loss_pixel = PixelLoss()
+        self.validation_loss_kl_div = KLDivergenceLoss()
+
+    def _validate(
+        self,
+        heliostat_group: HeliostatGroup,
+        data_split: training.TrainTestSplit,
+        evaluation_points: torch.Tensor,
+        device: torch.device | None = None,
+    ) -> torch.Tensor:
+        """
+        Validate the surface reconstruction for a specified heliostat group on the test data.
+
+        Parameters
+        ----------
+        heliostat_group : HeliostatGroup
+            Heliostat group to validate.
+        data_split : training.TrainTestSplit
+            Train/test split containing all test tensors and metadata.
+        evaluation_points : torch.Tensor
+            Evaluation points for the NURBS surface sampling.
+        device : torch.device | None
+            The device on which to perform computations or load tensors and models (default is None).
+            If None, ARTIST will automatically select the most appropriate
+            device (CUDA or CPU) based on availability and OS.
+
+        Returns
+        -------
+        torch.Tensor
+            Predicted flux distributions for the local validation samples.
+            Shape is ``[number_of_local_test_samples, height, width]``.
+        """
+        device = get_device(device=device)
+
+        heliostat_group.activate_heliostats(
+            active_heliostats_mask=data_split.active_heliostats_mask_test,
+            device=device,
+        )
+
+        nurbs_surfaces = NURBSSurfaces(
+            degrees=heliostat_group.nurbs_degrees,
+            control_points=heliostat_group.active_nurbs_control_points,
+            device=device,
+        )
+
+        (
+            new_surface_points,
+            new_surface_normals,
+        ) = nurbs_surfaces.calculate_surface_points_and_normals(
+            evaluation_points=evaluation_points[data_split.test_indices],
+            canting=heliostat_group.active_canting,
+            facet_translations=heliostat_group.active_facet_translations,
+            device=device,
+        )
+
+        heliostat_group.active_surface_points = new_surface_points.reshape(
+            heliostat_group.active_surface_points.shape[indices.heliostat_dimension],
+            -1,
+            4,
+        )
+        heliostat_group.active_surface_normals = new_surface_normals.reshape(
+            heliostat_group.active_surface_normals.shape[indices.heliostat_dimension],
+            -1,
+            4,
+        )
+
+        heliostat_group.align_surfaces_with_incident_ray_directions(
+            aim_points=self.scenario.solar_tower.get_centers_of_target_areas(
+                target_area_indices=data_split.target_area_indices_test, device=device
+            ),
+            incident_ray_directions=data_split.incident_ray_directions_test,
+            active_heliostats_mask=data_split.active_heliostats_mask_test,
+            device=device,
+        )
+
+        ray_tracer = HeliostatRayTracer(
+            scenario=self.scenario,
+            heliostat_group=heliostat_group,
+            blocking_active=False,
+            world_size=self.ddp_setup["heliostat_group_world_size"],
+            rank=self.ddp_setup["heliostat_group_rank"],
+            batch_size=self.optimizer_dict[constants.batch_size],
+            random_seed=self.ddp_setup["heliostat_group_rank"],
+            dni=self.dni,
+            bitmap_resolution=self.bitmap_resolution,
+        )
+
+        flux_prediction, _, _, _ = ray_tracer.trace_rays(
+            incident_ray_directions=data_split.incident_ray_directions_test,
+            active_heliostats_mask=data_split.active_heliostats_mask_test,
+            target_area_indices=data_split.target_area_indices_test,
+            device=device,
+        )
+
+        cropped_flux_distributions = bitmap.crop_flux_distributions_around_center(
+            flux_distributions=flux_prediction,
+            solar_tower=self.scenario.solar_tower,
+            target_area_indices=data_split.target_area_indices_test,
+            device=device,
+        )
+
+        indices_for_local_rank = ray_tracer.get_sampler_indices()
+
+        loss_pixel_per_sample = self.validation_loss_pixel(
+            prediction=cropped_flux_distributions,
+            ground_truth=data_split.flux_measured_test[indices_for_local_rank],
+            reduction_dimensions=(
+                1,
+                2,
+            ),
+        )
+        loss_kl_div_per_sample = self.validation_loss_kl_div(
+            prediction=cropped_flux_distributions,
+            ground_truth=data_split.flux_measured_test[indices_for_local_rank],
+            reduction_dimensions=(
+                1,
+                2,
+            ),
+        )
+
+        test_loss_pixel = reduce_loss_per_sample(
+            loss_per_sample=loss_pixel_per_sample,
+            number_of_samples_per_heliostat=data_split.number_of_test_samples,
+            reduction=partial(torch.mean),
+        )
+        test_loss_kl_div = reduce_loss_per_sample(
+            loss_per_sample=loss_kl_div_per_sample,
+            number_of_samples_per_heliostat=data_split.number_of_test_samples,
+            reduction=partial(torch.mean),
+        )
+
+        log.info(
+            "pixel mean: %.5f, kl div mean: %.5f",
+            torch.mean(test_loss_pixel).item(),
+            torch.mean(test_loss_kl_div).item(),
+        )
+
+        return flux_prediction
+
+    def _plot_fluxes(
+        self,
+        flux_measured: torch.Tensor,
+        flux_prediction_train: torch.Tensor,
+        flux_prediction_test: torch.Tensor,
+        data_split: training.TrainTestSplit,
+        plot_name: str,
+    ) -> None:
+        """
+        Plot predicted and measured flux maps for each heliostat sample.
+
+        Each row in the generated figure corresponds to one sample of a
+        heliostat, where the left column contains the predicted flux and
+        the right column contains the measured flux.
+        The subplot borders are color-coded to indicate whether a sample
+        belongs to the training samples (green) or testing samples (red)
+        One figure is generated per heliostat.
+
+        Parameters
+        ----------
+        flux_measured : torch.Tensor
+            Ground-truth measured flux maps.
+        flux_prediction_train : torch.Tensor
+            Predicted flux maps of the training samples.
+        flux_prediction_test : torch.Tensor
+            Predicted flux maps of the testing samples.
+        data_split : training.TrainTestSplit
+            Information on the train/test split.
+        plot_name : str
+            Name suffix used when saving the generated plot files.
+        """
+        device = torch.device("cpu")
+
+        flux_predicted = torch.zeros_like(flux_measured, device=device)
+        flux_predicted[data_split.train_indices] = flux_prediction_train
+        flux_predicted[data_split.test_indices] = flux_prediction_test
+        samples_per_heliostat = data_split.number_of_samples_per_heliostat
+        total_samples = flux_measured.shape[0]
+        train_indices = set(data_split.train_indices.tolist())
+        test_indices = set(data_split.test_indices.tolist())
+
+        for heliostat_start_index in range(0, total_samples, samples_per_heliostat):
+            fig, axes = plt.subplots(
+                samples_per_heliostat,
+                2,
+                figsize=(8, samples_per_heliostat * 4),
+            )
+            heliostat_index = heliostat_start_index // samples_per_heliostat
+
+            for sample_offset in range(samples_per_heliostat):
+                sample_index = heliostat_start_index + sample_offset
+
+                if sample_index in train_indices:
+                    border_color = "green"
+                    split_name = "TRAIN"
+                elif sample_index in test_indices:
+                    border_color = "red"
+                    split_name = "TEST"
+
+                axes[sample_offset, 0].imshow(flux_predicted[sample_index])
+                axes[sample_offset, 0].set_title(
+                    f"Predicted Flux - Heliostat {heliostat_index} ({split_name})"
+                )
+
+                axes[sample_offset, 1].imshow(flux_measured[sample_index])
+                axes[sample_offset, 1].set_title(
+                    f"Measured Flux - Heliostat {heliostat_index} ({split_name})"
+                )
+
+                for spine in axes[sample_offset, :].flat:
+                    for spines in spine.spines.values():
+                        spines.set_edgecolor(border_color)
+                        spines.set_linewidth(4)
+
+            plt.tight_layout()
+            plt.savefig(f"heliostat_{heliostat_index}_{plot_name}")
+            plt.close(fig)
 
     def reconstruct_surfaces(
         self,
@@ -207,19 +436,17 @@ def reconstruct_surfaces(
                 self.scenario.heliostat_field.heliostat_groups[heliostat_group_index]
             )
 
-            # Load calibration parser and input file mapping.
+            # Load data parser and input file mapping, then parse the calibration data.
             parser = cast(CalibrationDataParser, self.data[constants.data_parser])
             heliostat_mapping = cast(
                 list[tuple[str, list[pathlib.Path], list[pathlib.Path]]],
                 self.data[constants.heliostat_data_mapping],
             )
-
-            # Obtain measured flux + metadata for this group.
             (
-                measured_flux_distributions,
-                _,
+                flux_measured,
+                focal_spots_measured,
                 incident_ray_directions,
-                _,
+                motor_positions,
                 active_heliostats_mask,
                 target_area_indices,
             ) = parser.parse_data_for_reconstruction(
@@ -232,7 +459,16 @@ def reconstruct_surfaces(
 
             # Skip groups with no active heliostats.
             if active_heliostats_mask.sum() > 0:
-                # Create NURBS evaluation points: Build a UV grid for NURBS sampling and expand to required shape.
+                data_split: training.TrainTestSplit = training.train_test_split(
+                    active_heliostats_mask=active_heliostats_mask,
+                    flux_measured=flux_measured,
+                    focal_spots_measured=focal_spots_measured,
+                    incident_ray_directions=incident_ray_directions,
+                    motor_positions=motor_positions,
+                    target_area_indices=target_area_indices,
+                    test_fraction=0.25,
+                    device=device,
+                )
                 evaluation_points = (
                     create_nurbs_evaluation_grid(
                         number_of_evaluation_points=self.number_of_surface_points,
@@ -247,14 +483,13 @@ def reconstruct_surfaces(
                         -1,
                     )
                 )
-
                 # Keep a frozen copy of original control points for regularization terms.
                 with torch.no_grad():
                     original_control_points = heliostat_group.nurbs_control_points[
                         active_heliostats_mask > 0
                     ].clone()
 
-                # Create the optimizer: Optimize NURBS control points directly.
+                # Create the optimizer.
                 optimizer = torch.optim.Adam(
                     [heliostat_group.nurbs_control_points.requires_grad_()],
                     lr=float(self.optimizer_dict[constants.initial_learning_rate]),
@@ -310,23 +545,16 @@ def reconstruct_surfaces(
                     if self.optimizer_dict[constants.log_step] == 0
                     else self.optimizer_dict[constants.log_step]
                 )
-                max_epoch = int(
-                    torch.tensor(
-                        [self.optimizer_dict[constants.max_epoch]],
-                        device=device,
-                    )
-                )
-
-                # Optimization loop.
                 while (
                     total_loss > float(self.optimizer_dict[constants.tolerance])
-                    and epoch <= max_epoch
+                    and epoch <= self.optimizer_dict[constants.max_epoch]
                 ):
                     optimizer.zero_grad()
 
-                    # Activate heliostats according to current mask.
+                    # Activate heliostats.
                     heliostat_group.activate_heliostats(
-                        active_heliostats_mask=active_heliostats_mask, device=device
+                        active_heliostats_mask=data_split.active_heliostats_mask_train,
+                        device=device,
                     )
 
                     # Build NURBS surface from current control points.
@@ -341,7 +569,7 @@ def reconstruct_surfaces(
                         new_surface_points,
                         new_surface_normals,
                     ) = nurbs_surfaces.calculate_surface_points_and_normals(
-                        evaluation_points=evaluation_points,
+                        evaluation_points=evaluation_points[data_split.train_indices],
                         canting=heliostat_group.active_canting,
                         facet_translations=heliostat_group.active_facet_translations,
                         device=device,
@@ -368,10 +596,11 @@ def reconstruct_surfaces(
                     # Align heliostat surfaces toward target under current incident ray directions.
                     heliostat_group.align_surfaces_with_incident_ray_directions(
                         aim_points=self.scenario.solar_tower.get_centers_of_target_areas(
-                            target_area_indices=target_area_indices, device=device
+                            target_area_indices=data_split.target_area_indices_train,
+                            device=device,
                         ),
-                        incident_ray_directions=incident_ray_directions,
-                        active_heliostats_mask=active_heliostats_mask,
+                        incident_ray_directions=data_split.incident_ray_directions_train,
+                        active_heliostats_mask=data_split.active_heliostats_mask_train,
                         device=device,
                     )
 
@@ -389,41 +618,38 @@ def reconstruct_surfaces(
                     )
 
                     # Perform heliostat-based ray tracing to obtain simulated flux from current reconstructed surfaces.
-                    flux_distributions, _, _, _ = ray_tracer.trace_rays(
-                        incident_ray_directions=incident_ray_directions,
-                        active_heliostats_mask=active_heliostats_mask,
-                        target_area_indices=target_area_indices,
+                    flux_prediction_train, _, _, _ = ray_tracer.trace_rays(
+                        incident_ray_directions=data_split.incident_ray_directions_train,
+                        active_heliostats_mask=data_split.active_heliostats_mask_train,
+                        target_area_indices=data_split.target_area_indices_train,
                         device=device,
                     )
 
-                    # Recover local sampler ordering and samples-per-heliostat factor.
-                    sample_indices_for_local_rank = ray_tracer.get_sampler_indices()
-                    number_of_samples_per_heliostat = int(
-                        heliostat_group.active_heliostats_mask.sum()
-                        / (heliostat_group.active_heliostats_mask > 0).sum()
-                    )
-                    local_indices = (
-                        sample_indices_for_local_rank[::number_of_samples_per_heliostat]
-                        // number_of_samples_per_heliostat
-                    )
-
                     # Crop predictions around center before comparing to measurements.
-                    cropped_flux_distributions = (
+                    cropped_flux_predictions = (
                         bitmap.crop_flux_distributions_around_center(
-                            flux_distributions=flux_distributions,
+                            flux_distributions=flux_prediction_train,
                             solar_tower=self.scenario.solar_tower,
-                            target_area_indices=target_area_indices,
+                            target_area_indices=data_split.target_area_indices_train,
                             device=device,
                         )
                     )
 
-                    # Flux loss (data-fit term): Compare predicted and measured flux maps.
+                    sample_indices_for_local_rank = ray_tracer.get_sampler_indices()
+                    local_indices = (
+                        sample_indices_for_local_rank[
+                            :: data_split.number_of_train_samples
+                        ]
+                        // data_split.number_of_train_samples
+                    )
+
+                    # Compute loss from prediction vs. measured flux.
                     flux_loss_per_sample = loss_definition(
-                        prediction=cropped_flux_distributions,
-                        ground_truth=measured_flux_distributions[
+                        prediction=cropped_flux_predictions,
+                        ground_truth=data_split.flux_measured_train[
                             sample_indices_for_local_rank
                         ],
-                        target_area_indices=target_area_indices[
+                        target_area_indices=data_split.target_area_indices_train[
                             sample_indices_for_local_rank
                         ],
                         reduction_dimensions=(
@@ -432,27 +658,32 @@ def reconstruct_surfaces(
                         ),
                         device=device,
                     )
-                    flux_loss_per_heliostat = mean_loss_per_heliostat(
+
+                    flux_loss_per_heliostat = reduce_loss_per_sample(
                         loss_per_sample=flux_loss_per_sample,
-                        number_of_samples_per_heliostat=number_of_samples_per_heliostat,
+                        number_of_samples_per_heliostat=data_split.number_of_train_samples,
+                        reduction=partial(torch.mean),
                     )
 
                     # Add Augmented-Lagrangian constraint to ensure that flux integral is conserved,
                     # i.e., intensity does not get lost.
                     if epoch == 0:
-                        flux_integrals_reference = cropped_flux_distributions.sum(
-                            dim=(1, 2)
+                        flux_integrals_reference = cropped_flux_predictions.sum(
+                            dim=(indices.batched_bitmap_e, indices.batched_bitmap_u)
                         ).detach()
                     flux_integrals_relative_differences = (
-                        cropped_flux_distributions.sum(dim=(1, 2))
+                        cropped_flux_predictions.sum(
+                            dim=(indices.batched_bitmap_e, indices.batched_bitmap_u)
+                        )
                         - flux_integrals_reference
                     ) / (flux_integrals_reference + torch.tensor(self.epsilon))
                     flux_constraint_per_sample = torch.clamp(
                         -energy_tolerance - flux_integrals_relative_differences, min=0.0
                     )
-                    flux_constraint_per_heliostat = mean_loss_per_heliostat(
+                    flux_constraint_per_heliostat = reduce_loss_per_sample(
                         loss_per_sample=flux_constraint_per_sample,
-                        number_of_samples_per_heliostat=number_of_samples_per_heliostat,
+                        number_of_samples_per_heliostat=data_split.number_of_train_samples,
+                        reduction=partial(torch.mean),
                     )
                     flux_integrals_constraint = (
                         lambda_flux_integral * flux_constraint_per_heliostat
@@ -469,7 +700,7 @@ def reconstruct_surfaces(
                     if weight_smoothness > 0:
                         smoothness_loss_per_heliostat = smoothness_regularizer(
                             current_control_points=heliostat_group.active_nurbs_control_points[
-                                ::number_of_samples_per_heliostat
+                                :: data_split.number_of_train_samples
                             ][local_indices],
                             original_control_points=original_control_points[
                                 local_indices
@@ -479,7 +710,7 @@ def reconstruct_surfaces(
                     if weight_ideal_surface > 0:
                         ideal_surface_loss_per_heliostat = ideal_surface_regularizer(
                             current_control_points=heliostat_group.active_nurbs_control_points[
-                                ::number_of_samples_per_heliostat
+                                :: data_split.number_of_train_samples
                             ][local_indices],
                             original_control_points=original_control_points[
                                 local_indices
@@ -511,10 +742,9 @@ def reconstruct_surfaces(
                         + alpha * smoothness_loss_per_heliostat
                         + beta * ideal_surface_loss_per_heliostat
                     )
-                    # Average over all heliostats in group.
-                    total_loss = total_loss_per_heliostat.mean()
 
-                    # Back-propagate through ray tracing and NURBS evaluation pipeline.
+                    total_loss = torch.mean(total_loss_per_heliostat)
+
                     total_loss.backward()
 
                     # Update Augmented-Lagrangian multiplier.
@@ -559,11 +789,38 @@ def reconstruct_surfaces(
                     else:
                         scheduler.step()
 
-                    if epoch % log_step == 0 and rank == 0:
+                    is_last_epoch = (
+                        epoch == self.optimizer_dict[constants.max_epoch] - 1
+                    )
+                    stop = early_stopper.step(total_loss.item())
+
+                    if epoch % log_step == 0 or is_last_epoch or stop:
                         log.info(
-                            f"Rank: {rank}, Epoch: {epoch}, Loss: {total_loss}, LR: {optimizer.param_groups[indices.optimizer_param_group_0]['lr']}"
+                            f"Rank: {rank}, Epoch: {epoch}, Loss: {total_loss}",
+                        )
+
+                        with torch.no_grad():
+                            flux_prediction_test = self._validate(
+                                heliostat_group=heliostat_group,
+                                data_split=data_split,
+                                evaluation_points=evaluation_points,
+                                device=device,
+                            )
+
+                    if self.plot_results and (is_last_epoch or stop):
+                        self._plot_fluxes(
+                            flux_measured=flux_measured.cpu().detach(),
+                            flux_prediction_train=flux_prediction_train.cpu().detach(),
+                            flux_prediction_test=flux_prediction_test.cpu().detach(),
+                            data_split=data_split,
+                            plot_name=f"{epoch}",
                         )
 
+                    # Early stopping when loss did not improve for a predefined number of epochs.
+                    if stop:
+                        log.info(f"Early stopping at epoch {epoch}.")
+                        break
+
                     total_loss_history.append(total_loss.detach().cpu().item())
                     flux_loss_history.append(
                         flux_loss_per_heliostat.mean().detach().cpu().item()
@@ -589,13 +846,6 @@ def reconstruct_surfaces(
                         flux_integrals_constraint.mean().detach().cpu().item()
                     )
 
-                    # Early stopping when loss did not improve for a predefined number of epochs.
-                    stop = early_stopper.step(total_loss.item())
-
-                    if stop:
-                        log.info(f"Early stopping at epoch {epoch}.")
-                        break
-
                     epoch += 1
 
                 loss_history.append(
diff --git a/artist/optim/training.py b/artist/optim/training.py
index ad30d1286..8c20c3ed9 100644
--- a/artist/optim/training.py
+++ b/artist/optim/training.py
@@ -1,4 +1,5 @@
 from collections import deque
+from dataclasses import dataclass
 from typing import Deque
 
 import torch
@@ -6,6 +7,7 @@
 from torch.optim.lr_scheduler import LRScheduler
 
 from artist.util import constants
+from artist.util.env import get_device
 
 
 def exponential(
@@ -181,3 +183,180 @@ def step(self, loss: float) -> bool:
             self.counter += 1
 
         return self.counter >= self.patience
+
+
+@dataclass
+class TrainTestSplit:
+    """
+    Container holding the train/test split for heliostat reconstruction data.
+
+    Attributes
+    ----------
+    flux_measured_train : torch.Tensor
+        Measured flux distributions for the training set.
+        Shape is ``[number_of_train_samples_total, height, width]``.
+    focal_spots_measured_train : torch.Tensor
+        Measured focal spot coordinates for the training set.
+        Shape is ``[number_of_train_samples_total, 4]``.
+    incident_ray_directions_train : torch.Tensor
+        Incident ray directions for the training set.
+        Shape is ``[number_of_train_samples_total, 4]``.
+    motor_positions_train : torch.Tensor
+        Motor positions for the training set.
+        Shape is ``[number_of_train_samples_total, 2]``.
+    target_area_indices_train : torch.Tensor
+        Target area indices for the training set.
+        Shape is ``[number_of_train_samples_total]``.
+    flux_measured_test : torch.Tensor
+        Measured flux distributions for the test set.
+        Shape is ``[number_of_test_samples_total, height, width]``.
+    focal_spots_measured_test : torch.Tensor
+        Measured focal spot coordinates for the test set.
+        Shape is ``[number_of_test_samples_total, 4]``.
+    incident_ray_directions_test : torch.Tensor
+        Incident ray directions for the test set.
+        Shape is ``[number_of_test_samples_total, 4]``.
+    motor_positions_test : torch.Tensor
+        Motor positions for the test set.
+        Shape is ``[number_of_test_samples_total, 2]``.
+    target_area_indices_test : torch.Tensor
+        Target area indices for the test set.
+        Shape is ``[number_of_test_samples_total]``.
+    active_heliostats_mask_train : torch.Tensor
+        Mask for active training samples available per heliostat after splitting.
+        Shape is ``[number_of_heliostats]``.
+    active_heliostats_mask_test : torch.Tensor
+        Mask for active test samples available per heliostat after splitting.
+        Shape is ``[number_of_heliostats]``.
+    train_indices : torch.Tensor
+        Indices of training samples from the original dataset.
+        Shape is ``[number_of_train_samples_total]``.
+    test_indices : torch.Tensor
+        Indices of test samples from the original dataset.
+        Shape is ``[number_of_test_samples_total]``.
+    number_of_train_samples : int
+        Number of training samples per heliostat.
+    number_of_test_samples : int
+        Number of test samples per heliostat.
+    number_of_samples_per_heliostat : int
+        Total number of samples available per heliostat before splitting.
+    """
+
+    flux_measured_train: torch.Tensor
+    focal_spots_measured_train: torch.Tensor
+    incident_ray_directions_train: torch.Tensor
+    motor_positions_train: torch.Tensor
+    target_area_indices_train: torch.Tensor
+
+    flux_measured_test: torch.Tensor
+    focal_spots_measured_test: torch.Tensor
+    incident_ray_directions_test: torch.Tensor
+    motor_positions_test: torch.Tensor
+    target_area_indices_test: torch.Tensor
+
+    active_heliostats_mask_train: torch.Tensor
+    active_heliostats_mask_test: torch.Tensor
+
+    train_indices: torch.Tensor
+    test_indices: torch.Tensor
+
+    number_of_train_samples: int
+    number_of_test_samples: int
+    number_of_samples_per_heliostat: int
+
+
+def train_test_split(
+    active_heliostats_mask: torch.Tensor,
+    flux_measured: torch.Tensor,
+    focal_spots_measured: torch.Tensor,
+    incident_ray_directions: torch.Tensor,
+    motor_positions: torch.Tensor,
+    target_area_indices: torch.Tensor,
+    test_fraction: float = 0.25,
+    device: torch.device | None = None,
+) -> TrainTestSplit:
+    """
+    Split heliostat reconstruction data into training and test subsets.
+
+    The split is performed independently for each heliostat while preserving
+    the original ordering of samples. Training samples are taken from the
+    beginning of each heliostat block, while test samples are taken from the
+    end of each heliostat block.
+
+    Parameters
+    ----------
+    active_heliostats_mask : torch.Tensor
+        Mask for active samples available per heliostat.
+        Shape is ``[number_of_heliostats]``.
+    flux_measured : torch.Tensor
+        Measured flux distributions for all heliostats.
+        Shape is ``[total_number_of_samples, height, width]``.
+    focal_spots_measured : torch.Tensor
+        Measured focal spot coordinates for all samples.
+        Shape is ``[total_number_of_samples, 4]``.
+    incident_ray_directions : torch.Tensor
+        Incident ray directions for all samples.
+        Shape is ``[total_number_of_samples, 4]``.
+    motor_positions : torch.Tensor
+        Motor positions for all samples.
+        Shape is ``[total_number_of_samples, 2]``.
+    target_area_indices : torch.Tensor
+        Integer target area indices for all samples.
+        Shape is ``[total_number_of_samples]``.
+    test_fraction : float
+        Fraction of samples per heliostat assigned to the test set (default is 0.25).
+    device : torch.device | None
+        The device on which to perform computations or load tensors and models (default is None).
+        If None, ARTIST will automatically select the most appropriate
+        device (CUDA or CPU) based on availability and OS.
+
+    Returns
+    -------
+    TrainTestSplit
+        Dataclass containing: train/test tensors, train/test indices, updated active heliostat masks.
+    """
+    device = get_device(device=device)
+
+    total_samples = int(active_heliostats_mask.sum().item())
+    number_of_heliostats = int((active_heliostats_mask > 0).sum().item())
+    number_of_samples_per_heliostat = int(total_samples / number_of_heliostats)
+    number_of_test_samples = max(
+        1, int(number_of_samples_per_heliostat * test_fraction)
+    )
+    number_of_train_samples = number_of_samples_per_heliostat - number_of_test_samples
+    starts = torch.arange(number_of_heliostats) * number_of_samples_per_heliostat
+    offsets = torch.arange(number_of_train_samples, number_of_samples_per_heliostat)
+    train_indices = (
+        torch.arange(0, total_samples, number_of_samples_per_heliostat, device=device)[
+            :, None
+        ]
+        + torch.arange(number_of_train_samples, device=device)
+    ).reshape(-1)
+    test_indices = (starts[:, None] + offsets).reshape(-1)
+
+    active_heliostats_mask_train = torch.clamp(
+        active_heliostats_mask - number_of_test_samples, min=0
+    )
+    active_heliostats_mask_test = torch.clamp(
+        active_heliostats_mask - number_of_train_samples, min=0
+    )
+
+    return TrainTestSplit(
+        flux_measured_train=flux_measured[train_indices],
+        focal_spots_measured_train=focal_spots_measured[train_indices],
+        incident_ray_directions_train=incident_ray_directions[train_indices],
+        motor_positions_train=motor_positions[train_indices],
+        target_area_indices_train=target_area_indices[train_indices],
+        flux_measured_test=flux_measured[test_indices],
+        focal_spots_measured_test=focal_spots_measured[test_indices],
+        incident_ray_directions_test=incident_ray_directions[test_indices],
+        motor_positions_test=motor_positions[test_indices],
+        target_area_indices_test=target_area_indices[test_indices],
+        active_heliostats_mask_train=active_heliostats_mask_train,
+        active_heliostats_mask_test=active_heliostats_mask_test,
+        train_indices=train_indices,
+        test_indices=test_indices,
+        number_of_train_samples=number_of_train_samples,
+        number_of_test_samples=number_of_test_samples,
+        number_of_samples_per_heliostat=number_of_samples_per_heliostat,
+    )
diff --git a/artist/util/constants.py b/artist/util/constants.py
index 09ce5e04d..7841e646c 100644
--- a/artist/util/constants.py
+++ b/artist/util/constants.py
@@ -165,6 +165,8 @@
 
 kinematics_reconstruction_raytracing = "raytracing"
 """Defines that the kinematics reconstructor uses ray tracing."""
+kinematics_reconstruction_alignment = "alignment"
+"""Defines that the kinematics reconstructor uses alignment."""
 
 names = "names"
 """Key to access heliostat names."""
diff --git a/examples/field_optimizations/config.yaml b/examples/field_optimizations/config.yaml
index 4d44d85b7..81448801b 100644
--- a/examples/field_optimizations/config.yaml
+++ b/examples/field_optimizations/config.yaml
@@ -19,12 +19,12 @@ heliostats_for_plots:
 random_seed: 7
 device: "cuda"
 surface_reconstruction_optimization_configuration:
-  max_epoch: 200
+  max_epoch: 2
   batch_size: 48
   batch_size_outer: 32
   number_of_surface_points: 50
   number_of_control_points: 6
-  number_of_rays: 170
+  number_of_rays: 190
   nurbs_degree: 3
   sample_limit: 3
   initial_learning_rate: 1e-06
@@ -42,10 +42,10 @@ surface_reconstruction_optimization_configuration:
   weight_smoothness: 0.00
   weight_ideal_surface: 0.10
 kinematics_reconstruction_optimization_configuration:
-  max_epoch: 200
+  max_epoch: 2
   batch_size: 960
-  sample_limit: 18
-  number_of_rays: 7
+  sample_limit: 10
+  number_of_rays: 16
   initial_learning_rate_rotation_deviation: 6.7e-04
   initial_learning_rate_initial_angles: 9.29e-3
   initial_learning_rate_initial_stroke_length: 6.46e-2
@@ -59,9 +59,9 @@ kinematics_reconstruction_optimization_configuration:
   cooldown: 5
   gamma: 0.92
 aim_point_optimization_configuration:
-  max_epoch: 200
+  max_epoch: 2
   batch_size: 96
-  number_of_rays: 5
+  number_of_rays: 8
   initial_learning_rate: 1e-03
   scheduler: "cyclic"
   min_learning_rate: 1e-06
@@ -78,13 +78,13 @@ aim_point_optimization_configuration:
   max_flux_density: 1000000 # W/m^2
 basic_config:
   batch_size: 96
-  number_of_rays: 5
+  number_of_rays: 8
   dni: 500
   bitmap_resolution: [230, 276]
   baseline_incident_ray_direction: [0.411, 0.706, -0.576, 0.0]
   baseline_target_area: "receiver"
-  baseline_aim_point: [-0.19720931, -0.03458419, 54.929, 1.0]
-  log_step: 1
+  baseline_aim_point: [-0.19720931, -0.03458419, 55.349, 1.0] #[-0.19720931, -0.03458419, 54.929, 1.0]
+  log_step: 20
 heliostat_list_baseline:
   - "AG48"
   - "AG49"
diff --git a/examples/field_optimizations/generate_plots.py b/examples/field_optimizations/generate_plots.py
index 936b329ba..8bb2f867e 100644
--- a/examples/field_optimizations/generate_plots.py
+++ b/examples/field_optimizations/generate_plots.py
@@ -10,6 +10,7 @@
 from matplotlib.gridspec import GridSpec
 from scipy.stats import gaussian_kde
 
+from artist.flux import bitmap
 from artist.geometry import transforms
 from artist.optim.loss import KLDivergenceLoss
 from artist.util.env import get_device
@@ -105,24 +106,26 @@ def plot_surface_reconstruction_flux(
     plot_data = results["surface_reconstruction"]["flux_plot_data"]
     number_of_heliostats = len(plot_data)
     fig, axes = plt.subplots(number_of_heliostats, 7, figsize=(35, 15))
-    for index, heliostat_name in enumerate(plot_data):
+    fontsize = 28
+    for index in range(1):
+        heliostat_name = "AK54"
         heliostat_data = plot_data[heliostat_name]
         axes[index, 0].imshow(
             heliostat_data["fluxes"][0].cpu().detach(), cmap="inferno"
         )
-        axes[index, 0].set_title("Calibration Flux")
+        axes[index, 0].set_title("Calibration Flux", fontsize=fontsize)
         axes[index, 0].axis("off")
 
         axes[index, 1].imshow(
             heliostat_data["fluxes"][1].cpu().detach(), cmap="inferno"
         )
-        axes[index, 1].set_title("Surface not reconstructed")
+        axes[index, 1].set_title("Surface not reconstructed", fontsize=fontsize)
         axes[index, 1].axis("off")
 
         axes[index, 2].imshow(
             heliostat_data["fluxes"][2].cpu().detach(), cmap="inferno"
         )
-        axes[index, 2].set_title("Surface reconstructed")
+        axes[index, 2].set_title("Surface reconstructed", fontsize=fontsize)
         axes[index, 2].axis("off")
 
         reference_direction = torch.tensor([0.0, 0.0, 1.0], device=torch.device("cpu"))
@@ -164,13 +167,14 @@ def plot_surface_reconstruction_flux(
             vmin=0.0345,
             vmax=0.036,
         )
-        axes[index, 3].set_title("Deflectometry Points original")
+        axes[index, 3].set_title("Deflectometry Points", fontsize=fontsize)
         axes[index, 3].axis("off")
         axes[index, 3].set_aspect("equal", adjustable="box")
         cbar3 = fig.colorbar(
             sc3, ax=axes[index, 3], orientation="horizontal", fraction=0.046, pad=0.1
         )
-        cbar3.set_label("m")
+        cbar3.set_label("m", fontsize=fontsize)
+        cbar3.ax.tick_params(labelsize=23)
 
         sc4 = axes[index, 4].scatter(
             x=deflectometry_points_original[:, 0],
@@ -180,13 +184,14 @@ def plot_surface_reconstruction_flux(
             vmin=0.0,
             vmax=0.005,
         )
-        axes[index, 4].set_title("Deflectometry normals")
+        axes[index, 4].set_title("Deflectometry Normals", fontsize=fontsize)
         axes[index, 4].axis("off")
         axes[index, 4].set_aspect("equal", adjustable="box")
         cbar4 = fig.colorbar(
             sc4, ax=axes[index, 4], orientation="horizontal", fraction=0.046, pad=0.1
         )
-        cbar4.set_label("Angle (rad)")
+        cbar4.set_label("Angle (rad)", fontsize=fontsize)
+        cbar4.ax.tick_params(labelsize=23)
 
         # Process reconstructed data.
         points_uncanted = transforms.perform_canting(
@@ -217,13 +222,14 @@ def plot_surface_reconstruction_flux(
             vmin=0.0345,
             vmax=0.036,
         )
-        axes[index, 5].set_title("Reconstructed Surface (Points)")
+        axes[index, 5].set_title("Reconstructed Points", fontsize=fontsize)
         axes[index, 5].axis("off")
         axes[index, 5].set_aspect("equal", adjustable="box")
         cbar5 = fig.colorbar(
             sc5, ax=axes[index, 5], orientation="horizontal", fraction=0.046, pad=0.1
         )
-        cbar5.set_label("m")
+        cbar5.set_label("m", fontsize=fontsize)
+        cbar5.ax.tick_params(labelsize=23)
 
         sc6 = axes[index, 6].scatter(
             x=reconstructed_points[:, 0],
@@ -233,22 +239,23 @@ def plot_surface_reconstruction_flux(
             vmin=0.0,
             vmax=0.005,
         )
-        axes[index, 6].set_title("Reconstructed normals")
+        axes[index, 6].set_title("Reconstructed Normals", fontsize=fontsize)
         axes[index, 6].axis("off")
         axes[index, 6].set_aspect("equal", adjustable="box")
         cbar6 = fig.colorbar(
             sc6, ax=axes[index, 6], orientation="horizontal", fraction=0.046, pad=0.1
         )
-        cbar6.set_label("Angle (rad)")
+        cbar6.set_label("Angle (rad)", fontsize=fontsize)
+        cbar6.ax.tick_params(labelsize=23)
 
     save_dir = save_dir / f"run_{results_number}"
     save_dir.mkdir(parents=True, exist_ok=True)
-    filename = save_dir / "surface_reconstruction_flux.pdf"
+    filename = save_dir / "surface_reconstruction_flux.png"
     fig.tight_layout()
     fig.savefig(filename, dpi=300, bbox_inches="tight")
     plt.close(fig)
 
-    print(f"Saved surface reconstruction error distribution plot at: {filename}.")
+    print(f"Saved surface reconstruction flux plot at: {filename}.")
 
 
 def plot_surface_error_analysis(
@@ -438,6 +445,125 @@ def plot_surface_loss_history(
         print(f"Saved surface reconstruction loss history plot at: {filename}.")
 
 
+def plot_surface_validation(
+    results: dict[str, Any],
+    save_dir: pathlib.Path,
+    results_number: int,
+) -> None:
+    """
+    Plot the surface reconstruction validation.
+
+    Parameters
+    ----------
+    results : dict[str, Any]
+        Results of the study.
+    save_dir : pathlib.Path
+        Path to the location where the plots are saved.
+    results_number : int
+        Identifier of the results run.
+    """
+    plot_data = results["surface_reconstruction"]["validation_data"]
+
+    kl_divergence = KLDivergenceLoss()
+    kl_losses = []
+    mse_losses = []
+    l1_norm_losses = []
+
+    keys = list(plot_data.keys())
+
+    for k in keys:
+        measured_flux = torch.nn.functional.normalize(
+            plot_data[k]["measured_flux"],
+            p=1,
+            dim=(0, 1),
+            eps=1e-12,
+        )
+        artist_flux = torch.nn.functional.normalize(
+            plot_data[k]["artist_flux"],
+            p=1,
+            dim=(0, 1),
+            eps=1e-12,
+        )
+
+        kl = kl_divergence(
+            measured_flux.unsqueeze(0),
+            artist_flux.unsqueeze(0),
+            reduction_dimensions=(1, 2),
+        )
+        kl_losses.append(float(kl))
+
+        # MSE
+        mse = torch.nn.functional.mse_loss(measured_flux, artist_flux).item()
+        mse_losses.append(mse)
+
+        # normalized L1-like
+        l1_norm = (
+            torch.abs(artist_flux - measured_flux).sum() / measured_flux.sum()
+        ).item()
+        l1_norm_losses.append(l1_norm)
+
+    x = list(range(len(keys)))
+
+    plt.figure(figsize=(10, 5))
+    plt.scatter(x, kl_losses, label="KL Divergence")
+    plt.scatter(x, mse_losses, label="MSE")
+    plt.scatter(x, l1_norm_losses, label="L1 normalized")
+    plt.title("Loss comparison per heliostat")
+    plt.xlabel("Heliostat index")
+    plt.ylabel("Loss value")
+    plt.legend()
+    plt.tight_layout()
+
+    save_dir.mkdir(parents=True, exist_ok=True)
+    filename = (
+        save_dir / f"run_{results_number}/surface_reconstruction_validation_loss.pdf"
+    )
+    plt.tight_layout()
+    plt.savefig(filename, dpi=300, bbox_inches="tight")
+    plt.clf()
+
+    print(f"Saved surface reconstruction validation loss at: {filename}.")
+
+    selected_indices = list(range(0, len(keys), 1))
+    n_rows = len(selected_indices)
+
+    fig, axes = plt.subplots(
+        n_rows,
+        2,
+        figsize=(8, 3 * n_rows),
+    )
+
+    if n_rows == 1:
+        axes = [axes]
+
+    for row, idx in enumerate(selected_indices):
+        k = keys[idx]
+
+        measured_flux = plot_data[k]["measured_flux"].detach().cpu()
+        artist_flux = plot_data[k]["artist_flux"].detach().cpu()
+
+        ax1 = axes[row][0]
+        ax2 = axes[row][1]
+
+        _ = ax1.imshow(measured_flux)
+        ax1.set_title(f"Measured {k}")
+        ax1.axis("off")
+
+        _ = ax2.imshow(artist_flux)
+        ax2.set_title(f"Artist {k}")
+        ax2.axis("off")
+
+    save_dir.mkdir(parents=True, exist_ok=True)
+    filename = (
+        save_dir / f"run_{results_number}/surface_reconstruction_validation_fluxes.pdf"
+    )
+    fig.tight_layout()
+    fig.savefig(filename, dpi=300, bbox_inches="tight")
+    plt.close(fig)
+
+    print(f"Saved surface reconstruction validation fluxes at: {filename}.")
+
+
 def plot_kinematics_reconstruction_flux(
     results: dict[str, Any],
     save_dir: pathlib.Path,
@@ -786,6 +912,112 @@ def plot_kinematics_loss_history(
         print(f"Saved kinematics reconstructions loss history plot at: {filename}.")
 
 
+def plot_kinematics_validation(
+    results: dict[str, Any],
+    save_dir: pathlib.Path,
+    results_number: int,
+) -> None:
+    """
+    Plot the surface reconstruction validation.
+
+    Parameters
+    ----------
+    results : dict[str, Any]
+        Results of the study.
+    save_dir : pathlib.Path
+        Path to the location where the plots are saved.
+    results_number : int
+        Identifier of the results run.
+    """
+    plot_data = results["kinematics_reconstruction_with_reconstructed_surfaces"][
+        "validation_data"
+    ]
+
+    keys = list(plot_data.keys())
+
+    # for k in keys:
+    #     measured_flux = plot_data[k]["measured_flux"]
+    #     artist_flux = plot_data[k]["artist_flux"]
+
+    #     centers = utils.get_center_of_mass(
+    #         bitmaps=torch.stack([measured_flux, artist_flux]),
+    #     )
+
+    # x = list(range(len(keys)))
+    # plt.figure(figsize=(10, 5))
+    # plt.scatter(x, kl_losses, label="KL Divergence")
+    # plt.scatter(x, mse_losses, label="MSE")
+    # plt.scatter(x, l1_norm_losses, label="L1 normalized")
+    # plt.title("Loss comparison per heliostat")
+    # plt.xlabel("Heliostat index")
+    # plt.ylabel("Loss value")
+    # plt.legend()
+    # plt.tight_layout()
+
+    # save_dir.mkdir(parents=True, exist_ok=True)
+    # filename = (
+    #     save_dir
+    #     / f"run_{results_number}/kinematics_reconstruction_validation_loss.pdf"
+    # )
+    # plt.tight_layout()
+    # plt.savefig(filename, dpi=300, bbox_inches="tight")
+    # plt.clf()
+
+    # print(f"Saved kinematics reconstruction validation loss at: {filename}.")
+
+    selected_indices = list(range(0, len(keys), 1))
+    n_rows = len(selected_indices)
+
+    fig, axes = plt.subplots(
+        n_rows,
+        2,
+        figsize=(8, 3 * n_rows),
+    )
+
+    if n_rows == 1:
+        axes = [axes]
+
+    for row, idx in enumerate(selected_indices):
+        k = keys[idx]
+
+        measured_flux = plot_data[k]["measured_flux"]
+        artist_flux = plot_data[k]["artist_flux"]
+
+        centers = (
+            bitmap.get_center_of_mass(
+                bitmaps=torch.stack([measured_flux, artist_flux]),
+            )
+            .detach()
+            .cpu()
+        )
+
+        ax1 = axes[row][0]
+        ax2 = axes[row][1]
+
+        _ = ax1.imshow(measured_flux.detach().cpu())
+        ax1.set_title(f"Measured {k}")
+        ax1.scatter(x=centers[0, 0], y=centers[0, 1], c="black", marker="o", s=30)
+        ax1.scatter(x=centers[1, 0], y=centers[1, 1], c="red", marker="x", s=30)
+        ax1.axis("off")
+
+        _ = ax2.imshow(artist_flux.detach().cpu())
+        ax2.set_title(f"Artist {k}")
+        ax2.scatter(x=centers[0, 0], y=centers[0, 1], c="black", marker="o", s=30)
+        ax2.scatter(x=centers[1, 0], y=centers[1, 1], c="red", marker="x", s=30)
+        ax2.axis("off")
+
+    save_dir.mkdir(parents=True, exist_ok=True)
+    filename = (
+        save_dir
+        / f"run_{results_number}/kinematics_reconstruction_validation_fluxes.pdf"
+    )
+    fig.tight_layout()
+    fig.savefig(filename, dpi=300, bbox_inches="tight")
+    plt.close(fig)
+
+    print(f"Saved kinematics reconstruction validation fluxes at: {filename}.")
+
+
 def plot_model_reconstruction(
     results: dict[str, Any],
     save_dir: pathlib.Path,
@@ -914,7 +1146,7 @@ def plot_model_reconstruction(
     fig, axes = plt.subplots(
         1,
         n_images,
-        figsize=(4 * n_images, 5),
+        figsize=(4 * n_images, 8),
         gridspec_kw={"bottom": 0.15, "top": 0.95, "wspace": 0.05},
     )
     if n_images == 1:
@@ -928,56 +1160,57 @@ def plot_model_reconstruction(
         r"\textbf{Combined Reconstructions}",
     ]
 
-    points = {
-        "focal spot aim point": 256
-        - torch.tensor([104.7413, 139.8122], device="cuda:0"),
-        "aim point 'measured'": 256
-        - torch.tensor([112.9362, 111.4531], device="cuda:0"),
-    }
+    # height, width = measured_flux.shape
+    # points = {
+    #     "focal spot aim point": 256
+    #     - torch.tensor([104.7413, 139.8122], device="cuda:0"),
+    #     "aim point 'measured'": 256
+    #     - torch.tensor([112.9362, 111.4531], device="cuda:0"),
+    # }
 
-    measured = images[0]
-    y_grid, x_grid = torch.meshgrid(torch.arange(256), torch.arange(256), indexing="ij")
-    x_com_measured = (x_grid * measured).sum() / measured.sum()
-    y_com_measured = (y_grid * measured).sum() / measured.sum()
+    # measured = images[0]
+    # y_grid, x_grid = torch.meshgrid(torch.arange(256), torch.arange(256), indexing="ij")
+    # x_com_measured = (x_grid * measured).sum() / measured.sum()
+    # y_com_measured = (y_grid * measured).sum() / measured.sum()
 
     for idx, (ax, flux) in enumerate(zip(axes, images)):
         im = ax.imshow(flux, cmap=cmap, vmin=vmin, vmax=vmax)
         ax.axis("off")
         ax.set_title(titles[idx], fontsize=13)
-        ax.scatter(
-            x_com_measured,
-            y_com_measured,
-            s=80,
-            marker="x",
-            color="blue",
-            label="center of mass reference",
-        )
+        # ax.scatter(
+        #     x_com_measured,
+        #     y_com_measured,
+        #     s=80,
+        #     marker="x",
+        #     color="blue",
+        #     label="center of mass reference",
+        # )
 
-        y_grid, x_grid = torch.meshgrid(
-            torch.arange(256), torch.arange(256), indexing="ij"
-        )
-        x_com = (x_grid * flux).sum() / flux.sum()
-        y_com = (y_grid * flux).sum() / flux.sum()
-        ax.scatter(
-            x_com, y_com, s=80, marker="o", color="green", label="center of mass"
-        )
+        # y_grid, x_grid = torch.meshgrid(
+        #     torch.arange(256), torch.arange(256), indexing="ij"
+        # )
+        # x_com = (x_grid * flux).sum() / flux.sum()
+        # y_com = (y_grid * flux).sum() / flux.sum()
+        # ax.scatter(
+        #     x_com, y_com, s=80, marker="o", color="green", label="center of mass"
+        # )
 
-        ax.scatter(
-            256 / 2,
-            256 / 2,
-            s=40,
-            marker="x",
-            color="white",
-            label="geometric center of bitmap",
-        )
-        ax.scatter(
-            points["aim point 'measured'"][0].cpu(),
-            points["aim point 'measured'"][1].cpu(),
-            s=80,
-            marker="x",
-            color="white",
-            label="aim point from protocol",
-        )
+        # ax.scatter(
+        #     256 / 2,
+        #     256 / 2,
+        #     s=40,
+        #     marker="x",
+        #     color="white",
+        #     label="geometric center of bitmap",
+        # )
+        # ax.scatter(
+        #     points["aim point 'measured'"][0].cpu(),
+        #     points["aim point 'measured'"][1].cpu(),
+        #     s=80,
+        #     marker="x",
+        #     color="white",
+        #     label="aim point from protocol",
+        # )
 
         ax.legend(fontsize=9, loc="upper right")
 
@@ -1000,13 +1233,14 @@ def plot_model_reconstruction(
             fontsize=12,
         )
 
-    cbar_ax = fig.add_axes([0.15, 0.01, 0.7, 0.03])
+    fig.subplots_adjust(bottom=0.15)
+    cbar_ax = fig.add_axes([0.15, 0.05, 0.7, 0.03])
     cbar = fig.colorbar(im, cax=cbar_ax, orientation="horizontal")
     cbar.ax.tick_params(labelsize=14)
 
     fig.text(
         0.95,
-        0.02,
+        0.01,
         "M = Measured flux\nU = Unoptimized flux\nO = Optimized flux\nH = Homogeneous flux",
         ha="right",
         va="bottom",
@@ -1352,26 +1586,32 @@ def _make_abs(p: str | pathlib.Path) -> pathlib.Path:
         #     results=results, results_ftp=results_ftp, save_dir=plots_path
         # )
 
-        # plot_surface_reconstruction_flux(
+        # plot_kinematics_reconstruction_flux(
         #     results=results, save_dir=plots_path, results_number=results_number
         # )
-        plot_surface_error_analysis(
+        # plot_kinematics_error_analysis(
+        #     results=results,
+        #     save_dir=plots_path,
+        #     split_left=False,
+        #     results_number=results_number,
+        # )
+        # plot_kinematics_loss_history(
+        #     results=results, save_dir=plots_path, results_number=results_number
+        # )
+        # plot_kinematics_validation(
+        #     results=results, save_dir=plots_path, results_number=results_number
+        # )
+
+        plot_surface_reconstruction_flux(
             results=results, save_dir=plots_path, results_number=results_number
         )
-        plot_surface_loss_history(
+        plot_surface_error_analysis(
             results=results, save_dir=plots_path, results_number=results_number
         )
-
-        plot_kinematics_reconstruction_flux(
+        plot_surface_loss_history(
             results=results, save_dir=plots_path, results_number=results_number
         )
-        plot_kinematics_error_analysis(
-            results=results,
-            save_dir=plots_path,
-            split_left=False,
-            results_number=results_number,
-        )
-        plot_kinematics_loss_history(
+        plot_surface_validation(
             results=results, save_dir=plots_path, results_number=results_number
         )
 
diff --git a/examples/field_optimizations/generate_results.py b/examples/field_optimizations/generate_results.py
index e221722f8..ec81f7b72 100644
--- a/examples/field_optimizations/generate_results.py
+++ b/examples/field_optimizations/generate_results.py
@@ -15,9 +15,9 @@
 from artist.flux import bitmap
 from artist.io.calibration_parser import CalibrationDataParser
 from artist.io.paint_calibration_parser import PaintCalibrationDataParser
+from artist.optim.aim_point_optimizer import AimPointOptimizer
 from artist.optim.kinematics_reconstructor import KinematicsReconstructor
 from artist.optim.loss import FocalSpotLoss, KLDivergenceLoss
-from artist.optim.motor_position_optimizer import MotorPositionsOptimizer
 from artist.optim.surface_reconstructor import SurfaceReconstructor
 from artist.raytracing.heliostat_ray_tracer import HeliostatRayTracer
 from artist.scenario.scenario import Scenario
@@ -667,7 +667,7 @@ def aim_point_plots(
             )
             if label == "before":
                 heliostat_group.align_surfaces_with_incident_ray_directions(
-                    aim_points=aim_point,
+                    aim_points=aim_point[None],
                     incident_ray_directions=incident_ray_directions,
                     active_heliostats_mask=active_heliostats_mask,
                     device=device,
@@ -767,6 +767,9 @@ def full_field_optimizations(
     results_dict = cast(
         dict[str, dict[str, torch.Tensor]], torch.load(results_path, weights_only=False)
     )
+    control_points_path = (
+        results_path.parent / f"reconstructed_nurbs_control_points_{results_number}.pt"
+    )
 
     number_of_heliostat_groups = Scenario.get_number_of_heliostat_groups_from_hdf5(
         scenario_path=scenario_path
@@ -777,14 +780,6 @@ def full_field_optimizations(
         device=device,
     ) as ddp_setup:
         device = ddp_setup["device"]
-        control_points_path = (
-            results_path.parent / "reconstructed_nurbs_control_points.pt"
-        )
-
-        assert data_mappings is not None, "data_mappings must be provided."
-        assert surface_config is not None, "surface_config must be provided."
-        assert kinematics_config is not None, "kinematics_config must be provided."
-        assert aim_point_config is not None, "aim_point_config must be provided."
 
         bitmap_resolution = torch.tensor(
             basic_config["bitmap_resolution"], device=device
@@ -1131,13 +1126,13 @@ def full_field_optimizations(
         if ddp_setup["is_distributed"]:
             torch.distributed.barrier()
         kinematics_validation = kinematics_evaluation(
-            scenario=scenario_kinematics_ideal_surfaces,
+            scenario=scenario_kinematics,
             ddp_setup=ddp_setup,
             heliostat_data=data_mappings["kinematics_validation"],
             device=device,
         )
         kinematics_data_plot_after = kinematics_evaluation(
-            scenario=scenario_kinematics_ideal_surfaces,
+            scenario=scenario_kinematics,
             ddp_setup=ddp_setup,
             heliostat_data=data_mappings["kinematics_plot"],
             device=device,
@@ -1246,7 +1241,7 @@ def full_field_optimizations(
                 constants.max_flux_density: aim_point_config["max_flux_density"],
             },
         }
-        motor_positions_optimizer = MotorPositionsOptimizer(
+        aim_point_optimizer = AimPointOptimizer(
             ddp_setup=ddp_setup,
             scenario=scenario_aim_points,
             optimization_configuration=optimization_configuration_aim_points,
@@ -1260,7 +1255,7 @@ def full_field_optimizations(
             device=device,
         )
         aimpoint_optimization_final_loss, loss_history_aim_points, _, _, _ = (
-            motor_positions_optimizer.optimize(
+            aim_point_optimizer.optimize(
                 loss_definition=KLDivergenceLoss(), device=device
             )
         )
diff --git a/tests/data/expected_bitmaps_blocking/bitmaps_cpu.pt b/tests/data/expected_bitmaps_blocking/bitmaps_cpu.pt
index fcebe334d..8382ea30d 100644
Binary files a/tests/data/expected_bitmaps_blocking/bitmaps_cpu.pt and b/tests/data/expected_bitmaps_blocking/bitmaps_cpu.pt differ
diff --git a/tests/data/expected_bitmaps_blocking/bitmaps_cuda.pt b/tests/data/expected_bitmaps_blocking/bitmaps_cuda.pt
index 19c232fee..1e1b1e09c 100644
Binary files a/tests/data/expected_bitmaps_blocking/bitmaps_cuda.pt and b/tests/data/expected_bitmaps_blocking/bitmaps_cuda.pt differ
diff --git a/tests/data/expected_bitmaps_integration/test_scenario_paint_mix_ideal_prototype_deflectometry_cpu.pt b/tests/data/expected_bitmaps_integration/test_scenario_paint_mix_ideal_prototype_deflectometry_cpu.pt
index 4621d873a..ea2a004af 100644
Binary files a/tests/data/expected_bitmaps_integration/test_scenario_paint_mix_ideal_prototype_deflectometry_cpu.pt and b/tests/data/expected_bitmaps_integration/test_scenario_paint_mix_ideal_prototype_deflectometry_cpu.pt differ
diff --git a/tests/data/expected_bitmaps_integration/test_scenario_paint_mix_ideal_prototype_deflectometry_cuda.pt b/tests/data/expected_bitmaps_integration/test_scenario_paint_mix_ideal_prototype_deflectometry_cuda.pt
index 88aa68442..e751addcd 100644
Binary files a/tests/data/expected_bitmaps_integration/test_scenario_paint_mix_ideal_prototype_deflectometry_cuda.pt and b/tests/data/expected_bitmaps_integration/test_scenario_paint_mix_ideal_prototype_deflectometry_cuda.pt differ
diff --git a/tests/data/expected_bitmaps_integration/test_scenario_paint_single_heliostat_cpu.pt b/tests/data/expected_bitmaps_integration/test_scenario_paint_single_heliostat_cpu.pt
index 78ec0e6cf..1b092fea9 100644
Binary files a/tests/data/expected_bitmaps_integration/test_scenario_paint_single_heliostat_cpu.pt and b/tests/data/expected_bitmaps_integration/test_scenario_paint_single_heliostat_cpu.pt differ
diff --git a/tests/data/expected_bitmaps_integration/test_scenario_paint_single_heliostat_cuda.pt b/tests/data/expected_bitmaps_integration/test_scenario_paint_single_heliostat_cuda.pt
index 44c1831c8..4e4eb6458 100644
Binary files a/tests/data/expected_bitmaps_integration/test_scenario_paint_single_heliostat_cuda.pt and b/tests/data/expected_bitmaps_integration/test_scenario_paint_single_heliostat_cuda.pt differ
diff --git a/tests/data/expected_bitmaps_integration/test_scenario_stral_single_heliostat_cpu.pt b/tests/data/expected_bitmaps_integration/test_scenario_stral_single_heliostat_cpu.pt
index be8cebea6..983fbc296 100644
Binary files a/tests/data/expected_bitmaps_integration/test_scenario_stral_single_heliostat_cpu.pt and b/tests/data/expected_bitmaps_integration/test_scenario_stral_single_heliostat_cpu.pt differ
diff --git a/tests/data/expected_bitmaps_integration/test_scenario_stral_single_heliostat_cuda.pt b/tests/data/expected_bitmaps_integration/test_scenario_stral_single_heliostat_cuda.pt
index 3578b0476..182000eb9 100644
Binary files a/tests/data/expected_bitmaps_integration/test_scenario_stral_single_heliostat_cuda.pt and b/tests/data/expected_bitmaps_integration/test_scenario_stral_single_heliostat_cuda.pt differ
diff --git a/tests/data/expected_bitmaps_ray_extinction/bitmaps_cpu.pt b/tests/data/expected_bitmaps_ray_extinction/bitmaps_cpu.pt
index ce9d6f50d..1c17bb1c7 100644
Binary files a/tests/data/expected_bitmaps_ray_extinction/bitmaps_cpu.pt and b/tests/data/expected_bitmaps_ray_extinction/bitmaps_cpu.pt differ
diff --git a/tests/data/expected_bitmaps_ray_extinction/bitmaps_cuda.pt b/tests/data/expected_bitmaps_ray_extinction/bitmaps_cuda.pt
index 958b7914a..fedfc383a 100644
Binary files a/tests/data/expected_bitmaps_ray_extinction/bitmaps_cuda.pt and b/tests/data/expected_bitmaps_ray_extinction/bitmaps_cuda.pt differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_0_10_cpu.pt b/tests/data/expected_optimized_motor_positions/distribution_group_0_10_cpu.pt
deleted file mode 100644
index 544826323..000000000
Binary files a/tests/data/expected_optimized_motor_positions/distribution_group_0_10_cpu.pt and /dev/null differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_0_10_cuda.pt b/tests/data/expected_optimized_motor_positions/distribution_group_0_10_cuda.pt
deleted file mode 100644
index 8bb9d65e7..000000000
Binary files a/tests/data/expected_optimized_motor_positions/distribution_group_0_10_cuda.pt and /dev/null differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_0_30_cpu.pt b/tests/data/expected_optimized_motor_positions/distribution_group_0_30_cpu.pt
new file mode 100644
index 000000000..e784266a2
Binary files /dev/null and b/tests/data/expected_optimized_motor_positions/distribution_group_0_30_cpu.pt differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_0_30_cuda.pt b/tests/data/expected_optimized_motor_positions/distribution_group_0_30_cuda.pt
new file mode 100644
index 000000000..c3789536a
Binary files /dev/null and b/tests/data/expected_optimized_motor_positions/distribution_group_0_30_cuda.pt differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_0_50_cpu.pt b/tests/data/expected_optimized_motor_positions/distribution_group_0_50_cpu.pt
deleted file mode 100644
index 320074399..000000000
Binary files a/tests/data/expected_optimized_motor_positions/distribution_group_0_50_cpu.pt and /dev/null differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_0_50_cuda.pt b/tests/data/expected_optimized_motor_positions/distribution_group_0_50_cuda.pt
deleted file mode 100644
index f7b004cf9..000000000
Binary files a/tests/data/expected_optimized_motor_positions/distribution_group_0_50_cuda.pt and /dev/null differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_1_10_cpu.pt b/tests/data/expected_optimized_motor_positions/distribution_group_1_10_cpu.pt
deleted file mode 100644
index 5b5f41edb..000000000
Binary files a/tests/data/expected_optimized_motor_positions/distribution_group_1_10_cpu.pt and /dev/null differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_1_10_cuda.pt b/tests/data/expected_optimized_motor_positions/distribution_group_1_10_cuda.pt
deleted file mode 100644
index fe84b4ea6..000000000
Binary files a/tests/data/expected_optimized_motor_positions/distribution_group_1_10_cuda.pt and /dev/null differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_1_30_cpu.pt b/tests/data/expected_optimized_motor_positions/distribution_group_1_30_cpu.pt
new file mode 100644
index 000000000..729dee786
Binary files /dev/null and b/tests/data/expected_optimized_motor_positions/distribution_group_1_30_cpu.pt differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_1_30_cuda.pt b/tests/data/expected_optimized_motor_positions/distribution_group_1_30_cuda.pt
new file mode 100644
index 000000000..7bcf27c74
Binary files /dev/null and b/tests/data/expected_optimized_motor_positions/distribution_group_1_30_cuda.pt differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_1_50_cpu.pt b/tests/data/expected_optimized_motor_positions/distribution_group_1_50_cpu.pt
deleted file mode 100644
index f2521a050..000000000
Binary files a/tests/data/expected_optimized_motor_positions/distribution_group_1_50_cpu.pt and /dev/null differ
diff --git a/tests/data/expected_optimized_motor_positions/distribution_group_1_50_cuda.pt b/tests/data/expected_optimized_motor_positions/distribution_group_1_50_cuda.pt
deleted file mode 100644
index 09718bea5..000000000
Binary files a/tests/data/expected_optimized_motor_positions/distribution_group_1_50_cuda.pt and /dev/null differ
diff --git a/tests/data/expected_optimized_motor_positions/focal_spot_group_0_50_cpu.pt b/tests/data/expected_optimized_motor_positions/focal_spot_group_0_50_cpu.pt
deleted file mode 100644
index 5f0375af4..000000000
Binary files a/tests/data/expected_optimized_motor_positions/focal_spot_group_0_50_cpu.pt and /dev/null differ
diff --git a/tests/data/expected_optimized_motor_positions/focal_spot_group_0_50_cuda.pt b/tests/data/expected_optimized_motor_positions/focal_spot_group_0_50_cuda.pt
deleted file mode 100644
index 3f82c6cfe..000000000
Binary files a/tests/data/expected_optimized_motor_positions/focal_spot_group_0_50_cuda.pt and /dev/null differ
diff --git a/tests/data/expected_optimized_motor_positions/focal_spot_group_1_50_cpu.pt b/tests/data/expected_optimized_motor_positions/focal_spot_group_1_50_cpu.pt
deleted file mode 100644
index 2706abd52..000000000
Binary files a/tests/data/expected_optimized_motor_positions/focal_spot_group_1_50_cpu.pt and /dev/null differ
diff --git a/tests/data/expected_optimized_motor_positions/focal_spot_group_1_50_cuda.pt b/tests/data/expected_optimized_motor_positions/focal_spot_group_1_50_cuda.pt
deleted file mode 100644
index e31fb85f4..000000000
Binary files a/tests/data/expected_optimized_motor_positions/focal_spot_group_1_50_cuda.pt and /dev/null differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_0_10_cpu.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_0_10_cpu.pt
deleted file mode 100644
index 22b51df9c..000000000
Binary files a/tests/data/expected_reconstructed_kinematics_parameters/group_0_10_cpu.pt and /dev/null differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_0_10_cuda.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_0_10_cuda.pt
deleted file mode 100644
index 46cfaf65e..000000000
Binary files a/tests/data/expected_reconstructed_kinematics_parameters/group_0_10_cuda.pt and /dev/null differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_0_50_cpu.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_0_50_cpu.pt
deleted file mode 100644
index 347812962..000000000
Binary files a/tests/data/expected_reconstructed_kinematics_parameters/group_0_50_cpu.pt and /dev/null differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_0_50_cuda.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_0_50_cuda.pt
deleted file mode 100644
index 333e0f182..000000000
Binary files a/tests/data/expected_reconstructed_kinematics_parameters/group_0_50_cuda.pt and /dev/null differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_0_alignment_50_cpu.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_0_alignment_50_cpu.pt
new file mode 100644
index 000000000..c97e201ad
Binary files /dev/null and b/tests/data/expected_reconstructed_kinematics_parameters/group_0_alignment_50_cpu.pt differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_0_alignment_50_cuda.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_0_alignment_50_cuda.pt
new file mode 100644
index 000000000..deea3f9e5
Binary files /dev/null and b/tests/data/expected_reconstructed_kinematics_parameters/group_0_alignment_50_cuda.pt differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_0_raytracing_50_cpu.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_0_raytracing_50_cpu.pt
new file mode 100644
index 000000000..75993792d
Binary files /dev/null and b/tests/data/expected_reconstructed_kinematics_parameters/group_0_raytracing_50_cpu.pt differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_0_raytracing_50_cuda.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_0_raytracing_50_cuda.pt
new file mode 100644
index 000000000..4e4b8c1b3
Binary files /dev/null and b/tests/data/expected_reconstructed_kinematics_parameters/group_0_raytracing_50_cuda.pt differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_1_10_cpu.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_1_10_cpu.pt
deleted file mode 100644
index 8b686a532..000000000
Binary files a/tests/data/expected_reconstructed_kinematics_parameters/group_1_10_cpu.pt and /dev/null differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_1_10_cuda.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_1_10_cuda.pt
deleted file mode 100644
index 07aebeb7b..000000000
Binary files a/tests/data/expected_reconstructed_kinematics_parameters/group_1_10_cuda.pt and /dev/null differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_1_50_cpu.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_1_50_cpu.pt
deleted file mode 100644
index 18df89a0d..000000000
Binary files a/tests/data/expected_reconstructed_kinematics_parameters/group_1_50_cpu.pt and /dev/null differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_1_50_cuda.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_1_50_cuda.pt
deleted file mode 100644
index 74cedbaaa..000000000
Binary files a/tests/data/expected_reconstructed_kinematics_parameters/group_1_50_cuda.pt and /dev/null differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_1_alignment_50_cpu.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_1_alignment_50_cpu.pt
new file mode 100644
index 000000000..975da2a2c
Binary files /dev/null and b/tests/data/expected_reconstructed_kinematics_parameters/group_1_alignment_50_cpu.pt differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_1_alignment_50_cuda.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_1_alignment_50_cuda.pt
new file mode 100644
index 000000000..30f04037a
Binary files /dev/null and b/tests/data/expected_reconstructed_kinematics_parameters/group_1_alignment_50_cuda.pt differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_1_raytracing_50_cpu.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_1_raytracing_50_cpu.pt
new file mode 100644
index 000000000..53bc35700
Binary files /dev/null and b/tests/data/expected_reconstructed_kinematics_parameters/group_1_raytracing_50_cpu.pt differ
diff --git a/tests/data/expected_reconstructed_kinematics_parameters/group_1_raytracing_50_cuda.pt b/tests/data/expected_reconstructed_kinematics_parameters/group_1_raytracing_50_cuda.pt
new file mode 100644
index 000000000..de3526246
Binary files /dev/null and b/tests/data/expected_reconstructed_kinematics_parameters/group_1_raytracing_50_cuda.pt differ
diff --git a/tests/data/expected_reconstructed_surfaces/kl_divergence_group_1_40_cpu.pt b/tests/data/expected_reconstructed_surfaces/kl_divergence_group_1_40_cpu.pt
index 88552dcac..5b0e38598 100644
Binary files a/tests/data/expected_reconstructed_surfaces/kl_divergence_group_1_40_cpu.pt and b/tests/data/expected_reconstructed_surfaces/kl_divergence_group_1_40_cpu.pt differ
diff --git a/tests/data/expected_reconstructed_surfaces/kl_divergence_group_1_40_cuda.pt b/tests/data/expected_reconstructed_surfaces/kl_divergence_group_1_40_cuda.pt
index ed4ae29c9..f2b67d952 100644
Binary files a/tests/data/expected_reconstructed_surfaces/kl_divergence_group_1_40_cuda.pt and b/tests/data/expected_reconstructed_surfaces/kl_divergence_group_1_40_cuda.pt differ
diff --git a/tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_20_cpu.pt b/tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_10_cpu.pt
similarity index 89%
rename from tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_20_cpu.pt
rename to tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_10_cpu.pt
index 55d192edd..488567797 100644
Binary files a/tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_20_cpu.pt and b/tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_10_cpu.pt differ
diff --git a/tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_20_cuda.pt b/tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_10_cuda.pt
similarity index 89%
rename from tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_20_cuda.pt
rename to tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_10_cuda.pt
index 32a2d8c39..86a0f0e0e 100644
Binary files a/tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_20_cuda.pt and b/tests/data/expected_reconstructed_surfaces/pixel_loss_group_0_10_cuda.pt differ
diff --git a/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_10_cpu.pt b/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_10_cpu.pt
new file mode 100644
index 000000000..5fbaa3796
Binary files /dev/null and b/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_10_cpu.pt differ
diff --git a/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_10_cuda.pt b/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_10_cuda.pt
new file mode 100644
index 000000000..7a3744f47
Binary files /dev/null and b/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_10_cuda.pt differ
diff --git a/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_20_cpu.pt b/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_20_cpu.pt
deleted file mode 100644
index 1537cdc29..000000000
Binary files a/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_20_cpu.pt and /dev/null differ
diff --git a/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_20_cuda.pt b/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_20_cuda.pt
deleted file mode 100644
index 63c3024c2..000000000
Binary files a/tests/data/expected_reconstructed_surfaces/pixel_loss_group_1_20_cuda.pt and /dev/null differ
diff --git a/tests/data/field_data/AA39-calibration-properties_2.json b/tests/data/field_data/AA39-calibration-properties_2.json
index aa7ef5981..2077848b7 100644
--- a/tests/data/field_data/AA39-calibration-properties_2.json
+++ b/tests/data/field_data/AA39-calibration-properties_2.json
@@ -1,21 +1,21 @@
 {
     "motor_position": {
-        "axis_1_motor_position": 20634,
-        "axis_2_motor_position": 40816
+        "axis_1_motor_position": 25702,
+        "axis_2_motor_position": 33826
     },
-    "target_name": "multi_focus_tower",
-    "sun_elevation": 48.68939631461994,
-    "sun_azimuth": 62.865328091974916,
+    "target_name": "solar_tower_juelich_lower",
+    "sun_elevation": 43.81453437734031,
+    "sun_azimuth": 68.52002947921484,
     "focal_spot": {
         "HeliOS": [
-            50.91339643166456,
-            6.387580043820685,
-            138.3399014968552
+            50.91339201932565,
+            6.3878282653934475,
+            122.50069327110258
         ],
         "UTIS": [
-            50.91339643109615,
-            6.387581065254336,
-            138.3357048034668
+            50.91339202742253,
+            6.3878294207790125,
+            122.50684356689452
         ]
     }
 }
diff --git a/tests/data/field_data/AA39-calibration-properties_3.json b/tests/data/field_data/AA39-calibration-properties_3.json
new file mode 100644
index 000000000..aa7ef5981
--- /dev/null
+++ b/tests/data/field_data/AA39-calibration-properties_3.json
@@ -0,0 +1,21 @@
+{
+    "motor_position": {
+        "axis_1_motor_position": 20634,
+        "axis_2_motor_position": 40816
+    },
+    "target_name": "multi_focus_tower",
+    "sun_elevation": 48.68939631461994,
+    "sun_azimuth": 62.865328091974916,
+    "focal_spot": {
+        "HeliOS": [
+            50.91339643166456,
+            6.387580043820685,
+            138.3399014968552
+        ],
+        "UTIS": [
+            50.91339643109615,
+            6.387581065254336,
+            138.3357048034668
+        ]
+    }
+}
diff --git a/tests/data/field_data/AA39-flux-centered_2.png b/tests/data/field_data/AA39-flux-centered_2.png
index 6f227094f..a74bb8919 100644
Binary files a/tests/data/field_data/AA39-flux-centered_2.png and b/tests/data/field_data/AA39-flux-centered_2.png differ
diff --git a/tests/data/field_data/AA39-flux-centered_3.png b/tests/data/field_data/AA39-flux-centered_3.png
new file mode 100644
index 000000000..6f227094f
Binary files /dev/null and b/tests/data/field_data/AA39-flux-centered_3.png differ
diff --git a/tests/data/field_data/AA39-flux_2.png b/tests/data/field_data/AA39-flux_2.png
index 404a97aa2..bde65ef4e 100644
Binary files a/tests/data/field_data/AA39-flux_2.png and b/tests/data/field_data/AA39-flux_2.png differ
diff --git a/tests/data/field_data/AA39-flux_3.png b/tests/data/field_data/AA39-flux_3.png
new file mode 100644
index 000000000..404a97aa2
Binary files /dev/null and b/tests/data/field_data/AA39-flux_3.png differ
diff --git a/tests/field/test_kinematics.py b/tests/field/test_kinematics.py
index 50090307c..289f03705 100644
--- a/tests/field/test_kinematics.py
+++ b/tests/field/test_kinematics.py
@@ -11,22 +11,22 @@
             [
                 [
                     [
-                        9.999455809593e-01,
-                        -4.559969624118e-10,
-                        -1.043199468404e-02,
-                        -1.852722256444e-03,
+                        1.000000000000e00,
+                        7.034584910024e-14,
+                        1.609325408936e-06,
+                        -2.858161849417e-07,
                     ],
                     [
-                        -7.413947489113e-03,
-                        7.035020589828e-01,
-                        -7.106545567513e-01,
-                        -1.891756802797e-01,
+                        1.130841724262e-06,
+                        7.115054130554e-01,
+                        -7.026806473732e-01,
+                        1.247960701585e-01,
                     ],
                     [
-                        7.338930387050e-03,
-                        7.106932401657e-01,
-                        7.034637928009e-01,
-                        6.132812798023e-02,
+                        -1.145043825090e-06,
+                        7.026806473732e-01,
+                        7.115054130554e-01,
+                        -1.263633668423e-01,
                     ],
                     [
                         0.000000000000e00,
@@ -66,10 +66,10 @@ def test_kinematics_forward(expected: torch.Tensor, device: torch.device) -> Non
                 0.0000,
                 0.0000,
                 0.0000,
-                0.3150,
                 0.0000,
-                -0.1776,
-                -0.4045,
+                0.0000,
+                0.1776,
+                0.0000,
             ]
         ],
         device=device,
diff --git a/tests/field/test_kinematics_rigid_body.py b/tests/field/test_kinematics_rigid_body.py
index d6bf3ef06..37ceb41a8 100644
--- a/tests/field/test_kinematics_rigid_body.py
+++ b/tests/field/test_kinematics_rigid_body.py
@@ -63,9 +63,7 @@ def kinematics_parameters(device: torch.device) -> tuple[torch.Tensor, torch.Ten
     """
     translation_deviation_parameters = torch.zeros((6, 9), device=device)
 
-    translation_deviation_parameters[:, 5] = 0.315
-    translation_deviation_parameters[:, 7] = -0.17755
-    translation_deviation_parameters[:, 8] = -0.4045
+    translation_deviation_parameters[:, 7] = 0.175
 
     rotation_deviation_parameters = torch.zeros((6, 4), device=device)
 
@@ -107,6 +105,8 @@ def kinematics_model_linear(
 
     actuator_parameters_non_optimizable = torch.zeros((6, 7, 2), device=device)
     actuator_parameters_non_optimizable[:, 1, 1] = 1
+    actuator_parameters_non_optimizable[:, 3, 0] = 68745
+    actuator_parameters_non_optimizable[:, 3, 1] = 75308
     actuator_parameters_non_optimizable[:, 4, 0] = 154166.666
     actuator_parameters_non_optimizable[:, 4, 1] = 154166.666
     actuator_parameters_non_optimizable[:, 5, 0] = 0.34061
@@ -182,6 +182,8 @@ def kinematics_model_ideal_1(
     actuator_parameters_non_optimizable = torch.zeros((10, 4, 2), device=device)
     actuator_parameters_non_optimizable[:, 0, :] = 1
     actuator_parameters_non_optimizable[:, 1, 1] = 1
+    actuator_parameters_non_optimizable[:, 2, :] = -torch.pi
+    actuator_parameters_non_optimizable[:, 3, :] = torch.pi
 
     return RigidBody(
         number_of_heliostats=10,
@@ -272,22 +274,22 @@ def kinematics_model_ideal_2(
                 [
                     [
                         [
-                            9.999456405640e-01,
-                            -4.558640964714e-10,
-                            -1.042895484716e-02,
-                            -1.851661014371e-03,
+                            1.000000000000e00,
+                            6.774044351713e-14,
+                            1.549720764160e-06,
+                            -2.712011166750e-07,
                         ],
                         [
-                            -7.411775179207e-03,
-                            7.035031914711e-01,
-                            -7.106534838676e-01,
-                            -1.891400665045e-01,
+                            1.089058969228e-06,
+                            7.114415168762e-01,
+                            -7.027453184128e-01,
+                            1.229804158211e-01,
                         ],
                         [
-                            7.336803711951e-03,
-                            7.106921076775e-01,
-                            7.034649252892e-01,
-                            6.129327416420e-02,
+                            -1.102535748032e-06,
+                            7.027453184128e-01,
+                            7.114415168762e-01,
+                            -1.245022714138e-01,
                         ],
                         [
                             0.000000000000e00,
@@ -298,22 +300,22 @@ def kinematics_model_ideal_2(
                     ],
                     [
                         [
-                            7.122755050659e-01,
-                            3.068102216730e-08,
-                            7.019000053406e-01,
-                            1.246223449707e-01,
+                            7.027451992035e-01,
+                            3.109810009505e-08,
+                            7.114416360855e-01,
+                            -1.245022863150e-01,
                         ],
                         [
-                            7.018629312515e-01,
-                            -1.027521397918e-02,
-                            -7.122378945351e-01,
-                            -1.255382150412e-01,
+                            7.114416360855e-01,
+                            1.399793518431e-06,
+                            -7.027451992035e-01,
+                            1.229804083705e-01,
                         ],
                         [
-                            7.212151307613e-03,
-                            9.999471902847e-01,
-                            -7.318804971874e-03,
-                            -9.079471230507e-02,
+                            -1.017725480779e-06,
+                            1.000000000000e00,
+                            9.615737326385e-07,
+                            -1.759248817734e-07,
                         ],
                         [
                             0.000000000000e00,
@@ -324,22 +326,22 @@ def kinematics_model_ideal_2(
                     ],
                     [
                         [
-                            9.999730587006e-01,
-                            -3.207012433393e-10,
-                            -7.336789276451e-03,
-                            -1.302646938711e-03,
+                            1.000000000000e00,
+                            2.605401856173e-14,
+                            5.960464477539e-07,
+                            -1.043081283569e-07,
                         ],
                         [
-                            -7.336692418903e-03,
-                            -5.136868450791e-03,
-                            -9.999598860741e-01,
-                            -1.770831346512e-01,
+                            5.960464477539e-07,
+                            5.523350523617e-07,
+                            -1.000000000000e00,
+                            1.749999970198e-01,
                         ],
                         [
-                            -3.768780152313e-05,
-                            9.999868273735e-01,
-                            -5.136730149388e-03,
-                            -9.041082859039e-02,
+                            -3.552713678801e-13,
+                            1.000000000000e00,
+                            5.523350523617e-07,
+                            -1.043081283569e-07,
                         ],
                         [
                             0.000000000000e00,
@@ -350,22 +352,22 @@ def kinematics_model_ideal_2(
                     ],
                     [
                         [
-                            7.018998265266e-01,
-                            -3.113456159554e-08,
-                            -7.122757434845e-01,
-                            -1.264645606279e-01,
+                            7.027456760406e-01,
+                            -3.109808233148e-08,
+                            -7.114411592484e-01,
+                            1.245022043586e-01,
                         ],
                         [
-                            -7.122381329536e-01,
-                            -1.027521397918e-02,
-                            -7.018627524376e-01,
-                            -1.236961036921e-01,
+                            -7.114411592484e-01,
+                            1.757421387083e-06,
+                            -7.027456760406e-01,
+                            1.229804903269e-01,
                         ],
                         [
-                            -7.318763993680e-03,
-                            9.999471902847e-01,
-                            -7.212193217129e-03,
-                            -9.077578783035e-02,
+                            1.272155941479e-06,
+                            1.000000000000e00,
+                            1.212895881508e-06,
+                            -2.199062549835e-07,
                         ],
                         [
                             0.000000000000e00,
@@ -376,22 +378,22 @@ def kinematics_model_ideal_2(
                     ],
                     [
                         [
-                            9.999683499336e-01,
-                            -3.478643761934e-10,
-                            -7.958209142089e-03,
-                            -1.412980025634e-03,
+                            1.000000000000e00,
+                            3.126481956358e-14,
+                            7.152557373047e-07,
+                            -1.251697483440e-07,
                         ],
                         [
-                            -7.367914076895e-03,
-                            3.779509067535e-01,
-                            -9.257963299751e-01,
-                            -1.982017606497e-01,
+                            6.598978075090e-07,
+                            3.857483565807e-01,
+                            -9.226040244102e-01,
+                            1.614557057619e-01,
                         ],
                         [
-                            3.007812658325e-03,
-                            9.258256554604e-01,
-                            3.779389560223e-01,
-                            -1.575836539268e-02,
+                            -2.759087465165e-07,
+                            9.226040244102e-01,
+                            3.857483565807e-01,
+                            -6.750596314669e-02,
                         ],
                         [
                             0.000000000000e00,
@@ -402,22 +404,22 @@ def kinematics_model_ideal_2(
                     ],
                     [
                         [
-                            9.999456405640e-01,
-                            -4.558640964714e-10,
-                            -1.042895484716e-02,
-                            -1.851661014371e-03,
+                            1.000000000000e00,
+                            6.774044351713e-14,
+                            1.549720764160e-06,
+                            -2.712011166750e-07,
                         ],
                         [
-                            -7.411775179207e-03,
-                            7.035031914711e-01,
-                            -7.106534838676e-01,
-                            8.108599185944e-01,
+                            1.089058969228e-06,
+                            7.114415168762e-01,
+                            -7.027453184128e-01,
+                            1.122980356216e00,
                         ],
                         [
-                            7.336803711951e-03,
-                            7.106921076775e-01,
-                            7.034649252892e-01,
-                            6.129327416420e-02,
+                            -1.102535748032e-06,
+                            7.027453184128e-01,
+                            7.114415168762e-01,
+                            -1.245022714138e-01,
                         ],
                         [
                             0.000000000000e00,
@@ -752,19 +754,19 @@ def test_incident_ray_direction_to_orientation(
                             5.735765099525e-01,
                             3.580627350175e-08,
                             8.191520571709e-01,
-                            1.454404443502e-01,
+                            -1.433516144753e-01,
                         ],
                         [
                             2.571453308065e-07,
                             1.000000000000e00,
                             -2.237665057692e-07,
-                            -8.950003981590e-02,
+                            3.150964289489e-08,
                         ],
                         [
                             -8.191520571709e-01,
                             3.389882863303e-07,
                             5.735765099525e-01,
-                            1.018384844065e-01,
+                            -1.003758907318e-01,
                         ],
                         [
                             0.000000000000e00,
@@ -778,19 +780,19 @@ def test_incident_ray_direction_to_orientation(
                             5.922092795372e-01,
                             3.522194802486e-08,
                             8.057842254639e-01,
-                            1.430669873953e-01,
+                            -1.410122364759e-01,
                         ],
                         [
                             8.305405208375e-05,
                             1.000000000000e00,
                             -6.108410161687e-05,
-                            -8.951085805893e-02,
+                            1.068206802302e-05,
                         ],
                         [
                             -8.057842254639e-01,
                             1.030982093653e-04,
                             5.922092795372e-01,
-                            1.051375344396e-01,
+                            -1.036366224289e-01,
                         ],
                         [
                             0.000000000000e00,
@@ -804,19 +806,19 @@ def test_incident_ray_direction_to_orientation(
                             5.743377208710e-01,
                             3.578294993645e-08,
                             8.186184763908e-01,
-                            1.453457176685e-01,
+                            -1.432582288980e-01,
                         ],
                         [
                             1.681064895820e-04,
                             1.000000000000e00,
                             -1.179862010758e-04,
-                            -8.952096104622e-02,
+                            2.063993451884e-05,
                         ],
                         [
                             -8.186184763908e-01,
                             2.053789939964e-04,
                             5.743377208710e-01,
-                            1.019552871585e-01,
+                            -1.005090996623e-01,
                         ],
                         [
                             0.000000000000e00,
@@ -830,19 +832,19 @@ def test_incident_ray_direction_to_orientation(
                             5.735765099525e-01,
                             3.580627350175e-08,
                             8.191520571709e-01,
-                            1.454404443502e-01,
+                            -1.433516144753e-01,
                         ],
                         [
                             2.571453308065e-07,
                             1.000000000000e00,
                             -2.237665057692e-07,
-                            -8.950003981590e-02,
+                            3.150964289489e-08,
                         ],
                         [
                             -8.191520571709e-01,
                             3.389882863303e-07,
                             5.735765099525e-01,
-                            1.018384844065e-01,
+                            -1.003758907318e-01,
                         ],
                         [
                             0.000000000000e00,
@@ -856,19 +858,19 @@ def test_incident_ray_direction_to_orientation(
                             5.922092795372e-01,
                             3.522194802486e-08,
                             8.057842254639e-01,
-                            1.430669873953e-01,
+                            -1.410122364759e-01,
                         ],
                         [
                             8.305405208375e-05,
                             1.000000000000e00,
                             -6.108410161687e-05,
-                            -8.951085805893e-02,
+                            1.068206802302e-05,
                         ],
                         [
                             -8.057842254639e-01,
                             1.030982093653e-04,
                             5.922092795372e-01,
-                            1.051375344396e-01,
+                            -1.036366224289e-01,
                         ],
                         [
                             0.000000000000e00,
@@ -882,19 +884,19 @@ def test_incident_ray_direction_to_orientation(
                             5.743377208710e-01,
                             3.578294993645e-08,
                             8.186184763908e-01,
-                            1.453457176685e-01,
+                            -1.432582288980e-01,
                         ],
                         [
                             1.681064895820e-04,
                             1.000000000000e00,
                             -1.179862010758e-04,
-                            9.104790687561e-01,
+                            1.000020623207e00,
                         ],
                         [
                             -8.186184763908e-01,
                             2.053789939964e-04,
                             5.743377208710e-01,
-                            1.019552871585e-01,
+                            -1.005090996623e-01,
                         ],
                         [
                             0.000000000000e00,
diff --git a/tests/io/test_paint_calibration_parser.py b/tests/io/test_paint_calibration_parser.py
index e627fbcc9..c16e5c6a3 100644
--- a/tests/io/test_paint_calibration_parser.py
+++ b/tests/io/test_paint_calibration_parser.py
@@ -19,7 +19,7 @@
                         pathlib.Path(ARTIST_ROOT)
                         / "tests/data/field_data/AA39-calibration-properties_1.json",
                         pathlib.Path(ARTIST_ROOT)
-                        / "tests/data/field_data/AA39-calibration-properties_2.json",
+                        / "tests/data/field_data/AA39-calibration-properties_3.json",
                     ],
                 )
             ],
@@ -55,7 +55,7 @@
                         pathlib.Path(ARTIST_ROOT)
                         / "tests/data/field_data/AA39-calibration-properties_1.json",
                         pathlib.Path(ARTIST_ROOT)
-                        / "tests/data/field_data/AA39-calibration-properties_2.json",
+                        / "tests/data/field_data/AA39-calibration-properties_3.json",
                     ],
                 )
             ],
@@ -110,7 +110,7 @@
                         pathlib.Path(ARTIST_ROOT)
                         / "tests/data/field_data/AA39-calibration-properties_1.json",
                         pathlib.Path(ARTIST_ROOT)
-                        / "tests/data/field_data/AA39-calibration-properties_2.json",
+                        / "tests/data/field_data/AA39-calibration-properties_3.json",
                     ],
                 )
             ],
diff --git a/tests/optim/test_motor_position_optimizer.py b/tests/optim/test_aim_point_optimizer.py
similarity index 82%
rename from tests/optim/test_motor_position_optimizer.py
rename to tests/optim/test_aim_point_optimizer.py
index b6755e111..fd2b37434 100644
--- a/tests/optim/test_motor_position_optimizer.py
+++ b/tests/optim/test_aim_point_optimizer.py
@@ -5,8 +5,8 @@
 import torch
 
 from artist import ARTIST_ROOT
-from artist.optim import MotorPositionsOptimizer
-from artist.optim.loss import FocalSpotLoss, KLDivergenceLoss, Loss
+from artist.optim import AimPointOptimizer
+from artist.optim.loss import KLDivergenceLoss, Loss
 from artist.scenario import Scenario
 from artist.util import constants
 from artist.util.env import DdpSetup
@@ -57,15 +57,13 @@ def distribution(device: torch.device) -> torch.Tensor:
 
 
 @pytest.mark.parametrize(
-    "loss_class, ground_truth_fixture_name, early_stopping_window, scheduler",
+    "loss, ground_truth_fixture_name, early_stopping_window, scheduler",
     [
-        (FocalSpotLoss, "focal_spot", 50, constants.cyclic),
-        (KLDivergenceLoss, "distribution", 50, constants.reduce_on_plateau),
-        (KLDivergenceLoss, "distribution", 10, constants.reduce_on_plateau),
+        (KLDivergenceLoss(), "distribution", 30, constants.reduce_on_plateau),
     ],
 )
-def test_motor_positions_optimizer(
-    loss_class: Loss,
+def test_aim_point_optimizer(
+    loss: Loss,
     ground_truth_fixture_name: str,
     early_stopping_window: int,
     scheduler: str,
@@ -74,12 +72,12 @@ def test_motor_positions_optimizer(
     device: torch.device,
 ) -> None:
     """
-    Test the motor positions optimizer.
+    Test the aim point optimizer.
 
     Parameters
     ----------
-    loss_class : Loss
-        The loss class.
+    loss : Loss
+        The loss function.
     ground_truth_fixture_name : str
         A fixture to retrieve the ground truth.
     early_stopping_window : int
@@ -114,7 +112,7 @@ def test_motor_positions_optimizer(
     optimizer_dict = {
         constants.initial_learning_rate: 1e-3,
         constants.tolerance: 0.0005,
-        constants.max_epoch: 50,
+        constants.max_epoch: 40,
         constants.batch_size: 50,
         constants.log_step: 1,
         constants.early_stopping_delta: 1.0,
@@ -145,8 +143,8 @@ def test_motor_positions_optimizer(
 
     ddp_setup_for_testing["device"] = device
 
-    # Create the motor positions optimizer.
-    motor_positions_optimizer = MotorPositionsOptimizer(
+    # Create the aim point optimizer.
+    aim_point_optimizer = AimPointOptimizer(
         ddp_setup=ddp_setup_for_testing,
         scenario=scenario,
         optimization_configuration=optimization_configuration,
@@ -158,16 +156,8 @@ def test_motor_positions_optimizer(
         device=device,
     )
 
-    loss_definition = (
-        FocalSpotLoss(scenario=scenario)
-        if loss_class is FocalSpotLoss
-        else KLDivergenceLoss()
-    )
-
-    # Optimize the motor positions.
-    _, _, _, _, _ = motor_positions_optimizer.optimize(
-        loss_definition=loss_definition, device=device
-    )
+    # Optimize the aim points.
+    _, _, _, _, _ = aim_point_optimizer.optimize(loss_definition=loss, device=device)
 
     for index, heliostat_group in enumerate(scenario.heliostat_field.heliostat_groups):
         expected_path = (
@@ -176,6 +166,8 @@ def test_motor_positions_optimizer(
             / f"{ground_truth_fixture_name}_group_{index}_{early_stopping_window}_{device.type}.pt"
         )
 
+        torch.save(heliostat_group.kinematics.motor_positions, expected_path)
+
         expected = torch.load(expected_path, map_location=device, weights_only=True)
 
         torch.testing.assert_close(
diff --git a/tests/optim/test_kinematics_reconstructor.py b/tests/optim/test_kinematics_reconstructor.py
deleted file mode 100644
index 8d4fff94a..000000000
--- a/tests/optim/test_kinematics_reconstructor.py
+++ /dev/null
@@ -1,259 +0,0 @@
-import pathlib
-
-import h5py
-import paint.util.paint_mappings as paint_mappings
-import pytest
-import torch
-
-from artist import ARTIST_ROOT
-from artist.io import CalibrationDataParser, PaintCalibrationDataParser
-from artist.optim import KinematicsReconstructor
-from artist.optim.loss import FocalSpotLoss
-from artist.scenario import Scenario
-from artist.util import constants
-from artist.util.env import DdpSetup
-
-
-@pytest.mark.parametrize(
-    "reconstruction_method, data_parser, centroid_extraction_method, early_stopping_window, scheduler",
-    [
-        # Test normal behavior.
-        (
-            constants.kinematics_reconstruction_raytracing,
-            PaintCalibrationDataParser(),
-            paint_mappings.UTIS_KEY,
-            50,
-            constants.exponential,
-        ),
-        # Test early stopping.
-        (
-            constants.kinematics_reconstruction_raytracing,
-            PaintCalibrationDataParser(),
-            paint_mappings.UTIS_KEY,
-            10,
-            constants.reduce_on_plateau,
-        ),
-        # Test invalid centroid extraction.
-        (
-            constants.kinematics_reconstruction_raytracing,
-            PaintCalibrationDataParser(),
-            "invalid",
-            10,
-            constants.reduce_on_plateau,
-        ),
-        # Test invalid reconstruction method.
-        (
-            "invalid",
-            PaintCalibrationDataParser(),
-            paint_mappings.UTIS_KEY,
-            10,
-            constants.reduce_on_plateau,
-        ),
-        # Test invalid parser.
-        (
-            constants.kinematics_reconstruction_raytracing,
-            CalibrationDataParser(),
-            paint_mappings.UTIS_KEY,
-            10,
-            constants.reduce_on_plateau,
-        ),
-    ],
-)
-def test_kinematics_reconstructor(
-    reconstruction_method: str,
-    data_parser: CalibrationDataParser,
-    centroid_extraction_method: str,
-    early_stopping_window: int,
-    scheduler: str,
-    ddp_setup_for_testing: DdpSetup,
-    device: torch.device,
-) -> None:
-    """
-    Test the kinematics reconstruction methods.
-
-    Parameters
-    ----------
-    reconstruction_method : str
-        The name of the reconstruction method.
-    data_parser : CalibrationDataParser
-        The data parser used to load calibration data from files.
-    centroid_extraction_method : str
-        The method used to extract the focal spot centroids.
-    early_stopping_window : int
-        Early stopping window size.
-    scheduler : str
-        The scheduler to be used.
-    ddp_setup_for_testing : DdpSetup
-        Information about the distributed environment, process_groups, devices, ranks, world_size, heliostat group to ranks mapping.
-    device : torch.device
-        The device on which to initialize tensors.
-
-    Raises
-    ------
-    AssertionError
-        If test does not complete as expected.
-    """
-    torch.manual_seed(7)
-    torch.cuda.manual_seed(7)
-
-    scheduler_dict = {
-        constants.scheduler_type: scheduler,
-        constants.gamma: 0.99,
-        constants.lr_min: 1e-4,
-        constants.reduce_factor: 0.9,
-        constants.patience: 100,
-        constants.threshold: 1e-3,
-        constants.cooldown: 20,
-    }
-    optimizer_dict = {
-        constants.initial_learning_rate_rotation_deviation: 1e-4,
-        constants.initial_learning_rate_initial_angles: 1e-3,
-        constants.initial_learning_rate_initial_stroke_length: 1e-2,
-        constants.tolerance: 0.0005,
-        constants.max_epoch: 250,
-        constants.batch_size: 50,
-        constants.log_step: 1,
-        constants.early_stopping_delta: 1.0,
-        constants.early_stopping_patience: 2,
-        constants.early_stopping_window: early_stopping_window,
-    }
-    optimization_configuration = {
-        constants.optimization: optimizer_dict,
-        constants.scheduler: scheduler_dict,
-    }
-
-    scenario_path = (
-        pathlib.Path(ARTIST_ROOT)
-        / "tests/data/scenarios/test_scenario_paint_four_heliostats.h5"
-    )
-
-    heliostat_data_mapping = [
-        (
-            "AA39",
-            [
-                pathlib.Path(ARTIST_ROOT)
-                / "tests/data/field_data/AA39-calibration-properties_1.json",
-                pathlib.Path(ARTIST_ROOT)
-                / "tests/data/field_data/AA39-calibration-properties_2.json",
-            ],
-            [
-                pathlib.Path(ARTIST_ROOT) / "tests/data/field_data/AA39-flux_1.png",
-                pathlib.Path(ARTIST_ROOT) / "tests/data/field_data/AA39-flux_2.png",
-            ],
-        ),
-        (
-            "AA31",
-            [
-                pathlib.Path(ARTIST_ROOT)
-                / "tests/data/field_data/AA31-calibration-properties_1.json",
-                pathlib.Path(ARTIST_ROOT)
-                / "tests/data/field_data/AA31-calibration-properties_2.json",
-            ],
-            [
-                pathlib.Path(ARTIST_ROOT) / "tests/data/field_data/AA31-flux_1.png",
-                pathlib.Path(ARTIST_ROOT) / "tests/data/field_data/AA31-flux_2.png",
-            ],
-        ),
-    ]
-
-    if centroid_extraction_method == "invalid":
-        with pytest.raises(ValueError) as exc_info:
-            data_parser = (
-                PaintCalibrationDataParser(
-                    centroid_extraction_method=centroid_extraction_method
-                )
-                if isinstance(data_parser, PaintCalibrationDataParser)
-                else CalibrationDataParser()
-            )
-            assert (
-                f"The selected centroid extraction method {centroid_extraction_method} is not yet supported. Please use either {paint_mappings.UTIS_KEY} or {paint_mappings.HELIOS_KEY}!"
-                in str(exc_info.value)
-            )
-    else:
-        data: dict[
-            str,
-            CalibrationDataParser
-            | list[tuple[str, list[pathlib.Path], list[pathlib.Path]]],
-        ] = {
-            constants.data_parser: data_parser,
-            constants.heliostat_data_mapping: heliostat_data_mapping,
-        }
-
-        with h5py.File(scenario_path, "r") as scenario_file:
-            scenario = Scenario.load_scenario_from_hdf5(
-                scenario_file=scenario_file, device=device
-            )
-
-        ddp_setup_for_testing["device"] = device
-
-        if reconstruction_method == "invalid":
-            with pytest.raises(ValueError) as exc_info:
-                _ = KinematicsReconstructor(
-                    ddp_setup=ddp_setup_for_testing,
-                    scenario=scenario,
-                    data=data,
-                    optimization_configuration=optimization_configuration,
-                    reconstruction_method=reconstruction_method,
-                )
-                assert (
-                    f"ARTIST currently only supports the {constants.kinematics_reconstruction_raytracing} reconstruction method. The reconstruction method {reconstruction_method} is not recognized. Please select another reconstruction method and try again!"
-                    in str(exc_info.value)
-                )
-        else:
-            kinematics_reconstructor = KinematicsReconstructor(
-                ddp_setup=ddp_setup_for_testing,
-                scenario=scenario,
-                data=data,
-                optimization_configuration=optimization_configuration,
-                reconstruction_method=reconstruction_method,
-            )
-
-            loss_definition = FocalSpotLoss(scenario=scenario)
-
-            # Reconstruct the kinematics.
-            if not isinstance(data_parser, PaintCalibrationDataParser):
-                with pytest.raises(NotImplementedError) as exc_info:
-                    _ = kinematics_reconstructor.reconstruct_kinematics(
-                        loss_definition=loss_definition, device=device
-                    )
-
-                    assert "Must be overridden!" in str(exc_info.value)
-            else:
-                _ = kinematics_reconstructor.reconstruct_kinematics(
-                    loss_definition=loss_definition, device=device
-                )
-
-                for index, heliostat_group in enumerate(
-                    scenario.heliostat_field.heliostat_groups
-                ):
-                    expected_path = (
-                        pathlib.Path(ARTIST_ROOT)
-                        / "tests/data/expected_reconstructed_kinematics_parameters"
-                        / f"group_{index}_{early_stopping_window}_{device.type}.pt"
-                    )
-
-                    expected = torch.load(
-                        expected_path, map_location=device, weights_only=True
-                    )
-
-                    tol = 1e-6 + 5e-3 * torch.maximum(
-                        heliostat_group.kinematics.rotation_deviation_parameters.abs(),
-                        expected["rotation_deviations"].abs(),
-                    )
-                    diff = (
-                        heliostat_group.kinematics.rotation_deviation_parameters
-                        - expected["rotation_deviations"]
-                    ).abs()
-
-                    assert torch.all(diff <= tol)
-
-                    tol = 1e-6 + 5e-3 * torch.maximum(
-                        heliostat_group.kinematics.actuators.optimizable_parameters.abs(),
-                        expected["optimizable_parameters"].abs(),
-                    )
-                    diff = (
-                        heliostat_group.kinematics.actuators.optimizable_parameters
-                        - expected["optimizable_parameters"]
-                    ).abs()
-
-                    assert torch.all(diff <= tol)
diff --git a/tests/optim/test_loss_functions.py b/tests/optim/test_loss_functions.py
index e39f8574c..5d6774937 100644
--- a/tests/optim/test_loss_functions.py
+++ b/tests/optim/test_loss_functions.py
@@ -1,3 +1,4 @@
+import math
 from unittest import mock
 
 import pytest
@@ -5,12 +6,12 @@
 
 from artist.field import (
     SolarTower,
-    TowerTargetAreas,
     TowerTargetAreasCylindrical,
     TowerTargetAreasPlanar,
 )
 from artist.optim.loss import (
     AngleLoss,
+    CosineSimilarityLoss,
     FocalSpotLoss,
     KLDivergenceLoss,
     Loss,
@@ -18,7 +19,6 @@
     VectorLoss,
 )
 from artist.scenario import Scenario
-from artist.scene import LightSource, LightSourceArray
 
 
 def test_base_loss(
@@ -140,22 +140,22 @@ def test_vector_loss(
         (
             torch.ones((1, 2, 2)),
             torch.tensor([[0.0, 0.0, 0.0, 1.0]]),
-            torch.tensor([[0.0, 0.0, 0.0, 1.0]]),
+            torch.ones((1, 2, 2)),
             torch.tensor([0.0]),
             True,
         ),
         (
             torch.ones((1, 2, 2)),
             torch.tensor([[1.5, 0.0, 0.0, 1.0]]),
-            torch.tensor([[1.5, 0.0, 0.0, 1.0]]),
+            torch.ones((1, 2, 2)),
             torch.tensor([0.0]),
             True,
         ),
         (
             torch.ones((1, 2, 2)),
             torch.tensor([[0.0, 0.0, 0.0, 1.0]]),
-            torch.tensor([[1.0, 1.0, 1.0, 0.0]]),
-            torch.tensor([1.732050776482]),
+            torch.tensor([[[0.0, 1.0], [0.0, 0.0]]]),
+            torch.tensor([0.7071]),
             True,
         ),
         (
@@ -259,33 +259,27 @@ def test_focal_spot_loss(
             device=device,
         )
 
-        torch.testing.assert_close(result, expected.to(device), atol=1e-6, rtol=1e-6)
+        torch.testing.assert_close(result, expected.to(device), atol=1e-5, rtol=1e-6)
 
 
 @pytest.mark.parametrize(
-    "prediction, ground_truth, target_area_dimensions, number_of_rays, expected, kwargs",
+    "prediction, ground_truth, expected, kwargs",
     [
         (
             torch.tensor([[[1.0, 2.0], [3.0, 4.0]]]),
             torch.tensor([[[1.0, 2.0], [3.0, 4.0]]]),
-            torch.tensor([[2.0, 2.0]]),
-            100,
             torch.tensor([0.0]),
             True,
         ),
         (
             torch.tensor([[[2.0, 3.0], [9.0, 12.0]]]),
             torch.tensor([[[1.0, 2.0], [8.0, 6.0]]]),
-            torch.tensor([[2.0, 2.0]]),
-            100,
-            torch.tensor([39.0]),
+            torch.tensor([2.2941176470588234]),
             True,
         ),
         (
             torch.tensor([[[2.0, 3.0], [9.0, 12.0]]]),
             torch.tensor([[[1.0, 2.0], [8.0, 6.0]]]),
-            torch.tensor([[2.0, 2.0]]),
-            100,
             torch.tensor([1.0]),
             False,
         ),
@@ -294,8 +288,6 @@ def test_focal_spot_loss(
 def test_pixel_loss(
     prediction: torch.Tensor,
     ground_truth: torch.Tensor,
-    target_area_dimensions: torch.Tensor,
-    number_of_rays: int,
     expected: torch.Tensor,
     kwargs: bool,
     device: torch.device,
@@ -311,11 +303,6 @@ def test_pixel_loss(
     ground_truth : torch.Tensor
         The ground truth.
         Tensor of shape [number_of_samples, bitmap_resolution_e, bitmap_resolution_u].
-    target_area_dimensions : torch.Tensor
-        The dimensions of the tower target areas aimed at.
-        Tensor of shape [number_of_samples, 2].
-    number_of_rays : int
-        The number of rays used to generate the flux.
     expected : torch.Tensor
         The expected pixel loss.
         Tensor of shape [number_of_samples].
@@ -329,21 +316,7 @@ def test_pixel_loss(
     AssertionError
         If test does not complete as expected.
     """
-    mock_scenario = mock.MagicMock(spec=Scenario)
-
-    target_areas = mock.MagicMock(spec=TowerTargetAreas)
-    target_areas.dimensions = target_area_dimensions.to(device)
-    mock_scenario.target_areas = target_areas
-
-    light_sources = mock.MagicMock(spec=LightSourceArray)
-    light_source = mock.MagicMock(spec=LightSource)
-    light_source.number_of_rays = number_of_rays
-    light_sources.light_source_list = [light_source]
-    mock_scenario.light_sources = light_sources
-
-    target_area_indices = torch.tensor([0], device=device)
-
-    pixel_loss = PixelLoss(scenario=mock_scenario)
+    pixel_loss = PixelLoss()
 
     if not kwargs:
         with pytest.raises(ValueError) as exc_info:
@@ -352,7 +325,7 @@ def test_pixel_loss(
                 ground_truth=ground_truth.to(device),
             )
         assert (
-            "The vector loss expects ['reduction_dimensions', 'device', 'target_area_indices'] as keyword arguments. Please add reduction_dimensions as keyword argument. Please add device as keyword argument. Please add target_area_indices as keyword argument."
+            "The vector loss expects ['reduction_dimensions'] as keyword arguments. Please add reduction_dimensions as keyword argument."
             in str(exc_info.value)
         )
 
@@ -360,9 +333,7 @@ def test_pixel_loss(
         result = pixel_loss(
             prediction=prediction.to(device),
             ground_truth=ground_truth.to(device),
-            target_area_indices=target_area_indices,
             reduction_dimensions=(1, 2),
-            device=device,
         )
 
         torch.testing.assert_close(result, expected.to(device), atol=1e-6, rtol=1e-6)
@@ -495,7 +466,7 @@ def test_kl_divergence(
         ),
     ],
 )
-def test_angle_loss(
+def test_cosine_similarity_loss(
     prediction: torch.Tensor,
     ground_truth: torch.Tensor,
     reduction_dimensions: tuple[int],
@@ -503,7 +474,7 @@ def test_angle_loss(
     device: torch.device,
 ) -> None:
     """
-    Test the angle loss.
+    Test the cosine similarity loss.
 
     Parameters
     ----------
@@ -526,7 +497,7 @@ def test_angle_loss(
     AssertionError
         If test does not complete as expected.
     """
-    vector_loss = AngleLoss()
+    vector_loss = CosineSimilarityLoss()
     result = vector_loss(
         prediction=prediction.to(device),
         ground_truth=ground_truth.to(device),
@@ -536,3 +507,62 @@ def test_angle_loss(
     )
 
     torch.testing.assert_close(result, expected.to(device), atol=1e-6, rtol=1e-6)
+
+
+@pytest.mark.parametrize(
+    "prediction, ground_truth, expected",
+    [
+        (
+            torch.tensor([[1.0, 2.0, 3.0, 0.0], [1.0, 1.0, 0.0, 1.0]]),
+            torch.tensor([[1.0, 2.0, 3.0, 0.0], [0.0, 1.0, 0.0, 0.0]]),
+            torch.tensor([0.0, math.pi / 4]),
+        ),
+        (
+            torch.tensor([[0.0, 1.0, 0.0]]),
+            torch.tensor([[0.0, -1.0, 0.0]]),
+            torch.tensor([math.pi]),
+        ),
+        (
+            torch.tensor([[0.0, 1.0, 0.0]]),
+            torch.tensor([[1.0, 0.0, 0.0]]),
+            torch.tensor([math.pi / 2]),
+        ),
+    ],
+)
+def test_angle_loss(
+    prediction: torch.Tensor,
+    ground_truth: torch.Tensor,
+    expected: torch.Tensor,
+    device: torch.device,
+) -> None:
+    """
+    Test the angle loss.
+
+    Parameters
+    ----------
+    prediction : torch.Tensor
+        The predicted values.
+        Tensor of variable shape.
+    ground_truth : torch.Tensor
+        The ground truth.
+        Tensor of variable shape.
+    expected : torch.Tensor
+        The expected loss.
+        Tensor of shape [number_of_samples].
+    device : torch.device
+        The device on which to initialize tensors.
+
+    Raises
+    ------
+    AssertionError
+        If test does not complete as expected.
+    """
+    vector_loss = AngleLoss()
+    result = vector_loss(
+        prediction=prediction.to(device),
+        ground_truth=ground_truth.to(device),
+        target_area_indices=torch.tensor([0, 1], device=device),
+        device=device,
+    )
+
+    torch.testing.assert_close(result, expected.to(device), atol=1e-3, rtol=1e-4)
diff --git a/tests/optim/test_surface_reconstructor.py b/tests/optim/test_surface_reconstructor.py
index 43d4c58a7..7c01eb727 100644
--- a/tests/optim/test_surface_reconstructor.py
+++ b/tests/optim/test_surface_reconstructor.py
@@ -14,20 +14,20 @@
 
 
 @pytest.mark.parametrize(
-    "loss_class, early_stopping_window, data_parser, scheduler",
+    "loss, early_stopping_window, data_parser, scheduler",
     [
         (
-            KLDivergenceLoss,
+            KLDivergenceLoss(),
             40,
             PaintCalibrationDataParser(),
             constants.reduce_on_plateau,
         ),
-        (PixelLoss, 20, PaintCalibrationDataParser(), constants.cyclic),
-        (PixelLoss, 10, CalibrationDataParser(), constants.cyclic),
+        (PixelLoss(), 10, PaintCalibrationDataParser(), constants.cyclic),
+        (PixelLoss(), 10, CalibrationDataParser(), constants.cyclic),
     ],
 )
 def test_surface_reconstructor(
-    loss_class: Loss,
+    loss: Loss,
     early_stopping_window: int,
     data_parser: CalibrationDataParser | PaintCalibrationDataParser,
     scheduler: str,
@@ -39,8 +39,8 @@ def test_surface_reconstructor(
 
     Parameters
     ----------
-    loss_class : Loss
-        The loss class.
+    loss: Loss
+        The loss definition.
     early_stopping_window : int
         Number of epochs used to estimate loss trend.
     data_parser : CalibrationDataParser
@@ -105,27 +105,16 @@ def test_surface_reconstructor(
                 / "tests/data/field_data/AA39-calibration-properties_1.json",
                 pathlib.Path(ARTIST_ROOT)
                 / "tests/data/field_data/AA39-calibration-properties_2.json",
+                pathlib.Path(ARTIST_ROOT)
+                / "tests/data/field_data/AA39-calibration-properties_3.json",
             ],
             [
                 pathlib.Path(ARTIST_ROOT)
                 / "tests/data/field_data/AA39-flux-centered_1.png",
                 pathlib.Path(ARTIST_ROOT)
                 / "tests/data/field_data/AA39-flux-centered_2.png",
-            ],
-        ),
-        (
-            "AA31",
-            [
-                pathlib.Path(ARTIST_ROOT)
-                / "tests/data/field_data/AA31-calibration-properties_1.json",
                 pathlib.Path(ARTIST_ROOT)
-                / "tests/data/field_data/AA31-calibration-properties_2.json",
-            ],
-            [
-                pathlib.Path(ARTIST_ROOT)
-                / "tests/data/field_data/AA31-flux-centered_1.png",
-                pathlib.Path(ARTIST_ROOT)
-                / "tests/data/field_data/AA31-flux-centered_2.png",
+                / "tests/data/field_data/AA39-flux-centered_3.png",
             ],
         ),
     ]
@@ -157,25 +146,21 @@ def test_surface_reconstructor(
         data=data,
         optimization_configuration=optimization_configuration,
         device=device,
-    )
-
-    loss_definition = (
-        PixelLoss(scenario=scenario) if loss_class is PixelLoss else KLDivergenceLoss()
+        plot_results=True,
     )
 
     if not isinstance(data_parser, PaintCalibrationDataParser):
         with pytest.raises(NotImplementedError) as exc_info:
             _ = surface_reconstructor.reconstruct_surfaces(
-                loss_definition=loss_definition, device=device
+                loss_definition=loss, device=device
             )
-
-            assert "Must be overridden!" in str(exc_info.value)
+        assert "Must be overridden!" in str(exc_info.value)
     else:
         old_state = torch.are_deterministic_algorithms_enabled()
         torch.use_deterministic_algorithms(False)
         try:
             _ = surface_reconstructor.reconstruct_surfaces(
-                loss_definition=loss_definition, device=device
+                loss_definition=loss, device=device
             )
         finally:
             torch.use_deterministic_algorithms(old_state)
@@ -183,7 +168,7 @@ def test_surface_reconstructor(
         for index, heliostat_group in enumerate(
             scenario.heliostat_field.heliostat_groups
         ):
-            loss_name = "pixel_loss" if loss_class is PixelLoss else "kl_divergence"
+            loss_name = "pixel_loss" if isinstance(loss, PixelLoss) else "kl_divergence"
             expected_path = (
                 pathlib.Path(ARTIST_ROOT)
                 / "tests/data/expected_reconstructed_surfaces"
diff --git a/tutorials/03_nurbs_surface_reconstruction_tutorial.py b/tutorials/03_nurbs_surface_reconstruction_tutorial.py
index 4b7c99c36..b59e62b8b 100644
--- a/tutorials/03_nurbs_surface_reconstruction_tutorial.py
+++ b/tutorials/03_nurbs_surface_reconstruction_tutorial.py
@@ -26,7 +26,7 @@
 
 #############################################################################################################
 # Define helper functions for the plots.
-# Skip to line 348 for the tutorial code.
+# Skip to line 353 for the tutorial code.
 #############################################################################################################
 
 
@@ -402,7 +402,7 @@ def create_flux_plots(
 # heliostat_data_mapping = paint_scenario_parser.build_heliostat_data_mapping(
 #     base_path=base_path_data,
 #     heliostat_names=heliostat_names_reconstruction,
-#     number_of_measurements=2,
+#     number_of_measurements=4,
 #     image_variant="flux-centered",
 #     randomize=True,
 # )
@@ -411,9 +411,9 @@ def create_flux_plots(
 optimizer_dict = {
     constants.initial_learning_rate: 1e-4,
     constants.tolerance: 1e-5,
-    constants.max_epoch: 100,
+    constants.max_epoch: 200,
     constants.batch_size: 30,
-    constants.log_step: 1,
+    constants.log_step: 3,
     constants.early_stopping_delta: 1e-4,
     constants.early_stopping_patience: 100,
     constants.early_stopping_window: 100,
@@ -449,7 +449,7 @@ def create_flux_plots(
     str,
     CalibrationDataParser | list[tuple[str, list[pathlib.Path], list[pathlib.Path]]],
 ] = {
-    constants.data_parser: PaintCalibrationDataParser(sample_limit=2),
+    constants.data_parser: PaintCalibrationDataParser(sample_limit=4),
     constants.heliostat_data_mapping: heliostat_data_mapping,
 }
 
@@ -502,6 +502,7 @@ def create_flux_plots(
         data=data,
         optimization_configuration=optimization_configuration,
         bitmap_resolution=resolution,
+        plot_results=True,
         device=device,
     )
 
diff --git a/tutorials/04_kinematics_reconstruction_tutorial.py b/tutorials/04_kinematics_reconstruction_tutorial.py
index e411f22a9..6a3bf7fd4 100644
--- a/tutorials/04_kinematics_reconstruction_tutorial.py
+++ b/tutorials/04_kinematics_reconstruction_tutorial.py
@@ -16,7 +16,7 @@
     paint_scenario_parser,
 )
 from artist.optim import KinematicsReconstructor
-from artist.optim.loss import FocalSpotLoss
+from artist.optim.loss import AngleLoss
 from artist.raytracing import HeliostatRayTracer
 from artist.scenario import Scenario
 from artist.util import constants, set_logger_config
@@ -35,6 +35,7 @@ def create_fluxes(
     data_parser: CalibrationDataParser,
     heliostat_data_mapping: list[tuple[str, list[pathlib.Path], list[pathlib.Path]]],
     resolution: torch.Tensor,
+    method: str,
 ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
     """
     Create data to plot the heliostat fluxes.
@@ -57,7 +58,7 @@ def create_fluxes(
     """
     bitmaps = []
     measured_bitmaps = []
-    scenario.set_number_of_rays(number_of_rays=500)
+    scenario.set_number_of_rays(number_of_rays=100)
     for heliostat_group_index in range(len(scenario.heliostat_field.heliostat_groups)):
         heliostat_group: HeliostatGroup = scenario.heliostat_field.heliostat_groups[
             heliostat_group_index
@@ -66,7 +67,7 @@ def create_fluxes(
             measured_flux,
             _,
             incident_ray_directions,
-            _,
+            motor_positions,
             active_heliostats_mask,
             target_area_indices,
         ) = data_parser.parse_data_for_reconstruction(
@@ -86,15 +87,23 @@ def create_fluxes(
                 device=device,
             )
 
-            # Align heliostats.
-            heliostat_group.align_surfaces_with_incident_ray_directions(
-                aim_points=scenario.solar_tower.get_centers_of_target_areas(
-                    target_area_indices=target_area_indices, device=device
-                ),
-                incident_ray_directions=incident_ray_directions,
-                active_heliostats_mask=active_heliostats_mask,
-                device=device,
-            )
+            if method == "motor_pos":
+                # Align heliostats.
+                heliostat_group.align_surfaces_with_motor_positions(
+                    motor_positions=motor_positions,
+                    active_heliostats_mask=active_heliostats_mask,
+                    device=device,
+                )
+            elif method == "incident_ray":
+                # Align heliostats.
+                heliostat_group.align_surfaces_with_incident_ray_directions(
+                    aim_points=scenario.solar_tower.get_centers_of_target_areas(
+                        target_area_indices=target_area_indices, device=device
+                    ),
+                    incident_ray_directions=incident_ray_directions,
+                    active_heliostats_mask=active_heliostats_mask,
+                    device=device,
+                )
 
             # Create a ray tracer.
             ray_tracer = HeliostatRayTracer(
@@ -122,6 +131,7 @@ def create_fluxes(
 def create_plots(
     fluxes_before: torch.Tensor,
     fluxes_after: torch.Tensor,
+    fluxes_ray: torch.Tensor,
     fluxes_measured: torch.Tensor,
 ) -> None:
     """
@@ -136,8 +146,8 @@ def create_plots(
     fluxes_measured : torch.Tensor
         Measured flux references.
     """
-    for group_index, (flux_before, flux_after, flux_measured) in enumerate(
-        zip(fluxes_before, fluxes_after, fluxes_measured)
+    for group_index, (flux_before, flux_after, flux_measured, flux_ray) in enumerate(
+        zip(fluxes_before, fluxes_after, fluxes_measured, fluxes_ray)
     ):
         center_measured = (
             bitmap.get_center_of_mass(flux_measured, device=device).cpu().detach()
@@ -148,29 +158,86 @@ def create_plots(
         center_after = (
             bitmap.get_center_of_mass(flux_after, device=device).cpu().detach()
         )
+        center_ray = bitmap.get_center_of_mass(flux_ray, device=device).cpu().detach()
 
         for i in range(len(flux_before)):
-            _, axes = plt.subplots(nrows=1, ncols=3, figsize=(15, 5))
+            _, axes = plt.subplots(nrows=1, ncols=4, figsize=(20, 5))
 
             axes[0].imshow(flux_before[i].cpu().detach(), cmap="gray")
             axes[0].set_title("Before reconstruction", fontsize=16)
-            axes[0].scatter(x=center_measured[i, 0], y=center_measured[i, 1], c="r")
-            axes[0].scatter(x=center_before[i, 0], y=center_before[i, 1], c="g")
+            axes[0].scatter(
+                x=center_measured[i, 0],
+                y=center_measured[i, 1],
+                c="green",
+                marker="o",
+                label="Measured",
+            )
+            axes[0].scatter(
+                x=center_before[i, 0],
+                y=center_before[i, 1],
+                c="red",
+                marker="x",
+                label="Predicted",
+            )
+            # axes[0].scatter(256/2, 256/2, c="blue")
             axes[0].axis("off")
 
             axes[1].imshow(flux_after[i].cpu().detach(), cmap="gray")
             axes[1].set_title("After reconstruction", fontsize=16)
-            axes[1].scatter(x=center_measured[i, 0], y=center_measured[i, 1], c="r")
-            axes[1].scatter(x=center_after[i, 0], y=center_after[i, 1], c="g")
+            axes[1].scatter(
+                x=center_measured[i, 0],
+                y=center_measured[i, 1],
+                c="green",
+                marker="o",
+                label="Measured",
+            )
+            axes[1].scatter(
+                x=center_after[i, 0],
+                y=center_after[i, 1],
+                c="red",
+                marker="x",
+                label="Predicted",
+            )
+            # axes[1].scatter(256/2, 256/2, c="blue")
             axes[1].axis("off")
 
-            axes[2].imshow(flux_measured[i].cpu().detach(), cmap="gray")
-            axes[2].set_title("Measured", fontsize=16)
+            axes[2].imshow(flux_ray[i].cpu().detach(), cmap="gray")
+            axes[2].set_title(
+                "After reconstruction align only with sun pos", fontsize=16
+            )
+            axes[2].scatter(
+                x=center_measured[i, 0],
+                y=center_measured[i, 1],
+                c="green",
+                marker="o",
+                label="Measured",
+            )
+            axes[2].scatter(
+                x=center_ray[i, 0],
+                y=center_ray[i, 1],
+                c="red",
+                marker="x",
+                label="Predicted",
+            )
+            # axes[2].scatter(256/2, 256/2, c="blue")
             axes[2].axis("off")
 
+            axes[3].imshow(flux_measured[i].cpu().detach(), cmap="gray")
+            axes[3].set_title("Measured", fontsize=16)
+            axes[3].scatter(
+                x=center_measured[i, 0], y=center_measured[i, 1], c="green", marker="o"
+            )
+            axes[3].axis("off")
+
+            axes[0].legend()
+            axes[1].legend()
+            axes[2].legend()
+
             plt.subplots_adjust(wspace=0.05)
             plt.show()
-            plt.savefig(f"heliostat_{i}_in_group_{group_index}_calibration.png")
+            plt.savefig(
+                f"./ignored/fixed/heliostat_{i}_in_group_{group_index}_calibration.png"
+            )
 
 
 #############################################################################################################
@@ -225,10 +292,11 @@ def create_plots(
     number_of_measurements=5,
     image_variant="flux",
     randomize=True,
+    seed=7,
 )
 
 # Configure the optimization.
-optimizer_dict = {
+optimizer_dict: dict[str, str | float | int] = {
     constants.initial_learning_rate_rotation_deviation: 1e-4,
     constants.initial_learning_rate_initial_angles: 1e-3,
     constants.initial_learning_rate_initial_stroke_length: 1e-2,
@@ -241,7 +309,7 @@ def create_plots(
     constants.early_stopping_window: 2000,
 }
 # Configure the learning rate scheduler.
-scheduler_dict = {
+scheduler_dict: dict[str, str | float | int] = {
     constants.scheduler_type: constants.reduce_on_plateau,
     constants.gamma: 0.9,
     constants.lr_min: 1e-6,
@@ -253,13 +321,13 @@ def create_plots(
     constants.cooldown: 10,
 }
 # Combine configurations.
-optimization_configuration = {
+optimization_configuration: dict[str, dict[str, str | float | int]] = {
     constants.optimization: optimizer_dict,
     constants.scheduler: scheduler_dict,
 }
 
 data_parser = PaintCalibrationDataParser(
-    sample_limit=10, centroid_extraction_method=paint_mappings.UTIS_KEY
+    sample_limit=15, centroid_extraction_method=paint_mappings.UTIS_KEY
 )
 data_parser_plots = PaintCalibrationDataParser(
     sample_limit=1, centroid_extraction_method=paint_mappings.UTIS_KEY
@@ -307,10 +375,15 @@ def create_plots(
             for heliostat in heliostat_data_mapping
         ],
         resolution=resolution,
+        method="motor_pos",
     )
 
-    loss_definition = FocalSpotLoss(scenario=scenario)
+    scenario.set_number_of_rays(number_of_rays=4)
 
+    optimization_configuration[constants.optimization][
+        constants.initial_learning_rate_rotation_deviation
+    ] = 3e-4
+    optimization_configuration[constants.optimization][constants.max_epoch] = 200
     # Create the kinematics reconstructor.
     kinematics_reconstructor = KinematicsReconstructor(
         ddp_setup=ddp_setup,
@@ -318,15 +391,36 @@ def create_plots(
         data=data,
         dni=500,
         optimization_configuration=optimization_configuration,
-        reconstruction_method=constants.kinematics_reconstruction_raytracing,
+        reconstruction_method=constants.kinematics_reconstruction_alignment,
         bitmap_resolution=resolution,
+        plot_results=True,
     )
-
     # Reconstruct the kinematics.
     final_loss_per_heliostat = kinematics_reconstructor.reconstruct_kinematics(
-        loss_definition=loss_definition, device=device
+        loss_definition=AngleLoss(), device=device
     )
 
+    # Uncomment the code below to add a further reconstruction step using raytracing to refine the kinematics reconstruction.
+    # optimization_configuration[constants.optimization][
+    #     constants.initial_learning_rate_rotation_deviation
+    # ] = 1e-4
+    # optimization_configuration[constants.optimization][constants.max_epoch] = 400
+    # # Create the kinematics reconstructor.
+    # kinematics_reconstructor = KinematicsReconstructor(
+    #     ddp_setup=ddp_setup,
+    #     scenario=scenario,
+    #     data=data,
+    #     dni=500,
+    #     optimization_configuration=optimization_configuration,
+    #     reconstruction_method=constants.kinematics_reconstruction_raytracing,
+    #     bitmap_resolution=resolution,
+    #     plot_results=True,
+    # )
+    # # Reconstruct the kinematics.
+    # final_loss_per_heliostat = kinematics_reconstructor.reconstruct_kinematics(
+    #     loss_definition=FocalSpotLoss(scenario=scenario), device=device
+    # )
+
 # Inspect the synchronized loss per heliostat. Heliostats that have not been optimized have an infinite loss.
 print(f"rank {ddp_setup['rank']}, final loss per heliostat {final_loss_per_heliostat}")
 
@@ -337,9 +431,22 @@ def create_plots(
         for heliostat in heliostat_data_mapping
     ],
     resolution=resolution,
+    method="motor_pos",
+)
+
+bitmaps_ray, bitmaps_measured = create_fluxes(
+    data_parser=data_parser_plots,
+    heliostat_data_mapping=[
+        (heliostat[0], [heliostat[1][-1]], [heliostat[2][-1]])
+        for heliostat in heliostat_data_mapping
+    ],
+    resolution=resolution,
+    method="incident_ray",
 )
+
 create_plots(
     fluxes_before=bitmaps_before,
     fluxes_after=bitmaps_after,
+    fluxes_ray=bitmaps_ray,
     fluxes_measured=bitmaps_measured,
 )
diff --git a/tutorials/05_motor_positions_optimizer_tutorial.py b/tutorials/05_motor_positions_optimizer_tutorial.py
index 9b96e5de8..0a1eac6de 100644
--- a/tutorials/05_motor_positions_optimizer_tutorial.py
+++ b/tutorials/05_motor_positions_optimizer_tutorial.py
@@ -1,4 +1,4 @@
-"""Motor positions optimizer tutorial."""
+"""Aim point optimization tutorial."""
 
 import pathlib
 
@@ -7,7 +7,7 @@
 from matplotlib import pyplot as plt
 
 from artist.flux import bitmap
-from artist.optim import MotorPositionsOptimizer
+from artist.optim import AimPointOptimizer
 from artist.optim.loss import KLDivergenceLoss
 from artist.raytracing import HeliostatRayTracer
 from artist.scenario import Scenario
@@ -19,7 +19,7 @@
 
 #############################################################################################################
 # Define helper functions for the plots.
-# Skip to line 115 for the tutorial code.
+# Skip to line 124 for the tutorial code.
 #############################################################################################################
 
 
@@ -238,8 +238,8 @@ def create_flux_plot(label: str, resolution: torch.Tensor) -> None:
 
     create_flux_plot(label="before", resolution=bitmap_resolution)
 
-    # Create the motor positions optimizer.
-    motor_positions_optimizer = MotorPositionsOptimizer(
+    # Create the aim point optimizer.
+    aim_point_optimizer = AimPointOptimizer(
         ddp_setup=ddp_setup,
         scenario=scenario,
         optimization_configuration=optimization_configuration,
@@ -252,11 +252,11 @@ def create_flux_plot(label: str, resolution: torch.Tensor) -> None:
     )
 
     # Optimize the motor positions.
-    final_loss, _, _, _, _ = motor_positions_optimizer.optimize(
+    final_loss, _, _, _, _ = aim_point_optimizer.optimize(
         loss_definition=loss_definition, device=device
     )
 
-# Inspect the synchronized loss per heliostat. Heliostats that have not been optimized have an infinite loss.
+# Inspect the synchronized loss per heliostat.
 print(f"rank {ddp_setup['rank']}, final loss {final_loss}")
 
 create_flux_plot(label="after", resolution=bitmap_resolution)