mirror of
https://github.com/TREX-CoE/trexio.git
synced 2024-08-25 14:41:44 +02:00
153 lines
4.3 KiB
Python
153 lines
4.3 KiB
Python
import os
|
|
import shutil
|
|
#import numpy as np
|
|
|
|
from pytrexio import *
|
|
|
|
# TODO: make a user-friendly more pythonic API that will have to be autogenerated
|
|
# add Exception handling
|
|
# check of dimensions and call to safe API
|
|
# conversion to and from numpy arrays
|
|
|
|
# automatically download (hopefully the latest version) numpy.i using
|
|
# wget https://raw.githubusercontent.com/numpy/numpy/main/tools/swig/numpy.i
|
|
# but also keep the original copy in case if wget fails
|
|
|
|
#=========================================================#
|
|
#======== SETUP THE BACK END AND OUTPUT FILE NAME ========#
|
|
#=========================================================#
|
|
|
|
# 0: TREXIO_HDF5 ; 1: TREXIO_TEXT
|
|
TEST_TREXIO_BACKEND = TREXIO_TEXT
|
|
OUTPUT_FILENAME_TEXT = 'test_py_swig.dir'
|
|
OUTPUT_FILENAME_HDF5 = 'test_py_swig.h5'
|
|
|
|
|
|
if TEST_TREXIO_BACKEND == TREXIO_HDF5:
|
|
output_filename = OUTPUT_FILENAME_HDF5
|
|
elif TEST_TREXIO_BACKEND == TREXIO_TEXT:
|
|
output_filename = OUTPUT_FILENAME_TEXT
|
|
else:
|
|
raise ValueError ('Specify one of the supported back ends as TEST_TREXIO_BACKEND')
|
|
|
|
|
|
try:
|
|
if TEST_TREXIO_BACKEND == TREXIO_HDF5:
|
|
os.remove(output_filename)
|
|
elif TEST_TREXIO_BACKEND == TREXIO_TEXT:
|
|
shutil.rmtree(output_filename)
|
|
except:
|
|
print (f'Test file {output_filename} does not exist')
|
|
|
|
#=========================================================#
|
|
#============ WRITE THE DATA IN THE TEST FILE ============#
|
|
#=========================================================#
|
|
|
|
test_file = trexio_open(output_filename, 'w', TEST_TREXIO_BACKEND)
|
|
|
|
nucleus_num = 12
|
|
|
|
rc = trexio_write_nucleus_num(test_file, nucleus_num)
|
|
assert rc==TREXIO_SUCCESS
|
|
|
|
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)
|
|
assert rc==TREXIO_SUCCESS
|
|
|
|
point_group = 'B3U'
|
|
|
|
rc = trexio_write_nucleus_point_group(test_file, point_group, 10)
|
|
assert rc==TREXIO_SUCCESS
|
|
|
|
labels = [
|
|
'C',
|
|
'C',
|
|
'C',
|
|
'C',
|
|
'C',
|
|
'C',
|
|
'H',
|
|
'H',
|
|
'H',
|
|
'H',
|
|
'H',
|
|
'H']
|
|
|
|
rc = trexio_write_nucleus_label(test_file, labels, 10)
|
|
assert rc==TREXIO_SUCCESS
|
|
|
|
rc = trexio_close(test_file)
|
|
assert rc==TREXIO_SUCCESS
|
|
|
|
#==========================================================#
|
|
#============ READ THE DATA FROM THE TEST FILE ============#
|
|
#==========================================================#
|
|
|
|
test_file2 = trexio_open(output_filename, 'r', TEST_TREXIO_BACKEND)
|
|
|
|
# TODO: think about adding exception handling if rc != TREXIO_SUCCESS throw an exception but do not return the code itself and then there is less arguments
|
|
# TODO: maybe also for the write
|
|
|
|
result = trexio_read_nucleus_num(test_file2)
|
|
assert result[0]==TREXIO_SUCCESS
|
|
assert result[1]==nucleus_num
|
|
#print(result)
|
|
|
|
charges2 = doubleArray(nucleus_num)
|
|
print(charges2[3])
|
|
#for i in range(nucleus_num):
|
|
# charges2[i] = -1.
|
|
|
|
rc = trexio_read_nucleus_charge(test_file2, charges2)
|
|
assert rc==TREXIO_SUCCESS
|
|
for i in range(nucleus_num):
|
|
assert charges2[i]==charges[i]
|
|
|
|
#charge_numpy = np.zeros(nucleus_num, dtype=np.float64)
|
|
#print(charge_numpy)
|
|
|
|
rc, charge_numpy = trexio_read_safe_nucleus_charge(test_file2, 12)
|
|
|
|
print(charge_numpy)
|
|
print(charge_numpy[11])
|
|
assert rc==TREXIO_SUCCESS
|
|
|
|
# unsafe call to read_safe should not only have return code = TREXIO_UNSAFE_ARRAY_DIM
|
|
# but also should not return numpy array filled with garbage
|
|
rc, charge_numpy = trexio_read_safe_nucleus_charge(test_file2, 12*5)
|
|
|
|
#print(charge_numpy)
|
|
assert rc==TREXIO_UNSAFE_ARRAY_DIM
|
|
|
|
# [WIP]: ideally, the list of strings should be returned as below
|
|
#rc, label_2d = trexio_read_nucleus_label(test_file2, 10)
|
|
# [WIP]: 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==TREXIO_SUCCESS
|
|
|
|
labels_2d = [label for label in labels_1d.split(TREXIO_DELIM) if label]
|
|
print(labels_2d)
|
|
for i in range(nucleus_num):
|
|
assert labels_2d[i]==labels[i]
|
|
|
|
rc = trexio_close(test_file2)
|
|
assert rc==TREXIO_SUCCESS
|
|
|
|
try:
|
|
if TEST_TREXIO_BACKEND == TREXIO_HDF5:
|
|
os.remove(output_filename)
|
|
elif TEST_TREXIO_BACKEND == TREXIO_TEXT:
|
|
shutil.rmtree(output_filename)
|
|
except:
|
|
print (f'No output file {output_filename} has been produced')
|
|
|
|
#==========================================================#
|
|
|