GeneralPropertyFittingParameters

Included in QATK.MLFF

class GeneralPropertyFittingParameters(model_parameters=None, dataset_parameters=None, training_parameters=None)

Constructor for GeneralPropertyFittingParameters.

Parameters:
datasetParameters()
Returns:

Dataset-specific parameters.

Return type:

GeneralPropertyDatasetParameters

modelParameters()
Returns:

Model-specific parameters.

Return type:

MACEModelParameters

nlinfo()
Returns:

The nlinfo.

Return type:

dict

trainingParameters()
Returns:

Training process parameters.

Return type:

TrainingParameters

uniqueString()

Return a unique string representing the state of the object.

Usage Examples

Configure fitting parameters for training a property prediction model:

# Define model parameters
model_params = MACEModelParameters(
    number_of_channels=128,
    max_l_equivariance=1,
    distance_cutoff=6.0*Angstrom,
)

# Define training parameters
training_params = TrainingParameters(
    experiment_name='bandgap_model',
    random_seed=42,
    max_number_of_epochs=200,
    batch_size=8,
    learning_rate=0.005,
)

# Define dataset parameters for band gap property
dataset_params = GeneralPropertyDatasetParameters(
    property_key='bandgap',
    task_type=MLParameterOptions.TASK_TYPE.GENERAL,  # Configuration-level property
    validation_fraction=0.1,
    loss_weight=1.0,
    pooling=MLParameterOptions.POOLING.MEAN,
)

# Combine into GeneralPropertyFittingParameters
fitting_params = GeneralPropertyFittingParameters(
    model_parameters=model_params,
    dataset_parameters=dataset_params,
    training_parameters=training_params,
)

# Use with MachineLearnedPropertyTrainer
trainer = MachineLearnedPropertyTrainer(
    fitting_parameters=fitting_params,
    training_sets=training_set,
    train_test_split=0.8,
)
trainer.train()

Atom-wise Properties

For atom-wise properties like atomic charges, use TASK_TYPE.ATOM_WISE:

# Configure for atom-wise property prediction
dataset_params = GeneralPropertyDatasetParameters(
    property_key='atomic_charges',
    task_type=MLParameterOptions.TASK_TYPE.ATOM_WISE,  # Per-atom property
    validation_fraction=0.1,
)

fitting_params = GeneralPropertyFittingParameters(
    model_parameters=model_params,
    training_parameters=training_params(experiment_name='charge_model'),
    dataset_parameters=dataset_params,
)

Notes

The GeneralPropertyFittingParameters class organizes all parameters needed for training property prediction models. It consists of three main components:

1. Model Parameters (MACEModelParameters)

Controls the MACE neural network architecture including embedding dimensions, interaction layers, and optional fine-tuning from foundation models. See MACEModelParameters for details.

2. Training Parameters (TrainingParameters)

Controls the training process including learning rate, batch size, number of epochs, and early stopping criteria. See TrainingParameters for details.

3. Dataset Parameters (GeneralPropertyDatasetParameters)

Controls property-specific settings:

  • property_key: Name of the property in the training data (must match the key used in TrainingSet)

  • task_type: Property type - MLParameterOptions.TASK_TYPE.GENERAL for configuration-level scalars or MLParameterOptions.TASK_TYPE.ATOM_WISE for per-atom arrays

  • validation_fraction: Fraction of training data for validation

  • loss_weight: Weight for this property in the loss function

  • pooling: How to pool values from graph nodes (only used when task_type=MLParameterOptions.TASK_TYPE.GENERAL):

    • MLParameterOptions.POOLING.MEAN: For intensive properties that don’t scale with system size (e.g., band gap)

    • MLParameterOptions.POOLING.SUM: For extensive properties that scale with system size (e.g., total energy)

    • MLParameterOptions.POOLING.MAX: For properties determined by the strongest local feature (e.g., catalytic reactivity)

Property Types

Choose the appropriate task type:

  • Configuration-level (MLParameterOptions.TASK_TYPE.GENERAL): Single value per configuration

    • Examples: band gap, formation energy, total magnetization

    • Training data: scalar values in TrainingSet additional properties

    • Output: scalar value from PropertyPredictor.predict()

  • Atom-wise (MLParameterOptions.TASK_TYPE.ATOM_WISE): One value for each atom in a configuration

    • Examples: atomic charges, magnetic moments, local energies

    • Training data: arrays in TrainingSet additional properties

    • Output: array with one value per atom from PropertyPredictor.predict()