# Load the configuration
configuration = nlread("equilibrated_pmma.hdf5")[0]

# Create and set the calculator
builder = OPLSPotentialBuilder()
calculator = SoftMatterCalculator(builder, use_lennard_jones_pme=True, assign_atom_types=True)
configuration.setCalculator(calculator)

# Set the simulation
atomic_constraints = [FixCenterOfMass()]
method = NVEDynamics()
simulation = SoftMatterDynamicsSimulation(configuration, method, atomic_constraints=atomic_constraints)
simulation.setLogging(1000)
simulation.setTrajectory(1000, "pmma_trajectory.hdf5")

# Set the non-equilibrium hook function
non_equilibrium = SoftMatterDynamicsNonEquilibrium(
    heat_source=(0.0, 5.0 * Angstrom),
    heat_sink=(0.5, 5.0 * Angstrom),
    call_interval=100,
)

# Run the simulation
simulation.simulate(100000, hook_functions=[non_equilibrium])