Allow setting multi-universe fills

doc/source/changelog.rst

@@ -38,6 +38,7 @@ MontePy Changelog
 * Fixed bug that wouldn't allow cloning most surfaces (:issue:`704`).
 * Fixed bug that crashed when some cells were not assigned to any universes (:issue:`705`).
 * Fixed bug where setting ``surf.is_reflecting`` to ``False`` did not always get exported properly (:issue:`709`).
+* Fixed bug where setting multiple universes for a cell fill not being properly exported (:issue:`714`).
 
 **Breaking Changes**
 

montepy/data_inputs/fill.py

@@ -1,5 +1,8 @@
-# Copyright 2024, Battelle Energy Alliance, LLC All Rights Reserved.
+# Copyright 2024 - 2025, Battelle Energy Alliance, LLC All Rights Reserved.
 import itertools as it
+from numbers import Integral, Real
+import numpy as np
+
 from montepy.data_inputs.cell_modifier import CellModifierInput, InitInput
 from montepy.data_inputs.transform import Transform
 from montepy.errors import *
@@ -9,7 +12,14 @@
 from montepy.mcnp_object import MCNP_Object
 from montepy.universe import Universe
 from montepy.utilities import *
-import numpy as np
+
+
+def _verify_3d_index(self, indices):
+    for index in indices:
+        if not isinstance(index, Integral):
+            raise TypeError(f"Index values for fill must be an int. {index} given.")
+    if len(indices) != 3:
+        raise ValueError(f"3 values must be given for fill. {indices} given")
 
 
 class Fill(CellModifierInput):
@@ -45,6 +55,8 @@ def __init__(
         self._hidden_transform = None
         self._old_transform_number = None
         self._multi_universe = False
+        self._min_index = None
+        self._max_index = None
         super().__init__(input, in_cell_block, key, value)
         if self.in_cell_block:
             if key:
@@ -71,7 +83,7 @@ def _generate_default_cell_tree(self):
         list_node.append(self._generate_default_node(float, None))
         classifier = syntax_node.ClassifierNode()
         classifier.prefix = self._generate_default_node(
-            str, self._class_prefix().upper(), None
+            str, self._class_prefix().upper(), None, never_pad=True
         )
         self._tree = syntax_node.SyntaxNode(
             "fill",
@@ -259,17 +271,35 @@ def universes(self):
     def universes(self, value):
         if not isinstance(value, (np.ndarray, type(None))):
             raise TypeError(f"Universes must be set to an array. {value} given.")
-        if not self.multiple_universes:
+        if value.ndim != 3:
             raise ValueError(
-                "Multiple universes can only be set when multiple_universes is True."
+                f"3D array must be given for fill.universes. Array of shape: {value.shape} given."
+            )
+
+        def is_universes(array):
+            type_checker = lambda x: isinstance(x, (Universe, type(None)))
+            return map(type_checker, array.flat)
+
+        if value.dtype != np.object_ or not all(is_universes(value)):
+            raise TypeError(
+                f"All values in array must be a Universe (or None). {value} given."
             )
+        self.multiple_universes = True
+        if self.min_index is None:
+            self.min_index = np.array([0] * 3)
+        self.max_index = self.min_index + np.array(value.shape) - 1
         self._universes = value
 
     @universes.deleter
     def universes(self):
         self._universes = None
 
-    @property
+    @make_prop_pointer(
+        "_min_index",
+        (list, np.ndarray),
+        validator=_verify_3d_index,
+        deletable=True,
+    )
     def min_index(self):
         """The minimum indices of the matrix in each dimension.
 
@@ -280,9 +310,14 @@ def min_index(self):
         :class:`numpy.ndarry`
             the minimum indices of the matrix for complex fills
         """
-        return self._min_index
+        pass
 
-    @property
+    @make_prop_pointer(
+        "_max_index",
+        (list, np.ndarray),
+        validator=_verify_3d_index,
+        deletable=True,
+    )
     def max_index(self):
         """The maximum indices of the matrix in each dimension.
 
@@ -293,7 +328,7 @@ def max_index(self):
         :class:`numpy.ndarry`
             the maximum indices of the matrix for complex fills
         """
-        return self._max_index
+        pass
 
     @property
     def multiple_universes(self):
@@ -500,13 +535,15 @@ def __repr__(self):
         )
 
     def _update_cell_values(self):
-        # Todo update matrix fills
         new_vals = list(self._tree["data"])
         if self.transform and self.transform.is_in_degrees:
             self._tree["classifier"].modifier = "*"
         else:
             self._tree["classifier"].modifier = None
         new_vals = self._update_cell_universes(new_vals)
+        self._update_cell_transform_values(new_vals)
+
+    def _update_cell_transform_values(self, new_vals):
         if self.transform is None:
             try:
                 values = [val.value for val in self._tree["data"]]
@@ -579,7 +616,29 @@ def _value_node_generator():
         for universe, value in zip(payload, value_nodes):
             value.value = universe
             buffer.append(value)
-        buffer = new_vals[:start_matrix] + buffer
+        buffer = self._update_multi_index_limits(new_vals[:start_matrix]) + buffer
         if start_transform:
             buffer += new_vals[start_transform:]
         return buffer
+
+    def _update_multi_index_limits(self, new_vals):
+        if not self.multiple_universes and ":" not in new_vals:
+            return new_vals
+        if ":" not in new_vals:
+            for min_idx, max_idx in zip(self.min_index, self.max_index):
+                new_vals.extend(
+                    [
+                        self._generate_default_node(int, str(min_idx), padding=None),
+                        syntax_node.PaddingNode(":"),
+                        self._generate_default_node(int, str(max_idx), padding=" "),
+                    ]
+                )
+            return new_vals
+        vals_iter = iter(new_vals)
+        for min_idx, max_idx in zip(self.min_index, self.max_index):
+            min_val = next(vals_iter)
+            min_val.value = min_idx
+            next(vals_iter)
+            max_val = next(vals_iter)
+            max_val.value = max_idx
+        return new_vals

tests/test_universe.py

@@ -1,4 +1,6 @@
-# Copyright 2024, Battelle Energy Alliance, LLC All Rights Reserved.
+# Copyright 2024 - 2025, Battelle Energy Alliance, LLC All Rights Reserved.
+from hypothesis import given, strategies as st
+import pytest
 from unittest import TestCase
 
 import copy
@@ -141,7 +143,7 @@ def setUp(self):
         list_node = syntax_node.ListNode("numbers")
         list_node.append(syntax_node.ValueNode("1", float))
         classifier = syntax_node.ClassifierNode()
-        classifier.prefix = "lat"
+        classifier.prefix = syntax_node.ValueNode("lat", str)
         tree = syntax_node.SyntaxNode(
             "lattice",
             {
@@ -221,10 +223,12 @@ class TestFill(TestCase):
     def setUp(self):
         list_node = syntax_node.ListNode("num")
         list_node.append(syntax_node.ValueNode("5", float))
+        classifier = syntax_node.ClassifierNode()
+        classifier.prefix = syntax_node.ValueNode("fill", str)
         tree = syntax_node.SyntaxNode(
             "fill",
             {
-                "classifier": "",
+                "classifier": classifier,
                 "seperator": syntax_node.ValueNode("=", str),
                 "data": list_node,
             },
@@ -370,8 +374,10 @@ def test_fill_universes_setter(self):
         with self.assertRaises(TypeError):
             fill.universes = "hi"
         fill.multiple_universes = False
-        with self.assertRaises(ValueError):
-            fill.universes = fill_array
+        with pytest.raises(ValueError):
+            fill.universes = np.array([1, 2])
+        with pytest.raises(TypeError):
+            fill.universes = np.array([[[1]]])
 
     def test_fill_str(self):
         input = Input(["1 0 -1 fill=0:1 0:1 0:1 1 2 3 4 5 6 7 8"], BlockType.CELL)
@@ -388,3 +394,75 @@ def test_fill_merge(self):
         fill2 = Fill(card)
         with self.assertRaises(MalformedInputError):
             fill1.merge(fill2)
+
+    @given(
+        indices=st.lists(st.integers(), min_size=3, max_size=3),
+        width=st.lists(st.integers(1), min_size=3, max_size=3),
+    )
+    def test_fill_index_setter(self, indices, width):
+        fill = self.simple_fill.clone()
+        fill.multiple_universes = True
+        fill.min_index = indices
+        end = np.array(indices) + np.array(width)
+        fill.max_index = end
+        assert fill.min_index == indices
+        assert (fill.max_index == end).all()
+
+    def test_fill_index_bad_setter(self):
+        fill = self.simple_fill
+        with pytest.raises(TypeError):
+            fill.min_index = "hi"
+        with pytest.raises(TypeError):
+            fill.max_index = "hi"
+        with pytest.raises(TypeError):
+            fill.min_index = ["hi"]
+        with pytest.raises(TypeError):
+            fill.max_index = ["hi"]
+        with pytest.raises(ValueError):
+            fill.min_index = [1]
+        with pytest.raises(ValueError):
+            fill.max_index = [1]
+
+    @given(
+        universes=st.lists(st.integers(0, 1_000_000), min_size=1, max_size=10),
+        y_len=st.integers(1, 10),
+        z_len=st.integers(1, 10),
+    )
+    @pytest.mark.filterwarnings("ignore")
+    def test_fill_multi_unis(self, universes, y_len, z_len):
+        fill = self.simple_fill.clone()
+        universes = np.array([[[Universe(u) for u in universes]] * y_len] * z_len)
+        fill.multiple_universes = True
+        fill.universes = universes
+        assert (fill.universes == universes).all()
+        assert (fill.min_index == np.array([0, 0, 0])).all()
+        assert (fill.max_index == np.array(universes.shape) - np.array([1, 1, 1])).all()
+        self.verify_export(fill)
+
+    def verify_export(self, fill):
+        output = fill.format_for_mcnp_input((6, 3, 0))
+        print(output)
+        cell = montepy.Cell("1 0 -2 " + "\n".join(output))
+        new_fill = cell.fill
+        for attr in [
+            "multiple_universes",
+            "old_universe_numbers",
+            "old_universe_number",
+        ]:
+            old_val = getattr(fill, attr)
+            if "old" in attr:
+                if attr.endswith("s"):
+                    old_val = getattr(fill, "universes")
+                    if old_val is not None:
+                        numberer = np.vectorize(lambda u: u.number)
+                        old_val = numberer(old_val)
+                else:
+                    old_val = getattr(fill, "universe")
+                    if old_val is not None:
+                        old_val = old_val.number
+            new_val = getattr(new_fill, attr)
+            print(attr, old_val, new_val)
+            if isinstance(old_val, np.ndarray):
+                assert (old_val == new_val).all()
+            else:
+                assert old_val == new_val