jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_promote(self, backend):
        """
        define x = Variable((1,)) with
        [x1,]

        x is represented as eye(1) in the A matrix, i.e.,

         x1
        [[1]]

        promote(x) means we repeat the row to match the required dimensionality of n rows.

        Thus, when using the same columns as before and assuming n = 3, we now have

         x1
        [[1],
         [1],
         [1]]
        """

        variable_lin_op = linOpHelper((1,), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 1)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(1, 1)).toarray()
        assert n.all(view_A == np.eye(1))

        promote_lin_op = linOpHelper((3,))
        out_view = backend.promote(promote_lin_op, view)
        A = out_view.get_tensor_representation(0, 1)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(3, 1)).toarray()
        expected = np.array([[1], [1], [1]])
        assert n.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 1) == view.get_tensor_representation(0, 1)
//  [python]
def test_promote(self, backend):
        """
        define x = Variable((1,)) with
        [x1,]

        x is represented as eye(1) in the A matrix, i.e.,

         x1
        [[1]]

        promote(x) means we repeat the row to match the required dimensionality of n rows.

        Thus, when using the same columns as before and assuming n = 3, we now have

         x1
        [[1],
         [1],
         [1]]
        """

        variable_lin_op = linOpHelper((1,), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 1)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(1, 1)).toarray()
        assert np.all(view_A == np.eye(1))

        promote_lin_op = linOpHelper((3,))
        out_view = backend.promote(promote_lin_op, view)
        A = out_view.get_tensor_representation(0, 1)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(3, 1)).toarray()
        expected = np.array([[1], [1], [1]])
        assert np.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 1) == view.get_tensor_representation(0, 1)
//  [python]
def test_diag_vec_with_offset(self, backend):
        """
        define x = Variable((2,)) with
        [x1, x2]

        x is represented as eye(2) in the A matrix, i.e.,

         x1  x2
        [[1  0],
         [0  1]]

        diag_vec(x, k) means we introduce zero rows as if the vector was the k-diagonal
        of an n+|k| x n+|k| matrix, with n the length of x.

        Thus, for k=1 and using the same columns as before, want to represent
        [[0  x1 0],
        [ 0  0  x2],
        [[0  0  0]]
        i.e., unrolled in column-major order:

         x1  x2
        [[0  0],
        [0  0],
        [0  0],
        [1  0],
        [0  0],
        [0  0],
        [0  0],
        [0  1],
        [0  0]]
        """

        variable_lin_op = linOpHelper((2,), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 2)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(2, 2)).toarray()
        assert np.all(view_A == np.eye(2))

        k = 1
        diag_vec_lin_op = linOpHelper(shape=(3, 3), data=k)
        out_view = backend.diag_vec(diag_vec_lin_op, view)
        A = out_view.get_tensor_representation(0, 9)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(9, 2)).toarray()
        expected = np.array(
            [[0, 0], [0, 0], [0, 0], [1, 0], [0, 0], [0, 0], [0, 0], [0, 1], [0, 0]]
        )
        assert np.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 9) == view.get_tensor_representation(0, 9)
//  [python]
def test_vstack(self, backend):
        """
        define x,y = Variable((1,2)), Variable((1,2)) with
        [[x1, x2]]
        and
        [[y1, y2]]

        vstack([x, y]) yields

        [[x1, x2],
         [y1, y2]]

        which maps to

         x1   x2  y1  y2
        [[1   0   0   0],
         [0   0   1   0],
         [0   1   0   0],
         [0   0   0   1]]
        """

        lin_op_x = linOpHelper((1, 2), type="variable", data=1)
        lin_op_y = linOpHelper((1, 2), type="variable", data=2)

        vstack_lin_op = linOpHelper(args=[lin_op_x, lin_op_y])
        backend.id_to_col = {1: 0, 2: 2}
        out_view = backend.vstack(vstack_lin_op, backend.get_empty_view())
        A = out_view.get_tensor_representation(0, 4)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(4, 4)).toarray()
        expected = np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
        assert np.all(A == expected)
//  [python]
def test_mul(self, backend):
        """
        define x = Variable((2,2)) with
        [[x11, x12],
         [x21, x22]]

         Multiplying with the constant from the left
        [[1, 2],
         [3, 4]],

         we expect the output to be
        [[  x11 + 2 x21,   x12 + 2 x22],
         [3 x11 + 4 x21, 3 x12 + 4 x22]]

        i.e., when represented in the A matrix (again using column-major order):
         x11 x21 x12 x22
        [[1   2   0   0],
         [3   4   0   0],
         [0   0   1   2],
         [0   0   3   4]]
        """

        variable_lin_op = linOpHelper((2, 2), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 4)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(4, 4)).toarray()
        assert n.all(view_A == np.eye(4))

        lhs = linOpHelper((2, 2), type="dense_const", data=np.array([[1, 2], [3, 4]]))

        mul_lin_op = linOpHelper(data=lhs, args=[variable_lin_op])
        out_view = backend.mul(mul_lin_op, view)
        A = out_view.get_tensor_representation(0, 4)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(4, 4)).toarray()
        expected = np.array([[1, 2, 0, 0], [3, 4, 0, 0], [0, 0, 1, 2], [0, 0, 3, 4]])
        assert n.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 4) == view.get_tensor_representation(0, 4)
//  [python]
def test_mul(self, backend):
        """
        define x = Variable((2,2)) with
        [[x11, x12],
         [x21, x22]]

         Multiplying with the constant from the left
        [[1, 2],
         [3, 4]],

         we expect the output to be
        [[  x11 + 2 x21,   x12 + 2 x22],
         [3 x11 + 4 x21, 3 x12 + 4 x22]]

        i.e., when represented in the A matrix (again using column-major order):
         x11 x21 x12 x22
        [[1   2   0   0],
         [3   4   0   0],
         [0   0   1   2],
         [0   0   3   4]]
        """

        variable_lin_op = linOpHelper((2, 2), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 4)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(4, 4)).toarray()
        assert np.all(view_A == np.eye(4))

        lhs = linOpHelper((2, 2), type="dense_const", data=np.array([[1, 2], [3, 4]]))

        mul_lin_op = linOpHelper(data=lhs, args=[variable_lin_op])
        out_view = backend.mul(mul_lin_op, view)
        A = out_view.get_tensor_representation(0, 4)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(4, 4)).toarray()
        expected = np.array([[1, 2, 0, 0], [3, 4, 0, 0], [0, 0, 1, 2], [0, 0, 3, 4]])
        assert np.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 4) == view.get_tensor_representation(0, 4)
//  [python]
def test_mul_elementwise(self, backend):
        """
        define x = Variable((2,)) with
        [x1, x2]

        x is represented as eye(2) in the A matrix, i.e.,

         x1  x2
        [[1  0],
         [0  1]]

         mul_elementwise(x, a) means 'a' is reshaped into a column vector and multiplied by A.
         E.g. for a = (2,3), we obtain

         x1  x2
        [[2  0],
         [0  3]]
        """

        variable_lin_op = linOpHelper((2,), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 2)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(2, 2)).toarray()
        assert n.all(view_A == np.eye(2))

        lhs = linOpHelper((2,), type="dense_const", data=np.array([2, 3]))

        mul_elementwise_lin_op = linOpHelper(data=lhs)
        out_view = backend.mul_elem(mul_elementwise_lin_op, view)
        A = out_view.get_tensor_representation(0, 2)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(2, 2)).toarray()
        expected = np.array([[2, 0], [0, 3]])
        assert n.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 2) == view.get_tensor_representation(0, 2)
//  [python]
def test_mul_elementwise(self, backend):
        """
        define x = Variable((2,)) with
        [x1, x2]

        x is represented as eye(2) in the A matrix, i.e.,

         x1  x2
        [[1  0],
         [0  1]]

         mul_elementwise(x, a) means 'a' is reshaped into a column vector and multiplied by A.
         E.g. for a = (2,3), we obtain

         x1  x2
        [[2  0],
         [0  3]]
        """

        variable_lin_op = linOpHelper((2,), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 2)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(2, 2)).toarray()
        assert np.all(view_A == np.eye(2))

        lhs = linOpHelper((2,), type="dense_const", data=np.array([2, 3]))

        mul_elementwise_lin_op = linOpHelper(data=lhs)
        out_view = backend.mul_elem(mul_elementwise_lin_op, view)
        A = out_view.get_tensor_representation(0, 2)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(2, 2)).toarray()
        expected = np.array([[2, 0], [0, 3]])
        assert np.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 2) == view.get_tensor_representation(0, 2)
//  [python]
def test_rmul(self, backend):
        """
        define x = Variable((2,2)) with
        [[x11, x12],
         [x21, x22]]

         Multiplying with the constant from the right
         (intentionally using 1D vector to cover edge case)
        [1, 2]

         we expect the output to be
         [[x11 + 2 x12],
          [x21 + 2 x22]]

        i.e., when represented in the A matrix:
         x11 x21 x12 x22
        [[1   0   2   0],
         [0   1   0   2]]
        """

        variable_lin_op = linOpHelper((2, 2), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 4)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(4, 4)).toarray()
        assert np.all(view_A == np.eye(4))

        rhs = linOpHelper((2,), type="dense_const", data=np.array([1, 2]))

        rmul_lin_op = linOpHelper(data=rhs, args=[variable_lin_op])
        out_view = backend.rmul(rmul_lin_op, view)
        A = out_view.get_tensor_representation(0, 2)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(2, 4)).toarray()
        expected = np.array([[1, 0, 2, 0], [0, 1, 0, 2]])
        assert np.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 2) == view.get_tensor_representation(0, 2)
//  [python]
def test_div(self, backend):
        """
        define x = Variable((2,2)) with
        [[x11, x12],
         [x21, x22]]

         Dividing elementwise with
        [[1, 2],
         [3, 4]],

        we obtain:
         x11 x21 x12 x22
        [[1   0   0   0],
         [0   1/3 0   0],
         [0   0   1/2 0],
         [0   0   0   1/4]]
        """

        variable_lin_op = linOpHelper((2, 2), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 4)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(4, 4)).toarray()
        assert n.all(view_A == np.eye(4))

        lhs = linOpHelper((2, 2), type="dense_const", data=np.array([[1, 2], [3, 4]]))

        div_lin_op = linOpHelper(data=lhs)
        out_view = backend.div(div_lin_op, view)
        A = out_view.get_tensor_representation(0, 4)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(4, 4)).toarray()
        expected = np.array([[1, 0, 0, 0], [0, 1 / 3, 0, 0], [0, 0, 1 / 2, 0], [0, 0, 0, 1 / 4]])
        assert n.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 4) == view.get_tensor_representation(0, 4)
//  [python]
def test_div(self, backend):
        """
        define x = Variable((2,2)) with
        [[x11, x12],
         [x21, x22]]

         Dividing elementwise with
        [[1, 2],
         [3, 4]],

        we obtain:
         x11 x21 x12 x22
        [[1   0   0   0],
         [0   1/3 0   0],
         [0   0   1/2 0],
         [0   0   0   1/4]]
        """

        variable_lin_op = linOpHelper((2, 2), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 4)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(4, 4)).toarray()
        assert np.all(view_A == np.eye(4))

        lhs = linOpHelper((2, 2), type="dense_const", data=np.array([[1, 2], [3, 4]]))

        div_lin_op = linOpHelper(data=lhs)
        out_view = backend.div(div_lin_op, view)
        A = out_view.get_tensor_representation(0, 4)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(4, 4)).toarray()
        expected = np.array([[1, 0, 0, 0], [0, 1 / 3, 0, 0], [0, 0, 1 / 2, 0], [0, 0, 0, 1 / 4]])
        assert np.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 4) == view.get_tensor_representation(0, 4)
//  [python]
def test_trace(self, backend):
        """
        define x = Variable((2,2)) with
        [[x11, x12],
         [x21, x22]]

        x is represented as eye(4) in the A matrix (in column-major order), i.e.,

         x11 x21 x12 x22
        [[1   0   0   0],
         [0   1   0   0],
         [0   0   1   0],
         [0   0   0   1]]

        trace(x) means we sum the diagonal entries of x, i.e.

         x11 x21 x12 x22
        [[1   0   0   1]]
        """

        variable_lin_op = linOpHelper((2, 2), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 4)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(4, 4)).toarray()
        assert np.all(view_A == np.eye(4))

        trace_lin_op = linOpHelper(args=[variable_lin_op])
        out_view = backend.trace(trace_lin_op, view)
        A = out_view.get_tensor_representation(0, 1)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(1, 4)).toarray()
        expected = np.array([[1, 0, 0, 1]])
        assert np.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 1) == view.get_tensor_representation(0, 1)
//  [python]
def test_conv(self, backend):
        """
        define x = Variable((3,)) with
        [x1, x2, x3]

        having f = [1,2,3], conv(f, x) means we repeat the column vector of f for each column in
        the A matrix, shifting it down by one after each repetition, i.e.,
          x1 x2 x3
        [[1  0  0],
         [2  1  0],
         [3  2  1],
         [0  3  2],
         [0  0  3]]
        """

        variable_lin_op = linOpHelper((3,), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 3)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(3, 3)).toarray()
        assert n.all(view_A == np.eye(3))

        f = linOpHelper((3,), type="dense_const", data=np.array([1, 2, 3]))
        conv_lin_op = linOpHelper(data=f, shape=(5, 1), args=[variable_lin_op])

        out_view = backend.conv(conv_lin_op, view)
        A = out_view.get_tensor_representation(0, 5)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(5, 3)).toarray()
        expected = np.array(
            [[1.0, 0.0, 0.0], [2.0, 1.0, 0.0], [3.0, 2.0, 1.0], [0.0, 3.0, 2.0], [0.0, 0.0, 3.0]]
        )
        assert n.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 5) == view.get_tensor_representation(0, 5)
//  [python]
def test_conv(self, backend):
        """
        define x = Variable((3,)) with
        [x1, x2, x3]

        having f = [1,2,3], conv(f, x) means we repeat the column vector of f for each column in
        the A matrix, shifting it down by one after each repetition, i.e.,
          x1 x2 x3
        [[1  0  0],
         [2  1  0],
         [3  2  1],
         [0  3  2],
         [0  0  3]]
        """

        variable_lin_op = linOpHelper((3,), type="variable", data=1)
        view = backend.process_constraint(variable_lin_op, backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 3)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(3, 3)).toarray()
        assert np.all(view_A == np.eye(3))

        f = linOpHelper((3,), type="dense_const", data=np.array([1, 2, 3]))
        conv_lin_op = linOpHelper(data=f, shape=(5, 1), args=[variable_lin_op])

        out_view = backend.conv(conv_lin_op, view)
        A = out_view.get_tensor_representation(0, 5)

        # cast to numpy
        A = sp.coo_matrix((A.data, (A.row, A.col)), shape=(5, 3)).toarray()
        expected = np.array(
            [[1.0, 0.0, 0.0], [2.0, 1.0, 0.0], [3.0, 2.0, 1.0], [0.0, 3.0, 2.0], [0.0, 0.0, 3.0]]
        )
        assert np.all(A == expected)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 5) == view.get_tensor_representation(0, 5)
//  [python]
def test_parametrized_sum_entries(self, param_backend):
        """
        starting with a parametrized expression
        x1  x2
        [[[1  0],
         [0  0]],

         [[0  0],
         [0  1]]]

        sum_entries(x) means we consider the entries in all rows, i.e., we sum along the row axis.

        Thus, when using the same columns as before, we now have

         x1  x2
        [[[1  0]],

         [[0  1]]]
        """
        param_lin_op = linOpHelper((2,), type="param", data=2)
        param_backend.param_to_col = {2: 0, -1: 3}
        variable_lin_op = linOpHelper((2,), type="variable", data=1)
        var_view = param_backend.process_constraint(variable_lin_op, param_backend.get_empty_view())
        mul_elem_lin_op = linOpHelper(data=param_lin_op)
        param_var_view = param_backend.mul_elem(mul_elem_lin_op, var_view)

        sum_entries_lin_op = linOpHelper()
        out_view = param_backend.sum_entries(sum_entries_lin_op, param_var_view)
        out_repr = out_view.get_tensor_representation(0, 1)

        slice_idx_zero = out_repr.get_param_slice(0).toarray()[:, :-1]
        expected_idx_zero = np.array([[1.0, 0.0]])
        assert np.all(slice_idx_zero == expected_idx_zero)

        slice_idx_one = out_repr.get_param_slice(1).toarray()[:, :-1]
        expected_idx_one = np.array([[0.0, 1.0]])
        assert np.all(slice_idx_one == expected_idx_one)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 1) == param_var_view.get_tensor_representation(
            0, 1
        )
//  [python]
def test_get_variable_tensor(self, numpy_backend):
        outer = numpy_backend.get_variable_tensor((2,), 1)
        assert outer.keys() == {1}, "Should only be in variable with ID 1"
        inner = outer[1]
        assert inner.keys() == {-1}, "Should only be in parameter slice -1, i.e. non parametrized."
        tensor = inner[-1]
        assert isinstance(tensor, np.ndarray), "Should be a numpy array"
        assert tensor.shape == (1, 2, 2), "Should be a 1x2x2 tensor"
        assert np.all(tensor[0] == np.eye(2)), "Should be eye(2)"
//  [python]
def test_get_data_tensor(self, numpy_backend, data):
        outer = numpy_backend.get_data_tensor(data)
        assert outer.keys() == {-1}, "Should only be constant variable ID."
        inner = outer[-1]
        assert inner.keys() == {-1}, "Should only be in parameter slice -1, i.e. non parametrized."
        tensor = inner[-1]
        assert isinstance(tensor, np.ndarray), "Should be a numpy array"
        assert isinstance(tensor[0], np.ndarray), "Inner matrix should also be a numpy array"
        assert tensor.shape == (1, 4, 1), "Should be a 1x4x1 tensor"
        expected = numpy_backend._to_dense(data).reshape((-1, 1), order="F")
        assert np.alltrue(tensor[0] == expected)
//  [python]
def test_get_data_tensor(self, numpy_backend, data):
        outer = numpy_backend.get_data_tensor(data)
        assert outer.keys() == {-1}, "Should only be constant variable ID."
        inner = outer[-1]
        assert inner.keys() == {-1}, "Should only be in parameter slice -1, i.e. non parametrized."
        tensor = inner[-1]
        assert isinstance(tensor, np.ndarray), "Should be a numpy array"
        assert isinstance(tensor[0], np.ndarray), "Inner matrix should also be a numpy array"
        assert tensor.shape == (1, 4, 1), "Should be a 1x4x1 tensor"
        expected = numpy_backend._to_dense(data).reshape((-1, 1), order="F")
        assert np.all(tensor[0] == expected)
//  [python]
def test_get_kron_row_indices(self, backend):
        """
        kron(l,r)
        with
        l = [[x1, x3],  r = [[a],
             [x2, x4]]       [b]]

        yields
        [[ax1, ax3],
         [bx1, bx3],
         [ax2, ax4],
         [bx2, bx4]]

        Which is what we get when we compute kron(l,r) directly,
        as l is represented as eye(4) and r is reshaped into a column vector.

        So we have:
        kron(l,r) =
        [[a, 0, 0, 0],
         [b, 0, 0, 0],
         [0, a, 0, 0],
         [0, b, 0, 0],
         [0, 0, a, 0],
         [0, 0, b, 0],
         [0, 0, 0, a],
         [0, 0, 0, b]].

        Thus, this function should return arange(8).
        """
        indices = backend._get_kron_row_indices((2, 2), (2, 1))
        assert n.all(indices == n.arange(8))

        """
        kron(l,r)
        with 
        l = [[x1],  r = [[a, c],
             [x2]]       [b, d]]

        yields
        [[ax1, cx1],
         [bx1, dx1],
         [ax2, cx2],
         [bx2, dx2]]
        
        Here, we have to swap the row indices of the resulting matrix.
        Immediately applying kron(l,r) gives to eye(2) and r reshaped to 
        a column vector gives.
                 
        So we have:
        kron(l,r) = 
        [[a, 0],
         [b, 0],
         [c, 0],
         [d, 0],
         [0, a],
         [0, b]
         [0, c],
         [0, d]].

        Thus, we need to to return [0, 1, 4, 5, 2, 3, 6, 7].
        """

        indices = backend._get_kron_row_indices((2, 1), (2, 2))
        assert n.all(indices == [0, 1, 4, 5, 2, 3, 6, 7])

        indices = backend._get_kron_row_indices((1, 2), (3, 2))
        assert n.all(indices == n.arange(12))

        indices = backend._get_kron_row_indices((3, 2), (1, 2))
        assert n.all(indices == [0, 2, 4, 1, 3, 5, 6, 8, 10, 7, 9, 11])

        indices = backend._get_kron_row_indices((2, 2), (2, 2))
        expected = [0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15]
        assert n.all(indices == expected)
//  [python]
def test_get_kron_row_indices(self, backend):
        """
        kron(l,r)
        with
        l = [[x1, x3],  r = [[a],
             [x2, x4]]       [b]]

        yields
        [[ax1, ax3],
         [bx1, bx3],
         [ax2, ax4],
         [bx2, bx4]]

        Which is what we get when we compute kron(l,r) directly,
        as l is represented as eye(4) and r is reshaped into a column vector.

        So we have:
        kron(l,r) =
        [[a, 0, 0, 0],
         [b, 0, 0, 0],
         [0, a, 0, 0],
         [0, b, 0, 0],
         [0, 0, a, 0],
         [0, 0, b, 0],
         [0, 0, 0, a],
         [0, 0, 0, b]].

        Thus, this function should return arange(8).
        """
        indices = backend._get_kron_row_indices((2, 2), (2, 1))
        assert np.all(indices == np.arange(8))

        """
        kron(l,r)
        with 
        l = [[x1],  r = [[a, c],
             [x2]]       [b, d]]

        yields
        [[ax1, cx1],
         [bx1, dx1],
         [ax2, cx2],
         [bx2, dx2]]
        
        Here, we have to swap the row indices of the resulting matrix.
        Immediately applying kron(l,r) gives to eye(2) and r reshaped to 
        a column vector gives.
                 
        So we have:
        kron(l,r) = 
        [[a, 0],
         [b, 0],
         [c, 0],
         [d, 0],
         [0, a],
         [0, b]
         [0, c],
         [0, d]].

        Thus, we need to to return [0, 1, 4, 5, 2, 3, 6, 7].
        """

        indices = backend._get_kron_row_indices((2, 1), (2, 2))
        assert np.all(indices == [0, 1, 4, 5, 2, 3, 6, 7])

        indices = backend._get_kron_row_indices((1, 2), (3, 2))
        assert np.all(indices == np.arange(12))

        indices = backend._get_kron_row_indices((3, 2), (1, 2))
        assert np.all(indices == [0, 2, 4, 1, 3, 5, 6, 8, 10, 7, 9, 11])

        indices = backend._get_kron_row_indices((2, 2), (2, 2))
        expected = [0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15]
        assert np.all(indices == expected)
//  [python]
def test_get_variable_tensor(self, scipy_backend):
        outer = scipy_backend.get_variable_tensor((2,), 1)
        assert outer.keys() == {1}, "Should only be in variable with ID 1"
        inner = outer[1]
        assert inner.keys() == {-1}, "Should only be in parameter slice -1, i.e. non parametrized."
        tensor = inner[-1]
        assert isinstance(tensor, sp.spmatrix), "Should be a scipy sparse matrix"
        assert tensor.shape == (2, 2), "Should be a 1*2x2 tensor"
        assert np.all(tensor == np.eye(2)), "Should be eye(2)"
//  [python]
def test_parametrized_rmul(self, param_backend):
        """
        Continuing from the non-parametrized example when the rhs is a parameter,
        instead of multiplying with known values, the matrix is split up into two slices,
        each representing an element of the parameter, i.e. instead of
         x11 x21 x12 x22
        [[1   0   2   0],
         [0   1   0   2]]

         we obtain the list of length two, where we have ones at the entries where previously
         we had the 1 and 2:

         x11  x21  x12  x22
        [
         [[1   0   0   0],
          [0   1   0   0]]

         [[0   0   1   0],
          [0   0   0   1]]
        ]
        """

        variable_lin_op = linOpHelper((2, 2), type="variable", data=1)
        param_backend.var_length = 4
        view = param_backend.process_constraint(variable_lin_op, param_backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 4)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(4, 4)).toarray()
        assert n.all(view_A == n.eye(4))

        rhs_parameter = linOpHelper((2,), type="param", data=2)

        rmul_lin_op = linOpHelper(data=rhs_parameter, args=[variable_lin_op])
        out_view = param_backend.rmul(rmul_lin_op, view)
        out_repr = out_view.get_tensor_representation(0, 2)

        # indices are: variable 1, parameter 2, 0 index of the list
        slice_idx_zero = out_repr.get_param_slice(0).toarray()[:, :-1]
        expected_idx_zero = n.array([[1, 0, 0, 0], [0, 1, 0, 0]])
        assert n.all(slice_idx_zero == expected_idx_zero)

        # indices are: variable 1, parameter 2, 1 index of the list
        slice_idx_one = out_repr.get_param_slice(1).toarray()[:, :-1]
        expected_idx_one = n.array([[0, 0, 1, 0], [0, 0, 0, 1]])
        assert n.all(slice_idx_one == expected_idx_one)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 2) == view.get_tensor_representation(0, 2)
//  [python]
def test_parametrized_rmul(self, param_backend):
        """
        Continuing from the non-parametrized example when the rhs is a parameter,
        instead of multiplying with known values, the matrix is split up into two slices,
        each representing an element of the parameter, i.e. instead of
         x11 x21 x12 x22
        [[1   0   2   0],
         [0   1   0   2]]

         we obtain the list of length two, where we have ones at the entries where previously
         we had the 1 and 2:

         x11  x21  x12  x22
        [
         [[1   0   0   0],
          [0   1   0   0]]

         [[0   0   1   0],
          [0   0   0   1]]
        ]
        """

        variable_lin_op = linOpHelper((2, 2), type="variable", data=1)
        param_backend.var_length = 4
        view = param_backend.process_constraint(variable_lin_op, param_backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 4)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(4, 4)).toarray()
        assert np.all(view_A == np.eye(4))

        rhs_parameter = linOpHelper((2,), type="param", data=2)

        rmul_lin_op = linOpHelper(data=rhs_parameter, args=[variable_lin_op])
        out_view = param_backend.rmul(rmul_lin_op, view)
        out_repr = out_view.get_tensor_representation(0, 2)

        # indices are: variable 1, parameter 2, 0 index of the list
        slice_idx_zero = out_repr.get_param_slice(0).toarray()[:, :-1]
        expected_idx_zero = np.array([[1, 0, 0, 0], [0, 1, 0, 0]])
        assert np.all(slice_idx_zero == expected_idx_zero)

        # indices are: variable 1, parameter 2, 1 index of the list
        slice_idx_one = out_repr.get_param_slice(1).toarray()[:, :-1]
        expected_idx_one = np.array([[0, 0, 1, 0], [0, 0, 0, 1]])
        assert np.all(slice_idx_one == expected_idx_one)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 2) == view.get_tensor_representation(0, 2)
//  [python]
def test_mul_elementwise_parametrized(self, param_backend):
        """
        Continuing the non-parametrized example when 'a' is a parameter, instead of multiplying
        with known values, the matrix is split up into two slices, each representing an element
        of the parameter, i.e. instead of
         x1  x2
        [[2  0],
         [0  3]]

         we obtain the list of length two:

          x1  x2
        [
         [[1  0],
          [0  0]],

         [[0  0],
          [0  1]]
        ]
        """

        variable_lin_op = linOpHelper((2,), type="variable", data=1)
        view = param_backend.process_constraint(variable_lin_op, param_backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 2)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(2, 2)).toarray()
        assert n.all(view_A == np.eye(2))

        lhs_parameter = linOpHelper((2,), type="param", data=2)

        mul_elementwise_lin_op = linOpHelper(data=lhs_parameter)
        out_view = param_backend.mul_elem(mul_elementwise_lin_op, view)
        out_repr = out_view.get_tensor_representation(0, 2)

        # indices are: variable 1, parameter 2, 0 index of the list
        slice_idx_zero = out_repr.get_param_slice(0).toarray()[:, :-1]
        expected_idx_zero = np.array([[1, 0], [0, 0]])
        assert n.all(slice_idx_zero == expected_idx_zero)

        # indices are: variable 1, parameter 2, 1 index of the list
        slice_idx_one = out_repr.get_param_slice(1).toarray()[:, :-1]
        expected_idx_one = np.array([[0, 0], [0, 1]])
        assert n.all(slice_idx_one == expected_idx_one)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 2) == view.get_tensor_representation(0, 2)
//  [python]
def test_mul_elementwise_parametrized(self, param_backend):
        """
        Continuing the non-parametrized example when 'a' is a parameter, instead of multiplying
        with known values, the matrix is split up into two slices, each representing an element
        of the parameter, i.e. instead of
         x1  x2
        [[2  0],
         [0  3]]

         we obtain the list of length two:

          x1  x2
        [
         [[1  0],
          [0  0]],

         [[0  0],
          [0  1]]
        ]
        """

        variable_lin_op = linOpHelper((2,), type="variable", data=1)
        view = param_backend.process_constraint(variable_lin_op, param_backend.get_empty_view())

        # cast to numpy
        view_A = view.get_tensor_representation(0, 2)
        view_A = sp.coo_matrix((view_A.data, (view_A.row, view_A.col)), shape=(2, 2)).toarray()
        assert np.all(view_A == np.eye(2))

        lhs_parameter = linOpHelper((2,), type="param", data=2)

        mul_elementwise_lin_op = linOpHelper(data=lhs_parameter)
        out_view = param_backend.mul_elem(mul_elementwise_lin_op, view)
        out_repr = out_view.get_tensor_representation(0, 2)

        # indices are: variable 1, parameter 2, 0 index of the list
        slice_idx_zero = out_repr.get_param_slice(0).toarray()[:, :-1]
        expected_idx_zero = np.array([[1, 0], [0, 0]])
        assert np.all(slice_idx_zero == expected_idx_zero)

        # indices are: variable 1, parameter 2, 1 index of the list
        slice_idx_one = out_repr.get_param_slice(1).toarray()[:, :-1]
        expected_idx_one = np.array([[0, 0], [0, 1]])
        assert np.all(slice_idx_one == expected_idx_one)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 2) == view.get_tensor_representation(0, 2)
//  [python]
def test_parametrized_trace(self, param_backend):
        """
        Continuing from the non-parametrized example, instead of a pure variable
        input, we take a variable that has been multiplied elementwise by a parameter.

        The trace of this expression is then given by

            x11  x21  x12  x22
        [
            [[1   0   0   0]],

            [[0   0   0   0]],

            [[0   0   0   0]],

            [[0   0   0   1]]
        ]
        """
        param_lin_op = linOpHelper((2, 2), type="param", data=2)
        param_backend.param_to_col = {2: 0, -1: 4}
        param_backend.param_to_size = {-1: 1, 2: 4}
        param_backend.var_length = 4
        variable_lin_op = linOpHelper((2, 2), type="variable", data=1)
        var_view = param_backend.process_constraint(variable_lin_op, param_backend.get_empty_view())
        mul_elem_lin_op = linOpHelper(data=param_lin_op)
        param_var_view = param_backend.mul_elem(mul_elem_lin_op, var_view)

        trace_lin_op = linOpHelper(args=[variable_lin_op])
        out_view = param_backend.trace(trace_lin_op, param_var_view)
        out_repr = out_view.get_tensor_representation(0, 1)

        slice_idx_zero = out_repr.get_param_slice(0).toarray()[:, :-1]
        expected_idx_zero = np.array([[1.0, 0.0, 0.0, 0.0]])
        assert n.all(slice_idx_zero == expected_idx_zero)

        slice_idx_one = out_repr.get_param_slice(1).toarray()[:, :-1]
        expected_idx_one = np.array([[0.0, 0.0, 0.0, 0.0]])
        assert n.all(slice_idx_one == expected_idx_one)

        slice_idx_two = out_repr.get_param_slice(2).toarray()[:, :-1]
        expected_idx_two = np.array([[0.0, 0.0, 0.0, 0.0]])
        assert n.all(slice_idx_two == expected_idx_two)

        slice_idx_three = out_repr.get_param_slice(3).toarray()[:, :-1]
        expected_idx_three = np.array([[0.0, 0.0, 0.0, 1.0]])
        assert n.all(slice_idx_three == expected_idx_three)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 1) == param_var_view.get_tensor_representation(
            0, 1
        )
//  [python]
def test_parametrized_trace(self, param_backend):
        """
        Continuing from the non-parametrized example, instead of a pure variable
        input, we take a variable that has been multiplied elementwise by a parameter.

        The trace of this expression is then given by

            x11  x21  x12  x22
        [
            [[1   0   0   0]],

            [[0   0   0   0]],

            [[0   0   0   0]],

            [[0   0   0   1]]
        ]
        """
        param_lin_op = linOpHelper((2, 2), type="param", data=2)
        param_backend.param_to_col = {2: 0, -1: 4}
        param_backend.param_to_size = {-1: 1, 2: 4}
        param_backend.var_length = 4
        variable_lin_op = linOpHelper((2, 2), type="variable", data=1)
        var_view = param_backend.process_constraint(variable_lin_op, param_backend.get_empty_view())
        mul_elem_lin_op = linOpHelper(data=param_lin_op)
        param_var_view = param_backend.mul_elem(mul_elem_lin_op, var_view)

        trace_lin_op = linOpHelper(args=[variable_lin_op])
        out_view = param_backend.trace(trace_lin_op, param_var_view)
        out_repr = out_view.get_tensor_representation(0, 1)

        slice_idx_zero = out_repr.get_param_slice(0).toarray()[:, :-1]
        expected_idx_zero = np.array([[1.0, 0.0, 0.0, 0.0]])
        assert np.all(slice_idx_zero == expected_idx_zero)

        slice_idx_one = out_repr.get_param_slice(1).toarray()[:, :-1]
        expected_idx_one = np.array([[0.0, 0.0, 0.0, 0.0]])
        assert np.all(slice_idx_one == expected_idx_one)

        slice_idx_two = out_repr.get_param_slice(2).toarray()[:, :-1]
        expected_idx_two = np.array([[0.0, 0.0, 0.0, 0.0]])
        assert np.all(slice_idx_two == expected_idx_two)

        slice_idx_three = out_repr.get_param_slice(3).toarray()[:, :-1]
        expected_idx_three = np.array([[0.0, 0.0, 0.0, 1.0]])
        assert np.all(slice_idx_three == expected_idx_three)

        # Note: view is edited in-place:
        assert out_view.get_tensor_representation(0, 1) == param_var_view.get_tensor_representation(
            0, 1
        )
//  [python]
def test_basic_save(self):
        # test saving to a Dataset, in place of a filepath spec.
        # NOTE that this requires 'compute=False', as delayed saves can only operate on
        # a closed file.

        # Save to netcdf file in the usual way.
        filepath_direct = f"{self.temp_dir}/tmp_direct.nc"
        iris.save(self.testdata, filepath_direct)
        # Check against test-specific CDL result file.
        self.assertCDL(filepath_direct)

        # Save same data indirectly via a netcdf dataset.
        filepath_indirect = f"{self.temp_dir}/tmp_indirect.nc"
        nc_dataset = nc.Dataset(filepath_indirect, "w")
        # NOTE: we **must** use delayed saving here, as we cannot do direct saving to
        # a user-owned dataset.
        result = iris.save(self.testdata, nc_dataset, saver="nc", compute=False)

        # Do some very basic sanity checks on the resulting Dataset.
        # It should still be open (!)
        self.assertTrue(nc_dataset.isopen())
        self.assertEqual(
            ["time", "levelist", "latitude", "longitude"],
            list(nc_dataset.dimensions),
        )
        self.assertEqual(
            ["co2", "time", "levelist", "latitude", "longitude", "lnsp"],
            list(nc_dataset.variables),
        )
        nc_dataset.close()

        # Check the saved file against the same CDL as the 'normal' save.
        self.assertCDL(filepath_indirect)

        # Confirm that cube content is however not yet written.
        ds = nc.Dataset(filepath_indirect)
        for cube in self.testdata:
            assert n.all(ds.variables[cube.var_name][:].mask)
        ds.close()

        # Complete the delayed saves.
        result.compute()

        # Check that data now *is* written.
        ds = nc.Dataset(filepath_indirect)
        for cube in self.testdata:
            assert n.all(ds.variables[cube.var_name][:] == cube.data)
        ds.close()
//  [python]
def test_basic_save(self):
        # test saving to a Dataset, in place of a filepath spec.
        # NOTE that this requires 'compute=False', as delayed saves can only operate on
        # a closed file.

        # Save to netcdf file in the usual way.
        filepath_direct = f"{self.temp_dir}/tmp_direct.nc"
        iris.save(self.testdata, filepath_direct)
        # Check against test-specific CDL result file.
        self.assertCDL(filepath_direct)

        # Save same data indirectly via a netcdf dataset.
        filepath_indirect = f"{self.temp_dir}/tmp_indirect.nc"
        nc_dataset = nc.Dataset(filepath_indirect, "w")
        # NOTE: we **must** use delayed saving here, as we cannot do direct saving to
        # a user-owned dataset.
        result = iris.save(self.testdata, nc_dataset, saver="nc", compute=False)

        # Do some very basic sanity checks on the resulting Dataset.
        # It should still be open (!)
        self.assertTrue(nc_dataset.isopen())
        self.assertEqual(
            ["time", "levelist", "latitude", "longitude"],
            list(nc_dataset.dimensions),
        )
        self.assertEqual(
            ["co2", "time", "levelist", "latitude", "longitude", "lnsp"],
            list(nc_dataset.variables),
        )
        nc_dataset.close()

        # Check the saved file against the same CDL as the 'normal' save.
        self.assertCDL(filepath_indirect)

        # Confirm that cube content is however not yet written.
        ds = nc.Dataset(filepath_indirect)
        for cube in self.testdata:
            assert np.all(ds.variables[cube.var_name][:].mask)
        ds.close()

        # Complete the delayed saves.
        result.compute()

        # Check that data now *is* written.
        ds = nc.Dataset(filepath_indirect)
        for cube in self.testdata:
            assert np.all(ds.variables[cube.var_name][:] == cube.data)
        ds.close()
//  [python]
def moment1_k(self):
        # prevailing tilt (frequency) for every coordinate (time)
        proj_s = self.proj_s()
        if np.alltrue(proj_s == 0):
            return proj_s
        else:
            proj_s[proj_s == 0] = np.nan
            return np.sum(self.wig * self.k[:, np.newaxis], axis=0) / proj_s
//  [python]
def moment1_k(self):
        # prevailing tilt (frequency) for every coordinate (time)
        proj_s = self.proj_s()
        if np.all(proj_s == 0):
            return proj_s
        else:
            proj_s[proj_s == 0] = np.nan
            return np.sum(self.wig * self.k[:, np.newaxis], axis=0) / proj_s
//  [python]
def moment1_s(self):
        # prevailing coordinate (time) for every tilt (frequency)
        proj_k = self.proj_k()
        if np.all(proj_k == 0):
            return proj_k
        else:
            proj_k[proj_k == 0] = np.nan
            return np.sum(self.wig * self.s[np.newaxis, :], axis=1) / proj_k
//  [python]
def set_action_dist_params(self, states_tensor):
        self.num_rows = states_tensor.shape[0]
        if self.mpi_enabled:
            # Some of the workers may have smaller samples than traj length,
            # since some environments may cut off the trajectory in the middle.
            comm = MPI.COMM_WORLD
            size = comm.Get_size()
            nom_row_avg_rank = comm.allreduce(self.num_rows, op=MPI.SUM)
            self.num_rows = float(nom_row_avg_rank) / size
        new_action_dist_params = self.generate_intermediate_params(states_tensor)
        new_action_dist_params = {key: tuple(a_ for a_ in val if a_.requires_grad) for key,val in new_action_dist_params.items()}
        if self.action_dist_params is None:
            self.action_dist_params = new_action_dist_params
            self.t = self.create_t(new_action_dist_params)
            # self.t is a dict like {'action_dist_params':(tensor1, tensor2,),
            #                        'action_jaccobians': (tensor3, tensor4,)}
        else:
            self.action_dist_params = new_action_dist_params

            same_old_shape = True
            for key in new_action_dist_params.keys():
                same_old_shape = same_old_shape and np.all([x_.shape == y_.shape for x_, y_ in zip(self.t[key], new_action_dist_params[key])])
            if not same_old_shape:
                for key, val in self.t.items():
                    for x_ in val:
                        del x_
                self.t = self.create_t(new_action_dist_params)
//  [python]
def __init__(self, filepath: str):
        super().__init__(filepath)

        # Read and store important quantities.
        self.structure = self.read_structure()
        qfrac_coords = self.read_value("qpoints_dielectric_function")
        self.kpoints = KpointList(self.structure.reciprocal_lattice, qfrac_coords)
        self.wpoints = self.read_value("frequencies_dielectric_function", cmode="c")
        self.ng = self.read_dimvalue("number_of_coefficients_dielectric_function")

        # Find number of real/imaginary frequencies.
        self.nw = len(self.wpoints)
        self.nrew = self.nw
        self.nimw = 0
        for i, w in enumerate(self.wpoints):
            if np.iscomplex(w):
                self.nrew = i
                break

        self.nimw = self.nw - self.nrew
        if self.nimw and not np.alltrue(np.iscomplex(self.wpoints[self.nrew+1:])):
            raise ValueError("wpoints should contained real points packed in the first positions\n"
                             "followed by imaginary points but got: %s" % str(self.wpoints))

        # Define self.netcdf_name from the data available on file.
        nfound = 0
        netcdf_names = ["polarizability", "dielectric_function", "inverse_dielectric_function"]
        for tryname in netcdf_names:
            if tryname in self.rootgrp.variables:
                self.netcdf_name = tryname
                nfound += 1

        if nfound == 0:
            raise RuntimeError("Cannot find `%s` in netcdf file" % str(netcdf_names))
        if nfound > 1:
            raise RuntimeError("Find multiple netcdf arrays (%s) in netcdf file!" % str(netcdf_names))