Testing Utilities#

Utility functions for writings tests for AtomArray objects.

atomworks.io.utils.testing.assert_same_annotation_cardinality(arr1: AtomArray, arr2: AtomArray, annotations: list[str]) → None[source]#

Assert annotations have the same number of unique values in both arrays.

Parameters:

  • arr1 – First atom array.

  • arr2 – Second atom array.

  • annotations – List of annotation names to check cardinality for.

atomworks.io.utils.testing.assert_same_atom_array_or_stack(arr1: AtomArray | AtomArrayStack, arr2: AtomArray | AtomArrayStack, compare_coords: bool = True, compare_bonds: bool = True, compare_box: bool = False, annotations_to_compare: list[str] | Literal['arr1'] | None = None, enforce_order: bool = True, compare_bond_order: bool = True, cast_to_common_dtype: bool = False, _n_mismatches_to_show: int = 5) → None[source]#

Asserts that two AtomArray or AtomArrayStack objects are equal.

Parameters:

  • arr1 (AtomArray) – The first AtomArray or AtomArrayStack to compare.

  • arr2 (AtomArray) – The second AtomArray or AtomArrayStack to compare.

  • compare_coords (bool, optional) – Whether to compare coordinates. Defaults to True.

  • compare_bonds (bool, optional) – Whether to compare bonds. Defaults to True.

  • compare_box (bool, optional) – Whether to compare the box attribute. Defaults to False.

  • annotations_to_compare (list[str] | Literal["arr1"] | None, optional) – List of annotation categories to compare, or “arr1” to compare on all annotations in arr1. Defaults to None, in which case all annotations are compared.

  • enforce_order (bool, optional) – Whether to enforce the order of the atoms. Defaults to True. NOTE: Enforcing order is much faster; use False only when strictly necessary.

  • compare_bond_order (bool, optional) – Whether to compare bond order. Defaults to True.

  • cast_to_common_dtype (bool, optional) – Whether to cast numeric annotations to a common dtype before comparison. Useful for tests where dtype preservation is not guaranteed (e.g., CIF roundtrips). Defaults to False.

  • _n_mismatches_to_show (int, optional) – Number of mismatches to show. Defaults to 5.

Raises:

AssertionError – If the AtomArray or AtomArrayStack objects are not equal.

atomworks.io.utils.testing.get_pdb_path(pdbid: str, mirror_path: str | PathLike = '/mnt/data/frozen_pdb_copies/2026_01_06_pdb') → str[source]#

Get the local path to a PDB file based on the provided mirror path.

Parameters:

  • pdbid (str) – The PDB ID.

  • mirror_path (str | os.PathLike, optional) – Path to the PDB mirror directory. Defaults to PDB_MIRROR_PATH constant.

Returns:

The local path to the PDB file.

Return type:

str

Raises:

FileNotFoundError – If the file does not exist at the expected location or if no mirror path is provided.

atomworks.io.utils.testing.get_pdb_path_or_buffer(pdb_id: str) → str | StringIO[source]#

Returns a local file path or an in-memory buffer for a given PDB ID.

Parameters:

pdb_id (str) – The PDB identifier of the structure.

Returns:

The local file path to the structure file if available, otherwise an in-memory buffer containing the fetched file.

Return type:

str | io.StringIO

atomworks.io.utils.testing.has_ambiguous_annotation_set(atom_array: AtomArray, annotation_set: Iterable[str] = ('chain_id', 'res_id', 'res_name', 'atom_name', 'ins_code')) → bool[source]#

Detect whether a given set of annotations is insufficient to distinguish all atoms.

Used to detect ambiguous annotations that would lose information on CIF write, since struct_conn distinguishes bonds by the 5-tuple (chain_id, res_id, res_name, atom_name, ins_code).

atomworks.io.utils.testing.is_same_in_group(groups: ndarray, data: ndarray) → ndarray[source]#

Check if all elements in data are the same within each group defined by groups.

Parameters:

  • groups – 1D array of group identifiers, same length as data.

  • data – 1D array of data values to check for sameness within each group.

Returns:

Boolean array of shape (n_groups,) indicating whether all elements in each group are the same.

Return type:

np.ndarray

Example

>>> groups = np.array([1, 1, 2, 2, 2, 3])
>>> data = np.array([5, 5, 7, 7, 7, 9])
>>> is_same_in_group(groups, data)
array([ True,  True,  True])
>>> data = np.array([5, 5, 7, 8, 7, 9])
>>> is_same_in_group(groups, data)
array([ True, False,  True])
atomworks.io.utils.testing.is_same_in_segment(segment_start_stop: ndarray, data: ndarray, raise_if_false: bool = False) → ndarray[source]#

Check if all elements in a segment are the same.

Parameters:

  • segment_start_stop (np.ndarray) – Array of segment start and stop indices (end of segment is inclusive), as obtained from struc.get_residue_starts(… add_exclusive_stop=True) for example.

  • data (np.ndarray) – Data array to check for sameness within segments.

Returns:

Boolean array indicating whether all elements in each segment are the same.

Return type:

np.ndarray

atomworks.io.utils.testing.verify_atom_array_chain_info_consistency(chain_info: dict[str, Any], atom_array: AtomArray, verify_sequences: bool = True) → None[source]#

Verify that atom array and chain_info are consistent.

Parameters:

  • chain_info – The chain_info dictionary.

  • atom_array – The processed atom array.

  • verify_sequences – Whether to verify 1-letter sequences match. Defaults to True.

Raises:

AssertionError – If inconsistencies are detected between atom array and chain_info.