[BUG] Bug in `optimal_geometry` for KS case

[sofroniewn]

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__

[sofroniewn]

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