Repository: NREL/infrasys

Issue:
Time series drop units of `pint.Quantity`s that do not subclass `infrasys.BaseQuantity`
When an `infrasys.BaseQuantity` is added to a `SingleTimeSeries`, its units are preserved, but when a non-`BaseQuantity` `pint.Quantity` is added, its units are dropped. I believe this is the culprit:
https://github.com/NREL/infrasys/blob/169accf25e5be536ef86dd2c078c5d658fcdf3ec/src/infrasys/time_series_models.py#L344
.

Hints from issue comments:

Solution patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 3786802..2b9203e 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -9,6 +9,6 @@ repos:
     # Run the formatter.
     - id: ruff-format
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v1.8.0
+  rev: v1.13.0
   hooks:
   - id: mypy
diff --git a/pyproject.toml b/pyproject.toml
index 84d08ba..3536377 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,7 +27,10 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
 dependencies = [
-    "numpy~=1.26.4",
+    # This is what we want, but are currently be limited by pyarrow.
+    # Leave unbounded until pyarrow upgrades numpy.
+    #"numpy >= 2, < 3",
+    "numpy",
     "pyarrow~=15.0.2",
     "pint~=0.23",
     "loguru~=0.7.2",
@@ -37,7 +40,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = [
     "furo",
-    "mypy",
+    "mypy >=1.13, < 2",
     "myst_parser",
     "pre-commit",
     "pyarrow-stubs",
@@ -56,6 +59,13 @@ Documentation = "https://github.com/NREL/infrasys#readme"
 Issues = "https://github.com/NREL/infrasys/issues"
 Source = "https://github.com/NREL/infrasys"
 
+[tool.mypy]
+check_untyped_defs = true
+files = [
+  "src",
+  "tests",
+]
+
 [tool.pytest.ini_options]
 pythonpath = "src"
 minversion = "6.0"
diff --git a/src/infrasys/arrow_storage.py b/src/infrasys/arrow_storage.py
index c09c4de..29e726b 100644
--- a/src/infrasys/arrow_storage.py
+++ b/src/infrasys/arrow_storage.py
@@ -8,6 +8,7 @@
 from typing import Any, Optional
 from uuid import UUID
 
+import numpy as np
 import pyarrow as pa
 import pint
 from loguru import logger
@@ -112,13 +113,17 @@ def _get_single_time_series(
         index, length = metadata.get_range(start_time=start_time, length=length)
         data = base_ts[metadata.variable_name][index : index + length]
         if metadata.quantity_metadata is not None:
-            data = metadata.quantity_metadata.quantity_type(data, metadata.quantity_metadata.units)
+            np_array = metadata.quantity_metadata.quantity_type(
+                data, metadata.quantity_metadata.units
+            )
+        else:
+            np_array = np.array(data)
         return SingleTimeSeries(
             uuid=metadata.time_series_uuid,
             variable_name=metadata.variable_name,
             resolution=metadata.resolution,
             initial_time=start_time or metadata.initial_time,
-            data=data,
+            data=np_array,
             normalization=metadata.normalization,
         )
 
@@ -126,10 +131,10 @@ def _convert_to_record_batch(
         self, time_series: SingleTimeSeries, variable_name: str
     ) -> pa.RecordBatch:
         """Create record batch to save array to disk."""
-        pa_array = time_series.data
-        if not isinstance(pa_array, pa.Array) and isinstance(pa_array, pint.Quantity):
-            pa_array = pa.array(pa_array.magnitude)
-        assert isinstance(pa_array, pa.Array)
+        if isinstance(time_series.data, pint.Quantity):
+            pa_array = pa.array(time_series.data.magnitude)
+        else:
+            pa_array = pa.array(time_series.data)
         schema = pa.schema([pa.field(variable_name, pa_array.type)])
         return pa.record_batch([pa_array], schema=schema)
 
diff --git a/src/infrasys/function_data.py b/src/infrasys/function_data.py
index 4a5df0d..01898aa 100644
--- a/src/infrasys/function_data.py
+++ b/src/infrasys/function_data.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 import pint
+from numpy.typing import NDArray
 from pydantic import Field, model_validator
 from pydantic.functional_validators import AfterValidator
 from typing_extensions import Annotated
@@ -186,7 +187,7 @@ def validate_piecewise_xy(self):
         return self
 
 
-def get_x_lengths(x_coords: List[float]) -> List[float]:
+def get_x_lengths(x_coords: List[float]) -> NDArray[np.float64]:
     return np.subtract(x_coords[1:], x_coords[:-1])
 
 
diff --git a/src/infrasys/normalization.py b/src/infrasys/normalization.py
index 742d9ef..da6c166 100644
--- a/src/infrasys/normalization.py
+++ b/src/infrasys/normalization.py
@@ -5,6 +5,7 @@
 from typing import Literal, Optional, Annotated, Union
 
 import numpy as np
+from numpy.typing import NDArray
 from pydantic import Field
 
 from infrasys.models import InfraSysBaseModel
@@ -29,7 +30,7 @@ class NormalizationMax(NormalizationBase):
     max_value: Optional[float] = None
     normalization_type: Literal[NormalizationType.MAX] = NormalizationType.MAX
 
-    def normalize_array(self, data: np.ndarray) -> np.ndarray:
+    def normalize_array(self, data: NDArray) -> NDArray:
         self.max_value = np.max(data)
         return data / self.max_value
 
@@ -40,7 +41,7 @@ class NormalizationByValue(NormalizationBase):
     value: float
     normalization_type: Literal[NormalizationType.BY_VALUE] = NormalizationType.BY_VALUE
 
-    def normalize_array(self, data: np.ndarray) -> np.ndarray:
+    def normalize_array(self, data: NDArray) -> NDArray:
         return data / self.value
 
 
diff --git a/src/infrasys/time_series_models.py b/src/infrasys/time_series_models.py
index 0e6c2e4..c057397 100644
--- a/src/infrasys/time_series_models.py
+++ b/src/infrasys/time_series_models.py
@@ -17,8 +17,8 @@
 from uuid import UUID
 
 import numpy as np
-import pyarrow as pa
 import pint
+from numpy.typing import NDArray
 from pydantic import (
     Field,
     WithJsonSchema,
@@ -29,7 +29,6 @@
 )
 from typing_extensions import Annotated
 
-from infrasys.base_quantity import BaseQuantity
 from infrasys.exceptions import (
     ISConflictingArguments,
     InconsistentTimeseriesAggregation,
@@ -42,7 +41,7 @@
 VALUE_COLUMN = "value"
 
 
-ISArray: TypeAlias = Sequence | pa.Array | np.ndarray | pint.Quantity
+ISArray: TypeAlias = Sequence | NDArray | pint.Quantity
 
 
 class TimeSeriesStorageType(str, Enum):
@@ -74,7 +73,7 @@ def get_time_series_metadata_type() -> Type:
 class SingleTimeSeries(TimeSeriesData):
     """Defines a time array with a single dimension of floats."""
 
-    data: pa.Array | pint.Quantity
+    data: NDArray | pint.Quantity
     resolution: timedelta
     initial_time: datetime
 
@@ -83,21 +82,36 @@ def length(self) -> int:
         """Return the length of the data."""
         return len(self.data)
 
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, SingleTimeSeries):
+            raise NotImplementedError
+        is_equal = True
+        for field in self.model_fields_set:
+            if field == "data":
+                if not (self.data == other.data).all():
+                    is_equal = False
+                    break
+            else:
+                if not getattr(self, field) == getattr(other, field):
+                    is_equal = False
+                    break
+        return is_equal
+
     @field_validator("data", mode="before")
     @classmethod
-    def check_data(
-        cls, data
-    ) -> pa.Array | pa.ChunkedArray | pint.Quantity:  # Standarize what object we receive.
+    def check_data(cls, data) -> NDArray | pint.Quantity:  # Standarize what object we receive.
         """Check time series data."""
         if len(data) < 2:
             msg = f"SingleTimeSeries length must be at least 2: {len(data)}"
             raise ValueError(msg)
 
         if isinstance(data, pint.Quantity):
+            if not isinstance(data.magnitude, np.ndarray):
+                return type(data)(np.array(data.magnitude), units=data.units)
             return data
 
-        if not isinstance(data, pa.Array):
-            return pa.array(data)
+        if not isinstance(data, np.ndarray):
+            return np.array(data)
 
         return data
 
@@ -136,7 +150,7 @@ def aggregate(cls, ts_data: list[Self]) -> Self:
             raise InconsistentTimeseriesAggregation(msg)
 
         # Aggregate data
-        is_quantity = issubclass(next(iter(unique_props["data_type"])), BaseQuantity)
+        is_quantity = issubclass(next(iter(unique_props["data_type"])), pint.Quantity)
         magnitude_type = (
             type(ts_data[0].data.magnitude)
             if is_quantity
@@ -144,12 +158,9 @@ def aggregate(cls, ts_data: list[Self]) -> Self:
         )
 
         # Aggregate data based on magnitude type
-        if issubclass(magnitude_type, pa.Array):
+        if issubclass(magnitude_type, np.ndarray):
             new_data = sum(
-                [
-                    data.data.to_numpy() * (data.data.units if is_quantity else 1)
-                    for data in ts_data
-                ]
+                [data.data * (data.data.units if is_quantity else 1) for data in ts_data]
             )
         elif issubclass(magnitude_type, np.ndarray):
             new_data = sum([data.data for data in ts_data])
@@ -274,7 +285,7 @@ class SingleTimeSeriesScalingFactor(SingleTimeSeries):
 
 
 class QuantityMetadata(InfraSysBaseModel):
-    """Contains the metadata needed to de-serialize time series stored within a BaseQuantity."""
+    """Contains the metadata needed to de-serialize time series stored within a pint.Quantity."""
 
     module: str
     quantity_type: Annotated[Type, WithJsonSchema({"type": "string"})]
@@ -341,7 +352,7 @@ def from_data(cls, time_series: SingleTimeSeries, **user_attributes) -> Any:
                 quantity_type=type(time_series.data),
                 units=str(time_series.data.units),
             )
-            if isinstance(time_series.data, BaseQuantity)
+            if isinstance(time_series.data, pint.Quantity)
             else None
         )
         return cls(
diff --git a/src/infrasys/value_curves.py b/src/infrasys/value_curves.py
index 7fdcad6..7a8a7d8 100644
--- a/src/infrasys/value_curves.py
+++ b/src/infrasys/value_curves.py
@@ -7,6 +7,7 @@
     QuadraticFunctionData,
     PiecewiseLinearData,
     PiecewiseStepData,
+    XYCoords,
     running_sum,
 )
 from pydantic import Field
@@ -117,7 +118,9 @@ def to_input_output(self) -> InputOutputCurve:
                 points = running_sum(self.function_data)
 
                 return InputOutputCurve(
-                    function_data=PiecewiseLinearData(points=[(p.x, p.y + c) for p in points]),
+                    function_data=PiecewiseLinearData(
+                        points=[XYCoords(p.x, p.y + c) for p in points]
+                    ),
                     input_at_zero=self.input_at_zero,
                 )
 
@@ -188,7 +191,10 @@ def to_input_output(self) -> InputOutputCurve:
                 else:
                     return InputOutputCurve(
                         function_data=QuadraticFunctionData(
-                            quadratic_term=p, proportional_term=m, constant_term=c
+                            # issue 53
+                            quadratic_term=p,
+                            proportional_term=m,
+                            constant_term=c,  # type: ignore
                         ),
                         input_at_zero=self.input_at_zero,
                     )
@@ -203,7 +209,9 @@ def to_input_output(self) -> InputOutputCurve:
                 ys.insert(0, c)
 
                 return InputOutputCurve(
-                    function_data=PiecewiseLinearData(points=list(zip(xs, ys))),
+                    function_data=PiecewiseLinearData(
+                        points=[XYCoords(x, y) for x, y in zip(xs, ys)]
+                    ),
                     input_at_zero=self.input_at_zero,
                 )
             case _:


Test patch:
diff --git a/tests/test_arrow_storage.py b/tests/test_arrow_storage.py
index 1529dec..6a7d447 100644
--- a/tests/test_arrow_storage.py
+++ b/tests/test_arrow_storage.py
@@ -3,7 +3,7 @@
 from datetime import datetime, timedelta
 from pathlib import Path
 
-import pyarrow as pa
+import numpy as np
 from loguru import logger
 
 from infrasys.arrow_storage import ArrowTimeSeriesStorage
@@ -88,5 +88,7 @@ def test_read_deserialize_time_series(tmp_path):
     assert isinstance(deserialize_ts, SingleTimeSeries)
     assert deserialize_ts.resolution == ts.resolution
     assert deserialize_ts.initial_time == ts.initial_time
-    assert isinstance(deserialize_ts.data, pa.Array)
-    assert deserialize_ts.data[-1].as_py() == ts.length - 1
+    assert isinstance(deserialize_ts.data, np.ndarray)
+    length = ts.length
+    assert isinstance(length, int)
+    assert np.array_equal(deserialize_ts.data, np.array(range(length)))
diff --git a/tests/test_base_quantity.py b/tests/test_base_quantity.py
index e133efa..4190901 100644
--- a/tests/test_base_quantity.py
+++ b/tests/test_base_quantity.py
@@ -76,7 +76,7 @@ class _(BaseQuantity):
     with pytest.raises(ValidationError):
         BaseQuantityComponent(name="test", voltage=Voltage(test_magnitude, "meter"))
 
-    test_component = BaseQuantityComponent(name="test", voltage=[0, 1])
+    test_component = BaseQuantityComponent(name="test", voltage=Voltage([0, 1], units="volt"))
     assert type(test_component.voltage) is Voltage
     assert test_component.voltage.magnitude.tolist() == [0, 1]
     assert test_component.voltage.units == test_unit
@@ -91,7 +91,7 @@ def test_different_validate(input_unit):
 
 
 def test_custom_serialization():
-    component = BaseQuantityComponent(name="test", voltage=10.0)
+    component = BaseQuantityComponent(name="test", voltage=Voltage(10.0, units="volt"))
 
     model_dump = component.model_dump(mode="json")
 
diff --git a/tests/test_serialization.py b/tests/test_serialization.py
index d696645..9841d54 100644
--- a/tests/test_serialization.py
+++ b/tests/test_serialization.py
@@ -4,7 +4,6 @@
 from datetime import datetime, timedelta
 
 import numpy as np
-import pyarrow as pa
 import pytest
 from pydantic import WithJsonSchema
 from typing_extensions import Annotated
@@ -98,7 +97,7 @@ def test_serialize_time_series(tmp_path, time_series_in_memory):
         system2.remove_time_series(gen1b, variable_name=variable_name)
 
     ts2 = system.get_time_series(gen1b, variable_name=variable_name)
-    assert ts2.data == ts.data
+    assert np.array_equal(ts2.data, ts.data)
 
     system3 = SimpleSystem.from_json(filename, time_series_read_only=False)
     assert system3.get_time_series_directory() != SimpleSystem._make_time_series_directory(
@@ -121,6 +120,7 @@ def test_serialize_quantity(tmp_path, distance):
     system = SimpleSystem()
     gen = SimpleGenerator.example()
     component = ComponentWithPintQuantity(name="test", distance=distance)
+    assert gen.bus.coordinates is not None
     system.add_components(gen.bus.coordinates, gen.bus, gen, component)
     sys_file = tmp_path / "system.json"
     system.to_json(sys_file)
@@ -156,9 +156,8 @@ def test_with_time_series_quantity(tmp_path):
     assert ts.length == length
     assert ts.resolution == resolution
     assert ts.initial_time == initial_time
-    assert isinstance(ts2.data.magnitude, pa.Array)
-    assert ts2.data[-1].as_py() == length - 1
-    assert ts2.data.magnitude == pa.array(range(length))
+    assert isinstance(ts2.data.magnitude, np.ndarray)
+    assert np.array_equal(ts2.data.magnitude, np.array(range(length)))
 
 
 @pytest.mark.parametrize("in_memory", [True, False])
diff --git a/tests/test_single_time_series.py b/tests/test_single_time_series.py
index be65210..680c0a6 100644
--- a/tests/test_single_time_series.py
+++ b/tests/test_single_time_series.py
@@ -1,8 +1,9 @@
 """Test related to arrow storage module."""
 from datetime import datetime, timedelta
 
+import pint
 import pytest
-import pyarrow as pa
+import numpy as np
 
 from infrasys.normalization import NormalizationMax
 from infrasys.quantities import ActivePower
@@ -21,8 +22,8 @@ def test_single_time_series_attributes():
     assert ts.length == length
     assert ts.resolution == resolution
     assert ts.initial_time == start
-    assert isinstance(ts.data, pa.Array)
-    assert ts.data[-1].as_py() == length - 1
+    assert isinstance(ts.data, np.ndarray)
+    assert ts.data[-1] == length - 1
 
 
 def test_from_array_construction():
@@ -37,8 +38,8 @@ def test_from_array_construction():
     assert ts.length == length
     assert ts.resolution == resolution
     assert ts.initial_time == start
-    assert isinstance(ts.data, pa.Array)
-    assert ts.data[-1].as_py() == length - 1
+    assert isinstance(ts.data, np.ndarray)
+    assert ts.data[-1] == length - 1
 
 
 def test_invalid_sequence_length():
@@ -65,8 +66,8 @@ def test_from_time_array_constructor():
     assert ts.length == length
     assert ts.resolution == resolution
     assert ts.initial_time == initial_time
-    assert isinstance(ts.data, pa.Array)
-    assert ts.data[-1].as_py() == length - 1
+    assert isinstance(ts.data, np.ndarray)
+    assert ts.data[-1] == length - 1
 
 
 def test_with_quantity():
@@ -99,7 +100,7 @@ def test_normalization():
     assert isinstance(ts, SingleTimeSeries)
     assert ts.length == len(data)
     for i, val in enumerate(ts.data):
-        assert val.as_py() == data[i] / max_val
+        assert val == data[i] / max_val
 
 
 def test_normal_array_aggregate():
@@ -113,7 +114,7 @@ def test_normal_array_aggregate():
     )
     ts_agg = SingleTimeSeries.aggregate([ts1, ts2])
     assert isinstance(ts_agg, SingleTimeSeries)
-    assert list([el.as_py() for el in ts_agg.data]) == [2.2, 4.4, 6.6, 9, 11]
+    assert list([el for el in ts_agg.data]) == [2.2, 4.4, 6.6, 9, 11]
 
 
 def test_pint_array_aggregate():
@@ -127,4 +128,5 @@ def test_pint_array_aggregate():
     )
     ts_agg = SingleTimeSeries.aggregate([ts1, ts2])
     assert isinstance(ts_agg, SingleTimeSeries)
+    assert isinstance(ts_agg.data, pint.Quantity)
     assert list(ts_agg.data.magnitude) == [2.2, 4.4, 6.6, 9, 11]
diff --git a/tests/test_system.py b/tests/test_system.py
index 8f2c2e2..aa8e557 100644
--- a/tests/test_system.py
+++ b/tests/test_system.py
@@ -31,7 +31,7 @@ def test_system():
     gen = SimpleGenerator(name="test-gen", active_power=1.0, rating=1.0, bus=bus, available=True)
     subsystem = SimpleSubsystem(name="test-subsystem", generators=[gen])
     system.add_components(geo, bus, gen, subsystem)
-    assert system.add_components() is None
+    assert system.add_components() is None  # type: ignore
 
     gen2 = system.get_component(SimpleGenerator, "test-gen")
     assert gen2 is gen
@@ -214,7 +214,7 @@ def test_time_series_attach_from_array():
     system.add_time_series(ts, gen1, gen2)
     assert system.has_time_series(gen1, variable_name=variable_name)
     assert system.has_time_series(gen2, variable_name=variable_name)
-    assert system.get_time_series(gen1, variable_name=variable_name).data == ts.data
+    assert np.array_equal(system.get_time_series(gen1, variable_name=variable_name).data, ts.data)
 
 
 def test_time_series():
@@ -232,7 +232,7 @@ def test_time_series():
     ts = SingleTimeSeries.from_time_array(data, variable_name, time_array)
     with pytest.raises(ValueError, match="The first argument must"):
         # Test a common mistake.
-        system.add_time_series(gen1, ts)
+        system.add_time_series(gen1, ts)  # type: ignore
 
     system.add_time_series(ts, gen1, gen2)
     assert system.has_time_series(gen1, variable_name=variable_name)
@@ -265,7 +265,7 @@ def test_time_series_retrieval(in_memory, use_quantity, sql_json):
         data = (
             [ActivePower(np.random.rand(length), "watts") for _ in range(4)]
             if use_quantity
-            else [np.random.rand(length) for _ in range(4)]
+            else [np.random.rand(length) for _ in range(4)]  # type: ignore
         )
         variable_name = "active_power"
         ts1 = SingleTimeSeries.from_time_array(data[0], variable_name, time_array)
@@ -435,22 +435,22 @@ def test_time_series_slices(in_memory):
     first_timestamp = start
     second_timestamp = start + resolution
     last_timestamp = start + (length - 1) * resolution
-    assert len(system.time_series.get(gen, variable_name=variable_name).data) == length
-    assert len(system.time_series.get(gen, variable_name=variable_name, length=10).data) == 10
+    ts_tmp = system.time_series.get(gen, variable_name=variable_name)
+    assert isinstance(ts_tmp, SingleTimeSeries)
+    assert len(ts_tmp.data) == length
+    ts_tmp = system.time_series.get(gen, variable_name=variable_name, length=10)
+    assert isinstance(ts_tmp, SingleTimeSeries)
+    assert len(ts_tmp.data) == 10
     ts2 = system.time_series.get(
         gen, variable_name=variable_name, start_time=second_timestamp, length=5
     )
+    assert isinstance(ts2, SingleTimeSeries)
     assert len(ts2.data) == 5
     assert ts2.data.tolist() == data[1:6]
 
-    assert (
-        len(
-            system.time_series.get(
-                gen, variable_name=variable_name, start_time=second_timestamp
-            ).data
-        )
-        == len(data) - 1
-    )
+    ts_tmp = system.time_series.get(gen, variable_name=variable_name, start_time=second_timestamp)
+    assert isinstance(ts_tmp, SingleTimeSeries)
+    assert len(ts_tmp.data) == len(data) - 1
 
     with pytest.raises(ISConflictingArguments, match="is less than"):
         system.time_series.get(
diff --git a/tests/test_values_curves.py b/tests/test_values_curves.py
index 651a4e6..28b00cc 100644
--- a/tests/test_values_curves.py
+++ b/tests/test_values_curves.py
@@ -86,6 +86,7 @@ def test_average_rate_conversion():
     assert isinstance(new_curve.function_data, QuadraticFunctionData)
     assert new_curve.function_data.quadratic_term == 1.0
 
+    assert isinstance(curve.function_data, LinearFunctionData)
     curve.function_data.proportional_term = 0.0
     new_curve = curve.to_input_output()
     assert isinstance(new_curve, InputOutputCurve)
@@ -117,6 +118,7 @@ def test_incremental_conversion():
     assert isinstance(new_curve.function_data, QuadraticFunctionData)
     assert new_curve.function_data.quadratic_term == 0.5
 
+    assert isinstance(curve.function_data, LinearFunctionData)
     curve.function_data.proportional_term = 0.0
     new_curve = curve.to_input_output()
     assert isinstance(new_curve, InputOutputCurve)