Skip to content

Commit

Permalink
Merge pull request #490 from MichaelBroughton/mps_sample
Browse files Browse the repository at this point in the history 
Adds working bitstring Sample function to MPS.
  • Loading branch information
@MichaelBroughton
MichaelBroughton authored Dec 22, 2021
2 parents b945b9d + fc1fcc8 commit 79197f9
Show file tree
Hide file tree
Showing 2 changed files with 397 additions and 1 deletion.
184 changes: 183 additions & 1 deletion lib/mps_statespace.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
#include <cstdlib>
#include <cstring>
#include <memory>
#include <random>

#include "../eigen/Eigen/Dense"
#include "../eigen/unsupported/Eigen/CXX11/Tensor"
Expand Down Expand Up @@ -323,7 +324,7 @@ class MPSStateSpace {
// Merge top into partial_contract2.
new (&top) ConstMatrixMap((Complex*)(scratch_raw + offset), bond_dim,
2 * bond_dim);
// [bd, bd] = [2bd, bd] @ [bd, 2bd]
// [bd, bd] = [bd, 2bd] @ [bd, 2bd]
partial_contract.noalias() = top * partial_contract2.adjoint();
}

Expand Down Expand Up @@ -372,6 +373,187 @@ class MPSStateSpace {
out = t_4d.contract(t_2d, product_dims);
}

// Draw a single bitstring sample from state using scratch and scratch2
// as working space.
static void SampleOnce(MPS& state, MPS& scratch, MPS& scratch2,
std::mt19937* random_gen, std::vector<bool>* sample) {
// TODO: carefully profile with perf and optimize temp storage
// locations for cache friendliness.
const auto bond_dim = state.bond_dim();
const auto num_qubits = state.num_qubits();
const auto end = Size(state);
const auto left_frontier_offset = GetBlockOffset(state, num_qubits + 1);
std::default_random_engine generator;
fp_type* state_raw = state.get();
fp_type* scratch_raw = scratch.get();
fp_type* scratch2_raw = scratch2.get();
fp_type rdm[8];

sample->reserve(num_qubits);
Copy(state, scratch);
Copy(state, scratch2);

// Store prefix contractions in scratch2.
auto offset = GetBlockOffset(state, num_qubits - 1);
ConstMatrixMap top((Complex*)(state_raw + offset), bond_dim, 2);
ConstMatrixMap bot((Complex*)(scratch_raw + offset), bond_dim, 2);
MatrixMap partial_contract((Complex*)(scratch2_raw + offset), bond_dim,
bond_dim);
MatrixMap partial_contract2((Complex*)(scratch_raw + end), bond_dim,
2 * bond_dim);
partial_contract.noalias() = top * bot.adjoint();

for (unsigned i = num_qubits - 2; i > 0; --i) {
offset = GetBlockOffset(state, i);
// reshape:
new (&partial_contract2)
MatrixMap((Complex*)(scratch_raw + end), 2 * bond_dim, bond_dim);

// Merge bot into left boundary merged tensor.
new (&bot) ConstMatrixMap((Complex*)(scratch_raw + offset), 2 * bond_dim,
bond_dim);
partial_contract2.noalias() = bot * partial_contract.adjoint();

// reshape:
new (&partial_contract2)
MatrixMap((Complex*)(scratch_raw + end), bond_dim, 2 * bond_dim);

// Merge top into partial_contract2.
new (&top) ConstMatrixMap((Complex*)(state_raw + offset), bond_dim,
2 * bond_dim);

// merge into partial_contract -> scracth2_raw.
new (&partial_contract)
MatrixMap((Complex*)(scratch2_raw + offset), bond_dim, bond_dim);
partial_contract.noalias() = top * partial_contract2.adjoint();
}

// Compute RDM-0 and draw first sample.
offset = GetBlockOffset(state, 1);
new (&top) ConstMatrixMap((Complex*)state_raw, 2, bond_dim);
new (&bot) ConstMatrixMap((Complex*)scratch_raw, 2, bond_dim);
new (&partial_contract)
MatrixMap((Complex*)(scratch2_raw + offset), bond_dim, bond_dim);
new (&partial_contract2)
MatrixMap((Complex*)(scratch_raw + end), 2, bond_dim);

partial_contract2.noalias() = bot * partial_contract.adjoint();

new (&partial_contract) MatrixMap((Complex*)rdm, 2, 2);
partial_contract.noalias() = top * partial_contract2.adjoint();
auto p0 = rdm[0] / (rdm[0] + rdm[6]);
std::bernoulli_distribution distribution(1 - p0);
auto bit_val = distribution(*random_gen);
sample->push_back(bit_val);

// collapse state.
new (&partial_contract) MatrixMap((Complex*)scratch_raw, 2, bond_dim);
partial_contract.row(!bit_val).setZero();

// Prepare left contraction frontier.
new (&partial_contract2) MatrixMap(
(Complex*)(scratch2_raw + left_frontier_offset), bond_dim, bond_dim);
partial_contract2.noalias() =
partial_contract.transpose() * partial_contract.conjugate();

// Compute RDM-i and draw internal tensor samples.
for (unsigned i = 1; i < num_qubits - 1; i++) {
// Get leftmost [bd, bd] contraction and contract with top.
offset = GetBlockOffset(state, i);
new (&partial_contract) MatrixMap(
(Complex*)(scratch2_raw + left_frontier_offset), bond_dim, bond_dim);
new (&top) ConstMatrixMap((Complex*)(state_raw + offset), bond_dim,
2 * bond_dim);
new (&partial_contract2)
MatrixMap((Complex*)(state_raw + end), bond_dim, 2 * bond_dim);
partial_contract2.noalias() = partial_contract * top.conjugate();

// Contract top again for correct shape.
MatrixMap partial_contract3((Complex*)(scratch_raw + end), 2 * bond_dim,
2 * bond_dim);
partial_contract3.noalias() = top.transpose() * partial_contract2;

// Conduct final tensor contraction operations. Cannot be easily compiled
// to matmul. Perf reports shows only ~6% of runtime spent here on large
// systems.
offset = GetBlockOffset(state, i + 1);
const Eigen::TensorMap<const Eigen::Tensor<Complex, 4, Eigen::RowMajor>>
t_4d((Complex*)(scratch_raw + end), 2, bond_dim, 2, bond_dim);
const Eigen::TensorMap<const Eigen::Tensor<Complex, 2, Eigen::RowMajor>>
t_2d((Complex*)(scratch2_raw + offset), bond_dim, bond_dim);

const Eigen::array<Eigen::IndexPair<int>, 2> product_dims = {
Eigen::IndexPair<int>(1, 0),
Eigen::IndexPair<int>(3, 1),
};
Eigen::TensorMap<Eigen::Tensor<Complex, 2, Eigen::RowMajor>> out(
(Complex*)rdm, 2, 2);
out = t_4d.contract(t_2d, product_dims);

// Sample bit and collapse state.
p0 = rdm[0] / (rdm[0] + rdm[6]);
distribution = std::bernoulli_distribution(1 - p0);
bit_val = distribution(*random_gen);

sample->push_back(bit_val);
offset = GetBlockOffset(state, i);
new (&partial_contract)
MatrixMap((Complex*)(scratch_raw + offset), bond_dim * 2, bond_dim);
for (unsigned j = !bit_val; j < 2 * bond_dim; j += 2) {
partial_contract.row(j).setZero();
}

// Update left frontier.
new (&partial_contract) MatrixMap(
(Complex*)(scratch2_raw + left_frontier_offset), bond_dim, bond_dim);

// reshape:
new (&partial_contract2)
MatrixMap((Complex*)(state_raw + end), bond_dim, 2 * bond_dim);

// Merge bot into left boundary merged tensor.
new (&bot) ConstMatrixMap((Complex*)(scratch_raw + offset), bond_dim,
2 * bond_dim);
partial_contract2.noalias() = partial_contract * bot.conjugate();

// reshape:
new (&partial_contract2)
MatrixMap((Complex*)(state_raw + end), 2 * bond_dim, bond_dim);

// Merge top into partial_contract2.
new (&top) ConstMatrixMap((Complex*)(scratch_raw + offset), 2 * bond_dim,
bond_dim);
partial_contract.noalias() = top.transpose() * partial_contract2;
}

// Compute RDM-(n-1) and sample.
offset = GetBlockOffset(state, num_qubits - 1);
new (&partial_contract2)
MatrixMap((Complex*)(state_raw + end), bond_dim, 2);

new (&top) ConstMatrixMap((Complex*)(state_raw + offset), bond_dim, 2);
partial_contract2.noalias() = partial_contract * top.conjugate();
new (&partial_contract) MatrixMap((Complex*)rdm, 2, 2);
partial_contract.noalias() = top.transpose() * partial_contract2;

p0 = rdm[0] / (rdm[0] + rdm[6]);
distribution = std::bernoulli_distribution(1 - p0);
bit_val = distribution(*random_gen);
sample->push_back(bit_val);
}

// Draw num_samples bitstring samples from state and store the result
// bit vectors in results. Uses scratch and scratch2 as workspace.
static void Sample(MPS& state, MPS& scratch, MPS& scratch2,
unsigned num_samples, unsigned seed,
std::vector<std::vector<bool>>* results) {
std::mt19937 rand_source(seed);
results->reserve(num_samples);
for (unsigned i = 0; i < num_samples; i++) {
SampleOnce(state, scratch, scratch2, &rand_source, &(*results)[i]);
}
}

// Testing only. Convert the MPS to a wavefunction under "normal" ordering.
// Requires: wf be allocated beforehand with bond_dim * 2 ^ num_qubits -1
// memory.
Expand Down
Loading

0 comments on commit 79197f9

Please sign in to comment.