allow GGA-level semilocal feature sets

ciderpress/dft/plans.py

@@ -155,25 +155,58 @@ def __init__(self, settings, nspin):
         """
         self.settings = settings
         self.nspin = nspin
-        if self.settings.mode != "npa":
+        if self.settings.mode not in ["npa", "nst", "np", "ns"]:
             raise NotImplementedError
 
-    def get_feat(self, rho):
-        assert rho.ndim == 3
+    def _fill_feat_npa_(self, rho, feat):
+        assert feat.shape[1] == 3
         sigma = np.einsum("sxg,sxg->sg", rho[:, 1:4], rho[:, 1:4])
-        # nfeat is always 3 for npa SL settings
-        feat = np.empty((rho.shape[0], self.settings.nfeat, rho.shape[2]))
         feat[:, 0] = rho[:, 0]
         feat[:, 1] = get_s2(rho[:, 0], sigma)
         feat[:, 2] = get_alpha(rho[:, 0], sigma, rho[:, 4])
         feat[:, 0] *= self.nspin
         feat[:, 1:3] /= self.nspin ** (2.0 / 3)
-        return feat
 
-    def get_vxc(self, rho, vfeat, vxc=None):
+    def _fill_feat_nst_(self, rho, feat):
+        assert feat.shape[1] == 3
+        sigma = np.einsum("sxg,sxg->sg", rho[:, 1:4], rho[:, 1:4])
+        feat[:, 0] = rho[:, 0]
+        feat[:, 1] = sigma
+        feat[:, 2] = rho[:, 4]
+        feat[:, 0] *= self.nspin
+        feat[:, 1] *= self.nspin * self.nspin
+        feat[:, 2] *= self.nspin
+
+    def _fill_feat_np_(self, rho, feat):
+        assert feat.shape[1] == 2
+        sigma = np.einsum("sxg,sxg->sg", rho[:, 1:4], rho[:, 1:4])
+        feat[:, 0] = rho[:, 0]
+        feat[:, 1] = get_s2(rho[:, 0], sigma)
+        feat[:, 0] *= self.nspin
+        feat[:, 1] /= self.nspin ** (2.0 / 3)
+
+    def _fill_feat_ns_(self, rho, feat):
+        assert feat.shape[1] == 2
+        sigma = np.einsum("sxg,sxg->sg", rho[:, 1:4], rho[:, 1:4])
+        feat[:, 0] = rho[:, 0]
+        feat[:, 1] = sigma
+        feat[:, 0] *= self.nspin
+        feat[:, 1] *= self.nspin * self.nspin
+
+    def get_feat(self, rho):
         assert rho.ndim == 3
-        if vxc is None:
-            vxc = np.zeros_like(rho[:, :5])
+        feat = np.empty((rho.shape[0], self.settings.nfeat, rho.shape[2]))
+        if self.settings.mode == "npa":
+            self._fill_feat_npa_(rho, feat)
+        elif self.settings.mode == "nst":
+            self._fill_feat_nst_(rho, feat)
+        elif self.settings.mode == "np":
+            self._fill_feat_np_(rho, feat)
+        else:
+            self._fill_feat_ns_(rho, feat)
+        return feat
+
+    def _fill_vxc_npa_(self, rho, vfeat, vxc):
         sigma = np.einsum("sxg,sxg->sg", rho[:, 1:4], rho[:, 1:4])
         sm23 = self.nspin ** (-2.0 / 3)
         dpdn, dpdsigma = ds2(rho[:, 0], sigma)
@@ -187,6 +220,39 @@ def get_vxc(self, rho, vfeat, vxc=None):
             * rho[:, 1:4]
         )
         vxc[:, 4] += sm23 * vfeat[:, 2] * dadtau
+
+    def _fill_vxc_nst_(self, rho, vfeat, vxc):
+        vxc[:, 0] += self.nspin * vfeat[:, 0]
+        vxc[:, 1:4] += (self.nspin * self.nspin * 2 * vfeat[:, 1])[:, None, :] * rho[
+            :, 1:4
+        ]
+        vxc[:, 4] += self.nspin * vfeat[:, 2]
+
+    def _fill_vxc_np_(self, rho, vfeat, vxc):
+        sigma = np.einsum("sxg,sxg->sg", rho[:, 1:4], rho[:, 1:4])
+        sm23 = self.nspin ** (-2.0 / 3)
+        dpdn, dpdsigma = ds2(rho[:, 0], sigma)
+        vxc[:, 0] += self.nspin * vfeat[:, 0] + sm23 * vfeat[:, 1] * dpdn
+        vxc[:, 1:4] += (sm23 * vfeat[:, 1] * dpdsigma)[:, None, :] * 2 * rho[:, 1:4]
+
+    def _fill_vxc_ns_(self, rho, vfeat, vxc):
+        vxc[:, 0] += self.nspin * vfeat[:, 0]
+        vxc[:, 1:4] += (self.nspin * self.nspin * 2 * vfeat[:, 1])[:, None, :] * rho[
+            :, 1:4
+        ]
+
+    def get_vxc(self, rho, vfeat, vxc=None):
+        assert rho.ndim == 3
+        if vxc is None:
+            vxc = np.zeros_like(rho[:, :5])
+        if self.settings.mode == "npa":
+            self._fill_vxc_npa_(rho, vfeat, vxc)
+        elif self.settings.mode == "nst":
+            self._fill_vxc_nst_(rho, vfeat, vxc)
+        elif self.settings.mode == "np":
+            self._fill_vxc_np_(rho, vfeat, vxc)
+        else:
+            self._fill_vxc_ns_(rho, vfeat, vxc)
         return vxc
 
 

ciderpress/dft/settings.py

@@ -1451,34 +1451,41 @@ class SemilocalSettings(BaseSettings):
     Should not be edited from default in general.
     """
 
-    def __init__(self, level="MGGA", mode="nst"):
-        # mode should be nst or npa
-        self.level = level.upper()
-        # Currently only support meta-GGA because we don't currently
-        # have an elegant way of handling whether various semilocal
-        # ingredients are available to the more complex features
-        # like NLDF and Fractional Laplacian.
-        if self.level != "MGGA":
-            raise NotImplementedError("Only MGGA implemented currently")
-        if mode not in ["nst", "npa"]:
-            raise ValueError("Mode must be nst or npa")
+    def __init__(self, mode="nst"):
+        # mode should be nst or npa (MGGA) or ns or np (GGA).
+        self.level = None
+        if mode in ["nst", "npa"]:
+            self.level = "MGGA"
+        elif mode in ["ns", "np"]:
+            self.level = "GGA"
+        else:
+            raise ValueError("Mode must be nst, npa, ns, or np.")
         self.mode = mode
 
     @property
     def nfeat(self):
-        return 3
+        if self.level == "MGGA":
+            return 3
+        else:
+            return 2
 
     def get_feat_usps(self):
         if self.mode == "nst":
             return [3, 8, 5]
-        else:
+        elif self.mode == "npa":
             return [3, 0, 0]
+        elif self.mode == "ns":
+            return [3, 8]
+        else:
+            return [3, 0]
 
     def ueg_vector(self, rho=1.0):
         if self.mode == "nst":
             return np.array([rho, 0.0, CFC * rho ** (5.0 / 3)])
-        else:
+        elif self.mode == "npa":
             return np.array([rho, 0.0, 1.0])
+        else:
+            return np.array([rho, 0.0])
 
     def get_reasonable_normalizer(self):
         return [None] * self.nfeat

ciderpress/dft/tests/test_plans.py

@@ -1,15 +1,16 @@
 import unittest
 
 import numpy as np
-from numpy.testing import assert_almost_equal, assert_equal
+from numpy.testing import assert_allclose, assert_almost_equal, assert_equal
 
-from ciderpress.dft.plans import NLDFGaussianPlan, NLDFSplinePlan
+from ciderpress.dft.plans import NLDFGaussianPlan, NLDFSplinePlan, SemilocalPlan
 from ciderpress.dft.settings import (
     CFC,
     NLDFSettingsVI,
     NLDFSettingsVIJ,
     NLDFSettingsVJ,
     NLDFSettingsVK,
+    SemilocalSettings,
 )
 
 DELTA = 1e-7
@@ -30,6 +31,103 @@ def ovlp_reference(a, b, featid, extra_args=None):
         raise ValueError
 
 
+class TestSemilocalPlan(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        N = 50
+        cls._x = np.linspace(0, 1, N)
+
+        # construct a reasonable set of densities
+        np.random.seed(53)
+        directions = 0.25 * np.random.normal(size=(3, N))
+        np.random.seed(53)
+        alpha = np.random.normal(size=N) ** 2
+        rho = 1 + 0.5 * np.sin(np.pi * cls._x)
+        drho = directions * rho
+        tauw = np.einsum("xg,xg->g", drho, drho) / (8 * rho)
+        tau = tauw + alpha * CFC * rho ** (5.0 / 3)
+        rho_data = np.concatenate([[rho], drho, [tau]], axis=0)
+        assert rho_data.shape[0] == 5
+        cls.rho = rho_data[None, :, :]
+        cls.sprho = 0.5 * np.stack([rho_data, rho_data])
+        cls.gga_rho = rho_data[None, :4, :]
+        cls.gga_sprho = 0.5 * np.stack([rho_data[:4], rho_data[:4]])
+
+        cls.npa = SemilocalSettings(mode="npa")
+        cls.nst = SemilocalSettings(mode="nst")
+        cls.np = SemilocalSettings(mode="np")
+        cls.ns = SemilocalSettings(mode="ns")
+
+    def check_settings(self, settings):
+        plan = SemilocalPlan(settings, 1)
+        if settings.level == "MGGA":
+            rho = self.rho
+        else:
+            rho = self.gga_rho
+        rho = rho.copy()
+        feat = plan.get_feat(rho)
+
+        def get_e_and_vfeat(feat):
+            energy = (1 + (feat**2).sum(axis=1)) ** 0.5
+            return (energy, feat / energy[:, None, :])
+
+        vfeat = get_e_and_vfeat(feat)[1]
+        vxc1 = plan.get_vxc(rho, vfeat, vxc=None)
+        vxc2 = np.zeros_like(vxc1)
+        vxc3 = plan.get_vxc(rho, vfeat, vxc=vxc2)
+        assert_allclose(vxc2, vxc1, atol=1e-14)
+        assert_allclose(vxc3, vxc1, atol=1e-14)
+
+        delta = 1e-5
+        for i in range(rho.shape[1]):
+            rho[:, i] += 0.5 * delta
+            f = plan.get_feat(rho)
+            ep = get_e_and_vfeat(f)[0]
+            rho[:, i] -= delta
+            f = plan.get_feat(rho)
+            em = get_e_and_vfeat(f)[0]
+            rho[:, i] += 0.5 * delta
+            assert_allclose(vxc1[:, i], (ep - em) / delta, rtol=1e-8, atol=1e-8)
+
+        nsp_feat = feat
+        plan = SemilocalPlan(settings, 2)
+        if settings.level == "MGGA":
+            rho = self.sprho
+        else:
+            rho = self.gga_sprho
+        rho = rho.copy()
+        feat = plan.get_feat(rho)
+
+        assert feat.shape[0] == 2
+        assert feat.shape[1:] == nsp_feat.shape[1:]
+        assert_allclose(feat[0], nsp_feat[0], atol=1e-14)
+        assert_allclose(feat[1], nsp_feat[0], atol=1e-14)
+
+        vfeat = get_e_and_vfeat(feat)[1]
+        vxc1 = plan.get_vxc(rho, vfeat, vxc=None)
+        vxc2 = np.zeros_like(vxc1)
+        vxc3 = plan.get_vxc(rho, vfeat, vxc=vxc2)
+        assert_allclose(vxc2, vxc1, atol=1e-14)
+        assert_allclose(vxc3, vxc1, atol=1e-14)
+
+        delta = 1e-5
+        for i in range(rho.shape[1]):
+            rho[:, i] += 0.5 * delta
+            f = plan.get_feat(rho)
+            ep = get_e_and_vfeat(f)[0]
+            rho[:, i] -= delta
+            f = plan.get_feat(rho)
+            em = get_e_and_vfeat(f)[0]
+            rho[:, i] += 0.5 * delta
+            assert_allclose(vxc1[:, i], (ep - em) / delta, rtol=1e-8, atol=1e-8)
+
+    def test_settings(self):
+        self.check_settings(self.npa)
+        self.check_settings(self.np)
+        self.check_settings(self.nst)
+        self.check_settings(self.ns)
+
+
 class TestNLDFGaussianPlan(unittest.TestCase):
 
     _x = None
@@ -295,7 +393,7 @@ def get_e_and_vfeat(feat):
             energy = (1 + (feat**2).sum(axis=0)) ** 0.5
             return (energy, feat / energy)
 
-        energy, vfeat = get_e_and_vfeat(feat)
+        vfeat = get_e_and_vfeat(feat)[1]
         vrho = np.zeros_like(rho)
         vdrho = np.zeros_like(drho)
         vtau = np.zeros_like(tau)