# Set up configuration
bulk_configuration = NanoTube(3,1)

# Add tags
bulk_configuration.addTags('hot',  [24, 30, 31, 35, 36, 37, 41, 42, 45, 46, 47, 48, 49, 50, 51])
bulk_configuration.addTags('cold', [0, 1, 2, 3, 5, 6, 7, 10, 11, 16, 17])

# Set calculator
calculator = TremoloXCalculator(parameters=Tersoff_C_2010())
bulk_configuration.setCalculator(calculator)

method = NVEVelocityVerlet(
    time_step=0.5*femtoSecond,
    initial_velocity=MaxwellBoltzmannDistribution()
)

momentum_exchange_hook = NonEquilibriumMomentumExchange(
    configuration=bulk_configuration,
    exchange_interval=100,
    heat_source='hot',
    heat_sink='cold',
    update_profile_interval=0,
    profile_resolution=2.0*Ang
)

md_trajectory = MolecularDynamics(
    bulk_configuration,
    constraints=[],
    trajectory_filename='md_trajectory.hdf5',
    steps=5000,
    log_interval=50,
    measurement_hook=momentum_exchange_hook,
    method=method
)