[FEATURE] Some wandb log improvements

[tulga-rdn]

Is your feature request related to a problem? Please describe.

I think the wandb logging needs some small changes that I think would be useful.

1. Saving every output from local logs to wandb. Currently, wandb doesn't record everything, and most crucially, it misses the output folder, so the user has to check the local log to get it.
2. The losses and the energies are in a single group in wandb workspace, but each of the 4 forces metrics (training/val rmse/mse) is in a separate group.

These issues are small, but it makes some analysis workflows a bit more annoying.

Describe the solution you'd like

For 1), I think we can just initialize the wandb logger a bit earlier. For 2), I have no idea

Describe alternatives you've considered

Currently, I just grep the output folder from the local log files, and move the graphs around manually on wandb workspace.

Additional context

No response

[Luthaf]

Saving every output from local logs to wandb. Currently, wandb doesn't record everything, and most crucially, it misses the output folder, so the user has to check the local log to get it.

What exactly is missing from the output folder? Do you want to (can we even?) send checkpoints to wandb? Or are you talking about data that's in the log files but not in wandb?

[tulga-rdn]

I meant that the local log starts with smth like

[2025-07-07 15:35:50][INFO] - This log is also available at '/Users/runner/work/metatrain/metatrain/tests/resources/outputs/2025-07-07/15-35-50/train.log'.
[2025-07-07 15:35:50][INFO] - Package version: 0.1.dev1+ga38649d
[2025-07-07 15:35:50][INFO] - Package directory: /Users/runner/work/metatrain/metatrain/src/metatrain
[2025-07-07 15:35:50][INFO] - Working directory: /Users/runner/work/metatrain/metatrain/tests/resources
[2025-07-07 15:35:50][INFO] - Executed command: mtt train options.yaml -o model-32-bit.pt -r base_precision=32
[2025-07-07 15:35:50][INFO] - Running training for 'soap_bpnn' architecture
[2025-07-07 15:35:50][INFO] - Random seed of this run is 42
[2025-07-07 15:35:50][INFO] - Setting up training set
Warning: 07 15:35:51][WARNING] - No forces found in section 'U0'.
Warning: 07 15:35:51][WARNING] - No stress found in section 'U0'.
[2025-07-07 15:35:51][INFO] - Setting up validation set
[2025-07-07 15:35:51][INFO] - Setting up test set
[2025-07-07 15:35:51][INFO] - Atomic types inferred from training and validation datasets
[2025-07-07 15:35:51][INFO] - Model defined for atomic types: [1, 6, 7, 8]
[2025-07-07 15:35:51][INFO] - Training dataset:

but the wandb log starts with

[2025-07-07 15:35:50][INFO] - Setting up training set
Warning: 07 15:35:51][WARNING] - No forces found in section 'U0'.
Warning: 07 15:35:51][WARNING] - No stress found in section 'U0'.
[2025-07-07 15:35:51][INFO] - Setting up validation set
[2025-07-07 15:35:51][INFO] - Setting up test set
[2025-07-07 15:35:51][INFO] - Atomic types inferred from training and validation datasets
[2025-07-07 15:35:51][INFO] - Model defined for atomic types: [1, 6, 7, 8]
[2025-07-07 15:35:51][INFO] - Training dataset:

so it doesn't save the dir where the checkpoints etc are stored, which is a bit annoying, so you have to grep the slurm output file, for example.

[tulga-rdn]

and the workspace looks like this

[PicoCentauri]

Regarding 1: Very valid point. Very difficult (or ugly) technically. When the local logger is started the options.yaml is not parsed yet and we have no idea if the user wants us to log on wandb. We could hack it... I will think

What exactly is missing from the output folder? Do you want to (can we even?) send checkpoints to wandb? Or are you talking about data that's in the log files but not in wandb?

We could: https://docs.wandb.ai/ref/python/save/

Regarding 2: I have no idea either. We use the log function of wandb and create the dictionary like this

metatrain/src/metatrain/utils/logging.py

Lines 83 to 98 in 84cc28f

	def emit_data(self, keys: List[str], values: List[str], units: List[str]):
	"""Log structured data to Weights & Biases.
	:param keys: Column header names
	:param values: Data values to write
	:param units: Units for each column
	"""
	_validate_length(keys, values, units)
	data = {}
	for key, value, unit in zip(keys, values, units):
	name = f"{key} [{unit}]" if unit else key
	data[name] = float(value)
	epoch = int(data.pop("Epoch"))
	self.run.log(data, step=epoch, commit=True)

So every metric gets a key in the dictionary that we send. Maybe @sirmarcel has more insights from his projects.

[sirmarcel]

You can make groups by including / in the key. Here's the logger for marathon:

class WandB:
    def __init__(self, run, keys=["energy", "forces"], metrics=None):
        if metrics is None:
            metrics = {key: ["r2", "mae", "rmse"] for key in keys}
        if keys is None:
            assert metrics is not None

        self.metrics = metrics

        self.run = run

        for key, metrics in self.metrics.items():
            for metric in metrics:
                if "r2" == metric:
                    summary = "max"
                else:
                    summary = "min"

            self.run.define_metric(f"train/{key} {metric}", summary=summary)
            self.run.define_metric(f"val/{key} {metric}", summary=summary)

    def __call__(self, epoch, train_loss, train_metrics, val_loss, val_metrics, other=None):
        import numpy as np

        data = {}
        data["train/loss"] = train_loss
        data["val/loss"] = val_loss

        if np.isnan(train_loss):
            # we don't log NaNs
            return

        for key, metrics in self.metrics.items():
            for metric in metrics:
                if key in train_metrics and metric in train_metrics[key]:
                    data[f"train/{key} {metric}"] = train_metrics[key][metric]
                if key in val_metrics and metric in val_metrics[key]:
                    data[f"val/{key} {metric}"] = val_metrics[key][metric]

        if other is not None:
            for k, v in other.items():
                data[k] = v

        self.run.log(step=epoch, data=data, commit=True)

the output in a workspace looks like this

[sirmarcel]

Regarding the folder, wandb.init accepts a config argument (a dict). I just put all the available information there.