Skip to content

Commit

Permalink
Allow kwargs to be passed through in universe.transfer_to_memory()
Browse files Browse the repository at this point in the history
…and `in_memory=True`` (#3938)

Fixes #3684
* pass kwargs to memoryreader on calling transfer_to_memory() and `in_memory=True`
  • Loading branch information
hmacdope authored Nov 29, 2022
1 parent 5df4a55 commit 2d98ad5
Show file tree
Hide file tree
Showing 3 changed files with 17 additions and 4 deletions.
2 changes: 2 additions & 0 deletions package/CHANGELOG
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@ The rules for this file:
* 2.4.0

Fixes
* Kwargs passed through to `MemoryReader` when using `transfer_to_memory()`
and `in_memory=True`.
* NoDataError raised on empty atomgroup provided to `DCDReader.timeseries`
was changed to ValueError (see #3901). A matching ValueError was added to
`MemoryReader.timeseries`(PR #3890).
Expand Down
11 changes: 7 additions & 4 deletions package/MDAnalysis/core/universe.py
Original file line number Diff line number Diff line change
Expand Up @@ -541,6 +541,8 @@ def load_new(self, filename, format=None, in_memory=False,
.. versionchanged:: 0.17.0
Now returns a :class:`Universe` instead of the tuple of file/array
and detected file type.
.. versionchanged:: 2.4.0
Passes through kwargs if `in_memory=True`.
"""
# filename==None happens when only a topology is provided
Expand Down Expand Up @@ -574,12 +576,12 @@ def load_new(self, filename, format=None, in_memory=False,
trj_n_atoms=self.trajectory.n_atoms))

if in_memory:
self.transfer_to_memory(step=in_memory_step)
self.transfer_to_memory(step=in_memory_step, **kwargs)

return self

def transfer_to_memory(self, start=None, stop=None, step=None,
verbose=False):
verbose=False, **kwargs):
"""Transfer the trajectory to in memory representation.
Replaces the current trajectory reader object with one of type
Expand All @@ -600,6 +602,8 @@ def transfer_to_memory(self, start=None, stop=None, step=None,
.. versionadded:: 0.16.0
.. versionchanged:: 2.4.0
Passes through kwargs to MemoryReader
"""
from ..coordinates.memory import MemoryReader

Expand Down Expand Up @@ -641,8 +645,7 @@ def transfer_to_memory(self, start=None, stop=None, step=None,
dt=self.trajectory.ts.dt * step,
filename=self.trajectory.filename,
velocities=velocities,
forces=forces,
)
forces=forces, **kwargs)

# python 2 doesn't allow an efficient splitting of kwargs in function
# argument signatures.
Expand Down
8 changes: 8 additions & 0 deletions testsuite/MDAnalysisTests/core/test_universe.py
Original file line number Diff line number Diff line change
Expand Up @@ -602,6 +602,14 @@ def test_slicing_negative_stop(self):
(3341, 78, 3),
err_msg="Unexpected shape of trajectory timeseries")

def test_transfer_to_memory_kwargs(self):
u = mda.Universe(PSF, DCD)
u.transfer_to_memory(example_kwarg=True)
assert(u.trajectory._kwargs['example_kwarg'])

def test_in_memory_kwargs(self):
u = mda.Universe(PSF, DCD, in_memory=True, example_kwarg=True)
assert(u.trajectory._kwargs['example_kwarg'])

class TestCustomReaders(object):
"""
Expand Down

0 comments on commit 2d98ad5

Please sign in to comment.