Repository: yt-project/unyt (python)

Issue:
`@` operator not supported by `unyt_array`
Performing matrix multiplication on `unyt_array`s using the `@` operator throws a `KeyError`

```
>>> x = unyt.unyt_array([[1]])
>>> y = unyt.unyt_array([[2]])
>>> x @ y
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-17-93e3eb3e5b56> in <module>
----> 1 x @ y

~\AppData\Local\Continuum\Anaconda3\lib\site-packages\unyt\array.py in __array_ufunc__(self, ufunc, method, *inputs, **kwargs)
   1662                 else:
   1663                     u1 = 1.0
-> 1664             unit_operator = self._ufunc_registry[ufunc]
   1665             if unit_operator in (_preserve_units, _comparison_unit, _arctan2_unit):
   1666                 # check "is" equality first for speed

KeyError: <ufunc 'matmul'>
```

`unyt` version = 2.0.0

Pull request description:
add a matmul implementation. fixes #75
I didn't add an `__imatmul__` implementation because that's not supported by upstream numpy yet.

Patch:
diff --git a/.travis.yml b/.travis.yml
index 074bb9a..f114aa7 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -36,7 +36,7 @@ before_install:
   - pip install tox tox-pyenv codecov twine
 
 # Command to run tests, e.g. python setup.py test
-script: tox -p auto -vvv
+script: tox -vvv
 
 after_success: codecov
 
diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst
index 1aa93f1..b09880b 100644
--- a/CONTRIBUTING.rst
+++ b/CONTRIBUTING.rst
@@ -98,7 +98,7 @@ Ready to contribute? Here's how to set up ``unyt`` for local development.
     $ flake8 unyt
     $ black ./
     $ pytest --doctest-modules --doctest-rst --doctest-plus
-    $ tox -p auto
+    $ tox
 
    To get ``flake8``, ``black``, ``pytest``, ``pytest-doctestplus``, and
    ``tox``, just pip install them into your virtualenv.
@@ -127,9 +127,9 @@ run this command. Some tests depend on ``h5py``, ``Pint``, ``astropy``,
 If you would like to run the tests on multiple python versions, first ensure that you have multiple python versions visible on your ``$PATH``, then simply execute ``tox`` in the root of the ``unyt`` repository::
 
    $ cd unyt
-   $ tox -p auto
+   $ tox
 
-The ``tox`` package itself can be installed using the ``pip`` associated with one of the python installations. See the ``tox.ini`` file in the root of the repository for more details about our ``tox`` setup. Note that you do not need to install anything besides ``tox`` and ``python`` for this to work, ``tox`` will handle setting up the test environment, including installing any necessary dependencies via ``pip``. Running ``tox`` with ``-p auto`` will run the tests in parallel and is not required but may make the tests run more quickly.
+The ``tox`` package itself can be installed using the ``pip`` associated with one of the python installations. See the ``tox.ini`` file in the root of the repository for more details about our ``tox`` setup. Note that you do not need to install anything besides ``tox`` and ``python`` for this to work, ``tox`` will handle setting up the test environment, including installing any necessary dependencies via ``pip``.
 
 Pull Request Guidelines
 -----------------------
diff --git a/docs/usage.rst b/docs/usage.rst
index 79d0953..70ffbf3 100644
--- a/docs/usage.rst
+++ b/docs/usage.rst
@@ -938,8 +938,6 @@ contains units that are compatible with the ``Simulation`` instance we named
   >>> units = UnitContainer(s_cgs.registry)
   >>> units.kilometer
   km
-  >>> units.code_length
-  code_length
   >>> (10.0 * units.kilometer).in_base()
   unyt_quantity(1000000., 'cm')
   >>> (10.0 * units.kilometer).in_units('code_length')
diff --git a/unyt/_unit_lookup_table.py b/unyt/_unit_lookup_table.py
index ec24b2c..362e6ea 100644
--- a/unyt/_unit_lookup_table.py
+++ b/unyt/_unit_lookup_table.py
@@ -425,8 +425,8 @@ default_unit_name_alternatives = OrderedDict(
         ("yr", ("year",)),
         # Solar units
         ("Msun", ("solar_mass", "solMass")),
-        ("Rsun", ("rsun", "r_sun", "solar_radius", "solRadius")),
-        ("Lsun", ("lsun", "l_sun", "solar_luminosity", "solLuminosity")),
+        ("Rsun", ("r_sun", "solar_radius", "solRadius")),
+        ("Lsun", ("l_sun", "solar_luminosity", "solLuminosity")),
         ("Tsun", ("t_sun", "solar_temperature", "solTemperature")),
         ("Zsun", ("z_sun", "solar_metallicity", "solMetallicity")),
         ("Mjup", ("m_jup", "jupiter_mass")),
@@ -434,7 +434,7 @@ default_unit_name_alternatives = OrderedDict(
         ("Rjup", ("r_jup", "jupiter_radius")),
         ("Rearth", ("r_earth", "earth_radius")),
         # astro distances
-        ("AU", ("au", "astronomical_unit")),
+        ("AU", ("astronomical_unit",)),
         ("pc", ("parsec",)),
         ("ly", ("light_year",)),
         # angles
@@ -514,8 +514,7 @@ def generate_name_alternatives():
                 append_name(names[key], key, alt)
                 if not alt.islower() or len(alt) < 4:
                     continue
-                if alt.title() not in names[key]:
-                    append_name(names[key], key, alt.title())
+                append_name(names[key], key, alt.title())
     return names, inv_names
 
 
diff --git a/unyt/array.py b/unyt/array.py
index 004a7f0..8aefad3 100644
--- a/unyt/array.py
+++ b/unyt/array.py
@@ -104,6 +104,7 @@ from numpy import (
     isnat,
     heaviside,
     ones_like,
+    matmul,
 )
 from numpy.core.umath import _ones_like
 from sympy import Rational
@@ -469,6 +470,7 @@ class unyt_array(np.ndarray):
         isnat: _return_without_unit,
         heaviside: _preserve_units,
         _ones_like: _preserve_units,
+        matmul: _multiply_units,
     }
 
     __array_priority__ = 2.0
@@ -1803,7 +1805,6 @@ class unyt_array(np.ndarray):
         # numpy issue #9081
         return type(self)(super(unyt_array, self).__pos__(), self.units)
 
-    @_return_arr
     def dot(self, b, out=None):
         """dot product of two arrays.
 
@@ -1828,7 +1829,11 @@ class unyt_array(np.ndarray):
         [[8. 8.]
          [8. 8.]] km*s**2
         """
-        return super(unyt_array, self).dot(b), self.units * b.units
+        res_units = self.units * getattr(b, "units", NULL_UNIT)
+        ret = self.view(np.ndarray).dot(np.asarray(b), out=out) * res_units
+        if out is not None:
+            out.units = res_units
+        return ret
 
     def __reduce__(self):
         """Pickle reduction method
diff --git a/unyt/unit_object.py b/unyt/unit_object.py
index 400d53f..379a5c4 100644
--- a/unyt/unit_object.py
+++ b/unyt/unit_object.py
@@ -15,7 +15,6 @@ A class that represents a unit symbol.
 
 import copy
 import itertools
-import math
 
 from functools import lru_cache
 from keyword import iskeyword as _iskeyword
@@ -548,16 +547,13 @@ class Unit(object):
         """ Test unit equality. """
         if not isinstance(u, Unit):
             return False
-        return (
-            math.isclose(self.base_value, u.base_value)
-            and self.dimensions == u.dimensions
-        )
+        return self.base_value == u.base_value and self.dimensions == u.dimensions
 
     def __ne__(self, u):
         """ Test unit inequality. """
         if not isinstance(u, Unit):
             return True
-        if not math.isclose(self.base_value, u.base_value):
+        if self.base_value != u.base_value:
             return True
         # use 'is' comparison dimensions to avoid expensive sympy operation
         if self.dimensions is u.dimensions:
diff --git a/unyt/unit_registry.py b/unyt/unit_registry.py
index 9b2edb5..46ff81a 100644
--- a/unyt/unit_registry.py
+++ b/unyt/unit_registry.py
@@ -34,7 +34,7 @@ def _sanitize_unit_system(unit_system, obj):
     elif hasattr(unit_system, "unit_registry"):
         unit_system = unit_system.unit_registry.unit_system_id
     elif unit_system == "code":
-        unit_system = obj.units.registry.unit_system_id
+        unit_system = obj.registry.unit_system_id
     return unit_system_registry[str(unit_system)]
 
 
diff --git a/unyt/unit_systems.py b/unyt/unit_systems.py
index e235600..ae770a0 100644
--- a/unyt/unit_systems.py
+++ b/unyt/unit_systems.py
@@ -55,9 +55,6 @@ def add_symbols(namespace, registry):
     for canonical_name, alt_names in name_alternatives.items():
         for alt_name in alt_names:
             namespace[alt_name] = Unit(canonical_name, registry=registry)
-    for name in registry.keys():
-        if name not in namespace:
-            namespace[name] = Unit(name, registry=registry)
 
 
 def add_constants(namespace, registry):
@@ -278,11 +275,6 @@ class UnitSystem(object):
                 repr += "  %s: %s\n" % (key, self.units_map[dim])
         return repr[:-1]
 
-    @property
-    def has_current_mks(self):
-        """Does this unit system have an MKS current dimension?"""
-        return self.units_map[cmks] is not None
-
 
 #: The CGS unit system
 cgs_unit_system = UnitSystem("cgs", "cm", "g", "s", current_mks_unit=None)


Test patch:
diff --git a/unyt/tests/test_unit_systems.py b/unyt/tests/test_unit_systems.py
index 58aa047..adf7541 100644
--- a/unyt/tests/test_unit_systems.py
+++ b/unyt/tests/test_unit_systems.py
@@ -17,12 +17,7 @@ import unyt.unit_symbols as us
 
 from unyt.exceptions import IllDefinedUnitSystem, MissingMKSCurrent
 from unyt.unit_object import Unit
-from unyt.unit_systems import (
-    UnitSystem,
-    cgs_unit_system,
-    mks_unit_system,
-    unit_system_registry,
-)
+from unyt.unit_systems import UnitSystem, cgs_unit_system, unit_system_registry
 from unyt.unit_registry import UnitRegistry
 from unyt import dimensions
 
@@ -74,30 +69,6 @@ def test_bad_unit_system():
         UnitSystem("atomic", us.nm, us.fs, us.nK, us.rad, registry=UnitRegistry())
 
 
-def test_code_unit_system():
-    ureg = UnitRegistry()
-    ureg.add("code_length", 2.0, dimensions.length)
-    ureg.add("code_mass", 3.0, dimensions.mass)
-    ureg.add("code_time", 4.0, dimensions.time)
-    ureg.add("code_temperature", 5.0, dimensions.temperature)
-    code_unit_system = UnitSystem(
-        "my_unit_system",
-        "code_length",
-        "code_mass",
-        "code_time",
-        "code_temperature",
-        registry=ureg,
-    )
-    assert code_unit_system["length"] == Unit("code_length", registry=ureg)
-    assert code_unit_system["length"].base_value == 2
-    assert code_unit_system["mass"] == Unit("code_mass", registry=ureg)
-    assert code_unit_system["mass"].base_value == 3
-    assert code_unit_system["time"] == Unit("code_time", registry=ureg)
-    assert code_unit_system["time"].base_value == 4
-    assert code_unit_system["temperature"] == Unit("code_temperature", registry=ureg)
-    assert code_unit_system["temperature"].base_value == 5
-
-
 def test_mks_current():
     with pytest.raises(MissingMKSCurrent):
         cgs_unit_system[dimensions.current_mks]
@@ -107,8 +78,6 @@ def test_mks_current():
         cgs_unit_system[dimensions.current_mks] = "foo"
     with pytest.raises(MissingMKSCurrent):
         cgs_unit_system[dimensions.magnetic_field] = "bar"
-    assert cgs_unit_system.has_current_mks is False
-    assert mks_unit_system.has_current_mks is True
 
 
 def test_create_unit_system_from_unit_objects():
diff --git a/unyt/tests/test_unyt_array.py b/unyt/tests/test_unyt_array.py
index 545fb77..8fc775f 100644
--- a/unyt/tests/test_unyt_array.py
+++ b/unyt/tests/test_unyt_array.py
@@ -1102,20 +1102,92 @@ def test_ufuncs():
             binary_ufunc_comparison(ufunc, pair[0], pair[1])
 
 
-def test_reductions():
-    arr = unyt_array([[1, 2, 3], [4, 5, 6]], "cm")
+@pytest.mark.skipif(
+    np.__version__ < "1.16", reason="matmul is broken on old numpy versions"
+)
+def test_dot_matmul():
+    arr = unyt_array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], "cm")
+
+    ev_result = arr.dot(unyt_array([1.0, 2.0, 3.0], "kg"))
+    matmul_result = arr @ unyt_array([1.0, 2.0, 3.0], "kg")
+    res = unyt_array([14.0, 32.0], "cm*kg")
+    assert_equal(ev_result, res)
+    assert_equal(ev_result.units, res.units)
+    assert_isinstance(ev_result, unyt_array)
+    assert_equal(matmul_result, res)
+    assert_equal(matmul_result.units, res.units)
+    assert_isinstance(matmul_result, unyt_array)
+
+    ev_result = arr.dot(np.array([1.0, 2.0, 3.0]))
+    matmul_result = arr @ np.array([1.0, 2.0, 3.0])
+    res = unyt_array([14.0, 32.0], "cm")
+    assert_equal(ev_result, res)
+    assert_equal(ev_result.units, res.units)
+    assert_isinstance(ev_result, unyt_array)
+    assert_equal(matmul_result, res)
+    assert_equal(matmul_result.units, res.units)
+    assert_isinstance(matmul_result, unyt_array)
+
+    ev_result = arr.dot(arr.T)
+    matmul_result = arr @ arr.T
+    res = unyt_array([[14.0, 32.0], [32.0, 77.0]], "cm**2")
+    assert_equal(ev_result, res)
+    assert_equal(ev_result.units, res.units)
+    assert_isinstance(ev_result, unyt_array)
+    assert_equal(matmul_result, res)
+    assert_equal(matmul_result.units, res.units)
+    assert_isinstance(matmul_result, unyt_array)
 
-    ev_result = arr.dot(unyt_array([1, 2, 3], "cm"))
-    res = unyt_array([14.0, 32.0], "cm**2")
+    ev_result = arr.v.dot(arr.T)
+    matmul_result = arr.v @ arr.T
+    res = unyt_array([[14.0, 32.0], [32.0, 77.0]], "cm")
     assert_equal(ev_result, res)
     assert_equal(ev_result.units, res.units)
     assert_isinstance(ev_result, unyt_array)
+    assert_equal(matmul_result, res)
+    assert_equal(matmul_result.units, res.units)
+    assert_isinstance(matmul_result, unyt_array)
+
+    ev_result = arr.dot(arr.T.v)
+    matmul_result = arr @ arr.T.v
+    res = unyt_array([[14.0, 32.0], [32.0, 77.0]], "cm")
+    assert_equal(ev_result, res)
+    assert_equal(ev_result.units, res.units)
+    assert_isinstance(ev_result, unyt_array)
+    assert_equal(matmul_result, res)
+    assert_equal(matmul_result.units, res.units)
+    assert_isinstance(matmul_result, unyt_array)
+
+    arr = unyt_array([[1.0, 2.0], [3.0, 4.0]], "kg")
+    arr.dot(arr.T, out=arr)
+    res = unyt_array([[5.0, 11.0], [11.0, 25.0]], "kg**2")
+    assert_equal(arr, res)
+    assert_equal(arr.units, res.units)
+    assert_isinstance(arr, unyt_array)
 
     qv = unyt_array([1, 2, 3], "cm").dot(unyt_array([1, 2, 3], "cm"))
+    mv = unyt_array([1, 2, 3], "cm") @ unyt_array([1, 2, 3], "cm")
     qa = unyt_quantity(14, "cm**2")
     assert qv == qa
     assert qv.units == qa.units
     assert_isinstance(qv, unyt_quantity)
+    assert mv == qa
+    assert mv.units == qa.units
+    assert_isinstance(mv, unyt_quantity)
+
+    qv = unyt_array([1, 2, 3], "cm").dot(np.array([1, 2, 3]))
+    mv = unyt_array([1, 2, 3], "cm") @ np.array([1, 2, 3])
+    qa = unyt_quantity(14, "cm")
+    assert qv == qa
+    assert qv.units == qa.units
+    assert_isinstance(qv, unyt_quantity)
+    assert mv == qa
+    assert mv.units == qa.units
+    assert_isinstance(mv, unyt_quantity)
+
+
+def test_reductions():
+    arr = unyt_array([[1, 2, 3], [4, 5, 6]], "cm")
 
     answers = {
         "prod": (