crabnet package
Subpackages
- crabnet.data package
- crabnet.utils package
- Submodules
- crabnet.utils.composition module
- crabnet.utils.data module
- crabnet.utils.estimatorselectionhelper module
- crabnet.utils.figures module
- crabnet.utils.get_compute_device module
- crabnet.utils.get_core_count module
- crabnet.utils.modelselectionhelper module
- crabnet.utils.optim module
- crabnet.utils.utils module
- Module contents
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-6weight_decay (float, optional) – L2 penalty in
Lamb
, by default 0adam (bool, optional) – Whether to constrain the
Lamb
model to be the Adam model, by default Falsemin_trust (float, optional) – [description], by default None
alpha (float, optional) –
Lookahead
“slow update” rate, by default 0.5k (int, optional) – Number of
Lookahead
steps, by default 6base_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.”
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
- 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
- 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.
Module contents
Predict materials properties using only the composition information.