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

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

# Import a pre-calculated training dataset.
training_data = nlread('training_data_precalculated.hdf5')[-1]
training_set_precalc = TrainingSet(training_data, recalculate_training_data=False)

# Import a training dataset with a calculator to check for calculation consistency.
# In this case the calculator is different to the one given in the MomentTensorPotentialTraining
# object, and so the training data will be re-calculated.
calculator = LCAOCalculator(exchange_correlation=GGA.PBE)
training_data = nlread('training_data_recalculate.hdf5')[-1]
training_set_recalc = TrainingSet(training_data, calculator=calculator)

# Import a training dataset with a calculator to check for calculation consistency.
# In this case the calculator is the same as the one given in the MomentTensorPotentialTraining
# object, and so the training data will not be re-calculated.
calculator = LCAOCalculator(exchange_correlation=HybridGGA.HSE06)
training_data = nlread('training_data_same_calculator.hdf5')[-1]
training_set_same_calc = TrainingSet(training_data, calculator=calculator)

# Set up MTP training and run the training data calculation and MTP training.
moment_tensor_potential_training = MomentTensorPotentialTraining(
    filename='mtp_study.hdf5',
    object_id='training',
    training_sets=[training_set_precalc, training_set_recalc, training_set_same_calc],
    calculator=LCAOCalculator(exchange_correlation=HybridGGA.HSE06),
    calculate_stress=True,
    fitting_parameters_list=fitting_parameters,
)
moment_tensor_potential_training.update()