diff --git a/brainunit/math/_compat_numpy_array_creation.py b/brainunit/math/_compat_numpy_array_creation.py index 32219a8..68873dc 100644 --- a/brainunit/math/_compat_numpy_array_creation.py +++ b/brainunit/math/_compat_numpy_array_creation.py @@ -498,6 +498,7 @@ def asarray( Convert the input to a quantity or array. If unit is provided, the input will be checked whether it has the same unit as the provided unit. + (If they have same dimension but different magnitude, the input will be converted to the provided unit.) If unit is not provided, the input will be converted to an array. Args: @@ -513,9 +514,13 @@ def asarray( if unit is not None: assert isinstance(unit, Unit) fail_for_dimension_mismatch(a, unit, error_message="a and unit have to have the same units.") - return Quantity(jnp.asarray(a.value, dtype=dtype, order=order), dim=a.dim) + if a.dim == unit: + return a + else: + # Convert to the magnitude of the provided unit + return Quantity(a.value / unit.value, dim=unit) else: - return Quantity(jnp.asarray(a.value, dtype=dtype, order=order), dim=a.dim) + return Quantity(jnp.asarray(a.value, dtype=dtype, order=order) / unit.value, dim=a.dim) elif isinstance(a, (jax.Array, np.ndarray)): if unit is not None: assert isinstance(unit, Unit) diff --git a/brainunit/math/_compat_numpy_test.py b/brainunit/math/_compat_numpy_test.py index 2e73403..0bd8c10 100644 --- a/brainunit/math/_compat_numpy_test.py +++ b/brainunit/math/_compat_numpy_test.py @@ -32,7 +32,7 @@ def assert_quantity(q, values, unit): values = jnp.asarray(values) if isinstance(q, Quantity): - assert q.dim == unit.dim, f"Unit mismatch: {q.dim} != {unit}" + assert q.dim == unit.dim or q.dim == unit, f"Unit mismatch: {q.dim} != {unit}" assert jnp.allclose(q.value, values), f"Values do not match: {q.value} != {values}" else: assert jnp.allclose(q, values), f"Values do not match: {q} != {values}" @@ -162,6 +162,10 @@ def test_asarray(self): result_q = bu.math.asarray([1, 2, 3], unit=bu.second) assert_quantity(result_q, jnp.asarray([1, 2, 3]), bu.second) + q1 = [1, 2, 3] * bu.second + result_q = bu.math.asarray(q1, unit=bu.ms) + assert_quantity(result_q, jnp.asarray([1, 2, 3]) * 1000, bu.ms) + def test_arange(self): result = bu.math.arange(5) self.assertEqual(result.shape, (5,)) @@ -171,7 +175,7 @@ def test_arange(self): assert_quantity(result_q, jnp.arange(5, step=1), bu.second) result_q = bu.math.arange(3 * bu.second, 9 * bu.second, 1 * bu.second) - assert_quantity(result_q, jnp.arange(3, 9, 1), bu.second) + assert_quantity(result_q, jnp.arange(3, 9, 1), bu.ms) def test_linspace(self): result = bu.math.linspace(0, 10, 5)