fix xarray 1D bug, add test, add example

LBL-EESA · Mar 27, 2024 · 81d0e35 · 81d0e35
1 parent 560df37
commit 81d0e35
Show file tree

Hide file tree

Showing 5 changed files with 246 additions and 114 deletions.
diff --git a/examples/Plotting.ipynb b/examples/Plotting.ipynb
diff --git a/examples/fastKDE and xarray.ipynb b/examples/fastKDE and xarray.ipynb
diff --git a/examples/readme_test.ipynb b/examples/readme_test.ipynb
diff --git a/src/fastkde/fastKDE.py b/src/fastkde/fastKDE.py
@@ -1217,7 +1217,10 @@ def pdf(*args, **kwargs):
             # set the variable name
             varname = var_names[i]
             # set the coordinate
-            coords[varname] = pdf_tuple[1][i]
+            if num_variables == 1:
+                coords[varname] = pdf_tuple[1]
+            else:
+                coords[varname] = pdf_tuple[1][i]
             # set the dimension
             dims.append(varname)
 

diff --git a/tests/test_execution.py b/tests/test_execution.py
@@ -34,6 +34,15 @@ def test_simple_3D():
     # the numPoints argument results in a coarser PDF--but one that is calculated
     # faster (and with less memory)
 
+def test_simple_1D_with_xarray():
+    np.random.seed(42)
+    N = int(2e5)
+    var1 = 50 * np.random.normal(size=N) + 0.1
+
+    # Do the self-consistent density estimate
+    myPDF = fastkde.pdf(var1, use_xarray=True)
+
+
 def test_simple_2D_with_xarray():
     np.random.seed(42)
     N = int(2e5)