Add check that flux is conserved

Author: dschwoerer

zoidberg/zoidberg.py

@@ -237,7 +237,10 @@ def get_metric(grid, magnetic_field):
         metric["g_yy"][:, yindex, :] *= (Bmag[:, yindex, :] / By) ** 2
         metric["gyy"][:, yindex, :] *= (By / Bmag[:, yindex, :]) ** 2
 
-    return metric, Bmag, pressure
+    # B * J / sqrt(g22)
+    BJg = Bmag * np.sqrt(metric["g_xx"] * metric["g_zz"] - metric["g_xz"] ** 2)
+
+    return metric, Bmag, pressure, BJg
 
 
 class MapWriter:
@@ -273,6 +276,7 @@ def __init__(
         self.is_open = False
         self.create = create
 
+        self.BJg = None
         self.grid = None
         self.field = None
         self.metric_done = False
@@ -373,7 +377,7 @@ def _write_metric(self):
         if self.metric_done:
             return
 
-        metric, Bmag, pressure = get_metric(self.grid, self.field)
+        metric, Bmag, pressure, self.BJg = get_metric(self.grid, self.field)
 
         # Add Rxy, Bxy
         metric["Bxy"] = Bmag
@@ -450,7 +454,17 @@ def _write_par_metric(self, maps, nslice, ypar):
                 yperiodic=yperiodic,
             )
 
-            par_metric, _, _ = get_metric(par_grid, self.field)
+            par_metric, _, _, par_BJg = get_metric(par_grid, self.field)
+
+            if self.BJg is not None and self.BJg.shape[0] > 4:
+                mymax = np.max(np.abs(par_BJg / self.BJg - 1)[2:-2], axis=(0, 2))
+                if np.max(mymax) > 1e-6:
+                    print(
+                        "FluxError",
+                        np.max(np.abs(par_BJg / self.BJg - 1)[2:-2], axis=(0, 2)),
+                    )
+                    ysel = np.argmax(mymax)
+
             if not self.new_names:
                 par_metric = update_metric_names(par_metric)