예제 #1
0
def zcut_geom(geom, z_cut):
    """Split an ASE Object along a plane perpendicular to the vertical (last) dimension"""
    
    from ase import Atoms
    import numpy as np
    from numpy import r_, c_
    c_log = logger_setup(__name__)
    
    # Split top and bottom
    # It's a new cell, starting at origin, positions must be reset
    bottom = Atoms([atm for atm in geom if atm.position[-1] < z_cut])
    top = Atoms([atm for atm in geom if atm.position[-1] > z_cut])

    # Check that it makes sense: no empty objects
    if len(top) == 0 or len(bottom) == 0:
        c_log.error("Cutting plane defines an empty obj: n_top=%i, n_bottom=%i", len(top), len(bottom))
        raise ValueError
    
    # Set top layer positions and cell
    top.set_positions([x-r_[geom.cell[-1][:-1], z_cut]
                      for x in top.positions])
    top_cell = geom.get_cell()
    top_cell[-1] = geom.cell[-1] - [0, 0, z_cut]
    top.set_cell(top_cell)
#    top.wrap() # Should not be needed
    
    # Set bottom layer cell. Here positions should be fine already.
    bottom_cell = geom.get_cell()
    bottom_cell = np.array([v for v in [*geom.cell[:-1],
                                        r_[geom.cell[-1][:-1], z_cut]]
                           ])
    bottom.set_cell(bottom_cell)
    
    # Return the two objects
    return (top, bottom)
예제 #2
0
def expand_geom(geom, factor):
    """Expand a ASE Atoms structure isotropically by a given factor. Keeps relative coordinate fixed"""
    geom=geom.copy()
    c_log = logger_setup(__name__)
    if [sum(x) for x in geom.cell]==[0,0,0]:
        c_log.error("Supercell is not defined")
        raise ValueError
    
    tmp = geom.get_scaled_positions()
    geom.cell *= factor
    geom.set_scaled_positions(tmp)

    return geom
예제 #3
0
def vector_lin_stretch(vec, stretch, s_dir):
    """Stretch a vector by a given factor. Optionally along a direction.

    As input: the vector to strerch, the scaling factor, the direction of scaling.
    If direction is not give, the vector is multiplied by the factor (i.e. np_vect*factor)
    """
    c_log = logger_setup(__name__)
    
    norm = np.linalg.norm(s_dir)
    vec = np.array(vec)
    
    # Stretch of a in s dir: v_{stretched}=v_{orth-s}+a*(v.s/|s|)s/|s|
    if norm:
        c_log.info("Directional case: stretching along given direction ("+"%.5f"*len(s)+")", *s)
        s_dir = np.array(s_dir)/norm   # normalize direction, safer
        return vec + (stretch-1.)*dot(vec, s_dir)*s_dir
예제 #4
0
def geom_plane_cut(geom, n, p):
    """Return atoms in ASE structure (geom) above and below (tuple of Atoms object) the plane defined by the normal n and intercept p"""

    c_log = logger_setup(__name__)
    
    #-------------------------------------------------------------------------------
    # Divede the points 
    #-------------------------------------------------------------------------------
    p_up, p_down = [], []
    for i, pi in enumerate(geom.positions):
        if pi[-1] > plane_at_r(pi, n, p):
            p_up.append(geom[i])
        else:
            p_down.append(geom[i])
    p_up = Atoms(p_up)
    p_up.set_cell(geom.get_cell())
    p_down = Atoms(p_down)
    p_down.set_cell(geom.get_cell())
    c_log.debug("Up and down")
    c_log.debug(p_up)
    c_log.debug(p_down)

    return p_up, p_down
예제 #5
0
def pretty_columns(argv):
    """Adjust the width of data lines in given file or stdin.

    Comment lines are start with given comment character.
    Comments can be grouped at the top or left at original position.

    Return stringIO with output (Python func) or prints on stdout (bash script)."""

    #-------------------------------------------------------------------------------
    # Argument parser
    #-------------------------------------------------------------------------------
    parser = argparse.ArgumentParser(description=pretty_columns.__doc__,
                                     epilog="Silva 03-11-2019")
    # Positional arguments
    parser.add_argument('filename',
                        default=None,
                        type=str,
                        nargs="?",
                        help='filename (if blank use stdin);')
    # Optional args
    parser.add_argument(
        '-c',
        '--comment',
        dest='comment_c',
        default="#",
        metavar="char",
        help="""set comment character (default='#'). Use -1 to disable filter).
                                NOTE: must be #col file == #col comments.""")
    parser.add_argument('--split',
                        action='store_true',
                        dest='split_flg',
                        help='put comments at beginning of file;')
    parser.add_argument('--debug',
                        action='store_true',
                        dest='debug',
                        help='show debug information.')

    #-------------------------------------------------------------------------------
    # Initialize and check variables
    #-------------------------------------------------------------------------------
    args = parser.parse_args(argv)  # Process arguments

    # Set up logger and debug options
    c_log = logger_setup(__name__)
    c_log.setLevel(logging.INFO)
    debug_opt = []  # Pass debug flag down to other functions
    if args.debug:
        c_log.setLevel(logging.DEBUG)
        debug_opt = ['-d']
    c_log.debug(args)

    # Initialize input stream
    if args.filename is None:
        in_stream = sys.stdin
    else:
        in_stream = open(args.filename, 'r')

    if args.comment_c == "-1": args.comment_c = None

    #-------------------------------------------------------------------------------
    # Initialise output string stream
    #-------------------------------------------------------------------------------
    output = io.StringIO()

    #-------------------------------------------------------------------------------
    # Process stream
    #-------------------------------------------------------------------------------
    if args.split_flg:
        # Split comments and data
        c_log.debug("Splitting")
        comments, data = load_stream(in_stream, comment_char=args.comment_c)
        c_log.debug(comments)

        for l in comments + adjust_col_width(data):
            print(l, file=output)
    else:
        c_log.debug("Keep order")
        # Keep ordered
        c_num, data = load_stream(in_stream,
                                  comment_char=args.comment_c,
                                  split=False)
        c_log.debug(c_num)
        c_log.debug("Comments at " + "%i " * len(c_num), *c_num)
        # Adjust only data lines
        data_adj = adjust_col_width(
            [l for i, l in enumerate(data) if i not in c_num])

        di = 0
        for i, l in enumerate(data):
            if i in c_num:
                print(l, file=output)
            else:
                print(data_adj[di], file=output)
                di += 1

    #-------------------------------------------------------------------------------
    # Close input and return output string stream
    #-------------------------------------------------------------------------------
    if args.filename: in_stream.close()
    return output
예제 #6
0
def plane_cut_wrap(argv):
    """Cut a structure according to given plane

    Valid ASE input geometry from filename or stdin. Plane defined by normal vector n and intercept p.
    Returns a ASE atoms object with the atoms below (or above) the plane.
    If used as script prints an xyz file."""

    #-------------------------------------------------------------------------------
    # Argument parser
    #-------------------------------------------------------------------------------
    parser = argparse.ArgumentParser(description=plane_cut_wrap.__doc__)
    # Positional arguments
    parser.add_argument('filename',
                        default=None,
                        type=str, nargs="?",
                        help='file with initial structure in xzy format. If black use stdin;')
    # Optional args
    parser.add_argument('-n', '--norm',
                        dest='normal',
                        nargs=3, type=float, required=True,
                        help='normal to the plane.')
    parser.add_argument('-p', '--point',
                        dest='point',
                        nargs=3, type=float, required=True,
                        help='intercept of the plane.')
    parser.add_argument('--format',
                        dest='format', default="vasp",
                        help='set ASE-supported format for output.')
    parser.add_argument('-a',
                        action='store_true', dest='get_above',
                        help='get atoms above rather than below the plane.')
    parser.add_argument('--plot',
                        action='store_true', dest='plot_flg',
                        help='plot structure and cutting plane.')
    parser.add_argument('--debug',
                        action='store_true', dest='debug',
                        help='show debug informations.')

    #-------------------------------------------------------------------------------
    # Initialize and check variables
    #-------------------------------------------------------------------------------
    args = parser.parse_args(argv) # Process arguments

    # Set up logger and debug options
    c_log = logger_setup(__name__)
    c_log.setLevel(logging.INFO)
    debug_opt = [] # Pass debug flag down to other functions
    if args.debug:
        c_log.setLevel(logging.DEBUG)
        debug_opt = ['-d']
    c_log.debug(args)

    # Load data from the right source
    # FIXME: there is something broken here. Gets broken pipe.
    if args.filename is None:
        geom = ase.io.read(sys.stdin) #, format="xyz")
    else:
        geom = ase.io.read(args.filename) # , format="xyz")

    # Define the plane
    n = np.array(args.normal)
    n = n/np.linalg.norm(n) # Normalize the normal vector
    p = np.array(args.point)

    #-------------------------------------------------------------------------------
    # Divede the points 
    #-------------------------------------------------------------------------------
    p_up, p_down = geom_plane_cut(geom, n, p)

    xx = np.linspace(min(geom.positions[:,0]), max(geom.positions[:,0]), 5)
    yy = np.linspace(min(geom.positions[:,1]), max(geom.positions[:,1]), 5)
    xm, ym = np.meshgrid(xx, yy)
    zm = -((xm-p[0])*n[0]+(ym-p[1])*n[1])/n[2]+p[2]
    
    #-------------------------------------------------------------------------------
    # Plot the thing
    #-------------------------------------------------------------------------------
    if args.plot_flg:
        import matplotlib.pyplot as plt
        from mpl_toolkits.mplot3d import Axes3D
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        #ax.set_aspect('equal') # Square plot, do not distort angles.
         
        o = np.array([0, 0, 0])
        ax.quiver([0],[0],[0],
                  n[0],n[1],n[2])
        ax.scatter(p[0], p[1], p[2], c="red")


        ax.plot_surface(xm, ym, zm, alpha=0.7)
        
        ax.scatter(p_up.positions[:,0],
                   p_up.positions[:,1],
                   p_up.positions[:,2],
                   alpha=0.8, color="green")
        ax.scatter(p_down.positions[:,0],
                   p_down.positions[:,1],
                   p_down.positions[:,2],
                   alpha=0.8, color="purple")
        
#        ax.set_xlim3d(min(geom.positions[:,0]), max(geom.positions[:,0]))
#        ax.set_xlim3d(min(geom.positions[:,1]), max(geom.positions[:,1]))
#        ax.set_xlim3d(min(geom.positions[:,2]), max(geom.positions[:,2]))
        plt.show()

    # Return points above the plane or below
    if args.get_above:
        res = p_up
    else:
        res = p_down

    if __name__ == "__main__":
        ase.io.write('-', res, format=args.format)

    return res