1 files changed, 32 insertions, 1 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[];