code: # -*- coding: utf-8 -*-
  """
  Test symbolic unit handling.
  
  
  
  
  """
  
  # -----------------------------------------------------------------------------
  # Copyright (c) 2018, yt Development Team.
  #
  # Distributed under the terms of the Modified BSD License.
  #
  # The full license is in the LICENSE file, distributed with this software.
  # -----------------------------------------------------------------------------
  
  
  import numpy as np
  from numpy.testing import (
      assert_almost_equal,
      assert_allclose,
      assert_array_almost_equal_nulp,
      assert_equal,
  )
  import operator
  import pickle
  import pytest
  from sympy import Symbol
  
  from unyt.testing import assert_allclose_units
  from unyt.unit_registry import UnitRegistry
  from unyt.dimensions import (
      mass,
      length,
      time,
      temperature,
      energy,
      magnetic_field_cgs,
      magnetic_field_mks,
      power,
      rate,
  )
  from unyt.exceptions import InvalidUnitOperation, UnitsNotReducible, UnitConversionError
  from unyt.unit_object import default_unit_registry, Unit, UnitParseError
  from unyt.unit_systems import cgs_unit_system, UnitSystem
  from unyt._unit_lookup_table import (
      default_unit_symbol_lut,
      name_alternatives,
      unit_prefixes,
  )
  import unyt.unit_symbols as unit_symbols
  from unyt._physical_ratios import (
      m_per_pc,
      sec_per_year,
      m_per_km,
      m_per_mpc,
      mass_sun_kg,
  )
  
  
  def test_no_conflicting_symbols():
      """
      Check unit symbol definitions for conflicts.
  
      """
      full_set = set(default_unit_symbol_lut.keys())
  
      # go through all possible prefix combos
      for symbol in default_unit_symbol_lut.keys():
          if default_unit_symbol_lut[symbol][4]:
              keys = unit_prefixes.keys()
          else:
              keys = [symbol]
          for prefix in keys:
              new_symbol = "%s%s" % (prefix, symbol)
  
              # test if we have seen this symbol
              assert new_symbol not in full_set, "Duplicate symbol: %s" % new_symbol
  
              full_set.add(new_symbol)
  
  
  def test_dimensionless():
      """
      Create dimensionless unit and check attributes.
  
      """
      u1 = Unit()
  
      assert u1.is_dimensionless
      assert u1.expr == 1
      assert u1.base_value == 1
      assert u1.dimensions == 1
      assert u1 != "hello!"
      assert (u1 == "hello") is False
  
      u2 = Unit("")
  
      assert u2.is_dimensionless
      assert u2.expr == 1
      assert u2.base_value == 1
      assert u2.dimensions == 1
  
      assert_equal(u1.latex_repr, "")
      assert_equal(u2.latex_repr, "")
  
  
  def test_create_from_string():
      """
      Create units with strings and check attributes.
  
      """
  
      u1 = Unit("kg * m**2 * s**-2")
      assert u1.dimensions == energy
      assert u1.base_value == 1.0
  
      # make sure order doesn't matter
      u2 = Unit("m**2 * s**-2 * kg")
      assert u2.dimensions == energy
      assert u2.base_value == 1.0
  
      # Test rationals
      u3 = Unit("kg**0.5 * m**-0.5 * s**-1")
      assert u3.dimensions == magnetic_field_cgs
      assert u3.base_value == 1.0
  
      # sqrt functions
      u4 = Unit("sqrt(kg)/sqrt(m)/s")
      assert u4.dimensions == magnetic_field_cgs
      assert u4.base_value == 1.0
  
      # commutative sqrt function
      u5 = Unit("sqrt(kg/m)/s")
      assert u5.dimensions == magnetic_field_cgs
      assert u5.base_value == 1.0
  
      # nonzero CGS conversion factor
      u6 = Unit("Msun/pc**3")
      assert u6.dimensions == mass / length ** 3
      assert_array_almost_equal_nulp(
          np.array([u6.base_value]), np.array([mass_sun_kg / m_per_pc ** 3])
      )
  
      with pytest.raises(UnitParseError):
          Unit("m**m")
      with pytest.raises(UnitParseError):
          Unit("m**g")
      with pytest.raises(UnitParseError):
          Unit("m+g")
      with pytest.raises(UnitParseError):
          Unit("m-g")
      with pytest.raises(UnitParseError):
          Unit("hello!")
      with pytest.raises(UnitParseError):
          Unit("True")
      with pytest.raises(UnitParseError):
          Unit("else")
      with pytest.raises(UnitParseError):
          Unit("hello(37)")
      with pytest.raises(UnitParseError):
          Unit("hello(foo=37)")
  
      cm = Unit("cm")
      data = 1 * cm
  
      assert Unit(data) == cm
      assert Unit(b"cm") == cm
  
  
  def test_create_from_expr():
      """
      Create units from sympy Exprs and check attributes.
  
      """
      pc_mks = m_per_pc
      yr_mks = sec_per_year
  
      # Symbol expr
      s1 = Symbol("pc", positive=True)
      s2 = Symbol("yr", positive=True)
      # Mul expr
      s3 = s1 * s2
      # Pow expr
      s4 = s1 ** 2 * s2 ** (-1)
  
      u1 = Unit(s1)
      u2 = Unit(s2)
      u3 = Unit(s3)
      u4 = Unit(s4)
  
      assert u1.expr == s1
      assert u2.expr == s2
      assert u3.expr == s3
      assert u4.expr == s4
  
      assert_allclose_units(u1.base_value, pc_mks, 1e-12)
      assert_allclose_units(u2.base_value, yr_mks, 1e-12)
      assert_allclose_units(u3.base_value, pc_mks * yr_mks, 1e-12)
      assert_allclose_units(u4.base_value, pc_mks ** 2 / yr_mks, 1e-12)
  
      assert u1.dimensions == length
      assert u2.dimensions == time
      assert u3.dimensions == length * time
      assert u4.dimensions == length ** 2 / time
  
  
  def test_create_with_duplicate_dimensions():
      """
      Create units with overlapping dimensions. Ex: km/Mpc.
  
      """
  
      u1 = Unit("J * s**-1")
      u2 = Unit("km/s/Mpc")
      km_mks = m_per_km
      Mpc_mks = m_per_mpc
  
      assert u1.base_value == 1
      assert u1.dimensions == power
  
      assert_allclose_units(u2.base_value, km_mks / Mpc_mks, 1e-12)
      assert u2.dimensions == rate
  
  
  def test_create_new_symbol():
      """
      Create unit with unknown symbol.
  
      """
      u1 = Unit("abc", base_value=42, dimensions=(mass / time))
  
      assert u1.expr == Symbol("abc", positive=True)
      assert u1.base_value == 42
      assert u1.dimensions == mass / time
  
      u1 = Unit("abc", base_value=42, dimensions=length ** 3)
  
      assert u1.expr == Symbol("abc", positive=True)
      assert u1.base_value == 42
      assert u1.dimensions == length ** 3
  
      u1 = Unit("abc", base_value=42, dimensions=length * (mass * length))
  
      assert u1.expr == Symbol("abc", positive=True)
      assert u1.base_value == 42
      assert u1.dimensions == length ** 2 * mass
  
      with pytest.raises(UnitParseError):
          Unit("abc", base_value=42, dimensions=length ** length)
      with pytest.raises(UnitParseError):
          Unit("abc", base_value=42, dimensions=length ** (length * length))
      with pytest.raises(UnitParseError):
          Unit("abc", base_value=42, dimensions=length - mass)
      with pytest.raises(UnitParseError):
          Unit("abc", base_value=42, dimensions=length + mass)
  
  
  def test_create_fail_on_unknown_symbol():
      """
      Fail to create unit with unknown symbol, without base_value and dimensions.
  
      """
      with pytest.raises(UnitParseError):
          Unit(Symbol("jigawatts"))
  
  
  def test_create_fail_on_bad_symbol_type():
      """
      Fail to create unit with bad symbol type.
  
      """
      with pytest.raises(UnitParseError):
          Unit([1])  # something other than Expr and str
  
  
  def test_create_fail_on_bad_dimensions_type():
      """
      Fail to create unit with bad dimensions type.
  
      """
      with pytest.raises(UnitParseError):
          Unit("a", base_value=1, dimensions="(mass)")
  
  
  def test_create_fail_on_dimensions_content():
      """
      Fail to create unit with bad dimensions expr.
  
      """
      a = Symbol("a")
      with pytest.raises(UnitParseError):
          Unit("a", base_value=1, dimensions=a)
  
  
  def test_create_fail_on_base_value_type():
      """
      Fail to create unit with bad base_value type.
  
      """
      with pytest.raises(UnitParseError):
          Unit("a", base_value="a", dimensions=(mass / time))
  
  
  def test_string_representation():
      """
      Check unit string representation.
  
      """
      pc = Unit("pc")
      Myr = Unit("Myr")
      speed = pc / Myr
      dimensionless = Unit()
  
      assert str(pc) == "pc"
      assert str(Myr) == "Myr"
      assert str(speed) == "pc/Myr"
      assert repr(speed) == "pc/Myr"
      assert str(dimensionless) == "dimensionless"
      assert repr(dimensionless) == "(dimensionless)"
  
  
  def test_multiplication():
      """
      Multiply two units.
  
      """
      msun_mks = mass_sun_kg
      pc_mks = m_per_pc
  
      # Create symbols
      msun_sym = Symbol("Msun", positive=True)
      pc_sym = Symbol("pc", positive=True)
      s_sym = Symbol("s", positive=True)
  
      # Create units
      u1 = Unit("Msun")
      u2 = Unit("pc")
  
      # Mul operation
      u3 = u1 * u2
  
      assert u3.expr == msun_sym * pc_sym
      assert_allclose_units(u3.base_value, msun_mks * pc_mks, 1e-12)
      assert u3.dimensions == mass * length
  
      # Pow and Mul operations
      u4 = Unit("pc**2")
      u5 = Unit("Msun * s")
  
      u6 = u4 * u5
  
      assert u6.expr == pc_sym ** 2 * msun_sym * s_sym
      assert_allclose_units(u6.base_value, pc_mks ** 2 * msun_mks, 1e-12)
      assert u6.dimensions == length ** 2 * mass * time
  
  
  def test_division():
      """
      Divide two units.
  
      """
      pc_mks = m_per_pc
      km_mks = m_per_km
  
      # Create symbols
      pc_sym = Symbol("pc", positive=True)
      km_sym = Symbol("km", positive=True)
      s_sym = Symbol("s", positive=True)
  
      # Create units
      u1 = Unit("pc")
      u2 = Unit("km * s")
  
      u3 = u1 / u2
  
      assert u3.expr == pc_sym / (km_sym * s_sym)
      assert_allclose_units(u3.base_value, pc_mks / km_mks, 1e-12)
      assert u3.dimensions == 1 / time
  
  
  def test_power():
      """
      Take units to some power.
  
      """
      from sympy import nsimplify
  
      pc_mks = m_per_pc
      mK_mks = 1e-3
      u1_dims = mass * length ** 2 * time ** -3 * temperature ** 4
      u1 = Unit("kg * pc**2 * s**-3 * mK**4")
  
      u2 = u1 ** 2
  
      assert u2.dimensions == u1_dims ** 2
      assert_allclose_units(u2.base_value, (pc_mks ** 2 * mK_mks ** 4) ** 2, 1e-12)
  
      u3 = u1 ** (-1.0 / 3)
  
      assert u3.dimensions == nsimplify(u1_dims ** (-1.0 / 3))
      assert_allclose_units(
          u3.base_value, (pc_mks ** 2 * mK_mks ** 4) ** (-1.0 / 3), 1e-12
      )
  
  
  def test_equality():
      """
      Check unit equality with different symbols, but same dimensions and
      base_value.
  
      """
      u1 = Unit("km * s**-1")
      u2 = Unit("m * ms**-1")
  
      assert u1 == u2
      assert u1.copy() == u2
  
  
  def test_invalid_operations():
      u1 = Unit("cm")
      u2 = Unit("m")
  
      with pytest.raises(InvalidUnitOperation):
          u1 + u2
      with pytest.raises(InvalidUnitOperation):
          u1 += u2
      with pytest.raises(InvalidUnitOperation):
          1 + u1
      with pytest.raises(InvalidUnitOperation):
          u1 + 1
      with pytest.raises(InvalidUnitOperation):
          u1 - u2
      with pytest.raises(InvalidUnitOperation):
          u1 -= u2
      with pytest.raises(InvalidUnitOperation):
          1 - u1
      with pytest.raises(InvalidUnitOperation):
          u1 - 1
      with pytest.raises(InvalidUnitOperation):
          u1 *= u2
      with pytest.raises(InvalidUnitOperation):
          u1 * "hello!"
      with pytest.raises(InvalidUnitOperation):
          u1 /= u2
      with pytest.raises(InvalidUnitOperation):
          u1 / "hello!"
  
  
  def test_base_equivalent():
      """
      Check base equivalent of a unit.
  
      """
      Msun_mks = mass_sun_kg
      Mpc_mks = m_per_mpc
  
      u1 = Unit("Msun * Mpc**-3")
      u2 = Unit("kg * m**-3")
      u3 = u1.get_base_equivalent()
  
      assert u2.expr == u3.expr
      assert u2 == u3
  
      assert_allclose_units(u1.base_value, Msun_mks / Mpc_mks ** 3, 1e-12)
      assert u2.base_value == 1
      assert u3.base_value == 1
  
      mass_density = mass / length ** 3
  
      assert u1.dimensions == mass_density
      assert u2.dimensions == mass_density
      assert u3.dimensions == mass_density
  
      assert_allclose_units(
          u1.get_conversion_factor(u3)[0], Msun_mks / Mpc_mks ** 3, 1e-12
      )
  
      with pytest.raises(UnitConversionError):
          u1.get_conversion_factor(Unit("m"))
  
      with pytest.raises(UnitConversionError):
          u1.get_conversion_factor(Unit("degF"))
  
      reg = UnitRegistry(unit_system=cgs_unit_system)
  
      u = Unit("kg", registry=reg)
  
      assert u.get_base_equivalent() == Unit("g")
  
      u = Unit("kg")
  
      assert u.get_base_equivalent() == Unit("kg")
  
      u = Unit("A")
      assert u.get_base_equivalent(unit_system="mks") == Unit("A")
  
  
  def test_temperature_offsets():
      u1 = Unit("degC")
      u2 = Unit("degF")
  
      with pytest.raises(InvalidUnitOperation):
          operator.mul(u1, u2)
      with pytest.raises(InvalidUnitOperation):
          operator.truediv(u1, u2)
  
  
  def test_latex_repr():
      registry = UnitRegistry()
  
      # create a fake comoving unit
      registry.add(
          "pccm",
          registry.lut["pc"][0] / (1 + 2),
          length,
          "\\rm{pc}/(1+z)",
          prefixable=True,
      )
  
      test_unit = Unit("Mpccm", registry=registry)
      assert_almost_equal(test_unit.base_value, m_per_mpc / 3)
      assert_equal(test_unit.latex_repr, r"\rm{Mpc}/(1+z)")
  
      test_unit = Unit("cm**-3", base_value=1.0, registry=registry)
      assert_equal(test_unit.latex_repr, "\\frac{1}{\\rm{cm}^{3}}")
  
      test_unit = Unit("m_geom/l_geom**3")
      assert_equal(test_unit.latex_repr, "\\frac{1}{M_\\odot^{2}}")
  
      test_unit = Unit("1e9*cm")
      assert_equal(test_unit.latex_repr, "1.0 \\times 10^{9}\\ \\rm{cm}")
  
      test_unit = Unit("1.0*cm")
      assert_equal(test_unit.latex_repr, "\\rm{cm}")
  
  
  def test_latitude_longitude():
      lat = unit_symbols.lat
      lon = unit_symbols.lon
      deg = unit_symbols.deg
      assert_equal(lat.units.base_offset, 90.0)
      assert_equal((deg * 90.0).in_units("lat").value, 0.0)
      assert_equal((deg * 180).in_units("lat").value, -90.0)
      assert_equal((lat * 0.0).in_units("deg"), deg * 90.0)
      assert_equal((lat * -90).in_units("deg"), deg * 180)
  
      assert_equal(lon.units.base_offset, -180.0)
      assert_equal((deg * 0.0).in_units("lon").value, -180.0)
      assert_equal((deg * 90.0).in_units("lon").value, -90.0)
      assert_equal((deg * 180).in_units("lon").value, 0.0)
      assert_equal((deg * 360).in_units("lon").value, 180.0)
  
      assert_equal((lon * -180.0).in_units("deg"), deg * 0.0)
      assert_equal((lon * -90.0).in_units("deg"), deg * 90.0)
      assert_equal((lon * 0.0).in_units("deg"), deg * 180.0)
      assert_equal((lon * 180.0).in_units("deg"), deg * 360)
  
  
  def test_creation_from_ytarray():
      from unyt import electrostatic_unit, elementary_charge_cgs
  
      u1 = Unit(electrostatic_unit)
      assert_equal(str(u1), "statC")
      assert_equal(u1, Unit("esu"))
      assert_equal(u1, electrostatic_unit.units)
  
      u2 = Unit(elementary_charge_cgs)
      assert_equal(str(u2), "4.80320467299766e-10*statC")
      assert_equal(u2, Unit("4.80320467299766e-10*statC"))
      assert_equal(u1, elementary_charge_cgs.units)
  
      assert_allclose((u1 / u2).base_value, electrostatic_unit / elementary_charge_cgs)
  
      with pytest.raises(UnitParseError):
          Unit([1, 2, 3] * elementary_charge_cgs)
  
  
  def test_list_same_dimensions():
      from unyt import m
  
      reg = default_unit_registry
      for equiv in reg.list_same_dimensions(m):
          assert Unit(equiv).dimensions is length
  
  
  def test_decagram():
      dag = Unit("dag")
      g = Unit("g")
      assert dag.get_conversion_factor(g) == (10.0, None)
  
  
  def test_pickle():
      cm = Unit("cm")
      assert cm == pickle.loads(pickle.dumps(cm))
  
  
  def test_preserve_offset():
      from unyt import degF, dimensionless
  
      new_unit = degF * dimensionless
  
      assert new_unit is not degF
      assert new_unit == degF
      assert new_unit.base_offset == degF.base_offset
  
      new_unit = degF / dimensionless
  
      assert new_unit is not degF
      assert new_unit == degF
      assert new_unit.base_offset == degF.base_offset
  
      with pytest.raises(InvalidUnitOperation):
          dimensionless / degF
  
  
  def test_code_unit():
      from unyt import UnitRegistry
  
      ureg = UnitRegistry()
      ureg.add("code_length", 10.0, length)
      ureg.add("code_magnetic_field", 2.0, magnetic_field_mks)
      u = Unit("code_length", registry=ureg)
      assert u.is_code_unit is True
      assert u.get_base_equivalent() == Unit("m")
      u = Unit("cm")
      assert u.is_code_unit is False
  
      u = Unit("code_magnetic_field", registry=ureg)
      assert u.get_base_equivalent("mks") == Unit("T")
      with pytest.raises(UnitsNotReducible):
          assert u.get_base_equivalent("cgs")
  
      # see issue #60
      u = Unit("s/m")
      assert u.get_mks_equivalent() == Unit("s/m")
      assert u.get_mks_equivalent() != Unit("ohm")
      assert u.get_cgs_equivalent() == Unit("s/cm")
  
      u = Unit("kC")
      assert u.get_cgs_equivalent() == Unit("kesu")
      assert u.get_cgs_equivalent().get_mks_equivalent() == u
  
      UnitSystem(ureg.unit_system_id, "code_length", "kg", "s", registry=ureg)
  
      u = Unit("cm", registry=ureg)
      ue = u.get_base_equivalent("code")
  
      assert str(ue) == "code_length"
      assert ue.base_value == 10
      assert ue.dimensions is length
  
      class FakeDataset(object):
          unit_registry = ureg
  
      ds = FakeDataset()
  
      UnitSystem(ds, "code_length", "kg", "s", registry=ureg)
  
      u = Unit("cm", registry=ureg)
      ue = u.get_base_equivalent(ds)
  
      assert str(ue) == "code_length"
      assert ue.base_value == 10
      assert ue.dimensions is length
  
      with pytest.raises(UnitParseError):
          Unit("code_length")
  
  
  def test_bad_equivalence():
      from unyt import cm
  
      with pytest.raises(KeyError):
          cm.has_equivalent("dne")
  
  
  def test_em_unit_base_equivalent():
      from unyt import A, cm
  
      with pytest.raises(UnitsNotReducible):
          (A / cm).get_base_equivalent("cgs")
  
  
  def test_symbol_lut_length():
      for v in default_unit_symbol_lut.values():
          assert len(v) == 5
  
  
  def test_simplify():
      import unyt as u
  
      answers = {
          u.Hz * u.s: "dimensionless",
          u.kg / u.g: "1000",
          u.Hz * u.s * u.km: "km",
          u.kHz * u.s: "1000",
          u.kHz * u.s * u.km: "1000*km",
          u.kHz * u.s ** 2: "1000*s",
          u.kHz * u.s ** 2 * u.km: "1000*km*s",
          u.Hz ** -1 * u.s: "s/Hz",
          u.Hz ** -1 * u.s * u.km: "km*s/Hz",
          u.Hz ** 1.5 * u.s ** 1.7: "sqrt(Hz)*s**(7/10)",
          u.Hz ** 1.5 * u.s ** 1.7 * u.km: "sqrt(Hz)*km*s**(7/10)",
          u.m ** 2 / u.cm ** 2: "10000",
      }
  
      for unit, answer in answers.items():
          assert str(unit.simplify()) == answer
  
  
  def test_micro_prefix():
      import unyt as u
  
      # both versions of unicode mu work correctly
      assert u.um == u.µm
      assert u.um == u.μm
  
      # parsing both versions works as well
      assert u.ug == u.Unit("µg")
      assert u.ug == u.Unit("μg")
  
  
  def test_name_alternatives():
      import unyt
      from unyt._unit_lookup_table import (
          default_unit_name_alternatives,
          name_alternatives,
          inv_name_alternatives,
      )
  
      # concatenated list of all alternative unit names
      allowed_names = sum(name_alternatives.values(), [])
  
      # ensure the values are all tuples and not e.g. strings
      for val in default_unit_name_alternatives.values():
          assert isinstance(val, tuple)
  
      # all names are unique
      assert len(set(allowed_names)) == len(allowed_names)
      # each allowed name has a key in the inverse dict
      assert len(inv_name_alternatives.keys()) == len(allowed_names)
      assert set(inv_name_alternatives.keys()) == set(allowed_names)
  
      for name in allowed_names:
          assert hasattr(unyt, name)
          assert hasattr(unyt.unit_symbols, name)
  
  
  def test_attosecond():
      from unyt import Unit, attosecond, second
  
      assert Unit("as") == attosecond
      assert str(Unit("as")) == "as"
      assert Unit("as/s") == attosecond / second
  
  
  def test_micro():
      from unyt import Unit
  
      assert str(Unit("um")) == "µm"
      assert str(Unit("us")) == "µs"
  
  
  def test_show_all_units_doc_table_ops():
      for name in set(name_alternatives.keys()):
          u = Unit(name)
          (1 * u).in_mks()
          try:
              (1 * u).in_cgs()
          except UnitsNotReducible:
              pass
  
apis: numpy.testing.assert_equal numpy.array numpy.testing.assert_almost_equal numpy.testing.assert_allclose