From 5aadbf488fd471a499e47d3a45f928d00507d804 Mon Sep 17 00:00:00 2001 From: Qinqing Liu Date: Mon, 16 Nov 2020 16:49:23 -0500 Subject: [PATCH] test on linux test build, generation, parse, visualization --- README.md | 8 +- octree/.DS_Store | Bin 10244 -> 10244 bytes octree/tools/build_octree.cpp | 4 +- octree/tools/transform_points.cpp | 2 +- pdbbind/bash_scripts/correct_index0.sh | 2 +- pdbbind/python/.DS_Store | Bin 6148 -> 6148 bytes pdbbind/python/octree_parse.py | 138 ++++ pdbbind/python/utils.py | 1036 ++++++++++++++++++++++++ 8 files changed, 1184 insertions(+), 6 deletions(-) create mode 100644 pdbbind/python/octree_parse.py create mode 100644 pdbbind/python/utils.py diff --git a/README.md b/README.md index baaee51..81c1a7b 100644 --- a/README.md +++ b/README.md @@ -27,7 +27,9 @@ conda install -c conda-forge yacs tqdm pip install sklearn pip install matplotlib pip install seaborn -pip install -U openbabel +# uncomment if want to visualize +# pip install vtk +conda install -c conda-forge openbabel cd octree/external && git clone --recursive https://github.com/wang-ps/octree-ext.git cd .. && mkdir build && cd build @@ -43,7 +45,9 @@ cd ../../ #### Octree Generation Example ```angular2 -cd pdbbind/data_folder/refined-set/1a1e +cd pdbbind/data_example/pdbbind +cp ../../bash_scripts/octree_for_single.sh . +cd 1a1e bash ../octree_for_single.sh cd octree_folder ls -lh diff --git a/octree/.DS_Store b/octree/.DS_Store index c376d1a98175b31f94e815b69bf87e00ee3023b6..54d263c8d76a4e665667c4a25290399d2249cfb2 100644 GIT binary patch delta 198 zcmZn(XbIR5E;#v&^ts6@LLr=DqT&*wl2TGKlN*JKWa9<+i<67;b8-Umi!)0y^YbF3 z(~B}w{L|8kQ%fcb2>VKdRAl6rdnV`S1r(*0Wu}%#WCGd2nN_KiD}>cTlpRtQFn|Du z0AomAeqJgkr?Nu;SOh4>pv~aO;L8xkkj+rVFp*&{!+M6@3?~@wGQ4N_$;i$q$*9Dr q&1k}C%V^K&!5GLG#2CUD$rv>`MNW0IxX63%&Fl)lST>0=GXntc$2J%M delta 194 zcmZn(XbIR5F35Ot@;T`SL5b>WV*@iC1w#X~S{;RIbMwgwGUALgCs)YGGtQc%}At@y-HMvQsXmXBhu~57Ke{n{Bxo2{IUO-W5S!Qba7GhPw>!8GbUdGYT*&F={iK qFxoQOGkP$3GI}$HFh(*)F(yw=ms8#>A@ZJkGn>M1md#9}%*+7VC@_%# diff --git a/octree/tools/build_octree.cpp b/octree/tools/build_octree.cpp index 08c8f86..af42630 100644 --- a/octree/tools/build_octree.cpp +++ b/octree/tools/build_octree.cpp @@ -17,10 +17,10 @@ const float kPI = 3.14159265f; DEFINE_string(filenames, kRequired, "", "The input filenames"); DEFINE_string(output_path, kOptional, ".", "The output path"); -DEFINE_string(axis, kOptional, "y", "The upright axis of the input model"); +DEFINE_string(axis, kOptional, "xyz", "The upright axis of the input model"); DEFINE_int(depth, kOptional, 6, "The maximum depth of the octree"); DEFINE_int(full_depth, kOptional, 2, "The full layer of the octree"); -DEFINE_int(rot_num, kOptional, 12, "Number of poses rotated along the upright axis"); +DEFINE_int(rot_num, kOptional, 1, "Number of poses rotated along the upright axis"); DEFINE_float(offset, kOptional, 0.0f, "The offset value for handing thin shapes"); DEFINE_bool(node_dis, kOptional, false, "Output per-node displacement"); DEFINE_bool(node_feature, kOptional, false, "Compute per node feature"); diff --git a/octree/tools/transform_points.cpp b/octree/tools/transform_points.cpp index 47c5553..b0b7c50 100644 --- a/octree/tools/transform_points.cpp +++ b/octree/tools/transform_points.cpp @@ -12,7 +12,7 @@ DEFINE_string(output_path, kOptional, ".", "The output path"); DEFINE_float(scale, kOptional, 1.0, "The scale factor"); DEFINE_float(trans, kOptional, 0.0, "The translation factor"); DEFINE_float(offset, kOptional, 0.0, "Offset the points along its normal"); -DEFINE_float(ratio, kOptional, 0.5, "The dropout ratio"); +DEFINE_float(ratio, kOptional, 0.0, "The dropout ratio"); DEFINE_float(dim, kOptional, 0, "The resolution for dropout"); DEFINE_float(std_nm, kOptional, 0.0, "The std of normal noise "); DEFINE_float(std_pt, kOptional, 0.0, "The std of posistion noise"); diff --git a/pdbbind/bash_scripts/correct_index0.sh b/pdbbind/bash_scripts/correct_index0.sh index b234a49..9381d8d 100644 --- a/pdbbind/bash_scripts/correct_index0.sh +++ b/pdbbind/bash_scripts/correct_index0.sh @@ -1,6 +1,6 @@ for D in ./*/; do cd "${D}"; - python ../../python/clean_index_error.py + python ../../../python/clean_index_error.py cd ..; done \ No newline at end of file diff --git a/pdbbind/python/.DS_Store b/pdbbind/python/.DS_Store index de8af17ec13ca487804a74e9878d63c28739817b..034c068f6ce1716411f79ff6c422b75a52bb82ea 100644 GIT binary patch delta 129 zcmZoMXfc@J&&a+pU^gQp`(_>{e#T^ehJ1!(h7yJ%hEyPoXDDDu1hR{PY&{^m(laMN zIVmSUiGhJZ0En%Dc*cJ)0J0buI2cNSdNUbvfI6_LUjbAzb+Z6-C(CAbj=%f>9z7sZ delta 32 ocmZoMXfc@J&&awlU^gQp>t-G%e#Xt6%yldi8}c@@bNuB80H5Ot@Bjb+ diff --git a/pdbbind/python/octree_parse.py b/pdbbind/python/octree_parse.py new file mode 100644 index 0000000..5ca0453 --- /dev/null +++ b/pdbbind/python/octree_parse.py @@ -0,0 +1,138 @@ +import sys +import numpy as np +sys.path.insert(0, '../../octree/build/python') +import pyoctree +import numpy as np +import os + + + + +pro_lig_flag_pos = 18 +kLeaf_flag = -1 + +# Can modify below +visualize = False # True or False; If True, will generate vtk files, and can be visualized in Paraview +target_range = 32 # The radius of bounding box, default 32 generate a 64 * 64 * 64 Angstrom bounding box. +line_depth = 5 # Only plot octree division segment lines < line_depth. + + + +def convert_key_to_xyz(key, depth): + xyz = [0,0,0] + for i in range(depth): + for j in range(3): + mask = 1 << (3 * i + 2 - j) + xyz[j] |= (key & mask) >> (2 * i + 2 - j) + return xyz + +def octree_info_to_cubic_dic(depth, node_number, keys, children, features): + cubic_dic = {} + for node_idx_d in range(node_number): + child = children[node_idx_d] + if child == kLeaf_flag: + continue + key = keys[node_idx_d] + xyz = convert_key_to_xyz(key, depth) + [x,y,z] = xyz + key_for_cubic_dic = '{}_{}_{}_{}'.format(depth, x,y,z) + if features[node_idx_d][pro_lig_flag_pos] == 1: + cubic_dic[key_for_cubic_dic] = 'Protein' + elif features[node_idx_d][pro_lig_flag_pos] == 0: + cubic_dic[key_for_cubic_dic] = 'Ligand' + else: + cubic_dic[key_for_cubic_dic] = 'Mixture' + return cubic_dic + + + + + + + + +def vtk_from_cubic_dic(cubic_dic, total_depth, complex_id, differ=False, + shift = [0,0,0], save_folder = './results'): + vtk_line_dic = {} + for i in range(total_depth + 1): + current_depth = total_depth - i + base = target_range*2 / 2 ** current_depth + vtk_dic = {} + color = get_color_for_dep(current_depth + 1, differ) + for key, value in cubic_dic.items(): + [depth, x, y, z] = [int(item) for item in key.split('_')] + coord_x = x * base - target_range + coord_y = y * base - target_range + coord_z = z * base - target_range + min_coord = [coord_x, coord_y, coord_z] + max_coord = [coord_x + base, coord_y + base, coord_z + base] + cubic = Cubic([min_coord, max_coord], key) + cubic.set_color_new(value) + vtk_dic = cubic.load_to_vtk_opt(vtk_dic) + if current_depth < line_depth: # current_depth < depth for all. + vtk_line_dic = cubic.write_oct_line_opt(vtk_line_dic, color) + + if not os.path.isdir('{}/{}/'.format(save_folder, complex_id)): + os.makedirs('{}/{}/'.format(save_folder, complex_id)) + file_name = '{}/{}/{}_Octree_{}_{}'.format(save_folder, complex_id, complex_id, target_range, current_depth) + vtk_writer(vtk_dic=vtk_dic, file_name = file_name, shift = shift) + + if current_depth > 0: + upper_layer_dic = get_upper_layer_dic(cubic_dic) + cubic_dic = upper_layer_dic + + color = get_color_for_dep(0, differ) + vtk_line_dic = cubic.write_edge_line_opt(vtk_line_dic, color) + + pts, lines, colors = fill_vtk_lines(vtk_line_dic, shift) + + linesPolyData = vtk.vtkPolyData() + linesPolyData.SetPoints(pts) + linesPolyData.SetLines(lines) + linesPolyData.GetCellData().SetScalars(colors) + + ShellVolmapper = vtk.vtkDataSetMapper() + ShellVolmapper.SetScalarModeToUseCellData() + ShellVolmapper.UseLookupTableScalarRangeOn() + ShellVolmapper.SetInputData(linesPolyData) + ShellVolactor = vtk.vtkActor() + ShellVolactor.SetMapper(ShellVolmapper) + + modelwriter = vtk.vtkDataSetWriter() + modelwriter.SetFileName('{}/{}/{}_segline_{}_{}.vtk'.format(save_folder, complex_id, complex_id, target_range, line_depth)) + modelwriter.SetInputData(ShellVolmapper.GetInput()) + modelwriter.Write() + print('Save {}/{}/{}_segline_{}_{}.vtk'.format(save_folder, complex_id, complex_id, target_range, line_depth)) + + +if __name__ == "__main__": + # parse octree + id = '1a1e' + + oct = pyoctree.Octree() + print('Parse the 1a1e_points_10_2_000.octree just generated.') + oct.read_octree("../data_example/pdbbind/{}/octree_folder/{}_points_10_2_000.octree".format(id, id)) + + num_channel = oct.num_channel() + print("Number of channel: {}".format(num_channel)) + + depth = oct.depth() + features = oct.features(depth) + node_number = oct.num_nodes(depth) + children = oct.children(depth) + keys = oct.keys(depth) + + print("Depth: {}".format(depth)) + print("Number of nodes at depth {}: {}".format(depth, node_number)) + + features = np.array(features).reshape((num_channel, node_number)) + features = np.swapaxes(features, 0, 1) + + # for visualization + if visualize: + import vtk + from utils import * + print("Start Generate vtk files for Visualization in Paraview Software") + cubic_dic = octree_info_to_cubic_dic(depth, node_number, keys, children, features) + vtk_from_cubic_dic(cubic_dic, depth, id, differ = False, save_folder = './') + diff --git a/pdbbind/python/utils.py b/pdbbind/python/utils.py new file mode 100644 index 0000000..8925f5c --- /dev/null +++ b/pdbbind/python/utils.py @@ -0,0 +1,1036 @@ +# utils for adaptive_octree_opt.py. + + +import numpy as np +import vtk +import time + +def read_xyz_file(file_name): + with open('{}'.format(file_name), 'r') as data_fid: + points, x, y, z = [], [], [], [] + for row in data_fid.readlines(): + row = row.replace('\n', '') + numbers = row.split(' ') + try: + point_coord = [float(item) for item in numbers[0:3]] + point_x = point_coord[0] + point_y = point_coord[1] + point_z = point_coord[2] + points.append(point_coord) + x.append(point_x) + y.append(point_y) + z.append(point_z) + except ValueError: + print('ValueError: {}, row {}'.format(file_name, row)) + except IndexError: + print('IndexError: {}, row {}'.format(file_name, row)) + else: + pass + return np.array(x), np.array(y), np.array(z) + +def read_feature_file(file_name): + with open('{}'.format(file_name), 'r') as data_fid: + features = [] + for row in data_fid.readlines(): + row = row.replace('\n', '') + row_features = row.split(' ') + try: + row_features = [float(item) for item in row_features] + features.append(row_features) + except ValueError: + print('ValueError: {}, row {}'.format(file_name, row)) + except IndexError: + print('IndexError: {}, row {}'.format(file_name, row)) + else: + pass + return features + +class Point(object): + def __init__(self, coord): + self.coord = coord + self.string = '{}_{}_{}'.format(coord[0], coord[1], coord[2]) + self.order = None + + def get_coord(self): + return self.coord + +class Cubic(object): + """ + A class that describe a cubic in 3D space. + Can detect if this cubic contain points, then spilt into 8 child cubics and plot. + """ + def __init__(self, coord, id): + self.id = id + min_coord = coord[0] + max_coord = coord[1] + self.min_coord = min_coord + self.max_coord = max_coord + self.min_x = min_coord[0] + self.min_y = min_coord[1] + self.min_z = min_coord[2] + self.max_x = max_coord[0] + self.max_y = max_coord[1] + self.max_z = max_coord[2] + + self.dot1 = [self.min_x, self.min_y, self.min_z] + self.dot2 = [self.min_x, self.min_y, self.max_z] + self.dot3 = [self.min_x, self.max_y, self.min_z] + self.dot4 = [self.min_x, self.max_y, self.max_z] + self.dot5 = [self.max_x, self.min_y, self.min_z] + self.dot6 = [self.max_x, self.min_y, self.max_z] + self.dot7 = [self.max_x, self.max_y, self.min_z] + self.dot8 = [self.max_x, self.max_y, self.max_z] + + self.resolution = self.max_z - self.min_z + self.mid_point = [(self.min_x + self.max_x)/2, (self.min_y + self.max_y)/2, (self.min_z + self.max_z)/2] + + def set_color_new(self, value): + if value == 'Protein': + self.color = [.4, .4, 1, 1] + self.color_scale = 11 + elif value == 'Ligand': + self.color = [.4, 1, .4, 1] + self.color_scale = 12 + elif value == 'Mixture': + self.color = [.4, 1, 1, 1] + self.color_scale = 13 + elif value == 'Empty': + self.color = [1, 1, 1, 1] + self.color_scale = 14 + else: + print('error in assign color') + + def member_test(self, points): + if points[0] >= self.min_x and points[0] <= self.max_x and \ + points[1] >= self.min_y and points[1] <= self.max_y and \ + points[2] >= self.min_z and points[2] <= self.max_z: + return True + else: + return False + + def load_to_vtk(self, vtk_dic): + if 'points' not in vtk_dic: + vtk_dic['points'] = [] + vtk_dic['num_points'] = 0 + vtk_dic['cells'] = [] + vtk_dic['num_cells'] = 0 + vtk_dic['cell_color_scale'] = [] + + vtk_dic['Ligand'] = 0 + vtk_dic['Protein'] = 0 + vtk_dic['Mixture'] = 0 + vtk_dic['Empty'] = 0 + + vtk_dic['points'] += [self.dot1, self.dot2, self.dot3, self.dot4, self.dot5, self.dot6, self.dot7, self.dot8] + index_start = vtk_dic['num_points'] + + # use all 6 surfaces as cells. + new_cells = [[4, 0 + index_start, 1 + index_start, 3 + index_start, 2 + index_start], + [4, 4 + index_start, 5 + index_start, 7 + index_start, 6 + index_start], + [4, 0 + index_start, 4 + index_start, 6 + index_start, 2 + index_start], + [4, 2 + index_start, 6 + index_start, 7 + index_start, 3 + index_start], + [4, 3 + index_start, 7 + index_start, 5 + index_start, 1 + index_start], + [4, 0 + index_start, 4 + index_start, 5 + index_start, 1 + index_start]] + vtk_dic['cells'] += new_cells + vtk_dic['num_cells'] += 6 + for i in range(6): + vtk_dic['cell_color_scale'].append(self.color_scale) + + if self.color_scale == 11: + vtk_dic['Protein'] += 1 + elif self.color_scale == 12: + vtk_dic['Ligand'] += 1 + elif self.color_scale == 13: + vtk_dic['Mixture'] += 1 + elif self.color_scale == 14: + vtk_dic['Empty'] += 1 + + vtk_dic['num_points'] += 8 + if len(vtk_dic['points']) != vtk_dic['num_points']: + print('points error') + print(len(vtk_dic['points']), vtk_dic['num_points']) + if len(vtk_dic['cells']) != vtk_dic['num_cells']: + print('cell error') + return vtk_dic + + # print('number of ligand cells: {}'.format(vtk_dic['ligand'])) + # print('number of pocket cells: {}'.format(vtk_dic['pocket'])) + # print('number of mix cells: {}'.format(vtk_dic['mix'])) + # print('number of empty cells: {}'.format(vtk_dic['empty'])) + # print('number of total cells: {}'.format(vtk_dic['ligand'] + vtk_dic['pocket'] + vtk_dic['mix'] + vtk_dic['empty'])) + + def load_to_vtk_opt(self, vtk_dic_opt): + if 'points' not in vtk_dic_opt: + vtk_dic_opt['points'] = {} + vtk_dic_opt['num_points'] = 0 + vtk_dic_opt['cells'] = [] + vtk_dic_opt['num_cells'] = 0 + vtk_dic_opt['cell_color_scale'] = [] + + vtk_dic_opt['Ligand'] = 0 + vtk_dic_opt['Protein'] = 0 + vtk_dic_opt['Mixture'] = 0 + vtk_dic_opt['Empty'] = 0 + + # current_cubic_points = [self.dot1, self.dot2, self.dot3, self.dot4, self.dot5, self.dot6, self.dot7, self.dot8] + current_cubic_cells = [[self.dot1, self.dot2, self.dot4, self.dot3], + [self.dot5, self.dot6, self.dot8, self.dot7], + [self.dot1, self.dot5, self.dot7, self.dot3], + [self.dot3, self.dot7, self.dot8, self.dot4], + [self.dot4, self.dot8, self.dot6, self.dot2], + [self.dot1, self.dot5, self.dot6, self.dot2]] + + for cell in current_cubic_cells: + cell_order = [] + for point in cell: + point = Point(point) + if point.string in vtk_dic_opt['points']: + point_order = vtk_dic_opt['points'][point.string] + cell_order.append(point_order) + else: + point_order = len(vtk_dic_opt['points']) + vtk_dic_opt['points'][point.string] = point_order + cell_order.append(point_order) + vtk_dic_opt['cells'].append(cell_order) + vtk_dic_opt['cell_color_scale'].append(self.color_scale) + + vtk_dic_opt['num_cells'] += 6 + if self.color_scale == 11: + vtk_dic_opt['Protein'] += 1 + elif self.color_scale == 12: + vtk_dic_opt['Ligand'] += 1 + elif self.color_scale == 13: + vtk_dic_opt['Mixture'] += 1 + elif self.color_scale == 14: + vtk_dic_opt['Empty'] += 1 + else: + print('Cubic {} have no color scale settled.'.format(self.id)) + + if len(vtk_dic_opt['cells']) != vtk_dic_opt['num_cells']: + print('cell error') + return vtk_dic_opt + + def write_oct_line_opt(self, vtk_line_dic, color): + """ + :param pts_dic: {key, value}, where key is Point.string, value is order of the point. + :param lines_list: + :param colors_list: + :param color: + :return: + """ + if len(vtk_line_dic) == 0: + vtk_line_dic['pts_dic'] = {} + vtk_line_dic['lines_dic'] = {} + vtk_line_dic['colors_list'] = [] + pts_dic = vtk_line_dic['pts_dic'] + lines_dic = vtk_line_dic['lines_dic'] + colors_list = vtk_line_dic['colors_list'] + + pts_len = len(pts_dic) + pts_inc = 0 + + p0 = [self.min_x, (self.min_y + self.max_y) / 2, (self.min_z + self.max_z) / 2] + p1 = [self.max_x, (self.min_y + self.max_y) / 2, (self.min_z + self.max_z) / 2] + p2 = [(self.min_x + self.max_x) / 2, self.min_y, (self.min_z + self.max_z) / 2] + p3 = [(self.min_x + self.max_x) / 2, self.max_y, (self.min_z + self.max_z) / 2] + p4 = [(self.min_x + self.max_x) / 2, (self.min_y + self.max_y) / 2, self.min_z] + p5 = [(self.min_x + self.max_x) / 2, (self.min_y + self.max_y) / 2, self.max_z] + + p6 = [self.min_x, (self.min_y + self.max_y) / 2, self.min_z] + p7 = [self.min_x, (self.min_y + self.max_y) / 2, self.max_z] + p8 = [self.min_x, self.min_y, (self.min_z + self.max_z) / 2] + p9 = [self.min_x, self.max_y, (self.min_z + self.max_z) / 2] + + p10 = [(self.min_x + self.max_x) / 2, self.min_y, self.min_z] + p11 = [(self.min_x + self.max_x) / 2, self.min_y, self.max_z] + p12 = [(self.min_x + self.max_x) / 2, self.max_y, self.min_z] + p13 = [(self.min_x + self.max_x) / 2, self.max_y, self.max_z] + + p14 = [self.max_x, (self.min_y + self.max_y) / 2, self.min_z] + p15 = [self.max_x, (self.min_y + self.max_y) / 2, self.max_z] + p16 = [self.max_x, self.min_y, (self.min_z + self.max_z) / 2] + p17 = [self.max_x, self.max_y, (self.min_z + self.max_z) / 2] + + line_list = [[p0, p1], [p2,p3], [p4, p5], [p6, p7], [p8, p9], [p14, p15], [p16, p17], + [p8, p16], [p10, p11], [p9, p17], [p12, p13], [p11, p13], [p7, p15], + [p6, p14], [p10, p12]] + + for line in line_list: + start_point = Point(line[0]) + end_point = Point(line[1]) + if start_point.string in pts_dic and end_point.string in pts_dic: + start_point_order = pts_dic[start_point.string] + end_point_order = pts_dic[end_point.string] + elif start_point.string not in pts_dic and end_point.string in pts_dic: + start_point_order = pts_len + pts_inc + pts_dic[start_point.string] = start_point_order + pts_inc += 1 + end_point_order = pts_dic[end_point.string] + elif start_point.string in pts_dic and end_point.string not in pts_dic: + end_point_order = pts_len + pts_inc + pts_dic[end_point.string] = end_point_order + pts_inc += 1 + start_point_order = pts_dic[start_point.string] + else: + start_point_order = pts_len + pts_inc + pts_dic[start_point.string] = start_point_order + pts_inc += 1 + end_point_order = pts_len + pts_inc + pts_dic[end_point.string] = end_point_order + pts_inc += 1 + line_string = '{}_{}'.format(start_point_order, end_point_order) + if line_string not in lines_dic: + lines_dic[line_string] = len(lines_dic) + colors_list.append(color) + + vtk_line_dic['pts_dic'] = pts_dic + vtk_line_dic['lines_dic'], vtk_line_dic['colors_list'] = lines_dic, colors_list + + return vtk_line_dic + + def write_edge_line_opt(self, vtk_line_dic, color): + if 'pts_dic' not in vtk_line_dic: + vtk_line_dic['pts_dic'] = {} + vtk_line_dic['lines_dic'] = {} + vtk_line_dic['colors_list'] = [] + pts_dic = vtk_line_dic['pts_dic'] + lines_dic, colors_list = vtk_line_dic['lines_dic'], vtk_line_dic['colors_list'] + pts_len = len(pts_dic) + pts_inc = 0 + + p0 = self.dot1 #= [self.min_x, self.min_y, self.min_z] + p1 = self.dot2 #= [self.min_x, self.min_y, self.max_z] + p2 = self.dot3 #= [self.min_x, self.max_y, self.min_z] + p3 = self.dot4 #= [self.min_x, self.max_y, self.max_z] + p4 = self.dot5 #= [self.max_x, self.min_y, self.min_z] + p5 = self.dot6 #= [self.max_x, self.min_y, self.max_z] + p6 = self.dot7 #= [self.max_x, self.max_y, self.min_z] + p7 = self.dot8 #= [self.max_x, self.max_y, self.max_z] + + line_list = [[p0,p1], [p1,p3], [p2,p3], [p0,p2], [p4,p5], [p5,p7], [p6,p7], [p4,p6], + [p0,p4], [p2,p6], [p3,p7], [p1,p5]] + + for line in line_list: + start_point = Point(line[0]) + end_point = Point(line[1]) + if start_point.string in pts_dic and end_point.string in pts_dic: + start_point_order = pts_dic[start_point.string] + end_point_order = pts_dic[end_point.string] + elif start_point.string not in pts_dic and end_point.string in pts_dic: + start_point_order = pts_len + pts_inc + pts_dic[start_point.string] = start_point_order + pts_inc += 1 + end_point_order = pts_dic[end_point.string] + elif start_point.string in pts_dic and end_point.string not in pts_dic: + end_point_order = pts_len + pts_inc + pts_dic[end_point.string] = end_point_order + pts_inc += 1 + start_point_order = pts_dic[start_point.string] + else: + start_point_order = pts_len + pts_inc + pts_dic[start_point.string] = start_point_order + pts_inc += 1 + end_point_order = pts_len + pts_inc + pts_dic[end_point.string] = end_point_order + pts_inc += 1 + line_string = '{}_{}'.format(start_point_order, end_point_order) + if line_string not in lines_dic: + lines_dic[line_string] = len(lines_dic) + colors_list.append(color) + vtk_line_dic['pts_dic'] = pts_dic + vtk_line_dic['lines_dic'], vtk_line_dic['colors_list'] = lines_dic, colors_list + return vtk_line_dic + +class Adaptive_Cubics(object): + def __init__(self, cubic_dic, complex_id, focus_center = [0,0,0], center = [0,0,0], + adaptive_line = True, differ = False, + count = False): + self.cubic_dic = cubic_dic + self.complex_id = complex_id + self.differ = differ + self.adaptive_line = adaptive_line + self.focus_center = focus_center + self.center = center + self.count = count + + def get_adaptive_cubic(self, core_area, side_area1, side_area2, side_area3, + target_range, line_depth, total_depth, save_folder, light = False): + """ + :param core_area: the range core area that need the finest resolution + :param side_area1: the range side area that need less resolution + :param side_area2: more range side area that need much less resolution + :return: + """ + self.vtk_dic = {} + self.fine_cubic_id_set = set() + self.parent_cubic_id_set = set() + + self.fine_cubic_count = {} + # self.non_fine_cubic_id_set = set() + + for value in ['Protein', 'Ligand', 'Mixture', 'Empty']: + self.fine_cubic_count[value] = 0 + + side1_layer_dic = self.construct_adaptive_dic_V1(self.cubic_dic, core_area, + target_range, self.focus_center, light) + side2_layer_dic = self.construct_adaptive_dic_V1(side1_layer_dic, side_area1, + target_range, self.focus_center, light) + side3_layer_dic = self.construct_adaptive_dic_V1(side2_layer_dic, side_area2, + target_range, self.focus_center, light) + side4_layer_dic = self.construct_adaptive_dic_V1(side3_layer_dic, side_area3, + target_range, self.focus_center, light) + side5_layer_dic = self.construct_adaptive_dic_V1(side4_layer_dic, target_range, + target_range, self.focus_center, light) + + print('Fine Cubic Dic Length: {}'.format(len(self.fine_cubic_id_set))) + # add for adaptive count: + # count_cubic_id_set = set() + # parent_cubic_id_set = set() + + if self.adaptive_line: + segline_cubic_id_set = set() + vtk_line_dic = {} + for cubic_id in self.fine_cubic_id_set: + # if not light: + # segline_cubic_id_set.add(cubic.id) # why need this? + # pass + # self.get_complement_adaptive_id(cubic_id, count_cubic_id_set, target_range) + parent_cubic_id = get_parent_cubic_id(cubic_id, target_range) + if parent_cubic_id in self.fine_cubic_id_set: + print('Parent and Fine conflict: {}'.format(parent_cubic_id)) + + while int(parent_cubic_id.split('_')[0]) >= 0: + # self.get_complement_adaptive_id(parent_cubic_id, parent_cubic_id_set, target_range) + # parent_cubic_set.add(parent_cubic) + # self.get_complement_adaptive(parent_cubic, count_cubic_set, target_range) + if int(parent_cubic_id.split('_')[0]) < line_depth: + segline_cubic_id_set.add(parent_cubic_id) + parent_cubic_id = get_parent_cubic_id(parent_cubic_id, target_range) + + if not light: + # for parent_cubic_id in parent_cubic_id_set: + # if int(parent_cubic_id.split('_')[0]) < line_depth: + # segline_cubic_id_set.add(parent_cubic_id) + print('parent_cubic_id_set length: {}'.format(len(self.parent_cubic_id_set))) + print('segline_cubic_id_set: {}'.format(len(segline_cubic_id_set))) + for cubic_id in segline_cubic_id_set: + depth = int(cubic_id.split('_')[0]) + color = get_color_for_dep(depth + 1, self.differ) + cubic = cubic_id_to_cubic(cubic_id, target_range) + vtk_line_dic = cubic.write_oct_line_opt(vtk_line_dic, color) + + max_cubic = Cubic( + [[-target_range, -target_range, -target_range], [target_range, target_range, target_range]], + '0_0_0_0') + color = get_color_for_dep(0, self.differ) + vtk_line_dic = max_cubic.write_edge_line_opt(vtk_line_dic, color) + + pts, lines, colors = fill_vtk_lines(vtk_line_dic, self.center) + write_seg_vtk(pts, lines, colors, self.complex_id, target_range, line_depth, save_folder) + + # + # for cubic_id in self.parent_cubic_id_set: + # if cubic_id not in segline_cubic_id_set: + # print('parent cubic {} not in segline_cubic_id_set'.format(cubic_id)) + + print('Final Node: {}'.format(len(self.fine_cubic_id_set))) + print('Mid Node: {}'.format(len(segline_cubic_id_set))) + + if not light: + file_name = '{}/{}/{}_Octree_Adaptive_{}_{}'.format(save_folder, self.complex_id, self.complex_id, target_range, total_depth) + vtk_writer(vtk_dic=self.vtk_dic, file_name=file_name, shift=self.center) + + def get_adaptive_count(self,core_area, side_area1, side_area2, side_area3, + target_range, line_depth, total_depth, save_folder, light = False): + depth_cubic_dic = {} + for depth in range(0, total_depth + 1): + depth_cubic_dic[depth] = set() + for cubic_id, value in self.cubic_dic: + [depth, x, y, z] = [int(item) for item in cubic_id.split('_')] + while depth >= 0: + depth_cubic_dic[depth].add(cubic_id) + depth = depth - 1 + x >> 1 + y >> 1 + z >> 1 + cubic_id = '{}_{}_{}_{}'.format(depth, x, y, z) + + + + + + + + def get_complement_adaptive_id(self, cubic_id, count_cubic_id_set, target_range): + [depth, x, y, z] = [int(item) for item in cubic_id.split('_')] + if depth <= 0: + count_cubic_id_set.add(cubic_id) + return + x = x >> 1 << 1 + y = y >> 1 << 1 + z = z >> 1 << 1 + + for i in range(0, 2): + for j in range(0, 2): + for k in range(0, 2): + string = '{}_{}_{}_{}'.format(depth, x + i, y + j, z + k) + + # base = target_range * 2 / 2 ** depth + # coord_x = x * base - target_range + i + # coord_y = y * base - target_range + j + # coord_z = z * base - target_range + k + # min_coord = [coord_x, coord_y, coord_z] + # max_coord = [coord_x + base, coord_y + base, coord_z + base] + # comp_cubic = Cubic([min_coord, max_coord], string) + # count_cubic_set.add(comp_cubic) + count_cubic_id_set.add(string) + + def construct_adaptive_dic_V1(self, cubic_dic, area, target_range, + focus_center=[0, 0, 0], light=False): + ''' + :param cubic_dic: dictionary, key is cubic.id, value is Protein, Mixture, Ligand. + :param area: within this area, use current depth in cubic_dic, otherwise need to compute upper layer + :param vtk_dic: dictionary that used to plot cubic. + :param seg_dic: dictionary that used to plot segment line of adaptive cibics. + :return: + ''' + + side_dic = {} + for key, value in cubic_dic.items(): + [depth, x, y, z] = [int(item) for item in key.split('_')] + + # if cubic's parent in the core range, the all child of its parent should be in fine_cubic_set. + p_x = x >> 1 + p_y = y >> 1 + p_z = z >> 1 + p_depth = depth - 1 + + p_base = target_range * 2 / 2 ** p_depth + p_coord_x = p_x * p_base - target_range + p_coord_y = p_y * p_base - target_range + p_coord_z = p_z * p_base - target_range + [focus_x, focus_y, focus_z] = focus_center + + # test if this cubic is in core area, if so, dump into adaptive cubic_dic + core_flag_z = min(abs(p_coord_z - focus_z), abs(p_coord_z + p_base - focus_z)) <= area + core_flag_y = min(abs(p_coord_y - focus_y), abs(p_coord_y + p_base - focus_y)) <= area + core_flag_x = min(abs(p_coord_x - focus_x), abs(p_coord_x + p_base - focus_x)) <= area + core_flag = core_flag_x and core_flag_y and core_flag_z + + p_id = '{}_{}_{}_{}'.format(p_depth, p_x, p_y, p_z) + + if not core_flag: + if p_id not in side_dic: + side_dic[p_id] = value + elif value != side_dic[p_id]: + side_dic[p_id] = 'Mixture' + else: + pass + + else: + self.parent_cubic_id_set.add(p_id) + # self.non_fine_cubic_id_set.add(p_id) + b_x = p_x << 1 + b_y = p_y << 1 + b_z = p_z << 1 + for i in range(0, 2): + for j in range(0, 2): + for k in range(0, 2): + n_x = b_x + i + n_y = b_y + j + n_z = b_z + k + n_id = '{}_{}_{}_{}'.format(depth, n_x, n_y, n_z) + + if n_id in self.parent_cubic_id_set: + continue + if n_id in self.fine_cubic_id_set: + continue + self.fine_cubic_id_set.add(n_id) + + if n_id in cubic_dic: + base = target_range * 2 / 2 ** depth + n_coord_x = n_x * base - target_range + n_coord_y = n_y * base - target_range + n_coord_z = n_z * base - target_range + min_coord = [n_coord_x, n_coord_y, n_coord_z] + max_coord = [n_coord_x + base, n_coord_y + base, n_coord_z + base] + n_cubic = Cubic([min_coord, max_coord], n_id) + n_value = cubic_dic[n_cubic.id] + n_cubic.set_color_new(value) + if not light: + self.vtk_dic = n_cubic.load_to_vtk_opt(self.vtk_dic) + self.fine_cubic_count[n_value] += 1 + else: + self.fine_cubic_count['Empty'] += 1 + #cubic.set_color_new('Empty') + pass + return side_dic + +def cubic_id_to_cubic(id, target_range): + ''' + based on id, target range, the cubic is identicle. + :param id: + :param target_range: + :return: + ''' + [depth, x, y, z] = [int(item) for item in id.split('_')] + # (depth, x, y, z) = id + base = target_range * 2 / 2 ** depth + coord_x = x * base - target_range + coord_y = y * base - target_range + coord_z = z * base - target_range + min_coord = [coord_x, coord_y, coord_z] + max_coord = [coord_x + base, coord_y + base, coord_z + base] + return Cubic([min_coord, max_coord], id) + +def point_string_to_coord(string): + coord = string.split('_') + return [float(item) for item in coord] + +def write_seg_vtk(pts, lines, colors, complex_id, target_range, line_depth, save_folder): + linesPolyData = vtk.vtkPolyData() + linesPolyData.SetPoints(pts) + linesPolyData.SetLines(lines) + linesPolyData.GetCellData().SetScalars(colors) + + # clean = vtk.vtkCleanPolyData() + # clean.SetInputConnection(linesPolyData) + # clean.Update() + # pd = clean.GetOutput() + + # print('clean segline_Adaptive data.') + # then = time.time() + # # pd = linesPolyData.getPolyData() + # # clean = vtk.vtkCleanPolyData() + # # clean.SetInput(linesPolyData) + # # clean.Update() + # # pd = clean.GetOutput() + # + # # clean = vtk.vtkCleanPolyData() + # # clean.SetInputConnection(linesPolyData) + # # clean.Update() + # # pd = clean.GetOutput() + # + # now = time.time() + # diff = int(now - then) + # print('clean segline_Adaptive data done. spend {} seconds'.format(diff)) + + ShellVolmapper = vtk.vtkDataSetMapper() + ShellVolmapper.SetScalarModeToUseCellData() + ShellVolmapper.UseLookupTableScalarRangeOn() + # ShellVolmapper.SetLookupTable(lut) + ShellVolmapper.SetInputData(linesPolyData) + ShellVolactor = vtk.vtkActor() + ShellVolactor.SetMapper(ShellVolmapper) + + # modelwriter = vtk.vtkUnstructuredGridWriter() + modelwriter = vtk.vtkDataSetWriter() + modelwriter.SetFileName('{}/{}/{}_segline_Adaptive_{}_{}.vtk'.format(save_folder, complex_id, complex_id, + target_range, line_depth)) + # modelwriter.SetInputData(grid) + modelwriter.SetInputData(ShellVolmapper.GetInput()) + modelwriter.Write() + print('Save {}/{}/{}_segline_Adaptive_{}_{}.vtk'.format(save_folder, complex_id, complex_id, target_range, line_depth)) + +def get_parent_cubic(cubic, target_range): + id = cubic.id + [depth, x, y, z] = [int(item) for item in id.split('_')] + x = x >> 1 + y = y >> 1 + z = z >> 1 + depth -= 1 + base = target_range * 2 / 2 ** depth + coord_x = x * base - target_range + coord_y = y * base - target_range + coord_z = z * base - target_range + min_coord = [coord_x, coord_y, coord_z] + max_coord = [coord_x + base, coord_y + base, coord_z + base] + string = '{}_{}_{}_{}'.format(depth, x, y, z) + parent_cubic = Cubic([min_coord, max_coord], string) + return parent_cubic + +def get_parent_cubic_id(id, target_range): + [depth, x, y, z] = [int(item) for item in id.split('_')] + x = x >> 1 + y = y >> 1 + z = z >> 1 + depth -= 1 + # base = target_range * 2 / 2 ** depth + # coord_x = x * base - target_range + # coord_y = y * base - target_range + # coord_z = z * base - target_range + # min_coord = [coord_x, coord_y, coord_z] + # max_coord = [coord_x + base, coord_y + base, coord_z + base] + parent_id = '{}_{}_{}_{}'.format(depth, x, y, z) + # parent_cubic = Cubic([min_coord, max_coord], string) + return parent_id + +def get_color_for_dep(depth, differ = True): + # need to be different with the volume scaler (0,1,2,3). + if differ == False: + color = 0 + return color + if depth == 0: + color = 0 #[0, 0, 0] # white + elif depth == 1: + color = 1 # [255, 0, 0] # red + elif depth == 2: + color = 2 # [255, 128, 0] + elif depth == 3: + color = 3 #[255,255,0] + elif depth == 4: + color = 4 # [0, 255, 0] + elif depth == 5: + color = 5 #[0, 255, 255] + elif depth == 6: + color = 6 #[0, 0, 255] + elif depth == 7: + color = 7 # [255, 0 , 255] + elif depth == 8 : + color = 8 + else: + color = 9 + return color + +# A better way is use class in vtk package, this solve the color issue to distinguish ligand, pocket. +def vtk_writer(vtk_dic, file_name, shift = [0,0,0]): + points = vtk.vtkPoints() + grid = vtk.vtkUnstructuredGrid() + Color = vtk.vtkFloatArray() + Scale = vtk.vtkFloatArray() + + sorted_dic = sorted(vtk_dic['points'].items(), key=lambda x: x[1]) + pts = vtk.vtkPoints() + for item in sorted_dic: + # pts.InsertNextPoint(point_string_to_coord(item[0])) + coord = point_string_to_coord(item[0]) + coord = [coord[0] + shift[0], coord[1] + shift[1], coord[2] + shift[2]] + pts.InsertNextPoint(coord) + grid.SetPoints(pts) + + # for vtk_point in vtk_dic['points']: + # Coord = vtk_point + # test = points.InsertNextPoint(*Coord) + # #Color.InsertTuple1(test, 0) + # grid.SetPoints(points) + + for i in range(len(vtk_dic['cells'])): + vtk_node = vtk_dic['cells'][i] + elem = vtk.vtkQuad() + elem.GetPointIds().SetId(0, vtk_node[0]) + elem.GetPointIds().SetId(1, vtk_node[1]) + elem.GetPointIds().SetId(2, vtk_node[2]) + elem.GetPointIds().SetId(3, vtk_node[3]) + Quad4cell = grid.InsertNextCell(elem.GetCellType(), elem.GetPointIds()) + color = vtk_dic['cell_color_scale'][i] + Color.InsertTuple1(Quad4cell, color) + + lut = vtk.vtkLookupTable() + lut.SetNumberOfTableValues(5) + lut.SetTableRange(0, 5) + lut.Build() + lut.SetTableValue(0, 0, 0, 0, 1) # Black + lut.SetTableValue(1, 1, 0, 0, 1) # Red + lut.SetTableValue(2, 0, 1, 0, 1) # Green + lut.SetTableValue(3, 0, 0, 1, 1) # Blue + lut.SetTableValue(4, 1, 1, 1, 1) # White + + grid.GetCellData().SetScalars(Color) + ShellVolmapper = vtk.vtkDataSetMapper() + # ShellVolmapper = vtk.vtkUnstructuredGridVolumeRayCastMapper() + ShellVolmapper.SetScalarModeToUseCellData() + ShellVolmapper.UseLookupTableScalarRangeOn() + ShellVolmapper.SetLookupTable(lut) + ShellVolmapper.SetInputData(grid) + ShellVolactor = vtk.vtkActor() + ShellVolactor.SetMapper(ShellVolmapper) + + modelwriter = vtk.vtkUnstructuredGridWriter() + modelwriter.SetFileName('{}.vtk'.format(file_name)) + #modelwriter.SetInputData(grid) + modelwriter.SetInputData(ShellVolmapper.GetInput()) + modelwriter.Write() + print('Save {}.vtk'.format(file_name)) + + +def construct_adaptive_dic(cubic_dic, area, vtk_dic, fine_cubic_set, target_range, + focus_center = [0,0,0], light =False): + ''' + :param cubic_dic: dictionary, key is cubic.id, value is Protein, Mixture, Ligand. + :param area: within this area, use current depth in cubic_dic, otherwise need to compute upper layer + :param vtk_dic: dictionary that used to plot cubic. + :param seg_dic: dictionary that used to plot segment line of adaptive cibics. + :return: + ''' + side_dic = {} + for key, value in cubic_dic.items(): + [depth, x, y, z] = [int(item) for item in key.split('_')] + base = target_range * 2 / 2 ** depth + coord_x = x * base - target_range + coord_y = y * base - target_range + coord_z = z * base - target_range + min_coord = [coord_x, coord_y, coord_z] + max_coord = [coord_x + base, coord_y + base, coord_z + base] + [focus_x, focus_y, focus_z] = focus_center + + # test if this cubic is in core area, if so, dump into adaptive cubic_dic + core_flag_z = min(abs(coord_z - focus_z), abs(coord_z + base - focus_z)) <= area + core_flag_y = min(abs(coord_y - focus_y), abs(coord_y + base - focus_y)) <= area + core_flag_x = min(abs(coord_x - focus_x), abs(coord_x + base - focus_x)) <= area + core_flag = core_flag_x and core_flag_y and core_flag_z + + # find parent cubic or upper layer cubic. + # parent cubic for adaptive segment line, upper layer cubic for cubic plot. + x = x >> 1 + y = y >> 1 + z = z >> 1 + depth -= 1 + string = '{}_{}_{}_{}'.format(depth, x, y, z) + + if core_flag: + cubic = Cubic([min_coord, max_coord], key) + cubic.set_color_new(value) + if not light: + vtk_dic = cubic.load_to_vtk_opt(vtk_dic) + fine_cubic_set.add(cubic) + else: + if string not in side_dic: + side_dic[string] = value + elif value != side_dic[string]: + side_dic[string] = 'Mixture' + else: + pass + return side_dic + + + + +def centerize(pocket_xyz, ligand_xyz, by_ligand = True, protein_xyz = None): + if by_ligand: + [x, y, z] = ligand_xyz + center_x, center_y, center_z = x.mean(), y.mean(), z.mean() + x, y, z = x - center_x, y - center_y, z - center_z + ligand_xyz = [x, y, z] + [x, y, z] = pocket_xyz + x, y, z = x - center_x, y - center_y, z - center_z + pocket_xyz = [x, y, z] + ligand_center = [0, 0, 0] + center = [center_x, center_y, center_z] + + if protein_xyz is not None: + [x, y, z] = protein_xyz + x, y, z = x - center_x, y - center_y, z - center_z + protein_xyz = [x, y, z] + + else: + [x_ligand, y_ligand, z_ligand] = ligand_xyz + x_ligand_max, y_ligand_max, z_ligand_max = x_ligand.max(), y_ligand.max(), z_ligand.max() + x_ligand_min, y_ligand_min, z_ligand_min = x_ligand.min(), y_ligand.min(), z_ligand.min() + if protein_xyz is None: + [x_pocket, y_pocket, z_pocket] = pocket_xyz + x_protein_max, y_protein_max, z_protein_max = x_pocket.max(), y_pocket.max(), z_pocket.max() + x_protein_min, y_protein_min, z_protein_min = x_pocket.min(), y_pocket.min(), z_pocket.min() + else: + [x_protein, y_protein, z_protein] = protein_xyz + x_protein_max, y_protein_max, z_protein_max = x_protein.max(), y_protein.max(), z_protein.max() + x_protein_min, y_protein_min, z_protein_min = x_protein.min(), y_protein.min(), z_protein.min() + x_max, y_max, z_max = max(x_ligand_max, x_protein_max), max(y_ligand_max, y_protein_max), max(z_ligand_max, z_protein_max) + x_min, y_min, z_min = min(x_ligand_min, x_protein_min), min(y_ligand_min, y_protein_min), min(z_ligand_min, z_protein_min) + center = [(x_max+x_min)/2, (y_max+y_min)/2, (z_max+z_min)/2] + ligand_center = [x_ligand.mean() - center[0], y_ligand.mean() - center[1], z_ligand.mean()- center[2]] + ligand_xyz = [ligand_xyz[0] - center[0], ligand_xyz[1] - center[1], ligand_xyz[2] - center[2]] + pocket_xyz = [pocket_xyz[0] - center[0], pocket_xyz[1] - center[1], pocket_xyz[2] - center[2]] + if protein_xyz is not None: + [x, y, z] = protein_xyz + x, y, z = x - center[0], y - center[1], z - center[2] + protein_xyz = [x, y, z] + + return pocket_xyz, ligand_xyz, protein_xyz, center, ligand_center + +def filter_by_center_range(x, y, z, target_range, feature=None): # Here x, y, z is list + x = np.array(x) + y = np.array(y) + z = np.array(z) + + min_x, min_y, min_z = -target_range, -target_range, -target_range + max_x, max_y, max_z = target_range, target_range, target_range + + flags_x = np.logical_and(x > min_x, x < max_x) + flags_y = np.logical_and(y > min_y, y < max_y) + flags_z = np.logical_and(z > min_z, z < max_z) + flag1 = np.logical_and(flags_x, flags_y) + flag2 = np.logical_and(flag1, flags_z) + x = x[(flag2)] + y = y[(flag2)] + z = z[(flag2)] + + if feature is not None: + feature = np.array(feature) + feature = feature[(flag2)] + return x,y,z,feature + return x, y, z + +def get_vol_index(x, y, z, depth, target_range): + min_x, min_y, min_z = -target_range, -target_range, -target_range + x = x - min_x + y = y - min_y + z = z - min_z + resolution = target_range*2 / 2 ** depth + x = (x // resolution).astype(int) + y = (y // resolution).astype(int) + z = (z // resolution).astype(int) + # vol_set = set() + # for i in range(len(x)): + # vol_set.add((x[i], y[i], z[i])) + # print(vol_set) + # print(min(x), max(x)) + return x, y, z + +def get_comp_dic(cubic_dic, comp_layer_dic = {}): + for key, value in cubic_dic.items(): + [depth, x, y, z] = [int(item) for item in key.split('_')] + x = x >> 1 << 1 + y = y >> 1 << 1 + z = z >> 1 << 1 + + for i in range(0,2): + for j in range(0,2): + for k in range(0,2): + string = '{}_{}_{}_{}'.format(depth, x+i, y+j, z+k) + if string not in cubic_dic: + comp_layer_dic[string] = 'Empty' + else: + comp_layer_dic[string] = cubic_dic[string] + pass + + # for key, value in comp_layer_dic.items(): + # if value == 'Ligand': + # print(key) + return comp_layer_dic + + + +def get_upper_layer_dic(cubic_dic): + upper_layer_dic = {} + for key, value in cubic_dic.items(): + [depth, x, y, z] = [int(item) for item in key.split('_')] + x = x >> 1 + y = y >> 1 + z = z >> 1 + depth -= 1 + string = '{}_{}_{}_{}'.format(depth, x,y,z) + if string not in upper_layer_dic: + upper_layer_dic[string] = value + elif value != upper_layer_dic[string]: + upper_layer_dic[string] = 'Mixture' + else: + pass + return upper_layer_dic + + +def create_cubic_dic(protein_xyz, ligand_xyz, depth): + cubic_dic = {} + [x, y, z] = protein_xyz + print('Protein points number: {}'.format(len(x))) + for i in range(len(x)): + string = '{}_{}_{}_{}'.format(depth, x[i], y[i], z[i]) + if string not in cubic_dic: + cubic_dic[string] = 'Protein' + else: + pass + + [x, y, z] = ligand_xyz + print('Ligand points number: {}'.format(len(x))) + for i in range(len(x)): + string = '{}_{}_{}_{}'.format(depth, x[i], y[i], z[i]) + if string not in cubic_dic: + cubic_dic[string] = 'Ligand' + elif cubic_dic[string] == 'Protein': + cubic_dic[string] = 'Mixture' + else: + pass + return cubic_dic + +def create_cubic_dic_with_feature(protein_xyz, ligand_xyz, depth, pocket_feature, ligand_feature): + cubic_dic = {} + [x, y, z] = protein_xyz + for i in range(len(x)): + string = '{}_{}_{}_{}'.format(depth, x[i], y[i], z[i]) + feature = pocket_feature[i] + atom_type = feature[0] + atom_index = 'P' + str(feature[1]) + atom_detail_type = 'P' + str(feature[0]) + if string not in cubic_dic: + cubic_dic[string] = {} + cubic_dic[string]['PLM'] = 'Protein' + cubic_dic[string]['Atoms'] = [atom_type] + cubic_dic[string]['Atom_index'] = [atom_index] + cubic_dic[string]['Atom_Detail_Type'] = [atom_detail_type] + else: + cubic_dic[string]['Atom_Detail_Type'].append(atom_detail_type) + if atom_type not in cubic_dic[string]['Atoms']: + cubic_dic[string]['Atoms'].append(atom_type) + if atom_index not in cubic_dic[string]['Atom_index']: + cubic_dic[string]['Atom_index'].append(atom_index) + + [x, y, z] = ligand_xyz + for i in range(len(x)): + string = '{}_{}_{}_{}'.format(depth, x[i], y[i], z[i]) + feature = ligand_feature[i] + atom_type = feature[0] + atom_index = 'L' + str(feature[1]) + atom_detail_type = 'L' + str(feature[0]) + if string not in cubic_dic: + cubic_dic[string] = {} + cubic_dic[string]['PLM'] = 'Ligand' + cubic_dic[string]['Atoms'] = [atom_type] + cubic_dic[string]['Atom_index'] = [atom_index] + cubic_dic[string]['Atom_Detail_Type'] = [atom_detail_type] + elif cubic_dic[string]['PLM'] == 'Protein': + cubic_dic[string]['PLM'] = 'Mixture' + cubic_dic[string]['Atom_Detail_Type'].append(atom_detail_type) + else: + cubic_dic[string]['Atom_Detail_Type'].append(atom_detail_type) + if atom_type not in cubic_dic[string]['Atoms']: + cubic_dic[string]['Atoms'].append(atom_type) + if atom_index not in cubic_dic[string]['Atom_index']: + cubic_dic[string]['Atom_index'].append(atom_index) + return cubic_dic + +def line_string_to_idx(string): + contents = string.split('_') + contents = [int(item) for item in contents] + return contents + +def fill_vtk_lines(vtk_line_dic, shift = [0,0,0]): + pts_dic = vtk_line_dic['pts_dic'] + lines_dic, colors_list = vtk_line_dic['lines_dic'], vtk_line_dic['colors_list'] + + pts = vtk.vtkPoints() + sorted_dic = sorted(pts_dic.items(), key=lambda x: x[1]) + for item in sorted_dic: + coord = point_string_to_coord(item[0]) + coord = [coord[0] + shift[0], coord[1] + shift[1], coord[2] + shift[2]] + pts.InsertNextPoint(coord) + + + lines = vtk.vtkCellArray() + colors = vtk.vtkFloatArray() + sorted_line_dic = sorted(lines_dic.items(), key=lambda x: x[1]) + for item in sorted_line_dic: + line_idx = line_string_to_idx(item[0]) + line0 = vtk.vtkLine() + line0.GetPointIds().SetId(0, line_idx[0]) + line0.GetPointIds().SetId(1, line_idx[1]) + lines.InsertNextCell(line0) + + # for line in lines_list: + # points = line.split('_') + # points = [int(item) for item in points] + # line0 = vtk.vtkLine() + # line0.GetPointIds().SetId(0, points[0]) + # line0.GetPointIds().SetId(1, points[1]) + # lines.InsertNextCell(line0) + for color in colors_list: + colors.InsertNextTuple1(color) + return pts, lines, colors