content: """
  This module provides functions for parallel execution of loops and function
  mappings, using the builtin Python module multiprocessing or the loky parallel
  execution library.
  """
  __all__ = ['parallel_map', 'serial_map', 'loky_pmap', 'mpi_pmap']
  
  import multiprocessing
  import os
  import sys
  import time
  import threading
  import concurrent.futures
  import warnings
  from qutip.ui.progressbar import progress_bars
  from qutip.settings import available_cpu_count
  
  if sys.platform == 'darwin':
      mp_context = multiprocessing.get_context('fork')
  elif sys.platform == 'linux':
      # forkserver would handle threads better, but is much slower at starting
      # the executor and spawning tasks
      mp_context = multiprocessing.get_context('fork')
  else:
      mp_context = multiprocessing.get_context()
  
  
  default_map_kw = {
      'timeout': threading.TIMEOUT_MAX,
      'num_cpus': available_cpu_count(),
      'fail_fast': True,
  }
  
  
  def _read_map_kw(options):
      options = options or {}
      map_kw = default_map_kw.copy()
      map_kw.update({k: v for k, v in options.items() if v is not None})
      return map_kw
  
  
  class MapExceptions(Exception):
      def __init__(self, msg, errors, results):
          super().__init__(msg, errors, results)
          self.errors = errors
          self.results = results
  
  
  def serial_map(task, values, task_args=None, task_kwargs=None,
                 reduce_func=None, map_kw=None,
                 progress_bar=None, progress_bar_kwargs={}):
      """
      Serial mapping function with the same call signature as parallel_map, for
      easy switching between serial and parallel execution. This
      is functionally equivalent to::
  
          result = [task(value, *task_args, **task_kwargs) for value in values]
  
      This function work as a drop-in replacement of :func:`parallel_map`.
  
      Parameters
      ----------
      task : a Python function
          The function that is to be called for each value in ``task_vec``.
      values : array / list
          The list or array of values for which the ``task`` function is to be
          evaluated.
      task_args : list, optional
          The optional additional argument to the ``task`` function.
      task_kwargs : dictionary, optional
          The optional additional keyword argument to the ``task`` function.
      reduce_func : func, optional
          If provided, it will be called with the output of each tasks instead of
          storing a them in a list. It should return None or a number.
          When returning a number, it represent the estimation of the number of
          task left. On a return <= 0, the map will end early.
      progress_bar : str, optional
          Progress bar options's string for showing progress.
      progress_bar_kwargs : dict, optional
          Options for the progress bar.
      map_kw: dict, optional
          Dictionary containing:
          - timeout: float, Maximum time (sec) for the whole map.
          - fail_fast: bool, Raise an error at the first.
  
      Returns
      -------
      result : list
          The result list contains the value of
          ``task(value, *task_args, **task_kwargs)`` for each
          value in ``values``. If a ``reduce_func`` is provided, and empty list
          will be returned.
  
      """
      if task_args is None:
          task_args = ()
      if task_kwargs is None:
          task_kwargs = {}
      map_kw = _read_map_kw(map_kw)
      remaining_ntraj = None
      progress_bar = progress_bars[progress_bar](
          len(values), **progress_bar_kwargs
      )
      end_time = map_kw['timeout'] + time.time()
      results = None
      if reduce_func is None:
          results = [None] * len(values)
      errors = {}
      for n, value in enumerate(values):
          if time.time() > end_time:
              break
          progress_bar.update()
          try:
              result = task(value, *task_args, **task_kwargs)
          except Exception as err:
              if map_kw["fail_fast"]:
                  raise err
              else:
                  errors[n] = err
          else:
              if reduce_func is not None:
                  remaining_ntraj = reduce_func(result)
              else:
                  results[n] = result
          if remaining_ntraj is not None and remaining_ntraj <= 0:
              end_time = 0
      progress_bar.finished()
  
      if errors:
          raise MapExceptions(f"{len(errors)} iterations failed in serial_map",
                              errors, results)
      return results
  
  
  def _generic_pmap(task, values, task_args, task_kwargs, reduce_func,
                    timeout, fail_fast, num_workers,
                    progress_bar, progress_bar_kwargs,
                    setup_executor, extract_result, shutdown_executor):
      """
      Common functionality for parallel_map, loky_pmap and mpi_pmap.
      The parameters `setup_executor`, `extract_result` and `shutdown_executor`
      are callback functions with the following signatures:
  
      setup_executor: () -> ProcessPoolExecutor
  
      extract_result: Future -> (Any, BaseException)
          If there was an exception e, returns (None, e).
          Otherwise returns (result, None).
  
      shutdown_executor: (executor: ProcessPoolExecutor,
                          active_tasks: set[Future]) -> None
          executor: The ProcessPoolExecutor that was created in setup_executor
          active_tasks: A set of Futures that are currently still being executed
              (non-empty if: timeout, error, or reduce_func requesting exit)
      """
  
      if task_args is None:
          task_args = ()
      if task_kwargs is None:
          task_kwargs = {}
      end_time = timeout + time.time()
  
      progress_bar = progress_bars[progress_bar](
          len(values), **progress_bar_kwargs
      )
  
      errors = {}
      finished = []
      if reduce_func is not None:
          results = None
  
          def result_func(_, value):
              return reduce_func(value)
      else:
          results = [None] * len(values)
          result_func = results.__setitem__
  
      def _done_callback(future):
          if not future.cancelled():
              result, exception = extract_result(future)
              if isinstance(exception, KeyboardInterrupt):
                  # When a keyboard interrupt happens, it is raised in the main
                  # thread and in all worker threads. At this point in the code,
                  # the worker threads have already returned and the main thread
                  # is only waiting for the ProcessPoolExecutor to shutdown
                  # before exiting. We therefore return immediately.
                  return
              if exception is not None:
                  if isinstance(exception, Exception):
                      errors[future._i] = exception
                  else:
                      raise exception
              else:
                  remaining_ntraj = result_func(future._i, result)
                  if remaining_ntraj is not None and remaining_ntraj <= 0:
                      finished.append(True)
          progress_bar.update()
  
      os.environ['QUTIP_IN_PARALLEL'] = 'TRUE'
      try:
          with setup_executor() as executor:
              waiting = set()
              i = 0
              aborted = False
  
              while i < len(values):
                  # feed values to the executor, ensuring that there is at
                  # most one task per worker at any moment in time so that
                  # we can shutdown without waiting for greater than the time
                  # taken by the longest task
                  if len(waiting) >= num_workers:
                      # no space left, wait for a task to complete or
                      # the time to run out
                      timeout = max(0, end_time - time.time())
                      _done, waiting = concurrent.futures.wait(
                          waiting,
                          timeout=timeout,
                          return_when=concurrent.futures.FIRST_COMPLETED,
                      )
                  if (
                      time.time() >= end_time
                      or (errors and fail_fast)
                      or finished
                  ):
                      # no time left, exit the loop
                      aborted = True
                      break
                  while len(waiting) < num_workers and i < len(values):
                      # space and time available, add tasks
                      value = values[i]
                      future = executor.submit(
                          task, *((value,) + task_args), **task_kwargs,
                      )
                      # small hack to avoid add_done_callback not supporting
                      # extra arguments and closures inside loops retaining
                      # a reference not a value:
                      future._i = i
                      future.add_done_callback(_done_callback)
                      waiting.add(future)
                      i += 1
  
              if not aborted:
                  # all tasks have been submitted, timeout has not been reaches
                  # -> wait for all workers to finish before shutting down
                  timeout = max(0, end_time - time.time())
                  _done, waiting = concurrent.futures.wait(
                      waiting,
                      timeout=timeout,
                      return_when=concurrent.futures.ALL_COMPLETED
                  )
              shutdown_executor(executor, waiting)
      finally:
          os.environ['QUTIP_IN_PARALLEL'] = 'FALSE'
  
      progress_bar.finished()
      if errors and fail_fast:
          raise list(errors.values())[0]
      elif errors:
          raise MapExceptions(
              f"{len(errors)} iterations failed in parallel_map",
              errors, results
          )
  
      return results
  
  
  def parallel_map(task, values, task_args=None, task_kwargs=None,
                   reduce_func=None, map_kw=None,
                   progress_bar=None, progress_bar_kwargs={}):
      """
      Parallel execution of a mapping of ``values`` to the function ``task``.
      This is functionally equivalent to::
  
          result = [task(value, *task_args, **task_kwargs) for value in values]
  
      Parameters
      ----------
      task : a Python function
          The function that is to be called for each value in ``task_vec``.
      values : array / list
          The list or array of values for which the ``task`` function is to be
          evaluated.
      task_args : list, optional
          The optional additional arguments to the ``task`` function.
      task_kwargs : dictionary, optional
          The optional additional keyword arguments to the ``task`` function.
      reduce_func : func, optional
          If provided, it will be called with the output of each task instead of
          storing them in a list. Note that the order in which results are
          passed to ``reduce_func`` is not defined. It should return None or a
          number. When returning a number, it represents the estimation of the
          number of tasks left. On a return <= 0, the map will end early.
      progress_bar : str, optional
          Progress bar options's string for showing progress.
      progress_bar_kwargs : dict, optional
          Options for the progress bar.
      map_kw: dict, optional
          Dictionary containing entry for:
          - timeout: float, Maximum time (sec) for the whole map.
          - num_cpus: int, Number of jobs to run at once.
          - fail_fast: bool, Abort at the first error.
  
      Returns
      -------
      result : list
          The result list contains the value of
          ``task(value, *task_args, **task_kwargs)`` for
          each value in ``values``. If a ``reduce_func`` is provided, and empty
          list will be returned.
  
      """
  
      map_kw = _read_map_kw(map_kw)
      if sys.version_info >= (3, 7):
          # ProcessPoolExecutor only supports mp_context from 3.7 onwards
          ctx_kw = {"mp_context": mp_context}
      else:
          ctx_kw = {}
  
      def setup_executor():
          return concurrent.futures.ProcessPoolExecutor(
              max_workers=map_kw['num_cpus'], **ctx_kw,
          )
  
      def extract_result(future: concurrent.futures.Future):
          exception = future.exception()
          if exception is not None:
              return None, exception
          return future.result(), None
  
      def shutdown_executor(executor, _):
          # Since `ProcessPoolExecutor` leaves no other option,
          # we wait for all worker processes to finish their current task
          executor.shutdown()
  
      return _generic_pmap(
          task, values, task_args, task_kwargs, reduce_func,
          map_kw['timeout'], map_kw['fail_fast'], map_kw['num_cpus'],
          progress_bar, progress_bar_kwargs,
          setup_executor, extract_result, shutdown_executor
      )
  
  
  def loky_pmap(task, values, task_args=None, task_kwargs=None,
                reduce_func=None, map_kw=None,
                progress_bar=None, progress_bar_kwargs={}):
      """
      Parallel execution of a mapping of ``values`` to the function ``task``.
      This is functionally equivalent to::
  
          result = [task(value, *task_args, **task_kwargs) for value in values]
  
      Use the loky module instead of multiprocessing.
  
      Parameters
      ----------
      task : a Python function
          The function that is to be called for each value in ``task_vec``.
      values : array / list
          The list or array of values for which the ``task`` function is to be
          evaluated.
      task_args : list, optional
          The optional additional arguments to the ``task`` function.
      task_kwargs : dictionary, optional
          The optional additional keyword arguments to the ``task`` function.
      reduce_func : func, optional
          If provided, it will be called with the output of each task instead of
          storing them in a list. Note that the order in which results are
          passed to ``reduce_func`` is not defined. It should return None or a
          number. When returning a number, it represents the estimation of the
          number of tasks left. On a return <= 0, the map will end early.
      progress_bar : str, optional
          Progress bar options's string for showing progress.
      progress_bar_kwargs : dict, optional
          Options for the progress bar.
      map_kw: dict, optional
          Dictionary containing entry for:
          - timeout: float, Maximum time (sec) for the whole map.
          - num_cpus: int, Number of jobs to run at once.
          - fail_fast: bool, Abort at the first error.
  
      Returns
      -------
      result : list
          The result list contains the value of
          ``task(value, *task_args, **task_kwargs)`` for
          each value in ``values``. If a ``reduce_func`` is provided, and empty
          list will be returned.
  
      """
  
      from loky import get_reusable_executor
      from loky.process_executor import ShutdownExecutorError
      map_kw = _read_map_kw(map_kw)
  
      def setup_executor():
          return get_reusable_executor(max_workers=map_kw['num_cpus'])
  
      def extract_result(future: concurrent.futures.Future):
          exception = future.exception()
          if isinstance(exception, ShutdownExecutorError):
              # Task was aborted due to timeout etc
              return None, None
          if exception is not None:
              return None, exception
          return future.result(), None
  
      def shutdown_executor(executor, active_tasks):
          # If there are still tasks running, we kill all workers in  order to
          # return immediately. Otherwise, `kill_workers` is set to False so
          # that the worker threads can be reused in subsequent loky_pmap calls.
          kill_workers = len(active_tasks) > 0
          executor.shutdown(kill_workers=kill_workers)
  
      return _generic_pmap(
          task, values, task_args, task_kwargs, reduce_func,
          map_kw['timeout'], map_kw['fail_fast'], map_kw['num_cpus'],
          progress_bar, progress_bar_kwargs,
          setup_executor, extract_result, shutdown_executor
      )
  
  
  def mpi_pmap(task, values, task_args=None, task_kwargs=None,
               reduce_func=None, map_kw=None,
               progress_bar=None, progress_bar_kwargs={}):
      """
      Parallel execution of a mapping of ``values`` to the function ``task``.
      This is functionally equivalent to::
  
          result = [task(value, *task_args, **task_kwargs) for value in values]
  
      Uses the mpi4py module to execute the tasks asynchronously with MPI
      processes. For more information, consult the documentation of mpi4py and
      the mpi4py.MPIPoolExecutor class.
  
      Note: in keeping consistent with the API of `parallel_map`, the parameter
      determining the number of requested worker processes is called `num_cpus`.
      The value of `map_kw['num_cpus']` is passed to the MPIPoolExecutor as its
      `max_workers` argument.
      If this parameter is not provided, the environment variable
      `QUTIP_NUM_PROCESSES` is used instead. If this environment variable is not
      set either, QuTiP will use default values that might be unsuitable for MPI
      applications.
  
      Parameters
      ----------
      task : a Python function
          The function that is to be called for each value in ``task_vec``.
      values : array / list
          The list or array of values for which the ``task`` function is to be
          evaluated.
      task_args : list, optional
          The optional additional arguments to the ``task`` function.
      task_kwargs : dictionary, optional
          The optional additional keyword arguments to the ``task`` function.
      reduce_func : func, optional
          If provided, it will be called with the output of each task instead of
          storing them in a list. Note that the order in which results are
          passed to ``reduce_func`` is not defined. It should return None or a
          number. When returning a number, it represents the estimation of the
          number of tasks left. On a return <= 0, the map will end early.
      progress_bar : str, optional
          Progress bar options's string for showing progress.
      progress_bar_kwargs : dict, optional
          Options for the progress bar.
      map_kw: dict, optional
          Dictionary containing entry for:
          - timeout: float, Maximum time (sec) for the whole map.
          - num_cpus: int, Number of jobs to run at once.
          - fail_fast: bool, Abort at the first error.
          All remaining entries of map_kw will be passed to the
          mpi4py.MPIPoolExecutor constructor.
  
      Returns
      -------
      result : list
          The result list contains the value of
          ``task(value, *task_args, **task_kwargs)`` for
          each value in ``values``. If a ``reduce_func`` is provided, and empty
          list will be returned.
  
      """
  
      from mpi4py.futures import MPIPoolExecutor
  
      # If loky is used (see loky_pmap), it starts two "resource tracker" child
      # processes that are normally never shut down. However, using mpi4py
      # when these child processes are still alive may lead to a deadlock.
      # -- WE RECOMMEND NOT TO MIX loky_pmap AND mpi_pmap IN THE SAME PROGRAM --
      # However, it happens in our test runs. For this reason, we here try to
      # shut down the resource tracker processes manually
      try:
          import loky
          loky_resource_tracker = loky.backend.resource_tracker._resource_tracker
          if loky_resource_tracker:
              loky_resource_tracker._stop()
  
          if hasattr(multiprocessing, 'resource_tracker'):
              mp_resource_tracker =\
                  multiprocessing.resource_tracker._resource_tracker
              if mp_resource_tracker:
                  mp_resource_tracker._stop()
      except ImportError:
          pass
  
      # If the provided num_cpus is None, we use the default value instead.
      # We thus intentionally make it impossible to call
      #   MPIPoolExecutor(max_workers=None, ...)
      # in which case mpi4py would determine a default value. That would be
      # useful, but unfortunately mpi4py provides no public API to access the
      # actual number of workers that is used in that case, which we would need.
      worker_number_provided = (
          ((map_kw is not None) and ('num_cpus' in map_kw))
          or 'QUTIP_NUM_PROCESSES' in os.environ)
  
      map_kw = _read_map_kw(map_kw)
      timeout = map_kw.pop('timeout')
      num_workers = map_kw.pop('num_cpus')
      fail_fast = map_kw.pop('fail_fast')
  
      if not worker_number_provided:
          warnings.warn(f'mpi_pmap was called without specifying the number of '
                        f'worker processes, using the default {num_workers}')
  
      def setup_executor():
          return MPIPoolExecutor(max_workers=num_workers, **map_kw)
  
      def extract_result(future):
          exception = future.exception()
          if exception is not None:
              return None, exception
          return future.result(), None
  
      def shutdown_executor(executor, _):
          executor.shutdown()
  
      return _generic_pmap(
          task, values, task_args, task_kwargs, reduce_func,
          timeout, fail_fast, num_workers,
          progress_bar, progress_bar_kwargs,
          setup_executor, extract_result, shutdown_executor
      )
  
  
  _maps = {
      "parallel_map": parallel_map,
      "parallel": parallel_map,
      "serial_map": serial_map,
      "serial": serial_map,
      "loky": loky_pmap,
      "mpi": mpi_pmap
  }
  
  
  def _get_map(options):
      map_func = _maps[options['map']]
  
      if map_func == mpi_pmap:
          map_kw = options['mpi_options']
      else:
          map_kw = {}
  
      return map_func, map_kw
  