feat: Add n_samples parameter to state classifier methods for customi…

…zable plotting
amachino · Dec 28, 2024 · 5fe1c22 · 5fe1c22
1 parent 3b10f84
commit 5fe1c22
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 0 deletions.
diff --git a/src/qubex/measurement/state_classifier.py b/src/qubex/measurement/state_classifier.py
@@ -133,6 +133,7 @@ def plot(
         target: str,
         data: NDArray[np.complex128],
         labels: NDArray,
+        n_samples: int = 1000,
     ):
         """
         Plot the data and the predicted labels.
@@ -143,6 +144,8 @@ def plot(
             An array of complex numbers representing the data.
         labels : NDArray
             An array of predicted state labels.
+        n_samples : int, optional
+            The number of samples to plot, by default 1000.
         """
         raise NotImplementedError
 

diff --git a/src/qubex/measurement/state_classifier_gmm.py b/src/qubex/measurement/state_classifier_gmm.py
@@ -291,6 +291,7 @@ def plot(
         target: str,
         data: NDArray[np.complex128],
         labels: NDArray,
+        n_samples: int = 1000,
     ):
         """
         Plot the data and the predicted labels.
@@ -301,7 +302,12 @@ def plot(
             An array of complex numbers representing the data.
         labels : NDArray
             An array of predicted state labels.
+        n_samples : int, optional
+            The number of samples to plot, by default 1000.
         """
+        if len(data) > n_samples:
+            data = data[:n_samples]
+            labels = labels[:n_samples]
         x = data.real
         y = data.imag
         unique_labels = np.unique(labels)

diff --git a/src/qubex/measurement/state_classifier_kmeans.py b/src/qubex/measurement/state_classifier_kmeans.py
@@ -256,6 +256,7 @@ def plot(
         target: str,
         data: NDArray[np.complex128],
         labels: NDArray,
+        n_samples: int = 1000,
     ):
         """
         Plot the data and the predicted labels.
@@ -266,7 +267,12 @@ def plot(
             An array of complex numbers representing the data.
         labels : NDArray
             An array of predicted state labels.
+        n_samples : int, optional
+            The number of samples to plot, by default 1000.
         """
+        if len(data) > n_samples:
+            data = data[:n_samples]
+            labels = labels[:n_samples]
         x = data.real
         y = data.imag
         unique_labels = np.unique(labels)