8000 Compatibility with `pyscfad` by obackhouse · Pull Request #84 · BoothGroup/ebcc · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

Compatibility with pyscfad #84

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 11 commits into from
Nov 24, 2024
Merged

Compatibility with pyscfad #84

merged 11 commits into from
Nov 24, 2024

Conversation

obackhouse
Copy link
Contributor

Currently incredibly hacky support for pyscfad with jax backends. Will take a lot of refactoring to properly support.

@obackhouse
Copy link
Contributor Author

@mkrompiec here is an example script (copied from a pyscfad example) comparing RCCSD to the pyscfad implementation and then doing the same thing with UCCSD:

import jax
import sys
import os
os.environ["EBCC_BACKEND"] = "jax"
from pyscfad import gto, scf, cc, config
from ebcc import EBCC, default_log
from ebcc.util import Namespace

config.update("pyscfad_scf_implicit_diff", True)
config.update("pyscfad_ccsd_implicit_diff", True)

mol = gto.Mole(
    atom="H 0 0 0; F 0 0 1",
    basis="cc-pvdz",
    verbose=0,
)
mol.build()

def ccsd(mol, dm0=None):
    mf = scf.RHF(mol)
    mf.kernel(dm0)
    ccsd = cc.RCCSD(mf)
    ccsd.kernel()
    return ccsd.e_tot

jac = jax.jacrev(ccsd)(mol)
print("*** PySCF-AD ***")
print(f"Nuclear gradient:\n{jac.coords}")
print(f"Gradient wrt basis exponents:\n{jac.exp}")
print(f"Gradient wrt basis contraction coefficients:\n{jac.ctr_coeff}")

def ccsd(mol, dm0=None):
    mf = scf.RHF(mol)
    mf.kernel(dm0)
    ccsd = EBCC(mf, ansatz="CCSD")
    ccsd.kernel()
    return ccsd.e_tot

jac = jax.jacrev(ccsd)(mol)
print("*** PySCF-AD + ebcc ***")
print(f"Nuclear gradient:\n{jac.coords}")
print(f"Gradient wrt basis exponents:\n{jac.exp}")
print(f"Gradient wrt basis contraction coefficients:\n{jac.ctr_coeff}")

def uccsd(mol, dm0=None):
    mf = scf.UHF(mol)
    mf.kernel(dm0)
    ccsd = EBCC(mf, ansatz="CCSD")
    ccsd.kernel()
    return ccsd.e_tot

jac = jax.jacrev(uccsd)(mol)
print("*** PySCF-AD + ebcc (UHF) ***")
print(f"Nuclear gradient:\n{jac.coords}")
print(f"Gradient wrt basis exponents:\n{jac.exp}")
print(f"Gradient wrt basis contraction coefficients:\n{jac.ctr_coeff}")

The gradients are in good agreement with pyscfad for the restricted case.

Note that I default to 64 bit jax:

ebcc/ebcc/backend/_jax.py

Line 8 in b687c61

jax.config.update("jax_enable_x64", True)

Let me know if you have any thoughts, if this is useful, if you need any other functionality etc.

@mkrompiec
Copy link

Thanks a lot! The example script works for me.

@obackhouse obackhouse marked this pull request as ready for review November 23, 2024 13:21
@obackhouse
Copy link
Contributor Author

This can get merged, but I'll wait on fishjojo/pyscfad#60 to be fixed rather than using the github link for the installation.

@codecov Codecov
Copy link
codecov bot commented Nov 23, 2024 •
edited
Loading

Codecov Report

Attention: Patch coverage is 48.93617% with 24 lines in your changes missing coverage. Please review.

Project coverage is 89.66%. Comparing base (c212d48) to head (c0b9575).
Report is 12 commits behind head on master.

Files with missing lines Patch % Lines
ebcc/__init__.py 7.69% 11 Missing and 1 partial ⚠️
ebcc/ham/eris.py 50.00% 7 Missing ⚠️
ebcc/backend/__init__.py 0.00% 2 Missing and 1 partial ⚠️
ebcc/core/precision.py 60.00% 1 Missing and 1 partial ⚠️
Additional details and impacted files 
@@            Coverage Diff             @@
##           master      #84      +/-   ##
==========================================
- Coverage   89.67%   89.66%   -0.02%     
==========================================
  Files          94       94              
  Lines      112751   112768      +17     
  Branches      662      665       +3     
==========================================
+ Hits       101110   101111       +1     
- Misses      11436    11450      +14     
- Partials      205      207       +2

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚨 Try these New Features:

@obackhouse
Copy link
Contributor Author

This can get merged, but I'll wait on fishjojo/pyscfad#60 to be fixed rather than using the github link for the installation.

nvm my bad, just need to hack in the imports. I fixed the logging now and I'm happy to merge, it's not completely covered but I don't require that for non-numpy backends

@obackhouse obackhouse merged commit 3948c83 into master Nov 24, 2024
6 checks passed
@obackhouse obackhouse deleted the pyscfad_compatibility branch November 24, 2024 11:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@obackhouse @mkrompiec
0