files/ptrac: add processor class

src/pymcnp/files/ptrac/ptrac_filtered.py

@@ -0,0 +1,124 @@
+"""
+Contains the ``FilteredPtrac`` for filtering PTRAC file.
+"""
+
+from .event import Event, EventType
+from .ptrac import Ptrac
+
+
+class FilteredPtrac:
+    """
+    Filters PTRAC files.
+
+    ``FilteredPtrac`` provides helper methods and stubs for processing PTRAC.
+    Subclasses override its methods to specify how ``run`` processes PTRAC.
+
+    Attributes:
+        ptrac: ``Ptrac`` to filter.
+    """
+
+    def __init__(self, ptrac: Ptrac):
+        """
+        Initializes ``FilteredPtrac``.
+
+        Parameters:
+            ptrac: ``Ptrac`` to filter.
+        """
+
+        self.ptrac: Ptrac = ptrac
+
+    @staticmethod
+    def filter_source(event: Event) -> bool:
+        """
+        Runs when ``run`` filters PTRAC source events.
+
+        Returns:
+            True/False if the given event should/shouldn't be kept.
+        """
+
+        raise NotImplementedError
+
+    @staticmethod
+    def filter_bank(event: Event) -> bool:
+        """
+        Runs when ``run`` filters PTRAC bank events.
+
+        Returns:
+            True/False if the given event should/shouldn't be kept.
+        """
+
+        raise NotImplementedError
+
+    @staticmethod
+    def filter_surface(event: Event) -> bool:
+        """
+        Runs when ``run`` filters PTRAC surface events.
+
+        Returns:
+            True/False if the given event should/shouldn't be kept.
+        """
+
+        raise NotImplementedError
+
+    @staticmethod
+    def filter_collision(event: Event) -> bool:
+        """
+        Runs when ``run`` filters PTRAC collision events.
+
+        Returns:
+            True/False if the given event should/shouldn't be kept.
+        """
+
+        raise NotImplementedError
+
+    @staticmethod
+    def filter_terminal(event: Event) -> bool:
+        """
+        Runs when ``run`` filters PTRAC terminal events.
+
+        Returns:
+            True/False if the given event should/shouldn't be kept.
+        """
+
+        raise NotImplementedError
+
+    @staticmethod
+    def filter_flag(event: Event) -> bool:
+        """
+        Runs when ``run`` filters PTRAC flag events.
+
+        Returns:
+            True/False if the given event should/shouldn't be kept.
+        """
+
+        raise NotImplementedError
+
+    def run(self):
+        """
+        Filters PTRAC.
+        """
+
+        def histories():
+            for history in self.ptrac.history:
+                for event in self.ptracFile.history.events:
+                    match event.event_type:
+                        case EventType.SOURCE:
+                            _filter = self.filter_source
+                        case EventType.SURFACE:
+                            _filter = self.filter_surface
+                        case EventType.COLLISION:
+                            _filter = self.filter_collision
+                        case EventType.TERMINAL:
+                            _filter = self.filter_terminal
+                        case EventType.FLAG:
+                            _filter = self.filter_flag
+                        case _:
+                            _filter = self.filter_bank
+
+                    if _filter(event):
+                        yield event
+                    else:
+                        continue
+            return
+
+        return histories

src/pymcnp/files/ptrac/ptrac_processed.py

@@ -0,0 +1,109 @@
+"""
+Contains the ``ProcessedPtrac`` class for processing PTRAC.
+"""
+
+from .event import Event, EventType
+from .ptrac import Ptrac
+
+
+class ProcessedPtrac:
+    """
+    Processes PTRAC.
+
+    ``ProcessedPtrac`` provides helper methods and stubs for processing PTRAC.
+    Subclasses override its methods to specify how ``run`` processes PTRAC.
+
+    Attributes:
+        ptrac: ``Ptrac`` to process.
+    """
+
+    def __init__(self, ptrac: Ptrac):
+        """
+        Initializes ``ProcessedPtrac``.
+
+        Parameters:
+            ptrac: ``Ptrac`` to process.
+        """
+
+        self.ptrac: Ptrac = ptrac
+
+    def prehook(self):
+        """
+        Runs before ``run``.
+        """
+
+        pass
+
+    def posthook(self):
+        """
+        Runs after ``run``.
+        """
+
+        pass
+
+    def process_source(event: Event):
+        """
+        Runs when ``run`` processes PTRAC source events.
+        """
+
+        raise NotImplementedError
+
+    def process_bank(event: Event):
+        """
+        Runs when ``run`` processes PTRAC bank events.
+        """
+
+        raise NotImplementedError
+
+    def process_surface(event: Event):
+        """
+        Runs when ``run`` processes PTRAC surface events.
+        """
+
+        raise NotImplementedError
+
+    def process_collision(event: Event):
+        """
+        Runs when ``run`` processes PTRAC collision events.
+        """
+
+        raise NotImplementedError
+
+    def process_terminal(event: Event):
+        """
+        Runs when ``run`` processes PTRAC termianl events.
+        """
+
+        raise NotImplementedError
+
+    def process_flag(event: Event):
+        """
+        Runs when ``run`` processes PTRAC source events.
+        """
+
+        raise NotImplementedError
+
+    def run(self):
+        """
+        Processes PTRAC.
+        """
+
+        self.prehook()
+
+        for history in self.ptrac.history:
+            for event in history:
+                match event.event_type:
+                    case EventType.SOURCE:
+                        self.process_source(event)
+                    case EventType.SURFACE:
+                        self.process_surface(event)
+                    case EventType.COLLISION:
+                        self.process_collision(event)
+                    case EventType.TERMINAL:
+                        self.process_terminal(event)
+                    case EventType.FLAG:
+                        self.process_flag(event)
+                    case _:
+                        self.process_bank(event)
+
+        self.posthook()