Adding atom bond predictions to chemprop test

chemprop/cli/chemprop_hpopt/mol/ray_results/TorchTrainer_2024-12-20_16-42-36/0e7392c7/error.txt

@@ -0,0 +1,35 @@
+Failure # 1 (occurred at 2024-12-20_16-42-45)
+ray::_Inner.train() (pid=14210, ip=127.0.0.1, actor_id=2c32c4e80ea58095ad1be8d001000000, repr=TorchTrainer)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/ray/tune/trainable/trainable.py", line 331, in train
+    raise skipped from exception_cause(skipped)
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/ray/train/_internal/utils.py", line 57, in check_for_failure
+    ray.get(object_ref)
+           ^^^^^^^^^^^^^^^^^^^
+           ^^^^^^^^^^^^^^^^^^^^^
+                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+ray.exceptions.RayTaskError(ValueError): ray::_RayTrainWorker__execute.get_next() (pid=14211, ip=127.0.0.1, actor_id=ec31d849709d7a2717fe59d601000000, repr=<ray.train._internal.worker_group.RayTrainWorker object at 0x107848e50>)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/ray/train/_internal/worker_group.py", line 33, in __execute
+    raise skipped from exception_cause(skipped)
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/ray/train/_internal/utils.py", line 176, in discard_return_wrapper
+    train_func(*args, **kwargs)
+  File "/Users/brianli/Documents/chemprop/chemprop/cli/hpopt.py", line 394, in <lambda>
+    lambda config: train_model(
+                   ^^^^^^^^^^^^
+  File "/Users/brianli/Documents/chemprop/chemprop/cli/hpopt.py", line 333, in train_model
+    trainer = pl.Trainer(
+              ^^^^^^^^^^^
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/utilities/argparse.py", line 70, in insert_env_defaults
+    return fn(self, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 395, in __init__
+    self._accelerator_connector = _AcceleratorConnector(
+                                  ^^^^^^^^^^^^^^^^^^^^^^
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/accelerator_connector.py", line 130, in __init__
+    self._check_config_and_set_final_flags(
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/accelerator_connector.py", line 218, in _check_config_and_set_final_flags
+    raise ValueError(
+ValueError: You set `strategy=<ray.train.lightning._lightning_utils.RayDDPStrategy object at 0x15dcc6a90>` but strategies from the DDP family are not supported on the MPS accelerator. Either explicitly set `accelerator='cpu'` or change the strategy.

chemprop/cli/chemprop_hpopt/mol/ray_results/TorchTrainer_2024-12-20_16-42-36/f8faf821/error.txt

@@ -0,0 +1,35 @@
+Failure # 1 (occurred at 2024-12-20_16-42-43)
+ray::_Inner.train() (pid=14206, ip=127.0.0.1, actor_id=2f76ed1f85d26f7b310fd19201000000, repr=TorchTrainer)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/ray/tune/trainable/trainable.py", line 331, in train
+    raise skipped from exception_cause(skipped)
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/ray/train/_internal/utils.py", line 57, in check_for_failure
+    ray.get(object_ref)
+           ^^^^^^^^^^^^^^^^^^^
+           ^^^^^^^^^^^^^^^^^^^^^
+                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+ray.exceptions.RayTaskError(ValueError): ray::_RayTrainWorker__execute.get_next() (pid=14209, ip=127.0.0.1, actor_id=9faae5919004157f098a734001000000, repr=<ray.train._internal.worker_group.RayTrainWorker object at 0x10488cc10>)
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/ray/train/_internal/worker_group.py", line 33, in __execute
+    raise skipped from exception_cause(skipped)
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/ray/train/_internal/utils.py", line 176, in discard_return_wrapper
+    train_func(*args, **kwargs)
+  File "/Users/brianli/Documents/chemprop/chemprop/cli/hpopt.py", line 394, in <lambda>
+    lambda config: train_model(
+                   ^^^^^^^^^^^^
+  File "/Users/brianli/Documents/chemprop/chemprop/cli/hpopt.py", line 333, in train_model
+    trainer = pl.Trainer(
+              ^^^^^^^^^^^
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/utilities/argparse.py", line 70, in insert_env_defaults
+    return fn(self, **kwargs)
+           ^^^^^^^^^^^^^^^^^^
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 395, in __init__
+    self._accelerator_connector = _AcceleratorConnector(
+                                  ^^^^^^^^^^^^^^^^^^^^^^
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/accelerator_connector.py", line 130, in __init__
+    self._check_config_and_set_final_flags(
+  File "/opt/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/accelerator_connector.py", line 218, in _check_config_and_set_final_flags
+    raise ValueError(
+ValueError: You set `strategy=<ray.train.lightning._lightning_utils.RayDDPStrategy object at 0x150a70190>` but strategies from the DDP family are not supported on the MPS accelerator. Either explicitly set `accelerator='cpu'` or change the strategy.

chemprop/cli/common.py

@@ -28,7 +28,16 @@ def add_common_args(parser: ArgumentParser) -> ArgumentParser:
         action="store_true",
         help="Turn off using the first row in the input CSV as column names",
     )
-
+    data_args.add_argument(
+        "--is-mixed",
+        action="store_true",
+        help="If some of the targets are intended for molecules, some for atoms, and some for bonds",
+    )
+    data_args.add_argument(
+        "--mixed-columns",
+        nargs="+",
+        help="Target columns that correspond to molecule (first index), atom (second index), and bond (third index) predictions",
+    )
     dataloader_args = parser.add_argument_group("Dataloader args")
     dataloader_args.add_argument(
         "-n",
@@ -120,6 +129,11 @@ def add_common_args(parser: ArgumentParser) -> ArgumentParser:
     featurization_args.add_argument(
         "--no-bond-feature-scaling", action="store_true", help="Turn off extra bond feature scaling"
     )
+    featurization_args.add_argument(
+        "--no-bond-descriptor-scaling",
+        action="store_true",
+        help="Turn off extra bond descriptor scaling",
+    )
     featurization_args.add_argument(
         "--atom-features-path",
         nargs="+",

chemprop/cli/hpopt.py

@@ -316,7 +316,10 @@ def train_model(config, args, train_dset, val_dset, logger, output_transform, in
     logger.info(model)
 
     if args.tracking_metric == "val_loss":
-        T_tracking_metric = model.criterion.__class__
+        if args.is_mixed:
+            T_tracking_metric = model.criterion[0].__class__
+        else:
+            T_tracking_metric = model.criterion.__class__
     else:
         T_tracking_metric = MetricRegistry[args.tracking_metric]
         args.tracking_metric = "val/" + args.tracking_metric
@@ -479,13 +482,35 @@ def main(args: Namespace):
 
     input_transforms = normalize_inputs(train_dset, val_dset, args)
 
-    if "regression" in args.task_type:
+    if args.is_mixed and "regression" in args.task_type:
+        output_transform = []
+        mol_output_scaler = train_dset.mol_dataset.normalize_targets()
+        val_dset.mol_dataset.normalize_targets(mol_output_scaler)
+        output_transform.append(UnscaleTransform.from_standard_scaler(mol_output_scaler))
+        atom_output_scaler = train_dset.atom_dataset.normalize_targets()
+        val_dset.atom_dataset.normalize_targets(atom_output_scaler)
+        output_transform.append(UnscaleTransform.from_standard_scaler(atom_output_scaler))
+        bond_output_scaler = train_dset.bond_dataset.normalize_targets()
+        val_dset.bond_dataset.normalize_targets(bond_output_scaler)
+        output_transform.append(UnscaleTransform.from_standard_scaler(bond_output_scaler))
+        logger.info(
+            f"Train data: mean = {mol_output_scaler.mean_} | std = {mol_output_scaler.scale_}"
+        )
+        logger.info(
+            f"Train data: mean = {atom_output_scaler.mean_} | std = {atom_output_scaler.scale_}"
+        )
+        logger.info(
+            f"Train data: mean = {bond_output_scaler.mean_} | std = {bond_output_scaler.scale_}"
+        )
+    elif args.is_mixed:
+        output_transform = [None, None, None]
+    elif "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)
+        output_transform = [UnscaleTransform.from_standard_scaler(output_scaler)]
     else:
-        output_transform = None
+        output_transform = [None]
 
     train_loader = build_dataloader(
         train_dset, args.batch_size, args.num_workers, seed=args.data_seed