cleanup of STAGES; several new filters

dwhswenson · May 15, 2021 · May 17, 2021 · Jun 9, 2021 · Jun 17, 2021 · Jul 5, 2021
commit 3ce4a889dc956c0824e475770d7fcc4e6f9ac8ef
diff --git a/openpathsampling/analysis/filters.py b/openpathsampling/analysis/filters.py
@@ -1,26 +1,14 @@
 import openpathsampling as paths
 import enum
 
-class Stages(enum.Enum):
-    UNKNOWN = 0
-    STEP_FILTER = 1
-    SAMPLE_SELECTOR = 2
-    SAMPLE_FILTER = 3
-    POSTPROCESSING = 4
-
-class FilterTypes(enum.Enum):
-    UNKNOWN = 0
-    STEP = 1
-    SAMPLE = 2
-
 def _default_repr(obj):
     return obj.name if obj.name is not None \
             else super(obj.__class__, obj).__repr__()
 
 class _Filter(object):
     """Abstract base class for filters.
     """
-    STAGE = Stages.UNKNOWN
+    FILTER_TYPE = None
     def condition(self, inp):
         raise NotImplementedError()
 
@@ -46,8 +34,8 @@ def __call__(self, inputs):
 
 class NegationFilter(_Filter):
     @property
-    def STAGE(self):
-        return self.filter.STAGE
+    def FILTER_TYPE(self):
+        return self.filter.FILTER_TYPE
 
     def __init__(self, filt):
         self.filter = filt
@@ -58,42 +46,34 @@ def condition(self, inp):
     def __repr__(self):
         return "~" + repr(self.filter)
 
-
-class AndFilterCombination(_Filter):
+class _FilterCombination(_Filter):
     @property
-    def STAGE(self):
-        return self.filter1.STAGE
+    def FILTER_TYPE(self):
+        retrun self.filter1.FILTER_TYPE
 
     def __init__(self, filter1, filter2):
-        if filter1.STAGE != filter2.STAGE:
+        if filter1.FILTER_TYPE is not filter2.FILTER_TYPE:
             raise TypeError(f"Cannot combine filters {filter1} and {filter2}")
         self.filter1 = filter1
         self.filter2 = filter2
 
+
+class AndFilterCombination(_FilterCombination):
     def condition(self, inp):
         return self.filter1.condition(inp) and self.filter2.condition(inp)
 
     def __repr__(self):
         return "(" + repr(self.filter1) + " & " + repr(self.filter2) + ")"
 
 
-class OrFilterCombination(_Filter):
-    @property
-    def STAGE(self):
-        return self.filter1.STAGE
-
-    def __init__(self, filter1, filter2):
-        if filter1.STAGE != filter2.STAGE:
-            raise TypeError(f"Cannot combine filters {filter1} and {filter2}")
-        self.filter1 = filter1
-        self.filter2 = filter2
-
+class OrFilterCombination(_FilterCombination):
     def condition(self, inp):
         return self.filter1.condition(inp) or self.filter2.condition(inp)
 
     def __repr__(self):
         return "(" + repr(self.filter1) + " | " + repr(self.filter2) + ")"
 
+
 class GenericFilter(_Filter):
     def __init__(self, condition, name):
         self._condition = condition
@@ -108,7 +88,7 @@ def __repr__(self):
 ### OPS STEP FILTERS #######################################################
 
 class CanonicalMover(_Filter):
-    STAGE = Stages.STEP_FILTER
+    FILTER_TYPE = paths.MCStep
     def __init__(self, mover):
         if isinstance(mover, paths.PathMover):
             self.mover = mover
@@ -136,7 +116,7 @@ def __str__(self):
 
 
 class TrialReplica(_Filter):
-    STAGE = Stages.STEP_FILTER
+    FILTER_TYPE = paths.MCStep
     def __init__(self, replica):
         self.replica = replica
 
@@ -146,7 +126,7 @@ def condition(self, step):
 
 
 class TrialEnsemble(_Filter):
-    STAGE = Stages.STEP_FILTER
+    FILTER_TYPE = paths.MCStep
     def __init__(self, ensemble):
         self.ensemble = ensemble
 
@@ -156,7 +136,7 @@ def condition(self, step):
 
 
 class StepFilter(GenericFilter):
-    STAGE = Stages.STEP_FILTER
+    FILTER_TYPE = paths.MCStep
 
 RejectedSteps = StepFilter(lambda step: not step.change.accepted,
                            name="RejectedSteps")
@@ -169,12 +149,12 @@ class StepFilter(GenericFilter):
 ### SAMPLE FILTERS #########################################################
 
 class SampleFilterCondition(GenericFilter):
-    STAGE = Stages.SAMPLE_FILTER
+    FILTER_TYPE = paths.Sample
 
 AllSamples = SampleFilterCondition(lambda sample: True, name="AllSamples")
 
 class Ensemble(_Filter):
-    STAGE = Stages.SAMPLE_FILTER
+    FILTER_TYPE = paths.Sample
     def __init__(self, ensemble):
         self.ensemble = ensemble
 
@@ -186,7 +166,7 @@ def __repr__(self):
 
 
 class Replica(_Filter):
-    STAGE = Stages.SAMPLE_FILTER
+    FILTER_TYPE = paths.Sample
     def __init__(self, replica):
         self.replica = replica
 
@@ -196,12 +176,41 @@ def condition(self, sample):
     def __repr__(self):
         return f"Replica({self.replica})"
 
-### EXTRACTOR-FILTERS ######################################################
 
-def _pass_thru(items):
-    return items
-    for item in items:
-        yield item
+class _NetworkEnsemble(_Filter):
+    FILTER_TYPE = paths.Sample
+    def __init__(self, network):
+        super().__init__()
+        self.network = network
+        self.ensembles = set(self._get_ensembles())
+
+    def _get_ensembles(self):
+        raise NotImplementedError()
+
+    def condition(self, sample):
+        return sample.ensemble in self.ensembles
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.network})"
+
+
+class SamplingEnsemble(_NetworkEnsemble):
+    def _get_ensembles(self):
+        return self.network.sampling_ensembles
+
+
+class MinusEnsemble(_NetworkEnsemble):
+    def _get_ensembles(self):
+        return self.network.minus_ensembles
+
+
+class MSOuterEnsemble(_NetworkEnsemble):
+    def _get_ensembles(self):
+        return self.network.ms_outers
+
+
+
+### EXTRACTOR-FILTERS ######################################################
 
 NOT_EXTRACTED = object()  # flag for errors on extraction
 
@@ -215,7 +224,6 @@ def _list_per_step(items):
 
 class ExtractorFilter(object):
     _default_secondary = None
-    # _default_postprocessor = _pass_thru
     def __init__(self, step_filter, extractor):
         if step_filter is None:
             step_filter = AllSteps
@@ -255,54 +263,59 @@ def __call__(self, steps, flatten=None):
             for result in finalized:
                 yield result
 
-    def __and__(self, other):
-        return _AndExtractorFilter(self, other)
-
-    def __or__(self, other):
-        return _OrExtractorFilter(self, other)
-
-    def __invert__(self):
-        return _NegationExtractorFilter(self)
-
-    def _test_compatibility(self, other):
-        value_err_str = ("Can not combine ExtractorFilters that do not "
-                         "have the same {attr}: {self_attr} != "
-                         "{other_attr}")
-        if self.extractor != other.extractor:
-            raise ValueError(value_err_str.format(
-                attr='extractor',
-                self_attr=self.extractor,
-                other_attr=other.extractor
-            ))
-        if self.secondary_filter != other.secondary_filter:
-            raise ValueError(value_err_str.format(
-                attr='secondary filter',
-                self_attr=self.secondary_filter,
-                other_attr=other.secondary_filter
-            ))
-
-class _AndExtractorFilter(ExtractorFilter):
-    def __init__(self, ef1, ef2):
-        ef1._test_compatibility(ef2)
-        step_filter = ef1.step_filter & ef2.step_filter
-        self._default_secondary = ef1._default_secondary
-        super(_AndExtractorFilter, self).__init__(step_filter,
-                                                  ef1.extractor)
-
-class _OrExtractorFilter(ExtractorFilter):
-    def __init__(self, ef1, ef2):
-        ef1._test_compatibility(ef2)
-        step_filter = ef1.step_filter | ef2.step_filter
-        self._default_secondary = ef1._default_secondary
-        super(_OrExtractorFilter, self).__init__(step_filter,
-                                                  ef1.extractor)
-
-class _NegationExtractorFilter(ExtractorFilter):
-    def __init__(self, ef):
-        self._default_secondary = ef._default_secondary
-        super(_NegationExtractorFilter, self).__init__(~ef.step_filter,
-                                                       ef.extractor)
-        self.with_filter(ef.secondary_filter)
+# For the time being, leaving all this combination-related code commented
+# out. It's enough typing to keep it around in case we decide we want it,
+# but I have the feeling right now that it would be overly complicated from
+# a user standpoint.
+#
+#    def __and__(self, other):
+#        return _AndExtractorFilter(self, other)
+#
+#    def __or__(self, other):
+#        return _OrExtractorFilter(self, other)
+#
+#    def __invert__(self):
+#        return _NegationExtractorFilter(self)
+#
+#    def _test_compatibility(self, other):
+#        value_err_str = ("Can not combine ExtractorFilters that do not "
+#                         "have the same {attr}: {self_attr} != "
+#                         "{other_attr}")
+#        if self.extractor != other.extractor:
+#            raise ValueError(value_err_str.format(
+#                attr='extractor',
+#                self_attr=self.extractor,
+#                other_attr=other.extractor
+#            ))
+#        if self.secondary_filter != other.secondary_filter:
+#            raise ValueError(value_err_str.format(
+#                attr='secondary filter',
+#                self_attr=self.secondary_filter,
+#                other_attr=other.secondary_filter
+#            ))
+#
+#class _AndExtractorFilter(ExtractorFilter):
+#    def __init__(self, ef1, ef2):
+#        ef1._test_compatibility(ef2)
+#        step_filter = ef1.step_filter & ef2.step_filter
+#        self._default_secondary = ef1._default_secondary
+#        super(_AndExtractorFilter, self).__init__(step_filter,
+#                                                  ef1.extractor)
+#
+#class _OrExtractorFilter(ExtractorFilter):
+#    def __init__(self, ef1, ef2):
+#        ef1._test_compatibility(ef2)
+#        step_filter = ef1.step_filter | ef2.step_filter
+#        self._default_secondary = ef1._default_secondary
+#        super(_OrExtractorFilter, self).__init__(step_filter,
+#                                                  ef1.extractor)
+#
+#class _NegationExtractorFilter(ExtractorFilter):
+#    def __init__(self, ef):
+#        self._default_secondary = ef._default_secondary
+#        super(_NegationExtractorFilter, self).__init__(~ef.step_filter,
+#                                                       ef.extractor)
+#        self.with_filter(ef.secondary_filter)
 
 
 class SampleExtractorFilter(ExtractorFilter):
@@ -312,7 +325,6 @@ class SampleExtractorFilter(ExtractorFilter):
 ### EXTRACTORS #############################################################
 
 class Extractor(object):
-    hook = _pass_thru
     def __init__(self, extractor, name=None,
                  returns_iterable=False,
                  extractor_filter=ExtractorFilter):
@@ -324,7 +336,7 @@ def __init__(self, extractor, name=None,
     def __call__(self, step):
         return self.extractor(step)
 
-    def using(self, step_filter=None, sample_filter=None):
+    def using(self, step_filter=None):
         return self.extractor_filter(step_filter=step_filter,
                                      extractor=self)
 
@@ -353,23 +365,36 @@ def __init__(self, extractor, name):
         )
 
 ActiveSamples = SampleExtractor(lambda step: step.active.samples,
-                               name="ActiveSamples")
+                                name="ActiveSamples")
 TrialSamples = SampleExtractor(lambda step: step.change.canonical.trials,
-                              name="TrialSamples")
+                               name="TrialSamples")
 
 
 def _get_shooting_point(step):
     details = step.change.canonical.details
     try:
         shooting_pt = details.shooting_snapshot
     except AttributeError:
-        shooting_pt = None
+        shooting_pt = NOT_EXTRACTED
     return shooting_pt
 
 ShootingSteps = StepFilter(
-   lambda step: _get_shooting_point(step) is not None,
+   lambda step: _get_shooting_point(step) is not NOT_EXTRACTED,
    name="ShootingSteps"
 )
 
 ShootingPoints = Extractor(_get_shooting_point, name="ShootingPoints")
 
+def _get_modified_shooting_point(step):
+    if not ShootingSteps.condition(step):
+        return NOT_EXTRACTED
+    details = step.change.canonical.details
+    try:
+        shooting_pt = ...
+    except AttributeError:
+        # return the original shooting point for one-way
+        shooting_pt = details.shooting_snapshot
+    return shooting_pt
+
+ModifiedShootingPoints = Extractor(_get_modified_shooting_point,
+                                   name="ModifiedShootingPoints")