diff --git a/ciderpress/analyzers.py b/ciderpress/analyzers.py
deleted file mode 100644
index 98b6813..0000000
--- a/ciderpress/analyzers.py
+++ /dev/null
@@ -1,462 +0,0 @@
-from abc import ABC, abstractmethod
-
-import numpy as np
-from pyscf import dft, gto, lib, scf
-from pyscf.dft.gen_grid import Grids
-
-from ciderpress.external.sgx_tools import get_jk_densities, sgx_fit
-
-CALC_TYPES = {
- "RHF": scf.hf.RHF,
- "UHF": scf.uhf.UHF,
-}
-
-
-def recursive_remove_none(obj):
- if isinstance(obj, dict):
- return {k: recursive_remove_none(v) for k, v in obj.items() if v is not None}
- else:
- return obj
-
-
-def recursive_bytes_to_str(obj):
- if isinstance(obj, dict):
- return {k: recursive_bytes_to_str(v) for k, v in obj.items()}
- elif isinstance(obj, (list, tuple)):
- lst = [recursive_bytes_to_str(item) for item in obj]
- if isinstance(obj, tuple):
- lst = tuple(lst)
- return lst
- elif isinstance(obj, bytes):
- return obj.decode("utf-8")
- else:
- return obj
-
-
-class ElectronAnalyzer(ABC):
- """
- A class for generating and storing data derived from a PySCF
- electronic structure calculation, in particular distributions
- on a real-space grid such as the density, exchange (correlation)
- energy density, Coulomb energy density, etc.
-
- Attributes:
-
- TODO
-
- """
-
- _atype = None
-
- def __init__(
- self, mol, dm, grids_level=3, mo_occ=None, mo_coeff=None, mo_energy=None
- ):
- if mo_occ is None and hasattr(dm, "mo_occ"):
- self.mo_occ = dm.mo_occ
- self.mo_coeff = dm.mo_coeff
- self.mo_energy = dm.mo_energy
- else:
- self.mo_occ = mo_occ
- self.mo_coeff = mo_coeff
- self.mo_energy = mo_energy
- self.dm = dm
- self.mol = mol
- self.mol.verbose = 3
- self.mol.build()
- self.grids_level = grids_level
- self.grids = Grids(self.mol)
- self.grids.level = self.grids_level
- self._data = {}
- self.grids.build(with_non0tab=False)
-
- def keys(self):
- return self._data.keys()
-
- def set(self, k, v):
- self._data[k] = v
-
- def as_dict(self):
- return {
- "atype": self._atype,
- "data": self._data,
- "dm": self.dm,
- "mo_occ": self.mo_occ,
- "mo_coeff": self.mo_coeff,
- "mo_energy": self.mo_energy,
- "grids_level": self.grids_level,
- "mol": gto.mole.pack(self.mol),
- }
-
- @staticmethod
- def from_dict(d):
- args = [
- gto.mole.unpack(d["mol"]),
- d["dm"],
- ]
- kwargs = {
- "grids_level": d["grids_level"],
- "mo_occ": d.get("mo_occ"),
- "mo_coeff": d.get("mo_coeff"),
- "mo_energy": d.get("mo_energy"),
- }
- atype = d.get("atype")
- if atype is None:
- cls = ElectronAnalyzer
- elif atype == "RHF" or atype == "RKS":
- cls = RHFAnalyzer
- elif atype == "UHF" or atype == "UKS":
- cls = UHFAnalyzer
- else:
- raise ValueError("Unrecognized analyzer type {}".format(atype))
- analyzer = cls(*args, **kwargs)
- analyzer._data.update(d["data"])
- return analyzer
-
- def dump(self, fname):
- """
- Dump self to an hdf5 file called fname.
-
- Args:
- fname (str): Name of file to dump to
- """
- h5dict = recursive_remove_none(self.as_dict())
- lib.chkfile.dump(fname, "analyzer", h5dict)
-
- @staticmethod
- def load(fname):
- """
- Load instance of cls from hdf5
- Args:
- fname (str): Name of file from which to load
- max_mem: See __init__
- """
- analyzer_dict = lib.chkfile.load(fname, "analyzer")
- analyzer_dict = recursive_bytes_to_str(analyzer_dict)
- return ElectronAnalyzer.from_dict(analyzer_dict)
-
- def get(self, name):
- return self._data[name]
-
- def calculate_vxc(self, xcname, xcfunc=None, grids=None, xctype="MGGA"):
- """
- Computes the XC potential matrix for a functional,
- store it, and return it.
-
- Args:
- xcname (str): XC functional
- xcfunc (callable): XC function in case one
- wants to use an XC functional not in libxc
- grids (optional): Grid to use to calculate XC.
-
- Returns:
- vxc (nspin x nao x nao): XC potential matrix
- """
- if not isinstance(xcname, str):
- raise ValueError("XC must be string")
- if self._atype == "RHF":
- calc = dft.RKS(self.mol)
- elif self._atype == "UHF":
- calc = dft.UKS(self.mol)
- else:
- raise NotImplementedError
- calc.xc = xcname
- if grids is not None:
- calc.grids = grids
- if xcfunc is not None:
- calc.define_xc_(xcfunc, xctype=xctype)
- vxc = calc.get_veff(dm=self.rdm1)
- self._data["EXC_{}".format(xcname)] = vxc.exc
- vxc = vxc - vxc.vj
- self._data["VXC_{}".format(xcname)] = vxc
- return vxc
-
- def calculate_vxc_on_mo(self, xcname, orbs=None, **kwargs):
- """
- Compute contributions of XC potential to the eigenvalues. If
- VXC_{xcname} is not in _data, calculate_vxc is called first.
-
- Args:
- xcname (str): Name of XC functional (for libxc)
- orbs: Orbital index dictionary for computing orbital derivatives,
- in the format use for descriptors.get_descriptors.
- **kwargs: extra arguments to pass to calculate_vxc, if it is called.
-
- Returns:
- eigenvalue contributions. If orbs is provided, it will be in
- the same format used by descritpors.get_descriptors.
- Otherwise, it will be a numpy array with all the eigenvalue
- contributions.
- """
- if "VXC_{}".format(xcname) not in self._data.keys():
- vxc = self.calculate_vxc(xcname, **kwargs)
- else:
- vxc = self._data["VXC_{}".format(xcname)]
- assert self.mo_coeff is not None
- if orbs is None:
- orbvals = vxc.dot(self.mo_coeff)
- eigvals = np.einsum("...ui,ui->...i", self.mo_coeff, orbvals)
- self._data["ORBXC_{}".format(xcname)] = eigvals
- else:
- from ciderpress.descriptors import get_labels_and_coeffs, unpack_eigvals
-
- labels, coeffs, en_list, sep_spins = get_labels_and_coeffs(
- orbs, self.mo_coeff, self.mo_occ, self.mo_energy
- )
- if self.mo_occ.ndim == 2:
- orbvals = {}
- for s in range(2):
- c = np.array(coeffs[s])
- if len(c) == 0:
- orbvals[s] = None
- continue
- orbvals[s] = np.array(c).dot(vxc[s])
- orbvals[s] = np.einsum("iu,iu->i", c, orbvals[s])
- else:
- c = np.array(coeffs)
- orbvals = c.dot(vxc)
- orbvals = np.einsum("iu,iu->i", c, orbvals)
- if sep_spins:
- eigvals = (
- unpack_eigvals(orbs[0], labels[0], orbvals[0]),
- unpack_eigvals(orbs[1], labels[1], orbvals[1]),
- )
- elif self.mo_occ.ndim == 2:
- eigvals = {}
- for k in list(orbs.keys()):
- eigvals[k] = {}
- for s in range(2):
- eigvals_tmp = unpack_eigvals(orbs, labels[s], orbvals[s])
- for k in list(orbs.keys()):
- eigvals[k].update(eigvals_tmp[k])
- else:
- eigvals = unpack_eigvals(orbs, labels, orbvals)
- self._data["ORBXC_{}".format(xcname)] = eigvals
- return self._data["ORBXC_{}".format(xcname)]
-
- @staticmethod
- def from_calc(calc, grids_level=None, store_energy_orig=True):
- """
- NOTE: This has side effects on calc, see notes below.
-
- Args:
- calc: SCF object from PySCF
- grids_level: Size of PySCF grids for XC integration
- store_energy_orig: Whether to store original xc energy.
-
- Returns:
-
- """
- if isinstance(calc, scf.uhf.UHF):
- cls = UHFAnalyzer
- elif isinstance(calc, scf.hf.RHF):
- cls = RHFAnalyzer
- else:
- raise ValueError("Must be HF calculation")
-
- if grids_level is None:
- if isinstance(calc, dft.rks.KohnShamDFT):
- grids_level = calc.grids.level
- else:
- grids_level = 3
- analyzer = cls(
- calc.mol,
- calc.make_rdm1(),
- grids_level=grids_level,
- mo_occ=calc.mo_occ,
- mo_coeff=calc.mo_coeff,
- mo_energy=calc.mo_energy,
- )
- if store_energy_orig and isinstance(calc, dft.rks.KohnShamDFT):
- # save initial exc for reference
- if (calc.scf_summary.get("exc") is None) or (
- grids_level != calc.grids.level
- ):
- # NOTE this should leave calc almost the same
- # as it started, but might make some
- # minor changes if there are non-default
- # settings/screenings on the grids. Also
- # scf_summary will be overwritten
- old_level = calc.grids.level
- calc.grids.level = grids_level
- calc.grids.build()
- if hasattr(calc, "with_df") and hasattr(calc.with_df, "grids"):
- calc.with_df.build()
- e_tot = calc.energy_tot(analyzer.dm)
- calc.grids.level = old_level
- calc.grids.build()
- else:
- e_tot = calc.e_tot
- analyzer._data["xc_orig"] = calc.xc
- analyzer._data["exc_orig"] = calc.scf_summary["exc"]
- analyzer._data["e_tot_orig"] = e_tot
- return analyzer
-
- def get_rho_data(self, overwrite=False):
- if "rho_data" in self.keys() and not overwrite:
- return self.get("rho_data")
- ni = dft.numint.NumInt()
- spinpol = self.dm.ndim == 3
- if spinpol:
- rho_list_a = []
- rho_list_b = []
- else:
- rho_list = []
- for ao, non0, weight, coord in ni.block_loop(self.mol, self.grids, deriv=2):
- if spinpol:
- rho_list_a.append(ni.eval_rho(self.mol, ao, self.dm[0], xctype="MGGA"))
- rho_list_b.append(ni.eval_rho(self.mol, ao, self.dm[1], xctype="MGGA"))
- else:
- rho_list.append(ni.eval_rho(self.mol, ao, self.dm, xctype="MGGA"))
- if spinpol:
- rhoa = np.concatenate(rho_list_a, axis=-1)
- rhob = np.concatenate(rho_list_b, axis=-1)
- self._data["rho_data"] = np.stack([rhoa, rhob], axis=0)
- else:
- self._data["rho_data"] = np.concatenate(rho_list, axis=-1)
- return self._data["rho_data"]
-
- @abstractmethod
- def perform_full_analysis(self):
- pass
-
- @property
- def rdm1(self):
- return self.dm
-
- @property
- def rho_data(self):
- return self.get("rho_data")
-
-
-class RHFAnalyzer(ElectronAnalyzer):
-
- _atype = "RHF"
-
- def get_xc(self, xcname):
- """
- Function for accessing XC functional data from _data
- """
- if xcname in self._data["xc_list"]:
- return self._data["xc"][xcname]
- else:
- raise ValueError("Data not computed for xc functional")
-
- def get_rs(self, term, omega, tol=1e-9):
- """
- Function for accessing range-separated data from _data
- Finds match based on a tolerance, in case rounding
- issues arise with floats.
- """
- if not (term in ["ee", "ha", "ex"]):
- raise ValueError("Term must be ee, ha, or ex")
- name = "{}_energy_density_rs".format(term)
- for ind, omega_tst in enumerate(self._data["omega_list"]):
- if abs(omega_tst - omega) < tol:
- return self._data[name][ind]
- else:
- raise RuntimeError("RS data not found for omega={}".format(omega))
-
- def get_xc_energy(self, xcname):
- """
- Store and return the XC energy for the functional given by xcname
- """
- if self._data.get("xc") is None:
- self._data["xc_list"] = []
- self._data["xc"] = {}
- ni = dft.numint.NumInt()
- if isinstance(self, UHFAnalyzer):
- _, exc, _ = ni.nr_uks(
- self.mol, self.grids, xcname, self.dm, relativity=0, hermi=1
- )
- else:
- _, exc, _ = ni.nr_rks(
- self.mol, self.grids, xcname, self.dm, relativity=0, hermi=1
- )
- self._data["xc_list"].append(xcname)
- self._data["xc"][xcname] = exc
- return exc
-
- def get_ha_energy_density(self):
- self.get_ee_energy_density()
- return self._data["ha_energy_density"]
-
- def get_ex_energy_density(self):
- self.get_ee_energy_density()
- return self._data["ex_energy_density"]
-
- get_fx_energy_density = get_ex_energy_density
-
- def _get_ej_ek(self, omega=None):
- calc_tmp = sgx_fit(scf.RHF(self.mol))
- calc_tmp.build()
- calc_tmp.with_df.grids = self.grids
- if omega is None:
- ej, ek = get_jk_densities(calc_tmp.with_df, self.dm)
- else:
- with calc_tmp.mol.with_range_coulomb(omega):
- ej, ek = get_jk_densities(calc_tmp.with_df, self.dm)
- return ej, ek
-
- def _get_ee(self, omega=None):
- ej, ek = self._get_ej_ek(omega)
- ej = ej[0]
- ek = 0.5 * ek[0]
- return ej, ek, ej + ek
-
- def get_ee_energy_density(self):
- """
- Returns the sum of E_{Ha} and E_{X}, i.e. the Coulomb repulsion
- energy of the HF or KS Slater determinant.
- """
- ej, ek, ee = self._get_ee()
- self._data["ha_energy_density"] = ej
- self._data["ex_energy_density"] = ek
- self._data["ee_energy_density"] = ee
- return self._data["ee_energy_density"]
-
- def get_ee_energy_density_rs(self, omega):
- if self._data.get("omega_list") is None:
- self._data["omega_list"] = []
- self._data["ee_energy_density_rs"] = []
- self._data["ha_energy_density_rs"] = []
- self._data["ex_energy_density_rs"] = []
- ej, ek, ee = self._get_ee(omega)
- self._data["omega_list"].append(omega)
- self._data["ha_energy_density_rs"].append(ej)
- self._data["ex_energy_density_rs"].append(ek)
- self._data["ee_energy_density_rs"].append(ee)
- return self._data["ee_energy_density_rs"][-1]
-
- def perform_full_analysis(self):
- self.get_rho_data()
- self.get_ee_energy_density()
-
- @property
- def atype(self):
- return self._atype
-
-
-RKSAnalyzer = RHFAnalyzer
-
-
-class UHFAnalyzer(RHFAnalyzer):
-
- _atype = "UHF"
-
- def _get_ej_ek(self, omega=None):
- calc_tmp = sgx_fit(scf.UHF(self.mol))
- calc_tmp.build()
- calc_tmp.with_df.grids = self.grids
- if omega is None:
- ej, ek = get_jk_densities(calc_tmp.with_df, self.dm)
- else:
- with calc_tmp.mol.with_range_coulomb(omega):
- ej, ek = get_jk_densities(calc_tmp.with_df, self.dm)
- return ej, ek
-
- def _get_ee(self, omega=None):
- ej, ek = self._get_ej_ek(omega)
- return ej, ek, ej + ek
-
-
-UKSAnalyzer = UHFAnalyzer
diff --git a/ciderpress/density.py b/ciderpress/density.py
deleted file mode 100644
index b554141..0000000
--- a/ciderpress/density.py
+++ /dev/null
@@ -1,145 +0,0 @@
-#!/usr/bin/env python
-# CiderPress: Machine-learning based density functional theory calculations
-# Copyright (C) 2024 The President and Fellows of Harvard College
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program. If not, see
-#
-# Author: Kyle Bystrom
-#
-
-import numpy
-import numpy as np
-
-LDA_FACTOR = -3.0 / 4.0 * (3.0 / np.pi) ** (1.0 / 3)
-GG_SMUL = 1.0
-GG_AMUL = 1.0
-GG_AMIN = 1.0 / 18
-CFC = (3.0 / 10) * (3 * np.pi**2) ** (2.0 / 3)
-DESC_VERSION_LIST = ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k"]
-
-UEG_VECTORS = {
- "b": np.array([1.0, 0.0, 1.0, 2.0, 2.0, 2.0, 0.0, 0.0]),
- "d": np.array([1.0, 0.0, 1.0, 2.0, 2.0, 2.0, 0.0, 0.0]),
-}
-
-
-def ldax(n):
- """
- LDA exchange energy density
- Args:
- n: Density
- """
- return LDA_FACTOR * n ** (4.0 / 3)
-
-
-def ldaxp(n):
- """
- Fully spin-polarized LDA exchange energy density
- Args:
- n: Density
- """
- return 0.5 * ldax(2 * n)
-
-
-def lsda(nu, nd):
- """
- LSDA exchange energy density
- Args:
- nu: Spin-up Density
- nd: Spin-down Density
- """
- return ldaxp(nu) + ldaxp(nd)
-
-
-def get_ldax_dens(n):
- """
- LDA exchange energy density
- Args:
- n: Density
- """
- return LDA_FACTOR * n ** (4.0 / 3)
-
-
-def get_ldax(n):
- """
- LDA Exchange energy per particle
- Args:
- n: Density
- """
- return LDA_FACTOR * n ** (1.0 / 3)
-
-
-def get_xed_from_y(y, rho):
- """
- Get the exchange energy density (n * epsilon_x)
- from the exchange enhancement factor y
- and density rho.
- """
- return rho * get_x(y, rho)
-
-
-def get_x(y, rho):
- return (y + 1) * get_ldax(rho)
-
-
-def get_y_from_xed(xed, rho):
- """
- Get the exchange enhancement factor minus one.
- """
- return xed / (get_ldax_dens(rho) - 1e-12) - 1
-
-
-def get_gradient_magnitude(rho_data):
- return np.linalg.norm(rho_data[1:4, :], axis=0)
-
-
-def get_normalized_grad(rho, mag_grad):
- sprefac = 2 * (3 * np.pi * np.pi) ** (1.0 / 3)
- n43 = rho ** (4.0 / 3)
- s = mag_grad / (sprefac * n43 + 1e-16)
- return s
-
-
-def get_single_orbital_tau(rho, mag_grad):
- return mag_grad**2 / (8 * rho + 1e-16)
-
-
-def get_uniform_tau(rho):
- return (3.0 / 10) * (3 * np.pi**2) ** (2.0 / 3) * rho ** (5.0 / 3)
-
-
-def get_regularized_tau(tau, tau_w, tau_unif):
- alpha = (tau - tau_w) / (tau_unif + 1e-4)
- return alpha**3 / (alpha**2 + 1e-3)
-
-
-def get_normalized_tau(tau, tau_w, tau_unif):
- return (tau - tau_w) / (tau_unif + 1e-16)
-
-
-def get_dft_input(rho_data):
- rho = rho_data[0, :]
- mag_grad = get_gradient_magnitude(rho_data)
- s = get_normalized_grad(rho, mag_grad)
- tau_w = get_single_orbital_tau(rho, mag_grad)
- tau_unif = get_uniform_tau(rho)
- alpha = get_normalized_tau(rho_data[5], tau_w, tau_unif)
- return rho, s, alpha, tau_w, tau_unif
-
-
-def get_dft_input_gga(rho_data):
- rho = rho_data[0, :]
- mag_grad = get_gradient_magnitude(rho_data)
- s = get_normalized_grad(rho, mag_grad)
- return rho, s
diff --git a/ciderpress/lib/__init__.py b/ciderpress/lib/__init__.py
index 744da5d..7281493 100644
--- a/ciderpress/lib/__init__.py
+++ b/ciderpress/lib/__init__.py
@@ -2,6 +2,8 @@
from ciderpress.lib.load import load_library
+__all__ = ["load_library", "c_double_p", "c_int_p", "c_null_ptr"]
+
c_double_p = ctypes.POINTER(ctypes.c_double)
c_int_p = ctypes.POINTER(ctypes.c_int)
c_null_ptr = ctypes.POINTER(ctypes.c_void_p)