Allow overwriting arguments in interface (#167)

# Allow overwriting arguments in nbed function.
Where keyword arguments exist these are used to overwrite existing values, all inputs are then revalidated.
Doing this makes it much easier to quickly try options or programatically update inputs to `nbed()`, while ensuring that the NbedDriver class stays simple.

# Seperates out embedded-fci and dft-in-dft functions from driver class for ease of testing.
Added in some tests for dft-in-dft, ensuring that the energy matches whole ssytem dft where the same exchange correlation functional is used.

Author: MIWdlB

nbed/__init__.py

@@ -1,8 +1,9 @@
 """Init for Nbed package."""
 
+from .config import NbedConfig
 from .embed import nbed
 from .utils import setup_logs
 
-__all__ = ["nbed"]
+__all__ = ["nbed", "NbedConfig"]
 
 setup_logs()

nbed/driver.py

@@ -13,6 +13,29 @@
 logger = logging.getLogger(__name__)
 
 
+def overwrite_config_kwargs(config: NbedConfig, **config_kwargs) -> NbedConfig:
+    """Overwrites config values with key-words and revalidates.
+
+    Args:
+        config (NbedConfig): A config model.
+        config_kwargs (dict): Any possible key-word arguments.
+
+    Returns:
+        NbedConfig: A validated config model.
+
+    Raises:
+        ValidationError: If key-word arguments provided are not part of model.
+    """
+    if config_kwargs != {}:
+        logger.info("Overwriting select field with additonal config.")
+        config_dict = config.model_dump()
+        for k, v in config_kwargs.items():
+            config_dict[k] = v
+        return NbedConfig(**config_dict)
+    else:
+        return config
+
+
 def nbed(
     config: NbedConfig | str | None = None,
     **config_kwargs,
@@ -29,15 +52,20 @@ def nbed(
     Returns:
         NbedDriver: An embedded driver.
     """
+    logger.info(f"Running Nbed with:\n\tconfig\t{config}\n\tkeywords\t{config_kwargs}")
     match config:
         case NbedConfig():
             logger.info("Using validated config.")
+            config = overwrite_config_kwargs(config, **config_kwargs)
+
         case str() | Path():
             logger.info("Using config file %s", config)
+            logger.info("Validating config from file.")
             with open(FilePath(config)) as f:
-                logger.info("Validating config from file.")
                 data = json.load(f)
-                config = NbedConfig(**data)
+            config = NbedConfig(**data)
+            config = overwrite_config_kwargs(config, **config_kwargs)
+
         case None:
             logger.info("Validating config from passed arguments.")
             logger.debug(f"{config_kwargs=}")

nbed/embed.py

@@ -278,7 +278,7 @@ def reduce_virtuals(scf_method, n_frozen_virt: int) -> lib.StreamObject:
         reduced_scf_method.mo_coeff = reduced_scf_method.mo_coeff[:, :, :-n_frozen_virt]
         reduced_scf_method.mo_occ = reduced_scf_method.mo_occ[:, :-n_frozen_virt]
 
-    elif isinstance(reduced_scf_method, (scf.hf.RHF, scf.rohf.ROHF)):
+    elif isinstance(reduced_scf_method, (scf.hf.RHF)):
         reduced_scf_method.mo_coeff = reduced_scf_method.mo_coeff[:, :-n_frozen_virt]
         reduced_scf_method.mo_occ = reduced_scf_method.mo_occ[:-n_frozen_virt]
 

nbed/ham_builder.py

@@ -70,7 +70,7 @@ def localize_path(self) -> tuple[int, int]:
         logger.debug("Singular Values")
         logger.debug(singular_values)
 
-        if isinstance(scf_object, (scf.hf.RHF, dft.rks.RKS)):
+        if isinstance(scf_object, (scf.rhf.RHF, dft.rks.RKS)):
             alpha = self.localize_spin([s[0] for s in singular_values])
             beta = alpha
         elif isinstance(scf_object, (scf.uhf.UHF, dft.uks.UKS)):

nbed/localizers/ace.py

@@ -169,7 +169,7 @@ def huzinaga_scf(
             scf_energy = calculate_ks_energy(
                 scf_method, dft_potential, density_matrix, huzinaga_op_std
             )
-        elif isinstance(scf_method, (scf.hf.RHF, scf.uhf.UHF)):
+        elif isinstance(scf_method, (scf.rhf.RHF, scf.uhf.UHF)):
             hamiltonian = (
                 scf_method.get_hcore() + dft_potential + 0.5 * vhf + huzinaga_op_std
             )

nbed/scf/huzinaga_scf.py

@@ -34,7 +34,7 @@ def water_mol(water_filepath) -> gto.Mole:
 
 @pytest.fixture(scope="module")
 def water_rhf(water_molecule) -> StreamObject:
-    rhf = scf.RHF(water_molecule)
+    rhf = scf.rhf.RHF(water_molecule)
     rhf.kernel()
     return rhf
 

tests/conftest.py

@@ -43,18 +43,25 @@ def huz_unrestricted_driver(nbed_config) -> NbedDriver:
     driver.embed()
     return driver
 
+def test_restricted_dft_in_dft(mu_driver, huz_driver):
+    mu_did = mu_driver._dft_in_dft(ProjectorEnum.MU)
+    huz_did = huz_driver._dft_in_dft(ProjectorEnum.HUZ)
+    assert np.isclose(mu_did["e_dft_in_dft"], mu_driver._global_ks().e_tot)
+    assert np.isclose(huz_did["e_dft_in_dft"], huz_driver._global_ks().e_tot)
+    assert np.isclose(mu_did["e_dft_in_dft"], huz_did["e_dft_in_dft"])
+
 def test_embedded_fci(nbed_config, mu_driver, mu_unrestricted_driver, huz_driver, huz_unrestricted_driver):
-    assert(np.isclose(mu_driver._run_emb_FCI(mu_driver.embedded_scf).e_tot, -62.261794716560416))
-    assert(np.isclose(mu_unrestricted_driver._run_emb_FCI(mu_unrestricted_driver.embedded_scf).e_tot, -62.261794716560416))
+    assert(np.isclose(mu_driver._run_emb_fci(mu_driver.embedded_scf).e_tot, -62.261794716560416))
+    assert(np.isclose(mu_unrestricted_driver._run_emb_fci(mu_unrestricted_driver.embedded_scf).e_tot, -62.261794716560416))
     nbed_config.projector = ProjectorEnum.HUZ
     nbed_config.n_active_atoms=1
     huz_driver = NbedDriver(nbed_config)
     huz_driver.embed()
-    assert(np.isclose(huz_driver._run_emb_FCI(huz_driver.embedded_scf).e_tot, -51.61379094995273))
+    assert(np.isclose(huz_driver._run_emb_fci(huz_driver.embedded_scf).e_tot, -51.61379094995273))
     nbed_config.force_unrestricted = True
     huz_unrestricted_driver = NbedDriver(nbed_config)
     huz_unrestricted_driver.embed()
-    assert(np.isclose(huz_unrestricted_driver._run_emb_FCI(huz_unrestricted_driver.embedded_scf).e_tot, -51.61379094995273))
+    assert(np.isclose(huz_unrestricted_driver._run_emb_fci(huz_unrestricted_driver.embedded_scf).e_tot, -51.61379094995273))
 
 def test_restricted_projector_results_match(mu_driver, huz_driver) -> None:
     assert mu_driver._mu is not {} and mu_driver._huzinaga is None