Repository: PlasmaPy/PlasmaPy
Pull request #1027 (2021-02-21 00:03:22)

Grids 4
A place for more minor changes to `plasmapy.plasma.grids`. 

Changelog: 

- `grid.__getitem__(key)` now returns a `u.Quantity` array instead of an `xarray.DataArray` object. The units are more apparent this way (not buried in the attrs).

- Added `grid.require_quantities` method to make sure that a list of required quantities are present on a grid. @StanczakDominik  suggested this on #895 to replace some functionality in the proton radiography `__init__` function. This (very minor) refactor is also included in this PR.

- Fixed bugs in `grids.py` for non-uniform grids that arose when xarray upgraded to v0.17.0.

codecov[bot]: # [Codecov](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027?src=pr&el=h1) Report
> Merging [#1027](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027?src=pr&el=desc) (8793f8a) into [master](https://codecov.io/gh/PlasmaPy/PlasmaPy/commit/8a4d87d052bf59410becc480f409f8acf72f8f24?el=desc) (8a4d87d) will **decrease** coverage by `0.78%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027/graphs/tree.svg?width=650&height=150&src=pr&token=LzZ5mnXVad)](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #1027      +/-   ##
==========================================
- Coverage   96.66%   95.87%   -0.79%     
==========================================
  Files          64       64              
  Lines        6262     6262              
==========================================
- Hits         6053     6004      -49     
- Misses        209      258      +49     
```


| [Impacted Files](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [plasmapy/plasma/grids.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027/diff?src=pr&el=tree#diff-cGxhc21hcHkvcGxhc21hL2dyaWRzLnB5) | `94.67% <100.00%> (ø)` | |
| [plasmapy/simulation/particle\_integrators.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027/diff?src=pr&el=tree#diff-cGxhc21hcHkvc2ltdWxhdGlvbi9wYXJ0aWNsZV9pbnRlZ3JhdG9ycy5weQ==) | `33.33% <0.00%> (-66.67%)` | :arrow_down: |
| [plasmapy/simulation/particletracker.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027/diff?src=pr&el=tree#diff-cGxhc21hcHkvc2ltdWxhdGlvbi9wYXJ0aWNsZXRyYWNrZXIucHk=) | `36.20% <0.00%> (-62.07%)` | :arrow_down: |
| [plasmapy/plasma/sources/plasma3d.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027/diff?src=pr&el=tree#diff-cGxhc21hcHkvcGxhc21hL3NvdXJjZXMvcGxhc21hM2QucHk=) | `93.75% <0.00%> (-6.25%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027?src=pr&el=footer). Last update [8a4d87d...8793f8a](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1027?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

StanczakDominik: Would you mind if we commandeered this PR to also fix the issue of the new xarray 0.17 breaking our stuff, as found out by #1045?

Patch:
diff --git a/changelog/1027.breaking.rst b/changelog/1027.breaking.rst
new file mode 100644
index 00000000..aeccbe32
--- /dev/null
+++ b/changelog/1027.breaking.rst
@@ -0,0 +1,5 @@
+Several changes to `grids.py` (and one refactor in `diagnostics.proton_radiography`
+- (BREAKING)`grid.__getitem__` now returns a `u.Quantity` array instead of a reference to a `xarray.DataArray`.
+- Added the `grid.require_quantities` method that verifies that a list of quantities are present on a grid.
+- Used `grid.require_quantities` to refactor code in `diagnostics.proton_radiography.SyntheticProtonRadiograph`
+- Fixed a bug in `grids.py` for non-uniform grids that arose when `xarray` upgraded to `v0.17.0`
diff --git a/plasmapy/diagnostics/proton_radiography.py b/plasmapy/diagnostics/proton_radiography.py
index 52e21358..257b088e 100644
--- a/plasmapy/diagnostics/proton_radiography.py
+++ b/plasmapy/diagnostics/proton_radiography.py
@@ -213,22 +213,13 @@ def __init__(
         # ************************************************************************
 
         req_quantities = ["E_x", "E_y", "E_z", "B_x", "B_y", "B_z"]
-        for rq in req_quantities:
 
-            # Error check that grid contains E and B variables required
-            if rq not in self.grid.quantities:
-                warnings.warn(
-                    f"{rq} not specified for provided grid."
-                    "This quantity will be assumed to be zero.",
-                    RuntimeWarning,
-                )
-                # If missing, warn user and then replace with an array of zeros
-                unit = self.grid._recognized_quantities[rq].unit
-                arg = {rq: np.zeros(self.grid.shape) * unit}
-                self.grid.add_quantities(**arg)
+        self.grid.require_quantities(req_quantities, replace_with_zeros=True)
+
+        for rq in req_quantities:
 
             # Check that there are no infinite values
-            if not np.isfinite(self.grid[rq]).all():
+            if not np.isfinite(self.grid[rq].value).all():
                 raise ValueError(
                     f"Input arrays must be finite: {rq} contains "
                     "either NaN or infinite values."
@@ -236,7 +227,10 @@ def __init__(
 
             # Check that the max values on the edges of the arrays are
             # small relative to the maximum values on that grid
-            arr = np.abs(self.grid[rq])
+            #
+            # Array must be dimensionless to re-assemble it into an array
+            # of max values like this
+            arr = np.abs(self.grid[rq]).value
             edge_max = np.max(
                 np.array(
                     [
diff --git a/plasmapy/plasma/grids.py b/plasmapy/plasma/grids.py
index d04007b7..61a1bd34 100644
--- a/plasmapy/plasma/grids.py
+++ b/plasmapy/plasma/grids.py
@@ -111,6 +111,70 @@ def recognized_quantities(self):
         """
         return self._recognized_quantities
 
+    def require_quantities(self, req_quantities, replace_with_zeros=False):
+        r"""
+        Checks to make sure that a list of required quantities are present.
+        Optionally, can create missing quantities and fill them with
+        an array of zeros.
+
+        Parameters
+        ----------
+        req_quantities : list of str
+            A list of quantity keys that are required
+
+        replace_with_zeros : boolean, optional
+            If true, missing quantities will be replaced with an array
+            of zeros. If false, an exception will be raised instead.
+            The default is False.
+
+        Raises
+        ------
+        KeyError
+            If `replace_with_zeros` is False and a required quantity is missing,
+            raises a KeyError.
+
+        KeyError
+            If `replace_with_zeros` is True but the quantity is not in the
+            list of recognized quantities, raises a KeyError. This is because
+            in this case the units for the quantity are unknown, so an array
+            of zeros cannot be constructed.
+
+        Returns
+        -------
+        None.
+
+        """
+        for rq in req_quantities:
+
+            # Error check that grid contains E and B variables required
+            if rq not in self.quantities:
+
+                # If missing, warn user and then replace with an array of zeros
+                if replace_with_zeros:
+                    warnings.warn(
+                        f"{rq} is not specified for the provided grid."
+                        "This quantity will be assumed to be zero.",
+                        RuntimeWarning,
+                    )
+
+                    if rq in self.recognized_quantities.keys():
+                        unit = self.recognized_quantities[rq].unit
+                    else:
+                        raise KeyError(
+                            f"{rq} is not a recognized key, and "
+                            "so cannot be automatically assumed "
+                            "to be zero."
+                        )
+
+                    arg = {rq: np.zeros(self.shape) * unit}
+                    self.add_quantities(**arg)
+
+                else:
+                    raise KeyError(
+                        f"{rq} is not specified for the provided "
+                        "grid but is required."
+                    )
+
     # *************************************************************************
     # Fundamental properties of the grid
     # *************************************************************************
@@ -168,8 +232,14 @@ def __repr__(self):
         return s
 
     def __getitem__(self, key):
+        """
+        Given a key, return the corresponding array as an `astropy.Quantity`
 
-        return self.ds[key]
+        Returning with copy=False means that the array returned is a direct
+        reference to the underlying DataArray, so changes made will be reflected
+        in the underlying DataArray.
+        """
+        return u.Quantity(self.ds[key].data, self.ds[key].attrs["unit"], copy=False)
 
     @property
     def shape(self):
@@ -498,11 +568,17 @@ def add_quantities(self, **kwargs):
                 )
 
             if self.is_uniform:
-                axes = ["ax0", "ax1", "ax2"]
+                dims = ["ax0", "ax1", "ax2"]
+                coords = {
+                    "ax0": self.ds.coords["ax0"],
+                    "ax1": self.ds.coords["ax1"],
+                    "ax2": self.ds.coords["ax2"],
+                }
             # If grid is non-uniform, flatten quantity
             else:
                 quantity = quantity.flatten()
-                axes = ["ax"]
+                dims = ["ax"]
+                coords = {"ax": self.ds.coords["ax"]}
 
             if quantity.shape != self.shape:
                 raise ValueError(
@@ -510,7 +586,9 @@ def add_quantities(self, **kwargs):
                     f"does not match the grid shape {self.shape}."
                 )
 
-            data = xr.DataArray(quantity, dims=axes, attrs={"unit": quantity.unit})
+            data = xr.DataArray(
+                quantity, dims=dims, coords=coords, attrs={"unit": quantity.unit}
+            )
             self.ds[key] = data
 
     @property
diff --git a/requirements/environment.yml b/requirements/environment.yml
index 879416cb..8ed50e10 100644
--- a/requirements/environment.yml
+++ b/requirements/environment.yml
@@ -15,7 +15,7 @@ dependencies:
   - pandas
   - pillow
   - pip
-  - xarray
+  - xarray < 0.17
   - pip:
     - numpydoc
     - pytest
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 8b79507e..2228fac2 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -11,4 +11,4 @@ pandas >= 1.0.0
 scipy >= 1.2
 setuptools >= 41.2
 tqdm >= 4.56.0
-xarray  >= 0.14.0
+xarray  >= 0.14.0,<0.17.0
diff --git a/setup.cfg b/setup.cfg
index e7a3bbff..a3df6fdd 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -38,7 +38,7 @@ install_requires =
   scipy >= 1.2
   setuptools >= 41.2
   tqdm >= 4.56.0
-  xarray >= 0.14.0
+  xarray >= 0.14.0, <0.17.0
 
 [options.extras_require]
 classes =
@@ -83,6 +83,10 @@ doctest_optionflags =
     ELLIPSIS
     NUMBER
 addopts = --doctest-modules --doctest-continue-on-failure --ignore=docs/conf.py
+filterwarnings =
+    ignore:.*Creating a LegacyVersion.*:DeprecationWarning
+
+
 
 [pycodestyle]
 # E101 - mix of tabs and spaces
diff --git a/setup.py b/setup.py
index 25d3bd3f..e56ab9ab 100644
--- a/setup.py
+++ b/setup.py
@@ -1,8 +1,14 @@
 #!/usr/bin/env python
+# https://github.com/pypa/pip/issues/7953#issuecomment-645133255
+import site
+import sys
+
 from itertools import chain
 from setuptools import setup
 from setuptools.config import read_configuration
 
+site.ENABLE_USER_SITE = "--user" in sys.argv[1:]
+
 ################################################################################
 # Programmatically generate some extras combos.
 ################################################################################
diff --git a/tox.ini b/tox.ini
index 309b17b8..3fb142d6 100644
--- a/tox.ini
+++ b/tox.ini
@@ -14,6 +14,7 @@ extras = all,tests
 deps =
     numpydev: git+https://github.com/numpy/numpy
     astropydev: git+https://github.com/astropy/astropy
+    xarraydev: git+https://github.com/pydata/xarray
     cov: pytest-cov
     !minimal: pytest-xdist
     pytest-github-actions-annotate-failures
@@ -26,6 +27,7 @@ description =
     run tests
     numpydev: with the git master version of numpy
     astropydev: with the git master version of astropy
+    xarraydev: with the git master version of xarray
     minimal: with minimal versions of dependencies
     cov: with code coverage
 


Test patch:
diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml
index 059c244b..604ab28f 100644
--- a/.github/workflows/testing.yml
+++ b/.github/workflows/testing.yml
@@ -48,6 +48,11 @@ jobs:
             python: 3.8
             toxenv: py38-numpydev
 
+          - name: Python 3.9 with Xarray dev
+            os: ubuntu-latest
+            python: 3.9
+            toxenv: py39-numpydev
+
           - name: Python 3.9
             os: ubuntu-latest
             python: 3.9
diff --git a/plasmapy/diagnostics/tests/test_proton_radiography.py b/plasmapy/diagnostics/tests/test_proton_radiography.py
index 214f24a3..6b7e5ca8 100644
--- a/plasmapy/diagnostics/tests/test_proton_radiography.py
+++ b/plasmapy/diagnostics/tests/test_proton_radiography.py
@@ -194,21 +194,26 @@ def test_input_validation():
 
     # Check that an error is raised when an input grid has a nan or infty value
     # First check NaN
-    grid["E_x"][0, 0, 0] = np.nan * u.V / u.m
+    Ex = grid["E_x"]
+    Ex[0, 0, 0] = np.nan * u.V / u.m
+    grid.add_quantities(E_x=Ex)
     with pytest.raises(ValueError):
         sim = prad.SyntheticProtonRadiograph(grid, source, detector, verbose=False)
+    Ex[0, 0, 0] = 0 * u.V / u.m
 
-    grid["E_x"][0, 0, 0] = np.inf * u.V / u.m  # Reset element for the rest of the tests
+    Ex[0, 0, 0] = np.inf * u.V / u.m  # Reset element for the rest of the tests
+    grid.add_quantities(E_x=Ex)
     with pytest.raises(ValueError):
         sim = prad.SyntheticProtonRadiograph(grid, source, detector, verbose=False)
-    grid["E_x"][0, 0, 0] = 0 * u.V / u.m
+    Ex[0, 0, 0] = 0 * u.V / u.m
 
     # Check what happens if a value is large realtive to the rest of the array
-    grid["E_x"][0, 0, 0] = 0.5 * np.max(grid["E_x"])
+    Ex[0, 0, 0] = 0.5 * np.max(Ex)
+    grid.add_quantities(E_x=Ex)
     # with pytest.raises(ValueError):
     with pytest.warns(RuntimeWarning):
         sim = prad.SyntheticProtonRadiograph(grid, source, detector, verbose=False)
-    grid["E_x"][0, 0, 0] = 0 * u.V / u.m
+    Ex[0, 0, 0] = 0 * u.V / u.m
 
     # Raise error when source-to-detector vector doesn't pass through the
     # field grid
diff --git a/plasmapy/plasma/tests/test_grids.py b/plasmapy/plasma/tests/test_grids.py
index a32eeeaf..d58308f5 100644
--- a/plasmapy/plasma/tests/test_grids.py
+++ b/plasmapy/plasma/tests/test_grids.py
@@ -64,9 +64,35 @@ def test_AbstractGrid():
     q = np.random.randn(10, 10, 10) * u.T
     grid.add_quantities(B_x=q)
 
+    # Test setting a subset of a quantity array
+    grid["B_x"][0, 0, 0] = 21 * u.T
+    assert grid["B_x"][0, 0, 0] == 21 * u.T
+
     # Test accessing a quantity using __getitem__ or directly
     Bx = grid.ds["B_x"]
     Bx = grid["B_x"]
+    # Assert that the array returned is a u.Quantity
+    assert isinstance(Bx, u.Quantity)
+    # Assert that the array returned has the right shape
+    assert Bx.shape == grid.shape
+
+    # Test require_quantities
+    # Test with a key that is there
+    req_q = ["B_x"]
+    grid.require_quantities(req_q, replace_with_zeros=False)
+    req_q = ["B_x", "B_y"]
+    # Test with a key that is not there, but can be replaced
+    # Do not replace
+    with pytest.raises(KeyError):
+        grid.require_quantities(req_q, replace_with_zeros=False)
+    # Do replace
+    grid.require_quantities(req_q, replace_with_zeros=True)
+    req_q = ["B_x", "B_y"]
+    # Test with a key that is not there, but cannot be replaced because
+    # it's not a recognized key
+    req_q = ["B_x", "not_a_recognized_key"]
+    with pytest.raises(KeyError):
+        grid.require_quantities(req_q, replace_with_zeros=True)
 
     # Test adding a quantity with wrong units
     q = np.random.randn(10, 10, 10) * u.kg
@@ -221,6 +247,9 @@ def test_nearest_neighbor_interpolator():
 
     # Create a non-uniform grid
     grid = grids.NonUniformCartesianGrid(-1 * u.cm, 1 * u.cm, num=100)
+
+    print(grid.shape)
+
     grid.add_quantities(x=grid.grids[0], y=grid.grids[1])
 
     # One position
@@ -271,12 +300,13 @@ def test_volume_averaged_interpolator():
 
 def test_NonUniformCartesianGrid():
     grid = grids.NonUniformCartesianGrid(-1 * u.cm, 1 * u.cm, num=10)
-    grid.add_quantities(x=grid.grids[0])
 
     pts0, pts1, pts2 = grid.grids
+
     shape = grid.shape
     units = grid.units
 
+    grid.add_quantities(x=pts0)
     print(grid)
 
     # Grid should be non-uniform
@@ -316,11 +346,13 @@ def test_NonUniformCartesianGrid():
 
 
 if __name__ == "__main__":
-    # test_AbstractGrid()
+    """
+    test_AbstractGrid()
     test_CartesianGrid()
-    # test_grid_methods()
-    # test_interpolate_indices()
-    # test_nearest_neighbor_interpolator()
-    # test_volume_averaged_interpolator()
+    test_grid_methods()
+    test_interpolate_indices()
+    test_nearest_neighbor_interpolator()
+    test_volume_averaged_interpolator()
     test_NonUniformCartesianGrid()
+    """
     pass