Skip to content

Commit c311d16

Browse files
authored
fix: only apply comp_prec for floating dtypes (#711)
Fix #703. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Improved handling of data types during the reshaping process, ensuring type conversion only occurs for floating-point data. - **Bug Fixes** - Enhanced robustness of data processing by preventing unnecessary type casting for non-floating-point data types. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Jinzhe Zeng <[email protected]>
1 parent 6bf41e3 commit c311d16

File tree

2 files changed

+7
-4
lines changed

2 files changed

+7
-4
lines changed

dpdata/deepmd/comp.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -204,7 +204,9 @@ def dump(folder, data, set_size=5000, comp_prec=np.float32, remove_sets=True):
204204
f"Shape of {dtype.name} is not (nframes, ...), but {dtype.shape}. This type of data will not converted to deepmd/npy format."
205205
)
206206
continue
207-
ddata = np.reshape(data[dtype.name], [nframes, -1]).astype(comp_prec)
207+
ddata = np.reshape(data[dtype.name], [nframes, -1])
208+
if np.issubdtype(ddata.dtype, np.floating):
209+
ddata = ddata.astype(comp_prec)
208210
for ii in range(nsets):
209211
set_stt = ii * set_size
210212
set_end = (ii + 1) * set_size

dpdata/deepmd/hdf5.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -252,9 +252,10 @@ def dump(
252252
for dt, prop in data_types.items():
253253
if dt in data:
254254
if prop["dump"]:
255-
reshaped_data[dt] = np.reshape(data[dt], prop["shape"]).astype(
256-
comp_prec
257-
)
255+
ddata = np.reshape(data[dt], prop["shape"])
256+
if np.issubdtype(ddata.dtype, np.floating):
257+
ddata = ddata.astype(comp_prec)
258+
reshaped_data[dt] = ddata
258259

259260
# dump frame properties: cell, coord, energy, force and virial
260261
nsets = nframes // set_size

0 commit comments

Comments
 (0)