Skip to content

Commit

Permalink
Replace Panics with error returns to BernoulliNBClassifier Fit method…
Browse files Browse the repository at this point in the history
… to satisfy base.Classifier interface
  • Loading branch information
JustinJudd committed Jul 17, 2019
1 parent c3cae57 commit bffc4a5
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions naive/bernoulli_nb.go
Original file line number Diff line number Diff line change
Expand Up @@ -182,22 +182,22 @@ func NewBernoulliNBClassifier() *BernoulliNBClassifier {

// Fill data matrix with Bernoulli Naive Bayes model. All values
// necessary for calculating prior probability and p(f_i)
func (nb *BernoulliNBClassifier) Fit(X base.FixedDataGrid) {
func (nb *BernoulliNBClassifier) Fit(X base.FixedDataGrid) error {

// Check that all Attributes are binary
classAttrs := X.AllClassAttributes()
allAttrs := X.AllAttributes()
featAttrs := base.AttributeDifference(allAttrs, classAttrs)
for i := range featAttrs {
if _, ok := featAttrs[i].(*base.BinaryAttribute); !ok {
panic(fmt.Sprintf("%v: Should be BinaryAttribute", featAttrs[i]))
return fmt.Errorf("%v: Should be BinaryAttribute", featAttrs[i])
}
}
featAttrSpecs := base.ResolveAttributes(X, featAttrs)

// Check that only one classAttribute is defined
if len(classAttrs) != 1 {
panic("Only one class Attribute can be used")
return fmt.Errorf("Only one class Attribute can be used")
}

// Number of features and instances in this training set
Expand Down Expand Up @@ -258,6 +258,7 @@ func (nb *BernoulliNBClassifier) Fit(X base.FixedDataGrid) {
}

nb.fitOn = base.NewStructuralCopy(X)
return nil
}

// Use trained model to predict test vector's class. The following
Expand Down

0 comments on commit bffc4a5

Please sign in to comment.