Intra-Period Segmentation with cluster()¶

import timeit

import pandas as pd
import plotly.express as px

import flixopt as fx

fx.CONFIG.notebook()

from data.generate_example_systems import create_district_heating_system

flow_system = create_district_heating_system()
flow_system.connect_and_transform()

print(f'Timesteps: {len(flow_system.timesteps)}')
print(f'Duration: {(flow_system.timesteps[-1] - flow_system.timesteps[0]).days + 1} days')

# Visualize input data
heat_demand = flow_system.components['HeatDemand'].inputs[0].fixed_relative_profile
heat_demand.plotly.line(title='Heat Demand Profile')

solver = fx.solvers.HighsSolver(mip_gap=0.01)

start = timeit.default_timer()
fs_full = flow_system.copy()
fs_full.name = 'Full Optimization'
fs_full.optimize(solver)
time_full = timeit.default_timer() - start

print(f'Full optimization: {time_full:.2f} seconds')
print(f'Total cost: {fs_full.solution["costs"].item():,.0f} €')

from tsam import ExtremeConfig, SegmentConfig

start = timeit.default_timer()

# Cluster into 8 typical days with 6 segments each
fs_segmented = flow_system.transform.cluster(
    n_clusters=8,
    cluster_duration='1D',
    segments=SegmentConfig(n_segments=6),  # 6 segments per day instead of 96 quarter-hours
    extremes=ExtremeConfig(
        method='replace', max_value=['HeatDemand(Q_th)|fixed_relative_profile'], preserve_n_clusters=True
    ),
)

time_clustering = timeit.default_timer() - start

print(f'Clustering time: {time_clustering:.2f} seconds')
print(f'Original timesteps: {len(flow_system.timesteps)}')
print(
    f'Segmented timesteps: {len(fs_segmented.timesteps)} × {len(fs_segmented.clusters)} clusters = {len(fs_segmented.timesteps) * len(fs_segmented.clusters)}'
)

clustering = fs_segmented.clustering

print('Segmentation Properties:')
print(f'  is_segmented: {clustering.is_segmented}')
print(f'  n_segments: {clustering.n_segments}')
print(f'  n_clusters: {clustering.n_clusters}')
print(f'  timesteps_per_cluster (original): {clustering.timesteps_per_cluster}')
print(f'\nTime dimension uses RangeIndex: {type(fs_segmented.timesteps)}')

# Timestep duration is now a DataArray with (cluster, time) dimensions
timestep_duration = fs_segmented.timestep_duration

print(f'Timestep duration shape: {dict(timestep_duration.sizes)}')
print('\nSegment durations for cluster 0:')
cluster_0_durations = timestep_duration.sel(cluster=0).values
for i, dur in enumerate(cluster_0_durations):
    print(f'  Segment {i}: {dur:.2f} hours')
print(f'  Total: {cluster_0_durations.sum():.2f} hours (should be 24h)')

# Visualize segment durations across clusters
duration_df = timestep_duration.to_dataframe('duration').reset_index()
fig = px.bar(
    duration_df,
    x='time',
    y='duration',
    facet_col='cluster',
    facet_col_wrap=4,
    title='Segment Durations by Cluster',
    labels={'time': 'Segment', 'duration': 'Duration [hours]'},
)
fig.update_layout(height=400)
fig.show()

start = timeit.default_timer()
fs_segmented.optimize(solver)
time_segmented = timeit.default_timer() - start

print(f'Segmented optimization: {time_segmented:.2f} seconds')
print(f'Total cost: {fs_segmented.solution["costs"].item():,.0f} €')
print(f'\nSpeedup vs full: {time_full / (time_clustering + time_segmented):.1f}x')

# Duration curves show how well the distribution is preserved
fs_segmented.clustering.plot.compare(kind='duration_curve')

# Clustering quality metrics
fs_segmented.clustering.metrics.to_dataframe().style.format('{:.3f}')

start = timeit.default_timer()
fs_expanded = fs_segmented.transform.expand()
time_expand = timeit.default_timer() - start

print(f'Expansion time: {time_expand:.3f} seconds')
print(f'Expanded timesteps: {len(fs_expanded.timesteps)}')
print(f'Objective preserved: {fs_expanded.solution["costs"].item():,.0f} €')

# Compare flow rates: Full vs Expanded
import xarray as xr

flow_var = 'CHP(Q_th)|flow_rate'
comparison_ds = xr.concat(
    [fs_full.solution[flow_var], fs_expanded.solution[flow_var]],
    dim=pd.Index(['Full', 'Expanded'], name='method'),
)
comparison_ds.plotly.line(color='method', title='CHP Heat Output Comparison')

# Stage 1: Sizing with segmentation (already done)
SAFETY_MARGIN = 1.05
sizes_with_margin = {name: float(size.item()) * SAFETY_MARGIN for name, size in fs_segmented.stats.sizes.items()}

print('Optimized sizes with safety margin:')
for name, size in sizes_with_margin.items():
    print(f'  {name}: {size:.1f}')

# Stage 2: Full resolution dispatch with fixed sizes
start = timeit.default_timer()
fs_dispatch = flow_system.transform.fix_sizes(sizes_with_margin)
fs_dispatch.name = 'Two-Stage'
fs_dispatch.optimize(solver)
time_dispatch = timeit.default_timer() - start

print(f'Dispatch time: {time_dispatch:.2f} seconds')
print(f'Final cost: {fs_dispatch.solution["costs"].item():,.0f} €')

total_segmented = time_clustering + time_segmented
total_two_stage = total_segmented + time_dispatch

results = {
    'Full (baseline)': {
        'Time [s]': time_full,
        'Cost [€]': fs_full.solution['costs'].item(),
        'CHP': fs_full.stats.sizes['CHP(Q_th)'].item(),
        'Boiler': fs_full.stats.sizes['Boiler(Q_th)'].item(),
        'Storage': fs_full.stats.sizes['Storage'].item(),
    },
    'Segmented (8×6)': {
        'Time [s]': total_segmented,
        'Cost [€]': fs_segmented.solution['costs'].item(),
        'CHP': fs_segmented.stats.sizes['CHP(Q_th)'].item(),
        'Boiler': fs_segmented.stats.sizes['Boiler(Q_th)'].item(),
        'Storage': fs_segmented.stats.sizes['Storage'].item(),
    },
    'Two-Stage': {
        'Time [s]': total_two_stage,
        'Cost [€]': fs_dispatch.solution['costs'].item(),
        'CHP': sizes_with_margin['CHP(Q_th)'],
        'Boiler': sizes_with_margin['Boiler(Q_th)'],
        'Storage': sizes_with_margin['Storage'],
    },
}

comparison = pd.DataFrame(results).T
baseline_cost = comparison.loc['Full (baseline)', 'Cost [€]']
baseline_time = comparison.loc['Full (baseline)', 'Time [s]']
comparison['Cost Gap [%]'] = ((comparison['Cost [€]'] - baseline_cost) / abs(baseline_cost) * 100).round(2)
comparison['Speedup'] = (baseline_time / comparison['Time [s]']).round(1)

comparison.style.format(
    {
        'Time [s]': '{:.2f}',
        'Cost [€]': '{:,.0f}',
        'CHP': '{:.1f}',
        'Boiler': '{:.1f}',
        'Storage': '{:.0f}',
        'Cost Gap [%]': '{:.2f}',
        'Speedup': '{:.1f}x',
    }
)

from data.generate_example_systems import create_multiperiod_system

fs_multi = create_multiperiod_system()
# Use first week only for faster demo
fs_multi = fs_multi.transform.isel(time=slice(0, 168))

print(f'Periods: {list(fs_multi.periods.values)}')
print(f'Scenarios: {list(fs_multi.scenarios.values)}')

# Cluster with segmentation
fs_multi_seg = fs_multi.transform.cluster(
    n_clusters=3,
    cluster_duration='1D',
    segments=SegmentConfig(n_segments=6),
    extremes=ExtremeConfig(
        method='replace', max_value=['Building(Heat)|fixed_relative_profile'], preserve_n_clusters=True
    ),
)

print(f'Original: {len(fs_multi.timesteps)} timesteps')
print(f'Segmented: {len(fs_multi_seg.timesteps)} × {len(fs_multi_seg.clusters)} clusters')
print(f'is_segmented: {fs_multi_seg.clustering.is_segmented}')

# Cluster assignments have period/scenario dimensions
fs_multi_seg.clustering.cluster_assignments

# Optimize and expand
fs_multi_seg.optimize(solver)
fs_multi_expanded = fs_multi_seg.transform.expand()

print(f'Expanded timesteps: {len(fs_multi_expanded.timesteps)}')
print(f'Objective: {fs_multi_expanded.solution["objective"].item():,.0f} €')