Add support for tallies

montepy/data_inputs/__init__.py

@@ -1,6 +1,7 @@
 # Copyright 2024, Battelle Energy Alliance, LLC All Rights Reserved.
 __name__ = "montepy.data_inputs"
 from .data_input import DataInput
+from .data_parser import parse_data
 from .material import Material
+from .tally import Tally
 from .thermal_scattering import ThermalScatteringLaw
-from .data_parser import parse_data

montepy/data_inputs/tally.py

@@ -0,0 +1,122 @@
+# Copyright 2024, Battelle Energy Alliance, LLC All Rights Reserved.
+import copy
+
+import montepy
+from montepy.cells import Cells
+from montepy.data_inputs.data_input import DataInputAbstract
+from montepy.data_inputs.tally_type import TallyType
+from montepy.input_parser.tally_parser import TallyParser
+from montepy.input_parser import syntax_node
+from montepy.numbered_mcnp_object import Numbered_MCNP_Object
+from montepy.utilities import *
+
+_TALLY_TYPE_MODULUS = 10
+
+
+def _number_validator(self, number):
+    if number <= 0:
+        raise ValueError("number must be > 0")
+    if number % _TALL_TYPE_MODULUS != self._type.value:
+        raise ValueError(f"Tally Type cannot be changed.")
+    if self._problem:
+        self._problem.tallies.check_number(number)
+
+
+class Tally(DataInputAbstract, Numbered_MCNP_Object):
+    """ """
+
+    # todo type enforcement
+    _parser = TallyParser()
+
+    __slots__ = {"_groups", "_type", "_number", "_old_number", "_include_total"}
+
+    def __init__(self, input=None):
+        self._cells = Cells()
+        self._old_number = None
+        self._number = self._generate_default_node(int, -1)
+        super().__init__(input)
+        if input:
+            num = self._input_number
+            self._old_number = copy.deepcopy(num)
+            self._number = num
+            try:
+                tally_type = TallyType(self.number % _TALLY_TYPE_MODULUS)
+            except ValueError as e:
+                raise MalformedInputEror(input, f"Tally Type provided not allowed: {e}")
+            groups, has_total = TallyGroup.parse_tally_specification(
+                self._tree["tally"]
+            )
+            self._groups = groups
+            self._include_total = has_total
+
+    @staticmethod
+    def _class_prefix():
+        return "f"
+
+    @staticmethod
+    def _has_number():
+        return True
+
+    @staticmethod
+    def _has_classifier():
+        return 2
+
+    @make_prop_val_node("_old_number")
+    def old_number(self):
+        """
+        The material number that was used in the read file
+
+        :rtype: int
+        """
+        pass
+
+    @make_prop_val_node("_number", int, validator=_number_validator)
+    def number(self):
+        """
+        The number to use to identify the material by
+
+        :rtype: int
+        """
+        pass
+
+
+class TallyGroup:
+    __slots__ = {"_cells", "_old_numbers"}
+
+    def __init__(self, cells=None, nodes=None):
+        self._cells = montepy.cells.Cells()
+        self._old_numbers = []
+
+    @staticmethod
+    def parse_tally_specification(tally_spec):
+        # TODO type enforcement
+        ret = []
+        in_parens = False
+        buff = None
+        has_total = False
+        for node in tally_spec:
+            # TODO handle total
+            if in_parens:
+                if node.value == ")":
+                    in_parens = False
+                    buff._append_node(node)
+                    ret.append(buff)
+                    buff = None
+                else:
+                    buff._append_node(node)
+            else:
+                if node.value == "(":
+                    in_parens = True
+                    buff = TallyGroup()
+                    buff._append_node(node)
+                else:
+                    ret.append(TallyGroup(nodes=[node]))
+        return (ret, has_total)
+
+    def _append_node(self, node):
+        if not isinstance(node, syntax_node.ValueNode):
+            raise ValueError(f"Can only append ValueNode. {node} given")
+        self._old_numbers.append(node)
+
+    def append(self, cell):
+        self._cells.append(cell)

montepy/data_inputs/tally_type.py

@@ -0,0 +1,16 @@
+# Copyright 2024, Battelle Energy Alliance, LLC All Rights Reserved.
+
+from enum import unique, Enum
+
+
+@unique
+class TallyType(Enum):
+    """ """
+
+    CURRENT = 1
+    SURFACE_FLUX = 2
+    CELL_FLUX = 4
+    DETECTOR = 5
+    ENERGY_DEPOSITION = 6
+    FISSION_ENERGY_DEPOSITION = 7
+    ENERGY_DETECTOR_PULSE = 8

montepy/input_parser/cell_parser.py

@@ -197,15 +197,15 @@ def geometry_factory(self, p):
     def number_sequence(self, p):
         if isinstance(p[0], str):
             sequence = syntax_node.ListNode("parenthetical statement")
-            sequence.append(p[0])
+            sequence.append(syntax_node.ValueNode(p[0], str))
         else:
             sequence = p[0]
         for node in list(p)[1:]:
             if isinstance(node, syntax_node.ListNode):
                 for val in node.nodes:
                     sequence.append(val)
             elif isinstance(node, str):
-                sequence.append(syntax_node.PaddingNode(node))
+                sequence.append(syntax_node.ValueNode(node, str))
             else:
                 sequence.append(node)
         return sequence

montepy/input_parser/tally_parser.py

@@ -20,7 +20,7 @@ def tally(self, p):
         ret = {}
         for key, node in p.introduction.nodes.items():
             ret[key] = node
-        ret["tally"] = p.tally_specification
+        ret["tally"] = p.tally_specification["tally"]
         return syntax_node.SyntaxNode("data", ret)
 
     @_("tally_numbers", "tally_numbers end_phrase")
@@ -34,6 +34,11 @@ def tally_specification(self, p):
             "tally list", {"tally": p.tally_numbers, "end": text}
         )
 
+    @_('"("', '"(" padding', '")"', '")" padding')
+    def paren_phrase(self, p):
+        """ """
+        return self._flush_phrase(p, str)
+
     @_("PARTICLE", "PARTICLE padding")
     def end_phrase(self, p):
         """
@@ -48,31 +53,23 @@ def end_phrase(self, p):
         "tally_numbers tally_numbers",
         "number_sequence",
         "tally_group",
-        "tally_numbers padding",
     )
     def tally_numbers(self, p):
-        if hasattr(p, "tally_numbers"):
-            ret = p.tally_numbers
-            ret.nodes["right"] += p.padding
-            return ret
         if hasattr(p, "tally_numbers1"):
-            return syntax_node.SyntaxNode("tally tree", {"left": p[0], "right": p[1]})
+            ret = p.tally_numbers1
+            for node in p.tally_numbers2.nodes:
+                ret.append(node)
+            return ret
         else:
-            left = syntax_node.PaddingNode(None)
-            right = syntax_node.PaddingNode(None)
-            return syntax_node.SyntaxNode(
-                "tally set", {"left": left, "tally": p[0], "right": right}
-            )
+            return p[0]
 
     @_(
-        '"(" number_sequence ")"',
-        '"(" padding number_sequence ")"',
+        "paren_phrase number_sequence paren_phrase",
     )
     def tally_group(self, p):
-        left = syntax_node.PaddingNode(p[0])
-        if hasattr(p, "padding"):
-            left.append(p.padding)
-        right = syntax_node.PaddingNode(p[-1])
-        return syntax_node.SyntaxNode(
-            "tally set", {"left": left, "tally": p.number_sequence, "right": right}
-        )
+        ret = syntax_node.ListNode()
+        ret.append(p[0])
+        for node in p.number_sequence.nodes:
+            ret.append(node)
+        ret.append(p[2])
+        return ret

montepy/mcnp_problem.py

@@ -10,6 +10,7 @@
 from montepy.constants import DEFAULT_VERSION
 from montepy.materials import Materials
 from montepy.surfaces import surface_builder
+from montepy.tallies import Tallies
 from montepy.surface_collection import Surfaces
 from montepy.data_inputs import Material, parse_data
 from montepy.input_parser import input_syntax_reader, block_type, mcnp_input
@@ -37,6 +38,7 @@ def __init__(self, file_name):
         self._surfaces = Surfaces(problem=self)
         self._universes = Universes(problem=self)
         self._transforms = Transforms(problem=self)
+        self._tallies = Tallies(problem=self)
         self._data_inputs = []
         self._materials = Materials(problem=self)
         self._mcnp_version = DEFAULT_VERSION

montepy/surfaces/surface_type.py

@@ -2,7 +2,7 @@
 from enum import unique, Enum
 
 
-# @unique
+@unique
 class SurfaceType(str, Enum):
     """
     An enumeration of the surface types allowed.

montepy/tallies.py

@@ -0,0 +1,15 @@
+# Copyright 2024, Battelle Energy Alliance, LLC All Rights Reserved.
+import montepy
+from montepy.numbered_object_collection import NumberedObjectCollection
+
+
+class Tallies(NumberedObjectCollection):
+    """
+    A container of multiple :class:`~montepy.data_inputs.tally.Tally` instances.
+
+    :param objects: the list of tallies to start with if needed
+    :type objects: list
+    """
+
+    def __init__(self, objects=None, problem=None):
+        super().__init__(montepy.data_inputs.tally.Tally, objects, problem)

tests/test_tally.py

@@ -17,8 +17,8 @@ def test_parsing_tally_groups(self):
             "F4:n 1 2 3",
             "F4:n (1 3i 5) (7 8 9) T",
             "f4:n (1 3i 5) (7 8 9)",
-            "F7 (1 3i 5) (7 8 9)",
-            "F7 (1 3i 5) (7 8 9) ",
+            "F7:n (1 3i 5) (7 8 9)",
+            "F7:n (1 3i 5) (7 8 9) ",
         ]:
             print(line)
             input = Input([line], BlockType.DATA)