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)
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
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
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
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
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