# TrainingSet with precomputed energy, force (and stress if present and desired) data. This can be
# one or more TrainingSets
training_set = nlread(
    'training_data.hdf5', TrainingSet
)[0]

# Either fetch the calculator from the training set or if not present,
# set up analogously as to used in the training set generator to
# calculate isolated atom energies during training.
calculator = LCAOCalculator()

# Setup model specific fitting parameters object for the training
# Set up non-linear coefficients with optimization.
non_linear_coefficients_parameters = NonLinearCoefficientsParameters(
    perform_optimization=False,
    initial_coefficients=Random,
    random_seed=1234,
)

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

# Setup ML model training object
machine_learned_force_field_trainer = MachineLearnedForceFieldTrainer(
    fitting_parameters=fitting_parameters,
    training_sets=training_set,
    calculator=calculator,
    # Optional parameters can be set as desired
    train_test_split=0.8,
    random_seed=1234,
)

# Run the training
machine_learned_force_field_trainer.train()

# After training, the trained model evaluator can be retrieved for validation/analysis
model_evaluator = machine_learned_force_field_trainer.modelEvaluator()
nlprint(model_evaluator)