brain-score / brainio

Data management for quantitative comparison of brains and brain-inspired systems
MIT License
5 stars 12 forks source link

clip returns DataArray instead of original type #27

Open mschrimpf opened 2 years ago

mschrimpf commented 2 years ago

As of xarray 0.18, the clip method uses the apply_ufunc function (https://github.com/pydata/xarray/compare/v0.17.0...v0.18.0#diff-5e1c0e7a002249e30f1e688e4e87af2d67600ce8a240f8b27dccdb28f365c971R381) which in turn uses apply_dataarray_vfunc and ends up returning a DataArray https://github.com/pydata/xarray/blob/v0.18.0/xarray/core/computation.py#L280.

Is there a way to preserve the assembly's original data type? E.g. if we call <BehavioralAssembly>.clip(), we get BehavioralAssembly as output again.

Here's a unit test that fails in 0.18 but works in 0.17:

    def test_clip(self):
        assembly = BehavioralAssembly([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
                                      coords={'a': ['a', 'b', 'c', 'd'],
                                              'aa': ('a', ['a', 'a', 'b', 'b']),
                                              'b': ['x', 'y', 'z']},
                                      dims=['a', 'b'])
        clipped = assembly.clip(-5, +5)
        assert type(clipped) == type(assembly)