MCPcopy Index your code
hub / github.com/pydata/xarray / test_getitem_fancy

Method test_getitem_fancy

xarray/tests/test_variable.py:780–865  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

778 assert_array_equal(v_new, v_data[0])
779
780 def test_getitem_fancy(self):
781 v = self.cls(["x", "y"], [[0, 1, 2], [3, 4, 5]])
782 v_data = v.compute().data
783
784 ind = Variable(["a", "b"], [[0, 1, 1], [1, 1, 0]])
785 v_new = v[ind]
786 assert v_new.dims == ("a", "b", "y")
787 assert_array_equal(v_new, v_data[[[0, 1, 1], [1, 1, 0]], :])
788
789 # It would be ok if indexed with the multi-dimensional array including
790 # the same name
791 ind = Variable(["x", "b"], [[0, 1, 1], [1, 1, 0]])
792 v_new = v[ind]
793 assert v_new.dims == ("x", "b", "y")
794 assert_array_equal(v_new, v_data[[[0, 1, 1], [1, 1, 0]], :])
795
796 ind = Variable(["a", "b"], [[0, 1, 2], [2, 1, 0]])
797 v_new = v[dict(y=ind)]
798 assert v_new.dims == ("x", "a", "b")
799 assert_array_equal(v_new, v_data[:, ([0, 1, 2], [2, 1, 0])])
800
801 ind = Variable(["a", "b"], [[0, 0], [1, 1]])
802 v_new = v[dict(x=[1, 0], y=ind)]
803 assert v_new.dims == ("x", "a", "b")
804 assert_array_equal(v_new, v_data[[1, 0]][:, ind])
805
806 # along diagonal
807 ind = Variable(["a"], [0, 1])
808 v_new = v[ind, ind]
809 assert v_new.dims == ("a",)
810 assert_array_equal(v_new, v_data[[0, 1], [0, 1]])
811
812 # with integer
813 ind = Variable(["a", "b"], [[0, 0], [1, 1]])
814 v_new = v[dict(x=0, y=ind)]
815 assert v_new.dims == ("a", "b")
816 assert_array_equal(v_new[0], v_data[0][[0, 0]])
817 assert_array_equal(v_new[1], v_data[0][[1, 1]])
818
819 # with slice
820 ind = Variable(["a", "b"], [[0, 0], [1, 1]])
821 v_new = v[dict(x=slice(None), y=ind)]
822 assert v_new.dims == ("x", "a", "b")
823 assert_array_equal(v_new, v_data[:, [[0, 0], [1, 1]]])
824
825 ind = Variable(["a", "b"], [[0, 0], [1, 1]])
826 v_new = v[dict(x=ind, y=slice(None))]
827 assert v_new.dims == ("a", "b", "y")
828 assert_array_equal(v_new, v_data[[[0, 0], [1, 1]], :])
829
830 ind = Variable(["a", "b"], [[0, 0], [1, 1]])
831 v_new = v[dict(x=ind, y=slice(None, 1))]
832 assert v_new.dims == ("a", "b", "y")
833 assert_array_equal(v_new, v_data[[[0, 0], [1, 1]], slice(None, 1)])
834
835 # slice matches explicit dimension
836 ind = Variable(["y"], [0, 1])
837 v_new = v[ind, :2]

Callers

nothing calls this directly

Calls 5

computeMethod · 0.95
VariableClass · 0.90
assert_identicalFunction · 0.90
arangeMethod · 0.80
clsMethod · 0.45

Tested by

no test coverage detected