Models

Model Class Documentation

The Model class defines a framework for representing parameterized models and dense neural networks, especially for use in optimization workflows. This is intended to be subclassed with user-defined behavior. Several methods must be implemented by any subclass. The class provides a consistent interface for managing model parameters, for KIM-API portable models.

Note

For more complicated ML models, like NequIP, it is best to directly train the model without using the Model class. However, Model affords a better interface for KLIFF utilities, such as uq. In future releases of KLIFF, a more tightly integrated interface for GNNs will be provided for the Model class.

Constructor

__init__(model_name=None)

Initializes the model instance. You can optionally give your model a name. On initialization, it tries to set up parameters, supported species, and influence distance by calling methods you must define.

Required Methods (To Be Implemented in Subclass)

init_model_params()

You must define this method to initialize and return a dictionary of model parameters (usually custom Parameter objects).

init_influence_distance()

Define this to return the interaction or influence distance of your model (a float). For graph neural networks, this is typically the cutoff distance times the number of convolutions.

init_supported_species()

Define this to return a dictionary of atomic species and its corresponding index. For example, {"H": 0, "O": 1} for a water model. This is used to determine the mapping between atomic species and their indices in simulators.

get_compute_argument_class()

Specific to KIM-API models. This should return the ComputeArguments.

write_kim_model(path=None)

Implement this if your model can be exported in KIM-API format.

Model Info Methods

get_influence_distance()

Returns the model’s influence distance.

get_supported_species()

Returns a dictionary of supported atomic species.

get_model_params()

Returns the dictionary of parameters used in the model.

echo_model_params(filename=sys.stdout)

Prints or writes the current parameter values (both raw and transformed).

Parameter Configuration

set_params_mutable(list_of_params)

Marks specific parameters as optimizable (mutable). You must pass a list of parameter names.

set_opt_params(**kwargs)

Sets multiple parameters at once using keyword arguments. Usually, this will call set_one_opt_param() for each parameter.

set_one_opt_param(name, settings)

Allows fine-grained control of a single parameter’s value and bounds for optimization.

echo_opt_params(filename=sys.stdout)

Displays the values of parameters that are marked as optimizable.

get_num_opt_params()

Returns the number of total optimizable (mutable) parameter values (size of opt array).

get_opt_params()

Returns all optimizable values as a single NumPy array—useful for passing to an optimizer.

update_model_params(params)

Updates the model with a new set of parameter values (in the same shape as get_opt_params()).

get_opt_param_name_value_and_indices(index)

Given a global index, returns the parameter name, its value, and its index within the model. Helpful in debugging or tracking.

get_formatted_param_bounds()

Returns a tuple of bounds for all optimizable parameters—formatted for scipy.optimize.

opt_params_has_bounds()

Returns True if any of the optimizable parameters have bounds defined.

Saving and Loading

save(filename='trained_model.yaml')

Saves the mutable parameters of the model to disk as a YAML file.

load(filename='trained_model.yaml')

Loads the model’s parameters from a YAML file (previously saved via save()).

Parameter Utilities

named_parameters()

Returns a dictionary of all parameters currently marked as mutable/optimizable.

parameters()

Returns a list of mutable parameters for optimization (akin to torch.nn.Module.parameters()).

KIM models

Neural network models