Add memory-based thread calculation and --max-memory argument

docs/Writerside/topics/Command-Line-Usage.topic

@@ -91,6 +91,11 @@
                 <p>Use GPU acceleration. (Linux only)</p>
                 <b>Advanced</b>
             </def>
+            <def title="--max-memory">
+                <p>Set the number of threads based on the maximum amount of memory you would like NucleoFind to use. This option is
+                only active if `nthreads` is not specified.</p>
+                <warning>This feature does not explicitly limit memory usage. Use with caution.</warning>
+            </def>
 <!--            <def title="-model_path">-->
 <!--                <p>Path to ONNX model for use with a custom model.</p>-->
 <!--                <b>Advanced</b>-->

package/src/nucleofind/prediction/arguments.py

@@ -28,7 +28,7 @@ def parse_arguments() -> SimpleNamespace:
         "-n",
         "--nthreads",
         nargs="?",
-        default=1,
+        default=0,
         type=int,
         help="Number of threads to use",
     )
@@ -62,6 +62,12 @@ def parse_arguments() -> SimpleNamespace:
     parser.add_argument(
         "--gpu", action=argparse.BooleanOptionalAction, help="Use GPU (experimental)"
     )
+    parser.add_argument(
+        "--max-memory",
+        nargs="?",
+        help="Maximum memory to use in Gigabytes (experimental). This option is only active when nthreads is not specified.",
+        type=float,
+    )
     parser.add_argument(
         "--debug", action=argparse.BooleanOptionalAction, help="Turn on debug logging"
     )

package/src/nucleofind/prediction/load.py

@@ -63,7 +63,9 @@ def load_onnx_model(
     if use_gpu:
         providers.insert(0, "CUDAExecutionProvider")
     sess_options = rt.SessionOptions()
-    sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+    sess_options.graph_optimization_level = (
+        rt.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+    )
 
     try:
         return rt.InferenceSession(

package/src/nucleofind/prediction/memory.py

@@ -0,0 +1,47 @@
+import logging
+
+
+def calculate_n_from_max_memory(max_memory: float):
+    """
+    Calculates the number of processing cores (n) based on the provided maximum memory
+    constraint and a heuristic mapping of memory per core.
+
+    This function uses a predefined heuristic mapping of memory per core to determine
+    the number of cores that can fit within the specified maximum memory. The input
+    memory value is expected to be a string representing the maximum available memory.
+
+    Parameters:
+        max_memory (str): The maximum allowable memory in string format.
+
+    Returns:
+        int: The calculated number of cores based on the memory constraint.
+    """
+    if not max_memory:
+        return 1
+
+    try:
+        max_memory = float(max_memory)
+    except ValueError:
+        logging.error(
+            f"Invalid max_memory value: {max_memory}. Please provide a valid number in GB, ignoring for now."
+        )
+        return 1
+
+    # Check to ensure the max_memory isn't too big (i.e. passed in MB rather than GB)
+    if max_memory > 1024:
+        max_memory = max_memory // 1024
+        logging.warning(
+            f"NucleoFind has converted the max_memory argument to GB. Max Memory = {max_memory} GB"
+        )
+
+    if max_memory < 1:
+        logging.error(
+            f"Invalid max_memory value: {max_memory}. Please provide a valid number in GB, ignoring for now."
+        )
+        return 1
+
+    if max_memory < 4 :
+        return 1
+
+    max_cpus = (max_memory - 1.47) // 2.37
+    return max_cpus

package/src/nucleofind/prediction/predict.py

@@ -13,6 +13,7 @@
 from .save import save_grid
 from ..logs import setup_logging
 from .config import Configuration, MapType
+from .memory import calculate_n_from_max_memory
 
 
 class NucleoFind:
@@ -188,6 +189,17 @@ def run():
         n_threads=args.nthreads,
         **vars(model_configuration),
     )
+
+    # Set number of threads
+    # Default to 1 thread if no nthreads specified and no max_memory specified
+    # If max_memory specified, calculate nthreads based on max_memory
+    if args.nthreads == 0 and args.max_memory:
+        configuration.n_threads = calculate_n_from_max_memory(args.max_memory)
+    elif args.nthreads == 0 and not args.max_memory:
+        configuration.n_threads = 1
+    else:
+        configuration.n_threads = args.nthreads
+
     nucleofind = NucleoFind(model_path, configuration)
     nucleofind.predict(
         args.input,