diff --git a/flixopt/core.py b/flixopt/core.py index bec472aaa..ea447e652 100644 --- a/flixopt/core.py +++ b/flixopt/core.py @@ -255,7 +255,7 @@ def _broadcast_time_to_scenarios( return data.copy(deep=True) # Broadcast values - values = np.tile(data.values, (len(coords['scenario']), 1)).T # Tile seems to be faster than repeat() + values = np.repeat(data.values[:, np.newaxis], len(coords['scenario']), axis=1) return xr.DataArray(values.copy(), coords=coords, dims=dims) @staticmethod @@ -278,7 +278,7 @@ def _broadcast_scenario_to_time( raise ConversionError("Source scenario coordinates don't match target scenario coordinates") # Broadcast values - values = np.repeat(data.values[:, np.newaxis], len(coords['time']), axis=1) + values = np.repeat(data.values[:, np.newaxis], len(coords['time']), axis=1).T return xr.DataArray(values.copy(), coords=coords, dims=dims) @staticmethod @@ -361,7 +361,7 @@ def _convert_ndarray_two_dims(data: np.ndarray, coords: Dict[str, pd.Index], dim return xr.DataArray(values, coords=coords, dims=dims) elif data.shape[0] == scenario_length: # Broadcast across time - values = np.tile(data, (time_length, 1)) + values = np.repeat(data[np.newaxis, :], time_length, axis=0) return xr.DataArray(values, coords=coords, dims=dims) else: raise ConversionError(f"1D array length {data.shape[0]} doesn't match either dimension") @@ -414,7 +414,7 @@ def _convert_series(data: pd.Series, coords: Dict[str, pd.Index], dims: Tuple[st # 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']))) + values = np.repeat(data.values[:, np.newaxis], len(coords['scenario']), axis=1) return xr.DataArray(values.copy(), coords=coords, dims=dims) # Case 2: Series is indexed by scenario diff --git a/tests/test_dataconverter.py b/tests/test_dataconverter.py index a50754301..0484d4aac 100644 --- a/tests/test_dataconverter.py +++ b/tests/test_dataconverter.py @@ -451,6 +451,45 @@ def test_dataarray_dimension_mismatch(self, sample_time_index, sample_scenario_i with pytest.raises(ConversionError): DataConverter.as_dataarray(wrong_length, sample_time_index, sample_scenario_index) +class TestDataArrayBroadcasting: + """Tests for broadcasting DataArrays.""" + def test_broadcast_1d_array_to_2d(self, sample_time_index, sample_scenario_index): + """Test broadcasting a 1D array to all scenarios.""" + arr_1d = np.array([1, 2, 3, 4, 5]) + + xr.testing.assert_equal( + DataConverter.as_dataarray(arr_1d, sample_time_index, sample_scenario_index), + xr.DataArray( + np.array([arr_1d] * 3).T, + coords=(sample_time_index, sample_scenario_index) + ) + ) + + arr_1d = np.array([1, 2, 3]) + xr.testing.assert_equal( + DataConverter.as_dataarray(arr_1d, sample_time_index, sample_scenario_index), + xr.DataArray( + np.array([arr_1d] * 5), + coords=(sample_time_index, sample_scenario_index) + ) + ) + + def test_broadcast_1d_array_to_1d(self, sample_time_index,): + """Test broadcasting a 1D array to all scenarios.""" + arr_1d = np.array([1, 2, 3, 4, 5]) + + xr.testing.assert_equal( + DataConverter.as_dataarray(arr_1d, sample_time_index), + xr.DataArray( + arr_1d, + coords=(sample_time_index,) + ) + ) + + arr_1d = np.array([1, 2, 3]) + with pytest.raises(ConversionError): + DataConverter.as_dataarray(arr_1d, sample_time_index) + class TestEdgeCases: """Tests for edge cases and special scenarios."""