Repository: NREL/infrasys

Issue:
bug: pint unit definitions are not capable of performing multiplications
It looks like the `__new__` method is preventing us to multiply `ActivePower` * `Time` to obtain energy. The reason is that the `__new__` method of `BaseQuantity` is checking unit consistency even on the result multiplication. 

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/NREL/infrasys/pull/25?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=NREL) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 94.79%. Comparing base [(`fe2295d`)](https://app.codecov.io/gh/NREL/infrasys/commit/fe2295d1a9b561dda82fe5ed8ea5a8eb26c27ac7?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=NREL) to head [(`384dd4e`)](https://app.codecov.io/gh/NREL/infrasys/pull/25?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=NREL).


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main      #25      +/-   ##
==========================================
- Coverage   95.75%   94.79%   -0.96%     
==========================================
  Files          24       25       +1     
  Lines        2074     2095      +21     
==========================================
  Hits         1986     1986              
- Misses         88      109      +21     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/NREL/infrasys/pull/25?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=NREL).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=NREL).


Patch:
diff --git a/src/infrasys/base_quantity.py b/src/infrasys/base_quantity.py
index 22c4fac..5a3d9b4 100644
--- a/src/infrasys/base_quantity.py
+++ b/src/infrasys/base_quantity.py
@@ -1,25 +1,71 @@
-""" This module contains base class for handling pint quantity."""
+"""This module contains base class for handling pint quantity."""
 
-from abc import ABC
-from typing import Any
+from typing import Any, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from __main__ import BaseQuantity
 
 import numpy as np
 import pint
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import core_schema
+from typing_extensions import Annotated
 
 ureg = pint.UnitRegistry()
 
 
-class BaseQuantity(ureg.Quantity, ABC):  # type: ignore
+class BaseQuantity(ureg.Quantity):  # type: ignore
     """Interface for base quantity."""
 
-    def __new__(cls, value, units, **kwargs):
-        instance = super().__new__(cls, value, units, **kwargs)
-        if not hasattr(cls, "__compatible_unit__"):
-            raise ValueError("You should define __compatible_unit__ attribute in your class.")
-        if not instance.is_compatible_with(cls.__compatible_unit__):
-            message = f"{__class__} must be compatible with {cls.__compatible_unit__}, not {units}"
-            raise ValueError(message)
-        return instance
+    __base_unit__ = None
+
+    def __init_subclass__(cls, **kwargs):
+        if not cls.__base_unit__:
+            raise TypeError("__base_unit__ should be defined")
+        super().__init_subclass__(**kwargs)
+
+    # NOTE: This creates a type hint for the unit.
+    def __class_getitem__(cls):
+        return Annotated.__class_getitem__((cls, cls.__base_unit__))  # type: ignore
+
+    # Required for pydantic validation
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        return core_schema.with_info_after_validator_function(
+            cls.validate,
+            handler(pint.Quantity),
+            field_name=handler.field_name,
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                cls.serialize, info_arg=False, return_schema=core_schema.str_schema()
+            ),
+        )
+
+    # Required for pydantic validation
+    @classmethod
+    def validate(cls, value, *_) -> "BaseQuantity":
+        if isinstance(value, BaseQuantity):
+            if cls.__base_unit__:
+                assert value.check(
+                    cls.__base_unit__
+                ), f"Unit must be compatible with {cls.__base_unit__}"
+                return cls(value)
+        if isinstance(value, pint.Quantity):
+            if cls.__base_unit__:
+                assert value.check(
+                    cls.__base_unit__
+                ), f"Unit must be compatible with {cls.__base_unit__}"
+                return cls(value.magnitude, value.units)
+            else:
+                raise ValueError(f"Invalid type for BaseQuantity: {type(value)}")
+        if isinstance(value, cls):
+            return cls(value)
+        return value
+
+    @staticmethod
+    def serialize(value):
+        return str(value)
 
     def to_dict(self) -> dict[str, Any]:
         """Convert a quantity to a dictionary for serialization."""
diff --git a/src/infrasys/quantities.py b/src/infrasys/quantities.py
index ca75ed8..4be259d 100644
--- a/src/infrasys/quantities.py
+++ b/src/infrasys/quantities.py
@@ -1,22 +1,33 @@
-"""This module defines basic unit quantities."""
+"""This module defines basic unit quantities.
+
+To create new Quantities for a given base unit, we just need to specify the
+base unit as the second argument of `ureg.check`.
+"""
 
 from infrasys.base_quantity import BaseQuantity
+from infrasys.component import Component
 
 # ruff:noqa
 # fmt: off
 
-class Distance(BaseQuantity): __compatible_unit__ = "meter"
+class Distance(BaseQuantity): __base_unit__ = "meter"
+
+class Voltage(BaseQuantity): __base_unit__ = "volt"
+
+class Current(BaseQuantity): __base_unit__ = "ampere"
+
+class Angle(BaseQuantity): __base_unit__ = "degree"
 
-class Voltage(BaseQuantity): __compatible_unit__ = "volt"
+class ActivePower(BaseQuantity): __base_unit__ = "watt"
 
-class Current(BaseQuantity): __compatible_unit__ = "ampere"
+class Energy(BaseQuantity): __base_unit__ = "watthour"
 
-class Angle(BaseQuantity): __compatible_unit__ = "degree"
+class Time(BaseQuantity): __base_unit__ = "minute"
 
-class ActivePower(BaseQuantity): __compatible_unit__ = "watt"
+class Resistance(BaseQuantity): __base_unit__ = "ohm"
 
-class Energy(BaseQuantity): __compatible_unit__ = "watthour"
 
-class Time(BaseQuantity): __compatible_unit__ = "minute"
+class Test(Component):
+    voltage: Voltage
 
-class Resistance(BaseQuantity): __compatible_unit__ = "ohm"
+Test(name="test", voltage=Voltage(100, "kV"))
diff --git a/src/infrasys/system.py b/src/infrasys/system.py
index 6202d51..0e1a1fc 100644
--- a/src/infrasys/system.py
+++ b/src/infrasys/system.py
@@ -1064,7 +1064,7 @@ class System:
                     values[field] = composed_value
                 elif isinstance(metadata.fields, SerializedQuantityType):
                     quantity_type = cached_types.get_type(metadata.fields)
-                    values[field] = quantity_type.from_dict(value)
+                    values[field] = quantity_type(value=value["value"], units=value["units"])
                 else:
                     msg = f"Bug: unhandled type: {field=} {value=}"
                     raise NotImplementedError(msg)
diff --git a/src/infrasys/time_series_models.py b/src/infrasys/time_series_models.py
index a374df0..3fccbe4 100644
--- a/src/infrasys/time_series_models.py
+++ b/src/infrasys/time_series_models.py
@@ -9,6 +9,7 @@ from uuid import UUID
 
 import numpy as np
 import pyarrow as pa
+import pint
 from pydantic import (
     Field,
     WithJsonSchema,
@@ -62,7 +63,7 @@ class TimeSeriesData(InfraSysBaseModelWithIdentifers, abc.ABC):
 class SingleTimeSeries(TimeSeriesData):
     """Defines a time array with a single dimension of floats."""
 
-    data: pa.Array | BaseQuantity
+    data: pa.Array | pint.Quantity
     resolution: timedelta
     initial_time: datetime
 


Test patch:
diff --git a/tests/test_base_quantity.py b/tests/test_base_quantity.py
new file mode 100644
index 0000000..8107aee
--- /dev/null
+++ b/tests/test_base_quantity.py
@@ -0,0 +1,48 @@
+from infrasys.base_quantity import ureg, BaseQuantity
+from infrasys.quantities import ActivePower, Time
+from pint.errors import DimensionalityError
+import pytest
+import numpy as np
+
+
+def test_base_quantity():
+    distance_quantity = ureg.check(None, "meter")(BaseQuantity)
+
+    unit = distance_quantity(100, "meter")
+    assert isinstance(unit, BaseQuantity)
+
+    # Check that we can not assign units that are not-related.
+    with pytest.raises(DimensionalityError):
+        _ = distance_quantity(100, "kWh")
+
+    # Check unit multiplication
+    active_power_quantity = ActivePower(100, "kW")
+    hours = Time(2, "h")
+
+    result_quantity = active_power_quantity * hours
+    assert result_quantity.check("[energy]")
+    assert result_quantity.magnitude == 200
+
+    # Check to dict
+    assert result_quantity.to_dict() == {
+        "value": result_quantity.magnitude,
+        "units": str(result_quantity.units),
+    }
+
+
+def test_base_quantity_numpy():
+    array = np.arange(0, 10)
+    measurements = ActivePower(array, "kW")
+    assert isinstance(measurements, BaseQuantity)
+    assert measurements.to_dict()["value"] == array.tolist()
+
+
+def test_unit_deserialization():
+    test_units = {
+        "value": 100,
+        "units": "kilowatt",  # The unit name should be the pint default name
+    }
+    active_power = BaseQuantity.from_dict(test_units)
+    assert isinstance(active_power, BaseQuantity)
+    assert active_power.magnitude == 100
+    assert str(active_power.units) == "kilowatt"
diff --git a/tests/test_serialization.py b/tests/test_serialization.py
index b7cb98c..5043795 100644
--- a/tests/test_serialization.py
+++ b/tests/test_serialization.py
@@ -8,7 +8,8 @@ import pytest
 from pydantic import WithJsonSchema
 from typing_extensions import Annotated
 
-from infrasys import Component, Location, SingleTimeSeries
+from infrasys import Location, SingleTimeSeries
+from infrasys.component import Component
 from infrasys.quantities import Distance, ActivePower
 from infrasys.exceptions import ISOperationNotAllowed
 from infrasys.normalization import NormalizationMax