make warning a float instead of bool

sblunt · Jan 16, 2024 · c8facdf · c8facdf
1 parent 0b5f85b
commit c8facdf
Showing 1 changed file with 14 additions and 14 deletions.
diff --git a/orbitize/sampler.py b/orbitize/sampler.py
@@ -526,7 +526,7 @@ def _sampler_process(
         return (np.array(output_orbits), output_lnlikes)
 
     def run_sampler(
-        self, total_orbits, num_samples=10000, num_cores=None, OFTI_warning=True
+        self, total_orbits, num_samples=10000, num_cores=None, OFTI_warning=60.0
     ):
         """
         Runs OFTI in parallel on multiple cores until we get the number of total accepted orbits we want.
@@ -537,17 +537,16 @@ def run_sampler(
                 rejection sampling on. Defaults to 10000.
             num_cores (int): the number of cores to run OFTI on. Defaults to
                 number of cores availabe.
-            OFTI_warning (bool): if True, print a warning if OFTI doesn't accept
-                any orbits after 60 s.
+            OFTI_warning (float): if OFTI doesn't accept a single orbit before
+                this amount of time (in seconds), print a warning suggesting to
+                try MCMC. If None, don't print a warning.
         Return:
             np.array: array of accepted orbits. Size: total_orbits.
 
         Written by: Vighnesh Nagpal(2019)
 
         """
 
-        # a flag to print out a warning if no OFTI orbits are accepted in x seconds
-        time_warning = 60.0  # seconds
         start_time = time.time()
 
         if num_cores != 1:
@@ -579,14 +578,14 @@ def run_sampler(
 
             # print out the number of orbits generated every second
             while orbits_saved.value < total_orbits:
-                if OFTI_warning and orbits_saved.value == 0:
+                if OFTI_warning is not None and orbits_saved.value == 0:
                     check_time = time.time() - start_time
-                    if check_time > time_warning:
+                    if check_time > OFTI_warning:
                         print(
                             "Warning! OFTI is taking a while, and you may want to consider MCMC, check out the MCMC vs OFTI tutorial.",
                             end="\n",
                         )
-                        OFTI_warning = False
+                        OFTI_warning = None
                 print(
                     str(orbits_saved.value) + "/" + str(total_orbits) + " orbits found",
                     end="\r",
@@ -631,12 +630,13 @@ def run_sampler(
 
                 if len(accepted_orbits) == 0:
                     check_time = time.time() - start_time
-                    if OFTI_warning and check_time > time_warning:
-                        print(
-                            "Warning! OFTI is taking a while, and you may want to consider MCMC, check out the MCMC vs OFTI tutorial.",
-                            end="\n",
-                        )
-                        OFTI_warning = False
+                    if OFTI_warning is not None:
+                        if check_time > OFTI_warning:
+                            print(
+                                "Warning! OFTI is taking a while, and you may want to consider MCMC, check out the MCMC vs OFTI tutorial.",
+                                end="\n",
+                            )
+                            OFTI_warning = None
                     pass
                 else:
                     n_accepted = len(accepted_orbits)