dataset predictions fixed

1 files changed, 23 insertions, 0 deletions

diff --git a/lib/dataset.rb b/lib/dataset.rb
index 41d7b5c..78f5633 100644
--- a/lib/dataset.rb
+++ b/lib/dataset.rb
@@ -86,6 +86,10 @@ module OpenTox
       features.select{|f| f._type.match("SubstanceProperty")}
     end
 
+    def prediction_features
+      features.select{|f| f._type.match("Prediction")}
+    end
+
     # Writers
 
     # Add a value for a given substance and feature
@@ -352,6 +356,25 @@ module OpenTox
       sdf
     end
 
+    def predictions
+      predictions = {}
+      substances.each do |s| 
+        predictions[s] ||= {}
+        prediction_feature = prediction_features.first
+        predictions[s][:value] = values(s,prediction_feature).first
+        predictions[s][:warnings] = []
+        warnings_features.each { |w| predictions[s][:warnings] += values(s,w) }
+        if predictions[s][:value] and prediction_feature.is_a? NominalLazarPrediction
+          prediction_feature.accept_values.each do |v|
+            f = LazarPredictionProbability.find_by(:name => v, :model_id => prediction_feature.model_id, :training_feature_id => prediction_feature.training_feature_id)
+            predictions[s][:probabilities] ||= {}
+            predictions[s][:probabilities][v] = values(s,f).first
+          end
+        end
+      end
+      predictions
+    end
+
     # Dataset operations
 
     # Merge an array of datasets