Simplify QueryGraph interface (#118)

## What is the goal of this PR?

Simplify the user's job of building QueryGraphs as much as possible

## What are the changes implemented in this PR?

- Make `'solution'` class mandatory when building QueryGraphs, this reduces the possibility of error
- Make methods chainable and therefore more readable

Author: jmsfltchr

kglib/kgcn/examples/diagnosis/BUILD

@@ -8,6 +8,7 @@ py_test(
     ],
     deps = [
         "diagnosis",
+        "//kglib/utils/graph/test",
         requirement('numpy'),
         requirement('networkx'),
         requirement('decorator'),

kglib/kgcn/examples/diagnosis/diagnosis.py

@@ -34,13 +34,13 @@
 URI = "localhost:48555"
 
 # Existing elements in the graph are those that pre-exist in the graph, and should be predicted to continue to exist
-PREEXISTS = dict(solution=0)
+PREEXISTS = 0
 
 # Candidates are neither present in the input nor in the solution, they are negative samples
-CANDIDATE = dict(solution=1)
+CANDIDATE = 1
 
 # Elements to infer are the graph elements whose existence we want to predict to be true, they are positive samples
-TO_INFER = dict(solution=2)
+TO_INFER = 2
 
 # Categorical Attribute types and the values of their categories
 CATEGORICAL_ATTRIBUTES = {'name': ['Diabetes Type II', 'Multiple Sclerosis', 'Blurred vision', 'Fatigue', 'Cigarettes',
@@ -138,19 +138,23 @@ def create_concept_graphs(example_indices, grakn_session):
             # Build a graph from the queries, samplers, and query graphs
             graph = build_graph_from_queries(graph_query_handles, tx, infer=infer)
 
-        # Remove label leakage - change type labels that indicate candidates into non-candidates
-        for data in multidigraph_data_iterator(graph):
-            for label_to_obfuscate, with_label in TYPES_AND_ROLES_TO_OBFUSCATE.items():
-                if data['type'] == label_to_obfuscate:
-                    data.update(type=with_label)
-                    break
+        obfuscate_labels(graph, TYPES_AND_ROLES_TO_OBFUSCATE)
 
         graph.name = example_id
         graphs.append(graph)
 
     return graphs
 
 
+def obfuscate_labels(graph, types_and_roles_to_obfuscate):
+    # Remove label leakage - change type labels that indicate candidates into non-candidates
+    for data in multidigraph_data_iterator(graph):
+        for label_to_obfuscate, with_label in types_and_roles_to_obfuscate.items():
+            if data['type'] == label_to_obfuscate:
+                data.update(type=with_label)
+                break
+
+
 def get_query_handles(example_id):
     """
     Creates an iterable, each element containing a Graql query, a function to sample the answers, and a QueryGraph
@@ -174,15 +178,13 @@ def get_query_handles(example_id):
            get;''')
 
     vars = p, par, ps, d, diag, n = 'p', 'par', 'ps', 'd', 'diag', 'n'
-    g = QueryGraph()
-    g.add_vars(*vars, **PREEXISTS)
-    g.add_role_edge(ps, p, 'child', **PREEXISTS)
-    g.add_role_edge(ps, par, 'parent', **PREEXISTS)
-    g.add_role_edge(diag, par, 'patient', **PREEXISTS)
-    g.add_role_edge(diag, d, 'diagnosed-disease', **PREEXISTS)
-    g.add_has_edge(d, n, **PREEXISTS)
-
-    hereditary_query_graph = g
+    hereditary_query_graph = (QueryGraph()
+                              .add_vars(vars, PREEXISTS)
+                              .add_role_edge(ps, p, 'child', PREEXISTS)
+                              .add_role_edge(ps, par, 'parent', PREEXISTS)
+                              .add_role_edge(diag, par, 'patient', PREEXISTS)
+                              .add_role_edge(diag, d, 'diagnosed-disease', PREEXISTS)
+                              .add_has_edge(d, n, PREEXISTS))
 
     # === Consumption Feature ===
     consumption_query = inspect.cleandoc(f'''match
@@ -192,26 +194,22 @@ def get_query_handles(example_id):
            has units-per-week $u; get;''')
 
     vars = p, s, n, c, u = 'p', 's', 'n', 'c', 'u'
-    g = QueryGraph()
-    g.add_vars(*vars, **PREEXISTS)
-    g.add_has_edge(s, n, **PREEXISTS)
-    g.add_role_edge(c, p, 'consumer', **PREEXISTS)
-    g.add_role_edge(c, s, 'consumed-substance', **PREEXISTS)
-    g.add_has_edge(c, u, **PREEXISTS)
-
-    consumption_query_graph = g
+    consumption_query_graph = (QueryGraph()
+                               .add_vars(vars, PREEXISTS)
+                               .add_has_edge(s, n, PREEXISTS)
+                               .add_role_edge(c, p, 'consumer', PREEXISTS)
+                               .add_role_edge(c, s, 'consumed-substance', PREEXISTS)
+                               .add_has_edge(c, u, PREEXISTS))
 
     # === Age Feature ===
     person_age_query = inspect.cleandoc(f'''match 
             $p isa person, has example-id {example_id}, has age $a; 
             get;''')
 
     vars = p, a = 'p', 'a'
-    g = QueryGraph()
-    g.add_vars(*vars, **PREEXISTS)
-    g.add_has_edge(p, a, **PREEXISTS)
-
-    person_age_query_graph = g
+    person_age_query_graph = (QueryGraph()
+                              .add_vars(vars, PREEXISTS)
+                              .add_has_edge(p, a, PREEXISTS))
 
     # === Risk Factors Feature ===
     risk_factor_query = inspect.cleandoc(f'''match 
@@ -221,12 +219,10 @@ def get_query_handles(example_id):
             get;''')
 
     vars = p, d, r = 'p', 'd', 'r'
-    g = QueryGraph()
-    g.add_vars(*vars, **PREEXISTS)
-    g.add_role_edge(r, p, 'person-at-risk', **PREEXISTS)
-    g.add_role_edge(r, d, 'risked-disease', **PREEXISTS)
-
-    risk_factor_query_graph = g
+    risk_factor_query_graph = (QueryGraph()
+                               .add_vars(vars, PREEXISTS)
+                               .add_role_edge(r, p, 'person-at-risk', PREEXISTS)
+                               .add_role_edge(r, d, 'risked-disease', PREEXISTS))
 
     # === Diagnosis ===
     diagnosis_query = inspect.cleandoc(f'''match
@@ -239,26 +235,22 @@ def get_query_handles(example_id):
            get;''')
 
     vars = p, s, sn, d, dn, sp, sev, c = 'p', 's', 'sn', 'd', 'dn', 'sp', 'sev', 'c'
-    g = QueryGraph()
-    g.add_vars(*vars, **PREEXISTS)
-    g.add_has_edge(s, sn, **PREEXISTS)
-    g.add_has_edge(d, dn, **PREEXISTS)
-    g.add_role_edge(sp, s, 'presented-symptom', **PREEXISTS)
-    g.add_has_edge(sp, sev, **PREEXISTS)
-    g.add_role_edge(sp, p, 'symptomatic-patient', **PREEXISTS)
-    g.add_role_edge(c, s, 'effect', **PREEXISTS)
-    g.add_role_edge(c, d, 'cause', **PREEXISTS)
-
-    base_query_graph = g
-
-    g = copy.copy(base_query_graph)
+    base_query_graph = (QueryGraph()
+                        .add_vars(vars, PREEXISTS)
+                        .add_has_edge(s, sn, PREEXISTS)
+                        .add_has_edge(d, dn, PREEXISTS)
+                        .add_role_edge(sp, s, 'presented-symptom', PREEXISTS)
+                        .add_has_edge(sp, sev, PREEXISTS)
+                        .add_role_edge(sp, p, 'symptomatic-patient', PREEXISTS)
+                        .add_role_edge(c, s, 'effect', PREEXISTS)
+                        .add_role_edge(c, d, 'cause', PREEXISTS))
 
     diag, d, p = 'diag', 'd', 'p'
-    g.add_vars(diag, **TO_INFER)
-    g.add_role_edge(diag, d, 'diagnosed-disease', **TO_INFER)
-    g.add_role_edge(diag, p, 'patient', **TO_INFER)
 
-    diagnosis_query_graph = g
+    diagnosis_query_graph = (copy.copy(base_query_graph)
+                             .add_vars([diag], TO_INFER)
+                             .add_role_edge(diag, d, 'diagnosed-disease', TO_INFER)
+                             .add_role_edge(diag, p, 'patient', TO_INFER))
 
     # === Candidate Diagnosis ===
     candidate_diagnosis_query = inspect.cleandoc(f'''match
@@ -270,13 +262,10 @@ def get_query_handles(example_id):
            $diag(candidate-patient: $p, candidate-diagnosed-disease: $d) isa candidate-diagnosis; 
            get;''')
 
-    g = copy.copy(base_query_graph)
-
-    diag, d, p = 'diag', 'd', 'p'
-    g.add_vars(diag, **CANDIDATE)
-    g.add_role_edge(diag, d, 'candidate-diagnosed-disease', **CANDIDATE)
-    g.add_role_edge(diag, p, 'candidate-patient', **CANDIDATE)
-    candidate_diagnosis_query_graph = g
+    candidate_diagnosis_query_graph = (copy.copy(base_query_graph)
+                                       .add_vars([diag], CANDIDATE)
+                                       .add_role_edge(diag, d, 'candidate-diagnosed-disease', CANDIDATE)
+                                       .add_role_edge(diag, p, 'candidate-patient', CANDIDATE))
 
     return [
         (diagnosis_query, lambda x: x, diagnosis_query_graph),

kglib/kgcn/examples/diagnosis/diagnosis_test.py

@@ -24,8 +24,9 @@
 import networkx as nx
 import numpy as np
 
-from kglib.kgcn.examples.diagnosis.diagnosis import write_predictions_to_grakn
+from kglib.kgcn.examples.diagnosis.diagnosis import write_predictions_to_grakn, obfuscate_labels
 from kglib.utils.grakn.object.thing import Thing
+from kglib.utils.graph.test.case import GraphTestCase
 
 
 class TestWritePredictionsToGrakn(unittest.TestCase):
@@ -90,5 +91,34 @@ def test_query_made_only_if_relation_wins(self):
         tx.commit.assert_called()
 
 
+class TestObfuscateLabels(GraphTestCase):
+
+    def test_labels_obfuscated_as_expected(self):
+
+        graph = nx.MultiDiGraph()
+
+        graph.add_node(0, type='person')
+        graph.add_node(1, type='disease')
+        graph.add_node(2, type='candidate-diagnosis')
+
+        graph.add_edge(2, 0, type='candidate-patient')
+        graph.add_edge(2, 1, type='candidate-diagnosed-disease')
+
+        obfuscate_labels(graph, {'candidate-diagnosis': 'diagnosis',
+                                 'candidate-patient': 'patient',
+                                 'candidate-diagnosed-disease': 'diagnosed-disease'})
+
+        expected_graph = nx.MultiDiGraph()
+        expected_graph.add_node(0, type='person')
+        expected_graph.add_node(1, type='disease')
+        expected_graph.add_node(2, type='diagnosis')
+
+        expected_graph.add_edge(2, 0, type='patient')
+        expected_graph.add_edge(2, 1, type='diagnosed-disease')
+
+        self.assertGraphsEqual(graph, expected_graph)
+
+
+
 if __name__ == "__main__":
     unittest.main()

kglib/utils/graph/query/query_graph.py

@@ -25,42 +25,45 @@ class QueryGraph(nx.MultiDiGraph):
     A custom graph to represent a query. Has additional helper methods specific to adding Graql patterns.
     """
 
-    def add_vars(self, *vars, **attr):
+    def add_vars(self, vars, solution):
         """
         Add Graql variables, stored as nodes in the graph
         Args:
-            *vars: String variables
-            **attr: Properties to be added to the data stored on each variable node
+            vars: String variables
+            solution: Indicator of the ground truth class that the variables belongs to
 
         Returns:
-            None
+            self
         """
         for var in vars:
-            self.add_node(var, **attr)
+            self.add_node(var, solution=solution)
+        return self
 
-    def add_has_edge(self, owner_var, attribute_var, **attr):
+    def add_has_edge(self, owner_var, attribute_var, solution):
         """
         Add a "has" edge to represent ownership of an attribute
         Args:
             owner_var: The variable of the owner
             attribute_var: The variable of the owned attribute
-            **attr: Properties to be added to the data stored on the "has" edge added
+            solution: Indicator of the ground truth class that the edge belongs to
 
         Returns:
-            None
+            self
         """
-        self.add_edge(owner_var, attribute_var, type='has', **attr)
+        self.add_edge(owner_var, attribute_var, type='has', solution=solution)
+        return self
 
-    def add_role_edge(self, relation_var, roleplayer_var, role_label, **attr):
+    def add_role_edge(self, relation_var, roleplayer_var, role_label, solution):
         """
         Add an edge to represent the role a variable plays in a relation
         Args:
             relation_var: The variable of the relation
             roleplayer_var: The variable of the roleplayer in the relation
             role_label: The role the roleplayer plays in the relation
-            **attr: Properties to be added to the data stored on the role edge added
+            solution: Indicator of the ground truth class that the edge belongs to
 
         Returns:
-            None
+            self
         """
-        self.add_edge(relation_var, roleplayer_var, type=role_label, **attr)
+        self.add_edge(relation_var, roleplayer_var, type=role_label, solution=solution)
+        return self

kglib/utils/graph/query/query_graph_test.py

@@ -24,27 +24,32 @@
 
 class TestQueryGraph(unittest.TestCase):
 
+    def test_add_single_var_adds_variable_node_as_expected(self):
+        g = QueryGraph()
+        g.add_vars(['a'], 0)
+        self.assertDictEqual({'solution': 0}, g.nodes['a'])
+
     def test_add_vars_adds_variable_nodes_as_expected(self):
         g = QueryGraph()
-        g.add_vars('a', 'b')
+        g.add_vars(['a', 'b'], 0)
         nodes = {node for node in g.nodes}
         self.assertSetEqual({'a', 'b'}, nodes)
 
     def test_add_has_edge_adds_edge_as_expected(self):
         g = QueryGraph()
         g.add_vars('a', 'b')
-        g.add_has_edge('a', 'b')
+        g.add_has_edge('a', 'b', 0)
         edges = [edge for edge in g.edges]
         self.assertEqual(1, len(edges))
-        self.assertEqual('has', g.edges['a', 'b', 0]['type'])
+        self.assertDictEqual({'type': 'has', 'solution': 0}, g.edges['a', 'b', 0])
 
     def test_add_role_edge_adds_role_as_expected(self):
         g = QueryGraph()
         g.add_vars('a', 'b')
-        g.add_role_edge('a', 'b', 'role')
+        g.add_role_edge('a', 'b', 'role_label', 1)
         edges = [edge for edge in g.edges]
         self.assertEqual(1, len(edges))
-        self.assertEqual('role', g.edges['a', 'b', 0]['type'])
+        self.assertDictEqual({'type': 'role_label', 'solution': 1}, g.edges['a', 'b', 0])
 
 
 if __name__ == "__main__":