Python utils.aa_eq函数代码示例

    def test_drmsd(self):
        traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
        txt = '''
        parm data/tz2.parm7
        trajin data/tz2.nc
        drmsd drms_nofit out drmsd.dat
        rms rms_nofit out drmsd.dat nofit
        rms rms_fit out drmsd.dat
        drmsd drms_fit out drmsd.dat
        '''

        state = pt.load_cpptraj_state(txt)
        state.run()
        cpp_data = state.data[1:]

        # distance_rmsd
        data_drmsd = pt.distance_rmsd(traj)
        aa_eq(data_drmsd, cpp_data[0])
        aa_eq(pt.drmsd(traj), cpp_data[0])

        # rms_nofit
        aa_eq(cpp_data[1], pt.rmsd(traj, nofit=True))

        # rms_fit
        aa_eq(cpp_data[2], pt.rmsd(traj, nofit=False))

        # drmsd with rmsfit
        aa_eq(cpp_data[3], pt.distance_rmsd(traj(rmsfit=0), ref=traj[0]))

    def test_vs_list_tuple(self):
        traj = pt.iterload("./data/tz2.ortho.nc", "./data/tz2.ortho.parm7")
        traj0 = pt.replicate_cell(traj, direction='dir 001 dir 0-10')
        traj1 = pt.replicate_cell(traj, direction=('001', '0-10'))
        aa_eq(traj0.xyz, traj1.xyz)

        self.assertRaises(ValueError, lambda: pt.replicate_cell(traj, direction=traj[0]))

    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_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_0(self):
        import numpy as np
        from pytraj.datasetlist import stack as stack
        _ds1 = pt.calc_dssp(traj[:5], dtype='dataset')
        ds1 = _ds1.grep('LYS')
        _ds2 = pt.calc_dssp(traj[5:], dtype='dataset')
        ds2 = _ds2.grep('LYS')

        dstack = stack((ds1, ds2))
        _d12 = pt.calc_dssp(traj, dtype='dataset')
        d12 = _d12.grep("LYS")

        dstack_dict = dstack.to_dict()
        d12_dict = d12.to_dict()
        assert sorted(dstack_dict.keys()) == sorted(d12_dict)

        for key in dstack_dict.keys():
            arr0 = dstack_dict[key]
            arr1 = d12_dict[key]
            if np.any(arr0 == arr1) == False:
                pass

        arr1 = ds1.to_ndarray()
        arr2 = ds2.to_ndarray()
        arrstack = dstack.to_ndarray()
        arr12 = d12.to_ndarray()

        aa_eq(arrstack.flatten(), arr12.flatten())

    def test_crdframes(self):
        '''test crdframes in cpptraj
        '''
        max_frames = 50
        traj = pt.iterload('data/tz2.nc',
                           'data/tz2.parm7',
                           frame_slice=(0, max_frames, 2))

        state = pt.load_cpptraj_state('''
                parm {0}
                trajin {1} 1 {2} 2
                trajin {1} 1 {2} 2
                loadtraj name traj
                crdaction traj rms crdframes 1,10
                crdaction traj rms crdframes 1,30,2
                '''.format(traj.top.filename, traj.filename, max_frames))
        state.run()

        rmsd_0 = pt.rmsd(traj, ref=0, frame_indices=range(10))
        rmsd_crdframes = state.data[2].values
        aa_eq(rmsd_0, rmsd_crdframes)

        traj2 = traj.copy()
        assert traj2.n_frames == traj.n_frames, 'must have the same n_frames'
        traj2._load(traj.filename, frame_slice=(0, max_frames, 2))
        assert traj2.n_frames == 2 * traj.n_frames, 'n_frames must be doubled after reload'

        rmsd_1 = pt.rmsd(traj2, ref=0, frame_indices=range(0, 30, 2))
        rmsd_crdframes = state.data[3].values
        aa_eq(rmsd_1, rmsd_crdframes)

    def test_rmsfit_with_autoimage_compared_to_cpptraj(self):
        # assert to cpptraj: need to set mass
        traj = self.traj.copy()

        txt = '''
        parm {0}
        trajin {1}
        autoimage
        rms first {2} mass
        trajout tmp.nc
        '''.format(traj.top.filename, traj.filename, self.mask)

        with tempfolder():
            state = pt.datafiles.load_cpptraj_output(txt, dtype='state')
            state.run()
            # need to load to memory (not iterload)
            saved_traj = pt.load('tmp.nc', traj.top)

        fa1 = traj[:]
        fa1.autoimage()
        pt.superpose(fa1, ref=0, mask=self.mask, mass=True)

        aa_eq(fa1.xyz, saved_traj.xyz)

        fa_saved_nowat = saved_traj['!:WAT']
        fa1_nowat = fa1['!:WAT']
        aa_eq(fa_saved_nowat.xyz, fa1_nowat.xyz)

    def test_multiple_cores(self):
        from multiprocessing import Pool
        traj = pt.iterload('data/tz2.nc', 'data/tz2.parm7')
        for _ in range(10):
            traj._load(traj.filelist)
        saved_angle = pt.angle(traj, ':3 :10 :11')
        saved_dist = pt.distance(traj, ':3 :10')

        for n_cores in [2, 3]:
            lines = ['angle :3 :10 :11', 'distance :3 :10']
            pfuncs = partial(worker_state,
                             n_cores=n_cores,
                             traj=traj,
                             dtype='dict',
                             lines=lines)
            p = Pool(n_cores)
            data_list = p.map(pfuncs, [rank for rank in range(n_cores)])
            p.close()
            p.join()
            data_list_sorted_rank = (data[1]
                                     for data in sorted(data_list,
                                                        key=lambda x: x[0]))
            final_data = concat_dict(data_list_sorted_rank)
            aa_eq(final_data['Ang_00002'], saved_angle)
            aa_eq(final_data['Dis_00003'], saved_dist)

    def test_iterframe(self):
        '''test iterframe for both Trajectory and TrajectoryIterator
        '''

        orig_traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")
        # iterframe (already in doctest), just throwing raise to increase coverage score

        for traj in [orig_traj, orig_traj[:]]:
            self.assertRaises(ValueError,
                              lambda: traj.iterframe(rmsfit='crazy'))

            # rmsfit is an int
            t0 = orig_traj[:].rmsfit(ref=3)
            aa_eq(
                pt.rmsd_nofit(
                    traj(rmsfit=3),
                    ref=orig_traj[-1]),
                pt.rmsd_nofit(t0, ref=orig_traj[-1]))

        # test TypeError if not has n_frames info
        t0 = orig_traj[:]

        def int_gen(k):
            for i in range(k):
                yield i

        fi = pt.iterframe(t0, frame_indices=int_gen(3))
        aa_eq(pt.radgyr(fi, top=traj.top), pt.radgyr(orig_traj[:3]))

    def test_xyz(self):
        traj = pt.datafiles.load_tz2_ortho()
        t0 = traj[:]

        def set_xyz_not_c_contiguous():
            t0.xyz = np.asfortranarray(traj.xyz)

        def append_2d():
            traj1 = pt.load_sample_data('ala3')

        def set_xyz_not_same_n_atoms():
            traj1 = pt.load_sample_data('ala3')
            t0.xyz = traj1.xyz

        def append_2d():
            traj1 = pt.load_sample_data('ala3')
            t0.append_xyz(pt.tools.as_2darray(traj))

        # fortran order, autoconvert
        # make sure there is no TypeError
        set_xyz_not_c_contiguous()

        self.assertRaises(ValueError, lambda: set_xyz_not_same_n_atoms())
        self.assertRaises(ValueError, lambda: append_2d())

        # make sure to autoconvert from f4 to f8
        xyz_f4 = np.array(traj.xyz, dtype='f4')
        assert xyz_f4.itemsize == 4, 'must be f4'
        t0.xyz = xyz_f4
        aa_eq(t0.xyz, xyz_f4)
        assert t0.xyz.itemsize == 8, 'must be converted from f4 to f8'

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

        for n_cores in [2, 3]:
            for func in [matrix.dist, matrix.idea]:
                x = pt.pmap(func, traj, '@CA', n_cores=n_cores)
                aa_eq(x, func(traj, '@CA'))

    def test_DatasetMatrix3x3(self):
        # test _append_from_array
        mat0 = pt.calc_rotation_matrix(self.traj, ref=0)

        shape2d = (mat0.shape[0], mat0.shape[1] * mat0.shape[2])
        dmat3x3 = c_datasets.DatasetMatrix3x3()
        dmat3x3._append_from_array(mat0.reshape(shape2d))
        aa_eq(mat0, dmat3x3.values)

    def test_volume(self):
        traj = pt.iterload("data/tz2.ortho.nc", "data/tz2.ortho.parm7")
        state = pt.load_cpptraj_state('''
        volume''', traj)
        state.run()

        vol = pt.volume(traj)
        aa_eq(state.data[-1], vol)

    def test_0(self):
        traj = pt.iterload("./data/Tc5b.x", "./data/Tc5b.top")
        from pytraj.utils import convert as cv
        arange = cv.array_to_cpptraj_range(range(7))

        a0 = pt.multidihedral(traj, resrange='1-7').values
        a1 = pt.multidihedral(traj, resrange=range(7)).values
        aa_eq(a0.flatten(), a1.flatten())

    def test_issue_172_cpptraj(self):
        '''Internal Error: Attempting to assign to Frame with external memory
        '''
        traj = pt.iterload("./data/tz2.nc", "./data/tz2.parm7")

        arr_out_of_memory = pt.atomiccorr(traj(0, 8, 2), '@CA')
        arr_in_memory = pt.atomiccorr(traj[0: 8: 2], '@CA')
        aa_eq(arr_out_of_memory, arr_in_memory)