# 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
training_parameters = TrainingParameters(
    experiment_name='mace_experiment',
    # Other training parameters can be set as desired ...
)
model_parameters = MACEModelParameters(
    max_l_equivariance=0,
    # Other model parameters can be set as desired ...
)
dataset_parameters = ForceFieldDatasetParameters(
    validation_fraction=0.2,
    # Other dataset parameters can be set as desired ...
)

fitting_parameters = MACEFittingParameters(
    model_parameters=model_parameters,
    dataset_parameters=dataset_parameters,
    training_parameters=training_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)