Merge branch 'main' into joss_paper

Author: MarleneBusch

.gitignore

@@ -132,6 +132,9 @@ venv.bak/
 .spyderproject
 .spyproject
 
+# VS Code project settings
+.vscode/settings.json
+
 # Rope project settings
 .ropeproject
 

.pre-commit-config.yaml

@@ -1,18 +1,18 @@
 fail_fast: false
 repos:
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.17.0 # Use the sha / tag you want to point at
+    rev: v1.17.1 # Use the sha / tag you want to point at
     hooks:
       - id: mypy
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
       - id: check-yaml
       - id: check-added-large-files
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.12.4
+    rev: v0.12.11
     hooks:
       - id: ruff
         args: [ --fix ]

artist/core/__init__.py

@@ -5,13 +5,15 @@
     HeliostatRayTracer,
     RestrictedDistributedSampler,
 )
-from artist.core.kinematic_optimizer import KinematicOptimizer
-from artist.scene.rays import Rays
+from artist.core.kinematic_calibrator import KinematicCalibrator
+from artist.core.motor_position_optimizer import MotorPositionsOptimizer
+from artist.core.surface_reconstructor import SurfaceReconstructor
 
 __all__ = [
     "HeliostatRayTracer",
     "DistortionsDataset",
     "RestrictedDistributedSampler",
-    "Rays",
-    "KinematicOptimizer",
+    "KinematicCalibrator",
+    "SurfaceReconstructor",
+    "MotorPositionsOptimizer",
 ]

artist/core/core_utils.py

@@ -0,0 +1,93 @@
+import torch
+
+from artist.util.environment_setup import get_device
+
+
+def per_heliostat_reduction(
+    per_sample_values: torch.Tensor,
+    active_heliostats_mask: torch.Tensor,
+    device: torch.device | None = None,
+) -> torch.Tensor:
+    """
+    Compute mean losses for each heliostat with multiple samples.
+
+    If the active heliostats of one group have different amounts of samples to train on, i.e.
+    one heliostat is trained with more samples than another, this function makes sure that
+    each heliostat still contributes equally to the overall loss of the group. This function
+    computes the mean loss for each heliostat.
+
+    Parameters
+    ----------
+    per_sample_values : torch.Tensor
+        The per sample values to be reduced.
+        Tensor of shape [number_of_samples]
+    active_heliostats_mask : torch.Tensor
+        A mask defining which heliostats are activated.
+        Tensor of shape [number_of_heliostats].
+    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 mean loss per heliostat.
+        Tensor of shape [number_of_heliostats].
+    """
+    device = get_device(device=device)
+
+    # A sample to heliostat index mapping.
+    heliostat_ids = torch.repeat_interleave(
+        torch.arange(len(active_heliostats_mask), device=device),
+        active_heliostats_mask,
+    )
+
+    loss_sum_per_heliostat = torch.zeros(len(active_heliostats_mask), device=device)
+    loss_sum_per_heliostat = loss_sum_per_heliostat.index_add(
+        0, heliostat_ids, per_sample_values
+    )
+
+    # Compute mean MSE per heliostat on each rank.
+    number_of_samples_per_heliostat = torch.zeros(
+        len(active_heliostats_mask), device=device
+    )
+    number_of_samples_per_heliostat.index_add_(
+        0, heliostat_ids, torch.ones_like(per_sample_values, device=device)
+    )
+
+    counts_clamped = number_of_samples_per_heliostat.clamp_min(1.0)
+    mean_loss_per_heliostat = loss_sum_per_heliostat / counts_clamped
+    mean_loss_per_heliostat = mean_loss_per_heliostat * (
+        number_of_samples_per_heliostat > 0
+    )
+
+    return mean_loss_per_heliostat
+
+
+def scale_loss(
+    loss: torch.Tensor, reference: torch.Tensor, weight: float
+) -> torch.Tensor:
+    """
+    Scale one loss so that its weighted contribution is a ratio of the reference loss.
+
+    Parameters
+    ----------
+    loss : torch.Tensor
+        The loss to be scaled.
+        Tensor of shape [1].
+    reference_loss :  torch.Tensor
+        The reference loss.
+        Tensor of shape [1].
+    weight : float
+        The weight or ratio used for the scaling.
+
+    Returns
+    -------
+    torch.Tensor
+        The scaled loss.
+        Tensor of shape [1].
+    """
+    epsilon = 1e-12
+    scale = (reference * weight) / (loss + epsilon)
+    return loss * scale

artist/core/heliostat_ray_tracer.py

@@ -194,10 +194,9 @@ class HeliostatRayTracer:
         The distortion sampler.
     distortions_loader : DataLoader
         The dataloader that loads the distortions.
-    bitmap_resolution_e : int
-        The resolution of the bitmap in the east dimension.
-    bitmap_resolution_u : int
-        The resolution of the bitmap in the up dimension.
+    bitmap_resolution : int
+        The resolution of the bitmap in both directions.
+        Tensor of shape [2].
 
     Methods
     -------
@@ -219,8 +218,7 @@ def __init__(
         rank: int = 0,
         batch_size: int = 1,
         random_seed: int = 7,
-        bitmap_resolution_e: int = 256,
-        bitmap_resolution_u: int = 256,
+        bitmap_resolution: torch.Tensor = torch.tensor([256, 256]),
     ) -> None:
         """
         Initialize the heliostat ray tracer.
@@ -245,10 +243,9 @@ def __init__(
             The amount of samples (Heliostats) processed parallel within a single rank (default is 1).
         random_seed : int
             The random seed used for generating the distortions (default is 7).
-        bitmap_resolution_e : int
-            The resolution of the bitmap in the east dimension (default is 256).
-        bitmap_resolution_u : int
-            The resolution of the bitmap in the up dimension (default is 256).
+        bitmap_resolution : torch.Tensor
+            The resolution of the bitmap in both directions. (default is torch.tensor([256,256])).
+            Tensor of shape [2].
         """
         self.scenario = scenario
         self.heliostat_group = heliostat_group
@@ -282,8 +279,7 @@ def __init__(
             sampler=self.distortions_sampler,
         )
 
-        self.bitmap_resolution_e = bitmap_resolution_e
-        self.bitmap_resolution_u = bitmap_resolution_u
+        self.bitmap_resolution = bitmap_resolution
 
     def trace_rays(
         self,
@@ -296,19 +292,22 @@ def trace_rays(
         Perform heliostat ray tracing.
 
         Scatter the rays according to the distortions, calculate the intersections with the target planes,
-        and sample the resulting bitmaps on the target areas. The bitmaps are generated seperatly for each
+        and sample the resulting bitmaps on the target areas. The bitmaps are generated separately for each
         active heliostat and can be accessed individually or they can be combined to get the total flux
         density distribution for all heliostats on all target areas.
 
         Parameters
         ----------
         incident_ray_directions : torch.Tensor
             The direction of the incident rays as seen from the heliostats.
+            Tensor of shape [number_of_active_heliostats, 4].
         active_heliostats_mask : torch.Tensor
             A mask where 0 indicates a deactivated heliostat and 1 an activated one.
             An integer greater than 1 indicates that this heliostat is regarded multiple times.
+            Tensor of shape [number_of_heliostats].
         target_area_mask : torch.Tensor
             The indices of the target areas for each active heliostat.
+            Tensor of shape [number_of_active_heliostats].
         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
@@ -323,6 +322,7 @@ def trace_rays(
         -------
         torch.Tensor
             The resulting bitmaps per heliostat.
+            Tensor of shape [number_of_active_heliostats, bitmap_resolution_e, bitmap_resolution_u].
         """
         device = get_device(device=device)
 
@@ -333,8 +333,8 @@ def trace_rays(
         flux_distributions = torch.zeros(
             (
                 self.heliostat_group.number_of_active_heliostats,
-                self.bitmap_resolution_u,
-                self.bitmap_resolution_e,
+                self.bitmap_resolution[1],
+                self.bitmap_resolution[0],
             ),
             device=device,
         )
@@ -405,10 +405,13 @@ def scatter_rays(
         ----------
         distortion_u : torch.Tensor
             The distortions in up direction (angles for scattering).
+            Tensor of shape [number_of_active_heliostats, number_of_rays, number_of_combined_surface_normals_all_facets].
         distortion_e : torch.Tensor
             The distortions in east direction (angles for scattering).
+            Tensor of shape [number_of_active_heliostats, number_of_rays, number_of_combined_surface_normals_all_facets].
         original_ray_direction : torch.Tensor
             The ray direction around which to scatter.
+            Tensor of shape [number_of_active_heliostats, number_of_combined_surface_normals_all_facets, 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
@@ -451,12 +454,16 @@ def sample_bitmaps(
         ----------
         intersections : torch.Tensor
             The intersections of rays on the target area planes for each heliostat.
+            Tensor of shape [number_of_active_heliostats, number_of_rays, number_of_combined_surface_points_all_facets, 4].
         absolute_intensities : torch.Tensor
             The absolute intensities of the rays hitting the target planes for each heliostat.
+            Tensor of shape [number_of_active_heliostats, number_of_rays, number_of_combined_surface_points_all_facets].
         active_heliostats_mask : torch.Tensor
             Used to map bitmaps per heliostat to correct index.
+            Tensor of shape [number_of_heliostats].
         target_area_mask : torch.Tensor
-            The indices of target areas on which each heliostat should be raytraced.
+            The indices of target areas on which each heliostat is raytraced.
+            Tensor of shape [number_of_active_heliostats].
         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
@@ -466,6 +473,7 @@ def sample_bitmaps(
         -------
         torch.Tensor
             The flux density distributions of the reflected rays on the target areas for each active heliostat.
+            Tensor of shape [number_of_active_heliostats, bitmap_resolution_e, bitmap_resolution_u].
         """
         device = get_device(device=device)
 
@@ -512,10 +520,10 @@ def sample_bitmaps(
         # x_intersections and y_intersections contain those intersection coordinates scaled to a range from 0 to bitmap_resolution_e/_u.
         # Additionally a mask is applied, only the intersections where intersection_indices == True are kept, the tensors are flattened.
         x_intersections = (
-            dx_intersections / plane_widths * self.bitmap_resolution_e
+            dx_intersections / plane_widths * self.bitmap_resolution[0]
         ).reshape(-1, total_intersections)
         y_intersections = (
-            dy_intersections / plane_heights * self.bitmap_resolution_u
+            dy_intersections / plane_heights * self.bitmap_resolution[1]
         ).reshape(-1, total_intersections)
 
         # We assume a continuously positioned value in-between four
@@ -614,9 +622,9 @@ def sample_bitmaps(
         # in the bitmap (i.e. we prevent out-of-bounds access).
         intersection_indices_2 = (
             (0 <= x_indices)
-            & (x_indices < self.bitmap_resolution_e)
+            & (x_indices < self.bitmap_resolution[0])
             & (0 <= y_indices)
-            & (y_indices < self.bitmap_resolution_u)
+            & (y_indices < self.bitmap_resolution[1])
         )
 
         final_intersection_indices = (
@@ -636,8 +644,8 @@ def sample_bitmaps(
         bitmaps_per_heliostat = torch.zeros(
             (
                 self.heliostat_group.number_of_active_heliostats,
-                self.bitmap_resolution_u,
-                self.bitmap_resolution_e,
+                self.bitmap_resolution[1],
+                self.bitmap_resolution[0],
             ),
             dtype=dx_intersections.dtype,
             device=device,
@@ -647,8 +655,8 @@ def sample_bitmaps(
         bitmaps_per_heliostat.index_put_(
             (
                 heliostat_indices[mask],
-                self.bitmap_resolution_u - 1 - y_indices[final_intersection_indices],
-                self.bitmap_resolution_e - 1 - x_indices[final_intersection_indices],
+                self.bitmap_resolution[1] - 1 - y_indices[final_intersection_indices],
+                self.bitmap_resolution[0] - 1 - x_indices[final_intersection_indices],
             ),
             intensities[final_intersection_indices],
             accumulate=True,
@@ -669,8 +677,10 @@ def get_bitmaps_per_target(
         ----------
         bitmaps_per_heliostat : torch.Tensor
             Bitmaps per heliostat.
+            Tensor of shape [number_of_active_heliostats, bitmap_resolution_e, bitmap_resolution_u].
         target_area_mask : torch.Tensor
             The mapping from heliostat to target area.
+            Tensor of shape [number_of_active_heliostats].
         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
@@ -680,14 +690,15 @@ def get_bitmaps_per_target(
         -------
         torch.Tensor
             Bitmaps per target area.
+            Tensor of shape [number_of_target_areas, bitmap_resolution_e, bitmap_resolution_u].
         """
         device = get_device(device=device)
 
         group_bitmaps_per_target = torch.zeros(
             (
                 self.scenario.target_areas.number_of_target_areas,
-                self.bitmap_resolution_e,
-                self.bitmap_resolution_u,
+                self.bitmap_resolution[0],
+                self.bitmap_resolution[1],
             ),
             device=device,
         )