clean up, fix issues

desc/__init__.py

@@ -2,10 +2,13 @@
 
 import importlib
 import os
+import platform
 import re
+import subprocess
 import warnings
 
 import colorama
+import psutil
 from termcolor import colored
 
 from ._version import get_versions
@@ -61,6 +64,23 @@ def __getattr__(name):
 config = {"device": None, "avail_mem": None, "kind": None, "num_device": None}
 
 
+def _get_processor_name():
+    """Get the processor name of the current system."""
+    if platform.system() == "Windows":
+        return platform.processor()
+    elif platform.system() == "Darwin":
+        os.environ["PATH"] = os.environ["PATH"] + os.pathsep + "/usr/sbin"
+        command = "sysctl -n machdep.cpu.brand_string"
+        return subprocess.check_output(command).strip()
+    elif platform.system() == "Linux":
+        command = "cat /proc/cpuinfo"
+        all_info = subprocess.check_output(command, shell=True).decode().strip()
+        for line in all_info.split("\n"):
+            if "model name" in line:
+                return re.sub(".*model name.*:", "", line, 1)
+    return ""
+
+
 def set_device(kind="cpu", gpuid=None, num_device=1):
     """Sets the device to use for computation.
 
@@ -85,15 +105,21 @@ def set_device(kind="cpu", gpuid=None, num_device=1):
         number of devices to use. Default is 1.
 
     """
+    if kind == "cpu" and num_device > 1:
+        # TODO: implement multi-CPU support
+        raise ValueError("Cannot request multiple CPUs")
+
     config["kind"] = kind
     config["num_device"] = num_device
+
+    cpu_mem = psutil.virtual_memory().available / 1024**3  # RAM in GB
+    cpu_info = _get_processor_name()
+    config["cpu_info"] = f"{cpu_info} CPU"
+    config["cpu_mem"] = cpu_mem
     if kind == "cpu":
         os.environ["JAX_PLATFORMS"] = "cpu"
         os.environ["CUDA_VISIBLE_DEVICES"] = ""
-        import psutil
-
-        cpu_mem = psutil.virtual_memory().available / 1024**3  # RAM in GB
-        config["devices"] = ["CPU"]
+        config["devices"] = [f"{cpu_info} CPU"]
         config["avail_mems"] = [cpu_mem]
 
     elif kind == "gpu":

desc/backend.py

@@ -69,12 +69,17 @@
             )
         )
 
-print(f"Using {desc_config['num_device']} device:")
-for i, dev in enumerate(desc_config["devices"]):
-    print(
-        f"\t Device {i}: {dev} with {desc_config['avail_mems'][i]:.2f} "
-        "GB available memory"
-    )
+print(
+    f"CPU Info: {desc_config['cpu_info']} with {desc_config['cpu_mem']:.2f} "
+    "GB available memory"
+)
+if desc_config["kind"] == "gpu":
+    print(f"Using {desc_config['num_device']} device:")
+    for i, dev in enumerate(desc_config["devices"]):
+        print(
+            f"\t Device {i}: {dev} with {desc_config['avail_mems'][i]:.2f} "
+            "GB available memory"
+        )
 
 if use_jax:  # noqa: C901
     from jax import custom_jvp, jit, vmap

desc/objectives/objective_funs.py

@@ -310,6 +310,7 @@ def build(self, use_jit=None, verbose=1):  # noqa:  C901
             Level of output.
 
         """
+        use_jit_wrapper = True
         if use_jit is not None:
             self._use_jit = use_jit
             # use_jit_wrapper is used to determine if we jit the ObjectiveFunction
@@ -373,24 +374,23 @@ def build(self, use_jit=None, verbose=1):  # noqa:  C901
         if self._jac_chunk_size == "auto":
             # Heuristic estimates of fwd mode Jacobian memory usage,
             # slightly conservative, based on using ForceBalance as the objective
-            if self._deriv_mode == "batched":
-                estimated_memory_usage = 2.4e-7 * self.dim_f * self.dim_x + 1  # in GB
-                mem_avail = desc_config["avail_mems"][0]  # in GB
-                max_chunk_size = round(
-                    (mem_avail / estimated_memory_usage - 0.22) / 0.85 * self.dim_x
-                )
+            estimated_memory_usage = 2.4e-7 * self.dim_f * self.dim_x + 1  # in GB
+            mem_avail = desc_config["avail_mems"][0]  # in GB
+            max_chunk_size = round(
+                (mem_avail / estimated_memory_usage - 0.22) / 0.85 * self.dim_x
+            )
             self._jac_chunk_size = max([1, max_chunk_size])
             if self._deriv_mode == "blocked":
                 for obj in self.objectives:
                     if obj._jac_chunk_size is None:
                         estimated_memory_usage = (
-                            2.4e-7 * obj.dim_f * obj.dim_x + 1
+                            2.4e-7 * obj.dim_f * obj.things[0].dim_x + 1
                         )  # in GB
                         mem_avail = desc_config["avail_mems"][obj._device_id]  # in GB
                         max_chunk_size = round(
                             (mem_avail / estimated_memory_usage - 0.22)
                             / 0.85
-                            * obj.dim_x
+                            * obj.things[0].dim_x
                         )
                         obj._jac_chunk_size = max([1, max_chunk_size])