Using JAX we optimize the type of atoms given an adjacency matrix of a molecular graph and the target observable. All observables were computed with the Hückel model.
Otpimization of HOMO-LUMO gap , and polarizability
,
Our code only considers three different optimizers,
- BFGS
- Gradient descent
- Adam
To extended to other optimization methods, check wrapper_opt_method
in huxel/minimize.py
.
import jax.numpy as jnp
from huxel import myMolecule
from huxel import optimization as _opt
atom_types = ["X", "X", "X", "X", "X", "X"]
smile = "C6" # name label
connectivity_matrix = jnp.array(
[
[0, 1, 0, 0, 0, 1],
[1, 0, 1, 0, 0, 0],
[0, 1, 0, 1, 0, 0],
[0, 0, 1, 0, 1, 0],
[0, 0, 0, 1, 0, 1],
[1, 0, 0, 0, 1, 0],
],
dtype=int,
)
xyz = jnp.array([[ 1.40000000e+00, 3.70074342e-17, 0.00000000e+00],
[ 7.00000000e-01, -1.21243557e+00, 0.00000000e+00],
[-7.00000000e-01, -1.21243557e+00, 0.00000000e+00],
[-1.40000000e+00, 2.08457986e-16, 0.00000000e+00],
[-7.00000000e-01, 1.21243557e+00, 0.00000000e+00],
[ 7.00000000e-01, 1.21243557e+00, 0.00000000e+00]])
molec = myMolecule(
"benzene",
smile,
atom_types,
connectivity_matrix,
xyz
)
l = 0
_opt(l, molec,'homo_lumo','BFGS')
We considered eight different molecules.
execute run_molecule_i.py
where the options are,
--l
, integer (to initializejax.random.PRNGKey(l)
)--s
, integer for smile data set (range[1, ..., 8]
)--obj
, objective to optimize options[homo_lumo, polarizability]
--opt
, optimization method[adam, GD, BFGS]
--extfield
, external field magnitude (only for polarizability)
jax, optax, jaxopt