diff --git a/.github/workflows/all_push.yml b/.github/workflows/all_push.yml
index 32044da..4e8ec41 100644
--- a/.github/workflows/all_push.yml
+++ b/.github/workflows/all_push.yml
@@ -1,19 +1,16 @@
 name: All push
 
-on:
-  push:
-    branches: [ main ]
-  pull_request:
-    branches: [ main ]
+on: [push]
 
 jobs:
   build:
     runs-on: ${{ matrix.os }}
     strategy:
+      fail-fast: false
       matrix:
         os: [ubuntu-latest]
         compiler: [g++-12, clang++-17]
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.10"]
     steps:
     - uses: actions/checkout@v2
       with:
@@ -33,14 +30,24 @@ jobs:
     - name: Configure CMake
       env:
         CXX: ${{ matrix.compiler }}
-      run: cmake -B ${{github.workspace}}/build -DCMAKE_CXX_COMPILER=${CXX} -DCMAKE_BUILD_TYPE=Release
+      run: cmake -B ${{github.workspace}}/build -DCMAKE_CXX_COMPILER=${CXX}
 
     - name: Build
       # Build your program with the given configuration
-      run: cmake --build ${{github.workspace}}/build --parallel
+      run: |
+        cmake --build ${{github.workspace}}/build --verbose --parallel
+        du build/tests/test_pauli_op
 
     - name: Test C++
-      run: ctest --test-dir build --parallel
+      env:
+        OMP_NUM_THREADS: 2
+      run: |
+        # ctest --test-dir build --verbose # TODO not using bc of PauliOp problems on CI
+        ./build/tests/test_factory
+        ./build/tests/test_pauli
+        ./build/tests/test_pauli_op --test-case-exclude="*multistring*"
+        ./build/tests/test_pauli_string
+        ./build/tests/test_summed_pauli_op
 
     # - name: Test Python
-      # run: PYTHONPATH=build:$PYTHONPATH pytest -v test
\ No newline at end of file
+      # run: PYTHONPATH=build:$PYTHONPATH pytest -v test
diff --git a/.gitignore b/.gitignore
index bd3e442..aaf460a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -79,3 +79,4 @@ __pycache__
 logs
 scratch
 notes
+.cache
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 20652c1..a5dd3f5 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -17,7 +17,6 @@ include(cmake/CPM.cmake)
 
 # Dependencies
 cpmaddpackage("gh:doctest/doctest@2.4.11")
-# cpmaddpackage("gh:p-ranav/argparse@3.0")
 cpmaddpackage("gh:pybind/pybind11@2.11.1")
 cpmaddpackage("gh:fmtlib/fmt#10.2.1")
 cpmaddpackage("gh:kokkos/mdspan#b885a2c60ad42f9e1aaa0d317a38105b950cbed0")
@@ -60,11 +59,9 @@ target_compile_options(
             -Wall
             -Wextra
             -Werror
+            -Wpedantic
             # -stdlib=libc++
-            ${FAST_PAULI_EXTRA_CXX_COMPILE_FLAGS})
-target_link_options(fast_pauli INTERFACE ${FAST_PAULI_EXTRA_CXX_LD_FLAGS}
-                    -fuse-ld=mold)
-# target_compile_definitions(fast_pauli INTERFACE)
+)
 
 # Testing
 include(CTest)
diff --git a/README.md b/README.md
index a32b36c..4ad36cc 100644
--- a/README.md
+++ b/README.md
@@ -23,7 +23,7 @@
 
 ## Requirements
 
-- CMake >= 3.20
+- CMake >= 3.25
 - C++ compiler with OpenMP and C++20 support (LLVM recommended)
   - Tested Compilers GCC@12, LLVM@17, LLVM@18
 - Python >= 3.10
diff --git a/include/__factory.hpp b/include/__factory.hpp
index 7e7469f..ea9e255 100644
--- a/include/__factory.hpp
+++ b/include/__factory.hpp
@@ -1,6 +1,7 @@
 #ifndef __FAST_PAULI_FACTORY_HPP
 #define __FAST_PAULI_FACTORY_HPP
 
+#include <algorithm>
 #include <array>
 #include <complex>
 #include <experimental/mdspan>
@@ -74,7 +75,7 @@ auto rand(std::vector<T> &blob, std::array<size_t, n_dim> extents) {
   for (auto ei : extents) {
     total_size *= ei;
   }
-  blob.reserve(total_size);
+  blob = std::vector<T>(total_size);
 
   // Fill with random numbers
   std::random_device rd;
diff --git a/include/__pauli.hpp b/include/__pauli.hpp
index 1c01443..fbb94a3 100644
--- a/include/__pauli.hpp
+++ b/include/__pauli.hpp
@@ -21,24 +21,24 @@ struct Pauli {
    * @brief Default constructor, initializes to I.
    *
    */
-  constexpr Pauli() : code(0) {}
+  Pauli() : code(0) {}
 
   /**
    * @brief Constructor given a numeric code. TODO add input checking
    *
    * @tparam T Any type convertible to uint8_t
    * @param code 0: I, 1: X, 2: Y, 3: Z
-   * @return requires constexpr
+   * @return requires
    */
   template <class T>
     requires std::convertible_to<T, uint8_t>
-  constexpr Pauli(T const code) : code(code) {}
+  Pauli(T const code) : code(code) {}
 
   // Copy ctor
-  constexpr Pauli(Pauli const &other) = default;
+  Pauli(Pauli const &other) = default;
 
   // Copy assignment
-  constexpr Pauli &operator=(Pauli const &other) noexcept = default;
+  Pauli &operator=(Pauli const &other) noexcept = default;
 
   // The default operator does everything we want and make this more composable
   friend auto operator<=>(Pauli const &, Pauli const &) = default;
@@ -48,10 +48,10 @@ struct Pauli {
    *
    * @param lhs
    * @param rhs
-   * @return constexpr std::pair<std::complex<double>, Pauli>
+   * @return  std::pair<std::complex<double>, Pauli>
    */
-  friend constexpr std::pair<std::complex<double>, Pauli>
-  operator*(Pauli lhs, Pauli const &rhs) {
+  friend std::pair<std::complex<double>, Pauli> operator*(Pauli lhs,
+                                                          Pauli const &rhs) {
     switch (lhs.code) {
     case 0:
       switch (rhs.code) {
@@ -130,10 +130,10 @@ struct Pauli {
    * @brief Returns the pauli matrix as a 2D vector of complex numbers.
    *
    * @tparam T floating point type
-   * @return constexpr std::vector<std::vector<std::complex<T>>>
+   * @return  std::vector<std::vector<std::complex<T>>>
    */
   template <std::floating_point T>
-  constexpr std::vector<std::vector<std::complex<T>>> to_tensor() const {
+  std::vector<std::vector<std::complex<T>>> to_tensor() const {
     std::vector<std::vector<std::complex<T>>> result;
     switch (code) {
     case 0:
@@ -160,7 +160,6 @@ struct Pauli {
 } // namespace fast_pauli
 
 // Adding specialization to the fmt library so we can easily print Pauli
-namespace fmt {
 template <> struct fmt::formatter<fast_pauli::Pauli> {
   constexpr auto parse(format_parse_context &ctx) { return ctx.begin(); }
 
@@ -199,6 +198,4 @@ template <> struct fmt::formatter<std::complex<double>> {
   }
 };
 
-} // namespace fmt
-
 #endif // __PAULI_HPP
diff --git a/include/__pauli_op.hpp b/include/__pauli_op.hpp
index 8d77265..e886f45 100644
--- a/include/__pauli_op.hpp
+++ b/include/__pauli_op.hpp
@@ -19,11 +19,11 @@ template <std::floating_point T> struct PauliOp {
   // possible combinations up to and including weight 3 strings?)
   std::vector<PauliString> pauli_strings;
 
-  constexpr PauliOp() = default;
+  PauliOp() = default;
 
   //
-  constexpr PauliOp(std::vector<std::complex<T>> const &coeffs,
-                    std::vector<PauliString> const &pauli_strings)
+  PauliOp(std::vector<std::complex<T>> const &coeffs,
+          std::vector<PauliString> const &pauli_strings)
       : coeffs(coeffs), pauli_strings(pauli_strings) {
     // TODO may want to wrap this in a #IFDEF DEBUG block to avoid the overhead
     // input check
@@ -44,7 +44,7 @@ template <std::floating_point T> struct PauliOp {
     }
   }
 
-  constexpr size_t dims() const {
+  size_t dims() const {
     if (pauli_strings.size() > 0) {
       return pauli_strings[0].dims();
     } else {
@@ -52,7 +52,7 @@ template <std::floating_point T> struct PauliOp {
     }
   }
 
-  constexpr std::vector<std::complex<T>>
+  std::vector<std::complex<T>>
   apply(std::vector<std::complex<T>> const &state) const {
     // input check
     if (state.size() != dims()) {
@@ -110,6 +110,10 @@ template <std::floating_point T> struct PauliOp {
   void
   apply(mdspan<std::complex<T>, std::dextents<size_t, 2>> new_states,
         mdspan<std::complex<T>, std::dextents<size_t, 2>> const states) const {
+
+    fmt::println(
+        "[WARNING] Apply function causes problems on CI, use with caution.");
+
     // input check
     if (states.extent(0) != this->dims()) {
       throw std::invalid_argument(
@@ -143,25 +147,29 @@ template <std::floating_point T> struct PauliOp {
     std::mdspan<std::complex<T>, std::dextents<size_t, 3>> new_states_thr(
         new_states_thr_raw.data(), n_threads, n_dim, n_data);
 
-//
-#pragma omp parallel for schedule(static)
-    for (size_t i = 0; i < pauli_strings.size(); ++i) {
-      size_t const tid = omp_get_thread_num();
+    //
 
-      PauliString const &ps = pauli_strings[i];
-      std::complex<T> c = coeffs[i];
-      std::mdspan<std::complex<T>, std::dextents<size_t, 2>> new_states_local =
-          std::submdspan(new_states_thr, tid, std::full_extent,
-                         std::full_extent);
-      ps.apply_batch(new_states_local, states, c);
-    }
+#pragma omp parallel
+    {
+#pragma omp for schedule(static)
+      for (size_t i = 0; i < pauli_strings.size(); ++i) {
+        size_t const tid = omp_get_thread_num();
+
+        PauliString const &ps = pauli_strings[i];
+        std::complex<T> c = coeffs[i];
+        std::mdspan<std::complex<T>, std::dextents<size_t, 2>>
+            new_states_local = std::submdspan(
+                new_states_thr, tid, std::full_extent, std::full_extent);
+        ps.apply_batch(new_states_local, states, c);
+      }
 
-    // Do the reduction and tranpose back
-#pragma omp parallel for schedule(static) collapse(2)
-    for (size_t i = 0; i < new_states.extent(0); ++i) {
-      for (size_t t = 0; t < new_states.extent(1); ++t) {
-        for (size_t th = 0; th < n_threads; ++th) {
-          new_states(i, t) += new_states_thr(th, i, t);
+      // Do the reduction and tranpose back
+#pragma omp for schedule(static) collapse(2)
+      for (size_t i = 0; i < new_states.extent(0); ++i) {
+        for (size_t t = 0; t < new_states.extent(1); ++t) {
+          for (size_t th = 0; th < n_threads; ++th) {
+            new_states(i, t) += new_states_thr(th, i, t);
+          }
         }
       }
     }
@@ -170,7 +178,7 @@ template <std::floating_point T> struct PauliOp {
   //
   // Helpers (mostly for debugging)
   //
-  constexpr std::vector<std::vector<std::complex<T>>> get_dense_repr() const {
+  std::vector<std::vector<std::complex<T>>> get_dense_repr() const {
     size_t const dim = 1UL << pauli_strings[0].paulis.size();
 
     std::vector<std::vector<std::complex<T>>> res(
diff --git a/include/__pauli_string.hpp b/include/__pauli_string.hpp
index 4f59038..7e121d4 100644
--- a/include/__pauli_string.hpp
+++ b/include/__pauli_string.hpp
@@ -32,14 +32,14 @@ struct PauliString {
    * @brief Default constructor, initialize weight and empty vector for paulis.
    *
    */
-  constexpr PauliString() noexcept = default;
+  PauliString() noexcept = default;
 
   /**
    * @brief Constructs a PauliString from a span of pauli operators and
    * calculates the weight.
    *
    */
-  constexpr PauliString(std::span<fast_pauli::Pauli> const &paulis)
+  PauliString(std::span<fast_pauli::Pauli> const &paulis)
       : weight(0), paulis(paulis.begin(), paulis.end()) {
     for (auto const &pauli : paulis) {
       weight += pauli.code > 0;
@@ -51,7 +51,7 @@ struct PauliString {
    * This is often the most compact way to initialize a PauliString.
    *
    */
-  constexpr PauliString(std::string const &str) : weight(0) {
+  PauliString(std::string const &str) : weight(0) {
     for (auto const &c : str) {
       switch (c) {
       case 'I':
@@ -81,11 +81,11 @@ struct PauliString {
    * @brief
    *
    */
-  constexpr PauliString(char const *str) : PauliString(std::string(str)) {}
+  PauliString(char const *str) : PauliString(std::string(str)) {}
 
-  constexpr PauliString(PauliString const &other)
+  PauliString(PauliString const &other)
       : weight(other.weight), paulis(other.paulis){};
-  constexpr PauliString &operator=(PauliString const &other) {
+  PauliString &operator=(PauliString const &other) {
     this->weight = other.weight;
     this->paulis = other.paulis;
     return *this;
@@ -98,16 +98,16 @@ struct PauliString {
   /**
    * @brief Return the number of qubits in the PauliString.
    *
-   * @return constexpr size_t
+   * @return  size_t
    */
-  constexpr size_t n_qubits() const noexcept { return paulis.size(); }
+  size_t n_qubits() const noexcept { return paulis.size(); }
 
   /**
    * @brief Return the dimension (2^n_qubits) of the PauliString.
    *
-   * @return constexpr size_t
+   * @return  size_t
    */
-  constexpr size_t dims() const noexcept { return 1UL << paulis.size(); }
+  size_t dims() const noexcept { return 1UL << paulis.size(); }
 
   /**
    * @brief Get the sparse representation of the pauli string matrix.
@@ -126,8 +126,8 @@ struct PauliString {
    * @param m The values of the matrix
    */
   template <std::floating_point T>
-  constexpr void get_sparse_repr(std::vector<size_t> &j, std::vector<size_t> &k,
-                                 std::vector<std::complex<T>> &m) const {
+  void get_sparse_repr(std::vector<size_t> &j, std::vector<size_t> &k,
+                       std::vector<std::complex<T>> &m) const {
     // We reverse the view here because the tensor product is taken from right
     // to left
     auto ps = paulis | std::views::reverse;
@@ -184,11 +184,11 @@ struct PauliString {
    * @tparam T The floating point base to use for all the complex numbers
    * @param v The input vector to apply the PauliString to. Must be the same
    * size as PauliString.dims().
-   * @return constexpr std::vector<std::complex<T>> The output state after
+   * @return  std::vector<std::complex<T>> The output state after
    * applying the PauliString.
    */
   template <std::floating_point T>
-  constexpr std::vector<std::complex<T>>
+  std::vector<std::complex<T>>
   apply(std::vector<std::complex<T>> const &v) const {
     // Input check
     if (v.size() != dims()) {
@@ -214,11 +214,11 @@ struct PauliString {
    * @tparam T The floating point base to use for all the complex numbers
    * @param v The input vector to apply the PauliString to. Must be the same
    * size as PauliString.dims().
-   * @return constexpr std::vector<std::complex<T>> The output state after
+   * @return  std::vector<std::complex<T>> The output state after
    * applying the PauliString.
    */
   template <std::floating_point T>
-  constexpr std::vector<std::complex<T>>
+  std::vector<std::complex<T>>
   apply(std::mdspan<std::complex<T>, std::dextents<size_t, 1>> v) const {
     // Input check
     if (v.size() != dims()) {
@@ -242,6 +242,7 @@ struct PauliString {
    * @brief Apply the PauliString to a batch of states. This function takes a
    * different shape of the states than the other apply functions. here all the
    * states (new and old) are transposed so their shape is (n_dims x n_states).
+   * All the new_stats are overwritten, no need to initialize.
    *
    * @tparam T The floating point base to use for all the complex numbers
    * @param new_states_T The outpus states after applying the PauliString
@@ -251,12 +252,11 @@ struct PauliString {
    * @param c Multiplication factor to apply to the PauliString
    */
   template <std::floating_point T>
-  constexpr void
-  apply_batch(std::mdspan<std::complex<T>, std::dextents<size_t, 2>>
-                  new_states_T, // extent(0) = dims, extent(1) = n_states
-              std::mdspan<std::complex<T>, std::dextents<size_t, 2>> const
-                  states_T, // extent(0) = dims, extent(1) = n_states
-              std::complex<T> const c) const {
+  void apply_batch(std::mdspan<std::complex<T>, std::dextents<size_t, 2>>
+                       new_states_T, // extent(0) = dims, extent(1) = n_states
+                   std::mdspan<std::complex<T>, std::dextents<size_t, 2>> const
+                       states_T, // extent(0) = dims, extent(1) = n_states
+                   std::complex<T> const c) const {
     // Input check
     if (states_T.extent(0) != dims()) {
       auto error_msg =
@@ -297,10 +297,10 @@ struct PauliString {
    * @brief Get the dense representation of the object as a 2D-std::vector
    *
    * @tparam T The floating point base to use for all the complex numbers
-   * @return constexpr std::vector<std::vector<std::complex<T>>>
+   * @return  std::vector<std::vector<std::complex<T>>>
    */
   template <std::floating_point T>
-  constexpr std::vector<std::vector<std::complex<T>>> get_dense_repr() const {
+  std::vector<std::vector<std::complex<T>>> get_dense_repr() const {
     //
     std::vector<size_t> j, k;
     std::vector<std::complex<T>> m;
@@ -450,7 +450,6 @@ std::vector<PauliString> calculate_pauli_strings_max_weight(size_t n_qubits,
 // fmt::formatter specialization
 //
 
-namespace fmt {
 //
 template <> struct fmt::formatter<fast_pauli::PauliString> {
   constexpr auto parse(format_parse_context &ctx) { return ctx.begin(); }
@@ -462,6 +461,4 @@ template <> struct fmt::formatter<fast_pauli::PauliString> {
   }
 };
 
-} // namespace fmt
-
 #endif // __PAULI_STRING_HPP
diff --git a/include/__summed_pauli_op.hpp b/include/__summed_pauli_op.hpp
index 56fb0d6..800bc42 100644
--- a/include/__summed_pauli_op.hpp
+++ b/include/__summed_pauli_op.hpp
@@ -22,11 +22,11 @@ template <std::floating_point T> struct SummedPauliOp {
   size_t _dim;
   size_t _n_operators;
 
-  constexpr SummedPauliOp() noexcept = default;
+  SummedPauliOp() noexcept = default;
 
   //
-  constexpr SummedPauliOp(std::vector<PauliString> const &pauli_strings,
-                          std::vector<std::complex<T>> const &coeffs_raw)
+  SummedPauliOp(std::vector<PauliString> const &pauli_strings,
+                std::vector<std::complex<T>> const &coeffs_raw)
       : pauli_strings(pauli_strings), coeffs_raw(coeffs_raw) {
 
     // TODO add more checks
@@ -47,8 +47,8 @@ template <std::floating_point T> struct SummedPauliOp {
     }
   }
 
-  constexpr SummedPauliOp(std::vector<PauliString> const &pauli_strings,
-                          Tensor<2> const coeffs)
+  SummedPauliOp(std::vector<PauliString> const &pauli_strings,
+                Tensor<2> const coeffs)
       : pauli_strings(pauli_strings) {
 
     // Check that the dims are all the same
@@ -77,11 +77,9 @@ template <std::floating_point T> struct SummedPauliOp {
   //
   // Accessors/helpers
   //
-  constexpr size_t n_dimensions() const noexcept { return _dim; }
-  constexpr size_t n_operators() const noexcept { return _n_operators; }
-  constexpr size_t n_pauli_strings() const noexcept {
-    return pauli_strings.size();
-  }
+  size_t n_dimensions() const noexcept { return _dim; }
+  size_t n_operators() const noexcept { return _n_operators; }
+  size_t n_pauli_strings() const noexcept { return pauli_strings.size(); }
 
   //
   // Primary (TODO "primary" is vague here) functions
diff --git a/tests/test_pauli_op.cpp b/tests/test_pauli_op.cpp
index cdc434d..f7082f6 100644
--- a/tests/test_pauli_op.cpp
+++ b/tests/test_pauli_op.cpp
@@ -8,6 +8,65 @@
 using namespace std::literals;
 using namespace fast_pauli;
 
+//
+// Helpers
+//
+
+void __check_apply(PauliOp<double> &pauli_op, size_t n_states) {
+  size_t const dims = pauli_op.dims();
+
+  // Set up random states
+  std::vector<std::complex<double>> states_raw;
+  std::mdspan states =
+      fast_pauli::rand<std::complex<double>, 2>(states_raw, {dims, n_states});
+  // fmt::println("states: \n{}\n", fmt::join(states_raw, ",\n "));
+
+  // Apply the PauliOp to a batch of states
+  std::vector<std::complex<double>> new_states_raw(dims * n_states, 0);
+  std::mdspan<std::complex<double>, std::dextents<size_t, 2>> new_states(
+      new_states_raw.data(), dims, n_states);
+  pauli_op.apply(new_states, states);
+  // fmt::println("new_states: \n{}\n", fmt::join(new_states_raw, ",\n "));
+
+  //
+  // Calculate the expected new states
+  //
+  std::vector<std::vector<std::complex<double>>> pop_dense =
+      pauli_op.get_dense_repr();
+
+  // pop_dense : d x d
+  // states : n x d
+  // expected_states : n x d
+
+  std::vector<std::complex<double>> expected(dims * n_states, 0);
+  std::mdspan<std::complex<double>, std::dextents<size_t, 2>> expected_span(
+      expected.data(), dims, n_states);
+
+  // Calculate the expected new states
+  for (size_t t = 0; t < n_states; ++t) {
+    for (size_t i = 0; i < dims; ++i) {
+      for (size_t j = 0; j < dims; ++j) {
+        expected_span(i, t) += pop_dense[i][j] * states(j, t);
+      }
+    }
+  }
+
+  // Check
+  for (size_t t = 0; t < n_states; ++t) {
+    for (size_t i = 0; i < dims; ++i) {
+      CHECK(abs(new_states(i, t) - expected_span(i, t)) < 1e-6);
+    }
+  }
+}
+
+//
+// Tests
+//
+
+//
+// Constructors
+//
+
 TEST_CASE("test pauli op") {
   std::vector<PauliString> pauli_strings = {"IXYZ", "XIII", "XXYZ"};
   std::vector<std::complex<double>> coeffs = {1i, 2., -1i};
@@ -30,6 +89,29 @@ TEST_CASE("test bad apply") {
   CHECK_THROWS(pauli_op.apply(state));
 }
 
+//
+// Member functions
+//
+
+TEST_CASE("test get_dense_repr") {
+  std::vector<PauliString> pauli_strings = {"III", "III", "III"};
+  std::vector<std::complex<double>> coeffs = {1, 2., 1};
+  PauliOp<double> pauli_op(coeffs, pauli_strings);
+  auto pop_dense = pauli_op.get_dense_repr();
+
+  size_t const dim = pauli_op.dims();
+
+  for (size_t i = 0; i < dim; ++i) {
+    for (size_t j = 0; j < dim; ++j) {
+      if (i == j) {
+        CHECK(abs(pop_dense[i][j] - std::complex<double>(4)) < 1e-6);
+      } else {
+        CHECK(abs(pop_dense[i][j]) < 1e-6);
+      }
+    }
+  }
+}
+
 // TODO add tests for multiple pauli strings
 TEST_CASE("test apply simple") {
   // Set up PauliOp
@@ -58,68 +140,59 @@ TEST_CASE("test apply simple") {
   }
 }
 
-TEST_CASE("test apply multistate") {
+TEST_CASE("test apply single state") {
   // Set up PauliOp
   std::vector<PauliString> pauli_strings = {"IXYZ"};
   std::vector<std::complex<double>> coeffs = {1i};
   PauliOp<double> pauli_op(coeffs, pauli_strings);
+  __check_apply(pauli_op, 1);
+}
 
-  size_t const dims = pauli_strings[0].dims();
-
-  // Set up random states
-  size_t const n_states = 10; // TODO change back to more than 1
-  std::vector<std::complex<double>> states_raw;
-  std::mdspan states =
-      fast_pauli::rand<std::complex<double>, 2>(states_raw, {dims, n_states});
-
-  // Apply the PauliOp to a batch of states
-  std::vector<std::complex<double>> new_states_raw(dims * n_states, 0);
-  std::mdspan<std::complex<double>, std::dextents<size_t, 2>> new_states(
-      new_states_raw.data(), dims, n_states);
-  pauli_op.apply(new_states, states);
-
-  //
-  // Calculate the expected new states
-  //
-  auto pop_dense = pauli_op.get_dense_repr();
-
-  // pop_dense : d x d
-  // states : n x d
-  // expected_states : n x d
-
-  std::vector<std::complex<double>> expected(dims * n_states, 0);
-  std::mdspan<std::complex<double>, std::dextents<size_t, 2>> expected_span(
-      expected.data(), dims, n_states);
+TEST_CASE("test apply two states") {
+  // Set up PauliOp
+  std::vector<PauliString> pauli_strings = {"IXYZ"};
+  std::vector<std::complex<double>> coeffs = {1i};
+  PauliOp<double> pauli_op(coeffs, pauli_strings);
+  __check_apply(pauli_op, 2);
+}
 
-  // Calculate the expected new states
-  for (size_t t = 0; t < n_states; ++t) {
-    for (size_t i = 0; i < dims; ++i) {
-      for (size_t j = 0; j < dims; ++j) {
-        expected_span(i, t) += pop_dense[i][j] * states(j, t);
-      }
-    }
-  }
+TEST_CASE("test apply multistate") {
+  // Set up PauliOp
+  std::vector<PauliString> pauli_strings = {"IXYZ"};
+  std::vector<std::complex<double>> coeffs = {1i};
+  PauliOp<double> pauli_op(coeffs, pauli_strings);
+  __check_apply(pauli_op, 10);
+}
 
-  // Check
-  for (size_t t = 0; t < n_states; ++t) {
-    for (size_t i = 0; i < dims; ++i) {
-      CHECK(abs(new_states(i, t) - expected_span(i, t)) < 1e-6);
-    }
-  }
+TEST_CASE("test apply single state multistring") {
+  // Set up PauliOp
+  std::vector<PauliString> pauli_strings = {"XXY", "YYZ"};
+  std::vector<std::complex<double>> coeffs = {1i, -1.23};
+  PauliOp<double> pauli_op(coeffs, pauli_strings);
+  __check_apply(pauli_op, 1);
 }
 
 TEST_CASE("test apply multistate multistring") {
   // Set up PauliOp
-  std::vector<PauliString> pauli_strings = {"IXYZ", "XXXX", "YZZZ", "ZIII",
-                                            "IIII"};
-  std::vector<std::complex<double>> coeffs = {1i, -2., 42i, 0.5,
-                                              std::complex<double>(1e-3, -1e1)};
+  std::vector<PauliString> pauli_strings = {"IXXZYZ", "XXXXZX", "YZZXZZ",
+                                            "ZXZIII", "IIIXZI", "ZZXZYY"};
+  std::vector<std::complex<double>> coeffs = {1i, -2., 42i, 0.5, 1, 0.1};
   PauliOp<double> pauli_op(coeffs, pauli_strings);
+  __check_apply(pauli_op, 10);
+}
 
+TEST_CASE("test apply multistate multistring identity") {
+  fmt::println("test apply multistate multistring identity");
+
+  // Set up PauliOp
+  std::vector<PauliString> pauli_strings(16, "IIII");
+  std::vector<std::complex<double>> coeffs(pauli_strings.size(),
+                                           1. / pauli_strings.size());
+  PauliOp<double> pauli_op(coeffs, pauli_strings);
   size_t const dims = pauli_strings[0].dims();
 
   // Set up random states
-  size_t const n_states = 10; // TODO change back to more than 1
+  size_t const n_states = 10;
   std::vector<std::complex<double>> states_raw;
   std::mdspan states =
       fast_pauli::rand<std::complex<double>, 2>(states_raw, {dims, n_states});
@@ -130,32 +203,16 @@ TEST_CASE("test apply multistate multistring") {
       new_states_raw.data(), dims, n_states);
   pauli_op.apply(new_states, states);
 
-  //
-  // Calculate the expected new states
-  //
-  auto pop_dense = pauli_op.get_dense_repr();
-
-  // pop_dense : d x d
-  // states : n x d
-  // expected_states : n x d
-
-  std::vector<std::complex<double>> expected(dims * n_states, 0);
-  std::mdspan<std::complex<double>, std::dextents<size_t, 2>> expected_span(
-      expected.data(), dims, n_states);
-
-  // Calculate the expected new states
-  for (size_t t = 0; t < n_states; ++t) {
-    for (size_t i = 0; i < dims; ++i) {
-      for (size_t j = 0; j < dims; ++j) {
-        expected_span(i, t) += pop_dense[i][j] * states(j, t);
-      }
-    }
-  }
-
+  // States should be unchanged
   // Check
   for (size_t t = 0; t < n_states; ++t) {
     for (size_t i = 0; i < dims; ++i) {
-      CHECK(abs(new_states(i, t) - expected_span(i, t)) < 1e-6);
+      CHECK(abs(new_states(i, t) - states(i, t)) < 1e-6);
+      if (abs(new_states(i, t) - states(i, t)) > 1e-6) {
+        fmt::print("new_states(i, t): {}, states(i, t): {}", new_states(i, t),
+                   states(i, t));
+        fmt::println(" ratio {}", new_states(i, t) / states(i, t));
+      }
     }
   }
 }
diff --git a/tests/test_summed_pauli_op.cpp b/tests/test_summed_pauli_op.cpp
index 5c48c37..38f2aa5 100644
--- a/tests/test_summed_pauli_op.cpp
+++ b/tests/test_summed_pauli_op.cpp
@@ -1,6 +1,8 @@
+#include "__pauli.hpp"
+#include "__pauli_string.hpp"
+#include <cstddef>
 #define DOCTEST_CONFIG_IMPLEMENT_WITH_MAIN
 
-// #include "__summed_pauli_op.hpp"
 #include <doctest/doctest.h>
 #include <fmt/core.h>
 #include <fmt/ranges.h>
@@ -10,6 +12,72 @@
 using namespace std::literals;
 using namespace fast_pauli;
 
+/**
+ * @brief Helper function to check the apply function on a set of states for a
+ * given set of pauli strings and coefficients.
+ *
+ * @param pauli_strings
+ * @param coeff
+ */
+void __check_apply(
+    std::vector<PauliString> &pauli_strings,
+    std::mdspan<std::complex<double>, std::dextents<size_t, 2>> coeff,
+    size_t const n_states = 10) {
+
+  SummedPauliOp<double> summed_op{pauli_strings, coeff};
+
+  // Setup states
+  size_t const dim = summed_op.n_dimensions();
+  size_t const n_ops = summed_op.n_operators();
+
+  std::vector<std::complex<double>> states_raw;
+  std::mdspan states =
+      fast_pauli::rand<std::complex<double>, 2>(states_raw, {dim, n_states});
+
+  std::vector<std::complex<double>> new_states_raw;
+  std::mdspan new_states = fast_pauli::zeros<std::complex<double>, 2>(
+      new_states_raw, {dim, n_states});
+
+  // Init data (aka data)
+  std::vector<double> data_raw;
+  std::mdspan data = fast_pauli::rand<double, 2>(data_raw, {n_ops, n_states});
+
+  // Apply the summed operator
+  summed_op.apply(new_states, states, data);
+
+  // Check the check
+  std::vector<std::complex<double>> expected_raw;
+  std::mdspan expected =
+      fast_pauli::zeros<std::complex<double>, 2>(expected_raw, {dim, n_states});
+
+  for (size_t j = 0; j < summed_op.n_pauli_strings(); ++j) {
+    auto ps = summed_op.pauli_strings[j];
+
+    PauliOp<double> pop{{std::complex<double>(1.)}, {ps}};
+    std::vector<std::complex<double>> tmp_raw;
+    std::mdspan tmp =
+        fast_pauli::zeros<std::complex<double>, 2>(tmp_raw, {dim, n_states});
+
+    pop.apply(tmp, states);
+
+    // Manually calculate the sum over the different pauli operators
+    // This is specific to the coefficients we've chosen above
+    for (size_t i = 0; i < dim; ++i) {
+      for (size_t t = 0; t < n_states; ++t) {
+        for (size_t k = 0; k < n_ops; ++k) {
+          expected(i, t) += data(k, t) * summed_op.coeffs(j, k) * tmp(i, t);
+        }
+      }
+    }
+  }
+
+  for (size_t t = 0; t < n_states; ++t) {
+    for (size_t i = 0; i < dim; ++i) {
+      CHECK(abs(new_states(i, t) - expected(i, t)) < 1e-6);
+    }
+  }
+}
+
 TEST_CASE("ctors") {
   {
     // Default
@@ -57,144 +125,34 @@ TEST_CASE("accessors") {
 TEST_CASE("apply 1 operator 1 PauliString") {
   fmt::print("\n\napply 1 operator 1 PauliString\n");
   // Setup operator
-  SummedPauliOp<double> summed_op{{"XYZ"}, {std::complex<double>(1.)}};
-
-  // Setup states
-  size_t const n_states = 1;
-  size_t const dim = summed_op.n_dimensions();
-  size_t const n_ops = summed_op.n_operators();
 
-  std::vector<std::complex<double>> states_raw;
-  std::mdspan states =
-      fast_pauli::rand<std::complex<double>, 2>(states_raw, {dim, n_states});
+  std::vector<PauliString> pauli_strings = {"XYZ"};
+  std::vector<std::complex<double>> coeff_raw = {1i};
+  std::mdspan<std::complex<double>, std::dextents<size_t, 2>> coeff(
+      coeff_raw.data(), 1, 1);
 
-  std::vector<std::complex<double>> new_states_raw;
-  std::mdspan new_states = fast_pauli::zeros<std::complex<double>, 2>(
-      new_states_raw, {dim, n_states});
-
-  // Init data
-  std::vector<double> data_raw(n_ops * n_states, 1);
-  std::mdspan<double, std::dextents<size_t, 2>> data(data_raw.data(), n_ops,
-                                                     n_states);
-  summed_op.apply(new_states, states, data);
-
-  // Calculated the expected answer and check
-  std::vector<std::complex<double>> expected_raw;
-  std::mdspan expected =
-      fast_pauli::zeros<std::complex<double>, 2>(expected_raw, {dim, n_states});
-
-  PauliOp<double> pop{{std::complex<double>(1.)}, {"XYZ"}};
-  pop.apply(expected, states);
-
-  for (size_t t = 0; t < n_states; ++t) {
-    for (size_t i = 0; i < dim; ++i) {
-      CHECK(abs(new_states(i, t) - expected(i, t)) < 1e-6);
-    }
-  }
+  __check_apply(pauli_strings, coeff);
 }
 
 TEST_CASE("apply 2 operators 1 PauliString") {
   fmt::print("\n\napply 2 operators 1 PauliString\n");
 
   // Setup operator
-  SummedPauliOp<double> summed_op{{"XYZ"}, {std::complex<double>(1.), 1i}};
-
-  // Setup states
-  size_t const n_states = 1;
-  size_t const dim = summed_op.n_dimensions();
-  size_t const n_ops = summed_op.n_operators();
-
-  std::vector<std::complex<double>> states_raw;
-  std::mdspan states =
-      fast_pauli::rand<std::complex<double>, 2>(states_raw, {dim, n_states});
-
-  std::vector<std::complex<double>> new_states_raw;
-  std::mdspan new_states = fast_pauli::zeros<std::complex<double>, 2>(
-      new_states_raw, {dim, n_states});
-
-  // Init data (aka data)
-  std::vector<double> data_raw(n_ops * n_states, 1);
-  std::mdspan<double, std::dextents<size_t, 2>> data(data_raw.data(), n_ops,
-                                                     n_states);
-  summed_op.apply(new_states, states, data);
-
-  // Calculated the expected answer and check
-  std::vector<std::complex<double>> expected_raw;
-  std::mdspan expected =
-      fast_pauli::zeros<std::complex<double>, 2>(expected_raw, {dim, n_states});
-
-  // Manually calculate the sum over the different pauli operators
-  for (size_t i = 0; i < dim; ++i) {
-    for (size_t t = 0; t < n_states; ++t) {
-      expected(i, t) += 1i * expected(i, t);
-    }
-  }
-
-  for (size_t t = 0; t < n_states; ++t) {
-    for (size_t i = 0; i < dim; ++i) {
-      CHECK(abs(new_states(i, t) - expected(i, t)) < 1e-6);
-    }
-  }
+  std::vector<PauliString> pauli_strings = {"XYZ"};
+  std::vector<std::complex<double>> coeff_raw = {1i, 1};
+  std::mdspan<std::complex<double>, std::dextents<size_t, 2>> coeff(
+      coeff_raw.data(), 1, 2);
+  __check_apply(pauli_strings, coeff);
 }
 
 TEST_CASE("apply 2 operators 2 PauliString") {
   fmt::print("\n\napply 2 operators 2 PauliString\n");
   // Setup operator
-  SummedPauliOp<double> summed_op{{"XYZ", "YYZ"}, {1i, 1, 0.5i, -0.99}};
-
-  // Setup states
-  size_t const n_states = 10;
-  size_t const dim = summed_op.n_dimensions();
-  size_t const n_ops = summed_op.n_operators();
-
-  std::vector<std::complex<double>> states_raw;
-  std::mdspan states =
-      fast_pauli::rand<std::complex<double>, 2>(states_raw, {dim, n_states});
-
-  std::vector<std::complex<double>> new_states_raw;
-  std::mdspan new_states = fast_pauli::zeros<std::complex<double>, 2>(
-      new_states_raw, {dim, n_states});
-
-  // Init data
-  std::vector<double> data_raw(n_ops * n_states, 1);
-  std::mdspan<double, std::dextents<size_t, 2>> data(data_raw.data(), n_ops,
-                                                     n_states);
-  summed_op.apply(new_states, states, data);
-
-  // Calculated expected answer
-  std::vector<std::complex<double>> expected_raw;
-  std::mdspan expected =
-      fast_pauli::zeros<std::complex<double>, 2>(expected_raw, {dim, n_states});
-
-  for (size_t j = 0; j < summed_op.n_pauli_strings(); ++j) {
-
-    auto ps = summed_op.pauli_strings[j];
-
-    PauliOp<double> pop{{std::complex<double>(1.)}, {ps}};
-    std::vector<std::complex<double>> tmp_raw;
-    std::mdspan tmp =
-        fast_pauli::zeros<std::complex<double>, 2>(tmp_raw, {dim, n_states});
-
-    pop.apply(tmp, states);
-
-    // Manually calculate the sum over the different pauli operators
-    // This is specific to the coefficients we've chosen above
-    for (size_t i = 0; i < dim; ++i) {
-      for (size_t t = 0; t < n_states; ++t) {
-        // expected(i, t) += 1i * tmp(i, t);
-        for (size_t k = 0; k < n_ops; ++k) {
-          expected(i, t) += data(k, t) * summed_op.coeffs(j, k) * tmp(i, t);
-        }
-        // expected(i, t) += std::complex<double>(2) * tmp(i, t);
-      }
-    }
-  }
-
-  for (size_t t = 0; t < n_states; ++t) {
-    for (size_t i = 0; i < dim; ++i) {
-      CHECK(abs(new_states(i, t) - expected(i, t)) < 1e-6);
-    }
-  }
+  std::vector<std::complex<double>> coeff_raw = {1i, 1, 0.5i, -0.99};
+  std::mdspan<std::complex<double>, std::dextents<size_t, 2>> coeff(
+      coeff_raw.data(), 2, 2);
+  std::vector<PauliString> pauli_strings = {"XYZ", "YYZ"};
+  __check_apply(pauli_strings, coeff);
 }
 
 TEST_CASE("apply many operators many PauliString") {
@@ -207,59 +165,7 @@ TEST_CASE("apply many operators many PauliString") {
   std::mdspan coeff = fast_pauli::rand<std::complex<double>, 2>(
       coeff_raw, {pauli_strings.size(), 100});
 
-  SummedPauliOp<double> summed_op{pauli_strings, coeff};
-
-  // Setup states
-  size_t const n_states = 10;
-  size_t const dim = summed_op.n_dimensions();
-  size_t const n_ops = summed_op.n_operators();
-
-  std::vector<std::complex<double>> states_raw;
-  std::mdspan states =
-      fast_pauli::rand<std::complex<double>, 2>(states_raw, {dim, n_states});
-
-  std::vector<std::complex<double>> new_states_raw;
-  std::mdspan new_states = fast_pauli::zeros<std::complex<double>, 2>(
-      new_states_raw, {dim, n_states});
-
-  // Init data (aka data)
-  std::vector<double> data_raw;
-  std::mdspan data = fast_pauli::rand<double, 2>(data_raw, {n_ops, n_states});
-
-  // Apply the summed operator
-  summed_op.apply(new_states, states, data);
-
-  // Check the check
-  std::vector<std::complex<double>> expected_raw;
-  std::mdspan expected =
-      fast_pauli::zeros<std::complex<double>, 2>(expected_raw, {dim, n_states});
-
-  for (size_t j = 0; j < summed_op.n_pauli_strings(); ++j) {
-    auto ps = summed_op.pauli_strings[j];
-
-    PauliOp<double> pop{{std::complex<double>(1.)}, {ps}};
-    std::vector<std::complex<double>> tmp_raw;
-    std::mdspan tmp =
-        fast_pauli::zeros<std::complex<double>, 2>(tmp_raw, {dim, n_states});
-
-    pop.apply(tmp, states);
-
-    // Manually calculate the sum over the different pauli operators
-    // This is specific to the coefficients we've chosen above
-    for (size_t i = 0; i < dim; ++i) {
-      for (size_t t = 0; t < n_states; ++t) {
-        for (size_t k = 0; k < n_ops; ++k) {
-          expected(i, t) += data(k, t) * summed_op.coeffs(j, k) * tmp(i, t);
-        }
-      }
-    }
-  }
-
-  for (size_t t = 0; t < n_states; ++t) {
-    for (size_t i = 0; i < dim; ++i) {
-      CHECK(abs(new_states(i, t) - expected(i, t)) < 1e-6);
-    }
-  }
+  __check_apply(pauli_strings, coeff);
 }
 
 TEST_CASE("apply many operators many PauliString") {
@@ -272,58 +178,5 @@ TEST_CASE("apply many operators many PauliString") {
   std::mdspan coeff = fast_pauli::rand<std::complex<double>, 2>(
       coeff_raw, {pauli_strings.size(), 100});
 
-  SummedPauliOp<double> summed_op{pauli_strings, coeff};
-
-  // Setup states
-  size_t const n_states = 10;
-  size_t const dim = summed_op.n_dimensions();
-  size_t const n_ops = summed_op.n_operators();
-
-  std::vector<std::complex<double>> states_raw;
-  std::mdspan states =
-      fast_pauli::rand<std::complex<double>, 2>(states_raw, {dim, n_states});
-
-  std::vector<std::complex<double>> new_states_raw;
-  std::mdspan new_states = fast_pauli::zeros<std::complex<double>, 2>(
-      new_states_raw, {dim, n_states});
-
-  // Init data
-  std::vector<double> data_raw;
-  std::mdspan data = fast_pauli::rand<double, 2>(data_raw, {n_ops, n_states});
-
-  // Apply the summed operator
-  summed_op.apply_parallel(new_states, states, data);
-
-  // Calculate the expected answer
-  std::vector<std::complex<double>> expected_raw;
-  std::mdspan expected =
-      fast_pauli::zeros<std::complex<double>, 2>(expected_raw, {dim, n_states});
-
-  for (size_t j = 0; j < summed_op.n_pauli_strings(); ++j) {
-    auto ps = summed_op.pauli_strings[j];
-
-    PauliOp<double> pop{{std::complex<double>(1.)}, {ps}};
-    std::vector<std::complex<double>> tmp_raw;
-    std::mdspan tmp =
-        fast_pauli::zeros<std::complex<double>, 2>(tmp_raw, {dim, n_states});
-
-    pop.apply(tmp, states);
-
-    // Manually calculate the sum over the different pauli operators
-    // This is specific to the coefficients we've chosen above
-    for (size_t i = 0; i < dim; ++i) {
-      for (size_t t = 0; t < n_states; ++t) {
-        // expected(i, t) += 1i * tmp(i, t);
-        for (size_t k = 0; k < n_ops; ++k) {
-          expected(i, t) += data(k, t) * summed_op.coeffs(j, k) * tmp(i, t);
-        }
-      }
-    }
-  }
-
-  for (size_t t = 0; t < n_states; ++t) {
-    for (size_t i = 0; i < dim; ++i) {
-      CHECK(abs(new_states(i, t) - expected(i, t)) < 1e-6);
-    }
-  }
+  __check_apply(pauli_strings, coeff);
 }
\ No newline at end of file