How to use the biotite.structure.io.load_structure function in biotite

def test_remove_pbc_unsegmented():
    """
    `remove_pbc()` should not alter unsegmented structures,
    when the structure is entirely in the box.
    Exclude the solvent, due to high distances between each atom. 
    """
    ref_array = load_structure(join(data_dir("structure"), "3o5r.mmtf"))
    # Center structure in box
    centroid = struc.centroid(ref_array)
    box_center = np.diag(ref_array.box) / 2
    ref_array = struc.translate(ref_array, box_center-centroid)
    # Remove solvent
    ref_array = ref_array[~struc.filter_solvent(ref_array)]
    array = struc.remove_pbc(ref_array)

    assert ref_array.equal_annotation_categories(array)
    assert np.allclose(ref_array.coord, array.coord)

def test_nucleotide_filter():
    nuc_sample_array = strucio.load_structure(
        join(data_dir("structure"), "5ugo.pdb")
    )

    assert len(nuc_sample_array[struc.filter_nucleotides(nuc_sample_array)
            ]) == 651

def test_rdf_periodic():
    """ Test if the periodic argument gives the correct results"""
    test_file = TEST_FILE
    stack = load_structure(test_file)

    # calculate oxygen RDF for water
    oxygen = stack[:, stack.atom_name == 'OW']
    interval = np.array([0, 10])
    n_bins = 100
    bins, g_r = rdf(oxygen[:, 0].coord, oxygen[:, 1:], interval=interval,
                    bins=n_bins, periodic=True)

    # Compare with MDTraj
    import mdtraj
    traj = mdtraj.load(TEST_FILE)
    ow = [a.index for a in traj.topology.atoms if a.name == 'O']
    pairs = itertools.product([ow[0]], ow[1:])
    mdt_bins, mdt_g_r = mdtraj.compute_rdf(traj, list(pairs),
                                           r_range=interval/10, n_bins=n_bins,
                                           periodic=True)

def test_rdf_atom_argument():
    """ Test if the first argument allows to use AtomArrayStack """
    stack = load_structure(TEST_FILE)

    # calculate oxygen RDF for water with and without a selection
    oxygen = stack[:, stack.atom_name == 'OW']
    interval = np.array([0, 10])
    n_bins = 100

    bins, g_r = rdf(oxygen[:, 0], stack, interval=interval,
                    bins=n_bins, periodic=False)

    atom_bins, atoms_g_r = rdf(oxygen[:, 0].coord, stack, interval=interval,
                    bins=n_bins, periodic=False)

    assert np.allclose(g_r, atoms_g_r)

# Atom array stacks require a *(m,3,3)*-shaped :class:`ndarray`, 
# that contains the box vectors for each model.
# The box is accessed via the `box` attribute, which is ``None`` by
# default.
# When loaded from a structure file, the box described in the file is
# automatically used.

import biotite.database.rcsb as rcsb
import biotite.structure.io as strucio

array = struc.AtomArray(length=100)
print(array.box)
array.box = struc.vectors_from_unitcell(10, 20, 30, np.pi/2, np.pi/2, np.pi/2)
print(array.box)
file_path = rcsb.fetch("3o5r", "mmtf", gettempdir())
array = strucio.load_structure(file_path)
print(array.box)

########################################################################
# When loading a trajectory from an MD simulation, the molecules are
# often fragmented over the periodic boundary.
# While a lot of analysis functions can handle such periodic boundary
# conditions automatically, some require completed molecules.
# In this case you should use :func:`remove_pbc()`.

array = struc.remove_pbc(array)

########################################################################
# Structure manipulation
# ----------------------
#
# The most basic way to manipulate a structure is to edit the

The structure was resolved using NMR, so multiple models are present
in the structure.
Hence, we can also calculate the frequency of each bond.
"""

# Code source: Daniel Bauer
# License: BSD 3 clause

import biotite
import matplotlib.pyplot as plt
import biotite.structure as struc
import biotite.structure.io as strucio
import biotite.database.rcsb as rcsb

file_name = rcsb.fetch("2KB1", "mmtf", biotite.temp_dir())
stack = strucio.load_structure(file_name)
# Four identical chains, consider only chain A
chain_a = stack[:, stack.chain_id == "A"]
# Selection for p-helix
p_helix = (chain_a.res_id >= 40) & (chain_a.res_id <= 52)
# Selection for selectivity filter
sf = (chain_a.res_id >= 53) & (chain_a.res_id <= 58)

# Calculate the hydrogen bonds and the frequency of each bond
triplets, mask = struc.hbond(chain_a, selection1=p_helix, selection2=sf)
freq = struc.hbond_frequency(mask)

# Create names of bonds
label = "{d_resid}{d_resnm}-{d_a} -- {a_resid}{a_resnm}-{a_a}"
names = [label.format(
    d_resid=chain_a.res_id[donor],
    d_resnm=chain_a.res_name[donor],

# residue
    leaflet_masks = np.empty(
        (len(head_leaflets), structure.array_length()),
        dtype=bool
    )
    for i, head_leaflet in enumerate(head_leaflets):
        leaflet_masks[i] = struc.get_residue_masks(structure, head_leaflet) \
                                .any(axis=0)
    return leaflet_masks


# Suppress warning that elements were guessed,
# as this PDB file omits the 'chemical element' column
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    structure = strucio.load_structure(PDB_FILE_PATH)
# We cannot go over periodic boundaries in this case,
# because the input PDB does not define a box -> periodic=False
# However, as we have a planer lipid bilayer,
# periodicity should not matter
leaflets = find_leaflets(
    structure,
    head_atom_mask=(structure.res_name == "DPP") & (structure.atom_name == "P")
)
# Bilayer -> Expect two leaflets
assert len(leaflets) == 2
# Mark leaflets using different chain IDs
for chain_id, leaflet_mask in zip(("A", "B"), leaflets):
    structure.chain_id[leaflet_mask] = chain_id

# Save marked lipids to structure file
temp = NamedTemporaryFile(suffix=".pdb")

# Since programmers are usually lazy and do not want to write more code
# than necessary, there are two convenient function for loading and
# saving atom arrays or stacks, unifying the forementioned file formats:
# :func:`load_structure()` takes a file path and outputs an array
# (or stack, if the file contains multiple models).
# Internally, this function uses the appropriate :class:`File` class,
# depending on the file format.
# The analogous :func:`save_structure()` function provides a shortcut
# for writing to structure files.
# The desired file format is inferred from the the extension of the
# provided file name.

import biotite.structure.io as strucio

stack_from_pdb = strucio.load_structure(pdb_file_path)
stack_from_cif = strucio.load_structure(cif_file_path)
temp_file = NamedTemporaryFile(suffix=".cif")
strucio.save_structure(temp_file.name, stack_from_pdb)
temp_file.close()

########################################################################
# Reading trajectory files
# ^^^^^^^^^^^^^^^^^^^^^^^^
# 
# If the package *MDtraj* is installed, *Biotite* provides a read/write
# interface for different trajectory file formats.
# All supported trajectory formats have in common, that they store
# only coordinates.
# These can be extracted as :class:`ndarray` with the
# :func:`get_coord()` method.

from tempfile import NamedTemporaryFile