Utilities¶

DeepChem has a broad collection of utility functions. Many of these maybe be of independent interest to users since they deal with some tricky aspects of processing scientific datatypes.

Data Utilities¶

Array Utilities¶

pad_array(x: ndarray, shape: Tuple | int, fill: float = 0.0, both: bool = False) → ndarray[source]¶

Pad an array with a fill value.

Parameters:

x (np.ndarray) – A numpy array.
shape (Tuple or int) – Desired shape. If int, all dimensions are padded to that size.
fill (float, optional (default 0.0)) – The padded value.
both (bool, optional (default False)) – If True, split the padding on both sides of each axis. If False, padding is applied to the end of each axis.

Returns:

A padded numpy array

Return type:

np.ndarray

Data Directory¶

The DeepChem data directory is where downloaded MoleculeNet datasets are stored.

get_data_dir() → str[source]¶

Get the DeepChem data directory.

Returns:: The default path to store DeepChem data. If you want to change this path, please set your own path to DEEPCHEM_DATA_DIR as an environment variable.
Return type:: str

URL Handling¶

download_url(url: str, dest_dir: str = '/tmp', name: str | None = None)[source]¶

Download a file to disk.

Parameters:

url (str) – The URL to download from
dest_dir (str) – The directory to save the file in
name (str) – The file name to save it as. If omitted, it will try to extract a file name from the URL

File Handling¶

untargz_file(file: str, dest_dir: str = '/tmp', name: str | None = None)[source]¶

Untar and unzip a .tar.gz file to disk.

Parameters:

file (str) – The filepath to decompress
dest_dir (str) – The directory to save the file in
name (str) – The file name to save it as. If omitted, it will use the file name

unzip_file(file: str, dest_dir: str = '/tmp', name: str | None = None)[source]¶

Unzip a .zip file to disk.

Parameters:

file (str) – The filepath to decompress
dest_dir (str) – The directory to save the file in
name (str) – The directory name to unzip it to. If omitted, it will use the file name

load_data(input_files: List[str], shard_size: int | None = None) → Iterator[Any][source]¶

Loads data from files.

Parameters:

input_files (List[str]) – List of filenames.
shard_size (int, default None) – Size of shard to yield

Returns:

Iterator which iterates over provided files.

Return type:

Iterator[Any]

Notes

The supported file types are SDF, CSV and Pickle.

load_sdf_files(input_files: List[str], clean_mols: bool = True, tasks: List[str] = [], shard_size: int | None = None) → Iterator[DataFrame][source]¶

Load SDF file into dataframe.

Parameters:

input_files (List[str]) – List of filenames
clean_mols (bool, default True) – Whether to sanitize molecules.
tasks (List[str], default []) – Each entry in tasks is treated as a property in the SDF file and is retrieved with mol.GetProp(str(task)) where mol is the RDKit mol loaded from a given SDF entry.
shard_size (int, default None) – The shard size to yield at one time.

Returns:

Generator which yields the dataframe which is the same shard size.

Return type:

Iterator[pd.DataFrame]

Notes

This function requires RDKit to be installed.

load_csv_files(input_files: List[str], shard_size: int | None = None) → Iterator[DataFrame][source]¶

Load data as pandas dataframe from CSV files.

Parameters:

input_files (List[str]) – List of filenames
shard_size (int, default None) – The shard size to yield at one time.

Returns:

Generator which yields the dataframe which is the same shard size.

Return type:

Iterator[pd.DataFrame]

load_json_files(input_files: List[str], shard_size: int | None = None) → Iterator[DataFrame][source]¶

Load data as pandas dataframe.

Parameters:

input_files (List[str]) – List of json filenames.
shard_size (int, default None) – Chunksize for reading json files.

Returns:

Generator which yields the dataframe which is the same shard size.

Return type:

Iterator[pd.DataFrame]

Notes

To load shards from a json file into a Pandas dataframe, the file must be originally saved with df.to_json('filename.json', orient='records', lines=True)

load_pickle_files(input_files: List[str]) → Iterator[Any][source]¶

Load dataset from pickle files.

Parameters:: input_files (List[str]) – The list of filenames of pickle file. This function can load from gzipped pickle file like XXXX.pkl.gz.
Returns:: Generator which yields the objects which is loaded from each pickle file.
Return type:: Iterator[Any]

load_from_disk(filename: str) → Any[source]¶

Load a dataset from file.

Parameters:: filename (str) – A filename you want to load data.
Returns:: A loaded object from file.
Return type:: Any

save_to_disk(dataset: Any, filename: str, compress: int = 3)[source]¶

Save a dataset to file.

Parameters:

dataset (str) – A data saved
filename (str) – Path to save data.
compress (int, default 3) – The compress option when dumping joblib file.

load_dataset_from_disk(save_dir: str) → Tuple[bool, Tuple[DiskDataset, DiskDataset, DiskDataset] | None, List[Transformer]][source]¶

Loads MoleculeNet train/valid/test/transformers from disk.

Expects that data was saved using save_dataset_to_disk below. Expects the following directory structure for save_dir: save_dir/

—> train_dir/ | —> valid_dir/ | —> test_dir/ | —> transformers.pkl

Parameters:

save_dir (str) – Directory name to load datasets.

Returns:

loaded (bool) – Whether the load succeeded
all_dataset (Tuple[DiskDataset, DiskDataset, DiskDataset]) – The train, valid, test datasets
transformers (Transformer) – The transformers used for this dataset

Molecular Utilities¶

class ConformerGenerator(max_conformers: int = 1, rmsd_threshold: float = 0.5, force_field: str = 'uff', pool_multiplier: int = 10)[source]¶

Generate molecule conformers.

Notes

Procedure 1. Generate a pool of conformers. 2. Minimize conformers. 3. Prune conformers using an RMSD threshold.

Note that pruning is done _after_ minimization, which differs from the protocol described in the references [1]_ [2]_.

References

Notes

This class requires RDKit to be installed.

__init__(max_conformers: int = 1, rmsd_threshold: float = 0.5, force_field: str = 'uff', pool_multiplier: int = 10)[source]¶

Parameters:

max_conformers (int, optional (default 1)) – Maximum number of conformers to generate (after pruning).
rmsd_threshold (float, optional (default 0.5)) – RMSD threshold for pruning conformers. If None or negative, no pruning is performed.
force_field (str, optional (default 'uff')) – Force field to use for conformer energy calculation and minimization. Options are ‘uff’, ‘mmff94’, and ‘mmff94s’.
pool_multiplier (int, optional (default 10)) – Factor to multiply by max_conformers to generate the initial conformer pool. Since conformers are pruned after energy minimization, increasing the size of the pool increases the chance of identifying max_conformers unique conformers.

generate_conformers(mol: Any) → Any[source]¶

Generate conformers for a molecule.

This function returns a copy of the original molecule with embedded conformers.

Parameters:: mol (rdkit.Chem.rdchem.Mol) – RDKit Mol object
Returns:: mol – A new RDKit Mol object containing the chosen conformers, sorted by increasing energy.
Return type:: rdkit.Chem.rdchem.Mol

embed_molecule(mol: Any) → Any[source]¶

Generate conformers, possibly with pruning.

Parameters:: mol (rdkit.Chem.rdchem.Mol) – RDKit Mol object
Returns:: mol – RDKit Mol object with embedded multiple conformers.
Return type:: rdkit.Chem.rdchem.Mol

get_molecule_force_field(mol: Any, conf_id: int | None = None, **kwargs) → Any[source]¶

Get a force field for a molecule.

Parameters:

mol (rdkit.Chem.rdchem.Mol) – RDKit Mol object with embedded conformers.
conf_id (int, optional) – ID of the conformer to associate with the force field.
kwargs (dict, optional) – Keyword arguments for force field constructor.

Returns:

ff – RDKit force field instance for a molecule.

Return type:

rdkit.ForceField.rdForceField.ForceField

minimize_conformers(mol: Any) → None[source]¶

Minimize molecule conformers.

Parameters:: mol (rdkit.Chem.rdchem.Mol) – RDKit Mol object with embedded conformers.

get_conformer_energies(mol: Any) → ndarray[source]¶

Calculate conformer energies.

Parameters:: mol (rdkit.Chem.rdchem.Mol) – RDKit Mol object with embedded conformers.
Returns:: energies – Minimized conformer energies.
Return type:: np.ndarray

prune_conformers(mol: Any) → Any[source]¶

Prune conformers from a molecule using an RMSD threshold, starting with the lowest energy conformer.

Parameters:: mol (rdkit.Chem.rdchem.Mol) – RDKit Mol object
Returns:: new_mol – A new rdkit.Chem.rdchem.Mol containing the chosen conformers, sorted by increasing energy.
Return type:: rdkit.Chem.rdchem.Mol

static get_conformer_rmsd(mol: Any) → ndarray[source]¶

Calculate conformer-conformer RMSD.

Parameters:: mol (rdkit.Chem.rdchem.Mol) – RDKit Mol object
Returns:: rmsd – A conformer-conformer RMSD value. The shape is (NumConformers, NumConformers)
Return type:: np.ndarray

class MoleculeLoadException(*args, **kwargs)[source]¶

__init__(*args, **kwargs)[source]¶

get_xyz_from_mol(mol)[source]¶

Extracts a numpy array of coordinates from a molecules.

Returns a (N, 3) numpy array of 3d coords of given rdkit molecule

Parameters:: mol (rdkit Molecule) – Molecule to extract coordinates for
Return type:: Numpy ndarray of shape (N, 3) where N = mol.GetNumAtoms().

add_hydrogens_to_mol(mol, is_protein=False)[source]¶

Add hydrogens to a molecule object

Parameters:

mol (Rdkit Mol) – Molecule to hydrogenate
is_protein (bool, optional (default False)) – Whether this molecule is a protein.

Return type:

Rdkit Mol

Note

This function requires RDKit and PDBFixer to be installed.

compute_charges(mol)[source]¶

Attempt to compute Gasteiger Charges on Mol

This also has the side effect of calculating charges on mol. The mol passed into this function has to already have been sanitized

Parameters:: mol (rdkit molecule) –
Return type:: No return since updates in place.

Note

This function requires RDKit to be installed.

load_molecule(molecule_file, add_hydrogens=True, calc_charges=True, sanitize=True, is_protein=False)[source]¶

Converts molecule file to (xyz-coords, obmol object)

Given molecule_file, returns a tuple of xyz coords of molecule and an rdkit object representing that molecule in that order (xyz, rdkit_mol). This ordering convention is used in the code in a few places.

Parameters:

molecule_file (str) – filename for molecule
add_hydrogens (bool, optional (default True)) – If True, add hydrogens via pdbfixer
calc_charges (bool, optional (default True)) – If True, add charges via rdkit
sanitize (bool, optional (default False)) – If True, sanitize molecules via rdkit
is_protein (bool, optional (default False)) – If True`, this molecule is loaded as a protein. This flag will affect some of the cleanup procedures applied.

Returns:

Tuple (xyz, mol) if file contains single molecule. Else returns a
list of the tuples for the separate molecules in this list.

Note

This function requires RDKit to be installed.

write_molecule(mol, outfile, is_protein=False)[source]¶

Write molecule to a file

This function writes a representation of the provided molecule to the specified outfile. Doesn’t return anything.

Parameters:

mol (rdkit Mol) – Molecule to write
outfile (str) – Filename to write mol to
is_protein (bool, optional) – Is this molecule a protein?

Note

This function requires RDKit to be installed.

Raises:: ValueError – if outfile isn’t of a supported format.:

Molecular Fragment Utilities¶

It’s often convenient to manipulate subsets of a molecule. The MolecularFragment class aids in such manipulations.

class MolecularFragment(atoms: Sequence[Any], coords: ndarray)[source]¶

A class that represents a fragment of a molecule.

It’s often convenient to represent a fragment of a molecule. For example, if two molecules form a molecular complex, it may be useful to create two fragments which represent the subsets of each molecule that’s close to the other molecule (in the contact region).

Ideally, we’d be able to do this in RDKit direct, but manipulating molecular fragments doesn’t seem to be supported functionality.

Examples

>>> import numpy as np
>>> from rdkit import Chem
>>> mol = Chem.MolFromSmiles("C")
>>> coords = np.array([[0.0, 0.0, 0.0]])
>>> atom = mol.GetAtoms()[0]
>>> fragment = MolecularFragment([atom], coords)

__init__(atoms: Sequence[Any], coords: ndarray)[source]¶

Initialize this object.

Parameters:

atoms (Iterable[rdkit.Chem.rdchem.Atom]) – Each entry in this list should be a RDKit Atom.
coords (np.ndarray) – Array of locations for atoms of shape (N, 3) where N == len(atoms).

GetAtoms() → List[AtomShim][source]¶

Returns the list of atoms

Returns:: list of atoms in this fragment.
Return type:: List[AtomShim]

GetNumAtoms() → int[source]¶

Returns the number of atoms

Returns:: Number of atoms in this fragment.
Return type:: int

GetCoords() → ndarray[source]¶

Returns 3D coordinates for this fragment as numpy array.

Returns:: A numpy array of shape (N, 3) with coordinates for this fragment. Here, N is the number of atoms.
Return type:: np.ndarray

class AtomShim(atomic_num: int, partial_charge: float, atom_coords: ndarray)[source]¶

This is a shim object wrapping an atom.

We use this class instead of raw RDKit atoms since manipulating a large number of rdkit Atoms seems to result in segfaults. Wrapping the basic information in an AtomShim seems to avoid issues.

__init__(atomic_num: int, partial_charge: float, atom_coords: ndarray)[source]¶

Initialize this object

Parameters:

atomic_num (int) – Atomic number for this atom.
partial_charge (float) – The partial Gasteiger charge for this atom
atom_coords (np.ndarray) – Of shape (3,) with the coordinates of this atom

GetAtomicNum() → int[source]¶

Returns atomic number for this atom.

Returns:: Atomic number for this atom.
Return type:: int

GetPartialCharge() → float[source]¶

Returns partial charge for this atom.

Returns:: A partial Gasteiger charge for this atom.
Return type:: float

GetCoords() → ndarray[source]¶

Returns 3D coordinates for this atom as numpy array.

Returns:: Numpy array of shape (3,) with coordinates for this atom.
Return type:: np.ndarray

strip_hydrogens(coords: ndarray, mol: Any | MolecularFragment) → Tuple[ndarray, MolecularFragment][source]¶

Strip the hydrogens from input molecule

Parameters:

coords (np.ndarray) – The coords must be of shape (N, 3) and correspond to coordinates of mol.
mol (rdkit.Chem.rdchem.Mol or MolecularFragment) – The molecule to strip

Returns:

A tuple of (coords, mol_frag) where coords is a numpy array of coordinates with hydrogen coordinates. mol_frag is a MolecularFragment.

Return type:

Tuple[np.ndarray, MolecularFragment]

Notes

This function requires RDKit to be installed.

merge_molecular_fragments(molecules: List[MolecularFragment]) → MolecularFragment | None[source]¶

Helper method to merge two molecular fragments.

Parameters:: molecules (List[MolecularFragment]) – List of MolecularFragment objects.
Returns:: Returns a merged MolecularFragment
Return type:: Optional[MolecularFragment]

get_contact_atom_indices(fragments: List[Tuple[ndarray, Any]], cutoff: float = 4.5) → List[List[int]][source]¶

Compute that atoms close to contact region.

Molecular complexes can get very large. This can make it unwieldy to compute functions on them. To improve memory usage, it can be very useful to trim out atoms that aren’t close to contact regions. This function computes pairwise distances between all pairs of molecules in the molecular complex. If an atom is within cutoff distance of any atom on another molecule in the complex, it is regarded as a contact atom. Otherwise it is trimmed.

Parameters:

fragments (List[Tuple[np.ndarray, rdkit.Chem.rdchem.Mol]]) – As returned by rdkit_utils.load_complex, a list of tuples of (coords, mol) where coords is a (N_atoms, 3) array and mol is the rdkit molecule object.
cutoff (float, optional (default 4.5)) – The cutoff distance in angstroms.

Returns:

A list of length len(molecular_complex). Each entry in this list is a list of atom indices from that molecule which should be kept, in sorted order.

Return type:

List[List[int]]

reduce_molecular_complex_to_contacts(fragments: List[Tuple[ndarray, Any]], cutoff: float = 4.5) → List[Tuple[ndarray, MolecularFragment]][source]¶

Reduce a molecular complex to only those atoms near a contact.

Molecular complexes can get very large. This can make it unwieldy to compute functions on them. To improve memory usage, it can be very useful to trim out atoms that aren’t close to contact regions. This function takes in a molecular complex and returns a new molecular complex representation that contains only contact atoms. The contact atoms are computed by calling get_contact_atom_indices under the hood.

Parameters:

fragments (List[Tuple[np.ndarray, rdkit.Chem.rdchem.Mol]]) – As returned by rdkit_utils.load_complex, a list of tuples of (coords, mol) where coords is a (N_atoms, 3) array and mol is the rdkit molecule object.
cutoff (float) – The cutoff distance in angstroms.

Returns:

A list of length len(molecular_complex). Each entry in this list is a tuple of (coords, MolecularFragment). The coords is stripped down to (N_contact_atoms, 3) where N_contact_atoms is the number of contact atoms for this complex. MolecularFragment is used since it’s tricky to make a RDKit sub-molecule.

Return type:

List[Tuple[np.ndarray, MolecularFragment]]

Coordinate Box Utilities¶

class CoordinateBox(x_range: Tuple[float, float], y_range: Tuple[float, float], z_range: Tuple[float, float])[source]¶

A coordinate box that represents a block in space.

Molecular complexes are typically represented with atoms as coordinate points. Each complex is naturally associated with a number of different box regions. For example, the bounding box is a box that contains all atoms in the molecular complex. A binding pocket box is a box that focuses in on a binding region of a protein to a ligand. A interface box is the region in which two proteins have a bulk interaction.

The CoordinateBox class is designed to represent such regions of space. It consists of the coordinates of the box, and the collection of atoms that live in this box alongside their coordinates.

__init__(x_range: Tuple[float, float], y_range: Tuple[float, float], z_range: Tuple[float, float])[source]¶

Initialize this box.

Parameters:

x_range (Tuple[float, float]) – A tuple of (x_min, x_max) with max and min x-coordinates.
y_range (Tuple[float, float]) – A tuple of (y_min, y_max) with max and min y-coordinates.
z_range (Tuple[float, float]) – A tuple of (z_min, z_max) with max and min z-coordinates.

Raises:

ValueError –

__contains__(point: Sequence[float]) → bool[source]¶

Check whether a point is in this box.

Parameters:: point (Sequence[float]) – 3-tuple or list of length 3 or np.ndarray of shape (3,). The (x, y, z) coordinates of a point in space.
Returns:: True if other is contained in this box.
Return type:: bool

center() → Tuple[float, float, float][source]¶

Computes the center of this box.

Returns:: (x, y, z) the coordinates of the center of the box.
Return type:: Tuple[float, float, float]

Examples

>>> box = CoordinateBox((0, 1), (0, 1), (0, 1))
>>> box.center()
(0.5, 0.5, 0.5)

volume() → float[source]¶

Computes and returns the volume of this box.

Returns:: The volume of this box. Can be 0 if box is empty
Return type:: float

Examples

>>> box = CoordinateBox((0, 1), (0, 1), (0, 1))
>>> box.volume()
1

contains(other: CoordinateBox) → bool[source]¶

Test whether this box contains another.

This method checks whether other is contained in this box.

Parameters:: other (CoordinateBox) – The box to check is contained in this box.
Returns:: True if other is contained in this box.
Return type:: bool
Raises:: ValueError –

intersect_interval(interval1: Tuple[float, float], interval2: Tuple[float, float]) → Tuple[float, float][source]¶

Computes the intersection of two intervals.

Parameters:

interval1 (Tuple[float, float]) – Should be (x1_min, x1_max)
interval2 (Tuple[float, float]) – Should be (x2_min, x2_max)

Returns:

x_intersect – Should be the intersection. If the intersection is empty returns (0, 0) to represent the empty set. Otherwise is (max(x1_min, x2_min), min(x1_max, x2_max)).

Return type:

Tuple[float, float]

union(box1: CoordinateBox, box2: CoordinateBox) → CoordinateBox[source]¶

Merges provided boxes to find the smallest union box.

This method merges the two provided boxes.

Parameters:

box1 (CoordinateBox) – First box to merge in
box2 (CoordinateBox) – Second box to merge into this box

Returns:

Smallest CoordinateBox that contains both box1 and box2

Return type:

CoordinateBox

merge_overlapping_boxes(boxes: List[CoordinateBox], threshold: float = 0.8) → List[CoordinateBox][source]¶

Merge boxes which have an overlap greater than threshold.

Parameters:

boxes (list[CoordinateBox]) – A list of CoordinateBox objects.
threshold (float, default 0.8) – The volume fraction of the boxes that must overlap for them to be merged together.

Returns:

List[CoordinateBox] of merged boxes. This list will have length less than or equal to the length of boxes.

Return type:

List[CoordinateBox]

get_face_boxes(coords: ndarray, pad: float = 5.0) → List[CoordinateBox][source]¶

For each face of the convex hull, compute a coordinate box around it.

The convex hull of a macromolecule will have a series of triangular faces. For each such triangular face, we construct a bounding box around this triangle. Think of this box as attempting to capture some binding interaction region whose exterior is controlled by the box. Note that this box will likely be a crude approximation, but the advantage of this technique is that it only uses simple geometry to provide some basic biological insight into the molecule at hand.

The pad parameter is used to control the amount of padding around the face to be used for the coordinate box.

Parameters:

coords (np.ndarray) – A numpy array of shape (N, 3). The coordinates of a molecule.
pad (float, optional (default 5.0)) – The number of angstroms to pad.

Returns:

boxes – List of CoordinateBox

Return type:

List[CoordinateBox]

Examples

>>> coords = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]])
>>> boxes = get_face_boxes(coords, pad=5)

Evaluation Utils¶

class Evaluator(model, dataset: Dataset, transformers: List[Transformer])[source]¶

Class that evaluates a model on a given dataset.

The evaluator class is used to evaluate a dc.models.Model class on a given dc.data.Dataset object. The evaluator is aware of dc.trans.Transformer objects so will automatically undo any transformations which have been applied.

Examples

Evaluators allow for a model to be evaluated directly on a Metric for sklearn. Let’s do a bit of setup constructing our dataset and model.

>>> import deepchem as dc
>>> import numpy as np
>>> X = np.random.rand(10, 5)
>>> y = np.random.rand(10, 1)
>>> dataset = dc.data.NumpyDataset(X, y)
>>> model = dc.models.MultitaskRegressor(1, 5)
>>> transformers = []

Then you can evaluate this model as follows >>> import sklearn >>> evaluator = Evaluator(model, dataset, transformers) >>> multitask_scores = evaluator.compute_model_performance( … sklearn.metrics.mean_absolute_error)

Evaluators can also be used with dc.metrics.Metric objects as well in case you want to customize your metric further.

>>> evaluator = Evaluator(model, dataset, transformers)
>>> metric = dc.metrics.Metric(dc.metrics.mae_score)
>>> multitask_scores = evaluator.compute_model_performance(metric)

__init__(model, dataset: Dataset, transformers: List[Transformer])[source]¶

Initialize this evaluator

Parameters:

model (Model) – Model to evaluate. Note that this must be a regression or classification model and not a generative model.
dataset (Dataset) – Dataset object to evaluate model on.
transformers (List[Transformer]) – List of dc.trans.Transformer objects. These transformations must have been applied to dataset previously. The dataset will be untransformed for metric evaluation.

output_statistics(scores: Dict[str, float], stats_out: str)[source]¶

Write computed stats to file.

Parameters:

scores (dict) – Dictionary mapping names of metrics to scores.
stats_out (str) – Name of file to write scores to.

output_predictions(y_preds: ndarray, csv_out: str)[source]¶

Writes predictions to file.

Writes predictions made on self.dataset to a specified file on disk. self.dataset.ids are used to format predictions.

Parameters:

y_preds (np.ndarray) – Predictions to output
csv_out (str) – Name of file to write predictions to.

compute_model_performance(metrics: Metric | Callable[[...], Any] | List[Metric] | List[Callable[[...], Any]], csv_out: str | None = None, stats_out: str | None = None, per_task_metrics: bool = False, use_sample_weights: bool = False, n_classes: int = 2) → Dict[str, float] | Tuple[Dict[str, float], Dict[str, float]][source]¶

Computes statistics of model on test data and saves results to csv.

Parameters:

metrics (dc.metrics.Metric/list[dc.metrics.Metric]/function) – The set of metrics provided. This class attempts to do some intelligent handling of input. If a single dc.metrics.Metric object is provided or a list is provided, it will evaluate self.model on these metrics. If a function is provided, it is assumed to be a metric function that this method will attempt to wrap in a dc.metrics.Metric object. A metric function must accept two arguments, y_true, y_pred both of which are np.ndarray objects and return a floating point score. The metric function may also accept a keyword argument sample_weight to account for per-sample weights.
csv_out (str, optional (DEPRECATED)) – Filename to write CSV of model predictions.
stats_out (str, optional (DEPRECATED)) – Filename to write computed statistics.
per_task_metrics (bool, optional) – If true, return computed metric for each task on multitask dataset.
use_sample_weights (bool, optional (default False)) – If set, use per-sample weights w.
n_classes (int, optional (default None)) – If specified, will use n_classes as the number of unique classes in self.dataset. Note that this argument will be ignored for regression metrics.

Returns:

multitask_scores (dict) – Dictionary mapping names of metrics to metric scores.
all_task_scores (dict, optional) – If per_task_metrics == True, then returns a second dictionary of scores for each task separately.

class GeneratorEvaluator(model, generator: Iterable[Tuple[Any, Any, Any]], transformers: List[Transformer], labels: List | None = None, weights: List | None = None)[source]¶

Evaluate models on a stream of data.

This class is a partner class to Evaluator. Instead of operating over datasets this class operates over a generator which yields batches of data to feed into provided model.

Examples

>>> import deepchem as dc
>>> import numpy as np
>>> X = np.random.rand(10, 5)
>>> y = np.random.rand(10, 1)
>>> dataset = dc.data.NumpyDataset(X, y)
>>> model = dc.models.MultitaskRegressor(1, 5)
>>> generator = model.default_generator(dataset, pad_batches=False)
>>> transformers = []

Then you can evaluate this model as follows

>>> import sklearn
>>> evaluator = GeneratorEvaluator(model, generator, transformers)
>>> multitask_scores = evaluator.compute_model_performance(
...     sklearn.metrics.mean_absolute_error)

Evaluators can also be used with dc.metrics.Metric objects as well in case you want to customize your metric further. (Note that a given generator can only be used once so we have to redefine the generator here.)

>>> generator = model.default_generator(dataset, pad_batches=False)
>>> evaluator = GeneratorEvaluator(model, generator, transformers)
>>> metric = dc.metrics.Metric(dc.metrics.mae_score)
>>> multitask_scores = evaluator.compute_model_performance(metric)

__init__(model, generator: Iterable[Tuple[Any, Any, Any]], transformers: List[Transformer], labels: List | None = None, weights: List | None = None)[source]¶

Parameters:

model (Model) – Model to evaluate.
generator (generator) – Generator which yields batches to feed into the model. For a KerasModel, it should be a tuple of the form (inputs, labels, weights). The “correct” way to create this generator is to use model.default_generator as shown in the example above.
transformers (List[Transformer]) – Tranformers to “undo” when applied to the models outputs
labels (list of Layer) – layers which are keys in the generator to compare to outputs
weights (list of Layer) – layers which are keys in the generator for weight matrices

compute_model_performance(metrics: Metric | Callable[[...], Any] | List[Metric] | List[Callable[[...], Any]], per_task_metrics: bool = False, use_sample_weights: bool = False, n_classes: int = 2) → Dict[str, float] | Tuple[Dict[str, float], Dict[str, float]][source]¶

Computes statistics of model on test data and saves results to csv.

Parameters:

metrics (dc.metrics.Metric/list[dc.metrics.Metric]/function) – The set of metrics provided. This class attempts to do some intelligent handling of input. If a single dc.metrics.Metric object is provided or a list is provided, it will evaluate self.model on these metrics. If a function is provided, it is assumed to be a metric function that this method will attempt to wrap in a dc.metrics.Metric object. A metric function must accept two arguments, y_true, y_pred both of which are np.ndarray objects and return a floating point score.
per_task_metrics (bool, optional) – If true, return computed metric for each task on multitask dataset.
use_sample_weights (bool, optional (default False)) – If set, use per-sample weights w.
n_classes (int, optional (default None)) – If specified, will assume that all metrics are classification metrics and will use n_classes as the number of unique classes in self.dataset.

Returns:

multitask_scores (dict) – Dictionary mapping names of metrics to metric scores.
all_task_scores (dict, optional) – If per_task_metrics == True, then returns a second dictionary of scores for each task separately.

relative_difference(x: ndarray, y: ndarray) → ndarray[source]¶

Compute the relative difference between x and y

The two argument arrays must have the same shape.

Parameters:

x (np.ndarray) – First input array
y (np.ndarray) – Second input array

Returns:

z – We will have z == np.abs(x-y) / np.abs(max(x, y)).

Return type:

np.ndarray

Genomic Utilities¶

seq_one_hot_encode(sequences, letters: str = 'ATCGN') → ndarray[source]¶

One hot encodes list of genomic sequences.

Sequences encoded have shape (N_sequences, N_letters, sequence_length, 1). These sequences will be processed as images with one color channel.

Parameters:

sequences (np.ndarray or Iterator[Bio.SeqRecord]) – Iterable object of genetic sequences
letters (str, optional (default "ATCGN")) – String with the set of possible letters in the sequences.

Raises:

ValueError: – If sequences are of different lengths.

Returns:

A numpy array of shape (N_sequences, N_letters, sequence_length, 1).

Return type:

np.ndarray

encode_bio_sequence(fname: str, file_type: str = 'fasta', letters: str = 'ATCGN') → ndarray[source]¶

Loads a sequence file and returns an array of one-hot sequences.

Parameters:

fname (str) – Filename of fasta file.
file_type (str, optional (default "fasta")) – The type of file encoding to process, e.g. fasta or fastq, this is passed to Biopython.SeqIO.parse.
letters (str, optional (default "ATCGN")) – The set of letters that the sequences consist of, e.g. ATCG.

Returns:

A numpy array of shape (N_sequences, N_letters, sequence_length, 1).

Return type:

np.ndarray

Notes

This function requires BioPython to be installed.

hhblits(dataset_path, database=None, data_dir=None, evalue=0.001, num_iterations=2, num_threads=4)[source]¶

Run hhblits multisequence alignment search on a dataset. This function requires the hhblits binary to be installed and in the path. This function also requires a Hidden Markov Model reference database to be provided. Both can be found here: https://github.com/soedinglab/hh-suite

The database should be in the deepchem data directory or specified as an argument. To set the deepchem data directory, run this command in your environment:

export DEEPCHEM_DATA_DIR=<path to data directory>

Parameters:

dataset_path (str) – Path to single sequence or multiple sequence alignment (MSA) dataset. Results will be saved in this directory.
database (str) – Name of database to search against. Note this is not the path, but the name of the database.
data_dir (str) – Path to database directory.
evalue (float) – E-value cutoff.
num_iterations (int) – Number of iterations.
num_threads (int, default 4) – Number of threads.

Returns:

results (.a3m file) – MSA file containing the results of the hhblits search.
results (.hhr file) – hhsuite results file containing the results of the hhblits search.

Examples

>>> from deepchem.utils.sequence_utils import hhblits
>>> msa_path = hhblits('test/data/example.fasta', database='example_db', data_dir='test/data/', evalue=0.001, num_iterations=2, num_threads=4)

hhsearch(dataset_path, database=None, data_dir=None, evalue=0.001, num_iterations=2, num_threads=4)[source]¶

Run hhsearch multisequence alignment search on a dataset. This function requires the hhblits binary to be installed and in the path. This function also requires a Hidden Markov Model reference database to be provided. Both can be found here: https://github.com/soedinglab/hh-suite

The database should be in the deepchem data directory or specified as an argument. To set the deepchem data directory, run this command in your environment:

export DEEPCHEM_DATA_DIR=<path to data directory>

Examples

>>> from deepchem.utils.sequence_utils import hhsearch
>>> msa_path = hhsearch('test/data/example.fasta', database='example_db', data_dir='test/data/', evalue=0.001, num_iterations=2, num_threads=4)

Parameters:

dataset_path (str) – Path to multiple sequence alignment dataset. Results will be saved in this directory.
database (str) – Name of database to search against. Note this is not the path, but the name of the database.
data_dir (str) – Path to database directory.
evalue (float) – E-value cutoff.
num_iterations (int) – Number of iterations.
num_threads (int) – Number of threads.

Returns:

results (.a3m file) – MSA file containing the results of the hhblits search.
results (.hhr file) – hhsuite results file containing the results of the hhblits search.

MSA_to_dataset(msa_path)[source]¶: Convert a multiple sequence alignment to a NumpyDataset object.

Geometry Utilities¶

unit_vector(vector: ndarray) → ndarray[source]¶

Returns the unit vector of the vector.

Parameters:: vector (np.ndarray) – A numpy array of shape (3,), where 3 is (x,y,z).
Returns:: A numpy array of shape (3,). The unit vector of the input vector.
Return type:: np.ndarray

angle_between(vector_i: ndarray, vector_j: ndarray) → float[source]¶

Returns the angle in radians between vectors “vector_i” and “vector_j”

Note that this function always returns the smaller of the two angles between the vectors (value between 0 and pi).

Parameters:

vector_i (np.ndarray) – A numpy array of shape (3,), where 3 is (x,y,z).
vector_j (np.ndarray) – A numpy array of shape (3,), where 3 is (x,y,z).

Returns:

The angle in radians between the two vectors.

Return type:

np.ndarray

Examples

>>> print("%0.06f" % angle_between((1, 0, 0), (0, 1, 0)))
1.570796
>>> print("%0.06f" % angle_between((1, 0, 0), (1, 0, 0)))
0.000000
>>> print("%0.06f" % angle_between((1, 0, 0), (-1, 0, 0)))
3.141593

generate_random_unit_vector() → ndarray[source]¶

Generate a random unit vector on the sphere S^2.

Citation: http://mathworld.wolfram.com/SpherePointPicking.html

Pseudocode:

Choose random theta element [0, 2*pi]
Choose random z element [-1, 1]
Compute output vector u: (x,y,z) = (sqrt(1-z^2)*cos(theta), sqrt(1-z^2)*sin(theta),z)

Returns:: u – A numpy array of shape (3,). u is an unit vector
Return type:: np.ndarray

generate_random_rotation_matrix() → ndarray[source]¶

Generates a random rotation matrix.

Generate a random unit vector u, randomly sampled from the

unit sphere (see function generate_random_unit_vector() for details)

Generate a second random unit vector v

If absolute value of u dot v > 0.99, repeat.

(This is important for numerical stability. Intuition: we want them to be as linearly independent as possible or else the orthogonalized version of v will be much shorter in magnitude compared to u. I assume in Stack they took this from Gram-Schmidt orthogonalization?)

v” = v - (u dot v)*u, i.e. subtract out the component of

v that’s in u’s direction

normalize v” (this isn”t in Stack but I assume it must be

done)

find w = u cross v”
u, v”, and w will form the columns of a rotation matrix, R.

The intuition is that u, v” and w are, respectively, what the standard basis vectors e1, e2, and e3 will be mapped to under the transformation.

Returns:: R – A numpy array of shape (3, 3). R is a rotation matrix.
Return type:: np.ndarray

is_angle_within_cutoff(vector_i: ndarray, vector_j: ndarray, angle_cutoff: float) → bool[source]¶

A utility function to compute whether two vectors are within a cutoff from 180 degrees apart.

Parameters:

vector_i (np.ndarray) – A numpy array of shape (3,)`, where 3 is (x,y,z).
vector_j (np.ndarray) – A numpy array of shape (3,), where 3 is (x,y,z).
cutoff (float) – The deviation from 180 (in degrees)

Returns:

Whether two vectors are within a cutoff from 180 degrees apart

Return type:

bool

Graph Utilities¶

fourier_encode_dist(x, num_encodings=4, include_self=True)[source]¶

Fourier encode the input tensor x based on the specified number of encodings.

This function applies a Fourier encoding to the input tensor x by dividing it by a range of scales (2^i for i in range(num_encodings)) and then concatenating the sine and cosine of the scaled values. Optionally, the original input tensor can be included in the output.

Parameters:

x (torch.Tensor) – Input tensor to be Fourier encoded.
num_encodings (int, optional, default=4) – Number of Fourier encodings to apply.
include_self (bool, optional, default=True) – Whether to include the original input tensor in the output.

Returns:

Fourier encoded tensor.

Return type:

torch.Tensor

Examples

>>> import torch
>>> x = torch.tensor([1.0, 2.0, 3.0])
>>> encoded_x = fourier_encode_dist(x, num_encodings=4, include_self=True)

aggregate_mean(h, **kwargs)[source]¶

Compute the mean of the input tensor along the second to last dimension.

Parameters:: h (torch.Tensor) – Input tensor.
Returns:: Mean of the input tensor along the second to last dimension.
Return type:: torch.Tensor

aggregate_max(h, **kwargs)[source]¶

Compute the max of the input tensor along the second to last dimension.

Parameters:: h (torch.Tensor) – Input tensor.
Returns:: Max of the input tensor along the second to last dimension.
Return type:: torch.Tensor

aggregate_min(h, **kwargs)[source]¶

Compute the min of the input tensor along the second to last dimension.

Parameters:

h (torch.Tensor) – Input tensor.
**kwargs – Additional keyword arguments.

Returns:

Min of the input tensor along the second to last dimension.

Return type:

torch.Tensor

aggregate_std(h, **kwargs)[source]¶

Compute the standard deviation of the input tensor along the second to last dimension.

Parameters:: h (torch.Tensor) – Input tensor.
Returns:: Standard deviation of the input tensor along the second to last dimension.
Return type:: torch.Tensor

aggregate_var(h, **kwargs)[source]¶

Compute the variance of the input tensor along the second to last dimension.

Parameters:: h (torch.Tensor) – Input tensor.
Returns:: Variance of the input tensor along the second to last dimension.
Return type:: torch.Tensor

aggregate_moment(h, n=3, **kwargs)[source]¶

Compute the nth moment of the input tensor along the second to last dimension.

Parameters:

h (torch.Tensor) – Input tensor.
n (int, optional, default=3) – The order of the moment to compute.

Returns:

Nth moment of the input tensor along the second to last dimension.

Return type:

torch.Tensor

aggregate_sum(h, **kwargs)[source]¶

Compute the sum of the input tensor along the second to last dimension.

Parameters:: h (torch.Tensor) – Input tensor.
Returns:: Sum of the input tensor along the second to last dimension.
Return type:: torch.Tensor

scale_identity(h, D=None, avg_d=None)[source]¶

Identity scaling function.

Parameters:

h (torch.Tensor) – Input tensor.
D (torch.Tensor, optional) – Degree tensor.
avg_d (dict, optional) – Dictionary containing averages over the training set.

Returns:

Scaled input tensor.

Return type:

torch.Tensor

scale_amplification(h, D, avg_d)[source]¶

Amplification scaling function. log(D + 1) / d * h where d is the average of the log(D + 1) in the training set

Parameters:

h (torch.Tensor) – Input tensor.
D (torch.Tensor) – Degree tensor.
avg_d (dict) – Dictionary containing averages over the training set.

Returns:

Scaled input tensor.

Return type:

torch.Tensor

scale_attenuation(h, D, avg_d)[source]¶

Attenuation scaling function. (log(D + 1))^-1 / d * X where d is the average of the log(D + 1))^-1 in the training set

Parameters:

h (torch.Tensor) – Input tensor.
D (torch.Tensor) – Degree tensor.
avg_d (dict) – Dictionary containing averages over the training set.

Returns:

Scaled input tensor.

Return type:

torch.Tensor

Hash Function Utilities¶

hash_ecfp(ecfp: str, size: int = 1024) → int[source]¶

Returns an int < size representing given ECFP fragment.

Input must be a string. This utility function is used for various ECFP based fingerprints.

Parameters:

ecfp (str) – String to hash. Usually an ECFP fragment.
size (int, optional (default 1024)) – Hash to an int in range [0, size)

Returns:

ecfp_hash – An int < size representing given ECFP fragment

Return type:

int

hash_ecfp_pair(ecfp_pair: Tuple[str, str], size: int = 1024) → int[source]¶

Returns an int < size representing that ECFP pair.

Input must be a tuple of strings. This utility is primarily used for spatial contact featurizers. For example, if a protein and ligand have close contact region, the first string could be the protein’s fragment and the second the ligand’s fragment. The pair could be hashed together to achieve one hash value for this contact region.

Parameters:

ecfp_pair (Tuple[str, str]) – Pair of ECFP fragment strings
size (int, optional (default 1024)) – Hash to an int in range [0, size)

Returns:

ecfp_hash – An int < size representing given ECFP pair.

Return type:

int

vectorize(hash_function: Callable[[Any, int], int], feature_dict: Dict[int, str] | None = None, size: int = 1024, feature_list: List | None = None) → ndarray[source]¶

Helper function to vectorize a spatial description from a hash.

Hash functions are used to perform spatial featurizations in DeepChem. However, it’s necessary to convert backwards from the hash function to feature vectors. This function aids in this conversion procedure. It creates a vector of zeros of length size. It then loops through feature_dict, uses hash_function to hash the stored value to an integer in range [0, size) and bumps that index.

Parameters:

hash_function (Function, Callable[[str, int], int]) – Should accept two arguments, feature, and size and return a hashed integer. Here feature is the item to hash, and size is an int. For example, if size=1024, then hashed values must fall in range [0, 1024).
feature_dict (Dict, optional (default None)) – Maps unique keys to features computed.
size (int (default 1024)) – Length of generated bit vector
feature_list (List, optional (default None)) – List of features.

Returns:

feature_vector – A numpy array of shape (size,)

Return type:

np.ndarray

Voxel Utils¶

convert_atom_to_voxel(coordinates: ndarray, atom_index: int, box_width: float, voxel_width: float) → ndarray[source]¶

Converts atom coordinates to an i,j,k grid index.

This function offsets molecular atom coordinates by (box_width/2, box_width/2, box_width/2) and then divides by voxel_width to compute the voxel indices.

Parameters:

coordinates (np.ndarray) – Array with coordinates of all atoms in the molecule, shape (N, 3).
atom_index (int) – Index of an atom in the molecule.
box_width (float) – Size of the box in Angstroms.
voxel_width (float) – Size of a voxel in Angstroms

Returns:

indices – A 1D numpy array of length 3 with [i, j, k], the voxel coordinates of specified atom.

Return type:

np.ndarray

convert_atom_pair_to_voxel(coordinates_tuple: Tuple[ndarray, ndarray], atom_index_pair: Tuple[int, int], box_width: float, voxel_width: float) → ndarray[source]¶

Converts a pair of atoms to i,j,k grid indexes.

Parameters:

coordinates_tuple (Tuple[np.ndarray, np.ndarray]) – A tuple containing two molecular coordinate arrays of shapes (N, 3) and (M, 3).
atom_index_pair (Tuple[int, int]) – A tuple of indices for the atoms in the two molecules.
box_width (float) – Size of the box in Angstroms.
voxel_width (float) – Size of a voxel in Angstroms

Returns:

indices_list – A numpy array of shape (2, 3), where 3 is [i, j, k] of the voxel coordinates of specified atom.

Return type:

np.ndarray

voxelize(get_voxels: Callable[[...], Any], coordinates: Any, box_width: float = 16.0, voxel_width: float = 1.0, hash_function: Callable[[...], Any] | None = None, feature_dict: Dict[Any, Any] | None = None, feature_list: List[int | Tuple[int]] | None = None, nb_channel: int = 16, dtype: str = 'int') → ndarray[source]¶

Helper function to voxelize inputs.

This helper function helps convert a hash function which specifies spatial features of a molecular complex into a voxel tensor. This utility is used by various featurizers that generate voxel grids.

Parameters:

get_voxels (Function) – Function that voxelizes inputs
coordinates (Any) – Contains the 3D coordinates of a molecular system. This should have whatever type get_voxels() expects as its first argument.
box_width (float, optional (default 16.0)) – Size of a box in which voxel features are calculated. Box is centered on a ligand centroid.
voxel_width (float, optional (default 1.0)) – Size of a 3D voxel in a grid in Angstroms.
hash_function (Function) – Used to map feature choices to voxel channels.
feature_dict (Dict, optional (default None)) – Keys are atom indices or tuples of atom indices, the values are computed features. If hash_function is not None, then the values are hashed using the hash function into [0, nb_channels) and this channel at the voxel for the given key is incremented by 1 for each dictionary entry. If hash_function is None, then the value must be a vector of size (n_channels,) which is added to the existing channel values at that voxel grid.
feature_list (List, optional (default None)) – List of atom indices or tuples of atom indices. This can only be used if nb_channel==1. Increments the voxels corresponding to these indices by 1 for each entry.
nb_channel (int, , optional (default 16)) – The number of feature channels computed per voxel. Should be a power of 2.
dtype (str ('int' or 'float'), optional (default 'int')) – The type of the numpy ndarray created to hold features.

Returns:

feature_tensor – The voxel of the input with the shape (voxels_per_edge, voxels_per_edge, voxels_per_edge, nb_channel).

Return type:

np.ndarray

Graph Convolution Utilities¶

one_hot_encode(val: int | str, allowable_set: List[str] | List[int], include_unknown_set: bool = False) → List[float][source]¶

One hot encoder for elements of a provided set.

Examples

>>> one_hot_encode("a", ["a", "b", "c"])
[1.0, 0.0, 0.0]
>>> one_hot_encode(2, [0, 1, 2])
[0.0, 0.0, 1.0]
>>> one_hot_encode(3, [0, 1, 2])
[0.0, 0.0, 0.0]
>>> one_hot_encode(3, [0, 1, 2], True)
[0.0, 0.0, 0.0, 1.0]

Parameters:

val (int or str) – The value must be present in allowable_set.
allowable_set (List[int] or List[str]) – List of allowable quantities.
include_unknown_set (bool, default False) – If true, the index of all values not in allowable_set is len(allowable_set).

Returns:

An one-hot vector of val. If include_unknown_set is False, the length is len(allowable_set). If include_unknown_set is True, the length is len(allowable_set) + 1.

Return type:

List[float]

Raises:

ValueError – If include_unknown_set is False and val is not in allowable_set.

get_atom_type_one_hot(atom: Any, allowable_set: List[str] = ['C', 'N', 'O', 'F', 'P', 'S', 'Cl', 'Br', 'I'], include_unknown_set: bool = True) → List[float][source]¶

Get an one-hot feature of an atom type.

Parameters:

atom (rdkit.Chem.rdchem.Atom) – RDKit atom object
allowable_set (List[str]) – The atom types to consider. The default set is [“C”, “N”, “O”, “F”, “P”, “S”, “Cl”, “Br”, “I”].
include_unknown_set (bool, default True) – If true, the index of all atom not in allowable_set is len(allowable_set).

Returns:

An one-hot vector of atom types. If include_unknown_set is False, the length is len(allowable_set). If include_unknown_set is True, the length is len(allowable_set) + 1.

Return type:

List[float]

construct_hydrogen_bonding_info(mol: Any) → List[Tuple[int, str]][source]¶

Construct hydrogen bonding infos about a molecule.

Parameters:: mol (rdkit.Chem.rdchem.Mol) – RDKit mol object
Returns:: A list of tuple (atom_index, hydrogen_bonding_type). The hydrogen_bonding_type value is “Acceptor” or “Donor”.
Return type:: List[Tuple[int, str]]

get_atom_hydrogen_bonding_one_hot(atom: Any, hydrogen_bonding: List[Tuple[int, str]]) → List[float][source]¶

Get an one-hot feat about whether an atom accepts electrons or donates electrons.

Parameters:

atom (rdkit.Chem.rdchem.Atom) – RDKit atom object
hydrogen_bonding (List[Tuple[int, str]]) – The return value of construct_hydrogen_bonding_info. The value is a list of tuple (atom_index, hydrogen_bonding) like (1, “Acceptor”).

Returns:

A one-hot vector of the ring size type. The first element indicates “Donor”, and the second element indicates “Acceptor”.

Return type:

List[float]

get_atom_is_in_aromatic_one_hot(atom: Any) → List[float][source]¶

Get ans one-hot feature about whether an atom is in aromatic system or not.

Parameters:: atom (rdkit.Chem.rdchem.Atom) – RDKit atom object
Returns:: A vector of whether an atom is in aromatic system or not.
Return type:: List[float]

get_atom_hybridization_one_hot(atom: Any, allowable_set: List[str] = ['SP', 'SP2', 'SP3'], include_unknown_set: bool = False) → List[float][source]¶

Get an one-hot feature of hybridization type.

Parameters:

atom (rdkit.Chem.rdchem.Atom) – RDKit atom object
allowable_set (List[str]) – The hybridization types to consider. The default set is [“SP”, “SP2”, “SP3”]
include_unknown_set (bool, default False) – If true, the index of all types not in allowable_set is len(allowable_set).

Returns:

An one-hot vector of the hybridization type. If include_unknown_set is False, the length is len(allowable_set). If include_unknown_set is True, the length is len(allowable_set) + 1.

Return type:

List[float]

get_atom_total_num_Hs_one_hot(atom: Any, allowable_set: List[int] = [0, 1, 2, 3, 4], include_unknown_set: bool = True) → List[float][source]¶

Get an one-hot feature of the number of hydrogens which an atom has.

Parameters:

atom (rdkit.Chem.rdchem.Atom) – RDKit atom object
allowable_set (List[int]) – The number of hydrogens to consider. The default set is [0, 1, …, 4]
include_unknown_set (bool, default True) – If true, the index of all types not in allowable_set is len(allowable_set).

Returns:

A one-hot vector of the number of hydrogens which an atom has. If include_unknown_set is False, the length is len(allowable_set). If include_unknown_set is True, the length is len(allowable_set) + 1.

Return type:

List[float]

get_atom_chirality_one_hot(atom: Any) → List[float][source]¶

Get an one-hot feature about an atom chirality type.

Parameters:: atom (rdkit.Chem.rdchem.Atom) – RDKit atom object
Returns:: A one-hot vector of the chirality type. The first element indicates “R”, and the second element indicates “S”.
Return type:: List[float]

get_atom_formal_charge(atom: Any) → List[float][source]¶

Get a formal charge of an atom.

Parameters:: atom (rdkit.Chem.rdchem.Atom) – RDKit atom object
Returns:: A vector of the formal charge.
Return type:: List[float]

get_atom_partial_charge(atom: Any) → List[float][source]¶

Get a partial charge of an atom.

Parameters:: atom (rdkit.Chem.rdchem.Atom) – RDKit atom object
Returns:: A vector of the parital charge.
Return type:: List[float]

Notes

Before using this function, you must calculate GasteigerCharge like AllChem.ComputeGasteigerCharges(mol).

get_atom_total_degree_one_hot(atom: Any, allowable_set: List[int] = [0, 1, 2, 3, 4, 5], include_unknown_set: bool = True) → List[float][source]¶

Get an one-hot feature of the degree which an atom has.

Parameters:

atom (rdkit.Chem.rdchem.Atom) – RDKit atom object
allowable_set (List[int]) – The degree to consider. The default set is [0, 1, …, 5]
include_unknown_set (bool, default True) – If true, the index of all types not in allowable_set is len(allowable_set).

Returns:

A one-hot vector of the degree which an atom has. If include_unknown_set is False, the length is len(allowable_set). If include_unknown_set is True, the length is len(allowable_set) + 1.

Return type:

List[float]

get_bond_type_one_hot(bond: Any, allowable_set: List[str] = ['SINGLE', 'DOUBLE', 'TRIPLE', 'AROMATIC'], include_unknown_set: bool = False) → List[float][source]¶

Get an one-hot feature of bond type.

Parameters:

bond (rdkit.Chem.rdchem.Bond) – RDKit bond object
allowable_set (List[str]) – The bond types to consider. The default set is [“SINGLE”, “DOUBLE”, “TRIPLE”, “AROMATIC”].
include_unknown_set (bool, default False) – If true, the index of all types not in allowable_set is len(allowable_set).

Returns:

A one-hot vector of the bond type. If include_unknown_set is False, the length is len(allowable_set). If include_unknown_set is True, the length is len(allowable_set) + 1.

Return type:

List[float]

get_bond_is_in_same_ring_one_hot(bond: Any) → List[float][source]¶

Get an one-hot feature about whether atoms of a bond is in the same ring or not.

Parameters:: bond (rdkit.Chem.rdchem.Bond) – RDKit bond object
Returns:: A one-hot vector of whether a bond is in the same ring or not.
Return type:: List[float]

get_bond_is_conjugated_one_hot(bond: Any) → List[float][source]¶

Get an one-hot feature about whether a bond is conjugated or not.

Parameters:: bond (rdkit.Chem.rdchem.Bond) – RDKit bond object
Returns:: A one-hot vector of whether a bond is conjugated or not.
Return type:: List[float]

get_bond_stereo_one_hot(bond: Any, allowable_set: List[str] = ['STEREONONE', 'STEREOANY', 'STEREOZ', 'STEREOE'], include_unknown_set: bool = True) → List[float][source]¶

Get an one-hot feature of the stereo configuration of a bond.

Parameters:

bond (rdkit.Chem.rdchem.Bond) – RDKit bond object
allowable_set (List[str]) – The stereo configuration types to consider. The default set is [“STEREONONE”, “STEREOANY”, “STEREOZ”, “STEREOE”].
include_unknown_set (bool, default True) – If true, the index of all types not in allowable_set is len(allowable_set).

Returns:

A one-hot vector of the stereo configuration of a bond. If include_unknown_set is False, the length is len(allowable_set). If include_unknown_set is True, the length is len(allowable_set) + 1.

Return type:

List[float]

get_bond_graph_distance_one_hot(bond: Any, graph_dist_matrix: ndarray, allowable_set: List[int] = [1, 2, 3, 4, 5, 6, 7], include_unknown_set: bool = True) → List[float][source]¶

Get an one-hot feature of graph distance.

Parameters:

bond (rdkit.Chem.rdchem.Bond) – RDKit bond object
graph_dist_matrix (np.ndarray) – The return value of Chem.GetDistanceMatrix(mol). The shape is (num_atoms, num_atoms).
allowable_set (List[int]) – The graph distance types to consider. The default set is [1, 2, …, 7].
include_unknown_set (bool, default False) – If true, the index of all types not in allowable_set is len(allowable_set).

Returns:

A one-hot vector of the graph distance. If include_unknown_set is False, the length is len(allowable_set). If include_unknown_set is True, the length is len(allowable_set) + 1.

Return type:

List[float]

Grover Utilities¶

extract_grover_attributes(molgraph: BatchGraphData)[source]¶

Utility to extract grover attributes for grover model

Parameters:: molgraph (BatchGraphData) – A batched graph data representing a collection of molecules.
Returns:: graph_attributes – A tuple containing atom features, bond features, atom to bond mapping, bond to atom mapping, bond to reverse bond mapping, atom to atom mapping, atom scope, bond scope, functional group labels and other additional features.
Return type:: Tuple

Example

>>> import deepchem as dc
>>> from deepchem.feat.graph_data import BatchGraphData, GraphData
>>> import numpy as np
>>> batch_graphs = []
>>> smiles = ['CC', 'CCC', 'CC(=O)C']
>>> featurizer = dc.feat.GroverFeaturizer(features_generator=dc.feat.CircularFingerprint())
>>> graphs = featurizer.featurize(smiles)
>>> for g in graphs:
...     batch_graphs.append(GraphData(
...        g.node_features,
...        g.edge_index,
...        g.edge_features,
...        fg_labels=np.asarray(g.fg_labels),
...        additional_features=np.asarray(g.additional_features)
...    ))
>>> molgraph = BatchGraphData(batch_graphs)
>>> attributes = extract_grover_attributes(molgraph)

Debug Utilities¶

Docking Utilities¶

These utilities assist in file preparation and processing for molecular docking.

write_vina_conf(protein_filename: str, ligand_filename: str, centroid: ndarray, box_dims: ndarray, conf_filename: str, num_modes: int = 9, exhaustiveness: int | None = None) → None[source]¶

Writes Vina configuration file to disk.

Autodock Vina accepts a configuration file which provides options under which Vina is invoked. This utility function writes a vina configuration file which directs Autodock vina to perform docking under the provided options.

Parameters:

protein_filename (str) – Filename for protein
ligand_filename (str) – Filename for the ligand
centroid (np.ndarray) – A numpy array with shape (3,) holding centroid of system
box_dims (np.ndarray) – A numpy array of shape (3,) holding the size of the box to dock
conf_filename (str) – Filename to write Autodock Vina configuration to.
num_modes (int, optional (default 9)) – The number of binding modes Autodock Vina should find
exhaustiveness (int, optional) – The exhaustiveness of the search to be performed by Vina

write_gnina_conf(protein_filename: str, ligand_filename: str, conf_filename: str, num_modes: int = 9, exhaustiveness: int | None = None, **kwargs) → None[source]¶

Writes GNINA configuration file to disk.

GNINA accepts a configuration file which provides options under which GNINA is invoked. This utility function writes a configuration file which directs GNINA to perform docking under the provided options.

Parameters:

protein_filename (str) – Filename for protein
ligand_filename (str) – Filename for the ligand
conf_filename (str) – Filename to write Autodock Vina configuration to.
num_modes (int, optional (default 9)) – The number of binding modes GNINA should find
exhaustiveness (int, optional) – The exhaustiveness of the search to be performed by GNINA
kwargs – Args supported by GNINA documented here https://github.com/gnina/gnina#usage

load_docked_ligands(pdbqt_output: str) → Tuple[List[Any], List[float]][source]¶

This function loads ligands docked by autodock vina.

Autodock vina writes outputs to disk in a PDBQT file format. This PDBQT file can contain multiple docked “poses”. Recall that a pose is an energetically favorable 3D conformation of a molecule. This utility function reads and loads the structures for multiple poses from vina’s output file.

Parameters:: pdbqt_output (str) – Should be the filename of a file generated by autodock vina’s docking software.
Returns:: Tuple of molecules, scores. molecules is a list of rdkit molecules with 3D information. scores is the associated vina score.
Return type:: Tuple[List[rdkit.Chem.rdchem.Mol], List[float]]

Notes

This function requires RDKit to be installed.

prepare_inputs(protein: str, ligand: str, replace_nonstandard_residues: bool = True, remove_heterogens: bool = True, remove_water: bool = True, add_hydrogens: bool = True, pH: float = 7.0, optimize_ligand: bool = True, pdb_name: str | None = None) → Tuple[Any, Any][source]¶

This prepares protein-ligand complexes for docking.

Autodock Vina requires PDB files for proteins and ligands with sensible inputs. This function uses PDBFixer and RDKit to ensure that inputs are reasonable and ready for docking. Default values are given for convenience, but fixing PDB files is complicated and human judgement is required to produce protein structures suitable for docking. Always inspect the results carefully before trying to perform docking.

Parameters:

protein (str) – Filename for protein PDB file or a PDBID.
ligand (str) – Either a filename for a ligand PDB file or a SMILES string.
replace_nonstandard_residues (bool (default True)) – Replace nonstandard residues with standard residues.
remove_heterogens (bool (default True)) – Removes residues that are not standard amino acids or nucleotides.
remove_water (bool (default True)) – Remove water molecules.
add_hydrogens (bool (default True)) – Add missing hydrogens at the protonation state given by pH.
pH (float (default 7.0)) – Most common form of each residue at given pH value is used.
optimize_ligand (bool (default True)) – If True, optimize ligand with RDKit. Required for SMILES inputs.
pdb_name (Optional[str]) – If given, write sanitized protein and ligand to files called “pdb_name.pdb” and “ligand_pdb_name.pdb”

Returns:

Tuple of protein_molecule, ligand_molecule with 3D information.

Return type:

Tuple[RDKitMol, RDKitMol]

Note

This function requires RDKit and OpenMM to be installed. Read more about PDBFixer here: https://github.com/openmm/pdbfixer.

Examples

>>> p, m = prepare_inputs('3cyx', 'CCC')

>> p.GetNumAtoms() >> m.GetNumAtoms()

>>> p, m = prepare_inputs('3cyx', 'CCC', remove_heterogens=False)

>> p.GetNumAtoms()

read_gnina_log(log_file: str) → ndarray[source]¶

Read GNINA logfile and get docking scores.

GNINA writes computed binding affinities to a logfile.

Parameters:: log_file (str) – Filename of logfile generated by GNINA.
Returns:: scores – Array of binding affinity (kcal/mol), CNN pose score, and CNN affinity for each binding mode.
Return type:: np.array, dimension (num_modes, 3)

Print Threshold¶

The printing threshold controls how many dataset elements are printed when dc.data.Dataset objects are converted to strings or represnted in the IPython repl.

get_print_threshold() → int[source]¶

Return the printing threshold for datasets.

The print threshold is the number of elements from ids/tasks to print when printing representations of Dataset objects.

Returns:: threshold – Number of elements that will be printed
Return type:: int

set_print_threshold(threshold: int)[source]¶

Set print threshold

The print threshold is the number of elements from ids/tasks to print when printing representations of Dataset objects.

Parameters:: threshold (int) – Number of elements to print.

get_max_print_size() → int[source]¶

Return the max print size for a dataset.

If a dataset is large, printing self.ids as part of a string representation can be very slow. This field controls the maximum size for a dataset before ids are no longer printed.

Returns:: max_print_size – Maximum length of a dataset for ids to be printed in string representation.
Return type:: int

set_max_print_size(max_print_size: int)[source]¶

Set max_print_size

If a dataset is large, printing self.ids as part of a string representation can be very slow. This field controls the maximum size for a dataset before ids are no longer printed.

Parameters:: max_print_size (int) – Maximum length of a dataset for ids to be printed in string representation.

Fake Data Generator¶

The utilities here are used to generate random sample data which can be used for testing model architectures or other purposes.

class FakeGraphGenerator(min_nodes: int = 10, max_nodes: int = 10, n_node_features: int = 5, avg_degree: int = 4, n_edge_features: int = 3, n_classes: int = 2, task: str = 'graph', **kwargs)[source]¶

Generates a random graphs which can be used for testing or other purposes.

The generated graph supports both node-level and graph-level labels.

Example

>>> from deepchem.utils.fake_data_generator import FakeGraphGenerator
>>> fgg  = FakeGraphGenerator(min_nodes=8, max_nodes=10,  n_node_features=5, avg_degree=8, n_edge_features=3, n_classes=2, task='graph', z=5)
>>> graphs = fgg.sample(n_graphs=10)
>>> type(graphs)
<class 'deepchem.data.datasets.NumpyDataset'>
>>> type(graphs.X[0])
<class 'deepchem.feat.graph_data.GraphData'>
>>> len(graphs) == 10  # num_graphs
True

Note

The FakeGraphGenerator class is based on torch_geometric.dataset.FakeDataset class.

__init__(min_nodes: int = 10, max_nodes: int = 10, n_node_features: int = 5, avg_degree: int = 4, n_edge_features: int = 3, n_classes: int = 2, task: str = 'graph', **kwargs)[source]¶

Parameters:

min_nodes (int, default 10) – Minimum number of permissible nodes in a graph
max_nodes (int, default 10) – Maximum number of permissible nodes in a graph
n_node_features (int, default 5) – Average number of node features in a graph
avg_degree (int, default 4) – Average degree of the graph (avg_degree should be a positive number greater than the min_nodes)
n_edge_features (int, default 3) – Average number of features in the edge
task (str, default 'graph') – Indicates node-level labels or graph-level labels
kwargs (optional) – Additional graph attributes and their shapes , e.g. global_features = 5

sample(n_graphs: int = 100) → NumpyDataset[source]¶

Samples graphs

Parameters:: n_graphs (int, default 100) – Number of graphs to generate
Returns:: graphs – Generated Graphs
Return type:: NumpyDataset

Electron Sampler¶

The utilities here are used to sample electrons in a given molecule and update it using monte carlo methods, which can be used for methods like Variational Monte Carlo, etc.

class ElectronSampler(central_value: ndarray, f: Callable[[ndarray], ndarray], batch_no: int = 10, x: ndarray = array([], dtype=float64), steps: int = 200, steps_per_update: int = 10, seed: int | None = None, symmetric: bool = True, simultaneous: bool = True)[source]¶

This class enables to initialize electron’s position using gauss distribution around a nucleus and update using Markov Chain Monte-Carlo(MCMC) moves.

Using the probability obtained from the square of magnitude of wavefunction of a molecule/atom, MCMC steps can be performed to get the electron’s positions and further update the wavefunction. This method is primarily used in methods like Variational Monte Carlo to sample electrons around the nucleons. Sampling can be done in 2 ways: -Simultaneous: All the electrons’ positions are updated all at once.

-Single-electron: MCMC steps are performed only a particular electron, given their index value.

Further these moves can be done in 2 methods: -Symmetric: In this configuration, the standard deviation for all the steps are uniform.

-Asymmetric: In this configuration, the standard deviation are not uniform and typically the standard deviation is obtained a function like harmonic distances, etc.

Irrespective of these methods, the initialization is done uniformly around the respective nucleus and the number of electrons specified.

Example

>>> from deepchem.utils.electron_sampler import ElectronSampler
>>> test_f = lambda x: 2*np.log(np.random.uniform(low=0,high=1.0,size=np.shape(x)[0]))
>>> distribution=ElectronSampler(central_value=np.array([[1,1,3],[3,2,3]]),f=test_f,seed=0,batch_no=2,steps=1000,)
>>> distribution.gauss_initialize_position(np.array([[1],[2]]))

>> print(distribution.x) [[[[1.03528105 1.00800314 3.01957476]]

[[3.01900177 1.99697286 2.99793562]]

[[3.00821197 2.00288087 3.02908547]]]

[[[1.04481786 1.03735116 2.98045444]]

[[3.01522075 2.0024335 3.00887726]]

[[3.00667349 2.02988158 2.99589683]]]]

>>> distribution.move()
0.5115

>> print(distribution.x) [[[[-0.32441754 1.23330263 2.67927645]]

[[ 3.42250997 2.23617126 3.55806632]]

[[ 3.37491385 1.54374006 3.13575241]]]

[[[ 0.49067726 1.03987841 3.70277884]]

[[ 3.5631939 1.68703947 2.5685874 ]]

[[ 2.84560249 1.73998364 3.41274181]]]]

__init__(central_value: ndarray, f: Callable[[ndarray], ndarray], batch_no: int = 10, x: ndarray = array([], dtype=float64), steps: int = 200, steps_per_update: int = 10, seed: int | None = None, symmetric: bool = True, simultaneous: bool = True)[source]¶

Parameters:

central_value (np.ndarray) – Contains each nucleus’ coordinates in a 2D array. The shape of the array should be(number_of_nucleus,3).Ex: [[1,2,3],[3,4,5],..]
f (Callable[[np.ndarray],np.ndarray]) – A function that should give the twice the log probability of wavefunction of the molecular system when called. Should taken in a 4D array of electron’s positions(x) as argument and return a numpy array containing the log probabilities of each batch.
batch_no (int, optional (default 10)) – Number of batches of the electron’s positions to be initialized.
x (np.ndarray, optional (default np.ndarray([]))) – Contains the electron’s coordinates in a 4D array. The shape of the array should be(batch_no,no_of_electrons,1,3). Can be a 1D empty array, when electron’s positions are yet to be initialized.
steps (int, optional (default 10)) – The number of MCMC steps to be performed when the moves are called.
steps_per_update (int (default 10)) – The number of steps after which the parameters of the MCMC gets updated.
seed (int, optional (default None)) – Random seed to use.
symmetric (bool, optional(default True)) – If true, symmetric moves will be used, else asymmetric moves will be followed.
simultaneous (bool, optional(default True)) – If true, MCMC steps will be performed on all the electrons, else only a single electron gets updated.

sampled_electrons[source]¶

Keeps track of the sampled electrons at every step, must be empty at start.

Type:: np.ndarray

harmonic_mean(y: ndarray) → ndarray[source]¶

Calculates the harmonic mean of the value ‘y’ from the self.central value. The numpy array returned is typically scaled up to get the standard deviation matrix.

Parameters:: y (np.ndarray) – Containing the data distribution. Shape of y should be (batch,no_of_electron,1,3)
Returns:: Contains the harmonic mean of the data distribution of each batch. Shape of the array obtained (batch_no, no_of_electrons,1,1)
Return type:: np.ndarray

log_prob_gaussian(y: ndarray, mu: ndarray, sigma: ndarray) → ndarray[source]¶

Calculates the log probability of a gaussian distribution, given the mean and standard deviation

Parameters:

y (np.ndarray) – data for which the log normal distribution is to be found
mu (np.ndarray) – Means wrt which the log normal is calculated. Same shape as x or should be brodcastable to x
sigma (np.ndarray,) – The standard deviation of the log normal distribution. Same shape as x or should be brodcastable to x

Returns:

Log probability of gaussian distribution, with the shape - (batch_no,).

Return type:

np.ndarray

gauss_initialize_position(no_sample: ndarray, stddev: float = 0.02)[source]¶: Initializes the position around a central value as mean sampled from a gauss distribution and updates self.x. :param no_sample: Contains the number of samples to initialize under each mean. should be in the form [[3],[2]..], where here it means 3 samples and 2 samples around the first entry and second entry,respectively in self.central_value is taken. :type no_sample: np.ndarray, :param stddev: contains the stddev with which the electrons’ coordinates are initialized :type stddev: float, optional (default 0.02)

electron_update(lp1, lp2, move_prob, ratio, x2) → ndarray[source]¶

Performs sampling & parameter updates of electrons and appends the sampled electrons to self.sampled_electrons.

Parameters:

lp1 (np.ndarray) – Log probability of initial parameter state.
lp2 (np.ndarray) – Log probability of the new sampled state.
move_prob (np.ndarray) – Sampled log probabilty of the electron moving from the initial to final state, sampled assymetrically or symetrically.
ratio (np.ndarray) – Ratio of lp1 and lp2 state.
x2 (np.ndarray) – Numpy array of the new sampled electrons.

Returns:

lp1 – The update log probability of initial parameter state.

Return type:

np.ndarray

move(stddev: float = 0.02, asymmetric_func: Callable[[ndarray], ndarray] | None = None, index: int | None = None) → float[source]¶

Performs Metropolis-Hasting move for self.x(electrons). The type of moves to be followed -(simultaneous or single-electron, symmetric or asymmetric) have been specified when calling the class. The self.x array is replaced with a new array at the end of each step containing the new electron’s positions.

Parameters:

asymmetric_func (Callable[[np.ndarray],np.ndarray], optional(default None)) – Should be specified for an asymmetric move.The function should take in only 1 argument- y: a numpy array wrt to which mean should be calculated. This function should return the mean for the asymmetric proposal. For ferminet, this function is the harmonic mean of the distance between the electron and the nucleus.
stddev (float, optional (default 0.02)) – Specifies the standard deviation in the case of symmetric moves and the scaling factor of the standard deviation matrix in the case of asymmetric moves.
index (int, optional (default None)) – Specifies the index of the electron to be updated in the case of a single electron move.

Returns:

accepted move ratio of the MCMC steps.

Return type:

float

Density Functional Theory Utilities¶

The utilites here are used to create an object that contains information about a system’s self-consistent iteration steps and other processes.

class Lattice(a: Tensor)[source]¶

Lattice is an object that describe the periodicity of the lattice. Note that this object does not know about atoms. For the integrated object between the lattice and atoms, please see Sol

Examples

>>> import torch
>>> from deepchem.utils.dft_utils import Lattice
>>> a = torch.tensor([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])
>>> lattice = Lattice(a)
>>> lattice.lattice_vectors()
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]])
>>> lattice.recip_vectors()
tensor([[6.2832, 0.0000, 0.0000],
        [0.0000, 6.2832, 0.0000],
        [0.0000, 0.0000, 6.2832]])
>>> lattice.volume() # volume of the unit cell
tensor(1.)
>>> lattice.get_lattice_ls(1.0) # get the neighboring lattice vectors
tensor([[ 0.,  0., -1.],
        [ 0., -1.,  0.],
        [-1.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 1.,  0.,  0.],
        [ 0.,  1.,  0.],
        [ 0.,  0.,  1.]])
>>> lattice.get_gvgrids(6.0) # get the neighboring G-vectors
(tensor([[ 0.0000,  0.0000, -6.2832],
        [ 0.0000, -6.2832,  0.0000],
        [-6.2832,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000],
        [ 6.2832,  0.0000,  0.0000],
        [ 0.0000,  6.2832,  0.0000],
        [ 0.0000,  0.0000,  6.2832]]), tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]))
>>> lattice.estimate_ewald_eta(1e-5) # estimate the ewald's sum eta
1.8

__init__(a: Tensor)[source]¶

Initialize the lattice object.

2D or 1D repetition are not implemented yet

Parameters:: a (torch.Tensor) – The lattice vectors with shape (ndim, ndim) with ndim == 3

lattice_vectors() → Tensor[source]¶: Returns the 3D lattice vectors (nv, ndim) with nv == 3

recip_vectors() → Tensor[source]¶

Returns the 3D reciprocal vectors with norm == 2 * pi with shape (nv, ndim) with nv == 3

Note: torch.det(self.a) should not be equal to zero.

volume() → Tensor[source]¶: Returns the volume of a lattice.

property params: Tuple[Tensor, ...][source]¶: Returns the list of parameters of this object

get_lattice_ls(rcut: float, exclude_zeros: bool = False) → Tensor[source]¶

Returns a tensor that contains the coordinates of the neighboring lattices.

Parameters:

rcut (float) – The threshold of the distance from the main cell to be included in the neighbor.
exclude_zeros (bool (default: False)) – If True, then it will exclude the vector that are all zeros.

Returns:

ls – Tensor with size (nb, ndim) containing the coordinates of the neighboring cells.

Return type:

torch.Tensor

get_gvgrids(gcut: float, exclude_zeros: bool = False) → Tuple[Tensor, Tensor][source]¶

Returns a tensor that contains the coordinate in reciprocal space of the neighboring Brillouin zones.

Parameters:

gcut (float) – Cut off for generating the G-points.
exclude_zeros (bool (default: False)) – If True, then it will exclude the vector that are all zeros.

Returns:

gvgrids (torch.Tensor) – Tensor with size (ng, ndim) containing the G-coordinates of the Brillouin zones.
weights (torch.Tensor) – Tensor with size (ng) representing the weights of the G-points.

estimate_ewald_eta(precision: float) → float[source]¶

estimate the ewald’s sum eta for nuclei interaction energy the precision is assumed to be relative precision this formula is obtained by estimating the sum as an integral.

Parameters:: precision (float) – The precision of the ewald’s sum.
Returns:: eta – The estimated eta.
Return type:: float

class SpinParam(u: T, d: T)[source]¶

Data structure to store different values for spin-up and spin-down electrons.

Examples

>>> import torch
>>> from deepchem.utils.dft_utils import SpinParam
>>> dens_u = torch.ones(1)
>>> dens_d = torch.zeros(1)
>>> sp = SpinParam(u=dens_u, d=dens_d)
>>> sp.u
tensor([1.])
>>> sp.sum()
tensor([1.])
>>> sp.reduce(torch.multiply)
tensor([0.])

__init__(u: T, d: T)[source]¶

Initialize the SpinParam object.

Parameters:

u (any type) – The parameters that corresponds to the spin-up electrons.
d (any type) – The parameters that corresponds to the spin-down electrons.

sum() → Any[source]¶: Returns the sum of up and down parameters.

reduce(fcn: Callable[[T, T], T]) → T[source]¶: Reduce up and down parameters with the given function.

static apply_fcn(fcn: Callable[[...], P], *a)[source]¶: “Apply the function for each up and down elements of a

class ValGrad(value: Tensor, grad: Tensor | None = None, lapl: Tensor | None = None, kin: Tensor | None = None)[source]¶

Data structure that contains local information about density profiles. Data structure used as a umbrella class for density profiles and the derivative of the potential w.r.t. density profiles.

Examples

>>> import torch
>>> from deepchem.utils.dft_utils import ValGrad
>>> dens = torch.ones(1)
>>> grad = torch.zeros(1)
>>> lapl = torch.ones(1)
>>> kin = torch.ones(1)
>>> vg = ValGrad(value=dens, grad=grad, lapl=lapl, kin=kin)
>>> vg + vg
ValGrad(value=tensor([2.]), grad=tensor([0.]), lapl=tensor([2.]), kin=tensor([2.]))
>>> vg * 5
ValGrad(value=tensor([5.]), grad=tensor([0.]), lapl=tensor([5.]), kin=tensor([5.]))

__init__(value: Tensor, grad: Tensor | None = None, lapl: Tensor | None = None, kin: Tensor | None = None)[source]¶

Initialize the ValGrad object.

Parameters:

value (torch.Tensor) – Tensors containing the value of the local information.
grad (torch.Tensor or None) – If tensor, it represents the gradient of the local information with shape (..., 3) where ... should be the same shape as value.
lapl (torch.Tensor or None) – If tensor, represents the laplacian value of the local information. It should have the same shape as value.
kin (torch.Tensor or None) – If tensor, represents the local kinetic energy density. It should have the same shape as value.

__add__(b)[source]¶: Add two ValGrad objects together.

__mul__(f: float | int | Tensor)[source]¶: Multiply the ValGrad object with a scalar.

class CGTOBasis(angmom: int, alphas: Tensor, coeffs: Tensor, normalized: bool = False)[source]¶

Data structure that contains information about a contracted gaussian type orbital (CGTO).

Examples

>>> import torch
>>> from deepchem.utils.dft_utils import CGTOBasis
>>> alphas = torch.ones(1)
>>> coeffs = torch.ones(1)
>>> cgto = CGTOBasis(angmom=0, alphas=alphas, coeffs=coeffs)
>>> cgto.wfnormalize_()
CGTOBasis(angmom=0, alphas=tensor([1.]), coeffs=tensor([2.5265]), normalized=True)

__init__(angmom: int, alphas: Tensor, coeffs: Tensor, normalized: bool = False)[source]¶

Initialize the CGTOBasis object.

Parameters:

angmom (int) – The angular momentum of the basis.
alphas (torch.Tensor) – The gaussian exponents of the basis. Shape: (nbasis,)
coeffs (torch.Tensor) – The coefficients of the basis. Shape: (nbasis,)

wfnormalize_() → CGTOBasis[source]¶

Wavefunction normalization

The normalization is obtained from CINTgto_norm from libcint/src/misc.c, or https://github.com/sunqm/libcint/blob/b8594f1d27c3dad9034984a2a5befb9d607d4932/src/misc.c#L80

Please note that the square of normalized wavefunctions do not integrate to 1, but e.g. for s: 4*pi, p: (4*pi/3)

class AtomCGTOBasis(atomz: int | float | Tensor, bases: List[CGTOBasis], pos: List[List[float]] | ndarray | Tensor)[source]¶

Data structure that contains information about a atom and its contracted gaussian type orbital (CGTO).

Examples

>>> import torch
>>> from deepchem.utils.dft_utils import AtomCGTOBasis, CGTOBasis
>>> alphas = torch.ones(1)
>>> coeffs = torch.ones(1)
>>> cgto = CGTOBasis(angmom=0, alphas=alphas, coeffs=coeffs)
>>> atomcgto = AtomCGTOBasis(atomz=1, bases=[cgto], pos=[[0.0, 0.0, 0.0]])
>>> atomcgto
AtomCGTOBasis(atomz=1, bases=[CGTOBasis(angmom=0, alphas=tensor([1.]), coeffs=tensor([1.]), normalized=False)], pos=tensor([[0., 0., 0.]]))

__init__(atomz: int | float | Tensor, bases: List[CGTOBasis], pos: List[List[float]] | ndarray | Tensor)[source]¶

Initialize the AtomCGTOBasis object.

Parameters:

atomz (ZType) – Atomic number of the atom.
bases (List[CGTOBasis]) – List of CGTOBasis objects.
pos (AtomPosType) – Position of the atom. Shape: (ndim,)

class MulBaseXC(a: BaseXC, b: float | Tensor)[source]¶

Multiply a BaseXC with a float or a tensor

Examples

>>> import torch
>>> from deepchem.utils.dft_utils import ValGrad, SpinParam
>>> from deepchem.utils.dft_utils import BaseXC, MulBaseXC
>>> class MyXC(BaseXC):
...     @property
...     def family(self) -> int:
...         return 1
...     def get_edensityxc(self, densinfo: Union[ValGrad, SpinParam[ValGrad]]) -> torch.Tensor:
...         if isinstance(densinfo, ValGrad):
...             return densinfo.value.pow(2)
...         else:
...             return densinfo.u.value.pow(2) + densinfo.d.value.pow(2)
...     def get_vxc(self, densinfo: Union[ValGrad, SpinParam[ValGrad]]) -> Union[ValGrad, SpinParam[ValGrad]]:
...         if isinstance(densinfo, ValGrad):
...             return ValGrad(value=2*densinfo.value)
...         else:
...             return SpinParam(u=ValGrad(value=2*densinfo.u.value),
...                              d=ValGrad(value=2*densinfo.d.value))
>>> xc = MyXC()
>>> densinfo = ValGrad(value=torch.tensor([1., 2., 3.], requires_grad=True))
>>> xc.get_edensityxc(densinfo)
tensor([1., 4., 9.], grad_fn=<PowBackward0>)
>>> xc.get_vxc(densinfo)
ValGrad(value=tensor([2., 4., 6.], grad_fn=<MulBackward0>), grad=None, lapl=None, kin=None)
>>> densinfo = SpinParam(u=ValGrad(value=torch.tensor([1., 2., 3.], requires_grad=True)),
...                      d=ValGrad(value=torch.tensor([4., 5., 6.], requires_grad=True)))
>>> xc.get_edensityxc(densinfo)
tensor([17., 29., 45.], grad_fn=<AddBackward0>)
>>> xc.get_vxc(densinfo)
SpinParam(u=ValGrad(value=tensor([2., 4., 6.], grad_fn=<MulBackward0>), grad=None, lapl=None, kin=None), d=ValGrad(value=tensor([ 8., 10., 12.], grad_fn=<MulBackward0>), grad=None, lapl=None, kin=None))
>>> xc2 = MulBaseXC(xc, 2.)
>>> xc2.get_edensityxc(densinfo)
tensor([34., 58., 90.], grad_fn=<MulBackward0>)
>>> xc2.get_vxc(densinfo)
SpinParam(u=ValGrad(value=tensor([ 4.,  8., 12.], grad_fn=<MulBackward0>), grad=None, lapl=None, kin=None), d=ValGrad(value=tensor([16., 20., 24.], grad_fn=<MulBackward0>), grad=None, lapl=None, kin=None))
>>> xc3 = xc * 2.
>>> xc3.get_edensityxc(densinfo)
tensor([34., 58., 90.], grad_fn=<MulBackward0>)
>>> xc3.get_vxc(densinfo)
SpinParam(u=ValGrad(value=tensor([ 4.,  8., 12.], grad_fn=<MulBackward0>), grad=None, lapl=None, kin=None), d=ValGrad(value=tensor([16., 20., 24.], grad_fn=<MulBackward0>), grad=None, lapl=None, kin=None))

__init__(a: BaseXC, b: float | Tensor) → None[source]¶

Initialize the MulBaseXC

Parameters:

a (BaseXC) – BaseXC to be multiplied to.
b (Union[float, torch.Tensor]) – float or tensor to be multiplied with.

property family[source]¶: Returns 1 for LDA, 2 for GGA, and 4 for Meta-GGA.

get_vxc(densinfo: ValGrad | SpinParam[ValGrad]) → ValGrad | SpinParam[ValGrad][source]¶

Returns the ValGrad for the xc potential given the density info for unpolarized case.

Parameters:: densinfo (Union[ValGrad, SpinParam[ValGrad]]) – The density information. If the XC is unpolarized, then densinfo is ValGrad. If the XC is polarized, then densinfo is SpinParam[ValGrad]. The ValGrad contains the value and gradient of the density. The SpinParam[ValGrad] contains the value and gradient of the density for each spin channel.
Returns:: The ValGrad for the xc potential. If the XC is unpolarized, then the return is ValGrad. If the XC is polarized, then the return is SpinParam[ValGrad].
Return type:: Union[ValGrad, SpinParam[ValGrad]]

get_edensityxc(densinfo: ValGrad | SpinParam[ValGrad]) → Tensor[source]¶

Returns the xc energy density (energy per unit volume)

Parameters:: densinfo (Union[ValGrad, SpinParam[ValGrad]]) – The density information. If the XC is unpolarized, then densinfo is ValGrad. If the XC is polarized, then densinfo is SpinParam[ValGrad]. The ValGrad contains the value and gradient of the density. The SpinParam[ValGrad] contains the value and gradient of the density for each spin channel.
Returns:: The energy density of the XC.
Return type:: torch.Tensor

getparamnames(methodname: str, prefix: str = '') → List[str][source]¶

This method should list tensor names that affect the output of the method with name indicated in methodname. If the methodname is not on the list in this function, it should raise KeyError.

Parameters:

methodname (str) – The name of the method of the class.
prefix (str) – The prefix to be appended in front of the parameters name. This usually contains the dots.

Returns:

Sequence of name of parameters affecting the output of the method.

Return type:

List[str]

Raises:

KeyError – If the list in this function does not contain methodname.

class CalcLDALibXCPol(*args, **kwargs)[source]¶

Local-density approximations (LDA) are a class of approximations to the exchange–correlation (XC) energy functional in density functional theory (DFT) that depend solely upon the value of the electronic density at each point in space (and not, for example, derivatives of the density or the Kohn–Sham orbitals).

Examples

>>> import torch
>>> import pylibxc
>>> libxcfcn = pylibxc.LibXCFunctional("lda_x", "polarized")
>>> rho_u = torch.tensor([0.1, 0.2, 0.3])
>>> rho_d = torch.tensor([0.1, 0.2, 0.3])
>>> res = CalcLDALibXCPol.apply(rho_u, rho_d, 0, libxcfcn)[0]
>>> print(res)
tensor([[-0.0864, -0.2177, -0.3738]], dtype=torch.float64)

static forward(ctx, rho_u: torch.Tensor, rho_d: torch.Tensor, deriv: int, libxcfcn: pylibxc.functional.LibXCFunctional) → Tuple[torch.Tensor, ...][source]¶

Calculates and returns the energy density or its derivative w.r.t. density for polarized LDA.

Parameters:

rho_u (torch.Tensor) – Density tensor for spin-up with shape (ninps)
rho_d (torch.Tensor) – Density tensor for spin-down with shape (ninps)
deriv (int) – Derivative order. 0 for energy density, 1 for derivative w.r.t. density, 2 for second derivative w.r.t. density, etc.
libxcfcn (pylibxc.functional.LibXCFunctional) – libxc functional to use

Returns:

Result is a tensor with shape (nderiv, ninps) where the first dimension indicates the result for derivatives of spin-up and spin-down and some of its combination.

Return type:

Tuple[torch.Tensor]

static backward(ctx, *grad_res: Tensor) → Tuple[Tensor | None, ...][source]¶

Calculates the gradient w.r.t. the input rho.

Parameters:: grad_res (torch.Tensor) – Gradient of the result w.r.t. the result itself.
Returns:: Gradient w.r.t. the input rho.
Return type:: Tuple[torch.Tensor]

class CalcLDALibXCUnpol(*args, **kwargs)[source]¶

Calculates the energy density or its derivative w.r.t. density for unpolarized LDA.

Local-density approximations (LDA) are a class of approximations to the exchange–correlation (XC) energy functional in density functional theory (DFT) that depend solely upon the value of the electronic density at each point in space (and not, for example, derivatives of the density or the Kohn–Sham orbitals).

The result is a tensor with shape (ninps).

Examples

>>> import torch
>>> import pylibxc
>>> libxcfcn = pylibxc.LibXCFunctional("lda_x", "unpolarized")
>>> rho = torch.tensor([0.1, 0.2, 0.3])
>>> res = CalcLDALibXCUnpol.apply(rho, 0, libxcfcn)[0]
>>> print(res)
tensor([[-0.0343, -0.0864, -0.1483]], dtype=torch.float64)

static forward(ctx, rho: torch.Tensor, deriv: int, libxcfcn: pylibxc.functional.LibXCFunctional) → Tuple[torch.Tensor, ...][source]¶

Calculates and returns the energy density or its derivative w.r.t. density.

Parameters:

rho (torch.Tensor) – Density tensor with shape (ninps)
deriv (int) – Derivative order. 0 for energy density, 1 for derivative w.r.t. density, 2 for second derivative w.r.t. density, etc.
libxcfcn (pylibxc.functional.LibXCFunctional) – libxc functional to use

Returns:

Result is a tensor with shape (ninps)

Return type:

Tuple[torch.Tensor]

static backward(ctx, *grad_res: Tensor) → Tuple[Tensor | None, ...][source]¶

Calculates the gradient w.r.t. the input rho.

Parameters:: grad_res (torch.Tensor) – Gradient of the result w.r.t. the result itself.
Returns:: Gradient w.r.t. the input rho.
Return type:: Tuple[torch.Tensor]

class CalcGGALibXCUnpol(*args, **kwargs)[source]¶

Calculates the energy density or its derivative w.r.t. density for unpolarized GGA.

Generalized-gradient approximations (GGA) are a class of approximations to the exchange–correlation (XC) energy functional in density functional theory (DFT) that depend not only upon the value of the electronic density at each point in space, but also upon its gradient.

Examples

>>> import torch
>>> import pylibxc
>>> libxcfcn = pylibxc.LibXCFunctional("gga_c_pbe", "unpolarized")
>>> rho = torch.tensor([0.1, 0.2, 0.3])
>>> sigma = torch.tensor([0.1, 0.2, 0.3])
>>> res = CalcGGALibXCUnpol.apply(rho, sigma, 0, libxcfcn)[0]
>>> print(res)
tensor([[-0.0016, -0.0070, -0.0137]], dtype=torch.float64)

static forward(ctx, rho: torch.Tensor, sigma: torch.Tensor, deriv: int, libxcfcn: pylibxc.functional.LibXCFunctional) → Tuple[torch.Tensor, ...][source]¶

Calculates and returns the energy density or its derivative w.r.t. density and contracted gradient.

Every element in the tuple is a tensor with shape (ninps)

Parameters:

rho (torch.Tensor) – Density tensor with shape (ninps)
sigma (torch.Tensor) – Contracted gradient tensor with shape (ninps)
deriv (int) – Derivative order. 0 for energy density, 1 for derivative w.r.t. density, 2 for second derivative w.r.t. density, etc.
libxcfcn (pylibxc.functional.LibXCFunctional) – libxc functional to use

static backward(ctx, *grad_res: Tensor) → Tuple[Tensor | None, ...][source]¶

Calculates the gradient w.r.t. the input rho and sigma.

Parameters:: grad_res (torch.Tensor) – Gradient of the result w.r.t. the result itself.
Returns:: Gradient w.r.t. the input rho and sigma.
Return type:: Tuple[torch.Tensor]

class CalcGGALibXCPol(*args, **kwargs)[source]¶

Calculates the energy density or its derivative w.r.t. density for polarized GGA.

Generalized-gradient approximations (GGA) are a class of approximations to the exchange–correlation (XC) energy functional in density functional theory (DFT) that depend not only upon the value of the electronic density at each point in space, but also upon its gradient.

Examples

>>> import torch
>>> import pylibxc
>>> libxcfcn = pylibxc.LibXCFunctional("gga_c_pbe", "polarized")
>>> rho_u = torch.tensor([0.1, 0.2, 0.3])
>>> rho_d = torch.tensor([0.1, 0.2, 0.3])
>>> sigma_uu = torch.tensor([0.1, 0.2, 0.3])
>>> sigma_ud = torch.tensor([0.1, 0.2, 0.3])
>>> sigma_dd = torch.tensor([0.1, 0.2, 0.3])
>>> res = CalcGGALibXCPol.apply(rho_u, rho_d, sigma_uu, sigma_ud, sigma_dd, 0, libxcfcn)[0]
>>> print(res)
tensor([[-0.0047, -0.0175, -0.0322]], dtype=torch.float64)

static forward(ctx, rho_u: torch.Tensor, rho_d: torch.Tensor, sigma_uu: torch.Tensor, sigma_ud: torch.Tensor, sigma_dd: torch.Tensor, deriv: int, libxcfcn: pylibxc.functional.LibXCFunctional) → Tuple[torch.Tensor, ...][source]¶

Calculates and returns the energy density or its derivative w.r.t. density and contracted gradient.

Parameters:

rho_u (torch.Tensor) – Density tensor for spin-up with shape (ninps)
rho_d (torch.Tensor) – Density tensor for spin-down with shape (ninps)
sigma_uu (torch.Tensor) – Contracted gradient tensor for spin-up with shape (ninps)
sigma_ud (torch.Tensor) – Contracted gradient tensor for spin-up and spin-down with shape (ninps)
sigma_dd (torch.Tensor) – Contracted gradient tensor for spin-down with shape (ninps)
deriv (int) – Derivative order. 0 for energy density, 1 for derivative w.r.t. density, 2 for second derivative w.r.t. density, etc.
libxcfcn (pylibxc.functional.LibXCFunctional) – libxc functional to use

Returns:

Result is a tensor with shape (nderiv, ninps) where the first dimension indicates the result for derivatives of spin-up and spin-down and some of its combination.

Return type:

Tuple[torch.Tensor]

static backward(ctx, *grad_res: Tensor) → Tuple[Tensor | None, ...][source]¶

Returns the gradient w.r.t. the input rho and sigma.

Parameters:: grad_res (torch.Tensor) – Gradient of the result w.r.t. the result itself.
Returns:: Gradient w.r.t. the input rho and sigma.
Return type:: Tuple[torch.Tensor]

class CalcMGGALibXCUnpol(*args, **kwargs)[source]¶

Calculates the energy density or its derivative w.r.t. density for unpolarized meta-GGA.

Meta-generalized-gradient approximations (meta-GGA) are a class of approximations to the exchange–correlation (XC) energy functional in density functional theory (DFT) that depend not only upon the value of the electronic density at each point in space and its gradient, but also upon the Laplacian of the density and the kinetic energy density.

Examples

>>> import torch
>>> import pylibxc
>>> libxcfcn = pylibxc.LibXCFunctional("mgga_c_m06_l", "unpolarized")
>>> rho = torch.tensor([0.1, 0.2, 0.3])
>>> sigma = torch.tensor([0.1, 0.2, 0.3])
>>> lapl = torch.tensor([0.1, 0.2, 0.3])
>>> kin = torch.tensor([0.1, 0.2, 0.3])
>>> res = CalcMGGALibXCUnpol.apply(rho, sigma, lapl, kin, 0, libxcfcn)[0]
>>> print(res)
tensor([[-0.0032, -0.0066, -0.0087]], dtype=torch.float64)

static forward(ctx, rho: torch.Tensor, sigma: torch.Tensor, lapl: torch.Tensor, kin: torch.Tensor, deriv: int, libxcfcn: pylibxc.functional.LibXCFunctional) → Tuple[torch.Tensor, ...][source]¶

Calculates and returns the energy density or its derivative w.r.t. density and contracted gradient.

Parameters:

rho (torch.Tensor) – Density tensor with shape (ninps)
sigma (torch.Tensor) – Contracted gradient tensor with shape (ninps)
lapl (torch.Tensor) – Laplacian tensor with shape (ninps)
kin (torch.Tensor) – Kinetic energy density tensor with shape (ninps)
deriv (int) – Derivative order. 0 for energy density, 1 for derivative w.r.t. density, 2 for second derivative w.r.t. density, etc.
libxcfcn (pylibxc.functional.LibXCFunctional) – libxc functional to use

Returns:

The result is a tensor with shape (nderiv, ninps)

Return type:

Tuple[torch.Tensor]

static backward(ctx, *grad_res: Tensor) → Tuple[Tensor | None, ...][source]¶

Returns the gradient w.r.t. the inputs

Parameters:: grad_res (torch.Tensor) – The gradient of the result w.r.t. the result itself.
Returns:: The gradient w.r.t. the inputs
Return type:: Tuple[torch.Tensor]

class CalcMGGALibXCPol(*args, **kwargs)[source]¶

Calculates the energy density or its derivative w.r.t. density for polarized meta-GGA.

Meta-generalized-gradient approximations (meta-GGA) are a class of approximations to the exchange–correlation (XC) energy functional in density functional theory (DFT) that depend not only upon the value of the electronic density at each point in space and its gradient, but also upon the Laplacian of the density and the kinetic energy density.

Examples

>>> import torch
>>> import pylibxc
>>> libxcfcn = pylibxc.LibXCFunctional("mgga_c_m06_l", "polarized")
>>> rho_u = torch.tensor([0.1, 0.2, 0.3])
>>> rho_d = torch.tensor([0.1, 0.2, 0.3])
>>> sigma_uu = torch.tensor([0.1, 0.2, 0.3])
>>> sigma_ud = torch.tensor([0.1, 0.2, 0.3])
>>> sigma_dd = torch.tensor([0.1, 0.2, 0.3])
>>> lapl_u = torch.tensor([0.1, 0.2, 0.3])
>>> lapl_d = torch.tensor([0.1, 0.2, 0.3])
>>> kin_u = torch.tensor([0.1, 0.2, 0.3])
>>> kin_d = torch.tensor([0.1, 0.2, 0.3])
>>> res = CalcMGGALibXCPol.apply(rho_u, rho_d, sigma_uu, sigma_ud, sigma_dd,
...                               lapl_u, lapl_d, kin_u, kin_d, 0, libxcfcn)[0]
>>> print(res)
tensor([[-0.0065, -0.0115, -0.0162]], dtype=torch.float64)

static forward(ctx, rho_u: torch.Tensor, rho_d: torch.Tensor, sigma_uu: torch.Tensor, sigma_ud: torch.Tensor, sigma_dd: torch.Tensor, lapl_u: torch.Tensor, lapl_d: torch.Tensor, kin_u: torch.Tensor, kin_d: torch.Tensor, deriv: int, libxcfcn: pylibxc.functional.LibXCFunctional) → Tuple[torch.Tensor, ...][source]¶

Calculates and returns the energy density or its derivative w.r.t. density and contracted gradient and laplacian and kinetic energy density. Every element in the tuple is a tensor with shape of (nderiv, ninps) where nderiv depends on the number of derivatives for spin-up and spin-down combinations, e.g. nderiv == 3 for vsigma (see libxc manual)

Parameters:

rho_u (torch.Tensor) – The density tensor for spin-up with shape (ninps)
rho_d (torch.Tensor) – The density tensor for spin-down with shape (ninps)
sigma_uu (torch.Tensor) – The contracted gradient tensor for spin-up with shape (ninps)
sigma_ud (torch.Tensor) – The contracted gradient tensor for spin-up and spin-down with shape (ninps)
sigma_dd (torch.Tensor) – The contracted gradient tensor for spin-down with shape (ninps)
lapl_u (torch.Tensor) – The laplacian tensor for spin-up with shape (ninps)
lapl_d (torch.Tensor) – The laplacian tensor for spin-down with shape (ninps)
kin_u (torch.Tensor) – The kinetic energy density tensor for spin-up with shape (ninps)
kin_d (torch.Tensor) – The kinetic energy density tensor for spin-down with shape (ninps)
deriv (int) – The derivative order. 0 for energy density, 1 for derivative w.r.t. density, 2 for second derivative w.r.t. density, etc.
libxcfcn (pylibxc.functional.LibXCFunctional) – The libxc functional to use

Returns:

The result is a tensor with shape (nderiv, ninps) The result is a tensor with shape (nderiv, ninps) where the first dimension indicates the result for derivatives of spin-up and spin-down and some of its combination.

Return type:

Tuple[torch.Tensor]

static backward(ctx, *grad_res: Tensor) → Tuple[Tensor | None, ...][source]¶

Returns the gradient w.r.t. the input rho and sigma.

Parameters:: grad_res (torch.Tensor) – The gradient of the result w.r.t. the result itself.
Returns:: The gradient w.r.t. the input rho and sigma.
Return type:: Tuple[torch.Tensor]

loadbasis(cmd: str, dtype: dtype = torch.float64, device: device = device(type='cpu'), requires_grad: bool = False) → List[CGTOBasis][source]¶

Load basis from a file and return the list of CGTOBasis.

Examples

>>> H = loadbasis("1:3-21G")
>>> H[0].alphas.shape
torch.Size([2])

Parameters:

cmd (str) – This can be a file path where the basis is stored or a string in format "atomz:basis", e.g. "1:6-311++G**".
dtype (torch.dtype) – Tensor data type for alphas and coeffs of the GTO basis
device (torch.device) – Tensor device for alphas and coeffs
requires_grad (bool) – If True, the alphas and coeffs tensors become differentiable

Returns:

List of GTO basis loaded from the given file

Return type:

list of CGTOBasis

_read_float(s: str) → float[source]¶

Read a float from a string, replacing D with E.

Examples

>>> _read_float("1.0D+00")
1.0

Parameters:: s (str) – String to be converted to float
Returns:: Float number
Return type:: float

_get_basis_file(cmd: str) → str[source]¶

parse the string command, check if the basis has already been downloaded (download if not), and return the file name

Examples

>>> path = _get_basis_file("1:6-311++G**")

Parameters:: cmd (str) – Command string in the format “atomz:basis”
Returns:: File path to the basis set
Return type:: str

_normalize_basisname(basisname: str) → str[source]¶

Normalize the basis name to a file name

Examples

>>> _normalize_basisname("6-311++G**")
'6-311ppgss'

Parameters:: basisname (str) – Basis name to be normalized
Returns:: Normalized basis name
Return type:: str

_download_basis(fname: str, atomz: int, basisname: str) → str[source]¶

Download the basis set from the basis set exchange and save it to a file

Examples

>>> import pathlib
>>> import os
>>> path = pathlib.Path(__file__).parent.resolve()
>>> dir = str(path) + "/.database/6-311ppgss"
>>> if not os.path.exists(dir):
...     os.makedirs(dir)
>>> path = _download_basis(str(path) + "/01.gaussian94", 1, "6-311++G**")

Parameters:

fname (str) – File name to save the basis set
atomz (int) – Atomic number of the atom
basisname (str) – Basis name to be downloaded

_expand_angmoms(s: str, n: int) → List[int][source]¶

convert the angular momentum characters into angmom and returns a list of n integer containing the angular momentums

Examples

>>> _expand_angmoms("SP", 2)
[0, 1]

Parameters:

s (str) – Angular momentum characters
n (int) – Number of coefficients

Returns:

List of angular momentums

Return type:

List[int]

class BaseDF[source]¶

BaseDF represents the density fitting object used in calculating the electron repulsion (and xc energy) in Hamiltonian.

Density fitting in density functional theory (DFT) is a technique used to reduce the computational cost of evaluating electron repulsion integrals. In DFT, the key quantity is the electron density rather than the wave function, and the electron repulsion integrals involve four-electron interactions, making them computationally demanding.

Density fitting exploits the fact that many-electron integrals can be expressed as a sum of two-electron integrals involving auxiliary basis functions. By approximating these auxiliary basis functions, often referred to as fitting functions, the computational cost can be significantly reduced.

Examples

>>> from deepchem.utils.dft_utils import BaseDF
>>> import torch
>>> class MyDF(BaseDF):
...     def __init__(self):
...         super(MyDF, self).__init__()
...     def get_j2c(self):
...         return torch.ones((3, 3))
...     def get_j3c(self):
...         return torch.ones((3, 3, 3))
>>> df = MyDF()
>>> df.get_j2c()
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]])

abstract build() → BaseDF[source]¶

Construct the matrices required to perform the calculation and return self.

Returns:: The constructed density fitting object.
Return type:: BaseDF

abstract get_elrep(dm: Tensor) → LinearOperator[source]¶

Construct the electron repulsion linear operator from the given density matrix using the density fitting method.

Parameters:: dm (torch.Tensor) – The density matrix.
Returns:: The electron repulsion linear operator.
Return type:: LinearOperator

abstract property j2c: Tensor[source]¶

Returns the 2-centre 2-electron integrals of the auxiliary basis.

Returns:: The 2-centre 2-electron integrals of the auxiliary basis.
Return type:: torch.Tensor

abstract property j3c: Tensor[source]¶

Return the 3-centre 2-electron integrals of the auxiliary basis and the basis.

Returns:: The 3-centre 2-electron integrals of the auxiliary basis and the basis.
Return type:: torch.Tensor

abstract getparamnames(methodname: str, prefix: str = '') → List[str][source]¶

This method should list tensor names that affect the output of the method with name indicated in methodname.

Parameters:

methodname (str) – The name of the method of the class.
prefix (str (default="")) – The prefix to be appended in front of the parameters name. This usually contains the dots.

Returns:

Sequence of name of parameters affecting the output of the method.

Return type:

List[str]

class _Config(THRESHOLD_MEMORY: int = 10737418240, CHUNK_MEMORY: int = 16777216, VERBOSE: int = 0)[source]¶

Contains the configuration for the DFT module

Examples

>>> from deepchem.utils.dft_utils.config import config
>>> Memory_usage = 1024**4 # Sample Memory usage by some Object/Matrix
>>> if Memory_usage > config.THRESHOLD_MEMORY :
...     print("Overload")
Overload

THRESHOLD_MEMORY[source]¶

Threshold memory (matrix above this size should not be constructed)

Type:: int (default=10*1024**3)

CHUNK_MEMORY[source]¶

The memory for splitting big tensors into chunks.

Type:: int (default=16*1024**2)

VERBOSE[source]¶

Allowed Verbosity level (Defines the level of detail) Used by Looger for maintaining Logs.

Type:: int (default=0)

Usage[source]¶

-----

1. HamiltonCGTO

Type:: Usage it for splitting big tensors into chunks.

__init__(THRESHOLD_MEMORY: int = 10737418240, CHUNK_MEMORY: int = 16777216, VERBOSE: int = 0) → None[source]¶

class BaseOrbParams[source]¶

Class that provides free-parameterization of orthogonal orbitals.

Examples

>>> import torch
>>> from deepchem.utils.dft_utils import BaseOrbParams
>>> class MyOrbParams(BaseOrbParams):
...     @staticmethod
...     def params2orb(params, coeffs, with_penalty):
...         return params, coeffs
...     @staticmethod
...     def orb2params(orb):
...         return orb, torch.tensor([0], dtype=orb.dtype, device=orb.device)
>>> params = torch.randn(3, 4, 5)
>>> coeffs = torch.randn(3, 4, 5)
>>> with_penalty = 0.1
>>> orb, penalty = MyOrbParams.params2orb(params, coeffs, with_penalty)
>>> params2, coeffs2 = MyOrbParams.orb2params(orb)
>>> torch.allclose(params, params2)
True

static params2orb(params: Tensor, coeffs: Tensor, with_penalty: float = 0.0) → List[Tensor][source]¶

Convert the parameters & coefficients to the orthogonal orbitals. params is the tensor to be optimized in variational method, while coeffs is a tensor that is needed to get the orbital, but it is not optimized in the variational method.

Parameters:

params (torch.Tensor) – The free parameters to be optimized.
coeffs (torch.Tensor) – The coefficients to get the orthogonal orbitals.
with_penalty (float (default 0.0)) – If not 0.0, return the penalty term for the free parameters.

Returns:

orb (torch.Tensor) – The orthogonal orbitals.
penalty (torch.Tensor) – The penalty term for the free parameters. If with_penalty is 0.0, this is not returned.

static orb2params(orb: Tensor) → List[Tensor][source]¶

Get the free parameters from the orthogonal orbitals. Returns params and coeffs described in params2orb.

Parameters:

orb (torch.Tensor) – The orthogonal orbitals.

Returns:

params (torch.Tensor) – The free parameters to be optimized.
coeffs (torch.Tensor) – The coefficients to get the orthogonal orbitals.

class QROrbParams[source]¶

Orthogonal orbital parameterization using QR decomposition. The orthogonal orbital is represented by:

P = QR

Where Q is the parameters defining the rotation of the orthogonal tensor, and R is the coefficients tensor.

Examples

>>> import torch
>>> from deepchem.utils.dft_utils import QROrbParams
>>> params = torch.randn(3, 3)
>>> coeffs = torch.randn(4, 3)
>>> with_penalty = 0.1
>>> orb, penalty = QROrbParams.params2orb(params, coeffs, with_penalty)
>>> params2, coeffs2 = QROrbParams.orb2params(orb)

static params2orb(params: Tensor, coeffs: Tensor, with_penalty: float = 0.0) → List[Tensor][source]¶

Convert the parameters & coefficients to the orthogonal orbitals. params is the tensor to be optimized in variational method, while coeffs is a tensor that is needed to get the orbital, but it is not optimized in the variational method.

Parameters:

params (torch.Tensor) – The free parameters to be optimized.
coeffs (torch.Tensor) – The coefficients to get the orthogonal orbitals.
with_penalty (float (default 0.0)) – If not 0.0, return the penalty term for the free parameters.

Returns:

orb (torch.Tensor) – The orthogonal orbitals.
penalty (torch.Tensor) – The penalty term for the free parameters. If with_penalty is 0.0, this is not returned.

static orb2params(orb: Tensor) → List[Tensor][source]¶

Get the free parameters from the orthogonal orbitals. Returns params and coeffs described in params2orb.

Parameters:

orb (torch.Tensor) – The orthogonal orbitals.

Returns:

params (torch.Tensor) – The free parameters to be optimized.
coeffs (torch.Tensor) – The coefficients to get the orthogonal orbitals.

class MatExpOrbParams[source]¶

Orthogonal orbital parameterization using matrix exponential. The orthogonal orbital is represented by:

P = matrix_exp(Q) @ C

where C is an orthogonal coefficient tensor, and Q is the parameters defining the rotation of the orthogonal tensor.

Examples

>>> from deepchem.utils.dft_utils import MatExpOrbParams
>>> params = torch.randn(3, 3)
>>> coeffs = torch.randn(4, 3)
>>> with_penalty = 0.1
>>> orb, penalty = MatExpOrbParams.params2orb(params, coeffs, with_penalty)
>>> params2, coeffs2 = MatExpOrbParams.orb2params(orb)

static params2orb(params: Tensor, coeffs: Tensor, with_penalty: float = 0.0) → List[Tensor][source]¶

Convert the parameters & coefficients to the orthogonal orbitals. params is the tensor to be optimized in variational method, while coeffs is a tensor that is needed to get the orbital, but it is not optimized in the variational method.

Parameters:

params (torch.Tensor) – The free parameters to be optimized. (*, nparams)
coeffs (torch.Tensor) – The coefficients to get the orthogonal orbitals. (*, nao, norb)
with_penalty (float (default 0.0)) – If not 0.0, return the penalty term for the free parameters.

Returns:

orb (torch.Tensor) – The orthogonal orbitals.
penalty (torch.Tensor) – The penalty term for the free parameters. If with_penalty is 0.0, this is not returned.

static orb2params(orb: Tensor) → List[Tensor][source]¶

Get the free parameters from the orthogonal orbitals. Returns params and coeffs described in params2orb.

Parameters:

orb (torch.Tensor) – The orthogonal orbitals.

Returns:

params (torch.Tensor) – The free parameters to be optimized.
coeffs (torch.Tensor) – The coefficients to get the orthogonal orbitals.

Parse the string of molecular descriptor and returns tensors of atomzs and atom positions.

Examples

>>> from deepchem.utils.dft_utils import parse_moldesc
>>> system = {
...     'type': 'mol',
...     'kwargs': {
...         'moldesc': 'H 0.86625 0 0; F -0.86625 0 0',
...         'basis': '6-311++G(3df,3pd)'
...     }
... }
>>> atomzs, atomposs = parse_moldesc(system["kwargs"]["moldesc"])
>>> atomzs
tensor([1, 9])
>>> atomposs
tensor([[ 0.8662,  0.0000,  0.0000],
        [-0.8662,  0.0000,  0.0000]], dtype=torch.float64)

Parameters:

moldesc (Union[str, Tuple[AtomZsType, AtomPosType]]) – String that describes the system, e.g. "H -1 0 0; H 1 0 0" for H2 molecule separated by 2 Bohr.
dtype (torch.dtype (default torch.float64)) – The datatype of the returned atomic positions.
device (torch.device (default torch.device('cpu'))) – The device to store the returned tensors.

Returns:

atomzs (torch.Tensor) – The tensor of atomzs [Atom Number].
atompos (torch.Tensor) – The tensor of atomic positions [Bohr].

class EditableModule[source]¶

EditableModule is a base class to enable classes that it inherits be converted to pure functions for higher order derivatives purpose.

Usage¶

To use this class, the user must implement the getparamnames method which returns a list of tensor names that affect the output of the method with name indicated in methodname.

Used in:

Classes of Density Functional Theory (DFT).
It can also be used in other classes that need to be converted to pure functions for higher order derivatives purpose.

Examples

>>> import torch
>>> from deepchem.utils.differentiation_utils import EditableModule
>>> class A(EditableModule):
...     def __init__(self, a):
...         self.b = a*a
...
...     def mult(self, x):
...         return self.b * x
...
...     def getparamnames(self, methodname, prefix=""):
...         if methodname == "mult":
...             return [prefix+"b"]
...         else:
...             raise KeyError()
>>> a = torch.tensor(2.0).requires_grad_()
>>> x = torch.tensor(0.4).requires_grad_()
>>> alpha = A(a)
>>> alpha.mult(x)
tensor(1.6000, grad_fn=<MulBackward0>)
>>> alpha.getparamnames("mult")
['b']
>>> alpha.assertparams(alpha.mult, x)
"mult" method check done

getparams(methodname: str) → Sequence[Tensor][source]¶

Returns a list of tensor parameters used in the object’s operations. Requires the getparamnames method to be implemented.

Parameters:: methodname (str) – The name of the method of the class.
Returns:: Sequence of tensors that are involved in the specified method of the object.
Return type:: Sequence[torch.Tensor]

setparams(methodname: str, *params) → int[source]¶

Set the input parameters to the object’s parameters to make a copy of the operations.

Parameters:

methodname (str) – The name of the method of the class.
*params – The parameters to be set to the object’s parameters.

Returns:

The number of parameters that are set to the object’s parameters.

Return type:

int

cached_getparamnames(methodname: str, refresh: bool = False) → List[str][source]¶

getparamnames, but cached, so it is only called once

Parameters:

methodname (str) – The name of the method of the class.
refresh (bool) – If True, the cache is refreshed.

Returns:

Sequence of name of parameters affecting the output of the method.

Return type:

List[str]

abstract getparamnames(methodname: str, prefix: str = '') → List[str][source]¶

This method should list tensor names that affect the output of the method with name indicated in methodname. If the methodname is not on the list in this function, it should raise KeyError.

Parameters:

methodname (str) – The name of the method of the class.
prefix (str) – The prefix to be appended in front of the parameters name. This usually contains the dots.

Returns:

Sequence of name of parameters affecting the output of the method.

Return type:

List[str]

Raises:

KeyError – If the list in this function does not contain methodname.

getuniqueparams(methodname: str, onlyleaves: bool = False) → List[Tensor][source]¶

Returns the list of unique parameters involved in the method specified by methodname.

Parameters:

methodname (str) – Name of the method where the returned parameters play roles.
onlyleaves (bool) – If True, only returns leaf tensors. Otherwise, returns all tensors.

Returns:

List of tensors that are involved in the specified method of the object.

Return type:

List[torch.Tensor]

setuniqueparams(methodname: str, *uniqueparams) → int[source]¶

Set the input parameters to the object’s parameters to make a copy of the operations. The input parameters are unique parameters, i.e. they are not necessarily the same tensors as the object’s parameters.

Note: This function can only be run after running getuniqueparams.

Parameters:

methodname (str) – The name of the method of the class.
*uniqueparams – The parameters to be set to the object’s parameters. The number of parameters must be the same as the number of unique parameters returned by getuniqueparams.

Returns:

The number of parameters that are set to the object’s parameters.

Return type:

int

assertparams(method: Callable, *args, **kwargs)[source]¶

Perform a rigorous check on the implemented getparamnames in the class for a given method and its parameters as well as keyword Parameters. It raises warnings if there are missing or excess parameters in the getparamnames implementation.

Parameters:

method (Callable) – The method of this class to be tested
*args – Parameters of the method
**kwargs – Keyword parameters of the method

normalize_bcast_dims(*shapes)[source]¶

Normalize the lengths of the input shapes to have the same length. The shapes are padded at the front by 1 to make the lengths equal.

Examples

>>> normalize_bcast_dims([1, 2, 3], [2, 3])
[[1, 2, 3], [1, 2, 3]]

Parameters:: shapes (List[List[int]]) – The shapes to normalize.
Returns:: The normalized shapes.
Return type:: List[List[int]]

get_bcasted_dims(*shapes)[source]¶

Return the broadcasted shape of the given shapes.

Examples

>>> get_bcasted_dims([1, 2, 5], [2, 3, 4])
[2, 3, 5]

Parameters:: shapes (List[List[int]]) – The shapes to broadcast.
Returns:: The broadcasted shape.
Return type:: List[int]

match_dim(*xs: Tensor, contiguous: bool = False) → Tuple[Tensor, ...][source]¶

match the N-1 dimensions of x and xq for searchsorted and gather with dim=-1

Examples

>>> x = torch.randn(10, 5)
>>> xq = torch.randn(10, 3)
>>> x_new, xq_new = match_dim(x, xq)
>>> x_new.shape
torch.Size([10, 5])
>>> xq_new.shape
torch.Size([10, 3])

class LinearOperator(*args, **kwargs)[source]¶

LinearOperator is a base class designed to behave as a linear operator without explicitly determining the matrix. This LinearOperator should be able to operate as batched linear operators where its shape is (B1,B2,...,Bb,p,q) with B* as the (optional) batch dimensions. For a user-defined class to behave as LinearOperator, it must use LinearOperator as one of the parent and it has to have ._mv() method implemented and ._getparamnames() if used in xitorch’s functionals with torch grad enabled.

Examples

>>> import torch
>>> seed = torch.manual_seed(100)
>>> class MyLinOp(LinearOperator):
...     def __init__(self, shape):
...         super(MyLinOp, self).__init__(shape)
...         self.param = torch.rand(shape)
...     def _getparamnames(self, prefix=""):
...         return [prefix + "param"]
...     def _mv(self, x):
...         return torch.matmul(self.param, x)
...     def _rmv(self, x):
...         return torch.matmul(self.param.transpose(-2,-1).conj(), x)
...     def _mm(self, x):
...         return torch.matmul(self.param, x)
...     def _rmm(self, x):
...         return torch.matmul(self.param.transpose(-2,-1).conj(), x)
...     def _fullmatrix(self):
...         return self.param
>>> linop = MyLinOp((1,3,1,2))
>>> print(linop)
LinearOperator (MyLinOp) with shape (1, 3, 1, 2), dtype = torch.float32, device = cpu
>>> x = torch.rand(1,3,2,2)
>>> linop.mv(x)
tensor([[[[0.1991, 0.1011]],

         [[0.3764, 0.5742]],

         [[1.0345, 1.1802]]]])
>>> x = torch.rand(1,3,1,1)
>>> linop.rmv(x)
tensor([[[[0.0250],
          [0.1827]],

         [[0.0794],
          [0.1463]],

         [[0.1207],
          [0.1345]]]])
>>> x = torch.rand(1,3,2,2)
>>> linop.mm(x)
tensor([[[[0.8891, 0.4243]],

         [[0.4856, 0.3128]],

         [[0.6601, 0.9532]]]])
>>> x = torch.rand(1,3,1,2)
>>> linop.rmm(x)
tensor([[[[0.0473, 0.0019],
          [0.3455, 0.0138]],

         [[0.0580, 0.2504],
          [0.1069, 0.4614]],

         [[0.4779, 0.1102],
          [0.5326, 0.1228]]]])
>>> linop.fullmatrix()
tensor([[[[0.1117, 0.8158]],

         [[0.2626, 0.4839]],

         [[0.6765, 0.7539]]]])

static __new__(self, *args, **kwargs)[source]¶: Check the implemented functions in the class.

classmethod m(mat: Tensor, is_hermitian: bool | None = None)[source]¶

Class method to wrap a matrix into LinearOperator.

Parameters:

mat (torch.Tensor) – Matrix to be wrapped in the LinearOperator.
is_hermitian (bool or None) – Indicating if the matrix is Hermitian. If None, the symmetry will be checked. If supplied as a bool, there is no check performed.

Returns:

Linear operator object that represents the matrix.

Return type:

LinearOperator

Example

>>> import torch
>>> from deepchem.utils.differentiation_utils import LinearOperator
>>> seed = torch.manual_seed(100)
>>> mat = torch.rand(1,3,1,2)  # 1x2 matrix with (1,3) batch dimensions
>>> linop = LinearOperator.m(mat)
>>> print(linop)
MatrixLinearOperator with shape (1, 3, 1, 2):
   tensor([[[[0.1117, 0.8158]],

            [[0.2626, 0.4839]],

            [[0.6765, 0.7539]]]])

__init__(shape: Sequence[int], is_hermitian: bool = False, dtype: dtype | None = None, device: device | None = None, _suppress_hermit_warning: bool = False) → None[source]¶

Initialize the LinearOperator.

Parameters:

shape (Sequence[int]) – The shape of the linear operator.
is_hermitian (bool) – Whether the linear operator is Hermitian.
dtype (torch.dtype or None) – The dtype of the linear operator.
device (torch.device or None) – The device of the linear operator.
_suppress_hermit_warning (bool) – Whether to suppress the warning when the linear operator is Hermitian but the .rmv() or .rmm() is implemented.

getlinopparams() → Sequence[Tensor][source]¶: Get the parameters that affects most of the methods (i.e. mm, mv, rmm, rmv).

uselinopparams(*params)[source]¶: Context manager to temporarily set the parameters that affects most of the methods (i.e. mm, mv, rmm, rmv).

mv(x: Tensor) → Tensor[source]¶

Apply the matrix-vector operation to vector x with shape (...,q). The batch dimensions of x need not be the same as the batch dimensions of the LinearOperator, but it must be broadcastable.

Parameters:: x (torch.tensor) – The vector with shape (...,q) where the linear operation is operated on
Returns:: y – The result of the linear operation with shape (...,p)
Return type:: torch.tensor

mm(x: Tensor) → Tensor[source]¶

Apply the matrix-matrix operation to matrix x with shape (...,q,r). The batch dimensions of x need not be the same as the batch dimensions of the LinearOperator, but it must be broadcastable.

Parameters:: x (torch.tensor) – The matrix with shape (...,q,r) where the linear operation is operated on.
Returns:: y – The result of the linear operation with shape (...,p,r)
Return type:: torch.tensor

rmv(x: Tensor) → Tensor[source]¶

Apply the matrix-vector adjoint operation to vector x with shape (...,p), i.e. A^H x. The batch dimensions of x need not be the same as the batch dimensions of the LinearOperator, but it must be broadcastable.

Parameters:: x (torch.tensor) – The vector of shape (...,p) where the adjoint linear operation is operated at.
Returns:: y – The result of the adjoint linear operation with shape (...,q)
Return type:: torch.tensor

rmm(x: Tensor) → Tensor[source]¶

Apply the matrix-matrix adjoint operation to matrix x with shape (...,p,r), i.e. A^H X. The batch dimensions of x need not be the same as the batch dimensions of the LinearOperator, but it must be broadcastable.

Parameters:: x (torch.Tensor) – The matrix of shape (...,p,r) where the adjoint linear operation is operated on.
Returns:: y – The result of the adjoint linear operation with shape (...,q,r).
Return type:: torch.Tensor

fullmatrix() → Tensor[source]¶: Full matrix representation of the linear operator.

scipy_linalg_op()[source]¶: Return the scipy.sparse.linalg.LinearOperator object of the linear operator.

getparamnames(methodname: str, prefix: str = '') → List[str][source]¶: Get the parameter names that affects the method.

property H[source]¶

Returns a LinearOperator representing the Hermite / transposed of the self LinearOperator.

Returns:: The Hermite / transposed LinearOperator
Return type:: LinearOperator

matmul(b: LinearOperator, is_hermitian: bool = False)[source]¶

Returns a LinearOperator representing self @ b.

Examples

>>> import torch
>>> seed = torch.manual_seed(100)
>>> class MyLinOp(LinearOperator):
...     def __init__(self, shape):
...         super(MyLinOp, self).__init__(shape)
...         self.param = torch.rand(shape)
...     def _getparamnames(self, prefix=""):
...         return [prefix + "param"]
...     def _mv(self, x):
...         return torch.matmul(self.param, x)
>>> linop1 = MyLinOp((1,3,1,2))
>>> linop2 = MyLinOp((1,3,2,1))
>>> linop = linop1.matmul(linop2)
>>> print(linop)
MatmulLinearOperator with shape (1, 3, 1, 1) of:
 * LinearOperator (MyLinOp) with shape (1, 3, 1, 2), dtype = torch.float32, device = cpu
 * LinearOperator (MyLinOp) with shape (1, 3, 2, 1), dtype = torch.float32, device = cpu
>>> x = torch.rand(1,3,1,1)
>>> linop.mv(x)
tensor([[[[0.0458]],

         [[0.0880]],

         [[0.2664]]]])

Parameters:

b (LinearOperator) – Other linear operator
is_hermitian (bool) – Flag to indicate if the resulting LinearOperator is Hermitian.

Returns:

LinearOperator representing self @ b

Return type:

LinearOperator

__add__(b: LinearOperator)[source]¶

Addition with another linear operator.

Examples

>>> class Operator(LinearOperator):
...     def __init__(self, mat: torch.Tensor, is_hermitian: bool) -> None:
...         super(Operator, self).__init__(
...             shape=mat.shape,
...             is_hermitian=is_hermitian,
...             dtype=mat.dtype,
...             device=mat.device,
...             _suppress_hermit_warning=True,
...         )
...         self.mat = mat
...     def _mv(self, x: torch.Tensor) -> torch.Tensor:
...         return torch.matmul(self.mat, x.unsqueeze(-1)).squeeze(-1)
...     def _mm(self, x: torch.Tensor) -> torch.Tensor:
...         return torch.matmul(self.mat, x)
...     def _rmv(self, x: torch.Tensor) -> torch.Tensor:
...         return torch.matmul(self.mat.transpose(-3, -1).conj(), x.unsqueeze(-1)).squeeze(-1)
...     def _rmm(self, x: torch.Tensor) -> torch.Tensor:
...         return torch.matmul(self.mat.transpose(-2, -1).conj(), x)
...     def _fullmatrix(self) -> torch.Tensor:
...         return self.mat
...     def _getparamnames(self, prefix: str = "") -> List[str]:
...         return [prefix + "mat"]
>>> op = Operator(torch.tensor([[1, 2.],
...                             [3, 4]]), is_hermitian=False)
>>> x = torch.tensor([[2, 2],
...                   [1, 2.]])
>>> op.mm(x)
tensor([[ 4.,  6.],
        [10., 14.]])
>>> op2 = op + op
>>> op2.mm(x)
tensor([[ 8., 12.],
        [20., 28.]])

Parameters:: b (LinearOperator) – The linear operator to be added.
Returns:: The result of the addition.
Return type:: LinearOperator

__sub__(b: LinearOperator)[source]¶

Subtraction with another linear operator.

Examples

>>> class Operator(LinearOperator):
...     def __init__(self, mat: torch.Tensor, is_hermitian: bool) -> None:
...         super(Operator, self).__init__(
...             shape=mat.shape,
...             is_hermitian=is_hermitian,
...             dtype=mat.dtype,
...             device=mat.device,
...             _suppress_hermit_warning=True,
...         )
...         self.mat = mat
...     def _mv(self, x: torch.Tensor) -> torch.Tensor:
...         return torch.matmul(self.mat, x.unsqueeze(-1)).squeeze(-1)
...     def _mm(self, x: torch.Tensor) -> torch.Tensor:
...         return torch.matmul(self.mat, x)
...     def _rmv(self, x: torch.Tensor) -> torch.Tensor:
...         return torch.matmul(self.mat.transpose(-3, -1).conj(), x.unsqueeze(-1)).squeeze(-1)
...     def _rmm(self, x: torch.Tensor) -> torch.Tensor:
...         return torch.matmul(self.mat.transpose(-2, -1).conj(), x)
...     def _fullmatrix(self) -> torch.Tensor:
...         return self.mat
...     def _getparamnames(self, prefix: str = "") -> List[str]:
...         return [prefix + "mat"]
>>> op = Operator(torch.tensor([[1, 2.],
...                             [3, 4]]), is_hermitian=False)
>>> op1 = Operator(torch.tensor([[0, 1.],
...                              [1, 2]]), is_hermitian=False)
>>> x = torch.tensor([[2, 2],
...                   [1, 2.]])
>>> op.mm(x)
tensor([[ 4.,  6.],
        [10., 14.]])
>>> op2 = op - op1
>>> op2.mm(x)
tensor([[3., 4.],
        [6., 8.]])

Parameters:: b (LinearOperator) – The linear operator to be subtracted.
Returns:: The result of the subtraction.
Return type:: LinearOperator

property dtype: dtype[source]¶: The dtype of the linear operator.

property device: device[source]¶: The device of the linear operator.

property shape: Sequence[int][source]¶: The shape of the linear operator.

property is_hermitian: bool[source]¶: Whether the linear operator is Hermitian.

property is_mv_implemented: bool[source]¶: Whether the .mv() method is implemented.

property is_mm_implemented: bool[source]¶: Whether the .mm() method is implemented.

property is_rmv_implemented: bool[source]¶: Whether the .rmv() method is implemented.

property is_rmm_implemented: bool[source]¶: Whether the .rmm() method is implemented.

property is_fullmatrix_implemented: bool[source]¶: Whether the .fullmatrix() method is implemented.

property is_getparamnames_implemented: bool[source]¶: Whether the ._getparamnames() method is implemented.

class AddLinearOperator(*args, **kwargs)[source]¶

Adds two linear operators.

Examples

>>> import torch
>>> seed = torch.manual_seed(100)
>>> class MyLinOp(LinearOperator):
...     def __init__(self, shape):
...         super(MyLinOp, self).__init__(shape)
...         self.param = torch.rand(shape)
...     def _getparamnames(self, prefix=""):
...         return [prefix + "param"]
...     def _mv(self, x):
...         return torch.matmul(self.param, x)
...     def _rmv(self, x):
...         return torch.matmul(self.param.transpose(-2,-1).conj(), x)
...     def _mm(self, x):
...         return torch.matmul(self.param, x)
...     def _rmm(self, x):
...         return torch.matmul(self.param.transpose(-2,-1).conj(), x)
...     def _fullmatrix(self):
...         return self.param
>>> linop1 = MyLinOp((1,3,1,2))
>>> linop2 = MyLinOp((1,3,1,2))
>>> linop = AddLinearOperator(linop1, linop2)
>>> print(linop)
AddLinearOperator with shape (1, 3, 1, 2) of:
 * LinearOperator (MyLinOp) with shape (1, 3, 1, 2), dtype = torch.float32, device = cpu
 * LinearOperator (MyLinOp) with shape (1, 3, 1, 2), dtype = torch.float32, device = cpu
>>> x = torch.rand(1,3,2,2)
>>> linop.mv(x)
tensor([[[[0.6256, 1.0689]],

         [[0.6039, 0.5380]],

         [[0.9702, 2.1129]]]])
>>> x = torch.rand(1,3,1,1)
>>> linop.rmv(x)
tensor([[[[0.1662],
          [0.3813]],

         [[0.4460],
          [0.5705]],

         [[0.5942],
          [1.1089]]]])
>>> x = torch.rand(1,2,2,1)
>>> linop.mm(x)
tensor([[[[0.7845],
          [0.5439]]],


        [[[0.6518],
          [0.4318]]],


        [[[1.4336],
          [0.9796]]]])

__init__(a: LinearOperator, b: LinearOperator, mul: int = 1)[source]¶

Initialize the AddLinearOperator.

Parameters:

a (LinearOperator) – The first linear operator to be added.
b (LinearOperator) – The second linear operator to be added.
mul (int) – The multiplier of the second linear operator. Default to 1. If -1, then the second linear operator will be subtracted.

class MulLinearOperator(*args, **kwargs)[source]¶

Multiply a linear operator with a scalar.

Examples

>>> import torch
>>> seed = torch.manual_seed(100)
>>> class MyLinOp(LinearOperator):
...     def __init__(self, shape):
...         super(MyLinOp, self).__init__(shape)
...         self.param = torch.rand(shape)
...     def _getparamnames(self, prefix=""):
...         return [prefix + "param"]
...     def _mv(self, x):
...         return torch.matmul(self.param, x)
>>> linop = MyLinOp((1,3,1,2))
>>> print(linop)
LinearOperator (MyLinOp) with shape (1, 3, 1, 2), dtype = torch.float32, device = cpu
>>> x = torch.rand(1,3,2,2)
>>> linop.mv(x)
tensor([[[[0.1991, 0.1011]],

         [[0.3764, 0.5742]],

         [[1.0345, 1.1802]]]])
>>> linop2 = linop * 2
>>> linop2.mv(x)
tensor([[[[0.3981, 0.2022]],

         [[0.7527, 1.1485]],

         [[2.0691, 2.3603]]]])

__init__(a: LinearOperator, f: int | float)[source]¶

Initialize the MulLinearOperator.

Parameters:

a (LinearOperator) – Linear operator to be multiplied.
f (Union[int, float]) – Integer or floating point number to be multiplied.

class AdjointLinearOperator(*args, **kwargs)[source]¶

Adjoint of a LinearOperator.

This is used to calculate the adjoint of a LinearOperator without explicitly constructing the adjoint matrix. This is useful when the adjoint matrix is not explicitly constructed, e.g. when the LinearOperator is a function of other parameters.

Examples

>>> import torch
>>> seed = torch.manual_seed(100)
>>> class MyLinOp(LinearOperator):
...     def __init__(self, shape):
...         super(MyLinOp, self).__init__(shape)
...         self.param = torch.rand(shape)
...     def _getparamnames(self, prefix=""):
...         return [prefix + "param"]
...     def _mv(self, x):
...         return torch.matmul(self.param, x)
...     def _rmv(self, x):
...         return torch.matmul(self.param.transpose(-2,-1).conj(), x)
>>> linop = MyLinOp((1,3,1,2))
>>> print(linop)
LinearOperator (MyLinOp) with shape (1, 3, 1, 2), dtype = torch.float32, device = cpu
>>> x = torch.rand(1,3,1,1)
>>> linop.rmv(x)
tensor([[[[0.0293],
          [0.2143]],

         [[0.0112],
          [0.0207]],

         [[0.1407],
          [0.1568]]]])
>>> linop2 = linop.H
>>> linop2.mv(x)
tensor([[[[0.0293],
          [0.2143]],

         [[0.0112],
          [0.0207]],

         [[0.1407],
          [0.1568]]]])

__init__(obj: LinearOperator)[source]¶

Initialize the AdjointLinearOperator.

Parameters:: obj (LinearOperator) – The linear operator to be adjointed.

property H[source]¶

Adjoint of the linear operator.

Returns:: Adjoint of the linear operator.
Return type:: LinearOperator

class MatmulLinearOperator(*args, **kwargs)[source]¶

Matrix-matrix multiplication of two linear operators.

Examples

>>> import torch
>>> seed = torch.manual_seed(100)
>>> class MyLinOp(LinearOperator):
...     def __init__(self, shape):
...         super(MyLinOp, self).__init__(shape)
...         self.param = torch.rand(shape)
...     def _getparamnames(self, prefix=""):
...         return [prefix + "param"]
...     def _mv(self, x):
...         return torch.matmul(self.param, x)
>>> linop1 = MyLinOp((1,3,2,2))
>>> linop2 = MyLinOp((1,3,2,2))
>>> linop = MatmulLinearOperator(linop1, linop2)
>>> print(linop)
MatmulLinearOperator with shape (1, 3, 2, 2) of:
 * LinearOperator (MyLinOp) with shape (1, 3, 2, 2), dtype = torch.float32, device = cpu
 * LinearOperator (MyLinOp) with shape (1, 3, 2, 2), dtype = torch.float32, device = cpu
>>> x = torch.rand(1,2,2,1)
>>> linop.mm(x)
tensor([[[[0.7998],
          [0.8016]],

         [[0.6515],
          [0.6835]]],


        [[[0.9251],
          [1.1611]],

         [[0.2781],
          [0.3609]]],


        [[[0.2591],
          [0.2376]],

         [[0.8009],
          [0.8087]]]])

__init__(a: LinearOperator, b: LinearOperator, is_hermitian: bool = False)[source]¶

Initialize the MatmulLinearOperator.

Parameters:

a (LinearOperator) – The first linear operator to be multiplied.
b (LinearOperator) – The second linear operator to be multiplied.
is_hermitian (bool) – Whether the result is Hermitian. Default to False.

class MatrixLinearOperator(*args, **kwargs)[source]¶

Class method to wrap a matrix into LinearOperator. It is a standard linear operator, used in many operations.

Examples

>>> import torch
>>> seed = torch.manual_seed(100)
>>> mat = torch.rand(3, 2)
>>> linop = MatrixLinearOperator(mat, is_hermitian=False)
>>> print(linop)
MatrixLinearOperator with shape (3, 2):
   tensor([[0.1117, 0.8158],
           [0.2626, 0.4839],
           [0.6765, 0.7539]])
>>> x = torch.rand(2, 2)
>>> linop.mm(x)
tensor([[0.1991, 0.1011],
        [0.1696, 0.0684],
        [0.3345, 0.1180]])
>>> x = torch.rand(3, 2)
>>> linop.mv(x)
tensor([[0.6137, 0.3879, 0.6369],
        [0.7220, 0.5680, 1.0753],
        [0.7821, 0.5460, 0.9626]])

__init__(mat: Tensor, is_hermitian: bool) → None[source]¶

Initialize the MatrixLinearOperator.

Parameters:

mat (torch.Tensor) – The matrix to be wrapped.
is_hermitian (bool) – Indicating if the matrix is Hermitian. If None, the symmetry will be checked. If supplied as a bool, there is no check performed.

class PureFunction(fcntocall: Callable)[source]¶

PureFunction class wraps methods to make it stateless and expose the pure function to take inputs of the original inputs (params) and the object’s states (objparams). For functions, this class only acts as a thin wrapper.

Restore stack stores list of (objparams, identical) everytime the objparams are set, it will store the old objparams and indication if the old and new objparams are identical.

For Using this Class we first need to implement _get_all_obj_params_init and _set_all_obj_params.

Examples

>>> class WrapperFunction(PureFunction):
...     def _get_all_obj_params_init(self):
...         return []
...     def _set_all_obj_params(self, objparams):
...         pass
>>> def fcn(x, y):
...    return x + y
>>> pfunc = WrapperFunction(fcn)
>>> pfunc(1, 2)
3

__init__(fcntocall: Callable)[source]¶

Initialize the PureFunction.

Parameters:: fcntocall (Callable) – The function to be wrapped

objparams() → List[source]¶

Get the current object parameters.

Returns:: The current object parameters
Return type:: List

set_objparams(objparams: List)[source]¶

Set the object parameters.

Parameters:

objparams (List) – The object parameters to be set
TODO (check if identical with current object parameters) –

restore_objparams()[source]¶: Restore the object parameters to the previous state.

useobjparams(objparams: List)[source]¶

Context manager to temporarily set the object parameters.

Parameters:: objparams (List) – The object parameters to be set temporarily

disable_state_change()[source]¶: Context manager to temporarily disable the state change.

class FunctionPureFunction(fcntocall: Callable)[source]¶

Implementation of PureFunction for functions. It just acts as a thin wrapper for the function.

Examples

>>> def fcn(x, y):
...     return x + y
>>> pfunc = FunctionPureFunction(fcn)
>>> pfunc(1, 2)
3

class EditableModulePureFunction(obj: EditableModule, method: Callable)[source]¶

Implementation of PureFunction for EditableModule.

Examples

>>> import torch
>>> from deepchem.utils.differentiation_utils import EditableModule, get_pure_function
>>> class A(EditableModule):
...     def __init__(self, a):
...         self.b = a*a
...     def mult(self, x):
...         return self.b * x
...     def getparamnames(self, methodname, prefix=""):
...         if methodname == "mult":
...             return [prefix+"b"]
...         else:
...             raise KeyError()
>>> B = A(4)
>>> m = get_pure_function(B.mult)
>>> m.set_objparams([3])
>>> m(2)
6

__init__(obj: EditableModule, method: Callable)[source]¶

Initialize the EditableModulePureFunction.

Parameters:

obj (EditableModule) – The object to be wrapped
method (Callable) – The method to be wrapped

class TorchNNPureFunction(obj: Module, method: Callable)[source]¶

Implementation of PureFunction for torch.nn.Module.

Examples

>>> import torch
>>> from deepchem.utils.differentiation_utils import get_pure_function
>>> class A(torch.nn.Module):
...     def __init__(self, a):
...         super().__init__()
...         self.b = torch.nn.Parameter(torch.tensor(a*a))
...     def forward(self, x):
...         return self.b * x
>>> B = A(4.)
>>> m = get_pure_function(B.forward)
>>> m.set_objparams([3.])
>>> m(2)
6.0

__init__(obj: Module, method: Callable)[source]¶

Initialize the TorchNNPureFunction.

Parameters:

obj (torch.nn.Module) – Object to be wrapped
method (Callable) – Method to be wrapped

class PureFunction(fcntocall: Callable)[source]¶

PureFunction class wraps methods to make it stateless and expose the pure function to take inputs of the original inputs (params) and the object’s states (objparams). For functions, this class only acts as a thin wrapper.

Restore stack stores list of (objparams, identical) everytime the objparams are set, it will store the old objparams and indication if the old and new objparams are identical.

For Using this Class we first need to implement _get_all_obj_params_init and _set_all_obj_params.

Examples

>>> class WrapperFunction(PureFunction):
...     def _get_all_obj_params_init(self):
...         return []
...     def _set_all_obj_params(self, objparams):
...         pass
>>> def fcn(x, y):
...    return x + y
>>> pfunc = WrapperFunction(fcn)
>>> pfunc(1, 2)
3

__init__(fcntocall: Callable)[source]¶

Initialize the PureFunction.

Parameters:: fcntocall (Callable) – The function to be wrapped

objparams() → List[source]¶

Get the current object parameters.

Returns:: The current object parameters
Return type:: List

set_objparams(objparams: List)[source]¶

Set the object parameters.

Parameters:

objparams (List) – The object parameters to be set
TODO (check if identical with current object parameters) –

restore_objparams()[source]¶: Restore the object parameters to the previous state.

useobjparams(objparams: List)[source]¶

Context manager to temporarily set the object parameters.

Parameters:: objparams (List) – The object parameters to be set temporarily

disable_state_change()[source]¶: Context manager to temporarily disable the state change.

_check_identical_objs(objs1: List, objs2: List) → bool[source]¶

Check if the two lists of objects are identical.

Examples

>>> l1 = [2, 2, 3]
>>> l2 = [1, 2, 3]
>>> _check_identical_objs(l1, l2)
False

Parameters:

objs1 (List) – The first list of objects
objs2 (List) – The second list of objects

Returns:

True if the two lists of objects are identical, False otherwise

Return type:

bool

get_pure_function(fcn) → PureFunction[source]¶

Get the pure function form of the function or method fcn.

Examples

>>> import torch
>>> from deepchem.utils.differentiation_utils import get_pure_function
>>> def fcn(x, y):
...     return x + y
>>> pfunc = get_pure_function(fcn)
>>> pfunc(1, 2)
3

Parameters:: fcn (function or method) – Function or method to be converted into a PureFunction by exposing the hidden parameters affecting its outputs.
Returns:: The pure function wrapper
Return type:: PureFunction

set_default_option(defopt: Dict, opt: Dict) → Dict[source]¶

return a dictionary based on the options and if no item from option, take it from defopt make a shallow copy to detach the results from defopt

Examples

>>> set_default_option({'a': 1, 'b': 2}, {'a': 3})
{'a': 3, 'b': 2}

Parameters:

defopt (dict) – Default options
opt (dict) – Options

Returns:

Merged options

Return type:

dict

get_and_pop_keys(dct: Dict, keys: List) → Dict[source]¶

Get and pop keys from a dictionary

Examples

>>> get_and_pop_keys({'a': 1, 'b': 2}, ['a'])
{'a': 1}

Parameters:

dct (dict) – Dictionary to pop from
keys (list) – Keys to pop

Returns:

Dictionary containing the popped keys

Return type:

dict

get_method(algname: str, methods: Mapping[str, Callable], method: str | Callable) → Callable[source]¶

Get a method from a dictionary of methods

Examples

>>> get_method('foo', {'bar': lambda: 1}, 'bar')()
1

Parameters:

algname (str) – Name of the algorithm
methods (dict) – Dictionary of methods
method (str or callable) – Method to get

Returns:

The method

Return type:

callable

dummy_context_manager()[source]¶: Dummy context manager

assert_runtime(cond, msg='')[source]¶

Assert at runtime

Examples

>>> assert_runtime(False, "This is a test")
Traceback (most recent call last):
...
RuntimeError: This is a test

Parameters:

cond (bool) – Condition to assert
msg (str) – Message to raise if condition is not met

Raises:

RuntimeError – If condition is not met

_set_initial_v(vinit_type: str, dtype: dtype, device: device, batch_dims: Sequence, na: int, nguess: int, M: LinearOperator | None = None) → Tensor[source]¶

Set the initial guess for the eigenvectors.

Examples

>>> import torch
>>> vinit_type = "eye"
>>> dtype = torch.float64
>>> device = torch.device("cpu")
>>> batch_dims = (2, 3)
>>> na = 4
>>> nguess = 2
>>> M = None
>>> V = _set_initial_v(vinit_type, dtype, device, batch_dims, na, nguess, M)
>>> V
tensor([[[[1., 0.],
          [0., 1.],
          [0., 0.],
          [0., 0.]],

         [[1., 0.],
          [0., 1.],
          [0., 0.],
          [0., 0.]],

         [[1., 0.],
          [0., 1.],
          [0., 0.],
          [0., 0.]]],


        [[[1., 0.],
          [0., 1.],
          [0., 0.],
          [0., 0.]],

         [[1., 0.],
          [0., 1.],
          [0., 0.],
          [0., 0.]],

         [[1., 0.],
          [0., 1.],
          [0., 0.],
          [0., 0.]]]], dtype=torch.float64)

Parameters:

vinit_type (str) – Mode of the initial guess ("randn", "rand", "eye")
dtype (torch.dtype) – Data type of the initial guess.
device (torch.device) – Device of the initial guess.
batch_dims (Sequence) – Batch dimensions of the initial guess.
na (int) – Number of basis functions.
nguess (int) – Number of initial guesses.
M (Optional[LinearOperator] (default None)) – The overlap matrix. If None, identity matrix is used.

Returns:

V – Initial guess for the eigenvectors.

Return type:

torch.Tensor

_take_eigpairs(eival: Tensor, eivec: Tensor, neig: int, mode: str)[source]¶

Take the eigenpairs from the eigendecomposition.

Examples

>>> import torch
>>> eival = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
>>> eivec = torch.tensor([[[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]],
...                       [[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]])
>>> neig = 2
>>> mode = "lowest"
>>> eival, eivec = _take_eigpairs(eival, eivec, neig, mode)
>>> eival
tensor([[1., 2.],
        [4., 5.]])
>>> eivec
tensor([[[1., 2.],
         [4., 5.],
         [7., 8.]],

        [[1., 2.],
         [4., 5.],
         [7., 8.]]])

Parameters:

eival (torch.Tensor) – Eigenvalues of the linear operator. Shape: (*BV, na).
eivec (torch.Tensor) – Eigenvectors of the linear operator. Shape: (*BV, na, na).
neig (int) – Number of eigenvalues and eigenvectors to be calculated.
mode (str) – Mode of the eigenvalues to be calculated ("lowest", "uppest")

Returns:

eival (torch.Tensor) – Eigenvalues of the linear operator.
eivec (torch.Tensor) – Eigenvectors of the linear operator.

exacteig(A: LinearOperator, neig: int, mode: str, M: LinearOperator | None) → Tuple[Tensor, Tensor][source]¶

Eigendecomposition using explicit matrix construction. No additional option for this method.

Examples

>>> import torch
>>> import numpy as np
>>> from deepchem.utils.differentiation_utils import LinearOperator
>>> A = LinearOperator.m(torch.rand(2, 2))
>>> neig = 2
>>> mode = "lowest"
>>> M = None
>>> evals, evecs = exacteig(A, neig, mode, M)
>>> evals.shape
torch.Size([2])
>>> evecs.shape
torch.Size([2, 2])

Parameters:

A (LinearOperator) – Linear operator to be diagonalized. Shape: (*BA, q, q).
neig (int) – Number of eigenvalues and eigenvectors to be calculated.
mode (str) – Mode of the eigenvalues to be calculated ("lowest", "uppest")
M (Optional[LinearOperator] (default None)) – The overlap matrix. If None, identity matrix is used. Shape: (*BM, q, q).

Returns:

evals (torch.Tensor) – Eigenvalues of the linear operator.
evecs (torch.Tensor) – Eigenvectors of the linear operator.

Warning

As this method construct the linear operators explicitly, it might requires a large memory.

degen_symeig(*args, **kwargs)[source]¶

A wrapper for torch.linalg.eigh to avoid complex eigenvalues for degenerate case.

Examples

>>> import torch
>>> import numpy as np
>>> from deepchem.utils.differentiation_utils import LinearOperator
>>> A = LinearOperator.m(torch.rand(2, 2))
>>> evals, evecs = degen_symeig.apply(A.fullmatrix())
>>> evals.shape
torch.Size([2])
>>> evecs.shape
torch.Size([2, 2])

davidson(A: LinearOperator, neig: int, mode: str, M: LinearOperator | None = None, max_niter: int = 1000, nguess: int | None = None, v_init: str = 'randn', max_addition: int | None = None, min_eps: float = 1e-06, verbose: bool = False, **unused) → Tuple[Tensor, Tensor][source]¶

Using Davidson method for large sparse matrix eigendecomposition [2]_.

Examples

>>> import torch
>>> import numpy as np
>>> from deepchem.utils.differentiation_utils import LinearOperator
>>> A = LinearOperator.m(torch.rand(2, 2))
>>> neig = 2
>>> mode = "lowest"
>>> eigen_val, eigen_vec = davidson(A, neig, mode)

Parameters:

A (LinearOperator) – Linear operator to be diagonalized. Shape: (*BA, q, q).
neig (int) – Number of eigenvalues and eigenvectors to be calculated.
mode (str) – Mode of the eigenvalues to be calculated ("lowest", "uppest")
M (Optional[LinearOperator] (default None)) – The overlap matrix. If None, identity matrix is used. Shape: (*BM, q, q).
max_niter (int) – Maximum number of iterations
v_init (str) – Mode of the initial guess ("randn", "rand", "eye")
max_addition (int or None) – Maximum number of new guesses to be added to the collected vectors. If None, set to neig.
min_eps (float) – Minimum residual error to be stopped
verbose (bool) – Option to be verbose

Returns:

evals (torch.Tensor) – Eigenvalues of the linear operator.
evecs (torch.Tensor) – Eigenvectors of the linear operator.

References

lsymeig(A: LinearOperator, neig: int | None = None, M: LinearOperator | None = None, bck_options: Mapping[str, Any] = {}, method: Callable | str | None = None, **fwd_options) → Tuple[Tensor, Tensor][source]¶: Obtain neig lowest eigenvalues and eigenvectors of a linear operator

usymeig(A: LinearOperator, neig: int | None = None, M: LinearOperator | None = None, bck_options: Mapping[str, Any] = {}, method: Callable | str | None = None, **fwd_options) → Tuple[Tensor, Tensor][source]¶: Obtain neig uppest eigenvalues and eigenvectors of a linear operator

symeig(A: LinearOperator, neig: int | None = None, mode: str = 'lowest', M: LinearOperator | None = None, bck_options: Mapping[str, Any] = {}, method: Callable | str | None = None, **fwd_options) → Tuple[Tensor, Tensor][source]¶

Obtain neig lowest eigenvalues and eigenvectors of a linear operator,

Examples

>>> import torch
>>> from deepchem.utils.differentiation_utils import LinearOperator
>>> A = LinearOperator.m(torch.tensor([[3, -1j], [1j, 4]]))
>>> evals, evecs = symeig(A)
>>> evals.shape
torch.Size([2])
>>> evecs.shape
torch.Size([2, 2])

\[\mathbf{AX = MXE}\]

where $\mathbf{A}, \mathbf{M}$ are linear operators, $\mathbf{E}$ is a diagonal matrix containing the eigenvalues, and $\mathbf{X}$ is a matrix containing the eigenvectors. This function can handle derivatives for degenerate cases by setting non-zero degen_atol and degen_rtol in the backward option using the expressions in [1]_.

Parameters:

A (LinearOperator) – The linear operator object on which the eigenpairs are constructed. It must be a Hermitian linear operator with shape (*BA, q, q)
neig (int or None) – The number of eigenpairs to be retrieved. If None, all eigenpairs are retrieved
mode (str) – "lowest" or "uppermost"/"uppest". If "lowest", it will take the lowest neig eigenpairs. If "uppest", it will take the uppermost neig.
M (LinearOperator) – The transformation on the right hand side. If None, then M=I. If specified, it must be a Hermitian with shape (*BM, q, q).
bck_options (dict) –
Method-specific options for solve() which used in backpropagation calculation with some additional arguments for computing the backward derivatives:
- degen_atol (float or None): Minimum absolute difference between two eigenvalues to be treated as degenerate. If None, it is torch.finfo(dtype).eps**0.6. If 0.0, no special treatment on degeneracy is applied. (default: None)
- degen_rtol (float or None): Minimum relative difference between two eigenvalues to be treated as degenerate. If None, it is torch.finfo(dtype).eps**0.4. If 0.0, no special treatment on degeneracy is applied. (default: None)
Note: the default values of degen_atol and degen_rtol are going to change in the future. So, for future compatibility, please specify the specific values.
method (str or callable or None) – Method for the eigendecomposition. If None, it will choose "exacteig".
**fwd_options – Method-specific options (see method section below).

Returns:

It will return eigenvalues and eigenvectors with shapes respectively (*BAM, neig) and (*BAM, na, neig), where *BAM is the broadcasted shape of *BA and *BM.

Return type:

tuple of tensors (eigenvalues, eigenvectors)

References

class symeig_torchfcn(*args, **kwargs)[source]¶

A wrapper for symeig to be used in torch.autograd.Function

static forward(ctx, A, neig, mode, M, fwd_options, bck_options, na, *amparams)[source]¶

Calculate the eigenvalues and eigenvectors of a linear operator

Parameters:

A (LinearOperator) – The linear operator object on which the eigenpairs are constructed. It must be a Hermitian linear operator with shape (*BA, q, q)
neig (int) – The number of eigenpairs to be retrieved. If None, all eigenpairs are retrieved
mode (str) – "lowest" or "uppermost"/"uppest". If "lowest", it will take the lowest neig eigenpairs. If "uppest", it will take the uppermost neig.
M (xitorch.LinearOperator) – The transformation on the right hand side. If None, then M=I. If specified, it must be a Hermitian with shape (*BM, q, q).
fwd_options (dict) – Method-specific options (see method section below).
bck_options (dict) – Method-specific options for solve() which used in backpropagation calculation with some additional arguments for computing the backward derivatives: degen_atol and degen_rtol.
na (int) – Number of parameters of A (and M if M is not None)
*amparams (torch.Tensor) – Parameters of A (and M if M is not None)

static backward(ctx, grad_evals, grad_evecs)[source]¶

Calculate the gradient of the eigenvalues and eigenvectors of a linear operator

Parameters:

grad_evals (torch.Tensor) – The gradient of the eigenvalues. Shape: (*BAM, neig)
grad_evecs (torch.Tensor) – The gradient of the eigenvectors. Shape: (*BAM, na, neig)

_check_degen(evals: Tensor, degen_atol: float, degen_rtol: float) → Tuple[Tensor, bool][source]¶

Check the degeneracy of the eigenvalues

Examples

>>> import torch
>>> evals = torch.tensor([1, 1, 2, 3, 3, 3, 4, 5, 5])
>>> degen_atol = 0.1
>>> degen_rtol = 0.1
>>> idx_degen, isdegenerate = _check_degen(evals, degen_atol, degen_rtol)
>>> idx_degen.shape
torch.Size([9, 9])
>>> isdegenerate
True

Parameters:

evals (torch.Tensor) – Eigenvalues of the linear operator. Shape: (*BAM, neig)
degen_atol (float) – Minimum absolute difference between two eigenvalues to be treated as degenerate.
degen_rtol (float) – Minimum relative difference between two eigenvalues to be treated as degenerate.

Returns:

idx_degen (torch.Tensor) – The degeneracy map. Shape: (*BAM, neig, neig)
isdegenerate (bool) – Whether the eigenvalues are degenerate

ortho(A: Tensor, B: Tensor, *, D: Tensor | None = None, M: LinearOperator | None = None, mright: bool = False) → Tensor[source]¶

Orthogonalize A w.r.t. B

Examples

>>> import torch
>>> A = torch.tensor([[1, 2], [3, 4]])
>>> B = torch.tensor([[1, 0], [0, 1]])
>>> ortho(A, B)
tensor([[0, 2],
        [3, 0]])

Parameters:

A (torch.Tensor) – The tensor to be orthogonalized. Shape: (*BAM, na, neig)
B (torch.Tensor) – The tensor to be orthogonalized against. Shape: (*BAM, na, neig)
D (torch.Tensor or None) – The degeneracy map. If None, it is identity matrix. Shape: (*BAM, neig, neig)
M (LinearOperator or None) – The overlap matrix. If None, identity matrix is used. Shape: (*BM, q, q)
mright (bool) – Whether to operate M at the right or at the left

Returns:

The orthogonalized tensor. Shape: (*BAM, na, neig)

Return type:

torch.Tensor

jac(fcn: Callable[[...], Tensor], params: Sequence[Any], idxs: None | int | Sequence[int] = None)[source]¶

Returns the LinearOperator that acts as the jacobian of the params. The shape of LinearOperator is (nout, nin) where nout and nin are the total number of elements in the output and the input, respectively.

Examples

>>> import torch
>>> from deepchem.utils.differentiation_utils import jac
>>> def fcn(x, y):
...     return x * y
>>> x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
>>> y = torch.tensor([4.0, 5.0, 6.0], requires_grad=True)
>>> jac(fcn, [x, y])
[LinearOperator (_Jac) with shape (3, 3), dtype = torch.float32, device = cpu, LinearOperator (_Jac) with shape (3, 3), dtype = torch.float32, device = cpu]

Parameters:

fcn (Callable[...,torch.Tensor]) – Callable with tensor output and arbitrary numbers of input parameters.
params (Sequence[Any]) – List of input parameters of the function.
idxs (int or list of int or None) – List of the parameters indices to get the jacobian. The pointed parameters in params must be tensors and requires_grad. If it is None, then it will return all jacobian for all parameters that are tensor which requires_grad.

Returns:

linops – List of LinearOperator of the jacobian

Return type:

Union[LinearOperator, List]

class _Jac(*args, **kwargs)[source]¶

Jacobian of a function with respect to a parameter in the function.

Examples

>>> import torch
>>> def fcn(x, y):
...     return x * y
>>> x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
>>> y = torch.tensor([4.0, 5.0, 6.0], requires_grad=True)
>>> pfcn = get_pure_function(fcn)
>>> _Jac(pfcn, [x, y], 1)
LinearOperator (_Jac) with shape (3, 3), dtype = torch.float32, device = cpu

__init__(fcn: PureFunction, params: Sequence[Any], idx: int, is_hermitian=False) → None[source]¶

Initialize the _Jac object.

Parameters:

fcn (PureFunction) – The function that will be differentiated.
params (Sequence[Any]) – List of input parameters of the function.
idx (int) – The index of the parameter to be differentiated.
is_hermitian (bool) – If True, then the LinearOperator is hermitian.

_setup_idxs(idxs: None | int | Sequence[int], params: Sequence[Any]) → Sequence[int][source]¶

Check the idxs and return the list of indices.

Examples

>>> import torch
>>> import numpy as np
>>> x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
>>> y = torch.tensor([4.0, 5.0, 6.0], requires_grad=True)
>>> _setup_idxs(None, [x, y])
[0, 1]

Parameters:

idxs (int or list of int or None) – List of the parameters indices to get the jacobian. The pointed parameters in params must be tensors and requires_grad. If it is None, then it will return all jacobian for all parameters that are tensor which requires_grad.
params (Sequence[Any]) – List of input parameters of the function.

Returns:

idxs – List of the parameters indices to get the jacobian.

Return type:

list of int

connect_graph(out: Tensor, params: Sequence[Any])[source]¶

Just to have a dummy graph, in case there is a parameter that is disconnected in calculating df/dy.

Examples

>>> import torch
>>> x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
>>> y = torch.tensor([4.0, 5.0, 6.0], requires_grad=True)
>>> out = x * y
>>> connect_graph(out, [x, y])
tensor([ 4., 10., 18.], grad_fn=<AddBackward0>)

Parameters:

out (torch.Tensor) – The output tensor. It will be added with a dummy graph.
params (Sequence[Any]) – List of parameters that will be added with a dummy graph.

Returns:

out – The output tensor with a dummy graph.

Return type:

torch.Tensor

wrap_gmres(A, B, E=None, M=None, min_eps=1e-09, max_niter=None, **unused)[source]¶

Using SciPy’s gmres method to solve the linear equation.

Examples

>>> import torch
>>> from deepchem.utils.differentiation_utils import LinearOperator
>>> A = LinearOperator.m(torch.tensor([[1., 2], [3, 4]]))
>>> B = torch.tensor([[[5., 6], [7, 8]]])
>>> wrap_gmres(A, B, None, None)
tensor([[[-3.0000, -4.0000],
         [ 4.0000,  5.0000]]])

Parameters:

A (LinearOperator) – The linear operator A to be solved. Shape: (*BA, na, na)
B (torch.Tensor) – Batched matrix B. Shape: (*BB, na, ncols)
E (torch.Tensor or None) – Batched vector E. Shape: (*BE, ncols)
M (LinearOperator or None) – The linear operator M. Shape: (*BM, na, na)
min_eps (float) – Relative tolerance for stopping conditions
max_niter (int or None) – Maximum number of iterations. If None, default to twice of the number of columns of A.

Returns:

The Solution matrix X. Shape: (*BBE, na, ncols)

Return type:

torch.Tensor

exactsolve(A: LinearOperator, B: Tensor, E: Tensor | None, M: LinearOperator | None)[source]¶

Solve the linear equation by contructing the full matrix of LinearOperators.

Examples

>>> import torch
>>> from deepchem.utils.differentiation_utils import LinearOperator
>>> A = LinearOperator.m(torch.tensor([[1., 2], [3, 4]]))
>>> B = torch.tensor([[5., 6], [7, 8]])
>>> exactsolve(A, B, None, None)
tensor([[-3., -4.],
        [ 4.,  5.]])

Parameters:

A (LinearOperator) – The linear operator A to be solved. Shape: (*BA, na, na)
B (torch.Tensor) – Batched matrix B. Shape: (*BB, na, ncols)
E (torch.Tensor or None) – Batched vector E. Shape: (*BE, ncols)
M (LinearOperator or None) – The linear operator M. Shape: (*BM, na, na)

Returns:

The Solution matrix X. Shape: (*BBE, na, ncols)

Return type:

torch.Tensor

Warning

As this method construct the linear operators explicitly, it might requires a large memory.

solve_ABE(A: Tensor, B: Tensor, E: Tensor)[source]¶

Solve the linear equation AX = B - diag(E)X.

Examples

>>> import torch
>>> A = torch.tensor([[1., 2], [3, 4]])
>>> B = torch.tensor([[5., 6], [7, 8]])
>>> E = torch.tensor([1., 2])
>>> solve_ABE(A, B, E)
tensor([[-0.1667,  0.5000],
        [ 2.5000,  3.2500]])

Parameters:

A (torch.Tensor) – The batched matrix A. Shape: (*BA, na, na)
B (torch.Tensor) – The batched matrix B. Shape: (*BB, na, ncols)
E (torch.Tensor) – The batched vector E. Shape: (*BE, ncols)

Returns:

The batched matrix X.

Return type:

torch.Tensor

get_batchdims(A: LinearOperator, B: Tensor, E: Tensor | None, M: LinearOperator | None)[source]¶

Get the batch dimensions of the linear operator and the matrix B

Examples

>>> from deepchem.utils.differentiation_utils import MatrixLinearOperator
>>> import torch
>>> A = MatrixLinearOperator(torch.randn(4, 3, 3), True)
>>> B = torch.randn(3, 3, 2)
>>> get_batchdims(A, B, None, None)
[4]

Parameters:

A (LinearOperator) – The linear operator. It can be a batched linear operator.
B (torch.Tensor) – The matrix B. It can be a batched matrix.
E (Union[torch.Tensor, None]) – The matrix E. It can be a batched matrix.
M (Union[LinearOperator, None]) – The linear operator M. It can be a batched linear operator.

Returns:

The batch dimensions of the linear operator and the matrix B

Return type:

List[int]

setup_precond(precond: LinearOperator | None = None) → Callable[[Tensor], Tensor][source]¶

Setup the preconditioning function

Examples

>>> from deepchem.utils.differentiation_utils import MatrixLinearOperator
>>> import torch
>>> A = MatrixLinearOperator(torch.randn(4, 3, 3), True)
>>> B = torch.randn(4, 3, 2)
>>> cond = setup_precond(A)
>>> cond(B).shape
torch.Size([4, 3, 2])

Parameters:: precond (Optional[LinearOperator]) – The preconditioning linear operator. If None, no preconditioning is applied.
Returns:: The preconditioning function. It takes a tensor and returns a tensor.
Return type:: Callable[[torch.Tensor], torch.Tensor]

dot(r: Tensor, z: Tensor) → Tensor[source]¶

Dot product of two vectors. r and z must have the same shape. Then sums it up across the last dimension.

Examples

>>> import torch
>>> r = torch.tensor([[1, 2], [3, 4]])
>>> z = torch.tensor([[5, 6], [7, 8]])
>>> dot(r, z)
tensor([[26, 44]])

Parameters:

r (torch.Tensor) – The first vector. Shape: (*BR, nr, nc)
z (torch.Tensor) – The second vector. Shape: (*BR, nr, nc)

Returns:

The dot product of r and z. Shape: (*BR, 1, nc)

Return type:

torch.Tensor

gmres(A: LinearOperator, B: Tensor, E: Tensor | None = None, M: LinearOperator | None = None, posdef: bool | None = None, max_niter: int | None = None, rtol: float = 1e-06, atol: float = 1e-08, eps: float = 1e-12, **unused) → Tensor[source]¶

Solve the linear equations using Generalised minial residual method.

Examples

>>> import torch
>>> from deepchem.utils.differentiation_utils import LinearOperator
>>> A = LinearOperator.m(torch.tensor([[1., 2], [3, 4]]))
>>> B = torch.tensor([[5., 6], [7, 8]])
>>> gmres(A, B)
tensor([[0.8959, 1.0697],
        [1.2543, 1.4263]])

Parameters:

A (LinearOperator) – The linear operator A to be solved. Shape: (*BA, na, na)
B (torch.Tensor) – Batched matrix B. Shape: (*BB, na, ncols)
E (torch.Tensor or None) – Batched vector E. Shape: (*BE, ncols)
M (LinearOperator or None) – The linear operator M. Shape: (*BM, na, na)
posdef (bool or None) – Indicating if the operation $\mathbf{AX-MXE}$ a positive definite for all columns and batches. If None, it will be determined by power iterations.
max_niter (int or None) – Maximum number of iteration. If None, it is set to int(1.5 * A.shape[-1])
rtol (float) – Relative tolerance for stopping condition w.r.t. norm of B
atol (float) – Absolute tolerance for stopping condition w.r.t. norm of B
eps (float) – Substitute the absolute zero in the algorithm’s denominator with this value to avoid nan.

Returns:

The solution matrix X. Shape: (*BBE, na, ncols)

Return type:

torch.Tensor

setup_linear_problem(A: LinearOperator, B: Tensor, E: Tensor | None, M: LinearOperator | None, batchdims: Sequence[int], posdef: bool | None, need_hermit: bool) → Tuple[Callable[[Tensor], Tensor], Callable[[Tensor], Tensor], Tensor, bool][source]¶

Setup the linear problem for solving AX = B

Examples

>>> from deepchem.utils.differentiation_utils import MatrixLinearOperator
>>> import torch
>>> A = MatrixLinearOperator(torch.randn(4, 3, 3), True)
>>> B = torch.randn(4, 3, 2)
>>> A_fcn, AT_fcn, B_new, col_swapped = setup_linear_problem(A, B, None, None, [4], None, False)
>>> A_fcn(B).shape
torch.Size([4, 3, 2])

Parameters:

A (LinearOperator) – The linear operator A. It can be a batched linear operator.
B (torch.Tensor) – The matrix B. It can be a batched matrix.
E (Optional[torch.Tensor]) – The matrix E. It can be a batched matrix.
M (Optional[LinearOperator]) – The linear operator M. It can be a batched linear operator.
batchdims (Sequence[int]) – The batch dimensions of the linear operator and the matrix B
posdef (Optional[bool]) – Whether the linear operator is positive definite. If None, it will be estimated.
need_hermit (bool) – Whether the linear operator is Hermitian. If True, it will be estimated.

Returns:

Callable[[torch.Tensor], torch.Tensor],: torch.Tensor, bool]

The function A, its transposed function, the matrix B, and whether the columns of B are swapped.

Return type:

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

safedenom(r: Tensor, eps: float) → Tensor[source]¶

Make sure the denominator is not zero

Examples

>>> import torch
>>> r = torch.tensor([[0., 2], [3, 4]])
>>> safedenom(r, 1e-9)
tensor([[1.0000e-09, 2.0000e+00],
        [3.0000e+00, 4.0000e+00]])

Parameters:

r (torch.Tensor) – The input tensor. Shape: (*BR, nr, nc)
eps (float) – The small number to replace the zero denominator

Returns:

The tensor with non-zero denominator. Shape: (*BR, nr, nc)

Return type:

torch.Tensor

get_largest_eival(Afcn: Callable, x: Tensor) → Tensor[source]¶

Get the largest eigenvalue of the linear operator Afcn

Examples

>>> import torch
>>> def Afcn(x):
...     return 10 * x
>>> x = torch.tensor([[1., 2], [3, 4]])
>>> get_largest_eival(Afcn, x)
tensor([[10., 10.]])

Parameters:

Afcn (Callable) – The linear operator A. It takes a tensor and returns a tensor.
x (torch.Tensor) – The input tensor. Shape: (*, nr, nc)

Returns:

The largest eigenvalue. Shape: (*, 1, nc)

Return type:

torch.Tensor

solve(A: LinearOperator, B: Tensor, E: Tensor | None = None, M: LinearOperator | None = None, bck_options: Mapping[str, Any] = {}, method: Callable | str | None = None, **fwd_options) → Tensor[source]¶

Performing iterative method to solve the equation.

Examples

>>> import torch
>>> from deepchem.utils.differentiation_utils import LinearOperator
>>> A = LinearOperator.m(torch.tensor([[1., 2], [3, 4]]))
>>> B = torch.tensor([[5., 6], [7, 8]])
>>> solve(A, B)
tensor([[-3., -4.],
        [ 4.,  5.]])

\[\mathbf{AX=B}\]

or

\[\mathbf{AX-MXE=B}\]

where $\mathbf{E}$ is a diagonal matrix. This function can also solve batched multiple inverse equation at the same time by applying $\mathbf{A}$ to a tensor $\mathbf{X}$ with shape (...,na,ncols). The applied $\mathbf{E}$ are not necessarily identical for each column.

Parameters:

A (LinearOperator) – A linear operator that takes an input X and produce the vectors in the same space as B. It should have the shape of (*BA, na, na)
B (torch.Tensor) – The tensor on the right hand side with shape (*BB, na, ncols)
E (Union[torch.Tensor, None]) – If a tensor, it will solve $\mathbf{AX-MXE = B}$. It will be regarded as the diagonal of the matrix. Otherwise, it just solves $\mathbf{AX = B}$ and M is ignored. If it is a tensor, it should have shape of (*BE, ncols).
M (Optional[LinearOperator]) – The transformation on the E side. If E is None, then this argument is ignored. If E is not None and M is None, then M=I. If LinearOperator, it must be Hermitian with shape (*BM, na, na).
bck_options (dict) – Options of the iterative solver in the backward calculation.
method (Union[str, Callable, None]) – The method of linear equation solver. If None, it will choose "cg" or "bicgstab" based on the matrices symmetry. Note: default method will be changed quite frequently, so if you want future compatibility, please specify a method.
**fwd_options – Method-specific options

Returns:

The tensor $\mathbf{X}$ that satisfies $\mathbf{AX-MXE=B}$.

Return type:

torch.Tensor

broyden1_solve(fcn: Callable, x0: Tensor, params, method: str, alpha=None, uv0=None, max_rank=None, maxiter=None, f_tol=None, f_rtol=None, x_tol=None, x_rtol=None, line_search=True, verbose=False, custom_terminator=None, **unused)[source]¶

Solve the root finder or linear equation using the first Broyden method [1]_. It can be used to solve minimization by finding the root of the function’s gradient.

Examples

>>> def fcn(x):
...    return x**2 - 4
>>> x0 = torch.tensor(0.0, requires_grad=True)
>>> x = broyden1(fcn, x0)
>>> x
tensor(-2.0000, grad_fn=<ViewBackward0>)

Parameters:

fcn (callable) – The function to solve. It should take a tensor and return a tensor.
x0 (torch.Tensor) – The initial guess of the solution.
params (tuple) – The parameters to pass to the function.

References

_rootfinder_solve(alg: str, A: LinearOperator, B: Tensor, E: Tensor | None = None, M: LinearOperator | None = None, **options)[source]¶

Solve the linear equations using rootfinder algorithm

Examples

>>> import torch
>>> A = torch.tensor([[1., 2], [3, 4]])
>>> B = torch.tensor([[5., 6], [7, 8]])
>>> _rootfinder_solve("broyden1", A, B)
tensor([[-3.0000, -4.0000],
        [ 4.0000,  5.0000]])

Parameters:

alg (str) – The algorithm to use. Currently, only “broyden1” is supported.
A (torch.Tensor) – The matrix A. Shape: (*BA, nr, nr)
B (torch.Tensor) – The matrix B. Shape: (*BB, nr, ncols)
E (torch.Tensor or None) – The matrix E. Shape: (*BE, ncols)
M (torch.Tensor or None) – The matrix M. Shape: (*BM, nr, nr)
options (dict) – The options for the rootfinder algorithm

Returns:

The solution matrix X. Shape: (*BBE, nr, ncols)

Return type:

torch.Tensor

cg(A: LinearOperator, B: Tensor, E: Tensor | None = None, M: LinearOperator | None = None, posdef: bool | None = None, precond: LinearOperator | None = None, max_niter: int | None = None, rtol: float = 1e-06, atol: float = 1e-08, eps: float = 1e-12, resid_calc_every: int = 10, verbose: bool = False, **unused) → Tensor[source]¶

Solve the linear equations using Conjugate-Gradient (CG) method.

Parameters:

A (LinearOperator) – A linear operator that takes an input X and produce the vectors in the same space as B. It should have the shape of (*BA, na, na)
B (torch.Tensor) – The tensor on the right hand side with shape (*BB, na, ncols)
E (Union[torch.Tensor, None]) – If a tensor, it will solve $\mathbf{AX-MXE = B}$. It will be regarded as the diagonal of the matrix. Otherwise, it just solves $\mathbf{AX = B}$ and M is ignored. If it is a tensor, it should have shape of (*BE, ncols).
M (Optional[LinearOperator]) – The transformation on the E side. If E is None, then this argument is ignored. If E is not None and M is None, then M=I.
posdef (bool or None) – Indicating if the operation $\mathbf{AX-MXE}$ a positive definite for all columns and batches. If None, it will be determined by power iterations.
precond (LinearOperator or None) – LinearOperator for the preconditioning. If None, no preconditioner is applied.
max_niter (int or None) – Maximum number of iteration. If None, it is set to int(1.5 * A.shape[-1])
rtol (float) – Relative tolerance for stopping condition w.r.t. norm of B
atol (float) – Absolute tolerance for stopping condition w.r.t. norm of B
eps (float) – Substitute the absolute zero in the algorithm’s denominator with this value to avoid nan.
resid_calc_every (int) – Calculate the residual in its actual form instead of substitution form with this frequency, to avoid rounding error accummulation. If your linear operator has bad numerical precision, set this to be low. If 0, then never calculate the residual in its actual form.
verbose (bool) – Verbosity of the algorithm.

bicgstab(A: LinearOperator, B: Tensor, E: Tensor | None = None, M: LinearOperator | None = None, posdef: bool | None = None, precond_l: LinearOperator | None = None, precond_r: LinearOperator | None = None, max_niter: int | None = None, rtol: float = 1e-06, atol: float = 1e-08, eps: float = 1e-12, verbose: bool = False, resid_calc_every: int = 10, **unused) → Tensor[source]¶

Solve the linear equations using stabilized Biconjugate-Gradient method.

Parameters:

posdef (bool or None) – Indicating if the operation $\mathbf{AX-MXE}$ a positive definite for all columns and batches. If None, it will be determined by power iterations.
precond_l (LinearOperator or None) – LinearOperator for the left preconditioning. If None, no preconditioner is applied.
precond_r (LinearOperator or None) – LinearOperator for the right preconditioning. If None, no preconditioner is applied.
max_niter (int or None) – Maximum number of iteration. If None, it is set to int(1.5 * A.shape[-1])
rtol (float) – Relative tolerance for stopping condition w.r.t. norm of B
atol (float) – Absolute tolerance for stopping condition w.r.t. norm of B
eps (float) – Substitute the absolute zero in the algorithm’s denominator with this value to avoid nan.
resid_calc_every (int) – Calculate the residual in its actual form instead of substitution form with this frequency, to avoid rounding error accummulation. If your linear operator has bad numerical precision, set this to be low. If 0, then never calculate the residual in its actual form.
verbose (bool) – Verbosity of the algorithm.

class solve_torchfcn(*args, **kwargs)[source]¶

static forward(ctx, A, B, E, M, method, fwd_options, bck_options, na, *all_params)[source]¶

Forward calculation of the solve function.

Parameters:

A (LinearOperator) – A linear operator that takes an input X and produce the vectors in the same space as B. It should have the shape of (*BA, na, na)
B (torch.Tensor) – The tensor on the right hand side with shape (*BB, na, ncols)
E (Union[torch.Tensor, None]) – If a tensor, it will solve $\mathbf{AX-MXE = B}$. It will be regarded as the diagonal of the matrix. Otherwise, it just solves $\mathbf{AX = B}$ and M is ignored. If it is a tensor, it should have shape of (*BE, ncols).
M (Optional[LinearOperator]) – The transformation on the E side. If E is None, then this argument is ignored. If E is not None and M is None, then M=I.
method (Union[str, Callable, None]) – The method of linear equation solver. If None, it will choose "cg" or "bicgstab" based on the matrices symmetry. Note: default method will be changed quite frequently, so if you want future compatibility, please specify a method.
fwd_options – Method-specific options
bck_options (dict) – Options of the iterative solver in the backward calculation.
na (int) – Number of parameters of A
all_params (Sequence[torch.Tensor]) – All the parameters of M and A

static backward(ctx, grad_x)[source]¶

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

anderson_acc(fcn: Callable[[...], Tensor], x0: Tensor, params: List, feat_ndims: int = 1, msize: int = 5, beta: float = 1.0, lmbda: float = 0.0001, maxiter=None, f_tol=None, f_rtol=None, x_tol=None, x_rtol=None, custom_terminator=None, verbose: bool = False) → Tensor[source]¶

Solve the equilibrium (or fixed-point iteration) problem using Anderson acceleration.

Examples

>>> import torch
>>> def fcn(x):
...     return x
>>> x0 = torch.tensor([0.0], requires_grad=True)
>>> x = anderson_acc(fcn, x0, [], 2, 10, maxiter=1000)
>>> x
tensor([0.], requires_grad=True)

Parameters:

feat_ndims (int) – The number of dimensions at the end that describe the features (i.e. non-batch dimensions)
msize (int) – The maximum number of previous iterations we should save for the algorithm
beta (float) – The damped or overcompensated parameters
lmbda (float) – Small number to ensure invertability of the matrix
maxiter (int or None) – Maximum number of iterations, or inf if it is set to None.
f_tol (float or None) – The absolute tolerance of the norm of the output f - x.
f_rtol (float or None) – The relative tolerance of the norm of the output f - x.
x_tol (float or None) – The absolute tolerance of the norm of the input x.
x_rtol (float or None) – The relative tolerance of the norm of the input x.
verbose (bool) – Options for verbosity

References

gd(fcn: Callable[[...], Tensor], x0: Tensor, params: List, step: float = 0.001, gamma: float = 0.9, maxiter: int = 1000, f_tol: float = 0.0, f_rtol: float = 1e-08, x_tol: float = 0.0, x_rtol: float = 1e-08, verbose=False, terminate=False, **unused)[source]¶

Vanilla gradient descent with momentum. The stopping conditions use OR criteria. The update step is following the equations below.

Examples

>>> import torch
>>> from deepchem.utils.differentiation_utils.optimize.minimizer import gd
>>> def fcn(x):
...     return (x - 2) ** 2, 2 * (x - 2)
>>> x0 = torch.tensor(0.0, requires_grad=True)
>>> x = gd(fcn, x0, [])
>>> x
tensor(2.0000)

\[\begin{split}\mathbf{v}_{t+1} &= \gamma \mathbf{v}_t - \eta \nabla_{\mathbf{x}} f(\mathbf{x}_t) \\ \mathbf{x}_{t+1} &= \mathbf{x}_t + \mathbf{v}_{t+1}\end{split}\]

Parameters:

fcn (callable) – The objective function to minimize. It should take a tensor and return a tensor and its gradient.
x0 (torch.Tensor) – The initial guess.
step (float) – The step size towards the steepest descent direction, i.e. $\eta$ in the equations above.
gamma (float) – The momentum factor, i.e. $\gamma$ in the equations above.
maxiter (int) – Maximum number of iterations.
f_tol (float or None) – The absolute tolerance of the output f.
f_rtol (float or None) – The relative tolerance of the output f.
x_tol (float or None) – The absolute tolerance of the norm of the input x.
x_rtol (float or None) – The relative tolerance of the norm of the input x.
terminate (bool (default False)) – Whether to use termination condition or, keep on running the minimizer.

adam(fcn: Callable[[...], Tensor], x0: Tensor, params: List, step: float = 0.001, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-08, maxiter: int = 1000, f_tol: float = 0.0, f_rtol: float = 1e-08, x_tol: float = 0.0, x_rtol: float = 1e-08, verbose=False, terminate=False, **unused)[source]¶

Adam optimizer by Kingma & Ba (2015). The stopping conditions use OR criteria. The update step is following the equations below.

Examples

>>> from deepchem.utils.differentiation_utils.optimize.minimizer import adam
>>> def fcn(x):
...     return (x - 4) * 2, (x * 2) + 3
>>> x0 = torch.tensor(0.0, requires_grad=True)
>>> x = adam(fcn, x0, [], maxiter=10000)
>>> x
tensor(-1.5000)

\[\begin{split}\mathbf{g}_t &= \nabla_{\mathbf{x}} f(\mathbf{x}_{t-1}) \\ \mathbf{m}_t &= \beta_1 \mathbf{m}_{t-1} + (1 - \beta_1) \mathbf{g}_t \\ \mathbf{v}_t &= \beta_2 \mathbf{v}_{t-1} + (1 - \beta_2) \mathbf{g}_t^2 \\ \hat{\mathbf{m}}_t &= \mathbf{m}_t / (1 - \beta_1^t) \\ \hat{\mathbf{v}}_t &= \mathbf{v}_t / (1 - \beta_2^t) \\ \mathbf{x}_t &= \mathbf{x}_{t-1} - \alpha \hat{\mathbf{m}}_t / (\sqrt{\hat{\mathbf{v}}_t} + \epsilon)\end{split}\]

Parameters:

step (float) – The step size towards the descent direction, i.e. $\alpha$ in the equations above.
beta1 (float) – Exponential decay rate for the first moment estimate.
beta2 (float) – Exponential decay rate for the first moment estimate.
eps (float) – Small number to prevent division by 0.
maxiter (int) – Maximum number of iterations.
f_tol (float or None) – The absolute tolerance of the output f.
f_rtol (float or None) – The relative tolerance of the output f.
x_tol (float or None) – The absolute tolerance of the norm of the input x.
x_rtol (float or None) – The relative tolerance of the norm of the input x.
terminate (bool (default False)) – Whether to use the termination condition, or keep running the minimizer.

TerminationCondition(f_tol: float, f_rtol: float, x_tol: float, x_rtol: float, verbose: bool)[source]¶

The class to handle the stopping conditions.

Examples

>>> stop_cond = TerminationCondition(1e-8, 1e-8, 1e-8, 1e-8, True)

_nonlin_solver(fcn: Callable, x0: Tensor, params, method: str, alpha=None, uv0=None, max_rank=None, maxiter=None, f_tol=None, f_rtol=None, x_tol=None, x_rtol=None, line_search=True, verbose=False, custom_terminator=None, **unused)[source]¶

Parameters:

alpha (float or None) – The initial guess of inverse Jacobian is - alpha * I + u v^T.
uv0 (tuple of tensors or str or None) – The initial guess of inverse Jacobian is - alpha * I + u v^T. If "svd", then it uses 1-rank svd to obtain u and v. If None, then u and v are zeros.
max_rank (int or None) – The maximum rank of inverse Jacobian approximation. If None, it is inf.
maxiter (int or None) – Maximum number of iterations, or inf if it is set to None.
f_tol (float or None) – The absolute tolerance of the norm of the output f.
f_rtol (float or None) – The relative tolerance of the norm of the output f.
x_tol (float or None) – The absolute tolerance of the norm of the input x.
x_rtol (float or None) – The relative tolerance of the norm of the input x.
line_search (bool or str) – Options to perform line search. If True, it is set to "armijo".
verbose (bool) – Options for verbosity

broyden1(fcn: Callable, x0: Tensor, params, method: str, alpha=None, uv0=None, max_rank=None, maxiter=None, f_tol=None, f_rtol=None, x_tol=None, x_rtol=None, line_search=True, verbose=False, custom_terminator=None, **unused)[source]¶

Solve the root finder or linear equation using the first Broyden method [1]_. It can be used to solve minimization by finding the root of the function’s gradient.

Examples

>>> def fcn(x):
...    return x**2 - 4
>>> x0 = torch.tensor(0.0, requires_grad=True)
>>> x = broyden1(fcn, x0)
>>> x
tensor(-2.0000, grad_fn=<ViewBackward0>)

Parameters:

fcn (callable) – The function to solve. It should take a tensor and return a tensor.
x0 (torch.Tensor) – The initial guess of the solution.
params (tuple) – The parameters to pass to the function.

References

broyden2(fcn: Callable, x0: Tensor, params, method: str, alpha=None, uv0=None, max_rank=None, maxiter=None, f_tol=None, f_rtol=None, x_tol=None, x_rtol=None, line_search=True, verbose=False, custom_terminator=None, **unused)[source]¶

Solve the root finder or linear equation using the second Broyden method [2]_. It can be used to solve minimization by finding the root of the function’s gradient.

Examples

>>> def fcn(x):
...    return x**2 - 4
>>> x0 = torch.tensor(0.0, requires_grad=True)
>>> x = broyden1(fcn, x0)
>>> x
tensor(-2.0000, grad_fn=<ViewBackward0>)

Parameters:

fcn (callable) – The function to solve. It should take a tensor and return a tensor.
x0 (torch.Tensor) – The initial guess of the solution.
params (tuple) – The parameters to pass to the function.

References

linearmixing(fcn: Callable, x0: Tensor, params=(), alpha=None, maxiter=None, f_tol=None, f_rtol=None, x_tol=None, x_rtol=None, line_search=True, verbose=False, **unused)[source]¶

Solve the root finding problem by approximating the inverse of Jacobian to be a constant scalar.

Examples

>>> def fcn(x):
...    return x**2 - 4
>>> x0 = torch.tensor(0.0, requires_grad=True)
>>> x = broyden1(fcn, x0)
>>> x
tensor(-2.0000, grad_fn=<ViewBackward0>)

Parameters:

fcn (Callable) – The function to solve. It should take a tensor and return a tensor.
x0 (torch.Tensor) – The initial guess of the solution.
params (tuple) – The parameters to pass to the function.
alpha (float or None) – The initial guess of inverse Jacobian is -alpha * I.
maxiter (int or None) – Maximum number of iterations, or inf if it is set to None.
f_tol (float or None) – The absolute tolerance of the norm of the output f.
f_rtol (float or None) – The relative tolerance of the norm of the output f.
x_tol (float or None) – The absolute tolerance of the norm of the input x.
x_rtol (float or None) – The relative tolerance of the norm of the input x.
line_search (bool or str) – Options to perform line search. If True, it is set to "armijo".
verbose (bool) – Options for verbosity

References

_safe_norm(v)[source]¶: Compute the norm of a vector, checking for finite values.

_nonline_line_search(func: Callable, x: Tensor, y: Tensor, dx: Tensor, search_type='armijo', rdiff=1e-08, smin=0.01)[source]¶

Find a suitable step length for a line search.

Parameters:

func (Callable) – The function to minimize.
x (torch.Tensor) – The current point.
y (torch.Tensor) – The function value at the current point.
dx (torch.Tensor) – The search direction.
search_type (str) – The type of line search to perform. Currently, only “armijo” is supported.
rdiff (float) – The relative difference to compute the derivative.
smin (float) – The minimum step length to take.

Returns:

s (float) – The step length.
x (torch.Tensor) – The new point.
y (torch.Tensor) – The function value at the new point.
y_norm (float) – The norm of the function value at the new point.

_scalar_search_armijo(phi: Callable, phi0: float, derphi0: float, c1: float = 0.0001, alpha0=1, amin=0, max_niter=20)[source]¶

Minimize over alpha, the function phi(s) at the current point and the derivative derphi(s) at the current point.

Parameters:

phi (callable) – The function to minimize.
phi0 (float) – The value of phi at 0.
derphi0 (float) – The value of the derivative of phi at 0.
c1 (float) – The Armijo condition parameter.
alpha0 (float) – The initial guess of the step length.
amin (float) – The minimum step length to take.
max_niter (int) – The maximum number of iterations to take.

Returns:

alpha (float) – The step length.
phi (float) – The value of the function at the step length.

TerminationCondition(f_tol: float, f_rtol: float, f0_norm: float, x_tol: float, x_rtol: float)[source]¶: Class to check the termination condition of the root finder.

class Jacobian[source]¶

Base class for the Jacobians used in rootfinder algorithms.

A Jacobian can best be defined as a determinant which is defined for a finite number of functions of the same number of variables in which each row consists of the first partial derivatives of the same function with respect to each of the variables.

References

[1].. https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant [2].. Kasim, Muhammad & Vinko, Sam. (2020). xi$-torch: differentiable scientific computing library.

abstract setup(x0: Tensor, y0: Tensor, func: Callable)[source]¶: Setup the Jacobian for the rootfinder.

abstract solve(v: Tensor, tol: Any = 0)[source]¶: Solve the linear system J dx = v.

abstract update(x: Tensor, y: Tensor)[source]¶: Update the Jacobian approximation.

class BroydenFirst(alpha: Tensor | None = None, uv0: Any | None = None, max_rank: float | None = None)[source]¶

Approximating the Jacobian based on Broyden’s first approximation.

Examples

>>> from deepchem.utils.differentiation_utils.optimize.jacobian import BroydenFirst
>>> jacobian = BroydenFirst()
>>> x0 = torch.tensor([1.0, 1.0], requires_grad=True)
>>> def func(x):
...     return torch.tensor([x[0]**2 + x[1]**2 - 1.0, x[0] - x[1]])
>>> y0 = func(x0)
>>> v = torch.tensor([1.0, 1.0])
>>> jacobian.setup(x0, y0, func)
>>> jacobian.solve(v)
tensor([-0.7071, -0.7071], grad_fn=<MulBackward0>)

References

[1].. B.A. van der Rotten, PhD thesis,: “A limited memory Broyden method to solve high-dimensional systems of nonlinear equations”. Mathematisch Instituut, Universiteit Leiden, The Netherlands (2003).

__init__(alpha: Tensor | None = None, uv0: Any | None = None, max_rank: float | None = None)[source]¶

The initial guess of inverse Jacobian is -alpha * I + u v^T. max_rank indicates the maximum rank of the Jacoabian before reducing it

Parameters:

alpha (Union[torch.Tensor, None]) – The initial guess of inverse Jacobian is -alpha * I. If None, it is set to -1.0.
uv0 (tuple, optional) – The initial guess of the inverse Jacobian.
max_rank (Union[float, None]) – The maximum rank of the Jacobian before reducing it. If None, it is set to inf.

setup(x0: Tensor, y0: Tensor, func: Callable)[source]¶

Setup the Jacobian for the rootfinder.

Parameters:

x0 – The initial guess of the root.
y0 – The function value at the initial guess.
func – The function to find the root.

solve(v: Tensor, tol=0) → Tensor[source]¶

Solve the linear system J dx = v.

Parameters:

v (torch.Tensor) – The right-hand side of the linear system.
tol (torch.Tensor) – The tolerance for the linear system.

Returns:

res – The solution of the linear system.

Return type:

torch.Tensor

update(x: Tensor, y: Tensor)[source]¶

Update the Jacobian approximation.

Parameters:

x (torch.Tensor) – The current point.
y (torch.Tensor) – The function value at the current point.

class BroydenSecond(alpha: Tensor | None = None, uv0: Any | None = None, max_rank: float | None = None)[source]¶

Inverse Jacobian approximation based on Broyden’s second method.

Examples

>>> from deepchem.utils.differentiation_utils.optimize.jacobian import BroydenSecond
>>> jacobian = BroydenSecond()
>>> x0 = torch.tensor([1.0, 1.0], requires_grad=True)
>>> def func(x):
...     return torch.tensor([x[0]**2 + x[1]**2 - 1.0, x[0] - x[1]])
>>> y0 = func(x0)
>>> v = torch.tensor([1.0, 1.0])
>>> jacobian.setup(x0, y0, func)
>>> jacobian.solve(v)
tensor([-0.7071, -0.7071], grad_fn=<MulBackward0>)

References

[1] B.A. van der Rotten, PhD thesis,: “A limited memory Broyden method to solve high-dimensional systems of nonlinear equations”. Mathematisch Instituut, Universiteit Leiden, The Netherlands (2003).

class LinearMixing(alpha: float | None = None)[source]¶

Approximating the Jacobian based on linear mixing. It acts as a simple check for the functionality of the rootfinder.

Examples

>>> from deepchem.utils.differentiation_utils.optimize.jacobian import LinearMixing
>>> jacobian = LinearMixing()
>>> x0 = torch.tensor([1.0, 1.0], requires_grad=True)
>>> def func(x):
...     return torch.tensor([x[0]**2 + x[1]**2 - 1.0, x[0] - x[1]])
>>> y0 = func(x0)
>>> v = torch.tensor([1.0, 1.0])
>>> jacobian.setup(x0, y0, func)
>>> jacobian.solve(v)
tensor([1., 1.])

__init__(alpha: float | None = None)[source]¶

The initial guess of inverse Jacobian is -alpha * I

Parameters:: alpha (float, optional) – The initial guess of inverse Jacobian is -alpha * I. If None, it is set to -1.0.

setup(x0: Tensor, y0: Tensor, func: Callable)[source]¶

Setup the Jacobian for the rootfinder.

Parameters:

x0 (torch.Tensor) – The initial guess of the root.
y0 (torch.Tensor) – The function value at the initial guess.
func (Callable) – The function to find the root.

solve(v: Tensor, tol=0) → Tensor[source]¶

Solve the linear system J dx = v.

Parameters:

v (torch.Tensor) – The right-hand side of the linear system.
tol – The tolerance for the linear system.

update(x: Tensor, y: Tensor)[source]¶

Update the Jacobian approximation.

Parameters:

x (torch.Tensor) – The current point.
y (torch.Tensor) – The function value at the current point.

class LowRankMatrix(alpha: Tensor, uv0, reduce_method: str)[source]¶

represents a matrix of lpha * I + sum_n c_n d_n^T

Examples

>>> from deepchem.utils.differentiation_utils.optimize.jacobian import LowRankMatrix
>>> import torch
>>> alpha = 1.0
>>> uv0 = (torch.tensor([1.0, 1.0]), torch.tensor([1.0, 1.0]))
>>> reduce_method = "restart"
>>> matrix = LowRankMatrix(alpha, uv0, reduce_method)
>>> v = torch.tensor([1.0, 1.0])
>>> matrix.mv(v)
tensor([3., 3.])
>>> matrix.rmv(v)
tensor([3., 3.])

__init__(alpha: Tensor, uv0, reduce_method: str)[source]¶

initialize the matrix

Parameters:

alpha (torch.Tensor) – The coefficient of the identity matrix
uv0 (tuple) – The initial guess of the inverse Jacobian
reduce_method (str) – The method to reduce the rank of the matrix

mv(v: Tensor) → Tensor[source]¶

multiply the matrix with a vector

Parameters:: v (torch.Tensor) – Vector to multiply
Returns:: res – Result of the multiplication
Return type:: torch.Tensor

rmv(v: Tensor) → Tensor[source]¶

multiply the transpose of the matrix with a vector

Parameters:: v (torch.Tensor) – Vector to multiply
Returns:: res – Result of the multiplication
Return type:: torch.Tensor

append(c: Tensor, d: Tensor)[source]¶

append a rank-1 matrix to the matrix

Parameters:

c (torch.Tensor) – The first vector
d (torch.Tensor) – The second vector

Returns:

res – The matrix after appending the rank-1 matrix

Return type:

Union[‘LowRankMatrix’, ‘FullRankMatrix’]

reduce(max_rank: int, **otherparams)[source]¶

reduce the rank of the matrix

Parameters:

max_rank (int) – The maximum rank of the matrix
otherparams – Other parameters

class FullRankMatrix(alpha: Tensor, cns: Any, dns: Any)[source]¶

represents a full rank matrix of lpha * I + sum_n c_n d_n^T

Examples

>>> from deepchem.utils.differentiation_utils.optimize.jacobian import FullRankMatrix
>>> import torch
>>> alpha = 1.0
>>> cns = [torch.tensor([1.0, 1.0]), torch.tensor([1.0, 1.0])]
>>> dns = [torch.tensor([1.0, 1.0]), torch.tensor([1.0, 1.0])]
>>> matrix = FullRankMatrix(alpha, cns, dns)
>>> v = torch.tensor([1.0, 1.0])
>>> matrix.mv(v)
tensor([5., 5.])
>>> matrix.rmv(v)
tensor([5., 5.])

__init__(alpha: Tensor, cns: Any, dns: Any)[source]¶

initialize the matrix

Parameters:

alpha (torch.Tensor) – Coefficient of the identity matrix
cns (List) – List of the first vectors
dns (List) – List of the second vectors

mv(v: Tensor) → Tensor[source]¶

multiply the matrix with a vector

Parameters:: v (torch.Tensor) – The vector to multiply
Returns:: res – The result of the multiplication
Return type:: torch.Tensor

rmv(v: Tensor) → Tensor[source]¶

multiply the transpose of the matrix with a vector

Parameters:: v (torch.Tensor) – The vector to multiply
Returns:: The result of the multiplication
Return type:: torch.Tensor

append(c: Tensor, d: Tensor)[source]¶

append a rank-1 matrix to the matrix

Parameters:

c (torch.Tensor) – The first vector
d (torch.Tensor) – The second vector

Returns:

The matrix after appending the rank-1 matrix

Return type:

FullRankMatrix

reduce(max_rank: int, **kwargs)[source]¶

reduce the rank of the matrix

Parameters:

max_rank (int) – The maximum rank of the matrix
otherparams – Other parameters

rootfinder(fcn: Callable[[...], Tensor], y0: Tensor, params: Sequence[Any] = [], bck_options: Mapping[str, Any] = {}, method: Callable | str | None = None, **fwd_options) → Tensor[source]¶

Solving the rootfinder equation of a given function,

\[\mathbf{0} = \mathbf{f}(\mathbf{y}, \theta)\]

where $\mathbf{f}$ is a function that can be non-linear and produce output of the same shape of $\mathbf{y}$, and $\theta$ is other parameters required in the function. The output of this block is $\mathbf{y}$ that produces the $\mathbf{0}$ as the output.

Parameters:

fcn (callable) – The function $\mathbf{f}$ with output tensor (*ny)
y0 (torch.tensor) – Initial guess of the solution with shape (*ny)
params (list) – Sequence of any other parameters to be put in fcn
bck_options (dict) – Method-specific options for the backward solve (see xitorch.linalg.solve())
method (str or callable or None) – Rootfinder method. If None, it will choose "broyden1".
**fwd_options – Method-specific options (see method section)

Returns:

The solution which satisfies $\mathbf{0} = \mathbf{f}(\mathbf{y},\theta)$ with shape (*ny)

Return type:

torch.tensor

Example

>>> import torch
>>> def func1(y, A):  # example function
...     return torch.tanh(A @ y + 0.1) + y / 2.0
>>> A = torch.tensor([[1.1, 0.4], [0.3, 0.8]]).requires_grad_()
>>> y0 = torch.zeros((2,1))  # zeros as the initial guess
>>> yroot = rootfinder(func1, y0, params=(A,))
>>> print(yroot)
tensor([[-0.0459],
        [-0.0663]], grad_fn=<_RootFinderBackward>)

equilibrium(fcn: Callable[[...], Tensor], y0: Tensor, params: Sequence[Any] = [], bck_options: Mapping[str, Any] = {}, method: Callable | str | None = None, **fwd_options) → Tensor[source]¶

Solving the equilibrium equation of a given function,

\[\mathbf{y} = \mathbf{f}(\mathbf{y}, \theta)\]

where $\mathbf{f}$ is a function that can be non-linear and produce output of the same shape of $\mathbf{y}$, and $\theta$ is other parameters required in the function. The output of this block is $\mathbf{y}$ that produces the same $\mathbf{y}$ as the output.

Parameters:

fcn (callable) – The function $\mathbf{f}$ with output tensor (*ny)
y0 (torch.tensor) – Initial guess of the solution with shape (*ny)
params (list) – Sequence of any other parameters to be put in fcn
bck_options (dict) – Method-specific options for the backward solve (see xitorch.linalg.solve())
method (str or None) – Rootfinder method. If None, it will choose "broyden1".
**fwd_options – Method-specific options (see method section)

Returns:

The solution which satisfies $\mathbf{y} = \mathbf{f}(\mathbf{y},\theta)$ with shape (*ny)

Return type:

torch.tensor

Example

>>> import torch
>>> def func1(y, A):  # example function
...     return torch.tanh(A @ y + 0.1) + y / 2.0
>>> A = torch.tensor([[1.1, 0.4], [0.3, 0.8]]).requires_grad_()
>>> y0 = torch.zeros((2,1))  # zeros as the initial guess
>>> yequil = equilibrium(func1, y0, params=(A,))
>>> print(yequil)
tensor([[ 0.2313],
        [-0.5957]], grad_fn=<_RootFinderBackward>)

Note

This is a direct implementation of finding the root of $\mathbf{g}(\mathbf{y}, \theta) = \mathbf{y} - \mathbf{f}(\mathbf{y}, \theta)$

minimize(fcn: Callable[[...], Tensor], y0: Tensor, params: Sequence[Any] = [], bck_options: Mapping[str, Any] = {}, method: Callable | str | None = None, **fwd_options) → Tensor[source]¶

Solve the unbounded minimization problem:

\[\mathbf{y^*} = \arg\min_\mathbf{y} f(\mathbf{y}, \theta)\]

to find the best $\mathbf{y}$ that minimizes the output of the function $f$.

Parameters:

fcn (callable) – The function to be optimized with output tensor with 1 element.
y0 (torch.tensor) – Initial guess of the solution with shape (*ny)
params (list) – Sequence of any other parameters to be put in fcn
bck_options (dict) – Method-specific options for the backward solve (see xitorch.linalg.solve())
method (str or callable or None) – Minimization method. If None, it will choose "broyden1".
**fwd_options – Method-specific options (see method section)

Returns:

The solution of the minimization with shape (*ny)

Return type:

torch.tensor

Example

>>> import torch
>>> def func1(y, A):  # example function
...     return torch.sum((A @ y)**2 + y / 2.0)
>>> A = torch.tensor([[1.1, 0.4], [0.3, 0.8]]).requires_grad_()
>>> y0 = torch.zeros((2,1))  # zeros as the initial guess
>>> ymin = minimize(func1, y0, params=(A,))
>>> print(ymin)
tensor([[-0.0519],
        [-0.2684]], grad_fn=<_RootFinderBackward>)

class _RootFinder(*args, **kwargs)[source]¶

static forward(ctx, fcn, y0, fwd_fcn, alg_type, options, bck_options, nparams, *allparams)[source]¶

Forward method for the rootfinder, minimizer, and equilibrium

Parameters:

fcn – a function that returns what has to be 0 (will be used in the backward, not used in the forward). For minimization, it is the gradient
y0 – initial guess
fwd_fcn – a function that will be executed in the forward method (unused in the backward)
alg_type – the type of algorithm: “rootfinder”, “minimizer”, or “equilibrium”
options – options for the forward method
bck_options – options for the backward method
nparams – number of parameters
allparams – all parameters (including the non-tensor parameters)
minimization (This class is also used for) –
might (where fcn and fwd_fcn) –
different (be slightly) –

Returns:

The solution of the rootfinder, minimizer, or equilibrium

Return type:

torch.tensor

static backward(ctx, grad_yout)[source]¶

Backward method for the rootfinder, minimizer, and equilibrium

Parameters:: grad_yout (torch.tensor) – the gradient of the output of the rootfinder, minimizer, or equilibrium
Returns:: The gradients of the parameters
Return type:: tuple

_get_rootfinder_default_method(method: Callable | str | None = None) → str | Callable | None[source]¶

Get the default method for the rootfinder, minimizer, and equilibrium

Examples

>>> _get_rootfinder_default_method(None)
'broyden1'

Parameters:: method (str or None) – The method name
Returns:: The method name
Return type:: str

_get_equilibrium_default_method(method: Callable | str | None = None) → str | Callable | None[source]¶

Get the default method for the equilibrium

Examples

>>> _get_equilibrium_default_method(None)
'broyden1'

Parameters:: method (str or None) – The method name
Returns:: The method name
Return type:: str

_get_minimizer_default_method(method: Callable | str | None = None) → str | Callable | None[source]¶

Get the default method for the minimizer

Examples

>>> _get_minimizer_default_method(None)
'broyden1'

Parameters:: method (str or None) – The method name
Returns:: The method name
Return type:: str

class _Tableau(c: List[float], b: List[float], a: List[List[float]])[source]¶

To specify a particular method, one needs to provide the integer s (the number of stages), and the coefficients a[i,j] (for 1 ≤ j < i ≤ s), b[i] (for i = 1, 2, …, s) and c[i] (for i = 2, 3, …, s). The matrix [aij] is called the Runge–Kutta matrix, while the b[i] and c[i] are known as the weights and the nodes. These data are usually arranged in a mnemonic device, known as a Butcher tableau (after John C. Butcher):

Examples

>>> euler = _Tableau(c=[0.0],
...                  b=[1.0],
...                  a=[[0.0]]
... )
>>> euler.c
[0.0]

c[source]¶

The nodes

Type:: List[float]

b[source]¶

The weights

Type:: List[float]

a[source]¶

The Runge-Kutta matrix

Type:: List[List[float]]

c: List[float][source]¶: Alias for field number 0

b: List[float][source]¶: Alias for field number 1

a: List[List[float]][source]¶: Alias for field number 2

__getnewargs__()[source]¶: Return self as a plain tuple. Used by copy and pickle.

static __new__(_cls, c: List[float], b: List[float], a: List[List[float]])[source]¶: Create new instance of _Tableau(c, b, a)

explicit_rk(tableau: _Tableau, fcn: Callable[[...], Tensor], y0: Tensor, t: Tensor, params: Tensor, batch_size: int = 1, device='cpu')[source]¶

The family of explicit Runge–Kutta methods is a generalization of the RK4 method mentioned above.

Examples

>>> def lotka_volterra(t, y, params):
...     X, Y = y
...     alpha, beta, delta, gamma = params
...     dx_dt = alpha * X - beta * X * Y
...     dy_dt = delta * X * Y - gamma * Y
...     return torch.stack([dx_dt, dy_dt])
>>> t = torch.linspace(0, 50, 100)
>>> y_init = torch.rand(2, 1)
>>> solver_param = [rk4_tableau,
...                 lotka_volterra,
...                 y_init,
...                 t,
...                 torch.tensor([1.1, 0.4, 0.1, 0.4])]
>>> solution = explicit_rk(*solver_param)
>>> print(solution.shape)
torch.Size([100, 2, 1])

For solving multiple ODEs, we can feed in multiple initial conditions and parameters >>> def lotka_volterra(t, y, params): … X, Y = y … alpha, beta, delta, gamma = params … dx_dt = alpha * X - beta * X * Y … dy_dt = delta * X * Y - gamma * Y … return torch.stack([dx_dt, dy_dt]) >>> t = torch.linspace(0, 50, 100) >>> batch_size = 10 >>> y_init = torch.rand(2, batch_size) >>> params = torch.rand(4, batch_size) >>> solver_param = [rk4_tableau, … lotka_volterra, … y_init, … t, … params] >>> solution = explicit_rk(*solver_param, batch_size=batch_size, device=’cpu’) >>> print(solution.shape) torch.Size([100, 2, 10])

Parameters:

fcn (callable dy/dt = fcn(t, y, *params)) – The function to be integrated. It should produce output of list of tensors following the shapes of tuple y. t should be a single element.
t (torch.Tensor (nt,)) – The time values to integrate over. Can be generated by linspace (Refer example above).
y0 (list of torch.Tensor (*ny)) – The list of initial values.
params (list) – List of input parameters for the function fcn.
batch_size (int) – The batch size to compute the RK method. Default is 1.
device (str) – The device to compute the RK method. Default is “cpu”.

Returns:

yt – The value of y at the given time t for each batch.

Return type:

list of torch.Tensor (nt, *ny, batch_size)

Notes

We require the function fcn to have t as the first argument, y as the second argument and the rest of the arguments as params. The function should return a tensor of the same shape as y. Even though t may not be used in the function, it is required to have it as the first argument. For multiple derivatives, pass the y as a list of tensors. If multiple parameters are required, they should be passed as a list. Refer example above.

References

[1].. https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods#Explicit_Runge.E2.80.93Kutta_methods

rk38_ivp(fcn: Callable[[...], Tensor], y0: Tensor, t: Tensor, params: Tensor, batch_size: int = 1, device='cpu', **kwargs)[source]¶

A slight variation of “the” Runge–Kutta method is also due to Kutta in 1901 and is called the 3/8-rule.[19] The primary advantage this method has is that almost all of the error coefficients are smaller than in the popular method, but it requires slightly more FLOPs (floating-point operations) per time step.

Examples

>>> def lotka_volterra(t, y, params):
...     X, Y = y
...     alpha, beta, delta, gamma = params
...     dx_dt = alpha * X - beta * X * Y
...     dy_dt = delta * X * Y - gamma * Y
...     return torch.stack([dx_dt, dy_dt])
>>> t = torch.linspace(0, 50, 100)
>>> y_init = torch.rand(2, 1)
>>> solver_param = [lotka_volterra,
...                 y_init,
...                 t,
...                 torch.tensor([1.1, 0.4, 0.1, 0.4])]
>>> solution = rk38_ivp(*solver_param)
>>> print(solution.shape)
torch.Size([100, 2, 1])

Parameters:

fcn (callable dy/dt = fcn(t, y, *params)) – The function to be integrated. It should produce output of list of tensors following the shapes of tuple y. t should be a single element.
t (torch.Tensor (nt,)) – The integrated times
y0 (list of torch.Tensor (*ny)) – The list of initial values
params (list) – List of any other parameters
batch_size (int) – The batch size to compute the RK method. Default is 1.
device (str) – The device to compute the RK method. Default is “cpu”.
**kwargs (dict) – Any other keyword arguments

Returns:

yt – The value of y at the given time t

Return type:

list of torch.Tensor (nt,*ny)

fwd_euler_ivp(fcn: Callable[[...], Tensor], y0: Tensor, t: Tensor, params: Tensor, batch_size: int = 1, device='cpu', **kwargs)[source]¶

However, the simplest Runge–Kutta method is the (forward) Euler method, given by the formula $y_{n+1} = y_{n} + hf(t_{n}, y_{n}). This is the only consistent explicit Runge–Kutta method with one stage.

Examples

>>> def lotka_volterra(t, y, params):
...     X, Y = y
...     alpha, beta, delta, gamma = params
...     dx_dt = alpha * X - beta * X * Y
...     dy_dt = delta * X * Y - gamma * Y
...     return torch.stack([dx_dt, dy_dt])
>>> t = torch.linspace(0, 50, 1000)
>>> y_init = torch.randn(2, 1)
>>> solver_param = [lotka_volterra,
...                 y_init,
...                 t,
...                 torch.tensor([1.1, 0.4, 0.1, 0.4])]
>>> solution = fwd_euler_ivp(*solver_param)
>>> print(solution.shape)
torch.Size([1000, 2, 1])

Parameters:

fcn (callable dy/dt = fcn(t, y, *params)) – The function to be integrated. It should produce output of list of tensors following the shapes of tuple y. t should be a single element.
t (torch.Tensor (nt,)) – The integrated times
y0 (list of torch.Tensor (*ny)) – The list of initial values
params (list) – List of any other parameters
batch_size (int) – The batch size to compute the RK method. Default is 1.
device (str) – The device to compute the RK method. Default is “cpu”.
**kwargs (dict) – Any other keyword arguments

Returns:

yt – The value of y at the given time t

Return type:

list of torch.Tensor (nt,*ny)

rk4_ivp(fcn: Callable[[...], Tensor], y0: Tensor, t: Tensor, params: Tensor, batch_size: int = 1, device='cpu', **kwargs)[source]¶

The most commonly used Runge Kutta method to find the solution of a differential equation is the RK4 method, i.e., the fourth-order Runge-Kutta method. The Runge-Kutta method provides the approximate value of y for a given point x. Only the first order ODEs can be solved using the Runge Kutta RK4 method.

Examples

>>> def lotka_volterra(t, y, params):
...     X, Y = y
...     alpha, beta, delta, gamma = params
...     dx_dt = alpha * X - beta * X * Y
...     dy_dt = delta * X * Y - gamma * Y
...     return torch.stack([dx_dt, dy_dt])
>>> t = torch.linspace(0, 50, 100)
>>> y_init = torch.rand(2, 1)
>>> solver_param = [lotka_volterra,
...                 y_init,
...                 t,
...                 torch.tensor([1.1, 0.4, 0.1, 0.4])]
>>> solution = rk4_ivp(*solver_param)
>>> print(solution.shape)
torch.Size([100, 2, 1])

Parameters:

fcn (callable dy/dt = fcn(t, y, *params)) – The function to be integrated. It should produce output of list of tensors following the shapes of tuple y. t should be a single element.
t (torch.Tensor (nt,)) – The integrated times
y0 (list of torch.Tensor (*ny)) – The list of initial values
params (list) – List of any other parameters
batch_size (int) – The batch size to compute the RK method. Default is 1.
device (str) – The device to compute the RK method. Default is “cpu”.
**kwargs (dict) – Any other keyword arguments

Returns:

yt – The value of y at the given time t

Return type:

list of torch.Tensor (nt,*ny)

mid_point_ivp(fcn: Callable[[...], Tensor], y0: Tensor, t: Tensor, params: Tensor, batch_size: int = 1, device='cpu', **kwargs)[source]¶

The explicit midpoint method is sometimes also known as the modified Euler method, the implicit method is the most simple collocation method, and, applied to Hamiltonian dynamics, a symplectic integrator.

Examples

>>> def lotka_volterra(t, y, params):
...     X, Y = y
...     alpha, beta, delta, gamma = params
...     dx_dt = alpha * X - beta * X * Y
...     dy_dt = delta * X * Y - gamma * Y
...     return torch.stack([dx_dt, dy_dt])
>>> t = torch.linspace(0, 50, 100)
>>> y_init = torch.rand(2, 1)
>>> solver_param = [lotka_volterra,
...                 y_init,
...                 t,
...                 torch.tensor([1.1, 0.4, 0.1, 0.4])]
>>> solution = rk4_ivp(*solver_param)
>>> print(solution.shape)
torch.Size([100, 2, 1])

Parameters:

fcn (callable dy/dt = fcn(t, y, *params)) – The function to be integrated. It should produce output of list of tensors following the shapes of tuple y. t should be a single element.
t (torch.Tensor (nt,)) – The integrated times
y0 (list of torch.Tensor (*ny)) – The list of initial values
params (list) – List of any other parameters
batch_size (int) – The batch size to compute the RK method. Default is 1.
device (str) – The device to compute the RK method. Default is “cpu”.
**kwargs (dict) – Any other keyword arguments

Returns:

yt – The value of y at the given time t

Return type:

list of torch.Tensor (nt,*ny)

Attribute Utilities¶

The utilities here are used to modify the attributes of the classes. Used by differentiation_utils.

class get_attr(obj: object, name: str)[source]¶

Get the attribute of an object.

Examples

>>> from deepchem.utils.attribute_utils import get_attr
>>> class MyClass:
...     def __init__(self):
...         self.a = 1
...         self.b = 2
>>> obj = MyClass()
>>> get_attr(obj, "a")
1

Parameters:

obj (object) – The object to get the attribute from.
name (str) – The name of the attribute.

Returns:

val – The value of the attribute.

Return type:

object

class set_attr(obj: object, name: str, val: object)[source]¶

Set the attribute of an object.

Examples

>>> from deepchem.utils import set_attr
>>> class MyClass:
...     def __init__(self):
...         self.a = 1
...         self.b = 2
>>> obj = MyClass()
>>> set_attr(obj, "a", 3)
>>> set_attr(obj, "c", 4)
>>> obj.a
3
>>> obj.c
4

Parameters:

obj (object) – The object to set the attribute to.
name (str) – The name of the attribute.
val (object) – The value to set the attribute to.

class del_attr(obj: Any, name: str)[source]¶

Delete the attribute of an object.

Examples

>>> from deepchem.utils import del_attr
>>> class MyClass:
...     def __init__(self):
...         self.a = 1
...         self.b = 2
>>> obj = MyClass()
>>> del_attr(obj, "a")
>>> try:
...     obj.a
... except AttributeError:
...     print("AttributeError")
AttributeError

Polymer Weighted Directed Graph Data Utilities¶

These classes and functions are required to handle converstion of string data to graph data and validation of the same.

handle_hydrogen(smiles: str, keep_h: bool = False, add_h: bool = False) → Mol[source]¶

Builds an RDKit molecule from a SMILES string by conditionally handling addition or removal of hydrogens.

This function is useful to process organic molecules properly.

Examples

>>> import deepchem as dc
>>> mol = dc.utils.handle_hydrogen('C', keep_h=True, add_h=False)
>>> mol.GetNumAtoms()
1
>>> mol = dc.utils.handle_hydrogen('C', keep_h=True, add_h=True)
>>> mol.GetNumAtoms()
5

Parameters:

smiles (str) – SMILES string.
keep_h (bool, optional, default=False) – Whether to keep hydrogens in the molecule.
add_h (bool, optional, default=False) – Whether to add hydrogens to the molecule.

Returns:

RDKit molecule object.

Return type:

rdkit.Chem.rdchem.Mol

make_polymer_mol(smiles: str, fragment_weights: List[float], keep_h: bool, add_h: bool) → Mol[source]¶

Builds an RDKit joined molecule from a SMILES string of monomer molecules. The weight of each monomer is stored as a metadata property of each atom.

Examples

>>> import deepchem as dc
>>> mol = dc.utils.make_polymer_mol('C.C', [1, 1], True, False)
>>> for atom in mol.GetAtoms():
...     print(atom.GetDoubleProp('w_frag'))
1.0
1.0

Parameters:

smiles (str) – SMILES string of monomer molecules joined by “.”.
fragment_weights (list) – List of weights for each monomer.
keep_h (bool) – Whether to keep hydrogens in the molecule.
add_h (bool) – Whether to add hydrogens to the molecule.

Returns:

RDKit polymer molecule object.

Return type:

rdkit.Chem.rdchem.Mol

parse_polymer_rules(rules: List[str]) → Tuple[List[tuple], float][source]¶

This function extracts probabilty weight distribution details for bonds from string to list of tuples in following format. (start, end, weight_forward, weight_reverse) The start and end is in string. (As it will be mapped with open bond string) The weight_forward and weight_reverse in float It also returns the degree of polymerization of the polymer (DoP) as 1. + np.log10(DoP). The DoP can be mentioned at the end of the rules separated by “~”. If DoP is not given, the function assumes DoP = 1.

Examples

>>> import deepchem as dc
>>> polymer_info, degree_of_polymerization = dc.utils.parse_polymer_rules(
...     ['1-2:0.5:0.5'])
>>> polymer_info
[('1', '2', 0.5, 0.5)]
>>> degree_of_polymerization
1.0

Parameters:

rules (list[str]) – List of strings containing bond rules in the format “start-end:weight_forward:weight_reverse”
Returns –
-------- –
polymer_info (list[tuple]) – List of tuples containing bond rules in the format (start, end, weight_forward, weight_reverse)
degree_of_polymerization (float) – Degree of polymerization of the polymer.

tag_atoms_in_repeating_unit(mol: RWMol) → Tuple[RWMol, dict][source]¶

This function tags atoms that are part of the core units, as well as atoms serving to identify attachment points. In addition, create a map of bond types based on what bonds are connected to R groups in the input. The input molecules must be of Chem.rdchem.RWMol type to be editable.

Examples

>>> import deepchem as dc
>>> from rdkit import Chem
>>> rw_mol = Chem.rdchem.RWMol(Chem.MolFromSmiles('[1*]CC.C[2*]'))
>>> mol, _ = dc.utils.tag_atoms_in_repeating_unit(rw_mol)
>>> mol.GetAtomWithIdx(0).GetBoolProp('core')
False
>>> mol.GetAtomWithIdx(1).GetBoolProp('core')
True
>>> mol.GetAtomWithIdx(3).GetProp('R')
'2*'

Parameters:

mol (rdkit.Chem.rdchem.RWMol) – RDKit read and write enabled molecule object.

Returns:

rdkit.Chem.rdchem.RWMol – RDKit read and write enabled molecule object.
dict – Map of R group to bond type.

onek_encoding_unk(value: int, choices: list) → list[source]¶

This function generates the vector for a value as one hot encoded list. If there is an unknown value, it will be encoded as 1 in the last index

Examples

>>> import deepchem as dc
>>> dc.utils.onek_encoding_unk(1, [1, 2, 3])
[1, 0, 0, 0]
>>> dc.utils.onek_encoding_unk(69, [1, 2, 3])
[0, 0, 0, 1]

Parameters:

value (int) – Value to be encoded.
choices (list) – List of choices.

Returns:

One hot encoded vector.

Return type:

list

remove_wildcard_atoms(rwmol: RWMol) → RWMol[source]¶

This function removes the connection virtual atoms for open bonds in a molecule. This is necessary for molecules with wildcard notations.

Examples

>>> import deepchem as dc
>>> from rdkit import Chem
>>> mol = Chem.MolFromSmiles("[*]CC")
>>> rwmol = Chem.RWMol(mol)
>>> rwmol = dc.utils.remove_wildcard_atoms(rwmol)
>>> Chem.MolToSmiles(rwmol)
'CC'

Parameters:: rwmol (rdkit.Chem.rdchem.RWMol) – Read and Writable RDKit molecule object.
Returns:: Read and writable RDKit molecule object.
Return type:: rdkit.Chem.rdchem.RWMol

Polymer Weighted Directed Graph String Validator¶

This class validates the input string representation for graphical conversion of polymer data. It splits the input strings, validates their order and values. Upon recieving error it can return meaningful error message to the user indicating the type of error that hinders the conversion process.

The format is an extended string representation that is used as input format for Weighted Directed Message Passing Neural Network (wDMPNN) implementation using polymer-chemprop python module. For more information and understanding, visit the github repo <https://github.com/coleygroup/polymer-chemprop>. The format development and implementation is done by Matteo Aldeghi and Connor W. Coley for their work on “A graph representation of molecular ensembles for polymer property prediction”.

The dc.utils.PolyWDGStringValidator class is explicitly useful for validating a Weighted Directed Graph Representaion within a string data for polymers. It validates atom notations in monomer, valid Fragment weights for monomers, and valid polymer rules within the string representation.

References:

`Aldeghi M, Coley CW. A graph representation of molecular ensembles for polymer property prediction. Chemical Science. 2022;13(35):10486-98.`_

class PolyWDGStringValidator[source]¶

Class for validating the string format of weighted directed graph data.This class provides methods to validate the format of a datapoint string. This is a specific string format that is used for storing weighted directed polymer data in a parsable format.

The format is as follows:: [monomer1].[monomer2]|[fraction_of_monomer1]|[fraction_of_monomer2]|<[polymer_rule1]<[polymer_rule2]
The polymer rule has an own format in it. Which is as follows:: [[atom_index1]-[atom_index2]]:[fraction_of_bond_between_atom1_to_atom2]:[fraction_of_bond_between_atom2_to_atom1]

This format is explicitly used for formatting the input for Weighted Directed Message Passing Neural Networks (wD-MPNN). The input format holds a SMART notation and regular expression formatting to keep molecular data with corresponding bonds and weights. Irrespective of this explicit usecase, the formatting can allow featurization of same data for other graph based neural networks.

The validate method validates the proper formatting for monomer molecules, proper value of the fractions and valid atom indicies and corresponding weights in the polymer rules.

Example

>>> from deepchem.utils import PolyWDGStringValidator
>>> validator = PolyWDGStringValidator()
>>> validator.validate("[1*]C.C[2*]|0.5|0.5|<1-2:0.5:0.5")
True

References

static get_parsed_vals(datapoint: str) → Tuple[str, list, str][source]¶

This static method parses the datapoint string into 3 parts: 1. Monomer molecules 2. Fragments 3. Polymer rules

Parameters:: datapoint (str) – The datapoint string to parse
Returns:: A tuple containing the 3 parts of the datapoint string
Return type:: Tuple[str, list, str]

static get_polymer_rules(rules_str: str) → List[str][source]¶

This static method parses the polymer rules string into a list of rules.

Parameters:: rules_str (str) – The polymer rules string to parse
Returns:: A list containing the parsed rule strings
Return type:: List[str]

validate(datapoint: str)[source]¶

This method validates the string format of weighted directed graph data. To validate the string format it checks for following conditions:

The number of fragments and the number of monomer molecules should match.
The wild card indexes should be present in the monomer molecules string and should be in the correct sequence.
The polymer rules should be in the correct format.
The atom indexes in the polymer rules should be valid and present in the monomer molecules string.

It raises ValueError if the string format is invalid.

Parameters:: datapoint (str) – The datapoint string to validate
Returns:: True if the string format is valid, None otherwise (Error will be raised otherwise)
Return type:: bool

Polymer Conversion Utilities¶

These classes and functions are required to handle converstion of PSMILES string data format to WDG string data format and vice versa.

class PSMILES2WDGConverter(conversion_types: list = ['alternate'])[source]¶

This class is used to convert a PSMILES string to a to corresponding WDGraph string. The conversion can be done in two ways:

With metadata: The metadata contains the indices of the atoms to be bonded and the residue.

Without metadata: The conversion is done using the conversion types specified in the constructor.

The mechanism is as follows:

The PSMILES string is split into two parts.
The bond is broken between the specified indices.
The fragments are converted to SMILES strings.
The SMILES strings are combined with the residue to form the final WDGraph string.

This specific mechanism is utilized to convert a PSMILES string to a WDGraph string such that the representational variation is tested with neural network architectural differences in our research paper “Open-Source Polymer Generative Pipeline” [1]_.

References

Examples

>>> from rdkit import Chem
>>> from deepchem.utils.poly_converters import PSMILES2WDGConverter
>>> psmiles = "*CCCC*"
>>> metadata = {
...     "indicies": [2, 3],
...     "seq_index": 4,
...     "residue": "0.5|0.5|",
...     "smiles_type": "SMARTS"
... }
>>> converter = PSMILES2WDGConverter()
>>> result = converter([psmiles], [metadata])
>>> print(result)
[['[*:1]CC.[*:2]C[*:3][*:4]C|0.5|0.5|']]

__init__(conversion_types: list = ['alternate']) → None[source]¶

Parameters:: conversion_types (list) – The conversion types to be used for the conversion. The default is [“alternate”].

index_wildcards(psmiles_part: str) → str[source]¶

This function is used to index the wildcard atoms in the PSMILES string.

Parameters:: psmiles_part (str) – The PSMILES string to be indexed.
Returns:: The indexed PSMILES string.
Return type:: str

add_indicies_to_smiles_from_meta(smiles: str, seq_index: int) → str[source]¶

This function is used to add the indices to the SMILES string from the metadata.

Parameters:

smiles (str) – The SMILES string to be modified.
seq_index (int) – The index of the sequence in the SMILES string.

Returns:

The modified SMILES string.

Return type:

str

make_wdgraph_string_from_meta(psmiles: str, indicies: list) → str[source]¶

This function is used to make the WDGraph string from the metadata.

Parameters:

psmiles (str) – The PSMILES string to be converted.
indicies (list) – The indices of the atoms to be bonded.

Returns:

The WDGraph string.

Return type:

str

convert_smiles_to_SMARTS(smiles_string: str) → str[source]¶

This function is used to convert the SMILES string to SMARTS string.

Parameters:: smiles_string (str) – The SMILES string to be converted.
Returns:: The converted SMARTS string.
Return type:: str

convert_smiles_to_MOD_SMARTS(smiles_string: str) → str[source]¶

This function is used to convert the SMILES string to MOD SMARTS string.

Parameters:: smiles_string (str) – The SMILES string to be converted.
Returns:: The converted MOD SMARTS string.
Return type:: str

compose_from_meta(psmiles: str, metadata: dict) → str[source]¶

This function is used to compose the WDGraph string from the metadata.

Parameters:

psmiles (str) – The PSMILES string to be converted.
metadata (dict) – The metadata containing the indices of the atoms to be bonded, the residue, and the type of SMILES string.

Returns:

The composed WDGraph string.

Return type:

str

convert(psmiles_string: str, metadata: dict | None = None) → list[source]¶

This function is used to convert the PSMILES string to a WDGraph string.

Parameters:

psmiles_string (str) – The PSMILES string to be converted.
metadata (dict) – The metadata containing the indices of the atoms to be bonded, the residue, and the type of SMILES string.

Returns:

The converted WDGraph string.

Return type:

list

class WDG2PSMILESConverter(return_metadata: bool = True)[source]¶

This class is used to convert a WDGraph string to a corresponding PSMILES string. The mechanism returns two formats:

With metadata: Gives metadata containing the indices of the atoms to be bonded and the residue.

Without metadata: Gives only the converted PSMILES string. (the exact reverse conversion of PSMILES to WDG will be done considering alternating structure only)

The mechanism is as follows:

The WDGraph string is split into two parts (PSMILES part, Residue part)
The PSMILES part is converted from SMART notation to normal notation.
The PSMILES part is converted to a single molecule keeping the bond breaking indicies as metadata.
The final PSMILES is returned with the metadata of the bond breaking indices and initial residue.

This specific mechanism is utilized to convert a PSMILES string to a WDGraph string such that the representational variation is tested with neural network architectural differences in our research paper “Open-Source Polymer Generative Pipeline” [1]_.

References

Examples

>>> from rdkit import Chem
>>> from deepchem.utils.poly_converters import WDG2PSMILESConverter
>>> wd_graph_string = "[*:1]CC[*:2].[*:3]CC[*:4]|0.5|0.5|<1-3:0.5:0.5<1-4:0.5:0.5<2-3:0.5:0.5<2-4:0.5:0.5"
>>> converter = WDG2PSMILESConverter(return_metadata = True)
>>> result, meta = converter([wd_graph_string])
>>> print(result)
['*CCCC*']
>>> print(meta)
[{'indicies': [3, 4], 'seq_index': 3, 'residue': '0.5|0.5|<1-3:0.5:0.5<1-4:0.5:0.5<2-3:0.5:0.5<2-4:0.5:0.5', 'smiles_type': 'SMARTS'}]

__init__(return_metadata: bool = True) → None[source]¶

Parameters:: return_metadata (bool) – Whether to return metadata or not. The default is True.

get_wildcard_bond_indecies(combined_psmiles: str) → Tuple[int, int][source]¶

This function is used to get the wildcard bond indices in the combined PSMILES string.

Parameters:: combined_psmiles (str) – The combined PSMILES string.
Returns:: The wildcard bond indices.
Return type:: Tuple(int, int)

convert_smiles_part(smiles: str) → Tuple[str, int, int, int][source]¶

This function is used to convert the SMILES part of the WDG string to a PSMILES string and returns the indicies along side.

Parameters:: smiles (str) – The SMILES part of the WDG string.
Returns:: The converted PSMILES string, the bond breaking indicies, and the sequence index.
Return type:: Tuple(str, int, int, int)

replace_SMARTS(smiles: str) → str[source]¶

This function is used to replace the SMARTS notation in the SMILES string.

Parameters:: smiles (str) – The SMILES string to be modified.
Returns:: The modified SMILES string
Return type:: str

replace_mod_SMARTS(mod_smiles: str) → str[source]¶

This function is used to replace the MOD SMARTS notation in the SMILES string.

Parameters:: mod_smiles (str) – The SMILES string to be modified.
Returns:: The modified SMILES string.
Return type:: str

convert(wdg_graph_string: str) → Tuple[str, dict][source]¶

This is used to convert a single WDGraph string to a PSMILES string and form the metadata.

Parameters:: wdg_graph_string (str) – The WDGraph string to be converted.
Returns:: The converted PSMILES string and the metadata.
Return type:: Tuple(str, dict)

Pytorch Utilities¶

unsorted_segment_sum(data: Tensor, segment_ids: Tensor, num_segments: int) → Tensor[source]¶

Computes the sum along segments of a tensor. Analogous to tf.unsorted_segment_sum.

Parameters:

data (torch.Tensor) – A tensor whose segments are to be summed.
segment_ids (torch.Tensor) – The segment indices tensor.
num_segments (int) – The number of segments.

Returns:

tensor

Return type:

torch.Tensor

Examples

>>> segment_ids = torch.Tensor([0, 1, 0]).to(torch.int64)
>>> data = torch.Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [4, 3, 2, 1]])
>>> num_segments = 2
>>> result = unsorted_segment_sum(data=data,
...                               segment_ids=segment_ids,
...                               num_segments=num_segments)
>>> data.shape[0]
3
>>> segment_ids.shape[0]
3
>>> len(segment_ids.shape)
1
>>> result
tensor([[5., 5., 5., 5.],
        [5., 6., 7., 8.]])

segment_sum(data: Tensor, segment_ids: Tensor) → Tensor[source]¶

This function computes the sum of values along segments within a tensor. It is useful when you have a tensor with segment IDs and you want to compute the sum of values for each segment. This function is analogous to tf.segment_sum. (https://www.tensorflow.org/api_docs/python/tf/math/segment_sum).

Parameters:

data (torch.Tensor) – A pytorch tensor containing the values to be summed. It can have any shape, but its rank (number of dimensions) should be at least 1.
segment_ids (torch.Tensor) – A 1-D tensor containing the indices for the segmentation. The segments can be any non-negative integer values, but they must be sorted in non-decreasing order.

Returns:

out_tensor – Tensor with the same shape as data, where each value corresponds to the sum of values within the corresponding segment.

Return type:

torch.Tensor

Examples

>>> data = torch.Tensor([[1, 2, 3, 4], [4, 3, 2, 1], [5, 6, 7, 8]])
>>> segment_ids = torch.Tensor([0, 0, 1]).to(torch.int64)
>>> result = segment_sum(data=data, segment_ids=segment_ids)
>>> data.shape[0]
3
>>> segment_ids.shape[0]
3
>>> len(segment_ids.shape)
1
>>> result
tensor([[5., 5., 5., 5.],
        [5., 6., 7., 8.]])

chunkify(a: Tensor, dim: int, maxnumel: int) → Generator[Tuple[Tensor, int, int], None, None][source]¶

Splits the tensor a into several chunks of size maxnumel along the dimension given by dim.

Examples

>>> import torch
>>> from deepchem.utils.pytorch_utils import chunkify
>>> a = torch.arange(10)
>>> for chunk, istart, iend in chunkify(a, 0, 3):
...     print(chunk, istart, iend)
tensor([0, 1, 2]) 0 3
tensor([3, 4, 5]) 3 6
tensor([6, 7, 8]) 6 9
tensor([9]) 9 12

Parameters:

a (torch.Tensor) – The big tensor to be splitted into chunks.
dim (int) – The dimension where the tensor would be splitted.
maxnumel (int) – Maximum number of elements in a chunk.

Returns:

chunks – A generator that yields a tuple of three elements: the chunk tensor, the starting index of the chunk and the ending index of the chunk.

Return type:

Generator[Tuple[torch.Tensor, int, int], None, None]

get_memory(a: Tensor) → int[source]¶

Returns the size of the tensor in bytes.

Examples

>>> import torch
>>> from deepchem.utils.pytorch_utils import get_memory
>>> a = torch.randn(100, 100, dtype=torch.float64)
>>> get_memory(a)
80000

Parameters:: a (torch.Tensor) – The tensor to be measured.
Returns:: size – The size of the tensor in bytes.
Return type:: int

gaussian_integral(n: int, alpha: float | Tensor) → float | Tensor[source]¶

Performs the gaussian integration.

Examples

>>> gaussian_integral(5, 1.0)
1.0

Parameters:

n (int) – The order of the integral
alpha (Union[float, torch.Tensor]) – The parameter of the gaussian

Returns:

The value of the integral

Return type:

Union[float, torch.Tensor]

TensorNonTensorSeparator(params: Sequence, varonly: bool = True)[source]¶

Class that provides function to separate/combine tensors and nontensors parameters.

Examples

>>> import torch
>>> from deepchem.utils.pytorch_utils import TensorNonTensorSeparator
>>> a = torch.tensor([1.,2,3])
>>> b = 4.
>>> c = torch.tensor([5.,6,7], requires_grad=True)
>>> params = [a, b, c]
>>> separator = TensorNonTensorSeparator(params)
>>> tensor_params = separator.get_tensor_params()
>>> tensor_params
[tensor([5., 6., 7.], requires_grad=True)]

tallqr(V, MV=None)[source]¶

QR decomposition for tall and skinny matrix.

Examples

>>> import torch
>>> from deepchem.utils.pytorch_utils import tallqr
>>> V = torch.randn(3, 2)
>>> Q, R = tallqr(V)
>>> Q.shape
torch.Size([3, 2])
>>> R.shape
torch.Size([2, 2])
>>> torch.allclose(Q @ R, V)
True

Parameters:

V (torch.Tensor) – V is a matrix to be decomposed. (*BV, na, nguess)
MV (torch.Tensor) – (*BM, na, nguess) where M is the basis to make Q M-orthogonal if MV is None, then MV=V (default=None)

Returns:

Q (torch.Tensor) – The Orthogonal Part. Shape: (*BV, na, nguess)
R (torch.Tensor) – The (*BM, nguess, nguess) where M is the basis to make Q M-orthogonal

to_fortran_order(V)[source]¶

Convert a tensor to Fortran order. (The last two dimensions are made Fortran order.) Fortran order/ array is a special case in which all elements of an array are stored in column-major order.

Examples

>>> import torch
>>> from deepchem.utils.pytorch_utils import to_fortran_order
>>> V = torch.randn(3, 2)
>>> V.is_contiguous()
True
>>> V = to_fortran_order(V)
>>> V.is_contiguous()
False
>>> V.shape
torch.Size([3, 2])
>>> V = torch.randn(3, 2).transpose(-2, -1)
>>> V.is_contiguous()
False
>>> V = to_fortran_order(V)
>>> V.is_contiguous()
False
>>> V.shape
torch.Size([2, 3])

Parameters:: V (torch.Tensor) – V is a matrix to be converted. (*BV, na, nguess)
Returns:: outV – (*BV, nguess, na)
Return type:: torch.Tensor

get_np_dtype(dtype: dtype) → Any[source]¶

corresponding numpy dtype from the input pytorch’s tensor dtype

Examples

>>> import torch
>>> from deepchem.utils.pytorch_utils import get_np_dtype
>>> get_np_dtype(torch.float32)
<class 'numpy.float32'>
>>> get_np_dtype(torch.float64)
<class 'numpy.float64'>

Parameters:: dtype (torch.dtype) – pytorch’s tensor dtype
Returns:: corresponding numpy dtype
Return type:: np.dtype

unsorted_segment_max(data: Tensor, segment_ids: Tensor, num_segments: int) → Tensor[source]¶

Computes the maximum along segments of a tensor. Analogous to tf.unsorted_segment_max.

Parameters:

data (torch.Tensor) – A tensor whose segments are to be maximized.
segment_ids (torch.Tensor) – The segment indices tensor.
num_segments (int) – The number of segments.

Returns:

tensor

Return type:

torch.Tensor

Examples

>>> segment_ids = torch.Tensor([0, 1, 0]).to(torch.int64)
>>> data = torch.Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [4, 3, 2, 1]])
>>> num_segments = 2
>>> result = unsorted_segment_max(data=data,
...                               segment_ids=segment_ids,
...                               num_segments=num_segments)
>>> data.shape[0]
3
>>> segment_ids.shape[0]
3
>>> len(segment_ids.shape)
1
>>> result
tensor([[4., 3., 3., 4.],
        [5., 6., 7., 8.]])

unweighted_coul_ft(gvgrids: FloatTensor) → Tensor[source]¶

Unweighted fourier transform of the coulomb kernel: 4*pi/|gv|^2

Examples

>>> from deepchem.utils import unweighted_coul_ft
>>> unweighted_coul_ft(torch.tensor([[1., 2], [3, 4]]))
tensor([2.5133, 0.5027])

Parameters:: gvgrids (torch.FloatTensor) – Reciprocal space vector. Quantities like electron density and hartree potential are evaluated on this. (ngv, ndim)
Returns:: Unweighted fourier transform of gvgrids. (ngv,)
Return type:: torch.FloatTensor

get_complex_dtype(dtype: dtype) → dtype[source]¶

Corresponding complex type given the real floating point datatype.

Examples

>>> import torch
>>> from deepchem.utils import get_complex_dtype
>>> a = torch.tensor([1.2, 23.4])
>>> a_new = a.to(get_complex_dtype(a.dtype))
>>> a_new
tensor([ 1.2000+0.j, 23.4000+0.j])

Batch Utilities¶

The utilites here are used for computing features on batch of data. Can be used inside of default_generator function.

batch_coulomb_matrix_features(X_b: ndarray, distance_max: float = -1, distance_min: float = 18, n_distance: int = 100)[source]¶

Computes the values for different Feature on given batch. It works as a helper function to coulomb matrix.

This function takes in a batch of Molecules represented as Coulomb Matrix.

It proceeds as follows:

It calculates the Number of atoms per molecule by counting all the non zero elements(numbers) of every molecule layer in matrix in one dimension.
The Gaussian distance is calculated using the Euclidean distance between the Cartesian coordinates of two atoms. The distance value is then passed through a Gaussian function, which transforms it into a continuous value.
Then using number of atom per molecule, calculates the atomic charge by looping over the molecule layer in the Coulomb matrix and takes the 2.4 root of the diagonal of 2X of each molecule layer. Undoing the Equation of coulomb matrix.
Atom_membership is assigned as a commomn repeating integers for all the atoms for a specific molecule.
Distance Membership encodes spatial information, assigning closer values to atoms that are in that specific molecule. All initial Distances are added a start value to them which are unique to each molecule.

Models Used in:

DTNN

Parameters:

X_b (np.ndarray) – It is a 3d Matrix containing information of each the atom’s ionic interaction with other atoms in the molecule.
distance_min (float (default -1)) – minimum distance of atom pairs (in Angstrom)
distance_max (float (default = 18)) – maximum distance of atom pairs (in Angstrom)
n_distance (int (default 100)) – granularity of distance matrix step size will be (distance_max-distance_min)/n_distance

Returns:

atom_number (np.ndarray) – Atom numbers are assigned to each atom based on their atomic properties. The atomic numbers are derived from the periodic table of elements. For example, hydrogen -> 1, carbon -> 6, and oxygen -> 8.
gaussian_dist (np.ndarray) – Gaussian distance refers to the method of representing the pairwise distances between atoms in a molecule using Gaussian functions. The Gaussian distance is calculated using the Euclidean distance between the Cartesian coordinates of two atoms. The distance value is then passed through a Gaussian function, which transforms it into a continuous value.
atom_mem (np.ndarray) – Atom membership refers to the binary representation of whether an atom belongs to a specific group or property within a molecule. It allows the model to incorporate domain-specific information and enhance its understanding of the molecule’s properties and interactions.
dist_mem_i (np.ndarray) – Distance membership i are utilized to encode spatial information and capture the influence of atom distances on the properties and interactions within a molecule. The inner membership function assigns higher values to atoms that are closer to the atoms’ interaction region, thereby emphasizing the impact of nearby atoms.
dist_mem_j (np.ndarray) – It captures the long-range effects and influences between atoms that are not in direct proximity but still contribute to the overall molecular properties. Distance membership j are utilized to encode spatial information and capture the influence of atom distances on the properties and interactions outside a molecule. The outer membership function assigns higher values to atoms that are farther to the atoms’ interaction region, thereby emphasizing the impact of farther atoms.

Examples

>>> import os
>>> import deepchem as dc
>>> current_dir = os.path.dirname(os.path.abspath(__file__))
>>> dataset_file = os.path.join(current_dir, 'test/assets/qm9_mini.sdf')
>>> TASKS = ["alpha", "homo"]
>>> loader = dc.data.SDFLoader(tasks=TASKS,
...                            featurizer=dc.feat.CoulombMatrix(29),
...                            sanitize=True)
>>> data = loader.create_dataset(dataset_file, shard_size=100)
>>> inputs = dc.utils.batch_utils.batch_coulomb_matrix_features(data.X)

References

batch_elements(elements: List[Any], batch_size: int)[source]¶

Combine elements into batches.

Parameters:

elements (List[Any]) – List of Elements to be combined into batches.
batch_size (int) – Batch size in which to divide.

Returns:

batch – List of Lists of elements divided into batches.

Return type:

List[Any]

Examples

>>> import deepchem as dc
>>> # Prepare Data
>>> inputs = [[i, i**2, i**3] for i in range(10)]
>>> # Run
>>> output = list(dc.utils.batch_utils.batch_elements(inputs, 3))
>>> len(output)
4

create_input_array(sequences: Collection, max_input_length: int, reverse_input: bool, batch_size: int, input_dict: Dict, end_mark: Any)[source]¶

Create the array describing the input sequences.

It creates a 2d Matrix empty matrix according to batch size and max_length. Then iteratively fills it with the key-values from the input dictionary.

Many NLP Models like SeqToSeq has sentences as there inputs. We need to convert these sentences into numbers so that the model can do computation on them.

This function takes in the sentence then using the input_dict dictionary picks up the words/letters equivalent numerical represntation. Then makes an numpy array of it.

If the reverse_input is True, then the order of the input sequences is reversed before sending them into the encoder. This can improve performance when working with long sequences.

These values can be used to generate embeddings for further processing.

Models used in:

SeqToSeq

Parameters:

sequences (Collection) – List of sequences to be converted into input array.
reverse_input (bool) – If True, reverse the order of input sequences before sending them into the encoder. This can improve performance when working with long sequences.
batch_size (int) – Batch size of the input array.
input_dict (dict) – Dictionary containing the key-value pairs of input sequences.
end_mark (Any) – End mark for the input sequences.

Returns:

features – Numeric Representation of the given sequence according to input_dict.

Return type:

np.Array

Examples

>>> import deepchem as dc
>>> # Prepare Data
>>> inputs = [["a", "b"], ["b", "b", "b"]]
>>> input_dict = {"c": 0, "a": 1, "b": 2}
>>> # Inputs property
>>> max_length = max([len(x) for x in inputs])
>>> # Without reverse input
>>> output_1 = dc.utils.batch_utils.create_input_array(inputs, max_length,
...                                                    False, 2, input_dict,
...                                                    "c")
>>> output_1.shape
(2, 4)
>>> # With revercse input
>>> output_2 = dc.utils.batch_utils.create_input_array(inputs, max_length,
...                                                    True, 2, input_dict,
...                                                    "c")
>>> output_2.shape
(2, 4)

create_output_array(sequences: Collection, max_output_length: int, batch_size: int, output_dict: Dict, end_mark: Any)[source]¶

Create the array describing the target sequences.

It creates a 2d Matrix empty matrix according to batch size and max_length. Then iteratively fills it with the key-values from the output dictionary.

This function is similar to create_input_array function. The only difference is that it is used for output sequences and does not have the reverse_input parameter as it is not required for output sequences.

It is used in NLP Models like SeqToSeq where the output is also a sentence and we need to convert it into numbers so that the model can do computation on them. This function takes in the sentence then using the output_dict dictionary picks up the words/letters equivalent numerical represntation. Then makes an numpy array of it.

These values can be used to generate embeddings for further processing.

Models used in:

SeqToSeq

Parameters:

sequences (Collection) – List of sequences to be converted into output array.
max_output_length (bool) – Maximum length of output sequence that may be generated
batch_size (int) – Batch size of the output array.
output_dict (dict) – Dictionary containing the key-value pairs of output sequences.
end_mark (Any) – End mark for the output sequences.

Returns:

features – Numeric Representation of the given sequence according to output_dict.

Return type:

np.Array

Examples

>>> import deepchem as dc
>>> # Prepare Data
>>> inputs = [["a", "b"], ["b", "b", "b"]]
>>> output_dict = {"c": 0, "a": 1, "b": 2}
>>> # Inputs property
>>> max_length = max([len(x) for x in inputs])
>>> output = dc.utils.batch_utils.create_output_array(inputs, max_length, 2,
...                                                   output_dict, "c")
>>> output.shape
(2, 3)

Periodic Table Utilities¶

The Utilities here are used to computing atomic mass and radii data. These can be used by DFT and many other Molecular Models.

Returns the atomic number for the given element

Examples

>>> from deepchem.utils import get_atomz
>>> element_symbol = "Al" # Aluminium
>>> get_atomz(element_symbol)
13
>>> get_atomz(17)
17

Parameters:: element (Union[str, ZType]) – String symbol of Element or Atomic Number. Ex: H, He, C
Returns:: atom_n – Atomic Number of the given Element.
Return type:: ZType

References

Kasim, Muhammad F., and Sam M. Vinko. “Learning the exchange-correlation functional from nature with fully differentiable density functional theory.” Physical Review Letters 127.12 (2021): 126403. https://github.com/diffqc/dqc/blob/master/dqc/utils/periodictable.py

Equivariance Utilities¶

The utilities here refer to equivariance tools that play a vital role in mathematics and applied sciences. They excel in preserving the relationships between objects or data points when undergoing transformations such as rotations or scaling.

You can refer to the tutorials for additional information regarding equivariance and Deepchem’s support for equivariance.

su2_generators(k: int) → Tensor[source]¶

Generate the generators of the special unitary group SU(2) in a given representation.

The function computes the generators of the SU(2) group for a specific representation determined by the value of ‘k’. These generators are commonly used in the study of quantum mechanics, angular momentum, and related areas of physics and mathematics. The generators are represented as matrices.

The SU(2) group is a fundamental concept in quantum mechanics and symmetry theory. The generators of the group, denoted as J_x, J_y, and J_z, represent the three components of angular momentum operators. These generators play a key role in describing the transformation properties of physical systems under rotations.

The returned tensor contains three matrices corresponding to the x, y, and z generators, usually denoted as J_x, J_y, and J_z. These matrices form a basis for the Lie algebra of the SU(2) group.

In linear algebra, specifically within the context of quantum mechanics, lowering and raising operators are fundamental concepts that play a crucial role in altering the eigenvalues of certain operators while acting on quantum states. These operators are often referred to collectively as “ladder operators.”

A lowering operator is an operator that, when applied to a quantum state, reduces the eigenvalue associated with a particular observable. In the context of SU(2), the lowering operator corresponds to J_-.

Conversely, a raising operator is an operator that increases the eigenvalue of an observable when applied to a quantum state. In the context of SU(2), the raising operator corresponds to J_+.

The z-generator matrix represents the component of angular momentum along the z-axis, often denoted as J_z. It commutes with both J_x and J_y and is responsible for quantizing the angular momentum.

Note that the dimensions of the returned tensor will be (3, 2j+1, 2j+1), where each matrix has a size of (2j+1) x (2j+1). :param k: The representation index, which determines the order of the representation. :type k: int

Returns:: A stack of three SU(2) generators, corresponding to J_x, J_z, and J_y.
Return type:: torch.Tensor

Notes

A generating set of a group is a subset $S$ of the group $G$ such that every element of $G$ can be expressed as a combination (under the group operation) of finitely many elements of the subset $S$ and their inverses.

The special unitary group $SU_n(q)$ is the set of $n*n$ unitary matrices with determinant +1. $SU(2)$ is homeomorphic with the orthogonal group $O_3^+(2)$. It is also called the unitary unimodular group and is a Lie group.

References

Examples

>>> su2_generators(1)
tensor([[[ 0.0000+0.0000j,  0.7071+0.0000j,  0.0000+0.0000j],
         [-0.7071+0.0000j,  0.0000+0.0000j,  0.7071+0.0000j],
         [ 0.0000+0.0000j, -0.7071+0.0000j,  0.0000+0.0000j]],

        [[-0.0000-1.0000j,  0.0000+0.0000j,  0.0000+0.0000j],
         [ 0.0000+0.0000j,  0.0000+0.0000j,  0.0000+0.0000j],
         [ 0.0000+0.0000j,  0.0000+0.0000j,  0.0000+1.0000j]],

        [[ 0.0000-0.0000j,  0.0000+0.7071j,  0.0000-0.0000j],
         [ 0.0000+0.7071j,  0.0000-0.0000j,  0.0000+0.7071j],
         [ 0.0000-0.0000j,  0.0000+0.7071j,  0.0000-0.0000j]]])

so3_generators(k: int) → Tensor[source]¶

Construct the generators of the SO(3) Lie algebra for a given quantum angular momentum.

The function generates the generators of the special orthogonal group SO(3), which represents the group of rotations in three-dimensional space. Its Lie algebra, which consists of the generators of infinitesimal rotations, is often used in physics to describe angular momentum operators. The generators of the Lie algebra can be related to the SU(2) group, and this function uses a transformation to convert the SU(2) generators to the SO(3) basis.

The primary significance of the SO(3) group lies in its representation of three-dimensional rotations. Each matrix in SO(3) corresponds to a unique rotation, capturing the intricate ways in which objects can be oriented in 3D space. This concept finds application in numerous fields, ranging from physics to engineering.

Parameters:: k (int) – The representation index, which determines the order of the representation.
Returns:: A stack of three SO(3) generators, corresponding to J_x, J_z, and J_y.
Return type:: torch.Tensor

Notes

The special orthogonal group $SO_n(q)$ is the subgroup of the elements of general orthogonal group $GO_n(q)$ with determinant 1. $SO_3$ (often written $SO(3)$) is the rotation group for three-dimensional space.

These matrices are orthogonal, which means their rows and columns form mutually perpendicular unit vectors. This preservation of angles and lengths makes orthogonal matrices fundamental in various mathematical and practical applications.

The “special” part of $SO(3)$ refers to the determinant of these matrices being $+1$. The determinant is a scalar value that indicates how much a matrix scales volumes. A determinant of $+1$ ensures that the matrix represents a rotation in three-dimensional space without involving any reflection or scaling operations that would reverse the orientation of space.

References

Examples

>>> so3_generators(1)
tensor([[[ 0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000, -1.0000],
         [ 0.0000,  1.0000,  0.0000]],

        [[ 0.0000,  0.0000,  1.0000],
         [ 0.0000,  0.0000,  0.0000],
         [-1.0000,  0.0000,  0.0000]],

        [[ 0.0000, -1.0000,  0.0000],
         [ 1.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000]]])

change_basis_real_to_complex(k: int, dtype: dtype | None = None, device: device | None = None) → Tensor[source]¶

Construct a transformation matrix to change the basis from real to complex spherical harmonics.

This function constructs a transformation matrix Q that converts real spherical harmonics into complex spherical harmonics. It operates on the basis functions $Y_{ell m}$ and $Y_{ell}^{m}$, and accounts for the relationship between the real and complex forms of these harmonics as defined in the provided mathematical expressions.

The resulting transformation matrix Q is used to change the basis of vectors or tensors of real spherical harmonics to their complex counterparts.

Parameters:

k (int) – The representation index, which determines the order of the representation.
dtype (torch.dtype, optional) – The data type for the output tensor. If not provided, the function will infer it. Default is None.
device (torch.device, optional) – The device where the output tensor will be placed. If not provided, the function will use the default device. Default is None.

Returns:

A transformation matrix Q that changes the basis from real to complex spherical harmonics.

Return type:

torch.Tensor

Notes

Spherical harmonics Y_l^m are a family of functions that are defined on the surface of a unit sphere. They are used to represent various physical and mathematical phenomena that exhibit spherical symmetry. The indices l and m represent the degree and order of the spherical harmonics, respectively.

The conversion from real to complex spherical harmonics is achieved by applying specific transformation coefficients to the real-valued harmonics. These coefficients are derived from the properties of spherical harmonics.

References

Examples

# The transformation matrix generated is used to change the basis of a vector of # real spherical harmonics with representation index 1 to complex spherical harmonics. >>> change_basis_real_to_complex(1) tensor([[-0.7071+0.0000j, 0.0000+0.0000j, 0.0000-0.7071j],

[ 0.0000+0.0000j, 0.0000-1.0000j, 0.0000+0.0000j], [-0.7071+0.0000j, 0.0000+0.0000j, 0.0000+0.7071j]])

wigner_D(k: int, alpha: Tensor, beta: Tensor, gamma: Tensor) → Tensor[source]¶

Wigner D matrix representation of the SO(3) rotation group.

The function computes the Wigner D matrix representation of the SO(3) rotation group for a given representation index ‘k’ and rotation angles ‘alpha’, ‘beta’, and ‘gamma’. The resulting matrix satisfies properties of the SO(3) group representation.

Parameters:

k (int) – The representation index, which determines the order of the representation.
alpha (torch.Tensor) – Rotation angles (in radians) around the Y axis, applied third.
beta (torch.Tensor) – Rotation angles (in radians) around the X axis, applied second.
gamma (torch.Tensor)) – Rotation angles (in radians) around the Y axis, applied first.

Returns:

The Wigner D matrix of shape (#angles, 2k+1, 2k+1).

Return type:

torch.Tensor

Notes

The Wigner D-matrix is a unitary matrix in an irreducible representation of the groups SU(2) and SO(3).

The Wigner D-matrix is used in quantum mechanics to describe the action of rotations on states of particles with angular momentum. It is a key concept in the representation theory of the rotation group SO(3), and it plays a crucial role in various physical contexts.

Examples

>>> k = 1
>>> alpha = torch.tensor([0.1, 0.2])
>>> beta = torch.tensor([0.3, 0.4])
>>> gamma = torch.tensor([0.5, 0.6])
>>> wigner_D_matrix = wigner_D(k, alpha, beta, gamma)
>>> wigner_D_matrix
tensor([[[ 0.8275,  0.1417,  0.5433],
         [ 0.0295,  0.9553, -0.2940],
         [-0.5607,  0.2593,  0.7863]],

        [[ 0.7056,  0.2199,  0.6737],
         [ 0.0774,  0.9211, -0.3817],
         [-0.7044,  0.3214,  0.6329]]])

semifactorial(x: int) → float[source]¶

Compute the semifactorial function x!!, defined as: x!! = x * (x-2) * (x-4) * … :param x: A positive integer. :type x: int

Returns:: The value of x!!, computed iteratively.
Return type:: float

Examples

>>> semifactorial(5)
15.0
>>> semifactorial(6)
48.0

pochhammer(x: int, k: int) → float[source]¶

Compute the Pochhammer symbol (x)_k , defined as: (x)_k = x * (x+1) * (x+2) * … * (x+k-1). :param x: The starting integer of the sequence. :type x: int :param k: The number of terms in the product. :type k: int

Returns:: The Pochhammer symbol value.
Return type:: float

Examples

>>> pochhammer(3, 4)
360.0
>>> pochhammer(5, 2)
30.0

lpmv(d: int, m: int, x: Tensor) → Tensor[source]¶

Compute the associated Legendre function P_l^m(x), including the Condon-Shortley phase. The implementation includes the base case for P_m^m(x), P_{m+1}^m(x), and uses a recurrence relation for higher degrees. :param l: Degree of the Legendre polynomial. :type l: int :param m: Order of the Legendre polynomial. :type m: int :param x: A tensor of shape (N,) representing the input values for x,

where ( -1 <= x <= 1).

Returns:: A tensor of shape (N,) containing the values of P_l^m(x).
Return type:: torch.Tensor

Examples

>>> lpmv(2, 1, torch.tensor([0.5]))
tensor([-1.2990])

SphericalHarmonics()[source]¶

A class for computing real tesseral spherical harmonics, including the Condon-Shortley phase. This implementation is inspired by the SE(3)-Transformer library by Fabian Fuchs, which provides efficient computation of spherical harmonics and related transformations. For more details, see the SE(3)-Transformer repository: https://github.com/FabianFuchsML/se3-transformer-public

get_element(l, m, theta, phi)[source]¶: Compute the tesseral spherical harmonic Y_l^m(theta, phi).

get(l, theta, phi)[source]¶: Compute all spherical harmonics of degree l for given angles.

get_matrix_kernel(A: Tensor, eps: float = 1e-10) → Tensor[source]¶

Compute an orthonormal basis of the kernel (x_1, x_2, …).

This function calculates the null space (kernel) of the given matrix A, such that:: A x_i = 0 scalar_product(x_i, x_j) = delta_ij

Parameters:

A (torch.Tensor) – The input matrix.
eps (float, optional) – Tolerance for singular values considered as zero (default is 1e-10).

Returns:

A matrix where each row is a basis vector of the kernel of A.

Return type:

torch.Tensor

Examples

>>> from deepchem.utils.equivariance_utils import get_matrix_kernel
>>> A = torch.tensor([[1.0, 2.0, 3.0],
...                   [2.0, 4.0, 6.0],
...                   [3.0, 6.0, 9.0]])
>>> get_matrix_kernel(A)
tensor([[ 0.0000, -0.8321,  0.5547],
        [ 0.9636, -0.1482, -0.2224]])

basis_transformation_Q_J(J: int, order_in: int, order_out: int) → Tensor[source]¶

Compute one part of the Q^-1 matrix for the article.

This function computes the spherical harmonics projection matrix for the Sylvester equation in the subspace J needed in for the weight basis in SE(3)-Transformer model.

References:¶

SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks

Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, Max Welling NeurIPS 2020, https://arxiv.org/abs/2006.10503

param J:: Order of the spherical harmonics.
type J:: int
param order_in:: Order of the input representation.
type order_in:: int
param order_out:: Order of the output representation.
type order_out:: int
param version:: Version of the computation (default is 3).
type version:: int, optional
returns:: A tensor of shape [(m_out * m_in), m], where m = 2 * J + 1.
rtype:: torch.Tensor

Examples

>>> from deepchem.utils.equivariance_utils import basis_transformation_Q_J
>>> basis_transformation_Q_J(1, 1, 1).shape
torch.Size([9, 3])

get_spherical_from_cartesian(cartesian: Tensor, divide_radius_by: float = 1.0) → Tensor[source]¶

Convert Cartesian coordinates to spherical coordinates.

Parameters:

cartesian (torch.Tensor) – Cartesian coordinates tensor of shape […, 3].
divide_radius_by (float, optional) – Factor by which to divide the radius (default is 1.0).

Returns:

Spherical coordinates tensor of shape […, 3] with [radius, azimuth (phi), elevation (theta)].

Return type:

torch.Tensor

Examples

>>> from deepchem.utils.equivariance_utils import get_spherical_from_cartesian
>>> cartesian = torch.tensor([[1.0, 1.0, 1.0]])
>>> get_spherical_from_cartesian(cartesian)
tensor([[1.7321, 0.7854, 0.9553]])

>>> cartesian = torch.tensor([[0.0, 0.0, 1.0]])  # Point on Z-axis
>>> get_spherical_from_cartesian(cartesian)
tensor([[1.0000, 0.0000, 1.5708]])

>>> cartesian = torch.tensor([[0.0, 0.0, -1.0]])  # Point on negative Z-axis
>>> get_spherical_from_cartesian(cartesian)
tensor([[1.0000, 3.1416, 1.5708]])

kron(a: Tensor, b: Tensor) → Tensor[source]¶

Computes the Kronecker product of two tensors needed to comput Q_J matrix.

References:¶

SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks

Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, Max Welling NeurIPS 2020, https://arxiv.org/abs/2006.10503

param a:: First input tensor of shape [*, m, n].
type a:: torch.Tensor
param b:: Second input tensor of shape [*, p, q].
type b:: torch.Tensor
returns:: The Kronecker product of a and b.
rtype:: torch.Tensor

Examples

>>> from deepchem.utils.equivariance_utils import kron
>>> A = torch.tensor([[1, 2], [3, 4]])
>>> B = torch.tensor([[0, 5], [6, 7]])
>>> kron(A, B)
tensor([[ 0,  5,  0, 10],
        [ 6,  7, 12, 14],
        [ 0, 15,  0, 20],
        [18, 21, 24, 28]])

precompute_sh(r_ij: Tensor, max_J: int) → dict[source]¶

Precompute spherical harmonics up to a given order used in the forward pass of SE(3)-Transformer model.

References:¶

SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks

Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, Max Welling NeurIPS 2020, https://arxiv.org/abs/2006.10503

param r_ij:: Relative positions tensor.
type r_ij:: torch.Tensor
param max_J:: Maximum order of the spherical harmonics.
type max_J:: int
returns:: A dictionary where each key corresponds to an order J and the value is a tensor of shape [B, N, K, 2J+1].
rtype:: dict

Examples

>>> from deepchem.utils.equivariance_utils import precompute_sh
>>> r_ij = torch.tensor([[1.0, 0.5, 1.0]])  # Example spherical coordinates (radius, phi, theta)
>>> precompute_sh(r_ij, max_J=2)
{0: tensor([[0.2821]]), 1: tensor([[-0.1971, -0.2640, -0.3608]]), 2: tensor([[ 0.3255,  0.2381, -0.0392,  0.4359,  0.2090]])}

get_basis(G, max_degree: int, compute_gradients: bool = False) → Dict[str, Tensor][source]¶

Precompute the SE(3)-equivariant weight basis, ( W_J^{lk}(x) ).

This function computes the equivariant weight basis used in SE(3)-equivariant convolutions, enabling the model to learn **rotation-equivariant filters.

It is called internally by get_basis_and_r(), which computes both: - Equivariant weight basis ( W_J^{lk}(x) ). - Inter-nodal distances ( r_{ij} ) (used in attention & convolution).

— ### Mathematical Background This function follows Equation (8) from [SE(3)-Transformer paper](https://arxiv.org/pdf/2006.10503.pdf):

[ W_J^{lk}(x) = Q_J cdot Y_J(x) ]

( Y_J(x) ): Spherical Harmonic Basis
( Q_J ): Basis transformation matrix
( J ): Angular momentum index
( d_{in}, d_{out} ): Feature degrees

This function precomputes basis functions for SE(3)-equivariant convolutions by calcutating spherical harmonic projections.

Parameters:

G (dgl.DGLGraph) – DGL graph where G.edata[‘edge_attr’] stores edge displacement vectors.
max_degree (int) – Maximum degree (l_max) of equivariant tensors.
compute_gradients (bool, optional, default=`False`) – If True: Enables gradient tracking for backpropagation. If False: Uses torch.no_grad() for efficiency.

Returns:

basis – Dictionary mapping ‘<d_in>,<d_out>’: precomputed basis tensor. Shape: (batch_size, 1, 2*d_out+1, 1, 2*d_in+1, num_bases)

Return type:

Dict[str, torch.Tensor]

Example

>>> import torch
>>> import dgl
>>> import shutil
>>> import os
>>> from deepchem.utils.equivariance_utils import get_basis
>>> from rdkit import Chem
>>> import deepchem as dc
>>> mol = Chem.MolFromSmiles('CCO')
>>> featurizer = dc.feat.EquivariantGraphFeaturizer(fully_connected=False, embeded=True)
>>> features = featurizer.featurize([mol])[0]
>>> G = features.to_dgl_graph()
>>> G.edata['w'] = torch.tensor(features.edge_weights, dtype=torch.float32)
>>> basis = get_basis(G, max_degree=2)
>>> print(basis.keys())
dict_keys(['0,0', '0,1', '0,2', '1,0', '1,1', '1,2', '2,0', '2,1', '2,2'])
>>> dir_path = "cache"
>>> if os.path.exists(dir_path) and os.path.isdir(dir_path):
...     shutil.rmtree(dir_path)

get_r(G) → Tensor[source]¶

Compute inter-nodal distances for a given DGL graph.

This function computes the Euclidean distance between connected nodes based on edge feature d, which represents relative displacements.

Parameters:: G (dgl.DGLGraph) – The input graph where G.edata[‘edge_attr’] contains edge displacement vectors.
Returns:: A tensor containing the computed inter-nodal distances for each edge, with shape (num_edges, 1).
Return type:: torch.Tensor

Example

>>> import torch
>>> import dgl
>>> from deepchem.utils.equivariance_utils import get_r
>>> from rdkit import Chem
>>> import deepchem as dc
>>> mol = Chem.MolFromSmiles('CCO')
>>> featurizer = dc.feat.EquivariantGraphFeaturizer(fully_connected=False, embeded=True)
>>> features = featurizer.featurize([mol])[0]
>>> G = features.to_dgl_graph()
>>> G.edata['w'] = torch.tensor(features.edge_weights, dtype=torch.float32)
>>> # Compute internodal distances
>>> r = get_r(G)
>>> print(r.shape)  # (num_edges, 1)
torch.Size([4, 1])

get_equivariant_basis_and_r(G, max_degree: int, compute_gradients: bool = False) → Tuple[Dict[str, Tensor], Tensor][source]¶

Compute equivariant weight basis and inter-nodal distances for SE(3)-Transformers.

This function computes: 1. Equivariant weight basis: Required for SE(3)-equivariant convolutions. 2. Inter-nodal distances: Used in attention mechanisms and basis function computation.

This should be called once per forward pass to compute shared equivariant representations.

Parameters:

G (dgl.DGLGraph) – A DGL graph containing edge features.
max_degree (int) – The maximum degree of the SE(3) tensor representation.
compute_gradients (bool, optional (default=False)) – If True, enables gradient computation for the basis.

Returns:

basis (dict[str, torch.Tensor]): Dictionary of equivariant bases, indexed by ‘<d_in><d_out>’, where d_in and d_out are feature degrees.
r (torch.Tensor): Inter-nodal distance tensor with shape (num_edges, 1).

Return type:

Tuple[Dict[str, torch.Tensor], torch.Tensor]

Example

>>> import torch
>>> import dgl
>>> from deepchem.utils.equivariance_utils import get_equivariant_basis_and_r
>>> from rdkit import Chem
>>> import deepchem as dc
>>> mol = Chem.MolFromSmiles('CCO')
>>> featurizer = dc.feat.EquivariantGraphFeaturizer(fully_connected=False, embeded=True)
>>> features = featurizer.featurize([mol])[0]
>>> G = features.to_dgl_graph()
>>> G.edata['w'] = torch.tensor(features.edge_weights, dtype=torch.float32)
>>> # Compute basis and distances
>>> basis, r = get_equivariant_basis_and_r(G, max_degree=2)
>>> print(r.shape)  # Expected: (num_edges, 1)
torch.Size([4, 1])
>>> print(basis.keys())  # Expected: dict
dict_keys(['0,0', '0,1', '0,2', '1,0', '1,1', '1,2', '2,0', '2,1', '2,2'])

fiber2head(F: Dict[str, Tensor], h: int, structure: Fiber, squeeze: bool = False) → Tensor[source]¶

Converts SE(3)-equivariant features into multi-head format for attention. This function reshapes and concatenates fiber-based features into a format suitable for multi-head attention.

Input Fiber Representation: $$ F[d] in mathbb{R}^{B imes N imes M_d imes (2d+1)} $$
After Multi-Head Transformation: $$ F[d] in mathbb{R}^{B imes N imes H imes (M_d / H) imes (2d+1)} $$

where: - B: Batch size. - N: Number of nodes. - M_d: Multiplicity (number of channels for degree d). - 2d+1: Number of components in degree d. - H: Number of attention heads. - M_d / H: Channels per head.

Parameters:

(Dict[str (- F) – Dictionary of SE(3)-equivariant features, where keys are **degrees (d).
torch.Tensor]) – Dictionary of SE(3)-equivariant features, where keys are **degrees (d).
(int) (- h) – Number of attention heads.
(Fiber) (- structure) – Fiber representation defining feature multiplicities and degrees.
(bool (- squeeze) –
- If True: Concatenates along last feature dimension.
- If False: Concatenates along second-last dimension.
optional –
- If True: Concatenates along last feature dimension.
- If False: Concatenates along second-last dimension.
default=`False`) –
- If True: Concatenates along last feature dimension.
- If False: Concatenates along second-last dimension.

Returns:

Reshaped tensor with multi-head representation.
Shape: (batch, num_nodes, num_heads, channels_per_head, 2d+1)

Return type:

torch.Tensor

Example

>>> import torch
>>> from deepchem.models.torch_models.layers import Fiber
>>> from deepchem.utils.equivariance_utils import fiber2head

# Create Fiber Structure >>> fiber_structure = Fiber(dictionary={0: 16, 1: 32}) # Scalars & Vectors

# Create Input Feature Dictionary >>> F = { … ‘0’: torch.randn(10, 16, 1), # Degree 0 features (scalars) … ‘1’: torch.randn(10, 32, 3) # Degree 1 features (vectors) … }

# Define Number of Attention Heads >>> num_heads = 4 >>> output = fiber2head(F, num_heads, fiber_structure, squeeze=True) >>> print(output.shape) torch.Size([10, 4, 28])

References

class LieGroup(alpha: float = 0.2)[source]¶

Abstract base class for Lie groups used in LieConv [1] (https://github.com/mfinzi/LieConv).

This class represents continuous Lie groups acting on Euclidean data, as required by LieConv-style equivariant neural networks [1]. Specific groups, like translations (T) or 3d roto-translations (SE(3)), or subclasses must implement the group exponential and logarithm maps, as well as a lifting procedure from Euclidean coordinates to group-aligned representations.

Each group is characterized by: - A representation space on which the group acts. - A Lie algebra parameterization. - An optional quotient (orbit) embedding.

Subclasses must implement at least: - exp() - the group exponential map - log() - the group logarithm map - lifted_elems() - lifting from Euclidean coordinates to Lie algebra elements

and define the attributes rep_dim, lie_dim, and q_dim.

rep_dim[source]¶

Dimension of the representation space on which the group acts (e.g., 2 for SO(2) acting on R^2).

Type:: int

lie_dim[source]¶

Dimension of the Lie algebra parameterization (e.g., 1 for SO(2)).

Type:: int

q_dim[source]¶

Dimension of the quotient space embedding X/G, if applicable (e.g., 1 for SO(2) acting on R^2).

Type:: int

References

__init__(alpha: float = 0.2)[source]¶

Initialize the Lie group.

Parameters:: alpha (float, optional (default 0.2)) – Parameter controlling the contribution of Lie group displacement versus orbit (quotient-space) displacement in the distance metric. Larger values emphasize group motion (orientation/rotation), while smaller values emphasize orbit or spatial displacement.

exp(a: Tensor) → Tensor[source]¶

Exponential map from the Lie algebra to the Lie group.

Computes the group element $expleft(sum_i a_i A_iright)$, where ${A_i}$ are the generators of the Lie algebra and $a in mathbb{R}^{text{lie_dim}}$ are the corresponding coefficients.

Parameters:: a (torch.Tensor) – Lie algebra coefficients taking values in $mathbb{R}^{text{lie_dim}}$.
Returns:: Group elements represented as matrices in $mathbb{R}^{text{rep_dim} times text{rep_dim}}$.
Return type:: torch.Tensor

log(u: Tensor) → Tensor[source]¶

Logarithm map from the Lie group to the Lie algebra.

Computes the Lie algebra element corresponding to a group element, expressed in the chosen Lie algebra basis.

Parameters:: u (torch.Tensor) – Group elements represented as matrices in $mathbb{R}^{text{rep_dim} times text{rep_dim}}$.
Returns:: Lie algebra coefficients taking values in $mathbb{R}^{text{lie_dim}}$.
Return type:: torch.Tensor

lifted_elems(points: Tensor, n_samples: int, **kwargs) → Tuple[Tensor, Tensor | None][source]¶

Lift Euclidean coordinates into Lie algebra and quotient representations.

This method maps input points in Euclidean space to elements of the Lie algebra, together with optional quotient (orbit) embeddings. For groups with multivalued lifts, multiple samples may be generated per input point.

Parameters:

points (torch.Tensor) – Input points taking values in $mathbb{R}^{text{rep_dim}}$.
n_samples (int) – Number of lifted samples per point.

Returns:

a (torch.Tensor) – Lie algebra elements taking values in $mathbb{R}^{text{lie_dim}}$.
q (torch.Tensor or None) – Quotient (orbit) embeddings taking values in $mathbb{R}^{text{q_dim}}$, or None if the quotient is trivial.

inv(g: Tensor) → Tensor[source]¶

Compute the inverse of a group element.

Uses the identity $g^{-1} = exp(-log(g))$.

Parameters:: g (torch.Tensor) – Group elements represented as matrices in $mathbb{R}^{text{rep_dim} times text{rep_dim}}$.
Returns:: Inverse group elements in $mathbb{R}^{text{rep_dim} times text{rep_dim}}$.
Return type:: torch.Tensor

distance(abq_pairs: Tensor) → Tensor[source]¶

Compute the combined group and orbit distance.

The distance between two lifted points is defined as $$ alpha , lVert Delta a rVert + (1 - alpha) , lVert q_1 - q_2 rVert, $$ where $Delta a = log(v^{-1} u)$ is the relative Lie algebra element and $q_1, q_2$ are the corresponding quotient embeddings.

Parameters:: abq_pairs (torch.Tensor) – Concatenated relative Lie algebra and quotient components taking values in $mathbb{R}^{text{lie_dim} + 2,text{q_dim}}$.
Returns:: Scalar distance values.
Return type:: ArrayLike

lift(x: Tuple[Tensor, Tensor, Tensor], nsamples: int, **kwargs) → Tuple[Tensor, Tensor, Tensor][source]¶

Lift inputs to pairwise Lie group representations.

This method lifts Euclidean points into the Lie algebra (and optional quotient space), expands associated features and masks, and constructs all pairwise relative group elements via the group logarithm.

Parameters:

x (Tuple) – Tuple (p, v, m) consisting of: - p: input points taking values in $mathbb{R}^{text{rep_dim}}$ - v: feature values associated with each point - m: binary mask indicating valid points
nsamples (int) – Number of lifted samples per point.

Returns:

embedded_locations (torch.Tensor) – Pairwise relative group elements, optionally augmented with quotient embeddings, taking values in $mathbb{R}^{text{lie_dim} + 2,text{q_dim}}$.
expanded_v (torch.Tensor) – Feature values expanded to align with pairwise group elements.
expanded_mask (torch.Tensor) – Mask expanded to align with pairwise group elements.

expand_like(v: Tensor, m: Tensor, a: Tensor) → Tuple[Tensor, Tensor][source]¶

Expand features and masks to align with lifted Lie algebra elements.

This method replicates per-point feature values and masks so that they correspond to the lifted Lie algebra elements produced by the group lifting procedure.

Parameters:

v (torch.Tensor) – Feature values associated with input points.
m (torch.Tensor) – Mask indicating valid input points.
a (torch.Tensor) – Lifted Lie algebra elements used to determine the expansion.

Returns:

expanded_v (torch.Tensor) – Feature values expanded to align with lifted Lie algebra elements.
expanded_mask (torch.Tensor) – Mask expanded to align with lifted Lie algebra elements.

elems2pairs(a: Tensor) → Tensor[source]¶

Compute pairwise relative Lie algebra elements.

For lifted elements $a_i, a_j in mathbb{R}^{ ext{lie_dim}}$, this method computes the relative displacement $log!left(exp(-a_j)exp(a_i)

ight)$

for all ordered pairs $(i, j)$.

atorch.Tensor: Lie algebra elements taking values in $mathbb{R}^{ ext{lie_dim}}$.

pairwise_lie_algebratorch.Tensor: Pairwise relative Lie algebra elements.

class T(k: int, alpha: float = 0.2)[source]¶

k-dimensional translation group R^k adapted from [1] (https://github.com/mfinzi/LieConv).

This Lie group (T) represents pure translations acting on R^k.

Note that for the translation group the exponential and logarithm maps are identities. Consequently, relative group elements reduce to simple subtraction. In particular, in two dimensions this corresponds to the trivial difference x_i − x_j for $x in mathbb{R}^2$.

The quotient space is trivial (q_dim = 0).

Examples

>>> from deepchem.utils.equivariance_utils import T
>>> G = T(k=2)
>>> x = torch.tensor([[0.0, 0.0],
...                   [1.0, 2.0]])
>>> a, q = G.lifted_elems(x, n_samples=1)
>>> a
tensor([[0., 0.],
        [1., 2.]])
>>> pairs = G.elems2pairs(a)
>>> pairs
tensor([[[ 0.,  0.],
         [-1., -2.]],

        [[ 1.,  2.],
         [ 0.,  0.]]])
>>> ab = pairs.reshape(-1, 2)
>>> G.distance(ab)
tensor([0.0000, 0.4472, 0.4472, 0.0000])

References

__init__(k: int, alpha: float = 0.2) → None[source]¶

Initialize the k-dimensional translation group.

Parameters:

k (int) – Dimension of the translation space R^k.
alpha (float, optional) – Weighting factor for group displacement in the distance metric.

exp(a: Tensor) → Tensor[source]¶

Exponential map (identity for translations).

Parameters:: a (torch.Tensor) – Lie algebra coefficients taking values in $mathbb{R}^{ ext{lie_dim}}$.
Returns:: Group elements of shape $mathbb{R}^{ ext{k}}$.
Return type:: torch.Tensor

log(g: Tensor) → Tensor[source]¶

Logarithm map from the Lie group to the Lie algebra.

For the translation group, this map is the identity.

Parameters:: g (torch.Tensor) – Group elements taking values in $mathbb{R}^k$.
Returns:: Lie algebra elements taking values in $mathbb{R}^k$.
Return type:: torch.Tensor

lifted_elems(points: Tensor, n_samples: int, **kwargs) → Tuple[Tensor, Tensor | None][source]¶

Lift Euclidean points into the Lie algebra.

For the translation group, the lift is unique since the group is Abelian and simply connected; therefore nsamples must be 1.

Parameters:

points (torch.Tensor) – Input points taking values in $mathbb{R}^k$.
n_samples (int) – Number of lifted samples per point (must be 1).

Returns:

a (torch.Tensor) – Lie algebra elements taking values in $mathbb{R}^k$.
q (None) – No quotient embedding for the translation group.

elems2pairs(a: Tensor) → Tensor[source]¶

Compute pairwise Lie algebra differences.

For translations:: log(exp(-b) exp(a)) = a - b

Parameters:: a (torch.Tensor) – Lie algebra elements of shape (*, n, k).
Returns:: Pairwise differences of shape (*, n, n, k).
Return type:: torch.Tensor

sinc(x: Tensor, tresh: float = 0.07) → Tensor[source]¶

Compute the normalized sinc function $\mathrm{sinc}(x) = \sin(x) / x$. This code was adapted from [1] (https://github.com/mfinzi/LieConv).

This implementation is numerically stable near $x = 0$ by using a Taylor series expansion for small values of $|x|$.

The Taylor approximation used is:

$\mathrm{sinc}(x) \approx 1 - \frac{x^2}{6} \left(1 - \frac{x^2}{20} \left(1 - \frac{x^2}{42}\right)\right)$

Parameters:

x (torch.Tensor) – Input tensor.
tresh (float, optional (default 7e-02)) – Threshold below which the Taylor expansion is used.

Returns:

Elementwise values of $\sin(x) / x$.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from deepchem.utils.equivariance_utils import sinc
>>> x = torch.tensor([0.0, 1e-3, 1.0])
>>> sinc(x)
tensor([1.0000, 1.0000, 0.8415])

References

sincc(x: Tensor, tresh: float = 0.07) → Tensor[source]¶

Compute $(1 - \mathrm{sinc}(x)) / x^2$.

This function arises frequently in Lie group exponential map expansions, particularly for rotation groups such as $\mathrm{SO}(3)$. This code was adapted from [1] (https://github.com/mfinzi/LieConv).

Near $x = 0$, the Taylor approximation is:

$\frac{1 - \mathrm{sinc}(x)}{x^2} \approx \frac{1}{6} \left(1 - \frac{x^2}{20} \left(1 - \frac{x^2}{42} \left(1 - \frac{x^2}{72}\right)\right)\right)$

Parameters:

x (torch.Tensor) – Input tensor.
tresh (float, optional (default 7e-02)) – Threshold below which the Taylor expansion is used.

Returns:

Elementwise values of $(1 - \mathrm{sinc}(x)) / x^2$.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from deepchem.utils.equivariance_utils import sincc
>>> x = torch.tensor([0.0, 1e-3, 1.0])
>>> sincc(x)
tensor([0.1667, 0.1667, 0.1585])

References

cosc(x: Tensor, tresh: float = 0.07) → Tensor[source]¶

Compute $(1 - \cos(x)) / x^2$. This code was adapted from [1] (https://github.com/mfinzi/LieConv).

This function appears in second-order expansions of rotation matrices and Lie group Jacobians.

Near $x = 0$, the Taylor approximation is:

$\frac{1 - \cos(x)}{x^2} \approx \frac{1}{2} \left(1 - \frac{x^2}{12} \left(1 - \frac{x^2}{30} \left(1 - \frac{x^2}{56}\right)\right)\right)$

Parameters:

x (torch.Tensor) – Input tensor.
tresh (float, optional (default 7e-02)) – Threshold below which the Taylor expansion is used.

Returns:

Elementwise values of $(1 - \cos(x)) / x^2$.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from deepchem.utils.equivariance_utils import cosc
>>> x = torch.tensor([0.0, 1e-3, 1.0])
>>> cosc(x)
tensor([0.5000, 0.5000, 0.4597])

References

coscc(x: Tensor, tresh: float = 0.07) → Tensor[source]¶

Compute a higher-order cosine correction term used in $\mathrm{SO}(3)$ Jacobians. This code was adapted from [1] (https://github.com/mfinzi/LieConv).

This function corresponds to stabilized expressions involving $(1 - \cos(x))$ and $\sin(x)$ that appear in inverse Jacobians of the exponential map.

Near $x = 0$, the Taylor approximation is:

$\mathrm{coscc}(x) \approx \frac{1}{12} \left(1 + \frac{x^2}{60} \left(1 + \frac{x^2}{42} \left(1 + \frac{x^2}{40}\right)\right)\right)$

Parameters:

x (torch.Tensor) – Input tensor.
tresh (float, optional (default 7e-02)) – Threshold below which the Taylor expansion is used.

Returns:

Elementwise correction values.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from deepchem.utils.equivariance_utils import coscc
>>> x = torch.tensor([0.0, 1e-3, 1.0])
>>> coscc(x)
tensor([0.0833, 0.0833, 0.0848])

Notes

The full expression includes safeguards against numerical instability when $\cos(x) \approx 1$.

References

sinc_inv(x: Tensor, tresh: float = 0.07) → Tensor[source]¶

Compute the inverse sinc function $x / \sin(x)$. This code was adapted from [1] (https://github.com/mfinzi/LieConv).

Near $x = 0$, the Taylor approximation is:

$\frac{x}{\sin(x)} \approx 1 + \frac{x^2}{6} + \frac{7 x^4}{360}$

Parameters:

x (torch.Tensor) – Input tensor.
tresh (float, optional (default 7e-02)) – Threshold below which the Taylor expansion is used.

Returns:

Elementwise values of $x / \sin(x)$.

Return type:

torch.Tensor

Examples

>>> import torch
>>> from deepchem.utils.equivariance_utils import sinc_inv
>>> x = torch.tensor([0.0, 1e-3, 1.0])
>>> sinc_inv(x)
tensor([1.0000, 1.0000, 1.1884])

References

cross_matrix(k: Tensor) → Tensor[source]¶

Construct the skew-symmetric cross-product matrix for vectors in $\mathbb{R}^3$. This code was adapted from [1] (https://github.com/mfinzi/LieConv).

This corresponds to the Hodge star operator mapping $\Lambda^1 \mathbb{R}^3 \rightarrow \Lambda^2 \mathbb{R}^3$, such that for any $v \in \mathbb{R}^3$, $k \times v = [k]_\times v$.

Parameters:: k (torch.Tensor) – Input vectors of shape $(n, 3)$, where n is the dimension of k.
Returns:: Skew-symmetric matrices of shape $(n, 3, 3)$.
Return type:: torch.Tensor

Examples

>>> import torch
>>> from deepchem.utils.equivariance_utils import cross_matrix
>>> k = torch.tensor([1.0, 2.0, 3.0])
>>> cross_matrix(k)
tensor([[ 0., -3.,  2.],
        [ 3.,  0., -1.],
        [-2.,  1.,  0.]])

References

uncross_matrix(K: Tensor) → Tensor[source]¶

Recover a vector in $\mathbb{R}^3$ from a skew-symmetric matrix.

This code was adapted from [1] (https://github.com/mfinzi/LieConv). This is the inverse operation of cross_matrix(), corresponding to the Hodge star mapping $\Lambda^2 \mathbb{R}^3 \rightarrow \Lambda^1 \mathbb{R}^3$.

Parameters:: K (torch.Tensor) – Skew-symmetric matrices of shape $(n, 3, 3)$, where n is the dimension of K.
Returns:: Vectors of shape $(n, 3)$.
Return type:: torch.Tensor

Examples

>>> import torch
>>> from deepchem.utils.equivariance_utils import uncross_matrix
>>> k = torch.tensor([1.0, 2.0, 3.0])
>>> uncross_matrix(cross_matrix(k))
tensor([1., 2., 3.])

References

class SO3(alpha: float = 0.2)[source]¶

Special Orthogonal Group in 3D $\mathrm{SO}(3)$. This code was adapted from [1] (https://github.com/mfinzi/LieConv).

This class implements the Lie group of 3D rotations with exponential and logarithmic maps between the Lie algebra $\mathfrak{so}(3) \cong \mathbb{R}^3$ and rotation matrices in $\mathbb{R}^{3 \times 3}$.

lie_dim[source]¶

Dimension of the Lie algebra, equal to 3.

Type:: int

rep_dim[source]¶

Representation dimension, equal to 3.

Type:: int

q_dim[source]¶

Orbit identifier dimension, equal to 1.

Type:: int

Examples

>>> import torch
>>> from deepchem.utils.equivariance_utils import SO3
>>> G = SO3()
>>> w = torch.tensor([0.0, 0.0, 1.0])
>>> R = G.exp(w)
>>> R.shape
torch.Size([3, 3])
>>> # Recover the Lie algebra element using the logarithmic map:
>>> w_rec = G.log(R)
>>> torch.allclose(w, w_rec, atol=1e-5)
True
>>> # Sample a batch of random rotations:
>>> R = G.sample(8)
>>> R.shape
torch.Size([8, 3, 3])
>>> # Lift 3D points into so3 using stabilizer sampling:
>>> points = torch.randn(4, 3)
>>> a, q = G.lifted_elems(points, n_samples=6)
>>> a.shape
torch.Size([24, 3])
>>> q.shape
torch.Size([24, 1])

References

__init__(alpha: float = 0.2)[source]¶

Initialize the $\mathrm{SO}(3)$ group.

Parameters:: alpha (float, optional (default 0.2)) – Weighting factor for the group displacement component in the distance metric.

exp(w: Tensor) → Tensor[source]¶

Exponential map $\exp : \mathfrak{so}(3) \rightarrow \mathrm{SO}(3)$.

Implements Rodrigues’ rotation formula. The input vectors encode axis–angle rotations, where the direction specifies the rotation axis and the magnitude specifies the rotation angle.

Parameters:: w (torch.Tensor) – Lie algebra elements of shape $(n, 3)$, where n is the dimension of w.
Returns:: Rotation matrices of shape $(n, 3, 3)$.
Return type:: torch.Tensor

Examples

>>> from deepchem.utils.equivariance_utils import SO3
>>> import torch
>>> G = SO3()
>>> w = torch.tensor([0.0, 0.0, 1.0])
>>> R = G.exp(w)
>>> R.shape
torch.Size([3, 3])

log(R: Tensor) → Tensor[source]¶

Logarithmic map $\log : \mathrm{SO}(3) \rightarrow \mathfrak{so}(3)$.

Converts rotation matrices into axis–angle coordinates.

Parameters:: R (torch.Tensor) – Rotation matrices of shape $(n, 3, 3)$, where n is the dimension of R.
Returns:: Lie algebra elements of shape $(n, 3)$.
Return type:: torch.Tensor

Examples

>>> from deepchem.utils.equivariance_utils import SO3
>>> import torch
>>> G = SO3()
>>> w = torch.tensor([0.3, -0.2, 0.1])
>>> R = G.exp(w)
>>> w_rec = G.log(R)
>>> torch.allclose(w, w_rec, atol=1e-5)
True

sample(*shape: int, device: device = device(type='cpu'), dtype: dtype = torch.float32) → Tensor[source]¶

Sample random rotations from $\mathrm{SO}(3)$.

Rotations are sampled by drawing random unit quaternions and converting them to axis–angle form.

Parameters:

*shape (int) – Batch shape.
device (torch.device, optional) – Target device.
dtype (torch.dtype, optional) – Tensor data type.

Returns:

Rotation matrices of shape $(*shape, 3, 3)$.

Return type:

torch.Tensor

Examples

>>> from deepchem.utils.equivariance_utils import SO3
>>> G = SO3()
>>> R = G.sample(8)
>>> R.shape
torch.Size([8, 3, 3])

lifted_elems(points: Tensor, n_samples: int, **kwargs)[source]¶

Lift points in $\mathbb{R}^3$ to $\mathrm{SO}(3)$ elements.

Each point is lifted by aligning a reference axis with the direction of the point and composing with samples from the stabilizer subgroup.

Parameters:

points (torch.Tensor) – Input points of shape $(*, 3)$.
n_samples (int) – Number of lifted group elements per point.

Returns:

a (torch.Tensor) – Lie algebra elements of shape, where n is the dimension of points $(n, n_samples, 3)$.
q (torch.Tensor) – Orbit identifiers of shape $(n, n_samples, 1)$.

Examples

>>> from deepchem.utils.equivariance_utils import SO3
>>> import torch
>>> G = SO3()
>>> points = torch.randn(4, 3)
>>> a, q = G.lifted_elems(points, n_samples=6)
>>> a.shape
torch.Size([24, 3])
>>> q.shape
torch.Size([24, 1])

class SE3(alpha: float = 0.2, per_point: bool = True)[source]¶

Special Euclidean Group in 3D, $\mathrm{SE}(3)$. This code was adapted from [1] (https://github.com/mfinzi/LieConv).

This class represents rigid body transformations in 3D, combining rotations in $\mathrm{SO}(3)$ and translations in $\mathbb{R}^3$.

Elements of the Lie algebra $\mathfrak{se}(3)$ are represented as 6D vectors $(\omega, v)$, where $\omega \in \mathbb{R}^3$ is the rotation vector and $v \in \mathbb{R}^3$ is the translation vector.

The group elements are represented as homogeneous transformation matrices in $\mathbb{R}^{4 \times 4}$.

lie_dim[source]¶

Dimension of the Lie algebra, equal to 6.

Type:: int

rep_dim[source]¶

Representation dimension, equal to 4.

Type:: int

q_dim[source]¶

Quotient dimension, equal to 0 (trivial quotient).

Type:: int

Examples

>>> import torch
>>> from deepchem.utils.equivariance_utils import SE3
>>> G = SE3()
>>> w = torch.tensor([0.0, 0.0, 1.0, 1.0, 0.0, 0.0])
>>> T = G.exp(w)
>>> T.shape
torch.Size([4, 4])
>>> # Recover the Lie algebra element using the logarithmic map:
>>> w_rec = G.log(T)
>>> torch.allclose(w, w_rec, atol=1e-5)
True
>>> # Lift a batch of 3D points into se3 elements:
>>> points = torch.randn(2, 5, 3)
>>> a, q = G.lifted_elems(points, n_samples=4)
>>> a.shape
torch.Size([2, 20, 6])
>>> q is None
True
>>> # Compute a weighted distance between relative transformations:
>>> G = SE3(alpha=0.5)
>>> rel = torch.randn(10, 6)
>>> d = G.distance(rel)
>>> d.shape
torch.Size([10])

References

__init__(alpha: float = 0.2, per_point: bool = True)[source]¶

Initialize the $\mathrm{SE}(3)$ group.

Parameters:

alpha (float, optional (default 0.2)) – Weighting factor between rotational and translational components in the distance metric.
per_point (bool, optional (default True)) – If True, use a distinct random group element per point when lifting inputs.

exp(w: Tensor) → Tensor[source]¶

Exponential map $\exp : \mathfrak{se}(3) \rightarrow \mathrm{SE}(3)$.

Converts Lie algebra elements into homogeneous transformation matrices using the closed-form Rodrigues expansion and the associated left Jacobian.

Parameters:: w (torch.Tensor) – Lie algebra elements of shape $(n, 6)$, where the first three components correspond to rotation and the last three to translation and n is the dimension of w.
Returns:: Homogeneous transformation matrices of shape $(n, 4, 4)$.
Return type:: torch.Tensor

Examples

>>> from deepchem.utils.equivariance_utils import SE3
>>> import torch
>>> G = SE3()
>>> w = torch.tensor([0.0, 0.0, 1.0, 1.0, 0.0, 0.0])
>>> T = G.exp(w)
>>> T.shape
torch.Size([4, 4])

log(U: Tensor) → Tensor[source]¶

Logarithmic map $\log : \mathrm{SE}(3) \rightarrow \mathfrak{se}(3)$.

Converts homogeneous transformation matrices into Lie algebra elements using the inverse left Jacobian.

Parameters:: U (torch.Tensor) – Homogeneous transformation matrices of shape $(n, 4, 4)$, where n is the dimension of U.
Returns:: Lie algebra elements of shape $(n, 6)$.
Return type:: torch.Tensor

Examples

>>> from deepchem.utils.equivariance_utils import SE3
>>> import torch
>>> G = SE3()
>>> w = torch.randn(6)
>>> T = G.exp(w)
>>> w_rec = G.log(T)
>>> torch.allclose(w, w_rec, atol=1e-5)
True

lifted_elems(points: Tensor, n_samples: int, **kwargs) → Tuple[Tensor, Tensor | None][source]¶

Lift Euclidean points into $\mathfrak{se}(3)$ elements.

This method generates random rotations and combines them with translations defined by the input points.

Parameters:

points (torch.Tensor) – Input points of shape $(B, N, 3)$.
n_samples (int) – Number of lifted samples per point.

Returns:

a (torch.Tensor) – Lie algebra elements of shape $(B, N \cdot \text{n_samples}, 6)$.
q (None) – No quotient embedding for $\mathrm{SE}(3)$.

Examples

>>> from deepchem.utils.equivariance_utils import SE3
>>> import torch
>>> G = SE3()
>>> points = torch.randn(2, 5, 3)
>>> a, q = G.lifted_elems(points, n_samples=4)
>>> a.shape
torch.Size([2, 20, 6])
>>> q is None
True

distance(abq_pairs: Tensor) → Tensor[source]¶

Compute a weighted distance in $\mathfrak{se}(3)$.

The distance is defined as: $\alpha |\omega| + (1 - \alpha) |v|$,

where $\omega$ is the rotational component and $v$ is the translational component.

Parameters:: abq_pairs (torch.Tensor) – Relative Lie algebra elements of shape $(n, 6)$, where n is the dimension of abq_pairs.
Returns:: Scalar distance values of shape $(n)$.
Return type:: torch.Tensor

Examples

>>> from deepchem.utils.equivariance_utils import SE3
>>> G = SE3(alpha=0.5)
>>> d = G.distance(torch.randn(10, 6))
>>> d.shape
torch.Size([10])

Miscellaneous Utilities¶

The utilities here are used for miscellaneous purposes. Initial usecases are for improving the printing format of __repr__.

indent(s, nspace)[source]¶

Gives indentation of the second line and next lines. It is used to format the string representation of an object. Which might be containing multiples objects in it. Usage: LinearOperator

Parameters:

s (str) – The string to be indented.
nspace (int) – The number of spaces to be indented.

Returns:

The indented string.

Return type:

str

shape2str(shape)[source]¶

Convert the shape to string representation. It also nicely formats the shape to be readable.

Parameters:: shape (Sequence[int]) – The shape to be converted to string representation.
Returns:: The string representation of the shape.
Return type:: str

memoize_method(fcn: Callable[[Any], T]) → Callable[[Any], T][source]¶

Memoize a method without any arguments using a cache in the object

Examples

>>> class A:
...     @memoize_method
...     def foo(self):
...         print("foo")
...         return 1
>>> a = A()
>>> a.foo()
foo
1
>>> a.foo()
1

Parameters:: fcn (callable) – Function to memoize
Returns:: Memoized function
Return type:: callable

class UnimplementedError[source]¶: Raised if a method is not implemented.

class GetSetParamsError[source]¶: Raised if there is an error in getting or setting parameters.

class ConvergenceWarning[source]¶: Warning to be raised if the convergence of an algorithm is not achieved.

class MathWarning[source]¶: Raised if there are mathematical conditions that are not satisfied.

class Uniquifier(allobjs: List)[source]¶

Identifies and tracks unique objects within a list, even if they are duplicates based on internal memory addresses (using id()). It Optimizes operations involving unique objects by avoiding redundant processing.

Examples

>>> from deepchem.utils import Uniquifier
>>> a = 1
>>> b = 2
>>> c = 3
>>> d = 1
>>> u = Uniquifier([a, b, c, a, d])
>>> u.get_unique_objs()
[1, 2, 3]

__init__(allobjs: List)[source]¶

Initialize the uniquifier.

Parameters:: allobjs (List) – The list of objects to be uniquified.

get_unique_objs(allobjs: List | None = None) → List[source]¶

Get the unique objects.

Parameters:: allobjs (Optional[List]) – The list of objects to be uniquified.
Returns:: The list of unique objects.
Return type:: List

map_unique_objs(uniqueobjs: List) → List[source]¶

Map the unique objects to the original objects.

Parameters:: uniqueobjs (List) – The list of unique objects.

SafeOperations Utilities¶

The utilities here are used for safe operations on tensors. These are used to avoid NaNs and Infs in the output.

safepow(a: Tensor, p: float | Tensor, eps: float = 1e-12) → Tensor[source]¶

Safely calculate the power of a tensor with a small eps to avoid nan.

Examples

>>> import torch
>>> a = torch.tensor([1e-35, 2e-40])
>>> p = torch.tensor([2., 3])
>>> safepow(a, p)
tensor([1.0000e-24, 1.0000e-36])
>>> a**p
tensor([0., 0.])

Parameters:

a (torch.Tensor) – Base tensor on which to calculate the power. Must be positive.
p (Union[float, torch.Tensor]) – Power value or tensor, by which to calculate the power.
eps (float (default 1e-12)) – The eps to add to the base tensor.

Returns:

The result tensor.

Return type:

torch.Tensor

Raises:

RuntimeError – If the base tensor contains negative values.

safenorm(a: Tensor, dim: int, eps: float = 1e-15) → Tensor[source]¶

Calculate the 2-norm safely. The square root of the inner product of a vector with itself.

Examples

>>> import torch
>>> a = torch.tensor([1e-35, 2e-40])
>>> safenorm(a, 0)
tensor(1.4142e-15)
>>> a.norm()
tensor(0.)

Parameters:

a (torch.Tensor) – The tensor to calculate the norm.
dim (int) – The dimension to calculate the norm.
eps (float (default 1e-15)) – The eps to add to the base tensor.

Returns:

The result tensor.

Return type:

torch.Tensor

occnumber(a: int | float | Tensor, n: int | None = None, dtype: dtype = torch.float64, device: device = device(type='cpu')) → Tensor[source]¶

Occupation number (maxed at 1) where the total sum of the output equals to a with length of the output is n.

Examples

>>> import torch
>>> occnumber(torch.tensor(2.5), 3, torch.double, torch.device('cpu'))
tensor([1.0000, 1.0000, 0.5000], dtype=torch.float64)
>>> occnumber(2.5)
tensor([1.0000, 1.0000, 0.5000], dtype=torch.float64)

Parameters:

a (ZType) – Total sum of the output
n (Optional[int] (default None)) – Length of the output
dtype (torch.dtype (default torch.double)) – Data type of the output
device (torch.device (default torch.device('cpu'))) – Device of the output

Returns:

The constructed occupation number

Return type:

torch.Tensor

get_floor_and_ceil(aa: int | float) → Tuple[int, int][source]¶

get the ceiling and flooring of aa.

Examples

>>> get_floor_and_ceil(2.5)
(2, 3)

Parameters:: aa (Union[int, float]) – The input number
Returns:: The flooring and ceiling of aa
Return type:: Tuple[int, int]

safe_cdist(a: Tensor, b: Tensor, add_diag_eps: bool = False, diag_inf: bool = False)[source]¶

L2 pairwise distance of a and b. The diagonal is either replaced with a small eps or infinite.

Examples

>>> import torch
>>> a = torch.tensor([[1., 2], [3, 4]])
>>> b = torch.tensor([[1., 2], [3, 4]])
>>> safe_cdist(a, b)
tensor([[0.0000, 2.8284],
        [2.8284, 0.0000]])
>>> safe_cdist(a, b, add_diag_eps=True)
tensor([[1.4142e-12, 2.8284e+00],
        [2.8284e+00, 1.4142e-12]])
>>> safe_cdist(a, b, diag_inf=True)
tensor([[   inf, 2.8284],
        [2.8284,    inf]])

Parameters:

a (torch.Tensor) – First Tensor. Shape: (*BA, na, ndim)
n (torch.Tensor) – Second Tensor. Shape: (*BB, nb, ndim)

Returns:

Pairwise distance. Shape: (*BAB, na, nb)

Return type:

torch.Tensor

Lightning Utilities¶

The utilities here are used for lightning specific operations and act as componenets for lightning modules.

collate_dataset_fn(batch_data: List[Tuple[Any, Any, Any, Any]], model)[source]¶

Default Collate function for DeepChem datasets to work with PyTorch DataLoader.

This function takes a batch of data from a PyTorch DataLoader and converts it into a format compatible with DeepChem models by wrapping it in a DeepChemBatch class that processes the data through the model’s default generator and batch preparation methods.

It does 3 important operations: 1. Extracts the features (X), labels (Y), weights (W), and ids from the batch and arranges them correctly. 2. Creates a NumpyDataset from these components and passes it to the model’s default generator. 3. Calls the model’s _prepare_batch method that outputs the processed batch as a tuple of tensors.

Parameters:

batch (List[Tuple[Any, Any, Any, Any]]) – Batch of data from DataLoader containing tuples of (X, y, w, ids).
model (TorchModel) – DeepChem model instance used for batch processing.

Returns:

Processed batch tuple prepared by the model’s _prepare_batch method.

Return type:

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

Examples

>>> import deepchem as dc
>>> import torch
>>> from torch.utils.data import DataLoader
>>> from deepchem.data import _TorchIndexDiskDataset as TorchIndexDiskDataset
>>> from deepchem.utils.lightning_utils import collate_dataset_fn
>>>
>>> # Load a dataset and create a model
>>> tasks, datasets, _ = dc.molnet.load_clintox()
>>> _, valid_dataset, _ = datasets
>>> model = dc.models.MultitaskClassifier(
...     n_tasks=len(tasks),
...     n_features=1024,
...     layer_sizes=[1000],
...     device="cpu",
...     batch_size=16
... )
>>>
>>> # Create DataLoader with custom collate function
>>> wrapped_dataset = TorchIndexDiskDataset(valid_dataset)
>>> dataloader = DataLoader(
...     wrapped_dataset,
...     batch_size=16,
...     collate_fn=lambda batch: collate_dataset_fn(batch, model)
... )
>>>
>>> # Use in training loop
>>> for batch in dataloader:
...     inputs, labels, weights = batch
...     # inputs, labels, weights are now properly formatted torch tensors
...     break

Caching Utils¶

Caching utilities to reuse computation results across active sessions. These prevent redundant recalculations by saving function outputs to disk or memory, improving performance for expensive and data-heavy operations (e.g., equivariant basis).

FileSystemMutex(filename: str)[source]¶

A cross-platform file-based mutual exclusion (mutex) object.

This class ensures that only one process at a time can access or modify cache files. It uses the file system as a locking mechanism [1].

On Unix: uses fcntl.lockf
On Windows: uses msvcrt.locking

This lock works across multiple CPUs on the same session and prevents race conditions when several processes try to read or write the same file. The lock is automatically released when leaving the with block when called or when the process terminates.

Parameters:: filename (str) – Path to the lock file.

Example

>>> from time import sleep
>>> import os
>>> mutex = FileSystemMutex("example.lock")
>>> # Using the mutex directly
>>> with mutex:
...     print("Critical section start")
...     sleep(2)  # Simulate protected work
...     print("Critical section end")
...
Critical section start
Critical section end
>>> os.remove("example.lock")  # Clean up lock file

# In multi-process or multi-thread scenarios, # only one process can hold the lock at a time.

References

cached_dirpklgz(dirname: str = '', maxsize: int = 128, verbose: bool = False)[source]¶

Decorator that caches the output of a function on memory and on disk [1].

It uses: - A Least Recently Used (LRU) cache that uses in-memory storage

for fast repeated calls.

A persistent on-disk cache using compressed pickle files (.pkl.gz) so that results are available as long as the session exists.

Cache are protected by a FileSystemMutex to allow multi-process safe access.

Parameters:

dirname (str, optional) – Directory path where cache files will be stored. If empty, a temporary directory is created.
maxsize (int, optional) – Maximum number of cached items in memory (default = 128). When exceeded, LRU entries are removed.
verbose (bool, optional) – If True, prints debug information.

Returns:

decorator – A decorator that can be applied to any function to cache results.

Return type:

Callable

Example

>>> from deepchem.utils.cache_utils import cached_dirpklgz
>>> import shutil
>>> import time

>>> @cached_dirpklgz("cache/", verbose=True)
... def slow_add(a, b):
...     time.sleep(2)
...     return a + b
...
>>> t0 = time.time()
>>> slow_add(2, 3)
compute 0.pkl.gz... save 0.pkl.gz... done
5
>>> t1 = time.time()
>>> if t1 - t0 >= 2:
...     print(f"First call took more than 2 seconds")
First call took more than 2 seconds
>>> t2 = time.time()
>>> slow_add(2, 3)
5
>>> t3 = time.time()
>>> if t3 - t2 <= 1:
...     print(f"First call took less than 1 second")
First call took less than 1 second
>>> shutil.rmtree("cache/")  # Clean up cache directory

References