lindi/tests/test_numpy_array_conversion.py at fccbdf6cf7af1736c71809908d944a1a709f1740 · NeurodataWithoutBorders/lindi · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import tempfile
import numpy as np
import h5py
import lindi


def test_numpy_array_conversion():
    """Test that LindiH5pyDataset supports np.asarray() and np.atleast_2d().

    Regression test for https://github.com/NeurodataWithoutBorders/lindi/issues/120
    """
    with tempfile.TemporaryDirectory() as tmpdir:
        h5_fname = f'{tmpdir}/test.h5'
        lindi_json_fname = f'{tmpdir}/test.lindi.json'

        # Create a simple HDF5 file with a 1D dataset
        with h5py.File(h5_fname, 'w') as f:
            f.create_dataset('data', data=np.arange(100, dtype=np.float64))
            f.create_dataset('scalar', data=42.0)

        # Convert to lindi
        with lindi.LindiH5pyFile.from_hdf5_file(h5_fname, url=h5_fname) as f:
            f.write_lindi_file(lindi_json_fname)

        # Open the lindi file and test numpy conversions
        with lindi.LindiH5pyFile.from_lindi_file(lindi_json_fname) as f:
            ds = f['data']

            # Test _is_empty
            assert ds._is_empty is False

            # Test np.asarray - this triggers __array__
            arr = np.asarray(ds)
            assert arr.shape == (100,)
            np.testing.assert_array_equal(arr, np.arange(100, dtype=np.float64))

            # Test np.atleast_2d - this is what failed in the issue
            arr2d = np.atleast_2d(ds)
            assert arr2d.shape == (1, 100)

            # Test np.array with dtype conversion
            arr_int = np.array(ds, dtype=np.int32)
            assert arr_int.dtype == np.int32

            # Test scalar dataset
            sc = f['scalar']
            arr_sc = np.asarray(sc)
            assert arr_sc.shape == ()
            assert float(arr_sc) == 42.0


def test_numpy_array_conversion_compound():
    """Test that __array__ works for compound dtype datasets."""
    with tempfile.TemporaryDirectory() as tmpdir:
        h5_fname = f'{tmpdir}/test.h5'
        lindi_json_fname = f'{tmpdir}/test.lindi.json'

        compound_dtype = np.dtype([('x', np.int32), ('y', np.float64)])
        data = np.array([(1, 2.5), (3, 4.5), (5, 6.5)], dtype=compound_dtype)

        with h5py.File(h5_fname, 'w') as f:
            f.create_dataset('compound', data=data)

        with lindi.LindiH5pyFile.from_hdf5_file(h5_fname, url=h5_fname) as f:
            f.write_lindi_file(lindi_json_fname)

        with lindi.LindiH5pyFile.from_lindi_file(lindi_json_fname) as f:
            ds = f['compound']
            assert ds._is_empty is False

            # Test field access followed by np.asarray
            x_vals = np.asarray(ds['x'][:])
            np.testing.assert_array_equal(x_vals, np.array([1, 3, 5], dtype=np.int32))

            y_vals = np.asarray(ds['y'][:])
            np.testing.assert_array_equal(y_vals, np.array([2.5, 4.5, 6.5], dtype=np.float64))

            # Test integer indexing - returns np.void
            row0 = ds[0]
            assert isinstance(row0, np.void)
            assert row0['x'] == 1
            assert row0['y'] == 2.5

            # Test slice indexing - returns structured array
            rows = ds[0:2]
            assert rows.dtype == compound_dtype
            assert len(rows) == 2
            np.testing.assert_array_equal(rows['x'], np.array([1, 3], dtype=np.int32))
            np.testing.assert_array_equal(rows['y'], np.array([2.5, 4.5], dtype=np.float64))