Update

aiida_aurora/utils/analyzers.py

@@ -41,6 +41,7 @@ def __init__(
         check_type="discharge_capacity",
         threshold=0.8,
         consecutive_cycles=2,
+        keep_last=10,
     ) -> None:
         """`CapacityAnalyzer` constructor.
 
@@ -54,6 +55,8 @@ def __init__(
         `consecutive_cycles` : `int`
             The number of required consecutive cycles,
             `2` by default.
+        `keep_last` : `int`
+            The number of cycles to keep in snapshot.
 
         Raises
         ------
@@ -67,6 +70,7 @@ def __init__(
         self.threshold = threshold
         self.consecutive = consecutive_cycles
         self.is_discharge = check_type == "discharge_capacity"
+        self.keep_last = keep_last
 
         self.flag = ""
         self.status = ""
@@ -85,6 +89,7 @@ def analyze(self, snapshot: dict) -> None:
         """
         self._extract_capacities(snapshot)
         self._check_capacity()
+        self._truncate_snapshot()
 
     ###########
     # PRIVATE #
@@ -158,3 +163,21 @@ def _filter_consecutive(self, cycles: list[int]) -> list[int]:
             if i >= self.consecutive - 1 and \
             all(cycles[i - j] == cycle - j for j in range(1, self.consecutive))
         ]
+
+    def _truncate_snapshot(self) -> None:
+        """Truncate the snapshot to user defined size."""
+
+        truncated = {}
+
+        size = min(self.keep_last, len(self.snapshot["cn"]))
+
+        for key, value in self.snapshot.items():
+
+            if key in ("time", "I", "Ewe", "Q"):
+                index = self.snapshot["cycle-index"][-size]
+                truncated[key] = value[index:]
+
+            elif key in ("cn", "Qc", "Qd"):
+                truncated[key] = value[-size:]
+
+        self.snapshot = truncated

aiida_aurora/utils/cycling_analysis.py

@@ -349,10 +349,7 @@ def get_data_from_remote(source: RemoteData) -> dict:
 
 def get_data_from_snapshot(snapshot: dict) -> dict:
     """docstring"""
-    return {
-        k: np.array(v) if isinstance(v, list) else v \
-        for k, v in snapshot.items()
-    }
+    return {k: np.array(v) for k, v in snapshot.items()}
 
 
 def add_analysis(data: dict) -> Styler | str:

aiida_aurora/utils/parsers.py

@@ -8,10 +8,10 @@ def get_data_from_raw(jsdata) -> dict:
     "Extract raw data from json file."
 
     if not isinstance(jsdata, dict):
-        raise TypeError('jsdata should be a dictionary')
+        raise TypeError("jsdata should be a dictionary")
 
     if len(jsdata["steps"]) > 1:
-        raise NotImplementedError('Analysis of multiple steps is not implemented.')
+        raise NotImplementedError("multi-step analysis not implemented.")
 
     raw_data = jsdata["steps"][0]["data"]
 
@@ -27,13 +27,13 @@ def get_data_from_results(array_node) -> dict:
     "Extract data from parsed ArrayData node."
 
     if not isinstance(array_node, ArrayData):
-        raise TypeError('array_node should be an ArrayData')
+        raise TypeError("array_node should be an ArrayData")
 
     # collect data
-    t = array_node.get_array('step0_uts')
+    t = array_node.get_array("step0_uts")
     t -= t[0]
-    Ewe = array_node.get_array('step0_Ewe_n')
-    I = array_node.get_array('step0_I_n')
+    Ewe = array_node.get_array("step0_Ewe_n")
+    I = array_node.get_array("step0_I_n")
 
     return post_process_data(t, Ewe, I)
 
@@ -49,24 +49,28 @@ def post_process_data(t: np.ndarray, Ewe: np.ndarray, I: np.ndarray) -> dict:
     idx = np.append(idx, len(I))
 
     # integrate and store charge and discharge currents
-    Qc, Qd = [], []
+    Qc, Qd, cycle_idx = [], [], []
+
     for ii in range(len(idx) - 1):
+
         i0, ie = idx[ii], idx[ii + 1]
+
         if ie - i0 < 10:
             continue
-        q = np.trapz(I[i0:ie], t[i0:ie])
-        if q > 0:
+
+        if (q := np.trapz(I[i0:ie], t[i0:ie])) > 0:
+            cycle_idx.append(i0)
             Qc.append(q)
         else:
             Qd.append(abs(q))
 
     return {
-        'time': t,
-        'Ewe': Ewe,
-        'I': I,
-        'cn': len(Qd),
-        'time-cycles': t[idx[2::2]],
-        'Q': cumtrapz(I, t, axis=0, initial=0) / 3.6,
-        'Qd': np.array(Qd) / 3.6,
-        'Qc': np.array(Qc) / 3.6,
+        "time": t,
+        "Ewe": Ewe,
+        "I": I,
+        "cycle-index": cycle_idx,
+        "cn": np.arange(len(Qd)),
+        "Q": cumtrapz(I, t, axis=0, initial=0) / 3.6,
+        "Qd": np.array(Qd) / 3.6,
+        "Qc": np.array(Qc) / 3.6,
     }