def _getCellsWithFewestSegments(cls, connections, rng, columns,
                                  cellsPerColumn):
    """
    For each column, get the cell that has the fewest total basal segments.
    Break ties randomly.

    @param connections (SparseMatrixConnections)
    @param rng (Random)
    @param columns (numpy array) Columns to check

    @return (numpy array)
    One cell for each of the provided columns
    """
    candidateCells = np2.getAllCellsInColumns(columns, cellsPerColumn)

    # Arrange the segment counts into one row per minicolumn.
    segmentCounts = np.reshape(connections.getSegmentCounts(candidateCells),
                               newshape=(len(columns),
                                         cellsPerColumn))

    # Filter to just the cells that are tied for fewest in their minicolumn.
    minSegmentCounts = np.amin(segmentCounts, axis=1, keepdims=True)
    candidateCells = candidateCells[np.flatnonzero(segmentCounts ==
                                                   minSegmentCounts)]

    # Filter to one cell per column, choosing randomly from the minimums.
    # To do the random choice, add a random offset to each index in-place, using
    # casting to floor the result.
    (_,
     onePerColumnFilter,
     numCandidatesInColumns) = np.unique(candidateCells / cellsPerColumn,
                                         return_index=True, return_counts=True)

    offsetPercents = np.empty(len(columns), dtype="float32")
    rng.initializeReal32Array(offsetPercents)

    np.add(onePerColumnFilter,
           offsetPercents*numCandidatesInColumns,
           out=onePerColumnFilter,
           casting="unsafe")

    return candidateCells[onePerColumnFilter]
  def _getCellsWithFewestSegments(self, columns):
    """
    For each column, get the cell that has the fewest total segments (basal or
    apical). Break ties randomly.

    @param columns (numpy array)
    Columns to check

    @return (numpy array)
    One cell for each of the provided columns
    """
    candidateCells = np2.getAllCellsInColumns(columns, self.cellsPerColumn)

    # Arrange the segment counts into one row per minicolumn.
    segmentCounts = np.reshape(
      self.basalConnections.getSegmentCounts(candidateCells) +
      self.apicalConnections.getSegmentCounts(candidateCells),
      newshape=(len(columns),
                self.cellsPerColumn))

    # Filter to just the cells that are tied for fewest in their minicolumn.
    minSegmentCounts = np.amin(segmentCounts, axis=1, keepdims=True)
    candidateCells = candidateCells[np.flatnonzero(segmentCounts ==
                                                   minSegmentCounts)]

    # Filter to one cell per column, choosing randomly from the minimums.
    # To do the random choice, add a random offset to each index in-place, using
    # casting to floor the result.
    (_,
     onePerColumnFilter,
     numCandidatesInColumns) = np.unique(candidateCells / self.cellsPerColumn,
                                         return_index=True, return_counts=True)

    offsetPercents = np.empty(len(columns), dtype="float32")
    self.rng.initializeReal32Array(offsetPercents)

    np.add(onePerColumnFilter,
           offsetPercents*numCandidatesInColumns,
           out=onePerColumnFilter,
           casting="unsafe")

    return candidateCells[onePerColumnFilter]
  def compute(self,
              activeColumns,
              basalInput,
              basalGrowthCandidates,
              apicalInput=EMPTY_UINT_ARRAY,
              apicalGrowthCandidates=EMPTY_UINT_ARRAY,
              learn=True):
    """
    @param activeColumns (numpy array)
    @param basalInput (numpy array)
    @param basalGrowthCandidates (numpy array)
    @param apicalInput (numpy array)
    @param apicalGrowthCandidates (numpy array)
    @param learn (bool)
    """
    # Calculate predictions for this timestep
    (activeBasalSegments,
     matchingBasalSegments,
     basalPotentialOverlaps) = self._calculateSegmentActivity(
       self.basalConnections, basalInput, self.connectedPermanence,
       self.activationThreshold, self.minThreshold)

    (activeApicalSegments,
     matchingApicalSegments,
     apicalPotentialOverlaps) = self._calculateSegmentActivity(
       self.apicalConnections, apicalInput, self.connectedPermanence,
       self.activationThreshold, self.minThreshold)

    predictedCells = self._calculatePredictedCells(activeBasalSegments,
                                                   activeApicalSegments)

    # Calculate active cells
    (correctPredictedCells,
     burstingColumns) = np2.setCompare(predictedCells, activeColumns,
                                       predictedCells / self.cellsPerColumn,
                                       rightMinusLeft=True)
    newActiveCells = np.concatenate((correctPredictedCells,
                                     np2.getAllCellsInColumns(
                                       burstingColumns, self.cellsPerColumn)))

    # Calculate learning
    (learningActiveBasalSegments,
     learningMatchingBasalSegments,
     basalSegmentsToPunish,
     newBasalSegmentCells,
     learningCells) = self._calculateBasalLearning(
       activeColumns, burstingColumns, correctPredictedCells,
       activeBasalSegments, matchingBasalSegments, basalPotentialOverlaps)

    (learningActiveApicalSegments,
     learningMatchingApicalSegments,
     apicalSegmentsToPunish,
     newApicalSegmentCells) = self._calculateApicalLearning(
       learningCells, activeColumns, activeApicalSegments,
       matchingApicalSegments, apicalPotentialOverlaps)

    # Learn
    if learn:
      # Learn on existing segments
      for learningSegments in (learningActiveBasalSegments,
                               learningMatchingBasalSegments):
        self._learn(self.basalConnections, self.rng, learningSegments,
                    basalInput, basalGrowthCandidates, basalPotentialOverlaps,
                    self.initialPermanence, self.sampleSize,
                    self.permanenceIncrement, self.permanenceDecrement,
                    self.maxSynapsesPerSegment)

      for learningSegments in (learningActiveApicalSegments,
                               learningMatchingApicalSegments):

        self._learn(self.apicalConnections, self.rng, learningSegments,
                    apicalInput, apicalGrowthCandidates,
                    apicalPotentialOverlaps, self.initialPermanence,
                    self.sampleSize, self.permanenceIncrement,
                    self.permanenceDecrement, self.maxSynapsesPerSegment)

      # Punish incorrect predictions
      if self.predictedSegmentDecrement != 0.0:
        self.basalConnections.adjustActiveSynapses(
          basalSegmentsToPunish, basalInput, -self.predictedSegmentDecrement)
        self.apicalConnections.adjustActiveSynapses(
          apicalSegmentsToPunish, apicalInput, -self.predictedSegmentDecrement)

      # Grow new segments
      if len(basalGrowthCandidates) > 0:
        self._learnOnNewSegments(self.basalConnections, self.rng,
                                 newBasalSegmentCells, basalGrowthCandidates,
                                 self.initialPermanence, self.sampleSize,
                                 self.maxSynapsesPerSegment)

      if len(apicalGrowthCandidates) > 0:
        self._learnOnNewSegments(self.apicalConnections, self.rng,
                                 newApicalSegmentCells, apicalGrowthCandidates,
                                 self.initialPermanence, self.sampleSize,
                                 self.maxSynapsesPerSegment)

    # Save the results
    self.activeCells = newActiveCells
    self.winnerCells = learningCells
    self.prevPredictedCells = predictedCells
  def activateCells(self,
                    activeColumns,
                    basalReinforceCandidates,
                    apicalReinforceCandidates,
                    basalGrowthCandidates,
                    apicalGrowthCandidates,
                    learn=True):
    """
    Activate cells in the specified columns, using the result of the previous
    'depolarizeCells' as predictions. Then learn.

    @param activeColumns (numpy array)
    List of active columns

    @param basalReinforceCandidates (numpy array)
    List of bits that the active cells may reinforce basal synapses to.

    @param apicalReinforceCandidates (numpy array)
    List of bits that the active cells may reinforce apical synapses to.

    @param basalGrowthCandidates (numpy array or None)
    List of bits that the active cells may grow new basal synapses to.

    @param apicalGrowthCandidates (numpy array or None)
    List of bits that the active cells may grow new apical synapses to

    @param learn (bool)
    Whether to grow / reinforce / punish synapses
    """

    # Calculate active cells
    (correctPredictedCells,
     burstingColumns) = np2.setCompare(self.predictedCells, activeColumns,
                                       self.predictedCells / self.cellsPerColumn,
                                       rightMinusLeft=True)

    newActiveCells = np.concatenate((correctPredictedCells,
                                     np2.getAllCellsInColumns(
                                       burstingColumns, self.cellsPerColumn)))

    # Calculate learning
    (learningActiveBasalSegments,
     learningActiveApicalSegments,
     learningMatchingBasalSegments,
     learningMatchingApicalSegments,
     basalSegmentsToPunish,
     apicalSegmentsToPunish,
     newSegmentCells,
     learningCells) = self._calculateLearning(activeColumns,
                                              burstingColumns,
                                              correctPredictedCells,
                                              self.activeBasalSegments,
                                              self.activeApicalSegments,
                                              self.matchingBasalSegments,
                                              self.matchingApicalSegments,
                                              self.basalPotentialOverlaps,
                                              self.apicalPotentialOverlaps)

    if learn:
      # Learn on existing segments
      for learningSegments in (learningActiveBasalSegments,
                               learningMatchingBasalSegments):
        self._learn(self.basalConnections, self.rng, learningSegments,
                    basalReinforceCandidates, basalGrowthCandidates,
                    self.basalPotentialOverlaps,
                    self.initialPermanence, self.sampleSize,
                    self.permanenceIncrement, self.permanenceDecrement,
                    self.maxSynapsesPerSegment)

      for learningSegments in (learningActiveApicalSegments,
                               learningMatchingApicalSegments):
        self._learn(self.apicalConnections, self.rng, learningSegments,
                    apicalReinforceCandidates, apicalGrowthCandidates,
                    self.apicalPotentialOverlaps, self.initialPermanence,
                    self.sampleSize, self.permanenceIncrement,
                    self.permanenceDecrement, self.maxSynapsesPerSegment)

      # Punish incorrect predictions
      if self.basalPredictedSegmentDecrement != 0.0:
        self.basalConnections.adjustActiveSynapses(
          basalSegmentsToPunish, basalReinforceCandidates,
          -self.basalPredictedSegmentDecrement)

      if self.apicalPredictedSegmentDecrement != 0.0:
        self.apicalConnections.adjustActiveSynapses(
          apicalSegmentsToPunish, apicalReinforceCandidates,
          -self.apicalPredictedSegmentDecrement)

      # Only grow segments if there is basal *and* apical input.
      if len(basalGrowthCandidates) > 0 and len(apicalGrowthCandidates) > 0:
        self._learnOnNewSegments(self.basalConnections, self.rng,
                                 newSegmentCells, basalGrowthCandidates,
                                 self.initialPermanence, self.sampleSize,
                                 self.maxSynapsesPerSegment)
        self._learnOnNewSegments(self.apicalConnections, self.rng,
                                 newSegmentCells, apicalGrowthCandidates,
                                 self.initialPermanence, self.sampleSize,
                                 self.maxSynapsesPerSegment)


    # Save the results
    newActiveCells.sort()
    learningCells.sort()
    self.activeCells = newActiveCells
    self.winnerCells = learningCells
    self.predictedActiveCells = correctPredictedCells
    def compute(self,
                activeColumns,
                basalInput,
                apicalInput,
                basalGrowthCandidates=None,
                apicalGrowthCandidates=None,
                learn=True):
        """
    Perform one timestep. Use the basal and apical input to form a set of
    predictions, then activate the specified columns.

    @param activeColumns (numpy array)
    List of active columns

    @param basalInput (numpy array)
    List of active input bits for the basal dendrite segments

    @param apicalInput (numpy array)
    List of active input bits for the apical dendrite segments

    @param basalGrowthCandidates (numpy array or None)
    List of bits that the active cells may grow new basal synapses to.
    If None, the basalInput is assumed to be growth candidates.

    @param apicalGrowthCandidates (numpy array or None)
    List of bits that the active cells may grow new apical synapses to
    If None, the apicalInput is assumed to be growth candidates.

    @param learn (bool)
    Whether to grow / reinforce / punish synapses
    """

        if basalGrowthCandidates is None:
            basalGrowthCandidates = basalInput

        if apicalGrowthCandidates is None:
            apicalGrowthCandidates = apicalInput

        # Calculate predictions for this timestep
        (activeBasalSegments, matchingBasalSegments,
         basalPotentialOverlaps) = self._calculateSegmentActivity(
             self.basalConnections, basalInput, self.connectedPermanence,
             self.activationThresholdBasal, self.minThresholdBasal)

        (activeApicalSegments, matchingApicalSegments,
         apicalPotentialOverlaps) = self._calculateSegmentActivity(
             self.apicalConnections, apicalInput, self.connectedPermanence,
             self.activationThresholdApical, self.minThresholdApical)

        # Take union of predictions
        predictedCells = np.union1d(
            self.basalConnections.mapSegmentsToCells(activeBasalSegments),
            self.apicalConnections.mapSegmentsToCells(activeApicalSegments))

        # Calculate active cells
        (correctPredictedCells, burstingColumns) = np2.setCompare(
            predictedCells,
            activeColumns,
            predictedCells / self.cellsPerColumn,
            rightMinusLeft=True)
        newActiveCells = np.concatenate(
            (correctPredictedCells,
             np2.getAllCellsInColumns(burstingColumns, self.cellsPerColumn)))

        # Calculate learning
        (learningActiveBasalSegments, learningActiveApicalSegments,
         learningMatchingBasalSegments, learningMatchingApicalSegments,
         basalSegmentsToPunish, apicalSegmentsToPunish,
         newSegmentCells, learningCells) = self._calculateLearning(
             activeColumns, burstingColumns, correctPredictedCells,
             activeBasalSegments, activeApicalSegments, matchingBasalSegments,
             matchingApicalSegments, basalPotentialOverlaps,
             apicalPotentialOverlaps)

        if learn:
            # Learn on existing segments
            for learningSegments in (learningActiveBasalSegments,
                                     learningMatchingBasalSegments):
                self._learn(self.basalConnections, self.rng, learningSegments,
                            basalInput, basalGrowthCandidates,
                            basalPotentialOverlaps, self.initialPermanence,
                            self.sampleSize, self.permanenceIncrement,
                            self.permanenceDecrement,
                            self.maxSynapsesPerSegment)

            for learningSegments in (learningActiveApicalSegments,
                                     learningMatchingApicalSegments):
                self._learn(self.apicalConnections, self.rng, learningSegments,
                            apicalInput, apicalGrowthCandidates,
                            apicalPotentialOverlaps, self.initialPermanence,
                            self.sampleSize, self.permanenceIncrement,
                            self.permanenceDecrement,
                            self.maxSynapsesPerSegment)

            # Punish incorrect predictions
            if self.basalPredictedSegmentDecrement != 0.0:
                self.basalConnections.adjustActiveSynapses(
                    basalSegmentsToPunish, basalInput,
                    -self.basalPredictedSegmentDecrement)

            if self.apicalPredictedSegmentDecrement != 0.0:
                self.apicalConnections.adjustActiveSynapses(
                    apicalSegmentsToPunish, apicalInput,
                    -self.apicalPredictedSegmentDecrement)

            # Only grow segments if there is basal *and* apical input.
            if len(basalGrowthCandidates) > 0 and len(
                    apicalGrowthCandidates) > 0:
                self._learnOnNewSegments(self.basalConnections, self.rng,
                                         newSegmentCells,
                                         basalGrowthCandidates,
                                         self.initialPermanence,
                                         self.sampleSize,
                                         self.maxSynapsesPerSegment)
                self._learnOnNewSegments(self.apicalConnections, self.rng,
                                         newSegmentCells,
                                         apicalGrowthCandidates,
                                         self.initialPermanence,
                                         self.sampleSize,
                                         self.maxSynapsesPerSegment)

        # Save the results
        newActiveCells.sort()
        learningCells.sort()
        self.activeCells = newActiveCells
        self.winnerCells = learningCells
        self.predictedCells = predictedCells
        self.activeBasalSegments = activeBasalSegments
        self.activeApicalSegments = activeApicalSegments