change default from cpu to cuda and transform string-device to torch.…

…device
ARTIST-Association · Nov 6, 2024 · 1ec9497 · 1ec9497
1 parent 672ee41
commit 1ec9497
Show file tree

Hide file tree

Showing 29 changed files with 523 additions and 167 deletions.
diff --git a/artist/field/actuator.py b/artist/field/actuator.py
@@ -1,3 +1,5 @@
+from typing import Union
+
 import torch
 
 
@@ -81,7 +83,7 @@ def __init__(
         self.phi_0 = phi_0
 
     def motor_steps_to_angles(
-        self, motor_steps: torch.Tensor, device: torch.device = "cpu"
+        self, motor_steps: torch.Tensor, device: Union[torch.device, str] = "cuda"
     ) -> torch.Tensor:
         """
         Translate motor steps to a joint angle.
@@ -90,6 +92,8 @@ def motor_steps_to_angles(
         ----------
         motor_steps : torch.Tensor
             The motor steps.
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Raises
         ------
@@ -99,7 +103,7 @@ def motor_steps_to_angles(
         raise NotImplementedError("Must be overridden!")
 
     def angles_to_motor_steps(
-        self, angles: torch.Tensor, device: torch.device = "cpu"
+        self, angles: torch.Tensor, device: Union[torch.device, str] = "cuda"
     ) -> torch.Tensor:
         """
         Translate a joint angle to motor steps.
@@ -108,6 +112,8 @@ def angles_to_motor_steps(
         ----------
         angles : torch.Tensor
             The joint angles.
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Raises
         ------
@@ -116,14 +122,18 @@ def angles_to_motor_steps(
         """
         raise NotImplementedError("Must be overridden!")
 
-    def forward(self, actuator_pos: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, actuator_pos: torch.Tensor, device: Union[torch.device, str] = "cuda"
+    ) -> torch.Tensor:
         """
         Perform forward kinematic.
 
         Parameters
         ----------
         actuator_pos : torch.Tensor
             The position of the actuator.
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Raises
         ------

diff --git a/artist/field/actuator_array.py b/artist/field/actuator_array.py
@@ -1,4 +1,5 @@
 import logging
+from typing import Union
 
 import torch
 
@@ -30,7 +31,7 @@ class ActuatorArray(torch.nn.Module):
     def __init__(
         self,
         actuator_list_config: ActuatorListConfig,
-        device: torch.device = "cpu",
+        device: Union[torch.device, str] = "cuda",
     ) -> None:
         """
         Initialize the actuator array.
@@ -44,10 +45,11 @@ def __init__(
         ----------
         actuator_list_config : ActuatorListConfig
             The configuration parameters for the actuators.
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
         """
         super().__init__()
+        device = torch.device(device)
         actuator_array = []
         # Iterate through each actuator configuration in the list of actuator configurations.
         for i, actuator_config in enumerate(actuator_list_config.actuator_list):

diff --git a/artist/field/actuator_ideal.py b/artist/field/actuator_ideal.py
@@ -1,3 +1,5 @@
+from typing import Union
+
 import torch
 
 from artist.field.actuator import (
@@ -24,7 +26,7 @@ class IdealActuator(Actuator):
     """
 
     def motor_steps_to_angles(
-        self, motor_steps: torch.Tensor, device: torch.device = "cpu"
+        self, motor_steps: torch.Tensor, device: Union[torch.device, str] = "cuda"
     ) -> torch.Tensor:
         """
         Translate motor steps to a joint angle.
@@ -33,8 +35,8 @@ def motor_steps_to_angles(
         ----------
         motor_steps : torch.Tensor
             The motor steps.
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Returns
         -------
@@ -44,7 +46,7 @@ def motor_steps_to_angles(
         return motor_steps
 
     def angles_to_motor_steps(
-        self, angles: torch.Tensor, device: torch.device = "cpu"
+        self, angles: torch.Tensor, device: Union[torch.device, str] = "cuda"
     ) -> torch.Tensor:
         """
         Translate a joint angle to motor steps.
@@ -53,8 +55,8 @@ def angles_to_motor_steps(
         ----------
         angles : torch.Tensor
             The joint angles.
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Returns
         -------
@@ -64,7 +66,7 @@ def angles_to_motor_steps(
         return angles
 
     def forward(
-        self, actuator_pos: torch.Tensor, device: torch.device = "cpu"
+        self, actuator_pos: torch.Tensor, device: Union[torch.device, str] = "cuda"
     ) -> torch.Tensor:
         """
         Perform the forward kinematic for an ideal actuator.
@@ -73,8 +75,8 @@ def forward(
         ----------
         actuator_pos : torch.Tensor
             The position of the actuator.
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Returns
         -------

diff --git a/artist/field/actuator_linear.py b/artist/field/actuator_linear.py
@@ -1,3 +1,5 @@
+from typing import Union
+
 import torch
 
 from artist.field.actuator import Actuator
@@ -120,7 +122,7 @@ def steps_to_phi(self, actuator_pos: torch.Tensor) -> torch.Tensor:
         return angle
 
     def motor_steps_to_angles(
-        self, actuator_pos: torch.Tensor, device: torch.device = "cpu"
+        self, actuator_pos: torch.Tensor, device: Union[torch.device, str] = "cuda"
     ) -> torch.Tensor:
         """
         Calculate the angles given the motor steps.
@@ -129,14 +131,15 @@ def motor_steps_to_angles(
         ----------
         actuator_pos : torch.Tensor
             The actuator (motor) position.
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Returns
         -------
         torch.Tensor
             The angles corresponding to the motor steps.
         """
+        device = torch.device(device)
         phi = self.steps_to_phi(actuator_pos=actuator_pos)
         phi_0 = self.steps_to_phi(
             actuator_pos=torch.zeros(actuator_pos.shape, device=device)
@@ -147,7 +150,7 @@ def motor_steps_to_angles(
         return angles
 
     def angles_to_motor_steps(
-        self, angles: torch.Tensor, device: torch.device = "cpu"
+        self, angles: torch.Tensor, device: Union[torch.device, str] = "cuda"
     ) -> torch.Tensor:
         """
         Calculate the motor steps given the angles.
@@ -156,14 +159,15 @@ def angles_to_motor_steps(
         ----------
         angles : torch.Tensor
             The angles.
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Returns
         -------
         torch.Tensor
             The motor steps.
         """
+        device = torch.device(device)
         delta_phi = angles - self.phi_0 if self.clockwise else self.phi_0 - angles
 
         phi_0 = self.steps_to_phi(
@@ -179,7 +183,7 @@ def angles_to_motor_steps(
         return actuator_steps
 
     def forward(
-        self, actuator_pos: torch.Tensor, device: torch.device = "cpu"
+        self, actuator_pos: torch.Tensor, device: Union[torch.device, str] = "cuda"
     ) -> torch.Tensor:
         """
         Perform the forward kinematic.
@@ -188,12 +192,14 @@ def forward(
         ----------
         actuator_pos : torch.Tensor
             The actuator (motor) position.
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Returns
         -------
         torch.Tensor
             The angles.
         """
-        return self.motor_steps_to_angles(actuator_pos=actuator_pos, device=device)
+        return self.motor_steps_to_angles(
+            actuator_pos=actuator_pos, device=torch.device(device)
+        )
diff --git a/artist/field/facets_nurbs.py b/artist/field/facets_nurbs.py
@@ -1,3 +1,5 @@
+from typing import Union
+
 import torch
 
 from artist.util.nurbs import NURBSSurface
@@ -90,20 +92,23 @@ def __init__(
         self.canting_e = canting_e
         self.canting_n = canting_n
 
-    def create_nurbs_surface(self, device: torch.device = "cpu") -> NURBSSurface:
+    def create_nurbs_surface(
+        self, device: Union[torch.device, str] = "cuda"
+    ) -> NURBSSurface:
         """
         Create a NURBS surface to model a facet.
 
         Parameters
         ----------
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Returns
         -------
         NURBSSurface
             The NURBS surface of one facet.
         """
+        device = torch.device(device)
         # Since NURBS are only defined between (0,1), a small offset is required to exclude the boundaries from the
         # defined evaluation points.
         evaluation_points_rows = torch.linspace(

diff --git a/artist/field/heliostat.py b/artist/field/heliostat.py
@@ -1,5 +1,5 @@
 import logging
-from typing import Optional
+from typing import Optional, Union
 
 import h5py
 import torch
@@ -70,7 +70,7 @@ def __init__(
         surface_config: SurfaceConfig,
         kinematic_config: KinematicLoadConfig,
         actuator_config: ActuatorListConfig,
-        device: torch.device = "cpu",
+        device: Union[torch.device, str] = "cuda",
     ) -> None:
         """
         Implement the behavior of a heliostat.
@@ -95,10 +95,11 @@ def __init__(
             The configuration parameters to use for the heliostat kinematic.
         actuator_config : ActuatorListConfig
             The configuration parameters to use for the list of actuators.
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
         """
         super().__init__()
+        device = torch.device(device)
         self.heliostat_id = heliostat_id
         self.position = position
         self.aim_point = aim_point
@@ -130,7 +131,7 @@ def from_hdf5(
         prototype_kinematic: Optional[KinematicLoadConfig] = None,
         prototype_actuator: Optional[ActuatorListConfig] = None,
         heliostat_name: Optional[str] = None,
-        device: torch.device = "cpu",
+        device: Union[torch.device, str] = "cuda",
     ) -> Self:
         """
         Class method to initialize a heliostat from an HDF5 file.
@@ -147,8 +148,8 @@ def from_hdf5(
             An optional prototype for the actuator configuration.
         heliostat_name : str, optional
             The name of the heliostat being loaded - used for logging.
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Returns
         -------
@@ -157,6 +158,7 @@ def from_hdf5(
         """
         if heliostat_name:
             log.info(f"Loading {heliostat_name} from an HDF5 file.")
+        device = torch.device(device)
         heliostat_id = int(config_file[config_dictionary.heliostat_id][()])
         position = torch.tensor(
             config_file[config_dictionary.heliostat_position][()],
@@ -752,7 +754,9 @@ def from_hdf5(
         )
 
     def set_aligned_surface(
-        self, incident_ray_direction: torch.Tensor, device: torch.device = "cpu"
+        self,
+        incident_ray_direction: torch.Tensor,
+        device: Union[torch.device, str] = "cuda",
     ) -> None:
         """
         Compute the aligned surface points and aligned surface normals of the heliostat.
@@ -761,9 +765,11 @@ def set_aligned_surface(
         ----------
         incident_ray_direction : torch.Tensor
             The incident ray direction.
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
         """
+        device = torch.device(device)
+
         surface_points, surface_normals = self.surface.get_surface_points_and_normals(
             device=device
         )

diff --git a/artist/field/heliostat_field.py b/artist/field/heliostat_field.py
@@ -1,4 +1,5 @@
 import logging
+from typing import Union
 
 import h5py
 import torch.nn
@@ -54,7 +55,7 @@ def from_hdf5(
         prototype_surface: SurfaceConfig,
         prototype_kinematic: KinematicLoadConfig,
         prototype_actuator: ActuatorListConfig,
-        device: torch.device = "cpu",
+        device: Union[torch.device, str] = "cuda",
     ) -> Self:
         """
         Load a heliostat field from an HDF5 file.
@@ -69,15 +70,16 @@ def from_hdf5(
             The prototype for the kinematic configuration to be used if the heliostat has no individual kinematic.
         prototype_actuator : ActuatorListConfig
             The prototype for the actuator configuration to be used if the heliostat has no individual actuators.
-        device : torch.device
-            The device on which to initialize tensors (default is CPU).
+        device : Union[torch.device, str]
+            The device on which to initialize tensors (default is cuda).
 
         Returns
         -------
         HeliostatField
             The heliostat field loaded from the HDF5 file.
         """
         log.info("Loading a heliostat field from an HDF5 file.")
+        device = torch.device(device)
         heliostat_list = [
             Heliostat.from_hdf5(
                 config_file=config_file[config_dictionary.heliostat_key][