def test_do_mapping_mods(modified_molecule, modifications):
    """
    Test do_mapping on a molecule with modifications
    """
    mappings = {}
    for name in 'ABCDEFGHIJ':
        block = Block(force_field=FF_UNIVERSAL, name=name)
        block.add_node(name, atomname=name, resid=1)
        mappings[name] = Mapping(block, block, mapping={name: {name: 1}},
                                 references={}, ff_from=FF_UNIVERSAL,
                                 ff_to=FF_UNIVERSAL, names=(name,))
    for mod in modifications.values():
        name = (mod.name,) if isinstance(mod.name, str) else mod.name
        mappings[name] = Mapping(mod, mod, mapping={idx: {idx: 1} for idx in mod},
                                 references={}, ff_from=FF_UNIVERSAL,
                                 ff_to=FF_UNIVERSAL, names=name, type='modification')
    mappings = {'universal': {'universal': mappings}}

    out = do_mapping(modified_molecule, mappings, FF_UNIVERSAL,
                     attribute_keep=('chain', 'resid'))
    pprint.pprint(list(out.nodes(data=True)))
    pprint.pprint(list(out.edges))

    expected = modified_molecule.copy()
    for node in expected.nodes:
        expected.nodes[node]['mapping_weights'] = {node: 1}
        expected.nodes[node]['graph'] = expected.subgraph([node])

    expected = nx.relabel_nodes(expected, {idx: idx+1 for idx in expected})
    pprint.pprint(list(expected.nodes(data='atomname')))
    pprint.pprint(list(expected.edges))

    assert equal_graphs(expected, out, node_attrs=['atomname', 'resid', 'mapping_weights'])
    assert out.interactions['bond'] == [Interaction(atoms=(10, 11), parameters=(3, 4), meta={})]
def test_no_residue_crossing():
    """
    Make sure we don't cross residue boundaries
    """
    mapping = {'C1': {'B1': 1}, 'C2': {'B1': 1}, 'C3': {'B1': 1}}
    extra = ()
    mappings = {'universal': {'martini22': {'IPO': Mapping(FF_UNIVERSAL.blocks['IPO'],
                                                           FF_MARTINI.blocks['IPO'],
                                                           mapping=mapping,
                                                           references={},
                                                           ff_from=FF_UNIVERSAL,
                                                           ff_to=FF_MARTINI,
                                                           names=('IPO',),
                                                           extra=extra)}}}

    cg = do_mapping(AA_MOL, mappings, FF_MARTINI, attribute_keep=['chain'])

    expected = Molecule(force_field=FF_MARTINI)
    expected.add_nodes_from((
        (0, {'resid': 1, 'resname': 'IPO', 'atomname': 'B1', 'chain': 'A', 'charge_group': 1}),
        (1, {'resid': 2, 'resname': 'IPO', 'atomname': 'B1', 'chain': 'A', 'charge_group': 2}),
        (2, {'resid': 3, 'resname': 'IPO', 'atomname': 'B1', 'chain': 'A', 'charge_group': 3}),
    ))
    expected.add_edges_from(([0, 1], [1, 2]))

    print(cg.nodes(data=True))
    print(cg.edges())
    print('-'*80)
    print(expected.nodes(data=True))
    print(expected.edges())
    assert equal_graphs(cg, expected)
def test_mod_matches(modified_molecule, modifications):
    """
    Test modification matches
    """
    mappings = []
    for name in 'ABCDEFGHIJ':
        block = Block(force_field=FF_UNIVERSAL, name=name)
        block.add_node(name, atomname=name)
        mappings.append(Mapping(block, block, mapping={name: {name: 1}},
                                references={}, ff_from=FF_UNIVERSAL,
                                ff_to=FF_UNIVERSAL, names=(name,)))
    for mod in modifications.values():
        name = (mod.name,) if isinstance(mod.name, str) else mod.name
        mappings.append(Mapping(mod, mod, mapping={idx: {idx: 1} for idx in mod},
                                references={}, ff_from=FF_UNIVERSAL,
                                ff_to=FF_UNIVERSAL, names=name, type='modification'))

    print([m.names for m in mappings])
    found = list(modification_matches(modified_molecule, mappings))
    expected = [
        ({1: {'mA': 1}}, modifications['mA'], {}),
        ({4: {'mC': 1}}, modifications['mC'], {}),
        ({6: {'mD': 1}}, modifications['mD'], {}),
        ({9: {'mF': 1}, 11: {'mG': 1}}, modifications['mFG'], {}),
        ({14: {'mI': 1}}, modifications['mI'], {}),
        ({15: {'mI2': 1}}, modifications['mI2'], {}),
        ({16: {'J': 1}, 17: {'mJ': 1}, 18: {'mJ2': 1}}, modifications['mJ', 'mJ2'], {}),
    ]
    pprint.pprint(found)
    pprint.pprint(expected)
    print([e[1].name for e in expected])
    print([e[1].name for e in found])
    assert len(expected) == len(found)
    e_matches, e_mods, e_refs = zip(*expected)
    f_matches, f_mods, f_refs = zip(*found)

    for f_match in f_matches:
        assert f_match in e_matches
    for f_ref in f_refs:
        assert f_ref in e_refs
    e_mod_names = [(mod.name,) if isinstance(mod.name, str) else mod.name for mod in e_mods]
    for f_mod in f_mods:
        assert f_mod.name in e_mod_names
def test_residue_crossing():
    '''
    Make sure we do cross residue boundaries and can rename residues
    '''
    mapping = {'C1': {'B1': 1}, 'C2': {'B1': 1}, 'C3': {'B1': 1},
               'C4': {'B2': 1}, 'C5': {'B2': 1}, 'C6': {'B2': 1},
               'C7': {'B3': 1}, 'C8': {'B3': 1}, 'C9': {'B3': 1},
               }
    extra = ()
    mappings = {
        'universal': {
            'martini22': {
                'IPO_large': Mapping(FF_UNIVERSAL.blocks['IPO_large'],
                                     FF_MARTINI.blocks['IPO_large'],
                                     mapping=mapping,
                                     references={},
                                     ff_from=FF_UNIVERSAL,
                                     ff_to=FF_MARTINI,
                                     names=('IPO', 'IPO', 'IPO'),
                                     extra=extra)
            }
        }
    }

    cg = do_mapping(AA_MOL, mappings, FF_MARTINI, attribute_keep=('chain',))

    expected = Molecule()
    expected.add_nodes_from((
        (0, {'resid': 1, 'resname': 'IPO_large', 'atomname': 'B1', 'chain': 'A', 'charge_group': 1}),
        (1, {'resid': 2, 'resname': 'IPO_large', 'atomname': 'B1', 'chain': 'A', 'charge_group': 1}),
        (2, {'resid': 3, 'resname': 'IPO_large', 'atomname': 'B1', 'chain': 'A', 'charge_group': 1}),
    ))
    expected.add_edges_from(([0, 1], [1, 2]))

    print(cg.nodes(data=True))
    print(cg.edges())
    print('-'*80)
    print(expected.nodes(data=True))
    print(expected.edges())
    assert equal_graphs(cg, expected)
def test_peptide():
    """
    Test multiple cg beads in residue
    """
    gly = {'C': {'BB': 1}, 'N': {'BB': 1}, 'O': {'BB': 1}, 'CA': {'BB': 1}}
    ile = {'C': {'BB': 1}, 'N': {'BB': 1}, 'O': {'BB': 1}, 'CA': {'BB': 1},
           'CB': {'SC1': 1}, 'CG1': {'SC1': 1}, 'CG2': {'SC1': 1}, 'CD': {'SC1': 1}}
    leu = {'C': {'BB': 1}, 'N': {'BB': 1}, 'O': {'BB': 1}, 'CA': {'BB': 1},
           'CB': {'SC1': 1}, 'CG': {'SC1': 1}, 'CD1': {'SC1': 1}, 'CD2': {'SC1': 1}}
    extra = ()

    mappings = {'universal': {'martini22': {}}}
    mappings['universal']['martini22']['GLY'] = Mapping(FF_UNIVERSAL.blocks['GLY'],
                                                        FF_MARTINI.blocks['GLY'],
                                                        mapping=gly,
                                                        references={},
                                                        ff_from=FF_UNIVERSAL,
                                                        ff_to=FF_MARTINI,
                                                        names=('GLY',),
                                                        extra=extra)
    mappings['universal']['martini22']['ILE'] = Mapping(FF_UNIVERSAL.blocks['ILE'],
                                                        FF_MARTINI.blocks['ILE'],
                                                        mapping=ile,
                                                        references={},
                                                        ff_from=FF_UNIVERSAL,
                                                        ff_to=FF_MARTINI,
                                                        names=('ILE',),
                                                        extra=extra)
    mappings['universal']['martini22']['LEU'] = Mapping(FF_UNIVERSAL.blocks['LEU'],
                                                        FF_MARTINI.blocks['LEU'],
                                                        mapping=leu,
                                                        references={},
                                                        ff_from=FF_UNIVERSAL,
                                                        ff_to=FF_MARTINI,
                                                        names=('LEU',),
                                                        extra=extra)

    peptide = Molecule(force_field=FF_UNIVERSAL)
    aa = FF_UNIVERSAL.blocks['GLY'].to_molecule()
    peptide.merge_molecule(aa)

    aa = FF_UNIVERSAL.blocks['ILE'].to_molecule()
    peptide.merge_molecule(aa)

    aa = FF_UNIVERSAL.blocks['LEU'].to_molecule()
    peptide.merge_molecule(aa)

    peptide.add_edge(list(peptide.find_atoms(atomname='N', resid=1))[0],
                     list(peptide.find_atoms(atomname='C', resid=2))[0])
    peptide.add_edge(list(peptide.find_atoms(atomname='N', resid=2))[0],
                     list(peptide.find_atoms(atomname='C', resid=3))[0])

    for node in peptide:
        peptide.nodes[node]['atomid'] = node + 1
        peptide.nodes[node]['chain'] = ''

    cg = do_mapping(peptide, mappings, FF_MARTINI, attribute_keep=('chain',))

    expected = Molecule(force_field=FF_MARTINI)
    expected.add_nodes_from({1: {'atomname': 'BB',
                                 'atype': 'P5',
                                 'chain': '',
                                 'charge': 0.0,
                                 'charge_group': 1,
                                 'resid': 1,
                                 'resname': 'GLY'},
                             2: {'atomname': 'BB',
                                 'atype': 'P5',
                                 'chain': '',
                                 'charge': 0.0,
                                 'charge_group': 2,
                                 'resid': 2,
                                 'resname': 'ILE'},
                             3: {'atomname': 'SC1',
                                 'atype': 'AC1',
                                 'chain': '',
                                 'charge': 0.0,
                                 'charge_group': 3,
                                 'resid': 2,
                                 'resname': 'ILE'},
                             4: {'atomname': 'BB',
                                 'atype': 'P5',
                                 'chain': '',
                                 'charge': 0.0,
                                 'charge_group': 4,
                                 'resid': 3,
                                 'resname': 'LEU'},
                             5: {'atomname': 'SC1',
                                 'atype': 'AC1',
                                 'chain': '',
                                 'charge': 0.0,
                                 'charge_group': 5,
                                 'resid': 3,
                                 'resname': 'LEU'}}.items()
                            )
    expected.add_edges_from([(1, 2), (2, 3), (2, 4), (4, 5)])

    for node in expected:
        expected.nodes[node]['atomid'] = node + 1
    print(cg.nodes(data=True))
    print(cg.edges())
    print('-'*80)
    print(expected.nodes(data=True))
    print(expected.edges())

    assert equal_graphs(cg, expected)