dqc
dqc copied to clipboard
diffqc
[BUG] Bug in `optimal_geometry` for KS case
Describe the bug
There is a bug in the optimal_geometry function add in #11 for DFT case
To Reproduce Steps to reproduce the behavior:
import dqc
moldesc = "O 0.0000 0.0000 0.2217; H 0.0000 1.4309 -0.8867; H 0.0000 -1.4309 -0.8867"
m = dqc.Mol(moldesc, basis="cc-pvdz", grid=4, efield=efield)
m.atompos.requires_grad_()
# sys = dqc.HF(m).run() # WORKS
sys = dqc.KS(m, xc="lda_x+lda_c_pw").run() # DOESN'T WORK
dqc.optimal_geometry(sys)
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
/Users/nsofroniew/Documents/code/chem/dqc_0.ipynb Cell 24' in <module>
[6](vscode-notebook-cell:/Users/nsofroniew/Documents/code/chem/dqc_0.ipynb#ch0000030?line=5)[ # sys = dqc.HF(m).run() # WORKS
]()[7](vscode-notebook-cell:/Users/nsofroniew/Documents/code/chem/dqc_0.ipynb#ch0000030?line=6)[ sys = dqc.KS(m, xc="lda_x+lda_c_pw").run() # DOESN'T WORK
----> ]()[8](vscode-notebook-cell:/Users/nsofroniew/Documents/code/chem/dqc_0.ipynb#ch0000030?line=7)[ dqc.optimal_geometry(sys)
File ~/GitHub/dqc/dqc/api/properties.py:339, in optimal_geometry(qc, length_unit)
]()[321](file:///~/GitHub/dqc/dqc/api/properties.py?line=320)[ def optimal_geometry(qc: BaseQCCalc, length_unit: Optional[str] = None) -> torch.Tensor:
]()[322](file:///~/GitHub/dqc/dqc/api/properties.py?line=321)[ """
]()[323](file:///~/GitHub/dqc/dqc/api/properties.py?line=322)[ Compute the optimal atomic positions of the system.
]()[324](file:///~/GitHub/dqc/dqc/api/properties.py?line=323)[
(...)
]()[337](file:///~/GitHub/dqc/dqc/api/properties.py?line=336)[ of atoms at the optimal geometry.
]()[338](file:///~/GitHub/dqc/dqc/api/properties.py?line=337)[ """
--> ]()[339](file:///~/GitHub/dqc/dqc/api/properties.py?line=338)[ atompos = _optimal_geometry(qc)
]()[340](file:///~/GitHub/dqc/dqc/api/properties.py?line=339)[ atompos = convert_length(atompos, to_unit=length_unit)
]()[341](file:///~/GitHub/dqc/dqc/api/properties.py?line=340)[ return atompos
File ~/GitHub/dqc/dqc/utils/misc.py:32, in memoize_method.<locals>.new_fcn(self)
]()[30](file:///~/GitHub/dqc/dqc/utils/misc.py?line=29)[ return self.__dict__[cachename]
]()[31](file:///~/GitHub/dqc/dqc/utils/misc.py?line=30)[ else:
---> ]()[32](file:///~/GitHub/dqc/dqc/utils/misc.py?line=31)[ res = fcn(self)
]()[33](file:///~/GitHub/dqc/dqc/utils/misc.py?line=32)[ self.__dict__[cachename] = res
]()[34](file:///~/GitHub/dqc/dqc/utils/misc.py?line=33)[ return res
File ~/GitHub/dqc/dqc/api/properties.py:503, in _optimal_geometry(qc)
]()[500](file:///~/GitHub/dqc/dqc/api/properties.py?line=499)[ return ene
]()[502](file:///~/GitHub/dqc/dqc/api/properties.py?line=501)[ # get the minimal enery position
--> ]()[503](file:///~/GitHub/dqc/dqc/api/properties.py?line=502)[ minpos = xitorch.optimize.minimize(_get_energy, atompos, method="gd", maxiter=200,
]()[504](file:///~/GitHub/dqc/dqc/api/properties.py?line=503)[ step=1e-2)
]()[506](file:///~/GitHub/dqc/dqc/api/properties.py?line=505)[ return minpos
File ~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py:275, in minimize(fcn, y0, params, bck_options, method, **fwd_options)
]()[270](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=269)[ # if it is just using the rootfinder algorithm, then the forward function
]()[271](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=270)[ # is the one that returns only the grad
]()[272](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=271)[ else:
]()[273](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=272)[ _fwd_fcn = _rf_fcn
--> ]()[275](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=274)[ return _RootFinder.apply(_rf_fcn, y0, _fwd_fcn, opt_method, fwd_options, bck_options,
]()[276](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=275)[ len(params), *params, *pfunc.objparams())
File ~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py:302, in _RootFinder.forward(ctx, fcn, y0, fwd_fcn, is_opt_method, options, bck_options, nparams, *allparams)
]()[300](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=299)[ name = "rootfinder" if not is_opt_method else "minimizer"
]()[301](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=300)[ method_fcn = get_method(name, methods, method)
--> ]()[302](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=301)[ y = method_fcn(fwd_fcn, y0, params, **config)
]()[304](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=303)[ ctx.fcn = fcn
]()[305](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=304)[ ctx.is_opt_method = is_opt_method
File ~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_impls/optimize/minimizer.py:50, in gd(fcn, x0, params, step, gamma, maxiter, f_tol, f_rtol, x_tol, x_rtol, verbose, **unused)
]()[48](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_impls/optimize/minimizer.py?line=47)[ v = torch.zeros_like(x)
]()[49](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_impls/optimize/minimizer.py?line=48)[ for i in range(maxiter):
---> ]()[50](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_impls/optimize/minimizer.py?line=49)[ f, dfdx = fcn(x, *params)
]()[52](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_impls/optimize/minimizer.py?line=51)[ # update the step
]()[53](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_impls/optimize/minimizer.py?line=52)[ v = (gamma * v - step * dfdx).detach()
File ~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_core/pure_function.py:34, in PureFunction.__call__(self, *params)
]()[33](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_core/pure_function.py?line=32)[ def __call__(self, *params):
---> ]()[34](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_core/pure_function.py?line=33)[ return self._fcntocall(*params)
File ~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py:256, in minimize.<locals>._min_fwd_fcn(y, *params)
]()[254](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=253)[ with torch.enable_grad():
]()[255](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=254)[ y1 = y.clone().requires_grad_()
--> ]()[256](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=255)[ z = pfunc(y1, *params)
]()[257](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=256)[ grady, = torch.autograd.grad(z, (y1,), retain_graph=True,
]()[258](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=257)[ create_graph=torch.is_grad_enabled())
]()[259](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/optimize/rootfinder.py?line=258)[ return z, grady
File ~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_core/pure_function.py:34, in PureFunction.__call__(self, *params)
]()[33](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_core/pure_function.py?line=32)[ def __call__(self, *params):
---> ]()[34](file:///~/opt/anaconda3/envs/chem/lib/python3.9/site-packages/xitorch/_core/pure_function.py?line=33)[ return self._fcntocall(*params)
File ~/GitHub/dqc/dqc/api/properties.py:498, in _optimal_geometry.<locals>._get_energy(atompos)
]()[496](file:///~/GitHub/dqc/dqc/api/properties.py?line=495)[ def _get_energy(atompos: torch.Tensor) -> torch.Tensor:
]()[497](file:///~/GitHub/dqc/dqc/api/properties.py?line=496)[ new_system = system.make_copy(moldesc=(system.atomzs, atompos))
--> ]()[498](file:///~/GitHub/dqc/dqc/api/properties.py?line=497)[ new_qc = qc.__class__(new_system).run()
]()[499](file:///~/GitHub/dqc/dqc/api/properties.py?line=498)[ ene = new_qc.energy() # calculate the energy
]()[500](file:///~/GitHub/dqc/dqc/api/properties.py?line=499)[ return ene
TypeError: __init__() missing 1 required positional argument: 'xc']()
Expected behavior The method should work
Desktop (please complete the following information):
- DQC version: master
- Python version: 3.9
- PyTorch version: 1.10.2
- OS: macOS
Additional context
I will try and fix and add a test when I get the chance - maybe later tonight or tomorrow depending on timing. Hopefully it will be an easy fix, just looks like something missing from an __init__
So the issue here is the new_qc = qc.__class__(new_system).run() call which is missing the xc="lda_x+lda_c_pw" argument for the KS case. I could try and fix it with some sort of if/else logic, but a much cleaner solution might be to add a new method make_copy to BaseQCCalc, which could take the same kwargs as HF / KS etc. Alternatively we could add a set_system method, but I think for consistency I prefer the make_copy approach.
the function inside the optimal_geometry method then becomes
# get the energy for a given geometry
def _get_energy(atompos: torch.Tensor) -> torch.Tensor:
new_system = system.make_copy(moldesc=(system.atomzs, atompos))
new_qc = qc.make_copy(system=new_system).run()
ene = new_qc.energy() # calculate the energy
return ene
which I quite like.
I will note the reason this wasn't caught in tests is that right now all the properties are only tested with h2o_qc which is HF. If you wanted I could rename that h2o_qc_hf and add an h2o_qc_ks and then make all the tests using h2o_qc parametric tests.
Let me know what approach you'd like both for fixing the bug and for adding tests