[FEATURE] Support datasets with partly missing labels

[PicoCentauri]

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

For Born effective charge (BEC) training of the next PET-MAD model, only a fraction of the structures will have this label. The reason is that for metals, BECs are undefined and can't be computed. I further expect that partly missing labels will occur more often in the future. For example, for structures with a spin. Since spin polarized calculations are expensive, one might only carry them out for some systems in large datasets but still want to train on them.

Describe the solution you'd like

Allow training even if labels are missing in some structures. The question is if we could technically go down to if only a single structure has the label. The training would be a bit useless, but I think we could support this.

For the UI, I suggest adding a switch in options.yaml. Inside the targets sections, the user could enable the support of partly missing labels. By default, this should be turned off.

Describe alternatives you've considered

One could split the dataset into two parts. One with complete labels and one without. This might already work.

Additional context

No response

[pfebrer]

This kind of happens with Hamiltonians as well. So what would you do, you set the missing labels to NaN or you let the user choose? (e.g. maybe they want them to be 0)

[Luthaf]

I kinda feel like this should be done by splitting into separate datasets, especially since metatrain already handles this well. So maybe what we need are tools/examples on how to split existing dataset by properties.

Unless there is a ML reason not to split them?

This also look very similar to the issue of datasets with and without stress/virial, and I think that in this case we are splitting the datasets, also because setting the stress/virial to a default value makes no sense (and NaN would have to be filtered from the loss anyway, and might confuse the gradients tracking).

[PicoCentauri]

I don't see an ML reason to split. It will all go in the same model. But then I agree that we should have an example and some quality of life functions to split a dataset.

[pfebrer]

I don't care what we do but

and NaN would have to be filtered from the loss anyway

this already happens

[ceriottm]

I think that setting to NaN is the way to go. Setting them to zero would be wrong in the specific case of the BECs, as if anything in a metal the polarizability is infinite and so they should be huge, not zero - at least for atoms at the surface... Michele

…

On Mon, 13 Jul 2026 at 11:32, Pol Febrer Calabozo ***@***.***> wrote: *pfebrer* left a comment (metatensor/metatrain#1198) <#1198 (comment)> I don't care what we do but and NaN would have to be filtered from the loss anyway this already happens — Reply to this email directly, view it on GitHub <#1198?email_source=notifications&email_token=AAIREZ3C7PRLDPUBLPET3Q35ESUD7A5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOJVGY2TAMRSHE3KM4TFMFZW63VKON2WE43DOJUWEZLEUVSXMZLOOSWGM33PORSXEX3DNRUWG2Y#issuecomment-4956502296>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAIREZ6ZR3CUGN7IEHLZ5QT5ESUD7AVCNFSNUABFKJSXA33TNF2G64TZHM3TEMJRHE2TGOBUHNEXG43VMU5TIOBXGEYDMNRWGU42C5QC> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[pfebrer]

Btw I just remembered, you can already treat a single file as two separate datasets in metatrain, because you can pass a file with the indices 🙂 So as long as you have metals.txt and non_metals.txt you are good to go in current metatrain, no need for a splitting utility

[PicoCentauri]

Yes, but you have to figure them out. The "issue" is not writing a file or creating some indices but having some useful docs or a function that explains this to someone who has this problem.

[pfebrer]

@ppegolo also has this problem, where he has a dataset with charge density and spin density, but some calculations are not spin polarized so the spin density is not present.

I don't know, if you have tried to make it work with multiple datasets and indices files we can create a doc page and think about a command. But do you have a command in mind? Would it be just for this or a more general command?