From 664b2d7e380d11571b85f41d955c216344531ec9 Mon Sep 17 00:00:00 2001
From: Remi Gau <remi_gau@hotmail.com>
Date: Mon, 19 Feb 2024 09:54:26 +0100
Subject: [PATCH] extract loop

---
 neurodesign/classes.py | 37 +++++++++++++------------------------
 1 file changed, 13 insertions(+), 24 deletions(-)

diff --git a/neurodesign/classes.py b/neurodesign/classes.py
index cc7c402..4c286e5 100644
--- a/neurodesign/classes.py
+++ b/neurodesign/classes.py
@@ -996,6 +996,17 @@ def clear(self):
 
         return self
 
+    def _run_optimisation_loop(self, weights):
+        with progress_bar(text="Optimizing") as progress:
+            task = progress.add_task(
+                description="optimize", total=len(range(self.preruncycles))
+            )
+            for _ in range(self.preruncycles):
+                self.to_next_generation(seed=self.seed, weights=weights)
+                progress.update(task, advance=1)
+                if self.finished:
+                    continue
+
     def optimise(self):
         """Run design optimization."""
         if self.exp.FcMax == 1 and self.exp.FfMax == 1:
@@ -1005,35 +1016,13 @@ def optimise(self):
             # add new designs
             self.clear()
             self.add_new_designs(weights=[1, 0, 0, 0])
-
-            # loop
-            with progress_bar(text="Optimizing") as progress:
-                task = progress.add_task(
-                    description="optimize", total=len(range(self.preruncycles))
-                )
-                for _ in range(self.preruncycles):
-                    self.to_next_generation(seed=self.seed, weights=[1, 0, 0, 0])
-                    progress.update(task, advance=1)
-                    if self.finished:
-                        continue
-
+            self._run_optimisation_loop(self, weights=[1, 0, 0, 0])
             self.exp.FeMax = np.max(self.bestdesign.F)
 
         if self.exp.FdMax == 1 and self.weights[1] > 0:
             self.clear()
             self.add_new_designs(weights=[0, 1, 0, 0])
-
-            # loop
-            with progress_bar(text="Optimizing") as progress:
-                task = progress.add_task(
-                    description="optimize", total=len(range(self.preruncycles))
-                )
-                for _ in range(self.preruncycles):
-                    self.to_next_generation(seed=self.seed, weights=[0, 1, 0, 0])
-                    progress.update(task, advance=1)
-                    if self.finished:
-                        continue
-
+            self._run_optimisation_loop(self, weights=[0, 1, 0, 0])
             self.exp.FdMax = np.max(self.bestdesign.F)
 
         # clear all attributes