added tests

Author: kosack

pixi.lock

@@ -14,6 +14,8 @@ pkoffee = { path = "." }
 ruff = ">=0.14.11,<0.15"
 numpy = ">=2.4.1,<3"
 matplotlib = ">=3.10.8,<4"
+pytest = ">=9.0.2,<10"
+pytest-cov = ">=7.0.0,<8"
 
 [package]
 name = "pkoffee"

pixi.toml

@@ -0,0 +1 @@
+"""Unit tests initialization."""

tests/__init__.py

@@ -0,0 +1,140 @@
+"""Unit tests for data loading and preprocessing."""
+
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from pkoffee.data import (
+    ColumnTypeError,
+    CSVReadError,
+    MissingColumnsError,
+    RequiredColumn,
+    curate,
+    data_dtype,
+    extract_arrays,
+    load_csv,
+    validate,
+)
+
+
+def test_validate() -> None:
+    """Test validate with valide DataFrame."""
+    assert validate(pd.DataFrame({"cups": [0], "productivity": [1.2]})) is None
+
+
+def test_validate_wrong_type() -> None:
+    """Test validate with incorrect required column type."""
+    with pytest.raises(ColumnTypeError):
+        validate(pd.DataFrame({"cups": [0], "productivity": ["a"]}))
+
+
+def test_validate_missing_column() -> None:
+    """Test validate with missing required column in DataFrame."""
+    with pytest.raises(MissingColumnsError):
+        validate(pd.DataFrame({f"{RequiredColumn.CUPS}": [1], "notproductivity": [1.2]}))
+
+
+@pytest.mark.parametrize(
+    ("data", "expected"),
+    [
+        (
+            pd.DataFrame({"cups": [1, np.nan, 2], "productivity": [1.2, 2.1, np.nan]}),
+            pd.DataFrame({"cups": [1.0], "productivity": [1.2]}),
+        ),
+        (
+            pd.DataFrame({"cups": [np.nan, np.nan, np.nan], "productivity": [1.2, 2.1, 3.4]}),
+            pd.DataFrame({"cups": [], "productivity": []}),
+        ),
+        (
+            pd.DataFrame({"cups": [1, 1, 4], "productivity": [1.2, 2.1, 0.5]}),
+            pd.DataFrame({"cups": [1, 1, 4], "productivity": [1.2, 2.1, 0.5]}),
+        ),
+    ],
+)
+def test_currate(data: pd.DataFrame, expected: pd.DataFrame) -> None:
+    """Test curate with different DataFrames containing nans."""
+    assert curate(data).equals(expected)
+
+
+def test_load_csv_valid_file(tmp_path: Path) -> None:
+    """Test loading valid CSV."""
+    data_file = tmp_path / "valid.csv"
+    cups = np.array([1, 2, 3], dtype=int)
+    prod = np.array([2.3, 1.2, 4.8], dtype=data_dtype)
+    np.savetxt(
+        data_file,
+        np.stack([cups, prod], axis=1),
+        fmt=["%d", "%10.4f"],
+        delimiter=",",
+        header=f"{RequiredColumn.CUPS},{RequiredColumn.PRODUCTIVITY}",
+        comments="",
+    )
+
+    data = load_csv(data_file)
+    assert RequiredColumn.CUPS in data.columns
+    assert RequiredColumn.PRODUCTIVITY in data.columns
+    assert np.isclose(data[RequiredColumn.CUPS].to_numpy(), cups).all()
+    assert np.isclose(data[RequiredColumn.PRODUCTIVITY].to_numpy(), prod).all()
+    assert data.dtypes[RequiredColumn.CUPS] == np.int64
+    assert data.dtypes[RequiredColumn.PRODUCTIVITY] == np.float64
+
+
+def test_load_csv_missing_file() -> None:
+    """Test that FileNotFoundError is raised for missing files."""
+    with pytest.raises(FileNotFoundError):
+        load_csv(Path("nonexistent_file.csv"))
+
+
+def test_load_csv_missing_columns(tmp_path: Path) -> None:
+    """Test MissingColumnsError is raised for missing required columns."""
+    wrong_col_file = tmp_path / "missing_columns.csv"
+    data = np.stack([[1], [2.3]], axis=1)
+    np.savetxt(
+        wrong_col_file,
+        data,
+        fmt=["%d", "%10.4f"],
+        delimiter=",",
+        header=f"{RequiredColumn.CUPS},wrong_column",
+        comments="",
+    )
+    with pytest.raises(MissingColumnsError, match="Missing required columns"):
+        load_csv(wrong_col_file)
+
+
+def test_load_data_with_nan_values(tmp_path: Path) -> None:
+    """Test that rows with NaN values are dropped."""
+    data_file = tmp_path / "valid_with_nan.csv"
+    with data_file.open("w") as fh:
+        fh.write(f"{RequiredColumn.CUPS},{RequiredColumn.PRODUCTIVITY}\n")
+        fh.write("1,10.5\n")
+        fh.write("2,\n")  # Missing productivity
+        fh.write("3,18.2\n")
+
+    data = load_csv(data_file)
+    expected = pd.DataFrame({RequiredColumn.CUPS: [1, 3], RequiredColumn.PRODUCTIVITY: [10.5, 18.2]})
+    assert data.equals(expected)
+
+
+def test_load_data_with_extra_values(tmp_path: Path) -> None:
+    """Test that rows with NaN values are dropped."""
+    data_file = tmp_path / "valid_with_nan.csv"
+    with data_file.open("w") as fh:
+        fh.write(f"{RequiredColumn.CUPS},{RequiredColumn.PRODUCTIVITY}\n")
+        fh.write("1,2.1\n")
+        # try to read the file while it is open for write
+        with pytest.raises(CSVReadError):
+            load_csv(data_file)
+
+
+def test_extract_arrays() -> None:
+    """Test extracting numpy arrays from DataFrame."""
+    cups_ref = np.array([1, 2, 3], dtype=int)
+    productivity_ref = np.array([10.5, 15.3, 18.2], dtype=np.float64)
+    data = pd.DataFrame({RequiredColumn.CUPS: cups_ref, RequiredColumn.PRODUCTIVITY: productivity_ref})
+
+    cups, productivity = extract_arrays(data)
+
+    assert np.allclose(cups_ref, cups)
+    assert np.allclose(productivity_ref, productivity)

tests/test_data.py

@@ -0,0 +1,236 @@
+"""Unit tests for model fitting."""
+
+import numpy as np
+import pytest
+
+from pkoffee.data import AnyShapeDataDtypeArray, neg_inf, pos_inf
+from pkoffee.data import data_dtype as dt
+from pkoffee.fit_model import (
+    FunctionIdNotFoundInMappingError,
+    FunctionNotFoundInMappingError,
+    Model,
+    ModelParsingError,
+    fit_model,
+)
+from pkoffee.parametric_function import ParametersBounds
+
+
+class Linear:
+    """Linear function."""
+
+    def __call__(self, x: AnyShapeDataDtypeArray, a: dt, b: dt) -> AnyShapeDataDtypeArray:
+        """Evaluate y = a * x + b."""
+        return a * x + b  # pyright: ignore[reportReturnType] return dtype is data_dtype alright.
+
+    @staticmethod
+    def param_guess() -> dict[str, dt]:
+        """Guess parameters for tests: a: 2.0, b: 1.0."""
+        return {"a": dt(2.0), "b": dt(1.0)}
+
+    @staticmethod
+    def param_bounds() -> ParametersBounds:
+        """Parameter bounds in [-inf, +inf]."""
+        return ParametersBounds(min={"a": neg_inf, "b": neg_inf}, max={"a": pos_inf, "b": pos_inf})
+
+
+class LinearSQRT:
+    """Passthrough model with square root of coefficient parameter."""
+
+    def __call__(self, x: AnyShapeDataDtypeArray, a: dt) -> AnyShapeDataDtypeArray:
+        """Evaluate y = sqrt(a) * x."""
+        return np.sqrt(a) * x
+
+    @staticmethod
+    def param_guess() -> dict[str, dt]:
+        """Guess parameter: a: 1.0."""
+        return {"a": dt(1.0)}
+
+    @staticmethod
+    def param_bounds() -> ParametersBounds:
+        """Parameter bounds: "a" should be positive."""
+        return ParametersBounds(min={"a": dt(5e-7)}, max={"a": pos_inf})
+
+
+class PassThrough:
+    """PassThrough empty model for testing."""
+
+    def __call__(self, x: AnyShapeDataDtypeArray) -> AnyShapeDataDtypeArray:
+        """Passthrough."""
+        return x
+
+    @staticmethod
+    def param_guess() -> dict[str, dt]:
+        """No parameter to guess."""
+        return {}
+
+    @staticmethod
+    def param_bounds() -> ParametersBounds:
+        """No parameter bounds."""
+        return ParametersBounds(min={}, max={})
+
+
+def test_model_result_predict() -> None:
+    """Test ModelResult prediction method."""
+    model_result = Model(
+        name="Linear",
+        function=Linear(),
+        params=Linear.param_guess(),
+        bounds=Linear.param_bounds(),
+        r_squared=dt(0.95),
+    )
+
+    x_new = np.array([0.0, 1.0, 2.0], dtype=dt)
+    predictions = model_result.predict(x_new)
+
+    expected = np.array([1.0, 3.0, 5.0])
+    np.testing.assert_allclose(predictions, expected)
+
+
+def test_fit_model_success() -> None:
+    """Test successful model fitting."""
+    # Create simple linear data
+    x = np.linspace(0, 10, 50)
+    rng = np.random.default_rng(1337)
+    y = 2.0 * x + 1.0 + rng.normal(0, 0.1, 50)
+
+    lin_model = Model(
+        name="Test Linear",
+        function=Linear(),
+        params={"a": dt(1.0), "b": dt(0.0)},
+        bounds=ParametersBounds(min=dict.fromkeys(["a", "b"], neg_inf), max=dict.fromkeys(["a", "b"], pos_inf)),
+    )
+
+    result, _ = fit_model(x, y, lin_model)
+
+    assert result is not None
+    assert result.name == "Test Linear"
+    assert result.r_squared > 0.9  # Should fit well  # noqa: PLR2004
+    assert len(result.params) == 2  # noqa: PLR2004
+    assert result.params["a"] == pytest.approx(2.0, abs=0.2)
+
+
+def test_fit_model_failure() -> None:
+    """Test that fit_model raises error on failure."""
+    x = np.array([1.0, 2.0, 3.0])
+    y = np.array([1.0, 2.0, 3.0])
+
+    # Create a config that will fail
+    config = Model(
+        name="Bad Model",
+        function=LinearSQRT(),
+        params={"a": dt(-10.0)},
+        bounds=ParametersBounds(min={"a": dt(-100.0)}, max={"a": dt(-1.0)}),
+    )
+
+    # Verify value error is raised, ignore RuntimeWarning from sqrt(0)
+    with pytest.raises(ValueError), pytest.warns(RuntimeWarning):  # noqa: PT011 error raised by scipy
+        _ = fit_model(x, y, config)
+
+
+def test_model_result_repr() -> None:
+    """Test ModelResult string representation."""
+    model = Model(
+        name="TestModel",
+        function=PassThrough(),
+        params={},
+        r_squared=dt(0.8765),
+        bounds=ParametersBounds({}, {}),
+    )
+
+    repr_str = repr(model)
+
+    assert repr_str == f"ModelFit(name='{model.name}', R²={model.r_squared:.3f})"
+
+
+def test_model_sort() -> None:
+    """Test Model.sort."""
+    r_squared_min = dt(0.1)
+    r_squared_max = dt(0.3)
+    model_list = [
+        Model(
+            name="a",
+            function=PassThrough(),
+            params={},
+            r_squared=r_squared_max,
+            bounds=ParametersBounds({}, {}),
+        ),
+        Model(
+            name="b",
+            function=PassThrough(),
+            params={},
+            r_squared=r_squared_min,
+            bounds=ParametersBounds({}, {}),
+        ),
+    ]
+    Model.sort(model_list)
+    assert model_list[0].r_squared == r_squared_max
+    assert model_list[1].r_squared == r_squared_min
+
+
+def test_model_to_dict() -> None:
+    """Test model conversion to dictionary."""
+    linear_model = Model(
+        name="test_linear",
+        function=Linear(),
+        params={"a": dt(1.0), "b": dt(0.5)},
+        bounds=ParametersBounds(min={"a": dt(-1.0), "b": dt(-5.0)}, max={"a": dt(1.0), "b": dt(5.0)}),
+        r_squared=dt(0.2),
+    )
+    assert linear_model.to_dict({Linear: "Linear"}) == {
+        "name": "test_linear",
+        "function": "Linear",
+        "params": {"a": 1.0, "b": 0.5},
+        "bounds": {"min": {"a": -1.0, "b": -5.0}, "max": {"a": 1.0, "b": 5.0}},
+        "r_squared": 0.20000000298023224,
+    }
+
+
+def test_model_to_dict_missing_mapping() -> None:
+    """Test Error raising when function mapping is missing during dict conversion."""
+    model = Model(
+        name="Passthrough",
+        function=PassThrough(),
+        params={},
+        bounds=ParametersBounds(min={}, max={}),
+        r_squared=neg_inf,
+    )
+    with pytest.raises(FunctionNotFoundInMappingError):
+        model.to_dict({})
+
+
+def test_model_from_dict() -> None:
+    """Test model creation from dictionary."""
+    d = {
+        "name": "test_linear",
+        "function": "Linear",
+        "params": {"a": 1.0, "b": 0.5},
+        "bounds": {"min": {"a": -1.0, "b": -5.0}, "max": {"a": 1.0, "b": 5.0}},
+        "r_squared": 0.20000000298023224,
+    }
+    linear_model = Model.from_dict(d, {"Linear": Linear})
+    assert linear_model.name == d["name"]
+    assert isinstance(linear_model.function, Linear)
+    assert linear_model.params == {p: dt(v) for p, v in d["params"].items()}
+    assert linear_model.bounds == ParametersBounds(
+        min={p: dt(v) for p, v in d["bounds"]["min"].items()}, max={p: dt(v) for p, v in d["bounds"]["max"].items()}
+    )
+    assert linear_model.r_squared == d["r_squared"]
+
+
+def test_model_from_dict_missing_mapping() -> None:
+    """Test model creation from dictionary error in case of missing mapping."""
+    d = {
+        "name": "passthrough",
+        "function": "PassThrough",
+        "params": {},
+        "bounds": {"min": {}, "max": {}},
+        "r_squared": neg_inf,
+    }
+    with pytest.raises(FunctionIdNotFoundInMappingError):
+        Model.from_dict(d, {})
+
+
+def test_model_from_dict_bad_dict() -> None:
+    """Test model creation error when the dictionary doesn't have the required content."""
+    with pytest.raises(ModelParsingError):
+        Model.from_dict({}, {})