fix oom bug; #6
fix oom bug; #6
Conversation
…ld parameters in potential_fn
feat(classical): create harmonic bond & angle generators
…holes into if pol:; api.py:396: remove U_ind from params, access it directly from pme_force
KuangYu
left a comment
KuangYu
left a comment
There was a problem hiding this comment.
Can you change the default setting of ethresh for both admp dispersion and electrostatics? Also, it may be good to start thinking about jit and vmap all the classical potentials (maybe do it in the next pr).
dmff/api.py
Outdated
| self._jaxPotential = None | ||
| self.types = [] | ||
| self.ethresh = 1.0e-5 | ||
| self.ethresh = 1e-5 |
There was a problem hiding this comment.
Change it to 5e-4, keep consistent with electrostatic PME
dmff/api.py
Outdated
| # self._input_params = defaultDict(list) | ||
| self._jaxPotential = None | ||
| self.types = [] | ||
| self.ethresh = 1.0e-5 |
|
|
||
| def generate_get_energy(self): | ||
| def get_energy(positions, box, pairs, k, theta0): | ||
| p1 = positions[self.p1idx] |
There was a problem hiding this comment.
Just a remind that there are distribute_xx functions in admp.pairwise module (maybe it is not the most appropriate place to put these functions ...), which were designed to do such parameter expansions. These functions are jitted and vmapped, and my tests show that they can be noticeably faster. It's okay for now, but in the next stage we may want to think about jit and vmap all classical potentials.
3 participants
fix OOM bug;
api.py:302: add default polarity;
api.py:420: move pol&tholes into if pol:;
api.py:396: remove U_ind from params, access it directly from pme_force