Implement a wrapper for heat flux calculation #104
Description
The idea is to implement the heat flux calculation as a torch script, and as a result it can be easily used in both ASE and LAMMPS.
The model will be implemented in this way (Thanks @Luthaf ):
class HeatFlux(torch.nn.Module):
def __init__(self, model: AtomisticModel):
super().__init__()
self.model = model.module
def forward(
self,
systems: List[System],
outputs: Dict[str, ModelOutput],
selected_atoms = None
) -> Dict[str, TensorMap]:
if "experimental::heat_flux" in outputs:
..
else:
return self.model(systems, outputs, selected_atoms)
wrapper = HeatFlux(pet_model)
wrapper.save("heat_flux.pt")
calc = MetatomicCalculator("heat_flux.pt")
outputs = calc.run_model(atoms, {"experimental::heat_flux": ModelOutput(per_atom=False)})
heat_flux = outputs["experimental::heat_flux"].block().valuesAnd will be located in the models folder under metatomic/python/metatomic_torch/metatomic/torch/.