add k warping function

sed/calibrator/hextof.py

@@ -152,3 +152,75 @@ def convert_to_ns(x):
         "tof_binwidth": tof_binwidth
     }
     return df, metadata
+
+
+
+
+def calibrate_k(
+        df: Union[pd.DataFrame, dask.dataframe.DataFrame],
+        warp_params: Sequence[float] = None,
+        x_column: str = "dldPosX",
+        y_column: str = "dldPosY",
+        kx_column: str = "kx",
+        ky_column: str = "ky",
+        config: dict = None,
+) -> Tuple[Union[pd.DataFrame, dask.dataframe.DataFrame], dict]:
+    """Converts the 8s time in steps to time in ns.
+
+    Args:
+        warp_params (Sequence[float], optional): warping parameters.
+            this function returns the distorted coordinates given the undistorted ones
+            a little complicated by the fact that (warp_params[6],warp_params[7]) needs to go to (0,0)
+            it uses a radial distortion model called division model (https://en.wikipedia.org/wiki/Distortion_(optics)#Software_correction)
+            commonly used to correct for lens artifacts
+            warp_params[0],warp_params[1] center of distortion in px
+            warp_params[6],warp_params[7] normal emission (Gamma) in px
+            warp_params[2],warp_params[3],warp_params[4] K_n; rk = rpx/(K_0 + K_1*rpx^2 + K_2*rpx^4)
+            warp_params[5] rotation in rad
+            Defaults to config["dataframe"]["warp_params"].
+        x_column (str, optional): Name of the column containing the
+            x steps. Defaults to config["dataframe"]["x_column"].
+        y_column (str, optional): Name of the column containing the
+            y steps. Defaults to config["dataframe"]["y_column"].
+        kx_column (str, optional): Name of the column containing the
+            x steps. Defaults to config["dataframe"]["kx_column"].
+        ky_column (str, optional): Name of the column containing the
+            y steps. Defaults to config["dataframe"]["ky_column"].
+    """
+    if warp_params is None:
+        if config is None:
+            raise ValueError("Either warp_params or config must be given.")
+        warp_params: float = config["dataframe"]["warp_params"]
+    if x_column is None:
+        if config is None:
+            raise ValueError("Either x_column or config must be given.")
+        x_column: str = config["dataframe"]["x_column"]
+    if kx_column is None:
+        if config is None:
+            raise ValueError("Either kx_column or config must be given.")
+        kx_column: str = config["dataframe"]["kx_column"]
+    if y_column is None:
+        if config is None:
+            raise ValueError("Either y_column or config must be given.")
+        y_column: str = config["dataframe"]["y_column"]
+    if ky_column is None:
+        if config is None:
+            raise ValueError("Either ky_column or config must be given.")
+        ky_column: str = config["dataframe"]["ky_column"]
+
+    def convert_to_kx(x):
+        return np.sqrt(np.power((x[x_column]-warp_params[0]),2)+np.power((x[y_column]-warp_params[1]),2))/(warp_params[2]+warp_params[3]*np.power(np.sqrt(np.power((x[x_column]-warp_params[0]),2)+np.power((x[y_column]-warp_params[1]),2)),2)+warp_params[4]*np.power(np.sqrt(np.power((x[x_column]-warp_params[0]),2)+np.power((x[y_column]-warp_params[1]),2)),4)) * np.cos(np.arctan2(x[y_column]-warp_params[1],x[x_column]-warp_params[0])-warp_params[5])-np.sqrt(np.power((warp_params[6]-warp_params[0]),2)+np.power((warp_params[7]-warp_params[1]),2))/(warp_params[2]+warp_params[3]*(np.power((warp_params[6]-warp_params[0]),2)+np.power((warp_params[7]-warp_params[1]),2))+warp_params[4]*np.power((np.power((warp_params[6]-warp_params[0]),2)+np.power((warp_params[7]-warp_params[1]),2)),2))*np.cos(np.arctan2(warp_params[7]-warp_params[1],warp_params[6]-warp_params[0])-warp_params[5])
+    def convert_to_ky(x):
+        return np.sqrt(np.power((x[x_column]-warp_params[0]),2)+np.power((x[y_column]-warp_params[1]),2))/(warp_params[2]+warp_params[3]*np.power(np.sqrt(np.power((x[x_column]-warp_params[0]),2)+np.power((x[y_column]-warp_params[1]),2)),2)+warp_params[4]*np.power(np.sqrt(np.power((x[x_column]-warp_params[0]),2)+np.power((x[y_column]-warp_params[1]),2)),4)) * np.sin(np.arctan2(x[y_column]-warp_params[1],x[x_column]-warp_params[0])-warp_params[5])-np.sqrt(np.power((warp_params[6]-warp_params[0]),2)+np.power((warp_params[7]-warp_params[1]),2))/(warp_params[2]+warp_params[3]*(np.power((warp_params[6]-warp_params[0]),2)+np.power((warp_params[7]-warp_params[1]),2))+warp_params[4]*np.power((np.power((warp_params[6]-warp_params[0]),2)+np.power((warp_params[7]-warp_params[1]),2)),2))*np.sin(np.arctan2(warp_params[7]-warp_params[1],warp_params[6]-warp_params[0])-warp_params[5])
+    df[kx_column] = df.map_partitions(
+        convert_to_kx, meta=(kx_column, np.float64)
+    )
+    df[ky_column] = df.map_partitions(
+        convert_to_ky, meta=(ky_column, np.float64)
+    )
+
+    metadata = {}
+    metadata["applied"] = True
+    metadata["warp_params"] = warp_params
+
+    return df, metadata