Implement _collection_ref to link objects to their NumberedObjectCollection parent

demo/0_IntroToMontePy.ipynb

@@ -228,8 +228,9 @@
    },
    "outputs": [],
    "source": [
-    "#note this install_montepy is only necessary for jupyterlite; You don't need to use it locally\n",
+    "# note this install_montepy is only necessary for jupyterlite; You don't need to use it locally\n",
     "from _config import install_montepy\n",
+    "\n",
     "install_montepy()\n",
     "\n",
     "import montepy\n",
@@ -428,7 +429,7 @@
    },
    "outputs": [],
    "source": [
-    "#note this %pip is only needed for running in jupyterlite online\n",
+    "# note this %pip is only needed for running in jupyterlite online\n",
     "%pip install ipython\n",
     "from _config import IFrame\n",
     "\n",

demo/1_PinCellCorrection_inter.ipynb

@@ -121,8 +121,9 @@
    },
    "outputs": [],
    "source": [
-    "#note this install_montepy is only necessary for jupyterlite; You don't need to use it locally\n",
+    "# note this install_montepy is only necessary for jupyterlite; You don't need to use it locally\n",
     "from _config import install_montepy\n",
+    "\n",
     "install_montepy()\n",
     "\n",
     "# actually needed\n",

demo/2_BuildAssembly_inter.ipynb

@@ -65,6 +65,7 @@
    "outputs": [],
    "source": [
     "from _config import install_montepy\n",
+    "\n",
     "install_montepy()\n",
     "\n",
     "import montepy\n",

demo/3_AddingGuideTubes_inter.ipynb

@@ -105,6 +105,7 @@
    "outputs": [],
    "source": [
     "from _config import install_montepy\n",
+    "\n",
     "install_montepy()\n",
     "\n",
     "import montepy\n",

demo/4_AxialDiscretize_inter.ipynb

@@ -71,6 +71,7 @@
    "outputs": [],
    "source": [
     "from _config import install_montepy\n",
+    "\n",
     "install_montepy()\n",
     "\n",
     "import montepy\n",

demo/answers/1_PinCellCorrection.ipynb

@@ -120,8 +120,9 @@
    },
    "outputs": [],
    "source": [
-    "#note this install_montepy is only necessary for jupyterlite; You don't need to use it locally\n",
+    "# note this install_montepy is only necessary for jupyterlite; You don't need to use it locally\n",
     "from _config import install_montepy\n",
+    "\n",
     "install_montepy()\n",
     "\n",
     "# actually needed\n",

demo/answers/2_BuildAssembly.ipynb

@@ -64,8 +64,9 @@
    },
    "outputs": [],
    "source": [
-    "#note this install_montepy is only necessary for jupyterlite; You don't need to use it locally\n",
+    "# note this install_montepy is only necessary for jupyterlite; You don't need to use it locally\n",
     "from _config import install_montepy\n",
+    "\n",
     "install_montepy()\n",
     "\n",
     "import montepy\n",

demo/answers/3_AddingGuideTubes.ipynb

@@ -104,6 +104,7 @@
    "outputs": [],
    "source": [
     "from _config import install_montepy\n",
+    "\n",
     "install_montepy()\n",
     "\n",
     "import montepy\n",

demo/answers/4_AxialDiscretize.ipynb

@@ -71,6 +71,7 @@
    "outputs": [],
    "source": [
     "from _config import install_montepy\n",
+    "\n",
     "install_montepy()\n",
     "\n",
     "import montepy\n",

doc/source/changelog.rst

@@ -5,6 +5,13 @@ MontePy Changelog
 1.2 releases
 ============
 
+#Next Version#
+--------------
+
+**Performance Improvement**
+
+* Implement _collection_ref to link objects to their NumberedObjectCollection parent (:issue:`867`).
+
 1.2.0
 --------------
 

montepy/cell.py

@@ -927,6 +927,10 @@ def num(obj):
         for key in keys:
             attr = getattr(self, key)
             setattr(result, key, copy.deepcopy(attr, memo))
+        # Clear weakrefs so the cloned cell isn't linked to original collection/problem
+        # This prevents number conflict checks against the original collection
+        result._collection_ref = None
+        result._problem_ref = None
         # copy geometry
         for special in special_keys:
             new_objs = []

montepy/numbered_mcnp_object.py

@@ -5,6 +5,7 @@
 import itertools
 from typing import Union
 from numbers import Integral
+import weakref
 
 from montepy.mcnp_object import MCNP_Object, InitInput
 import montepy
@@ -14,20 +15,30 @@
 def _number_validator(self, number):
     if number < 0:
         raise ValueError("number must be >= 0")
-    if self._problem:
-        obj_map = montepy.MCNP_Problem._NUMBERED_OBJ_MAP
-        try:
-            collection_type = obj_map[type(self)]
-        except KeyError as e:
-            found = False
-            for obj_class in obj_map:
-                if isinstance(self, obj_class):
-                    collection_type = obj_map[obj_class]
-                    found = True
-                    break
-            if not found:
-                raise e
-        collection = getattr(self._problem, collection_type.__name__.lower())
+
+    # Only validate against collection if linked to a problem
+    if self._problem is not None:
+        if self._collection is not None:
+            collection = self._collection
+        else:
+            # Find collection via _problem
+            obj_map = montepy.MCNP_Problem._NUMBERED_OBJ_MAP
+            collection_type = obj_map.get(type(self))
+
+            if collection_type is None:
+                # Finding via inheritance
+                for obj_class in obj_map:
+                    if isinstance(self, obj_class):
+                        collection_type = obj_map[obj_class]
+                        break
+
+            if collection_type is not None:
+                collection = getattr(self._problem, collection_type.__name__.lower())
+            else:
+                raise TypeError(
+                    f"Could not find collection type for {type(self).__name__} in problem."
+                )
+
         collection.check_number(number)
         collection._update_number(self.number, number, self)
 
@@ -59,6 +70,7 @@ def __init__(
             self._number = self._generate_default_node(int, -1)
         super().__init__(input, parser)
         self._load_init_num(number)
+        self._collection_ref = None
 
     def _load_init_num(self, number):
         if number is not None:
@@ -123,6 +135,24 @@ def _add_children_objs(self, problem):
                 except (TypeError, AssertionError):
                     prob_collect.append(child_collect)
 
+    @property
+    def _collection(self):
+        """Returns the parent collection this object belongs to, if any."""
+        if self._collection_ref is not None:
+            return self._collection_ref()
+        return None
+
+    def __getstate__(self):
+        state = super().__getstate__()
+        # Remove _collection_ref weakref as it can't be pickled
+        if "_collection_ref" in state:
+            del state["_collection_ref"]
+        return state
+
+    def __setstate__(self, crunchy_data):
+        crunchy_data["_collection_ref"] = None
+        super().__setstate__(crunchy_data)
+
     def clone(self, starting_number=None, step=None):
         """Create a new independent instance of this object with a new number.
 

montepy/numbered_object_collection.py

@@ -178,6 +178,7 @@ def __init__(
                         )
                     )
                 self.__num_cache[obj.number] = obj
+                self._link_to_collection(obj)
             self._objects = objects
 
     def link_to_problem(self, problem):
@@ -198,13 +199,38 @@ def link_to_problem(self, problem):
             self._problem_ref = weakref.ref(problem)
         for obj in self:
             obj.link_to_problem(problem)
+            # the _collection_ref that points to the main cells collection.
+            if problem is not None:
+                existing_coll = obj._collection
+                if existing_coll is None or existing_coll._problem is not problem:
+                    self._link_to_collection(obj)
 
     @property
     def _problem(self):
         if self._problem_ref is not None:
             return self._problem_ref()
         return None
 
+    def _link_to_collection(self, obj):
+        """Links the given object to this collection via a weakref.
+
+        Parameters
+        ----------
+        obj : Numbered_MCNP_Object
+            The object to link to this collection.
+        """
+        obj._collection_ref = weakref.ref(self)
+
+    def _unlink_from_collection(self, obj):
+        """Unlinks the given object from this collection.
+
+        Parameters
+        ----------
+        obj : Numbered_MCNP_Object
+            The object to unlink from this collection.
+        """
+        obj._collection_ref = None
+
     def __getstate__(self):
         state = self.__dict__.copy()
         weakref_key = "_problem_ref"
@@ -341,10 +367,8 @@ def extend(self, other_list):
                 )
             if obj.number in nums:
                 raise NumberConflictError(
-                    (
-                        f"When adding to {type(self).__name__} there was a number collision due to "
-                        f"adding {obj} which conflicts with {self[obj.number]}"
-                    )
+                    f"When adding to {type(self).__name__} there was a number collision due to "
+                    f"adding {obj} which conflicts with existing object number {obj.number}"
                 )
             nums.add(obj.number)
         for obj in other_list:
@@ -496,6 +520,7 @@ def __internal_append(self, obj, **kwargs):
                 )
         self.__num_cache[obj.number] = obj
         self._objects.append(obj)
+        self._link_to_collection(obj)
         self._append_hook(obj, **kwargs)
         if self._problem:
             obj.link_to_problem(self._problem)
@@ -507,6 +532,7 @@ def __internal_delete(self, obj, **kwargs):
         """
         self.__num_cache.pop(obj.number, None)
         self._objects.remove(obj)
+        self._unlink_from_collection(obj)
         self._delete_hook(obj, **kwargs)
 
     def add(self, obj: Numbered_MCNP_Object):
@@ -613,6 +639,7 @@ def append_renumber(self, obj, step=1):
         number = obj.number if obj.number > 0 else 1
         if self._problem:
             obj.link_to_problem(self._problem)
+        self._unlink_from_collection(obj)
         try:
             self.append(obj)
         except (NumberConflictError, ValueError) as e:

tests/test_numbered_collection.py

@@ -86,6 +86,7 @@ def test_extend(self, cp_simple_problem):
         extender = copy.deepcopy(extender)
         for surf in extender:
             surf._problem = None
+            surf._collection_ref = None
         surfaces[1000].number = 1
         extender[0].number = 1000
         extender[1].number = 70
@@ -161,6 +162,7 @@ def test_append_renumber(self, cp_simple_problem):
             cells.append_renumber(cell, "hi")
         cell = copy.deepcopy(cell)
         cell._problem = None
+        cell._collection_ref = None
         cell.number = 1
         cells.append_renumber(cell)
         assert cell.number == 4