diff --git a/src/pymcnp/files/ptrac/ptrac_filtered.py b/src/pymcnp/files/ptrac/ptrac_filtered.py new file mode 100644 index 0000000..d2a20fa --- /dev/null +++ b/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 diff --git a/src/pymcnp/files/ptrac/ptrac_processed.py b/src/pymcnp/files/ptrac/ptrac_processed.py new file mode 100644 index 0000000..699a9fc --- /dev/null +++ b/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()