From 1110993e1b19cd69e6f6b5cc8ec5857c9676d8e6 Mon Sep 17 00:00:00 2001
From: Ryan Abernathey <ryan.abernathey@gmail.com>
Date: Mon, 28 Sep 2020 16:45:41 -0400
Subject: [PATCH 1/2] back to dask=parallelized

---
 pop_tools/eos.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/pop_tools/eos.py b/pop_tools/eos.py
index 80075973..7498d652 100644
--- a/pop_tools/eos.py
+++ b/pop_tools/eos.py
@@ -1,10 +1,10 @@
 import dask
 import numpy as np
 import xarray as xr
-from numba import jit
+from numba import vectorize
 
 
-@jit(nopython=True)
+@vectorize(['float64(float64)', 'float32(float32)'], nopython=True)
 def compute_pressure(depth):
     """
     Convert depth in meters to pressure in bars.
@@ -145,7 +145,10 @@ def eos(salt, temp, return_coefs=False, **kwargs):
         return RHO
 
 
-@jit(nopython=True)
+@vectorize(
+    ['float64(float64, float64, float64)', 'float32(float32, float32, float32)'],
+    nopython=True,
+)
 def _compute_eos(salt, temp, pressure):
     # MWJF EOS coefficients
     # *** these constants will be used to construct the numerator

From 43d71b4b1ed62c52a24918af606566bcfaf2ed70 Mon Sep 17 00:00:00 2001
From: Ryan Abernathey <ryan.abernathey@gmail.com>
Date: Mon, 28 Sep 2020 16:46:37 -0400
Subject: [PATCH 2/2] remove unneeded if clause

---
 pop_tools/eos.py | 20 ++++++++------------
 1 file changed, 8 insertions(+), 12 deletions(-)

diff --git a/pop_tools/eos.py b/pop_tools/eos.py
index 7498d652..a08a3c58 100644
--- a/pop_tools/eos.py
+++ b/pop_tools/eos.py
@@ -113,18 +113,14 @@ def eos(salt, temp, return_coefs=False, **kwargs):
             dRHOdT.attrs['long_name'] = 'Thermal expansion coefficient'
 
         else:
-            if isinstance(salt.data, dask.array.Array):
-                RHO = xr.apply_ufunc(
-                    _compute_eos,
-                    salt,
-                    temp,
-                    pressure,
-                    dask='parallelized',
-                    output_dtypes=[salt.dtype],
-                )
-            else:
-                RHO = xr.full_like(salt, fill_value=np.nan)
-                RHO[:] = _compute_eos(salt.data, temp.data, pressure.data)
+            RHO = xr.apply_ufunc(
+                _compute_eos,
+                salt,
+                temp,
+                pressure,
+                dask='parallelized',
+                output_dtypes=[salt.dtype],
+            )
 
         RHO.name = 'density'
         RHO.attrs['units'] = 'kg/m^3'