def make_cg_mol(aa_mol, mapping, bead_names, bead_types): molname = aa_mol.graph['name'] cg_mol = Molecule(nrexcl=1, meta=dict(moltype=molname)) mapdict = defaultdict(list) for bd_idx, at_idxs in enumerate(mapping): name = bead_names[bd_idx] for member in at_idxs: mapdict[member].append(bd_idx) subgraph = aa_mol.subgraph(at_idxs) charge = sum(nx.get_node_attributes(subgraph, 'charge').values()) position = np.mean([ subgraph.nodes[idx].get('position', (np.nan, np.nan, np.nan)) for idx in subgraph ], axis=0) cg_mol.add_node(bd_idx, atomname=name, resname=molname, resid=1, atype=bead_types[bd_idx], charge_group=bd_idx + 1, graph=subgraph, charge=charge, position=position) for aa_idx, aa_jdx in aa_mol.edges: cg_idxs = mapdict[aa_idx] cg_jdxs = mapdict[aa_jdx] for cg_idx, cg_jdx in product(cg_idxs, cg_jdxs): if cg_idx != cg_jdx: cg_mol.add_edge(cg_idx, cg_jdx) for idx, jdx in cg_mol.edges: cg_mol.add_interaction('bonds', [idx, jdx], []) return cg_mol
def test_apply_mod_mapping(modified_molecule, modifications): """ Test apply_mod_mapping """ graph_out = Molecule(force_field=FF_UNIVERSAL) graph_out.add_nodes_from([ (0, {'atomname': 'A', 'resid': 1}) ]) mol_to_out = {0: {0: 1}} out_to_mol = {0: {0: 1}} match = ({1: {'mA': 1}}, modifications['mA'], {}) out = apply_mod_mapping(match, modified_molecule, graph_out, mol_to_out, out_to_mol) print(mol_to_out) print(out_to_mol) print(graph_out.nodes[1]) print(modifications['mA'].nodes['mA']) for key in modifications['mA'].nodes['mA']: assert graph_out.nodes[1][key] == modifications['mA'].nodes['mA'][key] assert out == ({}, {}) assert mol_to_out == {0: {0: 1}, 1: {1: 1}} assert out_to_mol == {0: {0: 1}, 1: {1: 1}} graph_out.add_node(2, atomname='J', resid=2) mol_to_out[16] = {2: 1} out_to_mol[2] = {16: 1} out = apply_mod_mapping(( {16: {'J': 1}, 17: {'mJ': 1}, 18: {'mJ2': 1}}, modifications['mJ', 'mJ2'], {} ), modified_molecule, graph_out, mol_to_out, out_to_mol) print(mol_to_out) print(out_to_mol) assert out == ({}, {}) assert mol_to_out == {0: {0: 1}, 1: {1: 1}, 16: {2: 1}, 17: {3: 1}, 18: {4: 1}} assert out_to_mol == {0: {0: 1}, 1: {1: 1}, 2: {16: 1}, 3: {17: 1}, 4: {18: 1}}