From b192d38331d25715c59ff26ba6e2a82aa57e7698 Mon Sep 17 00:00:00 2001
From: Mathias Soeken <mathias.soeken@outlook.com>
Date: Mon, 30 Sep 2024 11:58:04 -0700
Subject: [PATCH] Control how physical qubits are computed in factories (#1940)

This adds the possibility to change the behavior, how physical qubits
are computed in a round based factory. The default behavior is to take
the max among all distillation rounds, but one can set to calculate the
sum over all units instead. This change is accessible via the job
params.
---
 resource_estimator/src/estimates.rs         |  2 +-
 resource_estimator/src/estimates/factory.rs | 42 ++++++++++++++++++---
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/resource_estimator/src/estimates.rs b/resource_estimator/src/estimates.rs
index 6d9f93b985..4c33f22483 100644
--- a/resource_estimator/src/estimates.rs
+++ b/resource_estimator/src/estimates.rs
@@ -10,7 +10,7 @@ pub use error_correction::{CodeWithThresholdAndDistance, CodeWithThresholdAndDis
 mod factory;
 pub use factory::{
     BuilderDispatch2, DistillationRound, DistillationUnit, FactoryBuildError, FactoryDispatch2,
-    NoFactories, RoundBasedFactory,
+    NoFactories, PhysicalQubitCalculation, RoundBasedFactory,
 };
 mod physical_estimation;
 pub use physical_estimation::{
diff --git a/resource_estimator/src/estimates/factory.rs b/resource_estimator/src/estimates/factory.rs
index 7affcbde85..13560f18ad 100644
--- a/resource_estimator/src/estimates/factory.rs
+++ b/resource_estimator/src/estimates/factory.rs
@@ -142,6 +142,7 @@ pub struct RoundBasedFactory<P> {
     rounds: Vec<DistillationRound<P>>,
     input_error_rate_before_each_round: Vec<f64>,
     failure_probability_after_each_round: Vec<f64>,
+    physical_qubit_calculation: PhysicalQubitCalculation,
 }
 
 impl<P: Clone> RoundBasedFactory<P> {
@@ -159,6 +160,7 @@ impl<P: Clone> RoundBasedFactory<P> {
             rounds,
             input_error_rate_before_each_round,
             failure_probability_after_each_round,
+            physical_qubit_calculation: PhysicalQubitCalculation::default(),
         }
     }
 
@@ -178,6 +180,7 @@ impl<P: Clone> RoundBasedFactory<P> {
             rounds,
             input_error_rate_before_each_round,
             failure_probability_after_each_round,
+            physical_qubit_calculation: PhysicalQubitCalculation::default(),
         };
 
         pipeline.compute_units_per_round(units, 1)?;
@@ -207,6 +210,18 @@ impl<P: Clone> RoundBasedFactory<P> {
         Ok(())
     }
 
+    #[must_use]
+    pub fn physical_qubit_calculation(&self) -> PhysicalQubitCalculation {
+        self.physical_qubit_calculation
+    }
+
+    pub fn set_physical_qubit_calculation(
+        &mut self,
+        physical_qubit_calculation: PhysicalQubitCalculation,
+    ) {
+        self.physical_qubit_calculation = physical_qubit_calculation;
+    }
+
     #[must_use]
     pub fn rounds(&self) -> &[DistillationRound<P>] {
         &self.rounds
@@ -321,11 +336,19 @@ impl<P: Clone> Factory for RoundBasedFactory<P> {
     type Parameter = P;
 
     fn physical_qubits(&self) -> u64 {
-        self.rounds
-            .iter()
-            .map(DistillationRound::physical_qubits)
-            .max()
-            .unwrap_or(0)
+        match self.physical_qubit_calculation {
+            PhysicalQubitCalculation::Max => self
+                .rounds
+                .iter()
+                .map(DistillationRound::physical_qubits)
+                .max()
+                .unwrap_or(0),
+            PhysicalQubitCalculation::Sum => self
+                .rounds
+                .iter()
+                .map(DistillationRound::physical_qubits)
+                .sum::<u64>(),
+        }
     }
 
     fn duration(&self) -> u64 {
@@ -350,3 +373,12 @@ impl<P: Clone> Factory for RoundBasedFactory<P> {
             .map(|f| Cow::Borrowed(f))
     }
 }
+
+#[derive(Copy, Clone, Debug, Default)]
+pub enum PhysicalQubitCalculation {
+    /// physical qubits can be shared among rounds
+    #[default]
+    Max,
+    /// each round has its own physical qubits
+    Sum,
+}