def get_id_map_1(
functional_group: stk.FunctionalGroup,
) -> dict[int, int]:
new_id = max(functional_group.get_atom_ids()) + 1
num_atoms = sum(1 for _ in functional_group.get_atoms())
ids = range(new_id, new_id+num_atoms)
return dict(zip(functional_group.get_atom_ids(), ids))
def get_id_map_0(functional_group: stk.FunctionalGroup, ) -> dict[int, int]:
"""
Get an id_map with a single atom.
"""
# Make sure new_id is always valid, by making 1 larger than the
# biggest on in the functional group. This prevents two atoms in
# the functional group from having the same id.
new_id = max(functional_group.get_atom_ids()) + 1
ids = (new_id, )
return dict(zip(functional_group.get_atom_ids(), ids))
def test_with_ids(
generic_functional_group: stk.FunctionalGroup,
get_id_map: Callable[[stk.FunctionalGroup], dict[int, int]],
) -> None:
"""
Test :meth:`.FunctionalGroup.with_ids`.
Parameters:
functional_group:
The functional group to test.
get_id_map:
Takes a single parameter, `functional_group`, and returns a
valid `id_map` parameter for its
:meth:`.FunctionalGroup.with_atoms` method. This allows the
testing of different values of this parameter.
"""
# Save a clone to ensure that "functional_group" is not changed by
# the test, because it should be immutable.
before = generic_functional_group.clone()
_test_with_atoms(generic_functional_group, get_id_map)
is_clone_functional_group(before, generic_functional_group)
def _test_with_atoms(
functional_group: stk.FunctionalGroup,
get_id_map: abc.Callable[[stk.FunctionalGroup], dict[int, int]],
) -> None:
"""
Test :meth:`.FunctionalGroup.with_atoms`.
Parameters:
functional_group:
The functional group to test.
get_id_map:
Takes a single parameter, `functional_group`, and returns a
valid `id_map` parameter for its
:meth:`.FunctionalGroup.with_atoms` method. This allows the
testing of different values of this parameter.
"""
id_map = get_id_map(functional_group)
clone = functional_group.with_ids(id_map)
is_modified_id_sequence(
original_ids=functional_group.get_atom_ids(),
modified_ids=clone.get_atom_ids(),
id_map=id_map,
)
is_modified_id_sequence(
original_ids=functional_group.get_placer_ids(),
modified_ids=clone.get_placer_ids(),
id_map=id_map,
)
is_modified_id_sequence(
original_ids=functional_group.get_core_atom_ids(),
modified_ids=clone.get_core_atom_ids(),
id_map=id_map,
)