Added time derivative variable subtype and more latex improvements

Author: SMijin

RMK_support/common_variables.py

@@ -4,10 +4,60 @@
 from .variable_container import Variable, node
 from .derivations import Species
 from . import derivations
-from .remkit_context import RMKContext
 from abc import ABC, abstractmethod
 
 
+def timeDerivative(
+    name: str, timeNorm: float, variable: Variable, **kwargs
+) -> Variable:
+    """Generate a time derivative of a given variable as a variable, making sure units and grids are correct.
+
+    Args:
+        name (str): Variable name for the time derivative
+        timeNorm (float): Time normalisation
+        variable (Variable): Variable whose derivative should be taken
+
+    Returns:
+        Variable: Variable with correct time derivative units
+    """
+
+    for key in [
+        "units",
+        "unitSI",
+        "normSI",
+        "isOnDualGrid",
+        "isDistribution",
+        "isSingleHarmonic",
+        "isScalar",
+        "subtype",
+    ]:
+        assert key not in kwargs, key + " not allowed in time derivative variable call"
+
+    latexWrapped = "$" in variable.unitsSI
+    unitSI = variable.unitsSI.replace("$", "") + "/s"
+    if latexWrapped:
+        unitSI = "$" + unitSI + "$"
+    var = Variable(
+        name,
+        variable.grid,
+        **kwargs,
+        units=variable.unitsNorm + " / time norm.",
+        unitSI=unitSI,
+        normSI=variable.normSI / timeNorm,
+        isOnDualGrid=variable.isOnDualGrid,
+        isScalar=variable.isScalar,
+        isDistribution=variable.isDistribution,
+        isSingleHarmonics=variable.isSingleHarmonic,
+        subtype=variable.subtype + "_time_derivative"
+    )
+
+    return var
+
+
+# cyclic import workaround
+from .remkit_context import RMKContext
+
+
 def density(name: str, context: RMKContext, **kwargs) -> Variable:
     """Generate a density variable based on normalisation data within a RMKContext. Accepts regular Variable kwargs, except for those having to do with units, which are set automatically.
 

RMK_support/dashboard_support.py

@@ -9,7 +9,7 @@
 from typing_extensions import Self
 import param  # type: ignore
 from warnings import warn
-import matplotlib.pyplot as plt
+import matplotlib.pyplot as plt  # type: ignore
 
 
 def fluidVarX(var: Variable, time: int) -> hv.Curve:
@@ -603,18 +603,20 @@ def load_dist_curves(x: int, time: int, harmonic: int):
 
         return pn.Row(pn.Column(self.param, pos, t, h), pn.panel(dmap))
 
-def quickDash(vars:VariableContainer,runPaths: Dict[str, str]):
 
-    hv.extension('matplotlib')
-    plt.rcParams['figure.dpi'] = 150
-    hv.output(size=150,dpi=150)
+def quickDash(vars: VariableContainer, runPaths: Dict[str, str]):
 
-    loader = LazyLoader(vars,runPaths)
+    hv.extension("matplotlib")
+    plt.rcParams["figure.dpi"] = 150
+    hv.output(size=150, dpi=150)
+
+    loader = LazyLoader(vars, runPaths)
 
     element1 = ElementDisplay(loader)
     element2 = ElementDisplay(loader)
 
-    return pn.Row(element1.layout,element2.layout)
+    return pn.Row(element1.layout, element2.layout)
+
 
 class ReMKiT1DDashboard:
     def __init__(self, data: xr.Dataset, gridObj: Grid):

RMK_support/derivations.py

@@ -230,7 +230,7 @@ def addLatexToDoc(self, doc: tex.Document, **kwargs):
                                 "$"
                                 + species.latex()
                                 + "$"
-                                + f": ID: {species.speciesID}; A: {species.atomicA:.4e}; Z: {species.charge:.2f}; Associated vars: "
+                                + f": ID: {species.speciesID}; A: {species.atomicA:.4e}; Z: {species.charge:.2f} \\newline Associated vars: "
                                 + ", ".join(associatedVarNames)
                             )
                         )

RMK_support/remkit_context.py

@@ -792,6 +792,9 @@ def loadSimulation(
 
     def addTermDiagnostics(self, *args: Variable):
         """Add diagnostic variables for all terms evolving given variables. **NOTE**: Does not detect terms generated by TermGenerators - to include these use GroupEvaluators."""
+
+        from . import common_variables as cv
+
         for var in args:
             terms = self.models.getTermsThatEvolveVar(var)
             if not len(terms):
@@ -800,19 +803,43 @@ def addTermDiagnostics(self, *args: Variable):
                     + var.name
                     + " has no terms that evolve it"
                 )
-
+            if var.isStationary or var.unitsSI == "":
+                if var.isStationary:
+                    warnings.warn(
+                        var.name
+                        + " is stationary - cannot determine term units for diagnostics"
+                    )
+                else:
+                    warnings.warn(
+                        var.name
+                        + " doesn't have SI units set - cannot determine term units for diagnostics"
+                    )
             for pair in terms:
                 model, term = pair
-                self.variables.add(
-                    Variable(
-                        model + "_" + term,
-                        self.grid,
-                        isDerived=True,
-                        isDistribution=var.isDistribution,
-                        isOnDualGrid=var.isOnDualGrid,
-                        isCommunicated=False,
+
+                if var.isStationary or var.unitsSI == "":
+
+                    self.variables.add(
+                        Variable(
+                            model + "_" + term,
+                            self.grid,
+                            isDerived=True,
+                            isDistribution=var.isDistribution,
+                            isOnDualGrid=var.isOnDualGrid,
+                            isCommunicated=False,
+                        )
                     )
-                )
+                else:
+                    self.variables.add(
+                        cv.timeDerivative(
+                            model + "_" + term,
+                            self.norms["time"],
+                            var,
+                            isDerived=True,
+                            isCommunicated=False,
+                        )
+                    )
+
                 self.manipulators.add(
                     TermEvaluator(
                         model + "_" + term, [pair], self.variables[model + "_" + term]

RMK_support/tests/test_common_variables.py

@@ -78,3 +78,8 @@ def test_standard_variable_factory(context):
     assert factory.species["density"].name == "nn"
     assert factory.species["temperature"].name == "Tn"
     assert factory.species["heatflux"].name == "qn"
+
+    dndt = cv.timeDerivative("dndt", rk.norms["time"], n)
+    assert dndt.units == "norm. density / time norm."
+    assert dndt.normConst == rk.normDensity / rk.norms["time"]
+    assert dndt.unitsSI == "$m^{-3}/s$"

RMK_support/tests/test_examples.py

@@ -56,20 +56,21 @@ 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",

RMK_support/tests/test_remkit_context.py

@@ -219,7 +219,7 @@ def test_wrapper(grid: Grid):
     a, b = (Variable(name, rk.grid) for name in "ab")
 
     cDeriv = dv.NodeDerivation("cDeriv", node=vc.node(a) + vc.node(b))
-    c = Variable("c", rk.grid, isDerived=True, derivation=cDeriv)
+    c = Variable("c", rk.grid, isDerived=True, derivation=cDeriv, unitSI="$c$")
 
     rk.variables.add(a, b, c)
 
@@ -350,7 +350,8 @@ def test_wrapper(grid: Grid):
     # Now add term diagnostic manipulators for the model terms
 
     # Terms with evolved variables can have term diagnostics
-    rk.addTermDiagnostics(*[a, b])
+    with pytest.warns(UserWarning):
+        rk.addTermDiagnostics(*[a, b])
 
     # Adding a term diagnostic for a non-evolved term should raise a warning
     with pytest.warns(

RMK_support/variable_container.py

@@ -722,6 +722,7 @@ def latex(self, latexRemap: Dict[str, str] = {}) -> str:
             result += "single harmonic"
         if self.isStationary:
             result += ", stationary"
+        result += "\\newline Subtype: " + self.subtype.replace("_", r"\_")
 
         return result