Fix: Update jitter initialization logic and improve test assertions for add_jitter function

Author: SvenKlaassen

doubleml/did/utils/_plot.py

@@ -25,7 +25,10 @@ def add_jitter(data, x_col, is_datetime=None, jitter_value=None):
         is_datetime = pd.api.types.is_datetime64_any_dtype(data[x_col])
 
     # Initialize jittered_x with original values
-    data["jittered_x"] = data[x_col]
+    if is_datetime:
+        data["jittered_x"] = data[x_col]
+    else:
+        data["jittered_x"] = data[x_col].astype(float)
 
     for x_val in data[x_col].unique():
         mask = data[x_col] == x_val

doubleml/did/utils/tests/test_add_jitter.py

@@ -1,5 +1,6 @@
 from datetime import datetime, timedelta
 
+import numpy as np
 import pandas as pd
 import pytest
 
@@ -41,8 +42,7 @@ def test_add_jitter_numeric_no_duplicates(numeric_df_no_duplicates):
     """Test that no jitter is added when there are no duplicates."""
     result = add_jitter(numeric_df_no_duplicates, "x")
     # No jitter should be added when there are no duplicates
-    pd.testing.assert_series_equal(result["jittered_x"], result["x"], check_names=False)
-
+    np.testing.assert_allclose(result["jittered_x"], result["x"])
 
 @pytest.mark.ci
 def test_add_jitter_numeric_with_duplicates(numeric_df_with_duplicates):
@@ -121,7 +121,7 @@ def test_add_jitter_explicit_datetime_flag():
     df = pd.DataFrame({"x": ["2023-01-01", "2023-01-01", "2023-01-02"], "y": [10, 15, 20]})
 
     # Without specifying is_datetime, it would treat as strings
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         _ = add_jitter(df, "x")
 
     # With is_datetime=True, it should convert and jitter as datetimes