Repository: arviz-devs/arviz

Issue:
Error plotting a single variable with plot_density and bokeh backend
**Describe the bug**
Over in ArviZ.jl, we use the Julia equivalent to the below snippet to test Bokeh integration for `plot_density`. It worked fine until recently, where we now get an error with bokeh only but not matplotlib, though I'm not certain whether a change in arviz or bokeh is responsible.

**To Reproduce**
```python
>>> import arviz
>>> import numpy as np
>>> import matplotlib.pyplot as plt
>>> arr1 = np.random.randn(4, 100)
>>> arr2 = np.random.randn(4, 100)
>>> arviz.plot_density([{"x": arr1}, {"x": arr2}], var_names = ["x"]) # matplotlib works fine
>>> plt.show()
```
<img src=https://user-images.githubusercontent.com/8673634/94775414-9bce2480-0374-11eb-8938-f74a486f97de.png width=400></img>
```python
>>> arviz.plot_density([{"x": arr1}, {"x": arr2}], var_names = ["x"], backend="bokeh") # errors
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/saxen/.julia/conda/3/lib/python3.8/site-packages/arviz/plots/densityplot.py", line 252, in plot_density
    ax = plot(**plot_density_kwargs)
  File "/Users/saxen/.julia/conda/3/lib/python3.8/site-packages/arviz/plots/backends/bokeh/densityplot.py", line 74, in plot_density
    for label, ax_ in zip(all_labels, (item for item in ax.flatten() if item is not None))
AttributeError: 'Figure' object has no attribute 'flatten'
```

**Additional context**
Relevant package versions in the conda environment used:
```
arviz                     0.10.0                     py_0    conda-forge
bokeh                     2.2.1            py38h32f6830_0    conda-forge
matplotlib                3.1.3                    py38_0    conda-forge
numpy                     1.19.1           py38h3b9f5b6_0  
```

Discussion:
ahartikainen: I was trying to figure out, why there was a squeeze, but didn't have any good idea.

LGTM
ahartikainen: This will return 2D list of axes, I think that is fine.

Patch:
diff --git a/arviz/plots/backends/bokeh/densityplot.py b/arviz/plots/backends/bokeh/densityplot.py
index 7ede6ce..eb8dcec 100644
--- a/arviz/plots/backends/bokeh/densityplot.py
+++ b/arviz/plots/backends/bokeh/densityplot.py
@@ -63,7 +63,7 @@ def plot_density(
             rows,
             cols,
             figsize=figsize,
-            squeeze=True,
+            squeeze=False,
             backend_kwargs=backend_kwargs,
         )
     else:


Test patch:
diff --git a/arviz/tests/base_tests/test_plots_bokeh.py b/arviz/tests/base_tests/test_plots_bokeh.py
index 7b5bb21..4fbb31e 100644
--- a/arviz/tests/base_tests/test_plots_bokeh.py
+++ b/arviz/tests/base_tests/test_plots_bokeh.py
@@ -112,10 +112,26 @@ def test_plot_density_no_subset():
             "c": np.random.normal(size=200),
         }
     )
-    axes = plot_density([model_ab, model_bc])
+    axes = plot_density([model_ab, model_bc], backend="bokeh", show=False)
     assert axes.size == 3
 
 
+def test_plot_density_one_var():
+    """Test plot_density works when there is only one variable (#1401)."""
+    model_ab = from_dict(
+        {
+            "a": np.random.normal(size=200),
+        }
+    )
+    model_bc = from_dict(
+        {
+            "a": np.random.normal(size=200),
+        }
+    )
+    axes = plot_density([model_ab, model_bc], backend="bokeh", show=False)
+    assert axes.size == 1
+
+
 def test_plot_density_bad_kwargs(models):
     obj = [getattr(models, model_fit) for model_fit in ["model_1", "model_2"]]
     with pytest.raises(ValueError):