[ENH] Geometry optimization

[sofroniewn]

Is your feature request related to a problem? Please describe.
I would like to optimize the geometry of a molecule from an initial configuration. This would then allow me to compute vibrational frequencies and IR spectra at a local minimum of the potential energy surface.

Describe the solution you'd like

In PySCF I can use their geometry_optimizer to optimize the geometry of a molecule, for example something like

from pyscf import gto
from pyscf.geomopt.geometric_solver import optimize as geometric_opt

mol = gto.Mole(atom=atom, basis = 'sto-3g') # atom is some initial starting description of the molecule
rhf = scf.RHF(mol)
mol_opt = geometric_opt(rhf)

I would like to do something similar in dqc, for example something like

import dqc

mol = dqc.Mol(moldesc=moldesc, basis="sto-3g") # moldesc is some initial starting description of the molecule
mol.atompos.requires_grad_()
hf = dqc.HF(mol).run()
mol_opt = dqc.optimize_geometry(hf) # new proposed function

Describe alternatives you've considered
Alternatively this would have to be implemented outside of dqc.

Additional context
I know there are many algorithms possible for geometry optimization and I am very open to whichever is recommended.

[mfkasim1]

You can couple it with xitorch (like scipy for pytorch) with a few additional lines as shown here: https://github.com/diffqc/dqc/blob/master/examples/01-equil-pos.py.
I have a slight disfavor on having it specially implemented in dqc as it would be just a very thin wrapper of xitorch.optimize.minimize.

[sofroniewn]

You can couple it with xitorch (like scipy for pytorch) with a few additional lines as shown here: https://github.com/diffqc/dqc/blob/master/examples/01-equil-pos.py.

Ah great - somehow I missed that example, maybe it could be more prominently added to the documentation.

I have a slight disfavor on having it specially implemented in dqc as it would be just a very thin wrapper of xitorch.optimize.minimize.

I can understand that. I have given that example try now and I can confirm that it appears to work for me, though I havn't done any rigorous testing.

Having just given it a try, my personal preference would still be to add it to dqc as a single function call. The main reason is that I have to define my own function get_ene and then handle the passing of the initial atompos and creation of a new molecule with the optimized atompos. All in all I have to touch quite a few different parts of the API, and import a function xitorch.optimize from another library (which I love!) but which a new dqc user might now be so familiar with.

All of this could be wrapt inside a single mol_opt = dqc.optimize_geometry(hf) call, which would be much more robust and convenient. If you like I can try to implement this and submit a PR, but also fine if you don't want to expand the API or if you'd rather do it yourself.

Thanks again for the prompt and helpful reply.

[mfkasim1]

I think it can be put in the API along side with other properties: https://github.com/diffqc/dqc/blob/master/dqc/api/properties.py
One thing that I don't know yet is: should it returns the tensor of atom position or the molecule object. The other API properties returns tensor(s), so if it returns a molecule object, it would be the odd one out (a very good way to confuse new users). But I can see clunkiness of having to reconstruct the molecule object with a new position (especially if you have a lot of parameters such as efield etc). What do you think?

[sofroniewn]

Hmm - that's a good call out. I could imagine that sometimes people might just want the optimized positions and I can see how it is much more in keeping with the design on the overall API to return a tensor and favor two calls to get the molecule back. I also tend not to like APIs that change return type based on a flag, so I think you have to go with one or the other. Even though it's slightly less convenient for me the following might be fine

mol = dqc.Mol(moldesc=moldesc, basis="sto-3g") # moldesc is some initial starting description of the molecule
mol.atompos.requires_grad_()
hf = dqc.HF(mol).run()
geo_opt = dqc.optimimal_geometry(hf) # new proposed function
mol_opt = dqc.Mol(moldesc=(mol.atomzs, geo_opt), basis=mol.basis) # NOTE: mol.basis doesn't actually exist
# OR
mol.atompos = geo_opt # NOTE: attribute can't be set right now but update in place might be nice

I will call out that mol.basis doesn't exist - but some nice way to get the basis might be convenient.

I will also call out that mol.atompos isn't settable, but with an optimimal_geometry function it could be very nice to just call mol.atompos = dqc.optimimal_geometry(hf). I'm curious what you think about that / if you foresee potential problems?

Finally I suggested the name optimimal_geometry now as what is returned is a "geometry" (i.e. atompos) whereas optimize_geometry to me suggests more on an inplace operation/ operation that returns a full molecule, but maybe that is just me.

Curious what you think next step should be?

[mfkasim1]

Hmm - that's a good call out. I could imagine that sometimes people might just want the optimized positions and I can see how it is much more in keeping with the design on the overall API to return a tensor and favor two calls to get the molecule back.

Returning the tensor also works better if one wants to train a ML system that matches the atom positions to experimental data.

I will also call out that mol.atompos isn't settable, but with an optimimal_geometry function it could be very nice to just call mol.atompos = dqc.optimimal_geometry(hf). I'm curious what you think about that / if you foresee potential problems?

mol.atomposs = something is quite tricky, because it's basically reinitialize everything. It appears simple from the outside, but requires a lot of computation, e.g. basis integrals, in the background. I also avoid to use .build() like in PySCF because it adds some unclarity (e.g. it is unclear for new users if some parameters are assigned, do we need to do .build() or not?)

A solution that I can think of is to add a new method, say Mol.make_copy(**kwargs) where the new parameters are set in kwargs, e.g. for the above example: mol_opt = mol.make_copy(atomposs=geo_opt)

Finally I suggested the name optimimal_geometry now as what is returned is a "geometry" (i.e. atompos) whereas optimize_geometry to me suggests more on an inplace operation/ operation that returns a full molecule, but maybe that is just me.

Is there a special reason why it is optimimal_geometry instead of optimal_geometry?

[sofroniewn]

A solution that I can think of is to add a new method, say Mol.make_copy(**kwargs) where the new parameters are set in kwargs, e.g. for the above example: mol_opt = mol.make_copy(atomposs=geo_opt)

Makes sense, I like that API - very explicit.

Is there a special reason why it is optimimal_geometry instead of optimal_geometry?

Nope, typo! Sorry about that.

I think we're aligned on the API now. Is this something you're ok to implement or would you like me to take a pass at making a PR?

[mfkasim1]

You can take it if you want

[sofroniewn]

You can take it if you want

Great, thanks! I'd love to. I won't be able to make a PR till this weekend at the earliest though