def basic_compute_loop(compute_function,
looper,
run_parallel=True,
debug=None):
"""
Canonical form of the basic compute loop.
!!! remove this from contacts.py when it works
"""
#---send the frame as the debug argument
if debug != None and debug != False:
fr = debug
incoming = compute_function(**looper[fr])
import ipdb
ipdb.set_trace()
sys.quit()
start = time.time()
if run_parallel:
incoming = Parallel(n_jobs=8, verbose=10 if debug else 0)(
delayed(compute_function, has_shareable_memory)(**looper[ll])
for ll in framelooper(len(looper), start=start))
else:
incoming = []
for ll in framelooper(len(looper)):
incoming.append(compute_function(**looper[ll]))
return incoming
def basic_compute_loop(compute_function,looper,run_parallel=True,debug=False): """Canonical form of the basic compute loop.""" start = time.time() if run_parallel: incoming = Parallel(n_jobs=8,verbose=10 if debug else 0)( delayed(compute_function,has_shareable_memory)(**looper[ll]) for ll in framelooper(len(looper),start=start)) else: incoming = [] for ll in framelooper(len(looper)): incoming.append(compute_function(**looper[ll])) return incoming
def basic_compute_loop(compute_function,
looper,
run_parallel=True,
debug=False):
"""Canonical form of the basic compute loop."""
start = time.time()
if run_parallel:
incoming = Parallel(n_jobs=8, verbose=10 if debug else 0)(
delayed(compute_function, has_shareable_memory)(**looper[ll])
for ll in framelooper(len(looper), start=start))
else:
incoming = []
for ll in framelooper(len(looper)):
incoming.append(compute_function(**looper[ll]))
return incoming
def import_readymade_meso_v2_membrane(**kwargs):
"""
Compute bilayer midplane structures for studying undulations.
Adapted from `undulations.py`.
"""
import ipdb
ipdb.set_trace()
#---parameters
sn = kwargs['sn']
work = kwargs['workspace']
calc = kwargs['calc']
#---import mesh points
points = import_membrane_mesh(sn=sn, calc=calc, work=work)
#---ensure there are the same number of points
points_shapes = list(set([p.shape for p in points]))
if len(points_shapes) != 1:
raise Exception('some frames have a different number of points: %s' %
points_shapes)
else:
npoints, ncols = points_shapes[0]
if ncols != 4: raise Exception('expecting 4-column input on incoming data')
#---with a consistent number of points everything is an array
points = np.array(points)[:, :, :3]
#---previously checked that the minimum points were identically zero but this was not always true
#---box vectors are just the maximum points
#---! check that this assumption makes sense
vecs = points.max(axis=1)[:, :3]
#---debug the shapes in 3D
if False:
from codes import review3d
fr = 0
review3d.pbcbox(vecs[fr])
review3d.review3d(points=[points[fr][:, :3]], radius=10)
grid_spacing = calc['specs']['grid_spacing']
nframes = len(points)
#---choose grid dimensions
grid = np.array([round(i)
for i in np.mean(vecs, axis=0) / grid_spacing])[:2]
#---compute in parallel
start = time.time()
mesh = [[]]
mesh[0] = Parallel(n_jobs=work.nprocs, verbose=0)(
delayed(makemesh_regular, has_shareable_memory)(points[fr], vecs[fr],
grid)
for fr in framelooper(nframes, start=start, text='frame'))
checktime()
#---pack
attrs, result = {}, {}
result['mesh'] = np.array(mesh)
result['grid'] = np.array(grid)
result['nframes'] = np.array(nframes)
result['vecs'] = vecs
attrs['grid_spacing'] = grid_spacing
#---introduce a dated validator string to ensure that any changes to the pipeline do not overwrite
#---...other data and are fully propagated downstream
attrs['validator'] = '2017.08.16.1930'
return result, attrs
def lipid_mesh(**kwargs):
"""
Compute monolayer mesh objects.
"""
#---parameters
sn = kwargs['sn']
work = kwargs['workspace']
calc = kwargs['calc']
dat = kwargs['upstream']['lipid_abstractor']
resnames = dat['resnames']
monolayer_indices = dat['monolayer_indices']
nframes = dat['nframes']
debug = kwargs.pop('debug',False)
kwargs_out = dict(curvilinear=calc.get('specs',{}).get('curvilinear',False))
#---parallel
mesh = [[],[]]
if debug:
mn,fr = 0,10
makemesh(dat['points'][fr][where(monolayer_indices==mn)],dat['vecs'][fr],
debug=True,**kwargs_out)
sys.exit(1)
for mn in range(2):
start = time.time()
mesh[mn] = Parallel(n_jobs=work.nprocs,verbose=0)(
delayed(makemesh)(
dat['points'][fr][where(monolayer_indices==mn)],dat['vecs'][fr],**kwargs_out)
for fr in framelooper(nframes,start=start,text='monolayer %d, frame'%mn))
checktime()
#---pack
attrs,result = {},{}
result['nframes'] = array(nframes)
result['vecs'] = dat['vecs']
result['resnames'] = resnames
result['monolayer_indices'] = monolayer_indices
#---pack mesh objects
#---keys include: vertnorms simplices nmol facenorms gauss points vec ghost_ids mean principals areas
keylist = mesh[0][0].keys()
for key in keylist:
for mn in range(2):
for fr in range(nframes):
result['%d.%d.%s'%(mn,fr,key)] = mesh[mn][fr][key]
return result,attrs
def undulations(**kwargs):
"""
Compute bilayer midplane structures for studying undulations.
"""
#---parameters
sn = kwargs['sn']
work = kwargs['workspace']
calc = kwargs['calc']
upname = 'lipid_abstractor'
grid_spacing = calc['specs']['grid_spacing']
vecs = datmerge(kwargs, upname, 'vecs')
nframes = int(np.sum(datmerge(kwargs, upname, 'nframes')))
trajectory = datmerge(kwargs, upname, 'points')
attrs, result = {}, {}
#---! hacking through error with monolayer separation
try:
monolayer_indices = kwargs['upstream'][upname +
'0']['monolayer_indices']
except:
monolayer_indices = kwargs['upstream'][upname]['monolayer_indices']
#---choose grid dimensions
grid = np.array([round(i)
for i in np.mean(vecs, axis=0) / grid_spacing])[:2]
#---! removed timeseries from result for new version of omnicalc
#---parallel
mesh = [[], []]
for mn in range(2):
start = time.time()
mesh[mn] = Parallel(
n_jobs=work.nprocs, verbose=0, require='sharedmem')(
delayed(makemesh_regular)(trajectory[fr][np.where(
monolayer_indices == mn)], vecs[fr], grid)
for fr in framelooper(
nframes, start=start, text='monolayer %d, frame' % mn))
checktime()
#---pack
result['mesh'] = np.array(mesh)
result['grid'] = np.array(grid)
result['nframes'] = np.array(nframes)
result['vecs'] = vecs
attrs['grid_spacing'] = grid_spacing
return result, attrs
def undulations(**kwargs):
"""
Compute bilayer midplane structures for studying undulations.
"""
#---parameters
sn = kwargs['sn']
work = kwargs['workspace']
calc = kwargs['calc']
grid_spacing = calc['specs']['grid_spacing']
dat = kwargs['upstream']['lipid_abstractor']
nframes = dat['nframes']
#---choose grid dimensions
grid = array([round(i)
for i in mean(dat['vecs'], axis=0) / grid_spacing])[:2]
monolayer_indices = dat['monolayer_indices']
#---parallel
start = time.time()
mesh = [[], []]
for mn in range(2):
mesh[mn] = Parallel(n_jobs=work.nprocs, verbose=0)(
delayed(makemesh_regular)(dat['points'][fr][where(
monolayer_indices == mn)], dat['vecs'][fr], grid)
for fr in framelooper(
nframes, start=start, text='monolayer %d, frame' % mn))
checktime()
#---pack
attrs, result = {}, {}
result['mesh'] = array(mesh)
result['grid'] = array(grid)
result['nframes'] = array(nframes)
result['vecs'] = dat['vecs']
result['timeseries'] = work.slice(sn)[kwargs['slice_name']][
'all' if not kwargs['group'] else kwargs['group']]['timeseries']
attrs['grid_spacing'] = grid_spacing
return result, attrs
def electron_density_profiles_deprecated(grofile, trajfile, **kwargs):
"""
Compute the electron density profile
"""
#from calcs.codes.mesh import identify_lipid_leaflets
bin_size = kwargs['calc']['specs']['bin_size']
chargedict = {
'^N(?!A$)': 7,
'^C[0-9]+': 6,
'^CL$': 17,
'^H': 1,
'^O': 8,
'^P': 15,
'^Cal': 18,
'^MG': 10,
'^NA': 11,
'^S': 16,
'K': 18
}
group_regexes = ['.+', '^(OW)|(HW(1|2))$', '^C[0-9]+']
#---unpack
sn = kwargs['sn']
work = kwargs['workspace']
#---prepare universe
grofile, trajfile = kwargs['structure'], kwargs['trajectory']
uni = MDAnalysis.Universe(grofile, trajfile)
nframes = len(uni.trajectory)
#---MDAnalysis uses Angstroms not nm
lenscale = 10.
#---select residues of interest
selector = kwargs['calc']['specs']['selector']
monolayer_cutoff = kwargs['calc']['specs']['selector']['monolayer_cutoff']
#---center of mass over residues
if 'type' in selector and selector[
'type'] == 'com' and 'resnames' in selector:
resnames = selector['resnames']
selstring = '(' + ' or '.join(['resname %s' % i
for i in resnames]) + ')'
else:
raise Exception('\n[ERROR] unclear selection %s' % str(selector))
#---compute masses by atoms within the selection
sel = uni.select_atoms(selstring)
mass_table = {
'H': 1.008,
'C': 12.011,
'O': 15.999,
'N': 14.007,
'P': 30.974
}
masses = np.array([mass_table[i[0]] for i in sel.atoms.names])
resids = sel.resids
#---create lookup table of residue indices
divider = [np.where(resids == r) for r in np.unique(resids)]
#---load trajectory into memory
trajectory, vecs = [], []
for fr in range(nframes):
status('loading frame', tag='load', i=fr, looplen=nframes)
uni.trajectory[fr]
trajectory.append(sel.positions / lenscale)
vecs.append(sel.dimensions[:3])
#---parallel
start = time.time()
coms = Parallel(n_jobs=work.nprocs, verbose=0)(
delayed(centroid)(trajectory[fr], masses, divider)
for fr in framelooper(nframes, start=start))
#---identify monolayers
#---! why though?
#---note that this could just refer to the mesh object but this is very fast
if False:
monolayer_indices = identify_lipid_leaflets(
coms[0], vecs[0], monolayer_cutoff=monolayer_cutoff)
#---load trajectory into memory
allsel = uni.select_atoms('all')
trajectory, vecs = [], []
for fr in range(nframes):
status('loading frame', tag='load', i=fr, looplen=nframes)
uni.trajectory[fr]
trajectory.append(allsel.positions / lenscale)
vecs.append(allsel.dimensions[:3])
trajectory = np.array(trajectory)
vecs = np.array(vecs) / lenscale
#---center the mean of com positions at z=0
midplane_heights = np.array(
[np.mean(coms[fr], axis=0)[2] for fr in range(nframes)])
for fr in range(nframes):
trajectory[fr, :, 2] -= midplane_heights[fr]
#---correct for periodic boundaries
for fr in range(nframes):
trajectory[fr, :,
2] -= (trajectory[fr, :, 2] > vecs[fr, 2] / 2.) * vecs[fr,
2]
trajectory[
fr, :,
2] += (trajectory[fr, :, 2] < -1 * vecs[fr, 2] / 2.) * vecs[fr, 2]
#---offset by one bin width here so the use of astype(int) is symmetric later on
trajectory[..., 2] += bin_size / 2.
#---assign charges
namelist = uni.atoms.names
resnamelist = list(set(uni.atoms.resnames))
#---charge dictionary for the atoms in this particular system
chargedict_obs = dict([
(name, [chargedict[key] for key in chargedict if re.match(key, name)])
for name in np.unique(namelist)
])
unclear_charges = dict([(key, val) for key, val in chargedict_obs.items()
if len(val) != 1])
if any(unclear_charges):
raise Exception('charges for these atoms were not specified: %s' %
unclear_charges)
chargedict_obs = dict([(key, val[0])
for key, val in chargedict_obs.items()])
charges = [chargedict_obs[n] for n in namelist]
#---identify atoms for each residue type
groups = [[ii for ii, i in enumerate(uni.atoms) if i.resname == r]
for r in resnamelist]
groups += [
np.array([i for i, j in enumerate(namelist) if re.match(reg, j)])
for reg in group_regexes
]
#---bin the heights according to bin_size
xx = np.array(np.floor(trajectory[:, :, 2] / bin_size)).astype(int)
offset = xx.min()
xx -= xx.min()
bincounts_by_group = [np.zeros(xx.max() + 1) for grp in groups]
start = time.time()
for fr in range(nframes):
status('electron density',
i=fr,
looplen=nframes,
tag='compute',
start=start)
for g, grp in enumerate(groups):
for i, z in enumerate(xx[fr][grp]):
bincounts_by_group[g][z] += charges[grp[i]]
xvals = bin_size * np.arange(0, len(bincounts_by_group[0])) - len(
bincounts_by_group[0]) * bin_size / 2.
mvecs = np.mean(np.array([vecs[i] for i in range(nframes)]), axis=0)
scaleconst = np.product(mvecs[:2]) * (mvecs[2] / len(xvals)) * nframes
results, attrs = {}, {}
results['midplane_heights'] = midplane_heights
#---note that bincoungs_by_group is ordered first by resnamelist then group_regexes
results['bincounts_by_group'] = np.array(bincounts_by_group) / scaleconst
for index, group in enumerate(groups):
results['groups%d' % index] = np.array(group)
results['offset'] = np.array(offset)
attrs['selector'] = selector
attrs['resnamelist'] = resnamelist
attrs['bin_size'] = bin_size
attrs['group_regexes'] = group_regexes
return results, attrs
def ion_binding(grofile,trajfile,**kwargs):
"""
Analyze bound ion distances to the nearest lipids.
"""
#---unpack
sn = kwargs['sn']
work = kwargs['workspace']
nrank = kwargs['calc']['specs']['nrank']
#---! note that the parallel code is severely broken and caused wierd spikes in the distances!!!
#! on 2018.07.13 trying to revive parallel
compute_parallel = True
#---prepare universe
#! is this deprecated? grofile,trajfile = [work.slice(sn)['current']['all'][i] for i in ['gro','xtc']]
#! uni = MDAnalysis.Universe(work.postdir+grofile,work.postdir+trajfile)
uni = MDAnalysis.Universe(grofile,trajfile)
nframes = len(uni.trajectory)
#---MDAnalysis uses Angstroms not nm
lenscale = 10.
#---compute masses by atoms within the selection
sel_lipids = uni.select_atoms(' or '.join('resname %s'%r
for r in work.vars['selectors']['resnames_lipid_chol']))
sel_ions = uni.select_atoms(work.vars['selectors']['cations'])
#---load lipid points into memory
trajectory_ions = zeros((nframes,len(sel_ions),3))
trajectory_lipids = zeros((nframes,len(sel_lipids),3))
vecs = zeros((nframes,3))
for fr in range(nframes):
status('loading frame',tag='load',i=fr,looplen=nframes)
uni.trajectory[fr]
#! mdanalysis removed coordinates(), using positions with cast to be sure
trajectory_lipids[fr] = np.array(sel_lipids.positions)/lenscale
trajectory_ions[fr] = np.array(sel_ions.positions)/lenscale
vecs[fr] = sel_lipids.dimensions[:3]/lenscale
monolayer_indices = kwargs['upstream']['lipid_abstractor']['monolayer_indices']
resids = kwargs['upstream']['lipid_abstractor']['resids']
monolayer_residues = [resids[where(monolayer_indices==mn)[0]] for mn in range(2)]
group_lipid = uni.select_atoms(' or '.join(['resid '+str(i) for mononum in range(2)
for i in monolayer_residues[mononum]]))
lipid_resids = array([i.resid for i in group_lipid])
if work.meta[sn]['composition_name'] != 'asymmetric': lipid_resid_subselect = slice(None,None)
#---hack to account for asymmetric bilayer by analyzing only the first (top) monolayer
else: lipid_resid_subselect = where([i.resid in monolayer_residues[0] for i in group_lipid])[0]
#---parallel partnerfinder
start = time.time()
lipid_resids = array([i.resid for i in group_lipid])
if not compute_parallel:
incoming = []
start = time.time()
for fr in range(nframes):
status('frame',i=fr,looplen=nframes,start=start,tag='compute')
ans = partnerfinder(trajectory_lipids[fr],trajectory_ions[fr],vecs[fr],
lipid_resids,nrank,includes=lipid_resid_subselect)
incoming.append(ans)
else:
incoming = Parallel(n_jobs=4,verbose=0,require='sharedmem')(
delayed(partnerfinder)
(trajectory_lipids[fr],trajectory_ions[fr],vecs[fr],
lipid_resids,nrank,includes=lipid_resid_subselect)
for fr in framelooper(nframes,start=start))
n_ion_atoms,n_lipid_atoms = len(sel_ions),len(sel_lipids)
lipid_distances = zeros((nframes,n_ion_atoms,nrank))
partners_atoms = zeros((nframes,n_ion_atoms,nrank))
#---unpack
start = time.time()
for fr in range(nframes):
status('[UNPACK] frame',i=fr,looplen=nframes,start=start)
lipid_distances[fr] = incoming[fr][0]
partners_atoms[fr] = incoming[fr][1]
result,attrs = {},{}
attrs['nrank'] = nrank
result['lipid_distances'] = lipid_distances
result['partners_atoms'] = partners_atoms.astype(int)
result['names'] = array([i.name for i in group_lipid])
result['resnames'] = array([i.resname for i in group_lipid])
result['resids'] = array([i.resid for i in group_lipid])
result['nframes'] = array(nframes)
return result,attrs
def lipid_abstractor(grofile, trajfile, **kwargs):
"""
LIPID ABSTRACTOR
Reduce a bilayer simulation to a set of points.
"""
#---unpack
sn = kwargs['sn']
work = kwargs['workspace']
parallel = kwargs.get('parallel', False)
#---prepare universe
#---note that the universe throws a UserWarning on coarse-grained systems
#---...which is annoying to elevate to error stage and handled below without problems
uni = MDAnalysis.Universe(grofile, trajfile)
nframes = len(uni.trajectory)
#---MDAnalysis uses Angstroms not nm
lenscale = 10.
#---select residues of interest
selector = kwargs['calc']['specs']['selector']
nojumps = kwargs['calc']['specs'].get('nojumps', '')
#---center of mass over residues
if 'type' in selector and selector[
'type'] == 'com' and 'resnames' in selector:
resnames = selector['resnames']
selstring = '(' + ' or '.join(['resname %s' % i
for i in resnames]) + ')'
elif 'type' in selector and selector[
'type'] == 'select' and 'selection' in selector:
if 'resnames' not in selector:
raise Exception('add resnames to the selector')
selstring = selector['selection']
elif selector.get('type', None) == 'custom':
custom_exec_vars = dict(uni=uni, selector=selector)
exec(selector['custom'], globals(), custom_exec_vars)
selstring = custom_exec_vars['selstring']
else:
raise Exception('\n[ERROR] unclear selection %s' % str(selector))
#---compute masses by atoms within the selection
sel = uni.select_atoms(selstring)
if len(sel) == 0: raise Exception('empty selection')
mass_table = {
'H': 1.008,
'C': 12.011,
'O': 15.999,
'N': 14.007,
'P': 30.974,
'S': 32.065
}
missing_atoms_aamd = list(
set([i[0] for i in sel.atoms.names if i[0] not in mass_table]))
if any(missing_atoms_aamd):
print(
'[WARNING] missing mass for atoms %s so we assume this is coarse-grained'
% missing_atoms_aamd)
#---MARTINI masses
mass_table = {
'C': 72,
'N': 72,
'P': 72,
'S': 45,
'G': 72,
'D': 72,
'R': 72
}
missing_atoms_cgmd = list(
set([i[0] for i in sel.atoms.names if i[0] not in mass_table]))
if any(missing_atoms_cgmd):
raise Exception(
'we are trying to assign masses. if this simulation is atomistic then we are '
+
'missing atoms "%s". if it is MARTINI then we are missing atoms "%s"'
% (missing_atoms_aamd, missing_atoms_cgmd))
else:
masses = np.array([mass_table[i[0]] for i in sel.atoms.names])
else:
masses = np.array([mass_table[i[0]] for i in sel.atoms.names])
# note that the following sequence has been reworked to reflect apparent changes in the
# ... residue-handling. previously we used `if len(sel.resids)==len(np.unique(sel.resids)):` but this
# ... is now incompatible
resids = sel.residues.resids
# create lookup table of residue indices
if len(resids) == len(np.unique(resids)):
divider = [np.where(sel.resids == r) for r in np.unique(resids)]
# note that redundant residue numbering requires special treatment
else:
#! note that the resid handling change above may not have been implemented in the custom method below
if (('type' in selector) and (selector['type'] in ['com', 'select'])
and ('resnames' in selector)):
#---note that MDAnalysis sel.residues *cannot* handle redundant numbering
#---note also that some test cases have redundant residues *and* adjacent residues with
#---...the same numbering. previously we tried a method that used the following sequence:
#---......divider = [np.where(np.in1d(np.where(np.in1d(
#---..........uni.select_atoms('all').resnames,resnames))[0],d))[0] for d in divider_abs]
#---...however this method is flawed because it uses MDAnalysis sel.residues and in fact
#---...since it recently worked, RPB suspects that a recent patch to MDAnalysis has broken it
#---note that rpb started a method to correct this and found v inconsistent MDAnalysis behavior
#---the final fix is heavy-handed: leaving nothing to MDAnalysis subselections
allsel = uni.select_atoms('all')
lipids = np.where(
np.in1d(allsel.resnames, np.array(selector['resnames'])))[0]
resid_changes = np.concatenate(([
-1
], np.where(
allsel[lipids].resids[1:] != allsel[lipids].resids[:-1])[0]))
residue_atomcounts = resid_changes[1:] - resid_changes[:-1]
#---get the residue names for each lipid in our selection by the first atom in that lipid
#---the resid_changes is prepended with -1 in the unlikely (but it happened) event that
#---...a unique lipid leads this list (note that a blase comment dismissed this possibility at
#---...first!) and here we correct the resnames list to reflect this. resnames samples the last
#---...atom in each residue from allsel
resnames = np.concatenate(
(allsel[lipids].resnames[resid_changes[1:]],
[allsel[lipids].resnames[-1]]))
guess_atoms_per_residue = np.array(
zip(resnames, residue_atomcounts))
#---get consensus counts for each lipid name
atoms_per_residue = {}
for name in np.unique(resnames):
#---get the most common count
counts, obs_counts = np.unique(
guess_atoms_per_residue[:, 1][np.where(
guess_atoms_per_residue[:, 0] == name)[0]].astype(int),
return_counts=True)
atoms_per_residue[name] = counts[obs_counts.argmax()]
#---faster method
resid_to_start = np.transpose(
np.unique(allsel.resids, return_index=True))
resid_to_start = np.concatenate(
(resid_to_start, [[resid_to_start[-1][0] + 1,
len(lipids)]]))
divider = np.array([
np.arange(i, j)
for i, j in np.transpose((resid_to_start[:, 1][:-1],
resid_to_start[:, 1][1:]))
])
#---make sure no molecules have the wrong number of atoms
if not set(np.unique([len(i) for i in divider])) == set(
atoms_per_residue.values()):
status('checking lipid residue indices the careful way',
tag='warning')
#---the following method is slow on large systems. we use it when the fast method above fails
#---iterate over the list of lipid atoms and get the indices for each N-atoms for each lipid
counter, divider = 0, []
while counter < len(lipids):
status('indexing lipids',
i=counter,
looplen=len(lipids),
tag='compute')
#---until the end, get the next lipid resname
this_resname = allsel.resnames[lipids][counter]
if selector['type'] == 'select':
#---the only way to subselect here is to select on each divided lipid (since
#---...the procedure above has correctly divided the lipids). we perform the
#---...subselection by pivoting over indices
#---! this method needs checked
this_inds = np.arange(
counter, counter + atoms_per_residue[this_resname])
this_lipid = allsel[lipids][this_inds]
this_subsel = np.where(
np.in1d(
this_lipid.indices,
this_lipid.select_atoms(
selector['selection']).indices))[0]
divider.append(this_inds[this_subsel])
else:
divider.append(
np.arange(
counter,
counter + atoms_per_residue[this_resname]))
counter += atoms_per_residue[this_resname]
#---in the careful method the sel from above is broken but allsel[lipids] is correct
sel = allsel[lipids]
masses = np.array([mass_table[i[0]] for i in sel.atoms.names])
else:
import ipdb
ipdb.set_trace()
raise Exception(
'residues have redundant resids and selection is not the easy one'
)
#---load trajectory into memory
trajectory, vecs = [], []
for fr in range(nframes):
status('loading frame', tag='load', i=fr, looplen=nframes)
uni.trajectory[fr]
trajectory.append(sel.positions / lenscale)
#! critical fix: you must cast the dimensions or you get repeated vectors
vecs.append(np.array(uni.trajectory[fr].dimensions[:3]))
vecs = np.array(vecs) / lenscale
checktime()
#---parallel
start = time.time()
if parallel:
coms = Parallel(n_jobs=work.nprocs, verbose=0)(
delayed(codes.mesh.centroid)(trajectory[fr], masses, divider)
for fr in framelooper(nframes, start=start))
else:
coms = []
for fr in range(nframes):
status('computing centroid',
tag='compute',
i=fr,
looplen=nframes,
start=start)
coms.append(codes.mesh.centroid(trajectory[fr], masses, divider))
#---identify leaflets
status('identify leaflets', tag='compute')
separator = kwargs['calc']['specs'].get('separator', {})
leaflet_finder_trials = separator.get('trials', 3)
#---preselect a few frames, always including the zeroth
selected_frames = [0] + list(
np.random.choice(
np.arange(1, nframes), leaflet_finder_trials, replace=False))
#---alternate lipid representation is useful for separating monolayers
if 'lipid_tip' in separator:
tip_select = separator['lipid_tip']
sel = uni.select_atoms(tip_select)
atoms_separator = []
for fr in selected_frames:
uni.trajectory[fr]
atoms_separator.append(sel.positions / lenscale)
#---default is to use the centers of mass to distinguish leaflets
else:
atoms_separator = [coms[fr] for fr in selected_frames]
#---pass frames to the leaflet finder, which has legacy and cluster modes
leaflet_finder = codes.mesh.LeafletFinder(
atoms_separator=atoms_separator,
#---pass along the corresponding vectors for topologize
vecs=[vecs[i] for i in selected_frames],
cluster=separator.get('cluster', False),
cluster_neighbors=separator.get('cluster_neighbors', None),
topologize_tolerance=separator.get('topologize_tolerance', None))
#---get the indices from the leaflet finder
monolayer_indices = leaflet_finder.monolayer_indices
# for convenience when doing planar bilayers we put the zero index on top
top_mono = np.argmax([
atoms_separator[0][monolayer_indices == i][:, 2].mean()
for i in range(2)
])
if top_mono != 0: monolayer_indices = 1 - monolayer_indices
checktime()
coms_out = np.array(coms)
#---remove jumping in some directions if requested
if nojumps:
nojump_dims = ['xyz'.index(j) for j in nojumps]
nobjs = coms_out.shape[1]
displacements = np.array([(coms_out[1:] - coms_out[:-1])[..., i]
for i in range(3)])
for d in nojump_dims:
shift_binary = (
np.abs(displacements) * (1. - 2 * (displacements < 0)) /
(np.transpose(np.tile(vecs[:-1],
(nobjs, 1, 1))) / 2.))[d].astype(int)
shift = (np.cumsum(-1 * shift_binary, axis=0) *
np.transpose(np.tile(vecs[:-1, d], (nobjs, 1))))
coms_out[1:, :, d] += shift
#---pack
attrs, result = {}, {}
attrs['selector'] = selector
attrs['nojumps'] = nojumps
result['resnames'] = np.array(sel.residues.resnames)
result['monolayer_indices'] = np.array(monolayer_indices)
result['vecs'] = vecs
result['nframes'] = np.array(nframes)
result['points'] = coms_out
result['resids'] = np.array(np.unique(resids))
result['resids_exact'] = resids
attrs['separator'] = kwargs['calc']['specs']['separator']
return result, attrs
