Source code for lib.chkfile

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import json
import h5py

[docs]def load(chkfile, key): '''Load array(s) from chkfile Args: chkfile : str Name of chkfile. The chkfile needs to be saved in HDF5 format. key : str HDF5.dataset name or group name. If key is the name of a HDF5 group, the group will be loaded into a Python dict, recursively. Returns: whatever read from chkfile Examples: >>> from pyscf import gto, scf, lib >>> mol = gto.M(atom='He 0 0 0') >>> mf = scf.RHF(mol) >>> mf.chkfile = 'He.chk' >>> mf.kernel() >>> mo_coeff = lib.chkfile.load('He.chk', 'scf/mo_coeff') >>> mo_coeff.shape (1, 1) >>> scfdat = lib.chkfile.load('He.chk', 'scf') >>> scfdat.keys() ['e_tot', 'mo_occ', 'mo_energy', 'mo_coeff'] ''' def load_as_dic(key, group): if key in group: val = group[key] elif key + '__from_list__' in group: key = key + '__from_list__' val = group[key] else: return None if isinstance(val, h5py.Group): if key.endswith('__from_list__'): return [load_as_dic(k, val) for k in val] else: return dict([(k.replace('__from_list__', ''), load_as_dic(k, val)) for k in val]) else: return val.value with h5py.File(chkfile, 'r') as fh5: return load_as_dic(key, fh5)
load_chkfile_key = load def dump(chkfile, key, value): '''Save array(s) in chkfile Args: chkfile : str Name of chkfile. key : str key to be used in h5py object. It can contain "/" to represent the path in the HDF5 storage structure. value : array, vector, list ... or dict If value is a python dict or list, the key/value of the dict will be saved recursively as the HDF5 group/dataset structure. Returns: No return value Examples: >>> import h5py >>> from pyscf import lib >>> ci = {'Ci' : {'op': ('E', 'i'), 'irrep': ('Ag', 'Au')}} >>> lib.chkfile.save('symm.chk', 'symm', ci) >>> f = h5py.File('symm.chk') >>> f.keys() ['symm'] >>> f['symm'].keys() ['Ci'] >>> f['symm/Ci'].keys() ['op', 'irrep'] >>> f['symm/Ci/op'] <HDF5 dataset "op": shape (2,), type "|S1"> ''' def save_as_group(key, value, root): if isinstance(value, dict): root1 = root.create_group(key) for k in value: save_as_group(k, value[k], root1) elif isinstance(value, (tuple, list)): root1 = root.create_group(key + '__from_list__') for k, v in enumerate(value): save_as_group('%06d'%k, v, root1) else: try: root[key] = value except (TypeError, ValueError) as e: if not (e.args[0] == "Object dtype dtype('O') has no native HDF5 equivalent" or e.args[0].startswith('could not broadcast input array')): raise e root1 = root.create_group(key + '__from_list__') for k, v in enumerate(value): save_as_group('%06d'%k, v, root1) if h5py.is_hdf5(chkfile): with h5py.File(chkfile, 'r+') as fh5: if key in fh5: del(fh5[key]) elif key + '__from_list__' in fh5: del(fh5[key+'__from_list__']) save_as_group(key, value, fh5) else: with h5py.File(chkfile, 'w') as fh5: save_as_group(key, value, fh5) dump_chkfile_key = save = dump
[docs]def load_mol(chkfile): '''Load Mole object from chkfile. The save_mol/load_mol operation can be used a serialization method for Mole object. Args: chkfile : str Name of chkfile. Returns: A (initialized/built) Mole object Examples: >>> from pyscf import gto, lib >>> mol = gto.M(atom='He 0 0 0') >>> lib.chkfile.save_mol(mol, 'He.chk') >>> lib.chkfile.load_mol('He.chk') <pyscf.gto.mole.Mole object at 0x7fdcd94d7f50> ''' from numpy import array from pyscf import gto try: with h5py.File(chkfile, 'r') as fh5: mol = gto.loads(fh5['mol'].value) except: # Compatibility to the old serialization format # TODO: remove it in future release with h5py.File(chkfile, 'r') as fh5: mol = gto.Mole() mol.output = '/dev/null' moldic = eval(fh5['mol'].value) for key in ('mass', 'grids', 'light_speed'): if key in moldic: del(moldic[key]) mol.build(False, False, **moldic) return mol
[docs]def save_mol(mol, chkfile): '''Save Mole object in chkfile Args: chkfile str: Name of chkfile. Returns: No return value ''' dump(chkfile, 'mol', mol.dumps())
dump_mol = save_mol