Reaction relative loss

README.md

@@ -178,6 +178,8 @@ The preprocessing command accepts `.xyz`, `.lmdb`/`.aselmdb`, and `.h5` inputs;
 
 For MACE-POLAR/PolarMACE inputs, Equitrain preserves system-level `charge`, `spin`, and `external_field` metadata and maps them to the MACE graph keys `total_charge`, `total_spin`, and `external_field`; the XYZ key names can be changed with `--total-charge-key`, `--total-spin-key`, and `--external-field-key`.
 
+For reaction-relative training data, Equitrain also preserves integer `source_id`, `reaction_id`, and `state_id` metadata. Ordinary frames default to `source_id=0`, `reaction_id=-1`, and `state_id=-1`; reactive triplets can use `state_id=0` for reactant, `1` for transition state, and `2` for product. The XYZ key names can be changed with `--source-id-key`, `--reaction-id-key`, and `--state-id-key`.
+
 Under the hood, each processed file is organised as:
 
 - `/structures`: per-configuration metadata (cell, energy, stress, charge, spin, external field, weights, etc.) and pointers into the per-atom arrays.
@@ -268,6 +270,12 @@ HDF5 inputs can be a directory, a glob (e.g. `data/train_*.h5`), or a comma-sepa
 list of files; all shards are concatenated in order. This applies to
 `--train-file`, `--valid-file`, and `--test-file` when training with either backend.
 
+Torch training can add reaction-relative energy targets with `--barrier-weight` for
+`E_TS - E_reactant` and `--reaction-energy-weight` for
+`E_product - E_reactant`. These losses require complete reaction groups in the
+Torch batch, are averaged once per reaction rather than per frame, and are not
+currently available for the JAX backend.
+
 <!-- TODO: change this following a notebook style -->
 #### Python Script:
 
@@ -701,9 +709,25 @@ Delta fine-tuning is the simplest adapter method in the repository:
 - the forward pass uses `base_parameter + delta`
 - the base model stays frozen throughout optimisation
 
-This is effectively LoRA without any rank compression. It is useful when you
-want the simplest possible residual fine-tuning scheme and do not need to limit
-adapter size aggressively.
+This is the Equitrain residual-parameter implementation of L<sup>2</sup>-SP ("Starting
+Point") regularization from [Li, Grandvalet, and Davoine, 2018, *Explicit
+Inductive Bias for Transfer Learning with Convolutional
+Networks*](https://proceedings.mlr.press/v80/li18a.html). L<sup>2</sup>-SP regularizes
+fine-tuned parameters toward their pre-trained starting values instead of toward
+zero:
+
+```text
+Omega(theta) = lambda / 2 * ||theta - theta_0||_2^2
+```
+
+Equitrain parameterizes this as `theta = theta_0 + delta`, with `theta_0`
+frozen and `delta` initialized at zero. Optimizer weight decay on decayed delta
+tensors therefore regularizes `||delta||_2^2`, i.e. the distance between the
+effective fine-tuned parameters and the pre-trained parameters.
+
+Compared with LoRA, delta fine-tuning uses full-size residuals rather than
+low-rank residuals, so it is useful when you want the simplest residual
+fine-tuning scheme and do not need to limit adapter size aggressively.
 
 Implementation details:
 
@@ -740,6 +764,12 @@ zero-based layer indices or ranges. For MACE models, Equitrain groups deltas as
 `TorchDeltaFineTuneWrapper(base_model, freeze_layers="2-")` keeps only the
 node embedding and first interaction block trainable.
 
+When delta fine-tuning is combined with `freeze_layers`, Equitrain calls this
+targeted L<sup>2</sup>-SP (L<sup>2</sup>-TSP): the L<sup>2</sup>-SP
+penalty is applied only to the selected trainable delta layers, while frozen
+layers keep `delta = 0` and remain exactly at their pre-trained starting
+values.
+
 #### Freeze Fine-Tuning
 
 For Torch models, `TorchFreezeFineTuneWrapper` provides the same semantic layer

equitrain/argparser.py

@@ -274,6 +274,18 @@ def add_loss_weights_args(parser: argparse.ArgumentParser) -> argparse.ArgumentP
     parser.add_argument(
         '--stress-weight', help='Weight for stress loss', type=float, default=1.0
     )
+    parser.add_argument(
+        '--barrier-weight',
+        help='Weight for relative barrier loss E_TS - E_reactant',
+        type=float,
+        default=0.0,
+    )
+    parser.add_argument(
+        '--reaction-energy-weight',
+        help='Weight for reaction energy loss E_product - E_reactant',
+        type=float,
+        default=0.0,
+    )
     return parser
 
 
@@ -607,6 +619,24 @@ def get_args_parser(script_type: str) -> argparse.ArgumentParser:
             type=str,
             default='external_field',
         )
+        parser.add_argument(
+            '--source-id-key',
+            help='Key of integer source id in training xyz',
+            type=str,
+            default='source_id',
+        )
+        parser.add_argument(
+            '--reaction-id-key',
+            help='Key of integer reaction group id in training xyz',
+            type=str,
+            default='reaction_id',
+        )
+        parser.add_argument(
+            '--state-id-key',
+            help='Key of integer reaction state id in training xyz',
+            type=str,
+            default='state_id',
+        )
         parser.add_argument(
             '--output-dir', help='Output directory', type=str, default=''
         )
@@ -885,19 +915,57 @@ def _ensure_losses_defined(args, backend_name: str) -> None:
     energy = getattr(args, 'energy_weight', 0.0) or 0.0
     forces = getattr(args, 'forces_weight', 0.0) or 0.0
     stress = getattr(args, 'stress_weight', 0.0) or 0.0
-    if energy == 0.0 and forces == 0.0 and stress == 0.0:
+    barrier = getattr(args, 'barrier_weight', 0.0) or 0.0
+    reaction_energy = getattr(args, 'reaction_energy_weight', 0.0) or 0.0
+
+    if backend_name == 'jax' and (barrier != 0.0 or reaction_energy != 0.0):
+        raise ArgumentError(
+            'The JAX backend does not support relative reaction losses yet; '
+            'set --barrier-weight 0 and --reaction-energy-weight 0.'
+        )
+
+    if getattr(args, 'weighted_sampler', False) and (
+        barrier != 0.0 or reaction_energy != 0.0
+    ):
+        raise ArgumentError(
+            'The weighted sampler does not support relative reaction losses yet; '
+            'disable --weighted-sampler or set relative loss weights to zero.'
+        )
+
+    if (
+        energy == 0.0
+        and forces == 0.0
+        and stress == 0.0
+        and barrier == 0.0
+        and reaction_energy == 0.0
+    ):
         raise ArgumentError(
             f'{backend_name} backend requires at least one non-zero loss weight.'
         )
 
 
+def fine_tune_export_config(model):
+    config_fn = getattr(model, 'get_fine_tune_export_config', None)
+    if callable(config_fn):
+        return config_fn()
+    return None
+
+
+def args_dict_with_runtime_metadata(args):
+    args_dict = dict(vars(args))
+    fine_tune_config = fine_tune_export_config(args_dict.get('model'))
+    if fine_tune_config is not None:
+        args_dict['fine_tune_export'] = fine_tune_config
+    return args_dict
+
+
 class ArgsFormatter:
     def __init__(self, args):
         """
         Initialize the ArgsFormatter with parsed arguments.
         :param args: argparse.Namespace object
         """
-        self.args = vars(args)  # Convert Namespace to dictionary
+        self.args = args_dict_with_runtime_metadata(args)
 
     def format(self):
         """
@@ -932,6 +1000,10 @@ def is_simple(self, value):
 
     def filter(self, args):
         """Filter the list of arguments to include only allowed types."""
-        return {
-            key: value for key, value in vars(args).items() if self.is_simple(value)
+        args_dict = args_dict_with_runtime_metadata(args)
+        filtered = {
+            key: value for key, value in args_dict.items() if self.is_simple(value)
         }
+        if 'fine_tune_export' in args_dict:
+            filtered['fine_tune_export'] = args_dict['fine_tune_export']
+        return filtered

equitrain/backends/jax_backend.py

@@ -20,6 +20,7 @@
 from jax import tree_util as jtu
 
 from equitrain.argparser import (
+    ArgsFilterSimple,
     ArgsFormatter,
     check_args_consistency,
     validate_training_args,
@@ -921,6 +922,90 @@ def _run_eval_loop(
     return mean_loss, loss_collection
 
 
+def _jax_runtime_config(
+    args,
+    *,
+    requested_batch_size,
+    requested_batch_max_nodes,
+    multi_device: bool,
+    device_count: int,
+    effective_workers: int,
+    prefetch_batches: int,
+    process_count: int | None = None,
+    process_index: int | None = None,
+) -> dict[str, object]:
+    graph_multiple = device_count if multi_device else 1
+    config: dict[str, object] = {
+        'backend': 'jax',
+        'jax_runtime_batching': 'graph-packing',
+        'jax_requested_batch_size': requested_batch_size,
+        'jax_runtime_batch_size': getattr(args, 'batch_size', None),
+        'jax_requested_batch_max_nodes': requested_batch_max_nodes,
+        'jax_runtime_batch_max_nodes': getattr(args, 'batch_max_nodes', None),
+        'jax_runtime_batch_max_edges': getattr(args, 'batch_max_edges', None),
+        'jax_runtime_graph_multiple': graph_multiple,
+        'jax_runtime_multi_device': multi_device,
+        'jax_runtime_device_count': device_count,
+        'jax_runtime_num_workers': effective_workers,
+        'jax_runtime_prefetch_batches': prefetch_batches,
+    }
+    if process_count is not None:
+        config['jax_runtime_process_count'] = process_count
+    if process_index is not None:
+        config['jax_runtime_process_index'] = process_index
+    return config
+
+
+def _log_jax_runtime_summary(logger, runtime_config: dict[str, object]) -> None:
+    if logger is None:
+        return
+
+    logger.log(
+        1,
+        'JAX runtime batching      : '
+        f'{runtime_config["jax_runtime_batching"]} '
+        f'(requested batch_size={runtime_config["jax_requested_batch_size"]}, '
+        'runtime batch_size='
+        f'{runtime_config["jax_runtime_batch_size"]})',
+    )
+    logger.log(
+        1,
+        'JAX runtime node limit    : '
+        f'requested={runtime_config["jax_requested_batch_max_nodes"]}, '
+        f'runtime={runtime_config["jax_runtime_batch_max_nodes"]}',
+    )
+    logger.log(
+        1,
+        f'JAX runtime edge limit    : {runtime_config["jax_runtime_batch_max_edges"]}',
+    )
+    logger.log(
+        1,
+        f'JAX runtime graph multiple: {runtime_config["jax_runtime_graph_multiple"]}',
+    )
+    logger.log(
+        1,
+        'JAX runtime devices       : '
+        f'{runtime_config["jax_runtime_device_count"]} '
+        f'(multi_device={runtime_config["jax_runtime_multi_device"]})',
+    )
+    logger.log(
+        1,
+        'JAX runtime workers       : '
+        f'{runtime_config["jax_runtime_num_workers"]} '
+        f'(prefetch={runtime_config["jax_runtime_prefetch_batches"]})',
+    )
+    if (
+        'jax_runtime_process_index' in runtime_config
+        and 'jax_runtime_process_count' in runtime_config
+    ):
+        logger.log(
+            1,
+            'JAX runtime process       : '
+            f'{runtime_config["jax_runtime_process_index"]}/'
+            f'{runtime_config["jax_runtime_process_count"]}',
+        )
+
+
 def train(args):
     exit_code = _launch_local_processes(args)
     if exit_code is not None:
@@ -952,20 +1037,6 @@ def train(args):
     logger.log(1, ArgsFormatter(args))
 
     wandb_run = None
-    if is_primary and getattr(args, 'wandb_project', None):
-        try:
-            import wandb
-        except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
-            raise RuntimeError(
-                'wandb is required for the JAX backend when wandb_project is set.'
-            ) from exc
-
-        init_kwargs = {'project': args.wandb_project}
-        if getattr(args, 'wandb_name', None):
-            init_kwargs['name'] = args.wandb_name
-        if getattr(args, 'wandb_group', None):
-            init_kwargs['group'] = args.wandb_group
-        wandb_run = wandb.init(**init_kwargs, config={'backend': 'jax'})
 
     bundle = load_model_bundle(
         args.model,
@@ -988,6 +1059,8 @@ def train(args):
     local_devices = jax.local_devices()
     device_count = len(local_devices) if multi_device else 1
 
+    requested_batch_size = getattr(args, 'batch_size', None)
+    requested_batch_max_nodes = getattr(args, 'batch_max_nodes', None)
     args.batch_size = None
     if getattr(args, 'batch_max_edges', None) is None:
         raise ValueError(
@@ -1008,6 +1081,36 @@ def train(args):
     else:
         prefetch_batches = max(int(prefetch_requested or 0), 0)
 
+    runtime_config = _jax_runtime_config(
+        args,
+        requested_batch_size=requested_batch_size,
+        requested_batch_max_nodes=requested_batch_max_nodes,
+        multi_device=multi_device,
+        device_count=device_count,
+        effective_workers=effective_workers,
+        prefetch_batches=prefetch_batches,
+        process_count=process_count,
+        process_index=process_index,
+    )
+    _log_jax_runtime_summary(logger, runtime_config)
+
+    if is_primary and getattr(args, 'wandb_project', None):
+        try:
+            import wandb
+        except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
+            raise RuntimeError(
+                'wandb is required for the JAX backend when wandb_project is set.'
+            ) from exc
+
+        init_kwargs = {'project': args.wandb_project}
+        if getattr(args, 'wandb_name', None):
+            init_kwargs['name'] = args.wandb_name
+        if getattr(args, 'wandb_group', None):
+            init_kwargs['group'] = args.wandb_group
+        wandb_config = ArgsFilterSimple().filter(args)
+        wandb_config.update(runtime_config)
+        wandb_run = wandb.init(**init_kwargs, config=wandb_config)
+
     def _build_streaming_loader(path: str | None, shuffle: bool):
         if path in (None, 'None'):
             return None

equitrain/backends/jax_evaluate.py

@@ -10,7 +10,9 @@
 from equitrain.backends.jax_backend import (
     _build_eval_step,
     _initialize_distributed,
+    _jax_runtime_config,
     _launch_local_processes,
+    _log_jax_runtime_summary,
     _run_eval_loop,
     _shutdown_distributed,
 )
@@ -119,6 +121,8 @@ def _evaluate_initialized(args):
         )
     multi_device = _is_multi_device()
     device_count = jax.local_device_count() if multi_device else 1
+    requested_batch_size = getattr(args, 'batch_size', None)
+    requested_batch_max_nodes = getattr(args, 'batch_max_nodes', None)
     args.batch_size = None
     if getattr(args, 'batch_max_edges', None) is None:
         raise ValueError(
@@ -139,6 +143,19 @@ def _evaluate_initialized(args):
     else:
         prefetch_batches = max(int(prefetch_requested or 0), 0)
 
+    runtime_config = _jax_runtime_config(
+        args,
+        requested_batch_size=requested_batch_size,
+        requested_batch_max_nodes=requested_batch_max_nodes,
+        multi_device=multi_device,
+        device_count=device_count,
+        effective_workers=effective_workers,
+        prefetch_batches=prefetch_batches,
+        process_count=getattr(jax, 'process_count', lambda: 1)(),
+        process_index=process_index,
+    )
+    _log_jax_runtime_summary(logger, runtime_config)
+
     test_loader = get_dataloader(
         data_file=test_file,
         atomic_numbers=z_table,