add chemprop files

Dockerfile

@@ -87,8 +87,6 @@ USER user
 RUN --mount=target=requirements.txt,source=requirements.txt \
     pip install --no-cache-dir --upgrade -r requirements.txt
 
-# Pull chemprop from github
-RUN git clone https://github.com/chemprop/chemprop.git $HOME/app/chemprop
 
 # build an empty conda environment with appropriate Python version
 RUN conda create --name chemprop_env python=3.11*
@@ -96,11 +94,11 @@ RUN conda create --name chemprop_env python=3.11*
 SHELL ["conda", "run", "--no-capture-output", "-n", "chemprop_env", "/bin/bash", "-c"]
 
 # Follow the installation instructions then clear the cache
-ADD $HOME/app/chemprop /data/chemprop
+ADD chemprop chemprop
 ADD LICENSE.txt pyproject.toml README.md ./
 RUN conda install pytorch cpuonly -c pytorch && \
     conda clean --all --yes && \
-    python -m pip install ./chemprop && \
+    python -m pip install . && \
     python -m pip cache purge
 
 

chemprop/__init__.py

@@ -0,0 +1,5 @@
+from . import data, featurizers, models, nn, utils, conf, exceptions, schedulers
+
+__all__ = ["data", "featurizers", "models", "nn", "utils", "conf", "exceptions", "schedulers"]
+
+__version__ = "2.0.0"

chemprop/cli/common.py

@@ -0,0 +1,183 @@
+import logging
+from argparse import ArgumentParser, Namespace, ArgumentError
+from pathlib import Path
+
+from chemprop.cli.utils import LookupAction
+from chemprop.cli.utils.args import uppercase
+from chemprop.featurizers import MoleculeFeaturizerRegistry, RxnMode, AtomFeatureMode
+
+logger = logging.getLogger(__name__)
+
+
+def add_common_args(parser: ArgumentParser) -> ArgumentParser:
+    data_args = parser.add_argument_group("Shared input data args")
+    data_args.add_argument(
+        "-s",
+        "--smiles-columns",
+        nargs="+",
+        help="The column names in the input CSV containing SMILES strings. If unspecified, uses the the 0th column.",
+    )
+    data_args.add_argument(
+        "-r",
+        "--reaction-columns",
+        nargs="+",
+        help="The column names in the input CSV containing reaction SMILES in the format 'REACTANT>AGENT>PRODUCT', where 'AGENT' is optional.",
+    )
+    data_args.add_argument(
+        "--no-header-row",
+        action="store_true",
+        help="If specified, the first row in the input CSV will not be used as column names.",
+    )
+
+    dataloader_args = parser.add_argument_group("Dataloader args")
+    dataloader_args.add_argument(
+        "-n",
+        "--num-workers",
+        type=int,
+        default=0,
+        help="""Number of workers for parallel data loading (0 means sequential).
+Warning: setting num_workers>0 can cause hangs on Windows and MacOS.""",
+    )
+    dataloader_args.add_argument("-b", "--batch-size", type=int, default=64, help="Batch size.")
+
+    parser.add_argument(
+        "--accelerator", default="auto", help="Passed directly to the lightning Trainer()."
+    )
+    parser.add_argument(
+        "--devices",
+        default="auto",
+        help="Passed directly to the lightning Trainer(). If specifying multiple devices, must be a single string of comma separated devices, e.g. '1, 2'.",
+    )
+
+    featurization_args = parser.add_argument_group("Featurization args")
+    featurization_args.add_argument(
+        "--rxn-mode",
+        "--reaction-mode",
+        type=uppercase,
+        default="REAC_DIFF",
+        choices=list(RxnMode.keys()),
+        help="""Choices for construction of atom and bond features for reactions (case insensitive):
+- 'reac_prod': concatenates the reactants feature with the products feature.
+- 'reac_diff': concatenates the reactants feature with the difference in features between reactants and products. (Default)
+- 'prod_diff': concatenates the products feature with the difference in features between reactants and products.
+- 'reac_prod_balance': concatenates the reactants feature with the products feature, balances imbalanced reactions.
+- 'reac_diff_balance': concatenates the reactants feature with the difference in features between reactants and products, balances imbalanced reactions.
+- 'prod_diff_balance': concatenates the products feature with the difference in features between reactants and products, balances imbalanced reactions.""",
+    )
+    # TODO: Update documenation for multi_hot_atom_featurizer_mode
+    featurization_args.add_argument(
+        "--multi-hot-atom-featurizer-mode",
+        type=uppercase,
+        default="V2",
+        choices=list(AtomFeatureMode.keys()),
+        help="""Choices for multi-hot atom featurization scheme. This will affect both non-reatction and reaction feturization (case insensitive):
+- `V1`: Corresponds to the original configuration employed in the Chemprop V1.
+- `V2`: Tailored for a broad range of molecules, this configuration encompasses all elements in the first four rows of the periodic table, along with iodine. It is the default in Chemprop V2.
+- `ORGANIC`: Designed specifically for use with organic molecules for drug research and development, this configuration includes a subset of elements most common in organic chemistry, including H, B, C, N, O, F, Si, P, S, Cl, Br, and I.""",
+    )
+    featurization_args.add_argument(
+        "--keep-h",
+        action="store_true",
+        help="Whether hydrogens explicitly specified in input should be kept in the mol graph.",
+    )
+    featurization_args.add_argument(
+        "--add-h", action="store_true", help="Whether hydrogens should be added to the mol graph."
+    )
+    featurization_args.add_argument(
+        "--features-generators",
+        nargs="+",
+        action=LookupAction(MoleculeFeaturizerRegistry),
+        help="Method(s) of generating additional features.",
+    )
+    featurization_args.add_argument(
+        "--descriptors-path",
+        type=Path,
+        help="Path to extra descriptors to concatenate to learned representation.",
+    )
+    # TODO: Add in v2.1
+    # featurization_args.add_argument(
+    #     "--phase-features-path",
+    #     help="Path to features used to indicate the phase of the data in one-hot vector form. Used in spectra datatype.",
+    # )
+    featurization_args.add_argument(
+        "--no-descriptor-scaling", action="store_true", help="Turn off extra descriptor scaling."
+    )
+    featurization_args.add_argument(
+        "--no-atom-feature-scaling",
+        action="store_true",
+        help="Turn off extra atom feature scaling.",
+    )
+    featurization_args.add_argument(
+        "--no-atom-descriptor-scaling",
+        action="store_true",
+        help="Turn off extra atom descriptor scaling.",
+    )
+    featurization_args.add_argument(
+        "--no-bond-feature-scaling",
+        action="store_true",
+        help="Turn off extra bond feature scaling.",
+    )
+    featurization_args.add_argument(
+        "--atom-features-path",
+        nargs="+",
+        action="append",
+        help="If a single path is given, it's assumed to correspond to the 0-th molecule. Or, it can be a two-tuple of molecule index and path to additional atom features to supply before message passing. E.g., `--atom-features-path 0 /path/to/features_0.npz` indicates that the features at the given path should be supplied to the 0-th component. To supply additional features for multiple components, repeat this argument on the command line for each component's respective values, e.g., `--atom-features-path [...] --atom-features-path [...]`.",
+    )
+    featurization_args.add_argument(
+        "--atom-descriptors-path",
+        nargs="+",
+        action="append",
+        help="If a single path is given, it's assumed to correspond to the 0-th molecule. Or, it can be a two-tuple of molecule index and path to additional atom descriptors to supply after message passing. E.g., `--atom-descriptors-path 0 /path/to/descriptors_0.npz` indicates that the descriptors at the given path should be supplied to the 0-th component. To supply additional descriptors for multiple components, repeat this argument on the command line for each component's respective values, e.g., `--atom-descriptors-path [...] --atom-descriptors-path [...]`.",
+    )
+    featurization_args.add_argument(
+        "--bond-features-path",
+        nargs="+",
+        action="append",
+        help="If a single path is given, it's assumed to correspond to the 0-th molecule. Or, it can be a two-tuple of molecule index and path to additional bond features to supply before message passing. E.g., `--bond-features-path 0 /path/to/features_0.npz` indicates that the features at the given path should be supplied to the 0-th component. To supply additional features for multiple components, repeat this argument on the command line for each component's respective values, e.g., `--bond-features-path [...] --bond-features-path [...]`.",
+    )
+    # TODO: Add in v2.2
+    # parser.add_argument(
+    #     "--constraints-path",
+    #     help="Path to constraints applied to atomic/bond properties prediction.",
+    # )
+
+    return parser
+
+
+def process_common_args(args: Namespace) -> Namespace:
+    for key in ["atom_features_path", "atom_descriptors_path", "bond_features_path"]:
+        inds_paths = getattr(args, key)
+
+        if not inds_paths:
+            continue
+
+        ind_path_dict = {}
+
+        for ind_path in inds_paths:
+            if len(ind_path) > 2:
+                raise ArgumentError(
+                    argument=None,
+                    message="Too many arguments given for atom features/descriptors or bond features. It can be either a two-tuple of molecule index and a path, or a single path (assumed to be the 0-th molecule).",
+                )
+
+            if len(ind_path) == 1:
+                ind = 0
+                path = ind_path[0]
+            else:
+                ind, path = ind_path
+
+            if ind_path_dict.get(int(ind), None):
+                raise ArgumentError(
+                    argument=None,
+                    message=f"Duplicate atom features/descriptors or bond features given for molecule index {ind}.",
+                )
+
+            ind_path_dict[int(ind)] = Path(path)
+
+        setattr(args, key, ind_path_dict)
+
+    return args
+
+
+def validate_common_args(args):
+    pass

chemprop/cli/conf.py

@@ -0,0 +1,8 @@
+from datetime import datetime
+import logging
+import os
+from pathlib import Path
+
+LOG_DIR = Path(os.getenv("CHEMPROP_LOG_DIR", "chemprop_logs"))
+LOG_LEVELS = [logging.ERROR, logging.WARNING, logging.INFO, logging.DEBUG]
+NOW = datetime.now().strftime("%Y-%m-%dT%H-%M-%S")

chemprop/cli/convert.py

@@ -0,0 +1,55 @@
+from argparse import ArgumentError, ArgumentParser, Namespace
+import sys
+import logging
+from pathlib import Path
+
+from chemprop.cli.utils import Subcommand
+from chemprop.utils.v1_to_v2 import convert_model_file_v1_to_v2
+
+logger = logging.getLogger(__name__)
+
+
+class ConvertSubcommand(Subcommand):
+    COMMAND = "convert"
+    HELP = "convert a v1 model checkpoint (.pt) to a v2 model checkpoint (.ckpt)"
+
+    @classmethod
+    def add_args(cls, parser: ArgumentParser) -> ArgumentParser:
+        parser.add_argument(
+            "-i",
+            "--input-path",
+            required=True,
+            type=Path,
+            help="The path to a v1 model .pt checkpoint file.",
+        )
+        parser.add_argument(
+            "-o",
+            "--output-path",
+            type=Path,
+            help="The path to which the converted model will be saved. Defaults to 'CURRENT_DIRECTORY/STEM_OF_INPUT_v2.ckpt'",
+        )
+        return parser
+
+    @classmethod
+    def func(cls, args: Namespace):
+        if args.output_path is None:
+            args.output_path = Path(args.input_path.stem + "_v2.ckpt")
+        if args.output_path.suffix != ".ckpt":
+            raise ArgumentError(
+                argument=None, message=f"Output must be a `.ckpt` file. Got {args.output_path}"
+            )
+
+        logger.info(
+            f"Converting v1 model checkpoint '{args.input_path}' to v2 model checkpoint '{args.output_path}'..."
+        )
+        convert_model_file_v1_to_v2(args.input_path, args.output_path)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser = ConvertSubcommand.add_args(parser)
+
+    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, force=True)
+
+    args = parser.parse_args()
+    ConvertSubcommand.func(args)

chemprop/cli/fingerprint.py

@@ -0,0 +1,197 @@
+import logging
+import sys
+from argparse import ArgumentError, ArgumentParser, Namespace
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+
+from chemprop import data
+from chemprop.cli.common import add_common_args, process_common_args, validate_common_args
+from chemprop.cli.utils import Subcommand, build_data_from_files, make_dataset
+from chemprop.featurizers import MoleculeFeaturizerRegistry
+from chemprop.models import load_model
+from chemprop.nn.loss import LossFunctionRegistry
+from chemprop.utils import Factory
+
+logger = logging.getLogger(__name__)
+
+
+class FingerprintSubcommand(Subcommand):
+    COMMAND = "fingerprint"
+    HELP = "use a pretrained chemprop model for to calculate learned representations"
+
+    @classmethod
+    def add_args(cls, parser: ArgumentParser) -> ArgumentParser:
+        parser = add_common_args(parser)
+        parser.add_argument(
+            "-i",
+            "--test-path",
+            required=True,
+            type=Path,
+            help="Path to an input CSV file containing SMILES.",
+        )
+        parser.add_argument(
+            "-o",
+            "--output",
+            "--preds-path",
+            type=Path,
+            help="Path to which predictions will be saved. If the file extension is .npz, they will be saved as a npz file, respectively. Otherwise, will save predictions as a CSV. The index of the model will be appended to the filename's stem. By default, predictions will be saved to the same location as '--test-path' with '_fps' appended, i.e., 'PATH/TO/TEST_PATH_fps_0.csv'.",
+        )
+        parser.add_argument(
+            "--model-path",
+            required=True,
+            type=Path,
+            help="Path to either a single pretrained model checkpoint (.ckpt) or single pretrained model file (.pt) or to a directory that contains these files. If a directory, will recursively search and predict on all found models.",
+        )
+        parser.add_argument(
+            "--ffn-block-index",
+            required=True,
+            type=int,
+            default=-1,
+            help="The index indicates which linear layer returns the encoding in the FFN. An index of 0 denotes the post-aggregation representation through a 0-layer MLP, while an index of 1 represents the output from the first linear layer in the FFN, and so forth.",
+        )
+
+        return parser
+
+    @classmethod
+    def func(cls, args: Namespace):
+        args = process_common_args(args)
+        validate_common_args(args)
+        args = process_fingerprint_args(args)
+        main(args)
+
+
+def process_fingerprint_args(args: Namespace) -> Namespace:
+    if args.test_path.suffix not in [".csv"]:
+        raise ArgumentError(
+            argument=None, message=f"Input data must be a CSV file. Got {args.test_path}"
+        )
+    if args.output is None:
+        args.output = args.test_path.parent / (args.test_path.stem + "_fps.csv")
+    if args.output.suffix not in [".csv", ".npz"]:
+        raise ArgumentError(
+            argument=None, message=f"Output must be a CSV or NPZ file. Got '{args.output}'."
+        )
+    return args
+
+
+def find_models(model_path: Path):
+    if model_path.suffix in [".ckpt", ".pt"]:
+        return [model_path]
+    elif model_path.is_dir():
+        return list(model_path.rglob("*.ckpt")) + list(model_path.rglob("*.pt"))
+
+
+def make_fingerprint_for_model(
+    args: Namespace, model_path: Path, multicomponent: bool, output_path: Path
+):
+    model = load_model(model_path, multicomponent)
+    model.eval()
+
+    bounded = any(
+        isinstance(model.criterion, LossFunctionRegistry[loss_function])
+        for loss_function in LossFunctionRegistry.keys()
+        if "bounded" in loss_function
+    )
+
+    format_kwargs = dict(
+        no_header_row=args.no_header_row,
+        smiles_cols=args.smiles_columns,
+        rxn_cols=args.reaction_columns,
+        target_cols=None,
+        ignore_cols=None,
+        splits_col=None,
+        weight_col=None,
+        bounded=bounded,
+    )
+
+    if args.features_generators is not None:
+        # TODO: MorganFeaturizers take radius, length, and include_chirality as arguements. Should we expose these through the CLI?
+        features_generators = [
+            Factory.build(MoleculeFeaturizerRegistry[features_generator])
+            for features_generator in args.features_generators
+        ]
+    else:
+        features_generators = None
+
+    featurization_kwargs = dict(
+        features_generators=features_generators, keep_h=args.keep_h, add_h=args.add_h
+    )
+
+    test_data = build_data_from_files(
+        args.test_path,
+        **format_kwargs,
+        p_descriptors=args.descriptors_path,
+        p_atom_feats=args.atom_features_path,
+        p_bond_feats=args.bond_features_path,
+        p_atom_descs=args.atom_descriptors_path,
+        **featurization_kwargs,
+    )
+    logger.info(f"test size: {len(test_data[0])}")
+    test_dsets = [
+        make_dataset(d, args.rxn_mode, args.multi_hot_atom_featurizer_mode) for d in test_data
+    ]
+
+    if multicomponent:
+        test_dset = data.MulticomponentDataset(test_dsets)
+    else:
+        test_dset = test_dsets[0]
+
+    test_loader = data.build_dataloader(test_dset, args.batch_size, args.num_workers, shuffle=False)
+
+    logger.info(model)
+
+    with torch.no_grad():
+        if multicomponent:
+            encodings = [
+                model.encoding(batch.bmgs, batch.V_ds, batch.X_d, args.ffn_block_index)
+                for batch in test_loader
+            ]
+        else:
+            encodings = [
+                model.encoding(batch.bmg, batch.V_d, batch.X_d, args.ffn_block_index)
+                for batch in test_loader
+            ]
+        H = torch.cat(encodings, 0).numpy()
+
+    if output_path.suffix in [".npz"]:
+        np.savez(output_path, H=H)
+    elif output_path.suffix == ".csv":
+        fingerprint_columns = [f"fp_{i}" for i in range(H.shape[1])]
+        df_fingerprints = pd.DataFrame(H, columns=fingerprint_columns)
+        df_fingerprints.to_csv(output_path, index=False)
+    else:
+        raise ArgumentError(
+            argument=None, message=f"Output must be a CSV or npz file. Got {args.output}."
+        )
+    logger.info(f"Fingerprints saved to '{output_path}'")
+
+
+def main(args):
+    match (args.smiles_columns, args.reaction_columns):
+        case [None, None]:
+            n_components = 1
+        case [_, None]:
+            n_components = len(args.smiles_columns)
+        case [None, _]:
+            n_components = len(args.reaction_columns)
+        case _:
+            n_components = len(args.smiles_columns) + len(args.reaction_columns)
+
+    multicomponent = n_components > 1
+
+    for i, model_path in enumerate(find_models(args.model_path)):
+        logger.info(f"Fingerprints with model at '{model_path}'")
+        output_path = args.output.parent / f"{args.output.stem}_{i}{args.output.suffix}"
+        make_fingerprint_for_model(args, model_path, multicomponent, output_path)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser = FingerprintSubcommand.add_args(parser)
+
+    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, force=True)
+    args = parser.parse_args()
+    args = FingerprintSubcommand.func(args)

chemprop/cli/hpopt.py

@@ -0,0 +1,442 @@
+import json
+import logging
+import sys
+from argparse import ArgumentParser, Namespace
+from copy import deepcopy
+from pathlib import Path
+import torch
+from lightning import pytorch as pl
+from lightning.pytorch.callbacks import EarlyStopping
+
+from chemprop.cli.common import add_common_args, process_common_args, validate_common_args
+from chemprop.cli.train import (
+    add_train_args,
+    build_datasets,
+    build_model,
+    build_splits,
+    normalize_inputs,
+    process_train_args,
+    validate_train_args,
+)
+from chemprop.cli.utils.command import Subcommand
+from chemprop.data import build_dataloader
+from chemprop.featurizers import MoleculeFeaturizerRegistry
+from chemprop.nn import AggregationRegistry
+from chemprop.nn.transforms import UnscaleTransform
+from chemprop.nn.utils import Activation
+from chemprop.utils import Factory
+
+NO_RAY = False
+DEFAULT_SEARCH_SPACE = {}
+try:
+    import ray
+    from ray import tune
+    from ray.train import CheckpointConfig, RunConfig, ScalingConfig
+    from ray.train.lightning import (
+        RayDDPStrategy,
+        RayLightningEnvironment,
+        RayTrainReportCallback,
+        prepare_trainer,
+    )
+    from ray.train.torch import TorchTrainer
+    from ray.tune.schedulers import ASHAScheduler
+
+    DEFAULT_SEARCH_SPACE = {
+        "activation": tune.choice(categories=list(Activation.keys())),
+        "aggregation": tune.choice(categories=list(AggregationRegistry.keys())),
+        "aggregation_norm": tune.quniform(lower=1, upper=200, q=1),
+        "batch_size": tune.choice([16, 32, 64, 128, 256]),
+        "depth": tune.qrandint(lower=2, upper=6, q=1),
+        "dropout": tune.choice([tune.choice([0.0]), tune.quniform(lower=0.05, upper=0.4, q=0.05)]),
+        "ffn_hidden_dim": tune.qrandint(lower=300, upper=2400, q=100),
+        "ffn_num_layers": tune.qrandint(lower=1, upper=3, q=1),
+        "final_lr_ratio": tune.loguniform(lower=1e-2, upper=1),
+        "message_hidden_dim": tune.qrandint(lower=300, upper=2400, q=100),
+        "init_lr_ratio": tune.loguniform(lower=1e-2, upper=1),
+        "max_lr": tune.loguniform(lower=1e-4, upper=1e-2),
+        "warmup_epochs": None,
+    }
+except ImportError:
+    NO_RAY = True
+
+NO_HYPEROPT = False
+try:
+    from ray.tune.search.hyperopt import HyperOptSearch
+except ImportError:
+    NO_HYPEROPT = True
+
+# NO_OPTUNA = False
+# try:
+#     from ray.tune.search.optuna import OptunaSearch
+# except ImportError:
+#     NO_OPTUNA = True
+
+
+logger = logging.getLogger(__name__)
+
+SEARCH_SPACE = DEFAULT_SEARCH_SPACE
+
+SEARCH_PARAM_KEYWORDS_MAP = {
+    "basic": ["depth", "ffn_num_layers", "dropout", "ffn_hidden_dim", "message_hidden_dim"],
+    "learning_rate": ["max_lr", "init_lr_ratio", "final_lr_ratio", "warmup_epochs"],
+    "all": list(DEFAULT_SEARCH_SPACE.keys()),
+}
+
+
+class HpoptSubcommand(Subcommand):
+    COMMAND = "hpopt"
+    HELP = "perform hyperparameter optimization on the given task"
+
+    @classmethod
+    def add_args(cls, parser: ArgumentParser) -> ArgumentParser:
+        parser = add_common_args(parser)
+        parser = add_train_args(parser)
+        return add_hpopt_args(parser)
+
+    @classmethod
+    def func(cls, args: Namespace):
+        args = process_common_args(args)
+        args = process_train_args(args)
+        args = process_hpopt_args(args)
+        validate_common_args(args)
+        validate_train_args(args)
+        main(args)
+
+
+def add_hpopt_args(parser: ArgumentParser) -> ArgumentParser:
+    hpopt_args = parser.add_argument_group("Chemprop hyperparameter optimization arguments")
+
+    hpopt_args.add_argument(
+        "--search-parameter-keywords",
+        type=str,
+        nargs="+",
+        default=["basic"],
+        help=f"""The model parameters over which to search for an optimal hyperparameter configuration.
+    Some options are bundles of parameters or otherwise special parameter operations.
+
+    Special keywords:
+        basic - the default set of hyperparameters for search: depth, ffn_num_layers, dropout, message_hidden_dim, and ffn_hidden_dim.
+        learning_rate - search for max_lr, init_lr_ratio, final_lr_ratio, and warmup_epochs. The search for init_lr and final_lr values
+            are defined as fractions of the max_lr value. The search for warmup_epochs is as a fraction of the total epochs used.
+        all - include search for all 13 inidividual keyword options
+
+    Individual supported parameters:
+        {list(DEFAULT_SEARCH_SPACE.keys())}
+    """,
+    )
+
+    hpopt_args.add_argument(
+        "--hpopt-save-dir",
+        type=Path,
+        help="Directory to save the hyperparameter optimization results",
+    )
+
+    raytune_args = parser.add_argument_group("Ray Tune arguments")
+
+    raytune_args.add_argument(
+        "--raytune-num-samples",
+        type=int,
+        default=10,
+        help="Passed directly to Ray Tune TuneConfig to control number of trials to run",
+    )
+
+    raytune_args.add_argument(
+        "--raytune-search-algorithm",
+        choices=["random", "hyperopt"],  # , "optuna"],
+        default="hyperopt",
+        help="Passed to Ray Tune TuneConfig to control search algorithm",
+    )
+
+    raytune_args.add_argument(
+        "--raytune-num-workers",
+        type=int,
+        default=1,
+        help="Passed directly to Ray Tune ScalingConfig to control number of workers to use",
+    )
+
+    raytune_args.add_argument(
+        "--raytune-use-gpu",
+        action="store_true",
+        help="Passed directly to Ray Tune ScalingConfig to control whether to use GPUs",
+    )
+
+    raytune_args.add_argument(
+        "--raytune-num-checkpoints-to-keep",
+        type=int,
+        default=1,
+        help="Passed directly to Ray Tune CheckpointConfig to control number of checkpoints to keep",
+    )
+
+    raytune_args.add_argument(
+        "--raytune-grace-period",
+        type=int,
+        default=10,
+        help="Passed directly to Ray Tune ASHAScheduler to control grace period",
+    )
+
+    raytune_args.add_argument(
+        "--raytune-reduction-factor",
+        type=int,
+        default=2,
+        help="Passed directly to Ray Tune ASHAScheduler to control reduction factor",
+    )
+
+    hyperopt_args = parser.add_argument_group("Hyperopt arguments")
+
+    hyperopt_args.add_argument(
+        "--hyperopt-n-initial-points",
+        type=int,
+        default=20,
+        help="Passed directly to HyperOptSearch to control number of initial points to sample",
+    )
+
+    hyperopt_args.add_argument(
+        "--hyperopt-random-state-seed",
+        type=int,
+        default=None,
+        help="Passed directly to HyperOptSearch to control random state seed",
+    )
+
+    return parser
+
+
+def process_hpopt_args(args: Namespace) -> Namespace:
+    if args.hpopt_save_dir is None:
+        args.hpopt_save_dir = Path(f"chemprop_hpopt/{args.data_path.stem}")
+
+    args.hpopt_save_dir.mkdir(exist_ok=True, parents=True)
+
+    search_parameters = set()
+
+    for keyword in args.search_parameter_keywords:
+        if keyword not in SEARCH_PARAM_KEYWORDS_MAP and keyword not in SEARCH_SPACE:
+            raise ValueError(
+                f"Search parameter keyword: {keyword} not in available options: {list(SEARCH_PARAM_KEYWORDS_MAP.keys()) + list(SEARCH_SPACE.keys())}."
+            )
+
+        search_parameters.update(
+            SEARCH_PARAM_KEYWORDS_MAP[keyword]
+            if keyword in SEARCH_PARAM_KEYWORDS_MAP
+            else [keyword]
+        )
+
+    args.search_parameter_keywords = list(search_parameters)
+
+    return args
+
+
+def build_search_space(search_parameters: list[str], train_epochs: int) -> dict:
+    if "warmup_epochs" in search_parameters and SEARCH_SPACE.get("warmup_epochs", None) is None:
+        SEARCH_SPACE["warmup_epochs"] = tune.qrandint(lower=1, upper=train_epochs // 2, q=1)
+
+    return {param: SEARCH_SPACE[param] for param in search_parameters}
+
+
+def update_args_with_config(args: Namespace, config: dict) -> Namespace:
+    args = deepcopy(args)
+
+    for key, value in config.items():
+        match key:
+            case "final_lr_ratio":
+                setattr(args, "final_lr", value * args.max_lr)
+
+            case "init_lr_ratio":
+                setattr(args, "init_lr", value * args.max_lr)
+
+            case _:
+                assert key in args, f"Key: {key} not found in args."
+                setattr(args, key, value)
+
+    return args
+
+
+def train_model(config, args, train_dset, val_dset, logger, output_transform, input_transforms):
+    update_args_with_config(args, config)
+
+    train_loader = build_dataloader(
+        train_dset, args.batch_size, args.num_workers, seed=args.data_seed
+    )
+    val_loader = build_dataloader(val_dset, args.batch_size, args.num_workers, shuffle=False)
+
+    seed = args.pytorch_seed if args.pytorch_seed is not None else torch.seed()
+
+    torch.manual_seed(seed)
+
+    model = build_model(args, train_loader.dataset, output_transform, input_transforms)
+    logger.info(model)
+
+    monitor_mode = "min" if model.metrics[0].minimize else "max"
+    logger.debug(f"Evaluation metric: '{model.metrics[0].alias}', mode: '{monitor_mode}'")
+
+    patience = args.patience if args.patience is not None else args.epochs
+    early_stopping = EarlyStopping("val_loss", patience=patience, mode=monitor_mode)
+
+    trainer = pl.Trainer(
+        accelerator=args.accelerator,
+        devices=args.devices,
+        max_epochs=args.epochs,
+        gradient_clip_val=args.grad_clip,
+        strategy=RayDDPStrategy(find_unused_parameters=True),
+        callbacks=[RayTrainReportCallback(), early_stopping],
+        plugins=[RayLightningEnvironment()],
+        deterministic=args.pytorch_seed is not None,
+    )
+    trainer = prepare_trainer(trainer)
+    trainer.fit(model, train_loader, val_loader)
+
+
+def tune_model(
+    args, train_dset, val_dset, logger, monitor_mode, output_transform, input_transforms
+):
+    scheduler = ASHAScheduler(
+        max_t=args.epochs,
+        grace_period=min(args.raytune_grace_period, args.epochs),
+        reduction_factor=args.raytune_reduction_factor,
+    )
+
+    scaling_config = ScalingConfig(
+        num_workers=args.raytune_num_workers, use_gpu=args.raytune_use_gpu
+    )
+
+    checkpoint_config = CheckpointConfig(
+        num_to_keep=args.raytune_num_checkpoints_to_keep,
+        checkpoint_score_attribute="val_loss",
+        checkpoint_score_order=monitor_mode,
+    )
+
+    run_config = RunConfig(
+        checkpoint_config=checkpoint_config,
+        storage_path=args.hpopt_save_dir.absolute() / "ray_results",
+    )
+
+    ray_trainer = TorchTrainer(
+        lambda config: train_model(
+            config, args, train_dset, val_dset, logger, output_transform, input_transforms
+        ),
+        scaling_config=scaling_config,
+        run_config=run_config,
+    )
+
+    match args.raytune_search_algorithm:
+        case "random":
+            search_alg = None
+        case "hyperopt":
+            if NO_HYPEROPT:
+                raise ImportError(
+                    "HyperOptSearch requires hyperopt to be installed. Use 'pip -U install hyperopt' to install."
+                )
+
+            search_alg = HyperOptSearch(
+                n_initial_points=args.hyperopt_n_initial_points,
+                random_state_seed=args.hyperopt_random_state_seed,
+            )
+        # case "optuna":
+        #     if NO_OPTUNA:
+        #         raise ImportError(
+        #             "OptunaSearch requires optuna to be installed. Use 'pip -U install optuna' to install."
+        #         )
+
+        #     search_alg = OptunaSearch()
+
+    tune_config = tune.TuneConfig(
+        metric="val_loss",
+        mode=monitor_mode,
+        num_samples=args.raytune_num_samples,
+        scheduler=scheduler,
+        search_alg=search_alg,
+    )
+
+    tuner = tune.Tuner(
+        ray_trainer,
+        param_space={
+            "train_loop_config": build_search_space(args.search_parameter_keywords, args.epochs)
+        },
+        tune_config=tune_config,
+    )
+
+    return tuner.fit()
+
+
+def main(args: Namespace):
+    if NO_RAY:
+        raise ImportError(
+            "Ray Tune requires ray to be installed. Use 'pip -U install ray[tune]' to install."
+        )
+
+    format_kwargs = dict(
+        no_header_row=args.no_header_row,
+        smiles_cols=args.smiles_columns,
+        rxn_cols=args.reaction_columns,
+        target_cols=args.target_columns,
+        ignore_cols=args.ignore_columns,
+        splits_col=args.splits_column,
+        weight_col=args.weight_column,
+        bounded=args.loss_function is not None and "bounded" in args.loss_function,
+    )
+
+    if args.features_generators is not None:
+        # TODO: MorganFeaturizers take radius, length, and include_chirality as arguements. Should we expose these through the CLI?
+        features_generators = [
+            Factory.build(MoleculeFeaturizerRegistry[features_generator])
+            for features_generator in args.features_generators
+        ]
+    else:
+        features_generators = None
+
+    featurization_kwargs = dict(
+        features_generators=features_generators, keep_h=args.keep_h, add_h=args.add_h
+    )
+
+    train_data, val_data, test_data = build_splits(args, format_kwargs, featurization_kwargs)
+    train_dset, val_dset, test_dset = build_datasets(args, train_data[0], val_data[0], test_data[0])
+
+    input_transforms = normalize_inputs(train_dset, val_dset, args)
+
+    if "regression" in args.task_type:
+        output_scaler = train_dset.normalize_targets()
+        val_dset.normalize_targets(output_scaler)
+        logger.info(f"Train data: mean = {output_scaler.mean_} | std = {output_scaler.scale_}")
+        output_transform = UnscaleTransform.from_standard_scaler(output_scaler)
+    else:
+        output_transform = None
+
+    train_loader = build_dataloader(
+        train_dset, args.batch_size, args.num_workers, seed=args.data_seed
+    )
+
+    model = build_model(args, train_loader.dataset, output_transform, input_transforms)
+    monitor_mode = "min" if model.metrics[0].minimize else "max"
+
+    results = tune_model(
+        args, train_dset, val_dset, logger, monitor_mode, output_transform, input_transforms
+    )
+
+    best_result = results.get_best_result()
+    best_config = best_result.config
+    best_checkpoint = best_result.checkpoint  # Get best trial's best checkpoint
+
+    logger.info(f"Saving best hyperparameter parameters: {best_config}")
+
+    with open(args.hpopt_save_dir / "best_params.json", "w") as f:
+        json.dump(best_config, f, indent=4)
+
+    logger.info(f"Saving best hyperparameter configuration checkpoint: {best_checkpoint}")
+
+    torch.save(best_checkpoint, args.hpopt_save_dir / "best_checkpoint.ckpt")
+
+    result_df = results.get_dataframe()
+
+    logger.info(f"Saving hyperparameter optimization results: {result_df}")
+
+    result_df.to_csv(args.hpopt_save_dir / "all_progress.csv", index=False)
+
+    ray.shutdown()
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser = HpoptSubcommand.add_args(parser)
+
+    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, force=True)
+    args = parser.parse_args()
+    HpoptSubcommand.func(args)

chemprop/cli/main.py

@@ -0,0 +1,80 @@
+from configargparse import ArgumentParser
+import logging
+import sys
+from pathlib import Path
+
+from chemprop.cli.train import TrainSubcommand
+from chemprop.cli.predict import PredictSubcommand
+from chemprop.cli.convert import ConvertSubcommand
+from chemprop.cli.fingerprint import FingerprintSubcommand
+from chemprop.cli.hpopt import HpoptSubcommand
+
+from chemprop.cli.utils import pop_attr
+from chemprop.cli.conf import LOG_DIR, LOG_LEVELS, NOW
+
+logger = logging.getLogger(__name__)
+
+SUBCOMMANDS = [
+    TrainSubcommand,
+    PredictSubcommand,
+    ConvertSubcommand,
+    FingerprintSubcommand,
+    HpoptSubcommand,
+]
+
+
+def construct_parser():
+    parser = ArgumentParser()
+    subparsers = parser.add_subparsers(title="mode", dest="mode", required=True)
+
+    parent = ArgumentParser(add_help=False)
+    parent.add_argument(
+        "--logfile",
+        "--log",
+        nargs="?",
+        const="default",
+        help=f"The path to which the log file should be written. Specifying just the flag (i.e., '--log/--logfile') will automatically log to a file '{LOG_DIR}/MODE/TIMESTAMP.log', where 'MODE' is the CLI mode chosen. An example 'TIMESTAMP' is {NOW}.",
+    )
+    parent.add_argument(
+        "-v",
+        "--verbose",
+        action="count",
+        default=0,
+        help="The verbosity level, specify the flag multiple times to increase verbosity.",
+    )
+
+    parents = [parent]
+    for subcommand in SUBCOMMANDS:
+        subcommand.add(subparsers, parents)
+
+    return parser
+
+
+def main():
+    parser = construct_parser()
+    args = parser.parse_args()
+    logfile, verbose, mode, func = (
+        pop_attr(args, attr) for attr in ["logfile", "verbose", "mode", "func"]
+    )
+
+    match logfile:
+        case None:
+            handler = logging.StreamHandler(sys.stderr)
+        case "default":
+            (LOG_DIR / mode).mkdir(parents=True, exist_ok=True)
+            handler = logging.FileHandler(str(LOG_DIR / mode / f"{NOW}.log"))
+        case _:
+            Path(logfile).parent.mkdir(parents=True, exist_ok=True)
+            handler = logging.FileHandler(logfile)
+
+    logging.basicConfig(
+        handlers=[handler],
+        format="%(asctime)s - %(levelname)s:%(name)s - %(message)s",
+        level=LOG_LEVELS[min(verbose, len(LOG_LEVELS) - 1)],
+        datefmt="%Y-%m-%dT%H:%M:%S",
+        force=True,
+    )
+
+    logger.info(f"Running in mode '{mode}' with args: {vars(args)}")
+
+    func(args)

chemprop/cli/predict.py

@@ -0,0 +1,315 @@
+from argparse import ArgumentError, ArgumentParser, Namespace
+import logging
+from pathlib import Path
+import sys
+import pandas as pd
+
+from lightning import pytorch as pl
+import torch
+
+from chemprop import data
+from chemprop.nn.loss import LossFunctionRegistry
+from chemprop.nn.predictors import MulticlassClassificationFFN
+from chemprop.models import load_model
+
+from chemprop.cli.utils import Subcommand, build_data_from_files, make_dataset
+from chemprop.cli.common import add_common_args, process_common_args, validate_common_args
+
+
+logger = logging.getLogger(__name__)
+
+
+class PredictSubcommand(Subcommand):
+    COMMAND = "predict"
+    HELP = "use a pretrained chemprop model for prediction"
+
+    @classmethod
+    def add_args(cls, parser: ArgumentParser) -> ArgumentParser:
+        parser = add_common_args(parser)
+        return add_predict_args(parser)
+
+    @classmethod
+    def func(cls, args: Namespace):
+        args = process_common_args(args)
+        validate_common_args(args)
+        args = process_predict_args(args)
+        main(args)
+
+
+def add_predict_args(parser: ArgumentParser) -> ArgumentParser:
+    parser.add_argument(
+        "-i",
+        "--test-path",
+        required=True,
+        type=Path,
+        help="Path to an input CSV file containing SMILES.",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        "--preds-path",
+        type=Path,
+        help="Path to which predictions will be saved. If the file extension is .pkl, will be saved as a pickle file. Otherwise, will save predictions as a CSV. The index of the model will be appended to the filename's stem. By default, predictions will be saved to the same location as '--test-path' with '_preds' appended, i.e., 'PATH/TO/TEST_PATH_preds_0.csv'.",
+    )
+    parser.add_argument(
+        "--drop-extra-columns",
+        action="store_true",
+        help="Whether to drop all columns from the test data file besides the SMILES columns and the new prediction columns.",
+    )
+    parser.add_argument(
+        "--model-path",
+        required=True,
+        type=Path,
+        help="Path to either a single pretrained model checkpoint (.ckpt) or single pretrained model file (.pt) or to a directory that contains these files. If a directory, will recursively search and predict on all found models.",
+    )
+    parser.add_argument(
+        "--target-columns",
+        nargs="+",
+        help="Column names to save the predictions to. If not provided, the predictions will be saved to columns named 'pred_0', 'pred_1', etc.",
+    )
+
+    # TODO: add uncertainty and calibration in v2.1
+    # unc_args = parser.add_argument_group("Uncertainty and calibration args")
+    # unc_args.add_argument("--cal-path")
+    # unc_args.add_argument("--cal-features-path")
+    # unc_args.add_argument("--cal-atom-features-path")
+    # unc_args.add_argument("--cal-bond-features-path")
+    # unc_args.add_argument("--cal-atom-descriptors-path")
+    # unc_args.add_argument(
+    #     "--ensemble-variance",
+    #     type=None,
+    #     help="Deprecated. Whether to calculate the variance of ensembles as a measure of epistemic uncertainty. If True, the variance is saved as an additional column for each target in the preds_path.",
+    # )
+    # unc_args.add_argument(
+    #     "--individual-ensemble-predictions",
+    #     type=bool,
+    #     action="store_true",
+    #     help="Whether to return the predictions made by each of the individual models rather than the average of the ensemble.",
+    # )
+    # unc_args.add_argument(
+    #     "--uncertainty-method",
+    #     #action=RegistryAction(TODO: make register for uncertainty methods)
+    #     help="The method of calculating uncertainty.",
+    # )
+    # unc_args.add_argument(
+    #     "--calibration-method",
+    #     #action=RegistryAction(TODO: make register for calibration methods)
+    #     help="Methods used for calibrating the uncertainty calculated with uncertainty method.",
+    # )
+    # unc_args.add_argument(
+    #     "--evaluation-method",
+    #     #action=RegistryAction(TODO: make register for evaluation methods)
+    #     type=list[str],
+    #     help="The methods used for evaluating the uncertainty performance if the test data provided includes targets. Available methods are [nll, miscalibration_area, ence, spearman] or any available classification or multiclass metric.",
+    # )
+    # unc_args.add_argument(
+    #     "--evaluation-scores-path",
+    #     help="Location to save the results of uncertainty evaluations.",
+    # )
+    # unc_args.add_argument(
+    #     "--uncertainty-dropout-p",
+    #     type=float,
+    #     default=0.1,
+    #     help="The probability to use for Monte Carlo dropout uncertainty estimation.",
+    # )
+    # unc_args.add_argument(
+    #     "--dropout-sampling-size",
+    #     type=int,
+    #     default=10,
+    #     help="The number of samples to use for Monte Carlo dropout uncertainty estimation. Distinct from the dropout used during training.",
+    # )
+    # unc_args.add_argument(
+    #     "--calibration-interval-percentile",
+    #     type=float,
+    #     default=95,
+    #     help="Sets the percentile used in the calibration methods. Must be in the range (1,100).",
+    # )
+    # unc_args.add_argument(
+    #     "--regression-calibrator-metric",
+    #     choices=['stdev', 'interval'],
+    #     help="Regression calibrators can output either a stdev or an inverval.",
+    # )
+    # unc_args.add_argument(
+    #     "--calibrationipath",
+    #     help="Path to data file to be used for uncertainty calibration.",
+    # )
+    # unc_args.add_argument(
+    #     "--calibration-features-path",
+    #     type=list[str],
+    #     help="Path to features data to be used with the uncertainty calibration dataset.",
+    # )
+    # unc_args.add_argument(
+    #     "--calibration-phase-features-path",
+    #     help=" ",
+    # )
+    # unc_args.add_argument(
+    #     "--calibration-atom-descriptors-path",
+    #     help="Path to the extra atom descriptors.",
+    # )
+    # unc_args.add_argument(
+    #     "--calibration-bond-descriptors-path",
+    #     help="Path to the extra bond descriptors that will be used as bond features to featurize a given molecule.",
+    # )
+
+    return parser
+
+
+def process_predict_args(args: Namespace) -> Namespace:
+    if args.test_path.suffix not in [".csv"]:
+        raise ArgumentError(
+            argument=None, message=f"Input data must be a CSV file. Got {args.test_path}"
+        )
+    if args.output is None:
+        args.output = args.test_path.parent / (args.test_path.stem + "_preds.csv")
+    if args.output.suffix not in [".csv", ".pkl"]:
+        raise ArgumentError(
+            argument=None, message=f"Output must be a CSV or Pickle file. Got {args.output}"
+        )
+    return args
+
+
+def find_models(model_path: Path):
+    if model_path.suffix in [".ckpt", ".pt"]:
+        return [model_path]
+    elif model_path.is_dir():
+        return list(model_path.rglob("*.ckpt")) + list(model_path.rglob("*.pt"))
+
+
+def make_prediction_for_model(
+    args: Namespace, model_path: Path, multicomponent: bool, output_path: Path
+):
+    model = load_model(model_path, multicomponent)
+
+    bounded = any(
+        isinstance(model.criterion, LossFunctionRegistry[loss_function])
+        for loss_function in LossFunctionRegistry.keys()
+        if "bounded" in loss_function
+    )
+
+    format_kwargs = dict(
+        no_header_row=args.no_header_row,
+        smiles_cols=args.smiles_columns,
+        rxn_cols=args.reaction_columns,
+        target_cols=None,
+        ignore_cols=None,
+        splits_col=None,
+        weight_col=None,
+        bounded=bounded,
+    )
+    featurization_kwargs = dict(
+        features_generators=args.features_generators, keep_h=args.keep_h, add_h=args.add_h
+    )
+
+    test_data = build_data_from_files(
+        args.test_path,
+        **format_kwargs,
+        p_descriptors=args.descriptors_path,
+        p_atom_feats=args.atom_features_path,
+        p_bond_feats=args.bond_features_path,
+        p_atom_descs=args.atom_descriptors_path,
+        **featurization_kwargs,
+    )
+    logger.info(f"test size: {len(test_data[0])}")
+    test_dsets = [
+        make_dataset(d, args.rxn_mode, args.multi_hot_atom_featurizer_mode) for d in test_data
+    ]
+
+    if multicomponent:
+        test_dset = data.MulticomponentDataset(test_dsets)
+    else:
+        test_dset = test_dsets[0]
+
+    # TODO: add uncertainty and calibration
+    # if args.cal_path is not None:
+    #     cal_data = build_data_from_files(
+    #         args.cal_path,
+    #         **format_kwargs,
+    #         target_columns=args.target_columns,
+    #         p_features=args.cal_features_path,
+    #         p_atom_feats=args.cal_atom_features_path,
+    #         p_bond_feats=args.cal_bond_features_path,
+    #         p_atom_descs=args.cal_atom_descriptors_path,
+    #         **featurization_kwargs,
+    #     )
+    #     logger.info(f"calibration size: {len(cal_data)}")
+    # else:
+    #     cal_data = None
+
+    test_loader = data.build_dataloader(test_dset, args.batch_size, args.num_workers, shuffle=False)
+    # TODO: add uncertainty and calibration
+    # if cal_data is not None:
+    #     cal_dset = make_dataset(cal_data, bond_messages, args.rxn_mode)
+    #     cal_loader = data.build_dataloader(cal_dset, args.batch_size, args.num_workers, shuffle=False)
+    # else:
+    #     cal_loader = None
+
+    logger.info(model)
+
+    trainer = pl.Trainer(
+        logger=False, enable_progress_bar=True, accelerator=args.accelerator, devices=args.devices
+    )
+
+    predss = trainer.predict(model, test_loader)
+
+    # TODO: add uncertainty and calibration
+    # if cal_dset is not None:
+    #     if args.task_type == "regression":
+    #         model.loc, model.scale = float(scaler.mean_), float(scaler.scale_)
+    #     predss_cal = trainer.predict(model, cal_loader)[0]
+
+    # TODO: might want to write a shared function for this as train.py might also want to do this.
+    df_test = pd.read_csv(args.test_path)
+    preds = torch.concat(predss, 0)
+
+    if isinstance(model.predictor, MulticlassClassificationFFN):
+        preds = torch.argmax(preds, dim=-1)
+
+    if args.target_columns is not None:
+        assert (
+            len(args.target_columns) == model.n_tasks
+        ), "Number of target columns must match the number of tasks."
+        target_columns = args.target_columns
+    else:
+        target_columns = [
+            f"pred_{i}" for i in range(preds.shape[1])
+        ]  # TODO: need to improve this for cases like multi-task MVE and multi-task multiclass
+
+    df_test[target_columns] = preds
+    if output_path.suffix == ".pkl":
+        df_test = df_test.reset_index(drop=True)
+        df_test.to_pickle(output_path)
+    else:
+        df_test.to_csv(output_path, index=False)
+    logger.info(f"Predictions saved to '{output_path}'")
+
+
+def main(args):
+    match (args.smiles_columns, args.reaction_columns):
+        case [None, None]:
+            n_components = 1
+        case [_, None]:
+            n_components = len(args.smiles_columns)
+        case [None, _]:
+            n_components = len(args.reaction_columns)
+        case _:
+            n_components = len(args.smiles_columns) + len(args.reaction_columns)
+
+    multicomponent = n_components > 1
+
+    model_paths = find_models(args.model_path)
+
+    for i, model_path in enumerate(model_paths):
+        logger.info(f"Predicting with model at '{model_path}'")
+        output_path = args.output.parent / Path(
+            str(args.output.stem) + f"_{i}" + str(args.output.suffix)
+        )
+        make_prediction_for_model(args, model_path, multicomponent, output_path)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser = PredictSubcommand.add_args(parser)
+
+    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, force=True)
+    args = parser.parse_args()
+    args = PredictSubcommand.func(args)

chemprop/cli/train.py

@@ -0,0 +1,1007 @@
+import json
+import logging
+import sys
+from copy import deepcopy
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+from configargparse import ArgumentError, ArgumentParser, Namespace
+from lightning import pytorch as pl
+from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint
+from lightning.pytorch.loggers import CSVLogger, TensorBoardLogger
+
+from chemprop.cli.common import add_common_args, process_common_args, validate_common_args
+from chemprop.cli.conf import NOW
+from chemprop.cli.utils import (
+    LookupAction,
+    Subcommand,
+    build_data_from_files,
+    get_column_names,
+    make_dataset,
+    parse_indices,
+)
+from chemprop.cli.utils.args import uppercase
+from chemprop.data import (
+    MoleculeDataset,
+    MolGraphDataset,
+    MulticomponentDataset,
+    ReactionDatapoint,
+    SplitType,
+    build_dataloader,
+    make_split_indices,
+    split_data_by_indices,
+)
+from chemprop.featurizers import MoleculeFeaturizerRegistry
+from chemprop.models import MPNN, MulticomponentMPNN, save_model
+from chemprop.nn import AggregationRegistry, LossFunctionRegistry, MetricRegistry, PredictorRegistry
+from chemprop.nn.message_passing import (
+    AtomMessagePassing,
+    BondMessagePassing,
+    MulticomponentMessagePassing,
+)
+from chemprop.nn.transforms import GraphTransform, ScaleTransform, UnscaleTransform
+from chemprop.nn.utils import Activation
+from chemprop.utils import Factory
+
+logger = logging.getLogger(__name__)
+
+
+class TrainSubcommand(Subcommand):
+    COMMAND = "train"
+    HELP = "train a chemprop model"
+    parser = None
+
+    @classmethod
+    def add_args(cls, parser: ArgumentParser) -> ArgumentParser:
+        parser = add_common_args(parser)
+        parser = add_train_args(parser)
+        cls.parser = parser
+        return parser
+
+    @classmethod
+    def func(cls, args: Namespace):
+        args = process_common_args(args)
+        validate_common_args(args)
+        args = process_train_args(args)
+        validate_train_args(args)
+
+        args.output_dir.mkdir(exist_ok=True, parents=True)
+        save_config(cls.parser, args)
+        main(args)
+
+
+def add_train_args(parser: ArgumentParser) -> ArgumentParser:
+    parser.add_argument(
+        "--config-path",
+        type=Path,
+        is_config_file=True,
+        help="Path to a configuration file. Command line arguments override values in the configuration file.",
+    )
+    parser.add_argument(
+        "-i",
+        "--data-path",
+        type=Path,
+        help="Path to an input CSV file containing SMILES and the associated target values.",
+    )
+    parser.add_argument(
+        "-o",
+        "--output-dir",
+        "--save-dir",
+        type=Path,
+        help="Directory where training outputs will be saved. Defaults to 'CURRENT_DIRECTORY/chemprop_training/STEM_OF_INPUT/TIME_STAMP'.",
+    )
+
+    # TODO: Add in v2.1
+    # parser.add_argument(
+    #     "--checkpoint-dir",
+    #     help="Directory from which to load model checkpoints (walks directory and ensembles all models that are found).",
+    # )
+    # parser.add_argument("--checkpoint-path", help="Path to model checkpoint (:code:`.pt` file).")
+    # parser.add_argument(
+    #     "--checkpoint-paths",
+    #     type=list[str],
+    #     help="List of paths to model checkpoints (:code:`.pt` files).",
+    # )
+    # # TODO: Is this a prediction only argument?
+    # parser.add_argument(
+    #     "--checkpoint",
+    #     help="Location of checkpoint(s) to use for ... If the location is a directory, chemprop walks it and ensembles all models that are found. If the location is a path or list of paths to model checkpoints (:code:`.pt` files), only those models will be loaded.",
+    # )
+
+    # TODO: Add in v2.1; see if we can tell lightning how often to log training loss
+    # parser.add_argument(
+    #     "--log-frequency",
+    #     type=int,
+    #     default=10,
+    #     help="The number of batches between each logging of the training loss.",
+    # )
+
+    transfer_args = parser.add_argument_group("transfer learning args")
+    transfer_args.add_argument(
+        "--model-frzn",
+        help="Path to model checkpoint file to be loaded for overwriting and freezing weights.",
+    )
+    transfer_args.add_argument(
+        "--frzn-ffn-layers",
+        type=int,
+        default=0,
+        help="Overwrites weights for the first n layers of the ffn from checkpoint model (specified checkpoint_frzn), where n is specified in the input. Automatically also freezes mpnn weights.",
+    )
+    # transfer_args.add_argument(
+    #     "--freeze-first-only",
+    #     action="store_true",
+    #     help="Determines whether or not to use checkpoint_frzn for just the first encoder. Default (False) is to use the checkpoint to freeze all encoders. (only relevant for number_of_molecules > 1, where checkpoint model has number_of_molecules = 1)",
+    # )
+
+    # TODO: Add in v2.1
+    # parser.add_argument(
+    #     "--resume-experiment",
+    #     action="store_true",
+    #     help="Whether to resume the experiment. Loads test results from any folds that have already been completed and skips training those folds.",
+    # )
+    # parser.add_argument(
+    #     "--config-path",
+    #     help="Path to a :code:`.json` file containing arguments. Any arguments present in the config file will override arguments specified via the command line or by the defaults.",
+    # )
+    parser.add_argument(
+        "--ensemble-size",
+        type=int,
+        default=1,
+        help="Number of models in ensemble for each splitting of data.",
+    )
+
+    # TODO: Add in v2.2
+    # abt_args = parser.add_argument_group("atom/bond target args")
+    # abt_args.add_argument(
+    #     "--is-atom-bond-targets",
+    #     action="store_true",
+    #     help="Whether this is atomic/bond properties prediction.",
+    # )
+    # abt_args.add_argument(
+    #     "--no-adding-bond-types",
+    #     action="store_true",
+    #     help="Whether the bond types determined by RDKit molecules added to the output of bond targets. This option is intended to be used with the :code:`is_atom_bond_targets`.",
+    # )
+    # abt_args.add_argument(
+    #     "--keeping-atom-map",
+    #     action="store_true",
+    #     help="Whether RDKit molecules keep the original atom mapping. This option is intended to be used when providing atom-mapped SMILES with the :code:`is_atom_bond_targets`.",
+    # )
+    # abt_args.add_argument(
+    #     "--no-shared-atom-bond-ffn",
+    #     action="store_true",
+    #     help="Whether the FFN weights for atom and bond targets should be independent between tasks.",
+    # )
+    # abt_args.add_argument(
+    #     "--weights-ffn-num-layers",
+    #     type=int,
+    #     default=2,
+    #     help="Number of layers in FFN for determining weights used in constrained targets.",
+    # )
+
+    mp_args = parser.add_argument_group("message passing")
+    mp_args.add_argument(
+        "--message-hidden-dim", type=int, default=300, help="hidden dimension of the messages"
+    )
+    mp_args.add_argument(
+        "--message-bias", action="store_true", help="add bias to the message passing layers"
+    )
+    mp_args.add_argument("--depth", type=int, default=3, help="Number of message passing steps.")
+    mp_args.add_argument(
+        "--undirected",
+        action="store_true",
+        help="Pass messages on undirected bonds/edges (always sum the two relevant bond vectors).",
+    )
+    mp_args.add_argument(
+        "--dropout",
+        type=float,
+        default=0.0,
+        help="dropout probability in message passing/FFN layers",
+    )
+    mp_args.add_argument(
+        "--mpn-shared",
+        action="store_true",
+        help="Whether to use the same message passing neural network for all input molecules. Only relevant if :code:`number_of_molecules > 1`",
+    )
+    mp_args.add_argument(
+        "--activation",
+        type=uppercase,
+        default="RELU",
+        choices=list(Activation.keys()),
+        help="activation function in message passing/FFN layers",
+    )
+    mp_args.add_argument(
+        "--aggregation",
+        "--agg",
+        default="mean",
+        action=LookupAction(AggregationRegistry),
+        help="the aggregation mode to use during graph predictor",
+    )
+    mp_args.add_argument(
+        "--aggregation-norm",
+        type=float,
+        default=100,
+        help="normalization factor by which to divide summed up atomic features for 'norm' aggregation",
+    )
+    mp_args.add_argument(
+        "--atom-messages", action="store_true", help="pass messages on atoms rather than bonds"
+    )
+
+    # TODO: Add in v2.1
+    # mpsolv_args = parser.add_argument_group("message passing with solvent")
+    # mpsolv_args.add_argument(
+    #     "--reaction-solvent",
+    #     action="store_true",
+    #     help="Whether to adjust the MPNN layer to take as input a reaction and a molecule, and to encode them with separate MPNNs.",
+    # )
+    # mpsolv_args.add_argument(
+    #     "--bias-solvent",
+    #     action="store_true",
+    #     help="Whether to add bias to linear layers for solvent MPN if :code:`reaction_solvent` is True.",
+    # )
+    # mpsolv_args.add_argument(
+    #     "--hidden-size-solvent",
+    #     type=int,
+    #     default=300,
+    #     help="Dimensionality of hidden layers in solvent MPN if :code:`reaction_solvent` is True.",
+    # )
+    # mpsolv_args.add_argument(
+    #     "--depth-solvent",
+    #     type=int,
+    #     default=3,
+    #     help="Number of message passing steps for solvent if :code:`reaction_solvent` is True.",
+    # )
+
+    ffn_args = parser.add_argument_group("FFN args")
+    ffn_args.add_argument(
+        "--ffn-hidden-dim", type=int, default=300, help="hidden dimension in the FFN top model"
+    )
+    ffn_args.add_argument(  # TODO: the default in v1 was 2. (see weights_ffn_num_layers option) Do we really want the default to now be 1?
+        "--ffn-num-layers", type=int, default=1, help="number of layers in FFN top model"
+    )
+    # TODO: Decide if we want to implment this in v2
+    # ffn_args.add_argument(
+    #     "--features-only",
+    #     action="store_true",
+    #     help="Use only the additional features in an FFN, no graph network.",
+    # )
+
+    extra_mpnn_args = parser.add_argument_group("extra MPNN args")
+    extra_mpnn_args.add_argument(
+        "--no-batch-norm",
+        action="store_true",
+        help="Don't use batch normalization after aggregation.",
+    )
+    extra_mpnn_args.add_argument(
+        "--multiclass-num-classes",
+        type=int,
+        default=3,
+        help="Number of classes when running multiclass classification.",
+    )
+    # TODO: Add in v2.1
+    # extra_mpnn_args.add_argument(
+    #     "--spectral-activation",
+    #     default="exp",
+    #     choices=["softplus", "exp"],
+    #     help="Indicates which function to use in task_type spectra training to constrain outputs to be positive.",
+    # )
+
+    train_data_args = parser.add_argument_group("training input data args")
+    train_data_args.add_argument(
+        "-w",
+        "--weight-column",
+        help="the name of the column in the input CSV containg individual data weights",
+    )
+    train_data_args.add_argument(
+        "--target-columns",
+        nargs="+",
+        help="Name of the columns containing target values. By default, uses all columns except the SMILES column and the :code:`ignore_columns`.",
+    )
+    train_data_args.add_argument(
+        "--ignore-columns",
+        nargs="+",
+        help="Name of the columns to ignore when :code:`target_columns` is not provided.",
+    )
+    # TODO: Add in v2.1
+    # train_data_args.add_argument(
+    #     "--spectra-phase-mask-path",
+    #     help="Path to a file containing a phase mask array, used for excluding particular regions in spectra predictions.",
+    # )
+
+    train_args = parser.add_argument_group("training args")
+    train_args.add_argument(
+        "-t",
+        "--task-type",
+        default="regression",
+        action=LookupAction(PredictorRegistry),
+        help="Type of dataset. This determines the default loss function used during training. Defaults to regression.",
+    )
+    train_args.add_argument(
+        "-l",
+        "--loss-function",
+        action=LookupAction(LossFunctionRegistry),
+        help="Loss function to use during training. If not specified, will use the default loss function for the given task type (see documentation).",
+    )
+    train_args.add_argument(
+        "--v-kl",
+        "--evidential-regularization",
+        type=float,
+        default=0.0,
+        help="Value used in regularization for evidential loss function. The default value recommended by Soleimany et al.(2021) is 0.2. Optimal value is dataset-dependent; it is recommended that users test different values to find the best value for their model.",
+    )
+
+    train_args.add_argument(
+        "--eps", type=float, default=1e-8, help="evidential regularization epsilon"
+    )
+    # TODO: Add in v2.1
+    # train_args.add_argument(  # TODO: Is threshold the same thing as the spectra target floor? I'm not sure but combined them.
+    #     "-T",
+    #     "--threshold",
+    #     "--spectra-target-floor",
+    #     type=float,
+    #     default=1e-8,
+    #     help="spectral threshold limit. v1 help string: Values in targets for dataset type spectra are replaced with this value, intended to be a small positive number used to enforce positive values.",
+    # )
+    train_args.add_argument(
+        "--metrics",
+        "--metric",
+        nargs="+",
+        action=LookupAction(MetricRegistry),
+        help="evaluation metrics. If unspecified, will use the following metrics for given dataset types: regression->rmse, classification->roc, multiclass->ce ('cross entropy'), spectral->sid. If multiple metrics are provided, the 0th one will be used for early stopping and checkpointing",
+    )
+    # TODO: Add in v2.1
+    # train_args.add_argument(
+    #     "--show-individual-scores",
+    #     action="store_true",
+    #     help="Show all scores for individual targets, not just average, at the end.",
+    # )
+    train_args.add_argument(
+        "--task-weights",
+        nargs="+",
+        type=float,
+        help="the weight to apply to an individual task in the overall loss",
+    )
+    train_args.add_argument(
+        "--warmup-epochs",
+        type=int,
+        default=2,
+        help="Number of epochs during which learning rate increases linearly from :code:`init_lr` to :code:`max_lr`. Afterwards, learning rate decreases exponentially from :code:`max_lr` to :code:`final_lr`.",
+    )
+
+    train_args.add_argument("--init-lr", type=float, default=1e-4, help="Initial learning rate.")
+    train_args.add_argument("--max-lr", type=float, default=1e-3, help="Maximum learning rate.")
+    train_args.add_argument("--final-lr", type=float, default=1e-4, help="Final learning rate.")
+    train_args.add_argument(
+        "--epochs", type=int, default=50, help="the number of epochs to train over"
+    )
+    train_args.add_argument(
+        "--patience",
+        type=int,
+        default=None,
+        help="Number of epochs to wait for improvement before early stopping.",
+    )
+    train_args.add_argument(
+        "--grad-clip",
+        type=float,
+        help="Passed directly to the lightning trainer which controls grad clipping. See the :code:`Trainer()` docstring for details.",
+    )
+    # TODO: Add in v2.1
+    # train_args.add_argument(
+    #     "--class-balance",
+    #     action="store_true",
+    #     help="Trains with an equal number of positives and negatives in each batch.",
+    # )
+
+    split_args = parser.add_argument_group("split args")
+    split_args.add_argument(
+        "--split",
+        "--split-type",
+        type=uppercase,
+        default="RANDOM",
+        choices=list(SplitType.keys()),
+        help="Method of splitting the data into train/val/test (case insensitive).",
+    )
+    split_args.add_argument(
+        "--split-sizes",
+        type=float,
+        nargs=3,
+        default=[0.8, 0.1, 0.1],
+        help="Split proportions for train/validation/test sets.",
+    )
+    split_args.add_argument(
+        "--split-key-molecule",
+        type=int,
+        default=0,
+        help="The index of the key molecule used for splitting when multiple molecules are present and constrained split_type is used (e.g., 'scaffold_balanced' or 'random_with_repeated_smiles'). Note that this index begins with zero for the first molecule.",
+    )
+    split_args.add_argument(
+        "-k",
+        "--num-folds",
+        type=int,
+        default=1,
+        help="Number of folds when performing cross validation.",
+    )
+    split_args.add_argument(
+        "--save-smiles-splits",
+        action="store_true",
+        help="Save smiles for each train/val/test splits for prediction convenience later.",
+    )
+    split_args.add_argument(
+        "--splits-file",
+        type=Path,
+        help="Path to a JSON file containing pre-defined splits for the input data, formatted as a list of dictionaries with keys 'train', 'val', and 'test' and values as lists of indices or strings formatted like '0-2,4'. See documentation for more details.",
+    )
+    train_data_args.add_argument(
+        "--splits-column",
+        help="Name of the column in the input CSV file containing 'train', 'val', or 'test' for each row.",
+    )
+    split_args.add_argument(
+        "--data-seed",
+        type=int,
+        default=0,
+        help="Random seed to use when splitting data into train/val/test sets. When :code`num_folds > 1`, the first fold uses this seed and all subsequent folds add 1 to the seed. Also used for shuffling data in :code:`build_dataloader` when :code:`shuffle` is True.",
+    )
+
+    parser.add_argument(
+        "--pytorch-seed",
+        type=int,
+        default=None,
+        help="Seed for PyTorch randomness (e.g., random initial weights).",
+    )
+
+    return parser
+
+
+def process_train_args(args: Namespace) -> Namespace:
+    if args.config_path is None and args.data_path is None:
+        raise ArgumentError(argument=None, message="Data path must be provided for training.")
+
+    if args.data_path.suffix not in [".csv"]:
+        raise ArgumentError(
+            argument=None, message=f"Input data must be a CSV file. Got {args.data_path}"
+        )
+    if args.output_dir is None:
+        args.output_dir = Path(f"chemprop_training/{args.data_path.stem}/{NOW}")
+
+    return args
+
+
+def validate_train_args(args):
+    pass
+
+
+def normalize_inputs(train_dset, val_dset, args):
+    multicomponent = isinstance(train_dset, MulticomponentDataset)
+    num_components = train_dset.n_components if multicomponent else 1
+
+    X_d_transform = None
+    V_f_transforms = [nn.Identity()] * num_components
+    E_f_transforms = [nn.Identity()] * num_components
+    V_d_transforms = [None] * num_components
+    graph_transforms = []
+
+    d_xd = train_dset.d_xd
+    d_vf = train_dset.d_vf
+    d_ef = train_dset.d_ef
+    d_vd = train_dset.d_vd
+
+    if d_xd > 0 and not args.no_descriptor_scaling:
+        scaler = train_dset.normalize_inputs("X_d")
+        val_dset.normalize_inputs("X_d", scaler)
+
+        scaler = scaler if not isinstance(scaler, list) else scaler[0]
+
+        if scaler is not None:
+            logger.info(
+                f"Descriptors: loc = {np.array2string(scaler.mean_, precision=3)}, scale = {np.array2string(scaler.scale_, precision=3)}"
+            )
+            X_d_transform = ScaleTransform.from_standard_scaler(scaler)
+
+    if d_vf > 0 and not args.no_atom_feature_scaling:
+        scaler = train_dset.normalize_inputs("V_f")
+        val_dset.normalize_inputs("V_f", scaler)
+
+        scalers = [scaler] if not isinstance(scaler, list) else scaler
+
+        for i, scaler in enumerate(scalers):
+            if scaler is None:
+                continue
+
+            logger.info(
+                f"Atom features for mol {i}: loc = {np.array2string(scaler.mean_, precision=3)}, scale = {np.array2string(scaler.scale_, precision=3)}"
+            )
+            featurizer = (
+                train_dset.datasets[i].featurizer if multicomponent else train_dset.featurizer
+            )
+            V_f_transforms[i] = ScaleTransform.from_standard_scaler(
+                scaler, pad=featurizer.atom_fdim - featurizer.extra_atom_fdim
+            )
+
+    if d_ef > 0 and not args.no_bond_feature_scaling:
+        scaler = train_dset.normalize_inputs("E_f")
+        val_dset.normalize_inputs("E_f", scaler)
+
+        scalers = [scaler] if not isinstance(scaler, list) else scaler
+
+        for i, scaler in enumerate(scalers):
+            if scaler is None:
+                continue
+
+            logger.info(
+                f"Bond features for mol {i}: loc = {np.array2string(scaler.mean_, precision=3)}, scale = {np.array2string(scaler.scale_, precision=3)}"
+            )
+            featurizer = (
+                train_dset.datasets[i].featurizer if multicomponent else train_dset.featurizer
+            )
+            E_f_transforms[i] = ScaleTransform.from_standard_scaler(
+                scaler, pad=featurizer.bond_fdim - featurizer.extra_bond_fdim
+            )
+
+    for V_f_transform, E_f_transform in zip(V_f_transforms, E_f_transforms):
+        graph_transforms.append(GraphTransform(V_f_transform, E_f_transform))
+
+    if d_vd > 0 and not args.no_atom_descriptor_scaling:
+        scaler = train_dset.normalize_inputs("V_d")
+        val_dset.normalize_inputs("V_d", scaler)
+
+        scalers = [scaler] if not isinstance(scaler, list) else scaler
+
+        for i, scaler in enumerate(scalers):
+            if scaler is None:
+                continue
+
+            logger.info(
+                f"Atom descriptors for mol {i}: loc = {np.array2string(scaler.mean_, precision=3)}, scale = {np.array2string(scaler.scale_, precision=3)}"
+            )
+            V_d_transforms[i] = ScaleTransform.from_standard_scaler(scaler)
+
+    return X_d_transform, graph_transforms, V_d_transforms
+
+
+def save_config(parser: ArgumentParser, args: Namespace):
+    config_args = deepcopy(args)
+    for key, value in vars(config_args).items():
+        if isinstance(value, Path):
+            setattr(config_args, key, str(value))
+
+    for key in ["atom_features_path", "atom_descriptors_path", "bond_features_path"]:
+        if getattr(config_args, key) is not None:
+            for index, path in getattr(config_args, key).items():
+                getattr(config_args, key)[index] = str(path)
+
+    config_path = str(args.output_dir / "config.toml")
+    parser.write_config_file(parsed_namespace=config_args, output_file_paths=[config_path])
+
+
+def save_smiles_splits(args: Namespace, output_dir, train_dset, val_dset, test_dset):
+    train_smis = train_dset.smiles
+    df_train = pd.DataFrame(train_smis, columns=args.smiles_columns)
+    df_train.to_csv(output_dir / "train_smiles.csv", index=False)
+
+    val_smis = val_dset.smiles
+    df_val = pd.DataFrame(val_smis, columns=args.smiles_columns)
+    df_val.to_csv(output_dir / "val_smiles.csv", index=False)
+
+    if test_dset is not None:
+        test_smis = test_dset.smiles
+        df_test = pd.DataFrame(test_smis, columns=args.smiles_columns)
+        df_test.to_csv(output_dir / "test_smiles.csv", index=False)
+
+
+def build_splits(args, format_kwargs, featurization_kwargs):
+    """build the train/val/test splits"""
+    logger.info(f"Pulling data from file: {args.data_path}")
+    all_data = build_data_from_files(
+        args.data_path,
+        p_descriptors=args.descriptors_path,
+        p_atom_feats=args.atom_features_path,
+        p_bond_feats=args.bond_features_path,
+        p_atom_descs=args.atom_descriptors_path,
+        **format_kwargs,
+        **featurization_kwargs,
+    )
+
+    if args.splits_column is not None:
+        df = pd.read_csv(
+            args.data_path, header=None if args.no_header_row else "infer", index_col=False
+        )
+        grouped = df.groupby(df[args.splits_column].str.lower())
+        train_indices = grouped.groups.get("train", pd.Index([])).tolist()
+        val_indices = grouped.groups.get("val", pd.Index([])).tolist()
+        test_indices = grouped.groups.get("test", pd.Index([])).tolist()
+        train_indices, val_indices, test_indices = [train_indices], [val_indices], [test_indices]
+
+    elif args.splits_file is not None:
+        with open(args.splits_file, "rb") as json_file:
+            split_idxss = json.load(json_file)
+        train_indices = [parse_indices(d["train"]) for d in split_idxss]
+        val_indices = [parse_indices(d["val"]) for d in split_idxss]
+        test_indices = [parse_indices(d["test"]) for d in split_idxss]
+
+    else:
+        splitting_data = all_data[args.split_key_molecule]
+        if isinstance(splitting_data[0], ReactionDatapoint):
+            splitting_mols = [datapoint.rct for datapoint in splitting_data]
+        else:
+            splitting_mols = [datapoint.mol for datapoint in splitting_data]
+        train_indices, val_indices, test_indices = make_split_indices(
+            splitting_mols, args.split, args.split_sizes, args.data_seed, args.num_folds
+        )
+        if not (
+            SplitType.get(args.split) == SplitType.CV_NO_VAL
+            or SplitType.get(args.split) == SplitType.CV
+        ):
+            train_indices, val_indices, test_indices = (
+                [train_indices],
+                [val_indices],
+                [test_indices],
+            )
+
+    train_data, val_data, test_data = split_data_by_indices(
+        all_data, train_indices, val_indices, test_indices
+    )
+    for i_split in range(len(train_data)):
+        sizes = [len(train_data[i_split][0]), len(val_data[i_split][0]), len(test_data[i_split][0])]
+        logger.info(f"train/val/test split_{i_split} sizes: {sizes}")
+
+    return train_data, val_data, test_data
+
+
+def build_datasets(args, train_data, val_data, test_data):
+    """build the train/val/test datasets, where :attr:`test_data` may be None"""
+    multicomponent = len(train_data) > 1
+    if multicomponent:
+        train_dsets = [
+            make_dataset(data, args.rxn_mode, args.multi_hot_atom_featurizer_mode)
+            for data in train_data
+        ]
+        val_dsets = [
+            make_dataset(data, args.rxn_mode, args.multi_hot_atom_featurizer_mode)
+            for data in val_data
+        ]
+        train_dset = MulticomponentDataset(train_dsets)
+        val_dset = MulticomponentDataset(val_dsets)
+        if len(test_data[0]) > 0:
+            test_dsets = [
+                make_dataset(data, args.rxn_mode, args.multi_hot_atom_featurizer_mode)
+                for data in test_data
+            ]
+            test_dset = MulticomponentDataset(test_dsets)
+        else:
+            test_dset = None
+    else:
+        train_data = train_data[0]
+        val_data = val_data[0]
+        test_data = test_data[0]
+
+        train_dset = make_dataset(train_data, args.rxn_mode, args.multi_hot_atom_featurizer_mode)
+        val_dset = make_dataset(val_data, args.rxn_mode, args.multi_hot_atom_featurizer_mode)
+        if len(test_data) > 0:
+            test_dset = make_dataset(test_data, args.rxn_mode, args.multi_hot_atom_featurizer_mode)
+        else:
+            test_dset = None
+
+    return train_dset, val_dset, test_dset
+
+
+def build_model(
+    args,
+    train_dset: MolGraphDataset | MulticomponentDataset,
+    output_transform: UnscaleTransform,
+    input_transforms: tuple[ScaleTransform, list[GraphTransform], list[ScaleTransform]],
+) -> MPNN:
+    mp_cls = AtomMessagePassing if args.atom_messages else BondMessagePassing
+
+    X_d_transform, graph_transforms, V_d_transforms = input_transforms
+
+    if isinstance(train_dset, MulticomponentDataset):
+        mp_blocks = [
+            mp_cls(
+                train_dset.datasets[i].featurizer.atom_fdim,
+                train_dset.datasets[i].featurizer.bond_fdim,
+                d_h=args.message_hidden_dim,
+                d_vd=(
+                    train_dset.datasets[i].d_vd
+                    if isinstance(train_dset.datasets[i], MoleculeDataset)
+                    else 0
+                ),
+                bias=args.message_bias,
+                depth=args.depth,
+                undirected=args.undirected,
+                dropout=args.dropout,
+                activation=args.activation,
+                V_d_transform=V_d_transforms[i],
+                graph_transform=graph_transforms[i],
+            )
+            for i in range(train_dset.n_components)
+        ]
+        if args.mpn_shared:
+            if args.reaction_columns is not None and args.smiles_columns is not None:
+                raise ArgumentError(
+                    argument=None,
+                    message="Cannot use shared MPNN with both molecule and reaction data.",
+                )
+
+        mp_block = MulticomponentMessagePassing(mp_blocks, train_dset.n_components, args.mpn_shared)
+        # NOTE(degraff): this if/else block should be handled by the init of MulticomponentMessagePassing
+        # if args.mpn_shared:
+        #     mp_block = MulticomponentMessagePassing(mp_blocks[0], n_components, args.mpn_shared)
+        # else:
+        d_xd = train_dset.datasets[0].d_xd
+        n_tasks = train_dset.datasets[0].Y.shape[1]
+        mpnn_cls = MulticomponentMPNN
+    else:
+        mp_block = mp_cls(
+            train_dset.featurizer.atom_fdim,
+            train_dset.featurizer.bond_fdim,
+            d_h=args.message_hidden_dim,
+            d_vd=train_dset.d_vd if isinstance(train_dset, MoleculeDataset) else 0,
+            bias=args.message_bias,
+            depth=args.depth,
+            undirected=args.undirected,
+            dropout=args.dropout,
+            activation=args.activation,
+            V_d_transform=V_d_transforms[0],
+            graph_transform=graph_transforms[0],
+        )
+        d_xd = train_dset.d_xd
+        n_tasks = train_dset.Y.shape[1]
+        mpnn_cls = MPNN
+
+    agg = Factory.build(AggregationRegistry[args.aggregation], norm=args.aggregation_norm)
+    predictor_cls = PredictorRegistry[args.task_type]
+    if args.loss_function is not None:
+        criterion = Factory.build(
+            LossFunctionRegistry[args.loss_function],
+            task_weights=args.task_weights,
+            v_kl=args.v_kl,
+            # threshold=args.threshold, TODO: Add in v2.1
+            eps=args.eps,
+        )
+    else:
+        criterion = None
+    if args.metrics is not None:
+        metrics = [Factory.build(MetricRegistry[metric]) for metric in args.metrics]
+    else:
+        metrics = None
+
+    predictor = Factory.build(
+        predictor_cls,
+        input_dim=mp_block.output_dim + d_xd,
+        n_tasks=n_tasks,
+        hidden_dim=args.ffn_hidden_dim,
+        n_layers=args.ffn_num_layers,
+        dropout=args.dropout,
+        activation=args.activation,
+        criterion=criterion,
+        n_classes=args.multiclass_num_classes,
+        output_transform=output_transform,
+        # spectral_activation=args.spectral_activation, TODO: Add in v2.1
+    )
+
+    if args.loss_function is None:
+        logger.info(
+            f"No loss function was specified! Using class default: {predictor_cls._T_default_criterion}"
+        )
+
+    if args.model_frzn is not None:
+        model = mpnn_cls.load_from_file(args.model_frzn)
+        model.message_passing.apply(lambda module: module.requires_grad_(False))
+        model.message_passing.apply(
+            lambda m: setattr(m, "p", 0.0) if isinstance(m, torch.nn.Dropout) else None
+        )
+        model.bn.apply(lambda module: module.requires_grad_(False))
+        for idx in range(args.frzn_ffn_layers):
+            model.predictor.ffn[idx].requires_grad_(False)
+            setattr(model.predictor.ffn[idx + 1][1], "p", 0.0)
+
+        return model
+
+    return mpnn_cls(
+        mp_block,
+        agg,
+        predictor,
+        not args.no_batch_norm,
+        metrics,
+        args.warmup_epochs,
+        args.init_lr,
+        args.max_lr,
+        args.final_lr,
+        X_d_transform=X_d_transform,
+    )
+
+
+def train_model(
+    args, train_loader, val_loader, test_loader, output_dir, output_transform, input_transforms
+):
+    for model_idx in range(args.ensemble_size):
+        model_output_dir = output_dir / f"model_{model_idx}"
+        model_output_dir.mkdir(exist_ok=True, parents=True)
+
+        if args.pytorch_seed is None:
+            seed = torch.seed()
+            deterministic = False
+        else:
+            seed = args.pytorch_seed + model_idx
+            deterministic = True
+
+        torch.manual_seed(seed)
+
+        model = build_model(args, train_loader.dataset, output_transform, input_transforms)
+        logger.info(model)
+
+        monitor_mode = "min" if model.metrics[0].minimize else "max"
+        logger.debug(f"Evaluation metric: '{model.metrics[0].alias}', mode: '{monitor_mode}'")
+
+        try:
+            trainer_logger = TensorBoardLogger(model_output_dir, "trainer_logs")
+        except ModuleNotFoundError:
+            trainer_logger = CSVLogger(model_output_dir, "trainer_logs")
+
+        checkpointing = ModelCheckpoint(
+            model_output_dir / "checkpoints",
+            "best-{epoch}-{val_loss:.2f}",
+            "val_loss",
+            mode=monitor_mode,
+            save_last=True,
+        )
+
+        patience = args.patience if args.patience is not None else args.epochs
+        early_stopping = EarlyStopping("val_loss", patience=patience, mode=monitor_mode)
+
+        trainer = pl.Trainer(
+            logger=trainer_logger,
+            enable_progress_bar=True,
+            accelerator=args.accelerator,
+            devices=args.devices,
+            max_epochs=args.epochs,
+            callbacks=[checkpointing, early_stopping],
+            gradient_clip_val=args.grad_clip,
+            deterministic=deterministic,
+        )
+        trainer.fit(model, train_loader, val_loader)
+
+        if test_loader is not None:
+            predss = trainer.predict(dataloaders=test_loader)
+            preds = torch.concat(predss, 0).numpy()
+
+            if isinstance(test_loader.dataset, MulticomponentDataset):
+                test_dset = test_loader.dataset.datasets[0]
+            else:
+                test_dset = test_loader.dataset
+            targets = test_dset.Y
+            mask = torch.from_numpy(np.isfinite(targets))
+            targets = np.nan_to_num(targets, nan=0.0)
+            weights = torch.from_numpy(test_dset.weights)
+            lt_mask = (
+                torch.from_numpy(test_dset.lt_mask) if test_dset.lt_mask[0] is not None else None
+            )
+            gt_mask = (
+                torch.from_numpy(test_dset.gt_mask) if test_dset.gt_mask[0] is not None else None
+            )
+            preds_losses = [
+                metric(
+                    torch.from_numpy(preds),
+                    torch.from_numpy(targets),
+                    mask,
+                    weights,
+                    lt_mask,
+                    gt_mask,
+                )
+                for metric in model.metrics
+            ]
+            preds_metrics = {
+                f"entire_test/{m.alias}": l.item() for m, l in zip(model.metrics, preds_losses)
+            }
+            print(f"Entire Test Set results: {preds_metrics}")
+
+            columns = get_column_names(
+                args.data_path,
+                args.smiles_columns,
+                args.reaction_columns,
+                args.target_columns,
+                args.ignore_columns,
+                args.splits_column,
+                args.weight_column,
+                args.no_header_row,
+            )
+            names = test_loader.dataset.names
+            if isinstance(test_loader.dataset, MulticomponentDataset):
+                namess = list(zip(*names))
+            else:
+                namess = [names]
+            if "multiclass" in args.task_type:
+                df_preds = pd.DataFrame(list(zip(*namess, preds)), columns=columns)
+            else:
+                df_preds = pd.DataFrame(list(zip(*namess, *preds.T)), columns=columns)
+            df_preds.to_csv(model_output_dir / "test_predictions.csv", index=False)
+
+        best_model_path = checkpointing.best_model_path
+        model = model.__class__.load_from_checkpoint(best_model_path)
+        p_model = model_output_dir / "best.pt"
+        save_model(p_model, model)
+        logger.info(f"Best model saved to '{p_model}'")
+
+
+def main(args):
+    format_kwargs = dict(
+        no_header_row=args.no_header_row,
+        smiles_cols=args.smiles_columns,
+        rxn_cols=args.reaction_columns,
+        target_cols=args.target_columns,
+        ignore_cols=args.ignore_columns,
+        splits_col=args.splits_column,
+        weight_col=args.weight_column,
+        bounded=args.loss_function is not None and "bounded" in args.loss_function,
+    )
+    if args.features_generators is not None:
+        # TODO: MorganFeaturizers take radius, length, and include_chirality as arguements. Should we expose these through the CLI?
+        features_generators = [
+            Factory.build(MoleculeFeaturizerRegistry[features_generator])
+            for features_generator in args.features_generators
+        ]
+    else:
+        features_generators = None
+
+    featurization_kwargs = dict(
+        features_generators=features_generators, keep_h=args.keep_h, add_h=args.add_h
+    )
+
+    splits = build_splits(args, format_kwargs, featurization_kwargs)
+
+    for fold_idx, (train_data, val_data, test_data) in enumerate(zip(*splits)):
+        if args.num_folds == 1:
+            output_dir = args.output_dir
+        else:
+            output_dir = args.output_dir / f"fold_{fold_idx}"
+
+        output_dir.mkdir(exist_ok=True, parents=True)
+
+        train_dset, val_dset, test_dset = build_datasets(args, train_data, val_data, test_data)
+
+        input_transforms = normalize_inputs(train_dset, val_dset, args)
+
+        if args.save_smiles_splits:
+            save_smiles_splits(args, output_dir, train_dset, val_dset, test_dset)
+
+        if "regression" in args.task_type:
+            output_scaler = train_dset.normalize_targets()
+            val_dset.normalize_targets(output_scaler)
+            logger.info(f"Train data: mean = {output_scaler.mean_} | std = {output_scaler.scale_}")
+            output_transform = UnscaleTransform.from_standard_scaler(output_scaler)
+        else:
+            output_transform = None
+
+        train_loader = build_dataloader(
+            train_dset, args.batch_size, args.num_workers, seed=args.data_seed
+        )
+        val_loader = build_dataloader(val_dset, args.batch_size, args.num_workers, shuffle=False)
+        if test_dset is not None:
+            test_loader = build_dataloader(
+                test_dset, args.batch_size, args.num_workers, shuffle=False
+            )
+        else:
+            test_loader = None
+
+        train_model(
+            args,
+            train_loader,
+            val_loader,
+            test_loader,
+            output_dir,
+            output_transform,
+            input_transforms,
+        )
+
+
+if __name__ == "__main__":
+    # TODO: update this old code or remove it.
+    parser = ArgumentParser()
+    parser = TrainSubcommand.add_args(parser)
+
+    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG, force=True)
+    args = parser.parse_args()
+    TrainSubcommand.func(args)

chemprop/cli/utils/__init__.py

@@ -0,0 +1,31 @@
+from .args import bounded
+from .actions import LookupAction
+from .command import Subcommand
+from .parsing import (
+    build_data_from_files,
+    make_datapoints,
+    make_dataset,
+    get_column_names,
+    parse_indices,
+)
+from .utils import pop_attr, _pop_attr, _pop_attr_d, validate_loss_function
+
+__all__ = [
+    "bounded",
+    "LookupAction",
+    "Subcommand",
+    "build_data_from_files",
+    "make_datapoints",
+    "make_dataset",
+    "get_column_names",
+    "parse_indices",
+    "actions",
+    "args",
+    "command",
+    "parsing",
+    "utils",
+    "pop_attr",
+    "_pop_attr",
+    "_pop_attr_d",
+    "validate_loss_function",
+]

chemprop/cli/utils/actions.py

@@ -0,0 +1,28 @@
+from argparse import Action, ArgumentParser, Namespace
+from typing import Any, Mapping, Sequence
+
+
+def LookupAction(obj: Mapping[str, Any]):
+    class LookupAction_(Action):
+        def __init__(self, option_strings, dest, default=None, choices=None, **kwargs):
+            if default not in obj.keys() and default is not None:
+                raise ValueError(
+                    f"Invalid value for arg 'default': '{default}'. "
+                    f"Expected one of {tuple(obj.keys())}"
+                )
+
+            kwargs["choices"] = choices if choices is not None else obj.keys()
+            kwargs["default"] = default
+
+            super().__init__(option_strings, dest, **kwargs)
+
+        def __call__(
+            self,
+            parser: ArgumentParser,
+            namespace: Namespace,
+            values: str | Sequence[Any] | None,
+            option_string: str | None = None,
+        ):
+            setattr(namespace, self.dest, values)
+
+    return LookupAction_

chemprop/cli/utils/args.py

@@ -0,0 +1,34 @@
+import functools
+
+__all__ = ["bounded"]
+
+
+def bounded(lo: float | None = None, hi: float | None = None):
+    if lo is None and hi is None:
+        raise ValueError("No bounds provided!")
+
+    def decorator(f):
+        @functools.wraps(f)
+        def wrapper(*args, **kwargs):
+            x = f(*args, **kwargs)
+
+            if (lo is not None and hi is not None) and not lo <= x <= hi:
+                raise ValueError(f"Parsed value outside of range [{lo}, {hi}]! got: {x}")
+            if hi is not None and x > hi:
+                raise ValueError(f"Parsed value below {hi}! got: {x}")
+            if lo is not None and x < lo:
+                raise ValueError(f"Parsed value above {lo}]! got: {x}")
+
+            return x
+
+        return wrapper
+
+    return decorator
+
+
+def uppercase(x: str):
+    return x.upper()
+
+
+def lowercase(x: str):
+    return x.lower()

chemprop/cli/utils/command.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from argparse import ArgumentParser, _SubParsersAction, Namespace
+
+
+class Subcommand(ABC):
+    COMMAND: str
+    HELP: str | None = None
+
+    @classmethod
+    def add(cls, subparsers: _SubParsersAction, parents) -> ArgumentParser:
+        parser = subparsers.add_parser(cls.COMMAND, help=cls.HELP, parents=parents)
+        cls.add_args(parser).set_defaults(func=cls.func)
+
+        return parser
+
+    @classmethod
+    @abstractmethod
+    def add_args(cls, parser: ArgumentParser) -> ArgumentParser:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def func(cls, args: Namespace):
+        pass

chemprop/cli/utils/parsing.py

@@ -0,0 +1,380 @@
+import logging
+from os import PathLike
+from typing import Mapping, Sequence
+
+import numpy as np
+import pandas as pd
+from rdkit.Chem import Mol
+
+from chemprop.data.datapoints import MoleculeDatapoint, ReactionDatapoint
+from chemprop.data.datasets import MoleculeDataset, ReactionDataset
+from chemprop.featurizers.base import VectorFeaturizer
+from chemprop.featurizers.molgraph import (
+    CondensedGraphOfReactionFeaturizer,
+    SimpleMoleculeMolGraphFeaturizer,
+)
+from chemprop.featurizers.atom import get_multi_hot_atom_featurizer
+
+logger = logging.getLogger(__name__)
+
+
+def parse_csv(
+    path: PathLike,
+    smiles_cols: Sequence[str] | None,
+    rxn_cols: Sequence[str] | None,
+    target_cols: Sequence[str] | None,
+    ignore_cols: Sequence[str] | None,
+    splits_col: str | None,
+    weight_col: str | None,
+    bounded: bool = False,
+    no_header_row: bool = False,
+):
+    df = pd.read_csv(path, header=None if no_header_row else "infer", index_col=False)
+
+    if smiles_cols is not None and rxn_cols is not None:
+        smiss = df[smiles_cols].T.values.tolist()
+        rxnss = df[rxn_cols].T.values.tolist()
+        input_cols = [*smiles_cols, *rxn_cols]
+    elif smiles_cols is not None and rxn_cols is None:
+        smiss = df[smiles_cols].T.values.tolist()
+        rxnss = None
+        input_cols = smiles_cols
+    elif smiles_cols is None and rxn_cols is not None:
+        smiss = None
+        rxnss = df[rxn_cols].T.values.tolist()
+        input_cols = rxn_cols
+    else:
+        smiss = df.iloc[:, [0]].T.values.tolist()
+        rxnss = None
+        input_cols = [df.columns[0]]
+
+    if target_cols is None:
+        target_cols = list(
+            set(df.columns)
+            - set(input_cols)
+            - set(ignore_cols or [])
+            - set(splits_col or [])
+            - set(weight_col or [])
+        )
+
+    Y = df[target_cols]
+    weights = None if weight_col is None else df[weight_col].to_numpy(np.single)
+
+    if bounded:
+        lt_mask = Y.applymap(lambda x: "<" in x).to_numpy()
+        gt_mask = Y.applymap(lambda x: ">" in x).to_numpy()
+        Y = Y.applymap(lambda x: x.strip("<").strip(">")).to_numpy(np.single)
+    else:
+        Y = Y.to_numpy(np.single)
+        lt_mask = None
+        gt_mask = None
+
+    return smiss, rxnss, Y, weights, lt_mask, gt_mask
+
+
+def get_column_names(
+    path: PathLike,
+    smiles_cols: Sequence[str] | None,
+    rxn_cols: Sequence[str] | None,
+    target_cols: Sequence[str] | None,
+    ignore_cols: Sequence[str] | None,
+    splits_col: str | None,
+    weight_col: str | None,
+    no_header_row: bool = False,
+):
+    df = pd.read_csv(path, header=None if no_header_row else "infer", index_col=False)
+
+    if no_header_row:
+        return ["SMILES"] + ["pred_" + str(i) for i in range((len(df.columns) - 1))]
+
+    input_cols = (smiles_cols or []) + (rxn_cols or [])
+
+    if len(input_cols) == 0:
+        input_cols = [df.columns[0]]
+
+    if target_cols is None:
+        target_cols = list(
+            set(df.columns)
+            - set(input_cols)
+            - set(ignore_cols or [])
+            - set(splits_col or [])
+            - set(weight_col or [])
+        )
+
+    return input_cols + target_cols
+
+
+def make_datapoints(
+    smiss: list[list[str]] | None,
+    rxnss: list[list[str]] | None,
+    Y: np.ndarray,
+    weights: np.ndarray | None,
+    lt_mask: np.ndarray | None,
+    gt_mask: np.ndarray | None,
+    X_d: np.ndarray | None,
+    V_fss: list[list[np.ndarray] | list[None]] | None,
+    E_fss: list[list[np.ndarray] | list[None]] | None,
+    V_dss: list[list[np.ndarray] | list[None]] | None,
+    features_generators: list[VectorFeaturizer[Mol]] | None,
+    keep_h: bool,
+    add_h: bool,
+) -> tuple[list[list[MoleculeDatapoint]], list[list[ReactionDatapoint]]]:
+    """Make the :class:`MoleculeDatapoint`s and :class:`ReactionDatapoint`s for a given
+    dataset.
+
+    Parameters
+    ----------
+    smiss : list[list[str]] | None
+        a list of ``j`` lists of ``n`` SMILES strings, where ``j`` is the number of molecules per
+        datapoint and ``n`` is the number of datapoints. If ``None``, the corresponding list of
+        :class:`MoleculeDatapoint`\s will be empty.
+    rxnss : list[list[str]] | None
+        a list of ``k`` lists of ``n`` reaction SMILES strings, where ``k`` is the number of
+        reactions per datapoint. If ``None``, the corresponding list of :class:`ReactionDatapoint`\s
+        will be empty.
+    Y : np.ndarray
+        the target values of shape ``n x m``, where ``m`` is the number of targets
+    weights : np.ndarray | None
+        the weights of the datapoints to use in the loss function of shape ``n x m``. If ``None``,
+        the weights all default to 1.
+    lt_mask : np.ndarray | None
+        a boolean mask of shape ``n x m`` indicating whether the targets are less than inequality
+        targets. If ``None``, ``lt_mask`` for all datapoints will be ``None``.
+    gt_mask : np.ndarray | None
+        a boolean mask of shape ``n x m`` indicating whether the targets are greater than inequality
+        targets. If ``None``, ``gt_mask`` for all datapoints will be ``None``.
+    X_d : np.ndarray | None
+        the extra descriptors of shape ``n x p``, where ``p`` is the number of extra descriptors. If
+        ``None``, ``x_d`` for all datapoints will be ``None``.
+    V_fss : list[list[np.ndarray] | list[None]] | None
+        a list of ``j`` lists of ``n`` np.ndarrays each of shape ``v_jn x q_j``, where ``v_jn`` is
+        the number of atoms in the j-th molecule of the n-th datapoint and ``q_j`` is the number of
+        extra atom features used for the j-th molecules. Any of the ``j`` lists can be a list of
+        None values if the corresponding component does not use extra atom features. If ``None``,
+        ``V_f`` for all datapoints will be ``None``.
+    E_fss : list[list[np.ndarray] | list[None]] | None
+        a list of ``j`` lists of ``n`` np.ndarrays each of shape ``e_jn x r_j``, where ``e_jn`` is
+        the number of bonds in the j-th molecule of the n-th datapoint and ``r_j`` is the number of
+        extra bond features used for the j-th molecules. Any of the ``j`` lists can be a list of
+        None values if the corresponding component does not use extra bond features. If ``None``,
+        ``E_f`` for all datapoints will be ``None``.
+    V_dss : list[list[np.ndarray] | list[None]] | None
+        a list of ``j`` lists of ``n`` np.ndarrays each of shape ``v_jn x s_j``, where ``s_j`` is
+        the number of extra atom descriptors used for the j-th molecules. Any of the ``j`` lists can
+        be a list of None values if the corresponding component does not use extra atom features. If
+        ``None``, ``V_d`` for all datapoints will be ``None``.
+    features_generators : list[MoleculeFeaturizer] | None
+        a list of :class:`MoleculeFeaturizer` instances to generate additional molecule features to
+        use as extra descriptors
+    keep_h : bool
+    add_h : bool
+
+    Returns
+    -------
+    list[list[MoleculeDatapoint]]
+        a list of ``j`` lists of ``n`` :class:`MoleculeDatapoint`\s
+    list[list[ReactionDatapoint]]
+        a list of ``k`` lists of ``n`` :class:`ReactionDatapoint`\s
+    .. note::
+        either ``j`` or ``k`` may be 0, in which case the corresponding list will be empty.
+
+    Raises
+    ------
+    ValueError
+        if both ``smiss`` and ``rxnss`` are ``None``.
+        if ``smiss`` and ``rxnss`` are both given and have different lengths.
+    """
+    if smiss is None and rxnss is None:
+        raise ValueError("args 'smiss' and 'rnxss' were both `None`!")
+    elif rxnss is None:
+        N = len(smiss[0])
+        rxnss = []
+    elif smiss is None:
+        N = len(rxnss[0])
+        smiss = []
+    elif len(smiss[0]) != len(rxnss[0]):
+        raise ValueError(
+            f"args 'smiss' and 'rxnss' must have same length! got {len(smiss[0])} and {len(rxnss[0])}"
+        )
+    else:
+        N = len(smiss[0])
+
+    weights = np.ones(N, dtype=np.single) if weights is None else weights
+    gt_mask = [None] * N if gt_mask is None else gt_mask
+    lt_mask = [None] * N if lt_mask is None else lt_mask
+
+    n_mols = len(smiss) if smiss else 0
+    X_d = [None] * N if X_d is None else X_d
+    V_fss = [[None] * N] * n_mols if V_fss is None else V_fss
+    E_fss = [[None] * N] * n_mols if E_fss is None else E_fss
+    V_dss = [[None] * N] * n_mols if V_dss is None else V_dss
+
+    mol_data = [
+        [
+            MoleculeDatapoint.from_smi(
+                smis[i],
+                keep_h=keep_h,
+                add_h=add_h,
+                y=Y[i],
+                weight=weights[i],
+                gt_mask=gt_mask[i],
+                lt_mask=lt_mask[i],
+                x_d=X_d[i],
+                mfs=features_generators,
+                x_phase=None,
+                V_f=V_fss[mol_idx][i],
+                E_f=E_fss[mol_idx][i],
+                V_d=V_dss[mol_idx][i],
+            )
+            for i in range(N)
+        ]
+        for mol_idx, smis in enumerate(smiss)
+    ]
+    rxn_data = [
+        [
+            ReactionDatapoint.from_smi(
+                rxns[i],
+                keep_h=keep_h,
+                add_h=add_h,
+                y=Y[i],
+                weight=weights[i],
+                gt_mask=gt_mask[i],
+                lt_mask=lt_mask[i],
+                x_d=X_d[i],
+                mfs=features_generators,
+                x_phase=None,
+            )
+            for i in range(N)
+        ]
+        for rxn_idx, rxns in enumerate(rxnss)
+    ]
+
+    return mol_data, rxn_data
+
+
+def build_data_from_files(
+    p_data: PathLike,
+    no_header_row: bool,
+    smiles_cols: Sequence[str] | None,
+    rxn_cols: Sequence[str] | None,
+    target_cols: Sequence[str] | None,
+    ignore_cols: Sequence[str] | None,
+    splits_col: str | None,
+    weight_col: str | None,
+    bounded: bool,
+    p_descriptors: PathLike,
+    p_atom_feats: dict[int, PathLike],
+    p_bond_feats: dict[int, PathLike],
+    p_atom_descs: dict[int, PathLike],
+    **featurization_kwargs: Mapping,
+) -> list[list[MoleculeDatapoint] | list[ReactionDatapoint]]:
+    smiss, rxnss, Y, weights, lt_mask, gt_mask = parse_csv(
+        p_data,
+        smiles_cols,
+        rxn_cols,
+        target_cols,
+        ignore_cols,
+        splits_col,
+        weight_col,
+        bounded,
+        no_header_row,
+    )
+    n_molecules = len(smiss) if smiss is not None else 0
+    n_datapoints = len(Y)
+
+    X_ds = load_input_feats_and_descs(p_descriptors, None, None, feat_desc="X_d")
+    V_fss = load_input_feats_and_descs(p_atom_feats, n_molecules, n_datapoints, feat_desc="V_f")
+    E_fss = load_input_feats_and_descs(p_bond_feats, n_molecules, n_datapoints, feat_desc="E_f")
+    V_dss = load_input_feats_and_descs(p_atom_descs, n_molecules, n_datapoints, feat_desc="V_d")
+
+    mol_data, rxn_data = make_datapoints(
+        smiss,
+        rxnss,
+        Y,
+        weights,
+        lt_mask,
+        gt_mask,
+        X_ds,
+        V_fss,
+        E_fss,
+        V_dss,
+        **featurization_kwargs,
+    )
+
+    return mol_data + rxn_data
+
+
+def load_input_feats_and_descs(
+    paths: dict[int, PathLike] | PathLike,
+    n_molecules: int | None,
+    n_datapoints: int | None,
+    feat_desc: str,
+):
+    if paths is None:
+        return None
+
+    match feat_desc:
+        case "X_d":
+            path = paths
+            loaded_feature = np.load(path)
+            features = loaded_feature["arr_0"]
+
+        case _:
+            for index in paths:
+                if index >= n_molecules:
+                    raise ValueError(
+                        f"For {n_molecules} molecules, atom/bond features/descriptors can only be specified for indices 0-{n_molecules - 1}! Got index {index}."
+                    )
+
+            features = []
+            for idx in range(n_molecules):
+                path = paths.get(idx, None)
+
+                if path is not None:
+                    loaded_feature = np.load(path)
+                    loaded_feature = [
+                        loaded_feature[f"arr_{i}"] for i in range(len(loaded_feature))
+                    ]
+                else:
+                    loaded_feature = [None] * n_datapoints
+
+                features.append(loaded_feature)
+    return features
+
+
+def make_dataset(
+    data: Sequence[MoleculeDatapoint] | Sequence[ReactionDatapoint],
+    reaction_mode: str,
+    multi_hot_atom_featurizer_mode: str = "V2",
+) -> MoleculeDataset | ReactionDataset:
+    atom_featurizer = get_multi_hot_atom_featurizer(multi_hot_atom_featurizer_mode)
+
+    if isinstance(data[0], MoleculeDatapoint):
+        extra_atom_fdim = data[0].V_f.shape[1] if data[0].V_f is not None else 0
+        extra_bond_fdim = data[0].E_f.shape[1] if data[0].E_f is not None else 0
+        featurizer = SimpleMoleculeMolGraphFeaturizer(
+            atom_featurizer=atom_featurizer,
+            extra_atom_fdim=extra_atom_fdim,
+            extra_bond_fdim=extra_bond_fdim,
+        )
+        return MoleculeDataset(data, featurizer)
+
+    featurizer = CondensedGraphOfReactionFeaturizer(
+        mode_=reaction_mode, atom_featurizer=atom_featurizer
+    )
+
+    return ReactionDataset(data, featurizer)
+
+
+def parse_indices(idxs):
+    """Parses a string of indices into a list of integers. e.g. '0,1,2-4' -> [0, 1, 2, 3, 4]"""
+    if isinstance(idxs, str):
+        indices = []
+        for idx in idxs.split(","):
+            if "-" in idx:
+                start, end = map(int, idx.split("-"))
+                indices.extend(range(start, end + 1))
+            else:
+                indices.append(int(idx))
+        return indices
+    return idxs

chemprop/cli/utils/utils.py

@@ -0,0 +1,76 @@
+from typing import Any, Type
+
+from chemprop.nn import loss, predictors
+
+__all__ = ["pop_attr"]
+
+
+def pop_attr(o: object, attr: str, *args) -> Any | None:
+    """like ``pop()`` but for attribute maps"""
+    match len(args):
+        case 0:
+            return _pop_attr(o, attr)
+        case 1:
+            return _pop_attr_d(o, attr, args[0])
+        case _:
+            raise TypeError(f"Expected at most 2 arguments! got: {len(args)}")
+
+
+def _pop_attr(o: object, attr: str) -> Any:
+    val = getattr(o, attr)
+    delattr(o, attr)
+
+    return val
+
+
+def _pop_attr_d(o: object, attr: str, default: Any | None = None) -> Any | None:
+    try:
+        val = getattr(o, attr)
+        delattr(o, attr)
+    except AttributeError:
+        val = default
+
+    return val
+
+
+def validate_loss_function(
+    predictor_ffn: Type[predictors._FFNPredictorBase], criterion: Type[loss.LossFunction]
+):
+    match predictor_ffn:
+        case predictors.RegressionFFN:
+            if criterion not in (loss.MSELoss, loss.BoundedMSELoss):
+                raise ValueError(f"Expected a regression loss function! got: {criterion.__name__}")
+        case predictors.MveFFN:
+            if criterion is not loss.MVELoss:
+                raise ValueError(f"Expected a MVE loss function! got: {criterion.__name__}")
+        case predictors.EvidentialFFN:
+            if criterion is not loss.EvidentialLoss:
+                raise ValueError(f"Expected an evidential loss function! got: {criterion.__name__}")
+        case predictors.BinaryClassificationFFN:
+            if criterion not in (loss.BCELoss, loss.BinaryMCCLoss):
+                raise ValueError(
+                    f"Expected a binary classification loss function! got: {criterion.__name__}"
+                )
+        case predictors.BinaryDirichletFFN:
+            if loss is not loss.BinaryDirichletLoss:
+                raise ValueError(
+                    f"Expected a binary Dirichlet loss function! got: {criterion.__name__}"
+                )
+        case predictors.MulticlassClassificationFFN:
+            if loss not in (loss.CrossEntropyLoss, loss.MulticlassMCCLoss):
+                raise ValueError(
+                    f"Expected a multiclass classification loss function! got: {criterion.__name__}"
+                )
+        case predictors.MulticlassDirichletFFN:
+            if loss is not loss.MulticlassDirichletLoss:
+                raise ValueError(
+                    f"Expected a multiclass Dirichlet loss function! got: {criterion.__name__}"
+                )
+        case predictors.SpectralFFN:
+            if loss not in (loss.SIDLoss, loss.WassersteinLoss):
+                raise ValueError(f"Expected a spectral loss function! got: {criterion.__name__}")
+        case _:
+            raise ValueError(
+                f"Unknown predictor function! got: {predictor_ffn}. "
+                f"Expected one of: {tuple(predictors.PredictorRegistry.values())}"
+            )

chemprop/conf.py

@@ -0,0 +1,7 @@
+"""Global configuration variables for chemprop"""
+
+from chemprop.featurizers.molgraph.molecule import SimpleMoleculeMolGraphFeaturizer
+
+
+DEFAULT_ATOM_FDIM, DEFAULT_BOND_FDIM = SimpleMoleculeMolGraphFeaturizer().shape
+DEFAULT_HIDDEN_DIM = 300

chemprop/data/__init__.py

@@ -0,0 +1,34 @@
+from .collate import BatchMolGraph, TrainingBatch, collate_batch, collate_multicomponent
+from .dataloader import build_dataloader
+from .datapoints import MoleculeDatapoint, ReactionDatapoint
+from .datasets import (
+    MoleculeDataset,
+    ReactionDataset,
+    Datum,
+    MulticomponentDataset,
+    MolGraphDataset,
+)
+from .molgraph import MolGraph
+from .samplers import ClassBalanceSampler, SeededSampler
+from .splitting import SplitType, make_split_indices, split_data_by_indices
+
+__all__ = [
+    "BatchMolGraph",
+    "TrainingBatch",
+    "collate_batch",
+    "collate_multicomponent",
+    "build_dataloader",
+    "MoleculeDatapoint",
+    "ReactionDatapoint",
+    "MoleculeDataset",
+    "ReactionDataset",
+    "Datum",
+    "MulticomponentDataset",
+    "MolGraphDataset",
+    "MolGraph",
+    "ClassBalanceSampler",
+    "SeededSampler",
+    "SplitType",
+    "make_split_indices",
+    "split_data_by_indices",
+]

chemprop/data/collate.py

@@ -0,0 +1,120 @@
+from dataclasses import dataclass, field, InitVar
+from typing import Iterable, NamedTuple, Sequence
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from chemprop.data.datasets import Datum
+from chemprop.data.molgraph import MolGraph
+
+
+@dataclass(repr=False, eq=False, slots=True)
+class BatchMolGraph:
+    """A :class:`BatchMolGraph` represents a batch of individual :class:`MolGraph`\s.
+
+    It has all the attributes of a ``MolGraph`` with the addition of the ``batch`` attribute. This
+    class is intended for use with data loading, so it uses :obj:`~torch.Tensor`\s to store data
+    """
+
+    mgs: InitVar[Sequence[MolGraph]]
+    """A list of individual :class:`MolGraph`\s to be batched together"""
+    V: Tensor = field(init=False)
+    """the atom feature matrix"""
+    E: Tensor = field(init=False)
+    """the bond feature matrix"""
+    edge_index: Tensor = field(init=False)
+    """an tensor of shape ``2 x E`` containing the edges of the graph in COO format"""
+    rev_edge_index: Tensor = field(init=False)
+    """A tensor of shape ``E`` that maps from an edge index to the index of the source of the
+    reverse edge in the ``edge_index`` attribute."""
+    batch: Tensor = field(init=False)
+    """the index of the parent :class:`MolGraph` in the batched graph"""
+
+    __size: int = field(init=False)
+
+    def __post_init__(self, mgs: Sequence[MolGraph]):
+        self.__size = len(mgs)
+
+        Vs = []
+        Es = []
+        edge_indexes = []
+        rev_edge_indexes = []
+        batch_indexes = []
+
+        num_nodes = 0
+        num_edges = 0
+        for i, mg in enumerate(mgs):
+            Vs.append(mg.V)
+            Es.append(mg.E)
+            edge_indexes.append(mg.edge_index + num_nodes)
+            rev_edge_indexes.append(mg.rev_edge_index + num_edges)
+            batch_indexes.append([i] * len(mg.V))
+
+            num_nodes += mg.V.shape[0]
+            num_edges += mg.edge_index.shape[1]
+
+        self.V = torch.from_numpy(np.concatenate(Vs)).float()
+        self.E = torch.from_numpy(np.concatenate(Es)).float()
+        self.edge_index = torch.from_numpy(np.hstack(edge_indexes)).long()
+        self.rev_edge_index = torch.from_numpy(np.concatenate(rev_edge_indexes)).long()
+        self.batch = torch.tensor(np.concatenate(batch_indexes)).long()
+
+    def __len__(self) -> int:
+        """the number of individual :class:`MolGraph`\s in this batch"""
+        return self.__size
+
+    def to(self, device: str | torch.device):
+        self.V = self.V.to(device)
+        self.E = self.E.to(device)
+        self.edge_index = self.edge_index.to(device)
+        self.rev_edge_index = self.rev_edge_index.to(device)
+        self.batch = self.batch.to(device)
+
+
+class TrainingBatch(NamedTuple):
+    bmg: BatchMolGraph
+    V_d: Tensor | None
+    X_d: Tensor | None
+    Y: Tensor | None
+    w: Tensor
+    lt_mask: Tensor | None
+    gt_mask: Tensor | None
+
+
+def collate_batch(batch: Iterable[Datum]) -> TrainingBatch:
+    mgs, V_ds, x_ds, ys, weights, lt_masks, gt_masks = zip(*batch)
+
+    return TrainingBatch(
+        BatchMolGraph(mgs),
+        None if V_ds[0] is None else torch.from_numpy(np.concatenate(V_ds)).float(),
+        None if x_ds[0] is None else torch.from_numpy(np.array(x_ds)).float(),
+        None if ys[0] is None else torch.from_numpy(np.array(ys)).float(),
+        torch.tensor(weights, dtype=torch.float).unsqueeze(1),
+        None if lt_masks[0] is None else torch.from_numpy(np.array(lt_masks)),
+        None if gt_masks[0] is None else torch.from_numpy(np.array(gt_masks)),
+    )
+
+
+class MulticomponentTrainingBatch(NamedTuple):
+    bmgs: list[BatchMolGraph]
+    V_ds: list[Tensor | None]
+    X_d: Tensor | None
+    Y: Tensor | None
+    w: Tensor
+    lt_mask: Tensor | None
+    gt_mask: Tensor | None
+
+
+def collate_multicomponent(batches: Iterable[Iterable[Datum]]) -> MulticomponentTrainingBatch:
+    tbs = [collate_batch(batch) for batch in zip(*batches)]
+
+    return MulticomponentTrainingBatch(
+        [tb.bmg for tb in tbs],
+        [tb.V_d for tb in tbs],
+        tbs[0].X_d,
+        tbs[0].Y,
+        tbs[0].w,
+        tbs[0].lt_mask,
+        tbs[0].gt_mask,
+    )

chemprop/data/dataloader.py

@@ -0,0 +1,69 @@
+import warnings
+
+from torch.utils.data import DataLoader
+
+from chemprop.data.collate import collate_batch, collate_multicomponent
+from chemprop.data.datasets import MoleculeDataset, MulticomponentDataset, ReactionDataset
+from chemprop.data.samplers import ClassBalanceSampler, SeededSampler
+
+
+def build_dataloader(
+    dataset: MoleculeDataset | ReactionDataset | MulticomponentDataset,
+    batch_size: int = 64,
+    num_workers: int = 0,
+    class_balance: bool = False,
+    seed: int | None = None,
+    shuffle: bool = True,
+    **kwargs,
+):
+    """Return a :obj:`~torch.utils.data.DataLoader` for :class:`MolGraphDataset`\s
+
+    Parameters
+    ----------
+    dataset : MoleculeDataset | ReactionDataset | MulticomponentDataset
+        The dataset containing the molecules or reactions to load.
+    batch_size : int, default=64
+        the batch size to load.
+    num_workers : int, default=0
+        the number of workers used to build batches.
+    class_balance : bool, default=False
+        Whether to perform class balancing (i.e., use an equal number of positive and negative
+        molecules). Class balance is only available for single task classification datasets. Set
+        shuffle to True in order to get a random subset of the larger class.
+    seed : int, default=None
+        the random seed to use for shuffling (only used when `shuffle` is `True`).
+    shuffle : bool, default=False
+        whether to shuffle the data during sampling.
+    """
+
+    if class_balance:
+        sampler = ClassBalanceSampler(dataset.Y, seed, shuffle)
+    elif shuffle and seed is not None:
+        sampler = SeededSampler(len(dataset), seed)
+    else:
+        sampler = None
+
+    if isinstance(dataset, MulticomponentDataset):
+        collate_fn = collate_multicomponent
+    else:
+        collate_fn = collate_batch
+
+    if len(dataset) % batch_size == 1:
+        warnings.warn(
+            f"Dropping last batch of size 1 to avoid issues with batch normalization \
+(dataset size = {len(dataset)}, batch_size = {batch_size})"
+        )
+        drop_last = True
+    else:
+        drop_last = False
+
+    return DataLoader(
+        dataset,
+        batch_size,
+        sampler is None and shuffle,
+        sampler,
+        num_workers=num_workers,
+        collate_fn=collate_fn,
+        drop_last=drop_last,
+        **kwargs,
+    )

chemprop/data/datapoints.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+from dataclasses import InitVar, dataclass
+
+import numpy as np
+from rdkit.Chem import AllChem as Chem
+
+from chemprop.featurizers import Featurizer
+from chemprop.utils import make_mol
+
+MoleculeFeaturizer = Featurizer[Chem.Mol, np.ndarray]
+
+
+@dataclass(slots=True)
+class _DatapointMixin:
+    """A mixin class for both molecule- and reaction- and multicomponent-type data"""
+
+    y: np.ndarray | None = None
+    """the targets for the molecule with unknown targets indicated by `nan`s"""
+    weight: float = 1.0
+    """the weight of this datapoint for the loss calculation."""
+    gt_mask: np.ndarray | None = None
+    """Indicates whether the targets are an inequality regression target of the form `<x`"""
+    lt_mask: np.ndarray | None = None
+    """Indicates whether the targets are an inequality regression target of the form `>x`"""
+    x_d: np.ndarray | None = None
+    """A vector of length ``d_f`` containing additional features (e.g., Morgan fingerprint) that
+    will be concatenated to the global representation *after* aggregation"""
+    mfs: InitVar[list[MoleculeFeaturizer] | None] = None
+    """A list of molecule featurizers to use"""
+    x_phase: list[float] = None
+    """A one-hot vector indicating the phase of the data, as used in spectra data."""
+    name: str | None = None
+    """A string identifier for the datapoint."""
+
+    def __post_init__(self, mfs: list[MoleculeFeaturizer] | None):
+        if self.x_d is not None and mfs is not None:
+            raise ValueError("Cannot provide both loaded features and molecular featurizers!")
+
+        if mfs is not None:
+            self.x_d = self.calc_features(mfs)
+
+        NAN_TOKEN = 0
+        if self.x_d is not None:
+            self.x_d[np.isnan(self.x_d)] = NAN_TOKEN
+
+    @property
+    def t(self) -> int | None:
+        return len(self.y) if self.y is not None else None
+
+
+@dataclass
+class _MoleculeDatapointMixin:
+    mol: Chem.Mol
+    """the molecule associated with this datapoint"""
+
+    @classmethod
+    def from_smi(
+        cls, smi: str, *args, keep_h: bool = False, add_h: bool = False, **kwargs
+    ) -> _MoleculeDatapointMixin:
+        mol = make_mol(smi, keep_h, add_h)
+
+        kwargs["name"] = smi if "name" not in kwargs else kwargs["name"]
+
+        return cls(mol, *args, **kwargs)
+
+
+@dataclass
+class MoleculeDatapoint(_DatapointMixin, _MoleculeDatapointMixin):
+    """A :class:`MoleculeDatapoint` contains a single molecule and its associated features and targets."""
+
+    V_f: np.ndarray | None = None
+    """a numpy array of shape ``V x d_vf``, where ``V`` is the number of atoms in the molecule, and
+    ``d_vf`` is the number of additional features that will be concatenated to atom-level features
+    *before* message passing"""
+    E_f: np.ndarray | None = None
+    """A numpy array of shape ``E x d_ef``, where ``E`` is the number of bonds in the molecule, and
+    ``d_ef`` is the number of additional features  containing additional features that will be
+    concatenated to bond-level features *before* message passing"""
+    V_d: np.ndarray | None = None
+    """A numpy array of shape ``V x d_vd``, where ``V`` is the number of atoms in the molecule, and
+    ``d_vd`` is the number of additional descriptors that will be concatenated to atom-level
+    descriptors *after* message passing"""
+
+    def __post_init__(self, mfs: list[MoleculeFeaturizer] | None):
+        if self.mol is None:
+            raise ValueError("Input molecule was `None`!")
+
+        NAN_TOKEN = 0
+
+        if self.V_f is not None:
+            self.V_f[np.isnan(self.V_f)] = NAN_TOKEN
+        if self.E_f is not None:
+            self.E_f[np.isnan(self.E_f)] = NAN_TOKEN
+        if self.V_d is not None:
+            self.V_d[np.isnan(self.V_d)] = NAN_TOKEN
+
+        super().__post_init__(mfs)
+
+    def __len__(self) -> int:
+        return 1
+
+    def calc_features(self, mfs: list[MoleculeFeaturizer]) -> np.ndarray:
+        if self.mol.GetNumHeavyAtoms() == 0:
+            return np.zeros(sum(len(mf) for mf in mfs))
+
+        return np.hstack([mf(self.mol) for mf in mfs])
+
+
+@dataclass
+class _ReactionDatapointMixin:
+    rct: Chem.Mol
+    """the reactant associated with this datapoint"""
+    pdt: Chem.Mol
+    """the product associated with this datapoint"""
+
+    @classmethod
+    def from_smi(
+        cls,
+        rxn_or_smis: str | tuple[str, str],
+        *args,
+        keep_h: bool = False,
+        add_h: bool = False,
+        **kwargs,
+    ) -> _ReactionDatapointMixin:
+        match rxn_or_smis:
+            case str():
+                rct_smi, agt_smi, pdt_smi = rxn_or_smis.split(">")
+                rct_smi = f"{rct_smi}.{agt_smi}" if agt_smi else rct_smi
+                name = rxn_or_smis
+            case tuple():
+                rct_smi, pdt_smi = rxn_or_smis
+                name = ">>".join(rxn_or_smis)
+            case _:
+                raise TypeError(
+                    "Must provide either a reaction SMARTS string or a tuple of reactant and product SMILES strings!"
+                )
+
+        rct = make_mol(rct_smi, keep_h, add_h)
+        pdt = make_mol(pdt_smi, keep_h, add_h)
+
+        kwargs["name"] = name if "name" not in kwargs else kwargs["name"]
+
+        return cls(rct, pdt, *args, **kwargs)
+
+
+@dataclass
+class ReactionDatapoint(_DatapointMixin, _ReactionDatapointMixin):
+    """A :class:`ReactionDatapoint` contains a single reaction and its associated features and targets."""
+
+    def __post_init__(self, mfs: list[MoleculeFeaturizer] | None):
+        if self.rct is None:
+            raise ValueError("Reactant cannot be `None`!")
+        if self.pdt is None:
+            raise ValueError("Product cannot be `None`!")
+
+        return super().__post_init__(mfs)
+
+    def __len__(self) -> int:
+        return 2
+
+    def calc_features(self, mfs: list[MoleculeFeaturizer]) -> np.ndarray:
+        x_ds = [
+            mf(mol) if mol.GetNumHeavyAtoms() > 0 else np.zeros(len(mf))
+            for mf in mfs
+            for mol in [self.rct, self.pdt]
+        ]
+
+        return np.hstack(x_ds)

chemprop/data/datasets.py

@@ -0,0 +1,459 @@
+from dataclasses import dataclass, field
+from functools import cached_property
+from typing import NamedTuple, TypeAlias
+
+import numpy as np
+from numpy.typing import ArrayLike
+from rdkit import Chem
+from rdkit.Chem import Mol
+from sklearn.preprocessing import StandardScaler
+from torch.utils.data import Dataset
+
+from chemprop.types import Rxn
+from chemprop.data.datapoints import MoleculeDatapoint, ReactionDatapoint
+from chemprop.data.molgraph import MolGraph
+from chemprop.featurizers.base import Featurizer
+from chemprop.featurizers.molgraph.cache import MolGraphCache, MolGraphCacheOnTheFly
+from chemprop.featurizers.molgraph import SimpleMoleculeMolGraphFeaturizer, CGRFeaturizer
+
+
+class Datum(NamedTuple):
+    """a singular training data point"""
+
+    mg: MolGraph
+    V_d: np.ndarray | None
+    x_d: np.ndarray | None
+    y: np.ndarray | None
+    weight: float
+    lt_mask: np.ndarray | None
+    gt_mask: np.ndarray | None
+
+
+MolGraphDataset: TypeAlias = Dataset[Datum]
+
+
+class _MolGraphDatasetMixin:
+    def __len__(self) -> int:
+        return len(self.data)
+
+    @cached_property
+    def _Y(self) -> np.ndarray:
+        """the raw targets of the dataset"""
+        return np.array([d.y for d in self.data], float)
+
+    @property
+    def Y(self) -> np.ndarray:
+        """the (scaled) targets of the dataset"""
+        return self.__Y
+
+    @Y.setter
+    def Y(self, Y: ArrayLike):
+        self._validate_attribute(Y, "targets")
+
+        self.__Y = np.array(Y, float)
+
+    @cached_property
+    def _X_d(self) -> np.ndarray:
+        """the raw extra descriptors of the dataset"""
+        return np.array([d.x_d for d in self.data])
+
+    @property
+    def X_d(self) -> np.ndarray:
+        """the (scaled) extra descriptors of the dataset"""
+        return self.__X_d
+
+    @X_d.setter
+    def X_d(self, X_d: ArrayLike):
+        self._validate_attribute(X_d, "extra descriptors")
+
+        self.__X_d = np.array(X_d)
+
+    @property
+    def weights(self) -> np.ndarray:
+        return np.array([d.weight for d in self.data])
+
+    @property
+    def gt_mask(self) -> np.ndarray:
+        return np.array([d.gt_mask for d in self.data])
+
+    @property
+    def lt_mask(self) -> np.ndarray:
+        return np.array([d.lt_mask for d in self.data])
+
+    @property
+    def t(self) -> int | None:
+        return self.data[0].t if len(self.data) > 0 else None
+
+    @property
+    def d_xd(self) -> int:
+        """the extra molecule descriptor dimension, if any"""
+        return 0 if self.X_d[0] is None else self.X_d.shape[1]
+
+    @property
+    def names(self) -> list[str]:
+        return [d.name for d in self.data]
+
+    def normalize_targets(self, scaler: StandardScaler | None = None) -> StandardScaler:
+        """Normalizes the targets of this dataset using a :obj:`StandardScaler`
+
+        The :obj:`StandardScaler` subtracts the mean and divides by the standard deviation for
+        each task independently. NOTE: This should only be used for regression datasets.
+
+        Returns
+        -------
+        StandardScaler
+            a scaler fit to the targets.
+        """
+
+        if scaler is None:
+            scaler = StandardScaler().fit(self._Y)
+
+        self.Y = scaler.transform(self._Y)
+
+        return scaler
+
+    def normalize_inputs(
+        self, key: str = "X_d", scaler: StandardScaler | None = None
+    ) -> StandardScaler:
+        VALID_KEYS = {"X_d"}
+        if key not in VALID_KEYS:
+            raise ValueError(f"Invalid feature key! got: {key}. expected one of: {VALID_KEYS}")
+
+        X = self.X_d if self.X_d[0] is not None else None
+
+        if X is None:
+            return scaler
+
+        if scaler is None:
+            scaler = StandardScaler().fit(X)
+
+        self.X_d = scaler.transform(X)
+
+        return scaler
+
+    def reset(self):
+        """Reset the atom and bond features; atom and extra descriptors; and targets of each
+        datapoint to their initial, unnormalized values."""
+        self.__Y = self._Y
+        self.__X_d = self._X_d
+
+    def _validate_attribute(self, X: np.ndarray, label: str):
+        if not len(self.data) == len(X):
+            raise ValueError(
+                f"number of molecules ({len(self.data)}) and {label} ({len(X)}) "
+                "must have same length!"
+            )
+
+
+@dataclass
+class MoleculeDataset(_MolGraphDatasetMixin, MolGraphDataset):
+    """A :class:`MoleculeDataset` composed of :class:`MoleculeDatapoint`\s
+
+    A :class:`MoleculeDataset` produces featurized data for input to a
+    :class:`MPNN` model. Typically, data featurization is performed on-the-fly
+    and parallelized across multiple workers via the :class:`~torch.utils.data
+    DataLoader` class. However, for small datasets, it may be more efficient to
+    featurize the data in advance and cache the results. This can be done by
+    setting ``MoleculeDataset.cache=True``.
+
+    Parameters
+    ----------
+    data : Iterable[MoleculeDatapoint]
+        the data from which to create a dataset
+    featurizer : MoleculeFeaturizer
+        the featurizer with which to generate MolGraphs of the molecules
+    """
+
+    data: list[MoleculeDatapoint]
+    featurizer: Featurizer[Mol, MolGraph] = field(default_factory=SimpleMoleculeMolGraphFeaturizer)
+
+    def __post_init__(self):
+        if self.data is None:
+            raise ValueError("Data cannot be None!")
+
+        self.reset()
+        self.cache = False
+
+    def __getitem__(self, idx: int) -> Datum:
+        d = self.data[idx]
+        mg = self.mg_cache[idx]
+
+        return Datum(mg, self.V_ds[idx], self.X_d[idx], self.Y[idx], d.weight, d.lt_mask, d.gt_mask)
+
+    @property
+    def cache(self) -> bool:
+        return self.__cache
+
+    @cache.setter
+    def cache(self, cache: bool = False):
+        self.__cache = cache
+        self._init_cache()
+
+    def _init_cache(self):
+        """initialize the cache"""
+        self.mg_cache = (MolGraphCache if self.cache else MolGraphCacheOnTheFly)(
+            self.mols, self.V_fs, self.E_fs, self.featurizer
+        )
+
+    @property
+    def smiles(self) -> list[str]:
+        """the SMILES strings associated with the dataset"""
+        return [Chem.MolToSmiles(d.mol) for d in self.data]
+
+    @property
+    def mols(self) -> list[Chem.Mol]:
+        """the molecules associated with the dataset"""
+        return [d.mol for d in self.data]
+
+    @property
+    def _V_fs(self) -> list[np.ndarray]:
+        """the raw atom features of the dataset"""
+        return [d.V_f for d in self.data]
+
+    @property
+    def V_fs(self) -> list[np.ndarray]:
+        """the (scaled) atom descriptors of the dataset"""
+        return self.__V_fs
+
+    @V_fs.setter
+    def V_fs(self, V_fs: list[np.ndarray]):
+        """the (scaled) atom features of the dataset"""
+        self._validate_attribute(V_fs, "atom features")
+
+        self.__V_fs = V_fs
+        self._init_cache()
+
+    @property
+    def _E_fs(self) -> list[np.ndarray]:
+        """the raw bond features of the dataset"""
+        return [d.E_f for d in self.data]
+
+    @property
+    def E_fs(self) -> list[np.ndarray]:
+        """the (scaled) bond features of the dataset"""
+        return self.__E_fs
+
+    @E_fs.setter
+    def E_fs(self, E_fs: list[np.ndarray]):
+        self._validate_attribute(E_fs, "bond features")
+
+        self.__E_fs = E_fs
+        self._init_cache()
+
+    @property
+    def _V_ds(self) -> list[np.ndarray]:
+        """the raw atom descriptors of the dataset"""
+        return [d.V_d for d in self.data]
+
+    @property
+    def V_ds(self) -> list[np.ndarray]:
+        """the (scaled) atom descriptors of the dataset"""
+        return self.__V_ds
+
+    @V_ds.setter
+    def V_ds(self, V_ds: list[np.ndarray]):
+        self._validate_attribute(V_ds, "atom descriptors")
+
+        self.__V_ds = V_ds
+
+    @property
+    def d_vf(self) -> int:
+        """the extra atom feature dimension, if any"""
+        return 0 if self.V_fs[0] is None else self.V_fs[0].shape[1]
+
+    @property
+    def d_ef(self) -> int:
+        """the extra bond feature dimension, if any"""
+        return 0 if self.E_fs[0] is None else self.E_fs[0].shape[1]
+
+    @property
+    def d_vd(self) -> int:
+        """the extra atom descriptor dimension, if any"""
+        return 0 if self.V_ds[0] is None else self.V_ds[0].shape[1]
+
+    def normalize_inputs(
+        self, key: str = "X_d", scaler: StandardScaler | None = None
+    ) -> StandardScaler:
+        VALID_KEYS = {"X_d", "V_f", "E_f", "V_d"}
+
+        match key:
+            case "X_d":
+                X = None if self.d_xd == 0 else self.X_d
+            case "V_f":
+                X = None if self.d_vf == 0 else np.concatenate(self.V_fs, axis=0)
+            case "E_f":
+                X = None if self.d_ef == 0 else np.concatenate(self.E_fs, axis=0)
+            case "V_d":
+                X = None if self.d_vd == 0 else np.concatenate(self.V_ds, axis=0)
+            case _:
+                raise ValueError(f"Invalid feature key! got: {key}. expected one of: {VALID_KEYS}")
+
+        if X is None:
+            return scaler
+
+        if scaler is None:
+            scaler = StandardScaler().fit(X)
+
+        match key:
+            case "X_d":
+                self.X_d = scaler.transform(X)
+            case "V_f":
+                self.V_fs = [scaler.transform(V_f) if V_f.size > 0 else V_f for V_f in self.V_fs]
+            case "E_f":
+                self.E_fs = [scaler.transform(E_f) if E_f.size > 0 else E_f for E_f in self.E_fs]
+            case "V_d":
+                self.V_ds = [scaler.transform(V_d) if V_d.size > 0 else V_d for V_d in self.V_ds]
+            case _:
+                raise RuntimeError("unreachable code reached!")
+
+        return scaler
+
+    def reset(self):
+        """Reset the atom and bond features; atom and extra descriptors; and targets of each
+        datapoint to their initial, unnormalized values."""
+        super().reset()
+        self.__V_fs = self._V_fs
+        self.__E_fs = self._E_fs
+        self.__V_ds = self._V_ds
+
+
+@dataclass
+class ReactionDataset(_MolGraphDatasetMixin, MolGraphDataset):
+    """A :class:`ReactionDataset` composed of :class:`ReactionDatapoint`\s
+
+    .. note::
+        The featurized data provided by this class may be cached, simlar to a
+        :class:`MoleculeDataset`. To enable the cache, set ``ReactionDataset
+        cache=True``.
+    """
+
+    data: list[ReactionDatapoint]
+    """the dataset from which to load"""
+    featurizer: Featurizer[Rxn, MolGraph] = field(default_factory=CGRFeaturizer)
+    """the featurizer with which to generate MolGraphs of the input"""
+
+    def __post_init__(self):
+        if self.data is None:
+            raise ValueError("Data cannot be None!")
+
+        self.reset()
+        self.cache = False
+
+    @property
+    def cache(self) -> bool:
+        return self.__cache
+
+    @cache.setter
+    def cache(self, cache: bool = False):
+        self.__cache = cache
+        self.mg_cache = (MolGraphCache if cache else MolGraphCacheOnTheFly)(
+            self.mols, [None] * len(self), [None] * len(self), self.featurizer
+        )
+
+    def __getitem__(self, idx: int) -> Datum:
+        d = self.data[idx]
+        mg = self.mg_cache[idx]
+
+        return Datum(mg, None, self.X_d[idx], self.Y[idx], d.weight, d.lt_mask, d.gt_mask)
+
+    @property
+    def smiles(self) -> list[tuple]:
+        return [(Chem.MolToSmiles(d.rct), Chem.MolToSmiles(d.pdt)) for d in self.data]
+
+    @property
+    def mols(self) -> list[Rxn]:
+        return [(d.rct, d.pdt) for d in self.data]
+
+    @property
+    def d_vf(self) -> int:
+        return 0
+
+    @property
+    def d_ef(self) -> int:
+        return 0
+
+    @property
+    def d_vd(self) -> int:
+        return 0
+
+
+@dataclass(repr=False, eq=False)
+class MulticomponentDataset(_MolGraphDatasetMixin, Dataset):
+    """A :class:`MulticomponentDataset` is a :class:`Dataset` composed of parallel
+    :class:`MoleculeDatasets` and :class:`ReactionDataset`\s"""
+
+    datasets: list[MoleculeDataset | ReactionDataset]
+    """the parallel datasets"""
+
+    def __post_init__(self):
+        sizes = [len(dset) for dset in self.datasets]
+        if not all(sizes[0] == size for size in sizes[1:]):
+            raise ValueError(f"Datasets must have all same length! got: {sizes}")
+
+    def __len__(self) -> int:
+        return len(self.datasets[0])
+
+    @property
+    def n_components(self) -> int:
+        return len(self.datasets)
+
+    def __getitem__(self, idx: int) -> list[Datum]:
+        return [dset[idx] for dset in self.datasets]
+
+    @property
+    def smiles(self) -> list[list[str]]:
+        return list(zip(*[dset.smiles for dset in self.datasets]))
+
+    @property
+    def names(self) -> list[list[str]]:
+        return list(zip(*[dset.names for dset in self.datasets]))
+
+    @property
+    def mols(self) -> list[list[Chem.Mol]]:
+        return list(zip(*[dset.mols for dset in self.datasets]))
+
+    def normalize_targets(self, scaler: StandardScaler | None = None) -> StandardScaler:
+        return self.datasets[0].normalize_targets(scaler)
+
+    def normalize_inputs(
+        self, key: str = "X_d", scaler: list[StandardScaler] | None = None
+    ) -> list[StandardScaler]:
+        RXN_VALID_KEYS = {"X_d"}
+        match scaler:
+            case None:
+                return [
+                    dset.normalize_inputs(key)
+                    if isinstance(dset, MoleculeDataset) or key in RXN_VALID_KEYS
+                    else None
+                    for dset in self.datasets
+                ]
+            case _:
+                assert len(scaler) == len(
+                    self.datasets
+                ), "Number of scalers must match number of datasets!"
+
+                return [
+                    dset.normalize_inputs(key, s)
+                    if isinstance(dset, MoleculeDataset) or key in RXN_VALID_KEYS
+                    else None
+                    for dset, s in zip(self.datasets, scaler)
+                ]
+
+    def reset(self):
+        return [dset.reset() for dset in self.datasets]
+
+    @property
+    def d_xd(self) -> list[int]:
+        return self.datasets[0].d_xd
+
+    @property
+    def d_vf(self) -> list[int]:
+        return sum(dset.d_vf for dset in self.datasets)
+
+    @property
+    def d_ef(self) -> list[int]:
+        return sum(dset.d_ef for dset in self.datasets)
+
+    @property
+    def d_vd(self) -> list[int]:
+        return sum(dset.d_vd for dset in self.datasets)

chemprop/data/molgraph.py

@@ -0,0 +1,16 @@
+from typing import NamedTuple
+
+import numpy as np
+
+
+class MolGraph(NamedTuple):
+    """A :class:`MolGraph` represents the graph featurization of a molecule."""
+
+    V: np.ndarray
+    """an array of shape ``V x d_v`` containing the atom features of the molecule"""
+    E: np.ndarray
+    """an array of shape ``E x d_e`` containing the bond features of the molecule"""
+    edge_index: np.ndarray
+    """an array of shape ``2 x E`` containing the edges of the graph in COO format"""
+    rev_edge_index: np.ndarray
+    """A array of shape ``E`` that maps from an edge index to the index of the source of the reverse edge in :attr:`edge_index` attribute."""

chemprop/data/samplers.py

@@ -0,0 +1,66 @@
+from itertools import chain
+from typing import Iterator, Optional
+
+import numpy as np
+from torch.utils.data import Sampler
+
+
+class SeededSampler(Sampler):
+    """A :class`SeededSampler` is a class for iterating through a dataset in a randomly seeded
+    fashion"""
+
+    def __init__(self, N: int, seed: int):
+        if seed is None:
+            raise ValueError("arg 'seed' was `None`! A SeededSampler must be seeded!")
+
+        self.idxs = np.arange(N)
+        self.rg = np.random.default_rng(seed)
+
+    def __iter__(self) -> Iterator[int]:
+        """an iterator over indices to sample."""
+        self.rg.shuffle(self.idxs)
+
+        return iter(self.idxs)
+
+    def __len__(self) -> int:
+        """the number of indices that will be sampled."""
+        return len(self.idxs)
+
+
+class ClassBalanceSampler(Sampler):
+    """A :class:`ClassBalanceSampler` samples data from a :class:`MolGraphDataset` such that
+    positive and negative classes are equally sampled
+
+    Parameters
+    ----------
+    dataset : MolGraphDataset
+        the dataset from which to sample
+    seed : int
+        the random seed to use for shuffling (only used when `shuffle` is `True`)
+    shuffle : bool, default=False
+        whether to shuffle the data during sampling
+    """
+
+    def __init__(self, Y: np.ndarray, seed: Optional[int] = None, shuffle: bool = False):
+        self.shuffle = shuffle
+        self.rg = np.random.default_rng(seed)
+
+        idxs = np.arange(len(Y))
+        actives = Y.any(1)
+
+        self.pos_idxs = idxs[actives]
+        self.neg_idxs = idxs[~actives]
+
+        self.length = 2 * min(len(self.pos_idxs), len(self.neg_idxs))
+
+    def __iter__(self) -> Iterator[int]:
+        """an iterator over indices to sample."""
+        if self.shuffle:
+            self.rg.shuffle(self.pos_idxs)
+            self.rg.shuffle(self.neg_idxs)
+
+        return chain(*zip(self.pos_idxs, self.neg_idxs))
+
+    def __len__(self) -> int:
+        """the number of indices that will be sampled."""
+        return self.length

chemprop/data/splitting.py

@@ -0,0 +1,257 @@
+import copy
+import logging
+from enum import auto
+from collections.abc import Sequence, Iterable
+import numpy as np
+from astartes import train_test_split, train_val_test_split
+from astartes.molecules import train_test_split_molecules, train_val_test_split_molecules
+from rdkit import Chem
+
+from chemprop.data.datapoints import MoleculeDatapoint, ReactionDatapoint
+from chemprop.utils.utils import EnumMapping
+
+logger = logging.getLogger(__name__)
+
+Datapoints = Sequence[MoleculeDatapoint] | Sequence[ReactionDatapoint]
+MulticomponentDatapoints = Sequence[Datapoints]
+
+
+class SplitType(EnumMapping):
+    CV_NO_VAL = auto()
+    CV = auto()
+    SCAFFOLD_BALANCED = auto()
+    RANDOM_WITH_REPEATED_SMILES = auto()
+    RANDOM = auto()
+    KENNARD_STONE = auto()
+    KMEANS = auto()
+
+
+def make_split_indices(
+    mols: Sequence[Chem.Mol],
+    split: SplitType | str = "random",
+    sizes: tuple[float, float, float] = (0.8, 0.1, 0.1),
+    seed: int = 0,
+    num_folds: int = 1,
+):
+    """Splits data into training, validation, and test splits.
+
+    Parameters
+    ----------
+    mols : Sequence[Chem.Mol]
+        Sequence of RDKit molecules to use for structure based splitting
+    split : SplitType | str, optional
+        Split type, one of ~chemprop.data.utils.SplitType, by default "random"
+    sizes : tuple[float, float, float], optional
+        3-tuple with the proportions of data in the train, validation, and test sets, by default
+        (0.8, 0.1, 0.1). Set the middle value to 0 for a two way split.
+    seed : int, optional
+        The random seed passed to astartes, by default 0
+    num_folds : int, optional
+        Number of folds to create (only needed for "cv" and "cv-no-test"), by default 1
+
+    Returns
+    -------
+    tuple[list[int], list[int], list[int]] | tuple[list[list[int], ...], list[list[int], ...], list[list[int], ...]]
+        A tuple of list of indices corresponding to the train, validation, and test splits of the
+        data. If the split type is "cv" or "cv-no-test", returns a tuple of lists of lists of
+        indices corresponding to the train, validation, and test splits of each fold.
+            .. important::
+                validation may or may not be present
+
+    Raises
+    ------
+    ValueError
+        Requested split sizes tuple not of length 3
+    ValueError
+        Innapropriate number of folds requested
+    ValueError
+        Unsupported split method requested
+    """
+    if (num_splits := len(sizes)) != 3:
+        raise ValueError(
+            f"Specify sizes for train, validation, and test (got {num_splits} values)."
+        )
+    # typically include a validation set
+    include_val = True
+    split_fun = train_val_test_split
+    mol_split_fun = train_val_test_split_molecules
+    # default sampling arguments for astartes sampler
+    astartes_kwargs = dict(
+        train_size=sizes[0], test_size=sizes[2], return_indices=True, random_state=seed
+    )
+    # if no validation set, reassign the splitting functions
+    if sizes[1] == 0.0:
+        include_val = False
+        split_fun = train_test_split
+        mol_split_fun = train_test_split_molecules
+    else:
+        astartes_kwargs["val_size"] = sizes[1]
+
+    n_datapoints = len(mols)
+    train, val, test = None, None, None
+    match SplitType.get(split):
+        case SplitType.CV_NO_VAL | SplitType.CV:
+            min_folds = 2 if SplitType.get(split) == SplitType.CV_NO_VAL else 3
+            if not (min_folds <= num_folds <= n_datapoints):
+                raise ValueError(
+                    f"invalid number of folds requested! got: {num_folds}, but expected between "
+                    f"{min_folds} and {n_datapoints} (i.e., number of datapoints), inclusive, "
+                    f"for split type: {repr(split)}"
+                )
+
+            # returns nested lists of indices
+            train, val, test = [], [], []
+            random = np.random.default_rng(seed)
+
+            indices = np.tile(np.arange(num_folds), 1 + n_datapoints // num_folds)[:n_datapoints]
+            random.shuffle(indices)
+
+            for fold_idx in range(num_folds):
+                test_index = fold_idx
+                val_index = (fold_idx + 1) % num_folds
+
+                if split != SplitType.CV_NO_VAL:
+                    i_val = np.where(indices == val_index)[0]
+                    i_test = np.where(indices == test_index)[0]
+                    i_train = np.where((indices != val_index) & (indices != test_index))[0]
+                else:
+                    i_val = []
+                    i_test = np.where(indices == test_index)[0]
+                    i_train = np.where(indices != test_index)[0]
+
+                train.append(i_train)
+                val.append(i_val)
+                test.append(i_test)
+
+        case SplitType.SCAFFOLD_BALANCED:
+            mols_without_atommaps = []
+            for mol in mols:
+                copied_mol = copy.deepcopy(mol)
+                for atom in copied_mol.GetAtoms():
+                    atom.SetAtomMapNum(0)
+                mols_without_atommaps.append(copied_mol)
+            result = mol_split_fun(
+                np.array(mols_without_atommaps), sampler="scaffold", **astartes_kwargs
+            )
+            train, val, test = _unpack_astartes_result(result, include_val)
+
+        # Use to constrain data with the same smiles go in the same split.
+        case SplitType.RANDOM_WITH_REPEATED_SMILES:
+            # get two arrays: one of all the smiles strings, one of just the unique
+            all_smiles = np.array([Chem.MolToSmiles(mol) for mol in mols])
+            unique_smiles = np.unique(all_smiles)
+
+            # save a mapping of smiles -> all the indices that it appeared at
+            smiles_indices = {}
+            for smiles in unique_smiles:
+                smiles_indices[smiles] = np.where(all_smiles == smiles)[0].tolist()
+
+            # randomly split the unique smiles
+            result = split_fun(np.arange(len(unique_smiles)), sampler="random", **astartes_kwargs)
+            train_idxs, val_idxs, test_idxs = _unpack_astartes_result(result, include_val)
+
+            # convert these to the 'actual' indices from the original list using the dict we made
+            train = sum((smiles_indices[unique_smiles[i]] for i in train_idxs), [])
+            val = sum((smiles_indices[unique_smiles[j]] for j in val_idxs), [])
+            test = sum((smiles_indices[unique_smiles[k]] for k in test_idxs), [])
+
+        case SplitType.RANDOM:
+            result = split_fun(np.arange(n_datapoints), sampler="random", **astartes_kwargs)
+            train, val, test = _unpack_astartes_result(result, include_val)
+
+        case SplitType.KENNARD_STONE:
+            result = mol_split_fun(
+                np.array(mols),
+                sampler="kennard_stone",
+                hopts=dict(metric="jaccard"),
+                fingerprint="morgan_fingerprint",
+                fprints_hopts=dict(n_bits=2048),
+                **astartes_kwargs,
+            )
+            train, val, test = _unpack_astartes_result(result, include_val)
+
+        case SplitType.KMEANS:
+            result = mol_split_fun(
+                np.array(mols),
+                sampler="kmeans",
+                hopts=dict(metric="jaccard"),
+                fingerprint="morgan_fingerprint",
+                fprints_hopts=dict(n_bits=2048),
+                **astartes_kwargs,
+            )
+            train, val, test = _unpack_astartes_result(result, include_val)
+
+        case _:
+            raise RuntimeError("Unreachable code reached!")
+
+    return train, val, test
+
+
+def _unpack_astartes_result(
+    result: tuple, include_val: bool
+) -> tuple[list[list[int]], list[list[int]], list[list[int]]]:
+    """Helper function to partition input data based on output of astartes sampler
+
+    Parameters
+    -----------
+    result: tuple
+        Output from call to astartes containing the split indices
+    include_val: bool
+        True if a validation set is included, False otherwise.
+
+    Returns
+    ---------
+    train: list[int]
+    val: list[int]
+    .. important::
+        validation possibly empty
+    test: list[int]
+    """
+    train_idxs, val_idxs, test_idxs = [], [], []
+    # astartes returns a set of lists containing the data, clusters (if applicable)
+    # and indices (always last), so we pull out the indices
+    if include_val:
+        train_idxs, val_idxs, test_idxs = result[-3], result[-2], result[-1]
+    else:
+        train_idxs, test_idxs = result[-2], result[-1]
+    return list(train_idxs), list(val_idxs), list(test_idxs)
+
+
+def split_data_by_indices(
+    data: Datapoints | MulticomponentDatapoints,
+    train_indices: Iterable[Iterable[int]] | Iterable[int] | None = None,
+    val_indices: Iterable[Iterable[int]] | Iterable[int] | None = None,
+    test_indices: Iterable[Iterable[int]] | Iterable[int] | None = None,
+):
+    """Splits data into training, validation, and test groups based on split indices given."""
+
+    train_data = _splitter_helper(data, train_indices) if train_indices is not None else None
+    val_data = _splitter_helper(data, val_indices) if val_indices is not None else None
+    test_data = _splitter_helper(data, test_indices) if test_indices is not None else None
+
+    return train_data, val_data, test_data
+
+
+def _splitter_helper(data, indices):
+    nested_component = not isinstance(data[0], (MoleculeDatapoint, ReactionDatapoint))
+    nested_split = isinstance(indices[0], Iterable)
+
+    match (nested_component, nested_split):
+        case (False, False):
+            datapoints = data
+            idxs = indices
+            return [datapoints[idx] for idx in idxs]
+        case (False, True):
+            datapoints = data
+            idxss = indices
+            return [[datapoints[idx] for idx in idxs] for idxs in idxss]
+        case (True, False):
+            datapointss = data
+            idxs = indices
+            return [[datapoints[idx] for idx in idxs] for datapoints in datapointss]
+        case (True, True):
+            datapointss = data
+            idxss = indices
+            return [
+                [[datapoints[idx] for idx in idxs] for datapoints in datapointss] for idxs in idxss
+            ]

chemprop/exceptions.py

@@ -0,0 +1,12 @@
+from typing import Iterable
+
+from chemprop.utils import pretty_shape
+
+
+class InvalidShapeError(ValueError):
+    def __init__(self, var_name: str, received: Iterable[int], expected: Iterable[int]):
+        message = (
+            f"arg '{var_name}' has incorrect shape! "
+            f"got: `{pretty_shape(received)}`. expected: `{pretty_shape(expected)}`"
+        )
+        super().__init__(message)

chemprop/featurizers/__init__.py

@@ -0,0 +1,46 @@
+from .base import Featurizer, S, T, VectorFeaturizer, GraphFeaturizer
+from .atom import MultiHotAtomFeaturizer, AtomFeatureMode, get_multi_hot_atom_featurizer
+from .bond import MultiHotBondFeaturizer
+from .molgraph import (
+    MolGraphCacheFacade,
+    MolGraphCache,
+    MolGraphCacheOnTheFly,
+    SimpleMoleculeMolGraphFeaturizer,
+    CondensedGraphOfReactionFeaturizer,
+    CGRFeaturizer,
+    RxnMode,
+)
+from .molecule import (
+    MorganFeaturizerMixin,
+    BinaryFeaturizerMixin,
+    CountFeaturizerMixin,
+    MorganBinaryFeaturizer,
+    MorganCountFeaturizer,
+    MoleculeFeaturizerRegistry,
+)
+
+__all__ = [
+    "Featurizer",
+    "S",
+    "T",
+    "VectorFeaturizer",
+    "GraphFeaturizer",
+    "MultiHotAtomFeaturizer",
+    "AtomFeatureMode",
+    "get_multi_hot_atom_featurizer",
+    "MultiHotBondFeaturizer",
+    "MolGraphCacheFacade",
+    "MolGraphCache",
+    "MolGraphCacheOnTheFly",
+    "SimpleMoleculeMolGraphFeaturizer",
+    "CondensedGraphOfReactionFeaturizer",
+    "CGRFeaturizer",
+    "RxnMode",
+    "MoleculeFeaturizer",
+    "MorganFeaturizerMixin",
+    "BinaryFeaturizerMixin",
+    "CountFeaturizerMixin",
+    "MorganBinaryFeaturizer",
+    "MorganCountFeaturizer",
+    "MoleculeFeaturizerRegistry",
+]

chemprop/featurizers/atom.py

@@ -0,0 +1,222 @@
+from typing import Sequence
+from enum import auto
+
+import numpy as np
+from rdkit.Chem.rdchem import Atom, HybridizationType
+
+from chemprop.utils.utils import EnumMapping
+from chemprop.featurizers.base import VectorFeaturizer
+
+
+class MultiHotAtomFeaturizer(VectorFeaturizer[Atom]):
+    """A :class:`MultiHotAtomFeaturizer` uses a multi-hot encoding to featurize atoms.
+
+    .. seealso::
+        The class provides three default parameterization schemes:
+
+        * :meth:`MultiHotAtomFeaturizer.v1`
+        * :meth:`MultiHotAtomFeaturizer.v2`
+        * :meth:`MultiHotAtomFeaturizer.organic`
+
+    The generated atom features are ordered as follows:
+    * atomic number
+    * degree
+    * formal charge
+    * chiral tag
+    * number of hydrogens
+    * hybridization
+    * aromaticity
+    * mass
+
+    .. important::
+        Each feature, except for aromaticity and mass, includes a pad for unknown values.
+
+    Parameters
+    ----------
+    atomic_nums : Sequence[int]
+        the choices for atom type denoted by atomic number. Ex: ``[4, 5, 6]`` for C, N and O.
+    degrees : Sequence[int]
+        the choices for number of bonds an atom is engaged in.
+    formal_charges : Sequence[int]
+        the choices for integer electronic charge assigned to an atom.
+    chiral_tags : Sequence[int]
+        the choices for an atom's chiral tag. See :class:`rdkit.Chem.rdchem.ChiralType` for possible integer values.
+    num_Hs : Sequence[int]
+        the choices for number of bonded hydrogen atoms.
+    hybridizations : Sequence[int]
+        the choices for an atom’s hybridization type. See :class:`rdkit.Chem.rdchem.HybridizationType` for possible integer values.
+    """
+
+    def __init__(
+        self,
+        atomic_nums: Sequence[int],
+        degrees: Sequence[int],
+        formal_charges: Sequence[int],
+        chiral_tags: Sequence[int],
+        num_Hs: Sequence[int],
+        hybridizations: Sequence[int],
+    ):
+        self.atomic_nums = {j: i for i, j in enumerate(atomic_nums)}
+        self.degrees = {i: i for i in degrees}
+        self.formal_charges = {j: i for i, j in enumerate(formal_charges)}
+        self.chiral_tags = {i: i for i in chiral_tags}
+        self.num_Hs = {i: i for i in num_Hs}
+        self.hybridizations = {ht: i for i, ht in enumerate(hybridizations)}
+
+        self._subfeats: list[dict] = [
+            self.atomic_nums,
+            self.degrees,
+            self.formal_charges,
+            self.chiral_tags,
+            self.num_Hs,
+            self.hybridizations,
+        ]
+        subfeat_sizes = [
+            1 + len(self.atomic_nums),
+            1 + len(self.degrees),
+            1 + len(self.formal_charges),
+            1 + len(self.chiral_tags),
+            1 + len(self.num_Hs),
+            1 + len(self.hybridizations),
+            1,
+            1,
+        ]
+        self.__size = sum(subfeat_sizes)
+
+    def __len__(self) -> int:
+        return self.__size
+
+    def __call__(self, a: Atom | None) -> np.ndarray:
+        x = np.zeros(self.__size)
+
+        if a is None:
+            return x
+
+        feats = [
+            a.GetAtomicNum(),
+            a.GetTotalDegree(),
+            a.GetFormalCharge(),
+            int(a.GetChiralTag()),
+            int(a.GetTotalNumHs()),
+            a.GetHybridization(),
+        ]
+        i = 0
+        for feat, choices in zip(feats, self._subfeats):
+            j = choices.get(feat, len(choices))
+            x[i + j] = 1
+            i += len(choices) + 1
+        x[i] = int(a.GetIsAromatic())
+        x[i + 1] = 0.01 * a.GetMass()
+
+        return x
+
+    def num_only(self, a: Atom) -> np.ndarray:
+        """featurize the atom by setting only the atomic number bit"""
+        x = np.zeros(len(self))
+
+        if a is None:
+            return x
+
+        i = self.atomic_nums.get(a.GetAtomicNum(), len(self.atomic_nums))
+        x[i] = 1
+
+        return x
+
+    @classmethod
+    def v1(cls, max_atomic_num: int = 100):
+        """The original implementation used in Chemprop V1 [1]_, [2]_.
+
+        Parameters
+        ----------
+        max_atomic_num : int, default=100
+            Include a bit for all atomic numbers in the interval :math:`[1, \mathtt{max_atomic_num}]`
+
+        References
+        -----------
+        .. [1] Yang, K.; Swanson, K.; Jin, W.; Coley, C.; Eiden, P.; Gao, H.; Guzman-Perez, A.; Hopper, T.;
+        Kelley, B.; Mathea, M.; Palmer, A. "Analyzing Learned Molecular Representations for Property Prediction."
+        J. Chem. Inf. Model. 2019, 59 (8), 3370–3388. https://doi.org/10.1021/acs.jcim.9b00237
+        .. [2] Heid, E.; Greenman, K.P.; Chung, Y.; Li, S.C.; Graff, D.E.; Vermeire, F.H.; Wu, H.; Green, W.H.; McGill,
+        C.J. "Chemprop: A machine learning package for chemical property prediction." J. Chem. Inf. Model. 2024,
+        64 (1), 9–17. https://doi.org/10.1021/acs.jcim.3c01250
+        """
+
+        return cls(
+            atomic_nums=list(range(1, max_atomic_num + 1)),
+            degrees=list(range(6)),
+            formal_charges=[-1, -2, 1, 2, 0],
+            chiral_tags=list(range(4)),
+            num_Hs=list(range(5)),
+            hybridizations=[
+                HybridizationType.SP,
+                HybridizationType.SP2,
+                HybridizationType.SP3,
+                HybridizationType.SP3D,
+                HybridizationType.SP3D2,
+            ],
+        )
+
+    @classmethod
+    def v2(cls):
+        """An implementation that includes an atom type bit for all elements in the first four rows of the periodic table plus iodine."""
+
+        return cls(
+            atomic_nums=list(range(1, 37)) + [53],
+            degrees=list(range(6)),
+            formal_charges=[-1, -2, 1, 2, 0],
+            chiral_tags=list(range(4)),
+            num_Hs=list(range(5)),
+            hybridizations=[
+                HybridizationType.S,
+                HybridizationType.SP,
+                HybridizationType.SP2,
+                HybridizationType.SP2D,
+                HybridizationType.SP3,
+                HybridizationType.SP3D,
+                HybridizationType.SP3D2,
+            ],
+        )
+
+    @classmethod
+    def organic(cls):
+        r"""A specific parameterization intended for use with organic or drug-like molecules.
+
+        This parameterization features:
+            1. includes an atomic number bit only for H, B, C, N, O, F, Si, P, S, Cl, Br, and I atoms
+            2. a hybridization bit for :math:`s, sp, sp^2` and :math:`sp^3` hybridizations.
+        """
+
+        return cls(
+            atomic_nums=[1, 5, 6, 7, 8, 9, 14, 15, 16, 17, 35, 53],
+            degrees=list(range(6)),
+            formal_charges=[-1, -2, 1, 2, 0],
+            chiral_tags=list(range(4)),
+            num_Hs=list(range(5)),
+            hybridizations=[
+                HybridizationType.S,
+                HybridizationType.SP,
+                HybridizationType.SP2,
+                HybridizationType.SP3,
+            ],
+        )
+
+
+class AtomFeatureMode(EnumMapping):
+    """The mode of an atom is used for featurization into a `MolGraph`"""
+
+    V1 = auto()
+    V2 = auto()
+    ORGANIC = auto()
+
+
+def get_multi_hot_atom_featurizer(mode: str | AtomFeatureMode) -> MultiHotAtomFeaturizer:
+    """Build the corresponding multi-hot atom featurizer."""
+    match AtomFeatureMode.get(mode):
+        case AtomFeatureMode.V1:
+            return MultiHotAtomFeaturizer.v1()
+        case AtomFeatureMode.V2:
+            return MultiHotAtomFeaturizer.v2()
+        case AtomFeatureMode.ORGANIC:
+            return MultiHotAtomFeaturizer.organic()
+        case _:
+            raise RuntimeError("unreachable code reached!")

chemprop/featurizers/base.py

@@ -0,0 +1,30 @@
+from abc import abstractmethod
+from collections.abc import Sized
+from typing import Generic, TypeVar
+
+import numpy as np
+
+from chemprop.data.molgraph import MolGraph
+
+S = TypeVar("S")
+T = TypeVar("T")
+
+
+class Featurizer(Generic[S, T]):
+    """An :class:`Featurizer` featurizes inputs type ``S`` into outputs of
+    type ``T``."""
+
+    @abstractmethod
+    def __call__(self, input: S, *args, **kwargs) -> T:
+        """featurize an input"""
+
+
+class VectorFeaturizer(Featurizer[S, np.ndarray], Sized):
+    ...
+
+
+class GraphFeaturizer(Featurizer[S, MolGraph]):
+    @property
+    @abstractmethod
+    def shape(self) -> tuple[int, int]:
+        ...

chemprop/featurizers/bond.py

@@ -0,0 +1,92 @@
+from typing import Sequence
+
+import numpy as np
+from rdkit.Chem.rdchem import Bond, BondType
+
+from chemprop.featurizers.base import VectorFeaturizer
+
+
+class MultiHotBondFeaturizer(VectorFeaturizer[Bond]):
+    """A :class:`MultiHotBondFeaturizer` feauturizes bonds based on the following attributes:
+
+    * ``null``-ity (i.e., is the bond ``None``?)
+    * bond type
+    * conjugated?
+    * in ring?
+    * stereochemistry
+
+    The feature vectors produced by this featurizer have the following (general) signature:
+
+    +---------------------+-----------------+--------------+
+    | slice [start, stop) | subfeature      | unknown pad? |
+    +=====================+=================+==============+
+    | 0-1                 | null?           | N            |
+    +---------------------+-----------------+--------------+
+    | 1-5                 | bond type       | N            |
+    +---------------------+-----------------+--------------+
+    | 5-6                 | conjugated?     | N            |
+    +---------------------+-----------------+--------------+
+    | 6-8                 | in ring?        | N            |
+    +---------------------+-----------------+--------------+
+    | 7-14                | stereochemistry | Y            |
+    +---------------------+-----------------+--------------+
+
+    **NOTE**: the above signature only applies for the default arguments, as the bond type and
+    sterochemistry slices can increase in size depending on the input arguments.
+
+    Parameters
+    ----------
+    bond_types : Sequence[BondType] | None, default=[SINGLE, DOUBLE, TRIPLE, AROMATIC]
+        the known bond types
+    stereos : Sequence[int] | None, default=[0, 1, 2, 3, 4, 5]
+        the known bond stereochemistries. See [1]_ for more details
+
+    References
+    ----------
+    .. [1] https://www.rdkit.org/docs/source/rdkit.Chem.rdchem.html#rdkit.Chem.rdchem.BondStereo.values
+    """
+
+    def __init__(
+        self, bond_types: Sequence[BondType] | None = None, stereos: Sequence[int] | None = None
+    ):
+        self.bond_types = bond_types or [
+            BondType.SINGLE,
+            BondType.DOUBLE,
+            BondType.TRIPLE,
+            BondType.AROMATIC,
+        ]
+        self.stereo = stereos or range(6)
+
+    def __len__(self):
+        return 1 + len(self.bond_types) + 2 + (len(self.stereo) + 1)
+
+    def __call__(self, b: Bond) -> np.ndarray:
+        x = np.zeros(len(self), int)
+
+        if b is None:
+            x[0] = 1
+            return x
+
+        i = 1
+        bond_type = b.GetBondType()
+        bt_bit, size = self.one_hot_index(bond_type, self.bond_types)
+        if bt_bit != size:
+            x[i + bt_bit] = 1
+        i += size - 1
+
+        x[i] = int(b.GetIsConjugated())
+        x[i + 1] = int(b.IsInRing())
+        i += 2
+
+        stereo_bit, _ = self.one_hot_index(int(b.GetStereo()), self.stereo)
+        x[i + stereo_bit] = 1
+
+        return x
+
+    @classmethod
+    def one_hot_index(cls, x, xs: Sequence) -> tuple[int, int]:
+        """Returns a tuple of the index of ``x`` in ``xs`` and ``len(xs) + 1`` if ``x`` is in ``xs``.
+        Otherwise, returns a tuple with ``len(xs)`` and ``len(xs) + 1``."""
+        n = len(xs)
+
+        return xs.index(x) if x in xs else n, n + 1

chemprop/featurizers/molecule.py

@@ -0,0 +1,43 @@
+import numpy as np
+from rdkit import Chem
+from rdkit.Chem import Mol
+from rdkit.Chem.rdFingerprintGenerator import GetMorganGenerator
+
+from chemprop.featurizers.base import VectorFeaturizer
+from chemprop.utils import ClassRegistry
+
+MoleculeFeaturizerRegistry = ClassRegistry[VectorFeaturizer[Mol]]()
+
+
+class MorganFeaturizerMixin:
+    def __init__(self, radius: int = 2, length: int = 2048, include_chirality: bool = True):
+        if radius < 0:
+            raise ValueError(f"arg 'radius' must be >= 0! got: {radius}")
+
+        self.length = length
+        self.F = GetMorganGenerator(
+            radius=radius, fpSize=length, includeChirality=include_chirality
+        )
+
+    def __len__(self) -> int:
+        return self.length
+
+
+class BinaryFeaturizerMixin:
+    def __call__(self, mol: Chem.Mol) -> np.ndarray:
+        return self.F.GetFingerprintAsNumPy(mol)
+
+
+class CountFeaturizerMixin:
+    def __call__(self, mol: Chem.Mol) -> np.ndarray:
+        return self.F.GetCountFingerprintAsNumPy(mol).astype(np.int32)
+
+
+@MoleculeFeaturizerRegistry("morgan_binary")
+class MorganBinaryFeaturizer(MorganFeaturizerMixin, BinaryFeaturizerMixin, VectorFeaturizer[Mol]):
+    pass
+
+
+@MoleculeFeaturizerRegistry("morgan_count")
+class MorganCountFeaturizer(MorganFeaturizerMixin, CountFeaturizerMixin, VectorFeaturizer[Mol]):
+    pass

chemprop/featurizers/molgraph/__init__.py

@@ -0,0 +1,13 @@
+from .cache import MolGraphCacheFacade, MolGraphCache, MolGraphCacheOnTheFly
+from .molecule import SimpleMoleculeMolGraphFeaturizer
+from .reaction import CondensedGraphOfReactionFeaturizer, CGRFeaturizer, RxnMode
+
+__all__ = [
+    "MolGraphCacheFacade",
+    "MolGraphCache",
+    "MolGraphCacheOnTheFly",
+    "SimpleMoleculeMolGraphFeaturizer",
+    "CondensedGraphOfReactionFeaturizer",
+    "CGRFeaturizer",
+    "RxnMode",
+]

chemprop/featurizers/molgraph/cache.py

@@ -0,0 +1,89 @@
+from abc import abstractmethod
+from collections.abc import Sequence
+from typing import Generic, Iterable
+
+import numpy as np
+
+from chemprop.featurizers.base import S, Featurizer
+from chemprop.data.molgraph import MolGraph
+
+
+class MolGraphCacheFacade(Sequence[MolGraph], Generic[S]):
+    """
+    A :class:`MolGraphCacheFacade` provided an interface for caching
+    :class:`~chemprop.data.molgraph.MolGraph`\s.
+
+    .. note::
+        This class only provides a facade for a cached dataset, but it _does not guarantee_
+        whether the underlying data is truly cached.
+
+
+    Parameters
+    ----------
+    inputs : Iterable[S]
+        The inputs to be featurized.
+    V_fs : Iterable[np.ndarray]
+        The node features for each input.
+    E_fs : Iterable[np.ndarray]
+        The edge features for each input.
+    featurizer : Featurizer[S, MolGraph]
+        The featurizer with which to generate the
+        :class:`~chemprop.data.molgraph.MolGraph`\s.
+    """
+
+    @abstractmethod
+    def __init__(
+        self,
+        inputs: Iterable[S],
+        V_fs: Iterable[np.ndarray],
+        E_fs: Iterable[np.ndarray],
+        featurizer: Featurizer[S, MolGraph],
+    ):
+        pass
+
+
+class MolGraphCache(MolGraphCacheFacade):
+    """
+    A :class:`MolGraphCache` precomputes the corresponding
+    :class:`~chemprop.data.molgraph.MolGraph`\s and caches them in memory.
+    """
+
+    def __init__(
+        self,
+        inputs: Iterable[S],
+        V_fs: Iterable[np.ndarray | None],
+        E_fs: Iterable[np.ndarray | None],
+        featurizer: Featurizer[S, MolGraph],
+    ):
+        self._mgs = [featurizer(input, V_f, E_f) for input, V_f, E_f in zip(inputs, V_fs, E_fs)]
+
+    def __len__(self) -> int:
+        return len(self._mgs)
+
+    def __getitem__(self, index: int) -> MolGraph:
+        return self._mgs[index]
+
+
+class MolGraphCacheOnTheFly(MolGraphCacheFacade):
+    """
+    A :class:`MolGraphCacheOnTheFly` computes the corresponding
+    :class:`~chemprop.data.molgraph.MolGraph`\s as they are requested.
+    """
+
+    def __init__(
+        self,
+        inputs: Iterable[S],
+        V_fs: Iterable[np.ndarray | None],
+        E_fs: Iterable[np.ndarray | None],
+        featurizer: Featurizer[S, MolGraph],
+    ):
+        self._inputs = list(inputs)
+        self._V_fs = list(V_fs)
+        self._E_fs = list(E_fs)
+        self._featurizer = featurizer
+
+    def __len__(self) -> int:
+        return len(self._inputs)
+
+    def __getitem__(self, index: int) -> MolGraph:
+        return self._featurizer(self._inputs[index], self._V_fs[index], self._E_fs[index])

chemprop/featurizers/molgraph/mixins.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass, field
+
+from rdkit.Chem.rdchem import Atom, Bond
+
+from chemprop.featurizers.base import VectorFeaturizer
+from chemprop.featurizers.atom import MultiHotAtomFeaturizer
+from chemprop.featurizers.bond import MultiHotBondFeaturizer
+
+
+@dataclass
+class _MolGraphFeaturizerMixin:
+    atom_featurizer: VectorFeaturizer[Atom] = field(default_factory=MultiHotAtomFeaturizer.v2)
+    bond_featurizer: VectorFeaturizer[Bond] = field(default_factory=MultiHotBondFeaturizer)
+
+    def __post_init__(self):
+        self.atom_fdim = len(self.atom_featurizer)
+        self.bond_fdim = len(self.bond_featurizer)
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        """the feature dimension of the atoms and bonds, respectively, of `MolGraph`s generated by
+        this featurizer"""
+        return self.atom_fdim, self.bond_fdim

chemprop/featurizers/molgraph/molecule.py

@@ -0,0 +1,95 @@
+from dataclasses import InitVar, dataclass
+
+import numpy as np
+from rdkit import Chem
+from rdkit.Chem import Mol
+
+from chemprop.data.molgraph import MolGraph
+from chemprop.featurizers.base import GraphFeaturizer
+from chemprop.featurizers.molgraph.mixins import _MolGraphFeaturizerMixin
+
+
+@dataclass
+class SimpleMoleculeMolGraphFeaturizer(_MolGraphFeaturizerMixin, GraphFeaturizer[Mol]):
+    """A :class:`SimpleMoleculeMolGraphFeaturizer` is the default implementation of a
+    :class:`MoleculeMolGraphFeaturizer`
+
+    Parameters
+    ----------
+    atom_featurizer : AtomFeaturizer, default=MultiHotAtomFeaturizer()
+        the featurizer with which to calculate feature representations of the atoms in a given
+        molecule
+    bond_featurizer : BondFeaturizer, default=MultiHotBondFeaturizer()
+        the featurizer with which to calculate feature representations of the bonds in a given
+        molecule
+    extra_atom_fdim : int, default=0
+        the dimension of the additional features that will be concatenated onto the calculated
+        features of each atom
+    extra_bond_fdim : int, default=0
+        the dimension of the additional features that will be concatenated onto the calculated
+        features of each bond
+    """
+
+    extra_atom_fdim: InitVar[int] = 0
+    extra_bond_fdim: InitVar[int] = 0
+
+    def __post_init__(self, extra_atom_fdim: int = 0, extra_bond_fdim: int = 0):
+        super().__post_init__()
+
+        self.extra_atom_fdim = extra_atom_fdim
+        self.extra_bond_fdim = extra_bond_fdim
+        self.atom_fdim += self.extra_atom_fdim
+        self.bond_fdim += self.extra_bond_fdim
+
+    def __call__(
+        self,
+        mol: Chem.Mol,
+        atom_features_extra: np.ndarray | None = None,
+        bond_features_extra: np.ndarray | None = None,
+    ) -> MolGraph:
+        n_atoms = mol.GetNumAtoms()
+        n_bonds = mol.GetNumBonds()
+
+        if atom_features_extra is not None and len(atom_features_extra) != n_atoms:
+            raise ValueError(
+                "Input molecule must have same number of atoms as `len(atom_features_extra)`!"
+                f"got: {n_atoms} and {len(atom_features_extra)}, respectively"
+            )
+        if bond_features_extra is not None and len(bond_features_extra) != n_bonds:
+            raise ValueError(
+                "Input molecule must have same number of bonds as `len(bond_features_extra)`!"
+                f"got: {n_bonds} and {len(bond_features_extra)}, respectively"
+            )
+
+        if n_atoms == 0:
+            V = np.zeros((1, self.atom_fdim), dtype=np.single)
+        else:
+            V = np.array([self.atom_featurizer(a) for a in mol.GetAtoms()], dtype=np.single)
+        E = np.empty((2 * n_bonds, self.bond_fdim))
+        edge_index = [[], []]
+
+        if atom_features_extra is not None:
+            V = np.hstack((V, atom_features_extra))
+
+        i = 0
+        for u in range(n_atoms):
+            for v in range(u + 1, n_atoms):
+                bond = mol.GetBondBetweenAtoms(u, v)
+                if bond is None:
+                    continue
+
+                x_e = self.bond_featurizer(bond)
+                if bond_features_extra is not None:
+                    x_e = np.concatenate((x_e, bond_features_extra[bond.GetIdx()]), dtype=np.single)
+
+                E[i : i + 2] = x_e
+
+                edge_index[0].extend([u, v])
+                edge_index[1].extend([v, u])
+
+                i += 2
+
+        rev_edge_index = np.arange(len(E)).reshape(-1, 2)[:, ::-1].ravel()
+        edge_index = np.array(edge_index, int)
+
+        return MolGraph(V, E, edge_index, rev_edge_index)

chemprop/featurizers/molgraph/reaction.py

@@ -0,0 +1,333 @@
+from dataclasses import InitVar, dataclass
+from enum import auto
+from typing import Iterable, Sequence, TypeAlias
+import warnings
+
+import numpy as np
+from rdkit import Chem
+from rdkit.Chem.rdchem import Bond, Mol
+from chemprop.featurizers.base import GraphFeaturizer
+
+from chemprop.types import Rxn
+from chemprop.data.molgraph import MolGraph
+from chemprop.featurizers.molgraph.mixins import _MolGraphFeaturizerMixin
+from chemprop.utils.utils import EnumMapping
+
+
+class RxnMode(EnumMapping):
+    """The mode by which a reaction should be featurized into a `MolGraph`"""
+
+    REAC_PROD = auto()
+    """concatenate the reactant features with the product features."""
+    REAC_PROD_BALANCE = auto()
+    """concatenate the reactant features with the products feature and balances imbalanced
+    reactions"""
+    REAC_DIFF = auto()
+    """concatenates the reactant features with the difference in features between reactants and
+    products"""
+    REAC_DIFF_BALANCE = auto()
+    """concatenates the reactant features with the difference in features between reactants and
+    product and balances imbalanced reactions"""
+    PROD_DIFF = auto()
+    """concatenates the product features with the difference in features between reactants and
+    products"""
+    PROD_DIFF_BALANCE = auto()
+    """concatenates the product features with the difference in features between reactants and
+    products and balances imbalanced reactions"""
+
+
+@dataclass
+class CondensedGraphOfReactionFeaturizer(_MolGraphFeaturizerMixin, GraphFeaturizer[Rxn]):
+    """A :class:`CondensedGraphOfReactionFeaturizer` featurizes reactions using the condensed
+    reaction graph method utilized in [1]_
+
+    **NOTE**: This class *does not* accept a :class:`AtomFeaturizer` instance. This is because
+    it requries the :meth:`num_only()` method, which is only implemented in the concrete
+    :class:`AtomFeaturizer` class
+
+    Parameters
+    ----------
+    atom_featurizer : AtomFeaturizer, default=AtomFeaturizer()
+        the featurizer with which to calculate feature representations of the atoms in a given
+        molecule
+    bond_featurizer : BondFeaturizerBase, default=BondFeaturizer()
+        the featurizer with which to calculate feature representations of the bonds in a given
+        molecule
+    mode_ : Union[str, ReactionMode], default=ReactionMode.REAC_DIFF
+        the mode by which to featurize the reaction as either the string code or enum value
+
+    References
+    ----------
+    .. [1] Heid, E.; Green, W.H. "Machine Learning of Reaction Properties via Learned
+        Representations of the Condensed Graph of Reaction." J. Chem. Inf. Model. 2022, 62,
+        2101-2110. https://doi.org/10.1021/acs.jcim.1c00975
+    """
+
+    mode_: InitVar[str | RxnMode] = RxnMode.REAC_DIFF
+
+    def __post_init__(self, mode_: str | RxnMode):
+        super().__post_init__()
+
+        self.mode = mode_
+        self.atom_fdim += len(self.atom_featurizer) - len(self.atom_featurizer.atomic_nums) - 1
+        self.bond_fdim *= 2
+
+    @property
+    def mode(self) -> RxnMode:
+        return self.__mode
+
+    @mode.setter
+    def mode(self, m: str | RxnMode):
+        self.__mode = RxnMode.get(m)
+
+    def __call__(
+        self,
+        rxn: tuple[Chem.Mol, Chem.Mol],
+        atom_features_extra: np.ndarray | None = None,
+        bond_features_extra: np.ndarray | None = None,
+    ) -> MolGraph:
+        """Featurize the input reaction into a molecular graph
+
+        Parameters
+        ----------
+        rxn : Rxn
+            a 2-tuple of atom-mapped rdkit molecules, where the 0th element is the reactant and the
+            1st element is the product
+        atom_features_extra : np.ndarray | None, default=None
+            *UNSUPPORTED* maintained only to maintain parity with the method signature of the
+            `MoleculeFeaturizer`
+        bond_features_extra : np.ndarray | None, default=None
+            *UNSUPPORTED* maintained only to maintain parity with the method signature of the
+            `MoleculeFeaturizer`
+
+        Returns
+        -------
+        MolGraph
+            the molecular graph of the reaction
+        """
+
+        if atom_features_extra is not None:
+            warnings.warn("'atom_features_extra' is currently unsupported for reactions")
+        if bond_features_extra is not None:
+            warnings.warn("'bond_features_extra' is currently unsupported for reactions")
+
+        reac, pdt = rxn
+        r2p_idx_map, pdt_idxs, reac_idxs = self.map_reac_to_prod(reac, pdt)
+
+        V = self._calc_node_feature_matrix(reac, pdt, r2p_idx_map, pdt_idxs, reac_idxs)
+        E = []
+        edge_index = [[], []]
+
+        n_atoms_tot = len(V)
+        n_atoms_reac = reac.GetNumAtoms()
+
+        i = 0
+        for u in range(n_atoms_tot):
+            for v in range(u + 1, n_atoms_tot):
+                b_reac, b_prod = self._get_bonds(
+                    reac, pdt, r2p_idx_map, pdt_idxs, n_atoms_reac, u, v
+                )
+                if b_reac is None and b_prod is None:
+                    continue
+
+                x_e = self._calc_edge_feature(b_reac, b_prod)
+                E.extend([x_e, x_e])
+                edge_index[0].extend([u, v])
+                edge_index[1].extend([v, u])
+
+                i += 2
+
+        E = np.array(E)
+        rev_edge_index = np.arange(len(E)).reshape(-1, 2)[:, ::-1].ravel()
+        edge_index = np.array(edge_index, int)
+
+        return MolGraph(V, E, edge_index, rev_edge_index)
+
+    def _calc_node_feature_matrix(
+        self,
+        rct: Mol,
+        pdt: Mol,
+        r2p_idx_map: dict[int, int],
+        pdt_idxs: Iterable[int],
+        reac_idxs: Iterable[int],
+    ) -> np.ndarray:
+        """Calculate the node feature matrix for the reaction"""
+        X_v_r1 = np.array([self.atom_featurizer(a) for a in rct.GetAtoms()])
+        X_v_p2 = np.array([self.atom_featurizer(pdt.GetAtomWithIdx(i)) for i in pdt_idxs])
+        X_v_p2 = X_v_p2.reshape(-1, X_v_r1.shape[1])
+
+        if self.mode in [RxnMode.REAC_DIFF, RxnMode.PROD_DIFF, RxnMode.REAC_PROD]:
+            # Reactant:
+            # (1) regular features for each atom in the reactants
+            # (2) zero features for each atom that's only in the products
+            X_v_r2 = [self.atom_featurizer.num_only(pdt.GetAtomWithIdx(i)) for i in pdt_idxs]
+            X_v_r2 = np.array(X_v_r2).reshape(-1, X_v_r1.shape[1])
+
+            # Product:
+            # (1) either (a) product-side features for each atom in both
+            #         or (b) zero features for each atom only in the reatants
+            # (2) regular features for each atom only in the products
+            X_v_p1 = np.array(
+                [
+                    (
+                        self.atom_featurizer(pdt.GetAtomWithIdx(r2p_idx_map[a.GetIdx()]))
+                        if a.GetIdx() not in reac_idxs
+                        else self.atom_featurizer.num_only(a)
+                    )
+                    for a in rct.GetAtoms()
+                ]
+            )
+        else:
+            # Reactant:
+            # (1) regular features for each atom in the reactants
+            # (2) regular features for each atom only in the products
+            X_v_r2 = [self.atom_featurizer(pdt.GetAtomWithIdx(i)) for i in pdt_idxs]
+            X_v_r2 = np.array(X_v_r2).reshape(-1, X_v_r1.shape[1])
+
+            # Product:
+            # (1) either (a) product-side features for each atom in both
+            #         or (b) reactant-side features for each atom only in the reatants
+            # (2) regular features for each atom only in the products
+            X_v_p1 = np.array(
+                [
+                    (
+                        self.atom_featurizer(pdt.GetAtomWithIdx(r2p_idx_map[a.GetIdx()]))
+                        if a.GetIdx() not in reac_idxs
+                        else self.atom_featurizer(a)
+                    )
+                    for a in rct.GetAtoms()
+                ]
+            )
+
+        X_v_r = np.concatenate((X_v_r1, X_v_r2))
+        X_v_p = np.concatenate((X_v_p1, X_v_p2))
+
+        m = min(len(X_v_r), len(X_v_p))
+
+        if self.mode in [RxnMode.REAC_PROD, RxnMode.REAC_PROD_BALANCE]:
+            X_v = np.hstack((X_v_r[:m], X_v_p[:m, len(self.atom_featurizer.atomic_nums) + 1 :]))
+        else:
+            X_v_d = X_v_p[:m] - X_v_r[:m]
+            if self.mode in [RxnMode.REAC_DIFF, RxnMode.REAC_DIFF_BALANCE]:
+                X_v = np.hstack((X_v_r[:m], X_v_d[:m, len(self.atom_featurizer.atomic_nums) + 1 :]))
+            else:
+                X_v = np.hstack((X_v_p[:m], X_v_d[:m, len(self.atom_featurizer.atomic_nums) + 1 :]))
+
+        return X_v
+
+    def _get_bonds(
+        self,
+        rct: Bond,
+        pdt: Bond,
+        ri2pj: dict[int, int],
+        pids: Sequence[int],
+        n_atoms_r: int,
+        u: int,
+        v: int,
+    ) -> tuple[Bond, Bond]:
+        """get the corresponding reactant- and product-side bond, respectively, betweeen atoms `u` and `v`"""
+        if u >= n_atoms_r and v >= n_atoms_r:
+            b_prod = pdt.GetBondBetweenAtoms(pids[u - n_atoms_r], pids[v - n_atoms_r])
+
+            if self.mode in [
+                RxnMode.REAC_PROD_BALANCE,
+                RxnMode.REAC_DIFF_BALANCE,
+                RxnMode.PROD_DIFF_BALANCE,
+            ]:
+                b_reac = b_prod
+            else:
+                b_reac = None
+        elif u < n_atoms_r and v >= n_atoms_r:  # One atom only in product
+            b_reac = None
+
+            if u in ri2pj:
+                b_prod = pdt.GetBondBetweenAtoms(ri2pj[u], pids[v - n_atoms_r])
+            else:  # Atom atom only in reactant, the other only in product
+                b_prod = None
+        else:
+            b_reac = rct.GetBondBetweenAtoms(u, v)
+
+            if u in ri2pj and v in ri2pj:  # Both atoms in both reactant and product
+                b_prod = pdt.GetBondBetweenAtoms(ri2pj[u], ri2pj[v])
+            elif self.mode in [
+                RxnMode.REAC_PROD_BALANCE,
+                RxnMode.REAC_DIFF_BALANCE,
+                RxnMode.PROD_DIFF_BALANCE,
+            ]:
+                b_prod = None if (u in ri2pj or v in ri2pj) else b_reac
+            else:  # One or both atoms only in reactant
+                b_prod = None
+
+        return b_reac, b_prod
+
+    def _calc_edge_feature(self, b_reac: Bond, b_pdt: Bond):
+        """Calculate the global features of the two bonds"""
+        x_e_r = self.bond_featurizer(b_reac)
+        x_e_p = self.bond_featurizer(b_pdt)
+        x_e_d = x_e_p - x_e_r
+
+        if self.mode in [RxnMode.REAC_PROD, RxnMode.REAC_PROD_BALANCE]:
+            x_e = np.hstack((x_e_r, x_e_p))
+        elif self.mode in [RxnMode.REAC_DIFF, RxnMode.REAC_DIFF_BALANCE]:
+            x_e = np.hstack((x_e_r, x_e_d))
+        else:
+            x_e = np.hstack((x_e_p, x_e_d))
+
+        return x_e
+
+    @classmethod
+    def map_reac_to_prod(
+        cls, reacs: Chem.Mol, pdts: Chem.Mol
+    ) -> tuple[dict[int, int], list[int], list[int]]:
+        """Map atom indices between corresponding atoms in the reactant and product molecules
+
+        Parameters
+        ----------
+        reacs : Chem.Mol
+            An RDKit molecule of the reactants
+        pdts : Chem.Mol
+            An RDKit molecule of the products
+
+        Returns
+        -------
+        ri2pi : dict[int, int]
+            A dictionary of corresponding atom indices from reactant atoms to product atoms
+        pdt_idxs : list[int]
+            atom indices of poduct atoms
+        rct_idxs : list[int]
+            atom indices of reactant atoms
+        """
+        pdt_idxs = []
+        mapno2pj = {}
+        reac_atommap_nums = {a.GetAtomMapNum() for a in reacs.GetAtoms()}
+
+        for a in pdts.GetAtoms():
+            map_num = a.GetAtomMapNum()
+            j = a.GetIdx()
+
+            if map_num > 0:
+                mapno2pj[map_num] = j
+                if map_num not in reac_atommap_nums:
+                    pdt_idxs.append(j)
+            else:
+                pdt_idxs.append(j)
+
+        rct_idxs = []
+        r2p_idx_map = {}
+
+        for a in reacs.GetAtoms():
+            map_num = a.GetAtomMapNum()
+            i = a.GetIdx()
+
+            if map_num > 0:
+                try:
+                    r2p_idx_map[i] = mapno2pj[map_num]
+                except KeyError:
+                    rct_idxs.append(i)
+            else:
+                rct_idxs.append(i)
+
+        return r2p_idx_map, pdt_idxs, rct_idxs
+
+
+CGRFeaturizer: TypeAlias = CondensedGraphOfReactionFeaturizer

chemprop/models/__init__.py

@@ -0,0 +1,5 @@
+from .model import MPNN
+from .multi import MulticomponentMPNN
+from .utils import load_model, save_model
+
+__all__ = ["MPNN", "MulticomponentMPNN", "load_model", "save_model"]

chemprop/models/model.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+from typing import Iterable
+
+from lightning import pytorch as pl
+import torch
+from torch import nn, Tensor, optim
+
+from chemprop.data import TrainingBatch, BatchMolGraph
+from chemprop.nn.metrics import Metric
+from chemprop.nn import MessagePassing, Aggregation, Predictor, LossFunction
+from chemprop.schedulers import NoamLR
+from chemprop.nn.transforms import ScaleTransform
+
+
+class MPNN(pl.LightningModule):
+    r"""An :class:`MPNN` is a sequence of message passing layers, an aggregation routine, and a
+    predictor routine.
+
+    The first two modules calculate learned fingerprints from an input molecule
+    reaction graph, and the final module takes these learned fingerprints as input to calculate a
+    final prediction. I.e., the following operation:
+
+    .. math::
+        \mathtt{MPNN}(\mathcal{G}) =
+            \mathtt{predictor}(\mathtt{agg}(\mathtt{message\_passing}(\mathcal{G})))
+
+    The full model is trained end-to-end.
+
+    Parameters
+    ----------
+    message_passing : MessagePassing
+        the message passing block to use to calculate learned fingerprints
+    agg : Aggregation
+        the aggregation operation to use during molecule-level predictor
+    predictor : Predictor
+        the function to use to calculate the final prediction
+    batch_norm : bool, default=True
+        if `True`, apply batch normalization to the output of the aggregation operation
+    metrics : Iterable[Metric] | None, default=None
+        the metrics to use to evaluate the model during training and evaluation
+    warmup_epochs : int, default=2
+        the number of epochs to use for the learning rate warmup
+    init_lr : int, default=1e-4
+        the initial learning rate
+    max_lr : float, default=1e-3
+        the maximum learning rate
+    final_lr : float, default=1e-4
+        the final learning rate
+
+    Raises
+    ------
+    ValueError
+        if the output dimension of the message passing block does not match the input dimension of
+        the predictor function
+    """
+
+    def __init__(
+        self,
+        message_passing: MessagePassing,
+        agg: Aggregation,
+        predictor: Predictor,
+        batch_norm: bool = True,
+        metrics: Iterable[Metric] | None = None,
+        warmup_epochs: int = 2,
+        init_lr: float = 1e-4,
+        max_lr: float = 1e-3,
+        final_lr: float = 1e-4,
+        X_d_transform: ScaleTransform | None = None,
+    ):
+        super().__init__()
+
+        self.save_hyperparameters(ignore=["message_passing", "agg", "predictor"])
+        self.hparams.update(
+            {
+                "message_passing": message_passing.hparams,
+                "agg": agg.hparams,
+                "predictor": predictor.hparams,
+            }
+        )
+
+        self.message_passing = message_passing
+        self.agg = agg
+        self.bn = nn.BatchNorm1d(self.message_passing.output_dim) if batch_norm else nn.Identity()
+        self.predictor = predictor
+
+        self.X_d_transform = X_d_transform if X_d_transform is not None else nn.Identity()
+
+        self.metrics = (
+            [*metrics, self.criterion]
+            if metrics
+            else [self.predictor._T_default_metric(), self.criterion]
+        )
+
+        self.warmup_epochs = warmup_epochs
+        self.init_lr = init_lr
+        self.max_lr = max_lr
+        self.final_lr = final_lr
+
+    @property
+    def output_dim(self) -> int:
+        return self.predictor.output_dim
+
+    @property
+    def n_tasks(self) -> int:
+        return self.predictor.n_tasks
+
+    @property
+    def n_targets(self) -> int:
+        return self.predictor.n_targets
+
+    @property
+    def criterion(self) -> LossFunction:
+        return self.predictor.criterion
+
+    def fingerprint(
+        self, bmg: BatchMolGraph, V_d: Tensor | None = None, X_d: Tensor | None = None
+    ) -> Tensor:
+        """the learned fingerprints for the input molecules"""
+        H_v = self.message_passing(bmg, V_d)
+        H = self.agg(H_v, bmg.batch)
+        H = self.bn(H)
+
+        return H if X_d is None else torch.cat((H, self.X_d_transform(X_d)), 1)
+
+    def encoding(
+        self, bmg: BatchMolGraph, V_d: Tensor | None = None, X_d: Tensor | None = None, i: int = -1
+    ) -> Tensor:
+        """Calculate the :attr:`i`-th hidden representation"""
+        return self.predictor.encode(self.fingerprint(bmg, V_d, X_d), i)
+
+    def forward(
+        self, bmg: BatchMolGraph, V_d: Tensor | None = None, X_d: Tensor | None = None
+    ) -> Tensor:
+        """Generate predictions for the input molecules/reactions"""
+        return self.predictor(self.fingerprint(bmg, V_d, X_d))
+
+    def training_step(self, batch: TrainingBatch, batch_idx):
+        bmg, V_d, X_d, targets, weights, lt_mask, gt_mask = batch
+
+        mask = targets.isfinite()
+        targets = targets.nan_to_num(nan=0.0)
+
+        Z = self.fingerprint(bmg, V_d, X_d)
+        preds = self.predictor.train_step(Z)
+        l = self.criterion(preds, targets, mask, weights, lt_mask, gt_mask)
+
+        self.log("train_loss", l, prog_bar=True)
+
+        return l
+
+    def on_validation_model_eval(self) -> None:
+        self.eval()
+        self.predictor.output_transform.train()
+
+    def validation_step(self, batch: TrainingBatch, batch_idx: int = 0):
+        losses = self._evaluate_batch(batch)
+        metric2loss = {f"val/{m.alias}": l for m, l in zip(self.metrics, losses)}
+
+        self.log_dict(metric2loss, batch_size=len(batch[0]))
+        self.log("val_loss", losses[0], batch_size=len(batch[0]), prog_bar=True)
+
+    def test_step(self, batch: TrainingBatch, batch_idx: int = 0):
+        losses = self._evaluate_batch(batch)
+        metric2loss = {f"batch_averaged_test/{m.alias}": l for m, l in zip(self.metrics, losses)}
+
+        self.log_dict(metric2loss, batch_size=len(batch[0]))
+
+    def _evaluate_batch(self, batch) -> list[Tensor]:
+        bmg, V_d, X_d, targets, _, lt_mask, gt_mask = batch
+
+        mask = targets.isfinite()
+        targets = targets.nan_to_num(nan=0.0)
+        preds = self(bmg, V_d, X_d)
+
+        return [
+            metric(preds, targets, mask, None, lt_mask, gt_mask) for metric in self.metrics[:-1]
+        ]
+
+    def predict_step(self, batch: TrainingBatch, batch_idx: int, dataloader_idx: int = 0) -> Tensor:
+        """Return the predictions of the input batch
+
+        Parameters
+        ----------
+        batch : TrainingBatch
+            the input batch
+
+        Returns
+        -------
+        Tensor
+            a tensor of varying shape depending on the task type:
+
+            * regression/binary classification: ``n x (t * s)``, where ``n`` is the number of input
+            molecules/reactions, ``t`` is the number of tasks, and ``s`` is the number of targets
+            per task. The final dimension is flattened, so that the targets for each task are
+            grouped. I.e., the first ``t`` elements are the first target for each task, the second
+            ``t`` elements the second target, etc.
+            * multiclass classification: ``n x t x c``, where ``c`` is the number of classes
+        """
+        bmg, X_vd, X_d, *_ = batch
+
+        return self(bmg, X_vd, X_d)
+
+    def configure_optimizers(self):
+        opt = optim.Adam(self.parameters(), self.init_lr)
+
+        lr_sched = NoamLR(
+            opt,
+            self.warmup_epochs,
+            self.trainer.max_epochs,
+            self.trainer.estimated_stepping_batches // self.trainer.max_epochs,
+            self.init_lr,
+            self.max_lr,
+            self.final_lr,
+        )
+        lr_sched_config = {
+            "scheduler": lr_sched,
+            "interval": "step" if isinstance(lr_sched, NoamLR) else "batch",
+        }
+
+        return {"optimizer": opt, "lr_scheduler": lr_sched_config}
+
+    @classmethod
+    def load_submodules(cls, checkpoint_path, **kwargs):
+        hparams = torch.load(checkpoint_path)["hyper_parameters"]
+
+        kwargs |= {
+            key: hparams[key].pop("cls")(**hparams[key])
+            for key in ("message_passing", "agg", "predictor")
+            if key not in kwargs
+        }
+        return kwargs
+
+    @classmethod
+    def load_from_checkpoint(
+        cls, checkpoint_path, map_location=None, hparams_file=None, strict=True, **kwargs
+    ) -> MPNN:
+        kwargs = cls.load_submodules(checkpoint_path, **kwargs)
+        return super().load_from_checkpoint(
+            checkpoint_path, map_location, hparams_file, strict, **kwargs
+        )
+
+    @classmethod
+    def load_from_file(cls, model_path, map_location=None, strict=True) -> MPNN:
+        d = torch.load(model_path, map_location=map_location)
+
+        try:
+            hparams = d["hyper_parameters"]
+            state_dict = d["state_dict"]
+        except KeyError:
+            raise KeyError(f"Could not find hyper parameters and/or state dict in {model_path}. ")
+
+        for key in ["message_passing", "agg", "predictor"]:
+            hparam_kwargs = hparams[key]
+            hparam_cls = hparam_kwargs.pop("cls")
+            hparams[key] = hparam_cls(**hparam_kwargs)
+
+        model = cls(**hparams)
+        model.load_state_dict(state_dict, strict=strict)
+
+        return model

chemprop/models/multi.py

@@ -0,0 +1,92 @@
+from typing import Iterable
+
+import torch
+from torch import Tensor
+
+from chemprop.data import BatchMolGraph
+from chemprop.nn import MulticomponentMessagePassing, Aggregation, Predictor
+from chemprop.models.model import MPNN
+from chemprop.nn.metrics import Metric
+from chemprop.nn.transforms import ScaleTransform
+
+
+class MulticomponentMPNN(MPNN):
+    def __init__(
+        self,
+        message_passing: MulticomponentMessagePassing,
+        agg: Aggregation,
+        predictor: Predictor,
+        batch_norm: bool = True,
+        metrics: Iterable[Metric] | None = None,
+        warmup_epochs: int = 2,
+        init_lr: float = 1e-4,
+        max_lr: float = 1e-3,
+        final_lr: float = 1e-4,
+        X_d_transform: ScaleTransform | None = None,
+    ):
+        super().__init__(
+            message_passing,
+            agg,
+            predictor,
+            batch_norm,
+            metrics,
+            warmup_epochs,
+            init_lr,
+            max_lr,
+            final_lr,
+            X_d_transform,
+        )
+        self.message_passing: MulticomponentMessagePassing
+
+    def fingerprint(
+        self,
+        bmgs: Iterable[BatchMolGraph],
+        V_ds: Iterable[Tensor | None],
+        X_d: Tensor | None = None,
+    ) -> Tensor:
+        H_vs: list[Tensor] = self.message_passing(bmgs, V_ds)
+        Hs = [self.agg(H_v, bmg.batch) for H_v, bmg in zip(H_vs, bmgs)]
+        H = torch.cat(Hs, 1)
+        H = self.bn(H)
+
+        return H if X_d is None else torch.cat((H, self.X_d_transform(X_d)), 1)
+
+    @classmethod
+    def load_submodules(cls, checkpoint_path, **kwargs):
+        hparams = torch.load(checkpoint_path)["hyper_parameters"]
+
+        hparams["message_passing"]["blocks"] = [
+            block_hparams.pop("cls")(**block_hparams)
+            for block_hparams in hparams["message_passing"]["blocks"]
+        ]
+        kwargs |= {
+            key: hparams[key].pop("cls")(**hparams[key])
+            for key in ("message_passing", "agg", "predictor")
+            if key not in kwargs
+        }
+        return kwargs
+
+    @classmethod
+    def load_from_file(cls, model_path, map_location=None, strict=True) -> MPNN:
+        d = torch.load(model_path, map_location=map_location)
+
+        try:
+            hparams = d["hyper_parameters"]
+            state_dict = d["state_dict"]
+        except KeyError:
+            raise KeyError(f"Could not find hyper parameters and/or state dict in {model_path}. ")
+
+        for key in ["message_passing", "agg", "predictor"]:
+            hparam_kwargs = hparams[key]
+            if key == "message_passing":
+                hparam_kwargs["blocks"] = [
+                    block_hparams.pop("cls")(**block_hparams)
+                    for block_hparams in hparam_kwargs["blocks"]
+                ]
+            hparam_cls = hparam_kwargs.pop("cls")
+            hparams[key] = hparam_cls(**hparam_kwargs)
+
+        model = cls(**hparams)
+        model.load_state_dict(state_dict, strict=strict)
+
+        return model

chemprop/models/utils.py

@@ -0,0 +1,18 @@
+from os import PathLike
+import torch
+
+from chemprop.models.model import MPNN
+from chemprop.models.multi import MulticomponentMPNN
+
+
+def save_model(path: PathLike, model: MPNN) -> None:
+    torch.save({"hyper_parameters": model.hparams, "state_dict": model.state_dict()}, path)
+
+
+def load_model(path: PathLike, multicomponent: bool) -> MPNN:
+    if multicomponent:
+        model = MulticomponentMPNN.load_from_file(path)
+    else:
+        model = MPNN.load_from_file(path)
+
+    return model

chemprop/nn/__init__.py

@@ -0,0 +1,133 @@
+from .agg import (
+    Aggregation,
+    AggregationRegistry,
+    MeanAggregation,
+    SumAggregation,
+    NormAggregation,
+    AttentiveAggregation,
+)
+from .loss import (
+    LossFunction,
+    LossFunctionRegistry,
+    MSELoss,
+    BoundedMSELoss,
+    MVELoss,
+    EvidentialLoss,
+    BCELoss,
+    CrossEntropyLoss,
+    MccMixin,
+    BinaryMCCLoss,
+    MulticlassMCCLoss,
+    DirichletMixin,
+    BinaryDirichletLoss,
+    MulticlassDirichletLoss,
+    SIDLoss,
+    WassersteinLoss,
+)
+from .metrics import (
+    Metric,
+    MetricRegistry,
+    ThresholdedMixin,
+    MAEMetric,
+    MSEMetric,
+    RMSEMetric,
+    BoundedMixin,
+    BoundedMAEMetric,
+    BoundedMSEMetric,
+    BoundedRMSEMetric,
+    R2Metric,
+    BinaryAUROCMetric,
+    BinaryAUPRCMetric,
+    BinaryAccuracyMetric,
+    BinaryF1Metric,
+    BCEMetric,
+    CrossEntropyMetric,
+    BinaryMCCMetric,
+    MulticlassMCCMetric,
+    SIDMetric,
+    WassersteinMetric,
+)
+from .message_passing import (
+    MessagePassing,
+    AtomMessagePassing,
+    BondMessagePassing,
+    MulticomponentMessagePassing,
+)
+from .predictors import (
+    Predictor,
+    PredictorRegistry,
+    RegressionFFN,
+    MveFFN,
+    EvidentialFFN,
+    BinaryClassificationFFNBase,
+    BinaryClassificationFFN,
+    BinaryDirichletFFN,
+    MulticlassClassificationFFN,
+    MulticlassDirichletFFN,
+    SpectralFFN,
+)
+from .utils import Activation
+from .transforms import UnscaleTransform
+
+__all__ = [
+    "Aggregation",
+    "AggregationRegistry",
+    "MeanAggregation",
+    "SumAggregation",
+    "NormAggregation",
+    "AttentiveAggregation",
+    "LossFunction",
+    "LossFunctionRegistry",
+    "MSELoss",
+    "BoundedMSELoss",
+    "MVELoss",
+    "EvidentialLoss",
+    "BCELoss",
+    "CrossEntropyLoss",
+    "MccMixin",
+    "BinaryMCCLoss",
+    "MulticlassMCCLoss",
+    "DirichletMixin",
+    "BinaryDirichletLoss",
+    "MulticlassDirichletLoss",
+    "SIDLoss",
+    "WassersteinLoss",
+    "Metric",
+    "MetricRegistry",
+    "ThresholdedMixin",
+    "MAEMetric",
+    "MSEMetric",
+    "RMSEMetric",
+    "BoundedMixin",
+    "BoundedMAEMetric",
+    "BoundedMSEMetric",
+    "BoundedRMSEMetric",
+    "R2Metric",
+    "BinaryAUROCMetric",
+    "BinaryAUPRCMetric",
+    "BinaryAccuracyMetric",
+    "BinaryF1Metric",
+    "BCEMetric",
+    "CrossEntropyMetric",
+    "BinaryMCCMetric",
+    "MulticlassMCCMetric",
+    "SIDMetric",
+    "WassersteinMetric",
+    "MessagePassing",
+    "AtomMessagePassing",
+    "BondMessagePassing",
+    "MulticomponentMessagePassing",
+    "Predictor",
+    "PredictorRegistry",
+    "RegressionFFN",
+    "MveFFN",
+    "EvidentialFFN",
+    "BinaryClassificationFFNBase",
+    "BinaryClassificationFFN",
+    "BinaryDirichletFFN",
+    "MulticlassClassificationFFN",
+    "MulticlassDirichletFFN",
+    "SpectralFFN",
+    "Activation",
+    "UnscaleTransform",
+]

chemprop/nn/agg.py

@@ -0,0 +1,132 @@
+from abc import abstractmethod
+import torch
+from torch import Tensor, nn
+
+from chemprop.utils import ClassRegistry
+from chemprop.nn.hparams import HasHParams
+
+
+__all__ = [
+    "Aggregation",
+    "AggregationRegistry",
+    "MeanAggregation",
+    "SumAggregation",
+    "NormAggregation",
+    "AttentiveAggregation",
+]
+
+
+class Aggregation(nn.Module, HasHParams):
+    """An :class:`Aggregation` aggregates the node-level representations of a batch of graphs into
+    a batch of graph-level representations
+
+    .. note::
+        this class is abstract and cannot be instantiated.
+
+    See also
+    --------
+    :class:`~chemprop.v2.models.modules.agg.MeanAggregation`
+    :class:`~chemprop.v2.models.modules.agg.SumAggregation`
+    :class:`~chemprop.v2.models.modules.agg.NormAggregation`
+    """
+
+    def __init__(self, dim: int = 0, *args, **kwargs):
+        super().__init__()
+
+        self.dim = dim
+        self.hparams = {"dim": dim, "cls": self.__class__}
+
+    @abstractmethod
+    def forward(self, H: Tensor, batch: Tensor) -> Tensor:
+        """Aggregate the graph-level representations of a batch of graphs into their respective
+        global representations
+
+        NOTE: it is possible for a graph to have 0 nodes. In this case, the representation will be
+        a zero vector of length `d` in the final output.
+
+        Parameters
+        ----------
+        H : Tensor
+            a tensor of shape ``V x d`` containing the batched node-level representations of ``b``
+            graphs
+        batch : Tensor
+            a tensor of shape ``V`` containing the index of the graph a given vertex corresponds to
+
+        Returns
+        -------
+        Tensor
+            a tensor of shape ``b x d`` containing the graph-level representations
+        """
+
+
+AggregationRegistry = ClassRegistry[Aggregation]()
+
+
+@AggregationRegistry.register("mean")
+class MeanAggregation(Aggregation):
+    r"""Average the graph-level representation:
+
+    .. math::
+        \mathbf h = \frac{1}{|V|} \sum_{v \in V} \mathbf h_v
+    """
+
+    def forward(self, H: Tensor, batch: Tensor) -> Tensor:
+        index_torch = batch.unsqueeze(1).repeat(1, H.shape[1])
+        dim_size = batch.max().int() + 1
+        return torch.zeros(dim_size, H.shape[1], dtype=H.dtype, device=H.device).scatter_reduce_(
+            self.dim, index_torch, H, reduce="mean", include_self=False
+        )
+
+
+@AggregationRegistry.register("sum")
+class SumAggregation(Aggregation):
+    r"""Sum the graph-level representation:
+
+    .. math::
+        \mathbf h = \sum_{v \in V} \mathbf h_v
+
+    """
+
+    def forward(self, H: Tensor, batch: Tensor) -> Tensor:
+        index_torch = batch.unsqueeze(1).repeat(1, H.shape[1])
+        dim_size = batch.max().int() + 1
+        return torch.zeros(dim_size, H.shape[1], dtype=H.dtype, device=H.device).scatter_reduce_(
+            self.dim, index_torch, H, reduce="sum", include_self=False
+        )
+
+
+@AggregationRegistry.register("norm")
+class NormAggregation(SumAggregation):
+    r"""Sum the graph-level representation and divide by a normalization constant:
+
+    .. math::
+        \mathbf h = \frac{1}{c} \sum_{v \in V} \mathbf h_v
+    """
+
+    def __init__(self, dim: int = 0, *args, norm: float = 100.0, **kwargs):
+        super().__init__(dim, **kwargs)
+
+        self.norm = norm
+        self.hparams["norm"] = norm
+
+    def forward(self, H: Tensor, batch: Tensor) -> Tensor:
+        return super().forward(H, batch) / self.norm
+
+
+class AttentiveAggregation(Aggregation):
+    def __init__(self, dim: int = 0, *args, output_size: int, **kwargs):
+        super().__init__(dim, *args, **kwargs)
+
+        self.W = nn.Linear(output_size, 1)
+
+    def forward(self, H: Tensor, batch: Tensor) -> Tensor:
+        dim_size = batch.max().int() + 1
+        attention_logits = self.W(H).exp()
+        Z = torch.zeros(dim_size, 1, dtype=H.dtype, device=H.device).scatter_reduce_(
+            self.dim, batch.unsqueeze(1), attention_logits, reduce="sum", include_self=False
+        )
+        alphas = attention_logits / Z[batch]
+        index_torch = batch.unsqueeze(1).repeat(1, H.shape[1])
+        return torch.zeros(dim_size, H.shape[1], dtype=H.dtype, device=H.device).scatter_reduce_(
+            self.dim, index_torch, alphas * H, reduce="sum", include_self=False
+        )

chemprop/nn/ffn.py

@@ -0,0 +1,63 @@
+from abc import abstractmethod
+
+from torch import nn, Tensor
+
+from chemprop.nn.utils import get_activation_function
+
+
+class FFN(nn.Module):
+    r"""A :class:`FFN` is a differentiable function
+    :math:`f_\theta : \mathbb R^i \mapsto \mathbb R^o`"""
+
+    input_dim: int
+    output_dim: int
+
+    @abstractmethod
+    def forward(self, X: Tensor) -> Tensor:
+        pass
+
+
+class MLP(nn.Sequential, FFN):
+    r"""An :class:`MLP` is an FFN that implements the following function:
+
+    .. math::
+        \mathbf h_0 &= \mathbf W_0 \mathbf x \,+ \mathbf b_{0} \\
+        \mathbf h_l &= \mathbf W_l \left( \mathtt{dropout} \left( \sigma ( \,\mathbf h_{l-1}\, ) \right) \right) + \mathbf b_l\\
+
+    where :math:`\mathbf x` is the input tensor, :math:`\mathbf W_l` and :math:`\mathbf b_l`
+    are the learned weight matrix and bias, respectively, of the :math:`l`-th layer,
+    :math:`\mathbf h_l` is the hidden representation after layer :math:`l`, and :math:`\sigma`
+    is the activation function.
+    """
+
+    @classmethod
+    def build(
+        cls,
+        input_dim: int,
+        output_dim: int,
+        hidden_dim: int = 300,
+        n_layers: int = 1,
+        dropout: float = 0.0,
+        activation: str = "relu",
+    ):
+        dropout = nn.Dropout(dropout)
+        act = get_activation_function(activation)
+        dims = [input_dim] + [hidden_dim] * n_layers + [output_dim]
+        blocks = [nn.Sequential(nn.Linear(dims[0], dims[1]))]
+        if len(dims) > 2:
+            blocks.extend(
+                [
+                    nn.Sequential(act, dropout, nn.Linear(d1, d2))
+                    for d1, d2 in zip(dims[1:-1], dims[2:])
+                ]
+            )
+
+        return cls(*blocks)
+
+    @property
+    def input_dim(self) -> int:
+        return self[0][-1].in_features
+
+    @property
+    def output_dim(self) -> int:
+        return self[-1][-1].out_features

chemprop/nn/hparams.py

@@ -0,0 +1,38 @@
+from typing import Protocol, Type, TypedDict
+
+
+class HParamsDict(TypedDict):
+    """A dictionary containing a module's class and it's hyperparameters
+
+    Using this type should essentially allow for initializing a module via::
+
+        module = hparams.pop('cls')(**hparams)
+    """
+
+    cls: Type
+
+
+class HasHParams(Protocol):
+    """:class:`HasHParams` is a protocol for clases which possess an :attr:`hparams` attribute which is a dictionary containing the object's class and arguments required to initialize it.
+
+    That is, any object which implements :class:`HasHParams` should be able to be initialized via::
+
+        class Foo(HasHParams):
+            def __init__(self, *args, **kwargs):
+                ...
+
+        foo1 = Foo(...)
+        foo1_cls = foo1.hparams['cls']
+        foo1_kwargs = {k: v for k, v in foo1.hparams.items() if k != "cls"}
+        foo2 = foo1_cls(**foo1_kwargs)
+        # code to compare foo1 and foo2 goes here and they should be equal
+    """
+
+    hparams: HParamsDict
+
+
+def from_hparams(hparams: HParamsDict):
+    cls = hparams["cls"]
+    kwargs = {k: v for k, v in hparams.items() if k != "cls"}
+
+    return cls(**kwargs)

chemprop/nn/loss.py

@@ -0,0 +1,344 @@
+from abc import abstractmethod
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+from numpy.typing import ArrayLike
+
+from chemprop.utils import ClassRegistry
+
+
+__all__ = [
+    "LossFunction",
+    "LossFunctionRegistry",
+    "MSELoss",
+    "BoundedMSELoss",
+    "MVELoss",
+    "EvidentialLoss",
+    "BCELoss",
+    "CrossEntropyLoss",
+    "MccMixin",
+    "BinaryMCCLoss",
+    "MulticlassMCCLoss",
+    "DirichletMixin",
+    "BinaryDirichletLoss",
+    "MulticlassDirichletLoss",
+    "SIDLoss",
+    "WassersteinLoss",
+]
+
+
+class LossFunction(nn.Module):
+    def __init__(self, task_weights: ArrayLike = 1.0):
+        """
+        Parameters
+        ----------
+        task_weights :  ArrayLike, default=1.0
+            the per-task weights of shape `t` or `1 x t`. Defaults to all tasks having a weight of 1.
+        """
+        super().__init__()
+        task_weights = torch.as_tensor(task_weights, dtype=torch.float).view(1, -1)
+        self.register_buffer("task_weights", task_weights)
+
+    def forward(
+        self,
+        preds: Tensor,
+        targets: Tensor,
+        mask: Tensor,
+        weights: Tensor,
+        lt_mask: Tensor,
+        gt_mask: Tensor,
+    ):
+        """Calculate the mean loss function value given predicted and target values
+
+        Parameters
+        ----------
+        preds : Tensor
+            a tensor of shape `b x (t * s)` (regression), `b x t` (binary classification), or
+            `b x t x c` (multiclass classification) containing the predictions, where `b` is the
+            batch size, `t` is the number of tasks to predict, `s` is the number of
+            targets to predict for each task, and `c` is the number of classes.
+        targets : Tensor
+            a float tensor of shape `b x t` containing the target values
+        mask : Tensor
+            a boolean tensor of shape `b x t` indicating whether the given prediction should be
+            included in the loss calculation
+        weights : Tensor
+            a tensor of shape `b` or `b x 1` containing the per-sample weight
+        lt_mask: Tensor
+        gt_mask: Tensor
+
+        Returns
+        -------
+        Tensor
+            a scalar containing the fully reduced loss
+        """
+        L = self._calc_unreduced_loss(preds, targets, mask, weights, lt_mask, gt_mask)
+        L = L * weights.view(-1, 1) * self.task_weights.view(1, -1) * mask
+
+        return L.sum() / mask.sum()
+
+    @abstractmethod
+    def _calc_unreduced_loss(self, preds, targets, mask, weights, lt_mask, gt_mask) -> Tensor:
+        """Calculate a tensor of shape `b x t` containing the unreduced loss values."""
+
+    def extra_repr(self) -> str:
+        return f"task_weights={self.task_weights.tolist()}"
+
+
+LossFunctionRegistry = ClassRegistry[LossFunction]()
+
+
+@LossFunctionRegistry.register("mse")
+class MSELoss(LossFunction):
+    def _calc_unreduced_loss(self, preds: Tensor, targets: Tensor, *args) -> Tensor:
+        return F.mse_loss(preds, targets, reduction="none")
+
+
+@LossFunctionRegistry.register("bounded-mse")
+class BoundedMSELoss(MSELoss):
+    def _calc_unreduced_loss(
+        self, preds: Tensor, targets: Tensor, mask, weights, lt_mask: Tensor, gt_mask: Tensor
+    ) -> Tensor:
+        preds = torch.where((preds < targets) & lt_mask, targets, preds)
+        preds = torch.where((preds > targets) & gt_mask, targets, preds)
+
+        return super()._calc_unreduced_loss(preds, targets)
+
+
+@LossFunctionRegistry.register("mve")
+class MVELoss(LossFunction):
+    """Calculate the loss using Eq. 9 from [nix1994]_
+
+    References
+    ----------
+    .. [nix1994] Nix, D. A.; Weigend, A. S. "Estimating the mean and variance of the target
+        probability distribution." Proceedings of 1994 IEEE International Conference on Neural
+        Networks, 1994 https://doi.org/10.1109/icnn.1994.374138
+    """
+
+    def _calc_unreduced_loss(self, preds: Tensor, targets: Tensor, *args) -> Tensor:
+        mean, var = torch.chunk(preds, 2, 1)
+
+        L_sos = (mean - targets) ** 2 / (2 * var)
+        L_kl = (2 * torch.pi * var).log() / 2
+
+        return L_sos + L_kl
+
+
+@LossFunctionRegistry.register("evidential")
+class EvidentialLoss(LossFunction):
+    """Calculate the loss using Eqs. 8, 9, and 10 from [amini2020]_
+
+    References
+    ----------
+    .. [amini2020] Amini, A; Schwarting, W.; Soleimany, A.; Rus, D.;
+        "Deep Evidential Regression" Advances in Neural Information Processing Systems;2020; Vol.33.
+        https://proceedings.neurips.cc/paper_files/paper/2020/file/aab085461de182608ee9f607f3f7d18f-Paper.pdf
+    .. [soleimany2021] Soleimany, A.P.; Amini, A.; Goldman, S.; Rus, D.; Bhatia, S.N.; Coley, C.W.;
+        "Evidential Deep Learning for Guided Molecular Property Prediction and Discovery." ACS
+        Cent. Sci. 2021, 7, 8, 1356-1367. https://doi.org/10.1021/acscentsci.1c00546
+    """
+
+    def __init__(self, task_weights: Tensor | None = None, v_kl: float = 0.2, eps: float = 1e-8):
+        super().__init__(task_weights)
+        self.v_kl = v_kl
+        self.eps = eps
+
+    def _calc_unreduced_loss(self, preds: Tensor, targets: Tensor, *args) -> Tensor:
+        mean, v, alpha, beta = torch.chunk(preds, 4, 1)
+
+        residuals = targets - mean
+        twoBlambda = 2 * beta * (1 + v)
+
+        L_nll = (
+            0.5 * (torch.pi / v).log()
+            - alpha * twoBlambda.log()
+            + (alpha + 0.5) * torch.log(v * residuals**2 + twoBlambda)
+            + torch.lgamma(alpha)
+            - torch.lgamma(alpha + 0.5)
+        )
+
+        L_reg = (2 * v + alpha) * residuals.abs()
+
+        return L_nll + self.v_kl * (L_reg - self.eps)
+
+    def extra_repr(self) -> str:
+        parent_repr = super().extra_repr()
+        return parent_repr + f", v_kl={self.v_kl}, eps={self.eps}"
+
+
+@LossFunctionRegistry.register("bce")
+class BCELoss(LossFunction):
+    def _calc_unreduced_loss(self, preds: Tensor, targets: Tensor, *args) -> Tensor:
+        return F.binary_cross_entropy_with_logits(preds, targets, reduction="none")
+
+
+@LossFunctionRegistry.register("ce")
+class CrossEntropyLoss(LossFunction):
+    def _calc_unreduced_loss(self, preds: Tensor, targets: Tensor, *args) -> Tensor:
+        preds = preds.transpose(1, 2)
+        targets = targets.long()
+
+        return F.cross_entropy(preds, targets, reduction="none")
+
+
+class MccMixin:
+    """Calculate a soft Matthews correlation coefficient ([mccWiki]_) loss for multiclass
+    classification based on the implementataion of [mccSklearn]_
+
+    References
+    ----------
+    .. [mccWiki] https://en.wikipedia.org/wiki/Phi_coefficient#Multiclass_case
+    .. [mccSklearn] https://scikit-learn.org/stable/modules/generated/sklearn.metrics.matthews_corrcoef.html
+    """
+
+    def __call__(self, preds: Tensor, targets: Tensor, mask: Tensor, weights: Tensor, *args):
+        if not (0 <= preds.min() and preds.max() <= 1):  # assume logits
+            preds = preds.softmax(2)
+
+        L = self._calc_unreduced_loss(preds, targets.long(), mask, weights, *args)
+        L = L * self.task_weights
+
+        return L.mean()
+
+
+@LossFunctionRegistry.register("binary-mcc")
+class BinaryMCCLoss(LossFunction, MccMixin):
+    def _calc_unreduced_loss(self, preds, targets, mask, weights, *args) -> Tensor:
+        TP = (targets * preds * weights * mask).sum(0, keepdim=True)
+        FP = ((1 - targets) * preds * weights * mask).sum(0, keepdim=True)
+        TN = ((1 - targets) * (1 - preds) * weights * mask).sum(0, keepdim=True)
+        FN = (targets * (1 - preds) * weights * mask).sum(0, keepdim=True)
+
+        MCC = (TP * TN - FP * FN) / ((TP + FP) * (TP + FN) * (TN + FP) * (TN + FN)).sqrt()
+
+        return 1 - MCC
+
+
+@LossFunctionRegistry.register("multiclass-mcc")
+class MulticlassMCCLoss(LossFunction, MccMixin):
+    def _calc_unreduced_loss(self, preds, targets, mask, weights, *args) -> Tensor:
+        device = preds.device
+
+        C = preds.shape[2]
+        bin_targets = torch.eye(C, device=device)[targets]
+        bin_preds = torch.eye(C, device=device)[preds.argmax(-1)]
+        masked_data_weights = weights.unsqueeze(2) * mask.unsqueeze(2)
+
+        p = (bin_preds * masked_data_weights).sum(0)
+        t = (bin_targets * masked_data_weights).sum(0)
+        c = (bin_preds * bin_targets * masked_data_weights).sum()
+        s = (preds * masked_data_weights).sum()
+        s2 = s.square()
+
+        # the `einsum` calls amount to calculating the batched dot product
+        cov_ytyp = c * s - torch.einsum("ij,ij->i", p, t).sum()
+        cov_ypyp = s2 - torch.einsum("ij,ij->i", p, p).sum()
+        cov_ytyt = s2 - torch.einsum("ij,ij->i", t, t).sum()
+
+        x = cov_ypyp * cov_ytyt
+        MCC = torch.tensor(0.0, device=device) if x == 0 else cov_ytyp / x.sqrt()
+
+        return 1 - MCC
+
+
+class DirichletMixin:
+    """Uses the loss function from [sensoy2018]_ based on the implementation at [sensoyGithub]_
+
+    References
+    ----------
+    .. [sensoy2018] Sensoy, M.; Kaplan, L.; Kandemir, M. "Evidential deep learning to quantify
+        classification uncertainty." NeurIPS, 2018, 31. https://doi.org/10.48550/arXiv.1806.01768
+    .. [sensoyGithub] https://muratsensoy.github.io/uncertainty.html#Define-the-loss-function
+    """
+
+    def __init__(self, task_weights: Tensor | None = None, v_kl: float = 0.2):
+        super().__init__(task_weights)
+        self.v_kl = v_kl
+
+    def _calc_unreduced_loss(self, preds, targets, *args) -> Tensor:
+        S = preds.sum(-1, keepdim=True)
+        p = preds / S
+
+        A = (targets - p).square().sum(-1, keepdim=True)
+        B = ((p * (1 - p)) / (S + 1)).sum(-1, keepdim=True)
+
+        L_mse = A + B
+
+        alpha = targets + (1 - targets) * preds
+        beta = torch.ones_like(alpha)
+        S_alpha = alpha.sum(-1, keepdim=True)
+        S_beta = beta.sum(-1, keepdim=True)
+
+        ln_alpha = S_alpha.lgamma() - alpha.lgamma().sum(-1, keepdim=True)
+        ln_beta = beta.lgamma().sum(-1, keepdim=True) - S_beta.lgamma()
+
+        dg0 = torch.digamma(alpha)
+        dg1 = torch.digamma(S_alpha)
+
+        L_kl = ln_alpha + ln_beta + torch.sum((alpha - beta) * (dg0 - dg1), -1, keepdim=True)
+
+        return (L_mse + self.v_kl * L_kl).mean(-1)
+
+    def extra_repr(self) -> str:
+        return f"v_kl={self.v_kl}"
+
+
+@LossFunctionRegistry.register("binary-dirichlet")
+class BinaryDirichletLoss(DirichletMixin, LossFunction):
+    def _calc_unreduced_loss(self, preds: Tensor, targets: Tensor, *args) -> Tensor:
+        N_CLASSES = 2
+        n_tasks = targets.shape[1]
+        preds = preds.reshape(len(preds), n_tasks, N_CLASSES)
+        y_one_hot = torch.eye(N_CLASSES, device=preds.device)[targets.long()]
+
+        return super()._calc_unreduced_loss(preds, y_one_hot, *args)
+
+
+@LossFunctionRegistry.register("multiclass-dirichlet")
+class MulticlassDirichletLoss(DirichletMixin, LossFunction):
+    def _calc_unreduced_loss(self, preds: Tensor, targets: Tensor, mask: Tensor, *args) -> Tensor:
+        y_one_hot = torch.eye(preds.shape[2], device=preds.device)[targets.long()]
+
+        return super()._calc_unreduced_loss(preds, y_one_hot, mask)
+
+
+@LossFunctionRegistry.register("sid")
+class SIDLoss(LossFunction):
+    def __init__(self, task_weights: Tensor | None = None, threshold: float | None = None):
+        super().__init__(task_weights)
+
+        self.threshold = threshold
+
+    def _calc_unreduced_loss(self, preds: Tensor, targets: Tensor, mask: Tensor, *args) -> Tensor:
+        if self.threshold is not None:
+            preds = preds.clamp(min=self.threshold)
+
+        preds_norm = preds / (preds * mask).sum(1, keepdim=True)
+
+        targets = targets.masked_fill(~mask, 1)
+        preds_norm = preds_norm.masked_fill(~mask, 1)
+
+        return (preds_norm / targets).log() * preds_norm + (targets / preds_norm).log() * targets
+
+    def extra_repr(self) -> str:
+        return f"threshold={self.threshold}"
+
+
+@LossFunctionRegistry.register(["earthmovers", "wasserstein"])
+class WassersteinLoss(LossFunction):
+    def __init__(self, task_weights: Tensor | None = None, threshold: float | None = None):
+        super().__init__(task_weights)
+
+        self.threshold = threshold
+
+    def _calc_unreduced_loss(self, preds: Tensor, targets: Tensor, mask: Tensor, *args) -> Tensor:
+        if self.threshold is not None:
+            preds = preds.clamp(min=self.threshold)
+
+        preds_norm = preds / (preds * mask).sum(1, keepdim=True)
+
+        return (targets.cumsum(1) - preds_norm.cumsum(1)).abs()
+
+    def extra_repr(self) -> str:
+        return f"threshold={self.threshold}"

chemprop/nn/message_passing/__init__.py

@@ -0,0 +1,10 @@
+from .proto import MessagePassing
+from .base import AtomMessagePassing, BondMessagePassing
+from .multi import MulticomponentMessagePassing
+
+__all__ = [
+    "MessagePassing",
+    "AtomMessagePassing",
+    "BondMessagePassing",
+    "MulticomponentMessagePassing",
+]

chemprop/nn/message_passing/base.py

@@ -0,0 +1,311 @@
+from abc import abstractmethod
+
+from lightning.pytorch.core.mixins import HyperparametersMixin
+import torch
+from torch import Tensor, nn
+
+from chemprop.conf import DEFAULT_ATOM_FDIM, DEFAULT_BOND_FDIM, DEFAULT_HIDDEN_DIM
+from chemprop.exceptions import InvalidShapeError
+from chemprop.data import BatchMolGraph
+from chemprop.nn.utils import Activation, get_activation_function
+from chemprop.nn.message_passing.proto import MessagePassing
+from chemprop.nn.transforms import ScaleTransform, GraphTransform
+
+
+class _MessagePassingBase(MessagePassing, HyperparametersMixin):
+    """The base message-passing block for atom- and bond-based message-passing schemes
+
+    NOTE: this class is an abstract base class and cannot be instantiated
+
+    Parameters
+    ----------
+    d_v : int, default=DEFAULT_ATOM_FDIM
+        the feature dimension of the vertices
+    d_e : int, default=DEFAULT_BOND_FDIM
+        the feature dimension of the edges
+    d_h : int, default=DEFAULT_HIDDEN_DIM
+        the hidden dimension during message passing
+    bias : bool, defuault=False
+        if `True`, add a bias term to the learned weight matrices
+    depth : int, default=3
+        the number of message passing iterations
+    undirected : bool, default=False
+        if `True`, pass messages on undirected edges
+    dropout : float, default=0.0
+        the dropout probability
+    activation : str, default="relu"
+        the activation function to use
+    d_vd : int | None, default=None
+        the dimension of additional vertex descriptors that will be concatenated to the hidden features before readout
+
+    See also
+    --------
+    * :class:`AtomMessagePassing`
+
+    * :class:`BondMessagePassing`
+    """
+
+    def __init__(
+        self,
+        d_v: int = DEFAULT_ATOM_FDIM,
+        d_e: int = DEFAULT_BOND_FDIM,
+        d_h: int = DEFAULT_HIDDEN_DIM,
+        bias: bool = False,
+        depth: int = 3,
+        dropout: float = 0.0,
+        activation: str | Activation = Activation.RELU,
+        undirected: bool = False,
+        d_vd: int | None = None,
+        V_d_transform: ScaleTransform | None = None,
+        graph_transform: GraphTransform | None = None,
+        # layers_per_message: int = 1,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+        self.hparams["cls"] = self.__class__
+
+        self.W_i, self.W_h, self.W_o, self.W_d = self.setup(d_v, d_e, d_h, d_vd, bias)
+        self.depth = depth
+        self.undirected = undirected
+        self.dropout = nn.Dropout(dropout)
+        self.tau = get_activation_function(activation)
+        self.V_d_transform = V_d_transform if V_d_transform is not None else nn.Identity()
+        self.graph_transform = graph_transform if graph_transform is not None else nn.Identity()
+
+    @property
+    def output_dim(self) -> int:
+        return self.W_d.out_features if self.W_d is not None else self.W_o.out_features
+
+    @abstractmethod
+    def setup(
+        self,
+        d_v: int = DEFAULT_ATOM_FDIM,
+        d_e: int = DEFAULT_BOND_FDIM,
+        d_h: int = DEFAULT_HIDDEN_DIM,
+        d_vd: int | None = None,
+        bias: bool = False,
+    ) -> tuple[nn.Module, nn.Module, nn.Module, nn.Module | None]:
+        """setup the weight matrices used in the message passing update functions
+
+        Parameters
+        ----------
+        d_v : int
+            the vertex feature dimension
+        d_e : int
+            the edge feature dimension
+        d_h : int, default=300
+            the hidden dimension during message passing
+        d_vd : int | None, default=None
+            the dimension of additional vertex descriptors that will be concatenated to the hidden
+            features before readout, if any
+        bias: bool, default=False
+            whether to add a learned bias to the matrices
+
+        Returns
+        -------
+        W_i, W_h, W_o, W_d : tuple[nn.Module, nn.Module, nn.Module, nn.Module | None]
+            the input, hidden, output, and descriptor weight matrices, respectively, used in the
+            message passing update functions. The descriptor weight matrix is `None` if no vertex
+            dimension is supplied
+        """
+
+    @abstractmethod
+    def initialize(self, bmg: BatchMolGraph) -> Tensor:
+        """initialize the message passing scheme by calculating initial matrix of hidden features"""
+
+    @abstractmethod
+    def message(self, H_t: Tensor, bmg: BatchMolGraph):
+        """Calculate the message matrix"""
+
+    def update(self, M_t, H_0):
+        """Calcualte the updated hidden for each edge"""
+        H_t = self.W_h(M_t)
+        H_t = self.tau(H_0 + H_t)
+        H_t = self.dropout(H_t)
+
+        return H_t
+
+    def finalize(self, M: Tensor, V: Tensor, V_d: Tensor | None) -> Tensor:
+        r"""Finalize message passing by (1) concatenating the final message ``M`` and the original
+        vertex features ``V`` and (2) if provided, further concatenating additional vertex
+        descriptors ``V_d``.
+
+        This function implements the following operation:
+
+        .. math::
+            H &= \mathtt{dropout} \left( \tau(\mathbf{W}_o(V \mathbin\Vert M)) \right) \\
+            H &= \mathtt{dropout} \left( \tau(\mathbf{W}_d(H \mathbin\Vert V_d)) \right),
+
+        where :math:`\tau` is the activation function, :math:`\Vert` is the concatenation operator,
+        :math:`\mathbf{W}_o` and :math:`\mathbf{W}_d` are learned weight matrices, :math:`M` is
+        the message matrix, :math:`V` is the original vertex feature matrix, and :math:`V_d` is an
+        optional vertex descriptor matrix.
+
+        Parameters
+        ----------
+        M : Tensor
+            a tensor of shape ``V x d_h`` containing the message vector of each vertex
+        V : Tensor
+            a tensor of shape ``V x d_v`` containing the original vertex features
+        V_d : Tensor | None
+            an optional tensor of shape ``V x d_vd`` containing additional vertex descriptors
+
+        Returns
+        -------
+        Tensor
+            a tensor of shape ``V x (d_h + d_v [+ d_vd])`` containing the final hidden
+            representations
+
+        Raises
+        ------
+        InvalidShapeError
+            if ``V_d`` is not of shape ``b x d_vd``, where ``b`` is the batch size and ``d_vd`` is
+            the vertex descriptor dimension
+        """
+        H = self.W_o(torch.cat((V, M), dim=1))  # V x d_o
+        H = self.tau(H)
+        H = self.dropout(H)
+
+        if V_d is not None:
+            V_d = self.V_d_transform(V_d)
+            try:
+                H = self.W_d(torch.cat((H, V_d), dim=1))  # V x (d_o + d_vd)
+                H = self.dropout(H)
+            except RuntimeError:
+                raise InvalidShapeError("V_d", V_d.shape, [len(H), self.W_d.in_features])
+
+        return H
+
+    def forward(self, bmg: BatchMolGraph, V_d: Tensor | None = None) -> Tensor:
+        """Encode a batch of molecular graphs.
+
+        Parameters
+        ----------
+        bmg: BatchMolGraph
+            a batch of :class:`BatchMolGraph`s to encode
+        V_d : Tensor | None, default=None
+            an optional tensor of shape ``V x d_vd`` containing additional descriptors for each atom
+            in the batch. These will be concatenated to the learned atomic descriptors and
+            transformed before the readout phase.
+
+        Returns
+        -------
+        Tensor
+            a tensor of shape ``V x d_h`` or ``V x (d_h + d_vd)`` containing the encoding of each
+            molecule in the batch, depending on whether additional atom descriptors were provided
+        """
+        bmg = self.graph_transform(bmg)
+        H_0 = self.initialize(bmg)
+
+        H = self.tau(H_0)
+        for _ in range(1, self.depth):
+            if self.undirected:
+                H = (H + H[bmg.rev_edge_index]) / 2
+
+            M = self.message(H, bmg)
+            H = self.update(M, H_0)
+
+        index_torch = bmg.edge_index[1].unsqueeze(1).repeat(1, H.shape[1])
+        M = torch.zeros(len(bmg.V), H.shape[1], dtype=H.dtype, device=H.device).scatter_reduce_(
+            0, index_torch, H, reduce="sum", include_self=False
+        )
+        return self.finalize(M, bmg.V, V_d)
+
+
+class BondMessagePassing(_MessagePassingBase):
+    r"""A :class:`BondMessagePassing` encodes a batch of molecular graphs by passing messages along
+    directed bonds.
+
+    It implements the following operation:
+
+    .. math::
+
+        h_{vw}^{(0)} &= \tau \left( \mathbf W_i(e_{vw}) \right) \\
+        m_{vw}^{(t)} &= \sum_{u \in \mathcal N(v)\setminus w} h_{uv}^{(t-1)} \\
+        h_{vw}^{(t)} &= \tau \left(h_v^{(0)} + \mathbf W_h m_{vw}^{(t-1)} \right) \\
+        m_v^{(T)} &= \sum_{w \in \mathcal N(v)} h_w^{(T-1)} \\
+        h_v^{(T)} &= \tau \left (\mathbf W_o \left( x_v \mathbin\Vert m_{v}^{(T)} \right) \right),
+
+    where :math:`\tau` is the activation function; :math:`\mathbf W_i`, :math:`\mathbf W_h`, and
+    :math:`\mathbf W_o` are learned weight matrices; :math:`e_{vw}` is the feature vector of the
+    bond between atoms :math:`v` and :math:`w`; :math:`x_v` is the feature vector of atom :math:`v`;
+    :math:`h_{vw}^{(t)}` is the hidden representation of the bond :math:`v \rightarrow w` at
+    iteration :math:`t`; :math:`m_{vw}^{(t)}` is the message received by the bond :math:`v
+    \to w` at iteration :math:`t`; and :math:`t \in \{1, \dots, T-1\}` is the number of
+    message passing iterations.
+    """
+
+    def setup(
+        self,
+        d_v: int = DEFAULT_ATOM_FDIM,
+        d_e: int = DEFAULT_BOND_FDIM,
+        d_h: int = DEFAULT_HIDDEN_DIM,
+        d_vd: int | None = None,
+        bias: bool = False,
+    ):
+        W_i = nn.Linear(d_v + d_e, d_h, bias)
+        W_h = nn.Linear(d_h, d_h, bias)
+        W_o = nn.Linear(d_v + d_h, d_h)
+        W_d = nn.Linear(d_h + d_vd, d_h + d_vd) if d_vd is not None else None
+
+        return W_i, W_h, W_o, W_d
+
+    def initialize(self, bmg: BatchMolGraph) -> Tensor:
+        return self.W_i(torch.cat([bmg.V[bmg.edge_index[0]], bmg.E], dim=1))
+
+    def message(self, H: Tensor, bmg: BatchMolGraph) -> Tensor:
+        index_torch = bmg.edge_index[1].unsqueeze(1).repeat(1, H.shape[1])
+        M_all = torch.zeros(len(bmg.V), H.shape[1], dtype=H.dtype, device=H.device).scatter_reduce_(
+            0, index_torch, H, reduce="sum", include_self=False
+        )[bmg.edge_index[0]]
+        M_rev = H[bmg.rev_edge_index]
+
+        return M_all - M_rev
+
+
+class AtomMessagePassing(_MessagePassingBase):
+    r"""A :class:`AtomMessagePassing` encodes a batch of molecular graphs by passing messages along
+    atoms.
+
+    It implements the following operation:
+
+    .. math::
+
+        h_v^{(0)} &= \tau \left( \mathbf{W}_i(x_v) \right) \\
+        m_v^{(t)} &= \sum_{u \in \mathcal{N}(v)} h_u^{(t-1)} \mathbin\Vert e_{uv} \\
+        h_v^{(t)} &= \tau\left(h_v^{(0)} + \mathbf{W}_h m_v^{(t-1)}\right) \\
+        m_v^{(T)} &= \sum_{w \in \mathcal{N}(v)} h_w^{(T-1)} \\
+        h_v^{(T)} &= \tau \left (\mathbf{W}_o \left( x_v \mathbin\Vert m_{v}^{(T)} \right)  \right),
+
+    where :math:`\tau` is the activation function; :math:`\mathbf{W}_i`, :math:`\mathbf{W}_h`, and
+    :math:`\mathbf{W}_o` are learned weight matrices; :math:`e_{vw}` is the feature vector of the
+    bond between atoms :math:`v` and :math:`w`; :math:`x_v` is the feature vector of atom :math:`v`;
+    :math:`h_v^{(t)}` is the hidden representation of atom :math:`v` at iteration :math:`t`;
+    :math:`m_v^{(t)}` is the message received by atom :math:`v` at iteration :math:`t`; and
+    :math:`t \in \{1, \dots, T\}` is the number of message passing iterations.
+    """
+
+    def setup(
+        self,
+        d_v: int = DEFAULT_ATOM_FDIM,
+        d_e: int = DEFAULT_BOND_FDIM,
+        d_h: int = DEFAULT_HIDDEN_DIM,
+        d_vd: int | None = None,
+        bias: bool = False,
+    ):
+        W_i = nn.Linear(d_v, d_h, bias)
+        W_h = nn.Linear(d_e + d_h, d_h, bias)
+        W_o = nn.Linear(d_v + d_h, d_h)
+        W_d = nn.Linear(d_h + d_vd, d_h + d_vd) if d_vd is not None else None
+
+        return W_i, W_h, W_o, W_d
+
+    def initialize(self, bmg: BatchMolGraph) -> Tensor:
+        return self.W_i(bmg.V[bmg.edge_index[0]])
+
+    def message(self, H: Tensor, bmg: BatchMolGraph):
+        H = torch.cat((H, bmg.E), dim=1)
+        index_torch = bmg.edge_index[1].unsqueeze(1).repeat(1, H.shape[1])
+        return torch.zeros(len(bmg.V), H.shape[1], dtype=H.dtype, device=H.device).scatter_reduce_(
+            0, index_torch, H, reduce="sum", include_self=False
+        )[bmg.edge_index[0]]

chemprop/nn/message_passing/multi.py

@@ -0,0 +1,78 @@
+from typing import Iterable, Sequence
+import warnings
+
+from torch import Tensor, nn
+
+from chemprop.data import BatchMolGraph
+from chemprop.nn.message_passing.proto import MessagePassing
+from chemprop.nn.hparams import HasHParams
+
+
+class MulticomponentMessagePassing(nn.Module, HasHParams):
+    """A `MulticomponentMessagePassing` performs message-passing on each individual input in a
+    multicomponent input then concatenates the representation of each input to construct a
+    global representation
+
+    Parameters
+    ----------
+    blocks : Sequence[MessagePassing]
+        the invidual message-passing blocks for each input
+    n_components : int
+        the number of components in each input
+    shared : bool, default=False
+        whether one block will be shared among all components in an input. If not, a separate
+        block will be learned for each component.
+    """
+
+    def __init__(self, blocks: Sequence[MessagePassing], n_components: int, shared: bool = False):
+        super().__init__()
+        self.hparams = {
+            "cls": self.__class__,
+            "blocks": [block.hparams for block in blocks],
+            "n_components": n_components,
+            "shared": shared,
+        }
+
+        if len(blocks) == 0:
+            raise ValueError("arg 'blocks' was empty!")
+        if shared and len(blocks) > 1:
+            warnings.warn(
+                "More than 1 block was supplied but 'shared' was True! Using only the 0th block..."
+            )
+        elif not shared and len(blocks) != n_components:
+            raise ValueError(
+                "arg 'n_components' must be equal to `len(blocks)` if 'shared' is False! "
+                f"got: {n_components} and {len(blocks)}, respectively."
+            )
+
+        self.n_components = n_components
+        self.shared = shared
+        self.blocks = nn.ModuleList([blocks[0]] * self.n_components if shared else blocks)
+
+    def __len__(self) -> int:
+        return len(self.blocks)
+
+    @property
+    def output_dim(self) -> int:
+        d_o = sum(block.output_dim for block in self.blocks)
+
+        return d_o
+
+    def forward(self, bmgs: Iterable[BatchMolGraph], V_ds: Iterable[Tensor | None]) -> list[Tensor]:
+        """Encode the multicomponent inputs
+
+        Parameters
+        ----------
+        bmgs : Iterable[BatchMolGraph]
+        V_ds : Iterable[Tensor | None]
+
+        Returns
+        -------
+        list[Tensor]
+            a list of tensors of shape `V x d_i` containing the respective encodings of the `i`\th
+            component, where `d_i` is the output dimension of the `i`\th encoder
+        """
+        if V_ds is None:
+            return [block(bmg) for block, bmg in zip(self.blocks, bmgs)]
+        else:
+            return [block(bmg, V_d) for block, bmg, V_d in zip(self.blocks, bmgs, V_ds)]

chemprop/nn/message_passing/proto.py

@@ -0,0 +1,35 @@
+from abc import abstractmethod
+
+from torch import nn, Tensor
+
+from chemprop.data import BatchMolGraph
+from chemprop.nn.hparams import HasHParams
+
+
+class MessagePassing(nn.Module, HasHParams):
+    """A :class:`MessagePassing` module encodes a batch of molecular graphs
+    using message passing to learn vertex-level hidden representations."""
+
+    input_dim: int
+    output_dim: int
+
+    @abstractmethod
+    def forward(self, bmg: BatchMolGraph, V_d: Tensor | None = None) -> Tensor:
+        """Encode a batch of molecular graphs.
+
+        Parameters
+        ----------
+        bmg: BatchMolGraph
+            the batch of :class:`~chemprop.featurizers.molgraph.MolGraph`\s to encode
+        V_d : Tensor | None, default=None
+            an optional tensor of shape `V x d_vd` containing additional descriptors for each atom
+            in the batch. These will be concatenated to the learned atomic descriptors and
+            transformed before the readout phase.
+
+        Returns
+        -------
+        Tensor
+            a tensor of shape `V x d_h` or `V x (d_h + d_vd)` containing the hidden representation
+            of each vertex in the batch of graphs. The feature dimension depends on whether
+            additional atom descriptors were provided
+        """

chemprop/nn/metrics.py

@@ -0,0 +1,209 @@
+from abc import abstractmethod
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor
+from torchmetrics import functional as F
+from torchmetrics.utilities.compute import auc
+
+from chemprop.utils.registry import ClassRegistry
+from chemprop.nn.loss import (
+    BCELoss,
+    BinaryMCCLoss,
+    CrossEntropyLoss,
+    LossFunction,
+    MSELoss,
+    MulticlassMCCLoss,
+    SIDLoss,
+    WassersteinLoss,
+)
+
+__all__ = [
+    "Metric",
+    "MetricRegistry",
+    "ThresholdedMixin",
+    "MAEMetric",
+    "MSEMetric",
+    "RMSEMetric",
+    "BoundedMixin",
+    "BoundedMAEMetric",
+    "BoundedMSEMetric",
+    "BoundedRMSEMetric",
+    "R2Metric",
+    "BinaryAUROCMetric",
+    "BinaryAUPRCMetric",
+    "BinaryAccuracyMetric",
+    "BinaryF1Metric",
+    "BCEMetric",
+    "CrossEntropyMetric",
+    "BinaryMCCMetric",
+    "MulticlassMCCMetric",
+    "SIDMetric",
+    "WassersteinMetric",
+]
+
+
+class Metric(LossFunction):
+    """
+    Parameters
+    ----------
+    task_weights :  ArrayLike = 1.0
+        .. important::
+            Ignored. Maintained for compatibility with :class:`~chemprop.nn.loss.LossFunction`
+    """
+
+    minimize: bool = True
+
+    def forward(
+        self,
+        preds: Tensor,
+        targets: Tensor,
+        mask: Tensor,
+        weights: Tensor,
+        lt_mask: Tensor,
+        gt_mask: Tensor,
+    ):
+        return self._calc_unreduced_loss(preds, targets, mask, lt_mask, gt_mask)[mask].mean()
+
+    @abstractmethod
+    def _calc_unreduced_loss(self, preds, targets, mask, lt_mask, gt_mask) -> Tensor:
+        pass
+
+
+MetricRegistry = ClassRegistry[Metric]()
+
+
+@dataclass
+class ThresholdedMixin:
+    threshold: float | None = 0.5
+
+    def extra_repr(self) -> str:
+        return f"threshold={self.threshold}"
+
+
+@MetricRegistry.register("mae")
+class MAEMetric(Metric):
+    def _calc_unreduced_loss(self, preds, targets, *args) -> Tensor:
+        return (preds - targets).abs()
+
+
+@MetricRegistry.register("mse")
+class MSEMetric(MSELoss, Metric):
+    pass
+
+
+@MetricRegistry.register("rmse")
+class RMSEMetric(MSEMetric):
+    def forward(
+        self,
+        preds: Tensor,
+        targets: Tensor,
+        mask: Tensor,
+        weights: Tensor,
+        lt_mask: Tensor,
+        gt_mask: Tensor,
+    ):
+        squared_errors = super()._calc_unreduced_loss(preds, targets, mask, lt_mask, gt_mask)
+
+        return squared_errors[mask].mean().sqrt()
+
+
+class BoundedMixin:
+    def _calc_unreduced_loss(self, preds, targets, mask, lt_mask, gt_mask) -> Tensor:
+        preds = torch.where((preds < targets) & lt_mask, targets, preds)
+        preds = torch.where((preds > targets) & gt_mask, targets, preds)
+
+        return super()._calc_unreduced_loss(preds, targets, mask, lt_mask, gt_mask)
+
+
+@MetricRegistry.register("bounded-mae")
+class BoundedMAEMetric(MAEMetric, BoundedMixin):
+    pass
+
+
+@MetricRegistry.register("bounded-mse")
+class BoundedMSEMetric(MSEMetric, BoundedMixin):
+    pass
+
+
+@MetricRegistry.register("bounded-rmse")
+class BoundedRMSEMetric(RMSEMetric, BoundedMixin):
+    pass
+
+
+@MetricRegistry.register("r2")
+class R2Metric(Metric):
+    minimize = False
+
+    def forward(self, preds: Tensor, targets: Tensor, mask: Tensor, *args, **kwargs):
+        return F.r2_score(preds[mask], targets[mask])
+
+
+@MetricRegistry.register("roc")
+class BinaryAUROCMetric(Metric):
+    minimize = False
+
+    def forward(self, preds: Tensor, targets: Tensor, mask: Tensor, *args, **kwargs):
+        return self._calc_unreduced_loss(preds, targets, mask)
+
+    def _calc_unreduced_loss(self, preds, targets, mask, *args) -> Tensor:
+        return F.auroc(preds[mask], targets[mask].long(), task="binary")
+
+
+@MetricRegistry.register("prc")
+class BinaryAUPRCMetric(Metric):
+    minimize = False
+
+    def forward(self, preds: Tensor, targets: Tensor, *args, **kwargs):
+        p, r, _ = F.precision_recall_curve(preds, targets.long(), task="binary")
+        return auc(r, p)
+
+
+@MetricRegistry.register("accuracy")
+class BinaryAccuracyMetric(Metric, ThresholdedMixin):
+    minimize = False
+
+    def forward(self, preds: Tensor, targets: Tensor, mask: Tensor, *args, **kwargs):
+        return F.accuracy(
+            preds[mask], targets[mask].long(), threshold=self.threshold, task="binary"
+        )
+
+
+@MetricRegistry.register("f1")
+class BinaryF1Metric(Metric, ThresholdedMixin):
+    minimize = False
+
+    def forward(self, preds: Tensor, targets: Tensor, mask: Tensor, *args, **kwargs):
+        return F.f1_score(
+            preds[mask], targets[mask].long(), threshold=self.threshold, task="binary"
+        )
+
+
+@MetricRegistry.register("bce")
+class BCEMetric(BCELoss, Metric):
+    pass
+
+
+@MetricRegistry.register("ce")
+class CrossEntropyMetric(CrossEntropyLoss, Metric):
+    pass
+
+
+@MetricRegistry.register("binary-mcc")
+class BinaryMCCMetric(BinaryMCCLoss, Metric):
+    pass
+
+
+@MetricRegistry.register("multiclass-mcc")
+class MulticlassMCCMetric(MulticlassMCCLoss, Metric):
+    pass
+
+
+@MetricRegistry.register("sid")
+class SIDMetric(SIDLoss, Metric):
+    pass
+
+
+@MetricRegistry.register("wasserstein")
+class WassersteinMetric(WassersteinLoss, Metric):
+    pass