Add qk_obs_compose(_map) (#13950)

* add compose functions

* add a test

* null handling and less off-putting arg names

Co-authored-by: Jake Lishman <jake.lishman@ibm.com>

---------

Co-authored-by: Jake Lishman <jake.lishman@ibm.com>

crates/cext/src/sparse_observable.rs

@@ -631,6 +631,96 @@ pub unsafe extern "C" fn qk_obs_add(
     Box::into_raw(Box::new(result))
 }
 
+/// @ingroup QkObs
+/// Compose (multiply) two observables.
+///
+/// @param first One observable.
+/// @param second The other observable.
+///
+/// @return ``first.compose(second)`` which equals the observable ``result = second @ first``,
+///     in terms of the matrix multiplication ``@``.
+///
+/// # Example
+///
+///     QkObs *first = qk_obs_zero(100);
+///     QkObs *second = qk_obs_identity(100);
+///     QkObs *result = qk_obs_compose(first, second);
+///
+/// # Safety
+///
+/// Behavior is undefined if ``first`` or ``second`` are not valid, non-null pointers to
+/// ``QkObs``\ s.
+#[no_mangle]
+#[cfg(feature = "cbinding")]
+pub unsafe extern "C" fn qk_obs_compose(
+    first: *const SparseObservable,
+    second: *const SparseObservable,
+) -> *mut SparseObservable {
+    // SAFETY: Per documentation, the pointers are non-null and aligned.
+    let first = unsafe { const_ptr_as_ref(first) };
+    let second = unsafe { const_ptr_as_ref(second) };
+
+    let result = first.compose(second);
+    Box::into_raw(Box::new(result))
+}
+
+/// @ingroup QkObs
+/// Compose (multiply) two observables according to a custom qubit order.
+///
+/// Notably, this allows composing two observables of different size.
+///
+/// @param first One observable.
+/// @param second The other observable. The number of qubits must match the length of ``qargs``.
+/// @param qargs The qubit arguments specified which indices in ``first`` to associate with
+///     the ones in ``second``.
+///
+/// @return ``first.compose(second)`` which equals the observable ``result = second @ first``,
+///     in terms of the matrix multiplication ``@``.
+///
+/// # Example
+///
+///     QkObs *first = qk_obs_zero(100);
+///     QkObs *second = qk_obs_identity(100);
+///     QkObs *result = qk_obs_compose(first, second);
+///
+/// # Safety
+///
+/// To call this function safely
+///
+///   * ``first`` and ``second`` must be valid, non-null pointers to ``QkObs``\ s
+///   * ``qargs`` must point to an array of ``uint32_t``, readable for ``qk_obs_num_qubits(second)``
+///     elements (meaning the number of qubits in ``second``)
+#[no_mangle]
+#[cfg(feature = "cbinding")]
+pub unsafe extern "C" fn qk_obs_compose_map(
+    first: *const SparseObservable,
+    second: *const SparseObservable,
+    qargs: *const u32,
+) -> *mut SparseObservable {
+    // SAFETY: Per documentation, the pointers are non-null and aligned.
+    let first = unsafe { const_ptr_as_ref(first) };
+    let second = unsafe { const_ptr_as_ref(second) };
+
+    let qargs = if qargs.is_null() {
+        if second.num_qubits() != 0 {
+            panic!("If qargs is null, then second must have 0 qubits.");
+        }
+        &[]
+    } else {
+        if !qargs.is_aligned() {
+            panic!("qargs pointer is not aligned to u32");
+        }
+        // SAFETY: Per documentation, qargs is safe to read up to ``second.num_qubits()`` elements,
+        // which is the maximal value of ``index`` here.
+        unsafe { ::std::slice::from_raw_parts(qargs, second.num_qubits() as usize) }
+    };
+
+    let qargs_map = |index: u32| qargs[index as usize];
+
+    let result = first.compose_map(second, qargs_map);
+    Box::into_raw(Box::new(result))
+}
+
 /// @ingroup QkObs
 /// Calculate the canonical representation of the observable.
 ///

test/c/test_sparse_observable.c

@@ -75,7 +75,128 @@ int test_add() {
 }
 
 /**
- * Test multiplying two observables.
+ * Test composing two observables.
+ */
+int test_compose() {
+    u_int32_t num_qubits = 100;
+
+    QkObs *op1 = qk_obs_zero(num_qubits);
+    complex double coeff = 1;
+    QkBitTerm op1_bits[3] = {QkBitTerm_X, QkBitTerm_Y, QkBitTerm_Z};
+    uint32_t op1_indices[3] = {0, 1, 2};
+    QkObsTerm term1 = {coeff, 3, op1_bits, op1_indices, num_qubits};
+    qk_obs_add_term(op1, &term1);
+
+    QkObs *op2 = qk_obs_zero(num_qubits);
+    coeff = 2;
+    QkBitTerm op2_bits[3] = {QkBitTerm_Plus, QkBitTerm_X, QkBitTerm_Z};
+    uint32_t op2_indices[3] = {0, 1, 3};
+    QkObsTerm term2 = {coeff, 3, op2_bits, op2_indices, num_qubits};
+    qk_obs_add_term(op2, &term2);
+
+    QkObs *result = qk_obs_compose(op1, op2);
+
+    QkObs *expected = qk_obs_zero(num_qubits);
+    coeff = 2 * I;
+    QkBitTerm expected_bits[4] = {QkBitTerm_Plus, QkBitTerm_Z, QkBitTerm_Z, QkBitTerm_Z};
+    uint32_t expected_indices[4] = {0, 1, 2, 3};
+    QkObsTerm expected_term = {coeff, 4, expected_bits, expected_indices, num_qubits};
+    qk_obs_add_term(expected, &expected_term);
+
+    bool is_equal = qk_obs_equal(expected, result);
+
+    qk_obs_free(op1);
+    qk_obs_free(op2);
+    qk_obs_free(result);
+    qk_obs_free(expected);
+
+    if (!is_equal) {
+        return EqualityError;
+    }
+    return Ok;
+}
+
+/**
+ * Test composing two observables and specifying the qargs argument.
+ */
+int test_compose_map() {
+    u_int32_t num_qubits = 100;
+
+    QkObs *op1 = qk_obs_zero(num_qubits);
+    complex double coeff = 1;
+    QkBitTerm op1_bits[3] = {QkBitTerm_X, QkBitTerm_Y, QkBitTerm_Z};
+    uint32_t op1_indices[3] = {97, 98, 99};
+    QkObsTerm term1 = {coeff, 3, op1_bits, op1_indices, num_qubits};
+    qk_obs_add_term(op1, &term1);
+
+    QkObs *op2 = qk_obs_zero(2);
+    coeff = 2;
+    QkBitTerm op2_bits[3] = {QkBitTerm_Right, QkBitTerm_X};
+    uint32_t op2_indices[3] = {0, 1};
+    QkObsTerm term2 = {coeff, 2, op2_bits, op2_indices, 2};
+    qk_obs_add_term(op2, &term2);
+
+    uint32_t qargs[2] = {98, 97}; // compose op2 onto these indices in op1
+
+    QkObs *result = qk_obs_compose_map(op1, op2, qargs);
+
+    QkObs *expected = qk_obs_zero(num_qubits);
+    QkBitTerm expected_bits[2] = {QkBitTerm_Right, QkBitTerm_Z};
+    uint32_t expected_indices[2] = {98, 99};
+    QkObsTerm expected_term = {coeff, 2, expected_bits, expected_indices, num_qubits};
+    qk_obs_add_term(expected, &expected_term);
+
+    bool is_equal = qk_obs_equal(expected, result);
+
+    qk_obs_free(op1);
+    qk_obs_free(op2);
+    qk_obs_free(result);
+    qk_obs_free(expected);
+
+    if (!is_equal) {
+        return EqualityError;
+    }
+    return Ok;
+}
+
+/**
+ * Test composing an observables with a scalar observable.
+ */
+int test_compose_scalar() {
+    u_int32_t num_qubits = 100;
+
+    QkObs *op = qk_obs_zero(num_qubits);
+    complex double coeff = 1;
+    QkBitTerm bits[3] = {QkBitTerm_X, QkBitTerm_Y, QkBitTerm_Z};
+    uint32_t indices[3] = {97, 98, 99};
+    QkObsTerm term = {coeff, 3, bits, indices, num_qubits};
+    qk_obs_add_term(op, &term);
+
+    QkObs *scalar = qk_obs_identity(0);
+    coeff = 2;
+    QkObs *mult = qk_obs_multiply(scalar, &coeff);
+    uint32_t qargs[0];
+
+    QkObs *result = qk_obs_compose_map(op, mult, qargs);
+
+    QkObs *expected = qk_obs_multiply(op, &coeff);
+
+    bool is_equal = qk_obs_equal(expected, result);
+
+    qk_obs_free(op);
+    qk_obs_free(scalar);
+    qk_obs_free(mult);
+    qk_obs_free(result);
+    qk_obs_free(expected);
+
+    if (!is_equal) {
+        return EqualityError;
+    }
+    return Ok;
+}
+
+/**
+ * Test multiplying an observable by a complex coefficient.
  */
 int test_mult() {
     complex double coeffs[3] = {2, 2 * I, 2 + 2 * I};
@@ -604,6 +725,9 @@ int test_sparse_observable() {
     num_failed += RUN_TEST(test_zero);
     num_failed += RUN_TEST(test_identity);
     num_failed += RUN_TEST(test_add);
+    num_failed += RUN_TEST(test_compose);
+    num_failed += RUN_TEST(test_compose_map);
+    num_failed += RUN_TEST(test_compose_scalar);
     num_failed += RUN_TEST(test_mult);
     num_failed += RUN_TEST(test_canonicalize);
     num_failed += RUN_TEST(test_copy);