reorganize operators in brainpy.math (#357)

reorganize operators in brainpy.math (#357)

Author: ztqakita

brainpy/_src/math/environment.py

@@ -324,66 +324,6 @@ def clone(self):
     return self.__class__()
 
 
-def set(
-    mode: modes.Mode = None,
-    dt: float = None,
-    x64: bool = None,
-    complex_: type = None,
-    float_: type = None,
-    int_: type = None,
-    bool_: type = None,
-):
-  """Set the default computation environment.
-
-  Parameters
-  ----------
-  mode: Mode
-    The computing mode.
-  dt: float
-    The numerical integration precision.
-  x64: bool
-    Enable x64 computation.
-  complex_: type
-    The complex data type.
-  float_
-    The floating data type.
-  int_
-    The integer data type.
-  bool_
-    The bool data type.
-  """
-  if dt is not None:
-    assert isinstance(dt, float), '"dt" must a float.'
-    set_dt(dt)
-
-  if mode is not None:
-    assert isinstance(mode, modes.Mode), f'"mode" must a {modes.Mode}.'
-    set_mode(mode)
-
-  if x64 is not None:
-    assert isinstance(x64, bool), f'"x64" must be a bool.'
-    set_x64(x64)
-
-  if float_ is not None:
-    assert isinstance(float_, type), '"float_" must a float.'
-    set_float(float_)
-
-  if int_ is not None:
-    assert isinstance(int_, type), '"int_" must a type.'
-    set_int(int_)
-
-  if bool_ is not None:
-    assert isinstance(bool_, type), '"bool_" must a type.'
-    set_bool(bool_)
-
-  if complex_ is not None:
-    assert isinstance(complex_, type), '"complex_" must a type.'
-    set_complex(complex_)
-
-
-set_environment = set
-
-
 class environment(_DecoratorContextManager):
   r"""Context-manager that sets a computing environment for brain dynamics computation.
 
@@ -541,6 +481,66 @@ def __init__(
                      mode=modes.BatchingMode(batch_size))
 
 
+def set(
+    mode: modes.Mode = None,
+    dt: float = None,
+    x64: bool = None,
+    complex_: type = None,
+    float_: type = None,
+    int_: type = None,
+    bool_: type = None,
+):
+  """Set the default computation environment.
+
+  Parameters
+  ----------
+  mode: Mode
+    The computing mode.
+  dt: float
+    The numerical integration precision.
+  x64: bool
+    Enable x64 computation.
+  complex_: type
+    The complex data type.
+  float_
+    The floating data type.
+  int_
+    The integer data type.
+  bool_
+    The bool data type.
+  """
+  if dt is not None:
+    assert isinstance(dt, float), '"dt" must a float.'
+    set_dt(dt)
+
+  if mode is not None:
+    assert isinstance(mode, modes.Mode), f'"mode" must a {modes.Mode}.'
+    set_mode(mode)
+
+  if x64 is not None:
+    assert isinstance(x64, bool), f'"x64" must be a bool.'
+    set_x64(x64)
+
+  if float_ is not None:
+    assert isinstance(float_, type), '"float_" must a float.'
+    set_float(float_)
+
+  if int_ is not None:
+    assert isinstance(int_, type), '"int_" must a type.'
+    set_int(int_)
+
+  if bool_ is not None:
+    assert isinstance(bool_, type), '"bool_" must a type.'
+    set_bool(bool_)
+
+  if complex_ is not None:
+    assert isinstance(complex_, type), '"complex_" must a type.'
+    set_complex(complex_)
+
+
+set_environment = set
+
+
 def enable_x64():
   config.update("jax_enable_x64", True)
   set_int(jnp.int64)

brainpy/_src/math/ndarray.py

@@ -90,7 +90,7 @@ class Array(object):
 
   is_brainpy_array = True
   _need_check_context = True
-  __slots__ = ("_value", "_transform_context")
+  __slots__ = ("_value", )
 
   def __init__(self, value, dtype=None):
     # array value
@@ -101,26 +101,6 @@ def __init__(self, value, dtype=None):
     if dtype is not None:
       value = jnp.asarray(value, dtype=dtype)
     self._value = value
-    # jit mode
-    self._transform_context = get_context()
-
-  def __check_context(self) -> None:
-    # raise error when in-place updating a
-    if self._need_check_context:
-      if self._transform_context is None:
-        if len(_jax_transformation_context_) > 0:
-          raise MathError(f'Array created outside of the transformation functions '
-                          f'({_jax_transformation_context_[-1]}) cannot be updated. '
-                          f'You should mark it as a brainpy.math.Variable instead.')
-      else:
-        if len(_jax_transformation_context_) > 0:
-          if self._transform_context != _jax_transformation_context_[-1]:
-            raise MathError(f'Array context "{self._transform_context}" differs from the JAX '
-                            f'transformation context "{_jax_transformation_context_[-1]}"'
-                            '\n\n'
-                            'Array created in one transformation function '
-                            'cannot be updated another transformation function. '
-                            'You should mark it as a brainpy.math.Variable instead.')
 
   @property
   def value(self):
@@ -1455,7 +1435,7 @@ class Variable(Array):
   """
 
   _need_check_context = False
-  __slots__ = ('_value', '_batch_axis')
+  __slots__ = ('_value', '_batch_axis', '_env')
 
   def __init__(
       self,
@@ -1487,6 +1467,8 @@ def __init__(
         raise MathError(f'This variables has {self.ndim} dimension, '
                         f'but the batch axis is set to be {batch_axis}.')
 
+    self._env = True
+
   @property
   def nobatch_shape(self) -> TupleType[int, ...]:
     """Shape without batch axis."""

brainpy/math/__init__.py

@@ -4,7 +4,7 @@
 # data structure
 from .ndarray import *
 from .delayvars import *
-from .arrayinterporate import *
+from .interoperability import *
 from .datatypes import *
 from .compat_numpy import *
 from .compat_tensorflow import *
@@ -15,7 +15,11 @@
 from . import activations
 
 # operators
-from .operators import *
+from .event_ops import *
+from .jitconn_ops import *
+from .pre_syn_post import *
+from .sparse_ops import *
+from .op_register import *
 from . import surrogate
 
 # Variable and Objects for object-oriented JAX transformations

brainpy/math/event_ops.py

@@ -0,0 +1,6 @@
+
+from brainpy._src.math.operators.event_ops import (
+  event_csr_matvec,
+  event_info
+)
+

brainpy/math/arrayinterporate.py renamed to brainpy/math/interoperability.py

@@ -0,0 +1,14 @@
+
+from brainpy._src.math.operators.jitconn_ops import (
+  event_matvec_prob_conn_homo_weight,
+  event_matvec_prob_conn_uniform_weight,
+  event_matvec_prob_conn_normal_weight,
+
+  matmat_prob_conn_homo_weight,
+  matmat_prob_conn_uniform_weight,
+  matmat_prob_conn_normal_weight,
+
+  matvec_prob_conn_homo_weight,
+  matvec_prob_conn_uniform_weight,
+  matvec_prob_conn_normal_weight,
+)

brainpy/math/jitconn_ops.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+
+
+from brainpy._src.math.operators.op_registers import (
+  XLACustomOp,
+  compile_cpu_signature_with_numba,
+)
+
+
+

brainpy/math/op_register.py

@@ -0,0 +1,20 @@
+from brainpy._src.math.operators.pre_syn_post import (
+  pre2post_sum,
+  pre2post_prod,
+  pre2post_max,
+  pre2post_min,
+  pre2post_mean,
+
+  pre2post_event_sum,
+  pre2post_coo_event_sum,
+  pre2post_event_prod,
+
+  pre2syn,
+
+  syn2post_sum, syn2post,
+  syn2post_prod,
+  syn2post_max,
+  syn2post_min,
+  syn2post_mean,
+  syn2post_softmax,
+)

brainpy/math/operators.py

@@ -0,0 +1,9 @@
+from brainpy._src.math.operators.sparse_ops import (
+  cusparse_csr_matvec,
+  cusparse_coo_matvec,
+  csr_matvec,
+  sparse_matmul,
+  coo_to_csr,
+  csr_to_coo,
+  csr_to_dense
+)