jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def add_adaptive_distance_restraints(self, restraints):
        '''
        Add a set of adaptive distance restraints to the simulation. Sets
        :attr:`sim_index` for each restraint so it knows its place in the
        simulation. Automatically calls :func:`context_reinit_needed`.

        Args:
            * restraints:
                - a :py:class:`AdaptiveDistanceRestraints` instance
        '''
        force = self._adaptive_distance_restraints_force
        all_atoms = self._atoms
        dr_atoms = restraints.atoms
        indices = numpy.array([all_atoms.indices(atoms) for atoms in dr_atoms])
        ifilter = n.all(indices!=-1, axis=0)
        indices = [i[ifilter] for i in indices]
        restraints = restraints[ifilter]
        restraints.sim_indices = force.add_bonds(indices,
            restraints.enableds, restraints.kappas, restraints.cs/10,
            restraints.targets/10, restraints.tolerances/10, restraints.alphas
        )
        self.context_reinit_needed()
//  [python]
def add_adaptive_distance_restraints(self, restraints):
        '''
        Add a set of adaptive distance restraints to the simulation. Sets
        :attr:`sim_index` for each restraint so it knows its place in the
        simulation. Automatically calls :func:`context_reinit_needed`.

        Args:
            * restraints:
                - a :py:class:`AdaptiveDistanceRestraints` instance
        '''
        force = self._adaptive_distance_restraints_force
        all_atoms = self._atoms
        dr_atoms = restraints.atoms
        indices = numpy.array([all_atoms.indices(atoms) for atoms in dr_atoms])
        ifilter = numpy.all(indices!=-1, axis=0)
        indices = [i[ifilter] for i in indices]
        restraints = restraints[ifilter]
        restraints.sim_indices = force.add_bonds(indices,
            restraints.enableds, restraints.kappas, restraints.cs/10,
            restraints.targets/10, restraints.tolerances/10, restraints.alphas
        )
        self.context_reinit_needed()
//  [python]
def add_adaptive_dihedral_restraints(self, restraints):
        '''
        Add a set of adaptive dihedral restraints to the simulation. This sets the
        :attr:`sim_index` for each restraint so it knows its own place in the
        simulation for later updating purposes. Automatically calls
        :func:`context_reinit_needed()`.

        Args:
            * restraints:
                - A :py:class:`AdaptiveDihedralRestraints` object.
        '''
        force = self._adaptive_dihedral_restraint_force
        all_atoms = self._atoms
        dihedral_atoms = restraints.dihedrals.atoms
        atom_indices = numpy.array([all_atoms.indices(atoms) for atoms in dihedral_atoms])
        # Filter out those which don't have all atoms in the simulation
        ifilter = numpy.all(atom_indices!=-1, axis=0)
        atom_indices = [a[ifilter] for a in atom_indices]
        #atom_indices = atom_indices[ifilter]
        restraints = restraints[ifilter]
        restraints.sim_indices = force.add_torsions(atom_indices,
            restraints.enableds, restraints.spring_constants,
            restraints.targets, restraints.kappas, restraints.alphas)
        self.context_reinit_needed()
//  [python]
def merge(datasets: Sequence["Dataset"]) -> "Dataset":
        """Merge Datasets along the item dimension

        Parameters
        ---------
        datasets: sequence of Datasets

        Returns
        -------
        Dataset
            merged dataset
        """
        ds = datasets[0].copy()
        for other in datasets[1:]:
            item_names = {item.name for item in ds.items}
            other_names = {item.name for item in other.items}

            overlap = other_names.intersection(item_names)
            if len(overlap) != 0:
                raise ValueError("Can not append items, names are not unique")

            if not n.all(ds.time == other.time):
                raise ValueError("All timesteps must match")

            for key, value in other._data_vars.items():
                if key != "Z coordinate":
                    ds[key] = value

        return ds
//  [python]
def merge(datasets: Sequence["Dataset"]) -> "Dataset":
        """Merge Datasets along the item dimension

        Parameters
        ---------
        datasets: sequence of Datasets

        Returns
        -------
        Dataset
            merged dataset
        """
        ds = datasets[0].copy()
        for other in datasets[1:]:
            item_names = {item.name for item in ds.items}
            other_names = {item.name for item in other.items}

            overlap = other_names.intersection(item_names)
            if len(overlap) != 0:
                raise ValueError("Can not append items, names are not unique")

            if not np.all(ds.time == other.time):
                raise ValueError("All timesteps must match")

            for key, value in other._data_vars.items():
                if key != "Z coordinate":
                    ds[key] = value

        return ds
//  [python]
def _append_items(self, other: "Dataset") -> "Dataset":
        item_names = {item.name for item in self.items}
        other_names = {item.name for item in other.items}

        overlap = other_names.intersection(item_names)
        if len(overlap) != 0:
            raise ValueError("Can not append items, names are not unique")

        if not np.all(self.time == other.time):
            # if not: create common time?
            raise ValueError("All timesteps must match")

        for key, value in other._data_vars.items():
            if key != "Z coordinate":
                self[key] = value

        return self
//  [python]
def complete_cases(x):
    return numpy.all(numpy.invert(numpy.isnan(x)), axis=1)
//  [python]
def create_pcd(
        points: T.Union[torch.Tensor, np.ndarray],
        colors: T.Union[torch.Tensor, np.ndarray] = None,
        remove_nan_inf: bool = True,
) -> o3d.geometry.PointCloud:
    """
    Create o3d point cloud from points
    Args:
        points:
            (*, 3)
        colors:
            (*, 3) optional
        remove_nan_inf:

    Returns:
        an o3d point cloud
    """

    if isinstance(points, torch.Tensor):
        points = points.detach().cpu().numpy()

    points = points.reshape(-1, 3)  # (n, 3)
    if colors is not None:
        colors = colors.reshape(-1, 3)  # (n, 3)

    # remove any inf or nan points
    if remove_nan_inf:
        idxs = np.isfinite(points).all(axis=-1)  # (n,)
        points = points[idxs]
        if colors is not None:
            colors = colors[idxs]

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    if colors is not None:
        pcd.colors = o3d.utility.Vector3dVector(colors)

    return pcd
//  [python]
def create_pcd(
        points: T.Union[torch.Tensor, np.ndarray],
        colors: T.Union[torch.Tensor, np.ndarray] = None,
        remove_nan_inf: bool = True,
) -> o3d.geometry.PointCloud:
    """
    Create o3d point cloud from points
    Args:
        points:
            (*, 3)
        colors:
            (*, 3) optional
        remove_nan_inf:

    Returns:
        an o3d point cloud
    """

    if isinstance(points, torch.Tensor):
        points = points.detach().cpu().numpy()

    points = points.reshape(-1, 3)  # (n, 3)
    if colors is not None:
        colors = colors.reshape(-1, 3)  # (n, 3)

    # remove any inf or nan points
    if remove_nan_inf:
        idxs = np.all(np.isfinite(points), axis=-1)  # (n,)
        points = points[idxs]
        if colors is not None:
            colors = colors[idxs]

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    if colors is not None:
        pcd.colors = o3d.utility.Vector3dVector(colors)

    return pcd
//  [python]
def test_period_list_cross_a_period():
    """Tests the PeriodList.cross_a_period method."""
    periods = numpy.array(
        [['2019-12-01T02:33:57.989', '2019-12-01T03:24:59.504'],
         ['2019-12-01T04:16:09.242', '2019-12-01T05:07:10.757'],
         ['2019-12-01T05:58:20.495', '2019-12-01T06:49:22.010'],
         ['2019-12-01T07:40:31.748', '2019-12-01T08:31:33.263'],
         ['2019-12-02T00:42:24.278', '2019-12-02T01:33:25.793']],
        dtype='datetime64[ms]',
    )
    handler = PeriodList(periods.T)

    dates = numpy.arange(numpy.datetime64('2019-12-02T00:42', 'ms'),
                         numpy.datetime64('2019-12-02T01:40', 'ms'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    indices = numpy.where(~flags)
    assert n.all(
        dates[indices] > numpy.datetime64('2019-12-02T01:33:25.794', 'ms'))

    dates = numpy.arange(numpy.datetime64('2019-12-02T00:42', 'ms'),
                         numpy.datetime64('2019-12-02T01:33', 'ms'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    assert n.all(flags)

    dates = numpy.arange(numpy.datetime64('2019-12-01T08:32', 'ms'),
                         numpy.datetime64('2019-12-02T00:40', 'ms'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    assert n.all(~flags)

    dates = numpy.arange(numpy.datetime64('2019-12-01T08:30', 'ms'),
                         numpy.datetime64('2019-12-02T00:40', 'us'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    assert not n.all(~flags)
    index = numpy.max(numpy.where(flags)[0])
    assert dates[index] > numpy.datetime64('2019-12-01T08:30')

    other = pickle.loads(pickle.dumps(handler))
    assert n.all(other.belong_to_a_period(dates) == flags)

    dates = numpy.arange(numpy.datetime64('2019-12-01T03:30:00.000', 'ms'),
                         numpy.datetime64('2019-12-01T04:00:00.000', 'ms'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    assert n.all(~flags)

    dates = numpy.arange(numpy.datetime64('2019-12-01T03:20:00.000', 'ms'),
                         numpy.datetime64('2019-12-01T04:00:00.000', 'ms'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    assert dates[~flags][0] > numpy.datetime64('2019-12-01T03:24:59.504', 'ms')
//  [python]
def test_period_list_cross_a_period():
    """Tests the PeriodList.cross_a_period method."""
    periods = numpy.array(
        [['2019-12-01T02:33:57.989', '2019-12-01T03:24:59.504'],
         ['2019-12-01T04:16:09.242', '2019-12-01T05:07:10.757'],
         ['2019-12-01T05:58:20.495', '2019-12-01T06:49:22.010'],
         ['2019-12-01T07:40:31.748', '2019-12-01T08:31:33.263'],
         ['2019-12-02T00:42:24.278', '2019-12-02T01:33:25.793']],
        dtype='datetime64[ms]',
    )
    handler = PeriodList(periods.T)

    dates = numpy.arange(numpy.datetime64('2019-12-02T00:42', 'ms'),
                         numpy.datetime64('2019-12-02T01:40', 'ms'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    indices = numpy.where(~flags)
    assert numpy.all(
        dates[indices] > numpy.datetime64('2019-12-02T01:33:25.794', 'ms'))

    dates = numpy.arange(numpy.datetime64('2019-12-02T00:42', 'ms'),
                         numpy.datetime64('2019-12-02T01:33', 'ms'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    assert numpy.all(flags)

    dates = numpy.arange(numpy.datetime64('2019-12-01T08:32', 'ms'),
                         numpy.datetime64('2019-12-02T00:40', 'ms'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    assert numpy.all(~flags)

    dates = numpy.arange(numpy.datetime64('2019-12-01T08:30', 'ms'),
                         numpy.datetime64('2019-12-02T00:40', 'us'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    assert not numpy.all(~flags)
    index = numpy.max(numpy.where(flags)[0])
    assert dates[index] > numpy.datetime64('2019-12-01T08:30')

    other = pickle.loads(pickle.dumps(handler))
    assert numpy.all(other.belong_to_a_period(dates) == flags)

    dates = numpy.arange(numpy.datetime64('2019-12-01T03:30:00.000', 'ms'),
                         numpy.datetime64('2019-12-01T04:00:00.000', 'ms'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    assert numpy.all(~flags)

    dates = numpy.arange(numpy.datetime64('2019-12-01T03:20:00.000', 'ms'),
                         numpy.datetime64('2019-12-01T04:00:00.000', 'ms'),
                         numpy.timedelta64(1, 's'))
    flags = handler.cross_a_period(dates)
    assert dates[~flags][0] > numpy.datetime64('2019-12-01T03:24:59.504', 'ms')
//  [python]
def _check_datasets_match(self, other: "Dataset") -> None:
        if self.n_items != other.n_items:
            raise ValueError(
                f"Number of items must match ({self.n_items} and {other.n_items})"
            )
        for j in range(self.n_items):
            if self.items[j].type != other.items[j].type:
                raise ValueError(
                    f"Item types must match. Item {j}: {self.items[j].type} != {other.items[j].type}"
                )
            if self.items[j].unit != other.items[j].unit:
                raise ValueError(
                    f"Item units must match. Item {j}: {self.items[j].unit} != {other.items[j].unit}"
                )
        if not np.all(self.time == other.time):
            raise ValueError("All timesteps must match")
        if self.shape != other.shape:
            raise ValueError("shape must match")
//  [python]
def _find_duplicated_spikes_random(spike_train: np.ndarray, censored_period: int, seed: int) -> np.ndarray:
    # random seed
    rng = np.random.RandomState(seed=seed)

    indices_of_duplicates = []
    while not np.all(np.diff(np.delete(spike_train, indices_of_duplicates)) > censored_period):
        (duplicates,) = np.where(np.diff(spike_train) <= censored_period)
        duplicates = np.unique(np.concatenate((duplicates, duplicates + 1)))
        duplicate = rng.choice(duplicates)
        indices_of_duplicates.append(duplicate)

    return np.array(indices_of_duplicates, dtype=np.int64)
//  [python]
def _concat_time(self, other: "Dataset", copy: bool = True) -> "Dataset":
        self._check_all_items_match(other)
        # assuming time is always first dimension we can skip / keep it by bool
        start_dim = int("time" in self.dims)
        if not n.all(
            self.shape[start_dim:] == other.shape[int("time" in other.dims) :]
        ):
            # if not np.all(self.shape[1:] == other.shape[1:]):
            raise ValueError("Shape of the datasets must match (except time dimension)")
        if hasattr(self, "time"):  # using attribute instead of dim checking. Works
            ds = self.copy() if copy else self
        else:
            raise ValueError(
                "Datasets cannot be concatenated as they have no time attribute!"
            )

        s1 = pd.Series(np.arange(len(ds.time)), index=ds.time, name="idx1")
        s2 = pd.Series(np.arange(len(other.time)), index=other.time, name="idx2")
        df12 = pd.concat([s1, s2], axis=1)

        newtime = df12.index
        newdata = self.create_empty_data(
            n_items=ds.n_items, n_timesteps=len(newtime), shape=ds.shape[start_dim:]
        )
        idx1 = np.where(~df12["idx1"].isna())
        idx2 = np.where(~df12["idx2"].isna())
        for j in range(ds.n_items):
            #    # if there is an overlap "other" data will be used!
            newdata[j][idx1] = ds[j].to_numpy()
            newdata[j][idx2] = other[j].to_numpy()

        zn = None
        if self._zn is not None:
            zshape = (len(newtime), self._zn.shape[start_dim])
            zn = np.zeros(shape=zshape, dtype=self._zn.dtype)
            zn[idx1, :] = self._zn
            zn[idx2, :] = other._zn

        return Dataset(
            newdata, time=newtime, items=ds.items, geometry=ds.geometry, zn=zn
        )
//  [python]
def _concat_time(self, other: "Dataset", copy: bool = True) -> "Dataset":
        self._check_all_items_match(other)
        # assuming time is always first dimension we can skip / keep it by bool
        start_dim = int("time" in self.dims)
        if not np.all(
            self.shape[start_dim:] == other.shape[int("time" in other.dims) :]
        ):
            # if not np.all(self.shape[1:] == other.shape[1:]):
            raise ValueError("Shape of the datasets must match (except time dimension)")
        if hasattr(self, "time"):  # using attribute instead of dim checking. Works
            ds = self.copy() if copy else self
        else:
            raise ValueError(
                "Datasets cannot be concatenated as they have no time attribute!"
            )

        s1 = pd.Series(np.arange(len(ds.time)), index=ds.time, name="idx1")
        s2 = pd.Series(np.arange(len(other.time)), index=other.time, name="idx2")
        df12 = pd.concat([s1, s2], axis=1)

        newtime = df12.index
        newdata = self.create_empty_data(
            n_items=ds.n_items, n_timesteps=len(newtime), shape=ds.shape[start_dim:]
        )
        idx1 = np.where(~df12["idx1"].isna())
        idx2 = np.where(~df12["idx2"].isna())
        for j in range(ds.n_items):
            #    # if there is an overlap "other" data will be used!
            newdata[j][idx1] = ds[j].to_numpy()
            newdata[j][idx2] = other[j].to_numpy()

        zn = None
        if self._zn is not None:
            zshape = (len(newtime), self._zn.shape[start_dim])
            zn = np.zeros(shape=zshape, dtype=self._zn.dtype)
            zn[idx1, :] = self._zn
            zn[idx2, :] = other._zn

        return Dataset(
            newdata, time=newtime, items=ds.items, geometry=ds.geometry, zn=zn
        )
//  [python]
def _union(X, Y): #TODO: add test for _union
    """
    Wrapper for np.vstack and pd.vstack that returns unique locations

    Parameters
    ----------
    X : numpy array or pandas dataframe with electrode locations

    Y : numpy array or pandas dataframe with electrode locations

    Returns
    ----------
    XY: the unique values of X and Y stacked together in the same format.  If either X or Y is a dataframe, then XY
        is a dataframe with the same columsn.  Otherwise XY is a numpy array.
    """

    if X is None:
        return Y
    elif Y is None:
        return X

    dataframeX = type(X) is pd.DataFrame
    if dataframeX:
        columnsX = X.columns
        X = X.values

    dataframeY = type(Y) is pd.DataFrame
    if dataframeY:
        columnsY = Y.columns
        Y = Y.values

    if dataframeX and dataframeY:
        assert n.all(columnsX == columnsY), 'Input dataframes have mismatched columns'

    assert X.shape[1] == Y.shape[1], 'Input data must have the same number of columns'

    XY = np.vstack((X, Y))
    XY_unique, tmp = _unique(XY)

    if dataframeX or dataframeY:
        if dataframeX:
            columns = columnsX
        else:
            columns = columnsY
        return pd.DataFrame(data=XY_unique, columns=columns, index=np.arange(XY_unique.shape[0]))
    else:
        return XY
//  [python]
def _union(X, Y): #TODO: add test for _union
    """
    Wrapper for np.vstack and pd.vstack that returns unique locations

    Parameters
    ----------
    X : numpy array or pandas dataframe with electrode locations

    Y : numpy array or pandas dataframe with electrode locations

    Returns
    ----------
    XY: the unique values of X and Y stacked together in the same format.  If either X or Y is a dataframe, then XY
        is a dataframe with the same columsn.  Otherwise XY is a numpy array.
    """

    if X is None:
        return Y
    elif Y is None:
        return X

    dataframeX = type(X) is pd.DataFrame
    if dataframeX:
        columnsX = X.columns
        X = X.values

    dataframeY = type(Y) is pd.DataFrame
    if dataframeY:
        columnsY = Y.columns
        Y = Y.values

    if dataframeX and dataframeY:
        assert np.all(columnsX == columnsY), 'Input dataframes have mismatched columns'

    assert X.shape[1] == Y.shape[1], 'Input data must have the same number of columns'

    XY = np.vstack((X, Y))
    XY_unique, tmp = _unique(XY)

    if dataframeX or dataframeY:
        if dataframeX:
            columns = columnsX
        else:
            columns = columnsY
        return pd.DataFrame(data=XY_unique, columns=columns, index=np.arange(XY_unique.shape[0]))
    else:
        return XY
//  [python]
def _reduce_qconv_shared_weights(
    model: ModelProto, orig_qconv_weight_name_to_node_ids: Dict[str, List[NodeProto]]
):
    graph = ONNXGraph(model)
    for weight_name, node_ids in orig_qconv_weight_name_to_node_ids.items():
        if len(node_ids) < 2:
            continue

        qconv_nodes = [graph.get_node_by_output_id(id) for id in node_ids]
        if any(
            node.op_type not in ["QLinearConv", "ConvInteger"] for node in qconv_nodes
        ):
            continue

        weights = [
            graph.get_init_by_name(
                node.input[3 if node.op_type == "QLinearConv" else 1]
            )
            for node in qconv_nodes
        ]
        if any(weight is None for weight in weights):
            continue

        weights = [numpy_helper.to_array(weight) for weight in weights]
        if not all(numpy.array_equal(weight, weights[0]) for weight in weights):
            continue

        shared_weight = numpy_helper.from_array(
            weights[0], name=f"{weight_name}_quantized"
        )
        for node in qconv_nodes:
            target_dim = 3 if node.op_type == "QLinearConv" else 1
            node.input[target_dim] = shared_weight.name
        model.graph.initializer.append(shared_weight)

    graph.update()
    graph.delete_unused_initializers()
//  [python]
def _reduce_qconv_shared_weights(
    model: ModelProto, orig_qconv_weight_name_to_node_ids: Dict[str, List[NodeProto]]
):
    graph = ONNXGraph(model)
    for weight_name, node_ids in orig_qconv_weight_name_to_node_ids.items():
        if len(node_ids) < 2:
            continue

        qconv_nodes = [graph.get_node_by_output_id(id) for id in node_ids]
        if any(
            node.op_type not in ["QLinearConv", "ConvInteger"] for node in qconv_nodes
        ):
            continue

        weights = [
            graph.get_init_by_name(
                node.input[3 if node.op_type == "QLinearConv" else 1]
            )
            for node in qconv_nodes
        ]
        if any(weight is None for weight in weights):
            continue

        weights = [numpy_helper.to_array(weight) for weight in weights]
        if not all(numpy.all(weight == weights[0]) for weight in weights):
            continue

        shared_weight = numpy_helper.from_array(
            weights[0], name=f"{weight_name}_quantized"
        )
        for node in qconv_nodes:
            target_dim = 3 if node.op_type == "QLinearConv" else 1
            node.input[target_dim] = shared_weight.name
        model.graph.initializer.append(shared_weight)

    graph.update()
    graph.delete_unused_initializers()
//  [python]
def test_dfsetup_surfaceSigmaProfile():
    df = dehnendf(profileParams=(0.25, 0.75, 0.1), beta=0.0, correct=False)
    from galpy.df import expSurfaceSigmaProfile

    essp = expSurfaceSigmaProfile(params=(0.25, 0.75, 0.1))
    df_alt = dehnendf(surfaceSigma=essp, beta=0.0, correct=False)
    assert n.all(
        n.fabs(
            n.array(df._surfaceSigmaProfile._params)
            - n.array(df_alt._surfaceSigmaProfile._params)
        )
        < 10.0**-10.0
    ), "diskdf setup with explicit surfaceSigmaProfile class does not give the same profile as with parameters only"
    return None
//  [python]
def test_dfsetup_surfaceSigmaProfile():
    df = dehnendf(profileParams=(0.25, 0.75, 0.1), beta=0.0, correct=False)
    from galpy.df import expSurfaceSigmaProfile

    essp = expSurfaceSigmaProfile(params=(0.25, 0.75, 0.1))
    df_alt = dehnendf(surfaceSigma=essp, beta=0.0, correct=False)
    assert numpy.all(
        numpy.fabs(
            numpy.array(df._surfaceSigmaProfile._params)
            - numpy.array(df_alt._surfaceSigmaProfile._params)
        )
        < 10.0**-10.0
    ), "diskdf setup with explicit surfaceSigmaProfile class does not give the same profile as with parameters only"
    return None
//  [python]
def _to_own_unit(self, value, check_precision=True):
        try:
            _value = value.to_value(self.unit)
        except AttributeError:
            # We're not a Quantity, so let's try a more general conversion.
            # Plain arrays will be converted to dimensionless in the process,
            # but anything with a unit attribute will use that.
            try:
                as_quantity = Quantity(value)
                _value = as_quantity.to_value(self.unit)
            except TypeError:
                # Could not make a Quantity.  Maybe masked printing?
                # Note: masked quantities do not work very well, but no reason
                # to break even repr and str.
                if (value is np.ma.masked_print_option and
                        self.dtype.kind == 'O'):
                    return value
                else:
                    raise

            except UnitsError:
                # last chance: if this was not something with a unit
                # and is all 0, inf, or nan, we treat it as arbitrary unit.
                if (not hasattr(value, 'unit') and
                        can_have_arbitrary_unit(as_quantity.value)):
                    _value = as_quantity.value
                else:
                    raise

        if check_precision:
            # If, e.g., we are casting double to float, we want to fail if
            # precision is lost, but let things pass if it works.
            _value = np.array(_value, copy=False)
            if not np.can_cast(_value.dtype, self.dtype):
                self_dtype_array = np.array(_value, self.dtype)
                if not np.alltrue(np.logical_or(self_dtype_array == _value,
                                            np.isnan(_value))):
                    raise TypeError("cannot convert value type to array type "
                                    "without precision loss")
        return _value
//  [python]
def _to_own_unit(self, value, check_precision=True):
        try:
            _value = value.to_value(self.unit)
        except AttributeError:
            # We're not a Quantity, so let's try a more general conversion.
            # Plain arrays will be converted to dimensionless in the process,
            # but anything with a unit attribute will use that.
            try:
                as_quantity = Quantity(value)
                _value = as_quantity.to_value(self.unit)
            except TypeError:
                # Could not make a Quantity.  Maybe masked printing?
                # Note: masked quantities do not work very well, but no reason
                # to break even repr and str.
                if (value is np.ma.masked_print_option and
                        self.dtype.kind == 'O'):
                    return value
                else:
                    raise

            except UnitsError:
                # last chance: if this was not something with a unit
                # and is all 0, inf, or nan, we treat it as arbitrary unit.
                if (not hasattr(value, 'unit') and
                        can_have_arbitrary_unit(as_quantity.value)):
                    _value = as_quantity.value
                else:
                    raise

        if check_precision:
            # If, e.g., we are casting double to float, we want to fail if
            # precision is lost, but let things pass if it works.
            _value = np.array(_value, copy=False)
            if not np.can_cast(_value.dtype, self.dtype):
                self_dtype_array = np.array(_value, self.dtype)
                if not np.all(np.logical_or(self_dtype_array == _value,
                                            np.isnan(_value))):
                    raise TypeError("cannot convert value type to array type "
                                    "without precision loss")
        return _value
//  [python]
def after_epoch(self, epoch):
    """Updates probe_images and probe_labels dictionaries."""
    # Get top-k ranking per class
    if not np.alltrue(self.ds_labels >= 0):
      t = np.sum(self.ds_labels < 0)
      logging.warning('ds_labels has {} unfilled items'.format(t))

    best_ids = self._get_top_indices()
    images, labels, hit = [], [], []
    ds_next_example = self.ds_iter.get_next()
    self.sess.run(self.ds_iter.initializer)
    for _ in range(int(np.ceil(self.dataset_size / self.batch_size))):
      # look through the dataset to get data with best_ids
      try:
        b_imgs, b_labels, b_ids, b_gold_labels = self.sess.run(ds_next_example)
      except tf.errors.OutOfRangeError:
        break
      keep_ids = np.intersect1d(b_ids, best_ids, return_indices=True)[1]

      if np.size(keep_ids) > 0:
        images.append(b_imgs[keep_ids])
        labels.append(b_labels[keep_ids])
        # Orcal used only for monitoring
        hit.append(b_labels[keep_ids] == b_gold_labels[keep_ids])
    images = np.concatenate(images)
    labels = np.concatenate(labels)
    hit = np.concatenate(hit).astype(np.float32)
    mean_hit = np.mean(hit)
    if images.shape[0] != self.capacity:
      logging.warning('{} != {}'.format(images.shape[0], self.capacity))

    # Put to queue
    self._enqueue(images, labels)

    max_score = self.ds_rates.max()
    min_score = self.ds_rates.min()
    self.purity_log.load(mean_hit)
    self.capacity_log.load(self.capacity)
    if best_ids.shape[0] >= self.init_capacity:
      # To prevent best_ids has shorter size and makes error.
      self.queue_data_indices.load(best_ids[:self.init_capacity])
    ratio = np.mean(self.ds_labels >= 0)

    logging.info(
        'Queue info at epoch {}: '
        'purity:{:.2f} ratio:{:.1f} cap:{} score(max/min):{:.2f}/{:.2f}'.format(
            epoch, mean_hit, ratio * 100., self.capacity, max_score, min_score))
//  [python]
def after_epoch(self, epoch):
    """Updates probe_images and probe_labels dictionaries."""
    # Get top-k ranking per class
    if not np.all(self.ds_labels >= 0):
      t = np.sum(self.ds_labels < 0)
      logging.warning('ds_labels has {} unfilled items'.format(t))

    best_ids = self._get_top_indices()
    images, labels, hit = [], [], []
    ds_next_example = self.ds_iter.get_next()
    self.sess.run(self.ds_iter.initializer)
    for _ in range(int(np.ceil(self.dataset_size / self.batch_size))):
      # look through the dataset to get data with best_ids
      try:
        b_imgs, b_labels, b_ids, b_gold_labels = self.sess.run(ds_next_example)
      except tf.errors.OutOfRangeError:
        break
      keep_ids = np.intersect1d(b_ids, best_ids, return_indices=True)[1]

      if np.size(keep_ids) > 0:
        images.append(b_imgs[keep_ids])
        labels.append(b_labels[keep_ids])
        # Orcal used only for monitoring
        hit.append(b_labels[keep_ids] == b_gold_labels[keep_ids])
    images = np.concatenate(images)
    labels = np.concatenate(labels)
    hit = np.concatenate(hit).astype(np.float32)
    mean_hit = np.mean(hit)
    if images.shape[0] != self.capacity:
      logging.warning('{} != {}'.format(images.shape[0], self.capacity))

    # Put to queue
    self._enqueue(images, labels)

    max_score = self.ds_rates.max()
    min_score = self.ds_rates.min()
    self.purity_log.load(mean_hit)
    self.capacity_log.load(self.capacity)
    if best_ids.shape[0] >= self.init_capacity:
      # To prevent best_ids has shorter size and makes error.
      self.queue_data_indices.load(best_ids[:self.init_capacity])
    ratio = np.mean(self.ds_labels >= 0)

    logging.info(
        'Queue info at epoch {}: '
        'purity:{:.2f} ratio:{:.1f} cap:{} score(max/min):{:.2f}/{:.2f}'.format(
            epoch, mean_hit, ratio * 100., self.capacity, max_score, min_score))
//  [python]
def setupOutputs(self):
        assert self.LabelImage.meta.axistags == self.RawImage.meta.axistags

        taggedOutputShape = self.LabelImage.meta.getTaggedShape()
        taggedRawShape = self.RawImage.meta.getTaggedShape()

        if not all(list(taggedOutputShape.get(k, 0) == taggedRawShape.get(k, 0) for k in "txyz")):
            raise DatasetConstraintError(
                "Object Extraction",
                "Raw Image and Label Image shapes do not match.\n"
                "Label Image shape: {}. Raw Image shape: {}"
                "".format(self.LabelImage.meta.shape, self.RawVolume.meta.shape),
            )

        # Every value in the regionfeatures output is cached seperately as it's own "block"
        blockshape = (1,) * len(self._opRegionFeatures.Output.meta.shape)
        self._opCache.BlockShape.setValue(blockshape)
//  [python]
def setupOutputs(self):
        assert self.LabelImage.meta.axistags == self.RawImage.meta.axistags

        taggedOutputShape = self.LabelImage.meta.getTaggedShape()
        taggedRawShape = self.RawImage.meta.getTaggedShape()

        if not numpy.all(list(taggedOutputShape.get(k, 0) == taggedRawShape.get(k, 0) for k in "txyz")):
            raise DatasetConstraintError(
                "Object Extraction",
                "Raw Image and Label Image shapes do not match.\n"
                "Label Image shape: {}. Raw Image shape: {}"
                "".format(self.LabelImage.meta.shape, self.RawVolume.meta.shape),
            )

        # Every value in the regionfeatures output is cached seperately as it's own "block"
        blockshape = (1,) * len(self._opRegionFeatures.Output.meta.shape)
        self._opCache.BlockShape.setValue(blockshape)
//  [python]
def analyze_wfnsym(tdobj, x_sym, x):
    '''Guess the wfn symmetry of TDDFT X amplitude'''
    possible_sym = x_sym[(x > 1e-7) | (x < -1e-7)]
    if n.all(possible_sym == symm.MULTI_IRREPS):
        if tdobj.wfnsym is None:
            wfnsym = '???'
        else:
            wfnsym = tdobj.wfnsym
    else:
        ids = possible_sym[possible_sym != symm.MULTI_IRREPS]
        ids = numpy.unique(ids)
        if ids.size == 1:
            wfnsym = symm.irrep_id2name(tdobj.mol.groupname, ids[0])
        else:
            wfnsym = '???'
    return wfnsym
//  [python]
def analyze_wfnsym(tdobj, x_sym, x):
    '''Guess the wfn symmetry of TDDFT X amplitude'''
    possible_sym = x_sym[(x > 1e-7) | (x < -1e-7)]
    if numpy.all(possible_sym == symm.MULTI_IRREPS):
        if tdobj.wfnsym is None:
            wfnsym = '???'
        else:
            wfnsym = tdobj.wfnsym
    else:
        ids = possible_sym[possible_sym != symm.MULTI_IRREPS]
        ids = numpy.unique(ids)
        if ids.size == 1:
            wfnsym = symm.irrep_id2name(tdobj.mol.groupname, ids[0])
        else:
            wfnsym = '???'
    return wfnsym
//  [python]
def test_dehnendf_sample_sampleLOS():
    # Test that the samples returned through sample with los are the same as those returned with sampleLOS
    beta = 0.0
    dfc = dehnendf(beta=beta, profileParams=(1.0 / 4.0, 1.0, 0.2))
    # Sample a large number of points, then check some moments against the analytic distribution
    numpy.random.seed(1)
    os = dfc.sampleLOS(45.0, n=2, targetSurfmass=False, deg=True)
    rs = numpy.array([o.R() for o in os])
    vrs = numpy.array([o.vR() for o in os])
    vts = numpy.array([o.vT() for o in os])
    numpy.random.seed(1)
    os2 = dfc.sample(los=45.0, n=2, targetSurfmass=False, losdeg=True)
    rs2 = numpy.array([o.R() for o in os2])
    vrs2 = numpy.array([o.vR() for o in os2])
    vts2 = numpy.array([o.vT() for o in os2])
    assert n.all(
        n.fabs(rs - rs2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    assert n.all(
        n.fabs(vrs - vrs2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    assert n.all(
        n.fabs(vts - vts2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    return None
//  [python]
def test_dehnendf_sample_sampleLOS():
    # Test that the samples returned through sample with los are the same as those returned with sampleLOS
    beta = 0.0
    dfc = dehnendf(beta=beta, profileParams=(1.0 / 4.0, 1.0, 0.2))
    # Sample a large number of points, then check some moments against the analytic distribution
    numpy.random.seed(1)
    os = dfc.sampleLOS(45.0, n=2, targetSurfmass=False, deg=True)
    rs = numpy.array([o.R() for o in os])
    vrs = numpy.array([o.vR() for o in os])
    vts = numpy.array([o.vT() for o in os])
    numpy.random.seed(1)
    os2 = dfc.sample(los=45.0, n=2, targetSurfmass=False, losdeg=True)
    rs2 = numpy.array([o.R() for o in os2])
    vrs2 = numpy.array([o.vR() for o in os2])
    vts2 = numpy.array([o.vT() for o in os2])
    assert numpy.all(
        numpy.fabs(rs - rs2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    assert numpy.all(
        numpy.fabs(vrs - vrs2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    assert numpy.all(
        numpy.fabs(vts - vts2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    return None
//  [python]
def test_shudf_sample_sampleLOS():
    # Test that the samples returned through sample with los are the same as those returned with sampleLOS
    beta = 0.0
    dfc = shudf(beta=beta, profileParams=(1.0 / 4.0, 1.0, 0.2))
    # Sample a large number of points, then check some moments against the analytic distribution
    numpy.random.seed(1)
    os = dfc.sampleLOS(45.0, n=2, targetSurfmass=False, deg=True)
    rs = numpy.array([o.R() for o in os])
    vrs = numpy.array([o.vR() for o in os])
    vts = numpy.array([o.vT() for o in os])
    numpy.random.seed(1)
    os2 = dfc.sample(los=45.0, n=2, targetSurfmass=False, losdeg=True)
    rs2 = numpy.array([o.R() for o in os2])
    vrs2 = numpy.array([o.vR() for o in os2])
    vts2 = numpy.array([o.vT() for o in os2])
    assert n.all(
        n.fabs(rs - rs2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    assert n.all(
        n.fabs(vrs - vrs2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    assert n.all(
        n.fabs(vts - vts2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    return None
//  [python]
def test_shudf_sample_sampleLOS():
    # Test that the samples returned through sample with los are the same as those returned with sampleLOS
    beta = 0.0
    dfc = shudf(beta=beta, profileParams=(1.0 / 4.0, 1.0, 0.2))
    # Sample a large number of points, then check some moments against the analytic distribution
    numpy.random.seed(1)
    os = dfc.sampleLOS(45.0, n=2, targetSurfmass=False, deg=True)
    rs = numpy.array([o.R() for o in os])
    vrs = numpy.array([o.vR() for o in os])
    vts = numpy.array([o.vT() for o in os])
    numpy.random.seed(1)
    os2 = dfc.sample(los=45.0, n=2, targetSurfmass=False, losdeg=True)
    rs2 = numpy.array([o.R() for o in os2])
    vrs2 = numpy.array([o.vR() for o in os2])
    vts2 = numpy.array([o.vT() for o in os2])
    assert numpy.all(
        numpy.fabs(rs - rs2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    assert numpy.all(
        numpy.fabs(vrs - vrs2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    assert numpy.all(
        numpy.fabs(vts - vts2) < 10.0**-10.0
    ), "Samples returned from dehnendf.sample with los set are not the same as those returned with sampleLOS"
    return None
//  [python]
def test_dehnendf_flat_DFcorrection_mag():
    # Test that the call is not too different from before
    tcorr = ddf_correct2_flat._corr.correct(1.1, log=True)
    assert numpy.fabs(tcorr[0]) < 0.15, "dehnendf correction is larger than expected"
    assert numpy.fabs(tcorr[1]) < 0.1, "dehnendf correction is larger than expected"
    # small R
    tcorr = numpy.log(ddf_correct2_flat._corr.correct(10.0**-12.0))
    assert numpy.fabs(tcorr[0]) < 0.4, "dehnendf correction is larger than expected"
    assert numpy.fabs(tcorr[1]) < 1.0, "dehnendf correction is larger than expected"
    # large R
    tcorr = numpy.log(ddf_correct2_flat._corr.correct(12.0))
    assert numpy.fabs(tcorr[0]) < 0.01, "dehnendf correction is larger than expected"
    assert numpy.fabs(tcorr[1]) < 0.01, "dehnendf correction is larger than expected"
    # small R, array
    tcorr = numpy.log(ddf_correct2_flat._corr.correct(10.0**-12.0 * numpy.ones(2)))
    assert numpy.alltrue(
        numpy.fabs(tcorr[0]) < 0.4
    ), "dehnendf correction is larger than expected"
    assert numpy.alltrue(
        numpy.fabs(tcorr[1]) < 1.0
    ), "dehnendf correction is larger than expected"
    # large R
    tcorr = numpy.log(ddf_correct2_flat._corr.correct(12.0 * numpy.ones(2)))
    assert numpy.alltrue(
        numpy.fabs(tcorr[0]) < 0.01
    ), "dehnendf correction is larger than expected"
    assert numpy.alltrue(
        numpy.fabs(tcorr[1]) < 0.01
    ), "dehnendf correction is larger than expected"
    # large R, log
    tcorr = ddf_correct2_flat._corr.correct(12.0 * numpy.ones(2), log=True)
    assert numpy.alltrue(
        numpy.fabs(tcorr[0]) < 0.01
    ), "dehnendf correction is larger than expected"
    assert numpy.alltrue(
        numpy.fabs(tcorr[1]) < 0.01
    ), "dehnendf correction is larger than expected"
    return None
//  [python]
def test_dehnendf_flat_DFcorrection_mag():
    # Test that the call is not too different from before
    tcorr = ddf_correct2_flat._corr.correct(1.1, log=True)
    assert numpy.fabs(tcorr[0]) < 0.15, "dehnendf correction is larger than expected"
    assert numpy.fabs(tcorr[1]) < 0.1, "dehnendf correction is larger than expected"
    # small R
    tcorr = numpy.log(ddf_correct2_flat._corr.correct(10.0**-12.0))
    assert numpy.fabs(tcorr[0]) < 0.4, "dehnendf correction is larger than expected"
    assert numpy.fabs(tcorr[1]) < 1.0, "dehnendf correction is larger than expected"
    # large R
    tcorr = numpy.log(ddf_correct2_flat._corr.correct(12.0))
    assert numpy.fabs(tcorr[0]) < 0.01, "dehnendf correction is larger than expected"
    assert numpy.fabs(tcorr[1]) < 0.01, "dehnendf correction is larger than expected"
    # small R, array
    tcorr = numpy.log(ddf_correct2_flat._corr.correct(10.0**-12.0 * numpy.ones(2)))
    assert numpy.all(
        numpy.fabs(tcorr[0]) < 0.4
    ), "dehnendf correction is larger than expected"
    assert numpy.all(
        numpy.fabs(tcorr[1]) < 1.0
    ), "dehnendf correction is larger than expected"
    # large R
    tcorr = numpy.log(ddf_correct2_flat._corr.correct(12.0 * numpy.ones(2)))
    assert numpy.all(
        numpy.fabs(tcorr[0]) < 0.01
    ), "dehnendf correction is larger than expected"
    assert numpy.all(
        numpy.fabs(tcorr[1]) < 0.01
    ), "dehnendf correction is larger than expected"
    # large R, log
    tcorr = ddf_correct2_flat._corr.correct(12.0 * numpy.ones(2), log=True)
    assert numpy.all(
        numpy.fabs(tcorr[0]) < 0.01
    ), "dehnendf correction is larger than expected"
    assert numpy.all(
        numpy.fabs(tcorr[1]) < 0.01
    ), "dehnendf correction is larger than expected"
    return None
//  [python]
def test_binary_convex_hull():
    #
    # Set the four points of a square to True
    #
    image = numpy.zeros((20, 15), bool)
    image[2, 3] = True
    image[17, 3] = True
    image[2, 12] = True
    image[17, 12] = True
    expected = numpy.zeros((20, 15), bool)
    expected[2:18, 3:13] = True
    result = execute(image, "convex hull")
    assert n.all(result == expected)
//  [python]
def test_binary_convex_hull():
    #
    # Set the four points of a square to True
    #
    image = numpy.zeros((20, 15), bool)
    image[2, 3] = True
    image[17, 3] = True
    image[2, 12] = True
    image[17, 12] = True
    expected = numpy.zeros((20, 15), bool)
    expected[2:18, 3:13] = True
    result = execute(image, "convex hull")
    assert numpy.all(result == expected)
//  [python]
def setupOutputs(self):
        if self.LabelVolume.meta.axistags != self.RawVolume.meta.axistags:
            raise Exception("raw and label axis tags do not match")

        taggedOutputShape = self.LabelVolume.meta.getTaggedShape()
        taggedRawShape = self.RawVolume.meta.getTaggedShape()

        if not all(list(taggedOutputShape.get(k, 0) == taggedRawShape.get(k, 0) for k in "txyz")):
            raise Exception(
                "shapes do not match. label volume shape: {}."
                " raw data shape: {}".format(self.LabelVolume.meta.shape, self.RawVolume.meta.shape)
            )

        self.Output.meta.shape = (taggedOutputShape["t"],)
        self.Output.meta.axistags = vigra.defaultAxistags("t")
        # The features for the entire block (in xyz) are provided for the requested tc coordinates.
        self.Output.meta.dtype = object
//  [python]
def setupOutputs(self):
        if self.LabelVolume.meta.axistags != self.RawVolume.meta.axistags:
            raise Exception("raw and label axis tags do not match")

        taggedOutputShape = self.LabelVolume.meta.getTaggedShape()
        taggedRawShape = self.RawVolume.meta.getTaggedShape()

        if not numpy.all(list(taggedOutputShape.get(k, 0) == taggedRawShape.get(k, 0) for k in "txyz")):
            raise Exception(
                "shapes do not match. label volume shape: {}."
                " raw data shape: {}".format(self.LabelVolume.meta.shape, self.RawVolume.meta.shape)
            )

        self.Output.meta.shape = (taggedOutputShape["t"],)
        self.Output.meta.axistags = vigra.defaultAxistags("t")
        # The features for the entire block (in xyz) are provided for the requested tc coordinates.
        self.Output.meta.dtype = object
//  [python]
def binary_tteesstt(
    function_name, function, gray_out=False, scale=None, custom_repeats=None
):
    numpy.random.seed(list(map(ord, function_name)))
    input = numpy.random.uniform(size=(20, 20)) > 0.7
    output = execute(input, function_name, scale=scale, custom_repeats=custom_repeats)
    if scale is None:
        expected = function(input)
    else:
        footprint = centrosome.cpmorphology.strel_disk(float(scale) / 2.0)
        expected = function(input, footprint=footprint)
    if not gray_out:
        expected = expected > 0
        assert numpy.alltrue(output == expected)
    else:
        assert numpy.alltrue(numpy.abs(output - expected) < numpy.finfo(numpy.float32).eps)
//  [python]
def binary_tteesstt(
    function_name, function, gray_out=False, scale=None, custom_repeats=None
):
    numpy.random.seed(list(map(ord, function_name)))
    input = numpy.random.uniform(size=(20, 20)) > 0.7
    output = execute(input, function_name, scale=scale, custom_repeats=custom_repeats)
    if scale is None:
        expected = function(input)
    else:
        footprint = centrosome.cpmorphology.strel_disk(float(scale) / 2.0)
        expected = function(input, footprint=footprint)
    if not gray_out:
        expected = expected > 0
        assert numpy.all(output == expected)
    else:
        assert numpy.all(numpy.abs(output - expected) < numpy.finfo(numpy.float32).eps)
//  [python]
def test_nans_like():
    shapeInput = np.random.randint(
        20,
        100,
        [
            2,
        ],
    )
    shape = NumCpp.Shape(shapeInput[0].item(), shapeInput[1].item())
    cArray1 = NumCpp.NdArray(shape)
    cArray2 = NumCpp.nans_like(cArray1)
    assert (
        cArray2.shape().rows == shape.rows
        and cArray2.shape().cols == shape.cols
        and cArray2.size() == shapeInput.prod()
        and np.isnan(cArray2.getNumpyArray()).all()
    )