add 'fft' and 'linalg' shortcut in brainunit.math (#76)

* add 'fft' and 'linalg' shortcut in `brainunit.math`

* update README

.gitignore

@@ -6,6 +6,7 @@ io_test_tmp*
 brainpy/math/brainpy_object/tests/io_test_tmp*
 
 lib/
+.VSCodeCounter
 
 development
 

README.md

@@ -3,23 +3,23 @@
 # Physical units and unit-aware mathematical system in JAX
 
 <p align="center">
-  	<img alt="Header image of brainunit." src="https://github.com/chaoming0625/brainunit/blob/main/docs/_static/brainunit.png" width=50%>
+  	<img alt="Header image of brainunit." src="https://github.com/chaobrain/brainunit/blob/main/docs/_static/brainunit.png" width=50%>
 </p> 
 
 
 
 <p align="center">
 	<a href="https://pypi.org/project/brainunit/"><img alt="Supported Python Version" src="https://img.shields.io/pypi/pyversions/brainunit"></a>
-	<a href="https://github.com/chaoming0625/brainunit/blob/main/LICENSE"><img alt="LICENSE" src="https://img.shields.io/badge/License-Apache%202.0-blue.svg"></a>
+	<a href="https://github.com/chaobrain/brainunit/blob/main/LICENSE"><img alt="LICENSE" src="https://img.shields.io/badge/License-Apache%202.0-blue.svg"></a>
     <a href='https://brainunit.readthedocs.io/en/latest/?badge=latest'>
         <img src='https://readthedocs.org/projects/brainunit/badge/?version=latest' alt='Documentation Status' />
     </a>  	
     <a href="https://badge.fury.io/py/brainunit"><img alt="PyPI version" src="https://badge.fury.io/py/brainunit.svg"></a>
-    <a href="https://github.com/chaoming0625/brainunit/actions/workflows/CI.yml"><img alt="Continuous Integration" src="https://github.com/chaoming0625/brainunit/actions/workflows/CI.yml/badge.svg"></a>
+    <a href="https://github.com/chaobrain/brainunit/actions/workflows/CI.yml"><img alt="Continuous Integration" src="https://github.com/chaobrain/brainunit/actions/workflows/CI.yml/badge.svg"></a>
 </p>
 
 
-[``brainunit``](https://github.com/chaoming0625/brainunit) provides physical units and unit-aware mathematical system in JAX for brain dynamics and AI4Science
+[``brainunit``](https://github.com/chaor/brainunit) provides physical units and unit-aware mathematical system in JAX for brain dynamics and AI4Science
 
 
 ## Installation
@@ -35,7 +35,25 @@ 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
+
+
+`brainunit` has been deeply integrated into following diverse projects, such as:
+
+- [``brainstate``](https://github.com/chaobrain/brainstate): A State-based Transformation System for Program Compilation and Augmentation
+- [``braintaichi``](https://github.com/chaobrain/braintaichi): Leveraging Taichi Lang to customize brain dynamics operators
+- [``braintools``](https://github.com/chaobrain/braintools): The Common Toolbox for Brain Dynamics Programming.
+- [``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.
+- [``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
+- ...
+
+
 
 ## See also the BDP ecosystem
 
-We are building the [brain dynamics programming ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/).
+We are building the [brain dynamics programming ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/). 
+[``brainunit``](https://github.com/chaobrain/brainunit) has been deeply integrated into our BDP ecosystem.

brainunit/_base.py

@@ -4476,7 +4476,8 @@ def new_f(*args, **kwds):
                         if isinstance(v, bool):
                             newkeyset[n] = v
                         else:
-                            raise TypeError(f"Function '{f.__name__}' expected a boolean value for argument '{n}' but got '{v}'")
+                            raise TypeError(
+                                f"Function '{f.__name__}' expected a boolean value for argument '{n}' but got '{v}'")
 
                     elif specific_unit == 1:
                         if isinstance(v, Quantity):
@@ -4521,14 +4522,15 @@ def new_f(*args, **kwds):
                     )
 
                 result = jax.tree.map(
-                    partial(_assign_unit, f), result, expected_result
+                    partial(_assign_unit, f), result, expected_result,
                 )
             return result
 
         return new_f
 
     return do_assign_units
 
+
 def _check_unit(f, val, unit):
     unit = UNITLESS if unit is None else unit
     if not has_same_unit(val, unit):
@@ -4540,8 +4542,9 @@ def _check_unit(f, val, unit):
         )
         raise UnitMismatchError(error_message, get_unit(val))
 
+
 def _assign_unit(f, val, unit):
-    if unit is None or unit == bool or unit == 1:
+    if unit is None or unit is bool:
         return val
     return Quantity(val, unit=unit)
 

brainunit/fft/_fft_change_unit.py

@@ -14,7 +14,6 @@
 # ==============================================================================
 from __future__ import annotations
 
-import sys
 from typing import Callable, Union, Sequence
 
 import jax
@@ -1046,6 +1045,7 @@ def irfftn(
     24: (u.Ysecond, u.yhertz),
 }
 
+
 def _find_closet_scale(scale):
     values = list(_time_freq_map.keys())
 
@@ -1113,7 +1113,6 @@ def fftfreq(
         return jnpfft.fftfreq(n, d, dtype=dtype)
 
 
-
 @set_module_as('brainunit.fft')
 def rfftfreq(
     n: int,

brainunit/fft/_fft_keep_unit_test.py

@@ -13,29 +13,19 @@
 # limitations under the License.
 # ==============================================================================
 
-import jax.numpy as jnp
-import pytest
-from absl.testing import parameterized
-
-import brainunit as bu
-import brainunit.math as bm
-from brainunit import second, meter, ms
-from brainunit._base import assert_quantity
-
 import jax.numpy as jnp
 import jax.numpy.fft as jnpfft
-import pytest
 from absl.testing import parameterized
 
-import brainunit as u
 import brainunit.fft as ufft
 from brainunit import meter, second
-from brainunit._base import assert_quantity, Unit, get_or_create_dimension
+from brainunit._base import assert_quantity
 
 fft_keep_unit = [
     'fftshift', 'ifftshift',
 ]
 
+
 class TestFftKeepUnit(parameterized.TestCase):
     def __init__(self, *args, **kwargs):
         super(TestFftKeepUnit, self).__init__(*args, **kwargs)
@@ -67,4 +57,4 @@ def test_fft_keep_unit(self, value_axes, unit):
             q = value * unit
             result = ufft_fun(q, axes=axes)
             expected = ufft_fun(jnp.array(value), axes=axes)
-            assert_quantity(result, expected, unit=unit)
+            assert_quantity(result, expected, unit=unit)

brainunit/math/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+from . import linalg, fft
 from ._activation import *
 from ._activation import __all__ as _activation_all
 from ._alias import *
@@ -32,22 +33,29 @@
 from ._misc import *
 from ._misc import __all__ as _compat_misc_all
 
-__all__ = (_compat_array_creation_all +
-           _alias_all +
-           _compat_funcs_change_unit_all +
-           _compat_funcs_keep_unit_all +
-           _compat_funcs_accept_unitless_all +
-           _compat_funcs_remove_unit_all +
-           _compat_misc_all +
-           _einops_all +
-           _activation_all)
+__all__ = (
+    _compat_array_creation_all +
+    _alias_all +
+    _compat_funcs_change_unit_all +
+    _compat_funcs_keep_unit_all +
+    _compat_funcs_accept_unitless_all +
+    _compat_funcs_remove_unit_all +
+    _compat_misc_all +
+    _einops_all +
+    _activation_all +
+    [
+        'linalg', 'fft'
+    ]
+)
 
-del (_compat_array_creation_all,
-     _alias_all,
-     _compat_funcs_change_unit_all,
-     _compat_funcs_keep_unit_all,
-     _compat_funcs_accept_unitless_all,
-     _compat_funcs_remove_unit_all,
-     _compat_misc_all,
-     _einops_all,
-     _activation_all)
+del (
+    _compat_array_creation_all,
+    _alias_all,
+    _compat_funcs_change_unit_all,
+    _compat_funcs_keep_unit_all,
+    _compat_funcs_accept_unitless_all,
+    _compat_funcs_remove_unit_all,
+    _compat_misc_all,
+    _einops_all,
+    _activation_all
+)

brainunit/math/fft.py

@@ -0,0 +1,18 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from brainunit.fft import *
+
+

brainunit/math/linalg.py

@@ -0,0 +1,18 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+from brainunit.linalg import *
+