remove the default unit register

brainunit/_base.py

@@ -51,20 +51,9 @@
 _all_slice = slice(None, None, None)
 random = None
 _unit_checking = True
-_automatically_register_units = True
 _allow_python_scalar_value = False
 
 
-@contextmanager
-def turn_off_unit_register():
-  try:
-    global _automatically_register_units
-    _automatically_register_units = False
-    yield
-  finally:
-    _automatically_register_units = True
-
-
 @contextmanager
 def allow_python_scalar():
   try:
@@ -2513,9 +2502,6 @@ def __init__(
 
     super().__init__(value, dtype=dtype, dim=dim)
 
-    if _automatically_register_units:
-      register_new_unit(self)
-
   @staticmethod
   def create(unit: Dimension, name: str, dispname: str, scale: int = 0):
     """

brainunit/_unit_common.py

@@ -20,7 +20,6 @@
   additional_unit_register,
   get_or_create_dimension,
   standard_unit_register,
-  turn_off_unit_register,
   allow_python_scalar
 )
 
@@ -2097,7 +2096,7 @@
   "celsius"  # Dummy object raising an error
 ]
 
-with turn_off_unit_register(), allow_python_scalar():
+with allow_python_scalar():
   #### FUNDAMENTAL UNITS
   metre = Unit.create(get_or_create_dimension(m=1), "metre", "m")
   meter = Unit.create(get_or_create_dimension(m=1), "meter", "m")

brainunit/_unit_constants.py

@@ -41,8 +41,6 @@
 
 import numpy as np
 
-from ._base import (turn_off_unit_register,
-                    allow_python_scalar)
 from ._unit_common import (
   amp,
   coulomb,
@@ -69,24 +67,23 @@
   'zero_celsius',
 ]
 
-with turn_off_unit_register(), allow_python_scalar():
-  #: Avogadro constant (http://physics.nist.gov/cgi-bin/cuu/Value?na)
-  avogadro_constant = 6.022140857e23 / mole
-  #: Boltzmann constant (physics.nist.gov/cgi-bin/cuu/Value?k)
-  boltzmann_constant = 1.38064852e-23 * joule / kelvin
-  #: electric constant (http://physics.nist.gov/cgi-bin/cuu/Value?ep0)
-  electric_constant = 8.854187817e-12 * farad / meter
-  #: Electron rest mass (physics.nist.gov/cgi-bin/cuu/Value?me)
-  electron_mass = 9.10938356e-31 * kilogram
-  #: Elementary charge (physics.nist.gov/cgi-bin/cuu/Value?e)
-  elementary_charge = 1.6021766208e-19 * coulomb
-  #: Faraday constant (http://physics.nist.gov/cgi-bin/cuu/Value?f)
-  faraday_constant = 96485.33289 * coulomb / mole
-  #: gas constant (http://physics.nist.gov/cgi-bin/cuu/Value?r)
-  gas_constant = 8.3144598 * joule / mole / kelvin
-  #: Magnetic constant (http://physics.nist.gov/cgi-bin/cuu/Value?mu0)
-  magnetic_constant = 4 * np.pi * 1e-7 * newton / amp ** 2
-  #: Molar mass constant (http://physics.nist.gov/cgi-bin/cuu/Value?mu)
-  molar_mass_constant = 1 * gram / mole
-  #: zero degree Celsius
-  zero_celsius = 273.15 * kelvin
+#: Avogadro constant (http://physics.nist.gov/cgi-bin/cuu/Value?na)
+avogadro_constant = 6.022140857e23 / mole
+#: Boltzmann constant (physics.nist.gov/cgi-bin/cuu/Value?k)
+boltzmann_constant = 1.38064852e-23 * joule / kelvin
+#: electric constant (http://physics.nist.gov/cgi-bin/cuu/Value?ep0)
+electric_constant = 8.854187817e-12 * farad / meter
+#: Electron rest mass (physics.nist.gov/cgi-bin/cuu/Value?me)
+electron_mass = 9.10938356e-31 * kilogram
+#: Elementary charge (physics.nist.gov/cgi-bin/cuu/Value?e)
+elementary_charge = 1.6021766208e-19 * coulomb
+#: Faraday constant (http://physics.nist.gov/cgi-bin/cuu/Value?f)
+faraday_constant = 96485.33289 * coulomb / mole
+#: gas constant (http://physics.nist.gov/cgi-bin/cuu/Value?r)
+gas_constant = 8.3144598 * joule / mole / kelvin
+#: Magnetic constant (http://physics.nist.gov/cgi-bin/cuu/Value?mu0)
+magnetic_constant = 4 * np.pi * 1e-7 * newton / amp ** 2
+#: Molar mass constant (http://physics.nist.gov/cgi-bin/cuu/Value?mu)
+molar_mass_constant = 1 * gram / mole
+#: zero degree Celsius
+zero_celsius = 273.15 * kelvin

brainunit/_unit_test.py

@@ -17,12 +17,14 @@
 import warnings
 
 import brainstate as bst
+import jax
 import jax.numpy as jnp
 import numpy as np
 import pytest
 from numpy.testing import assert_equal
 
-array = np.array
+import brainunit as bu
+
 bst.environ.set(precision=64)
 
 from brainunit._unit_common import *
@@ -1623,6 +1625,7 @@ def test_constants():
     (constants.avogadro_constant * constants.elementary_charge).value,
   )
 
+
 # if __name__ == "__main__":
 #     test_construction()
 #     test_get_dimensions()
@@ -1659,3 +1662,23 @@ def test_constants():
 #     test_units_vs_quantities()
 #     test_all_units_list()
 #     test_constants()
+
+
+def test_jit_array():
+  @jax.jit
+  def f1(a):
+    b = a * bu.siemens / bu.cm ** 2
+    return b
+
+  val = np.random.rand(3)
+  r = f1(val)
+  bu.math.allclose(val * bu.siemens / bu.cm ** 2, r)
+
+  @jax.jit
+  def f2(a):
+    a = a + 1. * bu.siemens / bu.cm ** 2
+    return a
+
+  val = np.random.rand(3) * bu.siemens / bu.cm ** 2
+  r = f2(val)
+  bu.math.allclose(val + 1 * bu.siemens / bu.cm ** 2, r)