Further dual variable support (automatic adding to containers, renami…

…ng, etc.)

RMK_support/tests/test_examples.py

@@ -54,29 +54,30 @@ def test_solkit_fluid_janev():
     heatingPower = 3.5464790894703255
 
     # these tests only check for latex construction errors
+    with pytest.warns(UserWarning):
+
+        rk = solkit_mijin_thesis.generatorSKThesis(
+            dx0=dx0 / 4,
+            dxN=dxN / 4,
+            Nx=256,
+            Nh=17,
+            lmax=1,
+            mpiProcs=16,
+            initialTimestep=2.0,
+            nu=0.8 / 1.09345676,
+            heatingPower=heatingPower,
+            includedJanevTransitions=["ex", "deex", "ion", "recomb3b"],
+            numNeutrals=10,
+        )
 
-    rk = solkit_mijin_thesis.generatorSKThesis(
-        dx0=dx0 / 4,
-        dxN=dxN / 4,
-        Nx=256,
-        Nh=17,
-        lmax=1,
-        mpiProcs=16,
-        initialTimestep=2.0,
-        nu=0.8 / 1.09345676,
-        heatingPower=heatingPower,
-        includedJanevTransitions=["ex", "deex", "ion", "recomb3b"],
-        numNeutrals=10,
-    )
-
-    rk.setPETScOptions(
-        cliOpts="-pc_type bjacobi -sub_pc_factor_shift_type nonzero",
-        kspSolverType="gmres",
-    )
-
-    rk.generatePDF("SOL-KiT Fluid Test")
-
-    rk.writeConfigFile()
+        rk.setPETScOptions(
+            cliOpts="-pc_type bjacobi -sub_pc_factor_shift_type nonzero",
+            kspSolverType="gmres",
+        )
+
+        rk.generatePDF("SOL-KiT Fluid Test")
+
+        rk.writeConfigFile()
 
     os.remove(os.path.curdir + "/SOL-KiT_Fluid_Test.pdf")
     os.remove(os.path.curdir + "/config.json")

RMK_support/tests/test_remkit_context.py

@@ -60,7 +60,7 @@ def test_add_var(grid: Grid):
 
     # Add the variables to another variable container and set the rk.variables container
     varCont = vc.VariableContainer(grid)
-    varCont.add(a, b, var, var_dual)
+    varCont.add(a, b, var)
 
     rk.variables = varCont
 

RMK_support/tests/test_variable_container.py

@@ -26,7 +26,7 @@ def grid():
 
 @pytest.fixture
 def vCont(grid):
-    return VariableContainer(grid)
+    return VariableContainer(grid, autoAddDuals=False)
 
 
 def test_fluid_zero_var(grid):
@@ -108,7 +108,7 @@ def test_var_and_dual(grid):
 
     c = a.onDualGrid().rename("c")
     assert c.name == "c"
-    assert c.dual.name == "a_dual"
+    assert c.dual.name == "c_dual"
     assert c.isOnDualGrid
 
     b_dual, b = varAndDual("b", grid, primaryOnDualGrid=True)

RMK_support/variable_container.py

@@ -432,9 +432,12 @@ def name(self, name: str):
         self.__name__ = name
 
     def rename(self, name: str):
-        """Create a copy of this variable and rename it"""
+        """Create a copy of this variable and rename it. If the variable has a dual, it will be remade to avoid name clashes"""
         newVar = deepcopy(self)
         newVar.name = name
+        if newVar.dual is not None:
+            newVar.dual = None
+            _ = newVar.withDual()
         return newVar
 
     @property
@@ -561,9 +564,12 @@ def inOutput(self):
         return self.__inOutput__
 
     def onDualGrid(self, dual=True):
-        """Return a copy of this variable setting it to live on the dual or regular grid"""
+        """Return a copy of this variable setting it to live on the dual or regular grid. If the variable has a dual associated with it, it will be remade to avoid having duals that live on the same grid"""
         newVar = deepcopy(self)
         newVar.__isOnDualGrid__ = dual
+        if newVar.dual is not None:
+            newVar.dual = None
+            _ = newVar.withDual()
         return newVar
 
     def switchUnits(self):
@@ -747,12 +753,15 @@ def varAndDual(
 class VariableContainer:
     """Container object for Variables"""
 
-    def __init__(self, gridObj: Grid, timestamps: np.ndarray = np.array([])) -> None:
+    def __init__(
+        self, gridObj: Grid, timestamps: np.ndarray = np.array([]), autoAddDuals=True
+    ) -> None:
         """Container object for Variables
 
         Args:
             gridObj (Grid): Grid on which all the variables will live
             timestamps (np.ndarray, optional): Timestamps in case a time dimension should be added. Defaults to np.array([]).
+            autoAddDuals (bool, optional): If true will automatically add dual variables if associated at time of adding. Defaults to True.
         """
         self.__coords__ = {
             "x": gridObj.xGrid,
@@ -786,6 +795,8 @@ def __init__(self, gridObj: Grid, timestamps: np.ndarray = np.array([])) -> None
             )
         ]
 
+        self.__autoAddDuals__ = autoAddDuals
+
     def add(self, *args: Variable):
         """Add any number of variables to the container"""
         for var in args:
@@ -805,6 +816,17 @@ def add(self, *args: Variable):
             else:
                 self.__variables__.append(var)
 
+            if self.__autoAddDuals__ and var.dual is not None:
+                if var.dual.name in self.varNames:
+                    warnings.warn(
+                        "Variable "
+                        + var.dual.name
+                        + " already in VariableContainer. Overwriting."
+                    )
+                    self.__variables__[self.varNames.index(var.dual.name)] = var.dual
+                else:
+                    self.__variables__.append(var.dual)
+
     def setVar(
         self,
         name: str,

examples/epperlein_short_test.py

@@ -92,7 +92,7 @@ def esTestGenerator(**kwargs) -> rmk.RMKContext:
     q_dual,q = rmk.varAndDual("q",rk.grid,primaryOnDualGrid=True,derivation=rk.textbook["heatFluxMoment"],derivationArgs=["f"])
 
     logLee = rmk.Variable("logLee",rk.grid,isDerived=True)
-    rk.variables.add(f,f_dual,W,n,n_dual,ni,ni_dual,T,T_dual,zeroVar,E_dual,E,q_dual,q,logLee)
+    rk.variables.add(f,W,n,ni,T,zeroVar,E,q,logLee)
     # ### Braginskii flux derivation objects
 
     nConstGradT = 12*np.pi**1.5*epsilon0**2/np.sqrt(elMass*elCharge)  # Comes from e-i collision time

examples/flowing_cold_ion_ei_collisions.ipynb

@@ -153,7 +153,7 @@
     "T,T_dual = rmk.varAndDual(\"T\",rk.grid,derivation=rk.textbook[\"tempFromEnergye\"],derivationArgs=[\"W\",\"n\",\"zeroVar\"])\n",
     "zeroVar = rmk.Variable(\"zeroVar\",rk.grid,isDerived=True,inOutput=False)\n",
     "\n",
-    "rk.variables.add(f,f_dual,W,n,n_dual,ni,ni_dual,Gi_dual,Gi,ui_dual,ui,G_dual,G,u_dual,u,T,T_dual,zeroVar)\n"
+    "rk.variables.add(f,W,n,ni,Gi,ui,G,u,T,zeroVar)\n"
    ]
   },
   {

examples/gaussian_advection.ipynb

@@ -94,7 +94,7 @@
     "T = rmk.Variable(\"T\",rk.grid,data=TInit,isDerived=True,isCommunicated=False)\n",
     "G_dual,G = rmk.varAndDual(\"G\",rk.grid,primaryOnDualGrid=True) #the first return value is the primary, so here it is the dual\n",
     "\n",
-    "rk.variables.add(n,n_dual,T,G_dual,G)"
+    "rk.variables.add(n,T,G)"
    ]
   },
   {

examples/logical_boundary_condition.ipynb

@@ -117,7 +117,7 @@
     "f,f_dual = rmk.varAndDual(\"f\",rk.grid,isDistribution=True,data=fInit)\n",
     "n,n_dual = rmk.varAndDual(\"n\",rk.grid,derivation=rk.textbook[\"densityMoment\"],derivationArgs=[\"f\"])\n",
     "\n",
-    "rk.variables.add(f,f_dual,n,n_dual)\n",
+    "rk.variables.add(f,n)\n",
     "\n",
     "nb_right = rmk.derivations.BoundedExtrapolationDerivation(\"nb_right\",extrapolationType=\"linlog\")(n)\n",
     "nb_left = rmk.derivations.BoundedExtrapolationDerivation(\"nb_left\",extrapolationType=\"linlog\",leftBoundary=True)(n)\n",

examples/mms_test_2_fluid.ipynb

@@ -187,7 +187,7 @@
     "\n",
     "ui_dual,ui = rmk.varAndDual(\"ui\",rk.grid,primaryOnDualGrid=True,derivation=rk.textbook[\"flowSpeedFromFlux\"],derivationArgs=[\"Gi_dual\",\"ni_dual\"])\n",
     "\n",
-    "rk.variables.add(ne,ne_dual,ni,ni_dual,Ge_dual,Ge,Gi_dual,Te,Te_dual,E_dual,E,ue_dual,ue,ui_dual,ui)"
+    "rk.variables.add(ne,ni,Ge,Te,E,ue,ui)"
    ]
   },
   {

examples/solkit_mijin_thesis.py

@@ -134,36 +134,36 @@ def generatorSKThesis(**kwargs) -> rmk.RMKContext:
     ni,ni_dual = rmk.varAndDual("ni",rk.grid,data=nInit,units="$10^{19} m^{-3}$")
     ionSpecies.associateVar(ni)
 
-    rk.variables.add(ne,ne_dual,ni,ni_dual)
+    rk.variables.add(ne,ni)
     Ge_dual,Ge = rmk.varAndDual("Ge",rk.grid,primaryOnDualGrid=True)
     Gi_dual,Gi = rmk.varAndDual("Gi",rk.grid,primaryOnDualGrid=True)
 
-    rk.variables.add(Ge_dual,Ge,Gi_dual,Gi)
+    rk.variables.add(Ge,Gi)
 
     We, We_dual = rmk.varAndDual("We",rk.grid,units="$10^{20} eV m^{-3}$",data=WInit)
     Te,Te_dual = rmk.varAndDual("Te",rk.grid,units="$10eV$",isStationary=True,data=TInit)
 
-    rk.variables.add(We,We_dual,Te,Te_dual)
+    rk.variables.add(We,Te)
 
     # Set heat fluxes
 
     qe_dual,qe = rmk.varAndDual("qe",rk.grid,primaryOnDualGrid=True,isStationary=True)
 
-    rk.variables.add(qe_dual,qe)
+    rk.variables.add(qe)
 
     # Set E field
 
     E_dual, E = rmk.varAndDual("E",rk.grid,primaryOnDualGrid=True)
 
-    rk.variables.add(E_dual,E)
+    rk.variables.add(E)
     # Set derived fluid quantities
 
     rk.textbook = rmk.Textbook(rk.grid,tempDerivSpeciesIDs=[0])
 
     ue_dual,ue = rmk.varAndDual("ue",rk.grid,primaryOnDualGrid=True,derivation=rk.textbook["flowSpeedFromFlux"],derivationArgs=[Ge_dual.name,ne_dual.name])
     ui_dual,ui = rmk.varAndDual("ui",rk.grid,primaryOnDualGrid=True,derivation=rk.textbook["flowSpeedFromFlux"],derivationArgs=[Gi_dual.name,ni_dual.name])
 
-    rk.variables.add(ue_dual,ue,ui_dual,ui)
+    rk.variables.add(ue,ui)
 
     cs = rmk.Variable("cs",rk.grid,derivation=rk.textbook["sonicSpeedD+"],derivationArgs=[Te.name,Te.name])
 
@@ -223,7 +223,7 @@ def generatorSKThesis(**kwargs) -> rmk.RMKContext:
         nn.append(n)
         neutralSpecies[i].associateVar(n)
         nn_dual.append(n_dual)
-        rk.variables.add(n,n_dual)
+        rk.variables.add(n)
 
     # We need a distribution function to calculate rates from cross-sections built into the code
     fInit = np.zeros([rk.grid.numX, rk.grid.numH, rk.grid.numV])

examples/solkit_mijin_thesis_kin.py

@@ -154,11 +154,11 @@ def generatorSKThesisKin(**kwargs) -> rmk.RMKContext:
     electronSpecies.associateVar(ne,ne_dual)
     ionSpecies.associateVar(ni,ni_dual)
 
-    rk.variables.add(f,f_dual,We,ne,ne_dual,Ge_dual,Ge,qe_tot_dual,qe_tot,Te,Te_dual,qe_dual,qe,ue_dual,ue,ni,ni_dual,Gi,Gi_dual,ui,ui_dual)
+    rk.variables.add(f,We,ne,Ge,qe_tot,Te,qe,ue,ni,Gi,ui)
 
     E_dual, E = rmk.varAndDual("E",rk.grid,primaryOnDualGrid=True)
 
-    rk.variables.add(E_dual,E)
+    rk.variables.add(E)
 
     cs = rmk.Variable("cs",rk.grid,derivation=rk.textbook["sonicSpeedD+"],derivationArgs=[Te.name,Te.name])
 
@@ -198,7 +198,7 @@ def generatorSKThesisKin(**kwargs) -> rmk.RMKContext:
     qt_dual = qt.makeDual()
 
     q_ratio = varFromNode("qRatio",rk.grid,node(qe_dual)/node(qt_dual),isOnDualGrid=True)
-    rk.variables.add(qt,qt_dual,q_ratio,logLee)
+    rk.variables.add(qt,q_ratio,logLee)
 
     # Set neutral densities
 
@@ -209,7 +209,7 @@ def generatorSKThesisKin(**kwargs) -> rmk.RMKContext:
         nn.append(n)
         neutralSpecies[i].associateVar(n,n_dual)
         nn_dual.append(n_dual)
-        rk.variables.add(n,n_dual)
+        rk.variables.add(n)
 
     # ### Models
 
@@ -456,11 +456,11 @@ def generatorSKThesisKin(**kwargs) -> rmk.RMKContext:
 
     rk.addTermDiagnostics(Gi_dual,ni,E_dual,f)
 
-    # rk.variables.add(rmk.Variable("gammaRight",rk.grid,isScalar=True,scalarHostProcess=rk.mpiContext.fluidProcs[-1]))
-    # rk.manipulators.add(rmk.MBDataExtractor("gammaExtRight",rk.models["lbc_right"],rk.models["lbc_right"].mbData["gamma"],rk.variables["gammaRight"]))
+    rk.variables.add(rmk.Variable("gammaRight",rk.grid,isScalar=True,scalarHostProcess=rk.mpiContext.fluidProcs[-1]))
+    rk.manipulators.add(rmk.MBDataExtractor("gammaExtRight",rk.models["lbc_right"],rk.models["lbc_right"].mbData["gamma"],rk.variables["gammaRight"]))
 
-    # rk.variables.add(rmk.Variable("TRight",rk.grid,isScalar=True,scalarHostProcess=rk.mpiContext.fluidProcs[-1]))
-    # rk.manipulators.add(rmk.MBDataExtractor("tempExtRight",rk.models["lbc_right"],rk.models["lbc_right"].mbData["shTemp"],rk.variables["TRight"]))
+    rk.variables.add(rmk.Variable("TRight",rk.grid,isScalar=True,scalarHostProcess=rk.mpiContext.fluidProcs[-1]))
+    rk.manipulators.add(rmk.MBDataExtractor("tempExtRight",rk.models["lbc_right"],rk.models["lbc_right"].mbData["shTemp"],rk.variables["TRight"]))
 
     rk.manipulators.add(rmk.MBDataExtractor("logLeeExtract",rk.models["e-e0"],logLee))
 

tutorials/02_variables_IO.ipynb

@@ -358,9 +358,9 @@
    "source": [
     "vc = VariableContainer(grid) # An empty variable container\n",
     "\n",
-    "# We can add multiple variables to the container \n",
+    "# We can add multiple variables to the container (note that by default the VariableContainer adds the duals of any added variables)\n",
     "\n",
-    "vc.add(a,a_dual,b,b_dual)\n",
+    "vc.add(a,b)\n",
     "\n",
     "# Variable containers provide accessors that require knowing the variable name\n",
     "\n",
@@ -493,7 +493,9 @@
       "    b        (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n",
       "    b_dual   (x_dual) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0\n",
       "    c2       (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n",
-      "    c3       (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n"
+      "    c2_dual  (x_dual) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0\n",
+      "    c3       (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n",
+      "    c3_dual  (x_dual) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0\n"
      ]
     }
    ],
@@ -532,7 +534,9 @@
       "    b        (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n",
       "    b_dual   (x_dual) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0\n",
       "    c2       (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n",
-      "    c3       (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n"
+      "    c2_dual  (x_dual) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0\n",
+      "    c3       (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n",
+      "    c3_dual  (x_dual) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0\n"
      ]
     }
    ],
@@ -653,14 +657,32 @@
       "    b        (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n",
       "    b_dual   (x_dual) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0\n",
       "    c2       (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n",
-      "    c3       (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n"
+      "    c2_dual  (x_dual) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0\n",
+      "    c3       (x) float64 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0 0.0\n",
+      "    c3_dual  (x_dual) float64 0.0 0.0 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0 0.0\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/home/stefan/.local/lib/python3.8/site-packages/RMK_support/variable_container.py:799: UserWarning: Variable time already in VariableContainer. Overwriting.\n",
+      "/home/stefan/.local/lib/python3.8/site-packages/RMK_support/variable_container.py:807: UserWarning: Variable time already in VariableContainer. Overwriting.\n",
+      "  warnings.warn(\n",
+      "/home/stefan/.local/lib/python3.8/site-packages/RMK_support/variable_container.py:807: UserWarning: Variable a_dual already in VariableContainer. Overwriting.\n",
+      "  warnings.warn(\n",
+      "/home/stefan/.local/lib/python3.8/site-packages/RMK_support/variable_container.py:818: UserWarning: Variable a already in VariableContainer. Overwriting.\n",
+      "  warnings.warn(\n",
+      "/home/stefan/.local/lib/python3.8/site-packages/RMK_support/variable_container.py:807: UserWarning: Variable b_dual already in VariableContainer. Overwriting.\n",
+      "  warnings.warn(\n",
+      "/home/stefan/.local/lib/python3.8/site-packages/RMK_support/variable_container.py:818: UserWarning: Variable b already in VariableContainer. Overwriting.\n",
+      "  warnings.warn(\n",
+      "/home/stefan/.local/lib/python3.8/site-packages/RMK_support/variable_container.py:807: UserWarning: Variable c2_dual already in VariableContainer. Overwriting.\n",
+      "  warnings.warn(\n",
+      "/home/stefan/.local/lib/python3.8/site-packages/RMK_support/variable_container.py:818: UserWarning: Variable c2 already in VariableContainer. Overwriting.\n",
+      "  warnings.warn(\n",
+      "/home/stefan/.local/lib/python3.8/site-packages/RMK_support/variable_container.py:807: UserWarning: Variable c3_dual already in VariableContainer. Overwriting.\n",
+      "  warnings.warn(\n",
+      "/home/stefan/.local/lib/python3.8/site-packages/RMK_support/variable_container.py:818: UserWarning: Variable c3 already in VariableContainer. Overwriting.\n",
       "  warnings.warn(\n"
      ]
     }

tutorials/07_rmk_context.ipynb

@@ -199,7 +199,7 @@
     "T = rmk.Variable(\"T\",rk.grid,data=TInit,isDerived=True,isCommunicated=False)\n",
     "G_dual,G = rmk.varAndDual(\"G\",rk.grid,primaryOnDualGrid=True) \n",
     "\n",
-    "rk.variables.add(n,n_dual,T,G_dual,G) # rk.variables is a VariableContainer\n",
+    "rk.variables.add(n,T,G) # rk.variables is a VariableContainer\n",
     "\n",
     "model = rmk.Model(name=\"adv\")\n",
     "\n",