Skip to content

Commit

Permalink
Test run 5
Browse files Browse the repository at this point in the history
  • Loading branch information
@choglass
choglass committed Jun 5, 2024
1 parent b9a41c5 commit eda8773
Show file tree
Hide file tree
Showing 5 changed files with 416 additions and 74 deletions.
92 changes: 82 additions & 10 deletions cell2mol/charge_assignment.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,17 @@
#######################################################
def get_possible_charge_state(spec: object, debug: int=0):
if not hasattr(spec,"protonation_states"): spec.get_protonation_states(debug=debug)
if spec.protonation_states is None: return None
if spec.subtype == "group" or (spec.subtype == 'molecule' and spec.iscomplex): return None
if spec.protonation_states is None:
return None
if spec.formula in ["O4-Cl", "N3", "I3"]:
ch_state = get_charge_manual(spec, debug=debug)
possible_cs = [ch_state]
return possible_cs

if spec.subtype == "group" or (spec.subtype == 'molecule' and spec.iscomplex):
return None
print(f"GET_POSSIBLE_CHARGE_STATE: {spec.formula} ({spec.subtype}) {spec.cov_factor=}")
charge_states = []
charge_states = []
### Evaluates possible charges for each protonation state ###
for prot in spec.protonation_states:
# charge_states = []
Expand Down Expand Up @@ -178,8 +185,8 @@ def get_protonation_states_specie(specie: object, debug: int=0) -> list:
##############################
if specie.type != "specie": return None
if specie.subtype == "group": return None
elif specie.subtype == "molecule" and specie.iscomplex == True: return None
elif specie.subtype == "molecule" and specie.iscomplex == False:
elif specie.subtype == "molecule" and specie.iscomplex == True: return None
elif (specie.subtype == "molecule" and specie.iscomplex == False) or specie.formula in ["O4-Cl", "N3", "I3"]:
if debug >= 2: print(f"\nPOSCHARGE: doing empty PROTONATION for this specie {specie.formula} ({specie.subtype})")
#empty_list = list([np.zeros((len(specie.labels)))])
empty_list = []
Expand Down Expand Up @@ -391,10 +398,11 @@ def get_protonation_states_specie(specie: object, debug: int=0) -> list:
if a.connec == 0:
elemlist[idx] = "Cl"
addedlist[idx] = 1
else:
# block[idx] = 1
elif a.connec == 1:
elemlist[idx] = "Cl"
addedlist[idx] = 1
else:
block[idx] = 1
# Nitrogen
elif a.label == "N":
# Nitrosyl
Expand Down Expand Up @@ -593,6 +601,70 @@ def get_protonation_states_specie(specie: object, debug: int=0) -> list:
if debug >= 2: print(f" GET_PROTONATION_STATES: Protonation DISCARDED. Steric Clashes found when adding atoms. status={new_prot.status}")
if debug >= 2: print(f" GET_PROTONATION_STATES:{protonation_states=}")
return protonation_states
#######################################################
def move_to_front(lst, index):
element = lst.pop(index) # Remove the element from its current position
lst.insert(0, element) # Insert the element at the beginning
return lst
#######################################################
def move_element(lst, old_index, new_index):
element = lst.pop(old_index) # Remove the element from its current position
lst.insert(new_index, element) # Insert the element at the new position
return lst
#######################################################
def get_charge_manual(spec, debug: int=0):

if spec.formula == "O4-Cl":
smiles = "[O-]Cl(=O)(=O)=O"
charge = -1
order = [0, 1, 2, 3, 4] # Cl index is 1
for idx, a in enumerate(spec.atoms):
if a.label == "Cl":
new_order = move_element(order, 1, idx)
if debug >= 2: print(f" O4-Cl: {new_order=}")
elif spec.formula == "N3":
smiles = "[N-]=[N+]=[N-]"
charge = -1
order = [0, 1, 2]
for idx, a in enumerate(spec.atoms):
list_of_adj_atoms = []
for adj in a.adjacency:
if debug >= 2: print(f" N3: {adj=}", spec.get_parent("molecule").labels[adj])
list_of_adj_atoms.append(spec.get_parent("molecule").labels[adj])
numN = list_of_adj_atoms.count("N")
if numN == 2:
new_order = move_element(order, 1, idx)
if debug >= 2: print(f" N3: {new_order=}")

elif spec.formula == "I3":
smiles = "I[I-]I"
charge = -1
order = [0, 1, 2]
for idx, a in enumerate(spec.atoms):
list_of_adj_atoms = []
for adj in a.adjacency:
if debug >= 2: print(f" I3: {adj=}", spec.get_parent("molecule").labels[adj])
list_of_adj_atoms.append(spec.get_parent("molecule").labels[adj])
numI = list_of_adj_atoms.count("I")
if numI == 2:
new_order = move_element(order, 1, idx)
if debug >= 2: print(f" I3: {new_order=}")

temp_mol = Chem.MolFromSmiles(smiles, sanitize=False)
mol = Chem.RenumberAtoms(temp_mol, new_order)
atom_charge = []
total_charge = 0
for i in range(spec.natoms):
a = mol.GetAtomWithIdx(i) # Returns a particular Atom
atom_charge.append(a.GetFormalCharge())
total_charge += a.GetFormalCharge()

# iscorrect = check_rdkit_obj_connectivity(mol, spec.natoms, charge, debug=debug)
iscorrect = True
allow = True
prot = spec.get_protonation_states()[0]
ch_state = charge_state(iscorrect, total_charge, atom_charge, mol, smiles, charge, allow, prot)
return ch_state

######################################################
def get_charge(charge: int, prot: object, allow: bool=True, embed_chiral: bool=True, debug: int=0):
Expand Down Expand Up @@ -628,9 +700,9 @@ def get_charge(charge: int, prot: object, allow: bool=True, embed_chiral: bool=T
for mol in mols:
for i in range(natoms):
a = mol.GetAtomWithIdx(i)
if a.GetExplicitValence() != a.GetTotalValence():
if debug >=2 : print(f"GET_CHARGE. {i} {a.GetSymbol()=}, {a.GetFormalCharge()=}, \
{a.GetImplicitValence()=}, {a.GetExplicitValence()=} {a.GetTotalValence()=}")
# if a.GetExplicitValence() != a.GetTotalValence():
if debug >=2 : print(f"GET_CHARGE. {i} {a.GetSymbol()=}, {a.GetFormalCharge()=}, \
{a.GetImplicitValence()=}, {a.GetExplicitValence()=} {a.GetTotalValence()=}")
return None

# Smiles are generated with rdkit
Expand Down
2 changes: 1 addition & 1 deletion cell2mol/new_cell_reconstruction.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -417,7 +417,7 @@ def get_updated_indices(sp_idx, new, cell_labels, cell_pos, cell_fracs, debug: i
for kdx, (l, p, f) in enumerate(zip(cell_labels, cell_pos, cell_fracs)):
if n_l == l and np.allclose(n_p, p, atol=1e-4, rtol=1e-2):
if np.allclose(np.remainder(n_f, 1), np.remainder(f, 1), atol=1e-4, rtol=1e-2):
if debug >= 2:
if debug > 2:
print(f"symmtry operation {sp_idx}:", f"atom of new (index: {jdx})", n_l, n_p, n_f, \
f"is the same as the atom of the unit cell (index: {kdx})", l, p, f)
indices_lists.append((jdx, kdx))
Expand Down
85 changes: 22 additions & 63 deletions cell2mol/new_charge_assignment.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import numpy as np
import copy
from cell2mol.charge_assignment import protonation, get_charge
from cell2mol.charge_assignment import protonation, get_charge, get_charge_manual
import itertools

#######################################################
Expand Down Expand Up @@ -101,79 +101,38 @@ def print_possible_and_selected_cs (newcell, refcell, debug: int=0):
if idx != specie.unique_index:
print(f"WARNING: {specie.formula=} {specie.unique_index=} {idx=} from newcell unique indices")

######################################################
def set_charge_state_simple (reference, target, debug: int=0):

final_charge = reference.totcharge
print("SET_CHARGE_STATE:", reference.charge_state)

if target.subtype == "molecule" and target.iscomplex == False:
if debug >=1 : print(f"({target.subtype}) {target.formula} {final_charge=} Create Empty PROTONATION for this specie")
empty_list = [int(0)]*len(target.labels)
empty_prot = protonation(target.labels, target.coord, target.cov_factor,
int(0), empty_list, empty_list, empty_list, empty_list, typ="Empty", parent=target)
cs = get_charge(final_charge, empty_prot)

elif target.subtype == "ligand":
if debug >=1 : print(f"({target.subtype}) {target.formula} {reference.charge_state.uncorr_total_charge=} Ligand")
target.get_protonation_states(debug=debug)
prot = target.protonation_states[0]
cs = get_charge(reference.charge_state.uncorr_total_charge, prot)

if len(target.protonation_states) != 1 :
if debug >=1 : print("WARNING:", target.protonation_states)

target.charge_state = cs
if final_charge != cs.corr_total_charge:
print(f"WARNING: {target.formula=} {final_charge=} {cs.corr_total_charge=} final_charge != cs.corr_total_charge")

target.set_charges(cs.corr_total_charge, cs.corr_atom_charges, cs.smiles, cs.rdkit_obj)
print(f"SET_CHARGE_STATE. {target.formula=} {target.totcharge=} {target.smiles=}")

######################################################
def set_charge_state(reference, target, mode, debug: int=0):

final_charge = reference.totcharge
print("SET_CHARGE_STATE:", reference.charge_state)

if target.subtype == "molecule" and target.iscomplex == False:
if debug >=1 : print(f"SET_CHARGE_STATE:({target.subtype}) {target.formula} {final_charge=} Create Empty PROTONATION for this specie")
empty_list = [int(0)]*len(target.labels)
empty_prot = protonation(target.labels, target.coord, target.cov_factor,
int(0), empty_list, empty_list, empty_list, empty_list, typ="Empty", parent=target)
cs = get_charge(final_charge, empty_prot)

elif target.subtype == "ligand":
if debug >=1 : print(f"SET_CHARGE_STATE:({target.subtype}) {target.formula} {reference.charge_state.uncorr_total_charge=} Ligand")
prot = reference.charge_state.protonation
temp_prot = copy.deepcopy(prot)
temp_prot.parent = target

if mode == 1:
# For "reference" cell
target.get_protonation_states(debug=debug)
if mode == 1 : # For "reference" cell. Their possible charge states are already calculated
if target.formula in ["O4-Cl", "N3", "I3"]:
cs = get_charge_manual(target, debug=debug)
else :
target.get_possible_cs(debug=debug)
charge_list = [cs.corr_total_charge for cs in target.possible_cs]
idx = charge_list.index(final_charge)
cs = target.possible_cs[idx]
# if len(target.protonation_states) != 1 :
# if debug >=1 : print(f"WARNING: {target.protonation_states=}")

# ref_data, target_data = arrange_data_for_reorder(reference, target)
# if debug >=2 : print(ref_data, target_data)
# dummy1, dummy2, map12 = reorder(ref_data, target_data, reference.coord, target.coord)

# if np.array_equal(map12, np.arange(len(target_data))):
# if debug >=1 : print(f"({target.subtype}) {target.formula} {reference.charge_state.uncorr_total_charge=} No need to reorder")
# # temp_prot.coords = target.coord
# cs = get_charge(reference.charge_state.uncorr_total_charge, temp_prot)
# else:
# reordered_prot = temp_prot.reorder(map12)
# # reordered_prot.coords = target.coord
# if debug >=1 : print(f"({target.subtype}) {target.formula} {reference.charge_state.uncorr_total_charge=} Reordered {map12=}")
# cs = get_charge(reference.charge_state.uncorr_total_charge, reordered_prot)

elif mode == 2:

elif mode == 2 : # For "unit" cell. Their charge state are not calculated
if target.formula in ["O4-Cl", "N3", "I3"]:
cs = get_charge_manual(target, debug=debug)
elif (target.subtype == "molecule" and target.iscomplex == False) :
if debug >=1 : print(f"SET_CHARGE_STATE:({target.subtype}) {target.formula} {final_charge=} Create Empty PROTONATION for this specie")
empty_list = [int(0)]*len(target.labels)
empty_prot = protonation(target.labels, target.coord, target.cov_factor,
int(0), empty_list, empty_list, empty_list, empty_list, typ="Empty", parent=target)
cs = get_charge(final_charge, empty_prot)

elif target.subtype == "ligand":
if debug >=1 : print(f"SET_CHARGE_STATE:({target.subtype}) {target.formula} {reference.charge_state.uncorr_total_charge=} Ligand")
prot = reference.charge_state.protonation
temp_prot = copy.deepcopy(prot)
temp_prot.parent = target

print(f"SET_CHARGE_STATE:{temp_prot.labels=} {len(temp_prot.labels)=} {len(temp_prot.coords)=} {len(temp_prot.block)=} {temp_prot.added_atoms=}")
# For "unit" cell
ref_data = reference.get_parent_indices("reference")
Expand Down
Loading

0 comments on commit eda8773

Please sign in to comment.