API reference

Autodocumentation of the frozen public surface. Behavioral contracts are documented in API stability policy.

Package version

graphchem.__version__

str(object=’’) -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.__str__() (if defined) or repr(object). encoding defaults to sys.getdefaultencoding(). errors defaults to ‘strict’.

graphchem.data

class graphchem.data.MoleculeGraph(atom_attr, bond_attr, connectivity, target=None)[source]

Bases: Data

A custom graph class representing a molecular structure.

This class extends the Data class from PyTorch Geometric to represent molecules with node attributes (atoms), edge attributes (bonds), and connectivity information. It also includes an optional target value.

Parameters:
x

The node features (atom attributes).

Type:

torch.Tensor

edge_index

A 2D tensor describing the connectivity between atoms.

Type:

torch.Tensor

edge_attr

Edge features (bond attributes).

Type:

torch.Tensor

y

Target value(s) of the molecule.

Type:

torch.Tensor

__init__(atom_attr, bond_attr, connectivity, target=None)[source]

Initialize the MoleculeGraph object.

Parameters:
  • atom_attr (torch.Tensor) – A 2D tensor of shape (num_atoms, num_atom_features) representing the attributes of each atom in the molecule.

  • bond_attr (torch.Tensor) – A 2D tensor of shape (num_bonds, num_bond_features) representing the attributes of each bond in the molecule.

  • connectivity (torch.Tensor) – A 2D tensor of shape (2, num_bonds) where each column represents an edge (bond) between two atoms. The first row contains the source atom indices and the second row contains the target atom indices.

  • target (Optional[torch.Tensor]) – An optional 1D or 2D tensor representing the target value(s) of the molecule. If not provided, it defaults to a tensor with a single element set to 0.0.

class graphchem.data.MoleculeDataset(graphs)[source]

Bases: Dataset

A custom dataset class for molecular graphs.

This class extends the Dataset class from PyTorch Geometric to create a dataset of molecular graphs. Each graph in the dataset is an instance of MoleculeGraph.

Parameters:

graphs (Iterable[MoleculeGraph])

_graphs

A list containing all the molecule graphs in the dataset.

Type:

List[MoleculeGraph]

__init__(graphs)[source]

Initialize the MoleculeDataset object.

Parameters:

graphs (Iterable[MoleculeGraph]) – An iterable of MoleculeGraph instances representing the molecules in the dataset.

len()[source]

Returns the number of molecules in the dataset.

Returns:

The number of molecule graphs in the dataset.

Return type:

int

get(idx)[source]

Retrieves a molecule graph from the dataset by index.

Returns:

The molecule graph at the specified index.

Return type:

MoleculeGraph

Parameters:

idx (int)

graphchem.datasets

graphchem.datasets.load_cn()[source]

Loads cetane number data.

Returns:

A tuple containing two elements: - List[str]: A list of SMILES strings. - torch.Tensor: A tensor of CN values with dtype float32.

Return type:

Tuple[List[str], torch.Tensor]

graphchem.datasets.load_lhv()[source]

Loads lower heating value data.

Returns:

A tuple containing two elements: - List[str]: A list of SMILES strings. - torch.Tensor: A tensor of LHV values with dtype float32.

Return type:

Tuple[List[str], torch.Tensor]

graphchem.datasets.load_mon()[source]

Loads motor octane number data.

Returns:

A tuple containing two elements: - List[str]: A list of SMILES strings. - torch.Tensor: A tensor of MON values with dtype float32.

Return type:

Tuple[List[str], torch.Tensor]

graphchem.datasets.load_ron()[source]

Loads research octane number data.

Returns:

A tuple containing two elements: - List[str]: A list of SMILES strings. - torch.Tensor: A tensor of RON values with dtype float32.

Return type:

Tuple[List[str], torch.Tensor]

graphchem.datasets.load_ysi()[source]

Loads yield sooting index data.

Returns:

A tuple containing two elements: - List[str]: A list of SMILES strings. - torch.Tensor: A tensor of YSI values with dtype float32.

Return type:

Tuple[List[str], torch.Tensor]

graphchem.nn

class graphchem.nn.MoleculeGCN(atom_vocab_size, bond_vocab_size, output_dim, embedding_dim=128, n_messages=2, n_readout=2, readout_dim=64, p_dropout=0.0, aggr='add', act_fn=<built-in function softplus>)[source]

Bases: Module

A Graph Convolutional Network (GCN) model for molecular property prediction.

Parameters:
  • atom_vocab_size (int)

  • bond_vocab_size (int)

  • output_dim (int | None)

  • embedding_dim (int)

  • n_messages (int)

  • n_readout (int)

  • readout_dim (int)

  • p_dropout (float)

  • aggr (str)

  • act_fn (Callable[..., Any])

_p_dropout

Probability of an element to be zeroed in dropout layers.

Type:

float

_n_messages

Number of message passing steps.

Type:

int

act_fn

Activation function, e.g. torch.nn.functional.softplus.

Type:

Callable[..., Any]

emb_atom

Embedding layer for atoms.

Type:

nn.Embedding

emb_bond

Embedding layer for bonds.

Type:

nn.Embedding

atom_conv

General convolution layer for atoms.

Type:

GeneralConv

bond_conv

Edge convolution layer for bonds.

Type:

EdgeConv

readout

Readout network consisting of fully connected layers, or None when n_readout=0.

Type:

nn.ModuleList | None

__init__(atom_vocab_size, bond_vocab_size, output_dim, embedding_dim=128, n_messages=2, n_readout=2, readout_dim=64, p_dropout=0.0, aggr='add', act_fn=<built-in function softplus>)[source]

Initialize the MoleculeGCN object.

Parameters:
  • atom_vocab_size (int) – Number of unique atom representations in the dataset.

  • bond_vocab_size (int) – Number of unique bond representations in the dataset.

  • output_dim (int or None) – Dimensionality of the output space. May be None when n_readout=0 (pooled embeddings are returned directly).

  • embedding_dim (int, optional (default=128)) – Dimensionality of the atom and bond embeddings.

  • n_messages (int, optional (default=2)) – Number of message passing steps.

  • n_readout (int, optional (default=2)) – Number of fully connected layers in the readout network.

  • readout_dim (int, optional (default=64)) – Dimensionality of the hidden layers in the readout network.

  • p_dropout (float, optional (default=0.0)) – Dropout probability for the dropout layers.

  • aggr (str, optional (default add)) – Aggregation scheme to use in the GeneralConv layer.

  • act_fn (callable, optional) – Activation function (default torch.nn.functional.softplus). Other examples: torch.nn.functional.sigmoid, torch.nn.functional.relu.

forward(data)[source]

Forward pass of the MoleculeGCN.

Parameters:

data (torch_geometric.data.Data) – Input data containing node features (x), edge attributes (edge_attr), edge indices (edge_index), and batch vector (batch).

Return type:

tuple[Tensor, Tensor, Tensor]

Returns:

  • out (torch.Tensor) – The final output predictions for the input molecules.

  • out_atom (torch.Tensor) – Atom-level representations after message passing.

  • out_bond (torch.Tensor) – Bond-level representations after message passing.

graphchem.preprocessing

graphchem.preprocessing.get_ring_size(obj, max_size=12)[source]

Determine the size of the smallest ring that an atom or bond is part of.

Parameters:
  • obj (Union[rdkit.Chem.Atom, rdkit.Chem.Bond]) – An RDKit Atom or Bond object to check for ring membership.

  • max_size (Optional[int], default 12) – The maximum size of the ring to consider. If no ring is found with a size less than or equal to max_size, this value will be returned.

Returns:

The size of the smallest ring that the atom or bond is part of, or max_size if no smaller ring is found.

Return type:

int

graphchem.preprocessing.atom_to_str(atom)[source]

Convert an RDKit Atom object to a string representation.

The string representation includes various properties of the atom, such as chiral tag, degree, explicit valence, formal charge, hybridization, implicit valence, aromaticity, number of implicit hydrogen atoms, and more.

Parameters:

atom (rdkit.Chem.Atom) – An RDKit Atom object representing a single atom in a molecule.

Returns:

A string representation of the atom, including its properties.

Return type:

str

Examples

>>> from rdkit import Chem
>>> mol = Chem.MolFromSmiles('C1=CC=CC=C1')
>>> atom = mol.GetAtomWithIdx(0)
>>> atom_to_str(atom)
'(CHIRAL_NONE, 3, 4, 0, <Hybridization.SP2: 6>, 0, True, False, 0, 0, 0,
 'C', 3, 1, 4, 5)'
graphchem.preprocessing.bond_to_str(bond)[source]

Convert an RDKit Bond object to a string representation.

The string representation includes various properties of the bond, including bond type, conjugation, stereochemistry, ring size, and connected atom symbols.

Parameters:

bond (rdkit.Chem.Bond) – An RDKit Bond object representing a single bond in a molecule.

Returns:

A string representation of the bond, including its properties.

Return type:

str

Examples

>>> from rdkit import Chem
>>> mol = Chem.MolFromSmiles('C=C')
>>> bond = mol.GetBondWithIndices(0, 1)
>>> bond_to_str(bond)
"(DOUBLE, False, NONE, None, ['C', 'C'])"
class graphchem.preprocessing.Tokenizer[source]

Bases: object

A simple tokenizer that assigns a unique integer to each token (word) in the input data. If the tokenizer is in training mode, it will add new tokens to the vocabulary. Otherwise, it will return the integer corresponding to ‘unk’ for unknown tokens.

_data

A dictionary mapping each token to a unique integer. Initialized with {“unk”: 1}.

Type:

dict

num_classes

The number of unique classes (tokens) in the vocabulary, including ‘unk’.

Type:

int

train

A flag indicating whether the tokenizer is in training mode.

Type:

bool

unknown

A list to store tokens that were encountered during inference but are not in the vocabulary.

Type:

list

__init__()[source]

Initialize the Tokenizer with default values.

__call__(item)[source]

Tokenizes a given string by returning its corresponding integer from the vocabulary.

Parameters:

item (str) – The token (word) to be tokenized.

Returns:

The unique integer assigned to the token. If the token is not in the vocabulary and the tokenizer is in training mode, it will add the token and return its corresponding integer. Otherwise, it returns 1, which corresponds to ‘unk’.

Return type:

int

property vocab_size: int

Returns the size of the vocabulary, which is the number of unique tokens plus one.

Returns:

The total number of classes (tokens) in the vocabulary plus one.

Return type:

int

class graphchem.preprocessing.MoleculeEncoder(smiles)[source]

Bases: object

A class to encode molecular SMILES strings into numerical (integer) representations using tokenized atom and bond information.

Parameters:

smiles (list[str])

_atom_tokenizer

integer Tokenizer for atom representations.

Type:

Tokenizer

_bond_tokenizer

integer Tokenizer for bond representations.

Type:

Tokenizer

__init__(smiles)[source]

Initializes the MoleculeEncoder with a list of SMILES strings and creates/trains integer tokenizers for atoms and bonds.

Parameters:

smiles (List[str]) – A list of SMILES strings representing molecules used for tokenizer creation/training.

Raises:

ValueError – If any provided SMILES string cannot be parsed by RDKit.

property vocab_sizes: tuple[int, int]

Returns the vocabulary sizes of the atom and bond tokenizers.

Returns:

A tuple containing two integers representing the sizes of the atom and bond tokenizers’ vocabularies respectively.

Return type:

Tuple[int, int]

encode(smiles)[source]

Encodes a single SMILES string into three tensors representing atoms, bonds, and connectivity.

Parameters:

smiles (str) – A SMILES string representing the molecule to be encoded.

Returns:

A tuple containing:
  • A tensor of atom encodings, shape (n_atoms,).

  • A tensor of bond encodings, shape (n_bonds,).

  • A connectivity matrix as a tensor, shape (2, n_bonds).

Return type:

Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Raises:

ValueError – If the provided SMILES string cannot be parsed by RDKit.

encode_many(smiles)[source]

Encodes a list of SMILES strings into tensors representing atoms, bonds, and connectivities.

Parameters:

smiles (Iterable[str]) – An iterable collection of SMILES strings representing molecules to be encoded.

Returns:

A list containing tuples with three elements:
  • A tensor of atom encodings, shape (n_atoms,).

  • A tensor of bond encodings, shape (n_bonds,).

  • A connectivity matrix as a tensor, shape (2, n_bonds).

Return type:

List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]

Raises:

ValueError – If any provided SMILES string cannot be parsed by RDKit.

save(filename)[source]

Save the encoder to a file.

Parameters:

filename (str) – filename/path to save the encoder to.

Return type:

None

graphchem.preprocessing.load_encoder(filename)[source]

Loads a pre-saved MoleculeEncoder object from a file.

Parameters:

filename (str) – The path to the saved encoder file.

Returns:

The loaded MoleculeEncoder object.

Return type:

MoleculeEncoder