Download example files

The download module contains utilities for downloading pre-trained models, using load_pretrained(), and example data, using load_data().

unlocknn.download.load_pretrained(model_name: Literal['binary_e_form'], branch: str = 'master', save_dir: PathLike = PosixPath('/home/docs/checkouts/readthedocs.org/user_builds/unlocknn/checkouts/latest/unlocknn/models')) → MEGNetProbModel

Download a pre-trained model.

A list of available models and their descriptions can be found at https://github.com/a-ws-m/unlockNN/tree/master/models.

Parameters

  • model_name – The name of the model to download.

  • branch – Which branch of the unlockNN repository to download from.

  • save_dir – The directory to check for already-downloaded models and in which to save newly downloaded models.

Returns

The downloaded model.

Examples

Download model for predicting binary compounds’ formation energies:

>>> model = load_pretrained("binary_e_form")
unlocknn.download.load_data(data_name: Literal['binary_e_form'], branch: str = 'master', save_dir: PathLike = PosixPath('/home/docs/checkouts/readthedocs.org/user_builds/unlocknn/checkouts/latest/unlocknn/data'))

Download sample data.

A list of available data, their sources and descriptions can be found at https://github.com/a-ws-m/unlockNN/tree/master/data.

Parameters

  • data_name – The name of the data to download.

  • branch – Which branch of the unlockNN repository to download from.

  • save_dir – The directory to check for already-downloaded data and in which to save newly downloaded data.

Returns

The downloaded data.

Examples

Download binary compounds and their formation energies, then print the first dataset entry:

>>> data = load_data("binary_e_form")
>>> print(data.iloc[0])
structure                    [[ 1.982598   -4.08421341  3.2051745 ] La, [1....
formation_energy_per_atom                                            -0.737439
Name: 0, dtype: object

Adding new example files

New pre-trained models and example data can be added to unlockNN by uploading the serialized data to the GitHub repository. When doing so, make sure to add metadata to the appropriate README.md file (in the data or models directories), as well as the AVAILABLE_MODELS or AVAILABLE_DATA type variables in unlocknn.download.

For models, simply save the model, then compress the resulting folder to a tar.gz file format.

For data, unlockNN provides a convenience function for serializing pandas DataFrame s that contain pymatgen Structure s in a “structure” column: save_struct_data(). The resulting output file can be uploaded to the repository.

unlocknn.download.save_struct_data(df: DataFrame, fname: PathLike)

Save data containing Structure s to a file.

Serializes the “structure” column to a string. The converse of _load_struct_data().

Parameters

  • df – The pd.DataFrame to serialize.

  • fname – Where to save the file.