blablub

bla.py

#!/usr/bin/python
###This script is to convert the xml file after picking to .coords files (txt).
###It can take the rotation applied before picking.
###It will also split the coordinates based on boundary line drawn during picking.


import argparse
import numpy as np
import subprocess
import os


# argparser.add_argument('--separate', help="Normalisation when generating the stacks.")


class CoordTransform:

    def __init__(
            self,
            thickness_rot,
            thickness_ori,
            rot_angles_xyz,
            tomogram_dimensions_original,
            tomogram_width_rot,
    ):
        self.thickness_rot = thickness_rot
        self.thickness_ori = thickness_ori
        self.rot_angles_xyz = rot_angles_xyz
        self.tomogram_dimensions_original = tomogram_dimensions_original #[width,height], default 479.5, 463.5
        self.tomogram_width_rot = tomogram_width_rot # default 1400
        print("Thickness_rot:", thickness_rot)
        print("thickness_ori:", thickness_ori)
        print("Angles:", rot_angles_xyz)
        print("Dimension Original:", tomogram_dimensions_original)
        print("tomo_width_rot:", tomogram_width_rot)


    def rotate_points(self, coords_raw):
        # rotation matrix
        # g = open(info_file)
        # lines = g.readlines()
        # rotation = -float(lines[0].rstrip().split(": ")[1])
        # g.close()
        # rotang = np.radians(rotation)
        ################ This is the matrix
        # rotM = np.array([[np.cos(rotang), -np.sin(rotang), 0],[np.sin(rotang), np.cos(rotang),0],[0, 0, 1]])

        # Translation to shift the origin to the center

        coords_columnsplit = np.hsplit(coords_raw, 4)
        coords = np.column_stack(
            (coords_columnsplit[0], coords_columnsplit[1], coords_columnsplit[2])
        )
        recoords_matrix = np.array([-self.tomogram_width_rot/2, -self.tomogram_width_rot/2, -float(self.thickness_rot / 2)])
        coords_shifted = np.add(coords, recoords_matrix)

        # Invert coords array for rotation
        coords_inv = coords_shifted.swapaxes(0, 1)

        # Calculate rotation matrix
        rotx, roty, rotz = self.rot_angles_xyz
        R_z = np.array(
            [[np.cos(rotz), -np.sin(rotz), 0], [np.sin(rotz), np.cos(rotz), 0], [0, 0, 1]]
        )
        R_y = np.array(
            [[np.cos(roty), 0, np.sin(roty)], [0, 1, 0], [-np.sin(roty), 0, np.cos(roty)]]
        )
        R_x = np.array(
            [[1, 0, 0], [0, np.cos(rotx), -np.sin(rotx)], [0, np.sin(rotx), np.cos(rotx)]]
        )
        R = np.dot(R_x, np.dot(R_y, R_z))
        R_inv = np.linalg.inv(R)

        # Rotate the points
        new_coords_inv = np.dot(R_inv, coords_inv)

        # Revert coords and shift back the origin to the corner
        new_coords = new_coords_inv.swapaxes(0, 1)
        recoords_matrix_inv = np.array(
            [self.tomogram_dimensions_original[0]/2, self.tomogram_dimensions_original[1]/2, float(self.thickness_ori / 2)]
        )  # Half dimension orignal tomogram
        new_coords_recoords = np.add(new_coords, recoords_matrix_inv)
        final = np.column_stack((new_coords_recoords, coords_columnsplit[3]))
        return final

    # separate particles before any rotation
    def separate_particles(self, info_file, coords_array):
        """

        :param info_file: Defines roughly a plane
        :param coords_array: Tracing results as array. Columns: x,y,z,filament_id
        :return:
        """
        f = open(info_file)
        lines = f.readlines()
        x11, y11, x12, y12, z1 = [float(i) for i in lines[3].rstrip().split(", ")]
        y11 = self.tomogram_width_rot - y11
        y12 = self.tomogram_width_rot - y12
        x21, y21, x22, y22, z2 = [float(i) for i in lines[6].rstrip().split(", ")]
        y21 = self.tomogram_width_rot - y21
        y22 = self.tomogram_width_rot - y22
        m1 = (y12 - y11) / (x12 - x11)
        m2 = (y22 - y21) / (x22 - x21)
        c1 = y11 - m1 * x11
        c2 = y21 - m2 * x21
        if z1 >= z2:
            zp = [z2, z1]
            mp = [m2, m1]
            cp = [c2, c1]
        else:
            zp = [z1, z2]
            mp = [m1, m2]
            cp = [c1, c2]

        top_class = np.array([])
        bot_class = np.array([])

        for i in coords_array:
            x = i[0]
            y = i[1]
            z = i[2]
            fid = i[3]
            current_m = np.interp(z, zp, mp)
            current_c = np.interp(z, zp, cp)
            critical_y = current_m * x + current_c
            if y >= critical_y:
                if top_class.any():
                    top_class = np.concatenate((top_class, np.array([i])))
                else:
                    top_class = np.array([i])
            else:
                if bot_class.any():
                    bot_class = np.concatenate((bot_class, np.array([i])))
                else:
                    bot_class = np.array([i])
        return top_class, bot_class

    def writefile(self, coords_raw, info_file, tomoname, output_path = ""):
        all = coords_raw
        #all = self.rotate_points(coords_raw)
        #top_raw, bot_raw = self.separate_particles(info_file, coords_raw)
        #top = self.rotate_points(top_raw)
        #bot = self.rotate_points(bot_raw)
        all_withoutfid = np.column_stack(
            (np.hsplit(all, 4)[0], np.hsplit(all, 4)[1], np.hsplit(all, 4)[2])
        )
        tomoname_out = os.path.join(output_path,tomoname)
        ###Write out all files
        #np.savetxt("%s_bin4_top_fid.coords" % (tomoname_out), top, fmt="%.2f")
        #np.savetxt("%s_bin4_bot_fid.coords" % (tomoname_out), bot, fmt="%.2f")

        np.savetxt("%s_bin4.coords" % (tomoname_out), all_withoutfid, fmt="%.2f")
        np.savetxt("%s_bin4_fid.coords" % (tomoname_out), all, fmt="%.2f")

        try:
            imod_command = (
                    "point2model -in %s_bin4.coords -ou %s_bin4.mod -scat -sphere 4"
                    % (tomoname_out, tomoname_out)
            )
            subprocess.call(imod_command, shell=True)
        except Exception as e:
            print("Point2Model cannot be executed.")
            print(e)
            pass

    def convert(self, results_tracing_dataframe, windowwidth, boxdistance):
        """
        Converted the tracing results into coordinates for export
        :param results_tracing_dataframe: Pandas dataframe with columns x,y,frame,particle
        :param windowwidth: Window with for moving average
        :param boxdistance: Box distance for resampling
        :return: converted coordinates
        """

        num_particles = np.max(results_tracing_dataframe["particle"])

        all_filament_xcoords = []
        all_filament_ycoords = []
        all_filament_zcoords = []
        all_filament_ids = []
        for trace_id in range(num_particles):
            one_trace = results_tracing_dataframe[results_tracing_dataframe["particle"] == trace_id]
            fil_xcoords = one_trace["x"].tolist()
            fil_ycoords = one_trace["frame"].tolist()

            fil_zcoords = (self.thickness_rot-one_trace["y"]).tolist()

            fil_xcoords, fil_ycoords, fil_zcoords = moving_average(
                fil_xcoords, fil_ycoords, fil_zcoords, windowwidth
            )
            fil_xcoords, fil_ycoords, fil_zcoords, fil_fids = resample_filament(
                fil_xcoords, fil_ycoords, fil_zcoords, trace_id, boxdistance
            )

            all_filament_xcoords.extend(fil_xcoords)
            all_filament_ycoords.extend(fil_ycoords)
            all_filament_zcoords.extend(fil_zcoords)
            all_filament_ids.extend(fil_fids)

        coords_raw = np.column_stack(
            (
                np.array(all_filament_xcoords),
                np.array(all_filament_ycoords),
                np.array(all_filament_zcoords),
                np.array(all_filament_ids),
            )
        )
        return coords_raw


def dist_squared(ax, ay, az, bx, by, bz):
    return np.square(ax - bx) + np.square(ay - by) + np.square(az - bz)


def moving_average(fil_xcoords, fil_ycoords, fil_zcoords, window_width):
    len_filament = len(fil_xcoords)
    if len_filament < (window_width + 1):
        return fil_xcoords, fil_ycoords, fil_zcoords

    new_x = []
    new_y = []
    new_z = []
    offset = int((window_width - 1) / 2)
    for i in range(offset):
        new_x.append(fil_xcoords[i])
        new_y.append(fil_ycoords[i])
        new_z.append(fil_zcoords[i])

    for i in range(offset, len_filament - offset):
        mean_x = 0
        mean_y = 0
        mean_z = 0
        mean_window = range(i - offset, i + offset + 1)
        for j in mean_window:
            mean_x += fil_xcoords[j]
            mean_y += fil_ycoords[j]
            mean_z += fil_zcoords[j]
        mean_x = mean_x / len(mean_window)
        mean_y = mean_y / len(mean_window)
        mean_z = mean_z / len(mean_window)

        new_x.append(mean_x)
        new_y.append(mean_y)
        new_z.append(mean_z)

    for i in range(len_filament - offset, len(fil_xcoords)):
        new_x.append(fil_xcoords[i])
        new_y.append(fil_ycoords[i])
        new_z.append(fil_zcoords[i])

    return new_x, new_y, new_z


def getEquidistantFilamentPoints(xcoords, ycoords, zcoords, posi, posj, parts):
    points = zip(
        np.linspace(xcoords[posi], xcoords[posj], parts + 1, endpoint=True),
        np.linspace(ycoords[posi], ycoords[posj], parts + 1, endpoint=True),
        np.linspace(zcoords[posi], zcoords[posj], parts + 1, endpoint=True),
    )
    newx = []
    newy = []
    newz = []

    for point in points:
        newx.append(point[0])
        newy.append(point[1])
        newz.append(point[2])
    return newx, newy, newz


def resample_filament(fil_xcoords, fil_ycoords, fil_zcoords, fil_id, distance):
    len_filament = len(fil_xcoords)
    newxcand = []
    newycand = []
    newzcand = []
    for i in range(len_filament - 1):
        eqx, eqy, eqz = getEquidistantFilamentPoints(
            fil_xcoords, fil_ycoords, fil_zcoords, i, i + 1, parts=100
        )
        for k in range(len(eqx)):
            newxcand.append(eqx[k])
            newycand.append(eqy[k])
            newzcand.append(eqz[k])

    len_filament = len(newxcand)
    newx = []
    newy = []
    newz = []
    newfilid = []
    dist = -1
    sqdistance = distance * distance
    for posi in range(len_filament):
        if dist == -1:
            newx.append(newxcand[posi])
            newy.append(newycand[posi])
            newz.append(newzcand[posi])
            newfilid.append(fil_id)
            dist = 0
        else:
            dist = dist_squared(
                newx[-1],
                newy[-1],
                newz[-1],
                newxcand[posi],
                newycand[posi],
                newzcand[posi],
            )
            if dist > sqdistance:

                if not (
                    newxcand[posi] == newx[-1]
                    and newycand[posi] == newy[-1]
                    and newzcand[posi] == newz[-1]
                ):
                    newx.append(newxcand[posi])
                    newy.append(newycand[posi])
                    newz.append(newzcand[posi])
                    newfilid.append(fil_id)

    return newx, newy, newz, newfilid




'''
args = argparser.parse_args()

filename = args.inputxml
boxdistance = args.boxdistance
windowwidth = args.windowsize
thickness_ori = args.thickness_ori
thickness_rot = args.thickness_rot
info_file = args.infofile
tomoname = args.tomoname
rotx = np.radians(args.rotzyx[2])
roty = np.radians(args.rotzyx[1])
rotz = np.radians(args.rotzyx[0])

xcoords = []
ycoords = []
zcoords = []
fids = []

xmldoc = minidom.parse(filename)

filaments = xmldoc.getElementsByTagName("particle")

for fil_id, fil in enumerate(filaments):

    print(fil_id)
    positions = fil.getElementsByTagName("detection")
    fil_xcoords = []
    fil_ycoords = []
    fil_zcoords = []
    for pos_id, pos in enumerate(positions):
        fil_ycoords.append(float(pos.attributes["t"].value))
        fil_xcoords.append(float(pos.attributes["x"].value))
        fil_zcoords.append(float(pos.attributes["y"].value))

    fil_xcoords, fil_ycoords, fil_zcoords = moving_average(
        fil_xcoords, fil_ycoords, fil_zcoords, windowwidth
    )
    fil_xcoords, fil_ycoords, fil_zcoords, fil_fids = resample_filament(
        fil_xcoords, fil_ycoords, fil_zcoords, fil_id, boxdistance
    )

    # for y in fil_ycoords:
    #    print("After resampling:", y)

    xcoords.extend(fil_xcoords)
    ycoords.extend(fil_ycoords)
    zcoords.extend(fil_zcoords)
    fids.extend(fil_fids)

# create matrix of coordinates
coords_raw = np.column_stack(
    (np.array(xcoords), np.array(ycoords), np.array(zcoords), np.array(fids))
)

# x_newarray = np.array(xcoords)-700
# y_newarray = np.array(ycoords)-700
# z_newarray = np.array(zcoords)-float(thickness/2)

writefile()
'''