mirror of
https://github.com/TREX-CoE/trexio.git
synced 2024-10-31 11:13:56 +01:00
230 lines
7.5 KiB
Python
230 lines
7.5 KiB
Python
import os
|
|
import shutil
|
|
import numpy as np
|
|
|
|
from pytrexio.pytrexio import *
|
|
|
|
#=========================================================#
|
|
#======== SETUP THE BACK END AND OUTPUT FILE NAME ========#
|
|
#=========================================================#
|
|
|
|
# 0: TREXIO_HDF5 ; 1: TREXIO_TEXT
|
|
TEST_TREXIO_BACKEND = 0
|
|
OUTPUT_FILENAME_TEXT = 'test_py_swig.dir'
|
|
OUTPUT_FILENAME_HDF5 = 'test_py_swig.h5'
|
|
|
|
|
|
if TEST_TREXIO_BACKEND == 0:
|
|
output_filename = OUTPUT_FILENAME_HDF5
|
|
elif TEST_TREXIO_BACKEND == 1:
|
|
output_filename = OUTPUT_FILENAME_TEXT
|
|
else:
|
|
raise ValueError ('Specify one of the supported back ends as TEST_TREXIO_BACKEND')
|
|
|
|
|
|
try:
|
|
if TEST_TREXIO_BACKEND == 0:
|
|
os.remove(output_filename)
|
|
elif TEST_TREXIO_BACKEND == 1:
|
|
shutil.rmtree(output_filename)
|
|
except:
|
|
print ('Nothing to remove.')
|
|
|
|
#=========================================================#
|
|
#============ WRITE THE DATA IN THE TEST FILE ============#
|
|
#=========================================================#
|
|
|
|
return_obj = trexio_open(output_filename, 'w', TEST_TREXIO_BACKEND)
|
|
assert return_obj is not None
|
|
if isinstance(return_obj, int):
|
|
print(trexio_string_of_error(return_obj))
|
|
assert return_obj==0
|
|
else:
|
|
rc_open = return_obj[1]
|
|
assert rc_open==0
|
|
test_file = return_obj[0]
|
|
assert test_file is not None
|
|
|
|
nucleus_num = 12
|
|
|
|
rc = trexio_write_nucleus_num(test_file, nucleus_num)
|
|
assert rc==0
|
|
|
|
# initialize charge arrays as a list and convert it to numpy array
|
|
charges = [6., 6., 6., 6., 6., 6., 1., 1., 1., 1., 1., 1.]
|
|
charges_np = np.array(charges, dtype=np.float64)
|
|
|
|
# function call below works with both lists and numpy arrays, dimension needed for memory-safety is derived
|
|
# from the size of the list/array by SWIG using typemaps from numpy.i
|
|
rc = trexio_write_safe_nucleus_charge(test_file, charges_np)
|
|
assert rc==0
|
|
|
|
# less Python-ic way to read/write arrays using Array classes (probably more portable to other languages)
|
|
#charges = doubleArray(nucleus_num)
|
|
#for i in range(nucleus_num):
|
|
# if i < nucleus_num/2:
|
|
# charges[i] = 6.
|
|
# else:
|
|
# charges[i] = 1.
|
|
#rc = trexio_write_nucleus_charge(test_file, charges)
|
|
|
|
basis_num = 24
|
|
# initialize arrays of nuclear indices as a list and convert it to numpy array
|
|
indices = [i for i in range(basis_num)]
|
|
# type cast is important here because by default numpy transforms a list of integers into int64 array
|
|
indices_np = np.array(indices, dtype=np.int32)
|
|
|
|
# first write basis_num because it is needed to check dimensions of basis_nucleus_index in TREXIO >= 2.0.0
|
|
rc = trexio_write_basis_shell_num(test_file, basis_num)
|
|
assert rc==0
|
|
|
|
# function call below works with both lists and numpy arrays, dimension needed for memory-safety is derived
|
|
# from the size of the list/array by SWIG using typemacs from numpy.i
|
|
rc = trexio_write_safe_basis_nucleus_index(test_file, indices_np)
|
|
assert rc==0
|
|
|
|
# test writing of sparse data
|
|
rc = trexio_write_mo_num(test_file, 600)
|
|
assert rc==0
|
|
|
|
num_integrals = 100
|
|
indices = [i for i in range(num_integrals*4)]
|
|
values = [(3.14 + float(i)) for i in range(num_integrals)]
|
|
|
|
rc = trexio_write_safe_mo_2e_int_eri(test_file, 0, num_integrals, indices, values)
|
|
assert rc==0
|
|
|
|
# test writing of single string
|
|
point_group = 'B3U'
|
|
|
|
rc = trexio_write_nucleus_point_group(test_file, point_group, 10)
|
|
assert rc==0
|
|
|
|
# test writing of array of strings
|
|
labels = [
|
|
'C',
|
|
'C',
|
|
'C',
|
|
'C',
|
|
'C',
|
|
'C',
|
|
'H',
|
|
'H',
|
|
'H',
|
|
'H',
|
|
'H',
|
|
'H']
|
|
|
|
rc = trexio_write_nucleus_label(test_file, labels, 10)
|
|
assert rc==0
|
|
|
|
rc = trexio_close(test_file)
|
|
assert rc==0
|
|
|
|
#==========================================================#
|
|
#============ READ THE DATA FROM THE TEST FILE ============#
|
|
#==========================================================#
|
|
|
|
return_obj = trexio_open(output_filename, 'r', TEST_TREXIO_BACKEND)
|
|
assert return_obj is not None
|
|
if isinstance(return_obj, int):
|
|
print(trexio_string_of_error(return_obj))
|
|
assert return_obj==0
|
|
else:
|
|
rc_open = return_obj[1]
|
|
assert rc_open==0
|
|
test_file2 = return_obj[0]
|
|
assert test_file is not None
|
|
|
|
result = trexio_read_nucleus_num(test_file2)
|
|
assert result[0]==0
|
|
assert result[1]==nucleus_num
|
|
|
|
# safe call to read_safe array of float values
|
|
rc, rcharges_np = trexio_read_safe_nucleus_charge(test_file2, nucleus_num)
|
|
assert rc==0
|
|
assert rcharges_np.dtype is np.dtype(np.float64)
|
|
np.testing.assert_array_almost_equal(rcharges_np, charges_np, decimal=8)
|
|
|
|
# unsafe call to read_safe should not only have return code = TREXIO_UNSAFE_ARRAY_DIM
|
|
# TODO: it should not return numpy array filled with garbage
|
|
rc, rcharges_fail = trexio_read_safe_nucleus_charge(test_file2, nucleus_num*5)
|
|
assert rc==23
|
|
|
|
# less Python-ic way to read/write arrays using Array classes (probably more portable to other languages)
|
|
#charges2 = doubleArray(nucleus_num)
|
|
#for i in range(nucleus_num):
|
|
# charges2[i] = -1.
|
|
#rc = trexio_read_nucleus_charge(test_file2, charges2)
|
|
#assert rc==0
|
|
#for i in range(nucleus_num):
|
|
# assert charges2[i]==charges[i]
|
|
|
|
result = trexio_read_basis_shell_num(test_file2)
|
|
assert result[0]==0
|
|
assert result[1]==basis_num
|
|
|
|
# safe call to read_safe array of int values
|
|
rc, rindices_np = trexio_read_safe_basis_nucleus_index(test_file2, basis_num)
|
|
assert rc==0
|
|
assert rindices_np.dtype is np.dtype(np.int32)
|
|
for i in range(nucleus_num):
|
|
assert rindices_np[i]==indices_np[i]
|
|
|
|
# currently only low-level routines (return one long string instead of an array of strings) work
|
|
rc, labels_1d = trexio_read_nucleus_label_low(test_file2, 10)
|
|
assert rc==0
|
|
|
|
labels_2d = [label for label in labels_1d.split(TREXIO_DELIM) if label]
|
|
print('Read and parsed nuclear labels:\n', labels_2d)
|
|
for i in range(nucleus_num):
|
|
assert labels_2d[i]==labels[i]
|
|
|
|
rc, rpoint_group = trexio_read_nucleus_point_group(test_file2, 3)
|
|
print(f'Read point group: {rpoint_group}')
|
|
assert rc==0
|
|
assert rpoint_group==point_group
|
|
|
|
# test reasing sparse quantities
|
|
rc, mo_2e_int_size = trexio_read_mo_2e_int_eri_size(test_file2)
|
|
assert rc==0
|
|
assert mo_2e_int_size==num_integrals
|
|
|
|
buf_size = 60
|
|
offset_file = 0
|
|
# read full buf_size (i.e. the one that does not reach EOF)
|
|
rc, read_buf_size, indices_sparse_np, value_sparse_np = trexio_read_safe_mo_2e_int_eri(test_file2, offset_file, buf_size, buf_size*4, buf_size)
|
|
print(f'First complete sparse read size: {read_buf_size}')
|
|
#print(indices_sparse_np)
|
|
assert rc==0
|
|
assert read_buf_size==buf_size
|
|
assert indices_sparse_np[0]==0
|
|
assert indices_sparse_np[read_buf_size*4-1]==read_buf_size*4-1
|
|
offset_file += buf_size
|
|
|
|
# read incomplete buf_size (i.e. the one that does reach EOF)
|
|
rc, read_buf_size, indices_sparse_np, value_sparse_np = trexio_read_safe_mo_2e_int_eri(test_file2, offset_file, buf_size, buf_size*4, buf_size)
|
|
print(f'Second incomplete sparse read size: {read_buf_size}')
|
|
# Incomplete read still allocates NumPy array of buf_size=60 elements but only 40 elements read upon encounter of EOF,
|
|
# Thus the remaining 20 elements are filled with garbage rather than zeros. Handle this in the front end ?
|
|
print(indices_sparse_np)
|
|
# trexio_exit_code = 6 correspond to TREXIO_END
|
|
assert rc==6
|
|
assert read_buf_size==(num_integrals - buf_size)
|
|
assert indices_sparse_np[0]==offset_file*4
|
|
assert indices_sparse_np[read_buf_size*4-1]==(offset_file+read_buf_size)*4-1
|
|
|
|
|
|
rc = trexio_close(test_file2)
|
|
assert rc==0
|
|
|
|
try:
|
|
if TEST_TREXIO_BACKEND == 0:
|
|
os.remove(output_filename)
|
|
elif TEST_TREXIO_BACKEND == 1:
|
|
shutil.rmtree(output_filename)
|
|
except:
|
|
print (f'No output file {output_filename} has been produced')
|
|
|
|
#==========================================================#
|