meshtal: added to_mcnp methods

Author: BitterB0NG0

src/pymcnp/Meshtal.py

@@ -1,3 +1,4 @@
+import re
 import typing
 
 from . import meshtal
@@ -14,6 +15,8 @@ class Meshtal(_object.McnpFile_):
         tallies: MESTHAL tallies.
     """
 
+    _REGEX = re.compile(rf'({meshtal.Header._REGEX.pattern})((?:{meshtal.Tally._REGEX.pattern})+)')
+
     def __init__(
         self, header: meshtal.Header, tallies: typing.Generator[meshtal.Tally, None, None]
     ):
@@ -55,15 +58,16 @@ def from_mcnp(source: str):
             ``Meshtal``.
         """
 
-        header, lines = meshtal.Header.from_mcnp(source)
+        tokens = Meshtal._REGEX.match(source)
 
-        def tallies(lines):
-            while lines:
-                tally, lines = meshtal.Tally.from_mcnp(lines, header)
-                yield tally
-            return
+        if not tokens:
+            raise errors.MeshtalError(errors.MeshtalCode.SYNTAX_MESHTAL, source)
 
-        tallies = tallies(lines)
+        header = meshtal.Header.from_mcnp(tokens[1])
+        tallies = (
+            meshtal.Tally.from_mcnp(match[0], header)
+            for match in meshtal.Tally._REGEX.finditer(tokens[12])
+        )
 
         return Meshtal(header, tallies)
 
@@ -75,4 +79,4 @@ def to_mcnp(self):
             INP for ``Meshtal``.
         """
 
-        return header.to_mcnp() + '\n'.join(tallies.to_mcnp())
+        return self.header.to_mcnp() + ' ' + '\n '.join(tally.to_mcnp() for tally in self.tallies)

src/pymcnp/meshtal/Header.py

@@ -29,7 +29,7 @@ class Header(_block.Block):
     """
 
     _REGEX = re.compile(
-        r'(.{7})version (.{6})ld=(.{10})probid =(.{20})\n\s(.{80})\n\sNumber of histories used for normalizing tallies =(.{17})\n\n\sMesh Tally Number(.{10})\n\s(.{8}) mesh tally[.]\n This mesh tally is modified by a dose response function[.]\n\n Tally bin boundaries:\n((?:    (?:X direction|Y direction|Z direction|Energy bin boundaries|Time bin boundaries):.+\n)+)\n(.+)\n([\s\S]+)'
+        r'(.{7})version (.{6})ld=(.{10})probid =(.{20})\n\s(.{80})\n\sNumber of histories used for normalizing tallies =(.{17})\n\n\sMesh Tally Number(.{10})\n\s(.{8}) mesh tally[.]\n This mesh tally is modified by a dose response function[.]\n\n Tally bin boundaries:\n((?:    .+\n)+)\n(.+)\n'
     )
 
     def __init__(
@@ -103,6 +103,7 @@ def __init__(
         self.version: typing.Final[types.String] = version
         self.ld: typing.Final[types.String] = ld
         self.probid: typing.Final[types.String] = probid
+        self.title: typing.Final[types.String] = title
         self.histories: typing.Final[types.Integer] = histories
         self.number: typing.Final[types.Integer] = number
         self.particle: typing.Final[types.String] = particle
@@ -114,7 +115,7 @@ def __init__(
         self.columns: typing.Final[types.String] = columns
 
     @staticmethod
-    def from_mcnp(source: str) -> tuple[Header, str]:
+    def from_mcnp(source: str):
         """
         Generates ``Header`` from MESHTAL.
 
@@ -157,24 +158,21 @@ def from_mcnp(source: str) -> tuple[Header, str]:
 
         columns = types.String.from_mcnp(tokens[10])
 
-        return (
-            Header(
-                code,
-                version,
-                ld,
-                probid,
-                title,
-                histories,
-                number,
-                particle,
-                bins_x,
-                bins_y,
-                bins_z,
-                bins_energy,
-                bins_time,
-                columns,
-            ),
-            tokens[11],
+        return Header(
+            code,
+            version,
+            ld,
+            probid,
+            title,
+            histories,
+            number,
+            particle,
+            bins_x,
+            bins_y,
+            bins_z,
+            bins_energy,
+            bins_time,
+            columns,
         )
 
     def to_mcnp(self):
@@ -185,4 +183,19 @@ def to_mcnp(self):
             INP for ``Header``.
         """
 
-        assert False, "I'm working on it!"
+        bins = ''
+        bins += f"    X direction:{''.join(map(str, self.bins_x))}\n" if self.bins_x else ''
+        bins += f"    Y direction:{''.join(map(str, self.bins_y))}\n" if self.bins_y else ''
+        bins += f"    Z direction:{''.join(map(str, self.bins_z))}\n" if self.bins_z else ''
+        bins += (
+            f"    Energy bin boundaries:{''.join(map(str, self.bins_energy))}\n"
+            if self.bins_energy
+            else ''
+        )
+        bins += (
+            f"    Time bin boundaries:{''.join(map(str, self.bins_time))}\n"
+            if self.bins_time
+            else ''
+        )
+
+        return f'{self.code:>7}version {self.version:>6}ld={self.ld:>10}probid ={self.probid:>20}\n {self.title:>80}\n Number of histories used for normalizing tallies ={self.histories:>17.2F}\n\n Mesh Tally Number{self.number:>10}\n {self.particle:>8} mesh tally.\n This mesh tally is modified by a dose response function.\n\n Tally bin boundaries:\n{bins}\n{self.columns}\n'

src/pymcnp/meshtal/Tally.py

@@ -21,7 +21,7 @@ class Tally(_line.Line):
         error: Tally relative error.
     """
 
-    _REGEX = re.compile(r'\s*(\S*)\s*(\S*)\s*(\S*)\s*(\S*)\s*(\S*)\s*(\S*)\s*(\S*)\n([\s\S]+)')
+    _REGEX = re.compile(r'\s(.+)(\n|\Z)')
 
     def __init__(
         self,
@@ -52,13 +52,13 @@ def __init__(
         if error is None:
             raise errors.MeshtalError(errors.MeshtalCode.SEMANTICS_LINE, error)
 
+        self.result: typing.Final[types.Real] = result
+        self.error: typing.Final[types.Real] = error
         self.x: typing.Final[types.Real] = x
         self.y: typing.Final[types.Real] = y
         self.z: typing.Final[types.Real] = z
         self.energy: typing.Final[types.Real] = energy
         self.time: typing.Final[types.Real] = time
-        self.result: typing.Final[types.Real] = result
-        self.error: typing.Final[types.Real] = error
 
     @staticmethod
     def from_mcnp(source: str, header: Header):
@@ -71,9 +71,9 @@ def from_mcnp(source: str, header: Header):
         """
 
         tokens = Tally._REGEX.match(source)
-        headings = re.split(r'\s+', header.columns)
+        headings = re.split(r'\s+', header.columns.strip())
 
-        if not tokens or (len(tokens.groups()) != len(headings) + 2):
+        if not tokens:
             raise errors.MeshtalError(errors.MeshtalCode.SYNTAX_LINE, source)
 
         x = None
@@ -84,27 +84,53 @@ def from_mcnp(source: str, header: Header):
         result = None
         error = None
 
+        tokens = re.split(r'\s+', tokens[1].strip())
+
         for i, heading in enumerate(headings):
-            if heading == 'x':
-                x = tokens[i + 1]
-            elif heading == 'y':
-                y = tokens[i + 1]
-            elif heading == 'z':
-                z = tokens[i + 1]
-            elif heading == 'energy':
-                energy = tokens[i + 1]
-            elif heading == 'time':
-                time = tokens[i + 1]
-            elif heading == 'result':
-                result = tokens[i + 1]
-            elif heading == 'rel':
-                error = tokens[i + 1]
-            elif heading == 'error':
+            if heading == 'X':
+                x = tokens[i]
+            elif heading == 'Y':
+                y = tokens[i]
+            elif heading == 'Z':
+                z = tokens[i]
+            elif heading == 'Energy':
+                energy = tokens[i]
+            elif heading == 'Time':
+                time = tokens[i]
+            elif heading == 'Result':
+                result = tokens[i]
+            elif heading == 'Rel':
+                error = tokens[i]
+            elif heading == 'Error':
                 pass
             else:
                 raise errors.MeshtalError(errors.MeshtalCode.SYNTAX_LINE, source)
 
-        return (
-            Tally(x, y, z, energy, time, result, error),
-            tokens[8],
+        return Tally(
+            result,
+            error,
+            x=x,
+            y=y,
+            z=z,
+            energy=energy,
+            time=time,
         )
+
+    def to_mcnp(self):
+        """
+        Generates MESTHAL from ``Tally``.
+
+        Returns:
+            INP for ``Tally``.
+        """
+
+        line = ''
+        line += f'{self.energy:>11}' if self.energy else ''
+        line += f'{self.time:>11}' if self.time else ''
+        line += f'{self.x:>10}' if self.x else ''
+        line += f'{self.y:>10}' if self.y else ''
+        line += f'{self.z:>10}' if self.z else ''
+        line += f'{self.result:>12}' if self.result else ''
+        line += f'{self.error:>12}' if self.error else ''
+
+        return line