Default dtype MACEHamiltonian
The default dtype of macehamiltonian is float32, this is hardcoded in the parameters method of MACEHamiltonian. For certain methods, such as optimize, this default is overwritten to float64, but at the moment you cannot change the dtype to compute energy, forces, and/or stresses via the compute method of MACEHamiltonian. Ideally the user can set dtype when initializing the MACEHamiltonian. However, when set, should we still overwrite the dtype for methods such as optimize?
Yeah, it's a little awkward that way.
I actually think dtype should not be an attribute of the Hamiltonian (because it fundamentally isn't -- it's an attribute of the type of sampling/optimization being performed). I'd suggest to give this function a generic dtype argument which can be set to either float32 or float64 by the code that lives in either sample.py or optimize.py. Does that make sense?
https://github.com/molmod/psiflow/blob/b9573589fd14c5950d884a4f70a3b028e3d2afb5/psiflow/sampling/sampling.py#L383
I agree with this philosophy. Besides sample.py and optimize.py, also the Hamiltonian.compute() and Dataset.evaluate() methods should be changed such that the dtype can be specified in all model evaluations currently implemented in psiflow. Correct me if I forgot a method.
I guess the dtype of compute and evaluate should just be defined in ModelEvaluation?