Repository: PlasmaPy/PlasmaPy
Pull request #1041 (2021-02-25 01:02:15)

Formulary Documentation: Adding Missing Formulas to Docstrings
<!--
Thank you for contributing to PlasmaPy! Here's a bunch of pointers to
make things easier for all of us:

* If this PR will solve an issue tracked by GitHub, then please add
  "Closes #42" so the issue automatically closes once this pull request
  is merged.  In this example, issue #42 would have been closed.
  If your PR will not completely solve the issue, then please still reference
  still reference the issue like "issue #42".  (For more info see [GitHub's primer
  on linking issues](https://help.github.com/en/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword).)

* Remember to add some description of your changes in this text box.

* If your pull request is not yet ready for review - either because
  you're just looking for feedback on a change or because you're not
  perfectly satisfied with your change - submit it as a draft pull request.
  Remember to change its' status once it's ready.

* If this is your first contribution, please add your name to the author
  list in `docs/about/credits.rst`.

* Feel free to chat with other developers on our Matrix channel at:
   https://app.element.io/#/room/#plasmapy:openastronomy.org

* We have a developer's guide to help answer some of your questions.
  http://docs.plasmapy.org/en/latest/development/index.html

Many thanks in advance for following these pointers and for being willing to contribute!

When submitting a pull request, please ensure that you can (eventually,
sometime before it is merged) check the following basic requirements:

-->

- [x] I have added a changelog entry for this pull request.

<!--

In short: A changelog entry is a short description of your PR's changes.
Each entry is written in a `<PULL REQUEST>.<TYPE>.rst` file and stored in
the `changelog` directory,  where `<PULL REQUEST>` is a pull request
number and `<TYPE>` is one of:

* `breaking`: A change which requires users to change code and is not backwards compatible. (Not to be used for removal of deprecated features.)
* `feature`: New user facing features and any new behavior.
* `bugfix`: Fixes a reported bug.
* `doc`: Documentation addition or improvement, like rewording an entire session or adding missing docs.
* `removal`: Feature deprecation and/or feature removal.
* `trivial`: A change which has no user facing effect or is tiny change.

A PR number is generated after you successfully submit a new PR, so the changelog
file can only be added after you open the PR.```

For more information, see:
https://github.com/PlasmaPy/PlasmaPy/blob/master/changelog/README.rst

--->

- [ ] If adding new functionality, I have added tests and
      docstrings.

<!--
(Tests pop up at the bottom, in the checks box.)
-->

- [ ] I have fixed any newly failing tests.

<!--
(If you're unsure why they're failing, ask!)
-->

This PR is for Issue #935.

I went through the documentation of `plasmapy.formulary`, and from what I can tell, outside of `plasmapy.formulary.braginskii.py`, the only two (user-facing) functions that lack formulas are `dimensionless.quantum_theta` and `dimensionless.beta`. These quantities are simple ratios, but I believe they still should get complete documentation. Formulas for these two functions are added in this PR.

I plan to add formulas to the docstrings in `braginskii.py` in this PR although these may be more complicated to document because of multiple supported models of classical transport.

On a related note, I am not sure whether the methods of classes, such as the Classical Transport class in `braginskii.py`, should be documented in the same way (or at all) as standalone / free functions, and I would like to hear what other contributors think.

codecov[bot]: # [Codecov](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041?src=pr&el=h1) Report
> Merging [#1041](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041?src=pr&el=desc) (3240297) into [master](https://codecov.io/gh/PlasmaPy/PlasmaPy/commit/6d5651cb9cfb7bd010c8e81c45a1bffeb9c4e95d?el=desc) (6d5651c) will **increase** coverage by `0.64%`.
> The diff coverage is `100.00%`.

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

```diff
@@            Coverage Diff             @@
##           master    #1041      +/-   ##
==========================================
+ Coverage   96.21%   96.86%   +0.64%     
==========================================
  Files          60       70      +10     
  Lines        5417     6855    +1438     
==========================================
+ Hits         5212     6640    +1428     
- Misses        205      215      +10     
```


| [Impacted Files](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [plasmapy/diagnostics/thomson.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041/diff?src=pr&el=tree#diff-cGxhc21hcHkvZGlhZ25vc3RpY3MvdGhvbXNvbi5weQ==) | `100.00% <ø> (ø)` | |
| [plasmapy/dispersion/dispersionfunction.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041/diff?src=pr&el=tree#diff-cGxhc21hcHkvZGlzcGVyc2lvbi9kaXNwZXJzaW9uZnVuY3Rpb24ucHk=) | `100.00% <ø> (ø)` | |
| [plasmapy/formulary/\_\_init\_\_.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041/diff?src=pr&el=tree#diff-cGxhc21hcHkvZm9ybXVsYXJ5L19faW5pdF9fLnB5) | `100.00% <ø> (ø)` | |
| [plasmapy/formulary/collisions.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041/diff?src=pr&el=tree#diff-cGxhc21hcHkvZm9ybXVsYXJ5L2NvbGxpc2lvbnMucHk=) | `97.19% <ø> (+0.43%)` | :arrow_up: |
| [plasmapy/formulary/dimensionless.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041/diff?src=pr&el=tree#diff-cGxhc21hcHkvZm9ybXVsYXJ5L2RpbWVuc2lvbmxlc3MucHk=) | `100.00% <ø> (ø)` | |
| [plasmapy/particles/decorators.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041/diff?src=pr&el=tree#diff-cGxhc21hcHkvcGFydGljbGVzL2RlY29yYXRvcnMucHk=) | `100.00% <ø> (ø)` | |
| [plasmapy/utils/datatype\_factory\_base.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041/diff?src=pr&el=tree#diff-cGxhc21hcHkvdXRpbHMvZGF0YXR5cGVfZmFjdG9yeV9iYXNlLnB5) | `100.00% <ø> (ø)` | |
| [plasmapy/utils/pytest\_helpers/\_\_init\_\_.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041/diff?src=pr&el=tree#diff-cGxhc21hcHkvdXRpbHMvcHl0ZXN0X2hlbHBlcnMvX19pbml0X18ucHk=) | `100.00% <ø> (ø)` | |
| [plasmapy/dispersion/\_\_init\_\_.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041/diff?src=pr&el=tree#diff-cGxhc21hcHkvZGlzcGVyc2lvbi9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |
| [plasmapy/formulary/braginskii.py](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041/diff?src=pr&el=tree#diff-cGxhc21hcHkvZm9ybXVsYXJ5L2JyYWdpbnNraWkucHk=) | `99.72% <100.00%> (ø)` | |
| ... and [48 more](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041?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/1041?src=pr&el=footer). Last update [6d5651c...3240297](https://codecov.io/gh/PlasmaPy/PlasmaPy/pull/1041?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


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/changelog/1041.doc.rst b/changelog/1041.doc.rst
new file mode 100644
index 00000000..141684c5
--- /dev/null
+++ b/changelog/1041.doc.rst
@@ -0,0 +1,1 @@
+Adds formulas (which were missing) to the docstrings of plasmapy.formulary.dimensionless.quantum_theta and plasmapy.formulary.dimensionless.beta.
diff --git a/changelog/1050.trivial.rst b/changelog/1050.trivial.rst
new file mode 100644
index 00000000..e752724e
--- /dev/null
+++ b/changelog/1050.trivial.rst
@@ -0,0 +1,1 @@
+Properly handled warnings in `test_proton_radiography.py`
diff --git a/docs/conf.py b/docs/conf.py
index 8b62a920..3688c3c4 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -44,6 +44,7 @@
     "sphinx.ext.graphviz",
     "sphinx.ext.mathjax",
     "sphinx.ext.napoleon",
+    "sphinx.ext.todo",
     "sphinx_automodapi.automodapi",
     "sphinx_automodapi.smart_resolver",
     "nbsphinx",
diff --git a/docs/development/code_guide.rst b/docs/development/code_guide.rst
index e4091e25..7c5668a6 100644
--- a/docs/development/code_guide.rst
+++ b/docs/development/code_guide.rst
@@ -13,6 +13,42 @@ later.
 Coding Style
 ============
 
+TL;DR: use pre-commit
+---------------------
+
+PlasmaPy has a configuration for the `pre-commit framework
+<https://pre-commit.com/>`_ that takes care of style mostly automatically.
+Install it with `pip install pre-commit`, then use `pre-commit install` within
+the repository.
+
+This will cause pre-commit to download the right versions of linters we use,
+then run an automated style checking suite on every commit.  Do note that this
+works better with a `git add`, then `git commit` workflow than a `git commit
+-a` workflow — that way, you can check via `git diff` what the automated
+changes actually did.
+
+Note that the "Style linters / pre-commit (pull_request)" part of our
+Continuous Integration system can and will (metaphorically) shout at you if it
+finds you didn't apply the linters. Also note that the linters' output may vary
+with version, so, rather than apply `black` and `isort` manually, let
+pre-commit do the version management for you instead!
+
+Our pre-commit suite can be found in `.pre-commit-config.yaml
+<https://github.com/PlasmaPy/PlasmaPy/blob/master/.pre-commit-config.yaml>`_.
+It includes
+
+* `black <https://black.readthedocs.io/en/stable/>`_ to automatically
+  format code and ensure a consistent code style throughout the
+  package
+* `isort <https://pycqa.github.io/isort/>`_ to
+  automatically sort imports.
+* `nbqa <https://github.com/nbQA-dev/nbQA>`_ to automatically apply the above
+  to example notebooks as well.
+* a few tools for `requirements.txt`, `.yml` files and the like.
+
+PlasmaPy Code Style Guide, codified
+-----------------------------------
+
 * PlasmaPy follows the `PEP8 Style Guide for Python Code
   <http://www.python.org/dev/peps/pep-0008/>`_.  This style choice
   helps ensure that the code will be consistent and readable.
@@ -49,25 +85,11 @@ Coding Style
 * Use ``Optional[type]`` for type hinted keyword arguments with a
   default value of ``None``.
 
+* There should be at least one pun per 1284 lines of code.
+
 * Avoid using `lambda` to define functions, as this notation may be
   unfamiliar to newcomers to Python.
 
-* There should be at most one pun per 1284 lines of code.
-
-pre-commit hooks
-----------------
-
-PlasmaPy has a configuration for `pre-commit` that takes care of style mostly
-automatically. Install it with `pip install pre-commit`, then use `pre-commit
-install` within the repository.
-
-This will cause git to run an automated style checking suite, mostly composed
-of `flake8` and `black`, on every commit. Do note that this works better with
-a `git add`, then `git commit` workflow than a `git commit -a` workflow - that
-way, you can check via `git diff` what the automated changes actually did.
-
-Note that `flake8` does not change your files automatically (help in setting
-that up, if possible, most welcome!).
 
 Branches, commits, and pull requests
 ====================================
diff --git a/plasmapy/diagnostics/proton_radiography.py b/plasmapy/diagnostics/proton_radiography.py
index 52e21358..fa8cbccc 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(
                     [
@@ -996,8 +990,8 @@ def synthetic_radiograph(
         if percentage < 0.5:
             warnings.warn(
                 f"Only {percentage:.2%} of the particles are shown "
-                " on this synthetic radiograph. Consider increasing "
-                " the size to include more.",
+                "on this synthetic radiograph. Consider increasing "
+                "the size to include more.",
                 RuntimeWarning,
             )
 
diff --git a/plasmapy/formulary/braginskii.py b/plasmapy/formulary/braginskii.py
index fd42440f..d2381b76 100644
--- a/plasmapy/formulary/braginskii.py
+++ b/plasmapy/formulary/braginskii.py
@@ -457,9 +457,20 @@ def __init__(
     @property
     @validate_quantities
     def resistivity(self) -> u.Ohm * u.m:
-        """
+        r"""
         Calculate the resistivity.
 
+        The resistivity (:math:`\alpha`) of a plasma is defined by
+
+        .. math::
+            \alpha = \frac{\hat{\alpha}}{n_{e} * e^2 * \frac{\tau_{e}}{m_{e}}}
+
+        where :math:`\hat{\alpha}` is the non-dimensional resistivity of the plasma,
+        :math:`n_e` is the electron number density of the plasma,
+        :math:`e` is Euler's number,
+        :math:`\tau_{e}` is the fundamental electron collision period of the plasma,
+        and :math:`m_{e}` is the mass of an electron.
+
         Notes
         -----
 
@@ -492,9 +503,12 @@ def resistivity(self) -> u.Ohm * u.m:
 
     @property
     def thermoelectric_conductivity(self):
-        """
+        r"""
         Calculate the thermoelectric conductivity.
 
+        .. todo::
+            The thermoelectric conductivity (:math:`\hat{\beta}`) of a plasma is defined by...
+
         Notes
         -----
         To be improved.
@@ -512,9 +526,21 @@ def thermoelectric_conductivity(self):
     @property
     @validate_quantities
     def ion_thermal_conductivity(self) -> u.W / u.m / u.K:
-        """
+        r"""
         Calculate the thermal conductivity for ions.
 
+        The ion thermal conductivity (:math:`\kappa`) of a plasma is defined by
+
+        .. math::
+            \kappa = \hat{\kappa} * \frac{n_{i} * k_{B}^2 * T_{i} * \tau_{i}}{m_{i}}
+
+        where :math:`\hat{\kappa}` is the non-dimensional ion thermal conductivity of the plasma,
+        :math:`n_{i}` is the ion number density of the plasma,
+        :math:`k_{B}` is the Boltzmann constant,
+        :math:`T_{i}` is the ion temperature of the plasma,
+        :math:`\tau_{i}` is the fundamental ion collision period of the plasma,
+        and :math:`m_{i}` is the mass of an ion of the plasma.
+
         Notes
         -----
         This is the classical plasma ions' ability to conduct energy and heat,
@@ -553,9 +579,21 @@ def ion_thermal_conductivity(self) -> u.W / u.m / u.K:
     @property
     @validate_quantities
     def electron_thermal_conductivity(self) -> u.W / u.m / u.K:
-        """
+        r"""
         Calculate the thermal conductivity for electrons.
 
+        The electron thermal conductivity (:math:`\kappa`) of a plasma is defined by
+
+        .. math::
+            \kappa = \hat{\kappa} * \frac{n_{e} * k_{B}^2 * T_{e} * \tau_{e}}{m_{e}}
+
+        where :math:`\hat{\kappa}` is the non-dimensional electron thermal conductivity of the plasma,
+        :math:`n_{e}` is the electron number density of the plasma,
+        :math:`k_{B}` is the Boltzmann constant,
+        :math:`T_{e}` is the electron temperature of the plasma,
+        :math:`\tau_{e}` is the fundamental electron collision period of the plasma,
+        and :math:`m_{e}` is the mass of an electron.
+
         Notes
         -----
         This is quite similar to the ion thermal conductivity, except that it's
@@ -606,9 +644,12 @@ def electron_thermal_conductivity(self) -> u.W / u.m / u.K:
     @property
     @validate_quantities
     def ion_viscosity(self) -> u.Pa * u.s:
-        """
+        r"""
         Calculate the ion viscosity.
 
+        .. todo::
+            The ion viscosity (:math:`\eta`) of a plasma is defined by...
+
         Notes
         -----
         This is the dynamic viscosity that you find for ions in the classical
@@ -651,9 +692,12 @@ def ion_viscosity(self) -> u.Pa * u.s:
     @property
     @validate_quantities
     def electron_viscosity(self) -> u.Pa * u.s:
-        """
+        r"""
         Calculate the electron viscosity.
 
+        .. todo::
+            The electron viscosity (:math:`\eta`) of a plasma is defined by...
+
         Notes
         -----
         This is the dynamic viscosity that you find for electrons in the
@@ -753,9 +797,21 @@ def resistivity(
     theta=None,
     coulomb_log_method="classical",
 ) -> u.Ohm * u.m:
-    """
+    r"""
     Calculate the resistivity.
 
+    The resistivity (:math:`\alpha`) of a plasma is defined by
+
+    .. math::
+
+        \alpha = \frac{\hat{\alpha}}{n_{e} * e^2 * \frac{\tau_{e}}{m_{e}}}
+
+    where :math:`\hat{\alpha}` is the non-dimensional resistivity of the plasma,
+    :math:`n_e` is the electron number density of the plasma,
+    :math:`e` is Euler's number,
+    :math:`\tau_{e}` is the fundamental electron collision period of the plasma,
+    and :math:`m_{e}` is the mass of an electron.
+
     Notes
     -----
 
@@ -810,7 +866,12 @@ def thermoelectric_conductivity(
     theta=None,
     coulomb_log_method="classical",
 ):
-    """Calculate the thermoelectric conductivity."""
+    r"""
+    Calculate the thermoelectric conductivity.
+
+    .. todo::
+        The thermoelectric conductivity (:math:`\hat{\beta}`) of a plasma is defined by...
+    """
     ct = ClassicalTransport(
         T_e,
         n_e,
@@ -845,9 +906,22 @@ def ion_thermal_conductivity(
     theta=None,
     coulomb_log_method="classical",
 ) -> u.W / u.m / u.K:
-    """
+    r"""
     Calculate the thermal conductivity for ions.
 
+    The ion thermal conductivity (:math:`\kappa`) of a plasma is defined by
+
+    .. math::
+
+        \kappa = \hat{\kappa} * \frac{n_{i} * k_{B}^2 * T_{i} * \tau_{i}}{m_{i}}
+
+    where :math:`\hat{\kappa}` is the non-dimensional ion thermal conductivity of the plasma,
+    :math:`n_{i}` is the ion number density of the plasma,
+    :math:`k_{B}` is the Boltzmann constant,
+    :math:`T_{i}` is the ion temperature of the plasma,
+    :math:`\tau_{i}` is the fundamental ion collision period of the plasma,
+    and :math:`m_{i}` is the mass of an ion of the plasma.
+
     Notes
     -----
     This is the classical plasma ions' ability to conduct energy and heat,
@@ -902,9 +976,22 @@ def electron_thermal_conductivity(
     theta=None,
     coulomb_log_method="classical",
 ) -> u.W / u.m / u.K:
-    """
+    r"""
     Calculate the thermal conductivity for electrons.
 
+    The electron thermal conductivity (:math:`\kappa`) of a plasma is defined by
+
+    .. math::
+
+        \kappa = \hat{\kappa} * \frac{n_{e} * k_{B}^2 * T_{e} * \tau_{e}}{m_{e}}
+
+    where :math:`\hat{\kappa}` is the non-dimensional electron thermal conductivity of the plasma,
+    :math:`n_{e}` is the electron number density of the plasma,
+    :math:`k_{B}` is the Boltzmann constant,
+    :math:`T_{e}` is the electron temperature of the plasma,
+    :math:`\tau_{e}` is the fundamental electron collision period of the plasma,
+    and :math:`m_{e}` is the mass of an electron.
+
     Notes
     -----
     This is quite similar to the ion thermal conductivity, except that it's for
@@ -971,9 +1058,12 @@ def ion_viscosity(
     theta=None,
     coulomb_log_method="classical",
 ) -> u.Pa * u.s:
-    """
+    r"""
     Calculate the ion viscosity.
 
+    .. todo::
+        The ion viscosity (:math:`\eta`) of a plasma is defined by...
+
     Notes
     -----
     This is the dynamic viscosity that you find for ions in the classical
@@ -1025,9 +1115,12 @@ def electron_viscosity(
     theta=None,
     coulomb_log_method="classical",
 ) -> u.Pa * u.s:
-    """
+    r"""
     Calculate the electron viscosity.
 
+    .. todo::
+        The electron viscosity (:math:`\eta`) of a plasma is defined by...
+
     Notes
     -----
     This is the dynamic viscosity that you find for electrons in the
diff --git a/plasmapy/formulary/dimensionless.py b/plasmapy/formulary/dimensionless.py
index f4d297a9..8465ff6c 100644
--- a/plasmapy/formulary/dimensionless.py
+++ b/plasmapy/formulary/dimensionless.py
@@ -31,10 +31,18 @@
     n_e={"can_be_negative": False},
 )
 def quantum_theta(T: u.K, n_e: u.m ** -3) -> u.dimensionless_unscaled:
-    """
+    r"""
     Compares Fermi energy to thermal kinetic energy to check if quantum
     effects are important.
 
+    The quantum theta (:math:`\theta`) of a plasma is defined by
+
+    .. math::
+        \theta = \frac{E_{T}}{E_{F}}
+
+    where :math:`E_{T}` is the thermal energy of the plasma
+    and :math:`E_{F}` is the Fermi energy of the plasma.
+
     Parameters
     ----------
     T : ~astropy.units.Quantity
@@ -58,6 +66,19 @@ def quantum_theta(T: u.K, n_e: u.m ** -3) -> u.dimensionless_unscaled:
     -------
     theta: ~astropy.units.Quantity
 
+    Notes
+    -----
+    The thermal energy of the plasma (:math:`E_{T}`) is defined by
+
+    .. math::
+        E_{T} = k_{B} * T
+
+    where :math:`k_{B}` is the Boltzmann constant
+    and :math:`T` is the temperature of the plasma.
+
+    See Also
+    --------
+    plasmapy.formulary.quantum.Fermi_energy
     """
     fermi_energy = quantum.Fermi_energy(n_e)
     thermal_energy = constants.k_B * T
@@ -70,9 +91,17 @@ def quantum_theta(T: u.K, n_e: u.m ** -3) -> u.dimensionless_unscaled:
     n={"can_be_negative": False},
 )
 def beta(T: u.K, n: u.m ** -3, B: u.T) -> u.dimensionless_unscaled:
-    """
+    r"""
     The ratio of thermal pressure to magnetic pressure.
 
+    The beta (:math:`\beta`) of a plasma is defined by
+
+    .. math::
+        \beta = \frac{p_{th}}{p_{mag}}
+
+    where :math:`p_{th}` is the thermal pressure of the plasma
+    and :math:`p_{mag}` is the magnetic pressure of the plasma.
+
     Parameters
     ----------
     T : ~astropy.units.Quantity
@@ -95,6 +124,10 @@ def beta(T: u.K, n: u.m ** -3, B: u.T) -> u.dimensionless_unscaled:
     beta: ~astropy.units.Quantity
         Dimensionless quantity.
 
+    See Also
+    --------
+    plasmapy.formulary.parameters.thermal_pressure
+    plasmapy.formulary.parameters.magnetic_pressure
     """
     thermal_pressure = parameters.thermal_pressure(T, n)
     magnetic_pressure = parameters.magnetic_pressure(B)
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..f657c6db 100644
--- a/plasmapy/diagnostics/tests/test_proton_radiography.py
+++ b/plasmapy/diagnostics/tests/test_proton_radiography.py
@@ -123,9 +123,11 @@ def run_1D_example(name):
     source = (0 * u.mm, -10 * u.mm, 0 * u.mm)
     detector = (0 * u.mm, 200 * u.mm, 0 * u.mm)
 
-    # TODO: Catch warnings (test fields aren't well behaved at edges)
-
-    sim = prad.SyntheticProtonRadiograph(grid, source, detector, verbose=False)
+    # Expect warnings because these fields aren't well-behaved at the edges
+    with pytest.warns(
+        RuntimeWarning, match="Fields should go to zero at edges of grid to avoid "
+    ):
+        sim = prad.SyntheticProtonRadiograph(grid, source, detector, verbose=False)
     sim.create_particles(1e4, 3 * u.MeV, max_theta=0.1 * u.deg)
     sim.run()
 
@@ -194,21 +196,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
@@ -221,7 +228,7 @@ def test_input_validation():
 
     # Test raises warning when one (or more) of the required fields is missing
     grid_bad = CartesianGrid(-1 * u.mm, 1 * u.mm, num=50)
-    with pytest.warns(RuntimeWarning):
+    with pytest.warns(RuntimeWarning, match="is not specified for the provided grid."):
         sim = prad.SyntheticProtonRadiograph(grid_bad, source, detector, verbose=True)
 
     # ************************************************************************
@@ -249,7 +256,9 @@ def test_input_validation():
 
     # Choose a very small synthetic radiograph size that misses most of the
     # particles
-    with pytest.warns(RuntimeWarning):
+    with pytest.warns(
+        RuntimeWarning, match="of the particles are shown on this synthetic radiograph."
+    ):
         size = np.array([[-1, 1], [-1, 1]]) * 1 * u.mm
         hax, vax, values = sim.synthetic_radiograph(size=size)
 
@@ -341,7 +350,7 @@ def test_run_options():
     # Cartesian
     source = (0 * u.mm, -10 * u.mm, 0 * u.mm)
     detector = (0 * u.mm, 200 * u.mm, 0 * u.mm)
-    sim = prad.SyntheticProtonRadiograph(grid, source, detector, verbose=False)
+    sim = prad.SyntheticProtonRadiograph(grid, source, detector, verbose=True)
 
     # Test that trying to call run() without creating particles
     # raises an exception
@@ -359,7 +368,7 @@ def test_run_options():
 
     # Test way too big of a max_theta
     sim.create_particles(1e4, 3 * u.MeV, max_theta=89 * u.deg)
-    with pytest.warns(RuntimeWarning):
+    with pytest.warns(RuntimeWarning, match="of " "particles entered the field grid"):
         sim.run(field_weighting="nearest neighbor", dt=1e-12 * u.s)
 
     # Test extreme deflections -> warns user
@@ -368,9 +377,17 @@ def test_run_options():
     grid = _test_grid("constant_bz", num=50, B0=250 * u.T)
     source = (0 * u.mm, -10 * u.mm, 0 * u.mm)
     detector = (0 * u.mm, 200 * u.mm, 0 * u.mm)
-    sim = prad.SyntheticProtonRadiograph(grid, source, detector, verbose=False)
+
+    # Expectwarnings because these fields aren't well-behaved at the edges
+    with pytest.warns(
+        RuntimeWarning, match="Fields should go to zero at edges of grid to avoid "
+    ):
+        sim = prad.SyntheticProtonRadiograph(grid, source, detector, verbose=False)
     sim.create_particles(1e4, 3 * u.MeV, max_theta=0.1 * u.deg)
-    with pytest.warns(RuntimeWarning):
+    with pytest.warns(
+        RuntimeWarning,
+        match="particles have been " "deflected away from the detector plane",
+    ):
         sim.run(field_weighting="nearest neighbor", dt=1e-12 * u.s)
     # Calc max deflection: should be between 0 and pi/2
     # Note: that's only true because max_theta is very small
@@ -401,6 +418,7 @@ def test_synthetic_radiograph():
 
 
 if __name__ == "__main__":
+    """
     test_coordinate_systems()
     test_input_validation()
     test_1D_deflections()
@@ -409,4 +427,5 @@ def test_synthetic_radiograph():
     test_load_particles()
     test_run_options()
     test_synthetic_radiograph()
+    """
     pass
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