Merge pull request #51 from XanaduAI/16-fix-failing-examples

16 fix failing examples

examples/advanced_examples/readme.md

@@ -0,0 +1,16 @@
+These last two examples are pretty slow and require to have precomputed the molecule in an hdf5 file in advance. Please modify accordingly lines
+
+```python
+training_data_dirpath = os.path.normpath(dirpath + "/data/training/dissociation/")
+training_files = ["H2_dissociation_omegas.h5"]
+```
+
+in those two files. In the case of `train_complex_model.py`, the hdf5 file should have processed the molecules with
+
+```python
+from grad_dft.interface.pyscf import molecule_from_pyscf
+omegas = jnp.array([0., 0.4])
+molecule = molecule_from_pyscf(mf, omegas = omegas)
+```
+
+The `test_constraints.py` file executes an example showing that constraints run. However, note that they have not been tested thoroughly.

examples/train_complex_model.py → examples/advanced_examples/train_complex_model.py

@@ -51,9 +51,9 @@
 # In this example we explain how to replicate the experiments that train
 # the functional in some points of the dissociation curve of H2 or H2^+.
 
-dirpath = os.path.dirname(os.path.dirname(__file__))
+dirpath = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
 training_data_dirpath = os.path.normpath(dirpath + "/data/training/dissociation/")
-training_files = ["H2_extrapolation_train.h5"]
+training_files = ["H2_dissociation_omegas.h5"]
 # alternatively, use "H2plus_extrapolation.h5". You will have needed to execute in data_processing.py
 # distances = [0.5, 0.75, 1, 1.25, 1.5]
 # process_dissociation(atom1 = 'H', atom2 = 'H', charge = 0, spin = 0, file = 'H2_dissociation.xlsx', energy_column_name='cc-pV5Z', training_distances=distances)
@@ -160,7 +160,7 @@ def combine_densities(densities, ehf):
     return jnp.concatenate((densities, ehf), axis=1)
 
 
-omegas = [0.0, 0.4]
+omegas = jnp.array([0.0, 0.4])
 functional = NeuralFunctional(
     coefficients=nn_coefficients,
     energy_densities=partial(densities, functional_type="MGGA"),

examples/train_scf_example.py → examples/advanced_examples/train_scf_loop.py

@@ -115,9 +115,9 @@ def coefficients(instance, rhoinputs, *_, **__):
 num_epochs = 50
 cost_val = jnp.inf
 
-dirpath = os.path.dirname(os.path.dirname(__file__))
-training_data_dirpath = os.path.normpath(dirpath + "/data/training/")
-training_files = "/dissociation/H2_extrapolation_molecules.hdf5"
+dirpath = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+training_data_dirpath = os.path.join(dirpath, "data/training/")
+training_files = "dissociation/H2_extrapolation_train.hdf5"
 
 ####### Loss function and train kernel #######
 
@@ -158,7 +158,7 @@ def train_epoch(state, training_files, training_data_dirpath):
 
     batch_metrics = []
     params, opt_state, cost_val = state
-    fpath = os.path.join(training_data_dirpath + training_files)
+    fpath = os.path.join(training_data_dirpath, training_files)
     print("Training on file: ", fpath, "\n")
 
     load = loader(fname=fpath, randomize=True, training=True, config_omegas=[])

grad_dft/evaluate.py

@@ -28,21 +28,23 @@
 import time
 from scipy.optimize import bisect
 
-from grad_dft.external import Functional
-from grad_dft.utils import PyTree, Array, Scalar, Optimizer
 from grad_dft.functional import Functional
 
 from grad_dft.molecule import Molecule, abs_clip, make_rdm1, orbital_grad
 from grad_dft.train import molecule_predictor
-from grad_dft.utils import PyTree, Array, Scalar, safe_fock_solver
 from grad_dft.interface.pyscf import (
     generate_chi_tensor,
     mol_from_Molecule,
     process_mol,
     mol_from_Molecule,
 )
+
+from grad_dft.external import Functional
+from grad_dft.utils import Optimizer, safe_fock_solver
 from grad_dft.utils.types import Hartree2kcalmol
 
+from jaxtyping import PyTree, Array, Scalar
+
 
 ######## Test kernel ########
 

grad_dft/functional.py

@@ -37,7 +37,7 @@
 from typeguard import typechecked
 from grad_dft.molecule import abs_clip
 
-from grad_dft.utils import Scalar, Array, PyTree, DType, default_dtype
+from grad_dft.utils import DType, default_dtype
 from grad_dft.molecule import Grid, Molecule
 
 import sys

grad_dft/interface/pyscf.py

@@ -29,15 +29,14 @@
 
 # from qdft.reaction import Reaction, make_reaction, get_grad
 from grad_dft.molecule import Grid, Molecule, Reaction, make_reaction
-from grad_dft.utils import DType, default_dtype
+from grad_dft.utils import DType, default_dtype, DensityFunctional, HartreeFock
 from jax.tree_util import tree_map
 from grad_dft.external import NeuralNumInt
 from grad_dft.external import Functional as ExternalFunctional
 
 import h5py
 from pyscf import cc, dft, scf
 
-from grad_dft.utils import DensityFunctional, HartreeFock
 from jaxtyping import Array, Scalar, Int
 from grad_dft.external import _nu_chunk
 
@@ -322,7 +321,7 @@ def loader(
                         # select the indices from the omegas array and load the corresponding chi matrix
                         if config_omegas is None:
                             args[key] = jnp.asarray(value)
-                        elif config_omegas == []:
+                        elif list(config_omegas) == []:
                             args[key] = None
                         else:
                             omegas = list(group["omegas"])

setup.py

@@ -38,7 +38,8 @@
     extras_require={
         "examples": [
             "openfermion>=1.5.1", "tqdm>=4.66.1", "torch>=2.0.1",
-            "matplotlib>=3.7.2", "pandas>=2.0.3", "seaborn>=0.12.2"
+            "matplotlib>=3.7.2", "pandas>=2.0.3", "seaborn>=0.12.2",
+            "openpyxl"
         ],
     },  
     classifiers=[

tests/unit/test_eigenproblem.py

@@ -29,7 +29,7 @@
 import numpy as np
 
 from grad_dft.utils.eigenproblem import safe_eigh
-from grad_dft.utils import Array, Scalar
+from jaxtyping import Array, Scalar
 
 import pytest