Add SO(3)/O(3)-averaging calculator wrappers #85
Conversation
| :py:meth:`ase.Atoms.get_forces`, …); | ||
| - arbitrary outputs can be computed for any :py:class:`ase.Atoms` using | ||
| :py:meth:`MetatomicCalculator.run_model`; | ||
| - for non-equivariant architectures like PET, rotatonally-averaged energies, forces, |
There was a problem hiding this comment.
Worth a link to PET in metatrain, and maybe also link to another non-equivariant architecture?
| self, | ||
| base_calculator: MetatomicCalculator, | ||
| l_max: int = 3, | ||
| batch_size: Optional[int] = None, | ||
| include_inversion: bool = True, |
There was a problem hiding this comment.
| self, | |
| base_calculator: MetatomicCalculator, | |
| l_max: int = 3, | |
| batch_size: Optional[int] = None, | |
| include_inversion: bool = True, | |
| self, | |
| base_calculator: MetatomicCalculator, | |
| *, | |
| l_max: int = 3, | |
| batch_size: Optional[int] = None, | |
| include_inversion: bool = True, |
This forces people to pass all args as kwargs instead of allowing for positional arguments
| l_max: int = 3, | ||
| batch_size: Optional[int] = None, | ||
| include_inversion: bool = True, | ||
| apply_group_symmetry: bool = False, |
There was a problem hiding this comment.
So this is about crystal symmetry group, correct? Or can it also handle molecule point group?
We may want to rename this because I initially read "group" as the SO(3) group and did not understand the point of this ^_^
| "Out of memory error encountered during rotational averaging. " | ||
| "Please reduce the batch size or use lower rotational " | ||
| "averaging parameters. This can be done by setting the " | ||
| "`batch_size`, `lebedev_order`, and `n_inplane_rotations` " |
There was a problem hiding this comment.
n_inplane_rotations does not exists anymore
| :param return_samples: if ``True``, the results of the base calculator on each | ||
| rotated system will be returned. Most useful for debugging. |
There was a problem hiding this comment.
Do we really want to keep this around? Which use case do you foresee for this?
| if "energy" in results: | ||
| E = np.asarray(results["energy"], dtype=float) | ||
| out["energy"] = _wmean(E) | ||
| out["energy_rot_std"] = _wstd(E) |
There was a problem hiding this comment.
These does not seem to be documented, and I'm mildly against adding custom output to a standard interface. What's the use case for these?
IMO if we keep them around they should at least be behind an option store_rotational_std=True (or another similar name).
| # Forces (B,N,3) from rotated structures: back-rotate with F' R | ||
| if "forces" in results: | ||
| F = np.asarray(results["forces"], dtype=float) # (B,N,3) | ||
| F_back = np.einsum("bnj,bjk->bnk", F, R, optimize=True) # F' R |
There was a problem hiding this comment.
einsum (here and for stress) tend to be very slow, was the speed fine for you here? Maybe the model is slower anyway.
| - arbitrary outputs can be computed for any :py:class:`ase.Atoms` using | ||
| :py:meth:`MetatomicCalculator.run_model`; | ||
| - for non-equivariant architectures like PET, rotatonally-averaged energies, forces, | ||
| and stresses can be computed using :py:class:`SymmetrizedCalculator`. |
There was a problem hiding this comment.
You should also add a link to this page: https://docs.metatensor.org/metatomic/latest/torch/reference/ase.html
2 participants
Contributor (creator of pull-request) checklist
Reviewer checklist
📚 Download documentation for this pull-request