flt17 / confined-water-analysis Goto Github PK

def load_with_cell(filename_or_filenames, **kwargs):
    """Load a trajectory and inject cell dimensions from a topology PDB file if not present.
    All arguments and keyword arguments are passed on to `mdtraj.load`. The `top`
    keyword argument is used load a PDB file and get cell information from it.
    """

    # load the "topology frame" to get cell dimensions
    top = kwargs.get("top")
    if top is not None and isinstance(top, str):
        # load topology frame - just the first one from the file, in case there are more frames
        frame_top = mdt.load_frame(top, 0)
        unitcell_lengths = frame_top.unitcell_lengths
        unitcell_angles = frame_top.unitcell_angles
        if (unitcell_lengths is None) or (unitcell_angles is None):
            raise ValueError("Frame providing topology is missing cell information.")
    else:
        raise ValueError("Provide a PDB with cell dimensions.")

    # load the trajectory itself
    trj = mdt.load(filename_or_filenames, **kwargs)

    # inject the cell information
    len_trj = len(trj)
    trj.unitcell_lengths = unitcell_lengths.repeat(len_trj, axis=0)
    trj.unitcell_angles = unitcell_angles.repeat(len_trj, axis=0)

    return trj