def test_fit_and_then_nofit(self):
     traj = pt.iterload("data/Tc5b.x", "data/Tc5b.top")
     t0 = traj[:]
     pt.superpose(t0, ref=traj[3], mask='@CA')
     rmsd_0 = pt.rmsd_nofit(traj, ref=traj[3], mask='@CB')
     rmsd_1 = pt.rmsd(traj, ref=traj[3], mask='@CB', nofit=True)
     aa_eq(rmsd_1, rmsd_0)

    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_buffer_not_c_contiguous(self):
        # source code was lightly adapted from jmborr
        # https://github.com/Amber-MD/pytraj/issues/991
        traj = pt.load('data/issue991/short.dcd',
                       'data/issue991/pdb.gz',
                       mask='(!:1-256)&(@H1,@H2,@H3,@H4,@H5)')

        # Trajectory of the center of mass of the first two residues
        minitraj = np.empty((2, traj.n_frames, 3))
        minitraj[0] = pt.center_of_mass(traj, mask=':1')
        minitraj[1] = pt.center_of_mass(traj, mask=':2')
        minitraj = minitraj.transpose((1, 0, 2))

        # "topology" using the first two atoms
        minitop = traj.top['@1,2']

        # Save trajectory
        # make sure there is no ValueError
        # something is wrong with pdb, crd extension when loading with
        # minitop (poor topology?)
        # make issue in cpptraj too?
        for ext in ['nc', 'dcd', 'mdcrd', 'crd', 'pdb', 'trr']:
            fn = 'output/junk.' + ext
            pt.write_traj(filename=fn,
                          traj=minitraj,
                          top=minitop,
                          overwrite=True)

            # load coord back to make sure we correctly write it
            new_traj = pt.iterload(fn, minitop)
            # mdcrd, crd, pdb has only 3 digits after decimal
            decimal = 5 if ext in ['nc', 'dcd', 'trr'] else 3
            aa_eq(minitraj, new_traj.xyz, decimal=decimal)

 def testsuperpose_alias(self):
     '''testsuperpose_alias'''
     t0 = self.traj[:]
     t1 = self.traj[:]
     pt.transform(t0, ['superpose'])
     pt.transform(t1, ['rms'])
     aa_eq(t0.xyz, t1.xyz)

def main():
    # require 'sanderapi' and 'parmed' packages in AmberTools15
    # http://ambermd.org/#AmberTools
    # more info about `sanderapi` in Amber15 manual
    # http://ambermd.org/doc12/Amber15.pdf (page 341)
    traj = io.iterload("../tests/data/Tc5b.x",
                       "../tests/data/Tc5b.top")
    print(traj)
    edict0 = pyca.energy_decomposition(parm="../tests/data/Tc5b.top",
                                       traj=traj,
                                       igb=8,
                                       input_options=None)
    print(edict0)

    # edict1: use default igb=8
    edict1 = pyca.energy_decomposition(traj=traj)
    print(edict1)

    aa_eq(edict0['tot'], edict1['tot'])

    # update `input` and get minimal output
    import sander
    inp = sander.gas_input(6)
    edict2 = pyca.energy_decomposition(traj, input_options=inp, mode='minimal')
    print(edict2)

    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_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_vs_cpptraj(self):
     data = pt.dssp(self.traj, "*", dtype='cpptraj_dataset')
     data_int = np.array(
         [d0.values for d0 in data if d0.dtype == 'integer'],
         dtype='i4')
     # load cpptraj output
     cpp_data = np.loadtxt("./data/dssp.Tc5b.dat", skiprows=1)[:, 1:].T
     aa_eq(data_int.flatten(), cpp_data.flatten())

 def testsuperpose_vs_rmsd(self):
     # load frames to immutable traj
     traj = pt.iterload("data/tz2.nc", "data/tz2.parm7")
     t0 = traj[:]
     t1 = traj[:]
     pt.rmsd(t0, ref=traj[0], mask='@CA')
     pt.superpose(t1, ref=traj[0], mask='@CA')
     aa_eq(t0.xyz, t1.xyz)

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

        t0.autoimage()
        avg_0 = pt.mean_structure(t0, '@CA')
        avg_1 = pt.mean_structure(traj(autoimage=True), '@CA')
        aa_eq(avg_0.xyz, avg_1.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_not_update_coordinates(self):
        traj = self.traj[:]
        data = pt.rmsd(traj, ref=3, update_coordinate=False)

        # make sure coordinates are not updated
        aa_eq(traj.xyz, self.traj.xyz)

        # make sure give the same rmsd values
        aa_eq(pt.rmsd(traj, ref=3), data)

        def order_(modes):
            data = _ired(state_vecs, modes=modes)
            order_s2_v0 = data['IRED_00127[S2]']
            # make sure the S2 values is 1st array

            # load cpptraj's output and compare to pytraj' values for S2 order paramters
            cpp_order_s2 = np.loadtxt(os.path.join(
                cpptraj_test_dir, 'Test_IRED', 'orderparam.save')).T[-1]
            aa_eq(order_s2_v0.values, cpp_order_s2, decimal=4)

    def test_1(self):
        # status: OK
        from pytraj.compat import zip
        # note: use `load` instead of `iterload`
        traj = pt.load("./data/tz2.ortho.nc", "./data/tz2.ortho.parm7")
        traj.autoimage()
        traj.rmsfit(mask='@CA,C,N')
        saved_traj = pt.load('data/tz2.autoimage_with_rmsfit.nc', traj.top)

        # PASSED
        aa_eq(saved_traj.xyz, traj.xyz)

    def test_0(self):
        trajcpp = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
        farray = trajcpp[:]

        # bugs: trajcpp[0, 0, 0] != farray[0, 0, 0] (must be equal)
        assert farray[0][0, 0] == trajcpp[0][0, 0]

        farray[0, 0, 0] = 100.10
        assert farray[0, 0, 0] == 100.10
        farray[0, 0, :] = [100.10, 101.1, 102.1]
        aa_eq(farray[0, 0, :], [100.10, 101.1, 102.1])

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

        cout = pt.datafiles.load_cpptraj_output("""
        parm ./data/tz2.ortho.parm7
        trajin ./data/tz2.ortho.nc
        rms first nofit
        rms first mass
        """)
        aa_eq(pt.rmsd(traj, nofit=True), cout[1])
        aa_eq(pt.rmsd(traj, mass=True), cout[2])

 def test_noreference(self):
     from pytraj.datafiles import load_cpptraj_output, tz2_ortho_trajin
     traj = pt.iterload("./data/tz2.ortho.nc", "./data/tz2.ortho.parm7")
     cout = load_cpptraj_output(tz2_ortho_trajin + """
     rmsd first @CA perres range 2-7""")
     d = pt.rmsd_perres(traj,
                        ref=0,
                        mask='@CA',
                        resrange='2-7',
                        dtype='ndarray')
     aa_eq(cout[1:].values, d)

    def test_load_samples(self):
        traj = pt.load_sample_data()[:]
        assert isinstance(traj, pt.Trajectory) == True
        assert traj.top.n_atoms == 34
        assert traj.shape == (1, 34, 3)

        traj2 = pt.load_sample_data()
        assert isinstance(traj2, pt.TrajectoryIterator) == True
        assert traj2.top.n_atoms == 34
        assert traj2.shape == (1, 34, 3)
        aa_eq(traj.xyz, traj2.xyz)

    def test_watershell(self):
        traj = pt.iterload("data/tz2.truncoct.nc", "data/tz2.truncoct.parm7")
        state = pt.load_batch(traj, '''
        watershell :1-7
        ''')

        d0 = pt.watershell(traj, solute_mask=':1-7')
        state.run()
        aa_eq(d0.values, state.data[[1, 2]].values)

        # need to provide solute_mask
        self.assertRaises(ValueError, lambda: pt.watershell(traj))

    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)