saiunit 0.3.2__tar.gz → 0.4.0__tar.gz

{saiunit-0.3.2 → saiunit-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: saiunit
-Version: 0.3.2
+Version: 0.4.0
 Summary: Enabling Unit-aware Computations for AI-driven Scientific Computing.
 Author-email: SAIUnit Developers <chao.brain@qq.com>
 License-Expression: Apache-2.0

{saiunit-0.3.2 → saiunit-0.4.0}/pyproject.toml

@@ -91,6 +91,8 @@ all = ["saiunit[jax,cupy,torch,dask,ndonnx]"]
 # core requirements (numpy + jax) installed.
 [[tool.mypy.overrides]]
 module = [
+    "array_api_compat",
+    "array_api_compat.*",
     "brainstate",
     "brainstate.*",
     "cupy",

{saiunit-0.3.2 → saiunit-0.4.0}/saiunit/__init__.py

@@ -131,13 +131,12 @@ __all__ = [
               '__version__',
               '__version_info__',
 
-              # submodules
+              # submodules ('autograd', 'lax' and 'sparse' are appended below
+              # together with the lazy-loading machinery)
               'math',
               'linalg',
-              'autograd',
               'fft',
               'constants',
-              'sparse',
               'typing',
 
               # misc

{saiunit-0.3.2 → saiunit-0.4.0}/saiunit/_base_decorators.py

@@ -177,6 +177,17 @@ def check_dims(**au):
                             )
                 newkeyset[n] = v
 
+            # validate default values of parameters the caller did not
+            # supply, so a default that violates the specification does not
+            # silently bypass the check
+            defaults = f.__defaults__ or ()
+            for n, v in zip(arg_names[len(arg_names) - len(defaults):], defaults):
+                if n not in newkeyset and n in au:
+                    newkeyset[n] = v
+            for n, v in (f.__kwdefaults__ or {}).items():
+                if n not in newkeyset and n in au:
+                    newkeyset[n] = v
+
             for k in newkeyset:
                 # string variables are allowed to pass, the presumption is they
                 # name another variable. None is also allowed, useful for
@@ -307,7 +318,13 @@ def _check_dim(f, val, dim):
 @set_module_as('saiunit')
 def check_units(**au):
     """
-    Decorator to check units of arguments passed to a function
+    Decorator to check units of arguments passed to a function.
+
+    The check is **exact**: the argument must carry the very unit given in
+    the specification, including its scale. ``1 * mV`` does *not* pass a
+    ``volt`` specification even though both are voltages. Use
+    :func:`check_dims` if you only want to constrain the physical dimension
+    and accept any scaled unit of that dimension.
 
     Examples
     --------
@@ -319,12 +336,12 @@ def check_units(**au):
     You don't have to check the units of every variable in the function, and
     you can define what the units should be for variables that aren't
     explicitly named in the definition of the function. For example, the code
-    above checks that the variable wibble should be a length, so writing
+    above checks that the variable wibble should be in metre, so writing
 
     >>> getvoltage(1*amp, 1*ohm, wibble=1)  # doctest: +IGNORE_EXCEPTION_DETAIL
     Traceback (most recent call last):
     ...
-    DimensionMismatchError: Function "getvoltage" variable "wibble" has wrong dimensions, dimensions were (1) (m)
+    UnitMismatchError: Function 'getvoltage' expected a array with unit Unit("m") for argument 'wibble' but got '1'
 
     fails, but
 
@@ -360,7 +377,7 @@ def check_units(**au):
     a function that takes the units of all the arguments as its inputs (in the
     order specified in the function header):
 
-    >>> @check_units(result=lambda d: d**2)
+    >>> @check_units(result=lambda u: u**2)
     ... def square(value):
     ...     return value**2
 
@@ -379,14 +396,17 @@ def check_units(**au):
     >>> multiply_sum(3*nA, 4*mV, 5*nA)  # doctest: +IGNORE_EXCEPTION_DETAIL
     Traceback (most recent call last):
     ...
-    DimensionMismatchError: Function 'multiply_sum' expected the same arguments for arguments 'summand_1', 'summand_2', but argument 'summand_1' has unit V, while argument 'summand_2' has unit A.
+    UnitMismatchError: Function 'multiply_sum' expected the argument 'summand_2' to have the same units as argument 'summand_1', but argument 'summand_2' has unit nA, while argument 'summand_1' has unit mV.
+
+    Default values of parameters that the caller does not supply are
+    validated as well.
 
     Raises
     ------
 
-    DimensionMismatchError
+    UnitMismatchError
         In case the input arguments or the return value do not have the
-        expected dimensions.
+        expected units.
     TypeError
         If an input argument or return value was expected to be a boolean but
         is not.
@@ -432,6 +452,17 @@ def check_units(**au):
                             )
                 newkeyset[n] = v
 
+            # validate default values of parameters the caller did not
+            # supply, so a default that violates the specification does not
+            # silently bypass the check
+            defaults = f.__defaults__ or ()
+            for n, v in zip(arg_names[len(arg_names) - len(defaults):], defaults):
+                if n not in newkeyset and n in au:
+                    newkeyset[n] = v
+            for n, v in (f.__kwdefaults__ or {}).items():
+                if n not in newkeyset and n in au:
+                    newkeyset[n] = v
+
             for k in newkeyset:
                 # string variables are allowed to pass, the presumption is they
                 # name another variable. None is also allowed, useful for
@@ -488,7 +519,10 @@ def check_units(**au):
             result = f(*args, **kwds)
             if "result" in au:
                 if isinstance(au["result"], Callable) and au["result"] != bool:
-                    expected_result = au["result"](*[get_dim(a) for a in args])
+                    # Pass the argument *units* (not dimensions): the result is
+                    # validated by unit with ``_check_unit`` below, so a Dimension
+                    # here would always mismatch the returned Quantity's unit.
+                    expected_result = au["result"](*[get_unit(a) for a in args])
                 else:
                     expected_result = au["result"]
 

{saiunit-0.3.2 → saiunit-0.4.0}/saiunit/_base_decorators_test.py

@@ -266,6 +266,27 @@ class TestCheckUnits:
         with pytest.raises(UnitMismatchError):
             f()
 
+    def test_result_callable(self):
+        # Documented usage: the result unit is derived from the input unit.
+        # The callable must receive the argument *unit* (not its dimension),
+        # otherwise the unit comparison spuriously fails.
+        @check_units(result=lambda un: un ** 2)
+        def square(x):
+            return x ** 2
+
+        q = Quantity(3.0, unit=_metre)
+        result = square(q)
+        assert result.unit == _metre ** 2
+        assert jnp.allclose(result.mantissa, 9.0)
+
+    def test_result_callable_mismatch_raises(self):
+        @check_units(result=lambda un: un ** 2)
+        def bad(x):
+            return x  # wrong: returns metre, not metre**2
+
+        with pytest.raises(UnitMismatchError):
+            bad(Quantity(3.0, unit=_metre))
+
     def test_stores_metadata(self):
         @check_units(x=_metre, result=_metre)
         def f(x):
@@ -912,3 +933,53 @@ class TestDocstringExampleAssignUnits:
         assert v.unit == u.volt
         assert jnp.allclose(t.mantissa, 5.0)
         assert jnp.allclose(v.mantissa, 3.0)
+
+
+class TestDefaultValueValidation:
+    """Default values must be validated like explicitly passed arguments."""
+
+    def test_check_units_bad_default_rejected(self):
+        @check_units(v=_volt)
+        def f(v=1 * _msecond):
+            return v
+
+        with pytest.raises(UnitMismatchError):
+            f()
+
+    def test_check_units_good_default_passes(self):
+        @check_units(v=_volt)
+        def f(v=1 * _volt):
+            return v
+
+        assert f() == 1 * _volt
+
+    def test_check_units_none_default_passes(self):
+        @check_units(v=_volt)
+        def f(v=None):
+            return v
+
+        assert f() is None
+
+    def test_check_units_kwonly_default_rejected(self):
+        @check_units(w=_volt)
+        def f(*, w=2 * _second):
+            return w
+
+        with pytest.raises(UnitMismatchError):
+            f()
+        assert f(w=1 * _volt) == 1 * _volt
+
+    def test_check_dims_bad_default_rejected(self):
+        @check_dims(v=_voltage_dim)
+        def f(v=1 * _second):
+            return v
+
+        with pytest.raises(DimensionMismatchError):
+            f()
+
+    def test_check_dims_scaled_default_passes(self):
+        @check_dims(v=_time_dim)
+        def f(v=1 * _msecond):
+            return v
+
+        assert f() == 1 * _msecond

{saiunit-0.3.2 → saiunit-0.4.0}/saiunit/_base_dimension.py

@@ -162,7 +162,7 @@ class Dimension:
 
         This property memoizes the hash value for efficiency. Once calculated,
         the hash value is stored in the `_hash` attribute for future access.
-        The hash is based on the binary representation of the dimensions array.
+        The hash is based on the values of the dimension exponents.
 
         Returns
         -------
@@ -176,7 +176,12 @@ class Dimension:
         their use as dictionary keys and in sets.
         """
         if self._hash is None:
-            self._hash = hash(self._dims.tobytes())
+            # Hash by *value*, not by byte representation: ``tobytes()``
+            # differs between dtypes (int vs float dims) and between 0.0
+            # and -0.0, while ``__eq__`` (np.array_equal) treats those as
+            # equal.  Python scalar hashing guarantees hash(1) == hash(1.0)
+            # and hash(0.0) == hash(-0.0), keeping the hash/eq contract.
+            self._hash = hash(tuple(self._dims.tolist()))
         return self._hash
 
     @hash.setter
@@ -463,6 +468,14 @@ class Dimension:
         value = np.array(value)
         if value.size > 1:
             raise TypeError("Too many exponents")
+        if value.size == 0:
+            raise TypeError("Expected a scalar exponent, got an empty array")
+        if not np.all(np.isfinite(value)):
+            # A NaN/inf exponent would create NaN dimension entries; NaN
+            # compares unequal to itself, so every such Dimension would
+            # miss the singleton cache, growing it unboundedly and
+            # breaking the ``is``-identity the library relies on.
+            raise ValueError(f"Dimension exponent must be finite, got {value!r}")
         return get_or_create_dimension(self._dims * value)
 
     def __imul__(self, value):  # type: ignore[misc]
@@ -563,17 +576,19 @@ class Dimension:
         -------
         bool
             True if the dimensions are equal (have the same exponents),
-            False otherwise or if the provided value is not a Dimension object.
+            False otherwise.
 
         Notes
         -----
         The comparison uses ``numpy.array_equal``. Equal Dimensions are
         produced by ``get_or_create_dimension`` from the same exponent tuple,
         so callers that round their exponents identically observe equal
-        instances. If value is not a Dimension object, returns False.
+        instances. If value is not a Dimension object, ``NotImplemented``
+        is returned so that the other operand's reflected comparison is
+        attempted before Python falls back to ``False``.
         """
         if not isinstance(value, Dimension):
-            return False
+            return NotImplemented
         try:
             return np.array_equal(self._dims, value._dims)
         except (AttributeError, _TracerArrayConversionError):
@@ -597,7 +612,10 @@ class Dimension:
         bool
             True if the dimensions are not equal, False otherwise.
         """
-        return not self.__eq__(value)
+        result = self.__eq__(value)
+        if result is NotImplemented:
+            return result
+        return not result
 
     # MAKE DIMENSION PICKABLE #
     def __getstate__(self):
@@ -750,6 +768,11 @@ def get_or_create_dimension(*args, **kwds) -> Dimension:
     if len(args):
         if len(args) != 1:
             raise TypeError(f"get_or_create_dimension() takes at most 1 positional argument, got {len(args)}")
+        if kwds:
+            raise TypeError(
+                "get_or_create_dimension() accepts either a positional sequence "
+                "or keyword arguments, not both"
+            )
         # initialisation by list
         dims = args[0]
         try:
@@ -759,12 +782,25 @@ def get_or_create_dimension(*args, **kwds) -> Dimension:
             raise TypeError("Need a sequence of exactly 7 items")
     else:
         # initialisation by keywords
-        dims = np.asarray([0, 0, 0, 0, 0, 0, 0])
+        # Use a plain list (not an int array) so fractional exponents such
+        # as metre=0.5 are preserved; np.asarray below picks the dtype.
+        dims = [0, 0, 0, 0, 0, 0, 0]
         for k in kwds:
             # _dim2index stores the index of the dimension with name 'k'
-            dims[_dim2index[k]] = kwds[k]
+            try:
+                index = _dim2index[k]
+            except KeyError:
+                raise TypeError(
+                    f"Unknown dimension name {k!r}. Valid names are: "
+                    f"{sorted(_dim2index)}"
+                ) from None
+            dims[index] = kwds[k]
 
     dims = np.asarray(dims)
+    if not np.issubdtype(dims.dtype, np.number):
+        raise TypeError(
+            f"Dimension exponents must be numeric, got dtype {dims.dtype!r}"
+        )
     key = tuple(dims)
     cached = _dimension_cache.get(key)
     if cached is not None:

{saiunit-0.3.2 → saiunit-0.4.0}/saiunit/_base_dimension_test.py

@@ -280,6 +280,14 @@ class TestGetOrCreateDimension:
         assert d.get_dimension("kg") == 1
         assert d.get_dimension("s") == -2
 
+    def test_fractional_exponent_by_keyword(self):
+        # Fractional exponents passed by keyword must be preserved, not
+        # truncated to int. The list-form and ** paths already do this.
+        d = get_or_create_dimension(metre=0.5)
+        assert not d.is_dimensionless
+        assert d.get_dimension("m") == 0.5
+        assert d == get_or_create_dimension([0.5, 0, 0, 0, 0, 0, 0])
+
     def test_singleton_by_list(self):
         d1 = get_or_create_dimension([1, 0, 0, 0, 0, 0, 0])
         d2 = get_or_create_dimension([1, 0, 0, 0, 0, 0, 0])
@@ -496,3 +504,77 @@ def test_docstring_example_unit_mismatch_error():
     import saiunit as u
     e = u.UnitMismatchError("Addition", u.mvolt, u.volt)
     assert 'Addition' in str(e)
+
+
+# =========================================================================
+# Regression tests: audited Dimension API fixes (2026-06)
+# =========================================================================
+
+class TestDimensionHashEqContract:
+    def test_int_vs_float_dims_hash_consistent(self):
+        # Direct construction is public API; equal Dimensions must hash
+        # equal regardless of the dtype of the exponent array.
+        d_int = Dimension(np.array([1, 0, 0, 0, 0, 0, 0]))
+        d_flt = Dimension(np.array([1., 0., 0., 0., 0., 0., 0.]))
+        assert d_int == d_flt
+        assert hash(d_int) == hash(d_flt)
+
+    def test_negative_zero_hash_consistent(self):
+        d_pos = Dimension(np.array([1., 0., 0., 0., 0., 0., 0.]))
+        d_neg = Dimension(np.array([1., -0., 0., 0., 0., 0., 0.]))
+        assert d_pos == d_neg
+        assert hash(d_pos) == hash(d_neg)
+
+    def test_eq_non_dimension_returns_notimplemented(self):
+        d = get_or_create_dimension([1, 0, 0, 0, 0, 0, 0])
+        assert d.__eq__(5) is NotImplemented
+        assert (d == 5) is False
+        assert (d != 5) is True
+
+
+class TestDimensionPowGuards:
+    def test_pow_nan_raises(self):
+        d = get_or_create_dimension([1, 0, 0, 0, 0, 0, 0])
+        with pytest.raises(ValueError, match="finite"):
+            d ** float('nan')
+
+    def test_pow_inf_raises(self):
+        d = get_or_create_dimension([1, 0, 0, 0, 0, 0, 0])
+        with pytest.raises(ValueError, match="finite"):
+            d ** float('inf')
+
+    def test_pow_nan_does_not_grow_cache(self):
+        # A NaN exponent must not leak NaN-keyed entries into the
+        # singleton cache (NaN != NaN would make every lookup miss).
+        from saiunit._base_dimension import _dimension_cache
+        d = get_or_create_dimension([1, 0, 0, 0, 0, 0, 0])
+        n0 = len(_dimension_cache)
+        for _ in range(3):
+            with pytest.raises(ValueError):
+                d ** float('nan')
+        assert len(_dimension_cache) == n0
+
+    def test_pow_empty_array_raises(self):
+        d = get_or_create_dimension([1, 0, 0, 0, 0, 0, 0])
+        with pytest.raises(TypeError, match="empty"):
+            d ** np.array([])
+
+
+class TestFactoryValidation:
+    def test_positional_and_kwargs_raises(self):
+        with pytest.raises(TypeError, match="not both"):
+            get_or_create_dimension([1, 0, 0, 0, 0, 0, 0], length=5)
+
+    def test_unknown_kwarg_raises_typeerror(self):
+        with pytest.raises(TypeError, match="Unknown dimension name"):
+            get_or_create_dimension(bad_kw=1)
+
+    def test_non_numeric_string_raises(self):
+        # A 7-character string passes the len() check but is not a
+        # sequence of exponents.
+        with pytest.raises(TypeError, match="numeric"):
+            get_or_create_dimension("abcdefg")
+
+    def test_non_numeric_sequence_raises(self):
+        with pytest.raises(TypeError, match="numeric"):
+            get_or_create_dimension(['a'] * 7)

{saiunit-0.3.2 → saiunit-0.4.0}/saiunit/_base_getters.py

@@ -162,6 +162,10 @@ def get_dim(obj) -> Dimension:
         return obj
     if isinstance(obj, Quantity):
         return obj.dim
+    if isinstance(obj, (str, bytes)):
+        # Quantity() rejects str/bytes mantissas, but get_dim keeps its
+        # duck-typing contract: anything without units is dimensionless.
+        return DIMENSIONLESS
     try:
         return Quantity(obj).dim
     except TypeError:
@@ -207,6 +211,10 @@ def get_unit(obj) -> 'Unit':
         return obj
     if isinstance(obj, Quantity):
         return obj.unit
+    if isinstance(obj, (str, bytes)):
+        # Quantity() rejects str/bytes mantissas, but get_unit keeps its
+        # duck-typing contract: anything without units is unitless.
+        return UNITLESS
     try:
         return Quantity(obj).unit
     except TypeError:
@@ -731,12 +739,16 @@ def unit_scale_align_to_first(*args) -> 'list[Quantity]':
         if not isinstance(items[0], Quantity):
             items[0] = Quantity(items[0])
         for i in range(1, len(items)):
-            fail_for_unit_mismatch(items[i], items[0], 'Non-matching unit for function "unit_scale_align_to_first"')
-            if not isinstance(items[i], Quantity):
-                items[i] = Quantity(items[i])
+            # ``in_unit`` both checks the dimension and rescales
+            # scaled-dimensionless quantities (e.g. mV/volt) onto the
+            # unitless first item, instead of only checking dimensions.
+            items[i] = _to_quantity(items[i]).in_unit(first_unit)
     else:
         for i in range(1, len(items)):
-            items[i] = items[i].in_unit(first_unit)
+            # Route through _to_quantity so raw values (plain numbers, bare
+            # arrays, Unit objects) raise a proper UnitMismatchError instead
+            # of an AttributeError on ``.in_unit``.
+            items[i] = _to_quantity(items[i]).in_unit(first_unit)
     return items