transform coordinates

transform.py

from scipy.spatial.transform import Rotation as R
import numpy as np
import star
import transphire_transform.dump_load.star as trs
import os


input_path = "data/run_it024_data.star"
results_path = "results/test_results/"
result_filename = "run_it024_data"

try:
    os.makedirs(results_path)
except FileExistsError:
    pass

star_data = trs.load_star(input_path)

transformed_data = np.zeros((star_data.shape[0],3))
shift_x = 46
shift_y = -21
shift_z = 0
index_for_tomo = [i for i in range(star_data.shape[0]) if "d02t02" in star_data["MicrographName"][i]]
psis = star_data["AnglePsi"]
rots = star_data["AngleRot"]
tilts = star_data["AngleTilt"]
offx =  star_data["OriginX"]
offy = star_data["OriginY"]
offz = star_data["OriginZ"]

print("Mean OffX:", np.mean(offx[index_for_tomo]))
print("Mean OffY:", np.mean(offy[index_for_tomo]))
print("Mean OffZ:", np.mean(offz[index_for_tomo]))

transforms = [R.from_euler('ZYZ', [psis[i], rots[i], tilts[i]], degrees=True) for i in index_for_tomo]
applied_shift = [r.apply([shift_x,shift_y,shift_z]) for r in transforms]
 
star_data["CoordinateX"].loc[index_for_tomo] = offx[index_for_tomo] + star_data["CoordinateX"][index_for_tomo] + [l[0] for l in applied_shift]
star_data["CoordinateY"].loc[index_for_tomo] = offy[index_for_tomo] + star_data["CoordinateY"][index_for_tomo] +[l[1] for l in applied_shift]
star_data["CoordinateZ"].loc[index_for_tomo] = offz[index_for_tomo] + star_data["CoordinateZ"][index_for_tomo] + [l[2] for l in applied_shift]

print("done")


trs.dump_star(os.path.join(results_path,result_filename+".star"),star_data,"relion_3")
arr = np.array([star_data["CoordinateX"].loc[index_for_tomo] ,star_data["CoordinateY"].loc[index_for_tomo],star_data["CoordinateZ"].loc[index_for_tomo]])
np.savetxt(os.path.join(results_path,result_filename+"_onlycoords.star"),arr.transpose(),fmt="%5.4f")