Speed up tests

Author: FBumann

tests/test_resample_equivalence.py

@@ -23,7 +23,7 @@ def naive_dataset_resample(dataset: xr.Dataset, freq: str, method: str) -> xr.Da
     return getattr(dataset.resample(time=freq), method)()
 
 
-def create_dataset_with_mixed_dimensions(n_timesteps=100):
+def create_dataset_with_mixed_dimensions(n_timesteps=48):
     """
     Create a dataset with variables having different dimension structures.
 
@@ -37,7 +37,7 @@ def create_dataset_with_mixed_dimensions(n_timesteps=100):
     ds = xr.Dataset(
         coords={
             'time': timesteps,
-            'component': ['comp1', 'comp2', 'comp3'],
+            'component': ['comp1', 'comp2'],
             'bus': ['bus1', 'bus2'],
             'scenario': ['base', 'alt'],
         }
@@ -51,7 +51,7 @@ def create_dataset_with_mixed_dimensions(n_timesteps=100):
 
     # Variable with time + component
     ds['component_flow'] = xr.DataArray(
-        np.random.randn(n_timesteps, 3),
+        np.random.randn(n_timesteps, 2),
         dims=['time', 'component'],
     )
 
@@ -63,7 +63,7 @@ def create_dataset_with_mixed_dimensions(n_timesteps=100):
 
     # Variable with time + component + bus
     ds['flow_on_bus'] = xr.DataArray(
-        np.random.randn(n_timesteps, 3, 2),
+        np.random.randn(n_timesteps, 2, 2),
         dims=['time', 'component', 'bus'],
     )
 
@@ -75,7 +75,7 @@ def create_dataset_with_mixed_dimensions(n_timesteps=100):
 
     # Variable with time + component + scenario
     ds['component_scenario_flow'] = xr.DataArray(
-        np.random.randn(n_timesteps, 3, 2),
+        np.random.randn(n_timesteps, 2, 2),
         dims=['time', 'component', 'scenario'],
     )
 
@@ -109,12 +109,12 @@ def test_resample_equivalence_single_dimension(method):
     """
     Test with variables having only time dimension.
     """
-    timesteps = pd.date_range('2020-01-01', periods=100, freq='h')
+    timesteps = pd.date_range('2020-01-01', periods=48, freq='h')
 
     ds = xr.Dataset(coords={'time': timesteps})
-    ds['var1'] = xr.DataArray(np.random.randn(100), dims=['time'])
-    ds['var2'] = xr.DataArray(np.random.randn(100) * 10, dims=['time'])
-    ds['var3'] = xr.DataArray(np.random.randn(100) / 5, dims=['time'])
+    ds['var1'] = xr.DataArray(np.random.randn(48), dims=['time'])
+    ds['var2'] = xr.DataArray(np.random.randn(48) * 10, dims=['time'])
+    ds['var3'] = xr.DataArray(np.random.randn(48) / 5, dims=['time'])
 
     # Optimized approach
     result_optimized = fx.FlowSystem._resample_by_dimension_groups(ds, '2h', method)
@@ -130,7 +130,7 @@ def test_resample_equivalence_empty_dataset():
     """
     Test with an empty dataset (edge case).
     """
-    timesteps = pd.date_range('2020-01-01', periods=100, freq='h')
+    timesteps = pd.date_range('2020-01-01', periods=48, freq='h')
     ds = xr.Dataset(coords={'time': timesteps})
 
     # Both should handle empty dataset gracefully
@@ -144,9 +144,9 @@ def test_resample_equivalence_single_variable():
     """
     Test with a single variable.
     """
-    timesteps = pd.date_range('2020-01-01', periods=100, freq='h')
+    timesteps = pd.date_range('2020-01-01', periods=48, freq='h')
     ds = xr.Dataset(coords={'time': timesteps})
-    ds['single_var'] = xr.DataArray(np.random.randn(100), dims=['time'])
+    ds['single_var'] = xr.DataArray(np.random.randn(48), dims=['time'])
 
     # Test multiple methods
     for method in ['mean', 'sum', 'max', 'min']:
@@ -160,14 +160,14 @@ def test_resample_equivalence_with_nans():
     """
     Test with NaN values to ensure they're handled consistently.
     """
-    timesteps = pd.date_range('2020-01-01', periods=100, freq='h')
+    timesteps = pd.date_range('2020-01-01', periods=48, freq='h')
 
     ds = xr.Dataset(coords={'time': timesteps, 'component': ['a', 'b']})
 
     # Create variable with some NaN values
-    data = np.random.randn(100, 2)
-    data[10:20, 0] = np.nan
-    data[50:55, 1] = np.nan
+    data = np.random.randn(48, 2)
+    data[5:10, 0] = np.nan
+    data[20:25, 1] = np.nan
 
     ds['var_with_nans'] = xr.DataArray(data, dims=['time', 'component'])
 
@@ -183,31 +183,31 @@ def test_resample_equivalence_different_dimension_orders():
     """
     Test that dimension order doesn't affect the equivalence.
     """
-    timesteps = pd.date_range('2020-01-01', periods=100, freq='h')
+    timesteps = pd.date_range('2020-01-01', periods=48, freq='h')
 
     ds = xr.Dataset(
         coords={
             'time': timesteps,
-            'x': ['x1', 'x2', 'x3'],
+            'x': ['x1', 'x2'],
             'y': ['y1', 'y2'],
         }
     )
 
     # Variable with time first
     ds['var_time_first'] = xr.DataArray(
-        np.random.randn(100, 3, 2),
+        np.random.randn(48, 2, 2),
         dims=['time', 'x', 'y'],
     )
 
     # Variable with time in middle
     ds['var_time_middle'] = xr.DataArray(
-        np.random.randn(3, 100, 2),
+        np.random.randn(2, 48, 2),
         dims=['x', 'time', 'y'],
     )
 
     # Variable with time last
     ds['var_time_last'] = xr.DataArray(
-        np.random.randn(3, 2, 100),
+        np.random.randn(2, 2, 48),
         dims=['x', 'y', 'time'],
     )
 
@@ -225,14 +225,14 @@ def test_resample_equivalence_multiple_variables_same_dims():
     This is the key optimization case - variables with same dims should be
     grouped and resampled together.
     """
-    timesteps = pd.date_range('2020-01-01', periods=100, freq='h')
+    timesteps = pd.date_range('2020-01-01', periods=48, freq='h')
 
     ds = xr.Dataset(coords={'time': timesteps, 'location': ['A', 'B', 'C']})
 
     # Multiple variables with same dimensions (time, location)
-    for i in range(5):
+    for i in range(3):
         ds[f'var_{i}'] = xr.DataArray(
-            np.random.randn(100, 3),
+            np.random.randn(48, 3),
             dims=['time', 'location'],
         )
 
@@ -247,21 +247,21 @@ def test_resample_equivalence_large_dataset():
     """
     Test with a larger, more realistic dataset.
     """
-    timesteps = pd.date_range('2020-01-01', periods=8760, freq='h')  # Full year
+    timesteps = pd.date_range('2020-01-01', periods=168, freq='h')  # One week
 
     ds = xr.Dataset(
         coords={
             'time': timesteps,
-            'component': [f'comp_{i}' for i in range(10)],
-            'bus': [f'bus_{i}' for i in range(5)],
+            'component': [f'comp_{i}' for i in range(5)],
+            'bus': [f'bus_{i}' for i in range(3)],
         }
     )
 
     # Various variable types
-    ds['simple_var'] = xr.DataArray(np.random.randn(8760), dims=['time'])
-    ds['component_var'] = xr.DataArray(np.random.randn(8760, 10), dims=['time', 'component'])
-    ds['bus_var'] = xr.DataArray(np.random.randn(8760, 5), dims=['time', 'bus'])
-    ds['complex_var'] = xr.DataArray(np.random.randn(8760, 10, 5), dims=['time', 'component', 'bus'])
+    ds['simple_var'] = xr.DataArray(np.random.randn(168), dims=['time'])
+    ds['component_var'] = xr.DataArray(np.random.randn(168, 5), dims=['time', 'component'])
+    ds['bus_var'] = xr.DataArray(np.random.randn(168, 3), dims=['time', 'bus'])
+    ds['complex_var'] = xr.DataArray(np.random.randn(168, 5, 3), dims=['time', 'component', 'bus'])
 
     # Test with a subset of methods (to keep test time reasonable)
     for method in ['mean', 'sum', 'first']: