Update code for removing of dimensions to a np.squeeze()

Author: jurjendejong

xarray_selafin/xarray_backend.py

@@ -244,16 +244,8 @@ def _raw_indexing_method(self, key):
                     values = temp[plan_indices][:, node_indices]
                     data[it] = np.reshape(values, (len(plan_indices), len(node_indices)))
 
-        # Remove dimension if key was an integer
-        if isinstance(node_key, int):
-            if plan_key is None:
-                data = data[:, 0]
-            else:
-                data = data[:, :, 0]
-        if isinstance(plan_key, int):
-            data = data[:, 0, :]
-        if isinstance(time_key, int):
-            data = data[0, ...]
+        # Remove dimension if the size was reduced to 1 due to indexing of integer
+        data = data.squeeze()
         return data