def _link_ax_w_pos_2_vmd(ax, pos, geoms, **customVMD_kwargs):
r"""
Initial idea and key VMD-interface for this function comes from @fabian-paul
#TODO: CLEAN THE TEMPFILE
"""
# Prepare tempdir
tmpdir = tempfile.mkdtemp('vmd_interface')
print("please remember to: rm -r %s"%tmpdir)
# Prepare files
topfile = os.path.join(tmpdir,'top.pdb')
trjfile = os.path.join(tmpdir,'trj.xtc')
geoms[0].save(topfile)
geoms[1:].superpose(geoms[0]).save(trjfile)
# Create pipe
pipefile = os.path.join(tmpdir,'vmd_cmds.tmp.vmd')
os.mkfifo(pipefile)
os.system("vmd < %s & "%pipefile)
vmdpipe = open(pipefile,'w')
[vmdpipe.write(l) for l in customvmd(topfile, trajfile=trjfile, vmdout=None,
**customVMD_kwargs)]
myflush(vmdpipe)
kdtree = _cKDTree(pos)
x, y = pos.T
lineh = ax.axhline(ax.get_ybound()[0], c="black", ls='--')
linev = ax.axvline(ax.get_xbound()[0], c="black", ls='--')
dot, = ax.plot(pos[0,0],pos[0,1], 'o', c='red', ms=7)
def _onclick(event):
linev.set_xdata((event.xdata, event.xdata))
lineh.set_ydata((event.ydata, event.ydata))
data = [event.xdata, event.ydata]
_, index = kdtree.query(x=data, k=1)
dot.set_xdata((x[index]))
dot.set_ydata((y[index]))
vmdpipe.write(" animate goto %u;\nlist;\n\n"%index)
#myflush(vmdpipe,
#size=1e4
# )
# Connect axes to widget
axes_widget = _AxesWidget(ax)
axes_widget.connect_event('button_release_event', onclick)
return vmdpipe
def setUpClass(self):
self.MD_trajectory_files = glob(
molpx._molpxdir(
join='notebooks/data/c-alpha_centered.stride.1000*xtc'))
self.MD_topology_file = molpx._molpxdir(
join='notebooks/data/bpti-c-alpha_centered.pdb')
self.MD_topology = md.load(self.MD_topology_file).top
self.MD_trajectories = [
md.load(ff, top=self.MD_topology_file)
for ff in self.MD_trajectory_files
]
self.MD_trajectory = self.MD_trajectories[0]
self.pos = np.random.rand(self.MD_trajectory.n_frames, 2)
self.kdtree = _cKDTree(self.pos)
def _naturalNumbering(self, coords):
"""
Natural numbering based on DMPlex distribution
Args:
coords: mesh coordinates
"""
self.lcoords = self.dm.getCoordinatesLocal().array.reshape(-1, 3)
self.npoints = self.lcoords.shape[0]
if MPIsize > 1:
tree = _cKDTree(coords[:, :2])
distances, nat2loc = tree.query(self.lcoords[:, :2], 1)
self.natural2local = nat2loc.copy()
del tree, distances, nat2loc
else:
self.natural2local = np.arange(0, self.npoints, dtype=int)
return
def cKDTree(
data, leafsize=16, compact_nodes=True, balanced_tree=False, **kwargs):
return _cKDTree(
data, leafsize=leafsize, compact_nodes=compact_nodes,
balanced_tree=balanced_tree, **kwargs)
def __init__(self, data, **kwargs):
self.tree = _cKDTree(data, **kwargs)
def _meshAdvectorSphere(self, tectonicX, tectonicY, tectonicZ, timer):
"""
Advect spherical mesh horizontally and interpolate mesh information
"""
# Move coordinates
XYZ = np.zeros((self.gpoints, 3))
XYZ[:, 0] = tectonicX #self.gcoords[:,0] + tectonicX*timer
XYZ[:, 1] = tectonicY #self.gcoords[:,1] + tectonicY*timer
XYZ[:, 2] = tectonicZ #self.gcoords[:,2] + tectonicZ*timer
XYZ = XYZ[1:]
# Get mesh variables to interpolate
# Elevation
elev = np.zeros(self.gpoints)
elev.fill(-1.e8)
elev[self.natural2local] = self.hLocal.getArray().copy()
MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, elev, op=MPI.MAX)
elev = elev[1:]
# Erosion/deposition
erodep = np.zeros(self.gpoints)
erodep.fill(-1.e8)
erodep[self.natural2local] = self.cumEDLocal.getArray().copy()
MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, erodep, op=MPI.MAX)
erodep = erodep[1:]
# Soil thickness
if self.Ksed > 0.:
soil = np.zeros(self.gpoints)
soil.fill(-1.e8)
soil[self.natural2local] = self.HsoilLocal.getArray().copy()
MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, soil, op=MPI.MAX)
soil = soil[1:]
# Build kd-tree
tree = _cKDTree(XYZ)
distances, indices = tree.query(self.lcoords, k=10)
# Inverse weighting distance...
weights = 1.0 / distances**2
onIDs = np.where(distances[:, 0] == 0)[0]
nelev = np.sum(weights * elev[indices], axis=1) / np.sum(weights,
axis=1)
nerodep = np.sum(weights * erodep[indices], axis=1) / np.sum(weights,
axis=1)
if self.Ksed > 0.:
nsoil = np.sum(weights * soil[indices], axis=1) / np.sum(weights,
axis=1)
if len(onIDs) > 0:
nelev[onIDs] = elev[indices[onIDs, 0]]
nerodep[onIDs] = erodep[indices[onIDs, 0]]
if self.Ksed > 0.:
nsoil[onIDs] = soil[indices[onIDs, 0]]
self.hLocal.setArray(nelev)
self.dm.localToGlobal(self.hLocal, self.hGlobal, 1)
self.dm.globalToLocal(self.hGlobal, self.hLocal, 1)
del elev, nelev
self.cumEDLocal.setArray(nerodep)
self.dm.localToGlobal(self.cumEDLocal, self.cumED, 1)
self.dm.globalToLocal(self.cumED, self.cumEDLocal, 1)
del erodep, nerodep
if self.Ksed > 0.:
self.HsoilLocal.setArray(nsoil)
self.dm.localToGlobal(self.HsoilLocal, self.Hsoil, 1)
self.dm.globalToLocal(self.Hsoil, self.HsoilLocal, 1)
del soil, nsoil
del XYZ, tree, weights, distances, indices
return
def _meshAdvector(self, tectonicX, tectonicY, timer):
"""
Advect the mesh horizontally and interpolate mesh information
"""
# Move horizontally coordinates
XYZ = np.zeros((self.gpoints, 3))
XYZ[:, 0] = self.gcoords[:, 0] + tectonicX * timer
XYZ[:, 1] = self.gcoords[:, 1] + tectonicY * timer
# Get mesh variables to interpolate
# Elevation
elev = np.zeros(self.gpoints)
elev.fill(-1.e8)
elev[self.natural2local] = self.hLocal.getArray().copy()
MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, elev, op=MPI.MAX)
interpolator = CloughTocher2DInterpolator(XYZ[:, :2], elev)
nelev = interpolator(self.lcoords[:, :2])
id_NaNs = np.isnan(nelev)
nelev[id_NaNs] = 0.
# Erosion/deposition
erodep = np.zeros(self.gpoints)
erodep.fill(-1.e8)
erodep[self.natural2local] = self.cumEDLocal.getArray().copy()
MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, erodep, op=MPI.MAX)
interpolator = CloughTocher2DInterpolator(XYZ[:, :2], erodep)
nerodep = interpolator(self.lcoords[:, :2])
nerodep[id_NaNs] = 0.
# Soil thickness
if self.Ksed > 0.:
soil = np.zeros(self.gpoints)
soil.fill(-1.e8)
soil[self.natural2local] = self.HsoilLocal.getArray().copy()
MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, soil, op=MPI.MAX)
interpolator = CloughTocher2DInterpolator(XYZ[:, :2], soil)
nsoil = interpolator(self.lcoords[:, :2])
nsoil[id_NaNs] = 0.
# Build kd-tree
tree = _cKDTree(XYZ)
distances, indices = tree.query(self.lcoords[id_NaNs, :], k=10)
# Inverse weighting distance...
weights = 1.0 / distances**2
onIDs = np.where(distances[:, 0] == 0)[0]
nelev[id_NaNs] = np.sum(weights * elev[indices], axis=1) / np.sum(
weights, axis=1)
self.hLocal.setArray(nelev)
self.dm.localToGlobal(self.hLocal, self.hGlobal, 1)
self.dm.globalToLocal(self.hGlobal, self.hLocal, 1)
del elev, nelev
self.cumEDLocal.setArray(nerodep)
self.dm.localToGlobal(self.cumEDLocal, self.cumED, 1)
self.dm.globalToLocal(self.cumED, self.cumEDLocal, 1)
del erodep, nerodep
if self.Ksed > 0.:
self.HsoilLocal.setArray(nsoil)
self.dm.localToGlobal(self.HsoilLocal, self.Hsoil, 1)
self.dm.globalToLocal(self.Hsoil, self.HsoilLocal, 1)
del soil, nsoil
del XYZ, tree, weights, distances, indices
return
def build_tree(self):
# Use ax.transData to compute distance in pixels
# regardelss of the axes units (http://matplotlib.org/users/transforms_tutorial.html)
# Corresponds to the visual distance between clicked point and target point
self.kdtree = _cKDTree(self.ax.transData.transform(self.pos))