From 1db311db86697bb85f17bcdc9afd501eb0e21e34 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= <benoit.sierro@bluewin.ch>
Date: Fri, 27 Jan 2023 09:17:39 +0100
Subject: [PATCH] Tracked value pickling fix

---
 src/scgenerator/physics/simulate.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py
index a173fd1..2e83658 100644
--- a/src/scgenerator/physics/simulate.py
+++ b/src/scgenerator/physics/simulate.py
@@ -23,13 +23,19 @@ except ModuleNotFoundError:
 
 
 class TrackedValues(defaultdict):
-    def __init__(self):
+    def __init__(self, data=None):
         super().__init__(list)
 
     def append(self, d: dict[str, Any]):
         for k, v in d.items():
             self[k].append(v)
 
+    def __getstate__(self):
+        return self
+
+    def __setstate__(self, state):
+        self.update(state)
+
 
 class RK4IP:
     params: Parameters
@@ -186,7 +192,7 @@ class RK4IP:
 
                 new_tracked_values = integrator.all_values()
                 self.logger.debug(f"tracked values at z={state.z} : {new_tracked_values}")
-                self.tracked_values.append(new_tracked_values)
+                self.tracked_values.append(new_tracked_values | dict(z=state.z))
 
                 # Whether the current spectrum has to be stored depends on previous step
                 if store: