Mismatch in the condition of adding properties

[ksj20]

Hi,
First of all thank you for the great project.
I have a question regarding the calculator class in forcefield.
In https://github.com/orbital-materials/orb-models/blob/3c79a26ba28ff64a6522b8db30f52e015d0b9064/orb_models/forcefield/calculator.py#L47C1-L51C1, which is

        if model.node_head is not None:
            properties += ["energy", "free_energy"]
        if model.graph_head is not None:
            properties += ["forces"]

Shouldn't energy/free_energy be added to properties when model.graph_head exists, and forces be added when model.node_head exists?
This is contrary to how energy and forces are mapped after prediction in the same python file, which requires graph_pred and node_pred, respectively. This is shown in https://github.com/orbital-materials/orb-models/blob/3c79a26ba28ff64a6522b8db30f52e015d0b9064/orb_models/forcefield/calculator.py#L79C1-L85C1 which is

        if "energy" in self.implemented_properties:
            self.results["energy"] = float(out["graph_pred"].detach().cpu().item())
            self.results["free_energy"] = self.results["energy"]

        if "forces" in self.implemented_properties:
            self.results["forces"] = out["node_pred"].detach().cpu().numpy()

My solution is to correct the if/else conditions to

        if model.graph_head is not None:
            properties += ["energy", "free_energy"]
        if model.node_head is not None:
            properties += ["forces"]

Thank you

[DeNeutoy]

Ah good spot, these are the wrong way around!

        if model.node_head is not None:
            properties += ["energy", "free_energy"]
        if model.graph_head is not None:
            properties += ["forces"]

This is a bug, but doesn't affect any models in practice because all of the models we have released have both an energy and a forces head. If you would like to send a PR for this, please do.

[ksj20]

Sorry I missed it. Thanks for the PR.