content: import numpy as np
  import time
  import pytest
  import threading
  
  from qutip.solver.parallel import (
      parallel_map, serial_map, loky_pmap, mpi_pmap, MapExceptions
  )
  
  
  def _func1(x):
      return x**2
  
  
  def _func2(x, a, b, c, d=0, e=0, f=0):
      assert d > 0
      assert e > 0
      assert f > 0
      time.sleep(np.random.rand() * 0.1)  # random delay
      return x**2
  
  @pytest.mark.parametrize('map', [
      pytest.param(parallel_map, id='parallel_map'),
      pytest.param(loky_pmap, id='loky_pmap'),
      pytest.param(mpi_pmap, id='mpi_pmap'),
      pytest.param(serial_map, id='serial_map'),
  ])
  @pytest.mark.parametrize('num_cpus',
                           [1, 2],
                           ids=['1', '2'])
  def test_map(map, num_cpus):
      if map is loky_pmap:
          pytest.importorskip("loky")
      if map is mpi_pmap:
          pytest.importorskip("mpi4py")
  
      args = (1, 2, 3)
      kwargs = {'d': 4, 'e': 5, 'f': 6}
      map_kw = {
          'timeout': threading.TIMEOUT_MAX,
          'num_cpus': num_cpus,
      }
  
      x = np.arange(10)
      y1 = [_func1(xx) for xx in x]
  
      y2 = map(_func2, x, args, kwargs, map_kw=map_kw)
      assert ((np.array(y1) == np.array(y2)).all())
  
  
  @pytest.mark.parametrize('map', [
      pytest.param(parallel_map, id='parallel_map'),
      pytest.param(loky_pmap, id='loky_pmap'),
      pytest.param(mpi_pmap, id='mpi_pmap'),
      pytest.param(serial_map, id='serial_map'),
  ])
  @pytest.mark.parametrize('num_cpus',
                           [1, 2],
                           ids=['1', '2'])
  def test_map_accumulator(map, num_cpus):
      if map is loky_pmap:
          pytest.importorskip("loky")
      if map is mpi_pmap:
          pytest.importorskip("mpi4py")
  
      args = (1, 2, 3)
      kwargs = {'d': 4, 'e': 5, 'f': 6}
      map_kw = {
          'timeout': threading.TIMEOUT_MAX,
          'num_cpus': num_cpus,
      }
      y2 = []
  
      x = np.arange(10)
      y1 = [_func1(xx) for xx in x]
  
      map(_func2, x, args, kwargs, reduce_func=y2.append, map_kw=map_kw)
      assert ((np.array(sorted(y1)) == np.array(sorted(y2))).all())
  
  
  class CustomException(Exception):
      pass
  
  
  def func(i):
      if i % 2 == 1:
          raise CustomException(f"Error in subprocess {i}")
      return i
  
  
  @pytest.mark.parametrize('map', [
      pytest.param(parallel_map, id='parallel_map'),
      pytest.param(loky_pmap, id='loky_pmap'),
      pytest.param(mpi_pmap, id='mpi_pmap'),
      pytest.param(serial_map, id='serial_map'),
  ])
  def test_map_pass_error(map):
      kwargs = {}
      if map is loky_pmap:
          pytest.importorskip("loky")
      if map is mpi_pmap:
          pytest.importorskip("mpi4py")
          # do not use default value for num_cpus for mpi_pmap
          kwargs = {'map_kw': {'num_cpus': 1}}
  
      with pytest.raises(CustomException) as err:
          map(func, range(10), **kwargs)
      assert "Error in subprocess" in str(err.value)
  
  
  @pytest.mark.parametrize('map', [
      pytest.param(parallel_map, id='parallel_map'),
      pytest.param(loky_pmap, id='loky_pmap'),
      pytest.param(mpi_pmap, id='mpi_pmap'),
      pytest.param(serial_map, id='serial_map'),
  ])
  def test_map_store_error(map):
      map_kw = {"fail_fast": False}
      if map is loky_pmap:
          pytest.importorskip("loky")
      if map is mpi_pmap:
          pytest.importorskip("mpi4py")
          # do not use default value for num_cpus for mpi_pmap
          map_kw.update({'num_cpus': 1})
  
      with pytest.raises(MapExceptions) as err:
          map(func, range(10), map_kw=map_kw)
      map_error = err.value
      assert "iterations failed" in str(map_error)
      for iter, error in map_error.errors.items():
          assert isinstance(error, CustomException)
          assert f"Error in subprocess {iter}" == str(error)
      for n, result in enumerate(map_error.results):
          if n % 2 == 0:
              # Passed
              assert result == n
          else:
              assert result is None
  
  
  @pytest.mark.parametrize('map', [
      pytest.param(parallel_map, id='parallel_map'),
      pytest.param(loky_pmap, id='loky_pmap'),
      pytest.param(mpi_pmap, id='mpi_pmap'),
      pytest.param(serial_map, id='serial_map'),
  ])
  def test_map_early_end(map):
      kwargs = {}
      if map is loky_pmap:
          pytest.importorskip("loky")
      if map is mpi_pmap:
          pytest.importorskip("mpi4py")
          # do not use default value for num_cpus for mpi_pmap
          kwargs = {'map_kw': {'num_cpus': 1}}
  
      results = []
  
      def reduce_func(result):
          results.append(result)
          return 5 - len(results)
  
      map(_func1, range(100), reduce_func=reduce_func, **kwargs)
  
      assert len(results) < 100
  