add compatible_with_equinox (#77)

* add `compatible_with_equinox`

* update readme

.gitignore

@@ -54,13 +54,6 @@ docs/apis/generated
 docs/apis/math/generated
 docs/apis/integrators/generated
 
-develop/benchmark/COBA/brian2*
-develop/benchmark/COBA/annarchy*
-develop/benchmark/COBAHH/brian2*
-develop/benchmark/COBAHH/annarchy*
-develop/benchmark/CUBA/annarchy*
-develop/benchmark/CUBA/brian2*
-
 
 *~
 \#*\#
@@ -227,3 +220,8 @@ cython_debug/
 /my_tests/
 /examples/dynamics_simulation/Joglekar_2018_data/
 /docs/apis/deprecated/generated/
+/docs/apis/brainunit.fft.rst
+/docs/apis/brainunit.lax.rst
+/docs/apis/brainunit.linalg.rst
+/docs/apis/brainunit.math.rst
+/docs/apis/changelog.md

README.md

@@ -35,7 +35,7 @@ pip install brainunit --upgrade
 The official documentation is hosted on Read the Docs: [https://brainunit.readthedocs.io](https://brainunit.readthedocs.io)
 
 
-## Unit-Aware Computation Ecosystem
+## Unit-aware computation ecosystem
 
 
 `brainunit` has been deeply integrated into following diverse projects, such as:
@@ -46,7 +46,10 @@ The official documentation is hosted on Read the Docs: [https://brainunit.readth
 - [``dendritex``](https://github.com/chaobrain/dendritex): Dendritic Modeling in JAX
 - [``pinnx``](https://github.com/chaobrain/pinnx): Physics-Informed Neural Networks for Scientific Machine Learning in JAX.
 
-- [``diffrax``](https://github.com/chaobrain/diffrax): Numerical differential equation solvers in JAX.
+
+Other unofficial projects include:
+
+- [``diffrax``](https://github.com/chaoming0625/diffrax): Numerical differential equation solvers in JAX.
 - [``jax-md``](https://github.com/Routhleck/jax-md): Differentiable Molecular Dynamics in JAX
 - [``Catalax``](https://github.com/Routhleck/Catalax): JAX-based framework to model biological systems
 - ...

brainunit/_base.py

@@ -70,6 +70,12 @@
 ]
 PyTree = Any
 _all_slice = slice(None, None, None)
+compat_with_equinox = False
+
+
+def compatible_with_equinox(mode: bool = True):
+    global compat_with_equinox
+    compat_with_equinox = mode
 
 
 def _to_quantity(array) -> 'Quantity':
@@ -809,8 +815,6 @@ def fail_for_dimension_mismatch(
         # ):
         #   return dim1, dim2
 
-        # We do another check here, this should allow Brian1 units to pass as
-        # having the same dimensions as a Brian2 unit
         if dim1 == dim2:
             return dim1, dim2
 
@@ -877,8 +881,6 @@ def fail_for_unit_mismatch(
     else:
         unit2 = get_unit(obj2)
 
-    # We do another check here, this should allow Brian1 units to pass as
-    # having the same dimensions as a Brian2 unit
     if unit1.has_same_dim(unit2):
         return unit1, unit2
 
@@ -1273,15 +1275,22 @@ class Unit:
     r"""
      A physical unit.
 
-     Normally, you do not need to worry about the implementation of
-     units. They are derived from the `Array` object with
-     some additional information (name and string representation).
-
      Basically, a unit is just a number with given dimensions, e.g.
      mvolt = 0.001 with the dimensions of voltage. The units module
      defines a large number of standard units, and you can also define
      your own (see below).
 
+     Mathematically, a unit represents:
+
+        .. math::
+
+            \text{{factor}} \times \text{{base}}^{\text{{scale}}} \times \text{{dimension}}
+
+     where the ``factor`` is the conversion factor of the unit (e.g. ``1 calorie = 4.18400 Joule``,
+     so the factor is 4.18400), the ``base`` is the base of the exponent (e.g. 10 for the kilo prefix),
+     the ``scale`` is the exponent of the base (e.g. 3 for the kilo prefix), and the ``dimension`` is
+     the physical dimensions of the unit (e.g. ``joule`` for energy).
+
      The unit class also keeps track of various things that were used
      to define it so as to generate a nice string representation of it.
      See below.
@@ -1337,7 +1346,7 @@ class Unit:
      >>> area = 20000 * U.um**2
 
      If you now ask for the conductance density, you will get an "ugly" display
-     in basic SI dimensions, as Brian does not know of a corresponding unit:
+     in basic SI dimensions, as  does not know of a corresponding unit:
 
      >>> conductance/area
      0.5 * metre ** -4 * kilogram ** -1 * second ** 3 * amp ** 2
@@ -1347,11 +1356,11 @@ class Unit:
 
      >>> U.usiemens/U.cm**2
      usiemens / (cmetre ** 2)
-     >>> conductance/area  # same as before, but now Brian knows about uS/cm^2
+     >>> conductance/area  # same as before, but now knows about uS/cm^2
      50. * usiemens / (cmetre ** 2)
 
      Note that user-defined units cannot override the standard units (`volt`,
-     `second`, etc.) that are predefined by Brian. For example, the unit
+     `second`, etc.) that are predefined. For example, the unit
      ``Nm`` has the dimensions "length²·mass/time²", and therefore the same
      dimensions as the standard unit `joule`. The latter will be used for display
      purposes:
@@ -2927,6 +2936,14 @@ def __imatmul__(self, oc):
     # -------------------- #
 
     def __pow__(self, oc):
+        if compat_with_equinox:
+            try:
+                from equinox.internal._omega import ω  # noqa
+                if isinstance(oc, ω):
+                    return ω(self)
+            except (ImportError, ModuleNotFoundError):
+                pass
+
         if isinstance(oc, Quantity):
             assert oc.is_unitless, f"Cannot calculate {self} ** {oc}, the exponent has to be dimensionless"
             oc = oc.mantissa

brainunit/_unit_constants.py

@@ -15,7 +15,7 @@
 
 r"""
 A module providing some physical units as `Quantity` objects. Note that these
-units are not imported by wildcard imports (e.g. `from brian2 import *`), they
+units are not imported by wildcard imports, they
 have to be imported explicitly. You can use ``import ... as ...`` to import them
 with shorter names, e.g.::
 
@@ -24,7 +24,7 @@
 The available constants are:
 
 ==================== ================== ======================= ==================================================================
-Constant             Symbol(s)          Brian name              Value
+Constant             Symbol(s)          name                    Value
 ==================== ================== ======================= ==================================================================
 Avogadro constant    :math:`N_A, L`     ``avogadro_constant``   :math:`6.022140857\times 10^{23}\,\mathrm{mol}^{-1}`
 Boltzmann constant   :math:`k`          ``boltzmann_constant``  :math:`1.38064852\times 10^{-23}\,\mathrm{J}\,\mathrm{K}^{-1}`