# Set up non-linear coefficients.
non_linear_coefficients_parameters = NonLinearCoefficientsParameters(
    perform_optimization=True,
    energy_only=True,
)

# Set up parameters to use in the MTP fitting.
fitting_parameters = MomentTensorPotentialFittingParameters(
    basis_size=300,
    outer_cutoff_radii=4.5*Angstrom,
    mtp_filename='active_learning_mtp.mtp',
    non_linear_coefficients_parameters=non_linear_coefficients_parameters,
)

active_learning = ActiveLearningSimulation(
    fitting_parameters=fitting_parameters,
    initial_training_data=initial_training_data,
    mtp_study_filename='mtp_study',
    mtp_study_object_id='mtp',
    reference_calculator=reference_calculator,
    candidate_threshold=0.1,
    retrain_threshold=1.0,
    check_interval=1,
    use_stress=True,
    candidate_trajectory_filename='active_learning_candidates.hdf5',
    restart_simulation=True,
    extrapolation_selection_parameters=ExtrapolationSelectionParameters(
        extrapolation_grade_algorithm=QueryByCommitteeForces,
        descriptor_cutoff=0.1,
    ),
)

constraints = [FixStrain(x=True, y=True, z=True)]

# Run the geometry optimization through the active learning object.
optimzation_configuration = active_learning.runOptimizeGeometry(
    bulk_configuration,
    max_steps=1000,
    constraints=constraints,
    trajectory_filename='active_learning_optimization_trajectory.hdf5',
    optimize_cell=False,
    trajectory_interval=1,
)

# Extract the additional training data that has been added during active learning, as TrainingSet
# object.
additional_training_data = active_learning.additionalTrainingSet()