from __future__ import absolute_import, division, print_function from mmtbx.utils import get_atom_selections import mmtbx.utils from cctbx.array_family import flex from scitbx import lbfgs from libtbx.str_utils import format_value, make_sub_header from libtbx.utils import Sorry, null_out from libtbx import adopt_init_args from six.moves import range class manager(object): def __init__(self, fmodels, model, max_number_of_iterations = 25, number_of_macro_cycles = 3, occupancy_max = None, occupancy_min = None, log = None, exclude_hd = False): self.show(fmodels=fmodels, log= log, message="occupancy refinement: start") fmodels.update_xray_structure(xray_structure = model.get_xray_structure(), update_f_calc = True) selections = model.refinement_flags.s_occupancies # exclude H or D from refinement if requested if(exclude_hd): hd_sel = model.get_hd_selection() tmp_sel = [] for sel in selections: tmp_sel_ = [] for sel_ in sel: tmp_sel__ = flex.size_t() for sel__ in sel_: if(not hd_sel[sel__]): tmp_sel__.append(sel__) if(tmp_sel__.size()>0): tmp_sel_.append(tmp_sel__) if(len(tmp_sel_)>0): tmp_sel.append(tmp_sel_) selections = tmp_sel # if(len(selections)>0): i_selection = flex.size_t() for s in selections: for ss in s: i_selection.extend(ss) fmodels.fmodel_xray().xray_structure.scatterers().flags_set_grads( state=False) fmodels.fmodel_xray().xray_structure.scatterers().flags_set_grad_occupancy( iselection = i_selection) fmodels.fmodel_xray().xray_structure.adjust_occupancy( occ_max = occupancy_max, occ_min = occupancy_min, selection = i_selection) xray_structure_dc = fmodels.fmodel_xray().xray_structure.\ deep_copy_scatterers() par_initial = flex.double() occupancies = xray_structure_dc.scatterers().extract_occupancies() constrained_groups_selections = [] group_counter = 0 for sel in selections: ss = [] for sel_ in sel: ss.append(group_counter) group_counter += 1 val = flex.mean(occupancies.select(sel_)) par_initial.append(val) constrained_groups_selections.append(ss) minimized = None for macro_cycle in range(number_of_macro_cycles): if(minimized is not None): par_initial = minimized.par_min minimized = minimizer( fmodels = fmodels, selections = selections, constrained_groups_selections = constrained_groups_selections, par_initial = par_initial, max_number_of_iterations = max_number_of_iterations) if(minimized is not None): par_initial = minimized.par_min set_refinable_parameters( xray_structure = fmodels.fmodel_xray().xray_structure, parameters = par_initial, selections = selections, enforce_positivity = (occupancy_min>=0)) fmodels.fmodel_xray().xray_structure.adjust_occupancy( occ_max = occupancy_max, occ_min = occupancy_min, selection = i_selection) xray_structure_final = fmodels.fmodel_xray().xray_structure model.set_xray_structure(xray_structure_final) fmodels.update_xray_structure(xray_structure = xray_structure_final, update_f_calc = True) refined_occ = xray_structure_final.scatterers().extract_occupancies().\ select(i_selection) assert flex.min(refined_occ) >= occupancy_min assert flex.max(refined_occ) <= occupancy_max self.show(fmodels= fmodels, log = log, message="occupancy refinement: end") def show(self, fmodels, message, log): if(log is not None): print("|-"+message+"-"*(79-len("|-"+message+"|"))+"|", file=log) fm_x, fm_n = fmodels.fmodel_xray(), fmodels.fmodel_neutron() if(fm_n is not None): print("|"+" "*36+"X-ray"+" "*36+"|", file=log) self.show_helper(fmodel = fm_x, log = log) if(fm_n is not None): print("|"+" "*35+"neutron"+" "*35+"|", file=log) self.show_helper(fmodel = fm_n, log = log) occupancies = fm_x.xray_structure.scatterers().extract_occupancies() occ_max = format_value("%4.2f", flex.max(occupancies)) occ_min = format_value("%4.2f", flex.min(occupancies)) number_small = format_value("%8d", (occupancies < 0.1).count(True)) print("| occupancies: max = %s min = %s number of occupancies < 0.1: %s |"%( occ_max, occ_min, number_small), file=log) print("|"+"-"*77+"|", file=log) def show_helper(self, fmodel, log): r_work = format_value("%6.4f", fmodel.r_work()) r_free = format_value("%6.4f", fmodel.r_free()) target = format_value("%-13.3f", fmodel.target_w()) target_name = format_value("%s", fmodel.target_name) p1 = "| r_work = %s r_free = %s" % (r_work, r_free) p2 = "target_work(%s) = %s |" % (target_name, target) print(p1+" "*(79-len(p1+p2))+p2, file=log) class minimizer(object): def __init__(self, fmodels, constrained_groups_selections, selections, par_initial, max_number_of_iterations): adopt_init_args(self, locals()) self.fmodels.create_target_functors() self.fmodels.prepare_target_functors_for_minimization() from phenix.refinement import weight_xray_chem self.weights = weight_xray_chem.weights(wx = 1, wx_scale = 1, angle_x = None, wn = 1, wn_scale = 1, angle_n = None, w = 0, wxn = 1) # XXX self.par_min = self.par_initial.deep_copy() self.x = self.pack(self.par_min) self.n = self.x.size() self.minimizer = lbfgs.run( target_evaluator = self, termination_params = lbfgs.termination_parameters( max_iterations = max_number_of_iterations), exception_handling_params = lbfgs.exception_handling_parameters( ignore_line_search_failed_step_at_lower_bound = True, ignore_line_search_failed_step_at_upper_bound = True)) self.compute_functional_and_gradients() del self.x def pack(self, par): return pack_unpack(x = par, table = self.constrained_groups_selections) def unpack_x(self): self.par_min = pack_unpack(x = self.x, table = self.constrained_groups_selections) def compute_functional_and_gradients(self): self.unpack_x() set_refinable_parameters( xray_structure = self.fmodels.fmodel_xray().xray_structure, parameters = self.par_min, selections = self.selections) self.fmodels.update_xray_structure(update_f_calc = True) fmodels_target_and_gradients = self.fmodels.target_and_gradients( weights = self.weights, compute_gradients = True, occupancy = True) self.f = fmodels_target_and_gradients.target() g = fmodels_target_and_gradients.gradients() # do group grads first then pack for constraints grads_wrt_par = flex.double() for sel in self.selections: for sel_ in sel: grads_wrt_par.append(flex.sum( g.select(sel_) )) # now apply constraints self.g = pack_gradients(x = grads_wrt_par, table = self.constrained_groups_selections) return self.f, self.g def set_refinable_parameters(xray_structure, parameters, selections, enforce_positivity=False): # XXX PVA: Code below is terribly inefficient and MUST be moved into C++ sz = xray_structure.scatterers().size() i = 0 for sel in selections: # pre-check for positivity begin # spread negative occupancies across i_seqs having positive ones par_all = flex.double() par_neg = flex.double() i_p = i for sel_ in sel: p = parameters[i_p] par_all.append(p) if(p<0): par_neg.append(p) i_p += 1 if(enforce_positivity and par_neg.size()>0): par_all = par_all - flex.min(par_all) fs = flex.sum(par_all) if(fs != 0): par_all = par_all / fs # pre-check for positivity end for j, sel_ in enumerate(sel): sel__b = flex.bool(sz, flex.size_t(sel_)) xray_structure.set_occupancies(par_all[j], sel__b) i+=1 def pack_unpack(x, table): result = x.deep_copy() for indices in table : if (len(indices) == 1): i0 = indices[0] result[i0] = x[i0] else : xsum = 0 for i in indices[0:-1] : result[i] = x[i] xsum += x[i] result[indices[-1]] = 1. - xsum return result def pack_gradients(x, table): result = flex.double(x.size(), 0) for indices in table : if(len(indices) == 1): i0 = indices[0] result[i0] = x[i0] elif(len(indices) == 2): i0,i1 = indices result[i0] = x[i0] - x[i1] result[i1] =-x[i1] # ??? zero or this value? else : for i in indices[0:-1] : result[i] = x[i] - x[indices[-1]] result[indices[-1]] = 0 return result #----------------------------------------------------------------------- # SELECTION HANDLING def list_3d_as_bool_selection(list_3d, size): result = flex.bool(size, False) for i in list_3d: for j in i: for k in j: if (result[k]): raise Sorry("Duplicate selection for occupancies.") result[k] = True return result def add_occupancy_selection(result, size, selection, hd_special=None): result_as_1d_array_b = list_3d_as_bool_selection(list_3d=result, size=size) sel_b = selection if(isinstance(selection, flex.size_t)): sel_b = flex.bool(size, selection) if(hd_special is not None): not_common = ((sel_b != result_as_1d_array_b) & (sel_b == True)) not_common_ = ((not_common != hd_special) & (not_common == True)).iselection() not_common = not_common_ else: not_common = ((sel_b != result_as_1d_array_b) & (sel_b == True)).iselection() sel_checked = [] for i in not_common: sel_checked.append([[i]]) if(len(sel_checked) > 0): result.extend(sel_checked) return result def extract_partial_occupancy_selections(hierarchy): result = [] for model in hierarchy.models(): for chain in model.chains(): for residue_group in chain.residue_groups(): if(not residue_group.have_conformers()): assert len(residue_group.atom_groups()) == 1 occs = flex.double() i_seqs = [] for atom in residue_group.atoms(): occs.append(atom.occ) i_seqs.append(atom.i_seq) if(occs[0]<1 and occs[0]!=0 and occs.all_eq(occs[0]) and occs.size()>1): result.append([i_seqs]) return result def occupancy_selections( model, add_water = False, other_individual_selection_strings = None, other_constrained_groups = None, remove_selection = None, as_flex_arrays = True, constrain_correlated_3d_groups = False, log = None): # set up defaults if(other_individual_selection_strings is not None and len(other_individual_selection_strings) == 0): other_individual_selection_strings = None if(other_constrained_groups is not None and len(other_constrained_groups) == 0): other_constrained_groups = None if(remove_selection is not None and len(remove_selection) == 0): remove_selection = None result = model.get_hierarchy().occupancy_groups_simple( common_residue_name_class_only = None, always_group_adjacent = False, ignore_hydrogens = False) exchangable_hd_pairs = model.get_hierarchy().exchangeable_hd_selections() if(len(exchangable_hd_pairs)==0 and result is not None): occupancy_regroupping( pdb_hierarchy = model.get_hierarchy(), cgs = result) result = mmtbx.utils.remove_selections(selection = result, other = exchangable_hd_pairs, size = model.get_number_of_atoms()) result.extend(exchangable_hd_pairs) # extract group-[0,1]-constrained atoms withing a residue pogl = extract_partial_occupancy_selections(hierarchy = model.get_hierarchy()) rm_duplicate_with_pogl = [] for t_ in pogl: for t__ in t_: for t___ in t__: rm_duplicate_with_pogl.append(t___) result = mmtbx.utils.remove_selections(selection = result, other = pogl, size = model.get_number_of_atoms()) result.extend(pogl) # add partial occupancies occupancies = model.get_xray_structure().scatterers().extract_occupancies() sel = (occupancies != 1.) & (occupancies != 0.) result = add_occupancy_selection( result = result, size = model.get_number_of_atoms(), selection = sel, hd_special = None) # check user's input all_sel_strgs = [] if(other_individual_selection_strings is not None): all_sel_strgs = all_sel_strgs + other_individual_selection_strings if(other_constrained_groups is not None): for other_constrained_group in other_constrained_groups: for other_constrained_group in other_constrained_groups: if(len(other_constrained_group.selection)>0): all_sel_strgs = all_sel_strgs + other_constrained_group.selection if(len(all_sel_strgs) > 0): for sel_str in all_sel_strgs: sel_str_sel = get_atom_selections( model = model, selection_strings = [sel_str], iselection = True, one_selection_array = True) if(sel_str_sel.size() == 0): raise Sorry("Empty selection: %s"%sel_str) # if([other_individual_selection_strings, other_constrained_groups].count(None) == 0): sel1 = get_atom_selections( model = model, selection_strings = other_individual_selection_strings, iselection = True, one_selection_array = True) for other_constrained_group in other_constrained_groups: for other_constrained_group in other_constrained_groups: for cg_sel_strs in other_constrained_group.selection: sel2 = get_atom_selections( model = model, selection_strings = cg_sel_strs, iselection = True, one_selection_array = True) if(sel1.intersection(sel2).size() > 0): raise Sorry("Duplicate selection: same atoms selected for individual and group occupancy refinement.") # check user's input and apply remove_selection to default selection if(remove_selection is not None): sel1 = get_atom_selections( model = model, selection_strings = remove_selection, iselection = True, one_selection_array = True) if(sel1.size() == 0): # XXX check all and not total. raise Sorry("Empty selection: remove_selection.") if(other_individual_selection_strings is not None): sel2 = get_atom_selections( model = model, selection_strings = other_individual_selection_strings, iselection = True, one_selection_array = True) if(sel1.intersection(sel2).size() > 0): raise Sorry("Duplicate selection: occupancies of same atoms selected to be fixed and to be refined.") if(other_constrained_groups is not None): for other_constrained_group in other_constrained_groups: for cg_sel_strs in other_constrained_group.selection: sel2 = get_atom_selections( model = model, selection_strings = cg_sel_strs, iselection = True, one_selection_array = True) if(sel1.intersection(sel2).size() > 0): raise Sorry("Duplicate selection: occupancies of same atoms selected to be fixed and to be refined.") result = mmtbx.utils.remove_selections(selection = result, other = sel1, size = model.get_number_of_atoms()) # if(other_individual_selection_strings is not None): sel = get_atom_selections( model = model, selection_strings = other_individual_selection_strings, iselection = True, one_selection_array = True) result = mmtbx.utils.remove_selections(selection = result, other = sel, size = model.get_number_of_atoms()) result = add_occupancy_selection( result = result, size = model.get_number_of_atoms(), selection = sel, hd_special = None) if(other_constrained_groups is not None): for other_constrained_group in other_constrained_groups: cg_sel = [] for cg_sel_strs in other_constrained_group.selection: sel = get_atom_selections( model = model, selection_strings = cg_sel_strs, iselection = True, one_selection_array = True) result = mmtbx.utils.remove_selections(selection = result, other = sel, size = model.get_number_of_atoms()) if(sel.size() > 0): cg_sel.append(list(sel)) if(len(cg_sel) > 0): result.append(cg_sel) if(add_water): water_selection = get_atom_selections( model = model, selection_strings = ['water'], iselection = True, allow_empty_selection = True, one_selection_array = True) result = add_occupancy_selection( result = result, size = model.get_number_of_atoms(), selection = water_selection, hd_special = None) list_3d_as_bool_selection( list_3d=result, size=model.get_number_of_atoms()) if(len(result) == 0): result = None if(as_flex_arrays and result is not None): result_ = [] for gsel in result: result__ = [] for sel in gsel: result__.append(flex.size_t(sel)) result_.append(result__) result = result_ if (constrain_correlated_3d_groups) and (len(result) > 0): result = assemble_constraint_groups_3d( xray_structure=model.get_xray_structure(), pdb_atoms=model.get_atoms(), constraint_groups=result, log=log) return result def occupancy_regroupping(pdb_hierarchy, cgs): h = pdb_hierarchy awl = list(h.atoms_with_labels()) elements = h.atoms().extract_element() rgs = list(h.residue_groups()) for cg in cgs: # loop over constraint groups for c in cg: # loop over conformers of constrained group altloc_h = awl[c[0]].altloc ind=None for ind_ in c: if(elements[ind_].strip().upper() in ["H","D"] and awl[ind_].name.strip().upper() in ["H","D"] and altloc_h.strip() != ""): ind = ind_ break if(ind is not None): # find "-1" (previous to given constraint group) residue group rg_prev = None for i_rg, rg in enumerate(rgs): if(i_rg > 0 and ind in rg.atoms().extract_i_seq()): rg_prev = rgs[i_rg-1] break if(rg_prev is None): continue rg_prev_i_seqs = rg_prev.atoms().extract_i_seq() rg_i_seqs = rg.atoms().extract_i_seq() # find i_seq of C of previous residue i_seqs_c_prev=[] for a_prev in rg_prev.atoms(): if(a_prev.name.strip().upper()=="C"): i_seqs_c_prev.append(a_prev.i_seq) if(i_seqs_c_prev == []): continue # find constarint group corresponding to rg_prev cg_prev=None for cg2 in cgs: for c2 in cg2: i_seqs_inter = set(c2).intersection(set(rg_prev_i_seqs)) if(len(i_seqs_inter)>1 or (len(i_seqs_inter)==1 and not awl[list(i_seqs_inter)[0]].name.strip().upper() in ["H","D"])): for cg2_ in cg2: for i_seq_c_prev in i_seqs_c_prev: if(i_seq_c_prev in cg2_): cg_prev = cg2 break if(cg_prev is not None): break if(cg_prev is None): continue # identify to which constraint group H belongs and move it there for cg2 in cgs: for c2 in cg2: if(ind in c2): c2.remove(ind) break found = False for conformer_prev in rg_prev.conformers(): if(conformer_prev.altloc == altloc_h): conformer_prev_i_seqs = conformer_prev.atoms().extract_i_seq() for cg2 in cgs: for c2 in cg2: i_seqs_inter = set(c2).intersection(set(conformer_prev_i_seqs)) if(len(i_seqs_inter)>1 or (len(i_seqs_inter)==1 and not awl[list(i_seqs_inter)[0]].name.strip().upper() in ["H","D"])): for i_seq_c_prev in i_seqs_c_prev: if(i_seq_c_prev in c2): c2.append(ind) found = True for cg in cgs: while cg.count([None])>0: cg.remove([None]) while cg.count([])>0: cg.remove([]) while cgs.count([])>0: cgs.remove([]) def assemble_constraint_groups_3d( xray_structure, pdb_atoms, constraint_groups, interaction_distance_cutoff=4.0, verbose=False, log=None): """ Re-sorts occupancy constraint groups so that conformers whose motion is correlated (i.e. they interact in 3D, without necessarily being part of the same fragment/molecule/ASU) are grouped together. As input, it expects the constraint groups output by mmtbx.utils.occupancy_selections(), which will already have connectivity taken into account. This function will exit with an error if the occupancies for the new groups are not consistent. """ if (log is None): log = null_out() make_sub_header("Correlated occupancy grouping", out=log) print(""" !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! !! WARNING - EXPERIMENTAL FEATURE !! !! !! !! Grouping of occupancy constraints in 3D is experimental and not fully !! !! tested. Use at your own risk! For bug reports, etc. contact us by !! !! email at help@phenix-online.org. !! !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! """, file=log) occupancies = pdb_atoms.extract_occ() pair_asu_table = xray_structure.pair_asu_table( distance_cutoff=interaction_distance_cutoff) pair_sym_table = pair_asu_table.extract_pair_sym_table() k = 0 n_groups_start = len(constraint_groups) while (k < len(constraint_groups)): groups = constraint_groups[k] print("Constraint group %d: %d conformers" % (k+1, len(groups)), file=log) merge_constraints = [] for i_sel, selection in enumerate(groups): occ = occupancies.select(selection) altloc = pdb_atoms[selection[0]].fetch_labels().altloc print(" conformer '%s': %d atoms" % (altloc, len(selection)), file=log) if (not occ.all_eq(occ[0])): raise Sorry("At least one occupancy constraint group has "+ "inconsistent occupancies for atoms in a single conformer. To use "+ "the automatic 3D constraints, the starting occupancies must be "+ "uniform within each selection.") for i_seq in selection : labels = pdb_atoms[i_seq].fetch_labels() if (labels.altloc.strip() == ''): continue pair_sym_dict = pair_sym_table[i_seq] if (verbose): print("%s (group %d):" % (pdb_atoms[i_seq].id_str(), k+1)) for j_seq, sym_ops in pair_sym_dict.items(): kk = k + 1 while (kk < len(constraint_groups)): combine_group = False for other_selection in constraint_groups[kk] : if (j_seq in other_selection): if (verbose): print(" %s (group %d)" % (pdb_atoms[j_seq].id_str(), kk+1)) merge_constraints.append(constraint_groups[kk]) del constraint_groups[kk] combine_group = True break if (not combine_group): kk += 1 if (len(merge_constraints) > 0): print("Merging %d constraint groups with group %d" % ( len(merge_constraints), (k+1)), file=log) for selection in groups : first_atom = pdb_atoms[selection[0]] altloc = first_atom.fetch_labels().altloc if (altloc.strip() == ''): raise RuntimeError(("Atom '%s' in occupancy constraint group has "+ "blank altloc ID") % first_atom.id_str()) for merge_groups in merge_constraints : kk = 0 while (kk < len(merge_groups)): other_selection = merge_groups[kk] altloc2 = pdb_atoms[other_selection[0]].fetch_labels().altloc if (altloc2 == altloc): print(" combining %d atoms with altloc %s" % \ (len(other_selection), altloc), file=log) occ1 = occupancies.select(selection) occ2 = occupancies.select(other_selection) if (not occ1.all_eq(occ2[0])) or (not occ2.all_eq(occ1[0])): raise Sorry( ("Inconsistent occupancies in spatially related groups "+ "(%.2f versus %.2f). To use automatic 3D occupancy "+ "restraints, the correlated conformers must start with "+ "the same initial occupancy.") % (occ1[0], occ2[0])) selection.extend(other_selection) del merge_groups[kk] else : kk += 1 for merge_groups in merge_constraints : if (len(merge_groups) > 0): for other_selection in merge_groups : altloc = pdb_atoms[other_selection[0]].fetch_labels().altloc print((" warning: %d atoms with altloc %s do not "+ "correspond to an existing group") % (len(other_selection), altloc), file=log) groups.append(other_selection) k += 1 if (len(constraint_groups) != n_groups_start): print("New occupancy constraint groups:", file=log) for i_group, constraint_group in enumerate(constraint_groups): print(" group %d:" % (i_group+1), file=log) for selection in constraint_group : resids = [] altlocs = set() for i_seq in selection : atom_group = pdb_atoms[i_seq].parent() ag_id = atom_group.id_str() altlocs.add(atom_group.altloc) if (not ag_id in resids): resids.append(ag_id) assert len(altlocs) == 1 print(" conformer '%s' (%d atoms):" % (list(altlocs)[0], len(selection)), file=log) for ag_id in resids : print(" atom_group %s" % ag_id, file=log) else : print("Occupancy constraint groups unmodified.", file=log) print("", file=log) return constraint_groups