<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 US energy consumption
This example is based on the Sankey diagrams of US energy consumption from the Lawrence Livermore National Laboratory (thanks to John Muth for the suggestion and transcribing the data). We jump straight to the final result – for more explanation of the steps and concepts, see the tutorials.
Step1:  Load the dataset
Step2:  This defines the order the nodes appear in
Step3:  Now define the Sankey diagram definition.
Step4:  Define the colours to roughly imitate the original Sankey diagram
Step5:  And here's the result!
<ASSISTANT_TASK:>
Python Code:
from floweaver import *
Explanation: US energy consumption
This example is based on the Sankey diagrams of US energy consumption from the Lawrence Livermore National Laboratory (thanks to John Muth for the suggestion and transcribing the data). We jump straight to the final result – for more explanation of the steps and concepts, see the tutorials.
End of explanation
dataset = Dataset.from_csv('us-energy-consumption.csv',
                           dim_process_filename='us-energy-consumption-processes.csv')
Explanation: Load the dataset:
End of explanation
sources = ['Solar', 'Nuclear', 'Hydro', 'Wind', 'Geothermal',
           'Natural_Gas', 'Coal', 'Biomass', 'Petroleum']
uses = ['Residential', 'Commercial', 'Industrial', 'Transportation']
Explanation: This defines the order the nodes appear in:
End of explanation
nodes = {
    'sources': ProcessGroup('type == "source"', Partition.Simple('process', sources), title='Sources'),
    'imports': ProcessGroup(['Net_Electricity_Import'], title='Net electricity imports'),
    'electricity': ProcessGroup(['Electricity_Generation'], title='Electricity Generation'),
    'uses': ProcessGroup('type == "use"', partition=Partition.Simple('process', uses)),
    
    'energy_services': ProcessGroup(['Energy_Services'], title='Energy services'),
    'rejected': ProcessGroup(['Rejected_Energy'], title='Rejected energy'),
    
    'direct_use': Waypoint(Partition.Simple('source', [
        # This is a hack to hide the labels of the partition, there should be a better way...
        (' '*i, [k]) for i, k in enumerate(sources)
    ])),
}
ordering = [
    [[], ['sources'], []],
    [['imports'], ['electricity', 'direct_use'], []],
    [[], ['uses'], []],
    [[], ['rejected', 'energy_services'], []]
]
bundles = [
    Bundle('sources', 'electricity'),
    Bundle('sources', 'uses', waypoints=['direct_use']),
    Bundle('electricity', 'uses'),
    Bundle('imports', 'uses'),
    Bundle('uses', 'energy_services'),
    Bundle('uses', 'rejected'),
    Bundle('electricity', 'rejected'),
]
Explanation: Now define the Sankey diagram definition.
End of explanation
palette = {
    'Solar': 'gold',
    'Nuclear': 'red',
    'Hydro': 'blue',
    'Wind': 'purple',
    'Geothermal': 'brown',
    'Natural_Gas': 'steelblue',
    'Coal': 'black',
    'Biomass': 'lightgreen',
    'Petroleum': 'green',
    'Electricity': 'orange',
    'Rejected energy': 'lightgrey',
    'Energy services': 'dimgrey',
}
Explanation: Define the colours to roughly imitate the original Sankey diagram:
End of explanation
sdd = SankeyDefinition(nodes, bundles, ordering,
                       flow_partition=dataset.partition('type'))
weave(sdd, dataset, palette=palette) \
    .to_widget(width=700, height=450, margins=dict(left=100, right=120), debugging=True)
Explanation: And here's the result!
End of explanation
<END_TASK>