From 1b978f0788739704b77dd06e0677924785063443 Mon Sep 17 00:00:00 2001
From: church89 <l.chierici89@gmail.com>
Date: Mon, 30 Oct 2023 09:59:47 +0100
Subject: [PATCH] add authomatic refuel scheme class

---
 openmc/deplete/abc.py       |  7 +--
 openmc/deplete/batchwise.py | 94 +++++++++++++++++++++++++++++++++++--
 2 files changed, 95 insertions(+), 6 deletions(-)

diff --git a/openmc/deplete/abc.py b/openmc/deplete/abc.py
index 300d4d7b025..43a7689e43b 100644
--- a/openmc/deplete/abc.py
+++ b/openmc/deplete/abc.py
@@ -27,7 +27,8 @@
 from .transfer_rates import TransferRates
 from openmc import Material, Cell
 from .batchwise import (BatchwiseCellGeometrical, BatchwiseCellTemperature,
-    BatchwiseMaterialRefuel, BatchwiseMaterialDilute, BatchwiseSchemeStd)
+    BatchwiseMaterialRefuel, BatchwiseMaterialDilute, BatchwiseSchemeStd,
+    BatchwiseSchemeRefuel)
 
 __all__ = [
     "OperatorResult", "TransportOperator",
@@ -916,8 +917,8 @@ def add_batchwise_scheme(self, scheme_name, **kwargs):
         """
         if scheme_name == 'std':
             self.batchwise = BatchwiseSchemeStd(self.batchwise, len(self), **kwargs)
-        #elif scheme_name == 'flex':
-        #    self.batchwise = BatchwiseSchemeFlex(self.batchwise, **kwargs)
+        elif scheme_name == 'refuel':
+            self.batchwise = BatchwiseSchemeRefuel(self.batchwise, **kwargs)
 
     def add_density_function(self, mats, density_func, oxidation_states):
         self.batchwise.set_density_function(mats, density_func, oxidation_states)
diff --git a/openmc/deplete/batchwise.py b/openmc/deplete/batchwise.py
index 96f556e2d8a..f8426657f7f 100644
--- a/openmc/deplete/batchwise.py
+++ b/openmc/deplete/batchwise.py
@@ -286,7 +286,7 @@ def _search_for_keff(self, val):
 
             else:
                 raise ValueError('ERROR: Search_for_keff output is not valid')
-        
+
         return root
 
     def _get_materials(self, vals):
@@ -818,7 +818,7 @@ def __init__(self, cell, attrib_name, operator, model, bracket,
         check_value('attrib_name', attrib_name,
                     ('rotation', 'translation'))
         self.attrib_name = attrib_name
-        
+
         # check if cell is filled with 2 cells
         if not isinstance(self.cell.fill, openmc.universe.DAGMCUniverse):
 
@@ -831,7 +831,7 @@ def __init__(self, cell, attrib_name, operator, model, bracket,
         # Initialize vector
         self.vector = np.zeros(3)
 
-        
+
         check_type('samples', samples, int)
         self.samples = samples
 
@@ -1565,3 +1565,91 @@ def search_for_keff(self, x, step_index):
                 for rank in range(comm.size):
                     comm.Abort()
         return x
+
+class BatchwiseSchemeRefuel():
+    """
+    Batchwise wrapper class, it wraps BatchwiseGeom and BatchwiseMat instances,
+    with some user defined logic.
+
+    This class should probably not be defined here, but we can keep it now for
+    convenience
+
+    The loop logic of this wrapper class is the following:
+
+    1. Run BatchwiseGeom and return geometrical coefficient
+    2. check if geometrical coefficient hit upper limit
+    3.1 if not, continue
+    3.2 if yes, refuel and reset geometrical coefficient
+
+    An instance of this class can be passed directly to an instance of the
+    integrator class, such as :class:`openmc.deplete.CECMIntegrator`.
+
+    Parameters
+    ----------
+    bw_geom : BatchwiseGeom
+        openmc.deplete.batchwise.BatchwiseGeom object
+    bw_mat : BatchwiseMat
+        openmc.deplete.batchwise.BatchwiseMat object
+    restart_level : int
+        Geometrical coefficient after reset
+    """
+
+    def __init__(self, bw_list, restart_level):
+
+        if isinstance(bw_list, list):
+            for bw in bw_list:
+                if isinstance(bw, BatchwiseCell):
+                    self.bw_geom = bw
+                elif isinstance(bw, BatchwiseMaterial):
+                    self.bw_mat = bw
+                else:
+                    raise ValueError(f'{bw} is not a valid instance of'
+                                      ' Batchwise class')
+        else:
+            raise ValueError(f'{bw_list} is not a list')
+
+        self.bw_list = bw_list
+        
+        if not isinstance(restart_level, (float, int)):
+            raise ValueError(f'{restart_level} is of type {type(restart_level)},'
+                             ' while it should be int or float')
+        else:
+            self.restart_level = restart_level
+        
+    def set_density_function(self, mats, density_func, oxidation_states):
+        for bw in self.bw_list:
+            bw.set_density_function(mats, density_func, oxidation_states)
+
+    def _update_volumes_after_depletion(self, x):
+        """
+        This is for a restart. TODO update abc class
+        Parameters
+        ----------
+        x : list of numpy.ndarray
+            Total atom concentrations
+        """
+        self.bw_geom._update_volumes(x)
+
+    def search_for_keff(self, x, step_index):
+        """
+        Perform the criticality search on the parametric material model.
+        Will set the root of the `search_for_keff` function to the atoms
+        concentrations vector.
+        Parameters
+        ----------
+        x : list of numpy.ndarray
+            Total atoms concentrations
+        Returns
+        -------
+        x : list of numpy.ndarray
+            Updated total atoms concentrations
+        """
+        #Start by doing a geometrical search§
+        x,root = self.bw_geom.search_for_keff(x, step_index)
+        #check if upper geometrical limit gets hit
+        if root >= self.bw_geom.bracket_limit[1]:
+            # Reset geometry and refuel 
+            self.bw_geom._set_cell_attrib(self.restart_level)
+            x,root = self.bw_mat.search_for_keff(x, step_index)
+
+        return x,root