diff --git a/examples/04_Scenarios/scenario_example.py b/examples/04_Scenarios/scenario_example.py index 8e3349a4a..03c2a5be0 100644 --- a/examples/04_Scenarios/scenario_example.py +++ b/examples/04_Scenarios/scenario_example.py @@ -9,15 +9,16 @@ import flixopt as fx if __name__ == '__main__': - # --- Create Time Series Data --- - # Heat demand profile (e.g., kW) over time and corresponding power prices - heat_demand_per_h = np.array([[30, 0, 90, 110, 110, 20, 20, 20, 20], - [30, 0, 100, 118, 125, 20, 20, 20, 20]]).T - power_prices = np.ones(9) * 0.08 - # Create datetime array starting from '2020-01-01' for the given time period timesteps = pd.date_range('2020-01-01', periods=9, freq='h') scenarios = pd.Index(['Base Case', 'High Demand']) + + # --- Create Time Series Data --- + # Heat demand profile (e.g., kW) over time and corresponding power prices + heat_demand_per_h = pd.DataFrame({'Base Case':[30, 0, 90, 110, 110, 20, 20, 20, 20], + 'High Demand':[30, 0, 100, 118, 125, 20, 20, 20, 20]}, index=timesteps) + power_prices = np.array([0.08, 0.09]) + flow_system = fx.FlowSystem(timesteps=timesteps, scenarios=scenarios) # --- Define Energy Buses --- diff --git a/flixopt/core.py b/flixopt/core.py index 386a1d873..304048201 100644 --- a/flixopt/core.py +++ b/flixopt/core.py @@ -83,6 +83,12 @@ def as_dataarray( elif isinstance(data, np.ndarray): return DataConverter._convert_ndarray(data, coords, dims) + elif isinstance(data, pd.Series): + return DataConverter._convert_series(data, coords, dims) + + elif isinstance(data, pd.DataFrame): + return DataConverter._convert_dataframe(data, coords, dims) + else: raise ConversionError(f'Unsupported data type: {type(data).__name__}') @@ -171,6 +177,8 @@ def _convert_scalar( Returns: DataArray with the scalar value """ + if isinstance(data, (np.integer, np.floating)): + data = data.item() return xr.DataArray(data, coords=coords, dims=dims) @staticmethod @@ -192,7 +200,7 @@ def _convert_dataarray(data: xr.DataArray, coords: Dict[str, pd.Index], dims: Tu raise ConversionError('When converting to dimensionless DataArray, source must be scalar') return xr.DataArray(data.values.item()) - # Check if data already has matching dimensions + # Check if data already has matching dimensions and coordinates if set(data.dims) == set(dims): # Check if coordinates match is_compatible = True @@ -202,8 +210,13 @@ def _convert_dataarray(data: xr.DataArray, coords: Dict[str, pd.Index], dims: Tu break if is_compatible: - # Return existing DataArray if compatible - return data.copy(deep=True) + # Ensure dimensions are in the correct order + if data.dims != dims: + # Transpose to get dimensions in the right order + return data.transpose(*dims).copy(deep=True) + else: + # Return existing DataArray if compatible and order is correct + return data.copy(deep=True) # Handle dimension broadcasting if len(data.dims) == 1 and len(dims) == 2: @@ -216,8 +229,9 @@ def _convert_dataarray(data: xr.DataArray, coords: Dict[str, pd.Index], dims: Tu # Broadcast scenario dimension to include time return DataConverter._broadcast_scenario_to_time(data, coords, dims) - raise ConversionError(f'Cannot convert {data.dims} to {dims}') - + raise ConversionError( + f'Cannot convert {data.dims} to {dims}. Source coordinates: {data.coords}, Target coordinates: {coords}' + ) @staticmethod def _broadcast_time_to_scenarios( data: xr.DataArray, coords: Dict[str, pd.Index], dims: Tuple[str, ...] @@ -239,7 +253,7 @@ def _broadcast_time_to_scenarios( # Broadcast values values = np.tile(data.values, (len(coords['scenario']), 1)) - return xr.DataArray(values, coords=coords, dims=dims) + return xr.DataArray(values.copy(), coords=coords, dims=dims) @staticmethod def _broadcast_scenario_to_time( @@ -262,7 +276,7 @@ def _broadcast_scenario_to_time( # Broadcast values values = np.repeat(data.values[:, np.newaxis], len(coords['time']), axis=1) - return xr.DataArray(values, coords=coords, dims=dims) + return xr.DataArray(values.copy(), coords=coords, dims=dims) @staticmethod def _convert_ndarray(data: np.ndarray, coords: Dict[str, pd.Index], dims: Tuple[str, ...]) -> xr.DataArray: @@ -359,6 +373,113 @@ def _convert_ndarray_two_dims(data: np.ndarray, coords: Dict[str, pd.Index], dim else: raise ConversionError(f'Expected 1D or 2D array for two dimensions, got {data.ndim}D') + @staticmethod + def _convert_series(data: pd.Series, coords: Dict[str, pd.Index], dims: Tuple[str, ...]) -> xr.DataArray: + """ + Convert pandas Series to xarray DataArray. + + Args: + data: pandas Series to convert + coords: Target coordinates + dims: Target dimensions + + Returns: + DataArray from the pandas Series + """ + # Handle single dimension case + if len(dims) == 1: + dim_name = dims[0] + + # Check if series index matches the dimension + if data.index.equals(coords[dim_name]): + return xr.DataArray(data.values.copy(), coords=coords, dims=dims) + else: + raise ConversionError( + f"Series index doesn't match {dim_name} coordinates.\n" + f'Series index: {data.index}\n' + f'Target {dim_name} coordinates: {coords[dim_name]}' + ) + + # Handle two dimensions case + elif len(dims) == 2: + # Check if dimensions are time and scenario + if dims != ('time', 'scenario'): + raise ConversionError( + f'Two-dimensional conversion only supports time and scenario dimensions, got {dims}' + ) + + # Case 1: Series is indexed by time + if data.index.equals(coords['time']): + # Broadcast across scenarios + values = np.tile(data.values[:, np.newaxis], (1, len(coords['scenario']))) + return xr.DataArray(values.copy(), coords=coords, dims=dims) + + # Case 2: Series is indexed by scenario + elif data.index.equals(coords['scenario']): + # Broadcast across time + values = np.repeat(data.values[np.newaxis, :], len(coords['time']), axis=0) + return xr.DataArray(values.copy(), coords=coords, dims=dims) + + else: + raise ConversionError( + "Series index must match either 'time' or 'scenario' coordinates.\n" + f'Series index: {data.index}\n' + f'Target time coordinates: {coords["time"]}\n' + f'Target scenario coordinates: {coords["scenario"]}' + ) + + else: + raise ConversionError(f'Maximum 2 dimensions supported, got {len(dims)}') + + @staticmethod + def _convert_dataframe(data: pd.DataFrame, coords: Dict[str, pd.Index], dims: Tuple[str, ...]) -> xr.DataArray: + """ + Convert pandas DataFrame to xarray DataArray. + Only allows time as index and scenarios as columns. + + Args: + data: pandas DataFrame to convert + coords: Target coordinates + dims: Target dimensions + + Returns: + DataArray from the pandas DataFrame + """ + # Single dimension case + if len(dims) == 1: + # If DataFrame has one column, treat it like a Series + if len(data.columns) == 1: + series = data.iloc[:, 0] + return DataConverter._convert_series(series, coords, dims) + + raise ConversionError( + f'When converting DataFrame to single-dimension DataArray, DataFrame must have exactly one column, got {len(data.columns)}' + ) + + # Two dimensions case + elif len(dims) == 2: + # Check if dimensions are time and scenario + if dims != ('time', 'scenario'): + raise ConversionError( + f'Two-dimensional conversion only supports time and scenario dimensions, got {dims}' + ) + + # DataFrame must have time as index and scenarios as columns + if data.index.equals(coords['time']) and data.columns.equals(coords['scenario']): + # Create DataArray with proper dimension order + return xr.DataArray(data.values.copy(), coords=coords, dims=dims) + else: + raise ConversionError( + 'DataFrame must have time as index and scenarios as columns.\n' + f'DataFrame index: {data.index}\n' + f'DataFrame columns: {data.columns}\n' + f'Target time coordinates: {coords["time"]}\n' + f'Target scenario coordinates: {coords["scenario"]}' + ) + + else: + raise ConversionError(f'Maximum 2 dimensions supported, got {len(dims)}') + class TimeSeriesData: # TODO: Move to Interface.py @@ -913,8 +1034,8 @@ def clear_selection(self, timesteps: bool = True, scenarios: bool = True) -> Non if scenarios: self._selected_scenarios = None - # Apply the selection to all TimeSeries objects - self._propagate_selection_to_time_series() + for ts in self._time_series.values(): + ts.clear_selection(timesteps=timesteps, scenarios=scenarios) def set_selection(self, timesteps: Optional[pd.DatetimeIndex] = None, scenarios: Optional[pd.Index] = None) -> None: """ diff --git a/flixopt/structure.py b/flixopt/structure.py index 6830dbb1c..a1c9ffa0d 100644 --- a/flixopt/structure.py +++ b/flixopt/structure.py @@ -583,6 +583,9 @@ def copy_and_convert_datatypes(data: Any, use_numpy: bool = True, use_element_la return copy_and_convert_datatypes(data.selected_data, use_numpy, use_element_label) elif isinstance(data, TimeSeriesData): return copy_and_convert_datatypes(data.data, use_numpy, use_element_label) + elif isinstance(data, (pd.Series, pd.DataFrame)): + #TODO: This can be improved + return copy_and_convert_datatypes(data.values, use_numpy, use_element_label) elif isinstance(data, Interface): if use_element_label and isinstance(data, Element): diff --git a/tests/test_dataconverter.py b/tests/test_dataconverter.py index a023b8e58..61adcb284 100644 --- a/tests/test_dataconverter.py +++ b/tests/test_dataconverter.py @@ -101,8 +101,8 @@ def test_scalar_with_scenarios(self, sample_time_index, sample_scenario_index): result = DataConverter.as_dataarray(42, sample_time_index, sample_scenario_index) assert isinstance(result, xr.DataArray) - assert result.shape == (len(sample_scenario_index), len(sample_time_index)) - assert result.dims == ('scenario', 'time') + assert result.shape == (len(sample_time_index), len(sample_scenario_index)) + assert result.dims == ('time', 'scenario') assert np.all(result.values == 42) assert set(result.coords['scenario'].values) == set(sample_scenario_index.values) assert set(result.coords['time'].values) == set(sample_time_index.values) @@ -119,8 +119,8 @@ def test_1d_array_with_scenarios(self, sample_time_index, sample_scenario_index) # Convert with scenarios result = DataConverter.as_dataarray(arr_1d, sample_time_index, sample_scenario_index) - assert result.shape == (len(sample_scenario_index), len(sample_time_index)) - assert result.dims == ('scenario', 'time') + assert result.shape == (len(sample_time_index), len(sample_scenario_index)) + assert result.dims == ('time', 'scenario') # Each scenario should have the same values (broadcasting) for scenario in sample_scenario_index: @@ -139,10 +139,10 @@ def test_2d_array_with_scenarios(self, sample_time_index, sample_scenario_index) ) # Convert to DataArray - result = DataConverter.as_dataarray(arr_2d, sample_time_index, sample_scenario_index) + result = DataConverter.as_dataarray(arr_2d.T, sample_time_index, sample_scenario_index) - assert result.shape == (3, 5) - assert result.dims == ('scenario', 'time') + assert result.shape == (5, 3) + assert result.dims == ('time', 'scenario') # Check that each scenario has correct values assert np.array_equal(result.sel(scenario='baseline').values, arr_2d[0]) @@ -161,28 +161,181 @@ def test_dataarray_with_scenarios(self, sample_time_index, sample_scenario_index # Test conversion result = DataConverter.as_dataarray(original, sample_time_index, sample_scenario_index) - assert result.shape == (3, 5) - assert result.dims == ('scenario', 'time') - assert np.array_equal(result.values, original.values) + assert result.shape == (5, 3) + assert result.dims == ('time', 'scenario') + assert np.array_equal(result.values, original.values.T) # Ensure it's a copy - result.loc['baseline'] = 999 + result.loc[:, 'baseline'] = 999 assert original.sel(scenario='baseline')[0].item() == 1 # Original should be unchanged - def test_time_only_dataarray_with_scenarios(self, sample_time_index, sample_scenario_index): - """Test broadcasting a time-only DataArray to scenarios.""" - # Create a DataArray with only time dimension - time_only = xr.DataArray(data=np.array([1, 2, 3, 4, 5]), coords={'time': sample_time_index}, dims=['time']) - # Convert with scenarios - should broadcast to all scenarios - result = DataConverter.as_dataarray(time_only, sample_time_index, sample_scenario_index) +class TestSeriesConversion: + """Tests for converting pandas Series to DataArray.""" + + def test_series_single_dimension(self, sample_time_index): + """Test converting a pandas Series with time index.""" + # Create a Series with matching time index + series = pd.Series([10, 20, 30, 40, 50], index=sample_time_index) + + # Convert and check + result = DataConverter.as_dataarray(series, sample_time_index) + assert isinstance(result, xr.DataArray) + assert result.shape == (5,) + assert result.dims == ('time',) + assert np.array_equal(result.values, series.values) + assert np.array_equal(result.coords['time'].values, sample_time_index.values) + + # Test with scenario index + scenario_index = pd.Index(['baseline', 'high_demand', 'low_price'], name='scenario') + series = pd.Series([100, 200, 300], index=scenario_index) + + result = DataConverter.as_dataarray(series, scenarios=scenario_index) + assert result.shape == (3,) + assert result.dims == ('scenario',) + assert np.array_equal(result.values, series.values) + assert np.array_equal(result.coords['scenario'].values, scenario_index.values) + + def test_series_mismatched_index(self, sample_time_index): + """Test converting a Series with mismatched index.""" + # Create Series with different time index + different_times = pd.date_range('2025-01-01', periods=5, freq='D', name='time') + series = pd.Series([10, 20, 30, 40, 50], index=different_times) + + # Should raise error for mismatched index + with pytest.raises(ConversionError): + DataConverter.as_dataarray(series, sample_time_index) + + def test_series_broadcast_to_scenarios(self, sample_time_index, sample_scenario_index): + """Test broadcasting a time-indexed Series across scenarios.""" + # Create a Series with time index + series = pd.Series([10, 20, 30, 40, 50], index=sample_time_index) + + # Convert with scenarios + result = DataConverter.as_dataarray(series, sample_time_index, sample_scenario_index) - assert result.shape == (3, 5) - assert result.dims == ('scenario', 'time') + assert result.shape == (5, 3) + assert result.dims == ('time', 'scenario') - # Each scenario should have same values + # Check broadcasting - each scenario should have the same values for scenario in sample_scenario_index: - assert np.array_equal(result.sel(scenario=scenario).values, time_only.values) + scenario_slice = result.sel(scenario=scenario) + assert np.array_equal(scenario_slice.values, series.values) + + def test_series_broadcast_to_time(self, sample_time_index, sample_scenario_index): + """Test broadcasting a scenario-indexed Series across time.""" + # Create a Series with scenario index + series = pd.Series([100, 200, 300], index=sample_scenario_index) + + # Convert with time + result = DataConverter.as_dataarray(series, sample_time_index, sample_scenario_index) + + assert result.shape == (5, 3) + assert result.dims == ('time', 'scenario') + + # Check broadcasting - each time should have the same scenario values + for i, time in enumerate(sample_time_index): + time_slice = result.sel(time=time) + assert np.array_equal(time_slice.values, series.values) + + def test_series_dimension_order(self, sample_time_index, sample_scenario_index): + """Test that dimension order is respected with Series conversions.""" + # Create custom dimensions tuple with reversed order + dims = ('scenario', 'time',) + coords = {'time': sample_time_index, 'scenario': sample_scenario_index} + + # Time-indexed series + series = pd.Series([10, 20, 30, 40, 50], index=sample_time_index) + with pytest.raises(ConversionError, match="only supports time and scenario dimensions"): + _ = DataConverter._convert_series(series, coords, dims) + + # Scenario-indexed series + series = pd.Series([100, 200, 300], index=sample_scenario_index) + with pytest.raises(ConversionError, match="only supports time and scenario dimensions"): + _ = DataConverter._convert_series(series, coords, dims) + + +class TestDataFrameConversion: + """Tests for converting pandas DataFrame to DataArray.""" + + def test_dataframe_single_column(self, sample_time_index): + """Test converting a DataFrame with a single column.""" + # Create DataFrame with one column + df = pd.DataFrame({'value': [10, 20, 30, 40, 50]}, index=sample_time_index) + + # Convert and check + result = DataConverter.as_dataarray(df, sample_time_index) + assert isinstance(result, xr.DataArray) + assert result.shape == (5,) + assert result.dims == ('time',) + assert np.array_equal(result.values, df['value'].values) + + def test_dataframe_multi_column_fails(self, sample_time_index): + """Test that converting a multi-column DataFrame to 1D fails.""" + # Create DataFrame with multiple columns + df = pd.DataFrame({'val1': [10, 20, 30, 40, 50], 'val2': [15, 25, 35, 45, 55]}, index=sample_time_index) + + # Should raise error + with pytest.raises(ConversionError): + DataConverter.as_dataarray(df, sample_time_index) + + def test_dataframe_time_scenario(self, sample_time_index, sample_scenario_index): + """Test converting a DataFrame with time index and scenario columns.""" + # Create DataFrame with time as index and scenarios as columns + data = {'baseline': [10, 20, 30, 40, 50], 'high_demand': [15, 25, 35, 45, 55], 'low_price': [5, 15, 25, 35, 45]} + df = pd.DataFrame(data, index=sample_time_index) + + # Make sure columns are named properly + df.columns.name = 'scenario' + + # Convert and check + result = DataConverter.as_dataarray(df, sample_time_index, sample_scenario_index) + + assert result.shape == (5, 3) + assert result.dims == ('time', 'scenario') + assert np.array_equal(result.values, df.values) + + # Check values for specific scenarios + assert np.array_equal(result.sel(scenario='baseline').values, df['baseline'].values) + assert np.array_equal(result.sel(scenario='high_demand').values, df['high_demand'].values) + + def test_dataframe_mismatched_coordinates(self, sample_time_index, sample_scenario_index): + """Test conversion fails with mismatched coordinates.""" + # Create DataFrame with different time index + different_times = pd.date_range('2025-01-01', periods=5, freq='D', name='time') + data = {'baseline': [10, 20, 30, 40, 50], 'high_demand': [15, 25, 35, 45, 55], 'low_price': [5, 15, 25, 35, 45]} + df = pd.DataFrame(data, index=different_times) + df.columns = sample_scenario_index + + # Should raise error + with pytest.raises(ConversionError): + DataConverter.as_dataarray(df, sample_time_index, sample_scenario_index) + + # Create DataFrame with different scenario columns + different_scenarios = pd.Index(['scenario1', 'scenario2', 'scenario3'], name='scenario') + data = {'scenario1': [10, 20, 30, 40, 50], 'scenario2': [15, 25, 35, 45, 55], 'scenario3': [5, 15, 25, 35, 45]} + df = pd.DataFrame(data, index=sample_time_index) + df.columns = different_scenarios + + # Should raise error + with pytest.raises(ConversionError): + DataConverter.as_dataarray(df, sample_time_index, sample_scenario_index) + + def test_ensure_copy(self, sample_time_index, sample_scenario_index): + """Test that the returned DataArray is a copy.""" + # Create DataFrame + data = {'baseline': [10, 20, 30, 40, 50], 'high_demand': [15, 25, 35, 45, 55], 'low_price': [5, 15, 25, 35, 45]} + df = pd.DataFrame(data, index=sample_time_index) + df.columns = sample_scenario_index + + # Convert + result = DataConverter.as_dataarray(df, sample_time_index, sample_scenario_index) + + # Modify the result + result.loc[dict(time=sample_time_index[0], scenario='baseline')] = 999 + + # Original should be unchanged + assert df.loc[sample_time_index[0], 'baseline'] == 10 class TestInvalidInputs: @@ -314,8 +467,8 @@ def test_single_timestep(self, sample_scenario_index): # With scenarios result_with_scenarios = DataConverter.as_dataarray(42, single_timestep, sample_scenario_index) - assert result_with_scenarios.shape == (len(sample_scenario_index), 1) - assert result_with_scenarios.dims == ('scenario', 'time') + assert result_with_scenarios.shape == (1, len(sample_scenario_index)) + assert result_with_scenarios.dims == ('time', 'scenario') def test_single_scenario(self, sample_time_index): """Test with a single scenario.""" @@ -324,19 +477,19 @@ def test_single_scenario(self, sample_time_index): # Scalar conversion with single scenario result = DataConverter.as_dataarray(42, sample_time_index, single_scenario) - assert result.shape == (1, len(sample_time_index)) - assert result.dims == ('scenario', 'time') + assert result.shape == (len(sample_time_index), 1) + assert result.dims == ('time', 'scenario') # Array conversion with single scenario arr = np.array([1, 2, 3, 4, 5]) result_arr = DataConverter.as_dataarray(arr, sample_time_index, single_scenario) - assert result_arr.shape == (1, 5) + assert result_arr.shape == (5, 1) assert np.array_equal(result_arr.sel(scenario='baseline').values, arr) # 2D array with single scenario arr_2d = np.array([[1, 2, 3, 4, 5]]) # Note the extra dimension - result_arr_2d = DataConverter.as_dataarray(arr_2d, sample_time_index, single_scenario) - assert result_arr_2d.shape == (1, 5) + result_arr_2d = DataConverter.as_dataarray(arr_2d.T, sample_time_index, single_scenario) + assert result_arr_2d.shape == (5, 1) assert np.array_equal(result_arr_2d.sel(scenario='baseline').values, arr_2d[0]) def test_different_scenario_order(self, sample_time_index): @@ -352,7 +505,7 @@ def test_different_scenario_order(self, sample_time_index): [6, 7, 8, 9, 10], # b [11, 12, 13, 14, 15], # c ] - ) + ).T result1 = DataConverter.as_dataarray(data, sample_time_index, scenarios1) assert np.array_equal(result1.sel(scenario='a').values, [1, 2, 3, 4, 5]) @@ -374,7 +527,7 @@ def test_all_nan_data(self, sample_time_index, sample_scenario_index): # With scenarios result = DataConverter.as_dataarray(all_nan_array, sample_time_index, sample_scenario_index) - assert result.shape == (len(sample_scenario_index), len(sample_time_index)) + assert result.shape == (len(sample_time_index), len(sample_scenario_index)) assert np.all(np.isnan(result.values)) # Series of all NaNs @@ -417,11 +570,11 @@ def test_large_dataset(self, sample_scenario_index): large_data = np.random.rand(len(sample_scenario_index), len(large_timesteps)) # Convert and check - result = DataConverter.as_dataarray(large_data, large_timesteps, sample_scenario_index) + result = DataConverter.as_dataarray(large_data.T, large_timesteps, sample_scenario_index) - assert result.shape == (len(sample_scenario_index), len(large_timesteps)) - assert result.dims == ('scenario', 'time') - assert np.array_equal(result.values, large_data) + assert result.shape == (len(large_timesteps), len(sample_scenario_index)) + assert result.dims == ('time', 'scenario') + assert np.array_equal(result.values, large_data.T) class TestMultiScenarioArrayConversion: @@ -432,7 +585,7 @@ def test_1d_array_broadcasting(self, sample_time_index, sample_scenario_index): arr_1d = np.array([1, 2, 3, 4, 5]) result = DataConverter.as_dataarray(arr_1d, sample_time_index, sample_scenario_index) - assert result.shape == (len(sample_scenario_index), len(sample_time_index)) + assert result.shape == (len(sample_time_index), len(sample_scenario_index)) # Each scenario should have identical values for i, scenario in enumerate(sample_scenario_index): @@ -451,15 +604,15 @@ def test_2d_array_different_shapes(self, sample_time_index): single_scenario = pd.Index(['baseline'], name='scenario') arr_1_scenario = np.array([[1, 2, 3, 4, 5]]) - result = DataConverter.as_dataarray(arr_1_scenario, sample_time_index, single_scenario) - assert result.shape == (1, len(sample_time_index)) + result = DataConverter.as_dataarray(arr_1_scenario.T, sample_time_index, single_scenario) + assert result.shape == (len(sample_time_index), 1) # Test with 2 scenarios two_scenarios = pd.Index(['baseline', 'high_demand'], name='scenario') arr_2_scenarios = np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]) - result = DataConverter.as_dataarray(arr_2_scenarios, sample_time_index, two_scenarios) - assert result.shape == (2, len(sample_time_index)) + result = DataConverter.as_dataarray(arr_2_scenarios.T, sample_time_index, two_scenarios) + assert result.shape == (len(sample_time_index), 2) assert np.array_equal(result.sel(scenario='baseline').values, arr_2_scenarios[0]) assert np.array_equal(result.sel(scenario='high_demand').values, arr_2_scenarios[1]) @@ -474,7 +627,7 @@ def test_array_handling_edge_cases(self, sample_time_index, sample_scenario_inde bool_array = np.array([True, False, True, False, True]) result = DataConverter.as_dataarray(bool_array, sample_time_index, sample_scenario_index) assert result.dtype == bool - assert result.shape == (len(sample_scenario_index), len(sample_time_index)) + assert result.shape == (len(sample_time_index), len(sample_scenario_index)) # Test with array containing infinite values inf_array = np.array([1, np.inf, 3, -np.inf, 5]) @@ -504,7 +657,7 @@ def test_preserving_scenario_order(self, sample_time_index): ) # Convert to DataArray - result = DataConverter.as_dataarray(data, sample_time_index, scenarios) + result = DataConverter.as_dataarray(data.T, sample_time_index, scenarios) # Verify order of scenarios is preserved assert list(result.coords['scenario'].values) == list(scenarios) diff --git a/tests/test_timeseries.py b/tests/test_timeseries.py index d64c13d85..8237cf293 100644 --- a/tests/test_timeseries.py +++ b/tests/test_timeseries.py @@ -66,7 +66,7 @@ def test_initialization_validation(self, sample_timesteps): """Test validation during initialization.""" # Test missing time dimension invalid_data = xr.DataArray([1, 2, 3], dims=['invalid_dim']) - with pytest.raises(ValueError, match='must have a "time" index'): + with pytest.raises(ValueError, match='DataArray dimensions must be subset of'): TimeSeries(invalid_data, name='Invalid Series') # Test multi-dimensional data @@ -356,7 +356,7 @@ def test_initialization_with_scenarios(self, simple_scenario_dataarray): # Check basic properties assert ts.name == 'Scenario Series' - assert ts._has_scenarios is True + assert ts.has_scenario_dim is True assert ts._selected_scenarios is None # No selection initially # Check data initialization @@ -615,29 +615,29 @@ def test_add_time_series_with_scenarios(self, sample_scenario_allocator): """Test creating time series with scenarios.""" # Test scalar (broadcasts to all scenarios) ts1 = sample_scenario_allocator.add_time_series('scalar_series', 42) - assert ts1._has_scenarios + assert ts1.has_scenario_dim assert ts1.name == 'scalar_series' - assert ts1.selected_data.shape == (3, 5) # 3 scenarios, 5 timesteps + assert ts1.selected_data.shape == (5, 3) # 5 timesteps, 3 scenarios assert np.all(ts1.selected_data.values == 42) # Test 1D array (broadcasts to all scenarios) data = np.array([1, 2, 3, 4, 5]) ts2 = sample_scenario_allocator.add_time_series('array_series', data) - assert ts2._has_scenarios - assert ts2.selected_data.shape == (3, 5) + assert ts2.has_scenario_dim + assert ts2.selected_data.shape == (5, 3) # Each scenario should have the same values for scenario in sample_scenario_allocator.scenarios: assert np.array_equal(ts2.sel(scenario=scenario).values, data) # Test 2D array (one row per scenario) - data_2d = np.array([[10, 20, 30, 40, 50], [15, 25, 35, 45, 55], [5, 15, 25, 35, 45]]) + data_2d = np.array([[10, 20, 30, 40, 50], [15, 25, 35, 45, 55], [5, 15, 25, 35, 45]]).T ts3 = sample_scenario_allocator.add_time_series('scenario_specific_series', data_2d) - assert ts3._has_scenarios - assert ts3.selected_data.shape == (3, 5) + assert ts3.has_scenario_dim + assert ts3.selected_data.shape == (5, 3) # Each scenario should have its own values - assert np.array_equal(ts3.sel(scenario='baseline').values, data_2d[0]) - assert np.array_equal(ts3.sel(scenario='high_demand').values, data_2d[1]) - assert np.array_equal(ts3.sel(scenario='low_price').values, data_2d[2]) + assert np.array_equal(ts3.sel(scenario='baseline').values, data_2d[:,0]) + assert np.array_equal(ts3.sel(scenario='high_demand').values, data_2d[:,1]) + assert np.array_equal(ts3.sel(scenario='low_price').values, data_2d[:,2]) def test_selection_propagation_with_scenarios( self, sample_scenario_allocator, sample_timesteps, sample_scenario_index @@ -660,8 +660,8 @@ def test_selection_propagation_with_scenarios( assert ts2._selected_scenarios.equals(subset_scenarios) # Check data is filtered - assert ts1.selected_data.shape == (2, 5) # 2 scenarios, 5 timesteps - assert ts2.selected_data.shape == (2, 5) + assert ts1.selected_data.shape == (5, 2) # 5 timesteps, 2 scenarios + assert ts2.selected_data.shape == (5, 2) # Apply combined selection subset_timesteps = sample_timesteps[1:3] @@ -670,20 +670,22 @@ def test_selection_propagation_with_scenarios( # Check combined selection applied assert ts1._selected_timesteps.equals(subset_timesteps) assert ts1._selected_scenarios.equals(subset_scenarios) - assert ts1.selected_data.shape == (2, 2) # 2 scenarios, 2 timesteps + assert ts1.selected_data.shape == (2, 2) # 2 timesteps, 2 scenarios # Clear selections sample_scenario_allocator.clear_selection() assert ts1._selected_timesteps is None + assert ts1.active_timesteps.equals(sample_scenario_allocator.timesteps) assert ts1._selected_scenarios is None - assert ts1.selected_data.shape == (3, 5) # Back to full shape + assert ts1.active_scenarios.equals(sample_scenario_allocator.scenarios) + assert ts1.selected_data.shape == (5, 3) # Back to full shape def test_as_dataset_with_scenarios(self, sample_scenario_allocator): """Test as_dataset method with scenarios.""" # Add some time series sample_scenario_allocator.add_time_series('scalar_series', 42) sample_scenario_allocator.add_time_series( - 'varying_series', np.array([[10, 20, 30, 40, 50], [15, 25, 35, 45, 55], [5, 15, 25, 35, 45]]) + 'varying_series', np.array([[10, 20, 30, 40, 50], [15, 25, 35, 45, 55], [5, 15, 25, 35, 45]]).T ) # Get dataset @@ -723,21 +725,21 @@ def test_update_time_series_with_scenarios(self, sample_scenario_allocator, samp """Test updating a time series with scenarios.""" # Add a time series ts = sample_scenario_allocator.add_time_series('series', 42) - assert ts._has_scenarios + assert ts.has_scenario_dim assert np.all(ts.selected_data.values == 42) # Update with scenario-specific data - new_data = np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]]) + new_data = np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]]).T sample_scenario_allocator.update_time_series('series', new_data) # Check update was applied assert np.array_equal(ts.selected_data.values, new_data) - assert ts._has_scenarios + assert ts.has_scenario_dim # Check scenario-specific values - assert np.array_equal(ts.sel(scenario='baseline').values, new_data[0]) - assert np.array_equal(ts.sel(scenario='high_demand').values, new_data[1]) - assert np.array_equal(ts.sel(scenario='low_price').values, new_data[2]) + assert np.array_equal(ts.sel(scenario='baseline').values, new_data[:,0]) + assert np.array_equal(ts.sel(scenario='high_demand').values, new_data[:,1]) + assert np.array_equal(ts.sel(scenario='low_price').values, new_data[:,2]) if __name__ == '__main__':