From 166b5dc934eb62a12baffab5c436e8dbdd7b3c5d Mon Sep 17 00:00:00 2001 From: Alia Lescoulie Date: Thu, 23 Sep 2021 23:24:36 -0700 Subject: [PATCH 1/6] Add ensemble_wrapper function --- mdpow/analysis/ensemble.py | 31 ++++++++++++++++++++++++++++++- mdpow/tests/test_ensemble.py | 27 ++++++++++++++++++++++++++- 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/mdpow/analysis/ensemble.py b/mdpow/analysis/ensemble.py index 00dfbd27..7964c527 100644 --- a/mdpow/analysis/ensemble.py +++ b/mdpow/analysis/ensemble.py @@ -3,7 +3,7 @@ import os import errno -from typing import Optional, List +from typing import Optional, List, Union import numpy as np @@ -550,3 +550,32 @@ def check_groups_from_common_ensemble(groups: List[EnsembleAtomGroup]): from the same Ensemble.''' logger.error(msg) raise ValueError(msg) + + +def ensemble_wrapper(cls): + """ A decorator for :class:`mdanalysis.analysis.base.AnalysisBase` based classes modifying them to + accept an ensemble. + """ + class EnsembleWrapper: + def __init__(self, ensemble: Union[Ensemble, EnsembleAtomGroup], *args, **kwargs): + self._ensemble = ensemble + self._args = args + self._kwargs = kwargs + self._Analysis = cls + + def _prepare_ensemble(self): + # Defined separately so user can modify behavior + self._results_dict = {x: None for x in self._ensemble.keys()} + + def _conclude_system(self): + # Defined separately so user can modify behavior + self._results_dict[self._key] = self._SystemRun.results + + def run(self, start=0, stop=0, step=1): + self._prepare_ensemble() + for self._key in self._ensemble.keys(): + self._SystemRun = self._Analysis(self._ensemble[self._key], *self._args, **self._kwargs) + self._SystemRun.run(start=start, step=step, stop=stop) + self._conclude_system() + + return EnsembleWrapper diff --git a/mdpow/tests/test_ensemble.py b/mdpow/tests/test_ensemble.py index 9d8b2462..c047792b 100644 --- a/mdpow/tests/test_ensemble.py +++ b/mdpow/tests/test_ensemble.py @@ -13,10 +13,11 @@ import MDAnalysis as mda from MDAnalysis.exceptions import NoDataError, SelectionError +from MDAnalysis.analysis.base import AnalysisBase from gromacs.utilities import in_dir -from ..analysis.ensemble import Ensemble, EnsembleAnalysis, EnsembleAtomGroup +from ..analysis.ensemble import Ensemble, EnsembleAnalysis, EnsembleAtomGroup, ensemble_wrapper from ..analysis.dihedral import DihedralAnalysis from pkg_resources import resource_filename @@ -161,3 +162,27 @@ def test_value_error(self): dh4 = ens.select_atoms('name C4 or name C17 or name S2 or name N3') with pytest.raises(ValueError): dh_run = DihedralAnalysis([dh1, dh2, dh4, dh3]).run(start=0, stop=4, step=1) + + def test_ensemble_wrapper(self): + + class BaseTest(AnalysisBase): + def __init__(self, system: mda.Universe): + super(BaseTest, self).__init__(system.trajectory) + self.system = system + + def _prepare(self): + self._res_arr = [] + + def _single_frame(self): + self._res_arr.append(len(self.system.select_atoms('not resname SOL'))) + assert self._res_arr[-1] == 42 + + def _conclude(self): + self.results = self._res_arr + + @ensemble_wrapper + class EnsembleBaseTest(BaseTest): + pass + + Sim = Ensemble(dirname=self.tmpdir.name, solvents=['water']) + SolvCount = EnsembleBaseTest(Sim).run(stop=10) From d7e899971f9a86ebab7ac72ceaa4e6fa073e4c68 Mon Sep 17 00:00:00 2001 From: Alia Lescoulie Date: Sat, 25 Sep 2021 22:47:59 -0700 Subject: [PATCH 2/6] fix ensemble_wrapper --- mdpow/analysis/ensemble.py | 5 +++++ mdpow/tests/test_ensemble.py | 8 +++++--- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/mdpow/analysis/ensemble.py b/mdpow/analysis/ensemble.py index 7964c527..e2d81b59 100644 --- a/mdpow/analysis/ensemble.py +++ b/mdpow/analysis/ensemble.py @@ -571,11 +571,16 @@ def _conclude_system(self): # Defined separately so user can modify behavior self._results_dict[self._key] = self._SystemRun.results + def _conclude_ensemble(self): + self.results = self._results_dict + def run(self, start=0, stop=0, step=1): self._prepare_ensemble() for self._key in self._ensemble.keys(): self._SystemRun = self._Analysis(self._ensemble[self._key], *self._args, **self._kwargs) self._SystemRun.run(start=start, step=step, stop=stop) self._conclude_system() + self._conclude_ensemble() + return self return EnsembleWrapper diff --git a/mdpow/tests/test_ensemble.py b/mdpow/tests/test_ensemble.py index c047792b..f295d882 100644 --- a/mdpow/tests/test_ensemble.py +++ b/mdpow/tests/test_ensemble.py @@ -163,7 +163,7 @@ def test_value_error(self): with pytest.raises(ValueError): dh_run = DihedralAnalysis([dh1, dh2, dh4, dh3]).run(start=0, stop=4, step=1) - def test_ensemble_wrapper(self): + def test_ensemble_wrapper1(self): class BaseTest(AnalysisBase): def __init__(self, system: mda.Universe): @@ -181,8 +181,10 @@ def _conclude(self): self.results = self._res_arr @ensemble_wrapper - class EnsembleBaseTest(BaseTest): + class EnsembleTest(BaseTest): pass Sim = Ensemble(dirname=self.tmpdir.name, solvents=['water']) - SolvCount = EnsembleBaseTest(Sim).run(stop=10) + SolvCount = EnsembleTest(Sim).run(stop=10) + assert isinstance(SolvCount, EnsembleTest) + From fd3de18f729996bcb44ad19e5e3fcf979c3ff1a4 Mon Sep 17 00:00:00 2001 From: Alia Lescoulie Date: Sat, 25 Sep 2021 23:04:26 -0700 Subject: [PATCH 3/6] add docs --- mdpow/analysis/ensemble.py | 15 +++++++++++++-- mdpow/tests/test_ensemble.py | 2 +- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/mdpow/analysis/ensemble.py b/mdpow/analysis/ensemble.py index e2d81b59..9cacdd73 100644 --- a/mdpow/analysis/ensemble.py +++ b/mdpow/analysis/ensemble.py @@ -553,8 +553,19 @@ def check_groups_from_common_ensemble(groups: List[EnsembleAtomGroup]): def ensemble_wrapper(cls): - """ A decorator for :class:`mdanalysis.analysis.base.AnalysisBase` based classes modifying them to - accept an ensemble. + """A decorator for :class:`MDAnalysis.Universe ` subclasses modifying + them to accept an :class:`~mdpow.analysis.ensemble.Ensemble` or + :class:`~mdpow.analysis.ensemble.EnsembleAtomGroup`. + + .. rubric:: Example Analysis + + @ensemble_wrapper + class Example(AnalysisBase): + pass + + Ens = Ensemble(dirname='mol_dir) + ExRun = Example(Ens) + """ class EnsembleWrapper: def __init__(self, ensemble: Union[Ensemble, EnsembleAtomGroup], *args, **kwargs): diff --git a/mdpow/tests/test_ensemble.py b/mdpow/tests/test_ensemble.py index f295d882..1d9c0caf 100644 --- a/mdpow/tests/test_ensemble.py +++ b/mdpow/tests/test_ensemble.py @@ -187,4 +187,4 @@ class EnsembleTest(BaseTest): Sim = Ensemble(dirname=self.tmpdir.name, solvents=['water']) SolvCount = EnsembleTest(Sim).run(stop=10) assert isinstance(SolvCount, EnsembleTest) - + assert SolvCount.results[('water', 'VDW', '0000')] == ([42] * 10) From aa4cec142963be0e63b8896e7db68ae51c549bb6 Mon Sep 17 00:00:00 2001 From: Alia Lescoulie Date: Sat, 25 Sep 2021 23:07:13 -0700 Subject: [PATCH 4/6] add decorator docs to ensemble.txt --- doc/sphinx/source/analysis/ensemble.txt | 2 ++ 1 file changed, 2 insertions(+) diff --git a/doc/sphinx/source/analysis/ensemble.txt b/doc/sphinx/source/analysis/ensemble.txt index 50f6122f..3ff24f69 100644 --- a/doc/sphinx/source/analysis/ensemble.txt +++ b/doc/sphinx/source/analysis/ensemble.txt @@ -43,3 +43,5 @@ object when the :meth:`~mdpow.analysis.ensemble.EnsembleAtomGroup.ensemble` is r .. autoclass:: mdpow.analysis.ensemble.EnsembleAtomGroup :members: + +.. autodoc:: mdpow.analysis.ensemble_wrapper From 848a0ed547edb603dd42914bc34c7fc3b3dbf2d7 Mon Sep 17 00:00:00 2001 From: Alia Lescoulie Date: Sat, 25 Sep 2021 23:13:42 -0700 Subject: [PATCH 5/6] add decorator docs to ensemble.txt --- doc/sphinx/source/analysis/ensemble.txt | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/doc/sphinx/source/analysis/ensemble.txt b/doc/sphinx/source/analysis/ensemble.txt index 3ff24f69..345f06f2 100644 --- a/doc/sphinx/source/analysis/ensemble.txt +++ b/doc/sphinx/source/analysis/ensemble.txt @@ -44,4 +44,8 @@ object when the :meth:`~mdpow.analysis.ensemble.EnsembleAtomGroup.ensemble` is r .. autoclass:: mdpow.analysis.ensemble.EnsembleAtomGroup :members: -.. autodoc:: mdpow.analysis.ensemble_wrapper + +ensemble_wrapper Decorator +__________________________ + +.. automethod:: mdpow.analysis.ensemble_wrapper From 9afc298f3f8175d39d80a744215e889433394d6a Mon Sep 17 00:00:00 2001 From: Alia Lescoulie Date: Sat, 25 Sep 2021 23:18:13 -0700 Subject: [PATCH 6/6] add decorator docs to ensemble.txt --- doc/sphinx/source/analysis/ensemble.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/sphinx/source/analysis/ensemble.txt b/doc/sphinx/source/analysis/ensemble.txt index 345f06f2..5e9a907a 100644 --- a/doc/sphinx/source/analysis/ensemble.txt +++ b/doc/sphinx/source/analysis/ensemble.txt @@ -48,4 +48,4 @@ object when the :meth:`~mdpow.analysis.ensemble.EnsembleAtomGroup.ensemble` is r ensemble_wrapper Decorator __________________________ -.. automethod:: mdpow.analysis.ensemble_wrapper +.. autofunction:: mdpow.analysis.ensemble_wrapper