simple implementation for other feature types #120

tests/test_dataset.py

@@ -3,6 +3,7 @@ from tempfile import TemporaryDirectory
 from django.test import TestCase
 from epdb.logic import PackageManager
 from epdb.models import Reaction, Compound, User, Rule, Package
+from utilities.chem import FormatConverter
 from utilities.ml import RuleBasedDataset, EnviFormerDataset
 
 
@@ -102,6 +103,12 @@ class DatasetTest(TestCase):
         print(class_ds.df.head(5))
         print(products[:5])
 
+    def test_extra_features(self):
+        reactions = [r for r in Reaction.objects.filter(package=self.BBD_SUBSET)]
+        applicable_rules = [r for r in Rule.objects.filter(package=self.BBD_SUBSET)]
+        ds = RuleBasedDataset.generate_dataset(reactions, applicable_rules, feat_funcs=[FormatConverter.maccs, FormatConverter.morgan])
+        print(ds.shape)
+
     def test_to_arff(self):
         """Test exporting the arff version of the dataset"""
         ds, reactions, rules = self.generate_rule_dataset()