Allow arbitrary laser antenna normal in picmi (#3477)

* allow arbitrary laser antenna normal - picmi

* always check that laser propagation direction is along the normal direction of the antenna if the antenna normal is given

1 files changed, 22 insertions, 3 deletions

diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py
index 1c6ad317d..c112f0157 100644
--- a/Python/pywarpx/picmi.py
+++ b/Python/pywarpx/picmi.py
@@ -1054,13 +1054,32 @@ class AnalyticLaser(picmistandard.PICMI_AnalyticLaser):
         expression = pywarpx.my_constants.mangle_expression(self.field_expression, self.mangle_dict)
         self.laser.__setattr__('field_function(X,Y,t)', expression)
 
+
 class LaserAntenna(picmistandard.PICMI_LaserAntenna):
     def initialize_inputs(self, laser):
         laser.laser.position = self.position  # This point is on the laser plane
-        laser.laser.direction = self.normal_vector  # The plane normal direction
+        if (
+            self.normal_vector is not None
+            and not np.allclose(laser.laser.direction, self.normal_vector)
+        ):
+            raise AttributeError(
+                'The specified laser direction does not match the '
+                'specified antenna normal.'
+            )
+        self.normal_vector = laser.laser.direction # The plane normal direction
         if isinstance(laser, GaussianLaser):
-            laser.laser.profile_focal_distance = laser.focal_position[2] - self.position[2]  # Focal distance from the antenna (in meters)
-            laser.laser.profile_t_peak = (self.position[2] - laser.centroid_position[2])/constants.c  # The time at which the laser reaches its peak (in seconds)
+            # Focal distance from the antenna (in meters)
+            laser.laser.profile_focal_distance = np.sqrt(
+                (laser.focal_position[0] - self.position[0])**2 +
+                (laser.focal_position[1] - self.position[1])**2 +
+                (laser.focal_position[2] - self.position[2])**2
+            )
+            # The time at which the laser reaches its peak (in seconds)
+            laser.laser.profile_t_peak = np.sqrt(
+                (self.position[0] - laser.centroid_position[0])**2 +
+                (self.position[1] - laser.centroid_position[1])**2 +
+                (self.position[2] - laser.centroid_position[2])**2
+            ) / constants.c
 
 
 class ConstantAppliedField(picmistandard.PICMI_ConstantAppliedField):