Skip to content

Commit

Permalink
Refactor in light of discussion in #48.
Browse files Browse the repository at this point in the history 
We will encourage setting a Sim universe directly with an existing
Universe, with the hope that we can eventually hide UniverseDefinition.
At the moment, though, we leave it exposed for any manual interventions
that can't be handled from the Universe itself.

Also, AtomSelection getitem no longer returns an AtomGroup, but instead
gives back the definition of the selection (that which was stored). It
is possible to store only strings or numpy arrays of atom indices,
or tuples/lists of these in any combination. One cannot directly store
AtomGroups, but can store atom indices with `AG.indices`.

Getting back an AtomGroup from the stored definition requires using the
`create` method. This makes it clear that a new AtomGroup is being
generated, and that there is a cost to this.
  • Loading branch information
@dotsdl
dotsdl committed Apr 4, 2016
1 parent cd40641 commit 2a5e6a6
Show file tree
Hide file tree
Showing 4 changed files with 205 additions and 149 deletions.
184 changes: 99 additions & 85 deletions src/mdsynthesis/limbs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ class UniverseDefinition(Limb):
this universe directly available via ``Sim.atomselections``.
"""
_name = 'udef'
_name = 'universedef'
_filepaths = ['abs', 'rel']

def __init__(self, treant):
Expand Down Expand Up @@ -114,7 +114,8 @@ def topology(self):
"""
with self._treant._read:
topstate = self._treant._state['mdsynthesis']['udef']['topology']
mdsdict = self._treant._state['mdsynthesis']
topstate = mdsdict['universedef']['topology']

if not topstate:
return None
Expand All @@ -140,10 +141,11 @@ def topology(self, path):

def _set_topology(self, path):
with self._treant._write:
topstate = self._treant._state['mdsynthesis']['udef']['topology']
mdsdict = self._treant._state['mdsynthesis']
topstate = mdsdict['universedef']['topology']

if path is None:
self._treant._state['mdsynthesis']['udef']['topology'] = dict()
mdsdict['universedef']['topology'] = dict()
else:
topstate['abspath'] = os.path.abspath(path)
topstate['relpath'] = os.path.relpath(path,
Expand All @@ -161,13 +163,14 @@ def trajectory(self):
"""
with self._treant._read:
traj = self._treant._state['mdsynthesis']['udef']['trajectory']
mdsdict = self._treant._state['mdsynthesis']
traj = mdsdict['universedef']['trajectory']
if not traj:
return None
elif len(traj) == 1:
return traj[0][0]
else:
return [t[0] for t in traj]
return tuple([t[0] for t in traj])

@trajectory.setter
def trajectory(self, path):
Expand All @@ -191,9 +194,9 @@ def trajectory(self, path):

def _set_trajectory(self, trajs):
with self._treant._write:
self._treant._state['mdsynthesis']['udef']['trajectory'] = []
mdsdict = self._treant._state['mdsynthesis']
trajstate = mdsdict['udef']['trajectory']
mdsdict['universedef']['trajectory'] = []
trajstate = mdsdict['universedef']['trajectory']

for traj in trajs:
trajstate.append(
Expand All @@ -205,7 +208,7 @@ def _apply_resnums(self):
"""
with self._treant._read:
simdict = self._treant._state['mdsynthesis']['udef']
simdict = self._treant._state['mdsynthesis']['universedef']
try:
resnums = simdict['resnums']
except KeyError:
Expand All @@ -232,7 +235,7 @@ def _set_resnums(self, resnums):
resnum definition
"""
with self._treant._write:
simdict = self._treant._state['mdsynthesis']['udef']
simdict = self._treant._state['mdsynthesis']['universedef']
if resnums is None:
simdict['resnums'] = None
else:
Expand Down Expand Up @@ -263,11 +266,12 @@ def _define(self, pathtype='abs'):
"""
with self._treant._read:
top = self._treant._state['mdsynthesis']['udef']['topology']
mdsdict = self._treant._state['mdsynthesis']
top = mdsdict['universedef']['topology']
outtop = {'abs': top['abspath'],
'rel': top['relpath']}

trajs = self._treant._state['mdsynthesis']['udef']['trajectory']
trajs = mdsdict['universedef']['trajectory']
outtraj = {key: [] for key in self._filepaths}

for traj in trajs:
Expand All @@ -288,23 +292,26 @@ def kwargs(self):
"""
with self._treant._read:
mdsdict = self._treant._state['mdsynthesis']
return mdsdict['udef']['kwargs']
return mdsdict['universedef']['kwargs']

@kwargs.setter
def kwargs(self, kwargs):
if not isinstance(kwargs, dict):
raise TypeError("Must be a dictionary")

# check that values are serializable
for key, value in kwargs.items():
if not (isinstance(value, (string_types, bool, int, float)) or
value is None):
raise ValueError("Cannot store keyword '{}' for Universe; "
"value must be a string, bool, int, float, "
"or None, not '{}'".format(key, type(value)))
if kwargs is None:
pass
elif isinstance(kwargs, dict):
# check that values are serializable
for key, value in kwargs.items():
if not (isinstance(value, (string_types, bool, int, float)) or
value is None):
raise ValueError("Cannot store keyword '{}' for Universe; "
"value must be a string, bool, int, "
"float, or ``None``, "
"not '{}'".format(key, type(value)))
else:
raise TypeError("Must be a dictionary or ``None``")

with self._treant._write:
simdict = self._treant._state['mdsynthesis']['udef']
simdict = self._treant._state['mdsynthesis']['universedef']
simdict['kwargs'] = kwargs


Expand Down Expand Up @@ -340,53 +347,29 @@ def __init__(self, treant):

def __repr__(self):
return "<AtomSelections({})>".format(
{x: self.define(x) for x in self.keys()})

def __str__(self):
selections = self.keys()
agg = "Selections"
majsep = "="
minsep = "-"
subsep = "| "
seplength = len(agg)

if not self._treant._uname:
out = "No universe attached; no Selections to show"
elif not selections:
out = "No selections for universe '{}'".format(
self._treant._uname)
else:
out = agg + '\n'
out = out + majsep * seplength + '\n'
for selection in selections:
out = out + "'{}'\n".format(selection)
for item in self.define(selection):
out = out + subsep + "'{}'\n".format(item)
out = out + minsep * seplength + '\n'

return out
{x: self.get(x) for x in self.keys()})

def __getitem__(self, handle):
"""Get selection as an AtomGroup for given handle and the universe.
"""Get selection for given handle.
Parameters
----------
handle : str
Name of selection to return as an AtomGroup.
Name of selection to return.
Returns
-------
AtomGroup
The named selection as an AtomGroup of the universe.
selection : str, list, array_like
The named selection definition.
"""
return self._asAtomGroup(handle)
return self.get(handle)

def __setitem__(self, handle, selection):
"""Selection for the given handle and the active universe.
"""Selection for the given handle.
"""
if isinstance(selection, (string_types, AtomGroup)):
if isinstance(selection, (string_types, np.ndarray)):
selection = [selection]
self.add(handle, *selection)

Expand All @@ -406,32 +389,36 @@ def add(self, handle, *selection):
data. It is useful to store AtomGroup selections for later use, since
they can be complex and atom order may matter.
If a selection with the given *handle* already exists, it is replaced.
If a selection with the given `handle` already exists, it is replaced.
Parameters
----------
handle : str
Name to use for the selection.
selection : str, AtomGroup
Selection string or AtomGroup; multiple selections may be given and
their order will be preserved, which is useful for e.g. structural
alignments.
"""
# Conversion function, leave strings alone,
# turn AtomGroups into their indices
def conv(x):
return x if isinstance(x, string_types) else x.indices

selection = map(conv, selection)
Selection string or AtomGroup indices; multiple selections may be
given and their order will be preserved, which is useful for e.g.
structural alignments.
with self._treant._write:
"""
if len(selection) == 1:
sel = selection[0]
if isinstance(sel, np.ndarray):
outsel = sel.tolist()
elif isinstance(sel, string_types):
outsel = sel
else:
outsel = list()
for sel in selection:
if isinstance(sel, np.ndarray):
outsel.append(sel.tolist())
elif isinstance(sel, string_types):
outsel.append(sel)
else:
raise ValueError("Selections must be strings, arrays of "
"atom indices, or tuples/lists of these.")

with self._treant._write:
seldict = self._treant._state['mdsynthesis']['atomselections']
seldict[handle] = outsel

Expand Down Expand Up @@ -462,8 +449,8 @@ def keys(self):
seldict = self._treant._state['mdsynthesis']['atomselections']
return seldict.keys()

def _asAtomGroup(self, handle):
"""Get AtomGroup from universe from the given named selection.
def create(self, handle):
"""Generate AtomGroup from universe from the given named selection.
If named selection doesn't exist, :exc:`KeyError` raised.
Expand All @@ -476,29 +463,39 @@ def _asAtomGroup(self, handle):
-------
AtomGroup
The named selection as an AtomGroup of the universe.
"""
sel = self.define(handle)
sel = self.get(handle)

# Selections might be either
# - a single string
# - a single list of indices
# - a list of strings
# - a list of indices

ag = None
for item in sel:
if isinstance(item, string_types):
if ag:
ag += self._treant.universe.select_atoms(item)
else:
ag = self._treant.universe.select_atoms(item)
else:
if ag:
ag += self._treant.universe.atoms[item]
if isinstance(sel, string_types):
# if we have a single string
ag = self._treant.universe.select_atoms(sel)
elif all([isinstance(i, int) for i in sel]):
# if we have a single array_like of indices
ag = self._treant.universe.atoms[sel]
else:
ag = None
for item in sel:
if isinstance(item, string_types):
if ag:
ag += self._treant.universe.select_atoms(item)
else:
ag = self._treant.universe.select_atoms(item)
else:
ag = self._treant.universe.atoms[item]
if ag:
ag += self._treant.universe.atoms[item]
else:
ag = self._treant.universe.atoms[item]

return ag

def define(self, handle):
def get(self, handle):
"""Get selection definition for given handle.
If named selection doesn't exist, :exc:`KeyError` raised.
Expand All @@ -512,13 +509,30 @@ def define(self, handle):
-------
definition : list
list of strings defining the atom selection
"""
with self._treant._read:
seldict = self._treant._state['mdsynthesis']['atomselections']

try:
selstring = seldict[handle]
seldef = seldict[handle]
except KeyError:
raise KeyError("No such selection '{}'".format(handle))

return selstring
if isinstance(seldef, string_types):
# if we have a single string
out = seldef
elif all([isinstance(i, int) for i in seldef]):
# if we have a single list of indices
out = np.array(seldef)
else:
out = []
for item in seldef:
if isinstance(item, string_types):
out.append(item)
else:
out.append(np.array(item))

out = tuple(out)

return out
18 changes: 9 additions & 9 deletions src/mdsynthesis/tests/test_filesystem.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,8 @@ def sim_external(self, tmpdir):
py.path.local(GRO).copy(GRO_t)
py.path.local(XTC).copy(XTC_t)

s.udef.topology = GRO_t.strpath
s.udef.trajectory = XTC_t.strpath
s.universedef.topology = GRO_t.strpath
s.universedef.trajectory = XTC_t.strpath
return s

@pytest.fixture
Expand All @@ -45,8 +45,8 @@ def sim_internal(self, tmpdir):
py.path.local(GRO).copy(GRO_t)
py.path.local(XTC).copy(XTC_t)

s.udef.topology = GRO_t.strpath
s.udef.trajectory = XTC_t.strpath
s.universedef.topology = GRO_t.strpath
s.universedef.trajectory = XTC_t.strpath
return s

def test_move_Sim_external(self, sim_external, tmpdir):
Expand All @@ -57,24 +57,24 @@ def test_move_Sim_external(self, sim_external, tmpdir):
assert isinstance(sim_external.universe, MDAnalysis.Universe)

sim = sim_external
assert sim.universe.filename == sim.udef.topology
assert sim.universe.trajectory.filename == sim.udef.trajectory
assert sim.universe.filename == sim.universedef.topology
assert sim.universe.trajectory.filename == sim.universedef.trajectory

def test_move_Sim_internal(self, sim_internal, tmpdir):
sim_internal.location = tmpdir.mkdir('test').strpath
assert isinstance(sim_internal.universe, MDAnalysis.Universe)

sim = sim_internal
assert sim.universe.filename == sim.udef.topology
assert sim.universe.trajectory.filename == sim.udef.trajectory
assert sim.universe.filename == sim.universedef.topology
assert sim.universe.trajectory.filename == sim.universedef.trajectory

def test_move_Sim_internal_copy(self, sim_internal, tmpdir):
"""We move the whole Sim, including its internal trajectory files, and
we put a copy of one trajectory file where the internal ones used to
be.
"""
grofile = tmpdir.join(sim_internal.name, 'sub', self.GRO)
assert sim_internal.udef.topology == grofile.strpath
assert sim_internal.universedef.topology == grofile.strpath

sim_internal.location = tmpdir.mkdir('test').strpath
tmpdir.join('test', sim_internal.name, 'sub', self.GRO).copy(
Expand Down
Loading

0 comments on commit 2a5e6a6

Please sign in to comment.