Move conservation calculation code processing from KrancThorn.m to ConservationCalculation.m

2 files changed, 33 insertions, 26 deletions

diff --git a/Tools/CodeGen/ConservationCalculation.m b/Tools/CodeGen/ConservationCalculation.m
index 91926ad..3fa6696 100644
--- a/Tools/CodeGen/ConservationCalculation.m
+++ b/Tools/CodeGen/ConservationCalculation.m
@@ -19,7 +19,7 @@
 *)
 
 BeginPackage["ConservationCalculation`", {"Errors`", "Helpers`", "Kranc`",
-  "MapLookup`", "KrancGroups`", "CodeGenCalculation`", "Differencing`"}];
+  "MapLookup`", "KrancGroups`", "CodeGenCalculation`", "Differencing`", "Code`", "Object`"}];
 
 ProcessConservationCalculation;
 ConservationDifferencingOperators;
@@ -32,6 +32,7 @@ ConservationDifferencingRealParameters;
 hlleAlpha;
 PrimitiveEquations;
 ConservedEquations;
+ConservationCalculationProcessCode;
 
 Begin["`Private`"];
 
@@ -218,6 +219,36 @@ ConservationCalculationDeclaredGroups[calc_] :=
       ToString[#]<>"_flux_group",
       {fluxSymbol[#]}, {}] &, consVars[calc]]];
 
+Options[ConservationCalculationProcessCode] = ThornOptions;
+
+DefFn[
+  ConservationCalculationProcessCode[cIn_Code,opts:OptionsPattern[]] :=
+  Module[
+    {inputConsCalcs, outputConsCalcs, consGroups, c=cIn},
+
+    inputConsCalcs = Map[Append[#, Groups -> GetObjectField[c, "Groups"]] &,
+                         OptionValue[ConservationCalculations]];
+
+    outputConsCalcs = 
+    Flatten[
+      Map[
+        ProcessConservationCalculation[#, GetObjectField[c, "Name"]] &,
+        inputConsCalcs],
+      1];
+
+    outputConsCalcs =
+    Map[Join[#, {PartialDerivatives -> GetObjectField[c, "PartialDerivatives"],
+                 Implementation -> GetObjectField[c, "Implementation"]}] &,
+        outputConsCalcs];
+
+    consGroups = Union@Flatten[
+      Map[ConservationCalculationDeclaredGroups, inputConsCalcs],1];
+    
+    c = JoinObjectField[c, "Calculations", outputConsCalcs];
+    c = JoinObjectField[c, "Groups", consGroups];
+    c = JoinObjectField[c, "DeclaredGroups", Map[groupName, consGroups]];
+    c]];
+
 End[];
 
 EndPackage[];
diff --git a/Tools/CodeGen/KrancThorn.m b/Tools/CodeGen/KrancThorn.m
index e684f2b..de89f28 100644
--- a/Tools/CodeGen/KrancThorn.m
+++ b/Tools/CodeGen/KrancThorn.m
@@ -192,31 +192,7 @@ CreateKrancThorn[groupsOrig_, parentDirectory_, thornName_, opts:OptionsPattern[
        Conservation Calculations
        ------------------------------------------------------------------------ *)
 
-    Module[
-      {inputConsCalcs, outputConsCalcs, consGroups},
-
-      inputConsCalcs = Map[Append[#, Groups -> GetObjectField[c, "Groups"]] &,
-                           OptionValue[ConservationCalculations]];
-
-      outputConsCalcs = 
-      Flatten[
-        Map[
-          ProcessConservationCalculation[#, GetObjectField[c, "Name"]] &,
-          inputConsCalcs],
-        1];
-
-      outputConsCalcs =
-      Map[Join[#, {PartialDerivatives -> GetObjectField[c, "PartialDerivatives"],
-                   Implementation -> GetObjectField[c, "Implementation"]}] &,
-          outputConsCalcs];
-
-      consGroups = Union@Flatten[
-        Map[ConservationCalculationDeclaredGroups, inputConsCalcs],1];
-      
-      c = JoinObjectField[c, "Calculations", outputConsCalcs];
-      c = JoinObjectField[c, "Groups", consGroups];
-      c = JoinObjectField[c, "DeclaredGroups", Map[groupName, consGroups]]];
-
+    c = ConservationCalculationProcessCode[c, opts];
 
     (* ------------------------------------------------------------------------ 
        ODEs