Bug in nbath keyword in DMET_Bath

[obackhouse]

The nbath keyword in DMET_Bath is broken. Example:

from pyscf.pbc import gto, scf, tools
from vayesta.core import QEmbeddingMethod
from vayesta.core.bath import DMET_Bath
from vayesta.misc import gdf
import numpy as np

L = 5.0
n = 11
cell = gto.Cell()
cell.a = np.eye(3) * L
cell.mesh = [n, n, n]
cell.atom = 'He 3 2 3; He 1 1 1'
cell.basis = 'cc-pvdz'
cell.exp_to_discard = 0.1
cell.verbose = 0
cell.max_memory = 1e9
cell.build()

kpts = cell.make_kpts([2,2,2])

mf = scf.KRHF(cell, kpts)
mf.conv_tol = 1e-12
mf.with_df = gdf.GDF(cell, kpts)
mf.with_df.linear_dep_threshold = 1e-7
mf.with_df.build()
mf.kernel()

qemb = QEmbeddingMethod(mf)
qemb.sao_fragmentation()
frag = qemb.add_atomic_fragment([0])

bath = DMET_Bath(frag, frag.opts.dmet_threshold)
c_bath, c_occenv, c_virenv = bath.make_dmet_bath(frag.c_env)

nbath = c_bath.shape[-1]

c_bath, c_occenv, c_virenv = bath.make_dmet_bath(frag.c_env, nbath=nbath)

Which raises the error:

Traceback (most recent call last):
  File "bugs/dmet_nbath.py", line 37, in <module>
    c_bath, c_occenv, c_virenv = bath.make_dmet_bath(frag.c_env, nbath=nbath)
  File "/home/olli/git/vayesta/vayesta/core/bath/dmet.py", line 163, in make_dmet_bath
    assert (nbath + noccenv + nvirenv == c_env.shape[-1])
AssertionError