Update

pynapple/core/time_series.py

@@ -36,9 +36,7 @@
     _get_terminal_size,
     _split_tsd,
     _TsdFrameSliceHelper,
-    convert_to_jax_array,
-    convert_to_numpy_array,
-    get_backend,
+    convert_to_array,    
     is_array_like,
 )
 
@@ -73,11 +71,7 @@ class BaseTsd(Base, NDArrayOperatorsMixin, abc.ABC):
     def __init__(self, t, d, time_units="s", time_support=None):
         super().__init__(t, time_units, time_support)
 
-        # Check if jax backend
-        if get_backend() == "jax":
-            self.values = convert_to_jax_array(d, "d")
-        else:
-            self.values = convert_to_numpy_array(d, "d")
+        self.values = convert_to_array(d, "d")
 
         assert len(self.index) == len(
             self.values
@@ -455,7 +449,7 @@ def dropna(self, update_time_support=True):
         if hasattr(self, "columns"):
             kwargs["columns"] = self.columns
 
-        return self.__class__(t=t, d=d, time_support=ep)
+        return self.__class__(t=t, d=d, time_support=ep, **kwargs)
 
     def convolve(self, array, ep=None, trim="both"):
         """Return the discrete linear convolution of the time series with a one dimensional sequence.
@@ -472,7 +466,7 @@ def convolve(self, array, ep=None, trim="both"):
         ----------
         array : array-like
             One dimensional input array-like.
-            
+
         ep : None, optional
             The epochs to apply the convolution
         trim : str, optional

pynapple/core/utils.py

@@ -12,6 +12,16 @@
 from .config import nap_config
 
 
+def convert_to_array(array, array_name):
+    # Check if jax backend
+    if get_backend() == "jax":
+        from pynajax.utils import convert_to_jax_array
+
+        return convert_to_jax_array(array, array_name)
+    else:
+        return convert_to_numpy_array(array, array_name)
+
+
 def convert_to_numpy_array(array, array_name):
     """Convert any array like object to numpy ndarray.
 
@@ -48,46 +58,6 @@ def convert_to_numpy_array(array, array_name):
         )
 
 
-def convert_to_jax_array(array, array_name):
-    """Convert any array like object to jax Array.
-
-    Parameters
-    ----------
-    array : ArrayLike
-
-    array_name : str
-        Array name if RuntimeError is raised or object is casted to numpy
-
-    Returns
-    -------
-    jax.Array
-        Jax array object
-
-    Raises
-    ------
-    RuntimeError
-        If input can't be converted to jax array
-    """
-    import jax.numpy as jnp
-
-    if isinstance(array, Number):
-        return jnp.array([array])
-    elif isinstance(array, (list, tuple)):
-        return jnp.array(array)
-    elif isinstance(array, jnp.ndarray):
-        return array
-    elif isinstance(array, np.ndarray):
-        return cast_to_jax(array, array_name)
-    elif is_array_like(array):
-        return cast_to_jax(array, array_name)
-    else:
-        raise RuntimeError(
-            "Unknown format for {}. Accepted formats are numpy.ndarray, list, tuple or any array-like objects.".format(
-                array_name
-            )
-        )
-
-
 def get_backend():
     """
     Return the current backend of pynapple. Possible backends are

tests/test_correlograms.py

@@ -1,9 +1,3 @@
-# -*- coding: utf-8 -*-
-# @Author: gviejo
-# @Date:   2022-03-30 11:16:22
-# @Last Modified by:   Guillaume Viejo
-# @Last Modified time: 2024-04-04 10:46:31
-#!/usr/bin/env python
 
 """Tests of correlograms for `pynapple` package."""
 

tests/test_numpy_compatibility.py

@@ -1,11 +1,3 @@
-# -*- coding: utf-8 -*-
-# @Author: Guillaume Viejo
-# @Date:   2023-09-18 18:11:24
-# @Last Modified by:   Guillaume Viejo
-# @Last Modified time: 2024-04-04 11:34:35
-
-
-
 import pynapple as nap
 import numpy as np
 import pytest