예제 #1
0
def test_extend_array():
    arr = rand(3,3)
    nrep = 5
    a0 = num.extend_array(arr, nrep, axis=0)
    a1 = num.extend_array(arr, nrep, axis=1)
    a2 = num.extend_array(arr, nrep, axis=2)    
    am1 = num.extend_array(arr, nrep, axis=-1)    
    assert a0.shape == (nrep,3,3)
    assert a1.shape == (3,nrep,3)
    assert a2.shape == (3,3,nrep)
    assert am1.shape == (3,3,nrep)
    equal(a2, am1)

    for axis, aa in enumerate([a0, a1, a2]):
        for ii in range(nrep):
            # slicetake(a0, 3, 0) -> a0[3,:,:]
            equal(arr, num.slicetake(aa, ii, axis=axis))
def test_cart2frac_traj():
    nstep = 100
    coords = rand(nstep,20,3)
    coords_copy = coords.copy()
    cell = rand(nstep,3,3)
    c1 = np.array([np.dot(coords[ii,...], np.linalg.inv(cell[ii,...])) for ii in \
                   range(nstep)])
    c2 = _flib.cart2frac_traj(coords, cell)
    c3 = crys.coord_trans3d(coords, new=cell, 
                            old=num.extend_array(np.identity(3), 
                                                 nstep=nstep, axis=0))
    assert (coords == coords_copy).all()
    assert np.allclose(c1, c2)
    assert np.allclose(c1, c3)
    assert c2.flags.f_contiguous
def test_frac2cart_traj():
    nstep = 100
    coords_frac = rand(nstep,20,3)
    coords_frac_copy = coords_frac.copy()
    cell = rand(nstep,3,3)
    c1 = np.array([np.dot(coords_frac[ii,...], cell[ii,...]) for ii in \
                   range(nstep)])
    c2 = _flib.frac2cart_traj(coords_frac, cell)
    c3 = crys.coord_trans3d(coords_frac, old=cell, 
                            new=num.extend_array(np.identity(3), 
                                                 nstep=nstep, axis=0))
    assert (coords_frac == coords_frac_copy).all()
    assert np.allclose(c1, c2)
    assert np.allclose(c1, c3)
    assert c2.flags.f_contiguous
def test_cart2frac_traj():
    nstep = 100
    coords = rand(nstep, 20, 3)
    coords_copy = coords.copy()
    cell = rand(nstep, 3, 3)
    c1 = np.array([np.dot(coords[ii,...], np.linalg.inv(cell[ii,...])) for ii in \
                   range(nstep)])
    c2 = _flib.cart2frac_traj(coords, cell)
    c3 = crys.coord_trans3d(coords,
                            new=cell,
                            old=num.extend_array(np.identity(3),
                                                 nstep=nstep,
                                                 axis=0))
    assert (coords == coords_copy).all()
    assert np.allclose(c1, c2)
    assert np.allclose(c1, c3)
    assert c2.flags.f_contiguous
def test_frac2cart_traj():
    nstep = 100
    coords_frac = rand(nstep, 20, 3)
    coords_frac_copy = coords_frac.copy()
    cell = rand(nstep, 3, 3)
    c1 = np.array([np.dot(coords_frac[ii,...], cell[ii,...]) for ii in \
                   range(nstep)])
    c2 = _flib.frac2cart_traj(coords_frac, cell)
    c3 = crys.coord_trans3d(coords_frac,
                            old=cell,
                            new=num.extend_array(np.identity(3),
                                                 nstep=nstep,
                                                 axis=0))
    assert (coords_frac == coords_frac_copy).all()
    assert np.allclose(c1, c2)
    assert np.allclose(c1, c3)
    assert c2.flags.f_contiguous
예제 #6
0
def test_dist_traj():
    natoms = 10
    nstep = 100
    cell = rand(nstep, 3, 3)
    stress = rand(nstep, 3, 3)
    forces = rand(nstep, natoms, 3)
    etot = rand(nstep)
    cryst_const = crys.cell2cc3d(cell, axis=0)
    coords_frac = np.random.rand(nstep, natoms, 3)
    coords = crys.coord_trans3d(coords=coords_frac,
                                old=cell,
                                new=num.extend_array(np.identity(3),
                                                     nstep,
                                                     axis=0),
                                axis=1,
                                timeaxis=0)
    assert cryst_const.shape == (nstep, 6)
    assert coords.shape == (nstep, natoms, 3)
    symbols = ['H'] * natoms

    traj = Trajectory(
        coords_frac=coords_frac,
        cell=cell,
        symbols=symbols,
        forces=forces,
        stress=stress,
        etot=etot,
        timestep=1,
    )

    for pbc in [True, False]:
        # (nstep, natoms, natoms, 3)
        distvecs_frac = traj.coords_frac[:,:,None,:] - \
                        traj.coords_frac[:,None,:,:]
        assert distvecs_frac.shape == (nstep, natoms, natoms, 3)
        if pbc:
            distvecs_frac = crys.min_image_convention(distvecs_frac)
        distvecs = np.empty((nstep, natoms, natoms, 3))
        for ii in range(traj.nstep):
            distvecs[ii, ...] = np.dot(distvecs_frac[ii, ...], traj.cell[ii,
                                                                         ...])
        # (nstep, natoms, natoms)
        dists = np.sqrt((distvecs**2.0).sum(axis=-1))
        assert np.allclose(dists, crys.distances_traj(traj, pbc=pbc))
예제 #7
0
def test_coords_trans():
    natoms = 10
    nstep = 100
    cell = rand(nstep, 3, 3)
    cryst_const = crys.cell2cc3d(cell, axis=0)
    coords_frac = rand(nstep, natoms, 3)
    coords = crys.coord_trans3d(coords=coords_frac,
                                old=cell,
                                new=num.extend_array(np.identity(3),
                                                     nstep,
                                                     axis=0),
                                axis=1,
                                timeaxis=0)

    traj = Trajectory(coords_frac=coords_frac, cell=cell)
    assert np.allclose(cryst_const, traj.cryst_const)
    assert np.allclose(coords, traj.coords)

    traj = Trajectory(coords=coords, cell=cell)
    assert np.allclose(coords_frac, traj.coords_frac)
예제 #8
0
def test_dist_traj():
    natoms = 10
    nstep = 100
    cell = rand(nstep,3,3)
    stress = rand(nstep,3,3)
    forces = rand(nstep,natoms,3)
    etot=rand(nstep)
    cryst_const = crys.cell2cc3d(cell, axis=0)
    coords_frac = np.random.rand(nstep,natoms,3)
    coords = crys.coord_trans3d(coords=coords_frac,
                                old=cell,
                                new=num.extend_array(np.identity(3),
                                                     nstep,axis=0),
                                axis=1,
                                timeaxis=0)                                                    
    assert cryst_const.shape == (nstep, 6)
    assert coords.shape == (nstep,natoms,3)
    symbols = ['H']*natoms
    
    traj = Trajectory(coords_frac=coords_frac,
                      cell=cell,
                      symbols=symbols,
                      forces=forces,
                      stress=stress,
                      etot=etot,
                      timestep=1,
                      )
    
    for pbc in [True, False]:
        # (nstep, natoms, natoms, 3)
        distvecs_frac = traj.coords_frac[:,:,None,:] - \
                        traj.coords_frac[:,None,:,:]
        assert distvecs_frac.shape == (nstep, natoms, natoms, 3)
        if pbc:
            distvecs_frac = crys.min_image_convention(distvecs_frac)
        distvecs = np.empty((nstep, natoms, natoms, 3))
        for ii in range(traj.nstep):
            distvecs[ii,...] = np.dot(distvecs_frac[ii,...], traj.cell[ii,...])
        # (nstep, natoms, natoms)
        dists = np.sqrt((distvecs**2.0).sum(axis=-1))
        assert np.allclose(dists, crys.distances_traj(traj, pbc=pbc))
예제 #9
0
def test_coords_trans():
    natoms = 10
    nstep = 100
    cell = rand(nstep,3,3)
    cryst_const = crys.cell2cc3d(cell, axis=0)
    coords_frac = rand(nstep,natoms,3)
    coords = crys.coord_trans3d(coords=coords_frac,
                                old=cell,
                                new=num.extend_array(np.identity(3),
                                                     nstep,axis=0),
                                axis=1,
                                timeaxis=0)                                                    
    
    traj = Trajectory(coords_frac=coords_frac,
                      cell=cell)
    assert np.allclose(cryst_const, traj.cryst_const)
    assert np.allclose(coords, traj.coords)
    
    traj = Trajectory(coords=coords,
                      cell=cell)
    assert np.allclose(coords_frac, traj.coords_frac)
예제 #10
0
def test_traj():
    natoms = 10
    nstep = 100
    cell = rand(nstep,3,3)
    stress = rand(nstep,3,3)
    forces = rand(nstep,natoms,3)
    etot = rand(nstep)
    cryst_const = crys.cell2cc3d(cell, axis=0)
    coords_frac = rand(nstep,natoms,3)
    coords = crys.coord_trans3d(coords=coords_frac,
                                old=cell,
                                new=num.extend_array(np.identity(3),
                                                     nstep,axis=0),
                                axis=1,
                                timeaxis=0)                                                    
    assert cryst_const.shape == (nstep, 6)
    assert coords.shape == (nstep,natoms,3)
    symbols = ['H']*natoms
    
    # automatically calculated:
    #   coords
    #   cell
    #   pressure
    #   velocity (from coords)
    #   temperature (from ekin)
    #   ekin (from velocity)
    traj = Trajectory(coords_frac=coords_frac,
                    cell=cell,
                    symbols=symbols,
                    forces=forces,
                    stress=stress,
                    etot=etot,
                    timestep=1,
                    )
    # Test if all getters work.
    for name in traj.attr_lst:
        print "test if getters work:", name
        traj.try_set_attr(name)
        assert getattr(traj, name) is not None, "attr None: %s" %name
        assert eval('traj.get_%s()'%name) is not None, "getter returns None: %s" %name
        print "test if getters work:", name, "... ok"
    aaae(coords_frac, traj.coords_frac)
    aaae(coords, traj.coords)
    aaae(cryst_const, traj.cryst_const)
    aaae(np.trace(stress, axis1=1, axis2=2)/3.0, traj.pressure)
    assert traj.coords.shape == (nstep,natoms,3)
    assert traj.cell.shape == (nstep,3,3)
    assert traj.velocity.shape == (nstep, natoms, 3)
    assert traj.temperature.shape == (nstep,)
    assert traj.ekin.shape == (nstep,)
    assert traj.nstep == nstep
    assert traj.natoms == natoms

    traj = Trajectory(coords_frac=coords_frac,
                    symbols=symbols,
                    cell=cell)
    aaae(coords, traj.coords)
    
    # Cell calculated from cryst_const has defined orientation in space which may be
    # different from the original `cell`, but the volume and underlying cryst_const
    # must be the same.
    traj = Trajectory(coords_frac=coords_frac,
                    symbols=symbols,
                    cryst_const=cryst_const)
    np.testing.assert_almost_equal(crys.volume_cell3d(cell),
                                   crys.volume_cell3d(traj.cell))
    aaae(cryst_const, crys.cell2cc3d(traj.cell))
    
    # extend arrays
    cell2d = rand(3,3)
    cc2d = crys.cell2cc(cell2d)
    traj = Trajectory(coords_frac=coords_frac,
                      cell=cell2d,
                      symbols=symbols)
    assert traj.cell.shape == (nstep,3,3)
    assert traj.cryst_const.shape == (nstep,6)
    for ii in range(traj.nstep):
        assert (traj.cell[ii,...] == cell2d).all()
        assert (traj.cryst_const[ii,:] == cc2d).all()
    
    traj = Trajectory(coords_frac=coords_frac,
                      cryst_const=cc2d,
                      symbols=symbols)
    assert traj.cell.shape == (nstep,3,3)
    assert traj.cryst_const.shape == (nstep,6)
    for ii in range(traj.nstep):
        assert (traj.cryst_const[ii,:] == cc2d).all()

    # units
    traj = Trajectory(coords_frac=coords_frac,
                    cell=cell,
                    symbols=symbols,
                    stress=stress,
                    forces=forces,
                    units={'length': 2, 'forces': 3, 'stress': 4})
    aaae(2*coords, traj.coords)                    
    aaae(3*forces, traj.forces)                    
    aaae(4*stress, traj.stress)                    
    
    # iterate, check if Structures are complete
    traj = Trajectory(coords=coords, 
                      symbols=symbols,
                      cell=cell,
                      forces=forces,
                      stress=stress,
                      etot=etot,
                      timestep=1.0)
    for struct in traj:
        assert struct.is_struct, "st is not Structure"
        assert not struct.is_traj, "st is Trajectory"
        assert_attrs_not_none(struct)
    struct = traj[0]
    for attr_name in traj.attr_lst:
        if attr_name in struct.attrs_only_traj:
            msg = "tr[0] %s is not None" %attr_name
            assert getattr(struct,attr_name) is None, msg
        else:            
            msg = "tr[0] %s is None" %attr_name
            assert getattr(struct,attr_name) is not None, msg
    
    # slices, return traj
    keys = traj.attr_lst[:]
    tsl = traj[10:80:2]
    assert tsl.nstep == traj.nstep / 2 - 15
    assert_attrs_not_none(tsl, attr_lst=keys)
    tsl = traj[slice(10,80,2)]
    assert tsl.nstep == traj.nstep / 2 - 15
    assert_attrs_not_none(tsl, attr_lst=keys)
    tsl = traj[np.s_[10:80:2]]
    assert tsl.nstep == traj.nstep / 2 - 15
    assert_attrs_not_none(tsl, attr_lst=keys)
    assert tsl.is_traj
    
    # iteration over sliced traj
    tsl = traj[10:80:2]
    for x in tsl:
        pass
    for x in tsl.copy():
        pass

    # repeat iter
    for i in range(2):
        cnt = 0
        for st in traj:
            cnt += 1
        assert cnt == nstep, "%i, %i" %(cnt, nstep)    
    
    # copy
    traj2 = traj.copy()
    for name in traj.attr_lst:
        val = getattr(traj,name)
        if val is not None and not (isinstance(val, types.IntType) or \
            isinstance(val, types.FloatType)):
            val2 = getattr(traj2,name)
            print "test copy:", name, type(val), type(val2)
            assert id(val2) != id(val)
            assert_all_types_equal(val2, val)
    assert_dict_with_all_types_equal(traj.__dict__, traj2.__dict__,
                                     keys=traj.attr_lst)
예제 #11
0
def test_scell():
    cell = np.identity(3)
    coords_frac = np.array([[0.5, 0.5, 0.5],
                       [1,1,1]])
    symbols = ['Al', 'N']
    sc = crys.scell(Structure(coords_frac=coords_frac,
                              cell=cell,
                              symbols=symbols), (2,2,2))

    sc_coords_frac = \
        np.array([[ 0.25,  0.25,  0.25],
                  [ 0.25,  0.25,  0.75],
                  [ 0.25,  0.75,  0.25],
                  [ 0.25,  0.75,  0.75],
                  [ 0.75,  0.25,  0.25],
                  [ 0.75,  0.25,  0.75],
                  [ 0.75,  0.75,  0.25],
                  [ 0.75,  0.75,  0.75],
                  [ 0.5 ,  0.5 ,  0.5 ],
                  [ 0.5 ,  0.5 ,  1.  ],
                  [ 0.5 ,  1.  ,  0.5 ],
                  [ 0.5 ,  1.  ,  1.  ],
                  [ 1.  ,  0.5 ,  0.5 ],
                  [ 1.  ,  0.5 ,  1.  ],
                  [ 1.  ,  1.  ,  0.5 ],
                  [ 1.  ,  1.  ,  1.  ]])

    sc_symbols = ['Al']*8 + ['N']*8
    sc_cell = \
        np.array([[ 2.,  0.,  0.],
                  [ 0.,  2.,  0.],
                  [ 0.,  0.,  2.]])

    assert sc.symbols == sc_symbols
    np.testing.assert_array_almost_equal(sc.coords_frac, sc_coords_frac)
    np.testing.assert_array_almost_equal(sc.cell, sc_cell)
    
    # non-orthorhombic cell
    cell = \
        np.array([[ 1.,  0.5,  0.5],
                  [ 0.25,  1.,  0.2],
                  [ 0.2,  0.5,  1.]])

    sc = crys.scell(Structure(coords_frac=coords_frac,
                              cell=cell,
                              symbols=symbols), (2,2,2))
    sc_cell = \
        np.array([[ 2. ,  1. ,  1. ],
                  [ 0.5,  2. ,  0.4],
                  [ 0.4,  1. ,  2. ]])
    np.testing.assert_array_almost_equal(sc.cell, sc_cell)
    # crystal coords are cell-independent
    np.testing.assert_array_almost_equal(sc.coords_frac, sc_coords_frac)

    
    # slab
    #
    # Test if old and new implementation behave for a tricky case: natoms == 2
    # mask.shape[0], i.e. if reshape() behaves correctly. 
    # Reference generated with old implementation. Default is new.
    cell = np.identity(3)
    coords_frac = np.array([[0.5, 0.5, 0.5],
                       [1,1,1]])
    symbols = ['Al', 'N']
    sc = crys.scell(Structure(coords_frac=coords_frac,
                              cell=cell,
                              symbols=symbols), (1,1,2))
    sc_coords_frac = \
        np.array([[ 0.5 ,  0.5 ,  0.25],
                  [ 0.5 ,  0.5 ,  0.75],
                  [ 1.  ,  1.  ,  0.5 ],
                  [ 1.  ,  1.  ,  1.  ]])
    sc_cell = \
        np.array([[ 1.,  0.,  0.],
                  [ 0.,  1.,  0.],
                  [ 0.,  0.,  2.]])
    sc_symbols = ['Al', 'Al', 'N', 'N']
    assert sc.symbols == sc_symbols
    np.testing.assert_array_almost_equal(sc.cell, sc_cell)
    np.testing.assert_array_almost_equal(sc.coords_frac, sc_coords_frac)
    
    sc = crys.scell(Structure(coords_frac=coords_frac,
                              cell=cell,
                              symbols=symbols), (1,2,1))
    sc_coords_frac = \
        np.array([[ 0.5 ,  0.25,  0.5 ],
                  [ 0.5 ,  0.75,  0.5 ],
                  [ 1.  ,  0.5 ,  1.  ],
                  [ 1.  ,  1.  ,  1.  ]])
    sc_cell = \
        np.array([[ 1.,  0.,  0.],
                  [ 0.,  2.,  0.],
                  [ 0.,  0.,  1.]])
    assert sc.symbols == sc_symbols
    np.testing.assert_array_almost_equal(sc.cell, sc_cell)
    np.testing.assert_array_almost_equal(sc.coords_frac, sc_coords_frac)

    sc = crys.scell(Structure(coords_frac=coords_frac,
                              cell=cell,
                              symbols=symbols), (2,1,1))
    sc_coords_frac = \
        np.array([[ 0.25,  0.5 ,  0.5 ],
                  [ 0.75,  0.5 ,  0.5 ],
                  [ 0.5 ,  1.  ,  1.  ],
                  [ 1.  ,  1.  ,  1.  ]])
    sc_cell = \
        np.array([[ 2.,  0.,  0.],
                  [ 0.,  1.,  0.],
                  [ 0.,  0.,  1.]])
    assert sc.symbols == sc_symbols
    np.testing.assert_array_almost_equal(sc.cell, sc_cell)
    np.testing.assert_array_almost_equal(sc.coords_frac, sc_coords_frac)
    
    # symbols = None
    sc = crys.scell(Structure(coords_frac=coords_frac,
                              cell=cell,
                              symbols=None), (2,2,2))
    assert sc.symbols is None
    
    # Trajectory
    natoms = 4
    nstep = 100
    symbols = [syms.next() for ii in range(natoms)]
    # cell 2d
    coords_frac = rand(nstep,natoms,3)
    cell = rand(3,3)
    dims = (2,3,4)
    nmask = np.prod(dims)
    sc = crys.scell(Trajectory(coords_frac=coords_frac,
                               cell=cell,
                               symbols=symbols), 
                      dims=dims)
    assert sc.coords_frac.shape == (nstep, nmask*natoms, 3)
    assert sc.symbols == common.flatten([[sym]*nmask for sym in symbols])
    assert sc.cell.shape == (nstep,3,3)                                            
    np.testing.assert_array_almost_equal(sc.cell, 
                                         num.extend_array(cell * np.asarray(dims)[:,None], 
                                                          sc.nstep,
                                                          axis=0))
    # cell 3d
    cell = rand(nstep,3,3)
    sc = crys.scell(Trajectory(coords_frac=coords_frac,
                               cell=cell,
                               symbols=symbols), 
                      dims=dims)
    assert sc.coords_frac.shape == (nstep, nmask*natoms, 3)
    coords_frac2 = np.array([crys.scell(Structure(coords_frac=coords_frac[ii,...,], 
                                                  cell=cell[ii,...],
                                                  symbols=symbols), dims=dims).coords_frac \
                             for ii in range(nstep)])
    np.testing.assert_array_almost_equal(sc.coords_frac, coords_frac2) 
    assert sc.symbols == common.flatten([[sym]*nmask for sym in symbols])
    assert sc.cell.shape == (nstep,3,3) 
    np.testing.assert_array_almost_equal(sc.cell, 
                                         cell * np.asarray(dims)[None,:,None])
    
    # methods
    natoms = 20
    coords_frac = rand(natoms,3)
    cell = rand(3,3)
    dims = (2,3,4)
    symbols = [syms.next() for ii in range(natoms)]
    struct = Structure(coords_frac=coords_frac,
                       cell=cell,
                       symbols=symbols)
    sc1 = crys.scell(struct, dims=dims, method=1)
    sc2 = crys.scell(struct, dims=dims, method=2)
    d1 = dict([(key, getattr(sc1, key)) for key in sc1.attr_lst])
    d2 = dict([(key, getattr(sc2, key)) for key in sc2.attr_lst])
    tools.assert_dict_with_all_types_almost_equal(d1, d2)
예제 #12
0
def test_traj():
    natoms = 10
    nstep = 100
    cell = rand(nstep, 3, 3)
    stress = rand(nstep, 3, 3)
    forces = rand(nstep, natoms, 3)
    etot = rand(nstep)
    cryst_const = crys.cell2cc3d(cell, axis=0)
    coords_frac = rand(nstep, natoms, 3)
    coords = crys.coord_trans3d(coords=coords_frac,
                                old=cell,
                                new=num.extend_array(np.identity(3),
                                                     nstep,
                                                     axis=0),
                                axis=1,
                                timeaxis=0)
    assert cryst_const.shape == (nstep, 6)
    assert coords.shape == (nstep, natoms, 3)
    symbols = ['H'] * natoms

    # automatically calculated:
    #   coords
    #   cell
    #   pressure
    #   velocity (from coords)
    #   temperature (from ekin)
    #   ekin (from velocity)
    traj = Trajectory(
        coords_frac=coords_frac,
        cell=cell,
        symbols=symbols,
        forces=forces,
        stress=stress,
        etot=etot,
        timestep=1,
    )
    # Test if all getters work.
    for name in traj.attr_lst:
        print("test if getters work:", name)
        traj.try_set_attr(name)
        assert getattr(traj, name) is not None, "attr None: %s" % name
        assert eval('traj.get_%s()' %
                    name) is not None, "getter returns None: %s" % name
        print("test if getters work:", name, "... ok")
    aaae(coords_frac, traj.coords_frac)
    aaae(coords, traj.coords)
    aaae(cryst_const, traj.cryst_const)
    aaae(np.trace(stress, axis1=1, axis2=2) / 3.0, traj.pressure)
    assert traj.coords.shape == (nstep, natoms, 3)
    assert traj.cell.shape == (nstep, 3, 3)
    assert traj.velocity.shape == (nstep, natoms, 3)
    assert traj.temperature.shape == (nstep, )
    assert traj.ekin.shape == (nstep, )
    assert traj.nstep == nstep
    assert traj.natoms == natoms

    traj = Trajectory(coords_frac=coords_frac, symbols=symbols, cell=cell)
    aaae(coords, traj.coords)

    # Cell calculated from cryst_const has defined orientation in space which may be
    # different from the original `cell`, but the volume and underlying cryst_const
    # must be the same.
    traj = Trajectory(coords_frac=coords_frac,
                      symbols=symbols,
                      cryst_const=cryst_const)
    np.testing.assert_almost_equal(crys.volume_cell3d(cell),
                                   crys.volume_cell3d(traj.cell))
    aaae(cryst_const, crys.cell2cc3d(traj.cell))

    # extend arrays
    cell2d = rand(3, 3)
    cc2d = crys.cell2cc(cell2d)
    traj = Trajectory(coords_frac=coords_frac, cell=cell2d, symbols=symbols)
    assert traj.cell.shape == (nstep, 3, 3)
    assert traj.cryst_const.shape == (nstep, 6)
    for ii in range(traj.nstep):
        assert (traj.cell[ii, ...] == cell2d).all()
        assert (traj.cryst_const[ii, :] == cc2d).all()

    traj = Trajectory(coords_frac=coords_frac,
                      cryst_const=cc2d,
                      symbols=symbols)
    assert traj.cell.shape == (nstep, 3, 3)
    assert traj.cryst_const.shape == (nstep, 6)
    for ii in range(traj.nstep):
        assert (traj.cryst_const[ii, :] == cc2d).all()

    # units
    traj = Trajectory(coords_frac=coords_frac,
                      cell=cell,
                      symbols=symbols,
                      stress=stress,
                      forces=forces,
                      units={
                          'length': 2,
                          'forces': 3,
                          'stress': 4
                      })
    aaae(2 * coords, traj.coords)
    aaae(3 * forces, traj.forces)
    aaae(4 * stress, traj.stress)

    # iterate, check if Structures are complete
    traj = Trajectory(coords=coords,
                      symbols=symbols,
                      cell=cell,
                      forces=forces,
                      stress=stress,
                      etot=etot,
                      timestep=1.0)
    for struct in traj:
        assert struct.is_struct, "st is not Structure"
        assert not struct.is_traj, "st is Trajectory"
        assert_attrs_not_none(struct)
    struct = traj[0]
    for attr_name in traj.attr_lst:
        if attr_name in struct.attrs_only_traj:
            msg = "tr[0] %s is not None" % attr_name
            assert getattr(struct, attr_name) is None, msg
        else:
            msg = "tr[0] %s is None" % attr_name
            assert getattr(struct, attr_name) is not None, msg

    # slices, return traj
    keys = traj.attr_lst[:]
    tsl = traj[10:80:2]
    assert tsl.nstep == traj.nstep / 2 - 15
    assert_attrs_not_none(tsl, attr_lst=keys)
    tsl = traj[slice(10, 80, 2)]
    assert tsl.nstep == traj.nstep / 2 - 15
    assert_attrs_not_none(tsl, attr_lst=keys)
    tsl = traj[np.s_[10:80:2]]
    assert tsl.nstep == traj.nstep / 2 - 15
    assert_attrs_not_none(tsl, attr_lst=keys)
    assert tsl.is_traj

    # iteration over sliced traj
    tsl = traj[10:80:2]
    for x in tsl:
        pass
    for x in tsl.copy():
        pass

    # repeat iter
    for i in range(2):
        cnt = 0
        for st in traj:
            cnt += 1
        assert cnt == nstep, "%i, %i" % (cnt, nstep)

    # copy
    traj2 = traj.copy()
    for name in traj.attr_lst:
        val = getattr(traj, name)
        if val is not None and not (isinstance(val, int) or \
            isinstance(val, float)):
            val2 = getattr(traj2, name)
            print("test copy:", name, type(val), type(val2))
            assert id(val2) != id(val)
            assert_all_types_equal(val2, val)
    assert_dict_with_all_types_equal(traj.__dict__,
                                     traj2.__dict__,
                                     keys=traj.attr_lst)
예제 #13
0
def test_scell():
    cell = np.identity(3)
    coords_frac = np.array([[0.5, 0.5, 0.5], [1, 1, 1]])
    symbols = ['Al', 'N']
    sc = crys.scell(
        Structure(coords_frac=coords_frac, cell=cell, symbols=symbols),
        (2, 2, 2))

    sc_coords_frac = \
        np.array([[ 0.25,  0.25,  0.25],
                  [ 0.25,  0.25,  0.75],
                  [ 0.25,  0.75,  0.25],
                  [ 0.25,  0.75,  0.75],
                  [ 0.75,  0.25,  0.25],
                  [ 0.75,  0.25,  0.75],
                  [ 0.75,  0.75,  0.25],
                  [ 0.75,  0.75,  0.75],
                  [ 0.5 ,  0.5 ,  0.5 ],
                  [ 0.5 ,  0.5 ,  1.  ],
                  [ 0.5 ,  1.  ,  0.5 ],
                  [ 0.5 ,  1.  ,  1.  ],
                  [ 1.  ,  0.5 ,  0.5 ],
                  [ 1.  ,  0.5 ,  1.  ],
                  [ 1.  ,  1.  ,  0.5 ],
                  [ 1.  ,  1.  ,  1.  ]])

    sc_symbols = ['Al'] * 8 + ['N'] * 8
    sc_cell = \
        np.array([[ 2.,  0.,  0.],
                  [ 0.,  2.,  0.],
                  [ 0.,  0.,  2.]])

    assert sc.symbols == sc_symbols
    np.testing.assert_array_almost_equal(sc.coords_frac, sc_coords_frac)
    np.testing.assert_array_almost_equal(sc.cell, sc_cell)

    # non-orthorhombic cell
    cell = \
        np.array([[ 1.,  0.5,  0.5],
                  [ 0.25,  1.,  0.2],
                  [ 0.2,  0.5,  1.]])

    sc = crys.scell(
        Structure(coords_frac=coords_frac, cell=cell, symbols=symbols),
        (2, 2, 2))
    sc_cell = \
        np.array([[ 2. ,  1. ,  1. ],
                  [ 0.5,  2. ,  0.4],
                  [ 0.4,  1. ,  2. ]])
    np.testing.assert_array_almost_equal(sc.cell, sc_cell)
    # crystal coords are cell-independent
    np.testing.assert_array_almost_equal(sc.coords_frac, sc_coords_frac)

    # slab
    #
    # Test if old and new implementation behave for a tricky case: natoms == 2
    # mask.shape[0], i.e. if reshape() behaves correctly.
    # Reference generated with old implementation. Default is new.
    cell = np.identity(3)
    coords_frac = np.array([[0.5, 0.5, 0.5], [1, 1, 1]])
    symbols = ['Al', 'N']
    sc = crys.scell(
        Structure(coords_frac=coords_frac, cell=cell, symbols=symbols),
        (1, 1, 2))
    sc_coords_frac = \
        np.array([[ 0.5 ,  0.5 ,  0.25],
                  [ 0.5 ,  0.5 ,  0.75],
                  [ 1.  ,  1.  ,  0.5 ],
                  [ 1.  ,  1.  ,  1.  ]])
    sc_cell = \
        np.array([[ 1.,  0.,  0.],
                  [ 0.,  1.,  0.],
                  [ 0.,  0.,  2.]])
    sc_symbols = ['Al', 'Al', 'N', 'N']
    assert sc.symbols == sc_symbols
    np.testing.assert_array_almost_equal(sc.cell, sc_cell)
    np.testing.assert_array_almost_equal(sc.coords_frac, sc_coords_frac)

    sc = crys.scell(
        Structure(coords_frac=coords_frac, cell=cell, symbols=symbols),
        (1, 2, 1))
    sc_coords_frac = \
        np.array([[ 0.5 ,  0.25,  0.5 ],
                  [ 0.5 ,  0.75,  0.5 ],
                  [ 1.  ,  0.5 ,  1.  ],
                  [ 1.  ,  1.  ,  1.  ]])
    sc_cell = \
        np.array([[ 1.,  0.,  0.],
                  [ 0.,  2.,  0.],
                  [ 0.,  0.,  1.]])
    assert sc.symbols == sc_symbols
    np.testing.assert_array_almost_equal(sc.cell, sc_cell)
    np.testing.assert_array_almost_equal(sc.coords_frac, sc_coords_frac)

    sc = crys.scell(
        Structure(coords_frac=coords_frac, cell=cell, symbols=symbols),
        (2, 1, 1))
    sc_coords_frac = \
        np.array([[ 0.25,  0.5 ,  0.5 ],
                  [ 0.75,  0.5 ,  0.5 ],
                  [ 0.5 ,  1.  ,  1.  ],
                  [ 1.  ,  1.  ,  1.  ]])
    sc_cell = \
        np.array([[ 2.,  0.,  0.],
                  [ 0.,  1.,  0.],
                  [ 0.,  0.,  1.]])
    assert sc.symbols == sc_symbols
    np.testing.assert_array_almost_equal(sc.cell, sc_cell)
    np.testing.assert_array_almost_equal(sc.coords_frac, sc_coords_frac)

    # symbols = None
    sc = crys.scell(
        Structure(coords_frac=coords_frac, cell=cell, symbols=None), (2, 2, 2))
    assert sc.symbols is None

    # Trajectory
    natoms = 4
    nstep = 100
    symbols = [next(syms) for ii in range(natoms)]
    # cell 2d
    coords_frac = rand(nstep, natoms, 3)
    cell = rand(3, 3)
    dims = (2, 3, 4)
    nmask = np.prod(dims)
    sc = crys.scell(Trajectory(coords_frac=coords_frac,
                               cell=cell,
                               symbols=symbols),
                    dims=dims)
    assert sc.coords_frac.shape == (nstep, nmask * natoms, 3)
    assert sc.symbols == common.flatten([[sym] * nmask for sym in symbols])
    assert sc.cell.shape == (nstep, 3, 3)
    np.testing.assert_array_almost_equal(
        sc.cell,
        num.extend_array(cell * np.asarray(dims)[:, None], sc.nstep, axis=0))
    # cell 3d
    cell = rand(nstep, 3, 3)
    sc = crys.scell(Trajectory(coords_frac=coords_frac,
                               cell=cell,
                               symbols=symbols),
                    dims=dims)
    assert sc.coords_frac.shape == (nstep, nmask * natoms, 3)
    coords_frac2 = np.array([crys.scell(Structure(coords_frac=coords_frac[ii,...,],
                                                  cell=cell[ii,...],
                                                  symbols=symbols), dims=dims).coords_frac \
                             for ii in range(nstep)])
    np.testing.assert_array_almost_equal(sc.coords_frac, coords_frac2)
    assert sc.symbols == common.flatten([[sym] * nmask for sym in symbols])
    assert sc.cell.shape == (nstep, 3, 3)
    np.testing.assert_array_almost_equal(
        sc.cell,
        cell * np.asarray(dims)[None, :, None])

    # methods
    natoms = 20
    coords_frac = rand(natoms, 3)
    cell = rand(3, 3)
    dims = (2, 3, 4)
    symbols = [next(syms) for ii in range(natoms)]
    struct = Structure(coords_frac=coords_frac, cell=cell, symbols=symbols)
    sc1 = crys.scell(struct, dims=dims, method=1)
    sc2 = crys.scell(struct, dims=dims, method=2)
    d1 = dict([(key, getattr(sc1, key)) for key in sc1.attr_lst])
    d2 = dict([(key, getattr(sc2, key)) for key in sc2.attr_lst])
    tools.assert_dict_with_all_types_almost_equal(d1, d2)