basis add sort_states

Author: johannes-moegerle

src/rydstate/basis/basis_base.py

@@ -57,6 +57,17 @@ def filter_states(self, qn: str, value: float | tuple[float, float], *, delta: f
 
         return self
 
+    def sort_states(self, *qns: str) -> Self:
+        """Sort the basis states according to the given quantum numbers.
+
+        The first quantum number given is the primary sorting key, the second quantum number
+        is the secondary sorting key, and so on.
+        """
+        values = np.array([self.calc_exp_qn(qn) for qn in qns])
+        sorted_indices = np.lexsort(values[::-1])
+        self.states = [self.states[i] for i in sorted_indices]
+        return self
+
     def calc_exp_qn(self, qn: str) -> list[float]:
         if is_angular_momentum_quantum_number(qn):
             return [state.angular.calc_exp_qn(qn) for state in self.states]