refactor: Replace Array with BaseArray in multiple files for consistency

Author: chaoming0625

brainpy/_src/_delay.py

@@ -92,7 +92,7 @@ def __init__(
 
     # delay data
     if before_t0 is not None:
-      assert isinstance(before_t0, (int, float, bool, bm.Array, jax.Array, Callable))
+      assert isinstance(before_t0, (int, float, bool, bm.BaseArray, jax.Array, Callable))
     self._before_t0 = before_t0
     if length > 0:
       self._init_data(length)
@@ -139,7 +139,7 @@ def register_entry(
       delay_type = 'none'
     elif isinstance(delay_step, int):
       delay_type = 'homo'
-    elif isinstance(delay_step, (bm.Array, jax.Array, np.ndarray)):
+    elif isinstance(delay_step, (bm.BaseArray, jax.Array, np.ndarray)):
       if delay_step.size == 1 and delay_step.ndim == 0:
         delay_type = 'homo'
       else:
@@ -296,7 +296,7 @@ def _init_data(self, length, batch_size: int = None):
                             batch_axis=batch_axis)
     # update delay data
     self.data[0] = self.latest.value
-    if isinstance(self._before_t0, (bm.Array, jax.Array, float, int, bool)):
+    if isinstance(self._before_t0, (bm.BaseArray, jax.Array, float, int, bool)):
       self.data[1:] = self._before_t0
     elif callable(self._before_t0):
       self.data[1:] = self._before_t0((length,) + self.latest.shape, dtype=self.latest.dtype)

brainpy/_src/analysis/highdim/slow_points.py

@@ -340,13 +340,13 @@ def find_fps_with_gd_method(
     num_candidate = self._check_candidates(candidates)
     if not (isinstance(candidates, (bm.ndarray, jnp.ndarray, np.ndarray)) or isinstance(candidates, dict)):
       raise ValueError('Candidates must be instance of ArrayType or dict of ArrayType.')
-    fixed_points = tree_map(lambda a: bm.TrainVar(a), candidates, is_leaf=lambda x: isinstance(x, bm.Array))
+    fixed_points = tree_map(lambda a: bm.TrainVar(a), candidates, is_leaf=lambda x: isinstance(x, bm.BaseArray))
     f_eval_loss = self._get_f_eval_loss()
 
     def f_loss():
       return f_eval_loss(tree_map(lambda a: bm.as_jax(a),
                                   fixed_points,
-                                  is_leaf=lambda x: isinstance(x, bm.Array))).mean()
+                                  is_leaf=lambda x: isinstance(x, bm.BaseArray))).mean()
 
     grad_f = bm.grad(f_loss, grad_vars=fixed_points, return_value=True)
     optimizer.register_train_vars(fixed_points if isinstance(fixed_points, dict) else {'a': fixed_points})
@@ -390,10 +390,10 @@ def batch_train(start_i, n_batch):
     self._opt_losses = jnp.concatenate(opt_losses)
     self._losses = f_eval_loss(tree_map(lambda a: bm.as_jax(a),
                                         fixed_points,
-                                        is_leaf=lambda x: isinstance(x, bm.Array)))
+                                        is_leaf=lambda x: isinstance(x, bm.BaseArray)))
     self._fixed_points = tree_map(lambda a: bm.as_jax(a),
                                   fixed_points,
-                                  is_leaf=lambda x: isinstance(x, bm.Array))
+                                  is_leaf=lambda x: isinstance(x, bm.BaseArray))
     self._selected_ids = jnp.arange(num_candidate)
 
     if isinstance(self.target, DynamicalSystem):
@@ -429,7 +429,7 @@ def find_fps_with_opt_solver(
       print(f"Optimizing with {opt_solver} to find fixed points:")
 
     # optimizing
-    res = f_opt(tree_map(lambda a: bm.as_jax(a), candidates, is_leaf=lambda a: isinstance(a, bm.Array)))
+    res = f_opt(tree_map(lambda a: bm.as_jax(a), candidates, is_leaf=lambda a: isinstance(a, bm.BaseArray)))
 
     # results
     valid_ids = jnp.where(res.success)[0]
@@ -562,7 +562,7 @@ def compute_jacobians(
     """
     # check data
     info = np.asarray([(l.ndim, l.shape[0])
-                       for l in tree_flatten(points, is_leaf=lambda a: isinstance(a, bm.Array))[0]])
+                       for l in tree_flatten(points, is_leaf=lambda a: isinstance(a, bm.BaseArray))[0]])
     ndim = np.unique(info[:, 0])
     if len(ndim) != 1: raise ValueError(f'Get multiple dimension of the evaluated points. {ndim}')
     if ndim[0] == 1:

brainpy/_src/analysis/lowdim/lowdim_analyzer.py

@@ -359,9 +359,9 @@ def _get_fixed_points(self, candidates, *args, num_seg=None, tol_aux=1e-7, loss_
     """
     # candidates: xs, a vector with the length of self.resolutions[self.x_var]
     # args: parameters, a list/tuple of vectors
-    candidates = candidates.value if isinstance(candidates, bm.Array) else candidates
+    candidates = candidates.value if isinstance(candidates, bm.BaseArray) else candidates
     selected_ids = np.arange(len(candidates))
-    args = tuple(a.value if isinstance(candidates, bm.Array) else a for a in args)
+    args = tuple(a.value if isinstance(candidates, bm.BaseArray) else a for a in args)
     for a in args: assert len(a) == len(candidates)
     if num_seg is None:
       num_seg = len(self.resolutions[self.x_var])
@@ -557,7 +557,7 @@ def f_jacobian(*var_and_pars):
         return self.F_fx(*var_and_pars), self.F_fy(*var_and_pars)
 
       def call(*var_and_pars):
-        var_and_pars = tuple((vp.value if isinstance(vp, bm.Array) else vp) for vp in var_and_pars)
+        var_and_pars = tuple((vp.value if isinstance(vp, bm.BaseArray) else vp) for vp in var_and_pars)
         return jnp.array(jax.jit(f_jacobian, device=self.jit_device)(*var_and_pars))
 
       self.analyzed_results[C.F_jacobian] = call
@@ -879,7 +879,7 @@ def _get_fp_candidates_by_aux_rank(self, num_segments=1, num_rank=100):
 
       ps = tuple(p[ids[i]: ids[i] + arg_pre_len[i]] for i, p in enumerate(P))
       # change the position of meshgrid values
-      vps = tuple((v.value if isinstance(v, bm.Array) else v) for v in ((xs, ys) + ps))
+      vps = tuple((v.value if isinstance(v, bm.BaseArray) else v) for v in ((xs, ys) + ps))
       mesh_values = jnp.meshgrid(*vps)
       mesh_values = tuple(jnp.moveaxis(m, 0, 1) for m in mesh_values)
       mesh_values = tuple(m.flatten() for m in mesh_values)
@@ -934,9 +934,9 @@ def _get_fixed_points(self, candidates, *args, tol_aux=1e-7,
 
       # candidates: xs, a vector with the length of self.resolutions[self.x_var]
       # args: parameters, a list/tuple of vectors
-      candidates = candidates.value if isinstance(candidates, bm.Array) else candidates
+      candidates = candidates.value if isinstance(candidates, bm.BaseArray) else candidates
       selected_ids = np.arange(len(candidates))
-      args = tuple(a.value if isinstance(candidates, bm.Array) else a for a in args)
+      args = tuple(a.value if isinstance(candidates, bm.BaseArray) else a for a in args)
       for a in args: assert len(a) == len(candidates)
 
       if self.convert_type() == C.x_by_y:

brainpy/_src/analysis/utils/function.py

@@ -17,7 +17,7 @@
 def f_without_jaxarray_return(f):
   def f2(*args, **kwargs):
     r = f(*args, **kwargs)
-    return r.value if isinstance(r, bm.Array) else r
+    return r.value if isinstance(r, bm.BaseArray) else r
 
   return f2
 

brainpy/_src/analysis/utils/optimization.py

@@ -34,21 +34,21 @@
 
 
 def _logical_or(a, b):
-  a = a.value if isinstance(a, bm.Array) else a
-  b = b.value if isinstance(b, bm.Array) else b
+  a = a.value if isinstance(a, bm.BaseArray) else a
+  b = b.value if isinstance(b, bm.BaseArray) else b
   return jnp.logical_or(a, b)
 
 
 def _logical_and(a, b):
-  a = a.value if isinstance(a, bm.Array) else a
-  b = b.value if isinstance(b, bm.Array) else b
+  a = a.value if isinstance(a, bm.BaseArray) else a
+  b = b.value if isinstance(b, bm.BaseArray) else b
   return jnp.logical_and(a, b)
 
 
 def _where(p, a, b):
-  p = p.value if isinstance(p, bm.Array) else p
-  a = a.value if isinstance(a, bm.Array) else a
-  b = b.value if isinstance(b, bm.Array) else b
+  p = p.value if isinstance(p, bm.BaseArray) else p
+  a = a.value if isinstance(a, bm.BaseArray) else a
+  b = b.value if isinstance(b, bm.BaseArray) else b
   return jnp.where(p, a, b)
 
 
@@ -175,7 +175,7 @@ def get_brentq_candidates(f, xs, ys):
 
 def brentq_candidates(vmap_f, *values, args=()):
   # change the position of meshgrid values
-  values = tuple((v.value if isinstance(v, bm.Array) else v) for v in values)
+  values = tuple((v.value if isinstance(v, bm.BaseArray) else v) for v in values)
   xs = values[0]
   mesh_values = jnp.meshgrid(*values)
   if jnp.ndim(mesh_values[0]) > 1:
@@ -348,7 +348,7 @@ def scipy_minimize_with_jax(fun, x0,
   def fun_wrapper(x_flat, *args):
     x = unravel(x_flat)
     r = fun(x, *args)
-    r = r.value if isinstance(r, bm.Array) else r
+    r = r.value if isinstance(r, bm.BaseArray) else r
     return float(r)
 
   # Wrap the gradient in a similar manner
@@ -386,8 +386,8 @@ def roots_of_1d_by_x(f, candidates, args=()):
   """Find the roots of the given function by numerical methods.
   """
   f = f_without_jaxarray_return(f)
-  candidates = candidates.value if isinstance(candidates, bm.Array) else candidates
-  args = tuple(a.value if isinstance(candidates, bm.Array) else a for a in args)
+  candidates = candidates.value if isinstance(candidates, bm.BaseArray) else candidates
+  args = tuple(a.value if isinstance(candidates, bm.BaseArray) else a for a in args)
   vals = f(candidates, *args)
   signs = jnp.sign(vals)
   zero_sign_idx = jnp.where(signs == 0)[0]

brainpy/_src/analysis/utils/others.py

@@ -72,16 +72,16 @@ def check_plot_durations(plot_durations, duration, initials):
 
 def get_sign(f, xs, ys):
   f = f_without_jaxarray_return(f)
-  xs = xs.value if isinstance(xs, bm.Array) else xs
-  ys = ys.value if isinstance(ys, bm.Array) else ys
+  xs = xs.value if isinstance(xs, bm.BaseArray) else xs
+  ys = ys.value if isinstance(ys, bm.BaseArray) else ys
   Y, X = jnp.meshgrid(ys, xs)
   return jnp.sign(f(X, Y))
 
 
 def get_sign2(f, *xyz, args=()):
   in_axes = tuple(range(len(xyz))) + tuple([None] * len(args))
   f = jax.jit(jax.vmap(f_without_jaxarray_return(f), in_axes=in_axes))
-  xyz = tuple((v.value if isinstance(v, bm.Array) else v) for v in xyz)
+  xyz = tuple((v.value if isinstance(v, bm.BaseArray) else v) for v in xyz)
   XYZ = jnp.meshgrid(*xyz)
   XYZ = tuple(jnp.moveaxis(v, 1, 0).flatten() for v in XYZ)
   shape = (len(v) for v in xyz)
@@ -116,7 +116,7 @@ def keep_unique(candidates: Union[np.ndarray, Dict[str, np.ndarray]],
     return candidates, keep_ids
   if num_fps <= 1:
     return candidates, keep_ids
-  candidates = tree_map(lambda a: np.asarray(a), candidates, is_leaf=lambda a: isinstance(a, bm.Array))
+  candidates = tree_map(lambda a: np.asarray(a), candidates, is_leaf=lambda a: isinstance(a, bm.BaseArray))
 
   # If point A and point B are within identical_tol of each other, and the
   # A is first in the list, we keep A.

brainpy/_src/checkpoints/serialization.py

@@ -33,7 +33,7 @@
 except ModuleNotFoundError:
   msgpack = None
 
-from brainpy._src.math.ndarray import Array
+from brainpy._src.math.ndarray import BaseArray
 from brainpy.errors import (AlreadyExistsError,
                             MPACheckpointingRequiredError,
                             MPARestoreTargetRequiredError,
@@ -225,11 +225,11 @@ def _restore_list(xs, state_dict: Dict[str, Any]) -> List[Any]:
   return ys
 
 
-def _array_dict_state(x: Array) -> Dict[str, jax.Array]:
+def _array_dict_state(x: BaseArray) -> Dict[str, jax.Array]:
   return x.value
 
 
-def _restore_array(x, state_dict: jax.Array) -> Array:
+def _restore_array(x, state_dict: jax.Array) -> BaseArray:
   x.value = state_dict
   return x
 
@@ -276,7 +276,7 @@ def _restore_namedtuple(xs, state_dict: Dict[str, Any]):
   return type(xs)(**fields)
 
 
-register_serialization_state(Array, _array_dict_state, _restore_array)
+register_serialization_state(BaseArray, _array_dict_state, _restore_array)
 register_serialization_state(dict, _dict_state_dict, _restore_dict)
 # register_serialization_state(DotDict, _dict_state_dict, _restore_dict)
 # register_serialization_state(Collector, _dict_state_dict, _restore_dict)

brainpy/_src/connect/custom_conn.py

@@ -22,7 +22,7 @@ class MatConn(TwoEndConnector):
   def __init__(self, conn_mat, **kwargs):
     super(MatConn, self).__init__(**kwargs)
 
-    assert isinstance(conn_mat, (np.ndarray, bm.Array, jax.Array)) and conn_mat.ndim == 2
+    assert isinstance(conn_mat, (np.ndarray, bm.BaseArray, jax.Array)) and conn_mat.ndim == 2
     self.pre_num, self.post_num = conn_mat.shape
     self.pre_size, self.post_size = (self.pre_num,), (self.post_num,)
 
@@ -45,8 +45,8 @@ class IJConn(TwoEndConnector):
   def __init__(self, i, j, **kwargs):
     super(IJConn, self).__init__(**kwargs)
 
-    assert isinstance(i, (np.ndarray, bm.Array, jnp.ndarray)) and i.ndim == 1
-    assert isinstance(j, (np.ndarray, bm.Array, jnp.ndarray)) and j.ndim == 1
+    assert isinstance(i, (np.ndarray, bm.BaseArray, jnp.ndarray)) and i.ndim == 1
+    assert isinstance(j, (np.ndarray, bm.BaseArray, jnp.ndarray)) and j.ndim == 1
     assert i.size == j.size
 
     # initialize the class via "pre_ids" and "post_ids"

brainpy/_src/delay.py

@@ -102,7 +102,7 @@ def __init__(
 
     # delay data
     if init is not None:
-      assert isinstance(init, (numbers.Number, bm.Array, jax.Array, Callable))
+      assert isinstance(init, (numbers.Number, bm.BaseArray, jax.Array, Callable))
     self._init = init
 
     # other info
@@ -421,7 +421,7 @@ def _init_data(self, length: int, batch_size: int = None):
     else:
       self.data._value = data
     # update delay data
-    if isinstance(self._init, (bm.Array, jax.Array, numbers.Number)):
+    if isinstance(self._init, (bm.BaseArray, jax.Array, numbers.Number)):
       self.data[:] = self._init
     elif callable(self._init):
       self.data[:] = self._init((length,) + self.target.shape, dtype=self.target.dtype)
@@ -534,7 +534,7 @@ def init_delay_by_return(info: Union[bm.Variable, ReturnInfo], initial_delay_dat
     # init
     if isinstance(info.data, Callable):
       init = info.data(shape)
-    elif isinstance(info.data, (bm.Array, jax.Array)):
+    elif isinstance(info.data, (bm.BaseArray, jax.Array)):
       init = info.data
     else:
       raise TypeError

brainpy/_src/dnn/activations.py

@@ -14,7 +14,7 @@
 
 def _inplace(inp, val, inplace):
   if inplace:
-    assert isinstance(inp, bm.Array), 'input must be instance of brainpy.math.Array if inplace=True'
+    assert isinstance(inp, bm.BaseArray), 'input must be instance of brainpy.math.Array if inplace=True'
     inp.value = val
     return inp
   else: