From c358aa43cfeb55bf21f857996ca3641f8bae6bbc Mon Sep 17 00:00:00 2001
From: Hunter Smith <hsmith@terrapower.com>
Date: Mon, 28 Oct 2024 15:31:37 -0700
Subject: [PATCH] add test for detailedNDensity with detailed axial expansion

---
 armi/reactor/composites.py                    |  2 +-
 .../tests/test_axialExpansionChanger.py       | 32 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/armi/reactor/composites.py b/armi/reactor/composites.py
index 327228091..4b6e514a0 100644
--- a/armi/reactor/composites.py
+++ b/armi/reactor/composites.py
@@ -1576,7 +1576,7 @@ def changeNDensByFactor(self, factor):
         }
         self.setNumberDensities(densitiesScaled)
         # Update detailedNDens
-        if self.p.detailedNDens:
+        if self.p.detailedNDens is not None:
             self.p.detailedNDens *= factor
 
     def clearNumberDensities(self):
diff --git a/armi/reactor/converters/tests/test_axialExpansionChanger.py b/armi/reactor/converters/tests/test_axialExpansionChanger.py
index bf69ac64a..cfea92d0d 100644
--- a/armi/reactor/converters/tests/test_axialExpansionChanger.py
+++ b/armi/reactor/converters/tests/test_axialExpansionChanger.py
@@ -16,6 +16,7 @@
 import collections
 import os
 import unittest
+import copy
 from statistics import mean
 
 from numpy import array, linspace, zeros
@@ -267,6 +268,8 @@ def test_thermalExpansionContractionConservation_simple(self):
         a = buildTestAssemblyWithFakeMaterial(name="HT9")
         origMesh = a.getAxialMesh()[:-1]
         origMasses, origNDens = self._getComponentMassAndNDens(a)
+        # set component detailedNDens
+        origDetailedNDens = self._setComponentDetailedNDens(a, origNDens)
         axialExpChngr = AxialExpansionChanger(detailedAxialExpansion=True)
 
         tempGrid = linspace(0.0, a.getHeight())
@@ -281,16 +284,20 @@ def test_thermalExpansionContractionConservation_simple(self):
             # Set new isothermal temp and expand
             tempField = array([temp] * len(tempGrid))
             oldMasses, oldNDens = self._getComponentMassAndNDens(a)
+            oldDetailedNDens = self._getComponentDetailedNDens(a)
             axialExpChngr.performThermalAxialExpansion(a, tempGrid, tempField)
             newMasses, newNDens = self._getComponentMassAndNDens(a)
+            newDetailedNDens = self._getComponentDetailedNDens(a)
             self._checkMass(oldMasses, newMasses)
             self._checkNDens(oldNDens, newNDens, totGrowthFrac)
+            self._checkDetailedNDens(oldDetailedNDens, newDetailedNDens, totGrowthFrac)
 
         # make sure that the assembly returned to the original state
         for orig, new in zip(origMesh, a.getAxialMesh()):
             self.assertAlmostEqual(orig, new, places=12)
         self._checkMass(origMasses, newMasses)
         self._checkNDens(origNDens, newNDens, 1.0)
+        self._checkDetailedNDens(origDetailedNDens, newDetailedNDens, 1.0)
 
     def test_thermalExpansionContractionConservation_complex(self):
         """Thermally expand and then contract to ensure original state is recovered.
@@ -412,7 +419,13 @@ def _checkNDens(self, prevNDen, newNDens, ratio):
             for prev, new in zip(prevComp.values(), newComp.values()):
                 if prev:
                     self.assertAlmostEqual(prev / new, ratio, msg=f"{prev} / {new}")
-
+    
+    def _checkDetailedNDens(self, prevDetailedNDen, newDetailedNDens, ratio):
+        for prevComp, newComp in zip(prevDetailedNDen.values(), newDetailedNDens.values()):
+            for prev, new in zip(prevComp, newComp):
+                if prev:
+                    self.assertAlmostEqual(prev / new, ratio, msg=f"{prev} / {new}")
+    
     @staticmethod
     def _getComponentMassAndNDens(a):
         masses = {}
@@ -422,6 +435,23 @@ def _getComponentMassAndNDens(a):
                 masses[c] = c.getMass()
                 nDens[c] = c.getNumberDensities()
         return masses, nDens
+    
+    @staticmethod
+    def _setComponentDetailedNDens(a, nDens):
+        detailedNDens = {}
+        for b in a:
+            for c in getSolidComponents(b):
+                c.p.detailedNDens = copy.deepcopy([val for val in nDens[c].values()])
+                detailedNDens[c] = c.p.detailedNDens
+        return detailedNDens
+    
+    @staticmethod
+    def _getComponentDetailedNDens(a):
+        detailedNDens = {}
+        for b in a:
+            for c in getSolidComponents(b):
+                detailedNDens[c] = copy.deepcopy(c.p.detailedNDens)
+        return detailedNDens
 
     def test_targetComponentMassConservation(self):
         """Tests mass conservation for target components."""