crabnet package

Subpackages

Submodules

crabnet.crabnet_ module

Top-level module for instantiating a CrabNet model to predict properties.

class crabnet.crabnet_.CrabNet(*args: Any, **kwargs: Any)[source]

Bases: torch.nn.Module

Model class for instantiating, training, and predicting with CrabNet models.

__init__(model: Optional[Union[str, crabnet.kingcrab.SubCrab]] = None, model_name: str = 'UnnamedModel', n_elements: Union[str, int] = 'infer', classification: bool = False, verbose: bool = True, force_cpu: bool = False, prefer_last: bool = True, batch_size: Optional[int] = None, epochs: Optional[int] = None, epochs_step: int = 10, checkin: Optional[int] = None, fudge: float = 0.02, out_dims: int = 3, d_model: int = 512, extend_features: Optional[List[str]] = None, N: int = 3, heads: int = 4, elem_prop: str = 'mat2vec', compute_device: Optional[Union[str, torch.device]] = None, out_hidden: List[int] = [1024, 512, 256, 128], pe_resolution: int = 5000, ple_resolution: int = 5000, bias=False, emb_scaler: float = 1.0, pos_scaler: float = 1.0, pos_scaler_log: float = 1.0, dim_feedforward: int = 2048, dropout: float = 0.1, val_size: float = 0.2, criterion: Optional[Union[str, Callable]] = None, lr: float = 0.001, betas: Tuple[float, float] = (0.9, 0.999), eps: float = 1e-06, weight_decay: float = 0, adam: bool = False, min_trust: Optional[float] = None, alpha: float = 0.5, k: int = 6, base_lr: float = 0.0001, max_lr: float = 0.006, random_state: Optional[int] = None, mat_prop: Optional[Union[str, os.PathLike]] = None, losscurve: bool = True, learningcurve: bool = True, save: bool = True)[source]

Instantiate a CrabNet model.

Parameters

  • model (_CrabNet) – Instantiated CrabNet class, by default None.

  • model_name (str, optional) – The name of your model, by default “UnnamedModel”

  • n_elements (str, optional) – The maximum number of elements to consider during featurization, by default “infer”

  • classification (bool, optional) – Whether to perform classification. If False, then assume regression. By default, False

  • verbose (bool, optional) – Whether model information and progress should be printed, by default True

  • force_cpu (bool, optional) – Put all models on the cpu regardless of other available devices CPU, by default False

  • prefer_last (bool, optional) – Whether to prefer last used compute_device, by default True

  • batch_size (int) – The batch size to use during training. If not None, then used as-is. If specified, then it is assigned either 2 ** 7 == 128 or 2 ** 12 == 4096 based on the value of data_size.

  • epochs (int, optional) – How many epochs (# of passes through entire dataset). If None, then this is automatically assigned based on the dataset size using get_epochs_checkin_stepsize. The number of epochs must be even. By default None

  • checkin (int, optional) – When to do the checkin step. If None, then automatically assigned as half the number of epochs, by default None

  • fudge (float, optional) – The “fudge” (i.e. noise) applied to the fractional encodings, by default 0.02

  • out_dims (int, optional) – Output dimensions for Residual Network, by default 3

  • d_model (int, optional) – Size of the Model, see paper, by default 512

  • extend_features (_type_, optional) – Whether extended features will be included, by default None

  • N (int, optional) – Number of attention layers, by default 3

  • heads (int, optional) – Number of attention heads to use, by default 4

  • elem_prop (str, optional) – Which elemental feature vector to use. Possible values are “jarvis”, “magpie”, “mat2vec”, “oliynyk”, “onehot”, “ptable”, and “random_200”, by default “mat2vec”

  • compute_device (_type_, optional) – Computing device to run model on, by default None

  • out_hidden (list(int), optional) – Architecture of hidden layers in the Residual Network, by default [1024, 512, 256, 128]

  • pe_resolution (int, optional) – Number of discretizations for the prevalence encoding, by default 5000

  • ple_resolution (int, optional) – Number of discretizations for the prevalence log encoding, by default 5000

  • elem_prop – Which elemental feature vector to use. Possible values are “jarvis”, “magpie”, “mat2vec”, “oliynyk”, “onehot”, “ptable”, and “random_200”, by default “mat2vec”

  • bias (bool, optional) – Whether to bias the Residual Network, by default False

  • emb_scaler (float, optional) – Float value by which to scale the elemental embeddings, by default 1.0

  • pos_scaler (float, optional) – Float value by which to scale the fractional encodings, by default 1.0

  • pos_scaler_log (float, optional) – Float value by which to scale the log fractional encodings, by default 1.0

  • dim_feedforward (int, optional) – Dimenions of the feed forward network following transformer, by default 2048

  • dropout (float, optional) – Percent dropout in the feed forward network following the transformer, by default 0.1

  • val_size (float, optional) – fraction of validation data to take from training data only if val_df is None. By default, 0.2

  • criterion (torch.nn Module, optional) – Or in other words the loss function (e.g. BCEWithLogitsLoss for classification or RobustL1 for regression), by default None. Possible values are BCEWithLogitsLoss, RobustL1, and RobustL2.

  • lr (float, optional) – Learning rate, by default 1e-3

  • betas (tuple, optional) – Coefficients on gradient and squared gradient during Lamb optimization, by default (0.9, 0.999)

  • eps (float, optional) – Value added to the denominator during Lamb optimization, by default 1e-6

  • weight_decay (float, optional) – L2 penalty in Lamb, by default 0

  • adam (bool, optional) – Whether to constrain the Lamb model to be the Adam model, by default False

  • min_trust (float, optional) – [description], by default None

  • alpha (float, optional) – Lookahead “slow update” rate, by default 0.5

  • k (int, optional) – Number of Lookahead steps, by default 6

  • base_lr (float, optional) – Base learning rate, by default 1e-4

  • max_lr (float, optional) – Max learning rate, by default 6e-3

  • random_state (int, optional) – The seed to use for both torch and numpy random number generators. If None, then this has no effect. By default None.

  • mat_prop (str, optional) – name of material property (doesn’t affect computation), by default None

  • losscurve (bool, optional) – Whether to plot a loss curve, by default False

  • learningcurve (bool, optional) – Whether to plot a learning curve, by default True

  • save (bool, optional) – Whether to save the weights of the model, by default True

fit(train_df: Optional[pandas.core.frame.DataFrame] = None, val_df: Optional[pandas.core.frame.DataFrame] = None, extend_features: Optional[List[str]] = None, data_dir: Union[str, os.PathLike] = '/home/docs/checkouts/readthedocs.org/user_builds/crabnet/checkouts/latest/crabnet/data/materials_data', transfer: Optional[str] = None)[source]

Fit CrabNet to training data and update hyperparams with validation data.

Parameters

  • train_df (pd.DataFrame, optional) – Training and validation data with at minimum, “formula” and “target” columns and optionally, “extra features” (based on names in extend_features). If val_df is None, then test_size determines the amount of training data to be split into val_df. By default None

  • val_df (pd.DataFrame, optional) – Training and validation data with at minimum, “formula” and “target” columns and optionally, “extra features” (based on names in extend_features). If val_df is None, then test_size determines the amount of training data to be split into val_df. By default None

  • extend_features (List[str], optional) – Names of columns to use as extra features from train_df and val_df, by default None

  • data_dir (str, optional) – The directory from which to load data if loading from a file rather than a DataFrame. data_dir is only used if both train_df and val_df are None. It is assumed that the files in the data directory will be named train.csv, val.csv, and test.csv. By default join(dirname(__file__), “data”, “materials_data”)

  • transfer (str, optional) – Path to the saved weights to use for transfer learning. If None, then no transfer learning is performed. By default None

load_data(data: Union[str, pandas.core.frame.DataFrame], extra_features: Optional[pandas.core.frame.DataFrame] = None, batch_size: int = 512, train: bool = False)[source]

Load data using PyTorch Dataloader.

Parameters

  • data (Union[str, pd.DataFrame]) – The data to load, which can be the folder in which the .csv file resides or a pandas DataFrame.

  • extra_features (pd.DataFrame, optional) – DataFrame containing the extra features that will be used for training (e.g. state variables) that were extracted based on the column names in CrabNet().extend_features, by default None

  • batch_size (int, optional) – The batch size to use during training. By default 2 ** 9

  • train (bool, optional) – Whether this is the training data, by default False

load_model(model: Union[str, crabnet.kingcrab.SubCrab], data: Union[str, pandas.core.frame.DataFrame], classification: bool = False, verbose: bool = True)[source]

Load a _CrabNet model.

Parameters

  • model (Union[str, _CrabNet]) – The CrabNet model to load or the filename of the saved network.

  • data (Union[str, pd.DataFrame]) – The data to load, which can be the folder in which the .csv file resides or a pandas DataFrame.

  • classification (bool, optional) – Whether to perform classification. If False, then assume regression. By default, False

  • verbose (bool, optional) – Whether model information and progress should be printed, by default True

load_network(model_data: Union[str, dict])[source]

Load network weights from a .pth file.

Parameters

model_data (Union[str, Any]) –

Either the filename of the saved model or the network (see self.network) as a dictionary of the form:

{ “weights”: self.model.state_dict(), “scaler_state”: self.scaler.state_dict(), “model_name”: model_name, }

predict(test_df: Optional[pandas.core.frame.DataFrame] = None, loader=None, return_uncertainty=False, return_true=False)[source]

Predict on new data using a fitted CrabNet model.

Parameters

  • test_df (pd.DataFrame, optional) – _description_, by default None

  • loader (torch.Dataloader, optional) – The Dataloader corresponding to the test data, by default None

  • return_uncertainty (bool, optional) – Whether to return standard deviation uncertainties. If return_true, then return_uncertainty takes precendence and is returned as the second output. By default False

  • return_true (bool, optional) – Whether to return the true values (used for comparison with the predicted values). If return_uncertainty is also specified, then the uncertainties appear before the true values (i.e. pred, std, true), by default False

Returns

  • pred (np.array) – Predicted values. Always returned.

  • uncert (np.array) – Standard deviation uncertainty. Returned if return_uncertainty. Precedes act if act is also returned.

  • act (np.array) – True values. Returned if return_true. uncert precedes act if both uncert and act are returned.

Raises

  • SyntaxError – “Specify either data or loader, not neither.”

  • SyntaxError – “Specify either data or loader, not both.”

save_network(model_name: Optional[str] = None)[source]

Save network weights to a .pth file.

Parameters

model_name (str, optional) – The name of the .pth file. If None, then use self.model_name. By default None

crabnet.kingcrab module

Contains classes for transformer architecture within CrabNet.

class crabnet.kingcrab.Embedder(*args: Any, **kwargs: Any)[source]

Bases: torch.nn.Module

Perform composition-based embeddings of elemental features.

__init__(d_model: int, compute_device: Optional[str] = None, elem_prop: str = 'mat2vec')[source]

Embed elemental features, similar to CBFV.

Parameters

  • d_model (int) – Row dimenions of elemental emeddings, by default 512

  • compute_device (str) – Name of device which the model will be run on

  • elem_prop (str) – Which elemental feature vector to use. Possible values are “jarvis”, “magpie”, “mat2vec”, “oliynyk”, “onehot”, “ptable”, and “random_200”, by default “mat2vec”

forward(src)[source]

Compute forward call for embedder class to perform elemental embeddings.

Parameters

src (torch.tensor) – Tensor containing element numbers corresponding to elements in compound

Returns

Tensor containing elemental embeddings for compounds, reduced to d_model dimensions

Return type

torch.tensor

class crabnet.kingcrab.Encoder(*args: Any, **kwargs: Any)[source]

Bases: torch.nn.Module

Create elemental descriptor matrix via element embeddings and frac. encodings.

See the CrabNet paper for further details: https://www.nature.com/articles/s41524-021-00545-1

__init__(d_model, N, heads, extend_features=None, fractional=True, attention=True, compute_device=None, pe_resolution=5000, ple_resolution=5000, elem_prop='mat2vec', emb_scaler=1.0, pos_scaler=1.0, pos_scaler_log=1.0, dim_feedforward=2048, dropout=0.1)[source]

Instantiate the Encoder class to create elemental descriptor matrix (EDM).

Parameters

  • d_model (_type_) – _description_

  • N (int, optional) – Number of encoder layers, by default 3

  • heads (int, optional) – Number of attention heads to use, by default 4

  • extend_features (Optional[List[str]]) – Additional features to grab from columns of the other DataFrames (e.g. state variables such as temperature or applied load), by default None

  • fractional (bool, optional) – Whether to weight each element by its fractional contribution, by default True.

  • attention (bool, optional) – Whether to perform self attention, by default True

  • pe_resolution (int, optional) – Number of discretizations for the prevalence encoding, by default 5000

  • ple_resolution (int, optional) – Number of discretizations for the prevalence log encoding, by default 5000

  • elem_prop (str, optional) – Which elemental feature vector to use. Possible values are “jarvis”, “magpie”, “mat2vec”, “oliynyk”, “onehot”, “ptable”, and “random_200”, by default “mat2vec”

  • emb_scaler (float, optional) – _description_, by default 1.0

  • pos_scaler (float, optional) – Scaling factor applied to fractional encoder, by default 1.0

  • pos_scaler_log (float, optional) – Scaling factor applied to log fractional encoder, by default 1.0

  • dim_feedforward (int, optional) – Dimenions of the feed forward network following transformer, by default 2048

  • dropout (float, optional) – Percent dropout in the feed forward network following the transformer, by default 0.1

forward(src, frac, extra_features=None)[source]

Compute the forward pass for encoding the elemental descriptor matrix.

Parameters

  • src (torch.tensor) – Tensor containing integers corresponding to elements in compound

  • frac (torch.tensor) – Tensor containing the fractions of each element in compound

  • extra_features (bool, optional) – Whether to append extra features after encoding, by default None

Returns

Tensor containing flattened transformer representations of compounds concatenated with extended features.

Return type

torch.tensor

class crabnet.kingcrab.FractionalEncoder(*args: Any, **kwargs: Any)[source]

Bases: torch.nn.Module

Encode element fractional amount using a “fractional encoding”.

This is inspired by the positional encoder discussed by Vaswani. https://arxiv.org/abs/1706.03762

__init__(d_model, resolution=100, log10=False, compute_device=None)[source]

Instantiate the FractionalEncoder.

Parameters

  • d_model (int) – Model size, see paper, by default 512

  • resolution (int) – Number of discretizations for the fractional prevalence encoding, by default 100

  • log10 (bool) – Whether to apply a log operation to fraction prevalence encoding, by default False

  • compute_device (str) – The compute device to store and run the FractionalEncoder class

forward(x)[source]

Perform the forward pass of the fractional encoding.

Parameters

x (torch.tensor) – Tensor of linear spaced values based on defined resolution

Returns

Sinusoidal expansions of elemental fractions

Return type

out

class crabnet.kingcrab.ResidualNetwork(*args: Any, **kwargs: Any)[source]

Bases: torch.nn.Module

Feed forward Residual Neural Network as seen in Roost.

https://doi.org/10.1038/s41467-020-19964-7

__init__(input_dim, output_dim, hidden_layer_dims, bias=False)[source]

Instantiate a ResidualNetwork model.

Parameters

  • input_dim (int) – Input dimensions for the Residual Network, specified in SubCrab() model class, by default 512

  • output_dim (int) – Output dimensions for Residual Network, by default 3

  • hidden_layer_dims (list(int)) – Hidden layer architecture for the Residual Network, by default [1024, 512, 256, 128]

  • bias (bool) – Whether to bias the linear network, by default False

forward(fea)[source]

Propagate Residual Network weights forward.

Parameters

fea (torch.tensor (n_dim)) – Tensor output of self attention block

Returns

The output of the Residual Network

Return type

fc_out

class crabnet.kingcrab.SubCrab(*args: Any, **kwargs: Any)[source]

Bases: torch.nn.Module

SubCrab model class which implements the transformer architecture.

__init__(out_dims=3, d_model=512, extend_features=None, d_extend=0, N=3, heads=4, fractional=False, attention=True, compute_device=None, out_hidden=[1024, 512, 256, 128], pe_resolution=5000, ple_resolution=5000, elem_prop='mat2vec', bias=False, emb_scaler=1.0, pos_scaler=1.0, pos_scaler_log=1.0, dim_feedforward=2048, dropout=0.1)[source]

Instantiate a SubCrab class to be used within CrabNet.

Parameters

  • out_dims (int, optional) – Output dimensions for Residual Network, by default 3

  • d_model (int, optional) – Model size. See paper, by default 512

  • extend_features (_type_, optional) – Additional features to grab from columns of the other DataFrames (e.g. state variables such as temperature or applied load), by default None

  • d_extend (int, optional) – Number of extended features, by default 0

  • N (int, optional) – Number of attention layers, by default 3

  • heads (int, optional) – Number of attention heads, by default 4

  • frac (bool, optional) – Whether to multiply x by the fractional amounts for each element, by default False

  • attn (bool, optional) – Whether to perform self attention, by default True

  • compute_device (_type_, optional) – Computing device to run model on, by default None

  • out_hidden (list(int), optional) – Architecture of hidden layers in the Residual Network, by default [1024, 512, 256, 128]

  • pe_resolution (int, optional) – Number of discretizations for the prevalence encoding, by default 5000

  • ple_resolution (int, optional) – Number of discretizations for the prevalence log encoding, by default 5000

  • elem_prop (str, optional) – Which elemental feature vector to use. Possible values are “jarvis”, “magpie”, “mat2vec”, “oliynyk”, “onehot”, “ptable”, and “random_200”, by default “mat2vec”

  • bias (bool, optional) – Whether to bias the Residual Network, by default False

  • emb_scaler (float, optional) – Float value by which to scale the elemental embeddings, by default 1.0

  • pos_scaler (float, optional) – Float value by which to scale the fractional encodings, by default 1.0

  • pos_scaler_log (float, optional) – Float value by which to scale the log fractional encodings, by default 1.0

  • dim_feedforward (int, optional) – Dimenions of the feed forward network following transformer, by default 2048

  • dropout (float, optional) – Percent dropout in the feed forward network following the transformer, by default 0.1

forward(src, frac, extra_features=None)[source]

Compute forward pass of the SubCrab model class (i.e. transformer).

Parameters

  • src (torch.tensor) – Tensor containing element numbers corresponding to elements in compound

  • frac (torch.tensor) – Tensor containing fractional amounts of each element in compound

  • extra_features (bool, optional) – Whether to append extra features after encoding, by default None

Returns

Model output containing predicted value and uncertainty for that value

Return type

torch.tensor

class crabnet.kingcrab.TransferNetwork(*args: Any, **kwargs: Any)[source]

Bases: torch.nn.Module

Learn extended representations of materials during transfer learning.

This network was designed to have little impact on predictions during training and enhance learning with the inclusion of extended features.

__init__(input_dims, output_dims)[source]

Instantiate a TransferNetwork to learn extended representations.

Parameters

  • input_dims (int) – Dimensions of input layer

  • output_dims (int) – Dimensions of output layer

forward(x)[source]

Perform a forward pass of the TransferNetwork.

Parameters

x (_type_) – _description_

Returns

_description_

Return type

_type_

Module contents

Predict materials properties using only the composition information.