Question about species-specific force weighting and force masking

[yonghyuklee]

Dear Developers,

I would like to ask whether it is possible to assign different force weights for different species in MACE. My training set contains structures with varying species compositions, for example, a large A-species substrate and small B-species clusters.

Additionally, is it possible to mask forces for specific atoms in the training set? I recall that GAP allows this using the force_mask key, but I am unsure if a similar option exists in MACE.

Looking forward to your insights. Many thanks!

[gabor1]

I think @LarsSchaaf worked on this.

[LarsSchaaf]

@steffen-wedig was interested in this, will coordinate with him about who implements this.

[bernstei]

When I implemented stuff like this, I applied to each energy-related row energy_type_weight[atoms.info["config_type"]] * atoms.info["energy_weight"] and to each force-related row force_type_weight[atoms.info["config_type"]] * atoms.info["force_weight"] * atoms.arrays["weights"][atom_i] when setting up the design matrix

That gave maximum flexibility, at the cost of the more elaborate possibilities requiring the caller to preprocess configurations to mask out specific forces, etc. But the flexibility means that it'd be trivial to have utility functions to weight each atom's force by a species-specific weight, e.g.

However, I'm not sure how hard it'd be to do in pytorch - when I did it, it was pure python explicitly looping over atoms.