Make diversity checking work with optimization

Proof of concept implementation, mixing diversity checking and optimization
may not be very effective in practice

5 files changed, 216 insertions, 146 deletions

diff --git a/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/ViatraReasoner.xtend b/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/ViatraReasoner.xtend
index c022beac..edcca676 100644
--- a/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/ViatraReasoner.xtend
+++ b/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/ViatraReasoner.xtend
@@ -19,6 +19,7 @@ import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.par
 import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.statecoder.IdentifierBasedStateCoderFactory

 import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.statecoder.NeighbourhoodBasedStateCoderFactory

 import hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner.dse.BestFirstStrategyForModelGeneration

+import hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner.dse.DiversityChecker

 import hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner.dse.LoggerSolutionFoundHandler

 import hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner.dse.ModelGenerationCompositeObjective

 import hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner.dse.PartialModelAsLogicInterpretation

@@ -122,7 +123,9 @@ class ViatraReasoner extends LogicReasoner {
                                 new SolutionStore()

                         }

                 solutionStore.registerSolutionFoundHandler(new LoggerSolutionFoundHandler(viatraConfig))

-                val solutionSaver = new ViatraReasonerSolutionSaver(newArrayList(extremalObjectives), numberOfRequiredSolutions)

+                val diversityChecker = DiversityChecker.of(viatraConfig.diversityRequirement)

+                val solutionSaver = new ViatraReasonerSolutionSaver(newArrayList(extremalObjectives), numberOfRequiredSolutions,

+                        diversityChecker)

                 val solutionCopier = solutionSaver.solutionCopier

                 solutionStore.withSolutionSaver(solutionSaver)

                 dse.solutionStore = solutionStore

@@ -197,14 +200,14 @@ class ViatraReasoner extends LogicReasoner {
                                 it.name = "SolutionCopyTime"

                                 it.value = (solutionCopier.getTotalCopierRuntime / 1000000) as int

                         ]

-                        if (strategy.solutionStoreWithDiversityDescriptor.isActive) {

+                        if (diversityChecker.isActive) {

                                 it.entries += createIntStatisticEntry => [

                                         it.name = "SolutionDiversityCheckTime"

-                                        it.value = (strategy.solutionStoreWithDiversityDescriptor.sumRuntime / 1000000) as int

+                                        it.value = (diversityChecker.totalRuntime / 1000000) as int

                                 ]

                                 it.entries += createRealStatisticEntry => [

                                         it.name = "SolutionDiversitySuccessRate"

-                                        it.value = strategy.solutionStoreWithDiversityDescriptor.successRate

+                                        it.value = diversityChecker.successRate

                                 ]

                         }

                 ]

diff --git a/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/BestFirstStrategyForModelGeneration.java b/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/BestFirstStrategyForModelGeneration.java
index 1234d54b..6ff867d7 100644
--- a/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/BestFirstStrategyForModelGeneration.java
+++ b/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/BestFirstStrategyForModelGeneration.java
@@ -75,7 +75,6 @@ public class BestFirstStrategyForModelGeneration implements IStrategy {
         // Running

         private PriorityQueue<TrajectoryWithFitness> trajectoiresToExplore;

         private SolutionStore solutionStore;

-        private SolutionStoreWithDiversityDescriptor solutionStoreWithDiversityDescriptor;

         private volatile boolean isInterrupted = false;

         private ModelResult modelResultByInternalSolver = null;

         private Random random = new Random();

@@ -96,9 +95,6 @@ public class BestFirstStrategyForModelGeneration implements IStrategy {
                 this.method = method;

         }

         

-        public SolutionStoreWithDiversityDescriptor getSolutionStoreWithDiversityDescriptor() {

-                return solutionStoreWithDiversityDescriptor;

-        }

         public int getNumberOfStatecoderFail() {

                 return numberOfStatecoderFail;

         }

@@ -117,8 +113,6 @@ public class BestFirstStrategyForModelGeneration implements IStrategy {
                         matchers.add(matcher);

                 }

                 

-                this.solutionStoreWithDiversityDescriptor = new SolutionStoreWithDiversityDescriptor(configuration.diversityRequirement);

-

                 final ObjectiveComparatorHelper objectiveComparatorHelper = context.getObjectiveComparatorHelper();

                 this.comparator = new Comparator<TrajectoryWithFitness>() {

                         @Override

@@ -142,7 +136,7 @@ public class BestFirstStrategyForModelGeneration implements IStrategy {
                 }

                 

                 final Fitness firstfitness = context.calculateFitness();

-                checkForSolution(firstfitness);

+                solutionStore.newSolution(context);

                 

                 final ObjectiveComparatorHelper objectiveComparatorHelper = context.getObjectiveComparatorHelper();

                 final Object[] firstTrajectory = context.getTrajectory().toArray(new Object[0]);

@@ -215,7 +209,7 @@ public class BestFirstStrategyForModelGeneration implements IStrategy {
                                         context.backtrack();

                                 } else {

                                         final Fitness nextFitness = context.calculateFitness();

-                                        checkForSolution(nextFitness);

+                                        solutionStore.newSolution(context);

                                         if (context.getDepth() > configuration.searchSpaceConstraints.maxDepth) {

                                                 logger.debug("Reached max depth.");

                                                 context.backtrack();

@@ -264,15 +258,6 @@ public class BestFirstStrategyForModelGeneration implements IStrategy {
                 return activationIds;

         }

 

-        private void checkForSolution(final Fitness fitness) {

-                if (fitness.isSatisifiesHardObjectives()) {

-                        if (solutionStoreWithDiversityDescriptor.isDifferent(context)) {

-                                solutionStoreWithDiversityDescriptor.newSolution(context);

-                                solutionStore.newSolution(context);

-                        }

-                }

-        }

-

         @Override

         public void interruptStrategy() {

                 isInterrupted = true;

diff --git a/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/DiversityChecker.xtend b/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/DiversityChecker.xtend
new file mode 100644
index 00000000..fb1b2066
--- /dev/null
+++ b/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/DiversityChecker.xtend
@@ -0,0 +1,184 @@
+package hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner.dse

+

+import com.google.common.collect.HashMultiset

+import com.google.common.collect.ImmutableSet

+import com.google.common.collect.Multiset

+import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.neighbourhood.AbstractNodeDescriptor

+import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.neighbourhood.NeighbourhoodWithTraces

+import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.neighbourhood.PartialInterpretation2ImmutableTypeLattice

+import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.PartialInterpretation

+import hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner.DiversityDescriptor

+import java.util.Collection

+import java.util.HashSet

+import java.util.Map

+import java.util.Set

+import org.eclipse.viatra.dse.base.ThreadContext

+import org.eclipse.xtend.lib.annotations.Accessors

+

+interface DiversityChecker {

+        public static val NO_DIVERSITY_CHECKER = new DiversityChecker {

+                override isActive() {

+                        false

+                }

+

+                override getTotalRuntime() {

+                        0

+                }

+

+                override getSuccessRate() {

+                        1.0

+                }

+

+                override newSolution(ThreadContext threadContext, Object solutionId, Collection<Object> dominatedSolutionIds) {

+                        true

+                }

+        }

+

+        def boolean isActive()

+

+        def long getTotalRuntime()

+

+        def double getSuccessRate()

+

+        def boolean newSolution(ThreadContext threadContext, Object solutionId, Collection<Object> dominatedSolutionIds)

+

+        static def of(DiversityDescriptor descriptor) {

+                if (descriptor.ensureDiversity) {

+                        new NodewiseDiversityChecker(descriptor)

+                } else {

+                        NO_DIVERSITY_CHECKER

+                }

+        }

+}

+

+abstract class AbstractDiversityChecker implements DiversityChecker {

+        val DiversityDescriptor descriptor

+        val PartialInterpretation2ImmutableTypeLattice solutionCoder = new PartialInterpretation2ImmutableTypeLattice

+

+        @Accessors(PUBLIC_GETTER) var long totalRuntime = 0

+        var int allCheckCount = 0

+        var int successfulCheckCount = 0

+

+        protected new(DiversityDescriptor descriptor) {

+                if (!descriptor.ensureDiversity) {

+                        throw new IllegalArgumentException(

+                                "Diversity description should enforce diversity or NO_DIVERSITY_CHECKER should be used instead.")

+                }

+                this.descriptor = descriptor

+        }

+

+        override isActive() {

+                true

+        }

+

+        override getTotalRuntime() {

+                throw new UnsupportedOperationException("TODO: auto-generated method stub")

+        }

+

+        override getSuccessRate() {

+                (allCheckCount as double) / (successfulCheckCount as double)

+        }

+

+        override newSolution(ThreadContext threadContext, Object solutionId, Collection<Object> dominatedSolutionIds) {

+                val start = System.nanoTime

+                val model = threadContext.model as PartialInterpretation

+                val representation = solutionCoder.createRepresentation(model, descriptor.range, descriptor.parallels,

+                        descriptor.maxNumber, descriptor.relevantTypes, descriptor.relevantRelations)

+                val isDifferent = internalNewSolution(representation, solutionId, dominatedSolutionIds)

+                totalRuntime += System.nanoTime - start

+                allCheckCount++

+                if (isDifferent) {

+                        successfulCheckCount++

+                }

+                isDifferent

+        }

+

+        protected abstract def boolean internalNewSolution(

+                NeighbourhoodWithTraces<Map<? extends AbstractNodeDescriptor, Integer>, AbstractNodeDescriptor> representation,

+                Object solutionId, Collection<Object> dominatedSolutionIds)

+}

+

+class NodewiseDiversityChecker extends AbstractDiversityChecker {

+        var Multiset<Integer> nodeCodes = HashMultiset.create

+        val Map<Object, Set<Integer>> tracedNodeCodes = newHashMap

+

+        new(DiversityDescriptor descriptor) {

+                super(descriptor)

+        }

+

+        override protected internalNewSolution(

+                NeighbourhoodWithTraces<Map<? extends AbstractNodeDescriptor, Integer>, AbstractNodeDescriptor> representation,

+                Object solutionId, Collection<Object> dominatedSolutionIds) {

+                val nodeCodesInSolution = ImmutableSet.copyOf(representation.modelRepresentation.keySet.map[hashCode])

+                val remainingNodeCodes = if (dominatedSolutionIds.empty) {

+                                nodeCodes

+                        } else {

+                                getRemainingNodeCodes(dominatedSolutionIds)

+                        }

+                val hasNewCode = nodeCodesInSolution.exists[!remainingNodeCodes.contains(it)]

+                if (hasNewCode) {

+                        nodeCodes = remainingNodeCodes

+                        nodeCodes.addAll(nodeCodesInSolution)

+                        for (dominatedSolutionId : dominatedSolutionIds) {

+                                tracedNodeCodes.remove(dominatedSolutionId)

+                        }

+                        tracedNodeCodes.put(solutionId, nodeCodesInSolution)

+                }

+                hasNewCode

+        }

+

+        private def getRemainingNodeCodes(Collection<Object> dominatedSolutionIds) {

+                // TODO Optimize multiset operations.

+                val copyOfNodeCodes = HashMultiset.create(nodeCodes)

+                for (dominatedSolutionId : dominatedSolutionIds) {

+                        val dominatedModelCode = tracedNodeCodes.get(dominatedSolutionId)

+                        if (dominatedModelCode === null) {

+                                throw new IllegalArgumentException("Unknown dominated solution: " + dominatedSolutionId)

+                        }

+                        copyOfNodeCodes.removeAll(dominatedModelCode)

+                }

+                copyOfNodeCodes

+        }

+}

+

+class GraphwiseDiversityChecker extends AbstractDiversityChecker {

+        var Set<Integer> modelCodes = newHashSet

+        val Map<Object, Integer> tracedModelCodes = newHashMap

+

+        new(DiversityDescriptor descriptor) {

+                super(descriptor)

+        }

+

+        override protected internalNewSolution(

+                NeighbourhoodWithTraces<Map<? extends AbstractNodeDescriptor, Integer>, AbstractNodeDescriptor> representation,

+                Object solutionId, Collection<Object> dominatedSolutionIds) {

+                val modelCodeOfSolution = representation.modelRepresentation.hashCode

+                val remainingModelCodes = if (dominatedSolutionIds.empty) {

+                                modelCodes

+                        } else {

+                                getRemainingModelCodes(dominatedSolutionIds)

+                        }

+                val isNewCode = !remainingModelCodes.contains(modelCodeOfSolution)

+                if (isNewCode) {

+                        modelCodes = remainingModelCodes

+                        modelCodes += modelCodeOfSolution

+                        for (dominatedSolutionId : dominatedSolutionIds) {

+                                tracedModelCodes.remove(dominatedSolutionId)

+                        }

+                        tracedModelCodes.put(solutionId, modelCodeOfSolution)

+                }

+                isNewCode

+        }

+

+        private def getRemainingModelCodes(Collection<Object> dominatedSolutionIds) {

+                val copyOfModelCodes = new HashSet(modelCodes)

+                for (dominatedSolutionId : dominatedSolutionIds) {

+                        val dominatedModelCode = tracedModelCodes.get(dominatedSolutionId)

+                        if (dominatedModelCode === null) {

+                                throw new IllegalArgumentException("Unknown dominated solution: " + dominatedSolutionId)

+                        }

+                        copyOfModelCodes -= dominatedModelCode

+                }

+                copyOfModelCodes

+        }

+}

diff --git a/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/SolutionStoreWithDiversityDescriptor.xtend b/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/SolutionStoreWithDiversityDescriptor.xtend
deleted file mode 100644
index 1e7f18a8..00000000
--- a/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/SolutionStoreWithDiversityDescriptor.xtend
+++ /dev/null
@@ -1,120 +0,0 @@
-package hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner.dse

-

-import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.neighbourhood.PartialInterpretation2ImmutableTypeLattice

-import hu.bme.mit.inf.dslreasoner.viatrasolver.partialinterpretationlanguage.partialinterpretation.PartialInterpretation

-import hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner.DiversityDescriptor

-import java.util.LinkedList

-import java.util.List

-import org.eclipse.viatra.dse.base.ThreadContext

-import java.util.HashSet

-import java.util.Set

-

-enum DiversityGranularity {

-        Nodewise, Graphwise

-}

-

-class SolutionStoreWithDiversityDescriptor {

-        val DiversityDescriptor descriptor

-        DiversityGranularity granularity

-        val PartialInterpretation2ImmutableTypeLattice solutionCoder = new PartialInterpretation2ImmutableTypeLattice

-        val Set<Integer> solutionCodeList = new HashSet

-        

-        var long runtime

-        var int allCheck

-        var int successfulCheck

-        

-        new(DiversityDescriptor descriptor) {

-                if(descriptor.ensureDiversity) {

-                        this.descriptor = descriptor

-                        this.granularity = DiversityGranularity::Nodewise

-                } else {

-                        this.descriptor = null

-                        this.granularity = DiversityGranularity::Nodewise

-                }

-        }

-        

-        def public isActive() {

-                descriptor!==null

-        }

-        

-        def getSumRuntime() {

-                return runtime

-        }

-        def getSuccessRate() {

-                return successfulCheck as double / allCheck

-        }

-        

-        def isDifferent(ThreadContext context) {

-                if(active) {

-                        val start = System.nanoTime

-                        val model = context.model as PartialInterpretation

-                        var boolean isDifferent

-                        if(this.granularity == DiversityGranularity::Graphwise) {

-                                val code = solutionCoder.createRepresentation(model,

-                                descriptor.range,

-                                descriptor.parallels,

-                                descriptor.maxNumber,

-                                descriptor.relevantTypes,

-                                descriptor.relevantRelations).modelRepresentation.hashCode

-                        

-                                isDifferent = !solutionCodeList.contains(code)

-                        } else if(this.granularity == DiversityGranularity::Nodewise){

-                                val codes = solutionCoder.createRepresentation(model,

-                                descriptor.range,

-                                descriptor.parallels,

-                                descriptor.maxNumber,

-                                descriptor.relevantTypes,

-                                descriptor.relevantRelations).modelRepresentation.keySet.map[hashCode].toList

-                                val differentCodes = codes.filter[!solutionCodeList.contains(it)]

-                                //println(differentCodes.size)

-                                

-                                isDifferent = differentCodes.size>=1

-                        } else {

-                                throw new UnsupportedOperationException('''Unsupported diversity type: «this.granularity»''')

-                        }

-                        

-                        runtime += System.nanoTime - start

-                        allCheck++

-                        if(isDifferent) { successfulCheck++ }

-                        return isDifferent

-                } else {

-                        allCheck++

-                        successfulCheck++

-                        return true

-                }

-        }

-        

-        def canBeDifferent(ThreadContext context) {

-                return true

-        }

-        

-        def newSolution(ThreadContext context) {

-                if(active) {

-                        val start = System.nanoTime

-                        val model = context.model as PartialInterpretation

-                        if(this.granularity == DiversityGranularity::Graphwise) {

-                                val code = solutionCoder.createRepresentation(model,

-                                descriptor.range,

-                                descriptor.parallels,

-                                descriptor.maxNumber,

-                                descriptor.relevantTypes,

-                                descriptor.relevantRelations).modelRepresentation.hashCode

-                        

-                                solutionCodeList += code.hashCode

-                        } else if(this.granularity == DiversityGranularity::Nodewise){

-                                val codes = solutionCoder.createRepresentation(model,

-                                descriptor.range,

-                                descriptor.parallels,

-                                descriptor.maxNumber,

-                                descriptor.relevantTypes,

-                                descriptor.relevantRelations).modelRepresentation.keySet.map[hashCode].toList

-                                

-                                solutionCodeList += codes.map[it.hashCode]

-                        } else {

-                                throw new UnsupportedOperationException('''Unsupported diversity type: «this.granularity»''')

-                        }

-                        

-                        runtime += System.nanoTime - start

-                }

-        }

-}
\ No newline at end of file
diff --git a/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/ViatraReasonerSolutionSaver.xtend b/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/ViatraReasonerSolutionSaver.xtend
index 17df6c55..6bffeb59 100644
--- a/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/ViatraReasonerSolutionSaver.xtend
+++ b/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/ViatraReasonerSolutionSaver.xtend
@@ -13,10 +13,11 @@ import org.eclipse.viatra.dse.solutionstore.SolutionStore.ISolutionSaver
 import org.eclipse.xtend.lib.annotations.Accessors
 
 /**
- * Based on {@link org.eclipse.viatra.dse.solutionstore.SolutionStore.BestSolutionSaver}.
+ * Based on {@link SolutionStore.BestSolutionSaver}.
  */
 class ViatraReasonerSolutionSaver implements ISolutionSaver {
         @Accessors val solutionCopier = new SolutionCopier
+        @Accessors val DiversityChecker diversityChecker
         val boolean hasExtremalObjectives
         val int numberOfRequiredSolutions
         val ObjectiveComparatorHelper comparatorHelper
@@ -24,7 +25,8 @@ class ViatraReasonerSolutionSaver implements ISolutionSaver {
 
         @Accessors(PUBLIC_SETTER) var Map<Object, Solution> solutionsCollection
 
-        new(IObjective[][] leveledExtremalObjectives, int numberOfRequiredSolutions) {
+        new(IObjective[][] leveledExtremalObjectives, int numberOfRequiredSolutions, DiversityChecker diversityChecker) {
+                this.diversityChecker = diversityChecker
                 comparatorHelper = new ObjectiveComparatorHelper(leveledExtremalObjectives)
                 hasExtremalObjectives = leveledExtremalObjectives.exists[!empty]
                 this.numberOfRequiredSolutions = numberOfRequiredSolutions
@@ -34,12 +36,15 @@ class ViatraReasonerSolutionSaver implements ISolutionSaver {
                 if (hasExtremalObjectives) {
                         saveBestSolutionOnly(context, id, solutionTrajectory)
                 } else {
-                        basicSaveSolution(context, id, solutionTrajectory)
+                        saveAnyDiverseSolution(context, id, solutionTrajectory)
                 }
         }
 
         private def saveBestSolutionOnly(ThreadContext context, Object id, SolutionTrajectory solutionTrajectory) {
                 val fitness = context.lastFitness
+                if (!fitness.satisifiesHardObjectives) {
+                        return false
+                }
                 val dominatedTrajectories = newArrayList
                 for (entry : trajectories.entrySet) {
                         val isLastFitnessBetter = comparatorHelper.compare(fitness, entry.value)
@@ -54,9 +59,12 @@ class ViatraReasonerSolutionSaver implements ISolutionSaver {
                 if (dominatedTrajectories.size == 0 && !needsMoreSolutionsWithSameFitness) {
                         return false
                 }
+                if (!diversityChecker.newSolution(context, id, dominatedTrajectories.map[solution.stateCode])) {
+                        return false
+                }
                 // We must save the new trajectory before removing dominated trajectories
                 // to avoid removing the current solution when it is reachable only via dominated trajectories.
-                val solutionSaved = basicSaveSolution(context, id, solutionTrajectory)
+                val solutionSaved = basicSaveSolution(context, id, solutionTrajectory, fitness)
                 for (dominatedTrajectory : dominatedTrajectories) {
                         trajectories -= dominatedTrajectory
                         val dominatedSolution = dominatedTrajectory.solution
@@ -73,8 +81,18 @@ class ViatraReasonerSolutionSaver implements ISolutionSaver {
                 solutionSaved
         }
 
-        private def basicSaveSolution(ThreadContext context, Object id, SolutionTrajectory solutionTrajectory) {
+        private def saveAnyDiverseSolution(ThreadContext context, Object id, SolutionTrajectory solutionTrajectory) {
                 val fitness = context.lastFitness
+                if (!fitness.satisifiesHardObjectives) {
+                        return false
+                }
+                if (!diversityChecker.newSolution(context, id, emptyList)) {
+                        return false
+                }
+                basicSaveSolution(context, id, solutionTrajectory, fitness)
+        }
+
+        private def basicSaveSolution(ThreadContext context, Object id, SolutionTrajectory solutionTrajectory, Fitness fitness) {
                 var boolean solutionSaved = false
                 var dseSolution = solutionsCollection.get(id)
                 if (dseSolution === null) {