From 2e81b6bd238c61c1757d18ac9f3d35c2a75da7a4 Mon Sep 17 00:00:00 2001 From: Davide Curcio Date: Wed, 11 Oct 2023 14:46:32 +0200 Subject: [PATCH] add k warping function --- sed/calibrator/hextof.py | 72 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/sed/calibrator/hextof.py b/sed/calibrator/hextof.py index 58e2aee2..31e70243 100644 --- a/sed/calibrator/hextof.py +++ b/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 \ No newline at end of file