def test_frame_indices(self):
        traj = pt.iterload("data/tz2.truncoct.nc", "data/tz2.truncoct.parm7")
        traj2 = pt.iterload("data/tz2.truncoct.nc",
                            "data/tz2.truncoct.parm7",
                            frame_slice=(2, 8))

        txt = '''
        reference ./data/tz2.truncoct.nc 2 2
        rmsd :2-11 refindex 0 perres perresout center.agr range 1 perrescenter
        '''

        state = pt.load_batch(traj2, txt)
        state.run()

        frame_indices = range(2, 8)
        rmsd0 = pt.rmsd(traj, ref=1, mask=':2-11', frame_indices=frame_indices)
        rmsdperres = pt.rmsd_perres(traj,
                                    ref=1,
                                    mask=':2-11',
                                    perres_mask='*',
                                    resrange='1',
                                    perres_center=True,
                                    frame_indices=frame_indices)
        aa_eq(rmsd0, state.data[2])
        aa_eq(rmsdperres[1], state.data[3].values)

    def test_frame_fit(self):
        traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
        f0 = traj[0]
        f1 = traj[1]

        arr0 = list(f0[0])
        arr1 = list(f1[0])

        f0.rmsd(f1)
        aa_eq(arr0, f0[0])
        aa_eq(arr1, f1[0])

        f1.rmsfit(ref=f0)

        # expect reference `f0` xyz are not changed
        aa_eq(arr0, f0[0])

        trajsaved = pt.iterload("./data/fit_to_1stframe.Tc5b.x",
                                "./data/Tc5b.top")
        f1saved = trajsaved[1]

        # make sure we reproduce cpptraj output
        aa_eq(f1.xyz, f1saved.xyz, decimal=3)

        farray = traj[:]
        farray.rmsfit(ref=traj[0])
        aa_eq(farray[1].xyz, f1saved.xyz, decimal=3)

    def test_split_and_write_traj(self):
        fn = "data/Tc5b.x"
        traj = pt.iterload([fn, fn], "./data/Tc5b.top")
        # duplcate
        assert traj.n_frames == 20
        top = traj.top

        # test TrajectoryIterator object
        pt.tools.split_and_write_traj(traj,
                                      n_chunks=4,
                                      root_name='./output/trajiterx',
                                      overwrite=True)
        flist = sorted(glob("./output/trajiterx*"))
        traj4 = pt.iterload(flist, top)
        aa_eq(traj4.xyz, traj.xyz)

        # dcd ext
        pt.tools.split_and_write_traj(traj,
                                      4,
                                      root_name='./output/ts',
                                      ext='dcd',
                                      overwrite=True)
        flist = sorted(glob("./output/ts.*.dcd"))
        traj4 = pt.iterload(flist, top)
        aa_eq(traj4.xyz, traj.xyz)

    def test_0(self):
        traj = pt.iterload("./data/tz2.ortho.nc", "./data/tz2.ortho.parm7")
        trajectory_traj = traj[:]

        # test xyz
        aa_eq(trajectory_traj.xyz, traj.xyz)

        # test object lifetime
        aa_eq(trajectory_traj[0].xyz, trajectory_traj.xyz[0])

        # test Box
        assert (trajectory_traj.top.has_box() == True)
        boxes = traj.unitcells
        for i, frame in enumerate(trajectory_traj):
            assert (frame.has_box() == True)
            f_blist = frame.box.tolist()
            aa_eq(f_blist, boxes[i].tolist())

        # test autoimage
        # make Trajectory from TrajectoryIterator
        fa = traj[:]
        fa.autoimage()
        saved_traj = pt.iterload("./data/tz2.autoimage.nc",
                                 "./data/tz2.ortho.parm7")

        # make sure to reproduce cpptraj's output too
        aa_eq(saved_traj.xyz, fa.xyz)

    def test_1(self):

        # load frames to immutable traj
        traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
        trajsaved = pt.iterload("./data/fit_to_1stframe.Tc5b.x",
                                "./data/Tc5b.top")

        for _f1 in trajsaved:
            pass

        f0saved = traj[0].copy()
        first = traj[0].copy()

        # make mutable traj
        farray = traj[:]

        aa_eq(farray[0].xyz, first.xyz)
        farray.rmsfit(ref=first, mask="*", mass=False)
        farray2 = traj[:]
        farray2.superpose(ref=first, mask="*", mass=False)

        for i, _f0 in enumerate(farray):
            _f1 = trajsaved[i]
            aa_eq(_f0.xyz, _f1.xyz, decimal=3)

        for i, _f0 in enumerate(farray2):
            _f1 = trajsaved[i]
            aa_eq(_f0.xyz, _f1.xyz, decimal=3)

    def test_fancy_save(self):
        traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
        traj[1:8].save("./output/test_fancy_save_frame1_7.x", overwrite=True)

        fanew = pt.iterload("./output/test_fancy_save_frame1_7.x", traj.top)

        for idx, f0 in enumerate(traj[1:8]):
            f0new = fanew[idx]
            aa_eq(f0.xyz, f0new.xyz)

 def test_1(self):
     traj = pt.iterload("./data/tz2.truncoct.nc",
                        "./data/tz2.truncoct.parm7")
     f0 = traj[0]
     f0cp = f0.copy()
     adict['autoimage']("", f0, traj.top)
     fsaved = pt.iterload("./data/tz2.truncoct.autoiamge.save.r",
                          "./data/tz2.truncoct.parm7")[0]
     aa_eq(fsaved.xyz, f0.xyz, decimal=3)

    def test_volmap(self):
        traj = pt.iterload("./data/tz2.ortho.nc", "./data/tz2.ortho.parm7")[:1]
        size = ''
        center = ''
        state = pt.load_cpptraj_state(txt.format(cm, size, center))
        state.run()
        cpp_data = state.data[-1].values

        traj = traj.superpose(mask=':1-13').center(':1-13 mass origin')
        ds = pt.volmap(traj,
                       mask=':[email protected]',
                       grid_spacing=(0.5, 0.5, 0.5),
                       buffer=2.0,
                       centermask='!:1-13',
                       radscale=1.36)

        aa_eq(cpp_data, ds)

        # assert
        self.assertRaises(AssertionError, lambda: pt.volmap(traj, mask=':[email protected]', grid_spacing='0.5 0.5 0.5'))
        self.assertRaises(AssertionError, lambda: pt.volmap(traj, mask=':[email protected]', grid_spacing=(0.5, 0.5)))
        self.assertRaises(ValueError, lambda: pt.volmap(traj, mask=':[email protected]', grid_spacing=(0.5, 0.5, 0.5), size='20 20 20'))

        # test size
        cm_no_buffer = cm.replace('buffer 2.0', '')
        state = pt.load_cpptraj_state(txt.format(cm_no_buffer, 'size 20,20,20', ''))
        state.run()
        cpp_data = state.data[-1].values
        ds = volmap(traj,
                    mask=':[email protected]',
                    grid_spacing=(0.5, 0.5, 0.5),
                    size=(20, 20, 20),
                    buffer=2.0,
                    centermask='!:1-13',
                    radscale=1.36)
        aa_eq(cpp_data, ds)

        # test center
        state = pt.load_cpptraj_state(txt.format(cm_no_buffer, 'size 20,20,20', 'center 0.5,0.5,0.5'))
        state.run()
        cpp_data = state.data[-1].values
        ds = volmap(traj,
                    mask=':[email protected]',
                    grid_spacing=(0.5, 0.5, 0.5),
                    size=(20, 20, 20),
                    center=(0.5, 0.5, 0.5),
                    buffer=2.0,
                    centermask='!:1-13',
                    radscale=1.36)
        aa_eq(cpp_data, ds)

        # raise RuntimeError
        dry_traj = pt.iterload('data/tz2.nc', 'data/tz2.parm7')
        self.assertRaises(RuntimeError, lambda: pt.volmap(dry_traj, mask=':[email protected]',
                                                          grid_spacing=(0.5, 0.5, 0.5)))

    def test_randomize_ions(self):
        fn = os.path.join(cpptraj_test_dir, 'Test_RandomizeIons', 'adh206.tip3p.rst7.gz')
        tn = os.path.join(cpptraj_test_dir, 'Test_RandomizeIons', 'adh206.ff10.tip3p.parm7.gz')
        saved_traj_name = os.path.join(cpptraj_test_dir, 'Test_RandomizeIons', 'random.crd.save')

        traj = pt.iterload(fn, tn)
        traj_mut = traj[:]
        saved_traj = pt.iterload(saved_traj_name, traj.top)

        pt.randomize_ions(traj_mut, mask='@Na+', around=':1-16', by=5.0, overlap=3.0, seed=113698)
        aa_eq(traj_mut.xyz, saved_traj.xyz, decimal=2)

    def test_load_from_list(self):
        from glob import glob
        pattern = "./data/Test_RemdTraj/rem.nc.*"
        flist = sorted(glob(pattern))
        top = glob("./data/Test_RemdTraj/ala*parm7")[0]
        traj0 = pt.iterload(flist, top)
        traj1 = pt.iterload(pattern, top)
        aa_eq(traj0.xyz, traj1.xyz)

        # raise if not find files
        self.assertRaises(ValueError,
                          lambda: pt.iterload("./data/xyz_cool*.x", traj0.top))

    def test_2(self):
        from pytraj.all_actions import do_autoimage
        # test do_autoimage
        traj = pt.iterload("./data/tz2.truncoct.nc",
                           "./data/tz2.truncoct.parm7")
        f0 = traj[0]
        f0cp = f0.copy()
        do_autoimage(traj=f0, top=traj.top)

        fsaved = pt.iterload("./data/tz2.truncoct.autoiamge.save.r",
                             "./data/tz2.truncoct.parm7")[0]
        aa_eq(fsaved.xyz, f0.xyz, decimal=3)

    def test_comprehensive(self):
        traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
        # make sure we DO reproducing cpptraj output
        f_saved = pt.iterload("./data/avg.Tc5b.pdb", traj.top)[0]

        # shorter
        frame2 = mean_structure(traj)
        aa_eq(frame2.xyz, f_saved.xyz, decimal=3)

        frame3 = mean_structure(traj=traj)
        aa_eq(frame3.xyz, f_saved.xyz, decimal=3)

        # test list
        frame4 = mean_structure(traj=[traj, traj[:3]], top=traj.top)

        # test iter
        frame5 = mean_structure(traj=traj(1, 8, 2), top=traj.top)
        f5_saved = pt.iterload("./data/avg.Tc5b.frame_2_to_8_skip_2.pdb",
                               traj.top)[0]
        aa_eq(frame5.xyz, f5_saved.xyz, decimal=3)

        # test iter CA
        frame5 = mean_structure(traj[[0, 3, 7]], '@CA', top=traj.top)

        # use atom_indices
        ca_indices = pt.select('@CA', traj.top)
        frame5_1 = mean_structure(traj[[0, 3, 7]], ca_indices, top=traj.top)

        # test frame_indices
        frame6 = mean_structure(traj, mask='@CA', frame_indices=[0, 3, 7])
        aa_eq(frame5.xyz, frame6.xyz, decimal=3)
        aa_eq(frame5_1.xyz, frame6.xyz, decimal=3)

        xyz_0 = pt.get_coordinates(traj(1, 8, 2))
        xyz_1 = np.array([frame.xyz.copy(
        ) for frame in traj.iterframe(frame_indices=range(1, 8, 2))])
        aa_eq(xyz_0, xyz_1, decimal=3)

        # test as traj
        out_traj = mean_structure(traj,
                                  mask='@CA',
                                  frame_indices=[0, 3, 7],
                                  dtype='traj')
        assert isinstance(out_traj, Trajectory), 'must be Trajectory'
        aa_eq(out_traj.xyz, frame6.xyz, decimal=3)

        # raise if not trajectory, traj or frame
        self.assertRaises(ValueError, lambda: pt.mean_structure(traj, dtype='trajxyz'))

    def test_1(self):
        # merge_frames
        import numpy as np

        # load 2 frames
        traj = pt.iterload("./data/Tc5b.x",
                           "./data/Tc5b.top",
                           frame_slice=(0, 2))

        # test mutable traj
        t0 = traj[:]
        frame = pt.tools.merge_frames(t0)
        r0 = pt.tools.rmsd(frame.xyz, pt.get_coordinates(traj), True)
        assert r0 == 0.

        # test immutable traj
        assert np.any(pt.tools.merge_frames(traj).xyz.flatten() ==
                      pt.get_coordinates(traj).flatten())

        # tuple
        assert np.any(pt.tools.merge_frames((
            frame for frame in traj)).xyz.flatten() == pt.get_coordinates(
                traj).flatten())

        # list
        assert np.any(pt.tools.merge_frames(
            [frame for frame in traj]).xyz.flatten() == pt.get_coordinates(
                traj).flatten())

        # frame_iter: all atoms
        assert np.any(pt.tools.merge_frames(traj()).xyz.flatten() ==
                      pt.get_coordinates(traj()).flatten())

    def test_comprehensive(self):
        traj = pt.iterload("data/Test_RemdTraj/rem.nc.000",
                           "data/Test_RemdTraj/ala2.99sb.mbondi2.parm7")
        # temperature
        aa_eq(traj.temperatures, [300., 630.5, 630.5, 630.5, 630.5, 630.5,
                                  630.5, 630.5, 492.2, 384.3])

        # iterframe (already in doctest), just throwing raise to increase coverage score
        self.assertRaises(ValueError, lambda: traj.iterframe(rmsfit='crazy'))

        # raise
        # memory error if load larger than 1GB for xyz
        traj = pt.datafiles.load_tz2_ortho()
        for _ in range(11):
            traj._load(traj.filelist)

        self.assertRaises(MemoryError, lambda: traj.xyz)

        # can not find filename
        self.assertRaises(ValueError, lambda: traj._load(None))
        # has filename but does not have Topology
        self.assertRaises(
            ValueError,
            lambda: pt.TrajectoryIterator("data/Test_RemdTraj/rem.nc.000", top=None))
        self.assertRaises(
            ValueError,
            lambda: pt.TrajectoryIterator("data/Test_RemdTraj/rem.nc.000"))
        # empty Topology
        self.assertRaises(
            ValueError,
            lambda: pt.TrajectoryIterator("data/Test_RemdTraj/rem.nc.000", top=pt.Topology()))
        # weird Topology
        self.assertRaises(
            ValueError,
            lambda: pt.TrajectoryIterator("data/Test_RemdTraj/rem.nc.000", top=pt.Frame))

 def test_2(self):
     traj = pt.iterload("./data/Tc5b.nat.crd", "./data/Tc5b.top")
     t0 = traj[:1]
     pt.set_dihedral(t0, resid='4', dihedral_type='phi', deg=120)
     dih = pt.calc_phi(t0, resrange='4').values[0]
     assert abs(dih - 120) < 1E-3
     t0.save('test.pdb', options='model', overwrite=True)

def main(argv):
    args = sys.argv
    scorefile = 'AmberScores.sc'

    decoy_list = glob( "{base}/5.AmberMinimize/min*.rst7".format( base=BASE_DIR ) )

    decoy_list = []
    for d in decoy_list_long:
        decoy_list.append(d.strip())
    print("Decoy List: " + str(len(decoy_list)))
    parmfile = "{base}/5.AmberMinimize/decoy.parm7".format( base=BASE_DIR )
    
    print("Making Traj.")
    traj = pt.iterload( decoy_list, parmfile )
    print("Getting energies...")
    energy_data = pt.pmap_mpi(pt.energy_decomposition, traj, igb=8)
    print("done!")
    if energy_data:
        ekeys = energy_data.keys()
        ekeys.sort()

        print('Scorefile: ' + scorefile)
        with open(scorefile,'w') as myscorefile:

            header = 'description\t'
            for key in ekeys:
                header += key + '\t'
            myscorefile.write(header +"\n")

            for decoy in range(len(decoy_list)):
                scoreline = decoy_list[decoy]+'\t'
                for key in ekeys:
                    scoreline += "%s\t" % str(energy_data[key][decoy])
                myscorefile.write(scoreline+"\n")

    def test_0(self):
        import numpy as np

        # load 2 frames
        traj = pt.iterload("./data/Tc5b.x",
                           "./data/Tc5b.top",
                           frame_slice=(0, 2))

        # test mutable traj
        big_frame = pt.tools.merge_frames(traj[:])
        assert pt.tools.rmsd(big_frame.xyz.flatten(),
                             pt.get_coordinates(traj).flatten()) < 1E-5

        # test immutable traj
        big_frame = pt.tools.merge_frames(traj)
        assert pt.tools.rmsd(big_frame.xyz.flatten(),
                             pt.get_coordinates(traj).flatten()) < 1E-5

        # test tuple
        big_frame = pt.tools.merge_frames((frame for frame in traj))
        assert pt.tools.rmsd(big_frame.xyz.flatten(),
                             pt.get_coordinates(traj).flatten()) < 1E-5

        # test list
        big_frame = pt.tools.merge_frames([frame.copy() for frame in traj])
        assert pt.tools.rmsd(big_frame.xyz.flatten(),
                             pt.get_coordinates(traj).flatten()) < 1E-5

    def test_radgyr_dispatch(self):
        traj = pt.iterload("./data/tz2.ortho.nc", "./data/tz2.ortho.parm7")

        # mask is an array-like
        aa_eq(sb._toy_radgyr(traj, [0, 3]), pt.radgyr(traj, '@1,4'))

        # frame_indices
        aa_eq(
            sb._toy_radgyr(traj,
                           '@CA',
                           frame_indices=[0, 3]),
            pt.radgyr(traj, '@CA')[[0, 3]])

        # frame_indices, mask in kwd
        aa_eq(
            sb._toy_radgyr(traj,
                           mask='@CA',
                           frame_indices=[0, 3]),
            pt.radgyr(traj, '@CA')[[0, 3]])

        # frame iterator
        aa_eq(
            sb._toy_radgyr(
                traj(0, 3),
                mask='@CA'),
            pt.radgyr(traj, '@CA')[[0, 1, 2]])

    def test_general(self):
        traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")

        # with mask
        saved_data = pt.radgyr(traj, '@CA')
        data = pt.pmap(pt.radgyr, traj, '@CA')
        data = pt.tools.dict_to_ndarray(data)
        aa_eq(saved_data, data)

        # with a series of functions
        func_list = [pt.radgyr, pt.molsurf, pt.rmsd]
        ref = traj[-3]

        for n_cores in [2, 3]:
            for func in func_list:
                if func in [pt.rmsd, ]:
                    pout = pt.tools.dict_to_ndarray(pt.pmap(func=func,
                                                            traj=traj,
                                                            ref=ref,
                                                            n_cores=n_cores))
                    serial_out = flatten(func(traj, ref=ref))
                else:
                    pout = pt.tools.dict_to_ndarray(pt.pmap(n_cores=n_cores,
                                                            func=func,
                                                            traj=traj))
                    serial_out = flatten(func(traj))
                aa_eq(pout[0], serial_out)

         # test worker
         # need to test this since coverages seems not recognize partial func
        from pytraj.parallel.multiprocessing_ import worker_byfunc
        data = worker_byfunc(rank=2, n_cores=8, func=pt.radgyr, traj=traj, args=(), kwd={'mask': '@CA'}, iter_options={})
        assert data[0] == 2, 'rank must be 2'
        assert data[2] == 1, 'n_frames for rank=2 should be 1 (only 10 frames in total)'

    def test_0(self):
        traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")

        traj(0, 8, 2, mask='@CA').save('output/test0.nc', overwrite=True)
        pt.write_traj('./output/test1.nc',
                      traj(0,
                           8,
                           2,
                           mask='@CA'),
                      overwrite=True)

        new_top = traj.top._get_new_from_mask('@CA')
        t0 = pt.iterload('./output/test0.nc', new_top)
        t1 = pt.iterload('./output/test1.nc', new_top)

        aa_eq(t0.xyz, t1.xyz)