move variables into brainpy.math.object_transform

Author: chaoming0625

brainpy/__init__.py

@@ -132,10 +132,9 @@
 
 # deprecated
 from brainpy._src.math.object_transform.base import (Base as Base,
-                                                     DynVarCollector,
+                                                     ArrayCollector,
                                                      Collector as Collector, )
-globals()['ArrayCollector'] = DynVarCollector
-globals()['TensorCollector'] = DynVarCollector
+globals()['TensorCollector'] = ArrayCollector
 
 train.__dict__['DSTrainer'] = DSTrainer
 train.__dict__['BPTT'] = BPTT
@@ -150,8 +149,8 @@
 base.base.__dict__['BrainPyObject'] = BrainPyObject
 base.base.__dict__['Base'] = Base
 base.collector.__dict__['Collector'] = Collector
-base.collector.__dict__['ArrayCollector'] = DynVarCollector
-base.collector.__dict__['TensorCollector'] = DynVarCollector
+base.collector.__dict__['ArrayCollector'] = ArrayCollector
+base.collector.__dict__['TensorCollector'] = ArrayCollector
 base.function.__dict__['FunAsObject'] = math.FunAsObject
 base.function.__dict__['Function'] = math.FunAsObject
 base.io.__dict__['save_as_h5'] = checkpoints.io.save_as_h5
@@ -166,8 +165,8 @@
 base.__dict__['BrainPyObject'] = BrainPyObject
 base.__dict__['Base'] = Base
 base.__dict__['Collector'] = Collector
-base.__dict__['ArrayCollector'] = DynVarCollector
-base.__dict__['TensorCollector'] = DynVarCollector
+base.__dict__['ArrayCollector'] = ArrayCollector
+base.__dict__['TensorCollector'] = ArrayCollector
 base.__dict__['FunAsObject'] = math.FunAsObject
 base.__dict__['Function'] = math.FunAsObject
 base.__dict__['save_as_h5'] = checkpoints.io.save_as_h5

brainpy/_src/analysis/highdim/slow_points.py

@@ -132,11 +132,11 @@ def __init__(
 
     # update function
     if target_vars is None:
-      self.target_vars = bm.DynVarCollector()
+      self.target_vars = bm.ArrayCollector()
     else:
       if not isinstance(target_vars, dict):
         raise TypeError(f'"target_vars" must be a dict but we got {type(target_vars)}')
-      self.target_vars = bm.DynVarCollector(target_vars)
+      self.target_vars = bm.ArrayCollector(target_vars)
     excluded_vars = () if excluded_vars is None else excluded_vars
     if isinstance(excluded_vars, dict):
       excluded_vars = tuple(excluded_vars.values())

brainpy/_src/analysis/utils/model.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 
-from brainpy._src.math.ndarray import Variable
+from brainpy._src.math.object_transform import Variable
 from brainpy._src.math.environment import get_float
 from brainpy._src.math.interoperability import as_jax
 from brainpy._src.dyn.base import DynamicalSystem

brainpy/_src/checkpoints/io.py

@@ -8,7 +8,7 @@
 
 from brainpy import errors
 import brainpy.math as bm
-from brainpy._src.math.object_transform.base import BrainPyObject, DynVarCollector
+from brainpy._src.math.object_transform.base import BrainPyObject, ArrayCollector
 
 
 logger = logging.getLogger('brainpy.brainpy_object.io')
@@ -120,7 +120,7 @@ def _load(
 
 
 def _unique_and_duplicate(collector: dict):
-  gather = DynVarCollector()
+  gather = ArrayCollector()
   id2name = dict()
   duplicates = ([], [])
   for k, v in collector.items():

brainpy/_src/dyn/base.py

@@ -14,7 +14,7 @@
 from brainpy._src.connect import TwoEndConnector, MatConn, IJConn, One2One, All2All
 from brainpy._src.initialize import Initializer, parameter, variable, Uniform, noise as init_noise
 from brainpy._src.integrators import odeint, sdeint
-from brainpy._src.math.ndarray import Variable, VariableView
+from brainpy._src.math.object_transform.variables import Variable, VariableView
 from brainpy._src.math.object_transform.base import BrainPyObject, Collector
 from brainpy.errors import NoImplementationError, UnsupportedError
 from brainpy.types import ArrayType, Shape

brainpy/_src/integrators/fde/Caputo.py

@@ -329,16 +329,16 @@ def __init__(
 
     # initial values
     inits = check_inits(inits, self.variables)
-    self.inits = bm.VarDict({v: bm.asarray(inits[v]) for v in self.variables})
+    self.inits = bm.VarDict({v: bm.Variable(inits[v]) for v in self.variables})
 
     # coefficients
     ranges = bm.asarray([bm.arange(1, num_memory + 2) for _ in self.variables]).T
     coef = bm.diff(bm.power(ranges, 1 - self.alpha), axis=0)
     self.coef = bm.flip(coef, axis=0)
 
     # used to save the difference of two adjacent states
-    self.diff_states = bm.VarDict({v + "_diff": bm.zeros((num_memory,) + self.inits[v].shape,
-                                                         dtype=self.inits[v].dtype)
+    self.diff_states = bm.VarDict({v + "_diff": bm.Variable(bm.zeros((num_memory,) + self.inits[v].shape,
+                                                            dtype=self.inits[v].dtype))
                                    for v in self.variables})
     self.idx = bm.Variable(bm.asarray([self.num_memory - 1]))
 

brainpy/_src/math/delayvars.py

@@ -13,9 +13,10 @@
 from brainpy.errors import UnsupportedError
 from .interoperability import as_jax
 from .compat_numpy import vstack, broadcast_to
-from .environment import get_dt, get_float, get_int
-from .ndarray import ndarray, Variable, Array
+from .environment import get_dt, get_float
+from .ndarray import ndarray, Array
 from .object_transform.base import BrainPyObject
+from .object_transform.variables import Variable
 
 __all__ = [
   'AbstractDelay',

brainpy/_src/math/interoperability.py

@@ -3,7 +3,8 @@
 import jax.numpy as jnp
 import numpy as np
 
-from .ndarray import Array, Variable
+from .ndarray import Array
+
 
 __all__ = [
   'as_device_array', 'as_jax', 'as_ndarray', 'as_numpy', 'as_variable',
@@ -90,4 +91,5 @@ def as_variable(tensor, dtype=None):
     Array interpretation of `tensor`.  No copy is performed if the input
     is already an ndarray with matching dtype.
   """
+  from .object_transform.variables import Variable
   return Variable(tensor, dtype=dtype)