Merge pull request #195 from boutproject/staggered-integrate

Support staggered variables in BoutDataset.integrate_midpoints()

Author: johnomotani

setup.cfg

@@ -28,7 +28,7 @@ setup_requires =
     setuptools_scm_git_archive
 install_requires =
     xarray>=0.17.0
-    boutdata>=0.1.2
+    boutdata>=0.1.4
     dask[array]>=2.10.0
     natsort>=5.5.0
     matplotlib>=3.1.1,!=3.3.0,!=3.3.1,!=3.3.2

xbout/boutdataset.py

@@ -305,6 +305,12 @@ def integrate_midpoints(self, variable, *, dims=None):
         """
         ds = self.data
 
+        if isinstance(variable, str):
+            variable = ds[variable]
+
+        location = variable.cell_location
+        suffix = "" if location == "CELL_CENTRE" else f"_{location}"
+
         tcoord = ds.metadata["bout_tdim"]
         xcoord = ds.metadata["bout_xdim"]
         ycoord = ds.metadata["bout_ydim"]
@@ -331,55 +337,61 @@ def integrate_midpoints(self, variable, *, dims=None):
         elif isinstance(dims, str):
             dims = [dims]
 
-        dx = ds["dx"]
-        dy = ds["dy"]
+        dx = ds[f"dx{suffix}"]
+        dy = ds[f"dy{suffix}"]
         dz = ds.metadata["dz"]
 
         # Work out the spatial volume element
         if xcoord in dims and ycoord in dims and zcoord in dims:
             # Volume integral, use the 3d Jacobian "J"
-            spatial_volume_element = ds["J"] * dx * dy * dz
+            spatial_volume_element = ds[f"J{suffix}"] * dx * dy * dz
         elif xcoord in dims and ycoord in dims:
             # 2d integral on poloidal planes
-            if ds[variable].direction_y == "Standard":
+            if variable.direction_y == "Standard":
                 # Need to use a metric constructed from basis vectors within the
                 # poloidal plane, so use 'reciprocal basis vectors' Grad(x^i)
                 # J = 1/sqrt(det(g_2d))
                 # det(g_2d) = g11*g22 - g12**2
-                g = ds["g11"] * ds["g22"] - ds["g12"] ** 2
+                g = ds[f"g11{suffix}"] * ds[f"g22{suffix}"] - ds[f"g12{suffix}"] ** 2
                 J = 1.0 / np.sqrt(g)
-            elif ds[variable].direction_y == "Aligned":
+            elif variable.direction_y == "Aligned":
                 # Need to work out area element from metric coefficients. See book by
                 # D'haeseleer, Hitchon, Callen and Shohet eq. (2.5.51).
                 # Need to use a metric constructed from basis vectors within the
                 # field-aligned x-y plane, so use 'tangent basis vectors' e_i
                 # J = sqrt(g_11*g_22 - g_12**2)
-                J = np.sqrt(ds["g_11"] * ds["g_22"] - ds["g_12"] ** 2)
+                J = np.sqrt(
+                    ds[f"g_11{suffix}"] * ds[f"g_22{suffix}"] - ds[f"g_12{suffix}"] ** 2
+                )
             spatial_volume_element = J * dx * dy
         elif xcoord in dims and zcoord in dims:
             # 2d integral on toroidal planes
             # Need to work out area element from metric coefficients. See book by
             # D'haeseleer, Hitchon, Callen and Shohet eq. (2.5.51)
             # J = sqrt(g_11*g_33 - g_13**2)
-            J = np.sqrt(ds["g_11"] * ds["g_33"] - ds["g_13"] ** 2)
+            J = np.sqrt(
+                ds[f"g_11{suffix}"] * ds[f"g_33{suffix}"] - ds[f"g_13{suffix}"] ** 2
+            )
             spatial_volume_element = J * dx * dz
         elif ycoord in dims and zcoord in dims:
             # 2d integral on flux-surfaces
             # Need to work out area element from metric coefficients. See book by
             # D'haeseleer, Hitchon, Callen and Shohet eq. (2.5.51)
             # J = sqrt(g_22*g_33 - g_23**2)
-            J = np.sqrt(ds["g_22"] * ds["g_33"] - ds["g_23"] ** 2)
+            J = np.sqrt(
+                ds[f"g_22{suffix}"] * ds[f"g_33{suffix}"] - ds[f"g_23{suffix}"] ** 2
+            )
             spatial_volume_element = J * dy * dz
         elif xcoord in dims:
-            if ds[variable].direction_y == "Aligned":
+            if variable.direction_y == "Aligned":
                 raise ValueError(
                     "Variable is field-aligned, but radial integral along coordinate "
                     "line in globally field-aligned coordinates not supported"
                 )
             # 1d radial integral, line element is sqrt(g_11)*dx
-            spatial_volume_element = np.sqrt(ds["g_11"]) * dx
+            spatial_volume_element = np.sqrt(ds[f"g_11{suffix}"]) * dx
         elif ycoord in dims:
-            if ds[variable].direction_y == "Standard":
+            if variable.direction_y == "Standard":
                 # Poloidal integral, line element is e_y projected onto a unit vector in
                 # the poloidal direction. e_z is in the toroidal direction and Grad(x)
                 # is orthogonal to flux surfaces, so their cross product is in the
@@ -398,34 +410,35 @@ def integrate_midpoints(self, variable, *, dims=None):
                 # For 'orthogonal' coordinates (radial and poloidal directions are
                 # orthogonal) this is equal to 1/sqrt(g22)
                 spatial_volume_element = (
-                    (ds["g_22"] * ds["g_33"] - ds["g_23"] ** 2)
-                    / (ds["J"] * np.sqrt(ds["g11"] * ds["g_33"]))
+                    (
+                        ds[f"g_22{suffix}"] * ds[f"g_33{suffix}"]
+                        - ds[f"g_23{suffix}"] ** 2
+                    )
+                    / (
+                        ds[f"J{suffix}"]
+                        * np.sqrt(ds[f"g11{suffix}"] * ds[f"g_33{suffix}"])
+                    )
                     * dy
                 )
-            elif ds[variable].direction_y == "Aligned":
+            elif variable.direction_y == "Aligned":
                 # Parallel integral, line element is sqrt(g_22)*dy
-                spatial_volume_element = np.sqrt(ds["g_22"]) * dy
+                spatial_volume_element = np.sqrt(ds[f"g_22{suffix}"]) * dy
         elif zcoord in dims:
             # Toroidal integral, line element is sqrt(g_33)*dz
-            spatial_volume_element = np.sqrt(ds["g_33"]) * dz
+            spatial_volume_element = np.sqrt(ds[f"g_33{suffix}"]) * dz
         else:
             # No spatial integral
             spatial_volume_element = 1.0
 
-        if isinstance(variable, xr.DataArray):
-            integrand = variable
-        else:
-            integrand = ds[variable]
-
         spatial_dims = set(dims) - set([tcoord])
 
         # Need to check if the variable being integrated is a Field2D, which does not
         # have a z-dimension to sum over. Other variables are OK because metric
         # coefficients, dx and dy all have both x- and y-dimensions so variable would be
         # broadcast to include them if necessary
-        missing_z_sum = zcoord in dims and zcoord not in integrand.dims
+        missing_z_sum = zcoord in dims and zcoord not in variable.dims
 
-        integrand = integrand * spatial_volume_element
+        integrand = variable * spatial_volume_element
 
         integral = integrand.sum(dim=spatial_dims)
 

xbout/tests/test_boutdataset.py

@@ -1304,11 +1304,17 @@ def test_integrate_midpoints_slab(self, bout_xyt_example_files):
             rtol=1.4e-4,
         )
 
-    def test_integrate_midpoints_salpha(self, bout_xyt_example_files):
+    @pytest.mark.parametrize(
+        "location", ["CELL_CENTRE", "CELL_XLOW", "CELL_YLOW", "CELL_ZLOW"]
+    )
+    def test_integrate_midpoints_salpha(self, bout_xyt_example_files, location):
         # Create data
+        nx = 100
+        ny = 110
+        nz = 120
         dataset_list = bout_xyt_example_files(
             None,
-            lengths=(4, 100, 110, 120),
+            lengths=(4, nx, ny, nz),
             nxpe=1,
             nype=1,
             nt=1,
@@ -1323,6 +1329,16 @@ def test_integrate_midpoints_salpha(self, bout_xyt_example_files):
 
         # Integrate 1 so we just get volume, areas and lengths
         ds["n"].values[:] = 1.0
+        ds["n"].attrs["cell_location"] = location
+
+        # remove boundary cells (don't want to integrate over those)
+        ds = ds.bout.remove_yboundaries()
+        if ds.metadata["keep_xboundaries"] and ds.metadata["MXG"] > 0:
+            mxg = ds.metadata["MXG"]
+            xslice = slice(mxg, -mxg)
+        else:
+            xslice = slice(None)
+        ds = ds.isel(x=xslice)
 
         # Test geometry has major radius R and goes between minor radii a-Lr/2 and
         # a+Lr/2
@@ -1331,18 +1347,14 @@ def test_integrate_midpoints_salpha(self, bout_xyt_example_files):
         Lr = options.evaluate_scalar("mesh:Lr")
         rinner = a - Lr / 2.0
         router = a + Lr / 2.0
+        r = options.evaluate("mesh:r").squeeze()[xslice]
+        if location == "CELL_XLOW":
+            rinner = rinner - Lr / (2.0 * nx)
+            router = router - Lr / (2.0 * nx)
+            r = r - Lr / (2.0 * nx)
         q = options.evaluate_scalar("mesh:q")
         T_total = (ds.sizes["t"] - 1) * (ds["t"][1] - ds["t"][0]).values
 
-        # remove boundary cells (don't want to integrate over those)
-        ds = ds.bout.remove_yboundaries()
-        if ds.metadata["keep_xboundaries"] and ds.metadata["MXG"] > 0:
-            mxg = ds.metadata["MXG"]
-            xslice = slice(mxg, -mxg)
-        else:
-            xslice = slice(None)
-        ds = ds.isel(x=xslice)
-
         # Volume of torus with circular cross-section of major radius R and minor radius
         # a is 2*pi*R*pi*a^2
         # https://en.wikipedia.org/wiki/Torus
@@ -1378,12 +1390,6 @@ def test_integrate_midpoints_salpha(self, bout_xyt_example_files):
         # Area of torus with circular cross-section of major radius R and minor radius a
         # is 2*pi*R*2*pi*a
         # https://en.wikipedia.org/wiki/Torus
-        mxg = options._keys.get("MXG", 2)
-        if mxg == 0:
-            xslice = slice(None)
-        else:
-            xslice = slice(mxg, -mxg)
-        r = options.evaluate("mesh:r").squeeze()[xslice]
         npt.assert_allclose(
             ds.bout.integrate_midpoints("n", dims=["theta", "zeta"]),
             (2.0 * np.pi * R * 2.0 * np.pi * r)[np.newaxis, :]
@@ -1418,6 +1424,8 @@ def test_integrate_midpoints_salpha(self, bout_xyt_example_files):
         # x-z planes are 'conical frustrums', with area pi*(Rinner + Router)*Lr
         # https://en.wikipedia.org/wiki/Frustum
         theta = _1d_coord_from_spacing(ds["dy"], "theta").values
+        if location == "CELL_YLOW":
+            theta = theta - 2.0 * np.pi / (2.0 * ny)
         Rinner = R + rinner * np.cos(theta)
         Router = R + router * np.cos(theta)
         npt.assert_allclose(
@@ -1549,17 +1557,21 @@ def func(theta):
         )
 
         # Toroidal lines have length 2*pi*Rxy
+        if location == "CELL_CENTRE":
+            R_2d = ds["R"]
+        else:
+            R_2d = ds[f"Rxy_{location}"]
         npt.assert_allclose(
             ds.bout.integrate_midpoints("n", dims=["zeta"]),
-            (2.0 * np.pi * ds["R"]).values[np.newaxis, :, :]
+            (2.0 * np.pi * R_2d).values[np.newaxis, :, :]
             * np.ones(ds.sizes["t"])[:, np.newaxis, np.newaxis],
             rtol=1.0e-5,
             atol=0.0,
         )
         # Integrate in time too
         npt.assert_allclose(
             ds.bout.integrate_midpoints("n", dims=["t", "zeta"]),
-            T_total * (2.0 * np.pi * ds["R"]),
+            T_total * (2.0 * np.pi * R_2d),
             rtol=1.0e-5,
             atol=0.0,
         )

xbout/tests/utils_for_tests.py

@@ -68,14 +68,19 @@ def set_geometry_from_input_file(ds, name):
         "dx",
         "dy",
     ]:
-        # Need all arrays returned from options.evaluate() to be the right shape.
-        # Recommend adding '0*x' or '0*y' in the input file expressions if the
-        # expression would be 1d otherwise.
-        ds[v] = ds[v].copy(
-            data=np.broadcast_to(
-                options.evaluate(f"mesh:{v}").squeeze(axis=2)[slices], shape_2d
+        for location in ["CELL_CENTRE", "CELL_XLOW", "CELL_YLOW", "CELL_ZLOW"]:
+            suffix = "" if location == "CELL_CENTRE" else f"_{location}"
+            # Need all arrays returned from options.evaluate() to be the right shape.
+            # Recommend adding '0*x' or '0*y' in the input file expressions if the
+            # expression would be 1d otherwise.
+            ds[v + suffix] = ds[v].copy(
+                data=np.broadcast_to(
+                    options.evaluate(f"mesh:{v}", location=location).squeeze(axis=2)[
+                        slices
+                    ],
+                    shape_2d,
+                )
             )
-        )
 
     # Set dz as it would be calculated by BOUT++ (but don't support zmin, zmax or
     # zperiod here)
@@ -86,14 +91,19 @@ def set_geometry_from_input_file(ds, name):
 
     # Add extra fields needed by "toroidal" geometry
     for v in ["Rxy", "Zxy", "psixy"]:
-        # Need all arrays returned from options.evaluate() to be the right shape.
-        # Recommend adding '0*x' or '0*y' in the input file expressions if the
-        # expression would be 1d otherwise.
-        ds[v] = ds["g11"].copy(
-            data=np.broadcast_to(
-                options.evaluate(f"mesh:{v}").squeeze(axis=2)[slices], shape_2d
+        for location in ["CELL_CENTRE", "CELL_XLOW", "CELL_YLOW", "CELL_ZLOW"]:
+            suffix = "" if location == "CELL_CENTRE" else f"_{location}"
+            # Need all arrays returned from options.evaluate() to be the right shape.
+            # Recommend adding '0*x' or '0*y' in the input file expressions if the
+            # expression would be 1d otherwise.
+            ds[v + suffix] = ds["g11"].copy(
+                data=np.broadcast_to(
+                    options.evaluate(f"mesh:{v}", location=location).squeeze(axis=2)[
+                        slices
+                    ],
+                    shape_2d,
+                )
             )
-        )
 
     # Set fields that don't have to be in input files to NaN
     for v in ["G1", "G2", "G3"]: