def write_phmsd(fh5, occa, occb, nelec, norb, init=None, orbmat=None):
"""Write NOMSD to HDF.
Parameters
----------
fh5 : h5py group
Wavefunction group to write to file.
nelec : tuple
Number of alpha and beta electrons.
"""
# TODO: Update if we ever wanted "mixed" phmsd type wavefunctions.
na, nb = nelec
if init is not None:
add_dataset(fh5, 'Psi0_alpha', to_qmcpack_complex(init[0]))
add_dataset(fh5, 'Psi0_beta', to_qmcpack_complex(init[1]))
else:
init = numpy.eye(norb, dtype=numpy.complex128)
add_dataset(fh5, 'Psi0_alpha',
to_qmcpack_complex(init[:,occa[0]].copy()))
add_dataset(fh5, 'Psi0_beta',
to_qmcpack_complex(init[:,occb[0]].copy()))
if orbmat is not None:
fh5['type'] = 1
# Expects conjugate transpose.
oa = scipy.sparse.csr_matrix(orbmat[0].conj().T)
write_nomsd_single(fh5, oa, 0)
ob = scipy.sparse.csr_matrix(orbmat[1].conj().T)
write_nomsd_single(fh5, ob, 1)
else:
fh5['type'] = 0
occs = numpy.zeros((len(occa), na+nb), dtype=numpy.int32)
occs[:,:na] = numpy.array(occa)
occs[:,na:] = norb+numpy.array(occb)
# Reading 1D array currently in qmcpack.
fh5['occs'] = occs.ravel()
def write_nomsd(fh5, wfn, uhf, nelec, thresh=1e-8, init=None):
"""Write NOMSD to HDF.
Parameters
----------
fh5 : h5py group
Wavefunction group to write to file.
wfn : :class:`numpy.ndarray`
NOMSD trial wavefunctions.
uhf : bool
UHF style wavefunction.
nelec : tuple
Number of alpha and beta electrons.
thresh : float
Threshold for writing wavefunction elements.
"""
nalpha, nbeta = nelec
nmo = wfn.shape[1]
wfn[abs(wfn) < thresh] = 0.0
if init is not None:
add_dataset(fh5, 'Psi0_alpha', to_qmcpack_complex(init[0]))
add_dataset(fh5, 'Psi0_beta', to_qmcpack_complex(init[1]))
else:
add_dataset(fh5, 'Psi0_alpha',
to_qmcpack_complex(wfn[0,:,:nalpha].copy()))
if uhf:
add_dataset(fh5, 'Psi0_beta',
to_qmcpack_complex(wfn[0,:,nalpha:].copy()))
for idet, w in enumerate(wfn):
# QMCPACK stores this internally as a csr matrix, so first convert.
ix = 2*idet if uhf else idet
psia = scipy.sparse.csr_matrix(w[:,:nalpha].conj().T)
write_nomsd_single(fh5, psia, ix)
if uhf:
ix = 2*idet + 1
psib = scipy.sparse.csr_matrix(w[:,nalpha:].conj().T)
write_nomsd_single(fh5, psib, ix)