Initial work on geometry rotate method.

gprMax/cmds_geometry/cmds_geometry.py

@@ -18,6 +18,10 @@
 
 import logging
 
+import gprMax.config as config
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
 logger = logging.getLogger(__name__)
 
 
@@ -43,6 +47,80 @@ class UserObjectGeometry:
         """Create the object and add it to the grid."""
         pass
 
-    def rotate(self):
+    def rotate(self, axis, angle):
-        """Rotate geometry object."""
+        """Rotate geometry object.
-        logger.debug('Must complete rotate method.')
+        
+        Args:
+            axis (str): axis about which to perform rotation (x, y, or z)
+            angle (int): angle of rotation (degrees)
+        """
+
+        orig_p1 = self.kwargs['p1']
+        orig_p2 = self.kwargs['p2']
+        p1 = np.array([self.kwargs['p1']])
+        p2 = np.array([self.kwargs['p2']])
+
+        # Check angle value is suitable
+        if angle < 0 or angle > 360:
+            logger.exception(self.__str__() + ' angle of rotation must be between 0-360 degrees')
+            raise ValueError
+        if angle % 90 != 0:
+            logger.exception(self.__str__() + ' angle of rotation must be a multiple of 90 degrees')
+            raise ValueError
+
+        if axis != 'x' and axis != 'y' and axis != 'z':
+            logger.exception(self.__str__() + ' axis of rotation must be x, y, or z')
+            raise ValueError
+
+        # Coordinates for axis of rotation (centre of object)
+        offset = p1 + (p2 - p1) / 2
+
+        # Move object to axis of rotation
+        p1 -= offset
+        p2 -= offset
+
+        # Calculate rotation matrix
+        r = R.from_euler(axis, angle, degrees=True)
+
+        # Apply rotation
+        p1 = r.apply(p1)
+        p2 = r.apply(p2)
+
+        # Move object back to original axis
+        p1 += offset
+        p2 += offset
+
+        # Get lower left and upper right coordinates to define new object
+        tmp = np.concatenate((p1, p2), axis=0)
+        p1 = np.min(tmp, axis=0)
+        p2 = np.max(tmp, axis=0)
+
+        # For 2D modes check axis of rotation against mode 
+        # and correct invariant coordinate
+        # mode = config.get_model_config().mode
+        mode = 'TMz'
+        if mode == 'TMx':
+            if axis == 'y' or axis =='z':
+                logger.exception(self.__str__() +
+                                 ' axis of rotation must be x for TMx mode models')
+                raise ValueError
+            p1[2] = orig_p1[0]
+            p2[2] = orig_p2[0]
+        elif mode == 'TMy':
+            if axis == 'x' or axis == 'z':
+                logger.exception(self.__str__() +
+                                 ' axis of rotation must be x for TMy mode models')
+                raise ValueError
+            p1[2] = orig_p1[1]
+            p2[2] = orig_p2[1]
+        elif mode == 'TMz':
+            if axis == 'x' or axis == 'y':
+                logger.exception(self.__str__() +
+                                 ' axis of rotation must be x for TMz mode models')
+                raise ValueError
+            p1[2] = orig_p1[2]
+            p2[2] = orig_p2[2]
+
+        # Write points back to original tuple
+        self.kwargs['p1'] = tuple(p1)
+        self.kwargs['p2'] = tuple(p2)