def sortStructure(self, sort = "type", reset_idx = False, verbose = 1):
"""Function for sorting the structure
sort = str("type"/"index"), Sort by type-z-y-x or sort by index as loaded
from simulation file.
reset_idx = bool, If True then reset the index of the structure
verbose = int, Print extra information
"""
if sort.lower() == "type":
"""Sorts the structure based on type_i then z then y then x"""
si = np.lexsort((self.pos[:, 0], self.pos[:, 1], self.pos[:, 2], self.type_i))
elif sort.lower() == "index":
"""Sort by the idx property, (index as written by the simulation programs)"""
si = np.argsort(self.idx)
if verbose > 0:
string = "Sorting structure by %s" % sort.lower()
ut.infoPrint(string)
"""Sort as specified"""
self.pos = self.pos[si]
self.type_i = self.type_i[si]
self.type_n = self.type_n[si]
self.frozen = self.frozen[si]
self.mass = self.mass[si]
if reset_idx:
self.resetIndex(verbose = verbose)
else:
self.idx = self.idx[si]
def writeStructure(self, filename = None, format = "lammps",\
direct = False, sd = False, verbose = 1):
"""Function for writing the structure to specified file format
filename = str(), Filename to write to
format = str("lammps"/"vasp"/"eon"/"xyz"), Format to write to
direct = bool, Write in direct coordinates (VASP)
sd = bool, consider selective dynamics (VASP) and frozen atoms EON
verbose = int, Print extra information
"""
if filename is None:
if self.filename is None:
filename = "Structure.%s" % format
else:
filename = self.filename
"""Write the structure object to specified file"""
file_io.writeData(filename = filename, atoms = self, format = format,\
sd = sd, direct = direct, verbose = verbose - 1)
if verbose > 0:
string = "Structure written to file: %s (%s-format)" % (filename, format)
ut.infoPrint(string)
def resetIndex(self, verbose = 1):
"""Function for reseting the atomic indicies
verbose = int, Print extra information
"""
self.idx = np.arange(self.pos.shape[0])
if verbose > 0:
string = "Reseting atom index"
ut.infoPrint(string)
def getNearestNeighborCollection(self, idx = None, idx_to = None, NN = 8,\
verbose = 1, limit = np.array([5, 5, 5]),\
extend = np.array([1, 1, 1], dtype = bool)):
"""Function for getting nearest neighbors around specified atoms to
specified atoms collected as an average with a standard deviation
idx = int, [int,], Index from which to calculate nearest neighbors
NN = int, Number of nearest neighbors to keep
idx_to = int, [int,], Calculate the NN considering only these atoms
limit = np.ndarray([float, float, float]), Limit around the maximum extent
of the included atoms, to speed up calculations
extend = np.ndarray([1/0, 1/0, 1/0]), Extend the cell if needed in
specified x, y, z directions
verbose = int, Print extra information
"""
"""Check some defaults"""
if idx is None: idx = np.arange(self.pos.shape[0])
if isinstance(idx, (int, np.integer)): idx = np.array([idx])
if idx_to is None: idx_to = np.arange(self.pos.shape[0])
if isinstance(idx_to, (int, np.integer)): idx_to = np.array([idx_to])
"""Change to cartesian coordinates"""
self.dir2car()
distance = self.getNearestNeighbors(idx = idx, idx_to = idx_to, NN = NN,\
verbose = verbose, limit = limit,\
extend = extend)
"""Return if less than NN neighbors could be found"""
if distance is None:
return
if verbose > 0:
string = "Calculated %i nearest neighbors for %i atoms" % (NN, distance.shape[0])
ut.infoPrint(string)
dist_mean = np.mean(distance, axis = 0)
dist_std = np.std(distance, axis = 0)
return dist_mean, dist_std
def getAtoms(self, mode = "layer_z", opt = [0, 0.2], coordinates = "c"):
"""get index of atoms that match supplied criteria
Mode is supplied as a dict with keyword specified as below
and containing a list of
Mode and options
----------------
layer_x - [x_center, dx], slice out atoms within x_center +- dx
layer_y - [y_center, dy], slice out atoms within y_center +- dy
layer_z - [z_center, dz], slice out atoms within z_center +- dz
"""
if "c".startswith(coordinates.lower()):
self.dir2car()
string = "Changed to cartesian coordinates"
ut.infoPrint(string)
elif "d".startswith(coordinates.lower()):
self.car2dir()
string = "Changed to direct coordinates"
ut.infoPrint(string)
else:
string = "Unrecognized option coordinates = %s, (can be c or d)" % coordinates
ut.infoPrint(string)
return
if mode.lower() == "layer_x":
mask = (self.pos[:, 0] >= (opt[0] - opt[1])) *\
(self.pos[:, 0] <= (opt[0] + opt[1]))
index = np.arange(self.pos.shape[0])[mask]
elif mode.lower() == "layer_y":
mask = (self.pos[:, 1] >= (opt[0] - opt[1])) *\
(self.pos[:, 1] <= (opt[0] + opt[1]))
index = np.arange(self.pos.shape[0])[mask]
elif mode.lower() == "layer_z":
mask = (self.pos[:, 2] >= (opt[0] - opt[1])) *\
(self.pos[:, 2] <= (opt[0] + opt[1]))
index = np.arange(self.pos.shape[0])[mask]
return index
def plotRDF(self, idx = None, idx_to = None, r = 6, dr = 0.1, bins = None,\
extend = np.array([1, 1, 1], dtype = bool), cumulative = False, legend = None,\
row = 1, col = 1, N = 1, handle = False, save = False, format = "pdf",\
dpi = 100, verbose = 1, **kwargs):
"""Function for ploting the RDF and cumulative distribution
idx = int, [int,], Index from which to calculate nearest neighbors
idx_to = int, [int,], Calculate the NN considering only these atoms
r = float, Cut off used in the RDF
dr = float, bins size used in the RDF
bins = int, Alternative to dr, specify the total number of bins
extend = np.ndarray([1/0, 1/0, 1/0]), Allow the cell to be extended in
the specified x, y, z directions
edges = bool, Include both edges
cumulative = bool, Add the cumulative RDF value to a right y-axis
legend = [str,] legend inputs
handle = bool, If True only prepare the axis don't draw the plot
row = int, Rows if used in subplots
col = int, Columns if used in subplots
N = int, Nr of plot if used in subplots
save = bool or str, Name to save the file to or save to default name
if True
format = valid matplotlib format, Format to save the plot in
dpi = int, DPI used when saving the plot
kwargs = valid matplotlib plot kwargs
"""
lbl_1 = None
if idx is None:
idx = [np.arange(self.pos.shape[0])]
elif isinstance(idx, (int, np.integer)):
idx = [np.array([idx])]
elif isinstance(idx[0], (int, np.integer)):
idx = [idx]
elif isinstance(idx, str) and idx.lower() == "species":
idx, lbl_1 = self.getElementIdx()[:2]
lbl_2 = None
if idx_to is None:
idx_to = [np.arange(self.pos.shape[0])]
elif isinstance(idx_to, (int, np.integer)):
idx_to = [np.array([idx_to])]
elif isinstance(idx_to[0], (int, np.integer)):
idx_to = [idx_to]
elif isinstance(idx_to, str) and idx_to.lower() == "species":
idx_to, lbl_2 = self.getElementIdx()[:2]
if len(idx) == 1:
l_idx = np.zeros(len(idx_to), dtype = np.int)
else:
l_idx = np.arange(len(idx), dtype = np.int)
if len(idx_to) == 1:
l_idx_to = np.zeros(len(idx), dtype = np.int)
else:
l_idx_to = np.arange(len(idx_to), dtype = np.int)
if l_idx.shape[0] != l_idx_to.shape[0]:
string = "Length of idx and idx_to does not match (%i, %i). "\
"Can be (1,N), (N,1) or (N,N)"\
% (l_idx.shape[0], l_idx_to.shape[0])
ut.infoPrint(string)
return
y = []; yt = []
for i in range(l_idx.shape[0]):
cnt, bin, tot = self.getRDF(idx = idx[l_idx[i]], idx_to = idx_to[l_idx_to[i]], r = r,\
dr = 0.1, bins = bins, extend = extend,\
edges = False, verbose = verbose)
y.append(cnt)
yt.append(tot)
if not handle:
hFig = plt.figure()
hAx = plt.subplot(row, col, N)
if cumulative:
hAxR = hAx.twinx()
label = "_ignore"
for i, item in enumerate(y):
if legend is not None:
if legend.lower() == "idx":
label = "%i -> %i" % (l_idx[i], l_idx_to[i])
else:
label = legend[i]
elif lbl_1 is not None and lbl_2 is not None:
label = "%2s -> %2s" % (lbl_1[i], lbl_2[i])
elif lbl_1 is not None:
label = "%2s -> %i" % (lbl_1[i], l_idx_to[i])
elif lbl_2 is not None:
label = "%i -> %2s" % (l_idx[i], lbl_2[i])
hL = hAx.plot(bin, item, linestyle = "-", label = label, **kwargs)
if cumulative:
hAxR.plot(bin, yt[i], linestyle = "--", color = hL[-1].get_color(), **kwargs)
hAx.set_xlabel("Radius, $\AA$")
hAx.set_ylabel("Atoms / (Atom * Volume), ($1/\AA^3$)")
if label != "_ignore":
hAx.legend(framealpha = 1, loc = "upper left")
hAx.set_title("RDF")
plt.tight_layout()
if save:
if save is True:
ut.save_fig(filename = "RDF.%s" % (format), format = format,\
dpi = dpi, verbose = verbose)
else:
ut.save_fig(filename = save, format = format, dpi = dpi,\
verbose = verbose)
plt.close()
else:
plt.show()
def plotNNC(self, idx, idx_to = None, NN = 8, verbose = True,\
handle = False, row = 1, col = 1, N = 1, save = False,\
format = "pdf", dpi = 100, legend = None, **kwargs):
"""Function for plotting a nearest neighbor collection with std
idx = int, [int,], Index from which to calculate nearest neighbors
NN = int, Number of nearest neighbors to keep
idx_to = int, [int,], Calculate the NN considering only these atoms
handle = bool, If True only prepare the axis don't draw the plot
row = int, Rows if used in subplots
col = int, Columns if used in subplots
N = int, Nr of plot if used in subplots
save = bool or str, Name to save the file to or save to default name
if True
format = valid matplotlib format, Format to save the plot in
dpi = int, DPI used when saving the plot
legend = [str,], Legend for the different entries
**kvargs = valid matplotlib errorbar kwargs
"""
lbl_1 = None
if idx is None:
idx = [np.arange(self.pos.shape[0])]
elif isinstance(idx, (int, np.integer)):
idx = [np.array([idx])]
elif isinstance(idx[0], (int, np.integer)):
idx = [idx]
elif isinstance(idx, str) and idx.lower() == "species":
idx, lbl_1 = self.getElementIdx()[:2]
lbl_2 = None
if idx_to is None:
idx_to = [np.arange(self.pos.shape[0])]
elif isinstance(idx_to, (int, np.integer)):
idx_to = [np.array([idx_to])]
elif isinstance(idx_to[0], (int, np.integer)):
idx_to = [idx_to]
elif isinstance(idx_to, str) and idx_to.lower() == "species":
idx_to, lbl_2 = self.getElementIdx()[:2]
if len(idx) == 1:
l_idx = np.zeros(len(idx_to), dtype = np.int)
else:
l_idx = np.arange(len(idx), dtype = np.int)
if len(idx_to) == 1:
l_idx_to = np.zeros(len(idx), dtype = np.int)
else:
l_idx_to = np.arange(len(idx_to), dtype = np.int)
if l_idx.shape[0] != l_idx_to.shape[0]:
string = "Length of idx and idx_to does not match (%i, %i). "\
"Can be (1,N), (N,1) or (N,N)"\
% (l_idx.shape[0], l_idx_to.shape[0])
ut.infoPrint(string)
return
x = []; y = []; s = []
for i in range(l_idx.shape[0]):
d_mean, d_std = self.getNearestNeighborCollection(idx = idx[l_idx[i]],\
idx_to = idx_to[l_idx_to[i]],\
NN = NN, verbose = verbose)
x.append(np.arange(1, d_mean.shape[0] + 1))
y.append(d_mean)
s.append(d_std)
if not handle:
hFig = plt.figure()
hAx = plt.subplot(row, col, N)
label = "_ignore"
for i, item in enumerate(y):
if legend is not None:
if legend.lower() == "idx":
label = "%i -> %i" % (l_idx[i], l_idx_to[i])
else:
label = legend[i]
elif lbl_1 is not None and lbl_2 is not None:
label = "%2s -> %2s" % (lbl_1[i], lbl_2[i])
elif lbl_1 is not None:
label = "%2s -> %i" % (lbl_1[i], l_idx_to[i])
elif lbl_2 is not None:
label = "%i -> %2s" % (l_idx[i], lbl_2[i])
hAx.errorbar(x[i], y[i], yerr = s[i], linestyle = ls, marker = m, capsize = cs,\
elinewidth = elw, markersize = ms, linewidth = lw, label = label, **kwargs)
hAx.set_xlabel("Neighbor")
hAx.set_ylabel("Distance, $(\AA)$")
if label != "_ignore":
hAx.legend(framealpha = 1, loc = "upper left")
hAx.set_title("Nearest Neighbor Distances")
plt.tight_layout()
if save:
if save is True:
ut.save_fig(filename = "NNC.%s" % (format), format = format,\
dpi = dpi, verbose = verbose)
else:
ut.save_fig(filename = save, format = format, dpi = dpi,\
verbose = verbose)
plt.close()
else:
plt.show()
def getNearestNeighbors(self, idx = None, idx_to = None, NN = 8,\
verbose = 1, limit = np.array([5, 5, 5]),\
extend = np.array([1, 1, 1], dtype = bool)):
"""Function for getting index and distance to nearest neighbors
of specified atoms
idx = int, [int,], Index from which to calculate nearest neighbors
NN = int, Number of nearest neighbors to keep
idx_to = int, [int,], Calculate the NN considering only these atoms
limit = np.ndarray([float, float, float]), Limit around the maximum extent
of the included atoms, to speed up calculations
verbose = int, Print extra information
"""
"""Check some defaults"""
if idx is None: idx = np.arange(self.pos.shape[0])
if isinstance(idx, (int, np.integer)): idx = np.array([idx])
if idx_to is None: idx_to = np.arange(self.pos.shape[0])
if isinstance(idx_to, (int, np.integer)): idx_to = np.array([idx_to])
"""Change to cartesian coordinates"""
self.dir2car()
"""Do a rough check which atoms must be included in the NN search"""
max_pos = np.max(self.pos[idx, :], axis = 0) + limit
min_pos = np.min(self.pos[idx, :], axis = 0) - limit
lim = np.all(self.pos < max_pos, axis = 1) *\
np.all(self.pos > min_pos, axis = 1)
idx_to = np.intersect1d(self.idx[lim], idx_to)
if verbose > 0:
string = "Considering idx_to within [%.2f, %.2f] (x), [%.2f, "\
" %.2f] (y), [%.2f, %.2f] (z)" % (min_pos[0], max_pos[0],\
min_pos[1], max_pos[1], min_pos[2], max_pos[2])
ut.infoPrint(string)
"""Cell extension to comply with the sepecified limit"""
box = self.getBoxLengths()
rep = np.ceil(limit / box) - 1
rep = rep.astype(np.int)
rep[np.logical_not(extend)] = 0
if np.any(rep > 0):
if verbose > 0:
string = "Replicating cell by %i, %i, %i (x, y, z)"\
% (rep[0] + 1, rep[1] + 1, rep[2] + 1)
ut.infoPrint(string)
"""Extend teh cell"""
pos_to, cell = self.getExtendedPositions(x = rep[0], y = rep[1], z = rep[2],\
idx = idx_to, return_cart = True,\
verbose = verbose - 1)
"""Change to cartesian coordinates"""
self.dir2car()
"""Change back to direct coordinates using the new extended cell"""
pos_to = np.matmul(np.linalg.inv(cell), pos_to.T)
pos_from = np.matmul(np.linalg.inv(cell), self.pos.T)
else:
"""Change to direct coordinates"""
self.car2dir()
if verbose > 0:
string = "Cell is only wrapped, not extended"
ut.infoPrint(string)
pos_to = self.pos.copy()[idx_to, :].T
pos_from = self.pos.copy().T
cell = self.cell
if pos_to.shape[1] - 1 < NN:
string = "Within current limits (%.2f, %.2f, %.2f) fewer NN (%i) "\
"are present than specified (%i)" % (limit[0], limit[1],\
limit[2], pos_to.shape[1] - 1, NN)
ut.infoPrint(string)
return
distance = np.zeros((np.shape(idx)[0], NN))
"""Measure distances between all specified atoms, wrap cell"""
for i, item in enumerate(idx):
d = pos_to - pos_from[:, [item]]
d[d > 0.5] -= 1
d[d < -0.5] += 1
"""Convert to cartesian coordinates"""
c = np.matmul(cell, d)
"""Calculate distances"""
dist = np.sqrt(c[0, :]**2 + c[1, :]**2 + c[2, :]**2)
"""Remove distance to the same atom"""
mask = dist > 0
dist = dist[mask]
"""Sort distances"""
si = np.argsort(dist)
"""Pick out the NN nearest in all variables"""
distance[i, :] = dist[si][:NN]
return distance
def getNeighborDistance(self, idx = None, r = 6, idx_to = None,\
extend = np.array([1, 1, 1], dtype = bool),\
verbose = 1):
"""Function for getting the distance between specified atoms within
radius r
idx = int, [int,], Index from which to calculate nearest neighbors
r = float, Radius or NN calculation
idx_to = int, [int,], Calculate the NN considering only these atoms
extend = np.ndarray([1/0, 1/0, 1/0]), Extend the cell if needed in
specified x, y, z directions
verbose = int, Print extra information
"""
"""Check some defaults"""
if idx is None: idx = np.arange(self.pos.shape[0])
if isinstance(idx, (int, np.integer)): idx = np.array([idx])
if idx_to is None: idx_to = np.arange(self.pos.shape[0])
if isinstance(idx_to, (int, np.integer)): idx_to = np.array([idx_to])
extend = np.array(extend, dtype = bool)
"""Change to cartesian coordinates"""
self.dir2car()
"""Check the rough maximum possible extent for relevant atoms"""
max_pos = np.max(self.pos[idx, :], axis = 0) + r
min_pos = np.min(self.pos[idx, :], axis = 0) - r
lim = np.all(self.pos < max_pos, axis = 1) *\
np.all(self.pos > min_pos, axis = 1)
idx_to = np.intersect1d(self.idx[lim], idx_to)
if verbose > 0:
string = "Considering idx_to within [%.2f, %.2f] (x), [%.2f, "\
" %.2f] (y), [%.2f, %.2f] (z)" % (min_pos[0], max_pos[0],\
min_pos[1], max_pos[1], min_pos[2], max_pos[2])
ut.infoPrint(string)
"""Cell extension to comply with the sepecified r value"""
box = self.getBoxLengths()
rep = np.ceil(r / box) - 1
rep = rep.astype(np.int)
rep[np.logical_not(extend)] = 0
"""If the box < r then extend the box. Otherwise wrap the cell"""
if np.any(box < r):
if verbose > 0:
string = "Replicating cell by %i, %i, %i (x, y, z)"\
% (rep[0] + 1, rep[1] + 1, rep[2] + 1)
ut.infoPrint(string)
pos_to, cell = self.getExtendedPositions(x = rep[0], y = rep[1], z = rep[2],\
idx = idx_to, return_cart = True,\
verbose = verbose - 1)
"""Change to cartesian coordinates"""
self.dir2car()
"""Change back to direct coordinates using the new extended cell"""
pos_to = np.matmul(np.linalg.inv(cell), pos_to.T).T
pos_from = np.matmul(np.linalg.inv(cell), self.pos.T).T
else:
"""Change to direct coordinates"""
self.car2dir()
if verbose > 0:
string = "Cell is only wrapped, not extended"
ut.infoPrint(string)
pos_to = self.pos.copy()[idx_to, :]
pos_from = self.pos.copy()
cell = self.cell
ps = pos_to.shape[0]
dist = np.zeros((np.shape(idx)[0] * ps, 3))
"""Measure distances between all specified atoms, wrap cell"""
for i, item in enumerate(idx):
d = pos_to - pos_from[[item], :]
d[d > 0.5] -= 1
d[d < -0.5] += 1
dist[i * ps : (i + 1) * ps, :] = d
"""Convert to cartesian coordinates"""
dist = np.matmul(cell, dist.T).T
"""Calculate the distances"""
dist = np.linalg.norm(dist, axis = 1)
"""Remove distances outside of radius r"""
dist = dist[dist < r]
"""Remove the 0 distances, (atom to it self)"""
dist = dist[dist > 0]
return dist