make units an enum

nrel/hive/model/vehicle/mechatronics/bev.py

@@ -151,7 +151,7 @@ def consume_energy(self, vehicle: Vehicle, route: Route) -> Vehicle:
         """
         energy_used = self.powertrain.energy_cost(route)
         energy_used_kwh = energy_used * get_unit_conversion(
-            self.powertrain.energy_units, "kilowatthour"
+            self.powertrain.energy_units, Unit.KILOWATT_HOUR
         )
         vehicle_energy_kwh = vehicle.energy[EnergyType.ELECTRIC]
         new_energy_kwh = max(0.0, vehicle_energy_kwh - energy_used_kwh)

nrel/hive/model/vehicle/mechatronics/ice.py

@@ -134,7 +134,7 @@ def consume_energy(self, vehicle: Vehicle, route: Route) -> Vehicle:
         """
         energy_used = self.powertrain.energy_cost(route)
         energy_used_gal_gas = energy_used * get_unit_conversion(
-            self.powertrain.energy_units, "gal_gas"
+            self.powertrain.energy_units, Unit.GALLON_GASOLINE
         )
 
         vehicle_energy_gal_gas = vehicle.energy[EnergyType.GASOLINE]

nrel/hive/model/vehicle/mechatronics/powertrain/powertrain.py

@@ -4,13 +4,14 @@
 from dataclasses import dataclass
 
 from nrel.hive.model.roadnetwork.route import Route
+from nrel.hive.util.units import Unit
 
 
 @dataclass(frozen=True)
 class PowertrainMixin:
-    speed_units: str
-    distance_units: str
-    energy_units: str
+    speed_units: Unit
+    distance_units: Unit
+    energy_units: Unit
 
 
 class PowertrainABC(ABC):

nrel/hive/model/vehicle/mechatronics/powertrain/tabular_powertrain.py

@@ -3,11 +3,10 @@
 
 import numpy as np
 
-from nrel.hive.model.roadnetwork.link import Link
 from nrel.hive.model.roadnetwork.linktraversal import LinkTraversal
 from nrel.hive.model.roadnetwork.routetraversal import Route
 from nrel.hive.model.vehicle.mechatronics.powertrain.powertrain import Powertrain
-from nrel.hive.util.units import valid_unit, get_unit_conversion
+from nrel.hive.util.units import Unit, get_unit_conversion
 
 
 @dataclass(frozen=True)
@@ -16,9 +15,9 @@ class TabularPowertrain(Powertrain):
     builds a tabular, interpolated lookup model for energy consumption
     """
 
-    speed_units: str
-    distance_units: str
-    energy_units: str
+    speed_units: Unit
+    distance_units: Unit
+    energy_units: Unit
 
     consumption_speed: np.ndarray
     consumption_energy_per_distance: np.ndarray
@@ -43,16 +42,14 @@ def from_data(
             if key not in data:
                 raise IOError(f"invalid input file for tabular power train model missing key {key}")
 
-        if not valid_unit(data["speed_units"]):
-            raise TypeError(f"{data['speed_units']} not a recognized unit in hive")
-        elif not valid_unit(data["distance_units"]):
-            raise TypeError(f"{data['distance_units']} not a recognized unit in hive")
-        elif not valid_unit(data["energy_units"]):
-            raise TypeError(f"{data['energy_units']} not a recognized unit in hive")
-
-        speed_units = data["speed_units"]
-        energy_units = data["energy_units"]
-        distance_units = data["distance_units"]
+        try:
+            speed_units = Unit.from_string(data["speed_units"])
+            energy_units = Unit.from_string(data["energy_units"])
+            distance_units = Unit.from_string(data["distance_units"])
+        except ValueError as e:
+            raise ValueError(
+                "Failed to parse incoming units when building TabularPowertrain"
+            ) from e
 
         # linear interpolation function approximation via these lookup values
         consumption_model = sorted(data["consumption_model"], key=lambda x: x["speed"])
@@ -79,7 +76,7 @@ def link_cost(self, link: LinkTraversal) -> float:
         :return: energy in units captured by self.energy_units
         """
         # convert kilometers per hour to whatever units are used by this powertrain
-        link_speed = link.speed_kmph * get_unit_conversion("kmph", self.speed_units)
+        link_speed = link.speed_kmph * get_unit_conversion(Unit.KMPH, self.speed_units)
 
         energy_per_distance = float(
             np.interp(
@@ -89,7 +86,7 @@ def link_cost(self, link: LinkTraversal) -> float:
             )
         )
         # link distance is in kilometers
-        link_distance = link.distance_km * get_unit_conversion("kilometer", self.distance_units)
+        link_distance = link.distance_km * get_unit_conversion(Unit.KILOMETERS, self.distance_units)
         energy = energy_per_distance * link_distance
         return energy
 

nrel/hive/util/units.py

@@ -1,7 +1,40 @@
-# Energy
-from typing import Dict
-
+from __future__ import annotations
 
+from typing import Dict
+from enum import Enum, auto
+
+
+class Unit(Enum):
+    MPH = auto()
+    KMPH = auto()
+    MILES = auto()
+    KILOMETERS = auto()
+    WATT_HOUR = auto()
+    KILOWATT_HOUR = auto()
+    GALLON_GASOLINE = auto()
+
+    @classmethod
+    def from_string(cls, string: str) -> Unit:
+        s = string.strip().lower()
+        if s in ["mph", "miles_per_hour"]:
+            return Unit.MPH
+        if s in ["kmph", "kilomters_per_hour"]:
+            return Unit.KMPH
+        if s in ["mile", "miles", "mi"]:
+            return Unit.MILES
+        if s in ["kilometers", "kilometer", "km"]:
+            return Unit.KILOMETERS
+        if s in ["watthour", "watt-hour", "watt_hour", "wh"]:
+            return Unit.WATT_HOUR
+        if s in ["kilowatthour", "kilowatt-hour", "kilowatt_hour", "kwh"]:
+            return Unit.KILOWATT_HOUR
+        if s in ["gge", "gallon_gasoline", "gal_gas"]:
+            return Unit.GALLON_GASOLINE
+        else:
+            raise ValueError(f"Could not find unit from {string}")
+
+
+## TYPE ALIAS
 KwH = float  # kilowatt-hours
 J = float  # joules
 KwH_per_H = float
@@ -33,7 +66,7 @@
 Percentage = float  # between 0-100
 Ratio = float  # between 0-1
 
-# Conversions
+## CONVERSIONS
 #    Time
 HOURS_TO_SECONDS = 3600
 
@@ -62,39 +95,30 @@ def hours_to_seconds(hours: Hours) -> Seconds:
 KWH_TO_WH = 1 / WH_TO_KWH
 MilesPerGallon = float
 
-UNIT_CONVERSIONS: Dict[str, Dict[str, float]] = {
-    "mph": {
-        "kmph": MPH_TO_KMPH,
+UNIT_CONVERSIONS: Dict[Unit, Dict[Unit, float]] = {
+    Unit.MPH: {
+        Unit.KMPH: MPH_TO_KMPH,
     },
-    "kmph": {
-        "mph": KMPH_TO_MPH,
+    Unit.KMPH: {
+        Unit.MPH: KMPH_TO_MPH,
     },
-    "mile": {
-        "kilometer": M_TO_KM,
+    Unit.MILES: {
+        Unit.KILOMETERS: M_TO_KM,
     },
-    "kilometer": {
-        "mile": KM_TO_MILE,
+    Unit.KILOMETERS: {
+        Unit.MILES: KM_TO_MILE,
     },
-    "watthour": {
-        "kilowatthour": WH_TO_KWH,
+    Unit.WATT_HOUR: {
+        Unit.KILOWATT_HOUR: WH_TO_KWH,
     },
-    "kilowatthour": {
-        "watthour": KWH_TO_WH,
+    Unit.KILOWATT_HOUR: {
+        Unit.WATT_HOUR: KWH_TO_WH,
     },
-    "gal_gas": {},
 }
 
 
-def valid_unit(unit: str) -> bool:
-    return unit in UNIT_CONVERSIONS.keys()
-
-
-def get_unit_conversion(from_unit: str, to_unit: str) -> float:
-    if not valid_unit(from_unit):
-        raise TypeError(f"{from_unit} not a recognized unit in hive")
-    elif not valid_unit(to_unit):
-        raise TypeError(f"{to_unit} not a recognized unit in hive")
-    elif from_unit == to_unit:
+def get_unit_conversion(from_unit: Unit, to_unit: Unit) -> float:
+    if from_unit == to_unit:
         return 1
 
     from_conversion = UNIT_CONVERSIONS.get(from_unit)