refactor(model-input): use assign instead of [] for time column

#45
kausalflow · Aug 12, 2024 · 4dd093c · 4dd093c
1 parent f70b6ac
commit 4dd093c
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/hamilflow/models/brownian_motion.py b/hamilflow/models/brownian_motion.py
@@ -206,8 +206,6 @@ def __call__(self, t: TypeTime, seed: int = 42) -> pd.DataFrame:
         n_steps = np.array(t).size
         trajectory = self._trajectory(n_new_steps=n_steps - 1, seed=seed)
 
-        df = pd.DataFrame(trajectory, columns=self._axis_names)
-
-        df["t"] = t
+        df = pd.DataFrame(trajectory, columns=self._axis_names).assign(t=t)
 
         return df