def __init__(self, bnd_dihedral, bnd_rmsd, support_dihedral, support_rmsd):

self.bnd_dihedral = bnd_dihedral

self.bnd_rmsd = bnd_rmsd

self.support_dihedral = support_dihedral # in percent of the pool

def __init__(self, pdb_path, scores_file, n_sym = 1):

self.pdb_path = pdb_path

self.scores_file = scores_file

self.n_sym = n_sym

self.pdb_files = os.listdir(self.pdb_path)

self.pdb_files = [f for f in self.pdb_files if f.endswith(".pdb")]

self.pdb_files.sort()

self.n_models = len(self.pdb_files)

parsed_file_name = self.pdb_path+"_parsed.pkl"

if os.path.isfile(parsed_file_name):

self.CA, self.DIH = cPickle.load(open(parsed_file_name, "rb"))

else:

self.CA, self.DIH = self.parse_structures()

cPickle.dump((self.CA, self.DIH), open(parsed_file_name, "wb"))

self.n_res = self.DIH.shape[1]

if self.scores_file != "":

self.scores = []

scores_txt = open(self.scores_file, "rb").readlines()

for line in scores_txt[1:]:

line_arr = line.split()

score = float(line_arr[1])

if score > 0:

score = -1

self.scores.append(score)

assert len(self.scores) == len(self.pdb_files)

self.use_scores = True

else:

self.use_scores = False

def parse_structures(self):

phipsi_list = []

ca_all = []

warnings.filterwarnings("ignore")

for pdb_file in self.pdb_files:

print "Parsing", pdb_file

structure = PDB.PDBParser().get_structure(pdb_file, os.path.join(self.pdb_path,pdb_file))

model = structure[0]

phipsi = []

ca_list = []

for chain in model.child_list:

polypeptide = PDB.PPBuilder().build_peptides(chain)[0]

phipsi = phipsi + polypeptide.get_phi_psi_list()

for residue in chain:

atom = residue['CA']

ca_list.append(atom.coord)

phipsi_list.append(phipsi)

ca_all.append(ca_list)

n_res_total = len(phipsi_list[0])

DIH = np.zeros((self.n_models, n_res_total, 2))

CA = np.zeros((self.n_models, n_res_total, 3))

for i,(phipsi,model) in enumerate(zip(phipsi_list,ca_all)):

for j,(angle,ca_atom) in enumerate(zip(phipsi,model)):

DIH[i, j, :] = angle[:]

CA[i,j,:] = ca_atom

n_models = CA.shape[0]

working_CA = np.copy(CA)

sup=SVDSuperimposer()

ref_model = working_CA[0, :, :]

rms_total = 0

for i_model in range(1, n_models):

sup.set(ref_model, working_CA[i_model])

sup.run()

rms_total += sup.get_rms()**2

working_CA[i_model] = sup.get_transformed()

rms_best = float("inf")

epsilon = 0.001

while rms_best - rms_total > epsilon:

rms_best = rms_total

mean_model = np.mean(working_CA,0)

rms_total = 0

for i_model in range(n_models):

sup.set(mean_model, working_CA[i_model])

sup.run()

rms_total += sup.get_rms()**2

working_CA[i_model] = sup.get_transformed()

transformations = []

for start_model, result_model in zip(CA, working_CA):

sup.set(result_model, start_model)

sup.run()

transformations.append(sup.get_rotran())

n_models = CA.shape[0]

distances = np.zeros((n_models, n_models))

sup=SVDSuperimposer()

for i in range(n_models):

model1 = CA[i,:,:]

for j in range(i+1,n_models):

model2 = CA[j,:,:]

sup.set(model1, model2)

sup.run()

rms=sup.get_rms()

distances[i,j] = rms

distances[j,i] = rms

n_models = CA.shape[0]

distances = np.zeros((n_models, n_models))

centered_X_all = []

det_XX_all = []

for i in range(n_models):

model_orig = CA[i,:,:]

mean_v = np.mean(model_orig,0)

model_centered = model_orig - mean_v

det_XX = np.linalg.det(np.dot(np.transpose(model_centered), model_centered))

centered_X_all.append(model_centered)

det_XX_all.append(det_XX)

for i in range(n_models):

model1 = np.transpose(centered_X_all[i])

for j in range(i+1,n_models):

det_XY = np.linalg.det(np.dot(model1, centered_X_all[j]))

BC = 1-det_XY/math.sqrt(det_XX_all[i]*det_XX_all[j])

distances[i,j] = BC

distances[j,i] = BC

warnings.filterwarnings("ignore")

score,residues,support = ens

if len(support)>10:

support=support[:10]

aln_CA = work.CA.take(support, 0).take(residues, 1)

transformations,rms = align_models(aln_CA)

ens_structure = PDB.Structure.Structure(" ")

ens_name = prefix+"_{0:03d}_{1:03d}_{2:03d}_{3:02.2f}".format(min(residues), max(residues), len(residues), rms)

ala_atom_names = {"N", "CA", "C", "O", "CB"}

bfac_base = 10

for file_ind in support:

file_name = work.pdb_files[file_ind]

structure=PDB.PDBParser().get_structure(file_name, os.path.join(work.pdb_path,file_name))

model = structure.child_list[0]

new_model = PDB.Model.Model(len(ens_structure)+1)

offset = 0

for chain in model.child_list:

offset_next = len(chain.child_list)

for res_ind,residue in reversed(list(enumerate(chain))):

if res_ind+offset not in residues:

id = residue.id

chain.detach_child(id)

else:

bfac_mult = [i for i in range(len(levels)) if res_ind+offset in levels[i]][0]

for atom in reversed(residue.child_list):

if atom.get_name() not in ala_atom_names:

residue.detach_child(atom.id)

else:

atom.set_bfactor(bfac_base*bfac_mult)

offset += offset_next

new_model.add(chain)

ens_structure.add(new_model)

for aln_structure, rotran in zip(ens_structure.child_list, transformations):

atom_list = aln_structure.get_atoms()

for atom in atom_list:

atom.transform(rotran[0].astype('f'), rotran[1].astype('f'))

io=PDB.PDBIO()

io.set_structure(ens_structure)

combined = []

top_score,top_ind,top_support = top_frag

cand_score,cand_ind,cand_support = candidate

min_support = options.support_rmsd

comb_ind = sorted(list(set(top_ind) | set(cand_ind)))

comb_support = sorted(list(set(top_support) & set(cand_support)))

n_comb_support=len(comb_support)

if n_comb_support < min_support:

return []

if work.use_scores:

comb_scores = [work.scores[i] for i in comb_support]

aln_models = work.CA.take(comb_support, 0).take(comb_ind, 1)

calculator = RMSDCalculator.RMSDCalculator("QCP_OMP_CALCULATOR", aln_models)

dist = squareform(calculator.pairwiseRMSDMatrix())

mds = manifold.MDS(n_components=2, dissimilarity="precomputed", n_jobs=1, n_init = 5)

pos = mds.fit(dist).embedding_

try:

ms = MeanShift(bandwidth=options.bnd_rmsd, cluster_all=False, bin_seeding=True, min_bin_freq = min_support)

ms.fit(pos)

except:

try:

ms = MeanShift(bandwidth=options.bnd_rmsd, cluster_all=False, bin_seeding=False)

ms.fit(pos)

except:

return []

labels = ms.labels_

labels_unique = np.unique(labels)

for label in labels_unique:

if label == -1:

continue

class_members = [index[0] for index in np.argwhere(labels == label)]

class_support = [comb_support[i] for i in class_members]

n_class_support = len(class_support)

if n_class_support < min_support:

continue

class_dist = dist.take(class_members,0).take(class_members,1)

mean_dist = np.mean(squareform(class_dist))

if work.use_scores:

class_scores = [comb_scores[i] for i in class_members]

class_score = sum(class_scores)/(1+mean_dist)

else:

class_score = (-n_class_support)/(1+mean_dist)

heapq.heappush(combined, (class_score, comb_ind, class_support))

if combined:

return heapq.heappop(combined)

n_peptide = len(peptides[0][1])*2

peptides_out=[]

min_support_seeds = int(work.n_models*options.support_dihedral)

for res_ind in range(work.n_res-n_peptide+1):

indices = [range(res_ind,res_ind+int(n_peptide/2)), range(res_ind+int(n_peptide/2),res_ind+n_peptide)]

peptide_collection = []

for ind in indices:

singletons_ind = [residue for residue in peptides if residue[1] == ind]

peptide_collection.append(singletons_ind)

peptide_candidates = itertools.product(*peptide_collection)

for peptide in peptide_candidates:

singles_support = [set(residue[2]) for residue in peptide]

peptide_support = list(set.intersection(*singles_support))

n_peptide_support = len(peptide_support)

if n_peptide_support >= min_support_seeds:

indices_flat = [item for sublist in indices for item in sublist]

heapq.heappush(peptides_out, (-n_peptide_support, indices_flat, sorted(peptide_support)))

options, work, all_peptides, seed = args

n_level = len(all_peptides) - 2

seed_score, seed_ind, seed_support = seed

growing = (seed_score, seed_ind, seed_support)

is_growing = True

cur_peptides = all_peptides[n_level]

bfac_level = [set(seed_ind)]

while is_growing:

combined = []

for candidate in cur_peptides:

if set(candidate[1]) & set(growing[1]):

continue

cand_ev = evaluate_candidate(options, work, growing, candidate)

if cand_ev != []:

heapq.heappush(combined, cand_ev)

if combined != []:

prev_ind = set(growing[1])

growing = heapq.heappop(combined)

next_ind = set(growing[1])

bfac_level.append(set.difference(next_ind,prev_ind))

else:

n_level -= 1

if n_level == 0:

is_growing = False

else:

cur_peptides = all_peptides[n_level]

print "Finished seed", seed

print "Result:", growing

print 'Dihedral bandwidth:', options.bnd_dihedral

print 'RMSD bandwidth:', options.bnd_rmsd

singletons = []

min_bin_freq = options.support_rmsd

for res_ind in range(work.n_res):

distr_res = work.DIH[:,res_ind,:]

if distr_res[0,0] == None or distr_res[0,1] == None:

continue

try:

ms = MeanShift(bandwidth=options.bnd_dihedral, cluster_all=False, bin_seeding=True, min_bin_freq = min_bin_freq)

ms.fit(distr_res)

except:

try:

ms = MeanShift(bandwidth=options.bnd_dihedral, cluster_all=False, bin_seeding=False)

ms.fit(distr_res)

except:

continue

labels = ms.labels_

labels_unique = np.unique(labels)

for label in labels_unique:

if label == -1:

continue

class_members = [index[0] for index in np.argwhere(labels == label)]

singletons.append((-len(class_members), [res_ind], class_members))

print len(singletons), "singletons"

all_peptides = [singletons]

while len(all_peptides[-1]) > 0 and len(all_peptides) < 5:

doubled = double_peptides(options, work, all_peptides[-1])

print len(doubled), 2**len(all_peptides), "-peptides"

all_peptides.append(doubled)

if len(all_peptides[-1]) == 0:

all_peptides.pop()

seed_candidates = sorted(all_peptides[-1])

seeds = []

sym_offset = int(work.n_res/work.n_sym)

for seed_candidate in seed_candidates:

independent = True

cand_ind = set(seed_candidate[1])

cand_support = set(seed_candidate[2])

for seed in seeds:

seed_ind = set(seed[1])

seed_support = set(seed[2])

sim_ind = len(cand_ind & seed_ind)/len(cand_ind | seed_ind)

sim_support = len(cand_support & seed_support)/len(cand_support | seed_support)

if min(seed_ind) < min(cand_ind):

first_ind = set([i + sym_offset for i in seed_ind])

second_ind = cand_ind

else:

first_ind = set([i + sym_offset for i in cand_ind])

second_ind = seed_ind

sym_sim_ind = len(first_ind & second_ind)/len(first_ind | second_ind)

if (sim_ind > 0.5 or sym_sim_ind > 0.5) and sim_support > 0.5:

independent = False

break

if independent:

seeds.append(seed_candidate)

print len(seeds), "seeds from", 2**(len(all_peptides)-1), "-peptides"

n_proc = min(max(1, multiprocessing.cpu_count()), len(seeds))

pool = multiprocessing.Pool(n_proc)

work_seeds = itertools.product([options], [work], [all_peptides], seeds)

results = pool.map(grow_seed, work_seeds, 1)

pool.close()

pool.join()

# TODO: find pairs with most coverage