jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def load_lammps(lammps_file, tol=_TOL):
  """Load a lammpstraj file.
  Args:
    lammps_file: file path
  Returns:
    coords: shape (F, N, 3)
    lattices: shape (F, 3)
    types: shape (N,)
    raw_types: shape (N,)
    unwrapped_coords: shape (F, N, 3)
  """
  # Loads lammpstraj file.
  with ExtendedLAMMPSTrajectoryFile(lammps_file) as f:
    coords, lattices, angles, types, ixyz, masses = f.read()
  assert ixyz is not None, 'needs to have ixyz'
  assert np.allclose(angles, 90.), 'trajs are not in orthogonal boxes'
  assert np.isclose(types, types[0]).all(), 'atom type changes'
  assert np.isclose(masses, masses[0]).all(), 'atom mass changes'
  if not np.isclose(lattices, lattices[0]).all():
    warnings.warn('box size changes.')
  ixyz = ixyz.astype('int32')
  types = types[0]
  masses = masses[0]
  assert types.shape[0] == coords.shape[1] == ixyz.shape[1] == masses.shape[0]
  assert coords.shape[0] == lattices.shape[0] == ixyz.shape[0]

  unwrapped_coords = coords + lattices[:, np.newaxis] * ixyz

  # Wrap coords
  imag_loc = np.floor_divide(coords, lattices[:, np.newaxis])
  wrapped_coords = coords - imag_loc * lattices[:, np.newaxis]
  assert np.all((np.max(wrapped_coords, axis=1) -
             np.min(wrapped_coords, axis=1)) <= lattices)

  atom_types = []
  for mass in masses:
    diffs = np.abs(mass - _ATOM_MASSES)
    atomic_number = np.argmin(diffs)
    assert diffs[atomic_number] <= tol, 'diff is {} for {}'.format(
      diffs[atomic_number], mass)
    atom_types.append(atomic_number)
  atom_types = np.array(atom_types, dtype=np.int32)
  return wrapped_coords, unwrapped_coords, lattices, types, atom_types
//  [python]
def __eq__(a, b):
        """ True if a and b have same elements, size and names """
        if b.__class__ == numpy.ndarray:
        # in case b is a just a numarray and not a Table instance
        # in this case, variable should absoltely be at the same order
        # otherwise the Table and numArray are considered as different
            return (numpy.alltrue(a.cpt.flat == b.flat) \
                    and a.shape == b.shape)

        elif b == None:
        # in case b is None type
            return False

        elif isinstance(b, (int, float, long)):
        # b is just a number, int, float, long
            return a.cpt == b

        else:
        # the b class should better be a Table or something like that
        # order of variables is not important
            # put the variables in the same order
            # first must get the correspondance vector :
            bcpt = a.prepareOther(b)
            return (a.names == b.names and \
                    bcpt.shape == a.shape and \
                    numpy.allclose(bcpt, a.cpt))
//  [python]
def __eq__(a, b):
        """ True if a and b have same elements, size and names """
        if b.__class__ == numpy.ndarray:
        # in case b is a just a numarray and not a Table instance
        # in this case, variable should absoltely be at the same order
        # otherwise the Table and numArray are considered as different
            return (numpy.all(a.cpt.flat == b.flat) \
                    and a.shape == b.shape)

        elif b == None:
        # in case b is None type
            return False

        elif isinstance(b, (int, float, long)):
        # b is just a number, int, float, long
            return a.cpt == b

        else:
        # the b class should better be a Table or something like that
        # order of variables is not important
            # put the variables in the same order
            # first must get the correspondance vector :
            bcpt = a.prepareOther(b)
            return (a.names == b.names and \
                    bcpt.shape == a.shape and \
                    numpy.allclose(bcpt, a.cpt))
//  [python]
def seq_allequal(seq1, seq2):
    """
    seq1 and seq2 are either None or sequences or arrays
    Return True if both are None or both are seqs with identical
    elements
    """
    if seq1 is None:
        return seq2 is None

    if seq2 is None:
        return False
    #ok, neither are None:, assuming iterable

    if len(seq1) != len(seq2): return False
    return np.alltrue(np.equal(seq1, seq2))
//  [python]
def seq_allequal(seq1, seq2):
    """
    seq1 and seq2 are either None or sequences or arrays
    Return True if both are None or both are seqs with identical
    elements
    """
    if seq1 is None:
        return seq2 is None

    if seq2 is None:
        return False
    #ok, neither are None:, assuming iterable

    if len(seq1) != len(seq2): return False
    return np.all(np.equal(seq1, seq2))
//  [python]
def read(self,fid_list,itype='auto',all_mo=True,nosym=False, sort=True, **kwargs_all):
    '''Reads a list of input files.
    
    **Parameters:**
    
      fid_list : list of str
        List of input file names.
      itype : str, choices={'auto', 'tar', 'molden', 'gamess', 'gaussian.log', 'gaussian.fchk'}
        Specifies the type of the input files.
      sort: bool
        Sort input files by name.
    '''
    # self.geo_info and ao_info have to stay unchanged
    geo_old = []
    ao_old = []
    
    sym_list = {}
    n_ao = {}

    #Check if fname poits to a tar archive and
    #read all files from archive if that is the case 
    if is_tar_file(fid_list):
      fid_list, itypes = get_all_files_from_tar(fid_list, sort=sort)
    else:
      itypes = [[itype]*len(fid_list)][0]

    for i,fname in enumerate(fid_list):

      kwargs = kwargs_all['kwargs'][i] if 'kwargs' in kwargs_all.keys() else kwargs_all
      qc = main_read(fname, itype=itypes[i], all_mo=all_mo, **kwargs)
      # Geo Section
      if i > 0 and (geo_old != qc.geo_info).sum():
        raise IOError('qc.geo_info has changed!')
      else:
        geo_old = deepcopy(qc.geo_info)
      self.geo_spec_all.append(qc.geo_spec)
      # AO Section
      if (i > 0 and not
          numpy.alltrue([numpy.allclose(ao_old[j]['coeffs'],qc.ao_spec[j]['coeffs'])
                        for j in range(len(ao_old))]
                        )):
          raise IOError('qc.ao_spec has changed!')
      else:
          ao_old = deepcopy(qc.ao_spec)

      # MO Section
      sym_tmp = {}    
      self.MO_Spec.append(qc.mo_spec)
      for i,mo in enumerate(qc.mo_spec):
        if nosym:
          qc.mo_spec[i]['sym'] = '%d.1' % (i+1)
        key = mo['sym'].split('.')
        if key[1] not in sym_tmp.keys():
          sym_tmp[key[1]] = 0
          n_ao[key[1]] = len(qc.mo_spec[0]['coeffs'])
        sym_tmp[key[1]] += 1

      for k,it in sym_tmp.items():
        if k in sym_list:
          sym_list[k] = max(sym_list[k],it)
        else:
          sym_list[k] = it
    
    self.geo_spec_all = numpy.array(self.geo_spec_all)
    self.geo_info = qc.geo_info
    self.ao_spec = qc.ao_spec

    # Presorting of the MOs according to their self.symmetry 
    n_r = len(fid_list)
    self.sym = []
    for k in sorted(sym_list.keys()):
      it = sym_list[k]
      self.sym.append((k,len(self.sym)))
      self.mo_coeff_all.append(numpy.zeros((n_r,it,n_ao[k])))
      self.mo_energy_all.append(numpy.zeros((n_r,it)))
      self.mo_occ_all.append(numpy.zeros((n_r,it)))

    self.sym = odict(self.sym)

    for i,spec in enumerate(self.MO_Spec):
      for j,mo in enumerate(spec):
        index,k = mo['sym'].split('.')
        
        index = int(index)-1
        
        self.mo_coeff_all[self.sym[k]][i,index,:] = mo['coeffs']
        self.mo_energy_all[self.sym[k]][i,index] = mo['energy']
        self.mo_occ_all[self.sym[k]][i,index] = mo['occ_num']
    return
//  [python]
def read(self,fid_list,itype='auto',all_mo=True,nosym=False, sort=True, **kwargs_all):
    '''Reads a list of input files.
    
    **Parameters:**
    
      fid_list : list of str
        List of input file names.
      itype : str, choices={'auto', 'tar', 'molden', 'gamess', 'gaussian.log', 'gaussian.fchk'}
        Specifies the type of the input files.
      sort: bool
        Sort input files by name.
    '''
    # self.geo_info and ao_info have to stay unchanged
    geo_old = []
    ao_old = []
    
    sym_list = {}
    n_ao = {}

    #Check if fname poits to a tar archive and
    #read all files from archive if that is the case 
    if is_tar_file(fid_list):
      fid_list, itypes = get_all_files_from_tar(fid_list, sort=sort)
    else:
      itypes = [[itype]*len(fid_list)][0]

    for i,fname in enumerate(fid_list):

      kwargs = kwargs_all['kwargs'][i] if 'kwargs' in kwargs_all.keys() else kwargs_all
      qc = main_read(fname, itype=itypes[i], all_mo=all_mo, **kwargs)
      # Geo Section
      if i > 0 and (geo_old != qc.geo_info).sum():
        raise IOError('qc.geo_info has changed!')
      else:
        geo_old = deepcopy(qc.geo_info)
      self.geo_spec_all.append(qc.geo_spec)
      # AO Section
      if (i > 0 and not
          numpy.all([numpy.allclose(ao_old[j]['coeffs'],qc.ao_spec[j]['coeffs'])
                        for j in range(len(ao_old))]
                        )):
          raise IOError('qc.ao_spec has changed!')
      else:
          ao_old = deepcopy(qc.ao_spec)

      # MO Section
      sym_tmp = {}    
      self.MO_Spec.append(qc.mo_spec)
      for i,mo in enumerate(qc.mo_spec):
        if nosym:
          qc.mo_spec[i]['sym'] = '%d.1' % (i+1)
        key = mo['sym'].split('.')
        if key[1] not in sym_tmp.keys():
          sym_tmp[key[1]] = 0
          n_ao[key[1]] = len(qc.mo_spec[0]['coeffs'])
        sym_tmp[key[1]] += 1

      for k,it in sym_tmp.items():
        if k in sym_list:
          sym_list[k] = max(sym_list[k],it)
        else:
          sym_list[k] = it
    
    self.geo_spec_all = numpy.array(self.geo_spec_all)
    self.geo_info = qc.geo_info
    self.ao_spec = qc.ao_spec

    # Presorting of the MOs according to their self.symmetry 
    n_r = len(fid_list)
    self.sym = []
    for k in sorted(sym_list.keys()):
      it = sym_list[k]
      self.sym.append((k,len(self.sym)))
      self.mo_coeff_all.append(numpy.zeros((n_r,it,n_ao[k])))
      self.mo_energy_all.append(numpy.zeros((n_r,it)))
      self.mo_occ_all.append(numpy.zeros((n_r,it)))

    self.sym = odict(self.sym)

    for i,spec in enumerate(self.MO_Spec):
      for j,mo in enumerate(spec):
        index,k = mo['sym'].split('.')
        
        index = int(index)-1
        
        self.mo_coeff_all[self.sym[k]][i,index,:] = mo['coeffs']
        self.mo_energy_all[self.sym[k]][i,index] = mo['energy']
        self.mo_occ_all[self.sym[k]][i,index] = mo['occ_num']
    return
//  [python]
def test_sel(comm):

    dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))

    # no exact match fails
    with pytest.raises(IndexError):
        sliced = dataset.sel(k=0.1)

    # this should be squeezed
    sliced = dataset.sel(k=0.1, method='nearest')
    assert len(sliced.dims) == 1

    # this is not squeezed
    sliced = dataset.sel(k=[0.1], method='nearest')
    assert sliced.shape[0] == 1

    # this return empty k with arbitary edges.
    sliced = dataset.sel(k=[], method='nearest')
    assert sliced.shape[0] == 0

    # slice in a specific k-range
    sliced = dataset.sel(k=slice(0.02, 0.15), mu=[0.5], method='nearest')
    assert sliced.shape[1] == 1
    assert numpy.alltrue((sliced['k'] >= 0.02)&(sliced['k'] <= 0.15))
//  [python]
def test_sel(comm):

    dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))

    # no exact match fails
    with pytest.raises(IndexError):
        sliced = dataset.sel(k=0.1)

    # this should be squeezed
    sliced = dataset.sel(k=0.1, method='nearest')
    assert len(sliced.dims) == 1

    # this is not squeezed
    sliced = dataset.sel(k=[0.1], method='nearest')
    assert sliced.shape[0] == 1

    # this return empty k with arbitary edges.
    sliced = dataset.sel(k=[], method='nearest')
    assert sliced.shape[0] == 0

    # slice in a specific k-range
    sliced = dataset.sel(k=slice(0.02, 0.15), mu=[0.5], method='nearest')
    assert sliced.shape[1] == 1
    assert numpy.all((sliced['k'] >= 0.02)&(sliced['k'] <= 0.15))
//  [python]
def test_reindex(comm):
    import warnings

    dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))

    with pytest.raises(ValueError):
        new, spacing = dataset.reindex('k', 0.005, force=True, return_spacing=True)

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)

        weights = numpy.random.random(dataset.shape)
        dataset['weights'] = weights
        new, spacing = dataset.reindex('k', 0.02, weights='weights', force=True, return_spacing=True)

        diff = numpy.diff(new.coords['k'])
        assert numpy.alltrue(diff > numpy.diff(dataset.coords['k'])[0])

        with pytest.raises(ValueError):
            new = dataset.reindex('mu', 0.4, force=False)
        new = dataset.reindex('mu', 0.4, force=True)
//  [python]
def test_reindex(comm):
    import warnings

    dataset = BinnedStatistic.from_json(os.path.join(data_dir, 'dataset_2d.json'))

    with pytest.raises(ValueError):
        new, spacing = dataset.reindex('k', 0.005, force=True, return_spacing=True)

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)

        weights = numpy.random.random(dataset.shape)
        dataset['weights'] = weights
        new, spacing = dataset.reindex('k', 0.02, weights='weights', force=True, return_spacing=True)

        diff = numpy.diff(new.coords['k'])
        assert numpy.all(diff > numpy.diff(dataset.coords['k'])[0])

        with pytest.raises(ValueError):
            new = dataset.reindex('mu', 0.4, force=False)
        new = dataset.reindex('mu', 0.4, force=True)
//  [python]
def _make_acyclic(self):
        """If the contour is cyclic (last point == first point), strip the last point off."""
        if numpy.alltrue(self.points[-1] == self.points[0]):
            self.points = numpy.resize(self.points, [self.points.shape[0] - 1, self.points.shape[1]])
//  [python]
def _make_acyclic(self):
        """If the contour is cyclic (last point == first point), strip the last point off."""
        if numpy.all(self.points[-1] == self.points[0]):
            self.points = numpy.resize(self.points, [self.points.shape[0] - 1, self.points.shape[1]])
//  [python]
def check_data_loader_is_not_random(data_loader: DataLoader, max_iter: int = 3) -> None:
    flag = False
    for i, ((x1, y1), (x2, y2)) in enumerate(zip(data_loader, data_loader)):
        if i > max_iter:
            break
        if isinstance(x1, dict):
            if not all(
                [np.alltrue(x1[key] == x2[key]) for key in x1.keys()]
            ) or not np.alltrue(y1 == y2):
                flag = True
                break
        else:
            if not np.alltrue(x1 == x2) or not np.alltrue(y1 == y2):
                flag = True
                break

    if flag:
        raise ValueError(
            """The data loader randomizes at every iteration. To perform this method, please provide a data loader that
            generates the same sequence of data when called multiple times."""
        )
//  [python]
def check_data_loader_is_not_random(data_loader: DataLoader, max_iter: int = 3) -> None:
    flag = False
    for i, ((x1, y1), (x2, y2)) in enumerate(zip(data_loader, data_loader)):
        if i > max_iter:
            break
        if isinstance(x1, dict):
            if not all(
                [np.all(x1[key] == x2[key]) for key in x1.keys()]
            ) or not np.all(y1 == y2):
                flag = True
                break
        else:
            if not np.all(x1 == x2) or not np.all(y1 == y2):
                flag = True
                break

    if flag:
        raise ValueError(
            """The data loader randomizes at every iteration. To perform this method, please provide a data loader that
            generates the same sequence of data when called multiple times."""
        )
//  [python]
def finite_model(self):
        for shp in self.shapelets:
            if not numpy.alltrue(numpy.isfinite(shp)):
                return False
        return True
//  [python]
def finite_model(self):
        for shp in self.shapelets:
            if not numpy.all(numpy.isfinite(shp)):
                return False
        return True
//  [python]
def runTest(self):
        users_cnt, documents_cnt = self.users, self.documents
        config = RecommenderConfiguration().get_hyperparameters()
        config['n_factors'] = 5
        initializer = ModelInitializer(config, 1)
        path = initializer._create_path('user_v', (users_cnt, documents_cnt))
        self.assertTrue(path.endswith('n_iterations-1,n_rows-10user_v.dat'))
        matrix_shape = (users_cnt, config['n_factors'])
        users_mat = numpy.random.random(matrix_shape)
        initializer.save_matrix(users_mat, 'user_v')
        self.assertTrue(os.path.isfile(path))
        loaded, loaded_matrix = initializer.load_matrix(config, 'user_v', matrix_shape)
        self.assertTrue(loaded)
        self.assertTrue(numpy.alltrue(loaded_matrix == users_mat))
//  [python]
def runTest(self):
        users_cnt, documents_cnt = self.users, self.documents
        config = RecommenderConfiguration().get_hyperparameters()
        config['n_factors'] = 5
        initializer = ModelInitializer(config, 1)
        path = initializer._create_path('user_v', (users_cnt, documents_cnt))
        self.assertTrue(path.endswith('n_iterations-1,n_rows-10user_v.dat'))
        matrix_shape = (users_cnt, config['n_factors'])
        users_mat = numpy.random.random(matrix_shape)
        initializer.save_matrix(users_mat, 'user_v')
        self.assertTrue(os.path.isfile(path))
        loaded, loaded_matrix = initializer.load_matrix(config, 'user_v', matrix_shape)
        self.assertTrue(loaded)
        self.assertTrue(numpy.all(loaded_matrix == users_mat))
//  [python]
def __add__(self, it):
        if isinstance(it, DirectoryIterator):
            assert self.nb_sample == it.nb_sample
            assert len(self.filenames) == len(it.filenames)
            assert np.alltrue(self.classes == it.classes)
            assert self.image_reader == it.image_reader
            if inspect.isgeneratorfunction(self.image_reader):
                self._reader_generator = []
                it._reader_generator = []
        if isinstance(it, NumpyArrayIterator):
            assert self.nb_sample == self.X.shape[0]
        it.image_data_generator.sync(self.image_data_generator)
        return super(DirectoryIterator, self).__add__(it)
//  [python]
def __add__(self, it):
        if isinstance(it, DirectoryIterator):
            assert self.nb_sample == it.nb_sample
            assert len(self.filenames) == len(it.filenames)
            assert np.all(self.classes == it.classes)
            assert self.image_reader == it.image_reader
            if inspect.isgeneratorfunction(self.image_reader):
                self._reader_generator = []
                it._reader_generator = []
        if isinstance(it, NumpyArrayIterator):
            assert self.nb_sample == self.X.shape[0]
        it.image_data_generator.sync(self.image_data_generator)
        return super(DirectoryIterator, self).__add__(it)
//  [python]
def subblocks(a, *ranges):
    """
    This function produces a distributed array from a subset of the blocks in
    the `a`. The result and `a` will have the same number of dimensions. For
    example,
        subblocks(a, [0, 1], [2, 4])
    will produce a DistArray whose object_refs are
        [[a.object_refs[0, 2], a.object_refs[0, 4]],
         [a.object_refs[1, 2], a.object_refs[1, 4]]]
    We allow the user to pass in an empty list [] to indicate the full range.
    """
    ranges = list(ranges)
    if len(ranges) != a.ndim:
        raise Exception(
            "sub_blocks expects to receive a number of ranges "
            "equal to a.ndim, but it received {} ranges and "
            "a.ndim = {}.".format(len(ranges), a.ndim)
        )
    for i in range(len(ranges)):
        # We allow the user to pass in an empty list to indicate the full
        # range.
        if ranges[i] == []:
            ranges[i] = range(a.num_blocks[i])
        if not np.alltrue(ranges[i] == np.sort(ranges[i])):
            raise Exception(
                "Ranges passed to sub_blocks must be sorted, but "
                "the {}th range is {}.".format(i, ranges[i])
            )
        if ranges[i][0] < 0:
            raise Exception(
                "Values in the ranges passed to sub_blocks must "
                "be at least 0, but the {}th range is {}.".format(i, ranges[i])
            )
        if ranges[i][-1] >= a.num_blocks[i]:
            raise Exception(
                "Values in the ranges passed to sub_blocks must "
                "be less than the relevant number of blocks, but "
                "the {}th range is {}, and a.num_blocks = {}.".format(
                    i, ranges[i], a.num_blocks
                )
            )
    last_index = [r[-1] for r in ranges]
    last_block_shape = DistArray.compute_block_shape(last_index, a.shape)
    shape = [
        (len(ranges[i]) - 1) * BLOCK_SIZE + last_block_shape[i] for i in range(a.ndim)
    ]
    result = DistArray(shape)
    for index in np.ndindex(*result.num_blocks):
        result.object_refs[index] = a.object_refs[
            tuple(ranges[i][index[i]] for i in range(a.ndim))
        ]
    return result
//  [python]
def subblocks(a, *ranges):
    """
    This function produces a distributed array from a subset of the blocks in
    the `a`. The result and `a` will have the same number of dimensions. For
    example,
        subblocks(a, [0, 1], [2, 4])
    will produce a DistArray whose object_refs are
        [[a.object_refs[0, 2], a.object_refs[0, 4]],
         [a.object_refs[1, 2], a.object_refs[1, 4]]]
    We allow the user to pass in an empty list [] to indicate the full range.
    """
    ranges = list(ranges)
    if len(ranges) != a.ndim:
        raise Exception(
            "sub_blocks expects to receive a number of ranges "
            "equal to a.ndim, but it received {} ranges and "
            "a.ndim = {}.".format(len(ranges), a.ndim)
        )
    for i in range(len(ranges)):
        # We allow the user to pass in an empty list to indicate the full
        # range.
        if ranges[i] == []:
            ranges[i] = range(a.num_blocks[i])
        if not np.all(ranges[i] == np.sort(ranges[i])):
            raise Exception(
                "Ranges passed to sub_blocks must be sorted, but "
                "the {}th range is {}.".format(i, ranges[i])
            )
        if ranges[i][0] < 0:
            raise Exception(
                "Values in the ranges passed to sub_blocks must "
                "be at least 0, but the {}th range is {}.".format(i, ranges[i])
            )
        if ranges[i][-1] >= a.num_blocks[i]:
            raise Exception(
                "Values in the ranges passed to sub_blocks must "
                "be less than the relevant number of blocks, but "
                "the {}th range is {}, and a.num_blocks = {}.".format(
                    i, ranges[i], a.num_blocks
                )
            )
    last_index = [r[-1] for r in ranges]
    last_block_shape = DistArray.compute_block_shape(last_index, a.shape)
    shape = [
        (len(ranges[i]) - 1) * BLOCK_SIZE + last_block_shape[i] for i in range(a.ndim)
    ]
    result = DistArray(shape)
    for index in np.ndindex(*result.num_blocks):
        result.object_refs[index] = a.object_refs[
            tuple(ranges[i][index[i]] for i in range(a.ndim))
        ]
    return result
//  [python]
def is_identity(self):
        A = self.A
        I = numpy.identity(len(A), dtype=int)
        return numpy.alltrue(A==I)
//  [python]
def is_identity(self):
        A = self.A
        I = numpy.identity(len(A), dtype=int)
        return numpy.all(A==I)
//  [python]
def test_properties(self):
        rc = ctracer.RayCollection(10)
        rays = [self.make_ray() for i in range(7)]
        for ray in rays:
            rc.add_ray(ray)
        data = rc.copy_as_array()
        self.assertTrue( numpy.alltrue( rc.E1_amp==data['E1_amp'] ) )
        self.assertTrue( numpy.alltrue( rc.E2_amp==data['E2_amp'] ) )
        self.assertTrue( numpy.alltrue( rc.refractive_index==data['refractive_index'] ) )
//  [python]
def test_properties(self):
        rc = ctracer.RayCollection(10)
        rays = [self.make_ray() for i in range(7)]
        for ray in rays:
            rc.add_ray(ray)
        data = rc.copy_as_array()
        self.assertTrue( numpy.all( rc.E1_amp==data['E1_amp'] ) )
        self.assertTrue( numpy.all( rc.E2_amp==data['E2_amp'] ) )
        self.assertTrue( numpy.all( rc.refractive_index==data['refractive_index'] ) )
//  [python]
def calculate_mean_contour(contours):
    """Calculate the average of a set of contours, while retaining units and
    scaling information, if possible. If all contours have associated landmarks,
    then the average will be such a contour as well."""
    all_points = [c.points for c in contours]
    if not utility_tools.all_same_shape(all_points):
        raise ContourError('Cannot calculate mean of contours with different numbers of points.')
    mean_points = numpy.mean(all_points, axis=0)
    units = [c.units for c in contours]
    if not numpy.alltrue([u == units[0] for u in units]):
        raise ContourError('All contours must have the same units in order calculate their mean.')
    units = contours[0].units
    scales = [utility_tools.decompose_homogenous_transform(c.to_world_transform)[1] for c in contours]
    if numpy.alltrue([numpy.allclose(scales[0], s) for s in scales[1:]]):
        transform = utility_tools.make_homogenous_transform(transform=scales[0])
    else:
        transform = numpy.eye(3)
    if numpy.alltrue([isinstance(c, ContourAndLandmarks) for c in contours]):
        # if they're all landmark'd contours
        all_landmarks = [c.landmarks for c in contours]
        if not utility_tools.all_same_shape(all_landmarks):
            raise ContourError('Cannot calculate mean of contours with different numbers of landmarks.')
        mean_landmarks = numpy.mean(all_landmarks, axis=0)
        mean_weights = numpy.mean([c.weights for c in contours], axis=0)
        return ContourAndLandmarks(points=mean_points, units=units, landmarks=mean_landmarks,
            weights=mean_weights, to_world_transform=transform)
    else:
        return Contour(points=mean_points, units=units, to_world_transform=transform)
//  [python]
def calculate_mean_contour(contours):
    """Calculate the average of a set of contours, while retaining units and
    scaling information, if possible. If all contours have associated landmarks,
    then the average will be such a contour as well."""
    all_points = [c.points for c in contours]
    if not utility_tools.all_same_shape(all_points):
        raise ContourError('Cannot calculate mean of contours with different numbers of points.')
    mean_points = numpy.mean(all_points, axis=0)
    units = [c.units for c in contours]
    if not numpy.all([u == units[0] for u in units]):
        raise ContourError('All contours must have the same units in order calculate their mean.')
    units = contours[0].units
    scales = [utility_tools.decompose_homogenous_transform(c.to_world_transform)[1] for c in contours]
    if numpy.all([numpy.allclose(scales[0], s) for s in scales[1:]]):
        transform = utility_tools.make_homogenous_transform(transform=scales[0])
    else:
        transform = numpy.eye(3)
    if numpy.all([isinstance(c, ContourAndLandmarks) for c in contours]):
        # if they're all landmark'd contours
        all_landmarks = [c.landmarks for c in contours]
        if not utility_tools.all_same_shape(all_landmarks):
            raise ContourError('Cannot calculate mean of contours with different numbers of landmarks.')
        mean_landmarks = numpy.mean(all_landmarks, axis=0)
        mean_weights = numpy.mean([c.weights for c in contours], axis=0)
        return ContourAndLandmarks(points=mean_points, units=units, landmarks=mean_landmarks,
            weights=mean_weights, to_world_transform=transform)
    else:
        return Contour(points=mean_points, units=units, to_world_transform=transform)
//  [python]
raw_input("Press return when ready!")
for i in range(10):
    mlab.sum(myMat,nout=0)
for i in range(1000):
    s = mlab.sin(rand(100,100))
#print loads of crap to make sure this won't screw up
for i in range(20):
    mlab.sin(rand(100,100), nout=0)
assert numpy.alltrue(myMat.flat == mlab._get('mymat',1).flat)
//  [python]
raw_input("Press return when ready!")
for i in range(10):
    mlab.sum(myMat,nout=0)
for i in range(1000):
    s = mlab.sin(rand(100,100))
#print loads of crap to make sure this won't screw up
for i in range(20):
    mlab.sin(rand(100,100), nout=0)
assert numpy.all(myMat.flat == mlab._get('mymat',1).flat)
//  [python]
def test_expand_dims_0(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, 0)
      a = xdl.execute(a)
      self.assertTrue(np.alltrue(a == [[1, 2], [3, 4]]))
//  [python]
def test_expand_dims_0(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, 0)
      a = xdl.execute(a)
      self.assertTrue(np.all(a == [[1, 2], [3, 4]]))
//  [python]
def from_contours(cls, contours, required_variance_explained = 0.95, return_positions = False):
        """This class method should be used to construct a PCAContour object from a
        set of contours. The proncipal components of the contours are caltulated, and
        enough retained to account for the 'required_variance_explained' fraction.

        The mean shape is stored in the instance variable 'mean' and the principal components
        in 'modes'. The square root of the variance explained by each mode is stored in
        'standard_deviations', while the total variance of all of the modes (even those
        not retained) is in 'total_variance'.

        If the 'return_positions' parameter is true, then a tuple is returned containing
        (pca_contour, positions, normalized_positions), where 'pca_contour' is the new
        PCAContour instance, 'positions' is the position of each of the input contours
        along each principal shape mode, and 'normalized_positions' is the position
        along each mode in terms of standard deviations along that mode.
        """
        from celltool.numerics import pca
        data = [c.points for c in contours]
        if not utility_tools.all_same_shape(data):
            raise ValueError('All contours must have the same number of points in order to perform PCA.')
        units = [c.units for c in contours]
        if not numpy.alltrue([u == units[0] for u in units]):
            raise ValueError('All contours must have the same units in order to produce a PCA shape model from them.')
        units = units[0]
        scales = [utility_tools.decompose_homogenous_transform(c.to_world_transform)[1] for c in contours]
        if numpy.alltrue([numpy.allclose(scales[0], s) for s in scales[1:]]):
            transform = utility_tools.make_homogenous_transform(transform=scales[0])
        else:
            transform = numpy.eye(3)
        vals = pca.pca_dimensionality_reduce(numpy.array(data, dtype=numpy.float32), required_variance_explained)
        mean, pcs, norm_pcs, variances, total_variance, positions, norm_positions = vals
        c = cls(points=mean, mean=mean, modes=pcs, standard_deviations=numpy.sqrt(variances),
            total_variance=total_variance, position=numpy.zeros(len(norm_pcs)), units=units,
            to_world_transform=transform)
        if return_positions:
            return c, positions, norm_positions
        else:
            return c
//  [python]
def from_contours(cls, contours, required_variance_explained = 0.95, return_positions = False):
        """This class method should be used to construct a PCAContour object from a
        set of contours. The proncipal components of the contours are caltulated, and
        enough retained to account for the 'required_variance_explained' fraction.

        The mean shape is stored in the instance variable 'mean' and the principal components
        in 'modes'. The square root of the variance explained by each mode is stored in
        'standard_deviations', while the total variance of all of the modes (even those
        not retained) is in 'total_variance'.

        If the 'return_positions' parameter is true, then a tuple is returned containing
        (pca_contour, positions, normalized_positions), where 'pca_contour' is the new
        PCAContour instance, 'positions' is the position of each of the input contours
        along each principal shape mode, and 'normalized_positions' is the position
        along each mode in terms of standard deviations along that mode.
        """
        from celltool.numerics import pca
        data = [c.points for c in contours]
        if not utility_tools.all_same_shape(data):
            raise ValueError('All contours must have the same number of points in order to perform PCA.')
        units = [c.units for c in contours]
        if not numpy.all([u == units[0] for u in units]):
            raise ValueError('All contours must have the same units in order to produce a PCA shape model from them.')
        units = units[0]
        scales = [utility_tools.decompose_homogenous_transform(c.to_world_transform)[1] for c in contours]
        if numpy.all([numpy.allclose(scales[0], s) for s in scales[1:]]):
            transform = utility_tools.make_homogenous_transform(transform=scales[0])
        else:
            transform = numpy.eye(3)
        vals = pca.pca_dimensionality_reduce(numpy.array(data, dtype=numpy.float32), required_variance_explained)
        mean, pcs, norm_pcs, variances, total_variance, positions, norm_positions = vals
        c = cls(points=mean, mean=mean, modes=pcs, standard_deviations=numpy.sqrt(variances),
            total_variance=total_variance, position=numpy.zeros(len(norm_pcs)), units=units,
            to_world_transform=transform)
        if return_positions:
            return c, positions, norm_positions
        else:
            return c
//  [python]
def test_expand_dims_2(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, 2)
      a = xdl.execute(a)
      self.assertTrue(np.alltrue(a == [[[1], [2]], [[3], [4]]]))
//  [python]
def test_expand_dims_2(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, 2)
      a = xdl.execute(a)
      self.assertTrue(np.all(a == [[[1], [2]], [[3], [4]]]))
//  [python]
def test_expand_dims_1(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, 1)
      a = xdl.execute(a)
      self.assertTrue(np.alltrue(a == [[[1, 2]], [[3, 4]]]))
//  [python]
def test_expand_dims_1(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, 1)
      a = xdl.execute(a)
      self.assertTrue(np.all(a == [[[1, 2]], [[3, 4]]]))
//  [python]
def test_expand_dims_x3(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, -3)
      a = xdl.execute(a)
      self.assertTrue(np.alltrue(a == [[1, 2], [3, 4]]))
//  [python]
def test_expand_dims_x3(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, -3)
      a = xdl.execute(a)
      self.assertTrue(np.all(a == [[1, 2], [3, 4]]))
//  [python]
def add(self, thisType, value, position=None):
        """Add data to an existing data type (and add a new one if necess)
        """
        if not thisType in self:
            self.addDataType(thisType)
        if position is None:
            # 'ran' is always the first thing to update
            repN = sum(self['ran'][self.trials.thisIndex])
            if thisType != 'ran':
                # because it has already been updated
                repN -= 1
            # make a list where 1st digit is trial number
            position = [self.trials.thisIndex]
            position.append(repN)

        # check whether data falls within bounds
        posArr = np.asarray(position)
        shapeArr = np.asarray(self.dataShape)
        if not np.alltrue(posArr < shapeArr):
            # array isn't big enough
            logging.warning('need a bigger array for: ' + thisType)
            # not implemented yet!
            self[thisType] = extendArr(self[thisType], posArr)
        # check for ndarrays with more than one value and for non-numeric data
        if (self.isNumeric[thisType] and
                ((type(value) == np.ndarray and len(value) > 1) or
                     (type(value) not in [float, int]))):
            self._convertToObjectArray(thisType)
        # insert the value
        self[thisType][position[0], int(position[1])] = value
//  [python]
def add(self, thisType, value, position=None):
        """Add data to an existing data type (and add a new one if necess)
        """
        if not thisType in self:
            self.addDataType(thisType)
        if position is None:
            # 'ran' is always the first thing to update
            repN = sum(self['ran'][self.trials.thisIndex])
            if thisType != 'ran':
                # because it has already been updated
                repN -= 1
            # make a list where 1st digit is trial number
            position = [self.trials.thisIndex]
            position.append(repN)

        # check whether data falls within bounds
        posArr = np.asarray(position)
        shapeArr = np.asarray(self.dataShape)
        if not np.all(posArr < shapeArr):
            # array isn't big enough
            logging.warning('need a bigger array for: ' + thisType)
            # not implemented yet!
            self[thisType] = extendArr(self[thisType], posArr)
        # check for ndarrays with more than one value and for non-numeric data
        if (self.isNumeric[thisType] and
                ((type(value) == np.ndarray and len(value) > 1) or
                     (type(value) not in [float, int]))):
            self._convertToObjectArray(thisType)
        # insert the value
        self[thisType][position[0], int(position[1])] = value
//  [python]
def test_expand_dims_x2(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, -2)
      a = xdl.execute(a)
      self.assertTrue(np.alltrue(a == [[[1, 2]], [[3, 4]]]))
//  [python]
def test_expand_dims_x2(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, -2)
      a = xdl.execute(a)
      self.assertTrue(np.all(a == [[[1, 2]], [[3, 4]]]))
//  [python]
def test_expand_dims_x1(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, -1)
      a = xdl.execute(a)
      self.assertTrue(np.alltrue(a == [[[1], [2]], [[3], [4]]]))
//  [python]
def test_expand_dims_x1(self):
      value = np.array([[1, 2], [3, 4]])
      a = xdl.expand_dims(value, -1)
      a = xdl.execute(a)
      self.assertTrue(np.all(a == [[[1], [2]], [[3], [4]]]))
//  [python]
def test_pixfun_imag_r():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_imag_r.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    assert numpy.alltrue(data == 0)
//  [python]
def test_pixfun_imag_r():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_imag_r.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    assert numpy.all(data == 0)
//  [python]
def test_pixfun_real_r():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_real_r.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/int32.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    assert numpy.alltrue(data == refdata.real)
//  [python]
def test_pixfun_real_r():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_real_r.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/int32.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    assert numpy.all(data == refdata.real)
//  [python]
def test_pixfun_real_c():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_real_c.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/cint_sar.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    assert numpy.alltrue(data == refdata.real)
//  [python]
def test_pixfun_real_c():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_real_c.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/cint_sar.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    assert numpy.all(data == refdata.real)
//  [python]
def test_pixfun_imag_c():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_imag_c.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/cint_sar.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    assert numpy.alltrue(data == refdata.imag)

    # Test bugfix of #6599
    copied_ds = gdal.Translate('', filename, format='MEM')
    data_ds = copied_ds.GetRasterBand(1).ReadAsArray()
    copied_ds = None

    assert numpy.alltrue(data == data_ds)
//  [python]
def test_pixfun_imag_c():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_imag_c.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/cint_sar.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    assert numpy.all(data == refdata.imag)

    # Test bugfix of #6599
    copied_ds = gdal.Translate('', filename, format='MEM')
    data_ds = copied_ds.GetRasterBand(1).ReadAsArray()
    copied_ds = None

    assert numpy.all(data == data_ds)
//  [python]
def test_pixfun_mod_c():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_mod_c.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/cint_sar.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    res = numpy.alltrue(data == numpy.abs(refdata))
    if gdaltest.is_travis_branch('sanitize') and not res:
        print(data)
        print(numpy.abs(refdata))
        pytest.xfail()

    assert res
//  [python]
def test_pixfun_mod_c():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_mod_c.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/cint_sar.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    res = numpy.all(data == numpy.abs(refdata))
    if gdaltest.is_travis_branch('sanitize') and not res:
        print(data)
        print(numpy.abs(refdata))
        pytest.xfail()

    assert res
//  [python]
def test_pixfun_mod_r():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_mod_r.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/int32.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    res = numpy.alltrue(data == numpy.abs(refdata))
    if gdaltest.is_travis_branch('sanitize') and not res:
        print(data)
        print(numpy.abs(refdata))
        pytest.xfail()

    assert res
//  [python]
def test_pixfun_mod_r():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_mod_r.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/int32.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    res = numpy.all(data == numpy.abs(refdata))
    if gdaltest.is_travis_branch('sanitize') and not res:
        print(data)
        print(numpy.abs(refdata))
        pytest.xfail()

    assert res
//  [python]
def test_pixfun_phase_r():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_phase_r.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/pixfun_imag_c.vrt'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    assert numpy.alltrue(data == numpy.arctan2(0, refdata))
//  [python]
def test_pixfun_phase_r():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_phase_r.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/pixfun_imag_c.vrt'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    assert numpy.all(data == numpy.arctan2(0, refdata))
//  [python]
def test_pixfun_conj_c():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_conj_c.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/cint_sar.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    assert numpy.alltrue(data == numpy.conj(refdata))
//  [python]
def test_pixfun_conj_c():

    if not numpy_available:
        pytest.skip()

    filename = 'data/pixfun_conj_c.vrt'
    ds = gdal.OpenShared(filename, gdal.GA_ReadOnly)
    assert ds is not None, ('Unable to open "%s" dataset.' % filename)
    data = ds.GetRasterBand(1).ReadAsArray()

    reffilename = 'data/cint_sar.tif'
    refds = gdal.Open(reffilename)
    assert refds is not None, ('Unable to open "%s" dataset.' % reffilename)
    refdata = refds.GetRasterBand(1).ReadAsArray()

    assert numpy.all(data == numpy.conj(refdata))