def num_threads(n=None): '''Set the number of OMP threads. If argument is not specified, the function will return the total number of available OMP threads. It's recommended to call this function to set OMP threads than "os.environ['OMP_NUM_THREADS'] = int(n)". This is because environment variables like OMP_NUM_THREADS were read when a module was imported. They cannot be reset through os.environ after the module was loaded. Examples: >>> from pyscf import lib >>> print(lib.num_threads()) 8 >>> lib.num_threads(4) 4 >>> print(lib.num_threads()) 4 ''' from pyscf.lib.numpy_helper import _np_helper if n is not None: _np_helper.set_omp_threads.restype = ctypes.c_int threads = _np_helper.set_omp_threads(ctypes.c_int(int(n))) if threads == 0: warnings.warn('OpenMP is not available. ' 'Setting omp_threads to %s has no effects.' % n) return threads else: _np_helper.get_omp_threads.restype = ctypes.c_int return _np_helper.get_omp_threads()
def num_threads(n=None): '''Set the number of OMP threads. If argument is not specified, the function will return the total number of available OMP threads. Examples: >>> from pyscf import lib >>> print(lib.num_threads()) 8 >>> lib.num_threads(4) 4 >>> print(lib.num_threads()) 4 ''' from pyscf.lib.numpy_helper import _np_helper if n is not None: _np_helper.set_omp_threads.restype = ctypes.c_int threads = _np_helper.set_omp_threads(ctypes.c_int(int(n))) if threads == 0: warnings.warn('OpenMP is not available. ' 'Setting omp_threads to %s has no effects.' % n) return threads else: _np_helper.get_omp_threads.restype = ctypes.c_int return _np_helper.get_omp_threads()