Working gain pattern class

spatz/sensors/antenna/tx_gain.py

@@ -2,62 +2,113 @@ from numpy.typing import ArrayLike
 from typing import List, AnyStr
 from numpy import matrix
 from typing import List
+import re
+from io import StringIO
+import numpy as np
 
 import pandas as pd
+import math
 
-from spatz.sensors import Sensor
-from spatz.transforms import Transform
-from spatz.dataset import Dataset
-from spatz.logger import Logger
+# from spatz.sensors import Sensor
+# from spatz.transforms import Transform
+# from spatz.dataset import Dataset
+# from spatz.logger import Logger
 
 '''
 Sensor to simulate TX antenna gain in direction of ground station
-You will need to supply a gain pattern in the form of a R^3 matrix in the following form:
 
 Returns the gain in dBi per timestep.
 
-
 gain_pattern: matrix, groundstation_offset_vector
 
 '''
 
 class GainPattern():
     
-    def __init__(self, pattern_file):
-        self._df = pd.read_csv(pattern_file,delimiter='\t')
-        print(self._df)
+    def __init__(self, filepath: str): 
+        # This is a cursed parser. If it breaks, though luck.
+        with open(filepath,"r") as file:
+            # Read Header
+            header = file.readline()
+            header = re.sub(r'\[(.*?)\]',",",header).replace(" ","").replace(",\n",'\n')
 
-    def get_gain(self, phi, theta):
-        phi_left = round(phi,-1)
-        phi_right = round(phi ,-1)
+            # Discard ---- line
+            file.readline()
 
-        theta_left = theta
-        theta_right = theta
-
-class AntennaTxGain(Sensor):
+            # Parse to DF
+            lines = file.readlines()
+            clean_csv = header
+            for line in lines:
+                cleaned = re.sub(r'\s+',',',line).removeprefix(',').removesuffix(',').strip()
+                clean_csv = clean_csv + cleaned + '\n'
+            filelike = StringIO(clean_csv)
+            self._df = pd.read_csv(filelike)
+            print(self._df.head())
 
 
-
-    def __init__(self, dataset: Dataset, logger: Logger, transforms: List[Transform] = []):
-        super().__init__(dataset, logger, transforms)
-
-
-    def _get_data(self) -> ArrayLike | float:
-        # Get current position of rocket
-        [x,y,z] = self._dataset.fetch_values(['x', 'y', 'z'])
         
-        # Get current rotation of rocket
-        [pitch,roll,yaw] = self._dataset.fetch_values(['pitch','roll','yaw'])
+    def __get_gain_internal(self,phi5:float,theta5:float):
+        assert phi5%5 ==0
+        assert theta5%5==0
 
-        # Calculate angle between the vectors
+        row = self._df.loc[(self._df["Theta"] == theta5) & (self._df["Phi"] == phi5)].iloc[0]
+        return row["Abs(Dir.)"]
 
-        # Fetch gain in this direction
+    def get_gain(self, phi, theta) -> float:
+        assert 0 <= phi <= 180
+        assert 0 <= theta <= 360
+        
+        #Interpolate using binlinear interpolation https://en.wikipedia.org/wiki/Bilinear_interpolation
+        phi_lower = math.floor(phi/5)*5
+        phi_upper = phi_lower + 5
+        theta_lower = math.floor(theta/5)*5
+        theta_upper = theta_lower + 5
 
-        return 0
+        G11 = self.__get_gain_internal(phi_lower,theta_lower)
+        G12 = self.__get_gain_internal(phi_lower,theta_upper)
+        G21 = self.__get_gain_internal(phi_upper,theta_lower)
+        G22 = self.__get_gain_internal(phi_upper,theta_upper)
 
-    def _sensor_specific_effects(self, x: ArrayLike) -> ArrayLike:
-        return x
+        v1 = np.array([phi_upper-phi,phi-phi_lower])
+        v2 = np.array([[theta_upper-theta],[theta-theta_lower]])
+        A = np.array([[G11,G12],[G21,G22]])
+
+        interpolated = 1/25 * v1 @ A @ v2
+
+        return interpolated[0]
+        
+
+
+# class AntennaTxGain(Sensor):
+
+#     def __init__(self, dataset: Dataset, logger: Logger, transforms: List[Transform] = []):
+#         super().__init__(dataset, logger, transforms)
+
+
+#     def _get_data(self) -> ArrayLike | float:
+#         # Get current position of rocket
+#         [x,y,z] = self._dataset.fetch_values(['x', 'y', 'z'])
+        
+#         # Get current rotation of rocket
+#         [pitch,roll,yaw] = self._dataset.fetch_values(['pitch','roll','yaw'])
+
+#         # Calculate angle between the vectors
+
+#         # Fetch gain in this direction
+
+#         return 0
+
+#     def _sensor_specific_effects(self, x: ArrayLike) -> ArrayLike:
+#         return x
     
 
-    def _get_name(self) -> AnyStr:
-        return 'Generic Antenna TX'
+#     def _get_name(self) -> AnyStr:
+#         return 'Generic Antenna TX'
+        
+
+if  __name__ == '__main__':
+    pattern = GainPattern("data/gain_pattern/farfield_all.txt")
+    print(pattern.get_gain(0,12))
+    print(pattern.get_gain(0,16))
+    print(pattern.get_gain(6,12))
+    print(pattern.get_gain(0,10))