Fixed inverse kinematics

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",

artist/optim/motor_position_optimizer.py → 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,19 +210,17 @@ 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,
                     dtype=torch.int32,
                     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,15 +453,13 @@ 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[
                         heliostat_group_index
                     ][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(),