From e10f984f9caa2cad735b5e019717b54bd63b6798 Mon Sep 17 00:00:00 2001
From: Alexandra Udaltsova <43303448+AUdaltsova@users.noreply.github.com>
Date: Tue, 6 Aug 2024 13:43:07 +0100
Subject: [PATCH] Add 2D array test test_datetime_features.py

---
 .../numpy_batch/test_datetime_features.py     | 37 +++++++++----------
 1 file changed, 18 insertions(+), 19 deletions(-)

diff --git a/tests/transform/numpy_batch/test_datetime_features.py b/tests/transform/numpy_batch/test_datetime_features.py
index 54df6b5c..c487a25c 100644
--- a/tests/transform/numpy_batch/test_datetime_features.py
+++ b/tests/transform/numpy_batch/test_datetime_features.py
@@ -1,28 +1,21 @@
 import numpy as np
-from datetime import datetime
-
-from ocf_datapipes.transform.numpy_batch import AddTrigonometricDateTime
 
 from ocf_datapipes.transform.numpy_batch.datetime_features import _get_date_time_in_pi
 
 
 def test_get_date_time_in_pi():
-    times = [
-        "2020-01-01T00:00:01",
-        "2020-04-01T06:00:00",
-        "2020-07-01T12:00:00",
-        "2020-09-30T18:00:00",
-        "2020-12-31T23:59:59",
-        "2021-01-01T00:00:01",
-        "2021-04-02T06:00:00",
-        "2021-07-02T12:00:00",
-        "2021-10-01T18:00:00",
-        "2021-12-31T23:59:59",
-    ]
-
-    expected_times_in_pi = [0, 0.5 * np.pi, np.pi, 1.5 * np.pi, 2 * np.pi] * 2
-
-    times = np.array([datetime.fromisoformat(time) for time in times], dtype="datetime64[s]")
+    times = np.array([
+              "2020-01-01T00:00:00", "2020-04-01T06:00:00",
+              "2020-07-01T12:00:00", "2020-09-30T18:00:00",
+              "2020-12-31T23:59:59",
+              "2021-01-01T00:00:00", "2021-04-02T06:00:00",
+              "2021-07-02T12:00:00", "2021-10-01T18:00:00",
+              "2021-12-31T23:59:59"
+             ]).reshape((2, 5))
+
+    expected_times_in_pi = np.array([0, 0.5*np.pi, np.pi, 1.5*np.pi, 2*np.pi] * 2).reshape((2,5))
+
+    times = times.astype("datetime64[s]")
 
     date_in_pi, time_in_pi = _get_date_time_in_pi(times)
 
@@ -34,3 +27,9 @@ def test_get_date_time_in_pi():
     assert np.isclose(np.sin(time_in_pi), np.sin(expected_times_in_pi), atol=7.3e-05).all()
     assert np.isclose(np.cos(date_in_pi), np.cos(expected_times_in_pi), atol=0.02182).all()
     assert np.isclose(np.sin(date_in_pi), np.sin(expected_times_in_pi), atol=0.02182).all()
+
+    # 1D array test
+    assert np.isclose(np.cos(time_in_pi[0]), np.cos(expected_times_in_pi[0]), atol=7.3e-05).all()
+    assert np.isclose(np.sin(time_in_pi[0]), np.sin(expected_times_in_pi[0]), atol=7.3e-05).all()
+    assert np.isclose(np.cos(date_in_pi[0]), np.cos(expected_times_in_pi[0]), atol=0.02182).all()
+    assert np.isclose(np.sin(date_in_pi[0]), np.sin(expected_times_in_pi[0]), atol=0.02182).all()